feat(checkpoint): hold-pause a harness when the sandbox snapshot keeps its processes

fern/versions/latest/pages/agent-server/checkpointable-harness.mdx

@@ -46,6 +46,12 @@ confirmed, or continuing fails, the harness is stopped and `run_until_parked` re
 saved conversation. A wait for the agent's own concurrency limit goes inside `activation.queued()`, so a checkpoint
 counts the activation as parked.
 
+On a provider whose snapshot keeps the sandbox's processes (`snapshot_keeps_processes`), create the session with
+`pause_mode="hold"`. The supervisor then acknowledges a pause without signalling anything: the harness keeps running,
+its model calls wait at the checkpoint, and the time held still does not count against the harness timeout. Each
+acknowledgement names the supervisor (`identity`, `PID.START`), and `SandboxSession.checkpoint_state()` records it
+with the pause it holds.
+
 ## Resume after a crash
 
 Per-harness policy stays in the adapter:

fern/versions/latest/pages/environment-tutorials/checkpointable-resources-server.mdx

@@ -55,14 +55,16 @@ class MyServer(SandboxOwnerResourcesServer):
 
 A commit snapshots each live sandbox through the provider and indexes the snapshot under
 `<checkpoint.state_dir>/<server>/snapshots/<checkpoint id>/`. The snapshot is the restore point: a restore forks it
-into a new sandbox, whatever happened to the original since. The snapshot keeps **files, not processes**.
+into a new sandbox, whatever happened to the original since. What the snapshot keeps depends on the provider:
+always the sandbox's **files**, and its **running processes** too when the provider's `snapshot_keeps_processes` says
+so (a VM memory snapshot). The checkpoint records which it was (`keeps_processes`).
 
 - **Grade an artifact, not a live sandbox.** Capture it with `capture_once` and declare `checkpoint_verify = "replay"`:
   a checkpoint does not wait for grading, and a `/verify` that runs again after a restore grades the same artifact.
   A session past its capture exports its data alone; the restore grades again from the saved artifact.
-- **In-place graders** that need processes the agent started declare `restore_policy = "before_agent_close"`. The seed
-  reply says so, and the Environment Server makes the episode a restart once its agent has closed, instead of
-  continuing it in a fork that has none of those processes.
+- **In-place graders** that need processes the agent started declare `restore_policy = "before_agent_close"` on a
+  provider whose snapshots keep files only. The seed reply says so, and the Environment Server makes the episode a
+  restart once its agent has closed, instead of continuing it in a fork that has none of those processes.
 - After a restore, the Environment Server seeds the session again under the replacement attempt. The base answers with
   the forked sandbox and does not run your setup again.
 

nemo_gym/agents/utils/process_supervisor.sh

@@ -34,10 +34,18 @@
 # spent paused does not count against the deadline. A pause requested before the launch holds the
 # launch until it is continued. A stop requested while paused continues the processes first, so the
 # harness receives TERM and can save its state. The receipt, written when the harness ends, answers
-# any pause the supervisor never acknowledged.
+# any pause the supervisor never acknowledged. Every acknowledgement also names the pause mode and the
+# supervisor's identity, PID.START (its /proc start time), which marks the processes it launched, and
+# records "wall" (epoch seconds) and "uptime" (seconds since boot, from /proc/uptime) at the moment it
+# was written.
+#
+# With --pause-mode hold, for a sandbox whose snapshot keeps its processes, a pause sends no signal: the
+# harness keeps running, the supervisor acknowledges at once and stops counting the deadline until the
+# continue, which sends nothing either. The checkpoint holds the harness's model calls instead.
 #
 # usage: process_supervisor.sh --timeout SECONDS --cleanup-timeout SECONDS --receipt PATH --stop-file PATH
-#        [--expire-file PATH] [--pause-file PATH --paused PATH --continue-file PATH --continued PATH]
+#        [--expire-file PATH] [--pause-file PATH --paused PATH --continue-file PATH --continued PATH
+#         [--pause-mode stop|hold]]
 #        -- COMMAND...
 
 # An inherited BASH_ENV or SHELLOPTS must not turn an expected failure into a missing receipt.
