diff --git a/.github/workflows/trl.yml b/.github/workflows/trl.yml new file mode 100644 index 000000000..b538526bc --- /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 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: + 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.16.*", "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..532bdf144 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 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 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..669755df3 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,14 +159,19 @@ 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: + 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()}) @@ -189,9 +195,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.""" @@ -225,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) @@ -234,18 +242,24 @@ 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) + 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)) @@ -282,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: @@ -296,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/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..b5c210f8d --- /dev/null +++ b/tst/trl/grpo_test.py @@ -0,0 +1,68 @@ +"""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 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 + +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", 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)): + 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..8ad568816 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,27 @@ 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_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) + 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") + await _until(lambda: len(teardowns.started) == 2) + assert teardowns.finished == [] + teardowns.hold.set() + await _until(lambda: 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"