From 15c8eb3696d58c129a2fa0793397636453b0057b Mon Sep 17 00:00:00 2001 From: "tongyuxuan.361" Date: Tue, 15 Sep 2026 13:56:50 +0800 Subject: [PATCH] [framework, gateway, sandbox] feat: add runtime integration for Laminar async rollouts --- .../concepts/gateway-and-trajectories.md | 15 +- docs/source/concepts/sandbox.md | 4 +- .../framework/test_entry_adapter_on_cpu.py | 313 +++++++++++ .../test_generate_sequences_on_cpu.py | 526 +++++++++++++++++- .../test_worker_concurrency_on_cpu.py | 231 ++++++++ .../gateway/test_gateway_actor_on_cpu.py | 78 +++ .../gateway/test_gateway_manager_on_cpu.py | 71 +++ .../test_session_multiple_chains_on_cpu.py | 69 +++ .../sandbox/plugin_probe_provider.py | 18 + .../test_sandbox_provider_discovery_on_cpu.py | 46 ++ .../tasks/test_swe_agent_error_on_cpu.py | 46 ++ .../uni_agent/tasks/test_swe_reward_on_cpu.py | 72 +++ uni_agent/framework/entry.py | 39 +- uni_agent/framework/framework.py | 178 ++++-- uni_agent/gateway/gateway.py | 23 +- uni_agent/gateway/manager.py | 8 +- uni_agent/gateway/session/codec.py | 1 + uni_agent/gateway/session/session.py | 43 +- uni_agent/sandbox/registry.py | 29 +- uni_agent/tasks/swe_bench/reward.py | 3 + uni_agent/tasks/swe_bench/task.py | 4 + 21 files changed, 1739 insertions(+), 78 deletions(-) create mode 100644 tests/uni_agent/framework/test_entry_adapter_on_cpu.py create mode 100644 tests/uni_agent/framework/test_worker_concurrency_on_cpu.py create mode 100644 tests/uni_agent/sandbox/plugin_probe_provider.py create mode 100644 tests/uni_agent/sandbox/test_sandbox_provider_discovery_on_cpu.py create mode 100644 tests/uni_agent/tasks/test_swe_agent_error_on_cpu.py create mode 100644 tests/uni_agent/tasks/test_swe_reward_on_cpu.py diff --git a/docs/source/concepts/gateway-and-trajectories.md b/docs/source/concepts/gateway-and-trajectories.md index a881d440b..b372a3a0e 100644 --- a/docs/source/concepts/gateway-and-trajectories.md +++ b/docs/source/concepts/gateway-and-trajectories.md @@ -289,7 +289,20 @@ Important knobs include: - `agent_runners`: Runner import paths and arguments. With multiple entries, each registry key must match the sample's `agent_name`. - `dispatch_mode`: inline async execution or Ray tasks. -- `max_concurrent_sessions`: per-Runner concurrency limit. +- `max_concurrent_sessions`: per-Runner concurrency limit. This is a wide + worker-side safety net. `AgentFrameworkWorker` derives its Ray + `max_concurrency` as the larger of 1000 (Ray's async default) and the sum + of positive Runner limits; each Runner's semaphore still enforces its own + session limit. Ray counts + unfinished batch RPCs, and a batch can remain live for one slow session, + so this value is not divided by batch size or `rollout.n`. + An omitted or zero limit keeps that Runner's + sessions unbounded and adds nothing to the derived RPC limit; it does not + make Ray RPC concurrency unlimited. + Async-RL dispatch pause is owned by Laminar: `LLMServerStateTracker` snapshots replica + `kv_cache_usage`, and `RolloutDispatcher` skips new batches when every + replica is at/above the KV threshold. Slight overshoot from the ~1s poll + lag is expected. Do not add a driver-side inflight cap in this adapter. - `log_dir`: runtime log root. Sessions with a global step write `framework.log`, `task.log`, and trajectory artifacts under `step_//`; Sessions whose `global_steps` is `None` write directly under `/`. - `rollout.n`: sessions per prompt. - `rollout.multi_turn.format`: model-specific Tool parser. diff --git a/docs/source/concepts/sandbox.md b/docs/source/concepts/sandbox.md index 20c6acf5d..ca50c265b 100644 --- a/docs/source/concepts/sandbox.md +++ b/docs/source/concepts/sandbox.md @@ -136,7 +136,7 @@ def from_config(cls, config: SandboxConfig): ) ``` -Register the lazy module in `SANDBOX_MODULES`: +In-tree providers also register the lazy module in `SANDBOX_MODULES`: ```python SANDBOX_MODULES["my_backend"] = "my_package.sandbox" @@ -144,6 +144,8 @@ SANDBOX_MODULES["my_backend"] = "my_package.sandbox" The module is imported only when the provider is selected, so optional SDKs do not affect other backends. +Out-of-tree overlays (private backends) must not edit `SANDBOX_MODULES`. Set `UNI_AGENT_SANDBOX_PLUGINS` to a comma-separated list of importable modules; each module calls `register_sandbox` at import time. Missing modules fail closed with `ImportError`. Ray Jobs are not pip-installed, so this variable must be present in `runtime_env.env_vars` on every worker — a driver-only export is not enough. + ## Implementation Rules - Implement `_exec()`, not `exec()`. diff --git a/tests/uni_agent/framework/test_entry_adapter_on_cpu.py b/tests/uni_agent/framework/test_entry_adapter_on_cpu.py new file mode 100644 index 000000000..2eb9b8c34 --- /dev/null +++ b/tests/uni_agent/framework/test_entry_adapter_on_cpu.py @@ -0,0 +1,313 @@ +"""CPU coverage for the Laminar worker and classic TransferQueue path.""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest +from omegaconf import OmegaConf + +from tests.uni_agent.support import FakeTokenizer +from uni_agent.framework import entry as entry_module +from uni_agent.framework.base import AgentFramework +from uni_agent.sandbox.registry import get_sandbox_cls +from verl.utils import tensordict_utils as tu +from verl.workers.rollout.replica import TokenOutput + + +class _InProcessGatewayManager: + def __init__(self, actor): + self._actor = actor + self.created_session_kwargs = [] + + async def create_session(self, session_id: str, **kwargs): + self.created_session_kwargs.append(dict(kwargs)) + return await self._actor.create_session(session_id, **kwargs) + + async def finalize_session(self, session_id: str): + return await self._actor.finalize_session(session_id) + + async def abort_session(self, session_id: str) -> None: + await self._actor.abort_session(session_id) + + +class _VersionedBackend: + def __init__(self): + self.calls = [] + + async def generate( + self, + request_id, + *, + prompt_ids, + sampling_params, + image_data=None, + video_data=None, + mm_processor_kwargs=None, + weight_version=None, + ) -> TokenOutput: + self.calls.append({"request_id": request_id, "weight_version": weight_version}) + text = r"\boxed{Paris}" + return TokenOutput( + token_ids=[ord(char) for char in text], + log_probs=[-0.1] * len(text), + stop_reason="completed", + extra_fields={ + "min_global_steps": weight_version, + "max_global_steps": weight_version, + }, + ) + + +def _versioned_prompts(task_config): + return tu.get_tensordict( + tensor_dict={ + "raw_prompt": [[{"role": "user", "content": "sample 0"}]], + "uid": ["uid-0"], + "data_source": ["local-echo"], + "reward_model": [{"ground_truth": "ok"}], + "extra_info": [{"index": 0}], + "tools_kwargs": [{"task": task_config}], + }, + non_tensor_dict={ + "global_steps": 7, + }, + ) + + +def _build_framework(gateway_manager): + from uni_agent.framework.framework import GatewayAgentFramework + + config = OmegaConf.create( + { + "actor_rollout_ref": { + "rollout": { + "n": 1, + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "calculate_log_probs": True, + "val_kwargs": {"n": 1, "temperature": 0, "top_p": 1.0, "top_k": -1}, + "custom": { + "agent_framework": { + "agent_runners": { + "mem_agent": { + "runner_fqn": "uni_agent.framework.task_runner.run_task", + "runner_kwargs": { + "model_name": "test-model", + "task_config_path": str( + Path(__file__).resolve().parents[3] / "examples/mem_agent/task_config.yaml" + ), + }, + "dispatch_mode": "inline_async", + } + } + } + }, + } + } + } + ) + return GatewayAgentFramework.from_config(config=config, gateway_manager=gateway_manager) + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_custom_framework_requires_explicit_admission_support(): + release = asyncio.Event() + started = asyncio.Event() + + class CustomFramework(AgentFramework): + @classmethod + def from_config(cls, **_kwargs): + return cls() + + async def generate_sequences(self, _prompts): + started.set() + await release.wait() + + worker_class = entry_module.AgentFrameworkWorker.__ray_metadata__.modified_class + worker = worker_class.__new__(worker_class) + worker.framework = CustomFramework() + with pytest.raises(AttributeError, match="submit_sessions"): + await worker.submit_sessions(object()) + assert not started.is_set() + + task = asyncio.create_task(worker.generate_sequences(object())) + await asyncio.wait_for(started.wait(), 2) + assert not task.done() + release.set() + await asyncio.wait_for(task, 2) + assert task.result() is None + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.parametrize( + ("runner_limits", "expected_options"), + [ + ([1536], {"max_concurrency": 1536}), + ([600, 700], {"max_concurrency": 1300}), + ([4, 6], {"max_concurrency": 1000}), + ([0], {"max_concurrency": 1000}), + ([None], {"max_concurrency": 1000}), + ([1536, 0], {"max_concurrency": 1536}), + ([], {"max_concurrency": 1000}), + ], +) +def test_adapter_derives_worker_concurrency_from_session_limits(monkeypatch, runner_limits, expected_options): + options = {} + worker = object() + gateway_manager = object() + + class _FrameworkWorkerClass: + @classmethod + def options(cls, **kwargs): + options.update(kwargs) + return cls + + @staticmethod + def remote(**kwargs): + assert kwargs["gateway_manager"] is gateway_manager + return worker + + runners = { + f"runner_{index}": {} if limit is None else {"max_concurrent_sessions": limit} + for index, limit in enumerate(runner_limits) + } + config = OmegaConf.create( + {"actor_rollout_ref": {"rollout": {"n": 8, "custom": {"agent_framework": {"agent_runners": runners}}}}} + ) + monkeypatch.setattr(entry_module, "build_gateway_manager", lambda **_: gateway_manager) + monkeypatch.setattr(entry_module, "AgentFrameworkWorker", _FrameworkWorkerClass) + + adapter = entry_module.AgentFrameworkRolloutAdapter.create(config=config, llm_client=object()) + + assert options == expected_options + assert adapter.framework_worker is worker + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.parametrize( + "method, expected_refs", + [ + ("generate_sequences", []), + ("generate_sequences_and_wait", ["completion"] * 3), + ("submit_sessions", ["admission"] * 3), + ], +) +def test_adapter_forwards_selected_versions_and_unversioned_batches(monkeypatch, method, expected_refs): + submitted_versions = [] + joined_refs = [] + + class _RemoteGeneration: + def __init__(self, name): + self.name = name + + def remote(self, prompts): + submitted_versions.append(tu.get(prompts, "global_steps")) + return self.name + + class _FrameworkWorker: + generate_sequences = _RemoteGeneration("completion") + + if method == "submit_sessions": + _FrameworkWorker.submit_sessions = _RemoteGeneration("admission") + + monkeypatch.setattr(entry_module.ray, "get", joined_refs.append) + adapter = entry_module.AgentFrameworkRolloutAdapter() + adapter.framework_worker = _FrameworkWorker() + generate = getattr(adapter, method) + + training_prompts = _versioned_prompts({}) + generate(training_prompts) + + validation_prompts = _versioned_prompts({}) + tu.assign_non_tensor_data(validation_prompts, "validate", True) + generate(validation_prompts) + generate(tu.get_tensordict(tensor_dict={"raw_prompt": [["unversioned"]]})) + + assert tu.get(training_prompts, "global_steps") == 7 + assert tu.get(validation_prompts, "global_steps") == 7 + assert submitted_versions == [7, 7, None] + assert joined_refs == expected_refs + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_local_task_writes_versioned_classic_tq_contract(monkeypatch): + from uni_agent.framework import framework as framework_module + from uni_agent.gateway.config import GatewayActorConfig + from uni_agent.gateway.gateway import _GatewayActor + + sandbox_lifecycle = [] + batch_writes = [] + group_writes = [] + + local_sandbox_cls = get_sandbox_cls("local") + original_start = local_sandbox_cls.start + original_stop = local_sandbox_cls.stop + + async def tracked_start(self): + sandbox_lifecycle.append("start") + await original_start(self) + + async def tracked_stop(self): + sandbox_lifecycle.append("stop") + await original_stop(self) + + monkeypatch.setattr(local_sandbox_cls, "start", tracked_start) + monkeypatch.setattr(local_sandbox_cls, "stop", tracked_stop) + + class _FakeTransferQueue: + async def async_kv_put(self, *, key, partition_id, tag): + group_writes.append({"key": key, "partition_id": partition_id, "tag": dict(tag)}) + + async def async_kv_batch_put(self, *, keys, fields, tags, partition_id): + batch_writes.append( + { + "keys": list(keys), + "partition_id": partition_id, + "fields": fields, + "tags": [dict(item) for item in tags], + } + ) + + monkeypatch.setattr(framework_module, "tq", _FakeTransferQueue()) + + backend = _VersionedBackend() + gateway_actor = _GatewayActor(GatewayActorConfig(tokenizer=FakeTokenizer()), backend) + await gateway_actor.start() + runtime = _InProcessGatewayManager(gateway_actor) + task_config = { + "name": "hotpotqa", + "prompt": [{"role": "user", "content": "What is the capital of France?"}], + "ground_truth": ["Paris"], + "metadata": {"instance_id": "cpu-local", "chunks": []}, + } + + try: + await _build_framework(runtime).generate_sequences(_versioned_prompts(task_config)) + finally: + await gateway_actor.shutdown() + + assert sandbox_lifecycle == ["start", "stop"] + assert runtime.created_session_kwargs[0]["weight_version"] == 7 + assert backend.calls[0]["weight_version"] == 7 + assert batch_writes[0]["keys"] == ["uid-0_0_0"] + assert batch_writes[0]["tags"][0]["status"] == "success" + assert tu.get(batch_writes[0]["fields"], "rm_scores")[0][-1].item() == 1.0 + assert group_writes == [ + {"key": "uid-0", "partition_id": "train", "tag": {"status": "running"}}, + { + "key": "uid-0", + "partition_id": "train", + "tag": { + "status": "finished", + }, + }, + ] diff --git a/tests/uni_agent/framework/test_generate_sequences_on_cpu.py b/tests/uni_agent/framework/test_generate_sequences_on_cpu.py index 23ec355f3..09122e25d 100644 --- a/tests/uni_agent/framework/test_generate_sequences_on_cpu.py +++ b/tests/uni_agent/framework/test_generate_sequences_on_cpu.py @@ -740,6 +740,11 @@ async def result_runner(**kwargs): ] assert trajectories[0].reward_score == 0.42 assert trajectories[0].reward_metrics == {"acc": 0.25, "format": 0.8} + assert trajectories[0].extra_fields["runner_reward_info"] == { + "reward": 0.5, + "metrics": {"acc": 1.0}, + "reward_context": {"case_id": "case-1", "nested": {"values": [1]}}, + } assert trajectories[0].finished is True @@ -1104,6 +1109,7 @@ async def test_framework_binds_sampling_defaults_to_gateway_sessions( ) assert [kwargs["sampling_params"] for kwargs in runtime.created_session_kwargs] == [expected_sampling_params] + assert [kwargs.get("weight_version") for kwargs in runtime.created_session_kwargs] == [7] @pytest.mark.cpu @@ -1371,7 +1377,83 @@ async def scored_runner(**kwargs): fields = fake_tq.batch_puts[0]["fields"] assert "reward_extra_info" not in fields.keys() extra_fields = tu.get(fields, "extra_fields") - assert extra_fields == [{"reward_extra_info": {"acc": 1.0}}] + assert extra_fields == [ + { + "reward_extra_info": {"acc": 1.0}, + "runner_reward_info": {"reward": 0.5, "metrics": {"acc": 1.0}, "reward_context": {}}, + } + ] + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +@pytest.mark.parametrize("validate", [False, True]) +@pytest.mark.parametrize("mask_unfinished", [False, True]) +async def test_runner_reward_context_survives_tq_without_changing_training_fields(fake_tq, validate, mask_unfinished): + context = { + "eval_completed": False, + "eval_exit_code": 1, + "eval_execution_time": 2.5, + "eval_report": { + "found_eval_status": True, + "status_map": {"test_fix": "FAILED", "test_existing": "PASSED"}, + "resolved": False, + }, + "agent_error": "RuntimeError: agent upload failed", + } + + async def scored_runner(**kwargs): + return TaskResult(reward=0.75, accuracy=0.5, finished=False, extra_info=context) + + gateway_extra = {"runner_reward_info": {"reward_context": {"eval_completed": True}}, "gateway_field": 17} + trajectories = [ + _trajectory(response_ids=[30, 31, 32], response_mask=[1, 0, 1], extra_fields=gateway_extra), + _trajectory(response_ids=[40, 41], response_mask=[1, 1], extra_fields=gateway_extra), + ] + runtime = _FakeGatewayManager({"session-sample-0-rollout-0": trajectories}) + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": _inline_runner_config(scored_runner)}, + gateway_manager=runtime, + mask_unfinished_episode=mask_unfinished, + ) + + await framework.generate_sequences(_build_prompts(count=1, global_steps=7, validate=validate)) + + assert len(fake_tq.batch_puts) == 1 + batch = fake_tq.batch_puts[0] + assert batch["partition_id"] == ("val" if validate else "train") + assert batch["keys"] == ["uid-0_0_0", "uid-0_0_1"] + assert [tag["uid"] for tag in batch["tags"]] == ["uid-0", "uid-0"] + assert tu.get(batch["fields"], "session_id") == [0, 0] + for index, trajectory in enumerate(trajectories): + fields = batch["fields"] + assert fields["rm_scores"][index].tolist() == [0.0] * (len(trajectory.response_ids) - 1) + [0.75] + expected_mask = [0] * len(trajectory.response_ids) if mask_unfinished else trajectory.response_mask + assert fields["response_mask"][index].tolist() == expected_mask + assert fields["loss_mask"][index].tolist() == expected_mask + for extra in tu.get(batch["fields"], "extra_fields"): + assert extra["runner_reward_info"] == { + "reward": 0.75, + "metrics": {"acc": 0.5}, + "reward_context": context, + } + assert extra["reward_extra_info"] == {"acc": 0.5} + assert extra["gateway_field"] == 17 + assert gateway_extra["runner_reward_info"] == {"reward_context": {"eval_completed": True}} + assert fake_tq.puts == [ + { + "key": "uid-0", + "partition_id": "val" if validate else "train", + "tag": {"status": "running"}, + }, + { + "key": "uid-0", + "partition_id": "val" if validate else "train", + "tag": {"status": "finished"}, + }, + ] + assert "runner_reward_info" not in framework._trajectory_meta(trajectories[0]) @pytest.mark.cpu @@ -1669,6 +1751,76 @@ async def agent_runner(*, raw_prompt, session, sample_index, tools_kwargs, **kwa assert runtime.aborted_sessions[0].startswith("session-sample-0-rollout-1-") +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_generate_sequences_keeps_successful_sessions_when_tq_write_fails(monkeypatch, fake_tq): + runtime = _FakeGatewayManager( + { + "session-sample-0-rollout-0": [_trajectory()], + "session-sample-0-rollout-1": [_trajectory()], + } + ) + original_batch_put = fake_tq.async_kv_batch_put + + async def fail_second_session(*, keys, fields, tags, partition_id): + if keys == ["uid-0_1_0"]: + raise OSError("trajectory sink unavailable") + await original_batch_put(keys=keys, fields=fields, tags=tags, partition_id=partition_id) + + monkeypatch.setattr(fake_tq, "async_kv_batch_put", fail_second_session) + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": _inline_runner_config(_async_noop_runner)}, + gateway_manager=runtime, + n=2, + val_n=2, + ) + + await framework.generate_sequences(_build_prompts(count=1, global_steps=8)) + + assert [put["keys"] for put in fake_tq.batch_puts] == [["uid-0_0_0"]] + assert fake_tq.puts == [ + {"key": "uid-0", "partition_id": "train", "tag": {"status": "running"}}, + { + "key": "uid-0", + "partition_id": "train", + "tag": {"status": "finished"}, + }, + ] + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_generate_sequences_keeps_other_prompts_when_terminal_write_fails(monkeypatch, fake_tq): + runtime = _FakeGatewayManager( + { + "session-sample-0-rollout-0": [_trajectory()], + "session-sample-1-rollout-0": [_trajectory()], + } + ) + original_put = fake_tq.async_kv_put + + async def fail_first_prompt(*, key, partition_id, tag): + if key == "uid-0" and tag["status"] == "finished": + raise OSError("terminal sink unavailable") + await original_put(key=key, partition_id=partition_id, tag=tag) + + monkeypatch.setattr(fake_tq, "async_kv_put", fail_first_prompt) + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": _inline_runner_config(_async_noop_runner)}, + gateway_manager=runtime, + ) + + await framework.generate_sequences(_build_prompts(count=2, global_steps=8)) + + assert fake_tq.puts == [ + {"key": "uid-0", "partition_id": "train", "tag": {"status": "running"}}, + {"key": "uid-1", "partition_id": "train", "tag": {"status": "running"}}, + {"key": "uid-1", "partition_id": "train", "tag": {"status": "finished"}}, + ] + + @pytest.mark.cpu @pytest.mark.level0 @pytest.mark.asyncio @@ -1926,3 +2078,375 @@ async def abort_session(self, session_id: str) -> None: # so no force-kill escalation happened. assert [call["force"] for call in cancel_calls] == [False] assert runtime.aborted_sessions, "terminated session must be aborted" + + +def _admission_worker(framework): + from uni_agent.framework.entry import AgentFrameworkWorker + + worker_class = AgentFrameworkWorker.__ray_metadata__.modified_class + worker = worker_class.__new__(worker_class) + worker.framework = framework + return worker + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +@pytest.mark.parametrize("dispatch_mode", ["inline_async", "ray_task"]) +async def test_admission_starts_sessions_incrementally_when_batch_exceeds_cap(monkeypatch, fake_tq, dispatch_mode): + from uni_agent.framework import framework as framework_module + + started = asyncio.Queue() + + async def runner(**_kwargs): + release = asyncio.Event() + started.put_nowait(release) + await release.wait() + return TaskResult(reward=1.0) + + if dispatch_mode == "ray_task": + monkeypatch.setattr( + framework_module._run_agent_runner_ray_task, + "remote", + lambda **kwargs: asyncio.create_task(runner(**kwargs)), + ) + runtime = _FakeGatewayManager({"session-sample-0": [_trajectory()]}) + framework = await _build_framework_with_agent_runners( + agent_runners={ + "runner": {**_inline_runner_config(runner, dispatch_mode=dispatch_mode), "max_concurrent_sessions": 1} + }, + gateway_manager=runtime, + n=3, + ) + worker = _admission_worker(framework) + submit = asyncio.create_task(worker.submit_sessions(_build_prompts(count=1))) + for _ in range(2): + release = await asyncio.wait_for(started.get(), 2) + assert not submit.done(), "the remaining sessions have not acquired their permits" + release.set() + + release = await asyncio.wait_for(started.get(), 2) + await asyncio.wait_for(submit, 2) + assert len(runtime.created_sessions) == 3 + assert len(runtime.finalized_sessions) == 2 + assert not fake_tq.batch_puts, "admission must not wait for the complete prompt group" + assert len(framework._rollout_tasks) == 1 + release.set() + await asyncio.wait_for(asyncio.gather(*framework._rollout_tasks), 2) + assert len(fake_tq.batch_puts) == 3 + assert framework._runner_semaphores["runner"]._value == 1 + assert not framework._rollout_tasks + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_failed_session_settles_admission_before_other_session_finishes(fake_tq): + release = asyncio.Event() + + async def runner(**_kwargs): + await release.wait() + return TaskResult(reward=1.0) + + class FailingGateway(_FakeGatewayManager): + async def create_session(self, session_id, **kwargs): + if "rollout-0-" in session_id: + raise RuntimeError("cannot create first session") + return await super().create_session(session_id, **kwargs) + + runtime = FailingGateway({"session-sample-0": [_trajectory()]}) + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": {**_inline_runner_config(runner), "max_concurrent_sessions": 1}}, + gateway_manager=runtime, + n=2, + ) + worker = _admission_worker(framework) + await asyncio.wait_for(worker.submit_sessions(_build_prompts(count=1)), 2) + assert not runtime.finalized_sessions + release.set() + await asyncio.wait_for(asyncio.gather(*framework._rollout_tasks), 2) + assert len(fake_tq.batch_puts) == 1 + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_admission_propagates_errors_before_sessions_are_expanded(fake_tq): + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": _inline_runner_config(_async_noop_runner)}, + gateway_manager=_FakeGatewayManager({}), + ) + worker = _admission_worker(framework) + with pytest.raises(ValueError, match="global_steps"): + await asyncio.wait_for(worker.submit_sessions(_build_prompts(count=1, global_steps=None)), 2) + assert not framework._rollout_tasks + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_background_failure_after_admission_is_consumed_and_logged(fake_tq, caplog): + release = asyncio.Event() + + async def runner(**_kwargs): + await release.wait() + raise RuntimeError("runner failed") + + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": _inline_runner_config(runner)}, + gateway_manager=_FakeGatewayManager({}), + ) + worker = _admission_worker(framework) + await asyncio.wait_for(worker.submit_sessions(_build_prompts(count=1)), 2) + pending = list(framework._rollout_tasks) + with caplog.at_level(logging.ERROR, logger="uni_agent.framework.framework"): + release.set() + await asyncio.wait_for(asyncio.gather(*pending, return_exceptions=True), 2) + failures = [record for record in caplog.records if "batch failed after admission" in record.getMessage()] + assert len(failures) == 1 + assert failures[0].exc_info[0] is RuntimeError + assert "All rollouts failed" in str(failures[0].exc_info[1]) + assert failures[0].exc_info[2] is not None + assert not framework._rollout_tasks + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_session_submission_returns_before_gateway_creation(fake_tq): + creating = asyncio.Event() + allow_create = asyncio.Event() + allow_finish = asyncio.Event() + + class BlockingGateway(_FakeGatewayManager): + async def create_session(self, session_id, **kwargs): + creating.set() + await allow_create.wait() + return await super().create_session(session_id, **kwargs) + + async def runner(**_kwargs): + await allow_finish.wait() + return TaskResult(reward=1.0) + + runtime = BlockingGateway({"session-sample-0": [_trajectory()]}) + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": {**_inline_runner_config(runner), "max_concurrent_sessions": 1}}, + gateway_manager=runtime, + ) + submit = asyncio.create_task(framework.submit_sessions(_build_prompts(count=1))) + await asyncio.wait_for(creating.wait(), 2) + await asyncio.wait_for(submit, 2) + assert not runtime.created_sessions + allow_create.set() + assert not runtime.finalized_sessions + allow_finish.set() + await asyncio.wait_for(asyncio.gather(*framework._rollout_tasks), 2) + assert len(runtime.created_sessions) == 1 + assert len(fake_tq.batch_puts) == 1 + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_session_capacity_is_held_through_scoring(monkeypatch, fake_tq): + scoring = asyncio.Queue() + + async def score(trajectories, *_args): + release = asyncio.Event() + scoring.put_nowait(release) + await release.wait() + return [(1.0, {}) for _ in trajectories] + + runtime = _FakeGatewayManager({"session-sample-0": [_trajectory()]}) + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": {**_inline_runner_config(_async_noop_runner), "max_concurrent_sessions": 1}}, + gateway_manager=runtime, + reward_loop_worker_handles=[object()], + ) + monkeypatch.setattr(framework, "_score_trajectories", score) + await asyncio.wait_for(framework.submit_sessions(_build_prompts(count=1)), 2) + first_release = await asyncio.wait_for(scoring.get(), 2) + second = asyncio.create_task(framework.submit_sessions(_build_prompts(count=1))) + done, _ = await asyncio.wait((second,), timeout=0.01) + assert not done + assert len(runtime.created_sessions) == 1 + assert len(runtime.finalized_sessions) == 1 + first_release.set() + second_release = await asyncio.wait_for(scoring.get(), 2) + await asyncio.wait_for(second, 2) + second_release.set() + await asyncio.wait_for(asyncio.gather(*framework._rollout_tasks), 2) + assert len(fake_tq.batch_puts) == 2 + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_stage", ["sampling", "runner"]) +async def test_invalid_prompt_does_not_block_other_prompt_admission(monkeypatch, fake_tq, failure_stage): + release = asyncio.Event() + started = asyncio.Queue() + + async def runner(**_kwargs): + started.put_nowait(True) + await release.wait() + return TaskResult(reward=1.0) + + runtime = _FakeGatewayManager({"session-sample-1": [_trajectory()]}) + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": _inline_runner_config(runner)}, + gateway_manager=runtime, + n=2, + ) + build_params = framework._build_session_sampling_params + + def sampling_params(*, partition_id, sample_fields): + if sample_fields["uid"] == "uid-0": + raise ValueError("invalid prompt sampling parameters") + return build_params(partition_id=partition_id, sample_fields=sample_fields) + + if failure_stage == "sampling": + monkeypatch.setattr(framework, "_build_session_sampling_params", sampling_params) + else: + submit_episode = framework._submit_agent_episode + + async def invalid_runner(**kwargs): + if kwargs["sample_fields"]["uid"] == "uid-0": + raise ValueError("Unknown agent runner") + return await submit_episode(**kwargs) + + monkeypatch.setattr(framework, "_submit_agent_episode", invalid_runner) + await asyncio.wait_for(framework.submit_sessions(_build_prompts(count=2)), 2) + for _ in range(2): + await asyncio.wait_for(started.get(), 2) + assert len(runtime.created_sessions) == 2 + assert all(session.startswith("session-sample-1-") for session in runtime.created_sessions) + assert not fake_tq.batch_puts + if failure_stage == "runner": + assert any(put["key"] == "uid-0" and put["tag"] == {"status": "failure"} for put in fake_tq.puts) + release.set() + await asyncio.wait_for(asyncio.gather(*framework._rollout_tasks), 2) + assert len(fake_tq.batch_puts) == 2 + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_all_prompt_failures_settle_submission_but_full_generation_raises(monkeypatch, fake_tq, caplog): + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": _inline_runner_config(_async_noop_runner)}, + gateway_manager=_FakeGatewayManager({}), + ) + + def invalid_sampling_params(**_kwargs): + raise ValueError("invalid prompt sampling parameters") + + monkeypatch.setattr(framework, "_build_session_sampling_params", invalid_sampling_params) + with caplog.at_level(logging.ERROR, logger="uni_agent.framework.framework"): + await asyncio.wait_for(framework.submit_sessions(_build_prompts(count=1)), 2) + await asyncio.wait_for(asyncio.gather(*framework._rollout_tasks), 2) + failures = [record for record in caplog.records if "batch failed after admission" in record.getMessage()] + assert len(failures) == 1 + assert failures[0].exc_info[0] is RuntimeError + with pytest.raises(RuntimeError, match="All rollouts failed"): + await framework.generate_sequences(_build_prompts(count=1)) + assert not framework._rollout_tasks + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_full_capacity_blocks_submission_before_creating_more_episode_tasks(monkeypatch, fake_tq): + started = asyncio.Event() + release = asyncio.Event() + + async def runner(**_kwargs): + started.set() + await release.wait() + return TaskResult(reward=1.0) + + runtime = _FakeGatewayManager({"session-sample-0": [_trajectory()]}) + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": {**_inline_runner_config(runner), "max_concurrent_sessions": 1}}, + gateway_manager=runtime, + n=3, + ) + scheduled = [] + run_episode = framework._run_agent_episode + + def record_episode(**kwargs): + scheduled.append(kwargs["session_index"]) + return run_episode(**kwargs) + + monkeypatch.setattr(framework, "_run_agent_episode", record_episode) + submit = asyncio.create_task(framework.submit_sessions(_build_prompts(count=1))) + await asyncio.wait_for(started.wait(), 2) + assert scheduled == [0] + semaphore = framework._runner_semaphores["runner"] + assert len(semaphore._waiters) == 1 + release.set() + await asyncio.wait_for(submit, 2) + await asyncio.wait_for(asyncio.gather(*framework._rollout_tasks), 2) + assert scheduled == [0, 1, 2] + assert len(fake_tq.batch_puts) == 3 + assert semaphore._value == 1 + assert not framework._rollout_tasks + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_session_cancelled_before_first_execution_releases_its_slot(monkeypatch, fake_tq): + runtime = _FakeGatewayManager({}) + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": {**_inline_runner_config(_async_noop_runner), "max_concurrent_sessions": 1}}, + gateway_manager=runtime, + ) + collect = framework._collect_prompt_rollouts + + def cancel_before_collecting(*, tasks, **kwargs): + assert framework._runner_semaphores["runner"]._value == 0 + assert not runtime.created_sessions + tasks[0].cancel() + return collect(tasks=tasks, **kwargs) + + monkeypatch.setattr(framework, "_collect_prompt_rollouts", cancel_before_collecting) + with pytest.raises(asyncio.CancelledError): + await framework.generate_sequences(_build_prompts(count=1)) + assert not runtime.created_sessions + assert framework._runner_semaphores["runner"]._value == 1 + assert not framework._rollout_tasks + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_cancelled_waiting_submission_keeps_other_batch_running(fake_tq): + started = asyncio.Event() + release = asyncio.Event() + + async def runner(**_kwargs): + started.set() + await release.wait() + return TaskResult(reward=1.0) + + runtime = _FakeGatewayManager({"session-sample-0": [_trajectory()]}) + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": {**_inline_runner_config(runner), "max_concurrent_sessions": 1}}, + gateway_manager=runtime, + ) + await framework.submit_sessions(_build_prompts(count=1)) + await asyncio.wait_for(started.wait(), 2) + waiting = asyncio.create_task(framework.submit_sessions(_build_prompts(count=1))) + done, _ = await asyncio.wait((waiting,), timeout=0.01) + assert not done + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + assert not runtime.aborted_sessions + assert len(runtime.created_sessions) == 1 + assert framework._runner_semaphores["runner"]._value == 0 + release.set() + await asyncio.wait_for(asyncio.gather(*framework._rollout_tasks), 2) + assert len(fake_tq.batch_puts) == 1 + assert framework._runner_semaphores["runner"]._value == 1 diff --git a/tests/uni_agent/framework/test_worker_concurrency_on_cpu.py b/tests/uni_agent/framework/test_worker_concurrency_on_cpu.py new file mode 100644 index 000000000..351a74a0f --- /dev/null +++ b/tests/uni_agent/framework/test_worker_concurrency_on_cpu.py @@ -0,0 +1,231 @@ +"""Real Ray scheduling and session limits without model or TransferQueue setup.""" + +from __future__ import annotations + +import asyncio +import inspect +import sys +import time +from types import SimpleNamespace + +import pytest +import ray +from omegaconf import OmegaConf + +import verl +from uni_agent.framework import entry +from uni_agent.framework import framework as framework_module +from uni_agent.framework.framework import GatewayAgentFramework, _RunnerConfig +from verl.utils import tensordict_utils as tu + + +@pytest.mark.cpu +@pytest.mark.level0 +def test_worker_starts_over_1000_batches_with_runner_limits(monkeypatch): + def source_paths(): + return (sys.executable, entry.__file__, inspect.getfile(GatewayAgentFramework), verl.__file__) + + class ControlledFramework(GatewayAgentFramework): + def __init__(self, runners): + super().__init__( + None, + runner_registry={name: _RunnerConfig.from_config(name, cfg) for name, cfg in runners.items()}, + rollout_config=OmegaConf.create( + {"n": 1, "temperature": 1.0, "top_p": 1.0, "top_k": -1, "calculate_log_probs": False} + ), + ) + self.release = asyncio.Event() + self.state = {"batches_started": 0, "episodes_completed": 0} + for name in runners: + self.state.update({f"active/{name}": 0, f"peak/{name}": 0}) + + async def generate_sequences(self, prompts): + self.state["batches_started"] += 1 + await super().generate_sequences( + tu.get_tensordict( + tensor_dict={"uid": [str(index) for index in range(len(prompts))], "agent_name": prompts}, + non_tensor_dict={"global_steps": 0}, + ) + ) + + async def _collect_prompt_rollouts(self, *, tasks, **_kwargs): + # Replace output persistence only; use the production session scheduler. + await asyncio.gather(*tasks) + return { + "num_success_sessions": len(tasks), + "num_failed_sessions": 0, + "num_success_outputs": len(tasks), + "num_unfinished_episodes": 0, + "num_failed_uids": 0, + "failure_reasons": [], + } + + async def _run_agent_episode(self, *, runner_name, sample_fields, **kwargs): + active, peak = f"active/{runner_name}", f"peak/{runner_name}" + self.state[active] += 1 + self.state[peak] = max(self.state[peak], self.state[active]) + try: + await self.release.wait() + self.state["episodes_completed"] += 1 + return [], sample_fields + finally: + self.state[active] -= 1 + + class ControlledWorker(entry.AgentFrameworkWorker.__ray_metadata__.modified_class): + def __init__(self, *, config, gateway_manager, reward_loop_worker_handles=None): + async def write_status(**_kwargs): + pass + + assert source_paths() == gateway_manager, "Ray worker must use the driver Python and product sources" + framework_module.tq = SimpleNamespace(async_kv_put=write_status) + self.framework = ControlledFramework(config.actor_rollout_ref.rollout.custom.agent_framework.agent_runners) + + @ray.method(concurrency_group="control") + def release(self): + self.framework._semaphore_loop.call_soon_threadsafe(self.framework.release.set) + + @ray.method(concurrency_group="control") + def read_state(self): + return dict(self.framework.state) + + runners = { + name: {"runner_fqn": "unused.runner", "dispatch_mode": "ray_task", "max_concurrent_sessions": 501} + for name in ("a", "b") + } + config = OmegaConf.create( + {"actor_rollout_ref": {"rollout": {"custom": {"agent_framework": {"agent_runners": runners}}}}} + ) + monkeypatch.setattr(entry, "build_gateway_manager", lambda **_: source_paths()) + # Test-only controls observe and release batches while all generation slots are occupied. + monkeypatch.setattr(entry, "AgentFrameworkWorker", ray.remote(concurrency_groups={"control": 1})(ControlledWorker)) + + owns_ray = not ray.is_initialized() + if owns_ray: + ray.init(address="local", num_cpus=1, include_dashboard=False) + worker = None + try: + worker = entry.AgentFrameworkRolloutAdapter.create(config=config, llm_client=None).framework_worker + batches = [worker.generate_sequences.remote([name, name]) for _ in range(501) for name in runners] + deadline = time.monotonic() + 60 + while True: + state = ray.get(worker.read_state.remote(), timeout=10) + if state["batches_started"] == 1002 and all(state[f"active/{name}"] == 501 for name in runners): + break + assert time.monotonic() < deadline, state + time.sleep(0.02) + assert not ray.wait(batches, num_returns=1, timeout=0)[0] + assert state["episodes_completed"] == 0 + assert all(state[f"peak/{name}"] == 501 for name in runners) + + ray.get(worker.release.remote(), timeout=10) + assert ray.get(batches, timeout=30) == [None] * 1002 + state = ray.get(worker.read_state.remote(), timeout=10) + assert state["episodes_completed"] == 2004 + assert all(state[f"peak/{name}"] == 501 and state[f"active/{name}"] == 0 for name in runners) + finally: + if worker is not None: + ray.kill(worker) + if owns_ray: + ray.shutdown() + + +@pytest.mark.cpu +@pytest.mark.level0 +def test_ray_admission_rpc_keeps_rollouts_alive_and_backpressures_next_batch(): + class Gateway: + async def create_session(self, session_id, **kwargs): + return None + + async def finalize_session(self, session_id): + return [] + + class ControlledFramework(GatewayAgentFramework): + def __init__(self): + runner_config = _RunnerConfig("unused.runner", {}, "ray_task", 1) + super().__init__( + Gateway(), + runner_registry={"runner": runner_config}, + rollout_config=OmegaConf.create( + {"n": 1, "temperature": 1.0, "top_p": 1.0, "top_k": -1, "calculate_log_probs": False} + ), + ) + runner_config.dispatch_mode = "inline_async" + self._inline_runners["runner"] = self.run + self.releases = {} + self.started = [] + self.completed = [] + + async def run(self, *, raw_prompt, **_kwargs): + release = self.releases[raw_prompt] = asyncio.Event() + self.started.append(raw_prompt) + await release.wait() + self.completed.append(raw_prompt) + + async def _collect_prompt_rollouts(self, *, tasks, **_kwargs): + # Replace output persistence only; admission and lifetime use the real scheduler. + await asyncio.gather(*tasks) + return { + "num_success_sessions": len(tasks), + "num_failed_sessions": 0, + "num_success_outputs": len(tasks), + "num_unfinished_episodes": 0, + "num_failed_uids": 0, + "failure_reasons": [], + } + + class ControlledWorker(entry.AgentFrameworkWorker.__ray_metadata__.modified_class): + def __init__(self, expected_sources): + async def write_status(**_kwargs): + pass + + assert (sys.executable, entry.__file__, verl.__file__) == expected_sources + framework_module.tq = SimpleNamespace(async_kv_put=write_status) + self.framework = ControlledFramework() + + async def release(self, name): + self.framework.releases[name].set() + + async def read_state(self): + return self.framework.started, self.framework.completed, len(self.framework._rollout_tasks) + + def prompts(name): + return tu.get_tensordict( + tensor_dict={"uid": [name], "raw_prompt": [name]}, + non_tensor_dict={"global_steps": 0}, + ) + + owns_ray = not ray.is_initialized() + if owns_ray: + ray.init(address="local", num_cpus=1, include_dashboard=False) + worker = None + try: + worker = ray.remote(ControlledWorker).remote((sys.executable, entry.__file__, verl.__file__)) + assert ray.get(worker.submit_sessions.remote(prompts("first")), timeout=30) is None + assert ray.get(worker.read_state.remote(), timeout=10) == (["first"], [], 1) + + second = worker.submit_sessions.remote(prompts("second")) + assert not ray.wait([second], timeout=0.1)[0] + assert ray.get(worker.read_state.remote(), timeout=10) == (["first"], [], 1) + ray.get(worker.release.remote("first"), timeout=10) + assert ray.get(second, timeout=10) is None + assert ray.get(worker.read_state.remote(), timeout=10)[:2] == (["first", "second"], ["first"]) + + ray.get(worker.release.remote("second"), timeout=10) + complete = worker.generate_sequences.remote(prompts("standalone")) + assert not ray.wait([complete], timeout=0.1)[0] + assert ray.get(worker.read_state.remote(), timeout=10)[:2] == ( + ["first", "second", "standalone"], + ["first", "second"], + ) + ray.get(worker.release.remote("standalone"), timeout=10) + assert ray.get(complete, timeout=10) is None + assert ray.get(worker.read_state.remote(), timeout=10) == ( + ["first", "second", "standalone"], + ["first", "second", "standalone"], + 0, + ) + finally: + if worker is not None: + ray.kill(worker) + if owns_ray: + ray.shutdown() diff --git a/tests/uni_agent/gateway/test_gateway_actor_on_cpu.py b/tests/uni_agent/gateway/test_gateway_actor_on_cpu.py index df955dcdd..d2ba6be09 100644 --- a/tests/uni_agent/gateway/test_gateway_actor_on_cpu.py +++ b/tests/uni_agent/gateway/test_gateway_actor_on_cpu.py @@ -112,6 +112,84 @@ async def test_gateway_actor_forwards_last_assistant_rollback_to_session(): assert state["rollback_count"] == 1 +class _RouteBackend(InspectingBackend): + def __init__(self, *, release_error=None): + super().__init__() + self.release_error = release_error + self.bind_calls = [] + self.release_calls = [] + + async def bind_route(self, **kwargs): + self.bind_calls.append(kwargs) + + async def release_route(self, **kwargs): + self.release_calls.append(kwargs) + if self.release_error is not None: + raise self.release_error + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +@pytest.mark.parametrize("weight_version", [None, 3]) +@pytest.mark.parametrize("close_method", ["finalize_session", "abort_session"]) +async def test_gateway_actor_owns_route_lifecycle(weight_version, close_method): + from uni_agent.gateway.config import GatewayActorConfig + from uni_agent.gateway.gateway import _GatewayActor + + backend = _RouteBackend() + actor = _GatewayActor(GatewayActorConfig(tokenizer=FakeTokenizer()), backend) + await actor.start() + try: + await actor.create_session("route-session", weight_version=weight_version) + + assert backend.bind_calls == [ + { + "session_id": "route-session", + "weight_version": weight_version, + } + ] + + await getattr(actor, close_method)("route-session") + + assert backend.release_calls == [{"session_id": "route-session"}] + assert "route-session" not in actor._sessions + finally: + await actor.shutdown() + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_gateway_actor_abort_retries_failed_route_release(): + from uni_agent.gateway.config import GatewayActorConfig + from uni_agent.gateway.gateway import _GatewayActor + + release_error = RuntimeError("route-release-failed") + backend = _RouteBackend(release_error=release_error) + actor = _GatewayActor(GatewayActorConfig(tokenizer=FakeTokenizer()), backend) + await actor.start() + try: + await actor.create_session("route-session") + + with pytest.raises(RuntimeError, match="route-release-failed") as raised: + await actor.finalize_session("route-session") + + assert raised.value is release_error + assert "route-session" not in actor._sessions + assert backend.release_calls == [{"session_id": "route-session"}] + + backend.release_error = None + await actor.abort_session("route-session") + + assert backend.release_calls == [ + {"session_id": "route-session"}, + {"session_id": "route-session"}, + ] + finally: + await actor.shutdown() + + @pytest.mark.cpu @pytest.mark.level0 @pytest.mark.asyncio diff --git a/tests/uni_agent/gateway/test_gateway_manager_on_cpu.py b/tests/uni_agent/gateway/test_gateway_manager_on_cpu.py index 4ae058f21..7fc643d77 100644 --- a/tests/uni_agent/gateway/test_gateway_manager_on_cpu.py +++ b/tests/uni_agent/gateway/test_gateway_manager_on_cpu.py @@ -31,10 +31,12 @@ class _FakeGateway: def __init__(self): self.created = [] + self.created_kwargs = [] async def _create(self, session_id, **kwargs): await asyncio.sleep(0) # yield the loop mid-create to expose select/increment races self.created.append(session_id) + self.created_kwargs.append(kwargs) return session_id @property @@ -42,6 +44,37 @@ def create_session(self): return _FakeRemoteMethod(self._create) +class _LifecycleGateway: + """Controllable remote-method stub for close failures.""" + + def __init__(self, *, finalize_error=None): + self.finalize_error = finalize_error + self.abort_calls = [] + self.finalize_calls = [] + + async def _abort(self, session_id): + self.abort_calls.append(session_id) + + async def _finalize(self, session_id): + self.finalize_calls.append(session_id) + if self.finalize_error is not None: + raise self.finalize_error + return ["trajectory"] + + abort_session = property(lambda self: _FakeRemoteMethod(self._abort)) + finalize_session = property(lambda self: _FakeRemoteMethod(self._finalize)) + + +def _new_manager(gateway, *, session_ids=()): + from uni_agent.gateway.manager import GatewayManager + + manager = GatewayManager.__new__(GatewayManager) + manager.gateways = [gateway] + manager.active_sessions_per_gateway = [len(session_ids)] + manager._session_to_gateway_index = dict.fromkeys(session_ids, 0) + return manager + + @pytest.mark.cpu @pytest.mark.level0 @pytest.mark.asyncio @@ -70,6 +103,44 @@ async def test_gateway_manager_balances_concurrent_session_creation(): assert [len(g.created) for g in gateways] == counts +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_gateway_manager_forwards_session_config(): + gateway = _FakeGateway() + manager = _new_manager(gateway) + + await manager.create_session("session-create", weight_version=3) + assert gateway.created == ["session-create"] + assert gateway.created_kwargs == [{"weight_version": 3}] + + assert manager._session_to_gateway_index == {"session-create": 0} + assert manager.active_sessions_per_gateway == [1] + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_gateway_manager_close_failure_retains_mapping_and_owner(): + close_error = RuntimeError("actor-close-unknown") + gateway = _LifecycleGateway(finalize_error=close_error) + manager = _new_manager(gateway, session_ids=("session-close",)) + + with pytest.raises(RuntimeError, match="actor-close-unknown") as raised: + await manager.finalize_session("session-close") + + assert raised.value is close_error + assert gateway.finalize_calls == ["session-close"] + assert manager._session_to_gateway_index == {"session-close": 0} + assert manager.active_sessions_per_gateway == [1] + + await manager.abort_session("session-close") + + assert gateway.abort_calls == ["session-close"] + assert manager._session_to_gateway_index == {} + assert manager.active_sessions_per_gateway == [0] + + @pytest.mark.cpu @pytest.mark.level0 def test_gateway_manager_rejects_zero_gateway_count(): diff --git a/tests/uni_agent/gateway/test_session_multiple_chains_on_cpu.py b/tests/uni_agent/gateway/test_session_multiple_chains_on_cpu.py index d4ffff889..e850833c5 100644 --- a/tests/uni_agent/gateway/test_session_multiple_chains_on_cpu.py +++ b/tests/uni_agent/gateway/test_session_multiple_chains_on_cpu.py @@ -131,6 +131,7 @@ class _VersionedBackend: def __init__(self, steps): # steps: list of (text, min_global_steps, max_global_steps) self.steps = list(steps) + self.calls = [] async def generate( self, @@ -141,7 +142,9 @@ async def generate( image_data=None, video_data=None, mm_processor_kwargs=None, + weight_version=None, ): + self.calls.append({"request_id": request_id, "weight_version": weight_version}) text, min_steps, max_steps = self.steps.pop(0) token_ids = _ids(text) return TokenOutput( @@ -1936,3 +1939,69 @@ async def test_weight_versions_absent_when_backend_omits_them(): [trajectory] = await session.finalize() assert trajectory.extra_fields == {} + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_session_rejects_backend_version_conflicting_with_weight_version(): + session = GatewaySession( + SessionHandle(session_id="session-version-conflict"), + MessageCodec(FakeTokenizer()), + weight_version=3, + ) + + with pytest.raises(RuntimeError, match="backend generation version conflicts with session weight_version"): + await _run( + session, + _VersionedBackend([("OK", 4, 4)]), + [{"role": "user", "content": "hello"}], + ) + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_versioned_sibling_chains_share_session_route_and_version(): + session = GatewaySession( + SessionHandle(session_id="session-versioned-siblings"), + MessageCodec(FakeTokenizer()), + weight_version=3, + ) + backend = _VersionedBackend( + [ + ("FIRST", 3, 3), + ("SIBLING", 3, 3), + ("FIRST-NEXT", 3, 3), + ("SIBLING-NEXT", 3, 3), + ] + ) + prompt = [{"role": "user", "content": "same prompt"}] + + await _run(session, backend, prompt) + await _run(session, backend, prompt) + await _run( + session, + backend, + [*prompt, {"role": "assistant", "content": "FIRST"}, {"role": "user", "content": "continue first"}], + ) + await _run( + session, + backend, + [ + *prompt, + {"role": "assistant", "content": "SIBLING"}, + {"role": "user", "content": "continue sibling"}, + ], + ) + + assert backend.calls == [{"request_id": "session-versioned-siblings", "weight_version": 3}] * 4 + trajectories = await session.finalize() + decoded = [_decode_response_ids(trajectory.response_ids) for trajectory in trajectories] + assert len(decoded) == 2 + assert any(text.startswith("FIRST") and text.endswith("FIRST-NEXT") for text in decoded) + assert any(text.startswith("SIBLING") and text.endswith("SIBLING-NEXT") for text in decoded) + assert { + (trajectory.extra_fields["min_global_steps"], trajectory.extra_fields["max_global_steps"]) + for trajectory in trajectories + } == {(3, 3)} diff --git a/tests/uni_agent/sandbox/plugin_probe_provider.py b/tests/uni_agent/sandbox/plugin_probe_provider.py new file mode 100644 index 000000000..ba088f128 --- /dev/null +++ b/tests/uni_agent/sandbox/plugin_probe_provider.py @@ -0,0 +1,18 @@ +"""Test-only overlay provider imported via ``UNI_AGENT_SANDBOX_PLUGINS``.""" + +from __future__ import annotations + +from uni_agent.sandbox.base import ExecResult, Sandbox +from uni_agent.sandbox.registry import register_sandbox + + +@register_sandbox("plugin_probe") +class PluginProbeSandbox(Sandbox): + async def start(self) -> None: + return None + + async def stop(self) -> None: + return None + + async def _exec(self, argv, *, timeout=None, workdir=None, env=None) -> ExecResult: + return ExecResult(exit_code=0, stdout="", stderr="") diff --git a/tests/uni_agent/sandbox/test_sandbox_provider_discovery_on_cpu.py b/tests/uni_agent/sandbox/test_sandbox_provider_discovery_on_cpu.py new file mode 100644 index 000000000..3b621f832 --- /dev/null +++ b/tests/uni_agent/sandbox/test_sandbox_provider_discovery_on_cpu.py @@ -0,0 +1,46 @@ +"""CPU checks for ``UNI_AGENT_SANDBOX_PLUGINS`` overlay discovery.""" + +from __future__ import annotations + +import pytest + +from uni_agent.sandbox.registry import ( + SANDBOX_PLUGIN_MODULES_ENV, + SANDBOX_REGISTRY, + get_sandbox_cls, +) + +_PROBE_MODULE = "tests.uni_agent.sandbox.plugin_probe_provider" +_PROBE_NAME = "plugin_probe" + + +@pytest.fixture(autouse=True) +def _isolate_probe_registry(): + SANDBOX_REGISTRY.pop(_PROBE_NAME, None) + yield + SANDBOX_REGISTRY.pop(_PROBE_NAME, None) + + +@pytest.mark.cpu +@pytest.mark.level0 +def test_env_plugin_registers_provider(monkeypatch): + monkeypatch.setenv(SANDBOX_PLUGIN_MODULES_ENV, _PROBE_MODULE) + cls = get_sandbox_cls(_PROBE_NAME) + assert cls.__name__ == "PluginProbeSandbox" + assert cls.__module__ == _PROBE_MODULE + + +@pytest.mark.cpu +@pytest.mark.level0 +def test_missing_plugin_module_fails_closed(monkeypatch): + monkeypatch.setenv(SANDBOX_PLUGIN_MODULES_ENV, "not_a_real_sandbox_plugin_module") + with pytest.raises(ImportError, match="not_a_real_sandbox_plugin_module"): + get_sandbox_cls(_PROBE_NAME) + + +@pytest.mark.cpu +@pytest.mark.level0 +def test_unknown_provider_stays_fail_closed(monkeypatch): + monkeypatch.delenv(SANDBOX_PLUGIN_MODULES_ENV, raising=False) + with pytest.raises(ValueError, match="Unknown sandbox provider"): + get_sandbox_cls("not-a-real-provider") diff --git a/tests/uni_agent/tasks/test_swe_agent_error_on_cpu.py b/tests/uni_agent/tasks/test_swe_agent_error_on_cpu.py new file mode 100644 index 000000000..87b4a56d5 --- /dev/null +++ b/tests/uni_agent/tasks/test_swe_agent_error_on_cpu.py @@ -0,0 +1,46 @@ +from contextlib import asynccontextmanager +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from uni_agent.agents.base import AgentResult +from uni_agent.tasks.swe_bench import reward +from uni_agent.tasks.swe_bench.task import SWEBenchTask + + +@pytest.mark.asyncio +@pytest.mark.parametrize("agent_error", [None, "", " ", "RuntimeError: upload destination must be absolute"]) +async def test_swe_task_preserves_explicit_agent_error_without_changing_reward(monkeypatch, agent_error): + sandbox = object() + + @asynccontextmanager + async def build_sandbox(): + yield sandbox + + agent = SimpleNamespace(run=AsyncMock(return_value=AgentResult(finished=False, info={"error": agent_error}))) + task = object.__new__(SWEBenchTask) + task.config = SimpleNamespace( + metadata={"instance_id": "test"}, + model_dump=lambda **kwargs: {}, + run_oracle_solution=False, + prompt=[{"role": "user", "content": "test"}], + eval_timeout=60, + ) + task.build_sandbox = build_sandbox + task.build_agent = lambda: agent + scorer_result = {"resolved": False, "eval_report": {"found_eval_status": True}} + scorer = AsyncMock(return_value=scorer_result) + monkeypatch.setattr(reward, "compute_reward", scorer) + + result = await task.run() + + scorer.assert_awaited_once_with(task.config.metadata, sandbox, eval_timeout=60) + assert result.reward == result.accuracy == 0.0 + assert result.finished is False + assert result.extra_info["eval_report"] == scorer_result["eval_report"] + assert "agent_error" not in scorer_result + if agent_error and agent_error.strip(): + assert result.extra_info["agent_error"] == agent_error + else: + assert "agent_error" not in result.extra_info diff --git a/tests/uni_agent/tasks/test_swe_reward_on_cpu.py b/tests/uni_agent/tasks/test_swe_reward_on_cpu.py new file mode 100644 index 000000000..d39f175b0 --- /dev/null +++ b/tests/uni_agent/tasks/test_swe_reward_on_cpu.py @@ -0,0 +1,72 @@ +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from uni_agent.sandbox.base import ExecResult +from uni_agent.tasks.swe_bench import reward + +UNRESOLVED_OUTPUT = "\n".join( + ( + reward.START_TEST_OUTPUT, + "FAILED test_requests.py::test_fix - AssertionError", + "PASSED test_requests.py::test_existing", + reward.END_TEST_OUTPUT, + ) +) +RESOLVED_OUTPUT = UNRESOLVED_OUTPUT.replace( + "FAILED test_requests.py::test_fix - AssertionError", "PASSED test_requests.py::test_fix" +) +EMPTY_TEST_OUTPUT = f"{reward.START_TEST_OUTPUT}\n{reward.END_TEST_OUTPUT}" +UNRESOLVED_STATUS_MAP = {"test_requests.py::test_fix": "FAILED", "test_requests.py::test_existing": "PASSED"} +RESOLVED_STATUS_MAP = {**UNRESOLVED_STATUS_MAP, "test_requests.py::test_fix": "PASSED"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "exit_code,output,found,status_map,resolved", + [ + pytest.param(0, UNRESOLVED_OUTPUT, True, UNRESOLVED_STATUS_MAP, False, id="normal-unresolved"), + pytest.param(1, UNRESOLVED_OUTPUT, True, UNRESOLVED_STATUS_MAP, False, id="nonzero-unresolved"), + pytest.param(1, RESOLVED_OUTPUT, True, RESOLVED_STATUS_MAP, True, id="nonzero-resolved"), + pytest.param(0, EMPTY_TEST_OUTPUT, True, {}, False, id="markers-without-tests"), + pytest.param(1, "collection failed", False, {}, False, id="missing-markers"), + pytest.param(-1, "", False, {}, False, id="sandbox-timeout-result"), + pytest.param(127, "", False, {}, False, id="sandbox-exec-failure-result"), + ], +) +async def test_swe_reward_preserves_execution_evidence_without_changing_resolution( + exit_code, output, found, status_map, resolved +): + metadata = { + "instance_id": "psf__requests-test", + "repo": "psf/requests", + "version": "2.2", + "base_commit": "0" * 40, + "test_patch": "", + "FAIL_TO_PASS": json.dumps(["test_requests.py::test_fix"]), + "PASS_TO_PASS": json.dumps(["test_requests.py::test_existing"]), + } + sandbox = SimpleNamespace( + write_file=AsyncMock(), + exec_shell=AsyncMock( + return_value=ExecResult( + exit_code=exit_code, + stdout=output, + stderr="exec timed out after 600.0s" if exit_code == -1 else "", + ) + ), + ) + + result = await reward.compute_reward(metadata, sandbox, eval_timeout=600.0) + + assert result["eval_exit_code"] == exit_code + assert result["eval_completed"] is (exit_code == 0) + assert result["eval_report"]["found_eval_status"] is found + assert result["eval_report"]["status_map"] == status_map + assert result["eval_report"]["resolved"] is resolved + assert result["resolved"] is resolved + if not found: + assert result["eval_report"]["test_status"] is None + assert sandbox.exec_shell.await_args.kwargs == {"workdir": "/testbed", "timeout": 600.0} diff --git a/uni_agent/framework/entry.py b/uni_agent/framework/entry.py index 029f7b51c..92ac63c96 100644 --- a/uni_agent/framework/entry.py +++ b/uni_agent/framework/entry.py @@ -14,6 +14,8 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import ray from omegaconf import OmegaConf @@ -26,6 +28,9 @@ from verl.utils.transferqueue_utils import tq from verl.workers.config import HFModelConfig, RolloutConfig +if TYPE_CHECKING: + from tensordict import TensorDict + _DEFAULT_FRAMEWORK_CLASS = "uni_agent.framework.framework.GatewayAgentFramework" @@ -43,6 +48,10 @@ def build_gateway_manager(*, config, llm_client) -> GatewayManager: # Match AgentLoopWorker pattern: self-load tokenizer/processor via HFModelConfig. rollout_config: RolloutConfig = omega_conf_to_dataclass(rollout_cfg) model_config: HFModelConfig = omega_conf_to_dataclass(model_cfg) + # TODO: Gateway session capacity is prompt_length + response_length, which can + # disagree with actor_rollout_ref.rollout.max_model_len (engine context window + # or extra slack). Forward that knob into GatewayActorConfig once the override + # path is restored so session clipping matches the engine budget. gateway_actor_config = GatewayActorConfig( tokenizer=model_config.tokenizer, processor=model_config.processor, @@ -104,13 +113,22 @@ def __init__(self, *, config, gateway_manager, reward_loop_worker_handles=None) async def generate_sequences(self, prompts) -> None: await self.framework.generate_sequences(prompts) + async def submit_sessions(self, prompts: TensorDict): + """Require the Framework's admission-only submission interface.""" + await self.framework.submit_sessions(prompts) + class AgentFrameworkRolloutAdapter: """Trainer-facing adapter satisfying the `agent_loop_manager_class` contract. Holds zero recipe-specific logic; every agent-framework recipe wires the same class in yaml. The adapter owns the gateway manager (driver-side) and - injects it into the framework worker. + injects it into the framework worker. The caller owns version availability; + the adapter forwards the selected version unchanged. Per-runner + ``max_concurrent_sessions`` is the worker-side concurrency cap. + Laminar training explicitly uses ``submit_sessions`` for admission backpressure. + Ordinary generation only requires ``generate_sequences`` on the framework; + there is no completion-based admission fallback. """ def __init__(self) -> None: @@ -135,8 +153,17 @@ def create( "disable teacher policy/distillation or use an AgentLoopManager that supports it." ) + agent_runners = ( + OmegaConf.select(config, "actor_rollout_ref.rollout.custom.agent_framework.agent_runners", default={}) or {} + ) + runner_limits = [int(runner.get("max_concurrent_sessions", 0) or 0) for runner in agent_runners.values()] + # Cover the session semaphore budgets without lowering Ray's async default. + # A batch can retain its RPC slot for just one straggling session, + # so dividing by batch size or rollout.n would throttle useful work. + worker_concurrency = max(1000, sum(limit for limit in runner_limits if limit > 0)) + gateway_manager = build_gateway_manager(config=config, llm_client=llm_client) - framework_worker = AgentFrameworkWorker.remote( + framework_worker = AgentFrameworkWorker.options(max_concurrency=worker_concurrency).remote( config=config, gateway_manager=gateway_manager, reward_loop_worker_handles=reward_loop_worker_handles, @@ -155,6 +182,14 @@ def generate_sequences(self, prompts) -> None: self.framework_worker.generate_sequences.remote(prompts) return None + def submit_sessions(self, prompts: TensorDict) -> None: + """Wait for session admission; require framework support for this capability.""" + if self.framework_worker is None: + raise RuntimeError("framework must be initialized before submit_sessions") + + ray.get(self.framework_worker.submit_sessions.remote(prompts)) + return None + def generate_sequences_and_wait(self, prompts) -> None: """Blocking variant of :meth:`generate_sequences` for standalone (non-trainer) runs. diff --git a/uni_agent/framework/framework.py b/uni_agent/framework/framework.py index f5cc06a2c..ad08bc07a 100644 --- a/uni_agent/framework/framework.py +++ b/uni_agent/framework/framework.py @@ -239,6 +239,25 @@ def _list_of_tq_fields_to_tensordict(fields: list[dict[str, object]]) -> TensorD return td +def _attach_runner_reward_info(trajectories: list[Trajectory], task_result: TaskResult) -> list[Trajectory]: + """Preserve runner scoring context separately from final reward metrics.""" + task_metrics = {} if task_result.accuracy is None else {"acc": task_result.accuracy} + return [ + replace( + traj, + extra_fields={ + **traj.extra_fields, + "runner_reward_info": { + "reward": task_result.reward, + "metrics": dict(task_metrics), + "reward_context": dict(task_result.extra_info), + }, + }, + ) + for traj in trajectories + ] + + def _trajectory_to_reward_dataproto(trajectory, sample_fields, task_result: TaskResult): """Build a single-sample DataProto for RewardLoopWorker.compute_score. @@ -336,6 +355,7 @@ def __init__( logger.info("No streaming reward worker handles; using the non-streaming reward path") self._processor = processor self._rollout_config = rollout_config + self._rollout_tasks: set[asyncio.Task[None]] = set() self._runner_semaphores: dict[str, asyncio.Semaphore] = {} self._semaphore_loop: asyncio.AbstractEventLoop | None = None self._log_dir = log_dir @@ -484,8 +504,15 @@ def _build_session_sampling_params( sampling_params.update(temperature=0, top_p=1.0, top_k=-1) return sampling_params + async def submit_sessions(self, prompts: TensorDict) -> None: + """Wait only for session admission; do not await rollout completion.""" + await self._generate_sequences(prompts, wait_for_completion=False) + async def generate_sequences(self, prompts: TensorDict) -> None: """Run rollout-manager generation and write outputs into TransferQueue.""" + await self._generate_sequences(prompts, wait_for_completion=True) + + async def _generate_sequences(self, prompts: TensorDict, *, wait_for_completion: bool) -> None: if self._rollout_config is None: raise RuntimeError("GatewayAgentFramework requires rollout_config for generate_sequences") @@ -505,12 +532,30 @@ async def generate_sequences(self, prompts: TensorDict) -> None: if uids is None: raise ValueError("GatewayAgentFramework requires prompts['uid'] for TransferQueue output") - stats = await self._run_batch_rollouts( + batch = await self._run_batch_rollouts( prompts, global_steps=global_steps, partition_id=partition_id, num_sessions=num_sessions, ) + + async def finish_batch() -> None: + try: + await self._finish_generation(batch, global_steps=global_steps) + except Exception: + if wait_for_completion: + raise + logger.exception("Agent framework batch failed after admission") + finally: + self._rollout_tasks.discard(asyncio.current_task()) + + task = asyncio.create_task(finish_batch()) + self._rollout_tasks.add(task) + if wait_for_completion: + await task + + async def _finish_generation(self, batch: asyncio.Task[dict], *, global_steps: int | None) -> None: + stats = await batch logger.info( "generate_sequences summary: num_input_prompts=%s num_success_sessions=%s " "num_failed_sessions=%s num_success_outputs=%s num_unfinished_episodes=%s " @@ -537,30 +582,50 @@ async def _run_batch_rollouts( global_steps: int | None, partition_id: str, num_sessions: int = 1, - ) -> dict: - """Run all prompts in a batch and aggregate prompt/session stats.""" + ) -> asyncio.Task[dict]: + """Submit all prompts in a batch and return their background collection.""" assert len(prompts) > 0, "generate_sequences requires a non-empty batch" if num_sessions <= 0: raise ValueError(f"num_sessions must be positive, got {num_sessions}") # Batch layer: each sample/prompt owns its own group of rollout.n sessions. # Prompt tasks are isolated so one prompt failure does not drop the whole batch. + # TODO: honor per-prompt ``__rollout_n__`` (verl agent_loop_tq optional hook) + # once uni_agent needs sample-level n. Until then, keep main's batch-level n. tasks = [] + outcomes: list[dict | BaseException] = [] for sample_index in range(len(prompts)): - tasks.append( - self._run_prompt_rollouts( - sample_fields=self._extract_sample_fields(prompts=prompts, sample_index=sample_index), - sample_index=sample_index, - global_steps=global_steps, - partition_id=partition_id, - num_sessions=num_sessions, + sample_fields = self._extract_sample_fields(prompts=prompts, sample_index=sample_index) + try: + tasks.append( + await self._run_prompt_rollouts( + sample_fields=sample_fields, + sample_index=sample_index, + global_steps=global_steps, + partition_id=partition_id, + num_sessions=num_sessions, + ) ) - ) - outcomes = await asyncio.gather(*tasks, return_exceptions=True) + except Exception as exc: + outcomes.append(exc) + task = asyncio.create_task( + self._collect_batch_rollouts(tasks, outcomes, num_prompts=len(prompts), num_sessions=num_sessions) + ) + return task + + async def _collect_batch_rollouts( + self, + tasks: list[asyncio.Task[dict]], + outcomes: list[dict | BaseException], + *, + num_prompts: int, + num_sessions: int, + ) -> dict: + outcomes.extend(await asyncio.gather(*tasks, return_exceptions=True)) failure_reasons: list[str] = [] stats = { - "num_input_prompts": len(prompts), + "num_input_prompts": num_prompts, "num_success_sessions": 0, "num_failed_sessions": 0, "num_success_outputs": 0, @@ -594,8 +659,8 @@ async def _run_prompt_rollouts( global_steps: int | None, partition_id: str, num_sessions: int, - ) -> dict: - """Run ``rollout.n`` independent sessions for one prompt and persist their outputs.""" + ) -> asyncio.Task[dict]: + """Submit rollout.n sessions and return their background collection.""" uid = sample_fields.get("uid") if uid is None: raise ValueError("GatewayAgentFramework requires prompts['uid'] for TransferQueue output") @@ -604,20 +669,35 @@ async def _run_prompt_rollouts( partition_id=partition_id, sample_fields=sample_fields, ) - - # Prompt layer: rollout.n sessions race independently for the same uid. - # Successful sessions are written to TQ; failed sessions only affect this uid's stats. await tq.async_kv_put(key=uid, partition_id=partition_id, tag={"status": "running"}) - tasks = [ - self._run_agent_episode_with_concurrency_limit( - sample_fields=sample_fields, - sample_index=sample_index, - session_index=session_index, - global_steps=global_steps, - sampling_params=sampling_params, - ) - for session_index in range(num_sessions) - ] + tasks = [] + try: + for session_index in range(num_sessions): + task = await self._submit_agent_episode( + sample_fields=sample_fields, + sample_index=sample_index, + session_index=session_index, + global_steps=global_steps, + sampling_params=sampling_params, + ) + tasks.append(task) + except Exception: + await tq.async_kv_put(key=uid, partition_id=partition_id, tag={"status": "failure"}) + raise + task = asyncio.create_task( + self._collect_prompt_rollouts(uid=uid, tasks=tasks, global_steps=global_steps, partition_id=partition_id) + ) + return task + + async def _collect_prompt_rollouts( + self, + *, + uid: str, + tasks: list[asyncio.Task[tuple[list[Trajectory], dict[str, object]]]], + global_steps: int | None, + partition_id: str, + ) -> dict: + """Persist successful sessions and aggregate the prompt's result.""" outcomes = await asyncio.gather(*tasks, return_exceptions=True) success_sessions = 0 @@ -678,7 +758,7 @@ async def _run_prompt_rollouts( "failure_reasons": failure_reasons, } - async def _run_agent_episode_with_concurrency_limit( + async def _submit_agent_episode( self, *, sample_fields: dict[str, object], @@ -686,7 +766,7 @@ async def _run_agent_episode_with_concurrency_limit( session_index: int, global_steps: int | None, sampling_params: dict[str, object], - ) -> tuple[list[Trajectory], dict[str, object]]: + ) -> asyncio.Task[tuple[list[Trajectory], dict[str, object]]]: # Lazy-init semaphores on first use and rebind if the running loop # changed: asyncio.Semaphore binds to the loop at construction, but # Ray actors may run sessions on a different loop than __init__. @@ -709,25 +789,16 @@ async def _run_agent_episode_with_concurrency_limit( except KeyError as exc: raise ValueError(f"Unknown agent runner: {agent_name}") from exc - runner_cap = runner_config.max_concurrent_sessions - if runner_cap <= 0: - return await self._run_agent_episode( - sample_fields=sample_fields, - sample_index=sample_index, - session_index=session_index, - global_steps=global_steps, - runner_name=runner_name, - runner_config=runner_config, - sampling_params=sampling_params, - ) - - runner_semaphore = self._runner_semaphores.get(runner_name) - if runner_semaphore is None: - runner_semaphore = asyncio.Semaphore(runner_cap) - self._runner_semaphores[runner_name] = runner_semaphore + runner_semaphore = None + if runner_config.max_concurrent_sessions > 0: + runner_semaphore = self._runner_semaphores.get(runner_name) + if runner_semaphore is None: + runner_semaphore = asyncio.Semaphore(runner_config.max_concurrent_sessions) + self._runner_semaphores[runner_name] = runner_semaphore + await runner_semaphore.acquire() - async with runner_semaphore: - return await self._run_agent_episode( + task = asyncio.create_task( + self._run_agent_episode( sample_fields=sample_fields, sample_index=sample_index, session_index=session_index, @@ -736,6 +807,11 @@ async def _run_agent_episode_with_concurrency_limit( runner_config=runner_config, sampling_params=sampling_params, ) + ) + if runner_semaphore is not None: + # Completion also returns the slot when cancelled before the episode starts. + task.add_done_callback(lambda _task: runner_semaphore.release()) + return task async def _run_agent_episode( self, @@ -775,11 +851,15 @@ async def _run_agent_episode( tools_kwargs = sample_fields.get("tools_kwargs") tools_kwargs = dict(tools_kwargs or {}) tools_kwargs["_trace_identity"] = trace_identity + # Pin the replica to this batch's trainer step. Validation without + # global_steps, or any non-int / negative value, stays unversioned. + weight_version = global_steps if type(global_steps) is int and global_steps >= 0 else None async with _log_scope(parent_log): session = await self.gateway_manager.create_session( session_id, metadata={"_trace_identity": trace_identity}, sampling_params=dict(sampling_params), + weight_version=weight_version, ) logger.info( "session %s start: runner=%s sample_index=%s session_index=%s global_steps=%s", @@ -916,6 +996,8 @@ async def _run_agent_episode( for traj, (score, extra) in zip(session_trajectories, annotations, strict=True) ] + result_trajectories = _attach_runner_reward_info(result_trajectories, task_result) + self._log_trajectory_summary(session_id, result_trajectories) if run_dir is not None: await asyncio.to_thread(self._dump_trajectories, run_dir, session_id, result_trajectories) diff --git a/uni_agent/gateway/gateway.py b/uni_agent/gateway/gateway.py index 55ca7a8b7..94376c9c4 100644 --- a/uni_agent/gateway/gateway.py +++ b/uni_agent/gateway/gateway.py @@ -230,8 +230,9 @@ async def create_session( session_id: str, metadata: dict[str, Any] | None = None, sampling_params: dict[str, Any] | None = None, + weight_version: int | None = None, ) -> SessionHandle: - """Create an actor-owned session and return its provider-compatible handle.""" + """Create an actor-owned session, optionally bind its route, and return its handle.""" self._require_started() if session_id in self._sessions: raise RuntimeError(f"Session {session_id} already exists") @@ -248,24 +249,34 @@ async def create_session( sampling_params=sampling_params, enable_last_assistant_rollback=self._enable_last_assistant_rollback, metadata=metadata, + weight_version=weight_version, ) + bind_route = getattr(self._backend, "bind_route", None) + if bind_route is not None: + await bind_route(session_id=session_id, weight_version=weight_version) return handle async def finalize_session(self, session_id: str) -> list[Trajectory]: """Finalize a session, remove it from the actor, and return its trajectories.""" session = self._get_session(session_id) trajectories = await session.finalize() - self._sessions.pop(session_id, None) + await self._remove_session(session_id) return trajectories async def abort_session(self, session_id: str) -> None: - """Abort a session and remove it from the actor if it still exists.""" + """Abort a live session and remove its route even if the session is already gone.""" session = self._sessions.get(session_id) - if session is None: - return # Already finalized or aborted — treat as idempotent. - await session.abort() + if session is not None: + await session.abort() + await self._remove_session(session_id) + + async def _remove_session(self, session_id: str) -> None: self._sessions.pop(session_id, None) + release_route = getattr(self._backend, "release_route", None) + if release_route is not None: + await release_route(session_id=session_id) + async def get_session_state(self, session_id: str) -> dict[str, Any]: """Return a snapshot of a live session's state.""" session = self._get_session(session_id) diff --git a/uni_agent/gateway/manager.py b/uni_agent/gateway/manager.py index b96f72278..588d28781 100644 --- a/uni_agent/gateway/manager.py +++ b/uni_agent/gateway/manager.py @@ -1,6 +1,6 @@ """Driver-side gateway manager: owns the gateway actor pool and routes sessions. -The manager spawns ``GatewayActor`` handles, injects the ``LLMServerClient`` +The manager spawns ``GatewayActor`` handles, injects the LLM client backend into each, and tracks which actor owns each session so lifecycle calls forward to the right actor through Ray remote methods. """ @@ -12,19 +12,18 @@ import ray from uni_agent.gateway.config import GatewayActorConfig -from verl.workers.rollout.llm_server import LLMServerClient class GatewayManager: """Owns gateway actors and routes sessions to them. - Spawns ``gateway_count`` actors over the injected ``LLMServerClient`` backend + Spawns ``gateway_count`` actors over the injected LLM client backend and tracks which actor owns each session so lifecycle calls reach it. """ def __init__( self, - llm_client: LLMServerClient, + llm_client, *, gateway_count: int, gateway_actor_config: GatewayActorConfig | None = None, @@ -91,6 +90,7 @@ async def create_session(self, session_id: str, **kwargs): self._session_to_gateway_index.pop(session_id, None) raise + # TODO: Consider merging finalize_session and abort_session in a future PR. async def finalize_session(self, session_id: str): """Finalize a session on its owning actor, release the route, and return its trajectories.""" gateway, gateway_index = self._get_gateway(session_id) diff --git a/uni_agent/gateway/session/codec.py b/uni_agent/gateway/session/codec.py index a2ca4cbda..400410ef5 100644 --- a/uni_agent/gateway/session/codec.py +++ b/uni_agent/gateway/session/codec.py @@ -302,6 +302,7 @@ def _process_tool_calls_vllm( parser_cls = ToolParserManager.get_tool_parser(parser_name) parser_parameters = inspect.signature(parser_cls).parameters + # TODO: Avoid unused Qwen3 streaming argument deltas when parsing complete responses. if "tools" in parser_parameters: parser = parser_cls(self._tokenizer, tools=vllm_tools) else: diff --git a/uni_agent/gateway/session/session.py b/uni_agent/gateway/session/session.py index 6183db9c8..a1321594b 100644 --- a/uni_agent/gateway/session/session.py +++ b/uni_agent/gateway/session/session.py @@ -186,12 +186,15 @@ def __init__( sampling_params: dict[str, Any] | None = None, enable_last_assistant_rollback: bool = True, metadata: dict[str, Any] | None = None, + weight_version: int | None = None, ): """Create an active session bound to a handle and model codec.""" if prompt_length is not None and prompt_length <= 0: raise ValueError(f"prompt_length must be positive when set, got {prompt_length}") if response_length is not None and response_length <= 0: raise ValueError(f"response_length must be positive when set, got {response_length}") + if weight_version is not None and (type(weight_version) is not int or weight_version < 0): + raise ValueError(f"weight_version must be a non-negative integer when set, got {weight_version!r}") self.handle = handle self._codec = codec @@ -204,6 +207,7 @@ def __init__( self._enable_last_assistant_rollback = enable_last_assistant_rollback self._metadata = dict(metadata or {}) self._trace_identity = dict(self._metadata.get("_trace_identity") or {}) + self._weight_version = weight_version self.active_chains: list[ChainState] = [] self.materialized_chains: list[MaterializedChain] = [] self.reserved_chain_ids: set[int] = set() @@ -266,20 +270,35 @@ async def run_generation(self, request: InternalGenerationRequest, backend) -> G reserved_chain_id = encoded.chain_id try: - output = await backend.generate( - request_id=self.handle.session_id, - prompt_ids=encoded.context_ids, - sampling_params=encoded.sampling_params, - image_data=encoded.image_data, - video_data=encoded.video_data, - mm_processor_kwargs=encoded.mm_processor_kwargs, - ) + backend_kwargs = { + "prompt_ids": encoded.context_ids, + "sampling_params": encoded.sampling_params, + "image_data": encoded.image_data, + "video_data": encoded.video_data, + "mm_processor_kwargs": encoded.mm_processor_kwargs, + } + if self._weight_version is not None: + backend_kwargs["weight_version"] = self._weight_version + # TODO: Rename the backend's session-routing request_id to session_id, + # keeping it distinct from per-generation engine request IDs. + output = await backend.generate(request_id=self.handle.session_id, **backend_kwargs) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) from e except Exception as e: raise HTTPException(status_code=500, detail=f"{e.__class__.__name__}: {e}") from e response_ids = list(output.token_ids) + + generation_version = ( + output.extra_fields.get("min_global_steps"), + output.extra_fields.get("max_global_steps"), + ) + if self._weight_version is not None and generation_version != (self._weight_version, self._weight_version): + raise RuntimeError( + "backend generation version conflicts with session weight_version: " + f"expected {(self._weight_version, self._weight_version)!r}, got {generation_version!r}" + ) + assistant_logprobs = None if encoded.sampling_params.get("logprobs", False): if output.log_probs is None: @@ -305,12 +324,7 @@ async def run_generation(self, request: InternalGenerationRequest, backend) -> G encoded.buffer.response_ids = list(merged_token_ids[prompt_length:]) encoded.buffer.response_mask = list(response_mask) encoded.buffer.response_logprobs = list(response_logprobs or []) - encoded.buffer.generation_versions.append( - ( - output.extra_fields.get("min_global_steps"), - output.extra_fields.get("max_global_steps"), - ) - ) + encoded.buffer.generation_versions.append(generation_version) self._assert_response_logprob_alignment(encoded.buffer) # R3 router replay: the backend returns routing for the full context @@ -702,6 +716,7 @@ def _snapshot_last_assistant_start( video_data: list[Any] | None, tip_hash: str, ) -> LastAssistantStart: + """TODO: This method seems redundant. Remove it later.""" return LastAssistantStart( response_ids_len=len(buffer.response_ids), message_history_len=message_history_len, diff --git a/uni_agent/sandbox/registry.py b/uni_agent/sandbox/registry.py index 140e2dce2..791583236 100644 --- a/uni_agent/sandbox/registry.py +++ b/uni_agent/sandbox/registry.py @@ -3,10 +3,16 @@ Providers live in their own module and self-register via :func:`register_sandbox`; :func:`build_sandbox` imports that module lazily on first use, so an uninstalled provider SDK never blocks importing this package. + +External overlays (private backends) can attach extra providers +without adding them to :data:`SANDBOX_MODULES`. Set ``UNI_AGENT_SANDBOX_PLUGINS`` +to a comma-separated list of importable module names. Each module must call +:func:`register_sandbox` at import time. Missing modules fail closed. """ from __future__ import annotations +import os from collections.abc import Callable from importlib import import_module @@ -14,6 +20,9 @@ SANDBOX_REGISTRY: dict[str, type[Sandbox]] = {} +#: Comma-separated importable modules that call :func:`register_sandbox`. +SANDBOX_PLUGIN_MODULES_ENV = "UNI_AGENT_SANDBOX_PLUGINS" + #: provider name -> module that defines (and registers) it, for lazy loading. SANDBOX_MODULES: dict[str, str] = { "local": "uni_agent.sandbox.local", @@ -37,8 +46,24 @@ def decorator(cls: type[Sandbox]) -> type[Sandbox]: return decorator +def _plugin_module_names() -> tuple[str, ...]: + raw = os.environ.get(SANDBOX_PLUGIN_MODULES_ENV, "") + return tuple(part.strip() for part in raw.split(",") if part.strip()) + + +def _load_plugin_modules() -> None: + """Import overlay modules so they can self-register. ``import_module`` is cached.""" + for module_name in _plugin_module_names(): + try: + import_module(module_name) + except ImportError as exc: + raise ImportError( + f"Failed to import sandbox plugin module {module_name!r} from {SANDBOX_PLUGIN_MODULES_ENV}." + ) from exc + + def _load_sandbox_module(name: str) -> None: - """Import the module that registers provider ``name`` (no-op if unknown).""" + """Import the in-tree module that registers provider ``name`` (no-op if unknown).""" module_name = SANDBOX_MODULES.get(name) if module_name is None: return @@ -55,6 +80,8 @@ def get_sandbox_cls(name: str) -> type[Sandbox]: """Return a registered provider class by name, importing its module on first use.""" if name not in SANDBOX_REGISTRY: _load_sandbox_module(name) + if name not in SANDBOX_REGISTRY: + _load_plugin_modules() if name not in SANDBOX_REGISTRY: available = sorted(set(SANDBOX_REGISTRY) | set(SANDBOX_MODULES)) raise ValueError(f"Unknown sandbox provider: {name!r}. Available: {available}") diff --git a/uni_agent/tasks/swe_bench/reward.py b/uni_agent/tasks/swe_bench/reward.py index 7c1ec369d..b7b4808b1 100644 --- a/uni_agent/tasks/swe_bench/reward.py +++ b/uni_agent/tasks/swe_bench/reward.py @@ -86,6 +86,7 @@ def _get_eval_report(metadata, eval_output: str): # step 1: get logs eval status_map, found = _get_logs_eval(metadata, eval_output) eval_report["found_eval_status"] = found + eval_report["status_map"] = status_map if not found: return eval_report @@ -143,7 +144,9 @@ async def compute_reward(metadata, sandbox, eval_timeout: float = 600.0) -> dict resp = await sandbox.exec_shell(f"bash {eval_script_container} 2>&1", workdir="/testbed", timeout=eval_timeout) output, exit_code = resp.stdout, resp.exit_code execution_time = time.perf_counter() - execution_t0 + # TODO: Remove eval_completed; it is redundant with eval_exit_code == 0. result["eval_completed"] = exit_code == 0 + result["eval_exit_code"] = exit_code result["eval_execution_time"] = execution_time logger.info(f"eval finished in {execution_time:.1f}s (exit_code={exit_code})") diff --git a/uni_agent/tasks/swe_bench/task.py b/uni_agent/tasks/swe_bench/task.py index b86e88ce4..c07fda6dc 100644 --- a/uni_agent/tasks/swe_bench/task.py +++ b/uni_agent/tasks/swe_bench/task.py @@ -41,6 +41,7 @@ async def run(self) -> TaskResult: f"task config: {json.dumps(task_config_dump, indent=2)}" ) async with self.build_sandbox() as sandbox: + agent_error = None if cfg.run_oracle_solution: logger.info("applying gold patch to /testbed") await sandbox.write_file("/tmp/gold_patch.patch", sample["patch"]) @@ -56,10 +57,13 @@ async def run(self) -> TaskResult: workdir="/testbed", ) finished = agent_result.finished + agent_error = agent_result.info.get("error") from .reward import compute_reward result = await compute_reward(sample, sandbox, eval_timeout=cfg.eval_timeout) + if isinstance(agent_error, str) and agent_error.strip(): + result = {**result, "agent_error": agent_error} logger.info(f"task done: resolved={result['resolved']}") return TaskResult(