Why turn-level snapshots stall, then end the run, with most Gym agents

NVIDIA-NeMo/RL #4500 (Gym turn-level rollout recovery). RL permalinks pin head 341b9ad7; Gym permalinks pin the submodule commit 120d516f.

Four pages on this PR, one question each. This one: which Gym setups can use turn.

Gym prepares the servers one stage at a time, model before agent (PREPARE_ORDER). Once prepared, the model holds every reply until resume (release_response), and RL never turns on Gym's generation cuts, so nothing lets a call finish early (checkpoint_generation_cuts: False). A restart-only agent counts every running /run as a blocker (readiness), and those runs are waiting for the model's held replies. So the agent can never become ready. RL treats that "not ready" as a passing problem (coordinator L89) and tries again at the next snapshot.

One snapshot attempt, two Gym setups

Example: rollout s1 has finished turn 1 and is in the middle of turn 2 when the snapshot starts.

How many Gym setups hit this: only simple_agent can save its sessions (default False; 1 of 57 agents at this Gym commit). Resources servers default to restart-only (checkpoint_mode, what it means), and 8 of 146 opt out. A restart-only resources session is live from seed to verify, so it blocks the same way. The functional test passes because Workplace is one of the setups that works. Nothing at startup or in the docs warns about the others.

🔴 review comment — gym_checkpoint_adapter.py:208

In the setup that works: where s1 parks, and what the snapshot keeps

simple_agent has two park points (Gym calls them boundaries) in every turn: right after the model's reply (app.py:398) and after that turn's tools return (app.py:482). There is none in the middle of a tool call. So where s1 parks depends on what it is doing in turn 2:

The tool server saves its state at commit, after every agent has parked, so no tool call is in flight when it saves. Its saved state always matches the turn the agent parked at. The second row happens only if Gym has already asked agents to park when the reply arrives; otherwise s1 starts its tools and ends up in the third row. * Gym can keep a call's partly generated tokens (a "generation cut"), but RL does not hook that up in this PR.

What would fix it — proposed, not in the PR

Where the fix lands: Gym and RL. Gym, in the unmerged branch this PR already pins: restart-only servers say so in their status, and in best-effort mode they stop blocking prepare (a–c below). RL: a new rollout_recovery.allow_restart_only_gym_servers setting, a check in _spinup_gym, and a snapshot that accepts rollouts Gym could not save.

Nothing below is implemented or run. The setting only means something when target_level needs Gym saves (turn here, prefix in #4506), so config validation rejects it with sibling or prompt_group. Its two values:

• false (the default): setup fails if any Gym agent or resources server is restart-only, and the error names them. You find out before the run starts, and you know which environments need checkpoint support.
• true (best effort): rollouts on servers that can save keep turn recovery. Rollouts on a restart-only server get sibling behavior: after a crash they restart from their prompt, and their finished siblings are kept. Setup warns and names those servers. If no server is restart-only, setup warns that the setting should be false, so a restart-only server added later is caught at setup.

Gym: what stops the stall in best-effort mode. Example: s1 runs on simple_agent, s2 on a restart-only agent, and a snapshot starts.

Gym also needs RestartOnlyAgentParticipant.status_extra to return {"restart_only": True}, so setup can find these agents. Resources servers already report their mode.

RL: the setting and the setup check.

@@ RL nemo_rl/algorithms/single_controller_utils/config.py L718-L719 · RolloutRecoveryConfig
     # Seconds Gym gets to park its in-flight steps; see the exemplar YAML.
     gym_prepare_timeout_s: Annotated[float, Field(gt=0)] = 300.0
+    # Needs a target_level that saves Gym state. false: a restart-only Gym server
+    # fails setup. true: its rollouts restart from the prompt after a crash.
+    allow_restart_only_gym_servers: bool = False

@@ RL nemo_rl/environments/gym_checkpoint_adapter.py L208-L221 · GymCheckpointAdapter.discover
-    async def discover(self) -> GymCheckpointParticipantSummary:
+    async def discover(self) -> tuple[str, ...]:
         ...
         participants = self._participants
-        return GymCheckpointParticipantSummary(
-            instance_id=self._instance.instance_id,
-            members=tuple(
-                (member.server_name, member.kind) for member in participants.members
-            ),
-        )
+        # The summary class goes (another review comment); return what setup needs.
+        return tuple([
+            member.server_name
+            for member in participants.members
+            if _is_restart_only(await self._status(member.server_name))
+        ])

@@ RL nemo_rl/algorithms/single_controller_utils/setup.py L801 · _spinup_gym
+    recovery = master_config.rollout_recovery
+    restart_only = shard_set.restart_only_servers  # from each replica's discover()
+    if recovery.turn_checkpointing_enabled:
+        if restart_only and not recovery.allow_restart_only_gym_servers:
+            raise ValueError(
+                "rollout_recovery.target_level='turn' needs every Gym agent and "
+                f"resources server to save its rollouts; these are restart-only: {restart_only}. "
+                "Add checkpoint support to them, or set "
+                "rollout_recovery.allow_restart_only_gym_servers=true to restart their "
+                "rollouts from the prompt after a crash."
+            )
+        if restart_only:
+            warnings.warn(
+                f"Gym servers {restart_only} are restart-only: after a crash, their "
+                "unfinished rollouts restart from the prompt. Rollouts on other servers "
+                "continue from their last saved turn.",
+                UserWarning,
+                stacklevel=2,
+            )
+        elif recovery.allow_restart_only_gym_servers:
+            warnings.warn(
+                "rollout_recovery.allow_restart_only_gym_servers=true, but every Gym "
+                "server can save its rollouts. Set it to false so a restart-only server "
+                "added later fails at setup.",
+                UserWarning,
+                stacklevel=2,
+            )
     return shard_set, time.perf_counter() - t0

_status(name) reads the same /ng-control/v1/checkpoint/status route Gym's own discover already calls, and _is_restart_only is true for {"restart_only": True} or {"mode": "restart_only"}. _initialize_checkpoint_adapters_concurrently already gets every replica's result from ray.get and drops it; it passes the names up instead.

RL: the snapshot accepts unsaved rollouts. The drain (L4622-4636) and the ownership check before the TQ save (L4666-4676) accept rollouts Gym reports as unsaved, instead of timing out or raising. Restore needs no change: it already restarts any rollout Gym did not restore from its prompt (rollout_recovery.py:527-548).

Not this. Gym's docs say the controller can resume, retire the blockers and prepare again (coordination.py:22-23). With a restart-only agent, every running rollout is a blocker, so each snapshot would cancel every rollout in flight. When the snapshot interval is shorter than a rollout, no rollout ever finishes.