diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 7cfbad76840..0d2df3ee5ad 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/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() diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index 5596e548b8b..737d57a1da5 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/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,