From cb917e0b0070be5e80685f93402e56e967915cd1 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 15:48:43 -0700 Subject: [PATCH 1/4] fix(openenv): a reset or end_turn past its deadline ends the episode, and CI runs TRL The server ended a tool call that missed its deadline, but waited on a reset's deploy and on end_turn with none, so a client that had gone away left its session open and its sandboxes up. reset_async now takes timeout_s and end_turn honors timeout_s metadata; the TRL helper sends reset_timeout with both. The new tst/trl tier, run by .github/workflows/trl.yml against TRL 1.15 and the latest, trains one GRPOTrainer step on CPU through the helper. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/trl.yml | 56 +++++++++++++++++++++++ AGENTS.md | 6 ++- CONTRIBUTING.md | 1 + Makefile | 7 ++- src/agent_env/openenv/environment.py | 59 +++++++++++++----------- src/agent_env/training/trl.py | 35 ++++++++------- tst/trl/__init__.py | 0 tst/trl/conftest.py | 7 +++ tst/trl/grpo_test.py | 67 ++++++++++++++++++++++++++++ tst/unit/openenv/environment_test.py | 18 +++++++- tst/unit/training/trl_test.py | 22 +++++++-- tst/util/capabilities.py | 4 ++ 12 files changed, 235 insertions(+), 47 deletions(-) create mode 100644 .github/workflows/trl.yml create mode 100644 tst/trl/__init__.py create mode 100644 tst/trl/conftest.py create mode 100644 tst/trl/grpo_test.py diff --git a/.github/workflows/trl.yml b/.github/workflows/trl.yml new file mode 100644 index 000000000..fdcaec2aa --- /dev/null +++ b/.github/workflows/trl.yml @@ -0,0 +1,56 @@ +name: TRL contract + +# One GRPOTrainer step on CPU through agent_env.training.trl (tst/trl), against the oldest TRL the helper supports and +# the latest. TRL marks environment_factory experimental, so a release can break the helper with nothing else in CI +# noticing. It runs on pull requests that touch the helper, the OpenEnv server it trains against 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/trl/**" + - "tst/unit/training/**" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/trl.yml" + schedule: + - cron: "23 7 * * 1" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: trl-${{ 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: + grpo: + name: grpo (${{ matrix.trl }}) + runs-on: ubuntu-latest + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + trl: ["trl==1.15.*", "trl"] + 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, then TRL and CPU torch + run: | + python -m pip install --quiet -r .github/ci-requirements.txt + uv sync --locked --extra dev --python "$(command -v python)" + uv pip install --python .venv/bin/python --torch-backend cpu torch "${{ matrix.trl }}" + .venv/bin/python -c "import torch, transformers, trl; print('trl', trl.__version__, 'transformers', transformers.__version__, 'torch', torch.__version__)" + - name: TRL contract tier + run: .venv/bin/python -m pytest tst/trl -rA --timeout=900 --junitxml=junit-trl.xml + - name: No skips + if: always() + run: .venv/bin/python .github/scripts/check_skip_policy.py junit-trl.xml diff --git a/AGENTS.md b/AGENTS.md index f07305138..b0e459606 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -44,6 +44,7 @@ or one of them only where it is used. | `make int-test-fast` | `tst/integration` minus `int_test_slow`, in parallel; needs Docker and a local OCI registry (`docker run -d -p 5000:5000 public.ecr.aws/docker/library/registry:2`) | | `make int-test-slow` | the `int_test_slow` tests, serially (they build the same Docker tags from module fixtures and race under xdist) | | `make installer-test` | `tst/installer`: `plugin add` / `remove` through the real pip, uv and pipx, offline against wheels built from the checkout, and the container user journey (`python:3.12-slim`, `--network none`) with the PEP 668 refusal; needs uv, pipx and Docker | +| `make trl-test` | `tst/trl`: one `GRPOTrainer` step on CPU through `agent_env.training.trl`, with a tiny model from the Hub; needs trl and CPU torch, which the `dev` extra leaves out | | `make 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 | @@ -266,7 +267,7 @@ orchestrators call it the same way. What callers rely on: - Skips are declared capability gaps, decided at collection time: the gates in `tst/util/capabilities.py` produce `skipif` marks with the reason `agentenv-capability-missing: ` (`model_endpoint_configured`, `remote_sandbox`, - `default_a2a_agent`, `mcp_server_sources`, `gnu_stat`). CI rejects any other skip reason; a broken configuration must fail loudly, not skip. + `default_a2a_agent`, `mcp_server_sources`, `gnu_stat`, `trl_installed`). CI rejects any other skip reason; a broken configuration must fail loudly, not skip. - The suite is backend-agnostic: a test about one particular backend lives with that backend, not here. A custom store must pass `tst/store/{conformance,object_conformance,image_conformance,secret_conformance}.py`, the same kits the built-ins pass, and a custom `[runner]` must pass `tst/runner/conformance.py`, which @@ -300,6 +301,9 @@ skips for the allowed capabilities pass, and the fast job allows no skips under touches the plugin CLI, `agent_env/plugins` or `agent_env/config`, the protocol package, the tier itself, `pyproject.toml`, `uv.lock`, the `Makefile` or the CI that runs the tier, and allows no skips; a `changes` job decides, and a skipped `installer` counts as passed. +`.github/workflows/trl.yml` runs `tst/trl` against the oldest TRL the training helper supports and the latest, on pull +requests that touch `agent_env/training`, `agent_env/openenv` or the tier, weekly, and on demand; it allows no skips +and isn't a required check. `.github/workflows/plugin-api.yml` runs the `plugin-api` job on every pull request, title edits included, and on `main`: `.github/scripts/check_plugin_api.py` compares the plugin surface (its `BASES` and `USED`) between `HEAD^1` and `HEAD` with griffe, pinned in the `dev` extra, and fails on a break the diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 2a1fba54c..b97278465 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -43,6 +43,7 @@ reinstalling. | `make int-test-slow` | tests that build images or deploy sandboxes, minutes each | Docker, and a sandbox backend for some | | `make clean-install-test` | both distributions built as the release builds them, installed into a fresh venv from public PyPI, and `agent-env run hello` run twice by name | Python 3.11, uv and git | | `make installer-test` | `plugin add` and `remove` through the real pip, uv and pipx, offline against wheels built from the checkout, and the new-user journey in containers with no network | uv, pipx and Docker; the first run downloads the dependencies once | +| `make trl-test` | one TRL `GRPOTrainer` step on CPU through the training helper, with a tiny model | trl and CPU torch installed, and the Hugging Face Hub | CI runs the same tiers as the `unit`, `integration-local`, `integration-local-slow` and `integration-local-gateway` jobs (the last runs the slow tier's gateway tests beside the rest); diff --git a/Makefile b/Makefile index 189b9d793..b767155b8 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 clean-install-test test clean +.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test trl-test clean-install-test test clean VENV := .venv VENV_BIN := $(abspath $(VENV))/bin @@ -64,6 +64,11 @@ int-test-slow: ## Run only the @int_test_slow integration tests (serial to avoid installer-test: ## Run plugin add/remove through the real pip, uv and pipx, and the container journey (needs uv, pipx, Docker) $(PYTHON) -m pytest tst/installer -n auto --timeout=900 $(PYTEST_ARGS) +# The TRL contract tier trains one GRPOTrainer step on CPU through agent_env.training.trl, with a tiny model from the +# Hub. trl and torch aren't in the dev extra: install them first, as .github/workflows/trl.yml does. +trl-test: ## Run one GRPOTrainer step on CPU through the TRL helper (needs trl and CPU torch, and the Hub) + $(PYTHON) -m pytest tst/trl --timeout=900 $(PYTEST_ARGS) + # The clean-install gate CI runs, on the same Python. The script gives every step an allowlisted # environment, so nothing in this shell reaches the install it checks. A failed run keeps its folder. clean-install-test: ## Build both distributions, install them clean from public PyPI and run hello twice (needs python3.11, uv, git) diff --git a/src/agent_env/openenv/environment.py b/src/agent_env/openenv/environment.py index 2d0a09290..74750b834 100644 --- a/src/agent_env/openenv/environment.py +++ b/src/agent_env/openenv/environment.py @@ -36,8 +36,9 @@ input_schema={"type": "object", "properties": {"response": {"type": "string", "description": "Your reply."}}}, ) ENDED = "the episode has ended; call reset" -#: The key of a tool call's deadline in seconds in its action's metadata. A call that runs past it ends the episode, since -#: cancelling the call doesn't stop what it does, and the session goes on. +#: The key of a call's deadline in seconds in its action's metadata, ``end_turn``'s included; ``reset_async`` takes one as +#: ``timeout_s``. A call that runs past it ends the episode, since cancelling the call doesn't stop what it does, and the +#: session goes on; so no call holds a session longer than its client waits. TIMEOUT_METADATA = "timeout_s" NOT_RESET = "no episode has started; call reset before stepping" @@ -137,13 +138,13 @@ def get_task_range(self, split: str, start: int | None = None, stop: int | None # The episode async def reset_async(self, seed: int | None = None, episode_id: str | None = None, *, split: str | None = None, - index: int | None = None, **kwargs: Any) -> Observation: + index: int | None = None, timeout_s: float | None = None, **kwargs: Any) -> Observation: """Plays task ``index`` of ``split``, or without either the first task served, as OpenEnv's ORS and verifiers - importers do.""" + importers do, and ends the episode if its first turn hasn't opened ``timeout_s`` seconds in.""" async with self._lock: - return await self._reset(split, index) + return await self._reset(split, index, timeout_s) - async def _reset(self, split: str | None, index: int | None) -> Observation: + async def _reset(self, split: str | None, index: int | None, timeout_s: float | None) -> Observation: await self._abandon("reset") self._loop = asyncio.get_running_loop() self._served_task, self._state = None, State(episode_id=None, step_count=0, done=False) @@ -158,15 +159,18 @@ async def _reset(self, split: str | None, index: int | None) -> Observation: self._state = State(episode_id=None, step_count=0, task_id=self._served_task.task.id, task_name=self._served_task.name, split=split, index=index, turn=None, done=False) self._served.live.add(self) + deadline = asyncio.timeout(timeout_s) try: - self._episode = await self._served_task.spec.start(sandbox=self._served.sandbox) - self._state.episode_id = self._episode.id - turn = await self._episode.next_turn() - if turn is None: - raise EpisodeError("the run ended before the agent's first turn") - await self._open(turn) + async with deadline: + self._episode = await self._served_task.spec.start(sandbox=self._served.sandbox) + self._state.episode_id = self._episode.id + turn = await self._episode.next_turn() + if turn is None: + raise EpisodeError("the run ended before the agent's first turn") + await self._open(turn) except Exception as e: - return await self._end_with_error(e) + return await self._end_with_error( + f"the episode didn't open its first turn within {timeout_s:g}s" if deadline.expired() else e) return Observation(done=False, reward=None, metadata={ "prompt": turn.prompt, "system_prompt": turn.system_prompt, "task": self._task_metadata()}) @@ -189,9 +193,11 @@ async def _step(self, action: Any, timeout_s: float | None) -> Observation: if (stopped := await self._why_stopped()) is not None: return CallToolObservation(tool_name=action.tool_name, done=True, reward=None, metadata=stopped) self._state.step_count += 1 + if timeout_s is None: + timeout_s = _deadline_of(action) if action.tool_name == END_TURN.name: - return await self._end_turn(action.tool_name, action.arguments.get("response") or "") - return await self._call(action, timeout_s if timeout_s is not None else _deadline_of(action)) + return await self._end_turn(action.tool_name, action.arguments.get("response") or "", timeout_s) + return await self._call(action, timeout_s) def step(self, action: Any, timeout_s: float | None = None, **kwargs: Any) -> Observation: """OpenEnv's web playground steps from a worker thread; the episode lives on the loop its reset ran on.""" @@ -234,18 +240,21 @@ async def _call(self, action: CallToolAction, timeout_s: float | None) -> CallTo return CallToolObservation(tool_name=name, result={ "content": result.content, "structured_content": None, "meta": None, "data": result.text, "is_error": False}) - async def _end_turn(self, name: str, response: str) -> CallToolObservation: + async def _end_turn(self, name: str, response: str, timeout_s: float | None) -> CallToolObservation: + deadline = asyncio.timeout(timeout_s) try: - await self._turn.end(response) - await self._close_tools() - turn = await self._episode.next_turn() - if turn is not None: - await self._open(turn) - return CallToolObservation(tool_name=name, done=False, reward=None, result={ - "prompt": turn.prompt, "system_prompt": turn.system_prompt}) - scores = await self._episode.result() + async with deadline: + await self._turn.end(response) + await self._close_tools() + turn = await self._episode.next_turn() + if turn is not None: + await self._open(turn) + return CallToolObservation(tool_name=name, done=False, reward=None, result={ + "prompt": turn.prompt, "system_prompt": turn.system_prompt}) + scores = await self._episode.result() except Exception as e: - observation = await self._end_with_error(e) + observation = await self._end_with_error( + f"{name} didn't finish within {timeout_s:g}s, and the episode ended" if deadline.expired() else e) return CallToolObservation(tool_name=name, done=True, reward=None, metadata=observation.metadata) await self._abandon("the episode ended") return CallToolObservation(tool_name=name, done=True, **self._scored(scores)) diff --git a/src/agent_env/training/trl.py b/src/agent_env/training/trl.py index 83efa5302..21710f739 100644 --- a/src/agent_env/training/trl.py +++ b/src/agent_env/training/trl.py @@ -47,7 +47,7 @@ END_TURN = END_TURN_TOOL.name CONNECT_SECONDS = 60.0 CLOSE_SECONDS = 300.0 -#: How much longer than a tool's deadline the client waits for the server to end the call, before dropping the session. +#: How much longer than a call's deadline the client waits for the server to end the call, before dropping the session. GRACE_SECONDS = 30.0 #: How long a server ``AgentEnv.local`` started is given to stop, at ``close`` and at exit: its episodes' teardowns run #: at once, and each saves the agents' changelogs, ends open recordings and terminates the sandboxes, every part within @@ -83,7 +83,7 @@ class AgentEnv: tool_timeout: One tool call's deadline in seconds. The server ends a call that misses it, and the episode with it, since what the call does may still be under way: the rollout is left unscored, and the session kept. reset_timeout: How long an episode may take to deploy and open its first turn, and each turn to be ended, - in seconds. + in seconds. The server ends an episode that misses it, so the rollout is left unscored. Example: @@ -167,7 +167,7 @@ async def _environment(self, toolset: _Toolset) -> type[_Environment]: client = await _opened(self.url, self._reset_timeout) first, tools = None, {} try: - turn = await client.reset(split=split, index=index) + turn = await client.reset(split=split, index=index, timeout_s=self._reset_timeout) for number in range(turns): if turn.done: break @@ -178,7 +178,7 @@ async def _environment(self, toolset: _Toolset) -> type[_Environment]: first = first or frozenset(tool["name"] for tool in listed.observation["tools"]) tools |= {tool["name"]: tool for tool in listed.observation["tools"] if tool["name"] not in tools} if number + 1 < turns: - turn = await _end_turn(client, "") + turn = await _end_turn(client, "", self._reset_timeout) if error := turn.observation.get("error"): raise RuntimeError(f"task {index} of split {split!r} refused an empty {END_TURN}: " f"{error.get('message')}") @@ -286,8 +286,11 @@ def get_reward(self) -> float | None: return rollout.reward def _call(self, name: str, arguments: dict[str, Any]) -> str: - return _on_loop(self._step(name, arguments), - self._shared.reset_timeout + self._shared.tool_timeout + GRACE_SECONDS + CLOSE_SECONDS) + """Waits at most for the deploy and the call, each to its deadline and a grace after, then the session's close: + the first call of a rollout waits for its deploy.""" + shared = self._shared + return _on_loop(self._step(name, arguments), shared.reset_timeout + max(shared.reset_timeout, shared.tool_timeout) + + 2 * GRACE_SECONDS + CLOSE_SECONDS) async def _reset(self, split: str, index: int, expected: str | None) -> None: async with self._lock: @@ -304,7 +307,7 @@ async def _reset(self, split: str, index: int, expected: str | None) -> None: async def _begin(self, rollout: _Rollout, split: str, index: int, expected: str | None) -> None: if self._client is None: self._client = await _reopened(self._shared.url, self._shared.reset_timeout) - reset = await self._client.reset(split=split, index=index) + reset = await self._client.reset(split=split, index=index, timeout_s=self._shared.reset_timeout) if reset.done: raise _Failed(_error_of(reset.observation)) metadata = reset.observation.get("metadata") or {} @@ -327,16 +330,16 @@ async def _step(self, name: str, arguments: dict[str, Any]) -> str: return ENDED try: await rollout.started - deadline = None if name == END_TURN else self._shared.tool_timeout + deadline = self._shared.reset_timeout if name == END_TURN else self._shared.tool_timeout result = await asyncio.wait_for(_call_tool(self._client, name, arguments, deadline), - None if deadline is None else deadline + GRACE_SECONDS) + deadline + GRACE_SECONDS) except _Failed as e: rollout.misgrouped = isinstance(e, _Misgrouped) await self._fail(rollout, str(e)) return FAILED except TimeoutError: - await self._fail(rollout, f"{name} didn't answer within {self._shared.tool_timeout:g}s, and the " - f"server didn't end the call in the {GRACE_SECONDS:g}s after") + await self._fail(rollout, f"{name} didn't answer within {deadline:g}s, and the server didn't end " + f"the call in the {GRACE_SECONDS:g}s after") return FAILED except Exception as e: await self._fail(rollout, f"{name}: {_why(e)}") @@ -351,7 +354,7 @@ async def _finish(self, rollout: _Rollout) -> None: try: await rollout.started while not (rollout.done or rollout.failure): - result = await _end_turn(self._client, "") + result = await _end_turn(self._client, "", self._shared.reset_timeout) if error := result.observation.get("error"): rollout.failure = f"{END_TURN}: {error.get('message')}" self._answer(rollout, END_TURN, result) @@ -487,8 +490,8 @@ def _close_all() -> None: async def _opened(url: str, timeout: float) -> GenericEnvClient: """A session on the server, opened now rather than at the first reset, so a server at capacity refuses it while - the trainer sets up.""" - client = GenericEnvClient(base_url=url, message_timeout_s=timeout) + the trainer sets up; each message waits ``GRACE_SECONDS`` longer than the ``timeout`` the server ends a call at.""" + client = GenericEnvClient(base_url=url, message_timeout_s=timeout + GRACE_SECONDS) try: await client.state() except Exception as e: @@ -526,8 +529,8 @@ async def _call_tool(client: GenericEnvClient, name: str, arguments: dict[str, A **({"metadata": {TIMEOUT_METADATA: deadline}} if deadline is not None else {})}) -async def _end_turn(client: GenericEnvClient, response: str) -> Any: - return await _call_tool(client, END_TURN, {"response": response}) +async def _end_turn(client: GenericEnvClient, response: str, deadline: float) -> Any: + return await _call_tool(client, END_TURN, {"response": response}, deadline) def _read_tasks(url: str) -> tuple[str, dict[str, list[dict[str, Any]]]]: diff --git a/tst/trl/__init__.py b/tst/trl/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tst/trl/conftest.py b/tst/trl/conftest.py new file mode 100644 index 000000000..38e814d11 --- /dev/null +++ b/tst/trl/conftest.py @@ -0,0 +1,7 @@ +"""The unit tier's bundle served on loopback over the Episode API's fake steps, without that tier's network block: +TRL downloads its tiny test model from the Hub.""" + +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.trl_test import served # noqa: F401 diff --git a/tst/trl/grpo_test.py b/tst/trl/grpo_test.py new file mode 100644 index 000000000..af1918366 --- /dev/null +++ b/tst/trl/grpo_test.py @@ -0,0 +1,67 @@ +"""TRL's GRPOTrainer trains a step through ``agent_env.training.trl`` on CPU, against the OpenEnv server ``build_app`` +makes over the Episode API's fake steps. TRL marks ``environment_factory`` experimental, so this pins what the helper +relies on: TRL resets an environment with its row, calls its tools, reads ``get_reward``, and trains on the rewards. +As in TRL's own environment tests, a tiny model's ``generate`` is stood in for. CI runs it against the oldest TRL the +helper supports and the latest; AsyncGRPO needs a vLLM server, so it isn't covered here.""" + +from unittest.mock import patch + +import pytest + +from tst.unit.training.trl_test import PLAYED +from tst.util.capabilities import TRL, missing_capability_reason + +torch = pytest.importorskip("torch", reason=missing_capability_reason(TRL)) +trl = pytest.importorskip("trl", reason=missing_capability_reason(TRL)) + +MODEL = "trl-internal-testing/tiny-Qwen3MoeForCausalLM" +ADD = '\n{"name": "add_item", "arguments": {"item": "crate"}}\n' +END = '\n{"name": "end_turn", "arguments": {"response": "Added a crate."}}\n' + + +def test_grpotrainer_trains_a_step_on_a_served_bundle_and_scores_every_rollout(served, teardowns, tmp_path, + monkeypatch): + monkeypatch.setenv("TRL_EXPERIMENTAL_SILENCE", "1") + env = served({"a": PLAYED, "b": PLAYED}) + rewards = [] + + class Recorded(env.factory): + def get_reward(self): + rewards.append(super().get_reward()) + return rewards[-1] + + trainer = trl.GRPOTrainer(model=MODEL, train_dataset=env.dataset, environment_factory=Recorded, args=trl.GRPOConfig( + output_dir=str(tmp_path), learning_rate=0.1, per_device_train_batch_size=4, num_generations=2, max_steps=1, + max_tool_calling_iterations=4, shuffle_dataset=False, logging_steps=1, report_to="none")) + before = {name: param.clone() for name, param in trainer.model.named_parameters()} + + with patch.object(trainer.model, "generate", side_effect=_Policy(trainer.processing_class)): + trainer.train() + + [logged] = [row for row in trainer.state.log_history if "tools/call_frequency" in row] + assert sorted(rewards) == [0.0, 0.0, 1.0, 1.0] + assert (logged["tools/call_frequency"], logged["tools/failure_frequency"], logged["rewards/Recorded/mean"]) == ( + pytest.approx(1.0), 0.0, pytest.approx(0.5)) + assert any(not torch.equal(param, trainer.model.get_parameter(name)) for name, param in before.items()) + env.close() + assert len(teardowns.finished) == len(teardowns.started) + + +class _Policy: + """Stands in for ``generate``: of each group, the first rollout adds a crate, then ends its turn, then stops, and the + second answers without a tool, so its turn is ended empty and scores 0.""" + + def __init__(self, tokenizer): + self.tokenizer = tokenizer + self.calls = 0 + + def __call__(self, input_ids, **_): + self.calls += 1 + if self.calls == 1: + texts = [ADD if row % 2 == 0 else "I won't." for row in range(len(input_ids))] + else: + texts = [END if self.calls == 2 else "Done."] * len(input_ids) + completions = [self.tokenizer(text + "<|im_end|>", add_special_tokens=False).input_ids for text in texts] + width = max(map(len, completions)) + padded = [ids + [self.tokenizer.pad_token_id] * (width - len(ids)) for ids in completions] + return torch.cat([input_ids, torch.tensor(padded, device=input_ids.device)], dim=-1) diff --git a/tst/unit/openenv/environment_test.py b/tst/unit/openenv/environment_test.py index ffe1abb5f..a8b893a21 100644 --- a/tst/unit/openenv/environment_test.py +++ b/tst/unit/openenv/environment_test.py @@ -13,7 +13,7 @@ import agent_env.task.episode.run as episode_run from agent_env.openenv.app import _score_at -from agent_env.openenv.environment import ENDED, NOT_RESET, AgentEnvEnvironment, Served, ServedTask +from agent_env.openenv.environment import ENDED, NOT_RESET, TIMEOUT_METADATA, AgentEnvEnvironment, Served, ServedTask from agent_env.task import EpisodeError, EpisodeSpec, Task from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN, _gate, _until @@ -229,6 +229,22 @@ async def test_a_run_that_fails_mid_turn_ends_the_episode_at_the_next_step(teard assert (await _call(env, "add_item")).metadata["error"] == ENDED +@pytest.mark.asyncio +async def test_a_reset_or_an_end_turn_past_its_deadline_ends_the_episode_in_an_error_observation(teardowns, tools): + held = {"id": "held", "type": "episode_test_gate"} + deploying, turning = _env(ITEMS, held, ALICE, TURN, REPLIES), _env(ITEMS, ALICE, TURN, held, REPLIES) + reset = await deploying.reset_async(timeout_s=0.5) + await turning.reset_async() + + ended = await turning.step_async(CallToolAction(tool_name="end_turn", arguments={"response": "Added a crate."}, + metadata={TIMEOUT_METADATA: 0.5})) + + assert (reset.done, reset.metadata["error"]) == (True, "the episode didn't open its first turn within 0.5s") + assert (ended.done, ended.reward, ended.metadata["error"]) == ( + True, None, "end_turn didn't finish within 0.5s, and the episode ended") + assert len(teardowns.finished) == 2 + + @pytest.mark.asyncio async def test_an_end_turn_the_proxy_refuses_ends_the_episode_in_an_error_observation(teardowns, tools, monkeypatch): env = _env(ITEMS, ALICE, TURN, REPLIES) diff --git a/tst/unit/training/trl_test.py b/tst/unit/training/trl_test.py index 0821abe8d..de9f2f328 100644 --- a/tst/unit/training/trl_test.py +++ b/tst/unit/training/trl_test.py @@ -16,7 +16,7 @@ from agent_env.config.runtime import Config from agent_env.openenv.app import _app, _score_at -from agent_env.openenv.environment import Served, ServedTask +from agent_env.openenv.environment import AgentEnvEnvironment, Served, ServedTask from agent_env.task import EpisodeSpec, Task from agent_env.task_step.task_step import TaskStep from agent_env.training import trl @@ -266,16 +266,32 @@ async def execute(self, context): return context -def test_end_turn_waits_for_the_episodes_score_past_a_tools_deadline(served, monkeypatch): +def test_a_reset_and_an_end_turn_carry_the_resets_deadline_so_end_turn_waits_for_the_score_past_a_tools( + served, monkeypatch): registry = Config().task_step_registry() monkeypatch.setattr(Config, "task_step_registry", lambda self: {**registry, _Slow.type: _Slow}) - env = served({"a": [ITEMS, ALICE, TURN, {"id": "slow", "type": _Slow.type}, REPLIES]}, tool_timeout=0.2) + deadlines = [] + reset_async, end_turn = AgentEnvEnvironment.reset_async, AgentEnvEnvironment._end_turn + + async def reset(self, *args, timeout_s=None, **kwargs): + deadlines.append(("reset", timeout_s)) + return await reset_async(self, *args, timeout_s=timeout_s, **kwargs) + + async def end(self, name, response, timeout_s): + deadlines.append((name, timeout_s)) + return await end_turn(self, name, response, timeout_s) + + monkeypatch.setattr(AgentEnvEnvironment, "reset_async", reset) + monkeypatch.setattr(AgentEnvEnvironment, "_end_turn", end) + env = served({"a": [ITEMS, ALICE, TURN, {"id": "slow", "type": _Slow.type}, REPLIES]}, tool_timeout=0.2, + reset_timeout=30) environment = env.factory() environment.reset(**env.dataset[0]) assert environment.end_turn(response="Added a crate.") == ENDED assert environment.get_reward() == 1.0 + assert deadlines == [("reset", 30), ("reset", 30), ("end_turn", 30)] def test_a_server_at_capacity_refuses_the_reset_and_says_how_to_serve_more(served): diff --git a/tst/util/capabilities.py b/tst/util/capabilities.py index e2b3dc369..70fe65823 100644 --- a/tst/util/capabilities.py +++ b/tst/util/capabilities.py @@ -33,6 +33,10 @@ #: found through ``AGENT_ENV_TEST_MCP_SERVERS_DIR`` (one ``/Dockerfile`` per server). MCP_SERVER_SOURCES = "mcp_server_sources" +#: The TRL contract tier (``tst/trl``) trains with TRL on CPU; trl and torch aren't in the dev extra, and only that +#: tier's CI job installs them. +TRL = "trl_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" From bf549dc56da22c528febc9683cfd7c9cd7a4e80b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 15:59:26 -0700 Subject: [PATCH 2/4] test(trl): train on CPU in fp32, as a runner without a GPU must Co-Authored-By: Claude Opus 5.5 --- tst/trl/grpo_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tst/trl/grpo_test.py b/tst/trl/grpo_test.py index af1918366..37f4bc897 100644 --- a/tst/trl/grpo_test.py +++ b/tst/trl/grpo_test.py @@ -32,7 +32,8 @@ def get_reward(self): trainer = trl.GRPOTrainer(model=MODEL, train_dataset=env.dataset, environment_factory=Recorded, args=trl.GRPOConfig( output_dir=str(tmp_path), learning_rate=0.1, per_device_train_batch_size=4, num_generations=2, max_steps=1, - max_tool_calling_iterations=4, shuffle_dataset=False, logging_steps=1, report_to="none")) + max_tool_calling_iterations=4, shuffle_dataset=False, logging_steps=1, report_to="none", bf16=False, + use_cpu=True)) before = {name: param.clone() for name, param in trainer.model.named_parameters()} with patch.object(trainer.model, "generate", side_effect=_Policy(trainer.processing_class)): From 444e23501f717fdce1c00082aa8adbe9bd053426 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 16:02:02 -0700 Subject: [PATCH 3/4] ci(trl): test TRL 1.16 and the latest; 1.15 can't train on CPU TRL 1.15's log-prob path is a Triton GPU kernel, which a CPU-only runner can't run. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/trl.yml | 8 ++++---- AGENTS.md | 2 +- tst/trl/grpo_test.py | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.github/workflows/trl.yml b/.github/workflows/trl.yml index fdcaec2aa..b538526bc 100644 --- a/.github/workflows/trl.yml +++ b/.github/workflows/trl.yml @@ -1,8 +1,8 @@ name: TRL contract -# One GRPOTrainer step on CPU through agent_env.training.trl (tst/trl), against the oldest TRL the helper supports and -# the latest. TRL marks environment_factory experimental, so a release can break the helper with nothing else in CI -# noticing. It runs on pull requests that touch the helper, the OpenEnv server it trains against or the tier, weekly, +# One GRPOTrainer step on CPU through agent_env.training.trl (tst/trl), against TRL 1.16 and the latest. TRL marks +# environment_factory experimental, so a release can break the helper with nothing else in CI noticing. TRL 1.15 can't +# train on CPU (its log-prob path is a Triton GPU kernel), so 1.16 is the oldest this job can run. It runs on pull requests that touch the helper, the OpenEnv server it trains against or the tier, weekly, # and on demand; it isn't a required check. on: @@ -37,7 +37,7 @@ jobs: strategy: fail-fast: false matrix: - trl: ["trl==1.15.*", "trl"] + trl: ["trl==1.16.*", "trl"] steps: - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 diff --git a/AGENTS.md b/AGENTS.md index b0e459606..532bdf144 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -301,7 +301,7 @@ skips for the allowed capabilities pass, and the fast job allows no skips under touches the plugin CLI, `agent_env/plugins` or `agent_env/config`, the protocol package, the tier itself, `pyproject.toml`, `uv.lock`, the `Makefile` or the CI that runs the tier, and allows no skips; a `changes` job decides, and a skipped `installer` counts as passed. -`.github/workflows/trl.yml` runs `tst/trl` against the oldest TRL the training helper supports and the latest, on pull +`.github/workflows/trl.yml` runs `tst/trl` against TRL 1.16 (1.15 can't train on CPU) and the latest, on pull requests that touch `agent_env/training`, `agent_env/openenv` or the tier, weekly, and on demand; it allows no skips and isn't a required check. `.github/workflows/plugin-api.yml` runs the `plugin-api` job on every pull diff --git a/tst/trl/grpo_test.py b/tst/trl/grpo_test.py index 37f4bc897..b5c210f8d 100644 --- a/tst/trl/grpo_test.py +++ b/tst/trl/grpo_test.py @@ -1,8 +1,8 @@ """TRL's GRPOTrainer trains a step through ``agent_env.training.trl`` on CPU, against the OpenEnv server ``build_app`` makes over the Episode API's fake steps. TRL marks ``environment_factory`` experimental, so this pins what the helper relies on: TRL resets an environment with its row, calls its tools, reads ``get_reward``, and trains on the rewards. -As in TRL's own environment tests, a tiny model's ``generate`` is stood in for. CI runs it against the oldest TRL the -helper supports and the latest; AsyncGRPO needs a vLLM server, so it isn't covered here.""" +As in TRL's own environment tests, a tiny model's ``generate`` is stood in for. CI runs it against TRL 1.16, the oldest +that trains on CPU, and the latest; AsyncGRPO needs a vLLM server, so it isn't covered here.""" from unittest.mock import patch From b625bfbb2ca43255bdbdbd8cd6f79b2bc6b765e8 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 16:16:01 -0700 Subject: [PATCH 4/4] fix(openenv): a call past its deadline answers at once and its episode tears down after Waiting for the teardown could outlast the client's grace, so the client dropped the session before it heard why. The teardown stays tracked: the next reset, a close and the server's shutdown wait for it. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/openenv/environment.py | 30 +++++++++++++++++----------- tst/unit/openenv/environment_test.py | 11 +++++++--- 2 files changed, 26 insertions(+), 15 deletions(-) diff --git a/src/agent_env/openenv/environment.py b/src/agent_env/openenv/environment.py index 74750b834..669755df3 100644 --- a/src/agent_env/openenv/environment.py +++ b/src/agent_env/openenv/environment.py @@ -169,8 +169,10 @@ async def _reset(self, split: str | None, index: int | None, timeout_s: float | raise EpisodeError("the run ended before the agent's first turn") await self._open(turn) except Exception as e: - return await self._end_with_error( - f"the episode didn't open its first turn within {timeout_s:g}s" if deadline.expired() else e) + if deadline.expired(): + return await self._end_with_error(f"the episode didn't open its first turn within {timeout_s:g}s", + wait=False) + return await self._end_with_error(e) return Observation(done=False, reward=None, metadata={ "prompt": turn.prompt, "system_prompt": turn.system_prompt, "task": self._task_metadata()}) @@ -231,7 +233,7 @@ async def _call(self, action: CallToolAction, timeout_s: float | None) -> CallTo return _tool_error(name, ToolErrorType.TOOL_NOT_FOUND, e) except TimeoutError: ended = await self._end_with_error(f"{name} didn't answer within {timeout_s:g}s, and the episode ended, " - "since what the call does may still be under way") + "since what the call does may still be under way", wait=False) return CallToolObservation(tool_name=name, done=True, reward=None, metadata=ended.metadata) except httpx.HTTPError as e: return _tool_error(name, ToolErrorType.TRANSPORT_ERROR, e) @@ -253,8 +255,11 @@ async def _end_turn(self, name: str, response: str, timeout_s: float | None) -> "prompt": turn.prompt, "system_prompt": turn.system_prompt}) scores = await self._episode.result() except Exception as e: - observation = await self._end_with_error( - f"{name} didn't finish within {timeout_s:g}s, and the episode ended" if deadline.expired() else e) + if deadline.expired(): + observation = await self._end_with_error(f"{name} didn't finish within {timeout_s:g}s, and the episode " + "ended", wait=False) + else: + observation = await self._end_with_error(e) return CallToolObservation(tool_name=name, done=True, reward=None, metadata=observation.metadata) await self._abandon("the episode ended") return CallToolObservation(tool_name=name, done=True, **self._scored(scores)) @@ -291,9 +296,9 @@ async def _why_stopped(self) -> dict[str, Any] | None: error = _described(e) return (await self._end_with_error(error)).metadata - async def _end_with_error(self, error: Exception | str) -> Observation: + async def _end_with_error(self, error: Exception | str, *, wait: bool = True) -> Observation: message = error if isinstance(error, str) else _described(error) - await self._abandon(message) + await self._abandon(message, wait=wait) return self._error_observation(message) def _error_observation(self, error: str) -> Observation: @@ -305,17 +310,18 @@ async def _abandon_when_free(self, reason: str) -> None: async with self._lock: await self._abandon(reason) - async def _abandon(self, reason: str) -> None: - """Abandon the episode and close the turn's tools, waiting for the teardown; the caller holds the lock. The - abandon is tracked before anything is awaited, so a server shutting down meanwhile waits for it too, and a - caller cancelled while it waits leaves it running for the next reset to wait for.""" + async def _abandon(self, reason: str, *, wait: bool = True) -> None: + """Abandon the episode and close the turn's tools, waiting for the teardown unless ``wait`` is false, as a call + past its deadline answers at once; the caller holds the lock. The abandon is tracked before anything is + awaited, so a server shutting down meanwhile waits for it too, and one not waited for, or whose caller is + cancelled while it waits, is left running for the next reset or close to wait for.""" episode, self._episode = self._episode, None if episode is not None: self._ending = self._served.end_later(episode.abandon(reason)) self._served.live.discard(self) await self._close_tools() self._turn = None - if (ending := self._ending) is not None: + if wait and (ending := self._ending) is not None: try: await asyncio.shield(ending) finally: diff --git a/tst/unit/openenv/environment_test.py b/tst/unit/openenv/environment_test.py index a8b893a21..8ad568816 100644 --- a/tst/unit/openenv/environment_test.py +++ b/tst/unit/openenv/environment_test.py @@ -230,19 +230,24 @@ async def test_a_run_that_fails_mid_turn_ends_the_episode_at_the_next_step(teard @pytest.mark.asyncio -async def test_a_reset_or_an_end_turn_past_its_deadline_ends_the_episode_in_an_error_observation(teardowns, tools): +async def test_a_reset_or_an_end_turn_past_its_deadline_answers_at_once_and_its_episode_is_torn_down_after( + teardowns, tools): held = {"id": "held", "type": "episode_test_gate"} deploying, turning = _env(ITEMS, held, ALICE, TURN, REPLIES), _env(ITEMS, ALICE, TURN, held, REPLIES) - reset = await deploying.reset_async(timeout_s=0.5) await turning.reset_async() + teardowns.hold = asyncio.Event() + reset = await deploying.reset_async(timeout_s=0.5) ended = await turning.step_async(CallToolAction(tool_name="end_turn", arguments={"response": "Added a crate."}, metadata={TIMEOUT_METADATA: 0.5})) assert (reset.done, reset.metadata["error"]) == (True, "the episode didn't open its first turn within 0.5s") assert (ended.done, ended.reward, ended.metadata["error"]) == ( True, None, "end_turn didn't finish within 0.5s, and the episode ended") - assert len(teardowns.finished) == 2 + await _until(lambda: len(teardowns.started) == 2) + assert teardowns.finished == [] + teardowns.hold.set() + await _until(lambda: len(teardowns.finished) == 2) @pytest.mark.asyncio