refactor(checkpoint): use NeMo Gym's generation-cut wire names

docs/guides/single-controller.md

@@ -113,6 +113,10 @@ With `checkpointing.save_data_plane: true`, each Single-Controller checkpoint co
 - A native TQ snapshot containing rollout tensor payloads and TQ state.
 - A metadata-only replay index describing the completed rollout groups stored in TQ.
 - A `rollout_recovery.pt` ownership ledger describing unfinished prompt groups that must be redispatched after a restart.
+- With `rollout_recovery.target_level: turn` or `prefix`, a
+  `gym_checkpoint.json` manifest and per-Gym-instance participant snapshots for
+  unfinished episodes. `prefix` additionally stages decoded tokens from active
+  model calls at each coordinated checkpoint.
 - A `replacement_reserve.pt` sidecar containing prompts held for dropped-rollout replacement, when applicable.
 - The sampler dispatch position needed to continue scheduling from the correct point.
 
@@ -237,26 +241,59 @@ checkpoint exists.
 
 :::{note}
 Completed groups are restored directly from the TQ snapshot. For unfinished
-token-capture groups, `rollout_recovery.default_granularity` controls both live
-failure and restart behavior:
+token-capture groups, `rollout_recovery.target_level` selects the requested
+restart boundary:
 
+- `turn` requests coordinated Gym participant recovery from the last saved turn.
 - `sibling` preserves each sealed sibling and redispatches only unfinished ones.
 - `prompt_group` retries every sibling in the group when any sibling is unfinished.
 
 `sibling` is the default and avoids regenerating completed work. Use
 `prompt_group` when every generation in a recovered group must come from the
 policy weights live at redispatch.
 
-`task_source_granularity_overrides` can select the policy using the Gym
+Turn recovery coordinates Gym and RL as one checkpoint cut. Each prompt group
+is pinned to the Gym shard replica that accepted it. SC first closes its narrow
+`/run` dispatch admission gate; already-submitted requests, completion callbacks,
+finalization, and TQ writes remain live. Gym parks its participants and commits
+the subset of candidate episodes it still owns into the checkpoint's
+`gym-instances/` tree. Candidate replies already on the wire drain through the
+ordinary RL/TQ completion path. SC then acquires the exclusive data-plane
+barrier, verifies that the remaining Gym-owned episodes and their TQ staging
+keys match the Gym commit, and captures TQ plus the rollout ledger. If a
+candidate is neither exported by Gym nor drained before the deadline, the
+checkpoint aborts. Gym leases are renewed throughout and all deployments are
+resumed after the cut succeeds or aborts.
+
+On restart, SC validates that `gym_checkpoint.json` names exactly the episodes
+owned by the rollout ledger and that the saved Gym actor topology matches the
+live one. It restores Gym before starting rollout dispatch, advances each
+restored episode to its next attempt, and routes it back to the same Gym actor.
+If one shard replica fails restore, every replica conservatively retires its
+replacement attempts before the outer restore fails. This also covers a lost
+Ray reply where the remote restore may actually have succeeded, so a partially
+restored topology is never released for execution. A second checkpoint taken
+after restore but before redispatch includes the unclaimed restored attempt, so
+another crash does not lose the saved turn boundary.
+
+The `turn` target currently requires vLLM generation, NeMo Gym, token capture,
+periodic rollout snapshots, `restore_mode: latest`, and the same named Gym
+shards and replica counts on restart. Only checkpoints using the current rollout
+recovery schema are accepted.
+
+`task_source_target_level_overrides` can coarsen the target using the Gym
 `task_source` embedded in the raw rollout row. Unlike `agent_ref`, this identity
 is available before Gym resolves the concrete agent and SC reserves the recovery
 group. When a row already carries an `agent_ref`, a matching
-`agent_granularity_overrides` entry wins over a matching task-source entry,
+`agent_target_level_overrides` entry wins over a matching task-source entry,
 mirroring Gym's concrete-route precedence. Otherwise the task-source override,
 then the global default, applies. The agent map also keeps datasets collated
 before Gym recorded `task_source` working, although re-collating them is
