diff --git a/.github/workflows/miles.yml b/.github/workflows/miles.yml new file mode 100644 index 00000000..c89cd628 --- /dev/null +++ b/.github/workflows/miles.yml @@ -0,0 +1,78 @@ +name: Miles contract + +# agent_env.training.miles on Miles itself (tst/miles): Miles' rollout plays rows the writer wrote through its agentic +# generate path, session server, router and mock SGLang server, scores them with the helper's reward hook and groups +# them through its filter, on Miles v0.1.3 and on its main branch. Miles isn't on PyPI: as Miles' own CPU CI does, its +# checkout, SGLang's and Megatron-LM's go on PYTHONPATH, 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 its CPU CI's are +# installed over the dev environment, torch the CPU build. 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/miles/**" + - "tst/unit/training/**" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/miles.yml" + schedule: + - cron: "17 8 * * 1" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: miles-${{ 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: + rollout: + name: rollout (Miles ${{ matrix.miles }}) + runs-on: ubuntu-latest + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + miles: ["v0.1.3", "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 miles extra + run: | + python -m pip install --quiet -r .github/ci-requirements.txt + uv sync --locked --extra dev --extra miles --python "$(command -v python)" + - name: Check out Miles, SGLang and Megatron-LM + env: + MILES_REF: ${{ matrix.miles }} + 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 miles https://github.com/radixark/miles "$MILES_REF" + sglang=sglang-miles megatron=miles-main + if [ -f "$RUNNER_TEMP/miles/release-lock.json" ]; then + sglang=$(jq -er .sglang_commit "$RUNNER_TEMP/miles/release-lock.json") + megatron=$(jq -er .megatron_commit "$RUNNER_TEMP/miles/release-lock.json") + fi + checkout sglang https://github.com/sgl-project/sglang "$sglang" + checkout Megatron-LM https://github.com/radixark/Megatron-LM "$megatron" + echo "PYTHONPATH=$RUNNER_TEMP/miles:$RUNNER_TEMP/sglang/python:$RUNNER_TEMP/Megatron-LM" >> "$GITHUB_ENV" + - name: Install Miles' requirements and CPU torch + run: | + uv pip install --python .venv/bin/python -r "$RUNNER_TEMP/miles/requirements.txt" \ + -r "$RUNNER_TEMP/miles/tests/ci/requirements-ci-cpu.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'))" + - name: Miles tier + run: .venv/bin/python -m pytest tst/miles -rA --timeout=900 --junitxml=junit-miles.xml + - name: No skips + if: always() + run: .venv/bin/python .github/scripts/check_skip_policy.py junit-miles.xml diff --git a/AGENTS.md b/AGENTS.md index 6c7dffc9..bf0168d1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -48,6 +48,7 @@ or one of them only where it is used. | `make nemo-gym-test` | `tst/nemo_gym`: rollouts through `agent_env.training.nemo_gym`'s resources server in-process; needs the `nemo-gym` extra, on Python 3.13.14 or later, which the `dev` extra leaves out (`uv sync --extra dev --extra nemo-gym`) | | `make verl-test` | `tst/verl`: rows played through `agent_env.training.verl`'s agent loop on verl itself, with Qwen3's tokenizer from the Hub and a stand-in rollout server; needs verl and CPU torch, which the `dev` extra leaves out | | `make skyrl-test` | `tst/skyrl`: rows played through `agent_env.training.skyrl`'s generator and env on SkyRL itself, with Qwen3's and Qwen3.5's tokenizers from the Hub and a stand-in inference engine; needs SkyRL, skyrl-gym and CPU torch on Python 3.12+, which the `dev` extra leaves out | +| `make 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 clean-install-test` | the `clean-install` CI gate: both distributions built as the release builds them, installed into a fresh venv from public PyPI with an allowlisted environment, and `agent-env run hello` run twice by name and checked through the local store; needs Python 3.11, uv and git | | `npm ci && npm run typecheck && npm run lint && npm run test:smoke` | in `src/agent_env/explorer/ui`: the explorer UI checks the `ui` job runs (it also runs `npm run build:static`); needs Node 22 | | `pytest tst/unit/env/store_test.py -v` | one file; add `--log-cli-level=DEBUG` for debug logs | @@ -123,7 +124,7 @@ and deserialized through a registry. | `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `smol_vm` (the `smol` extra), `tensorlake` (the `tensorlake` extra), `tenki` (the `tenki` extra), `vercel` (the `vercel` extra), `boxd` (the `boxd` extra), `coreweave` (the `coreweave` extra), `daytona` (the `daytona` extra), `boat` (the `boat` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store), `EnvironmentGatewayProvider` (`gateway`, an env's default), which renders a docker-compose for the gateway and its MCP servers inside the sandbox, and `EnvironmentServerProvider` (`server`), one MCP server in its own container with no gateway in front and no state store; `providers/env_state/` holds env-state providers (`local_postgres` built in). | | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `openenv/` | `build_app(bundle)`: the [OpenEnv](https://github.com/openenv/OpenEnv) server for a bundle's tasks, each session playing one episode through the Episode API, with the envs' MCP tools, an added `end_turn` tool and the task's verifier score as the reward. It prepares bundles as `agent-env run` does (`prepare_bundle(played=True)`, without the agent the client plays) and keeps namespace routing on for the app's life. `openenv` comes with the `openenv` extra, a floor with no cap; this package and `training/` are the only ones that import it. Each task's description in the Task API names its `toolset` (the played agent's envs, role, and whether it gets `bash`). With `token`, every request and WebSocket but `/health` must carry `Authorization: Bearer `, for a server others reach. | -| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, `tasks_and_tools` (what the row writers read: the tasks, the split names checked, each toolset's tools) and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and raising `RolloutFailed` for one that fails every time, which verl's V1 trainer (the default) leaves out of the loss, the group's baseline and the eval reward; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing, and warns once a tool of an answer verl cuts to `max_tool_response_length`. The writer reads a server that asks for a token with the one `--token-var` names, which the loop reads in each worker, and with `--tokenizer` counts each prompt with its tools as verl renders it (`--chat-template-kwargs`), refusing rows over `--max-prompt-length` before writing. `skyrl/` (`skyrl` extra, installed into SkyRL's own environment): `python -m agent_env.training.skyrl URL OUT --tokenizer MODEL` writes a server's tasks as SkyRL's `.parquet` rows (`env_class: agentenv`, the episode's split, index, task id and tools, the server's settings), each with its tools rendered into its system prompt by the model's chat template, checked to render as `tools=` would, since SkyRL's generator passes the template no tools, and refuses rows over `--max-prompt-length` counted as SkyRL's filter and generator count them; `calls.parse` reads a reply's tool calls in the `hermes` or `qwen3_coder` format, typed by the row's schema; `env.AgentEnvTextEnv` (registered as `agentenv`) is skyrl-gym's `BaseTextEnv` on the row's episode, run on the event loop the sessions share (`_openenv.on_session_loop`), its steps running the reply's calls and a reply without one ending the open turn, raising `EpisodeFailed` when the env fails it, with scores, tool calls and unparsed calls as its metrics; `generator.AgentEnvGenerator` is SkyRL's `SkyRLGymGenerator` playing a failed rollout again on a new episode up to `attempts` times within `reset_timeout`, and masking the group of one that fails every attempt in training (an eval scores it 0); `python -m agent_env.training.skyrl.train` (or `trainer.override_entrypoint=agent_env.training.skyrl.train`) runs SkyRL's trainer with both, its fully async one with `trainer.fully_async.enabled`, passing the rows' token variables to the trainer's Ray task. | +| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, `tasks_and_tools` (what the row writers read: the tasks, the split names checked, each toolset's tools) and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and raising `RolloutFailed` for one that fails every time, which verl's V1 trainer (the default) leaves out of the loss, the group's baseline and the eval reward; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing, and warns once a tool of an answer verl cuts to `max_tool_response_length`. The writer reads a server that asks for a token with the one `--token-var` names, which the loop reads in each worker, and with `--tokenizer` counts each prompt with its tools as verl renders it (`--chat-template-kwargs`), refusing rows over `--max-prompt-length` before writing. `skyrl/` (`skyrl` extra, installed into SkyRL's own environment): `python -m agent_env.training.skyrl URL OUT --tokenizer MODEL` writes a server's tasks as SkyRL's `.parquet` rows (`env_class: agentenv`, the episode's split, index, task id and tools, the server's settings), each with its tools rendered into its system prompt by the model's chat template, checked to render as `tools=` would, since SkyRL's generator passes the template no tools, and refuses rows over `--max-prompt-length` counted as SkyRL's filter and generator count them; `calls.parse` reads a reply's tool calls in the `hermes` or `qwen3_coder` format, typed by the row's schema; `env.AgentEnvTextEnv` (registered as `agentenv`) is skyrl-gym's `BaseTextEnv` on the row's episode, run on the event loop the sessions share (`_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. | | `runner/` | The `[runner]` seam: `Runner.submit()` returns a `RunHandle` (`run_id`, `instance_id`) before the run starts; `LocalRunner` is built in. `runner/store.py` is the run index every runner writes and the explorer reads. | | `bundle/` | A bundle is a folder of entities (`envs/`, `agents/`, `artifacts/`, `skills/`, `tasks/.json`, `evals/.toml`) that runs without a registry: `agent-env run` plans it, writes it into the `@local` namespace (reusing each version whose inputs haven't changed) and runs its evals or tasks; `agent-env bundle init` writes a stored task or eval as one, with every env, agent and artifact it uses in its own folder (each image as its build context, pruned to what its Dockerfile reads and pinned to the Python and npm versions of the stored image, or, given `--image-ref ID=REF`, as that registry image pinned by digest with its context beside it, `image = { ref, dockerfile }`, which a run pulls and builds only when it can't; `task_images` lists the images an import writes, downloading nothing, so an export can push them first) so it runs with no store; `agent-env bundle register` writes a folder into the configured stores under registry ids (``), each task and eval naming the versions it reads and each image it builds pushed, through a ledger in the target store so registering again writes only what changed. | | `eval/` | `Eval` (`eval`), a versioned list of task references, and its store. | @@ -347,6 +348,10 @@ the same triggers; it allows no skips and isn't a required check. environment, on the same triggers; it allows no skips and isn't a required check. `.github/workflows/skyrl.yml` runs `tst/skyrl` on SkyRL 0.4.0 and its main branch, with skyrl-gym and CPU torch installed over the dev environment on Python 3.12, on the same triggers; it allows no skips and isn't a required check. +`.github/workflows/miles.yml` runs `tst/miles` on Miles v0.1.3 and its main branch, checked out on `PYTHONPATH` with +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/plugin-api.yml` runs the `plugin-api` job on every pull request, title edits included, and on `main`: `.github/scripts/check_plugin_api.py` compares the plugin surface (its `BASES` and `USED`) between `HEAD^1` and `HEAD` with griffe, pinned in the `dev` extra, and fails on a break the diff --git a/Makefile b/Makefile index 1b557f90..54fef1e3 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test trl-test nemo-gym-test verl-test skyrl-test clean-install-test test clean +.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test trl-test nemo-gym-test verl-test skyrl-test miles-test clean-install-test test clean VENV := .venv VENV_BIN := $(abspath $(VENV))/bin @@ -85,6 +85,12 @@ verl-test: ## Play rows through the verl agent loop on verl itself (needs verl a skyrl-test: ## Play rows through the SkyRL env and generator on SkyRL itself (needs SkyRL and CPU torch, and the Hub) $(PYTHON) -m pytest tst/skyrl --timeout=900 $(PYTEST_ARGS) +# The Miles tier plays rows through Miles' own rollout: its agentic generate path, session server, router and mock +# SGLang server, with Qwen3's tokenizer from the Hub. Miles isn't on PyPI: put checkouts of Miles, SGLang and Megatron-LM +# on PYTHONPATH and install Miles' requirements and CPU torch first, as .github/workflows/miles.yml does. +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 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 44212efa..15115e45 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,6 +91,11 @@ skyrl = [ "pyarrow>=19", "transformers>=5", ] +# agent_env.training.miles: Miles' trainers on a bundle served as an OpenEnv env; installed into Miles' own environment, +# whose Miles, SGLang and torch are the trainer's to pin. It needs nothing beyond the openenv extra. +miles = [ + "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 ed79a783..0571146a 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, and ``agent_env.training.skyrl`` for -SkyRL.""" +environments, NeMo RL among them, ``agent_env.training.verl`` for verl, ``agent_env.training.skyrl`` for SkyRL, and +``agent_env.training.miles`` for Miles.""" diff --git a/src/agent_env/training/miles/__init__.py b/src/agent_env/training/miles/__init__.py new file mode 100644 index 00000000..133847d4 --- /dev/null +++ b/src/agent_env/training/miles/__init__.py @@ -0,0 +1,5 @@ +"""Miles' trainers (https://github.com/radixark/miles) on a bundle served as an OpenEnv env (``agent_env.openenv.build_app``, +or the package ``openenv import`` writes): ``python -m agent_env.training.miles URL OUT`` writes the server's tasks as +Miles' prompt data, and ``agent_env.training.miles.agent`` holds the agent function Miles calls for each sample, which +plays the row's task as one episode on the server through Miles' session server, with the reward hook and the +dynamic-sampling filter beside it.""" diff --git a/src/agent_env/training/miles/__main__.py b/src/agent_env/training/miles/__main__.py new file mode 100644 index 00000000..d63acd18 --- /dev/null +++ b/src/agent_env/training/miles/__main__.py @@ -0,0 +1,33 @@ +"""``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}") + + +if __name__ == "__main__": + main() diff --git a/src/agent_env/training/miles/agent.py b/src/agent_env/training/miles/agent.py new file mode 100644 index 00000000..61beba5f --- /dev/null +++ b/src/agent_env/training/miles/agent.py @@ -0,0 +1,113 @@ +"""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 +(``--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 + +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] diff --git a/src/agent_env/training/miles/prepare.py b/src/agent_env/training/miles/prepare.py new file mode 100644 index 00000000..fa41809d --- /dev/null +++ b/src/agent_env/training/miles/prepare.py @@ -0,0 +1,33 @@ +"""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 + +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 diff --git a/tst/miles/__init__.py b/tst/miles/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tst/miles/conftest.py b/tst/miles/conftest.py new file mode 100644 index 00000000..79113c9a --- /dev/null +++ b/tst/miles/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 Miles' +own rollout.""" + +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/miles/rollout_test.py b/tst/miles/rollout_test.py new file mode 100644 index 00000000..c09386eb --- /dev/null +++ b/tst/miles/rollout_test.py @@ -0,0 +1,113 @@ +"""Rows ``write_rows`` wrote, played by Miles' own rollout: its agentic generate path hands each sample to +``agent_env.training.miles.agent.run`` with a session on Miles' session server, in front of Miles' router and its mock +SGLang server, which plays the model from a script; Miles scores the samples with ``agent.reward`` and groups them +through ``agent.keep_scored_groups``.""" + +import json +from unittest.mock import patch + +import httpx +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.util.capabilities import MILES, missing_capability_reason + +fixtures = pytest.importorskip("tests.fast.fixtures.rollout_fixtures", reason=missing_capability_reason(MILES)) +integration = pytest.importorskip("tests.fast.rollout.inference_rollout.integration.utils", + reason=missing_capability_reason(MILES)) +arguments = pytest.importorskip("miles.utils.arguments", reason=missing_capability_reason(MILES)) +function_registry = pytest.importorskip("miles.utils.function_registry", reason=missing_capability_reason(MILES)) +http_utils = pytest.importorskip("miles.utils.http_utils", reason=missing_capability_reason(MILES)) +misc = pytest.importorskip("miles.utils.misc", reason=missing_capability_reason(MILES)) +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"]) +def session_server(request): + """Miles' session server, bare (v1, append-only) or ``v2`` (tree serving), which merges the agent function's + metadata into the samples itself.""" + return ["--use-session-server", *request.param] + + +def test_miles_plays_a_row_through_its_session_and_trains_on_the_episodes_reward(serve, tmp_path, session_server): + server = serve({"a": PLAYED}) + rows = write_rows(server.url, tmp_path / "rows") / "all.jsonl" + + played = _rollout(rows, session_server) + server.stop() + + [[sample]] = _groups(played) + assert (sample.status, sample.reward) == (sample_types.Sample.Status.COMPLETED, 1.0) + assert {key: sample.metadata[key] for key in ("agentenv_scores", "agentenv_tool_calls", "agentenv_invalid_calls")} \ + == {"agentenv_scores": {"replies": 1.0}, "agentenv_tool_calls": 1, "agentenv_invalid_calls": 0} + assert sample.response.startswith(ADDS_A_CRATE) and sample.response.endswith("Added a crate.") + + +def test_miles_leaves_out_a_group_the_env_failed_and_draws_rows_until_its_batch_is_full(serve, tmp_path, + session_server): + server = serve({"a": FAILS_AFTER_THE_TURN, "b": PLAYED}) + rows = write_rows(server.url, tmp_path / "rows") / "all.jsonl" + + played = _rollout(rows, session_server, groups=2) + server.stop() + + assert [[(sample.metadata["task_id"], sample.reward) for sample in group] for group in _groups(played)] == [ + [("b", 1.0)], [("b", 1.0)]] + + +def _rollout(rows, session_server, groups=1): + """A train rollout of ``groups`` groups of one sample from ``rows``, the model answering each chat request with + ``_model``.""" + misc.SingletonMeta.clear_all_instances() + try: + with mock_sglang.with_mock_server(process_fn=_model) as model: + args = _args(rows, session_server, groups, http_utils.find_available_port(20000)) + with fixtures._with_miles_router(args) as router: + httpx.post(f"{router.url}/add_worker", params={"url": model.url}, timeout=5.0).raise_for_status() + with fixtures._with_session_server(args, router.url): + data_source = function_registry.load_function(args.data_source_path)(args) + return integration.load_and_call_rollout(args, data_source, "train") + finally: + misc.SingletonMeta.clear_all_instances() + + +def _groups(played): + """The rollout's groups, each a flat list of samples: v2 hands back each generate call's samples as a list.""" + return [[sample for each in group for sample in (each if isinstance(each, list) else [each])] + for group in played.samples] + + +def _model(prompt): + """Adds a crate, then says so once the tool has answered.""" + return mock_sglang.ProcessResult(text="Added a crate." if "" in prompt else ADDS_A_CRATE) + + +def _args(rows, session_server, groups, router_port): + """The arguments Miles' own rollout tests start from, with the helper's agent function and hooks in place of a + label and a reward type.""" + argv = [ + "pytest", "--train-backend", "fsdp", "--ci-test", "--no-enable-sample-ownership-checker", + "--hf-checkpoint", "Qwen/Qwen3-0.6B", "--rollout-num-gpus", "1", "--rollout-num-gpus-per-engine", "1", + "--num-rollout", "1", "--rollout-batch-size", str(groups), "--n-samples-per-prompt", "1", + "--rollout-max-response-len", "256", + "--prompt-data", str(rows), "--input-key", "messages", "--metadata-key", "metadata", + "--apply-chat-template-kwargs", json.dumps({"enable_thinking": False}), + "--use-miles-router", "--sglang-router-ip", "127.0.0.1", "--sglang-router-port", str(router_port), + *integration.MODULAR_ROLLOUT_BASE_ARGV, + "--custom-generate-function-path", "miles.rollout.generate_hub.agentic_tool_call.generate", + "--custom-agent-function-path", "agent_env.training.miles.agent.run", + *session_server, "--tito-model", "qwen3", + "--custom-rm-path", "agent_env.training.miles.agent.reward", + "--dynamic-sampling-filter-path", "agent_env.training.miles.agent.keep_scored_groups", + ] + with patch("sys.argv", argv): + args = arguments.parse_args() + http_utils.init_http_client(args) + return args diff --git a/tst/unit/test_optional_extras_stay_optional.py b/tst/unit/test_optional_extras_stay_optional.py index 1b441bef..9254ce45 100644 --- a/tst/unit/test_optional_extras_stay_optional.py +++ b/tst/unit/test_optional_extras_stay_optional.py @@ -32,7 +32,7 @@ def find_spec(self, fullname, path=None, target=None): _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.harness", "agent_env.training.miles", "agent_env.env.gateway", "agent_env.task_step.task_steps.run_code_runner", ) @@ -174,7 +174,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"} +_COVERED = {"verl": "trl", "skyrl": "trl", "harness": "openenv", "miles": "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 new file mode 100644 index 00000000..d5559f6c --- /dev/null +++ b/tst/unit/training/miles_test.py @@ -0,0 +1,182 @@ +"""``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.""" + +import asyncio +import json +from types import SimpleNamespace + +import pytest +from click.testing import CliRunner +from starlette.applications import Starlette +from starlette.responses import JSONResponse +from starlette.routing import Route + +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.agent import keep_scored_groups, reward, run +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 + +ADD_ITEM = {"type": "function", "function": {"name": "add_item", "description": "add_item.", + "parameters": {"type": "object"}}} +FUNCTION_END_TURN = {"type": "function", "function": {"name": "end_turn", "description": END_TURN.description, + "parameters": END_TURN.input_schema}} +SAMPLING = {"temperature": 1.0, "top_p": 1.0, "max_tokens": 256} + + +def _call(arguments, id="call_1", name="add_item"): + return {"role": "assistant", "content": None, "tool_calls": [ + {"id": id, "type": "function", "function": {"name": name, "arguments": arguments}}]} + + +def _says(text): + return {"role": "assistant", "content": text} + + +@pytest.fixture +def session(): + """Miles' session endpoint, answering the n-th chat request with ``script[n]``, the last again after it ends, with + the HTTP ``status`` after ``delay`` seconds; ``.requests`` are the bodies it got.""" + servers = [] + + def serving(*script, status=200, delay=0.0): + requests = [] + + async def completions(request): + requests.append(await request.json()) + await asyncio.sleep(delay) + message = script[min(len(requests), len(script)) - 1] + return JSONResponse({"choices": [{"index": 0, "message": message}]}, status_code=status) + + server = Server(Starlette(routes=[Route("/sessions/s1/v1/chat/completions", completions, methods=["POST"])])) + server.requests, server.base_url = requests, f"{server.url}/sessions/s1" + servers.append(server) + return server + + yield serving + for server in servers: + server.stop() + + +def _row(server, tmp_path, **options): + folder = write_rows(server.url, tmp_path / "rows", **options) + return json.loads((folder / "all.jsonl").read_text().splitlines()[0]) + + +def test_a_servers_tasks_become_a_row_per_task_of_each_split(serve, tmp_path, teardowns, monkeypatch): + monkeypatch.setenv("DOCK_TOKEN", "unread") + server = serve({"a": PLAYED, "b": [ITEMS, ALICE, {**TURN, "prompt": "Add a barrel."}, REPLIES]}, {"smoke": ["b"]}, + name="Dock Yard") + + folder = write_rows(f"{server.url}/", tmp_path / "rows", tool_timeout=30.0, reset_timeout=90.0, model_timeout=120.0, + max_turns=3, token_var="DOCK_TOKEN") + server.stop() + + [row] = [json.loads(line) for line in (folder / "smoke.jsonl").read_text().splitlines()] + assert sorted(path.name for path in folder.iterdir()) == ["all.jsonl", "smoke.jsonl"] + assert row == { + "messages": [{"role": "system", "content": "You keep the items."}, {"role": "user", "content": "Add a barrel."}], + "metadata": {"split": "smoke", "index": 0, "task_id": "b", "tools": [ADD_ITEM, FUNCTION_END_TURN], + "agentenv": {"url": server.url, "tool_timeout": 30.0, "reset_timeout": 90.0, + "model_timeout": 120.0, "max_turns": 3, "token_var": "DOCK_TOKEN"}}, + } + assert [json.loads(line)["metadata"]["task_id"] for line in (folder / "all.jsonl").read_text().splitlines()] == [ + "a", "b"] + assert len(teardowns.finished) == 1 + + +def test_the_agent_plays_a_rows_task_through_the_session_and_returns_its_reward(serve, session, tmp_path): + server = serve({"a": PLAYED}) + 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"])) + server.stop() + + assert played == {"reward": 1.0, "agentenv_scores": {"replies": 1.0}, "agentenv_tool_calls": 1, + "agentenv_invalid_calls": 0} + first, second = model.requests + assert first == {"model": "default", "messages": row["messages"], "tools": row["metadata"]["tools"], **SAMPLING} + assert second["messages"] == [*row["messages"], _call('{"item": "crate"}'), + {"role": "tool", "tool_call_id": "call_1", "content": second["messages"][-1]["content"]}] + assert second["messages"][-1]["content"] + + +def test_a_call_whose_arguments_arent_an_object_is_answered_and_the_episode_goes_on(serve, session, tmp_path): + server = serve({"a": PLAYED}) + row = _row(server, tmp_path) + 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"])) + server.stop() + + assert (played["reward"], played["agentenv_tool_calls"], played["agentenv_invalid_calls"]) == (1.0, 0, 3) + answers = [message["content"] for message in model.requests[-1]["messages"] if message["role"] == "tool"] + assert answers[0] == "add_item wasn't called: its arguments must be a JSON object, not str" + assert answers[1].startswith("add_item wasn't called: its arguments aren't JSON") + assert answers[2].startswith("add_item wasn't called: its arguments aren't JSON") + assert "agentenv_failure" not in played + + +def test_a_rollout_that_reaches_its_turn_cap_ends_the_turn_without_a_reply(serve, session, tmp_path): + server = serve({"a": PLAYED}) + 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"])) + server.stop() + + assert (played["reward"], played["agentenv_tool_calls"], len(model.requests)) == (0.0, 2, 2) + + +@pytest.mark.parametrize(("failing", "why"), [("env", "couldn't open a session on the server"), + ("session", "HTTPStatusError: Server error '503"), + ("model", "ReadTimeout"), ("reward", "the episode ended without a reward")]) +def test_an_episode_that_fails_or_ends_without_a_reward_is_flagged_and_its_group_left_out(serve, session, tmp_path, + failing, why): + server = serve({"a": PLAYED}, **({"reward": _unscored} if failing == "reward" else {})) + row = _row(server, tmp_path, reset_timeout=5.0, **({"model_timeout": 0.5} if failing == "model" else {})) + model = session(_says("Added a crate."), status=503 if failing == "session" else 200, + delay=2.0 if failing == "model" else 0.0) + if failing == "env": + server.stop() + + played = asyncio.run(run(model.base_url, row["messages"], SAMPLING, row["metadata"])) + server.stop() + + assert played["reward"] == 0.0 and played["agentenv_failure"].startswith(why) + failed, scored = _sample({**played}), _sample({"reward": 1.0}) + assert keep_scored_groups(None, [scored, failed]) is False + assert keep_scored_groups(None, [[scored], [_sample({"reward": 0.0})]]) is True + assert keep_scored_groups(None, [scored, _sample({"reward": 0.0}, status="ABORTED")]) is False + assert keep_scored_groups(None, [scored, _sample({"session_id": "s1"})]) is False + + +def test_the_reward_hook_reads_each_samples_reward_from_its_metadata(): + one, two = _sample({"reward": 0.25}), _sample({}) + + assert asyncio.run(reward(None, one)) == 0.25 + assert asyncio.run(reward(None, [one, two])) == [0.25, 0.0] + + +def test_the_command_writes_the_rows(serve, tmp_path, teardowns): + server = serve({"a": PLAYED}) + + result = CliRunner().invoke(main, [server.url, str(tmp_path / "rows"), "--max-turns", "2"]) + server.stop() + + assert result.exit_code == 0, result.output + assert sorted(path.name for path in (tmp_path / "rows").iterdir()) == ["all.jsonl"] + [line] = (tmp_path / "rows" / "all.jsonl").read_text().splitlines() + assert json.loads(line)["metadata"]["agentenv"]["max_turns"] == 2 + assert len(teardowns.finished) == 1 + + +def _sample(metadata, status="COMPLETED"): + return SimpleNamespace(metadata=metadata, status=SimpleNamespace(name=status)) diff --git a/tst/util/capabilities.py b/tst/util/capabilities.py index 2eddbb72..b75bff11 100644 --- a/tst/util/capabilities.py +++ b/tst/util/capabilities.py @@ -53,6 +53,10 @@ #: and torch aren't in the dev extra, and only that tier's CI job installs them. SKYRL = "skyrl_installed" +#: The Miles tier (``tst/miles``) plays rows through Miles' own rollout; Miles, its test fixtures and CPU torch aren't in +#: the dev extra, and only that tier's CI job puts a Miles checkout on PYTHONPATH and installs its requirements. +MILES = "miles_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 36c341d6..fafd9862 100644 --- a/uv.lock +++ b/uv.lock @@ -140,6 +140,9 @@ gcp = [ harness = [ { name = "openenv" }, ] +miles = [ + { name = "openenv" }, +] nemo-gym = [ { name = "nemo-gym", marker = "python_full_version >= '3.13.14'" }, { name = "omegaconf", marker = "python_full_version >= '3.13.14'" }, @@ -189,6 +192,7 @@ requires-dist = [ { name = "a2a-sdk", specifier = "==0.3.26" }, { name = "agentenv-framework", extras = ["aws", "boat", "boxd", "coreweave", "daytona", "gcp", "harness", "openenv", "sail", "skyrl", "tensorlake", "tenki", "trl", "vercel", "verl"], marker = "extra == 'dev'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'harness'" }, + { 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 == 'trl'" }, @@ -255,7 +259,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", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] +provides-extras = ["openenv", "harness", "trl", "nemo-gym", "verl", "skyrl", "miles", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] [[package]] name = "agentenv-framework-protocol"