@@ -198,6 +206,8 @@ reached_alive() {
 stopped() { ((stopping)) || [[ -e $stop_file ]]; }
 
 expired() { [[ -n $expire_file && -e $expire_file ]]; }
+# How a pause treats the processes it reaches: stop sends STOP and CONT, hold sends nothing (see the header).
+pause_mode=stop
 
 write_receipt() {
     printf '{"cleanup_confirmed": %s, "return_code": %s, "error": %s, "timed_out": %s, "stopped": %s}' \
@@ -216,7 +226,17 @@ request_seq() {
 
 # Write acknowledgement $1 for request $2 atomically.
 acknowledge() {
-    printf '{"seq": %s, "confirmed": %s, "pids": [%s], "error": %s}' "$2" "$3" "$4" "$5" >"$1.tmp" && mv -f "$1.tmp" "$1"
+    local named=null wall uptime=null
+    [[ -n ${identity-} ]] && named="\"$identity\""
+    # The wall clock (epoch seconds) and uptime (seconds since boot) when the acknowledgement is written. A restore
+    # that keeps the processes but sets the wall clock to the host's can set it back to wall + (uptime now - uptime),
+    # as if the sandbox had not been stopped: uptime does not count the time a snapshot spent stopped.
+    wall=${EPOCHREALTIME-}
+    wall=${wall/,/.}
+    [[ -n $wall ]] || printf -v wall '%(%s)T' -1
+    read -r uptime _ </proc/uptime 2>/dev/null || uptime=null
+    printf '{"seq": %s, "confirmed": %s, "pids": [%s], "error": %s, "mode": "%s", "identity": %s, "wall": %s, "uptime": %s}' \
+        "$2" "$3" "$4" "$5" "$pause_mode" "$named" "$wall" "${uptime:-null}" >"$1.tmp" && mv -f "$1.tmp" "$1"
 }
 
 # Record in paused[PID]=START every process a pause reaches: what a stop would reach now, and
@@ -264,6 +284,10 @@ paused_settled() {
 pause() {
     local limit error=null pid list=
     paused=()
+    if [[ $pause_mode == hold ]]; then
+        acknowledge "$paused_ack" "$1" true "" null
+        return
+    fi
     if [[ -n ${harness-} ]]; then
         pause_scan
         [[ -n ${groups[harness]} ]] && owned_group "$harness" "${groups[harness]}" && kill -STOP -- "-$harness" 2>/dev/null
@@ -290,6 +314,10 @@ pause() {
 # Continue request $1: continue every process the pause stopped and acknowledge.
 resume_paused() {
     local limit error=null
+    if [[ $pause_mode == hold ]]; then
+        acknowledge "$continued_ack" "$1" true "" null
+        return
+    fi
     now
     limit=$((NOW + pause_wait))
     signal_paused CONT
@@ -345,6 +373,10 @@ while (($#)) && [[ $1 != -- ]]; do
     --paused) paused_ack=$2 ;;
     --continue-file) continue_file=$2 ;;
     --continued) continued_ack=$2 ;;
+    --pause-mode)
+        [[ $2 == stop || $2 == hold ]] || { echo "--pause-mode must be stop or hold" >&2; exit 2; }
+        pause_mode=$2
+        ;;
     *) echo "unexpected argument: $1" >&2; exit 2 ;;
     esac
     shift 2

nemo_gym/agents/utils/sandbox_session.py

@@ -11,7 +11,7 @@
 from pathlib import Path
 from shlex import quote
 from time import monotonic
-from typing import Literal, TypedDict, cast
+from typing import Literal, NotRequired, TypedDict, cast
 
 from pydantic import TypeAdapter
 
@@ -27,6 +27,7 @@
     CleanupReceipt,
     HarnessTerminationError,
     PauseAcknowledgement,
+    PauseMode,
     SupervisorRequestOutcome,
     check_stdlib_runner,
     remove_session_directory,
@@ -126,24 +127,40 @@ class HarnessPauseResult:
     ``confirmed``: every process the supervisor reached is paused (or running again), or no launch had started,
     and the launch will honour the request when it starts. ``finished``: the harness had already ended; its cleanup
     receipt answers instead. ``unconfirmed``: the supervisor did not confirm in time or reported ``error``.
-    ``pids`` are the processes a pause stopped.
+    ``pids`` are the processes a pause stopped. ``identity`` is the answering supervisor's ``PID.START`` and
+    ``mode`` its pause mode, when it reported them. ``wall`` and ``uptime`` are the sandbox's wall clock (epoch
+    seconds) and uptime (seconds since boot) when it answered.
     """
 
     status: Literal["confirmed", "finished", "unconfirmed"]
     pids: tuple[int, ...] = ()
     error: str | None = None
+    identity: str | None = None
+    mode: PauseMode | None = None
+    wall: float | None = None
+    uptime: float | None = None
 
     @property
     def confirmed(self) -> bool:
         return self.status == "confirmed"
 
 
 class SandboxSessionCheckpoint(TypedDict):
-    """What ``SandboxSession.reattach`` needs to rebind a session to a sandbox restored from a snapshot."""
+    """What ``SandboxSession.reattach`` needs to rebind a session to a sandbox restored from a snapshot.
+
+    ``pause_seq`` and ``supervisor_identity`` are there while the launch's supervisor holds a confirmed pause: on a
+    sandbox whose snapshot keeps its processes, that launch can still be running after the restore.
+    """
 
     session_root: str
     session_dir: str
     launch_started: bool
+    pause_seq: NotRequired[int]
+    supervisor_identity: NotRequired[str]
+    # The sandbox's wall clock and uptime when that pause was confirmed, for a restore that sets the wall clock back
+    # to where the snapshot left it.
+    pause_wall: NotRequired[float]
+    pause_uptime: NotRequired[float]
 
 
 _CHECKPOINT_STATE = TypeAdapter(SandboxSessionCheckpoint)
@@ -162,6 +179,10 @@ def _pause_result(
         status="confirmed" if acknowledgement["confirmed"] else "unconfirmed",
         pids=tuple(acknowledgement["pids"]),
         error=acknowledgement["error"],
+        identity=acknowledgement.get("identity"),
+        mode=acknowledgement.get("mode"),
+        wall=acknowledgement.get("wall"),
+        uptime=acknowledgement.get("uptime"),
     )
 
 
@@ -197,11 +218,13 @@ class SandboxSession[Artifacts]:
     HTTP request binding and activation retries belong to the agent session layer.
 
     A ``pausable`` session's supervisor also pauses and continues the harness in place for a
-    checkpoint (``pause_harness`` and ``continue_harness``): every process it reaches, and every
-    process carrying its marker, is stopped with SIGSTOP and continued with SIGCONT, and time spent
-    paused does not count against the harness timeout. Directories under ``session_root`` survive a
-    snapshot of a borrowed sandbox; ``checkpoint_state`` and ``reattach`` rebind a session to the
-    sandbox a restore forks from that snapshot.
+    checkpoint (``pause_harness`` and ``continue_harness``), and time spent paused does not count
+    against the harness timeout. In ``pause_mode`` ``stop``, every process it reaches, and every
+    process carrying its marker, is stopped with SIGSTOP and continued with SIGCONT. In ``hold``, for
+    a sandbox whose snapshot keeps its processes, nothing is signalled: the harness keeps running
+    while the checkpoint holds its model calls. Directories under ``session_root`` survive a snapshot
+    of a borrowed sandbox; ``checkpoint_state`` and ``reattach`` rebind a session to the sandbox a
+    restore forks from that snapshot.
     """
 
     sandbox: AsyncSandbox
@@ -210,6 +233,7 @@ class SandboxSession[Artifacts]:
     harness: str
     owns_sandbox: bool = False
     pausable: bool = False
+    pause_mode: PauseMode = "stop"
     launch_started: bool = field(default=False, init=False)
     cleanup: CleanupReceipt | None = field(default=None, init=False)
     artifacts: Artifacts | None = field(default=None, init=False)
@@ -235,6 +259,11 @@ class SandboxSession[Artifacts]:
     _pause_seq: int = field(default=0, init=False)
     _continue_seq: int = field(default=0, init=False)
     _continued_seq: int = field(default=0, init=False)
+    # The supervisor that confirmed pause ``_held_seq``, as ``PID.START``.
+    _held_identity: str | None = field(default=None, init=False)
+    _held_seq: int = field(default=0, init=False)
+    # The sandbox's wall clock and uptime when that supervisor confirmed pause ``_held_seq``.
+    _held_clocks: tuple[float, float] | None = field(default=None, init=False)
     # Holds the provider's polls of the launch's command while the harness is paused, so a sandbox being snapshotted
     # is not asked about it; released before anything else runs a command in the sandbox.
     _parked_polls: ExitStack | None = field(default=None, init=False, repr=False)
@@ -246,11 +275,17 @@ def session_root(self) -> str:
 
     def checkpoint_state(self) -> SandboxSessionCheckpoint:
         """Where this session's files are, for ``reattach`` after a restore; JSON-serializable."""
-        return {
+        state: SandboxSessionCheckpoint = {
             "session_root": self.session_root,
             "session_dir": self.session_dir,
             "launch_started": self.launch_started,
         }
+        if self.launch_started and self._held_identity is not None and self._continued_seq < self._held_seq:
+            state["pause_seq"] = self._held_seq
+            state["supervisor_identity"] = self._held_identity
+            if self._held_clocks is not None:
+                state["pause_wall"], state["pause_uptime"] = self._held_clocks
+        return state
 
     @classmethod
     def reattach(
@@ -262,13 +297,16 @@ def reattach(
         harness: str,
         owns_sandbox: bool = False,
         pausable: bool = False,
+        pause_mode: PauseMode = "stop",
     ) -> "SandboxSession[Artifacts]":
         """Rebind a session from ``checkpoint_state`` to ``sandbox``, restored from a snapshot of the original.
 
-        A snapshot keeps files, not processes: no launch recorded in ``state`` still runs in ``sandbox``, so the
-        reattached session has none. Rearm it with ``begin_activation`` and a fresh launch directory under
-        ``session_root`` before the next ``execute``; ``state["launch_started"]`` says whether the harness had
-        already started, for an adapter that resumes the harness's saved conversation.
+        The reattached session has no launch of its own. What a snapshot kept depends on the provider. When it
+        kept files only, no launch recorded in ``state`` still runs in ``sandbox``: rearm the session with
+        ``begin_activation`` and a fresh launch directory under ``session_root`` before the next ``execute``;
+        ``state["launch_started"]`` says whether the harness had already started, for an adapter that resumes
+        the harness's saved conversation. When it kept the processes too, a launch whose supervisor held a
+        confirmed pause (``state["pause_seq"]``) may still run in ``sandbox``.
         """
         state = _CHECKPOINT_STATE.validate_python(state, strict=True)
         session = cls(
@@ -278,6 +316,7 @@ def reattach(
             harness=harness,
             owns_sandbox=owns_sandbox,
             pausable=pausable,
+            pause_mode=pause_mode,
         )
         if state["session_root"] != state["session_dir"]:
             session._session_root = state["session_root"]
@@ -382,6 +421,7 @@ async def execute(
                 cleanup_timeout=cleanup_timeout,
                 stdlib_runner=command.stdlib_runner,
                 pausable=self.pausable,
+                pause_mode=self.pause_mode,
             )
             self.launch_started = True
             self._exec_task = asyncio.create_task(
@@ -451,7 +491,9 @@ async def begin_activation(self, *, session_dir: str, timeout: float) -> None:
         self._collect = None
         self._capture_attempted = False
         # The next launch has a fresh supervisor: nothing paused carries over.
-        self._pause_seq = self._continue_seq = self._continued_seq = 0
+        self._pause_seq = self._continue_seq = self._continued_seq = self._held_seq = 0
+        self._held_identity = None
+        self._held_clocks = None
         self._release_parked_polls()
 
     async def pause_harness(self, *, timeout: float = DEFAULT_PAUSE_TIMEOUT) -> HarnessPauseResult:
@@ -470,6 +512,11 @@ async def pause_harness(self, *, timeout: float = DEFAULT_PAUSE_TIMEOUT) -> Harn
             self.sandbox, session_dir=self.session_dir, workdir=self.workdir, seq=self._pause_seq, timeout=timeout
         )
         result = _pause_result(outcome, acknowledgement)
+        if result.confirmed and result.identity is not None:
+            self._held_seq, self._held_identity = self._pause_seq, result.identity
+            self._held_clocks = None
+            if result.wall is not None and result.uptime is not None:
+                self._held_clocks = (result.wall, result.uptime)
         if result.confirmed and self.launch_started and self._parked_polls is None:
             # A no-op scope on providers that do not poll command status.
             self._parked_polls = ExitStack()

nemo_gym/agents/utils/supervisor_client.py

@@ -24,6 +24,10 @@
 continued.json        Supervisor           Acknowledges continue N once the processes run again
 ===================== ==================== ===========================================
 
+A supervisor launched with ``pause_mode="hold"``, for a sandbox whose snapshot keeps its processes,
+acknowledges a pause without stopping anything; the harness keeps running and its model calls are held
+instead.
+
 Pause and continue are re-armable: each checkpoint writes a sequence number greater than the last,
 and an acknowledgement names the number it answers, so a request is never mistaken for an earlier
 one and no file has to be removed. A pause written before the launch holds the launch until its
@@ -102,12 +106,27 @@ class CleanupReceipt(TypedDict):
 
 
 class PauseAcknowledgement(TypedDict):
-    """A supervisor's answer to pause or continue request ``seq``, from ``paused.json`` or ``continued.json``."""
+    """A supervisor's answer to pause or continue request ``seq``, from ``paused.json`` or ``continued.json``.
+
+    ``mode`` is the supervisor's pause mode and ``identity`` its ``PID.START``, ``None`` before it launched;
+    acknowledgements written before either existed have neither. ``wall`` (epoch seconds) and ``uptime`` (seconds
+    since boot) are the sandbox's clocks when the acknowledgement was written; ``uptime`` is ``None`` without a
+    readable ``/proc/uptime``.
+    """
 
     seq: int
     confirmed: bool
     pids: list[int]
     error: str | None
+    mode: NotRequired[Literal["stop", "hold"]]
+    identity: NotRequired[str | None]
+    wall: NotRequired[float]
+    uptime: NotRequired[float | None]
+
+
+# How a supervisor pauses: ``stop`` sends SIGSTOP and SIGCONT to every process it reaches; ``hold`` sends nothing,
+# for a sandbox whose snapshot keeps its processes running.
+PauseMode = Literal["stop", "hold"]
 
 
 # What a pause or continue request found:
@@ -211,6 +230,7 @@ def supervised_launch_command(
     cleanup_timeout: float,
     stdlib_runner: bool = False,
     pausable: bool = False,
+    pause_mode: PauseMode = "stop",
 ) -> str:
     """Fence delayed launches and run a harness command under the shared supervisor.
 
@@ -229,7 +249,7 @@ def supervised_launch_command(
     When the link cannot be made, the argv runs unchanged.
 
     With ``pausable``, the supervisor also serves pause and continue requests in
-    ``session_dir`` (see ``request_pause`` and ``request_continue``).
+    ``session_dir`` (see ``request_pause`` and ``request_continue``), in ``pause_mode``.
     """
     check_stdlib_runner(command, stdlib_runner=stdlib_runner)
     supervisor = quote(f"{session_dir}/{SUPERVISOR_FILE}")
@@ -238,6 +258,7 @@ def supervised_launch_command(
         f"--paused {quote(f'{session_dir}/{PAUSED_FILE}')} "
         f"--continue-file {quote(f'{session_dir}/{CONTINUE_REQUEST_FILE}')} "
         f"--continued {quote(f'{session_dir}/{CONTINUED_FILE}')} "
+        + (f"--pause-mode {pause_mode} " if pause_mode != "stop" else "")
         if pausable
         else ""
     )