diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index 0b6a435a..c37dd6ef 100644 --- a/src/agent_env/training/_openenv.py +++ b/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], ...]]: diff --git a/src/agent_env/training/nemo_gym/__main__.py b/src/agent_env/training/nemo_gym/__main__.py index 50d69691..f426bdde 100644 --- a/src/agent_env/training/nemo_gym/__main__.py +++ b/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/, 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/.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__": diff --git a/src/agent_env/training/nemo_gym/server.py b/src/agent_env/training/nemo_gym/server.py index b40e926f..e245675f 100644 --- a/src/agent_env/training/nemo_gym/server.py +++ b/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({})), ) diff --git a/tst/nemo_gym/server_test.py b/tst/nemo_gym/server_test.py index 677b91eb..abc510c9 100644 --- a/tst/nemo_gym/server_test.py +++ b/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"} diff --git a/tst/unit/training/sessions_test.py b/tst/unit/training/sessions_test.py new file mode 100644 index 00000000..381247b6 --- /dev/null +++ b/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