From 0aace6afe0a3a1a0b261d1112b1542c6ff874652 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Wed, 20 May 2026 13:13:02 -0700 Subject: [PATCH 01/14] Add sandbox API and mini SWE agent 2 integration Signed-off-by: Hemil Desai --- nemo_gym/sandbox/__init__.py | 48 + nemo_gym/sandbox/api.py | 305 ++++ nemo_gym/sandbox/providers/__init__.py | 46 + nemo_gym/sandbox/providers/base.py | 136 ++ .../sandbox/providers/opensandbox/__init__.py | 42 + .../sandbox/providers/opensandbox/provider.py | 1299 +++++++++++++++++ nemo_gym/sandbox/providers/registry.py | 74 + pyproject.toml | 25 +- .../mini_swe_agent_2/.gitignore | 1 + .../mini_swe_agent_2/README.md | 347 +++++ .../mini_swe_agent_2/__init__.py | 0 responses_api_agents/mini_swe_agent_2/app.py | 790 ++++++++++ .../configs/mini_swe_agent_opensandbox.yaml | 64 + .../mini_swe_agent_2/requirements.txt | 3 + .../mini_swe_agent_2/sandbox_environment.py | 199 +++ .../mini_swe_agent_2/tests/test_app.py | 916 ++++++++++++ .../tests/test_sandbox_environment.py | 51 + tests/unit_tests/test_opensandbox_provider.py | 604 ++++++++ tests/unit_tests/test_sandbox.py | 766 ++++++++++ uv.lock | 32 +- 20 files changed, 5746 insertions(+), 2 deletions(-) create mode 100644 nemo_gym/sandbox/__init__.py create mode 100644 nemo_gym/sandbox/api.py create mode 100644 nemo_gym/sandbox/providers/__init__.py create mode 100644 nemo_gym/sandbox/providers/base.py create mode 100644 nemo_gym/sandbox/providers/opensandbox/__init__.py create mode 100644 nemo_gym/sandbox/providers/opensandbox/provider.py create mode 100644 nemo_gym/sandbox/providers/registry.py create mode 100644 responses_api_agents/mini_swe_agent_2/.gitignore create mode 100644 responses_api_agents/mini_swe_agent_2/README.md create mode 100644 responses_api_agents/mini_swe_agent_2/__init__.py create mode 100644 responses_api_agents/mini_swe_agent_2/app.py create mode 100644 responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml create mode 100644 responses_api_agents/mini_swe_agent_2/requirements.txt create mode 100644 responses_api_agents/mini_swe_agent_2/sandbox_environment.py create mode 100644 responses_api_agents/mini_swe_agent_2/tests/test_app.py create mode 100644 responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py create mode 100644 tests/unit_tests/test_opensandbox_provider.py create mode 100644 tests/unit_tests/test_sandbox.py diff --git a/nemo_gym/sandbox/__init__.py b/nemo_gym/sandbox/__init__.py new file mode 100644 index 0000000000..cd9e1cff9e --- /dev/null +++ b/nemo_gym/sandbox/__init__.py @@ -0,0 +1,48 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Public sandbox API for NeMo Gym.""" + +from nemo_gym.sandbox.api import AsyncSandbox, Sandbox, rewrite_image +from nemo_gym.sandbox.providers import ( + SandboxBatchCreateError, + SandboxCreateError, + SandboxCreateVerificationError, + SandboxExecResult, + SandboxHandle, + SandboxProvider, + SandboxSpec, + create_provider, + get_provider_class, + list_providers, + register_provider, +) + + +__all__ = [ + "Sandbox", + "AsyncSandbox", + "SandboxBatchCreateError", + "SandboxCreateError", + "SandboxCreateVerificationError", + "SandboxExecResult", + "SandboxHandle", + "SandboxProvider", + "SandboxSpec", + "create_provider", + "get_provider_class", + "list_providers", + "register_provider", + "rewrite_image", +] diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py new file mode 100644 index 0000000000..39444814a3 --- /dev/null +++ b/nemo_gym/sandbox/api.py @@ -0,0 +1,305 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Provider-neutral public sandbox API. + +This module is the boundary Gym code should use when it needs a sandbox. +Provider packages implement the lower-level async protocol; callers use +``AsyncSandbox`` in async code and ``Sandbox`` in synchronous integrations. +""" + +import asyncio +import threading +from collections.abc import Awaitable, Callable, Mapping +from concurrent.futures import Future +from pathlib import Path +from typing import Any, TypeVar + +from nemo_gym.sandbox.providers import ( + SandboxExecResult, + SandboxHandle, + SandboxProvider, + SandboxSpec, + create_provider, +) + + +T = TypeVar("T") + + +def rewrite_image(image: str | None, rewrites: list[dict[str, str]]) -> str | None: + """Apply ordered image-prefix rewrites used by sandbox configs.""" + if image is None: + return None + for rewrite in rewrites: + from_prefix = rewrite["from"] + to_prefix = rewrite["to"] + if image.startswith(from_prefix): + return to_prefix + image[len(from_prefix) :] + return image + + +class AsyncSandbox: + """Async public facade for provider-backed sandbox operations.""" + + def __init__(self, provider: Mapping[str, Any] | SandboxProvider) -> None: + self._provider = create_provider(provider) if isinstance(provider, Mapping) else provider + + @property + def provider_name(self) -> str: + return self._provider.name + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + return await self._provider.create(spec) + + async def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + return await self._provider.create_batch(spec, count, allow_partial=allow_partial) + + async def connect(self, sandbox_id: str) -> SandboxHandle: + return await self._provider.connect(sandbox_id) + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + return await self._provider.exec( + handle, + command, + cwd=cwd, + env=env, + timeout_s=timeout_s, + user=user, + ) + + async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + await self._provider.write_file(handle, target_path, data) + + async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + return await self._provider.read_file(handle, source_path) + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + await self._provider.upload_file(handle, source_path, target_path) + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + await self._provider.download_file(handle, source_path, target_path) + + async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: + await self._provider.close(handle, delete=delete) + + async def delete(self, handle: SandboxHandle) -> None: + await self.close(handle, delete=True) + + async def aclose(self) -> None: + close_provider = getattr(self._provider, "aclose", None) + if close_provider is not None: + await close_provider() + + async def shutdown(self) -> None: + await self.aclose() + + async def __aenter__(self) -> "AsyncSandbox": + return self + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + await self.aclose() + + def handle_reference(self, handle: SandboxHandle) -> Any: + make_reference = getattr(self._provider, "handle_reference", None) + if make_reference is None: + return handle + return make_reference(handle) + + async def materialize_handle(self, value: Any) -> SandboxHandle: + materialize = getattr(self._provider, "materialize_handle", None) + if materialize is None: + if isinstance(value, SandboxHandle): + return value + raise ValueError(f"Provider {self.provider_name!r} cannot materialize handle references") + result = materialize(value) + if hasattr(result, "__await__"): + result = await result + if not isinstance(result, SandboxHandle): + raise TypeError(f"materialize_handle must return SandboxHandle, got {type(result).__name__}") + return result + + +class _AsyncLoopRunner: + """Run async sandbox operations for sync integrations on one private loop.""" + + def __init__(self) -> None: + self._loop = asyncio.new_event_loop() + self._ready = threading.Event() + self._closed = False + self._thread = threading.Thread(target=self._run_loop, name="nemo-gym-sandbox-sync-loop", daemon=True) + self._thread.start() + self._ready.wait() + + def _run_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._ready.set() + self._loop.run_forever() + + def _ensure_can_block(self, operation: str) -> None: + if self._closed or self._loop.is_closed(): + raise RuntimeError("Sandbox sync loop is closed") + try: + asyncio.get_running_loop() + except RuntimeError: + return + raise RuntimeError(f"Sandbox.{operation}() is blocking; use AsyncSandbox in async code instead.") + + def call(self, operation: str, func: Callable[[], T]) -> T: + self._ensure_can_block(operation) + future: Future[T] = Future() + + def invoke() -> None: + try: + future.set_result(func()) + except BaseException as e: + future.set_exception(e) + + self._loop.call_soon_threadsafe(invoke) + return future.result() + + def run(self, operation: str, awaitable_factory: Callable[[], Awaitable[T]]) -> T: + self._ensure_can_block(operation) + future = asyncio.run_coroutine_threadsafe(awaitable_factory(), self._loop) + return future.result() + + def close(self) -> None: + if self._closed: + return + self._closed = True + if not self._loop.is_closed(): + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=5) + self._loop.close() + + +class Sandbox: + """Sync public facade for provider-backed sandbox operations.""" + + def __init__(self, provider: Mapping[str, Any] | SandboxProvider) -> None: + self._runner = _AsyncLoopRunner() + try: + self._async_sandbox = self._runner.call( + "__init__", + lambda: AsyncSandbox(provider), + ) + except BaseException: + self._runner.close() + raise + self._closed = False + + @property + def provider_name(self) -> str: + return self._runner.call("provider_name", lambda: self._async_sandbox.provider_name) + + def create(self, spec: SandboxSpec) -> SandboxHandle: + return self._runner.run("create", lambda: self._async_sandbox.create(spec)) + + def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + return self._runner.run( + "create_batch", + lambda: self._async_sandbox.create_batch(spec, count, allow_partial=allow_partial), + ) + + def connect(self, sandbox_id: str) -> SandboxHandle: + return self._runner.run("connect", lambda: self._async_sandbox.connect(sandbox_id)) + + def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + return self._runner.run( + "exec", + lambda: self._async_sandbox.exec( + handle, + command, + cwd=cwd, + env=env, + timeout_s=timeout_s, + user=user, + ), + ) + + def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + self._runner.run("write_file", lambda: self._async_sandbox.write_file(handle, target_path, data)) + + def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + return self._runner.run("read_file", lambda: self._async_sandbox.read_file(handle, source_path)) + + def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + self._runner.run("upload_file", lambda: self._async_sandbox.upload_file(handle, source_path, target_path)) + + def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + self._runner.run("download_file", lambda: self._async_sandbox.download_file(handle, source_path, target_path)) + + def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: + self._runner.run("close", lambda: self._async_sandbox.close(handle, delete=delete)) + + def delete(self, handle: SandboxHandle) -> None: + self.close(handle, delete=True) + + def shutdown(self) -> None: + if self._closed: + return + self._closed = True + try: + self._runner.run("shutdown", self._async_sandbox.shutdown) + finally: + self._runner.close() + + def handle_reference(self, handle: SandboxHandle) -> Any: + return self._runner.call("handle_reference", lambda: self._async_sandbox.handle_reference(handle)) + + def materialize_handle(self, value: Any) -> SandboxHandle: + return self._runner.run("materialize_handle", lambda: self._async_sandbox.materialize_handle(value)) + + def __enter__(self) -> "Sandbox": + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self.shutdown() + + def __del__(self) -> None: + if hasattr(self, "_closed") and not self._closed: + try: + self.shutdown() + except Exception: + pass diff --git a/nemo_gym/sandbox/providers/__init__.py b/nemo_gym/sandbox/providers/__init__.py new file mode 100644 index 0000000000..359e99c19b --- /dev/null +++ b/nemo_gym/sandbox/providers/__init__.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Sandbox provider registry.""" + +from nemo_gym.sandbox.providers.base import ( + SandboxBatchCreateError, + SandboxCreateError, + SandboxCreateVerificationError, + SandboxExecResult, + SandboxHandle, + SandboxProvider, + SandboxSpec, +) +from nemo_gym.sandbox.providers.registry import ( + create_provider, + get_provider_class, + list_providers, + register_provider, +) + + +__all__ = [ + "SandboxBatchCreateError", + "SandboxCreateError", + "SandboxCreateVerificationError", + "SandboxExecResult", + "SandboxHandle", + "SandboxProvider", + "SandboxSpec", + "create_provider", + "get_provider_class", + "list_providers", + "register_provider", +] diff --git a/nemo_gym/sandbox/providers/base.py b/nemo_gym/sandbox/providers/base.py new file mode 100644 index 0000000000..7430cb63aa --- /dev/null +++ b/nemo_gym/sandbox/providers/base.py @@ -0,0 +1,136 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Provider-facing sandbox protocol. + +Providers are the only layer that talks to runtime and infrastructure APIs. +Gym agents and external harnesses consume the public ``nemo_gym.sandbox`` API +instead of importing provider-specific modules. +""" + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Protocol + + +@dataclass(frozen=True) +class SandboxSpec: + """Provider-neutral sandbox creation request.""" + + image: str | None = None + snapshot_id: str | None = None + timeout_s: int | None = None + ready_timeout_s: int | None = None + env: dict[str, str] = field(default_factory=dict) + metadata: dict[str, str] = field(default_factory=dict) + resources: dict[str, str] = field(default_factory=dict) + entrypoint: list[str] | None = None + extensions: dict[str, str] = field(default_factory=dict) + platform: dict[str, Any] | None = None + volumes: list[dict[str, Any]] | None = None + skip_health_check: bool | None = None + + +@dataclass(frozen=True) +class SandboxHandle: + """Provider-neutral handle to a created sandbox.""" + + sandbox_id: str + provider_name: str + raw: Any + + +@dataclass(frozen=True) +class SandboxExecResult: + """Provider-neutral process execution result.""" + + stdout: str | None + stderr: str | None + return_code: int + + +class SandboxCreateError(RuntimeError): + """Raised when a provider cannot create a sandbox.""" + + +class SandboxBatchCreateError(SandboxCreateError): + """Raised when a provider cannot complete sandbox batch creation.""" + + +class SandboxCreateVerificationError(SandboxCreateError): + """Raised when a newly-created sandbox fails provider readiness checks.""" + + +class SandboxProvider(Protocol): + """Runtime/infra provider contract used by the public sandbox API.""" + + name: str + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + """Create a sandbox and return a provider-neutral handle.""" + ... + + async def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + """Create several equivalent sandboxes. + + Providers that have a native bulk-allocation primitive should use it. + Providers without one may fall back to calling ``create`` repeatedly. + When ``allow_partial`` is true, providers may return a smaller + contiguous prefix of successfully created handles instead of failing the + whole batch. + """ + ... + + async def connect(self, sandbox_id: str) -> SandboxHandle: + """Connect to an existing sandbox.""" + ... + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + """Run a command inside a sandbox.""" + ... + + async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + """Write a file into a sandbox.""" + ... + + async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + """Read a file from a sandbox.""" + ... + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + """Upload one local file into a sandbox.""" + ... + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + """Download one sandbox file to the local filesystem.""" + ... + + async def close(self, handle: SandboxHandle, *, delete: bool) -> None: + """Close provider resources and optionally delete the sandbox.""" + ... diff --git a/nemo_gym/sandbox/providers/opensandbox/__init__.py b/nemo_gym/sandbox/providers/opensandbox/__init__.py new file mode 100644 index 0000000000..2615667d73 --- /dev/null +++ b/nemo_gym/sandbox/providers/opensandbox/__init__.py @@ -0,0 +1,42 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""OpenSandbox provider package.""" + +from nemo_gym.sandbox.providers.opensandbox.provider import ( + OpenSandboxBatchCreateError, + OpenSandboxConnectionConfig, + OpenSandboxCreateConfig, + OpenSandboxCreateError, + OpenSandboxCreateTimeoutError, + OpenSandboxCreateVerificationError, + OpenSandboxOperationConfig, + OpenSandboxPoolConfig, + OpenSandboxProbeConfig, + OpenSandboxProvider, +) + + +__all__ = [ + "OpenSandboxBatchCreateError", + "OpenSandboxConnectionConfig", + "OpenSandboxCreateConfig", + "OpenSandboxCreateError", + "OpenSandboxCreateTimeoutError", + "OpenSandboxCreateVerificationError", + "OpenSandboxOperationConfig", + "OpenSandboxPoolConfig", + "OpenSandboxProbeConfig", + "OpenSandboxProvider", +] diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py new file mode 100644 index 0000000000..359deb12a7 --- /dev/null +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -0,0 +1,1299 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""OpenSandbox provider implementation.""" + +import asyncio +import logging +import re +import shlex +from collections.abc import Mapping +from dataclasses import dataclass, replace +from datetime import timedelta +from pathlib import Path +from typing import Any, Awaitable, Callable +from uuid import uuid4 + +from nemo_gym.sandbox.providers.base import ( + SandboxBatchCreateError, + SandboxCreateError, + SandboxCreateVerificationError, + SandboxExecResult, + SandboxHandle, + SandboxSpec, +) + + +LOGGER = logging.getLogger(__name__) + + +class OpenSandboxBatchCreateError(SandboxBatchCreateError): + """Raised when a batch sandbox preallocation cannot be completed.""" + + +class OpenSandboxCreateError(SandboxCreateError): + """Raised when OpenSandbox cannot create a sandbox.""" + + +class OpenSandboxCreateTimeoutError(OpenSandboxCreateError): + """Raised when OpenSandbox sandbox creation exceeds the client timeout.""" + + +class OpenSandboxCreateVerificationError(SandboxCreateVerificationError): + """Raised when a newly-created sandbox cannot execute a probe command.""" + + +RETRYABLE_HTTP_STATUS_CODES = {408, 409, 425, 429, 500, 502, 503, 504} +RETRYABLE_ERROR_MARKERS = ( + "all connection attempts failed", + "connection refused", + "connection reset", + "gateway timeout", + "http 408", + "http 409", + "http 425", + "http 429", + "http 500", + "http 502", + "http 503", + "http 504", + "incomplete chunked read", + "peer closed connection", + "pod ip is not yet available", + "pod may still be starting", + "errimagepull", + "get endpoint for sandbox", + "imagepullbackoff", + "pod failed", + "podfailed", + "remote protocol error", + "service unavailable", + "server disconnected", + "status code: 408", + "status code: 409", + "status code: 425", + "status code: 429", + "status code: 500", + "status code: 502", + "status code: 503", + "status code: 504", + "temporarily unavailable", + "timed out", + "timeout", +) +METADATA_VALUE_RE = re.compile(r"[^A-Za-z0-9_.-]+") +DEFAULT_IMAGE_PULL_POLICY = "IfNotPresent" +IMAGE_PULL_POLICY_EXTENSION_KEY = "imagePullPolicy" +IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY = "opensandbox.extensions.image-pull-policy" +VALID_IMAGE_PULL_POLICIES = {"Always", "IfNotPresent", "Never"} +STATUS_CODE_RE = re.compile(r"(?:status code|http)\D+(\d{3})", re.IGNORECASE) + + +def validate_image_pull_policy(image_pull_policy: str) -> str: + """Validate a Kubernetes-compatible container image pull policy.""" + if image_pull_policy not in VALID_IMAGE_PULL_POLICIES: + allowed = ", ".join(sorted(VALID_IMAGE_PULL_POLICIES)) + raise ValueError(f"image_pull_policy must be one of: {allowed}") + return image_pull_policy + + +def _require_opensandbox_sdk() -> tuple[Any, Any, Any, Any, Any]: + try: + from opensandbox import Sandbox + from opensandbox.config import ConnectionConfig + from opensandbox.models.execd import RunCommandOpts + from opensandbox.models.sandboxes import PlatformSpec, Volume + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "OpenSandbox SDK is required for the opensandbox sandbox provider. " + "Install it in the NeMo-RL runtime image before using " + "env.sandbox.provider.name=opensandbox." + ) from e + + return Sandbox, ConnectionConfig, RunCommandOpts, PlatformSpec, Volume + + +def _require_opensandbox_sdk_pool() -> tuple[Any, Any, Any, Any]: + try: + from opensandbox import ( + AcquirePolicy, + InMemoryAsyncPoolStateStore, + PoolCreationSpec, + SandboxPoolAsync, + ) + except ImportError as e: + raise ModuleNotFoundError( + "OpenSandbox SDK >=0.1.9 is required for native SDK pool batch creation. " + "Install opensandbox>=0.1.9 in the NeMo-RL runtime image." + ) from e + + return AcquirePolicy, InMemoryAsyncPoolStateStore, PoolCreationSpec, SandboxPoolAsync + + +def _require_tenacity() -> tuple[Any, Any, Any, Any]: + try: + from tenacity import AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "tenacity is required for OpenSandbox retry handling. Install nemo-gym[sandbox] before using " + "env.sandbox.provider.name=opensandbox." + ) from e + + return AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential + + +def _httpx_retryable_types() -> tuple[type[BaseException], ...]: + try: + import httpx + except ModuleNotFoundError: + return tuple() + return ( + httpx.RemoteProtocolError, + httpx.ReadError, + httpx.WriteError, + httpx.ConnectError, + httpx.ConnectTimeout, + httpx.ReadTimeout, + httpx.WriteTimeout, + httpx.PoolTimeout, + httpx.TimeoutException, + httpx.NetworkError, + ) + + +def _has_retryable_error_marker(exception: BaseException) -> bool: + message = str(exception).lower() + return any(marker in message for marker in RETRYABLE_ERROR_MARKERS) + + +def _exception_status_code(exception: BaseException) -> int | None: + status_code = getattr(exception, "status_code", None) + if isinstance(status_code, int): + return status_code + + match = STATUS_CODE_RE.search(str(exception)) + if match is None: + return None + return int(match.group(1)) + + +def _sdk_error_attributes( + exception: BaseException, + *, + operation: str, + sandbox_id: str, + attempt_number: int | None = None, + max_attempts: int | None = None, + sleep_s: float | None = None, +) -> dict[str, Any]: + attrs: dict[str, Any] = { + "provider": OpenSandboxProvider.name, + "operation": operation, + "sandbox_id": sandbox_id, + "error_type": type(exception).__name__, + "error_message": str(exception)[:500], + } + status_code = _exception_status_code(exception) + if status_code is not None: + attrs["status_code"] = status_code + if attempt_number is not None: + attrs["attempt_number"] = attempt_number + if max_attempts is not None: + attrs["max_attempts"] = max_attempts + if sleep_s is not None: + attrs["next_sleep_s"] = sleep_s + return attrs + + +def _is_retryable_create_error(exception: BaseException) -> bool: + """Return whether a sandbox create failure is likely transient.""" + if isinstance(exception, SandboxCreateVerificationError): + return True + if isinstance(exception, SandboxCreateError): + return True + if isinstance(exception, (ConnectionError, OSError, TimeoutError)): + return True + httpx_types = _httpx_retryable_types() + if httpx_types and isinstance(exception, httpx_types): + return True + + try: + from opensandbox.exceptions import ( + InvalidArgumentException, + SandboxApiException, + SandboxException, + SandboxInternalException, + SandboxReadyTimeoutException, + SandboxUnhealthyException, + ) + except ModuleNotFoundError: + return _has_retryable_error_marker(exception) + + if isinstance(exception, InvalidArgumentException): + return False + if isinstance( + exception, + ( + SandboxInternalException, + SandboxReadyTimeoutException, + SandboxUnhealthyException, + ), + ): + return True + if isinstance(exception, SandboxApiException): + status_code = getattr(exception, "status_code", None) + if status_code in RETRYABLE_HTTP_STATUS_CODES: + return True + if status_code is not None and status_code < 500: + return False + if not isinstance(exception, SandboxException): + return _has_retryable_error_marker(exception) + + return _has_retryable_error_marker(exception) + + +def _is_retryable_sdk_operation_error(exception: BaseException) -> bool: + """Return whether an SDK operation can be retried by Gym. + + The OpenSandbox Python SDK does not retry generated lifecycle, execd, or + filesystem HTTP calls. It converts network failures into SDK exceptions and + exposes API status codes, so classify both the wrapper and its original + cause here. + """ + if isinstance(exception, TimeoutError): + return False + cause = exception.__cause__ + if isinstance(cause, BaseException) and _is_retryable_sdk_operation_error(cause): + return True + if isinstance(exception, (ConnectionError, OSError)): + return True + httpx_types = _httpx_retryable_types() + if httpx_types and isinstance(exception, httpx_types): + return True + return _is_retryable_create_error(exception) + + +def _is_missing_sandbox_delete_error(exception: BaseException) -> bool: + message = str(exception).lower() + return "sandbox" in message and "not found" in message + + +def _log_create_retry(retry_state: Any) -> None: + exception = retry_state.outcome.exception() if retry_state.outcome else None + sleep_s = retry_state.next_action.sleep if retry_state.next_action else None + LOGGER.warning( + "Retrying OpenSandbox sandbox create after attempt %s; next_sleep_s=%s; error=%r", + retry_state.attempt_number, + sleep_s, + exception, + ) + + +def _log_operation_retry(retry_state: Any) -> None: + exception = retry_state.outcome.exception() if retry_state.outcome else None + sleep_s = retry_state.next_action.sleep if retry_state.next_action else None + LOGGER.warning( + "Retrying OpenSandbox SDK operation after attempt %s; next_sleep_s=%s; error=%r", + retry_state.attempt_number, + sleep_s, + exception, + ) + + +def _string_map(values: dict[str, Any]) -> dict[str, str]: + return {str(key): str(value) for key, value in values.items()} + + +def _metadata_value(value: Any) -> str: + normalized = METADATA_VALUE_RE.sub("_", str(value)).strip("._-") + normalized = normalized[:63].strip("._-") + return normalized or "metadata" + + +def _metadata_map(values: dict[str, Any]) -> dict[str, str]: + return {str(key): _metadata_value(value) for key, value in values.items()} + + +def _normalize_spec(spec: SandboxSpec) -> SandboxSpec: + return replace( + spec, + env=_string_map(spec.env), + metadata=_metadata_map(spec.metadata), + resources=_string_map(spec.resources), + extensions=_string_map(spec.extensions), + ) + + +def _to_platform_spec(platform: dict[str, Any]) -> Any: + _, _, _, PlatformSpec, _ = _require_opensandbox_sdk() + return PlatformSpec(**platform) + + +def _to_volumes(volumes: list[dict[str, Any]]) -> list[Any]: + _, _, _, _, Volume = _require_opensandbox_sdk() + return [Volume(**volume) for volume in volumes] + + +def _seconds_to_timedelta(seconds: int | float | None) -> timedelta | None: + if seconds is None: + return None + return timedelta(seconds=float(seconds)) + + +@dataclass(frozen=True) +class OpenSandboxConnectionConfig: + """OpenSandbox server connection settings.""" + + domain: str | None = None + api_key: str | None = None + protocol: str | None = None + use_server_proxy: bool | None = None + exec_use_server_proxy: bool | None = None + request_timeout_s: int | None = None + connect_timeout_s: int | float | None = None + + def __post_init__(self) -> None: + if self.connect_timeout_s is not None and self.connect_timeout_s <= 0: + raise ValueError("connection.connect_timeout_s must be > 0") + + +@dataclass(frozen=True) +class OpenSandboxCreateConfig: + """OpenSandbox create/reconnect retry settings.""" + + request_timeout_s: int | None = None + timeout_s: float | None = None + retries: int = 2 + retry_delay_s: float = 5.0 + retry_max_delay_s: float = 60.0 + image_pull_policy: str | None = DEFAULT_IMAGE_PULL_POLICY + skip_health_check: bool = False + connect_attempt_timeout_s: float = 30.0 + connect_poll_s: float = 2.0 + + def __post_init__(self) -> None: + if self.image_pull_policy is not None: + validate_image_pull_policy(self.image_pull_policy) + if self.timeout_s is not None and self.timeout_s <= 0: + raise ValueError("create.timeout_s must be > 0") + if self.retries < 0: + raise ValueError("create.retries must be >= 0") + if self.retry_delay_s < 0: + raise ValueError("create.retry_delay_s must be >= 0") + if self.retry_max_delay_s < 0: + raise ValueError("create.retry_max_delay_s must be >= 0") + if self.connect_attempt_timeout_s <= 0: + raise ValueError("create.connect_attempt_timeout_s must be > 0") + if self.connect_poll_s <= 0: + raise ValueError("create.connect_poll_s must be > 0") + + +@dataclass(frozen=True) +class OpenSandboxProbeConfig: + """Post-create probe settings.""" + + command: str | None = "printf nemo-rl-sandbox-ready" + expected_stdout: str | None = "nemo-rl-sandbox-ready" + timeout_s: int = 30 + deadline_s: float | None = None + sample_count: int | None = None + stable_count: int = 1 + stable_delay_s: float = 0.0 + + def __post_init__(self) -> None: + if self.command is not None and self.timeout_s <= 0: + raise ValueError("probe.timeout_s must be > 0") + if self.deadline_s is not None and self.deadline_s <= 0: + raise ValueError("probe.deadline_s must be > 0") + if self.sample_count is not None and self.sample_count < 1: + raise ValueError("probe.sample_count must be >= 1") + if self.stable_count < 1: + raise ValueError("probe.stable_count must be >= 1") + if self.stable_delay_s < 0: + raise ValueError("probe.stable_delay_s must be >= 0") + + +@dataclass(frozen=True) +class OpenSandboxOperationConfig: + """Retry and timeout settings for SDK operations after create.""" + + retries: int = 3 + retry_delay_s: float = 1.0 + retry_max_delay_s: float = 15.0 + command_retries: int | None = None + close_timeout_s: float | None = 30.0 + + def __post_init__(self) -> None: + if self.retries < 0: + raise ValueError("operations.retries must be >= 0") + if self.retry_delay_s < 0: + raise ValueError("operations.retry_delay_s must be >= 0") + if self.retry_max_delay_s < 0: + raise ValueError("operations.retry_max_delay_s must be >= 0") + if self.command_retries is not None and self.command_retries < 0: + raise ValueError("operations.command_retries must be >= 0") + if self.close_timeout_s is not None and self.close_timeout_s <= 0: + raise ValueError("operations.close_timeout_s must be > 0") + + +@dataclass(frozen=True) +class OpenSandboxPoolConfig: + """OpenSandbox SDK pool and batch fanout settings.""" + + concurrency: int = 4 + progress_timeout_s: float | None = None + reconcile_interval_s: float = 0.1 + acquire_poll_interval_s: float = 0.1 + idle_timeout_s: float | None = None + primary_lock_ttl_s: float | None = None + + def __post_init__(self) -> None: + if self.concurrency < 1: + raise ValueError("pool.concurrency must be >= 1") + if self.progress_timeout_s is not None and self.progress_timeout_s <= 0: + raise ValueError("pool.progress_timeout_s must be > 0") + if self.reconcile_interval_s <= 0: + raise ValueError("pool.reconcile_interval_s must be > 0") + if self.acquire_poll_interval_s <= 0: + raise ValueError("pool.acquire_poll_interval_s must be > 0") + if self.idle_timeout_s is not None and self.idle_timeout_s <= 0: + raise ValueError("pool.idle_timeout_s must be > 0") + if self.primary_lock_ttl_s is not None and self.primary_lock_ttl_s <= 0: + raise ValueError("pool.primary_lock_ttl_s must be > 0") + + +def _coerce_config(value: Any, config_cls: type[Any]) -> Any: + if value is None: + return config_cls() + if isinstance(value, config_cls): + return value + if isinstance(value, Mapping): + return config_cls(**value) + raise TypeError(f"{config_cls.__name__} must be a mapping or {config_cls.__name__} instance") + + +class OpenSandboxProvider: + """Provider backed by the OpenSandbox SDK/server API. + + Batch allocations use the official OpenSandbox SDK client-side pool. + """ + + name = "opensandbox" + + def __init__( + self, + *, + connection: OpenSandboxConnectionConfig | Mapping[str, Any] | None = None, + create: OpenSandboxCreateConfig | Mapping[str, Any] | None = None, + probe: OpenSandboxProbeConfig | Mapping[str, Any] | None = None, + operations: OpenSandboxOperationConfig | Mapping[str, Any] | None = None, + pool: OpenSandboxPoolConfig | Mapping[str, Any] | None = None, + ) -> None: + self._connection = _coerce_config(connection, OpenSandboxConnectionConfig) + self._create = _coerce_config(create, OpenSandboxCreateConfig) + self._probe = _coerce_config(probe, OpenSandboxProbeConfig) + self._operations = _coerce_config(operations, OpenSandboxOperationConfig) + self._pool = _coerce_config(pool, OpenSandboxPoolConfig) + + def _with_default_image_pull_policy(self, spec: SandboxSpec) -> SandboxSpec: + """Ensure SDK create requests carry the desired image pull policy.""" + if self._create.image_pull_policy is None: + return spec + + extensions = dict(spec.extensions) + image_pull_policy = extensions.get(IMAGE_PULL_POLICY_EXTENSION_KEY) or extensions.get( + IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY + ) + if image_pull_policy is None: + image_pull_policy = self._create.image_pull_policy + image_pull_policy = validate_image_pull_policy(image_pull_policy) + extensions.setdefault(IMAGE_PULL_POLICY_EXTENSION_KEY, image_pull_policy) + extensions.setdefault(IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY, image_pull_policy) + return replace(spec, extensions=extensions) + + def _connection_config( + self, + request_timeout_s: int | float | None = None, + *, + use_server_proxy: bool | None = None, + ) -> Any: + _, ConnectionConfig, _, _, _ = _require_opensandbox_sdk() + kwargs: dict[str, Any] = {} + if self._connection.domain is not None: + kwargs["domain"] = self._connection.domain + if self._connection.api_key is not None: + kwargs["api_key"] = self._connection.api_key + if self._connection.protocol is not None: + kwargs["protocol"] = self._connection.protocol + if use_server_proxy is None: + use_server_proxy = self._connection.use_server_proxy + if use_server_proxy is not None: + kwargs["use_server_proxy"] = use_server_proxy + if request_timeout_s is None: + request_timeout_s = self._connection.request_timeout_s + if request_timeout_s is not None: + kwargs["request_timeout"] = timedelta(seconds=request_timeout_s) + return ConnectionConfig(**kwargs) + + def _exec_connection_config(self, request_timeout_s: int | float | None = None) -> Any: + """Connection config for SDK handles that issue execd/filesystem calls. + + For clustered evaluations, exec traffic should normally use the + OpenSandbox server proxy. That keeps clients off pod IP routing and lets + the server resolve the sandbox backend for each request. + """ + use_server_proxy = self._connection.use_server_proxy + if self._connection.exec_use_server_proxy is not None: + use_server_proxy = self._connection.exec_use_server_proxy + return self._connection_config( + request_timeout_s=request_timeout_s, + use_server_proxy=use_server_proxy, + ) + + async def aclose(self) -> None: + """Close provider-owned resources. + + The provider intentionally does not inject or own OpenSandbox SDK + network clients. SDK handles are closed per sandbox in ``close``. + """ + return None + + async def _await_sdk_call( + self, + awaitable: Any, + *, + operation: str, + sandbox_id: str, + timeout_s: float | None, + ) -> Any: + if timeout_s is None: + return await awaitable + + try: + return await asyncio.wait_for(awaitable, timeout=timeout_s) + except asyncio.TimeoutError as e: + raise TimeoutError( + f"Timed out during OpenSandbox {operation} after {timeout_s:g}s; sandbox_id={sandbox_id!r}" + ) from e + + async def _await_sdk_operation( + self, + operation_factory: Callable[[], Awaitable[Any]], + *, + operation: str, + sandbox_id: str, + timeout_s: float | None, + retries: int | None = None, + ) -> Any: + AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential = _require_tenacity() + retry_count = self._operations.retries if retries is None else retries + max_attempts = retry_count + 1 + + def _before_sleep(retry_state: Any) -> None: + _log_operation_retry(retry_state) + + retry_policy = AsyncRetrying( + retry=retry_if_exception(_is_retryable_sdk_operation_error), + stop=stop_after_attempt(max_attempts), + wait=wait_random_exponential( + multiplier=self._operations.retry_delay_s, + max=self._operations.retry_max_delay_s, + ), + before_sleep=_before_sleep, + reraise=True, + ) + async for attempt in retry_policy: + with attempt: + return await self._await_sdk_call( + operation_factory(), + operation=operation, + sandbox_id=sandbox_id, + timeout_s=timeout_s, + ) + + raise RuntimeError("OpenSandbox SDK operation retry loop did not run") + + async def _verify_created_handle(self, handle: SandboxHandle) -> None: + if self._probe.command is None: + return + + loop = asyncio.get_running_loop() + deadline_s = self._probe.deadline_s or float(self._probe.timeout_s) + deadline = loop.time() + deadline_s + successful_probes = 0 + attempt_number = 0 + last_exception: BaseException | None = None + + while successful_probes < self._probe.stable_count: + remaining_s = deadline - loop.time() + if remaining_s <= 0: + error = OpenSandboxCreateVerificationError( + "OpenSandbox sandbox failed create probe command before " + "the startup deadline; " + f"sandbox_id={handle.sandbox_id!r}, " + f"command={self._probe.command!r}, " + f"successful_probes={successful_probes}/{self._probe.stable_count}, " + f"attempts={attempt_number}, deadline_s={deadline_s:g}" + ) + raise error from last_exception + + attempt_number += 1 + if self._probe.deadline_s is None: + command_timeout_s = float(self._probe.timeout_s) + else: + command_timeout_s = min(float(self._probe.timeout_s), remaining_s) + try: + result = await asyncio.wait_for( + self._exec( + handle, + self._probe.command, + timeout_s=command_timeout_s, + user="root", + ), + timeout=command_timeout_s, + ) + except asyncio.CancelledError: + raise + except Exception as e: + last_exception = e + successful_probes = 0 + sleep_s = min(self._create.connect_poll_s, max(deadline - loop.time(), 0.0)) + if sleep_s > 0: + await asyncio.sleep(sleep_s) + continue + + stdout = result.stdout or "" + expected = self._probe.expected_stdout + if result.return_code != 0 or (expected is not None and expected not in stdout): + last_exception = OpenSandboxCreateVerificationError( + "OpenSandbox sandbox create probe command returned an " + f"unexpected result; sandbox_id={handle.sandbox_id!r}, " + f"return_code={result.return_code}, expected_stdout={expected!r}, " + f"stdout={stdout[:200]!r}, stderr={(result.stderr or '')[:200]!r}, " + f"probe={successful_probes + 1}/{self._probe.stable_count}" + ) + successful_probes = 0 + sleep_s = min(self._create.connect_poll_s, max(deadline - loop.time(), 0.0)) + if sleep_s > 0: + await asyncio.sleep(sleep_s) + continue + + successful_probes += 1 + if successful_probes < self._probe.stable_count and self._probe.stable_delay_s: + await asyncio.sleep(self._probe.stable_delay_s) + + async def _verify_created_handles( + self, + handles: list[SandboxHandle], + ) -> None: + """Verify a batch of created handles with bounded probe concurrency.""" + if self._probe.command is None or not handles: + return + + handles_to_probe = handles + if self._probe.sample_count is not None and self._probe.sample_count < len(handles): + sample_count = self._probe.sample_count + if sample_count == 1: + sampled_indices = [0] + else: + sampled_indices = [ + round(index * (len(handles) - 1) / (sample_count - 1)) for index in range(sample_count) + ] + handles_to_probe = [handles[index] for index in sampled_indices] + + semaphore = asyncio.Semaphore(self._pool.concurrency) + + async def _verify_one(handle: SandboxHandle) -> None: + async with semaphore: + await self._verify_created_handle(handle) + + results = await asyncio.gather( + *(_verify_one(handle) for handle in handles_to_probe), + return_exceptions=True, + ) + errors = [result for result in results if isinstance(result, Exception)] + if errors: + raise OpenSandboxCreateVerificationError( + "One or more OpenSandbox sandboxes failed create probe " + f"verification; failed={len(errors)}, total={len(handles)}" + ) from errors[0] + + async def _cleanup_failed_create_handle(self, handle: SandboxHandle) -> None: + try: + await self.close(handle, delete=True) + except Exception as e: + LOGGER.warning( + "Failed to clean up OpenSandbox sandbox after create probe failure; sandbox_id=%s; error=%r", + handle.sandbox_id, + e, + ) + + async def _connect_after_create(self, handle: SandboxHandle, spec: SandboxSpec) -> SandboxHandle: + """Reconnect after SDK create so follow-up calls use a fresh SDK handle.""" + timeout_s = spec.ready_timeout_s + if timeout_s is None: + timeout_s = self._create.timeout_s + if timeout_s is None: + timeout_s = self._create.connect_attempt_timeout_s + + Sandbox, _, _, _, _ = _require_opensandbox_sdk() + loop = asyncio.get_running_loop() + deadline = loop.time() + float(timeout_s) + last_exception: BaseException | None = None + + while True: + remaining_s = deadline - loop.time() + if remaining_s <= 0: + error = OpenSandboxCreateTimeoutError( + "Timed out connecting to OpenSandbox sandbox after SDK create; " + f"sandbox_id={handle.sandbox_id!r}, timeout_s={timeout_s:g}" + ) + raise error from last_exception + + attempt_timeout_s = min(self._create.connect_attempt_timeout_s, remaining_s) + try: + sandbox = await asyncio.wait_for( + Sandbox.connect( + handle.sandbox_id, + connection_config=self._exec_connection_config(request_timeout_s=attempt_timeout_s), + connect_timeout=timedelta(seconds=attempt_timeout_s), + skip_health_check=True, + ), + timeout=attempt_timeout_s, + ) + return SandboxHandle(sandbox_id=str(sandbox.id), provider_name=self.name, raw=sandbox) + except asyncio.CancelledError: + raise + except BaseException as e: + last_exception = e + if not _is_retryable_create_error(e): + raise + sleep_s = min(self._create.connect_poll_s, max(deadline - loop.time(), 0.0)) + if sleep_s > 0: + await asyncio.sleep(sleep_s) + + async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: + """Create a sandbox through ``opensandbox.Sandbox.create``.""" + if spec.extensions.get("poolRef") and self._connection.use_server_proxy is False: + raise ValueError( + "OpenSandbox pooled creation requires " + "use_server_proxy=True so SDK calls are routed through the " + "server proxy and do not rely on stale cached pod endpoints." + ) + + Sandbox, _, _, _, _ = _require_opensandbox_sdk() + + kwargs: dict[str, Any] = { + "env": spec.env, + "metadata": spec.metadata, + "resource": spec.resources, + "extensions": spec.extensions, + "connection_config": self._exec_connection_config(request_timeout_s=self._create.request_timeout_s), + } + if spec.image is not None: + kwargs["image"] = spec.image + if spec.snapshot_id is not None: + kwargs["snapshot_id"] = spec.snapshot_id + if spec.timeout_s is not None: + kwargs["timeout"] = timedelta(seconds=spec.timeout_s) + if spec.ready_timeout_s is not None: + kwargs["ready_timeout"] = timedelta(seconds=spec.ready_timeout_s) + if spec.entrypoint is not None: + kwargs["entrypoint"] = spec.entrypoint + if spec.platform is not None: + kwargs["platform"] = _to_platform_spec(spec.platform) + if spec.volumes is not None: + kwargs["volumes"] = _to_volumes(spec.volumes) + if self._create.skip_health_check: + kwargs["skip_health_check"] = True + elif spec.skip_health_check is not None: + kwargs["skip_health_check"] = spec.skip_health_check + + timeout_s = self._create.timeout_s + if timeout_s is None and self._connection.request_timeout_s is not None: + timeout_s = float(self._connection.request_timeout_s) + + sandbox_id: str | None = None + sandbox: Any | None = None + try: + if timeout_s is None: + sandbox = await Sandbox.create(**kwargs) + else: + sandbox = await asyncio.wait_for( + Sandbox.create(**kwargs), + timeout=timeout_s, + ) + sandbox_id = str(sandbox.id) + except TimeoutError as e: + error = OpenSandboxCreateTimeoutError( + "Timed out creating OpenSandbox sandbox after " + f"{timeout_s:g}s; image={spec.image!r}, " + f"poolRef={spec.extensions.get('poolRef')!r}, " + f"ready_timeout_s={spec.ready_timeout_s!r}" + ) + raise error from e + if sandbox is None or sandbox_id is None: + raise RuntimeError("OpenSandbox SDK create returned no sandbox handle") + created_handle = SandboxHandle( + sandbox_id=sandbox_id, + provider_name=self.name, + raw=sandbox, + ) + handle = created_handle + try: + if self._create.skip_health_check: + handle = await self._connect_after_create(created_handle, spec) + await self._verify_created_handle(handle) + except Exception: + await self._cleanup_failed_create_handle(created_handle) + raise + return handle + + async def _create_with_retries( + self, + spec: SandboxSpec, + *, + semaphore: asyncio.Semaphore | None = None, + ) -> SandboxHandle: + AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential = _require_tenacity() + 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=_log_create_retry, + reraise=True, + ) + async for attempt in retry_policy: + with attempt: + if semaphore is None: + return await self._create_once(spec) + async with semaphore: + return await self._create_once(spec) + + raise OpenSandboxBatchCreateError("OpenSandbox create retry loop did not run") + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + """Create one sandbox through the configured OpenSandbox path.""" + spec = self._with_default_image_pull_policy(_normalize_spec(spec)) + return await self._create_with_retries(spec) + + async def _close_many( + self, + handles: list[SandboxHandle], + *, + delete: bool, + ) -> list[Any]: + semaphore = asyncio.Semaphore(self._pool.concurrency) + + async def _close_one(handle: SandboxHandle) -> Any: + async with semaphore: + return await self.close(handle, delete=delete) + + return list( + await asyncio.gather( + *(_close_one(handle) for handle in handles), + return_exceptions=True, + ) + ) + + def _validate_sdk_pool_spec(self, spec: SandboxSpec) -> None: + if spec.image is None: + raise ValueError("OpenSandbox SDK pool requires SandboxSpec.image") + if spec.snapshot_id is not None: + raise ValueError("OpenSandbox SDK pool does not support snapshot_id") + + def _to_pool_creation_spec(self, spec: SandboxSpec) -> Any: + self._validate_sdk_pool_spec(spec) + _, _, PoolCreationSpec, _ = _require_opensandbox_sdk_pool() + return PoolCreationSpec( + image=spec.image, + entrypoint=spec.entrypoint, + resource=spec.resources or None, + env=spec.env or None, + metadata=spec.metadata or None, + extensions=spec.extensions or None, + platform=_to_platform_spec(spec.platform) if spec.platform is not None else None, + volumes=_to_volumes(spec.volumes) if spec.volumes is not None else None, + ) + + async def _wait_sdk_pool_idle( + self, + pool: Any, + *, + spec: SandboxSpec, + requested: int, + timeout_s: float, + allow_partial: bool, + ) -> int: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + last_progress_at = loop.time() + last_idle = 0 + last_snapshot: Any = None + + while True: + last_snapshot = await pool.snapshot() + idle_count = int(getattr(last_snapshot, "idle_count", 0) or 0) + if idle_count >= requested: + return requested + if idle_count > last_idle: + last_idle = idle_count + last_progress_at = loop.time() + + now = loop.time() + progress_timeout_s = self._pool.progress_timeout_s + if progress_timeout_s is not None and now - last_progress_at >= progress_timeout_s: + if allow_partial and idle_count > 0: + return idle_count + error = OpenSandboxCreateTimeoutError( + "Timed out waiting for OpenSandbox SDK pool warmup progress " + f"after {progress_timeout_s:g}s; requested={requested}, " + f"idle={idle_count}, snapshot={last_snapshot!r}" + ) + raise error + if now >= deadline: + if allow_partial and idle_count > 0: + return idle_count + error = OpenSandboxCreateTimeoutError( + "Timed out waiting for OpenSandbox SDK pool warmup after " + f"{timeout_s:g}s; requested={requested}, idle={idle_count}, " + f"snapshot={last_snapshot!r}" + ) + raise error + await asyncio.sleep(self._pool.acquire_poll_interval_s) + + async def _direct_exec_handle_for_acquired_sandbox(self, sandbox: Any, spec: SandboxSpec) -> SandboxHandle: + handle = SandboxHandle(sandbox_id=str(sandbox.id), provider_name=self.name, raw=sandbox) + if self._connection.exec_use_server_proxy is None and not self._create.skip_health_check: + return handle + return await self._connect_after_create(handle, spec) + + async def _create_batch_sdk_pool( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool, + ) -> list[SandboxHandle]: + AcquirePolicy, InMemoryAsyncPoolStateStore, _, SandboxPoolAsync = _require_opensandbox_sdk_pool() + ready_timeout_s = float( + spec.ready_timeout_s or self._create.timeout_s or self._connection.request_timeout_s or 300.0 + ) + idle_timeout_s = float(self._pool.idle_timeout_s or spec.timeout_s or max(ready_timeout_s * 2.0, 3600.0)) + primary_lock_ttl_s = float(self._pool.primary_lock_ttl_s or max(ready_timeout_s + 60.0, 60.0)) + pool_name = f"nemo-gym-{uuid4().hex[:12]}" + + async def _warmup_preparer(sandbox: Any) -> None: + if self._probe.command is None: + return + handle = await self._direct_exec_handle_for_acquired_sandbox(sandbox, spec) + try: + await self._verify_created_handle(handle) + finally: + if handle.raw is not sandbox: + try: + await self._await_sdk_call( + handle.raw.close(), + operation="close warmup direct handle", + sandbox_id=handle.sandbox_id, + timeout_s=self._operations.close_timeout_s, + ) + except Exception as e: + LOGGER.warning( + "Failed to close temporary OpenSandbox direct exec handle for sandbox %r: %r", + handle.sandbox_id, + e, + ) + + pool = SandboxPoolAsync( + pool_name=pool_name, + max_idle=count, + warmup_concurrency=self._pool.concurrency, + state_store=InMemoryAsyncPoolStateStore(), + connection_config=self._connection_config(request_timeout_s=self._create.request_timeout_s), + creation_spec=self._to_pool_creation_spec(spec), + reconcile_interval=timedelta(seconds=self._pool.reconcile_interval_s), + primary_lock_ttl=timedelta(seconds=primary_lock_ttl_s), + acquire_ready_timeout=timedelta(seconds=ready_timeout_s), + warmup_ready_timeout=timedelta(seconds=ready_timeout_s), + warmup_sandbox_preparer=_warmup_preparer, + acquire_skip_health_check=bool(self._create.skip_health_check or spec.skip_health_check), + warmup_skip_health_check=bool(self._create.skip_health_check or spec.skip_health_check), + idle_timeout=timedelta(seconds=idle_timeout_s), + ) + handles: list[SandboxHandle] = [] + try: + await pool.start() + ready_count = await self._wait_sdk_pool_idle( + pool, + spec=spec, + requested=count, + timeout_s=ready_timeout_s, + allow_partial=allow_partial, + ) + await pool.resize(0) + sandbox_timeout = _seconds_to_timedelta(spec.timeout_s) + for index in range(ready_count): + sandbox = await pool.acquire( + sandbox_timeout=sandbox_timeout, + policy=AcquirePolicy.FAIL_FAST, + ) + handle = await self._direct_exec_handle_for_acquired_sandbox(sandbox, spec) + handles.append(handle) + LOGGER.info( + "Acquired OpenSandbox SDK pool sandbox %s/%s: %s", + index + 1, + ready_count, + sandbox.id, + ) + return handles + except Exception: + await self._close_many(handles, delete=True) + raise + finally: + try: + await pool.shutdown(graceful=False) + finally: + await pool.release_all_idle() + + async def _create_batch_sdk( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + """Create several sandboxes through the OpenSandbox SDK pool.""" + if count < 1: + raise ValueError("count must be >= 1") + return await self._create_batch_sdk_pool( + spec, + count, + allow_partial=allow_partial, + ) + + async def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + """Create several equivalent OpenSandbox sandboxes.""" + if count < 1: + raise ValueError("count must be >= 1") + spec = self._with_default_image_pull_policy(_normalize_spec(spec)) + return await self._create_batch_sdk( + spec, + count, + allow_partial=allow_partial, + ) + + def handle_reference(self, handle: SandboxHandle) -> dict[str, Any]: + """Build a loop-neutral reference for a sandbox handle. + + OpenSandbox SDK handles are bound to the event loop where they were + created. Prewarmed handles may cross from a FastAPI prewarm request into + a thread-pool runner, so only pass a serializable reference across that + boundary and re-materialize SDK adapters in the consuming event loop. + """ + return { + "kind": "sandbox_id", + "provider": self.name, + "sandbox_id": handle.sandbox_id, + } + + async def materialize_handle(self, reference: dict[str, Any]) -> SandboxHandle: + """Create a loop-local handle from ``handle_reference`` output.""" + kind = reference.get("kind") + if kind == "sandbox_id": + return await self.connect(str(reference["sandbox_id"])) + raise ValueError(f"Unsupported OpenSandbox handle reference kind: {kind!r}") + + async def connect(self, sandbox_id: str) -> SandboxHandle: + """Connect to an existing OpenSandbox sandbox.""" + Sandbox, _, _, _, _ = _require_opensandbox_sdk() + kwargs: dict[str, Any] = { + "connection_config": self._exec_connection_config(), + } + if self._connection.connect_timeout_s is not None: + kwargs["connect_timeout"] = timedelta(seconds=self._connection.connect_timeout_s) + sandbox = await Sandbox.connect(sandbox_id, **kwargs) + return SandboxHandle(sandbox_id=str(sandbox.id), provider_name=self.name, raw=sandbox) + + def _command_retry_count(self) -> int: + return ( + self._operations.retries if self._operations.command_retries is None else self._operations.command_retries + ) + + async def _exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + retries: int | None = None, + ) -> SandboxExecResult: + """Run a command inside an OpenSandbox sandbox.""" + _, _, RunCommandOpts, _, _ = _require_opensandbox_sdk() + + opts_kwargs: dict[str, Any] = {} + if cwd is not None: + opts_kwargs["working_directory"] = cwd + if env is not None: + opts_kwargs["envs"] = env + if timeout_s is not None: + opts_kwargs["timeout"] = timedelta(seconds=timeout_s) + + effective_command = command + if isinstance(user, int): + opts_kwargs["uid"] = user + elif isinstance(user, str) and user != "root": + effective_command = f"su -s /bin/sh -c {shlex.quote(command)} {shlex.quote(user)}" + + sdk_timeout_s = ( + float(timeout_s) + 60.0 + if timeout_s is not None + else ( + float(self._connection.request_timeout_s) if self._connection.request_timeout_s is not None else None + ) + ) + effective_retries = self._command_retry_count() if retries is None else retries + execution = await self._await_sdk_operation( + lambda: handle.raw.commands.run(effective_command, opts=RunCommandOpts(**opts_kwargs)), + operation="command run", + sandbox_id=handle.sandbox_id, + timeout_s=sdk_timeout_s, + retries=effective_retries, + ) + stdout = "\n".join(msg.text for msg in execution.logs.stdout) or None + stderr_parts = [msg.text for msg in execution.logs.stderr] + if execution.error is not None: + stderr_parts.append(f"{execution.error.name}: {execution.error.value}") + stderr = "\n".join(stderr_parts) or None + if execution.exit_code is not None: + return_code = execution.exit_code + elif execution.error is not None: + return_code = 1 + else: + return_code = 0 + + return SandboxExecResult(stdout=stdout, stderr=stderr, return_code=return_code) + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + """Run a command inside an OpenSandbox sandbox.""" + return await self._exec( + handle, + command, + cwd=cwd, + env=env, + timeout_s=timeout_s, + user=user, + retries=self._command_retry_count(), + ) + + async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + """Write one file into an OpenSandbox sandbox.""" + await self._await_sdk_operation( + lambda: handle.raw.files.write_file(target_path, data), + operation=f"write_file({target_path})", + sandbox_id=handle.sandbox_id, + timeout_s=float(self._connection.request_timeout_s) + if self._connection.request_timeout_s is not None + else None, + ) + + async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + """Read one file from an OpenSandbox sandbox.""" + return await self._await_sdk_operation( + lambda: handle.raw.files.read_bytes(source_path), + operation=f"read_file({source_path})", + sandbox_id=handle.sandbox_id, + timeout_s=float(self._connection.request_timeout_s) + if self._connection.request_timeout_s is not None + else None, + ) + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + """Upload one local file into an OpenSandbox sandbox.""" + await self.write_file(handle, target_path, source_path.read_bytes()) + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + """Download one file from an OpenSandbox sandbox.""" + target_path.parent.mkdir(parents=True, exist_ok=True) + target_path.write_bytes(await self.read_file(handle, source_path)) + + async def close(self, handle: SandboxHandle, *, delete: bool) -> None: + """Close local SDK resources and optionally terminate the sandbox.""" + kill_error: Exception | None = None + if delete: + try: + await self._await_sdk_operation( + lambda: handle.raw.kill(), + operation="kill", + sandbox_id=handle.sandbox_id, + timeout_s=self._operations.close_timeout_s, + ) + except Exception as e: + if not _is_missing_sandbox_delete_error(e): + kill_error = e + else: + LOGGER.info( + "OpenSandbox sandbox %r was already deleted during close", + handle.sandbox_id, + ) + + close_error: Exception | None = None + try: + await self._await_sdk_call( + handle.raw.close(), + operation="close", + sandbox_id=handle.sandbox_id, + timeout_s=self._operations.close_timeout_s, + ) + except Exception as e: + close_error = e + LOGGER.warning( + "Timed out or failed while closing local OpenSandbox SDK handle for sandbox %r: %r", + handle.sandbox_id, + e, + ) + + if kill_error is not None: + if close_error is not None: + raise RuntimeError( + "Failed to delete and close OpenSandbox sandbox " + f"{handle.sandbox_id!r}: delete_error={kill_error!r}, " + f"close_error={close_error!r}" + ) from kill_error + raise kill_error + if close_error is not None: + if delete: + return + raise close_error diff --git a/nemo_gym/sandbox/providers/registry.py b/nemo_gym/sandbox/providers/registry.py new file mode 100644 index 0000000000..8aecd6c471 --- /dev/null +++ b/nemo_gym/sandbox/providers/registry.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Provider registration utilities.""" + +from collections.abc import Mapping +from typing import Any, TypeAlias + +from nemo_gym.sandbox.providers.base import SandboxProvider + + +ProviderClass: TypeAlias = type[SandboxProvider] + +_PROVIDER_REGISTRY: dict[str, ProviderClass] = {} + + +def register_provider(name: str, provider_class: ProviderClass) -> None: + """Register a sandbox provider class.""" + if not name: + raise ValueError("Provider name must be non-empty") + if name in _PROVIDER_REGISTRY: + raise ValueError(f"Sandbox provider {name!r} is already registered") + _PROVIDER_REGISTRY[name] = provider_class + + +def get_provider_class(name: str) -> ProviderClass: + """Return a registered provider class.""" + try: + return _PROVIDER_REGISTRY[name] + except KeyError as e: + available = ", ".join(sorted(_PROVIDER_REGISTRY)) or "" + raise ValueError(f"Unknown sandbox provider {name!r}. Available providers: {available}") from e + + +def create_provider(config: Mapping[str, Any]) -> SandboxProvider: + """Instantiate a provider from a single-key provider config.""" + if len(config) != 1: + raise ValueError("Sandbox provider config must contain exactly one provider name") + provider_name, provider_kwargs = next(iter(config.items())) + if not isinstance(provider_name, str) or not provider_name: + raise ValueError("Sandbox provider name must be a non-empty string") + if provider_kwargs is None: + provider_kwargs = {} + if not isinstance(provider_kwargs, Mapping): + raise TypeError(f"Sandbox provider {provider_name!r} config must be a mapping") + + provider_class = get_provider_class(provider_name) + return provider_class(**dict(provider_kwargs)) + + +def list_providers() -> list[str]: + """List registered provider names.""" + return sorted(_PROVIDER_REGISTRY) + + +def _register_builtins() -> None: + from nemo_gym.sandbox.providers.opensandbox import OpenSandboxProvider + + if "opensandbox" not in _PROVIDER_REGISTRY: + register_provider("opensandbox", OpenSandboxProvider) + + +_register_builtins() diff --git a/pyproject.toml b/pyproject.toml index ac4f4ffacc..9790b58d52 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -220,6 +220,18 @@ docs = [ ] [project.optional-dependencies] +sandbox = [ + # Tenacity: Retry helpers used by sandbox providers. + # Updated: Sat May 09, 2026 with tenacity==9.1.4 + # License: Apache 2.0 https://github.com/jd/tenacity/blob/master/LICENSE + "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 + # License: Apache 2.0 + "opensandbox>=0.1.9", +] + # We include dev dependencies as an extra since technically each server module is a consumer (which means we cannot use dependency groups, which are intended to be within a project). dev = [ # Pre-commit: Used for pre-commit hooks. @@ -395,7 +407,18 @@ ng_reinstall = "nemo_gym.cli:reinstall" [tool.setuptools.packages.find] where = ["."] -include = ["benchmarks", "resources_servers", "responses_api_agents", "responses_api_models", "nemo_gym"] +include = [ + "benchmarks", + "benchmarks.*", + "resources_servers", + "resources_servers.*", + "responses_api_agents", + "responses_api_agents.*", + "responses_api_models", + "responses_api_models.*", + "nemo_gym", + "nemo_gym.*", +] ################################################ # Testing diff --git a/responses_api_agents/mini_swe_agent_2/.gitignore b/responses_api_agents/mini_swe_agent_2/.gitignore new file mode 100644 index 0000000000..68bcbc9609 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/.gitignore @@ -0,0 +1 @@ +results/ \ No newline at end of file diff --git a/responses_api_agents/mini_swe_agent_2/README.md b/responses_api_agents/mini_swe_agent_2/README.md new file mode 100644 index 0000000000..ff72d22281 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/README.md @@ -0,0 +1,347 @@ +# Mini-SWE-Agent 2 Sandbox Agent + +A NeMo Gym Responses API agent that integrates +[mini-swe-agent](https://github.com/SWE-agent/mini-swe-agent) v2 for evaluating +language models on SWE-bench style software engineering tasks through the public +`nemo_gym.sandbox` API. + +This agent intentionally keeps only the sandbox-backed path. It does not carry +over the older Docker/Singularity mini-SWE integration. + +## Contents + +- [Mini-SWE-Agent 2 Sandbox Agent](#mini-swe-agent-2-sandbox-agent) + - [Contents](#contents) + - [Overview](#overview) + - [Dataset Information](#dataset-information) + - [Configuration](#configuration) + - [Agent Configuration](#agent-configuration) + - [Model Parameters](#model-parameters) + - [Usage](#usage) + - [Server](#server) + - [Collect Rollouts](#collect-rollouts) + - [Sandbox Environment Adapter](#sandbox-environment-adapter) + - [Environment Lifecycle](#environment-lifecycle) + - [Contributing](#contributing) + - [Licensing Information](#licensing-information) + - [Dependencies](#dependencies) + +## Overview + +`mini_swe_agent_2` runs mini-swe-agent's synchronous SWE-bench harness while +creating and executing each task environment through Gym's provider-neutral +sandbox facade. The validated path in this directory is: + +- mini-swe-agent `2.1.0` +- SWE-bench task rows, including SWE-bench Verified +- `env: sandbox` +- `responses_api_agents.mini_swe_agent_2.sandbox_environment.MiniSWESandboxEnvironment` +- OpenSandbox through `nemo_gym.sandbox.providers.opensandbox` + +For each `/run` request, `MiniSWEAgent.run()` loads mini-swe-agent's built-in +`swebench.yaml`, injects sandbox settings, runs mini-swe-agent in a Ray remote +task, evaluates the generated patch with the SWE-bench harness, and returns a +Gym verify response with reward `1.0` only when the instance is resolved and the +evaluation report includes test status. + +`MiniSWEAgent.setup_webserver()` also registers `/v1/responses`, but +`MiniSWEAgent.responses()` is intentionally not implemented in this agent. The +supported eval path is `/run`, typically via `ng_collect_rollouts`. + +## Dataset Information + +- Eval data - [princeton-nlp/SWE-bench_Verified](https://huggingface.co/datasets/princeton-nlp/SWE-bench_Verified) + is the primary validation target. It contains 500 human-validated SWE-bench + test instances. +- The rollout input JSONL should preserve the SWE-bench instance fields needed + by `swebench`, such as `instance_id`, `repo`, `base_commit`, + `problem_statement`, `patch`, `test_patch`, `FAIL_TO_PASS`, `PASS_TO_PASS`, + and related version fields. +- Each row must also include `responses_create_params`. Extra top-level + SWE-bench fields are accepted by the agent request model and passed into + mini-swe-agent as the instance dictionary. + +Example row shape: + +```json +{ + "instance_id": "django__django-13410", + "repo": "django/django", + "base_commit": "...", + "problem_statement": "...", + "patch": "...", + "test_patch": "...", + "FAIL_TO_PASS": ["..."], + "PASS_TO_PASS": ["..."], + "responses_create_params": { + "input": [], + "temperature": 0.6, + "top_p": 1.0, + "max_output_tokens": 16384 + } +} +``` + +When `image_name` is present on a row, the agent uses it directly. Otherwise it +derives the SWE-bench image from `instance_id` and `subset`: + +- `subset: verified` uses `docker.io/swebench/sweb.eval.x86_64.:latest` + with `__` replaced by `_1776_`. +- Other subsets use `docker.io/xingyaoww/sweb.eval.x86_64.:latest` with + `__` replaced by `_s_`. + +The default OpenSandbox config uses explicit Docker Hub image refs so cluster +mirroring can happen in the container runtime instead of Gym-side image +rewrites. + +## Configuration + +### Agent Configuration + +Path - `responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml` + +```yaml +mini_swe_agent_2: + responses_api_agents: + mini_swe_agent_2: + entrypoint: app.py + domain: coding + description: Software engineering tasks driven by mini-swe-agent harness on OpenSandbox. + value: Improve agentic software engineering capabilities. + model_server: + type: responses_api_models + name: policy_model + concurrency: 64 + env: sandbox + sandbox_provider: + opensandbox: + connection: + domain: opensandbox-server.opensandbox-system.svc.cluster.local + api_key: ${oc.env:OPENSANDBOX_API_KEY} + protocol: http + use_server_proxy: true + exec_use_server_proxy: true + request_timeout_s: 300 + create: + request_timeout_s: 1200 + timeout_s: 1200 + skip_health_check: true + retries: 10 + retry_delay_s: 5.0 + retry_max_delay_s: 90.0 + probe: + timeout_s: 60 + deadline_s: 180 + stable_count: 2 + stable_delay_s: 1.0 + operations: + retries: 5 + retry_delay_s: 1.0 + retry_max_delay_s: 45.0 + command_retries: 3 + close_timeout_s: 30 + sandbox_spec: + timeout_s: 18000 + ready_timeout_s: 1200 + resources: + cpu: "2" + memory: 8Gi + ephemeral-storage: 20Gi + platform: + os: linux + arch: amd64 + metadata: + benchmark: swebench-verified + harness: mini-swe-agent + sandbox-api: opensandbox-sdk + sandbox_environment_kwargs: + cwd: /testbed + conda_env: testbed + activate_conda: true + user: root + delete: true + run_golden: false + step_timeout: 600 + eval_timeout: 1800 + skip_if_exists: false + step_limit: 250 +``` + +Optional `sandbox_resource_profiles` can be configured as a list of resource +maps. When present, the agent hashes `instance_id` and deterministically merges +one profile into `sandbox_spec.resources`. This is useful for spreading +SWE-bench tasks across a small set of resource sizes without changing the input +data. + +### Model Parameters + +`MiniSWEAgent.run()` maps supported Responses API fields into mini-swe-agent +chat-completions kwargs: + +- `temperature`, `top_p`, `top_logprobs`, and `parallel_tool_calls` pass through. +- `max_output_tokens` becomes `max_tokens`. +- `responses_create_params.metadata.extra_body` must be a JSON object and is + passed as `extra_body`. +- `responses_create_params.metadata.chat_template_kwargs` must be a JSON object + and is nested under `extra_body.chat_template_kwargs`. +- `tool_choice` comes from the agent config when set, otherwise from the request. + The special value `bash` expands to the OpenAI function choice for the `bash` + tool. + +Keep the requested generation budget compatible with the live vLLM deployment. +For example, a deployment served with `--max-model-len 32768` will reject +`max_output_tokens=49152`. In earlier smoke testing, that upstream vLLM rejection +surfaced in mini-swe-agent as repeated: + +```text +No tool calls found in the response. Every response MUST include at least one tool call. +``` + +That symptom was not a sandbox failure and was not a reason to force the `bash` +tool. The successful smoke kept `tool_choice=auto` and lowered +`max_output_tokens` to `16384`. + +## Usage + +### Server + +Set the policy model endpoint in `env.yaml` or with equivalent Hydra overrides: + +```yaml +policy_base_url: http://..svc.cluster.local:8000/v1 +policy_api_key: dummy-key +policy_model_name: +``` + +Start the mini-swe-agent 2 server with the OpenSandbox provider and a policy +model server. The values below show a representative SWE-bench eval setup: + +```bash +CONFIG_PATHS="responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml,responses_api_models/vllm_model/configs/vllm_model.yaml" + +ng_run "+config_paths=[$CONFIG_PATHS]" \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.concurrency=64 \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.step_timeout=600 \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.eval_timeout=1800 \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.step_limit=50 \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.run_golden=false \ + '+mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.sandbox_spec.resources={cpu: 500m, memory: 4Gi, ephemeral-storage: 8Gi}' \ + '+mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.sandbox_spec.metadata={benchmark: swebench-verified, harness: mini_swe_agent_2, endpoint_label: hosted-vllm, run_family: mini-swe-agent-2-pass8}' +``` + +Use a model server config that matches the policy endpoint you are serving. The +example above uses `vllm_model`, which is the common path for hosted vLLM +`/v1/chat/completions` endpoints. + +### Collect Rollouts + +Collect eval rollouts from a SWE-bench-style JSONL file: + +```bash +ng_collect_rollouts \ + +agent_name=mini_swe_agent_2 \ + +input_jsonl_fpath=data/mini_swe_verified_smoke8.jsonl \ + +output_jsonl_fpath=results/mini_swe_agent_2_pass8.jsonl \ + +limit=8 \ + +num_repeats=8 \ + +num_samples_in_parallel=64 \ + '+responses_create_params={max_output_tokens: 32768, temperature: 0.6, top_p: 0.95, metadata: {chat_template_kwargs: "{\"enable_thinking\": true}"}}' +``` + +`ng_collect_rollouts` also writes +`results/mini_swe_agent_2_pass8_aggregate_metrics.json` +with per-task eval status, pass@k, resolved task counts, and eval error rates. +After collecting repeated rollouts, run `ng_reward_profile` on the collected +output when you want the standalone profiler JSONL as well: + +```bash +ng_reward_profile \ + +input_jsonl_fpath=data/mini_swe_verified_smoke8.jsonl \ + +materialized_inputs_jsonl_fpath=results/mini_swe_agent_2_pass8_materialized_inputs.jsonl \ + +rollouts_jsonl_fpath=results/mini_swe_agent_2_pass8.jsonl \ + +pass_threshold=1.0 +``` + +The profiler writes `*_reward_profiling.jsonl` and `*_agent_metrics.json` +next to the rollouts file. + +The agent writes per-instance mini-swe-agent configs and result artifacts under +`results///`. + +Use the agent's `step_timeout` and `eval_timeout` overrides above to bound tool +and verifier execution. If you launch from a custom Kubernetes wrapper, add any +outer per-sample guard there. + +## Sandbox Environment Adapter + +`MiniSWESandboxEnvironment` adapts mini-swe-agent's synchronous environment +contract to `nemo_gym.sandbox.Sandbox`. + +When `env` is `sandbox`, Gym injects this environment config before calling +mini-swe-agent: + +```yaml +environment: + environment_class: responses_api_agents.mini_swe_agent_2.sandbox_environment.MiniSWESandboxEnvironment + image: + provider: + opensandbox: + connection: ... + spec: + resources: ... + platform: ... + metadata: ... +``` + +### Environment Lifecycle + +`MiniSWESandboxEnvironment.__init__()`: + +- Validates that a sandbox provider was configured. +- Builds a `SandboxSpec` from the task image, environment variables, metadata, + resources, platform, volumes, provider-specific extensions, and health-check + settings. +- Adds standard metadata such as `nemo_gym_agent=mini_swe_agent_2` and + `instance_id`. +- Creates a `Sandbox` facade and calls `Sandbox.create(...)`. + +`execute()`: + +- Receives mini-swe-agent's command action. +- Applies the configured working directory and timeout. +- Optionally wraps the command in `conda activate ` for SWE-bench images + that expect a prebuilt conda environment. +- Calls `Sandbox.exec(...)` as the configured user, root by default. +- Returns mini-swe-agent's expected sync response shape: + +```python +{ + "output": "...", + "returncode": 0, + "exception_info": "", +} +``` + +`_check_finished()` preserves mini-swe-agent's submit sentinel behavior. If the +command output begins with `COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT` and the +command succeeded, it raises `minisweagent.exceptions.Submitted` with the final +submission payload. + +`cleanup()` calls `Sandbox.close(..., delete=config.delete)` and then +`Sandbox.shutdown()` to release provider-owned async resources and stop the sync +facade's private loop. + +## Contributing + +Please refer to the main NeMo Gym documentation for contributing guidelines. + +## Licensing Information + +- **Code**: Apache 2.0 +- **SWE-bench Verified**: MIT + +### Dependencies + +- **nemo_gym**: Apache 2.0 +- **mini-swe-agent**: MIT +- **SWE-bench / swebench**: MIT diff --git a/responses_api_agents/mini_swe_agent_2/__init__.py b/responses_api_agents/mini_swe_agent_2/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/responses_api_agents/mini_swe_agent_2/app.py b/responses_api_agents/mini_swe_agent_2/app.py new file mode 100644 index 0000000000..2de794e1ec --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/app.py @@ -0,0 +1,790 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 sys +import time +import traceback +from asyncio import Semaphore +from pathlib import Path +from typing import Any, Callable, Literal, Optional, cast +from uuid import uuid4 + +import ray +import yaml +from fastapi import Body, FastAPI +from minisweagent.config import builtin_config_dir, get_config_path +from pydantic import ConfigDict + +from nemo_gym.base_resources_server import ( + BaseRunRequest, + BaseVerifyRequest, + BaseVerifyResponse, +) +from nemo_gym.base_responses_api_agent import ( + BaseResponsesAPIAgentConfig, + SimpleResponsesAPIAgent, +) +from nemo_gym.config_types import ModelServerRef +from nemo_gym.global_config import TASK_INDEX_KEY_NAME +from nemo_gym.openai_utils import ( + NeMoGymResponse, + NeMoGymResponseCreateParamsNonStreaming, +) +from nemo_gym.reward_profile import compute_pass_majority_metrics, highest_k_metrics +from nemo_gym.server_utils import ( + ServerClient, + get_first_server_config_dict, +) + + +class MiniSWEAgentConfig(BaseResponsesAPIAgentConfig): + model_server: ModelServerRef + env: Literal["sandbox"] + concurrency: int + sandbox_provider: Optional[dict[str, Any]] = None + sandbox_spec: Optional[dict[str, Any]] = None + sandbox_environment_kwargs: Optional[dict[str, Any]] = None + run_golden: bool = False + step_timeout: int = 600 + eval_timeout: int = 1800 + skip_if_exists: bool = False + step_limit: int = 250 + tool_choice: Optional[str | dict[str, Any]] = None + sandbox_resource_profiles: Optional[list[dict[str, str]]] = None + + +class MiniSWEAgentRunRequest(BaseRunRequest): + model_config = ConfigDict(extra="allow") + + +class MiniSWEAgentVerifyRequest(BaseVerifyRequest): + model_config = ConfigDict(extra="allow") + + +class MiniSWEAgentVerifyResponse(BaseVerifyResponse): + model_config = ConfigDict(extra="allow") + + +@ray.remote( + scheduling_strategy="SPREAD", + runtime_env={ + "py_executable": sys.executable, + }, +) +def runner_ray_remote(runner: Callable, params: dict[str, Any]) -> Any: + return runner(**params) + + +def _json_dict_from_metadata(value: Any, *, field_name: str) -> dict[str, Any]: + if value is None: + return {} + if isinstance(value, dict): + return value + if isinstance(value, str): + parsed = json.loads(value) + if isinstance(parsed, dict): + return parsed + raise ValueError(f"responses_create_params.metadata.{field_name} must be a JSON object") + + +def _responses_create_params_to_model_kwargs( + params: dict[str, Any], + *, + default_tool_choice: Any = None, +) -> dict[str, Any]: + """Convert Gym Responses API rollout params into mini-swe-agent chat-completions kwargs.""" + model_kwargs: dict[str, Any] = {} + for key in ("temperature", "top_p", "top_logprobs", "parallel_tool_calls"): + value = params.get(key) + if value is not None: + model_kwargs[key] = value + + max_output_tokens = params.get("max_output_tokens") + if max_output_tokens is not None: + model_kwargs["max_tokens"] = max_output_tokens + + metadata = params.get("metadata") or {} + extra_body = _json_dict_from_metadata(metadata.get("extra_body"), field_name="extra_body") + chat_template_kwargs = _json_dict_from_metadata( + metadata.get("chat_template_kwargs"), + field_name="chat_template_kwargs", + ) + if chat_template_kwargs: + extra_body["chat_template_kwargs"] = chat_template_kwargs + if extra_body: + model_kwargs["extra_body"] = extra_body + + tool_choice = default_tool_choice if default_tool_choice is not None else params.get("tool_choice") + if tool_choice == "bash": + model_kwargs["tool_choice"] = _bash_tool_choice() + elif tool_choice is not None: + model_kwargs["tool_choice"] = tool_choice + + return model_kwargs + + +def _bash_tool_choice() -> dict[str, Any]: + return {"type": "function", "function": {"name": "bash"}} + + +def _sandbox_spec_for_instance( + spec: dict[str, Any] | None, + *, + resource_profiles: list[dict[str, str]] | None, + instance_id: str, +) -> dict[str, Any]: + instance_spec = dict(spec or {}) + if not resource_profiles: + return instance_spec + + resources = dict(instance_spec.get("resources") or {}) + digest = hashlib.sha256(instance_id.encode("utf-8")).digest() + profile = resource_profiles[int.from_bytes(digest[:4], "big") % len(resource_profiles)] + resources.update(profile) + instance_spec["resources"] = resources + return instance_spec + + +def _swebench_config_path() -> Path: + for candidate in ( + builtin_config_dir / "extra" / "swebench.yaml", + builtin_config_dir / "benchmarks" / "swebench.yaml", + ): + if candidate.exists(): + return candidate + return builtin_config_dir / "extra" / "swebench.yaml" + + +def _swebench_image_name(instance: dict[str, Any], subset: str) -> str: + image_name = instance.get("image_name") + if image_name: + return str(image_name) + + instance_id = instance["instance_id"] + if subset == "verified": + docker_compatible_id = instance_id.replace("__", "_1776_") + return f"docker.io/swebench/sweb.eval.x86_64.{docker_compatible_id}:latest".lower() + + docker_compatible_id = instance_id.replace("__", "_s_") + return f"docker.io/xingyaoww/sweb.eval.x86_64.{docker_compatible_id}:latest".lower() + + +def _message_content_to_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [] + for item in content: + if isinstance(item, dict): + parts.append(str(item.get("text") or item.get("content") or "")) + else: + parts.append(str(item)) + return "\n".join(part for part in parts if part) + return "" if content is None else str(content) + + +def _strip_extra(item: Any) -> dict[str, Any]: + if hasattr(item, "model_dump"): + item = item.model_dump() + if not isinstance(item, dict): + return {"type": "message", "role": "user", "content": str(item)} + return {key: value for key, value in item.items() if key != "extra"} + + +def _split_trajectory_for_responses( + messages: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]]]: + input_messages: list[dict[str, Any]] = [] + output_items: list[dict[str, Any]] = [] + raw_responses: list[dict[str, Any]] = [] + in_initial_prompt = True + + for message in messages: + role = message.get("role") + if in_initial_prompt and role in {"system", "user"}: + input_messages.append( + {"type": "message", "role": role, "content": _message_content_to_text(message.get("content"))} + ) + continue + + in_initial_prompt = False + if message.get("object") == "response": + response = _strip_extra(message) + raw_responses.append(response) + output_items.extend(_strip_extra(item) for item in response.get("output", [])) + elif role == "assistant": + content = _message_content_to_text(message.get("content")) + if content: + output_items.append( + { + "id": message.get("id") or f"msg_{uuid4()}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": content, "annotations": []}], + } + ) + for tool_call in message.get("tool_calls") or []: + function = tool_call.get("function") or {} + output_items.append( + { + "id": tool_call.get("id") or f"fc_{uuid4()}", + "type": "function_call", + "name": function.get("name") or tool_call.get("name") or "", + "call_id": tool_call.get("id") or tool_call.get("call_id") or "", + "arguments": function.get("arguments") or tool_call.get("arguments") or "{}", + } + ) + elif role == "tool": + output_items.append( + { + "type": "function_call_output", + "call_id": message.get("tool_call_id") or message.get("call_id") or "", + "output": _message_content_to_text(message.get("content")), + } + ) + elif message.get("type") == "function_call_output": + output_items.append(_strip_extra(message)) + + return input_messages, output_items, raw_responses + + +def _default_response_object() -> dict[str, Any]: + return { + "id": f"resp_{str(uuid4())}", + "created_at": int(time.time()), + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "object": "response", + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "background": False, + "max_output_tokens": None, + "max_tool_calls": None, + "previous_response_id": None, + "prompt": None, + "reasoning": { + "effort": None, + "generate_summary": None, + "summary": None, + }, + "service_tier": "default", + "status": "completed", + "text": {"format": {"type": "text"}, "verbosity": "medium"}, + "top_logprobs": 0, + "truncation": "disabled", + "usage": { + "input_tokens": 0, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 0, + }, + "user": None, + "prompt_cache_key": None, + "safety_identifier": None, + "store": True, + } + + +def _is_resolved(instance_id: str, eval_report: dict[str, Any]) -> bool: + try: + if not eval_report: + return False + report = eval_report["eval_report"][instance_id] + resolved = bool(report["resolved"]) + if not report.get("tests_status"): + return False + + tests_status = report["tests_status"] + f2f = tests_status.get("FAIL_TO_PASS", {}) + p2p = tests_status.get("PASS_TO_PASS", {}) + total_reported = ( + len(f2f.get("success", [])) + + len(f2f.get("failure", [])) + + len(p2p.get("success", [])) + + len(p2p.get("failure", [])) + ) + return resolved and total_reported > 0 + except Exception as exc: + print(f"Error in _is_resolved: {exc}", flush=True) + return False + + +def _metadata_dict(verify_response: dict[str, Any]) -> dict[str, Any]: + metadata = verify_response.get("metadata") or {} + return metadata if isinstance(metadata, dict) else {} + + +def _eval_report_map(verify_response: dict[str, Any]) -> dict[str, Any]: + report = _metadata_dict(verify_response).get("eval_report") or {} + return report if isinstance(report, dict) else {} + + +def _eval_instance_report(verify_response: dict[str, Any]) -> dict[str, Any]: + report_map = _eval_report_map(verify_response) + instance_id = verify_response.get("instance_id") or _metadata_dict(verify_response).get("instance_id") + if instance_id is not None: + report = report_map.get(str(instance_id)) + if isinstance(report, dict): + return report + + for report in report_map.values(): + if isinstance(report, dict) and "resolved" in report: + return report + return {} + + +def _test_status_counts(verify_response: dict[str, Any]) -> dict[str, int]: + report = _eval_instance_report(verify_response) + tests_status = report.get("tests_status") if isinstance(report, dict) else None + if not isinstance(tests_status, dict): + return {} + + counts: dict[str, int] = {} + for suite_name, suite_report in tests_status.items(): + if not isinstance(suite_report, dict): + continue + prefix = str(suite_name).lower() + counts[f"{prefix}_success"] = len(suite_report.get("success") or []) + counts[f"{prefix}_failure"] = len(suite_report.get("failure") or []) + return counts + + +def _run_eval_v2( + *, + instance: dict[str, Any], + env: Any, + model_patch: str, + instance_dir: Path, + run_id: str, + is_golden: bool, +) -> dict[str, Any]: + from swebench.harness.constants import SWEbenchInstance + from swebench.harness.docker_build import setup_logger + from swebench.harness.grading import get_eval_report + from swebench.harness.test_spec.test_spec import make_test_spec + + swebench_instance = cast(SWEbenchInstance, instance) + test_spec = make_test_spec(swebench_instance) + pred = {"instance_id": test_spec.instance_id, "model_patch": model_patch} + + instance_dir.mkdir(parents=True, exist_ok=True) + log_file = instance_dir / f"run_instance_{run_id}.log" + report_path = instance_dir / f"report_{run_id}.json" + patch_file = instance_dir / f"patch_{run_id}.diff" + patch_file.write_text(model_patch) + + logger = setup_logger(test_spec.instance_id, log_file) + logger.info(f"DEBUG test_spec {test_spec}") + logger.info(f"DEBUG eval_script {test_spec.eval_script}") + + if is_golden: + env.execute(f"cat > patch.diff <<'EOF'\n{model_patch}\n\nEOF") + env.execute("git status --porcelain") + env.execute("git apply --check patch.diff") + env.execute("git apply patch.diff") + + eval_script = test_spec.eval_script.replace("#!/bin/bash", "") + result = env.execute(eval_script, is_eval=True) + test_output = result["output"] + returncode = result["returncode"] + print(f"[EVAL]{test_spec.instance_id} returncode: {returncode}", flush=True) + + test_output_path = instance_dir / f"test_output_{run_id}.txt" + test_output_path.write_text(test_output) + print(f"[EVAL]{test_spec.instance_id} Test output written to {test_output_path}", flush=True) + + report = get_eval_report( + test_spec=test_spec, + prediction=pred, + test_log_path=str(test_output_path), + include_tests_status=True, + ) + print(f"[EVAL]{test_spec.instance_id} Result: resolved: {report[test_spec.instance_id]['resolved']}", flush=True) + + report_path.write_text(json.dumps(report, indent=4)) + return { + "instance_id": test_spec.instance_id, + "model_patch": model_patch, + "eval_report": report, + } + + +def _run_mini_swe_v2(**params: Any) -> dict[str, Any]: + from minisweagent.agents.default import DefaultAgent + from minisweagent.environments import get_environment + from minisweagent.models import get_model + + instance = params.get("instance_dict") + if isinstance(instance, str): + instance = json.loads(instance) + if not isinstance(instance, dict): + raise ValueError("mini-swe-agent v2 path requires instance_dict") + + instance = dict(instance) + instance_id = str(params.get("instance_id") or instance["instance_id"]).lower() + instance["instance_id"] = instance_id + + output_dir = Path(params["output"]) + instance_dir = output_dir / instance_id + output_dir.mkdir(parents=True, exist_ok=True) + instance_dir.mkdir(parents=True, exist_ok=True) + + config = yaml.safe_load(get_config_path(params["config"]).read_text()) + model_config = config.setdefault("model", {}) + model_config["model_class"] = "litellm" + model_config["model_name"] = params["model"] + model_config.setdefault("cost_tracking", "ignore_errors") + model_kwargs = model_config.setdefault("model_kwargs", {}) + model_kwargs["api_key"] = params["api_key"] + model_kwargs["base_url"] = params["base_url"] + 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: + model_kwargs["max_tokens"] = max_output_tokens + + environment_config = config.setdefault("environment", {}) + environment_config["image"] = _swebench_image_name(instance, params["subset"]) + environment_config["step_timeout"] = params["step_timeout"] + environment_config["eval_timeout"] = params["eval_timeout"] + environment_config["instance_id"] = instance_id + environment_config["environment_class"] = ( + "responses_api_agents.mini_swe_agent_2.sandbox_environment.MiniSWESandboxEnvironment" + ) + + agent_config = config.get("agent", {}) + agent_config["step_limit"] = params["step_limit"] + agent_config.pop("collapse_limit", None) + + run_id = f"{int(time.time())}_{uuid4()}" + trajectory_path = instance_dir / f"{instance_id}_{run_id}.traj.json" + agent_config["output_path"] = trajectory_path + env = None + agent = None + try: + print(f"[EVAL]{instance_id} Creating environment...", flush=True) + env = get_environment(environment_config) + print(f"[EVAL]{instance_id} Environment created", flush=True) + + model = get_model(config=model_config) + agent = DefaultAgent(model, env, **agent_config) + + if params["run_golden"]: + exit_status = "Gold Patch Applied" + model_patch = instance.get("patch", "") + data = agent.save(None, {"messages": []}) + else: + print(f"[EVAL]{instance_id} Running mini-swe-agent v2...", flush=True) + info = agent.run(instance["problem_statement"]) + exit_status = info.get("exit_status", "") + model_patch = info.get("submission", "") + data = agent.save( + trajectory_path, + {"instance_id": instance_id}, + ) + + print(f"[EVAL]{instance_id} Running eval", flush=True) + eval_report = _run_eval_v2( + instance=instance, + env=env, + model_patch=model_patch, + instance_dir=instance_dir, + run_id=run_id, + is_golden=params["run_golden"], + ) + print(f"[EVAL]{instance_id} Eval completed", flush=True) + + input_messages, response_output, responses = _split_trajectory_for_responses(data.get("messages", [])) + + return { + instance_id: { + "input_messages": input_messages, + "response_output": response_output, + "responses": responses, + "eval_report": eval_report, + "exit_status": exit_status, + } + } + finally: + if env and hasattr(env, "cleanup"): + env.cleanup() + + +def run_mini_swe_with_sandbox(**params: Any) -> Any: + return _run_mini_swe_v2(**params) + + +class MiniSWEAgent(SimpleResponsesAPIAgent): + config: MiniSWEAgentConfig + sem: Semaphore = None + model_config = ConfigDict(arbitrary_types_allowed=True) + + def model_post_init(self, __context: Any) -> None: + self.sem = Semaphore(self.config.concurrency) + + def setup_webserver(self) -> FastAPI: + app = FastAPI() + self.setup_session_middleware(app) + app.post("/v1/responses")(self.responses) + app.post("/run")(self.run) + app.post("/aggregate_metrics")(self.aggregate_metrics) + return app + + def compute_metrics(self, tasks: list[list[dict[str, Any]]]) -> dict[str, Any]: + metrics, _, _, max_k = compute_pass_majority_metrics(tasks) + metrics.pop("per_sample_aggregate", None) + + all_rollouts = [rollout for task in tasks for rollout in task] + rollout_count = len(all_rollouts) + resolved_task_count = sum(1 for task in tasks if any(float(r.get("reward", 0.0) or 0.0) >= 1.0 for r in task)) + eval_error_count = sum(1 for rollout in all_rollouts if _metadata_dict(rollout).get("error")) + eval_report_count = sum(1 for rollout in all_rollouts if _eval_report_map(rollout)) + tests_status_count = sum(1 for rollout in all_rollouts if _eval_instance_report(rollout).get("tests_status")) + patch_applied_count = sum( + 1 for rollout in all_rollouts if _eval_instance_report(rollout).get("patch_successfully_applied") + ) + + metrics.update( + { + "task_count": len(tasks), + "rollout_count": rollout_count, + "max_rollouts_per_task": max_k, + "resolved_task_count": resolved_task_count, + "resolved_task_rate": 100.0 * resolved_task_count / len(tasks) if tasks else 0.0, + "eval_error_rollout_count": eval_error_count, + "eval_error_rate": 100.0 * eval_error_count / rollout_count if rollout_count else 0.0, + "eval_report_rollout_count": eval_report_count, + "eval_report_rate": 100.0 * eval_report_count / rollout_count if rollout_count else 0.0, + "tests_status_rollout_count": tests_status_count, + "tests_status_rate": 100.0 * tests_status_count / rollout_count if rollout_count else 0.0, + "patch_applied_rollout_count": patch_applied_count, + "patch_applied_rate": 100.0 * patch_applied_count / rollout_count if rollout_count else 0.0, + "per_task_metrics": self._compute_per_task_eval_metrics(tasks), + } + ) + + test_status_totals: dict[str, int] = {} + for rollout in all_rollouts: + for key, value in _test_status_counts(rollout).items(): + test_status_totals[key] = test_status_totals.get(key, 0) + value + metrics.update({f"tests_status/{key}": value for key, value in sorted(test_status_totals.items())}) + + return metrics + + def _compute_per_task_eval_metrics(self, tasks: list[list[dict[str, Any]]]) -> list[dict[str, Any]]: + per_task_metrics: list[dict[str, Any]] = [] + for fallback_idx, rollouts in enumerate(tasks): + if not rollouts: + continue + + first = rollouts[0] + task_index = first.get(TASK_INDEX_KEY_NAME, fallback_idx) + instance_id = first.get("instance_id") or _metadata_dict(first).get("instance_id") + resolved_count = sum(1 for rollout in rollouts if float(rollout.get("reward", 0.0) or 0.0) >= 1.0) + error_count = sum(1 for rollout in rollouts if _metadata_dict(rollout).get("error")) + eval_report_count = sum(1 for rollout in rollouts if _eval_report_map(rollout)) + tests_status_count = sum(1 for rollout in rollouts if _eval_instance_report(rollout).get("tests_status")) + patch_applied_count = sum( + 1 for rollout in rollouts if _eval_instance_report(rollout).get("patch_successfully_applied") + ) + + task_metrics: dict[str, Any] = { + TASK_INDEX_KEY_NAME: task_index, + "instance_id": instance_id, + "rollout_count": len(rollouts), + "resolved": resolved_count > 0, + "resolved_rollout_count": resolved_count, + "eval_error_rollout_count": error_count, + "eval_report_rollout_count": eval_report_count, + "tests_status_rollout_count": tests_status_count, + "patch_applied_rollout_count": patch_applied_count, + } + + test_status_totals: dict[str, int] = {} + for rollout in rollouts: + for key, value in _test_status_counts(rollout).items(): + test_status_totals[key] = test_status_totals.get(key, 0) + value + task_metrics.update({f"tests_status/{key}": value for key, value in sorted(test_status_totals.items())}) + per_task_metrics.append(task_metrics) + + return per_task_metrics + + def get_key_metrics(self, agent_metrics: dict[str, Any]) -> dict[str, Any]: + key_metrics: dict[str, Any] = {} + key_metrics.update(highest_k_metrics(agent_metrics, "pass@{k}", score_names=["accuracy"])) + key_metrics.update(highest_k_metrics(agent_metrics, "pass@1[avg-of-{k}]", score_names=["accuracy"])) + for key in ( + "mean/reward", + "resolved_task_count", + "task_count", + "resolved_task_rate", + "eval_error_rate", + "tests_status_rate", + ): + if key in agent_metrics: + key_metrics[key] = agent_metrics[key] + return key_metrics + + async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: + raise NotImplementedError + + async def run(self, body: MiniSWEAgentRunRequest) -> MiniSWEAgentVerifyResponse: + async with self.sem: + model_server_name = self.config.model_server.name + global_config_dict = ServerClient.load_from_global_config().global_config_dict + + model_server_config = get_first_server_config_dict( + global_config_dict, + model_server_name, + ) + + policy_model_name = global_config_dict["policy_model_name"] + + ##### MINI-SWE-AGENT CONFIG ##### + subset = body.subset + split = body.split + workers = 1 + run_golden = self.config.run_golden + base_url = f"http://{model_server_config['host']}:{model_server_config['port']}/v1" + dummy_key = "dummy_key" + model_name = f"hosted_vllm/{policy_model_name}" + step_timeout = self.config.step_timeout + eval_timeout = self.config.eval_timeout + step_limit = self.config.step_limit + + instance_id = body.instance_id + + mini_swe_config_path = _swebench_config_path() + config = yaml.safe_load(get_config_path(mini_swe_config_path).read_text()) + responses_create_params_dict = body.responses_create_params.model_dump(exclude_none=True) + + default_model_kwargs = config["model"]["model_kwargs"] + temperature = ( + body.responses_create_params.temperature + if body.responses_create_params.temperature is not None + else default_model_kwargs["temperature"] + ) + top_p = ( + body.responses_create_params.top_p + if body.responses_create_params.top_p is not None + else default_model_kwargs["top_p"] + ) + model_kwargs = _responses_create_params_to_model_kwargs( + responses_create_params_dict, + default_tool_choice=self.config.tool_choice, + ) + if model_kwargs: + config.setdefault("model", {}).setdefault("model_kwargs", {}).update(model_kwargs) + + output_file_dir = f"{Path.cwd()}/results/{subset}/{policy_model_name}" + config_path = mini_swe_config_path + should_write_config = bool(model_kwargs) + if self.config.sandbox_provider is None: + raise ValueError("mini_swe_agent_2 requires sandbox_provider") + config.setdefault("environment", {}).update(self.config.sandbox_environment_kwargs or {}) + config["environment"]["provider"] = self.config.sandbox_provider + config["environment"]["spec"] = _sandbox_spec_for_instance( + self.config.sandbox_spec, + resource_profiles=self.config.sandbox_resource_profiles, + instance_id=instance_id, + ) + should_write_config = True + + if should_write_config: + config_output_dir = Path(output_file_dir) / "_configs" + config_output_dir.mkdir(parents=True, exist_ok=True) + config_path = config_output_dir / f"{instance_id}.sandbox.yaml" + config_path.write_text(yaml.safe_dump(config, sort_keys=False)) + + if self.config.skip_if_exists: + if Path(f"{output_file_dir}/{instance_id}/{instance_id}.json").exists(): + with open(f"{output_file_dir}/{instance_id}/{instance_id}.json", "r") as f: + print(f"Skipping {instance_id} because it already exists") + verify_response = MiniSWEAgentVerifyResponse.model_validate_json(f.read()) + return verify_response + + #### RUN MINI-SWE-AGENT ##### + try: + params = dict( + subset=subset, + split=split, + workers=workers, + output=output_file_dir, + model=model_name, + api_key=dummy_key, + base_url=base_url, + env="sandbox", + run_golden=run_golden, + instance_id=instance_id, + config=config_path, + # TODO: add this later + instance_dict=body.model_dump(), + responses_create_params=json.dumps(responses_create_params_dict), + step_timeout=step_timeout, + eval_timeout=eval_timeout, + step_limit=step_limit, + ) + future = runner_ray_remote.remote(run_mini_swe_with_sandbox, params) + result = await asyncio.to_thread(ray.get, future) + result = result[instance_id] + input_messages = result["input_messages"] + response_output = result["response_output"] + responses = result["responses"] + reward = 1.0 if _is_resolved(instance_id, result["eval_report"]) else 0.0 + + except Exception as e: + error_info = {"error": str(e), "traceback": traceback.format_exc()} + print(f"Error running mini-swe-agent: {e}\n{error_info['traceback']}", flush=True) + result = {"eval_report": error_info} + input_messages = [] + response_output = [] + responses = [] + reward = 0.0 + + body.responses_create_params.input = input_messages + response = _default_response_object() + if responses: + response.update(dict(responses[-1])) + response.pop("extra", None) + response["model"] = policy_model_name + response["temperature"] = temperature + response["top_p"] = top_p + response["output"] = response_output + + verify_response = MiniSWEAgentVerifyResponse( + responses_create_params=body.responses_create_params, + reward=reward, + response=response, + instance_id=instance_id, + metadata=result.get("eval_report", {}) if result else {}, + ) + + output_path = Path(f"{output_file_dir}/{instance_id}") + output_path.mkdir(parents=True, exist_ok=True) + + with open(f"{output_file_dir}/{instance_id}/{instance_id}.json", "w") as f: + json.dump(verify_response.model_dump(), f) + + return verify_response + + +if __name__ == "__main__": + MiniSWEAgent.run_webserver() diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml new file mode 100644 index 0000000000..5ec99bfab1 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml @@ -0,0 +1,64 @@ +mini_swe_agent_2: + responses_api_agents: + mini_swe_agent_2: + entrypoint: app.py + domain: coding + description: Software engineering tasks driven by mini-swe-agent harness on OpenSandbox. + value: Improve agentic software engineering capabilities. + model_server: + type: responses_api_models + name: policy_model + concurrency: 64 + env: sandbox + sandbox_provider: + opensandbox: + connection: + domain: opensandbox-server.opensandbox-system.svc.cluster.local + api_key: ${oc.env:OPENSANDBOX_API_KEY} + protocol: http + use_server_proxy: true + exec_use_server_proxy: true + request_timeout_s: 300 + create: + request_timeout_s: 1200 + timeout_s: 1200 + skip_health_check: true + retries: 10 + retry_delay_s: 5.0 + retry_max_delay_s: 90.0 + probe: + timeout_s: 60 + deadline_s: 180 + stable_count: 2 + stable_delay_s: 1.0 + operations: + retries: 5 + retry_delay_s: 1.0 + retry_max_delay_s: 45.0 + command_retries: 3 + close_timeout_s: 30 + sandbox_spec: + timeout_s: 18000 + ready_timeout_s: 1200 + resources: + cpu: "2" + memory: 8Gi + ephemeral-storage: 20Gi + platform: + os: linux + arch: amd64 + metadata: + benchmark: swebench-verified + harness: mini-swe-agent + sandbox-api: opensandbox-sdk + sandbox_environment_kwargs: + cwd: /testbed + conda_env: testbed + activate_conda: true + user: root + delete: true + run_golden: false + step_timeout: 600 + eval_timeout: 1800 + skip_if_exists: false + step_limit: 250 diff --git a/responses_api_agents/mini_swe_agent_2/requirements.txt b/responses_api_agents/mini_swe_agent_2/requirements.txt new file mode 100644 index 0000000000..2f82dbd4c7 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/requirements.txt @@ -0,0 +1,3 @@ +-e nemo-gym[dev,sandbox] @ ../../ +mini-swe-agent==2.1.0 +swebench==4.1.0 diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py new file mode 100644 index 0000000000..34c60f21ab --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -0,0 +1,199 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""mini-swe-agent environment adapter backed by the Gym sandbox API.""" + +import os +import shlex +from dataclasses import dataclass, field +from typing import Any + + +try: + from minisweagent.exceptions import Submitted +except ModuleNotFoundError: + + class Submitted(Exception): + """Compatibility shim for local mini-swe-agent versions before v2.""" + + def __init__(self, *messages: dict[str, Any]) -> None: + self.messages = messages + super().__init__() + + +from nemo_gym.sandbox import Sandbox, SandboxSpec, rewrite_image + + +@dataclass +class MiniSWESandboxEnvironmentConfig: + """Configuration for mini-swe-agent runs inside a sandbox.""" + + image: str + cwd: str = "/workspace" + env: dict[str, str] = field(default_factory=dict) + forward_env: list[str] = field(default_factory=list) + timeout: int = 60 + step_timeout: int = 600 + eval_timeout: int = 1800 + interpreter: list[str] = field(default_factory=lambda: ["bash", "-c"]) + executable: str = "sandbox" + run_args: list[str] = field(default_factory=list) + start_args: list[str] = field(default_factory=list) + container_timeout: str = "2h" + instance_id: str | None = None + provider: dict[str, Any] = field(default_factory=dict) + spec: dict[str, Any] = field(default_factory=dict) + conda_env: str | None = None + activate_conda: bool = False + user: str | int | None = "root" + delete: bool = True + + +class MiniSWESandboxEnvironment: + """mini-swe-agent sync environment implemented with ``nemo_gym.sandbox.Sandbox``.""" + + def __init__( + self, + *, + config_class: type = MiniSWESandboxEnvironmentConfig, + **kwargs: Any, + ) -> None: + self.config = config_class(**kwargs) + if not self.config.provider: + raise ValueError("MiniSWESandboxEnvironment requires provider") + + self._handle: Any | None = None + self._closed = False + + spec_config = dict(self.config.spec) + image = spec_config.pop("image", None) or self.config.image + image = rewrite_image(image, spec_config.pop("image_rewrites", [])) + + env = dict(spec_config.pop("env", {})) + for key in self.config.forward_env: + value = os.getenv(key) + if value is not None: + env[key] = value + env.update(self.config.env) + + self._sandbox = Sandbox(self.config.provider) + self._handle = self._sandbox.create( + SandboxSpec( + image=image, + snapshot_id=spec_config.pop("snapshot_id", None), + timeout_s=spec_config.pop("timeout_s", None), + ready_timeout_s=spec_config.pop("ready_timeout_s", None), + env=env, + metadata={ + **spec_config.pop("metadata", {}), + "nemo_gym_agent": "mini_swe_agent_2", + "instance_id": (self.config.instance_id or "unknown")[:63], + }, + resources=spec_config.pop("resources", {}), + entrypoint=spec_config.pop("entrypoint", None), + extensions=spec_config.pop("extensions", {}), + platform=spec_config.pop("platform", None), + volumes=spec_config.pop("volumes", None), + skip_health_check=spec_config.pop("skip_health_check", None), + ) + ) + + def get_template_vars(self, **kwargs: Any) -> dict[str, Any]: + return {**self.config.__dict__, **kwargs} + + def serialize(self) -> dict[str, Any]: + return { + "info": { + "config": { + "environment": self.config.__dict__, + "environment_type": f"{self.__class__.__module__}.{self.__class__.__name__}", + } + } + } + + def _command(self, command: str, cwd: str) -> str: + if not self.config.activate_conda or not self.config.conda_env: + return command + quoted_cwd = shlex.quote(cwd) + quoted_env = shlex.quote(self.config.conda_env) + return ( + f"cd {quoted_cwd} && " + "source $(conda info --base)/etc/profile.d/conda.sh && " + f"conda activate {quoted_env} && " + f"{command}" + ) + + def execute( + self, + action: dict[str, Any] | str, + cwd: str = "", + is_eval: bool = False, + timeout: int | None = None, + ) -> dict[str, Any]: + command = action.get("command", "") if isinstance(action, dict) else action + timeout_s = timeout or (self.config.eval_timeout if is_eval else self.config.step_timeout) + exec_cwd = cwd or self.config.cwd + + result = self._sandbox.exec( + self._handle, + self._command(command, exec_cwd), + cwd="/", + timeout_s=timeout_s, + user=self.config.user, + ) + output = "\n".join(part for part in (result.stdout, result.stderr) if part) + response = { + "output": output, + "returncode": result.return_code, + "exception_info": "", + } + self._check_finished(response) + return response + + def _check_finished(self, output: dict[str, Any]) -> None: + """Match mini-swe-agent's submit sentinel handling for sandbox-backed runs.""" + lines = output.get("output", "").lstrip().splitlines(keepends=True) + if lines and lines[0].strip() == "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" and output["returncode"] == 0: + submission = "".join(lines[1:]) + raise Submitted( + { + "role": "exit", + "content": submission, + "extra": {"exit_status": "Submitted", "submission": submission}, + } + ) + + def cleanup(self) -> None: + if self._closed: + return + self._closed = True + try: + if self._handle is not None: + self._sandbox.close(self._handle, delete=self.config.delete) + self._handle = None + finally: + self._sandbox.shutdown() + + def __enter__(self) -> "MiniSWESandboxEnvironment": + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self.cleanup() + + def __del__(self) -> None: + if hasattr(self, "_closed") and not self._closed: + try: + self.cleanup() + except Exception: + pass 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 new file mode 100644 index 0000000000..77feb64695 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/tests/test_app.py @@ -0,0 +1,916 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 json +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any, Dict, Optional +from unittest.mock import MagicMock, patch + +import pytest +import yaml +from fastapi.testclient import TestClient + +from nemo_gym.config_types import AggregateMetricsRequest, ModelServerRef +from nemo_gym.global_config import ROLLOUT_INDEX_KEY_NAME, TASK_INDEX_KEY_NAME +from nemo_gym.openai_utils import ( + NeMoGymChatCompletionCreateParamsNonStreaming, + NeMoGymResponseCreateParamsNonStreaming, +) +from nemo_gym.server_utils import ServerClient + + +try: + __import__("minisweagent.config") +except ModuleNotFoundError as exc: + if exc.name not in {"minisweagent", "minisweagent.config"}: + raise + minisweagent_module = ModuleType("minisweagent") + minisweagent_module.__path__ = [] + minisweagent_config_module = ModuleType("minisweagent.config") + minisweagent_config_module.builtin_config_dir = Path("/tmp/minisweagent/config") + minisweagent_config_module.get_config_path = Path + sys.modules["minisweagent"] = minisweagent_module + sys.modules["minisweagent.config"] = minisweagent_config_module + +from responses_api_agents.mini_swe_agent_2 import app as mini_swe_app_module +from responses_api_agents.mini_swe_agent_2.app import ( + MiniSWEAgent, + MiniSWEAgentConfig, + MiniSWEAgentRunRequest, + MiniSWEAgentVerifyResponse, + _is_resolved, + _json_dict_from_metadata, + _message_content_to_text, + _responses_create_params_to_model_kwargs, + _run_mini_swe_v2, + _sandbox_spec_for_instance, + _split_trajectory_for_responses, + _swebench_config_path, + _swebench_image_name, + run_mini_swe_with_sandbox, +) + + +DEFAULT_RUN_MINI_SWE_RESULT = { + "test_instance_123": { + "input_messages": [ + {"type": "message", "role": "system", "content": "You are a helpful assistant."}, + {"type": "message", "role": "user", "content": "Fix this bug."}, + ], + "response_output": [ + { + "id": "msg-1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "I'll help you fix the bug.", "annotations": []}], + } + ], + "responses": [ + { + "id": "resp-1", + "object": "response", + "output": [], + } + ], + "eval_report": { + "eval_report": { + "test_instance_123": { + "resolved": True, + "tests_status": { + "FAIL_TO_PASS": {"success": ["test1"], "failure": []}, + "PASS_TO_PASS": {"success": ["test2"], "failure": []}, + }, + } + } + }, + } +} + +DEFAULT_CONFIG_YAML = """ +model: + model_kwargs: + temperature: 0.5 + top_p: 0.8 +""" + +DEFAULT_CHAT_COMPLETION = { + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "test_model", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! How can I help you today?", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21}, +} + + +def create_test_config( + host: str = "0.0.0.0", + port: int = 8080, + model_name: str = "test_model", +) -> MiniSWEAgentConfig: + return MiniSWEAgentConfig( + name="mini_swe_agent_2", + host=host, + port=port, + entrypoint="", + model_server=ModelServerRef( + type="responses_api_models", + name=model_name, + ), + env="sandbox", + concurrency=1, + sandbox_provider={"opensandbox": {}}, + sandbox_spec={}, + ) + + +def setup_server_client_mocks(mock_load_from_global_config, mock_get_first_server_config_dict): + mock_server_client_instance = MagicMock() + mock_server_client_instance.global_config_dict = {"policy_model_name": "test_model"} + 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, + } + + +def setup_config_path_mock(mock_get_config_path, config_yaml: str = DEFAULT_CONFIG_YAML): + mock_config_path = MagicMock() + mock_config_path.read_text.return_value = config_yaml + mock_get_config_path.return_value = mock_config_path + + +def setup_run_mini_swe_mock( + mock_to_thread, + mock_runner_ray_remote, + run_mini_swe_result: Dict[str, Any] = None, +): + """Setup mock for Ray-based run_mini_swe execution""" + 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() + mock_runner_ray_remote.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", + temperature: float = 0.5, + top_p: float = 0.8, + max_output_tokens: int | None = None, + metadata: dict[str, Any] | None = None, + subset: str = "gym", + split: str = "train", + input_data: list = None, +) -> MiniSWEAgentRunRequest: + """Create a test run request with default values.""" + if input_data is None: + input_data = [] + + return MiniSWEAgentRunRequest( + instance_id=instance_id, + subset=subset, + split=split, + responses_create_params=NeMoGymResponseCreateParamsNonStreaming( + temperature=temperature, + top_p=top_p, + max_output_tokens=max_output_tokens, + metadata=metadata, + input=input_data, + ), + ) + + +def create_chat_completion_request( + model: str = "test_model", + messages: list = None, + temperature: float = 0.7, + max_tokens: Optional[int] = None, +) -> NeMoGymChatCompletionCreateParamsNonStreaming: + if messages is None: + messages = [{"role": "user", "content": "Hello!"}] + + kwargs = {"model": model, "messages": messages, "temperature": temperature} + if max_tokens is not None: + kwargs["max_tokens"] = max_tokens + + return NeMoGymChatCompletionCreateParamsNonStreaming(**kwargs) + + +def assert_run_response( + response: MiniSWEAgentVerifyResponse, + expected_reward: float = 1.0, + expected_temperature: float = 0.5, + expected_top_p: float = 0.8, + expected_input_length: int = 2, +): + assert isinstance(response, MiniSWEAgentVerifyResponse) + assert response.reward == expected_reward + assert response.responses_create_params.temperature == expected_temperature + assert response.responses_create_params.top_p == expected_top_p + assert len(response.responses_create_params.input) == expected_input_length + + if expected_input_length >= 2: + assert response.responses_create_params.input[0]["role"] == "system" + assert response.responses_create_params.input[1]["role"] == "user" + + +def assert_run_mini_swe_called( + mock_to_thread, + 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 + + +class TestApp: + def test_sanity(self) -> None: + config = create_test_config(model_name="") + MiniSWEAgent(config=config, server_client=MagicMock(spec=ServerClient)) + + def test_response_param_helpers_cover_metadata_and_tool_choice_modes(self) -> None: + assert _json_dict_from_metadata(None, field_name="extra_body") == {} + assert _json_dict_from_metadata({"top_k": 20}, field_name="extra_body") == {"top_k": 20} + + kwargs = _responses_create_params_to_model_kwargs( + { + "temperature": 0.6, + "top_p": 0.95, + "max_output_tokens": 123, + "metadata": { + "extra_body": json.dumps({"top_k": 20}), + "chat_template_kwargs": json.dumps({"enable_thinking": True}), + }, + "tool_choice": {"type": "function", "function": {"name": "python"}}, + } + ) + + assert kwargs == { + "temperature": 0.6, + "top_p": 0.95, + "max_tokens": 123, + "extra_body": {"top_k": 20, "chat_template_kwargs": {"enable_thinking": True}}, + "tool_choice": {"type": "function", "function": {"name": "python"}}, + } + assert _responses_create_params_to_model_kwargs({"tool_choice": "bash"})["tool_choice"] == { + "type": "function", + "function": {"name": "bash"}, + } + assert ( + _responses_create_params_to_model_kwargs({"tool_choice": "auto"}, default_tool_choice="none")[ + "tool_choice" + ] + == "none" + ) + + with pytest.raises(ValueError, match="extra_body"): + _json_dict_from_metadata("[]", field_name="extra_body") + + def test_sandbox_resource_profiles_override_static_resources(self) -> None: + spec = _sandbox_spec_for_instance( + {"resources": {"cpu": "1", "memory": "8Gi", "ephemeral-storage": "20Gi"}}, + resource_profiles=[ + {"cpu": "250m", "memory": "3Gi", "ephemeral-storage": "1Gi"}, + {"cpu": "500m", "memory": "4Gi", "ephemeral-storage": "1Gi"}, + ], + instance_id="django__django-12345", + ) + + assert spec["resources"] in ( + {"cpu": "250m", "memory": "3Gi", "ephemeral-storage": "1Gi"}, + {"cpu": "500m", "memory": "4Gi", "ephemeral-storage": "1Gi"}, + ) + assert _sandbox_spec_for_instance(None, resource_profiles=None, instance_id="task") == {} + + def test_split_trajectory_and_resolution_helpers_cover_edge_cases(self) -> None: + input_messages, output_items, raw_responses = _split_trajectory_for_responses( + [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "user"}, + { + "role": "assistant", + "content": "answer", + "tool_calls": [{"id": "call-1", "function": {"name": "bash", "arguments": '{"command":"pwd"}'}}], + }, + {"role": "tool", "tool_call_id": "call-1", "content": "tool output"}, + {"type": "function_call_output", "call_id": "call-2", "output": "raw", "extra": {"ignored": True}}, + {"object": "response", "output": [{"type": "message", "content": "raw"}], "extra": {"ignored": True}}, + ] + ) + + assert input_messages == [ + {"type": "message", "role": "system", "content": "sys"}, + {"type": "message", "role": "user", "content": "user"}, + ] + assert any(item["type"] == "function_call" and item["call_id"] == "call-1" for item in output_items) + assert any(item["type"] == "function_call_output" and item["call_id"] == "call-1" for item in output_items) + assert any(item["type"] == "function_call_output" and item["call_id"] == "call-2" for item in output_items) + assert raw_responses == [{"object": "response", "output": [{"type": "message", "content": "raw"}]}] + + assert not _is_resolved("task", {}) + assert not _is_resolved("task", {"eval_report": {"task": {"resolved": True}}}) + assert not _is_resolved( + "task", + { + "eval_report": { + "task": { + "resolved": True, + "tests_status": {"FAIL_TO_PASS": {"success": [], "failure": []}}, + } + } + }, + ) + + def test_misc_mini_swe_helpers(self, monkeypatch, tmp_path) -> None: + assert _swebench_image_name({"instance_id": "django__django-1"}, "verified") == ( + "docker.io/swebench/sweb.eval.x86_64.django_1776_django-1:latest" + ) + assert _swebench_image_name({"instance_id": "django__django-1"}, "lite") == ( + "docker.io/xingyaoww/sweb.eval.x86_64.django_s_django-1:latest" + ) + assert _swebench_image_name({"instance_id": "x", "image_name": "custom:image"}, "verified") == "custom:image" + assert _message_content_to_text("hello") == "hello" + assert _message_content_to_text(None) == "" + assert _message_content_to_text([{"text": "one"}, {"content": "two"}, 3]) == "one\ntwo\n3" + + builtin_dir = tmp_path / "configs" + benchmark_dir = builtin_dir / "benchmarks" + benchmark_dir.mkdir(parents=True) + (benchmark_dir / "swebench.yaml").write_text("{}", encoding="utf-8") + monkeypatch.setattr(mini_swe_app_module, "builtin_config_dir", builtin_dir) + assert _swebench_config_path() == benchmark_dir / "swebench.yaml" + monkeypatch.setattr(mini_swe_app_module, "builtin_config_dir", tmp_path / "missing") + assert _swebench_config_path() == tmp_path / "missing" / "extra" / "swebench.yaml" + + def test_run_mini_swe_records_completion_and_errors(self, monkeypatch) -> None: + monkeypatch.setattr( + mini_swe_app_module, + "_run_mini_swe_v2", + lambda **_params: { + "task-1": { + "eval_report": { + "task-1": {"resolved": True}, + } + } + }, + ) + assert run_mini_swe_with_sandbox( + env="sandbox", + instance_id="task-1", + ) == {"task-1": {"eval_report": {"task-1": {"resolved": True}}}} + + def fail_runner(**_params): + raise RuntimeError("boom") + + monkeypatch.setattr(mini_swe_app_module, "_run_mini_swe_v2", fail_runner) + with pytest.raises(RuntimeError, match="boom"): + run_mini_swe_with_sandbox(env="sandbox", instance_id="task-1") + + monkeypatch.setattr( + mini_swe_app_module, + "_run_mini_swe_v2", + lambda **_params: {"task-1": {"eval_report": {"task-1": {"resolved": False}}}}, + ) + assert run_mini_swe_with_sandbox(env="sandbox", instance_id="task-1") == { + "task-1": {"eval_report": {"task-1": {"resolved": False}}} + } + + monkeypatch.setattr(mini_swe_app_module, "_run_mini_swe_v2", lambda **_params: {"task-1": "bad"}) + assert run_mini_swe_with_sandbox(env="sandbox", instance_id="task-1") == {"task-1": "bad"} + + def test_run_mini_swe_v2_success_and_golden_paths(self, monkeypatch, tmp_path) -> None: + holder: dict[str, Any] = {} + + class FakeLogger: + def info(self, _message: str) -> None: + return None + + def setup_logger(_instance_id: str, _log_file: Path) -> FakeLogger: + return FakeLogger() + + def make_test_spec(instance: dict[str, Any]) -> SimpleNamespace: + return SimpleNamespace( + instance_id=instance["instance_id"], + eval_script="#!/bin/bash\npytest -q", + ) + + def get_eval_report( + *, + test_spec: SimpleNamespace, + prediction: dict[str, Any], + test_log_path: str, + **_kwargs: Any, + ): + assert Path(test_log_path).exists() + return {test_spec.instance_id: {"resolved": True, "prediction": prediction}} + + class FakeEnv: + def __init__(self, config: dict[str, Any]) -> None: + self.config = config + self.commands: list[tuple[str, bool]] = [] + self.cleaned = False + + def execute(self, command: str, is_eval: bool = False) -> dict[str, Any]: + self.commands.append((command, is_eval)) + return {"output": "tests passed", "returncode": 0} + + def cleanup(self) -> None: + self.cleaned = True + + class FakeAgent: + def __init__(self, model: Any, env: FakeEnv, **agent_config: Any) -> None: + self.model = model + self.env = env + self.agent_config = agent_config + holder["agent_config"] = agent_config + + def run(self, problem_statement: str) -> dict[str, Any]: + assert problem_statement == "Fix the bug" + return {"exit_status": "submitted", "submission": "diff --git a/file b/file"} + + def save(self, path: Path | None, metadata: dict[str, Any]) -> dict[str, Any]: + holder["save_path"] = path + holder["save_metadata"] = metadata + if path is not None: + path.write_text("{}", encoding="utf-8") + return { + "messages": [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": [{"text": "problem"}]}, + { + "id": "resp-1", + "object": "response", + "output": [ + { + "id": "msg-1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "answer", "annotations": []}], + }, + { + "type": "function_call", + "name": "bash", + "call_id": "call-1", + "arguments": json.dumps({"command": "echo hi"}), + }, + ], + "extra": {"actions": [{"command": "echo hi", "tool_call_id": "call-1"}]}, + }, + { + "type": "function_call_output", + "call_id": "call-1", + "output": "tool output", + "extra": {"raw_output": "tool output"}, + }, + ] + } + + def get_environment(config: dict[str, Any]) -> FakeEnv: + env = FakeEnv(config) + holder["env"] = env + return env + + def get_model(config: dict[str, Any]) -> SimpleNamespace: + holder["model_config"] = config + return SimpleNamespace(config=config) + + module_specs = { + "swebench": ModuleType("swebench"), + "swebench.harness": ModuleType("swebench.harness"), + "swebench.harness.constants": ModuleType("swebench.harness.constants"), + "swebench.harness.docker_build": ModuleType("swebench.harness.docker_build"), + "swebench.harness.grading": ModuleType("swebench.harness.grading"), + "swebench.harness.test_spec": ModuleType("swebench.harness.test_spec"), + "swebench.harness.test_spec.test_spec": ModuleType("swebench.harness.test_spec.test_spec"), + "minisweagent.agents": ModuleType("minisweagent.agents"), + "minisweagent.agents.default": ModuleType("minisweagent.agents.default"), + "minisweagent.environments": ModuleType("minisweagent.environments"), + "minisweagent.models": ModuleType("minisweagent.models"), + } + module_specs["swebench.harness.constants"].SWEbenchInstance = dict + module_specs["swebench.harness.docker_build"].setup_logger = setup_logger + module_specs["swebench.harness.grading"].get_eval_report = get_eval_report + module_specs["swebench.harness.test_spec.test_spec"].make_test_spec = make_test_spec + module_specs["minisweagent.agents.default"].DefaultAgent = FakeAgent + module_specs["minisweagent.environments"].get_environment = get_environment + module_specs["minisweagent.models"].get_model = get_model + for name, module in module_specs.items(): + monkeypatch.setitem(sys.modules, name, module) + + config_path = tmp_path / "swebench.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": {"model_kwargs": {"max_output_tokens": 99}}, + "environment": {}, + "agent": {"step_limit": 1, "collapse_limit": 3}, + } + ), + encoding="utf-8", + ) + monkeypatch.setattr(mini_swe_app_module, "get_config_path", lambda _config: config_path) + monkeypatch.setattr(mini_swe_app_module, "uuid4", lambda: "uuid") + monkeypatch.setattr(mini_swe_app_module.time, "time", lambda: 1234) + + params = { + "instance_dict": { + "instance_id": "django__django-123", + "problem_statement": "Fix the bug", + "patch": "gold", + }, + "instance_id": "django__django-123", + "output": str(tmp_path / "out"), + "config": "swebench", + "model": "hosted/model", + "api_key": "key", # pragma: allowlist secret + "base_url": "http://model/v1", + "subset": "verified", + "step_timeout": 30, + "eval_timeout": 60, + "env": "sandbox", + "step_limit": 7, + "run_golden": False, + } + + result = _run_mini_swe_v2(**params) + + env = holder["env"] + assert env.cleaned is True + assert env.config["environment_class"].endswith("MiniSWESandboxEnvironment") + assert env.config["image"] == "docker.io/swebench/sweb.eval.x86_64.django_1776_django-123:latest" + assert holder["model_config"]["model_class"] == "litellm" + assert holder["model_config"]["model_name"] == "hosted/model" + assert holder["model_config"]["model_kwargs"]["max_tokens"] == 99 + assert holder["model_config"]["model_kwargs"]["base_url"] == "http://model/v1" + assert "api_base" not in holder["model_config"]["model_kwargs"] + assert holder["agent_config"]["step_limit"] == 7 + assert holder["save_metadata"] == {"instance_id": "django__django-123"} + assert result["django__django-123"]["input_messages"] == [ + {"type": "message", "role": "system", "content": "sys"}, + {"type": "message", "role": "user", "content": "problem"}, + ] + assert result["django__django-123"]["response_output"] == [ + { + "id": "msg-1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "answer", "annotations": []}], + }, + { + "type": "function_call", + "name": "bash", + "call_id": "call-1", + "arguments": json.dumps({"command": "echo hi"}), + }, + {"type": "function_call_output", "call_id": "call-1", "output": "tool output"}, + ] + assert result["django__django-123"]["responses"] == [ + { + "id": "resp-1", + "object": "response", + "output": [ + { + "id": "msg-1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "answer", "annotations": []}], + }, + { + "type": "function_call", + "name": "bash", + "call_id": "call-1", + "arguments": json.dumps({"command": "echo hi"}), + }, + ], + } + ] + + golden_params = params | {"run_golden": True} + result = _run_mini_swe_v2(**golden_params) + + env = holder["env"] + assert env.cleaned is True + assert env.config["environment_class"].endswith("MiniSWESandboxEnvironment") + assert [command for command, _ in env.commands[:4]] == [ + "cat > patch.diff <<'EOF'\ngold\n\nEOF", + "git status --porcelain", + "git apply --check patch.diff", + "git apply patch.diff", + ] + assert result["django__django-123"]["exit_status"] == "Gold Patch Applied" + + string_params = params | { + "instance_dict": json.dumps( + {"instance_id": "django__django-123", "problem_statement": "Fix the bug", "patch": "gold"} + ), + } + assert "django__django-123" in _run_mini_swe_v2(**string_params) + + with pytest.raises(ValueError, match="instance_dict"): + _run_mini_swe_v2(**(params | {"instance_dict": None})) + + @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_successful_execution( + self, + mock_to_thread, + mock_runner_ray_remote, + mock_get_config_path, + mock_get_first_server_config_dict, + mock_load_from_global_config, + ) -> None: + """Test successful execution of the run method with mocked run_mini_swe.""" + + config = create_test_config() + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + 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) + + run_request = create_run_request() + + response = await server.run(run_request) + + assert_run_response(response) + + assert_run_mini_swe_called(mock_to_thread) + + @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, + mock_load_from_global_config, + tmp_path, + monkeypatch, + ) -> None: + monkeypatch.chdir(tmp_path) + config = create_test_config() + config.tool_choice = "bash" + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + 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) + + run_request = create_run_request( + temperature=0.6, + top_p=0.95, + max_output_tokens=49152, + metadata={ + "extra_body": '{"top_k":20,"min_p":0.0,"presence_penalty":0.0,"repetition_penalty":1.0}', + "chat_template_kwargs": '{"enable_thinking":true}', + }, + ) + + await server.run(run_request) + + call_args = mock_runner_ray_remote.remote.call_args + params = call_args.args[1] + generated_config = yaml.safe_load(Path(params["config"]).read_text()) + model_kwargs = generated_config["model"]["model_kwargs"] + assert model_kwargs["temperature"] == 0.6 + assert model_kwargs["top_p"] == 0.95 + assert model_kwargs["max_tokens"] == 49152 + assert "max_output_tokens" not in model_kwargs + assert model_kwargs["tool_choice"] == {"type": "function", "function": {"name": "bash"}} + assert model_kwargs["extra_body"] == { + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 0.0, + "repetition_penalty": 1.0, + "chat_template_kwargs": {"enable_thinking": True}, + } + + @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_failed_execution( + self, + mock_to_thread, + mock_runner_ray_remote, + mock_get_config_path, + mock_get_first_server_config_dict, + mock_load_from_global_config, + ) -> None: + """Test run method when run_mini_swe fails.""" + + config = create_test_config() + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + 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") + + run_request = create_run_request(instance_id="test_instance_456", temperature=0.3, top_p=0.95) + + response = await server.run(run_request) + + assert_run_response( + response, + expected_reward=0.0, + expected_temperature=0.3, + expected_top_p=0.95, + expected_input_length=0, + ) + + assert_run_mini_swe_called(mock_to_thread, 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, + mock_load_from_global_config, + ) -> None: + config = create_test_config() + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + 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") + + run_request = create_run_request(instance_id="test_instance_789", temperature=0.2, top_p=1.0) + + response = await server.run(run_request) + + assert_run_response( + response, + expected_reward=0.0, + expected_temperature=0.2, + expected_top_p=1.0, + expected_input_length=0, + ) + + assert_run_mini_swe_called(mock_to_thread, instance_id="test_instance_789") + + async def test_responses_not_implemented(self) -> None: + config = create_test_config() + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + request_body = NeMoGymResponseCreateParamsNonStreaming(temperature=0.7, top_p=0.9, input=[]) + + with pytest.raises(NotImplementedError): + await server.responses(request_body) + + async def test_aggregate_metrics_includes_eval_results(self) -> None: + config = create_test_config() + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + responses = [ + { + TASK_INDEX_KEY_NAME: 0, + ROLLOUT_INDEX_KEY_NAME: 0, + "instance_id": "task-a", + "reward": 1.0, + "metadata": { + "instance_id": "task-a", + "eval_report": { + "task-a": { + "resolved": True, + "patch_successfully_applied": True, + "tests_status": { + "FAIL_TO_PASS": {"success": ["test-a"], "failure": []}, + "PASS_TO_PASS": {"success": ["test-b"], "failure": []}, + }, + } + }, + }, + }, + { + TASK_INDEX_KEY_NAME: 0, + ROLLOUT_INDEX_KEY_NAME: 1, + "instance_id": "task-a", + "reward": 0.0, + "metadata": { + "instance_id": "task-a", + "eval_report": { + "task-a": { + "resolved": False, + "patch_successfully_applied": True, + "tests_status": { + "FAIL_TO_PASS": {"success": [], "failure": ["test-a"]}, + "PASS_TO_PASS": {"success": ["test-b"], "failure": []}, + }, + } + }, + }, + }, + { + TASK_INDEX_KEY_NAME: 1, + ROLLOUT_INDEX_KEY_NAME: 0, + "instance_id": "task-b", + "reward": 0.0, + "metadata": {"error": "boom"}, + }, + { + TASK_INDEX_KEY_NAME: 1, + ROLLOUT_INDEX_KEY_NAME: 1, + "instance_id": "task-b", + "reward": 0.0, + "metadata": {"error": "boom"}, + }, + ] + + result = await server.aggregate_metrics(AggregateMetricsRequest(verify_responses=responses)) + + assert result.agent_metrics["pass@2/accuracy"] == pytest.approx(50.0) + assert result.agent_metrics["resolved_task_count"] == 1 + assert result.agent_metrics["eval_error_rollout_count"] == 2 + assert result.agent_metrics["tests_status/fail_to_pass_success"] == 1 + assert result.key_metrics["pass@2/accuracy"] == pytest.approx(50.0) + + groups = {group[TASK_INDEX_KEY_NAME]: group for group in result.group_level_metrics} + assert groups[0]["instance_id"] == "task-a" + assert groups[0]["resolved"] is True + assert groups[0]["tests_status_rollout_count"] == 2 + assert groups[1]["instance_id"] == "task-b" + assert groups[1]["eval_error_rollout_count"] == 2 + + def test_endpoints_registration(self) -> None: + config = create_test_config() + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + app = server.setup_webserver() + client = TestClient(app, raise_server_exceptions=False) + + response = client.post("/v1/responses", json={"temperature": 0.7, "top_p": 0.9, "input": []}) + assert response.status_code == 500 + + run_response = client.post("/run", json={}) + assert run_response.status_code != 404 + + aggregate_response = client.post("/aggregate_metrics", json={"verify_responses": []}) + assert aggregate_response.status_code == 200 diff --git a/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py new file mode 100644 index 0000000000..331d732641 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +from responses_api_agents.mini_swe_agent_2.sandbox_environment import MiniSWESandboxEnvironment, Submitted + + +def test_check_finished_raises_submitted_for_submit_sentinel() -> None: + env = MiniSWESandboxEnvironment.__new__(MiniSWESandboxEnvironment) + + try: + env._check_finished( + { + "output": "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\npatch contents\n", + "returncode": 0, + "exception_info": "", + } + ) + except Submitted as error: + assert error.messages == ( + { + "role": "exit", + "content": "patch contents\n", + "extra": {"exit_status": "Submitted", "submission": "patch contents\n"}, + }, + ) + else: + raise AssertionError("Expected Submitted") + + +def test_check_finished_ignores_nonzero_submit_sentinel() -> None: + env = MiniSWESandboxEnvironment.__new__(MiniSWESandboxEnvironment) + + env._check_finished( + { + "output": "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\npatch contents\n", + "returncode": 1, + "exception_info": "", + } + ) diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py new file mode 100644 index 0000000000..5d5a9756c4 --- /dev/null +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -0,0 +1,604 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 importlib.util +from dataclasses import dataclass +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +from nemo_gym.sandbox.providers.base import SandboxSpec +from nemo_gym.sandbox.providers.opensandbox import provider as opensandbox_provider + + +pytestmark = pytest.mark.skipif( + importlib.util.find_spec("tenacity") is None, + reason="tenacity optional sandbox dependency is not installed", +) + + +@dataclass(frozen=True) +class FakePlatformSpec: + os: str + arch: str + + +class FakeConnectionConfig: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + +class FakeSandbox: + created_kwargs: dict[str, Any] = {} + connected_args: tuple[Any, ...] = () + connected_kwargs: dict[str, Any] = {} + + def __init__(self, sandbox_id: str = "sandbox-1") -> None: + self.id = sandbox_id + + @classmethod + async def create(cls, *_args: Any, **kwargs: Any) -> "FakeSandbox": + cls.created_kwargs = kwargs + return cls() + + @classmethod + async def connect(cls, *args: Any, **kwargs: Any) -> "FakeSandbox": + cls.connected_args = args + cls.connected_kwargs = kwargs + return cls() + + +@dataclass +class FakePoolCreationSpec: + image: str + entrypoint: list[str] | None = None + resource: dict[str, str] | None = None + env: dict[str, str] | None = None + metadata: dict[str, str] | None = None + extensions: dict[str, str] | None = None + platform: Any | None = None + volumes: list[Any] | None = None + + +class FakeAcquirePolicy: + FAIL_FAST = "fail_fast" + + +class FakeStateStore: + pass + + +class FakeSnapshot: + idle_count = 1 + state = None + + +class FakeSandboxPoolAsync: + received_kwargs: dict[str, Any] = {} + + def __init__(self, **kwargs: Any) -> None: + self.received_kwargs = kwargs + type(self).received_kwargs = kwargs + + async def start(self) -> None: + return None + + async def snapshot(self) -> FakeSnapshot: + return FakeSnapshot() + + async def resize(self, _count: int) -> None: + return None + + async def acquire( + self, + *, + sandbox_timeout: timedelta | None, + policy: str, + ) -> FakeSandbox: + del sandbox_timeout, policy + creation_spec = self.received_kwargs["creation_spec"] + return await FakeSandbox.create( + creation_spec.image, + platform=creation_spec.platform, + ) + + async def shutdown(self, *, graceful: bool) -> None: + del graceful + + async def release_all_idle(self) -> None: + return None + + +@pytest.fixture +def fake_opensandbox_sdk(monkeypatch: pytest.MonkeyPatch) -> None: + def require_sdk() -> tuple[Any, Any, Any, Any, Any]: + return FakeSandbox, FakeConnectionConfig, object, FakePlatformSpec, object + + def require_sdk_pool() -> tuple[Any, Any, Any, Any]: + return ( + FakeAcquirePolicy, + FakeStateStore, + FakePoolCreationSpec, + FakeSandboxPoolAsync, + ) + + monkeypatch.setattr(opensandbox_provider, "_require_opensandbox_sdk", require_sdk) + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk_pool", + require_sdk_pool, + ) + + +async def test_sdk_pool_passes_platform_through_pool_creation_spec( + fake_opensandbox_sdk: None, +) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + connection={"request_timeout_s": 10}, + probe={"command": None}, + ) + + handles = await provider.create_batch( + SandboxSpec( + image="mirror.gcr.io/astral/uv:python3.12-bookworm-slim", + platform={"os": "linux", "arch": "amd64"}, + ), + 1, + ) + + assert len(handles) == 1 + assert handles[0].sandbox_id == "sandbox-1" + assert "sandbox_factory" not in FakeSandboxPoolAsync.received_kwargs + assert FakeSandbox.created_kwargs["platform"] == FakePlatformSpec( + os="linux", + arch="amd64", + ) + + +async def test_connect_passes_configured_connect_timeout( + fake_opensandbox_sdk: None, +) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + connection={"connect_timeout_s": 300, "request_timeout_s": 10}, + probe={"command": None}, + ) + + handle = await provider.connect("sandbox-123") + + assert handle.sandbox_id == "sandbox-1" + assert FakeSandbox.connected_args == ("sandbox-123",) + assert FakeSandbox.connected_kwargs["connect_timeout"] == timedelta(seconds=300) + + +def test_provider_validation_and_retry_helpers() -> None: + with pytest.raises(ValueError, match="image_pull_policy"): + opensandbox_provider.validate_image_pull_policy("Sometimes") + + invalid_kwargs = [ + {"pool": {"concurrency": 0}}, + {"connection": {"connect_timeout_s": 0}}, + {"pool": {"progress_timeout_s": 0}}, + {"create": {"timeout_s": 0}}, + {"probe": {"timeout_s": 0}}, + {"probe": {"deadline_s": 0}}, + {"probe": {"sample_count": 0}}, + {"probe": {"stable_count": 0}}, + {"probe": {"stable_delay_s": -1}}, + {"create": {"retries": -1}}, + {"create": {"retry_delay_s": -1}}, + {"create": {"retry_max_delay_s": -1}}, + {"operations": {"retries": -1}}, + {"operations": {"retry_delay_s": -1}}, + {"operations": {"retry_max_delay_s": -1}}, + {"operations": {"command_retries": -1}}, + {"pool": {"reconcile_interval_s": 0}}, + {"pool": {"acquire_poll_interval_s": 0}}, + {"pool": {"idle_timeout_s": 0}}, + {"pool": {"primary_lock_ttl_s": 0}}, + {"operations": {"close_timeout_s": 0}}, + {"create": {"connect_attempt_timeout_s": 0}}, + {"create": {"connect_poll_s": 0}}, + {"create": {"image_pull_policy": "Sometimes"}}, + ] + for kwargs in invalid_kwargs: + with pytest.raises(ValueError): + opensandbox_provider.OpenSandboxProvider(**kwargs) + with pytest.raises(TypeError): + opensandbox_provider.OpenSandboxProvider(**{"batch_" + "create_retries": 1}) + with pytest.raises(TypeError): + opensandbox_provider.OpenSandboxProvider(connection=object()) + + assert opensandbox_provider._exception_status_code(RuntimeError("HTTP status code: 503")) == 503 + assert opensandbox_provider._exception_status_code(RuntimeError("plain error")) is None + attrs = opensandbox_provider._sdk_error_attributes( + RuntimeError("HTTP 502 bad gateway"), + operation="exec", + sandbox_id="sandbox-1", + attempt_number=2, + max_attempts=3, + sleep_s=0.5, + ) + assert attrs["status_code"] == 502 + assert attrs["attempt_number"] == 2 + assert attrs["next_sleep_s"] == 0.5 + assert opensandbox_provider._seconds_to_timedelta(None) is None + assert opensandbox_provider._seconds_to_timedelta(1.5) == timedelta(seconds=1.5) + + +def test_connection_config_exec_proxy_and_image_policy(fake_opensandbox_sdk: None) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + connection={ + "domain": "sandbox.example", + "api_key": "key", # pragma: allowlist secret + "protocol": "https", + "use_server_proxy": True, + "exec_use_server_proxy": False, + "request_timeout_s": 10, + } + ) + + config = provider._connection_config() + assert config.kwargs == { + "domain": "sandbox.example", + "api_key": "key", # pragma: allowlist secret + "protocol": "https", + "use_server_proxy": True, + "request_timeout": timedelta(seconds=10), + } + exec_config = provider._exec_connection_config(request_timeout_s=3) + assert exec_config.kwargs["use_server_proxy"] is False + assert exec_config.kwargs["request_timeout"] == timedelta(seconds=3) + + spec = SandboxSpec(image="image:tag", extensions={"imagePullPolicy": "Never"}) + updated = provider._with_default_image_pull_policy(spec) + assert updated.extensions["imagePullPolicy"] == "Never" + assert updated.extensions["opensandbox.extensions.image-pull-policy"] == "Never" + + no_policy_provider = opensandbox_provider.OpenSandboxProvider(create={"image_pull_policy": None}) + assert no_policy_provider._with_default_image_pull_policy(spec) is spec + + +async def test_wait_sdk_pool_idle_success_partial_and_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + class Snapshot: + def __init__(self, idle_count: int) -> None: + self.idle_count = idle_count + self.state = SimpleNamespace(value="warming") + + class FakePool: + def __init__(self, counts: list[int]) -> None: + self.counts = counts + self.index = 0 + self._config = SimpleNamespace(pool_name="pool-1") + + async def snapshot(self) -> Snapshot: + count = self.counts[min(self.index, len(self.counts) - 1)] + self.index += 1 + return Snapshot(count) + + async def no_sleep(_seconds: float) -> None: + return None + + monkeypatch.setattr(opensandbox_provider.asyncio, "sleep", no_sleep) + + provider = opensandbox_provider.OpenSandboxProvider( + pool={"acquire_poll_interval_s": 0.01}, + probe={"command": None}, + ) + assert ( + await provider._wait_sdk_pool_idle( + FakePool([0, 1, 2]), + spec=SandboxSpec(image="image:tag"), + requested=2, + timeout_s=1, + allow_partial=False, + ) + == 2 + ) + assert ( + await provider._wait_sdk_pool_idle( + FakePool([1]), + spec=SandboxSpec(image="image:tag"), + requested=2, + timeout_s=0, + allow_partial=True, + ) + == 1 + ) + with pytest.raises(opensandbox_provider.OpenSandboxCreateTimeoutError): + await provider._wait_sdk_pool_idle( + FakePool([0]), + spec=SandboxSpec(image="image:tag"), + requested=2, + timeout_s=0, + allow_partial=False, + ) + + +async def test_exec_file_operations_and_batch_validation(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + class FakeRunCommandOpts: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + class FakeLog: + def __init__(self, text: str) -> None: + self.text = text + + class FakeCommands: + def __init__(self) -> None: + self.calls: list[tuple[str, FakeRunCommandOpts]] = [] + + async def run(self, command: str, *, opts: FakeRunCommandOpts) -> Any: + self.calls.append((command, opts)) + if "fail" in command: + return SimpleNamespace( + logs=SimpleNamespace(stdout=[], stderr=[FakeLog("stderr")]), + error=SimpleNamespace(name="CommandError", value="failed"), + exit_code=None, + ) + return SimpleNamespace( + logs=SimpleNamespace(stdout=[FakeLog("stdout")], stderr=[]), + error=None, + exit_code=None, + ) + + class FakeFiles: + def __init__(self) -> None: + self.writes: list[tuple[str, str | bytes]] = [] + + async def write_file(self, target_path: str, data: str | bytes) -> None: + self.writes.append((target_path, data)) + + async def read_bytes(self, source_path: str) -> bytes: + return f"bytes:{source_path}".encode() + + class FakeRaw: + def __init__(self) -> None: + self.commands = FakeCommands() + self.files = FakeFiles() + + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk", + lambda: (object, object, FakeRunCommandOpts, object, object), + ) + + provider = opensandbox_provider.OpenSandboxProvider( + connection={"request_timeout_s": 5}, + probe={"command": None}, + ) + raw = FakeRaw() + handle = opensandbox_provider.SandboxHandle(sandbox_id="sandbox-1", provider_name="opensandbox", raw=raw) + + result = await provider.exec( + handle, + "echo hello", + cwd="/repo", + env={"A": "B"}, + timeout_s=2, + user=1000, + ) + assert result == opensandbox_provider.SandboxExecResult(stdout="stdout", stderr=None, return_code=0) + command, opts = raw.commands.calls[0] + assert command == "echo hello" + assert opts.kwargs == { + "working_directory": "/repo", + "envs": {"A": "B"}, + "timeout": timedelta(seconds=2), + "uid": 1000, + } + + result = await provider.exec(handle, "fail", user="agent") + assert result.return_code == 1 + assert result.stderr == "stderr\nCommandError: failed" + assert raw.commands.calls[1][0] == "su -s /bin/sh -c fail agent" + + await provider.write_file(handle, "/tmp/file.txt", "contents") + assert await provider.read_file(handle, "/tmp/file.txt") == b"bytes:/tmp/file.txt" + upload_path = tmp_path / "upload.txt" + upload_path.write_text("upload", encoding="utf-8") + await provider.upload_file(handle, upload_path, "/remote/upload.txt") + download_path = tmp_path / "nested" / "download.txt" + await provider.download_file(handle, "/remote/download.txt", download_path) + assert raw.files.writes == [("/tmp/file.txt", "contents"), ("/remote/upload.txt", b"upload")] + assert download_path.read_bytes() == b"bytes:/remote/download.txt" + + with pytest.raises(ValueError, match="count"): + await provider._create_batch_sdk(SandboxSpec(image="image:tag"), 0) + with pytest.raises(ValueError, match="count"): + await provider.create_batch(SandboxSpec(image="image:tag"), 0) + with pytest.raises(ValueError, match="snapshot_id"): + provider._validate_sdk_pool_spec(SandboxSpec(image="image:tag", snapshot_id="snapshot")) + with pytest.raises(ValueError, match="Unsupported"): + await provider.materialize_handle({"kind": "other"}) + + +async def test_provider_create_probe_and_close_error_paths(monkeypatch: pytest.MonkeyPatch) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + create={"connect_poll_s": 0.01}, + probe={ + "command": "probe", + "expected_stdout": "ready", + "timeout_s": 1, + "deadline_s": 0.01, + }, + ) + handle = opensandbox_provider.SandboxHandle(sandbox_id="sandbox-1", provider_name="opensandbox", raw=object()) + + async def bad_probe(*_args: Any, **_kwargs: Any) -> opensandbox_provider.SandboxExecResult: + return opensandbox_provider.SandboxExecResult(stdout="not ready", stderr="bad", return_code=1) + + async def no_sleep(_seconds: float) -> None: + return None + + monkeypatch.setattr(opensandbox_provider.asyncio, "sleep", no_sleep) + monkeypatch.setattr(provider, "_exec", bad_probe) + with pytest.raises(opensandbox_provider.OpenSandboxCreateVerificationError): + await provider._verify_created_handle(handle) + + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": "probe"}) + + async def fail_verify(_handle: Any) -> None: + raise RuntimeError("probe failed") + + monkeypatch.setattr(provider, "_verify_created_handle", fail_verify) + with pytest.raises(opensandbox_provider.OpenSandboxCreateVerificationError): + await provider._verify_created_handles([handle, handle]) + + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + await provider._verify_created_handles([]) + + async def close_raises(_handle: Any, *, delete: bool) -> None: + del delete + raise RuntimeError("close failed") + + monkeypatch.setattr(provider, "close", close_raises) + await provider._cleanup_failed_create_handle(handle) + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + + class DeleteAlreadyGoneRaw: + async def kill(self) -> None: + raise RuntimeError("sandbox sandbox-1 not found") + + async def close(self) -> None: + return None + + await provider.close( + opensandbox_provider.SandboxHandle( + sandbox_id="sandbox-1", + provider_name="opensandbox", + raw=DeleteAlreadyGoneRaw(), + ), + delete=True, + ) + + class DeleteAndCloseFailRaw: + async def kill(self) -> None: + raise RuntimeError("delete failed") + + async def close(self) -> None: + raise RuntimeError("close failed") + + with pytest.raises(RuntimeError, match="Failed to delete and close"): + await provider.close( + opensandbox_provider.SandboxHandle( + sandbox_id="sandbox-2", + provider_name="opensandbox", + raw=DeleteAndCloseFailRaw(), + ), + delete=True, + ) + + +async def test_create_once_and_connect_after_create_error_paths( + fake_opensandbox_sdk: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + connection={"use_server_proxy": False}, + probe={"command": None}, + ) + with pytest.raises(ValueError, match="pooled creation"): + await provider._create_once(SandboxSpec(image="image:tag", extensions={"poolRef": "pool"})) + + provider = opensandbox_provider.OpenSandboxProvider( + create={"timeout_s": 1, "skip_health_check": True}, + probe={"command": None}, + ) + monkeypatch.setattr(opensandbox_provider, "_to_volumes", lambda volumes: volumes) + spec = SandboxSpec( + image="image:tag", + snapshot_id="snapshot-1", + timeout_s=10, + ready_timeout_s=20, + entrypoint=["/bin/sh"], + platform={"os": "linux", "arch": "amd64"}, + volumes=[{"name": "workspace"}], + skip_health_check=False, + ) + handle = await provider._create_once(spec) + assert handle.sandbox_id == "sandbox-1" + assert FakeSandbox.created_kwargs["snapshot_id"] == "snapshot-1" + assert FakeSandbox.created_kwargs["timeout"] == timedelta(seconds=10) + assert FakeSandbox.created_kwargs["ready_timeout"] == timedelta(seconds=20) + assert FakeSandbox.created_kwargs["entrypoint"] == ["/bin/sh"] + assert FakeSandbox.created_kwargs["platform"] == FakePlatformSpec(os="linux", arch="amd64") + assert FakeSandbox.created_kwargs["volumes"] == [{"name": "workspace"}] + assert FakeSandbox.created_kwargs["skip_health_check"] is True + + class FailingConnectSandbox(FakeSandbox): + @classmethod + async def connect(cls, *args: Any, **kwargs: Any) -> "FakeSandbox": + del args, kwargs + raise ConnectionError("pod may still be starting") + + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk", + lambda: (FailingConnectSandbox, FakeConnectionConfig, object, FakePlatformSpec, object), + ) + provider = opensandbox_provider.OpenSandboxProvider( + create={"connect_attempt_timeout_s": 0.01, "connect_poll_s": 0.01}, + probe={"command": None}, + ) + + async def no_sleep(_seconds: float) -> None: + return None + + monkeypatch.setattr(opensandbox_provider.asyncio, "sleep", no_sleep) + with pytest.raises(opensandbox_provider.OpenSandboxCreateTimeoutError): + await provider._connect_after_create( + opensandbox_provider.SandboxHandle(sandbox_id="sandbox-1", provider_name="opensandbox", raw=None), + SandboxSpec(image="image:tag"), + ) + + +async def test_retry_classification_and_await_sdk_helpers(monkeypatch: pytest.MonkeyPatch) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + operations={"retries": 0}, + probe={"command": None}, + ) + assert await provider.aclose() is None + assert await provider._await_sdk_call(_return_value("ok"), operation="op", sandbox_id="sandbox-1", timeout_s=None) + assert opensandbox_provider._is_retryable_sdk_operation_error(TimeoutError("command timeout")) is False + assert opensandbox_provider._is_retryable_sdk_operation_error(ConnectionError("proxy failed")) is True + wrapped = RuntimeError("wrapper") + wrapped.__cause__ = ConnectionError("connection reset") + assert opensandbox_provider._is_retryable_sdk_operation_error(wrapped) is True + + class FakeHttpxConnectError(Exception): + pass + + monkeypatch.setattr(opensandbox_provider, "_httpx_retryable_types", lambda: (FakeHttpxConnectError,)) + assert opensandbox_provider._is_retryable_create_error(FakeHttpxConnectError("temporary")) is True + assert opensandbox_provider._is_retryable_sdk_operation_error(FakeHttpxConnectError("temporary")) is True + + async def cancelled() -> None: + raise asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await provider._await_sdk_operation( + cancelled, + operation="cancelled", + sandbox_id="sandbox-1", + timeout_s=None, + ) + + +async def _return_value(value: Any) -> Any: + return value diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py new file mode 100644 index 0000000000..908ac7a43a --- /dev/null +++ b/tests/unit_tests/test_sandbox.py @@ -0,0 +1,766 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 importlib.util +from pathlib import Path +from typing import Any +from uuid import uuid4 + +import pytest + +from nemo_gym.sandbox import ( + AsyncSandbox, + Sandbox, + SandboxCreateError, + SandboxExecResult, + SandboxHandle, + SandboxSpec, + get_provider_class, + list_providers, + register_provider, + rewrite_image, +) +from nemo_gym.sandbox.providers.opensandbox import provider as opensandbox_provider_module +from nemo_gym.sandbox.providers.opensandbox.provider import ( + IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY, + IMAGE_PULL_POLICY_EXTENSION_KEY, + OpenSandboxCreateVerificationError, + OpenSandboxProvider, +) +from responses_api_agents.mini_swe_agent_2.sandbox_environment import MiniSWESandboxEnvironment + + +def _has_module(module_name: str) -> bool: + try: + return importlib.util.find_spec(module_name) is not None + except ModuleNotFoundError: + return False + + +requires_tenacity = pytest.mark.skipif( + not _has_module("tenacity"), + reason="tenacity optional sandbox dependency is not installed", +) + + +class FakeSandboxProvider: + name = "fake" + last_instance: "FakeSandboxProvider | None" = None + + def __init__(self, marker: str = "default") -> None: + self.marker = marker + self.created_specs: list[SandboxSpec] = [] + self.exec_calls: list[dict[str, Any]] = [] + self.write_calls: list[tuple[SandboxHandle, str, str | bytes]] = [] + self.read_calls: list[tuple[SandboxHandle, str]] = [] + self.upload_calls: list[tuple[SandboxHandle, Path, str]] = [] + self.download_calls: list[tuple[SandboxHandle, str, Path]] = [] + self.closed: list[tuple[SandboxHandle, bool]] = [] + self.aclosed = False + FakeSandboxProvider.last_instance = self + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + self.created_specs.append(spec) + return SandboxHandle(sandbox_id="fake-1", provider_name=self.name, raw={"spec": spec}) + + async def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + del allow_partial + return [await self.create(spec) for _ in range(count)] + + async def connect(self, sandbox_id: str) -> SandboxHandle: + return SandboxHandle(sandbox_id=sandbox_id, provider_name=self.name, raw={}) + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + self.exec_calls.append( + { + "handle": handle, + "command": command, + "cwd": cwd, + "env": env, + "timeout_s": timeout_s, + "user": user, + } + ) + return SandboxExecResult(stdout="ok", stderr=None, return_code=0) + + async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + self.write_calls.append((handle, target_path, data)) + + async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + self.read_calls.append((handle, source_path)) + return f"read:{source_path}".encode() + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + self.upload_calls.append((handle, source_path, target_path)) + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + self.download_calls.append((handle, source_path, target_path)) + target_path.parent.mkdir(parents=True, exist_ok=True) + target_path.write_bytes(b"downloaded") + + async def close(self, handle: SandboxHandle, *, delete: bool) -> None: + self.closed.append((handle, delete)) + + async def aclose(self) -> None: + self.aclosed = True + + def handle_reference(self, handle: SandboxHandle) -> dict[str, str]: + return {"kind": "fake", "sandbox_id": handle.sandbox_id} + + async def materialize_handle(self, value: Any) -> SandboxHandle: + return SandboxHandle(sandbox_id=value["sandbox_id"], provider_name=self.name, raw={"materialized": True}) + + +def test_sandbox_facade_uses_public_provider_api() -> None: + asyncio.run(_assert_sandbox_facade_uses_public_provider_api()) + + +async def _assert_sandbox_facade_uses_public_provider_api() -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + + sandbox = AsyncSandbox({provider_name: {"marker": "configured"}}) + handle = await sandbox.create(SandboxSpec(image="image:tag", metadata={"suite": "unit"})) + + provider = FakeSandboxProvider.last_instance + assert provider is not None + assert provider.marker == "configured" + assert provider.created_specs[0].image == "image:tag" + assert provider.created_specs[0].metadata == {"suite": "unit"} + + result = await sandbox.exec(handle, "pytest -q", cwd="/repo", timeout_s=60, user="agent") + assert result == SandboxExecResult(stdout="ok", stderr=None, return_code=0) + assert provider.exec_calls[0] == { + "handle": handle, + "command": "pytest -q", + "cwd": "/repo", + "env": None, + "timeout_s": 60, + "user": "agent", + } + + await sandbox.delete(handle) + assert provider.closed[0] == (handle, True) + assert sandbox.handle_reference(handle) == {"kind": "fake", "sandbox_id": "fake-1"} + assert await sandbox.materialize_handle({"sandbox_id": "fake-2"}) == SandboxHandle( + sandbox_id="fake-2", provider_name="fake", raw={"materialized": True} + ) + async with AsyncSandbox(provider) as context_sandbox: + assert context_sandbox.provider_name == "fake" + await sandbox.shutdown() + assert provider.aclosed is True + + +def test_rewrite_image_and_materialize_handle_validation() -> None: + asyncio.run(_assert_rewrite_image_and_materialize_handle_validation()) + + +def test_provider_registry_validation_and_listing() -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + + assert get_provider_class(provider_name) is FakeSandboxProvider + assert provider_name in list_providers() + with pytest.raises(ValueError, match="must be non-empty"): + register_provider("", FakeSandboxProvider) + with pytest.raises(ValueError, match="already registered"): + register_provider(provider_name, FakeSandboxProvider) + with pytest.raises(ValueError, match="Unknown sandbox provider"): + get_provider_class(f"missing-{uuid4().hex}") + + +async def _assert_rewrite_image_and_materialize_handle_validation() -> None: + assert rewrite_image(None, []) is None + assert rewrite_image("image:tag", [{"from": "other/", "to": "mirror/"}]) == "image:tag" + + class BadMaterializeProvider(FakeSandboxProvider): + async def materialize_handle(self, value: Any) -> object: + del value + return object() + + sandbox = AsyncSandbox(BadMaterializeProvider()) + try: + await sandbox.materialize_handle({"sandbox_id": "bad"}) + except TypeError as e: + assert "must return SandboxHandle" in str(e) + else: + raise AssertionError("expected invalid materialize_handle return type to fail") + + +def test_async_sandbox_batch_file_and_fallback_reference_operations(tmp_path: Path) -> None: + asyncio.run(_assert_async_sandbox_batch_file_and_fallback_reference_operations(tmp_path)) + + +async def _assert_async_sandbox_batch_file_and_fallback_reference_operations(tmp_path: Path) -> None: + provider = FakeSandboxProvider() + sandbox = AsyncSandbox(provider) + + handles = await sandbox.create_batch(SandboxSpec(image="image:tag"), 2, allow_partial=True) + connected = await sandbox.connect("connected-1") + await sandbox.write_file(connected, "/tmp/file.txt", "contents") + assert await sandbox.read_file(connected, "/tmp/file.txt") == b"read:/tmp/file.txt" + source_path = tmp_path / "source.txt" + target_path = tmp_path / "nested" / "target.txt" + source_path.write_text("local", encoding="utf-8") + await sandbox.upload_file(connected, source_path, "/remote/source.txt") + await sandbox.download_file(connected, "/remote/source.txt", target_path) + await sandbox.close(connected) + + assert [handle.sandbox_id for handle in handles] == ["fake-1", "fake-1"] + assert provider.write_calls == [(connected, "/tmp/file.txt", "contents")] + assert provider.read_calls == [(connected, "/tmp/file.txt")] + assert provider.upload_calls == [(connected, source_path, "/remote/source.txt")] + assert provider.download_calls == [(connected, "/remote/source.txt", target_path)] + assert target_path.read_bytes() == b"downloaded" + + plain_provider = FakeSandboxProvider() + plain_provider.handle_reference = None # type: ignore[method-assign] + plain_provider.materialize_handle = None # type: ignore[method-assign] + plain_sandbox = AsyncSandbox(plain_provider) + plain_handle = SandboxHandle(sandbox_id="plain-1", provider_name="fake", raw={}) + assert plain_sandbox.handle_reference(plain_handle) is plain_handle + assert await plain_sandbox.materialize_handle(plain_handle) is plain_handle + try: + await plain_sandbox.materialize_handle({"sandbox_id": "plain-2"}) + except ValueError as e: + assert "cannot materialize" in str(e) + else: + raise AssertionError("expected materialize_handle without provider support to fail") + + +def test_sync_sandbox_facade_uses_public_provider_api() -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + + with Sandbox({provider_name: {"marker": "configured"}}) as sandbox: + handle = sandbox.create(SandboxSpec(image="image:tag", metadata={"suite": "unit"})) + + provider = FakeSandboxProvider.last_instance + assert provider is not None + assert provider.marker == "configured" + assert provider.created_specs[0].image == "image:tag" + assert provider.created_specs[0].metadata == {"suite": "unit"} + + result = sandbox.exec(handle, "pytest -q", cwd="/repo", timeout_s=60, user="agent") + assert result == SandboxExecResult(stdout="ok", stderr=None, return_code=0) + assert provider.exec_calls[0] == { + "handle": handle, + "command": "pytest -q", + "cwd": "/repo", + "env": None, + "timeout_s": 60, + "user": "agent", + } + + sandbox.delete(handle) + assert provider.closed[0] == (handle, True) + assert sandbox.handle_reference(handle) == {"kind": "fake", "sandbox_id": "fake-1"} + assert sandbox.materialize_handle({"sandbox_id": "fake-3"}).sandbox_id == "fake-3" + assert sandbox.provider_name == "fake" + assert len(sandbox.create_batch(SandboxSpec(image="image:tag"), 2)) == 2 + sandbox.shutdown() + sandbox.shutdown() + assert provider.aclosed is True + try: + sandbox.provider_name + except RuntimeError as e: + assert "sync loop is closed" in str(e) + else: + raise AssertionError("expected closed sync sandbox to reject further calls") + + +def test_sync_sandbox_file_operations(tmp_path: Path) -> None: + provider = FakeSandboxProvider() + with Sandbox(provider) as sandbox: + handle = sandbox.connect("sync-1") + sandbox.write_file(handle, "/tmp/file.txt", b"contents") + assert sandbox.read_file(handle, "/tmp/file.txt") == b"read:/tmp/file.txt" + source_path = tmp_path / "source.txt" + target_path = tmp_path / "target.txt" + source_path.write_text("local", encoding="utf-8") + sandbox.upload_file(handle, source_path, "/remote/source.txt") + sandbox.download_file(handle, "/remote/source.txt", target_path) + + assert provider.write_calls == [(handle, "/tmp/file.txt", b"contents")] + assert provider.read_calls == [(handle, "/tmp/file.txt")] + assert provider.upload_calls == [(handle, source_path, "/remote/source.txt")] + assert provider.download_calls == [(handle, "/remote/source.txt", target_path)] + assert target_path.read_bytes() == b"downloaded" + + +def test_sync_sandbox_facade_rejects_async_context() -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + + async def _create_sync_sandbox_in_async_context() -> None: + Sandbox({provider_name: {}}) + + try: + asyncio.run(_create_sync_sandbox_in_async_context()) + except RuntimeError as e: + assert "use AsyncSandbox in async code" in str(e) + else: + raise AssertionError("expected sync Sandbox to reject async context") + + +@requires_tenacity +def test_opensandbox_sdk_create_receives_default_image_pull_policy(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_sdk_create_receives_default_image_pull_policy(monkeypatch)) + + +async def _assert_opensandbox_sdk_create_receives_default_image_pull_policy(monkeypatch) -> None: + class FakeSDKSandbox: + create_calls: list[dict[str, Any]] = [] + + def __init__(self, sandbox_id: str) -> None: + self.id = sandbox_id + + @classmethod + async def create(cls, **kwargs: Any) -> "FakeSDKSandbox": + cls.create_calls.append(kwargs) + return cls("sdk-sandbox-1") + + monkeypatch.setattr( + opensandbox_provider_module, + "_require_opensandbox_sdk", + lambda: (FakeSDKSandbox, object, object, object, object), + ) + + provider = OpenSandboxProvider(probe={"command": None}) + monkeypatch.setattr(provider, "_connection_config", lambda request_timeout_s=None, use_server_proxy=None: object()) + + handle = await provider.create( + SandboxSpec( + image="image:tag", + metadata={ + "harbor_instance_id": "swebench::django__django-10880", + "long": f"bad:{'x' * 80}:", + }, + ) + ) + + assert handle.sandbox_id == "sdk-sandbox-1" + metadata = FakeSDKSandbox.create_calls[0]["metadata"] + assert metadata["harbor_instance_id"] == "swebench_django__django-10880" + assert metadata["long"] == ("bad_" + "x" * 59) + extensions = FakeSDKSandbox.create_calls[0]["extensions"] + assert extensions[IMAGE_PULL_POLICY_EXTENSION_KEY] == "IfNotPresent" + assert extensions[IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY] == "IfNotPresent" + + +@requires_tenacity +def test_opensandbox_connect_after_create_can_use_direct_exec_endpoint(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_connect_after_create_can_use_direct_exec_endpoint(monkeypatch)) + + +async def _assert_opensandbox_connect_after_create_can_use_direct_exec_endpoint(monkeypatch) -> None: + class FakeConnectionConfig: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + class FakeSDKSandbox: + connect_calls: list[dict[str, Any]] = [] + + def __init__(self, sandbox_id: str) -> None: + self.id = sandbox_id + + @classmethod + async def connect(cls, sandbox_id: str, **kwargs: Any) -> "FakeSDKSandbox": + cls.connect_calls.append({"sandbox_id": sandbox_id, **kwargs}) + return cls(sandbox_id) + + monkeypatch.setattr( + opensandbox_provider_module, + "_require_opensandbox_sdk", + lambda: (FakeSDKSandbox, FakeConnectionConfig, object, object, object), + ) + + provider = OpenSandboxProvider( + connection={"use_server_proxy": True, "exec_use_server_proxy": False}, + create={"connect_attempt_timeout_s": 1}, + probe={"command": None}, + ) + handle = await provider._connect_after_create( + SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=None), + SandboxSpec(image="image:tag", ready_timeout_s=10), + ) + + assert handle.sandbox_id == "sdk-sandbox-1" + 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["use_server_proxy"] is False + + +@requires_tenacity +def test_opensandbox_create_probe_can_require_stable_successes(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_create_probe_can_require_stable_successes(monkeypatch)) + + +async def _assert_opensandbox_create_probe_can_require_stable_successes(monkeypatch) -> None: + provider = OpenSandboxProvider( + probe={ + "command": "true", + "expected_stdout": None, + "stable_count": 3, + "stable_delay_s": 0, + }, + ) + calls: list[dict[str, Any]] = [] + + async def fake_exec( + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + calls.append( + { + "handle": handle, + "command": command, + "cwd": cwd, + "env": env, + "timeout_s": timeout_s, + "user": user, + } + ) + return SandboxExecResult(stdout="", stderr="", return_code=0) + + monkeypatch.setattr(provider, "_exec", fake_exec) + handle = SandboxHandle(sandbox_id="sdk-sandbox-0", provider_name="opensandbox", raw=object()) + + await provider._verify_created_handle(handle) + + assert [call["command"] for call in calls] == ["true", "true", "true"] + assert all(call["timeout_s"] == 30 for call in calls) + assert all(call["user"] == "root" for call in calls) + + +@requires_tenacity +def test_opensandbox_create_probe_polls_same_sandbox_after_transient_errors(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_create_probe_polls_same_sandbox_after_transient_errors(monkeypatch)) + + +async def _assert_opensandbox_create_probe_polls_same_sandbox_after_transient_errors(monkeypatch) -> None: + provider = OpenSandboxProvider( + create={"connect_poll_s": 0.01}, + probe={ + "command": "true", + "expected_stdout": None, + "timeout_s": 1, + "deadline_s": 2, + "stable_count": 2, + "stable_delay_s": 0, + }, + ) + attempts = 0 + handles: list[SandboxHandle] = [] + + async def fake_exec( + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + del command, cwd, env, timeout_s, user + nonlocal attempts + attempts += 1 + handles.append(handle) + if attempts <= 2: + raise ConnectionError("direct execd endpoint is not accepting connections yet") + return SandboxExecResult(stdout="", stderr="", return_code=0) + + monkeypatch.setattr(provider, "_exec", fake_exec) + handle = SandboxHandle(sandbox_id="sdk-sandbox-0", provider_name="opensandbox", raw=object()) + + await provider._verify_created_handle(handle) + + assert attempts == 4 + assert {seen_handle.sandbox_id for seen_handle in handles} == {"sdk-sandbox-0"} + + +def test_opensandbox_create_probe_failures_are_retryable() -> None: + error = OpenSandboxCreateVerificationError("pod sdk-sandbox-0 failed create probe") + + assert isinstance(error, SandboxCreateError) + assert opensandbox_provider_module._is_retryable_create_error(error) is True + + +def test_opensandbox_starting_pod_endpoint_errors_are_retryable() -> None: + error = RuntimeError( + "Get endpoint for sandbox sdk-sandbox-0 port 44772 failed: " + "Pod IP is not yet available. The Pod may still be starting." + ) + + assert opensandbox_provider_module._is_retryable_create_error(error) is True + + +@requires_tenacity +def test_opensandbox_exec_retries_retryable_sdk_failures(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_exec_retries_retryable_sdk_failures(monkeypatch)) + + +async def _assert_opensandbox_exec_retries_retryable_sdk_failures(monkeypatch) -> None: + class FakeRunCommandOpts: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + class FakeLog: + def __init__(self, text: str) -> None: + self.text = text + + class FakeLogs: + stdout = [FakeLog("ok")] + stderr: list[FakeLog] = [] + + class FakeExecution: + logs = FakeLogs() + error = None + exit_code = 0 + + class FakeCommands: + def __init__(self) -> None: + self.calls = 0 + + async def run(self, command: str, *, opts: FakeRunCommandOpts) -> FakeExecution: + del command, opts + self.calls += 1 + if self.calls <= 2: + raise ConnectionError("transient proxy failure") + return FakeExecution() + + class FakeRaw: + def __init__(self) -> None: + self.commands = FakeCommands() + + monkeypatch.setattr( + opensandbox_provider_module, + "_require_opensandbox_sdk", + lambda: (object, object, FakeRunCommandOpts, object, object), + ) + + provider = OpenSandboxProvider( + operations={ + "retries": 2, + "retry_delay_s": 0, + "retry_max_delay_s": 0, + "command_retries": 2, + }, + probe={"command": None}, + ) + raw = FakeRaw() + handle = SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=raw) + + result = await provider.exec(handle, "echo hello", timeout_s=30) + + assert result.stdout == "ok" + assert result.return_code == 0 + assert raw.commands.calls == 3 + + +@requires_tenacity +def test_opensandbox_command_retries_can_be_disabled(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_command_retries_can_be_disabled(monkeypatch)) + + +async def _assert_opensandbox_command_retries_can_be_disabled(monkeypatch) -> None: + class FakeRunCommandOpts: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + class FakeCommands: + def __init__(self) -> None: + self.calls = 0 + + async def run(self, command: str, *, opts: FakeRunCommandOpts) -> None: + del command, opts + self.calls += 1 + raise ConnectionError("transient proxy failure") + + class FakeRaw: + def __init__(self) -> None: + self.commands = FakeCommands() + + monkeypatch.setattr( + opensandbox_provider_module, + "_require_opensandbox_sdk", + lambda: (object, object, FakeRunCommandOpts, object, object), + ) + + provider = OpenSandboxProvider( + operations={ + "retries": 2, + "retry_delay_s": 0, + "retry_max_delay_s": 0, + "command_retries": 0, + }, + probe={"command": None}, + ) + raw = FakeRaw() + handle = SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=raw) + + try: + await provider.exec(handle, "echo hello", timeout_s=30) + except ConnectionError: + pass + else: + raise AssertionError("expected provider.exec to propagate the command failure") + + assert raw.commands.calls == 1 + + +@requires_tenacity +def test_opensandbox_close_timeout_does_not_fail_after_delete() -> None: + asyncio.run(_assert_opensandbox_close_timeout_does_not_fail_after_delete()) + + +async def _assert_opensandbox_close_timeout_does_not_fail_after_delete() -> None: + class SlowCloseRaw: + def __init__(self) -> None: + self.killed = False + + async def kill(self) -> None: + self.killed = True + + async def close(self) -> None: + await asyncio.sleep(60) + + raw = SlowCloseRaw() + provider = OpenSandboxProvider( + operations={"close_timeout_s": 0.01}, + probe={"command": None}, + ) + handle = SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=raw) + + await provider.close(handle, delete=True) + + assert raw.killed is True + + +@requires_tenacity +def test_opensandbox_close_timeout_still_fails_without_delete() -> None: + asyncio.run(_assert_opensandbox_close_timeout_still_fails_without_delete()) + + +async def _assert_opensandbox_close_timeout_still_fails_without_delete() -> None: + class SlowCloseRaw: + async def close(self) -> None: + await asyncio.sleep(60) + + provider = OpenSandboxProvider( + operations={"close_timeout_s": 0.01}, + probe={"command": None}, + ) + handle = SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=SlowCloseRaw()) + + try: + await provider.close(handle, delete=False) + except TimeoutError: + pass + else: + raise AssertionError("expected close timeout to fail when delete=False") + + +def test_mini_swe_sandbox_environment_owns_conda_setup(monkeypatch) -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + monkeypatch.setenv("FORWARDED_KEY", "forwarded-value") + + env = MiniSWESandboxEnvironment( + image="upstream/image:tag", + cwd="/testbed", + provider={provider_name: {"marker": "configured"}}, + spec={ + "image_rewrites": [{"from": "upstream/", "to": "mirror/"}], + "metadata": {"suite": "unit"}, + "resources": {"cpu": "1"}, + }, + env={"STATIC_KEY": "static-value"}, + forward_env=["FORWARDED_KEY"], + conda_env="testbed", + activate_conda=True, + user="agent", + delete=True, + ) + + try: + assert env.get_template_vars(extra="value")["extra"] == "value" + serialized = env.serialize() + assert serialized["info"]["config"]["environment_type"].endswith("MiniSWESandboxEnvironment") + env.config.activate_conda = False + assert env._command("echo plain", "/tmp/work") == "echo plain" + env.config.activate_conda = True + + provider = FakeSandboxProvider.last_instance + assert provider is not None + assert provider.marker == "configured" + assert provider.created_specs[0].image == "mirror/image:tag" + assert provider.created_specs[0].env == { + "FORWARDED_KEY": "forwarded-value", + "STATIC_KEY": "static-value", + } + + result = env.execute("pytest -q", is_eval=True) + assert result == {"output": "ok", "returncode": 0, "exception_info": ""} + exec_call = provider.exec_calls[0] + assert exec_call["cwd"] == "/" + assert exec_call["timeout_s"] == 1800 + assert exec_call["user"] == "agent" + assert "conda activate testbed" in exec_call["command"] + assert exec_call["command"].endswith("pytest -q") + finally: + env.cleanup() + env.cleanup() + + assert FakeSandboxProvider.last_instance is not None + assert FakeSandboxProvider.last_instance.closed[0][1] is True + + +def test_mini_swe_sandbox_environment_validation_and_context_manager() -> None: + with pytest.raises(ValueError, match="requires provider"): + MiniSWESandboxEnvironment(image="image:tag") + + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + with MiniSWESandboxEnvironment( + image="image:tag", + provider={provider_name: {}}, + delete=False, + ) as env: + assert env._handle is not None + + assert FakeSandboxProvider.last_instance is not None + assert FakeSandboxProvider.last_instance.closed[-1][1] is False diff --git a/uv.lock b/uv.lock index c34d4a2c0c..5436871df9 100644 --- a/uv.lock +++ b/uv.lock @@ -1414,6 +1414,10 @@ dev = [ { name = "requests-mock" }, { name = "ruff" }, ] +sandbox = [ + { name = "opensandbox" }, + { name = "tenacity" }, +] [package.dev-dependencies] docs = [ @@ -1446,6 +1450,7 @@ requires-dist = [ { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.8.0" }, { name = "omegaconf" }, { name = "openai", specifier = "<=2.7.2" }, + { name = "opensandbox", marker = "extra == 'sandbox'", specifier = ">=0.1.9" }, { name = "orjson" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.6.0" }, { name = "psutil" }, @@ -1461,6 +1466,7 @@ requires-dist = [ { name = "requests-mock", marker = "extra == 'dev'" }, { name = "rich" }, { name = "ruff", marker = "extra == 'dev'" }, + { name = "tenacity", marker = "extra == 'sandbox'", specifier = ">=9.1.4" }, { name = "tqdm" }, { name = "urllib3", specifier = ">=2.7.0" }, { name = "uvicorn" }, @@ -1468,7 +1474,7 @@ requires-dist = [ { name = "wandb" }, { name = "yappi" }, ] -provides-extras = ["dev"] +provides-extras = ["sandbox", "dev"] [package.metadata.requires-dev] docs = [ @@ -1623,6 +1629,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/68/162c97ea78c957d68ecf78a5c5041d2e25bd5562bdf5d89a6cbf7f8429bf/opencensus_context-0.1.3-py2.py3-none-any.whl", hash = "sha256:073bb0590007af276853009fac7e4bab1d523c3f03baf4cb4511ca38967c6039", size = 5060, upload-time = "2022-08-03T22:20:20.352Z" }, ] +[[package]] +name = "opensandbox" +version = "0.1.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "httpx" }, + { 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" } +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" }, +] + [[package]] name = "opentelemetry-api" version = "1.36.0" @@ -2777,6 +2798,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/51/f0/1098f6628bbe04b086ce59692d09b116ec751286eb7d33e88c5bf0c2e210/swagger_plugin_for_sphinx-6.0.0-py3-none-any.whl", hash = "sha256:35dc646d759a44ce78aefde2fe34f54e7b8c3439d0a52541a6a8b9924a711832", size = 11253, upload-time = "2025-10-16T06:26:08.504Z" }, ] +[[package]] +name = "tenacity" +version = "9.1.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/47/c6/ee486fd809e357697ee8a44d3d69222b344920433d3b6666ccd9b374630c/tenacity-9.1.4.tar.gz", hash = "sha256:adb31d4c263f2bd041081ab33b498309a57c77f9acf2db65aadf0898179cf93a", size = 49413, upload-time = "2026-02-07T10:45:33.841Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl", hash = "sha256:6095a360c919085f28c6527de529e76a06ad89b23659fa881ae0649b867a9d55", size = 28926, upload-time = "2026-02-07T10:45:32.24Z" }, +] + [[package]] name = "tqdm" version = "4.67.1" From 9b1dc1c78b228092923418ddae50c0d72ef96cf7 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Wed, 20 May 2026 13:30:56 -0700 Subject: [PATCH 02/14] Fix sandbox unit test coverage Signed-off-by: Hemil Desai --- .github/workflows/unit-tests.yml | 2 +- nemo_gym/sandbox/api.py | 2 +- .../sandbox/providers/opensandbox/provider.py | 4 +- nemo_gym/sandbox/providers/registry.py | 20 +- .../mini_swe_agent_2/sandbox_environment.py | 2 +- tests/unit_tests/test_opensandbox_provider.py | 251 +++++++++++++++++- tests/unit_tests/test_sandbox.py | 125 ++++++++- 7 files changed, 380 insertions(+), 26 deletions(-) diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index 0056f3ed27..810f75fcc8 100644 --- a/.github/workflows/unit-tests.yml +++ b/.github/workflows/unit-tests.yml @@ -157,7 +157,7 @@ jobs: curl -LsSf "$UV_INSTALL_URL" | sh uv venv --python 3.12 source .venv/bin/activate - uv sync --extra dev + uv sync --extra dev --extra sandbox - name: Test run: | diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index 39444814a3..1348c00273 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -297,7 +297,7 @@ def __enter__(self) -> "Sandbox": def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: self.shutdown() - def __del__(self) -> None: + def __del__(self) -> None: # pragma: no cover if hasattr(self, "_closed") and not self._closed: try: self.shutdown() diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 359deb12a7..07f547eefd 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -834,6 +834,8 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: Sandbox.create(**kwargs), timeout=timeout_s, ) + if sandbox is None: + raise RuntimeError("OpenSandbox SDK create returned no sandbox handle") sandbox_id = str(sandbox.id) except TimeoutError as e: error = OpenSandboxCreateTimeoutError( @@ -843,7 +845,7 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: f"ready_timeout_s={spec.ready_timeout_s!r}" ) raise error from e - if sandbox is None or sandbox_id is None: + if sandbox_id is None: raise RuntimeError("OpenSandbox SDK create returned no sandbox handle") created_handle = SandboxHandle( sandbox_id=sandbox_id, diff --git a/nemo_gym/sandbox/providers/registry.py b/nemo_gym/sandbox/providers/registry.py index 8aecd6c471..efa36c5433 100644 --- a/nemo_gym/sandbox/providers/registry.py +++ b/nemo_gym/sandbox/providers/registry.py @@ -14,22 +14,24 @@ """Provider registration utilities.""" -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Any, TypeAlias from nemo_gym.sandbox.providers.base import SandboxProvider ProviderClass: TypeAlias = type[SandboxProvider] +ProviderLoader: TypeAlias = Callable[[], ProviderClass] _PROVIDER_REGISTRY: dict[str, ProviderClass] = {} +_BUILTIN_PROVIDER_LOADERS: dict[str, ProviderLoader] = {} def register_provider(name: str, provider_class: ProviderClass) -> None: """Register a sandbox provider class.""" if not name: raise ValueError("Provider name must be non-empty") - if name in _PROVIDER_REGISTRY: + if name in _PROVIDER_REGISTRY or name in _BUILTIN_PROVIDER_LOADERS: raise ValueError(f"Sandbox provider {name!r} is already registered") _PROVIDER_REGISTRY[name] = provider_class @@ -39,7 +41,10 @@ def get_provider_class(name: str) -> ProviderClass: try: return _PROVIDER_REGISTRY[name] except KeyError as e: - available = ", ".join(sorted(_PROVIDER_REGISTRY)) or "" + loader = _BUILTIN_PROVIDER_LOADERS.get(name) + if loader is not None: + return loader() + available = ", ".join(list_providers()) or "" raise ValueError(f"Unknown sandbox provider {name!r}. Available providers: {available}") from e @@ -61,14 +66,13 @@ def create_provider(config: Mapping[str, Any]) -> SandboxProvider: def list_providers() -> list[str]: """List registered provider names.""" - return sorted(_PROVIDER_REGISTRY) + return sorted({*_PROVIDER_REGISTRY, *_BUILTIN_PROVIDER_LOADERS}) -def _register_builtins() -> None: +def _load_opensandbox_provider() -> ProviderClass: from nemo_gym.sandbox.providers.opensandbox import OpenSandboxProvider - if "opensandbox" not in _PROVIDER_REGISTRY: - register_provider("opensandbox", OpenSandboxProvider) + return OpenSandboxProvider -_register_builtins() +_BUILTIN_PROVIDER_LOADERS["opensandbox"] = _load_opensandbox_provider diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index 34c60f21ab..8a8b8d845b 100644 --- a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -191,7 +191,7 @@ def __enter__(self) -> "MiniSWESandboxEnvironment": def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: self.cleanup() - def __del__(self) -> None: + def __del__(self) -> None: # pragma: no cover if hasattr(self, "_closed") and not self._closed: try: self.cleanup() diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index 5d5a9756c4..97e525d4d3 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -14,7 +14,6 @@ # limitations under the License. import asyncio -import importlib.util from dataclasses import dataclass from datetime import timedelta from pathlib import Path @@ -24,13 +23,11 @@ import pytest from nemo_gym.sandbox.providers.base import SandboxSpec -from nemo_gym.sandbox.providers.opensandbox import provider as opensandbox_provider -pytestmark = pytest.mark.skipif( - importlib.util.find_spec("tenacity") is None, - reason="tenacity optional sandbox dependency is not installed", -) +pytest.importorskip("tenacity", reason="tenacity optional sandbox dependency is not installed") + +from nemo_gym.sandbox.providers.opensandbox import provider as opensandbox_provider @dataclass(frozen=True) @@ -44,6 +41,11 @@ def __init__(self, **kwargs: Any) -> None: self.kwargs = kwargs +@dataclass(frozen=True) +class FakeVolume: + name: str + + class FakeSandbox: created_kwargs: dict[str, Any] = {} connected_args: tuple[Any, ...] = () @@ -146,6 +148,79 @@ def require_sdk_pool() -> tuple[Any, Any, Any, Any]: ) +def test_sdk_import_helpers_and_retry_classification() -> None: + assert len(opensandbox_provider._require_opensandbox_sdk()) == 5 + assert len(opensandbox_provider._require_opensandbox_sdk_pool()) == 4 + assert len(opensandbox_provider._require_tenacity()) == 4 + + class StatusCodeError(Exception): + status_code = 429 + + assert opensandbox_provider._exception_status_code(StatusCodeError("rate limited")) == 429 + assert opensandbox_provider._is_retryable_create_error( + opensandbox_provider.OpenSandboxCreateError("create failed") + ) + + from opensandbox.exceptions import ( # noqa: PLC0415 + InvalidArgumentException, + SandboxApiException, + SandboxException, + SandboxInternalException, + ) + + assert opensandbox_provider._is_retryable_create_error(InvalidArgumentException("bad input")) is False + assert opensandbox_provider._is_retryable_create_error(SandboxInternalException("server failed")) is True + + retryable_api_error = SandboxApiException("busy") + retryable_api_error.status_code = 503 + assert opensandbox_provider._is_retryable_create_error(retryable_api_error) is True + + nonretryable_api_error = SandboxApiException("not found") + nonretryable_api_error.status_code = 404 + assert opensandbox_provider._is_retryable_create_error(nonretryable_api_error) is False + assert opensandbox_provider._is_retryable_create_error(SandboxException("gateway timeout")) is True + + retry_state = SimpleNamespace( + outcome=SimpleNamespace(exception=lambda: RuntimeError("temporary")), + next_action=SimpleNamespace(sleep=0.5), + attempt_number=2, + ) + opensandbox_provider._log_create_retry(retry_state) + + +async def test_provider_reference_materialization_and_conversion_helpers( + fake_opensandbox_sdk: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + connection_config = opensandbox_provider.OpenSandboxConnectionConfig(domain="sandbox.example") + assert ( + opensandbox_provider._coerce_config(connection_config, opensandbox_provider.OpenSandboxConnectionConfig) + is connection_config + ) + + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk", + lambda: (object, object, object, FakePlatformSpec, FakeVolume), + ) + assert opensandbox_provider._to_volumes([{"name": "workspace"}]) == [FakeVolume(name="workspace")] + + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + handle = opensandbox_provider.SandboxHandle(sandbox_id="sandbox-1", provider_name="opensandbox", raw=object()) + assert provider.handle_reference(handle) == { + "kind": "sandbox_id", + "provider": "opensandbox", + "sandbox_id": "sandbox-1", + } + + async def connect(sandbox_id: str) -> opensandbox_provider.SandboxHandle: + return opensandbox_provider.SandboxHandle(sandbox_id=sandbox_id, provider_name="opensandbox", raw="connected") + + monkeypatch.setattr(provider, "connect", connect) + materialized = await provider.materialize_handle({"kind": "sandbox_id", "sandbox_id": "sandbox-2"}) + assert materialized.raw == "connected" + + async def test_sdk_pool_passes_platform_through_pool_creation_spec( fake_opensandbox_sdk: None, ) -> None: @@ -504,6 +579,23 @@ async def close(self) -> None: delete=True, ) + class DeleteFailsCloseSucceedsRaw: + async def kill(self) -> None: + raise RuntimeError("delete failed") + + async def close(self) -> None: + return None + + with pytest.raises(RuntimeError, match="delete failed"): + await provider.close( + opensandbox_provider.SandboxHandle( + sandbox_id="sandbox-3", + provider_name="opensandbox", + raw=DeleteFailsCloseSucceedsRaw(), + ), + delete=True, + ) + async def test_create_once_and_connect_after_create_error_paths( fake_opensandbox_sdk: None, @@ -567,6 +659,78 @@ async def no_sleep(_seconds: float) -> None: SandboxSpec(image="image:tag"), ) + provider = opensandbox_provider.OpenSandboxProvider( + connection={"request_timeout_s": 3}, + probe={"command": None}, + ) + handle = await provider._create_once(SandboxSpec(image="image:tag", skip_health_check=True)) + assert handle.sandbox_id == "sandbox-1" + assert FakeSandbox.created_kwargs["skip_health_check"] is True + + class TimeoutSandbox(FakeSandbox): + @classmethod + async def create(cls, **_kwargs: Any) -> "FakeSandbox": + await asyncio.get_running_loop().create_future() + return cls() + + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk", + lambda: (TimeoutSandbox, FakeConnectionConfig, object, FakePlatformSpec, object), + ) + provider = opensandbox_provider.OpenSandboxProvider( + create={"timeout_s": 0.01}, + probe={"command": None}, + ) + with pytest.raises(opensandbox_provider.OpenSandboxCreateTimeoutError): + await provider._create_once(SandboxSpec(image="image:tag")) + + class EmptyCreateSandbox(FakeSandbox): + @classmethod + async def create(cls, **_kwargs: Any) -> None: + return None + + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk", + lambda: (EmptyCreateSandbox, FakeConnectionConfig, object, FakePlatformSpec, object), + ) + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + with pytest.raises(RuntimeError, match="returned no sandbox handle"): + await provider._create_once(SandboxSpec(image="image:tag")) + + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk", + lambda: (FakeSandbox, FakeConnectionConfig, object, FakePlatformSpec, object), + ) + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": "probe"}) + cleanup_calls: list[str] = [] + + async def fail_verify(_handle: opensandbox_provider.SandboxHandle) -> None: + raise RuntimeError("probe failed") + + async def cleanup(handle: opensandbox_provider.SandboxHandle) -> None: + cleanup_calls.append(handle.sandbox_id) + + monkeypatch.setattr(provider, "_verify_created_handle", fail_verify) + monkeypatch.setattr(provider, "_cleanup_failed_create_handle", cleanup) + with pytest.raises(RuntimeError, match="probe failed"): + await provider._create_once(SandboxSpec(image="image:tag")) + assert cleanup_calls == ["sandbox-1"] + + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + + async def create_once(_spec: SandboxSpec) -> opensandbox_provider.SandboxHandle: + return opensandbox_provider.SandboxHandle( + sandbox_id="sandbox-semaphore", provider_name="opensandbox", raw=None + ) + + monkeypatch.setattr(provider, "_create_once", create_once) + assert ( + await provider._create_with_retries(SandboxSpec(image="image:tag"), semaphore=asyncio.Semaphore(1)) + ).sandbox_id == "sandbox-semaphore" + async def test_retry_classification_and_await_sdk_helpers(monkeypatch: pytest.MonkeyPatch) -> None: provider = opensandbox_provider.OpenSandboxProvider( @@ -600,5 +764,80 @@ async def cancelled() -> None: ) +async def test_probe_sampling_pool_progress_and_direct_exec_paths(monkeypatch: pytest.MonkeyPatch) -> None: + handles = [ + opensandbox_provider.SandboxHandle(sandbox_id=f"sandbox-{index}", provider_name="opensandbox", raw=object()) + for index in range(3) + ] + seen_handles: list[str] = [] + + provider = opensandbox_provider.OpenSandboxProvider( + probe={"command": "probe", "sample_count": 2}, + pool={"progress_timeout_s": 0.01, "acquire_poll_interval_s": 0.01}, + ) + + async def verify_created_handle(handle: opensandbox_provider.SandboxHandle) -> None: + seen_handles.append(handle.sandbox_id) + + monkeypatch.setattr(provider, "_verify_created_handle", verify_created_handle) + await provider._verify_created_handles(handles) + assert seen_handles == ["sandbox-0", "sandbox-2"] + + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": "probe", "sample_count": 1}) + seen_handles = [] + monkeypatch.setattr(provider, "_verify_created_handle", verify_created_handle) + await provider._verify_created_handles(handles) + assert seen_handles == ["sandbox-0"] + + with pytest.raises(ValueError, match="requires SandboxSpec.image"): + provider._validate_sdk_pool_spec(SandboxSpec(image=None)) + + class Snapshot: + idle_count = 0 + state = SimpleNamespace(value="warming") + + class NoProgressPool: + async def snapshot(self) -> Snapshot: + return Snapshot() + + async def no_sleep(_seconds: float) -> None: + return None + + provider = opensandbox_provider.OpenSandboxProvider( + pool={"progress_timeout_s": 0.001, "acquire_poll_interval_s": 0.001}, + probe={"command": None}, + ) + monkeypatch.setattr(opensandbox_provider.asyncio, "sleep", no_sleep) + with pytest.raises(opensandbox_provider.OpenSandboxCreateTimeoutError, match="warmup progress"): + await provider._wait_sdk_pool_idle( + NoProgressPool(), + spec=SandboxSpec(image="image:tag"), + requested=2, + timeout_s=1, + allow_partial=False, + ) + + provider = opensandbox_provider.OpenSandboxProvider( + connection={"exec_use_server_proxy": False}, + probe={"command": None}, + ) + + async def connect_after_create( + handle: opensandbox_provider.SandboxHandle, + _spec: SandboxSpec, + ) -> opensandbox_provider.SandboxHandle: + return opensandbox_provider.SandboxHandle( + sandbox_id=handle.sandbox_id, + provider_name="opensandbox", + raw="direct", + ) + + monkeypatch.setattr(provider, "_connect_after_create", connect_after_create) + direct_handle = await provider._direct_exec_handle_for_acquired_sandbox( + FakeSandbox("sandbox-direct"), SandboxSpec() + ) + assert direct_handle.raw == "direct" + + async def _return_value(value: Any) -> Any: return value diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index 908ac7a43a..90045cff8f 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -20,6 +20,7 @@ import pytest +import nemo_gym.sandbox.providers.registry as provider_registry from nemo_gym.sandbox import ( AsyncSandbox, Sandbox, @@ -27,18 +28,12 @@ SandboxExecResult, SandboxHandle, SandboxSpec, + create_provider, get_provider_class, list_providers, register_provider, rewrite_image, ) -from nemo_gym.sandbox.providers.opensandbox import provider as opensandbox_provider_module -from nemo_gym.sandbox.providers.opensandbox.provider import ( - IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY, - IMAGE_PULL_POLICY_EXTENSION_KEY, - OpenSandboxCreateVerificationError, - OpenSandboxProvider, -) from responses_api_agents.mini_swe_agent_2.sandbox_environment import MiniSWESandboxEnvironment @@ -55,6 +50,25 @@ def _has_module(module_name: str) -> bool: ) +def _require_opensandbox_provider() -> tuple[Any, Any, Any, str, str]: + pytest.importorskip("tenacity", reason="tenacity optional sandbox dependency is not installed") + from nemo_gym.sandbox.providers.opensandbox import provider as opensandbox_provider_module + from nemo_gym.sandbox.providers.opensandbox.provider import ( + IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY, + IMAGE_PULL_POLICY_EXTENSION_KEY, + OpenSandboxCreateVerificationError, + OpenSandboxProvider, + ) + + return ( + opensandbox_provider_module, + OpenSandboxProvider, + OpenSandboxCreateVerificationError, + IMAGE_PULL_POLICY_EXTENSION_KEY, + IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY, + ) + + class FakeSandboxProvider: name = "fake" last_instance: "FakeSandboxProvider | None" = None @@ -182,19 +196,51 @@ def test_rewrite_image_and_materialize_handle_validation() -> None: asyncio.run(_assert_rewrite_image_and_materialize_handle_validation()) -def test_provider_registry_validation_and_listing() -> None: +def test_provider_registry_validation_and_listing(monkeypatch: pytest.MonkeyPatch) -> None: provider_name = f"fake-{uuid4().hex}" register_provider(provider_name, FakeSandboxProvider) assert get_provider_class(provider_name) is FakeSandboxProvider + assert "opensandbox" in list_providers() assert provider_name in list_providers() with pytest.raises(ValueError, match="must be non-empty"): register_provider("", FakeSandboxProvider) with pytest.raises(ValueError, match="already registered"): register_provider(provider_name, FakeSandboxProvider) + with pytest.raises(ValueError, match="already registered"): + register_provider("opensandbox", FakeSandboxProvider) with pytest.raises(ValueError, match="Unknown sandbox provider"): get_provider_class(f"missing-{uuid4().hex}") + builtin_name = f"builtin-{uuid4().hex}" + monkeypatch.setitem(provider_registry._BUILTIN_PROVIDER_LOADERS, builtin_name, lambda: FakeSandboxProvider) + assert get_provider_class(builtin_name) is FakeSandboxProvider + assert builtin_name in list_providers() + + +def test_create_provider_validation_and_constructor_cleanup() -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + provider = create_provider({provider_name: None}) + assert isinstance(provider, FakeSandboxProvider) + assert provider.marker == "default" + + with pytest.raises(ValueError, match="exactly one provider name"): + create_provider({}) + with pytest.raises(ValueError, match="non-empty string"): + create_provider({"": {}}) + with pytest.raises(TypeError, match="must be a mapping"): + create_provider({provider_name: "not-a-mapping"}) + + class FailingProvider(FakeSandboxProvider): + def __init__(self) -> None: + raise RuntimeError("provider constructor failed") + + failing_provider_name = f"failing-{uuid4().hex}" + register_provider(failing_provider_name, FailingProvider) + with pytest.raises(RuntimeError, match="provider constructor failed"): + Sandbox({failing_provider_name: {}}) + async def _assert_rewrite_image_and_materialize_handle_validation() -> None: assert rewrite_image(None, []) is None @@ -336,6 +382,15 @@ def test_opensandbox_sdk_create_receives_default_image_pull_policy(monkeypatch) async def _assert_opensandbox_sdk_create_receives_default_image_pull_policy(monkeypatch) -> None: + ( + opensandbox_provider_module, + OpenSandboxProvider, + _OpenSandboxCreateVerificationError, + IMAGE_PULL_POLICY_EXTENSION_KEY, + IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY, + ) = _require_opensandbox_provider() + del _OpenSandboxCreateVerificationError + class FakeSDKSandbox: create_calls: list[dict[str, Any]] = [] @@ -381,6 +436,8 @@ def test_opensandbox_connect_after_create_can_use_direct_exec_endpoint(monkeypat async def _assert_opensandbox_connect_after_create_can_use_direct_exec_endpoint(monkeypatch) -> None: + opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() + class FakeConnectionConfig: def __init__(self, **kwargs: Any) -> None: self.kwargs = kwargs @@ -425,6 +482,8 @@ def test_opensandbox_create_probe_can_require_stable_successes(monkeypatch) -> N async def _assert_opensandbox_create_probe_can_require_stable_successes(monkeypatch) -> None: + _opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() + provider = OpenSandboxProvider( probe={ "command": "true", @@ -472,6 +531,8 @@ def test_opensandbox_create_probe_polls_same_sandbox_after_transient_errors(monk async def _assert_opensandbox_create_probe_polls_same_sandbox_after_transient_errors(monkeypatch) -> None: + _opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() + provider = OpenSandboxProvider( create={"connect_poll_s": 0.01}, probe={ @@ -513,6 +574,13 @@ async def fake_exec( def test_opensandbox_create_probe_failures_are_retryable() -> None: + ( + opensandbox_provider_module, + _OpenSandboxProvider, + OpenSandboxCreateVerificationError, + *_unused, + ) = _require_opensandbox_provider() + error = OpenSandboxCreateVerificationError("pod sdk-sandbox-0 failed create probe") assert isinstance(error, SandboxCreateError) @@ -520,6 +588,8 @@ def test_opensandbox_create_probe_failures_are_retryable() -> None: def test_opensandbox_starting_pod_endpoint_errors_are_retryable() -> None: + opensandbox_provider_module, *_unused = _require_opensandbox_provider() + error = RuntimeError( "Get endpoint for sandbox sdk-sandbox-0 port 44772 failed: " "Pod IP is not yet available. The Pod may still be starting." @@ -534,6 +604,8 @@ def test_opensandbox_exec_retries_retryable_sdk_failures(monkeypatch) -> None: async def _assert_opensandbox_exec_retries_retryable_sdk_failures(monkeypatch) -> None: + opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() + class FakeRunCommandOpts: def __init__(self, **kwargs: Any) -> None: self.kwargs = kwargs @@ -597,6 +669,8 @@ def test_opensandbox_command_retries_can_be_disabled(monkeypatch) -> None: async def _assert_opensandbox_command_retries_can_be_disabled(monkeypatch) -> None: + opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() + class FakeRunCommandOpts: def __init__(self, **kwargs: Any) -> None: self.kwargs = kwargs @@ -648,6 +722,8 @@ def test_opensandbox_close_timeout_does_not_fail_after_delete() -> None: async def _assert_opensandbox_close_timeout_does_not_fail_after_delete() -> None: + _opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() + class SlowCloseRaw: def __init__(self) -> None: self.killed = False @@ -676,6 +752,8 @@ def test_opensandbox_close_timeout_still_fails_without_delete() -> None: async def _assert_opensandbox_close_timeout_still_fails_without_delete() -> None: + _opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() + class SlowCloseRaw: async def close(self) -> None: await asyncio.sleep(60) @@ -764,3 +842,34 @@ def test_mini_swe_sandbox_environment_validation_and_context_manager() -> None: assert FakeSandboxProvider.last_instance is not None assert FakeSandboxProvider.last_instance.closed[-1][1] is False + + +def test_mini_swe_sandbox_environment_submit_sentinel() -> None: + class SubmitSandboxProvider(FakeSandboxProvider): + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + del handle, command, cwd, env, timeout_s, user + return SandboxExecResult( + stdout="COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\nfinal answer", + stderr=None, + return_code=0, + ) + + provider_name = f"submit-{uuid4().hex}" + register_provider(provider_name, SubmitSandboxProvider) + env = MiniSWESandboxEnvironment(image="image:tag", provider={provider_name: {}}) + + try: + with pytest.raises(Exception) as exc_info: + env.execute("submit") + assert exc_info.value.messages[0]["extra"]["submission"] == "final answer" + finally: + env.cleanup() From 55a8ba8d3e016a57f1835cb2ae741a29f40db82f Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Mon, 1 Jun 2026 14:59:12 -0700 Subject: [PATCH 03/14] Address sandbox API review comments Signed-off-by: Hemil Desai --- nemo_gym/sandbox/__init__.py | 5 +- nemo_gym/sandbox/api.py | 43 ++++------- nemo_gym/sandbox/providers/__init__.py | 2 + nemo_gym/sandbox/providers/base.py | 55 +++++++++++--- .../sandbox/providers/opensandbox/provider.py | 58 +++++++++++---- nemo_gym/sandbox/providers/registry.py | 4 +- nemo_gym/sandbox/utils.py | 27 +++++++ .../mini_swe_agent_2/README.md | 10 ++- .../configs/mini_swe_agent_opensandbox.yaml | 7 +- .../mini_swe_agent_2/sandbox_environment.py | 11 ++- tests/unit_tests/test_opensandbox_provider.py | 15 ++-- tests/unit_tests/test_sandbox.py | 74 ++++++++++++++++--- 12 files changed, 226 insertions(+), 85 deletions(-) create mode 100644 nemo_gym/sandbox/utils.py diff --git a/nemo_gym/sandbox/__init__.py b/nemo_gym/sandbox/__init__.py index cd9e1cff9e..6daf582e6f 100644 --- a/nemo_gym/sandbox/__init__.py +++ b/nemo_gym/sandbox/__init__.py @@ -14,13 +14,14 @@ """Public sandbox API for NeMo Gym.""" -from nemo_gym.sandbox.api import AsyncSandbox, Sandbox, rewrite_image +from nemo_gym.sandbox.api import AsyncSandbox, Sandbox from nemo_gym.sandbox.providers import ( SandboxBatchCreateError, SandboxCreateError, SandboxCreateVerificationError, SandboxExecResult, SandboxHandle, + SandboxHandleReferenceProvider, SandboxProvider, SandboxSpec, create_provider, @@ -28,6 +29,7 @@ list_providers, register_provider, ) +from nemo_gym.sandbox.utils import rewrite_image __all__ = [ @@ -38,6 +40,7 @@ "SandboxCreateVerificationError", "SandboxExecResult", "SandboxHandle", + "SandboxHandleReferenceProvider", "SandboxProvider", "SandboxSpec", "create_provider", diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index 1348c00273..7dcd1a1064 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -29,6 +29,7 @@ from nemo_gym.sandbox.providers import ( SandboxExecResult, SandboxHandle, + SandboxHandleReferenceProvider, SandboxProvider, SandboxSpec, create_provider, @@ -38,16 +39,10 @@ T = TypeVar("T") -def rewrite_image(image: str | None, rewrites: list[dict[str, str]]) -> str | None: - """Apply ordered image-prefix rewrites used by sandbox configs.""" - if image is None: - return None - for rewrite in rewrites: - from_prefix = rewrite["from"] - to_prefix = rewrite["to"] - if image.startswith(from_prefix): - return to_prefix + image[len(from_prefix) :] - return image +async def _maybe_await(value: T | Awaitable[T]) -> T: + if hasattr(value, "__await__"): + return await value + return value class AsyncSandbox: @@ -109,15 +104,11 @@ async def download_file(self, handle: SandboxHandle, source_path: str, target_pa async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: await self._provider.close(handle, delete=delete) - async def delete(self, handle: SandboxHandle) -> None: - await self.close(handle, delete=True) - async def aclose(self) -> None: - close_provider = getattr(self._provider, "aclose", None) - if close_provider is not None: - await close_provider() + await self._provider.aclose() async def shutdown(self) -> None: + """Close provider-scoped resources such as SDK clients or warm pools.""" await self.aclose() async def __aenter__(self) -> "AsyncSandbox": @@ -126,21 +117,17 @@ async def __aenter__(self) -> "AsyncSandbox": async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: await self.aclose() - def handle_reference(self, handle: SandboxHandle) -> Any: - make_reference = getattr(self._provider, "handle_reference", None) - if make_reference is None: + async def handle_reference(self, handle: SandboxHandle) -> Any: + if not isinstance(self._provider, SandboxHandleReferenceProvider): return handle - return make_reference(handle) + return await _maybe_await(self._provider.handle_reference(handle)) async def materialize_handle(self, value: Any) -> SandboxHandle: - materialize = getattr(self._provider, "materialize_handle", None) - if materialize is None: + if not isinstance(self._provider, SandboxHandleReferenceProvider): if isinstance(value, SandboxHandle): return value raise ValueError(f"Provider {self.provider_name!r} cannot materialize handle references") - result = materialize(value) - if hasattr(result, "__await__"): - result = await result + result = await _maybe_await(self._provider.materialize_handle(value)) if not isinstance(result, SandboxHandle): raise TypeError(f"materialize_handle must return SandboxHandle, got {type(result).__name__}") return result @@ -273,10 +260,8 @@ def download_file(self, handle: SandboxHandle, source_path: str, target_path: Pa def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: self._runner.run("close", lambda: self._async_sandbox.close(handle, delete=delete)) - def delete(self, handle: SandboxHandle) -> None: - self.close(handle, delete=True) - def shutdown(self) -> None: + """Close provider-scoped resources such as SDK clients or warm pools.""" if self._closed: return self._closed = True @@ -286,7 +271,7 @@ def shutdown(self) -> None: self._runner.close() def handle_reference(self, handle: SandboxHandle) -> Any: - return self._runner.call("handle_reference", lambda: self._async_sandbox.handle_reference(handle)) + return self._runner.run("handle_reference", lambda: self._async_sandbox.handle_reference(handle)) def materialize_handle(self, value: Any) -> SandboxHandle: return self._runner.run("materialize_handle", lambda: self._async_sandbox.materialize_handle(value)) diff --git a/nemo_gym/sandbox/providers/__init__.py b/nemo_gym/sandbox/providers/__init__.py index 359e99c19b..410f201ba3 100644 --- a/nemo_gym/sandbox/providers/__init__.py +++ b/nemo_gym/sandbox/providers/__init__.py @@ -20,6 +20,7 @@ SandboxCreateVerificationError, SandboxExecResult, SandboxHandle, + SandboxHandleReferenceProvider, SandboxProvider, SandboxSpec, ) @@ -37,6 +38,7 @@ "SandboxCreateVerificationError", "SandboxExecResult", "SandboxHandle", + "SandboxHandleReferenceProvider", "SandboxProvider", "SandboxSpec", "create_provider", diff --git a/nemo_gym/sandbox/providers/base.py b/nemo_gym/sandbox/providers/base.py index 7430cb63aa..bde9f09ee3 100644 --- a/nemo_gym/sandbox/providers/base.py +++ b/nemo_gym/sandbox/providers/base.py @@ -19,9 +19,10 @@ instead of importing provider-specific modules. """ +from collections.abc import Awaitable from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Protocol +from typing import Any, Protocol, runtime_checkable @dataclass(frozen=True) @@ -37,14 +38,18 @@ class SandboxSpec: resources: dict[str, str] = field(default_factory=dict) entrypoint: list[str] | None = None extensions: dict[str, str] = field(default_factory=dict) - platform: dict[str, Any] | None = None - volumes: list[dict[str, Any]] | None = None - skip_health_check: bool | None = None + provider_options: dict[str, Any] = field(default_factory=dict) -@dataclass(frozen=True) +@dataclass class SandboxHandle: - """Provider-neutral handle to a created sandbox.""" + """Provider-neutral handle to a created sandbox. + + ``raw`` is provider-owned opaque state, such as an SDK sandbox object, + transport session, or lightweight provider reference. Public Gym code + should pass it back to the provider through this handle rather than + inspecting or mutating it directly. + """ sandbox_id: str provider_name: str @@ -53,18 +58,25 @@ class SandboxHandle: @dataclass(frozen=True) class SandboxExecResult: - """Provider-neutral process execution result.""" + """Provider-neutral process execution result. + + ``return_code`` is the process exit code when the sandbox actually ran the + command. Providers may use a non-process sentinel with ``error_type`` set + when the sandbox runtime reports an execution failure without a process + exit code. + """ stdout: str | None stderr: str | None return_code: int + error_type: str | None = None class SandboxCreateError(RuntimeError): """Raised when a provider cannot create a sandbox.""" -class SandboxBatchCreateError(SandboxCreateError): +class SandboxBatchCreateError(RuntimeError): """Raised when a provider cannot complete sandbox batch creation.""" @@ -90,8 +102,12 @@ async def create_batch( ) -> list[SandboxHandle]: """Create several equivalent sandboxes. - Providers that have a native bulk-allocation primitive should use it. - Providers without one may fall back to calling ``create`` repeatedly. + Providers that have a native bulk-allocation primitive or warm-pool + implementation should use it. Providers without one may fall back to + calling ``create`` repeatedly. Long-lived pools are provider-owned and + configured through provider config or ``SandboxSpec.provider_options``, + rather than through a separate public pool handle. + When ``allow_partial`` is true, providers may return a smaller contiguous prefix of successfully created handles instead of failing the whole batch. @@ -131,6 +147,23 @@ async def download_file(self, handle: SandboxHandle, source_path: str, target_pa """Download one sandbox file to the local filesystem.""" ... - async def close(self, handle: SandboxHandle, *, delete: bool) -> None: + async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: """Close provider resources and optionally delete the sandbox.""" ... + + async def aclose(self) -> None: + """Close provider-scoped resources such as SDK clients or warm pools.""" + ... + + +@runtime_checkable +class SandboxHandleReferenceProvider(Protocol): + """Optional provider trait for loop-safe sandbox handle references.""" + + def handle_reference(self, handle: SandboxHandle) -> Any | Awaitable[Any]: + """Return a serializable or loop-safe reference for ``handle``.""" + ... + + def materialize_handle(self, value: Any) -> SandboxHandle | Awaitable[SandboxHandle]: + """Convert a value from ``handle_reference`` back into a local handle.""" + ... diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 07f547eefd..10423d8180 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -96,6 +96,9 @@ class OpenSandboxCreateVerificationError(SandboxCreateVerificationError): DEFAULT_IMAGE_PULL_POLICY = "IfNotPresent" IMAGE_PULL_POLICY_EXTENSION_KEY = "imagePullPolicy" IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY = "opensandbox.extensions.image-pull-policy" +PROVIDER_OPTION_PLATFORM = "platform" +PROVIDER_OPTION_SKIP_HEALTH_CHECK = "skip_health_check" +PROVIDER_OPTION_VOLUMES = "volumes" VALID_IMAGE_PULL_POLICIES = {"Always", "IfNotPresent", "Never"} STATUS_CODE_RE = re.compile(r"(?:status code|http)\D+(\d{3})", re.IGNORECASE) @@ -117,7 +120,7 @@ def _require_opensandbox_sdk() -> tuple[Any, Any, Any, Any, Any]: except ModuleNotFoundError as e: raise ModuleNotFoundError( "OpenSandbox SDK is required for the opensandbox sandbox provider. " - "Install it in the NeMo-RL runtime image before using " + "Install nemo-gym[sandbox] in the runtime image before using " "env.sandbox.provider.name=opensandbox." ) from e @@ -135,7 +138,7 @@ def _require_opensandbox_sdk_pool() -> tuple[Any, Any, Any, Any]: except ImportError as e: raise ModuleNotFoundError( "OpenSandbox SDK >=0.1.9 is required for native SDK pool batch creation. " - "Install opensandbox>=0.1.9 in the NeMo-RL runtime image." + "Install nemo-gym[sandbox] in the runtime image." ) from e return AcquirePolicy, InMemoryAsyncPoolStateStore, PoolCreationSpec, SandboxPoolAsync @@ -345,6 +348,15 @@ def _to_volumes(volumes: list[dict[str, Any]]) -> list[Any]: return [Volume(**volume) for volume in volumes] +def _provider_option_bool(provider_options: dict[str, Any], key: str) -> bool | None: + value = provider_options.get(key) + if value is None: + return None + if not isinstance(value, bool): + raise TypeError(f"OpenSandbox provider option {key!r} must be a bool") + return value + + def _seconds_to_timedelta(seconds: int | float | None) -> timedelta | None: if seconds is None: return None @@ -403,8 +415,8 @@ def __post_init__(self) -> None: class OpenSandboxProbeConfig: """Post-create probe settings.""" - command: str | None = "printf nemo-rl-sandbox-ready" - expected_stdout: str | None = "nemo-rl-sandbox-ready" + command: str | None = "printf nemo-gym-sandbox-ready" + expected_stdout: str | None = "nemo-gym-sandbox-ready" timeout_s: int = 30 deadline_s: float | None = None sample_count: int | None = None @@ -811,14 +823,18 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: kwargs["ready_timeout"] = timedelta(seconds=spec.ready_timeout_s) if spec.entrypoint is not None: kwargs["entrypoint"] = spec.entrypoint - if spec.platform is not None: - kwargs["platform"] = _to_platform_spec(spec.platform) - if spec.volumes is not None: - kwargs["volumes"] = _to_volumes(spec.volumes) + platform = spec.provider_options.get(PROVIDER_OPTION_PLATFORM) + volumes = spec.provider_options.get(PROVIDER_OPTION_VOLUMES) + if platform is not None: + kwargs["platform"] = _to_platform_spec(platform) + if volumes is not None: + kwargs["volumes"] = _to_volumes(volumes) if self._create.skip_health_check: kwargs["skip_health_check"] = True - elif spec.skip_health_check is not None: - kwargs["skip_health_check"] = spec.skip_health_check + else: + skip_health_check = _provider_option_bool(spec.provider_options, PROVIDER_OPTION_SKIP_HEALTH_CHECK) + if skip_health_check is not None: + kwargs["skip_health_check"] = skip_health_check timeout_s = self._create.timeout_s if timeout_s is None and self._connection.request_timeout_s is not None: @@ -928,8 +944,12 @@ def _to_pool_creation_spec(self, spec: SandboxSpec) -> Any: env=spec.env or None, metadata=spec.metadata or None, extensions=spec.extensions or None, - platform=_to_platform_spec(spec.platform) if spec.platform is not None else None, - volumes=_to_volumes(spec.volumes) if spec.volumes is not None else None, + platform=_to_platform_spec(spec.provider_options[PROVIDER_OPTION_PLATFORM]) + if PROVIDER_OPTION_PLATFORM in spec.provider_options + else None, + volumes=_to_volumes(spec.provider_options[PROVIDER_OPTION_VOLUMES]) + if PROVIDER_OPTION_VOLUMES in spec.provider_options + else None, ) async def _wait_sdk_pool_idle( @@ -998,6 +1018,10 @@ async def _create_batch_sdk_pool( idle_timeout_s = float(self._pool.idle_timeout_s or spec.timeout_s or max(ready_timeout_s * 2.0, 3600.0)) primary_lock_ttl_s = float(self._pool.primary_lock_ttl_s or max(ready_timeout_s + 60.0, 60.0)) pool_name = f"nemo-gym-{uuid4().hex[:12]}" + skip_health_check = bool( + self._create.skip_health_check + or _provider_option_bool(spec.provider_options, PROVIDER_OPTION_SKIP_HEALTH_CHECK) + ) async def _warmup_preparer(sandbox: Any) -> None: if self._probe.command is None: @@ -1033,8 +1057,8 @@ async def _warmup_preparer(sandbox: Any) -> None: acquire_ready_timeout=timedelta(seconds=ready_timeout_s), warmup_ready_timeout=timedelta(seconds=ready_timeout_s), warmup_sandbox_preparer=_warmup_preparer, - acquire_skip_health_check=bool(self._create.skip_health_check or spec.skip_health_check), - warmup_skip_health_check=bool(self._create.skip_health_check or spec.skip_health_check), + acquire_skip_health_check=skip_health_check, + warmup_skip_health_check=skip_health_check, idle_timeout=timedelta(seconds=idle_timeout_s), ) handles: list[SandboxHandle] = [] @@ -1190,14 +1214,16 @@ async def _exec( if execution.error is not None: stderr_parts.append(f"{execution.error.name}: {execution.error.value}") stderr = "\n".join(stderr_parts) or None + error_type = None if execution.exit_code is not None: return_code = execution.exit_code elif execution.error is not None: - return_code = 1 + return_code = 125 + error_type = "sandbox" else: return_code = 0 - return SandboxExecResult(stdout=stdout, stderr=stderr, return_code=return_code) + return SandboxExecResult(stdout=stdout, stderr=stderr, return_code=return_code, error_type=error_type) async def exec( self, diff --git a/nemo_gym/sandbox/providers/registry.py b/nemo_gym/sandbox/providers/registry.py index efa36c5433..8ec7ea81a8 100644 --- a/nemo_gym/sandbox/providers/registry.py +++ b/nemo_gym/sandbox/providers/registry.py @@ -27,11 +27,11 @@ _BUILTIN_PROVIDER_LOADERS: dict[str, ProviderLoader] = {} -def register_provider(name: str, provider_class: ProviderClass) -> None: +def register_provider(name: str, provider_class: ProviderClass, *, override: bool = False) -> None: """Register a sandbox provider class.""" if not name: raise ValueError("Provider name must be non-empty") - if name in _PROVIDER_REGISTRY or name in _BUILTIN_PROVIDER_LOADERS: + if not override and (name in _PROVIDER_REGISTRY or name in _BUILTIN_PROVIDER_LOADERS): raise ValueError(f"Sandbox provider {name!r} is already registered") _PROVIDER_REGISTRY[name] = provider_class diff --git a/nemo_gym/sandbox/utils.py b/nemo_gym/sandbox/utils.py new file mode 100644 index 0000000000..b25f0f6962 --- /dev/null +++ b/nemo_gym/sandbox/utils.py @@ -0,0 +1,27 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Sandbox utility helpers.""" + + +def rewrite_image(image: str | None, rewrites: list[dict[str, str]]) -> str | None: + """Apply ordered image-prefix rewrites used by sandbox configs.""" + if image is None: + return None + for rewrite in rewrites: + from_prefix = rewrite["from"] + to_prefix = rewrite["to"] + if image.startswith(from_prefix): + return to_prefix + image[len(from_prefix) :] + return image diff --git a/responses_api_agents/mini_swe_agent_2/README.md b/responses_api_agents/mini_swe_agent_2/README.md index ff72d22281..0e5dba44e7 100644 --- a/responses_api_agents/mini_swe_agent_2/README.md +++ b/responses_api_agents/mini_swe_agent_2/README.md @@ -147,9 +147,10 @@ mini_swe_agent_2: cpu: "2" memory: 8Gi ephemeral-storage: 20Gi - platform: - os: linux - arch: amd64 + provider_options: + platform: + os: linux + arch: amd64 metadata: benchmark: swebench-verified harness: mini-swe-agent @@ -289,7 +290,8 @@ environment: connection: ... spec: resources: ... - platform: ... + provider_options: + platform: ... metadata: ... ``` diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml index 5ec99bfab1..e2dcd5dec1 100644 --- a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml @@ -44,9 +44,10 @@ mini_swe_agent_2: cpu: "2" memory: 8Gi ephemeral-storage: 20Gi - platform: - os: linux - arch: amd64 + provider_options: + platform: + os: linux + arch: amd64 metadata: benchmark: swebench-verified harness: mini-swe-agent diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index 8a8b8d845b..fac486677f 100644 --- a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -32,7 +32,8 @@ def __init__(self, *messages: dict[str, Any]) -> None: super().__init__() -from nemo_gym.sandbox import Sandbox, SandboxSpec, rewrite_image +from nemo_gym.sandbox import Sandbox, SandboxSpec +from nemo_gym.sandbox.utils import rewrite_image @dataclass @@ -79,6 +80,10 @@ def __init__( spec_config = dict(self.config.spec) image = spec_config.pop("image", None) or self.config.image image = rewrite_image(image, spec_config.pop("image_rewrites", [])) + provider_options = dict(spec_config.pop("provider_options", {})) + for option_key in ("platform", "volumes", "skip_health_check"): + if option_key in spec_config: + provider_options[option_key] = spec_config.pop(option_key) env = dict(spec_config.pop("env", {})) for key in self.config.forward_env: @@ -103,9 +108,7 @@ def __init__( resources=spec_config.pop("resources", {}), entrypoint=spec_config.pop("entrypoint", None), extensions=spec_config.pop("extensions", {}), - platform=spec_config.pop("platform", None), - volumes=spec_config.pop("volumes", None), - skip_health_check=spec_config.pop("skip_health_check", None), + provider_options=provider_options, ) ) diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index 97e525d4d3..48bea2c32e 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -232,7 +232,7 @@ async def test_sdk_pool_passes_platform_through_pool_creation_spec( handles = await provider.create_batch( SandboxSpec( image="mirror.gcr.io/astral/uv:python3.12-bookworm-slim", - platform={"os": "linux", "arch": "amd64"}, + provider_options={"platform": {"os": "linux", "arch": "amd64"}}, ), 1, ) @@ -479,7 +479,8 @@ def __init__(self) -> None: } result = await provider.exec(handle, "fail", user="agent") - assert result.return_code == 1 + assert result.return_code == 125 + assert result.error_type == "sandbox" assert result.stderr == "stderr\nCommandError: failed" assert raw.commands.calls[1][0] == "su -s /bin/sh -c fail agent" @@ -619,9 +620,11 @@ async def test_create_once_and_connect_after_create_error_paths( timeout_s=10, ready_timeout_s=20, entrypoint=["/bin/sh"], - platform={"os": "linux", "arch": "amd64"}, - volumes=[{"name": "workspace"}], - skip_health_check=False, + provider_options={ + "platform": {"os": "linux", "arch": "amd64"}, + "volumes": [{"name": "workspace"}], + "skip_health_check": False, + }, ) handle = await provider._create_once(spec) assert handle.sandbox_id == "sandbox-1" @@ -663,7 +666,7 @@ async def no_sleep(_seconds: float) -> None: connection={"request_timeout_s": 3}, probe={"command": None}, ) - handle = await provider._create_once(SandboxSpec(image="image:tag", skip_health_check=True)) + handle = await provider._create_once(SandboxSpec(image="image:tag", provider_options={"skip_health_check": True})) assert handle.sandbox_id == "sandbox-1" assert FakeSandbox.created_kwargs["skip_health_check"] is True diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index 90045cff8f..6d997b2292 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -24,6 +24,7 @@ from nemo_gym.sandbox import ( AsyncSandbox, Sandbox, + SandboxBatchCreateError, SandboxCreateError, SandboxExecResult, SandboxHandle, @@ -32,8 +33,8 @@ get_provider_class, list_providers, register_provider, - rewrite_image, ) +from nemo_gym.sandbox.utils import rewrite_image from responses_api_agents.mini_swe_agent_2.sandbox_environment import MiniSWESandboxEnvironment @@ -152,6 +153,59 @@ async def materialize_handle(self, value: Any) -> SandboxHandle: return SandboxHandle(sandbox_id=value["sandbox_id"], provider_name=self.name, raw={"materialized": True}) +class PlainSandboxProvider: + name = "plain" + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + return SandboxHandle(sandbox_id="plain-1", provider_name=self.name, raw={"spec": spec}) + + async def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + del allow_partial + return [await self.create(spec) for _ in range(count)] + + async def connect(self, sandbox_id: str) -> SandboxHandle: + return SandboxHandle(sandbox_id=sandbox_id, provider_name=self.name, raw={}) + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + del handle, command, cwd, env, timeout_s, user + return SandboxExecResult(stdout="ok", stderr=None, return_code=0) + + async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + del handle, target_path, data + + async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + del handle + return f"read:{source_path}".encode() + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + del handle, source_path, target_path + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + del handle, source_path + target_path.write_bytes(b"downloaded") + + async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: + del handle, delete + + async def aclose(self) -> None: + return None + + def test_sandbox_facade_uses_public_provider_api() -> None: asyncio.run(_assert_sandbox_facade_uses_public_provider_api()) @@ -180,9 +234,9 @@ async def _assert_sandbox_facade_uses_public_provider_api() -> None: "user": "agent", } - await sandbox.delete(handle) + await sandbox.close(handle, delete=True) assert provider.closed[0] == (handle, True) - assert sandbox.handle_reference(handle) == {"kind": "fake", "sandbox_id": "fake-1"} + assert await sandbox.handle_reference(handle) == {"kind": "fake", "sandbox_id": "fake-1"} assert await sandbox.materialize_handle({"sandbox_id": "fake-2"}) == SandboxHandle( sandbox_id="fake-2", provider_name="fake", raw={"materialized": True} ) @@ -209,12 +263,15 @@ def test_provider_registry_validation_and_listing(monkeypatch: pytest.MonkeyPatc register_provider(provider_name, FakeSandboxProvider) with pytest.raises(ValueError, match="already registered"): register_provider("opensandbox", FakeSandboxProvider) + register_provider(provider_name, FakeSandboxProvider, override=True) with pytest.raises(ValueError, match="Unknown sandbox provider"): get_provider_class(f"missing-{uuid4().hex}") builtin_name = f"builtin-{uuid4().hex}" monkeypatch.setitem(provider_registry._BUILTIN_PROVIDER_LOADERS, builtin_name, lambda: FakeSandboxProvider) assert get_provider_class(builtin_name) is FakeSandboxProvider + register_provider(builtin_name, PlainSandboxProvider, override=True) + assert get_provider_class(builtin_name) is PlainSandboxProvider assert builtin_name in list_providers() @@ -286,12 +343,10 @@ async def _assert_async_sandbox_batch_file_and_fallback_reference_operations(tmp assert provider.download_calls == [(connected, "/remote/source.txt", target_path)] assert target_path.read_bytes() == b"downloaded" - plain_provider = FakeSandboxProvider() - plain_provider.handle_reference = None # type: ignore[method-assign] - plain_provider.materialize_handle = None # type: ignore[method-assign] + plain_provider = PlainSandboxProvider() plain_sandbox = AsyncSandbox(plain_provider) - plain_handle = SandboxHandle(sandbox_id="plain-1", provider_name="fake", raw={}) - assert plain_sandbox.handle_reference(plain_handle) is plain_handle + plain_handle = SandboxHandle(sandbox_id="plain-1", provider_name="plain", raw={}) + assert await plain_sandbox.handle_reference(plain_handle) is plain_handle assert await plain_sandbox.materialize_handle(plain_handle) is plain_handle try: await plain_sandbox.materialize_handle({"sandbox_id": "plain-2"}) @@ -325,7 +380,7 @@ def test_sync_sandbox_facade_uses_public_provider_api() -> None: "user": "agent", } - sandbox.delete(handle) + sandbox.close(handle, delete=True) assert provider.closed[0] == (handle, True) assert sandbox.handle_reference(handle) == {"kind": "fake", "sandbox_id": "fake-1"} assert sandbox.materialize_handle({"sandbox_id": "fake-3"}).sandbox_id == "fake-3" @@ -584,6 +639,7 @@ def test_opensandbox_create_probe_failures_are_retryable() -> None: error = OpenSandboxCreateVerificationError("pod sdk-sandbox-0 failed create probe") assert isinstance(error, SandboxCreateError) + assert not isinstance(SandboxBatchCreateError("batch failed"), SandboxCreateError) assert opensandbox_provider_module._is_retryable_create_error(error) is True From 893618fe1483f3059cf78ceab3ef8a93de96ba93 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Tue, 2 Jun 2026 15:38:02 -0700 Subject: [PATCH 04/14] Align sandbox API with evaluator semantics Signed-off-by: Hemil Desai --- nemo_gym/sandbox/__init__.py | 26 +- nemo_gym/sandbox/api.py | 303 +++++++++++++++++- nemo_gym/sandbox/providers/__init__.py | 22 ++ nemo_gym/sandbox/providers/base.py | 145 +++++++-- .../sandbox/providers/opensandbox/provider.py | 98 +++++- .../mini_swe_agent_2/sandbox_environment.py | 28 +- tests/unit_tests/test_opensandbox_provider.py | 19 +- tests/unit_tests/test_sandbox.py | 124 ++++++- 8 files changed, 708 insertions(+), 57 deletions(-) diff --git a/nemo_gym/sandbox/__init__.py b/nemo_gym/sandbox/__init__.py index 6daf582e6f..4f3ea39e23 100644 --- a/nemo_gym/sandbox/__init__.py +++ b/nemo_gym/sandbox/__init__.py @@ -14,16 +14,27 @@ """Public sandbox API for NeMo Gym.""" -from nemo_gym.sandbox.api import AsyncSandbox, Sandbox +from nemo_gym.sandbox.api import AsyncSandbox, AsyncSandboxInstance, Sandbox, SandboxInstance from nemo_gym.sandbox.providers import ( + ExecResult, + ImageBuildRequest, + ImageSpec, + OutsideEndpoint, + SandboxAddressProvider, + SandboxAttachProvider, SandboxBatchCreateError, SandboxCreateError, SandboxCreateVerificationError, SandboxExecResult, SandboxHandle, SandboxHandleReferenceProvider, + SandboxImageBuildProvider, + SandboxInlineFileProvider, SandboxProvider, SandboxSpec, + SandboxStatus, + SandboxStatusProvider, + VolumeMount, create_provider, get_provider_class, list_providers, @@ -34,15 +45,28 @@ __all__ = [ "Sandbox", + "SandboxInstance", "AsyncSandbox", + "AsyncSandboxInstance", + "ExecResult", + "ImageBuildRequest", + "ImageSpec", + "OutsideEndpoint", + "SandboxAddressProvider", + "SandboxAttachProvider", "SandboxBatchCreateError", "SandboxCreateError", "SandboxCreateVerificationError", "SandboxExecResult", "SandboxHandle", "SandboxHandleReferenceProvider", + "SandboxImageBuildProvider", + "SandboxInlineFileProvider", "SandboxProvider", "SandboxSpec", + "SandboxStatus", + "SandboxStatusProvider", + "VolumeMount", "create_provider", "get_provider_class", "list_providers", diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index 7dcd1a1064..17a8327ac3 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -20,18 +20,28 @@ """ import asyncio +import tempfile import threading from collections.abc import Awaitable, Callable, Mapping from concurrent.futures import Future +from dataclasses import replace from pathlib import Path from typing import Any, TypeVar from nemo_gym.sandbox.providers import ( + ImageBuildRequest, + OutsideEndpoint, + SandboxAddressProvider, + SandboxAttachProvider, SandboxExecResult, SandboxHandle, SandboxHandleReferenceProvider, + SandboxImageBuildProvider, + SandboxInlineFileProvider, SandboxProvider, SandboxSpec, + SandboxStatus, + SandboxStatusProvider, create_provider, ) @@ -55,8 +65,32 @@ def __init__(self, provider: Mapping[str, Any] | SandboxProvider) -> None: def provider_name(self) -> str: return self._provider.name + async def build_images(self, request: ImageBuildRequest) -> list[str]: + if not isinstance(self._provider, SandboxImageBuildProvider): + raise NotImplementedError(f"Provider {self.provider_name!r} does not support sandbox image builds") + return await self._provider.build_images(request) + + async def _resolve_image_build(self, spec: SandboxSpec) -> SandboxSpec: + if spec.image_build is None: + return spec + built_images = await self.build_images(ImageBuildRequest(specs=[spec.image_build])) + if not built_images: + raise ValueError("build_images returned no image references") + return replace(spec, image=spec.image or built_images[0]) + + async def _write_initial_files(self, handle: SandboxHandle, files: dict[str, str]) -> None: + for target_path, contents in files.items(): + await self.write_file(handle, target_path, contents) + async def create(self, spec: SandboxSpec) -> SandboxHandle: - return await self._provider.create(spec) + spec = await self._resolve_image_build(spec) + handle = await self._provider.create(spec) + try: + await self._write_initial_files(handle, spec.files) + except Exception: + await self.close(handle, delete=True) + raise + return handle async def create_batch( self, @@ -65,10 +99,44 @@ async def create_batch( *, allow_partial: bool = False, ) -> list[SandboxHandle]: - return await self._provider.create_batch(spec, count, allow_partial=allow_partial) + spec = await self._resolve_image_build(spec) + handles = await self._provider.create_batch(spec, count, allow_partial=allow_partial) + try: + await asyncio.gather(*(self._write_initial_files(handle, spec.files) for handle in handles)) + except Exception: + await asyncio.gather(*(self.close(handle, delete=True) for handle in handles), return_exceptions=True) + raise + return handles + + async def start( + self, + spec: SandboxSpec, + *, + outside_endpoints: list[OutsideEndpoint] | None = None, + delete_on_stop: bool = False, + ) -> "AsyncSandboxInstance": + if outside_endpoints: + endpoint_env = {endpoint.env_var: endpoint.url for endpoint in outside_endpoints} + spec = replace(spec, env={**spec.env, **endpoint_env}) + handle = await self.create(spec) + return AsyncSandboxInstance( + sandbox=self, + spec=spec, + handle=handle, + delete_on_stop=delete_on_stop, + ) + + async def attach(self, sandbox_id: str) -> SandboxHandle: + if isinstance(self._provider, SandboxAttachProvider): + return await self._provider.attach(sandbox_id) + connect = getattr(self._provider, "connect", None) + if connect is None: + raise NotImplementedError(f"Provider {self.provider_name!r} does not support attaching to sandboxes") + return await connect(sandbox_id) async def connect(self, sandbox_id: str) -> SandboxHandle: - return await self._provider.connect(sandbox_id) + """Compatibility alias for ``attach``.""" + return await self.attach(sandbox_id) async def exec( self, @@ -77,7 +145,7 @@ async def exec( *, cwd: str | None = None, env: dict[str, str] | None = None, - timeout_s: int | None = None, + timeout_s: int | float | None = None, user: str | int | None = None, ) -> SandboxExecResult: return await self._provider.exec( @@ -90,10 +158,24 @@ async def exec( ) async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: - await self._provider.write_file(handle, target_path, data) + if isinstance(self._provider, SandboxInlineFileProvider): + await self._provider.write_file(handle, target_path, data) + return + with tempfile.TemporaryDirectory(prefix="nemo-gym-sandbox-upload-") as tmp_dir: + source_path = Path(tmp_dir) / "contents" + if isinstance(data, str): + source_path.write_text(data, encoding="utf-8") + else: + source_path.write_bytes(data) + await self.upload_file(handle, source_path, target_path) async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: - return await self._provider.read_file(handle, source_path) + if isinstance(self._provider, SandboxInlineFileProvider): + return await self._provider.read_file(handle, source_path) + with tempfile.TemporaryDirectory(prefix="nemo-gym-sandbox-download-") as tmp_dir: + target_path = Path(tmp_dir) / "contents" + await self.download_file(handle, source_path, target_path) + return target_path.read_bytes() async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: await self._provider.upload_file(handle, source_path, target_path) @@ -104,6 +186,16 @@ async def download_file(self, handle: SandboxHandle, source_path: str, target_pa async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: await self._provider.close(handle, delete=delete) + async def status(self, handle: SandboxHandle) -> SandboxStatus: + if not isinstance(self._provider, SandboxStatusProvider): + return SandboxStatus.UNKNOWN + return await self._provider.status(handle) + + async def container_ip(self, handle: SandboxHandle) -> str | None: + if not isinstance(self._provider, SandboxAddressProvider): + return None + return await self._provider.container_ip(handle) + async def aclose(self) -> None: await self._provider.aclose() @@ -133,6 +225,90 @@ async def materialize_handle(self, value: Any) -> SandboxHandle: return result +class AsyncSandboxInstance: + """Evaluator-style async sandbox object returned by ``AsyncSandbox.start``.""" + + def __init__( + self, + *, + sandbox: AsyncSandbox, + spec: SandboxSpec, + handle: SandboxHandle, + delete_on_stop: bool, + ) -> None: + self._sandbox = sandbox + self._spec = spec + self._handle = handle + self._delete_on_stop = delete_on_stop + self._stopped = False + + @property + def spec(self) -> SandboxSpec: + return self._spec + + @property + def handle(self) -> SandboxHandle: + return self._handle + + async def exec( + self, + command: str, + timeout_sec: int | float | None = 180, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + user: str | int | None = None, + timeout_s: int | float | None = None, + ) -> SandboxExecResult: + return await self._sandbox.exec( + self._handle, + command, + cwd=cwd if cwd is not None else self._spec.workdir, + env=env, + timeout_s=timeout_s if timeout_s is not None else timeout_sec, + user=user, + ) + + async def write_file(self, target_path: str, data: str | bytes) -> None: + await self._sandbox.write_file(self._handle, target_path, data) + + async def read_file(self, source_path: str) -> bytes: + return await self._sandbox.read_file(self._handle, source_path) + + async def upload(self, local_path: Path | str, remote_path: str) -> None: + await self._sandbox.upload_file(self._handle, Path(local_path), remote_path) + + async def download(self, remote_path: str, local_path: Path | str) -> None: + await self._sandbox.download_file(self._handle, remote_path, Path(local_path)) + + async def status(self) -> SandboxStatus: + return await self._sandbox.status(self._handle) + + async def is_running(self) -> bool: + return await self.status() == SandboxStatus.RUNNING + + async def container_ip(self) -> str | None: + return await self._sandbox.container_ip(self._handle) + + async def stop(self, *, delete: bool | None = None) -> None: + if self._stopped: + return + self._stopped = True + await self._sandbox.close( + self._handle, + delete=self._delete_on_stop if delete is None else delete, + ) + + async def close(self, *, delete: bool | None = None) -> None: + await self.stop(delete=delete) + + async def __aenter__(self) -> "AsyncSandboxInstance": + return self + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + await self.stop() + + class _AsyncLoopRunner: """Run async sandbox operations for sync integrations on one private loop.""" @@ -205,6 +381,9 @@ def __init__(self, provider: Mapping[str, Any] | SandboxProvider) -> None: def provider_name(self) -> str: return self._runner.call("provider_name", lambda: self._async_sandbox.provider_name) + def build_images(self, request: ImageBuildRequest) -> list[str]: + return self._runner.run("build_images", lambda: self._async_sandbox.build_images(request)) + def create(self, spec: SandboxSpec) -> SandboxHandle: return self._runner.run("create", lambda: self._async_sandbox.create(spec)) @@ -220,8 +399,29 @@ def create_batch( lambda: self._async_sandbox.create_batch(spec, count, allow_partial=allow_partial), ) + def start( + self, + spec: SandboxSpec, + *, + outside_endpoints: list[OutsideEndpoint] | None = None, + delete_on_stop: bool = False, + ) -> "SandboxInstance": + async_instance = self._runner.run( + "start", + lambda: self._async_sandbox.start( + spec, + outside_endpoints=outside_endpoints, + delete_on_stop=delete_on_stop, + ), + ) + return SandboxInstance(self, async_instance) + + def attach(self, sandbox_id: str) -> SandboxHandle: + return self._runner.run("attach", lambda: self._async_sandbox.attach(sandbox_id)) + def connect(self, sandbox_id: str) -> SandboxHandle: - return self._runner.run("connect", lambda: self._async_sandbox.connect(sandbox_id)) + """Compatibility alias for ``attach``.""" + return self.attach(sandbox_id) def exec( self, @@ -230,7 +430,7 @@ def exec( *, cwd: str | None = None, env: dict[str, str] | None = None, - timeout_s: int | None = None, + timeout_s: int | float | None = None, user: str | int | None = None, ) -> SandboxExecResult: return self._runner.run( @@ -260,6 +460,12 @@ def download_file(self, handle: SandboxHandle, source_path: str, target_path: Pa def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: self._runner.run("close", lambda: self._async_sandbox.close(handle, delete=delete)) + def status(self, handle: SandboxHandle) -> SandboxStatus: + return self._runner.run("status", lambda: self._async_sandbox.status(handle)) + + def container_ip(self, handle: SandboxHandle) -> str | None: + return self._runner.run("container_ip", lambda: self._async_sandbox.container_ip(handle)) + def shutdown(self) -> None: """Close provider-scoped resources such as SDK clients or warm pools.""" if self._closed: @@ -288,3 +494,84 @@ def __del__(self) -> None: # pragma: no cover self.shutdown() except Exception: pass + + +class SandboxInstance: + """Evaluator-style sync sandbox object returned by ``Sandbox.start``.""" + + def __init__(self, owner: Sandbox, async_instance: AsyncSandboxInstance) -> None: + self._owner = owner + self._async_instance = async_instance + + @property + def spec(self) -> SandboxSpec: + return self._owner._runner.call("spec", lambda: self._async_instance.spec) + + @property + def handle(self) -> SandboxHandle: + return self._owner._runner.call("handle", lambda: self._async_instance.handle) + + def exec( + self, + command: str, + timeout_sec: int | float | None = 180, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + user: str | int | None = None, + timeout_s: int | float | None = None, + ) -> SandboxExecResult: + return self._owner._runner.run( + "instance.exec", + lambda: self._async_instance.exec( + command, + timeout_sec=timeout_sec, + cwd=cwd, + env=env, + user=user, + timeout_s=timeout_s, + ), + ) + + def write_file(self, target_path: str, data: str | bytes) -> None: + self._owner._runner.run( + "instance.write_file", + lambda: self._async_instance.write_file(target_path, data), + ) + + def read_file(self, source_path: str) -> bytes: + return self._owner._runner.run("instance.read_file", lambda: self._async_instance.read_file(source_path)) + + def upload(self, local_path: Path | str, remote_path: str) -> None: + self._owner._runner.run( + "instance.upload", + lambda: self._async_instance.upload(local_path, remote_path), + ) + + def download(self, remote_path: str, local_path: Path | str) -> None: + self._owner._runner.run( + "instance.download", + lambda: self._async_instance.download(remote_path, local_path), + ) + + def status(self) -> SandboxStatus: + return self._owner._runner.run("instance.status", self._async_instance.status) + + @property + def is_running(self) -> bool: + return self.status() == SandboxStatus.RUNNING + + def container_ip(self) -> str | None: + return self._owner._runner.run("instance.container_ip", self._async_instance.container_ip) + + def stop(self, *, delete: bool | None = None) -> None: + self._owner._runner.run("instance.stop", lambda: self._async_instance.stop(delete=delete)) + + def close(self, *, delete: bool | None = None) -> None: + self.stop(delete=delete) + + def __enter__(self) -> "SandboxInstance": + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self.stop() diff --git a/nemo_gym/sandbox/providers/__init__.py b/nemo_gym/sandbox/providers/__init__.py index 410f201ba3..3e646735b5 100644 --- a/nemo_gym/sandbox/providers/__init__.py +++ b/nemo_gym/sandbox/providers/__init__.py @@ -15,14 +15,25 @@ """Sandbox provider registry.""" from nemo_gym.sandbox.providers.base import ( + ExecResult, + ImageBuildRequest, + ImageSpec, + OutsideEndpoint, + SandboxAddressProvider, + SandboxAttachProvider, SandboxBatchCreateError, SandboxCreateError, SandboxCreateVerificationError, SandboxExecResult, SandboxHandle, SandboxHandleReferenceProvider, + SandboxImageBuildProvider, + SandboxInlineFileProvider, SandboxProvider, SandboxSpec, + SandboxStatus, + SandboxStatusProvider, + VolumeMount, ) from nemo_gym.sandbox.providers.registry import ( create_provider, @@ -33,14 +44,25 @@ __all__ = [ + "ExecResult", + "ImageBuildRequest", + "ImageSpec", + "OutsideEndpoint", + "SandboxAddressProvider", + "SandboxAttachProvider", "SandboxBatchCreateError", "SandboxCreateError", "SandboxCreateVerificationError", "SandboxExecResult", "SandboxHandle", "SandboxHandleReferenceProvider", + "SandboxImageBuildProvider", + "SandboxInlineFileProvider", "SandboxProvider", "SandboxSpec", + "SandboxStatus", + "SandboxStatusProvider", + "VolumeMount", "create_provider", "get_provider_class", "list_providers", diff --git a/nemo_gym/sandbox/providers/base.py b/nemo_gym/sandbox/providers/base.py index bde9f09ee3..c2ea4ed7d1 100644 --- a/nemo_gym/sandbox/providers/base.py +++ b/nemo_gym/sandbox/providers/base.py @@ -19,24 +19,87 @@ instead of importing provider-specific modules. """ -from collections.abc import Awaitable +from collections.abc import Awaitable, Callable from dataclasses import dataclass, field +from enum import Enum from pathlib import Path from typing import Any, Protocol, runtime_checkable +@dataclass(frozen=True) +class ImageSpec: + """Provider-neutral image build input. + + ``image`` is the target image reference the sandbox should run. ``source`` + describes where the provider or image builder can get the build context + from, for example a Git checkout, local path, archive, or provider-native + image recipe. + """ + + image: str + source: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ImageBuildRequest: + """Request for building one or more sandbox images before creation.""" + + specs: list[ImageSpec] + docker_build_fn: Callable[[ImageSpec], str] | None = None + codebuild_buildspec_fn: Callable[[ImageSpec], str] | None = None + + +@dataclass(frozen=True) +class OutsideEndpoint: + """Endpoint exposed outside the sandbox and passed in through an env var.""" + + url: str + env_var: str + + +@dataclass(frozen=True) +class VolumeMount: + """Provider-neutral volume mount description.""" + + host_path: str | None = None + container_path: str = "/workspace" + readonly: bool = False + efs_filesystem_id: str | None = None + efs_root_directory: str | None = None + efs_access_point_id: str | None = None + + @property + def is_efs(self) -> bool: + """Return whether this mount describes an EFS-backed volume.""" + return self.efs_filesystem_id is not None + + +class SandboxStatus(str, Enum): + """Provider-neutral sandbox lifecycle status.""" + + STARTING = "starting" + RUNNING = "running" + STOPPED = "stopped" + ERROR = "error" + UNKNOWN = "unknown" + + @dataclass(frozen=True) class SandboxSpec: """Provider-neutral sandbox creation request.""" image: str | None = None - snapshot_id: str | None = None - timeout_s: int | None = None - ready_timeout_s: int | None = None + image_build: ImageSpec | None = None + timeout_s: int | float | None = None + ready_timeout_s: int | float | None = None + workdir: str | None = None env: dict[str, str] = field(default_factory=dict) + files: dict[str, str] = field(default_factory=dict) metadata: dict[str, str] = field(default_factory=dict) resources: dict[str, str] = field(default_factory=dict) entrypoint: list[str] | None = None + volumes: list[VolumeMount] = field(default_factory=list) + environment_dir: str | None = None extensions: dict[str, str] = field(default_factory=dict) provider_options: dict[str, Any] = field(default_factory=dict) @@ -72,6 +135,9 @@ class SandboxExecResult: error_type: str | None = None +ExecResult = SandboxExecResult + + class SandboxCreateError(RuntimeError): """Raised when a provider cannot create a sandbox.""" @@ -90,7 +156,13 @@ class SandboxProvider(Protocol): name: str async def create(self, spec: SandboxSpec) -> SandboxHandle: - """Create a sandbox and return a provider-neutral handle.""" + """Create a ready sandbox and return a provider-neutral handle. + + Providers must return only after the sandbox is healthy enough to run + commands and transfer files. If the sandbox cannot become ready before + the configured timeout, providers should raise ``SandboxCreateError`` + or a provider-specific subclass. + """ ... async def create_batch( @@ -114,10 +186,6 @@ async def create_batch( """ ... - async def connect(self, sandbox_id: str) -> SandboxHandle: - """Connect to an existing sandbox.""" - ... - async def exec( self, handle: SandboxHandle, @@ -125,20 +193,12 @@ async def exec( *, cwd: str | None = None, env: dict[str, str] | None = None, - timeout_s: int | None = None, + timeout_s: int | float | None = None, user: str | int | None = None, ) -> SandboxExecResult: """Run a command inside a sandbox.""" ... - async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: - """Write a file into a sandbox.""" - ... - - async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: - """Read a file from a sandbox.""" - ... - async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: """Upload one local file into a sandbox.""" ... @@ -167,3 +227,52 @@ def handle_reference(self, handle: SandboxHandle) -> Any | Awaitable[Any]: def materialize_handle(self, value: Any) -> SandboxHandle | Awaitable[SandboxHandle]: """Convert a value from ``handle_reference`` back into a local handle.""" ... + + +@runtime_checkable +class SandboxAttachProvider(Protocol): + """Optional provider trait for attaching to an existing sandbox.""" + + async def attach(self, sandbox_id: str) -> SandboxHandle: + """Attach to an existing sandbox and return a loop-local handle.""" + ... + + +@runtime_checkable +class SandboxInlineFileProvider(Protocol): + """Optional provider trait for efficient inline file reads and writes.""" + + async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + """Write a small file into a sandbox without a local staging path.""" + ... + + async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + """Read a small file from a sandbox without a local staging path.""" + ... + + +@runtime_checkable +class SandboxStatusProvider(Protocol): + """Optional provider trait for sandbox lifecycle status.""" + + async def status(self, handle: SandboxHandle) -> SandboxStatus: + """Return the current sandbox lifecycle status.""" + ... + + +@runtime_checkable +class SandboxAddressProvider(Protocol): + """Optional provider trait for sandbox network addressing.""" + + async def container_ip(self, handle: SandboxHandle) -> str | None: + """Return the sandbox container IP when the provider exposes one.""" + ... + + +@runtime_checkable +class SandboxImageBuildProvider(Protocol): + """Optional provider trait for building images before sandbox creation.""" + + async def build_images(self, request: ImageBuildRequest) -> list[str]: + """Build images and return the image references that were produced.""" + ... diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 10423d8180..3467012c14 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -32,6 +32,8 @@ SandboxExecResult, SandboxHandle, SandboxSpec, + SandboxStatus, + VolumeMount, ) @@ -98,6 +100,7 @@ class OpenSandboxCreateVerificationError(SandboxCreateVerificationError): IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY = "opensandbox.extensions.image-pull-policy" PROVIDER_OPTION_PLATFORM = "platform" PROVIDER_OPTION_SKIP_HEALTH_CHECK = "skip_health_check" +PROVIDER_OPTION_SNAPSHOT_ID = "snapshot_id" PROVIDER_OPTION_VOLUMES = "volumes" VALID_IMAGE_PULL_POLICIES = {"Always", "IfNotPresent", "Never"} STATUS_CODE_RE = re.compile(r"(?:status code|http)\D+(\d{3})", re.IGNORECASE) @@ -343,9 +346,41 @@ def _to_platform_spec(platform: dict[str, Any]) -> Any: return PlatformSpec(**platform) -def _to_volumes(volumes: list[dict[str, Any]]) -> list[Any]: +def _volume_mount_name(volume: VolumeMount, index: int) -> str: + source = volume.container_path or volume.host_path or f"volume-{index}" + normalized = METADATA_VALUE_RE.sub("-", source.strip("/")).strip("-") + return (normalized or f"volume-{index}")[:63] + + +def _volume_to_mapping(volume: VolumeMount | Mapping[str, Any], index: int) -> dict[str, Any]: + if isinstance(volume, Mapping): + return dict(volume) + if not isinstance(volume, VolumeMount): + raise TypeError(f"OpenSandbox volume entries must be mappings or VolumeMount instances, got {type(volume)!r}") + if volume.is_efs: + raise ValueError("OpenSandbox does not support provider-neutral EFS VolumeMount entries") + if volume.host_path is None: + raise ValueError("OpenSandbox VolumeMount entries require host_path") + return { + "name": _volume_mount_name(volume, index), + "host": {"path": volume.host_path}, + "mount_path": volume.container_path, + "read_only": volume.readonly, + } + + +def _to_volumes(volumes: list[VolumeMount | Mapping[str, Any]]) -> list[Any]: _, _, _, _, Volume = _require_opensandbox_sdk() - return [Volume(**volume) for volume in volumes] + return [Volume(**_volume_to_mapping(volume, index)) for index, volume in enumerate(volumes)] + + +def _spec_volumes(spec: SandboxSpec) -> list[VolumeMount | Mapping[str, Any]] | None: + volumes: list[VolumeMount | Mapping[str, Any]] = [] + volumes.extend(spec.volumes) + provider_volumes = spec.provider_options.get(PROVIDER_OPTION_VOLUMES) + if provider_volumes is not None: + volumes.extend(provider_volumes) + return volumes or None def _provider_option_bool(provider_options: dict[str, Any], key: str) -> bool | None: @@ -363,6 +398,19 @@ def _seconds_to_timedelta(seconds: int | float | None) -> timedelta | None: return timedelta(seconds=float(seconds)) +def _to_sandbox_status(state: Any) -> SandboxStatus: + normalized = str(state or "").lower() + if normalized in {"active", "ready", "running"}: + return SandboxStatus.RUNNING + if normalized in {"creating", "initializing", "pending", "starting"}: + return SandboxStatus.STARTING + if normalized in {"completed", "deleted", "exited", "stopped", "terminated"}: + return SandboxStatus.STOPPED + if normalized in {"crashed", "error", "failed", "unhealthy"}: + return SandboxStatus.ERROR + return SandboxStatus.UNKNOWN + + @dataclass(frozen=True) class OpenSandboxConnectionConfig: """OpenSandbox server connection settings.""" @@ -815,8 +863,9 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: } if spec.image is not None: kwargs["image"] = spec.image - if spec.snapshot_id is not None: - kwargs["snapshot_id"] = spec.snapshot_id + snapshot_id = spec.provider_options.get(PROVIDER_OPTION_SNAPSHOT_ID) + if snapshot_id is not None: + kwargs["snapshot_id"] = snapshot_id if spec.timeout_s is not None: kwargs["timeout"] = timedelta(seconds=spec.timeout_s) if spec.ready_timeout_s is not None: @@ -824,7 +873,7 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: if spec.entrypoint is not None: kwargs["entrypoint"] = spec.entrypoint platform = spec.provider_options.get(PROVIDER_OPTION_PLATFORM) - volumes = spec.provider_options.get(PROVIDER_OPTION_VOLUMES) + volumes = _spec_volumes(spec) if platform is not None: kwargs["platform"] = _to_platform_spec(platform) if volumes is not None: @@ -931,12 +980,13 @@ async def _close_one(handle: SandboxHandle) -> Any: def _validate_sdk_pool_spec(self, spec: SandboxSpec) -> None: if spec.image is None: raise ValueError("OpenSandbox SDK pool requires SandboxSpec.image") - if spec.snapshot_id is not None: + if spec.provider_options.get(PROVIDER_OPTION_SNAPSHOT_ID) is not None: raise ValueError("OpenSandbox SDK pool does not support snapshot_id") def _to_pool_creation_spec(self, spec: SandboxSpec) -> Any: self._validate_sdk_pool_spec(spec) _, _, PoolCreationSpec, _ = _require_opensandbox_sdk_pool() + volumes = _spec_volumes(spec) return PoolCreationSpec( image=spec.image, entrypoint=spec.entrypoint, @@ -947,9 +997,7 @@ def _to_pool_creation_spec(self, spec: SandboxSpec) -> Any: platform=_to_platform_spec(spec.provider_options[PROVIDER_OPTION_PLATFORM]) if PROVIDER_OPTION_PLATFORM in spec.provider_options else None, - volumes=_to_volumes(spec.provider_options[PROVIDER_OPTION_VOLUMES]) - if PROVIDER_OPTION_VOLUMES in spec.provider_options - else None, + volumes=_to_volumes(volumes) if volumes is not None else None, ) async def _wait_sdk_pool_idle( @@ -1161,6 +1209,34 @@ async def connect(self, sandbox_id: str) -> SandboxHandle: sandbox = await Sandbox.connect(sandbox_id, **kwargs) return SandboxHandle(sandbox_id=str(sandbox.id), provider_name=self.name, raw=sandbox) + async def attach(self, sandbox_id: str) -> SandboxHandle: + """Attach to an existing OpenSandbox sandbox.""" + return await self.connect(sandbox_id) + + async def status(self, handle: SandboxHandle) -> SandboxStatus: + """Return the current OpenSandbox lifecycle status.""" + get_info = getattr(handle.raw, "get_info", None) + if get_info is None: + return SandboxStatus.UNKNOWN + info = await self._await_sdk_operation( + get_info, + operation="get_info", + sandbox_id=handle.sandbox_id, + timeout_s=float(self._connection.request_timeout_s) + if self._connection.request_timeout_s is not None + else None, + ) + raw_status = getattr(info, "status", None) + return _to_sandbox_status(getattr(raw_status, "state", None) if raw_status is not None else None) + + async def container_ip(self, handle: SandboxHandle) -> str | None: + """Return the container IP when the OpenSandbox SDK handle exposes it.""" + for attr in ("container_ip", "container_ip_address", "pod_ip", "ip"): + value = getattr(handle.raw, attr, None) + if value: + return str(value) + return None + def _command_retry_count(self) -> int: return ( self._operations.retries if self._operations.command_retries is None else self._operations.command_retries @@ -1173,7 +1249,7 @@ async def _exec( *, cwd: str | None = None, env: dict[str, str] | None = None, - timeout_s: int | None = None, + timeout_s: int | float | None = None, user: str | int | None = None, retries: int | None = None, ) -> SandboxExecResult: @@ -1232,7 +1308,7 @@ async def exec( *, cwd: str | None = None, env: dict[str, str] | None = None, - timeout_s: int | None = None, + timeout_s: int | float | None = None, user: str | int | None = None, ) -> SandboxExecResult: """Run a command inside an OpenSandbox sandbox.""" diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index fac486677f..88e35c1981 100644 --- a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -74,6 +74,8 @@ def __init__( if not self.config.provider: raise ValueError("MiniSWESandboxEnvironment requires provider") + self._sandbox_client: Sandbox | None = None + self._sandbox: Any | None = None self._handle: Any | None = None self._closed = False @@ -84,6 +86,8 @@ def __init__( for option_key in ("platform", "volumes", "skip_health_check"): if option_key in spec_config: provider_options[option_key] = spec_config.pop(option_key) + if "snapshot_id" in spec_config: + provider_options["snapshot_id"] = spec_config.pop("snapshot_id") env = dict(spec_config.pop("env", {})) for key in self.config.forward_env: @@ -92,14 +96,15 @@ def __init__( env[key] = value env.update(self.config.env) - self._sandbox = Sandbox(self.config.provider) - self._handle = self._sandbox.create( + self._sandbox_client = Sandbox(self.config.provider) + self._sandbox = self._sandbox_client.start( SandboxSpec( image=image, - snapshot_id=spec_config.pop("snapshot_id", None), timeout_s=spec_config.pop("timeout_s", None), ready_timeout_s=spec_config.pop("ready_timeout_s", None), + workdir=spec_config.pop("workdir", self.config.cwd), env=env, + files=spec_config.pop("files", {}), metadata={ **spec_config.pop("metadata", {}), "nemo_gym_agent": "mini_swe_agent_2", @@ -107,10 +112,13 @@ def __init__( }, resources=spec_config.pop("resources", {}), entrypoint=spec_config.pop("entrypoint", None), + environment_dir=spec_config.pop("environment_dir", None), extensions=spec_config.pop("extensions", {}), provider_options=provider_options, - ) + ), + delete_on_stop=self.config.delete, ) + self._handle = self._sandbox.handle def get_template_vars(self, **kwargs: Any) -> dict[str, Any]: return {**self.config.__dict__, **kwargs} @@ -149,10 +157,9 @@ def execute( exec_cwd = cwd or self.config.cwd result = self._sandbox.exec( - self._handle, self._command(command, exec_cwd), + timeout_sec=timeout_s, cwd="/", - timeout_s=timeout_s, user=self.config.user, ) output = "\n".join(part for part in (result.stdout, result.stderr) if part) @@ -182,11 +189,14 @@ def cleanup(self) -> None: return self._closed = True try: - if self._handle is not None: - self._sandbox.close(self._handle, delete=self.config.delete) + if self._sandbox is not None: + self._sandbox.stop() + self._sandbox = None self._handle = None finally: - self._sandbox.shutdown() + if self._sandbox_client is not None: + self._sandbox_client.shutdown() + self._sandbox_client = None def __enter__(self) -> "MiniSWESandboxEnvironment": return self diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index 48bea2c32e..a4d1e9d454 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -22,7 +22,7 @@ import pytest -from nemo_gym.sandbox.providers.base import SandboxSpec +from nemo_gym.sandbox.providers.base import SandboxSpec, SandboxStatus, VolumeMount pytest.importorskip("tenacity", reason="tenacity optional sandbox dependency is not installed") @@ -443,10 +443,15 @@ async def read_bytes(self, source_path: str) -> bytes: return f"bytes:{source_path}".encode() class FakeRaw: + container_ip = "10.1.2.3" + def __init__(self) -> None: self.commands = FakeCommands() self.files = FakeFiles() + async def get_info(self) -> Any: + return SimpleNamespace(status=SimpleNamespace(state="RUNNING")) + monkeypatch.setattr( opensandbox_provider, "_require_opensandbox_sdk", @@ -493,13 +498,15 @@ def __init__(self) -> None: await provider.download_file(handle, "/remote/download.txt", download_path) assert raw.files.writes == [("/tmp/file.txt", "contents"), ("/remote/upload.txt", b"upload")] assert download_path.read_bytes() == b"bytes:/remote/download.txt" + assert await provider.status(handle) == SandboxStatus.RUNNING + assert await provider.container_ip(handle) == "10.1.2.3" with pytest.raises(ValueError, match="count"): await provider._create_batch_sdk(SandboxSpec(image="image:tag"), 0) with pytest.raises(ValueError, match="count"): await provider.create_batch(SandboxSpec(image="image:tag"), 0) with pytest.raises(ValueError, match="snapshot_id"): - provider._validate_sdk_pool_spec(SandboxSpec(image="image:tag", snapshot_id="snapshot")) + provider._validate_sdk_pool_spec(SandboxSpec(image="image:tag", provider_options={"snapshot_id": "snapshot"})) with pytest.raises(ValueError, match="Unsupported"): await provider.materialize_handle({"kind": "other"}) @@ -616,11 +623,12 @@ async def test_create_once_and_connect_after_create_error_paths( monkeypatch.setattr(opensandbox_provider, "_to_volumes", lambda volumes: volumes) spec = SandboxSpec( image="image:tag", - snapshot_id="snapshot-1", timeout_s=10, ready_timeout_s=20, entrypoint=["/bin/sh"], + volumes=[VolumeMount(host_path="/host/workspace", container_path="/mnt/workspace", readonly=True)], provider_options={ + "snapshot_id": "snapshot-1", "platform": {"os": "linux", "arch": "amd64"}, "volumes": [{"name": "workspace"}], "skip_health_check": False, @@ -633,7 +641,10 @@ async def test_create_once_and_connect_after_create_error_paths( assert FakeSandbox.created_kwargs["ready_timeout"] == timedelta(seconds=20) assert FakeSandbox.created_kwargs["entrypoint"] == ["/bin/sh"] assert FakeSandbox.created_kwargs["platform"] == FakePlatformSpec(os="linux", arch="amd64") - assert FakeSandbox.created_kwargs["volumes"] == [{"name": "workspace"}] + assert FakeSandbox.created_kwargs["volumes"] == [ + VolumeMount(host_path="/host/workspace", container_path="/mnt/workspace", readonly=True), + {"name": "workspace"}, + ] assert FakeSandbox.created_kwargs["skip_health_check"] is True class FailingConnectSandbox(FakeSandbox): diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index 6d997b2292..b48e4b0cec 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -23,12 +23,16 @@ import nemo_gym.sandbox.providers.registry as provider_registry from nemo_gym.sandbox import ( AsyncSandbox, + ImageBuildRequest, + ImageSpec, + OutsideEndpoint, Sandbox, SandboxBatchCreateError, SandboxCreateError, SandboxExecResult, SandboxHandle, SandboxSpec, + SandboxStatus, create_provider, get_provider_class, list_providers, @@ -78,6 +82,7 @@ def __init__(self, marker: str = "default") -> None: self.marker = marker self.created_specs: list[SandboxSpec] = [] self.exec_calls: list[dict[str, Any]] = [] + self.image_build_requests: list[ImageBuildRequest] = [] self.write_calls: list[tuple[SandboxHandle, str, str | bytes]] = [] self.read_calls: list[tuple[SandboxHandle, str]] = [] self.upload_calls: list[tuple[SandboxHandle, Path, str]] = [] @@ -86,6 +91,10 @@ def __init__(self, marker: str = "default") -> None: self.aclosed = False FakeSandboxProvider.last_instance = self + async def build_images(self, request: ImageBuildRequest) -> list[str]: + self.image_build_requests.append(request) + return [spec.image for spec in request.specs] + async def create(self, spec: SandboxSpec) -> SandboxHandle: self.created_specs.append(spec) return SandboxHandle(sandbox_id="fake-1", provider_name=self.name, raw={"spec": spec}) @@ -110,7 +119,7 @@ async def exec( *, cwd: str | None = None, env: dict[str, str] | None = None, - timeout_s: int | None = None, + timeout_s: int | float | None = None, user: str | int | None = None, ) -> SandboxExecResult: self.exec_calls.append( @@ -140,6 +149,14 @@ async def download_file(self, handle: SandboxHandle, source_path: str, target_pa target_path.parent.mkdir(parents=True, exist_ok=True) target_path.write_bytes(b"downloaded") + async def status(self, handle: SandboxHandle) -> SandboxStatus: + del handle + return SandboxStatus.RUNNING + + async def container_ip(self, handle: SandboxHandle) -> str | None: + del handle + return "10.0.0.1" + async def close(self, handle: SandboxHandle, *, delete: bool) -> None: self.closed.append((handle, delete)) @@ -179,7 +196,7 @@ async def exec( *, cwd: str | None = None, env: dict[str, str] | None = None, - timeout_s: int | None = None, + timeout_s: int | float | None = None, user: str | int | None = None, ) -> SandboxExecResult: del handle, command, cwd, env, timeout_s, user @@ -206,6 +223,53 @@ async def aclose(self) -> None: return None +class TransferOnlySandboxProvider: + name = "transfer-only" + + def __init__(self) -> None: + self.upload_calls: list[tuple[SandboxHandle, Path, str]] = [] + self.download_calls: list[tuple[SandboxHandle, str, Path]] = [] + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + return SandboxHandle(sandbox_id="transfer-1", provider_name=self.name, raw={"spec": spec}) + + async def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + del allow_partial + return [await self.create(spec) for _ in range(count)] + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | float | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + del handle, command, cwd, env, timeout_s, user + return SandboxExecResult(stdout="ok", stderr=None, return_code=0) + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + self.upload_calls.append((handle, source_path, target_path)) + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + self.download_calls.append((handle, source_path, target_path)) + target_path.write_bytes(b"fallback") + + async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: + del handle, delete + + async def aclose(self) -> None: + return None + + def test_sandbox_facade_uses_public_provider_api() -> None: asyncio.run(_assert_sandbox_facade_uses_public_provider_api()) @@ -233,9 +297,39 @@ async def _assert_sandbox_facade_uses_public_provider_api() -> None: "timeout_s": 60, "user": "agent", } + assert await sandbox.status(handle) == SandboxStatus.RUNNING + assert await sandbox.container_ip(handle) == "10.0.0.1" + + built_handle = await sandbox.create( + SandboxSpec(image_build=ImageSpec(image="built:tag", source={"context": "repo"})) + ) + assert built_handle.sandbox_id == "fake-1" + assert provider.image_build_requests[-1].specs[0].source == {"context": "repo"} + assert provider.created_specs[-1].image == "built:tag" + + session = await sandbox.start( + SandboxSpec( + image="image:tag", + workdir="/session", + files={"/tmp/bootstrap.txt": "hello"}, + ), + outside_endpoints=[OutsideEndpoint(url="http://outside", env_var="OUTSIDE_URL")], + delete_on_stop=True, + ) + try: + assert session.spec.env["OUTSIDE_URL"] == "http://outside" + assert await session.is_running() is True + assert await session.container_ip() == "10.0.0.1" + session_result = await session.exec("pwd") + assert session_result.return_code == 0 + assert provider.exec_calls[-1]["cwd"] == "/session" + finally: + await session.stop() + assert provider.write_calls[-1] == (session.handle, "/tmp/bootstrap.txt", "hello") + assert provider.closed[-1] == (session.handle, True) await sandbox.close(handle, delete=True) - assert provider.closed[0] == (handle, True) + assert provider.closed[-1] == (handle, True) assert await sandbox.handle_reference(handle) == {"kind": "fake", "sandbox_id": "fake-1"} assert await sandbox.materialize_handle({"sandbox_id": "fake-2"}) == SandboxHandle( sandbox_id="fake-2", provider_name="fake", raw={"materialized": True} @@ -343,6 +437,17 @@ async def _assert_async_sandbox_batch_file_and_fallback_reference_operations(tmp assert provider.download_calls == [(connected, "/remote/source.txt", target_path)] assert target_path.read_bytes() == b"downloaded" + transfer_provider = TransferOnlySandboxProvider() + transfer_sandbox = AsyncSandbox(transfer_provider) + transfer_handle = SandboxHandle(sandbox_id="transfer-1", provider_name="transfer-only", raw={}) + await transfer_sandbox.write_file(transfer_handle, "/remote/inline.txt", b"fallback") + assert transfer_provider.upload_calls[0][0] == transfer_handle + assert transfer_provider.upload_calls[0][2] == "/remote/inline.txt" + assert await transfer_sandbox.read_file(transfer_handle, "/remote/inline.txt") == b"fallback" + assert transfer_provider.download_calls == [ + (transfer_handle, "/remote/inline.txt", transfer_provider.download_calls[0][2]) + ] + plain_provider = PlainSandboxProvider() plain_sandbox = AsyncSandbox(plain_provider) plain_handle = SandboxHandle(sandbox_id="plain-1", provider_name="plain", raw={}) @@ -386,6 +491,13 @@ def test_sync_sandbox_facade_uses_public_provider_api() -> None: assert sandbox.materialize_handle({"sandbox_id": "fake-3"}).sandbox_id == "fake-3" assert sandbox.provider_name == "fake" assert len(sandbox.create_batch(SandboxSpec(image="image:tag"), 2)) == 2 + session = sandbox.start(SandboxSpec(image="image:tag", workdir="/sync-session"), delete_on_stop=True) + assert session.is_running is True + assert session.container_ip() == "10.0.0.1" + assert session.exec("pwd").return_code == 0 + assert provider.exec_calls[-1]["cwd"] == "/sync-session" + session.stop() + assert provider.closed[-1] == (session.handle, True) sandbox.shutdown() sandbox.shutdown() assert provider.aclosed is True @@ -555,7 +667,7 @@ async def fake_exec( *, cwd: str | None = None, env: dict[str, str] | None = None, - timeout_s: int | None = None, + timeout_s: int | float | None = None, user: str | int | None = None, ) -> SandboxExecResult: calls.append( @@ -608,7 +720,7 @@ async def fake_exec( *, cwd: str | None = None, env: dict[str, str] | None = None, - timeout_s: int | None = None, + timeout_s: int | float | None = None, user: str | int | None = None, ) -> SandboxExecResult: del command, cwd, env, timeout_s, user @@ -909,7 +1021,7 @@ async def exec( *, cwd: str | None = None, env: dict[str, str] | None = None, - timeout_s: int | None = None, + timeout_s: int | float | None = None, user: str | int | None = None, ) -> SandboxExecResult: del handle, command, cwd, env, timeout_s, user From e83183ed15b684e4f5deae4192f001e05dafb05d Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Tue, 2 Jun 2026 15:56:04 -0700 Subject: [PATCH 05/14] Simplify sandbox public API Signed-off-by: Hemil Desai --- nemo_gym/sandbox/__init__.py | 6 +- nemo_gym/sandbox/api.py | 456 +++++------------- nemo_gym/sandbox/providers/__init__.py | 2 - nemo_gym/sandbox/providers/base.py | 9 - .../sandbox/providers/opensandbox/provider.py | 4 - .../mini_swe_agent_2/sandbox_environment.py | 24 +- tests/unit_tests/test_opensandbox_provider.py | 164 +++++++ tests/unit_tests/test_sandbox.py | 293 +++++------ 8 files changed, 453 insertions(+), 505 deletions(-) diff --git a/nemo_gym/sandbox/__init__.py b/nemo_gym/sandbox/__init__.py index 4f3ea39e23..36d54333d4 100644 --- a/nemo_gym/sandbox/__init__.py +++ b/nemo_gym/sandbox/__init__.py @@ -14,14 +14,13 @@ """Public sandbox API for NeMo Gym.""" -from nemo_gym.sandbox.api import AsyncSandbox, AsyncSandboxInstance, Sandbox, SandboxInstance +from nemo_gym.sandbox.api import AsyncSandbox, Sandbox from nemo_gym.sandbox.providers import ( ExecResult, ImageBuildRequest, ImageSpec, OutsideEndpoint, SandboxAddressProvider, - SandboxAttachProvider, SandboxBatchCreateError, SandboxCreateError, SandboxCreateVerificationError, @@ -45,15 +44,12 @@ __all__ = [ "Sandbox", - "SandboxInstance", "AsyncSandbox", - "AsyncSandboxInstance", "ExecResult", "ImageBuildRequest", "ImageSpec", "OutsideEndpoint", "SandboxAddressProvider", - "SandboxAttachProvider", "SandboxBatchCreateError", "SandboxCreateError", "SandboxCreateVerificationError", diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index 17a8327ac3..032256a2c6 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -12,12 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Provider-neutral public sandbox API. - -This module is the boundary Gym code should use when it needs a sandbox. -Provider packages implement the lower-level async protocol; callers use -``AsyncSandbox`` in async code and ``Sandbox`` in synchronous integrations. -""" +"""Provider-neutral public sandbox API.""" import asyncio import tempfile @@ -32,10 +27,8 @@ ImageBuildRequest, OutsideEndpoint, SandboxAddressProvider, - SandboxAttachProvider, SandboxExecResult, SandboxHandle, - SandboxHandleReferenceProvider, SandboxImageBuildProvider, SandboxInlineFileProvider, SandboxProvider, @@ -49,22 +42,41 @@ T = TypeVar("T") -async def _maybe_await(value: T | Awaitable[T]) -> T: - if hasattr(value, "__await__"): - return await value - return value - - class AsyncSandbox: - """Async public facade for provider-backed sandbox operations.""" + """Async sandbox object backed by a runtime provider.""" - def __init__(self, provider: Mapping[str, Any] | SandboxProvider) -> None: + def __init__( + self, + provider: Mapping[str, Any] | SandboxProvider, + spec: SandboxSpec | None = None, + *, + delete_on_stop: bool = False, + ) -> None: self._provider = create_provider(provider) if isinstance(provider, Mapping) else provider + self._spec = spec + self._handle: SandboxHandle | None = None + self._delete_on_stop = delete_on_stop + self._stopped = True @property def provider_name(self) -> str: return self._provider.name + @property + def spec(self) -> SandboxSpec: + if self._spec is None: + raise RuntimeError("Sandbox has not been configured") + return self._spec + + @property + def handle(self) -> SandboxHandle: + return self._require_handle() + + def _require_handle(self) -> SandboxHandle: + if self._handle is None or self._stopped: + raise RuntimeError("Sandbox has not been started") + return self._handle + async def build_images(self, request: ImageBuildRequest) -> list[str]: if not isinstance(self._provider, SandboxImageBuildProvider): raise NotImplementedError(f"Provider {self.provider_name!r} does not support sandbox image builds") @@ -78,86 +90,17 @@ async def _resolve_image_build(self, spec: SandboxSpec) -> SandboxSpec: raise ValueError("build_images returned no image references") return replace(spec, image=spec.image or built_images[0]) - async def _write_initial_files(self, handle: SandboxHandle, files: dict[str, str]) -> None: - for target_path, contents in files.items(): - await self.write_file(handle, target_path, contents) - - async def create(self, spec: SandboxSpec) -> SandboxHandle: - spec = await self._resolve_image_build(spec) - handle = await self._provider.create(spec) - try: - await self._write_initial_files(handle, spec.files) - except Exception: - await self.close(handle, delete=True) - raise - return handle - - async def create_batch( + def _with_outside_endpoints( self, spec: SandboxSpec, - count: int, - *, - allow_partial: bool = False, - ) -> list[SandboxHandle]: - spec = await self._resolve_image_build(spec) - handles = await self._provider.create_batch(spec, count, allow_partial=allow_partial) - try: - await asyncio.gather(*(self._write_initial_files(handle, spec.files) for handle in handles)) - except Exception: - await asyncio.gather(*(self.close(handle, delete=True) for handle in handles), return_exceptions=True) - raise - return handles - - async def start( - self, - spec: SandboxSpec, - *, - outside_endpoints: list[OutsideEndpoint] | None = None, - delete_on_stop: bool = False, - ) -> "AsyncSandboxInstance": - if outside_endpoints: - endpoint_env = {endpoint.env_var: endpoint.url for endpoint in outside_endpoints} - spec = replace(spec, env={**spec.env, **endpoint_env}) - handle = await self.create(spec) - return AsyncSandboxInstance( - sandbox=self, - spec=spec, - handle=handle, - delete_on_stop=delete_on_stop, - ) - - async def attach(self, sandbox_id: str) -> SandboxHandle: - if isinstance(self._provider, SandboxAttachProvider): - return await self._provider.attach(sandbox_id) - connect = getattr(self._provider, "connect", None) - if connect is None: - raise NotImplementedError(f"Provider {self.provider_name!r} does not support attaching to sandboxes") - return await connect(sandbox_id) - - async def connect(self, sandbox_id: str) -> SandboxHandle: - """Compatibility alias for ``attach``.""" - return await self.attach(sandbox_id) - - async def exec( - self, - handle: SandboxHandle, - command: str, - *, - cwd: str | None = None, - env: dict[str, str] | None = None, - timeout_s: int | float | None = None, - user: str | int | None = None, - ) -> SandboxExecResult: - return await self._provider.exec( - handle, - command, - cwd=cwd, - env=env, - timeout_s=timeout_s, - user=user, - ) + outside_endpoints: list[OutsideEndpoint] | None, + ) -> SandboxSpec: + if not outside_endpoints: + return spec + endpoint_env = {endpoint.env_var: endpoint.url for endpoint in outside_endpoints} + return replace(spec, env={**spec.env, **endpoint_env}) - async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + async def _write_inline_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: if isinstance(self._provider, SandboxInlineFileProvider): await self._provider.write_file(handle, target_path, data) return @@ -167,150 +110,105 @@ async def write_file(self, handle: SandboxHandle, target_path: str, data: str | source_path.write_text(data, encoding="utf-8") else: source_path.write_bytes(data) - await self.upload_file(handle, source_path, target_path) - - async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: - if isinstance(self._provider, SandboxInlineFileProvider): - return await self._provider.read_file(handle, source_path) - with tempfile.TemporaryDirectory(prefix="nemo-gym-sandbox-download-") as tmp_dir: - target_path = Path(tmp_dir) / "contents" - await self.download_file(handle, source_path, target_path) - return target_path.read_bytes() - - async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: - await self._provider.upload_file(handle, source_path, target_path) - - async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: - await self._provider.download_file(handle, source_path, target_path) - - async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: - await self._provider.close(handle, delete=delete) - - async def status(self, handle: SandboxHandle) -> SandboxStatus: - if not isinstance(self._provider, SandboxStatusProvider): - return SandboxStatus.UNKNOWN - return await self._provider.status(handle) - - async def container_ip(self, handle: SandboxHandle) -> str | None: - if not isinstance(self._provider, SandboxAddressProvider): - return None - return await self._provider.container_ip(handle) + await self._provider.upload_file(handle, source_path, target_path) - async def aclose(self) -> None: - await self._provider.aclose() - - async def shutdown(self) -> None: - """Close provider-scoped resources such as SDK clients or warm pools.""" - await self.aclose() - - async def __aenter__(self) -> "AsyncSandbox": - return self - - async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: - await self.aclose() - - async def handle_reference(self, handle: SandboxHandle) -> Any: - if not isinstance(self._provider, SandboxHandleReferenceProvider): - return handle - return await _maybe_await(self._provider.handle_reference(handle)) - - async def materialize_handle(self, value: Any) -> SandboxHandle: - if not isinstance(self._provider, SandboxHandleReferenceProvider): - if isinstance(value, SandboxHandle): - return value - raise ValueError(f"Provider {self.provider_name!r} cannot materialize handle references") - result = await _maybe_await(self._provider.materialize_handle(value)) - if not isinstance(result, SandboxHandle): - raise TypeError(f"materialize_handle must return SandboxHandle, got {type(result).__name__}") - return result - - -class AsyncSandboxInstance: - """Evaluator-style async sandbox object returned by ``AsyncSandbox.start``.""" + async def _write_initial_files(self, handle: SandboxHandle, files: dict[str, str]) -> None: + for target_path, contents in files.items(): + await self._write_inline_file(handle, target_path, contents) - def __init__( + async def start( self, + spec: SandboxSpec | None = None, *, - sandbox: AsyncSandbox, - spec: SandboxSpec, - handle: SandboxHandle, - delete_on_stop: bool, - ) -> None: - self._sandbox = sandbox - self._spec = spec + outside_endpoints: list[OutsideEndpoint] | None = None, + delete_on_stop: bool | None = None, + ) -> "AsyncSandbox": + if self._handle is not None and not self._stopped: + raise RuntimeError("Sandbox is already started") + requested_spec = spec if spec is not None else self._spec + if requested_spec is None: + raise ValueError("Sandbox.start() requires a SandboxSpec") + + requested_spec = self._with_outside_endpoints(requested_spec, outside_endpoints) + resolved_spec = await self._resolve_image_build(requested_spec) + handle = await self._provider.create(resolved_spec) + try: + await self._write_initial_files(handle, resolved_spec.files) + except Exception: + await self._provider.close(handle, delete=True) + raise + + self._spec = resolved_spec self._handle = handle - self._delete_on_stop = delete_on_stop + self._delete_on_stop = self._delete_on_stop if delete_on_stop is None else delete_on_stop self._stopped = False - - @property - def spec(self) -> SandboxSpec: - return self._spec - - @property - def handle(self) -> SandboxHandle: - return self._handle + return self async def exec( self, command: str, - timeout_sec: int | float | None = 180, *, cwd: str | None = None, env: dict[str, str] | None = None, + timeout_s: int | float | None = 180, user: str | int | None = None, - timeout_s: int | float | None = None, ) -> SandboxExecResult: - return await self._sandbox.exec( - self._handle, + return await self._provider.exec( + self._require_handle(), command, - cwd=cwd if cwd is not None else self._spec.workdir, + cwd=cwd if cwd is not None else self.spec.workdir, env=env, - timeout_s=timeout_s if timeout_s is not None else timeout_sec, + timeout_s=timeout_s, user=user, ) - async def write_file(self, target_path: str, data: str | bytes) -> None: - await self._sandbox.write_file(self._handle, target_path, data) - - async def read_file(self, source_path: str) -> bytes: - return await self._sandbox.read_file(self._handle, source_path) - async def upload(self, local_path: Path | str, remote_path: str) -> None: - await self._sandbox.upload_file(self._handle, Path(local_path), remote_path) + await self._provider.upload_file(self._require_handle(), Path(local_path), remote_path) async def download(self, remote_path: str, local_path: Path | str) -> None: - await self._sandbox.download_file(self._handle, remote_path, Path(local_path)) + await self._provider.download_file(self._require_handle(), remote_path, Path(local_path)) async def status(self) -> SandboxStatus: - return await self._sandbox.status(self._handle) + if self._handle is None: + return SandboxStatus.UNKNOWN + if self._stopped: + return SandboxStatus.STOPPED + if not isinstance(self._provider, SandboxStatusProvider): + return SandboxStatus.UNKNOWN + return await self._provider.status(self._handle) async def is_running(self) -> bool: return await self.status() == SandboxStatus.RUNNING async def container_ip(self) -> str | None: - return await self._sandbox.container_ip(self._handle) + if not isinstance(self._provider, SandboxAddressProvider): + return None + return await self._provider.container_ip(self._require_handle()) async def stop(self, *, delete: bool | None = None) -> None: - if self._stopped: + if self._handle is None or self._stopped: return self._stopped = True - await self._sandbox.close( + await self._provider.close( self._handle, delete=self._delete_on_stop if delete is None else delete, ) - async def close(self, *, delete: bool | None = None) -> None: - await self.stop(delete=delete) + async def aclose(self) -> None: + try: + await self.stop() + finally: + await self._provider.aclose() - async def __aenter__(self) -> "AsyncSandboxInstance": + async def __aenter__(self) -> "AsyncSandbox": return self async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: - await self.stop() + await self.aclose() class _AsyncLoopRunner: - """Run async sandbox operations for sync integrations on one private loop.""" + """Run async sandbox operations for sync callers.""" def __init__(self) -> None: self._loop = asyncio.new_event_loop() @@ -363,14 +261,20 @@ def close(self) -> None: class Sandbox: - """Sync public facade for provider-backed sandbox operations.""" + """Synchronous wrapper around ``AsyncSandbox``.""" - def __init__(self, provider: Mapping[str, Any] | SandboxProvider) -> None: + def __init__( + self, + provider: Mapping[str, Any] | SandboxProvider, + spec: SandboxSpec | None = None, + *, + delete_on_stop: bool = False, + ) -> None: self._runner = _AsyncLoopRunner() try: self._async_sandbox = self._runner.call( "__init__", - lambda: AsyncSandbox(provider), + lambda: AsyncSandbox(provider, spec, delete_on_stop=delete_on_stop), ) except BaseException: self._runner.close() @@ -381,32 +285,25 @@ def __init__(self, provider: Mapping[str, Any] | SandboxProvider) -> None: def provider_name(self) -> str: return self._runner.call("provider_name", lambda: self._async_sandbox.provider_name) - def build_images(self, request: ImageBuildRequest) -> list[str]: - return self._runner.run("build_images", lambda: self._async_sandbox.build_images(request)) + @property + def spec(self) -> SandboxSpec: + return self._runner.call("spec", lambda: self._async_sandbox.spec) - def create(self, spec: SandboxSpec) -> SandboxHandle: - return self._runner.run("create", lambda: self._async_sandbox.create(spec)) + @property + def handle(self) -> SandboxHandle: + return self._runner.call("handle", lambda: self._async_sandbox.handle) - def create_batch( - self, - spec: SandboxSpec, - count: int, - *, - allow_partial: bool = False, - ) -> list[SandboxHandle]: - return self._runner.run( - "create_batch", - lambda: self._async_sandbox.create_batch(spec, count, allow_partial=allow_partial), - ) + def build_images(self, request: ImageBuildRequest) -> list[str]: + return self._runner.run("build_images", lambda: self._async_sandbox.build_images(request)) def start( self, - spec: SandboxSpec, + spec: SandboxSpec | None = None, *, outside_endpoints: list[OutsideEndpoint] | None = None, - delete_on_stop: bool = False, - ) -> "SandboxInstance": - async_instance = self._runner.run( + delete_on_stop: bool | None = None, + ) -> "Sandbox": + self._runner.run( "start", lambda: self._async_sandbox.start( spec, @@ -414,29 +311,20 @@ def start( delete_on_stop=delete_on_stop, ), ) - return SandboxInstance(self, async_instance) - - def attach(self, sandbox_id: str) -> SandboxHandle: - return self._runner.run("attach", lambda: self._async_sandbox.attach(sandbox_id)) - - def connect(self, sandbox_id: str) -> SandboxHandle: - """Compatibility alias for ``attach``.""" - return self.attach(sandbox_id) + return self def exec( self, - handle: SandboxHandle, command: str, *, cwd: str | None = None, env: dict[str, str] | None = None, - timeout_s: int | float | None = None, + timeout_s: int | float | None = 180, user: str | int | None = None, ) -> SandboxExecResult: return self._runner.run( "exec", lambda: self._async_sandbox.exec( - handle, command, cwd=cwd, env=env, @@ -445,43 +333,34 @@ def exec( ), ) - def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: - self._runner.run("write_file", lambda: self._async_sandbox.write_file(handle, target_path, data)) - - def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: - return self._runner.run("read_file", lambda: self._async_sandbox.read_file(handle, source_path)) + def upload(self, local_path: Path | str, remote_path: str) -> None: + self._runner.run("upload", lambda: self._async_sandbox.upload(local_path, remote_path)) - def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: - self._runner.run("upload_file", lambda: self._async_sandbox.upload_file(handle, source_path, target_path)) + def download(self, remote_path: str, local_path: Path | str) -> None: + self._runner.run("download", lambda: self._async_sandbox.download(remote_path, local_path)) - def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: - self._runner.run("download_file", lambda: self._async_sandbox.download_file(handle, source_path, target_path)) + def status(self) -> SandboxStatus: + return self._runner.run("status", self._async_sandbox.status) - def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: - self._runner.run("close", lambda: self._async_sandbox.close(handle, delete=delete)) + @property + def is_running(self) -> bool: + return self.status() == SandboxStatus.RUNNING - def status(self, handle: SandboxHandle) -> SandboxStatus: - return self._runner.run("status", lambda: self._async_sandbox.status(handle)) + def container_ip(self) -> str | None: + return self._runner.run("container_ip", self._async_sandbox.container_ip) - def container_ip(self, handle: SandboxHandle) -> str | None: - return self._runner.run("container_ip", lambda: self._async_sandbox.container_ip(handle)) + def stop(self, *, delete: bool | None = None) -> None: + self._runner.run("stop", lambda: self._async_sandbox.stop(delete=delete)) def shutdown(self) -> None: - """Close provider-scoped resources such as SDK clients or warm pools.""" if self._closed: return self._closed = True try: - self._runner.run("shutdown", self._async_sandbox.shutdown) + self._runner.run("shutdown", self._async_sandbox.aclose) finally: self._runner.close() - def handle_reference(self, handle: SandboxHandle) -> Any: - return self._runner.run("handle_reference", lambda: self._async_sandbox.handle_reference(handle)) - - def materialize_handle(self, value: Any) -> SandboxHandle: - return self._runner.run("materialize_handle", lambda: self._async_sandbox.materialize_handle(value)) - def __enter__(self) -> "Sandbox": return self @@ -494,84 +373,3 @@ def __del__(self) -> None: # pragma: no cover self.shutdown() except Exception: pass - - -class SandboxInstance: - """Evaluator-style sync sandbox object returned by ``Sandbox.start``.""" - - def __init__(self, owner: Sandbox, async_instance: AsyncSandboxInstance) -> None: - self._owner = owner - self._async_instance = async_instance - - @property - def spec(self) -> SandboxSpec: - return self._owner._runner.call("spec", lambda: self._async_instance.spec) - - @property - def handle(self) -> SandboxHandle: - return self._owner._runner.call("handle", lambda: self._async_instance.handle) - - def exec( - self, - command: str, - timeout_sec: int | float | None = 180, - *, - cwd: str | None = None, - env: dict[str, str] | None = None, - user: str | int | None = None, - timeout_s: int | float | None = None, - ) -> SandboxExecResult: - return self._owner._runner.run( - "instance.exec", - lambda: self._async_instance.exec( - command, - timeout_sec=timeout_sec, - cwd=cwd, - env=env, - user=user, - timeout_s=timeout_s, - ), - ) - - def write_file(self, target_path: str, data: str | bytes) -> None: - self._owner._runner.run( - "instance.write_file", - lambda: self._async_instance.write_file(target_path, data), - ) - - def read_file(self, source_path: str) -> bytes: - return self._owner._runner.run("instance.read_file", lambda: self._async_instance.read_file(source_path)) - - def upload(self, local_path: Path | str, remote_path: str) -> None: - self._owner._runner.run( - "instance.upload", - lambda: self._async_instance.upload(local_path, remote_path), - ) - - def download(self, remote_path: str, local_path: Path | str) -> None: - self._owner._runner.run( - "instance.download", - lambda: self._async_instance.download(remote_path, local_path), - ) - - def status(self) -> SandboxStatus: - return self._owner._runner.run("instance.status", self._async_instance.status) - - @property - def is_running(self) -> bool: - return self.status() == SandboxStatus.RUNNING - - def container_ip(self) -> str | None: - return self._owner._runner.run("instance.container_ip", self._async_instance.container_ip) - - def stop(self, *, delete: bool | None = None) -> None: - self._owner._runner.run("instance.stop", lambda: self._async_instance.stop(delete=delete)) - - def close(self, *, delete: bool | None = None) -> None: - self.stop(delete=delete) - - def __enter__(self) -> "SandboxInstance": - return self - - def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: - self.stop() diff --git a/nemo_gym/sandbox/providers/__init__.py b/nemo_gym/sandbox/providers/__init__.py index 3e646735b5..7e2284b9a9 100644 --- a/nemo_gym/sandbox/providers/__init__.py +++ b/nemo_gym/sandbox/providers/__init__.py @@ -20,7 +20,6 @@ ImageSpec, OutsideEndpoint, SandboxAddressProvider, - SandboxAttachProvider, SandboxBatchCreateError, SandboxCreateError, SandboxCreateVerificationError, @@ -49,7 +48,6 @@ "ImageSpec", "OutsideEndpoint", "SandboxAddressProvider", - "SandboxAttachProvider", "SandboxBatchCreateError", "SandboxCreateError", "SandboxCreateVerificationError", diff --git a/nemo_gym/sandbox/providers/base.py b/nemo_gym/sandbox/providers/base.py index c2ea4ed7d1..be3f5301cd 100644 --- a/nemo_gym/sandbox/providers/base.py +++ b/nemo_gym/sandbox/providers/base.py @@ -229,15 +229,6 @@ def materialize_handle(self, value: Any) -> SandboxHandle | Awaitable[SandboxHan ... -@runtime_checkable -class SandboxAttachProvider(Protocol): - """Optional provider trait for attaching to an existing sandbox.""" - - async def attach(self, sandbox_id: str) -> SandboxHandle: - """Attach to an existing sandbox and return a loop-local handle.""" - ... - - @runtime_checkable class SandboxInlineFileProvider(Protocol): """Optional provider trait for efficient inline file reads and writes.""" diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 3467012c14..9afd93124c 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -1209,10 +1209,6 @@ async def connect(self, sandbox_id: str) -> SandboxHandle: sandbox = await Sandbox.connect(sandbox_id, **kwargs) return SandboxHandle(sandbox_id=str(sandbox.id), provider_name=self.name, raw=sandbox) - async def attach(self, sandbox_id: str) -> SandboxHandle: - """Attach to an existing OpenSandbox sandbox.""" - return await self.connect(sandbox_id) - async def status(self, handle: SandboxHandle) -> SandboxStatus: """Return the current OpenSandbox lifecycle status.""" get_info = getattr(handle.raw, "get_info", None) diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index 88e35c1981..55d265ebc6 100644 --- a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -74,9 +74,7 @@ def __init__( if not self.config.provider: raise ValueError("MiniSWESandboxEnvironment requires provider") - self._sandbox_client: Sandbox | None = None - self._sandbox: Any | None = None - self._handle: Any | None = None + self._sandbox: Sandbox | None = None self._closed = False spec_config = dict(self.config.spec) @@ -96,8 +94,7 @@ def __init__( env[key] = value env.update(self.config.env) - self._sandbox_client = Sandbox(self.config.provider) - self._sandbox = self._sandbox_client.start( + self._sandbox = Sandbox(self.config.provider).start( SandboxSpec( image=image, timeout_s=spec_config.pop("timeout_s", None), @@ -118,7 +115,6 @@ def __init__( ), delete_on_stop=self.config.delete, ) - self._handle = self._sandbox.handle def get_template_vars(self, **kwargs: Any) -> dict[str, Any]: return {**self.config.__dict__, **kwargs} @@ -155,10 +151,12 @@ def execute( command = action.get("command", "") if isinstance(action, dict) else action timeout_s = timeout or (self.config.eval_timeout if is_eval else self.config.step_timeout) exec_cwd = cwd or self.config.cwd + if self._sandbox is None: + raise RuntimeError("Sandbox is not available") result = self._sandbox.exec( self._command(command, exec_cwd), - timeout_sec=timeout_s, + timeout_s=timeout_s, cwd="/", user=self.config.user, ) @@ -188,15 +186,9 @@ def cleanup(self) -> None: if self._closed: return self._closed = True - try: - if self._sandbox is not None: - self._sandbox.stop() - self._sandbox = None - self._handle = None - finally: - if self._sandbox_client is not None: - self._sandbox_client.shutdown() - self._sandbox_client = None + if self._sandbox is not None: + self._sandbox.shutdown() + self._sandbox = None def __enter__(self) -> "MiniSWESandboxEnvironment": return self diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index a4d1e9d454..6cedb83961 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -14,6 +14,7 @@ # limitations under the License. import asyncio +import builtins from dataclasses import dataclass from datetime import timedelta from pathlib import Path @@ -99,6 +100,9 @@ def __init__(self, **kwargs: Any) -> None: type(self).received_kwargs = kwargs async def start(self) -> None: + preparer = self.received_kwargs.get("warmup_sandbox_preparer") + if preparer is not None: + await preparer(FakeSandbox("warmup-1")) return None async def snapshot(self) -> FakeSnapshot: @@ -188,6 +192,42 @@ class StatusCodeError(Exception): opensandbox_provider._log_create_retry(retry_state) +def test_missing_optional_dependency_import_helpers(monkeypatch: pytest.MonkeyPatch) -> None: + real_import = builtins.__import__ + + def block_imports(*blocked_names: str) -> None: + def fake_import( + name: str, + globals_: dict[str, Any] | None = None, + locals_: dict[str, Any] | None = None, + fromlist: tuple[str, ...] = (), + level: int = 0, + ) -> Any: + if any(name == blocked or name.startswith(f"{blocked}.") for blocked in blocked_names): + raise ModuleNotFoundError(name) + return real_import(name, globals_, locals_, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + block_imports("opensandbox") + with pytest.raises(ModuleNotFoundError, match="OpenSandbox SDK is required"): + opensandbox_provider._require_opensandbox_sdk() + + block_imports("opensandbox") + with pytest.raises(ModuleNotFoundError, match="OpenSandbox SDK >=0.1.9"): + opensandbox_provider._require_opensandbox_sdk_pool() + + block_imports("tenacity") + with pytest.raises(ModuleNotFoundError, match="tenacity is required"): + opensandbox_provider._require_tenacity() + + block_imports("httpx") + assert opensandbox_provider._httpx_retryable_types() == tuple() + + block_imports("opensandbox.exceptions") + assert opensandbox_provider._is_retryable_create_error(RuntimeError("gateway timeout")) is True + + async def test_provider_reference_materialization_and_conversion_helpers( fake_opensandbox_sdk: None, monkeypatch: pytest.MonkeyPatch, @@ -265,6 +305,32 @@ def test_provider_validation_and_retry_helpers() -> None: with pytest.raises(ValueError, match="image_pull_policy"): opensandbox_provider.validate_image_pull_policy("Sometimes") + volume_mapping = opensandbox_provider._volume_to_mapping( + VolumeMount(host_path="/host/workspace", container_path="/mnt/workspace", readonly=True), + 0, + ) + assert volume_mapping == { + "name": "mnt-workspace", + "host": {"path": "/host/workspace"}, + "mount_path": "/mnt/workspace", + "read_only": True, + } + assert opensandbox_provider._volume_to_mapping({"name": "raw-volume"}, 1) == {"name": "raw-volume"} + assert opensandbox_provider._volume_mount_name(VolumeMount(container_path="/"), 2) == "volume-2" + with pytest.raises(TypeError, match="VolumeMount"): + opensandbox_provider._volume_to_mapping(object(), 0) + with pytest.raises(ValueError, match="EFS"): + opensandbox_provider._volume_to_mapping(VolumeMount(efs_filesystem_id="fs-1"), 0) + with pytest.raises(ValueError, match="host_path"): + opensandbox_provider._volume_to_mapping(VolumeMount(), 0) + with pytest.raises(TypeError, match="must be a bool"): + opensandbox_provider._provider_option_bool({"skip_health_check": "true"}, "skip_health_check") + + assert opensandbox_provider._to_sandbox_status("starting") == SandboxStatus.STARTING + assert opensandbox_provider._to_sandbox_status("terminated") == SandboxStatus.STOPPED + assert opensandbox_provider._to_sandbox_status("failed") == SandboxStatus.ERROR + assert opensandbox_provider._to_sandbox_status(None) == SandboxStatus.UNKNOWN + invalid_kwargs = [ {"pool": {"concurrency": 0}}, {"connection": {"connect_timeout_s": 0}}, @@ -395,6 +461,20 @@ async def no_sleep(_seconds: float) -> None: ) == 1 ) + progress_timeout_provider = opensandbox_provider.OpenSandboxProvider( + pool={"progress_timeout_s": 0.000001, "acquire_poll_interval_s": 0.000001}, + probe={"command": None}, + ) + assert ( + await progress_timeout_provider._wait_sdk_pool_idle( + FakePool([1]), + spec=SandboxSpec(image="image:tag"), + requested=2, + timeout_s=1, + allow_partial=True, + ) + == 1 + ) with pytest.raises(opensandbox_provider.OpenSandboxCreateTimeoutError): await provider._wait_sdk_pool_idle( FakePool([0]), @@ -500,6 +580,9 @@ async def get_info(self) -> Any: assert download_path.read_bytes() == b"bytes:/remote/download.txt" assert await provider.status(handle) == SandboxStatus.RUNNING assert await provider.container_ip(handle) == "10.1.2.3" + bare_handle = opensandbox_provider.SandboxHandle(sandbox_id="sandbox-2", provider_name="opensandbox", raw=object()) + assert await provider.status(bare_handle) == SandboxStatus.UNKNOWN + assert await provider.container_ip(bare_handle) is None with pytest.raises(ValueError, match="count"): await provider._create_batch_sdk(SandboxSpec(image="image:tag"), 0) @@ -534,6 +617,31 @@ async def no_sleep(_seconds: float) -> None: with pytest.raises(opensandbox_provider.OpenSandboxCreateVerificationError): await provider._verify_created_handle(handle) + provider = opensandbox_provider.OpenSandboxProvider( + probe={"command": "probe", "expected_stdout": None, "stable_count": 2, "stable_delay_s": 0.01}, + ) + sleep_calls: list[float] = [] + + async def record_sleep(seconds: float) -> None: + sleep_calls.append(seconds) + + async def good_probe(*_args: Any, **_kwargs: Any) -> opensandbox_provider.SandboxExecResult: + return opensandbox_provider.SandboxExecResult(stdout="ready", stderr=None, return_code=0) + + monkeypatch.setattr(opensandbox_provider.asyncio, "sleep", record_sleep) + monkeypatch.setattr(provider, "_exec", good_probe) + await provider._verify_created_handle(handle) + assert sleep_calls == [0.01] + + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": "probe"}) + + async def cancelled_probe(*_args: Any, **_kwargs: Any) -> opensandbox_provider.SandboxExecResult: + raise asyncio.CancelledError() + + monkeypatch.setattr(provider, "_exec", cancelled_probe) + with pytest.raises(asyncio.CancelledError): + await provider._verify_created_handle(handle) + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": "probe"}) async def fail_verify(_handle: Any) -> None: @@ -570,6 +678,21 @@ async def close(self) -> None: delete=True, ) + class CloseSucceedsRaw: + async def close(self) -> None: + return None + + assert await provider._close_many( + [ + opensandbox_provider.SandboxHandle( + sandbox_id="sandbox-close", + provider_name="opensandbox", + raw=CloseSucceedsRaw(), + ) + ], + delete=False, + ) == [None] + class DeleteAndCloseFailRaw: async def kill(self) -> None: raise RuntimeError("delete failed") @@ -673,6 +796,47 @@ async def no_sleep(_seconds: float) -> None: SandboxSpec(image="image:tag"), ) + class CancelledConnectSandbox(FakeSandbox): + @classmethod + async def connect(cls, *args: Any, **kwargs: Any) -> "FakeSandbox": + del args, kwargs + raise asyncio.CancelledError() + + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk", + lambda: (CancelledConnectSandbox, FakeConnectionConfig, object, FakePlatformSpec, object), + ) + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + with pytest.raises(asyncio.CancelledError): + await provider._connect_after_create( + opensandbox_provider.SandboxHandle(sandbox_id="sandbox-1", provider_name="opensandbox", raw=None), + SandboxSpec(image="image:tag", ready_timeout_s=1), + ) + + class NonRetryableConnectSandbox(FakeSandbox): + @classmethod + async def connect(cls, *args: Any, **kwargs: Any) -> "FakeSandbox": + del args, kwargs + raise ValueError("bad connection request") + + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk", + lambda: (NonRetryableConnectSandbox, FakeConnectionConfig, object, FakePlatformSpec, object), + ) + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + with pytest.raises(ValueError, match="bad connection request"): + await provider._connect_after_create( + opensandbox_provider.SandboxHandle(sandbox_id="sandbox-1", provider_name="opensandbox", raw=None), + SandboxSpec(image="image:tag", ready_timeout_s=1), + ) + + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk", + lambda: (FakeSandbox, FakeConnectionConfig, object, FakePlatformSpec, object), + ) provider = opensandbox_provider.OpenSandboxProvider( connection={"request_timeout_s": 3}, probe={"command": None}, diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index b48e4b0cec..fb3ea80450 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -38,6 +38,7 @@ list_providers, register_provider, ) +from nemo_gym.sandbox.api import _AsyncLoopRunner from nemo_gym.sandbox.utils import rewrite_image from responses_api_agents.mini_swe_agent_2.sandbox_environment import MiniSWESandboxEnvironment @@ -109,9 +110,6 @@ async def create_batch( del allow_partial return [await self.create(spec) for _ in range(count)] - async def connect(self, sandbox_id: str) -> SandboxHandle: - return SandboxHandle(sandbox_id=sandbox_id, provider_name=self.name, raw={}) - async def exec( self, handle: SandboxHandle, @@ -163,12 +161,6 @@ async def close(self, handle: SandboxHandle, *, delete: bool) -> None: async def aclose(self) -> None: self.aclosed = True - def handle_reference(self, handle: SandboxHandle) -> dict[str, str]: - return {"kind": "fake", "sandbox_id": handle.sandbox_id} - - async def materialize_handle(self, value: Any) -> SandboxHandle: - return SandboxHandle(sandbox_id=value["sandbox_id"], provider_name=self.name, raw={"materialized": True}) - class PlainSandboxProvider: name = "plain" @@ -186,9 +178,6 @@ async def create_batch( del allow_partial return [await self.create(spec) for _ in range(count)] - async def connect(self, sandbox_id: str) -> SandboxHandle: - return SandboxHandle(sandbox_id=sandbox_id, provider_name=self.name, raw={}) - async def exec( self, handle: SandboxHandle, @@ -202,13 +191,6 @@ async def exec( del handle, command, cwd, env, timeout_s, user return SandboxExecResult(stdout="ok", stderr=None, return_code=0) - async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: - del handle, target_path, data - - async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: - del handle - return f"read:{source_path}".encode() - async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: del handle, source_path, target_path @@ -270,24 +252,48 @@ async def aclose(self) -> None: return None -def test_sandbox_facade_uses_public_provider_api() -> None: - asyncio.run(_assert_sandbox_facade_uses_public_provider_api()) +class EmptyImageBuildProvider(FakeSandboxProvider): + async def build_images(self, request: ImageBuildRequest) -> list[str]: + self.image_build_requests.append(request) + return [] + + +class FailingWriteProvider(FakeSandboxProvider): + async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + self.write_calls.append((handle, target_path, data)) + raise RuntimeError("write failed") + + +def test_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: + asyncio.run(_assert_sandbox_facade_uses_public_provider_api(tmp_path)) -async def _assert_sandbox_facade_uses_public_provider_api() -> None: +async def _assert_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: provider_name = f"fake-{uuid4().hex}" register_provider(provider_name, FakeSandboxProvider) sandbox = AsyncSandbox({provider_name: {"marker": "configured"}}) - handle = await sandbox.create(SandboxSpec(image="image:tag", metadata={"suite": "unit"})) + await sandbox.start( + SandboxSpec( + image="image:tag", + metadata={"suite": "unit"}, + workdir="/repo", + files={"/tmp/bootstrap.txt": "hello"}, + ), + outside_endpoints=[OutsideEndpoint(url="http://outside", env_var="OUTSIDE_URL")], + delete_on_stop=True, + ) + handle = sandbox.handle provider = FakeSandboxProvider.last_instance assert provider is not None assert provider.marker == "configured" assert provider.created_specs[0].image == "image:tag" assert provider.created_specs[0].metadata == {"suite": "unit"} + assert sandbox.spec.env["OUTSIDE_URL"] == "http://outside" + assert provider.write_calls == [(handle, "/tmp/bootstrap.txt", "hello")] - result = await sandbox.exec(handle, "pytest -q", cwd="/repo", timeout_s=60, user="agent") + result = await sandbox.exec("pytest -q", timeout_s=60, user="agent") assert result == SandboxExecResult(stdout="ok", stderr=None, return_code=0) assert provider.exec_calls[0] == { "handle": handle, @@ -297,51 +303,80 @@ async def _assert_sandbox_facade_uses_public_provider_api() -> None: "timeout_s": 60, "user": "agent", } - assert await sandbox.status(handle) == SandboxStatus.RUNNING - assert await sandbox.container_ip(handle) == "10.0.0.1" + assert await sandbox.status() == SandboxStatus.RUNNING + assert await sandbox.is_running() is True + assert await sandbox.container_ip() == "10.0.0.1" - built_handle = await sandbox.create( - SandboxSpec(image_build=ImageSpec(image="built:tag", source={"context": "repo"})) - ) - assert built_handle.sandbox_id == "fake-1" + source_path = tmp_path / "source.txt" + target_path = tmp_path / "nested" / "target.txt" + source_path.write_text("local", encoding="utf-8") + await sandbox.upload(source_path, "/remote/source.txt") + await sandbox.download("/remote/source.txt", target_path) + assert provider.upload_calls == [(handle, source_path, "/remote/source.txt")] + assert provider.download_calls == [(handle, "/remote/source.txt", target_path)] + assert target_path.read_bytes() == b"downloaded" + + await sandbox.stop() + await sandbox.stop() + assert provider.closed[-1] == (handle, True) + assert await sandbox.status() == SandboxStatus.STOPPED + + built_sandbox = AsyncSandbox(provider) + await built_sandbox.start(SandboxSpec(image_build=ImageSpec(image="built:tag", source={"context": "repo"}))) assert provider.image_build_requests[-1].specs[0].source == {"context": "repo"} assert provider.created_specs[-1].image == "built:tag" + await built_sandbox.stop(delete=True) - session = await sandbox.start( - SandboxSpec( - image="image:tag", - workdir="/session", - files={"/tmp/bootstrap.txt": "hello"}, - ), - outside_endpoints=[OutsideEndpoint(url="http://outside", env_var="OUTSIDE_URL")], - delete_on_stop=True, - ) - try: - assert session.spec.env["OUTSIDE_URL"] == "http://outside" - assert await session.is_running() is True - assert await session.container_ip() == "10.0.0.1" - session_result = await session.exec("pwd") - assert session_result.return_code == 0 - assert provider.exec_calls[-1]["cwd"] == "/session" - finally: - await session.stop() - assert provider.write_calls[-1] == (session.handle, "/tmp/bootstrap.txt", "hello") - assert provider.closed[-1] == (session.handle, True) - - await sandbox.close(handle, delete=True) - assert provider.closed[-1] == (handle, True) - assert await sandbox.handle_reference(handle) == {"kind": "fake", "sandbox_id": "fake-1"} - assert await sandbox.materialize_handle({"sandbox_id": "fake-2"}) == SandboxHandle( - sandbox_id="fake-2", provider_name="fake", raw={"materialized": True} - ) async with AsyncSandbox(provider) as context_sandbox: assert context_sandbox.provider_name == "fake" - await sandbox.shutdown() + await context_sandbox.start(SandboxSpec(image="image:tag"), delete_on_stop=True) + context_handle = context_sandbox.handle + assert provider.closed[-1] == (context_handle, True) + + await sandbox.aclose() assert provider.aclosed is True -def test_rewrite_image_and_materialize_handle_validation() -> None: - asyncio.run(_assert_rewrite_image_and_materialize_handle_validation()) +def test_async_sandbox_build_and_initial_file_error_paths() -> None: + asyncio.run(_assert_async_sandbox_build_and_initial_file_error_paths()) + + +async def _assert_async_sandbox_build_and_initial_file_error_paths() -> None: + empty_build_provider = EmptyImageBuildProvider() + empty_build_sandbox = AsyncSandbox(empty_build_provider) + with pytest.raises(ValueError, match="build_images returned no image references"): + await empty_build_sandbox.start(SandboxSpec(image_build=ImageSpec(image="missing:tag"))) + assert empty_build_provider.image_build_requests[0].specs[0].image == "missing:tag" + + failing_provider = FailingWriteProvider() + failing_sandbox = AsyncSandbox(failing_provider) + with pytest.raises(RuntimeError, match="write failed"): + await failing_sandbox.start(SandboxSpec(image="image:tag", files={"/tmp/bootstrap.txt": "hello"})) + assert failing_provider.closed == [ + ( + SandboxHandle( + sandbox_id="fake-1", + provider_name="fake", + raw={"spec": SandboxSpec(image="image:tag", files={"/tmp/bootstrap.txt": "hello"})}, + ), + True, + ) + ] + + unstarted = AsyncSandbox(FakeSandboxProvider()) + with pytest.raises(RuntimeError, match="not been started"): + await unstarted.exec("pwd") + + started = AsyncSandbox(FakeSandboxProvider()) + await started.start(SandboxSpec(image="image:tag")) + with pytest.raises(RuntimeError, match="already started"): + await started.start(SandboxSpec(image="image:tag")) + await started.stop() + + +def test_rewrite_image_validation() -> None: + assert rewrite_image(None, []) is None + assert rewrite_image("image:tag", [{"from": "other/", "to": "mirror/"}]) == "image:tag" def test_provider_registry_validation_and_listing(monkeypatch: pytest.MonkeyPatch) -> None: @@ -393,88 +428,59 @@ def __init__(self) -> None: Sandbox({failing_provider_name: {}}) -async def _assert_rewrite_image_and_materialize_handle_validation() -> None: - assert rewrite_image(None, []) is None - assert rewrite_image("image:tag", [{"from": "other/", "to": "mirror/"}]) == "image:tag" - - class BadMaterializeProvider(FakeSandboxProvider): - async def materialize_handle(self, value: Any) -> object: - del value - return object() - - sandbox = AsyncSandbox(BadMaterializeProvider()) - try: - await sandbox.materialize_handle({"sandbox_id": "bad"}) - except TypeError as e: - assert "must return SandboxHandle" in str(e) - else: - raise AssertionError("expected invalid materialize_handle return type to fail") - - -def test_async_sandbox_batch_file_and_fallback_reference_operations(tmp_path: Path) -> None: - asyncio.run(_assert_async_sandbox_batch_file_and_fallback_reference_operations(tmp_path)) - +def test_async_sandbox_transfer_fallback_and_unknown_status(tmp_path: Path) -> None: + asyncio.run(_assert_async_sandbox_transfer_fallback_and_unknown_status(tmp_path)) -async def _assert_async_sandbox_batch_file_and_fallback_reference_operations(tmp_path: Path) -> None: - provider = FakeSandboxProvider() - sandbox = AsyncSandbox(provider) - - handles = await sandbox.create_batch(SandboxSpec(image="image:tag"), 2, allow_partial=True) - connected = await sandbox.connect("connected-1") - await sandbox.write_file(connected, "/tmp/file.txt", "contents") - assert await sandbox.read_file(connected, "/tmp/file.txt") == b"read:/tmp/file.txt" - source_path = tmp_path / "source.txt" - target_path = tmp_path / "nested" / "target.txt" - source_path.write_text("local", encoding="utf-8") - await sandbox.upload_file(connected, source_path, "/remote/source.txt") - await sandbox.download_file(connected, "/remote/source.txt", target_path) - await sandbox.close(connected) - - assert [handle.sandbox_id for handle in handles] == ["fake-1", "fake-1"] - assert provider.write_calls == [(connected, "/tmp/file.txt", "contents")] - assert provider.read_calls == [(connected, "/tmp/file.txt")] - assert provider.upload_calls == [(connected, source_path, "/remote/source.txt")] - assert provider.download_calls == [(connected, "/remote/source.txt", target_path)] - assert target_path.read_bytes() == b"downloaded" +async def _assert_async_sandbox_transfer_fallback_and_unknown_status(tmp_path: Path) -> None: transfer_provider = TransferOnlySandboxProvider() transfer_sandbox = AsyncSandbox(transfer_provider) - transfer_handle = SandboxHandle(sandbox_id="transfer-1", provider_name="transfer-only", raw={}) - await transfer_sandbox.write_file(transfer_handle, "/remote/inline.txt", b"fallback") + await transfer_sandbox.start(SandboxSpec(image="image:tag", files={"/remote/inline.txt": "fallback"})) + transfer_handle = transfer_sandbox.handle assert transfer_provider.upload_calls[0][0] == transfer_handle assert transfer_provider.upload_calls[0][2] == "/remote/inline.txt" - assert await transfer_sandbox.read_file(transfer_handle, "/remote/inline.txt") == b"fallback" - assert transfer_provider.download_calls == [ - (transfer_handle, "/remote/inline.txt", transfer_provider.download_calls[0][2]) - ] + source_path = tmp_path / "source.txt" + target_path = tmp_path / "target.txt" + source_path.write_text("local", encoding="utf-8") + await transfer_sandbox.upload(source_path, "/remote/source.txt") + await transfer_sandbox.download("/remote/inline.txt", target_path) + assert transfer_provider.upload_calls[1] == (transfer_handle, source_path, "/remote/source.txt") + assert transfer_provider.download_calls == [(transfer_handle, "/remote/inline.txt", target_path)] + assert target_path.read_bytes() == b"fallback" plain_provider = PlainSandboxProvider() plain_sandbox = AsyncSandbox(plain_provider) - plain_handle = SandboxHandle(sandbox_id="plain-1", provider_name="plain", raw={}) - assert await plain_sandbox.handle_reference(plain_handle) is plain_handle - assert await plain_sandbox.materialize_handle(plain_handle) is plain_handle - try: - await plain_sandbox.materialize_handle({"sandbox_id": "plain-2"}) - except ValueError as e: - assert "cannot materialize" in str(e) - else: - raise AssertionError("expected materialize_handle without provider support to fail") + await plain_sandbox.start(SandboxSpec(image="image:tag")) + assert await plain_sandbox.status() == SandboxStatus.UNKNOWN + assert await plain_sandbox.container_ip() is None -def test_sync_sandbox_facade_uses_public_provider_api() -> None: +def test_sync_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: provider_name = f"fake-{uuid4().hex}" register_provider(provider_name, FakeSandboxProvider) with Sandbox({provider_name: {"marker": "configured"}}) as sandbox: - handle = sandbox.create(SandboxSpec(image="image:tag", metadata={"suite": "unit"})) + sandbox.start( + SandboxSpec( + image="image:tag", + metadata={"suite": "unit"}, + workdir="/repo", + files={"/tmp/bootstrap.txt": "hello"}, + ), + outside_endpoints=[OutsideEndpoint(url="http://outside", env_var="OUTSIDE_URL")], + delete_on_stop=True, + ) + handle = sandbox.handle provider = FakeSandboxProvider.last_instance assert provider is not None assert provider.marker == "configured" assert provider.created_specs[0].image == "image:tag" assert provider.created_specs[0].metadata == {"suite": "unit"} + assert sandbox.spec.env["OUTSIDE_URL"] == "http://outside" + assert provider.write_calls == [(handle, "/tmp/bootstrap.txt", "hello")] - result = sandbox.exec(handle, "pytest -q", cwd="/repo", timeout_s=60, user="agent") + result = sandbox.exec("pytest -q", timeout_s=60, user="agent") assert result == SandboxExecResult(stdout="ok", stderr=None, return_code=0) assert provider.exec_calls[0] == { "handle": handle, @@ -484,20 +490,23 @@ def test_sync_sandbox_facade_uses_public_provider_api() -> None: "timeout_s": 60, "user": "agent", } - - sandbox.close(handle, delete=True) - assert provider.closed[0] == (handle, True) - assert sandbox.handle_reference(handle) == {"kind": "fake", "sandbox_id": "fake-1"} - assert sandbox.materialize_handle({"sandbox_id": "fake-3"}).sandbox_id == "fake-3" + assert sandbox.build_images(ImageBuildRequest(specs=[ImageSpec(image="sync-built:tag")])) == ["sync-built:tag"] + assert sandbox.status() == SandboxStatus.RUNNING + assert sandbox.is_running is True + assert sandbox.container_ip() == "10.0.0.1" + + upload_path = tmp_path / "sync-upload.txt" + upload_path.write_text("sync", encoding="utf-8") + download_path = tmp_path / "sync-download.txt" + sandbox.upload(upload_path, "/tmp/sync-upload.txt") + sandbox.download("/tmp/sync-download.txt", download_path) + assert download_path.read_bytes() == b"downloaded" assert sandbox.provider_name == "fake" - assert len(sandbox.create_batch(SandboxSpec(image="image:tag"), 2)) == 2 - session = sandbox.start(SandboxSpec(image="image:tag", workdir="/sync-session"), delete_on_stop=True) - assert session.is_running is True - assert session.container_ip() == "10.0.0.1" - assert session.exec("pwd").return_code == 0 + sandbox.stop() + assert provider.closed[-1] == (handle, True) + sandbox.start(SandboxSpec(image="image:tag", workdir="/sync-session"), delete_on_stop=True) + assert sandbox.exec("pwd").return_code == 0 assert provider.exec_calls[-1]["cwd"] == "/sync-session" - session.stop() - assert provider.closed[-1] == (session.handle, True) sandbox.shutdown() sandbox.shutdown() assert provider.aclosed is True @@ -509,20 +518,23 @@ def test_sync_sandbox_facade_uses_public_provider_api() -> None: raise AssertionError("expected closed sync sandbox to reject further calls") +def test_sync_loop_runner_close_is_idempotent() -> None: + runner = _AsyncLoopRunner() + runner.close() + runner.close() + + def test_sync_sandbox_file_operations(tmp_path: Path) -> None: provider = FakeSandboxProvider() with Sandbox(provider) as sandbox: - handle = sandbox.connect("sync-1") - sandbox.write_file(handle, "/tmp/file.txt", b"contents") - assert sandbox.read_file(handle, "/tmp/file.txt") == b"read:/tmp/file.txt" + sandbox.start(SandboxSpec(image="image:tag")) + handle = sandbox.handle source_path = tmp_path / "source.txt" target_path = tmp_path / "target.txt" source_path.write_text("local", encoding="utf-8") - sandbox.upload_file(handle, source_path, "/remote/source.txt") - sandbox.download_file(handle, "/remote/source.txt", target_path) + sandbox.upload(source_path, "/remote/source.txt") + sandbox.download("/remote/source.txt", target_path) - assert provider.write_calls == [(handle, "/tmp/file.txt", b"contents")] - assert provider.read_calls == [(handle, "/tmp/file.txt")] assert provider.upload_calls == [(handle, source_path, "/remote/source.txt")] assert provider.download_calls == [(handle, "/remote/source.txt", target_path)] assert target_path.read_bytes() == b"downloaded" @@ -1006,7 +1018,8 @@ def test_mini_swe_sandbox_environment_validation_and_context_manager() -> None: provider={provider_name: {}}, delete=False, ) as env: - assert env._handle is not None + assert env._sandbox is not None + assert env._sandbox.handle.sandbox_id == "fake-1" assert FakeSandboxProvider.last_instance is not None assert FakeSandboxProvider.last_instance.closed[-1][1] is False From 22debc6a6545f5cd5d0ae667a2f41fa32244b57f Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Tue, 2 Jun 2026 16:09:10 -0700 Subject: [PATCH 06/14] Minimize sandbox public API Signed-off-by: Hemil Desai --- nemo_gym/sandbox/api.py | 87 +++++-------------- .../mini_swe_agent_2/sandbox_environment.py | 2 +- tests/unit_tests/test_sandbox.py | 80 ++++++++++------- 3 files changed, 72 insertions(+), 97 deletions(-) diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index 032256a2c6..a8132267ea 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -26,7 +26,6 @@ from nemo_gym.sandbox.providers import ( ImageBuildRequest, OutsideEndpoint, - SandboxAddressProvider, SandboxExecResult, SandboxHandle, SandboxImageBuildProvider, @@ -57,35 +56,25 @@ def __init__( self._handle: SandboxHandle | None = None self._delete_on_stop = delete_on_stop self._stopped = True + self._closed = False - @property - def provider_name(self) -> str: + def _provider_name(self) -> str: return self._provider.name - @property - def spec(self) -> SandboxSpec: - if self._spec is None: - raise RuntimeError("Sandbox has not been configured") - return self._spec - - @property - def handle(self) -> SandboxHandle: - return self._require_handle() - def _require_handle(self) -> SandboxHandle: if self._handle is None or self._stopped: raise RuntimeError("Sandbox has not been started") return self._handle - async def build_images(self, request: ImageBuildRequest) -> list[str]: + async def _build_images(self, request: ImageBuildRequest) -> list[str]: if not isinstance(self._provider, SandboxImageBuildProvider): - raise NotImplementedError(f"Provider {self.provider_name!r} does not support sandbox image builds") + raise NotImplementedError(f"Provider {self._provider_name()!r} does not support sandbox image builds") return await self._provider.build_images(request) async def _resolve_image_build(self, spec: SandboxSpec) -> SandboxSpec: if spec.image_build is None: return spec - built_images = await self.build_images(ImageBuildRequest(specs=[spec.image_build])) + built_images = await self._build_images(ImageBuildRequest(specs=[spec.image_build])) if not built_images: raise ValueError("build_images returned no image references") return replace(spec, image=spec.image or built_images[0]) @@ -123,6 +112,8 @@ async def start( outside_endpoints: list[OutsideEndpoint] | None = None, delete_on_stop: bool | None = None, ) -> "AsyncSandbox": + if self._closed: + raise RuntimeError("Sandbox has been stopped") if self._handle is not None and not self._stopped: raise RuntimeError("Sandbox is already started") requested_spec = spec if spec is not None else self._spec @@ -136,6 +127,8 @@ async def start( await self._write_initial_files(handle, resolved_spec.files) except Exception: await self._provider.close(handle, delete=True) + await self._provider.aclose() + self._closed = True raise self._spec = resolved_spec @@ -156,7 +149,7 @@ async def exec( return await self._provider.exec( self._require_handle(), command, - cwd=cwd if cwd is not None else self.spec.workdir, + cwd=cwd if cwd is not None else self._spec.workdir if self._spec is not None else None, env=env, timeout_s=timeout_s, user=user, @@ -177,34 +170,25 @@ async def status(self) -> SandboxStatus: return SandboxStatus.UNKNOWN return await self._provider.status(self._handle) - async def is_running(self) -> bool: - return await self.status() == SandboxStatus.RUNNING - - async def container_ip(self) -> str | None: - if not isinstance(self._provider, SandboxAddressProvider): - return None - return await self._provider.container_ip(self._require_handle()) - async def stop(self, *, delete: bool | None = None) -> None: - if self._handle is None or self._stopped: + if self._closed: return - self._stopped = True - await self._provider.close( - self._handle, - delete=self._delete_on_stop if delete is None else delete, - ) - - async def aclose(self) -> None: try: - await self.stop() + if self._handle is not None and not self._stopped: + self._stopped = True + await self._provider.close( + self._handle, + delete=self._delete_on_stop if delete is None else delete, + ) finally: await self._provider.aclose() + self._closed = True async def __aenter__(self) -> "AsyncSandbox": return self async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: - await self.aclose() + await self.stop() class _AsyncLoopRunner: @@ -281,21 +265,6 @@ def __init__( raise self._closed = False - @property - def provider_name(self) -> str: - return self._runner.call("provider_name", lambda: self._async_sandbox.provider_name) - - @property - def spec(self) -> SandboxSpec: - return self._runner.call("spec", lambda: self._async_sandbox.spec) - - @property - def handle(self) -> SandboxHandle: - return self._runner.call("handle", lambda: self._async_sandbox.handle) - - def build_images(self, request: ImageBuildRequest) -> list[str]: - return self._runner.run("build_images", lambda: self._async_sandbox.build_images(request)) - def start( self, spec: SandboxSpec | None = None, @@ -340,24 +309,16 @@ def download(self, remote_path: str, local_path: Path | str) -> None: self._runner.run("download", lambda: self._async_sandbox.download(remote_path, local_path)) def status(self) -> SandboxStatus: + if self._closed: + return SandboxStatus.STOPPED return self._runner.run("status", self._async_sandbox.status) - @property - def is_running(self) -> bool: - return self.status() == SandboxStatus.RUNNING - - def container_ip(self) -> str | None: - return self._runner.run("container_ip", self._async_sandbox.container_ip) - def stop(self, *, delete: bool | None = None) -> None: - self._runner.run("stop", lambda: self._async_sandbox.stop(delete=delete)) - - def shutdown(self) -> None: if self._closed: return self._closed = True try: - self._runner.run("shutdown", self._async_sandbox.aclose) + self._runner.run("stop", lambda: self._async_sandbox.stop(delete=delete)) finally: self._runner.close() @@ -365,11 +326,11 @@ def __enter__(self) -> "Sandbox": return self def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: - self.shutdown() + self.stop() def __del__(self) -> None: # pragma: no cover if hasattr(self, "_closed") and not self._closed: try: - self.shutdown() + self.stop() except Exception: pass diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index 55d265ebc6..e6d43d556d 100644 --- a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -187,7 +187,7 @@ def cleanup(self) -> None: return self._closed = True if self._sandbox is not None: - self._sandbox.shutdown() + self._sandbox.stop() self._sandbox = None def __enter__(self) -> "MiniSWESandboxEnvironment": diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index fb3ea80450..557d21e9cd 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -82,6 +82,7 @@ class FakeSandboxProvider: def __init__(self, marker: str = "default") -> None: self.marker = marker self.created_specs: list[SandboxSpec] = [] + self.created_handles: list[SandboxHandle] = [] self.exec_calls: list[dict[str, Any]] = [] self.image_build_requests: list[ImageBuildRequest] = [] self.write_calls: list[tuple[SandboxHandle, str, str | bytes]] = [] @@ -98,7 +99,13 @@ async def build_images(self, request: ImageBuildRequest) -> list[str]: async def create(self, spec: SandboxSpec) -> SandboxHandle: self.created_specs.append(spec) - return SandboxHandle(sandbox_id="fake-1", provider_name=self.name, raw={"spec": spec}) + handle = SandboxHandle( + sandbox_id=f"fake-{len(self.created_handles) + 1}", + provider_name=self.name, + raw={"spec": spec}, + ) + self.created_handles.append(handle) + return handle async def create_batch( self, @@ -165,8 +172,17 @@ async def aclose(self) -> None: class PlainSandboxProvider: name = "plain" + def __init__(self) -> None: + self.created_handles: list[SandboxHandle] = [] + async def create(self, spec: SandboxSpec) -> SandboxHandle: - return SandboxHandle(sandbox_id="plain-1", provider_name=self.name, raw={"spec": spec}) + handle = SandboxHandle( + sandbox_id=f"plain-{len(self.created_handles) + 1}", + provider_name=self.name, + raw={"spec": spec}, + ) + self.created_handles.append(handle) + return handle async def create_batch( self, @@ -209,11 +225,18 @@ class TransferOnlySandboxProvider: name = "transfer-only" def __init__(self) -> None: + self.created_handles: list[SandboxHandle] = [] self.upload_calls: list[tuple[SandboxHandle, Path, str]] = [] self.download_calls: list[tuple[SandboxHandle, str, Path]] = [] async def create(self, spec: SandboxSpec) -> SandboxHandle: - return SandboxHandle(sandbox_id="transfer-1", provider_name=self.name, raw={"spec": spec}) + handle = SandboxHandle( + sandbox_id=f"transfer-{len(self.created_handles) + 1}", + provider_name=self.name, + raw={"spec": spec}, + ) + self.created_handles.append(handle) + return handle async def create_batch( self, @@ -283,14 +306,14 @@ async def _assert_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> Non outside_endpoints=[OutsideEndpoint(url="http://outside", env_var="OUTSIDE_URL")], delete_on_stop=True, ) - handle = sandbox.handle provider = FakeSandboxProvider.last_instance assert provider is not None + handle = provider.created_handles[0] assert provider.marker == "configured" assert provider.created_specs[0].image == "image:tag" assert provider.created_specs[0].metadata == {"suite": "unit"} - assert sandbox.spec.env["OUTSIDE_URL"] == "http://outside" + assert provider.created_specs[0].env["OUTSIDE_URL"] == "http://outside" assert provider.write_calls == [(handle, "/tmp/bootstrap.txt", "hello")] result = await sandbox.exec("pytest -q", timeout_s=60, user="agent") @@ -304,8 +327,6 @@ async def _assert_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> Non "user": "agent", } assert await sandbox.status() == SandboxStatus.RUNNING - assert await sandbox.is_running() is True - assert await sandbox.container_ip() == "10.0.0.1" source_path = tmp_path / "source.txt" target_path = tmp_path / "nested" / "target.txt" @@ -320,21 +341,20 @@ async def _assert_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> Non await sandbox.stop() assert provider.closed[-1] == (handle, True) assert await sandbox.status() == SandboxStatus.STOPPED + assert provider.aclosed is True - built_sandbox = AsyncSandbox(provider) + build_provider = FakeSandboxProvider() + built_sandbox = AsyncSandbox(build_provider) await built_sandbox.start(SandboxSpec(image_build=ImageSpec(image="built:tag", source={"context": "repo"}))) - assert provider.image_build_requests[-1].specs[0].source == {"context": "repo"} - assert provider.created_specs[-1].image == "built:tag" + assert build_provider.image_build_requests[-1].specs[0].source == {"context": "repo"} + assert build_provider.created_specs[-1].image == "built:tag" await built_sandbox.stop(delete=True) - async with AsyncSandbox(provider) as context_sandbox: - assert context_sandbox.provider_name == "fake" + context_provider = FakeSandboxProvider() + async with AsyncSandbox(context_provider) as context_sandbox: await context_sandbox.start(SandboxSpec(image="image:tag"), delete_on_stop=True) - context_handle = context_sandbox.handle - assert provider.closed[-1] == (context_handle, True) - - await sandbox.aclose() - assert provider.aclosed is True + context_handle = context_provider.created_handles[0] + assert context_provider.closed[-1] == (context_handle, True) def test_async_sandbox_build_and_initial_file_error_paths() -> None: @@ -372,6 +392,8 @@ async def _assert_async_sandbox_build_and_initial_file_error_paths() -> None: with pytest.raises(RuntimeError, match="already started"): await started.start(SandboxSpec(image="image:tag")) await started.stop() + with pytest.raises(RuntimeError, match="has been stopped"): + await started.start(SandboxSpec(image="image:tag")) def test_rewrite_image_validation() -> None: @@ -436,7 +458,7 @@ async def _assert_async_sandbox_transfer_fallback_and_unknown_status(tmp_path: P transfer_provider = TransferOnlySandboxProvider() transfer_sandbox = AsyncSandbox(transfer_provider) await transfer_sandbox.start(SandboxSpec(image="image:tag", files={"/remote/inline.txt": "fallback"})) - transfer_handle = transfer_sandbox.handle + transfer_handle = transfer_provider.created_handles[0] assert transfer_provider.upload_calls[0][0] == transfer_handle assert transfer_provider.upload_calls[0][2] == "/remote/inline.txt" source_path = tmp_path / "source.txt" @@ -452,7 +474,6 @@ async def _assert_async_sandbox_transfer_fallback_and_unknown_status(tmp_path: P plain_sandbox = AsyncSandbox(plain_provider) await plain_sandbox.start(SandboxSpec(image="image:tag")) assert await plain_sandbox.status() == SandboxStatus.UNKNOWN - assert await plain_sandbox.container_ip() is None def test_sync_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: @@ -470,14 +491,14 @@ def test_sync_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: outside_endpoints=[OutsideEndpoint(url="http://outside", env_var="OUTSIDE_URL")], delete_on_stop=True, ) - handle = sandbox.handle provider = FakeSandboxProvider.last_instance assert provider is not None + handle = provider.created_handles[0] assert provider.marker == "configured" assert provider.created_specs[0].image == "image:tag" assert provider.created_specs[0].metadata == {"suite": "unit"} - assert sandbox.spec.env["OUTSIDE_URL"] == "http://outside" + assert provider.created_specs[0].env["OUTSIDE_URL"] == "http://outside" assert provider.write_calls == [(handle, "/tmp/bootstrap.txt", "hello")] result = sandbox.exec("pytest -q", timeout_s=60, user="agent") @@ -490,10 +511,7 @@ def test_sync_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: "timeout_s": 60, "user": "agent", } - assert sandbox.build_images(ImageBuildRequest(specs=[ImageSpec(image="sync-built:tag")])) == ["sync-built:tag"] assert sandbox.status() == SandboxStatus.RUNNING - assert sandbox.is_running is True - assert sandbox.container_ip() == "10.0.0.1" upload_path = tmp_path / "sync-upload.txt" upload_path.write_text("sync", encoding="utf-8") @@ -501,17 +519,12 @@ def test_sync_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: sandbox.upload(upload_path, "/tmp/sync-upload.txt") sandbox.download("/tmp/sync-download.txt", download_path) assert download_path.read_bytes() == b"downloaded" - assert sandbox.provider_name == "fake" sandbox.stop() assert provider.closed[-1] == (handle, True) - sandbox.start(SandboxSpec(image="image:tag", workdir="/sync-session"), delete_on_stop=True) - assert sandbox.exec("pwd").return_code == 0 - assert provider.exec_calls[-1]["cwd"] == "/sync-session" - sandbox.shutdown() - sandbox.shutdown() + assert sandbox.status() == SandboxStatus.STOPPED assert provider.aclosed is True try: - sandbox.provider_name + sandbox.exec("pwd") except RuntimeError as e: assert "sync loop is closed" in str(e) else: @@ -528,7 +541,7 @@ def test_sync_sandbox_file_operations(tmp_path: Path) -> None: provider = FakeSandboxProvider() with Sandbox(provider) as sandbox: sandbox.start(SandboxSpec(image="image:tag")) - handle = sandbox.handle + handle = provider.created_handles[0] source_path = tmp_path / "source.txt" target_path = tmp_path / "target.txt" source_path.write_text("local", encoding="utf-8") @@ -1019,7 +1032,8 @@ def test_mini_swe_sandbox_environment_validation_and_context_manager() -> None: delete=False, ) as env: assert env._sandbox is not None - assert env._sandbox.handle.sandbox_id == "fake-1" + assert FakeSandboxProvider.last_instance is not None + assert FakeSandboxProvider.last_instance.created_handles[0].sandbox_id == "fake-1" assert FakeSandboxProvider.last_instance is not None assert FakeSandboxProvider.last_instance.closed[-1][1] is False From 112051d1214afc19a0d0bedb6f65e286df35f269 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Tue, 2 Jun 2026 16:30:45 -0700 Subject: [PATCH 07/14] Remove OpenSandbox pool path Signed-off-by: Hemil Desai --- nemo_gym/sandbox/__init__.py | 4 - nemo_gym/sandbox/providers/__init__.py | 4 - nemo_gym/sandbox/providers/base.py | 42 +- .../sandbox/providers/opensandbox/__init__.py | 4 - .../sandbox/providers/opensandbox/provider.py | 374 +----------------- tests/unit_tests/test_opensandbox_provider.py | 333 +--------------- tests/unit_tests/test_sandbox.py | 32 -- 7 files changed, 10 insertions(+), 783 deletions(-) diff --git a/nemo_gym/sandbox/__init__.py b/nemo_gym/sandbox/__init__.py index 36d54333d4..3ff5dcd49a 100644 --- a/nemo_gym/sandbox/__init__.py +++ b/nemo_gym/sandbox/__init__.py @@ -21,12 +21,10 @@ ImageSpec, OutsideEndpoint, SandboxAddressProvider, - SandboxBatchCreateError, SandboxCreateError, SandboxCreateVerificationError, SandboxExecResult, SandboxHandle, - SandboxHandleReferenceProvider, SandboxImageBuildProvider, SandboxInlineFileProvider, SandboxProvider, @@ -50,12 +48,10 @@ "ImageSpec", "OutsideEndpoint", "SandboxAddressProvider", - "SandboxBatchCreateError", "SandboxCreateError", "SandboxCreateVerificationError", "SandboxExecResult", "SandboxHandle", - "SandboxHandleReferenceProvider", "SandboxImageBuildProvider", "SandboxInlineFileProvider", "SandboxProvider", diff --git a/nemo_gym/sandbox/providers/__init__.py b/nemo_gym/sandbox/providers/__init__.py index 7e2284b9a9..20405c18e2 100644 --- a/nemo_gym/sandbox/providers/__init__.py +++ b/nemo_gym/sandbox/providers/__init__.py @@ -20,12 +20,10 @@ ImageSpec, OutsideEndpoint, SandboxAddressProvider, - SandboxBatchCreateError, SandboxCreateError, SandboxCreateVerificationError, SandboxExecResult, SandboxHandle, - SandboxHandleReferenceProvider, SandboxImageBuildProvider, SandboxInlineFileProvider, SandboxProvider, @@ -48,12 +46,10 @@ "ImageSpec", "OutsideEndpoint", "SandboxAddressProvider", - "SandboxBatchCreateError", "SandboxCreateError", "SandboxCreateVerificationError", "SandboxExecResult", "SandboxHandle", - "SandboxHandleReferenceProvider", "SandboxImageBuildProvider", "SandboxInlineFileProvider", "SandboxProvider", diff --git a/nemo_gym/sandbox/providers/base.py b/nemo_gym/sandbox/providers/base.py index be3f5301cd..ae73564983 100644 --- a/nemo_gym/sandbox/providers/base.py +++ b/nemo_gym/sandbox/providers/base.py @@ -19,7 +19,7 @@ instead of importing provider-specific modules. """ -from collections.abc import Awaitable, Callable +from collections.abc import Callable from dataclasses import dataclass, field from enum import Enum from pathlib import Path @@ -142,10 +142,6 @@ class SandboxCreateError(RuntimeError): """Raised when a provider cannot create a sandbox.""" -class SandboxBatchCreateError(RuntimeError): - """Raised when a provider cannot complete sandbox batch creation.""" - - class SandboxCreateVerificationError(SandboxCreateError): """Raised when a newly-created sandbox fails provider readiness checks.""" @@ -165,27 +161,6 @@ async def create(self, spec: SandboxSpec) -> SandboxHandle: """ ... - async def create_batch( - self, - spec: SandboxSpec, - count: int, - *, - allow_partial: bool = False, - ) -> list[SandboxHandle]: - """Create several equivalent sandboxes. - - Providers that have a native bulk-allocation primitive or warm-pool - implementation should use it. Providers without one may fall back to - calling ``create`` repeatedly. Long-lived pools are provider-owned and - configured through provider config or ``SandboxSpec.provider_options``, - rather than through a separate public pool handle. - - When ``allow_partial`` is true, providers may return a smaller - contiguous prefix of successfully created handles instead of failing the - whole batch. - """ - ... - async def exec( self, handle: SandboxHandle, @@ -212,20 +187,7 @@ async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: ... async def aclose(self) -> None: - """Close provider-scoped resources such as SDK clients or warm pools.""" - ... - - -@runtime_checkable -class SandboxHandleReferenceProvider(Protocol): - """Optional provider trait for loop-safe sandbox handle references.""" - - def handle_reference(self, handle: SandboxHandle) -> Any | Awaitable[Any]: - """Return a serializable or loop-safe reference for ``handle``.""" - ... - - def materialize_handle(self, value: Any) -> SandboxHandle | Awaitable[SandboxHandle]: - """Convert a value from ``handle_reference`` back into a local handle.""" + """Close provider-scoped resources such as SDK clients.""" ... diff --git a/nemo_gym/sandbox/providers/opensandbox/__init__.py b/nemo_gym/sandbox/providers/opensandbox/__init__.py index 2615667d73..c676205dd4 100644 --- a/nemo_gym/sandbox/providers/opensandbox/__init__.py +++ b/nemo_gym/sandbox/providers/opensandbox/__init__.py @@ -15,28 +15,24 @@ """OpenSandbox provider package.""" from nemo_gym.sandbox.providers.opensandbox.provider import ( - OpenSandboxBatchCreateError, OpenSandboxConnectionConfig, OpenSandboxCreateConfig, OpenSandboxCreateError, OpenSandboxCreateTimeoutError, OpenSandboxCreateVerificationError, OpenSandboxOperationConfig, - OpenSandboxPoolConfig, OpenSandboxProbeConfig, OpenSandboxProvider, ) __all__ = [ - "OpenSandboxBatchCreateError", "OpenSandboxConnectionConfig", "OpenSandboxCreateConfig", "OpenSandboxCreateError", "OpenSandboxCreateTimeoutError", "OpenSandboxCreateVerificationError", "OpenSandboxOperationConfig", - "OpenSandboxPoolConfig", "OpenSandboxProbeConfig", "OpenSandboxProvider", ] diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 9afd93124c..7c394bf99c 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -23,10 +23,8 @@ from datetime import timedelta from pathlib import Path from typing import Any, Awaitable, Callable -from uuid import uuid4 from nemo_gym.sandbox.providers.base import ( - SandboxBatchCreateError, SandboxCreateError, SandboxCreateVerificationError, SandboxExecResult, @@ -40,10 +38,6 @@ LOGGER = logging.getLogger(__name__) -class OpenSandboxBatchCreateError(SandboxBatchCreateError): - """Raised when a batch sandbox preallocation cannot be completed.""" - - class OpenSandboxCreateError(SandboxCreateError): """Raised when OpenSandbox cannot create a sandbox.""" @@ -130,23 +124,6 @@ def _require_opensandbox_sdk() -> tuple[Any, Any, Any, Any, Any]: return Sandbox, ConnectionConfig, RunCommandOpts, PlatformSpec, Volume -def _require_opensandbox_sdk_pool() -> tuple[Any, Any, Any, Any]: - try: - from opensandbox import ( - AcquirePolicy, - InMemoryAsyncPoolStateStore, - PoolCreationSpec, - SandboxPoolAsync, - ) - except ImportError as e: - raise ModuleNotFoundError( - "OpenSandbox SDK >=0.1.9 is required for native SDK pool batch creation. " - "Install nemo-gym[sandbox] in the runtime image." - ) from e - - return AcquirePolicy, InMemoryAsyncPoolStateStore, PoolCreationSpec, SandboxPoolAsync - - def _require_tenacity() -> tuple[Any, Any, Any, Any]: try: from tenacity import AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential @@ -392,12 +369,6 @@ def _provider_option_bool(provider_options: dict[str, Any], key: str) -> bool | return value -def _seconds_to_timedelta(seconds: int | float | None) -> timedelta | None: - if seconds is None: - return None - return timedelta(seconds=float(seconds)) - - def _to_sandbox_status(state: Any) -> SandboxStatus: normalized = str(state or "").lower() if normalized in {"active", "ready", "running"}: @@ -421,11 +392,6 @@ class OpenSandboxConnectionConfig: use_server_proxy: bool | None = None exec_use_server_proxy: bool | None = None request_timeout_s: int | None = None - connect_timeout_s: int | float | None = None - - def __post_init__(self) -> None: - if self.connect_timeout_s is not None and self.connect_timeout_s <= 0: - raise ValueError("connection.connect_timeout_s must be > 0") @dataclass(frozen=True) @@ -467,7 +433,6 @@ class OpenSandboxProbeConfig: expected_stdout: str | None = "nemo-gym-sandbox-ready" timeout_s: int = 30 deadline_s: float | None = None - sample_count: int | None = None stable_count: int = 1 stable_delay_s: float = 0.0 @@ -476,8 +441,6 @@ def __post_init__(self) -> None: raise ValueError("probe.timeout_s must be > 0") if self.deadline_s is not None and self.deadline_s <= 0: raise ValueError("probe.deadline_s must be > 0") - if self.sample_count is not None and self.sample_count < 1: - raise ValueError("probe.sample_count must be >= 1") if self.stable_count < 1: raise ValueError("probe.stable_count must be >= 1") if self.stable_delay_s < 0: @@ -507,32 +470,6 @@ def __post_init__(self) -> None: raise ValueError("operations.close_timeout_s must be > 0") -@dataclass(frozen=True) -class OpenSandboxPoolConfig: - """OpenSandbox SDK pool and batch fanout settings.""" - - concurrency: int = 4 - progress_timeout_s: float | None = None - reconcile_interval_s: float = 0.1 - acquire_poll_interval_s: float = 0.1 - idle_timeout_s: float | None = None - primary_lock_ttl_s: float | None = None - - def __post_init__(self) -> None: - if self.concurrency < 1: - raise ValueError("pool.concurrency must be >= 1") - if self.progress_timeout_s is not None and self.progress_timeout_s <= 0: - raise ValueError("pool.progress_timeout_s must be > 0") - if self.reconcile_interval_s <= 0: - raise ValueError("pool.reconcile_interval_s must be > 0") - if self.acquire_poll_interval_s <= 0: - raise ValueError("pool.acquire_poll_interval_s must be > 0") - if self.idle_timeout_s is not None and self.idle_timeout_s <= 0: - raise ValueError("pool.idle_timeout_s must be > 0") - if self.primary_lock_ttl_s is not None and self.primary_lock_ttl_s <= 0: - raise ValueError("pool.primary_lock_ttl_s must be > 0") - - def _coerce_config(value: Any, config_cls: type[Any]) -> Any: if value is None: return config_cls() @@ -544,10 +481,7 @@ def _coerce_config(value: Any, config_cls: type[Any]) -> Any: class OpenSandboxProvider: - """Provider backed by the OpenSandbox SDK/server API. - - Batch allocations use the official OpenSandbox SDK client-side pool. - """ + """Provider backed by the OpenSandbox SDK/server API.""" name = "opensandbox" @@ -558,13 +492,11 @@ def __init__( create: OpenSandboxCreateConfig | Mapping[str, Any] | None = None, probe: OpenSandboxProbeConfig | Mapping[str, Any] | None = None, operations: OpenSandboxOperationConfig | Mapping[str, Any] | None = None, - pool: OpenSandboxPoolConfig | Mapping[str, Any] | None = None, ) -> None: self._connection = _coerce_config(connection, OpenSandboxConnectionConfig) self._create = _coerce_config(create, OpenSandboxCreateConfig) self._probe = _coerce_config(probe, OpenSandboxProbeConfig) self._operations = _coerce_config(operations, OpenSandboxOperationConfig) - self._pool = _coerce_config(pool, OpenSandboxPoolConfig) def _with_default_image_pull_policy(self, spec: SandboxSpec) -> SandboxSpec: """Ensure SDK create requests carry the desired image pull policy.""" @@ -753,42 +685,6 @@ async def _verify_created_handle(self, handle: SandboxHandle) -> None: if successful_probes < self._probe.stable_count and self._probe.stable_delay_s: await asyncio.sleep(self._probe.stable_delay_s) - async def _verify_created_handles( - self, - handles: list[SandboxHandle], - ) -> None: - """Verify a batch of created handles with bounded probe concurrency.""" - if self._probe.command is None or not handles: - return - - handles_to_probe = handles - if self._probe.sample_count is not None and self._probe.sample_count < len(handles): - sample_count = self._probe.sample_count - if sample_count == 1: - sampled_indices = [0] - else: - sampled_indices = [ - round(index * (len(handles) - 1) / (sample_count - 1)) for index in range(sample_count) - ] - handles_to_probe = [handles[index] for index in sampled_indices] - - semaphore = asyncio.Semaphore(self._pool.concurrency) - - async def _verify_one(handle: SandboxHandle) -> None: - async with semaphore: - await self._verify_created_handle(handle) - - results = await asyncio.gather( - *(_verify_one(handle) for handle in handles_to_probe), - return_exceptions=True, - ) - errors = [result for result in results if isinstance(result, Exception)] - if errors: - raise OpenSandboxCreateVerificationError( - "One or more OpenSandbox sandboxes failed create probe " - f"verification; failed={len(errors)}, total={len(handles)}" - ) from errors[0] - async def _cleanup_failed_create_handle(self, handle: SandboxHandle) -> None: try: await self.close(handle, delete=True) @@ -845,13 +741,6 @@ async def _connect_after_create(self, handle: SandboxHandle, spec: SandboxSpec) async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: """Create a sandbox through ``opensandbox.Sandbox.create``.""" - if spec.extensions.get("poolRef") and self._connection.use_server_proxy is False: - raise ValueError( - "OpenSandbox pooled creation requires " - "use_server_proxy=True so SDK calls are routed through the " - "server proxy and do not rely on stale cached pod endpoints." - ) - Sandbox, _, _, _, _ = _require_opensandbox_sdk() kwargs: dict[str, Any] = { @@ -906,7 +795,6 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: error = OpenSandboxCreateTimeoutError( "Timed out creating OpenSandbox sandbox after " f"{timeout_s:g}s; image={spec.image!r}, " - f"poolRef={spec.extensions.get('poolRef')!r}, " f"ready_timeout_s={spec.ready_timeout_s!r}" ) raise error from e @@ -930,8 +818,6 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: async def _create_with_retries( self, spec: SandboxSpec, - *, - semaphore: asyncio.Semaphore | None = None, ) -> SandboxHandle: AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential = _require_tenacity() retry_policy = AsyncRetrying( @@ -946,269 +832,15 @@ async def _create_with_retries( ) async for attempt in retry_policy: with attempt: - if semaphore is None: - return await self._create_once(spec) - async with semaphore: - return await self._create_once(spec) + return await self._create_once(spec) - raise OpenSandboxBatchCreateError("OpenSandbox create retry loop did not run") + raise OpenSandboxCreateError("OpenSandbox create retry loop did not run") async def create(self, spec: SandboxSpec) -> SandboxHandle: """Create one sandbox through the configured OpenSandbox path.""" spec = self._with_default_image_pull_policy(_normalize_spec(spec)) return await self._create_with_retries(spec) - async def _close_many( - self, - handles: list[SandboxHandle], - *, - delete: bool, - ) -> list[Any]: - semaphore = asyncio.Semaphore(self._pool.concurrency) - - async def _close_one(handle: SandboxHandle) -> Any: - async with semaphore: - return await self.close(handle, delete=delete) - - return list( - await asyncio.gather( - *(_close_one(handle) for handle in handles), - return_exceptions=True, - ) - ) - - def _validate_sdk_pool_spec(self, spec: SandboxSpec) -> None: - if spec.image is None: - raise ValueError("OpenSandbox SDK pool requires SandboxSpec.image") - if spec.provider_options.get(PROVIDER_OPTION_SNAPSHOT_ID) is not None: - raise ValueError("OpenSandbox SDK pool does not support snapshot_id") - - def _to_pool_creation_spec(self, spec: SandboxSpec) -> Any: - self._validate_sdk_pool_spec(spec) - _, _, PoolCreationSpec, _ = _require_opensandbox_sdk_pool() - volumes = _spec_volumes(spec) - return PoolCreationSpec( - image=spec.image, - entrypoint=spec.entrypoint, - resource=spec.resources or None, - env=spec.env or None, - metadata=spec.metadata or None, - extensions=spec.extensions or None, - platform=_to_platform_spec(spec.provider_options[PROVIDER_OPTION_PLATFORM]) - if PROVIDER_OPTION_PLATFORM in spec.provider_options - else None, - volumes=_to_volumes(volumes) if volumes is not None else None, - ) - - async def _wait_sdk_pool_idle( - self, - pool: Any, - *, - spec: SandboxSpec, - requested: int, - timeout_s: float, - allow_partial: bool, - ) -> int: - loop = asyncio.get_running_loop() - deadline = loop.time() + timeout_s - last_progress_at = loop.time() - last_idle = 0 - last_snapshot: Any = None - - while True: - last_snapshot = await pool.snapshot() - idle_count = int(getattr(last_snapshot, "idle_count", 0) or 0) - if idle_count >= requested: - return requested - if idle_count > last_idle: - last_idle = idle_count - last_progress_at = loop.time() - - now = loop.time() - progress_timeout_s = self._pool.progress_timeout_s - if progress_timeout_s is not None and now - last_progress_at >= progress_timeout_s: - if allow_partial and idle_count > 0: - return idle_count - error = OpenSandboxCreateTimeoutError( - "Timed out waiting for OpenSandbox SDK pool warmup progress " - f"after {progress_timeout_s:g}s; requested={requested}, " - f"idle={idle_count}, snapshot={last_snapshot!r}" - ) - raise error - if now >= deadline: - if allow_partial and idle_count > 0: - return idle_count - error = OpenSandboxCreateTimeoutError( - "Timed out waiting for OpenSandbox SDK pool warmup after " - f"{timeout_s:g}s; requested={requested}, idle={idle_count}, " - f"snapshot={last_snapshot!r}" - ) - raise error - await asyncio.sleep(self._pool.acquire_poll_interval_s) - - async def _direct_exec_handle_for_acquired_sandbox(self, sandbox: Any, spec: SandboxSpec) -> SandboxHandle: - handle = SandboxHandle(sandbox_id=str(sandbox.id), provider_name=self.name, raw=sandbox) - if self._connection.exec_use_server_proxy is None and not self._create.skip_health_check: - return handle - return await self._connect_after_create(handle, spec) - - async def _create_batch_sdk_pool( - self, - spec: SandboxSpec, - count: int, - *, - allow_partial: bool, - ) -> list[SandboxHandle]: - AcquirePolicy, InMemoryAsyncPoolStateStore, _, SandboxPoolAsync = _require_opensandbox_sdk_pool() - ready_timeout_s = float( - spec.ready_timeout_s or self._create.timeout_s or self._connection.request_timeout_s or 300.0 - ) - idle_timeout_s = float(self._pool.idle_timeout_s or spec.timeout_s or max(ready_timeout_s * 2.0, 3600.0)) - primary_lock_ttl_s = float(self._pool.primary_lock_ttl_s or max(ready_timeout_s + 60.0, 60.0)) - pool_name = f"nemo-gym-{uuid4().hex[:12]}" - skip_health_check = bool( - self._create.skip_health_check - or _provider_option_bool(spec.provider_options, PROVIDER_OPTION_SKIP_HEALTH_CHECK) - ) - - async def _warmup_preparer(sandbox: Any) -> None: - if self._probe.command is None: - return - handle = await self._direct_exec_handle_for_acquired_sandbox(sandbox, spec) - try: - await self._verify_created_handle(handle) - finally: - if handle.raw is not sandbox: - try: - await self._await_sdk_call( - handle.raw.close(), - operation="close warmup direct handle", - sandbox_id=handle.sandbox_id, - timeout_s=self._operations.close_timeout_s, - ) - except Exception as e: - LOGGER.warning( - "Failed to close temporary OpenSandbox direct exec handle for sandbox %r: %r", - handle.sandbox_id, - e, - ) - - pool = SandboxPoolAsync( - pool_name=pool_name, - max_idle=count, - warmup_concurrency=self._pool.concurrency, - state_store=InMemoryAsyncPoolStateStore(), - connection_config=self._connection_config(request_timeout_s=self._create.request_timeout_s), - creation_spec=self._to_pool_creation_spec(spec), - reconcile_interval=timedelta(seconds=self._pool.reconcile_interval_s), - primary_lock_ttl=timedelta(seconds=primary_lock_ttl_s), - acquire_ready_timeout=timedelta(seconds=ready_timeout_s), - warmup_ready_timeout=timedelta(seconds=ready_timeout_s), - warmup_sandbox_preparer=_warmup_preparer, - acquire_skip_health_check=skip_health_check, - warmup_skip_health_check=skip_health_check, - idle_timeout=timedelta(seconds=idle_timeout_s), - ) - handles: list[SandboxHandle] = [] - try: - await pool.start() - ready_count = await self._wait_sdk_pool_idle( - pool, - spec=spec, - requested=count, - timeout_s=ready_timeout_s, - allow_partial=allow_partial, - ) - await pool.resize(0) - sandbox_timeout = _seconds_to_timedelta(spec.timeout_s) - for index in range(ready_count): - sandbox = await pool.acquire( - sandbox_timeout=sandbox_timeout, - policy=AcquirePolicy.FAIL_FAST, - ) - handle = await self._direct_exec_handle_for_acquired_sandbox(sandbox, spec) - handles.append(handle) - LOGGER.info( - "Acquired OpenSandbox SDK pool sandbox %s/%s: %s", - index + 1, - ready_count, - sandbox.id, - ) - return handles - except Exception: - await self._close_many(handles, delete=True) - raise - finally: - try: - await pool.shutdown(graceful=False) - finally: - await pool.release_all_idle() - - async def _create_batch_sdk( - self, - spec: SandboxSpec, - count: int, - *, - allow_partial: bool = False, - ) -> list[SandboxHandle]: - """Create several sandboxes through the OpenSandbox SDK pool.""" - if count < 1: - raise ValueError("count must be >= 1") - return await self._create_batch_sdk_pool( - spec, - count, - allow_partial=allow_partial, - ) - - async def create_batch( - self, - spec: SandboxSpec, - count: int, - *, - allow_partial: bool = False, - ) -> list[SandboxHandle]: - """Create several equivalent OpenSandbox sandboxes.""" - if count < 1: - raise ValueError("count must be >= 1") - spec = self._with_default_image_pull_policy(_normalize_spec(spec)) - return await self._create_batch_sdk( - spec, - count, - allow_partial=allow_partial, - ) - - def handle_reference(self, handle: SandboxHandle) -> dict[str, Any]: - """Build a loop-neutral reference for a sandbox handle. - - OpenSandbox SDK handles are bound to the event loop where they were - created. Prewarmed handles may cross from a FastAPI prewarm request into - a thread-pool runner, so only pass a serializable reference across that - boundary and re-materialize SDK adapters in the consuming event loop. - """ - return { - "kind": "sandbox_id", - "provider": self.name, - "sandbox_id": handle.sandbox_id, - } - - async def materialize_handle(self, reference: dict[str, Any]) -> SandboxHandle: - """Create a loop-local handle from ``handle_reference`` output.""" - kind = reference.get("kind") - if kind == "sandbox_id": - return await self.connect(str(reference["sandbox_id"])) - raise ValueError(f"Unsupported OpenSandbox handle reference kind: {kind!r}") - - async def connect(self, sandbox_id: str) -> SandboxHandle: - """Connect to an existing OpenSandbox sandbox.""" - Sandbox, _, _, _, _ = _require_opensandbox_sdk() - kwargs: dict[str, Any] = { - "connection_config": self._exec_connection_config(), - } - if self._connection.connect_timeout_s is not None: - kwargs["connect_timeout"] = timedelta(seconds=self._connection.connect_timeout_s) - sandbox = await Sandbox.connect(sandbox_id, **kwargs) - return SandboxHandle(sandbox_id=str(sandbox.id), provider_name=self.name, raw=sandbox) - async def status(self, handle: SandboxHandle) -> SandboxStatus: """Return the current OpenSandbox lifecycle status.""" get_info = getattr(handle.raw, "get_info", None) diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index 6cedb83961..345c968d1c 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -67,94 +67,16 @@ async def connect(cls, *args: Any, **kwargs: Any) -> "FakeSandbox": return cls() -@dataclass -class FakePoolCreationSpec: - image: str - entrypoint: list[str] | None = None - resource: dict[str, str] | None = None - env: dict[str, str] | None = None - metadata: dict[str, str] | None = None - extensions: dict[str, str] | None = None - platform: Any | None = None - volumes: list[Any] | None = None - - -class FakeAcquirePolicy: - FAIL_FAST = "fail_fast" - - -class FakeStateStore: - pass - - -class FakeSnapshot: - idle_count = 1 - state = None - - -class FakeSandboxPoolAsync: - received_kwargs: dict[str, Any] = {} - - def __init__(self, **kwargs: Any) -> None: - self.received_kwargs = kwargs - type(self).received_kwargs = kwargs - - async def start(self) -> None: - preparer = self.received_kwargs.get("warmup_sandbox_preparer") - if preparer is not None: - await preparer(FakeSandbox("warmup-1")) - return None - - async def snapshot(self) -> FakeSnapshot: - return FakeSnapshot() - - async def resize(self, _count: int) -> None: - return None - - async def acquire( - self, - *, - sandbox_timeout: timedelta | None, - policy: str, - ) -> FakeSandbox: - del sandbox_timeout, policy - creation_spec = self.received_kwargs["creation_spec"] - return await FakeSandbox.create( - creation_spec.image, - platform=creation_spec.platform, - ) - - async def shutdown(self, *, graceful: bool) -> None: - del graceful - - async def release_all_idle(self) -> None: - return None - - @pytest.fixture def fake_opensandbox_sdk(monkeypatch: pytest.MonkeyPatch) -> None: def require_sdk() -> tuple[Any, Any, Any, Any, Any]: return FakeSandbox, FakeConnectionConfig, object, FakePlatformSpec, object - def require_sdk_pool() -> tuple[Any, Any, Any, Any]: - return ( - FakeAcquirePolicy, - FakeStateStore, - FakePoolCreationSpec, - FakeSandboxPoolAsync, - ) - monkeypatch.setattr(opensandbox_provider, "_require_opensandbox_sdk", require_sdk) - monkeypatch.setattr( - opensandbox_provider, - "_require_opensandbox_sdk_pool", - require_sdk_pool, - ) def test_sdk_import_helpers_and_retry_classification() -> None: assert len(opensandbox_provider._require_opensandbox_sdk()) == 5 - assert len(opensandbox_provider._require_opensandbox_sdk_pool()) == 4 assert len(opensandbox_provider._require_tenacity()) == 4 class StatusCodeError(Exception): @@ -213,10 +135,6 @@ def fake_import( with pytest.raises(ModuleNotFoundError, match="OpenSandbox SDK is required"): opensandbox_provider._require_opensandbox_sdk() - block_imports("opensandbox") - with pytest.raises(ModuleNotFoundError, match="OpenSandbox SDK >=0.1.9"): - opensandbox_provider._require_opensandbox_sdk_pool() - block_imports("tenacity") with pytest.raises(ModuleNotFoundError, match="tenacity is required"): opensandbox_provider._require_tenacity() @@ -228,7 +146,7 @@ def fake_import( assert opensandbox_provider._is_retryable_create_error(RuntimeError("gateway timeout")) is True -async def test_provider_reference_materialization_and_conversion_helpers( +async def test_provider_conversion_helpers( fake_opensandbox_sdk: None, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -245,23 +163,8 @@ async def test_provider_reference_materialization_and_conversion_helpers( ) assert opensandbox_provider._to_volumes([{"name": "workspace"}]) == [FakeVolume(name="workspace")] - provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) - handle = opensandbox_provider.SandboxHandle(sandbox_id="sandbox-1", provider_name="opensandbox", raw=object()) - assert provider.handle_reference(handle) == { - "kind": "sandbox_id", - "provider": "opensandbox", - "sandbox_id": "sandbox-1", - } - - async def connect(sandbox_id: str) -> opensandbox_provider.SandboxHandle: - return opensandbox_provider.SandboxHandle(sandbox_id=sandbox_id, provider_name="opensandbox", raw="connected") - - monkeypatch.setattr(provider, "connect", connect) - materialized = await provider.materialize_handle({"kind": "sandbox_id", "sandbox_id": "sandbox-2"}) - assert materialized.raw == "connected" - -async def test_sdk_pool_passes_platform_through_pool_creation_spec( +async def test_direct_create_passes_platform_to_sdk_create( fake_opensandbox_sdk: None, ) -> None: provider = opensandbox_provider.OpenSandboxProvider( @@ -269,38 +172,20 @@ async def test_sdk_pool_passes_platform_through_pool_creation_spec( probe={"command": None}, ) - handles = await provider.create_batch( + handle = await provider.create( SandboxSpec( image="mirror.gcr.io/astral/uv:python3.12-bookworm-slim", provider_options={"platform": {"os": "linux", "arch": "amd64"}}, ), - 1, ) - assert len(handles) == 1 - assert handles[0].sandbox_id == "sandbox-1" - assert "sandbox_factory" not in FakeSandboxPoolAsync.received_kwargs + assert handle.sandbox_id == "sandbox-1" assert FakeSandbox.created_kwargs["platform"] == FakePlatformSpec( os="linux", arch="amd64", ) -async def test_connect_passes_configured_connect_timeout( - fake_opensandbox_sdk: None, -) -> None: - provider = opensandbox_provider.OpenSandboxProvider( - connection={"connect_timeout_s": 300, "request_timeout_s": 10}, - probe={"command": None}, - ) - - handle = await provider.connect("sandbox-123") - - assert handle.sandbox_id == "sandbox-1" - assert FakeSandbox.connected_args == ("sandbox-123",) - assert FakeSandbox.connected_kwargs["connect_timeout"] == timedelta(seconds=300) - - def test_provider_validation_and_retry_helpers() -> None: with pytest.raises(ValueError, match="image_pull_policy"): opensandbox_provider.validate_image_pull_policy("Sometimes") @@ -332,13 +217,9 @@ def test_provider_validation_and_retry_helpers() -> None: assert opensandbox_provider._to_sandbox_status(None) == SandboxStatus.UNKNOWN invalid_kwargs = [ - {"pool": {"concurrency": 0}}, - {"connection": {"connect_timeout_s": 0}}, - {"pool": {"progress_timeout_s": 0}}, {"create": {"timeout_s": 0}}, {"probe": {"timeout_s": 0}}, {"probe": {"deadline_s": 0}}, - {"probe": {"sample_count": 0}}, {"probe": {"stable_count": 0}}, {"probe": {"stable_delay_s": -1}}, {"create": {"retries": -1}}, @@ -348,10 +229,6 @@ def test_provider_validation_and_retry_helpers() -> None: {"operations": {"retry_delay_s": -1}}, {"operations": {"retry_max_delay_s": -1}}, {"operations": {"command_retries": -1}}, - {"pool": {"reconcile_interval_s": 0}}, - {"pool": {"acquire_poll_interval_s": 0}}, - {"pool": {"idle_timeout_s": 0}}, - {"pool": {"primary_lock_ttl_s": 0}}, {"operations": {"close_timeout_s": 0}}, {"create": {"connect_attempt_timeout_s": 0}}, {"create": {"connect_poll_s": 0}}, @@ -378,8 +255,6 @@ def test_provider_validation_and_retry_helpers() -> None: assert attrs["status_code"] == 502 assert attrs["attempt_number"] == 2 assert attrs["next_sleep_s"] == 0.5 - assert opensandbox_provider._seconds_to_timedelta(None) is None - assert opensandbox_provider._seconds_to_timedelta(1.5) == timedelta(seconds=1.5) def test_connection_config_exec_proxy_and_image_policy(fake_opensandbox_sdk: None) -> None: @@ -415,77 +290,7 @@ def test_connection_config_exec_proxy_and_image_policy(fake_opensandbox_sdk: Non assert no_policy_provider._with_default_image_pull_policy(spec) is spec -async def test_wait_sdk_pool_idle_success_partial_and_timeout(monkeypatch: pytest.MonkeyPatch) -> None: - class Snapshot: - def __init__(self, idle_count: int) -> None: - self.idle_count = idle_count - self.state = SimpleNamespace(value="warming") - - class FakePool: - def __init__(self, counts: list[int]) -> None: - self.counts = counts - self.index = 0 - self._config = SimpleNamespace(pool_name="pool-1") - - async def snapshot(self) -> Snapshot: - count = self.counts[min(self.index, len(self.counts) - 1)] - self.index += 1 - return Snapshot(count) - - async def no_sleep(_seconds: float) -> None: - return None - - monkeypatch.setattr(opensandbox_provider.asyncio, "sleep", no_sleep) - - provider = opensandbox_provider.OpenSandboxProvider( - pool={"acquire_poll_interval_s": 0.01}, - probe={"command": None}, - ) - assert ( - await provider._wait_sdk_pool_idle( - FakePool([0, 1, 2]), - spec=SandboxSpec(image="image:tag"), - requested=2, - timeout_s=1, - allow_partial=False, - ) - == 2 - ) - assert ( - await provider._wait_sdk_pool_idle( - FakePool([1]), - spec=SandboxSpec(image="image:tag"), - requested=2, - timeout_s=0, - allow_partial=True, - ) - == 1 - ) - progress_timeout_provider = opensandbox_provider.OpenSandboxProvider( - pool={"progress_timeout_s": 0.000001, "acquire_poll_interval_s": 0.000001}, - probe={"command": None}, - ) - assert ( - await progress_timeout_provider._wait_sdk_pool_idle( - FakePool([1]), - spec=SandboxSpec(image="image:tag"), - requested=2, - timeout_s=1, - allow_partial=True, - ) - == 1 - ) - with pytest.raises(opensandbox_provider.OpenSandboxCreateTimeoutError): - await provider._wait_sdk_pool_idle( - FakePool([0]), - spec=SandboxSpec(image="image:tag"), - requested=2, - timeout_s=0, - allow_partial=False, - ) - - -async def test_exec_file_operations_and_batch_validation(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: +async def test_exec_file_operations_and_reference_validation(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: class FakeRunCommandOpts: def __init__(self, **kwargs: Any) -> None: self.kwargs = kwargs @@ -584,15 +389,6 @@ async def get_info(self) -> Any: assert await provider.status(bare_handle) == SandboxStatus.UNKNOWN assert await provider.container_ip(bare_handle) is None - with pytest.raises(ValueError, match="count"): - await provider._create_batch_sdk(SandboxSpec(image="image:tag"), 0) - with pytest.raises(ValueError, match="count"): - await provider.create_batch(SandboxSpec(image="image:tag"), 0) - with pytest.raises(ValueError, match="snapshot_id"): - provider._validate_sdk_pool_spec(SandboxSpec(image="image:tag", provider_options={"snapshot_id": "snapshot"})) - with pytest.raises(ValueError, match="Unsupported"): - await provider.materialize_handle({"kind": "other"}) - async def test_provider_create_probe_and_close_error_paths(monkeypatch: pytest.MonkeyPatch) -> None: provider = opensandbox_provider.OpenSandboxProvider( @@ -642,17 +438,7 @@ async def cancelled_probe(*_args: Any, **_kwargs: Any) -> opensandbox_provider.S with pytest.raises(asyncio.CancelledError): await provider._verify_created_handle(handle) - provider = opensandbox_provider.OpenSandboxProvider(probe={"command": "probe"}) - - async def fail_verify(_handle: Any) -> None: - raise RuntimeError("probe failed") - - monkeypatch.setattr(provider, "_verify_created_handle", fail_verify) - with pytest.raises(opensandbox_provider.OpenSandboxCreateVerificationError): - await provider._verify_created_handles([handle, handle]) - provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) - await provider._verify_created_handles([]) async def close_raises(_handle: Any, *, delete: bool) -> None: del delete @@ -678,21 +464,6 @@ async def close(self) -> None: delete=True, ) - class CloseSucceedsRaw: - async def close(self) -> None: - return None - - assert await provider._close_many( - [ - opensandbox_provider.SandboxHandle( - sandbox_id="sandbox-close", - provider_name="opensandbox", - raw=CloseSucceedsRaw(), - ) - ], - delete=False, - ) == [None] - class DeleteAndCloseFailRaw: async def kill(self) -> None: raise RuntimeError("delete failed") @@ -732,13 +503,6 @@ async def test_create_once_and_connect_after_create_error_paths( fake_opensandbox_sdk: None, monkeypatch: pytest.MonkeyPatch, ) -> None: - provider = opensandbox_provider.OpenSandboxProvider( - connection={"use_server_proxy": False}, - probe={"command": None}, - ) - with pytest.raises(ValueError, match="pooled creation"): - await provider._create_once(SandboxSpec(image="image:tag", extensions={"poolRef": "pool"})) - provider = opensandbox_provider.OpenSandboxProvider( create={"timeout_s": 1, "skip_health_check": True}, probe={"command": None}, @@ -897,18 +661,6 @@ async def cleanup(handle: opensandbox_provider.SandboxHandle) -> None: await provider._create_once(SandboxSpec(image="image:tag")) assert cleanup_calls == ["sandbox-1"] - provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) - - async def create_once(_spec: SandboxSpec) -> opensandbox_provider.SandboxHandle: - return opensandbox_provider.SandboxHandle( - sandbox_id="sandbox-semaphore", provider_name="opensandbox", raw=None - ) - - monkeypatch.setattr(provider, "_create_once", create_once) - assert ( - await provider._create_with_retries(SandboxSpec(image="image:tag"), semaphore=asyncio.Semaphore(1)) - ).sandbox_id == "sandbox-semaphore" - async def test_retry_classification_and_await_sdk_helpers(monkeypatch: pytest.MonkeyPatch) -> None: provider = opensandbox_provider.OpenSandboxProvider( @@ -942,80 +694,5 @@ async def cancelled() -> None: ) -async def test_probe_sampling_pool_progress_and_direct_exec_paths(monkeypatch: pytest.MonkeyPatch) -> None: - handles = [ - opensandbox_provider.SandboxHandle(sandbox_id=f"sandbox-{index}", provider_name="opensandbox", raw=object()) - for index in range(3) - ] - seen_handles: list[str] = [] - - provider = opensandbox_provider.OpenSandboxProvider( - probe={"command": "probe", "sample_count": 2}, - pool={"progress_timeout_s": 0.01, "acquire_poll_interval_s": 0.01}, - ) - - async def verify_created_handle(handle: opensandbox_provider.SandboxHandle) -> None: - seen_handles.append(handle.sandbox_id) - - monkeypatch.setattr(provider, "_verify_created_handle", verify_created_handle) - await provider._verify_created_handles(handles) - assert seen_handles == ["sandbox-0", "sandbox-2"] - - provider = opensandbox_provider.OpenSandboxProvider(probe={"command": "probe", "sample_count": 1}) - seen_handles = [] - monkeypatch.setattr(provider, "_verify_created_handle", verify_created_handle) - await provider._verify_created_handles(handles) - assert seen_handles == ["sandbox-0"] - - with pytest.raises(ValueError, match="requires SandboxSpec.image"): - provider._validate_sdk_pool_spec(SandboxSpec(image=None)) - - class Snapshot: - idle_count = 0 - state = SimpleNamespace(value="warming") - - class NoProgressPool: - async def snapshot(self) -> Snapshot: - return Snapshot() - - async def no_sleep(_seconds: float) -> None: - return None - - provider = opensandbox_provider.OpenSandboxProvider( - pool={"progress_timeout_s": 0.001, "acquire_poll_interval_s": 0.001}, - probe={"command": None}, - ) - monkeypatch.setattr(opensandbox_provider.asyncio, "sleep", no_sleep) - with pytest.raises(opensandbox_provider.OpenSandboxCreateTimeoutError, match="warmup progress"): - await provider._wait_sdk_pool_idle( - NoProgressPool(), - spec=SandboxSpec(image="image:tag"), - requested=2, - timeout_s=1, - allow_partial=False, - ) - - provider = opensandbox_provider.OpenSandboxProvider( - connection={"exec_use_server_proxy": False}, - probe={"command": None}, - ) - - async def connect_after_create( - handle: opensandbox_provider.SandboxHandle, - _spec: SandboxSpec, - ) -> opensandbox_provider.SandboxHandle: - return opensandbox_provider.SandboxHandle( - sandbox_id=handle.sandbox_id, - provider_name="opensandbox", - raw="direct", - ) - - monkeypatch.setattr(provider, "_connect_after_create", connect_after_create) - direct_handle = await provider._direct_exec_handle_for_acquired_sandbox( - FakeSandbox("sandbox-direct"), SandboxSpec() - ) - assert direct_handle.raw == "direct" - - async def _return_value(value: Any) -> Any: return value diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index 557d21e9cd..f7fdd5e571 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -27,7 +27,6 @@ ImageSpec, OutsideEndpoint, Sandbox, - SandboxBatchCreateError, SandboxCreateError, SandboxExecResult, SandboxHandle, @@ -107,16 +106,6 @@ async def create(self, spec: SandboxSpec) -> SandboxHandle: self.created_handles.append(handle) return handle - async def create_batch( - self, - spec: SandboxSpec, - count: int, - *, - allow_partial: bool = False, - ) -> list[SandboxHandle]: - del allow_partial - return [await self.create(spec) for _ in range(count)] - async def exec( self, handle: SandboxHandle, @@ -184,16 +173,6 @@ async def create(self, spec: SandboxSpec) -> SandboxHandle: self.created_handles.append(handle) return handle - async def create_batch( - self, - spec: SandboxSpec, - count: int, - *, - allow_partial: bool = False, - ) -> list[SandboxHandle]: - del allow_partial - return [await self.create(spec) for _ in range(count)] - async def exec( self, handle: SandboxHandle, @@ -238,16 +217,6 @@ async def create(self, spec: SandboxSpec) -> SandboxHandle: self.created_handles.append(handle) return handle - async def create_batch( - self, - spec: SandboxSpec, - count: int, - *, - allow_partial: bool = False, - ) -> list[SandboxHandle]: - del allow_partial - return [await self.create(spec) for _ in range(count)] - async def exec( self, handle: SandboxHandle, @@ -776,7 +745,6 @@ def test_opensandbox_create_probe_failures_are_retryable() -> None: error = OpenSandboxCreateVerificationError("pod sdk-sandbox-0 failed create probe") assert isinstance(error, SandboxCreateError) - assert not isinstance(SandboxBatchCreateError("batch failed"), SandboxCreateError) assert opensandbox_provider_module._is_retryable_create_error(error) is True From 7b193a693d3b589ee9f0494185c6dad134fd90eb Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Tue, 2 Jun 2026 17:00:31 -0700 Subject: [PATCH 08/14] Simplify sandbox provider protocol Signed-off-by: Hemil Desai --- nemo_gym/sandbox/__init__.py | 16 --- nemo_gym/sandbox/api.py | 48 +------ nemo_gym/sandbox/providers/__init__.py | 16 --- nemo_gym/sandbox/providers/base.py | 115 ++-------------- .../sandbox/providers/opensandbox/provider.py | 125 +++--------------- .../mini_swe_agent_2/README.md | 12 +- .../configs/mini_swe_agent_opensandbox.yaml | 2 - .../mini_swe_agent_2/sandbox_environment.py | 4 +- tests/unit_tests/test_opensandbox_provider.py | 64 ++------- tests/unit_tests/test_sandbox.py | 99 +++++--------- 10 files changed, 86 insertions(+), 415 deletions(-) diff --git a/nemo_gym/sandbox/__init__.py b/nemo_gym/sandbox/__init__.py index 3ff5dcd49a..4754ff483a 100644 --- a/nemo_gym/sandbox/__init__.py +++ b/nemo_gym/sandbox/__init__.py @@ -17,21 +17,13 @@ from nemo_gym.sandbox.api import AsyncSandbox, Sandbox from nemo_gym.sandbox.providers import ( ExecResult, - ImageBuildRequest, - ImageSpec, - OutsideEndpoint, - SandboxAddressProvider, SandboxCreateError, SandboxCreateVerificationError, SandboxExecResult, SandboxHandle, - SandboxImageBuildProvider, - SandboxInlineFileProvider, SandboxProvider, SandboxSpec, SandboxStatus, - SandboxStatusProvider, - VolumeMount, create_provider, get_provider_class, list_providers, @@ -44,21 +36,13 @@ "Sandbox", "AsyncSandbox", "ExecResult", - "ImageBuildRequest", - "ImageSpec", - "OutsideEndpoint", - "SandboxAddressProvider", "SandboxCreateError", "SandboxCreateVerificationError", "SandboxExecResult", "SandboxHandle", - "SandboxImageBuildProvider", - "SandboxInlineFileProvider", "SandboxProvider", "SandboxSpec", "SandboxStatus", - "SandboxStatusProvider", - "VolumeMount", "create_provider", "get_provider_class", "list_providers", diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index a8132267ea..1d58c931c7 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -19,21 +19,15 @@ import threading from collections.abc import Awaitable, Callable, Mapping from concurrent.futures import Future -from dataclasses import replace from pathlib import Path from typing import Any, TypeVar from nemo_gym.sandbox.providers import ( - ImageBuildRequest, - OutsideEndpoint, SandboxExecResult, SandboxHandle, - SandboxImageBuildProvider, - SandboxInlineFileProvider, SandboxProvider, SandboxSpec, SandboxStatus, - SandboxStatusProvider, create_provider, ) @@ -58,41 +52,12 @@ def __init__( self._stopped = True self._closed = False - def _provider_name(self) -> str: - return self._provider.name - def _require_handle(self) -> SandboxHandle: if self._handle is None or self._stopped: raise RuntimeError("Sandbox has not been started") return self._handle - async def _build_images(self, request: ImageBuildRequest) -> list[str]: - if not isinstance(self._provider, SandboxImageBuildProvider): - raise NotImplementedError(f"Provider {self._provider_name()!r} does not support sandbox image builds") - return await self._provider.build_images(request) - - async def _resolve_image_build(self, spec: SandboxSpec) -> SandboxSpec: - if spec.image_build is None: - return spec - built_images = await self._build_images(ImageBuildRequest(specs=[spec.image_build])) - if not built_images: - raise ValueError("build_images returned no image references") - return replace(spec, image=spec.image or built_images[0]) - - def _with_outside_endpoints( - self, - spec: SandboxSpec, - outside_endpoints: list[OutsideEndpoint] | None, - ) -> SandboxSpec: - if not outside_endpoints: - return spec - endpoint_env = {endpoint.env_var: endpoint.url for endpoint in outside_endpoints} - return replace(spec, env={**spec.env, **endpoint_env}) - async def _write_inline_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: - if isinstance(self._provider, SandboxInlineFileProvider): - await self._provider.write_file(handle, target_path, data) - return with tempfile.TemporaryDirectory(prefix="nemo-gym-sandbox-upload-") as tmp_dir: source_path = Path(tmp_dir) / "contents" if isinstance(data, str): @@ -109,7 +74,6 @@ async def start( self, spec: SandboxSpec | None = None, *, - outside_endpoints: list[OutsideEndpoint] | None = None, delete_on_stop: bool | None = None, ) -> "AsyncSandbox": if self._closed: @@ -120,18 +84,16 @@ async def start( if requested_spec is None: raise ValueError("Sandbox.start() requires a SandboxSpec") - requested_spec = self._with_outside_endpoints(requested_spec, outside_endpoints) - resolved_spec = await self._resolve_image_build(requested_spec) - handle = await self._provider.create(resolved_spec) + handle = await self._provider.create(requested_spec) try: - await self._write_initial_files(handle, resolved_spec.files) + await self._write_initial_files(handle, requested_spec.files) except Exception: await self._provider.close(handle, delete=True) await self._provider.aclose() self._closed = True raise - self._spec = resolved_spec + self._spec = requested_spec self._handle = handle self._delete_on_stop = self._delete_on_stop if delete_on_stop is None else delete_on_stop self._stopped = False @@ -166,8 +128,6 @@ async def status(self) -> SandboxStatus: return SandboxStatus.UNKNOWN if self._stopped: return SandboxStatus.STOPPED - if not isinstance(self._provider, SandboxStatusProvider): - return SandboxStatus.UNKNOWN return await self._provider.status(self._handle) async def stop(self, *, delete: bool | None = None) -> None: @@ -269,14 +229,12 @@ def start( self, spec: SandboxSpec | None = None, *, - outside_endpoints: list[OutsideEndpoint] | None = None, delete_on_stop: bool | None = None, ) -> "Sandbox": self._runner.run( "start", lambda: self._async_sandbox.start( spec, - outside_endpoints=outside_endpoints, delete_on_stop=delete_on_stop, ), ) diff --git a/nemo_gym/sandbox/providers/__init__.py b/nemo_gym/sandbox/providers/__init__.py index 20405c18e2..3614eac34f 100644 --- a/nemo_gym/sandbox/providers/__init__.py +++ b/nemo_gym/sandbox/providers/__init__.py @@ -16,21 +16,13 @@ from nemo_gym.sandbox.providers.base import ( ExecResult, - ImageBuildRequest, - ImageSpec, - OutsideEndpoint, - SandboxAddressProvider, SandboxCreateError, SandboxCreateVerificationError, SandboxExecResult, SandboxHandle, - SandboxImageBuildProvider, - SandboxInlineFileProvider, SandboxProvider, SandboxSpec, SandboxStatus, - SandboxStatusProvider, - VolumeMount, ) from nemo_gym.sandbox.providers.registry import ( create_provider, @@ -42,21 +34,13 @@ __all__ = [ "ExecResult", - "ImageBuildRequest", - "ImageSpec", - "OutsideEndpoint", - "SandboxAddressProvider", "SandboxCreateError", "SandboxCreateVerificationError", "SandboxExecResult", "SandboxHandle", - "SandboxImageBuildProvider", - "SandboxInlineFileProvider", "SandboxProvider", "SandboxSpec", "SandboxStatus", - "SandboxStatusProvider", - "VolumeMount", "create_provider", "get_provider_class", "list_providers", diff --git a/nemo_gym/sandbox/providers/base.py b/nemo_gym/sandbox/providers/base.py index ae73564983..f4629d9871 100644 --- a/nemo_gym/sandbox/providers/base.py +++ b/nemo_gym/sandbox/providers/base.py @@ -12,66 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Provider-facing sandbox protocol. +"""Provider-facing sandbox protocol.""" -Providers are the only layer that talks to runtime and infrastructure APIs. -Gym agents and external harnesses consume the public ``nemo_gym.sandbox`` API -instead of importing provider-specific modules. -""" - -from collections.abc import Callable from dataclasses import dataclass, field from enum import Enum from pathlib import Path -from typing import Any, Protocol, runtime_checkable - - -@dataclass(frozen=True) -class ImageSpec: - """Provider-neutral image build input. - - ``image`` is the target image reference the sandbox should run. ``source`` - describes where the provider or image builder can get the build context - from, for example a Git checkout, local path, archive, or provider-native - image recipe. - """ - - image: str - source: dict[str, Any] = field(default_factory=dict) - - -@dataclass(frozen=True) -class ImageBuildRequest: - """Request for building one or more sandbox images before creation.""" - - specs: list[ImageSpec] - docker_build_fn: Callable[[ImageSpec], str] | None = None - codebuild_buildspec_fn: Callable[[ImageSpec], str] | None = None - - -@dataclass(frozen=True) -class OutsideEndpoint: - """Endpoint exposed outside the sandbox and passed in through an env var.""" - - url: str - env_var: str - - -@dataclass(frozen=True) -class VolumeMount: - """Provider-neutral volume mount description.""" - - host_path: str | None = None - container_path: str = "/workspace" - readonly: bool = False - efs_filesystem_id: str | None = None - efs_root_directory: str | None = None - efs_access_point_id: str | None = None - - @property - def is_efs(self) -> bool: - """Return whether this mount describes an EFS-backed volume.""" - return self.efs_filesystem_id is not None +from typing import Any, Protocol class SandboxStatus(str, Enum): @@ -86,10 +32,9 @@ class SandboxStatus(str, Enum): @dataclass(frozen=True) class SandboxSpec: - """Provider-neutral sandbox creation request.""" + """Sandbox creation request.""" image: str | None = None - image_build: ImageSpec | None = None timeout_s: int | float | None = None ready_timeout_s: int | float | None = None workdir: str | None = None @@ -98,9 +43,6 @@ class SandboxSpec: metadata: dict[str, str] = field(default_factory=dict) resources: dict[str, str] = field(default_factory=dict) entrypoint: list[str] | None = None - volumes: list[VolumeMount] = field(default_factory=list) - environment_dir: str | None = None - extensions: dict[str, str] = field(default_factory=dict) provider_options: dict[str, Any] = field(default_factory=dict) @@ -108,10 +50,9 @@ class SandboxSpec: class SandboxHandle: """Provider-neutral handle to a created sandbox. - ``raw`` is provider-owned opaque state, such as an SDK sandbox object, - transport session, or lightweight provider reference. Public Gym code - should pass it back to the provider through this handle rather than - inspecting or mutating it directly. + ``raw`` is provider-owned opaque state. Public code should pass it back to + the provider through this handle rather than inspecting or mutating it + directly. """ sandbox_id: str @@ -182,50 +123,14 @@ async def download_file(self, handle: SandboxHandle, source_path: str, target_pa """Download one sandbox file to the local filesystem.""" ... - async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: - """Close provider resources and optionally delete the sandbox.""" - ... - - async def aclose(self) -> None: - """Close provider-scoped resources such as SDK clients.""" - ... - - -@runtime_checkable -class SandboxInlineFileProvider(Protocol): - """Optional provider trait for efficient inline file reads and writes.""" - - async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: - """Write a small file into a sandbox without a local staging path.""" - ... - - async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: - """Read a small file from a sandbox without a local staging path.""" - ... - - -@runtime_checkable -class SandboxStatusProvider(Protocol): - """Optional provider trait for sandbox lifecycle status.""" - async def status(self, handle: SandboxHandle) -> SandboxStatus: """Return the current sandbox lifecycle status.""" ... - -@runtime_checkable -class SandboxAddressProvider(Protocol): - """Optional provider trait for sandbox network addressing.""" - - async def container_ip(self, handle: SandboxHandle) -> str | None: - """Return the sandbox container IP when the provider exposes one.""" + async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: + """Close provider resources and optionally delete the sandbox.""" ... - -@runtime_checkable -class SandboxImageBuildProvider(Protocol): - """Optional provider trait for building images before sandbox creation.""" - - async def build_images(self, request: ImageBuildRequest) -> list[str]: - """Build images and return the image references that were produced.""" + async def aclose(self) -> None: + """Close provider-scoped resources such as SDK clients.""" ... diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 7c394bf99c..6bcbe895ac 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -31,7 +31,6 @@ SandboxHandle, SandboxSpec, SandboxStatus, - VolumeMount, ) @@ -136,25 +135,6 @@ def _require_tenacity() -> tuple[Any, Any, Any, Any]: return AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential -def _httpx_retryable_types() -> tuple[type[BaseException], ...]: - try: - import httpx - except ModuleNotFoundError: - return tuple() - return ( - httpx.RemoteProtocolError, - httpx.ReadError, - httpx.WriteError, - httpx.ConnectError, - httpx.ConnectTimeout, - httpx.ReadTimeout, - httpx.WriteTimeout, - httpx.PoolTimeout, - httpx.TimeoutException, - httpx.NetworkError, - ) - - def _has_retryable_error_marker(exception: BaseException) -> bool: message = str(exception).lower() return any(marker in message for marker in RETRYABLE_ERROR_MARKERS) @@ -207,9 +187,6 @@ def _is_retryable_create_error(exception: BaseException) -> bool: return True if isinstance(exception, (ConnectionError, OSError, TimeoutError)): return True - httpx_types = _httpx_retryable_types() - if httpx_types and isinstance(exception, httpx_types): - return True try: from opensandbox.exceptions import ( @@ -261,9 +238,6 @@ def _is_retryable_sdk_operation_error(exception: BaseException) -> bool: return True if isinstance(exception, (ConnectionError, OSError)): return True - httpx_types = _httpx_retryable_types() - if httpx_types and isinstance(exception, httpx_types): - return True return _is_retryable_create_error(exception) @@ -314,7 +288,6 @@ def _normalize_spec(spec: SandboxSpec) -> SandboxSpec: env=_string_map(spec.env), metadata=_metadata_map(spec.metadata), resources=_string_map(spec.resources), - extensions=_string_map(spec.extensions), ) @@ -323,41 +296,20 @@ def _to_platform_spec(platform: dict[str, Any]) -> Any: return PlatformSpec(**platform) -def _volume_mount_name(volume: VolumeMount, index: int) -> str: - source = volume.container_path or volume.host_path or f"volume-{index}" - normalized = METADATA_VALUE_RE.sub("-", source.strip("/")).strip("-") - return (normalized or f"volume-{index}")[:63] - - -def _volume_to_mapping(volume: VolumeMount | Mapping[str, Any], index: int) -> dict[str, Any]: - if isinstance(volume, Mapping): - return dict(volume) - if not isinstance(volume, VolumeMount): - raise TypeError(f"OpenSandbox volume entries must be mappings or VolumeMount instances, got {type(volume)!r}") - if volume.is_efs: - raise ValueError("OpenSandbox does not support provider-neutral EFS VolumeMount entries") - if volume.host_path is None: - raise ValueError("OpenSandbox VolumeMount entries require host_path") - return { - "name": _volume_mount_name(volume, index), - "host": {"path": volume.host_path}, - "mount_path": volume.container_path, - "read_only": volume.readonly, - } +def _to_volumes(volumes: list[Mapping[str, Any]]) -> list[Any]: + _, _, _, _, Volume = _require_opensandbox_sdk() + return [Volume(**dict(volume)) for volume in volumes] -def _to_volumes(volumes: list[VolumeMount | Mapping[str, Any]]) -> list[Any]: - _, _, _, _, Volume = _require_opensandbox_sdk() - return [Volume(**_volume_to_mapping(volume, index)) for index, volume in enumerate(volumes)] +def _spec_volumes(spec: SandboxSpec) -> list[Mapping[str, Any]] | None: + return spec.provider_options.get(PROVIDER_OPTION_VOLUMES) -def _spec_volumes(spec: SandboxSpec) -> list[VolumeMount | Mapping[str, Any]] | None: - volumes: list[VolumeMount | Mapping[str, Any]] = [] - volumes.extend(spec.volumes) - provider_volumes = spec.provider_options.get(PROVIDER_OPTION_VOLUMES) - if provider_volumes is not None: - volumes.extend(provider_volumes) - return volumes or None +def _spec_extensions(spec: SandboxSpec) -> dict[str, str]: + value = spec.provider_options.get("extensions", {}) + if not isinstance(value, Mapping): + raise TypeError("OpenSandbox provider option 'extensions' must be a mapping") + return _string_map(dict(value)) def _provider_option_bool(provider_options: dict[str, Any], key: str) -> bool | None: @@ -389,8 +341,6 @@ class OpenSandboxConnectionConfig: domain: str | None = None api_key: str | None = None protocol: str | None = None - use_server_proxy: bool | None = None - exec_use_server_proxy: bool | None = None request_timeout_s: int | None = None @@ -503,7 +453,8 @@ def _with_default_image_pull_policy(self, spec: SandboxSpec) -> SandboxSpec: if self._create.image_pull_policy is None: return spec - extensions = dict(spec.extensions) + provider_options = dict(spec.provider_options) + extensions = _spec_extensions(spec) image_pull_policy = extensions.get(IMAGE_PULL_POLICY_EXTENSION_KEY) or extensions.get( IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY ) @@ -512,13 +463,12 @@ def _with_default_image_pull_policy(self, spec: SandboxSpec) -> SandboxSpec: image_pull_policy = validate_image_pull_policy(image_pull_policy) extensions.setdefault(IMAGE_PULL_POLICY_EXTENSION_KEY, image_pull_policy) extensions.setdefault(IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY, image_pull_policy) - return replace(spec, extensions=extensions) + provider_options["extensions"] = extensions + return replace(spec, provider_options=provider_options) def _connection_config( self, request_timeout_s: int | float | None = None, - *, - use_server_proxy: bool | None = None, ) -> Any: _, ConnectionConfig, _, _, _ = _require_opensandbox_sdk() kwargs: dict[str, Any] = {} @@ -528,37 +478,14 @@ def _connection_config( kwargs["api_key"] = self._connection.api_key if self._connection.protocol is not None: kwargs["protocol"] = self._connection.protocol - if use_server_proxy is None: - use_server_proxy = self._connection.use_server_proxy - if use_server_proxy is not None: - kwargs["use_server_proxy"] = use_server_proxy if request_timeout_s is None: request_timeout_s = self._connection.request_timeout_s if request_timeout_s is not None: kwargs["request_timeout"] = timedelta(seconds=request_timeout_s) return ConnectionConfig(**kwargs) - def _exec_connection_config(self, request_timeout_s: int | float | None = None) -> Any: - """Connection config for SDK handles that issue execd/filesystem calls. - - For clustered evaluations, exec traffic should normally use the - OpenSandbox server proxy. That keeps clients off pod IP routing and lets - the server resolve the sandbox backend for each request. - """ - use_server_proxy = self._connection.use_server_proxy - if self._connection.exec_use_server_proxy is not None: - use_server_proxy = self._connection.exec_use_server_proxy - return self._connection_config( - request_timeout_s=request_timeout_s, - use_server_proxy=use_server_proxy, - ) - async def aclose(self) -> None: - """Close provider-owned resources. - - The provider intentionally does not inject or own OpenSandbox SDK - network clients. SDK handles are closed per sandbox in ``close``. - """ + """Close provider-owned resources.""" return None async def _await_sdk_call( @@ -722,7 +649,7 @@ async def _connect_after_create(self, handle: SandboxHandle, spec: SandboxSpec) sandbox = await asyncio.wait_for( Sandbox.connect( handle.sandbox_id, - connection_config=self._exec_connection_config(request_timeout_s=attempt_timeout_s), + connection_config=self._connection_config(request_timeout_s=attempt_timeout_s), connect_timeout=timedelta(seconds=attempt_timeout_s), skip_health_check=True, ), @@ -747,8 +674,8 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: "env": spec.env, "metadata": spec.metadata, "resource": spec.resources, - "extensions": spec.extensions, - "connection_config": self._exec_connection_config(request_timeout_s=self._create.request_timeout_s), + "extensions": _spec_extensions(spec), + "connection_config": self._connection_config(request_timeout_s=self._create.request_timeout_s), } if spec.image is not None: kwargs["image"] = spec.image @@ -857,14 +784,6 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: raw_status = getattr(info, "status", None) return _to_sandbox_status(getattr(raw_status, "state", None) if raw_status is not None else None) - async def container_ip(self, handle: SandboxHandle) -> str | None: - """Return the container IP when the OpenSandbox SDK handle exposes it.""" - for attr in ("container_ip", "container_ip_address", "pod_ip", "ip"): - value = getattr(handle.raw, attr, None) - if value: - return str(value) - return None - def _command_retry_count(self) -> int: return ( self._operations.retries if self._operations.command_retries is None else self._operations.command_retries @@ -950,7 +869,7 @@ async def exec( retries=self._command_retry_count(), ) - async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + async def _write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: """Write one file into an OpenSandbox sandbox.""" await self._await_sdk_operation( lambda: handle.raw.files.write_file(target_path, data), @@ -961,7 +880,7 @@ async def write_file(self, handle: SandboxHandle, target_path: str, data: str | else None, ) - async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + async def _read_file(self, handle: SandboxHandle, source_path: str) -> bytes: """Read one file from an OpenSandbox sandbox.""" return await self._await_sdk_operation( lambda: handle.raw.files.read_bytes(source_path), @@ -974,12 +893,12 @@ async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: """Upload one local file into an OpenSandbox sandbox.""" - await self.write_file(handle, target_path, source_path.read_bytes()) + await self._write_file(handle, target_path, source_path.read_bytes()) async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: """Download one file from an OpenSandbox sandbox.""" target_path.parent.mkdir(parents=True, exist_ok=True) - target_path.write_bytes(await self.read_file(handle, source_path)) + target_path.write_bytes(await self._read_file(handle, source_path)) async def close(self, handle: SandboxHandle, *, delete: bool) -> None: """Close local SDK resources and optionally terminate the sandbox.""" diff --git a/responses_api_agents/mini_swe_agent_2/README.md b/responses_api_agents/mini_swe_agent_2/README.md index 0e5dba44e7..8ae45b05ae 100644 --- a/responses_api_agents/mini_swe_agent_2/README.md +++ b/responses_api_agents/mini_swe_agent_2/README.md @@ -119,8 +119,6 @@ mini_swe_agent_2: domain: opensandbox-server.opensandbox-system.svc.cluster.local api_key: ${oc.env:OPENSANDBOX_API_KEY} protocol: http - use_server_proxy: true - exec_use_server_proxy: true request_timeout_s: 300 create: request_timeout_s: 1200 @@ -301,11 +299,10 @@ environment: - Validates that a sandbox provider was configured. - Builds a `SandboxSpec` from the task image, environment variables, metadata, - resources, platform, volumes, provider-specific extensions, and health-check - settings. + resources, and provider-specific options. - Adds standard metadata such as `nemo_gym_agent=mini_swe_agent_2` and `instance_id`. -- Creates a `Sandbox` facade and calls `Sandbox.create(...)`. +- Creates a `Sandbox` facade and calls `Sandbox.start(...)`. `execute()`: @@ -329,9 +326,8 @@ command output begins with `COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT` and the command succeeded, it raises `minisweagent.exceptions.Submitted` with the final submission payload. -`cleanup()` calls `Sandbox.close(..., delete=config.delete)` and then -`Sandbox.shutdown()` to release provider-owned async resources and stop the sync -facade's private loop. +`cleanup()` calls `Sandbox.stop(...)` to release provider-owned resources and +stop the sync facade's private loop. ## Contributing diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml index e2dcd5dec1..63f693f937 100644 --- a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml @@ -16,8 +16,6 @@ mini_swe_agent_2: domain: opensandbox-server.opensandbox-system.svc.cluster.local api_key: ${oc.env:OPENSANDBOX_API_KEY} protocol: http - use_server_proxy: true - exec_use_server_proxy: true request_timeout_s: 300 create: request_timeout_s: 1200 diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index e6d43d556d..e603d2dc8d 100644 --- a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -81,7 +81,7 @@ def __init__( image = spec_config.pop("image", None) or self.config.image image = rewrite_image(image, spec_config.pop("image_rewrites", [])) provider_options = dict(spec_config.pop("provider_options", {})) - for option_key in ("platform", "volumes", "skip_health_check"): + for option_key in ("platform", "volumes", "skip_health_check", "extensions"): if option_key in spec_config: provider_options[option_key] = spec_config.pop(option_key) if "snapshot_id" in spec_config: @@ -109,8 +109,6 @@ def __init__( }, resources=spec_config.pop("resources", {}), entrypoint=spec_config.pop("entrypoint", None), - environment_dir=spec_config.pop("environment_dir", None), - extensions=spec_config.pop("extensions", {}), provider_options=provider_options, ), delete_on_stop=self.config.delete, diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index 345c968d1c..b472eaafad 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -23,7 +23,7 @@ import pytest -from nemo_gym.sandbox.providers.base import SandboxSpec, SandboxStatus, VolumeMount +from nemo_gym.sandbox.providers.base import SandboxSpec, SandboxStatus pytest.importorskip("tenacity", reason="tenacity optional sandbox dependency is not installed") @@ -139,9 +139,6 @@ def fake_import( with pytest.raises(ModuleNotFoundError, match="tenacity is required"): opensandbox_provider._require_tenacity() - block_imports("httpx") - assert opensandbox_provider._httpx_retryable_types() == tuple() - block_imports("opensandbox.exceptions") assert opensandbox_provider._is_retryable_create_error(RuntimeError("gateway timeout")) is True @@ -189,25 +186,6 @@ async def test_direct_create_passes_platform_to_sdk_create( def test_provider_validation_and_retry_helpers() -> None: with pytest.raises(ValueError, match="image_pull_policy"): opensandbox_provider.validate_image_pull_policy("Sometimes") - - volume_mapping = opensandbox_provider._volume_to_mapping( - VolumeMount(host_path="/host/workspace", container_path="/mnt/workspace", readonly=True), - 0, - ) - assert volume_mapping == { - "name": "mnt-workspace", - "host": {"path": "/host/workspace"}, - "mount_path": "/mnt/workspace", - "read_only": True, - } - assert opensandbox_provider._volume_to_mapping({"name": "raw-volume"}, 1) == {"name": "raw-volume"} - assert opensandbox_provider._volume_mount_name(VolumeMount(container_path="/"), 2) == "volume-2" - with pytest.raises(TypeError, match="VolumeMount"): - opensandbox_provider._volume_to_mapping(object(), 0) - with pytest.raises(ValueError, match="EFS"): - opensandbox_provider._volume_to_mapping(VolumeMount(efs_filesystem_id="fs-1"), 0) - with pytest.raises(ValueError, match="host_path"): - opensandbox_provider._volume_to_mapping(VolumeMount(), 0) with pytest.raises(TypeError, match="must be a bool"): opensandbox_provider._provider_option_bool({"skip_health_check": "true"}, "skip_health_check") @@ -257,14 +235,12 @@ def test_provider_validation_and_retry_helpers() -> None: assert attrs["next_sleep_s"] == 0.5 -def test_connection_config_exec_proxy_and_image_policy(fake_opensandbox_sdk: None) -> None: +def test_connection_config_and_image_policy(fake_opensandbox_sdk: None) -> None: provider = opensandbox_provider.OpenSandboxProvider( connection={ "domain": "sandbox.example", "api_key": "key", # pragma: allowlist secret "protocol": "https", - "use_server_proxy": True, - "exec_use_server_proxy": False, "request_timeout_s": 10, } ) @@ -274,17 +250,16 @@ def test_connection_config_exec_proxy_and_image_policy(fake_opensandbox_sdk: Non "domain": "sandbox.example", "api_key": "key", # pragma: allowlist secret "protocol": "https", - "use_server_proxy": True, "request_timeout": timedelta(seconds=10), } - exec_config = provider._exec_connection_config(request_timeout_s=3) - assert exec_config.kwargs["use_server_proxy"] is False - assert exec_config.kwargs["request_timeout"] == timedelta(seconds=3) + short_timeout_config = provider._connection_config(request_timeout_s=3) + assert short_timeout_config.kwargs["request_timeout"] == timedelta(seconds=3) - spec = SandboxSpec(image="image:tag", extensions={"imagePullPolicy": "Never"}) + spec = SandboxSpec(image="image:tag", provider_options={"extensions": {"imagePullPolicy": "Never"}}) updated = provider._with_default_image_pull_policy(spec) - assert updated.extensions["imagePullPolicy"] == "Never" - assert updated.extensions["opensandbox.extensions.image-pull-policy"] == "Never" + extensions = updated.provider_options["extensions"] + assert extensions["imagePullPolicy"] == "Never" + assert extensions["opensandbox.extensions.image-pull-policy"] == "Never" no_policy_provider = opensandbox_provider.OpenSandboxProvider(create={"image_pull_policy": None}) assert no_policy_provider._with_default_image_pull_policy(spec) is spec @@ -328,8 +303,6 @@ async def read_bytes(self, source_path: str) -> bytes: return f"bytes:{source_path}".encode() class FakeRaw: - container_ip = "10.1.2.3" - def __init__(self) -> None: self.commands = FakeCommands() self.files = FakeFiles() @@ -374,20 +347,16 @@ async def get_info(self) -> Any: assert result.stderr == "stderr\nCommandError: failed" assert raw.commands.calls[1][0] == "su -s /bin/sh -c fail agent" - await provider.write_file(handle, "/tmp/file.txt", "contents") - assert await provider.read_file(handle, "/tmp/file.txt") == b"bytes:/tmp/file.txt" upload_path = tmp_path / "upload.txt" upload_path.write_text("upload", encoding="utf-8") await provider.upload_file(handle, upload_path, "/remote/upload.txt") download_path = tmp_path / "nested" / "download.txt" await provider.download_file(handle, "/remote/download.txt", download_path) - assert raw.files.writes == [("/tmp/file.txt", "contents"), ("/remote/upload.txt", b"upload")] + assert raw.files.writes == [("/remote/upload.txt", b"upload")] assert download_path.read_bytes() == b"bytes:/remote/download.txt" assert await provider.status(handle) == SandboxStatus.RUNNING - assert await provider.container_ip(handle) == "10.1.2.3" bare_handle = opensandbox_provider.SandboxHandle(sandbox_id="sandbox-2", provider_name="opensandbox", raw=object()) assert await provider.status(bare_handle) == SandboxStatus.UNKNOWN - assert await provider.container_ip(bare_handle) is None async def test_provider_create_probe_and_close_error_paths(monkeypatch: pytest.MonkeyPatch) -> None: @@ -513,7 +482,6 @@ async def test_create_once_and_connect_after_create_error_paths( timeout_s=10, ready_timeout_s=20, entrypoint=["/bin/sh"], - volumes=[VolumeMount(host_path="/host/workspace", container_path="/mnt/workspace", readonly=True)], provider_options={ "snapshot_id": "snapshot-1", "platform": {"os": "linux", "arch": "amd64"}, @@ -528,10 +496,7 @@ async def test_create_once_and_connect_after_create_error_paths( assert FakeSandbox.created_kwargs["ready_timeout"] == timedelta(seconds=20) assert FakeSandbox.created_kwargs["entrypoint"] == ["/bin/sh"] assert FakeSandbox.created_kwargs["platform"] == FakePlatformSpec(os="linux", arch="amd64") - assert FakeSandbox.created_kwargs["volumes"] == [ - VolumeMount(host_path="/host/workspace", container_path="/mnt/workspace", readonly=True), - {"name": "workspace"}, - ] + assert FakeSandbox.created_kwargs["volumes"] == [{"name": "workspace"}] assert FakeSandbox.created_kwargs["skip_health_check"] is True class FailingConnectSandbox(FakeSandbox): @@ -670,18 +635,11 @@ async def test_retry_classification_and_await_sdk_helpers(monkeypatch: pytest.Mo assert await provider.aclose() is None assert await provider._await_sdk_call(_return_value("ok"), operation="op", sandbox_id="sandbox-1", timeout_s=None) assert opensandbox_provider._is_retryable_sdk_operation_error(TimeoutError("command timeout")) is False - assert opensandbox_provider._is_retryable_sdk_operation_error(ConnectionError("proxy failed")) is True + assert opensandbox_provider._is_retryable_sdk_operation_error(ConnectionError("connection failed")) is True wrapped = RuntimeError("wrapper") wrapped.__cause__ = ConnectionError("connection reset") assert opensandbox_provider._is_retryable_sdk_operation_error(wrapped) is True - class FakeHttpxConnectError(Exception): - pass - - monkeypatch.setattr(opensandbox_provider, "_httpx_retryable_types", lambda: (FakeHttpxConnectError,)) - assert opensandbox_provider._is_retryable_create_error(FakeHttpxConnectError("temporary")) is True - assert opensandbox_provider._is_retryable_sdk_operation_error(FakeHttpxConnectError("temporary")) is True - async def cancelled() -> None: raise asyncio.CancelledError() diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index f7fdd5e571..bf068d0f60 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -14,6 +14,7 @@ import asyncio import importlib.util +from datetime import timedelta from pathlib import Path from typing import Any from uuid import uuid4 @@ -23,9 +24,6 @@ import nemo_gym.sandbox.providers.registry as provider_registry from nemo_gym.sandbox import ( AsyncSandbox, - ImageBuildRequest, - ImageSpec, - OutsideEndpoint, Sandbox, SandboxCreateError, SandboxExecResult, @@ -83,19 +81,12 @@ def __init__(self, marker: str = "default") -> None: self.created_specs: list[SandboxSpec] = [] self.created_handles: list[SandboxHandle] = [] self.exec_calls: list[dict[str, Any]] = [] - self.image_build_requests: list[ImageBuildRequest] = [] - self.write_calls: list[tuple[SandboxHandle, str, str | bytes]] = [] - self.read_calls: list[tuple[SandboxHandle, str]] = [] self.upload_calls: list[tuple[SandboxHandle, Path, str]] = [] self.download_calls: list[tuple[SandboxHandle, str, Path]] = [] self.closed: list[tuple[SandboxHandle, bool]] = [] self.aclosed = False FakeSandboxProvider.last_instance = self - async def build_images(self, request: ImageBuildRequest) -> list[str]: - self.image_build_requests.append(request) - return [spec.image for spec in request.specs] - async def create(self, spec: SandboxSpec) -> SandboxHandle: self.created_specs.append(spec) handle = SandboxHandle( @@ -128,13 +119,6 @@ async def exec( ) return SandboxExecResult(stdout="ok", stderr=None, return_code=0) - async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: - self.write_calls.append((handle, target_path, data)) - - async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: - self.read_calls.append((handle, source_path)) - return f"read:{source_path}".encode() - async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: self.upload_calls.append((handle, source_path, target_path)) @@ -147,10 +131,6 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: del handle return SandboxStatus.RUNNING - async def container_ip(self, handle: SandboxHandle) -> str | None: - del handle - return "10.0.0.1" - async def close(self, handle: SandboxHandle, *, delete: bool) -> None: self.closed.append((handle, delete)) @@ -193,6 +173,10 @@ async def download_file(self, handle: SandboxHandle, source_path: str, target_pa del handle, source_path target_path.write_bytes(b"downloaded") + async def status(self, handle: SandboxHandle) -> SandboxStatus: + del handle + return SandboxStatus.UNKNOWN + async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: del handle, delete @@ -237,6 +221,10 @@ async def download_file(self, handle: SandboxHandle, source_path: str, target_pa self.download_calls.append((handle, source_path, target_path)) target_path.write_bytes(b"fallback") + async def status(self, handle: SandboxHandle) -> SandboxStatus: + del handle + return SandboxStatus.RUNNING + async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: del handle, delete @@ -244,16 +232,10 @@ async def aclose(self) -> None: return None -class EmptyImageBuildProvider(FakeSandboxProvider): - async def build_images(self, request: ImageBuildRequest) -> list[str]: - self.image_build_requests.append(request) - return [] - - -class FailingWriteProvider(FakeSandboxProvider): - async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: - self.write_calls.append((handle, target_path, data)) - raise RuntimeError("write failed") +class FailingUploadProvider(FakeSandboxProvider): + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + self.upload_calls.append((handle, source_path, target_path)) + raise RuntimeError("upload failed") def test_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: @@ -272,7 +254,6 @@ async def _assert_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> Non workdir="/repo", files={"/tmp/bootstrap.txt": "hello"}, ), - outside_endpoints=[OutsideEndpoint(url="http://outside", env_var="OUTSIDE_URL")], delete_on_stop=True, ) @@ -282,8 +263,8 @@ async def _assert_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> Non assert provider.marker == "configured" assert provider.created_specs[0].image == "image:tag" assert provider.created_specs[0].metadata == {"suite": "unit"} - assert provider.created_specs[0].env["OUTSIDE_URL"] == "http://outside" - assert provider.write_calls == [(handle, "/tmp/bootstrap.txt", "hello")] + assert provider.upload_calls[0][0] == handle + assert provider.upload_calls[0][2] == "/tmp/bootstrap.txt" result = await sandbox.exec("pytest -q", timeout_s=60, user="agent") assert result == SandboxExecResult(stdout="ok", stderr=None, return_code=0) @@ -302,7 +283,7 @@ async def _assert_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> Non source_path.write_text("local", encoding="utf-8") await sandbox.upload(source_path, "/remote/source.txt") await sandbox.download("/remote/source.txt", target_path) - assert provider.upload_calls == [(handle, source_path, "/remote/source.txt")] + assert provider.upload_calls[1] == (handle, source_path, "/remote/source.txt") assert provider.download_calls == [(handle, "/remote/source.txt", target_path)] assert target_path.read_bytes() == b"downloaded" @@ -312,13 +293,6 @@ async def _assert_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> Non assert await sandbox.status() == SandboxStatus.STOPPED assert provider.aclosed is True - build_provider = FakeSandboxProvider() - built_sandbox = AsyncSandbox(build_provider) - await built_sandbox.start(SandboxSpec(image_build=ImageSpec(image="built:tag", source={"context": "repo"}))) - assert build_provider.image_build_requests[-1].specs[0].source == {"context": "repo"} - assert build_provider.created_specs[-1].image == "built:tag" - await built_sandbox.stop(delete=True) - context_provider = FakeSandboxProvider() async with AsyncSandbox(context_provider) as context_sandbox: await context_sandbox.start(SandboxSpec(image="image:tag"), delete_on_stop=True) @@ -326,20 +300,14 @@ async def _assert_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> Non assert context_provider.closed[-1] == (context_handle, True) -def test_async_sandbox_build_and_initial_file_error_paths() -> None: - asyncio.run(_assert_async_sandbox_build_and_initial_file_error_paths()) +def test_async_sandbox_initial_file_error_paths() -> None: + asyncio.run(_assert_async_sandbox_initial_file_error_paths()) -async def _assert_async_sandbox_build_and_initial_file_error_paths() -> None: - empty_build_provider = EmptyImageBuildProvider() - empty_build_sandbox = AsyncSandbox(empty_build_provider) - with pytest.raises(ValueError, match="build_images returned no image references"): - await empty_build_sandbox.start(SandboxSpec(image_build=ImageSpec(image="missing:tag"))) - assert empty_build_provider.image_build_requests[0].specs[0].image == "missing:tag" - - failing_provider = FailingWriteProvider() +async def _assert_async_sandbox_initial_file_error_paths() -> None: + failing_provider = FailingUploadProvider() failing_sandbox = AsyncSandbox(failing_provider) - with pytest.raises(RuntimeError, match="write failed"): + with pytest.raises(RuntimeError, match="upload failed"): await failing_sandbox.start(SandboxSpec(image="image:tag", files={"/tmp/bootstrap.txt": "hello"})) assert failing_provider.closed == [ ( @@ -457,7 +425,6 @@ def test_sync_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: workdir="/repo", files={"/tmp/bootstrap.txt": "hello"}, ), - outside_endpoints=[OutsideEndpoint(url="http://outside", env_var="OUTSIDE_URL")], delete_on_stop=True, ) @@ -467,8 +434,8 @@ def test_sync_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: assert provider.marker == "configured" assert provider.created_specs[0].image == "image:tag" assert provider.created_specs[0].metadata == {"suite": "unit"} - assert provider.created_specs[0].env["OUTSIDE_URL"] == "http://outside" - assert provider.write_calls == [(handle, "/tmp/bootstrap.txt", "hello")] + assert provider.upload_calls[0][0] == handle + assert provider.upload_calls[0][2] == "/tmp/bootstrap.txt" result = sandbox.exec("pytest -q", timeout_s=60, user="agent") assert result == SandboxExecResult(stdout="ok", stderr=None, return_code=0) @@ -570,7 +537,7 @@ async def create(cls, **kwargs: Any) -> "FakeSDKSandbox": ) provider = OpenSandboxProvider(probe={"command": None}) - monkeypatch.setattr(provider, "_connection_config", lambda request_timeout_s=None, use_server_proxy=None: object()) + monkeypatch.setattr(provider, "_connection_config", lambda request_timeout_s=None: object()) handle = await provider.create( SandboxSpec( @@ -592,11 +559,11 @@ async def create(cls, **kwargs: Any) -> "FakeSDKSandbox": @requires_tenacity -def test_opensandbox_connect_after_create_can_use_direct_exec_endpoint(monkeypatch) -> None: - asyncio.run(_assert_opensandbox_connect_after_create_can_use_direct_exec_endpoint(monkeypatch)) +def test_opensandbox_connect_after_create_uses_connection_config(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_connect_after_create_uses_connection_config(monkeypatch)) -async def _assert_opensandbox_connect_after_create_can_use_direct_exec_endpoint(monkeypatch) -> None: +async def _assert_opensandbox_connect_after_create_uses_connection_config(monkeypatch) -> None: opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() class FakeConnectionConfig: @@ -621,7 +588,7 @@ async def connect(cls, sandbox_id: str, **kwargs: Any) -> "FakeSDKSandbox": ) provider = OpenSandboxProvider( - connection={"use_server_proxy": True, "exec_use_server_proxy": False}, + connection={"domain": "sandbox.example", "protocol": "https"}, create={"connect_attempt_timeout_s": 1}, probe={"command": None}, ) @@ -634,7 +601,11 @@ 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["use_server_proxy"] is False + assert connect_call["connection_config"].kwargs == { + "domain": "sandbox.example", + "protocol": "https", + "request_timeout": timedelta(seconds=1), + } @requires_tenacity @@ -792,7 +763,7 @@ async def run(self, command: str, *, opts: FakeRunCommandOpts) -> FakeExecution: del command, opts self.calls += 1 if self.calls <= 2: - raise ConnectionError("transient proxy failure") + raise ConnectionError("transient connection failure") return FakeExecution() class FakeRaw: @@ -843,7 +814,7 @@ def __init__(self) -> None: async def run(self, command: str, *, opts: FakeRunCommandOpts) -> None: del command, opts self.calls += 1 - raise ConnectionError("transient proxy failure") + raise ConnectionError("transient connection failure") class FakeRaw: def __init__(self) -> None: From 1e3fc588035c7528fe237100636843f9cdba4b48 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Tue, 2 Jun 2026 17:25:19 -0700 Subject: [PATCH 09/14] Use SDK proxy for OpenSandbox smoke config Signed-off-by: Hemil Desai --- nemo_gym/sandbox/providers/opensandbox/provider.py | 3 +++ responses_api_agents/mini_swe_agent_2/README.md | 1 + .../mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml | 1 + tests/unit_tests/test_opensandbox_provider.py | 2 ++ 4 files changed, 7 insertions(+) diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 6bcbe895ac..cb9f12be5e 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -342,6 +342,7 @@ class OpenSandboxConnectionConfig: api_key: str | None = None protocol: str | None = None request_timeout_s: int | None = None + use_server_proxy: bool = False @dataclass(frozen=True) @@ -482,6 +483,8 @@ def _connection_config( request_timeout_s = self._connection.request_timeout_s if request_timeout_s is not None: kwargs["request_timeout"] = timedelta(seconds=request_timeout_s) + if self._connection.use_server_proxy: + kwargs["use_server_proxy"] = True return ConnectionConfig(**kwargs) async def aclose(self) -> None: diff --git a/responses_api_agents/mini_swe_agent_2/README.md b/responses_api_agents/mini_swe_agent_2/README.md index 8ae45b05ae..64e8beae32 100644 --- a/responses_api_agents/mini_swe_agent_2/README.md +++ b/responses_api_agents/mini_swe_agent_2/README.md @@ -120,6 +120,7 @@ mini_swe_agent_2: api_key: ${oc.env:OPENSANDBOX_API_KEY} protocol: http request_timeout_s: 300 + use_server_proxy: true create: request_timeout_s: 1200 timeout_s: 1200 diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml index 63f693f937..3c5e6f503b 100644 --- a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml @@ -17,6 +17,7 @@ mini_swe_agent_2: api_key: ${oc.env:OPENSANDBOX_API_KEY} protocol: http request_timeout_s: 300 + use_server_proxy: true create: request_timeout_s: 1200 timeout_s: 1200 diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index b472eaafad..92c7c49ed1 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -242,6 +242,7 @@ def test_connection_config_and_image_policy(fake_opensandbox_sdk: None) -> None: "api_key": "key", # pragma: allowlist secret "protocol": "https", "request_timeout_s": 10, + "use_server_proxy": True, } ) @@ -251,6 +252,7 @@ def test_connection_config_and_image_policy(fake_opensandbox_sdk: None) -> None: "api_key": "key", # pragma: allowlist secret "protocol": "https", "request_timeout": timedelta(seconds=10), + "use_server_proxy": True, } short_timeout_config = provider._connection_config(request_timeout_s=3) assert short_timeout_config.kwargs["request_timeout"] == timedelta(seconds=3) From 8dd151ca1fc659738ac3e093e67e6aaec213b47e Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Wed, 3 Jun 2026 00:25:54 -0700 Subject: [PATCH 10/14] Fix OpenSandbox retry classification cycles Signed-off-by: Hemil Desai --- .../sandbox/providers/opensandbox/provider.py | 25 ++++++++++--------- tests/unit_tests/test_opensandbox_provider.py | 9 +++++++ 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index cb9f12be5e..1ec7543b79 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -223,22 +223,23 @@ def _is_retryable_create_error(exception: BaseException) -> bool: return _has_retryable_error_marker(exception) -def _is_retryable_sdk_operation_error(exception: BaseException) -> bool: - """Return whether an SDK operation can be retried by Gym. - - The OpenSandbox Python SDK does not retry generated lifecycle, execd, or - filesystem HTTP calls. It converts network failures into SDK exceptions and - exposes API status codes, so classify both the wrapper and its original - cause here. - """ +def _is_retryable_sdk_operation_error(exception: BaseException, seen: set[int] | None = None) -> bool: + """Return whether an SDK operation can be retried.""" if isinstance(exception, TimeoutError): return False - cause = exception.__cause__ - if isinstance(cause, BaseException) and _is_retryable_sdk_operation_error(cause): - return True + seen = set() if seen is None else seen + exception_id = id(exception) + if exception_id in seen: + return False + seen.add(exception_id) if isinstance(exception, (ConnectionError, OSError)): return True - return _is_retryable_create_error(exception) + if _is_retryable_create_error(exception): + return True + cause = exception.__cause__ + if isinstance(cause, BaseException): + return _is_retryable_sdk_operation_error(cause, seen) + return False def _is_missing_sandbox_delete_error(exception: BaseException) -> bool: diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index 92c7c49ed1..36c4573007 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -641,6 +641,15 @@ async def test_retry_classification_and_await_sdk_helpers(monkeypatch: pytest.Mo wrapped = RuntimeError("wrapper") wrapped.__cause__ = ConnectionError("connection reset") assert opensandbox_provider._is_retryable_sdk_operation_error(wrapped) is True + wrapped.__cause__ = wrapped + assert opensandbox_provider._is_retryable_sdk_operation_error(wrapped) is False + + from opensandbox.exceptions import SandboxApiException # noqa: PLC0415 + + cyclic_api_error = SandboxApiException("proxy failed") + cyclic_api_error.status_code = 500 + cyclic_api_error.__cause__ = cyclic_api_error + assert opensandbox_provider._is_retryable_sdk_operation_error(cyclic_api_error) is True async def cancelled() -> None: raise asyncio.CancelledError() From d669369db2b9c0bbbc0524706e4752148fbcce4e Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Mon, 8 Jun 2026 16:24:54 -0700 Subject: [PATCH 11/14] Stop migrating provider options in mini-swe sandbox Signed-off-by: Hemil Desai --- .../mini_swe_agent_2/sandbox_environment.py | 5 --- tests/unit_tests/test_sandbox.py | 31 +++++++++++++++++++ 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index e603d2dc8d..91ab4b2404 100644 --- a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -81,11 +81,6 @@ def __init__( image = spec_config.pop("image", None) or self.config.image image = rewrite_image(image, spec_config.pop("image_rewrites", [])) provider_options = dict(spec_config.pop("provider_options", {})) - for option_key in ("platform", "volumes", "skip_health_check", "extensions"): - if option_key in spec_config: - provider_options[option_key] = spec_config.pop(option_key) - if "snapshot_id" in spec_config: - provider_options["snapshot_id"] = spec_config.pop("snapshot_id") env = dict(spec_config.pop("env", {})) for key in self.config.forward_env: diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index bf068d0f60..27b3c24f6d 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -959,6 +959,37 @@ def test_mini_swe_sandbox_environment_owns_conda_setup(monkeypatch) -> None: assert FakeSandboxProvider.last_instance.closed[0][1] is True +def test_mini_swe_sandbox_environment_only_uses_explicit_provider_options() -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + + env = MiniSWESandboxEnvironment( + image="image:tag", + provider={provider_name: {}}, + spec={ + "provider_options": { + "platform": {"os": "linux", "arch": "amd64"}, + "snapshot_id": "snapshot-1", + }, + "platform": {"os": "ignored", "arch": "ignored"}, + "extensions": {"imagePullPolicy": "Never"}, + "snapshot_id": "ignored-snapshot", + "skip_health_check": True, + "volumes": [{"name": "ignored"}], + }, + ) + + try: + provider = FakeSandboxProvider.last_instance + assert provider is not None + assert provider.created_specs[0].provider_options == { + "platform": {"os": "linux", "arch": "amd64"}, + "snapshot_id": "snapshot-1", + } + finally: + env.cleanup() + + def test_mini_swe_sandbox_environment_validation_and_context_manager() -> None: with pytest.raises(ValueError, match="requires provider"): MiniSWESandboxEnvironment(image="image:tag") From 768a8bd25a201e6f07b6a31789c9460c17a3f864 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Mon, 8 Jun 2026 22:52:53 -0700 Subject: [PATCH 12/14] Simplify initial sandbox file uploads Signed-off-by: Hemil Desai --- nemo_gym/sandbox/api.py | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index 1d58c931c7..5e5bb0c246 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -57,19 +57,6 @@ def _require_handle(self) -> SandboxHandle: raise RuntimeError("Sandbox has not been started") return self._handle - async def _write_inline_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: - with tempfile.TemporaryDirectory(prefix="nemo-gym-sandbox-upload-") as tmp_dir: - source_path = Path(tmp_dir) / "contents" - if isinstance(data, str): - source_path.write_text(data, encoding="utf-8") - else: - source_path.write_bytes(data) - await self._provider.upload_file(handle, source_path, target_path) - - async def _write_initial_files(self, handle: SandboxHandle, files: dict[str, str]) -> None: - for target_path, contents in files.items(): - await self._write_inline_file(handle, target_path, contents) - async def start( self, spec: SandboxSpec | None = None, @@ -86,7 +73,13 @@ async def start( handle = await self._provider.create(requested_spec) try: - await self._write_initial_files(handle, requested_spec.files) + if requested_spec.files: + with tempfile.TemporaryDirectory(prefix="nemo-gym-sandbox-upload-") as tmp_dir: + tmp_path = Path(tmp_dir) + for index, (target_path, contents) in enumerate(requested_spec.files.items()): + source_path = tmp_path / f"file-{index}" + source_path.write_text(contents, encoding="utf-8") + await self._provider.upload_file(handle, source_path, target_path) except Exception: await self._provider.close(handle, delete=True) await self._provider.aclose() From 64030426d84b4b7f881b363950d9ed491962eca1 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Tue, 9 Jun 2026 09:37:28 -0700 Subject: [PATCH 13/14] Clarify sandbox lifecycle and resources Signed-off-by: Hemil Desai --- nemo_gym/sandbox/__init__.py | 2 + nemo_gym/sandbox/api.py | 30 ++----- nemo_gym/sandbox/providers/__init__.py | 2 + nemo_gym/sandbox/providers/base.py | 42 +++++++++- .../sandbox/providers/opensandbox/provider.py | 82 +++++++++++-------- .../mini_swe_agent_2/README.md | 11 ++- responses_api_agents/mini_swe_agent_2/app.py | 2 +- .../configs/mini_swe_agent_opensandbox.yaml | 9 +- .../mini_swe_agent_2/sandbox_environment.py | 10 +-- .../mini_swe_agent_2/tests/test_app.py | 10 +-- tests/unit_tests/test_opensandbox_provider.py | 38 +++++---- tests/unit_tests/test_sandbox.py | 76 +++++++++-------- 12 files changed, 179 insertions(+), 135 deletions(-) diff --git a/nemo_gym/sandbox/__init__.py b/nemo_gym/sandbox/__init__.py index 4754ff483a..820fcbff71 100644 --- a/nemo_gym/sandbox/__init__.py +++ b/nemo_gym/sandbox/__init__.py @@ -22,6 +22,7 @@ SandboxExecResult, SandboxHandle, SandboxProvider, + SandboxResources, SandboxSpec, SandboxStatus, create_provider, @@ -41,6 +42,7 @@ "SandboxExecResult", "SandboxHandle", "SandboxProvider", + "SandboxResources", "SandboxSpec", "SandboxStatus", "create_provider", diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index 5e5bb0c246..5bb00a0cc2 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -42,13 +42,10 @@ def __init__( self, provider: Mapping[str, Any] | SandboxProvider, spec: SandboxSpec | None = None, - *, - delete_on_stop: bool = False, ) -> None: self._provider = create_provider(provider) if isinstance(provider, Mapping) else provider self._spec = spec self._handle: SandboxHandle | None = None - self._delete_on_stop = delete_on_stop self._stopped = True self._closed = False @@ -60,8 +57,6 @@ def _require_handle(self) -> SandboxHandle: async def start( self, spec: SandboxSpec | None = None, - *, - delete_on_stop: bool | None = None, ) -> "AsyncSandbox": if self._closed: raise RuntimeError("Sandbox has been stopped") @@ -81,14 +76,13 @@ async def start( source_path.write_text(contents, encoding="utf-8") await self._provider.upload_file(handle, source_path, target_path) except Exception: - await self._provider.close(handle, delete=True) + await self._provider.close(handle) await self._provider.aclose() self._closed = True raise self._spec = requested_spec self._handle = handle - self._delete_on_stop = self._delete_on_stop if delete_on_stop is None else delete_on_stop self._stopped = False return self @@ -123,16 +117,13 @@ async def status(self) -> SandboxStatus: return SandboxStatus.STOPPED return await self._provider.status(self._handle) - async def stop(self, *, delete: bool | None = None) -> None: + async def stop(self) -> None: if self._closed: return try: if self._handle is not None and not self._stopped: self._stopped = True - await self._provider.close( - self._handle, - delete=self._delete_on_stop if delete is None else delete, - ) + await self._provider.close(self._handle) finally: await self._provider.aclose() self._closed = True @@ -204,14 +195,12 @@ def __init__( self, provider: Mapping[str, Any] | SandboxProvider, spec: SandboxSpec | None = None, - *, - delete_on_stop: bool = False, ) -> None: self._runner = _AsyncLoopRunner() try: self._async_sandbox = self._runner.call( "__init__", - lambda: AsyncSandbox(provider, spec, delete_on_stop=delete_on_stop), + lambda: AsyncSandbox(provider, spec), ) except BaseException: self._runner.close() @@ -221,15 +210,10 @@ def __init__( def start( self, spec: SandboxSpec | None = None, - *, - delete_on_stop: bool | None = None, ) -> "Sandbox": self._runner.run( "start", - lambda: self._async_sandbox.start( - spec, - delete_on_stop=delete_on_stop, - ), + lambda: self._async_sandbox.start(spec), ) return self @@ -264,12 +248,12 @@ def status(self) -> SandboxStatus: return SandboxStatus.STOPPED return self._runner.run("status", self._async_sandbox.status) - def stop(self, *, delete: bool | None = None) -> None: + def stop(self) -> None: if self._closed: return self._closed = True try: - self._runner.run("stop", lambda: self._async_sandbox.stop(delete=delete)) + self._runner.run("stop", self._async_sandbox.stop) finally: self._runner.close() diff --git a/nemo_gym/sandbox/providers/__init__.py b/nemo_gym/sandbox/providers/__init__.py index 3614eac34f..d1b2eed496 100644 --- a/nemo_gym/sandbox/providers/__init__.py +++ b/nemo_gym/sandbox/providers/__init__.py @@ -21,6 +21,7 @@ SandboxExecResult, SandboxHandle, SandboxProvider, + SandboxResources, SandboxSpec, SandboxStatus, ) @@ -39,6 +40,7 @@ "SandboxExecResult", "SandboxHandle", "SandboxProvider", + "SandboxResources", "SandboxSpec", "SandboxStatus", "create_provider", diff --git a/nemo_gym/sandbox/providers/base.py b/nemo_gym/sandbox/providers/base.py index f4629d9871..b8bca468bb 100644 --- a/nemo_gym/sandbox/providers/base.py +++ b/nemo_gym/sandbox/providers/base.py @@ -14,6 +14,7 @@ """Provider-facing sandbox protocol.""" +from collections.abc import Mapping from dataclasses import dataclass, field from enum import Enum from pathlib import Path @@ -30,21 +31,54 @@ class SandboxStatus(str, Enum): UNKNOWN = "unknown" +@dataclass(frozen=True) +class SandboxResources: + """Provider-neutral resource request.""" + + cpu: float | None = None + memory_mib: int | None = None + disk_gib: int | None = None + gpu: int | None = None + gpu_type: str | None = None + + @classmethod + def from_mapping(cls, resources: Mapping[str, Any] | None) -> "SandboxResources": + if resources is None: + return cls() + allowed_keys = set(cls.__dataclass_fields__) + unknown_keys = set(resources) - allowed_keys + if unknown_keys: + unknown = ", ".join(sorted(unknown_keys)) + allowed = ", ".join(sorted(allowed_keys)) + raise ValueError(f"Unknown sandbox resource keys: {unknown}. Expected keys: {allowed}") + return cls( + cpu=float(resources["cpu"]) if resources.get("cpu") is not None else None, + memory_mib=int(resources["memory_mib"]) if resources.get("memory_mib") is not None else None, + disk_gib=int(resources["disk_gib"]) if resources.get("disk_gib") is not None else None, + gpu=int(resources["gpu"]) if resources.get("gpu") is not None else None, + gpu_type=str(resources["gpu_type"]) if resources.get("gpu_type") is not None else None, + ) + + @dataclass(frozen=True) class SandboxSpec: """Sandbox creation request.""" image: str | None = None - timeout_s: int | float | None = None + ttl_s: int | float | None = None ready_timeout_s: int | float | None = None workdir: str | None = None env: dict[str, str] = field(default_factory=dict) files: dict[str, str] = field(default_factory=dict) metadata: dict[str, str] = field(default_factory=dict) - resources: dict[str, str] = field(default_factory=dict) + resources: SandboxResources | Mapping[str, Any] = field(default_factory=SandboxResources) entrypoint: list[str] | None = None provider_options: dict[str, Any] = field(default_factory=dict) + def __post_init__(self) -> None: + if not isinstance(self.resources, SandboxResources): + object.__setattr__(self, "resources", SandboxResources.from_mapping(self.resources)) + @dataclass class SandboxHandle: @@ -127,8 +161,8 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: """Return the current sandbox lifecycle status.""" ... - async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: - """Close provider resources and optionally delete the sandbox.""" + async def close(self, handle: SandboxHandle) -> None: + """End the sandbox lifecycle and close provider resources for it.""" ... async def aclose(self) -> None: diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 1ec7543b79..b75befacb8 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -29,6 +29,7 @@ SandboxCreateVerificationError, SandboxExecResult, SandboxHandle, + SandboxResources, SandboxSpec, SandboxStatus, ) @@ -269,10 +270,31 @@ def _log_operation_retry(retry_state: Any) -> None: ) -def _string_map(values: dict[str, Any]) -> dict[str, str]: +def _string_map(values: Mapping[str, Any]) -> dict[str, str]: return {str(key): str(value) for key, value in values.items()} +def _resource_quantity(value: float | int) -> str: + if isinstance(value, float) and value.is_integer(): + return str(int(value)) + return str(value) + + +def _resource_map(resources: SandboxResources) -> dict[str, str]: + values: dict[str, str] = {} + if resources.cpu is not None: + values["cpu"] = _resource_quantity(resources.cpu) + if resources.memory_mib is not None: + values["memory"] = f"{resources.memory_mib}Mi" + if resources.disk_gib is not None: + values["ephemeral-storage"] = f"{resources.disk_gib}Gi" + if resources.gpu is not None: + values["gpu"] = str(resources.gpu) + if resources.gpu_type is not None: + values["gpu_type"] = resources.gpu_type + return values + + def _metadata_value(value: Any) -> str: normalized = METADATA_VALUE_RE.sub("_", str(value)).strip("._-") normalized = normalized[:63].strip("._-") @@ -288,7 +310,6 @@ def _normalize_spec(spec: SandboxSpec) -> SandboxSpec: spec, env=_string_map(spec.env), metadata=_metadata_map(spec.metadata), - resources=_string_map(spec.resources), ) @@ -618,7 +639,7 @@ async def _verify_created_handle(self, handle: SandboxHandle) -> None: async def _cleanup_failed_create_handle(self, handle: SandboxHandle) -> None: try: - await self.close(handle, delete=True) + await self.close(handle) except Exception as e: LOGGER.warning( "Failed to clean up OpenSandbox sandbox after create probe failure; sandbox_id=%s; error=%r", @@ -677,7 +698,7 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: kwargs: dict[str, Any] = { "env": spec.env, "metadata": spec.metadata, - "resource": spec.resources, + "resource": _resource_map(spec.resources), "extensions": _spec_extensions(spec), "connection_config": self._connection_config(request_timeout_s=self._create.request_timeout_s), } @@ -686,8 +707,8 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: snapshot_id = spec.provider_options.get(PROVIDER_OPTION_SNAPSHOT_ID) if snapshot_id is not None: kwargs["snapshot_id"] = snapshot_id - if spec.timeout_s is not None: - kwargs["timeout"] = timedelta(seconds=spec.timeout_s) + if spec.ttl_s is not None: + kwargs["timeout"] = timedelta(seconds=spec.ttl_s) if spec.ready_timeout_s is not None: kwargs["ready_timeout"] = timedelta(seconds=spec.ready_timeout_s) if spec.entrypoint is not None: @@ -904,25 +925,24 @@ async def download_file(self, handle: SandboxHandle, source_path: str, target_pa target_path.parent.mkdir(parents=True, exist_ok=True) target_path.write_bytes(await self._read_file(handle, source_path)) - async def close(self, handle: SandboxHandle, *, delete: bool) -> None: - """Close local SDK resources and optionally terminate the sandbox.""" - kill_error: Exception | None = None - if delete: - try: - await self._await_sdk_operation( - lambda: handle.raw.kill(), - operation="kill", - sandbox_id=handle.sandbox_id, - timeout_s=self._operations.close_timeout_s, + async def close(self, handle: SandboxHandle) -> None: + """Terminate the sandbox and close local SDK resources.""" + stop_error: Exception | None = None + try: + await self._await_sdk_operation( + lambda: handle.raw.kill(), + operation="kill", + sandbox_id=handle.sandbox_id, + timeout_s=self._operations.close_timeout_s, + ) + except Exception as e: + if not _is_missing_sandbox_delete_error(e): + stop_error = e + else: + LOGGER.info( + "OpenSandbox sandbox %r was already deleted during close", + handle.sandbox_id, ) - except Exception as e: - if not _is_missing_sandbox_delete_error(e): - kill_error = e - else: - LOGGER.info( - "OpenSandbox sandbox %r was already deleted during close", - handle.sandbox_id, - ) close_error: Exception | None = None try: @@ -940,15 +960,13 @@ async def close(self, handle: SandboxHandle, *, delete: bool) -> None: e, ) - if kill_error is not None: + if stop_error is not None: if close_error is not None: raise RuntimeError( - "Failed to delete and close OpenSandbox sandbox " - f"{handle.sandbox_id!r}: delete_error={kill_error!r}, " + "Failed to stop and close OpenSandbox sandbox " + f"{handle.sandbox_id!r}: stop_error={stop_error!r}, " f"close_error={close_error!r}" - ) from kill_error - raise kill_error + ) from stop_error + raise stop_error if close_error is not None: - if delete: - return - raise close_error + return diff --git a/responses_api_agents/mini_swe_agent_2/README.md b/responses_api_agents/mini_swe_agent_2/README.md index 64e8beae32..1d5bfe6ab7 100644 --- a/responses_api_agents/mini_swe_agent_2/README.md +++ b/responses_api_agents/mini_swe_agent_2/README.md @@ -140,12 +140,12 @@ mini_swe_agent_2: command_retries: 3 close_timeout_s: 30 sandbox_spec: - timeout_s: 18000 + ttl_s: 18000 ready_timeout_s: 1200 resources: - cpu: "2" - memory: 8Gi - ephemeral-storage: 20Gi + cpu: 2 + memory_mib: 8192 + disk_gib: 20 provider_options: platform: os: linux @@ -159,7 +159,6 @@ mini_swe_agent_2: conda_env: testbed activate_conda: true user: root - delete: true run_golden: false step_timeout: 600 eval_timeout: 1800 @@ -225,7 +224,7 @@ ng_run "+config_paths=[$CONFIG_PATHS]" \ +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.eval_timeout=1800 \ +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.step_limit=50 \ +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.run_golden=false \ - '+mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.sandbox_spec.resources={cpu: 500m, memory: 4Gi, ephemeral-storage: 8Gi}' \ + '+mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.sandbox_spec.resources={cpu: 0.5, memory_mib: 4096, disk_gib: 8}' \ '+mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.sandbox_spec.metadata={benchmark: swebench-verified, harness: mini_swe_agent_2, endpoint_label: hosted-vllm, run_family: mini-swe-agent-2-pass8}' ``` diff --git a/responses_api_agents/mini_swe_agent_2/app.py b/responses_api_agents/mini_swe_agent_2/app.py index 2de794e1ec..bc59d39968 100644 --- a/responses_api_agents/mini_swe_agent_2/app.py +++ b/responses_api_agents/mini_swe_agent_2/app.py @@ -144,7 +144,7 @@ def _bash_tool_choice() -> dict[str, Any]: def _sandbox_spec_for_instance( spec: dict[str, Any] | None, *, - resource_profiles: list[dict[str, str]] | None, + resource_profiles: list[dict[str, Any]] | None, instance_id: str, ) -> dict[str, Any]: instance_spec = dict(spec or {}) diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml index 3c5e6f503b..6276feceb6 100644 --- a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml @@ -37,12 +37,12 @@ mini_swe_agent_2: command_retries: 3 close_timeout_s: 30 sandbox_spec: - timeout_s: 18000 + ttl_s: 18000 ready_timeout_s: 1200 resources: - cpu: "2" - memory: 8Gi - ephemeral-storage: 20Gi + cpu: 2 + memory_mib: 8192 + disk_gib: 20 provider_options: platform: os: linux @@ -56,7 +56,6 @@ mini_swe_agent_2: conda_env: testbed activate_conda: true user: root - delete: true run_golden: false step_timeout: 600 eval_timeout: 1800 diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index 91ab4b2404..8e6acdf488 100644 --- a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -32,7 +32,7 @@ def __init__(self, *messages: dict[str, Any]) -> None: super().__init__() -from nemo_gym.sandbox import Sandbox, SandboxSpec +from nemo_gym.sandbox import Sandbox, SandboxResources, SandboxSpec from nemo_gym.sandbox.utils import rewrite_image @@ -58,7 +58,6 @@ class MiniSWESandboxEnvironmentConfig: conda_env: str | None = None activate_conda: bool = False user: str | int | None = "root" - delete: bool = True class MiniSWESandboxEnvironment: @@ -92,7 +91,7 @@ def __init__( self._sandbox = Sandbox(self.config.provider).start( SandboxSpec( image=image, - timeout_s=spec_config.pop("timeout_s", None), + ttl_s=spec_config.pop("ttl_s", None), ready_timeout_s=spec_config.pop("ready_timeout_s", None), workdir=spec_config.pop("workdir", self.config.cwd), env=env, @@ -102,11 +101,10 @@ def __init__( "nemo_gym_agent": "mini_swe_agent_2", "instance_id": (self.config.instance_id or "unknown")[:63], }, - resources=spec_config.pop("resources", {}), + resources=SandboxResources.from_mapping(spec_config.pop("resources", {})), entrypoint=spec_config.pop("entrypoint", None), provider_options=provider_options, - ), - delete_on_stop=self.config.delete, + ) ) def get_template_vars(self, **kwargs: Any) -> dict[str, Any]: 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 77feb64695..67693f4b7a 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 @@ -300,17 +300,17 @@ def test_response_param_helpers_cover_metadata_and_tool_choice_modes(self) -> No def test_sandbox_resource_profiles_override_static_resources(self) -> None: spec = _sandbox_spec_for_instance( - {"resources": {"cpu": "1", "memory": "8Gi", "ephemeral-storage": "20Gi"}}, + {"resources": {"cpu": 1, "memory_mib": 8192, "disk_gib": 20}}, resource_profiles=[ - {"cpu": "250m", "memory": "3Gi", "ephemeral-storage": "1Gi"}, - {"cpu": "500m", "memory": "4Gi", "ephemeral-storage": "1Gi"}, + {"cpu": 0.25, "memory_mib": 3072, "disk_gib": 1}, + {"cpu": 0.5, "memory_mib": 4096, "disk_gib": 1}, ], instance_id="django__django-12345", ) assert spec["resources"] in ( - {"cpu": "250m", "memory": "3Gi", "ephemeral-storage": "1Gi"}, - {"cpu": "500m", "memory": "4Gi", "ephemeral-storage": "1Gi"}, + {"cpu": 0.25, "memory_mib": 3072, "disk_gib": 1}, + {"cpu": 0.5, "memory_mib": 4096, "disk_gib": 1}, ) assert _sandbox_spec_for_instance(None, resource_profiles=None, instance_id="task") == {} diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index 36c4573007..1350525104 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -23,7 +23,7 @@ import pytest -from nemo_gym.sandbox.providers.base import SandboxSpec, SandboxStatus +from nemo_gym.sandbox.providers.base import SandboxResources, SandboxSpec, SandboxStatus pytest.importorskip("tenacity", reason="tenacity optional sandbox dependency is not installed") @@ -411,15 +411,14 @@ async def cancelled_probe(*_args: Any, **_kwargs: Any) -> opensandbox_provider.S provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) - async def close_raises(_handle: Any, *, delete: bool) -> None: - del delete + async def close_raises(_handle: Any) -> None: raise RuntimeError("close failed") monkeypatch.setattr(provider, "close", close_raises) await provider._cleanup_failed_create_handle(handle) provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) - class DeleteAlreadyGoneRaw: + class StopAlreadyGoneRaw: async def kill(self) -> None: raise RuntimeError("sandbox sandbox-1 not found") @@ -430,43 +429,40 @@ async def close(self) -> None: opensandbox_provider.SandboxHandle( sandbox_id="sandbox-1", provider_name="opensandbox", - raw=DeleteAlreadyGoneRaw(), + raw=StopAlreadyGoneRaw(), ), - delete=True, ) - class DeleteAndCloseFailRaw: + class StopAndCloseFailRaw: async def kill(self) -> None: - raise RuntimeError("delete failed") + raise RuntimeError("stop failed") async def close(self) -> None: raise RuntimeError("close failed") - with pytest.raises(RuntimeError, match="Failed to delete and close"): + with pytest.raises(RuntimeError, match="Failed to stop and close"): await provider.close( opensandbox_provider.SandboxHandle( sandbox_id="sandbox-2", provider_name="opensandbox", - raw=DeleteAndCloseFailRaw(), + raw=StopAndCloseFailRaw(), ), - delete=True, ) - class DeleteFailsCloseSucceedsRaw: + class StopFailsCloseSucceedsRaw: async def kill(self) -> None: - raise RuntimeError("delete failed") + raise RuntimeError("stop failed") async def close(self) -> None: return None - with pytest.raises(RuntimeError, match="delete failed"): + with pytest.raises(RuntimeError, match="stop failed"): await provider.close( opensandbox_provider.SandboxHandle( sandbox_id="sandbox-3", provider_name="opensandbox", - raw=DeleteFailsCloseSucceedsRaw(), + raw=StopFailsCloseSucceedsRaw(), ), - delete=True, ) @@ -481,8 +477,9 @@ async def test_create_once_and_connect_after_create_error_paths( monkeypatch.setattr(opensandbox_provider, "_to_volumes", lambda volumes: volumes) spec = SandboxSpec( image="image:tag", - timeout_s=10, + ttl_s=10, ready_timeout_s=20, + resources=SandboxResources(cpu=2, memory_mib=8192, disk_gib=20, gpu=1, gpu_type="H100"), entrypoint=["/bin/sh"], provider_options={ "snapshot_id": "snapshot-1", @@ -496,6 +493,13 @@ async def test_create_once_and_connect_after_create_error_paths( assert FakeSandbox.created_kwargs["snapshot_id"] == "snapshot-1" assert FakeSandbox.created_kwargs["timeout"] == timedelta(seconds=10) assert FakeSandbox.created_kwargs["ready_timeout"] == timedelta(seconds=20) + assert FakeSandbox.created_kwargs["resource"] == { + "cpu": "2", + "memory": "8192Mi", + "ephemeral-storage": "20Gi", + "gpu": "1", + "gpu_type": "H100", + } assert FakeSandbox.created_kwargs["entrypoint"] == ["/bin/sh"] assert FakeSandbox.created_kwargs["platform"] == FakePlatformSpec(os="linux", arch="amd64") assert FakeSandbox.created_kwargs["volumes"] == [{"name": "workspace"}] diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index 27b3c24f6d..6e8fe91ad8 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -28,6 +28,7 @@ SandboxCreateError, SandboxExecResult, SandboxHandle, + SandboxResources, SandboxSpec, SandboxStatus, create_provider, @@ -83,7 +84,7 @@ def __init__(self, marker: str = "default") -> None: self.exec_calls: list[dict[str, Any]] = [] self.upload_calls: list[tuple[SandboxHandle, Path, str]] = [] self.download_calls: list[tuple[SandboxHandle, str, Path]] = [] - self.closed: list[tuple[SandboxHandle, bool]] = [] + self.closed: list[SandboxHandle] = [] self.aclosed = False FakeSandboxProvider.last_instance = self @@ -131,8 +132,8 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: del handle return SandboxStatus.RUNNING - async def close(self, handle: SandboxHandle, *, delete: bool) -> None: - self.closed.append((handle, delete)) + async def close(self, handle: SandboxHandle) -> None: + self.closed.append(handle) async def aclose(self) -> None: self.aclosed = True @@ -177,8 +178,8 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: del handle return SandboxStatus.UNKNOWN - async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: - del handle, delete + async def close(self, handle: SandboxHandle) -> None: + del handle async def aclose(self) -> None: return None @@ -225,8 +226,8 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: del handle return SandboxStatus.RUNNING - async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: - del handle, delete + async def close(self, handle: SandboxHandle) -> None: + del handle async def aclose(self) -> None: return None @@ -254,7 +255,6 @@ async def _assert_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> Non workdir="/repo", files={"/tmp/bootstrap.txt": "hello"}, ), - delete_on_stop=True, ) provider = FakeSandboxProvider.last_instance @@ -289,15 +289,15 @@ async def _assert_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> Non await sandbox.stop() await sandbox.stop() - assert provider.closed[-1] == (handle, True) + assert provider.closed[-1] == handle assert await sandbox.status() == SandboxStatus.STOPPED assert provider.aclosed is True context_provider = FakeSandboxProvider() async with AsyncSandbox(context_provider) as context_sandbox: - await context_sandbox.start(SandboxSpec(image="image:tag"), delete_on_stop=True) + await context_sandbox.start(SandboxSpec(image="image:tag")) context_handle = context_provider.created_handles[0] - assert context_provider.closed[-1] == (context_handle, True) + assert context_provider.closed[-1] == context_handle def test_async_sandbox_initial_file_error_paths() -> None: @@ -315,8 +315,7 @@ async def _assert_async_sandbox_initial_file_error_paths() -> None: sandbox_id="fake-1", provider_name="fake", raw={"spec": SandboxSpec(image="image:tag", files={"/tmp/bootstrap.txt": "hello"})}, - ), - True, + ) ) ] @@ -338,6 +337,14 @@ def test_rewrite_image_validation() -> None: assert rewrite_image("image:tag", [{"from": "other/", "to": "mirror/"}]) == "image:tag" +def test_sandbox_resources_validation() -> None: + spec = SandboxSpec(resources={"cpu": "0.5", "memory_mib": "4096", "disk_gib": "8"}) + assert spec.resources == SandboxResources(cpu=0.5, memory_mib=4096, disk_gib=8) + + with pytest.raises(ValueError, match="Unknown sandbox resource keys"): + SandboxSpec(resources={"memory": "4Gi"}) + + def test_provider_registry_validation_and_listing(monkeypatch: pytest.MonkeyPatch) -> None: provider_name = f"fake-{uuid4().hex}" register_provider(provider_name, FakeSandboxProvider) @@ -425,7 +432,6 @@ def test_sync_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: workdir="/repo", files={"/tmp/bootstrap.txt": "hello"}, ), - delete_on_stop=True, ) provider = FakeSandboxProvider.last_instance @@ -456,7 +462,7 @@ def test_sync_sandbox_facade_uses_public_provider_api(tmp_path: Path) -> None: sandbox.download("/tmp/sync-download.txt", download_path) assert download_path.read_bytes() == b"downloaded" sandbox.stop() - assert provider.closed[-1] == (handle, True) + assert provider.closed[-1] == handle assert sandbox.status() == SandboxStatus.STOPPED assert provider.aclosed is True try: @@ -849,11 +855,11 @@ def __init__(self) -> None: @requires_tenacity -def test_opensandbox_close_timeout_does_not_fail_after_delete() -> None: - asyncio.run(_assert_opensandbox_close_timeout_does_not_fail_after_delete()) +def test_opensandbox_close_timeout_does_not_fail_after_stop() -> None: + asyncio.run(_assert_opensandbox_close_timeout_does_not_fail_after_stop()) -async def _assert_opensandbox_close_timeout_does_not_fail_after_delete() -> None: +async def _assert_opensandbox_close_timeout_does_not_fail_after_stop() -> None: _opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() class SlowCloseRaw: @@ -873,35 +879,34 @@ async def close(self) -> None: ) handle = SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=raw) - await provider.close(handle, delete=True) + await provider.close(handle) assert raw.killed is True @requires_tenacity -def test_opensandbox_close_timeout_still_fails_without_delete() -> None: - asyncio.run(_assert_opensandbox_close_timeout_still_fails_without_delete()) +def test_opensandbox_close_propagates_stop_failure() -> None: + asyncio.run(_assert_opensandbox_close_propagates_stop_failure()) -async def _assert_opensandbox_close_timeout_still_fails_without_delete() -> None: +async def _assert_opensandbox_close_propagates_stop_failure() -> None: _opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() - class SlowCloseRaw: + class StopFailureRaw: + async def kill(self) -> None: + raise RuntimeError("stop failed") + async def close(self) -> None: - await asyncio.sleep(60) + return None provider = OpenSandboxProvider( operations={"close_timeout_s": 0.01}, probe={"command": None}, ) - handle = SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=SlowCloseRaw()) + handle = SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=StopFailureRaw()) - try: - await provider.close(handle, delete=False) - except TimeoutError: - pass - else: - raise AssertionError("expected close timeout to fail when delete=False") + with pytest.raises(RuntimeError, match="stop failed"): + await provider.close(handle) def test_mini_swe_sandbox_environment_owns_conda_setup(monkeypatch) -> None: @@ -916,14 +921,13 @@ def test_mini_swe_sandbox_environment_owns_conda_setup(monkeypatch) -> None: spec={ "image_rewrites": [{"from": "upstream/", "to": "mirror/"}], "metadata": {"suite": "unit"}, - "resources": {"cpu": "1"}, + "resources": {"cpu": 1}, }, env={"STATIC_KEY": "static-value"}, forward_env=["FORWARDED_KEY"], conda_env="testbed", activate_conda=True, user="agent", - delete=True, ) try: @@ -942,6 +946,7 @@ def test_mini_swe_sandbox_environment_owns_conda_setup(monkeypatch) -> None: "FORWARDED_KEY": "forwarded-value", "STATIC_KEY": "static-value", } + assert provider.created_specs[0].resources == SandboxResources(cpu=1.0) result = env.execute("pytest -q", is_eval=True) assert result == {"output": "ok", "returncode": 0, "exception_info": ""} @@ -956,7 +961,7 @@ def test_mini_swe_sandbox_environment_owns_conda_setup(monkeypatch) -> None: env.cleanup() assert FakeSandboxProvider.last_instance is not None - assert FakeSandboxProvider.last_instance.closed[0][1] is True + assert FakeSandboxProvider.last_instance.closed[0].sandbox_id == "fake-1" def test_mini_swe_sandbox_environment_only_uses_explicit_provider_options() -> None: @@ -999,14 +1004,13 @@ def test_mini_swe_sandbox_environment_validation_and_context_manager() -> None: with MiniSWESandboxEnvironment( image="image:tag", provider={provider_name: {}}, - delete=False, ) as env: assert env._sandbox is not None assert FakeSandboxProvider.last_instance is not None assert FakeSandboxProvider.last_instance.created_handles[0].sandbox_id == "fake-1" assert FakeSandboxProvider.last_instance is not None - assert FakeSandboxProvider.last_instance.closed[-1][1] is False + assert FakeSandboxProvider.last_instance.closed[-1].sandbox_id == "fake-1" def test_mini_swe_sandbox_environment_submit_sentinel() -> None: From 3706a55fa740da1d8348d079185e923bce866cc1 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Wed, 10 Jun 2026 08:26:24 -0700 Subject: [PATCH 14/14] Address sandbox review edge cases Signed-off-by: Hemil Desai --- nemo_gym/sandbox/api.py | 43 ++++++++++++-- .../sandbox/providers/opensandbox/provider.py | 8 +-- .../mini_swe_agent_2/README.md | 2 +- responses_api_agents/mini_swe_agent_2/app.py | 56 ++++++++++++++++++- .../configs/mini_swe_agent_opensandbox.yaml | 2 +- .../mini_swe_agent_2/sandbox_environment.py | 14 ++--- .../mini_swe_agent_2/tests/test_app.py | 41 +++++++++++++- .../tests/test_sandbox_environment.py | 41 +++++++++++++- tests/unit_tests/test_sandbox.py | 30 ++++++++-- 9 files changed, 203 insertions(+), 34 deletions(-) diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index 5bb00a0cc2..329baa8ad9 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -19,6 +19,7 @@ import threading from collections.abc import Awaitable, Callable, Mapping from concurrent.futures import Future +from concurrent.futures import TimeoutError as FutureTimeoutError from pathlib import Path from typing import Any, TypeVar @@ -33,6 +34,8 @@ T = TypeVar("T") +SYNC_OPERATION_TIMEOUT_S = 3600.0 +SYNC_LOOP_CLOSE_TIMEOUT_S = 5.0 class AsyncSandbox: @@ -138,7 +141,14 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: class _AsyncLoopRunner: """Run async sandbox operations for sync callers.""" - def __init__(self) -> None: + def __init__( + self, + *, + wait_timeout_s: float = SYNC_OPERATION_TIMEOUT_S, + close_timeout_s: float = SYNC_LOOP_CLOSE_TIMEOUT_S, + ) -> None: + self._wait_timeout_s = wait_timeout_s + self._close_timeout_s = close_timeout_s self._loop = asyncio.new_event_loop() self._ready = threading.Event() self._closed = False @@ -160,23 +170,42 @@ def _ensure_can_block(self, operation: str) -> None: return raise RuntimeError(f"Sandbox.{operation}() is blocking; use AsyncSandbox in async code instead.") + def _wait_for_result(self, operation: str, future: Future[T]) -> T: + try: + return future.result(timeout=self._wait_timeout_s) + except FutureTimeoutError as e: + future.cancel() + raise TimeoutError( + f"Sandbox.{operation}() timed out waiting for the sync loop after {self._wait_timeout_s:g}s" + ) from e + def call(self, operation: str, func: Callable[[], T]) -> T: self._ensure_can_block(operation) future: Future[T] = Future() def invoke() -> None: try: - future.set_result(func()) + result = func() except BaseException as e: - future.set_exception(e) + if not future.cancelled(): + future.set_exception(e) + else: + if not future.cancelled(): + future.set_result(result) self._loop.call_soon_threadsafe(invoke) - return future.result() + return self._wait_for_result(operation, future) def run(self, operation: str, awaitable_factory: Callable[[], Awaitable[T]]) -> T: self._ensure_can_block(operation) future = asyncio.run_coroutine_threadsafe(awaitable_factory(), self._loop) - return future.result() + try: + return future.result(timeout=self._wait_timeout_s) + except FutureTimeoutError as e: + future.cancel() + raise TimeoutError( + f"Sandbox.{operation}() timed out waiting for the sync loop after {self._wait_timeout_s:g}s" + ) from e def close(self) -> None: if self._closed: @@ -184,7 +213,9 @@ def close(self) -> None: self._closed = True if not self._loop.is_closed(): self._loop.call_soon_threadsafe(self._loop.stop) - self._thread.join(timeout=5) + self._thread.join(timeout=self._close_timeout_s) + if self._thread.is_alive(): + return self._loop.close() diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index b75befacb8..8bac478ed8 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -427,7 +427,7 @@ class OpenSandboxOperationConfig: retries: int = 3 retry_delay_s: float = 1.0 retry_max_delay_s: float = 15.0 - command_retries: int | None = None + command_retries: int = 0 close_timeout_s: float | None = 30.0 def __post_init__(self) -> None: @@ -437,7 +437,7 @@ def __post_init__(self) -> None: raise ValueError("operations.retry_delay_s must be >= 0") if self.retry_max_delay_s < 0: raise ValueError("operations.retry_max_delay_s must be >= 0") - if self.command_retries is not None and self.command_retries < 0: + if self.command_retries < 0: raise ValueError("operations.command_retries must be >= 0") if self.close_timeout_s is not None and self.close_timeout_s <= 0: raise ValueError("operations.close_timeout_s must be > 0") @@ -810,9 +810,7 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: return _to_sandbox_status(getattr(raw_status, "state", None) if raw_status is not None else None) def _command_retry_count(self) -> int: - return ( - self._operations.retries if self._operations.command_retries is None else self._operations.command_retries - ) + return self._operations.command_retries async def _exec( self, diff --git a/responses_api_agents/mini_swe_agent_2/README.md b/responses_api_agents/mini_swe_agent_2/README.md index 1d5bfe6ab7..1170f8a215 100644 --- a/responses_api_agents/mini_swe_agent_2/README.md +++ b/responses_api_agents/mini_swe_agent_2/README.md @@ -137,7 +137,7 @@ mini_swe_agent_2: retries: 5 retry_delay_s: 1.0 retry_max_delay_s: 45.0 - command_retries: 3 + command_retries: 0 close_timeout_s: 30 sandbox_spec: ttl_s: 18000 diff --git a/responses_api_agents/mini_swe_agent_2/app.py b/responses_api_agents/mini_swe_agent_2/app.py index bc59d39968..febe8fe696 100644 --- a/responses_api_agents/mini_swe_agent_2/app.py +++ b/responses_api_agents/mini_swe_agent_2/app.py @@ -15,10 +15,12 @@ import asyncio import hashlib import json +import os import sys import time import traceback from asyncio import Semaphore +from copy import deepcopy from pathlib import Path from typing import Any, Callable, Literal, Optional, cast from uuid import uuid4 @@ -51,6 +53,10 @@ ) +OPENSANDBOX_PROVIDER_NAME = "opensandbox" +OPENSANDBOX_API_KEY_ENV = "OPENSANDBOX_API_KEY" # pragma: allowlist secret + + class MiniSWEAgentConfig(BaseResponsesAPIAgentConfig): model_server: ModelServerRef env: Literal["sandbox"] @@ -137,6 +143,47 @@ def _responses_create_params_to_model_kwargs( return model_kwargs +def _opensandbox_connection(provider: dict[str, Any] | None) -> dict[str, Any] | None: + if provider is None: + return None + provider_config = provider.get(OPENSANDBOX_PROVIDER_NAME) + if not isinstance(provider_config, dict): + return None + connection = provider_config.get("connection") + if not isinstance(connection, dict): + return None + return connection + + +def _sandbox_provider_for_config_dump(provider: dict[str, Any]) -> dict[str, Any]: + provider_for_disk = deepcopy(provider) + connection = _opensandbox_connection(provider_for_disk) + if connection is not None: + connection.pop("api_key", None) + return provider_for_disk + + +def _sandbox_runtime_env(provider: dict[str, Any] | None) -> dict[str, Any]: + runtime_env: dict[str, Any] = {"py_executable": sys.executable} + connection = _opensandbox_connection(provider) + if connection is None: + return runtime_env + api_key = connection.get("api_key") + if api_key: + runtime_env["env_vars"] = {OPENSANDBOX_API_KEY_ENV: str(api_key)} + return runtime_env + + +def _restore_sandbox_provider_secrets(config: dict[str, Any]) -> None: + provider = config.get("environment", {}).get("provider") + connection = _opensandbox_connection(provider if isinstance(provider, dict) else None) + if connection is None or connection.get("api_key"): + return + api_key = os.getenv(OPENSANDBOX_API_KEY_ENV) + if api_key: + connection["api_key"] = api_key + + def _bash_tool_choice() -> dict[str, Any]: return {"type": "function", "function": {"name": "bash"}} @@ -449,6 +496,7 @@ def _run_mini_swe_v2(**params: Any) -> dict[str, Any]: instance_dir.mkdir(parents=True, exist_ok=True) config = yaml.safe_load(get_config_path(params["config"]).read_text()) + _restore_sandbox_provider_secrets(config) model_config = config.setdefault("model", {}) model_config["model_class"] = "litellm" model_config["model_name"] = params["model"] @@ -700,7 +748,7 @@ async def run(self, body: MiniSWEAgentRunRequest) -> MiniSWEAgentVerifyResponse: if self.config.sandbox_provider is None: raise ValueError("mini_swe_agent_2 requires sandbox_provider") config.setdefault("environment", {}).update(self.config.sandbox_environment_kwargs or {}) - config["environment"]["provider"] = self.config.sandbox_provider + config["environment"]["provider"] = _sandbox_provider_for_config_dump(self.config.sandbox_provider) config["environment"]["spec"] = _sandbox_spec_for_instance( self.config.sandbox_spec, resource_profiles=self.config.sandbox_resource_profiles, @@ -742,7 +790,11 @@ async def run(self, body: MiniSWEAgentRunRequest) -> MiniSWEAgentVerifyResponse: eval_timeout=eval_timeout, step_limit=step_limit, ) - future = runner_ray_remote.remote(run_mini_swe_with_sandbox, params) + runner = runner_ray_remote + runtime_env = _sandbox_runtime_env(self.config.sandbox_provider) + 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) result = result[instance_id] input_messages = result["input_messages"] diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml index 6276feceb6..c6fffd12ca 100644 --- a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml @@ -34,7 +34,7 @@ mini_swe_agent_2: retries: 5 retry_delay_s: 1.0 retry_max_delay_s: 45.0 - command_retries: 3 + command_retries: 0 close_timeout_s: 30 sandbox_spec: ttl_s: 18000 diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index 8e6acdf488..704122b9ce 100644 --- a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -120,17 +120,11 @@ def serialize(self) -> dict[str, Any]: } } - def _command(self, command: str, cwd: str) -> str: + def _command(self, command: str) -> str: if not self.config.activate_conda or not self.config.conda_env: return command - quoted_cwd = shlex.quote(cwd) quoted_env = shlex.quote(self.config.conda_env) - return ( - f"cd {quoted_cwd} && " - "source $(conda info --base)/etc/profile.d/conda.sh && " - f"conda activate {quoted_env} && " - f"{command}" - ) + return f"source $(conda info --base)/etc/profile.d/conda.sh && conda activate {quoted_env} && {command}" def execute( self, @@ -146,9 +140,9 @@ def execute( raise RuntimeError("Sandbox is not available") result = self._sandbox.exec( - self._command(command, exec_cwd), + self._command(command), timeout_s=timeout_s, - cwd="/", + cwd=exec_cwd, user=self.config.user, ) output = "\n".join(part for part in (result.stdout, result.stderr) if part) 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 67693f4b7a..57ce79ad09 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 @@ -47,6 +47,7 @@ from responses_api_agents.mini_swe_agent_2 import app as mini_swe_app_module from responses_api_agents.mini_swe_agent_2.app import ( + OPENSANDBOX_API_KEY_ENV, MiniSWEAgent, MiniSWEAgentConfig, MiniSWEAgentRunRequest, @@ -56,6 +57,8 @@ _message_content_to_text, _responses_create_params_to_model_kwargs, _run_mini_swe_v2, + _sandbox_provider_for_config_dump, + _sandbox_runtime_env, _sandbox_spec_for_instance, _split_trajectory_for_responses, _swebench_config_path, @@ -176,6 +179,7 @@ def setup_run_mini_swe_mock( # Mock the Ray remote function to return a future-like object mock_future = MagicMock() 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 @@ -314,6 +318,23 @@ def test_sandbox_resource_profiles_override_static_resources(self) -> None: ) assert _sandbox_spec_for_instance(None, resource_profiles=None, instance_id="task") == {} + def test_sandbox_provider_config_dump_strips_api_key(self) -> None: + provider = { + "opensandbox": { + "connection": { + "domain": "sandbox.example", + "api_key": "fixture-value", # pragma: allowlist secret + } + } + } + + provider_for_disk = _sandbox_provider_for_config_dump(provider) + assert "api_key" not in provider_for_disk["opensandbox"]["connection"] + assert provider["opensandbox"]["connection"]["api_key"] == "fixture-value" # pragma: allowlist secret + assert _sandbox_runtime_env(provider)["env_vars"] == { + OPENSANDBOX_API_KEY_ENV: "fixture-value" # pragma: allowlist secret + } + def test_split_trajectory_and_resolution_helpers_cover_edge_cases(self) -> None: input_messages, output_items, raw_responses = _split_trajectory_for_responses( [ @@ -535,7 +556,7 @@ def get_model(config: dict[str, Any]) -> SimpleNamespace: yaml.safe_dump( { "model": {"model_kwargs": {"max_output_tokens": 99}}, - "environment": {}, + "environment": {"provider": {"opensandbox": {"connection": {}}}}, "agent": {"step_limit": 1, "collapse_limit": 3}, } ), @@ -544,6 +565,7 @@ def get_model(config: dict[str, Any]) -> SimpleNamespace: monkeypatch.setattr(mini_swe_app_module, "get_config_path", lambda _config: config_path) monkeypatch.setattr(mini_swe_app_module, "uuid4", lambda: "uuid") monkeypatch.setattr(mini_swe_app_module.time, "time", lambda: 1234) + monkeypatch.setenv(OPENSANDBOX_API_KEY_ENV, "worker-value") # pragma: allowlist secret params = { "instance_dict": { @@ -570,6 +592,10 @@ def get_model(config: dict[str, Any]) -> SimpleNamespace: env = holder["env"] assert env.cleaned is True assert env.config["environment_class"].endswith("MiniSWESandboxEnvironment") + assert ( + env.config["provider"]["opensandbox"]["connection"]["api_key"] + == "worker-value" # pragma: allowlist secret + ) assert env.config["image"] == "docker.io/swebench/sweb.eval.x86_64.django_1776_django-123:latest" assert holder["model_config"]["model_class"] == "litellm" assert holder["model_config"]["model_name"] == "hosted/model" @@ -693,6 +719,14 @@ async def test_run_writes_generation_params_to_config( monkeypatch.chdir(tmp_path) config = create_test_config() config.tool_choice = "bash" + config.sandbox_provider = { + "opensandbox": { + "connection": { + "domain": "sandbox.example", + "api_key": "fixture-value", # pragma: allowlist secret + } + } + } mock_server_client = MagicMock(spec=ServerClient) server = MiniSWEAgent(config=config, server_client=mock_server_client) @@ -712,9 +746,12 @@ async def test_run_writes_generation_params_to_config( await server.run(run_request) - call_args = mock_runner_ray_remote.remote.call_args + runtime_env = mock_runner_ray_remote.options.call_args.kwargs["runtime_env"] + assert runtime_env["env_vars"] == {OPENSANDBOX_API_KEY_ENV: "fixture-value"} # pragma: allowlist secret + call_args = mock_runner_ray_remote.options.return_value.remote.call_args params = call_args.args[1] generated_config = yaml.safe_load(Path(params["config"]).read_text()) + assert "api_key" not in generated_config["environment"]["provider"]["opensandbox"]["connection"] model_kwargs = generated_config["model"]["model_kwargs"] assert model_kwargs["temperature"] == 0.6 assert model_kwargs["top_p"] == 0.95 diff --git a/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py index 331d732641..0419e15b85 100644 --- a/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py @@ -13,7 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -from responses_api_agents.mini_swe_agent_2.sandbox_environment import MiniSWESandboxEnvironment, Submitted +from typing import Any + +from responses_api_agents.mini_swe_agent_2.sandbox_environment import ( + MiniSWESandboxEnvironment, + MiniSWESandboxEnvironmentConfig, + Submitted, +) def test_check_finished_raises_submitted_for_submit_sentinel() -> None: @@ -49,3 +55,36 @@ def test_check_finished_ignores_nonzero_submit_sentinel() -> None: "exception_info": "", } ) + + +def test_execute_passes_configured_cwd_without_conda_cd() -> None: + class FakeSandbox: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def exec(self, command: str, **kwargs: Any): + self.calls.append({"command": command, **kwargs}) + return type("Result", (), {"stdout": "ok", "stderr": None, "return_code": 0})() + + fake_sandbox = FakeSandbox() + env = MiniSWESandboxEnvironment.__new__(MiniSWESandboxEnvironment) + env.config = MiniSWESandboxEnvironmentConfig( + image="image:tag", + provider={"fake": {}}, + cwd="/default", + activate_conda=False, + ) + env._sandbox = fake_sandbox + + assert env.execute("pwd", cwd="/repo") == {"output": "ok", "returncode": 0, "exception_info": ""} + assert fake_sandbox.calls[-1]["command"] == "pwd" + assert fake_sandbox.calls[-1]["cwd"] == "/repo" + + env.config.activate_conda = True + env.config.conda_env = "testbed" + env.execute("python -V", cwd="/repo") + assert fake_sandbox.calls[-1]["command"] == ( + "source $(conda info --base)/etc/profile.d/conda.sh && conda activate testbed && python -V" + ) + assert "cd /repo" not in fake_sandbox.calls[-1]["command"] + assert fake_sandbox.calls[-1]["cwd"] == "/repo" diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index 6e8fe91ad8..46db01e491 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -14,6 +14,7 @@ import asyncio import importlib.util +import threading from datetime import timedelta from pathlib import Path from typing import Any @@ -479,6 +480,23 @@ def test_sync_loop_runner_close_is_idempotent() -> None: runner.close() +def test_sync_loop_runner_times_out_waits_and_skips_running_loop_close() -> None: + runner = _AsyncLoopRunner(wait_timeout_s=0.01, close_timeout_s=0.01) + release = threading.Event() + + try: + with pytest.raises(TimeoutError, match="timed out waiting for the sync loop"): + runner.call("blocked", release.wait) + + runner.close() + assert runner._thread.is_alive() + finally: + release.set() + runner._thread.join(timeout=1) + if not runner._loop.is_closed(): + runner._loop.close() + + def test_sync_sandbox_file_operations(tmp_path: Path) -> None: provider = FakeSandboxProvider() with Sandbox(provider) as sandbox: @@ -802,11 +820,11 @@ def __init__(self) -> None: @requires_tenacity -def test_opensandbox_command_retries_can_be_disabled(monkeypatch) -> None: - asyncio.run(_assert_opensandbox_command_retries_can_be_disabled(monkeypatch)) +def test_opensandbox_command_retries_default_to_disabled(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_command_retries_default_to_disabled(monkeypatch)) -async def _assert_opensandbox_command_retries_can_be_disabled(monkeypatch) -> None: +async def _assert_opensandbox_command_retries_default_to_disabled(monkeypatch) -> None: opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() class FakeRunCommandOpts: @@ -837,7 +855,6 @@ def __init__(self) -> None: "retries": 2, "retry_delay_s": 0, "retry_max_delay_s": 0, - "command_retries": 0, }, probe={"command": None}, ) @@ -935,7 +952,7 @@ def test_mini_swe_sandbox_environment_owns_conda_setup(monkeypatch) -> None: serialized = env.serialize() assert serialized["info"]["config"]["environment_type"].endswith("MiniSWESandboxEnvironment") env.config.activate_conda = False - assert env._command("echo plain", "/tmp/work") == "echo plain" + assert env._command("echo plain") == "echo plain" env.config.activate_conda = True provider = FakeSandboxProvider.last_instance @@ -951,9 +968,10 @@ def test_mini_swe_sandbox_environment_owns_conda_setup(monkeypatch) -> None: result = env.execute("pytest -q", is_eval=True) assert result == {"output": "ok", "returncode": 0, "exception_info": ""} exec_call = provider.exec_calls[0] - assert exec_call["cwd"] == "/" + assert exec_call["cwd"] == "/testbed" assert exec_call["timeout_s"] == 1800 assert exec_call["user"] == "agent" + assert "cd /testbed" not in exec_call["command"] assert "conda activate testbed" in exec_call["command"] assert exec_call["command"].endswith("pytest -q") finally: