diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py index d89351a3dd4..cfefeccd0e1 100644 --- a/miles/rollout/generate_hub/agentic_tool_call.py +++ b/miles/rollout/generate_hub/agentic_tool_call.py @@ -129,6 +129,7 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: return GenerateFnOutput(samples=sample) sample = samples[0] + _mark_limits_exceeded_truncated(sample, agent_metadata) logger.debug( f"{log_prefix} server-merged sample ready: " f"tokens={len(sample.tokens)} response_length={sample.response_length} " @@ -137,6 +138,16 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: return GenerateFnOutput(samples=sample) +def _mark_limits_exceeded_truncated( + samples: Sample | list[Sample], + agent_metadata: dict[str, Any] | None, +) -> None: + """Record turn-budget exhaustion even when the final model turn completed.""" + if (agent_metadata or {}).get("exit_status") == "LimitsExceeded": + final_sample = samples if isinstance(samples, Sample) else samples[-1] + final_sample.status = Sample.Status.TRUNCATED + + def build_agent_function_kwargs( args: argparse.Namespace, *, diff --git a/tests/fast/rollout/generate_hub/test_agentic_turn_budget.py b/tests/fast/rollout/generate_hub/test_agentic_turn_budget.py new file mode 100644 index 00000000000..61e870edd65 --- /dev/null +++ b/tests/fast/rollout/generate_hub/test_agentic_turn_budget.py @@ -0,0 +1,98 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from miles.rollout.base_types import GenerateFnInput +from miles.rollout.generate_hub import agentic_tool_call +from miles.rollout.generate_hub.agentic_tool_call import _mark_limits_exceeded_truncated +from miles.rollout.session.session_types import MergedSessionSample +from miles.utils.types import Sample + + +@pytest.mark.parametrize("as_list", [False, True], ids=["merged", "multi-sample"]) +def test_limits_exceeded_marks_only_final_sample_truncated(as_list: bool) -> None: + samples = [ + Sample(status=Sample.Status.COMPLETED), + Sample(status=Sample.Status.COMPLETED), + ] + value = samples if as_list else samples[-1] + + _mark_limits_exceeded_truncated(value, {"exit_status": "LimitsExceeded"}) + + assert samples[-1].status == Sample.Status.TRUNCATED + if as_list: + assert samples[0].status == Sample.Status.COMPLETED + + +@pytest.mark.parametrize("agent_metadata", [None, {}, {"exit_status": "Submitted"}]) +def test_other_exit_statuses_remain_completed(agent_metadata: dict | None) -> None: + sample = Sample(status=Sample.Status.COMPLETED) + + _mark_limits_exceeded_truncated(sample, agent_metadata) + + assert sample.status == Sample.Status.COMPLETED + + +@pytest.mark.parametrize("exit_status", ["LimitsExceeded", "Submitted"]) +@pytest.mark.parametrize("empty_session", [False, True]) +def test_generate_preserves_server_merged_sample(monkeypatch, exit_status, empty_session): + merged = MergedSessionSample( + tokens=[1, 2, 3, 4], + response="", + response_length=3, + loss_mask=[1, 0, 1], + rollout_log_probs=[-0.1, 0.0, -0.2], + status="completed", + weight_versions=["v1", "v2"], + ) + tracer = SimpleNamespace( + session_id="test-session", + session_server_instance_id="test-instance", + base_url="http://session-server/v1", + collect_merged_sample=AsyncMock( + return_value=(None if empty_session else merged, {"session_key": "preserved"}) + ), + ) + monkeypatch.setattr(agentic_tool_call.OpenAIEndpointTracer, "create", AsyncMock(return_value=tracer)) + agent = AsyncMock(return_value={"exit_status": exit_status, "reward": 1.0}) + monkeypatch.setattr(agentic_tool_call, "load_function", lambda _: agent) + args = SimpleNamespace( + session_server_ip="127.0.0.1", + session_server_port=1234, + custom_agent_function_path="test-agent", + generate_multi_samples=False, + ) + original = Sample(prompt="task", metadata={"task_id": "task-1"}) + result = asyncio.run( + agentic_tool_call.generate( + GenerateFnInput( + state=SimpleNamespace(args=args, tokenizer=None), + sample=original, + sampling_params={}, + evaluation=False, + ) + ) + ) + + tracer.collect_merged_sample.assert_awaited_once() + sample = result.samples + assert isinstance(sample, Sample) + assert original.status == Sample.Status.PENDING + assert original.metadata == {"task_id": "task-1"} + if empty_session: + assert sample.status == Sample.Status.ABORTED + return + + expected_status = Sample.Status.TRUNCATED if exit_status == "LimitsExceeded" else Sample.Status.COMPLETED + assert sample.status == expected_status + for field in ("tokens", "response", "response_length", "loss_mask", "rollout_log_probs", "weight_versions"): + assert getattr(sample, field) == getattr(merged, field) + assert sample.metadata == { + "task_id": "task-1", + "response_decoded": False, + "exit_status": exit_status, + "reward": 1.0, + "session_key": "preserved", + } diff --git a/tests/fast/rollout/generate_hub/test_multi_turn.py b/tests/fast/rollout/generate_hub/test_multi_turn.py index 38145f93d2c..96d01e7efda 100644 --- a/tests/fast/rollout/generate_hub/test_multi_turn.py +++ b/tests/fast/rollout/generate_hub/test_multi_turn.py @@ -668,6 +668,22 @@ def test_agent_returns_none_metadata_unchanged(self, variant, generation_env): assert s.metadata.get("instance_id") == "test-123" assert "reward" not in s.metadata + @pytest.mark.parametrize( + "generation_env", + [{"args_kwargs": {"agentic_return_metadata": {"exit_status": "LimitsExceeded"}}}], + indirect=True, + ) + @pytest.mark.parametrize("variant", ["agentic_tool_call_single_sample"]) + def test_limits_exceeded_marks_final_sample_truncated(self, variant, generation_env): + generation_env.mock_server.process_fn = TwoTurnStub.process_fn + + result = _run_generate(variant, generation_env, make_sample(prompt=TwoTurnStub.PROMPT)) + + samples = listify(result.sample) + assert all(sample.status == Sample.Status.COMPLETED for sample in samples[:-1]) + assert samples[-1].status == Sample.Status.TRUNCATED + assert samples[-1].metadata["exit_status"] == "LimitsExceeded" + def test_session_server_identity_forwarded_to_agent_metadata(self, variant, generation_env): from miles.utils.test_utils import mock_tools