-recommended. Non-default policies require `token_capture.enabled: true`. The
-task source and resolved policy are persisted in `rollout_recovery.pt`, so
+recommended. Overrides may only select a coarser level than the global target.
+Until Gym advertises group-scoring capabilities, configure group-scored routes
+such as GenRM explicitly as `prompt_group`. Non-default policies require
+`token_capture.enabled: true`. The task source and resolved policy are persisted
+in `rollout_recovery.pt`, so
 recovery does not reinterpret an existing group using changed configuration. A
 generation that already finished keeps its tokens in the token-capture staging
 area, so `sibling` reuses them unchanged; a redispatched sibling produces a new

examples/configs/grpo_math_1B_megatron_single_controller.yaml

@@ -219,20 +219,22 @@ token_capture:
   enabled: false
   staging_partition: rollout_staging
 
-# Retry and restore policy for unfinished token-capture groups. "sibling" reuses
-# already sealed generations; "prompt_group" regenerates every sibling when any
-# sibling fails in-process or was unfinished at a checkpoint.
+# Target recovery boundary for unfinished token-capture groups. "prefix" also
+# preserves decoded tokens from an active model call; "turn" enables coordinated
+# Gym participant snapshots; "sibling" preserves sealed siblings; "prompt_group"
+# regenerates every sibling in an unfinished group.
 rollout_recovery:
-  default_granularity: sibling
-  # Per-Gym-task-source override, keyed on extra_env_info.task_source.
-  #   task_source_granularity_overrides:
+  target_level: sibling
+  # Per-Gym-task-source coarsening override, keyed on extra_env_info.task_source.
+  # Until Gym advertises group-scoring capability, configure GenRM explicitly:
+  #   task_source_target_level_overrides:
   #     genrm_compare_resources_server: prompt_group
-  task_source_granularity_overrides: {}
+  task_source_target_level_overrides: {}
   # Agent-specific override for rows already carrying extra_env_info.agent_ref.
   # A matching agent entry wins over a matching task-source entry and also
   # provides compatibility for legacy datasets without task_source.
-  #   agent_granularity_overrides: {genrm_agent: prompt_group}
-  agent_granularity_overrides: {}
+  #   agent_target_level_overrides: {genrm_agent: prompt_group}
+  agent_target_level_overrides: {}
 
 # Leaving either map above non-empty requires the top-level token_capture section
 # to set enabled: true.

examples/configs/ppo_math_1B_megatron_single_controller.yaml

@@ -267,20 +267,22 @@ token_capture:
   enabled: false
   staging_partition: rollout_staging
 
-# Retry and restore policy for unfinished token-capture groups. "sibling" reuses
-# already sealed generations; "prompt_group" regenerates every sibling when any
-# sibling fails in-process or was unfinished at a checkpoint.
+# Target recovery boundary for unfinished token-capture groups. "prefix" also
+# preserves decoded tokens from an active model call; "turn" enables coordinated
+# Gym participant snapshots; "sibling" preserves sealed siblings; "prompt_group"
+# regenerates every sibling in an unfinished group.
 rollout_recovery:
-  default_granularity: sibling
-  # Per-Gym-task-source override, keyed on extra_env_info.task_source.
-  #   task_source_granularity_overrides:
+  target_level: sibling
+  # Per-Gym-task-source coarsening override, keyed on extra_env_info.task_source.
+  # Until Gym advertises group-scoring capability, configure GenRM explicitly:
+  #   task_source_target_level_overrides:
   #     genrm_compare_resources_server: prompt_group
-  task_source_granularity_overrides: {}
+  task_source_target_level_overrides: {}
   # Agent-specific override for rows already carrying extra_env_info.agent_ref.
   # A matching agent entry wins over a matching task-source entry and also
   # provides compatibility for legacy datasets without task_source.
-  #   agent_granularity_overrides: {genrm_agent: prompt_group}
-  agent_granularity_overrides: {}
+  #   agent_target_level_overrides: {genrm_agent: prompt_group}
+  agent_target_level_overrides: {}
 
 # Leaving either map above non-empty requires the top-level token_capture section
 # to set enabled: true.

examples/nemo_gym/grpo_workplace_assistant_two_replicas.yaml

