fix(opensandbox): retry a timed-out sandbox reconnect

nemo_gym/sandbox/providers/opensandbox/provider.py

@@ -1058,8 +1058,29 @@ async def connect(self, descriptor: Mapping[str, Any]) -> SandboxHandle:
         paused sandbox has no exec daemon to check; resume rebuilds its
         endpoints and performs the health check instead.
         """
-        Sandbox, _, _, _, _ = _require_opensandbox_sdk()
+        AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential = _require_tenacity()
         sandbox_id = str(descriptor["sandbox_id"])
+        # Reconnecting is read-only, so a connect timeout under load is safe to retry.
+        retry_policy = AsyncRetrying(
+            retry=retry_if_exception(_is_retryable_create_error),
+            stop=stop_after_attempt(self._create.retries + 1),
+            wait=wait_random_exponential(
+                multiplier=self._create.retry_delay_s,
+                max=self._create.retry_max_delay_s,
+            ),
+            before_sleep=lambda retry_state: _log_operation_retry(
+                retry_state, operation="connect", sandbox_id=sandbox_id
+            ),
+            reraise=True,
+        )
+        async for attempt in retry_policy:
+            with attempt:
+                return await self._connect_once(sandbox_id)
+
+        raise RuntimeError("OpenSandbox connect retry loop did not run")
+
+    async def _connect_once(self, sandbox_id: str) -> SandboxHandle:
+        Sandbox, _, _, _, _ = _require_opensandbox_sdk()
         timeout_s = self._create.connect_attempt_timeout_s
         # A cancelled SDK call skips its own transport cleanup; see resume().
         config = self._connection_config(request_timeout_s=timeout_s).with_transport_if_missing()

tests/unit_tests/test_opensandbox_provider.py

@@ -1955,7 +1955,7 @@ async def test_connect_preserves_failure_when_cleanup_stalls(
     close = AsyncMock(side_effect=asyncio.Event().wait)
     monkeypatch.setattr(FakeSandbox, "close", close)
     provider = opensandbox_provider.OpenSandboxProvider(
-        create={"connect_attempt_timeout_s": 0.05},
+        create={"connect_attempt_timeout_s": 0.05, "retries": 0},
         operations={"close_timeout_s": 0.05},
         probe={"command": None},
     )
@@ -2161,6 +2161,63 @@ async def test_resume_installs_the_new_handle_before_closing_the_old_one(fake_op
     raw.close.assert_awaited_once_with()
 
 
+class FlakyConnectSandbox(FakeSandbox):
+    failures: list[BaseException] = []
+    calls = 0
+
+    @classmethod
+    async def connect(cls, *args: Any, **kwargs: Any) -> "FakeSandbox":
+        cls.calls += 1
+        if cls.failures:
+            raise cls.failures.pop(0)
+        return cls(str(args[0]))
+
+
+@pytest.fixture
+def flaky_connect_sdk(monkeypatch: pytest.MonkeyPatch) -> type[FlakyConnectSandbox]:
+    FlakyConnectSandbox.failures = []
+    FlakyConnectSandbox.calls = 0
+    monkeypatch.setattr(
+        opensandbox_provider,
+        "_require_opensandbox_sdk",
+        lambda: (FlakyConnectSandbox, FakeConnectionConfig, object, FakePlatformSpec, object),
+    )
+    return FlakyConnectSandbox
+
+
+async def test_connect_retries_a_connect_timeout(flaky_connect_sdk: type[FlakyConnectSandbox]) -> None:
+    """One connect timeout under load used to fail the whole rollout."""
+    flaky_connect_sdk.failures = [TimeoutError(110, "Connect call failed ('10.0.0.1', 80)")]
+    provider = opensandbox_provider.OpenSandboxProvider(create={"retries": 2, "retry_delay_s": 0})
+
+    handle = await provider.connect({"sandbox_id": "sandbox-9"})
+
+    assert handle.sandbox_id == "sandbox-9"
+    assert flaky_connect_sdk.calls == 2
+
+
+async def test_connect_gives_up_after_create_retries(flaky_connect_sdk: type[FlakyConnectSandbox]) -> None:
+    flaky_connect_sdk.failures = [ConnectionError("connection refused")] * 3
+    provider = opensandbox_provider.OpenSandboxProvider(create={"retries": 1, "retry_delay_s": 0})
+
+    with pytest.raises(ConnectionError):
+        await provider.connect({"sandbox_id": "sandbox-9"})
+    assert flaky_connect_sdk.calls == 2
+
+
+async def test_connect_does_not_retry_a_missing_sandbox(flaky_connect_sdk: type[FlakyConnectSandbox]) -> None:
+    from opensandbox.exceptions import SandboxApiException  # noqa: PLC0415
+
+    not_found = SandboxApiException("not found")
+    not_found.status_code = 404
+    flaky_connect_sdk.failures = [not_found]
+    provider = opensandbox_provider.OpenSandboxProvider(create={"retries": 2, "retry_delay_s": 0})
+
+    with pytest.raises(SandboxApiException):
+        await provider.connect({"sandbox_id": "sandbox-9"})
+    assert flaky_connect_sdk.calls == 1
+
+
 @pytest.mark.asyncio
 async def test_exec_retries_backend_connect_502_despite_zero_command_retries(
     monkeypatch: pytest.MonkeyPatch,