diff --git a/.github/workflows/skyrl.yml b/.github/workflows/skyrl.yml new file mode 100644 index 00000000..7500c80f --- /dev/null +++ b/.github/workflows/skyrl.yml @@ -0,0 +1,69 @@ +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 + # 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 + 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 3adbc578..1e184964 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 87d9b1c2..1b557f90 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 70ffafe2..2101b0be 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 1440ad6e..ed79a783 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 2e226cac..7bb62668 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, Iterator, 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,63 @@ 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 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 new file mode 100644 index 00000000..0c43762b --- /dev/null +++ b/src/agent_env/training/skyrl/__init__.py @@ -0,0 +1,3 @@ +"""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/__main__.py b/src/agent_env/training/skyrl/__main__.py new file mode 100644 index 00000000..dd372a8e --- /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 00000000..a95997c5 --- /dev/null +++ b/src/agent_env/training/skyrl/calls.py @@ -0,0 +1,124 @@ +"""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) +_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``'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 + 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 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: + 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 00000000..43b93443 --- /dev/null +++ b/src/agent_env/training/skyrl/env.py @@ -0,0 +1,149 @@ +"""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 + +import json +import logging +from typing import Any + +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 + 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 AgentEnvTextEnv(BaseTextEnv): + """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__() + 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 + self._failed = False + + def init(self, prompt: ConversationType) -> tuple[ConversationType, dict[str, Any]]: + 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: + 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), 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), None) + if rollout.failure is not None: + 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 _Failed(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 _ended(float(rollout.reward)) + + 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: + try: + return on_session_loop(coroutine, timeout) + except TimeoutError as 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 _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): + 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 00000000..c61b4cd2 --- /dev/null +++ b/src/agent_env/training/skyrl/generator.py @@ -0,0 +1,72 @@ +"""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 +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._openenv import attempt_numbers +from agent_env.training.skyrl.prepare import AGENT + +logger = logging.getLogger(__name__) + + +class AgentEnvGenerator(SkyRLGymGenerator): + """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) + 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"] + 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 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 diff --git a/src/agent_env/training/skyrl/prepare.py b/src/agent_env/training/skyrl/prepare.py new file mode 100644 index 00000000..a1fe2ca4 --- /dev/null +++ b/src/agent_env/training/skyrl/prepare.py @@ -0,0 +1,118 @@ +"""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, 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. +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``, 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: + 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: + """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} + 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 new file mode 100644 index 00000000..e3c8c303 --- /dev/null +++ b/src/agent_env/training/skyrl/train.py @@ -0,0 +1,72 @@ +"""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 + +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, **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=cfg, **parts) + return FullyAsyncRayPPOTrainer(cfg=cfg, **parts) + + +@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 fd96f4f6..f595e5b8 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 f1371606..d7d69109 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, attempt_numbers, token_in logger = logging.getLogger(__name__) #: The tools whose answers verl was seen to cut, each warned of once a process. @@ -84,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) @@ -103,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 17d1b3da..ec9737f4 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, 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" @@ -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,32 +81,10 @@ 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``.""" lengths = {(split, row["extra_info"]["task_id"]): len(tokenizer.apply_chat_template( row["prompt"], tools=json.loads(row["extra_info"]["tools"]), add_generation_prompt=True, tokenize=True, return_dict=True, **chat_template_kwargs)["input_ids"]) for split, written in rows.items() for row in written} - if limit is not None and (over := {at: length for at, length in lengths.items() if length > limit}): - (split, task), length = max(over.items(), key=lambda item: item[1]) - raise ValueError(f"{len(over)} rows' prompts are over max_prompt_length={limit} tokens with their tools, the " - f"longest task {task} of split {split} at {length}, which verl would cut from the left: raise " - "it, and data.max_prompt_length with it") - return max(lengths.values(), default=0) - - -def _function(tool: dict[str, Any]) -> dict[str, Any]: - """An MCP tool as the function tool a chat template renders.""" - return {"type": "function", "function": { - "name": tool["name"], "description": tool.get("description") or "", - "parameters": tool.get("input_schema") or {"type": "object", "properties": {}}, - }} + 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/__init__.py b/tst/skyrl/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tst/skyrl/conftest.py b/tst/skyrl/conftest.py new file mode 100644 index 00000000..d37f30af --- /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/contract_test.py b/tst/skyrl/contract_test.py new file mode 100644 index 00000000..72f2c470 --- /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 new file mode 100644 index 00000000..d7e2c520 --- /dev/null +++ b/tst/skyrl/generator_test.py @@ -0,0 +1,335 @@ +"""``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 +from types import SimpleNamespace + +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/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) + 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 ((1 in output["loss_masks"][failed]), sum(output["rewards"][failed])) == (phase == "eval", 0.0) + 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) + 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_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) + + +@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.""" + + 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": [[0.0] * len(ids)]} + + 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, *, 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=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, 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: + 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, env_workers=env_workers, max_turns=max_turns, + step_wise_trajectories=step_wise, **settings) + 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 00cfc15a..79004823 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 00000000..35c5d433 --- /dev/null +++ b/tst/unit/training/skyrl_test.py @@ -0,0 +1,297 @@ +"""``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 + +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."), + ('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": "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( + "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." + + +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"})]), + ("\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 0335e1cc..21b31a17 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 e964c064..2eddbb72 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 21b05729..8fc2b1f6 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"