From f6aff95c7597c2c4e30f49c5275fb3d92b153961 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 2 Nov 2026 23:39:56 -0800 Subject: [PATCH 1/4] feat(training): train on a bundle with slime through its OpenAI adapter over the OpenEnv server agent_env.training.slime's generate function opens a session for each sample on slime's OpenAI adapter (slime.agent), which records the model's turns, plays the row's episode through it and returns the session's turns as samples carrying the episode's reward. The chat rows and the episode played over a trainer's chat endpoint move from the Miles helper into agent_env.training._chat, which both helpers share; the Miles helper's public names stay. tst/slime plays rows through slime's own adapter with slime's own arguments, and slime.yml runs it on slime 540ea91 and main. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/slime.yml | 81 +++++++++ AGENTS.md | 7 +- Makefile | 8 +- pyproject.toml | 5 + src/agent_env/training/__init__.py | 4 +- src/agent_env/training/_chat.py | 164 ++++++++++++++++++ src/agent_env/training/miles/__main__.py | 29 +--- src/agent_env/training/miles/agent.py | 113 +----------- src/agent_env/training/miles/prepare.py | 31 +--- src/agent_env/training/slime/__init__.py | 5 + src/agent_env/training/slime/__main__.py | 6 + src/agent_env/training/slime/generate.py | 59 +++++++ tst/slime/__init__.py | 0 tst/slime/conftest.py | 7 + tst/slime/generate_test.py | 149 ++++++++++++++++ .../test_optional_extras_stay_optional.py | 11 +- tst/unit/training/miles_test.py | 14 +- tst/util/capabilities.py | 4 + uv.lock | 6 +- 19 files changed, 522 insertions(+), 181 deletions(-) create mode 100644 .github/workflows/slime.yml create mode 100644 src/agent_env/training/_chat.py create mode 100644 src/agent_env/training/slime/__init__.py create mode 100644 src/agent_env/training/slime/__main__.py create mode 100644 src/agent_env/training/slime/generate.py create mode 100644 tst/slime/__init__.py create mode 100644 tst/slime/conftest.py create mode 100644 tst/slime/generate_test.py diff --git a/.github/workflows/slime.yml b/.github/workflows/slime.yml new file mode 100644 index 000000000..b3a1a3f46 --- /dev/null +++ b/.github/workflows/slime.yml @@ -0,0 +1,81 @@ +name: slime contract + +# agent_env.training.slime on slime itself (tst/slime): the helper's generate function plays rows the writer wrote +# through slime's own OpenAI adapter, with slime's own arguments, in front of a stand-in for SGLang's /generate, and +# slime's dynamic-sampling filter decides each group; on slime 540ea91 and on its main branch. slime isn't on PyPI: its +# checkout is installed with its requirements over the dev environment, torch the CPU build, and SGLang's and +# Megatron-LM's checkouts go on PYTHONPATH, at the SGLang tag and the Megatron-LM commit slime's Dockerfile builds on, +# with the transformers that SGLang tag pins. 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/slime/**" + - "tst/unit/training/**" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/slime.yml" + schedule: + - cron: "41 8 * * 1" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: slime-${{ 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: + generate: + name: generate (slime ${{ matrix.slime }}) + runs-on: ubuntu-latest + timeout-minutes: 40 + strategy: + fail-fast: false + matrix: + slime: ["540ea919e45b8b929198abce008a103be9cfc779", "main"] + steps: + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + - name: Install from uv.lock with the slime extra + run: | + python -m pip install --quiet -r .github/ci-requirements.txt + uv sync --locked --extra dev --extra slime --python "$(command -v python)" + - name: Check out slime, SGLang and Megatron-LM + env: + SLIME_REF: ${{ matrix.slime }} + run: | + checkout() { git init --quiet "$RUNNER_TEMP/$1" && git -C "$RUNNER_TEMP/$1" fetch --quiet --depth 1 "$2" "$3" \ + && git -C "$RUNNER_TEMP/$1" checkout --quiet FETCH_HEAD && echo "$1 $3 $(git -C "$RUNNER_TEMP/$1" rev-parse HEAD)"; } + checkout slime https://github.com/THUDM/slime "$SLIME_REF" + sglang=$(grep -oP '^ARG SGLANG_IMAGE_TAG=\K\S+' "$RUNNER_TEMP/slime/docker/Dockerfile" | sed -E 's/-cu[0-9]+$//') + megatron=$(grep -oP '^ARG MEGATRON_COMMIT=\K\S+' "$RUNNER_TEMP/slime/docker/Dockerfile") + checkout sglang https://github.com/sgl-project/sglang "$sglang" + checkout Megatron-LM https://github.com/NVIDIA/Megatron-LM "$megatron" + echo "PYTHONPATH=$RUNNER_TEMP/sglang/python:$RUNNER_TEMP/Megatron-LM" >> "$GITHUB_ENV" + - name: Install slime, its requirements and CPU torch + # mcp stays below 2, as the framework pins it; the rest beyond slime's requirements are what SGLang's tool-call + # parser and slime's argument parser import on CPU. + run: | + echo 'mcp>=1.25,<2' > "$RUNNER_TEMP/constraints.txt" + transformers=$(grep -oP '"\Ktransformers==[^"]+' "$RUNNER_TEMP/sglang/python/pyproject.toml" | head -1) + uv pip install --python .venv/bin/python -e "$RUNNER_TEMP/slime" --no-deps + uv pip install --python .venv/bin/python -r "$RUNNER_TEMP/slime/requirements.txt" torch "$transformers" \ + pybase64 triton msgspec pyzmq IPython partial-json-parser xgrammar sentencepiece \ + -c "$RUNNER_TEMP/constraints.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('torch', version('torch'), 'transformers', version('transformers'), 'mcp', version('mcp'))" + - name: slime tier + run: .venv/bin/python -m pytest tst/slime -rA --timeout=900 --junitxml=junit-slime.xml + - name: No skips + if: always() + run: .venv/bin/python .github/scripts/check_skip_policy.py junit-slime.xml diff --git a/AGENTS.md b/AGENTS.md index d22df7f7b..2507d8efe 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -49,6 +49,7 @@ or one of them only where it is used. | `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 miles-test` | `tst/miles`: rows played through Miles' own rollout (its agentic generate path, session server, router and mock SGLang server) with `agent_env.training.miles`' agent function, reward hook and filter, with Qwen3's tokenizer from the Hub; needs checkouts of Miles, SGLang and Megatron-LM on `PYTHONPATH`, Miles' requirements and CPU torch, which the `dev` extra leaves out | +| `make slime-test` | `tst/slime`: rows played through `agent_env.training.slime`'s generate function on slime's own `OpenAIAdapter`, with slime's own arguments and a stand-in for SGLang's `/generate` using Qwen3's tokenizer from the Hub, and its groups decided by slime's dynamic-sampling filter; needs a slime checkout and its requirements, SGLang's and Megatron-LM's checkouts on `PYTHONPATH` and CPU torch, which the `dev` extra leaves out | | `make clean-install-test` | the `clean-install` CI gate: both distributions built as the release builds them, installed into a fresh venv from public PyPI with an allowlisted environment, and `agent-env run hello` run twice by name and checked through the local store; needs Python 3.11, uv and git | | `npm ci && npm run typecheck && npm run lint && npm run test:smoke` | in `src/agent_env/explorer/ui`: the explorer UI checks the `ui` job runs (it also runs `npm run build:static`); needs Node 22 | | `pytest tst/unit/env/store_test.py -v` | one file; add `--log-cli-level=DEBUG` for debug logs | @@ -124,7 +125,7 @@ and deserialized through a registry. | `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `smol_vm` (the `smol` extra), `tensorlake` (the `tensorlake` extra), `tenki` (the `tenki` extra), `vercel` (the `vercel` extra), `boxd` (the `boxd` extra), `coreweave` (the `coreweave` extra), `daytona` (the `daytona` extra), `boat` (the `boat` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store), `EnvironmentGatewayProvider` (`gateway`, an env's default), which renders a docker-compose for the gateway and its MCP servers inside the sandbox, and `EnvironmentServerProvider` (`server`), one MCP server in its own container with no gateway in front and no state store; `providers/env_state/` holds env-state providers (`local_postgres` built in). | | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `openenv/` | `build_app(bundle)`: the [OpenEnv](https://github.com/openenv/OpenEnv) server for a bundle's tasks, each session playing one episode through the Episode API, with the envs' MCP tools, an added `end_turn` tool and the task's verifier score as the reward. It prepares bundles as `agent-env run` does (`prepare_bundle(played=True)`, without the agent the client plays) and keeps namespace routing on for the app's life. `openenv` comes with the `openenv` extra, a floor with no cap; this package and `training/` are the only ones that import it. Each task's description in the Task API names its `toolset` (the played agent's envs, role, and whether it gets `bash`). With `token`, every request and WebSocket but `/health` must carry `Authorization: Bearer `, for a server others reach. | -| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, `tasks_and_tools` (what the row writers read: the tasks, the split names checked, each toolset's tools) and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and raising `RolloutFailed` for one that fails every time, which verl's V1 trainer (the default) leaves out of the loss, the group's baseline and the eval reward; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing, and warns once a tool of an answer verl cuts to `max_tool_response_length`. The writer reads a server that asks for a token with the one `--token-var` names, which the loop reads in each worker, and with `--tokenizer` counts each prompt with its tools as verl renders it (`--chat-template-kwargs`), refusing rows over `--max-prompt-length` before writing. `skyrl/` (`skyrl` extra, installed into SkyRL's own environment): `python -m agent_env.training.skyrl URL OUT --tokenizer MODEL` writes a server's tasks as SkyRL's `.parquet` rows (`env_class: agentenv`, the episode's split, index, task id and tools, the server's settings), each with its tools rendered into its system prompt by the model's chat template, checked to render as `tools=` would, since SkyRL's generator passes the template no tools, and refuses rows over `--max-prompt-length` counted as SkyRL's filter and generator count them; `calls.parse` reads a reply's tool calls in the `hermes` or `qwen3_coder` format, typed by the row's schema; `env.AgentEnvTextEnv` (registered as `agentenv`) is skyrl-gym's `BaseTextEnv` on the row's episode, run on the event loop the sessions share (`_openenv.on_session_loop`), its steps running the reply's calls and a reply without one ending the open turn, raising `EpisodeFailed` when the env fails it, with scores, tool calls and unparsed calls as its metrics; `generator.AgentEnvGenerator` is SkyRL's `SkyRLGymGenerator` playing a failed rollout again on a new episode up to `attempts` times within `reset_timeout`, and masking the group of one that fails every attempt in training (an eval scores it 0); `python -m agent_env.training.skyrl.train` (or `trainer.override_entrypoint=agent_env.training.skyrl.train`) runs SkyRL's trainer with both, its fully async one with `trainer.fully_async.enabled`, passing the rows' token variables to the trainer's Ray task. `miles/` (`miles` extra, installed into Miles' own environment): `python -m agent_env.training.miles URL OUT` writes a server's tasks as Miles' prompt data, `.jsonl` rows of chat `messages` and `metadata` (the episode's split, index, task id and tools, and the server it plays on with its deadlines and turn cap); `agent.run` is Miles' agent function (`--custom-agent-function-path`), which plays the row's episode, the model chatting through the sample's session URL with the row's tools, answers a call whose arguments aren't a JSON object as the model's error, and returns the reward, or flags an episode the env or the session failed, whose model missed the row's `model_timeout` or that ended without a reward, which isn't played again since the session holds its turns; `agent.reward` is the reward hook (`--custom-rm-path`) and `agent.keep_scored_groups` the dynamic-sampling filter that leaves out the group of a flagged episode, of a sample Miles aborted, or of one whose agent function never returned. It imports nothing of Miles. | +| `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. `_chat.py` is what the Miles and slime helpers share: `write_rows` and its `command` write a server's tasks as `.jsonl` rows of chat `messages` and `metadata` (the episode's split, index, task id and tools, and the server it plays on with its deadlines and turn cap); `play` plays a row's episode, the model chatting at a trainer's OpenAI chat endpoint that records its turns, with the row's tools, answers a call whose arguments aren't a JSON object as the model's error, and returns the reward, or flags an episode the env or the session failed, whose model missed the row's `model_timeout` or that ended without a reward, which isn't played again since the session holds its turns; `reward` reads a sample's reward from its metadata, and `keep_scored_groups` leaves out the group of a flagged episode, of a sample the trainer aborted, or of one whose episode never returned. `miles/` (`miles` extra, installed into Miles' own environment): `python -m agent_env.training.miles URL OUT` writes the rows as Miles' prompt data; `agent.run` (`play`) is Miles' agent function (`--custom-agent-function-path`), chatting through the sample's session URL on Miles' session server; `agent.reward` is the reward hook (`--custom-rm-path`) and `agent.keep_scored_groups` the dynamic-sampling filter. It imports nothing of Miles. `slime/` (`slime` extra, installed into slime's own environment): `python -m agent_env.training.slime URL OUT` writes the rows as slime's prompt data; `generate.generate` is slime's generate function (`--custom-generate-function-path`), run on its rollout's event loop, which opens a session for the sample on slime's `OpenAIAdapter` (`slime.agent`, one per rollout process, served in-process and reading tool calls with `--sglang-tool-call-parser`, which it requires), plays the row with `play` through it, and returns the session's turns as samples carrying the reward, one per turn the adapter couldn't join to the turn before, or one aborted sample when no turn was recorded; `generate.keep_scored_groups` is the dynamic-sampling filter. | | `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. | @@ -352,6 +353,10 @@ installed over the dev environment on Python 3.12, on the same triggers; it allo SGLang and Megatron-LM at the commits a release's `release-lock.json` names or at the branches Miles' main tests against, and Miles' requirements 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/slime.yml` runs `tst/slime` on slime 540ea91 and its main branch (v0.4.0's adapter predates the +`rollout_args` and per-session chat-template settings the helper passes): the slime checkout installed with its requirements, SGLang's (at the tag slime's Dockerfile builds on, with +the transformers that tag pins) and Megatron-LM's (at the Dockerfile's commit) checkouts on `PYTHONPATH`, and CPU torch, +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 54fef1e3a..800fe086f 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test trl-test nemo-gym-test verl-test skyrl-test miles-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 miles-test slime-test clean-install-test test clean VENV := .venv VENV_BIN := $(abspath $(VENV))/bin @@ -91,6 +91,12 @@ skyrl-test: ## Play rows through the SkyRL env and generator on SkyRL itself (ne miles-test: ## Play rows through Miles' rollout and session server (needs Miles, SGLang and Megatron-LM on PYTHONPATH, CPU torch, and the Hub) $(PYTHON) -m pytest tst/miles --timeout=900 $(PYTEST_ARGS) +# The slime tier plays rows through the slime helper's generate function on slime's own OpenAI adapter, with slime's own +# arguments and Qwen3's tokenizer from the Hub. slime isn't on PyPI: install a checkout of it, its requirements and CPU +# torch, and put SGLang's and Megatron-LM's checkouts on PYTHONPATH first, as .github/workflows/slime.yml does. +slime-test: ## Play rows through the slime helper on slime's OpenAI adapter (needs slime, SGLang and Megatron-LM, CPU torch, and the Hub) + $(PYTHON) -m pytest tst/slime --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 12839d696..52c7592ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,6 +97,11 @@ skyrl = [ miles = [ "agentenv-framework[openenv]", ] +# agent_env.training.slime: slime's trainers on a bundle served as an OpenEnv env; installed into slime's own environment, +# whose slime, SGLang and torch are the trainer's to pin. It needs nothing beyond the openenv extra. +slime = [ + "agentenv-framework[openenv]", +] # The AWS store backends (S3, Secrets Manager, DynamoDB, ECR), selected by impl pointer in config.toml. aws = [ "boto3>=1.35.0", diff --git a/src/agent_env/training/__init__.py b/src/agent_env/training/__init__.py index 0571146a0..1cc2cabe4 100644 --- a/src/agent_env/training/__init__.py +++ b/src/agent_env/training/__init__.py @@ -1,4 +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, ``agent_env.training.verl`` for verl, ``agent_env.training.skyrl`` for SkyRL, and -``agent_env.training.miles`` for Miles.""" +environments, NeMo RL among them, ``agent_env.training.verl`` for verl, ``agent_env.training.skyrl`` for SkyRL, +``agent_env.training.miles`` for Miles, and ``agent_env.training.slime`` for slime.""" diff --git a/src/agent_env/training/_chat.py b/src/agent_env/training/_chat.py new file mode 100644 index 000000000..0d71d554c --- /dev/null +++ b/src/agent_env/training/_chat.py @@ -0,0 +1,164 @@ +"""What the Miles and slime helpers share: a server's tasks as chat rows, and a row's task played as one episode by a +model chatting at a trainer's OpenAI chat endpoint, a session of the trainer's that records the model's turns for +training. A row holds the task's system prompt and first prompt as ``messages``, and in ``metadata`` its episode, the +tools its toolset offers and the server it plays on with its settings, so the trainer's rollout processes need no +settings of their own.""" + +from __future__ import annotations + +import json +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import click +import httpx + +from agent_env.training._openenv import Session, Shared, function_tool, key, messages, tasks_and_tools, token_in + + +@click.command() +@click.argument("url") +@click.argument("out", type=click.Path(file_okay=False)) +@click.option("--tool-timeout", type=float, default=60.0, show_default=True, + help="One tool call's deadline in seconds; the server ends a call that misses it, and the episode.") +@click.option("--reset-timeout", type=float, default=600.0, show_default=True, + help="How long an episode may take to deploy, and each turn to be ended, in seconds; also how long a " + "rollout waits for room on the server.") +@click.option("--model-timeout", type=float, default=600.0, show_default=True, + help="How long a rollout waits for each of the model's replies, in seconds; an episode whose " + "model misses it fails, and its group is left out.") +@click.option("--max-turns", type=int, default=4, show_default=True, + help="The model turns a rollout may take; the episode is ended after the last.") +@click.option("--token-var", help="The environment variable that holds the token the server asks for, read now and in " + "each of the trainer's rollout processes.") +def command(url: str, out: str, tool_timeout: float, reset_timeout: float, model_timeout: float, max_turns: int, + token_var: str | None) -> None: + """Writes the tasks of the OpenEnv server at URL as OUT/.jsonl, a row per task, for the trainer's + --prompt-data with --input-key messages and --metadata-key metadata.""" + folder = write_rows(url, out, tool_timeout=tool_timeout, reset_timeout=reset_timeout, model_timeout=model_timeout, + max_turns=max_turns, token_var=token_var) + click.echo(f"wrote {', '.join(sorted(path.name for path in folder.iterdir()))} to {folder}") + + + +def write_rows(url: str, out: str | Path, *, tool_timeout: float = 60.0, reset_timeout: float = 600.0, + model_timeout: float = 600.0, max_turns: int = 4, token_var: str | None = None) -> Path: + """The tasks of the OpenEnv server at ``url`` as ``out/.jsonl``, a row per task of each split: its system + prompt and first prompt as ``messages``, and as ``metadata`` its split, index and task id, the tools its toolset + offers, and where and how ``play`` plays it: the server, its deadlines and the model's in seconds, the model turns a + rollout may take, and the environment variable holding the token a server given one asks for, read in each of the + trainer's rollout processes.""" + url = url.rstrip("/") + _, 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(tool) for tool in found] for each, found in probed.items()} + server = {"url": url, "tool_timeout": tool_timeout, "reset_timeout": reset_timeout, "model_timeout": model_timeout, + "max_turns": max_turns, **({"token_var": token_var} if token_var else {})} + out = Path(out) + out.mkdir(parents=True, exist_ok=True) + for split, described in tasks.items(): + (out / f"{split}.jsonl").write_text("".join(json.dumps({ + "messages": messages(task["system_prompt"], task["prompt"]), + "metadata": {"split": split, "index": task["index"], "task_id": task["id"], "tools": tools[key(task)], + "agentenv": server}, + }) + "\n" for task in described)) + return out + + +class _Episode(Session): + """One sample's episode on the server its row names.""" + + offered_by = "its row lists" + + def __init__(self, shared: Shared): + super().__init__() + self._shared = shared + + +async def play(base_url: str, prompt: list[dict[str, Any]], request_kwargs: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, **kwargs: Any) -> dict[str, Any]: + """Plays the row's task as one episode: the model chats at ``base_url``, a session of the trainer's that records its + turns, with the row's tools; each call is a call on the episode, and the model's last reply, or the turn cap, ends + the turn. A call whose arguments aren't a JSON object is answered with why, as the model's error, and the episode + goes on. Returns the sample's metadata: the episode's reward and scores, or, when the env or the session failed it, + the model missed the row's ``model_timeout`` or the episode ended without a reward, a reward of 0 and why, which + ``keep_scored_groups`` leaves out of training; and the calls made and refused. A failed episode isn't played again: + the session already holds its turns.""" + row = metadata or {} + server, tools = row["agentenv"], row["tools"] + episode = _Episode(Shared(server["url"], frozenset(tool["function"]["name"] for tool in tools), + server["tool_timeout"], server["reset_timeout"], patience=server["reset_timeout"], + waits_for_server=True, token=token_in(server.get("token_var")))) + counts = {"agentenv_tool_calls": 0, "agentenv_invalid_calls": 0} + try: + await episode._start(row["split"], int(row["index"]), row["task_id"]) + reply = await _converse(episode, base_url, list(prompt), tools, request_kwargs or {}, server, counts) + rollout = episode._rollout + if not (rollout.done or rollout.failure): + await episode._finish(rollout, reply) + if rollout.failure is None and rollout.reward is None: + rollout.failure = "the episode ended without a reward" + if rollout.failure is not None: + return {"reward": 0.0, "agentenv_failure": rollout.failure, **counts} + return {"reward": float(rollout.reward), "agentenv_scores": rollout.scores or {}, **counts} + except Exception as e: + return {"reward": 0.0, "agentenv_failure": f"{type(e).__name__}: {e}", **counts} + finally: + await episode._close() + + +async def reward(args: Any, samples: Any, **kwargs: Any) -> float | list[float]: + """Each sample's episode reward, which ``play`` put in its metadata.""" + if isinstance(samples, list): + return [float(sample.metadata.get("reward", 0.0)) for sample in samples] + return float(samples.metadata.get("reward", 0.0)) + + +def keep_scored_groups(args: Any, samples: list[Any], **kwargs: Any) -> bool: + """Whether to train on a group: not when an episode of it failed in the env, its episode never returned (a child + process that died mid-episode leaves the session's turns with no reward), or the trainer aborted a sample of it, + none of them the policy's doing.""" + return not any(sample.metadata.get("agentenv_failure") or "reward" not in sample.metadata + or getattr(sample.status, "name", None) == "ABORTED" for sample in _samples(samples)) + + +async def _converse(episode: _Episode, base_url: str, messages: list[dict[str, Any]], tools: list[dict[str, Any]], + request_kwargs: dict[str, Any], server: dict[str, Any], counts: dict[str, int]) -> str: + """The model's turns, until it replies without a call, the episode ends, or the server's ``max_turns`` pass; its + last reply.""" + async with httpx.AsyncClient(timeout=server["model_timeout"]) as client: + for _ in range(server["max_turns"]): + response = await client.post(f"{base_url}/v1/chat/completions", json={ + "model": "default", "messages": messages, "tools": tools, **request_kwargs}) + response.raise_for_status() + message = response.json()["choices"][0]["message"] + # Verbatim: the session server matches each message a request replays against the one it recorded. + messages.append(message) + if not (calls := message.get("tool_calls")): + return message.get("content") or "" + for call in calls: + messages.append({"role": "tool", "tool_call_id": call.get("id", ""), + "content": await _answer(episode, call, counts)}) + if episode._rollout.done or episode._rollout.failure: + return "" + return "" + + +async def _answer(episode: _Episode, call: dict[str, Any], counts: dict[str, int]) -> str: + name, arguments = call["function"]["name"], call["function"].get("arguments") + try: + arguments = json.loads(arguments) if isinstance(arguments, str) else arguments + except json.JSONDecodeError as e: + counts["agentenv_invalid_calls"] += 1 + return f"{name} wasn't called: its arguments aren't JSON ({e})" + if not isinstance(arguments, dict): + counts["agentenv_invalid_calls"] += 1 + return f"{name} wasn't called: its arguments must be a JSON object, not {type(arguments).__name__}" + counts["agentenv_tool_calls"] += 1 + return await episode._step(name, arguments) + + +def _samples(group: list[Any]) -> Iterator[Any]: + for sample in group: + yield from sample if isinstance(sample, list) else [sample] diff --git a/src/agent_env/training/miles/__main__.py b/src/agent_env/training/miles/__main__.py index d63acd182..f4f7e1a0c 100644 --- a/src/agent_env/training/miles/__main__.py +++ b/src/agent_env/training/miles/__main__.py @@ -1,33 +1,6 @@ """``python -m agent_env.training.miles URL OUT``: writes a bundle's tasks as Miles' prompt data.""" -import click - -from agent_env.training.miles.prepare import write_rows - - -@click.command() -@click.argument("url") -@click.argument("out", type=click.Path(file_okay=False)) -@click.option("--tool-timeout", type=float, default=60.0, show_default=True, - help="One tool call's deadline in seconds; the server ends a call that misses it, and the episode.") -@click.option("--reset-timeout", type=float, default=600.0, show_default=True, - help="How long an episode may take to deploy, and each turn to be ended, in seconds; also how long a " - "rollout waits for room on the server.") -@click.option("--model-timeout", type=float, default=600.0, show_default=True, - help="How long the agent function waits for each of the model's replies, in seconds; an episode whose " - "model misses it fails, and its group is left out.") -@click.option("--max-turns", type=int, default=4, show_default=True, - help="The model turns a rollout may take; the episode is ended after the last.") -@click.option("--token-var", help="The environment variable that holds the token the server asks for, read now and in " - "each of Miles' rollout processes.") -def main(url: str, out: str, tool_timeout: float, reset_timeout: float, model_timeout: float, max_turns: int, - token_var: str | None) -> None: - """Writes the tasks of the OpenEnv server at URL as OUT/.jsonl, a row per task, for Miles' --prompt-data - with --input-key messages and --metadata-key metadata.""" - folder = write_rows(url, out, tool_timeout=tool_timeout, reset_timeout=reset_timeout, model_timeout=model_timeout, - max_turns=max_turns, token_var=token_var) - click.echo(f"wrote {', '.join(sorted(path.name for path in folder.iterdir()))} to {folder}") - +from agent_env.training._chat import command as main if __name__ == "__main__": main() diff --git a/src/agent_env/training/miles/agent.py b/src/agent_env/training/miles/agent.py index 61beba5fd..6b10f0236 100644 --- a/src/agent_env/training/miles/agent.py +++ b/src/agent_env/training/miles/agent.py @@ -1,113 +1,10 @@ """What Miles calls for each sample of the rows ``python -m agent_env.training.miles`` writes: the agent function -(``--custom-agent-function-path agent_env.training.miles.agent.run``), the reward hook -(``--custom-rm-path agent_env.training.miles.agent.reward``) and the dynamic-sampling filter +(``--custom-agent-function-path agent_env.training.miles.agent.run``), whose session is Miles' session server, the +reward hook (``--custom-rm-path agent_env.training.miles.agent.reward``) and the dynamic-sampling filter (``--dynamic-sampling-filter-path agent_env.training.miles.agent.keep_scored_groups``). Imports nothing of Miles: it calls them by their arguments.""" -from __future__ import annotations +from agent_env.training._chat import keep_scored_groups, reward +from agent_env.training._chat import play as run -import json -from collections.abc import Iterator -from typing import Any - -import httpx - -from agent_env.training._openenv import Session, Shared, token_in - - -class _Episode(Session): - """One sample's episode on the server its row names.""" - - offered_by = "its row lists" - - def __init__(self, shared: Shared): - super().__init__() - self._shared = shared - - -async def run(base_url: str, prompt: list[dict[str, Any]], request_kwargs: dict[str, Any] | None = None, - metadata: dict[str, Any] | None = None, **kwargs: Any) -> dict[str, Any]: - """Plays the row's task as one episode: the model chats at ``base_url``, the sample's session, which Miles records, - with the row's tools; each call is a call on the episode, and the model's last reply, or the turn cap, ends the - turn. A call whose arguments aren't a JSON object is answered with why, as the model's error, and the episode goes - on. Returns what Miles merges into the sample's metadata: the episode's reward and scores, or, when the env or the - session failed it, the model missed the row's ``model_timeout`` or the episode ended without a reward, a reward of 0 - and why, which ``keep_scored_groups`` leaves out of training; and the calls made and refused. A failed episode - isn't played again: the session already holds its turns.""" - row = metadata or {} - server, tools = row["agentenv"], row["tools"] - episode = _Episode(Shared(server["url"], frozenset(tool["function"]["name"] for tool in tools), - server["tool_timeout"], server["reset_timeout"], patience=server["reset_timeout"], - waits_for_server=True, token=token_in(server.get("token_var")))) - counts = {"agentenv_tool_calls": 0, "agentenv_invalid_calls": 0} - try: - await episode._start(row["split"], int(row["index"]), row["task_id"]) - reply = await _converse(episode, base_url, list(prompt), tools, request_kwargs or {}, server, counts) - rollout = episode._rollout - if not (rollout.done or rollout.failure): - await episode._finish(rollout, reply) - if rollout.failure is None and rollout.reward is None: - rollout.failure = "the episode ended without a reward" - if rollout.failure is not None: - return {"reward": 0.0, "agentenv_failure": rollout.failure, **counts} - return {"reward": float(rollout.reward), "agentenv_scores": rollout.scores or {}, **counts} - except Exception as e: - return {"reward": 0.0, "agentenv_failure": f"{type(e).__name__}: {e}", **counts} - finally: - await episode._close() - - -async def reward(args: Any, samples: Any, **kwargs: Any) -> float | list[float]: - """Each sample's episode reward, which ``run`` put in its metadata.""" - if isinstance(samples, list): - return [float(sample.metadata.get("reward", 0.0)) for sample in samples] - return float(samples.metadata.get("reward", 0.0)) - - -def keep_scored_groups(args: Any, samples: list[Any], **kwargs: Any) -> bool: - """Whether to train on a group: not when an episode of it failed in the env, its agent function never returned - (a child process that died mid-episode leaves the session's turns with no reward), or Miles aborted a sample of - it, none of them the policy's doing.""" - return not any(sample.metadata.get("agentenv_failure") or "reward" not in sample.metadata - or getattr(sample.status, "name", None) == "ABORTED" for sample in _samples(samples)) - - -async def _converse(episode: _Episode, base_url: str, messages: list[dict[str, Any]], tools: list[dict[str, Any]], - request_kwargs: dict[str, Any], server: dict[str, Any], counts: dict[str, int]) -> str: - """The model's turns, until it replies without a call, the episode ends, or the server's ``max_turns`` pass; its - last reply.""" - async with httpx.AsyncClient(timeout=server["model_timeout"]) as client: - for _ in range(server["max_turns"]): - response = await client.post(f"{base_url}/v1/chat/completions", json={ - "model": "default", "messages": messages, "tools": tools, **request_kwargs}) - response.raise_for_status() - message = response.json()["choices"][0]["message"] - # Verbatim: the session server matches each message a request replays against the one it recorded. - messages.append(message) - if not (calls := message.get("tool_calls")): - return message.get("content") or "" - for call in calls: - messages.append({"role": "tool", "tool_call_id": call.get("id", ""), - "content": await _answer(episode, call, counts)}) - if episode._rollout.done or episode._rollout.failure: - return "" - return "" - - -async def _answer(episode: _Episode, call: dict[str, Any], counts: dict[str, int]) -> str: - name, arguments = call["function"]["name"], call["function"].get("arguments") - try: - arguments = json.loads(arguments) if isinstance(arguments, str) else arguments - except json.JSONDecodeError as e: - counts["agentenv_invalid_calls"] += 1 - return f"{name} wasn't called: its arguments aren't JSON ({e})" - if not isinstance(arguments, dict): - counts["agentenv_invalid_calls"] += 1 - return f"{name} wasn't called: its arguments must be a JSON object, not {type(arguments).__name__}" - counts["agentenv_tool_calls"] += 1 - return await episode._step(name, arguments) - - -def _samples(group: list[Any]) -> Iterator[Any]: - for sample in group: - yield from sample if isinstance(sample, list) else [sample] +__all__ = ["keep_scored_groups", "reward", "run"] diff --git a/src/agent_env/training/miles/prepare.py b/src/agent_env/training/miles/prepare.py index fa41809db..83a66552c 100644 --- a/src/agent_env/training/miles/prepare.py +++ b/src/agent_env/training/miles/prepare.py @@ -1,33 +1,6 @@ """The server's tasks as Miles' prompt data: a JSON line per task of each split, its chat ``messages`` and the ``metadata`` the agent function plays it from.""" -from __future__ import annotations +from agent_env.training._chat import write_rows -import json -from pathlib import Path - -from agent_env.training._openenv import function_tool, key, messages, tasks_and_tools, token_in - - -def write_rows(url: str, out: str | Path, *, tool_timeout: float = 60.0, reset_timeout: float = 600.0, - model_timeout: float = 600.0, max_turns: int = 4, token_var: str | None = None) -> Path: - """The tasks of the OpenEnv server at ``url`` as ``out/.jsonl``, a row per task of each split: its system - prompt and first prompt as ``messages``, and as ``metadata`` its split, index and task id, the tools its toolset - offers, and where and how the agent function plays it: the server, its deadlines and the model's in seconds, the - model turns a rollout may take, and the environment variable holding the token a server given one asks for, read in each of - Miles' rollout processes.""" - url = url.rstrip("/") - _, 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(tool) for tool in found] for each, found in probed.items()} - server = {"url": url, "tool_timeout": tool_timeout, "reset_timeout": reset_timeout, "model_timeout": model_timeout, - "max_turns": max_turns, **({"token_var": token_var} if token_var else {})} - out = Path(out) - out.mkdir(parents=True, exist_ok=True) - for split, described in tasks.items(): - (out / f"{split}.jsonl").write_text("".join(json.dumps({ - "messages": messages(task["system_prompt"], task["prompt"]), - "metadata": {"split": split, "index": task["index"], "task_id": task["id"], "tools": tools[key(task)], - "agentenv": server}, - }) + "\n" for task in described)) - return out +__all__ = ["write_rows"] diff --git a/src/agent_env/training/slime/__init__.py b/src/agent_env/training/slime/__init__.py new file mode 100644 index 000000000..881725444 --- /dev/null +++ b/src/agent_env/training/slime/__init__.py @@ -0,0 +1,5 @@ +"""slime's trainers (https://github.com/THUDM/slime) on a bundle served as an OpenEnv env (``agent_env.openenv.build_app``, +or the package ``openenv import`` writes): ``python -m agent_env.training.slime URL OUT`` writes the server's tasks as +slime's prompt data, and ``agent_env.training.slime.generate`` holds the generate function slime calls for each sample, +which plays the row's task as one episode on the server through a session on slime's OpenAI adapter, with the +dynamic-sampling filter beside it.""" diff --git a/src/agent_env/training/slime/__main__.py b/src/agent_env/training/slime/__main__.py new file mode 100644 index 000000000..754038416 --- /dev/null +++ b/src/agent_env/training/slime/__main__.py @@ -0,0 +1,6 @@ +"""``python -m agent_env.training.slime URL OUT``: writes a bundle's tasks as slime's prompt data.""" + +from agent_env.training._chat import command as main + +if __name__ == "__main__": + main() diff --git a/src/agent_env/training/slime/generate.py b/src/agent_env/training/slime/generate.py new file mode 100644 index 000000000..f3772eff4 --- /dev/null +++ b/src/agent_env/training/slime/generate.py @@ -0,0 +1,59 @@ +"""What slime calls for each sample of the rows ``python -m agent_env.training.slime`` writes: the generate function +(``--custom-generate-function-path agent_env.training.slime.generate.generate``) and the dynamic-sampling filter +(``--dynamic-sampling-filter-path agent_env.training.slime.generate.keep_scored_groups``). It runs in slime's own +environment, on its rollout's event loop.""" + +from __future__ import annotations + +import uuid +from typing import Any + +from slime.agent.adapters import OpenAIAdapter +from slime.agent.aiohttp_threaded import run_app_in_thread +from slime.utils.misc import SingletonMeta +from slime.utils.processing_utils import load_tokenizer +from slime.utils.types import Sample + +from agent_env.training._chat import keep_scored_groups, play + +__all__ = ["generate", "keep_scored_groups"] + + +async def generate(args: Any, sample: Sample, sampling_params: dict[str, Any], evaluation: bool = False) -> list[Sample]: + """Plays the sample's row as one episode, the model chatting through a session of its own on slime's OpenAI + adapter, which records its turns. Returns the session's turns as training samples carrying the episode's reward and + ``play``'s metadata: one sample, or one per turn the adapter couldn't join to the turn before. An episode with no + recorded turn comes back as a single aborted sample, which ``keep_scored_groups`` leaves out with its group.""" + adapter = _Adapter(args) + sid = uuid.uuid4().hex + adapter.adapter.open_session( + sid, sampling_defaults=sampling_params, max_context_tokens=args.rollout_max_context_len or 0, + apply_chat_template_kwargs={**args.apply_chat_template_kwargs, **(sample.apply_chat_template_kwargs or {})}) + played = await play(adapter.url, sample.prompt, {"user": sid}, sample.metadata) + samples = await adapter.adapter.finish_session(sid, base_sample=sample, reward=played["reward"], + extra_metadata=played) + if samples: + return samples + sample.tokens, sample.response, sample.response_length = [0, 0], "", 1 + sample.loss_mask, sample.rollout_log_probs = [0], [0.0] + sample.reward, sample.remove_sample, sample.status = played["reward"], True, Sample.Status.ABORTED + sample.metadata = {**sample.metadata, **played} + return [sample] + + +class _Adapter(metaclass=SingletonMeta): + """The rollout process's OpenAI adapter, in front of its SGLang router, reading the model's tool calls with + ``--sglang-tool-call-parser``.""" + + def __init__(self, args: Any): + if not args.sglang_tool_call_parser: + raise ValueError("slime's OpenAI adapter reads the model's tool calls with --sglang-tool-call-parser; pass " + "the parser for the model's chat template, qwen25 for Qwen2.5 and Qwen3") + self.adapter = OpenAIAdapter( + tokenizer=load_tokenizer(args.hf_checkpoint, trust_remote_code=True), + sglang_url=f"http://{args.sglang_router_ip}:{args.sglang_router_port}", + tool_parser=args.sglang_tool_call_parser, reasoning_parser=args.sglang_reasoning_parser, + rollout_args=args) + served = run_app_in_thread(self.adapter.app, host="127.0.0.1", thread_name="agentenv-adapter", + runner_kwargs={"handler_cancellation": True}) + self.url = f"http://127.0.0.1:{served.port}" diff --git a/tst/slime/__init__.py b/tst/slime/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tst/slime/conftest.py b/tst/slime/conftest.py new file mode 100644 index 000000000..cd046bb15 --- /dev/null +++ b/tst/slime/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 slime's +own OpenAI adapter.""" + +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/slime/generate_test.py b/tst/slime/generate_test.py new file mode 100644 index 000000000..5b01c8868 --- /dev/null +++ b/tst/slime/generate_test.py @@ -0,0 +1,149 @@ +"""Rows ``write_rows`` wrote, played by ``agent_env.training.slime.generate`` as slime's rollout calls it: with slime's +own arguments, through slime's own OpenAI adapter in front of a stand-in for SGLang's ``/generate``, which plays the +model from a script with Qwen3's tokenizer; slime's dynamic-sampling filter decides each group.""" + +import asyncio +import json +from unittest.mock import patch + +import pytest +from starlette.applications import Starlette +from starlette.responses import JSONResponse +from starlette.routing import Route + +from agent_env.training._chat import write_rows +from agent_env.training._openenv import Server +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 SLIME, missing_capability_reason + +slime_generate = pytest.importorskip("agent_env.training.slime.generate", reason=missing_capability_reason(SLIME), + exc_type=ImportError) +arguments = pytest.importorskip("slime.utils.arguments", reason=missing_capability_reason(SLIME)) +filters = pytest.importorskip("slime.rollout.filter_hub.base_types", reason=missing_capability_reason(SLIME)) +misc = pytest.importorskip("slime.utils.misc", reason=missing_capability_reason(SLIME)) +types = pytest.importorskip("slime.utils.types", reason=missing_capability_reason(SLIME)) +transformers = pytest.importorskip("transformers", reason=missing_capability_reason(SLIME)) + +CHECKPOINT = "Qwen/Qwen3-0.6B" +ADDS_A_CRATE = '\n{"name": "add_item", "arguments": {"item": "crate"}}\n<|im_end|>' +FAILS_AFTER_THE_TURN = [ITEMS, ALICE, TURN, {"id": "crane", "type": "episode_test_env", "env_id": "crane", "fail": True}, + REPLIES] +SAMPLING = {"temperature": 1.0, "top_p": 1.0, "max_new_tokens": 256} + + +@pytest.fixture(autouse=True) +def fresh_adapter(monkeypatch): + """Each test's own adapter, in front of its own stand-in SGLang: the module keeps one per process.""" + monkeypatch.delitem(misc.SingletonMeta._instances, slime_generate._Adapter, raising=False) + yield + misc.SingletonMeta._instances.pop(slime_generate._Adapter, None) + + +@pytest.fixture +def model(): + """A stand-in for SGLang's ``/generate``: adds a crate, then says so once the tool has answered; or answers 500 when + ``.broken``.""" + tokenizer = transformers.AutoTokenizer.from_pretrained(CHECKPOINT) + + async def generate(request): + if server.broken: + return JSONResponse({"error": "engine down"}, status_code=500) + prompt = tokenizer.decode((await request.json())["input_ids"]) + ids = tokenizer.encode("Added a crate.<|im_end|>" if "" in prompt else ADDS_A_CRATE, + add_special_tokens=False) + return JSONResponse({"meta_info": {"output_token_logprobs": [[-0.01, each] for each in ids], + "finish_reason": {"type": "stop"}}}) + + server = Server(Starlette(routes=[Route("/generate", generate, methods=["POST"])])) + server.broken, server.tokenizer = False, tokenizer + yield server + server.stop() + + +def test_slime_plays_a_row_through_its_adapter_and_trains_on_the_models_turns_with_the_episodes_reward( + serve, model, tmp_path): + server = serve({"a": PLAYED}) + row = _row(server, tmp_path) + + samples = _generate(row, _args(model, tmp_path)) + server.stop() + + assert {sample.reward for sample in samples} == {1.0} + assert all(sample.metadata["agentenv_scores"] == {"replies": 1.0} and sample.metadata["agentenv_tool_calls"] == 1 + for sample in samples) + trained = "".join(_trained(sample, model.tokenizer) for sample in samples) + assert '"crate"' in trained and "Added a crate." in trained and "" not in trained + + +def test_an_episode_the_env_failed_is_flagged_and_slime_leaves_its_group_out(serve, model, tmp_path): + server = serve({"a": FAILS_AFTER_THE_TURN}) + row = _row(server, tmp_path) + args = _args(model, tmp_path) + + samples = _generate(row, args) + server.stop() + + assert samples and all(sample.metadata["agentenv_failure"] for sample in samples) + assert filters.call_dynamic_filter(slime_generate.keep_scored_groups, args, samples).keep is False + + +def test_an_episode_with_no_recorded_turn_comes_back_as_one_aborted_sample_slime_leaves_out(serve, model, tmp_path): + server = serve({"a": PLAYED}) + row = _row(server, tmp_path) + args = _args(model, tmp_path) + model.broken = True + + [sample] = _generate(row, args) + server.stop() + + assert (sample.status, sample.remove_sample, sample.reward) == (types.Sample.Status.ABORTED, True, 0.0) + assert sample.metadata["agentenv_failure"] + assert filters.call_dynamic_filter(slime_generate.keep_scored_groups, args, [sample]).keep is False + + +def test_the_adapter_needs_slimes_tool_call_parser(serve, model, tmp_path): + server = serve({"a": PLAYED}) + row = _row(server, tmp_path) + args = _args(model, tmp_path) + args.sglang_tool_call_parser = None + + with pytest.raises(ValueError, match="--sglang-tool-call-parser"): + _generate(row, args) + server.stop() + + +def _row(server, tmp_path): + folder = write_rows(server.url, tmp_path / "rows") + return json.loads((folder / "all.jsonl").read_text().splitlines()[0]) + + +def _generate(row, args): + """The samples slime's rollout gets back for ``row``, as it calls the generate function.""" + sample = types.Sample(prompt=row["messages"], metadata=row["metadata"], index=0) + return asyncio.run(slime_generate.generate(args, sample, dict(SAMPLING))) + + +def _trained(sample, tokenizer): + """The text of the tokens slime trains on in ``sample``.""" + response = sample.tokens[len(sample.tokens) - sample.response_length:] + return tokenizer.decode([token for token, trained in zip(response, sample.loss_mask) if trained]) + + +def _args(model, tmp_path): + """slime's own arguments for a rollout of the rows with the helper, in front of the stand-in ``model``; a two-layer + model shape keeps Megatron's checks satisfied.""" + port = int(model.url.rsplit(":", 1)[1]) + argv = [ + "train.py", "--hf-checkpoint", CHECKPOINT, "--rollout-num-gpus", "1", "--rollout-num-gpus-per-engine", "1", + "--num-layers", "2", "--hidden-size", "64", "--num-attention-heads", "4", "--ffn-hidden-size", "128", + "--seq-length", "512", "--max-position-embeddings", "512", "--debug-rollout-only", + "--prompt-data", str(tmp_path / "rows" / "all.jsonl"), "--input-key", "messages", "--metadata-key", "metadata", + "--rollout-batch-size", "1", "--n-samples-per-prompt", "1", "--global-batch-size", "1", "--num-rollout", "1", + "--rollout-max-context-len", "8192", "--apply-chat-template-kwargs", json.dumps({"enable_thinking": False}), + "--sglang-router-ip", "127.0.0.1", "--sglang-router-port", str(port), "--sglang-tool-call-parser", "qwen25", + "--custom-generate-function-path", "agent_env.training.slime.generate.generate", + "--dynamic-sampling-filter-path", "agent_env.training.slime.generate.keep_scored_groups", + ] + with patch("sys.argv", argv): + return arguments.parse_args() diff --git a/tst/unit/test_optional_extras_stay_optional.py b/tst/unit/test_optional_extras_stay_optional.py index 9254ce459..526ff3879 100644 --- a/tst/unit/test_optional_extras_stay_optional.py +++ b/tst/unit/test_optional_extras_stay_optional.py @@ -26,13 +26,14 @@ 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 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: the Sail and Smol providers, the OpenEnv server, NeMo Gym's resources server and the verl, SkyRL, Miles and +# slime helpers, with the chat rows and episode the last two share, need their optional extras (and the verl agent loop +# verl itself, the SkyRL env and generator SkyRL, the slime generate function slime), 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.training.skyrl", - "agent_env.training.harness", "agent_env.training.miles", + "agent_env.training.harness", "agent_env.training.miles", "agent_env.training.slime", "agent_env.training._chat", "agent_env.env.gateway", "agent_env.task_step.task_steps.run_code_runner", ) @@ -174,7 +175,7 @@ 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 and SkyRL's pyarrow comes with trl's datasets. The test without any extra covers them. -_COVERED = {"verl": "trl", "skyrl": "trl", "harness": "openenv", "miles": "openenv"} +_COVERED = {"verl": "trl", "skyrl": "trl", "harness": "openenv", "miles": "openenv", "slime": "openenv"} @pytest.mark.parametrize("extra", sorted(set(_optional_requirements()) - set(_COVERED))) diff --git a/tst/unit/training/miles_test.py b/tst/unit/training/miles_test.py index d5559f6c5..a70aa2e80 100644 --- a/tst/unit/training/miles_test.py +++ b/tst/unit/training/miles_test.py @@ -1,7 +1,7 @@ -"""``agent_env.training.miles`` on the OpenEnv server ``build_app`` makes, served on loopback over the Episode API's fake -steps: the rows ``write_rows`` writes, and the agent function playing a row's task through a stand-in for Miles' -session endpoint, which answers each chat request from a script and records it; with the reward hook and the filter -Miles calls beside it.""" +"""``agent_env.training.miles``, and the chat rows and episode it shares with ``agent_env.training.slime``, on the +OpenEnv server ``build_app`` makes, served on loopback over the Episode API's fake steps: the rows ``write_rows`` writes +and both helpers' commands, and the agent function playing a row's task through a stand-in for a trainer's session +endpoint, which answers each chat request from a script and records it; with the reward hook and the filter beside it.""" import asyncio import json @@ -16,9 +16,10 @@ from agent_env.openenv.environment import END_TURN from agent_env.openenv.reward import _unscored from agent_env.training._openenv import Server -from agent_env.training.miles.__main__ import main +from agent_env.training.miles.__main__ import main as miles_main from agent_env.training.miles.agent import keep_scored_groups, reward, run from agent_env.training.miles.prepare import write_rows +from agent_env.training.slime.__main__ import main as slime_main from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN from tst.unit.training.nemo_gym_test import PLAYED @@ -165,7 +166,8 @@ def test_the_reward_hook_reads_each_samples_reward_from_its_metadata(): assert asyncio.run(reward(None, [one, two])) == [0.25, 0.0] -def test_the_command_writes_the_rows(serve, tmp_path, teardowns): +@pytest.mark.parametrize("main", [miles_main, slime_main], ids=["miles", "slime"]) +def test_the_command_writes_the_rows(serve, tmp_path, teardowns, main): server = serve({"a": PLAYED}) result = CliRunner().invoke(main, [server.url, str(tmp_path / "rows"), "--max-turns", "2"]) diff --git a/tst/util/capabilities.py b/tst/util/capabilities.py index b75bff118..24b1eaaa3 100644 --- a/tst/util/capabilities.py +++ b/tst/util/capabilities.py @@ -57,6 +57,10 @@ #: the dev extra, and only that tier's CI job puts a Miles checkout on PYTHONPATH and installs its requirements. MILES = "miles_installed" +#: The slime tier (``tst/slime``) plays rows through slime's own OpenAI adapter with slime's own arguments; slime, its +#: SGLang and Megatron-LM and CPU torch aren't in the dev extra, and only that tier's CI job installs them. +SLIME = "slime_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 85fe53809..37338756d 100644 --- a/uv.lock +++ b/uv.lock @@ -161,6 +161,9 @@ skyrl = [ { name = "pyarrow", version = "26.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13.14'" }, { name = "transformers" }, ] +slime = [ + { name = "openenv" }, +] smol = [ { name = "smolmachines" }, ] @@ -196,6 +199,7 @@ requires-dist = [ { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'miles'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'nemo-gym'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'skyrl'" }, + { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'slime'" }, { 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" }, @@ -261,7 +265,7 @@ requires-dist = [ { name = "uvicorn", specifier = ">=0.47" }, { name = "vercel-sandbox", marker = "extra == 'vercel'", specifier = ">=0.7.0,<0.8" }, ] -provides-extras = ["openenv", "harness", "trl", "nemo-gym", "verl", "skyrl", "miles", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] +provides-extras = ["openenv", "harness", "trl", "nemo-gym", "verl", "skyrl", "miles", "slime", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] [[package]] name = "agentenv-framework-protocol" From 4f5ef0b64ec95a84a60e36bbf864ddc84f79ce48 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 3 Nov 2026 06:01:05 -0800 Subject: [PATCH 2/4] fix(slime): an eval gives one sample per episode, and slime can cancel a rollout mid-episode Episodes come back as a sample per turn when the adapter can't join a turn to the one before (Qwen3's thinking-off template), and slime's eval reward averages over samples, so it weighted episodes by their turns. In an eval the generate function now drops the session and returns one sample carrying the reward, as slime's own agent examples do. generate opts into slime's request-level abort: when a step's batch is full, slime cancels the rollouts still playing, and each drops its session and closes its episode instead of playing on against the aborted model and running its verifiers. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/slime/generate.py | 33 +++++++++++---- tst/slime/generate_test.py | 54 +++++++++++++++++++++--- 2 files changed, 73 insertions(+), 14 deletions(-) diff --git a/src/agent_env/training/slime/generate.py b/src/agent_env/training/slime/generate.py index f3772eff4..25cb0bdf7 100644 --- a/src/agent_env/training/slime/generate.py +++ b/src/agent_env/training/slime/generate.py @@ -5,6 +5,7 @@ from __future__ import annotations +import asyncio import uuid from typing import Any @@ -22,21 +23,37 @@ async def generate(args: Any, sample: Sample, sampling_params: dict[str, Any], evaluation: bool = False) -> list[Sample]: """Plays the sample's row as one episode, the model chatting through a session of its own on slime's OpenAI adapter, which records its turns. Returns the session's turns as training samples carrying the episode's reward and - ``play``'s metadata: one sample, or one per turn the adapter couldn't join to the turn before. An episode with no - recorded turn comes back as a single aborted sample, which ``keep_scored_groups`` leaves out with its group.""" + ``play``'s metadata: one sample, or one per turn the adapter couldn't join to the turn before. In an eval it returns + one sample per episode, carrying only the reward, so slime's eval reward is a mean over episodes, not turns. An + episode with no recorded turn comes back as a single aborted sample, which ``keep_scored_groups`` leaves out with its + group. slime cancels a rollout it no longer needs (``abort_mode``): its session is dropped and its episode closed.""" adapter = _Adapter(args) sid = uuid.uuid4().hex adapter.adapter.open_session( sid, sampling_defaults=sampling_params, max_context_tokens=args.rollout_max_context_len or 0, apply_chat_template_kwargs={**args.apply_chat_template_kwargs, **(sample.apply_chat_template_kwargs or {})}) - played = await play(adapter.url, sample.prompt, {"user": sid}, sample.metadata) - samples = await adapter.adapter.finish_session(sid, base_sample=sample, reward=played["reward"], - extra_metadata=played) - if samples: - return samples + try: + played = await play(adapter.url, sample.prompt, {"user": sid}, sample.metadata) + if evaluation: + await adapter.adapter.drop_session(sid) + return _bare(sample, played, Sample.Status.COMPLETED) + samples = await adapter.adapter.finish_session(sid, base_sample=sample, reward=played["reward"], + extra_metadata=played) + except asyncio.CancelledError: + await adapter.adapter.drop_session(sid) + raise + return samples or _bare(sample, played, Sample.Status.ABORTED) + + +generate.abort_mode = "request" + + +def _bare(sample: Sample, played: dict[str, Any], status: Sample.Status) -> list[Sample]: + """``sample`` with no turns to train on, as slime's own agent examples return one: a masked placeholder token, the + episode's reward and metadata, and ``remove_sample``.""" sample.tokens, sample.response, sample.response_length = [0, 0], "", 1 sample.loss_mask, sample.rollout_log_probs = [0], [0.0] - sample.reward, sample.remove_sample, sample.status = played["reward"], True, Sample.Status.ABORTED + sample.reward, sample.remove_sample, sample.status = played["reward"], True, status sample.metadata = {**sample.metadata, **played} return [sample] diff --git a/tst/slime/generate_test.py b/tst/slime/generate_test.py index 5b01c8868..9bded874c 100644 --- a/tst/slime/generate_test.py +++ b/tst/slime/generate_test.py @@ -4,6 +4,7 @@ import asyncio import json +import threading from unittest.mock import patch import pytest @@ -42,21 +43,24 @@ def fresh_adapter(monkeypatch): @pytest.fixture def model(): - """A stand-in for SGLang's ``/generate``: adds a crate, then says so once the tool has answered; or answers 500 when - ``.broken``.""" + """A stand-in for SGLang's ``/generate``: adds a crate, then says so once the tool has answered, that reply held + back ``.stall`` seconds with ``.stalled`` set meanwhile; or answers 500 when ``.broken``.""" tokenizer = transformers.AutoTokenizer.from_pretrained(CHECKPOINT) async def generate(request): if server.broken: return JSONResponse({"error": "engine down"}, status_code=500) prompt = tokenizer.decode((await request.json())["input_ids"]) + if "" in prompt and server.stall: + server.stalled.set() + await asyncio.sleep(server.stall) ids = tokenizer.encode("Added a crate.<|im_end|>" if "" in prompt else ADDS_A_CRATE, add_special_tokens=False) return JSONResponse({"meta_info": {"output_token_logprobs": [[-0.01, each] for each in ids], "finish_reason": {"type": "stop"}}}) server = Server(Starlette(routes=[Route("/generate", generate, methods=["POST"])])) - server.broken, server.tokenizer = False, tokenizer + server.broken, server.stall, server.stalled, server.tokenizer = False, 0.0, threading.Event(), tokenizer yield server server.stop() @@ -76,6 +80,41 @@ def test_slime_plays_a_row_through_its_adapter_and_trains_on_the_models_turns_wi assert '"crate"' in trained and "Added a crate." in trained and "" not in trained +def test_an_eval_gives_slime_one_sample_per_episode_so_its_mean_is_over_episodes(serve, model, tmp_path): + server = serve({"a": PLAYED}) + row = _row(server, tmp_path) + args = _args(model, tmp_path) + + [sample] = _generate(row, args, evaluation=True) + server.stop() + + assert (sample.reward, sample.status, sample.loss_mask) == (1.0, types.Sample.Status.COMPLETED, [0]) + assert sample.metadata["agentenv_scores"] == {"replies": 1.0} + assert slime_generate._Adapter(args).adapter.store == {} + + +def test_a_rollout_slime_cancels_drops_its_session_and_closes_its_episode(serve, model, tmp_path, teardowns): + server = serve({"a": PLAYED}) + row = _row(server, tmp_path) + args = _args(model, tmp_path) + model.stall = 5.0 + + async def cancelled(): + rollout = asyncio.ensure_future(slime_generate.generate(args, _sample(row), dict(SAMPLING))) + while not model.stalled.is_set(): + await asyncio.sleep(0.05) + rollout.cancel() + with pytest.raises(asyncio.CancelledError): + await rollout + + asyncio.run(cancelled()) + server.stop() + + assert slime_generate.generate.abort_mode == "request" + assert slime_generate._Adapter(args).adapter.store == {} + assert len(teardowns.finished) == 2 + + def test_an_episode_the_env_failed_is_flagged_and_slime_leaves_its_group_out(serve, model, tmp_path): server = serve({"a": FAILS_AFTER_THE_TURN}) row = _row(server, tmp_path) @@ -118,10 +157,13 @@ def _row(server, tmp_path): return json.loads((folder / "all.jsonl").read_text().splitlines()[0]) -def _generate(row, args): +def _generate(row, args, evaluation=False): """The samples slime's rollout gets back for ``row``, as it calls the generate function.""" - sample = types.Sample(prompt=row["messages"], metadata=row["metadata"], index=0) - return asyncio.run(slime_generate.generate(args, sample, dict(SAMPLING))) + return asyncio.run(slime_generate.generate(args, _sample(row), dict(SAMPLING), evaluation=evaluation)) + + +def _sample(row): + return types.Sample(prompt=row["messages"], metadata=row["metadata"], index=0) def _trained(sample, tokenizer): From 10ac908e634ff0049c44906ac6597eb729512c59 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 3 Nov 2026 07:17:57 -0800 Subject: [PATCH 3/4] refactor(training): fewer slime internals, shared test tasks, and AGENTS.md's training helpers by module The slime helper depends on slime.agent's adapter and server helpers and slime's Sample only: its adapter is cached on the class, and its tokenizer loaded with transformers, rather than through slime's SingletonMeta and load_tokenizer, so the tier resets it without slime's private state. The tasks and model call the trainer tiers share (PLAYED, FAILS_AFTER_THE_TURN, ADDS_A_CRATE) live in tst/unit/training/_support.py, not copied between tiers or read from nemo_gym_test, and the shared rows and episode are tested in chat_test.py. AGENTS.md's training row becomes a Training helpers subsection, one entry per module, adding the harness helper it lacked. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 89 ++++++++++++++++++- src/agent_env/training/slime/generate.py | 27 +++--- tst/miles/rollout_test.py | 6 +- tst/nemo_gym/server_test.py | 3 +- tst/skyrl/generator_test.py | 2 +- tst/slime/generate_test.py | 17 ++-- tst/unit/training/_support.py | 14 ++- .../training/{miles_test.py => chat_test.py} | 21 +++-- tst/unit/training/nemo_gym_test.py | 2 +- tst/unit/training/skyrl_test.py | 2 +- tst/unit/training/verl_test.py | 2 +- tst/verifiers/taskset_test.py | 2 +- tst/verl/loop_test.py | 2 +- 13 files changed, 141 insertions(+), 48 deletions(-) rename tst/unit/training/{miles_test.py => chat_test.py} (89%) diff --git a/AGENTS.md b/AGENTS.md index 2507d8efe..f62b62a20 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -125,7 +125,7 @@ and deserialized through a registry. | `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `smol_vm` (the `smol` extra), `tensorlake` (the `tensorlake` extra), `tenki` (the `tenki` extra), `vercel` (the `vercel` extra), `boxd` (the `boxd` extra), `coreweave` (the `coreweave` extra), `daytona` (the `daytona` extra), `boat` (the `boat` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store), `EnvironmentGatewayProvider` (`gateway`, an env's default), which renders a docker-compose for the gateway and its MCP servers inside the sandbox, and `EnvironmentServerProvider` (`server`), one MCP server in its own container with no gateway in front and no state store; `providers/env_state/` holds env-state providers (`local_postgres` built in). | | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `openenv/` | `build_app(bundle)`: the [OpenEnv](https://github.com/openenv/OpenEnv) server for a bundle's tasks, each session playing one episode through the Episode API, with the envs' MCP tools, an added `end_turn` tool and the task's verifier score as the reward. It prepares bundles as `agent-env run` does (`prepare_bundle(played=True)`, without the agent the client plays) and keeps namespace routing on for the app's life. `openenv` comes with the `openenv` extra, a floor with no cap; this package and `training/` are the only ones that import it. Each task's description in the Task API names its `toolset` (the played agent's envs, role, and whether it gets `bash`). With `token`, every request and WebSocket but `/health` must carry `Authorization: Bearer `, for a server others reach. | -| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, `tasks_and_tools` (what the row writers read: the tasks, the split names checked, each toolset's tools) and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and raising `RolloutFailed` for one that fails every time, which verl's V1 trainer (the default) leaves out of the loss, the group's baseline and the eval reward; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing, and warns once a tool of an answer verl cuts to `max_tool_response_length`. The writer reads a server that asks for a token with the one `--token-var` names, which the loop reads in each worker, and with `--tokenizer` counts each prompt with its tools as verl renders it (`--chat-template-kwargs`), refusing rows over `--max-prompt-length` before writing. `skyrl/` (`skyrl` extra, installed into SkyRL's own environment): `python -m agent_env.training.skyrl URL OUT --tokenizer MODEL` writes a server's tasks as SkyRL's `.parquet` rows (`env_class: agentenv`, the episode's split, index, task id and tools, the server's settings), each with its tools rendered into its system prompt by the model's chat template, checked to render as `tools=` would, since SkyRL's generator passes the template no tools, and refuses rows over `--max-prompt-length` counted as SkyRL's filter and generator count them; `calls.parse` reads a reply's tool calls in the `hermes` or `qwen3_coder` format, typed by the row's schema; `env.AgentEnvTextEnv` (registered as `agentenv`) is skyrl-gym's `BaseTextEnv` on the row's episode, run on the event loop the sessions share (`_openenv.on_session_loop`), its steps running the reply's calls and a reply without one ending the open turn, raising `EpisodeFailed` when the env fails it, with scores, tool calls and unparsed calls as its metrics; `generator.AgentEnvGenerator` is SkyRL's `SkyRLGymGenerator` playing a failed rollout again on a new episode up to `attempts` times within `reset_timeout`, and masking the group of one that fails every attempt in training (an eval scores it 0); `python -m agent_env.training.skyrl.train` (or `trainer.override_entrypoint=agent_env.training.skyrl.train`) runs SkyRL's trainer with both, its fully async one with `trainer.fully_async.enabled`, passing the rows' token variables to the trainer's Ray task. `_chat.py` is what the Miles and slime helpers share: `write_rows` and its `command` write a server's tasks as `.jsonl` rows of chat `messages` and `metadata` (the episode's split, index, task id and tools, and the server it plays on with its deadlines and turn cap); `play` plays a row's episode, the model chatting at a trainer's OpenAI chat endpoint that records its turns, with the row's tools, answers a call whose arguments aren't a JSON object as the model's error, and returns the reward, or flags an episode the env or the session failed, whose model missed the row's `model_timeout` or that ended without a reward, which isn't played again since the session holds its turns; `reward` reads a sample's reward from its metadata, and `keep_scored_groups` leaves out the group of a flagged episode, of a sample the trainer aborted, or of one whose episode never returned. `miles/` (`miles` extra, installed into Miles' own environment): `python -m agent_env.training.miles URL OUT` writes the rows as Miles' prompt data; `agent.run` (`play`) is Miles' agent function (`--custom-agent-function-path`), chatting through the sample's session URL on Miles' session server; `agent.reward` is the reward hook (`--custom-rm-path`) and `agent.keep_scored_groups` the dynamic-sampling filter. It imports nothing of Miles. `slime/` (`slime` extra, installed into slime's own environment): `python -m agent_env.training.slime URL OUT` writes the rows as slime's prompt data; `generate.generate` is slime's generate function (`--custom-generate-function-path`), run on its rollout's event loop, which opens a session for the sample on slime's `OpenAIAdapter` (`slime.agent`, one per rollout process, served in-process and reading tool calls with `--sglang-tool-call-parser`, which it requires), plays the row with `play` through it, and returns the session's turns as samples carrying the reward, one per turn the adapter couldn't join to the turn before, or one aborted sample when no turn was recorded; `generate.keep_scored_groups` is the dynamic-sampling filter. | +| `training/` | One module per outside trainer, each behind its own extra, and the modules they share; [Training helpers](#training-helpers) lists them. | | `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. | @@ -138,6 +138,93 @@ and deserialized through a registry. | `drivers/` | The `drivers` bundle: the computer drivers agent-env ships as `computer_driver` envs, the reference driver (X11 on Linux, recording through `x11grab`) aliased `reference`, and trycua's cua-driver aliased `cua-driver` (its daemon from the pinned PyPI wheel, its own instructions, tools and skill resources served to agents through `@server_instructions`, `forwarded_tools` and `@resource`, its tools beside `computer_*`, `computer/v1` on its desktop-scope calls, its skill pack leading the computer skill; `tst/integration/computer/cua_driver_test.py` runs it). It has no tasks: a bundle with neither tasks nor evals is refused unless its `aliases.toml` names something for other bundles to use. The shipped `desktop` example composes it with the `ubuntu-24.04` base into a computer env, which `tst/integration/computer/` deploys, probes, records and has an agent use. | | `examples/` | The example bundles agent-env ships (`hello`, and `echo-mcp`, which needs Docker, and `desktop`, a computer env of `ubuntu-24.04` and `reference` with a probe task and an agent task). Tests run them, but they are examples for users, not fixtures. Every shipped bundle is registered under `agent_env.bundles` in `pyproject.toml`, ` = ""` shipping `src///`. | +### Training helpers + +`src/agent_env/training/` holds one module per outside trainer, each behind an extra of its own, and the modules they +share. + +Shared: +- `_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`. +- `_tools.py`: turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' + `get_json_schema` reads them. +- `_chat.py`, the rows and episode the Miles and slime helpers share: + - `write_rows` and its `command` write a server's tasks as `.jsonl` rows of chat `messages` and `metadata` + (the episode's split, index, task id and tools, and the server it plays on with its deadlines and turn cap). + - `play` plays a row's episode, the model chatting at a trainer's OpenAI chat endpoint that records its turns, with + the row's tools. It answers a call whose arguments aren't a JSON object as the model's error, and returns the reward, + or flags an episode the env or the session failed, whose model missed the row's `model_timeout` or that ended + without a reward; a flagged episode isn't played again, since the session holds its turns. + - `reward` reads a sample's reward from its metadata, and `keep_scored_groups` leaves out the group of a flagged + episode, of a sample the trainer aborted, or of one whose episode never returned. + +Per trainer: +- `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. +- `harness.py` (`harness` extra): trains a bundle's own agents with TRL's harness rollouts (`HarnessRolloutWorker` with + `harness_adapter=None`). `AgentEnvSessionFactory` runs each rollout's task as `agent-env run` does: the solver, the + agent a step prompts, reaches the trainer's vLLM through OpenEnv's capture proxy, which records the tokens TRL trains + on, while the task's other agents and judges keep the configured endpoint, and the task's verifiers score the run; + `prompt_rows` is the trainer's dataset, a row per task. +- `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`. It reads a server that asks + for a token with the one `--token-var` names, and with `--tokenizer` counts each prompt with its tools as verl + renders it (`--chat-template-kwargs`), refusing rows over `--max-prompt-length` before writing. + - `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`. It plays a rollout the env failed again + on a new episode up to `attempts` times within `reset_timeout` of its start, and raises `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. + - The loop hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them) and reads the token + variable in each worker. It 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`. +- `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 row has 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; rows over `--max-prompt-length`, counted as SkyRL's filter and + generator count them, are refused. + - `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 run the reply's calls, a reply without one + ends the open turn, it raises `EpisodeFailed` when the env fails it, and its metrics are the scores, tool calls and + unparsed calls. + - `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. +- `miles/` (`miles` extra, installed into Miles' own environment; it imports nothing of Miles): + `python -m agent_env.training.miles URL OUT` writes `_chat`'s rows as Miles' prompt data; `agent.run` (`_chat.play`) is + Miles' agent function (`--custom-agent-function-path`), chatting through the sample's session URL on Miles' session + server; `agent.reward` is the reward hook (`--custom-rm-path`) and `agent.keep_scored_groups` the dynamic-sampling + filter. +- `slime/` (`slime` extra, installed into slime's own environment): + - `python -m agent_env.training.slime URL OUT` writes `_chat`'s rows as slime's prompt data. + - `generate.generate` is slime's generate function (`--custom-generate-function-path`), run on its rollout's event + loop. It opens a session for the sample on slime's `OpenAIAdapter` (`slime.agent`: one per rollout process, made for + its first sample, served in-process and reading tool calls with `--sglang-tool-call-parser`, which it requires), + plays the row with `play` through it, and returns the session's turns as samples carrying the reward: one per turn + the adapter couldn't join to the turn before, or one aborted sample when no turn was recorded. In an eval it returns + one sample per episode, carrying only the reward, so slime's eval reward is a mean over episodes. + - It takes slime's request-level abort (`abort_mode = "request"`): a rollout slime cancels drops its session and + closes its episode. `generate.keep_scored_groups` is the dynamic-sampling filter. + ## Extension points Out-of-tree code plugs in either through `config.toml`, as a `module:Class` pointer constructed with diff --git a/src/agent_env/training/slime/generate.py b/src/agent_env/training/slime/generate.py index 25cb0bdf7..385f088a6 100644 --- a/src/agent_env/training/slime/generate.py +++ b/src/agent_env/training/slime/generate.py @@ -7,13 +7,12 @@ import asyncio import uuid -from typing import Any +from typing import Any, ClassVar from slime.agent.adapters import OpenAIAdapter from slime.agent.aiohttp_threaded import run_app_in_thread -from slime.utils.misc import SingletonMeta -from slime.utils.processing_utils import load_tokenizer from slime.utils.types import Sample +from transformers import AutoTokenizer from agent_env.training._chat import keep_scored_groups, play @@ -27,7 +26,7 @@ async def generate(args: Any, sample: Sample, sampling_params: dict[str, Any], e one sample per episode, carrying only the reward, so slime's eval reward is a mean over episodes, not turns. An episode with no recorded turn comes back as a single aborted sample, which ``keep_scored_groups`` leaves out with its group. slime cancels a rollout it no longer needs (``abort_mode``): its session is dropped and its episode closed.""" - adapter = _Adapter(args) + adapter = _Adapter.of(args) sid = uuid.uuid4().hex adapter.adapter.open_session( sid, sampling_defaults=sampling_params, max_context_tokens=args.rollout_max_context_len or 0, @@ -58,19 +57,27 @@ def _bare(sample: Sample, played: dict[str, Any], status: Sample.Status) -> list return [sample] -class _Adapter(metaclass=SingletonMeta): +class _Adapter: """The rollout process's OpenAI adapter, in front of its SGLang router, reading the model's tool calls with - ``--sglang-tool-call-parser``.""" + ``--sglang-tool-call-parser``; made for the process's first sample.""" + + served: ClassVar[_Adapter | None] = None + + @classmethod + def of(cls, args: Any) -> _Adapter: + if cls.served is None: + cls.served = cls(args) + return cls.served def __init__(self, args: Any): if not args.sglang_tool_call_parser: raise ValueError("slime's OpenAI adapter reads the model's tool calls with --sglang-tool-call-parser; pass " "the parser for the model's chat template, qwen25 for Qwen2.5 and Qwen3") self.adapter = OpenAIAdapter( - tokenizer=load_tokenizer(args.hf_checkpoint, trust_remote_code=True), + tokenizer=AutoTokenizer.from_pretrained(args.hf_checkpoint, trust_remote_code=True), sglang_url=f"http://{args.sglang_router_ip}:{args.sglang_router_port}", tool_parser=args.sglang_tool_call_parser, reasoning_parser=args.sglang_reasoning_parser, rollout_args=args) - served = run_app_in_thread(self.adapter.app, host="127.0.0.1", thread_name="agentenv-adapter", - runner_kwargs={"handler_cancellation": True}) - self.url = f"http://127.0.0.1:{served.port}" + app = run_app_in_thread(self.adapter.app, host="127.0.0.1", thread_name="agentenv-adapter", + runner_kwargs={"handler_cancellation": True}) + self.url = f"http://127.0.0.1:{app.port}" diff --git a/tst/miles/rollout_test.py b/tst/miles/rollout_test.py index c09386eba..84dcd02cc 100644 --- a/tst/miles/rollout_test.py +++ b/tst/miles/rollout_test.py @@ -10,8 +10,7 @@ import pytest from agent_env.training.miles.prepare import write_rows -from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN -from tst.unit.training.nemo_gym_test import PLAYED +from tst.unit.training._support import ADDS_A_CRATE, FAILS_AFTER_THE_TURN, PLAYED from tst.util.capabilities import MILES, missing_capability_reason fixtures = pytest.importorskip("tests.fast.fixtures.rollout_fixtures", reason=missing_capability_reason(MILES)) @@ -24,9 +23,6 @@ mock_sglang = pytest.importorskip("miles.utils.test_utils.mock_sglang_server", reason=missing_capability_reason(MILES)) sample_types = pytest.importorskip("miles.utils.types", reason=missing_capability_reason(MILES)) -ADDS_A_CRATE = '\n{"name": "add_item", "arguments": {"item": "crate"}}\n' -FAILS_AFTER_THE_TURN = [ITEMS, ALICE, TURN, {"id": "crane", "type": "episode_test_env", "env_id": "crane", "fail": True}, - REPLIES] @pytest.fixture(params=[[], ["v2"]], ids=["session-server-v1", "session-server-v2"]) diff --git a/tst/nemo_gym/server_test.py b/tst/nemo_gym/server_test.py index abc510c91..ecc588a42 100644 --- a/tst/nemo_gym/server_test.py +++ b/tst/nemo_gym/server_test.py @@ -25,7 +25,8 @@ from tst.unit.bundle._support import layout from tst.unit.openenv.environment_test import AGAIN from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN -from tst.unit.training.nemo_gym_test import ADD_ITEM, FUNCTION_END_TURN, PLAYED, STYLE +from tst.unit.training._support import PLAYED +from tst.unit.training.nemo_gym_test import ADD_ITEM, FUNCTION_END_TURN, STYLE from tst.util.capabilities import NEMO_GYM, missing_capability_reason server = pytest.importorskip("agent_env.training.nemo_gym.server", reason=missing_capability_reason(NEMO_GYM), diff --git a/tst/skyrl/generator_test.py b/tst/skyrl/generator_test.py index d7e2c5208..ec34f855e 100644 --- a/tst/skyrl/generator_test.py +++ b/tst/skyrl/generator_test.py @@ -12,7 +12,7 @@ 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.unit.training._support 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), diff --git a/tst/slime/generate_test.py b/tst/slime/generate_test.py index 9bded874c..7e6789bbe 100644 --- a/tst/slime/generate_test.py +++ b/tst/slime/generate_test.py @@ -14,31 +14,24 @@ from agent_env.training._chat import write_rows from agent_env.training._openenv import Server -from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN -from tst.unit.training.nemo_gym_test import PLAYED +from tst.unit.training._support import ADDS_A_CRATE, FAILS_AFTER_THE_TURN, PLAYED from tst.util.capabilities import SLIME, missing_capability_reason slime_generate = pytest.importorskip("agent_env.training.slime.generate", reason=missing_capability_reason(SLIME), exc_type=ImportError) arguments = pytest.importorskip("slime.utils.arguments", reason=missing_capability_reason(SLIME)) filters = pytest.importorskip("slime.rollout.filter_hub.base_types", reason=missing_capability_reason(SLIME)) -misc = pytest.importorskip("slime.utils.misc", reason=missing_capability_reason(SLIME)) types = pytest.importorskip("slime.utils.types", reason=missing_capability_reason(SLIME)) transformers = pytest.importorskip("transformers", reason=missing_capability_reason(SLIME)) CHECKPOINT = "Qwen/Qwen3-0.6B" -ADDS_A_CRATE = '\n{"name": "add_item", "arguments": {"item": "crate"}}\n<|im_end|>' -FAILS_AFTER_THE_TURN = [ITEMS, ALICE, TURN, {"id": "crane", "type": "episode_test_env", "env_id": "crane", "fail": True}, - REPLIES] SAMPLING = {"temperature": 1.0, "top_p": 1.0, "max_new_tokens": 256} @pytest.fixture(autouse=True) def fresh_adapter(monkeypatch): """Each test's own adapter, in front of its own stand-in SGLang: the module keeps one per process.""" - monkeypatch.delitem(misc.SingletonMeta._instances, slime_generate._Adapter, raising=False) - yield - misc.SingletonMeta._instances.pop(slime_generate._Adapter, None) + monkeypatch.setattr(slime_generate._Adapter, "served", None) @pytest.fixture @@ -54,7 +47,7 @@ async def generate(request): if "" in prompt and server.stall: server.stalled.set() await asyncio.sleep(server.stall) - ids = tokenizer.encode("Added a crate.<|im_end|>" if "" in prompt else ADDS_A_CRATE, + ids = tokenizer.encode(f"{'Added a crate.' if '' in prompt else ADDS_A_CRATE}<|im_end|>", add_special_tokens=False) return JSONResponse({"meta_info": {"output_token_logprobs": [[-0.01, each] for each in ids], "finish_reason": {"type": "stop"}}}) @@ -90,7 +83,7 @@ def test_an_eval_gives_slime_one_sample_per_episode_so_its_mean_is_over_episodes assert (sample.reward, sample.status, sample.loss_mask) == (1.0, types.Sample.Status.COMPLETED, [0]) assert sample.metadata["agentenv_scores"] == {"replies": 1.0} - assert slime_generate._Adapter(args).adapter.store == {} + assert slime_generate._Adapter.of(args).adapter.store == {} def test_a_rollout_slime_cancels_drops_its_session_and_closes_its_episode(serve, model, tmp_path, teardowns): @@ -111,7 +104,7 @@ async def cancelled(): server.stop() assert slime_generate.generate.abort_mode == "request" - assert slime_generate._Adapter(args).adapter.store == {} + assert slime_generate._Adapter.of(args).adapter.store == {} assert len(teardowns.finished) == 2 diff --git a/tst/unit/training/_support.py b/tst/unit/training/_support.py index 602373693..12ae529c8 100644 --- a/tst/unit/training/_support.py +++ b/tst/unit/training/_support.py @@ -1,5 +1,6 @@ -"""The fake bundle both trainer helpers' tests serve: the Episode API's fake steps as an OpenEnv env on loopback. Its -fixture takes the ``teardowns`` and ``tools`` fixtures from the conftest that imports it.""" +"""What the trainer helpers' tests share: the fake bundle they serve, the Episode API's fake steps as an OpenEnv env on +loopback, whose fixture takes the ``teardowns`` and ``tools`` fixtures from the conftest that imports it, and the tasks +and model replies they play it with.""" import pytest @@ -8,6 +9,15 @@ from agent_env.openenv.environment import Served, ServedTask from agent_env.task import EpisodeSpec, Task from agent_env.training._openenv import Server +from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN + +#: A task the model plays by adding a crate and saying so, which its verifier scores 1. +PLAYED = [ITEMS, ALICE, TURN, REPLIES] +#: The same task with an env that fails once the model's turn has ended. +FAILS_AFTER_THE_TURN = [ITEMS, ALICE, TURN, + {"id": "crane", "type": "episode_test_env", "env_id": "crane", "fail": True}, REPLIES] +#: The model's call that adds the crate, in Qwen's tool-call format. +ADDS_A_CRATE = '\n{"name": "add_item", "arguments": {"item": "crate"}}\n' @pytest.fixture diff --git a/tst/unit/training/miles_test.py b/tst/unit/training/chat_test.py similarity index 89% rename from tst/unit/training/miles_test.py rename to tst/unit/training/chat_test.py index a70aa2e80..ea9475be5 100644 --- a/tst/unit/training/miles_test.py +++ b/tst/unit/training/chat_test.py @@ -1,7 +1,7 @@ -"""``agent_env.training.miles``, and the chat rows and episode it shares with ``agent_env.training.slime``, on the -OpenEnv server ``build_app`` makes, served on loopback over the Episode API's fake steps: the rows ``write_rows`` writes -and both helpers' commands, and the agent function playing a row's task through a stand-in for a trainer's session -endpoint, which answers each chat request from a script and records it; with the reward hook and the filter beside it.""" +"""``agent_env.training._chat``, the chat rows and episode the Miles and slime helpers share, on the OpenEnv server +``build_app`` makes, served on loopback over the Episode API's fake steps: the rows ``write_rows`` writes and both +helpers' commands, and ``play`` playing a row's task through a stand-in for a trainer's session endpoint, which answers +each chat request from a script and records it; with the reward hook and the filter beside it.""" import asyncio import json @@ -17,11 +17,10 @@ from agent_env.openenv.reward import _unscored from agent_env.training._openenv import Server from agent_env.training.miles.__main__ import main as miles_main -from agent_env.training.miles.agent import keep_scored_groups, reward, run -from agent_env.training.miles.prepare import write_rows +from agent_env.training._chat import keep_scored_groups, play, reward, write_rows from agent_env.training.slime.__main__ import main as slime_main from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN -from tst.unit.training.nemo_gym_test import PLAYED +from tst.unit.training._support import PLAYED ADD_ITEM = {"type": "function", "function": {"name": "add_item", "description": "add_item.", "parameters": {"type": "object"}}} @@ -96,7 +95,7 @@ def test_the_agent_plays_a_rows_task_through_the_session_and_returns_its_reward( row = _row(server, tmp_path) model = session(_call('{"item": "crate"}'), _says("Added a crate.")) - played = asyncio.run(run(model.base_url, row["messages"], SAMPLING, row["metadata"])) + played = asyncio.run(play(model.base_url, row["messages"], SAMPLING, row["metadata"])) server.stop() assert played == {"reward": 1.0, "agentenv_scores": {"replies": 1.0}, "agentenv_tool_calls": 1, @@ -114,7 +113,7 @@ def test_a_call_whose_arguments_arent_an_object_is_answered_and_the_episode_goes model = session(_call('"{\\"item\\": \\"crate\\"}"'), _call("not json", id="call_2"), _call("", id="call_3"), _says("Added a crate.")) - played = asyncio.run(run(model.base_url, row["messages"], SAMPLING, row["metadata"])) + played = asyncio.run(play(model.base_url, row["messages"], SAMPLING, row["metadata"])) server.stop() assert (played["reward"], played["agentenv_tool_calls"], played["agentenv_invalid_calls"]) == (1.0, 0, 3) @@ -130,7 +129,7 @@ def test_a_rollout_that_reaches_its_turn_cap_ends_the_turn_without_a_reply(serve row = _row(server, tmp_path, max_turns=2) model = session(_call('{"item": "crate"}')) - played = asyncio.run(run(model.base_url, row["messages"], SAMPLING, row["metadata"])) + played = asyncio.run(play(model.base_url, row["messages"], SAMPLING, row["metadata"])) server.stop() assert (played["reward"], played["agentenv_tool_calls"], len(model.requests)) == (0.0, 2, 2) @@ -148,7 +147,7 @@ def test_an_episode_that_fails_or_ends_without_a_reward_is_flagged_and_its_group if failing == "env": server.stop() - played = asyncio.run(run(model.base_url, row["messages"], SAMPLING, row["metadata"])) + played = asyncio.run(play(model.base_url, row["messages"], SAMPLING, row["metadata"])) server.stop() assert played["reward"] == 0.0 and played["agentenv_failure"].startswith(why) diff --git a/tst/unit/training/nemo_gym_test.py b/tst/unit/training/nemo_gym_test.py index 2957897e3..ed31c0ad1 100644 --- a/tst/unit/training/nemo_gym_test.py +++ b/tst/unit/training/nemo_gym_test.py @@ -23,8 +23,8 @@ TURN, _tool_sessions, ) +from tst.unit.training._support import PLAYED -PLAYED = [ITEMS, ALICE, TURN, REPLIES] STYLE = {**REPLIES, "id": "style", "verifier_id": "style"} ADD_ITEM = {"type": "function", "name": "add_item", "description": "add_item.", "parameters": {"type": "object"}, "strict": False} diff --git a/tst/unit/training/skyrl_test.py b/tst/unit/training/skyrl_test.py index 35c5d4339..2e7a30fe2 100644 --- a/tst/unit/training/skyrl_test.py +++ b/tst/unit/training/skyrl_test.py @@ -24,7 +24,7 @@ TURN, _tool_sessions, ) -from tst.unit.training.nemo_gym_test import PLAYED +from tst.unit.training._support import PLAYED ADD_ITEM = {"type": "function", "function": {"name": "add_item", "description": "add_item.", "parameters": {"type": "object"}}} diff --git a/tst/unit/training/verl_test.py b/tst/unit/training/verl_test.py index 875ba993a..d472ce5c3 100644 --- a/tst/unit/training/verl_test.py +++ b/tst/unit/training/verl_test.py @@ -23,7 +23,7 @@ TURN, _tool_sessions, ) -from tst.unit.training.nemo_gym_test import PLAYED +from tst.unit.training._support import PLAYED ADD_ITEM = {"type": "function", "function": {"name": "add_item", "description": "add_item.", "parameters": {"type": "object"}}} diff --git a/tst/verifiers/taskset_test.py b/tst/verifiers/taskset_test.py index c3fbc9e28..2ab35de9c 100644 --- a/tst/verifiers/taskset_test.py +++ b/tst/verifiers/taskset_test.py @@ -16,7 +16,7 @@ from agent_env.training.nemo_gym.prepare import write_server from tst.nemo_gym.server_test import _free_port, _up -from tst.unit.training.nemo_gym_test import PLAYED +from tst.unit.training._support import PLAYED from tst.util.capabilities import NEMO_GYM, VERIFIERS, missing_capability_reason pytest.importorskip("agent_env.training.nemo_gym.server", reason=missing_capability_reason(NEMO_GYM), diff --git a/tst/verl/loop_test.py b/tst/verl/loop_test.py index e4e490c57..267377a6f 100644 --- a/tst/verl/loop_test.py +++ b/tst/verl/loop_test.py @@ -17,7 +17,7 @@ from agent_env.training.verl.prepare import LOOP, write_dataset from tst.unit.task.episode import run_test -from tst.unit.training.nemo_gym_test import PLAYED +from tst.unit.training._support import PLAYED from tst.util.capabilities import VERL, missing_capability_reason verl_loop = pytest.importorskip("agent_env.training.verl.loop", reason=missing_capability_reason(VERL), From b5e86db39b9d5b226d612501f2ae0312dafd3759 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 3 Nov 2026 07:43:18 -0800 Subject: [PATCH 4/4] fix(slime): load the adapter's tokenizer through slime, not transformers directly Importing transformers made the module read as using the extras that declare it, which don't declare slime; declaring transformers in the slime extra would let an install move slime's environment off the transformers its SGLang pins. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/slime/generate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/agent_env/training/slime/generate.py b/src/agent_env/training/slime/generate.py index 385f088a6..c87ae8f51 100644 --- a/src/agent_env/training/slime/generate.py +++ b/src/agent_env/training/slime/generate.py @@ -11,8 +11,8 @@ from slime.agent.adapters import OpenAIAdapter from slime.agent.aiohttp_threaded import run_app_in_thread +from slime.utils.processing_utils import load_tokenizer from slime.utils.types import Sample -from transformers import AutoTokenizer from agent_env.training._chat import keep_scored_groups, play @@ -74,7 +74,7 @@ def __init__(self, args: Any): raise ValueError("slime's OpenAI adapter reads the model's tool calls with --sglang-tool-call-parser; pass " "the parser for the model's chat template, qwen25 for Qwen2.5 and Qwen3") self.adapter = OpenAIAdapter( - tokenizer=AutoTokenizer.from_pretrained(args.hf_checkpoint, trust_remote_code=True), + tokenizer=load_tokenizer(args.hf_checkpoint, trust_remote_code=True), sglang_url=f"http://{args.sglang_router_ip}:{args.sglang_router_port}", tool_parser=args.sglang_tool_call_parser, reasoning_parser=args.sglang_reasoning_parser, rollout_args=args)