Retry PTY takeover attach when execd's eviction times out

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

nemo_gym/sandbox/providers/opensandbox/pty.py

@@ -55,6 +55,17 @@
 # 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)
+# 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)"'
@@ -586,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

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,78 @@ 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
+    # 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)