From e29d03ba69fbc1d9e2ba0d1aaf42ddbe64cf69d8 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Thu, 27 Aug 2026 23:43:19 -0700 Subject: [PATCH 1/4] Retry PTY takeover attach when execd's eviction times out A takeover attach evicts the currently attached client, and execd waits for that client to acknowledge. When the old client's socket is half-open (silently dropped along the NLB/proxy path) it cannot answer, so execd's eviction times out and the attach comes back as a policy-violation close - surfaced as 'PTY session already has an attached client' even though the client passed takeover=True. The stale client is torn down in the background, so a fresh attempt lands. attach_pty now re-dials up to three times (2s/5s/10s) on exactly that signature; rejections without takeover stay definitive. Observed on terminal-bench 2.1 verify re-attaches at 1k-sandbox scale. Co-Authored-By: Claude Fable 5 Signed-off-by: Hemil Desai --- .../sandbox/providers/opensandbox/provider.py | 34 ++++++++---- nemo_gym/sandbox/providers/opensandbox/pty.py | 5 ++ tests/unit_tests/test_opensandbox_pty.py | 52 +++++++++++++++++++ 3 files changed, 81 insertions(+), 10 deletions(-) diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index d1d4f1f2a7..487efc9d6f 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -32,6 +32,7 @@ SandboxEndpoint, SandboxExecResult, SandboxHandle, + SandboxPtyError, SandboxPtySession, SandboxPtySpec, SandboxResources, @@ -1489,18 +1490,31 @@ async def attach_pty( since: int | None = None, ) -> SandboxPtySession: """Re-attach to an existing execd PTY session by id.""" - from nemo_gym.sandbox.providers.opensandbox.pty import attach_pty_session + from nemo_gym.sandbox.providers.opensandbox.pty import _PTY_TAKEOVER_RETRY_DELAYS, attach_pty_session base_url, headers, request_timeout_s = await self._pty_target(handle) - session = await attach_pty_session( - client=self._pty_http_client(), - base_url=base_url, - headers=headers, - session_id=session_id, - takeover=takeover, - since=since, - request_timeout_s=request_timeout_s, - ) + # A takeover evicts the attached client, and execd waits for that + # client to acknowledge. A half-open peer (silently dropped along the + # proxy path) cannot answer, so execd's eviction times out and the + # attach comes back as a policy-violation close — reported as "already + # has an attached client" — while the stale client is torn down in the + # background. A fresh attempt then lands, so retry exactly that case. + for delay in (*_PTY_TAKEOVER_RETRY_DELAYS, None): + try: + session = await attach_pty_session( + client=self._pty_http_client(), + base_url=base_url, + headers=headers, + session_id=session_id, + takeover=takeover, + since=since, + request_timeout_s=request_timeout_s, + ) + break + except SandboxPtyError as e: + if not takeover or delay is None or "already has an attached client" not in str(e): + raise + await asyncio.sleep(delay) await self._retire_closed_pty_sessions() self._pty_sessions.add(session) return session diff --git a/nemo_gym/sandbox/providers/opensandbox/pty.py b/nemo_gym/sandbox/providers/opensandbox/pty.py index 85faf3a4ed..e789bd2482 100644 --- a/nemo_gym/sandbox/providers/opensandbox/pty.py +++ b/nemo_gym/sandbox/providers/opensandbox/pty.py @@ -55,6 +55,11 @@ # window; the tail rides out per-replica informer lag (each retry re-rolls the # load-balanced replica, so 404 "pod IP not yet available" clears quickly). _PTY_RETRY_DELAYS = (0.25, 0.5, 1.0, 2.0, 4.0, 8.0) +# A takeover that finds a half-open peer waits out execd's eviction timeout +# before the policy-violation close arrives, so these are spaced in seconds, +# not fractions: the stale client is torn down in the background and a later +# attempt lands. +_PTY_TAKEOVER_RETRY_DELAYS = (2.0, 5.0, 10.0) # Mirrors execd's shell pick (bash when available, else sh) for env-only specs. _DEFAULT_SHELL_SNIPPET = 'exec "$(command -v bash || echo sh)"' diff --git a/tests/unit_tests/test_opensandbox_pty.py b/tests/unit_tests/test_opensandbox_pty.py index a0163ec7c7..09175a8daf 100644 --- a/tests/unit_tests/test_opensandbox_pty.py +++ b/tests/unit_tests/test_opensandbox_pty.py @@ -584,6 +584,58 @@ async def get_endpoint(self, port: int) -> SimpleNamespace: await session.close() +async def test_provider_attach_pty_retries_timed_out_takeover(monkeypatch: pytest.MonkeyPatch) -> None: + # execd waits for the evicted client to acknowledge a takeover; a + # half-open peer cannot, so the first attach closes 1008 while the stale + # client is torn down in the background. The provider re-dials and lands. + pytest.importorskip("tenacity", reason="tenacity optional sandbox dependency is not installed") + pytest.importorskip("opensandbox", reason="opensandbox SDK is not installed") + from nemo_gym.sandbox.providers.opensandbox.provider import OpenSandboxProvider + + class FakeRaw: + async def get_endpoint(self, port: int) -> SimpleNamespace: + return SimpleNamespace(endpoint="server/v1/sandboxes/sb-1/proxy/44772", headers={}) + + provider = OpenSandboxProvider(connection={"domain": "server", "api_key": "k", "protocol": "https"}) + rejected = FakeWs([], close_code=1008) + rejected.closed = True + clients = [FakeHttpClient(ws=rejected), FakeHttpClient(ws=FakeWs([CONNECTED]))] + handed_out: list[FakeHttpClient] = [] + monkeypatch.setattr(provider, "_pty_http_client", lambda: handed_out.append(clients[len(handed_out)]) or handed_out[-1]) + monkeypatch.setattr(pty_module, "_PTY_TAKEOVER_RETRY_DELAYS", (0.0,)) + handle = SandboxHandle(sandbox_id="sb-1", provider_name="opensandbox", raw=FakeRaw()) + session = await provider.attach_pty(handle, "s-7", takeover=True) + assert len(handed_out) == 2, "the timed-out takeover must be re-dialed once" + assert handed_out[0].closed, "the failed attempt's client must be released" + await session.close() + + +async def test_provider_attach_pty_does_not_retry_without_takeover(monkeypatch: pytest.MonkeyPatch) -> None: + pytest.importorskip("tenacity", reason="tenacity optional sandbox dependency is not installed") + pytest.importorskip("opensandbox", reason="opensandbox SDK is not installed") + from nemo_gym.sandbox.providers.opensandbox.provider import OpenSandboxProvider + + class FakeRaw: + async def get_endpoint(self, port: int) -> SimpleNamespace: + return SimpleNamespace(endpoint="server/v1/sandboxes/sb-1/proxy/44772", headers={}) + + provider = OpenSandboxProvider(connection={"domain": "server", "api_key": "k", "protocol": "https"}) + rejected = FakeWs([], close_code=1008) + rejected.closed = True + clients_handed = 0 + + def _client() -> FakeHttpClient: + nonlocal clients_handed + clients_handed += 1 + return FakeHttpClient(ws=rejected) + + monkeypatch.setattr(provider, "_pty_http_client", _client) + handle = SandboxHandle(sandbox_id="sb-1", provider_name="opensandbox", raw=FakeRaw()) + with pytest.raises(SandboxPtyError, match="already has an attached client"): + await provider.attach_pty(handle, "s-7", takeover=False) + assert clients_handed == 1, "without takeover the rejection is definitive" + + async def test_create_rejected_before_connected_raises_and_cleans_up() -> None: # The session we created is torn down when the socket is rejected. ws = FakeWs([], close_code=1008) From f1974d2e31066a88ddd4f22c84e7a28ba642230d Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Fri, 28 Aug 2026 10:13:53 -0700 Subject: [PATCH 2/4] Detach the PTY baton instead of dangling it, and heartbeat the sockets The tb2.1 flow passes one PTY session through three holders (create, agent, verify), and every handoff relied on takeover-evicting the previous client. Eviction is only cheap while that client is alive; a socket idled through the whole agent phase is exactly the one the NLB/proxy path silently drops, turning the next takeover into a timeout. Two changes make handoffs clean instead of contested: - seed_session detaches its client right after creating the session, so the next attach finds nothing to evict; the server-side session keeps running and replays output on reattach. - every PTY dial requests websocket heartbeats (30s), so a dead socket surfaces as a failed ping and re-dials via the existing reattach path within a minute instead of dangling half-open. Takeover (with the retry from the previous commit) remains the recovery path for crashed holders rather than the routine handoff mechanism. Co-Authored-By: Claude Fable 5 Signed-off-by: Hemil Desai --- nemo_gym/sandbox/providers/opensandbox/pty.py | 11 +++++++++- resources_servers/terminal_bench_2_1/app.py | 7 ++++++ tests/unit_tests/test_opensandbox_pty.py | 22 ++++++++++++++++++- 3 files changed, 38 insertions(+), 2 deletions(-) diff --git a/nemo_gym/sandbox/providers/opensandbox/pty.py b/nemo_gym/sandbox/providers/opensandbox/pty.py index e789bd2482..e793374138 100644 --- a/nemo_gym/sandbox/providers/opensandbox/pty.py +++ b/nemo_gym/sandbox/providers/opensandbox/pty.py @@ -60,6 +60,12 @@ # not fractions: the stale client is torn down in the background and a later # attempt lands. _PTY_TAKEOVER_RETRY_DELAYS = (2.0, 5.0, 10.0) +# Long-held PTY sockets cross an NLB and the server proxy, either of which can +# drop them silently. Heartbeats keep intermediaries from idle-reaping the +# connection and surface a dead socket as a failed ping within about a minute +# (triggering the reattach path) instead of leaving a half-open client behind +# for the next takeover to time out against. +_PTY_WS_HEARTBEAT_S = 30.0 # Mirrors execd's shell pick (bash when available, else sh) for env-only specs. _DEFAULT_SHELL_SNIPPET = 'exec "$(command -v bash || echo sh)"' @@ -591,7 +597,10 @@ async def _connect_ws( # definitive answers (404 gone, 409 held) propagate immediately. for delay in (*_PTY_RETRY_DELAYS, None): try: - return await asyncio.wait_for(client.ws_connect(ws_url, headers=headers), timeout=request_timeout_s) + return await asyncio.wait_for( + client.ws_connect(ws_url, headers=headers, heartbeat=_PTY_WS_HEARTBEAT_S), + timeout=request_timeout_s, + ) except aiohttp.WSServerHandshakeError as e: if e.status not in (502, 503) or delay is None: raise diff --git a/resources_servers/terminal_bench_2_1/app.py b/resources_servers/terminal_bench_2_1/app.py index 50b2a7472d..57b8ebf5c3 100644 --- a/resources_servers/terminal_bench_2_1/app.py +++ b/resources_servers/terminal_bench_2_1/app.py @@ -104,6 +104,13 @@ async def _create_sandbox( async def seed_session(self, request: Request, body: TerminalBench21VerifyRequest) -> BaseSeedSessionResponse: eval_sandbox, pty_session = await self._create_sandbox(body) self._session_id_to_sandbox[request.session[SESSION_ID_KEY]] = eval_sandbox, pty_session + # Hand the terminal over without leaving our client attached: a socket + # idling through the whole agent phase is the one most likely to go + # half-open, and the next attach would then have to evict a peer that + # can no longer answer. The server-side session keeps running; a later + # user reattaches (or attaches fresh with takeover) when it needs it. + if hasattr(pty_session, "detach"): + await pty_session.detach() return TerminalBench21SeedSessionResponse(sandbox_handle=eval_sandbox._handle.sandbox_id) diff --git a/tests/unit_tests/test_opensandbox_pty.py b/tests/unit_tests/test_opensandbox_pty.py index 09175a8daf..73c1edf80a 100644 --- a/tests/unit_tests/test_opensandbox_pty.py +++ b/tests/unit_tests/test_opensandbox_pty.py @@ -125,8 +125,10 @@ def delete(self, url: str, *, headers: dict[str, str], timeout: Any = None) -> F self.delete_calls.append((url, headers)) return FakeResponse(200) - async def ws_connect(self, url: str, *, headers: dict[str, str]) -> FakeWs: + async def ws_connect(self, url: str, *, headers: dict[str, str], heartbeat: float | None = None) -> FakeWs: self.ws_calls.append((url, headers)) + self.ws_heartbeats: list[float | None] = getattr(self, "ws_heartbeats", []) + self.ws_heartbeats.append(heartbeat) if isinstance(self._ws_error, list): if self._ws_error: raise self._ws_error.pop(0) @@ -584,6 +586,24 @@ async def get_endpoint(self, port: int) -> SimpleNamespace: await session.close() +async def test_pty_sockets_dial_with_heartbeat() -> None: + # Silent half-open sockets are what takeover timeouts are made of; every + # PTY dial requests websocket heartbeats so a dead peer is detected and + # re-dialed instead of dangling until the next takeover. + from nemo_gym.sandbox.providers.opensandbox.pty import attach_pty_session + + client = FakeHttpClient(ws=FakeWs([CONNECTED])) + session = await attach_pty_session( + client=client, # type: ignore[arg-type] + base_url="http://server/base", + headers={}, + session_id="s-1", + request_timeout_s=5.0, + ) + assert client.ws_heartbeats == [pty_module._PTY_WS_HEARTBEAT_S] + await session.close() + + async def test_provider_attach_pty_retries_timed_out_takeover(monkeypatch: pytest.MonkeyPatch) -> None: # execd waits for the evicted client to acknowledge a takeover; a # half-open peer cannot, so the first attach closes 1008 while the stale From 46fd6295e56a600581fc9fd02274902e2cbb77b2 Mon Sep 17 00:00:00 2001 From: Brian Yu Date: Fri, 28 Aug 2026 15:16:32 -0700 Subject: [PATCH 3/4] remove detach Signed-off-by: Brian Yu --- resources_servers/terminal_bench_2_1/app.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/resources_servers/terminal_bench_2_1/app.py b/resources_servers/terminal_bench_2_1/app.py index 57b8ebf5c3..50b2a7472d 100644 --- a/resources_servers/terminal_bench_2_1/app.py +++ b/resources_servers/terminal_bench_2_1/app.py @@ -104,13 +104,6 @@ async def _create_sandbox( async def seed_session(self, request: Request, body: TerminalBench21VerifyRequest) -> BaseSeedSessionResponse: eval_sandbox, pty_session = await self._create_sandbox(body) self._session_id_to_sandbox[request.session[SESSION_ID_KEY]] = eval_sandbox, pty_session - # Hand the terminal over without leaving our client attached: a socket - # idling through the whole agent phase is the one most likely to go - # half-open, and the next attach would then have to evict a peer that - # can no longer answer. The server-side session keeps running; a later - # user reattaches (or attaches fresh with takeover) when it needs it. - if hasattr(pty_session, "detach"): - await pty_session.detach() return TerminalBench21SeedSessionResponse(sandbox_handle=eval_sandbox._handle.sandbox_id) From 5067394b0506e1101ed83b326f88a96c2aaca109 Mon Sep 17 00:00:00 2001 From: Brian Yu Date: Fri, 28 Aug 2026 15:21:31 -0700 Subject: [PATCH 4/4] fix formatting Signed-off-by: Brian Yu --- tests/unit_tests/test_opensandbox_pty.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/test_opensandbox_pty.py b/tests/unit_tests/test_opensandbox_pty.py index 73c1edf80a..1c4be62de4 100644 --- a/tests/unit_tests/test_opensandbox_pty.py +++ b/tests/unit_tests/test_opensandbox_pty.py @@ -621,7 +621,9 @@ async def get_endpoint(self, port: int) -> SimpleNamespace: rejected.closed = True clients = [FakeHttpClient(ws=rejected), FakeHttpClient(ws=FakeWs([CONNECTED]))] handed_out: list[FakeHttpClient] = [] - monkeypatch.setattr(provider, "_pty_http_client", lambda: handed_out.append(clients[len(handed_out)]) or handed_out[-1]) + monkeypatch.setattr( + provider, "_pty_http_client", lambda: handed_out.append(clients[len(handed_out)]) or handed_out[-1] + ) monkeypatch.setattr(pty_module, "_PTY_TAKEOVER_RETRY_DELAYS", (0.0,)) handle = SandboxHandle(sandbox_id="sb-1", provider_name="opensandbox", raw=FakeRaw()) session = await provider.attach_pty(handle, "s-7", takeover=True)