feat(sandbox): add suspend and resume lifecycle (#4765)

* feat(sandbox): add suspend and resume lifecycle Signed-off-by: Drew Newberry <anewberry@nvidia.com> * fix(sandbox): preserve V8 state across failed resume launches Signed-off-by: Drew Newberry <anewberry@nvidia.com> * fix(sandbox): handle suspend and resume recovery races Signed-off-by: Drew Newberry <anewberry@nvidia.com> * refactor(example): move V8 suspend support to a dependent change Signed-off-by: Drew Newberry <anewberry@nvidia.com> * fix(sandbox): align suspend tests with current supervisor APIs Signed-off-by: Drew Newberry <anewberry@nvidia.com> * fix(sandbox): isolate suspend requests and resume compensation Signed-off-by: Drew Newberry <anewberry@nvidia.com> * fix(supervisor): allow restoration to wait until agent admission Signed-off-by: Drew Newberry <anewberry@nvidia.com> --------- Signed-off-by: Drew Newberry <anewberry@nvidia.com>
2026-10-29 07:23:25 +08:00 · 2026-10-26 20:21:08 -07:00
parent 7b0f8ae07c
commit f07127e5d7
72 changed files with 5866 additions and 1731 deletions
@@ -639,6 +639,20 @@ class SandboxSession:
        self.sandbox = self._client.start(self.sandbox.name, workspace=self._workspace)
        return self.sandbox

+    def suspend(self) -> SandboxRef:
+        self.sandbox = self._client.suspend(
+            self.sandbox.name, workspace=self._workspace
+        )
+        return self.sandbox
+
+    def resume(self, *, restart_if_restore_fails: bool = False) -> SandboxRef:
+        self.sandbox = self._client.resume(
+            self.sandbox.name,
+            workspace=self._workspace,
+            restart_if_restore_fails=restart_if_restore_fails,
+        )
+        return self.sandbox
+

 class SandboxClient:
    """gRPC client for sandbox CRUD and command execution."""
@@ -1106,6 +1120,39 @@ class SandboxClient:
        )
        return _sandbox_ref(response.sandbox)

+    def suspend(self, name: str, *, workspace: str) -> SandboxRef:
+        """Suspend a ready sandbox, retaining its memory state.
+
+        Requires a compute driver and isolation backend that support suspend
+        and resume.
+        """
+        response = self._stub.SuspendSandbox(
+            openshell_pb2.SuspendSandboxRequest(
+                workspace_scope=_workspace_scope(workspace),
+                name=name,
+            ),
+            timeout=self._timeout,
+        )
+        return _sandbox_ref(response.sandbox)
+
+    def resume(
+        self,
+        name: str,
+        *,
+        workspace: str,
+        restart_if_restore_fails: bool = False,
+    ) -> SandboxRef:
+        """Resume a suspended sandbox in a new runtime generation."""
+        response = self._stub.ResumeSandbox(
+            openshell_pb2.ResumeSandboxRequest(
+                workspace_scope=_workspace_scope(workspace),
+                name=name,
+                restart_if_restore_fails=restart_if_restore_fails,
+            ),
+            timeout=self._timeout,
+        )
+        return _sandbox_ref(response.sandbox)
+
    def wait_deleted(
        self,
        name: str,
@@ -1152,6 +1199,17 @@ class SandboxClient:
            timeout_seconds=timeout_seconds,
        )

+    def wait_suspended(
+        self, name: str, *, workspace: str, timeout_seconds: float = 300.0
+    ) -> SandboxRef:
+        return self._wait_for_phase(
+            name,
+            workspace=workspace,
+            target_phase=openshell_pb2.SANDBOX_PHASE_SUSPENDED,
+            target_name="suspended",
+            timeout_seconds=timeout_seconds,
+        )
+
    def _wait_for_phase(
        self,
        name: str,
@@ -1176,6 +1234,16 @@ class SandboxClient:
                and sandbox.status.phase == openshell_pb2.SANDBOX_PHASE_STOPPED
            ):
                raise SandboxError(f"sandbox {name} stopped before becoming ready")
+            if target_phase == openshell_pb2.SANDBOX_PHASE_READY and (
+                sandbox.status.phase
+                in (
+                    openshell_pb2.SANDBOX_PHASE_SUSPENDING,
+                    openshell_pb2.SANDBOX_PHASE_SUSPENDED,
+                )
+            ):
+                raise SandboxError(
+                    f"sandbox {name} is suspended; resume it before waiting for ready"
+                )
            if sandbox.status.phase == openshell_pb2.SANDBOX_PHASE_ERROR:
                raise SandboxError(f"sandbox {name} entered error phase")
            time.sleep(1)
