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..e793374138 100644 --- a/nemo_gym/sandbox/providers/opensandbox/pty.py +++ b/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 diff --git a/tests/unit_tests/test_opensandbox_pty.py b/tests/unit_tests/test_opensandbox_pty.py index a0163ec7c7..1c4be62de4 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,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)