diff --git a/.github/workflows/areal.yml b/.github/workflows/areal.yml new file mode 100644 index 00000000..344091ba --- /dev/null +++ b/.github/workflows/areal.yml @@ -0,0 +1,60 @@ +name: AReaL contract + +# AgentEnvWorkflow in AReaL's agent workflow (tst/areal): AReaL's rollout engine resolves it from its import path into +# its proxy workflow in inline mode, grouped, and plays the unit tier's bundle through a stand-in model, on AReaL 2.1.0 +# and on its main branch. AReaL and torch are installed over the dev environment; torch is the CPU build, as nothing +# generates. It runs on pull requests that touch the helper or the tier, weekly, and on demand; it isn't a required +# check. + +on: + pull_request: + paths: + - "src/agent_env/training/**" + - "tst/areal/**" + - "tst/unit/training/**" + - "tst/util/fake_vllm.py" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/areal.yml" + schedule: + - cron: "53 7 * * 1" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: areal-${{ 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: + workflow: + name: agent workflow (AReaL ${{ matrix.areal }}) + runs-on: ubuntu-latest + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + areal: ["v2.1.0", "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 + run: | + python -m pip install --quiet -r .github/ci-requirements.txt + uv sync --locked --extra dev --python "$(command -v python)" + - name: Install AReaL and CPU torch + run: | + uv pip install --python .venv/bin/python "areal @ git+https://github.com/areal-project/AReaL@${{ matrix.areal }}" 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('areal', version('areal'), 'torch', version('torch'))" + - name: AReaL tier + run: .venv/bin/python -m pytest tst/areal -rA --timeout=900 --junitxml=junit-areal.xml + - name: No skips + if: always() + run: .venv/bin/python .github/scripts/check_skip_policy.py junit-areal.xml diff --git a/AGENTS.md b/AGENTS.md index 6c7dffc9..46c4dc1a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -48,6 +48,7 @@ or one of them only where it is used. | `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 skyrl-test` | `tst/skyrl`: rows played through `agent_env.training.skyrl`'s generator and env on SkyRL itself, with Qwen3's and Qwen3.5's tokenizers from the Hub and a stand-in inference engine; needs SkyRL, skyrl-gym and CPU torch on Python 3.12+, which the `dev` extra leaves out | +| `make areal-test` | `tst/areal`: rows played through `agent_env.training.areal`'s workflow in AReaL's agent workflow, resolved by AReaL's rollout engine, with a stand-in model; needs AReaL and CPU torch on Python 3.11 or 3.12, 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 | @@ -123,7 +124,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 as verl renders it (`--chat-template-kwargs`), refusing rows over `--max-prompt-length` before writing. `skyrl/` (`skyrl` extra, installed into SkyRL's own environment): `python -m agent_env.training.skyrl URL OUT --tokenizer MODEL` writes a server's tasks as SkyRL's `.parquet` rows (`env_class: agentenv`, the episode's split, index, task id and tools, the server's settings), each with its tools rendered into its system prompt by the model's chat template, checked to render as `tools=` would, since SkyRL's generator passes the template no tools, and refuses rows over `--max-prompt-length` counted as SkyRL's filter and generator count them; `calls.parse` reads a reply's tool calls in the `hermes` or `qwen3_coder` format, typed by the row's schema; `env.AgentEnvTextEnv` (registered as `agentenv`) is skyrl-gym's `BaseTextEnv` on the row's episode, run on the event loop the sessions share (`_openenv.on_session_loop`), its steps running the reply's calls and a reply without one ending the open turn, raising `EpisodeFailed` when the env fails it, with scores, tool calls and unparsed calls as its metrics; `generator.AgentEnvGenerator` is SkyRL's `SkyRLGymGenerator` playing a failed rollout again on a new episode up to `attempts` times within `reset_timeout`, and masking the group of one that fails every attempt in training (an eval scores it 0); `python -m agent_env.training.skyrl.train` (or `trainer.override_entrypoint=agent_env.training.skyrl.train`) runs SkyRL's trainer with both, its fully async one with `trainer.fully_async.enabled`, passing the rows' token variables to the trainer's Ray task. | +| `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. `skyrl/` (`skyrl` extra, installed into SkyRL's own environment): `python -m agent_env.training.skyrl URL OUT --tokenizer MODEL` writes a server's tasks as SkyRL's `.parquet` rows (`env_class: agentenv`, the episode's split, index, task id and tools, the server's settings), each with its tools rendered into its system prompt by the model's chat template, checked to render as `tools=` would, since SkyRL's generator passes the template no tools, and refuses rows over `--max-prompt-length` counted as SkyRL's filter and generator count them; `calls.parse` reads a reply's tool calls in the `hermes` or `qwen3_coder` format, typed by the row's schema; `env.AgentEnvTextEnv` (registered as `agentenv`) is skyrl-gym's `BaseTextEnv` on the row's episode, run on the event loop the sessions share (`_loop.on_session_loop`), its steps running the reply's calls and a reply without one ending the open turn, raising `EpisodeFailed` when the env fails it, with scores, tool calls and unparsed calls as its metrics; `generator.AgentEnvGenerator` is SkyRL's `SkyRLGymGenerator` playing a failed rollout again on a new episode up to `attempts` times within `reset_timeout`, and masking the group of one that fails every attempt in training (an eval scores it 0); `python -m agent_env.training.skyrl.train` (or `trainer.override_entrypoint=agent_env.training.skyrl.train`) runs SkyRL's trainer with both, its fully async one with `trainer.fully_async.enabled`, passing the rows' token variables to the trainer's Ray task. `harness.py` (`harness` extra) and `areal.py` (no extra) train a bundle's own agents instead, its tasks run as `agent-env run` runs them with the solver's model calls sent to the trainer: `harness.AgentEnvSessionFactory` is OpenEnv's session factory for TRL's harness rollouts, through OpenEnv's capture proxy; `areal.AgentEnvWorkflow` is the inline agent AReaL's proxy workflow runs, whose `run` returns the reward, raising for a run that fails or records no score, and whose `should_reject_failed_sample` leaves such a sample out of its group; `areal.task_rows` is its dataset. Both share `_solver.py` (the tasks whose solver can train, prepared for local Docker, and the overrides pointing the solver alone at the trainer's endpoint), `reward.py` (how a task's scores make its reward, which `build_app` applies too) and `_loop.py` (the event loop the sessions share, on a thread of its own). | | `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. | @@ -347,6 +348,9 @@ the same triggers; it allows no skips and isn't a required check. environment, on the same triggers; it allows no skips and isn't a required check. `.github/workflows/skyrl.yml` runs `tst/skyrl` on SkyRL 0.4.0 and its main branch, with skyrl-gym and CPU torch installed over the dev environment on Python 3.12, on the same triggers; it allows no skips and isn't a required check. +`.github/workflows/areal.yml` runs `tst/areal` on AReaL 2.1.0 and its main branch, with CPU torch installed over the dev +environment on Python 3.12, on pull requests that touch `agent_env/training` or the tier, weekly, and on demand; 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 1b557f90..be5cbe8c 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 verl-test skyrl-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 skyrl-test areal-test clean-install-test test clean VENV := .venv VENV_BIN := $(abspath $(VENV))/bin @@ -85,6 +85,12 @@ verl-test: ## Play rows through the verl agent loop on verl itself (needs verl a skyrl-test: ## Play rows through the SkyRL env and generator on SkyRL itself (needs SkyRL and CPU torch, and the Hub) $(PYTHON) -m pytest tst/skyrl --timeout=900 $(PYTEST_ARGS) +# The AReaL tier plays rows through AgentEnvWorkflow in AReaL's agent workflow, resolved by AReaL's rollout engine. +# AReaL and torch aren't in the dev extra, and AReaL needs Python 3.11 or 3.12: install them first, as +# .github/workflows/areal.yml does. +areal-test: ## Play rows through AgentEnvWorkflow in AReaL's agent workflow (needs AReaL and CPU torch) + $(PYTHON) -m pytest tst/areal --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/src/agent_env/openenv/app.py b/src/agent_env/openenv/app.py index 11acaa64..16fcce93 100644 --- a/src/agent_env/openenv/app.py +++ b/src/agent_env/openenv/app.py @@ -21,9 +21,9 @@ from agent_env.bundle.installed import locate_bundle from agent_env.store.routing import namespace_routing from agent_env.task import EpisodeError, EpisodeSpec, Task +from agent_env.training.reward import reward_of from .environment import AgentEnvEnvironment, Served, ServedTask -from .reward import reward_of logger = logging.getLogger(__name__) diff --git a/src/agent_env/training/__init__.py b/src/agent_env/training/__init__.py index ed79a783..53758839 100644 --- a/src/agent_env/training/__init__.py +++ b/src/agent_env/training/__init__.py @@ -1,4 +1,5 @@ """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, ``agent_env.training.nemo_gym`` for NeMo Gym and the trainers that take its environments, NeMo RL among them, ``agent_env.training.verl`` for verl, and ``agent_env.training.skyrl`` for -SkyRL.""" +SkyRL. ``agent_env.training.harness`` and ``agent_env.training.areal`` train a bundle's own agents instead, with TRL's +harness rollouts and AReaL's agent workflows; the second needs no extra.""" diff --git a/src/agent_env/training/_loop.py b/src/agent_env/training/_loop.py new file mode 100644 index 00000000..be0e2366 --- /dev/null +++ b/src/agent_env/training/_loop.py @@ -0,0 +1,31 @@ +"""The event loop this process's training sessions share, on a daemon thread of its own: a trainer calls its +environments from threads, or awaits them on an event loop of its own that a run's blocking calls mustn't stall.""" + +from __future__ import annotations + +import asyncio +import threading +from collections.abc import Coroutine +from typing import Any + +_loop: asyncio.AbstractEventLoop | None = None +_loop_lock = threading.Lock() + + +def on_session_loop(coroutine: Coroutine[Any, Any, Any], timeout: float | None) -> Any: + """``coroutine``'s result, run on the session loop; waits at most ``timeout`` seconds.""" + future = asyncio.run_coroutine_threadsafe(coroutine, session_loop()) + try: + return future.result(timeout) + except TimeoutError: + future.cancel() + raise + + +def session_loop() -> asyncio.AbstractEventLoop: + global _loop + with _loop_lock: + if _loop is None: + _loop = asyncio.new_event_loop() + threading.Thread(target=_loop.run_forever, name="agentenv-sessions", daemon=True).start() + return _loop diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index 8b96ae40..68436268 100644 --- a/src/agent_env/training/_openenv.py +++ b/src/agent_env/training/_openenv.py @@ -13,7 +13,7 @@ import threading import time import uuid -from collections.abc import Coroutine, Iterator, Mapping +from collections.abc import Iterator, Mapping from dataclasses import dataclass, field from typing import Any, ClassVar @@ -50,8 +50,6 @@ _waited: set[str] = set() #: The servers whose refused connections rollouts have been warned they wait out, once a process. _restarting: set[str] = set() -_loop: asyncio.AbstractEventLoop | None = None -_loop_lock = threading.Lock() @dataclass(frozen=True) @@ -426,26 +424,6 @@ def toolsets(name: str, tasks: dict[str, list[dict[str, Any]]]) -> dict[str, Nam return named -def on_session_loop(coroutine: Coroutine[Any, Any, Any], timeout: float | None) -> Any: - """``coroutine``'s result, run on the event loop this process's sessions share, on a daemon thread of its own, for - a trainer that calls its environments from threads; waits at most ``timeout`` seconds.""" - future = asyncio.run_coroutine_threadsafe(coroutine, _session_loop()) - try: - return future.result(timeout) - except TimeoutError: - future.cancel() - raise - - -def _session_loop() -> asyncio.AbstractEventLoop: - global _loop - with _loop_lock: - if _loop is None: - _loop = asyncio.new_event_loop() - threading.Thread(target=_loop.run_forever, name="agentenv-sessions", daemon=True).start() - return _loop - - 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: diff --git a/src/agent_env/training/_solver.py b/src/agent_env/training/_solver.py new file mode 100644 index 00000000..0aff3981 --- /dev/null +++ b/src/agent_env/training/_solver.py @@ -0,0 +1,89 @@ +"""What training a bundle's own agents shares, the harness and AReaL's helpers: the tasks whose solver, the agent a step +prompts, can train, prepared as ``agent-env run`` prepares them to run on local Docker, which alone reaches the +trainer's endpoint on this machine, a remote sandbox needing a tunnel; and the overrides that point the solver alone at +that endpoint, its other agents and judges keeping the configured one; and a finished run's scores, which a run +``agent-env run`` counts as failed has none of.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Any + +from agent_env.bundle.plan import solver_steps +from agent_env.bundle.resolve import build_step +from agent_env.bundle.run import planned_tasks, prepare_bundle, solver_params +from agent_env.store.routing import namespace_routing +from agent_env.task.task import Task, recorded_scores +from agent_env.task_step.context import TaskStepContext, sandbox_overrides +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.verifiers.rubrics_verifier import RubricsVerifierTaskStep +from agent_env.training.reward import reward_of + +logger = logging.getLogger(__name__) +SANDBOX = "local" + + +def trained(bundle: str | Path, tasks: Sequence[str], evals: Sequence[str], + agent: str | None) -> list[tuple[Task, dict[str, dict[str, Any]]]]: + """The selected tasks that can train their solver, prepared, each with its solver's deploys and their step params, + ``agent``'s if it swaps the solver. The others are left out before they're checked, so what they alone need can't + stop the rest.""" + with namespace_routing(): + kept = [] + for entry in planned_tasks(bundle, tasks=tasks, evals=evals, agent=agent): + if why := _untrainable(entry.config): + logger.warning("%s is left out: %s", entry.entry.name, why) + else: + kept.append(entry.entry.name) + if not kept: + raise ValueError("no task of the bundle has a solver to train: an agent a step prompts, deployed by a " + "deploy_agent step, that isn't also its judge") + prepared = prepare_bundle(bundle, tasks=kept, sandbox=SANDBOX, agent=agent, solver_endpoint=True) + swapped = solver_params(prepared.plan, prepared.materialization) + return [(task, {step: swapped.get(entry.id, {}).get(step, {}) + for step in solver_steps([step.to_dict() for step in task.steps])}) + for entry, task in prepared.tasks] + + +def solver_overrides(solver: dict[str, dict[str, Any]], endpoint: dict[str, str]) -> dict[str, Any]: + """A run's ``user_overrides``: every deploy on local Docker, and ``endpoint``, a ``litellm_base_url`` and + ``litellm_api_key``, on the solver's deploys alone.""" + return {**sandbox_overrides(SANDBOX), + "step_params": {step: {**params, **endpoint} for step, params in solver.items()}} + + +def rewarded(task: Task, reward: str | Callable[[dict[str, float]], float] | None + ) -> Callable[[dict[str, float]], float | None]: + of, why = reward_of(task, reward) + if of is None: + raise ValueError(why) + return of + + +def run_scores(task: Task, context: TaskStepContext) -> dict[str, float]: + """The scores a finished run of ``task`` recorded. Raises when a step failed and no retry recovered it, a run + ``agent-env run`` counts as failed, though a step that doesn't fail the task let it finish.""" + if failed := [failure for failure in context.metadata.get("failed_steps") or [] if not failure.get("retried")]: + raise RuntimeError("the run failed at " + "; ".join(f"{failure['step_id']}: {failure['error']}" + for failure in failed)) + return recorded_scores(task.steps, context) + + +def first_prompt(task: Task) -> str: + return next((step.prompt for step in task.steps if step.type == "prompt_agent"), "") + + +def _untrainable(config: list[dict[str, Any]]) -> str | None: + """Why a task's rollouts can't train its solver alone, or None.""" + solver = set(solver_steps(config)) + steps = [build_step(step) for step in config] + deploys = [step for step in steps if step.id in solver] + if not deploys or not all(isinstance(step, DeployAgentTaskStep) for step in deploys): + return "training needs the agent a step prompts deployed by a deploy_agent step" + agents = {step.agent_name for step in deploys} + if any(isinstance(step, RubricsVerifierTaskStep) and step.use_agent_judge and step.agent_name in agents + for step in steps): + return "its rubric judge is its solver, whose grading would be trained on" + return None diff --git a/src/agent_env/training/areal.py b/src/agent_env/training/areal.py new file mode 100644 index 00000000..8f0c4902 --- /dev/null +++ b/src/agent_env/training/areal.py @@ -0,0 +1,102 @@ +"""Train a bundle's own agents with AReaL's agent workflows (``rollout.agent.mode: inline``): ``AgentEnvWorkflow`` is +the agent AReaL's ``OpenAIProxyWorkflow`` runs for each sample, which plays the row's task as ``agent-env run`` does, +its solver's model calls going through AReaL's proxy, which records the tokens AReaL trains on, and returns the reward +the task's verifiers score. Needs no extra: AReaL calls it by its methods alone.""" + +from __future__ import annotations + +import asyncio +import functools +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Any + +from agent_env.store.routing import namespace_routing +from agent_env.task.task import Task +from agent_env.task.teardown import teardown_run +from agent_env.task_step.context import TaskStepContext +from agent_env.training._loop import session_loop +from agent_env.training._solver import first_prompt, rewarded, run_scores, solver_overrides, trained + +#: The model the solver names: AReaL's proxy serves its policy whatever a request names, and its own agents name this. +_MODEL = "default" + + +def task_rows(bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), + agent: str | None = None) -> list[dict[str, Any]]: + """The trainer's dataset, a row per task ``AgentEnvWorkflow`` plays with the same arguments: the task's id, which + is how the workflow finds it, and its first prompt as ``messages``, for AReaL's logs. The model never sees this + prompt: the task's own agent builds its own.""" + return [{"task_id": task.id, "messages": [{"role": "user", "content": first_prompt(task)}]} + for task, _ in trained(bundle, tasks, evals, agent)] + + +class AgentEnvWorkflow: + """The agent AReaL runs in inline mode, named as ``workflow: agent_env.training.areal.AgentEnvWorkflow`` with the + arguments as ``workflow_kwargs``. A process prepares the bundle to run on local Docker once, however often AReaL + builds the workflow; ``reward`` is as ``build_app``'s, and ``agent`` the agent the solver deploys, as + ``agent-env run --agent``'s.""" + + def __init__(self, bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), + reward: str | Callable[[dict[str, float]], float] | None = None, agent: str | None = None): + self._played = {task.id: (task, solver, rewarded(task, reward)) + for task, solver in _prepared(str(bundle), tuple(tasks), tuple(evals), agent)} + + async def run(self, data: dict[str, Any], *, base_url: str, api_key: str, **_: Any) -> float: + """Plays ``data``'s task, its solver calling the model at ``base_url`` with ``api_key``, the key of the session + AReaL records the calls in, and returns its reward. Raises, which has AReaL leave the sample out, when the run + fails or records no score its reward reads. Cancelled, it stops the run and passes the cancel on once the run + is torn down.""" + if (played := self._played.get(data.get("task_id"))) is None: + raise KeyError(f"the row names no task this workflow plays: {data.get('task_id')!r}; build the dataset " + "with task_rows, from the same bundle, tasks and evals") + task, solver, reward = played + run = _Run(task, solver_overrides(solver, {"litellm_base_url": base_url, "litellm_api_key": api_key})) + done = asyncio.wrap_future(asyncio.run_coroutine_threadsafe(run.scores(), session_loop())) + try: + scores = await asyncio.shield(done) + except asyncio.CancelledError: + session_loop().call_soon_threadsafe(run.stop) + await asyncio.wait([done]) + raise + if (score := reward(scores)) is None: + raise ValueError(f"task {task.id!r} records no score to train on") + return float(score) + + def should_reject_failed_sample(self, error: Exception, disposition: str | None) -> bool: + """AReaL's question, after 2.1.0, for a sample whose ``run`` raised: it is left out of its group, the group's + other samples still training, unless AReaL found the model at fault, which it trains on as a zero reward.""" + return disposition != "model_failure_zero" + + +class _Run: + """One run of a task on the session loop, which ``stop``, called there, cuts short; its teardown always runs to the + end.""" + + def __init__(self, task: Task, overrides: dict[str, Any]): + self._task, self._overrides = task, overrides + self._running: asyncio.Task | None = None + self._stopped = False + + async def scores(self) -> dict[str, float]: + context = TaskStepContext(metadata={"task_id": self._task.id, "user_overrides": self._overrides}) + with namespace_routing(): + try: + if self._stopped: + raise asyncio.CancelledError + self._running = asyncio.ensure_future(self._task.run(agent_model=_MODEL, context=context)) + await self._running + return run_scores(self._task, context) + finally: + await teardown_run(context) + + def stop(self) -> None: + self._stopped = True + if self._running is not None: + self._running.cancel() + + +@functools.cache +def _prepared(bundle: str, tasks: tuple[str, ...], evals: tuple[str, ...], + agent: str | None) -> list[tuple[Task, dict[str, dict[str, Any]]]]: + return trained(bundle, tasks, evals, agent) diff --git a/src/agent_env/training/harness.py b/src/agent_env/training/harness.py index 807eddde..d1271445 100644 --- a/src/agent_env/training/harness.py +++ b/src/agent_env/training/harness.py @@ -27,24 +27,16 @@ raise raise ImportError("agent_env.training.harness requires pip install 'agentenv-framework[harness]'") from exc -from agent_env.bundle.plan import solver_steps -from agent_env.bundle.resolve import build_step -from agent_env.bundle.run import planned_tasks, prepare_bundle, solver_params from agent_env.store.routing import namespace_routing -from agent_env.openenv.reward import reward_of -from agent_env.task.task import Task, recorded_scores +from agent_env.task.task import Task from agent_env.task.teardown import teardown_run -from agent_env.task_step.context import TaskStepContext, sandbox_overrides -from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep -from agent_env.task_step.task_steps.verifiers.rubrics_verifier import RubricsVerifierTaskStep -from agent_env.training._openenv import on_session_loop +from agent_env.task_step.context import TaskStepContext +from agent_env.training._loop import on_session_loop +from agent_env.training._solver import first_prompt, rewarded, run_scores, solver_overrides, trained logger = logging.getLogger(__name__) #: How a row's prompt names its task, since TRL hands the factory only the prompt. _NAMED = re.compile(r"^AgentEnv task (\d+): (.+)$") -#: Where every deploy runs: local Docker alone reaches the capture proxy on this machine, a remote sandbox needing a -#: tunnel. -_SANDBOX = "local" def prompt_rows(bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), @@ -53,8 +45,8 @@ def prompt_rows(bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequenc naming the task, which is how the factory finds it, then the task's first prompt, for TRL's logs. The model never sees this prompt: the task's own agent builds its own.""" return [{"prompt": [{"role": "system", "content": f"AgentEnv task {index}: {task.id}"}, - {"role": "user", "content": _first_prompt(task)}], "task_id": task.id} - for index, (task, _) in enumerate(_trained(bundle, tasks, evals, agent))] + {"role": "user", "content": first_prompt(task)}], "task_id": task.id} + for index, (task, _) in enumerate(trained(bundle, tasks, evals, agent))] class AgentEnvSessionFactory(ResourceSessionFactory): @@ -69,8 +61,8 @@ def __init__(self, bundle: str | Path, *, llm_url: str, model: str, sampling: di reward: str | Callable[[dict[str, float]], float] | None = None, agent: str | None = None): self.model = model self._llm_url, self._sampling = llm_url, sampling - self._trained = _trained(bundle, tasks, evals, agent) - self._rewards = [_reward(task, reward) for task, _ in self._trained] + self._trained = trained(bundle, tasks, evals, agent) + self._rewards = [rewarded(task, reward) for task, _ in self._trained] require_llm(llm_url, model, require_tokens=True) self._proxy = CaptureServer(llm_url=llm_url, model=model, port=_free_port(), admin_key=secrets.token_urlsafe(32)) @@ -82,9 +74,8 @@ def create(self, task: Any, seed: int | None = None, episode_id: str | None = No capture = self._proxy.registry.create( upstream=Upstream(llm_url=self._llm_url, model=self.model), capture_level="tokens", purpose="train", sampling=self._sampling, task=played.id, episode=episode_id or uuid.uuid4().hex) - endpoint = {"litellm_base_url": f"http://127.0.0.1:{self._proxy.port}", "litellm_api_key": capture.session_id} - overrides = {**sandbox_overrides(_SANDBOX), - "step_params": {step: {**params, **endpoint} for step, params in solver.items()}} + overrides = solver_overrides(solver, {"litellm_base_url": f"http://127.0.0.1:{self._proxy.port}", + "litellm_api_key": capture.session_id}) return AgentEnvSession(task=played, overrides=overrides, reward=self._rewards[index], registry=self._proxy.registry, capture=capture, model=self.model) @@ -103,7 +94,7 @@ def __init__(self, *, task: Task, overrides: dict[str, Any], reward: Callable[[d self._done: threading.Event | None = None # set once the run is torn down def initial_messages(self) -> list[dict[str, Any]]: - return [{"role": "user", "content": _first_prompt(self._task)}] + return [{"role": "user", "content": first_prompt(self._task)}] def list_tools(self) -> list[Any]: return [] @@ -159,7 +150,7 @@ async def _run(self, timeout: float | None) -> int: self._task.run(agent_model=self._model, context=context)) try: await asyncio.wait_for(task_run, timeout) - self._scores = recorded_scores(self._task.steps, context) + self._scores = run_scores(self._task, context) return 0 except TimeoutError: raise @@ -176,50 +167,6 @@ async def _run(self, timeout: float | None) -> int: self._done.set() -def _trained(bundle: str | Path, tasks: Sequence[str], evals: Sequence[str], - agent: str | None) -> list[tuple[Task, dict[str, dict[str, Any]]]]: - """The selected tasks that can train their solver, prepared as ``agent-env run`` prepares them, each with its - solver's deploys and their step params, ``agent``'s if it swaps the solver. The others are left out before they're - checked, so what they alone need can't stop the rest.""" - with namespace_routing(): - kept = [] - for entry in planned_tasks(bundle, tasks=tasks, evals=evals, agent=agent): - if why := _untrainable(entry.config): - logger.warning("%s is left out: %s", entry.entry.name, why) - else: - kept.append(entry.entry.name) - if not kept: - raise ValueError("no task of the bundle has a solver to train: an agent a step prompts, deployed by a " - "deploy_agent step, that isn't also its judge") - prepared = prepare_bundle(bundle, tasks=kept, sandbox=_SANDBOX, agent=agent, solver_endpoint=True) - swapped = solver_params(prepared.plan, prepared.materialization) - return [(task, {step: swapped.get(entry.id, {}).get(step, {}) - for step in solver_steps([step.to_dict() for step in task.steps])}) - for entry, task in prepared.tasks] - - -def _untrainable(config: list[dict[str, Any]]) -> str | None: - """Why a task's rollouts can't train its solver alone, or None.""" - solver = set(solver_steps(config)) - steps = [build_step(step) for step in config] - deploys = [step for step in steps if step.id in solver] - if not deploys or not all(isinstance(step, DeployAgentTaskStep) for step in deploys): - return "training needs the agent a step prompts deployed by a deploy_agent step" - agents = {step.agent_name for step in deploys} - if any(isinstance(step, RubricsVerifierTaskStep) and step.use_agent_judge and step.agent_name in agents - for step in steps): - return "its rubric judge is its solver, whose grading would be trained on" - return None - - -def _reward(task: Task, reward: str | Callable[[dict[str, float]], float] | None - ) -> Callable[[dict[str, float]], float | None]: - of, why = reward_of(task, reward) - if of is None: - raise ValueError(why) - return of - - def _index_named(prompt: Any, ids: list[str]) -> int: for message in prompt if isinstance(prompt, list) else []: if isinstance(message, dict) and (named := _NAMED.match(str(message.get("content") or ""))): @@ -230,10 +177,6 @@ def _index_named(prompt: Any, ids: list[str]) -> int: "same bundle, tasks and evals") -def _first_prompt(task: Task) -> str: - return next((step.prompt for step in task.steps if step.type == "prompt_agent"), "") - - def _free_port() -> int: with socket.socket() as probe: probe.bind(("", 0)) diff --git a/src/agent_env/openenv/reward.py b/src/agent_env/training/reward.py similarity index 95% rename from src/agent_env/openenv/reward.py rename to src/agent_env/training/reward.py index 08778e77..bf097763 100644 --- a/src/agent_env/openenv/reward.py +++ b/src/agent_env/training/reward.py @@ -1,4 +1,5 @@ -"""How a task's recorded scores make its reward: the rules ``build_app`` applies, which harness training shares.""" +"""How a task's recorded scores make its reward: the rules ``build_app`` applies, which training a bundle's own +agents shares.""" from __future__ import annotations diff --git a/src/agent_env/training/skyrl/env.py b/src/agent_env/training/skyrl/env.py index 43b93443..a6eb7318 100644 --- a/src/agent_env/training/skyrl/env.py +++ b/src/agent_env/training/skyrl/env.py @@ -16,7 +16,8 @@ raise ImportError("agent_env.training.skyrl.env runs in SkyRL's trainer: install 'agentenv-framework[skyrl]' into " "SkyRL's environment") from exc -from agent_env.training._openenv import CLOSE_SECONDS, GRACE_SECONDS, Rollout, Session, Shared, on_session_loop, token_in +from agent_env.training._loop import on_session_loop +from agent_env.training._openenv import CLOSE_SECONDS, GRACE_SECONDS, Rollout, Session, Shared, token_in from agent_env.training.skyrl.calls import Call, Unparsed, parse logger = logging.getLogger(__name__) diff --git a/src/agent_env/training/trl.py b/src/agent_env/training/trl.py index fc1de3f8..7627c17a 100644 --- a/src/agent_env/training/trl.py +++ b/src/agent_env/training/trl.py @@ -41,12 +41,12 @@ Shared, key, messages, - on_session_loop, probe, read_tasks, toolsets, why, ) +from agent_env.training._loop import on_session_loop from agent_env.training._tools import SchemaTool, tool_method logger = logging.getLogger(__name__) diff --git a/tst/areal/__init__.py b/tst/areal/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tst/areal/conftest.py b/tst/areal/conftest.py new file mode 100644 index 00000000..f06d6197 --- /dev/null +++ b/tst/areal/conftest.py @@ -0,0 +1,5 @@ +"""The unit tier's harness bundle, stand-in steps and stand-in model, which these tests play through AReaL's agent +workflow.""" + +from tst.unit.conftest import local_stores # noqa: F401 +from tst.unit.training.harness_test import bundle, registries # noqa: F401 diff --git a/tst/areal/workflow_test.py b/tst/areal/workflow_test.py new file mode 100644 index 00000000..56cac9ad --- /dev/null +++ b/tst/areal/workflow_test.py @@ -0,0 +1,107 @@ +"""``AgentEnvWorkflow`` on AReaL itself: AReaL's rollout engine resolves it from its import path into its proxy workflow +in inline mode, grouped as GRPO groups a row's samples, and each sample plays the unit tier's bundle, the solver +calling the stand-in model at the proxy address under that sample's session key; a played sample's reward goes on its +last completion. A sample whose run fails is left out of its group, its siblings still training, on an AReaL that asks +the agent; 2.1.0 doesn't, and rejects the group. CI runs it on AReaL 2.1.0 and main.""" + +import asyncio +import inspect +import logging +from unittest.mock import AsyncMock + +import pytest + +from agent_env.training.areal import task_rows +from tst.unit.training.harness_test import MODEL, asked +from tst.util.capabilities import AREAL, missing_capability_reason +from tst.util.fake_vllm import fake_vllm, says + +proxy = pytest.importorskip("areal.experimental.openai.proxy.workflow", reason=missing_capability_reason(AREAL)) +remote = pytest.importorskip("areal.infra.remote_inf_engine", reason=missing_capability_reason(AREAL)) +cli_args = pytest.importorskip("areal.api.cli_args", reason=missing_capability_reason(AREAL)) +workflow_context = pytest.importorskip("areal.infra.workflow_context", reason=missing_capability_reason(AREAL)) +openai_types = pytest.importorskip("areal.experimental.openai.types", reason=missing_capability_reason(AREAL)) + +WORKFLOW = "agent_env.training.areal.AgentEnvWorkflow" +#: Whether this AReaL asks an inline agent if a sample whose run failed leaves its group alone, or rejects the group. +REJECTS_A_SAMPLE = "should_reject_failed_sample" in inspect.getsource(proxy.OpenAIProxyWorkflow.arun_episode) + + +class _Session: + """AReaL's proxy client for one sample: the session key it hands the agent, and the reward set on the session's + last completion, which it exports.""" + + def __init__(self, *args, task_id, **kwargs): + self.session_api_key, self.session_id = f"key-{task_id}", f"session-{task_id}" + self.context_overflow = self.system_error = self.invalid_tool_json = False + self.interaction_count = 1 + self.reward = None + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc_info): + return None + + async def set_last_reward(self, reward): + self.reward = reward + + async def export_interactions(self, **kwargs): + return {self.session_id: openai_types.InteractionWithTokenLogpReward(reward=self.reward)} + + +@pytest.fixture +def areal(monkeypatch): + """Plays a row through the workflow AReaL's rollout engine resolves from ``WORKFLOW`` with ``kwargs``, a group of + ``group`` samples, its proxy at ``url``.""" + monkeypatch.setattr(proxy, "OpenAIProxyClient", _Session) + monkeypatch.setattr(proxy.OpenAIProxyWorkflow, "_grant_capacity", AsyncMock()) + monkeypatch.setattr(workflow_context, "get_httpx_client", AsyncMock(return_value=None)) + monkeypatch.setattr(workflow_context, "get_aiohttp_session", AsyncMock(return_value=None)) + config = cli_args.InferenceEngineConfig(agent=cli_args.AgentConfig(agent_cls_path=WORKFLOW, mode="inline")) + engine = remote.RemoteInfEngine(config, None) + engine.logger = logging.getLogger("areal-tier") + + def plays(row, url, group, **kwargs): + workflow = engine._resolve_workflow(WORKFLOW, kwargs, group_size=group, proxy_addr=url) + + async def episode(): + workflow_context.set(workflow_context.WorkflowContext(task_id=7)) + return await workflow.arun_episode(None, row) + + return asyncio.run(episode()) + + return plays + + +def _sum(bundle): + return next(row for row in task_rows(bundle) if row["task_id"].endswith("/sum")) + + +def test_each_sample_plays_the_task_under_its_own_session_and_its_reward_goes_on_its_last_completion(areal, bundle): + with fake_vllm(MODEL, says("It is 42.")) as engine: + played = areal(_sum(bundle), engine.url, 2, bundle=str(bundle)) + + assert {session: interaction.reward for session, interaction in played.items()} == { + "session-7:0": 1.0, "session-7:1": 1.0} + assert sorted(call["key"] for call in asked) == ["key-7:0", "key-7:1"] + assert {call["overrides"]["step_params"]["agent"]["litellm_base_url"] for call in asked} == {engine.url} + assert len(engine.requests) == 2 + + +def test_a_sample_whose_run_fails_is_rejected(areal, bundle): + answered = [] + + def second_fails(body): + answered.append(body) + if len(answered) == 2: + raise RuntimeError("the model server fell over") + return {"role": "assistant", "content": "It is 42."} + + with fake_vllm(MODEL, second_fails) as engine: + if REJECTS_A_SAMPLE: + played = areal(_sum(bundle), engine.url, 2, bundle=str(bundle)) + assert [interaction.reward for interaction in played.values()] == [1.0] + else: + with pytest.raises(ValueError): + areal(_sum(bundle), engine.url, 2, bundle=str(bundle)) diff --git a/tst/integration/training/areal_test.py b/tst/integration/training/areal_test.py new file mode 100644 index 00000000..8a9a2127 --- /dev/null +++ b/tst/integration/training/areal_test.py @@ -0,0 +1,31 @@ +"""``AgentEnvWorkflow`` on the opencode smoke task, as AReaL's agent workflow runs it in inline mode: opencode, in local +Docker, calls the model at the proxy address AReaL hands it, without /v1, under the sample's session key, a stand-in +model there having it run the task's shell command and then answer, and the task's verifier scores the run, with no +model endpoint configured. Needs a Docker daemon and network access for the agent's image.""" + +import asyncio +from pathlib import Path + +import pytest + +import agent_env +from agent_env.training.areal import AgentEnvWorkflow, task_rows +from tst.integration.training.harness_test import MODEL, state # noqa: F401 +from tst.util.fake_vllm import fake_vllm, shell_then_says + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] + + +def test_opencode_plays_its_smoke_task_against_the_proxy_address_and_the_run_returns_its_reward(state): # noqa: F811 + shipped = Path(agent_env.__file__).parent / "agents" / "tasks" / "opencode-smoke.json" + (state / "smoke" / "tasks").mkdir(parents=True) + (state / "smoke" / "tasks" / shipped.name).write_text(shipped.read_text()) + + with fake_vllm(MODEL, shell_then_says("echo $((6 * 7))", "42"), host="0.0.0.0") as engine: + [row] = task_rows(state / "smoke") + reward = asyncio.run(AgentEnvWorkflow(state / "smoke").run(row, base_url=engine.url, api_key="session-1", + http_client=None)) + + assert reward == 1.0 + assert any(message.get("role") == "tool" for body in engine.requests for message in body["messages"]) + assert {body["model"] for body in engine.requests} == {"default"} diff --git a/tst/unit/openenv/environment_test.py b/tst/unit/openenv/environment_test.py index 7f9ac0f3..dac58d33 100644 --- a/tst/unit/openenv/environment_test.py +++ b/tst/unit/openenv/environment_test.py @@ -12,9 +12,9 @@ from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction, ToolErrorType import agent_env.task.episode.run as episode_run -from agent_env.openenv.reward import _score_at from agent_env.openenv.environment import ENDED, NOT_RESET, TIMEOUT_METADATA, AgentEnvEnvironment, Served, ServedTask from agent_env.task import EpisodeError, EpisodeSpec, Task +from agent_env.training.reward import _score_at from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN, _gate, _until AGAIN = {**TURN, "id": "again", "prompt_id": "again", "prompt": "Add another.", "depends_on": ["turn"]} diff --git a/tst/unit/training/_support.py b/tst/unit/training/_support.py index 60237369..09ef9185 100644 --- a/tst/unit/training/_support.py +++ b/tst/unit/training/_support.py @@ -4,10 +4,10 @@ import pytest from agent_env.openenv.app import _app -from agent_env.openenv.reward import _score_at from agent_env.openenv.environment import Served, ServedTask from agent_env.task import EpisodeSpec, Task from agent_env.training._openenv import Server +from agent_env.training.reward import _score_at @pytest.fixture diff --git a/tst/unit/training/areal_test.py b/tst/unit/training/areal_test.py new file mode 100644 index 00000000..5bb8ae3d --- /dev/null +++ b/tst/unit/training/areal_test.py @@ -0,0 +1,133 @@ +"""``AgentEnvWorkflow`` running a bundle's task with its own agents, as AReaL's agent workflows run it in inline mode: +the solver's deploy alone is given AReaL's proxy, a stand-in here, as its model endpoint, under the key of the session +AReaL records the calls in, with no endpoint configured; the task's verifier scores the run.""" + +import asyncio +import json + +import pytest + +from agent_env.training import areal +from agent_env.training.areal import AgentEnvWorkflow, task_rows +from tst.unit.bundle._support import layout +from tst.unit.training.harness_test import ASIDE, ASK, CHECK, DEPLOY, asked, hanging + +#: What AReaL passes an inline agent besides its proxy and session key. +EXTRA = {"http_client": None, "proxy_gateway_api_key": "gateway"} + + +def _row(rows, name): + return next(row for row in rows if row["task_id"].endswith(f"/{name}")) + + +def test_the_solver_alone_calls_the_proxy_under_the_sessions_key_and_the_run_returns_its_reward(engine, bundle): + rows = task_rows(bundle) + + reward = asyncio.run(AgentEnvWorkflow(bundle).run(_row(rows, "sum"), base_url=engine.url, api_key="session-1", + **EXTRA)) + + [call] = asked + assert sorted(row["task_id"].rsplit("/", 1)[1] for row in rows) == ["broken", "hanging", "sum"] + assert _row(rows, "sum")["messages"] == [{"role": "user", "content": "What is 6 * 7?"}] + assert (call["model"], call["key"], call["answer"]) == ("default", "session-1", "It is 42.") + assert call["overrides"]["step_params"] == {"agent": {"litellm_base_url": engine.url, + "litellm_api_key": "session-1"}} + assert not {"litellm_base_url", "litellm_api_key"} & call["overrides"].keys() + assert reward == 1.0 and isinstance(reward, float) + + +def test_the_agent_named_is_the_one_the_solver_deploys(engine, bundle): + row = _row(task_rows(bundle, agent="mine"), "sum") + + asyncio.run(AgentEnvWorkflow(bundle, agent="mine").run(row, base_url=engine.url, api_key="session-1")) + + [call] = asked + assert call["overrides"]["step_params"]["agent"] == { + "a2a_agent_id": "mine", "a2a_agent_version": 1, "litellm_base_url": engine.url, "litellm_api_key": "session-1"} + + +def test_a_run_that_fails_raises_so_areal_leaves_the_sample_out(engine, bundle): + with pytest.raises(RuntimeError, match="the agent crashed"): + asyncio.run(AgentEnvWorkflow(bundle).run(_row(task_rows(bundle), "broken"), base_url=engine.url, api_key="k")) + + +def test_a_run_that_records_no_score_raises_so_areal_leaves_the_sample_out(engine, bundle, tmp_path): + unchecked = layout(tmp_path / "unchecked", {"tasks/sum.json": json.dumps([DEPLOY, ASK])}) + + with pytest.raises(ValueError, match="records no score to train on"): + asyncio.run(AgentEnvWorkflow(unchecked).run(_row(task_rows(unchecked), "sum"), base_url=engine.url, + api_key="k")) + assert len(asked) == 1 + + +def test_a_run_with_a_step_that_failed_without_failing_the_task_raises_so_areal_leaves_the_sample_out(engine, bundle, + tmp_path): + aside = layout(tmp_path / "aside", {"tasks/sum.json": json.dumps([DEPLOY, ASK, CHECK, ASIDE])}) + + with pytest.raises(RuntimeError, match="the run failed at aside: the side check fell over"): + asyncio.run(AgentEnvWorkflow(aside, reward="check").run(task_rows(aside)[0], base_url=engine.url, + api_key="k")) + assert asked[0]["answer"] == "It is 42." + + +def test_a_reward_function_need_not_be_hashable(bundle): + class Unhashable: + __hash__ = None + + def __call__(self, scores): + return scores["check"] + + assert AgentEnvWorkflow(bundle, reward=Unhashable())._played + + +@pytest.fixture +def torn(monkeypatch): + """The task ids of the runs torn down, each recorded once its teardown is done.""" + done, teardown = [], areal.teardown_run + + async def tears_down(context): + report = await teardown(context) + done.append(context.metadata["task_id"]) + return report + + monkeypatch.setattr(areal, "teardown_run", tears_down) + return done + + +def test_a_cancelled_run_is_stopped_and_torn_down_before_the_cancel_returns(engine, bundle, torn): + workflow = AgentEnvWorkflow(bundle) + row = _row(task_rows(bundle), "hanging") + + async def cancelled(): + running = asyncio.ensure_future(workflow.run(row, base_url=engine.url, api_key="k")) + assert await asyncio.to_thread(hanging.wait, 30) + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running + return list(torn) + + assert [task_id.rsplit("/", 1)[1] for task_id in asyncio.run(cancelled())] == ["hanging"] + + +def test_a_row_naming_no_task_of_the_workflows_is_refused(engine, bundle): + with pytest.raises(KeyError, match="build the dataset with task_rows"): + asyncio.run(AgentEnvWorkflow(bundle).run({"task_id": "another-bundle/sum"}, base_url=engine.url, api_key="k")) + + +def test_a_process_prepares_the_bundle_once_however_often_areal_builds_the_workflow(bundle, monkeypatch): + prepared, trained = [], areal.trained + monkeypatch.setattr(areal, "trained", lambda *args: prepared.append(args) or trained(*args)) + + AgentEnvWorkflow(bundle, tasks=["sum"]) + AgentEnvWorkflow(bundle, tasks=["sum"]) + AgentEnvWorkflow(bundle, tasks=["sum", "broken"]) + + assert [args[1] for args in prepared] == [("sum",), ("sum", "broken")] + + +def test_a_failed_sample_leaves_its_group_alone_unless_areal_found_the_model_at_fault(bundle): + workflow = AgentEnvWorkflow(bundle) + + assert workflow.should_reject_failed_sample(RuntimeError("the agent crashed"), "unknown_failure_reject") + assert workflow.should_reject_failed_sample(RuntimeError("the sandbox went"), None) + assert not workflow.should_reject_failed_sample(RuntimeError("context overflow"), "model_failure_zero") diff --git a/tst/unit/training/conftest.py b/tst/unit/training/conftest.py index 8183c3c0..8a95494f 100644 --- a/tst/unit/training/conftest.py +++ b/tst/unit/training/conftest.py @@ -1,7 +1,9 @@ """The Episode API's fake steps, stand-in teardown and tool sessions, which these tests serve as an OpenEnv env on -loopback, and the loopback connections that takes.""" +loopback, and the loopback connections that takes; and the bundle with stand-in steps and model that the helpers +training a bundle's own agents play.""" from tst.unit.openenv.environment_test import tools # noqa: F401 from tst.unit.task.episode.conftest import _disable_network # noqa: F401 from tst.unit.task.episode.run_test import teardowns # noqa: F401 from tst.unit.training._support import serve # noqa: F401 +from tst.unit.training.harness_test import bundle, engine, registries # noqa: F401 diff --git a/tst/unit/training/harness_test.py b/tst/unit/training/harness_test.py index defddc23..b877d84c 100644 --- a/tst/unit/training/harness_test.py +++ b/tst/unit/training/harness_test.py @@ -20,7 +20,7 @@ from agent_env.task_step.task_step import TaskStep from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep -from agent_env.training import harness +from agent_env.training import _solver, harness from agent_env.training.harness import AgentEnvSessionFactory, prompt_rows from tst.unit.bundle._support import layout from tst.unit.conftest import local_stores # noqa: F401 @@ -31,6 +31,9 @@ DEPLOY = {"id": "agent", "type": "deploy_agent", "env_ids": [], "a2a_agent_id": "theirs"} ASK = {"id": "ask", "type": "prompt_agent", "prompt_id": "ask", "prompt": "What is 6 * 7?", "depends_on": ["agent"]} CHECK = {"id": "check", "type": "harness_test_check", "verifier_id": "answer", "depends_on": ["ask"]} +#: A side check that fails without failing the task, after which the task still records its answer's score. +ASIDE = {"id": "aside", "type": "harness_test_check", "verifier_id": "fails", "depends_on": ["ask"], + "fail_task_on_error": False} JUDGE = {"id": "judge", "type": "rubrics_verifier", "verifier_id": "judge", "prompt_id": "ask", "judge_a2a_agent_id": "theirs", "agent_name": DeployAgentTaskStep.DEFAULT_AGENT_NAME, "criteria": [{"id": "c1", "description": "says 42"}], "depends_on": ["ask"]} @@ -84,6 +87,8 @@ def from_dict(cls, data): return cls(**cls._base_from_dict(data), verifier_id=data["verifier_id"]) async def execute(self, context): + if self.verifier_id == "fails": + raise RuntimeError("the side check fell over") context.metadata.setdefault("verifications", {})[self.verifier_id] = { "score": float("42" in context.metadata["answer"])} return context @@ -95,14 +100,14 @@ def engine(): yield server -@pytest.fixture(autouse=True) +@pytest.fixture def registries(monkeypatch): steps = {**Config().task_step_registry(), "deploy_agent": _Deploys, "prompt_agent": _Asks, _Checks.type: _Checks} monkeypatch.setattr(Config, "task_step_registry", lambda self: steps) @pytest.fixture -def bundle(tmp_path, monkeypatch, local_stores): # noqa: F811 +def bundle(tmp_path, monkeypatch, local_stores, registries): # noqa: F811 monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.setattr(preflight, "model_endpoint_missing", MODEL_ENDPOINT_MISSING) monkeypatch.delenv("LITELLM_BASE_URL", raising=False) @@ -123,7 +128,7 @@ def bundle(tmp_path, monkeypatch, local_stores): # noqa: F811 def test_the_solver_alone_calls_the_policy_through_the_proxy_and_the_run_is_captured_and_scored(engine, bundle, caplog): - with caplog.at_level(logging.WARNING, logger=harness.__name__): + with caplog.at_level(logging.WARNING, logger=_solver.__name__): rows = prompt_rows(bundle) factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, sampling={"temperature": 1.0}) engine.served.append("trl-policy-v3") @@ -176,6 +181,18 @@ def test_a_run_that_fails_is_unscored_rather_than_zero(engine, bundle, caplog): assert "failed, so its rollout is left unscored: " in caplog.text +def test_a_run_with_a_step_that_failed_without_failing_the_task_is_unscored(engine, bundle, tmp_path, caplog): + aside = layout(tmp_path / "aside", {"tasks/sum.json": json.dumps([DEPLOY, ASK, CHECK, ASIDE])}) + factory = AgentEnvSessionFactory(aside, llm_url=engine.url, model=MODEL, reward="check") + session = factory.create(prompt_rows(aside)[0]["prompt"]) + + with caplog.at_level(logging.WARNING, logger=harness.__name__): + assert session.wait_for_completion() == 1 + session.close() + assert session.verify([]).env_reward is None + assert "the run failed at aside: the side check fell over" in caplog.text + + @pytest.fixture def torn(monkeypatch): """The task ids of the runs torn down, each recorded once its teardown is done.""" diff --git a/tst/unit/training/trl_test.py b/tst/unit/training/trl_test.py index dd786834..f90c96f5 100644 --- a/tst/unit/training/trl_test.py +++ b/tst/unit/training/trl_test.py @@ -18,7 +18,8 @@ from agent_env.openenv.environment import AgentEnvEnvironment from agent_env.task_step.task_step import TaskStep from agent_env.training import _openenv, trl -from agent_env.training._openenv import ENDED, FAILED, on_session_loop +from agent_env.training._loop import on_session_loop +from agent_env.training._openenv import ENDED, FAILED from agent_env.training.trl import AgentEnv from tst.unit.bundle._support import layout from tst.unit.openenv.environment_test import AGAIN diff --git a/tst/util/capabilities.py b/tst/util/capabilities.py index 2eddbb72..4ebc8659 100644 --- a/tst/util/capabilities.py +++ b/tst/util/capabilities.py @@ -53,6 +53,10 @@ #: and torch aren't in the dev extra, and only that tier's CI job installs them. SKYRL = "skyrl_installed" +#: The AReaL tier (``tst/areal``) runs ``AgentEnvWorkflow`` in AReaL's agent workflow; AReaL and torch aren't in the dev +#: extra, and only that tier's CI job installs them. +AREAL = "areal_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/util/fake_vllm.py b/tst/util/fake_vllm.py index 9a57e132..4e596312 100644 --- a/tst/util/fake_vllm.py +++ b/tst/util/fake_vllm.py @@ -1,20 +1,22 @@ """A stand-in for vLLM served with ``--return-tokens-as-token-ids`` and processed logprobs, as OpenEnv's capture proxy needs its engine: it lists the models it serves and answers chat completions with token ids, a logprob per completion -token and the top two at each, scaled by the temperature as processed logprobs are. ``reply`` makes each answer's -message from the request body.""" +token and the top two at each, scaled by the temperature as processed logprobs are, or streamed, as an agent's own +client asks for them. ``reply`` makes each answer's message from the request body.""" import json +import threading +import time import uuid -from collections.abc import Callable, Iterator +from collections.abc import AsyncIterator, Callable, Iterator from contextlib import contextmanager +from types import SimpleNamespace from typing import Any +import uvicorn from starlette.applications import Starlette -from starlette.responses import JSONResponse +from starlette.responses import JSONResponse, StreamingResponse from starlette.routing import Route -from agent_env.training._openenv import Server - Reply = Callable[[dict[str, Any]], dict[str, Any]] @@ -49,9 +51,10 @@ def token_ids(message: dict[str, Any]) -> list[int]: @contextmanager -def fake_vllm(model: str, reply: Reply) -> Iterator[Server]: - """The engine, served on loopback; ``.requests`` are the completion bodies it got, and ``.served`` the models it - lists, which a test appends to as vLLM lists a LoRA adapter once it's loaded.""" +def fake_vllm(model: str, reply: Reply, host: str = "127.0.0.1") -> Iterator[SimpleNamespace]: + """The engine, listening on ``host``, every interface for a client in a container, at ``.url`` on loopback; + ``.requests`` are the completion bodies it got, and ``.served`` the models it lists, which a test appends to as vLLM + lists a LoRA adapter once it's loaded.""" requests: list[dict[str, Any]] = [] served = [model] @@ -69,19 +72,40 @@ async def completions(request): "top_logprobs": [{"token": f"token_id:{token}", "logprob": -0.5 / temperature}, {"token": f"token_id:{token + 1}", "logprob": -3.0 / temperature}]} for token in completion] + answer = {"id": f"chatcmpl-{uuid.uuid4().hex[:8]}", "model": body["model"]} + finish = "tool_calls" if message.get("tool_calls") else "stop" + usage = {"prompt_tokens": len(prompt), "completion_tokens": len(completion), + "total_tokens": len(prompt) + len(completion)} + if body.get("stream"): + return StreamingResponse(_chunks(answer, message, finish, usage), media_type="text/event-stream") return JSONResponse({ - "id": f"chatcmpl-{uuid.uuid4().hex[:8]}", "object": "chat.completion", "model": body["model"], - "prompt_token_ids": prompt, - "choices": [{"index": 0, "message": message, "finish_reason": "tool_calls" if message.get("tool_calls") - else "stop", "token_ids": completion, "logprobs": {"content": content}}], - "usage": {"prompt_tokens": len(prompt), "completion_tokens": len(completion), - "total_tokens": len(prompt) + len(completion)}, + **answer, "object": "chat.completion", "prompt_token_ids": prompt, + "choices": [{"index": 0, "message": message, "finish_reason": finish, "token_ids": completion, + "logprobs": {"content": content}}], + "usage": usage, }) - server = Server(Starlette(routes=[Route("/v1/models", models), - Route("/v1/chat/completions", completions, methods=["POST"])])) - server.requests, server.served = requests, served + app = Starlette(routes=[Route("/v1/models", models), Route("/v1/chat/completions", completions, methods=["POST"])]) + server = uvicorn.Server(uvicorn.Config(app, host=host, port=0, log_level="warning")) + thread = threading.Thread(target=server.run, name="fake-vllm", daemon=True) + thread.start() + while not server.started: + time.sleep(0.05) + port = server.servers[0].sockets[0].getsockname()[1] try: - yield server + yield SimpleNamespace(url=f"http://127.0.0.1:{port}", requests=requests, served=served) finally: - server.stop() + server.should_exit = True + thread.join() + + +async def _chunks(answer: dict[str, Any], message: dict[str, Any], finish: str, + usage: dict[str, int]) -> AsyncIterator[str]: + """The answer as a streamed completion: the whole message in one delta, then its finish reason and usage.""" + delta = {"role": "assistant", "content": message.get("content")} + if message.get("tool_calls"): + delta["tool_calls"] = [{"index": index, **call} for index, call in enumerate(message["tool_calls"])] + for choice, extra in (({"index": 0, "delta": delta, "finish_reason": None}, {}), + ({"index": 0, "delta": {}, "finish_reason": finish}, {"usage": usage})): + yield f"data: {json.dumps({**answer, 'object': 'chat.completion.chunk', 'choices': [choice], **extra})}\n\n" + yield "data: [DONE]\n\n"