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/AGENTS.md b/AGENTS.md index 22cd0a57..af70d4f5 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 | @@ -120,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 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. | +| `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 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. | @@ -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) diff --git a/pyproject.toml b/pyproject.toml index 38cb8b14..08c569f8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,14 @@ 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; 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 = [ "boto3>=1.35.0", @@ -117,7 +125,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/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/__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.""" diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index c37dd6ef..2e226cac 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() @@ -57,7 +58,8 @@ 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 + 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 @@ -65,6 +67,8 @@ class Shared: 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) @@ -122,7 +126,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 @@ -130,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) + 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: @@ -250,10 +256,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: @@ -265,15 +272,17 @@ 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, 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: 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 " @@ -285,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) @@ -293,19 +303,24 @@ 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 -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) @@ -347,10 +362,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: @@ -365,6 +380,19 @@ 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, 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, 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, token=token))[1] + for each, toolset in toolsets(name, tasks).items()} + + 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.""" @@ -388,6 +416,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/nemo_gym/prepare.py b/src/agent_env/training/nemo_gym/prepare.py index c2a450d4..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, 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"}) @@ -45,7 +44,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, @@ -69,13 +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") - 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 '.'") - 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/__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..71a078f1 --- /dev/null +++ b/src/agent_env/training/verl/__main__.py @@ -0,0 +1,43 @@ +"""``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 + + +@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, 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 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.") +@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, 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, + 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") + + +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..f1371606 --- /dev/null +++ b/src/agent_env/training/verl/loop.py @@ -0,0 +1,194 @@ +"""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. CI checks it on verl 0.9.1 and +main.""" + +from __future__ import annotations + +import json +import logging +import time +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 +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 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. + 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 + ``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. + 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, 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, token_var: str | None = None, **kwargs: Any): + super().__init__(*args, **kwargs) + 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 = "" + + 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, + 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) + try: + output, rollout = await self._play(session, schemas, sampling_params, kwargs) + finally: + 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 " + "reads, so it can't be trained on") + output.reward_score = float(rollout.reward) + 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: + break + 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, self.max_tool_response_length) + for schema in schemas} + self.tool_schemas = schemas + self._reply = "" + output = await super().run(sampling_params, **kwargs) + 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 + 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.""" + 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) + + def _task(self) -> str: + return f"task {self._row['task_id']} of split {self._row['split']}" + + +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], 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: + 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) + + async def finish(self, reply: str) -> Rollout: + rollout = self._rollout + await self._finish(rollout, reply) + return rollout + + async def close(self) -> None: + await self._close() + + +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/src/agent_env/training/verl/prepare.py b/src/agent_env/training/verl/prepare.py new file mode 100644 index 00000000..17d1b3da --- /dev/null +++ b/src/agent_env/training/verl/prepare.py @@ -0,0 +1,113 @@ +"""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 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 + +#: 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())])), +]) + + +@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, token_var: str | None = None, tokenizer: str | None = None, + 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 + 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 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 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)) + 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, + chat_template_kwargs or {}) + out = Path(out) + out.mkdir(parents=True, exist_ok=True) + 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, **({"token_var": token_var} if token_var else {}), + }], sort_keys=False)) + 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, + 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, **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 " + 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]: + """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/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/test_optional_extras_stay_optional.py b/tst/unit/test_optional_extras_stay_optional.py index b981a493..00cfc15a 100644 --- a/tst/unit/test_optional_extras_stay_optional.py +++ b/tst/unit/test_optional_extras_stay_optional.py @@ -26,11 +26,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", ) @@ -169,13 +170,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))) == [] diff --git a/tst/unit/training/_support.py b/tst/unit/training/_support.py index b9ef37a4..157db065 100644 --- a/tst/unit/training/_support.py +++ b/tst/unit/training/_support.py @@ -12,18 +12,19 @@ @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, 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"): + 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), _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))) + 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..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,15 +10,17 @@ 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): + 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}") @@ -30,8 +32,9 @@ async def flaky(url, timeout): 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 diff --git a/tst/unit/training/verl_test.py b/tst/unit/training/verl_test.py new file mode 100644 index 00000000..875ba993 --- /dev/null +++ b/tst/unit/training/verl_test.py @@ -0,0 +1,140 @@ +"""``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, 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 +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, ROW, 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") + + written = write_dataset(f"{server.url}/", tmp_path / "data", tool_timeout=30.0, reset_timeout=90.0, attempts=2) + server.stop() + + 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, + }] + [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).folder + 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_the_server_keeps_empty_has_a_rows_columns(serve, tmp_path): + server = serve({"a": PLAYED}, {"empty": []}) + + out = write_dataset(server.url, tmp_path).folder + 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"]}) + + 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_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 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): + 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()) == [] + + +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..e4e490c5 --- /dev/null +++ b/tst/verl/loop_test.py @@ -0,0 +1,269 @@ +"""``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, 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 +import logging +import re + +import pyarrow.parquet as pq +import pytest +from transformers import AutoTokenizer + +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 + +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)) +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" +DATA = omegaconf.OmegaConf.create({"apply_chat_template_kwargs": {"enable_thinking": False}}) +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") +def tokenizer(): + return AutoTokenizer.from_pretrained(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, parser) + finish, replies = verl_loop._Session.finish, [] + + async def finishing(self, reply): + replies.append(reply) + return await finish(self, reply) + + monkeypatch.setattr(verl_loop._Session, "finish", finishing) + output = asyncio.run(_play(server.url, row, model, tokenizer, parser=parser)) + server.stop() + + 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) + 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): + server = serve({"a": PLAYED}) + row = _row(write_dataset(server.url, tmp_path)) + deploy, deploys = run_test._Env.execute, [] + + 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(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["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): + healthy = serve({"a": PLAYED}) + row = _row(write_dataset(healthy.url, tmp_path)) + healthy.stop() + broken = serve({"a": [{**PLAYED[0], "fail": True}, *PLAYED[1:]]}) + + 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() + + +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 hanging(self, context): + await asyncio.sleep(2) + return await deploy(self, context) + + monkeypatch.setattr(run_test._Env, "execute", hanging) + 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() + + +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): + server = serve({"a": PLAYED}, reward=lambda scores: None) + row = _row(write_dataset(server.url, tmp_path)) + model = _Model(tokenizer) + + 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 + + +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_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_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() + + 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() + + +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.""" + + 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[self.parser] + return replica.TokenOutput(token_ids=self.tokenizer.encode(reply, add_special_tokens=False)) + + +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, 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": 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, token_var=token_var) + + +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") + + +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 61337f0f..a60483ed 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,20 @@ 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'" }, + { name = "transformers" }, +] [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 +216,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" }, @@ -226,11 +236,12 @@ 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" }, ] -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"