@@ -1992,6 +1992,8 @@ class _FakeSandboxStub:
        self.delete_request: openshell_pb2.DeleteSandboxRequest | None = None
        self.stop_request: openshell_pb2.StopSandboxRequest | None = None
        self.start_request: openshell_pb2.StartSandboxRequest | None = None
+        self.suspend_request: openshell_pb2.SuspendSandboxRequest | None = None
+        self.resume_request: openshell_pb2.ResumeSandboxRequest | None = None
        self.create_template_request: (
            openshell_pb2.CreateSandboxTemplateRequest | None
        ) = None
@@ -2063,6 +2065,38 @@ class _FakeSandboxStub:
            )
        )

+    def SuspendSandbox(
+        self,
+        request: openshell_pb2.SuspendSandboxRequest,
+        timeout: float | None = None,
+    ) -> Any:
+        self.suspend_request = request
+        _ = timeout
+        return SimpleNamespace(
+            sandbox=_make_sandbox_proto(
+                "sandbox-1",
+                _request_sandbox(request),
+                phase=openshell_pb2.SANDBOX_PHASE_SUSPENDING,
+                workspace=_request_workspace(request) or "default",
+            )
+        )
+
+    def ResumeSandbox(
+        self,
+        request: openshell_pb2.ResumeSandboxRequest,
+        timeout: float | None = None,
+    ) -> Any:
+        self.resume_request = request
+        _ = timeout
+        return SimpleNamespace(
+            sandbox=_make_sandbox_proto(
+                "sandbox-1",
+                _request_sandbox(request),
+                phase=openshell_pb2.SANDBOX_PHASE_STARTING,
+                workspace=_request_workspace(request) or "default",
+            )
+        )
+
    def CreateSandbox(
        self,
        request: openshell_pb2.CreateSandboxRequest,
@@ -2513,6 +2547,51 @@ def test_stop_and_start_forward_workspace_and_return_phase() -> None:
    assert starting.phase == openshell_pb2.SANDBOX_PHASE_STARTING


+def test_suspend_and_resume_forward_workspace_and_options() -> None:
+    stub = _FakeSandboxStub()
+    client = _client_with_fake_stub(stub)
+
+    suspending = client.suspend("job-1", workspace="team-a")
+    assert stub.suspend_request is not None
+    assert _request_sandbox(stub.suspend_request) == "job-1"
+    assert _request_workspace(stub.suspend_request) == "team-a"
+    assert suspending.phase == openshell_pb2.SANDBOX_PHASE_SUSPENDING
+
+    resumed = client.resume("job-1", workspace="team-a", restart_if_restore_fails=True)
+    assert stub.resume_request is not None
+    assert _request_sandbox(stub.resume_request) == "job-1"
+    assert _request_workspace(stub.resume_request) == "team-a"
+    assert stub.resume_request.restart_if_restore_fails
+    assert resumed.phase == openshell_pb2.SANDBOX_PHASE_STARTING
+
+    client.resume("job-1", workspace="team-a")
+    assert not stub.resume_request.restart_if_restore_fails
+
+
+def test_wait_ready_fails_fast_when_suspended() -> None:
+    class SuspendedStub(_FakeSandboxStub):
+        def GetSandbox(
+            self,
+            request: openshell_pb2.GetSandboxRequest,
+            timeout: float | None = None,
+        ) -> Any:
+            _ = timeout
+            return SimpleNamespace(
+                sandbox=_make_sandbox_proto(
+                    "sandbox-1",
+                    _request_sandbox(request),
+                    phase=openshell_pb2.SANDBOX_PHASE_SUSPENDED,
+                    workspace=_request_workspace(request) or "default",
+                )
+            )
+
+    client = _client_with_fake_stub(SuspendedStub())
+    with pytest.raises(SandboxError, match="is suspended"):
+        client.wait_ready("job-1", workspace="default", timeout_seconds=30)
+    suspended = client.wait_suspended("job-1", workspace="default", timeout_seconds=1)
+    assert suspended.phase == openshell_pb2.SANDBOX_PHASE_SUSPENDED
+
+
@pytest.mark.parametrize(
    ("phase", "should_succeed"),
    [