feat(training): train with Unsloth through the TRL helper, and wait out a restarting env server

.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

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.

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)
 

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:
@@ -209,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):

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

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,40 @@ 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)
+
+    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):
     monkeypatch.setattr(proxy_module, "ToolSession", _SecondTurnTools())
     env = served({"a": PLAYED})