From 21a74d6ad104f36bf0f49a1241276cc9878b739b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 07:22:16 -0700 Subject: [PATCH 01/13] feat(training): train with verl through an AgentEnv agent loop Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 8 +- src/agent_env/training/verl/__init__.py | 4 + src/agent_env/training/verl/__main__.py | 26 ++++++ src/agent_env/training/verl/loop.py | 106 ++++++++++++++++++++++++ src/agent_env/training/verl/prepare.py | 58 +++++++++++++ uv.lock | 13 ++- 6 files changed, 212 insertions(+), 3 deletions(-) create mode 100644 src/agent_env/training/verl/__init__.py create mode 100644 src/agent_env/training/verl/__main__.py create mode 100644 src/agent_env/training/verl/loop.py create mode 100644 src/agent_env/training/verl/prepare.py diff --git a/pyproject.toml b/pyproject.toml index 38cb8b14..5eaaa1ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,12 @@ nemo-gym = [ "nemo-gym>=0.7.0rc0; python_full_version >= '3.13.14'", "omegaconf>=2.3; python_full_version >= '3.13.14'", ] +# agent_env.training.verl: verl's trainers on a bundle served as an OpenEnv env; installed into verl's own environment, +# whose verl, vLLM and torch are the trainer's to pin. pyarrow writes the parquet rows verl reads. +verl = [ + "agentenv-framework[openenv]", + "pyarrow>=19", +] # The AWS store backends (S3, Secrets Manager, DynamoDB, ECR), selected by impl pointer in config.toml. aws = [ "boto3>=1.35.0", @@ -117,7 +123,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,tensorlake,tenki,trl,vercel]", # exercise portable optional providers and stores in CI + "agentenv-framework[aws,boat,boxd,coreweave,daytona,gcp,openenv,sail,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/training/verl/__init__.py b/src/agent_env/training/verl/__init__.py new file mode 100644 index 00000000..1b8c25dd --- /dev/null +++ b/src/agent_env/training/verl/__init__.py @@ -0,0 +1,4 @@ +"""verl's trainers on a bundle served as an OpenEnv env (``agent_env.openenv.build_app``, or the package ``openenv +import`` writes): ``python -m agent_env.training.verl URL OUT`` writes the server's tasks as verl's parquet rows with the +config of the agent loop that plays them, and ``agent_env.training.verl.loop.AgentEnvAgentLoop``, which runs in verl's +agent loop workers, plays each row's task as one episode on the server and takes its reward.""" diff --git a/src/agent_env/training/verl/__main__.py b/src/agent_env/training/verl/__main__.py new file mode 100644 index 00000000..23ed4d8d --- /dev/null +++ b/src/agent_env/training/verl/__main__.py @@ -0,0 +1,26 @@ +"""``python -m agent_env.training.verl URL OUT``: writes a bundle's tasks as verl's RL dataset and agent loop config.""" + +import click + +from agent_env.training.verl.prepare import write_dataset + + +@click.command() +@click.argument("url") +@click.argument("out", type=click.Path(file_okay=False)) +@click.option("--tool-timeout", type=float, default=60.0, show_default=True, + help="One tool call's deadline in seconds; the server ends a call that misses it, and the episode.") +@click.option("--reset-timeout", type=float, default=600.0, show_default=True, + help="How long an episode may take to deploy, and each turn to be ended, in seconds; also how long a " + "rollout waits for room on the server.") +@click.option("--attempts", type=int, default=3, show_default=True, + help="How many times a rollout the env fails is played, each time from the start on a new episode.") +def main(url: str, out: str, tool_timeout: float, reset_timeout: float, attempts: int) -> None: + """Writes the tasks of the OpenEnv server at URL as OUT/.parquet, a row per task, and OUT/agent_loop.yaml, + the agent loop that plays them on the server.""" + folder = write_dataset(url, out, tool_timeout=tool_timeout, reset_timeout=reset_timeout, attempts=attempts) + click.echo(f"wrote {', '.join(sorted(path.name for path in folder.iterdir()))} to {folder}") + + +if __name__ == "__main__": + main() diff --git a/src/agent_env/training/verl/loop.py b/src/agent_env/training/verl/loop.py new file mode 100644 index 00000000..8966f49e --- /dev/null +++ b/src/agent_env/training/verl/loop.py @@ -0,0 +1,106 @@ +"""verl's agent loop for the rows ``python -m agent_env.training.verl`` writes. ``AgentEnvAgentLoop`` is verl's +``ToolAgentLoop``, which generates on verl's rollout server, parses the model's tool calls with +``multi_turn.format``'s parser and adds their answers to the context, with the row's tools in place of verl's own: each +call is a call to the tool on the row's episode, the model's last reply ends the turn it left open, and the episode's +reward is the rollout's ``reward_score``, so verl's reward function isn't called. Verified with verl 0.9.1.""" + +from __future__ import annotations + +import json +import logging +from typing import Any + +try: + from verl.experimental.agent_loop.agent_loop import AgentLoopOutput + from verl.experimental.agent_loop.tool_agent_loop import ToolAgentLoop + from verl.tools.function_tool import FunctionTool + from verl.tools.schemas import OpenAIFunctionToolSchema +except ModuleNotFoundError as exc: + if exc.name is None or exc.name.split(".")[0] != "verl": + raise + raise ImportError("agent_env.training.verl.loop runs in verl's agent loop workers: install " + "'agentenv-framework[verl]' into verl's environment") from exc + +from agent_env.training._openenv import Rollout, Session, Shared + +logger = logging.getLogger(__name__) + + +class _Session(Session): + offered_by = "its row lists" + + def __init__(self, shared: Shared): + super().__init__() + self._shared = shared + + +class AgentEnvAgentLoop(ToolAgentLoop): + """Plays a row's task as one episode on the OpenEnv server at ``url``. A rollout the env fails, a deploy, a call + past its deadline or a server that went away, is played again from the start on a new episode, a fresh sample of + the same policy, up to ``attempts`` times; one that fails every time is kept out of the loss, its response masked, + with the reward 0. ``agent_loop.yaml`` sets the arguments; verl passes the rest. + + Args: + url: The OpenEnv server's base URL. + tool_timeout: One tool call's deadline in seconds. The server ends a call that misses it, and the episode. + reset_timeout: How long an episode may take to deploy, and each turn to be ended, in seconds; also how long + a rollout waits for room on a server at its limit. + attempts: How many times a rollout the env fails is played. + """ + + def __init__(self, *args: Any, url: str, tool_timeout: float = 60.0, reset_timeout: float = 600.0, + attempts: int = 3, **kwargs: Any): + super().__init__(*args, **kwargs) + self._url, self._attempts = url.rstrip("/"), attempts + self._tool_timeout, self._reset_timeout = tool_timeout, reset_timeout + self._reply = "" + + async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoopOutput: + row = kwargs["extra_info"] + schemas = json.loads(row["tools"]) + shared = Shared(self._url, frozenset(schema["function"]["name"] for schema in schemas), self._tool_timeout, + self._reset_timeout, patience=self._reset_timeout) + for attempt in range(1, self._attempts + 1): + session = _Session(shared) + try: + output, rollout = await self._play(session, schemas, row, sampling_params, kwargs) + finally: + await session._close() + if rollout.failure is None and rollout.reward is not None: + output.reward_score = float(rollout.reward) + output.extra_fields["agentenv"] = {"attempts": attempt, "scores": rollout.scores} + return output + failure = rollout.failure or "the episode ended without a reward" + logger.warning("task %s of split %s, attempt %d of %d: %s", row["task_id"], row["split"], attempt, + self._attempts, failure) + output.response_mask = [0] * len(output.response_mask) + output.reward_score = 0.0 + output.extra_fields["agentenv"] = {"attempts": self._attempts, "failure": failure} + return output + + async def _play(self, session: _Session, schemas: list[dict[str, Any]], row: dict[str, Any], + sampling_params: dict[str, Any], kwargs: dict[str, Any]) -> tuple[AgentLoopOutput, Rollout]: + await session._start(row["split"], int(row["index"]), row["task_id"]) + self.tools = {schema["function"]["name"]: _tool(session, schema) for schema in schemas} + self.tool_schemas = schemas + self._reply = "" + output = await super().run(sampling_params, **kwargs) + rollout = session._rollout + await session._finish(rollout, self._reply) + return output, rollout + + def _build_assistant_message(self, content: str, agent_data: Any) -> dict[str, Any]: + self._reply = content or "" + return super()._build_assistant_message(content, agent_data) + + +def _tool(session: _Session, schema: dict[str, Any]) -> FunctionTool: + """The episode's tool as the function tool ``ToolAgentLoop`` calls; its schema is built unvalidated, since verl's + schema models keep no nested JSON Schema, and the chat template renders the row's own.""" + name = schema["function"]["name"] + + async def call(**arguments: Any) -> str: + return await session._step(name, arguments) + + return FunctionTool(name=name, fn=call, tool_schema=OpenAIFunctionToolSchema.model_construct(**schema), + is_async=True) diff --git a/src/agent_env/training/verl/prepare.py b/src/agent_env/training/verl/prepare.py new file mode 100644 index 00000000..6a92cb62 --- /dev/null +++ b/src/agent_env/training/verl/prepare.py @@ -0,0 +1,58 @@ +"""Writes a bundle's tasks as verl's RL dataset: a parquet file of rows per split, which ``data.train_files`` and +``data.val_files`` take, and ``agent_loop.yaml``, which ``actor_rollout_ref.rollout.agent.agent_loop_config_path`` +takes.""" + +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from typing import Any + +import pyarrow as pa +import pyarrow.parquet as pq +import yaml + +from agent_env.training._openenv import key, messages, probe, read_tasks, toolsets + +#: The ``agent_name`` every row plays under, and the name ``agent_loop.yaml`` registers the loop at. +AGENT = "agentenv" +LOOP = "agent_env.training.verl.loop.AgentEnvAgentLoop" + + +def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, reset_timeout: float = 600.0, + attempts: int = 3) -> Path: + """The tasks of the OpenEnv server at ``url`` as ``out/.parquet``, a row per task of each split, and + ``out/agent_loop.yaml``, which points the rows' agent loop at ``url`` with the deadlines and the ``attempts`` a + rollout the env fails is played; returns ``out``. A row's prompt is its task's system prompt and first prompt, and + its tools every tool its task's turns offer, found by playing one task of each toolset to its last turn, unscored, + so writing deploys those.""" + if attempts < 1: + raise ValueError(f"attempts must be at least 1, not {attempts}") + url = url.rstrip("/") + env_name, tasks = read_tasks(url, local="serve the bundle with agent_env.openenv.build_app first") + tools = {each: [_function(tool) for tool in asyncio.run(probe(url, toolset, reset_timeout))[1]] + for each, toolset in toolsets(env_name, tasks).items()} + out = Path(out) + out.mkdir(parents=True, exist_ok=True) + for split, described in tasks.items(): + pq.write_table(pa.Table.from_pylist([{ + "data_source": f"agentenv/{env_name}", + "prompt": messages(task["system_prompt"], task["prompt"]), + "agent_name": AGENT, + "extra_info": {"split": split, "index": task["index"], "task_id": task["id"], + "tools": json.dumps(tools[key(task)])}, + } for task in described]), out / f"{split}.parquet") + (out / "agent_loop.yaml").write_text(yaml.safe_dump([{ + "name": AGENT, "_target_": LOOP, "url": url, "tool_timeout": tool_timeout, "reset_timeout": reset_timeout, + "attempts": attempts, + }], sort_keys=False)) + return out + + +def _function(tool: dict[str, Any]) -> dict[str, Any]: + """An MCP tool as the function tool a chat template renders.""" + return {"type": "function", "function": { + "name": tool["name"], "description": tool.get("description") or "", + "parameters": tool.get("input_schema") or {"type": "object", "properties": {}}, + }} diff --git a/uv.lock b/uv.lock index 61337f0f..b3974b21 100644 --- a/uv.lock +++ b/uv.lock @@ -112,6 +112,8 @@ dev = [ { name = "openenv" }, { name = "pillow" }, { name = "psycopg2-binary" }, + { name = "pyarrow", version = "25.0.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13.14'" }, + { name = "pyarrow", version = "26.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13.14'" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-dependency" }, @@ -166,13 +168,19 @@ vercel = [ { name = "httpx2" }, { name = "vercel-sandbox" }, ] +verl = [ + { name = "openenv" }, + { name = "pyarrow", version = "25.0.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13.14'" }, + { name = "pyarrow", version = "26.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13.14'" }, +] [package.metadata] requires-dist = [ { name = "a2a-sdk", specifier = "==0.3.26" }, - { name = "agentenv-framework", extras = ["aws", "boat", "boxd", "coreweave", "daytona", "gcp", "openenv", "sail", "tensorlake", "tenki", "trl", "vercel"], marker = "extra == 'dev'" }, + { name = "agentenv-framework", extras = ["aws", "boat", "boxd", "coreweave", "daytona", "gcp", "openenv", "sail", "tensorlake", "tenki", "trl", "vercel", "verl"], marker = "extra == 'dev'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'nemo-gym'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'trl'" }, + { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'verl'" }, { name = "agentenv-framework-protocol", extras = ["computer"], editable = "packages/agentenv-protocol" }, { name = "anyio", marker = "extra == 'vercel'", specifier = ">=4,<5" }, { name = "av", marker = "extra == 'dev'", specifier = ">=17" }, @@ -207,6 +215,7 @@ requires-dist = [ { name = "packaging", specifier = ">=22" }, { name = "pillow", marker = "extra == 'dev'", specifier = ">=10.4" }, { name = "psycopg2-binary", marker = "extra == 'dev'", specifier = ">=2.9.0" }, + { name = "pyarrow", marker = "extra == 'verl'", specifier = ">=19" }, { name = "pydantic", specifier = ">=2.12" }, { name = "pymongo", specifier = ">=4.7" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=9.0" }, @@ -230,7 +239,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", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] +provides-extras = ["openenv", "trl", "nemo-gym", "verl", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] [[package]] name = "agentenv-framework-protocol" From fb8c6107df27f231907424e123d391cd077d2ff1 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 07:27:51 -0700 Subject: [PATCH 02/13] test(training): the verl writer on the fake bundle, and the agent loop on verl itself Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/_openenv.py | 8 ++ src/agent_env/training/nemo_gym/prepare.py | 7 +- src/agent_env/training/verl/prepare.py | 3 +- tst/unit/training/verl_test.py | 100 ++++++++++++++++++ tst/util/capabilities.py | 4 + tst/verl/__init__.py | 0 tst/verl/conftest.py | 7 ++ tst/verl/loop_test.py | 115 +++++++++++++++++++++ 8 files changed, 238 insertions(+), 6 deletions(-) create mode 100644 tst/unit/training/verl_test.py create mode 100644 tst/verl/__init__.py create mode 100644 tst/verl/conftest.py create mode 100644 tst/verl/loop_test.py diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index c37dd6ef..b24a9e39 100644 --- a/src/agent_env/training/_openenv.py +++ b/src/agent_env/training/_openenv.py @@ -43,6 +43,7 @@ httpx.TimeoutException, httpx.ReadError, httpx.RemoteProtocolError) _CAPACITY = re.compile(r"Server at capacity: \d+/(\d+)") +_FILE = re.compile(r"[A-Za-z0-9_-][A-Za-z0-9._-]*") logger = logging.getLogger(__name__) _waited: set[str] = set() @@ -365,6 +366,13 @@ def read_tasks(url: str, *, local: str) -> tuple[str, dict[str, list[dict[str, A return name, tasks +def check_split_names(tasks: Mapping[str, Any]) -> None: + """Refuses splits whose names can't name the data file a writer writes for each.""" + if unfit := [split for split in tasks if not _FILE.fullmatch(split)]: + raise ValueError(f"the server's splits {unfit} can't name a data file: a split's name must be letters, " + "digits, '.', '_' and '-', not starting with '.'") + + def toolsets(name: str, tasks: dict[str, list[dict[str, Any]]]) -> dict[str, NamedToolset]: """Each toolset the tasks have, by its key, with the first task that has it, named after the server when it's the only one and otherwise after its envs, and ``bash`` when its turns offer it.""" diff --git a/src/agent_env/training/nemo_gym/prepare.py b/src/agent_env/training/nemo_gym/prepare.py index c2a450d4..05655762 100644 --- a/src/agent_env/training/nemo_gym/prepare.py +++ b/src/agent_env/training/nemo_gym/prepare.py @@ -13,7 +13,7 @@ import yaml from agent_env.openenv import build_app -from agent_env.training._openenv import Server, key, messages, probe, read_tasks, toolsets +from agent_env.training._openenv import Server, check_split_names, key, messages, probe, read_tasks, toolsets #: The routes a resources server serves itself, which a tool's route can't take. GYM_ROUTES = frozenset({"seed_session", "verify", "close_session", "aggregate_metrics", "reverify_mode"}) @@ -45,7 +45,6 @@ REWARD_AND_COMPONENTS = ("reward_components makes the reward the sum of the task's scores, as NeMo RL requires, so " "it takes no reward: weight the scores in GDPO's reward_weights instead") _ROUTE = re.compile(r"[A-Za-z0-9._~-]+") -_FILE = re.compile(r"[A-Za-z0-9_-][A-Za-z0-9._-]*") def write_server(source: str | Path, root: str | Path, *, name: str | None = None, sandbox: str | None = None, @@ -70,9 +69,7 @@ def write_server(source: str | Path, root: str | Path, *, name: str | None = Non max_concurrent_envs=1)) try: env_name, tasks = read_tasks(url or served.url, local="name the bundle's folder instead of its server") - if unfit := [split for split in tasks if not _FILE.fullmatch(split)]: - raise ValueError(f"the server's splits {unfit} can't name a data file: a split's name must be letters, " - "digits, '.', '_' and '-', not starting with '.'") + check_split_names(tasks) grouped = toolsets(env_name, tasks) tools = {each: [_function(tool) for tool in asyncio.run(probe(url or served.url, toolset, reset_timeout))[1]] for each, toolset in grouped.items()} diff --git a/src/agent_env/training/verl/prepare.py b/src/agent_env/training/verl/prepare.py index 6a92cb62..91ccc876 100644 --- a/src/agent_env/training/verl/prepare.py +++ b/src/agent_env/training/verl/prepare.py @@ -13,7 +13,7 @@ import pyarrow.parquet as pq import yaml -from agent_env.training._openenv import key, messages, probe, read_tasks, toolsets +from agent_env.training._openenv import check_split_names, key, messages, probe, read_tasks, toolsets #: The ``agent_name`` every row plays under, and the name ``agent_loop.yaml`` registers the loop at. AGENT = "agentenv" @@ -31,6 +31,7 @@ def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, rese raise ValueError(f"attempts must be at least 1, not {attempts}") url = url.rstrip("/") env_name, tasks = read_tasks(url, local="serve the bundle with agent_env.openenv.build_app first") + check_split_names(tasks) tools = {each: [_function(tool) for tool in asyncio.run(probe(url, toolset, reset_timeout))[1]] for each, toolset in toolsets(env_name, tasks).items()} out = Path(out) diff --git a/tst/unit/training/verl_test.py b/tst/unit/training/verl_test.py new file mode 100644 index 00000000..a3e29b61 --- /dev/null +++ b/tst/unit/training/verl_test.py @@ -0,0 +1,100 @@ +"""``write_dataset`` on the OpenEnv server ``build_app`` makes, served on loopback over the Episode API's fake steps: +verl's parquet rows for each split and the config of the agent loop that plays them, and the command that writes both.""" + +import json + +import pyarrow.parquet as pq +import pytest +import yaml +from click.testing import CliRunner + +from agent_env.openenv.environment import END_TURN +from agent_env.training.verl.__main__ import main +from agent_env.training.verl.prepare import LOOP, write_dataset +from tst.unit.task.episode.run_test import ( + ALICE, + ITEMS, + ITEMS_MCP, + PANTRY, + PANTRY_MCP, + REPLIES, + TURN, + _tool_sessions, +) +from tst.unit.training.nemo_gym_test import PLAYED + +ADD_ITEM = {"type": "function", "function": {"name": "add_item", "description": "add_item.", + "parameters": {"type": "object"}}} +FUNCTION_END_TURN = {"type": "function", "function": {"name": "end_turn", "description": END_TURN.description, + "parameters": END_TURN.input_schema}} + + +def test_a_servers_tasks_become_a_row_per_task_of_each_split_and_the_agent_loop_that_plays_them(serve, tmp_path, + teardowns): + server = serve({"a": PLAYED, "b": [ITEMS, ALICE, {**TURN, "prompt": "Add a barrel."}, REPLIES]}, {"smoke": ["b"]}, + name="Dock Yard") + + out = write_dataset(f"{server.url}/", tmp_path / "data", tool_timeout=30.0, reset_timeout=90.0, attempts=2) + server.stop() + + assert out == tmp_path / "data" + assert yaml.safe_load((out / "agent_loop.yaml").read_text()) == [{ + "name": "agentenv", "_target_": LOOP, "url": server.url, "tool_timeout": 30.0, "reset_timeout": 90.0, + "attempts": 2, + }] + [row] = _rows(out / "smoke.parquet") + assert row == { + "data_source": "agentenv/Dock Yard", + "prompt": [{"role": "system", "content": "You keep the items."}, {"role": "user", "content": "Add a barrel."}], + "agent_name": "agentenv", + "extra_info": {"split": "smoke", "index": 0, "task_id": "b", "tools": [ADD_ITEM, FUNCTION_END_TURN]}, + } + assert [(row["extra_info"]["index"], row["extra_info"]["task_id"]) for row in _rows(out / "all.parquet")] == [ + (0, "a"), (1, "b")] + assert len(teardowns.finished) == 1 + + +def test_each_toolset_is_probed_once_and_a_rows_tools_are_its_tasks(serve, tmp_path, monkeypatch, teardowns): + _tool_sessions(monkeypatch, {ITEMS_MCP: ["add_item"], PANTRY_MCP: ["take_item"]}) + server = serve({"a": PLAYED, "p": [PANTRY, {**ALICE, "env_ids": ["pantry"]}, TURN, REPLIES], "c": PLAYED}) + + out = write_dataset(server.url, tmp_path) + server.stop() + + assert [[tool["function"]["name"] for tool in row["extra_info"]["tools"]] for row in _rows(out / "all.parquet")] == [ + ["add_item", "end_turn"], ["take_item", "end_turn"], ["add_item", "end_turn"]] + assert len(teardowns.finished) == 2 + + +def test_a_split_whose_name_isnt_a_file_name_is_refused_before_anything_is_written(serve, tmp_path): + server = serve({"a": PLAYED}, {"../escape": ["a"]}) + + with pytest.raises(ValueError, match=r"the server's splits \['\.\./escape'\] can't name a data file"): + write_dataset(server.url, tmp_path / "data") + + assert list(tmp_path.iterdir()) == [] + + +def test_a_rollout_is_played_at_least_once(tmp_path): + with pytest.raises(ValueError, match="attempts must be at least 1, not 0"): + write_dataset("http://127.0.0.1:1", tmp_path / "data", attempts=0) + assert list(tmp_path.iterdir()) == [] + + +def test_the_command_writes_the_rows_and_the_agent_loop_config(serve, tmp_path, teardowns): + server = serve({"a": PLAYED}) + + result = CliRunner().invoke(main, [server.url, str(tmp_path / "data"), "--tool-timeout", "15", + "--reset-timeout", "120", "--attempts", "4"]) + server.stop() + + assert result.exit_code == 0, result.output + assert sorted(path.name for path in (tmp_path / "data").iterdir()) == ["agent_loop.yaml", "all.parquet"] + [config] = yaml.safe_load((tmp_path / "data" / "agent_loop.yaml").read_text()) + assert (config["tool_timeout"], config["reset_timeout"], config["attempts"]) == (15.0, 120.0, 4) + assert len(teardowns.finished) == 1 + + +def _rows(path): + return [{**row, "extra_info": {**row["extra_info"], "tools": json.loads(row["extra_info"]["tools"])}} + for row in pq.read_table(path).to_pylist()] diff --git a/tst/util/capabilities.py b/tst/util/capabilities.py index e0eeaff0..e964c064 100644 --- a/tst/util/capabilities.py +++ b/tst/util/capabilities.py @@ -45,6 +45,10 @@ #: framework mcp 1; only that tier's CI job installs it and puts ``vf-eval`` on PATH. VERIFIERS = "verifiers_installed" +#: The verl tier (``tst/verl``) runs ``AgentEnvAgentLoop`` on verl itself; verl and torch aren't in the dev extra, and +#: only that tier's CI job installs them. +VERL = "verl_installed" + #: ``collect_artifacts`` sizes a file with ``stat -c %s``, which only GNU stat takes (macOS ships BSD stat), #: so collecting off a local sandbox needs GNU stat first on PATH. GNU_STAT = "gnu_stat" diff --git a/tst/verl/__init__.py b/tst/verl/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tst/verl/conftest.py b/tst/verl/conftest.py new file mode 100644 index 00000000..58bc0751 --- /dev/null +++ b/tst/verl/conftest.py @@ -0,0 +1,7 @@ +"""The unit tier's bundle served on loopback over the Episode API's fake steps, which these tests play through verl's +agent loop.""" + +from tst.unit.conftest import local_stores # noqa: F401 +from tst.unit.openenv.environment_test import tools # noqa: F401 +from tst.unit.task.episode.run_test import teardowns # noqa: F401 +from tst.unit.training._support import serve # noqa: F401 diff --git a/tst/verl/loop_test.py b/tst/verl/loop_test.py new file mode 100644 index 00000000..4fd191c0 --- /dev/null +++ b/tst/verl/loop_test.py @@ -0,0 +1,115 @@ +"""``AgentEnvAgentLoop`` on verl itself: verl's ``ToolAgentLoop`` plays a row ``write_dataset`` wrote, with Qwen3's +tokenizer and chat template and a stand-in rollout server, on the unit tier's bundle served on loopback. A played row's +reward is its episode's, a rollout the env fails is played again on a new episode, and one it fails every time is kept +out of the loss. CI runs it on verl 0.9.1.""" + +import asyncio +import json + +import pyarrow.parquet as pq +import pytest +from transformers import AutoTokenizer + +from agent_env.training._openenv import Failed +from agent_env.training.verl.prepare import write_dataset +from tst.unit.training.nemo_gym_test import PLAYED +from tst.util.capabilities import VERL, missing_capability_reason + +verl_loop = pytest.importorskip("agent_env.training.verl.loop", reason=missing_capability_reason(VERL), + exc_type=ImportError) +omegaconf = pytest.importorskip("omegaconf", reason=missing_capability_reason(VERL)) +agent_loop = pytest.importorskip("verl.experimental.agent_loop.agent_loop", reason=missing_capability_reason(VERL)) +rl_dataset = pytest.importorskip("verl.utils.dataset.rl_dataset", reason=missing_capability_reason(VERL)) +replica = pytest.importorskip("verl.workers.rollout.replica", reason=missing_capability_reason(VERL)) + +TOKENIZER = "Qwen/Qwen3-0.6B" +CONFIG = omegaconf.OmegaConf.create({"actor_rollout_ref": {"rollout": { + "prompt_length": 2048, "response_length": 1024, "full_determinism": False, + "multi_turn": {"format": "hermes", "max_user_turns": None, "max_assistant_turns": 4, "max_parallel_calls": 1, + "max_tool_response_length": 2048, "tool_response_truncate_side": "middle"}, +}}}) +DATA = omegaconf.OmegaConf.create({"apply_chat_template_kwargs": {"enable_thinking": False}}) +ADD_A_CRATE = '\n{"name": "add_item", "arguments": {"item": "crate"}}\n<|im_end|>' + + +@pytest.fixture(scope="module") +def tokenizer(): + return AutoTokenizer.from_pretrained(TOKENIZER) + + +def test_a_played_rows_reward_is_its_episodes(serve, tmp_path, tokenizer, teardowns): + server = serve({"a": PLAYED}) + row = _row(write_dataset(server.url, tmp_path)) + model = _Model(tokenizer) + + output = asyncio.run(_play(server.url, row, model, tokenizer)) + server.stop() + + assert (output.reward_score, output.extra_fields["agentenv"]) == (1.0, {"attempts": 1, "scores": {"replies": 1.0}}) + assert model.calls == 2 + assert 0 in output.response_mask and 1 in output.response_mask + assert "Added a crate." in tokenizer.decode(output.response_ids) + assert len(teardowns.finished) == 2 + + +def test_a_rollout_the_env_fails_is_played_again_on_a_new_episode(serve, tmp_path, tokenizer, monkeypatch, teardowns): + server = serve({"a": PLAYED}) + row = _row(write_dataset(server.url, tmp_path)) + begin, began = verl_loop._Session._begin, [] + + async def failing_once(self, rollout, *args): + began.append(1) + if len(began) == 1: + raise Failed("the env's deploy failed") + return await begin(self, rollout, *args) + + monkeypatch.setattr(verl_loop._Session, "_begin", failing_once) + output = asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer)) + server.stop() + + assert (output.reward_score, output.extra_fields["agentenv"]["attempts"]) == (1.0, 2) + assert len(began) == 2 + + +def test_a_rollout_the_env_fails_every_time_is_kept_out_of_the_loss(serve, tmp_path, tokenizer, monkeypatch): + server = serve({"a": PLAYED}) + row = _row(write_dataset(server.url, tmp_path)) + + async def failing(self, rollout, *args): + raise Failed("the env's deploy failed") + + monkeypatch.setattr(verl_loop._Session, "_begin", failing) + output = asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer, attempts=2)) + server.stop() + + assert (output.reward_score, output.extra_fields["agentenv"]) == ( + 0.0, {"attempts": 2, "failure": "the env's deploy failed"}) + assert output.response_mask and set(output.response_mask) == {0} + + +class _Model: + """verl's rollout server for the loop: it adds a crate, then replies once the tool answered.""" + + def __init__(self, tokenizer): + self.tokenizer, self.calls = tokenizer, 0 + + async def generate(self, request_id, *, prompt_ids, sampling_params, **_): + self.calls += 1 + answered = "" in self.tokenizer.decode(prompt_ids) + reply = "Added a crate.<|im_end|>" if answered else ADD_A_CRATE + return replica.TokenOutput(token_ids=self.tokenizer.encode(reply, add_special_tokens=False)) + + +async def _play(url, row, model, tokenizer, attempts=3): + played = verl_loop.AgentEnvAgentLoop( + trainer_config=agent_loop.DictConfigWrap(CONFIG), server_manager=model, tokenizer=tokenizer, processor=None, + dataset_cls=rl_dataset.RLHFDataset, data_config=agent_loop.DictConfigWrap(DATA), hf_model_type="qwen3", + url=url, reset_timeout=60.0, attempts=attempts) + return await played.run({"temperature": 1.0}, raw_prompt=row["prompt"], extra_info=row["extra_info"], + tools_kwargs={}, index=0, agent_name="agentenv") + + +def _row(folder): + [row] = pq.read_table(folder / "all.parquet").to_pylist() + assert json.loads(row["extra_info"]["tools"])[0]["function"]["name"] == "add_item" + return row From fa2da056f3a48ab1733f8c11cba17d51c1d965ba Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 07:32:59 -0700 Subject: [PATCH 03/13] fix(training): end a verl rollout's turn with its reply's text, without special tokens; verl CI job Co-Authored-By: Claude Opus 5.5 --- .github/workflows/verl.yml | 63 +++++++++++++++++++++++++++++ src/agent_env/training/verl/loop.py | 7 +++- tst/verl/loop_test.py | 11 ++++- 3 files changed, 78 insertions(+), 3 deletions(-) create mode 100644 .github/workflows/verl.yml diff --git a/.github/workflows/verl.yml b/.github/workflows/verl.yml new file mode 100644 index 00000000..ad20854e --- /dev/null +++ b/.github/workflows/verl.yml @@ -0,0 +1,63 @@ +name: verl contract + +# AgentEnvAgentLoop on verl itself (tst/verl): verl's ToolAgentLoop plays a row the writer wrote, with Qwen3's +# tokenizer and a stand-in rollout server, on verl 0.9.1 and on its main branch. verl, torch and their dependencies +# are installed over the dev environment; torch is the CPU build, as nothing generates. It runs on pull requests that +# touch the helper, the OpenEnv server it plays or the tier, weekly, and on demand; it isn't a required check. + +on: + pull_request: + paths: + - "src/agent_env/training/**" + - "src/agent_env/openenv/**" + - "tst/verl/**" + - "tst/unit/training/**" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/verl.yml" + schedule: + - cron: "41 7 * * 1" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: verl-${{ 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: + loop: + name: agent loop (verl ${{ matrix.verl }}) + runs-on: ubuntu-latest + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + include: + - verl: "0.9.1" + requirement: "verl==0.9.1" + - verl: "main" + requirement: "verl @ git+https://github.com/volcengine/verl@main" + 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 verl extra + run: | + python -m pip install --quiet -r .github/ci-requirements.txt + uv sync --locked --extra dev --extra verl --python "$(command -v python)" + - name: Install verl and CPU torch + run: | + uv pip install --python .venv/bin/python "${{ matrix.requirement }}" torch \ + --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match + .venv/bin/python -c "from importlib.metadata import version; print('verl', version('verl'), 'torch', version('torch'))" + - name: verl tier + run: .venv/bin/python -m pytest tst/verl -rA --timeout=900 --junitxml=junit-verl.xml + - name: No skips + if: always() + run: .venv/bin/python .github/scripts/check_skip_policy.py junit-verl.xml diff --git a/src/agent_env/training/verl/loop.py b/src/agent_env/training/verl/loop.py index 8966f49e..0712047a 100644 --- a/src/agent_env/training/verl/loop.py +++ b/src/agent_env/training/verl/loop.py @@ -90,7 +90,12 @@ async def _play(self, session: _Session, schemas: list[dict[str, Any]], row: dic return output, rollout def _build_assistant_message(self, content: str, agent_data: Any) -> dict[str, Any]: - self._reply = content or "" + """Keeps the turn's text, without the special tokens the parser leaves in it, as the reply to end the turn + with.""" + reply = content or "" + for token in self.tokenizer.all_special_tokens: + reply = reply.replace(token, "") + self._reply = reply.strip() return super()._build_assistant_message(content, agent_data) diff --git a/tst/verl/loop_test.py b/tst/verl/loop_test.py index 4fd191c0..ffda9571 100644 --- a/tst/verl/loop_test.py +++ b/tst/verl/loop_test.py @@ -37,16 +37,23 @@ def tokenizer(): return AutoTokenizer.from_pretrained(TOKENIZER) -def test_a_played_rows_reward_is_its_episodes(serve, tmp_path, tokenizer, teardowns): +def test_a_played_rows_reward_is_its_episodes_and_its_last_reply_ends_the_turn(serve, tmp_path, tokenizer, + monkeypatch, teardowns): server = serve({"a": PLAYED}) row = _row(write_dataset(server.url, tmp_path)) model = _Model(tokenizer) + finish, replies = verl_loop._Session._finish, [] + async def finishing(self, rollout, response=""): + replies.append(response) + await finish(self, rollout, response) + + monkeypatch.setattr(verl_loop._Session, "_finish", finishing) output = asyncio.run(_play(server.url, row, model, tokenizer)) server.stop() assert (output.reward_score, output.extra_fields["agentenv"]) == (1.0, {"attempts": 1, "scores": {"replies": 1.0}}) - assert model.calls == 2 + assert (model.calls, replies) == (2, ["Added a crate."]) assert 0 in output.response_mask and 1 in output.response_mask assert "Added a crate." in tokenizer.decode(output.response_ids) assert len(teardowns.finished) == 2 From a17bf7232dd9aba71f041d362eb35d5e3acffafc Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 07:33:56 -0700 Subject: [PATCH 04/13] docs(training): the verl module and its tier in AGENTS.md and the Makefile Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 5 ++++- Makefile | 7 ++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 22cd0a57..035a309d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -46,6 +46,7 @@ or one of them only where it is used. | `make installer-test` | `tst/installer`: `plugin add` / `remove` through the real pip, uv and pipx, offline against wheels built from the checkout, and the container user journey (`python:3.12-slim`, `--network none`) with the PEP 668 refusal; needs uv, pipx and Docker | | `make trl-test` | `tst/trl`: one `GRPOTrainer` step on CPU through `agent_env.training.trl`, with a tiny model from the Hub; needs trl and CPU torch, which the `dev` extra leaves out | | `make nemo-gym-test` | `tst/nemo_gym`: rollouts through `agent_env.training.nemo_gym`'s resources server in-process; needs the `nemo-gym` extra, on Python 3.13.14 or later, which the `dev` extra leaves out (`uv sync --extra dev --extra nemo-gym`) | +| `make verl-test` | `tst/verl`: rows played through `agent_env.training.verl`'s agent loop on verl itself, with Qwen3's tokenizer from the Hub and a stand-in rollout server; needs verl and CPU torch, which the `dev` extra leaves out | | `make clean-install-test` | the `clean-install` CI gate: both distributions built as the release builds them, installed into a fresh venv from public PyPI with an allowlisted environment, and `agent-env run hello` run twice by name and checked through the local store; needs Python 3.11, uv and git | | `npm ci && npm run typecheck && npm run lint && npm run test:smoke` | in `src/agent_env/explorer/ui`: the explorer UI checks the `ui` job runs (it also runs `npm run build:static`); needs Node 22 | | `pytest tst/unit/env/store_test.py -v` | one file; add `--log-cli-level=DEBUG` for debug logs | @@ -121,7 +122,7 @@ and deserialized through a registry. | `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `smol_vm` (the `smol` extra), `tensorlake` (the `tensorlake` extra), `tenki` (the `tenki` extra), `vercel` (the `vercel` extra), `boxd` (the `boxd` extra), `coreweave` (the `coreweave` extra), `daytona` (the `daytona` extra), `boat` (the `boat` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store), `EnvironmentGatewayProvider` (`gateway`, an env's default), which renders a docker-compose for the gateway and its MCP servers inside the sandbox, and `EnvironmentServerProvider` (`server`), one MCP server in its own container with no gateway in front and no state store; `providers/env_state/` holds env-state providers (`local_postgres` built in). | | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `openenv/` | `build_app(bundle)`: the [OpenEnv](https://github.com/openenv/OpenEnv) server for a bundle's tasks, each session playing one episode through the Episode API, with the envs' MCP tools, an added `end_turn` tool and the task's verifier score as the reward. It prepares bundles as `agent-env run` does (`prepare_bundle(played=True)`, without the agent the client plays) and keeps namespace routing on for the app's life. `openenv` comes with the `openenv` extra, a floor with no cap; this package and `training/` are the only ones that import it. Each task's description in the Task API names its `toolset` (the played agent's envs, role, and whether it gets `bash`). | -| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that both play an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. | +| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, the split-name check the writers share and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times and masking one that fails every time. | | `runner/` | The `[runner]` seam: `Runner.submit()` returns a `RunHandle` (`run_id`, `instance_id`) before the run starts; `LocalRunner` is built in. `runner/store.py` is the run index every runner writes and the explorer reads. | | `bundle/` | A bundle is a folder of entities (`envs/`, `agents/`, `artifacts/`, `skills/`, `tasks/.json`, `evals/.toml`) that runs without a registry: `agent-env run` plans it, writes it into the `@local` namespace (reusing each version whose inputs haven't changed) and runs its evals or tasks; `agent-env bundle init` writes a stored task or eval as one, with every env, agent and artifact it uses in its own folder (each image as its build context, pruned to what its Dockerfile reads and pinned to the Python and npm versions of the stored image, or, given `--image-ref ID=REF`, as that registry image pinned by digest with its context beside it, `image = { ref, dockerfile }`, which a run pulls and builds only when it can't; `task_images` lists the images an import writes, downloading nothing, so an export can push them first) so it runs with no store; `agent-env bundle register` writes a folder into the configured stores under registry ids (``), each task and eval naming the versions it reads and each image it builds pushed, through a ledger in the target store so registering again writes only what changed. | | `eval/` | `Eval` (`eval`), a versioned list of task references, and its store. | @@ -336,6 +337,8 @@ extra from `uv.lock` and with the latest nemo-gym, on the same triggers; it allo `.github/workflows/verifiers.yml` runs `tst/verifiers`, `vf-eval` through verifiers' NeMo Gym taskset (prime-rl's) against `serve`, on verifiers 0.4.0 and its main branch, installed in a venv of its own since verifiers pins mcp 2, on the same triggers; it allows no skips and isn't a required check. +`.github/workflows/verl.yml` runs `tst/verl` on verl 0.9.1 and its main branch, with CPU torch installed over the dev +environment, on the same triggers; it allows no skips and isn't a required check. `.github/workflows/plugin-api.yml` runs the `plugin-api` job on every pull request, title edits included, and on `main`: `.github/scripts/check_plugin_api.py` compares the plugin surface (its `BASES` and `USED`) between `HEAD^1` and `HEAD` with griffe, pinned in the `dev` extra, and fails on a break the diff --git a/Makefile b/Makefile index b01b3a0c..87d9b1c2 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test trl-test nemo-gym-test clean-install-test test clean +.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test trl-test nemo-gym-test verl-test clean-install-test test clean VENV := .venv VENV_BIN := $(abspath $(VENV))/bin @@ -74,6 +74,11 @@ trl-test: ## Run one GRPOTrainer step on CPU through the TRL helper (needs trl a nemo-gym-test: ## Play rollouts through NeMo Gym's resources server for a bundle (needs nemo-gym, on Python 3.13.14+) $(PYTHON) -m pytest tst/nemo_gym --timeout=900 $(PYTEST_ARGS) +# The verl tier plays a row through AgentEnvAgentLoop on verl itself, with Qwen3's tokenizer from the Hub. verl and +# torch aren't in the dev extra: install them first, as .github/workflows/verl.yml does. +verl-test: ## Play rows through the verl agent loop on verl itself (needs verl and CPU torch, and the Hub) + $(PYTHON) -m pytest tst/verl --timeout=900 $(PYTEST_ARGS) + # The clean-install gate CI runs, on the same Python. The script gives every step an allowlisted # environment, so nothing in this shell reaches the install it checks. A failed run keeps its folder. clean-install-test: ## Build both distributions, install them clean from public PyPI and run hello twice (needs python3.11, uv, git) From c531443edfe9ce94b26901bdf9fb1c2a704690c4 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 07:42:40 -0700 Subject: [PATCH 05/13] docs(training): name the verl module in the training package Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/agent_env/training/__init__.py b/src/agent_env/training/__init__.py index dd3f29b1..1440ad6e 100644 --- a/src/agent_env/training/__init__.py +++ b/src/agent_env/training/__init__.py @@ -1,3 +1,3 @@ """Training on AgentEnv tasks with other projects' trainers, one module per trainer, each behind an extra of its own: -``agent_env.training.trl`` for TRL, and ``agent_env.training.nemo_gym`` for NeMo Gym and the trainers that take its -environments, NeMo RL among them.""" +``agent_env.training.trl`` for TRL, ``agent_env.training.nemo_gym`` for NeMo Gym and the trainers that take its +environments, NeMo RL among them, and ``agent_env.training.verl`` for verl.""" From 2156b44a7b1e62e454643b94505eabf70080c9d6 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 08:03:08 -0700 Subject: [PATCH 06/13] test: an extra's inclusion is what it declares; verl's distributions are trl's Co-Authored-By: Claude Opus 5.5 --- .../test_optional_extras_stay_optional.py | 35 ++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/tst/unit/test_optional_extras_stay_optional.py b/tst/unit/test_optional_extras_stay_optional.py index b981a493..1c304699 100644 --- a/tst/unit/test_optional_extras_stay_optional.py +++ b/tst/unit/test_optional_extras_stay_optional.py @@ -9,12 +9,13 @@ import json import subprocess import sys +import tomllib import pytest from packaging.requirements import Requirement from packaging.utils import canonicalize_name -from tst.unit.test_runtime_dependencies_are_imported import SELF, _module_paths, _optional_requirements +from tst.unit.test_runtime_dependencies_are_imported import PYPROJECT, SELF, _module_paths, _optional_requirements _BLOCKING = """ import importlib, importlib.abc, json, sys @@ -26,11 +27,12 @@ def find_spec(self, fullname, path=None, target=None): sys.meta_path.insert(0, Refuse()) """ -# Not core: the Sail and Smol providers, the OpenEnv server and NeMo Gym's resources server need their optional extras, -# and the gateway runs only in its container image, and the code runner is a script that reads its arguments on import. +# Not core: the Sail and Smol providers, the OpenEnv server, NeMo Gym's resources server and the verl helper need their +# optional extras (and the verl agent loop verl itself), and the gateway runs only in its container image, and the code +# runner is a script that reads its arguments on import. _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.env.gateway", + "agent_env.openenv", "agent_env.training.nemo_gym", "agent_env.training.verl", "agent_env.env.gateway", "agent_env.task_step.task_steps.run_code_runner", ) @@ -135,15 +137,22 @@ def _only_in(extra: str) -> list[str]: """Module paths the extra installs and neither a core install nor any other extra does, but for an extra that includes it (``trl`` includes ``openenv``): a distribution two extras share (grpcio, for gcp and tensorlake) stays for the other's modules.""" - mine = _closure(frozenset({extra})) - others = [other for other in _optional_requirements() - if other != extra and not mine <= _closure(frozenset({other}))] + others = [other for other in _optional_requirements() if other != extra and extra not in _included(other)] core = _closure(frozenset()) | set().union(*(_closure(frozenset({other})) for other in others)) provided_by_core = set().union(*(_module_paths(d) for d in core)) extra_only = _closure(frozenset({extra})) - core return sorted(set().union(*(_module_paths(d) for d in extra_only)) - provided_by_core) +def _included(extra: str) -> set[str]: + """The extras ``extra`` declares it includes, with those they include: an extra whose distributions another's + happen to cover, as ``verl``'s are ``trl``'s, isn't included by it.""" + declared = tomllib.loads(PYPROJECT.read_text())["project"]["optional-dependencies"][extra] + requirements = [Requirement(dep) for dep in declared] + return {name for requirement in requirements if canonicalize_name(requirement.name) == SELF + for included in requirement.extras for name in {included} | _included(included)} + + def _outside(extra: str) -> list[str]: """Module paths some extra installs and neither a core install nor ``extra``, with the extras it includes, does.""" installed = _closure(frozenset({extra})) @@ -169,13 +178,23 @@ def _run(script: str, blocked: list[str], *args: str) -> str: return done.stdout.strip() -@pytest.mark.parametrize("extra", sorted(_optional_requirements())) +#: Extras every distribution of which another extra installs too, so blocking what only they install blocks nothing: +#: verl's pyarrow comes with trl's datasets. The test without any extra covers them. +_COVERED = {"verl": "trl"} + + +@pytest.mark.parametrize("extra", sorted(set(_optional_requirements()) - set(_COVERED))) def test_core_and_every_store_module_import_without_the_extra(extra): blocked = _only_in(extra) assert blocked, f"the {extra!r} extra installs nothing core lacks, so this test proves nothing" assert json.loads(_run(_IMPORT_CORE, blocked, json.dumps(_NOT_CORE))) == [] +@pytest.mark.parametrize(("extra", "cover"), sorted(_COVERED.items())) +def test_a_covered_extra_installs_nothing_its_cover_doesnt(extra, cover): + assert _closure(frozenset({extra})) <= _closure(frozenset({cover})) + + def test_core_and_every_store_module_import_without_any_extra(): assert json.loads(_run(_IMPORT_CORE, _not_in_core(), json.dumps(_NOT_CORE))) == [] From 56931011e12a78ee20cb0767ef3ffe17a1d74f47 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 08:13:48 -0700 Subject: [PATCH 07/13] fix(training): a verl split the server keeps empty has a row's columns; keep the extras test's subset rule Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/verl/prepare.py | 10 +++++++++- tst/unit/test_optional_extras_stay_optional.py | 16 ++++------------ tst/unit/training/verl_test.py | 13 ++++++++++++- 3 files changed, 25 insertions(+), 14 deletions(-) diff --git a/src/agent_env/training/verl/prepare.py b/src/agent_env/training/verl/prepare.py index 91ccc876..2221264d 100644 --- a/src/agent_env/training/verl/prepare.py +++ b/src/agent_env/training/verl/prepare.py @@ -18,6 +18,14 @@ #: The ``agent_name`` every row plays under, and the name ``agent_loop.yaml`` registers the loop at. AGENT = "agentenv" LOOP = "agent_env.training.verl.loop.AgentEnvAgentLoop" +#: A row's columns, so a split the server keeps empty still has them. +ROW = pa.schema([ + ("data_source", pa.string()), + ("prompt", pa.list_(pa.struct([("role", pa.string()), ("content", pa.string())]))), + ("agent_name", pa.string()), + ("extra_info", pa.struct([("split", pa.string()), ("index", pa.int64()), ("task_id", pa.string()), + ("tools", pa.string())])), +]) def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, reset_timeout: float = 600.0, @@ -43,7 +51,7 @@ def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, rese "agent_name": AGENT, "extra_info": {"split": split, "index": task["index"], "task_id": task["id"], "tools": json.dumps(tools[key(task)])}, - } for task in described]), out / f"{split}.parquet") + } for task in described], schema=ROW), out / f"{split}.parquet") (out / "agent_loop.yaml").write_text(yaml.safe_dump([{ "name": AGENT, "_target_": LOOP, "url": url, "tool_timeout": tool_timeout, "reset_timeout": reset_timeout, "attempts": attempts, diff --git a/tst/unit/test_optional_extras_stay_optional.py b/tst/unit/test_optional_extras_stay_optional.py index 1c304699..00cfc15a 100644 --- a/tst/unit/test_optional_extras_stay_optional.py +++ b/tst/unit/test_optional_extras_stay_optional.py @@ -9,13 +9,12 @@ import json import subprocess import sys -import tomllib import pytest from packaging.requirements import Requirement from packaging.utils import canonicalize_name -from tst.unit.test_runtime_dependencies_are_imported import PYPROJECT, SELF, _module_paths, _optional_requirements +from tst.unit.test_runtime_dependencies_are_imported import SELF, _module_paths, _optional_requirements _BLOCKING = """ import importlib, importlib.abc, json, sys @@ -137,22 +136,15 @@ def _only_in(extra: str) -> list[str]: """Module paths the extra installs and neither a core install nor any other extra does, but for an extra that includes it (``trl`` includes ``openenv``): a distribution two extras share (grpcio, for gcp and tensorlake) stays for the other's modules.""" - others = [other for other in _optional_requirements() if other != extra and extra not in _included(other)] + mine = _closure(frozenset({extra})) + others = [other for other in _optional_requirements() + if other != extra and not mine <= _closure(frozenset({other}))] core = _closure(frozenset()) | set().union(*(_closure(frozenset({other})) for other in others)) provided_by_core = set().union(*(_module_paths(d) for d in core)) extra_only = _closure(frozenset({extra})) - core return sorted(set().union(*(_module_paths(d) for d in extra_only)) - provided_by_core) -def _included(extra: str) -> set[str]: - """The extras ``extra`` declares it includes, with those they include: an extra whose distributions another's - happen to cover, as ``verl``'s are ``trl``'s, isn't included by it.""" - declared = tomllib.loads(PYPROJECT.read_text())["project"]["optional-dependencies"][extra] - requirements = [Requirement(dep) for dep in declared] - return {name for requirement in requirements if canonicalize_name(requirement.name) == SELF - for included in requirement.extras for name in {included} | _included(included)} - - def _outside(extra: str) -> list[str]: """Module paths some extra installs and neither a core install nor ``extra``, with the extras it includes, does.""" installed = _closure(frozenset({extra})) diff --git a/tst/unit/training/verl_test.py b/tst/unit/training/verl_test.py index a3e29b61..515a682b 100644 --- a/tst/unit/training/verl_test.py +++ b/tst/unit/training/verl_test.py @@ -10,7 +10,7 @@ from agent_env.openenv.environment import END_TURN from agent_env.training.verl.__main__ import main -from agent_env.training.verl.prepare import LOOP, write_dataset +from agent_env.training.verl.prepare import LOOP, ROW, write_dataset from tst.unit.task.episode.run_test import ( ALICE, ITEMS, @@ -66,6 +66,17 @@ def test_each_toolset_is_probed_once_and_a_rows_tools_are_its_tasks(serve, tmp_p assert len(teardowns.finished) == 2 +def test_a_split_the_server_keeps_empty_has_a_rows_columns(serve, tmp_path): + server = serve({"a": PLAYED}, {"empty": []}) + + out = write_dataset(server.url, tmp_path) + server.stop() + + empty = pq.read_table(out / "empty.parquet") + assert (empty.num_rows, empty.schema) == (0, ROW) + assert pq.read_table(out / "all.parquet").schema == ROW + + def test_a_split_whose_name_isnt_a_file_name_is_refused_before_anything_is_written(serve, tmp_path): server = serve({"a": PLAYED}, {"../escape": ["a"]}) From 08cae59b93aa717f8b4531d32f68db51ba32b06c Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 10:00:39 -0700 Subject: [PATCH 08/13] fix(training): verl loop: typed tool schemas for verl's parsers, refuse a prompt verl would cut, raise on an unscored episode, bound replays qwen3_coder and deepseek_v4 read tool.function.* on the schemas; built unvalidated they were dicts, and qwen3_coder swallowed the AttributeError and dropped every tool call. verl left-truncates a first prompt over rollout.prompt_length, cutting the system prompt and tools, and its length filter doesn't count a row's tools. An episode that ends unscored without failing is a configuration error, not an infra failure to replay. Replays stop once reset_timeout has passed since the rollout began. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 2 +- src/agent_env/training/verl/loop.py | 69 ++++++++++++----- tst/verl/loop_test.py | 111 +++++++++++++++++++++++----- 3 files changed, 144 insertions(+), 38 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 035a309d..b033f789 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -122,7 +122,7 @@ and deserialized through a registry. | `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `smol_vm` (the `smol` extra), `tensorlake` (the `tensorlake` extra), `tenki` (the `tenki` extra), `vercel` (the `vercel` extra), `boxd` (the `boxd` extra), `coreweave` (the `coreweave` extra), `daytona` (the `daytona` extra), `boat` (the `boat` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store), `EnvironmentGatewayProvider` (`gateway`, an env's default), which renders a docker-compose for the gateway and its MCP servers inside the sandbox, and `EnvironmentServerProvider` (`server`), one MCP server in its own container with no gateway in front and no state store; `providers/env_state/` holds env-state providers (`local_postgres` built in). | | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `openenv/` | `build_app(bundle)`: the [OpenEnv](https://github.com/openenv/OpenEnv) server for a bundle's tasks, each session playing one episode through the Episode API, with the envs' MCP tools, an added `end_turn` tool and the task's verifier score as the reward. It prepares bundles as `agent-env run` does (`prepare_bundle(played=True)`, without the agent the client plays) and keeps namespace routing on for the app's life. `openenv` comes with the `openenv` extra, a floor with no cap; this package and `training/` are the only ones that import it. Each task's description in the Task API names its `toolset` (the played agent's envs, role, and whether it gets `bash`). | -| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, the split-name check the writers share and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times and masking one that fails every time. | +| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, the split-name check the writers share and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and masking one that fails every time; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), and raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing. | | `runner/` | The `[runner]` seam: `Runner.submit()` returns a `RunHandle` (`run_id`, `instance_id`) before the run starts; `LocalRunner` is built in. `runner/store.py` is the run index every runner writes and the explorer reads. | | `bundle/` | A bundle is a folder of entities (`envs/`, `agents/`, `artifacts/`, `skills/`, `tasks/.json`, `evals/.toml`) that runs without a registry: `agent-env run` plans it, writes it into the `@local` namespace (reusing each version whose inputs haven't changed) and runs its evals or tasks; `agent-env bundle init` writes a stored task or eval as one, with every env, agent and artifact it uses in its own folder (each image as its build context, pruned to what its Dockerfile reads and pinned to the Python and npm versions of the stored image, or, given `--image-ref ID=REF`, as that registry image pinned by digest with its context beside it, `image = { ref, dockerfile }`, which a run pulls and builds only when it can't; `task_images` lists the images an import writes, downloading nothing, so an export can push them first) so it runs with no store; `agent-env bundle register` writes a folder into the configured stores under registry ids (``), each task and eval naming the versions it reads and each image it builds pushed, through a ledger in the target store so registering again writes only what changed. | | `eval/` | `Eval` (`eval`), a versioned list of task references, and its store. | diff --git a/src/agent_env/training/verl/loop.py b/src/agent_env/training/verl/loop.py index 0712047a..fa1e95c6 100644 --- a/src/agent_env/training/verl/loop.py +++ b/src/agent_env/training/verl/loop.py @@ -8,6 +8,7 @@ import json import logging +import time from typing import Any try: @@ -37,15 +38,17 @@ def __init__(self, shared: Shared): class AgentEnvAgentLoop(ToolAgentLoop): """Plays a row's task as one episode on the OpenEnv server at ``url``. A rollout the env fails, a deploy, a call past its deadline or a server that went away, is played again from the start on a new episode, a fresh sample of - the same policy, up to ``attempts`` times; one that fails every time is kept out of the loss, its response masked, - with the reward 0. ``agent_loop.yaml`` sets the arguments; verl passes the rest. + the same policy, up to ``attempts`` times, but not once ``reset_timeout`` seconds have passed since it began; one + that fails every time is kept out of the loss, its response masked, with the reward 0. A task that ends unscored + without failing, and a prompt that doesn't fit ``rollout.prompt_length`` with its tools, are configuration errors no + replay changes, and raise. ``agent_loop.yaml`` sets the arguments; verl passes the rest. Args: url: The OpenEnv server's base URL. tool_timeout: One tool call's deadline in seconds. The server ends a call that misses it, and the episode. reset_timeout: How long an episode may take to deploy, and each turn to be ended, in seconds; also how long - a rollout waits for room on a server at its limit. - attempts: How many times a rollout the env fails is played. + a rollout waits for room on a server at its limit, and how long after it began it may be played again. + attempts: How many times a rollout the env fails is played at most. """ def __init__(self, *args: Any, url: str, tool_timeout: float = 60.0, reset_timeout: float = 600.0, @@ -53,34 +56,39 @@ def __init__(self, *args: Any, url: str, tool_timeout: float = 60.0, reset_timeo super().__init__(*args, **kwargs) self._url, self._attempts = url.rstrip("/"), attempts self._tool_timeout, self._reset_timeout = tool_timeout, reset_timeout + self._row: dict[str, Any] = {} self._reply = "" async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoopOutput: - row = kwargs["extra_info"] + row = self._row = kwargs["extra_info"] schemas = json.loads(row["tools"]) shared = Shared(self._url, frozenset(schema["function"]["name"] for schema in schemas), self._tool_timeout, self._reset_timeout, patience=self._reset_timeout) + began = time.monotonic() for attempt in range(1, self._attempts + 1): session = _Session(shared) try: - output, rollout = await self._play(session, schemas, row, sampling_params, kwargs) + output, rollout = await self._play(session, schemas, sampling_params, kwargs) finally: await session._close() - if rollout.failure is None and rollout.reward is not None: + if rollout.failure is None: + if rollout.reward is None: + raise ValueError(f"{self._task()} ended without a reward: it records no score its server's reward " + "reads, so it can't be trained on") output.reward_score = float(rollout.reward) output.extra_fields["agentenv"] = {"attempts": attempt, "scores": rollout.scores} return output - failure = rollout.failure or "the episode ended without a reward" - logger.warning("task %s of split %s, attempt %d of %d: %s", row["task_id"], row["split"], attempt, - self._attempts, failure) + logger.warning("%s, attempt %d of %d: %s", self._task(), attempt, self._attempts, rollout.failure) + if time.monotonic() - began >= self._reset_timeout: + break output.response_mask = [0] * len(output.response_mask) output.reward_score = 0.0 - output.extra_fields["agentenv"] = {"attempts": self._attempts, "failure": failure} + output.extra_fields["agentenv"] = {"attempts": attempt, "failure": rollout.failure} return output - async def _play(self, session: _Session, schemas: list[dict[str, Any]], row: dict[str, Any], - sampling_params: dict[str, Any], kwargs: dict[str, Any]) -> tuple[AgentLoopOutput, Rollout]: - await session._start(row["split"], int(row["index"]), row["task_id"]) + async def _play(self, session: _Session, schemas: list[dict[str, Any]], sampling_params: dict[str, Any], + kwargs: dict[str, Any]) -> tuple[AgentLoopOutput, Rollout]: + await session._start(self._row["split"], int(self._row["index"]), self._row["task_id"]) self.tools = {schema["function"]["name"]: _tool(session, schema) for schema in schemas} self.tool_schemas = schemas self._reply = "" @@ -89,6 +97,14 @@ async def _play(self, session: _Session, schemas: list[dict[str, Any]], row: dic await session._finish(rollout, self._reply) return output, rollout + def _cap_text_prompt_length(self, prompt_ids: list[int]) -> list[int]: + """Refuses the first prompt when it doesn't fit, where verl would cut it from the left, losing the system + prompt and the tools; verl's own length filter doesn't count a row's tools.""" + if len(prompt_ids) > (limit := self.rollout_config.prompt_length): + raise ValueError(f"{self._task()}'s prompt is {len(prompt_ids)} tokens with its tools, over " + f"rollout.prompt_length={limit}: raise data.max_prompt_length to fit it") + return prompt_ids + def _build_assistant_message(self, content: str, agent_data: Any) -> dict[str, Any]: """Keeps the turn's text, without the special tokens the parser leaves in it, as the reply to end the turn with.""" @@ -99,13 +115,30 @@ def _build_assistant_message(self, content: str, agent_data: Any) -> dict[str, A return super()._build_assistant_message(content, agent_data) + def _task(self) -> str: + return f"task {self._row['task_id']} of split {self._row['split']}" + + def _tool(session: _Session, schema: dict[str, Any]) -> FunctionTool: - """The episode's tool as the function tool ``ToolAgentLoop`` calls; its schema is built unvalidated, since verl's - schema models keep no nested JSON Schema, and the chat template renders the row's own.""" + """The episode's tool as the function tool ``ToolAgentLoop`` calls.""" name = schema["function"]["name"] async def call(**arguments: Any) -> str: return await session._step(name, arguments) - return FunctionTool(name=name, fn=call, tool_schema=OpenAIFunctionToolSchema.model_construct(**schema), - is_async=True) + return FunctionTool(name=name, fn=call, tool_schema=_parser_schema(schema), is_async=True) + + +def _parser_schema(schema: dict[str, Any]) -> OpenAIFunctionToolSchema: + """The tool as verl's parsers read it: those that convert arguments (``qwen3_coder``) or match a call to a tool by + its parameters (``deepseek_v4``) take each parameter's type, description and enum. verl's schema models keep no + nested JSON Schema, so the chat template renders the row's own, and a parameter without a type is left out, which + a parser takes as a string.""" + function = schema["function"] + parameters = function.get("parameters") or {} + return OpenAIFunctionToolSchema.model_validate({"type": "function", "function": { + "name": function["name"], "description": function.get("description") or "", + "parameters": {"type": "object", "required": parameters.get("required") or [], "properties": { + argument: {key: spec[key] for key in ("type", "description", "enum") if key in spec} + for argument, spec in (parameters.get("properties") or {}).items() if "type" in spec}}, + }}) diff --git a/tst/verl/loop_test.py b/tst/verl/loop_test.py index ffda9571..b680de8f 100644 --- a/tst/verl/loop_test.py +++ b/tst/verl/loop_test.py @@ -1,7 +1,8 @@ """``AgentEnvAgentLoop`` on verl itself: verl's ``ToolAgentLoop`` plays a row ``write_dataset`` wrote, with Qwen3's tokenizer and chat template and a stand-in rollout server, on the unit tier's bundle served on loopback. A played row's -reward is its episode's, a rollout the env fails is played again on a new episode, and one it fails every time is kept -out of the loss. CI runs it on verl 0.9.1.""" +reward is its episode's, through verl's hermes and qwen3_coder parsers; a rollout the env fails is played again on a new +episode, but not once its time is up, and one it fails every time is kept out of the loss; a prompt that doesn't fit and +an episode that ends unscored without failing raise. CI runs it on verl 0.9.1 and main.""" import asyncio import json @@ -19,17 +20,23 @@ exc_type=ImportError) omegaconf = pytest.importorskip("omegaconf", reason=missing_capability_reason(VERL)) agent_loop = pytest.importorskip("verl.experimental.agent_loop.agent_loop", reason=missing_capability_reason(VERL)) +tool_parser = pytest.importorskip("verl.experimental.agent_loop.tool_parser", reason=missing_capability_reason(VERL)) rl_dataset = pytest.importorskip("verl.utils.dataset.rl_dataset", reason=missing_capability_reason(VERL)) replica = pytest.importorskip("verl.workers.rollout.replica", reason=missing_capability_reason(VERL)) TOKENIZER = "Qwen/Qwen3-0.6B" -CONFIG = omegaconf.OmegaConf.create({"actor_rollout_ref": {"rollout": { - "prompt_length": 2048, "response_length": 1024, "full_determinism": False, - "multi_turn": {"format": "hermes", "max_user_turns": None, "max_assistant_turns": 4, "max_parallel_calls": 1, - "max_tool_response_length": 2048, "tool_response_truncate_side": "middle"}, -}}}) DATA = omegaconf.OmegaConf.create({"apply_chat_template_kwargs": {"enable_thinking": False}}) -ADD_A_CRATE = '\n{"name": "add_item", "arguments": {"item": "crate"}}\n<|im_end|>' +ADD_A_CRATE = { + "hermes": '\n{"name": "add_item", "arguments": {"item": "crate"}}\n<|im_end|>', + "qwen3_coder": ("\n\n\ncrate\n\n\n" + "<|im_end|>"), +} +SCHEDULE = {"type": "function", "function": {"name": "submit_schedule", "description": "Submit the day's schedule.", + "parameters": {"type": "object", "required": ["schedule"], "properties": { + "schedule": {"type": "array", "items": {"type": "integer"}, + "description": "One action an hour."}, + "hour": {"type": "integer"}, + "note": {"anyOf": [{"type": "string"}, {"type": "null"}]}}}}} @pytest.fixture(scope="module") @@ -37,11 +44,12 @@ def tokenizer(): return AutoTokenizer.from_pretrained(TOKENIZER) -def test_a_played_rows_reward_is_its_episodes_and_its_last_reply_ends_the_turn(serve, tmp_path, tokenizer, +@pytest.mark.parametrize("parser", ["hermes", "qwen3_coder"]) +def test_a_played_rows_reward_is_its_episodes_and_its_last_reply_ends_the_turn(serve, tmp_path, tokenizer, parser, monkeypatch, teardowns): server = serve({"a": PLAYED}) row = _row(write_dataset(server.url, tmp_path)) - model = _Model(tokenizer) + model = _Model(tokenizer, parser) finish, replies = verl_loop._Session._finish, [] async def finishing(self, rollout, response=""): @@ -49,7 +57,7 @@ async def finishing(self, rollout, response=""): await finish(self, rollout, response) monkeypatch.setattr(verl_loop._Session, "_finish", finishing) - output = asyncio.run(_play(server.url, row, model, tokenizer)) + output = asyncio.run(_play(server.url, row, model, tokenizer, parser=parser)) server.stop() assert (output.reward_score, output.extra_fields["agentenv"]) == (1.0, {"attempts": 1, "scores": {"replies": 1.0}}) @@ -59,6 +67,21 @@ async def finishing(self, rollout, response=""): assert len(teardowns.finished) == 2 +def test_a_tools_parameters_reach_verls_parsers_with_their_types(tokenizer): + schema = verl_loop._parser_schema(SCHEDULE) + call = ("\n\n\n[1, 0, -1]\n\n" + "\n7\n\n\nnone today\n\n\n") + parser = tool_parser.ToolParser.get_tool_parser("qwen3_coder", tokenizer) + + _, [parsed] = asyncio.run(parser.extract_tool_calls(tokenizer.encode(call, add_special_tokens=False), [schema])) + + assert (schema.function.name, schema.function.parameters.required) == ("submit_schedule", ["schedule"]) + assert {name: spec.model_dump(exclude_none=True) for name, spec in schema.function.parameters.properties.items()} == { + "schedule": {"type": "array", "description": "One action an hour."}, "hour": {"type": "integer"}} + assert (parsed.name, json.loads(parsed.arguments)) == ( + "submit_schedule", {"schedule": [1, 0, -1], "hour": 7, "note": "none today"}) + + def test_a_rollout_the_env_fails_is_played_again_on_a_new_episode(serve, tmp_path, tokenizer, monkeypatch, teardowns): server = serve({"a": PLAYED}) row = _row(write_dataset(server.url, tmp_path)) @@ -94,26 +117,76 @@ async def failing(self, rollout, *args): assert output.response_mask and set(output.response_mask) == {0} +def test_a_rollout_isnt_played_again_once_its_time_is_up(serve, tmp_path, tokenizer, monkeypatch): + server = serve({"a": PLAYED}) + row = _row(write_dataset(server.url, tmp_path)) + + async def failing_late(self, rollout, *args): + await asyncio.sleep(0.3) + raise Failed("the env's deploy missed its deadline") + + monkeypatch.setattr(verl_loop._Session, "_begin", failing_late) + output = asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer, reset_timeout=0.2)) + server.stop() + + assert (output.reward_score, output.extra_fields["agentenv"]) == ( + 0.0, {"attempts": 1, "failure": "the env's deploy missed its deadline"}) + + +def test_a_prompt_that_doesnt_fit_with_its_tools_is_refused_rather_than_cut(serve, tmp_path, tokenizer): + server = serve({"a": PLAYED}) + row = _row(write_dataset(server.url, tmp_path)) + + with pytest.raises(ValueError, match=r"task a of split all's prompt is \d+ tokens with its tools, over " + r"rollout\.prompt_length=64"): + asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer, prompt_length=64)) + server.stop() + + +def test_an_episode_that_ends_without_failing_or_a_reward_raises(serve, tmp_path, tokenizer, monkeypatch): + server = serve({"a": PLAYED}) + row = _row(write_dataset(server.url, tmp_path)) + model = _Model(tokenizer) + finish = verl_loop._Session._finish + + async def unscored(self, rollout, response=""): + await finish(self, rollout, response) + rollout.reward = None + + monkeypatch.setattr(verl_loop._Session, "_finish", unscored) + with pytest.raises(ValueError, match="task a of split all ended without a reward: it records no score"): + asyncio.run(_play(server.url, row, model, tokenizer)) + server.stop() + + assert model.calls == 2 + + class _Model: - """verl's rollout server for the loop: it adds a crate, then replies once the tool answered.""" + """verl's rollout server for the loop: it adds a crate in ``parser``'s format, then replies once the tool + answered.""" - def __init__(self, tokenizer): - self.tokenizer, self.calls = tokenizer, 0 + def __init__(self, tokenizer, parser="hermes"): + self.tokenizer, self.parser, self.calls = tokenizer, parser, 0 async def generate(self, request_id, *, prompt_ids, sampling_params, **_): self.calls += 1 answered = "" in self.tokenizer.decode(prompt_ids) - reply = "Added a crate.<|im_end|>" if answered else ADD_A_CRATE + reply = "Added a crate.<|im_end|>" if answered else ADD_A_CRATE[self.parser] return replica.TokenOutput(token_ids=self.tokenizer.encode(reply, add_special_tokens=False)) -async def _play(url, row, model, tokenizer, attempts=3): +async def _play(url, row, model, tokenizer, *, parser="hermes", prompt_length=2048, reset_timeout=60.0, attempts=3): + config = omegaconf.OmegaConf.create({"actor_rollout_ref": {"rollout": { + "prompt_length": prompt_length, "response_length": 1024, "full_determinism": False, + "multi_turn": {"format": parser, "max_user_turns": None, "max_assistant_turns": 4, "max_parallel_calls": 1, + "max_tool_response_length": 2048, "tool_response_truncate_side": "middle"}, + }}}) played = verl_loop.AgentEnvAgentLoop( - trainer_config=agent_loop.DictConfigWrap(CONFIG), server_manager=model, tokenizer=tokenizer, processor=None, + trainer_config=agent_loop.DictConfigWrap(config), server_manager=model, tokenizer=tokenizer, processor=None, dataset_cls=rl_dataset.RLHFDataset, data_config=agent_loop.DictConfigWrap(DATA), hf_model_type="qwen3", - url=url, reset_timeout=60.0, attempts=attempts) + url=url, reset_timeout=reset_timeout, attempts=attempts) return await played.run({"temperature": 1.0}, raw_prompt=row["prompt"], extra_info=row["extra_info"], - tools_kwargs={}, index=0, agent_name="agentenv") + tools_kwargs={}, index=0, agent_name="agentenv") def _row(folder): From 5e20c21d2ca9420f2d33e34f8b6588321a964b42 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 10:25:11 -0700 Subject: [PATCH 09/13] refactor(training): verl loop: session calls inside _Session, verl internals named and guarded; one read-and-probe for both writers The loop called Session's underscore methods from outside; _Session now owns start, tool, finish and close, as the TRL and NeMo Gym helpers' subclasses do. The class docstring names the verl internals the loop leans on, and run refuses a second rollout on one instance, since verl building a loop per trajectory is what keeps their state apart. tasks_and_tools in _openenv replaces the read, split-name check and probe both writers repeated. The verl tier fails envs through the fake env and a broken server rather than by patching Session, serves an unscored task through the serve fixture's new reward, and checks the config's loop path names the class. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 2 +- src/agent_env/training/_openenv.py | 10 +- src/agent_env/training/nemo_gym/prepare.py | 11 +-- src/agent_env/training/verl/loop.py | 63 +++++++----- src/agent_env/training/verl/prepare.py | 10 +- tst/unit/training/_support.py | 9 +- tst/verl/loop_test.py | 110 ++++++++++++--------- 7 files changed, 126 insertions(+), 89 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index b033f789..7367e4f1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -122,7 +122,7 @@ and deserialized through a registry. | `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `smol_vm` (the `smol` extra), `tensorlake` (the `tensorlake` extra), `tenki` (the `tenki` extra), `vercel` (the `vercel` extra), `boxd` (the `boxd` extra), `coreweave` (the `coreweave` extra), `daytona` (the `daytona` extra), `boat` (the `boat` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store), `EnvironmentGatewayProvider` (`gateway`, an env's default), which renders a docker-compose for the gateway and its MCP servers inside the sandbox, and `EnvironmentServerProvider` (`server`), one MCP server in its own container with no gateway in front and no state store; `providers/env_state/` holds env-state providers (`local_postgres` built in). | | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `openenv/` | `build_app(bundle)`: the [OpenEnv](https://github.com/openenv/OpenEnv) server for a bundle's tasks, each session playing one episode through the Episode API, with the envs' MCP tools, an added `end_turn` tool and the task's verifier score as the reward. It prepares bundles as `agent-env run` does (`prepare_bundle(played=True)`, without the agent the client plays) and keeps namespace routing on for the app's life. `openenv` comes with the `openenv` extra, a floor with no cap; this package and `training/` are the only ones that import it. Each task's description in the Task API names its `toolset` (the played agent's envs, role, and whether it gets `bash`). | -| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, the split-name check the writers share and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and masking one that fails every time; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), and raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing. | +| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, `tasks_and_tools` (what the row writers read: the tasks, the split names checked, each toolset's tools) and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and masking one that fails every time; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), and raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing. | | `runner/` | The `[runner]` seam: `Runner.submit()` returns a `RunHandle` (`run_id`, `instance_id`) before the run starts; `LocalRunner` is built in. `runner/store.py` is the run index every runner writes and the explorer reads. | | `bundle/` | A bundle is a folder of entities (`envs/`, `agents/`, `artifacts/`, `skills/`, `tasks/.json`, `evals/.toml`) that runs without a registry: `agent-env run` plans it, writes it into the `@local` namespace (reusing each version whose inputs haven't changed) and runs its evals or tasks; `agent-env bundle init` writes a stored task or eval as one, with every env, agent and artifact it uses in its own folder (each image as its build context, pruned to what its Dockerfile reads and pinned to the Python and npm versions of the stored image, or, given `--image-ref ID=REF`, as that registry image pinned by digest with its context beside it, `image = { ref, dockerfile }`, which a run pulls and builds only when it can't; `task_images` lists the images an import writes, downloading nothing, so an export can push them first) so it runs with no store; `agent-env bundle register` writes a folder into the configured stores under registry ids (``), each task and eval naming the versions it reads and each image it builds pushed, through a ledger in the target store so registering again writes only what changed. | | `eval/` | `Eval` (`eval`), a versioned list of task references, and its store. | diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index b24a9e39..98510e09 100644 --- a/src/agent_env/training/_openenv.py +++ b/src/agent_env/training/_openenv.py @@ -366,11 +366,17 @@ def read_tasks(url: str, *, local: str) -> tuple[str, dict[str, list[dict[str, A return name, tasks -def check_split_names(tasks: Mapping[str, Any]) -> None: - """Refuses splits whose names can't name the data file a writer writes for each.""" +def tasks_and_tools(url: str, *, local: str, timeout: float + ) -> tuple[str, dict[str, list[dict[str, Any]]], dict[str, tuple[dict[str, Any], ...]]]: + """What a writer of a split's rows reads: the server's env name, each split's tasks, refusing a split whose name + can't name a data file, and every tool each toolset's turns offer, by the toolset's key, found by playing its + first task through (``probe``), so reading deploys those.""" + name, tasks = read_tasks(url, local=local) if unfit := [split for split in tasks if not _FILE.fullmatch(split)]: raise ValueError(f"the server's splits {unfit} can't name a data file: a split's name must be letters, " "digits, '.', '_' and '-', not starting with '.'") + return name, tasks, {each: asyncio.run(probe(url, toolset, timeout))[1] + for each, toolset in toolsets(name, tasks).items()} def toolsets(name: str, tasks: dict[str, list[dict[str, Any]]]) -> dict[str, NamedToolset]: diff --git a/src/agent_env/training/nemo_gym/prepare.py b/src/agent_env/training/nemo_gym/prepare.py index 05655762..2b8f404f 100644 --- a/src/agent_env/training/nemo_gym/prepare.py +++ b/src/agent_env/training/nemo_gym/prepare.py @@ -3,7 +3,6 @@ from __future__ import annotations -import asyncio import json import re from importlib.metadata import version @@ -13,7 +12,7 @@ import yaml from agent_env.openenv import build_app -from agent_env.training._openenv import Server, check_split_names, key, messages, probe, read_tasks, toolsets +from agent_env.training._openenv import Server, key, messages, tasks_and_tools #: The routes a resources server serves itself, which a tool's route can't take. GYM_ROUTES = frozenset({"seed_session", "verify", "close_session", "aggregate_metrics", "reverify_mode"}) @@ -68,11 +67,9 @@ def write_server(source: str | Path, root: str | Path, *, name: str | None = Non served = None if url else Server(build_app(source, reward=summed if reward_components else reward, sandbox=sandbox, max_concurrent_envs=1)) try: - env_name, tasks = read_tasks(url or served.url, local="name the bundle's folder instead of its server") - check_split_names(tasks) - grouped = toolsets(env_name, tasks) - tools = {each: [_function(tool) for tool in asyncio.run(probe(url or served.url, toolset, reset_timeout))[1]] - for each, toolset in grouped.items()} + env_name, tasks, probed = tasks_and_tools(url or served.url, local="name the bundle's folder instead of its " + "server", timeout=reset_timeout) + tools = {each: [_function(tool) for tool in found] for each, found in probed.items()} finally: if served is not None: served.stop() diff --git a/src/agent_env/training/verl/loop.py b/src/agent_env/training/verl/loop.py index fa1e95c6..57cfd3ad 100644 --- a/src/agent_env/training/verl/loop.py +++ b/src/agent_env/training/verl/loop.py @@ -2,7 +2,8 @@ ``ToolAgentLoop``, which generates on verl's rollout server, parses the model's tool calls with ``multi_turn.format``'s parser and adds their answers to the context, with the row's tools in place of verl's own: each call is a call to the tool on the row's episode, the model's last reply ends the turn it left open, and the episode's -reward is the rollout's ``reward_score``, so verl's reward function isn't called. Verified with verl 0.9.1.""" +reward is the rollout's ``reward_score``, so verl's reward function isn't called. CI checks it on verl 0.9.1 and +main.""" from __future__ import annotations @@ -27,14 +28,6 @@ logger = logging.getLogger(__name__) -class _Session(Session): - offered_by = "its row lists" - - def __init__(self, shared: Shared): - super().__init__() - self._shared = shared - - class AgentEnvAgentLoop(ToolAgentLoop): """Plays a row's task as one episode on the OpenEnv server at ``url``. A rollout the env fails, a deploy, a call past its deadline or a server that went away, is played again from the start on a new episode, a fresh sample of @@ -43,6 +36,12 @@ class AgentEnvAgentLoop(ToolAgentLoop): without failing, and a prompt that doesn't fit ``rollout.prompt_length`` with its tools, are configuration errors no replay changes, and raise. ``agent_loop.yaml`` sets the arguments; verl passes the rest. + It leans on verl internals, which CI checks on verl 0.9.1 and main: verl builds an agent loop per trajectory, so an + instance holds one rollout's state, and ``run`` refuses a second; ``ToolAgentLoop.run`` reads the tools from + ``self.tools`` and ``self.tool_schemas``; ``_build_assistant_message`` is overridden to keep the reply, and + ``AgentLoopBase._cap_text_prompt_length`` to refuse a prompt verl would cut; and the tools are verl's + ``FunctionTool`` with ``OpenAIFunctionToolSchema``. + Args: url: The OpenEnv server's base URL. tool_timeout: One tool call's deadline in seconds. The server ends a call that misses it, and the episode. @@ -60,6 +59,9 @@ def __init__(self, *args: Any, url: str, tool_timeout: float = 60.0, reset_timeo self._reply = "" async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoopOutput: + if self._row: + raise RuntimeError("verl ran an AgentEnvAgentLoop a second time, but one holds a single rollout's state: " + "this verl doesn't build an agent loop per trajectory, as 0.9 does") row = self._row = kwargs["extra_info"] schemas = json.loads(row["tools"]) shared = Shared(self._url, frozenset(schema["function"]["name"] for schema in schemas), self._tool_timeout, @@ -70,7 +72,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop try: output, rollout = await self._play(session, schemas, sampling_params, kwargs) finally: - await session._close() + await session.close() if rollout.failure is None: if rollout.reward is None: raise ValueError(f"{self._task()} ended without a reward: it records no score its server's reward " @@ -88,14 +90,12 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop async def _play(self, session: _Session, schemas: list[dict[str, Any]], sampling_params: dict[str, Any], kwargs: dict[str, Any]) -> tuple[AgentLoopOutput, Rollout]: - await session._start(self._row["split"], int(self._row["index"]), self._row["task_id"]) - self.tools = {schema["function"]["name"]: _tool(session, schema) for schema in schemas} + await session.start(self._row) + self.tools = {schema["function"]["name"]: session.tool(schema) for schema in schemas} self.tool_schemas = schemas self._reply = "" output = await super().run(sampling_params, **kwargs) - rollout = session._rollout - await session._finish(rollout, self._reply) - return output, rollout + return output, await session.finish(self._reply) def _cap_text_prompt_length(self, prompt_ids: list[int]) -> list[int]: """Refuses the first prompt when it doesn't fit, where verl would cut it from the left, losing the system @@ -114,19 +114,38 @@ def _build_assistant_message(self, content: str, agent_data: Any) -> dict[str, A self._reply = reply.strip() return super()._build_assistant_message(content, agent_data) - def _task(self) -> str: return f"task {self._row['task_id']} of split {self._row['split']}" -def _tool(session: _Session, schema: dict[str, Any]) -> FunctionTool: - """The episode's tool as the function tool ``ToolAgentLoop`` calls.""" - name = schema["function"]["name"] +class _Session(Session): + """One rollout's session: it starts the row's episode, offers the row's tools as verl's function tools, and ends + the turn left open with the model's reply.""" + + offered_by = "its row lists" + + def __init__(self, shared: Shared): + super().__init__() + self._shared = shared + + async def start(self, row: dict[str, Any]) -> None: + await self._start(row["split"], int(row["index"]), row["task_id"]) + + def tool(self, schema: dict[str, Any]) -> FunctionTool: + name = schema["function"]["name"] + + async def call(**arguments: Any) -> str: + return await self._step(name, arguments) + + return FunctionTool(name=name, fn=call, tool_schema=_parser_schema(schema), is_async=True) - async def call(**arguments: Any) -> str: - return await session._step(name, arguments) + async def finish(self, reply: str) -> Rollout: + rollout = self._rollout + await self._finish(rollout, reply) + return rollout - return FunctionTool(name=name, fn=call, tool_schema=_parser_schema(schema), is_async=True) + async def close(self) -> None: + await self._close() def _parser_schema(schema: dict[str, Any]) -> OpenAIFunctionToolSchema: diff --git a/src/agent_env/training/verl/prepare.py b/src/agent_env/training/verl/prepare.py index 2221264d..30b560dc 100644 --- a/src/agent_env/training/verl/prepare.py +++ b/src/agent_env/training/verl/prepare.py @@ -4,7 +4,6 @@ from __future__ import annotations -import asyncio import json from pathlib import Path from typing import Any @@ -13,7 +12,7 @@ import pyarrow.parquet as pq import yaml -from agent_env.training._openenv import check_split_names, key, messages, probe, read_tasks, toolsets +from agent_env.training._openenv import key, messages, tasks_and_tools #: The ``agent_name`` every row plays under, and the name ``agent_loop.yaml`` registers the loop at. AGENT = "agentenv" @@ -38,10 +37,9 @@ def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, rese if attempts < 1: raise ValueError(f"attempts must be at least 1, not {attempts}") url = url.rstrip("/") - env_name, tasks = read_tasks(url, local="serve the bundle with agent_env.openenv.build_app first") - check_split_names(tasks) - tools = {each: [_function(tool) for tool in asyncio.run(probe(url, toolset, reset_timeout))[1]] - for each, toolset in toolsets(env_name, tasks).items()} + env_name, tasks, probed = tasks_and_tools(url, local="serve the bundle with agent_env.openenv.build_app first", + timeout=reset_timeout) + tools = {each: [_function(tool) for tool in found] for each, found in probed.items()} out = Path(out) out.mkdir(parents=True, exist_ok=True) for split, described in tasks.items(): diff --git a/tst/unit/training/_support.py b/tst/unit/training/_support.py index b9ef37a4..01df9602 100644 --- a/tst/unit/training/_support.py +++ b/tst/unit/training/_support.py @@ -12,15 +12,16 @@ @pytest.fixture def serve(teardowns, tools): """Serves ``tasks`` by name, every one in the split ``all`` and those ``splits`` name in theirs, as ``build_app`` - serves a bundle's under the env name ``name``; each server is stopped after the test, if the test hasn't stopped - it to wait out its episodes' teardowns.""" + serves a bundle's under the env name ``name``, each rewarded by ``reward`` of its scores, the ``replies`` step's by + default; each server is stopped after the test, if the test hasn't stopped it to wait out its episodes' + teardowns.""" servers = [] - def served(tasks, splits=None, max_concurrent_envs=8, name="dock"): + def served(tasks, splits=None, max_concurrent_envs=8, name="dock", reward=_score_at("replies")): by_name = {} for task_name, steps in tasks.items(): task = Task.from_dict({"id": task_name, "version": 1, "steps": steps}) - by_name[task_name] = ServedTask(task, task_name, EpisodeSpec.from_task(task), _score_at("replies")) + by_name[task_name] = ServedTask(task, task_name, EpisodeSpec.from_task(task), reward) named = {split: tuple(by_name[each] for each in names) for split, names in (splits or {}).items()} servers.append(Server(_app(Served(splits={"all": tuple(by_name.values()), **named}, sandbox=None), name, max_concurrent_envs))) diff --git a/tst/verl/loop_test.py b/tst/verl/loop_test.py index b680de8f..e9e72d36 100644 --- a/tst/verl/loop_test.py +++ b/tst/verl/loop_test.py @@ -1,8 +1,8 @@ """``AgentEnvAgentLoop`` on verl itself: verl's ``ToolAgentLoop`` plays a row ``write_dataset`` wrote, with Qwen3's tokenizer and chat template and a stand-in rollout server, on the unit tier's bundle served on loopback. A played row's reward is its episode's, through verl's hermes and qwen3_coder parsers; a rollout the env fails is played again on a new -episode, but not once its time is up, and one it fails every time is kept out of the loss; a prompt that doesn't fit and -an episode that ends unscored without failing raise. CI runs it on verl 0.9.1 and main.""" +episode, but not once its time is up, and one it fails every time is kept out of the loss; a prompt that doesn't fit, an +episode that ends unscored without failing and a loop run twice raise. CI runs it on verl 0.9.1 and main.""" import asyncio import json @@ -11,8 +11,8 @@ import pytest from transformers import AutoTokenizer -from agent_env.training._openenv import Failed -from agent_env.training.verl.prepare import write_dataset +from agent_env.training.verl.prepare import LOOP, write_dataset +from tst.unit.task.episode import run_test from tst.unit.training.nemo_gym_test import PLAYED from tst.util.capabilities import VERL, missing_capability_reason @@ -50,13 +50,13 @@ def test_a_played_rows_reward_is_its_episodes_and_its_last_reply_ends_the_turn(s server = serve({"a": PLAYED}) row = _row(write_dataset(server.url, tmp_path)) model = _Model(tokenizer, parser) - finish, replies = verl_loop._Session._finish, [] + finish, replies = verl_loop._Session.finish, [] - async def finishing(self, rollout, response=""): - replies.append(response) - await finish(self, rollout, response) + async def finishing(self, reply): + replies.append(reply) + return await finish(self, reply) - monkeypatch.setattr(verl_loop._Session, "_finish", finishing) + monkeypatch.setattr(verl_loop._Session, "finish", finishing) output = asyncio.run(_play(server.url, row, model, tokenizer, parser=parser)) server.stop() @@ -82,55 +82,52 @@ def test_a_tools_parameters_reach_verls_parsers_with_their_types(tokenizer): "submit_schedule", {"schedule": [1, 0, -1], "hour": 7, "note": "none today"}) -def test_a_rollout_the_env_fails_is_played_again_on_a_new_episode(serve, tmp_path, tokenizer, monkeypatch, teardowns): +def test_a_rollout_the_env_fails_is_played_again_on_a_new_episode(serve, tmp_path, tokenizer, monkeypatch): server = serve({"a": PLAYED}) row = _row(write_dataset(server.url, tmp_path)) - begin, began = verl_loop._Session._begin, [] + deploy, deploys = run_test._Env.execute, [] - async def failing_once(self, rollout, *args): - began.append(1) - if len(began) == 1: - raise Failed("the env's deploy failed") - return await begin(self, rollout, *args) + async def failing_once(self, context): + deploys.append(self.id) + if len(deploys) == 1: + raise RuntimeError("the env didn't come up") + return await deploy(self, context) - monkeypatch.setattr(verl_loop._Session, "_begin", failing_once) + monkeypatch.setattr(run_test._Env, "execute", failing_once) output = asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer)) server.stop() - assert (output.reward_score, output.extra_fields["agentenv"]["attempts"]) == (1.0, 2) - assert len(began) == 2 + assert (output.reward_score, output.extra_fields["agentenv"]["attempts"], len(deploys)) == (1.0, 2, 2) -def test_a_rollout_the_env_fails_every_time_is_kept_out_of_the_loss(serve, tmp_path, tokenizer, monkeypatch): - server = serve({"a": PLAYED}) - row = _row(write_dataset(server.url, tmp_path)) +def test_a_rollout_the_env_fails_every_time_is_kept_out_of_the_loss(serve, tmp_path, tokenizer): + healthy = serve({"a": PLAYED}) + row = _row(write_dataset(healthy.url, tmp_path)) + healthy.stop() + broken = serve({"a": [{**PLAYED[0], "fail": True}, *PLAYED[1:]]}) - async def failing(self, rollout, *args): - raise Failed("the env's deploy failed") + output = asyncio.run(_play(broken.url, row, _Model(tokenizer), tokenizer, attempts=2)) + broken.stop() - monkeypatch.setattr(verl_loop._Session, "_begin", failing) - output = asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer, attempts=2)) - server.stop() - - assert (output.reward_score, output.extra_fields["agentenv"]) == ( - 0.0, {"attempts": 2, "failure": "the env's deploy failed"}) + assert (output.reward_score, output.extra_fields["agentenv"]["attempts"]) == (0.0, 2) + assert "the env didn't come up" in output.extra_fields["agentenv"]["failure"] assert output.response_mask and set(output.response_mask) == {0} def test_a_rollout_isnt_played_again_once_its_time_is_up(serve, tmp_path, tokenizer, monkeypatch): server = serve({"a": PLAYED}) row = _row(write_dataset(server.url, tmp_path)) + deploy = run_test._Env.execute - async def failing_late(self, rollout, *args): - await asyncio.sleep(0.3) - raise Failed("the env's deploy missed its deadline") + async def hanging(self, context): + await asyncio.sleep(2) + return await deploy(self, context) - monkeypatch.setattr(verl_loop._Session, "_begin", failing_late) - output = asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer, reset_timeout=0.2)) + monkeypatch.setattr(run_test._Env, "execute", hanging) + output = asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer, reset_timeout=0.5)) server.stop() - assert (output.reward_score, output.extra_fields["agentenv"]) == ( - 0.0, {"attempts": 1, "failure": "the env's deploy missed its deadline"}) + assert (output.reward_score, output.extra_fields["agentenv"]["attempts"]) == (0.0, 1) def test_a_prompt_that_doesnt_fit_with_its_tools_is_refused_rather_than_cut(serve, tmp_path, tokenizer): @@ -143,17 +140,11 @@ def test_a_prompt_that_doesnt_fit_with_its_tools_is_refused_rather_than_cut(serv server.stop() -def test_an_episode_that_ends_without_failing_or_a_reward_raises(serve, tmp_path, tokenizer, monkeypatch): - server = serve({"a": PLAYED}) +def test_an_episode_that_ends_without_failing_or_a_reward_raises(serve, tmp_path, tokenizer): + server = serve({"a": PLAYED}, reward=lambda scores: None) row = _row(write_dataset(server.url, tmp_path)) model = _Model(tokenizer) - finish = verl_loop._Session._finish - async def unscored(self, rollout, response=""): - await finish(self, rollout, response) - rollout.reward = None - - monkeypatch.setattr(verl_loop._Session, "_finish", unscored) with pytest.raises(ValueError, match="task a of split all ended without a reward: it records no score"): asyncio.run(_play(server.url, row, model, tokenizer)) server.stop() @@ -161,6 +152,24 @@ async def unscored(self, rollout, response=""): assert model.calls == 2 +def test_an_agent_loop_plays_one_rollout(serve, tmp_path, tokenizer): + server = serve({"a": PLAYED}) + row = _row(write_dataset(server.url, tmp_path)) + + async def twice(): + played = _loop(server.url, _Model(tokenizer), tokenizer) + assert (await _run(played, row)).reward_score == 1.0 + await _run(played, row) + + with pytest.raises(RuntimeError, match="verl ran an AgentEnvAgentLoop a second time"): + asyncio.run(twice()) + server.stop() + + +def test_the_rows_config_names_this_agent_loop(): + assert LOOP == f"{verl_loop.AgentEnvAgentLoop.__module__}.{verl_loop.AgentEnvAgentLoop.__qualname__}" + + class _Model: """verl's rollout server for the loop: it adds a crate in ``parser``'s format, then replies once the tool answered.""" @@ -175,16 +184,23 @@ async def generate(self, request_id, *, prompt_ids, sampling_params, **_): return replica.TokenOutput(token_ids=self.tokenizer.encode(reply, add_special_tokens=False)) -async def _play(url, row, model, tokenizer, *, parser="hermes", prompt_length=2048, reset_timeout=60.0, attempts=3): +async def _play(url, row, model, tokenizer, **options): + return await _run(_loop(url, model, tokenizer, **options), row) + + +def _loop(url, model, tokenizer, *, parser="hermes", prompt_length=2048, reset_timeout=60.0, attempts=3): config = omegaconf.OmegaConf.create({"actor_rollout_ref": {"rollout": { "prompt_length": prompt_length, "response_length": 1024, "full_determinism": False, "multi_turn": {"format": parser, "max_user_turns": None, "max_assistant_turns": 4, "max_parallel_calls": 1, "max_tool_response_length": 2048, "tool_response_truncate_side": "middle"}, }}}) - played = verl_loop.AgentEnvAgentLoop( + return verl_loop.AgentEnvAgentLoop( trainer_config=agent_loop.DictConfigWrap(config), server_manager=model, tokenizer=tokenizer, processor=None, dataset_cls=rl_dataset.RLHFDataset, data_config=agent_loop.DictConfigWrap(DATA), hf_model_type="qwen3", url=url, reset_timeout=reset_timeout, attempts=attempts) + + +async def _run(played, row): return await played.run({"temperature": 1.0}, raw_prompt=row["prompt"], extra_info=row["extra_info"], tools_kwargs={}, index=0, agent_name="agentenv") From 3888a432a5cc06620c5150f93a8dd083a1f1cca8 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 10:57:21 -0700 Subject: [PATCH 10/13] feat(training): verl: a failed rollout is left out, a server token, a prompt check before writing, a warning on cut answers A rollout the env fails on every attempt now raises RolloutFailed: verl's V1 trainer (the default) records a session that raised as failed and leaves its trajectory out of the loss, its group's GRPO baseline and the eval reward, where the masked reward 0 still entered both. build_app(token=...) requires Authorization: Bearer on every request and WebSocket but /health, so the env server can run where other machines reach it. The session layer sends it; the verl writer and loop read it from the environment variable --token-var names, which agent_loop.yaml records, so the token is in no file. The writer with --tokenizer counts each row's prompt with its tools in the model's tokens and refuses rows over --max-prompt-length before writing; verl's own length filter doesn't count a row's tools. The loop warns once a tool when an env answer is over max_tool_response_length, which verl cuts silently. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 4 +- pyproject.toml | 4 +- src/agent_env/openenv/app.py | 35 +++++++++-- src/agent_env/training/_openenv.py | 41 ++++++++----- src/agent_env/training/verl/__main__.py | 21 +++++-- src/agent_env/training/verl/loop.py | 50 ++++++++++----- src/agent_env/training/verl/prepare.py | 75 +++++++++++++++++----- tst/unit/openenv/app_test.py | 26 +++++++- tst/unit/training/_support.py | 8 +-- tst/unit/training/sessions_test.py | 2 +- tst/unit/training/verl_test.py | 44 ++++++++++--- tst/verl/loop_test.py | 82 ++++++++++++++++++++----- uv.lock | 2 + 13 files changed, 303 insertions(+), 91 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 7367e4f1..f7618ff1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -121,8 +121,8 @@ and deserialized through a registry. | `config/` | The `Config` singleton (`get_config`, `configure`, `reset_config`) in `config/runtime.py`, file discovery in `config/loader.py`, and `load_impl`, which resolves `module:Class` pointers. `agent_env.store` re-exports the config names for compatibility. | | `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `smol_vm` (the `smol` extra), `tensorlake` (the `tensorlake` extra), `tenki` (the `tenki` extra), `vercel` (the `vercel` extra), `boxd` (the `boxd` extra), `coreweave` (the `coreweave` extra), `daytona` (the `daytona` extra), `boat` (the `boat` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store), `EnvironmentGatewayProvider` (`gateway`, an env's default), which renders a docker-compose for the gateway and its MCP servers inside the sandbox, and `EnvironmentServerProvider` (`server`), one MCP server in its own container with no gateway in front and no state store; `providers/env_state/` holds env-state providers (`local_postgres` built in). | | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | -| `openenv/` | `build_app(bundle)`: the [OpenEnv](https://github.com/openenv/OpenEnv) server for a bundle's tasks, each session playing one episode through the Episode API, with the envs' MCP tools, an added `end_turn` tool and the task's verifier score as the reward. It prepares bundles as `agent-env run` does (`prepare_bundle(played=True)`, without the agent the client plays) and keeps namespace routing on for the app's life. `openenv` comes with the `openenv` extra, a floor with no cap; this package and `training/` are the only ones that import it. Each task's description in the Task API names its `toolset` (the played agent's envs, role, and whether it gets `bash`). | -| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, `tasks_and_tools` (what the row writers read: the tasks, the split names checked, each toolset's tools) and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and masking one that fails every time; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), and raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing. | +| `openenv/` | `build_app(bundle)`: the [OpenEnv](https://github.com/openenv/OpenEnv) server for a bundle's tasks, each session playing one episode through the Episode API, with the envs' MCP tools, an added `end_turn` tool and the task's verifier score as the reward. It prepares bundles as `agent-env run` does (`prepare_bundle(played=True)`, without the agent the client plays) and keeps namespace routing on for the app's life. `openenv` comes with the `openenv` extra, a floor with no cap; this package and `training/` are the only ones that import it. Each task's description in the Task API names its `toolset` (the played agent's envs, role, and whether it gets `bash`). With `token`, every request and WebSocket but `/health` must carry `Authorization: Bearer `, for a server others reach. | +| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, `tasks_and_tools` (what the row writers read: the tasks, the split names checked, each toolset's tools) and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and raising `RolloutFailed` for one that fails every time, which verl's V1 trainer (the default) leaves out of the loss, the group's baseline and the eval reward; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing, and warns once a tool of an answer verl cuts to `max_tool_response_length`. The writer reads a server that asks for a token with the one `--token-var` names, which the loop reads in each worker, and with `--tokenizer` counts each prompt with its tools, refusing rows over `--max-prompt-length` before writing. | | `runner/` | The `[runner]` seam: `Runner.submit()` returns a `RunHandle` (`run_id`, `instance_id`) before the run starts; `LocalRunner` is built in. `runner/store.py` is the run index every runner writes and the explorer reads. | | `bundle/` | A bundle is a folder of entities (`envs/`, `agents/`, `artifacts/`, `skills/`, `tasks/.json`, `evals/.toml`) that runs without a registry: `agent-env run` plans it, writes it into the `@local` namespace (reusing each version whose inputs haven't changed) and runs its evals or tasks; `agent-env bundle init` writes a stored task or eval as one, with every env, agent and artifact it uses in its own folder (each image as its build context, pruned to what its Dockerfile reads and pinned to the Python and npm versions of the stored image, or, given `--image-ref ID=REF`, as that registry image pinned by digest with its context beside it, `image = { ref, dockerfile }`, which a run pulls and builds only when it can't; `task_images` lists the images an import writes, downloading nothing, so an export can push them first) so it runs with no store; `agent-env bundle register` writes a folder into the configured stores under registry ids (``), each task and eval naming the versions it reads and each image it builds pushed, through a ledger in the target store so registering again writes only what changed. | | `eval/` | `Eval` (`eval`), a versioned list of task references, and its store. | diff --git a/pyproject.toml b/pyproject.toml index 5eaaa1ba..08c569f8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,10 +70,12 @@ nemo-gym = [ "omegaconf>=2.3; python_full_version >= '3.13.14'", ] # agent_env.training.verl: verl's trainers on a bundle served as an OpenEnv env; installed into verl's own environment, -# whose verl, vLLM and torch are the trainer's to pin. pyarrow writes the parquet rows verl reads. +# whose verl, vLLM and torch are the trainer's to pin. pyarrow writes the parquet rows verl reads; transformers counts a +# row's prompt with its tools as the model's chat template renders them. verl = [ "agentenv-framework[openenv]", "pyarrow>=19", + "transformers>=5", ] # The AWS store backends (S3, Secrets Manager, DynamoDB, ECR), selected by impl pointer in config.toml. aws = [ diff --git a/src/agent_env/openenv/app.py b/src/agent_env/openenv/app.py index 71215307..becd7fd5 100644 --- a/src/agent_env/openenv/app.py +++ b/src/agent_env/openenv/app.py @@ -3,12 +3,16 @@ from __future__ import annotations import functools +import hmac import logging from collections.abc import Callable, Mapping from contextlib import ExitStack from pathlib import Path from fastapi import FastAPI +from starlette.responses import JSONResponse +from starlette.types import ASGIApp, Receive, Scope, Send +from starlette.websockets import WebSocketClose from openenv.core.env_server.http_server import create_app from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation from openenv.core.env_server.types import Observation @@ -31,6 +35,7 @@ def build_app( reward: str | Callable[[dict[str, float]], float] | None = None, sandbox: str | None = None, max_concurrent_envs: int = 8, + token: str | None = None, ) -> FastAPI: """The OpenEnv server for ``bundle``'s tasks: a folder or an installed bundle's name, as ``agent-env run BUNDLE`` reads it, served as the split ``all`` and a split per eval; or a mapping of split name to bundle. Each session @@ -40,19 +45,23 @@ def build_app( task that records one score has that one. Each bundle is prepared as ``agent-env run`` prepares it before its first task, without the agent a client plays, so every task must pass its checks; one that can't run as an episode is then left out and logged. Namespace routing is on from the app's startup to its shutdown, which tears - down every open episode. Raises BundleError when a bundle can't be prepared or serves no task, and ValueError + down every open episode. With ``token``, every request and WebSocket but ``/health`` must carry + ``Authorization: Bearer ``, for a server others can reach: without it, anyone who reaches the server + deploys and drives its envs. Raises BundleError when a bundle can't be prepared or serves no task, and ValueError when ``reward`` can't score every task.""" with namespace_routing(): prepared = {split: _prepared(source, sandbox) for split, source in _sources(bundle).items()} served = Served(splits=_served(_splits(bundle, prepared), reward), sandbox=sandbox) - return _app(served, _env_name(prepared), max_concurrent_envs) + return _app(served, _env_name(prepared), max_concurrent_envs, token) -def _app(served: Served, env_name: str, max_concurrent_envs: int) -> FastAPI: +def _app(served: Served, env_name: str, max_concurrent_envs: int, token: str | None = None) -> FastAPI: """OpenEnv's server for what ``served`` serves, with namespace routing on from its startup to its shutdown, which - tears down every open episode.""" + tears down every open episode, and ``token`` required of every request but ``/health`` when given.""" app = create_app(functools.partial(AgentEnvEnvironment, served), CallToolAction, CallToolObservation, env_name=env_name, max_concurrent_envs=max_concurrent_envs, reset_observation_cls=Observation) + if token: + app.add_middleware(_Bearer, token=token) routing = ExitStack() async def route() -> None: @@ -69,6 +78,24 @@ async def tear_down() -> None: return app +class _Bearer: + """Refuses every request and WebSocket handshake but ``/health``'s that doesn't carry ``Authorization: Bearer + ``: a request with 401, a WebSocket by closing it before it opens, which its client sees as HTTP 403.""" + + def __init__(self, app: ASGIApp, token: str): + self.app, self.expected = app, f"Bearer {token}".encode() + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] in ("http", "websocket") and scope["path"] != "/health": + given = dict(scope["headers"]).get(b"authorization", b"") + if not hmac.compare_digest(given, self.expected): + refusal = (JSONResponse({"detail": "this server needs Authorization: Bearer "}, 401) + if scope["type"] == "http" else WebSocketClose(code=1008)) + await refusal(scope, receive, send) + return + await self.app(scope, receive, send) + + def _sources(bundle: _Source | Mapping[str, _Source]) -> dict[str | None, _Source]: return dict(bundle) if isinstance(bundle, Mapping) else {None: bundle} diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index 98510e09..07698698 100644 --- a/src/agent_env/training/_openenv.py +++ b/src/agent_env/training/_openenv.py @@ -58,14 +58,15 @@ class NamedToolset: @dataclass(frozen=True) class Shared: """What a session needs of its server: where it is, the tools the model is shown, which an episode's first turn - must offer no more than, the deadlines, and how long to wait for room on the server. The key makes it equal to its - pickled copies only.""" + must offer no more than, the deadlines, how long to wait for room on the server, and the token it asks for. The key + makes it equal to its pickled copies only.""" url: str offered: frozenset[str] tool_timeout: float reset_timeout: float patience: float = REOPEN_PATIENCE + token: str | None = field(default=None, repr=False) key: str = field(default_factory=lambda: uuid.uuid4().hex) @@ -123,7 +124,8 @@ async def _start(self, split: str, index: int, expected: str | None) -> None: await asyncio.wait([rollout.started]) self._rollout = None if self.opens_at_start and not self._opened: - self._client, self._opened = await opened(self._shared.url, self._shared.reset_timeout), True + self._client, self._opened = await opened(self._shared.url, self._shared.reset_timeout, + token=self._shared.token), True rollout = Rollout() rollout.started = asyncio.ensure_future(self._begin(rollout, split, index, expected)) self._rollout = rollout @@ -131,7 +133,7 @@ async def _start(self, split: str, index: int, expected: str | None) -> None: async def _begin(self, rollout: Rollout, split: str, index: int, expected: str | None) -> None: if self._client is None: self._client = await reopened(self._shared.url, self._shared.reset_timeout, dropped=self._opened, - patience=self._shared.patience) + patience=self._shared.patience, token=self._shared.token) self._opened = True reset = await self._client.reset(split=split, index=index, timeout_s=self._shared.reset_timeout) if reset.done: @@ -251,10 +253,11 @@ def stop(self, timeout: float = STOP_SECONDS) -> bool: return not self._thread.is_alive() -async def opened(url: str, timeout: float) -> GenericEnvClient: +async def opened(url: str, timeout: float, *, token: str | None = None) -> GenericEnvClient: """A session on the server, opened now rather than at the first reset, so a server at capacity refuses it at once; - each message waits ``GRACE_SECONDS`` longer than the ``timeout`` the server ends a call at.""" - client = GenericEnvClient(base_url=url, message_timeout_s=timeout + GRACE_SECONDS) + each message waits ``GRACE_SECONDS`` longer than the ``timeout`` the server ends a call at. ``token`` goes with it + as ``Authorization: Bearer``, for a server that asks for one.""" + client = GenericEnvClient(base_url=url, message_timeout_s=timeout + GRACE_SECONDS, headers=authorization(token)) try: await client.state() except Exception as e: @@ -266,7 +269,8 @@ async def opened(url: str, timeout: float) -> GenericEnvClient: return client -async def reopened(url: str, timeout: float, *, dropped: bool, patience: float) -> GenericEnvClient: +async def reopened(url: str, timeout: float, *, dropped: bool, patience: float, + token: str | None = None) -> GenericEnvClient: """A session, tried again every ``REOPEN_SECONDS`` for ``patience`` seconds while the server is at its limit or too busy to answer, in place of a ``dropped`` one or for a rollout's own; raises Failed, which leaves the rollout unscored, when it can't open one, at once unless the server was at its limit or too busy to answer.""" @@ -274,7 +278,7 @@ async def reopened(url: str, timeout: float, *, dropped: bool, patience: float) deadline = time.monotonic() + patience while True: try: - return await opened(url, timeout) + return await opened(url, timeout, token=token) except AtCapacity as e: if time.monotonic() + REOPEN_SECONDS > deadline: raise Failed(f"{lead}the server stayed at its limit of {e.limit} sessions for " @@ -301,12 +305,13 @@ def _transient(error: BaseException | None) -> bool: return False -async def probe(url: str, toolset: NamedToolset, timeout: float) -> tuple[frozenset[str], tuple[dict[str, Any], ...]]: +async def probe(url: str, toolset: NamedToolset, timeout: float, *, + token: str | None = None) -> tuple[frozenset[str], tuple[dict[str, Any], ...]]: """The tools ``toolset``'s first task offers at its first turn, and every tool any of its turns offers, found by walking it through its turns, ending each but the last empty and listing its tools, then ending the episode unscored. Each call on the probe's session has the deadline ``timeout``.""" split, index, turns = toolset.first - client = await opened(url, timeout) + client = await opened(url, timeout, token=token) first, tools = None, {} try: turn = await client.reset(split=split, index=index, timeout_s=timeout) @@ -348,10 +353,10 @@ async def end_turn(client: GenericEnvClient, response: str, deadline: float) -> return await call_tool(client, END_TURN, {"response": response}, deadline) -def read_tasks(url: str, *, local: str) -> tuple[str, dict[str, list[dict[str, Any]]]]: +def read_tasks(url: str, *, local: str, token: str | None = None) -> tuple[str, dict[str, list[dict[str, Any]]]]: """The server's env name, and each split's tasks as its Task API describes them. ``local`` says how to serve the bundle in this process instead, for when no server answers.""" - with httpx.Client(base_url=url, timeout=CONNECT_SECONDS) as http: + with httpx.Client(base_url=url, timeout=CONNECT_SECONDS, headers=authorization(token)) as http: try: [name] = http.get("/list_environments").raise_for_status().json() except httpx.TransportError as e: @@ -366,16 +371,16 @@ def read_tasks(url: str, *, local: str) -> tuple[str, dict[str, list[dict[str, A return name, tasks -def tasks_and_tools(url: str, *, local: str, timeout: float +def tasks_and_tools(url: str, *, local: str, timeout: float, token: str | None = None ) -> tuple[str, dict[str, list[dict[str, Any]]], dict[str, tuple[dict[str, Any], ...]]]: """What a writer of a split's rows reads: the server's env name, each split's tasks, refusing a split whose name can't name a data file, and every tool each toolset's turns offer, by the toolset's key, found by playing its first task through (``probe``), so reading deploys those.""" - name, tasks = read_tasks(url, local=local) + name, tasks = read_tasks(url, local=local, token=token) if unfit := [split for split in tasks if not _FILE.fullmatch(split)]: raise ValueError(f"the server's splits {unfit} can't name a data file: a split's name must be letters, " "digits, '.', '_' and '-', not starting with '.'") - return name, tasks, {each: asyncio.run(probe(url, toolset, timeout))[1] + return name, tasks, {each: asyncio.run(probe(url, toolset, timeout, token=token))[1] for each, toolset in toolsets(name, tasks).items()} @@ -402,6 +407,10 @@ def toolsets(name: str, tasks: dict[str, list[dict[str, Any]]]) -> dict[str, Nam return named +def authorization(token: str | None) -> dict[str, str]: + return {"Authorization": f"Bearer {token}"} if token else {} + + def key(task: Mapping[str, Any]) -> str: return json.dumps(task["toolset"], sort_keys=True) diff --git a/src/agent_env/training/verl/__main__.py b/src/agent_env/training/verl/__main__.py index 23ed4d8d..f187204c 100644 --- a/src/agent_env/training/verl/__main__.py +++ b/src/agent_env/training/verl/__main__.py @@ -12,14 +12,25 @@ help="One tool call's deadline in seconds; the server ends a call that misses it, and the episode.") @click.option("--reset-timeout", type=float, default=600.0, show_default=True, help="How long an episode may take to deploy, and each turn to be ended, in seconds; also how long a " - "rollout waits for room on the server.") + "rollout waits for room on the server, and how long after it began it may be played again.") @click.option("--attempts", type=int, default=3, show_default=True, - help="How many times a rollout the env fails is played, each time from the start on a new episode.") -def main(url: str, out: str, tool_timeout: float, reset_timeout: float, attempts: int) -> None: + help="How many times a rollout the env fails is played at most, each time from the start on a new " + "episode.") +@click.option("--token-var", help="The environment variable that holds the token the server asks for, read now and in " + "each of verl's workers.") +@click.option("--tokenizer", help="The model's tokenizer, as transformers loads it, to count each prompt with its " + "tools.") +@click.option("--max-prompt-length", type=int, + help="Refuse rows whose prompt with its tools is longer, in tokens: verl's data.max_prompt_length.") +def main(url: str, out: str, tool_timeout: float, reset_timeout: float, attempts: int, token_var: str | None, + tokenizer: str | None, max_prompt_length: int | None) -> None: """Writes the tasks of the OpenEnv server at URL as OUT/.parquet, a row per task, and OUT/agent_loop.yaml, the agent loop that plays them on the server.""" - folder = write_dataset(url, out, tool_timeout=tool_timeout, reset_timeout=reset_timeout, attempts=attempts) - click.echo(f"wrote {', '.join(sorted(path.name for path in folder.iterdir()))} to {folder}") + written = write_dataset(url, out, tool_timeout=tool_timeout, reset_timeout=reset_timeout, attempts=attempts, + token_var=token_var, tokenizer=tokenizer, max_prompt_length=max_prompt_length) + click.echo(f"wrote {', '.join(sorted(path.name for path in written.folder.iterdir()))} to {written.folder}") + if written.longest_prompt is not None: + click.echo(f"the longest prompt is {written.longest_prompt} tokens with its tools") if __name__ == "__main__": diff --git a/src/agent_env/training/verl/loop.py b/src/agent_env/training/verl/loop.py index 57cfd3ad..ca92a1e5 100644 --- a/src/agent_env/training/verl/loop.py +++ b/src/agent_env/training/verl/loop.py @@ -24,23 +24,34 @@ "'agentenv-framework[verl]' into verl's environment") from exc from agent_env.training._openenv import Rollout, Session, Shared +from agent_env.training.verl.prepare import token_in logger = logging.getLogger(__name__) +#: The tools whose answers verl was seen to cut, each warned of once a process. +_cut: set[str] = set() + + +class RolloutFailed(RuntimeError): + """The env failed a rollout on every attempt. verl's V1 trainer (``trainer.use_v1``, the default) records the + session as failed and leaves the rollout out: out of the loss, its group's baseline and the eval reward.""" class AgentEnvAgentLoop(ToolAgentLoop): """Plays a row's task as one episode on the OpenEnv server at ``url``. A rollout the env fails, a deploy, a call past its deadline or a server that went away, is played again from the start on a new episode, a fresh sample of the same policy, up to ``attempts`` times, but not once ``reset_timeout`` seconds have passed since it began; one - that fails every time is kept out of the loss, its response masked, with the reward 0. A task that ends unscored - without failing, and a prompt that doesn't fit ``rollout.prompt_length`` with its tools, are configuration errors no - replay changes, and raise. ``agent_loop.yaml`` sets the arguments; verl passes the rest. + that fails every time raises RolloutFailed, so verl's V1 trainer leaves it out of training and of the eval reward. + A task that ends unscored without failing, and a prompt that doesn't fit ``rollout.prompt_length`` with its tools, + are configuration errors no replay changes, and raise too; the V1 trainer logs each such session and leaves it out. + An env answer longer than ``multi_turn.max_tool_response_length``, which verl cuts, is warned of once a tool. + ``agent_loop.yaml`` sets the arguments; verl passes the rest. It leans on verl internals, which CI checks on verl 0.9.1 and main: verl builds an agent loop per trajectory, so an instance holds one rollout's state, and ``run`` refuses a second; ``ToolAgentLoop.run`` reads the tools from - ``self.tools`` and ``self.tool_schemas``; ``_build_assistant_message`` is overridden to keep the reply, and - ``AgentLoopBase._cap_text_prompt_length`` to refuse a prompt verl would cut; and the tools are verl's - ``FunctionTool`` with ``OpenAIFunctionToolSchema``. + ``self.tools`` and ``self.tool_schemas``, and its limit on an answer from ``self.max_tool_response_length``; + ``_build_assistant_message`` is overridden to keep the reply, and ``AgentLoopBase._cap_text_prompt_length`` to + refuse a prompt verl would cut; the tools are verl's ``FunctionTool`` with ``OpenAIFunctionToolSchema``; and the V1 + trainer leaves a session that raised out of the batch. Args: url: The OpenEnv server's base URL. @@ -48,12 +59,13 @@ class AgentEnvAgentLoop(ToolAgentLoop): reset_timeout: How long an episode may take to deploy, and each turn to be ended, in seconds; also how long a rollout waits for room on a server at its limit, and how long after it began it may be played again. attempts: How many times a rollout the env fails is played at most. + token_var: The environment variable that holds the token a server given one asks for, set in every worker. """ def __init__(self, *args: Any, url: str, tool_timeout: float = 60.0, reset_timeout: float = 600.0, - attempts: int = 3, **kwargs: Any): + attempts: int = 3, token_var: str | None = None, **kwargs: Any): super().__init__(*args, **kwargs) - self._url, self._attempts = url.rstrip("/"), attempts + self._url, self._attempts, self._token = url.rstrip("/"), attempts, token_in(token_var) self._tool_timeout, self._reset_timeout = tool_timeout, reset_timeout self._row: dict[str, Any] = {} self._reply = "" @@ -65,7 +77,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop row = self._row = kwargs["extra_info"] schemas = json.loads(row["tools"]) shared = Shared(self._url, frozenset(schema["function"]["name"] for schema in schemas), self._tool_timeout, - self._reset_timeout, patience=self._reset_timeout) + self._reset_timeout, patience=self._reset_timeout, token=self._token) began = time.monotonic() for attempt in range(1, self._attempts + 1): session = _Session(shared) @@ -83,15 +95,14 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop logger.warning("%s, attempt %d of %d: %s", self._task(), attempt, self._attempts, rollout.failure) if time.monotonic() - began >= self._reset_timeout: break - output.response_mask = [0] * len(output.response_mask) - output.reward_score = 0.0 - output.extra_fields["agentenv"] = {"attempts": attempt, "failure": rollout.failure} - return output + raise RolloutFailed(f"{self._task()}: the env failed it on {attempt} attempt{'s' * (attempt > 1)}, so it's left " + f"out of training: {rollout.failure}") async def _play(self, session: _Session, schemas: list[dict[str, Any]], sampling_params: dict[str, Any], kwargs: dict[str, Any]) -> tuple[AgentLoopOutput, Rollout]: await session.start(self._row) - self.tools = {schema["function"]["name"]: session.tool(schema) for schema in schemas} + self.tools = {schema["function"]["name"]: session.tool(schema, self.max_tool_response_length) + for schema in schemas} self.tool_schemas = schemas self._reply = "" output = await super().run(sampling_params, **kwargs) @@ -131,11 +142,18 @@ def __init__(self, shared: Shared): async def start(self, row: dict[str, Any]) -> None: await self._start(row["split"], int(row["index"]), row["task_id"]) - def tool(self, schema: dict[str, Any]) -> FunctionTool: + def tool(self, schema: dict[str, Any], limit: int) -> FunctionTool: + """The tool as verl calls it, warning once a process when its answer is over ``limit``, where verl cuts it.""" name = schema["function"]["name"] async def call(**arguments: Any) -> str: - return await self._step(name, arguments) + answer = await self._step(name, arguments) + if len(answer) > limit and name not in _cut: + _cut.add(name) + logger.warning("%s answered %d characters, which verl cuts to " + "actor_rollout_ref.rollout.multi_turn.max_tool_response_length=%d, so the model doesn't " + "see all of it: raise it to cover the env's answers", name, len(answer), limit) + return answer return FunctionTool(name=name, fn=call, tool_schema=_parser_schema(schema), is_async=True) diff --git a/src/agent_env/training/verl/prepare.py b/src/agent_env/training/verl/prepare.py index 30b560dc..5d1dc87f 100644 --- a/src/agent_env/training/verl/prepare.py +++ b/src/agent_env/training/verl/prepare.py @@ -5,12 +5,15 @@ from __future__ import annotations import json +import os +from dataclasses import dataclass from pathlib import Path from typing import Any import pyarrow as pa import pyarrow.parquet as pq import yaml +from transformers import AutoTokenizer from agent_env.training._openenv import key, messages, tasks_and_tools @@ -27,34 +30,76 @@ ]) +@dataclass(frozen=True) +class Written: + """What ``write_dataset`` wrote: the folder, and the longest prompt in tokens with its tools when it counted + them.""" + + folder: Path + longest_prompt: int | None = None + + def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, reset_timeout: float = 600.0, - attempts: int = 3) -> Path: + attempts: int = 3, token_var: str | None = None, tokenizer: str | None = None, + max_prompt_length: int | None = None) -> Written: """The tasks of the OpenEnv server at ``url`` as ``out/.parquet``, a row per task of each split, and ``out/agent_loop.yaml``, which points the rows' agent loop at ``url`` with the deadlines and the ``attempts`` a - rollout the env fails is played; returns ``out``. A row's prompt is its task's system prompt and first prompt, and - its tools every tool its task's turns offer, found by playing one task of each toolset to its last turn, unscored, - so writing deploys those.""" + rollout the env fails is played. A row's prompt is its task's system prompt and first prompt, and its tools every + tool its task's turns offer, found by playing one task of each toolset to its last turn, unscored, so writing + deploys those. + + ``token_var`` names the environment variable that holds the token a server given one asks for + (``build_app(token=...)``): the writer reads it now and the loop in each of verl's workers, so the token is in no + file. With ``tokenizer``, the model's as transformers loads it, each row's prompt is counted with its tools as the + model's chat template renders them, and a row over ``max_prompt_length`` is refused before anything is written, + where verl would cut it from the left; verl's own length filter doesn't count a row's tools.""" if attempts < 1: raise ValueError(f"attempts must be at least 1, not {attempts}") + if max_prompt_length is not None and tokenizer is None: + raise ValueError("max_prompt_length is in tokens, so it takes the model's tokenizer") url = url.rstrip("/") env_name, tasks, probed = tasks_and_tools(url, local="serve the bundle with agent_env.openenv.build_app first", - timeout=reset_timeout) + timeout=reset_timeout, token=token_in(token_var)) tools = {each: [_function(tool) for tool in found] for each, found in probed.items()} + rows = {split: [{ + "data_source": f"agentenv/{env_name}", + "prompt": messages(task["system_prompt"], task["prompt"]), + "agent_name": AGENT, + "extra_info": {"split": split, "index": task["index"], "task_id": task["id"], + "tools": json.dumps(tools[key(task)])}, + } for task in described] for split, described in tasks.items()} + longest = None if tokenizer is None else _longest(rows, AutoTokenizer.from_pretrained(tokenizer), max_prompt_length) out = Path(out) out.mkdir(parents=True, exist_ok=True) - for split, described in tasks.items(): - pq.write_table(pa.Table.from_pylist([{ - "data_source": f"agentenv/{env_name}", - "prompt": messages(task["system_prompt"], task["prompt"]), - "agent_name": AGENT, - "extra_info": {"split": split, "index": task["index"], "task_id": task["id"], - "tools": json.dumps(tools[key(task)])}, - } for task in described], schema=ROW), out / f"{split}.parquet") + for split, written in rows.items(): + pq.write_table(pa.Table.from_pylist(written, schema=ROW), out / f"{split}.parquet") (out / "agent_loop.yaml").write_text(yaml.safe_dump([{ "name": AGENT, "_target_": LOOP, "url": url, "tool_timeout": tool_timeout, "reset_timeout": reset_timeout, - "attempts": attempts, + "attempts": attempts, **({"token_var": token_var} if token_var else {}), }], sort_keys=False)) - return out + return Written(out, longest) + + +def token_in(var: str | None) -> str | None: + """The token in the environment variable ``var``, which must be set when it's named.""" + if var is None: + return None + if not (token := os.environ.get(var)): + raise ValueError(f"{var} isn't set: it holds the token the OpenEnv server asks for") + return token + + +def _longest(rows: dict[str, list[dict[str, Any]]], tokenizer: Any, limit: int | None) -> int: + """The longest row's prompt in tokens with its tools, refusing rows over ``limit``.""" + lengths = {(split, row["extra_info"]["task_id"]): len(tokenizer.apply_chat_template( + row["prompt"], tools=json.loads(row["extra_info"]["tools"]), add_generation_prompt=True, tokenize=True, + return_dict=True)["input_ids"]) for split, written in rows.items() for row in written} + if limit is not None and (over := {at: length for at, length in lengths.items() if length > limit}): + (split, task), length = max(over.items(), key=lambda item: item[1]) + raise ValueError(f"{len(over)} rows' prompts are over max_prompt_length={limit} tokens with their tools, the " + f"longest task {task} of split {split} at {length}, which verl would cut from the left: raise " + "it, and data.max_prompt_length with it") + return max(lengths.values(), default=0) def _function(tool: dict[str, Any]) -> dict[str, Any]: diff --git a/tst/unit/openenv/app_test.py b/tst/unit/openenv/app_test.py index d2d7aff4..5ed1ab9b 100644 --- a/tst/unit/openenv/app_test.py +++ b/tst/unit/openenv/app_test.py @@ -1,6 +1,6 @@ """``build_app`` over bundle folders, served in process: the splits it serves and the tasks it leaves out, the reward -and bundles it refuses, an episode played over OpenEnv's own WebSocket after a step before reset is refused, and -namespace routing for each of the app's runs.""" +and bundles it refuses, an episode played over OpenEnv's own WebSocket after a step before reset is refused, +namespace routing for each of the app's runs, and the token a server given one asks of every request and WebSocket.""" import json import logging @@ -9,6 +9,7 @@ import pytest from fastapi.testclient import TestClient from pydantic import ValidationError +from starlette.websockets import WebSocketDisconnect from agent_env.bundle import BundleError from agent_env.openenv import build_app @@ -138,3 +139,24 @@ def test_a_dotted_bundle_name_is_served_dashed_as_openenvs_web_interface_reads_a with TestClient(build_app(root)) as client: assert [split["name"] for split in client.get("/dock-v1/splits").json()] == ["all"] + + +def test_a_server_given_a_token_asks_it_of_every_request_and_websocket_but_health(bundles): + root = bundles("dock", {"tasks/a.json": PLAYED}) + + with TestClient(build_app(str(root), token="s3cret")) as client: + bare = client.get("/list_environments") + wrong = client.get("/list_environments", headers={"Authorization": "Bearer guess"}) + right = client.get("/list_environments", headers={"Authorization": "Bearer s3cret"}) + health = client.get("/health") + with pytest.raises(WebSocketDisconnect) as refused, client.websocket_connect("/ws"): + pass + with client.websocket_connect("/ws", headers={"Authorization": "Bearer s3cret"}) as ws: + ws.send_json({"type": "reset", "data": {"split": "all", "index": 0}}) + reset = ws.receive_json()["data"] + + assert (bare.status_code, wrong.status_code, right.status_code, health.status_code) == (401, 401, 200, 200) + assert bare.json() == {"detail": "this server needs Authorization: Bearer "} + assert right.json() == ["dock"] + assert refused.value.code == 1008 + assert reset["observation"]["metadata"]["prompt"] == "Add a crate." diff --git a/tst/unit/training/_support.py b/tst/unit/training/_support.py index 01df9602..157db065 100644 --- a/tst/unit/training/_support.py +++ b/tst/unit/training/_support.py @@ -13,18 +13,18 @@ def serve(teardowns, tools): """Serves ``tasks`` by name, every one in the split ``all`` and those ``splits`` name in theirs, as ``build_app`` serves a bundle's under the env name ``name``, each rewarded by ``reward`` of its scores, the ``replies`` step's by - default; each server is stopped after the test, if the test hasn't stopped it to wait out its episodes' - teardowns.""" + default, asking for ``token`` when given; each server is stopped after the test, if the test hasn't stopped it to + wait out its episodes' teardowns.""" servers = [] - def served(tasks, splits=None, max_concurrent_envs=8, name="dock", reward=_score_at("replies")): + def served(tasks, splits=None, max_concurrent_envs=8, name="dock", reward=_score_at("replies"), token=None): by_name = {} for task_name, steps in tasks.items(): task = Task.from_dict({"id": task_name, "version": 1, "steps": steps}) by_name[task_name] = ServedTask(task, task_name, EpisodeSpec.from_task(task), reward) named = {split: tuple(by_name[each] for each in names) for split, names in (splits or {}).items()} servers.append(Server(_app(Served(splits={"all": tuple(by_name.values()), **named}, sandbox=None), name, - max_concurrent_envs))) + max_concurrent_envs, token))) return servers[-1] yield served diff --git a/tst/unit/training/sessions_test.py b/tst/unit/training/sessions_test.py index 381247b6..73658e97 100644 --- a/tst/unit/training/sessions_test.py +++ b/tst/unit/training/sessions_test.py @@ -18,7 +18,7 @@ def test_a_session_is_tried_again_while_its_server_is_busy_and_given_up_at_once_ opens, attempts): tried = [] - async def flaky(url, timeout): + async def flaky(url, timeout, *, token=None): tried.append(url) if len(tried) < 3: client_error = ConnectionError(f"Failed to connect to {url}/ws: {cause}") diff --git a/tst/unit/training/verl_test.py b/tst/unit/training/verl_test.py index 515a682b..a5132ea3 100644 --- a/tst/unit/training/verl_test.py +++ b/tst/unit/training/verl_test.py @@ -1,8 +1,10 @@ """``write_dataset`` on the OpenEnv server ``build_app`` makes, served on loopback over the Episode API's fake steps: -verl's parquet rows for each split and the config of the agent loop that plays them, and the command that writes both.""" +verl's parquet rows for each split and the config of the agent loop that plays them, from a server that asks for a +token too, and the command that writes both.""" import json +import httpx import pyarrow.parquet as pq import pytest import yaml @@ -34,10 +36,11 @@ def test_a_servers_tasks_become_a_row_per_task_of_each_split_and_the_agent_loop_ server = serve({"a": PLAYED, "b": [ITEMS, ALICE, {**TURN, "prompt": "Add a barrel."}, REPLIES]}, {"smoke": ["b"]}, name="Dock Yard") - out = write_dataset(f"{server.url}/", tmp_path / "data", tool_timeout=30.0, reset_timeout=90.0, attempts=2) + written = write_dataset(f"{server.url}/", tmp_path / "data", tool_timeout=30.0, reset_timeout=90.0, attempts=2) server.stop() - assert out == tmp_path / "data" + out = written.folder + assert (out, written.longest_prompt) == (tmp_path / "data", None) assert yaml.safe_load((out / "agent_loop.yaml").read_text()) == [{ "name": "agentenv", "_target_": LOOP, "url": server.url, "tool_timeout": 30.0, "reset_timeout": 90.0, "attempts": 2, @@ -58,7 +61,7 @@ def test_each_toolset_is_probed_once_and_a_rows_tools_are_its_tasks(serve, tmp_p _tool_sessions(monkeypatch, {ITEMS_MCP: ["add_item"], PANTRY_MCP: ["take_item"]}) server = serve({"a": PLAYED, "p": [PANTRY, {**ALICE, "env_ids": ["pantry"]}, TURN, REPLIES], "c": PLAYED}) - out = write_dataset(server.url, tmp_path) + out = write_dataset(server.url, tmp_path).folder server.stop() assert [[tool["function"]["name"] for tool in row["extra_info"]["tools"]] for row in _rows(out / "all.parquet")] == [ @@ -69,7 +72,7 @@ def test_each_toolset_is_probed_once_and_a_rows_tools_are_its_tasks(serve, tmp_p def test_a_split_the_server_keeps_empty_has_a_rows_columns(serve, tmp_path): server = serve({"a": PLAYED}, {"empty": []}) - out = write_dataset(server.url, tmp_path) + out = write_dataset(server.url, tmp_path).folder server.stop() empty = pq.read_table(out / "empty.parquet") @@ -86,9 +89,34 @@ def test_a_split_whose_name_isnt_a_file_name_is_refused_before_anything_is_writt assert list(tmp_path.iterdir()) == [] -def test_a_rollout_is_played_at_least_once(tmp_path): - with pytest.raises(ValueError, match="attempts must be at least 1, not 0"): - write_dataset("http://127.0.0.1:1", tmp_path / "data", attempts=0) +def test_a_server_that_asks_for_a_token_is_read_with_the_one_its_variable_holds_and_the_config_names_it( + serve, tmp_path, monkeypatch, teardowns): + server = serve({"a": PLAYED}, token="s3cret") + monkeypatch.setenv("DOCK_TOKEN", "s3cret") + + out = write_dataset(server.url, tmp_path / "data", token_var="DOCK_TOKEN").folder + monkeypatch.setenv("DOCK_TOKEN", "guess") + with pytest.raises(httpx.HTTPStatusError, match="401 Unauthorized"): + write_dataset(server.url, tmp_path / "refused", token_var="DOCK_TOKEN") + server.stop() + + [config] = yaml.safe_load((out / "agent_loop.yaml").read_text()) + assert config["token_var"] == "DOCK_TOKEN" + assert "s3cret" not in (out / "agent_loop.yaml").read_text() + assert len(_rows(out / "all.parquet")) == 1 + assert not (tmp_path / "refused").exists() + + +@pytest.mark.parametrize(("options", "message"), [ + ({"attempts": 0}, "attempts must be at least 1, not 0"), + ({"max_prompt_length": 1024}, "max_prompt_length is in tokens, so it takes the model's tokenizer"), + ({"token_var": "UNSET_DOCK_TOKEN"}, "UNSET_DOCK_TOKEN isn't set: it holds the token the OpenEnv server asks for"), +]) +def test_what_cant_be_written_is_refused_before_the_server_is_read(tmp_path, monkeypatch, options, message): + monkeypatch.delenv("UNSET_DOCK_TOKEN", raising=False) + + with pytest.raises(ValueError, match=message): + write_dataset("http://127.0.0.1:1", tmp_path / "data", **options) assert list(tmp_path.iterdir()) == [] diff --git a/tst/verl/loop_test.py b/tst/verl/loop_test.py index e9e72d36..fd0cf7b9 100644 --- a/tst/verl/loop_test.py +++ b/tst/verl/loop_test.py @@ -1,11 +1,15 @@ """``AgentEnvAgentLoop`` on verl itself: verl's ``ToolAgentLoop`` plays a row ``write_dataset`` wrote, with Qwen3's tokenizer and chat template and a stand-in rollout server, on the unit tier's bundle served on loopback. A played row's -reward is its episode's, through verl's hermes and qwen3_coder parsers; a rollout the env fails is played again on a new -episode, but not once its time is up, and one it fails every time is kept out of the loss; a prompt that doesn't fit, an -episode that ends unscored without failing and a loop run twice raise. CI runs it on verl 0.9.1 and main.""" +reward is its episode's, through verl's hermes and qwen3_coder parsers and from a server that asks for a token; a +rollout the env fails is played again on a new episode, but not once its time is up, and one it fails every time raises, +which verl's V1 trainer leaves out; a prompt that doesn't fit, an episode that ends unscored without failing and a loop +run twice raise; an answer verl cuts is warned of; and the writer counts prompts with their tools in the model's tokens. +CI runs it on verl 0.9.1 and main.""" import asyncio import json +import logging +import re import pyarrow.parquet as pq import pytest @@ -100,19 +104,17 @@ async def failing_once(self, context): assert (output.reward_score, output.extra_fields["agentenv"]["attempts"], len(deploys)) == (1.0, 2, 2) -def test_a_rollout_the_env_fails_every_time_is_kept_out_of_the_loss(serve, tmp_path, tokenizer): +def test_a_rollout_the_env_fails_every_time_raises_so_verl_leaves_it_out(serve, tmp_path, tokenizer): healthy = serve({"a": PLAYED}) row = _row(write_dataset(healthy.url, tmp_path)) healthy.stop() broken = serve({"a": [{**PLAYED[0], "fail": True}, *PLAYED[1:]]}) - output = asyncio.run(_play(broken.url, row, _Model(tokenizer), tokenizer, attempts=2)) + with pytest.raises(verl_loop.RolloutFailed, match=r"task a of split all: the env failed it on 2 attempts, so it's " + r"left out of training: .*the env didn't come up"): + asyncio.run(_play(broken.url, row, _Model(tokenizer), tokenizer, attempts=2)) broken.stop() - assert (output.reward_score, output.extra_fields["agentenv"]["attempts"]) == (0.0, 2) - assert "the env didn't come up" in output.extra_fields["agentenv"]["failure"] - assert output.response_mask and set(output.response_mask) == {0} - def test_a_rollout_isnt_played_again_once_its_time_is_up(serve, tmp_path, tokenizer, monkeypatch): server = serve({"a": PLAYED}) @@ -124,11 +126,10 @@ async def hanging(self, context): return await deploy(self, context) monkeypatch.setattr(run_test._Env, "execute", hanging) - output = asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer, reset_timeout=0.5)) + with pytest.raises(verl_loop.RolloutFailed, match="the env failed it on 1 attempt, so"): + asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer, reset_timeout=0.5)) server.stop() - assert (output.reward_score, output.extra_fields["agentenv"]["attempts"]) == (0.0, 1) - def test_a_prompt_that_doesnt_fit_with_its_tools_is_refused_rather_than_cut(serve, tmp_path, tokenizer): server = serve({"a": PLAYED}) @@ -166,6 +167,52 @@ async def twice(): server.stop() +def test_a_server_that_asks_for_a_token_is_played_with_the_one_its_variable_holds(serve, tmp_path, tokenizer, + monkeypatch): + server = serve({"a": PLAYED}, token="s3cret") + monkeypatch.setenv("DOCK_TOKEN", "s3cret") + row = _row(write_dataset(server.url, tmp_path, token_var="DOCK_TOKEN")) + + played = asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer, token_var="DOCK_TOKEN")) + with pytest.raises(verl_loop.RolloutFailed, match="couldn't open a session on the server"): + asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer, attempts=1)) + server.stop() + + assert played.reward_score == 1.0 + + +def test_an_answer_verl_cuts_is_warned_of_once_a_tool(serve, tmp_path, tokenizer, monkeypatch, caplog): + server = serve({"a": PLAYED}) + row = _row(write_dataset(server.url, tmp_path)) + monkeypatch.setattr(verl_loop, "_cut", set()) + + with caplog.at_level(logging.WARNING, logger=verl_loop.__name__): + for _ in range(2): + asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer, max_tool_response_length=4)) + server.stop() + + [warning] = [record.getMessage() for record in caplog.records if "verl cuts" in record.getMessage()] + assert re.fullmatch(r"add_item answered \d+ characters, which verl cuts to actor_rollout_ref\.rollout\.multi_turn\." + r"max_tool_response_length=4, so the model doesn't see all of it: raise it to cover the env's " + r"answers", warning) + + +def test_the_writer_counts_each_prompt_with_its_tools_and_refuses_one_over_the_limit(serve, tmp_path, tokenizer): + server = serve({"a": PLAYED}) + + written = write_dataset(server.url, tmp_path / "data", tokenizer=TOKENIZER) + row = _row(written) + with pytest.raises(ValueError, match=r"1 rows' prompts are over max_prompt_length=16 tokens with their tools, the " + r"longest task a of split all at \d+, which verl would cut"): + write_dataset(server.url, tmp_path / "refused", tokenizer=TOKENIZER, max_prompt_length=16) + server.stop() + + rendered = tokenizer.apply_chat_template(row["prompt"], tools=json.loads(row["extra_info"]["tools"]), + add_generation_prompt=True, tokenize=True, return_dict=True) + assert written.longest_prompt == len(rendered["input_ids"]) > 16 + assert not (tmp_path / "refused").exists() + + def test_the_rows_config_names_this_agent_loop(): assert LOOP == f"{verl_loop.AgentEnvAgentLoop.__module__}.{verl_loop.AgentEnvAgentLoop.__qualname__}" @@ -188,16 +235,17 @@ async def _play(url, row, model, tokenizer, **options): return await _run(_loop(url, model, tokenizer, **options), row) -def _loop(url, model, tokenizer, *, parser="hermes", prompt_length=2048, reset_timeout=60.0, attempts=3): +def _loop(url, model, tokenizer, *, parser="hermes", prompt_length=2048, max_tool_response_length=2048, + reset_timeout=60.0, attempts=3, token_var=None): config = omegaconf.OmegaConf.create({"actor_rollout_ref": {"rollout": { "prompt_length": prompt_length, "response_length": 1024, "full_determinism": False, "multi_turn": {"format": parser, "max_user_turns": None, "max_assistant_turns": 4, "max_parallel_calls": 1, - "max_tool_response_length": 2048, "tool_response_truncate_side": "middle"}, + "max_tool_response_length": max_tool_response_length, "tool_response_truncate_side": "middle"}, }}}) return verl_loop.AgentEnvAgentLoop( trainer_config=agent_loop.DictConfigWrap(config), server_manager=model, tokenizer=tokenizer, processor=None, dataset_cls=rl_dataset.RLHFDataset, data_config=agent_loop.DictConfigWrap(DATA), hf_model_type="qwen3", - url=url, reset_timeout=reset_timeout, attempts=attempts) + url=url, reset_timeout=reset_timeout, attempts=attempts, token_var=token_var) async def _run(played, row): @@ -205,7 +253,7 @@ async def _run(played, row): tools_kwargs={}, index=0, agent_name="agentenv") -def _row(folder): - [row] = pq.read_table(folder / "all.parquet").to_pylist() +def _row(written): + [row] = pq.read_table(written.folder / "all.parquet").to_pylist() assert json.loads(row["extra_info"]["tools"])[0]["function"]["name"] == "add_item" return row diff --git a/uv.lock b/uv.lock index b3974b21..a60483ed 100644 --- a/uv.lock +++ b/uv.lock @@ -172,6 +172,7 @@ verl = [ { name = "openenv" }, { name = "pyarrow", version = "25.0.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13.14'" }, { name = "pyarrow", version = "26.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13.14'" }, + { name = "transformers" }, ] [package.metadata] @@ -235,6 +236,7 @@ requires-dist = [ { name = "tensorlake", marker = "extra == 'tensorlake'", specifier = "~=0.5.140" }, { name = "tomli-w", marker = "extra == 'dev'", specifier = "==1.2.0" }, { name = "transformers", marker = "extra == 'trl'", specifier = ">=5.2" }, + { name = "transformers", marker = "extra == 'verl'", specifier = ">=5" }, { name = "urllib3", marker = "extra == 'boat'", specifier = ">=2.1.0,<3" }, { name = "uvicorn", specifier = ">=0.47" }, { name = "vercel-sandbox", marker = "extra == 'vercel'", specifier = ">=0.7.0,<0.8" }, From 4ab386d2113470ed12fd31ae3c5573492d5f8797 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 11:05:12 -0700 Subject: [PATCH 11/13] fix(training): a verl rollout waits out an env server refusing connections while it restarts A refused connection counted as nothing listening, so a server restarted mid-run failed every rollout in flight on all its attempts within seconds. The verl loop's sessions now wait it out within reset_timeout, as they wait out a server at its limit; the writer read the server first, so a refusal then is a server coming back. A server that refuses the token still fails at once. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/_openenv.py | 25 +++++++++++++++++-------- src/agent_env/training/verl/loop.py | 5 ++++- tst/unit/training/sessions_test.py | 23 +++++++++++++---------- 3 files changed, 34 insertions(+), 19 deletions(-) diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index 07698698..2e226cac 100644 --- a/src/agent_env/training/_openenv.py +++ b/src/agent_env/training/_openenv.py @@ -58,14 +58,16 @@ class NamedToolset: @dataclass(frozen=True) class Shared: """What a session needs of its server: where it is, the tools the model is shown, which an episode's first turn - must offer no more than, the deadlines, how long to wait for room on the server, and the token it asks for. The key - makes it equal to its pickled copies only.""" + must offer no more than, the deadlines, how long to wait for room on the server, whether to wait as long for a + server that refuses connections, as one restarting does, and the token it asks for. The key makes it equal to its + pickled copies only.""" url: str offered: frozenset[str] tool_timeout: float reset_timeout: float patience: float = REOPEN_PATIENCE + waits_for_server: bool = False token: str | None = field(default=None, repr=False) key: str = field(default_factory=lambda: uuid.uuid4().hex) @@ -133,7 +135,8 @@ async def _start(self, split: str, index: int, expected: str | None) -> None: async def _begin(self, rollout: Rollout, split: str, index: int, expected: str | None) -> None: if self._client is None: self._client = await reopened(self._shared.url, self._shared.reset_timeout, dropped=self._opened, - patience=self._shared.patience, token=self._shared.token) + patience=self._shared.patience, waits_for_server=self._shared.waits_for_server, + token=self._shared.token) self._opened = True reset = await self._client.reset(split=split, index=index, timeout_s=self._shared.reset_timeout) if reset.done: @@ -269,11 +272,12 @@ async def opened(url: str, timeout: float, *, token: str | None = None) -> Gener return client -async def reopened(url: str, timeout: float, *, dropped: bool, patience: float, +async def reopened(url: str, timeout: float, *, dropped: bool, patience: float, waits_for_server: bool = False, token: str | None = None) -> GenericEnvClient: """A session, tried again every ``REOPEN_SECONDS`` for ``patience`` seconds while the server is at its limit or too - busy to answer, in place of a ``dropped`` one or for a rollout's own; raises Failed, which leaves the rollout - unscored, when it can't open one, at once unless the server was at its limit or too busy to answer.""" + busy to answer, or with ``waits_for_server`` refuses the connection, as one restarting does, in place of a + ``dropped`` one or for a rollout's own; raises Failed, which leaves the rollout unscored, when it can't open one, at + once unless it was waiting.""" lead = "its session was dropped, and " if dropped else "" deadline = time.monotonic() + patience while True: @@ -290,7 +294,8 @@ async def reopened(url: str, timeout: float, *, dropped: bool, patience: float, "more at once", url, e.limit) await asyncio.sleep(REOPEN_SECONDS) except ConnectionError as e: - if not _transient(e) or time.monotonic() + REOPEN_SECONDS > deadline: + waits = _transient(e) or (waits_for_server and _caused_by(e, ConnectionRefusedError)) + if not waits or time.monotonic() + REOPEN_SECONDS > deadline: raise Failed(f"{lead}{e}") from e await asyncio.sleep(REOPEN_SECONDS) @@ -298,8 +303,12 @@ async def reopened(url: str, timeout: float, *, dropped: bool, patience: float, def _transient(error: BaseException | None) -> bool: """Whether ``error`` came of a server too busy to answer in time or a connection it dropped, rather than of nothing listening at its address or an answer that isn't a session.""" + return _caused_by(error, _TRANSIENT) + + +def _caused_by(error: BaseException | None, causes: type[BaseException] | tuple[type[BaseException], ...]) -> bool: while error is not None: - if isinstance(error, _TRANSIENT): + if isinstance(error, causes): return True error = error.__cause__ or error.__context__ return False diff --git a/src/agent_env/training/verl/loop.py b/src/agent_env/training/verl/loop.py index ca92a1e5..4ee43ab7 100644 --- a/src/agent_env/training/verl/loop.py +++ b/src/agent_env/training/verl/loop.py @@ -41,6 +41,9 @@ class AgentEnvAgentLoop(ToolAgentLoop): past its deadline or a server that went away, is played again from the start on a new episode, a fresh sample of the same policy, up to ``attempts`` times, but not once ``reset_timeout`` seconds have passed since it began; one that fails every time raises RolloutFailed, so verl's V1 trainer leaves it out of training and of the eval reward. + Opening its session, a rollout waits out a server at its limit or refusing connections, as one restarting does, + within ``reset_timeout``, so a restart delays rollouts rather than failing them: the writer read the server first, + so a refusal then is a server coming back, not a wrong address. A server that refuses the token fails at once. A task that ends unscored without failing, and a prompt that doesn't fit ``rollout.prompt_length`` with its tools, are configuration errors no replay changes, and raise too; the V1 trainer logs each such session and leaves it out. An env answer longer than ``multi_turn.max_tool_response_length``, which verl cuts, is warned of once a tool. @@ -77,7 +80,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop row = self._row = kwargs["extra_info"] schemas = json.loads(row["tools"]) shared = Shared(self._url, frozenset(schema["function"]["name"] for schema in schemas), self._tool_timeout, - self._reset_timeout, patience=self._reset_timeout, token=self._token) + self._reset_timeout, patience=self._reset_timeout, waits_for_server=True, token=self._token) began = time.monotonic() for attempt in range(1, self._attempts + 1): session = _Session(shared) diff --git a/tst/unit/training/sessions_test.py b/tst/unit/training/sessions_test.py index 73658e97..002fd962 100644 --- a/tst/unit/training/sessions_test.py +++ b/tst/unit/training/sessions_test.py @@ -1,6 +1,6 @@ """``agent_env.training._openenv`` opening a rollout's session: tried again while the server is at its limit or too busy -to answer its handshake in time, until the patience runs out, and given up at once when nothing listens at its -address.""" +to answer its handshake in time, or when asked to wait for it refuses connections as one restarting does, until the +patience runs out, and given up at once when nothing listens at its address or it refuses the token.""" import asyncio @@ -10,12 +10,14 @@ from agent_env.training._openenv import Failed, reopened -@pytest.mark.parametrize(("cause", "opens", "attempts"), [ - (TimeoutError("timed out"), True, 3), - (ConnectionRefusedError(61, "Connection refused"), False, 1), -], ids=["busy", "refused"]) -def test_a_session_is_tried_again_while_its_server_is_busy_and_given_up_at_once_when_none_listens(monkeypatch, cause, - opens, attempts): +@pytest.mark.parametrize(("cause", "waits_for_server", "opens", "attempts"), [ + (TimeoutError("timed out"), False, True, 3), + (ConnectionRefusedError(61, "Connection refused"), False, False, 1), + (ConnectionRefusedError(61, "Connection refused"), True, True, 3), + (ConnectionError("server rejected WebSocket connection: HTTP 403"), True, False, 1), +], ids=["busy", "refused", "restarting", "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, opens, attempts): tried = [] async def flaky(url, timeout, *, token=None): @@ -30,8 +32,9 @@ async def flaky(url, timeout, *, token=None): monkeypatch.setattr(_openenv, "REOPEN_SECONDS", 0.01) if opens: - assert asyncio.run(reopened("http://busy", 5, dropped=False, patience=5)) == "session" + assert asyncio.run(reopened("http://busy", 5, dropped=False, patience=5, + waits_for_server=waits_for_server)) == "session" else: 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=5)) + asyncio.run(reopened("http://busy", 5, dropped=False, patience=5, waits_for_server=waits_for_server)) assert len(tried) == attempts From 210888f25e246f058a980767248953cc778258b9 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 11:17:12 -0700 Subject: [PATCH 12/13] fix(training): the verl writer counts prompts with verl's chat template settings --chat-template-kwargs, verl's data.apply_chat_template_kwargs, so a prompt near the limit is counted as verl renders it (Qwen3 with thinking off renders differently). Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 2 +- src/agent_env/training/verl/__main__.py | 10 ++++++++-- src/agent_env/training/verl/prepare.py | 19 +++++++++++-------- tst/unit/training/verl_test.py | 3 ++- tst/verl/loop_test.py | 15 +++++++++++---- 5 files changed, 33 insertions(+), 16 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f7618ff1..af70d4f5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -122,7 +122,7 @@ and deserialized through a registry. | `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `smol_vm` (the `smol` extra), `tensorlake` (the `tensorlake` extra), `tenki` (the `tenki` extra), `vercel` (the `vercel` extra), `boxd` (the `boxd` extra), `coreweave` (the `coreweave` extra), `daytona` (the `daytona` extra), `boat` (the `boat` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store), `EnvironmentGatewayProvider` (`gateway`, an env's default), which renders a docker-compose for the gateway and its MCP servers inside the sandbox, and `EnvironmentServerProvider` (`server`), one MCP server in its own container with no gateway in front and no state store; `providers/env_state/` holds env-state providers (`local_postgres` built in). | | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `openenv/` | `build_app(bundle)`: the [OpenEnv](https://github.com/openenv/OpenEnv) server for a bundle's tasks, each session playing one episode through the Episode API, with the envs' MCP tools, an added `end_turn` tool and the task's verifier score as the reward. It prepares bundles as `agent-env run` does (`prepare_bundle(played=True)`, without the agent the client plays) and keeps namespace routing on for the app's life. `openenv` comes with the `openenv` extra, a floor with no cap; this package and `training/` are the only ones that import it. Each task's description in the Task API names its `toolset` (the played agent's envs, role, and whether it gets `bash`). With `token`, every request and WebSocket but `/health` must carry `Authorization: Bearer `, for a server others reach. | -| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, `tasks_and_tools` (what the row writers read: the tasks, the split names checked, each toolset's tools) and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and raising `RolloutFailed` for one that fails every time, which verl's V1 trainer (the default) leaves out of the loss, the group's baseline and the eval reward; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing, and warns once a tool of an answer verl cuts to `max_tool_response_length`. The writer reads a server that asks for a token with the one `--token-var` names, which the loop reads in each worker, and with `--tokenizer` counts each prompt with its tools, refusing rows over `--max-prompt-length` before writing. | +| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, `tasks_and_tools` (what the row writers read: the tasks, the split names checked, each toolset's tools) and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and raising `RolloutFailed` for one that fails every time, which verl's V1 trainer (the default) leaves out of the loss, the group's baseline and the eval reward; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing, and warns once a tool of an answer verl cuts to `max_tool_response_length`. The writer reads a server that asks for a token with the one `--token-var` names, which the loop reads in each worker, and with `--tokenizer` counts each prompt with its tools as verl renders it (`--chat-template-kwargs`), refusing rows over `--max-prompt-length` before writing. | | `runner/` | The `[runner]` seam: `Runner.submit()` returns a `RunHandle` (`run_id`, `instance_id`) before the run starts; `LocalRunner` is built in. `runner/store.py` is the run index every runner writes and the explorer reads. | | `bundle/` | A bundle is a folder of entities (`envs/`, `agents/`, `artifacts/`, `skills/`, `tasks/.json`, `evals/.toml`) that runs without a registry: `agent-env run` plans it, writes it into the `@local` namespace (reusing each version whose inputs haven't changed) and runs its evals or tasks; `agent-env bundle init` writes a stored task or eval as one, with every env, agent and artifact it uses in its own folder (each image as its build context, pruned to what its Dockerfile reads and pinned to the Python and npm versions of the stored image, or, given `--image-ref ID=REF`, as that registry image pinned by digest with its context beside it, `image = { ref, dockerfile }`, which a run pulls and builds only when it can't; `task_images` lists the images an import writes, downloading nothing, so an export can push them first) so it runs with no store; `agent-env bundle register` writes a folder into the configured stores under registry ids (``), each task and eval naming the versions it reads and each image it builds pushed, through a ledger in the target store so registering again writes only what changed. | | `eval/` | `Eval` (`eval`), a versioned list of task references, and its store. | diff --git a/src/agent_env/training/verl/__main__.py b/src/agent_env/training/verl/__main__.py index f187204c..71a078f1 100644 --- a/src/agent_env/training/verl/__main__.py +++ b/src/agent_env/training/verl/__main__.py @@ -1,5 +1,7 @@ """``python -m agent_env.training.verl URL OUT``: writes a bundle's tasks as verl's RL dataset and agent loop config.""" +import json + import click from agent_env.training.verl.prepare import write_dataset @@ -22,12 +24,16 @@ "tools.") @click.option("--max-prompt-length", type=int, help="Refuse rows whose prompt with its tools is longer, in tokens: verl's data.max_prompt_length.") +@click.option("--chat-template-kwargs", type=json.loads, + help="The chat template's settings as JSON, verl's data.apply_chat_template_kwargs, to count prompts as " + 'verl renders them, e.g. \'{"enable_thinking": false}\'.') def main(url: str, out: str, tool_timeout: float, reset_timeout: float, attempts: int, token_var: str | None, - tokenizer: str | None, max_prompt_length: int | None) -> None: + tokenizer: str | None, max_prompt_length: int | None, chat_template_kwargs: dict | None) -> None: """Writes the tasks of the OpenEnv server at URL as OUT/.parquet, a row per task, and OUT/agent_loop.yaml, the agent loop that plays them on the server.""" written = write_dataset(url, out, tool_timeout=tool_timeout, reset_timeout=reset_timeout, attempts=attempts, - token_var=token_var, tokenizer=tokenizer, max_prompt_length=max_prompt_length) + token_var=token_var, tokenizer=tokenizer, max_prompt_length=max_prompt_length, + chat_template_kwargs=chat_template_kwargs) click.echo(f"wrote {', '.join(sorted(path.name for path in written.folder.iterdir()))} to {written.folder}") if written.longest_prompt is not None: click.echo(f"the longest prompt is {written.longest_prompt} tokens with its tools") diff --git a/src/agent_env/training/verl/prepare.py b/src/agent_env/training/verl/prepare.py index 5d1dc87f..17d1b3da 100644 --- a/src/agent_env/training/verl/prepare.py +++ b/src/agent_env/training/verl/prepare.py @@ -41,7 +41,7 @@ class Written: def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, reset_timeout: float = 600.0, attempts: int = 3, token_var: str | None = None, tokenizer: str | None = None, - max_prompt_length: int | None = None) -> Written: + max_prompt_length: int | None = None, chat_template_kwargs: dict[str, Any] | None = None) -> Written: """The tasks of the OpenEnv server at ``url`` as ``out/.parquet``, a row per task of each split, and ``out/agent_loop.yaml``, which points the rows' agent loop at ``url`` with the deadlines and the ``attempts`` a rollout the env fails is played. A row's prompt is its task's system prompt and first prompt, and its tools every @@ -51,12 +51,13 @@ def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, rese ``token_var`` names the environment variable that holds the token a server given one asks for (``build_app(token=...)``): the writer reads it now and the loop in each of verl's workers, so the token is in no file. With ``tokenizer``, the model's as transformers loads it, each row's prompt is counted with its tools as the - model's chat template renders them, and a row over ``max_prompt_length`` is refused before anything is written, - where verl would cut it from the left; verl's own length filter doesn't count a row's tools.""" + model's chat template renders them with ``chat_template_kwargs``, verl's ``data.apply_chat_template_kwargs``, and a + row over ``max_prompt_length`` is refused before anything is written, where verl would cut it from the left; verl's + own length filter doesn't count a row's tools.""" if attempts < 1: raise ValueError(f"attempts must be at least 1, not {attempts}") - if max_prompt_length is not None and tokenizer is None: - raise ValueError("max_prompt_length is in tokens, so it takes the model's tokenizer") + if tokenizer is None and (max_prompt_length is not None or chat_template_kwargs): + raise ValueError("max_prompt_length and chat_template_kwargs count tokens, so they take the model's tokenizer") url = url.rstrip("/") env_name, tasks, probed = tasks_and_tools(url, local="serve the bundle with agent_env.openenv.build_app first", timeout=reset_timeout, token=token_in(token_var)) @@ -68,7 +69,8 @@ def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, rese "extra_info": {"split": split, "index": task["index"], "task_id": task["id"], "tools": json.dumps(tools[key(task)])}, } for task in described] for split, described in tasks.items()} - longest = None if tokenizer is None else _longest(rows, AutoTokenizer.from_pretrained(tokenizer), max_prompt_length) + longest = None if tokenizer is None else _longest(rows, AutoTokenizer.from_pretrained(tokenizer), max_prompt_length, + chat_template_kwargs or {}) out = Path(out) out.mkdir(parents=True, exist_ok=True) for split, written in rows.items(): @@ -89,11 +91,12 @@ def token_in(var: str | None) -> str | None: return token -def _longest(rows: dict[str, list[dict[str, Any]]], tokenizer: Any, limit: int | None) -> int: +def _longest(rows: dict[str, list[dict[str, Any]]], tokenizer: Any, limit: int | None, + chat_template_kwargs: dict[str, Any]) -> int: """The longest row's prompt in tokens with its tools, refusing rows over ``limit``.""" lengths = {(split, row["extra_info"]["task_id"]): len(tokenizer.apply_chat_template( row["prompt"], tools=json.loads(row["extra_info"]["tools"]), add_generation_prompt=True, tokenize=True, - return_dict=True)["input_ids"]) for split, written in rows.items() for row in written} + return_dict=True, **chat_template_kwargs)["input_ids"]) for split, written in rows.items() for row in written} if limit is not None and (over := {at: length for at, length in lengths.items() if length > limit}): (split, task), length = max(over.items(), key=lambda item: item[1]) raise ValueError(f"{len(over)} rows' prompts are over max_prompt_length={limit} tokens with their tools, the " diff --git a/tst/unit/training/verl_test.py b/tst/unit/training/verl_test.py index a5132ea3..875ba993 100644 --- a/tst/unit/training/verl_test.py +++ b/tst/unit/training/verl_test.py @@ -109,7 +109,8 @@ def test_a_server_that_asks_for_a_token_is_read_with_the_one_its_variable_holds_ @pytest.mark.parametrize(("options", "message"), [ ({"attempts": 0}, "attempts must be at least 1, not 0"), - ({"max_prompt_length": 1024}, "max_prompt_length is in tokens, so it takes the model's tokenizer"), + ({"max_prompt_length": 1024}, "max_prompt_length and chat_template_kwargs count tokens, so they take the model's"), + ({"chat_template_kwargs": {"enable_thinking": False}}, "count tokens, so they take the model's tokenizer"), ({"token_var": "UNSET_DOCK_TOKEN"}, "UNSET_DOCK_TOKEN isn't set: it holds the token the OpenEnv server asks for"), ]) def test_what_cant_be_written_is_refused_before_the_server_is_read(tmp_path, monkeypatch, options, message): diff --git a/tst/verl/loop_test.py b/tst/verl/loop_test.py index fd0cf7b9..98ab6bb1 100644 --- a/tst/verl/loop_test.py +++ b/tst/verl/loop_test.py @@ -197,19 +197,26 @@ def test_an_answer_verl_cuts_is_warned_of_once_a_tool(serve, tmp_path, tokenizer r"answers", warning) -def test_the_writer_counts_each_prompt_with_its_tools_and_refuses_one_over_the_limit(serve, tmp_path, tokenizer): +def test_the_writer_counts_each_prompt_with_its_tools_as_verl_renders_it_and_refuses_one_over_the_limit( + serve, tmp_path, tokenizer): server = serve({"a": PLAYED}) written = write_dataset(server.url, tmp_path / "data", tokenizer=TOKENIZER) + thinking_off = write_dataset(server.url, tmp_path / "off", tokenizer=TOKENIZER, + chat_template_kwargs={"enable_thinking": False}) row = _row(written) with pytest.raises(ValueError, match=r"1 rows' prompts are over max_prompt_length=16 tokens with their tools, the " r"longest task a of split all at \d+, which verl would cut"): write_dataset(server.url, tmp_path / "refused", tokenizer=TOKENIZER, max_prompt_length=16) server.stop() - rendered = tokenizer.apply_chat_template(row["prompt"], tools=json.loads(row["extra_info"]["tools"]), - add_generation_prompt=True, tokenize=True, return_dict=True) - assert written.longest_prompt == len(rendered["input_ids"]) > 16 + def rendered(**kwargs): + return len(tokenizer.apply_chat_template(row["prompt"], tools=json.loads(row["extra_info"]["tools"]), + add_generation_prompt=True, tokenize=True, return_dict=True, + **kwargs)["input_ids"]) + + assert (written.longest_prompt, thinking_off.longest_prompt) == (rendered(), rendered(enable_thinking=False)) + assert thinking_off.longest_prompt != written.longest_prompt > 16 assert not (tmp_path / "refused").exists() From 39791fada12be766ef4582855d225cc2ab166939 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 15:37:06 -0700 Subject: [PATCH 13/13] feat(training): a verl rollout's scores, attempts and model and env seconds in reward_extra_info verl reports reward_extra_info per eval as val-aux metrics and in its dumps: each of the episode's scores (score/), the attempts it took, and its seconds in the model and in the env (deploying, calling tools and scoring). The end-of-episode scoring is timed as verl's compute_score. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/verl/loop.py | 16 +++++++++++++--- tst/verl/loop_test.py | 17 ++++++++++------- 2 files changed, 23 insertions(+), 10 deletions(-) diff --git a/src/agent_env/training/verl/loop.py b/src/agent_env/training/verl/loop.py index 4ee43ab7..f1371606 100644 --- a/src/agent_env/training/verl/loop.py +++ b/src/agent_env/training/verl/loop.py @@ -47,7 +47,10 @@ class AgentEnvAgentLoop(ToolAgentLoop): A task that ends unscored without failing, and a prompt that doesn't fit ``rollout.prompt_length`` with its tools, are configuration errors no replay changes, and raise too; the V1 trainer logs each such session and leaves it out. An env answer longer than ``multi_turn.max_tool_response_length``, which verl cuts, is warned of once a tool. - ``agent_loop.yaml`` sets the arguments; verl passes the rest. + A scored rollout's ``reward_extra_info``, which verl reports per eval as ``val-aux`` metrics and in its dumps, + holds each of the episode's scores (``score/``), its attempts, and its seconds in the model and in the env + (deploying, calling tools and scoring, which is also verl's ``compute_score`` time). ``agent_loop.yaml`` sets the + arguments; verl passes the rest. It leans on verl internals, which CI checks on verl 0.9.1 and main: verl builds an agent loop per trajectory, so an instance holds one rollout's state, and ``run`` refuses a second; ``ToolAgentLoop.run`` reads the tools from @@ -93,7 +96,11 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop raise ValueError(f"{self._task()} ended without a reward: it records no score its server's reward " "reads, so it can't be trained on") output.reward_score = float(rollout.reward) - output.extra_fields["agentenv"] = {"attempts": attempt, "scores": rollout.scores} + output.extra_fields["reward_extra_info"] = { + "attempts": attempt, "model_seconds": output.metrics.generate_sequences, + "env_seconds": output.metrics.tool_calls + output.metrics.compute_score, + **{f"score/{step}": score for step, score in (rollout.scores or {}).items()}, + } return output logger.warning("%s, attempt %d of %d: %s", self._task(), attempt, self._attempts, rollout.failure) if time.monotonic() - began >= self._reset_timeout: @@ -109,7 +116,10 @@ async def _play(self, session: _Session, schemas: list[dict[str, Any]], sampling self.tool_schemas = schemas self._reply = "" output = await super().run(sampling_params, **kwargs) - return output, await session.finish(self._reply) + began = time.monotonic() + rollout = await session.finish(self._reply) + output.metrics.compute_score = time.monotonic() - began + return output, rollout def _cap_text_prompt_length(self, prompt_ids: list[int]) -> list[int]: """Refuses the first prompt when it doesn't fit, where verl would cut it from the left, losing the system diff --git a/tst/verl/loop_test.py b/tst/verl/loop_test.py index 98ab6bb1..e4e490c5 100644 --- a/tst/verl/loop_test.py +++ b/tst/verl/loop_test.py @@ -1,10 +1,10 @@ """``AgentEnvAgentLoop`` on verl itself: verl's ``ToolAgentLoop`` plays a row ``write_dataset`` wrote, with Qwen3's tokenizer and chat template and a stand-in rollout server, on the unit tier's bundle served on loopback. A played row's -reward is its episode's, through verl's hermes and qwen3_coder parsers and from a server that asks for a token; a -rollout the env fails is played again on a new episode, but not once its time is up, and one it fails every time raises, -which verl's V1 trainer leaves out; a prompt that doesn't fit, an episode that ends unscored without failing and a loop -run twice raise; an answer verl cuts is warned of; and the writer counts prompts with their tools in the model's tokens. -CI runs it on verl 0.9.1 and main.""" +reward is its episode's, and its scores, attempts and seconds its reward_extra_info, through verl's hermes and +qwen3_coder parsers and from a server that asks for a token; a rollout the env fails is played again on a new episode, +but not once its time is up, and one it fails every time raises, which verl's V1 trainer leaves out; a prompt that +doesn't fit, an episode that ends unscored without failing and a loop run twice raise; an answer verl cuts is warned of; +and the writer counts prompts with their tools in the model's tokens. CI runs it on verl 0.9.1 and main.""" import asyncio import json @@ -64,7 +64,10 @@ async def finishing(self, reply): output = asyncio.run(_play(server.url, row, model, tokenizer, parser=parser)) server.stop() - assert (output.reward_score, output.extra_fields["agentenv"]) == (1.0, {"attempts": 1, "scores": {"replies": 1.0}}) + info = output.extra_fields["reward_extra_info"] + assert (output.reward_score, info["attempts"], info["score/replies"]) == (1.0, 1, 1.0) + assert set(info) == {"attempts", "model_seconds", "env_seconds", "score/replies"} + assert info["env_seconds"] >= output.metrics.compute_score > 0 and info["model_seconds"] >= 0 assert (model.calls, replies) == (2, ["Added a crate."]) assert 0 in output.response_mask and 1 in output.response_mask assert "Added a crate." in tokenizer.decode(output.response_ids) @@ -101,7 +104,7 @@ async def failing_once(self, context): output = asyncio.run(_play(server.url, row, _Model(tokenizer), tokenizer)) server.stop() - assert (output.reward_score, output.extra_fields["agentenv"]["attempts"], len(deploys)) == (1.0, 2, 2) + assert (output.reward_score, output.extra_fields["reward_extra_info"]["attempts"], len(deploys)) == (1.0, 2, 2) def test_a_rollout_the_env_fails_every_time_raises_so_verl_leaves_it_out(serve, tmp_path, tokenizer):