fix(nemo-gym): sessions wait out a busy env server, and serve reclaims abandoned ones

src/agent_env/training/_openenv.py

@@ -38,6 +38,9 @@
 #: How long, and how far apart in seconds, a session tries to open while the server is at its limit, which it stays at
 #: until it has seen a closed session go.
 REOPEN_PATIENCE, REOPEN_SECONDS = 4.0, 1.0
+#: What a session's opening fails with when the server is too busy to answer in time, or drops the connection.
+_TRANSIENT = (TimeoutError, ConnectionResetError, ConnectionAbortedError, BrokenPipeError, EOFError,
+              httpx.TimeoutException, httpx.ReadError, httpx.RemoteProtocolError)
 
 _CAPACITY = re.compile(r"Server at capacity: \d+/(\d+)")
 
@@ -263,9 +266,9 @@ async def opened(url: str, timeout: float) -> GenericEnvClient:
 
 
 async def reopened(url: str, timeout: float, *, dropped: bool, patience: float) -> GenericEnvClient:
-    """A session, tried again every ``REOPEN_SECONDS`` for ``patience`` seconds while the server is at its limit, 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."""
+    """A session, tried again every ``REOPEN_SECONDS`` for ``patience`` seconds while the server is at its limit or too
+    busy to answer, 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 the server was at its limit or too busy to answer."""
     lead = "its session was dropped, and " if dropped else ""
     deadline = time.monotonic() + patience
     while True:
@@ -282,7 +285,19 @@ 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:
-            raise Failed(f"{lead}{e}") from e
+            if not _transient(e) or time.monotonic() + REOPEN_SECONDS > deadline:
+                raise Failed(f"{lead}{e}") from e
+            await asyncio.sleep(REOPEN_SECONDS)
+
+
+def _transient(error: BaseException | None) -> bool:
+    """Whether ``error`` came of a server too busy to answer in time or a connection it dropped, rather than of nothing
+    listening at its address or an answer that isn't a session."""
+    while error is not None:
+        if isinstance(error, _TRANSIENT):
+            return True
+        error = error.__cause__ or error.__context__
+    return False
 
 
 async def probe(url: str, toolset: NamedToolset, timeout: float) -> tuple[frozenset[str], tuple[dict[str, Any], ...]]:

src/agent_env/training/nemo_gym/__main__.py

@@ -57,13 +57,17 @@ def write(source: str, root: str, name: str | None, sandbox: str | None, reward:
               help="Answer verify for a rollout the environment failed with HTTP 503 rather than mask_sample, for a "
                    "client that doesn't read mask_sample, as verifiers' NeMo Gym taskset doesn't: it then fails the "
                    "rollout, and prime-rl trains on none of it.")
-def serve(folder: str, host: str, port: int, masked_as_error: bool) -> None:
+@click.option("--idle-timeout", type=float,
+              help="Seconds a session may go without a request before its episode is ended, the config's (600) by "
+                   "default. A client that abandons sessions without closing them, as verifiers' NeMo Gym taskset "
+                   "does with the rollouts prime-rl cancels, holds their episodes until then.")
+def serve(folder: str, host: str, port: int, masked_as_error: bool, idle_timeout: float | None) -> None:
     """Runs the resources server written into FOLDER, a Gym root's resources_servers/<name>, on its own: without Gym's
     head server or the agent servers its config names. Point verifiers' NeMo Gym taskset at it with
     env.taskset.task.resources_url, and its env.taskset.dataset at the folder's data/<split>.jsonl."""
     if run_server is None:
         raise click.ClickException(UNSERVABLE)
-    run_server(folder, host=host, port=port, masked_as_error=masked_as_error)
+    run_server(folder, host=host, port=port, masked_as_error=masked_as_error, idle_timeout=idle_timeout)
 
 
 if __name__ == "__main__":

src/agent_env/training/nemo_gym/server.py

@@ -17,7 +17,7 @@
 import yaml
 from fastapi import Body, FastAPI, HTTPException, Request
 from fastapi.responses import PlainTextResponse
-from pydantic import ConfigDict, PrivateAttr, model_validator
+from pydantic import ConfigDict, Field, PrivateAttr, model_validator
 
 try:
     from nemo_gym.base_resources_server import (
@@ -72,7 +72,7 @@ class AgentEnvResourcesServerConfig(BaseResourcesServerConfig):
     max_concurrent_envs: int = 8
     tool_timeout: float = 60.0
     reset_timeout: float = 600.0
-    idle_timeout: float = 600.0
+    idle_timeout: float = Field(600.0, gt=0, allow_inf_nan=False)
     reward_components: bool = False
     masked_as_error: bool = False
 
@@ -211,19 +211,21 @@ async def _shut_down(self) -> None:
                            "may be left running", served.url, STOP_SECONDS)
 
 
