feat: disaggregated sandbox backends on OpenSandbox (ns_tools + math_formal_lean)

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",

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.
 

nemo_gym/sandbox/providers/opensandbox/provider.py

@@ -39,6 +39,7 @@
     SandboxStatus,
 )
 from nemo_gym.sandbox.providers.utils import coerce_config as _coerce_config
+from nemo_gym.sandbox.utils import await_cleanup
 
 
 LOGGER = logging.getLogger(__name__)
@@ -661,17 +662,16 @@ def _connection_config(
             kwargs["request_timeout"] = timedelta(seconds=request_timeout_s)
         if self._connection.use_server_proxy:
             kwargs["use_server_proxy"] = True
-            # The SDK's execd-facing clients (health ping, commands, files)
-            # send only ConnectionConfig.headers — api_key alone never reaches
-            # proxied /proxy/* routes, so servers that enforce auth there 401
-            # every health ping and create times out at ready_timeout. Inject
-            # the key only in proxy mode: a direct sandbox endpoint runs
-            # untrusted code and must never see it.
             if self._connection.api_key is not None:
                 kwargs["headers"] = {"OPEN-SANDBOX-API-KEY": self._connection.api_key}
         if self._connection.keepalive_expiry_s is not None or self._connection.disable_connection_pooling:
             kwargs["transport"] = self._get_transport()
-        return ConnectionConfig(**kwargs)
+        config = ConnectionConfig(**kwargs)
+        if self._connection.use_server_proxy and (api_key := config.get_api_key()):
+            # Execd-facing SDK clients send only ConnectionConfig.headers.
+            # Direct endpoints run untrusted code and must never see this key.
+            config.headers.setdefault("OPEN-SANDBOX-API-KEY", api_key)
+        return config
 
     def _get_transport(self) -> Any:
         """Return the provider-owned shared transport, building it on first use."""
@@ -1004,29 +1004,35 @@ async def endpoint(
     ) -> SandboxEndpoint:
         """Resolve one client-reachable direct or server-proxied service URL."""
 
+        get_endpoint = getattr(handle.raw, "get_endpoint", None)
+        if get_endpoint is None:
+            raise NotImplementedError(
+                "The installed opensandbox SDK does not expose Sandbox.get_endpoint; "
+                "sandbox service endpoints require opensandbox>=0.1.15"
+            )
         resolved = await self._await_sdk_operation(
-            lambda: handle.raw.get_endpoint(port),
+            lambda: get_endpoint(port),
             operation="get_endpoint",
             sandbox_id=handle.sandbox_id,
             timeout_s=(
                 float(self._connection.request_timeout_s) if self._connection.request_timeout_s is not None else None
             ),
         )
-        endpoint_url = str(resolved.endpoint or "")
+        endpoint_url = str(getattr(resolved, "endpoint", "") or "")
         if not endpoint_url:
             raise RuntimeError(f"OpenSandbox returned an empty endpoint for sandbox {handle.sandbox_id!r} port {port}")
+        connection = handle.raw.connection_config
         if "://" not in endpoint_url:
             # Use the SDK handle's effective configuration so environment-
             # resolved domains and protocols match the lifecycle request.
-            scheme = urlsplit(handle.raw.connection_config.get_base_url()).scheme or "http"
+            scheme = urlsplit(connection.get_base_url()).scheme or "http"
             endpoint_url = f"{scheme}://{endpoint_url.lstrip('/')}"
-        headers = dict(handle.raw.connection_config.headers)
-        # Match the SDK's service adapters: connection-wide headers apply to
-        # every request, while endpoint-specific routing or auth headers win.
-        # The upstream proxy-auth fix adds the management API key to
-        # ConnectionConfig.headers only in server-proxy mode, so direct
-        # sandbox endpoints never receive it.
-        headers.update(resolved.headers)
+        headers = dict(getattr(connection, "headers", None) or {})
+        headers.update(getattr(resolved, "headers", None) or {})
+        if not getattr(connection, "use_server_proxy", self._connection.use_server_proxy):
+            # Direct endpoints terminate at untrusted sandbox code and must
+            # never receive the management credential.
+            headers.pop("OPEN-SANDBOX-API-KEY", None)
         return SandboxEndpoint(endpoint=endpoint_url, headers=headers)
 
     async def _create_once(self, spec: SandboxSpec) -> SandboxHandle:
@@ -1100,8 +1106,11 @@ async def _create_once(self, spec: SandboxSpec) -> SandboxHandle:
             if self._create.skip_health_check:
                 handle = await self._connect_after_create(created_handle, spec)
             await self._verify_created_handle(handle)
-        except Exception:
-            await self._cleanup_failed_create_handle(created_handle)
+        except BaseException:
+            # Once create returns an id, cancellation must not strand its
+            # remote sandbox. close() applies the configured cleanup bounds.
+            cleanup = asyncio.create_task(self._cleanup_failed_create_handle(created_handle))
+            await await_cleanup(cleanup)
             raise
         return handle
 
@@ -1496,15 +1505,15 @@ 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(),
                 operation="kill",
                 sandbox_id=handle.sandbox_id,
                 timeout_s=self._operations.close_timeout_s,
             )
