From 775715efecc227a04f2c146b21a279f6009586ba Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 01:23:17 -0700 Subject: [PATCH 1/2] feat(training): train with Unsloth through the TRL helper, and wait out a restarting env server Unsloth's GRPOTrainer, TRL's rewritten at import, takes AgentEnv's dataset and environment_factory as they are (Unsloth 2026.10.11 on TRL 1.15). Its GPU proof found that a TRL rollout whose session a restarting server dropped was left unscored at once, so every rollout of the steps the restart spanned was lost: the helper now waits for a server that refuses connections, as one restarting does, as long as a reset may take, keeping the short wait for room. The TRL tier also runs on transformers 5.17, the newest Unsloth allows. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/trl.yml | 19 +++++++++++------- AGENTS.md | 3 ++- src/agent_env/training/_openenv.py | 23 +++++++++++----------- src/agent_env/training/trl.py | 14 +++++++++----- tst/unit/training/sessions_test.py | 25 +++++++++++++----------- tst/unit/training/trl_test.py | 31 +++++++++++++++++++++++++++--- 6 files changed, 77 insertions(+), 38 deletions(-) diff --git a/.github/workflows/trl.yml b/.github/workflows/trl.yml index b538526b..cc614966 100644 --- a/.github/workflows/trl.yml +++ b/.github/workflows/trl.yml @@ -1,9 +1,10 @@ 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. +# One GRPOTrainer step on CPU through agent_env.training.trl (tst/trl), against TRL 1.16 and the latest, and on +# transformers 5.17, the newest Unsloth's GRPOTrainer allows. 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, though Unsloth caps TRL at 1.15. 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: @@ -31,13 +32,17 @@ env: jobs: grpo: - name: grpo (${{ matrix.trl }}) + name: grpo (${{ matrix.trl }}${{ matrix.transformers && format(', {0}', matrix.transformers) || '' }}) runs-on: ubuntu-latest timeout-minutes: 20 strategy: fail-fast: false matrix: - trl: ["trl==1.16.*", "trl"] + include: + - trl: "trl==1.16.*" + - trl: "trl" + - trl: "trl==1.16.*" + transformers: "transformers==5.17.*" steps: - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 @@ -47,7 +52,7 @@ jobs: 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 }}" + uv pip install --python .venv/bin/python --torch-backend cpu torch "${{ matrix.trl }}" ${{ matrix.transformers }} .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 diff --git a/AGENTS.md b/AGENTS.md index 1e184964..a14c25ca 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -331,7 +331,8 @@ 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 +`.github/workflows/trl.yml` runs `tst/trl` against TRL 1.16 (1.15 can't train on CPU) and the latest, and on +transformers 5.17, the newest Unsloth's `GRPOTrainer` allows, 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/nemo-gym.yml` runs `tst/nemo_gym` on Python 3.13 with the `nemo-gym` extra from `uv.lock` and with the latest nemo-gym, on the same triggers; it allows no skips and isn't a required check. diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index 7bb62668..1ceabbd2 100644 --- a/src/agent_env/training/_openenv.py +++ b/src/agent_env/training/_openenv.py @@ -61,9 +61,9 @@ class NamedToolset: @dataclass(frozen=True) class Shared: """What a session needs of its server: where it is, the tools the model is shown, which an episode's first turn - must offer no more than, the deadlines, how long to wait for room on the server, whether to wait as long for a - server that refuses connections, as one restarting does, and the token it asks for. The key makes it equal to its - pickled copies only.""" + must offer no more than, the deadlines, how long to wait for room on the server, whether to wait for a server that + refuses connections, as one restarting does, as long as a reset may take, and the token it asks for. The key makes + it equal to its pickled copies only.""" url: str offered: frozenset[str] @@ -277,17 +277,17 @@ async def opened(url: str, timeout: float, *, token: str | None = None) -> Gener async def reopened(url: str, timeout: float, *, dropped: bool, patience: float, waits_for_server: bool = False, token: str | None = None) -> GenericEnvClient: - """A session, tried again every ``REOPEN_SECONDS`` for ``patience`` seconds while the server is at its limit or too - busy to answer, or with ``waits_for_server`` refuses the connection, as one restarting does, in place of a - ``dropped`` one or for a rollout's own; raises Failed, which leaves the rollout unscored, when it can't open one, at - once unless it was waiting.""" + """A session, tried again every ``REOPEN_SECONDS`` while the server is at its limit or too busy to answer, for + ``patience`` seconds, or with ``waits_for_server`` while it refuses the connection, as one restarting does, for as + long as a reset may take, in place of a ``dropped`` one or for a rollout's own; raises Failed, which leaves the + rollout unscored, when it can't open one, at once unless it was waiting.""" lead = "its session was dropped, and " if dropped else "" - deadline = time.monotonic() + patience + began = time.monotonic() while True: try: return await opened(url, timeout, token=token) except AtCapacity as e: - if time.monotonic() + REOPEN_SECONDS > deadline: + if time.monotonic() + REOPEN_SECONDS > began + patience: raise Failed(f"{lead}the server stayed at its limit of {e.limit} sessions for " f"{patience:g}s{' after' if dropped else ''}") from e if url not in _waited: @@ -297,8 +297,9 @@ async def reopened(url: str, timeout: float, *, dropped: bool, patience: float, "more at once", url, e.limit) await asyncio.sleep(REOPEN_SECONDS) except ConnectionError as e: - waits = _transient(e) or (waits_for_server and _caused_by(e, ConnectionRefusedError)) - if not waits or time.monotonic() + REOPEN_SECONDS > deadline: + restarting = waits_for_server and _caused_by(e, ConnectionRefusedError) + waited = max(patience, timeout) if restarting else patience + if not (restarting or _transient(e)) or time.monotonic() + REOPEN_SECONDS > began + waited: raise Failed(f"{lead}{e}") from e await asyncio.sleep(REOPEN_SECONDS) diff --git a/src/agent_env/training/trl.py b/src/agent_env/training/trl.py index f595e5b8..7755489b 100644 --- a/src/agent_env/training/trl.py +++ b/src/agent_env/training/trl.py @@ -2,8 +2,9 @@ writes): ``AgentEnv`` reads the server's tasks into a dataset, and makes the ``environment_factory`` whose environments each play one episode of a task over OpenEnv's WebSocket, offer its turn's tools as methods and own its reward. ``GRPOTrainer`` takes both as they are, verified with TRL 1.16 and 1.17, and AsyncGRPO the same ``environment_factory``; -either needs TRL 1.15 or later, for an environment that owns its reward. Every session runs on one event-loop thread of -this module's own, since TRL calls an environment synchronously, from its own threads.""" +either needs TRL 1.15 or later, for an environment that owns its reward. Unsloth's ``GRPOTrainer``, TRL's rewritten when +``unsloth`` is imported first, takes them too, verified with Unsloth 2026.10.11 on TRL 1.15. Every session runs on one +event-loop thread of this module's own, since TRL calls an environment synchronously, from its own threads.""" from __future__ import annotations @@ -73,7 +74,8 @@ 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. The server ends an episode that misses it, so the rollout is left unscored. + in seconds. The server ends an episode that misses it, so the rollout is left unscored. A session dropped + by a server restarting waits as long for it to come back. Example: @@ -155,7 +157,8 @@ async def _environments(self) -> list[type[_Environment]]: async def _environment(self, toolset: NamedToolset) -> type[_Environment]: first, tools = await probe(self.url, toolset, self._reset_timeout) - shared = Shared(url=self.url, offered=first, tool_timeout=self._tool_timeout, reset_timeout=self._reset_timeout) + shared = Shared(url=self.url, offered=first, tool_timeout=self._tool_timeout, reset_timeout=self._reset_timeout, + waits_for_server=True) self._shared.append(shared) return _environment_class(toolset.name, tools, shared) @@ -187,7 +190,8 @@ def reset(self, *, split: str, index: int, task_id: str | None = None, **_: Any) """Starts the row's episode and returns before it is deployed, since ``GRPOTrainer`` resets its environments one at a time: the first tool call or ``get_reward`` waits for it. Raises when the server refuses this environment's first session, as one at capacity does; a later session it refuses, after one was dropped, - leaves only that rollout unscored.""" + leaves only that rollout unscored, unless the server is restarting, which it waits out within the reset + deadline.""" on_session_loop(self._start(split, index, task_id), CONNECT_SECONDS + CLOSE_SECONDS) def get_reward(self) -> float | None: diff --git a/tst/unit/training/sessions_test.py b/tst/unit/training/sessions_test.py index 002fd962..f515e115 100644 --- a/tst/unit/training/sessions_test.py +++ b/tst/unit/training/sessions_test.py @@ -1,6 +1,7 @@ """``agent_env.training._openenv`` opening a rollout's session: tried again while the server is at its limit or too busy -to answer its handshake in time, or when asked to wait for it refuses connections as one restarting does, until the -patience runs out, and given up at once when nothing listens at its address or it refuses the token.""" +to answer its handshake in time, until the patience runs out, or when asked to wait for it refuses connections as one +restarting does, for as long as a reset may take; given up at once when nothing listens at its address or it refuses +the token.""" import asyncio @@ -10,14 +11,16 @@ from agent_env.training._openenv import Failed, reopened -@pytest.mark.parametrize(("cause", "waits_for_server", "opens", "attempts"), [ - (TimeoutError("timed out"), False, True, 3), - (ConnectionRefusedError(61, "Connection refused"), False, False, 1), - (ConnectionRefusedError(61, "Connection refused"), True, True, 3), - (ConnectionError("server rejected WebSocket connection: HTTP 403"), True, False, 1), -], ids=["busy", "refused", "restarting", "token-refused"]) +@pytest.mark.parametrize(("cause", "waits_for_server", "patience", "opens", "attempts"), [ + (TimeoutError("timed out"), False, 5, True, 3), + (TimeoutError("timed out"), False, 0, False, 1), + (ConnectionRefusedError(61, "Connection refused"), False, 5, False, 1), + (ConnectionRefusedError(61, "Connection refused"), True, 5, True, 3), + (ConnectionRefusedError(61, "Connection refused"), True, 0, True, 3), + (ConnectionError("server rejected WebSocket connection: HTTP 403"), True, 5, False, 1), +], ids=["busy", "busy-past-patience", "refused", "restarting", "restarting-past-patience", "token-refused"]) def test_a_session_is_tried_again_while_its_server_is_busy_and_given_up_at_once_when_none_listens( - monkeypatch, cause, waits_for_server, opens, attempts): + monkeypatch, cause, waits_for_server, patience, opens, attempts): tried = [] async def flaky(url, timeout, *, token=None): @@ -32,9 +35,9 @@ async def flaky(url, timeout, *, token=None): monkeypatch.setattr(_openenv, "REOPEN_SECONDS", 0.01) if opens: - assert asyncio.run(reopened("http://busy", 5, dropped=False, patience=5, + assert asyncio.run(reopened("http://busy", 5, dropped=False, patience=patience, waits_for_server=waits_for_server)) == "session" else: with pytest.raises(Failed, match="couldn't open a session on the server at http://busy"): - asyncio.run(reopened("http://busy", 5, dropped=False, patience=5, waits_for_server=waits_for_server)) + asyncio.run(reopened("http://busy", 5, dropped=False, patience=patience, waits_for_server=waits_for_server)) assert len(tried) == attempts diff --git a/tst/unit/training/trl_test.py b/tst/unit/training/trl_test.py index 21b31a17..135d9bc9 100644 --- a/tst/unit/training/trl_test.py +++ b/tst/unit/training/trl_test.py @@ -1,7 +1,7 @@ """``AgentEnv`` on the OpenEnv server ``build_app`` makes, served on loopback over the Episode API's fake steps and used as TRL's trainers use it: the dataset and the environment classes, rollouts reset, stepped and scored on a reused -environment, the turns a model leaves open, episodes left unscored, a server at capacity or missing, tasks with other -tools, and a bundle served by ``AgentEnv.local`` whose episodes end at close and at exit.""" +environment, the turns a model leaves open, episodes left unscored, a server at capacity, missing or restarting, tasks +with other tools, and a bundle served by ``AgentEnv.local`` whose episodes end at close and at exit.""" import asyncio import inspect @@ -17,7 +17,7 @@ from agent_env.config.runtime import Config from agent_env.openenv.environment import AgentEnvEnvironment from agent_env.task_step.task_step import TaskStep -from agent_env.training import trl +from agent_env.training import _openenv, trl from agent_env.training._openenv import ENDED, FAILED, on_session_loop from agent_env.training.trl import AgentEnv from tst.unit.bundle._support import layout @@ -233,6 +233,31 @@ def test_a_session_refused_after_one_was_dropped_leaves_only_that_rollout_unscor assert "its session was dropped, and the server stayed at its limit of 1 sessions for 4s after" in caplog.text +def test_a_session_dropped_by_a_restarting_server_opens_again_once_it_is_back(served, monkeypatch): + env = served({"a": PLAYED}) + environment = env.factory() + environment.reset(**env.dataset[0]) + environment.end_turn(response="Added a crate.") + assert environment.get_reward() == 1.0 + on_session_loop(environment._close(), 30) + opened, refusals = _openenv.opened, [] + + async def restarting(url, timeout, *, token=None): + if len(refusals) < 3: + refusals.append(url) + refused = ConnectionError(f"Failed to connect to {url}/ws") + refused.__cause__ = ConnectionRefusedError(111, "Connection refused") + raise ConnectionError(f"couldn't open a session on the server at {url}") from refused + return await opened(url, timeout, token=token) + + monkeypatch.setattr(_openenv, "opened", restarting) + monkeypatch.setattr(_openenv, "REOPEN_SECONDS", 0.01) + environment.reset(**env.dataset[0]) + environment.end_turn(response="Added a crate.") + + assert (environment.get_reward(), len(refusals)) == (1.0, 3) + + def test_a_rollout_whose_tools_arent_its_classs_stops_training(served, monkeypatch): monkeypatch.setattr(proxy_module, "ToolSession", _SecondTurnTools()) env = served({"a": PLAYED}) From 2fad817d7338c54a1d1ab472183482a21360a648 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 1 Nov 2026 01:27:39 -0700 Subject: [PATCH 2/2] fix(training): a TRL environment's call waits out a restarting server's reopen within its overall deadline Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/trl.py | 9 +++++---- tst/unit/training/trl_test.py | 9 +++++++++ 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/agent_env/training/trl.py b/src/agent_env/training/trl.py index 7755489b..fc1de3f8 100644 --- a/src/agent_env/training/trl.py +++ b/src/agent_env/training/trl.py @@ -213,11 +213,12 @@ def get_reward(self) -> float | None: return rollout.reward def _call(self, name: str, arguments: dict[str, Any]) -> str: - """Waits at most for the deploy and the call, each to its deadline and a grace after, then the session's close: - the first call of a rollout waits for its deploy.""" + """Waits at most for a restarting server, 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, after its session reopens.""" shared = self._shared - return on_session_loop(self._step(name, arguments), shared.reset_timeout + max(shared.reset_timeout, shared.tool_timeout) - + 2 * GRACE_SECONDS + CLOSE_SECONDS) + reopen = max(shared.patience, shared.reset_timeout) + waits = reopen + shared.reset_timeout + max(shared.reset_timeout, shared.tool_timeout) + return on_session_loop(self._step(name, arguments), waits + 2 * GRACE_SECONDS + CLOSE_SECONDS) class _EnvironmentClass(type): diff --git a/tst/unit/training/trl_test.py b/tst/unit/training/trl_test.py index 135d9bc9..dd786834 100644 --- a/tst/unit/training/trl_test.py +++ b/tst/unit/training/trl_test.py @@ -250,12 +250,21 @@ async def restarting(url, timeout, *, token=None): raise ConnectionError(f"couldn't open a session on the server at {url}") from refused return await opened(url, timeout, token=token) + on_loop, budgets = trl.on_session_loop, [] + + def budgeted(coroutine, timeout): + budgets.append(timeout) + return on_loop(coroutine, timeout) + monkeypatch.setattr(_openenv, "opened", restarting) monkeypatch.setattr(_openenv, "REOPEN_SECONDS", 0.01) + monkeypatch.setattr(trl, "on_session_loop", budgeted) environment.reset(**env.dataset[0]) environment.end_turn(response="Added a crate.") + shared = environment._shared assert (environment.get_reward(), len(refusals)) == (1.0, 3) + assert budgets[1] > shared.reset_timeout + shared.reset_timeout + max(shared.reset_timeout, shared.tool_timeout) def test_a_rollout_whose_tools_arent_its_classs_stops_training(served, monkeypatch):