From c0367cc903d6ae8c6e8d2d22199d4ee74e75ff87 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Fri, 7 Aug 2026 17:47:32 -0700 Subject: [PATCH 1/6] feat: disaggregated sandbox backends on OpenSandbox MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds two opt-in sandbox backends that move code execution off the resources server's own host and onto OpenSandbox pods. Both default to today's behavior: with no environment variables set, the resolved configs are byte-identical to the current ones and neither backend is constructed or imported. ns_tools `sandbox_type: sandbox_pool` A fixed set of long-lived, SHARED pods each running 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 instead of one pod per session. New sessions pin to the least-loaded healthy pod and stay there, so stateful ipython state survives across tool calls. Slots are filled either by claiming a prewarmed pod from a server-side Pool (extensions.poolRef), which drops warmup to allocation time, or by a direct create; a failed or full pool claim degrades to a direct create by default (pool_fallback). A health loop evicts a pod after three consecutive failed probes and heals it in the same slot under a create-rate limit, and an idle sweep drops stale session pins. Every pod carries per-run attribution labels so an epilogue reaper can delete exactly one run's sandboxes. Transport rides a shared aiohttp session — httpx/httpcore's connection pooling collapses at the concurrency this backend targets — while keeping httpx exception types so the NeMo-Skills client contract is unchanged; infra failures degrade rewards rather than crashing the server. math_formal_lean `sandbox_backend: gym_sandbox` Lean4 compilation on OpenSandbox pods via provider exec, reproducing the NS server's lake/lean invocation and its process_status/stdout/stderr contract exactly. With pool_size=0 each verify gets a fresh pod under a bounded semaphore, destroyed in finally. With pool_size=N a warm pool is built at server startup and reused across verifies, because a cold pod's first `import Mathlib` lazy-pulls several GB of olean files one page fault at a time; pool pods bulk-prefetch that tree once at prepare and then serve verifies back-to-back. Failed pods are replaced in place and the verify retries once elsewhere before degrading. A third value, `ns_http_proxy`, speaks the NS protocol through a full base_url plus headers, which serves as a parity oracle against the default path. Also adds `endpoint()` to the OpenSandbox provider, implementing the existing SupportsSandboxEndpoint protocol: it resolves the SDK's server-proxy route for a declared port, absolutizes a scheme-less URL from the configured domain, and carries the auth header the proxy requires. Config surface is env-fed with empty defaults (domain, api_key, image, pool_ref, sizes), and values that must not arrive as strings go through oc.decode or explicit int()/float() coercion. Selecting either backend with an empty domain, api_key, or image is a hard startup error rather than a silent no-op. Tests: 21 sandbox_pool tests, 14 lean backend tests, 9 provider endpoint tests. Signed-off-by: Hemil Desai --- .../sandbox/providers/opensandbox/provider.py | 39 ++ resources_servers/math_formal_lean/app.py | 55 +- .../configs/math_formal_lean.yaml | 24 + .../configs/math_formal_lean_multi_turn.yaml | 24 + .../math_formal_lean/requirements.txt | 3 + .../math_formal_lean/sandbox_client.py | 251 ++++++++- .../tests/test_sandbox_backends.py | 237 +++++++++ resources_servers/ns_tools/app.py | 51 +- .../ns_tools/configs/ns_tools.yaml | 30 +- resources_servers/ns_tools/gym_sandbox.py | 140 +++++ resources_servers/ns_tools/requirements.txt | 2 + resources_servers/ns_tools/sandbox_pool.py | 502 ++++++++++++++++++ .../ns_tools/tests/test_sandbox_pool.py | 468 ++++++++++++++++ tests/unit_tests/test_opensandbox_endpoint.py | 145 +++++ 14 files changed, 1957 insertions(+), 14 deletions(-) create mode 100644 resources_servers/math_formal_lean/tests/test_sandbox_backends.py create mode 100644 resources_servers/ns_tools/gym_sandbox.py create mode 100644 resources_servers/ns_tools/sandbox_pool.py create mode 100644 resources_servers/ns_tools/tests/test_sandbox_pool.py create mode 100644 tests/unit_tests/test_opensandbox_endpoint.py diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 2986079228..86d4a0b690 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -1165,6 +1165,45 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: raw_status = getattr(info, "status", None) return _to_sandbox_status(getattr(raw_status, "state", None) if raw_status is not None else None) + async def endpoint(self, handle: SandboxHandle, port: int) -> SandboxEndpoint: + """Resolve an HTTP(S) endpoint for a declared sandbox port. + + The SDK returns the server-proxy route (`{domain}/v1/sandboxes/{id}/proxy/{port}`) + when the connection uses the server proxy, along with any headers the server + requires on every request to that endpoint (e.g. the API key header). + """ + 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: 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(getattr(resolved, "endpoint", "") or "") + if not endpoint_url: + raise RuntimeError(f"OpenSandbox returned an empty endpoint for sandbox {handle.sandbox_id} port {port}") + if "://" not in endpoint_url: + # The SDK returns the proxy endpoint without a scheme; borrow it from the + # configured server domain. + domain = str(self._connection.domain or "") + scheme = "https" if domain.startswith("https://") else "http" + endpoint_url = f"{scheme}://{endpoint_url}" + headers = dict(getattr(resolved, "headers", None) or {}) + if not headers and self._connection.use_server_proxy and self._connection.api_key: + # Defensive: the SDK currently returns no headers for proxy endpoints; include + # the API key so authenticated proxy deployments work either way. Proxy mode + # only — a direct endpoint terminates at the sandbox, which runs untrusted + # code and must never be handed the key. + headers["OPEN-SANDBOX-API-KEY"] = str(self._connection.api_key) + return SandboxEndpoint(endpoint=endpoint_url, headers=headers) + def _command_retry_count(self) -> int: return self._operations.command_retries diff --git a/resources_servers/math_formal_lean/app.py b/resources_servers/math_formal_lean/app.py index f81a261607..78de135847 100644 --- a/resources_servers/math_formal_lean/app.py +++ b/resources_servers/math_formal_lean/app.py @@ -20,7 +20,7 @@ from dataclasses import dataclass from typing import Any, Dict, List, Optional -from pydantic import BaseModel +from pydantic import BaseModel, Field from nemo_gym.base_resources_server import ( BaseResourcesServerConfig, @@ -29,7 +29,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 +341,15 @@ def build_correction_prompt( class MathFormalLeanResourcesServerConfig(BaseResourcesServerConfig): sandbox_host: str = "127.0.0.1" sandbox_port: int = 6000 + # Sandbox backend: "ns_http" (default — today's NS server over host/port), "gym_sandbox" + # (per-verify OpenSandbox pods via provider exec), or "ns_http_proxy" (the NS HTTP protocol + # through a full base_url + headers, e.g. an OpenSandbox proxied endpoint; parity oracle). + sandbox_backend: str = "ns_http" + sandbox_base_url: str = "" + sandbox_extra_headers: Dict[str, str] = Field(default_factory=dict) + # 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 +393,49 @@ 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, + ) + elif self.config.sandbox_backend == "ns_http_proxy": + self._sandbox_client = Lean4SandboxClient( + base_url=self.config.sandbox_base_url, + extra_headers=self.config.sandbox_extra_headers, + max_output_characters=self.config.max_output_characters, + ) + 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() + start_pool = getattr(self._sandbox_client, "start_pool", None) + if start_pool is None: + return app + from contextlib import asynccontextmanager + + main_app_lifespan = app.router.lifespan_context + + @asynccontextmanager + async def lifespan_wrapper(app): + # Warm lean pool pods from startup: a cold pod's first compile is ~15 min + # of nydus olean pulls, far beyond any verify's admission window. + start_pool() + async with main_app_lifespan(app) as maybe_state: + yield maybe_state + + 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..3daf78102d 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,30 @@ 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} + # Sandbox backend: 'ns_http' (default — today's path, byte-identical with zero env + # vars set), 'gym_sandbox' (per-verify OpenSandbox pods via provider exec), or + # 'ns_http_proxy' (NS protocol through a full base_url — parity oracle). + 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 = a fresh pod per verify. N > 0 = N warm pods reused across verifies — + # a cold pod's first compile lazy-pulls ~5GB of oleans (~15 min); warm ~4s. + 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..771ac9e88b 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,30 @@ 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} + # Sandbox backend: 'ns_http' (default — today's path, byte-identical with zero env + # vars set), 'gym_sandbox' (per-verify OpenSandbox pods via provider exec), or + # 'ns_http_proxy' (NS protocol through a full base_url — parity oracle). + 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 = a fresh pod per verify. N > 0 = N warm pods reused across verifies — + # a cold pod's first compile lazy-pulls ~5GB of oleans (~15 min); warm ~4s. + 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..d04abcbc14 100644 --- a/resources_servers/math_formal_lean/requirements.txt +++ b/resources_servers/math_formal_lean/requirements.txt @@ -1,2 +1,5 @@ -e nemo-gym[dev] @ ../../ httpx>=0.27.0 +## OpenSandbox SDK — used only by the opt-in opensandbox/ns_http_proxy sandbox backends +opensandbox==0.1.15 +tenacity diff --git a/resources_servers/math_formal_lean/sandbox_client.py b/resources_servers/math_formal_lean/sandbox_client.py index e49b0f46d7..1f39f8da36 100644 --- a/resources_servers/math_formal_lean/sandbox_client.py +++ b/resources_servers/math_formal_lean/sandbox_client.py @@ -22,6 +22,7 @@ import json import logging +import os from typing import Any, Dict import httpx @@ -38,6 +39,8 @@ def __init__( host: str = "127.0.0.1", port: int = 6000, max_output_characters: int = 1000, + base_url: str | None = None, + extra_headers: Dict[str, str] | None = None, ): """Initialize sandbox client. @@ -45,10 +48,15 @@ def __init__( host: Sandbox server hostname port: Sandbox server port max_output_characters: Maximum characters in output + base_url: Full base URL override (e.g. an OpenSandbox proxied endpoint); + when set, host/port are ignored. + extra_headers: Headers added to every request (e.g. proxy auth). """ self.host = host self.port = port self.max_output_characters = max_output_characters + self.base_url = base_url.rstrip("/") if base_url else None + self.extra_headers = dict(extra_headers or {}) self._client: httpx.AsyncClient | None = None async def _get_client(self) -> httpx.AsyncClient: @@ -67,6 +75,8 @@ async def close(self) -> None: def _get_execute_url(self) -> str: """Get the sandbox execute endpoint URL.""" + if self.base_url: + return f"{self.base_url}/execute" return f"http://{self.host}:{self.port}/execute" async def execute_lean4( @@ -97,7 +107,7 @@ async def execute_lean4( url=self._get_execute_url(), content=json.dumps(request_data), timeout=timeout + 5.0, # Add buffer for network overhead - headers={"Content-Type": "application/json"}, + headers={"Content-Type": "application/json", **self.extra_headers}, ) if response.status_code == 502: @@ -127,11 +137,246 @@ async def health_check(self, timeout: float = 5.0) -> bool: Returns: True if sandbox is healthy, False otherwise """ - url = f"http://{self.host}:{self.port}/health" + base = self.base_url if self.base_url else f"http://{self.host}:{self.port}" + url = f"{base}/health" client = await self._get_client() try: - response = await client.get(url=url, timeout=timeout) + response = await client.get(url=url, timeout=timeout, headers=self.extra_headers) return response.status_code == 200 except httpx.HTTPError: return False + + +class GymSandboxLean4Client: + """Lean4 compilation on per-verify OpenSandbox pods via provider exec. + + Reimplements the NS server's lean4 invocation exactly (reference frozen at + nemo_skills local_sandbox_server.py:631-685 @ da85a881): the proof lands in + /lean4/my_project, `lake env --dir /lean4/my_project lean ` runs with an + in-sandbox `timeout -s KILL`, and the exit code maps to the same + process_status/stdout/stderr contract as `Lean4SandboxClient.execute_lean4`. + Long compiles never hold an HTTP connection open (background/short exec + requests), so proxy read-timeout ceilings do not apply. + + With ``pool_size=0`` every verify gets a fresh pod (created under a bounded + semaphore, destroyed in finally). With ``pool_size=N`` a warm pool of N pods is + built lazily and reused across verifies: a fresh pod's first `import Mathlib` + lazy-pulls ~5GB of olean files through nydus one FUSE fault at a time (measured + at ~900s), while a warmed pod compiles in ~4s — so pool pods bulk-prefetch + the olean tree once at prepare and then serve verifies back-to-back. A pod that + hits an infra error (incl. TTL expiry) is killed and replaced in place; the + verify retries once on a second pod before degrading. Infra failures degrade to + the client's existing error/timeout shapes, never raise into verify(). + """ + + 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 = (next(iter(provider.values()), {}) 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: Any = None # bound lazily to the serving event loop + self._pool_size = int(pool_size) + self._prefetch_paths = prefetch_paths + self._pool_ref = pool_ref or "" + self._pool: Any = None # asyncio.Queue of warm AsyncSandbox pods, filled lazily + self._pool_started = False + + def _get_semaphore(self): + import asyncio + + if self._semaphore is None: + self._semaphore = asyncio.Semaphore(self._semaphore_size) + return self._semaphore + + def _new_sandbox(self, files: Dict[str, str] | None = None, use_pool: bool = True): + from nemo_gym.sandbox.api import AsyncSandbox + from nemo_gym.sandbox.providers.base import SandboxSpec + + # 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): + """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 _create_pool_pod(self): + """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 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: + self._pool_queue() + + @property + def pool_ready_count(self) -> int: + return self._pool.qsize() if self._pool is not None else 0 + + def _pool_queue(self): + import asyncio + + if self._pool is None: + self._pool = asyncio.Queue() + for _ in range(self._pool_size): + asyncio.get_running_loop().create_task(self._fill_one()) + return self._pool + + async def _fill_one(self): + try: + self._pool.put_nowait(await self._create_pool_pod()) + except Exception as e: + LOG.error("lean pool pod create failed (capacity reduced until next heal): %s", e) + + async def _execute_pooled(self, code: str, timeout: float) -> Dict[str, Any]: + import asyncio + import uuid + + pool = self._pool_queue() + 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: + import tempfile + + 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 Exception as e: + # Pod is suspect (TTL expiry, node loss): replace it, retry once elsewhere. + LOG.warning("lean pool pod failed (attempt %d), replacing: %s", attempt, e) + try: + await sandbox.stop() + except Exception: + pass + asyncio.get_running_loop().create_task(self._fill_one()) + if attempt == 1: + continue + return {"process_status": "error", "stdout": "", "stderr": str(e)} + 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 "")[: self.max_output_characters] + stderr = (result.stderr or "")[: self.max_output_characters] + if result.return_code == 0: + return {"process_status": "completed", "stdout": stdout, "stderr": stderr} + if result.return_code in (124, 137, -9): + return { + "process_status": "timeout", + "stdout": stdout, + "stderr": stderr + f"Execution timed out after {timeout} seconds\n", + } + return {"process_status": "failed", "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.""" + import asyncio + import uuid + + if self._pool_size > 0: + return await self._execute_pooled(code, timeout) + + try: + await asyncio.wait_for(self._get_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._get_semaphore().release() + if sandbox is not None: + try: + await sandbox.stop() + except Exception as exc: + LOG.warning("lean sandbox teardown failed (TTL will reap): %s", exc) 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..8a8b047219 --- /dev/null +++ b/resources_servers/math_formal_lean/tests/test_sandbox_backends.py @@ -0,0 +1,237 @@ +# 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: the additive base_url/header extension must +leave the default byte-identical, and the OpenSandbox exec client must reproduce the NS +server's process_status contract exactly (reference: local_sandbox_server.py:631-685).""" + +import asyncio +import sys +from pathlib import Path +from types import SimpleNamespace + +import httpx +import pytest + + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +from resources_servers.math_formal_lean.sandbox_client import ( # noqa: E402 + GymSandboxLean4Client, + Lean4SandboxClient, +) + + +PROVIDER = { + "opensandbox": { + "connection": {"domain": "http://sandbox.example", "api_key": "k", "use_server_proxy": True}, + } +} + + +class TestHttpClientStaysByteIdentical: + def test_default_url_is_unchanged(self): + client = Lean4SandboxClient() + assert client._get_execute_url() == "http://127.0.0.1:6000/execute" + + def test_base_url_override_and_headers(self): + seen = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["url"] = str(request.url) + seen["auth"] = request.headers.get("open-sandbox-api-key") + return httpx.Response(200, json={"process_status": "completed", "stdout": "", "stderr": ""}) + + client = Lean4SandboxClient( + base_url="http://sandbox.example/v1/sandboxes/sbx/proxy/6000", + extra_headers={"OPEN-SANDBOX-API-KEY": "k"}, + ) + client._client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + out = asyncio.run(client.execute_lean4("theorem t : True := trivial", timeout=5.0)) + assert out["process_status"] == "completed" + assert seen["url"] == "http://sandbox.example/v1/sandboxes/sbx/proxy/6000/execute" + assert seen["auth"] == "k" + + +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 nemo_gym.sandbox.api as api + + _FakeSandbox.instances = [] + _FakeSandbox.next_exec_result = None + _FakeSandbox.next_raise_on_exec = None + monkeypatch.setattr(api, "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_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._get_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" + + +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 + + +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..4eedaaa3af 100644 --- a/resources_servers/ns_tools/app.py +++ b/resources_servers/ns_tools/app.py @@ -78,6 +78,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: str = "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 @@ -88,6 +95,11 @@ class NSToolsConfig(BaseResourcesServerConfig): # The model receives a warning in stderr instead of restored state. disable_session_restore: bool = False + # Staged guard for tool calls arriving WITHOUT a gym session cookie (each such call + # silently mints a fresh sandbox session today). False = log + count (default, + # behavior-neutral); True = reject with 400 once the cookie path is proven end-to-end. + strict_session_cookie: bool = False + # ============================================================ # Run/Verify Request/Response Models @@ -131,6 +143,7 @@ class NSToolsResourcesServer(SimpleResourcesServer): _tool_name_map: Dict[str, str] = {} # Maps tool names to qualified names _python_tool_process: Optional[subprocess.Popen] = None _timing_by_session: Dict[str, list] = {} # session_id -> list of timing records + _missing_cookie_count: int = 0 _uses_python_tool_sidecar: bool = False def setup_webserver(self) -> FastAPI: @@ -144,6 +157,12 @@ def setup_webserver(self) -> FastAPI: @asynccontextmanager async def lifespan_wrapper(app): + if self.config.sandbox_type == "sandbox_pool": + import gym_sandbox + + if gym_sandbox.CURRENT_POOL is not None: + # Budgeted, non-blocking warmup: kicks pod creation without gating server boot. + await gym_sandbox.CURRENT_POOL.start() try: async with main_app_lifespan(app) as maybe_state: yield maybe_state @@ -262,12 +281,16 @@ def _initialize_nemo_skills_tools(self): context = { "sandbox": { - "sandbox_type": "local", + "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": + import gym_sandbox # noqa: F401 — registers the backend with the nemo_skills registry + + context["sandbox"]["pool"] = dict(self.config.sandbox_pool) overrides = { tool_name: dict(tool_config) for tool_name, tool_config in self.config.nemo_skills_tool_overrides.items() @@ -315,8 +338,19 @@ async def execute_tool(self, tool_name: str, request: Request) -> PlainTextRespo # Get session ID for stateful execution session_id = request.session.get(SESSION_ID_KEY) if not session_id: + # A missing gym cookie means every call mints a fresh sandbox session — silent + # per-call state loss. Staged fix: count + log by default; reject only once the + # cookie path is proven end-to-end and strict_session_cookie is flipped. + self._missing_cookie_count += 1 + if self.config.strict_session_cookie: + return PlainTextResponse( + json.dumps({"error": "missing session cookie; stateful tools require a session"}), + status_code=400, + ) session_id = str(uuid.uuid4()) - logger.warning(f"No session ID found, using fallback: {session_id}") + logger.warning( + f"No session ID found (occurrence {self._missing_cookie_count}), using fallback: {session_id}" + ) if session_id not in self._timing_by_session: self._timing_by_session[session_id] = [] @@ -463,7 +497,18 @@ 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() + try: + await self.tool_manager.shutdown() + except (asyncio.CancelledError, Exception): + # Tool-manager teardown must not skip pool teardown: unkilled pool + # pods stay allocated (and billed against pool capacity) until TTL. + logger.warning("tool_manager.shutdown failed; continuing to pool teardown", exc_info=True) + + if self.config.sandbox_type == "sandbox_pool": + import gym_sandbox + + if gym_sandbox.CURRENT_POOL is not None: + await gym_sandbox.CURRENT_POOL.aclose() # Terminate the python_tool subprocess if one was started. if self._python_tool_process: diff --git a/resources_servers/ns_tools/configs/ns_tools.yaml b/resources_servers/ns_tools/configs/ns_tools.yaml index 24bd4446be..23a7819f14 100644 --- a/resources_servers/ns_tools/configs/ns_tools.yaml +++ b/resources_servers/ns_tools/configs/ns_tools.yaml @@ -40,7 +40,35 @@ 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 + # First-ever create on a new image tag blocks on nydus conversion (40-55s); + # the SDK default request_timeout (30s) fails the whole cold warmup wave. + request_timeout: 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}} + 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..f62b14e994 --- /dev/null +++ b/resources_servers/ns_tools/gym_sandbox.py @@ -0,0 +1,140 @@ +# 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 sandbox backend that routes through an OpenSandbox pod pool. + +Registers ``sandbox_type: sandbox_pool`` with the nemo_skills sandbox registry. The +class IS a ``LocalSandbox`` — same request preparation, same session bookkeeping — with the +transport re-pointed: each request resolves (base_url, headers) from the pool by session +uuid, and rides a shared AIOHTTP session (httpx/httpcore's O(n^2) connection pooling +collapses at high concurrency — see CLAUDE.md; measured: health-only GETs fell +from 87 to 8 calls/s between 64 and 512 in-flight on httpx). Exception TYPES stay httpx +because the nemo_skills base class's execute_code catches those; anything non-200 or +transport-level is normalized to the NS timeout contract so infra failures degrade rewards +without new error shapes. + +Importing this module is the opt-in: the default ``local`` backend never imports it. +""" + +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__) + +# The pool the owning server can warm up at lifespan startup (set by the first construction). +CURRENT_POOL: Optional[SandboxPool] = None + + +class GymSandbox(ns_sandbox.LocalSandbox): + """LocalSandbox with the transport routed through an OpenSandbox pod pool over aiohttp.""" + + def __init__(self, pool: Optional[Dict[str, Any]] = None, **kwargs: Any) -> None: + super().__init__(**kwargs) + if not pool: + raise ValueError("sandbox_type=sandbox_pool requires a 'pool' config dict") + global CURRENT_POOL + self._pool = SandboxPool(**pool) + self._aiohttp: Optional[aiohttp.ClientSession] = None + CURRENT_POOL = self._pool + + 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) + if status == 502: + # A proxy-minted 502 means the pod never received the request, so ONE retry + # is idempotency-safe even for stateful ipython. + 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: + await self._pool.aclose() + finally: + if self._aiohttp is not None and not self._aiohttp.closed: + await self._aiohttp.close() + await super().close() + + +ns_sandbox.sandboxes["sandbox_pool"] = GymSandbox diff --git a/resources_servers/ns_tools/requirements.txt b/resources_servers/ns_tools/requirements.txt index 1a920f13ff..e6812c33c7 100644 --- a/resources_servers/ns_tools/requirements.txt +++ b/resources_servers/ns_tools/requirements.txt @@ -1,3 +1,5 @@ -e nemo-gym[dev] @ ../../ ## NeMo-Skills tools subpackage as of June 12, 2026 nemo-skills-tools @ git+https://github.com/NVIDIA-NeMo/Skills.git@da85a881d972e6fec847b90cf553a0bf9bf10638#subdirectory=tools +## OpenSandbox SDK — used only by the opt-in sandbox_pool sandbox backend +opensandbox==0.1.15 diff --git a/resources_servers/ns_tools/sandbox_pool.py b/resources_servers/ns_tools/sandbox_pool.py new file mode 100644 index 0000000000..59045fdb3b --- /dev/null +++ b/resources_servers/ns_tools/sandbox_pool.py @@ -0,0 +1,502 @@ +# 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: 16k concurrent sessions ride K pods +(each pod's NS server multiplexes many sessions), instead of 16k pods. Slots fill +and heal with direct async ``Sandbox.create`` calls — a full fan-out warm wave is +~5s per pod measured in production, so no warm-spare inventory layer is needed. + +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 datetime import timedelta +from typing import Any, Dict, Optional, Tuple + +import aiohttp +import httpx # exception types only: the nemo_skills client contract catches httpx errors + +from nemo_gym.sandbox.attribution import RUN_KEY, resolve_attribution, resolve_run_id +from nemo_gym.sandbox.providers.opensandbox.provider import ( + DEFAULT_ATTRIBUTION_KEY_PREFIX as _ATTRIBUTION_KEY_PREFIX, +) + + +LOGGER = logging.getLogger(__name__) + + +def _parse_connection(provider: Dict[str, Any]) -> Dict[str, Any]: + """Pull the connection kwargs out of a single-key provider config dict.""" + if not isinstance(provider, dict) or len(provider) != 1: + raise ValueError("sandbox_pool.provider must be a single-key provider config dict") + kwargs = next(iter(provider.values())) or {} + connection = dict(kwargs.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" + ) + return connection + + +@dataclass +class _Slot: + index: int + sandbox: Any = None # opensandbox.Sandbox + 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 = 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: Optional[float] = None, + env: Optional[Dict[str, str]] = None, + entrypoint: Optional[list] = None, + resources: Optional[Dict[str, str]] = None, + resource_requests: Optional[Dict[str, str]] = None, + setup_files: Optional[Dict[str, str]] = None, + setup_commands: Optional[list] = None, + service_command: Optional[str] = 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: + self._connection_kwargs = _parse_connection(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 + # Hydra/YAML overrides deliver bare numbers as ints; the create API's env map + # is string->string and the server 422s on anything else. + self._env = {k: str(v) for k, v in (env or {}).items()} + self._entrypoint = list(entrypoint) if entrypoint else None + self._resources = {k: str(v) for k, v in (resources or {}).items()} + # k8s schedules on REQUESTS; keeping them far below limits packs many more pods + # (= sessions) per node while bursts still get the limit headroom. + self._resource_requests = {k: str(v) for k, v in (resource_requests or {}).items()} + self._setup_files = dict(setup_files or {}) + self._setup_commands = list(setup_commands or []) + self._service_command = service_command + 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._connection_config: Any = None + 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 + attribution = resolve_attribution() + # No explicit label: NEMO_GYM_RUN_ID (set per job by the launch script) wins, + # else a per-process id — either way unique per run, so an epilogue reaper can + # delete exactly this run's sandboxes by the run attribution label. + attribution[RUN_KEY] = resolve_run_id() + self._metadata = {f"{_ATTRIBUTION_KEY_PREFIX}{k}": v for k, v in attribution.items()} + 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._tasks: list = [] + self._last_heal_create = 0.0 + self._http: Any = None # aiohttp session; created lazily on the serving loop + + # ------------------------------------------------------------------ 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 + from opensandbox.config.connection import ConnectionConfig + + kwargs = dict(self._connection_kwargs) + # Mirror of the provider's proxy-mode auth (PR 2462): the SDK's + # execd-facing clients (the create ready gate's health ping) send only + # ConnectionConfig.headers, so on servers that enforce auth on + # /proxy/* routes every ping 401s and the claim dies at ready_timeout. + # Inject the key only in proxy mode — a direct sandbox endpoint runs + # untrusted code and must never see it. + if kwargs.get("use_server_proxy") and kwargs.get("api_key"): + headers = dict(kwargs.get("headers") or {}) + headers.setdefault("OPEN-SANDBOX-API-KEY", str(kwargs["api_key"])) + kwargs["headers"] = headers + self._connection_config = ConnectionConfig(**kwargs) + 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: + self._closed = True + for task in self._tasks: + task.cancel() + for task in self._tasks: + try: + await task + except (asyncio.CancelledError, Exception): + pass + self._tasks.clear() + + async def _kill(slot: _Slot) -> None: + try: + await asyncio.wait_for(slot.sandbox.kill(), timeout=30.0) + except Exception as exc: + LOGGER.warning("pool slot %d teardown failed (TTL will reap): %s", slot.index, exc) + slot.sandbox = None + slot.healthy = False + + await asyncio.gather(*(_kill(slot) for slot in self._slots if slot.sandbox is not None)) + if self._http is not None and not self._http.closed: + await self._http.close() + + # ------------------------------------------------------------------ slot fill + + @property + def _needs_prepare(self) -> bool: + return bool(self._setup_files or self._setup_commands or self._service_command) + + async def _prepare(self, sandbox: Any) -> None: + """Bootstrap a pod. INVARIANT: the service start is the LAST execd command this pod + ever sees — execd reaps a completed command's backgrounded children when any later + command runs (probed empirically: service->touch->10s = dead listener; service-> + nothing = alive).""" + for target_path, local_path in self._setup_files.items(): + with open(local_path, "rb") as fh: + await sandbox.files.write_file(target_path, fh.read()) + for command in self._setup_commands: + execution = await sandbox.commands.run(command) + if (execution.exit_code or 0) != 0: + raise RuntimeError(f"setup command failed rc={execution.exit_code}: {command!r}") + if self._service_command: + # Plain shell backgrounding (cmd &): setsid and execd background:true both + # freeze the child during module import (probed); a plain & child reparents + # to the pod's PID 1 and survives — as long as nothing execs afterwards. + execution = await sandbox.commands.run(self._service_command) + if (execution.exit_code or 0) != 0: + raise RuntimeError(f"service command failed rc={execution.exit_code}") + + def _normalize_endpoint(self, resolved: Any) -> Tuple[str, Dict[str, str]]: + url = str(getattr(resolved, "endpoint", "") or "") + if not url: + raise RuntimeError("SDK returned an empty sandbox endpoint") + if "://" not in url: + domain = str(self._connection_kwargs.get("domain") or "") + scheme = "https" if domain.startswith("https://") else "http" + url = f"{scheme}://{url}" + headers = dict(getattr(resolved, "headers", None) or {}) + if not headers and self._connection_kwargs.get("use_server_proxy") and self._connection_kwargs.get("api_key"): + # Proxy mode only — a direct endpoint terminates at the sandbox, which runs + # untrusted code and must never be handed the key. + headers["OPEN-SANDBOX-API-KEY"] = str(self._connection_kwargs["api_key"]) + return url.rstrip("/"), headers + + async def _create_slot(self, slot: _Slot) -> None: + """Fill one slot. Single-flight per slot: a duplicate landing late would + overwrite base_url under pinned sessions and leak a pod until TTL.""" + if slot.creating: + return + slot.creating = True + try: + await self._create_slot_inner(slot) + finally: + slot.creating = False + + async def _acquire_sandbox(self) -> Tuple[Any, 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).""" + from opensandbox import Sandbox + + if self._pool_ref: + try: + sandbox = await Sandbox.create( + # The SDK's local validation requires an image even in pool mode; + # the pool template still defines what actually runs. + image=self._image, + extensions={"poolRef": self._pool_ref}, + metadata=dict(self._metadata), + timeout=timedelta(seconds=self._ttl_s or 14400.0), + ready_timeout=timedelta(seconds=self._ready_timeout_s), + connection_config=self._connection_config, + ) + 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 = await Sandbox.create( + image=self._image, + entrypoint=self._entrypoint, + env=self._env or None, + metadata=dict(self._metadata), + resource=self._resources or None, + resource_requests=self._resource_requests or None, + timeout=timedelta(seconds=self._ttl_s or 14400.0), + ready_timeout=timedelta(seconds=self._ready_timeout_s), + connection_config=self._connection_config, + ) + return sandbox, False + + async def _create_slot_inner(self, slot: _Slot) -> None: + sandbox, from_pool = await self._acquire_sandbox() + try: + # Pool pods are born with their service running; only direct creates + # (no pool, or fallback) need bootstrap. + if self._needs_prepare and not from_pool: + await self._prepare(sandbox) + resolved = await sandbox.get_endpoint(self._port) + base_url, headers = self._normalize_endpoint(resolved) + await self._wait_healthy(base_url, headers, budget_s=self._health_budget_s) + except Exception: + try: + await sandbox.kill() + except Exception: + pass + 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: Optional[str] = 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: + try: + await self._create_slot(slot) + except Exception as exc: + LOGGER.warning("pool warmup: slot %d failed (heal loop will retry): %s", slot.index, exc) + + 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) + if not self._warmup_done: + # Warmup owns every slot until it finishes; healing in parallel would + # race duplicate acquisitions into the same slot. + continue + + 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 True + 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: + try: + await slot.sandbox.kill() + except Exception: + pass + 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 + semaphore = asyncio.Semaphore(self._heal_concurrency) + + async def heal(slot: _Slot) -> None: + async with semaphore: + await self._heal_slot(slot) + + await asyncio.gather(*(heal(slot) for slot in to_heal)) + + 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.""" + 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) + try: + await self._create_slot(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, + ) + + 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: Optional[str]) -> 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_sandbox_pool.py b/resources_servers/ns_tools/tests/test_sandbox_pool.py new file mode 100644 index 0000000000..48e3cd7c94 --- /dev/null +++ b/resources_servers/ns_tools/tests/test_sandbox_pool.py @@ -0,0 +1,468 @@ +# 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 + +import httpx +import pytest + + +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) + 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") + + 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") + + def test_empty_image_is_a_hard_error(self): + with pytest.raises(ValueError, match="NS_SANDBOX_IMAGE"): + SandboxPool(provider=PROVIDER, image="") + + 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) + + +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=dict(provider=PROVIDER, image="img", size=1), + host="127.0.0.1", + port="6000", + disable_session_restore=True, + ) + _admit(sandbox._pool, 0) + return sandbox + + def test_backend_registers_with_the_nemo_skills_registry(self): + pytest.importorskip("nemo_skills") + import gym_sandbox + from nemo_skills.code_execution.sandbox import sandboxes + + assert sandboxes["sandbox_pool"] is gym_sandbox.GymSandbox + + 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_retries_exactly_once_then_succeeds(self, backend): + ok = '{"process_status": "completed", "stdout": "", "stderr": ""}' + backend._aiohttp = _FakeAiohttpSession( + [_FakeAiohttpResponse(502, "bad gateway"), _FakeAiohttpResponse(200, ok)] + ) + result = asyncio.run(backend._send_request({"generated_code": "1+1", "session_id": "s"}, timeout=10.0)) + assert result["process_status"] == "completed" + assert len(backend._aiohttp.calls) == 2 + + 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 TestHealWarmupRace: + """The heal loop must never race duplicate creates into a slot warmup still owns — + a late duplicate overwrites base_url under pinned sessions and breaks stickiness + (observed under slow pod creation).""" + + def test_create_slot_is_single_flight(self): + pool = _pool() + calls = {"n": 0} + + async def fake_inner(slot): + calls["n"] += 1 + await asyncio.sleep(0.05) + + pool._create_slot_inner = fake_inner + + async def main(): + await asyncio.gather(pool._create_slot(pool._slots[0]), pool._create_slot(pool._slots[0])) + + asyncio.run(main()) + assert calls["n"] == 1, "second concurrent create for the same slot must be a no-op" + + def test_heal_loop_waits_for_warmup(self): + pool = _pool() + assert pool._warmup_done is False + healed = [] + pool._heal_slot = lambda slot: healed.append(slot.index) + + async def one_heal_pass(): + pool._health_interval_s = 0.01 + task = asyncio.create_task(pool._heal_loop()) + await asyncio.sleep(0.05) + pool._closed = True + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(one_heal_pass()) + assert healed == [], "heal loop must not touch slots before warmup completes" + + +class TestEnvStringify: + def test_env_values_are_stringified(self): + from sandbox_pool import SandboxPool + + pool = SandboxPool( + provider={"opensandbox": {"connection": {"domain": "http://sandbox.example", "api_key": "k"}}}, + image="img:tag", + env={"NUM_WORKERS": 4, "FLAG": True}, + ) + # The create API's env map is string->string; ints/bools 422 server-side. + assert pool._env == {"NUM_WORKERS": "4", "FLAG": "True"} + + +class TestPoolRefFallback: + """pool_ref acquire semantics — SDK required (create is monkeypatched, no network).""" + + def _pool(self, **overrides): + from sandbox_pool import SandboxPool + + kwargs = dict( + provider={"opensandbox": {"connection": {"domain": "http://sandbox.example", "api_key": "k"}}}, + image="img:tag", + pool_ref="warm-pool", + size=1, + service_command="sh -c 'start & echo ok'", + ) + kwargs.update(overrides) + return SandboxPool(**kwargs) + + def test_pool_full_falls_back_to_direct_create(self, monkeypatch): + opensandbox = pytest.importorskip("opensandbox") + calls = [] + + async def fake_create(**kwargs): + calls.append(kwargs) + if "extensions" in kwargs: + raise RuntimeError("pool exhausted") + return object() + + monkeypatch.setattr(opensandbox.Sandbox, "create", staticmethod(fake_create)) + pool = self._pool() + pool._connection_config = object() + sandbox, from_pool = asyncio.run(pool._acquire_sandbox()) + assert from_pool is False and sandbox is not None + assert calls[0]["extensions"] == {"poolRef": "warm-pool"} + # The fallback create carries the full direct spec, not the pool claim shape. + assert "extensions" not in calls[1] and calls[1]["image"] == "img:tag" + + def test_pool_failure_raises_when_fallback_disabled(self, monkeypatch): + opensandbox = pytest.importorskip("opensandbox") + + async def fake_create(**kwargs): + raise RuntimeError("pool exhausted") + + monkeypatch.setattr(opensandbox.Sandbox, "create", staticmethod(fake_create)) + pool = self._pool(pool_fallback=False) + pool._connection_config = object() + with pytest.raises(RuntimeError, match="pool exhausted"): + asyncio.run(pool._acquire_sandbox()) + + def test_pool_claim_skips_prepare_and_fallback_does_not(self, monkeypatch): + opensandbox = pytest.importorskip("opensandbox") + prepared = [] + + async def fake_create(**kwargs): + if "extensions" in kwargs and fake_create.pool_ok: + return "pool-pod" + if "extensions" in kwargs: + raise RuntimeError("pool exhausted") + return "direct-pod" + + monkeypatch.setattr(opensandbox.Sandbox, "create", staticmethod(fake_create)) + pool = self._pool() + pool._connection_config = object() + + async def fake_prepare(sandbox): + prepared.append(sandbox) + + async def fake_endpoint_flow(slot, sandbox): + slot.sandbox = sandbox + slot.healthy = True + + pool._prepare = fake_prepare + + async def run_inner(pool_ok): + fake_create.pool_ok = pool_ok + sandbox, from_pool = await pool._acquire_sandbox() + if pool._needs_prepare and not from_pool: + await pool._prepare(sandbox) + return sandbox + + assert asyncio.run(run_inner(True)) == "pool-pod" + assert prepared == [] + assert asyncio.run(run_inner(False)) == "direct-pod" + assert prepared == ["direct-pod"] + + +class TestProxyAuthHeaders: + """Mirror of the provider's proxy-mode auth (PR 2462). The SDK's execd-facing + clients authenticate only via ConnectionConfig.headers, so without the key + there the create ready gate's health ping 401s on servers that enforce auth + on /proxy/* routes and every claim dies at ready_timeout.""" + + def _captured_kwargs(self, monkeypatch, connection) -> dict: + opensandbox = pytest.importorskip("opensandbox") + import opensandbox.config.connection as osb_connection + + captured: dict = {} + + class FakeConnectionConfig: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(osb_connection, "ConnectionConfig", FakeConnectionConfig) + + async def fake_create(**kwargs): + # Warmup must not reach the network; a fast failure leaves the slot + # empty and lets start() finish so we can inspect the config. + raise RuntimeError("no network in tests") + + monkeypatch.setattr(opensandbox.Sandbox, "create", staticmethod(fake_create)) + pool = _pool(provider={"opensandbox": {"connection": connection}}, size=1) + + async def main(): + await pool.start() + await pool.aclose() + + asyncio.run(main()) + return captured + + def test_proxy_mode_carries_the_api_key_as_a_header(self, monkeypatch): + captured = self._captured_kwargs( + monkeypatch, + { + "domain": "http://sandbox.example", + "api_key": "key", + "use_server_proxy": True, + }, # pragma: allowlist secret + ) + assert captured["headers"] == {"OPEN-SANDBOX-API-KEY": "key"} # pragma: allowlist secret + + def test_direct_mode_never_injects_the_key(self, monkeypatch): + # A direct sandbox endpoint runs untrusted code and must never see the key. + captured = self._captured_kwargs( + monkeypatch, + {"domain": "http://sandbox.example", "api_key": "key"}, # pragma: allowlist secret + ) + assert "headers" not in captured + + def test_caller_supplied_headers_survive_and_win(self, monkeypatch): + captured = self._captured_kwargs( + monkeypatch, + { + "domain": "http://sandbox.example", + "api_key": "key", # pragma: allowlist secret + "use_server_proxy": True, + "headers": {"X-Route": "r", "OPEN-SANDBOX-API-KEY": "explicit"}, # pragma: allowlist secret + }, + ) + assert captured["headers"] == { + "X-Route": "r", + "OPEN-SANDBOX-API-KEY": "explicit", # pragma: allowlist secret + } + + def test_normalize_endpoint_adds_the_key_only_in_proxy_mode(self): + from types import SimpleNamespace + + resolved = SimpleNamespace(endpoint="sandbox.example/v1/sandboxes/sbx/proxy/6000", headers={}) + + proxied = _pool() # PROVIDER sets use_server_proxy: True + _, headers = proxied._normalize_endpoint(resolved) + assert headers == {"OPEN-SANDBOX-API-KEY": "k"} + + # A direct endpoint terminates at the sandbox, which runs untrusted code. + direct = _pool( + provider={"opensandbox": {"connection": {"domain": "http://sandbox.example", "api_key": "k"}}}, + ) + _, headers = direct._normalize_endpoint(resolved) + assert headers == {} diff --git a/tests/unit_tests/test_opensandbox_endpoint.py b/tests/unit_tests/test_opensandbox_endpoint.py new file mode 100644 index 0000000000..02eefca102 --- /dev/null +++ b/tests/unit_tests/test_opensandbox_endpoint.py @@ -0,0 +1,145 @@ +# 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 + + +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 _RawWithEndpoint: + def __init__( + self, endpoint: str = "http://sandbox.example/v1/sandboxes/sbx-1/proxy/6000", headers: dict | None = None + ): + self._endpoint = endpoint + self._headers = {"OPEN-SANDBOX-API-KEY": "secret"} if headers is None else headers + 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"} + 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_is_absolutized_and_key_header_injected() -> None: + """The SDK returns proxy endpoints without a scheme and with empty headers (observed on + a production cluster, SDK 0.1.15); the provider must produce an absolute URL and carry the API key.""" + provider = opensandbox_provider.OpenSandboxProvider( + connection={ + "domain": "http://sandbox.example", + "api_key": "secret", # pragma: allowlist secret + "use_server_proxy": True, + }, + probe={"command": None}, + ) + raw = _RawWithEndpoint(endpoint="sandbox.example/v1/sandboxes/sbx-1/proxy/6000", headers={}) + resolved = asyncio.run(provider.endpoint(_handle(raw), 6000)) + assert resolved.endpoint == "http://sandbox.example/v1/sandboxes/sbx-1/proxy/6000" + assert resolved.headers == {"OPEN-SANDBOX-API-KEY": "secret"} # pragma: allowlist secret + + +def test_direct_mode_endpoint_never_carries_the_key() -> None: + """A direct endpoint terminates at the sandbox, which runs untrusted code; handing it + the API key would leak the key to the workload. Mirrors the provider's proxy-only rule.""" + provider = opensandbox_provider.OpenSandboxProvider( + connection={"domain": "http://sandbox.example", "api_key": "secret"}, # pragma: allowlist secret + probe={"command": None}, + ) + raw = _RawWithEndpoint(endpoint="http://pod.example:6000", headers={}) + resolved = asyncio.run(provider.endpoint(_handle(raw), 6000)) + assert resolved.headers == {} + + +def test_sdk_supplied_headers_are_never_overridden() -> None: + provider = opensandbox_provider.OpenSandboxProvider( + connection={"domain": "http://sandbox.example", "api_key": "secret"}, + probe={"command": None}, + ) + raw = _RawWithEndpoint(headers={"X-Route-Token": "t"}) + resolved = asyncio.run(provider.endpoint(_handle(raw), 6000)) + assert resolved.headers == {"X-Route-Token": "t"} + + +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()) From 6dc7a092d99ed56e6c2d5484761c29e359c4e1fd Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Wed, 12 Aug 2026 23:06:49 -0700 Subject: [PATCH 2/6] refactor: simplify disaggregated sandbox lifecycle Reuse the public sandbox API for pooled and Lean execution. Make resource ownership and cancellation cleanup explicit, validate backend configuration, and cover lifecycle regressions. Signed-off-by: Hemil Desai --- nemo_gym/sandbox/__init__.py | 3 +- nemo_gym/sandbox/api.py | 39 +- .../sandbox/providers/opensandbox/provider.py | 77 ++-- nemo_gym/sandbox/utils.py | 17 + resources_servers/math_formal_lean/app.py | 35 +- .../configs/math_formal_lean.yaml | 13 +- .../configs/math_formal_lean_multi_turn.yaml | 13 +- .../math_formal_lean/requirements.txt | 4 +- .../math_formal_lean/sandbox_client.py | 215 +++++----- .../math_formal_lean/tests/test_app.py | 13 + .../tests/test_sandbox_backends.py | 162 ++++++-- resources_servers/ns_tools/app.py | 111 +++--- .../ns_tools/configs/ns_tools.yaml | 8 +- resources_servers/ns_tools/gym_sandbox.py | 33 +- resources_servers/ns_tools/requirements.txt | 1 + resources_servers/ns_tools/sandbox_pool.py | 342 +++++++--------- .../tests/test_app_sandbox_pool_wiring.py | 116 ++++++ .../ns_tools/tests/test_sandbox_pool.py | 368 +++++++++--------- tests/unit_tests/test_opensandbox_endpoint.py | 84 ++-- tests/unit_tests/test_opensandbox_provider.py | 98 ++++- tests/unit_tests/test_sandbox.py | 69 ++++ 21 files changed, 1072 insertions(+), 749 deletions(-) create mode 100644 resources_servers/ns_tools/tests/test_app_sandbox_pool_wiring.py 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 86d4a0b690..ea18bf2bd6 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,14 @@ 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 +1002,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 +1104,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 @@ -1189,19 +1196,18 @@ async def endpoint(self, handle: SandboxHandle, port: int) -> SandboxEndpoint: endpoint_url = str(getattr(resolved, "endpoint", "") or "") if not endpoint_url: raise RuntimeError(f"OpenSandbox returned an empty endpoint for sandbox {handle.sandbox_id} port {port}") + connection = handle.raw.connection_config if "://" not in endpoint_url: - # The SDK returns the proxy endpoint without a scheme; borrow it from the - # configured server domain. - domain = str(self._connection.domain or "") - scheme = "https" if domain.startswith("https://") else "http" - endpoint_url = f"{scheme}://{endpoint_url}" + # Use the handle's SDK-resolved config so protocol and environment + # defaults match the connection that produced this endpoint. + scheme = connection.get_base_url().split("://", 1)[0] + endpoint_url = f"{scheme}://{endpoint_url.lstrip('/')}" headers = dict(getattr(resolved, "headers", None) or {}) - if not headers and self._connection.use_server_proxy and self._connection.api_key: - # Defensive: the SDK currently returns no headers for proxy endpoints; include - # the API key so authenticated proxy deployments work either way. Proxy mode - # only — a direct endpoint terminates at the sandbox, which runs untrusted - # code and must never be handed the key. - headers["OPEN-SANDBOX-API-KEY"] = str(self._connection.api_key) + api_key = connection.headers.get("OPEN-SANDBOX-API-KEY") + if connection.use_server_proxy and api_key: + # Direct endpoints terminate at untrusted sandbox code; only proxy + # endpoints may receive server credentials. + headers.setdefault("OPEN-SANDBOX-API-KEY", api_key) return SandboxEndpoint(endpoint=endpoint_url, headers=headers) def _command_retry_count(self) -> int: @@ -1535,7 +1541,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(), @@ -1543,7 +1549,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: @@ -1567,8 +1573,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 78de135847..f4edd96f7d 100644 --- a/resources_servers/math_formal_lean/app.py +++ b/resources_servers/math_formal_lean/app.py @@ -17,8 +17,9 @@ 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, Field @@ -341,12 +342,8 @@ def build_correction_prompt( class MathFormalLeanResourcesServerConfig(BaseResourcesServerConfig): sandbox_host: str = "127.0.0.1" sandbox_port: int = 6000 - # Sandbox backend: "ns_http" (default — today's NS server over host/port), "gym_sandbox" - # (per-verify OpenSandbox pods via provider exec), or "ns_http_proxy" (the NS HTTP protocol - # through a full base_url + headers, e.g. an OpenSandbox proxied endpoint; parity oracle). - sandbox_backend: str = "ns_http" - sandbox_base_url: str = "" - sandbox_extra_headers: Dict[str, str] = Field(default_factory=dict) + # 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) @@ -398,12 +395,6 @@ def model_post_init(self, context: Any) -> None: max_output_characters=self.config.max_output_characters, **self.config.opensandbox, ) - elif self.config.sandbox_backend == "ns_http_proxy": - self._sandbox_client = Lean4SandboxClient( - base_url=self.config.sandbox_base_url, - extra_headers=self.config.sandbox_extra_headers, - max_output_characters=self.config.max_output_characters, - ) else: self._sandbox_client = Lean4SandboxClient( host=self.config.sandbox_host, @@ -418,20 +409,18 @@ def model_post_init(self, context: Any) -> None: def setup_webserver(self): app = super().setup_webserver() - start_pool = getattr(self._sandbox_client, "start_pool", None) - if start_pool is None: - return app - from contextlib import asynccontextmanager - main_app_lifespan = app.router.lifespan_context @asynccontextmanager async def lifespan_wrapper(app): - # Warm lean pool pods from startup: a cold pod's first compile is ~15 min - # of nydus olean pulls, far beyond any verify's admission window. - start_pool() - async with main_app_lifespan(app) as maybe_state: - yield maybe_state + 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 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 3daf78102d..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,17 +4,15 @@ 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} - # Sandbox backend: 'ns_http' (default — today's path, byte-identical with zero env - # vars set), 'gym_sandbox' (per-verify OpenSandbox pods via provider exec), or - # 'ns_http_proxy' (NS protocol through a full base_url — parity oracle). + # 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,} + domain: ${oc.env:OPENSANDBOX_BASE_URL,""} + api_key: ${oc.env:OPENSANDBOX_API_KEY,""} use_server_proxy: true create: timeout_s: 90 @@ -23,10 +21,9 @@ math_formal_lean: # (survives proxy/LB stream caps; Gym PR 2296). operations: background_exec: true - image: ${oc.env:NS_SANDBOX_IMAGE,} + image: ${oc.env:NS_SANDBOX_IMAGE,""} max_concurrent: ${oc.decode:${oc.env:LEAN_SANDBOX_MAX_CONCURRENT,8}} - # 0 = a fresh pod per verify. N > 0 = N warm pods reused across verifies — - # a cold pod's first compile lazy-pulls ~5GB of oleans (~15 min); warm ~4s. + # 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 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 771ac9e88b..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,17 +4,15 @@ 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} - # Sandbox backend: 'ns_http' (default — today's path, byte-identical with zero env - # vars set), 'gym_sandbox' (per-verify OpenSandbox pods via provider exec), or - # 'ns_http_proxy' (NS protocol through a full base_url — parity oracle). + # 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,} + domain: ${oc.env:OPENSANDBOX_BASE_URL,""} + api_key: ${oc.env:OPENSANDBOX_API_KEY,""} use_server_proxy: true create: timeout_s: 90 @@ -23,10 +21,9 @@ math_formal_lean: # (survives proxy/LB stream caps; Gym PR 2296). operations: background_exec: true - image: ${oc.env:NS_SANDBOX_IMAGE,} + image: ${oc.env:NS_SANDBOX_IMAGE,""} max_concurrent: ${oc.decode:${oc.env:LEAN_SANDBOX_MAX_CONCURRENT,8}} - # 0 = a fresh pod per verify. N > 0 = N warm pods reused across verifies — - # a cold pod's first compile lazy-pulls ~5GB of oleans (~15 min); warm ~4s. + # 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 diff --git a/resources_servers/math_formal_lean/requirements.txt b/resources_servers/math_formal_lean/requirements.txt index d04abcbc14..d705047542 100644 --- a/resources_servers/math_formal_lean/requirements.txt +++ b/resources_servers/math_formal_lean/requirements.txt @@ -1,5 +1,5 @@ -e nemo-gym[dev] @ ../../ httpx>=0.27.0 -## OpenSandbox SDK — used only by the opt-in opensandbox/ns_http_proxy sandbox backends +## OpenSandbox SDK — used only by the opt-in gym_sandbox backend opensandbox==0.1.15 -tenacity +tenacity>=9.1.4 diff --git a/resources_servers/math_formal_lean/sandbox_client.py b/resources_servers/math_formal_lean/sandbox_client.py index 1f39f8da36..cfbda6fae3 100644 --- a/resources_servers/math_formal_lean/sandbox_client.py +++ b/resources_servers/math_formal_lean/sandbox_client.py @@ -20,13 +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__) @@ -39,8 +44,6 @@ def __init__( host: str = "127.0.0.1", port: int = 6000, max_output_characters: int = 1000, - base_url: str | None = None, - extra_headers: Dict[str, str] | None = None, ): """Initialize sandbox client. @@ -48,15 +51,10 @@ def __init__( host: Sandbox server hostname port: Sandbox server port max_output_characters: Maximum characters in output - base_url: Full base URL override (e.g. an OpenSandbox proxied endpoint); - when set, host/port are ignored. - extra_headers: Headers added to every request (e.g. proxy auth). """ self.host = host self.port = port self.max_output_characters = max_output_characters - self.base_url = base_url.rstrip("/") if base_url else None - self.extra_headers = dict(extra_headers or {}) self._client: httpx.AsyncClient | None = None async def _get_client(self) -> httpx.AsyncClient: @@ -75,8 +73,6 @@ async def close(self) -> None: def _get_execute_url(self) -> str: """Get the sandbox execute endpoint URL.""" - if self.base_url: - return f"{self.base_url}/execute" return f"http://{self.host}:{self.port}/execute" async def execute_lean4( @@ -107,7 +103,7 @@ async def execute_lean4( url=self._get_execute_url(), content=json.dumps(request_data), timeout=timeout + 5.0, # Add buffer for network overhead - headers={"Content-Type": "application/json", **self.extra_headers}, + headers={"Content-Type": "application/json"}, ) if response.status_code == 502: @@ -137,12 +133,11 @@ async def health_check(self, timeout: float = 5.0) -> bool: Returns: True if sandbox is healthy, False otherwise """ - base = self.base_url if self.base_url else f"http://{self.host}:{self.port}" - url = f"{base}/health" + url = f"http://{self.host}:{self.port}/health" client = await self._get_client() try: - response = await client.get(url=url, timeout=timeout, headers=self.extra_headers) + response = await client.get(url=url, timeout=timeout) return response.status_code == 200 except httpx.HTTPError: return False @@ -151,23 +146,9 @@ async def health_check(self, timeout: float = 5.0) -> bool: class GymSandboxLean4Client: """Lean4 compilation on per-verify OpenSandbox pods via provider exec. - Reimplements the NS server's lean4 invocation exactly (reference frozen at - nemo_skills local_sandbox_server.py:631-685 @ da85a881): the proof lands in - /lean4/my_project, `lake env --dir /lean4/my_project lean ` runs with an - in-sandbox `timeout -s KILL`, and the exit code maps to the same - process_status/stdout/stderr contract as `Lean4SandboxClient.execute_lean4`. - Long compiles never hold an HTTP connection open (background/short exec - requests), so proxy read-timeout ceilings do not apply. - - With ``pool_size=0`` every verify gets a fresh pod (created under a bounded - semaphore, destroyed in finally). With ``pool_size=N`` a warm pool of N pods is - built lazily and reused across verifies: a fresh pod's first `import Mathlib` - lazy-pulls ~5GB of olean files through nydus one FUSE fault at a time (measured - at ~900s), while a warmed pod compiles in ~4s — so pool pods bulk-prefetch - the olean tree once at prepare and then serve verifies back-to-back. A pod that - hits an infra error (incl. TTL expiry) is killed and replaced in place; the - verify retries once on a second pod before degrading. Infra failures degrade to - the client's existing error/timeout shapes, never raise into verify(). + 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__( @@ -186,7 +167,7 @@ def __init__( ): if not image: raise ValueError("sandbox_backend=gym_sandbox requires a non-empty image") - connection = (next(iter(provider.values()), {}) or {}).get("connection", {}) if provider else {} + 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 — " @@ -200,24 +181,17 @@ def __init__( self.max_output_characters = max_output_characters self._acquire_timeout_s = acquire_timeout_s self._semaphore_size = int(max_concurrent) - self._semaphore: Any = None # bound lazily to the serving event loop + 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: Any = None # asyncio.Queue of warm AsyncSandbox pods, filled lazily - self._pool_started = False - - def _get_semaphore(self): - import asyncio - - if self._semaphore is None: - self._semaphore = asyncio.Semaphore(self._semaphore_size) - return self._semaphore - - def _new_sandbox(self, files: Dict[str, str] | None = None, use_pool: bool = True): - from nemo_gym.sandbox.api import AsyncSandbox - from nemo_gym.sandbox.providers.base import SandboxSpec + 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. @@ -235,7 +209,7 @@ def _new_sandbox(self, files: Dict[str, str] | None = None, use_pool: bool = Tru ), ) - async def _start_sandbox(self, files: Dict[str, str] | None = None): + 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: @@ -249,7 +223,14 @@ async def _start_sandbox(self, files: Dict[str, str] | None = None): await sandbox.start() return sandbox - async def _create_pool_pod(self): + 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() @@ -259,6 +240,9 @@ async def _create_pool_pod(self): 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 @@ -266,33 +250,42 @@ async def _create_pool_pod(self): 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: - self._pool_queue() - - @property - def pool_ready_count(self) -> int: - return self._pool.qsize() if self._pool is not None else 0 - - def _pool_queue(self): - import asyncio - - if self._pool is None: - self._pool = asyncio.Queue() - for _ in range(self._pool_size): - asyncio.get_running_loop().create_task(self._fill_one()) - return self._pool - - async def _fill_one(self): - try: - self._pool.put_nowait(await self._create_pool_pod()) - except Exception as e: - LOG.error("lean pool pod create failed (capacity reduced until next heal): %s", e) + 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]: - import asyncio - import uuid - - pool = self._pool_queue() + 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) @@ -301,8 +294,6 @@ async def _execute_pooled(self, code: str, timeout: float) -> Dict[str, Any]: return {"process_status": "timeout", "stdout": "", "stderr": "Client timed out"} proof_name = f"proof_{uuid.uuid4().hex}.lean" try: - import tempfile - with tempfile.NamedTemporaryFile("w", suffix=".lean", delete=False) as fh: fh.write(code) try: @@ -316,44 +307,48 @@ async def _execute_pooled(self, code: str, timeout: float) -> Dict[str, Any]: f"rc=$?; rm -f {proof_name}; exit $rc" ) result = await sandbox.exec(command, timeout_s=timeout + 60) - except Exception as e: + 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, e) - try: - await sandbox.stop() - except Exception: - pass - asyncio.get_running_loop().create_task(self._fill_one()) + 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(e)} - pool.put_nowait(sandbox) + 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 "")[: self.max_output_characters] - stderr = (result.stderr or "")[: self.max_output_characters] + stdout = result.stdout or "" + stderr = result.stderr or "" if result.return_code == 0: - return {"process_status": "completed", "stdout": stdout, "stderr": stderr} - if result.return_code in (124, 137, -9): - return { - "process_status": "timeout", - "stdout": stdout, - "stderr": stderr + f"Execution timed out after {timeout} seconds\n", - } - return {"process_status": "failed", "stdout": stdout, "stderr": stderr} + 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.""" - import asyncio - import uuid - if self._pool_size > 0: return await self._execute_pooled(code, timeout) try: - await asyncio.wait_for(self._get_semaphore().acquire(), timeout=self._acquire_timeout_s) + 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"} @@ -374,9 +369,25 @@ async def execute_lean4(self, code: str, timeout: float = 30.0) -> Dict[str, Any LOG.error("OpenSandbox lean4 execution failed: %s", e) return {"process_status": "error", "stdout": "", "stderr": str(e)} finally: - self._get_semaphore().release() + self._semaphore.release() if sandbox is not None: - try: - await sandbox.stop() - except Exception as exc: - LOG.warning("lean sandbox teardown failed (TTL will reap): %s", exc) + 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 index 8a8b047219..1a38fe941c 100644 --- a/resources_servers/math_formal_lean/tests/test_sandbox_backends.py +++ b/resources_servers/math_formal_lean/tests/test_sandbox_backends.py @@ -12,22 +12,14 @@ # 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: the additive base_url/header extension must -leave the default byte-identical, and the OpenSandbox exec client must reproduce the NS -server's process_status contract exactly (reference: local_sandbox_server.py:631-685).""" +"""Backend tests for the Lean sandbox clients.""" import asyncio -import sys -from pathlib import Path from types import SimpleNamespace -import httpx import pytest - -sys.path.insert(0, str(Path(__file__).resolve().parents[3])) - -from resources_servers.math_formal_lean.sandbox_client import ( # noqa: E402 +from resources_servers.math_formal_lean.sandbox_client import ( GymSandboxLean4Client, Lean4SandboxClient, ) @@ -40,29 +32,11 @@ } -class TestHttpClientStaysByteIdentical: +class TestHttpClientDefaults: def test_default_url_is_unchanged(self): client = Lean4SandboxClient() assert client._get_execute_url() == "http://127.0.0.1:6000/execute" - def test_base_url_override_and_headers(self): - seen = {} - - def handler(request: httpx.Request) -> httpx.Response: - seen["url"] = str(request.url) - seen["auth"] = request.headers.get("open-sandbox-api-key") - return httpx.Response(200, json={"process_status": "completed", "stdout": "", "stderr": ""}) - - client = Lean4SandboxClient( - base_url="http://sandbox.example/v1/sandboxes/sbx/proxy/6000", - extra_headers={"OPEN-SANDBOX-API-KEY": "k"}, - ) - client._client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) - out = asyncio.run(client.execute_lean4("theorem t : True := trivial", timeout=5.0)) - assert out["process_status"] == "completed" - assert seen["url"] == "http://sandbox.example/v1/sandboxes/sbx/proxy/6000/execute" - assert seen["auth"] == "k" - class _FakeSandbox: """Stands in for AsyncSandbox; records lifecycle and returns a scripted exec result.""" @@ -100,12 +74,12 @@ async def stop(self): @pytest.fixture() def fake_sandbox(monkeypatch): - import nemo_gym.sandbox.api as api + 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(api, "AsyncSandbox", _FakeSandbox) + monkeypatch.setattr(sandbox_client, "AsyncSandbox", _FakeSandbox) return _FakeSandbox @@ -148,6 +122,15 @@ def test_timeout_rc_maps_to_the_ns_timeout_contract(self, fake_sandbox): 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)) @@ -159,7 +142,7 @@ 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._get_semaphore() + sem = client._semaphore await sem.acquire() # exhaust admission try: return await client.execute_lean4("x", timeout=5.0) @@ -170,6 +153,25 @@ async def 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): @@ -206,6 +208,102 @@ async def scenario(): 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.""" diff --git a/resources_servers/ns_tools/app.py b/resources_servers/ns_tools/app.py index 4eedaaa3af..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__) @@ -80,7 +84,7 @@ class NSToolsConfig(BaseResourcesServerConfig): # Sandbox backend: "local" (default — today's colocated server) or "sandbox_pool" # (disaggregated pods on OpenSandbox; requires the sandbox_pool block below). - sandbox_type: str = "local" + 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) @@ -95,11 +99,6 @@ class NSToolsConfig(BaseResourcesServerConfig): # The model receives a warning in stderr instead of restored state. disable_session_restore: bool = False - # Staged guard for tool calls arriving WITHOUT a gym session cookie (each such call - # silently mints a fresh sandbox session today). False = log + count (default, - # behavior-neutral); True = reject with 400 once the cookie path is proven end-to-end. - strict_session_cookie: bool = False - # ============================================================ # Run/Verify Request/Response Models @@ -143,8 +142,8 @@ class NSToolsResourcesServer(SimpleResourcesServer): _tool_name_map: Dict[str, str] = {} # Maps tool names to qualified names _python_tool_process: Optional[subprocess.Popen] = None _timing_by_session: Dict[str, list] = {} # session_id -> list of timing records - _missing_cookie_count: int = 0 _uses_python_tool_sidecar: bool = False + _sandbox_pool: Optional["SandboxPool"] = None def setup_webserver(self) -> FastAPI: app = super().setup_webserver() @@ -157,13 +156,10 @@ def setup_webserver(self) -> FastAPI: @asynccontextmanager async def lifespan_wrapper(app): - if self.config.sandbox_type == "sandbox_pool": - import gym_sandbox - - if gym_sandbox.CURRENT_POOL is not None: - # Budgeted, non-blocking warmup: kicks pod creation without gating server boot. - await gym_sandbox.CURRENT_POOL.start() 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: @@ -279,18 +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": self.config.sandbox_type, - "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": - import gym_sandbox # noqa: F401 — registers the backend with the nemo_skills registry + from gym_sandbox import GymSandbox + from nemo_skills.code_execution.sandbox import sandboxes + from sandbox_pool import SandboxPool - context["sandbox"]["pool"] = dict(self.config.sandbox_pool) + 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() @@ -338,19 +338,8 @@ async def execute_tool(self, tool_name: str, request: Request) -> PlainTextRespo # Get session ID for stateful execution session_id = request.session.get(SESSION_ID_KEY) if not session_id: - # A missing gym cookie means every call mints a fresh sandbox session — silent - # per-call state loss. Staged fix: count + log by default; reject only once the - # cookie path is proven end-to-end and strict_session_cookie is flipped. - self._missing_cookie_count += 1 - if self.config.strict_session_cookie: - return PlainTextResponse( - json.dumps({"error": "missing session cookie; stateful tools require a session"}), - status_code=400, - ) session_id = str(uuid.uuid4()) - logger.warning( - f"No session ID found (occurrence {self._missing_cookie_count}), using fallback: {session_id}" - ) + logger.warning(f"No session ID found, using fallback: {session_id}") if session_id not in self._timing_by_session: self._timing_by_session[session_id] = [] @@ -496,36 +485,32 @@ async def verify(self, request: Request, body: NSToolsVerifyRequest) -> NSToolsV async def shutdown(self): """Cleanup resources on server shutdown.""" - if self.tool_manager: - try: + try: + if self.tool_manager: await self.tool_manager.shutdown() - except (asyncio.CancelledError, Exception): - # Tool-manager teardown must not skip pool teardown: unkilled pool - # pods stay allocated (and billed against pool capacity) until TTL. - logger.warning("tool_manager.shutdown failed; continuing to pool teardown", exc_info=True) - - if self.config.sandbox_type == "sandbox_pool": - import gym_sandbox - - if gym_sandbox.CURRENT_POOL is not None: - await gym_sandbox.CURRENT_POOL.aclose() - - # 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() + 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 23a7819f14..1f578009f0 100644 --- a/resources_servers/ns_tools/configs/ns_tools.yaml +++ b/resources_servers/ns_tools/configs/ns_tools.yaml @@ -53,9 +53,8 @@ ns_tools: domain: ${oc.env:OPENSANDBOX_BASE_URL,} api_key: ${oc.env:OPENSANDBOX_API_KEY,} use_server_proxy: true - # First-ever create on a new image tag blocks on nydus conversion (40-55s); - # the SDK default request_timeout (30s) fails the whole cold warmup wave. - request_timeout: 180 + # 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 @@ -65,6 +64,9 @@ ns_tools: # 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} diff --git a/resources_servers/ns_tools/gym_sandbox.py b/resources_servers/ns_tools/gym_sandbox.py index f62b14e994..8eb855af76 100644 --- a/resources_servers/ns_tools/gym_sandbox.py +++ b/resources_servers/ns_tools/gym_sandbox.py @@ -12,19 +12,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""NeMo-Skills sandbox backend that routes through an OpenSandbox pod pool. +"""NeMo-Skills backend that routes sessions through a shared sandbox pool. -Registers ``sandbox_type: sandbox_pool`` with the nemo_skills sandbox registry. The -class IS a ``LocalSandbox`` — same request preparation, same session bookkeeping — with the -transport re-pointed: each request resolves (base_url, headers) from the pool by session -uuid, and rides a shared AIOHTTP session (httpx/httpcore's O(n^2) connection pooling -collapses at high concurrency — see CLAUDE.md; measured: health-only GETs fell -from 87 to 8 calls/s between 64 and 512 in-flight on httpx). Exception TYPES stay httpx -because the nemo_skills base class's execute_code catches those; anything non-200 or -transport-level is normalized to the NS timeout contract so infra failures degrade rewards -without new error shapes. - -Importing this module is the opt-in: the default ``local`` backend never imports it. +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 @@ -40,21 +32,14 @@ class IS a ``LocalSandbox`` — same request preparation, same session bookkeepi LOGGER = logging.getLogger(__name__) -# The pool the owning server can warm up at lifespan startup (set by the first construction). -CURRENT_POOL: Optional[SandboxPool] = None - class GymSandbox(ns_sandbox.LocalSandbox): """LocalSandbox with the transport routed through an OpenSandbox pod pool over aiohttp.""" - def __init__(self, pool: Optional[Dict[str, Any]] = None, **kwargs: Any) -> None: + def __init__(self, pool: SandboxPool, **kwargs: Any) -> None: super().__init__(**kwargs) - if not pool: - raise ValueError("sandbox_type=sandbox_pool requires a 'pool' config dict") - global CURRENT_POOL - self._pool = SandboxPool(**pool) + self._pool = pool self._aiohttp: Optional[aiohttp.ClientSession] = None - CURRENT_POOL = self._pool def _session(self) -> aiohttp.ClientSession: if self._aiohttp is None or self._aiohttp.closed: @@ -130,11 +115,7 @@ async def delete_session(self, session_id: str) -> None: async def close(self) -> None: try: - await self._pool.aclose() - finally: if self._aiohttp is not None and not self._aiohttp.closed: await self._aiohttp.close() + finally: await super().close() - - -ns_sandbox.sandboxes["sandbox_pool"] = GymSandbox diff --git a/resources_servers/ns_tools/requirements.txt b/resources_servers/ns_tools/requirements.txt index e6812c33c7..9c237657d7 100644 --- a/resources_servers/ns_tools/requirements.txt +++ b/resources_servers/ns_tools/requirements.txt @@ -3,3 +3,4 @@ nemo-skills-tools @ git+https://github.com/NVIDIA-NeMo/Skills.git@da85a881d972e6fec847b90cf553a0bf9bf10638#subdirectory=tools ## OpenSandbox SDK — used only by the opt-in sandbox_pool sandbox backend opensandbox==0.1.15 +tenacity>=9.1.4 diff --git a/resources_servers/ns_tools/sandbox_pool.py b/resources_servers/ns_tools/sandbox_pool.py index 59045fdb3b..797dd805e1 100644 --- a/resources_servers/ns_tools/sandbox_pool.py +++ b/resources_servers/ns_tools/sandbox_pool.py @@ -15,10 +15,9 @@ """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: 16k concurrent sessions ride K pods -(each pod's NS server multiplexes many sessions), instead of 16k pods. Slots fill -and heal with direct async ``Sandbox.create`` calls — a full fan-out warm wave is -~5s per pod measured in production, so no warm-spare inventory layer is needed. +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. @@ -28,46 +27,28 @@ import logging import time from dataclasses import dataclass, field -from datetime import timedelta -from typing import Any, Dict, Optional, Tuple +from typing import Any import aiohttp import httpx # exception types only: the nemo_skills client contract catches httpx errors -from nemo_gym.sandbox.attribution import RUN_KEY, resolve_attribution, resolve_run_id -from nemo_gym.sandbox.providers.opensandbox.provider import ( - DEFAULT_ATTRIBUTION_KEY_PREFIX as _ATTRIBUTION_KEY_PREFIX, -) +from nemo_gym.sandbox import AsyncSandbox, SandboxSpec, await_cleanup LOGGER = logging.getLogger(__name__) -def _parse_connection(provider: Dict[str, Any]) -> Dict[str, Any]: - """Pull the connection kwargs out of a single-key provider config dict.""" - if not isinstance(provider, dict) or len(provider) != 1: - raise ValueError("sandbox_pool.provider must be a single-key provider config dict") - kwargs = next(iter(provider.values())) or {} - connection = dict(kwargs.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" - ) - return connection - - @dataclass class _Slot: index: int - sandbox: Any = None # opensandbox.Sandbox + sandbox: AsyncSandbox | None = None base_url: str = "" - headers: Dict[str, str] = field(default_factory=dict) + headers: dict[str, str] = field(default_factory=dict) healthy: bool = False strikes: int = 0 creating: bool = False heal_failures: int = 0 - sessions: set = field(default_factory=set) + sessions: set[str] = field(default_factory=set) class SandboxPool: @@ -81,20 +62,20 @@ class SandboxPool: def __init__( self, *, - provider: Dict[str, Any], + provider: dict[str, Any], image: str, pool_ref: str = "", pool_fallback: bool = True, port: int = 6000, size: int = 8, - ttl_s: Optional[float] = None, - env: Optional[Dict[str, str]] = None, - entrypoint: Optional[list] = None, - resources: Optional[Dict[str, str]] = None, - resource_requests: Optional[Dict[str, str]] = None, - setup_files: Optional[Dict[str, str]] = None, - setup_commands: Optional[list] = None, - service_command: Optional[str] = None, + 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, @@ -105,7 +86,18 @@ def __init__( heal_concurrency: int = 16, session_idle_sweep_s: float = 7200.0, ) -> None: - self._connection_kwargs = _parse_connection(provider) + 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 = ( @@ -119,48 +111,44 @@ def __init__( self._port = int(port) self._size = int(size) self._ttl_s = float(ttl_s) if ttl_s else None - # Hydra/YAML overrides deliver bare numbers as ints; the create API's env map - # is string->string and the server 422s on anything else. - self._env = {k: str(v) for k, v in (env or {}).items()} + self._env = dict(env or {}) self._entrypoint = list(entrypoint) if entrypoint else None - self._resources = {k: str(v) for k, v in (resources or {}).items()} - # k8s schedules on REQUESTS; keeping them far below limits packs many more pods - # (= sessions) per node while bursts still get the limit headroom. - self._resource_requests = {k: str(v) for k, v in (resource_requests or {}).items()} + 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._connection_config: Any = None 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 - attribution = resolve_attribution() - # No explicit label: NEMO_GYM_RUN_ID (set per job by the launch script) wins, - # else a per-process id — either way unique per run, so an epilogue reaper can - # delete exactly this run's sandboxes by the run attribution label. - attribution[RUN_KEY] = resolve_run_id() - self._metadata = {f"{_ATTRIBUTION_KEY_PREFIX}{k}": v for k, v in attribution.items()} - self._metadata["purpose"] = "ns-tools-sandbox-pool" + 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._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._tasks: list = [] + 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: Any = None # aiohttp session; created lazily on the serving loop + self._http: aiohttp.ClientSession | None = None # ------------------------------------------------------------------ lifecycle @@ -177,151 +165,102 @@ async def start(self) -> None: if self._started or self._closed: return self._started = True - from opensandbox.config.connection import ConnectionConfig - - kwargs = dict(self._connection_kwargs) - # Mirror of the provider's proxy-mode auth (PR 2462): the SDK's - # execd-facing clients (the create ready gate's health ping) send only - # ConnectionConfig.headers, so on servers that enforce auth on - # /proxy/* routes every ping 401s and the claim dies at ready_timeout. - # Inject the key only in proxy mode — a direct sandbox endpoint runs - # untrusted code and must never see it. - if kwargs.get("use_server_proxy") and kwargs.get("api_key"): - headers = dict(kwargs.get("headers") or {}) - headers.setdefault("OPEN-SANDBOX-API-KEY", str(kwargs["api_key"])) - kwargs["headers"] = headers - self._connection_config = ConnectionConfig(**kwargs) 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: - self._closed = True - for task in self._tasks: - task.cancel() - for task in self._tasks: - try: - await task - except (asyncio.CancelledError, Exception): - pass - self._tasks.clear() + 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() - async def _kill(slot: _Slot) -> None: - try: - await asyncio.wait_for(slot.sandbox.kill(), timeout=30.0) - except Exception as exc: - LOGGER.warning("pool slot %d teardown failed (TTL will reap): %s", slot.index, exc) - slot.sandbox = None - slot.healthy = False + self._close_task = asyncio.create_task(cleanup()) + await await_cleanup(self._close_task) - await asyncio.gather(*(_kill(slot) for slot in self._slots if slot.sandbox is not None)) - if self._http is not None and not self._http.closed: - await self._http.close() + 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 - @property - def _needs_prepare(self) -> bool: - return bool(self._setup_files or self._setup_commands or self._service_command) - - async def _prepare(self, sandbox: Any) -> None: - """Bootstrap a pod. INVARIANT: the service start is the LAST execd command this pod - ever sees — execd reaps a completed command's backgrounded children when any later - command runs (probed empirically: service->touch->10s = dead listener; service-> - nothing = alive).""" - for target_path, local_path in self._setup_files.items(): - with open(local_path, "rb") as fh: - await sandbox.files.write_file(target_path, fh.read()) - for command in self._setup_commands: - execution = await sandbox.commands.run(command) - if (execution.exit_code or 0) != 0: - raise RuntimeError(f"setup command failed rc={execution.exit_code}: {command!r}") - if self._service_command: - # Plain shell backgrounding (cmd &): setsid and execd background:true both - # freeze the child during module import (probed); a plain & child reparents - # to the pod's PID 1 and survives — as long as nothing execs afterwards. - execution = await sandbox.commands.run(self._service_command) - if (execution.exit_code or 0) != 0: - raise RuntimeError(f"service command failed rc={execution.exit_code}") - - def _normalize_endpoint(self, resolved: Any) -> Tuple[str, Dict[str, str]]: - url = str(getattr(resolved, "endpoint", "") or "") - if not url: - raise RuntimeError("SDK returned an empty sandbox endpoint") - if "://" not in url: - domain = str(self._connection_kwargs.get("domain") or "") - scheme = "https" if domain.startswith("https://") else "http" - url = f"{scheme}://{url}" - headers = dict(getattr(resolved, "headers", None) or {}) - if not headers and self._connection_kwargs.get("use_server_proxy") and self._connection_kwargs.get("api_key"): - # Proxy mode only — a direct endpoint terminates at the sandbox, which runs - # untrusted code and must never be handed the key. - headers["OPEN-SANDBOX-API-KEY"] = str(self._connection_kwargs["api_key"]) - return url.rstrip("/"), headers - - async def _create_slot(self, slot: _Slot) -> None: - """Fill one slot. Single-flight per slot: a duplicate landing late would - overwrite base_url under pinned sessions and leak a pod until TTL.""" - if slot.creating: - return - slot.creating = True - try: - await self._create_slot_inner(slot) - finally: - slot.creating = False - - async def _acquire_sandbox(self) -> Tuple[Any, bool]: + 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).""" - from opensandbox import Sandbox - if self._pool_ref: + sandbox = AsyncSandbox(self._provider) try: - sandbox = await Sandbox.create( - # The SDK's local validation requires an image even in pool mode; - # the pool template still defines what actually runs. - image=self._image, - extensions={"poolRef": self._pool_ref}, - metadata=dict(self._metadata), - timeout=timedelta(seconds=self._ttl_s or 14400.0), - ready_timeout=timedelta(seconds=self._ready_timeout_s), - connection_config=self._connection_config, + 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 = await Sandbox.create( - image=self._image, - entrypoint=self._entrypoint, - env=self._env or None, - metadata=dict(self._metadata), - resource=self._resources or None, - resource_requests=self._resource_requests or None, - timeout=timedelta(seconds=self._ttl_s or 14400.0), - ready_timeout=timedelta(seconds=self._ready_timeout_s), - connection_config=self._connection_config, + 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: - # Pool pods are born with their service running; only direct creates - # (no pool, or fallback) need bootstrap. - if self._needs_prepare and not from_pool: - await self._prepare(sandbox) - resolved = await sandbox.get_endpoint(self._port) - base_url, headers = self._normalize_endpoint(resolved) + 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 Exception: - try: - await sandbox.kill() - except Exception: - pass + except BaseException: + await self._stop_sandbox(sandbox, slot.index) raise slot.sandbox = sandbox slot.base_url = base_url @@ -329,10 +268,10 @@ async def _create_slot_inner(self, slot: _Slot) -> None: slot.strikes = 0 slot.healthy = True - async def _wait_healthy(self, base_url: str, headers: Dict[str, str], budget_s: float) -> None: + 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: Optional[str] = None + 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: @@ -349,10 +288,13 @@ async def _warmup(self) -> None: async def one(slot: _Slot) -> None: async with semaphore: + slot.creating = True try: - await self._create_slot(slot) + 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) @@ -364,17 +306,13 @@ async def one(slot: _Slot) -> None: async def _heal_loop(self) -> None: while not self._closed: await asyncio.sleep(self._health_interval_s) - if not self._warmup_done: - # Warmup owns every slot until it finishes; healing in parallel would - # race duplicate acquisitions into the same slot. - continue 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 True + return self._warmup_done try: async with self._http_session().get( f"{slot.base_url}{self._health_path}", headers=slot.headers @@ -391,10 +329,7 @@ async def check(slot: _Slot) -> bool: slot.healthy = False await self._drop_slot_sessions(slot) if slot.sandbox is not None: - try: - await slot.sandbox.kill() - except Exception: - pass + await self._stop_sandbox(slot.sandbox, slot.index) slot.sandbox = None return True return False @@ -403,27 +338,32 @@ async def check(slot: _Slot) -> bool: to_heal = [slot for slot, needed in zip(self._slots, needs_heal) if needed] if not to_heal: continue - semaphore = asyncio.Semaphore(self._heal_concurrency) - - async def heal(slot: _Slot) -> None: - async with semaphore: - await self._heal_slot(slot) - - await asyncio.gather(*(heal(slot) for slot in to_heal)) + 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.""" - 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) + if slot.creating: + return + slot.creating = True try: - await self._create_slot(slot) + 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: @@ -440,6 +380,8 @@ async def _heal_slot(self, slot: _Slot) -> None: slot.heal_failures, exc, ) + finally: + slot.creating = False async def _drop_slot_sessions(self, slot: _Slot) -> None: async with self._lock: @@ -464,7 +406,7 @@ async def _sweep_loop(self) -> None: # ------------------------------------------------------------------ routing - async def route(self, session_id: Optional[str]) -> Tuple[str, Dict[str, str]]: + 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 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 index 48e3cd7c94..e03407664b 100644 --- a/resources_servers/ns_tools/tests/test_sandbox_pool.py +++ b/resources_servers/ns_tools/tests/test_sandbox_pool.py @@ -18,10 +18,13 @@ 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)) @@ -37,7 +40,7 @@ def _pool(**overrides) -> SandboxPool: - kwargs = dict(provider=PROVIDER, image="img", size=2) + kwargs = dict(provider=PROVIDER, image="img", size=2, entrypoint=["/start-with-nginx.sh"]) kwargs.update(overrides) return SandboxPool(**kwargs) @@ -54,16 +57,20 @@ 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") + 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") + 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="") + 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). @@ -170,6 +177,13 @@ 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).""" @@ -180,7 +194,7 @@ def backend(self): import gym_sandbox sandbox = gym_sandbox.GymSandbox( - pool=dict(provider=PROVIDER, image="img", size=1), + pool=_pool(size=1), host="127.0.0.1", port="6000", disable_session_restore=True, @@ -188,13 +202,6 @@ def backend(self): _admit(sandbox._pool, 0) return sandbox - def test_backend_registers_with_the_nemo_skills_registry(self): - pytest.importorskip("nemo_skills") - import gym_sandbox - from nemo_skills.code_execution.sandbox import sandboxes - - assert sandboxes["sandbox_pool"] is gym_sandbox.GymSandbox - def test_send_request_routes_with_pool_headers_and_session(self, backend): ok = '{"process_status": "completed", "stdout": "", "stderr": ""}' backend._aiohttp = _FakeAiohttpSession([_FakeAiohttpResponse(200, ok)]) @@ -245,224 +252,197 @@ async def main(): assert "sess-a" not in backend._pool._session_to_slot -class TestHealWarmupRace: - """The heal loop must never race duplicate creates into a slot warmup still owns — - a late duplicate overwrites base_url under pinned sessions and breaks stickiness - (observed under slow pod creation).""" - - def test_create_slot_is_single_flight(self): - pool = _pool() - calls = {"n": 0} - - async def fake_inner(slot): - calls["n"] += 1 - await asyncio.sleep(0.05) +class _FakeSandbox: + instances = [] + fail_claim = False - pool._create_slot_inner = fake_inner + def __init__(self, provider): + self.provider = provider + self.spec = None + self.stops = 0 + self.uploads = [] + self.commands = [] + self.instances.append(self) - async def main(): - await asyncio.gather(pool._create_slot(pool._slots[0]), pool._create_slot(pool._slots[0])) - - asyncio.run(main()) - assert calls["n"] == 1, "second concurrent create for the same slot must be a no-op" - - def test_heal_loop_waits_for_warmup(self): - pool = _pool() - assert pool._warmup_done is False - healed = [] - pool._heal_slot = lambda slot: healed.append(slot.index) - - async def one_heal_pass(): - pool._health_interval_s = 0.01 - task = asyncio.create_task(pool._heal_loop()) - await asyncio.sleep(0.05) - pool._closed = True - task.cancel() - try: - await task - except asyncio.CancelledError: - pass + async def start(self, spec): + self.spec = spec + if self.fail_claim and spec.provider_options.get("extensions"): + raise RuntimeError("pool exhausted") + return self - asyncio.run(one_heal_pass()) - assert healed == [], "heal loop must not touch slots before warmup completes" + async def stop(self): + self.stops += 1 + async def endpoint(self, port): + return SandboxEndpoint(endpoint=f"https://sandbox.example/{port}", headers={"X-Route": "r"}) -class TestEnvStringify: - def test_env_values_are_stringified(self): - from sandbox_pool import SandboxPool + async def upload(self, local_path, remote_path): + self.uploads.append((local_path, remote_path)) - pool = SandboxPool( - provider={"opensandbox": {"connection": {"domain": "http://sandbox.example", "api_key": "k"}}}, - image="img:tag", - env={"NUM_WORKERS": 4, "FLAG": True}, - ) - # The create API's env map is string->string; ints/bools 422 server-side. - assert pool._env == {"NUM_WORKERS": "4", "FLAG": "True"} + async def exec(self, command): + self.commands.append(command) + return SimpleNamespace(return_code=0) -class TestPoolRefFallback: - """pool_ref acquire semantics — SDK required (create is monkeypatched, no network).""" +class TestPoolSandboxApi: + @pytest.fixture(autouse=True) + def fake_sandbox(self, monkeypatch): + import sandbox_pool - def _pool(self, **overrides): - from sandbox_pool import SandboxPool + _FakeSandbox.instances = [] + _FakeSandbox.fail_claim = False + monkeypatch.setattr(sandbox_pool, "AsyncSandbox", _FakeSandbox) - kwargs = dict( - provider={"opensandbox": {"connection": {"domain": "http://sandbox.example", "api_key": "k"}}}, - image="img:tag", - pool_ref="warm-pool", + def test_pool_ref_and_direct_fallback_use_sandbox_specs(self): + pool = _pool( size=1, - service_command="sh -c 'start & echo ok'", + pool_ref="warm-pool", + env={"NUM_WORKERS": 4}, + resources={"cpu": 2, "memory_mib": 4096}, + resource_requests={"cpu": 0.5, "memory_mib": 1024}, ) - kwargs.update(overrides) - return SandboxPool(**kwargs) - - def test_pool_full_falls_back_to_direct_create(self, monkeypatch): - opensandbox = pytest.importorskip("opensandbox") - calls = [] - - async def fake_create(**kwargs): - calls.append(kwargs) - if "extensions" in kwargs: - raise RuntimeError("pool exhausted") - return object() - - monkeypatch.setattr(opensandbox.Sandbox, "create", staticmethod(fake_create)) - pool = self._pool() - pool._connection_config = object() - sandbox, from_pool = asyncio.run(pool._acquire_sandbox()) - assert from_pool is False and sandbox is not None - assert calls[0]["extensions"] == {"poolRef": "warm-pool"} - # The fallback create carries the full direct spec, not the pool claim shape. - assert "extensions" not in calls[1] and calls[1]["image"] == "img:tag" - def test_pool_failure_raises_when_fallback_disabled(self, monkeypatch): - opensandbox = pytest.importorskip("opensandbox") - - async def fake_create(**kwargs): - raise RuntimeError("pool exhausted") + 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 - monkeypatch.setattr(opensandbox.Sandbox, "create", staticmethod(fake_create)) - pool = self._pool(pool_fallback=False) - pool._connection_config = object() + _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_pool_claim_skips_prepare_and_fallback_does_not(self, monkeypatch): - opensandbox = pytest.importorskip("opensandbox") - prepared = [] - - async def fake_create(**kwargs): - if "extensions" in kwargs and fake_create.pool_ok: - return "pool-pod" - if "extensions" in kwargs: - raise RuntimeError("pool exhausted") - return "direct-pod" - - monkeypatch.setattr(opensandbox.Sandbox, "create", staticmethod(fake_create)) - pool = self._pool() - pool._connection_config = object() + 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 fake_prepare(sandbox): - prepared.append(sandbox) + 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 fake_endpoint_flow(slot, sandbox): - slot.sandbox = sandbox - slot.healthy = True + async def blocked_create(slot): + heal_started.set() + await asyncio.Future() - pool._prepare = fake_prepare + pool._create_slot_inner = blocked_create - async def run_inner(pool_ok): - fake_create.pool_ok = pool_ok - sandbox, from_pool = await pool._acquire_sandbox() - if pool._needs_prepare and not from_pool: - await pool._prepare(sandbox) - return sandbox + 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() - assert asyncio.run(run_inner(True)) == "pool-pod" - assert prepared == [] - assert asyncio.run(run_inner(False)) == "direct-pod" - assert prepared == ["direct-pod"] + 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 = [] -class TestProxyAuthHeaders: - """Mirror of the provider's proxy-mode auth (PR 2462). The SDK's execd-facing - clients authenticate only via ConnectionConfig.headers, so without the key - there the create ready gate's health ping 401s on servers that enforce auth - on /proxy/* routes and every claim dies at ready_timeout.""" + async def fail_fast(slot): + attempted.append(slot.index) - def _captured_kwargs(self, monkeypatch, connection) -> dict: - opensandbox = pytest.importorskip("opensandbox") - import opensandbox.config.connection as osb_connection + pool._heal_slot = fail_fast - captured: dict = {} + 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() - class FakeConnectionConfig: - def __init__(self, **kwargs): - captured.update(kwargs) + asyncio.run(main()) + assert set(attempted) == {0, 1, 2} - monkeypatch.setattr(osb_connection, "ConnectionConfig", FakeConnectionConfig) + def test_cancelled_admission_stops_the_sandbox(self): + pool = _pool(size=1) + waiting = asyncio.Event() - async def fake_create(**kwargs): - # Warmup must not reach the network; a fast failure leaves the slot - # empty and lets start() finish so we can inspect the config. - raise RuntimeError("no network in tests") + async def wait_forever(*args, **kwargs): + waiting.set() + await asyncio.Future() - monkeypatch.setattr(opensandbox.Sandbox, "create", staticmethod(fake_create)) - pool = _pool(provider={"opensandbox": {"connection": connection}}, size=1) + pool._wait_healthy = wait_forever async def main(): - await pool.start() - await pool.aclose() + 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()) - return captured - - def test_proxy_mode_carries_the_api_key_as_a_header(self, monkeypatch): - captured = self._captured_kwargs( - monkeypatch, - { - "domain": "http://sandbox.example", - "api_key": "key", - "use_server_proxy": True, - }, # pragma: allowlist secret - ) - assert captured["headers"] == {"OPEN-SANDBOX-API-KEY": "key"} # pragma: allowlist secret + assert _FakeSandbox.instances[0].stops == 1 + assert pool._slots[0].sandbox is None - def test_direct_mode_never_injects_the_key(self, monkeypatch): - # A direct sandbox endpoint runs untrusted code and must never see the key. - captured = self._captured_kwargs( - monkeypatch, - {"domain": "http://sandbox.example", "api_key": "key"}, # pragma: allowlist secret - ) - assert "headers" not in captured - - def test_caller_supplied_headers_survive_and_win(self, monkeypatch): - captured = self._captured_kwargs( - monkeypatch, - { - "domain": "http://sandbox.example", - "api_key": "key", # pragma: allowlist secret - "use_server_proxy": True, - "headers": {"X-Route": "r", "OPEN-SANDBOX-API-KEY": "explicit"}, # pragma: allowlist secret - }, - ) - assert captured["headers"] == { - "X-Route": "r", - "OPEN-SANDBOX-API-KEY": "explicit", # pragma: allowlist secret - } + def test_cancelled_aclose_finishes_cleanup(self): + class BlockingSandbox(_FakeSandbox): + stop_started = asyncio.Event() + finish_stop = asyncio.Event() - def test_normalize_endpoint_adds_the_key_only_in_proxy_mode(self): - from types import SimpleNamespace + async def stop(self): + self.stop_started.set() + await self.finish_stop.wait() + await super().stop() - resolved = SimpleNamespace(endpoint="sandbox.example/v1/sandboxes/sbx/proxy/6000", headers={}) + pool = _pool(size=1) + sandbox = BlockingSandbox(PROVIDER) + pool._slots[0].sandbox = sandbox + pool._http = _FakeAiohttpSession([]) - proxied = _pool() # PROVIDER sets use_server_proxy: True - _, headers = proxied._normalize_endpoint(resolved) - assert headers == {"OPEN-SANDBOX-API-KEY": "k"} + 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 - # A direct endpoint terminates at the sandbox, which runs untrusted code. - direct = _pool( - provider={"opensandbox": {"connection": {"domain": "http://sandbox.example", "api_key": "k"}}}, - ) - _, headers = direct._normalize_endpoint(resolved) - assert headers == {} + 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 index 02eefca102..444bff71d1 100644 --- a/tests/unit_tests/test_opensandbox_endpoint.py +++ b/tests/unit_tests/test_opensandbox_endpoint.py @@ -36,12 +36,26 @@ 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 | None = None + 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 + 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: @@ -58,7 +72,7 @@ def test_endpoint_returns_the_sdk_url_and_auth_headers() -> None: 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"} + assert resolved.headers == {"OPEN-SANDBOX-API-KEY": "secret"} # pragma: allowlist secret assert raw.requested_ports == [6000] @@ -80,43 +94,51 @@ def test_endpoint_headers_default_to_empty_dict_without_configured_key() -> None assert resolved.headers == {} -def test_schemeless_endpoint_is_absolutized_and_key_header_injected() -> None: - """The SDK returns proxy endpoints without a scheme and with empty headers (observed on - a production cluster, SDK 0.1.15); the provider must produce an absolute URL and carry the API key.""" - provider = opensandbox_provider.OpenSandboxProvider( - connection={ - "domain": "http://sandbox.example", - "api_key": "secret", # pragma: allowlist secret - "use_server_proxy": True, - }, - probe={"command": None}, +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, + ), ) - raw = _RawWithEndpoint(endpoint="sandbox.example/v1/sandboxes/sbx-1/proxy/6000", headers={}) - resolved = asyncio.run(provider.endpoint(_handle(raw), 6000)) - assert resolved.endpoint == "http://sandbox.example/v1/sandboxes/sbx-1/proxy/6000" + 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_never_carries_the_key() -> None: - """A direct endpoint terminates at the sandbox, which runs untrusted code; handing it - the API key would leak the key to the workload. Mirrors the provider's proxy-only rule.""" - provider = opensandbox_provider.OpenSandboxProvider( - connection={"domain": "http://sandbox.example", "api_key": "secret"}, # pragma: allowlist secret - probe={"command": None}, +def test_direct_mode_endpoint_does_not_inject_the_key() -> None: + raw = _RawWithEndpoint( + endpoint="http://pod.example:6000", + headers={}, + connection=_Connection(api_key="secret"), # pragma: allowlist secret ) - raw = _RawWithEndpoint(endpoint="http://pod.example:6000", headers={}) - resolved = asyncio.run(provider.endpoint(_handle(raw), 6000)) + resolved = asyncio.run(_provider().endpoint(_handle(raw), 6000)) assert resolved.headers == {} -def test_sdk_supplied_headers_are_never_overridden() -> None: - provider = opensandbox_provider.OpenSandboxProvider( - connection={"domain": "http://sandbox.example", "api_key": "secret"}, - probe={"command": None}, +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 ) - raw = _RawWithEndpoint(headers={"X-Route-Token": "t"}) - resolved = asyncio.run(provider.endpoint(_handle(raw), 6000)) - assert resolved.headers == {"X-Route-Token": "t"} + 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: diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index 8c93efa3ec..a4f30fb1b3 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,8 @@ 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 +542,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 +1285,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()) From 933e78fcd15f1fce1a29d024f0095d14d09b5266 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Wed, 12 Aug 2026 23:13:45 -0700 Subject: [PATCH 3/6] test: classify OpenSandbox endpoint coverage Signed-off-by: Hemil Desai --- tests/unit_tests/test_opensandbox_endpoint.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit_tests/test_opensandbox_endpoint.py b/tests/unit_tests/test_opensandbox_endpoint.py index 444bff71d1..e34cca679f 100644 --- a/tests/unit_tests/test_opensandbox_endpoint.py +++ b/tests/unit_tests/test_opensandbox_endpoint.py @@ -28,6 +28,9 @@ 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}) From d167ffe5d2a8305ac5f731f6abe85e8abc06f05a Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Thu, 13 Aug 2026 09:39:04 -0700 Subject: [PATCH 4/6] fix: avoid replaying stateful sandbox requests Signed-off-by: Hemil Desai --- resources_servers/ns_tools/gym_sandbox.py | 4 ---- .../ns_tools/tests/test_sandbox_pool.py | 13 +++++-------- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/resources_servers/ns_tools/gym_sandbox.py b/resources_servers/ns_tools/gym_sandbox.py index 8eb855af76..577df772f5 100644 --- a/resources_servers/ns_tools/gym_sandbox.py +++ b/resources_servers/ns_tools/gym_sandbox.py @@ -74,10 +74,6 @@ async def _send_request(self, request: Dict[str, Any], timeout: float): try: status, text = await self._post_execute(base_url, headers, payload, timeout) - if status == 502: - # A proxy-minted 502 means the pod never received the request, so ONE retry - # is idempotency-safe even for stateful ipython. - 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: diff --git a/resources_servers/ns_tools/tests/test_sandbox_pool.py b/resources_servers/ns_tools/tests/test_sandbox_pool.py index e03407664b..abab648b5e 100644 --- a/resources_servers/ns_tools/tests/test_sandbox_pool.py +++ b/resources_servers/ns_tools/tests/test_sandbox_pool.py @@ -217,14 +217,11 @@ def test_non_200_normalizes_to_the_timeout_contract(self, backend): with pytest.raises(httpx.TimeoutException): asyncio.run(backend._send_request({"generated_code": "1+1", "session_id": "s"}, timeout=10.0)) - def test_502_retries_exactly_once_then_succeeds(self, backend): - ok = '{"process_status": "completed", "stdout": "", "stderr": ""}' - backend._aiohttp = _FakeAiohttpSession( - [_FakeAiohttpResponse(502, "bad gateway"), _FakeAiohttpResponse(200, ok)] - ) - result = asyncio.run(backend._send_request({"generated_code": "1+1", "session_id": "s"}, timeout=10.0)) - assert result["process_status"] == "completed" - assert len(backend._aiohttp.calls) == 2 + 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 From ccbc7fba296bab81692a7ddd5ceac2527a04afe3 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Thu, 13 Aug 2026 09:52:16 -0700 Subject: [PATCH 5/6] fix: preserve OpenSandbox connection compatibility Signed-off-by: Hemil Desai --- .../sandbox/providers/opensandbox/provider.py | 40 +------------------ tests/unit_tests/test_opensandbox_endpoint.py | 2 +- tests/unit_tests/test_opensandbox_provider.py | 1 + 3 files changed, 4 insertions(+), 39 deletions(-) diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index ea18bf2bd6..dbd562b0ed 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -662,6 +662,8 @@ def _connection_config( kwargs["request_timeout"] = timedelta(seconds=request_timeout_s) if self._connection.use_server_proxy: kwargs["use_server_proxy"] = True + if self._connection.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() config = ConnectionConfig(**kwargs) @@ -1172,44 +1174,6 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: raw_status = getattr(info, "status", None) return _to_sandbox_status(getattr(raw_status, "state", None) if raw_status is not None else None) - async def endpoint(self, handle: SandboxHandle, port: int) -> SandboxEndpoint: - """Resolve an HTTP(S) endpoint for a declared sandbox port. - - The SDK returns the server-proxy route (`{domain}/v1/sandboxes/{id}/proxy/{port}`) - when the connection uses the server proxy, along with any headers the server - requires on every request to that endpoint (e.g. the API key header). - """ - 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: 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(getattr(resolved, "endpoint", "") or "") - if not endpoint_url: - raise RuntimeError(f"OpenSandbox returned an empty endpoint for sandbox {handle.sandbox_id} port {port}") - connection = handle.raw.connection_config - if "://" not in endpoint_url: - # Use the handle's SDK-resolved config so protocol and environment - # defaults match the connection that produced this endpoint. - scheme = connection.get_base_url().split("://", 1)[0] - endpoint_url = f"{scheme}://{endpoint_url.lstrip('/')}" - headers = dict(getattr(resolved, "headers", None) or {}) - api_key = connection.headers.get("OPEN-SANDBOX-API-KEY") - if connection.use_server_proxy and api_key: - # Direct endpoints terminate at untrusted sandbox code; only proxy - # endpoints may receive server credentials. - headers.setdefault("OPEN-SANDBOX-API-KEY", api_key) - return SandboxEndpoint(endpoint=endpoint_url, headers=headers) - def _command_retry_count(self) -> int: return self._operations.command_retries diff --git a/tests/unit_tests/test_opensandbox_endpoint.py b/tests/unit_tests/test_opensandbox_endpoint.py index e34cca679f..4a7e0bb9b2 100644 --- a/tests/unit_tests/test_opensandbox_endpoint.py +++ b/tests/unit_tests/test_opensandbox_endpoint.py @@ -116,7 +116,7 @@ def test_schemeless_endpoint_uses_the_sdk_resolved_url_and_key() -> None: def test_direct_mode_endpoint_does_not_inject_the_key() -> None: raw = _RawWithEndpoint( endpoint="http://pod.example:6000", - headers={}, + 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)) diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py index a4f30fb1b3..ed4b3db867 100644 --- a/tests/unit_tests/test_opensandbox_provider.py +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -532,6 +532,7 @@ def test_connection_config_and_image_policy( "protocol": "https", "request_timeout": timedelta(seconds=10), "use_server_proxy": True, + "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) From 5186de5e696197207df161be3e6daa8faa62a038 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Fri, 21 Aug 2026 14:14:28 -0700 Subject: [PATCH 6/6] build: install sandbox dependencies from Gym extra Signed-off-by: Hemil Desai --- resources_servers/math_formal_lean/requirements.txt | 5 +---- resources_servers/ns_tools/requirements.txt | 5 +---- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/resources_servers/math_formal_lean/requirements.txt b/resources_servers/math_formal_lean/requirements.txt index d705047542..0b42bad7a7 100644 --- a/resources_servers/math_formal_lean/requirements.txt +++ b/resources_servers/math_formal_lean/requirements.txt @@ -1,5 +1,2 @@ --e nemo-gym[dev] @ ../../ +-e nemo-gym[dev,sandbox] @ ../../ httpx>=0.27.0 -## OpenSandbox SDK — used only by the opt-in gym_sandbox backend -opensandbox==0.1.15 -tenacity>=9.1.4 diff --git a/resources_servers/ns_tools/requirements.txt b/resources_servers/ns_tools/requirements.txt index 9c237657d7..6dde21233a 100644 --- a/resources_servers/ns_tools/requirements.txt +++ b/resources_servers/ns_tools/requirements.txt @@ -1,6 +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 -## OpenSandbox SDK — used only by the opt-in sandbox_pool sandbox backend -opensandbox==0.1.15 -tenacity>=9.1.4