-def serve(folder: str | Path, *, host: str, port: int, masked_as_error: bool = False) -> None:
+def serve(folder: str | Path, *, host: str, port: int, masked_as_error: bool = False,
+          idle_timeout: float | None = None) -> None:
     """Runs the resources server ``write_server`` wrote into ``folder`` on its own, at ``host``:``port``, without Gym's
     head server or the agent servers its config names, for a trainer that plays the folder's rows against the server's
-    URL. ``host`` and ``port`` win over any the folder's config names, and ``masked_as_error``, as the config takes it,
-    is on when either it or the config sets it. As verifiers' launcher for Gym's own servers does, its client names a
-    head server nothing calls."""
+    URL. ``host``, ``port`` and an ``idle_timeout`` win over any the folder's config names, and ``masked_as_error``, as
+    the config takes it, is on when either it or the config sets it. As verifiers' launcher for Gym's own servers does,
+    its client names a head server nothing calls."""
     folder = Path(folder).resolve()
     written = yaml.safe_load((folder / "configs" / f"{folder.name}.yaml").read_text())
     config = written[folder.name]["resources_servers"][folder.name]
-    masked_as_error = masked_as_error or config.get("masked_as_error", False)
+    settings = {"name": folder.name, "host": host, "port": port,
+                "masked_as_error": masked_as_error or config.get("masked_as_error", False),
+                **({"idle_timeout": idle_timeout} if idle_timeout is not None else {})}
     resources = AgentEnvResourcesServer(
-        config=AgentEnvResourcesServerConfig(**{**config, "name": folder.name, "host": host, "port": port,
-                                                "masked_as_error": masked_as_error}),
+        config=AgentEnvResourcesServerConfig(**{**config, **settings}),
         server_client=ServerClient(head_server_config=BaseServerConfig(host=host, port=port),
                                    global_config_dict=OmegaConf.create({})),
     )

tst/nemo_gym/server_test.py

@@ -3,8 +3,8 @@
 the tool routes answer as text, verify ends the turns left open and scores the episode or masks the rollout, a rollout
 waits its turn at the server's limit, a closed or idle session ends its episode, the server takes one worker, a
 bundle named in the config is served for the server's lifetime, and ``serve`` runs a written server on its own,
-answering a masked rollout's verify with an HTTP error when asked. CI runs it on the nemo-gym ``uv.lock`` pins and on
-the latest nemo-gym."""
+answering a masked rollout's verify with an HTTP error when asked and ending abandoned episodes after the idle timeout
+it is given. CI runs it on the nemo-gym ``uv.lock`` pins and on the latest nemo-gym."""
 
 import asyncio
 import json
@@ -155,8 +155,11 @@ def test_a_rollout_gym_stopped_before_verifying_has_its_episode_ended_once_idle(
 @pytest.mark.parametrize(("config", "why"), [
     ({"num_workers": 2}, "the server runs with one worker"),
     ({"reward": "replies", "reward_components": True}, "reward_components makes the reward the sum"),
+    ({"idle_timeout": 0}, "greater than 0"),
+    ({"idle_timeout": float("nan")}, "finite number"),
 ])
-def test_a_config_with_more_than_one_worker_or_a_reward_beside_its_components_is_refused(config, why):
+def test_a_config_with_more_than_one_worker_a_reward_beside_its_components_or_no_finite_idle_timeout_is_refused(config,
+                                                                                                             why):
     with pytest.raises(pydantic.ValidationError, match=why):
         server.AgentEnvResourcesServerConfig(host="127.0.0.1", port=0, entrypoint="app.py", name="dock",
                                              url="http://127.0.0.1:8000", **config)
@@ -223,6 +226,32 @@ def test_serve_runs_a_written_server_on_its_own_at_the_port_it_is_given_answerin
     assert len(teardowns.finished) == 3
 
 
+def test_serve_ends_the_episode_of_a_session_its_client_abandons_after_the_idle_timeout_it_is_given(serve, tmp_path,
+                                                                                                 teardowns):
+    bundle_server = serve({"a": PLAYED})
+    folder = write_server(bundle_server.url, tmp_path)
+    row = json.loads((folder / "data" / "all.jsonl").read_text().splitlines()[0])
+    port = _free_port()
+
+    served = subprocess.Popen([sys.executable, "-m", "agent_env.training.nemo_gym", "serve", str(folder), "--port",
+                               str(port), "--idle-timeout", "1"],
+                              stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    try:
+        with httpx.Client(base_url=_up(served, f"http://127.0.0.1:{port}")) as client:
+            client.post("/seed_session", json=row)
+            assert client.post("/add_item", json={"item": "crate"}).text == "add_item {'item': 'crate'}"
+            time.sleep(3)
+            abandoned = client.post("/add_item", json={"item": "crate"})
+    finally:
+        served.terminate()
+        served.wait(timeout=30)
+    bundle_server.stop()
+
+    assert abandoned.status_code == 400
+    assert "ended the episode of a rollout with no request for 1s" in served.stdout.read()
+    assert len(teardowns.finished) == 2
+
+
 _EPISODE = {"rollout_id": "rollout"}
 
 

tst/unit/training/sessions_test.py

@@ -0,0 +1,37 @@
+"""``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, until the patience runs out, and given up at once when nothing listens at its
+address."""
+
+import asyncio
+
+import pytest
+
+from agent_env.training import _openenv
+from agent_env.training._openenv import Failed, reopened
+
+
+@pytest.mark.parametrize(("cause", "opens", "attempts"), [
+    (TimeoutError("timed out"), True, 3),
+    (ConnectionRefusedError(61, "Connection refused"), False, 1),
+], ids=["busy", "refused"])
+def test_a_session_is_tried_again_while_its_server_is_busy_and_given_up_at_once_when_none_listens(monkeypatch, cause,
+                                                                                                  opens, attempts):
+    tried = []
+
+    async def flaky(url, timeout):
+        tried.append(url)
+        if len(tried) < 3:
+            client_error = ConnectionError(f"Failed to connect to {url}/ws: {cause}")
+            client_error.__cause__ = cause
+            raise ConnectionError(f"couldn't open a session on the server at {url}") from client_error
+        return "session"
+
+    monkeypatch.setattr(_openenv, "opened", flaky)
+    monkeypatch.setattr(_openenv, "REOPEN_SECONDS", 0.01)
+
+    if opens:
+        assert asyncio.run(reopened("http://busy", 5, dropped=False, patience=5)) == "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))
+    assert len(tried) == attempts