diff --git a/nemo_gym/sandbox/__init__.py b/nemo_gym/sandbox/__init__.py index faa4583d6f..ac7fcceafc 100644 --- a/nemo_gym/sandbox/__init__.py +++ b/nemo_gym/sandbox/__init__.py @@ -39,7 +39,7 @@ list_providers, register_provider, ) -from nemo_gym.sandbox.utils import rewrite_image +from nemo_gym.sandbox.utils import await_cleanup, rewrite_image __all__ = [ @@ -64,6 +64,7 @@ "SupportsSandboxPty", "SupportsSandboxPtyAttach", "create_provider", + "await_cleanup", "get_provider_class", "list_providers", "register_provider", diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index e11d38bc4b..3001e6d132 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -40,6 +40,7 @@ SupportsSandboxPtyAttach, create_provider, ) +from nemo_gym.sandbox.utils import await_cleanup T = TypeVar("T") @@ -289,6 +290,7 @@ def __init__( self._stopped = True self._closed = False self.pty = SandboxPty(self) + self._stop_task: asyncio.Task[None] | None = None def _require_handle(self) -> SandboxHandle: if self._handle is None or self._stopped: @@ -307,7 +309,11 @@ async def start( if requested_spec is None: raise ValueError("Sandbox.start() requires a SandboxSpec") - handle = await self._provider.create(requested_spec) + try: + handle = await self._provider.create(requested_spec) + except BaseException: + await self._provider.aclose() + raise try: if requested_spec.files: with tempfile.TemporaryDirectory(prefix="nemo-gym-sandbox-upload-") as tmp_dir: @@ -316,15 +322,18 @@ async def start( 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) - await self._provider.aclose() - self._closed = True + except BaseException: + try: + await self._provider.close(handle) + finally: + await self._provider.aclose() + self._closed = True raise self._spec = requested_spec self._handle = handle self._stopped = False + return self async def exec( @@ -377,16 +386,20 @@ async def endpoint(self, port: int) -> SandboxEndpoint: return resolved 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) - finally: - await self._provider.aclose() + if self._stop_task is None: self._closed = True + async def cleanup() -> None: + try: + if self._handle is not None and not self._stopped: + self._stopped = True + await self._provider.close(self._handle) + finally: + await self._provider.aclose() + + self._stop_task = asyncio.create_task(cleanup()) + await await_cleanup(self._stop_task) + async def serialize(self, *, scope: str | None = None) -> dict[str, Any]: """Return a JSON descriptor another process can rebuild this box from. diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 2986079228..dbd562b0ed 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -39,6 +39,7 @@ SandboxStatus, ) from nemo_gym.sandbox.providers.utils import coerce_config as _coerce_config +from nemo_gym.sandbox.utils import await_cleanup LOGGER = logging.getLogger(__name__) @@ -661,17 +662,16 @@ def _connection_config( kwargs["request_timeout"] = timedelta(seconds=request_timeout_s) if self._connection.use_server_proxy: kwargs["use_server_proxy"] = True - # The SDK's execd-facing clients (health ping, commands, files) - # send only ConnectionConfig.headers — api_key alone never reaches - # proxied /proxy/* routes, so servers that enforce auth there 401 - # every health ping and create times out at ready_timeout. Inject - # the key only in proxy mode: a direct sandbox endpoint runs - # untrusted code and must never see it. if self._connection.api_key is not None: kwargs["headers"] = {"OPEN-SANDBOX-API-KEY": self._connection.api_key} if self._connection.keepalive_expiry_s is not None or self._connection.disable_connection_pooling: kwargs["transport"] = self._get_transport() - return ConnectionConfig(**kwargs) + config = ConnectionConfig(**kwargs) + if self._connection.use_server_proxy and (api_key := config.get_api_key()): + # Execd-facing SDK clients send only ConnectionConfig.headers. + # Direct endpoints run untrusted code and must never see this key. + config.headers.setdefault("OPEN-SANDBOX-API-KEY", api_key) + return config def _get_transport(self) -> Any: """Return the provider-owned shared transport, building it on first use.""" @@ -1004,29 +1004,35 @@ async def endpoint( ) -> SandboxEndpoint: """Resolve one client-reachable direct or server-proxied service URL.""" + get_endpoint = getattr(handle.raw, "get_endpoint", None) + if get_endpoint is None: + raise NotImplementedError( + "The installed opensandbox SDK does not expose Sandbox.get_endpoint; " + "sandbox service endpoints require opensandbox>=0.1.15" + ) resolved = await self._await_sdk_operation( - lambda: handle.raw.get_endpoint(port), + lambda: get_endpoint(port), operation="get_endpoint", sandbox_id=handle.sandbox_id, timeout_s=( float(self._connection.request_timeout_s) if self._connection.request_timeout_s is not None else None ), ) - endpoint_url = str(resolved.endpoint or "") + endpoint_url = str(getattr(resolved, "endpoint", "") or "") if not endpoint_url: raise RuntimeError(f"OpenSandbox returned an empty endpoint for sandbox {handle.sandbox_id!r} port {port}") + connection = handle.raw.connection_config if "://" not in endpoint_url: # Use the SDK handle's effective configuration so environment- # resolved domains and protocols match the lifecycle request. - scheme = urlsplit(handle.raw.connection_config.get_base_url()).scheme or "http" + scheme = urlsplit(connection.get_base_url()).scheme or "http" endpoint_url = f"{scheme}://{endpoint_url.lstrip('/')}" - headers = dict(handle.raw.connection_config.headers) - # Match the SDK's service adapters: connection-wide headers apply to - # every request, while endpoint-specific routing or auth headers win. - # The upstream proxy-auth fix adds the management API key to - # ConnectionConfig.headers only in server-proxy mode, so direct - # sandbox endpoints never receive it. - headers.update(resolved.headers) + headers = dict(getattr(connection, "headers", None) or {}) + headers.update(getattr(resolved, "headers", None) or {}) + if not getattr(connection, "use_server_proxy", self._connection.use_server_proxy): + # Direct endpoints terminate at untrusted sandbox code and must + # never receive the management credential. + headers.pop("OPEN-SANDBOX-API-KEY", None) return SandboxEndpoint(endpoint=endpoint_url, headers=headers) async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: @@ -1100,8 +1106,11 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: 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) + except BaseException: + # Once create returns an id, cancellation must not strand its + # remote sandbox. close() applies the configured cleanup bounds. + cleanup = asyncio.create_task(self._cleanup_failed_create_handle(created_handle)) + await await_cleanup(cleanup) raise return handle @@ -1496,7 +1505,7 @@ async def download_file(self, handle: SandboxHandle, source_path: str, target_pa async def close(self, handle: SandboxHandle) -> None: """Terminate the sandbox and close local SDK resources.""" - stop_error: Exception | None = None + stop_error: BaseException | None = None try: await self._await_sdk_operation( lambda: handle.raw.kill(), @@ -1504,7 +1513,7 @@ async def close(self, handle: SandboxHandle) -> None: sandbox_id=handle.sandbox_id, timeout_s=self._operations.close_timeout_s, ) - except Exception as e: + except BaseException as e: if not _is_missing_sandbox_delete_error(e): stop_error = e else: @@ -1528,8 +1537,9 @@ async def close(self, handle: SandboxHandle) -> None: handle.sandbox_id, e, ) - if stop_error is not None: + if not isinstance(stop_error, Exception): + raise stop_error if close_error is not None: raise RuntimeError( "Failed to stop and close OpenSandbox sandbox " diff --git a/nemo_gym/sandbox/utils.py b/nemo_gym/sandbox/utils.py index b25f0f6962..b5c6fb4175 100644 --- a/nemo_gym/sandbox/utils.py +++ b/nemo_gym/sandbox/utils.py @@ -14,6 +14,23 @@ """Sandbox utility helpers.""" +import asyncio + + +async def await_cleanup(task: asyncio.Task[None]) -> None: + """Finish owned cleanup before propagating caller cancellation.""" + cancellation: asyncio.CancelledError | None = None + while True: + try: + await asyncio.shield(task) + break + except asyncio.CancelledError as exc: + if task.cancelled(): + raise + cancellation = exc + if cancellation is not None: + raise cancellation + def rewrite_image(image: str | None, rewrites: list[dict[str, str]]) -> str | None: """Apply ordered image-prefix rewrites used by sandbox configs.""" diff --git a/resources_servers/math_formal_lean/app.py b/resources_servers/math_formal_lean/app.py index f81a261607..f4edd96f7d 100644 --- a/resources_servers/math_formal_lean/app.py +++ b/resources_servers/math_formal_lean/app.py @@ -17,10 +17,11 @@ import logging import re +from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Literal, Optional -from pydantic import BaseModel +from pydantic import BaseModel, Field from nemo_gym.base_resources_server import ( BaseResourcesServerConfig, @@ -29,7 +30,7 @@ BaseVerifyResponse, SimpleResourcesServer, ) -from resources_servers.math_formal_lean.sandbox_client import Lean4SandboxClient +from resources_servers.math_formal_lean.sandbox_client import GymSandboxLean4Client, Lean4SandboxClient LOG = logging.getLogger(__name__) @@ -341,6 +342,11 @@ def build_correction_prompt( class MathFormalLeanResourcesServerConfig(BaseResourcesServerConfig): sandbox_host: str = "127.0.0.1" sandbox_port: int = 6000 + # Sandbox backend: local NS HTTP (default) or provider-backed Gym sandboxes. + sandbox_backend: Literal["ns_http", "gym_sandbox"] = "ns_http" + # GymSandboxLean4Client kwargs (provider/image/max_concurrent/...) — read only when + # sandbox_backend == "gym_sandbox". + opensandbox: Dict[str, Any] = Field(default_factory=dict) compilation_timeout: float = 30.0 max_output_characters: int = 1000 extract_code_mode: str = "last" @@ -384,17 +390,41 @@ class MathFormalLeanResourcesServer(SimpleResourcesServer): def model_post_init(self, context: Any) -> None: super().model_post_init(context) - self._sandbox_client = Lean4SandboxClient( - host=self.config.sandbox_host, - port=self.config.sandbox_port, - max_output_characters=self.config.max_output_characters, - ) + if self.config.sandbox_backend == "gym_sandbox": + self._sandbox_client = GymSandboxLean4Client( + max_output_characters=self.config.max_output_characters, + **self.config.opensandbox, + ) + else: + self._sandbox_client = Lean4SandboxClient( + host=self.config.sandbox_host, + port=self.config.sandbox_port, + max_output_characters=self.config.max_output_characters, + ) self._proof_build_config = ProofBuildConfig( extract_code_mode=self.config.extract_code_mode, restate_formal_statement=self.config.restate_formal_statement, strip_theorem_from_proof=self.config.strip_theorem_from_proof, ) + def setup_webserver(self): + app = super().setup_webserver() + main_app_lifespan = app.router.lifespan_context + + @asynccontextmanager + async def lifespan_wrapper(app): + if isinstance(self._sandbox_client, GymSandboxLean4Client): + # A cold pod's first compile can exceed a verify's admission window. + self._sandbox_client.start_pool() + try: + async with main_app_lifespan(app) as maybe_state: + yield maybe_state + finally: + await self._sandbox_client.close() + + app.router.lifespan_context = lifespan_wrapper + return app + async def verify(self, body: MathFormalLeanVerifyRequest) -> MathFormalLeanVerifyResponse: """Verify a proof attempt with multi-turn self-correction support. diff --git a/resources_servers/math_formal_lean/configs/math_formal_lean.yaml b/resources_servers/math_formal_lean/configs/math_formal_lean.yaml index 478816b981..44147e8e2c 100644 --- a/resources_servers/math_formal_lean/configs/math_formal_lean.yaml +++ b/resources_servers/math_formal_lean/configs/math_formal_lean.yaml @@ -4,6 +4,27 @@ math_formal_lean: entrypoint: app.py sandbox_host: ${oc.env:NEMO_SKILLS_SANDBOX_HOST,127.0.0.1} sandbox_port: ${oc.env:NEMO_SKILLS_SANDBOX_PORT,6000} + # Default local NS HTTP or opt-in provider-backed Gym sandboxes. + sandbox_backend: ${oc.env:MATH_FORMAL_LEAN_BACKEND,ns_http} + # Read only when sandbox_backend == gym_sandbox; empty creds/image then = hard startup error. + opensandbox: + provider: + opensandbox: + connection: + domain: ${oc.env:OPENSANDBOX_BASE_URL,""} + api_key: ${oc.env:OPENSANDBOX_API_KEY,""} + use_server_proxy: true + create: + timeout_s: 90 + retries: 3 + # Long compiles poll short status requests instead of holding one stream + # (survives proxy/LB stream caps; Gym PR 2296). + operations: + background_exec: true + image: ${oc.env:NS_SANDBOX_IMAGE,""} + max_concurrent: ${oc.decode:${oc.env:LEAN_SANDBOX_MAX_CONCURRENT,8}} + # 0 creates a pod per verify; N reuses N warmed pods. + pool_size: ${oc.env:LEAN_SANDBOX_POOL_SIZE,0} compilation_timeout: 30.0 domain: math verified: false diff --git a/resources_servers/math_formal_lean/configs/math_formal_lean_multi_turn.yaml b/resources_servers/math_formal_lean/configs/math_formal_lean_multi_turn.yaml index 123e227aed..55c6521497 100644 --- a/resources_servers/math_formal_lean/configs/math_formal_lean_multi_turn.yaml +++ b/resources_servers/math_formal_lean/configs/math_formal_lean_multi_turn.yaml @@ -4,6 +4,27 @@ math_formal_lean: entrypoint: app.py sandbox_host: ${oc.env:NEMO_SKILLS_SANDBOX_HOST,127.0.0.1} sandbox_port: ${oc.env:NEMO_SKILLS_SANDBOX_PORT,6000} + # Default local NS HTTP or opt-in provider-backed Gym sandboxes. + sandbox_backend: ${oc.env:MATH_FORMAL_LEAN_BACKEND,ns_http} + # Read only when sandbox_backend == gym_sandbox; empty creds/image then = hard startup error. + opensandbox: + provider: + opensandbox: + connection: + domain: ${oc.env:OPENSANDBOX_BASE_URL,""} + api_key: ${oc.env:OPENSANDBOX_API_KEY,""} + use_server_proxy: true + create: + timeout_s: 90 + retries: 3 + # Long compiles poll short status requests instead of holding one stream + # (survives proxy/LB stream caps; Gym PR 2296). + operations: + background_exec: true + image: ${oc.env:NS_SANDBOX_IMAGE,""} + max_concurrent: ${oc.decode:${oc.env:LEAN_SANDBOX_MAX_CONCURRENT,8}} + # 0 creates a pod per verify; N reuses N warmed pods. + pool_size: ${oc.env:LEAN_SANDBOX_POOL_SIZE,0} compilation_timeout: 30.0 domain: math verified: false diff --git a/resources_servers/math_formal_lean/requirements.txt b/resources_servers/math_formal_lean/requirements.txt index 318081ec33..0b42bad7a7 100644 --- a/resources_servers/math_formal_lean/requirements.txt +++ b/resources_servers/math_formal_lean/requirements.txt @@ -1,2 +1,2 @@ --e nemo-gym[dev] @ ../../ +-e nemo-gym[dev,sandbox] @ ../../ httpx>=0.27.0 diff --git a/resources_servers/math_formal_lean/sandbox_client.py b/resources_servers/math_formal_lean/sandbox_client.py index e49b0f46d7..cfbda6fae3 100644 --- a/resources_servers/math_formal_lean/sandbox_client.py +++ b/resources_servers/math_formal_lean/sandbox_client.py @@ -20,12 +20,18 @@ - Dockerfile: https://github.com/NVIDIA-NeMo/NeMo-Skills/blob/main/dockerfiles/Dockerfile.sandbox """ +import asyncio import json import logging +import os +import tempfile +import uuid from typing import Any, Dict import httpx +from nemo_gym.sandbox import AsyncSandbox, SandboxSpec, await_cleanup + LOG = logging.getLogger(__name__) @@ -135,3 +141,253 @@ async def health_check(self, timeout: float = 5.0) -> bool: return response.status_code == 200 except httpx.HTTPError: return False + + +class GymSandboxLean4Client: + """Lean4 compilation on per-verify OpenSandbox pods via provider exec. + + Runs the same Lean command and preserves the NS server's result contract. A + configured warm pool reuses prepared sandboxes and replaces failed leases; + otherwise each verification gets a fresh sandbox. + """ + + def __init__( + self, + provider: Dict[str, Any], + image: str, + project_dir: str = "/lean4/my_project", + max_concurrent: int = 8, + acquire_timeout_s: float = 120.0, + create_ttl_s: float = 3600.0, + resources: Dict[str, Any] | None = None, + max_output_characters: int = 1000, + pool_size: int = 0, + prefetch_paths: str = "/root/.elan /lean4", + pool_ref: str = "", + ): + if not image: + raise ValueError("sandbox_backend=gym_sandbox requires a non-empty image") + connection = (provider.get("opensandbox") or {}).get("connection", {}) if provider else {} + if not connection.get("domain") or not connection.get("api_key"): + raise ValueError( + "sandbox_backend=gym_sandbox requires provider connection domain/api_key — " + "set OPENSANDBOX_BASE_URL / OPENSANDBOX_API_KEY" + ) + self._provider = provider + self._image = image + self._project_dir = project_dir.rstrip("/") + self._create_ttl_s = create_ttl_s + self._resources = dict(resources or {}) + self.max_output_characters = max_output_characters + self._acquire_timeout_s = acquire_timeout_s + self._semaphore_size = int(max_concurrent) + self._semaphore = asyncio.Semaphore(self._semaphore_size) + self._pool_size = int(pool_size) + self._prefetch_paths = prefetch_paths + self._pool_ref = pool_ref or "" + self._pool: asyncio.Queue[AsyncSandbox] | None = None + self._pool_sandboxes: set[AsyncSandbox] = set() + self._fill_tasks: set[asyncio.Task[None]] = set() + self._closed = False + self._close_task: asyncio.Task[None] | None = None + + def _new_sandbox(self, files: Dict[str, str] | None = None, use_pool: bool = True) -> AsyncSandbox: + # pool_ref claims a prewarmed pod from a server-side Pool whose template + # has already warmed the lean toolchain, so the prepare-time prefetch + # degrades to a fast cache hit. + provider_options = {"extensions": {"poolRef": self._pool_ref}} if (self._pool_ref and use_pool) else {} + return AsyncSandbox( + provider=dict(self._provider), + spec=SandboxSpec( + image=self._image, + entrypoint=["sleep", "infinity"], + ttl_s=self._create_ttl_s, + files=files or {}, + resources=self._resources, + metadata={"purpose": "math-formal-lean-verify"}, + provider_options=provider_options, + ), + ) + + async def _start_sandbox(self, files: Dict[str, str] | None = None) -> AsyncSandbox: + """Start a sandbox, degrading a failed pool claim to a direct create.""" + sandbox = self._new_sandbox(files) + try: + await sandbox.start() + return sandbox + except Exception as exc: + if not self._pool_ref: + raise + LOG.warning("lean pool '%s' claim failed (%s); falling back to a direct create", self._pool_ref, exc) + sandbox = self._new_sandbox(files, use_pool=False) + await sandbox.start() + return sandbox + + async def _stop_sandbox(self, sandbox: AsyncSandbox) -> None: + self._pool_sandboxes.discard(sandbox) + try: + await sandbox.stop() + except Exception as exc: + LOG.warning("lean sandbox teardown failed (TTL will reap): %s", exc) + + async def _create_pool_pod(self) -> AsyncSandbox: + """Create + warm one pool pod: a single bulk tar read pulls the olean tree at + line rate (concurrent chunk fetches) instead of the compile's serial faults.""" + sandbox = await self._start_sandbox() + try: + # NOT `tar cf /dev/null`: GNU tar detects the null sink and skips reading + # file contents, silently defeating the prefetch. + await sandbox.exec( + f"find {self._prefetch_paths} -type f -exec cat {{}} + > /dev/null 2>&1; true", timeout_s=1800 + ) + except asyncio.CancelledError: + await self._stop_sandbox(sandbox) + raise + except Exception: + pass # prefetch is an optimization; the first compile warms the rest + return sandbox + + def start_pool(self) -> None: + """Kick the pool fill early (call from server lifespan startup) so the first + verify does not pay pool warmup inside its admission window.""" + if self._pool_size <= 0 or self._pool is not None or self._closed: + return + self._pool = asyncio.Queue(maxsize=self._pool_size) + for _ in range(self._pool_size): + self._schedule_fill() + + def _schedule_fill(self) -> None: + if self._closed: + return + task = asyncio.create_task(self._fill_one()) + self._fill_tasks.add(task) + task.add_done_callback(self._fill_tasks.discard) + + async def _fill_one(self) -> None: + retry_s = 1.0 + while not self._closed: + try: + sandbox = await self._create_pool_pod() + except Exception as exc: + LOG.error("lean pool pod create failed; retrying in %.0fs: %s", retry_s, exc) + await asyncio.sleep(retry_s) + retry_s = min(retry_s * 2, 30.0) + continue + + if self._closed: + await self._stop_sandbox(sandbox) + else: + self._pool_sandboxes.add(sandbox) + self._pool.put_nowait(sandbox) + return + + async def _execute_pooled(self, code: str, timeout: float) -> Dict[str, Any]: + self.start_pool() + if self._pool is None: + raise RuntimeError("Lean sandbox pool is closed") + pool = self._pool + for attempt in (1, 2): + try: + sandbox = await asyncio.wait_for(pool.get(), timeout=self._acquire_timeout_s) + except asyncio.TimeoutError: + LOG.warning("Lean pool admission timed out after %.0fs", self._acquire_timeout_s) + return {"process_status": "timeout", "stdout": "", "stderr": "Client timed out"} + proof_name = f"proof_{uuid.uuid4().hex}.lean" + try: + with tempfile.NamedTemporaryFile("w", suffix=".lean", delete=False) as fh: + fh.write(code) + try: + await sandbox.upload(fh.name, f"{self._project_dir}/{proof_name}") + finally: + os.unlink(fh.name) + # rc must survive the cleanup rm: 124/137 keep mapping to timeout. + command = ( + f"cd {self._project_dir} && timeout -s KILL {timeout} " + f"lake env --dir {self._project_dir} lean {proof_name}; " + f"rc=$?; rm -f {proof_name}; exit $rc" + ) + result = await sandbox.exec(command, timeout_s=timeout + 60) + except asyncio.CancelledError: + await self._stop_sandbox(sandbox) + self._schedule_fill() + raise + except Exception as exc: + # Pod is suspect (TTL expiry, node loss): replace it, retry once elsewhere. + LOG.warning("lean pool pod failed (attempt %d), replacing: %s", attempt, exc) + await self._stop_sandbox(sandbox) + self._schedule_fill() + if attempt == 1: + continue + return {"process_status": "error", "stdout": "", "stderr": str(exc)} + if self._closed: + await self._stop_sandbox(sandbox) + else: + pool.put_nowait(sandbox) + return self._map_result(result, timeout) + return {"process_status": "error", "stdout": "", "stderr": "lean pool exhausted"} + + def _map_result(self, result, timeout: float) -> Dict[str, Any]: + stdout = result.stdout or "" + stderr = result.stderr or "" + if result.return_code == 0: + process_status = "completed" + elif result.return_code in (124, 137, -9): + process_status = "timeout" + stderr += f"Execution timed out after {timeout} seconds\n" + else: + process_status = "failed" + if len(stdout) > self.max_output_characters: + stdout = stdout[: self.max_output_characters] + "" + if len(stderr) > self.max_output_characters: + stderr = stderr[: self.max_output_characters] + "" + return {"process_status": process_status, "stdout": stdout, "stderr": stderr} + + async def execute_lean4(self, code: str, timeout: float = 30.0) -> Dict[str, Any]: + """Same signature and return contract as Lean4SandboxClient.execute_lean4.""" + if self._pool_size > 0: + return await self._execute_pooled(code, timeout) + + try: + await asyncio.wait_for(self._semaphore.acquire(), timeout=self._acquire_timeout_s) + except asyncio.TimeoutError: + LOG.warning("Lean sandbox admission timed out after %.0fs", self._acquire_timeout_s) + return {"process_status": "timeout", "stdout": "", "stderr": "Client timed out"} + + proof_name = f"proof_{uuid.uuid4().hex}.lean" + sandbox = None + try: + sandbox = await self._start_sandbox(files={f"{self._project_dir}/{proof_name}": code}) + # In-sandbox `timeout -s KILL` must always fire before the provider deadline so the + # partial-stdout + "Execution timed out..." contract is preserved (never a raw exec kill). + command = ( + f"cd {self._project_dir} && timeout -s KILL {timeout} " + f"lake env --dir {self._project_dir} lean {proof_name}" + ) + result = await sandbox.exec(command, timeout_s=timeout + 60) + return self._map_result(result, timeout) + except Exception as e: # infra failure -> degrade, never raise into verify() + LOG.error("OpenSandbox lean4 execution failed: %s", e) + return {"process_status": "error", "stdout": "", "stderr": str(e)} + finally: + self._semaphore.release() + if sandbox is not None: + await self._stop_sandbox(sandbox) + + async def close(self) -> None: + """Stop pool maintenance and every warm sandbox owned by this client.""" + if self._close_task is None: + self._closed = True + + async def cleanup() -> None: + tasks = tuple(self._fill_tasks) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + sandboxes = tuple(self._pool_sandboxes) + if sandboxes: + await asyncio.gather(*(self._stop_sandbox(sandbox) for sandbox in sandboxes)) + self._pool = None + + self._close_task = asyncio.create_task(cleanup()) + await await_cleanup(self._close_task) diff --git a/resources_servers/math_formal_lean/tests/test_app.py b/resources_servers/math_formal_lean/tests/test_app.py index ad9856bea5..c1e76fb48d 100644 --- a/resources_servers/math_formal_lean/tests/test_app.py +++ b/resources_servers/math_formal_lean/tests/test_app.py @@ -16,6 +16,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi.testclient import TestClient from nemo_gym.openai_utils import ( NeMoGymResponse, @@ -53,6 +54,18 @@ def config(self) -> MathFormalLeanResourcesServerConfig: def server(self, config) -> MathFormalLeanResourcesServer: return MathFormalLeanResourcesServer(config=config, server_client=MagicMock(spec=ServerClient)) + def test_unknown_sandbox_backend_is_rejected(self, config): + values = config.model_dump() + values["sandbox_backend"] = "gym_sandox" + with pytest.raises(ValueError, match="sandbox_backend"): + MathFormalLeanResourcesServerConfig.model_validate(values) + + def test_lifespan_closes_sandbox_client(self, server): + server._sandbox_client.close = AsyncMock() + with TestClient(server.setup_webserver()): + pass + server._sandbox_client.close.assert_awaited_once() + def _create_response(self, text: str, msg_id: str = "test_msg") -> NeMoGymResponse: return NeMoGymResponse( id="test_response_id", diff --git a/resources_servers/math_formal_lean/tests/test_sandbox_backends.py b/resources_servers/math_formal_lean/tests/test_sandbox_backends.py new file mode 100644 index 0000000000..1a38fe941c --- /dev/null +++ b/resources_servers/math_formal_lean/tests/test_sandbox_backends.py @@ -0,0 +1,335 @@ +# 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. +"""Backend tests for the Lean sandbox clients.""" + +import asyncio +from types import SimpleNamespace + +import pytest + +from resources_servers.math_formal_lean.sandbox_client import ( + GymSandboxLean4Client, + Lean4SandboxClient, +) + + +PROVIDER = { + "opensandbox": { + "connection": {"domain": "http://sandbox.example", "api_key": "k", "use_server_proxy": True}, + } +} + + +class TestHttpClientDefaults: + def test_default_url_is_unchanged(self): + client = Lean4SandboxClient() + assert client._get_execute_url() == "http://127.0.0.1:6000/execute" + + +class _FakeSandbox: + """Stands in for AsyncSandbox; records lifecycle and returns a scripted exec result.""" + + instances = [] + next_exec_result = None + next_raise_on_exec = None + + def __init__(self, provider=None, spec=None): + self.spec = spec + self.exec_result = _FakeSandbox.next_exec_result or SimpleNamespace(return_code=0, stdout="ok", stderr="") + self.raise_on_exec = _FakeSandbox.next_raise_on_exec + self.exec_command = None + self.exec_timeout_s = None + self.stopped = False + _FakeSandbox.instances.append(self) + + async def start(self): + return self + + async def exec(self, command, **kwargs): + self.exec_command = command + self.exec_commands = getattr(self, "exec_commands", []) + [command] + self.exec_timeout_s = kwargs.get("timeout_s") + if self.raise_on_exec: + raise self.raise_on_exec + return self.exec_result + + async def upload(self, local_path, remote_path): + self.uploaded = getattr(self, "uploaded", []) + [remote_path] + + async def stop(self): + self.stopped = True + + +@pytest.fixture() +def fake_sandbox(monkeypatch): + import resources_servers.math_formal_lean.sandbox_client as sandbox_client + + _FakeSandbox.instances = [] + _FakeSandbox.next_exec_result = None + _FakeSandbox.next_raise_on_exec = None + monkeypatch.setattr(sandbox_client, "AsyncSandbox", _FakeSandbox) + return _FakeSandbox + + +def _client(**overrides) -> GymSandboxLean4Client: + kwargs = dict(provider=PROVIDER, image="lean-img") + kwargs.update(overrides) + return GymSandboxLean4Client(**kwargs) + + +class TestGymSandboxLean4Client: + def test_empty_image_is_a_hard_error(self): + with pytest.raises(ValueError, match="image"): + GymSandboxLean4Client(provider=PROVIDER, image="") + + def test_empty_creds_is_a_hard_error(self): + bad = {"opensandbox": {"connection": {"domain": "", "api_key": ""}}} + with pytest.raises(ValueError, match="OPENSANDBOX"): + GymSandboxLean4Client(provider=bad, image="img") + + def test_completed_maps_rc_zero(self, fake_sandbox): + out = asyncio.run(_client().execute_lean4("theorem t : True := trivial", timeout=30.0)) + assert out == {"process_status": "completed", "stdout": "ok", "stderr": ""} + box = fake_sandbox.instances[0] + assert box.stopped, "per-verify pod must be destroyed in finally" + assert "timeout -s KILL 30.0 lake env --dir /lean4/my_project lean" in box.exec_command + assert box.exec_timeout_s == 90.0, "provider deadline must trail the in-sandbox timeout" + assert list(box.spec.files.values()) == ["theorem t : True := trivial"] + assert box.spec.entrypoint == ["sleep", "infinity"] + + def test_nonzero_rc_maps_to_failed(self, fake_sandbox): + fake_sandbox.next_exec_result = SimpleNamespace(return_code=1, stdout="", stderr="error: x") + out = asyncio.run(_client().execute_lean4("bad", timeout=30.0)) + assert out["process_status"] == "failed" + assert out["stderr"] == "error: x" + + def test_timeout_rc_maps_to_the_ns_timeout_contract(self, fake_sandbox): + fake_sandbox.next_exec_result = SimpleNamespace(return_code=137, stdout="partial", stderr="") + out = asyncio.run(_client().execute_lean4("slow", timeout=30.0)) + assert out["process_status"] == "timeout" + assert out["stdout"] == "partial", "partial stdout must survive, matching the NS server" + assert out["stderr"].endswith("Execution timed out after 30.0 seconds\n") + + def test_output_truncation_matches_the_ns_contract(self, fake_sandbox): + fake_sandbox.next_exec_result = SimpleNamespace(return_code=137, stdout="1234", stderr="abcd") + out = asyncio.run(_client(max_output_characters=3).execute_lean4("slow", timeout=30.0)) + assert out == { + "process_status": "timeout", + "stdout": "123", + "stderr": "abc", + } + + def test_infra_failure_degrades_to_error_and_still_tears_down(self, fake_sandbox): + fake_sandbox.next_raise_on_exec = RuntimeError("proxy exploded") + out = asyncio.run(_client().execute_lean4("x", timeout=5.0)) + assert out["process_status"] == "error" + assert "proxy exploded" in out["stderr"] + assert fake_sandbox.instances[0].stopped + + def test_admission_timeout_collapses_to_client_timed_out(self, fake_sandbox): + client = _client(max_concurrent=1, acquire_timeout_s=0.05) + + async def main(): + sem = client._semaphore + await sem.acquire() # exhaust admission + try: + return await client.execute_lean4("x", timeout=5.0) + finally: + sem.release() + + out = asyncio.run(main()) + assert out == {"process_status": "timeout", "stdout": "", "stderr": "Client timed out"} + assert not fake_sandbox.instances, "no pod may be created past a failed admission" + + def test_fresh_pod_is_stopped_when_exec_is_cancelled(self, fake_sandbox, monkeypatch): + entered = asyncio.Event() + + async def block(*args, **kwargs): + entered.set() + await asyncio.Event().wait() + + monkeypatch.setattr(fake_sandbox, "exec", block) + + async def scenario(): + task = asyncio.create_task(_client().execute_lean4("theorem t : True := trivial")) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(scenario()) + assert fake_sandbox.instances[0].stopped + + +class TestPooledMode: + def test_pool_reuses_pods_across_verifies(self, fake_sandbox): + client = _client(pool_size=2) + + async def scenario(): + return [await client.execute_lean4("theorem t : True := trivial", timeout=5.0) for _ in range(3)] + + results = asyncio.run(scenario()) + assert [r["process_status"] for r in results] == ["completed"] * 3 + # 2 pool pods serve 3 verifies — no per-verify creates. + assert len(fake_sandbox.instances) == 2 + pod = fake_sandbox.instances[0] + # First exec on a pool pod is the olean prefetch; compiles clean up their proof file. + assert pod.exec_commands[0].startswith("find ") + assert "rm -f" in pod.exec_commands[-1] and "timeout -s KILL 5.0" in pod.exec_commands[-1] + assert not pod.stopped + + def test_pool_replaces_failed_pod_and_retries(self, fake_sandbox): + client = _client(pool_size=1, acquire_timeout_s=5.0) + + async def scenario(): + # Warm the pool, then arm the NEXT pod acquisition's exec to fail once. + first = await client.execute_lean4("theorem t : True := trivial", timeout=5.0) + bad = fake_sandbox.instances[0] + bad.raise_on_exec = RuntimeError("pod lost") + second = await client.execute_lean4("theorem t : True := trivial", timeout=5.0) + return first, bad, second + + first, bad, second = asyncio.run(scenario()) + assert first["process_status"] == "completed" + # The dead pod was stopped and replaced; the retry ran on the replacement. + assert bad.stopped + assert second["process_status"] == "completed" + assert len(fake_sandbox.instances) >= 2 + + def test_pool_retries_initial_create_failure(self, fake_sandbox, monkeypatch): + client = _client(pool_size=1, acquire_timeout_s=1.0) + create = client._create_pool_pod + attempts = 0 + + async def flaky_create(): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("control plane unavailable") + return await create() + + real_sleep = asyncio.sleep + + async def fast_sleep(_): + await real_sleep(0) + + monkeypatch.setattr(client, "_create_pool_pod", flaky_create) + monkeypatch.setattr(asyncio, "sleep", fast_sleep) + + async def scenario(): + out = await client.execute_lean4("theorem t : True := trivial", timeout=5.0) + await client.close() + return out + + out = asyncio.run(scenario()) + assert out["process_status"] == "completed" + assert attempts == 2 + assert all(box.stopped for box in fake_sandbox.instances) + + @pytest.mark.parametrize("blocked_method", ["upload", "exec"]) + def test_cancelled_lease_is_stopped_and_replaced(self, fake_sandbox, blocked_method): + client = _client(pool_size=1, acquire_timeout_s=1.0) + + async def scenario(): + client.start_pool() + while client._pool is None or client._pool.qsize() == 0: + await asyncio.sleep(0) + leased = fake_sandbox.instances[0] + entered = asyncio.Event() + + async def blocked_operation(*args, **kwargs): + entered.set() + await asyncio.Event().wait() + + setattr(leased, blocked_method, blocked_operation) + task = asyncio.create_task(client.execute_lean4("theorem t : True := trivial", timeout=5.0)) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + while client._pool is None or client._pool.qsize() == 0: + await asyncio.sleep(0) + replacement = fake_sandbox.instances[-1] + await client.close() + return leased, replacement + + leased, replacement = asyncio.run(scenario()) + assert leased.stopped + assert replacement is not leased + assert replacement.stopped + + def test_cancelled_close_finishes_cleanup(self, fake_sandbox): + client = _client(pool_size=1) + + async def scenario(): + client.start_pool() + while client._pool is None or client._pool.qsize() == 0: + await asyncio.sleep(0) + sandbox = fake_sandbox.instances[0] + stop_started = asyncio.Event() + finish_stop = asyncio.Event() + original_stop = sandbox.stop + + async def blocked_stop(): + stop_started.set() + await finish_stop.wait() + await original_stop() + + sandbox.stop = blocked_stop + closing = asyncio.create_task(client.close()) + await stop_started.wait() + closing.cancel() + await asyncio.sleep(0) + closing.cancel() + await asyncio.sleep(0) + assert not closing.done() + finish_stop.set() + with pytest.raises(asyncio.CancelledError): + await closing + await client.close() + + asyncio.run(scenario()) + assert len(fake_sandbox.instances) == 1 + assert all(box.stopped for box in fake_sandbox.instances) + + +class TestLeanPoolRef: + """pool_ref rides SandboxSpec.provider_options into the provider's SDK extensions.""" + + def test_pool_ref_rides_provider_options(self, fake_sandbox): + out = asyncio.run(_client(pool_ref="math-lean-warm").execute_lean4("theorem t : True := trivial")) + assert out["process_status"] == "completed" + box = fake_sandbox.instances[0] + assert box.spec.provider_options == {"extensions": {"poolRef": "math-lean-warm"}} + + def test_no_pool_ref_means_no_extensions(self, fake_sandbox): + asyncio.run(_client().execute_lean4("theorem t : True := trivial")) + assert fake_sandbox.instances[0].spec.provider_options == {} + + def test_claim_failure_falls_back_to_direct_create(self, fake_sandbox, monkeypatch): + calls = {"n": 0} + + async def flaky_start(self): + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("pool exhausted") + return self + + monkeypatch.setattr(fake_sandbox, "start", flaky_start) + out = asyncio.run(_client(pool_ref="math-lean-warm").execute_lean4("theorem t : True := trivial")) + assert out["process_status"] == "completed" + # First attempt carried the claim; the fallback dropped it. + assert fake_sandbox.instances[0].spec.provider_options == {"extensions": {"poolRef": "math-lean-warm"}} + assert fake_sandbox.instances[1].spec.provider_options == {} diff --git a/resources_servers/ns_tools/app.py b/resources_servers/ns_tools/app.py index 70e841a168..c145a1369b 100644 --- a/resources_servers/ns_tools/app.py +++ b/resources_servers/ns_tools/app.py @@ -30,7 +30,7 @@ import time import uuid from contextlib import asynccontextmanager -from typing import Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional import httpx from fastapi import FastAPI, Request @@ -50,6 +50,10 @@ from nemo_gym.server_utils import SESSION_ID_KEY +if TYPE_CHECKING: + from sandbox_pool import SandboxPool + + logger = logging.getLogger(__name__) @@ -78,6 +82,13 @@ class NSToolsConfig(BaseResourcesServerConfig): sandbox_host: str = "127.0.0.1" sandbox_port: str = "6000" + # Sandbox backend: "local" (default — today's colocated server) or "sandbox_pool" + # (disaggregated pods on OpenSandbox; requires the sandbox_pool block below). + sandbox_type: Literal["local", "sandbox_pool"] = "local" + # SandboxPool constructor kwargs (see sandbox_pool.py). Only read when + # sandbox_type == "sandbox_pool"; the default backend never touches it. + sandbox_pool: Dict[str, Any] = Field(default_factory=dict) + # Legacy python_tool HTTP server port (only used for pre-main HTTP PythonTool variants) python_tool_port: int = 8765 @@ -132,6 +143,7 @@ class NSToolsResourcesServer(SimpleResourcesServer): _python_tool_process: Optional[subprocess.Popen] = None _timing_by_session: Dict[str, list] = {} # session_id -> list of timing records _uses_python_tool_sidecar: bool = False + _sandbox_pool: Optional["SandboxPool"] = None def setup_webserver(self) -> FastAPI: app = super().setup_webserver() @@ -145,6 +157,9 @@ def setup_webserver(self) -> FastAPI: @asynccontextmanager async def lifespan_wrapper(app): try: + if self._sandbox_pool is not None: + # Budgeted warmup launches creation tasks without gating server boot. + await self._sandbox_pool.start() async with main_app_lifespan(app) as maybe_state: yield maybe_state finally: @@ -260,14 +275,22 @@ def _initialize_nemo_skills_tools(self): logger.info(f"Initializing NeMo Skills ToolManager with tools: {self.config.nemo_skills_tools}") - context = { - "sandbox": { - "sandbox_type": "local", - "host": self.config.sandbox_host, - "port": self.config.sandbox_port, - "disable_session_restore": self.config.disable_session_restore, - } + sandbox_context: Dict[str, Any] = { + "sandbox_type": self.config.sandbox_type, + "host": self.config.sandbox_host, + "port": self.config.sandbox_port, + "disable_session_restore": self.config.disable_session_restore, } + if self.config.sandbox_type == "sandbox_pool": + from gym_sandbox import GymSandbox + from nemo_skills.code_execution.sandbox import sandboxes + from sandbox_pool import SandboxPool + + sandboxes["sandbox_pool"] = GymSandbox + self._sandbox_pool = SandboxPool(**self.config.sandbox_pool) + sandbox_context["pool"] = self._sandbox_pool + + context = {"sandbox": sandbox_context} overrides = { tool_name: dict(tool_config) for tool_name, tool_config in self.config.nemo_skills_tool_overrides.items() @@ -462,25 +485,32 @@ async def verify(self, request: Request, body: NSToolsVerifyRequest) -> NSToolsV async def shutdown(self): """Cleanup resources on server shutdown.""" - if self.tool_manager: - await self.tool_manager.shutdown() - - # Terminate the python_tool subprocess if one was started. - if self._python_tool_process: - pid: int = self._python_tool_process.pid - logger.info(f"Terminating python_tool server (PID: {pid})") - self._python_tool_process.terminate() + try: + if self.tool_manager: + await self.tool_manager.shutdown() + finally: try: - self._python_tool_process.wait(timeout=5) - except subprocess.TimeoutExpired: - logger.warning("python_tool server did not terminate gracefully, killing...") - self._python_tool_process.kill() - # Reap the child after SIGKILL so it doesn't linger as . - try: - self._python_tool_process.wait(timeout=5) - except subprocess.TimeoutExpired: - logger.error(f"python_tool server (PID: {pid}) did not exit after SIGKILL; may leak as a zombie") - self._python_tool_process = None + if self._sandbox_pool is not None: + await self._sandbox_pool.aclose() + finally: + # Terminate the python_tool subprocess if one was started. + if self._python_tool_process: + pid: int = self._python_tool_process.pid + logger.info(f"Terminating python_tool server (PID: {pid})") + self._python_tool_process.terminate() + try: + self._python_tool_process.wait(timeout=5) + except subprocess.TimeoutExpired: + logger.warning("python_tool server did not terminate gracefully, killing...") + self._python_tool_process.kill() + # Reap the child after SIGKILL so it doesn't linger as . + try: + self._python_tool_process.wait(timeout=5) + except subprocess.TimeoutExpired: + logger.error( + f"python_tool server (PID: {pid}) did not exit after SIGKILL; may leak as a zombie" + ) + self._python_tool_process = None if __name__ == "__main__": diff --git a/resources_servers/ns_tools/configs/ns_tools.yaml b/resources_servers/ns_tools/configs/ns_tools.yaml index 24bd4446be..1f578009f0 100644 --- a/resources_servers/ns_tools/configs/ns_tools.yaml +++ b/resources_servers/ns_tools/configs/ns_tools.yaml @@ -40,7 +40,37 @@ ns_tools: # Disable session replay after sandbox worker restarts (improves stability) disable_session_restore: true - + + # Sandbox backend switch: 'local' (default — today's colocated server, byte-identical + # with zero env vars set) or 'sandbox_pool' (disaggregated pods on OpenSandbox). + sandbox_type: ${oc.env:NS_TOOLS_SANDBOX_TYPE,local} + # Pool settings; read only when sandbox_type == sandbox_pool. Selecting that + # backend with an empty domain/api_key/image is a hard startup error, never a no-op. + sandbox_pool: + provider: + opensandbox: + connection: + domain: ${oc.env:OPENSANDBOX_BASE_URL,} + api_key: ${oc.env:OPENSANDBOX_API_KEY,} + use_server_proxy: true + # Cold image conversion can exceed the default request timeout. + request_timeout_s: 180 + image: ${oc.env:NS_SANDBOX_IMAGE,} + # Server-side Pool CRD name (pools.sandbox.opensandbox.io). When set, + # creates claim prewarmed pods via extensions.poolRef — image/resources come + # from the pool template and warmup drops to allocation time. + pool_ref: ${oc.env:NS_SANDBOX_POOL_REF,} + # With pool_ref set: fall back to a direct create when the pool is full or + # unavailable (true, default) or fail the slot instead (false). Fallback pods + # go through the prepare step, so keep setup/service settings configured. + pool_fallback: ${oc.decode:${oc.env:NS_SANDBOX_POOL_FALLBACK,true}} + # OpenSandbox otherwise replaces the image command with an idle execd + # bootstrap, so direct creates must start the NS service explicitly. + entrypoint: ["/start-with-nginx.sh"] + port: 6000 + size: ${oc.env:NS_SANDBOX_POOL_SIZE,8} + ttl_s: ${oc.env:NS_SANDBOX_TTL_S,14400} + domain: agent verified: false description: NeMo Skills tool execution with math verification diff --git a/resources_servers/ns_tools/gym_sandbox.py b/resources_servers/ns_tools/gym_sandbox.py new file mode 100644 index 0000000000..577df772f5 --- /dev/null +++ b/resources_servers/ns_tools/gym_sandbox.py @@ -0,0 +1,117 @@ +# 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. +"""NeMo-Skills backend that routes sessions through a shared sandbox pool. + +The resources server owns the pool lifetime and registers this backend when selected. +Subclassing ``LocalSandbox`` preserves its request and session behavior while replacing +only the HTTP transport. +""" + +import asyncio +import json +import logging +from typing import Any, Dict, Optional + +import aiohttp +import httpx +from nemo_skills.code_execution import sandbox as ns_sandbox +from sandbox_pool import SandboxPool + + +LOGGER = logging.getLogger(__name__) + + +class GymSandbox(ns_sandbox.LocalSandbox): + """LocalSandbox with the transport routed through an OpenSandbox pod pool over aiohttp.""" + + def __init__(self, pool: SandboxPool, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._pool = pool + self._aiohttp: Optional[aiohttp.ClientSession] = None + + def _session(self) -> aiohttp.ClientSession: + if self._aiohttp is None or self._aiohttp.closed: + connector = aiohttp.TCPConnector(limit=4096, limit_per_host=4096, ttl_dns_cache=300) + self._aiohttp = aiohttp.ClientSession(connector=connector) + return self._aiohttp + + @staticmethod + def _parse_output_text(text: str) -> Dict[str, Any]: + try: + return json.loads(text) + except json.JSONDecodeError: + LOGGER.error("Error during parsing output: %s", text[:500]) + return {"process_status": "error", "stdout": "", "stderr": "Unknown error"} + + async def _post_execute(self, base_url: str, headers: Dict[str, str], payload: str, timeout: float): + async with self._session().post( + f"{base_url}/execute", + data=payload, + headers=headers, + timeout=aiohttp.ClientTimeout(total=timeout + 5.0), + ) as response: + return response.status, await response.text() + + async def _send_request(self, request: Dict[str, Any], timeout: float): + session_id = request.pop("session_id", None) + base_url, pool_headers = await self._pool.route(str(session_id) if session_id is not None else None) + headers = {"Content-Type": "application/json", **pool_headers} + if session_id is not None: + headers["X-Session-ID"] = str(session_id) + payload = json.dumps(request) + + try: + status, text = await self._post_execute(base_url, headers, payload, timeout) + except (aiohttp.ClientError, asyncio.TimeoutError) as exc: + raise httpx.TimeoutException(f"sandbox pool transport error: {exc!r}") from exc + if status != 200: + # Normalize every infra failure to the shape the NS client already tolerates. + raise httpx.TimeoutException(f"sandbox pool transport returned HTTP {status}") + return self._parse_output_text(text) + + async def delete_session(self, session_id: str) -> None: + """Delete the session on the pod it is pinned to, then release the pin.""" + if str(session_id) not in self._pool._session_to_slot: + # Never pinned (or already swept): no pod holds this session's state, and + # route() would mint a fresh pin just to receive a guaranteed 404. + self._pool.release(str(session_id)) + self.session_histories.pop(str(session_id), None) + return + try: + base_url, pool_headers = await self._pool.route(str(session_id)) + except httpx.TimeoutException: + self._pool.release(str(session_id)) + self.session_histories.pop(str(session_id), None) + return + try: + async with self._session().delete( + f"{base_url}/sessions/{session_id}", + headers={**pool_headers, "X-Session-ID": str(session_id)}, + timeout=aiohttp.ClientTimeout(total=10.0), + ) as response: + if response.status not in (200, 404): + LOGGER.warning("delete_session %s returned HTTP %d", session_id, response.status) + except (aiohttp.ClientError, asyncio.TimeoutError) as exc: + LOGGER.warning("delete_session %s failed (pod TTL/idle reaper will clean up): %s", session_id, exc) + finally: + self._pool.release(str(session_id)) + self.session_histories.pop(str(session_id), None) + + async def close(self) -> None: + try: + if self._aiohttp is not None and not self._aiohttp.closed: + await self._aiohttp.close() + finally: + await super().close() diff --git a/resources_servers/ns_tools/requirements.txt b/resources_servers/ns_tools/requirements.txt index 1a920f13ff..6dde21233a 100644 --- a/resources_servers/ns_tools/requirements.txt +++ b/resources_servers/ns_tools/requirements.txt @@ -1,3 +1,3 @@ --e nemo-gym[dev] @ ../../ +-e nemo-gym[dev,sandbox] @ ../../ ## NeMo-Skills tools subpackage as of June 12, 2026 nemo-skills-tools @ git+https://github.com/NVIDIA-NeMo/Skills.git@da85a881d972e6fec847b90cf553a0bf9bf10638#subdirectory=tools diff --git a/resources_servers/ns_tools/sandbox_pool.py b/resources_servers/ns_tools/sandbox_pool.py new file mode 100644 index 0000000000..797dd805e1 --- /dev/null +++ b/resources_servers/ns_tools/sandbox_pool.py @@ -0,0 +1,444 @@ +# 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. +"""A fixed set of long-lived, SHARED OpenSandbox pods serving the NeMo-Skills sandbox +HTTP protocol, with sessions multiplexed across them by sticky routing. + +Sharing is what makes large batches feasible: many concurrent sessions ride K pods +(each pod's NS server multiplexes many sessions), instead of one pod per session. +Slots use the provider-neutral :mod:`nemo_gym.sandbox` lifecycle throughout. + +This module is imported only when ns_tools selects the ``sandbox_pool`` backend; +the default ``local`` backend never touches it. +""" + +import asyncio +import logging +import time +from dataclasses import dataclass, field +from typing import Any + +import aiohttp +import httpx # exception types only: the nemo_skills client contract catches httpx errors + +from nemo_gym.sandbox import AsyncSandbox, SandboxSpec, await_cleanup + + +LOGGER = logging.getLogger(__name__) + + +@dataclass +class _Slot: + index: int + sandbox: AsyncSandbox | None = None + base_url: str = "" + headers: dict[str, str] = field(default_factory=dict) + healthy: bool = False + strikes: int = 0 + creating: bool = False + heal_failures: int = 0 + sessions: set[str] = field(default_factory=set) + + +class SandboxPool: + """K shared NS-sandbox pods with sticky session routing. + + The constructor is pure (validation only); ``start()`` kicks a non-blocking + warmup and the health/idle-sweep loops. ``route()`` lazily starts everything + as a safety net if the owner never called ``start()``. + """ + + def __init__( + self, + *, + provider: dict[str, Any], + image: str, + pool_ref: str = "", + pool_fallback: bool = True, + port: int = 6000, + size: int = 8, + ttl_s: float | None = None, + env: dict[str, Any] | None = None, + entrypoint: list[str] | None = None, + resources: dict[str, Any] | None = None, + resource_requests: dict[str, Any] | None = None, + setup_files: dict[str, str] | None = None, + setup_commands: list[str] | None = None, + service_command: str | None = None, + health_path: str = "/health", + ready_timeout_s: float = 30.0, + health_budget_s: float = 300.0, + warmup_fill_concurrency: int = 0, # 0 = full fan-out (all slots at once) + health_interval_s: float = 15.0, + health_timeout_s: float = 10.0, + heal_creates_per_s: float = 4.0, + heal_concurrency: int = 16, + session_idle_sweep_s: float = 7200.0, + ) -> None: + if not isinstance(provider, dict) or set(provider) != {"opensandbox"}: + raise ValueError("sandbox_pool.provider must contain exactly one 'opensandbox' provider") + provider_config = provider["opensandbox"] or {} + if not isinstance(provider_config, dict): + raise TypeError("sandbox_pool.provider.opensandbox must be a mapping") + connection = provider_config.get("connection") or {} + if not connection.get("domain") or not connection.get("api_key"): + raise ValueError( + "sandbox_pool backend selected but the provider connection has an empty " + "domain or api_key — set OPENSANDBOX_BASE_URL / OPENSANDBOX_API_KEY" + ) + self._provider = provider + self._pool_ref = str(pool_ref or "") + # bool("false") is True; env-fed values arrive as strings. + self._pool_fallback = ( + pool_fallback if isinstance(pool_fallback, bool) else str(pool_fallback).lower() in ("true", "1", "yes") + ) + if not image: + raise ValueError("sandbox_pool backend selected but image is empty — set NS_SANDBOX_IMAGE") + if int(size) < 1: + raise ValueError(f"sandbox_pool.size must be >= 1, got {size}") + self._image = image + self._port = int(port) + self._size = int(size) + self._ttl_s = float(ttl_s) if ttl_s else None + self._env = dict(env or {}) + self._entrypoint = list(entrypoint) if entrypoint else None + self._resources = dict(resources or {}) + self._resource_requests = dict(resource_requests or {}) + self._setup_files = dict(setup_files or {}) + self._setup_commands = list(setup_commands or []) + self._service_command = service_command + if (not self._pool_ref or self._pool_fallback) and not (self._entrypoint or self._service_command): + raise ValueError( + "sandbox_pool direct creation requires entrypoint or service_command to start the NS server" + ) + self._health_path = health_path + # First pull of a large image on a fresh node can take minutes; the SDK default + # (30s) fails creates that would have succeeded. + self._ready_timeout_s = float(ready_timeout_s) + self._health_budget_s = float(health_budget_s) + self._warmup_fill_concurrency = int(warmup_fill_concurrency) or self._size + self._health_interval_s = health_interval_s + self._health_timeout_s = health_timeout_s + self._heal_min_interval_s = 1.0 / heal_creates_per_s if heal_creates_per_s > 0 else 0.0 + self._heal_concurrency = int(heal_concurrency) + self._heal_rate_lock = asyncio.Lock() + self._session_idle_sweep_s = session_idle_sweep_s + self._metadata = {"purpose": "ns-tools-sandbox-pool"} + + self._slots = [_Slot(index=i) for i in range(self._size)] + self._session_to_slot: dict[str, int] = {} + self._session_last_used: dict[str, float] = {} + self._lock = asyncio.Lock() + self._started = False + self._warmup_done = False + self._closed = False + self._close_task: asyncio.Task[None] | None = None + self._tasks: list[asyncio.Task[None]] = [] + self._heal_tasks: set[asyncio.Task[None]] = set() + self._next_heal_slot = 0 + self._last_heal_create = 0.0 + self._http: aiohttp.ClientSession | None = None + + # ------------------------------------------------------------------ lifecycle + + def _http_session(self) -> aiohttp.ClientSession: + if self._http is None or self._http.closed: + self._http = aiohttp.ClientSession( + connector=aiohttp.TCPConnector(limit=512, ttl_dns_cache=300), + timeout=aiohttp.ClientTimeout(total=self._health_timeout_s), + ) + return self._http + + async def start(self) -> None: + """Kick warmup and the maintenance loops; returns immediately.""" + if self._started or self._closed: + return + self._started = True + self._tasks.append(asyncio.create_task(self._warmup(), name="osb-pool-warmup")) + self._tasks.append(asyncio.create_task(self._heal_loop(), name="osb-pool-heal")) + self._tasks.append(asyncio.create_task(self._sweep_loop(), name="osb-pool-sweep")) + + async def aclose(self) -> None: + if self._close_task is None: + self._closed = True + + async def cleanup() -> None: + tasks = [*self._tasks, *self._heal_tasks] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + self._tasks.clear() + self._heal_tasks.clear() + + occupied = [slot for slot in self._slots if slot.sandbox is not None] + await asyncio.gather(*(self._stop_sandbox(slot.sandbox, slot.index) for slot in occupied)) + for slot in occupied: + slot.sandbox = None + slot.healthy = False + if self._http is not None and not self._http.closed: + await self._http.close() + + self._close_task = asyncio.create_task(cleanup()) + await await_cleanup(self._close_task) + + async def _stop_sandbox(self, sandbox: AsyncSandbox, slot_index: int) -> None: + try: + await sandbox.stop() + except Exception as exc: + LOGGER.warning("pool slot %d teardown failed (TTL will reap): %s", slot_index, exc) + + # ------------------------------------------------------------------ slot fill + + async def _acquire_sandbox(self) -> tuple[AsyncSandbox, bool]: + """Returns (sandbox, from_pool). Pool mode claims a prewarmed pod from the + server-side Pool CRD; when the pool is full or unavailable and pool_fallback + is on, it degrades to a direct create (which then needs the prepare step, so + pool configs should still carry setup/service settings for parity).""" + if self._pool_ref: + sandbox = AsyncSandbox(self._provider) + try: + await sandbox.start( + SandboxSpec( + image=self._image, + metadata=dict(self._metadata), + ttl_s=self._ttl_s or 14400.0, + ready_timeout_s=self._ready_timeout_s, + provider_options={"extensions": {"poolRef": self._pool_ref}}, + ports=(self._port,), + ) + ) + return sandbox, True + except Exception as exc: + if not self._pool_fallback: + raise + LOGGER.warning("pool %r allocation failed (%s); falling back to a direct create", self._pool_ref, exc) + sandbox = AsyncSandbox(self._provider) + await sandbox.start( + SandboxSpec( + image=self._image, + entrypoint=self._entrypoint, + env=dict(self._env), + metadata=dict(self._metadata), + resources=dict(self._resources), + provider_options={"resource_requests": dict(self._resource_requests)} + if self._resource_requests + else {}, + ttl_s=self._ttl_s or 14400.0, + ready_timeout_s=self._ready_timeout_s, + ports=(self._port,), + ) + ) + return sandbox, False + + async def _create_slot_inner(self, slot: _Slot) -> None: + sandbox, from_pool = await self._acquire_sandbox() + try: + if not from_pool: + for target_path, local_path in self._setup_files.items(): + await sandbox.upload(local_path, target_path) + for command in self._setup_commands: + execution = await sandbox.exec(command) + if execution.return_code != 0: + raise RuntimeError(f"setup command failed rc={execution.return_code}: {command!r}") + if self._service_command: + # This must be the last exec: a later command reaps the background service. + execution = await sandbox.exec(self._service_command) + if execution.return_code != 0: + raise RuntimeError(f"service command failed rc={execution.return_code}") + resolved = await sandbox.endpoint(self._port) + base_url, headers = resolved.endpoint.rstrip("/"), dict(resolved.headers) + await self._wait_healthy(base_url, headers, budget_s=self._health_budget_s) + except BaseException: + await self._stop_sandbox(sandbox, slot.index) + raise + slot.sandbox = sandbox + slot.base_url = base_url + slot.headers = headers + slot.strikes = 0 + slot.healthy = True + + async def _wait_healthy(self, base_url: str, headers: dict[str, str], budget_s: float) -> None: + """Gate admission on the ACTUAL traffic path: proxied GET /health must return 200.""" + deadline = time.monotonic() + budget_s + last_error: str | None = None + while time.monotonic() < deadline: + try: + async with self._http_session().get(f"{base_url}{self._health_path}", headers=headers) as response: + if response.status == 200: + return + last_error = f"HTTP {response.status}" + except (aiohttp.ClientError, asyncio.TimeoutError) as exc: + last_error = repr(exc) + await asyncio.sleep(1.0) + raise RuntimeError(f"pod never became healthy through the proxy: {last_error}") + + async def _warmup(self) -> None: + semaphore = asyncio.Semaphore(self._warmup_fill_concurrency) + + async def one(slot: _Slot) -> None: + async with semaphore: + slot.creating = True + try: + await self._create_slot_inner(slot) + except Exception as exc: + LOGGER.warning("pool warmup: slot %d failed (heal loop will retry): %s", slot.index, exc) + finally: + slot.creating = False + + await asyncio.gather(*(one(slot) for slot in self._slots)) + ready = sum(1 for slot in self._slots if slot.healthy) + self._warmup_done = True + LOGGER.info("pool ready %d/%d", ready, self._size) + + # ------------------------------------------------------------------ maintenance + + async def _heal_loop(self) -> None: + while not self._closed: + await asyncio.sleep(self._health_interval_s) + + async def check(slot: _Slot) -> bool: + """Returns True when the slot needs a heal. Health probes run concurrently.""" + if slot.creating: + return False + if slot.sandbox is None or not slot.healthy: + return self._warmup_done + try: + async with self._http_session().get( + f"{slot.base_url}{self._health_path}", headers=slot.headers + ) as response: + ok = response.status == 200 + except (aiohttp.ClientError, asyncio.TimeoutError): + ok = False + if ok: + slot.strikes = 0 + return False + slot.strikes += 1 + if slot.strikes >= 3: + LOGGER.warning("pool slot %d failed 3 health checks — evicting and healing in slot", slot.index) + slot.healthy = False + await self._drop_slot_sessions(slot) + if slot.sandbox is not None: + await self._stop_sandbox(slot.sandbox, slot.index) + slot.sandbox = None + return True + return False + + needs_heal = await asyncio.gather(*(check(slot) for slot in self._slots)) + to_heal = [slot for slot, needed in zip(self._slots, needs_heal) if needed] + if not to_heal: + continue + capacity = max(0, self._heal_concurrency - len(self._heal_tasks)) + ordered = sorted(to_heal, key=lambda slot: (slot.index - self._next_heal_slot) % self._size) + selected = ordered[:capacity] + for slot in selected: + task = asyncio.create_task(self._heal_slot(slot), name=f"osb-pool-heal-{slot.index}") + self._heal_tasks.add(task) + task.add_done_callback(self._heal_tasks.discard) + if selected: + self._next_heal_slot = (selected[-1].index + 1) % self._size + + async def _heal_slot(self, slot: _Slot) -> None: + """Replace a dead pod in the SAME slot. Heals run concurrently (bounded by + heal_concurrency); the rate lock spaces create STARTS so a mass heal cannot + storm the sandbox service's create path.""" + if slot.creating: + return + slot.creating = True + try: + async with self._heal_rate_lock: + now = time.monotonic() + start_at = max(now, self._last_heal_create + self._heal_min_interval_s) + self._last_heal_create = start_at + wait = start_at - time.monotonic() + if wait > 0: + await asyncio.sleep(wait) + await self._create_slot_inner(slot) + if slot.heal_failures: + LOGGER.info("pool slot %d healed after %d failed attempts", slot.index, slot.heal_failures) + else: + LOGGER.info("pool slot %d healed", slot.index) + slot.heal_failures = 0 + except Exception as exc: + slot.heal_failures += 1 + # A fully unreachable sandbox service spams 2 lines/slot/interval otherwise; warn on the first + # failure and every 10th, whisper the rest. + log = LOGGER.warning if slot.heal_failures == 1 or slot.heal_failures % 10 == 0 else LOGGER.debug + log( + "pool slot %d heal attempt %d failed (will retry next interval): %s", + slot.index, + slot.heal_failures, + exc, + ) + finally: + slot.creating = False + + async def _drop_slot_sessions(self, slot: _Slot) -> None: + async with self._lock: + for session_id in list(slot.sessions): + self._session_to_slot.pop(session_id, None) + self._session_last_used.pop(session_id, None) + slot.sessions.clear() + + async def _sweep_loop(self) -> None: + while not self._closed: + await asyncio.sleep(min(self._session_idle_sweep_s, 600.0)) + cutoff = time.monotonic() - self._session_idle_sweep_s + async with self._lock: + stale = [s for s, t in self._session_last_used.items() if t < cutoff] + for session_id in stale: + index = self._session_to_slot.pop(session_id, None) + self._session_last_used.pop(session_id, None) + if index is not None: + self._slots[index].sessions.discard(session_id) + if stale: + LOGGER.info("pool idle sweep dropped %d stale session pins", len(stale)) + + # ------------------------------------------------------------------ routing + + async def route(self, session_id: str | None) -> tuple[str, dict[str, str]]: + """Resolve (base_url, headers) for a session; pins new sessions to the least-loaded pod. + + Raises httpx.TimeoutException when no pod is healthy, which the NS client already + collapses into its timeout contract — total sandbox-service loss degrades rewards, never the server. + """ + if not self._started: + await self.start() + async with self._lock: + if session_id is not None: + index = self._session_to_slot.get(session_id) + if index is not None and self._slots[index].healthy: + self._session_last_used[session_id] = time.monotonic() + return self._slots[index].base_url, self._slots[index].headers + healthy = [slot for slot in self._slots if slot.healthy] + if not healthy: + raise httpx.TimeoutException("no healthy sandbox pods in the pool") + slot = min(healthy, key=lambda s: len(s.sessions)) + if session_id is not None: + previous = self._session_to_slot.get(session_id) + if previous is not None: + self._slots[previous].sessions.discard(session_id) + self._session_to_slot[session_id] = slot.index + self._session_last_used[session_id] = time.monotonic() + slot.sessions.add(session_id) + return slot.base_url, slot.headers + + def release(self, session_id: str) -> None: + index = self._session_to_slot.pop(session_id, None) + self._session_last_used.pop(session_id, None) + if index is not None: + self._slots[index].sessions.discard(session_id) + + @property + def ready_count(self) -> int: + return sum(1 for slot in self._slots if slot.healthy) diff --git a/resources_servers/ns_tools/tests/test_app_sandbox_pool_wiring.py b/resources_servers/ns_tools/tests/test_app_sandbox_pool_wiring.py new file mode 100644 index 0000000000..5d8026a97c --- /dev/null +++ b/resources_servers/ns_tools/tests/test_app_sandbox_pool_wiring.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from app import NSToolsConfig, NSToolsResourcesServer +from nemo_skills.code_execution import sandbox as ns_sandbox + +from nemo_gym.server_utils import ServerClient + + +class _FakePool: + def __init__(self, **_config): + self.start_calls = 0 + self.close_calls = 0 + + async def start(self): + self.start_calls += 1 + + async def aclose(self): + self.close_calls += 1 + + +class _FakeToolManager: + def __init__(self, module_specs, overrides, context): + sandbox_config = dict(context["sandbox"]) + sandbox_type = sandbox_config.pop("sandbox_type") + self.sandbox = ns_sandbox.get_sandbox(sandbox_type=sandbox_type, **sandbox_config) + + async def list_all_tools(self): + return [] + + async def shutdown(self): + await self.sandbox.close() + + +def _server() -> NSToolsResourcesServer: + config = NSToolsConfig( + host="0.0.0.0", + port=8080, + entrypoint="", + name="ns_tools", + nemo_skills_tools=["fake::DirectPythonTool"], + sandbox_type="sandbox_pool", + sandbox_pool={}, + ) + return NSToolsResourcesServer(config=config, server_client=MagicMock(spec=ServerClient)) + + +def test_unknown_sandbox_type_is_rejected(): + with pytest.raises(ValueError, match="sandbox_type"): + NSToolsConfig( + host="0.0.0.0", + port=8080, + entrypoint="", + name="ns_tools", + sandbox_type="sandbox_pol", + ) + + +def test_each_server_starts_and_closes_only_its_pool(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + with ( + patch("app.ToolManager", _FakeToolManager), + patch("sandbox_pool.SandboxPool", _FakePool), + patch.object(NSToolsResourcesServer, "_tool_uses_python_tool_sidecar", return_value=False), + patch.dict(ns_sandbox.sandboxes, {}, clear=True), + ): + first = _server() + second = _server() + first_app = first.setup_webserver() + second_app = second.setup_webserver() + from gym_sandbox import GymSandbox + + assert ns_sandbox.sandboxes["sandbox_pool"] is GymSandbox + finally: + loop.close() + asyncio.set_event_loop(None) + + first_pool = first._sandbox_pool + second_pool = second._sandbox_pool + assert first_pool is not second_pool + assert first.tool_manager.sandbox._pool is first_pool + assert second.tool_manager.sandbox._pool is second_pool + + async def run_lifespans(): + async with first_app.router.lifespan_context(first_app): + async with second_app.router.lifespan_context(second_app): + assert first_pool.start_calls == 1 + assert second_pool.start_calls == 1 + assert first_pool.close_calls == 0 + assert second_pool.close_calls == 0 + + assert first_pool.close_calls == 0 + assert second_pool.close_calls == 1 + + asyncio.run(run_lifespans()) + assert first_pool.close_calls == 1 + assert second_pool.close_calls == 1 + + +@pytest.mark.parametrize("failure", [RuntimeError("failed"), asyncio.CancelledError()]) +def test_shutdown_closes_pool_before_propagating_tool_manager_failure(failure): + server = _server() + server.tool_manager = MagicMock() + server.tool_manager.shutdown = AsyncMock(side_effect=failure) + server._sandbox_pool = _FakePool() + + with pytest.raises(type(failure)): + asyncio.run(server.shutdown()) + + assert server._sandbox_pool.close_calls == 1 diff --git a/resources_servers/ns_tools/tests/test_sandbox_pool.py b/resources_servers/ns_tools/tests/test_sandbox_pool.py new file mode 100644 index 0000000000..abab648b5e --- /dev/null +++ b/resources_servers/ns_tools/tests/test_sandbox_pool.py @@ -0,0 +1,445 @@ +# 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. +"""Unit tests for the sandbox_pool sandbox backend. No network, no live sandbox service: +routing/eviction logic is driven directly, the transport via a fake aiohttp session.""" + +import asyncio +import sys +from pathlib import Path +from types import SimpleNamespace + +import httpx +import pytest + +from nemo_gym.sandbox import SandboxEndpoint + + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from sandbox_pool import SandboxPool # noqa: E402 + + +PROVIDER = { + "opensandbox": { + "connection": {"domain": "http://sandbox.example", "api_key": "k", "use_server_proxy": True}, + "create": {"timeout_s": 90, "retries": 3}, + } +} + + +def _pool(**overrides) -> SandboxPool: + kwargs = dict(provider=PROVIDER, image="img", size=2, entrypoint=["/start-with-nginx.sh"]) + kwargs.update(overrides) + return SandboxPool(**kwargs) + + +def _admit(pool: SandboxPool, index: int) -> None: + slot = pool._slots[index] + slot.base_url = f"http://sandbox.example/v1/sandboxes/sbx-{index}/proxy/6000" + slot.headers = {"OPEN-SANDBOX-API-KEY": "k"} + slot.healthy = True + pool._started = True # skip lazy start in route() + + +class TestPoolConfigValidation: + def test_empty_domain_is_a_hard_error(self): + bad = {"opensandbox": {"connection": {"domain": "", "api_key": "k"}}} + with pytest.raises(ValueError, match="OPENSANDBOX_BASE_URL"): + SandboxPool(provider=bad, image="img", entrypoint=["start"]) + + def test_empty_api_key_is_a_hard_error(self): + bad = {"opensandbox": {"connection": {"domain": "http://sandbox.example", "api_key": ""}}} + with pytest.raises(ValueError, match="OPENSANDBOX_API_KEY"): + SandboxPool(provider=bad, image="img", entrypoint=["start"]) + + def test_empty_image_is_a_hard_error(self): + with pytest.raises(ValueError, match="NS_SANDBOX_IMAGE"): + SandboxPool(provider=PROVIDER, image="", entrypoint=["start"]) + + def test_direct_create_without_service_start_is_a_hard_error(self): + with pytest.raises(ValueError, match="requires entrypoint or service_command"): + SandboxPool(provider=PROVIDER, image="img") + + def test_ctor_is_pure_no_event_loop_required(self): + # Constructing outside any running loop must work (pure ctor rule). + pool = _pool() + assert pool.ready_count == 0 + + +class TestRouting: + def test_sessions_stick_to_their_assigned_slot(self): + pool = _pool() + _admit(pool, 0) + _admit(pool, 1) + + async def main(): + first = await pool.route("sess-a") + for _ in range(5): + again = await pool.route("sess-a") + assert again == first + + asyncio.run(main()) + + def test_new_sessions_go_to_the_least_loaded_slot(self): + pool = _pool() + _admit(pool, 0) + _admit(pool, 1) + + async def main(): + urls = {(await pool.route(f"sess-{i}"))[0] for i in range(4)} + per_slot = [len(s.sessions) for s in pool._slots] + assert per_slot == [2, 2], f"expected even spread, got {per_slot}" + assert len(urls) == 2 + + asyncio.run(main()) + + def test_total_outage_raises_the_timeout_contract(self): + pool = _pool() + pool._started = True # no healthy slots admitted + + async def main(): + with pytest.raises(httpx.TimeoutException): + await pool.route("sess-a") + + asyncio.run(main()) + + def test_dead_slot_reroutes_the_session_and_drops_the_old_pin(self): + pool = _pool() + _admit(pool, 0) + _admit(pool, 1) + + async def main(): + await pool.route("sess-a") + index = pool._session_to_slot["sess-a"] + pool._slots[index].healthy = False + url, _ = await pool.route("sess-a") + new_index = pool._session_to_slot["sess-a"] + assert new_index != index + assert "sess-a" not in pool._slots[index].sessions + assert url == pool._slots[new_index].base_url + + asyncio.run(main()) + + def test_release_unpins(self): + pool = _pool() + _admit(pool, 0) + _admit(pool, 1) + + async def main(): + await pool.route("sess-a") + + asyncio.run(main()) + pool.release("sess-a") + assert "sess-a" not in pool._session_to_slot + assert all("sess-a" not in s.sessions for s in pool._slots) + + +class _FakeAiohttpResponse: + def __init__(self, status: int, text: str = ""): + self.status = status + self._text = text + + async def text(self): + return self._text + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + +class _FakeAiohttpSession: + """Stands in for aiohttp.ClientSession; scripts responses per call.""" + + def __init__(self, responses): + self.responses = list(responses) + self.calls = [] + self.closed = False + + def post(self, url, data=None, headers=None, timeout=None): + self.calls.append(("POST", url, dict(headers or {}))) + return self.responses.pop(0) + + def delete(self, url, headers=None, timeout=None): + self.calls.append(("DELETE", url, dict(headers or {}))) + return self.responses.pop(0) + + def get(self, url, headers=None, timeout=None): + self.calls.append(("GET", url, dict(headers or {}))) + return self.responses.pop(0) + + async def close(self): + self.closed = True + + +class TestSandboxBackend: + """The nemo_skills subclass; skipped when nemo_skills is not installed (per-server dep).""" + + @pytest.fixture() + def backend(self): + pytest.importorskip("nemo_skills") + import gym_sandbox + + sandbox = gym_sandbox.GymSandbox( + pool=_pool(size=1), + host="127.0.0.1", + port="6000", + disable_session_restore=True, + ) + _admit(sandbox._pool, 0) + return sandbox + + def test_send_request_routes_with_pool_headers_and_session(self, backend): + ok = '{"process_status": "completed", "stdout": "", "stderr": ""}' + backend._aiohttp = _FakeAiohttpSession([_FakeAiohttpResponse(200, ok)]) + result = asyncio.run(backend._send_request({"generated_code": "1+1", "session_id": "sess-a"}, timeout=10.0)) + assert result["process_status"] == "completed" + method, url, headers = backend._aiohttp.calls[0] + assert method == "POST" and url.endswith("/proxy/6000/execute") + assert headers["OPEN-SANDBOX-API-KEY"] == "k" + assert headers["X-Session-ID"] == "sess-a" + + def test_non_200_normalizes_to_the_timeout_contract(self, backend): + backend._aiohttp = _FakeAiohttpSession([_FakeAiohttpResponse(500, "boom")]) + with pytest.raises(httpx.TimeoutException): + asyncio.run(backend._send_request({"generated_code": "1+1", "session_id": "s"}, timeout=10.0)) + + def test_502_does_not_replay_stateful_code(self, backend): + backend._aiohttp = _FakeAiohttpSession([_FakeAiohttpResponse(502, "bad gateway")]) + with pytest.raises(httpx.TimeoutException): + asyncio.run(backend._send_request({"generated_code": "1+1", "session_id": "s"}, timeout=10.0)) + assert len(backend._aiohttp.calls) == 1 + + def test_transport_errors_normalize_to_httpx_timeout(self, backend): + import aiohttp as _aiohttp + + class _Raising: + closed = False + + def post(self, *a, **k): + raise _aiohttp.ClientConnectionError("conn reset") + + backend._aiohttp = _Raising() + with pytest.raises(httpx.TimeoutException): + asyncio.run(backend._send_request({"generated_code": "1", "session_id": "s"}, timeout=10.0)) + + def test_delete_session_routes_to_the_pinned_pod_and_releases(self, backend): + async def main(): + await backend._pool.route("sess-a") + backend._aiohttp = _FakeAiohttpSession([_FakeAiohttpResponse(200)]) + await backend.delete_session("sess-a") + + asyncio.run(main()) + method, url, _ = backend._aiohttp.calls[0] + assert method == "DELETE" + assert url.endswith("/sessions/sess-a") + assert "sess-a" not in backend._pool._session_to_slot + + +class _FakeSandbox: + instances = [] + fail_claim = False + + def __init__(self, provider): + self.provider = provider + self.spec = None + self.stops = 0 + self.uploads = [] + self.commands = [] + self.instances.append(self) + + async def start(self, spec): + self.spec = spec + if self.fail_claim and spec.provider_options.get("extensions"): + raise RuntimeError("pool exhausted") + return self + + async def stop(self): + self.stops += 1 + + async def endpoint(self, port): + return SandboxEndpoint(endpoint=f"https://sandbox.example/{port}", headers={"X-Route": "r"}) + + async def upload(self, local_path, remote_path): + self.uploads.append((local_path, remote_path)) + + async def exec(self, command): + self.commands.append(command) + return SimpleNamespace(return_code=0) + + +class TestPoolSandboxApi: + @pytest.fixture(autouse=True) + def fake_sandbox(self, monkeypatch): + import sandbox_pool + + _FakeSandbox.instances = [] + _FakeSandbox.fail_claim = False + monkeypatch.setattr(sandbox_pool, "AsyncSandbox", _FakeSandbox) + + def test_pool_ref_and_direct_fallback_use_sandbox_specs(self): + pool = _pool( + size=1, + pool_ref="warm-pool", + env={"NUM_WORKERS": 4}, + resources={"cpu": 2, "memory_mib": 4096}, + resource_requests={"cpu": 0.5, "memory_mib": 1024}, + ) + + sandbox, from_pool = asyncio.run(pool._acquire_sandbox()) + assert from_pool is True + assert sandbox.spec.provider_options == {"extensions": {"poolRef": "warm-pool"}} + assert sandbox.spec.ports == (6000,) + assert sandbox.provider == PROVIDER + + _FakeSandbox.fail_claim = True + sandbox, from_pool = asyncio.run(pool._acquire_sandbox()) + assert from_pool is False + assert sandbox.spec.entrypoint == ["/start-with-nginx.sh"] + assert sandbox.spec.env == {"NUM_WORKERS": 4} + assert sandbox.spec.resources.cpu == 2 + assert sandbox.spec.resources.memory_mib == 4096 + assert sandbox.spec.provider_options == {"resource_requests": {"cpu": 0.5, "memory_mib": 1024}} + + def test_pool_failure_raises_when_fallback_disabled(self): + pool = _pool(size=1, pool_ref="warm-pool", pool_fallback=False) + _FakeSandbox.fail_claim = True + with pytest.raises(RuntimeError, match="pool exhausted"): + asyncio.run(pool._acquire_sandbox()) + + def test_claim_skips_prepare_and_direct_create_runs_it(self): + pool = _pool( + size=1, + pool_ref="warm-pool", + setup_files={"/opt/setup.py": "/tmp/setup.py"}, + setup_commands=["check"], + service_command="start &", + ) + + async def healthy(*args, **kwargs): + return None + + pool._wait_healthy = healthy + asyncio.run(pool._create_slot_inner(pool._slots[0])) + assert _FakeSandbox.instances[-1].uploads == [] + assert _FakeSandbox.instances[-1].commands == [] + assert pool._slots[0].base_url == "https://sandbox.example/6000" + assert pool._slots[0].headers == {"X-Route": "r"} + + pool._slots[0].sandbox = None + _FakeSandbox.fail_claim = True + asyncio.run(pool._create_slot_inner(pool._slots[0])) + assert _FakeSandbox.instances[-1].uploads == [("/tmp/setup.py", "/opt/setup.py")] + assert _FakeSandbox.instances[-1].commands == ["check", "start &"] + + def test_slow_heal_does_not_block_other_health_checks(self): + pool = _pool(size=2, health_interval_s=0.01) + pool._warmup_done = True + pool._http = _FakeAiohttpSession([_FakeAiohttpResponse(200) for _ in range(20)]) + _admit(pool, 1) + pool._slots[1].sandbox = _FakeSandbox(PROVIDER) + heal_started = asyncio.Event() + + async def blocked_create(slot): + heal_started.set() + await asyncio.Future() + + pool._create_slot_inner = blocked_create + + async def main(): + pool._tasks.append(asyncio.create_task(pool._heal_loop())) + await heal_started.wait() + await asyncio.sleep(0.04) + assert sum(method == "GET" for method, _, _ in pool._http.calls) >= 2 + await pool.aclose() + + asyncio.run(main()) + + def test_heal_attempts_rotate_across_unhealthy_slots(self): + pool = _pool(size=3, health_interval_s=0.01, heal_concurrency=1) + pool._warmup_done = True + attempted = [] + + async def fail_fast(slot): + attempted.append(slot.index) + + pool._heal_slot = fail_fast + + async def main(): + pool._tasks.append(asyncio.create_task(pool._heal_loop())) + for _ in range(30): + if set(attempted) == {0, 1, 2}: + break + await asyncio.sleep(0.01) + await pool.aclose() + + asyncio.run(main()) + assert set(attempted) == {0, 1, 2} + + def test_cancelled_admission_stops_the_sandbox(self): + pool = _pool(size=1) + waiting = asyncio.Event() + + async def wait_forever(*args, **kwargs): + waiting.set() + await asyncio.Future() + + pool._wait_healthy = wait_forever + + async def main(): + task = asyncio.create_task(pool._create_slot_inner(pool._slots[0])) + await waiting.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(main()) + assert _FakeSandbox.instances[0].stops == 1 + assert pool._slots[0].sandbox is None + + def test_cancelled_aclose_finishes_cleanup(self): + class BlockingSandbox(_FakeSandbox): + stop_started = asyncio.Event() + finish_stop = asyncio.Event() + + async def stop(self): + self.stop_started.set() + await self.finish_stop.wait() + await super().stop() + + pool = _pool(size=1) + sandbox = BlockingSandbox(PROVIDER) + pool._slots[0].sandbox = sandbox + pool._http = _FakeAiohttpSession([]) + + async def main(): + first = asyncio.create_task(pool.aclose()) + await sandbox.stop_started.wait() + second = asyncio.create_task(pool.aclose()) + first.cancel() + await asyncio.sleep(0) + first.cancel() + await asyncio.sleep(0) + assert not second.done() + sandbox.finish_stop.set() + with pytest.raises(asyncio.CancelledError): + await first + await second + + asyncio.run(main()) + assert sandbox.stops == 1 + assert pool._slots[0].sandbox is None + assert pool._http.closed is True diff --git a/tests/unit_tests/test_opensandbox_endpoint.py b/tests/unit_tests/test_opensandbox_endpoint.py new file mode 100644 index 0000000000..4a7e0bb9b2 --- /dev/null +++ b/tests/unit_tests/test_opensandbox_endpoint.py @@ -0,0 +1,170 @@ +# 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 +from types import SimpleNamespace + +import pytest + +from nemo_gym.sandbox.api import AsyncSandbox +from nemo_gym.sandbox.providers.base import ( + SandboxEndpoint, + SandboxHandle, + SandboxSpec, + SupportsSandboxEndpoint, +) +from nemo_gym.sandbox.providers.opensandbox import provider as opensandbox_provider + + +pytestmark = pytest.mark.sandbox + + +def _provider() -> opensandbox_provider.OpenSandboxProvider: + return opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + + +def _handle(raw: object) -> SandboxHandle: + return SandboxHandle(sandbox_id="sbx-1", provider_name="opensandbox", raw=raw) + + +class _Connection: + def __init__(self, *, base_url: str = "http://sandbox.example/v1", api_key: str = "", proxy: bool = False): + self._base_url = base_url + self.use_server_proxy = proxy + self.headers = {"OPEN-SANDBOX-API-KEY": api_key} if proxy and api_key else {} + + def get_base_url(self) -> str: + return self._base_url + + +class _RawWithEndpoint: + def __init__( + self, + endpoint: str = "http://sandbox.example/v1/sandboxes/sbx-1/proxy/6000", + headers: dict[str, str] | None = None, + connection: _Connection | None = None, + ): + self._endpoint = endpoint + self._headers = {"OPEN-SANDBOX-API-KEY": "secret"} if headers is None else headers # pragma: allowlist secret + self.connection_config = connection or _Connection(proxy=True) + self.requested_ports: list[int] = [] + + async def get_endpoint(self, port: int) -> SimpleNamespace: + self.requested_ports.append(port) + return SimpleNamespace(endpoint=self._endpoint, headers=self._headers) + + +def test_opensandbox_provider_satisfies_the_endpoint_protocol() -> None: + assert isinstance(_provider(), SupportsSandboxEndpoint) + + +def test_endpoint_returns_the_sdk_url_and_auth_headers() -> None: + raw = _RawWithEndpoint() + resolved = asyncio.run(_provider().endpoint(_handle(raw), 6000)) + assert isinstance(resolved, SandboxEndpoint) + assert resolved.endpoint == "http://sandbox.example/v1/sandboxes/sbx-1/proxy/6000" + assert resolved.headers == {"OPEN-SANDBOX-API-KEY": "secret"} # pragma: allowlist secret + assert raw.requested_ports == [6000] + + +def test_endpoint_requires_a_recent_sdk() -> None: + class RawWithoutEndpoint: + pass + + with pytest.raises(NotImplementedError, match="get_endpoint"): + asyncio.run(_provider().endpoint(_handle(RawWithoutEndpoint()), 6000)) + + +def test_endpoint_rejects_an_empty_url() -> None: + with pytest.raises(RuntimeError, match="empty endpoint"): + asyncio.run(_provider().endpoint(_handle(_RawWithEndpoint(endpoint="")), 6000)) + + +def test_endpoint_headers_default_to_empty_dict_without_configured_key() -> None: + resolved = asyncio.run(_provider().endpoint(_handle(_RawWithEndpoint(headers={})), 6000)) + assert resolved.headers == {} + + +def test_schemeless_endpoint_uses_the_sdk_resolved_url_and_key() -> None: + """A schemeless proxy endpoint inherits the resolved connection URL and credentials.""" + raw = _RawWithEndpoint( + endpoint="sandbox.example/v1/sandboxes/sbx-1/proxy/6000", + headers={}, + connection=_Connection( + base_url="https://sandbox.example/v1", + api_key="secret", # pragma: allowlist secret + proxy=True, + ), + ) + resolved = asyncio.run(_provider().endpoint(_handle(raw), 6000)) + assert resolved.endpoint == "https://sandbox.example/v1/sandboxes/sbx-1/proxy/6000" + assert resolved.headers == {"OPEN-SANDBOX-API-KEY": "secret"} # pragma: allowlist secret + + +def test_direct_mode_endpoint_does_not_inject_the_key() -> None: + raw = _RawWithEndpoint( + endpoint="http://pod.example:6000", + headers={"OPEN-SANDBOX-API-KEY": "secret"}, # pragma: allowlist secret + connection=_Connection(api_key="secret"), # pragma: allowlist secret + ) + resolved = asyncio.run(_provider().endpoint(_handle(raw), 6000)) + assert resolved.headers == {} + + +def test_proxy_key_is_merged_with_sdk_headers() -> None: + raw = _RawWithEndpoint( + headers={"X-Route-Token": "t"}, + connection=_Connection(api_key="secret", proxy=True), # pragma: allowlist secret + ) + resolved = asyncio.run(_provider().endpoint(_handle(raw), 6000)) + assert resolved.headers == { + "X-Route-Token": "t", + "OPEN-SANDBOX-API-KEY": "secret", # pragma: allowlist secret + } + + +def test_sdk_supplied_key_is_never_overridden() -> None: + raw = _RawWithEndpoint( + headers={"OPEN-SANDBOX-API-KEY": "signed"}, # pragma: allowlist secret + connection=_Connection(api_key="secret", proxy=True), # pragma: allowlist secret + ) + resolved = asyncio.run(_provider().endpoint(_handle(raw), 6000)) + assert resolved.headers == {"OPEN-SANDBOX-API-KEY": "signed"} # pragma: allowlist secret + + +def test_async_sandbox_endpoint_flows_through_the_provider() -> None: + """AsyncSandbox.endpoint() must accept the opensandbox provider once the port is declared.""" + + async def main() -> SandboxEndpoint: + provider = _provider() + sandbox = AsyncSandbox(provider=provider, spec=SandboxSpec(image="img", ports=(6000,))) + sandbox._handle = _handle(_RawWithEndpoint()) + sandbox._stopped = False + return await sandbox.endpoint(6000) + + resolved = asyncio.run(main()) + assert resolved.endpoint.endswith("/proxy/6000") + + +def test_async_sandbox_endpoint_still_rejects_undeclared_ports() -> None: + async def main() -> None: + provider = _provider() + sandbox = AsyncSandbox(provider=provider, spec=SandboxSpec(image="img", ports=(6000,))) + sandbox._handle = _handle(_RawWithEndpoint()) + sandbox._stopped = False + await sandbox.endpoint(8080) + + with pytest.raises(ValueError, match="not declared"): + asyncio.run(main()) diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index 8c93efa3ec..ed4b3db867 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -16,6 +16,7 @@ import asyncio import builtins import logging +import os import sys from dataclasses import dataclass from datetime import timedelta @@ -49,6 +50,11 @@ class FakePlatformSpec: class FakeConnectionConfig: def __init__(self, **kwargs: Any) -> None: self.kwargs = kwargs + self.headers = dict(kwargs.get("headers", {})) + self.use_server_proxy = bool(kwargs.get("use_server_proxy", False)) + + def get_api_key(self) -> str: + return self.kwargs.get("api_key") or os.getenv("OPEN_SANDBOX_API_KEY", "") @dataclass(frozen=True) @@ -181,7 +187,7 @@ async def test_provider_conversion_helpers( assert opensandbox_provider._to_volumes([{"name": "workspace"}]) == [FakeVolume(name="workspace")] -async def test_direct_create_passes_platform_to_sdk_create( +async def test_create_passes_provider_options_to_sdk( fake_opensandbox_sdk: None, ) -> None: provider = opensandbox_provider.OpenSandboxProvider( @@ -192,7 +198,10 @@ async def test_direct_create_passes_platform_to_sdk_create( handle = await provider.create( SandboxSpec( image="mirror.gcr.io/astral/uv:python3.12-bookworm-slim", - provider_options={"platform": {"os": "linux", "arch": "amd64"}}, + provider_options={ + "platform": {"os": "linux", "arch": "amd64"}, + "extensions": {"poolRef": "warm"}, + }, ), ) @@ -202,6 +211,7 @@ async def test_direct_create_passes_platform_to_sdk_create( arch="amd64", ) assert "network_policy" not in FakeSandbox.created_kwargs + assert FakeSandbox.created_kwargs["extensions"]["poolRef"] == "warm" async def test_direct_create_passes_network_policy_to_sdk_create(fake_opensandbox_sdk: None) -> None: @@ -499,7 +509,10 @@ def test_provider_options_from_mapping() -> None: options_cls.from_mapping({"volumes": ["workspace"]}) -def test_connection_config_and_image_policy(fake_opensandbox_sdk: None) -> None: +def test_connection_config_and_image_policy( + fake_opensandbox_sdk: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: provider = opensandbox_provider.OpenSandboxProvider( connection={ "domain": "sandbox.example/", @@ -519,11 +532,9 @@ def test_connection_config_and_image_policy(fake_opensandbox_sdk: None) -> None: "protocol": "https", "request_timeout": timedelta(seconds=10), "use_server_proxy": True, - # The API key must also travel as a header: the SDK's execd clients - # (health ping, commands, files) send only ConnectionConfig.headers, - # and proxied /proxy/* routes may enforce auth. "headers": {"OPEN-SANDBOX-API-KEY": "key"}, # pragma: allowlist secret } + assert config.headers == {"OPEN-SANDBOX-API-KEY": "key"} # pragma: allowlist secret short_timeout_config = provider._connection_config(request_timeout_s=3) assert short_timeout_config.kwargs["request_timeout"] == timedelta(seconds=3) @@ -532,7 +543,11 @@ def test_connection_config_and_image_policy(fake_opensandbox_sdk: None) -> None: direct = opensandbox_provider.OpenSandboxProvider( connection={"domain": "sandbox.example", "api_key": "key"} # pragma: allowlist secret ) - assert "headers" not in direct._connection_config().kwargs + assert direct._connection_config().headers == {} + + monkeypatch.setenv("OPEN_SANDBOX_API_KEY", "key-from-env") # pragma: allowlist secret + env_config = opensandbox_provider.OpenSandboxProvider(connection={"use_server_proxy": True})._connection_config() + assert env_config.headers == {"OPEN-SANDBOX-API-KEY": "key-from-env"} # pragma: allowlist secret def test_connection_transport_backends(fake_opensandbox_sdk: None, monkeypatch: pytest.MonkeyPatch) -> None: @@ -1271,6 +1286,74 @@ async def cleanup(handle: opensandbox_provider.SandboxHandle) -> None: assert cleanup_calls == ["sandbox-1"] +async def test_create_once_cleans_up_after_cancellation( + fake_opensandbox_sdk: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": "probe"}) + cleanup_calls: list[str] = [] + verify_started = asyncio.Event() + cleanup_started = asyncio.Event() + allow_cleanup = asyncio.Event() + + async def wait_in_verify(_handle: opensandbox_provider.SandboxHandle) -> None: + verify_started.set() + await asyncio.get_running_loop().create_future() + + async def cleanup(handle: opensandbox_provider.SandboxHandle) -> None: + cleanup_started.set() + await allow_cleanup.wait() + cleanup_calls.append(handle.sandbox_id) + + monkeypatch.setattr(provider, "_verify_created_handle", wait_in_verify) + monkeypatch.setattr(provider, "_cleanup_failed_create_handle", cleanup) + + create_task = asyncio.create_task(provider._create_once(SandboxSpec(image="image:tag"))) + await verify_started.wait() + create_task.cancel() + await cleanup_started.wait() + allow_cleanup.set() + with pytest.raises(asyncio.CancelledError): + await create_task + + assert cleanup_calls == ["sandbox-1"] + + +async def test_close_releases_local_resources_after_cancellation() -> None: + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + kill_started = asyncio.Event() + close_started = asyncio.Event() + allow_close = asyncio.Event() + close_finished = asyncio.Event() + + class Raw: + async def kill(self) -> None: + kill_started.set() + await asyncio.get_running_loop().create_future() + + async def close(self) -> None: + close_started.set() + await allow_close.wait() + close_finished.set() + + close_task = asyncio.create_task( + provider.close( + opensandbox_provider.SandboxHandle( + sandbox_id="sandbox-cancelled", + provider_name="opensandbox", + raw=Raw(), + ), + ) + ) + await kill_started.wait() + close_task.cancel() + await close_started.wait() + allow_close.set() + with pytest.raises(asyncio.CancelledError): + await close_task + assert close_finished.is_set() + + async def test_retry_classification_and_await_sdk_helpers(monkeypatch: pytest.MonkeyPatch) -> None: provider = opensandbox_provider.OpenSandboxProvider( operations={"retries": 0}, diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index 95515ecef4..30f1b1c023 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -352,6 +352,75 @@ async def _assert_async_sandbox_initial_file_error_paths() -> None: await started.start(SandboxSpec(image="image:tag")) +@pytest.mark.asyncio +async def test_async_sandbox_cancellation_during_initial_upload_closes_handle() -> None: + class BlockingUploadProvider(FakeSandboxProvider): + upload_started = asyncio.Event() + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + self.upload_started.set() + await asyncio.Future() + + async def close(self, handle: SandboxHandle) -> None: + await asyncio.sleep(0.01) + await super().close(handle) + + provider = BlockingUploadProvider() + sandbox = AsyncSandbox(provider) + start = asyncio.create_task(sandbox.start(SandboxSpec(image="image:tag", files={"/tmp/bootstrap.txt": "hello"}))) + await provider.upload_started.wait() + start.cancel() + + with pytest.raises(asyncio.CancelledError): + await start + + assert provider.closed == provider.created_handles + assert provider.aclosed is True + + +@pytest.mark.asyncio +async def test_async_sandbox_create_error_closes_provider_resources() -> None: + provider = FakeSandboxProvider() + + async def fail_create(_spec: SandboxSpec) -> SandboxHandle: + raise RuntimeError("create failed") + + provider.create = fail_create + sandbox = AsyncSandbox(provider) + with pytest.raises(RuntimeError, match="create failed"): + await sandbox.start(SandboxSpec(image="image:tag")) + assert provider.aclosed is True + + +@pytest.mark.asyncio +async def test_async_sandbox_cancelled_stop_finishes_cleanup() -> None: + provider = FakeSandboxProvider() + sandbox = await AsyncSandbox(provider).start(SandboxSpec(image="image:tag")) + close_started = asyncio.Event() + finish_close = asyncio.Event() + original_close = provider.close + + async def blocked_close(handle: SandboxHandle) -> None: + close_started.set() + await finish_close.wait() + await original_close(handle) + + provider.close = blocked_close + stopping = asyncio.create_task(sandbox.stop()) + await close_started.wait() + stopping.cancel() + await asyncio.sleep(0) + stopping.cancel() + await asyncio.sleep(0) + assert not stopping.done() + finish_close.set() + with pytest.raises(asyncio.CancelledError): + await stopping + await sandbox.stop() + assert provider.closed == provider.created_handles + assert provider.aclosed is True + + def test_async_sandbox_requires_spec_and_reports_unknown_status() -> None: asyncio.run(_assert_async_sandbox_requires_spec_and_reports_unknown_status())