@@ -0,0 +1,25 @@
+# Two Workplace Assistant Gym replicas for checkpoint/recovery tests. Replicas
+# share one port range, so NeMo-RL intentionally places them on separate Ray
+# nodes with the default STRICT_SPREAD strategy. The training shape and model
+# are supplied by the functional launcher.
+
+defaults: "grpo_qwen3_30ba3b_instruct.yaml"
+
+env:
+  nemo_gym:
+    # Shards replace the inherited unsharded server list. Policy-model config is
+    # an RL-owned overlay duplicated into every otherwise-identical replica.
+    config_paths: null
+    code_gen: null
+    shards:
+      - name: workplace
+        replicas: 2
+        port_range_low: 5000
+        port_range_high: 5999
+        config_paths:
+          - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+          - responses_api_agents/checkpoint_test_agent/configs/workplace_assistant.yaml
+    common_inherited_overlays:
+      - policy_model
+    allowed_duplicate_entries:
+      - policy_model

examples/nemo_gym/prefetch_venvs.py

@@ -29,6 +29,7 @@
 
 from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env
 from nemo_rl.distributed.virtual_cluster import init_ray
+from nemo_rl.environments.gym_checkpoint_adapter import GymCheckpointInstance
 from nemo_rl.environments.nemo_gym import (
     NemoGym,
     NemoGymConfig,
@@ -86,6 +87,8 @@ def prefetch_nemo_gym_venvs(config_paths: list[str]) -> None:
                 ray_namespace=None,
                 initial_global_config_dict=nemo_gym_dict,
                 invalid_tool_call_patterns=None,
+                checkpoint_instance=GymCheckpointInstance("nemo_gym", 0),
+                turn_checkpointing_enabled=False,
             )
 
             nemo_gym_opts = {

nemo_rl/algorithms/async_utils/replay_buffer.py

@@ -263,7 +263,11 @@ class DataPlaneCheckpointBarrier:
     """
 
     def __init__(self) -> None:
-        self._condition = asyncio.Condition()
+        lock = asyncio.Lock()
+        self._condition = asyncio.Condition(lock)
+        # Separate from _condition so a finished mutation wakes only waiters for
+        # it, not every mutation queued behind an active checkpoint.
+        self._mutation_finished = asyncio.Condition(lock)
         self._checkpoint_active = False
         self._active_mutations = 0
         self._section_holders: set[asyncio.Task[Any]] = set()
@@ -329,9 +333,25 @@ async def mutation(
                 # Count the section even when its body raised. A redundant
                 # snapshot is safe; skipping a partially applied mutation is not.
                 self._mutation_version += 1
+                self._mutation_finished.notify_all()
                 if self._active_mutations == 0:
                     self._condition.notify_all()
 
+    async def wait_for_mutation(self, since_version: int, timeout_s: float) -> None:
+        """Return once a mutation finishes after ``since_version``, or on timeout.
+
+        Read :attr:`mutation_version` before inspecting state, then pass it here:
+        a mutation that lands in between returns at once instead of being missed.
+        """
+        async with self._mutation_finished:
+            try:
+                async with asyncio.timeout(max(0.0, timeout_s)):
+                    await self._mutation_finished.wait_for(
+                        lambda: self._mutation_version != since_version
+                    )
+            except TimeoutError:
+                return
+
     async def drain_telemetry(self) -> DataPlaneCheckpointBarrierTelemetry:
         """Return and reset interval waits while preserving current state."""
         async with self._condition:

nemo_rl/algorithms/distillation.py

@@ -540,6 +540,7 @@ def init_nemo_gym():
                     # Distillation does not configure vLLM for router replay.
                     enable_router_replay=False,
                     use_fastokens=bool(policy_config["tokenizer"].get("use_fastokens")),
+                    turn_checkpointing_enabled=False,
                 )
                 try:
                     validate_dataset_agent_coverage(

nemo_rl/algorithms/grpo.py

@@ -893,6 +893,7 @@ def _spinup_nemo_gym(base_urls, model_name):
             tokenizer=tokenizer,
             enable_router_replay=router_replay_enabled(policy_config),
             use_fastokens=bool(policy_config["tokenizer"].get("use_fastokens")),
+            turn_checkpointing_enabled=False,
         )
         train_splits = (
             {f"train[{name}]": split for name, split in dataset.items()}