feat(sandbox): sandbox owner base for resources servers, with SWE-bench on it

nemo_gym/base_resources_server.py

@@ -429,7 +429,13 @@ def setup_resources_checkpoint(self, app: FastAPI) -> None:
         settings = checkpoint_settings(getattr(self.server_client, "global_config_dict", None))
         if settings is None:
             return
-        if self.checkpoint_mode == "exported":
+        mode = self.checkpoint_mode_in_effect()
+        if mode not in get_args(ResourcesCheckpointMode):
+            raise ValueError(
+                f"{type(self).__name__}.checkpoint_mode_in_effect() returned {mode!r}; "
+                f"expected one of {get_args(ResourcesCheckpointMode)}"
+            )
+        if mode == "exported":
             # Hooks from an intermediate base class or a mixin count; only the stubs here do not.
             missing = [
                 name for name in _SESSION_HOOKS if getattr(type(self), name) is getattr(SimpleResourcesServer, name)
@@ -438,7 +444,7 @@ def setup_resources_checkpoint(self, app: FastAPI) -> None:
                 raise ValueError(
                     f"{type(self).__name__} declares checkpoint_mode='exported' but does not implement {missing}"
                 )
-        participant = ResourcesParticipant(self, self.checkpoint_mode, verify_mode=self.checkpoint_verify)
+        participant = ResourcesParticipant(self, mode, verify_mode=self.checkpoint_verify)
         self._checkpoint = participant
         install_server_participant(
             app,
@@ -472,6 +478,14 @@ def setup_checkpoint_participant(self, app: FastAPI) -> None:
                 "build the app with SimpleResourcesServer.setup_webserver, or call setup_resources_checkpoint on it"
             )
 
+    def checkpoint_mode_in_effect(self) -> ResourcesCheckpointMode:
+        """The mode this instance takes part in checkpoints with; read once at startup, with checkpointing on.
+
+        The default is the declared ``checkpoint_mode``. A server whose ability to capture its sessions depends on
+        its configuration, such as a sandbox provider that cannot snapshot, returns a mode that captures less.
+        """
+        return self.checkpoint_mode
+
     @property
     def checkpoint_participant(self) -> Optional[ResourcesParticipant]:
         """The participant that takes part in partial-rollout checkpoints, or ``None`` with checkpointing off."""

nemo_gym/sandbox/checkpoint.py

@@ -82,6 +82,16 @@
 _UNRESTORED_PROVIDER_OPTIONS = ("image_auth", "snapshot_id")
 
 
+class SandboxCheckpointConfig(BaseModel):
+    """How a server checkpoints the sandboxes its sessions own: the ``sandbox_checkpoint`` block of its config."""
+
+    # Concurrent snapshot, re-create, and stop calls against the sandbox backend during a commit or restore.
+    parallelism: int = 16
+    # What happens to the agents' sandboxes when the controller stops after a commit: "kill" frees them (the
+    # restore forks from the snapshots), "none" leaves them to their TTL.
+    on_stop: OnStop = "kill"
+
+
 class SandboxCheckpointState(BaseModel):
     """What a restore needs to rebuild one session's sandbox as of a checkpoint: its snapshot and its spec.
 

resources_servers/swebench/app.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import sys
+from functools import partial
 from glob import glob
 from pathlib import Path
 from shutil import rmtree
@@ -28,17 +29,16 @@
 
 from docker.models.containers import ExecResult
 from nemo_gym.base_resources_server import (
-    BaseResourcesServerConfig,
     BaseSeedSessionRequest,
     BaseSeedSessionResponse,
     BaseVerifyRequest,
     BaseVerifyResponse,
-    SimpleResourcesServer,
 )
 from nemo_gym.global_config import get_global_config_dict
 from nemo_gym.rollout_observability import SandboxObservation
 from nemo_gym.sandbox import AsyncSandbox, SandboxResources, SandboxSpec
-from nemo_gym.sandbox.config import resolve_provider_config, resolve_provider_metadata
+from nemo_gym.sandbox.config import resolve_provider_metadata
+from nemo_gym.sandbox.owner import SandboxOwnerConfig, SandboxOwnerResourcesServer
 from nemo_gym.sandbox.utils import cpu_cap_env
 from nemo_gym.server_utils import SESSION_ID_KEY
 from resources_servers.swebench.swebench_patches import (
@@ -50,7 +50,7 @@
 )
 
 
-class SwebenchResourcesServerConfig(BaseResourcesServerConfig):
+class SwebenchResourcesServerConfig(SandboxOwnerConfig):
     is_verifying_golden_patch: bool = False
     apply_anti_cheating: bool = True
     # Write the Maven Central mirror to ~/.m2/settings.xml in the eval sandbox. Turn off for providers
@@ -59,8 +59,7 @@ class SwebenchResourcesServerConfig(BaseResourcesServerConfig):
 
     evaluation_timeout: Optional[int] = None
 
-    # Sandbox config
-    sandbox_provider: str
+    # Sandbox config; sandbox_provider and sandbox_checkpoint come from SandboxOwnerConfig.
     sandbox_config: Dict[str, Any]
 
     clear_swebench_debug_logs: bool = True
@@ -241,18 +240,21 @@ class SWEBenchSeedSessionResponse(BaseSeedSessionResponse):
     sandbox_handle: str  # @bxyu-nvidia: Just a plain string URI for now for OpenSandbox backend.
 
 
-class SwebenchResourcesServer(SimpleResourcesServer):
-    config: SwebenchResourcesServerConfig
+class SwebenchResourcesServer(SandboxOwnerResourcesServer):
+    """Each session owns the sandbox the agent works in, from /seed_session until /verify extracts the patch."""
 
-    def model_post_init(self, context: Any, /) -> None:
-        super().model_post_init(context)
+    config: SwebenchResourcesServerConfig
 
-        self._session_id_to_sandbox: Dict[str, AsyncSandbox] = dict()
+    def lends_sandboxes(self) -> bool:
+        # A golden-patch instance grades the dataset's own patch and never seeds: it lends no sandbox.
+        return not self.config.is_verifying_golden_patch
 
     async def _create_sandbox(self, test_spec: TestSpec) -> AsyncSandbox:
+        return await self._start_sandbox(self._sandbox_spec(test_spec), test_spec)
+
+    def _sandbox_spec(self, test_spec: TestSpec) -> SandboxSpec:
         # TODO @bxyu-nvidia: Refactor this after Hemil's swap from Python dataclass to Pydantic BaseModel
         global_config_dict = get_global_config_dict()
-        resolved_sandbox_provider = resolve_provider_config(self.config.sandbox_provider, global_config_dict)
         provider_default_metadata = resolve_provider_metadata(self.config.sandbox_provider, global_config_dict)
         resources = dict(self.config.sandbox_config.get("resources", {}))
 
@@ -282,20 +284,23 @@ async def _create_sandbox(self, test_spec: TestSpec) -> AsyncSandbox:
             entrypoint=None,
             provider_options=self.config.sandbox_config.get("provider_options", {}),
         )
-        eval_sandbox = AsyncSandbox(resolved_sandbox_provider)
-
-        async def _run_setup(sandbox: AsyncSandbox) -> None:
-            await patch_swebench_multilingual_sandbox(
-                test_spec.repo,
-                test_spec.instance_id,
-                sandbox,
-                overwrite_maven_settings=self.config.overwrite_maven_settings,
-            )
+        return eval_sandbox_spec
 
-        await eval_sandbox.start_with_setup(eval_sandbox_spec, _run_setup)
+    async def _start_sandbox(self, spec: SandboxSpec, test_spec: TestSpec) -> AsyncSandbox:
+        eval_sandbox = AsyncSandbox(self.shared_provider(), spec, owns_provider=False)
+
+        await eval_sandbox.start_with_setup(spec, partial(self._patch_multilingual, test_spec))
 
         return eval_sandbox
 
+    async def _patch_multilingual(self, test_spec: TestSpec, sandbox: AsyncSandbox) -> None:
+        await patch_swebench_multilingual_sandbox(
+            test_spec.repo,
+            test_spec.instance_id,
+            sandbox,
+            overwrite_maven_settings=self.config.overwrite_maven_settings,
+        )
+
     def _make_test_spec(self, body: SWEBenchVerifyRequest) -> TestSpec:
         return make_test_spec(
             # This accepts a SWEbenchInstance which is identically our body.
@@ -306,27 +311,39 @@ def _make_test_spec(self, body: SWEBenchVerifyRequest) -> TestSpec:
         )
 
     async def seed_session(self, request: Request, body: SWEBenchSeedSessionRequest) -> SWEBenchSeedSessionResponse:
-        test_spec = self._make_test_spec(body)
-        eval_sandbox = await self._create_sandbox(test_spec)
-        self._session_id_to_sandbox[request.session[SESSION_ID_KEY]] = eval_sandbox
-
-        if self.config.apply_anti_cheating:
-            # Remove the current Git repo's future history beyond the current commit to prevent the model from cheating.
-            wd = (await eval_sandbox.exec("pwd")).stdout.strip()
-            anti_cheat_setup_fpath = Path(__file__).parent / "anti_cheat_setup.sh"
-            await eval_sandbox.upload(anti_cheat_setup_fpath, f"{wd}/anti_cheat_setup.sh")
-            result = await eval_sandbox.exec(
-                f"""git reset --hard && WORKING_DIRECTORY={wd} bash anti_cheat_setup.sh && rm anti_cheat_setup.sh"""
+        session_id = request.session[SESSION_ID_KEY]
+        sandbox = self.session_sandbox(session_id)
+        if sandbox is None:
+            test_spec = self._make_test_spec(body)
+            sandbox = await self.start_session_sandbox(
+                session_id, self._sandbox_spec(test_spec), setup=partial(self._setup_agent_sandbox, test_spec)
             )
-            if result.return_code != 0:
-                print(f"""Failed to setup anti-cheating for {test_spec.instance_id}. Return code: {result.return_code}
+        # Otherwise a replacement attempt re-seeds the session a checkpoint restored: it continues in the restored
+        # sandbox, the fork of the checkpoint's snapshot, whose setup already ran.
+        return SWEBenchSeedSessionResponse(sandbox_handle=sandbox.handle.sandbox_id)
+
+    async def _setup_agent_sandbox(self, test_spec: TestSpec, sandbox: AsyncSandbox) -> None:
+        """The agent's sandbox before the agent starts: the multilingual patch, then the anti-cheat scrub.
+
+        Verify diffs the working tree in the sandbox, so the session needs no data beyond the sandbox itself.
+        """
+        await self._patch_multilingual(test_spec, sandbox)
+        if not self.config.apply_anti_cheating:
+            return
+        wd = (await sandbox.exec("pwd")).stdout.strip()
+        # Remove the current Git repo's future history beyond the current commit to prevent the model from cheating.
+        anti_cheat_setup_fpath = Path(__file__).parent / "anti_cheat_setup.sh"
+        await sandbox.upload(anti_cheat_setup_fpath, f"{wd}/anti_cheat_setup.sh")
+        result = await sandbox.exec(
+            f"""git reset --hard && WORKING_DIRECTORY={wd} bash anti_cheat_setup.sh && rm anti_cheat_setup.sh"""
+        )
+        if result.return_code != 0:
+            print(f"""Failed to setup anti-cheating for {test_spec.instance_id}. Return code: {result.return_code}
 Stdout:
 {result.stdout}
 Stderr:
 {result.stderr}""")
 
-        return SWEBenchSeedSessionResponse(sandbox_handle=eval_sandbox._handle.sandbox_id)
-
     async def verify(self, request: Request, body: SWEBenchVerifyRequest) -> SWEBenchVerifyResponse:
         """
         Key requirements:
@@ -351,6 +368,12 @@ async def verify(self, request: Request, body: SWEBenchVerifyRequest) -> SWEBenc
 
         test_spec = self._make_test_spec(body)
 
+        session_id = request.session[SESSION_ID_KEY]
+        original_sandbox = None
+        if not self.config.is_verifying_golden_patch:
+            # 409 for a session without a sandbox: an empty patch would grade as reward 0.
+            original_sandbox = self.require_session_sandbox(session_id)
+
         verifier_sandbox_lifecycle_started_at = monotonic()
         eval_sandbox = await self._create_sandbox(test_spec)
         eval_sandbox_start_time_taken = monotonic() - verifier_sandbox_lifecycle_started_at
@@ -359,17 +382,14 @@ async def verify(self, request: Request, body: SWEBenchVerifyRequest) -> SWEBenc
         if self.config.is_verifying_golden_patch:
             model_patch = body.patch
         else:
-            original_sandbox = self._session_id_to_sandbox.pop(request.session[SESSION_ID_KEY])
             try:
                 original_workdir = (await eval_sandbox.exec("pwd")).stdout.strip()
                 model_patch_result = await original_sandbox.exec(f"cd {original_workdir} && git --no-pager diff")
                 model_patch = model_patch_result.stdout
-            except:
+            except Exception:
                 print("Failed to extract patch from container", format_exc(), file=sys.stderr)
-            try:
-                await original_sandbox.stop()
-            except:
-                print("Failed to stop original sandbox", format_exc(), file=sys.stderr)
+            # The episode is over; its snapshots stay until a release drops the checkpoints that name them.
+            await self.end_session_sandbox(session_id)
 
         run_id = request.session[SESSION_ID_KEY]
         mock_container = DockerContainer(id=run_id, instance_id=test_spec.instance_id)

resources_servers/swebench/configs/swebench.yaml

@@ -16,6 +16,12 @@ swebench_resources_server:          # instance name — how agents/CLI refer to
       evaluation_timeout: 1800
 
       sandbox_provider: sandbox
+      # Partial-rollout checkpoints snapshot each session's sandbox at commit and fork it from the snapshot at
+      # restore; parallelism bounds the concurrent backend calls (size it with the commit deadline). on_stop is
+      # what happens to the sandboxes when the controller stops after a commit: kill (default) or none.
+      sandbox_checkpoint:
+        parallelism: 16
+        on_stop: kill
       sandbox_config:
         ttl_s: 18000
         ready_timeout_s: 1200

resources_servers/swebench/tests/test_app.py

@@ -125,7 +125,8 @@ async def test_create_sandbox_derives_cpu_cap_env_from_cpu_limit(self, monkeypat
         sandbox = MagicMock()
         sandbox.start_with_setup = AsyncMock(side_effect=lambda spec, setup: sandbox)
         monkeypatch.setattr("resources_servers.swebench.app.get_global_config_dict", lambda: {})
-        monkeypatch.setattr("resources_servers.swebench.app.resolve_provider_config", lambda *_: MagicMock())
+        monkeypatch.setattr("nemo_gym.sandbox.owner.get_global_config_dict", lambda: {})
+        monkeypatch.setattr("nemo_gym.sandbox.owner.resolve_provider_config", lambda *_: MagicMock())
         monkeypatch.setattr("resources_servers.swebench.app.resolve_provider_metadata", lambda *_: {})
         monkeypatch.setattr("resources_servers.swebench.app.AsyncSandbox", MagicMock(return_value=sandbox))
         monkeypatch.setattr("resources_servers.swebench.app.patch_swebench_multilingual_sandbox", AsyncMock())