From 73aa7f41b1adfc387c7c88db04947b6d1211b75d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 11:13:53 -0800 Subject: [PATCH 01/10] feat(training): train a bundle's own agents with TRL's harness rollouts AgentEnvSessionFactory is an OpenEnv ResourceSessionFactory for TRL's HarnessRolloutWorker with harness_adapter=None: each rollout runs a bundle task as agent-env run does, its agents' model calls going through OpenEnv's capture proxy to the trainer's vLLM, and the task's verifiers give the reward. prompt_rows builds the dataset, a row per task, which names its task so the factory finds it from the prompt TRL forwards. The factory refuses an engine that returns no token ids, pins each session to the policy version TRL published when it was created, and keeps the proxy's session routes behind an admin key. Rollouts waiting on a restarting server now say so once. Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 8 +- src/agent_env/openenv/app.py | 6 +- src/agent_env/training/_openenv.py | 6 + src/agent_env/training/harness.py | 164 ++++++++++++++++++ .../test_optional_extras_stay_optional.py | 3 +- tst/unit/training/harness_test.py | 156 +++++++++++++++++ tst/unit/training/sessions_test.py | 18 +- uv.lock | 9 +- 8 files changed, 355 insertions(+), 15 deletions(-) create mode 100644 src/agent_env/training/harness.py create mode 100644 tst/unit/training/harness_test.py diff --git a/pyproject.toml b/pyproject.toml index 0720a1de..7f9ff35d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,12 @@ dependencies = [ openenv = [ "openenv>=0.9", ] +# agent_env.training.harness: TRL's harness rollouts on a bundle's own agents, through OpenEnv's capture proxy (0.10 +# and later). TRL and vLLM are the trainer's to pin. +harness = [ + "agentenv-framework[openenv]", + "openenv>=0.10", +] # agent_env.training.trl: TRL's trainers on a bundle served as an OpenEnv env. TRL itself is the trainer's to pin. trl = [ "agentenv-framework[openenv]", @@ -133,7 +139,7 @@ coreweave = [ "cwsandbox>=1.19.6,<2", # Reject exec streams that close without an exit status. ] dev = [ - "agentenv-framework[aws,boat,boxd,coreweave,daytona,gcp,openenv,sail,skyrl,tensorlake,tenki,trl,vercel,verl]", # exercise portable optional providers and stores in CI + "agentenv-framework[aws,boat,boxd,coreweave,daytona,gcp,harness,openenv,sail,skyrl,tensorlake,tenki,trl,vercel,verl]", # exercise portable optional providers and stores in CI "fastmcp>=3.4,<4", # the standalone-fastmcp mount path of the protocol; <4 keeps mcp 1.x "Pillow>=10.4", # the protocol's computer extra, whose tests the unit tier runs "av>=17", # the protocol's computer extra, whose tests the unit tier runs and which the recording tests decode with diff --git a/src/agent_env/openenv/app.py b/src/agent_env/openenv/app.py index becd7fd5..8aac5934 100644 --- a/src/agent_env/openenv/app.py +++ b/src/agent_env/openenv/app.py @@ -142,11 +142,11 @@ def _served(splits: dict[str, list[tuple[str, Task]]], except EpisodeError as e: refused[task.id] = str(e) continue - reward_of, problem = _reward_of(task, reward) + rewarded, problem = reward_of(task, reward) if problem: unscorable[task.id] = problem continue - kept.append(ServedTask(task, name, spec, reward_of)) + kept.append(ServedTask(task, name, spec, rewarded)) served[split] = tuple(kept) for why in refused.values(): logger.warning("left out: %s", why) @@ -158,7 +158,7 @@ def _served(splits: dict[str, list[tuple[str, Task]]], return served -def _reward_of(task: Task, reward: str | Callable[[dict[str, float]], float] | None +def reward_of(task: Task, reward: str | Callable[[dict[str, float]], float] | None ) -> tuple[Callable[[dict[str, float]], float | None] | None, str | None]: """``task``'s reward as a function of its scores, or why ``reward`` can't make one.""" if callable(reward): diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index 1ceabbd2..8b96ae40 100644 --- a/src/agent_env/training/_openenv.py +++ b/src/agent_env/training/_openenv.py @@ -48,6 +48,8 @@ logger = logging.getLogger(__name__) _waited: set[str] = set() +#: The servers whose refused connections rollouts have been warned they wait out, once a process. +_restarting: set[str] = set() _loop: asyncio.AbstractEventLoop | None = None _loop_lock = threading.Lock() @@ -301,6 +303,10 @@ async def reopened(url: str, timeout: float, *, dropped: bool, patience: float, waited = max(patience, timeout) if restarting else patience if not (restarting or _transient(e)) or time.monotonic() + REOPEN_SECONDS > began + waited: raise Failed(f"{lead}{e}") from e + if restarting and url not in _restarting: + _restarting.add(url) + logger.warning("the server at %s refuses connections, so rollouts wait up to %gs for it to come back", + url, waited) await asyncio.sleep(REOPEN_SECONDS) diff --git a/src/agent_env/training/harness.py b/src/agent_env/training/harness.py new file mode 100644 index 00000000..c401987c --- /dev/null +++ b/src/agent_env/training/harness.py @@ -0,0 +1,164 @@ +"""Train a bundle's own agents with TRL's harness rollouts (``trl.experimental.async_grpo.HarnessRolloutWorker`` with +``harness_adapter=None``): each rollout runs a task as ``agent-env run`` does, its agents' model calls going through +OpenEnv's capture proxy to the trainer's vLLM, which records the tokens TRL trains on, and the task's verifiers score it.""" + +from __future__ import annotations + +import logging +import re +import secrets +import socket +import uuid +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Any + +try: + from openenv.core.harness.capture import CaptureServer, require_llm, to_training_trace + from openenv.core.harness.capture.export import export_session + from openenv.core.harness.capture.sessions import Upstream + from openenv.core.harness.rollout import ResourceSession, ResourceSessionFactory, ToolResult, VerifyResult + from openenv.core.harness.training import TrainingTrace +except ModuleNotFoundError as exc: + if exc.name is None or exc.name.split(".")[0] != "openenv": + raise + raise ImportError("agent_env.training.harness requires pip install 'agentenv-framework[harness]'") from exc + +from agent_env.bundle.run import prepare_bundle +from agent_env.openenv.app import reward_of +from agent_env.store.routing import namespace_routing +from agent_env.task.task import Task, recorded_scores +from agent_env.task.teardown import teardown_run +from agent_env.task_step.context import TaskStepContext, sandbox_overrides +from agent_env.training._openenv import on_session_loop + +logger = logging.getLogger(__name__) +#: How a row's prompt names its task, since TRL hands the factory only the prompt. +_NAMED = re.compile(r"^AgentEnv task (\d+): (.+)$") + + +def prompt_rows(bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), + sandbox: str = "local") -> list[dict[str, Any]]: + """The trainer's dataset, a row per task ``AgentEnvSessionFactory`` plays with the same arguments: a system line + naming the task, which is how the factory finds it, then the task's first prompt, for TRL's logs. The model never + sees this prompt: the task's own agent builds its own.""" + with namespace_routing(): + prepared = prepare_bundle(bundle, tasks=tasks, evals=evals, sandbox=sandbox) + return [{"prompt": [{"role": "system", "content": f"AgentEnv task {index}: {task.id}"}, + {"role": "user", "content": _first_prompt(task)}], "task_id": task.id} + for index, (_, task) in enumerate(prepared.tasks)] + + +class AgentEnvSessionFactory(ResourceSessionFactory): + """OpenEnv's session factory for TRL's ``HarnessRolloutWorker``, built in its rollout process. It prepares the + bundle, checks that ``llm_url``'s engine returns token ids, and serves OpenEnv's capture proxy in front of it, + where the task's sandboxes reach it. TRL sets ``model`` to the policy it publishes; ``reward`` is as + ``build_app``'s, and ``agent`` swaps the agent each task deploys.""" + + def __init__(self, bundle: str | Path, *, llm_url: str, model: str, sampling: dict[str, Any] | None = None, + tasks: Sequence[str] = (), evals: Sequence[str] = (), sandbox: str = "local", + reward: str | Callable[[dict[str, float]], float] | None = None, agent: str | None = None, + api_key: str | None = None, max_model_calls: int = 0): + self.model = model + self._llm_url, self._api_key, self._sampling = llm_url, api_key, sampling + self._sandbox, self._agent, self._max_model_calls = sandbox, agent, max_model_calls + with namespace_routing(): + self._tasks = [task for _, task in prepare_bundle(bundle, tasks=tasks, evals=evals, sandbox=sandbox).tasks] + self._rewards = [_reward(task, reward) for task in self._tasks] + require_llm(llm_url, model, api_key=api_key, require_tokens=True) + self._proxy = CaptureServer(llm_url=llm_url, model=model, port=_free_port(), api_key=api_key, + admin_key=secrets.token_urlsafe(32)) + self._proxy.start() + + def create(self, task: Any, seed: int | None = None, episode_id: str | None = None) -> AgentEnvSession: + index = _index_named(task, [each.id for each in self._tasks]) + capture = self._proxy.registry.create( + upstream=Upstream(llm_url=self._llm_url, model=self.model, api_key=self._api_key), + capture_level="tokens", purpose="train", sampling=self._sampling, max_model_calls=self._max_model_calls, + task=self._tasks[index].id, episode=episode_id or uuid.uuid4().hex) + return AgentEnvSession(self, self._tasks[index], self._rewards[index], capture, self.model) + + def _overrides(self, key: str) -> dict[str, Any]: + return {**sandbox_overrides(self._sandbox), "litellm_base_url": f"http://127.0.0.1:{self._proxy.port}", + "litellm_api_key": key, **({"a2a_agent_id": self._agent} if self._agent else {})} + + +class AgentEnvSession(ResourceSession): + """One run of a task with its own agents, their model calls captured under this session's key.""" + + def __init__(self, factory: AgentEnvSessionFactory, task: Task, reward: Callable[[dict[str, float]], float | None], + capture: Any, model: str): + self._factory, self._task, self._reward, self._capture, self._model = factory, task, reward, capture, model + self._scores: dict[str, float] | None = None + + def initial_messages(self) -> list[dict[str, Any]]: + return [{"role": "user", "content": _first_prompt(self._task)}] + + def list_tools(self) -> list[Any]: + return [] + + def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: + return ToolResult(error="the task's own agent calls its tools in its sandbox; none are offered here") + + def wait_for_completion(self, timeout_s: float | None = None) -> int: + """Runs the task and tears it down: 0 when it ran to the end, 1 when it failed.""" + return on_session_loop(self._run(), timeout_s) + + def fetch_training_trace(self) -> TrainingTrace: + document = export_session(self._capture, include_messages=True, capture_level="tokens") + return to_training_trace(self._capture.graph, document) + + def verify(self, transcript: list[dict[str, Any]], final_state: Any | None = None) -> VerifyResult: + """The task's reward from its recorded scores; ``None``, which TRL leaves out of the update, when the run + failed or recorded no score its reward reads.""" + try: + reward = None if self._scores is None else self._reward(self._scores) + except LookupError: + reward = None + return VerifyResult(env_reward=reward, done=True, metrics={f"score/{step}": score + for step, score in (self._scores or {}).items()}) + + def close(self) -> None: + self._factory._proxy.registry.delete(self._capture.session_id) + + async def _run(self) -> int: + context = TaskStepContext(metadata={"task_id": self._task.id, + "user_overrides": self._factory._overrides(self._capture.session_id)}) + with namespace_routing(): + try: + await self._task.run(agent_model=self._model, context=context) + self._scores = recorded_scores(self._task.steps, context) + return 0 + except Exception as e: + logger.warning("task %s failed, so its rollout is left unscored: %s", self._task.id, e) + return 1 + finally: + await teardown_run(context) + + +def _reward(task: Task, reward: str | Callable[[dict[str, float]], float] | None + ) -> Callable[[dict[str, float]], float | None]: + of, why = reward_of(task, reward) + if of is None: + raise ValueError(why) + return of + + +def _index_named(prompt: Any, ids: list[str]) -> int: + for message in prompt if isinstance(prompt, list) else []: + if isinstance(message, dict) and (named := _NAMED.match(str(message.get("content") or ""))): + index = int(named[1]) + if index < len(ids) and ids[index] == named[2]: + return index + raise KeyError("the row's prompt names no task of this factory's: build the dataset with prompt_rows, from the " + "same bundle, tasks and evals") + + +def _first_prompt(task: Task) -> str: + return next((step.prompt for step in task.steps if step.type == "prompt_agent"), "") + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("", 0)) + return probe.getsockname()[1] diff --git a/tst/unit/test_optional_extras_stay_optional.py b/tst/unit/test_optional_extras_stay_optional.py index 79004823..1b441bef 100644 --- a/tst/unit/test_optional_extras_stay_optional.py +++ b/tst/unit/test_optional_extras_stay_optional.py @@ -32,6 +32,7 @@ def find_spec(self, fullname, path=None, target=None): _NOT_CORE = ( "agent_env.providers.sandbox_providers.sail_vm", "agent_env.providers.sandbox_providers.smol_vm", "agent_env.openenv", "agent_env.training.nemo_gym", "agent_env.training.verl", "agent_env.training.skyrl", + "agent_env.training.harness", "agent_env.env.gateway", "agent_env.task_step.task_steps.run_code_runner", ) @@ -173,7 +174,7 @@ def _run(script: str, blocked: list[str], *args: str) -> str: #: Extras every distribution of which another extra installs too, so blocking what only they install blocks nothing: #: verl's and SkyRL's pyarrow comes with trl's datasets. The test without any extra covers them. -_COVERED = {"verl": "trl", "skyrl": "trl"} +_COVERED = {"verl": "trl", "skyrl": "trl", "harness": "openenv"} @pytest.mark.parametrize("extra", sorted(set(_optional_requirements()) - set(_COVERED))) diff --git a/tst/unit/training/harness_test.py b/tst/unit/training/harness_test.py new file mode 100644 index 00000000..117f717f --- /dev/null +++ b/tst/unit/training/harness_test.py @@ -0,0 +1,156 @@ +"""``AgentEnvSessionFactory`` running a bundle's task with its own agent, whose model calls go through OpenEnv's capture +proxy to a stand-in vLLM that returns token ids and processed logprobs, as TRL's harness rollouts run it.""" + +import json +import logging + +import httpx +import pytest +from starlette.applications import Starlette +from starlette.responses import JSONResponse +from starlette.routing import Route + +from agent_env.config.runtime import Config +from agent_env.task_step.task_step import TaskStep +from agent_env.training import harness +from agent_env.training._openenv import Server +from agent_env.training.harness import AgentEnvSessionFactory, prompt_rows +from tst.unit.bundle._support import layout +from tst.unit.conftest import local_stores # noqa: F401 + +MODEL = "Qwen/Qwen3-4B" +ASKS = {"id": "asks", "type": "harness_test_asks", "verifier_id": "answer"} +asked: list[dict] = [] + + +class _Asks(TaskStep): + """The task's own agent, as far as the factory can tell: it asks the model its run's endpoint names, with its key and + model, and scores the answer.""" + + type = "harness_test_asks" + entity_refs = () + + def __init__(self, id, version, verifier_id="answer", fail=False, depends_on=None, fail_task_on_error=True): + super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) + self.verifier_id, self.fail = verifier_id, fail + + def to_dict(self): + return {**super().to_dict(), "verifier_id": self.verifier_id, "fail": self.fail} + + @classmethod + def from_dict(cls, data): + return cls(**cls._base_from_dict(data), verifier_id=data["verifier_id"], fail=data.get("fail", False)) + + async def execute(self, context): + if self.fail: + raise RuntimeError("the agent crashed") + overrides = context.metadata["user_overrides"] + async with httpx.AsyncClient() as client: + answer = (await client.post( + f"{overrides['litellm_base_url']}/v1/chat/completions", + headers={"Authorization": f"Bearer {overrides['litellm_api_key']}"}, + json={"model": context.agent_model, "messages": [{"role": "user", "content": "What is 6 * 7?"}]}, + )).json()["choices"][0]["message"]["content"] + asked.append({"model": context.agent_model, "key": overrides["litellm_api_key"], "answer": answer}) + context.metadata.setdefault("verifications", {})[self.verifier_id] = {"score": float("42" in answer)} + return context + + +@pytest.fixture +def engine(): + """A stand-in for vLLM run with --return-tokens-as-token-ids and processed logprobs.""" + requests, served = [], [MODEL] + + async def models(request): + return JSONResponse({"data": [{"id": model, "object": "model"} for model in served]}) + + async def completions(request): + body = await request.json() + requests.append(body) + temperature = body.get("temperature") or 1.0 + prompt = [ord(char) % 500 + 1 for message in body["messages"] for char in str(message.get("content") or "")] + reply = "It is 42." + ids = [ord(char) % 500 + 600 for char in reply] + content = [{"token": f"token_id:{token}", "logprob": -0.5 / temperature, "bytes": None, + "top_logprobs": [{"token": f"token_id:{token}", "logprob": -0.5 / temperature}, + {"token": f"token_id:{token + 1}", "logprob": -3.0 / temperature}]} + for token in ids] + return JSONResponse({ + "id": "chatcmpl-1", "object": "chat.completion", "model": body["model"], "prompt_token_ids": prompt, + "choices": [{"index": 0, "message": {"role": "assistant", "content": reply}, "finish_reason": "stop", + "token_ids": ids, "logprobs": {"content": content}}], + "usage": {"prompt_tokens": len(prompt), "completion_tokens": len(ids), "total_tokens": len(prompt) + len(ids)}, + }) + + server = Server(Starlette(routes=[Route("/v1/models", models), Route("/v1/chat/completions", completions, + methods=["POST"])])) + server.requests, server.served = requests, served + yield server + server.stop() + + +@pytest.fixture(autouse=True) +def registries(monkeypatch): + steps = {**Config().task_step_registry(), _Asks.type: _Asks} + monkeypatch.setattr(Config, "task_step_registry", lambda self: steps) + + +@pytest.fixture +def bundle(tmp_path, monkeypatch, local_stores): # noqa: F811 + monkeypatch.setenv("HOME", str(tmp_path)) + asked.clear() + return layout(tmp_path / "quiz", {"tasks/sum.json": json.dumps([ASKS]), + "tasks/broken.json": json.dumps([{**ASKS, "fail": True}])}) + + +def test_the_rows_name_their_tasks_and_a_session_runs_its_task_with_its_calls_captured_and_scored(engine, bundle): + rows = prompt_rows(bundle) + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, sampling={"temperature": 1.0}) + engine.served.append("trl-policy-v3") + factory.model = "trl-policy-v3" + session = factory.create(next(row["prompt"] for row in rows if row["task_id"].endswith("sum"))) + factory.model = "trl-policy-v4" + + assert session.wait_for_completion() == 0 + trace = session.fetch_training_trace() + verified = session.verify([]) + session.close() + + [call] = asked + [turn] = trace.turns + assert (call["model"], call["answer"]) == ("trl-policy-v3", "It is 42.") + assert (verified.env_reward, verified.metrics) == (1.0, {"score/asks": 1.0}) + assert turn.completion_token_ids == [ord(char) % 500 + 600 for char in "It is 42."] + assert turn.loss_mask == [0] * len(turn.prompt_token_ids) + [1] * len(turn.completion_token_ids) + assert engine.requests[-1]["model"] == "trl-policy-v3" + assert factory._proxy.registry.get(call["key"]) is None + + +def test_a_run_that_fails_is_unscored_rather_than_zero(engine, bundle, caplog): + rows = prompt_rows(bundle) + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, sampling={"temperature": 1.0}) + session = factory.create(next(row["prompt"] for row in rows if row["task_id"].endswith("broken"))) + + with caplog.at_level(logging.WARNING, logger=harness.__name__): + assert session.wait_for_completion() == 1 + assert session.verify([]).env_reward is None + assert "failed, so its rollout is left unscored: " in caplog.text + + +def test_a_prompt_naming_no_task_of_the_factorys_is_refused(engine, bundle): + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL) + + for prompt in ([{"role": "user", "content": "What is 6 * 7?"}], + [{"role": "system", "content": "AgentEnv task 0: another-bundle/sum"}]): + with pytest.raises(KeyError, match="build the dataset with prompt_rows"): + factory.create(prompt) + + +def test_an_engine_without_token_ids_is_refused_before_any_rollout(bundle, monkeypatch): + def eval_only(llm_url, model, **_): + raise RuntimeError(f"{llm_url} is eval-only: no token ids") + + monkeypatch.setattr(harness, "require_llm", eval_only) + + with pytest.raises(RuntimeError, match="eval-only"): + AgentEnvSessionFactory(bundle, llm_url="http://127.0.0.1:9", model=MODEL) diff --git a/tst/unit/training/sessions_test.py b/tst/unit/training/sessions_test.py index f515e115..b74f27f9 100644 --- a/tst/unit/training/sessions_test.py +++ b/tst/unit/training/sessions_test.py @@ -11,16 +11,16 @@ from agent_env.training._openenv import Failed, reopened -@pytest.mark.parametrize(("cause", "waits_for_server", "patience", "opens", "attempts"), [ - (TimeoutError("timed out"), False, 5, True, 3), - (TimeoutError("timed out"), False, 0, False, 1), - (ConnectionRefusedError(61, "Connection refused"), False, 5, False, 1), - (ConnectionRefusedError(61, "Connection refused"), True, 5, True, 3), - (ConnectionRefusedError(61, "Connection refused"), True, 0, True, 3), - (ConnectionError("server rejected WebSocket connection: HTTP 403"), True, 5, False, 1), +@pytest.mark.parametrize(("cause", "waits_for_server", "patience", "opens", "attempts", "warned"), [ + (TimeoutError("timed out"), False, 5, True, 3, 0), + (TimeoutError("timed out"), False, 0, False, 1, 0), + (ConnectionRefusedError(61, "Connection refused"), False, 5, False, 1, 0), + (ConnectionRefusedError(61, "Connection refused"), True, 5, True, 3, 1), + (ConnectionRefusedError(61, "Connection refused"), True, 0, True, 3, 1), + (ConnectionError("server rejected WebSocket connection: HTTP 403"), True, 5, False, 1, 0), ], ids=["busy", "busy-past-patience", "refused", "restarting", "restarting-past-patience", "token-refused"]) def test_a_session_is_tried_again_while_its_server_is_busy_and_given_up_at_once_when_none_listens( - monkeypatch, cause, waits_for_server, patience, opens, attempts): + monkeypatch, caplog, cause, waits_for_server, patience, opens, attempts, warned): tried = [] async def flaky(url, timeout, *, token=None): @@ -33,6 +33,7 @@ async def flaky(url, timeout, *, token=None): monkeypatch.setattr(_openenv, "opened", flaky) monkeypatch.setattr(_openenv, "REOPEN_SECONDS", 0.01) + monkeypatch.setattr(_openenv, "_restarting", set()) if opens: assert asyncio.run(reopened("http://busy", 5, dropped=False, patience=patience, @@ -41,3 +42,4 @@ async def flaky(url, timeout, *, token=None): with pytest.raises(Failed, match="couldn't open a session on the server at http://busy"): asyncio.run(reopened("http://busy", 5, dropped=False, patience=patience, waits_for_server=waits_for_server)) assert len(tried) == attempts + assert caplog.text.count("refuses connections, so rollouts wait up to 5s for it to come back") == warned diff --git a/uv.lock b/uv.lock index 17e19491..26252fae 100644 --- a/uv.lock +++ b/uv.lock @@ -137,6 +137,9 @@ gcp = [ { name = "google-cloud-storage" }, { name = "requests" }, ] +harness = [ + { name = "openenv" }, +] nemo-gym = [ { name = "nemo-gym", marker = "python_full_version >= '3.13.14'" }, { name = "omegaconf", marker = "python_full_version >= '3.13.14'" }, @@ -184,7 +187,8 @@ verl = [ [package.metadata] requires-dist = [ { name = "a2a-sdk", specifier = "==0.3.26" }, - { name = "agentenv-framework", extras = ["aws", "boat", "boxd", "coreweave", "daytona", "gcp", "openenv", "sail", "skyrl", "tensorlake", "tenki", "trl", "vercel", "verl"], marker = "extra == 'dev'" }, + { name = "agentenv-framework", extras = ["aws", "boat", "boxd", "coreweave", "daytona", "gcp", "harness", "openenv", "sail", "skyrl", "tensorlake", "tenki", "trl", "vercel", "verl"], marker = "extra == 'dev'" }, + { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'harness'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'nemo-gym'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'skyrl'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'trl'" }, @@ -219,6 +223,7 @@ requires-dist = [ { name = "moto", marker = "extra == 'dev'", specifier = ">=5.0.0" }, { name = "nemo-gym", marker = "python_full_version >= '3.13.14' and extra == 'nemo-gym'", specifier = ">=0.7.0rc0" }, { name = "omegaconf", marker = "python_full_version >= '3.13.14' and extra == 'nemo-gym'", specifier = ">=2.3" }, + { name = "openenv", marker = "extra == 'harness'", specifier = ">=0.10" }, { name = "openenv", marker = "extra == 'openenv'", specifier = ">=0.9" }, { name = "packaging", specifier = ">=22" }, { name = "pillow", marker = "extra == 'dev'", specifier = ">=10.4" }, @@ -250,7 +255,7 @@ requires-dist = [ { name = "uvicorn", specifier = ">=0.47" }, { name = "vercel-sandbox", marker = "extra == 'vercel'", specifier = ">=0.7.0,<0.8" }, ] -provides-extras = ["openenv", "trl", "nemo-gym", "verl", "skyrl", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] +provides-extras = ["openenv", "harness", "trl", "nemo-gym", "verl", "skyrl", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] [[package]] name = "agentenv-framework-protocol" From ec820d09882af1762bf07286bb335dfc210094ec Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 11:27:07 -0800 Subject: [PATCH 02/10] fix(training): give the solver alone the capture proxy as its model endpoint The run-wide litellm_base_url and litellm_api_key overrides also reach the rubric judge and envs that call a model, so a judge would have asked the policy through the proxy and its tokens been trained on. The factory now sets them in step_params for the task's solver deploys (solver_steps), which deploy_agent reads over the run-wide ones, and agent= swaps the solver as agent-env run --agent does. prepare_bundle's solver_endpoint tells preflight the caller gives those deploys their endpoint, so training needs no endpoint configured while a judge still does. A task whose solver no deploy_agent step deploys is left out. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/preflight.py | 22 ++-- src/agent_env/bundle/run.py | 23 ++-- .../task_step/task_steps/deploy_agent.py | 20 ++- src/agent_env/training/harness.py | 70 ++++++---- tst/unit/bundle/preflight_test.py | 10 +- .../test_step_overrides_agent_and_model.py | 19 +++ tst/unit/training/harness_test.py | 120 +++++++++++++----- 7 files changed, 202 insertions(+), 82 deletions(-) diff --git a/src/agent_env/bundle/preflight.py b/src/agent_env/bundle/preflight.py index beb026b4..b4c8e30d 100644 --- a/src/agent_env/bundle/preflight.py +++ b/src/agent_env/bundle/preflight.py @@ -80,7 +80,7 @@ from agent_env.utils.docker.build import DEFAULT_BUILD_PLATFORM, docker_plugin_missing, docker_unreachable from .parse import BundleError, BundleKind -from .plan import AgentOverride, Plan +from .plan import AgentOverride, Plan, solver_steps from .resolve import BuiltImage, Reference, ReferencedImage, ResolvedEntry, build_step, played_deploys _SHOWN = 3 @@ -96,13 +96,16 @@ class Preflight: built_here: frozenset[str] = frozenset() # the ids of the images the bundle builds that a deploy here runs -def preflight_run(plan: Plan, tasks: Iterable[ResolvedEntry], sandbox: str | None, *, played: bool = False) -> Preflight: +def preflight_run(plan: Plan, tasks: Iterable[ResolvedEntry], sandbox: str | None, *, played: bool = False, + solver_endpoint: bool = False) -> Preflight: """Check where each of ``tasks`` deploys, with ``sandbox`` overriding every deploy's provider, and name the infra - envs the run builds; ``played`` skips each task's played agent's deploys (``played_deploys``). Raises BundleError - listing every problem found. Reads the store but writes nothing.""" + envs the run builds; ``played`` skips each task's played agent's deploys (``played_deploys``), and with + ``solver_endpoint`` the caller gives the solver's deploys (``solver_steps``) their model endpoint. Raises + BundleError listing every problem found. Reads the store but writes nothing.""" walk = _Walk(plan, sandbox) for task in tasks: - walk.task(task, skip=played_deploys(task) if played else frozenset()) + walk.task(task, skip=played_deploys(task) if played else frozenset(), + given=frozenset(solver_steps(task.config)) if solver_endpoint else frozenset()) return walk.finish() @@ -148,7 +151,7 @@ def __init__(self, plan: Plan, sandbox: str | None): # Each deploy of an image the bundle builds: (where, what, its provider's links, the images it accepts on each) self.uses: dict[str, list[tuple[str, str, list[SandboxProvider], Callable[[SandboxProvider], Accepts]]]] = {} - def task(self, task: ResolvedEntry, skip: frozenset[str]) -> None: + def task(self, task: ResolvedEntry, skip: frozenset[str], given: frozenset[str] = frozenset()) -> None: steps = [build_step(config) for config in task.config] sandboxes = {step.sandbox_name: step for step in steps if isinstance(step, DeploySandboxTaskStep)} solver = self.plan.agent.steps.get(task.entry.id, ()) if self.plan.agent else () @@ -161,7 +164,7 @@ def task(self, task: ResolvedEntry, skip: frozenset[str]) -> None: if isinstance(step, DeployEnvTaskStep): self._env(where, step) elif isinstance(step, DeployAgentTaskStep): - self._agent(where, step, sandboxes, override) + self._agent(where, step, sandboxes, override, endpoint_given=step.id in given) elif isinstance(step, DeploySandboxTaskStep): self._sandbox(where, step) elif isinstance(step, RubricsVerifierTaskStep): @@ -289,7 +292,7 @@ def _planned_version(self, kind: EntityKind, entity_id: str, version: int | None return version if version is not None else self.plan.store_latest.get((kind, entity_id)) def _agent(self, where: str, step: DeployAgentTaskStep, sandboxes: dict[str, DeploySandboxTaskStep], - override: AgentOverride | None = None) -> None: + override: AgentOverride | None = None, endpoint_given: bool = False) -> None: """Check a deploy of the step's agent, or of ``override``, the agent ``--agent`` names for the solver.""" accepts: Callable[[SandboxProvider], Accepts] = A2AAgent.accepts if step.sandbox_name: @@ -303,7 +306,8 @@ def _agent(self, where: str, step: DeployAgentTaskStep, sandboxes: dict[str, Dep accepts = functools.partial(A2AAgent.accepts, sandbox_mode=linked.sandbox_mode) # the VM loads its image else: provider = _provider(self.sandbox or step.sandbox_type, get_agent_sandbox_provider) - given = frozenset({*step.env_vars, *([ENV_BASE_URL] if step.litellm_base_url else [])}) + given = frozenset({*step.env_vars, *([ENV_BASE_URL] if step.litellm_base_url else []), + *(_MODEL_VARS if endpoint_given else ())}) agent_id, version = ((override.agent_id, override.version) if override is not None else (step.a2a_agent_id, step.a2a_agent_version)) self._agent_deploy(where, provider, agent_id, version, "names no agent", accepts, given=given) diff --git a/src/agent_env/bundle/run.py b/src/agent_env/bundle/run.py index 07bea46e..4bef58d2 100644 --- a/src/agent_env/bundle/run.py +++ b/src/agent_env/bundle/run.py @@ -251,7 +251,7 @@ def run_bundle( with Interrupts() as interrupts: overrides = _overrides(sandbox, user_model, judge_model) runs = interrupts.run(_run_all(plan, list(prepared.tasks), model, overrides, - _solver_params(plan, prepared.materialization), keep, say, interrupts)) + solver_params(plan, prepared.materialization), keep, say, interrupts)) _mark_cancelled(runs, interrupts.reason) result = _bundle_run(plan, prepared.materialization, runs, _skipped(plan, entries, every=not tasks and not evals)) @@ -274,17 +274,20 @@ def prepare_bundle( on_progress: Callable[[str], None] | None = None, id_root: str | None = None, played: bool = False, + agent: str | None = None, + solver_endpoint: bool = False, ) -> PreparedBundle: """Do what ``run_bundle`` with these arguments does before its first task runs, and stop there. ``played`` leaves out what each task's only prompted agent would deploy, for a caller that plays that agent itself - (``EpisodeSpec``), so the agent needn't exist. + (``EpisodeSpec``), so the agent needn't exist. With ``solver_endpoint`` the caller gives the solver's deploys their + model endpoint in each run's ``step_params``, so the config needn't set one. It doesn't turn namespace routing on: the caller keeps it on for as long as it runs the tasks (``agent_env.store.routing.namespace_routing``), which suits a service that runs nothing else. Raises what ``run_bundle`` raises before its first task runs, except for a running event loop, which it doesn't need.""" composing = _composing(sandbox) - return _prepare(root, tasks, evals, sandbox, id_root, None, played, _progress(on_progress), _hint(id_root), - composing) + return _prepare(root, tasks, evals, sandbox, id_root, agent, played, _progress(on_progress), _hint(id_root), + composing, solver_endpoint=solver_endpoint) def _write_bundle( @@ -373,8 +376,9 @@ def _bundle_run(plan: Plan, materialization: Materialization, runs: tuple[TaskRu def _prepare(root: Path | str, tasks: Sequence[str], evals: Sequence[str], sandbox: str | None, id_root: str | None, agent: str | None, played: bool, say: Callable[[str], None], declaring: Declaring, - composing: Composing = Composing()) -> PreparedBundle: - plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent, played=played) + composing: Composing = Composing(), *, solver_endpoint: bool = False) -> PreparedBundle: + plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent, played=played, + solver_endpoint=solver_endpoint) _say_plan(plan, say, declaring) materialization = _written_with_infra(plan, preflight, say, composing) to_run = tuple((entry, Task.get(entry.id, materialization.version_of("task", entry.id))) for entry in _to_run(plan)) @@ -382,7 +386,8 @@ def _prepare(root: Path | str, tasks: Sequence[str], evals: Sequence[str], sandb def _planned(root: Path | str, tasks: Sequence[str], evals: Sequence[str], id_root: str | None, - sandbox: str | None, agent: str | None = None, *, played: bool = False) -> tuple[Plan, Preflight]: + sandbox: str | None, agent: str | None = None, *, played: bool = False, + solver_endpoint: bool = False) -> tuple[Plan, Preflight]: resolved = resolve_bundle(parse_bundle(Path(root), id_root=id_root)) plan = plan_bundle(without_played_agents(resolved) if played else resolved, tasks=tasks, evals=evals, installed=installed_bundles, agent=agent) @@ -390,7 +395,7 @@ def _planned(root: Path | str, tasks: Sequence[str], evals: Sequence[str], id_ro to_run = set(_to_run(plan)) return plan, preflight_run(plan, [write.source for write in plan.writes if write.kind is BundleKind.TASK and write.source.entry in to_run], sandbox, - played=played) + played=played, solver_endpoint=solver_endpoint) def _composing(sandbox: str | None, emulate: bool = True) -> Composing: @@ -445,7 +450,7 @@ def _hint(id_root: str | None) -> Declaring: return Declaring.QUIET if id_root is not None else Declaring.HINT -def _solver_params(plan: Plan, materialization: Materialization) -> dict[str, dict[str, dict]]: +def solver_params(plan: Plan, materialization: Materialization) -> dict[str, dict[str, dict]]: """Each task's ``step_params`` that deploy ``Plan.agent`` at its solver steps: at the version this run wrote, for an agent it writes, or the store's version the plan read.""" if plan.agent is None: diff --git a/src/agent_env/task_step/task_steps/deploy_agent.py b/src/agent_env/task_step/task_steps/deploy_agent.py index 53baf931..8e763e70 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -365,13 +365,7 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas requested_disk = user_overrides.get("agent_disk_size_gb") or self.disk_size_gb min_disk = agent.metadata.get("min_disk_size_gb") agent_disk_size_gb = max(requested_disk or 0, min_disk or 0) or None - env_vars = dict(self.env_vars) if self.env_vars else {} - override_litellm_key = user_overrides.get("litellm_api_key") - if override_litellm_key: - env_vars["LITELLM_API_KEY"] = config.resolve_reference(override_litellm_key) - resolved_litellm_base_url = user_overrides.get("litellm_base_url") or self.litellm_base_url - if resolved_litellm_base_url: - env_vars["LITELLM_BASE_URL"] = resolved_litellm_base_url + env_vars = {**(self.env_vars or {}), **self._model_endpoint(context)} if self.role is not None: env_vars.setdefault(ROLE_ENV_VAR, self.role) # the generated MCP CLI sends its role from here @@ -509,6 +503,18 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas )) return context + def _model_endpoint(self, context: TaskStepContext) -> dict[str, str]: + """The agent's ``LITELLM_API_KEY`` and ``LITELLM_BASE_URL`` in this run: this step's ``step_params`` over the + run-wide overrides, so a trainer can point the solver alone at its own endpoint, and the base URL over the + step's.""" + overrides = {**context.metadata.get("user_overrides", {}), **self.step_param_overrides(context)} + endpoint = {} + if overrides.get("litellm_api_key"): + endpoint["LITELLM_API_KEY"] = get_config().resolve_reference(overrides["litellm_api_key"]) + if base_url := overrides.get("litellm_base_url") or self.litellm_base_url: + endpoint["LITELLM_BASE_URL"] = base_url + return endpoint + def _resolved_ttl_seconds(self, context: TaskStepContext) -> int: override = context.metadata.get("user_overrides", {}).get("ttl_seconds") if override is not None: diff --git a/src/agent_env/training/harness.py b/src/agent_env/training/harness.py index c401987c..f3bb9b36 100644 --- a/src/agent_env/training/harness.py +++ b/src/agent_env/training/harness.py @@ -24,12 +24,14 @@ raise raise ImportError("agent_env.training.harness requires pip install 'agentenv-framework[harness]'") from exc -from agent_env.bundle.run import prepare_bundle +from agent_env.bundle.plan import solver_steps +from agent_env.bundle.run import prepare_bundle, solver_params from agent_env.openenv.app import reward_of from agent_env.store.routing import namespace_routing from agent_env.task.task import Task, recorded_scores from agent_env.task.teardown import teardown_run from agent_env.task_step.context import TaskStepContext, sandbox_overrides +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep from agent_env.training._openenv import on_session_loop logger = logging.getLogger(__name__) @@ -37,23 +39,22 @@ _NAMED = re.compile(r"^AgentEnv task (\d+): (.+)$") -def prompt_rows(bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), - sandbox: str = "local") -> list[dict[str, Any]]: +def prompt_rows(bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), sandbox: str = "local", + agent: str | None = None) -> list[dict[str, Any]]: """The trainer's dataset, a row per task ``AgentEnvSessionFactory`` plays with the same arguments: a system line naming the task, which is how the factory finds it, then the task's first prompt, for TRL's logs. The model never sees this prompt: the task's own agent builds its own.""" - with namespace_routing(): - prepared = prepare_bundle(bundle, tasks=tasks, evals=evals, sandbox=sandbox) return [{"prompt": [{"role": "system", "content": f"AgentEnv task {index}: {task.id}"}, {"role": "user", "content": _first_prompt(task)}], "task_id": task.id} - for index, (_, task) in enumerate(prepared.tasks)] + for index, (task, _) in enumerate(_trained(bundle, tasks, evals, sandbox, agent))] class AgentEnvSessionFactory(ResourceSessionFactory): """OpenEnv's session factory for TRL's ``HarnessRolloutWorker``, built in its rollout process. It prepares the bundle, checks that ``llm_url``'s engine returns token ids, and serves OpenEnv's capture proxy in front of it, - where the task's sandboxes reach it. TRL sets ``model`` to the policy it publishes; ``reward`` is as - ``build_app``'s, and ``agent`` swaps the agent each task deploys.""" + where the task's solver reaches it: the agent a step prompts, which it deploys, its other agents and judges keeping + the configured endpoint. TRL sets ``model`` to the policy it publishes; ``reward`` is as ``build_app``'s, and + ``agent`` the agent the solver deploys, as ``agent-env run --agent``'s.""" def __init__(self, bundle: str | Path, *, llm_url: str, model: str, sampling: dict[str, Any] | None = None, tasks: Sequence[str] = (), evals: Sequence[str] = (), sandbox: str = "local", @@ -61,34 +62,34 @@ def __init__(self, bundle: str | Path, *, llm_url: str, model: str, sampling: di api_key: str | None = None, max_model_calls: int = 0): self.model = model self._llm_url, self._api_key, self._sampling = llm_url, api_key, sampling - self._sandbox, self._agent, self._max_model_calls = sandbox, agent, max_model_calls - with namespace_routing(): - self._tasks = [task for _, task in prepare_bundle(bundle, tasks=tasks, evals=evals, sandbox=sandbox).tasks] - self._rewards = [_reward(task, reward) for task in self._tasks] + self._sandbox, self._max_model_calls = sandbox, max_model_calls + self._trained = _trained(bundle, tasks, evals, sandbox, agent) + self._rewards = [_reward(task, reward) for task, _ in self._trained] require_llm(llm_url, model, api_key=api_key, require_tokens=True) self._proxy = CaptureServer(llm_url=llm_url, model=model, port=_free_port(), api_key=api_key, admin_key=secrets.token_urlsafe(32)) self._proxy.start() def create(self, task: Any, seed: int | None = None, episode_id: str | None = None) -> AgentEnvSession: - index = _index_named(task, [each.id for each in self._tasks]) + index = _index_named(task, [each.id for each, _ in self._trained]) + played, solver = self._trained[index] capture = self._proxy.registry.create( upstream=Upstream(llm_url=self._llm_url, model=self.model, api_key=self._api_key), capture_level="tokens", purpose="train", sampling=self._sampling, max_model_calls=self._max_model_calls, - task=self._tasks[index].id, episode=episode_id or uuid.uuid4().hex) - return AgentEnvSession(self, self._tasks[index], self._rewards[index], capture, self.model) - - def _overrides(self, key: str) -> dict[str, Any]: - return {**sandbox_overrides(self._sandbox), "litellm_base_url": f"http://127.0.0.1:{self._proxy.port}", - "litellm_api_key": key, **({"a2a_agent_id": self._agent} if self._agent else {})} + task=played.id, episode=episode_id or uuid.uuid4().hex) + endpoint = {"litellm_base_url": f"http://127.0.0.1:{self._proxy.port}", "litellm_api_key": capture.session_id} + overrides = {**sandbox_overrides(self._sandbox), + "step_params": {step: {**params, **endpoint} for step, params in solver.items()}} + return AgentEnvSession(played, overrides, self._rewards[index], self._proxy.registry, capture, self.model) class AgentEnvSession(ResourceSession): """One run of a task with its own agents, their model calls captured under this session's key.""" - def __init__(self, factory: AgentEnvSessionFactory, task: Task, reward: Callable[[dict[str, float]], float | None], - capture: Any, model: str): - self._factory, self._task, self._reward, self._capture, self._model = factory, task, reward, capture, model + def __init__(self, task: Task, overrides: dict[str, Any], reward: Callable[[dict[str, float]], float | None], + registry: Any, capture: Any, model: str): + self._task, self._overrides, self._reward, self._model = task, overrides, reward, model + self._registry, self._capture = registry, capture self._scores: dict[str, float] | None = None def initial_messages(self) -> list[dict[str, Any]]: @@ -119,11 +120,10 @@ def verify(self, transcript: list[dict[str, Any]], final_state: Any | None = Non for step, score in (self._scores or {}).items()}) def close(self) -> None: - self._factory._proxy.registry.delete(self._capture.session_id) + self._registry.delete(self._capture.session_id) async def _run(self) -> int: - context = TaskStepContext(metadata={"task_id": self._task.id, - "user_overrides": self._factory._overrides(self._capture.session_id)}) + context = TaskStepContext(metadata={"task_id": self._task.id, "user_overrides": self._overrides}) with namespace_routing(): try: await self._task.run(agent_model=self._model, context=context) @@ -136,6 +136,26 @@ async def _run(self) -> int: await teardown_run(context) +def _trained(bundle: str | Path, tasks: Sequence[str], evals: Sequence[str], sandbox: str, + agent: str | None) -> list[tuple[Task, dict[str, dict[str, Any]]]]: + """The bundle's tasks prepared as ``agent-env run`` prepares them, each with its solver's deploys and their step + params, ``agent``'s if it swaps the solver; a task whose solver no ``deploy_agent`` step deploys is left out.""" + with namespace_routing(): + prepared = prepare_bundle(bundle, tasks=tasks, evals=evals, sandbox=sandbox, agent=agent, solver_endpoint=True) + swapped = solver_params(prepared.plan, prepared.materialization) + trained = [] + for entry, task in prepared.tasks: + solver = solver_steps([step.to_dict() for step in task.steps]) + if solver and all(step.type == DeployAgentTaskStep.type for step in task.steps if step.id in solver): + trained.append((task, {step: swapped.get(entry.id, {}).get(step, {}) for step in solver})) + else: + logger.warning("%s is left out: training needs the agent a step prompts deployed by a deploy_agent step", + task.id) + if not trained: + raise ValueError("no task of the bundle deploys, with a deploy_agent step, an agent a step prompts to train") + return trained + + def _reward(task: Task, reward: str | Callable[[dict[str, float]], float] | None ) -> Callable[[dict[str, float]], float | None]: of, why = reward_of(task, reward) diff --git a/tst/unit/bundle/preflight_test.py b/tst/unit/bundle/preflight_test.py index 4c92736b..d60fcda0 100644 --- a/tst/unit/bundle/preflight_test.py +++ b/tst/unit/bundle/preflight_test.py @@ -18,7 +18,7 @@ from agent_env.artifact.artifacts.docker_image import DockerImageArtifact from agent_env.artifact.store import get_artifact_store from agent_env.bundle import BundleError, BundleKind, dry_run_bundle, run_bundle -from agent_env.bundle.run import _write_bundle +from agent_env.bundle.run import _write_bundle, prepare_bundle from agent_env.bundle.resolve import build_step from agent_env.cli import cli from agent_env.config import get_config @@ -1176,12 +1176,18 @@ def test_an_agent_or_agent_judge_needs_a_model_endpoint_unless_its_step_gives_on "judge_a2a_agent_id": "solver", "criteria": [{"id": "c1", "description": "done"}]} direct = {**judge, "id": "direct", "use_agent_judge": False} given = {**AGENT, "env_vars": {"LITELLM_API_KEY": "sk-1"}, "litellm_base_url": "http://models:4000"} - _task(bundle_dir, [AGENT, judge], name="judged") + asks = {"id": "ask", "type": "prompt_agent", "prompt_id": "agent", "prompt": "go", "depends_on": ["agent"]} + _task(bundle_dir, [AGENT, asks, judge], name="judged") _task(bundle_dir, [given, direct], name="given") assert _problems(lambda: dry_run_bundle(bundle_dir, sandbox="local")) == [ "tasks/judged.json: step 'agent', tasks/judged.json: step 'judge' deploy an agent that calls the model endpoint, " "and no model endpoint is set: set LITELLM_BASE_URL, or [model] base_url in .agentenv/config.toml"] + with namespace_routing(): # a trainer gives the solver its endpoint, and the judge still needs the configured one + trained = lambda: prepare_bundle(bundle_dir, tasks=["judged"], sandbox="local", solver_endpoint=True) # noqa: E731 + assert _problems(trained) == [ + "tasks/judged.json: step 'judge' deploys an agent that calls the model endpoint, and no model endpoint is " + "set: set LITELLM_BASE_URL, or [model] base_url in .agentenv/config.toml"] assert [entry.name for entry in dry_run_bundle(bundle_dir, tasks=["given"], sandbox="local").runs] == ["given"] monkeypatch.setenv("LITELLM_BASE_URL", "http://models:4000") [problem] = _problems(lambda: dry_run_bundle(bundle_dir, sandbox="local")) diff --git a/tst/unit/task_step/test_step_overrides_agent_and_model.py b/tst/unit/task_step/test_step_overrides_agent_and_model.py index 7a7fd976..d9eb11de 100644 --- a/tst/unit/task_step/test_step_overrides_agent_and_model.py +++ b/tst/unit/task_step/test_step_overrides_agent_and_model.py @@ -71,6 +71,25 @@ def test_other_steps_overrides_are_ignored(self): assert _deploy_step()._resolve_agent_ref(ctx) == ("stored-agent", 7) +class TestDeployAgentModelEndpoint: + def test_the_steps_own_url_without_overrides(self): + assert _deploy_step(litellm_base_url="http://stored")._model_endpoint(_ctx()) == { + "LITELLM_BASE_URL": "http://stored"} + + def test_the_run_wide_endpoint_applies(self): + ctx = _ctx(run_overrides={"litellm_base_url": "http://run", "litellm_api_key": "sk-run"}) + assert _deploy_step(litellm_base_url="http://stored")._model_endpoint(ctx) == { + "LITELLM_API_KEY": "sk-run", "LITELLM_BASE_URL": "http://run"} + + def test_the_steps_override_wins_and_another_steps_is_ignored(self): + ctx = _ctx({"deploy-solver": {"litellm_base_url": "http://trainer", "litellm_api_key": "session"}, + "deploy-judge": {"litellm_base_url": "http://elsewhere"}}, + {"litellm_base_url": "http://run", "litellm_api_key": "sk-run"}) + assert _deploy_step()._model_endpoint(ctx) == {"LITELLM_API_KEY": "session", "LITELLM_BASE_URL": "http://trainer"} + assert _deploy_step(id="deploy-judge")._model_endpoint(ctx) == { + "LITELLM_API_KEY": "sk-run", "LITELLM_BASE_URL": "http://elsewhere"} + + class TestPromptModelAndEffort: def test_stored_values_without_overrides(self): assert _prompt_step()._resolve_model_and_effort(_ctx()) == ("stored-model", "stored-effort") diff --git a/tst/unit/training/harness_test.py b/tst/unit/training/harness_test.py index 117f717f..aadd5d22 100644 --- a/tst/unit/training/harness_test.py +++ b/tst/unit/training/harness_test.py @@ -1,5 +1,7 @@ -"""``AgentEnvSessionFactory`` running a bundle's task with its own agent, whose model calls go through OpenEnv's capture -proxy to a stand-in vLLM that returns token ids and processed logprobs, as TRL's harness rollouts run it.""" +"""``AgentEnvSessionFactory`` running a bundle's task with its own agents, as TRL's harness rollouts run it: the +solver's deploy alone is given OpenEnv's capture proxy as its model endpoint, in front of a stand-in vLLM that returns +token ids and processed logprobs, with no endpoint configured; its calls are captured and the task's verifier scores +the run.""" import json import logging @@ -10,8 +12,15 @@ from starlette.responses import JSONResponse from starlette.routing import Route +from agent_env.a2a_agent import A2AAgent +from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.artifact.store import get_artifact_store +from agent_env.bundle import preflight from agent_env.config.runtime import Config +from agent_env.store.routing import namespace_routing from agent_env.task_step.task_step import TaskStep +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep from agent_env.training import harness from agent_env.training._openenv import Server from agent_env.training.harness import AgentEnvSessionFactory, prompt_rows @@ -19,40 +28,58 @@ from tst.unit.conftest import local_stores # noqa: F401 MODEL = "Qwen/Qwen3-4B" -ASKS = {"id": "asks", "type": "harness_test_asks", "verifier_id": "answer"} +MODEL_ENDPOINT_MISSING = preflight.model_endpoint_missing # the probe, which the unit tier's conftest stubs +DEPLOY = {"id": "agent", "type": "deploy_agent", "env_ids": [], "a2a_agent_id": "theirs"} +ASK = {"id": "ask", "type": "prompt_agent", "prompt_id": "ask", "prompt": "What is 6 * 7?", "depends_on": ["agent"]} +CHECK = {"id": "check", "type": "harness_test_check", "verifier_id": "answer", "depends_on": ["ask"]} asked: list[dict] = [] -class _Asks(TaskStep): - """The task's own agent, as far as the factory can tell: it asks the model its run's endpoint names, with its key and - model, and scores the answer.""" +class _Deploys(DeployAgentTaskStep): + """Deploys nothing: records the model endpoint the run gives the agent.""" - type = "harness_test_asks" + async def execute(self, context): + if self.a2a_agent_id == "crashes": + raise RuntimeError("the agent crashed") + context.metadata["endpoint"] = self._model_endpoint(context) + return context + + +class _Asks(PromptAgentTaskStep): + """The deployed agent's turn: it asks the model at its endpoint, with its key and the run's model.""" + + async def execute(self, context): + endpoint = context.metadata["endpoint"] + async with httpx.AsyncClient() as client: + answer = (await client.post( + f"{endpoint['LITELLM_BASE_URL']}/v1/chat/completions", + headers={"Authorization": f"Bearer {endpoint['LITELLM_API_KEY']}"}, + json={"model": context.agent_model, "messages": [{"role": "user", "content": self.prompt}]}, + )).json()["choices"][0]["message"]["content"] + asked.append({"model": context.agent_model, "key": endpoint["LITELLM_API_KEY"], "answer": answer, + "overrides": context.metadata["user_overrides"]}) + context.metadata["answer"] = answer + return context + + +class _Checks(TaskStep): + type = "harness_test_check" entity_refs = () - def __init__(self, id, version, verifier_id="answer", fail=False, depends_on=None, fail_task_on_error=True): + def __init__(self, id, version, verifier_id="answer", depends_on=None, fail_task_on_error=True): super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) - self.verifier_id, self.fail = verifier_id, fail + self.verifier_id = verifier_id def to_dict(self): - return {**super().to_dict(), "verifier_id": self.verifier_id, "fail": self.fail} + return {**super().to_dict(), "verifier_id": self.verifier_id} @classmethod def from_dict(cls, data): - return cls(**cls._base_from_dict(data), verifier_id=data["verifier_id"], fail=data.get("fail", False)) + return cls(**cls._base_from_dict(data), verifier_id=data["verifier_id"]) async def execute(self, context): - if self.fail: - raise RuntimeError("the agent crashed") - overrides = context.metadata["user_overrides"] - async with httpx.AsyncClient() as client: - answer = (await client.post( - f"{overrides['litellm_base_url']}/v1/chat/completions", - headers={"Authorization": f"Bearer {overrides['litellm_api_key']}"}, - json={"model": context.agent_model, "messages": [{"role": "user", "content": "What is 6 * 7?"}]}, - )).json()["choices"][0]["message"]["content"] - asked.append({"model": context.agent_model, "key": overrides["litellm_api_key"], "answer": answer}) - context.metadata.setdefault("verifications", {})[self.verifier_id] = {"score": float("42" in answer)} + context.metadata.setdefault("verifications", {})[self.verifier_id] = { + "score": float("42" in context.metadata["answer"])} return context @@ -91,20 +118,31 @@ async def completions(request): @pytest.fixture(autouse=True) def registries(monkeypatch): - steps = {**Config().task_step_registry(), _Asks.type: _Asks} + steps = {**Config().task_step_registry(), "deploy_agent": _Deploys, "prompt_agent": _Asks, _Checks.type: _Checks} monkeypatch.setattr(Config, "task_step_registry", lambda self: steps) @pytest.fixture def bundle(tmp_path, monkeypatch, local_stores): # noqa: F811 monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setattr(preflight, "model_endpoint_missing", MODEL_ENDPOINT_MISSING) + monkeypatch.delenv("LITELLM_BASE_URL", raising=False) + monkeypatch.delenv("LITELLM_API_KEY", raising=False) asked.clear() - return layout(tmp_path / "quiz", {"tasks/sum.json": json.dumps([ASKS]), - "tasks/broken.json": json.dumps([{**ASKS, "fail": True}])}) - - -def test_the_rows_name_their_tasks_and_a_session_runs_its_task_with_its_calls_captured_and_scored(engine, bundle): - rows = prompt_rows(bundle) + for agent_id in ("theirs", "crashes", "mine"): + with namespace_routing(): + image = get_artifact_store().put_document(DockerImageArtifact( + id=f"{agent_id}-image", description="d", image_name="ghcr.io/acme/agent:1", + tar_gz_s3_url="s3://b/x.tar.gz")) + A2AAgent.put(id=agent_id, docker_image_artifact=image) + return layout(tmp_path / "quiz", {"tasks/sum.json": json.dumps([DEPLOY, ASK, CHECK]), + "tasks/broken.json": json.dumps([{**DEPLOY, "a2a_agent_id": "crashes"}, ASK]), + "tasks/unprompted.json": json.dumps([DEPLOY])}) + + +def test_the_solver_alone_calls_the_policy_through_the_proxy_and_the_run_is_captured_and_scored(engine, bundle, caplog): + with caplog.at_level(logging.WARNING, logger=harness.__name__): + rows = prompt_rows(bundle) factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, sampling={"temperature": 1.0}) engine.served.append("trl-policy-v3") factory.model = "trl-policy-v3" @@ -118,14 +156,31 @@ def test_the_rows_name_their_tasks_and_a_session_runs_its_task_with_its_calls_ca [call] = asked [turn] = trace.turns + assert sorted(row["task_id"].rsplit("/", 1)[1] for row in rows) == ["broken", "sum"] + assert "unprompted is left out: training needs the agent a step prompts deployed by a deploy_agent step" in caplog.text assert (call["model"], call["answer"]) == ("trl-policy-v3", "It is 42.") - assert (verified.env_reward, verified.metrics) == (1.0, {"score/asks": 1.0}) + assert call["overrides"]["step_params"] == {"agent": {"litellm_base_url": f"http://127.0.0.1:{factory._proxy.port}", + "litellm_api_key": call["key"]}} + assert not {"litellm_base_url", "litellm_api_key"} & call["overrides"].keys() + assert (verified.env_reward, verified.metrics) == (1.0, {"score/check": 1.0}) assert turn.completion_token_ids == [ord(char) % 500 + 600 for char in "It is 42."] assert turn.loss_mask == [0] * len(turn.prompt_token_ids) + [1] * len(turn.completion_token_ids) assert engine.requests[-1]["model"] == "trl-policy-v3" assert factory._proxy.registry.get(call["key"]) is None +def test_the_agent_named_is_the_one_the_solver_deploys(engine, bundle): + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, agent="mine") + session = factory.create(next(row["prompt"] for row in prompt_rows(bundle, agent="mine") + if row["task_id"].endswith("sum"))) + + assert session.wait_for_completion() == 0 + [call] = asked + assert call["overrides"]["step_params"]["agent"] == { + "a2a_agent_id": "mine", "a2a_agent_version": 1, "litellm_base_url": f"http://127.0.0.1:{factory._proxy.port}", + "litellm_api_key": call["key"]} + + def test_a_run_that_fails_is_unscored_rather_than_zero(engine, bundle, caplog): rows = prompt_rows(bundle) factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, sampling={"temperature": 1.0}) @@ -146,6 +201,11 @@ def test_a_prompt_naming_no_task_of_the_factorys_is_refused(engine, bundle): factory.create(prompt) +def test_a_bundle_with_no_solver_to_train_is_refused(engine, bundle): + with pytest.raises(ValueError, match="no task of the bundle deploys"): + AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, tasks=["unprompted"]) + + def test_an_engine_without_token_ids_is_refused_before_any_rollout(bundle, monkeypatch): def eval_only(llm_url, model, **_): raise RuntimeError(f"{llm_url} is eval-only: no token ids") From 3bf983ea69d7a0c652b28a03f55bfbcfc5a04277 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 11:30:02 -0800 Subject: [PATCH 03/10] test(training): give the unprompted agent its own endpoint Co-Authored-By: Claude Opus 5.5 --- tst/unit/training/harness_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tst/unit/training/harness_test.py b/tst/unit/training/harness_test.py index aadd5d22..4e2869ff 100644 --- a/tst/unit/training/harness_test.py +++ b/tst/unit/training/harness_test.py @@ -32,6 +32,7 @@ DEPLOY = {"id": "agent", "type": "deploy_agent", "env_ids": [], "a2a_agent_id": "theirs"} ASK = {"id": "ask", "type": "prompt_agent", "prompt_id": "ask", "prompt": "What is 6 * 7?", "depends_on": ["agent"]} CHECK = {"id": "check", "type": "harness_test_check", "verifier_id": "answer", "depends_on": ["ask"]} +GIVEN = {"litellm_base_url": "http://models:4000", "env_vars": {"LITELLM_API_KEY": "sk-1"}} # its own endpoint asked: list[dict] = [] @@ -137,7 +138,7 @@ def bundle(tmp_path, monkeypatch, local_stores): # noqa: F811 A2AAgent.put(id=agent_id, docker_image_artifact=image) return layout(tmp_path / "quiz", {"tasks/sum.json": json.dumps([DEPLOY, ASK, CHECK]), "tasks/broken.json": json.dumps([{**DEPLOY, "a2a_agent_id": "crashes"}, ASK]), - "tasks/unprompted.json": json.dumps([DEPLOY])}) + "tasks/unprompted.json": json.dumps([{**DEPLOY, **GIVEN}])}) def test_the_solver_alone_calls_the_policy_through_the_proxy_and_the_run_is_captured_and_scored(engine, bundle, caplog): From 92469265635f102f2b10f3c60281886ab4f6ca48 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 11:37:44 -0800 Subject: [PATCH 04/10] fix(training): leave out untrainable tasks before preparing the bundle, and a solver that judges itself A task whose rubric judge is its solver would have its grading trained on, so it's left out. Tasks are now chosen from the plan (planned_tasks) before prepare_bundle checks and writes, so what a left-out task alone needs, such as a configured endpoint for an unprompted agent, can't stop the rest. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/run.py | 13 ++++++++- src/agent_env/training/harness.py | 48 +++++++++++++++++++++---------- tst/unit/bundle/run_test.py | 5 ++-- tst/unit/training/harness_test.py | 10 +++++-- 4 files changed, 54 insertions(+), 22 deletions(-) diff --git a/src/agent_env/bundle/run.py b/src/agent_env/bundle/run.py index 4bef58d2..9cb4207b 100644 --- a/src/agent_env/bundle/run.py +++ b/src/agent_env/bundle/run.py @@ -41,7 +41,8 @@ from .plan import AliasUse, Plan, Write, plan_bundle from .preflight import Preflight, preflight_run from .requires import Declaring, declare, declaring_for -from .resolve import BuiltImage, BundleImage, Reference, ReferencedImage, resolve_bundle, without_played_agents +from .resolve import (BuiltImage, BundleImage, Reference, ReferencedImage, ResolvedEntry, resolve_bundle, + without_played_agents) logger = logging.getLogger(__name__) @@ -404,6 +405,16 @@ def _composing(sandbox: str | None, emulate: bool = True) -> Composing: return Composing(build_sandbox_provider(sandbox) if sandbox else None, emulate) +def planned_tasks(root: Path | str, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), + agent: str | None = None) -> tuple[ResolvedEntry, ...]: + """The tasks ``run_bundle`` with these arguments would run, as resolved, without checking or writing anything. + Namespace routing must be on, as for ``plan_bundle``.""" + plan = plan_bundle(resolve_bundle(parse_bundle(Path(root))), tasks=tasks, evals=evals, installed=installed_bundles, + agent=agent) + selected = set(_to_run(plan)) + return tuple(write.source for write in plan.writes if write.kind is BundleKind.TASK and write.source.entry in selected) + + def _to_run(plan: Plan) -> tuple[BundleEntry, ...]: """The tasks selected on their own and every task a selected eval names, each once, in the plan's order.""" wanted = {entry.entry.id for entry in plan.tasks} diff --git a/src/agent_env/training/harness.py b/src/agent_env/training/harness.py index f3bb9b36..ac49c174 100644 --- a/src/agent_env/training/harness.py +++ b/src/agent_env/training/harness.py @@ -25,13 +25,15 @@ raise ImportError("agent_env.training.harness requires pip install 'agentenv-framework[harness]'") from exc from agent_env.bundle.plan import solver_steps -from agent_env.bundle.run import prepare_bundle, solver_params +from agent_env.bundle.resolve import build_step +from agent_env.bundle.run import planned_tasks, prepare_bundle, solver_params from agent_env.openenv.app import reward_of from agent_env.store.routing import namespace_routing from agent_env.task.task import Task, recorded_scores from agent_env.task.teardown import teardown_run from agent_env.task_step.context import TaskStepContext, sandbox_overrides from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.verifiers.rubrics_verifier import RubricsVerifierTaskStep from agent_env.training._openenv import on_session_loop logger = logging.getLogger(__name__) @@ -138,22 +140,38 @@ async def _run(self) -> int: def _trained(bundle: str | Path, tasks: Sequence[str], evals: Sequence[str], sandbox: str, agent: str | None) -> list[tuple[Task, dict[str, dict[str, Any]]]]: - """The bundle's tasks prepared as ``agent-env run`` prepares them, each with its solver's deploys and their step - params, ``agent``'s if it swaps the solver; a task whose solver no ``deploy_agent`` step deploys is left out.""" + """The selected tasks that can train their solver, prepared as ``agent-env run`` prepares them, each with its + solver's deploys and their step params, ``agent``'s if it swaps the solver. The others are left out before they're + checked, so what they alone need can't stop the rest.""" with namespace_routing(): - prepared = prepare_bundle(bundle, tasks=tasks, evals=evals, sandbox=sandbox, agent=agent, solver_endpoint=True) + kept = [] + for entry in planned_tasks(bundle, tasks=tasks, evals=evals, agent=agent): + if why := _untrainable(entry.config): + logger.warning("%s is left out: %s", entry.entry.name, why) + else: + kept.append(entry.entry.name) + if not kept: + raise ValueError("no task of the bundle has a solver to train: an agent a step prompts, deployed by a " + "deploy_agent step, that isn't also its judge") + prepared = prepare_bundle(bundle, tasks=kept, sandbox=sandbox, agent=agent, solver_endpoint=True) swapped = solver_params(prepared.plan, prepared.materialization) - trained = [] - for entry, task in prepared.tasks: - solver = solver_steps([step.to_dict() for step in task.steps]) - if solver and all(step.type == DeployAgentTaskStep.type for step in task.steps if step.id in solver): - trained.append((task, {step: swapped.get(entry.id, {}).get(step, {}) for step in solver})) - else: - logger.warning("%s is left out: training needs the agent a step prompts deployed by a deploy_agent step", - task.id) - if not trained: - raise ValueError("no task of the bundle deploys, with a deploy_agent step, an agent a step prompts to train") - return trained + return [(task, {step: swapped.get(entry.id, {}).get(step, {}) + for step in solver_steps([step.to_dict() for step in task.steps])}) + for entry, task in prepared.tasks] + + +def _untrainable(config: list[dict[str, Any]]) -> str | None: + """Why a task's rollouts can't train its solver alone, or None.""" + solver = set(solver_steps(config)) + steps = [build_step(step) for step in config] + deploys = [step for step in steps if step.id in solver] + if not deploys or not all(isinstance(step, DeployAgentTaskStep) for step in deploys): + return "training needs the agent a step prompts deployed by a deploy_agent step" + agents = {step.agent_name for step in deploys} + if any(isinstance(step, RubricsVerifierTaskStep) and step.use_agent_judge and step.agent_name in agents + for step in steps): + return "its rubric judge is its solver, whose grading would be trained on" + return None def _reward(task: Task, reward: str | Callable[[dict[str, float]], float] | None diff --git a/tst/unit/bundle/run_test.py b/tst/unit/bundle/run_test.py index 31d3572d..c51efa6f 100644 --- a/tst/unit/bundle/run_test.py +++ b/tst/unit/bundle/run_test.py @@ -359,9 +359,8 @@ def broken(line): def test_a_run_says_before_it_builds_an_image_and_names_the_image_apart_from_its_agent(bundle_dir, docker, monkeypatch): """As it does for an image a deploy on this machine runs, which preflight names.""" - monkeypatch.setattr(run_module, "preflight_run", - lambda plan, tasks, sandbox, played: Preflight(frozenset(), (), - frozenset({f"{ROOT}/solver__agent_image"}))) + monkeypatch.setattr(run_module, "preflight_run", lambda plan, tasks, sandbox, played, solver_endpoint: Preflight( + frozenset(), (), frozenset({f"{ROOT}/solver__agent_image"}))) layout(bundle_dir, BUILT_AGENT) lines = [] diff --git a/tst/unit/training/harness_test.py b/tst/unit/training/harness_test.py index 4e2869ff..c4aa8331 100644 --- a/tst/unit/training/harness_test.py +++ b/tst/unit/training/harness_test.py @@ -32,7 +32,9 @@ DEPLOY = {"id": "agent", "type": "deploy_agent", "env_ids": [], "a2a_agent_id": "theirs"} ASK = {"id": "ask", "type": "prompt_agent", "prompt_id": "ask", "prompt": "What is 6 * 7?", "depends_on": ["agent"]} CHECK = {"id": "check", "type": "harness_test_check", "verifier_id": "answer", "depends_on": ["ask"]} -GIVEN = {"litellm_base_url": "http://models:4000", "env_vars": {"LITELLM_API_KEY": "sk-1"}} # its own endpoint +JUDGE = {"id": "judge", "type": "rubrics_verifier", "verifier_id": "judge", "prompt_id": "ask", + "agent_name": DeployAgentTaskStep.DEFAULT_AGENT_NAME, "criteria": [{"id": "c1", "description": "says 42"}], + "depends_on": ["ask"]} asked: list[dict] = [] @@ -138,7 +140,8 @@ def bundle(tmp_path, monkeypatch, local_stores): # noqa: F811 A2AAgent.put(id=agent_id, docker_image_artifact=image) return layout(tmp_path / "quiz", {"tasks/sum.json": json.dumps([DEPLOY, ASK, CHECK]), "tasks/broken.json": json.dumps([{**DEPLOY, "a2a_agent_id": "crashes"}, ASK]), - "tasks/unprompted.json": json.dumps([{**DEPLOY, **GIVEN}])}) + "tasks/unprompted.json": json.dumps([DEPLOY]), + "tasks/self-judged.json": json.dumps([DEPLOY, ASK, JUDGE])}) def test_the_solver_alone_calls_the_policy_through_the_proxy_and_the_run_is_captured_and_scored(engine, bundle, caplog): @@ -159,6 +162,7 @@ def test_the_solver_alone_calls_the_policy_through_the_proxy_and_the_run_is_capt [turn] = trace.turns assert sorted(row["task_id"].rsplit("/", 1)[1] for row in rows) == ["broken", "sum"] assert "unprompted is left out: training needs the agent a step prompts deployed by a deploy_agent step" in caplog.text + assert "self-judged is left out: its rubric judge is its solver, whose grading would be trained on" in caplog.text assert (call["model"], call["answer"]) == ("trl-policy-v3", "It is 42.") assert call["overrides"]["step_params"] == {"agent": {"litellm_base_url": f"http://127.0.0.1:{factory._proxy.port}", "litellm_api_key": call["key"]}} @@ -203,7 +207,7 @@ def test_a_prompt_naming_no_task_of_the_factorys_is_refused(engine, bundle): def test_a_bundle_with_no_solver_to_train_is_refused(engine, bundle): - with pytest.raises(ValueError, match="no task of the bundle deploys"): + with pytest.raises(ValueError, match="no task of the bundle has a solver to train"): AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, tasks=["unprompted"]) From 1b6ddad99f89bcd845a824859e94d074200ac215 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 11:40:22 -0800 Subject: [PATCH 05/10] test(training): name the self-judged task's judge agent Co-Authored-By: Claude Opus 5.5 --- tst/unit/training/harness_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tst/unit/training/harness_test.py b/tst/unit/training/harness_test.py index c4aa8331..c56ff9de 100644 --- a/tst/unit/training/harness_test.py +++ b/tst/unit/training/harness_test.py @@ -32,7 +32,7 @@ DEPLOY = {"id": "agent", "type": "deploy_agent", "env_ids": [], "a2a_agent_id": "theirs"} ASK = {"id": "ask", "type": "prompt_agent", "prompt_id": "ask", "prompt": "What is 6 * 7?", "depends_on": ["agent"]} CHECK = {"id": "check", "type": "harness_test_check", "verifier_id": "answer", "depends_on": ["ask"]} -JUDGE = {"id": "judge", "type": "rubrics_verifier", "verifier_id": "judge", "prompt_id": "ask", +JUDGE = {"id": "judge", "type": "rubrics_verifier", "verifier_id": "judge", "prompt_id": "ask", "judge_a2a_agent_id": "theirs", "agent_name": DeployAgentTaskStep.DEFAULT_AGENT_NAME, "criteria": [{"id": "c1", "description": "says 42"}], "depends_on": ["ask"]} asked: list[dict] = [] From 98d3a3f9f0a9bd8b1548230d33d352c53ef1bcaa Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 11:43:13 -0800 Subject: [PATCH 06/10] test(training): the shipped harnesses' smoke tasks through the capture proxy, on Docker opencode and claude-code run their smoke tasks against a stand-in vLLM that has them run the shell command and then answer, with no endpoint configured: the rollout scores 1.0 and every captured turn trains. The stand-in moves to tst/util/fake_vllm.py, which the unit test shares. Co-Authored-By: Claude Opus 5.5 --- tst/integration/training/harness_test.py | 65 ++++++++++++++++++ tst/unit/training/harness_test.py | 39 ++--------- tst/util/fake_vllm.py | 86 ++++++++++++++++++++++++ 3 files changed, 155 insertions(+), 35 deletions(-) create mode 100644 tst/integration/training/harness_test.py create mode 100644 tst/util/fake_vllm.py diff --git a/tst/integration/training/harness_test.py b/tst/integration/training/harness_test.py new file mode 100644 index 00000000..b22c36bc --- /dev/null +++ b/tst/integration/training/harness_test.py @@ -0,0 +1,65 @@ +"""``AgentEnvSessionFactory`` on the shipped agents' smoke tasks, as TRL's harness rollouts run them: the harness, in +local Docker, calls the model through OpenEnv's capture proxy, in front of a stand-in vLLM that has it run the task's +shell command and then answer, so the rollout is captured as trainable turns and scored by the task's verifier, with no +model endpoint configured. Needs a Docker daemon and network access for the agents' images.""" + +from pathlib import Path + +import pytest + +import agent_env +from agent_env.artifact.store import reset_artifact_store +from agent_env.config import configure, reset_config +from agent_env.training.harness import AgentEnvSessionFactory, prompt_rows +from tst.util.fake_vllm import fake_vllm, shell_then_says + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] + +MODEL = "Qwen/Qwen3-4B" + + +@pytest.fixture +def state(monkeypatch, tmp_path): + """Local stores under this test's folder and no model endpoint; no rollout may leave a sandbox work folder.""" + sandboxes = tmp_path / "sandboxes" + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state")) + monkeypatch.setenv("AGENT_ENV_DOCUMENT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_OBJECT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(sandboxes)) + for name in ("LITELLM_API_KEY", "LITELLM_BASE_URL"): + monkeypatch.delenv(name, raising=False) + configure() + reset_artifact_store() + try: + yield tmp_path + assert not sandboxes.exists() or not any(sandboxes.iterdir()), "a rollout left a sandbox work folder" + finally: + reset_artifact_store() + reset_config() + + +@pytest.mark.parametrize("agent", ["opencode", "claude-code"]) +def test_a_harness_runs_its_smoke_task_through_the_capture_proxy_into_trainable_turns(state, agent): + shipped = Path(agent_env.__file__).parent / "agents" / "tasks" / f"{agent}-smoke.json" + (state / "smoke" / "tasks").mkdir(parents=True) + (state / "smoke" / "tasks" / shipped.name).write_text(shipped.read_text()) + + with fake_vllm(MODEL, shell_then_says("echo $((6 * 7))", "42")) as engine: + [row] = prompt_rows(state / "smoke") + factory = AgentEnvSessionFactory(state / "smoke", llm_url=engine.url, model=MODEL, + sampling={"temperature": 1.0}) + session = factory.create(row["prompt"]) + try: + assert session.wait_for_completion(900) == 0 + trace = session.fetch_training_trace() + verified = session.verify([]) + finally: + session.close() + + assert verified.env_reward == 1.0 + assert any(message.get("role") == "tool" for body in engine.requests for message in body["messages"]) + assert len(trace.turns) >= 2 + for turn in trace.turns: + completion = len(turn.completion_token_ids) + assert completion and len(turn.per_token_logps) == completion + assert turn.loss_mask == [0] * len(turn.prompt_token_ids) + [1] * completion diff --git a/tst/unit/training/harness_test.py b/tst/unit/training/harness_test.py index c56ff9de..bad94472 100644 --- a/tst/unit/training/harness_test.py +++ b/tst/unit/training/harness_test.py @@ -8,10 +8,6 @@ import httpx import pytest -from starlette.applications import Starlette -from starlette.responses import JSONResponse -from starlette.routing import Route - from agent_env.a2a_agent import A2AAgent from agent_env.artifact.artifacts.docker_image import DockerImageArtifact from agent_env.artifact.store import get_artifact_store @@ -22,10 +18,10 @@ from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep from agent_env.training import harness -from agent_env.training._openenv import Server from agent_env.training.harness import AgentEnvSessionFactory, prompt_rows from tst.unit.bundle._support import layout from tst.unit.conftest import local_stores # noqa: F401 +from tst.util.fake_vllm import fake_vllm, says, token_ids MODEL = "Qwen/Qwen3-4B" MODEL_ENDPOINT_MISSING = preflight.model_endpoint_missing # the probe, which the unit tier's conftest stubs @@ -88,35 +84,8 @@ async def execute(self, context): @pytest.fixture def engine(): - """A stand-in for vLLM run with --return-tokens-as-token-ids and processed logprobs.""" - requests, served = [], [MODEL] - - async def models(request): - return JSONResponse({"data": [{"id": model, "object": "model"} for model in served]}) - - async def completions(request): - body = await request.json() - requests.append(body) - temperature = body.get("temperature") or 1.0 - prompt = [ord(char) % 500 + 1 for message in body["messages"] for char in str(message.get("content") or "")] - reply = "It is 42." - ids = [ord(char) % 500 + 600 for char in reply] - content = [{"token": f"token_id:{token}", "logprob": -0.5 / temperature, "bytes": None, - "top_logprobs": [{"token": f"token_id:{token}", "logprob": -0.5 / temperature}, - {"token": f"token_id:{token + 1}", "logprob": -3.0 / temperature}]} - for token in ids] - return JSONResponse({ - "id": "chatcmpl-1", "object": "chat.completion", "model": body["model"], "prompt_token_ids": prompt, - "choices": [{"index": 0, "message": {"role": "assistant", "content": reply}, "finish_reason": "stop", - "token_ids": ids, "logprobs": {"content": content}}], - "usage": {"prompt_tokens": len(prompt), "completion_tokens": len(ids), "total_tokens": len(prompt) + len(ids)}, - }) - - server = Server(Starlette(routes=[Route("/v1/models", models), Route("/v1/chat/completions", completions, - methods=["POST"])])) - server.requests, server.served = requests, served - yield server - server.stop() + with fake_vllm(MODEL, says("It is 42.")) as server: + yield server @pytest.fixture(autouse=True) @@ -168,7 +137,7 @@ def test_the_solver_alone_calls_the_policy_through_the_proxy_and_the_run_is_capt "litellm_api_key": call["key"]}} assert not {"litellm_base_url", "litellm_api_key"} & call["overrides"].keys() assert (verified.env_reward, verified.metrics) == (1.0, {"score/check": 1.0}) - assert turn.completion_token_ids == [ord(char) % 500 + 600 for char in "It is 42."] + assert turn.completion_token_ids == token_ids({"role": "assistant", "content": "It is 42."}) assert turn.loss_mask == [0] * len(turn.prompt_token_ids) + [1] * len(turn.completion_token_ids) assert engine.requests[-1]["model"] == "trl-policy-v3" assert factory._proxy.registry.get(call["key"]) is None diff --git a/tst/util/fake_vllm.py b/tst/util/fake_vllm.py new file mode 100644 index 00000000..a5c9da15 --- /dev/null +++ b/tst/util/fake_vllm.py @@ -0,0 +1,86 @@ +"""A stand-in for vLLM served with ``--return-tokens-as-token-ids`` and processed logprobs, as OpenEnv's capture proxy +needs its engine: it lists the models it serves and answers chat completions with token ids, a logprob per completion +token and the top two at each, scaled by the temperature as processed logprobs are. ``reply`` makes each answer's +message from the request body.""" + +import json +import uuid +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from typing import Any + +from starlette.applications import Starlette +from starlette.responses import JSONResponse +from starlette.routing import Route + +from agent_env.training._openenv import Server + +Reply = Callable[[dict[str, Any]], dict[str, Any]] + + +def says(content: str) -> Reply: + return lambda body: {"role": "assistant", "content": content} + + +def shell_then_says(command: str, content: str) -> Reply: + """Calls the request's shell tool (``bash`` in any case) with ``command`` until a tool result follows the last user + message, then says ``content``; a request offering no shell gets ``content`` at once, as a harness's side calls do.""" + + def reply(body: dict[str, Any]) -> dict[str, Any]: + messages = body["messages"] + shell = next((tool["function"]["name"] for tool in body.get("tools") or [] + if tool["function"]["name"].lower() == "bash"), None) + last_user = max((i for i, message in enumerate(messages) if message.get("role") == "user"), default=-1) + if shell is None or any(message.get("role") == "tool" for message in messages[last_user + 1:]): + return {"role": "assistant", "content": content} + arguments = {"command": command, "description": "Run the command"} + return {"role": "assistant", "content": None, "tool_calls": [ + {"id": f"call_{uuid.uuid4().hex[:8]}", "type": "function", + "function": {"name": shell, "arguments": json.dumps(arguments)}}]} + + return reply + + +def token_ids(message: dict[str, Any]) -> list[int]: + """A message's tokens: one per character of its role, content and tool calls.""" + text = f"<{message.get('role')}>{json.dumps(message.get('content'))}{json.dumps(message.get('tool_calls') or [])}" + return [ord(char) % 50000 + 1 for char in text] + + +@contextmanager +def fake_vllm(model: str, reply: Reply) -> Iterator[Server]: + """The engine, served on loopback; ``.requests`` are the completion bodies it got, and ``.served`` the models it + lists, which a test appends to as vLLM lists a LoRA adapter once it's loaded.""" + requests: list[dict[str, Any]] = [] + served = [model] + + async def models(request): + return JSONResponse({"data": [{"id": name, "object": "model"} for name in served]}) + + async def completions(request): + body = await request.json() + requests.append(body) + temperature = body.get("temperature") or 1.0 + message = reply(body) + prompt = [token for each in body["messages"] for token in token_ids(each)] + completion = token_ids(message) + content = [{"token": f"token_id:{token}", "logprob": -0.5 / temperature, "bytes": None, + "top_logprobs": [{"token": f"token_id:{token}", "logprob": -0.5 / temperature}, + {"token": f"token_id:{token + 1}", "logprob": -3.0 / temperature}]} + for token in completion] + return JSONResponse({ + "id": f"chatcmpl-{uuid.uuid4().hex[:8]}", "object": "chat.completion", "model": body["model"], + "prompt_token_ids": prompt, + "choices": [{"index": 0, "message": message, "finish_reason": "tool_calls" if message.get("tool_calls") + else "stop", "token_ids": completion, "logprobs": {"content": content}}], + "usage": {"prompt_tokens": len(prompt), "completion_tokens": len(completion), + "total_tokens": len(prompt) + len(completion)}, + }) + + server = Server(Starlette(routes=[Route("/v1/models", models), + Route("/v1/chat/completions", completions, methods=["POST"])])) + server.requests, server.served = requests, served + try: + yield server + finally: + server.stop() From 46d6d9b56abaa01ec3d2df9c7e7492dbc861032b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 15:58:21 -0800 Subject: [PATCH 07/10] fix(training): stop and tear down a rollout in flight on close, and run only on local Docker TRL closes the sessions in flight when it stops, but close only dropped the captures: the run went on in the session loop's thread and its sandboxes outlived the trainer. close now cancels a run still in its task and returns once the run is torn down; a run already tearing down is waited for, not cut short. A rollout's timeout bounds its task alone, so its teardown also completes before TimeoutError reaches TRL. The factory and prompt_rows no longer take a sandbox: a remote sandbox can't reach the capture proxy on the trainer's loopback, so every deploy runs on local Docker. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/harness.py | 90 +++++++++++++++++++++---------- tst/unit/training/harness_test.py | 59 +++++++++++++++++++- 2 files changed, 120 insertions(+), 29 deletions(-) diff --git a/src/agent_env/training/harness.py b/src/agent_env/training/harness.py index ac49c174..5ee76a5f 100644 --- a/src/agent_env/training/harness.py +++ b/src/agent_env/training/harness.py @@ -6,10 +6,13 @@ import logging import re +import asyncio import secrets import socket +import threading import uuid from collections.abc import Callable, Sequence +from concurrent.futures import CancelledError from pathlib import Path from typing import Any @@ -39,33 +42,35 @@ logger = logging.getLogger(__name__) #: How a row's prompt names its task, since TRL hands the factory only the prompt. _NAMED = re.compile(r"^AgentEnv task (\d+): (.+)$") +#: Where every deploy runs: local Docker alone reaches the capture proxy on this machine; a remote sandbox needs a tunnel. +_SANDBOX = "local" -def prompt_rows(bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), sandbox: str = "local", +def prompt_rows(bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), agent: str | None = None) -> list[dict[str, Any]]: """The trainer's dataset, a row per task ``AgentEnvSessionFactory`` plays with the same arguments: a system line naming the task, which is how the factory finds it, then the task's first prompt, for TRL's logs. The model never sees this prompt: the task's own agent builds its own.""" return [{"prompt": [{"role": "system", "content": f"AgentEnv task {index}: {task.id}"}, {"role": "user", "content": _first_prompt(task)}], "task_id": task.id} - for index, (task, _) in enumerate(_trained(bundle, tasks, evals, sandbox, agent))] + for index, (task, _) in enumerate(_trained(bundle, tasks, evals, agent))] class AgentEnvSessionFactory(ResourceSessionFactory): """OpenEnv's session factory for TRL's ``HarnessRolloutWorker``, built in its rollout process. It prepares the - bundle, checks that ``llm_url``'s engine returns token ids, and serves OpenEnv's capture proxy in front of it, - where the task's solver reaches it: the agent a step prompts, which it deploys, its other agents and judges keeping - the configured endpoint. TRL sets ``model`` to the policy it publishes; ``reward`` is as ``build_app``'s, and + bundle to run on local Docker, checks that ``llm_url``'s engine returns token ids, and serves OpenEnv's capture + proxy in front of it, where the task's solver reaches it: the agent a step prompts, which it deploys, its other + agents and judges keeping the configured endpoint. TRL sets ``model`` to the policy it publishes; ``reward`` is as ``build_app``'s, and ``agent`` the agent the solver deploys, as ``agent-env run --agent``'s.""" def __init__(self, bundle: str | Path, *, llm_url: str, model: str, sampling: dict[str, Any] | None = None, - tasks: Sequence[str] = (), evals: Sequence[str] = (), sandbox: str = "local", + tasks: Sequence[str] = (), evals: Sequence[str] = (), reward: str | Callable[[dict[str, float]], float] | None = None, agent: str | None = None, api_key: str | None = None, max_model_calls: int = 0): self.model = model self._llm_url, self._api_key, self._sampling = llm_url, api_key, sampling - self._sandbox, self._max_model_calls = sandbox, max_model_calls - self._trained = _trained(bundle, tasks, evals, sandbox, agent) + self._max_model_calls = max_model_calls + self._trained = _trained(bundle, tasks, evals, agent) self._rewards = [_reward(task, reward) for task, _ in self._trained] require_llm(llm_url, model, api_key=api_key, require_tokens=True) self._proxy = CaptureServer(llm_url=llm_url, model=model, port=_free_port(), api_key=api_key, @@ -80,7 +85,7 @@ def create(self, task: Any, seed: int | None = None, episode_id: str | None = No capture_level="tokens", purpose="train", sampling=self._sampling, max_model_calls=self._max_model_calls, task=played.id, episode=episode_id or uuid.uuid4().hex) endpoint = {"litellm_base_url": f"http://127.0.0.1:{self._proxy.port}", "litellm_api_key": capture.session_id} - overrides = {**sandbox_overrides(self._sandbox), + overrides = {**sandbox_overrides(_SANDBOX), "step_params": {step: {**params, **endpoint} for step, params in solver.items()}} return AgentEnvSession(played, overrides, self._rewards[index], self._proxy.registry, capture, self.model) @@ -93,6 +98,8 @@ def __init__(self, task: Task, overrides: dict[str, Any], reward: Callable[[dict self._task, self._overrides, self._reward, self._model = task, overrides, reward, model self._registry, self._capture = registry, capture self._scores: dict[str, float] | None = None + self._lock = threading.Lock() + self._closed, self._running, self._done = False, None, None def initial_messages(self) -> list[dict[str, Any]]: return [{"role": "user", "content": _first_prompt(self._task)}] @@ -104,8 +111,16 @@ def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: return ToolResult(error="the task's own agent calls its tools in its sandbox; none are offered here") def wait_for_completion(self, timeout_s: float | None = None) -> int: - """Runs the task and tears it down: 0 when it ran to the end, 1 when it failed.""" - return on_session_loop(self._run(), timeout_s) + """Runs the task and tears it down: 0 when it ran to the end, 1 when it failed or ``close`` stopped it. Raises + TimeoutError, once the run is torn down, when the task runs past ``timeout_s``.""" + with self._lock: + if self._closed: + return 1 + self._done = threading.Event() + try: + return on_session_loop(self._run(timeout_s), None) + except CancelledError: + return 1 def fetch_training_trace(self) -> TrainingTrace: document = export_session(self._capture, include_messages=True, capture_level="tokens") @@ -122,23 +137,44 @@ def verify(self, transcript: list[dict[str, Any]], final_state: Any | None = Non for step, score in (self._scores or {}).items()}) def close(self) -> None: + """Stops a run still going, as TRL does to the rollouts in flight when it stops, and returns once the run is torn + down, so no sandbox outlives the trainer; then drops the session's captures.""" + with self._lock: + self._closed = True + running, done = self._running, self._done + if running is not None: + running.get_loop().call_soon_threadsafe(running.cancel) + if done is not None: + done.wait() self._registry.delete(self._capture.session_id) - async def _run(self) -> int: - context = TaskStepContext(metadata={"task_id": self._task.id, "user_overrides": self._overrides}) - with namespace_routing(): - try: - await self._task.run(agent_model=self._model, context=context) - self._scores = recorded_scores(self._task.steps, context) - return 0 - except Exception as e: - logger.warning("task %s failed, so its rollout is left unscored: %s", self._task.id, e) - return 1 - finally: - await teardown_run(context) - - -def _trained(bundle: str | Path, tasks: Sequence[str], evals: Sequence[str], sandbox: str, + async def _run(self, timeout: float | None) -> int: + try: + with self._lock: + if self._closed: + return 1 + self._running = asyncio.current_task() + context = TaskStepContext(metadata={"task_id": self._task.id, "user_overrides": self._overrides}) + with namespace_routing(): + try: + async with asyncio.timeout(timeout): + await self._task.run(agent_model=self._model, context=context) + self._scores = recorded_scores(self._task.steps, context) + return 0 + except TimeoutError: + raise + except Exception as e: + logger.warning("task %s failed, so its rollout is left unscored: %s", self._task.id, e) + return 1 + finally: + with self._lock: + self._running = None # a close from here on waits for the teardown rather than cutting it short + await teardown_run(context) + finally: + self._done.set() + + +def _trained(bundle: str | Path, tasks: Sequence[str], evals: Sequence[str], agent: str | None) -> list[tuple[Task, dict[str, dict[str, Any]]]]: """The selected tasks that can train their solver, prepared as ``agent-env run`` prepares them, each with its solver's deploys and their step params, ``agent``'s if it swaps the solver. The others are left out before they're @@ -153,7 +189,7 @@ def _trained(bundle: str | Path, tasks: Sequence[str], evals: Sequence[str], san if not kept: raise ValueError("no task of the bundle has a solver to train: an agent a step prompts, deployed by a " "deploy_agent step, that isn't also its judge") - prepared = prepare_bundle(bundle, tasks=kept, sandbox=sandbox, agent=agent, solver_endpoint=True) + prepared = prepare_bundle(bundle, tasks=kept, sandbox=_SANDBOX, agent=agent, solver_endpoint=True) swapped = solver_params(prepared.plan, prepared.materialization) return [(task, {step: swapped.get(entry.id, {}).get(step, {}) for step in solver_steps([step.to_dict() for step in task.steps])}) diff --git a/tst/unit/training/harness_test.py b/tst/unit/training/harness_test.py index bad94472..5f1d90c1 100644 --- a/tst/unit/training/harness_test.py +++ b/tst/unit/training/harness_test.py @@ -3,8 +3,11 @@ token ids and processed logprobs, with no endpoint configured; its calls are captured and the task's verifier scores the run.""" +import asyncio import json import logging +import threading +from concurrent.futures import ThreadPoolExecutor import httpx import pytest @@ -32,6 +35,7 @@ "agent_name": DeployAgentTaskStep.DEFAULT_AGENT_NAME, "criteria": [{"id": "c1", "description": "says 42"}], "depends_on": ["ask"]} asked: list[dict] = [] +hanging = threading.Event() class _Deploys(DeployAgentTaskStep): @@ -40,6 +44,9 @@ class _Deploys(DeployAgentTaskStep): async def execute(self, context): if self.a2a_agent_id == "crashes": raise RuntimeError("the agent crashed") + if self.a2a_agent_id == "hangs": + hanging.set() + await asyncio.sleep(3600) context.metadata["endpoint"] = self._model_endpoint(context) return context @@ -101,7 +108,8 @@ def bundle(tmp_path, monkeypatch, local_stores): # noqa: F811 monkeypatch.delenv("LITELLM_BASE_URL", raising=False) monkeypatch.delenv("LITELLM_API_KEY", raising=False) asked.clear() - for agent_id in ("theirs", "crashes", "mine"): + hanging.clear() + for agent_id in ("theirs", "crashes", "hangs", "mine"): with namespace_routing(): image = get_artifact_store().put_document(DockerImageArtifact( id=f"{agent_id}-image", description="d", image_name="ghcr.io/acme/agent:1", @@ -109,6 +117,7 @@ def bundle(tmp_path, monkeypatch, local_stores): # noqa: F811 A2AAgent.put(id=agent_id, docker_image_artifact=image) return layout(tmp_path / "quiz", {"tasks/sum.json": json.dumps([DEPLOY, ASK, CHECK]), "tasks/broken.json": json.dumps([{**DEPLOY, "a2a_agent_id": "crashes"}, ASK]), + "tasks/hanging.json": json.dumps([{**DEPLOY, "a2a_agent_id": "hangs"}, ASK]), "tasks/unprompted.json": json.dumps([DEPLOY]), "tasks/self-judged.json": json.dumps([DEPLOY, ASK, JUDGE])}) @@ -129,7 +138,7 @@ def test_the_solver_alone_calls_the_policy_through_the_proxy_and_the_run_is_capt [call] = asked [turn] = trace.turns - assert sorted(row["task_id"].rsplit("/", 1)[1] for row in rows) == ["broken", "sum"] + assert sorted(row["task_id"].rsplit("/", 1)[1] for row in rows) == ["broken", "hanging", "sum"] assert "unprompted is left out: training needs the agent a step prompts deployed by a deploy_agent step" in caplog.text assert "self-judged is left out: its rubric judge is its solver, whose grading would be trained on" in caplog.text assert (call["model"], call["answer"]) == ("trl-policy-v3", "It is 42.") @@ -166,6 +175,52 @@ def test_a_run_that_fails_is_unscored_rather_than_zero(engine, bundle, caplog): assert "failed, so its rollout is left unscored: " in caplog.text +@pytest.fixture +def torn(monkeypatch): + """The task ids of the runs torn down, each recorded once its teardown is done.""" + done, teardown = [], harness.teardown_run + + async def tears_down(context): + report = await teardown(context) + done.append(context.metadata["task_id"]) + return report + + monkeypatch.setattr(harness, "teardown_run", tears_down) + return done + + +def test_close_stops_a_rollout_in_flight_and_returns_once_it_is_torn_down(engine, bundle, torn): + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL) + session = factory.create(next(row["prompt"] for row in prompt_rows(bundle) if row["task_id"].endswith("hanging"))) + + with ThreadPoolExecutor(1) as pool: + waiting = pool.submit(session.wait_for_completion) + assert hanging.wait(30) + session.close() + assert [task_id.rsplit("/", 1)[1] for task_id in torn] == ["hanging"] + assert waiting.result(30) == 1 + assert session.verify([]).env_reward is None + + +def test_a_rollout_past_its_timeout_is_torn_down_before_the_timeout_is_raised(engine, bundle, torn): + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL) + session = factory.create(next(row["prompt"] for row in prompt_rows(bundle) if row["task_id"].endswith("hanging"))) + + with pytest.raises(TimeoutError): + session.wait_for_completion(0.5) + assert [task_id.rsplit("/", 1)[1] for task_id in torn] == ["hanging"] + + +def test_a_session_closed_before_it_runs_never_runs(engine, bundle, torn): + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL) + session = factory.create(next(row["prompt"] for row in prompt_rows(bundle) if row["task_id"].endswith("sum"))) + + session.close() + + assert session.wait_for_completion() == 1 + assert (asked, torn) == ([], []) + + def test_a_prompt_naming_no_task_of_the_factorys_is_refused(engine, bundle): factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL) From dc6d9bebc461cf2204832bf832b52d6e7793e5a1 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 16:04:37 -0800 Subject: [PATCH 08/10] fix(training): decide a stop's cancel on the session loop, and stop only once close read the running task under the lock but cancelled it later, so a run that left its task in between could have its teardown cancelled. The check and the cancel now happen together on the session loop, where the run clears its task before tearing down, and a second close sends no cancel. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/harness.py | 12 +++++++++--- tst/unit/training/harness_test.py | 2 ++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/agent_env/training/harness.py b/src/agent_env/training/harness.py index 5ee76a5f..e75f2db4 100644 --- a/src/agent_env/training/harness.py +++ b/src/agent_env/training/harness.py @@ -140,14 +140,20 @@ def close(self) -> None: """Stops a run still going, as TRL does to the rollouts in flight when it stops, and returns once the run is torn down, so no sandbox outlives the trainer; then drops the session's captures.""" with self._lock: - self._closed = True + stopping, self._closed = not self._closed, True running, done = self._running, self._done - if running is not None: - running.get_loop().call_soon_threadsafe(running.cancel) + if stopping and running is not None: + running.get_loop().call_soon_threadsafe(self._stop) if done is not None: done.wait() self._registry.delete(self._capture.session_id) + def _stop(self) -> None: + """Cancels the run if it's still in its task, decided on the session loop, where the run leaves its task for + its teardown, so the cancel can't land in the teardown.""" + if self._running is not None: + self._running.cancel() + async def _run(self, timeout: float | None) -> int: try: with self._lock: diff --git a/tst/unit/training/harness_test.py b/tst/unit/training/harness_test.py index 5f1d90c1..6369b20f 100644 --- a/tst/unit/training/harness_test.py +++ b/tst/unit/training/harness_test.py @@ -199,6 +199,8 @@ def test_close_stops_a_rollout_in_flight_and_returns_once_it_is_torn_down(engine session.close() assert [task_id.rsplit("/", 1)[1] for task_id in torn] == ["hanging"] assert waiting.result(30) == 1 + session.close() # as TRL closes it again once the rollout returns + assert len(torn) == 1 assert session.verify([]).env_reward is None From 9c54d8839467ea81dd77e1ce50096598b5773e14 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 16:10:49 -0800 Subject: [PATCH 09/10] refactor(training): simpler stop, reward rules in the task layer, fewer arguments; CI on the latest OpenEnv - A session runs Task.run as an asyncio task of its own, which a stop or the timeout cancels, and tears down in its own coroutine, which nothing cancels, so a stop can't cut a teardown short by construction. - reward_of moves beside recorded_scores in agent_env.task.task, so the harness no longer imports the OpenEnv server module. - The factory drops api_key and max_model_calls, which nothing used. - AgentEnvSession takes keyword arguments; imports and long lines tidied. - harness.yml runs the harness tier on the locked and the latest OpenEnv, whose capture internals the harness reads. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/harness.yml | 55 ++++++++++++++++++ src/agent_env/openenv/app.py | 30 +--------- src/agent_env/task/task.py | 29 ++++++++++ src/agent_env/training/harness.py | 83 +++++++++++++--------------- tst/unit/openenv/environment_test.py | 2 +- tst/unit/training/_support.py | 3 +- tst/unit/training/harness_test.py | 9 +-- tst/util/fake_vllm.py | 3 +- 8 files changed, 134 insertions(+), 80 deletions(-) create mode 100644 .github/workflows/harness.yml diff --git a/.github/workflows/harness.yml b/.github/workflows/harness.yml new file mode 100644 index 00000000..b4bc1000 --- /dev/null +++ b/.github/workflows/harness.yml @@ -0,0 +1,55 @@ +name: Harness contract + +# TRL's harness rollouts on a bundle's own agents (agent_env.training.harness, tst/unit/training/harness_test.py) +# through OpenEnv's capture proxy, on the uv.lock pins and on the latest OpenEnv. The harness reads the proxy's internals +# (its session registry, export and training trace), which OpenEnv can change in any minor release. It runs on pull +# requests that touch the harness or its tier, weekly, and on demand; it isn't a required check. + +on: + pull_request: + paths: + - "src/agent_env/training/**" + - "tst/unit/training/**" + - "tst/util/fake_vllm.py" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/harness.yml" + schedule: + - cron: "53 7 * * 1" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: harness-${{ github.event_name == 'pull_request' && github.ref || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +env: + PIP_DISABLE_PIP_VERSION_CHECK: "1" + +jobs: + capture: + name: capture (openenv ${{ matrix.openenv }}) + runs-on: ubuntu-latest + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + openenv: ["locked", "latest"] + steps: + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + - name: Install from uv.lock, with the latest OpenEnv for latest + run: | + python -m pip install --quiet -r .github/ci-requirements.txt + uv sync --locked --extra dev --python "$(command -v python)" + if [ "${{ matrix.openenv }}" = latest ]; then uv pip install --python .venv/bin/python --upgrade openenv; fi + .venv/bin/python -c "import importlib.metadata as m; print('openenv', m.version('openenv'))" + - name: Harness tier + run: .venv/bin/python -m pytest tst/unit/training/harness_test.py -rA --timeout=600 --junitxml=junit-harness.xml + - name: No skips + if: always() + run: .venv/bin/python .github/scripts/check_skip_policy.py junit-harness.xml diff --git a/src/agent_env/openenv/app.py b/src/agent_env/openenv/app.py index 8aac5934..01b6ec59 100644 --- a/src/agent_env/openenv/app.py +++ b/src/agent_env/openenv/app.py @@ -21,6 +21,7 @@ from agent_env.bundle.installed import locate_bundle from agent_env.store.routing import namespace_routing from agent_env.task import EpisodeError, EpisodeSpec, Task +from agent_env.task.task import reward_of from .environment import AgentEnvEnvironment, Served, ServedTask @@ -156,32 +157,3 @@ def _served(splits: dict[str, list[tuple[str, Task]]], raise BundleError([f"no task can be served as an OpenEnv env: {why}" for why in refused.values()] or ["no task can be served as an OpenEnv env: the bundle has none"]) return served - - -def reward_of(task: Task, reward: str | Callable[[dict[str, float]], float] | None - ) -> tuple[Callable[[dict[str, float]], float | None] | None, str | None]: - """``task``'s reward as a function of its scores, or why ``reward`` can't make one.""" - if callable(reward): - return reward, None - scoring = [step.id for step in task.steps if getattr(step, "verifier_id", None)] - if reward is None: - if len(scoring) > 1: - return None, (f"task {task.id!r} records a score at each of {', '.join(scoring)}; name the one that is " - "its reward with reward=") - return (_score_at(scoring[0]) if scoring else _unscored), None - if reward not in scoring: - return None, f"task {task.id!r} records no score at step {reward!r}; it records them at {scoring}" - return _score_at(reward), None - - -def _score_at(step: str) -> Callable[[dict[str, float]], float]: - def score(scores: dict[str, float]) -> float: - if step not in scores: - raise LookupError(f"the run recorded no score at step {step!r}") - return scores[step] - - return score - - -def _unscored(scores: dict[str, float]) -> None: - return None diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 6840a58d..4e8932ab 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -135,6 +135,35 @@ def recorded_scores(steps: list["TaskStep"], context: TaskStepContext) -> dict[s return scores +def reward_of(task: "Task", reward: str | Callable[[dict[str, float]], float] | None + ) -> tuple[Callable[[dict[str, float]], float | None] | None, str | None]: + """``task``'s reward as a function of its scores, or why ``reward`` can't make one.""" + if callable(reward): + return reward, None + scoring = [step.id for step in task.steps if getattr(step, "verifier_id", None)] + if reward is None: + if len(scoring) > 1: + return None, (f"task {task.id!r} records a score at each of {', '.join(scoring)}; name the one that is " + "its reward with reward=") + return (_score_at(scoring[0]) if scoring else _unscored), None + if reward not in scoring: + return None, f"task {task.id!r} records no score at step {reward!r}; it records them at {scoring}" + return _score_at(reward), None + + +def _score_at(step: str) -> Callable[[dict[str, float]], float]: + def score(scores: dict[str, float]) -> float: + if step not in scores: + raise LookupError(f"the run recorded no score at step {step!r}") + return scores[step] + + return score + + +def _unscored(scores: dict[str, float]) -> None: + return None + + def _dependency_ids(steps: list["TaskStep"]) -> dict[str, set[str]]: """Direct dependency ids per step; depends_on None = all prior (scheduler semantics).""" deps: dict[str, set[str]] = {} diff --git a/src/agent_env/training/harness.py b/src/agent_env/training/harness.py index e75f2db4..81dcec76 100644 --- a/src/agent_env/training/harness.py +++ b/src/agent_env/training/harness.py @@ -1,18 +1,18 @@ """Train a bundle's own agents with TRL's harness rollouts (``trl.experimental.async_grpo.HarnessRolloutWorker`` with ``harness_adapter=None``): each rollout runs a task as ``agent-env run`` does, its agents' model calls going through -OpenEnv's capture proxy to the trainer's vLLM, which records the tokens TRL trains on, and the task's verifiers score it.""" +OpenEnv's capture proxy to the trainer's vLLM, which records the tokens TRL trains on, and the task's verifiers score +it.""" from __future__ import annotations +import asyncio import logging import re -import asyncio import secrets import socket import threading import uuid from collections.abc import Callable, Sequence -from concurrent.futures import CancelledError from pathlib import Path from typing import Any @@ -30,9 +30,8 @@ from agent_env.bundle.plan import solver_steps from agent_env.bundle.resolve import build_step from agent_env.bundle.run import planned_tasks, prepare_bundle, solver_params -from agent_env.openenv.app import reward_of from agent_env.store.routing import namespace_routing -from agent_env.task.task import Task, recorded_scores +from agent_env.task.task import Task, recorded_scores, reward_of from agent_env.task.teardown import teardown_run from agent_env.task_step.context import TaskStepContext, sandbox_overrides from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep @@ -42,7 +41,8 @@ logger = logging.getLogger(__name__) #: How a row's prompt names its task, since TRL hands the factory only the prompt. _NAMED = re.compile(r"^AgentEnv task (\d+): (.+)$") -#: Where every deploy runs: local Docker alone reaches the capture proxy on this machine; a remote sandbox needs a tunnel. +#: Where every deploy runs: local Docker alone reaches the capture proxy on this machine, a remote sandbox needing a +#: tunnel. _SANDBOX = "local" @@ -60,20 +60,18 @@ class AgentEnvSessionFactory(ResourceSessionFactory): """OpenEnv's session factory for TRL's ``HarnessRolloutWorker``, built in its rollout process. It prepares the bundle to run on local Docker, checks that ``llm_url``'s engine returns token ids, and serves OpenEnv's capture proxy in front of it, where the task's solver reaches it: the agent a step prompts, which it deploys, its other - agents and judges keeping the configured endpoint. TRL sets ``model`` to the policy it publishes; ``reward`` is as ``build_app``'s, and - ``agent`` the agent the solver deploys, as ``agent-env run --agent``'s.""" + agents and judges keeping the configured endpoint. TRL sets ``model`` to the policy it publishes; ``reward`` is as + ``build_app``'s, and ``agent`` the agent the solver deploys, as ``agent-env run --agent``'s.""" def __init__(self, bundle: str | Path, *, llm_url: str, model: str, sampling: dict[str, Any] | None = None, tasks: Sequence[str] = (), evals: Sequence[str] = (), - reward: str | Callable[[dict[str, float]], float] | None = None, agent: str | None = None, - api_key: str | None = None, max_model_calls: int = 0): + reward: str | Callable[[dict[str, float]], float] | None = None, agent: str | None = None): self.model = model - self._llm_url, self._api_key, self._sampling = llm_url, api_key, sampling - self._max_model_calls = max_model_calls + self._llm_url, self._sampling = llm_url, sampling self._trained = _trained(bundle, tasks, evals, agent) self._rewards = [_reward(task, reward) for task, _ in self._trained] - require_llm(llm_url, model, api_key=api_key, require_tokens=True) - self._proxy = CaptureServer(llm_url=llm_url, model=model, port=_free_port(), api_key=api_key, + require_llm(llm_url, model, require_tokens=True) + self._proxy = CaptureServer(llm_url=llm_url, model=model, port=_free_port(), admin_key=secrets.token_urlsafe(32)) self._proxy.start() @@ -81,25 +79,27 @@ def create(self, task: Any, seed: int | None = None, episode_id: str | None = No index = _index_named(task, [each.id for each, _ in self._trained]) played, solver = self._trained[index] capture = self._proxy.registry.create( - upstream=Upstream(llm_url=self._llm_url, model=self.model, api_key=self._api_key), - capture_level="tokens", purpose="train", sampling=self._sampling, max_model_calls=self._max_model_calls, - task=played.id, episode=episode_id or uuid.uuid4().hex) + upstream=Upstream(llm_url=self._llm_url, model=self.model), capture_level="tokens", purpose="train", + sampling=self._sampling, task=played.id, episode=episode_id or uuid.uuid4().hex) endpoint = {"litellm_base_url": f"http://127.0.0.1:{self._proxy.port}", "litellm_api_key": capture.session_id} overrides = {**sandbox_overrides(_SANDBOX), "step_params": {step: {**params, **endpoint} for step, params in solver.items()}} - return AgentEnvSession(played, overrides, self._rewards[index], self._proxy.registry, capture, self.model) + return AgentEnvSession(task=played, overrides=overrides, reward=self._rewards[index], + registry=self._proxy.registry, capture=capture, model=self.model) class AgentEnvSession(ResourceSession): """One run of a task with its own agents, their model calls captured under this session's key.""" - def __init__(self, task: Task, overrides: dict[str, Any], reward: Callable[[dict[str, float]], float | None], + def __init__(self, *, task: Task, overrides: dict[str, Any], reward: Callable[[dict[str, float]], float | None], registry: Any, capture: Any, model: str): self._task, self._overrides, self._reward, self._model = task, overrides, reward, model self._registry, self._capture = registry, capture self._scores: dict[str, float] | None = None self._lock = threading.Lock() - self._closed, self._running, self._done = False, None, None + self._closed = False + self._task_run: asyncio.Task | None = None # Task.run, the one part of a run a stop cancels + self._done: threading.Event | None = None # set once the run is torn down def initial_messages(self) -> list[dict[str, Any]]: return [{"role": "user", "content": _first_prompt(self._task)}] @@ -117,10 +117,7 @@ def wait_for_completion(self, timeout_s: float | None = None) -> int: if self._closed: return 1 self._done = threading.Event() - try: - return on_session_loop(self._run(timeout_s), None) - except CancelledError: - return 1 + return on_session_loop(self._run(timeout_s), None) def fetch_training_trace(self) -> TrainingTrace: document = export_session(self._capture, include_messages=True, capture_level="tokens") @@ -137,44 +134,42 @@ def verify(self, transcript: list[dict[str, Any]], final_state: Any | None = Non for step, score in (self._scores or {}).items()}) def close(self) -> None: - """Stops a run still going, as TRL does to the rollouts in flight when it stops, and returns once the run is torn - down, so no sandbox outlives the trainer; then drops the session's captures.""" + """Stops a run still going, as TRL does to the rollouts in flight when it stops, and returns once the run is + torn down, so no sandbox outlives the trainer; then drops the session's captures.""" with self._lock: stopping, self._closed = not self._closed, True - running, done = self._running, self._done - if stopping and running is not None: - running.get_loop().call_soon_threadsafe(self._stop) + task_run, done = self._task_run, self._done + if stopping and task_run is not None: + task_run.get_loop().call_soon_threadsafe(task_run.cancel) if done is not None: done.wait() self._registry.delete(self._capture.session_id) - def _stop(self) -> None: - """Cancels the run if it's still in its task, decided on the session loop, where the run leaves its task for - its teardown, so the cancel can't land in the teardown.""" - if self._running is not None: - self._running.cancel() - async def _run(self, timeout: float | None) -> int: + """``Task.run`` in an asyncio task of its own, which a stop or the timeout cancels, then the teardown in this + coroutine, which nothing cancels, so a stop can't cut it short.""" + context = TaskStepContext(metadata={"task_id": self._task.id, "user_overrides": self._overrides}) try: - with self._lock: - if self._closed: - return 1 - self._running = asyncio.current_task() - context = TaskStepContext(metadata={"task_id": self._task.id, "user_overrides": self._overrides}) with namespace_routing(): + with self._lock: + if self._closed: + return 1 + task_run = self._task_run = asyncio.ensure_future( + self._task.run(agent_model=self._model, context=context)) try: - async with asyncio.timeout(timeout): - await self._task.run(agent_model=self._model, context=context) + await asyncio.wait_for(task_run, timeout) self._scores = recorded_scores(self._task.steps, context) return 0 except TimeoutError: raise + except asyncio.CancelledError: + if not task_run.cancelled(): + raise + return 1 except Exception as e: logger.warning("task %s failed, so its rollout is left unscored: %s", self._task.id, e) return 1 finally: - with self._lock: - self._running = None # a close from here on waits for the teardown rather than cutting it short await teardown_run(context) finally: self._done.set() diff --git a/tst/unit/openenv/environment_test.py b/tst/unit/openenv/environment_test.py index 046800be..cdf3811f 100644 --- a/tst/unit/openenv/environment_test.py +++ b/tst/unit/openenv/environment_test.py @@ -12,7 +12,7 @@ from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction, ToolErrorType import agent_env.task.episode.run as episode_run -from agent_env.openenv.app import _score_at +from agent_env.task.task import _score_at from agent_env.openenv.environment import ENDED, NOT_RESET, TIMEOUT_METADATA, AgentEnvEnvironment, Served, ServedTask from agent_env.task import EpisodeError, EpisodeSpec, Task from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN, _gate, _until diff --git a/tst/unit/training/_support.py b/tst/unit/training/_support.py index 157db065..c5ca86d8 100644 --- a/tst/unit/training/_support.py +++ b/tst/unit/training/_support.py @@ -3,7 +3,8 @@ import pytest -from agent_env.openenv.app import _app, _score_at +from agent_env.openenv.app import _app +from agent_env.task.task import _score_at from agent_env.openenv.environment import Served, ServedTask from agent_env.task import EpisodeSpec, Task from agent_env.training._openenv import Server diff --git a/tst/unit/training/harness_test.py b/tst/unit/training/harness_test.py index 6369b20f..defddc23 100644 --- a/tst/unit/training/harness_test.py +++ b/tst/unit/training/harness_test.py @@ -31,9 +31,9 @@ DEPLOY = {"id": "agent", "type": "deploy_agent", "env_ids": [], "a2a_agent_id": "theirs"} ASK = {"id": "ask", "type": "prompt_agent", "prompt_id": "ask", "prompt": "What is 6 * 7?", "depends_on": ["agent"]} CHECK = {"id": "check", "type": "harness_test_check", "verifier_id": "answer", "depends_on": ["ask"]} -JUDGE = {"id": "judge", "type": "rubrics_verifier", "verifier_id": "judge", "prompt_id": "ask", "judge_a2a_agent_id": "theirs", - "agent_name": DeployAgentTaskStep.DEFAULT_AGENT_NAME, "criteria": [{"id": "c1", "description": "says 42"}], - "depends_on": ["ask"]} +JUDGE = {"id": "judge", "type": "rubrics_verifier", "verifier_id": "judge", "prompt_id": "ask", + "judge_a2a_agent_id": "theirs", "agent_name": DeployAgentTaskStep.DEFAULT_AGENT_NAME, + "criteria": [{"id": "c1", "description": "says 42"}], "depends_on": ["ask"]} asked: list[dict] = [] hanging = threading.Event() @@ -139,7 +139,8 @@ def test_the_solver_alone_calls_the_policy_through_the_proxy_and_the_run_is_capt [call] = asked [turn] = trace.turns assert sorted(row["task_id"].rsplit("/", 1)[1] for row in rows) == ["broken", "hanging", "sum"] - assert "unprompted is left out: training needs the agent a step prompts deployed by a deploy_agent step" in caplog.text + assert "unprompted is left out: training needs the agent a step prompts deployed by a deploy_agent step" in ( + caplog.text) assert "self-judged is left out: its rubric judge is its solver, whose grading would be trained on" in caplog.text assert (call["model"], call["answer"]) == ("trl-policy-v3", "It is 42.") assert call["overrides"]["step_params"] == {"agent": {"litellm_base_url": f"http://127.0.0.1:{factory._proxy.port}", diff --git a/tst/util/fake_vllm.py b/tst/util/fake_vllm.py index a5c9da15..9a57e132 100644 --- a/tst/util/fake_vllm.py +++ b/tst/util/fake_vllm.py @@ -24,7 +24,8 @@ def says(content: str) -> Reply: def shell_then_says(command: str, content: str) -> Reply: """Calls the request's shell tool (``bash`` in any case) with ``command`` until a tool result follows the last user - message, then says ``content``; a request offering no shell gets ``content`` at once, as a harness's side calls do.""" + message, then says ``content``; a request offering no shell gets ``content`` at once, as a harness's side calls + do.""" def reply(body: dict[str, Any]) -> dict[str, Any]: messages = body["messages"] From 99e027ddda374faa5c351c5bbc0418926562f7f8 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 16:59:45 -0800 Subject: [PATCH 10/10] refactor(openenv): keep the reward rules in agent_env.openenv.reward reward_of and its helpers move from agent_env.task.task to a module of their own beside build_app, whose rules they are; the harness imports them there, and task.py is as on main. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/openenv/app.py | 2 +- src/agent_env/openenv/reward.py | 36 ++++++++++++++++++++++++++++ src/agent_env/task/task.py | 29 ---------------------- src/agent_env/training/harness.py | 3 ++- tst/unit/openenv/environment_test.py | 2 +- tst/unit/training/_support.py | 2 +- 6 files changed, 41 insertions(+), 33 deletions(-) create mode 100644 src/agent_env/openenv/reward.py diff --git a/src/agent_env/openenv/app.py b/src/agent_env/openenv/app.py index 01b6ec59..11acaa64 100644 --- a/src/agent_env/openenv/app.py +++ b/src/agent_env/openenv/app.py @@ -21,9 +21,9 @@ from agent_env.bundle.installed import locate_bundle from agent_env.store.routing import namespace_routing from agent_env.task import EpisodeError, EpisodeSpec, Task -from agent_env.task.task import reward_of from .environment import AgentEnvEnvironment, Served, ServedTask +from .reward import reward_of logger = logging.getLogger(__name__) diff --git a/src/agent_env/openenv/reward.py b/src/agent_env/openenv/reward.py new file mode 100644 index 00000000..08778e77 --- /dev/null +++ b/src/agent_env/openenv/reward.py @@ -0,0 +1,36 @@ +"""How a task's recorded scores make its reward: the rules ``build_app`` applies, which harness training shares.""" + +from __future__ import annotations + +from collections.abc import Callable + +from agent_env.task import Task + + +def reward_of(task: Task, reward: str | Callable[[dict[str, float]], float] | None + ) -> tuple[Callable[[dict[str, float]], float | None] | None, str | None]: + """``task``'s reward as a function of its scores, or why ``reward`` can't make one.""" + if callable(reward): + return reward, None + scoring = [step.id for step in task.steps if getattr(step, "verifier_id", None)] + if reward is None: + if len(scoring) > 1: + return None, (f"task {task.id!r} records a score at each of {', '.join(scoring)}; name the one that is " + "its reward with reward=") + return (_score_at(scoring[0]) if scoring else _unscored), None + if reward not in scoring: + return None, f"task {task.id!r} records no score at step {reward!r}; it records them at {scoring}" + return _score_at(reward), None + + +def _score_at(step: str) -> Callable[[dict[str, float]], float]: + def score(scores: dict[str, float]) -> float: + if step not in scores: + raise LookupError(f"the run recorded no score at step {step!r}") + return scores[step] + + return score + + +def _unscored(scores: dict[str, float]) -> None: + return None diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 4e8932ab..6840a58d 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -135,35 +135,6 @@ def recorded_scores(steps: list["TaskStep"], context: TaskStepContext) -> dict[s return scores -def reward_of(task: "Task", reward: str | Callable[[dict[str, float]], float] | None - ) -> tuple[Callable[[dict[str, float]], float | None] | None, str | None]: - """``task``'s reward as a function of its scores, or why ``reward`` can't make one.""" - if callable(reward): - return reward, None - scoring = [step.id for step in task.steps if getattr(step, "verifier_id", None)] - if reward is None: - if len(scoring) > 1: - return None, (f"task {task.id!r} records a score at each of {', '.join(scoring)}; name the one that is " - "its reward with reward=") - return (_score_at(scoring[0]) if scoring else _unscored), None - if reward not in scoring: - return None, f"task {task.id!r} records no score at step {reward!r}; it records them at {scoring}" - return _score_at(reward), None - - -def _score_at(step: str) -> Callable[[dict[str, float]], float]: - def score(scores: dict[str, float]) -> float: - if step not in scores: - raise LookupError(f"the run recorded no score at step {step!r}") - return scores[step] - - return score - - -def _unscored(scores: dict[str, float]) -> None: - return None - - def _dependency_ids(steps: list["TaskStep"]) -> dict[str, set[str]]: """Direct dependency ids per step; depends_on None = all prior (scheduler semantics).""" deps: dict[str, set[str]] = {} diff --git a/src/agent_env/training/harness.py b/src/agent_env/training/harness.py index 81dcec76..807eddde 100644 --- a/src/agent_env/training/harness.py +++ b/src/agent_env/training/harness.py @@ -31,7 +31,8 @@ from agent_env.bundle.resolve import build_step from agent_env.bundle.run import planned_tasks, prepare_bundle, solver_params from agent_env.store.routing import namespace_routing -from agent_env.task.task import Task, recorded_scores, reward_of +from agent_env.openenv.reward import reward_of +from agent_env.task.task import Task, recorded_scores from agent_env.task.teardown import teardown_run from agent_env.task_step.context import TaskStepContext, sandbox_overrides from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep diff --git a/tst/unit/openenv/environment_test.py b/tst/unit/openenv/environment_test.py index cdf3811f..7f9ac0f3 100644 --- a/tst/unit/openenv/environment_test.py +++ b/tst/unit/openenv/environment_test.py @@ -12,7 +12,7 @@ from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction, ToolErrorType import agent_env.task.episode.run as episode_run -from agent_env.task.task import _score_at +from agent_env.openenv.reward import _score_at from agent_env.openenv.environment import ENDED, NOT_RESET, TIMEOUT_METADATA, AgentEnvEnvironment, Served, ServedTask from agent_env.task import EpisodeError, EpisodeSpec, Task from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN, _gate, _until diff --git a/tst/unit/training/_support.py b/tst/unit/training/_support.py index c5ca86d8..60237369 100644 --- a/tst/unit/training/_support.py +++ b/tst/unit/training/_support.py @@ -4,7 +4,7 @@ import pytest from agent_env.openenv.app import _app -from agent_env.task.task import _score_at +from agent_env.openenv.reward import _score_at from agent_env.openenv.environment import Served, ServedTask from agent_env.task import EpisodeSpec, Task from agent_env.training._openenv import Server