diff --git a/nemo_gym/sandbox/providers/opensandbox/configs/opensandbox.yaml b/nemo_gym/sandbox/providers/opensandbox/configs/opensandbox.yaml index 7c187bbed7..0d8dd07120 100644 --- a/nemo_gym/sandbox/providers/opensandbox/configs/opensandbox.yaml +++ b/nemo_gym/sandbox/providers/opensandbox/configs/opensandbox.yaml @@ -26,10 +26,25 @@ sandbox: protocol: http request_timeout_s: 300 use_server_proxy: true + # Must stay below the server's keep-alive idle timeout (uvicorn: ~5s), + # else pooled sockets are reused after the server has closed them. + keepalive_expiry_s: 3.0 + # 0 disables connection reuse. + max_keepalive_connections: 20 + # null for no cap. + max_connections: 100 + connect_retries: 2 + # "aiohttp" routes through the optional httpx-aiohttp bridge, falling + # back to httpx with a warning when it is not installed. + transport_backend: httpx create: request_timeout_s: 1200 - timeout_s: 1200 - skip_health_check: true + # Must exceed the spec's ready_timeout_s (1200 for mini_swe_agent_2): this + # bounds the whole create call, which now includes the readiness wait. + timeout_s: 1500 + # Skipping the check lets the first command race a pod whose exec daemon + # is not listening yet, which returns a 502 and kills the rollout. + skip_health_check: false retries: 10 retry_delay_s: 5.0 retry_max_delay_s: 90.0 @@ -42,6 +57,8 @@ sandbox: retries: 5 retry_delay_s: 1.0 retry_max_delay_s: 45.0 + # A retried command the server already started runs twice; agent commands + # are usually mutating. Raise only for idempotent workloads. command_retries: 0 close_timeout_s: 30 # Job attribution merged into every sandbox's metadata. OpenSandbox propagates diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 2070fa2421..9cefd4d5dd 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -348,13 +348,27 @@ def _to_sandbox_status(state: Any) -> SandboxStatus: @dataclass(frozen=True) class OpenSandboxConnectionConfig: - """OpenSandbox server connection settings.""" + """OpenSandbox server connection settings. + + ``keepalive_expiry_s`` must stay below the server's own keep-alive idle + timeout (uvicorn defaults to 5s), or pooled sockets are reused after the + server has closed them; null falls back to the SDK's default transport. + ``transport_backend`` is "httpx" or "aiohttp" (via the optional + ``httpx-aiohttp`` bridge, falling back to httpx when it is absent). + The pool is shared, so ``max_connections`` also caps in-flight sandbox + operations per process; null means no cap. + """ domain: str | None = None api_key: str | None = None protocol: str | None = None request_timeout_s: int | None = None use_server_proxy: bool = False + keepalive_expiry_s: float | None = 3.0 + max_keepalive_connections: int = 20 + max_connections: int | None = 100 + connect_retries: int = 2 + transport_backend: str = "httpx" @dataclass(frozen=True) @@ -488,6 +502,9 @@ class OpenSandboxProviderOptions: volumes: tuple[Mapping[str, Any], ...] = () skip_health_check: bool | None = None extensions: Mapping[str, str] = field(default_factory=dict) + # Scheduling requests (same keys as SandboxSpec.resources, which become the + # limits). Unset, the server applies the single resources map as both. + resource_requests: Mapping[str, Any] | None = None @classmethod def from_mapping(cls, options: Mapping[str, Any] | None) -> "OpenSandboxProviderOptions": @@ -522,6 +539,9 @@ def from_mapping(cls, options: Mapping[str, Any] | None) -> "OpenSandboxProvider extensions = options.get("extensions", {}) if not isinstance(extensions, Mapping): raise TypeError("OpenSandbox provider option 'extensions' must be a mapping") + resource_requests = options.get("resource_requests") + if resource_requests is not None and not isinstance(resource_requests, Mapping): + raise TypeError("OpenSandbox provider option 'resource_requests' must be a mapping") return cls( image_auth=dict(image_auth) if image_auth is not None else None, @@ -530,6 +550,7 @@ def from_mapping(cls, options: Mapping[str, Any] | None) -> "OpenSandboxProvider volumes=tuple(dict(volume) for volume in volumes), skip_health_check=skip_health_check, extensions=_string_map(dict(extensions)), + resource_requests=dict(resource_requests) if resource_requests is not None else None, ) @@ -552,6 +573,10 @@ def __init__( self._probe = _coerce_config(probe, OpenSandboxProbeConfig) self._operations = _coerce_config(operations, OpenSandboxOperationConfig) self._attribution = _coerce_config(attribution, OpenSandboxAttributionConfig) + # Shared injected transport. The SDK never closes transports it did not + # create, so the provider owns this one: built once, reused by every + # ConnectionConfig, closed in aclose(). + self._transport: Any | None = None def _resolve_extensions(self, extensions: Mapping[str, str]) -> dict[str, str]: """Add the configured default image pull policy to SDK create extensions.""" @@ -587,11 +612,42 @@ def _connection_config( kwargs["request_timeout"] = timedelta(seconds=request_timeout_s) if self._connection.use_server_proxy: kwargs["use_server_proxy"] = True + if self._connection.keepalive_expiry_s is not None: + kwargs["transport"] = self._get_transport() return ConnectionConfig(**kwargs) + def _get_transport(self) -> Any: + """Return the provider-owned shared transport, building it on first use.""" + if self._transport is None: + self._transport = self._build_transport() + return self._transport + + def _build_transport(self) -> Any: + """Build the SDK transport with the configured pool limits.""" + import httpx + + limits = httpx.Limits( + max_connections=self._connection.max_connections, + max_keepalive_connections=self._connection.max_keepalive_connections, + keepalive_expiry=self._connection.keepalive_expiry_s, + ) + if self._connection.transport_backend == "aiohttp": + try: + from httpx_aiohttp import AiohttpTransport + + return AiohttpTransport(limits=limits, retries=self._connection.connect_retries) + except ImportError: + LOGGER.warning( + "connection.transport_backend=aiohttp requested but httpx-aiohttp " + "is not installed; falling back to the httpx transport" + ) + return httpx.AsyncHTTPTransport(limits=limits, retries=self._connection.connect_retries) + async def aclose(self) -> None: """Close provider-owned resources.""" - return None + transport, self._transport = self._transport, None + if transport is not None: + await transport.aclose() async def serialize_handle(self, handle: SandboxHandle, *, scope: str | None = None) -> dict[str, Any]: """Return a descriptor for reattaching to this sandbox by id. @@ -809,6 +865,8 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: "extensions": self._resolve_extensions(options.extensions), "connection_config": self._connection_config(request_timeout_s=self._create.request_timeout_s), } + if options.resource_requests is not None: + kwargs["resource_requests"] = _resource_map(SandboxResources.from_mapping(options.resource_requests)) if spec.image is not None: kwargs["image"] = _to_image_spec(spec.image, options.image_auth) if options.snapshot_id is not None: diff --git a/pyproject.toml b/pyproject.toml index 5d116063a4..c61fcadcad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -235,9 +235,10 @@ sandbox = [ "tenacity>=9.1.4", # OpenSandbox SDK: used by the OpenSandbox sandbox provider for create/exec/delete and SDK pool creation. - # Updated: Sat May 16, 2026 with opensandbox>=0.1.9 + # Lower bound 0.1.15: first version exposing separate `resource_requests` on Sandbox.create. + # Updated: Thu Jul 30, 2026 with opensandbox>=0.1.15 # License: Apache 2.0 - "opensandbox>=0.1.9", + "opensandbox>=0.1.15", # OpenShell SDK: used by the OpenShell sandbox provider for gateway create/exec/delete over gRPC. # Lower bound 0.0.92: the version that made `workspace` a required argument on sandbox diff --git a/responses_api_agents/mini_swe_agent_2/app.py b/responses_api_agents/mini_swe_agent_2/app.py index 248185797f..4073cdf889 100644 --- a/responses_api_agents/mini_swe_agent_2/app.py +++ b/responses_api_agents/mini_swe_agent_2/app.py @@ -12,11 +12,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import asyncio import hashlib import json import os import sys +import threading import time import traceback from asyncio import Semaphore @@ -89,6 +89,10 @@ class MiniSWEAgentVerifyResponse(BaseVerifyResponse): @ray.remote( + # Rollout tasks spend nearly all their time waiting on LLM calls and + # sandbox I/O; reserving a full CPU per task caps concurrent rollouts at + # the Ray cluster's core count long before any real resource limit. + num_cpus=0.25, scheduling_strategy="SPREAD", runtime_env={ "py_executable": sys.executable, @@ -509,6 +513,11 @@ def _run_mini_swe_v2(**params: Any) -> dict[str, Any]: model_kwargs = model_config.setdefault("model_kwargs", {}) model_kwargs["api_key"] = params["api_key"] model_kwargs["base_url"] = params["base_url"] + # Bounded retries for transient LLM-call failures (disconnects, resets): + # without any retry a single failed call kills the whole rollout, while a + # large value makes litellm retry silently for so long that the rollout + # looks hung. Config-provided model_kwargs take precedence. + model_kwargs.setdefault("num_retries", 5) model_kwargs.pop("api_base", None) max_output_tokens = model_kwargs.pop("max_output_tokens", None) if max_output_tokens is not None and "max_tokens" not in model_kwargs: @@ -578,7 +587,18 @@ def _run_mini_swe_v2(**params: Any) -> dict[str, Any]: } finally: if env and hasattr(env, "cleanup"): - env.cleanup() + # Off the critical path: this finally block runs before the task's + # return value becomes fetchable, so an in-band stop() delays every + # finished result and, on failure, re-raises over it. Orphans are + # covered by the provider's sandbox TTL. + threading.Thread(target=_cleanup_env_best_effort, args=(env,), daemon=True).start() + + +def _cleanup_env_best_effort(env: Any) -> None: + try: + env.cleanup() + except Exception as e: + print(f"[CLEANUP] best-effort sandbox teardown failed: {e}", flush=True) def run_mini_swe_with_sandbox(**params: Any) -> Any: @@ -807,7 +827,12 @@ async def run(self, body: MiniSWEAgentRunRequest) -> MiniSWEAgentVerifyResponse: if runtime_env.get("env_vars"): runner = runner.options(runtime_env=runtime_env) future = runner.remote(run_mini_swe_with_sandbox, params) - result = await asyncio.to_thread(ray.get, future) + # Ray ObjectRefs are awaitable: park on the event loop instead + # of pinning a thread in asyncio's default executor (capped at + # min(32, cpu+4) workers). With the thread-blocking ray.get, + # at most ~32 rollouts can be waiting on results at once and + # every other finished task queues behind them. + result = await future result = result[instance_id] input_messages = result["input_messages"] response_output = result["response_output"] diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml index 00d61d4489..3f4bcafc2e 100644 --- a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml @@ -20,6 +20,9 @@ mini_swe_agent_2: type: responses_api_models name: policy_model concurrency: 64 + # Declared (empty) so benchmark configs can inherit this server config and + # fill in their dataset list; the config merge rejects undeclared keys. + datasets: [] env: sandbox # Name of the sandbox to use; include a provider config that defines a # `sandbox` block (e.g. nemo_gym/sandbox/providers/opensandbox/configs/opensandbox.yaml). @@ -27,13 +30,19 @@ mini_swe_agent_2: sandbox_spec: ttl_s: 18000 ready_timeout_s: 1200 + # Limits (burst ceiling); memory-spiky test suites OOM-kill below this. resources: cpu: 2 memory_mib: 8192 # 30 GiB works on every provider: within Fargate's 21-200 GiB ephemeral range # (an explicit 20 is rejected there) and fine as an ephemeral request elsewhere. disk_gib: 30 - provider_options: {} + provider_options: + # Scheduling requests, kept below the limits so sandboxes pack densely. + resource_requests: + cpu: 0.5 + memory_mib: 2048 + disk_gib: 30 metadata: benchmark: swebench-verified harness: mini-swe-agent diff --git a/responses_api_agents/mini_swe_agent_2/tests/test_app.py b/responses_api_agents/mini_swe_agent_2/tests/test_app.py index 2e522aeb57..93a418157f 100644 --- a/responses_api_agents/mini_swe_agent_2/tests/test_app.py +++ b/responses_api_agents/mini_swe_agent_2/tests/test_app.py @@ -167,8 +167,23 @@ def setup_config_path_mock(mock_get_config_path, config_yaml: str = DEFAULT_CONF mock_get_config_path.return_value = mock_config_path +class FakeObjectRef: + """Stand-in for a Ray ObjectRef, which the agent awaits directly.""" + + def __init__(self, result: Any = None, error: BaseException | None = None): + self._result = result + self._error = error + + def __await__(self): + async def _resolve() -> Any: + if self._error is not None: + raise self._error + return self._result + + return _resolve().__await__() + + def setup_run_mini_swe_mock( - mock_to_thread, mock_runner_ray_remote, run_mini_swe_result: Dict[str, Any] = None, ): @@ -176,14 +191,11 @@ def setup_run_mini_swe_mock( if run_mini_swe_result is None: run_mini_swe_result = DEFAULT_RUN_MINI_SWE_RESULT - # Mock the Ray remote function to return a future-like object - mock_future = MagicMock() + # The Ray remote call returns an awaitable ref that resolves to the result. + mock_future = FakeObjectRef(run_mini_swe_result) mock_runner_ray_remote.remote.return_value = mock_future mock_runner_ray_remote.options.return_value.remote.return_value = mock_future - # Mock asyncio.to_thread (which calls ray.get) to return the result - mock_to_thread.return_value = run_mini_swe_result - def create_run_request( instance_id: str = "test_instance_123", @@ -248,15 +260,16 @@ def assert_run_response( def assert_run_mini_swe_called( - mock_to_thread, + mock_runner_ray_remote, subset: str = "gym", split: str = "train", instance_id: str = "test_instance_123", ): - mock_to_thread.assert_called_once() - call_args = mock_to_thread.call_args - args = call_args[0] - assert len(args) >= 1 + mock_runner_ray_remote.remote.assert_called_once() + params = mock_runner_ray_remote.remote.call_args.args[1] + assert params["subset"] == subset + assert params["split"] == split + assert params["instance_id"] == instance_id class TestApp: @@ -701,10 +714,8 @@ def get_model(config: dict[str, Any]) -> SimpleNamespace: @patch("responses_api_agents.mini_swe_agent_2.app.get_first_server_config_dict") @patch("responses_api_agents.mini_swe_agent_2.app.get_config_path") @patch("responses_api_agents.mini_swe_agent_2.app.runner_ray_remote") - @patch("asyncio.to_thread") async def test_run_successful_execution( self, - mock_to_thread, mock_runner_ray_remote, mock_get_config_path, mock_get_first_server_config_dict, @@ -720,7 +731,7 @@ async def test_run_successful_execution( setup_server_client_mocks(mock_load_from_global_config, mock_get_first_server_config_dict) setup_config_path_mock(mock_get_config_path) - setup_run_mini_swe_mock(mock_to_thread, mock_runner_ray_remote) + setup_run_mini_swe_mock(mock_runner_ray_remote) run_request = MiniSWEAgentRunRequest.model_validate( create_run_request().model_dump() | {TASK_INDEX_KEY_NAME: 2, ROLLOUT_INDEX_KEY_NAME: 1} @@ -730,17 +741,15 @@ async def test_run_successful_execution( assert_run_response(response) - assert_run_mini_swe_called(mock_to_thread) + assert_run_mini_swe_called(mock_runner_ray_remote) assert mock_runner_ray_remote.remote.call_args.args[1]["base_url"] == ("http://0.0.0.0:8080/ng-rollout/2-1/v1") @patch("responses_api_agents.mini_swe_agent_2.app.ServerClient.load_from_global_config") @patch("responses_api_agents.mini_swe_agent_2.app.get_first_server_config_dict") @patch("responses_api_agents.mini_swe_agent_2.app.get_config_path") @patch("responses_api_agents.mini_swe_agent_2.app.runner_ray_remote") - @patch("asyncio.to_thread") async def test_run_writes_generation_params_to_config( self, - mock_to_thread, mock_runner_ray_remote, mock_get_config_path, mock_get_first_server_config_dict, @@ -764,7 +773,7 @@ async def test_run_writes_generation_params_to_config( setup_server_client_mocks(mock_load_from_global_config, mock_get_first_server_config_dict) setup_config_path_mock(mock_get_config_path) - setup_run_mini_swe_mock(mock_to_thread, mock_runner_ray_remote) + setup_run_mini_swe_mock(mock_runner_ray_remote) run_request = create_run_request( temperature=0.6, @@ -802,10 +811,8 @@ async def test_run_writes_generation_params_to_config( @patch("responses_api_agents.mini_swe_agent_2.app.get_first_server_config_dict") @patch("responses_api_agents.mini_swe_agent_2.app.get_config_path") @patch("responses_api_agents.mini_swe_agent_2.app.runner_ray_remote") - @patch("asyncio.to_thread") async def test_run_resolves_named_sandbox_provider_reference( self, - mock_to_thread, mock_runner_ray_remote, mock_get_config_path, mock_get_first_server_config_dict, @@ -835,7 +842,7 @@ async def test_run_resolves_named_sandbox_provider_reference( mock_load_from_global_config.return_value = mock_server_client_instance mock_get_first_server_config_dict.return_value = {"host": "0.0.0.0", "port": 8080} setup_config_path_mock(mock_get_config_path) - setup_run_mini_swe_mock(mock_to_thread, mock_runner_ray_remote) + setup_run_mini_swe_mock(mock_runner_ray_remote) await server.run(create_run_request()) @@ -854,10 +861,8 @@ async def test_run_resolves_named_sandbox_provider_reference( @patch("responses_api_agents.mini_swe_agent_2.app.get_first_server_config_dict") @patch("responses_api_agents.mini_swe_agent_2.app.get_config_path") @patch("responses_api_agents.mini_swe_agent_2.app.runner_ray_remote") - @patch("asyncio.to_thread") async def test_run_failed_execution( self, - mock_to_thread, mock_runner_ray_remote, mock_get_config_path, mock_get_first_server_config_dict, @@ -872,12 +877,8 @@ async def test_run_failed_execution( setup_server_client_mocks(mock_load_from_global_config, mock_get_first_server_config_dict) setup_config_path_mock(mock_get_config_path) - # Mock Ray remote function - mock_future = MagicMock() - mock_runner_ray_remote.remote.return_value = mock_future - - # Mock asyncio.to_thread (ray.get) to raise an exception - mock_to_thread.side_effect = Exception("run_mini_swe failed") + # Awaiting the Ray result raises, standing in for a failed rollout task. + mock_runner_ray_remote.remote.return_value = FakeObjectRef(error=Exception("run_mini_swe failed")) run_request = create_run_request(instance_id="test_instance_456", temperature=0.3, top_p=0.95) @@ -891,16 +892,14 @@ async def test_run_failed_execution( expected_input_length=0, ) - assert_run_mini_swe_called(mock_to_thread, instance_id="test_instance_456") + assert_run_mini_swe_called(mock_runner_ray_remote, instance_id="test_instance_456") @patch("responses_api_agents.mini_swe_agent_2.app.ServerClient.load_from_global_config") @patch("responses_api_agents.mini_swe_agent_2.app.get_first_server_config_dict") @patch("responses_api_agents.mini_swe_agent_2.app.get_config_path") @patch("responses_api_agents.mini_swe_agent_2.app.runner_ray_remote") - @patch("asyncio.to_thread") async def test_run_mini_swe_not_found( self, - mock_to_thread, mock_runner_ray_remote, mock_get_config_path, mock_get_first_server_config_dict, @@ -913,12 +912,7 @@ async def test_run_mini_swe_not_found( setup_server_client_mocks(mock_load_from_global_config, mock_get_first_server_config_dict) setup_config_path_mock(mock_get_config_path) - # Mock Ray remote function - mock_future = MagicMock() - mock_runner_ray_remote.remote.return_value = mock_future - - # Mock asyncio.to_thread (ray.get) to raise FileNotFoundError - mock_to_thread.side_effect = FileNotFoundError("run_mini_swe not found") + mock_runner_ray_remote.remote.return_value = FakeObjectRef(error=FileNotFoundError("run_mini_swe not found")) run_request = create_run_request(instance_id="test_instance_789", temperature=0.2, top_p=1.0) @@ -932,7 +926,7 @@ async def test_run_mini_swe_not_found( expected_input_length=0, ) - assert_run_mini_swe_called(mock_to_thread, instance_id="test_instance_789") + assert_run_mini_swe_called(mock_runner_ray_remote, instance_id="test_instance_789") async def test_responses_not_implemented(self) -> None: config = create_test_config() diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index 5c6589e318..2f977f1c19 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -15,12 +15,14 @@ import asyncio import builtins +import sys from dataclasses import dataclass from datetime import timedelta from pathlib import Path from types import SimpleNamespace from typing import Any +import httpx import pytest from nemo_gym.sandbox.providers.base import SandboxResources, SandboxSpec, SandboxStatus @@ -189,6 +191,38 @@ async def test_direct_create_passes_platform_to_sdk_create( ) +async def test_direct_create_passes_resource_requests_to_sdk_create( + fake_opensandbox_sdk: None, +) -> None: + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + + await provider.create( + SandboxSpec( + image="mirror.gcr.io/astral/uv:python3.12-bookworm-slim", + resources={"cpu": 1, "memory_mib": 8192, "disk_gib": 30}, + provider_options={"resource_requests": {"cpu": 0.5, "memory_mib": 2048, "disk_gib": 30}}, + ), + ) + + assert FakeSandbox.created_kwargs["resource"] == {"cpu": "1", "memory": "8192Mi", "ephemeral-storage": "30Gi"} + assert FakeSandbox.created_kwargs["resource_requests"] == { + "cpu": "0.5", + "memory": "2048Mi", + "ephemeral-storage": "30Gi", + } + + with pytest.raises(TypeError, match="'resource_requests' must be a mapping"): + opensandbox_provider.OpenSandboxProviderOptions.from_mapping({"resource_requests": "big"}) + + with pytest.raises(ValueError, match="Unknown sandbox resource keys"): + await provider.create( + SandboxSpec( + image="mirror.gcr.io/astral/uv:python3.12-bookworm-slim", + provider_options={"resource_requests": {"memory_gib": 2}}, + ), + ) + + async def test_direct_create_passes_image_auth_to_sdk_create( fake_opensandbox_sdk: None, ) -> None: @@ -310,6 +344,8 @@ def test_connection_config_and_image_policy(fake_opensandbox_sdk: None) -> None: ) config = provider._connection_config() + transport = config.kwargs.pop("transport") + assert isinstance(transport, httpx.AsyncBaseTransport) assert config.kwargs == { "domain": "sandbox.example", "api_key": "key", # pragma: allowlist secret @@ -320,6 +356,84 @@ def test_connection_config_and_image_policy(fake_opensandbox_sdk: None) -> None: short_timeout_config = provider._connection_config(request_timeout_s=3) assert short_timeout_config.kwargs["request_timeout"] == timedelta(seconds=3) + +def test_connection_transport_backends(fake_opensandbox_sdk: None, monkeypatch: pytest.MonkeyPatch) -> None: + # Default backend is httpx, with the configured keepalive expiry on the pool. + provider = opensandbox_provider.OpenSandboxProvider() + transport = provider._build_transport() + assert isinstance(transport, httpx.AsyncHTTPTransport) + + # Custom pool settings still produce an httpx transport. + provider = opensandbox_provider.OpenSandboxProvider( + connection={ + "transport_backend": "httpx", + "keepalive_expiry_s": 2.5, + "max_connections": 7, + "max_keepalive_connections": 3, + "connect_retries": 1, + } + ) + transport = provider._build_transport() + assert isinstance(transport, httpx.AsyncHTTPTransport) + # connect_retries reaches the pool rather than silently falling back. + assert transport._pool._retries == 1 + + # aiohttp requested but httpx-aiohttp unavailable: falls back to httpx. + with pytest.MonkeyPatch.context() as mp: + mp.setitem(sys.modules, "httpx_aiohttp", None) + provider = opensandbox_provider.OpenSandboxProvider(connection={"transport_backend": "aiohttp"}) + transport = provider._build_transport() + assert isinstance(transport, httpx.AsyncHTTPTransport) + + # keepalive_expiry_s=null disables transport injection entirely. + provider = opensandbox_provider.OpenSandboxProvider(connection={"keepalive_expiry_s": None}) + config = provider._connection_config() + assert "transport" not in config.kwargs + + # max_connections=null uncaps the pool; max_keepalive_connections=0 disables reuse. + provider = opensandbox_provider.OpenSandboxProvider( + connection={"max_connections": None, "max_keepalive_connections": 0} + ) + transport = provider._build_transport() + assert isinstance(transport, httpx.AsyncHTTPTransport) + assert transport._pool._max_connections > 2**32 + assert transport._pool._max_keepalive_connections == 0 + + +async def test_connection_transport_is_shared_and_closed_by_provider(fake_opensandbox_sdk: None) -> None: + # The SDK never closes a transport it did not create, so the provider owns + # one: built on first use, reused by every ConnectionConfig rather than + # leaking a pool per call, and closed in aclose(). + class FakeTransport: + def __init__(self) -> None: + self.aclosed = False + + async def aclose(self) -> None: + self.aclosed = True + + provider = opensandbox_provider.OpenSandboxProvider() + provider._build_transport = FakeTransport + + transport = provider._connection_config().kwargs["transport"] + assert provider._connection_config().kwargs["transport"] is transport + + await provider.aclose() + assert transport.aclosed + assert provider._transport is None + + +def test_connection_transport_backend_aiohttp_opt_in(fake_opensandbox_sdk: None) -> None: + # Opt-in aiohttp backend via the httpx-aiohttp bridge; the package is not a + # declared dependency, so this coverage only runs where it is installed. + httpx_aiohttp = pytest.importorskip("httpx_aiohttp", reason="optional httpx-aiohttp is not installed") + provider = opensandbox_provider.OpenSandboxProvider(connection={"transport_backend": "aiohttp"}) + transport = provider._build_transport() + assert isinstance(transport, httpx_aiohttp.AiohttpTransport) + assert transport.limits.keepalive_expiry == 3.0 + # Both backends honor connect_retries; the bridge default is 0, so this + # would catch the option being dropped on the aiohttp path. + assert transport.retries == 2 + extensions = provider._resolve_extensions({"imagePullPolicy": "Never"}) assert extensions["imagePullPolicy"] == "Never" assert extensions["opensandbox.extensions.image-pull-policy"] == "Never" diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index 16b946007b..ed17d3507e 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -781,7 +781,10 @@ async def connect(cls, sandbox_id: str, **kwargs: Any) -> "FakeSDKSandbox": assert isinstance(handle.raw, FakeSDKSandbox) connect_call = FakeSDKSandbox.connect_calls[0] assert connect_call["skip_health_check"] is True - assert connect_call["connection_config"].kwargs == { + connection_kwargs = dict(connect_call["connection_config"].kwargs) + # Transport identity is asserted in test_opensandbox_provider.py. + connection_kwargs.pop("transport", None) + assert connection_kwargs == { "domain": "sandbox.example", "protocol": "https", "request_timeout": timedelta(seconds=300), diff --git a/uv.lock b/uv.lock index f93e9c8bc2..8cce6c3779 100644 --- a/uv.lock +++ b/uv.lock @@ -1800,7 +1800,7 @@ requires-dist = [ { name = "nemo-gym", extras = ["dev", "sandbox"], marker = "extra == 'all'" }, { name = "omegaconf" }, { name = "openai", specifier = "<=2.7.2" }, - { name = "opensandbox", marker = "extra == 'sandbox'", specifier = ">=0.1.9" }, + { name = "opensandbox", marker = "extra == 'sandbox'", specifier = ">=0.1.15" }, { name = "openshell", marker = "extra == 'sandbox'", specifier = ">=0.0.92,<0.1" }, { name = "orjson" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.6.0" }, @@ -2006,7 +2006,7 @@ wheels = [ [[package]] name = "opensandbox" -version = "0.1.9" +version = "0.1.15" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "attrs" }, @@ -2014,9 +2014,9 @@ dependencies = [ { name = "pydantic" }, { name = "python-dateutil" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/5a/2a/ab3cc141e041f71a373c97fcda8749dba9328f1b9bf80401378c0611556f/opensandbox-0.1.9.tar.gz", hash = "sha256:670fbf292c498f8467963d21e91ade9ea8b8f63f4ef18d18fff9581e0952ec03", size = 160034, upload-time = "2026-05-12T12:27:20.692Z" } +sdist = { url = "https://files.pythonhosted.org/packages/43/21/654a3d69815b09690e926d553f3f4a178640d1206000a1b49f5e22c8eb68/opensandbox-0.1.15.tar.gz", hash = "sha256:017abc9b399b88da51bf077d6fb94ee89b1783e601495e85ba85fed478fed1b0", size = 228729, upload-time = "2026-07-24T09:36:29.988Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4e/9b/553f8d7a30eddb12785711b2a1c682386878e2bb95450acd806f9fa62930/opensandbox-0.1.9-py3-none-any.whl", hash = "sha256:17faed35b60a982fee5a643fed8e4e12f041e5432d5ea0665d2828d1f2082759", size = 360945, upload-time = "2026-05-12T12:27:19.465Z" }, + { url = "https://files.pythonhosted.org/packages/d5/5c/ab87ea696531210790feb8f575471036fccba29dedaff201736f15bbb3a7/opensandbox-0.1.15-py3-none-any.whl", hash = "sha256:992b01490551f4d8e3f99caa25e34cb9d1690f0c5027eeebab912738291957d1", size = 538522, upload-time = "2026-07-24T09:36:28.277Z" }, ] [[package]]