-        except Exception as e:
+        except BaseException as e:
             if not _is_missing_sandbox_delete_error(e):
                 stop_error = e
             else:
@@ -1528,8 +1537,9 @@ async def close(self, handle: SandboxHandle) -> None:
                 handle.sandbox_id,
                 e,
             )
-
         if stop_error is not None:
+            if not isinstance(stop_error, Exception):
+                raise stop_error
             if close_error is not None:
                 raise RuntimeError(
                     "Failed to stop and close OpenSandbox sandbox "

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."""

resources_servers/math_formal_lean/app.py

@@ -17,10 +17,11 @@
 
 import logging
 import re
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Literal, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from nemo_gym.base_resources_server import (
     BaseResourcesServerConfig,
@@ -29,7 +30,7 @@
     BaseVerifyResponse,
     SimpleResourcesServer,
 )
-from resources_servers.math_formal_lean.sandbox_client import Lean4SandboxClient
+from resources_servers.math_formal_lean.sandbox_client import GymSandboxLean4Client, Lean4SandboxClient
 
 
 LOG = logging.getLogger(__name__)
@@ -341,6 +342,11 @@ def build_correction_prompt(
 class MathFormalLeanResourcesServerConfig(BaseResourcesServerConfig):
     sandbox_host: str = "127.0.0.1"
     sandbox_port: int = 6000
+    # Sandbox backend: local NS HTTP (default) or provider-backed Gym sandboxes.
+    sandbox_backend: Literal["ns_http", "gym_sandbox"] = "ns_http"
+    # GymSandboxLean4Client kwargs (provider/image/max_concurrent/...) — read only when
+    # sandbox_backend == "gym_sandbox".
+    opensandbox: Dict[str, Any] = Field(default_factory=dict)
     compilation_timeout: float = 30.0
     max_output_characters: int = 1000
     extract_code_mode: str = "last"
@@ -384,17 +390,41 @@ class MathFormalLeanResourcesServer(SimpleResourcesServer):
 
     def model_post_init(self, context: Any) -> None:
         super().model_post_init(context)
-        self._sandbox_client = Lean4SandboxClient(
-            host=self.config.sandbox_host,
-            port=self.config.sandbox_port,
-            max_output_characters=self.config.max_output_characters,
-        )
+        if self.config.sandbox_backend == "gym_sandbox":
+            self._sandbox_client = GymSandboxLean4Client(
+                max_output_characters=self.config.max_output_characters,
+                **self.config.opensandbox,
+            )
+        else:
+            self._sandbox_client = Lean4SandboxClient(
+                host=self.config.sandbox_host,
+                port=self.config.sandbox_port,
+                max_output_characters=self.config.max_output_characters,
+            )
         self._proof_build_config = ProofBuildConfig(
             extract_code_mode=self.config.extract_code_mode,
             restate_formal_statement=self.config.restate_formal_statement,
             strip_theorem_from_proof=self.config.strip_theorem_from_proof,
         )
 
+    def setup_webserver(self):
+        app = super().setup_webserver()
+        main_app_lifespan = app.router.lifespan_context
+
+        @asynccontextmanager
+        async def lifespan_wrapper(app):
+            if isinstance(self._sandbox_client, GymSandboxLean4Client):
+                # A cold pod's first compile can exceed a verify's admission window.
+                self._sandbox_client.start_pool()
+            try:
+                async with main_app_lifespan(app) as maybe_state:
+                    yield maybe_state
+            finally:
+                await self._sandbox_client.close()
+
+        app.router.lifespan_context = lifespan_wrapper
+        return app
+
     async def verify(self, body: MathFormalLeanVerifyRequest) -> MathFormalLeanVerifyResponse:
         """Verify a proof attempt with multi-turn self-correction support.
 

resources_servers/math_formal_lean/configs/math_formal_lean.yaml

@@ -4,6 +4,27 @@ math_formal_lean:
       entrypoint: app.py
       sandbox_host: ${oc.env:NEMO_SKILLS_SANDBOX_HOST,127.0.0.1}
       sandbox_port: ${oc.env:NEMO_SKILLS_SANDBOX_PORT,6000}
+      # Default local NS HTTP or opt-in provider-backed Gym sandboxes.
+      sandbox_backend: ${oc.env:MATH_FORMAL_LEAN_BACKEND,ns_http}
+      # Read only when sandbox_backend == gym_sandbox; empty creds/image then = hard startup error.
+      opensandbox:
+        provider:
+          opensandbox:
+            connection:
+              domain: ${oc.env:OPENSANDBOX_BASE_URL,""}
+              api_key: ${oc.env:OPENSANDBOX_API_KEY,""}
+              use_server_proxy: true
+            create:
+              timeout_s: 90
+              retries: 3
+            # Long compiles poll short status requests instead of holding one stream
+            # (survives proxy/LB stream caps; Gym PR 2296).
+            operations:
+              background_exec: true
+        image: ${oc.env:NS_SANDBOX_IMAGE,""}
+        max_concurrent: ${oc.decode:${oc.env:LEAN_SANDBOX_MAX_CONCURRENT,8}}
+        # 0 creates a pod per verify; N reuses N warmed pods.
+        pool_size: ${oc.env:LEAN_SANDBOX_POOL_SIZE,0}
       compilation_timeout: 30.0
       domain: math
       verified: false

resources_servers/math_formal_lean/configs/math_formal_lean_multi_turn.yaml

@@ -4,6 +4,27 @@ math_formal_lean:
       entrypoint: app.py
       sandbox_host: ${oc.env:NEMO_SKILLS_SANDBOX_HOST,127.0.0.1}
       sandbox_port: ${oc.env:NEMO_SKILLS_SANDBOX_PORT,6000}
+      # Default local NS HTTP or opt-in provider-backed Gym sandboxes.
+      sandbox_backend: ${oc.env:MATH_FORMAL_LEAN_BACKEND,ns_http}
+      # Read only when sandbox_backend == gym_sandbox; empty creds/image then = hard startup error.
+      opensandbox:
+        provider:
+          opensandbox:
+            connection:
+              domain: ${oc.env:OPENSANDBOX_BASE_URL,""}
+              api_key: ${oc.env:OPENSANDBOX_API_KEY,""}
+              use_server_proxy: true
+            create:
+              timeout_s: 90
+              retries: 3
+            # Long compiles poll short status requests instead of holding one stream
+            # (survives proxy/LB stream caps; Gym PR 2296).
+            operations:
+              background_exec: true
+        image: ${oc.env:NS_SANDBOX_IMAGE,""}
+        max_concurrent: ${oc.decode:${oc.env:LEAN_SANDBOX_MAX_CONCURRENT,8}}
+        # 0 creates a pod per verify; N reuses N warmed pods.
+        pool_size: ${oc.env:LEAN_SANDBOX_POOL_SIZE,0}
       compilation_timeout: 30.0
       domain: math
       verified: false

resources_servers/math_formal_lean/requirements.txt

@@ -1,2 +1,2 @@
--e nemo-gym[dev] @ ../../
+-e nemo-gym[dev,sandbox] @ ../../
 httpx>=0.27.0