From 5a7c4074334f499e50bebb70d12e7fdf8ed60ecd Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 18:30:54 -0700 Subject: [PATCH 1/8] feat(training): train with SkyRL through an AgentEnv env and generator agent_env.training.skyrl, behind a skyrl extra: a writer of SkyRL's parquet rows with each row's tools rendered into its system prompt by the model's chat template, a skyrl-gym BaseTextEnv that plays the row's task as one episode and parses the model's tool calls (hermes, qwen3_coder), a SkyRLGymGenerator that replays a rollout the env fails and masks the group of one that fails every attempt, and an entrypoint that registers both in SkyRL's trainer. The session loop, the token lookup and the MCP-to-function conversion move to the shared session layer. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/skyrl.yml | 66 ++++ AGENTS.md | 5 +- Makefile | 8 +- pyproject.toml | 10 +- src/agent_env/training/__init__.py | 3 +- src/agent_env/training/_openenv.py | 42 ++- src/agent_env/training/skyrl/__init__.py | 5 + src/agent_env/training/skyrl/__main__.py | 47 +++ src/agent_env/training/skyrl/calls.py | 105 +++++++ src/agent_env/training/skyrl/env.py | 137 ++++++++ src/agent_env/training/skyrl/generator.py | 92 ++++++ src/agent_env/training/skyrl/prepare.py | 132 ++++++++ src/agent_env/training/skyrl/train.py | 77 +++++ src/agent_env/training/trl.py | 35 +-- src/agent_env/training/verl/loop.py | 3 +- src/agent_env/training/verl/prepare.py | 21 +- tst/skyrl/__init__.py | 0 tst/skyrl/conftest.py | 7 + tst/skyrl/generator_test.py | 296 ++++++++++++++++++ .../test_optional_extras_stay_optional.py | 13 +- tst/unit/training/skyrl_test.py | 282 +++++++++++++++++ tst/unit/training/trl_test.py | 6 +- tst/util/capabilities.py | 4 + uv.lock | 13 +- 24 files changed, 1345 insertions(+), 64 deletions(-) create mode 100644 .github/workflows/skyrl.yml create mode 100644 src/agent_env/training/skyrl/__init__.py create mode 100644 src/agent_env/training/skyrl/__main__.py create mode 100644 src/agent_env/training/skyrl/calls.py create mode 100644 src/agent_env/training/skyrl/env.py create mode 100644 src/agent_env/training/skyrl/generator.py create mode 100644 src/agent_env/training/skyrl/prepare.py create mode 100644 src/agent_env/training/skyrl/train.py create mode 100644 tst/skyrl/__init__.py create mode 100644 tst/skyrl/conftest.py create mode 100644 tst/skyrl/generator_test.py create mode 100644 tst/unit/training/skyrl_test.py diff --git a/.github/workflows/skyrl.yml b/.github/workflows/skyrl.yml new file mode 100644 index 000000000..fcbc37dd9 --- /dev/null +++ b/.github/workflows/skyrl.yml @@ -0,0 +1,66 @@ +name: SkyRL contract + +# AgentEnvGenerator and AgentEnvTextEnv on SkyRL itself (tst/skyrl): SkyRL's PromptDataset loads rows the writer wrote, +# with Qwen3's and Qwen3.5's tokenizers, and its generator plays them through a stand-in inference engine, on SkyRL 0.4.0 +# and on its main branch. SkyRL, skyrl-gym, 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/skyrl/**" + - "tst/unit/training/**" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/skyrl.yml" + schedule: + - cron: "53 7 * * 1" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: skyrl-${{ 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: + generator: + name: env and generator (SkyRL ${{ matrix.skyrl }}) + runs-on: ubuntu-latest + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + include: + - skyrl: "0.4.0" + requirement: "skyrl[skyrl-train]==0.4.0 skyrl-gym==0.4.0" + - skyrl: "main" + requirement: >- + skyrl[skyrl-train]@git+https://github.com/NovaSky-AI/SkyRL@main + skyrl-gym@git+https://github.com/NovaSky-AI/SkyRL@main#subdirectory=skyrl-gym + 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 skyrl extra + run: | + python -m pip install --quiet -r .github/ci-requirements.txt + uv sync --locked --extra dev --extra skyrl --python "$(command -v python)" + - name: Install SkyRL, skyrl-gym 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('skyrl', version('skyrl'), 'skyrl-gym', version('skyrl-gym'), 'torch', version('torch'))" + - name: SkyRL tier + run: .venv/bin/python -m pytest tst/skyrl -rA --timeout=900 --junitxml=junit-skyrl.xml + - name: No skips + if: always() + run: .venv/bin/python .github/scripts/check_skip_policy.py junit-skyrl.xml diff --git a/AGENTS.md b/AGENTS.md index 3adbc5780..1e1849646 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -47,6 +47,7 @@ or one of them only where it is used. | `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 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 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 | @@ -122,7 +123,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. | +| `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. | | `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. | @@ -339,6 +340,8 @@ against `serve`, on verifiers 0.4.0 and its main branch, installed in a venv of 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/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/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 87d9b1c2d..1b557f909 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 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 clean-install-test test clean VENV := .venv VENV_BIN := $(abspath $(VENV))/bin @@ -79,6 +79,12 @@ nemo-gym-test: ## Play rollouts through NeMo Gym's resources server for a bundle 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 SkyRL tier plays rows through AgentEnvGenerator and AgentEnvTextEnv on SkyRL itself, with Qwen3's and Qwen3.5's +# tokenizers from the Hub. SkyRL, skyrl-gym and torch aren't in the dev extra, and SkyRL needs Python 3.12 or later: +# install them first, as .github/workflows/skyrl.yml does. +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 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 70ffafe29..2101b0be3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -77,6 +77,14 @@ verl = [ "pyarrow>=19", "transformers>=5", ] +# agent_env.training.skyrl: SkyRL's trainers on a bundle served as an OpenEnv env; installed into SkyRL's own +# environment, whose SkyRL, skyrl-gym, vLLM and torch are the trainer's to pin. pyarrow writes the parquet rows SkyRL +# reads; transformers renders a row's tools into its system prompt with the model's chat template. +skyrl = [ + "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", @@ -125,7 +133,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,verl]", # exercise portable optional providers and stores in CI + "agentenv-framework[aws,boat,boxd,coreweave,daytona,gcp,openenv,sail,skyrl,tensorlake,tenki,trl,vercel,verl]", # exercise portable optional providers and stores in CI "fastmcp>=3.4,<4", # the standalone-fastmcp mount path of the protocol; <4 keeps mcp 1.x "Pillow>=10.4", # the protocol's computer extra, whose tests the unit tier runs "av>=17", # the protocol's computer extra, whose tests the unit tier runs and which the recording tests decode with diff --git a/src/agent_env/training/__init__.py b/src/agent_env/training/__init__.py index 1440ad6e6..ed79a7834 100644 --- a/src/agent_env/training/__init__.py +++ b/src/agent_env/training/__init__.py @@ -1,3 +1,4 @@ """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, and ``agent_env.training.verl`` for verl.""" +environments, NeMo RL among them, ``agent_env.training.verl`` for verl, and ``agent_env.training.skyrl`` for +SkyRL.""" diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index 2e226cac4..fa7181a61 100644 --- a/src/agent_env/training/_openenv.py +++ b/src/agent_env/training/_openenv.py @@ -8,11 +8,12 @@ import contextlib import json import logging +import os import re import threading import time import uuid -from collections.abc import Mapping +from collections.abc import Coroutine, Mapping from dataclasses import dataclass, field from typing import Any, ClassVar @@ -47,6 +48,8 @@ logger = logging.getLogger(__name__) _waited: set[str] = set() +_loop: asyncio.AbstractEventLoop | None = None +_loop_lock = threading.Lock() @dataclass(frozen=True) @@ -416,6 +419,43 @@ 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: + 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 function_tool(tool: Mapping[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": {}}, + }} + + def authorization(token: str | None) -> dict[str, str]: return {"Authorization": f"Bearer {token}"} if token else {} diff --git a/src/agent_env/training/skyrl/__init__.py b/src/agent_env/training/skyrl/__init__.py new file mode 100644 index 000000000..b3b882f0f --- /dev/null +++ b/src/agent_env/training/skyrl/__init__.py @@ -0,0 +1,5 @@ +"""SkyRL'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.skyrl URL OUT`` writes the server's tasks as SkyRL's parquet rows, each +with its tools in its system prompt, ``agent_env.training.skyrl.env.AgentEnvTextEnv`` plays each row's task as one +episode on the server, ``agent_env.training.skyrl.generator.AgentEnvGenerator`` plays again a rollout the env fails, +and ``python -m agent_env.training.skyrl.train`` runs SkyRL's trainer with both.""" diff --git a/src/agent_env/training/skyrl/__main__.py b/src/agent_env/training/skyrl/__main__.py new file mode 100644 index 000000000..dd372a8e7 --- /dev/null +++ b/src/agent_env/training/skyrl/__main__.py @@ -0,0 +1,47 @@ +"""``python -m agent_env.training.skyrl URL OUT --tokenizer MODEL``: writes a bundle's tasks as SkyRL's RL dataset.""" + +import json + +import click + +from agent_env.training.skyrl.calls import FORMATS +from agent_env.training.skyrl.prepare import write_dataset + + +@click.command() +@click.argument("url") +@click.argument("out", type=click.Path(file_okay=False)) +@click.option("--tokenizer", required=True, + help="The model's tokenizer, as transformers loads it, whose chat template renders each row's tools into " + "its system prompt and counts its prompt.") +@click.option("--tool-format", type=click.Choice(FORMATS), default="hermes", show_default=True, + help="The tool-call format the model's chat template teaches: hermes for Qwen2.5 and Qwen3, qwen3_coder " + "for Qwen3.5 and Qwen3-Coder.") +@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 by " + "the trainer.") +@click.option("--max-prompt-length", type=int, + help="Refuse rows whose prompt with its tools is longer, in tokens: SkyRL's trainer.max_prompt_length.") +@click.option("--chat-template-kwargs", type=json.loads, + help="The chat template's settings as JSON, SkyRL's generator.chat_template_kwargs, to render prompts as " + 'SkyRL does, e.g. \'{"enable_thinking": false}\'.') +def main(url: str, out: str, tokenizer: str, tool_format: str, tool_timeout: float, reset_timeout: float, + attempts: int, token_var: 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, each with the tools its + task offers in its system prompt and the server's settings for the env that plays it.""" + written = write_dataset(url, out, tokenizer=tokenizer, tool_format=tool_format, tool_timeout=tool_timeout, + reset_timeout=reset_timeout, attempts=attempts, token_var=token_var, + 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}") + 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/skyrl/calls.py b/src/agent_env/training/skyrl/calls.py new file mode 100644 index 000000000..5de83e23b --- /dev/null +++ b/src/agent_env/training/skyrl/calls.py @@ -0,0 +1,105 @@ +"""The tool calls in a model's reply, in the format its chat template teaches it: ``hermes`` (Qwen2.5, Qwen3), a JSON +object of name and arguments in each ```` block, or ``qwen3_coder`` (Qwen3.5, Qwen3-Coder), a +```` with a ```` block per argument, each typed by the tool's JSON Schema.""" + +from __future__ import annotations + +import json +import re +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +FORMATS = ("hermes", "qwen3_coder") + +_BLOCK = re.compile(r"(.*?)(?:|$)", re.DOTALL) +_THINKING = re.compile(r".*?", re.DOTALL) +_FUNCTION = re.compile(r"\s]+)>(.*?)(?:|$)", re.DOTALL) +_PARAMETER = re.compile(r"\s]+)>(.*?)(?:|(?=`` block that isn't a call of one of the tools, and what the model is told of it.""" + + why: str + + +def parse(reply: str, tool_format: str, tools: Mapping[str, Mapping[str, Any]]) -> tuple[list[Call | Unparsed], str]: + """Each ```` block of ``reply`` in order, as a call of one of ``tools`` (each function by its name) or + why it isn't one, and the reply's text outside them and its thinking, which ends a turn.""" + if tool_format not in FORMATS: + raise ValueError(f"tool_format must be one of {', '.join(FORMATS)}, not {tool_format!r}") + reading = _hermes if tool_format == "hermes" else _qwen3_coder + calls = [_checked(reading(block.strip(), tools), tools) for block in _BLOCK.findall(reply)] + return calls, _THINKING.sub("", _BLOCK.sub("", reply)).strip() + + +def _hermes(block: str, tools: Mapping[str, Mapping[str, Any]]) -> Call | Unparsed: + try: + call = json.loads(block) + except json.JSONDecodeError as e: + return Unparsed(f"the tool call isn't JSON ({e}): write it as " + '{"name": , "arguments": }') + if not isinstance(call, dict) or not isinstance(call.get("name"), str): + return Unparsed('the tool call names no tool: write it as {"name": , "arguments": }') + arguments = call.get("arguments", {}) + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except json.JSONDecodeError: + pass + if not isinstance(arguments, dict): + return Unparsed(f"the arguments of {call['name']} aren't a JSON object") + return Call(call["name"], arguments) + + +def _qwen3_coder(block: str, tools: Mapping[str, Mapping[str, Any]]) -> Call | Unparsed: + if not (function := _FUNCTION.search(block)): + return Unparsed("the tool call names no tool: write it as with a " + " block per argument") + name, body = function[1], function[2] + properties = ((tools.get(name) or {}).get("parameters") or {}).get("properties") or {} + return Call(name, {argument: _typed(_trimmed(value), properties.get(argument) or {}) + for argument, value in _PARAMETER.findall(body)}) + + +def _checked(call: Call | Unparsed, tools: Mapping[str, Mapping[str, Any]]) -> Call | Unparsed: + if isinstance(call, Call) and call.name not in tools: + return Unparsed(f"there's no tool {call.name!r}; the tools are {', '.join(sorted(tools))}") + return call + + +def _trimmed(value: str) -> str: + """A parameter's value without the one newline the format puts on each side of it.""" + return value.removeprefix("\n").removesuffix("\n") + + +def _typed(value: str, schema: Mapping[str, Any]) -> Any: + """``value`` as the type its schema names, the first of a union that it parses as; as written when it parses as + none, or its schema names no type and it isn't JSON.""" + kinds = schema.get("type") + for kind in kinds if isinstance(kinds, list) else [kinds]: + if kind == "string": + return value + try: + parsed = json.loads(value) + except json.JSONDecodeError: + continue + if kind is None or _is(parsed, kind): + return parsed + return value + + +def _is(value: Any, kind: str) -> bool: + if kind == "integer": + return isinstance(value, int) and not isinstance(value, bool) + if kind == "number": + return isinstance(value, (int, float)) and not isinstance(value, bool) + return isinstance(value, {"boolean": bool, "array": list, "object": dict, "null": type(None)}.get(kind, object)) diff --git a/src/agent_env/training/skyrl/env.py b/src/agent_env/training/skyrl/env.py new file mode 100644 index 000000000..75427a995 --- /dev/null +++ b/src/agent_env/training/skyrl/env.py @@ -0,0 +1,137 @@ +"""SkyRL's env for the rows ``python -m agent_env.training.skyrl`` writes. ``AgentEnvTextEnv`` is a skyrl-gym +``BaseTextEnv`` playing a row's task as one episode on the bundle's OpenEnv server: each step runs the tool calls in the +model's reply on the episode and answers with their results, and a reply without one ends the turn it left open, which +scores the episode once its last turn ends. SkyRL calls an env from a thread of its pool, so the episode's session +runs on the event loop this process's sessions share.""" + +from __future__ import annotations + +import json +import logging +from typing import Any + +try: + from skyrl_gym.envs.base_text_env import BaseTextEnv, BaseTextEnvStepOutput, ConversationType +except ModuleNotFoundError as exc: + if exc.name is None or exc.name.split(".")[0] != "skyrl_gym": + raise + 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.skyrl.calls import Call, Unparsed, parse + +logger = logging.getLogger(__name__) + + +class EpisodeFailed(RuntimeError): + """The env failed the episode: a deploy, a call past its deadline, or a server that went away. AgentEnvGenerator + plays the rollout again on a new episode.""" + + +class AgentEnvTextEnv(BaseTextEnv): + """Plays the row's task as one episode on the OpenEnv server its row names. ``init`` starts the episode, and + returns before it is deployed, so the deploy overlaps the model's first turn. Each ``step`` runs the tool calls of + the model's reply in order, a tool call SkyRL's parse finds nothing to call with answered with why, and returns + their answers as ``tool`` messages; a reply without a tool call ends the turn it left open, and any after it end + empty, as does the reply of the ``generator.max_turns``-th step, and the step that ends the episode's last turn + returns its reward. A failure of the env raises EpisodeFailed; a task that ends unscored without failing, or offers + a tool its row doesn't list, is a configuration error no replay changes, and raises ValueError. + + ``get_metrics``, which SkyRL averages per step and per eval as ``environment/...``, holds each of the episode's + scores by the id of the step that recorded it (``score/``), the tool calls run and the unparsed ones.""" + + def __init__(self, env_config: Any = None, extras: dict[str, Any] | None = None): + super().__init__() + extras = extras or {} + self._episode, server = extras["episode"], extras["server"] + self.max_turns = int(extras.get("max_turns") or 1) + self._tools = {tool["function"]["name"]: tool["function"] for tool in json.loads(self._episode["tools"])} + self._format = server["tool_format"] + self._waits = server["reset_timeout"] + max(server["reset_timeout"], server["tool_timeout"]) \ + + 2 * GRACE_SECONDS + CLOSE_SECONDS + self._session = _Session(Shared(server["url"], frozenset(self._tools), server["tool_timeout"], + server["reset_timeout"], patience=server["reset_timeout"], + waits_for_server=True, token=token_in(server.get("token_var")))) + self._scores: dict[str, float] = {} + self._run = self._unparsed = 0 + + def init(self, prompt: ConversationType) -> tuple[ConversationType, dict[str, Any]]: + self._wait(self._session.start(self._episode)) + return prompt, {} + + def step(self, action: str) -> BaseTextEnvStepOutput: + self.turns += 1 + calls, reply = parse(action, self._format, self._tools) + answers = [] + for call in calls: + if isinstance(call, Unparsed): + self._unparsed += 1 + answers.append({"role": "tool", "content": f"Error: {call.why}"}) + continue + self._run += 1 + answers.append({"role": "tool", "content": self._wait(self._session.call(call))}) + if self._session.rollout.done or self._session.rollout.failure: + break + rollout = self._session.rollout + if not (rollout.done or rollout.failure) and (not calls or self.turns >= self.max_turns): + self._wait(self._session.finish(reply)) + if rollout.failure is not None: + self.close() + if rollout.misgrouped: + raise ValueError(f"{self._task()}: {rollout.failure}. Its row lists the tools of its toolset's first " + "task, so a task must offer no tool that one doesn't") + raise EpisodeFailed(f"{self._task()}: {rollout.failure}") + if not rollout.done: + return BaseTextEnvStepOutput(observations=answers, reward=0.0, done=False, metadata={}) + if rollout.reward is None: + self.close() + 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") + self._scores = rollout.scores or {} + return BaseTextEnvStepOutput(observations=[], reward=float(rollout.reward), done=True, metadata={}) + + def get_metrics(self) -> dict[str, Any]: + return {"tool_calls": self._run, "unparsed_calls": self._unparsed, + **{f"score/{step}": score for step, score in self._scores.items()}} + + def close(self) -> None: + on_session_loop(self._session.close(), CLOSE_SECONDS) + + def _wait(self, coroutine: Any) -> Any: + """The session's answer; one that takes past every deadline it has, and a grace, fails the episode.""" + try: + return on_session_loop(coroutine, self._waits) + except TimeoutError as e: + self.close() + raise EpisodeFailed(f"{self._task()}: the env didn't answer within {self._waits:g}s") from e + + def _task(self) -> str: + return f"task {self._episode['task_id']} of split {self._episode['split']}" + + +class _Session(Session): + """One rollout's session: it starts the row's episode, runs the model's tool calls on it, 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 + + @property + def rollout(self) -> Rollout: + return self._rollout + + async def start(self, episode: dict[str, Any]) -> None: + await self._start(episode["split"], int(episode["index"]), episode["task_id"]) + + async def call(self, call: Call) -> str: + return await self._step(call.name, call.arguments) + + async def finish(self, reply: str) -> None: + await self._finish(self._rollout, reply) + + async def close(self) -> None: + await self._close() diff --git a/src/agent_env/training/skyrl/generator.py b/src/agent_env/training/skyrl/generator.py new file mode 100644 index 000000000..a3a076212 --- /dev/null +++ b/src/agent_env/training/skyrl/generator.py @@ -0,0 +1,92 @@ +"""SkyRL's generator for the rows ``python -m agent_env.training.skyrl`` writes, which plays a rollout the env fails +again and, when every attempt fails, leaves its group out of training: SkyRL fails a whole batch on an env's exception, +and has no way to leave out one trajectory.""" + +from __future__ import annotations + +import logging +import time +from dataclasses import asdict +from typing import Any + +try: + from skyrl.train.generators.base import TRAINING_PHASE_EVAL, GeneratorInput, GeneratorOutput + from skyrl.train.generators.skyrl_gym_generator import SkyRLGymGenerator, StepWiseOutput, TrajectoryOutput +except ModuleNotFoundError as exc: + if exc.name is None or exc.name.split(".")[0] != "skyrl": + raise + raise ImportError("agent_env.training.skyrl.generator runs in SkyRL's trainer: install " + "'agentenv-framework[skyrl]' into SkyRL's environment") from exc + +from agent_env.training.skyrl.env import EpisodeFailed +from agent_env.training.skyrl.prepare import AGENT + +logger = logging.getLogger(__name__) + + +class AgentEnvGenerator(SkyRLGymGenerator): + """SkyRL's ``SkyRLGymGenerator`` that plays a rollout of an ``agentenv`` row the env fails (EpisodeFailed) again + from the start on a new episode, a fresh sample of the same policy: up to the row's ``attempts`` times, but not + once its ``reset_timeout`` has passed since the rollout began. A rollout that fails every attempt, in training, has + its whole group's loss mask and rewards zeroed, as SkyRL's Harbor generator does with a failed trial, so the group + is out of the gradient and of its own baseline; in an eval it scores 0. Each rollout's env metrics hold its + attempts and whether it ``failed``. + + It leans on SkyRL internals, which CI checks on SkyRL 0.4.0 and main: ``agent_loop`` builds a new env per call and + raises what the env raises, its output's ``env_metrics`` reach ``generate``'s, and ``generate``'s ``loss_masks``, + ``rewards`` and ``env_metrics`` are a list per trajectory, or per step, as its ``trajectory_ids`` or the input + batch's are.""" + + def __init__(self, *args: Any, **kwargs: Any): + super().__init__(*args, **kwargs) + if self.batched: + raise ValueError("agentenv rows play an episode over several turns, which SkyRL's batched generation " + "doesn't: set generator.batched=false") + + async def agent_loop(self, prompt: Any, env_class: str, env_extras: dict[str, Any], *args: Any, + **kwargs: Any) -> TrajectoryOutput | StepWiseOutput: + if env_class != AGENT: + return await super().agent_loop(prompt, env_class, env_extras, *args, **kwargs) + server, episode = env_extras["server"], env_extras["episode"] + began = time.monotonic() + for attempt in range(1, server["attempts"] + 1): + try: + output = await super().agent_loop(prompt, env_class, dict(env_extras), *args, **kwargs) + except EpisodeFailed as e: + failure = e + logger.warning("attempt %d of %d: %s", attempt, server["attempts"], e) + if time.monotonic() - began >= server["reset_timeout"]: + break + continue + last = output.step_outputs[-1] if isinstance(output, StepWiseOutput) else output + last.env_metrics.update(attempts=attempt, failed=0) + return output + logger.warning("task %s of split %s: the env failed it on %d attempt%s, so its group is left out of training: " + "%s", episode["task_id"], episode["split"], attempt, "s" * (attempt > 1), failure) + return self._left_out(attempt, time.monotonic() - began, kwargs.get("sampling_params")) + + async def generate(self, input_batch: GeneratorInput, *args: Any, **kwargs: Any) -> GeneratorOutput: + output = await super().generate(input_batch, *args, **kwargs) + failed = [row for row, metrics in enumerate(output["env_metrics"]) if (metrics or {}).get("failed")] + metadata = input_batch.get("batch_metadata") + ids = output.get("trajectory_ids") or input_batch.get("trajectory_ids") + if failed and ids and getattr(metadata, "training_phase", None) != TRAINING_PHASE_EVAL: + groups = {ids[row].instance_id for row in failed} + failed = [row for row, trajectory in enumerate(ids) if trajectory.instance_id in groups] + for row in failed: + output["loss_masks"][row] = [0] * len(output["loss_masks"][row]) + reward = output["rewards"][row] + output["rewards"][row] = [0.0] * len(reward) if isinstance(reward, list) else 0.0 + return output + + def _left_out(self, attempts: int, seconds: float, + sampling_params: dict[str, Any] | None) -> TrajectoryOutput | StepWiseOutput: + """A placeholder for a rollout every attempt of which the env failed: one token, masked, rewarded 0.""" + sampling = sampling_params if sampling_params is not None else asdict(self.generator_cfg.sampling_params) + token = self.tokenizer.eos_token_id + output = TrajectoryOutput(response_ids=[token], reward=[0.0], stop_reason="error", loss_mask=[0], + prompt_ids=[token], + rollout_logprobs=[0.0] if sampling.get("logprobs") is not None else None, + env_metrics={"attempts": attempts, "failed": 1}, e2e_time=seconds, + time_splits={"llm": 0.0, "env": 0.0}) + return StepWiseOutput(step_outputs=[output]) if self.generator_cfg.step_wise_trajectories else output diff --git a/src/agent_env/training/skyrl/prepare.py b/src/agent_env/training/skyrl/prepare.py new file mode 100644 index 000000000..7b9306d82 --- /dev/null +++ b/src/agent_env/training/skyrl/prepare.py @@ -0,0 +1,132 @@ +"""Writes a bundle's tasks as SkyRL's RL dataset: a parquet file of rows per split, which ``data.train_data`` and +``data.val_data`` take. SkyRL's generator shows the model no tools, so each row's system prompt carries its tools as +the model's chat template renders them with ``tools=``.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import pyarrow as pa +import pyarrow.parquet as pq +from transformers import AutoTokenizer + +from agent_env.training._openenv import function_tool, key, messages, tasks_and_tools, token_in +from agent_env.training.skyrl.calls import FORMATS + +#: The ``env_class`` every row plays under, and the id the env is registered at. +AGENT = "agentenv" +#: A row's columns, so a split the server keeps empty still has them. SkyRL hands the env every column but ``prompt`` +#: and ``env_class`` as its extras, and groups eval metrics by ``data_source``. +ROW = pa.schema([ + ("prompt", pa.list_(pa.struct([("role", pa.string()), ("content", pa.string())]))), + ("env_class", pa.string()), + ("data_source", pa.string()), + ("episode", pa.struct([("split", pa.string()), ("index", pa.int64()), ("task_id", pa.string()), + ("tools", pa.string())])), + ("server", pa.struct([("url", pa.string()), ("tool_timeout", pa.float64()), ("reset_timeout", pa.float64()), + ("attempts", pa.int64()), ("tool_format", pa.string()), ("token_var", pa.string())])), +]) +#: What stands in for a system prompt while finding where a chat template puts it. +_MARK = "AGENTENVSYSTEMPROMPT" + + +@dataclass(frozen=True) +class Written: + """What ``write_dataset`` wrote: the folder, and the longest prompt in tokens, its tools included.""" + + folder: Path + longest_prompt: int + + +def write_dataset(url: str, out: str | Path, *, tokenizer: str, tool_format: str = "hermes", + tool_timeout: float = 60.0, reset_timeout: float = 600.0, attempts: int = 3, + token_var: 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. A row's + prompt is its task's system prompt and first prompt, with every tool its task's turns offer, found by playing one + task of each toolset to its last turn, unscored, so writing deploys those. The tools go in the system prompt as + ``tokenizer``'s chat template renders them with ``tools=`` and ``chat_template_kwargs``, SkyRL's + ``generator.chat_template_kwargs``, checked to render the same text; ``tool_format`` names the format the template + teaches the model to call them in. Each row also holds the server's settings for its env: the deadlines, the + ``attempts`` a rollout the env fails is played, and ``token_var``, the environment variable that holds the token a + server given one asks for (``build_app(token=...)``), read now and in the trainer, so the token is in no file. + + A row's prompt is counted as SkyRL's prompt filter counts it, without ``chat_template_kwargs``, and as its generator + renders it, with them; a row over ``max_prompt_length``, SkyRL's ``trainer.max_prompt_length``, in either is + refused before anything is written, where the filter would drop it.""" + if attempts < 1: + raise ValueError(f"attempts must be at least 1, not {attempts}") + if tool_format not in FORMATS: + raise ValueError(f"tool_format must be one of {', '.join(FORMATS)}, not {tool_format!r}") + 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)) + template = AutoTokenizer.from_pretrained(tokenizer) + kwargs = chat_template_kwargs or {} + tools = {each: [function_tool(tool) for tool in found] for each, found in probed.items()} + server = {"url": url, "tool_timeout": tool_timeout, "reset_timeout": reset_timeout, "attempts": attempts, + "tool_format": tool_format, "token_var": token_var} + rows = {split: [{ + "prompt": with_tools(template, messages(task["system_prompt"], task["prompt"]), tools[key(task)], kwargs), + "env_class": AGENT, + "data_source": f"{AGENT}/{env_name}", + "episode": {"split": split, "index": task["index"], "task_id": task["id"], + "tools": json.dumps(tools[key(task)])}, + "server": server, + } for task in described] for split, described in tasks.items()} + longest = _longest(rows, template, max_prompt_length, kwargs) + 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") + return Written(out, longest) + + +def token_vars(files: list[str]) -> set[str]: + """The environment variables the rows in the local parquet ``files`` name for their servers' tokens; a file that + isn't on this machine, or holds other rows, names none.""" + ours = [file for file in files if Path(file).is_file() and "server" in pq.read_schema(file).names] + return {server["token_var"] for file in ours + for server in pq.read_table(file, columns=["server"]).column("server").to_pylist() if server["token_var"]} + + +def with_tools(tokenizer: Any, prompt: list[dict[str, str]], tools: list[dict[str, Any]], + chat_template_kwargs: dict[str, Any]) -> list[dict[str, str]]: + """``prompt`` with ``tools`` in its system prompt, as the chat template renders them with ``tools=``, which + SkyRL's generator doesn't pass: the system prompt that renders the same text without them. Raises for a template + that puts tools anywhere but the system turn.""" + target = _rendered(tokenizer, prompt, chat_template_kwargs, tools=tools) + rest = prompt[1:] if prompt[0]["role"] == "system" else prompt + marked = _rendered(tokenizer, [{"role": "system", "content": _MARK}, *rest], chat_template_kwargs) + if marked.count(_MARK) == 1: + before, after = marked.split(_MARK) + if target.startswith(before) and target.endswith(after) and len(target) >= len(before) + len(after): + found = [{"role": "system", "content": target[len(before):len(target) - len(after)]}, *rest] + if _rendered(tokenizer, found, chat_template_kwargs) == target: + return found + raise ValueError(f"{tokenizer.name_or_path}'s chat template doesn't render tools in the system prompt, where this " + "writer puts them, since SkyRL's generator passes the template no tools") + + +def _rendered(tokenizer: Any, prompt: list[dict[str, str]], chat_template_kwargs: dict[str, Any], **tools: Any) -> str: + return tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tokenize=False, **tools, + **chat_template_kwargs) + + +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, counted with and without ``chat_template_kwargs``, refusing rows over + ``limit``.""" + lengths = {(split, row["episode"]["task_id"]): max( + len(tokenizer.apply_chat_template(row["prompt"], add_generation_prompt=True, tokenize=True, return_dict=True, + **kwargs)["input_ids"]) + for kwargs in ({}, chat_template_kwargs)) 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 SkyRL's prompt filter would drop: " + "raise it, and trainer.max_prompt_length with it") + return max(lengths.values(), default=0) diff --git a/src/agent_env/training/skyrl/train.py b/src/agent_env/training/skyrl/train.py new file mode 100644 index 000000000..7726d5e80 --- /dev/null +++ b/src/agent_env/training/skyrl/train.py @@ -0,0 +1,77 @@ +"""``python -m agent_env.training.skyrl.train``, or SkyRL's ``main_base`` with +``trainer.override_entrypoint=agent_env.training.skyrl.train``: SkyRL's trainer on the rows +``python -m agent_env.training.skyrl`` writes, with the env registered as ``agentenv`` and AgentEnvGenerator +generating, and SkyRL's fully async trainer with ``trainer.fully_async.enabled=true``. It takes SkyRL's own +settings.""" + +from __future__ import annotations + +import os +import sys +from typing import Any + +try: + import ray + from skyrl.backends.skyrl_train.inference_servers.utils import resolve_policy_model_name + from skyrl.train.config import SkyRLTrainConfig + from skyrl.train.entrypoints.main_base import BasePPOExp + from skyrl.train.fully_async_trainer import FullyAsyncRayPPOTrainer + from skyrl.train.utils import initialize_ray, validate_cfg + from skyrl_gym.envs.registration import register as register_env + from skyrl_gym.envs.registration import registry +except ModuleNotFoundError as exc: + if exc.name is None or exc.name.split(".")[0] not in ("skyrl", "skyrl_gym", "ray"): + raise + raise ImportError("agent_env.training.skyrl.train runs in SkyRL's environment: install " + "'agentenv-framework[skyrl]' into it") from exc + +from agent_env.training._openenv import token_in +from agent_env.training.skyrl.generator import AgentEnvGenerator +from agent_env.training.skyrl.prepare import AGENT, token_vars + + +def register() -> None: + """Registers ``AgentEnvTextEnv`` with skyrl-gym as ``agentenv``, the ``env_class`` of every row, in this process: + a trainer of one's own calls it where SkyRL builds its envs, and generates with ``AgentEnvGenerator``.""" + if AGENT not in registry: + register_env(id=AGENT, entry_point="agent_env.training.skyrl.env:AgentEnvTextEnv") + + +class AgentEnvExp(BasePPOExp): + """SkyRL's experiment with ``AgentEnvGenerator``, and its fully async trainer when ``trainer.fully_async.enabled`` + is set.""" + + def get_generator(self, cfg: SkyRLTrainConfig, tokenizer: Any, inference_engine_client: Any) -> AgentEnvGenerator: + if cfg.generator.vision_language_generator: + raise ValueError("agentenv rows are text, so set generator.vision_language_generator=false") + return AgentEnvGenerator(generator_cfg=cfg.generator, skyrl_gym_cfg=cfg.environment.skyrl_gym, + inference_engine_client=inference_engine_client, tokenizer=tokenizer, + policy_model_name=resolve_policy_model_name(cfg)) + + def get_trainer(self, cfg: SkyRLTrainConfig, *args: Any, **kwargs: Any) -> Any: + if not cfg.trainer.fully_async.enabled: + return super().get_trainer(cfg, *args, **kwargs) + tracker, tokenizer, train_dataset, eval_dataset, inference_engine_client, generator, colocate_pg = args + return FullyAsyncRayPPOTrainer(cfg=cfg, tracker=tracker, tokenizer=tokenizer, train_dataset=train_dataset, + eval_dataset=eval_dataset, inference_engine_client=inference_engine_client, + generator=generator, colocate_pg=colocate_pg) + + +@ray.remote(num_cpus=1) +def _run(cfg: SkyRLTrainConfig, environ: dict[str, str]) -> None: + os.environ.update(environ) + register() + AgentEnvExp(cfg).run() + + +def main() -> None: + """Reads the tokens the rows' servers ask for here, so a trainer on any node of the Ray cluster has them.""" + cfg = SkyRLTrainConfig.from_cli_overrides(sys.argv[1:]) + validate_cfg(cfg) + environ = {var: token_in(var) for var in token_vars([*cfg.data.train_data, *cfg.data.val_data])} + initialize_ray(cfg) + ray.get(_run.remote(cfg, environ)) + + +if __name__ == "__main__": + main() diff --git a/src/agent_env/training/trl.py b/src/agent_env/training/trl.py index fd96f4f62..f595e5b87 100644 --- a/src/agent_env/training/trl.py +++ b/src/agent_env/training/trl.py @@ -13,7 +13,7 @@ import logging import threading import weakref -from collections.abc import Callable, Coroutine, Mapping +from collections.abc import Callable, Mapping from functools import cached_property from pathlib import Path from typing import Any, ClassVar @@ -40,6 +40,7 @@ Shared, key, messages, + on_session_loop, probe, read_tasks, toolsets, @@ -52,8 +53,6 @@ #: How long the exit hook waits for the sessions to close. EXIT_SECONDS = 30.0 -_loop: asyncio.AbstractEventLoop | None = None -_loop_lock = threading.Lock() _environments: weakref.WeakSet[_Environment] = weakref.WeakSet() _servers: set[Server] = set() @@ -138,7 +137,7 @@ def factory(self) -> type[_Environment] | dict[str, type[_Environment]]: the last empty and listing its tools, then ends the episode unscored, since TRL reads the tools from the class: its methods are every tool a turn offers, and one a turn doesn't offer answers so. Each call on a probe's session has the reset deadline.""" - classes = _on_loop(self._environments(), None) + classes = on_session_loop(self._environments(), None) return classes[0] if len(classes) == 1 else {each.__name__: each for each in classes} def close(self) -> None: @@ -146,7 +145,7 @@ def close(self) -> None: started.""" mine = [environment for environment in list(_environments) if environment._shared in self._shared] try: - _on_loop(_gathered(environment._close() for environment in mine), CLOSE_SECONDS) + on_session_loop(_gathered(environment._close() for environment in mine), CLOSE_SECONDS) finally: if self._server is not None: _stop(self._server) @@ -189,7 +188,7 @@ def reset(self, *, split: str, index: int, task_id: str | None = None, **_: Any) one at a time: the first tool call or ``get_reward`` waits for it. Raises when the server refuses this environment's first session, as one at capacity does; a later session it refuses, after one was dropped, leaves only that rollout unscored.""" - _on_loop(self._start(split, index, task_id), CONNECT_SECONDS + CLOSE_SECONDS) + on_session_loop(self._start(split, index, task_id), CONNECT_SECONDS + CLOSE_SECONDS) def get_reward(self) -> float | None: """The episode's reward, or ``None`` when it ended in an error. Raises when the episode's first turn offered a @@ -198,7 +197,7 @@ def get_reward(self) -> float | None: if rollout is None: return None try: - _on_loop(self._finish(rollout), None) + on_session_loop(self._finish(rollout), None) except Exception as e: rollout.failure = rollout.failure or why(e) if rollout.misgrouped: @@ -213,7 +212,7 @@ def _call(self, name: str, arguments: dict[str, Any]) -> str: """Waits at most for the deploy and the call, each to its deadline and a grace after, then the session's close: the first call of a rollout waits for its deploy.""" shared = self._shared - return _on_loop(self._step(name, arguments), shared.reset_timeout + max(shared.reset_timeout, shared.tool_timeout) + return on_session_loop(self._step(name, arguments), shared.reset_timeout + max(shared.reset_timeout, shared.tool_timeout) + 2 * GRACE_SECONDS + CLOSE_SECONDS) @@ -250,24 +249,6 @@ class does, so it keeps what it overrides.""" copyreg.pickle(_EnvironmentClass, _reduce_environment_class) -def _episode_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-trl-sessions", daemon=True).start() - return _loop - - -def _on_loop(coroutine: Coroutine[Any, Any, Any], timeout: float | None) -> Any: - future = asyncio.run_coroutine_threadsafe(coroutine, _episode_loop()) - try: - return future.result(timeout) - except TimeoutError: - future.cancel() - raise - - async def _gathered(coroutines: Any) -> list[Any]: return list(await asyncio.gather(*coroutines)) @@ -278,7 +259,7 @@ def _close_all() -> None: ``EXIT_SECONDS`` for the sessions and ``STOP_SECONDS`` for each server. It runs before threading's shutdown, while the thread pools an episode's teardown uses still take work.""" try: - _on_loop(_gathered(environment._close() for environment in list(_environments)), EXIT_SECONDS) + on_session_loop(_gathered(environment._close() for environment in list(_environments)), EXIT_SECONDS) except Exception as e: logger.warning("couldn't end the episodes left running: %s", why(e)) for server in list(_servers): diff --git a/src/agent_env/training/verl/loop.py b/src/agent_env/training/verl/loop.py index f13716066..d097d95fb 100644 --- a/src/agent_env/training/verl/loop.py +++ b/src/agent_env/training/verl/loop.py @@ -23,8 +23,7 @@ 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 +from agent_env.training._openenv import Rollout, Session, Shared, token_in logger = logging.getLogger(__name__) #: The tools whose answers verl was seen to cut, each warned of once a process. diff --git a/src/agent_env/training/verl/prepare.py b/src/agent_env/training/verl/prepare.py index 17d1b3da1..b1889860b 100644 --- a/src/agent_env/training/verl/prepare.py +++ b/src/agent_env/training/verl/prepare.py @@ -5,7 +5,6 @@ from __future__ import annotations import json -import os from dataclasses import dataclass from pathlib import Path from typing import Any @@ -15,7 +14,7 @@ import yaml from transformers import AutoTokenizer -from agent_env.training._openenv import key, messages, tasks_and_tools +from agent_env.training._openenv import function_tool, key, messages, tasks_and_tools, token_in #: The ``agent_name`` every row plays under, and the name ``agent_loop.yaml`` registers the loop at. AGENT = "agentenv" @@ -61,7 +60,7 @@ def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, rese 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()} + tools = {each: [function_tool(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"]), @@ -82,15 +81,6 @@ def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, rese 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``.""" @@ -104,10 +94,3 @@ def _longest(rows: dict[str, list[dict[str, Any]]], tokenizer: Any, limit: int | "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/skyrl/__init__.py b/tst/skyrl/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tst/skyrl/conftest.py b/tst/skyrl/conftest.py new file mode 100644 index 000000000..d37f30af8 --- /dev/null +++ b/tst/skyrl/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 SkyRL's +generator.""" + +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/skyrl/generator_test.py b/tst/skyrl/generator_test.py new file mode 100644 index 000000000..b1b1af0b3 --- /dev/null +++ b/tst/skyrl/generator_test.py @@ -0,0 +1,296 @@ +"""``AgentEnvGenerator`` and ``AgentEnvTextEnv`` on SkyRL itself: SkyRL's ``PromptDataset`` loads the rows +``write_dataset`` wrote, with Qwen3's and Qwen3.5's tokenizers and chat templates, and its generator plays them through +a stand-in inference engine, on the unit tier's bundle served on loopback. A played row's reward is its episode's, and +its scores, attempts and tool calls its env metrics, through the hermes and qwen3_coder formats, over two turns, and +from a server that asks for a token; the model sees its tools as its chat template renders them; an unparsed call is +answered with why; the turn cap ends the episode; a rollout the env fails is played again on a new episode, but not +once its time is up, and one it fails every attempt masks its group in training but not in an eval; an episode that +ends unscored without failing raises; and the entrypoint registers the env and reads the rows' token variables. CI runs +it on SkyRL 0.4.0 and main.""" + +import asyncio +import json + +import pyarrow.parquet as pq +import pytest +from transformers import AutoTokenizer + +from agent_env.training.skyrl.prepare import token_vars, write_dataset +from tst.unit.task.episode import run_test +from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN +from tst.unit.training.nemo_gym_test import PLAYED +from tst.util.capabilities import SKYRL, missing_capability_reason + +skyrl_env = pytest.importorskip("agent_env.training.skyrl.env", reason=missing_capability_reason(SKYRL), + exc_type=ImportError) +skyrl_generator = pytest.importorskip("agent_env.training.skyrl.generator", reason=missing_capability_reason(SKYRL), + exc_type=ImportError) +train = pytest.importorskip("agent_env.training.skyrl.train", reason=missing_capability_reason(SKYRL), + exc_type=ImportError) +config = pytest.importorskip("skyrl.train.config", reason=missing_capability_reason(SKYRL)) +dataset = pytest.importorskip("skyrl.train.dataset.dataset", reason=missing_capability_reason(SKYRL)) +generator_utils = pytest.importorskip("skyrl.train.generators.utils", reason=missing_capability_reason(SKYRL)) + +TOKENIZERS = {"hermes": "Qwen/Qwen3-0.6B", "qwen3_coder": "Qwen/Qwen3.5-0.8B"} +UNTHINKING = {"enable_thinking": False} +TWO_TURNS = [ITEMS, ALICE, TURN, {**TURN, "id": "turn-2", "prompt_id": "turn-2", "prompt": "Add a barrel."}, REPLIES] + + +def call(tool_format, name, **arguments): + """A tool call as ``tool_format`` writes it.""" + if tool_format == "hermes": + return f'\n{json.dumps({"name": name, "arguments": arguments})}\n' + parameters = "".join(f"\n{value}\n\n" for key, value in arguments.items()) + return f"\n\n{parameters}\n" + + +def adds_a_crate(tool_format): + """The model of a one-turn task: it adds a crate, then replies once the tool answered.""" + return lambda answered: call(tool_format, "add_item", item="crate") if answered == 0 else "Added a crate." + + +@pytest.fixture(scope="module", autouse=True) +def registered(): + train.register() + + +@pytest.fixture(scope="module", params=sorted(TOKENIZERS)) +def tool_format(request): + return request.param + + +@pytest.fixture(scope="module") +def tokenizer(tool_format): + return AutoTokenizer.from_pretrained(TOKENIZERS[tool_format]) + + +def test_a_played_rows_reward_is_its_episodes_and_its_last_reply_ends_the_turn(serve, tmp_path, tool_format, + tokenizer, monkeypatch, teardowns): + server = serve({"a": PLAYED}) + rows = _rows(server.url, tmp_path, tool_format) + engine, replies = _Engine(tokenizer, adds_a_crate(tool_format)), [] + finish = skyrl_env._Session.finish + + async def finishing(self, reply): + replies.append(reply) + return await finish(self, reply) + + monkeypatch.setattr(skyrl_env._Session, "finish", finishing) + output = _generate(rows, engine, tokenizer, samples=2) + server.stop() + + assert [sum(reward) for reward in output["rewards"]] == [1.0, 1.0] + assert all(1 in mask and 0 in mask for mask in output["loss_masks"]) + assert output["env_metrics"] == [{"tool_calls": 1, "unparsed_calls": 0, "score/replies": 1.0, "attempts": 1, + "failed": 0}] * 2 + assert output["rollout_metrics"]["environment/score/replies"] == 1.0 + assert (replies, engine.calls) == (["Added a crate."] * 2, 4) + assert len(teardowns.finished) == 3 + + +def test_the_model_sees_its_tools_as_its_chat_template_renders_them(serve, tmp_path, tool_format, tokenizer): + server = serve({"a": PLAYED}) + rows = _rows(server.url, tmp_path, tool_format) + engine = _Engine(tokenizer, adds_a_crate(tool_format)) + + _generate(rows, engine, tokenizer) + server.stop() + + [row] = pq.read_table(tmp_path / "all.parquet").to_pylist() + tools = json.loads(row["episode"]["tools"]) + native = tokenizer.apply_chat_template([{"role": "system", "content": "You keep the items."}, + {"role": "user", "content": "Add a crate."}], tools=tools, + add_generation_prompt=True, tokenize=True, return_dict=True, **UNTHINKING) + assert engine.first_prompt == native["input_ids"] + + +def test_a_task_of_two_turns_is_played_through_end_turn_and_scored_on_both_replies(serve, tmp_path, tool_format, + tokenizer): + server = serve({"a": TWO_TURNS}) + rows = _rows(server.url, tmp_path, tool_format) + + def two_turns(answered): + return [call(tool_format, "add_item", item="crate"), call(tool_format, "end_turn", response="Added a crate."), + call(tool_format, "add_item", item="barrel"), "Added a barrel."][answered] + + engine = _Engine(tokenizer, two_turns) + output = _generate(rows, engine, tokenizer, max_turns=6) + server.stop() + + assert [sum(reward) for reward in output["rewards"]] == [1.0] + assert output["env_metrics"][0]["tool_calls"] == 3 + assert "Add a barrel." in tokenizer.decode(output["response_ids"][0]) + + +def test_an_unparsed_call_is_answered_with_why_and_counted(serve, tmp_path, tool_format, tokenizer): + server = serve({"a": PLAYED}) + rows = _rows(server.url, tmp_path, tool_format) + + def misspelled(answered): + return [call(tool_format, "add_items", item="crate"), call(tool_format, "add_item", item="crate"), + "Added a crate."][answered] + + output = _generate(rows, _Engine(tokenizer, misspelled), tokenizer) + server.stop() + + assert [sum(reward) for reward in output["rewards"]] == [1.0] + assert (output["env_metrics"][0]["tool_calls"], output["env_metrics"][0]["unparsed_calls"]) == (1, 1) + assert "Error: there's no tool 'add_items'; the tools are add_item, end_turn" in tokenizer.decode( + output["response_ids"][0]) + + +def test_the_turn_cap_ends_the_episode_with_the_last_reply(serve, tmp_path, tool_format, tokenizer): + server = serve({"a": PLAYED}) + rows = _rows(server.url, tmp_path, tool_format) + engine = _Engine(tokenizer, adds_a_crate(tool_format)) + + output = _generate(rows, engine, tokenizer, max_turns=1) + server.stop() + + assert (engine.calls, [sum(reward) for reward in output["rewards"]]) == (1, [0.0]) + assert output["env_metrics"][0]["tool_calls"] == 1 + + +def test_a_rollout_the_env_fails_is_played_again_on_a_new_episode(serve, tmp_path, tool_format, tokenizer, + monkeypatch): + server = serve({"a": PLAYED}) + rows = _rows(server.url, tmp_path, tool_format) + 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 = _generate(rows, _Engine(tokenizer, adds_a_crate(tool_format)), tokenizer) + server.stop() + + assert [sum(reward) for reward in output["rewards"]] == [1.0] + assert (output["env_metrics"][0]["attempts"], output["env_metrics"][0]["failed"], len(deploys)) == (2, 0, 2) + + +@pytest.mark.parametrize("step_wise", [False, True]) +@pytest.mark.parametrize(("phase", "kept"), [("train", 0.0), ("eval", 1.0)]) +def test_a_rollout_the_env_fails_every_attempt_masks_its_group_in_training_but_not_in_an_eval( + serve, tmp_path, tool_format, tokenizer, monkeypatch, step_wise, phase, kept): + server = serve({"a": PLAYED}) + rows = _rows(server.url, tmp_path, tool_format, attempts=2) + deploy, deploys = run_test._Env.execute, [] + + async def failing_after_the_first(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_after_the_first) + output = _generate(rows, _Engine(tokenizer, adds_a_crate(tool_format)), tokenizer, samples=2, phase=phase, + step_wise=step_wise) + server.stop() + + last = [metrics for metrics in output["env_metrics"] if "failed" in metrics] + assert sorted((metrics["failed"], metrics["attempts"]) for metrics in last) == [(0, 1), (1, 2)] + played = next(row for row, metrics in enumerate(output["env_metrics"]) if metrics.get("failed") == 0) + assert sum(output["rewards"][played]) == kept + assert (1 in output["loss_masks"][played]) == (phase == "eval") + failed = next(row for row, metrics in enumerate(output["env_metrics"]) if metrics.get("failed") == 1) + assert (output["loss_masks"][failed], sum(output["rewards"][failed])) == ([0], 0.0) + assert len(deploys) == 3 + + +def test_a_rollout_isnt_played_again_once_its_time_is_up(serve, tmp_path, tool_format, tokenizer, monkeypatch): + server = serve({"a": PLAYED}) + rows = _rows(server.url, tmp_path, tool_format) + 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) + output = _generate(rows, _Engine(tokenizer, adds_a_crate(tool_format)), tokenizer, reset_timeout=0.5) + server.stop() + + assert (output["env_metrics"][0]["failed"], output["env_metrics"][0]["attempts"]) == (1, 1) + + +def test_an_episode_that_ends_without_failing_or_a_reward_raises(serve, tmp_path, tool_format, tokenizer): + server = serve({"a": PLAYED}, reward=lambda scores: None) + rows = _rows(server.url, tmp_path, tool_format) + + with pytest.raises(ValueError, match="task a of split all ended without a reward: it records no score"): + _generate(rows, _Engine(tokenizer, adds_a_crate(tool_format)), tokenizer) + server.stop() + + +def test_a_server_that_asks_for_a_token_is_played_with_the_one_its_variable_holds(serve, tmp_path, tool_format, + tokenizer, monkeypatch): + server = serve({"a": PLAYED}, token="s3cret") + monkeypatch.setenv("DOCK_TOKEN", "s3cret") + rows = _rows(server.url, tmp_path, tool_format, token_var="DOCK_TOKEN", attempts=1) + + played = _generate(rows, _Engine(tokenizer, adds_a_crate(tool_format)), tokenizer) + monkeypatch.setenv("DOCK_TOKEN", "guess") + refused = _generate(rows, _Engine(tokenizer, adds_a_crate(tool_format)), tokenizer) + server.stop() + + assert [sum(reward) for reward in played["rewards"]] == [1.0] + assert refused["env_metrics"][0]["failed"] == 1 + assert token_vars([str(tmp_path / "all.parquet"), str(tmp_path / "absent.parquet")]) == {"DOCK_TOKEN"} + + +def test_the_entrypoint_registers_the_env_once_and_batched_generation_is_refused(tokenizer): + train.register() + + with pytest.raises(ValueError, match="agentenv rows play an episode over several turns, which SkyRL's batched " + "generation doesn't: set generator.batched=false"): + _generator(_Engine(tokenizer, None), tokenizer, batched=True, chat_template_kwargs={}) + + +class _Engine: + """SkyRL's inference engine for the generator: ``script`` writes the reply to a trajectory's conversation from the + tool answers it holds, and the reply is its tokens and the end of the turn.""" + + def __init__(self, tokenizer, script): + self.tokenizer, self.script, self.calls, self.first_prompt = tokenizer, script, 0, None + + async def generate(self, engine_input, model=None): + self.calls += 1 + [prompt] = engine_input["prompt_token_ids"] + self.first_prompt = self.first_prompt or list(prompt) + reply = self.script(self.tokenizer.decode(prompt).count("")) + ids = [*self.tokenizer.encode(reply, add_special_tokens=False), self.tokenizer.eos_token_id] + return {"responses": [reply], "response_ids": [ids], "stop_reasons": ["stop"], "response_logprobs": None} + + async def finish_session(self, session_id): + pass + + +def _rows(url, folder, tool_format, **options): + return write_dataset(url, folder, tokenizer=TOKENIZERS[tool_format], tool_format=tool_format, + chat_template_kwargs=UNTHINKING, **options).folder / "all.parquet" + + +def _generator(engine, tokenizer, **settings): + generator_config = config.GeneratorConfig(**{"max_turns": 4, "max_input_length": 8192, + "chat_template_kwargs": UNTHINKING, **settings}) + return skyrl_generator.AgentEnvGenerator(generator_cfg=generator_config, + skyrl_gym_cfg=config.SkyRLGymConfig(max_env_workers=4), + inference_engine_client=engine, tokenizer=tokenizer, + policy_model_name="policy") + + +def _generate(rows, engine, tokenizer, *, samples=1, phase="train", step_wise=False, max_turns=4, + reset_timeout=None): + prompts = dataset.PromptDataset([str(rows)], tokenizer, max_prompt_length=8192, num_workers=1) + batch = prompts.collate_fn([prompts[index] for index in range(len(prompts.dataframe))]) + if reset_timeout is not None: + batch = [{**row, "env_extras": {**row["env_extras"], "server": {**row["env_extras"]["server"], + "reset_timeout": reset_timeout}}} + for row in batch] + generator_input, _ = generator_utils.prepare_generator_input(batch, samples, None, "agentenv", phase, 1) + generator = _generator(engine, tokenizer, max_turns=max_turns, step_wise_trajectories=step_wise) + return asyncio.run(generator.generate(generator_input)) + diff --git a/tst/unit/test_optional_extras_stay_optional.py b/tst/unit/test_optional_extras_stay_optional.py index 00cfc15a9..790048232 100644 --- a/tst/unit/test_optional_extras_stay_optional.py +++ b/tst/unit/test_optional_extras_stay_optional.py @@ -26,12 +26,13 @@ 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, 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: the Sail and Smol providers, the OpenEnv server, NeMo Gym's resources server and the verl and SkyRL helpers +# need their optional extras (and the verl agent loop verl itself, the SkyRL env and generator SkyRL), 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.training.verl", "agent_env.env.gateway", + "agent_env.openenv", "agent_env.training.nemo_gym", "agent_env.training.verl", "agent_env.training.skyrl", + "agent_env.env.gateway", "agent_env.task_step.task_steps.run_code_runner", ) @@ -171,8 +172,8 @@ def _run(script: str, blocked: list[str], *args: str) -> str: #: Extras every distribution of which another extra installs too, so blocking what only they install blocks nothing: -#: verl's pyarrow comes with trl's datasets. The test without any extra covers them. -_COVERED = {"verl": "trl"} +#: verl's and SkyRL's pyarrow comes with trl's datasets. The test without any extra covers them. +_COVERED = {"verl": "trl", "skyrl": "trl"} @pytest.mark.parametrize("extra", sorted(set(_optional_requirements()) - set(_COVERED))) diff --git a/tst/unit/training/skyrl_test.py b/tst/unit/training/skyrl_test.py new file mode 100644 index 000000000..88e6ea21d --- /dev/null +++ b/tst/unit/training/skyrl_test.py @@ -0,0 +1,282 @@ +"""``write_dataset`` on the OpenEnv server ``build_app`` makes, served on loopback over the Episode API's fake steps: +SkyRL's parquet rows for each split, each with its tools in its system prompt as a local tokenizer's chat template +renders them, from a server that asks for a token too, and the command that writes them; and the tool-call parsers the +env reads the model's replies with.""" + +import json + +import httpx +import pyarrow.parquet as pq +import pytest +from click.testing import CliRunner +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import AutoTokenizer, PreTrainedTokenizerFast + +from agent_env.openenv.environment import END_TURN +from agent_env.training.skyrl.__main__ import main +from agent_env.training.skyrl.calls import Call, Unparsed, parse +from agent_env.training.skyrl.prepare import 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}} +#: A chat template that renders tools in the system turn, after the system prompt, as Qwen3's does. +SYSTEM_TOOLS = ( + "{% if tools or messages[0]['role'] == 'system' %}{{ '<|im_start|>system\\n' }}" + "{% if messages[0]['role'] == 'system' %}{{ messages[0]['content'] }}{% if tools %}{{ '\\n\\n' }}{% endif %}" + "{% endif %}{% if tools %}{{ '# Tools\\n' }}{% for tool in tools %}{{ '\\n' + (tool | tojson) }}{% endfor %}" + "{{ '\\n' }}{% endif %}{{ '<|im_end|>\\n' }}{% endif %}" + "{% for message in messages %}{% if not (loop.first and message['role'] == 'system') %}" + "{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>\\n' }}{% endif %}{% endfor %}" + "{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}" + "{% if enable_thinking is defined and not enable_thinking %}{{ '\\n\\n\\n\\n' }}{% endif %}" + "{% endif %}" +) +#: One that renders them in the user's turn instead. +USER_TOOLS = ( + "{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] }}" + "{% if message['role'] == 'user' and tools %}{{ '\\n' + (tools | tojson) }}{% endif %}{{ '<|im_end|>\\n' }}" + "{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}" +) +SCHEDULE = {"name": "schedule", "parameters": {"type": "object", "properties": { + "hours": {"type": "array"}, "count": {"type": "integer"}, "price": {"type": "number"}, "sell": {"type": "boolean"}, + "note": {"type": "string"}, "plan": {"type": "object"}, "either": {"type": ["integer", "string"]}, "loose": {}}}} +TOOLS = {"schedule": SCHEDULE, "end_turn": {"name": "end_turn", "parameters": {"type": "object"}}} + + +@pytest.fixture +def tokenizer(tmp_path): + """A local tokenizer of whitespace-separated words with ``template`` as its chat template, saved where + ``AutoTokenizer`` loads it.""" + + def saved(template=SYSTEM_TOOLS): + words = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) + words.pre_tokenizer = pre_tokenizers.Whitespace() + folder = tmp_path / "tokenizer" + PreTrainedTokenizerFast(tokenizer_object=words, unk_token="[UNK]", eos_token="<|im_end|>", + additional_special_tokens=["<|im_start|>"], chat_template=template + ).save_pretrained(folder) + return str(folder) + + return saved + + +def test_a_servers_tasks_become_a_row_per_task_with_its_tools_in_its_system_prompt(serve, tmp_path, tokenizer, + 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", tokenizer=tokenizer(), tool_format="qwen3_coder", + tool_timeout=30.0, reset_timeout=90.0, attempts=2) + server.stop() + + tools = "\n".join(json.dumps(tool, ensure_ascii=False) for tool in (ADD_ITEM, FUNCTION_END_TURN)) + [row] = _rows(written.folder / "smoke.parquet") + assert row == { + "prompt": [{"role": "system", "content": f"You keep the items.\n\n# Tools\n\n{tools}\n"}, + {"role": "user", "content": "Add a barrel."}], + "env_class": "agentenv", + "data_source": "agentenv/Dock Yard", + "episode": {"split": "smoke", "index": 0, "task_id": "b", "tools": [ADD_ITEM, FUNCTION_END_TURN]}, + "server": {"url": server.url, "tool_timeout": 30.0, "reset_timeout": 90.0, "attempts": 2, + "tool_format": "qwen3_coder", "token_var": None}, + } + template = AutoTokenizer.from_pretrained(tokenizer()) + assert template.apply_chat_template(row["prompt"], add_generation_prompt=True, tokenize=False) == \ + template.apply_chat_template(row["prompt"][1:], tools=[ADD_ITEM, FUNCTION_END_TURN], + add_generation_prompt=True, tokenize=False).replace( + "<|im_start|>system\n", "<|im_start|>system\nYou keep the items.\n\n") + assert [(row["episode"]["index"], row["episode"]["task_id"]) for row in _rows(written.folder / "all.parquet")] == [ + (0, "a"), (1, "b")] + assert written.longest_prompt == max( + len(template.apply_chat_template(row["prompt"], add_generation_prompt=True, return_dict=True)["input_ids"]) + for row in _rows(written.folder / "all.parquet")) + assert len(teardowns.finished) == 1 + + +def test_a_task_without_a_system_prompt_gets_one_of_its_tools(serve, tmp_path, tokenizer): + server = serve({"a": [ITEMS, {**ALICE, "system_prompt": None}, TURN, REPLIES]}) + + out = write_dataset(server.url, tmp_path, tokenizer=tokenizer()).folder + server.stop() + + [row] = _rows(out / "all.parquet") + tools = "\n".join(json.dumps(tool, ensure_ascii=False) for tool in (ADD_ITEM, FUNCTION_END_TURN)) + assert row["prompt"] == [{"role": "system", "content": f"# Tools\n\n{tools}\n"}, + {"role": "user", "content": "Add a crate."}] + + +def test_each_toolset_is_probed_once_and_a_rows_tools_are_its_tasks(serve, tmp_path, tokenizer, 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, tokenizer=tokenizer()).folder + server.stop() + + rows = _rows(out / "all.parquet") + assert [[tool["function"]["name"] for tool in row["episode"]["tools"]] for row in rows] == [ + ["add_item", "end_turn"], ["take_item", "end_turn"], ["add_item", "end_turn"]] + assert ['"name": "take_item"' in row["prompt"][0]["content"] for row in rows] == [False, True, False] + assert len(teardowns.finished) == 2 + + +def test_a_split_the_server_keeps_empty_has_a_rows_columns(serve, tmp_path, tokenizer): + server = serve({"a": PLAYED}, {"empty": []}) + + out = write_dataset(server.url, tmp_path, tokenizer=tokenizer()).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_server_that_asks_for_a_token_is_read_with_the_one_its_variable_holds_and_the_rows_name_it( + serve, tmp_path, tokenizer, monkeypatch, teardowns): + server = serve({"a": PLAYED}, token="s3cret") + monkeypatch.setenv("DOCK_TOKEN", "s3cret") + + out = write_dataset(server.url, tmp_path / "data", tokenizer=tokenizer(), 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", tokenizer=tokenizer(), token_var="DOCK_TOKEN") + server.stop() + + [row] = _rows(out / "all.parquet") + assert row["server"]["token_var"] == "DOCK_TOKEN" + assert b"s3cret" not in (out / "all.parquet").read_bytes() + assert not (tmp_path / "refused").exists() + + +def test_rows_over_the_prompt_limit_are_refused_counted_with_and_without_the_template_settings(serve, tmp_path, + tokenizer): + server = serve({"a": PLAYED}) + plain = write_dataset(server.url, tmp_path / "plain", tokenizer=tokenizer()).longest_prompt + unthinking = write_dataset(server.url, tmp_path / "unthinking", tokenizer=tokenizer(), + chat_template_kwargs={"enable_thinking": False}).longest_prompt + + with pytest.raises(ValueError, match=rf"1 rows' prompts are over max_prompt_length={unthinking - 1} tokens with " + rf"their tools, the longest task a of split all at {unthinking}, which " + "SkyRL's prompt filter would drop"): + write_dataset(server.url, tmp_path / "data", tokenizer=tokenizer(), max_prompt_length=unthinking - 1, + chat_template_kwargs={"enable_thinking": False}) + server.stop() + + assert unthinking > plain + assert not (tmp_path / "data").exists() + + +def test_a_template_that_renders_tools_outside_the_system_prompt_is_refused(serve, tmp_path, tokenizer): + server = serve({"a": PLAYED}) + + with pytest.raises(ValueError, match="chat template doesn't render tools in the system prompt, where this writer " + "puts them, since SkyRL's generator passes the template no tools"): + write_dataset(server.url, tmp_path / "data", tokenizer=tokenizer(USER_TOOLS)) + server.stop() + + assert not (tmp_path / "data").exists() + + +@pytest.mark.parametrize(("options", "message"), [ + ({"attempts": 0}, "attempts must be at least 1, not 0"), + ({"tool_format": "pythonic"}, "tool_format must be one of hermes, qwen3_coder, not 'pythonic'"), + ({"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", tokenizer="unread", **options) + assert list(tmp_path.iterdir()) == [] + + +def test_the_command_writes_the_rows(serve, tmp_path, tokenizer, teardowns): + server = serve({"a": PLAYED}) + + result = CliRunner().invoke(main, [server.url, str(tmp_path / "data"), "--tokenizer", tokenizer(), + "--tool-format", "qwen3_coder", "--tool-timeout", "15", "--reset-timeout", + "120", "--attempts", "4", "--chat-template-kwargs", + '{"enable_thinking": false}']) + server.stop() + + assert result.exit_code == 0, result.output + [row] = _rows(tmp_path / "data" / "all.parquet") + assert {key: row["server"][key] for key in ("tool_format", "tool_timeout", "reset_timeout", "attempts")} == { + "tool_format": "qwen3_coder", "tool_timeout": 15.0, "reset_timeout": 120.0, "attempts": 4} + assert len(teardowns.finished) == 1 + + +@pytest.mark.parametrize(("reply", "calls", "text"), [ + ('Scheduling.\n\n{"name": "schedule", "arguments": {"hours": [1, -1], "count": 2}}\n', + [Call("schedule", {"hours": [1, -1], "count": 2})], "Scheduling."), + ('{"name": "schedule", "arguments": "{\\"count\\": 2}"}\n' + '{"name": "end_turn", "arguments": {"response": "Done."}}', + [Call("schedule", {"count": 2}), Call("end_turn", {"response": "Done."})], ""), + ('{"name": "schedule", "arguments": {}}', [Call("schedule", {})], ""), + ("Idle is safest.Idling all day.", [], "Idling all day."), + ('{"name": "sell"}', [Unparsed("there's no tool 'sell'; the tools are end_turn, schedule")], + ""), + ("schedule(count=2)", [Unparsed( + "the tool call isn't JSON (Expecting value: line 1 column 1 (char 0)): write it as " + '{"name": , "arguments": }')], ""), + ('{"arguments": {}}', + [Unparsed('the tool call names no tool: write it as {"name": , "arguments": }')], ""), + ('{"name": "schedule", "arguments": [2]}', + [Unparsed("the arguments of schedule aren't a JSON object")], ""), +]) +def test_hermes_calls_are_a_json_object_a_block(reply, calls, text): + assert parse(reply, "hermes", TOOLS) == (calls, text) + + +def test_qwen3_coder_calls_are_typed_by_the_tools_schema(): + reply = ("I'll sell at the peak.\n\n\n\n[0, 2, -1]\n\n" + "\n3\n\n\n41.5\n\n\ntrue\n" + "\n\n42\n\n\n{\"at\": 18}\n\n" + "\nseven\n\n\n[1]\n\n\n" + "\n\n\n\nSold.\n\n" + "\n") + + calls, text = parse(reply, "qwen3_coder", TOOLS) + + assert calls == [Call("schedule", {"hours": [0, 2, -1], "count": 3, "price": 41.5, "sell": True, "note": "42", + "plan": {"at": 18}, "either": "seven", "loose": [1]}), + Call("end_turn", {"response": "Sold."})] + assert text == "I'll sell at the peak." + + +@pytest.mark.parametrize(("reply", "calls"), [ + ("\n\n\nmany\n\n\n", + [Call("schedule", {"count": "many"})]), + ("\n\n\n4\n\nfalse\n", + [Call("schedule", {"count": 4, "sell": False})]), + ("\n\n\n", + [Unparsed("there's no tool 'refund'; the tools are end_turn, schedule")]), + ("\nschedule\n", [Unparsed( + "the tool call names no tool: write it as with a block per " + "argument")]), +]) +def test_qwen3_coder_calls_that_dont_fit_their_schema_are_passed_as_written_or_unparsed(reply, calls): + assert parse(reply, "qwen3_coder", TOOLS)[0] == calls + + +def test_an_unknown_format_is_refused(): + with pytest.raises(ValueError, match="tool_format must be one of hermes, qwen3_coder, not 'pythonic'"): + parse("", "pythonic", TOOLS) + + +def _rows(path): + return [{**row, "episode": {**row["episode"], "tools": json.loads(row["episode"]["tools"])}} + for row in pq.read_table(path).to_pylist()] diff --git a/tst/unit/training/trl_test.py b/tst/unit/training/trl_test.py index 0335e1ccc..21b31a176 100644 --- a/tst/unit/training/trl_test.py +++ b/tst/unit/training/trl_test.py @@ -18,7 +18,7 @@ from agent_env.openenv.environment import AgentEnvEnvironment from agent_env.task_step.task_step import TaskStep from agent_env.training import trl -from agent_env.training._openenv import ENDED, FAILED +from agent_env.training._openenv import ENDED, FAILED, on_session_loop from agent_env.training.trl import AgentEnv from tst.unit.bundle._support import layout from tst.unit.openenv.environment_test import AGAIN @@ -219,13 +219,13 @@ def test_a_session_refused_after_one_was_dropped_leaves_only_that_rollout_unscor first.reset(**env.dataset[0]) first.end_turn(response="Added a crate.") assert first.get_reward() == 1.0 - trl._on_loop(first._close(), 30) + on_session_loop(first._close(), 30) second.reset(**env.dataset[0]) with caplog.at_level(logging.WARNING, logger=trl.__name__): first.reset(**env.dataset[0]) assert first.get_reward() is None - trl._on_loop(second._close(), 30) + on_session_loop(second._close(), 30) first.reset(**env.dataset[0]) first.end_turn(response="Added a crate.") diff --git a/tst/util/capabilities.py b/tst/util/capabilities.py index e964c064c..2eddbb726 100644 --- a/tst/util/capabilities.py +++ b/tst/util/capabilities.py @@ -49,6 +49,10 @@ #: only that tier's CI job installs them. VERL = "verl_installed" +#: The SkyRL tier (``tst/skyrl``) runs ``AgentEnvGenerator`` and ``AgentEnvTextEnv`` on SkyRL itself; SkyRL, skyrl-gym +#: and torch aren't in the dev extra, and only that tier's CI job installs them. +SKYRL = "skyrl_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/uv.lock b/uv.lock index 21b05729d..8fc2b1f6a 100644 --- a/uv.lock +++ b/uv.lock @@ -148,6 +148,12 @@ openenv = [ sail = [ { name = "sail" }, ] +skyrl = [ + { 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" }, +] smol = [ { name = "smolmachines" }, ] @@ -178,8 +184,9 @@ verl = [ [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", "verl"], marker = "extra == 'dev'" }, + { name = "agentenv-framework", extras = ["aws", "boat", "boxd", "coreweave", "daytona", "gcp", "openenv", "sail", "skyrl", "tensorlake", "tenki", "trl", "vercel", "verl"], marker = "extra == 'dev'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'nemo-gym'" }, + { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'skyrl'" }, { 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" }, @@ -216,6 +223,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 == 'skyrl'", specifier = ">=19" }, { name = "pyarrow", marker = "extra == 'verl'", specifier = ">=19" }, { name = "pydantic", specifier = ">=2.12" }, { name = "pymongo", specifier = ">=4.7" }, @@ -235,13 +243,14 @@ requires-dist = [ { name = "tenki", extras = ["async"], marker = "extra == 'tenki'", specifier = "~=1.4.1" }, { name = "tensorlake", marker = "extra == 'tensorlake'", specifier = "~=0.5.140" }, { name = "tomli-w", marker = "extra == 'dev'", specifier = "==1.2.0" }, + { name = "transformers", marker = "extra == 'skyrl'", specifier = ">=5" }, { 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", "verl", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] +provides-extras = ["openenv", "trl", "nemo-gym", "verl", "skyrl", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] [[package]] name = "agentenv-framework-protocol" From 7c229b6c4780f14753a63385693842c90d404d97 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 18:41:57 -0700 Subject: [PATCH 2/8] ci(training): install SkyRL without flash-attn on the CPU tier From PyPI, SkyRL 0.4.0's flash-attn builds from source against torch, and nothing on CPU uses it. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/skyrl.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/skyrl.yml b/.github/workflows/skyrl.yml index fcbc37dd9..7500c80fe 100644 --- a/.github/workflows/skyrl.yml +++ b/.github/workflows/skyrl.yml @@ -55,8 +55,11 @@ jobs: python -m pip install --quiet -r .github/ci-requirements.txt uv sync --locked --extra dev --extra skyrl --python "$(command -v python)" - name: Install SkyRL, skyrl-gym and CPU torch + # flash-attn is SkyRL's on Linux, but nothing on CPU uses it, and from PyPI it builds from source against torch. run: | + echo 'flash-attn; sys_platform == "never"' > "$RUNNER_TEMP/overrides.txt" uv pip install --python .venv/bin/python ${{ matrix.requirement }} torch \ + --override "$RUNNER_TEMP/overrides.txt" \ --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match .venv/bin/python -c "from importlib.metadata import version; print('skyrl', version('skyrl'), 'skyrl-gym', version('skyrl-gym'), 'torch', version('torch'))" - name: SkyRL tier From bbf81bba649d4e873f829db544bce42fea76b2b1 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 18:45:14 -0700 Subject: [PATCH 3/8] fix(training): SkyRL's env runs no tool call from the model's reasoning, and lets each turn left open take its deadline The calls are read from what follows the reply's last , up to a left open, so a call the model only thinks about isn't run; and ending the turns left open waits on each turn's own deadline, not on one overall timeout that could cut a long episode short. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/skyrl/calls.py | 17 ++++++++++++----- src/agent_env/training/skyrl/env.py | 15 ++++++++------- tst/unit/training/skyrl_test.py | 6 ++++++ 3 files changed, 26 insertions(+), 12 deletions(-) diff --git a/src/agent_env/training/skyrl/calls.py b/src/agent_env/training/skyrl/calls.py index 5de83e23b..0d6f4e4c6 100644 --- a/src/agent_env/training/skyrl/calls.py +++ b/src/agent_env/training/skyrl/calls.py @@ -13,7 +13,6 @@ FORMATS = ("hermes", "qwen3_coder") _BLOCK = re.compile(r"(.*?)(?:|$)", re.DOTALL) -_THINKING = re.compile(r".*?", re.DOTALL) _FUNCTION = re.compile(r"\s]+)>(.*?)(?:|$)", re.DOTALL) _PARAMETER = re.compile(r"\s]+)>(.*?)(?:|(?= tuple[list[Call | Unparsed], str]: - """Each ```` block of ``reply`` in order, as a call of one of ``tools`` (each function by its name) or - why it isn't one, and the reply's text outside them and its thinking, which ends a turn.""" + """Each ```` block of ``reply``'s answer in order, as a call of one of ``tools`` (each function by its + name) or why it isn't one, and the answer's text outside them, which ends a turn. The answer is what follows the + reply's reasoning, so a call the model only thinks about isn't run.""" if tool_format not in FORMATS: raise ValueError(f"tool_format must be one of {', '.join(FORMATS)}, not {tool_format!r}") reading = _hermes if tool_format == "hermes" else _qwen3_coder - calls = [_checked(reading(block.strip(), tools), tools) for block in _BLOCK.findall(reply)] - return calls, _THINKING.sub("", _BLOCK.sub("", reply)).strip() + answer = _answer(reply) + calls = [_checked(reading(block.strip(), tools), tools) for block in _BLOCK.findall(answer)] + return calls, _BLOCK.sub("", answer).strip() + + +def _answer(reply: str) -> str: + """What follows the reply's last ````, which closes reasoning a chat template may have opened in the + prompt, up to a ```` left open.""" + return reply.rsplit("", 1)[-1].split("", 1)[0] def _hermes(block: str, tools: Mapping[str, Mapping[str, Any]]) -> Call | Unparsed: diff --git a/src/agent_env/training/skyrl/env.py b/src/agent_env/training/skyrl/env.py index 75427a995..2d58dc871 100644 --- a/src/agent_env/training/skyrl/env.py +++ b/src/agent_env/training/skyrl/env.py @@ -57,7 +57,7 @@ def __init__(self, env_config: Any = None, extras: dict[str, Any] | None = None) self._run = self._unparsed = 0 def init(self, prompt: ConversationType) -> tuple[ConversationType, dict[str, Any]]: - self._wait(self._session.start(self._episode)) + self._wait(self._session.start(self._episode), self._waits) return prompt, {} def step(self, action: str) -> BaseTextEnvStepOutput: @@ -70,12 +70,12 @@ def step(self, action: str) -> BaseTextEnvStepOutput: answers.append({"role": "tool", "content": f"Error: {call.why}"}) continue self._run += 1 - answers.append({"role": "tool", "content": self._wait(self._session.call(call))}) + answers.append({"role": "tool", "content": self._wait(self._session.call(call), self._waits)}) if self._session.rollout.done or self._session.rollout.failure: break rollout = self._session.rollout if not (rollout.done or rollout.failure) and (not calls or self.turns >= self.max_turns): - self._wait(self._session.finish(reply)) + self._wait(self._session.finish(reply), None) if rollout.failure is not None: self.close() if rollout.misgrouped: @@ -98,13 +98,14 @@ def get_metrics(self) -> dict[str, Any]: def close(self) -> None: on_session_loop(self._session.close(), CLOSE_SECONDS) - def _wait(self, coroutine: Any) -> Any: - """The session's answer; one that takes past every deadline it has, and a grace, fails the episode.""" + def _wait(self, coroutine: Any, timeout: float | None) -> Any: + """The session's answer within ``timeout``, past which the episode fails: a call's deadlines and a grace, or + none for ending the turns left open, each of which has its own deadline.""" try: - return on_session_loop(coroutine, self._waits) + return on_session_loop(coroutine, timeout) except TimeoutError as e: self.close() - raise EpisodeFailed(f"{self._task()}: the env didn't answer within {self._waits:g}s") from e + raise EpisodeFailed(f"{self._task()}: the env didn't answer within {timeout:g}s") from e def _task(self) -> str: return f"task {self._episode['task_id']} of split {self._episode['split']}" diff --git a/tst/unit/training/skyrl_test.py b/tst/unit/training/skyrl_test.py index 88e6ea21d..5c9e6aad2 100644 --- a/tst/unit/training/skyrl_test.py +++ b/tst/unit/training/skyrl_test.py @@ -227,6 +227,12 @@ def test_the_command_writes_the_rows(serve, tmp_path, tokenizer, teardowns): [Call("schedule", {"count": 2}), Call("end_turn", {"response": "Done."})], ""), ('{"name": "schedule", "arguments": {}}', [Call("schedule", {})], ""), ("Idle is safest.Idling all day.", [], "Idling all day."), + ('Not yet: {"name": "end_turn", "arguments": {}}Idling all day.', [], + "Idling all day."), + ('Charge first, then {"name": "end_turn", "arguments": {}}\n' + '{"name": "schedule", "arguments": {"count": 1}}', [Call("schedule", {"count": 1})], ""), + ('Done.{"name": "end_turn", "arguments": {}}Or {"name": "schedule"}', + [Call("end_turn", {})], "Done."), ('{"name": "sell"}', [Unparsed("there's no tool 'sell'; the tools are end_turn, schedule")], ""), ("schedule(count=2)", [Unparsed( From e14716f53a8afb448084b89cf0cb6a0db7e13b25 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 19:11:36 -0700 Subject: [PATCH 4/8] fix(training): SkyRL's generator plays no more agentenv rollouts at once than it has env threads A rollout waiting for room on the server holds the env thread SkyRL calls it from, so an eval of 200 rollouts on 64 threads and a server of 64 sessions stalled: the waiting rollouts held every thread, and the ones holding sessions couldn't get one to end them. At most environment.skyrl_gym.max_env_workers rollouts now play at once, and none threads (max_env_workers=0, which calls envs on the generator's loop) is refused. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/skyrl/generator.py | 14 ++++++++++++- tst/skyrl/generator_test.py | 25 ++++++++++++++++++----- 2 files changed, 33 insertions(+), 6 deletions(-) diff --git a/src/agent_env/training/skyrl/generator.py b/src/agent_env/training/skyrl/generator.py index a3a076212..f6ce28a05 100644 --- a/src/agent_env/training/skyrl/generator.py +++ b/src/agent_env/training/skyrl/generator.py @@ -4,6 +4,7 @@ from __future__ import annotations +import asyncio import logging import time from dataclasses import asdict @@ -30,7 +31,9 @@ class AgentEnvGenerator(SkyRLGymGenerator): once its ``reset_timeout`` has passed since the rollout began. A rollout that fails every attempt, in training, has its whole group's loss mask and rewards zeroed, as SkyRL's Harbor generator does with a failed trial, so the group is out of the gradient and of its own baseline; in an eval it scores 0. Each rollout's env metrics hold its - attempts and whether it ``failed``. + attempts and whether it ``failed``. At most ``environment.skyrl_gym.max_env_workers`` rollouts play at once, one + per thread SkyRL calls envs from: a rollout waiting for room on the server holds its thread, so more rollouts than + threads would starve the ones holding sessions of the threads that end them. It leans on SkyRL internals, which CI checks on SkyRL 0.4.0 and main: ``agent_loop`` builds a new env per call and raises what the env raises, its output's ``env_metrics`` reach ``generate``'s, and ``generate``'s ``loss_masks``, @@ -42,11 +45,20 @@ def __init__(self, *args: Any, **kwargs: Any): if self.batched: raise ValueError("agentenv rows play an episode over several turns, which SkyRL's batched generation " "doesn't: set generator.batched=false") + if self.skyrl_gym_cfg.max_env_workers < 1: + raise ValueError("agentenv rows' envs wait on their server, which SkyRL's env threads let them do without " + "stalling generation: set environment.skyrl_gym.max_env_workers to at least 1") + self._playing = asyncio.Semaphore(self.skyrl_gym_cfg.max_env_workers) async def agent_loop(self, prompt: Any, env_class: str, env_extras: dict[str, Any], *args: Any, **kwargs: Any) -> TrajectoryOutput | StepWiseOutput: if env_class != AGENT: return await super().agent_loop(prompt, env_class, env_extras, *args, **kwargs) + async with self._playing: + return await self._played(prompt, env_class, env_extras, *args, **kwargs) + + async def _played(self, prompt: Any, env_class: str, env_extras: dict[str, Any], *args: Any, + **kwargs: Any) -> TrajectoryOutput | StepWiseOutput: server, episode = env_extras["server"], env_extras["episode"] began = time.monotonic() for attempt in range(1, server["attempts"] + 1): diff --git a/tst/skyrl/generator_test.py b/tst/skyrl/generator_test.py index b1b1af0b3..ee8bcca78 100644 --- a/tst/skyrl/generator_test.py +++ b/tst/skyrl/generator_test.py @@ -200,6 +200,18 @@ async def failing_after_the_first(self, context): assert len(deploys) == 3 +def test_more_rollouts_than_env_threads_play_in_turn_on_a_server_at_its_limit(serve, tmp_path, tool_format, tokenizer): + server = serve({"a": PLAYED}, max_concurrent_envs=2) + rows = _rows(server.url, tmp_path, tool_format) + + output = _generate(rows, _Engine(tokenizer, adds_a_crate(tool_format)), tokenizer, samples=6, env_workers=2, + reset_timeout=30.0) + server.stop() + + assert [sum(reward) for reward in output["rewards"]] == [1.0] * 6 + assert [metrics["attempts"] for metrics in output["env_metrics"]] == [1] * 6 + + def test_a_rollout_isnt_played_again_once_its_time_is_up(serve, tmp_path, tool_format, tokenizer, monkeypatch): server = serve({"a": PLAYED}) rows = _rows(server.url, tmp_path, tool_format) @@ -241,12 +253,14 @@ def test_a_server_that_asks_for_a_token_is_played_with_the_one_its_variable_hold assert token_vars([str(tmp_path / "all.parquet"), str(tmp_path / "absent.parquet")]) == {"DOCK_TOKEN"} -def test_the_entrypoint_registers_the_env_once_and_batched_generation_is_refused(tokenizer): +def test_the_entrypoint_registers_the_env_once_and_batched_generation_and_no_env_threads_are_refused(tokenizer): train.register() with pytest.raises(ValueError, match="agentenv rows play an episode over several turns, which SkyRL's batched " "generation doesn't: set generator.batched=false"): _generator(_Engine(tokenizer, None), tokenizer, batched=True, chat_template_kwargs={}) + with pytest.raises(ValueError, match="set environment.skyrl_gym.max_env_workers to at least 1"): + _generator(_Engine(tokenizer, None), tokenizer, env_workers=0) class _Engine: @@ -273,17 +287,17 @@ def _rows(url, folder, tool_format, **options): chat_template_kwargs=UNTHINKING, **options).folder / "all.parquet" -def _generator(engine, tokenizer, **settings): +def _generator(engine, tokenizer, *, env_workers=4, **settings): generator_config = config.GeneratorConfig(**{"max_turns": 4, "max_input_length": 8192, "chat_template_kwargs": UNTHINKING, **settings}) return skyrl_generator.AgentEnvGenerator(generator_cfg=generator_config, - skyrl_gym_cfg=config.SkyRLGymConfig(max_env_workers=4), + skyrl_gym_cfg=config.SkyRLGymConfig(max_env_workers=env_workers), inference_engine_client=engine, tokenizer=tokenizer, policy_model_name="policy") def _generate(rows, engine, tokenizer, *, samples=1, phase="train", step_wise=False, max_turns=4, - reset_timeout=None): + reset_timeout=None, env_workers=4): prompts = dataset.PromptDataset([str(rows)], tokenizer, max_prompt_length=8192, num_workers=1) batch = prompts.collate_fn([prompts[index] for index in range(len(prompts.dataframe))]) if reset_timeout is not None: @@ -291,6 +305,7 @@ def _generate(rows, engine, tokenizer, *, samples=1, phase="train", step_wise=Fa "reset_timeout": reset_timeout}}} for row in batch] generator_input, _ = generator_utils.prepare_generator_input(batch, samples, None, "agentenv", phase, 1) - generator = _generator(engine, tokenizer, max_turns=max_turns, step_wise_trajectories=step_wise) + generator = _generator(engine, tokenizer, env_workers=env_workers, max_turns=max_turns, + step_wise_trajectories=step_wise) return asyncio.run(generator.generate(generator_input)) From d9169281d478b0c71fa45ff591fffbf372b2bf2d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 19:40:02 -0700 Subject: [PATCH 5/8] fix(training): SkyRL's experiment builds the fully async trainer from the parts SkyRL hands it by name SkyRL's experiment calls get_trainer with keyword arguments, so the fully async branch unpacked nothing. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/skyrl/train.py | 10 ++++------ tst/skyrl/generator_test.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/src/agent_env/training/skyrl/train.py b/src/agent_env/training/skyrl/train.py index 7726d5e80..830b98773 100644 --- a/src/agent_env/training/skyrl/train.py +++ b/src/agent_env/training/skyrl/train.py @@ -48,13 +48,11 @@ def get_generator(self, cfg: SkyRLTrainConfig, tokenizer: Any, inference_engine_ inference_engine_client=inference_engine_client, tokenizer=tokenizer, policy_model_name=resolve_policy_model_name(cfg)) - def get_trainer(self, cfg: SkyRLTrainConfig, *args: Any, **kwargs: Any) -> Any: + def get_trainer(self, cfg: SkyRLTrainConfig, **parts: Any) -> Any: + """SkyRL's trainer, or its fully async one, built from the parts SkyRL's experiment hands its own.""" if not cfg.trainer.fully_async.enabled: - return super().get_trainer(cfg, *args, **kwargs) - tracker, tokenizer, train_dataset, eval_dataset, inference_engine_client, generator, colocate_pg = args - return FullyAsyncRayPPOTrainer(cfg=cfg, tracker=tracker, tokenizer=tokenizer, train_dataset=train_dataset, - eval_dataset=eval_dataset, inference_engine_client=inference_engine_client, - generator=generator, colocate_pg=colocate_pg) + return super().get_trainer(cfg=cfg, **parts) + return FullyAsyncRayPPOTrainer(cfg=cfg, **parts) @ray.remote(num_cpus=1) diff --git a/tst/skyrl/generator_test.py b/tst/skyrl/generator_test.py index ee8bcca78..b6c162616 100644 --- a/tst/skyrl/generator_test.py +++ b/tst/skyrl/generator_test.py @@ -10,6 +10,7 @@ import asyncio import json +from types import SimpleNamespace import pyarrow.parquet as pq import pytest @@ -263,6 +264,21 @@ def test_the_entrypoint_registers_the_env_once_and_batched_generation_and_no_env _generator(_Engine(tokenizer, None), tokenizer, env_workers=0) +@pytest.mark.parametrize(("fully_async", "built"), [(False, "RayPPOTrainer"), (True, "FullyAsyncRayPPOTrainer")]) +def test_the_experiment_hands_skyrls_trainer_parts_to_the_trainer_its_settings_pick(monkeypatch, fully_async, built): + trainers = [] + monkeypatch.setattr(train.BasePPOExp, "get_trainer", + lambda self, **parts: trainers.append(("RayPPOTrainer", parts))) + monkeypatch.setattr(train, "FullyAsyncRayPPOTrainer", lambda **parts: trainers.append(("FullyAsyncRayPPOTrainer", parts))) + cfg = SimpleNamespace(trainer=SimpleNamespace(fully_async=SimpleNamespace(enabled=fully_async))) + parts = {"tracker": "tracker", "tokenizer": "tokenizer", "train_dataset": "rows", "eval_dataset": "eval rows", + "inference_engine_client": "engine", "generator": "generator", "colocate_pg": None} + + object.__new__(train.AgentEnvExp).get_trainer(cfg=cfg, **parts) + + assert trainers == [(built, {"cfg": cfg, **parts})] + + class _Engine: """SkyRL's inference engine for the generator: ``script`` writes the reply to a trajectory's conversation from the tool answers it holds, and the reply is its tokens and the end of the turn.""" From 8d06f7e6e034ffc414f87edc041c9fa09b8fa43b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 19:44:13 -0700 Subject: [PATCH 6/8] fix(training): SkyRL's env reads reasoning tags only outside tool calls A tool's argument may hold or as text, code or quoted output, so reasoning ends at the first outside a tool call and resumes at a outside one. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/skyrl/calls.py | 18 +++++++++++++++--- tst/unit/training/skyrl_test.py | 11 +++++++++++ 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/src/agent_env/training/skyrl/calls.py b/src/agent_env/training/skyrl/calls.py index 0d6f4e4c6..a95997c5d 100644 --- a/src/agent_env/training/skyrl/calls.py +++ b/src/agent_env/training/skyrl/calls.py @@ -43,9 +43,21 @@ def parse(reply: str, tool_format: str, tools: Mapping[str, Mapping[str, Any]]) def _answer(reply: str) -> str: - """What follows the reply's last ````, which closes reasoning a chat template may have opened in the - prompt, up to a ```` left open.""" - return reply.rsplit("", 1)[-1].split("", 1)[0] + """What follows the reply's first ```` outside a tool call, which closes reasoning the reply or its chat + template's prompt opened, up to a ```` outside a tool call that opens more: a tool's argument may hold + either tag as text.""" + calls = [match.span() for match in _BLOCK.finditer(reply)] + + def outside(tag: str, start: int) -> int: + found = reply.find(tag, start) + while found != -1 and any(begin <= found < end for begin, end in calls): + found = reply.find(tag, found + 1) + return found + + closed = outside("", 0) + start = 0 if closed == -1 else closed + len("") + opened = outside("", start) + return reply[start:] if opened == -1 else reply[start:opened] def _hermes(block: str, tools: Mapping[str, Mapping[str, Any]]) -> Call | Unparsed: diff --git a/tst/unit/training/skyrl_test.py b/tst/unit/training/skyrl_test.py index 5c9e6aad2..4ee735423 100644 --- a/tst/unit/training/skyrl_test.py +++ b/tst/unit/training/skyrl_test.py @@ -233,6 +233,10 @@ def test_the_command_writes_the_rows(serve, tmp_path, tokenizer, teardowns): '{"name": "schedule", "arguments": {"count": 1}}', [Call("schedule", {"count": 1})], ""), ('Done.{"name": "end_turn", "arguments": {}}Or {"name": "schedule"}', [Call("end_turn", {})], "Done."), + ('{"name": "schedule", "arguments": {"note": "keep and as written"}}', + [Call("schedule", {"note": "keep and as written"})], ""), + ('Quote it.{"name": "schedule", "arguments": {"note": ""}}', + [Call("schedule", {"note": ""})], ""), ('{"name": "sell"}', [Unparsed("there's no tool 'sell'; the tools are end_turn, schedule")], ""), ("schedule(count=2)", [Unparsed( @@ -263,6 +267,13 @@ def test_qwen3_coder_calls_are_typed_by_the_tools_schema(): assert text == "I'll sell at the peak." +def test_a_qwen3_coder_calls_argument_may_hold_the_reasoning_tags(): + reply = ("Write it down.\n\n\n\nkeep and " + " as written\n\n\n") + + assert parse(reply, "qwen3_coder", TOOLS) == ([Call("schedule", {"note": "keep and as written"})], "") + + @pytest.mark.parametrize(("reply", "calls"), [ ("\n\n\nmany\n\n\n", [Call("schedule", {"count": "many"})]), From b03f4743c3b971f2647a8f2832307365a0ffa40f Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 20:08:18 -0700 Subject: [PATCH 7/8] fix(training): SkyRL's env names its metrics' means as SkyRL names a mean Where SkyRL joins groups' metrics, as its fully async trainer does, it averages a metric named as a mean and adds up the rest, so the env's attempts, scores and tool calls were the sum over a step's groups. Each is now environment/avg_. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/skyrl/env.py | 12 ++++++++++-- tst/skyrl/generator_test.py | 23 +++++++++++++++++++---- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/src/agent_env/training/skyrl/env.py b/src/agent_env/training/skyrl/env.py index 2d58dc871..ec9f33011 100644 --- a/src/agent_env/training/skyrl/env.py +++ b/src/agent_env/training/skyrl/env.py @@ -12,6 +12,7 @@ try: from skyrl_gym.envs.base_text_env import BaseTextEnv, BaseTextEnvStepOutput, ConversationType + from skyrl_gym.metrics import default_aggregate_metrics except ModuleNotFoundError as exc: if exc.name is None or exc.name.split(".")[0] != "skyrl_gym": raise @@ -38,8 +39,9 @@ class AgentEnvTextEnv(BaseTextEnv): returns its reward. A failure of the env raises EpisodeFailed; a task that ends unscored without failing, or offers a tool its row doesn't list, is a configuration error no replay changes, and raises ValueError. - ``get_metrics``, which SkyRL averages per step and per eval as ``environment/...``, holds each of the episode's - scores by the id of the step that recorded it (``score/``), the tool calls run and the unparsed ones.""" + ``get_metrics`` holds each of the episode's scores by the id of the step that recorded it (``score/``), + the tool calls run and the unparsed ones; SkyRL logs their means per step and per eval as + ``environment/avg_``.""" def __init__(self, env_config: Any = None, extras: dict[str, Any] | None = None): super().__init__() @@ -95,6 +97,12 @@ def get_metrics(self) -> dict[str, Any]: return {"tool_calls": self._run, "unparsed_calls": self._unparsed, **{f"score/{step}": score for step, score in self._scores.items()}} + @staticmethod + def aggregate_metrics(metrics: list[dict[str, Any]]) -> dict[str, Any]: + """Each metric's mean over the rollouts, named ``avg_`` as SkyRL names a mean: where SkyRL joins + groups' metrics, as its fully async trainer does, it averages a mean so named and adds the others up.""" + return {f"avg_{key}": value for key, value in default_aggregate_metrics(metrics).items()} + def close(self) -> None: on_session_loop(self._session.close(), CLOSE_SECONDS) diff --git a/tst/skyrl/generator_test.py b/tst/skyrl/generator_test.py index b6c162616..3ba0fd3df 100644 --- a/tst/skyrl/generator_test.py +++ b/tst/skyrl/generator_test.py @@ -84,11 +84,25 @@ async def finishing(self, reply): assert all(1 in mask and 0 in mask for mask in output["loss_masks"]) assert output["env_metrics"] == [{"tool_calls": 1, "unparsed_calls": 0, "score/replies": 1.0, "attempts": 1, "failed": 0}] * 2 - assert output["rollout_metrics"]["environment/score/replies"] == 1.0 + assert output["rollout_metrics"]["environment/avg_score/replies"] == 1.0 assert (replies, engine.calls) == (["Added a crate."] * 2, 4) assert len(teardowns.finished) == 3 +def test_the_env_metrics_of_groups_joined_as_skyrls_fully_async_trainer_joins_them_stay_means(serve, tmp_path, + tool_format, tokenizer): + server = serve({"a": PLAYED}) + rows = _rows(server.url, tmp_path, tool_format) + + groups = [_generate(rows, _Engine(tokenizer, adds_a_crate(tool_format)), tokenizer, samples=2, + sampling_params=config.SamplingParams(logprobs=0)) for _ in range(2)] + server.stop() + + joined = generator_utils.concatenate_generator_outputs(groups)["rollout_metrics"] + assert {key: joined[f"environment/avg_{key}"] for key in ("attempts", "failed", "tool_calls", "score/replies")} == { + "attempts": 1.0, "failed": 0.0, "tool_calls": 1.0, "score/replies": 1.0} + + def test_the_model_sees_its_tools_as_its_chat_template_renders_them(serve, tmp_path, tool_format, tokenizer): server = serve({"a": PLAYED}) rows = _rows(server.url, tmp_path, tool_format) @@ -292,7 +306,8 @@ async def generate(self, engine_input, model=None): self.first_prompt = self.first_prompt or list(prompt) reply = self.script(self.tokenizer.decode(prompt).count("")) ids = [*self.tokenizer.encode(reply, add_special_tokens=False), self.tokenizer.eos_token_id] - return {"responses": [reply], "response_ids": [ids], "stop_reasons": ["stop"], "response_logprobs": None} + return {"responses": [reply], "response_ids": [ids], "stop_reasons": ["stop"], + "response_logprobs": [[0.0] * len(ids)]} async def finish_session(self, session_id): pass @@ -313,7 +328,7 @@ def _generator(engine, tokenizer, *, env_workers=4, **settings): def _generate(rows, engine, tokenizer, *, samples=1, phase="train", step_wise=False, max_turns=4, - reset_timeout=None, env_workers=4): + reset_timeout=None, env_workers=4, **settings): prompts = dataset.PromptDataset([str(rows)], tokenizer, max_prompt_length=8192, num_workers=1) batch = prompts.collate_fn([prompts[index] for index in range(len(prompts.dataframe))]) if reset_timeout is not None: @@ -322,6 +337,6 @@ def _generate(rows, engine, tokenizer, *, samples=1, phase="train", step_wise=Fa for row in batch] generator_input, _ = generator_utils.prepare_generator_input(batch, samples, None, "agentenv", phase, 1) generator = _generator(engine, tokenizer, env_workers=env_workers, max_turns=max_turns, - step_wise_trajectories=step_wise) + step_wise_trajectories=step_wise, **settings) return asyncio.run(generator.generate(generator_input)) From 43fcbf2442a95191f6da3d0b4ef8b3408f87a635 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 00:10:37 -0700 Subject: [PATCH 8/8] refactor(training): SkyRL's env ends a failed episode itself, and the replay rule is shared with verl The env ends an episode it fails with reward 0 and its failed metric set, and AgentEnvGenerator plays the rollout again from that metric, so SkyRL builds every trajectory and nothing here constructs SkyRL's output types. The attempts rule and the prompt-length refusal move to _openenv, where verl's loop and writer use them too. A contract test pins the SkyRL internals the generator and experiment build on, and the docstrings are cut to the repo's length. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/_openenv.py | 22 +++++- src/agent_env/training/skyrl/__init__.py | 8 +- src/agent_env/training/skyrl/env.py | 91 ++++++++++++----------- src/agent_env/training/skyrl/generator.py | 76 ++++++------------- src/agent_env/training/skyrl/prepare.py | 32 +++----- src/agent_env/training/skyrl/train.py | 7 +- src/agent_env/training/verl/loop.py | 7 +- src/agent_env/training/verl/prepare.py | 12 +-- tst/skyrl/contract_test.py | 63 ++++++++++++++++ tst/skyrl/generator_test.py | 13 +--- tst/unit/training/skyrl_test.py | 6 +- 11 files changed, 177 insertions(+), 160 deletions(-) create mode 100644 tst/skyrl/contract_test.py diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index fa7181a61..7bb62668d 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, Mapping +from collections.abc import Coroutine, Iterator, Mapping from dataclasses import dataclass, field from typing import Any, ClassVar @@ -456,6 +456,26 @@ def function_tool(tool: Mapping[str, Any]) -> dict[str, Any]: }} +def attempt_numbers(limit: int, reset_timeout: float) -> Iterator[int]: + """The attempts at a rollout the env fails, 1 to ``limit``, none starting once ``reset_timeout`` seconds have passed + since the first began.""" + began = time.monotonic() + for attempt in range(1, limit + 1): + if attempt > 1 and time.monotonic() - began >= reset_timeout: + return + yield attempt + + +def longest_prompt(lengths: Mapping[tuple[str, str], int], limit: int | None, *, cut: str) -> int: + """The longest of the rows' prompts, in tokens by split and task id, refusing rows over ``limit``, for which ``cut`` + says what the trainer would do and the setting to raise.""" + 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}, {cut}") + return max(lengths.values(), default=0) + + def authorization(token: str | None) -> dict[str, str]: return {"Authorization": f"Bearer {token}"} if token else {} diff --git a/src/agent_env/training/skyrl/__init__.py b/src/agent_env/training/skyrl/__init__.py index b3b882f0f..0c43762b3 100644 --- a/src/agent_env/training/skyrl/__init__.py +++ b/src/agent_env/training/skyrl/__init__.py @@ -1,5 +1,3 @@ -"""SkyRL'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.skyrl URL OUT`` writes the server's tasks as SkyRL's parquet rows, each -with its tools in its system prompt, ``agent_env.training.skyrl.env.AgentEnvTextEnv`` plays each row's task as one -episode on the server, ``agent_env.training.skyrl.generator.AgentEnvGenerator`` plays again a rollout the env fails, -and ``python -m agent_env.training.skyrl.train`` runs SkyRL's trainer with both.""" +"""SkyRL on a bundle served as an OpenEnv env: ``python -m agent_env.training.skyrl`` writes SkyRL's rows, +``env.AgentEnvTextEnv`` plays each row's task as one episode, ``generator.AgentEnvGenerator`` plays a failed rollout +again, and ``python -m agent_env.training.skyrl.train`` runs SkyRL's trainer with both.""" diff --git a/src/agent_env/training/skyrl/env.py b/src/agent_env/training/skyrl/env.py index ec9f33011..43b934439 100644 --- a/src/agent_env/training/skyrl/env.py +++ b/src/agent_env/training/skyrl/env.py @@ -1,8 +1,5 @@ -"""SkyRL's env for the rows ``python -m agent_env.training.skyrl`` writes. ``AgentEnvTextEnv`` is a skyrl-gym -``BaseTextEnv`` playing a row's task as one episode on the bundle's OpenEnv server: each step runs the tool calls in the -model's reply on the episode and answers with their results, and a reply without one ends the turn it left open, which -scores the episode once its last turn ends. SkyRL calls an env from a thread of its pool, so the episode's session -runs on the event loop this process's sessions share.""" +"""SkyRL's env for the rows ``python -m agent_env.training.skyrl`` writes, playing each row's task as one episode on the +bundle's OpenEnv server. SkyRL calls envs from threads, so sessions run on the event loop this process's sessions share.""" from __future__ import annotations @@ -25,23 +22,9 @@ logger = logging.getLogger(__name__) -class EpisodeFailed(RuntimeError): - """The env failed the episode: a deploy, a call past its deadline, or a server that went away. AgentEnvGenerator - plays the rollout again on a new episode.""" - - class AgentEnvTextEnv(BaseTextEnv): - """Plays the row's task as one episode on the OpenEnv server its row names. ``init`` starts the episode, and - returns before it is deployed, so the deploy overlaps the model's first turn. Each ``step`` runs the tool calls of - the model's reply in order, a tool call SkyRL's parse finds nothing to call with answered with why, and returns - their answers as ``tool`` messages; a reply without a tool call ends the turn it left open, and any after it end - empty, as does the reply of the ``generator.max_turns``-th step, and the step that ends the episode's last turn - returns its reward. A failure of the env raises EpisodeFailed; a task that ends unscored without failing, or offers - a tool its row doesn't list, is a configuration error no replay changes, and raises ValueError. - - ``get_metrics`` holds each of the episode's scores by the id of the step that recorded it (``score/``), - the tool calls run and the unparsed ones; SkyRL logs their means per step and per eval as - ``environment/avg_``.""" + """An episode the env fails ends with reward 0 and its ``failed`` metric set, for AgentEnvGenerator to play again, + so SkyRL builds every trajectory itself; a task that can't be scored at all raises.""" def __init__(self, env_config: Any = None, extras: dict[str, Any] | None = None): super().__init__() @@ -57,13 +40,39 @@ def __init__(self, env_config: Any = None, extras: dict[str, Any] | None = None) waits_for_server=True, token=token_in(server.get("token_var")))) self._scores: dict[str, float] = {} self._run = self._unparsed = 0 + self._failed = False def init(self, prompt: ConversationType) -> tuple[ConversationType, dict[str, Any]]: - self._wait(self._session.start(self._episode), self._waits) + try: + self._wait(self._session.start(self._episode), self._waits) + except _Failed as e: + self._fail(e) return prompt, {} def step(self, action: str) -> BaseTextEnvStepOutput: self.turns += 1 + if self._failed: + return _ended(0.0) + try: + return self._stepped(action) + except _Failed as e: + self._fail(e) + return _ended(0.0) + + def get_metrics(self) -> dict[str, Any]: + return {"tool_calls": self._run, "unparsed_calls": self._unparsed, "failed": int(self._failed), + **{f"score/{step}": score for step, score in self._scores.items()}} + + @staticmethod + def aggregate_metrics(metrics: list[dict[str, Any]]) -> dict[str, Any]: + """Named ``avg_`` as SkyRL names a mean, so where it joins groups' metrics, as its fully async trainer + does, it averages them rather than adding them up.""" + return {f"avg_{key}": value for key, value in default_aggregate_metrics(metrics).items()} + + def close(self) -> None: + on_session_loop(self._session.close(), CLOSE_SECONDS) + + def _stepped(self, action: str) -> BaseTextEnvStepOutput: calls, reply = parse(action, self._format, self._tools) answers = [] for call in calls: @@ -79,11 +88,11 @@ def step(self, action: str) -> BaseTextEnvStepOutput: if not (rollout.done or rollout.failure) and (not calls or self.turns >= self.max_turns): self._wait(self._session.finish(reply), None) if rollout.failure is not None: - self.close() if rollout.misgrouped: + self.close() raise ValueError(f"{self._task()}: {rollout.failure}. Its row lists the tools of its toolset's first " "task, so a task must offer no tool that one doesn't") - raise EpisodeFailed(f"{self._task()}: {rollout.failure}") + raise _Failed(rollout.failure) if not rollout.done: return BaseTextEnvStepOutput(observations=answers, reward=0.0, done=False, metadata={}) if rollout.reward is None: @@ -91,38 +100,32 @@ def step(self, action: str) -> BaseTextEnvStepOutput: 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") self._scores = rollout.scores or {} - return BaseTextEnvStepOutput(observations=[], reward=float(rollout.reward), done=True, metadata={}) + return _ended(float(rollout.reward)) - def get_metrics(self) -> dict[str, Any]: - return {"tool_calls": self._run, "unparsed_calls": self._unparsed, - **{f"score/{step}": score for step, score in self._scores.items()}} - - @staticmethod - def aggregate_metrics(metrics: list[dict[str, Any]]) -> dict[str, Any]: - """Each metric's mean over the rollouts, named ``avg_`` as SkyRL names a mean: where SkyRL joins - groups' metrics, as its fully async trainer does, it averages a mean so named and adds the others up.""" - return {f"avg_{key}": value for key, value in default_aggregate_metrics(metrics).items()} - - def close(self) -> None: - on_session_loop(self._session.close(), CLOSE_SECONDS) + def _fail(self, failure: _Failed) -> None: + self._failed = True + logger.warning("%s: %s", self._task(), failure) + self.close() def _wait(self, coroutine: Any, timeout: float | None) -> Any: - """The session's answer within ``timeout``, past which the episode fails: a call's deadlines and a grace, or - none for ending the turns left open, each of which has its own deadline.""" try: return on_session_loop(coroutine, timeout) except TimeoutError as e: - self.close() - raise EpisodeFailed(f"{self._task()}: the env didn't answer within {timeout:g}s") from e + raise _Failed(f"the env didn't answer within {timeout:g}s") from e def _task(self) -> str: return f"task {self._episode['task_id']} of split {self._episode['split']}" -class _Session(Session): - """One rollout's session: it starts the row's episode, runs the model's tool calls on it, and ends the turn left - open with the model's reply.""" +class _Failed(Exception): + pass + +def _ended(reward: float) -> BaseTextEnvStepOutput: + return BaseTextEnvStepOutput(observations=[], reward=reward, done=True, metadata={}) + + +class _Session(Session): offered_by = "its row lists" def __init__(self, shared: Shared): diff --git a/src/agent_env/training/skyrl/generator.py b/src/agent_env/training/skyrl/generator.py index f6ce28a05..c61b4cd2e 100644 --- a/src/agent_env/training/skyrl/generator.py +++ b/src/agent_env/training/skyrl/generator.py @@ -1,13 +1,10 @@ -"""SkyRL's generator for the rows ``python -m agent_env.training.skyrl`` writes, which plays a rollout the env fails -again and, when every attempt fails, leaves its group out of training: SkyRL fails a whole batch on an env's exception, -and has no way to leave out one trajectory.""" +"""SkyRL's generator for the rows ``python -m agent_env.training.skyrl`` writes. SkyRL can't leave one trajectory out +of a step, so a rollout the env fails is played again, and one it fails every time leaves its whole group out.""" from __future__ import annotations import asyncio import logging -import time -from dataclasses import asdict from typing import Any try: @@ -19,26 +16,15 @@ raise ImportError("agent_env.training.skyrl.generator runs in SkyRL's trainer: install " "'agentenv-framework[skyrl]' into SkyRL's environment") from exc -from agent_env.training.skyrl.env import EpisodeFailed +from agent_env.training._openenv import attempt_numbers from agent_env.training.skyrl.prepare import AGENT logger = logging.getLogger(__name__) class AgentEnvGenerator(SkyRLGymGenerator): - """SkyRL's ``SkyRLGymGenerator`` that plays a rollout of an ``agentenv`` row the env fails (EpisodeFailed) again - from the start on a new episode, a fresh sample of the same policy: up to the row's ``attempts`` times, but not - once its ``reset_timeout`` has passed since the rollout began. A rollout that fails every attempt, in training, has - its whole group's loss mask and rewards zeroed, as SkyRL's Harbor generator does with a failed trial, so the group - is out of the gradient and of its own baseline; in an eval it scores 0. Each rollout's env metrics hold its - attempts and whether it ``failed``. At most ``environment.skyrl_gym.max_env_workers`` rollouts play at once, one - per thread SkyRL calls envs from: a rollout waiting for room on the server holds its thread, so more rollouts than - threads would starve the ones holding sessions of the threads that end them. - - It leans on SkyRL internals, which CI checks on SkyRL 0.4.0 and main: ``agent_loop`` builds a new env per call and - raises what the env raises, its output's ``env_metrics`` reach ``generate``'s, and ``generate``'s ``loss_masks``, - ``rewards`` and ``env_metrics`` are a list per trajectory, or per step, as its ``trajectory_ids`` or the input - batch's are.""" + """At most ``max_env_workers`` agentenv rollouts play at once: a rollout waiting for room on the server holds the + thread it was called from, so more would starve the rollouts holding sessions of the threads that end them.""" def __init__(self, *args: Any, **kwargs: Any): super().__init__(*args, **kwargs) @@ -60,45 +46,27 @@ async def agent_loop(self, prompt: Any, env_class: str, env_extras: dict[str, An async def _played(self, prompt: Any, env_class: str, env_extras: dict[str, Any], *args: Any, **kwargs: Any) -> TrajectoryOutput | StepWiseOutput: server, episode = env_extras["server"], env_extras["episode"] - began = time.monotonic() - for attempt in range(1, server["attempts"] + 1): - try: - output = await super().agent_loop(prompt, env_class, dict(env_extras), *args, **kwargs) - except EpisodeFailed as e: - failure = e - logger.warning("attempt %d of %d: %s", attempt, server["attempts"], e) - if time.monotonic() - began >= server["reset_timeout"]: - break - continue - last = output.step_outputs[-1] if isinstance(output, StepWiseOutput) else output - last.env_metrics.update(attempts=attempt, failed=0) - return output - logger.warning("task %s of split %s: the env failed it on %d attempt%s, so its group is left out of training: " - "%s", episode["task_id"], episode["split"], attempt, "s" * (attempt > 1), failure) - return self._left_out(attempt, time.monotonic() - began, kwargs.get("sampling_params")) + for attempt in attempt_numbers(server["attempts"], server["reset_timeout"]): + output = await super().agent_loop(prompt, env_class, dict(env_extras), *args, **kwargs) + metrics = (output.step_outputs[-1] if isinstance(output, StepWiseOutput) else output).env_metrics + metrics["attempts"] = attempt + if not metrics.get("failed"): + return output + logger.warning("task %s of split %s: the env failed it on %d attempt%s, so its group is left out of training", + episode["task_id"], episode["split"], attempt, "s" * (attempt > 1)) + return output async def generate(self, input_batch: GeneratorInput, *args: Any, **kwargs: Any) -> GeneratorOutput: output = await super().generate(input_batch, *args, **kwargs) failed = [row for row, metrics in enumerate(output["env_metrics"]) if (metrics or {}).get("failed")] metadata = input_batch.get("batch_metadata") ids = output.get("trajectory_ids") or input_batch.get("trajectory_ids") - if failed and ids and getattr(metadata, "training_phase", None) != TRAINING_PHASE_EVAL: - groups = {ids[row].instance_id for row in failed} - failed = [row for row, trajectory in enumerate(ids) if trajectory.instance_id in groups] - for row in failed: - output["loss_masks"][row] = [0] * len(output["loss_masks"][row]) - reward = output["rewards"][row] - output["rewards"][row] = [0.0] * len(reward) if isinstance(reward, list) else 0.0 + if not failed or not ids or getattr(metadata, "training_phase", None) == TRAINING_PHASE_EVAL: + return output + groups = {ids[row].instance_id for row in failed} + for row, trajectory in enumerate(ids): + if trajectory.instance_id in groups: + output["loss_masks"][row] = [0] * len(output["loss_masks"][row]) + reward = output["rewards"][row] + output["rewards"][row] = [0.0] * len(reward) if isinstance(reward, list) else 0.0 return output - - def _left_out(self, attempts: int, seconds: float, - sampling_params: dict[str, Any] | None) -> TrajectoryOutput | StepWiseOutput: - """A placeholder for a rollout every attempt of which the env failed: one token, masked, rewarded 0.""" - sampling = sampling_params if sampling_params is not None else asdict(self.generator_cfg.sampling_params) - token = self.tokenizer.eos_token_id - output = TrajectoryOutput(response_ids=[token], reward=[0.0], stop_reason="error", loss_mask=[0], - prompt_ids=[token], - rollout_logprobs=[0.0] if sampling.get("logprobs") is not None else None, - env_metrics={"attempts": attempts, "failed": 1}, e2e_time=seconds, - time_splits={"llm": 0.0, "env": 0.0}) - return StepWiseOutput(step_outputs=[output]) if self.generator_cfg.step_wise_trajectories else output diff --git a/src/agent_env/training/skyrl/prepare.py b/src/agent_env/training/skyrl/prepare.py index 7b9306d82..a1fe2ca43 100644 --- a/src/agent_env/training/skyrl/prepare.py +++ b/src/agent_env/training/skyrl/prepare.py @@ -13,13 +13,12 @@ import pyarrow.parquet as pq from transformers import AutoTokenizer -from agent_env.training._openenv import function_tool, key, messages, tasks_and_tools, token_in +from agent_env.training._openenv import function_tool, key, longest_prompt, messages, tasks_and_tools, token_in from agent_env.training.skyrl.calls import FORMATS #: The ``env_class`` every row plays under, and the id the env is registered at. AGENT = "agentenv" -#: A row's columns, so a split the server keeps empty still has them. SkyRL hands the env every column but ``prompt`` -#: and ``env_class`` as its extras, and groups eval metrics by ``data_source``. +#: A row's columns, so a split the server keeps empty still has them. ROW = pa.schema([ ("prompt", pa.list_(pa.struct([("role", pa.string()), ("content", pa.string())]))), ("env_class", pa.string()), @@ -45,18 +44,9 @@ def write_dataset(url: str, out: str | Path, *, tokenizer: str, tool_format: str tool_timeout: float = 60.0, reset_timeout: float = 600.0, attempts: int = 3, token_var: 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. A row's - prompt is its task's system prompt and first prompt, with every tool its task's turns offer, found by playing one - task of each toolset to its last turn, unscored, so writing deploys those. The tools go in the system prompt as - ``tokenizer``'s chat template renders them with ``tools=`` and ``chat_template_kwargs``, SkyRL's - ``generator.chat_template_kwargs``, checked to render the same text; ``tool_format`` names the format the template - teaches the model to call them in. Each row also holds the server's settings for its env: the deadlines, the - ``attempts`` a rollout the env fails is played, and ``token_var``, the environment variable that holds the token a - server given one asks for (``build_app(token=...)``), read now and in the trainer, so the token is in no file. - - A row's prompt is counted as SkyRL's prompt filter counts it, without ``chat_template_kwargs``, and as its generator - renders it, with them; a row over ``max_prompt_length``, SkyRL's ``trainer.max_prompt_length``, in either is - refused before anything is written, where the filter would drop it.""" + """The tasks of the OpenEnv server at ``url`` as ``out/.parquet``, each row's tools in its system prompt as + ``tokenizer``'s chat template renders them, since SkyRL's generator passes the template none. ``token_var`` names + the variable holding the server's token, read here and in the trainer, so the token is in no file.""" if attempts < 1: raise ValueError(f"attempts must be at least 1, not {attempts}") if tool_format not in FORMATS: @@ -118,15 +108,11 @@ def _rendered(tokenizer: Any, prompt: list[dict[str, str]], chat_template_kwargs 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, counted with and without ``chat_template_kwargs``, refusing rows over - ``limit``.""" + """Counted as SkyRL's prompt filter counts a row, without ``chat_template_kwargs``, and as its generator renders it, + with them.""" lengths = {(split, row["episode"]["task_id"]): max( len(tokenizer.apply_chat_template(row["prompt"], add_generation_prompt=True, tokenize=True, return_dict=True, **kwargs)["input_ids"]) for kwargs in ({}, chat_template_kwargs)) 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 SkyRL's prompt filter would drop: " - "raise it, and trainer.max_prompt_length with it") - return max(lengths.values(), default=0) + return longest_prompt(lengths, limit, cut="which SkyRL's prompt filter would drop: raise it, and " + "trainer.max_prompt_length with it") diff --git a/src/agent_env/training/skyrl/train.py b/src/agent_env/training/skyrl/train.py index 830b98773..e3c8c3031 100644 --- a/src/agent_env/training/skyrl/train.py +++ b/src/agent_env/training/skyrl/train.py @@ -1,8 +1,5 @@ -"""``python -m agent_env.training.skyrl.train``, or SkyRL's ``main_base`` with -``trainer.override_entrypoint=agent_env.training.skyrl.train``: SkyRL's trainer on the rows -``python -m agent_env.training.skyrl`` writes, with the env registered as ``agentenv`` and AgentEnvGenerator -generating, and SkyRL's fully async trainer with ``trainer.fully_async.enabled=true``. It takes SkyRL's own -settings.""" +"""SkyRL's trainer with the env registered and AgentEnvGenerator generating: ``python -m agent_env.training.skyrl.train``, +or SkyRL's ``main_base`` with ``trainer.override_entrypoint=agent_env.training.skyrl.train``.""" from __future__ import annotations diff --git a/src/agent_env/training/verl/loop.py b/src/agent_env/training/verl/loop.py index d097d95fb..d7d69109c 100644 --- a/src/agent_env/training/verl/loop.py +++ b/src/agent_env/training/verl/loop.py @@ -23,7 +23,7 @@ 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, token_in +from agent_env.training._openenv import Rollout, Session, Shared, attempt_numbers, token_in logger = logging.getLogger(__name__) #: The tools whose answers verl was seen to cut, each warned of once a process. @@ -83,8 +83,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop 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): + for attempt in attempt_numbers(self._attempts, self._reset_timeout): session = _Session(shared) try: output, rollout = await self._play(session, schemas, sampling_params, kwargs) @@ -102,8 +101,6 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop } 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}") diff --git a/src/agent_env/training/verl/prepare.py b/src/agent_env/training/verl/prepare.py index b1889860b..ec9737f45 100644 --- a/src/agent_env/training/verl/prepare.py +++ b/src/agent_env/training/verl/prepare.py @@ -14,7 +14,7 @@ import yaml from transformers import AutoTokenizer -from agent_env.training._openenv import function_tool, key, messages, tasks_and_tools, token_in +from agent_env.training._openenv import function_tool, key, longest_prompt, messages, tasks_and_tools, token_in #: The ``agent_name`` every row plays under, and the name ``agent_loop.yaml`` registers the loop at. AGENT = "agentenv" @@ -83,14 +83,8 @@ def write_dataset(url: str, out: str | Path, *, tool_timeout: float = 60.0, rese 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) - + return longest_prompt(lengths, limit, cut="which verl would cut from the left: raise it, and data.max_prompt_length " + "with it") diff --git a/tst/skyrl/contract_test.py b/tst/skyrl/contract_test.py new file mode 100644 index 000000000..72f2c4701 --- /dev/null +++ b/tst/skyrl/contract_test.py @@ -0,0 +1,63 @@ +"""The parts of SkyRL that ``AgentEnvGenerator`` and ``AgentEnvExp`` build on, so a SkyRL that changes one fails here, +naming it, before a rollout does.""" + +import dataclasses +import inspect +import typing + +import pytest + +from tst.util.capabilities import SKYRL, missing_capability_reason + +train = pytest.importorskip("agent_env.training.skyrl.train", reason=missing_capability_reason(SKYRL), + exc_type=ImportError) +base = pytest.importorskip("skyrl.train.generators.base", reason=missing_capability_reason(SKYRL)) +gym_generator = pytest.importorskip("skyrl.train.generators.skyrl_gym_generator", + reason=missing_capability_reason(SKYRL)) +config = pytest.importorskip("skyrl.train.config", reason=missing_capability_reason(SKYRL)) +trainer = pytest.importorskip("skyrl.train.trainer", reason=missing_capability_reason(SKYRL)) + + +def test_the_agent_loop_takes_a_rows_prompt_env_class_and_extras_first(): + parameters = list(inspect.signature(gym_generator.SkyRLGymGenerator.agent_loop).parameters) + + assert parameters[1:4] == ["prompt", "env_class", "env_extras"] + + +def test_the_generator_finds_env_metrics_loss_masks_rewards_and_phases_where_it_reads_them(): + assert "env_metrics" in _fields(gym_generator.TrajectoryOutput) + assert "step_outputs" in _fields(gym_generator.StepWiseOutput) + assert {"env_metrics", "loss_masks", "rewards", "trajectory_ids"} <= set(typing.get_type_hints(base.GeneratorOutput)) + assert {"batch_metadata", "trajectory_ids"} <= set(typing.get_type_hints(base.GeneratorInput)) + assert "training_phase" in _fields(base.BatchMetadata) + + +def test_the_experiment_builds_its_generator_from_what_skyrls_own_experiment_passes(monkeypatch): + built = [] + + def recorded(self, *args, **kwargs): + built.append((args, kwargs)) + raise _Built + + monkeypatch.setattr(gym_generator.SkyRLGymGenerator, "__init__", recorded) + cfg, experiment = config.SkyRLTrainConfig(), object.__new__(train.AgentEnvExp) + for build in (train.BasePPOExp.get_generator, train.AgentEnvExp.get_generator): + with pytest.raises(_Built): + build(experiment, cfg, "tokenizer", "engine") + + assert built[0] == built[1] + + +def test_skyrls_trainers_take_the_parts_its_experiment_hands_get_trainer(): + parts = set(inspect.signature(train.BasePPOExp.get_trainer).parameters) - {"self"} + + assert parts <= set(inspect.signature(trainer.RayPPOTrainer.__init__).parameters) + assert issubclass(train.FullyAsyncRayPPOTrainer, trainer.RayPPOTrainer) + + +class _Built(Exception): + pass + + +def _fields(cls): + return {field.name for field in dataclasses.fields(cls)} diff --git a/tst/skyrl/generator_test.py b/tst/skyrl/generator_test.py index 3ba0fd3df..d7e2c5208 100644 --- a/tst/skyrl/generator_test.py +++ b/tst/skyrl/generator_test.py @@ -1,12 +1,5 @@ -"""``AgentEnvGenerator`` and ``AgentEnvTextEnv`` on SkyRL itself: SkyRL's ``PromptDataset`` loads the rows -``write_dataset`` wrote, with Qwen3's and Qwen3.5's tokenizers and chat templates, and its generator plays them through -a stand-in inference engine, on the unit tier's bundle served on loopback. A played row's reward is its episode's, and -its scores, attempts and tool calls its env metrics, through the hermes and qwen3_coder formats, over two turns, and -from a server that asks for a token; the model sees its tools as its chat template renders them; an unparsed call is -answered with why; the turn cap ends the episode; a rollout the env fails is played again on a new episode, but not -once its time is up, and one it fails every attempt masks its group in training but not in an eval; an episode that -ends unscored without failing raises; and the entrypoint registers the env and reads the rows' token variables. CI runs -it on SkyRL 0.4.0 and main.""" +"""``AgentEnvGenerator`` and ``AgentEnvTextEnv`` on SkyRL itself: its ``PromptDataset`` loads the written rows and its +generator plays them through a stand-in inference engine, with Qwen3's and Qwen3.5's tokenizers.""" import asyncio import json @@ -211,7 +204,7 @@ async def failing_after_the_first(self, context): assert sum(output["rewards"][played]) == kept assert (1 in output["loss_masks"][played]) == (phase == "eval") failed = next(row for row, metrics in enumerate(output["env_metrics"]) if metrics.get("failed") == 1) - assert (output["loss_masks"][failed], sum(output["rewards"][failed])) == ([0], 0.0) + assert ((1 in output["loss_masks"][failed]), sum(output["rewards"][failed])) == (phase == "eval", 0.0) assert len(deploys) == 3 diff --git a/tst/unit/training/skyrl_test.py b/tst/unit/training/skyrl_test.py index 4ee735423..35c5d4339 100644 --- a/tst/unit/training/skyrl_test.py +++ b/tst/unit/training/skyrl_test.py @@ -1,7 +1,5 @@ -"""``write_dataset`` on the OpenEnv server ``build_app`` makes, served on loopback over the Episode API's fake steps: -SkyRL's parquet rows for each split, each with its tools in its system prompt as a local tokenizer's chat template -renders them, from a server that asks for a token too, and the command that writes them; and the tool-call parsers the -env reads the model's replies with.""" +"""``write_dataset`` on the OpenEnv server ``build_app`` makes, served on loopback over the Episode API's fake steps, +with a local tokenizer whose template renders tools as Qwen3's does; and the parsers the env reads replies with.""" import json