From 6d1c0f8c78fb872910ca3c2810b49bbe3aabd431 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 19 Aug 2026 11:18:05 -0400 Subject: [PATCH 01/28] feat(rollout): add partial rollout lineage ledger Signed-off-by: Anish Mahishi (cherry picked from commit 52e2dbcf771cbdd5137a597a89488b4711b18e90) --- .../algorithms/async_utils/replay_buffer.py | 28 + nemo_rl/algorithms/single_controller.py | 282 ++-- nemo_rl/experience/rollout_manager.py | 356 +++-- nemo_rl/experience/rollout_reassembler.py | 12 +- .../experience/rollout_reassembler_actor.py | 7 +- nemo_rl/experience/rollout_recovery.py | 1256 ++++++++++++----- .../data_plane/test_rollout_reassembler.py | 34 +- tests/unit/experience/test_rollout_manager.py | 150 +- .../test_rollout_reassembler_actor.py | 1 + tests/unit/experience/test_rollouts.py | 2 + .../single_controller/test_checkpointing.py | 4 + .../test_finalizer_lifecycle.py | 9 +- 12 files changed, 1486 insertions(+), 655 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 52b5d570152..529a27f9f43 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -1052,6 +1052,11 @@ def set_post_write_enricher( """Install the required enrichment stage run before slots become ready.""" self._post_write_enricher = enricher + @property + def group_ids(self) -> tuple[str, ...]: + """Return a stable snapshot of controller-local replay ownership.""" + return tuple(self._group_ids) + def reserve( self, *, @@ -1228,6 +1233,29 @@ async def remove_group(self, group_id: str, *, remove_in_dp: bool = False) -> in raise ValueError(f"unknown group_id={group_id!r}") from error return await self._remove_unlocked([idx], clear_data_plane=remove_in_dp) + async def clear_staging_keys(self, staging_keys: list[str]) -> None: + """Clear known token-capture staging rows under the checkpoint barrier.""" + if not staging_keys: + return + if self._staging_partition_id is None: + raise RuntimeError( + "cannot clear token-capture staging keys without a staging partition" + ) + if self._data_plane_checkpoint_barrier is None: + raise RuntimeError( + "TQReplayBuffer must be bound to the controller data-plane " + "checkpoint barrier before clearing staging samples" + ) + unique_keys = list(dict.fromkeys(staging_keys)) + async with self._data_plane_checkpoint_barrier.mutation(): + await call_data_plane( + self._dp_client, + "clear_samples", + offload_sync=True, + sample_ids=unique_keys, + partition_id=self._staging_partition_id, + ) + async def commit_finalized( self, group_id: str, diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index f8524831887..e161178fee7 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -269,6 +269,7 @@ def __init__( # Rebind so writer and sampler share one buffer instance even # when Ray deserializes rollout_manager and tq_buffer separately. self._rollout_manager._tq_buffer = self._buffer + self._rollout_recovery_ledger = self._rollout_manager.recovery_ledger # Direct access, deliberately. A getattr default here reads as defensive but # buys a silent failure mode: rename or drop the field and @@ -1171,61 +1172,64 @@ def _request_staging_keys(request: "ReassemblyRequest") -> list[str]: keys.append(record["staging_key"]) return list(dict.fromkeys(keys)) - async def _cleanup_known_finalization_request( - self, request: "ReassemblyRequest" + async def _cleanup_known_finalization_request_unlocked( + self, + cut: DataPlaneMutationCut, + request: "ReassemblyRequest", ) -> None: - """Clear known request ownership after a pre-publication/known outcome.""" + """Clear known request ownership while holding a barrier mutation slot.""" errors: list[BaseException] = [] - # One shared mutation slot for both clears: they run outside the train - # pump task, so a live data-plane checkpoint must not interleave them. - # Offloaded so the rollout pump's event loop keeps serving other - # dispatches and finalizer handoffs while the clears are in flight. - async with self._data_plane_checkpoint_barrier.mutation(): + try: + await self._call_dp( + "clear_samples", + sample_ids=list(request.canonical_sample_ids), + partition_id=self._partition_id, + ) + except Exception as error: + errors.append( + RuntimeError( + "pre-publication canonical cleanup failed for " + f"group={request.group_id!r}, " + f"ids={request.canonical_sample_ids!r}" + ) + ) + errors[-1].__cause__ = error + staging_keys = self._request_staging_keys(request) + if staging_keys: try: - await call_data_plane( - self._dp_client, + await self._call_dp( "clear_samples", - offload_sync=True, - sample_ids=list(request.rollout_ids), - partition_id=self._partition_id, + sample_ids=staging_keys, + partition_id=self._master_config.token_capture.staging_partition, ) except Exception as error: errors.append( RuntimeError( - "pre-publication canonical cleanup failed for " - f"group={request.group_id!r}, ids={request.rollout_ids!r}" + "pre-publication staging cleanup failed for " + f"group={request.group_id!r}, keys={staging_keys!r}" ) ) errors[-1].__cause__ = error - staging_keys = self._request_staging_keys(request) - if staging_keys: - try: - await call_data_plane( - self._dp_client, - "clear_samples", - offload_sync=True, - sample_ids=staging_keys, - partition_id=self._master_config.token_capture.staging_partition, - ) - except Exception as error: - errors.append( - RuntimeError( - "pre-publication staging cleanup failed for " - f"group={request.group_id!r}, keys={staging_keys!r}" - ) - ) - errors[-1].__cause__ = error if errors: raise BaseExceptionGroup( f"known-outcome cleanup failed for group {request.group_id}", errors, ) self._buffer.abort(request.group_id) + if request.group_id in self._rollout_recovery_ledger: + self._rollout_recovery_ledger.discard_group(cut, request.group_id) + + async def _cleanup_known_finalization_request( + self, request: "ReassemblyRequest" + ) -> None: + """Clear a known request outcome without racing a native TQ snapshot.""" + async with self._data_plane_checkpoint_barrier.mutation() as cut: + await self._cleanup_known_finalization_request_unlocked(cut, request) async def _finalize_with_actor( self, request: "ReassemblyRequest" ) -> Optional["FinalizedGroup"]: - """Submit one metadata request to the bounded fixed actor pool. + """Finalize and index one group atomically with respect to TQ saves. Returns the committed FinalizedGroup once the group is committed to the replay buffer (callers may read valid_row_count/total_row_count @@ -1252,66 +1256,93 @@ async def _finalize_with_actor( self._active_finalizers += 1 active_actor_count = self._active_finalizers finalize_start = time.perf_counter() + rpc_submitted = False + actor_reusable = False try: - finalized = await actor.finalize.remote(request) - except BaseException: - self._finalizer_unknown_outcomes += 1 - print( - "FATAL: finalizer actor RPC failed after submission; canonical " - f"publication outcome is unknown for group {request.group_id}. " - "Stopping validation without actor replacement or retry.", - flush=True, - ) - raise - else: - self._available_finalizers.put_nowait(actor) + # The actor publishes canonical rows before returning metadata. Keep + # the remote write, local replay-index update, and lineage hand-off in + # one mutation cut so a TQ snapshot sees all of them or none of them. + async with self._data_plane_checkpoint_barrier.mutation() as cut: + ledger = self._rollout_recovery_ledger + ledger.mark_finalization_started(request.group_id) + try: + rpc_submitted = True + finalized = await actor.finalize.remote(request) + except BaseException: + self._finalizer_unknown_outcomes += 1 + ledger.mark_finalization_unknown(request.group_id) + print( + "FATAL: finalizer actor RPC failed after submission; canonical " + f"publication outcome is unknown for group {request.group_id}. " + "Stopping validation without actor replacement or retry.", + flush=True, + ) + raise + else: + actor_reusable = True + + if finalized.dropped: + try: + await self._cleanup_known_finalization_request_unlocked( + cut, request + ) + except BaseException as cleanup_error: + raise RuntimeError( + "finalizer dropped the group and known-key cleanup failed " + f"for group {request.group_id}" + ) from cleanup_error + print( + f" finalize: group {request.group_id} dropped " + f"({finalized.drop_reason or 'unspecified reason'})", + flush=True, + ) + committed = False + elif finalized.meta is None: + try: + await self._cleanup_known_finalization_request_unlocked( + cut, request + ) + except BaseException as cleanup_error: + raise RuntimeError( + "finalizer returned no metadata and known-key cleanup " + f"failed for group {request.group_id}" + ) from cleanup_error + raise RuntimeError( + "finalizer returned no metadata for non-dropped group " + f"{request.group_id}" + ) + else: + try: + await self._buffer.commit_finalized( + request.group_id, + finalized.meta, + finalized.group_min_wv, + finalized.group_max_wv, + staging_keys=finalized.staging_keys, + ) + except BaseException as commit_error: + try: + await self._cleanup_known_finalization_request_unlocked( + cut, request + ) + except BaseException as cleanup_error: + raise BaseExceptionGroup( + "finalizer commit and known-key cleanup failed for " + f"group {request.group_id}", + [commit_error, cleanup_error], + ) + raise + # Canonical TQ rows plus replay metadata now own the completed + # group; keep only unfinished work in the lineage sidecar. + ledger.discard_group(cut, request.group_id) + committed = True finally: self._active_finalizers -= 1 + if actor_reusable or not rpc_submitted: + self._available_finalizers.put_nowait(actor) finalize_total_ms = (time.perf_counter() - finalize_start) * 1000.0 - - if finalized.dropped: - try: - await self._cleanup_known_finalization_request(request) - except BaseException as cleanup_error: - raise RuntimeError( - "finalizer dropped the group and known-key cleanup failed " - f"for group {request.group_id}" - ) from cleanup_error - print( - f" finalize: group {request.group_id} dropped " - f"({finalized.drop_reason or 'unspecified reason'})", - flush=True, - ) + if not committed: return None - if finalized.meta is None: - try: - await self._cleanup_known_finalization_request(request) - except BaseException as cleanup_error: - raise RuntimeError( - "finalizer returned no metadata and known-key cleanup failed " - f"for group {request.group_id}" - ) from cleanup_error - raise RuntimeError( - f"finalizer returned no metadata for non-dropped group {request.group_id}" - ) - try: - await self._buffer.commit_finalized( - request.group_id, - finalized.meta, - finalized.group_min_wv, - finalized.group_max_wv, - staging_keys=finalized.staging_keys, - ) - except BaseException as commit_error: - try: - await self._cleanup_known_finalization_request(request) - except BaseException as cleanup_error: - raise BaseExceptionGroup( - f"finalizer commit and known-key cleanup failed for " - f"group {request.group_id}", - [commit_error, cleanup_error], - ) - raise finalized.metrics.update( { "finalize/queue_wait_ms": queue_wait_ms, @@ -1323,15 +1354,10 @@ async def _finalize_with_actor( self._finalizer_metrics_by_group[request.group_id] = dict(finalized.metrics) return finalized - async def _cleanup_consumed_metas(self, metas: list[KVBatchMeta]) -> None: - """Clear canonical rows and full-manifest staging keys after train success. - - Runs after ``finish_train_step`` because policy workers read the staged - capture deltas during training. One mutation cut spans both partitions - so a data-plane checkpoint sees either all of a step's rows or none, and - the clears are offloaded so the event loop keeps serving the rollout - pump and finalizer handoffs meanwhile. - """ + async def _cleanup_consumed_metas_unlocked( + self, metas: list[KVBatchMeta] + ) -> None: + """Clear consumed ownership while holding a barrier mutation slot.""" canonical_by_partition: dict[str, list[str]] = {} staging_by_partition: dict[str, list[str]] = {} for meta in metas: @@ -1348,31 +1374,33 @@ async def _cleanup_consumed_metas(self, metas: list[KVBatchMeta]) -> None: ) errors: list[BaseException] = [] - async with self._data_plane_checkpoint_barrier.mutation(): - for label, by_partition in ( - ("canonical", canonical_by_partition), - ("staging", staging_by_partition), - ): - for partition_id, ids in by_partition.items(): - unique_ids = list(dict.fromkeys(ids)) - try: - await call_data_plane( - self._dp_client, - "clear_samples", - offload_sync=True, - sample_ids=unique_ids, - partition_id=partition_id, - ) - except Exception as error: - cleanup_error = RuntimeError( - f"post-train {label} cleanup failed: " - f"partition={partition_id!r}, ids={unique_ids!r}" - ) - cleanup_error.__cause__ = error - errors.append(cleanup_error) + for label, by_partition in ( + ("canonical", canonical_by_partition), + ("staging", staging_by_partition), + ): + for partition_id, ids in by_partition.items(): + unique_ids = list(dict.fromkeys(ids)) + try: + await self._call_dp( + "clear_samples", + sample_ids=unique_ids, + partition_id=partition_id, + ) + except Exception as error: + cleanup_error = RuntimeError( + f"post-train {label} cleanup failed: " + f"partition={partition_id!r}, ids={unique_ids!r}" + ) + cleanup_error.__cause__ = error + errors.append(cleanup_error) if errors: raise BaseExceptionGroup("post-train DataPlane cleanup failed", errors) + async def _cleanup_consumed_metas(self, metas: list[KVBatchMeta]) -> None: + """Clear consumed rows without racing a native TQ checkpoint.""" + async with self._data_plane_checkpoint_barrier.mutation(): + await self._cleanup_consumed_metas_unlocked(metas) + # ── the three pumps + the inline advantage stage ─────────────────────── async def _rollout_pump(self) -> None: @@ -1432,13 +1460,19 @@ async def _dispatch_one_prompt( ownership_transferred = False try: while True: - request = ( - await self._rollout_manager.generate_for_finalization( + if lineage_group_id is None: + request = await self._rollout_manager.generate_for_finalization( prompt, target_step=target_step, inflight_registry=self._inflight_by_group_id, ) - ) + else: + request = await self._rollout_manager.generate_for_finalization( + prompt, + target_step=target_step, + inflight_registry=self._inflight_by_group_id, + lineage_group_id=lineage_group_id, + ) if not inflight_count_released: self._inflight_rollouts -= 1 inflight_count_released = True @@ -1473,6 +1507,7 @@ async def _dispatch_one_prompt( < min_valid_fraction ) if not below_threshold: + self._rollout_manager.stats.committed += 1 ownership_transferred = True break # Enough rows verified to publish, but too few to @@ -1507,6 +1542,7 @@ async def _dispatch_one_prompt( return replacements += 1 prompt = replacement + lineage_group_id = None # A substitution is a fresh rollout, not a # continuation of the one that fell short, so it # observes the same pause a first dispatch does diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index e5b5286fbab..7b5443b119f 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -17,6 +17,7 @@ import enum import json import uuid +from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Optional @@ -55,6 +56,8 @@ from nemo_rl.experience.metric_utils import calculate_single_metric, pct from nemo_rl.experience.rollout_recovery import ( PromptGroupPhase, + PromptGroupStatus, + RolloutAttemptStatus, RolloutRecoveryLedger, ) from nemo_rl.experience.rollouts import ( @@ -78,6 +81,7 @@ from nemo_rl.utils.timer import Timer TokenizerType = PreTrainedTokenizerBase +RolloutCompletionCallback = Callable[[int, Completion], Awaitable[None]] if TYPE_CHECKING: from nemo_rl.experience.rollout_reassembler_actor import ReassemblyRequest @@ -412,7 +416,12 @@ def __init__( self._timeouts = timeouts async def run_rollout( - self, input_sample: DatumSpec, *, rollout_ids: Optional[list[str]] = None + self, + input_sample: DatumSpec, + *, + rollout_ids: Optional[list[str]] = None, + generation_indices: Optional[list[int]] = None, + on_completion: Optional[RolloutCompletionCallback] = None, ) -> PromptGroupRecord: """Run num_generations_per_prompt rollouts for one prompt. @@ -426,6 +435,12 @@ async def run_rollout( assert rollout_ids is None, ( "token capture (rollout_ids) is only supported on the NeMo-Gym path" ) + assert generation_indices is None, ( + "partial sibling dispatch is only supported on the NeMo-Gym path" + ) + assert on_completion is None, ( + "streamed completion callbacks are only supported on the NeMo-Gym path" + ) timer = Timer() timer_prefix = "timing/rollout" timer.start(f"{timer_prefix}/total") @@ -828,7 +843,12 @@ def __init__( self._validate_init_params() async def run_rollout( - self, input_sample: DatumSpec, *, rollout_ids: Optional[list[str]] = None + self, + input_sample: DatumSpec, + *, + rollout_ids: Optional[list[str]] = None, + generation_indices: Optional[list[int]] = None, + on_completion: Optional[RolloutCompletionCallback] = None, ) -> PromptGroupRecord: """Run num_generations_per_prompt rollouts for one prompt. @@ -846,9 +866,16 @@ async def run_rollout( timer_prefix = "timing/rollout" timer.start(f"{timer_prefix}/total") - rollout_inputs = self._build_inputs(input_sample, rollout_ids=rollout_ids) + rollout_inputs = self._build_inputs( + input_sample, + rollout_ids=rollout_ids, + generation_indices=generation_indices, + ) completions, prompt_message_log, rollout_metrics = await self._run_rollouts( - rollout_inputs, timer, timer_prefix + rollout_inputs, + timer, + timer_prefix, + on_completion=on_completion, ) # Token-capture receipt rows carry empty message logs by design — the # canonical row (and any media it needs) is rebuilt by the finalizer @@ -905,7 +932,11 @@ def _validate_init_params(self) -> None: ) def _build_inputs( - self, input_sample: DatumSpec, *, rollout_ids: Optional[list[str]] = None + self, + input_sample: DatumSpec, + *, + rollout_ids: Optional[list[str]] = None, + generation_indices: Optional[list[int]] = None, ) -> list[dict]: """Build N row dicts from input_sample, applying generation config params.""" # Build a template row from the input_sample's extra_env_info, applying generation params. @@ -941,8 +972,19 @@ def _build_inputs( raise ValueError( f"{NEMO_GYM_GROUP_ATTEMPT_KEY} must be a non-negative integer" ) + indices = ( + list(range(self._num_generations_per_prompt)) + if generation_indices is None + else list(generation_indices) + ) + if len(indices) != len(set(indices)) or any( + not 0 <= index < self._num_generations_per_prompt for index in indices + ): + raise ValueError( + "generation_indices must be unique and within the prompt group" + ) rows = [] - for i in range(self._num_generations_per_prompt): + for i in indices: row = copy.deepcopy(template_row) row["_rowidx"] = i row[NEMO_GYM_GROUP_ID_KEY] = group_id @@ -963,6 +1005,7 @@ async def _stream_rows( results: list[Optional[dict]], total_rows: int, timer_prefix: str, + on_completion: Optional[RolloutCompletionCallback] = None, ) -> Optional[dict[str, Any]]: """Dispatch ``pending`` rows and fill their slots in ``results`` as they land. @@ -1002,6 +1045,8 @@ async def _stream_rows( received.add(rowidx) pending_by_rowidx[rowidx]["agent_ref"] = resolved_agent_ref results[rowidx] = result + if on_completion is not None: + await on_completion(rowidx, self._result_to_completion(result)) if timing_metrics is not None: env_timing_metrics = timing_metrics @@ -1012,6 +1057,8 @@ async def _run_rollouts( inputs: list[dict], timer: Timer, timer_prefix: str, + *, + on_completion: Optional[RolloutCompletionCallback] = None, ) -> tuple[list[Completion], LLMMessageLogType, dict[str, Any]]: """Dispatch rows to NeMo-Gym; return completions, prompt, and metrics. @@ -1022,21 +1069,26 @@ async def _run_rollouts( attempts, which is the same shape as the legacy collector's pending-group retry. """ nemo_gym_env = self._task_to_env["nemo_gym"] - total_rows = len(inputs) + if not inputs: + raise ValueError("NeMo-Gym rollout dispatch requires at least one row") + total_rows = self._num_generations_per_prompt # Re-dispatch maps NeMo-Gym's echoed _rowidx back onto the original group, so # the rows must carry the index _build_inputs stamped on them. Checked here # because the alternative is a KeyError several frames deeper. - for position, row in enumerate(inputs): - if row.get("_rowidx") != position: + for row in inputs: + rowidx = row.get("_rowidx") + if not isinstance(rowidx, int) or not 0 <= rowidx < total_rows: raise ValueError( - f"NeMo-Gym input row {position} carries _rowidx=" - f"{row.get('_rowidx')!r}; rows must be stamped with their own " - "position for re-dispatch to preserve ordering" + f"NeMo-Gym input row carries invalid _rowidx={rowidx!r}; " + f"expected an index within {total_rows} generations" ) + expected_indices = [row["_rowidx"] for row in inputs] + if len(expected_indices) != len(set(expected_indices)): + raise ValueError("NeMo-Gym input rows contain duplicate _rowidx values") # Run generation and restore input order as results stream back. with timer.time(f"{timer_prefix}/run_rollouts"): - results: list[dict | None] = [None for _ in inputs] + results: list[dict | None] = [None for _ in range(total_rows)] env_timing_metrics: dict[str, Any] = {} # One deadline for the whole prompt group, re-dispatches included -- it is # the group that has a budget, not each attempt. It also spans the stream @@ -1070,7 +1122,12 @@ async def _run_rollouts( self._stats.record_gym_row_redispatch(len(pending)) try: timing_metrics = await self._stream_rows( - nemo_gym_env, pending, results, total_rows, timer_prefix + nemo_gym_env, + pending, + results, + total_rows, + timer_prefix, + on_completion=on_completion, ) except Exception as error: last_error = error @@ -1085,7 +1142,7 @@ async def _run_rollouts( if timing_metrics is not None: env_timing_metrics = timing_metrics - missing = [i for i, result in enumerate(results) if result is None] + missing = [index for index in expected_indices if results[index] is None] if missing: failure = GymTransportError( "NeMo-Gym rollout stream ended before all rows arrived; missing " @@ -1484,12 +1541,24 @@ def set_weight_version(self, version: int) -> None: self._weight_version = int(version) async def run_rollout( - self, input_sample: DatumSpec, *, rollout_ids: Optional[list[str]] = None + self, + input_sample: DatumSpec, + *, + rollout_ids: Optional[list[str]] = None, + generation_indices: Optional[list[int]] = None, + on_completion: Optional[RolloutCompletionCallback] = None, ) -> PromptGroupRecord: if rollout_ids is None: + assert generation_indices is None + assert on_completion is None # Legacy path: keep the impl call signature byte-identical. return await self._impl.run_rollout(input_sample) - return await self._impl.run_rollout(input_sample, rollout_ids=rollout_ids) + return await self._impl.run_rollout( + input_sample, + rollout_ids=rollout_ids, + generation_indices=generation_indices, + on_completion=on_completion, + ) async def generate_and_push( self, @@ -1761,66 +1830,39 @@ async def generate_for_finalization( *, target_step: Optional[int] = None, inflight_registry: Optional[dict[str, tuple[asyncio.Task[None], int]]] = None, + lineage_group_id: Optional[str] = None, ) -> Optional["ReassemblyRequest"]: - """Run capture generation with the same retry budgets as the legacy path. - - ``generate_and_push`` re-dispatches infrastructure failures onto a - fresh shard up to ``max_infra_attempts``; without the same loop here, a - single gym-side HTTP 500 — typed ``GymTransportError``, INFRA by - definition — killed the whole capture run while the legacy arm absorbed - hundreds of them (job 6544554 died this way after 3 clean steps). - - Each attempt reserves a fresh group and fresh rollout ids, so a failed - attempt's staged rows can never collide with the retry's. Exhaustion - follows the same drop policy as ``generate_and_push``: under the - consecutive-drop budget the prompt is dropped (``None`` — the caller - owns the backpressure permit and the step shortfall), beyond it the - fleet is declared broken via ``RolloutRedispatchExhausted``. There is - no replacement-reserve support on this branch; a dropped prompt always - closes its step short. Data failures follow ``max_data_attempts`` then - re-raise. - """ + """Capture siblings with stable lineage and retry only unfinished work.""" + assert self._tq_buffer is not None, ( + "generate_for_finalization requires tq_buffer to be set at __init__" + ) + recovery_group_id = lineage_group_id + if recovery_group_id is None: + recovery_group_id = self.reserve_prompt_group( + input_sample, + target_step=target_step, + admitted=True, + ) + recovery_group = self._recovery_ledger.get_group(recovery_group_id) + if recovery_group.phase is not PromptGroupPhase.ADMITTED: + raise RuntimeError( + f"lineage group {recovery_group_id!r} must be admitted before dispatch" + ) + if recovery_group.expected_generations != self._num_generations_per_prompt: + raise ValueError( + f"lineage group {recovery_group_id!r} expects " + f"{recovery_group.expected_generations} generation(s), but the " + f"resumed configuration requests {self._num_generations_per_prompt}" + ) policy = self._retry_policy infra_attempts = 0 data_attempts = 0 last_infra_error: Optional[Exception] = None - # Same logical cohort identity as generate_and_push: group_id is minted - # once and survives retries, group_attempt increments per attempt, so - # Gym's genrm_compare can supersede a failed attempt's cohort instead - # of leaving it collecting forever. TQ/ledger ids stay fresh per - # attempt inside _generate_for_finalization_attempt — physical rows - # must never alias across attempts. - logical_group_id: Optional[str] = None - group_attempt = 0 - extra_env_info = input_sample.get("extra_env_info") - if isinstance(extra_env_info, dict): - configured_group_id = extra_env_info.get(NEMO_GYM_GROUP_ID_KEY) - if configured_group_id is not None and ( - not isinstance(configured_group_id, str) or not configured_group_id - ): - raise ValueError(f"{NEMO_GYM_GROUP_ID_KEY} must be a non-empty string") - logical_group_id = configured_group_id or uuid.uuid4().hex - configured_group_attempt = extra_env_info.get(NEMO_GYM_GROUP_ATTEMPT_KEY, 0) - if ( - not isinstance(configured_group_attempt, int) - or isinstance(configured_group_attempt, bool) - or configured_group_attempt < 0 - ): - raise ValueError( - f"{NEMO_GYM_GROUP_ATTEMPT_KEY} must be a non-negative integer" - ) - group_attempt = configured_group_attempt while infra_attempts < policy.max_infra_attempts: try: - attempt_input_sample = input_sample - if logical_group_id is not None: - attempt_input_sample = copy.deepcopy(input_sample) - attempt_extra_env_info = attempt_input_sample["extra_env_info"] - attempt_extra_env_info[NEMO_GYM_GROUP_ID_KEY] = logical_group_id - attempt_extra_env_info[NEMO_GYM_GROUP_ATTEMPT_KEY] = group_attempt request = await self._generate_for_finalization_attempt( - attempt_input_sample, - target_step=target_step, + input_sample, + recovery_group_id=recovery_group_id, inflight_registry=inflight_registry, ) except Exception as error: @@ -1832,20 +1874,15 @@ async def generate_for_finalization( break self._stats.record_redispatch(reason) await asyncio.sleep(policy.backoff_for(infra_attempts)) - group_attempt += 1 continue + data_attempts += 1 if data_attempts >= policy.max_data_attempts: + self._stats.record_data_failure(reason) raise - print( - f"retrying capture rollout idx={input_sample['idx']} after " - f"deterministic failure ({reason}: {error})", - flush=True, - ) - group_attempt += 1 + self._stats.record_data_retry(reason) continue - # Same placement as generate_and_push: a successful dispatch proves - # the fleet is answering, clearing the consecutive-drop run. + self._consecutive_infra_drops = 0 return request @@ -1855,22 +1892,13 @@ async def generate_for_finalization( if self._consecutive_infra_drops > policy.max_consecutive_dropped_prompts: raise RolloutRedispatchExhausted( f"prompt idx={input_sample['idx']} exhausted its infrastructure " - f"retry budget after {infra_attempts} capture attempt(s) " - f"(max_infra_attempts_per_prompt={policy.max_infra_attempts}), and " - f"this was drop {self._consecutive_infra_drops} with no rollout " - f"committed in between, exceeding max_consecutive_dropped_prompts=" - f"{policy.max_consecutive_dropped_prompts}; the generation fleet is " - f"not recovering. Last failure was {reason}: {last_infra_error}" + f"retry budget after {infra_attempts} capture attempt(s) and this " + f"was drop {self._consecutive_infra_drops}, exceeding " + f"max_consecutive_dropped_prompts=" + f"{policy.max_consecutive_dropped_prompts}; last failure was " + f"{reason}: {last_infra_error}" ) from last_infra_error self._stats.record_infra_drop(reason, self._consecutive_infra_drops) - print( - f"dropping capture prompt idx={input_sample['idx']} after " - f"{infra_attempts} infrastructure failure(s) ({reason}: " - f"{last_infra_error}) [consecutive drop " - f"{self._consecutive_infra_drops}/" - f"{policy.max_consecutive_dropped_prompts}]", - flush=True, - ) self._stats.skipped += 1 return None @@ -1878,66 +1906,116 @@ async def _generate_for_finalization_attempt( self, input_sample: DatumSpec, *, - target_step: Optional[int], + recovery_group_id: str, inflight_registry: Optional[dict[str, tuple[asyncio.Task[None], int]]], ) -> "ReassemblyRequest": - """One capture-generation attempt; the retry loop above owns budgets. - - The replay-buffer slot remains reserved and unready. The caller owns - finalizer submission and must either commit the returned group or stop - the validation run on an unknown publication outcome. - """ + """Dispatch only unfinished siblings and leave one reserved slot unready.""" from nemo_rl.experience.rollout_reassembler_actor import ReassemblyRequest - assert self._tq_buffer is not None, ( - "generate_for_finalization requires tq_buffer to be set at __init__" - ) - start_version = self._weight_version - group_id = str(uuid.uuid4()) - rollout_ids = tuple( - f"{group_id}_g{i}" for i in range(self._num_generations_per_prompt) - ) + assert self._tq_buffer is not None + recovery_group = self._recovery_ledger.get_group(recovery_group_id) + if recovery_group.status == PromptGroupStatus.GENERATING: + recovery_group = self._recovery_ledger.prepare_incomplete_retry( + recovery_group_id + ) + pending_indices = [ + sibling.generation_index + for sibling in recovery_group.siblings + if sibling.current_attempt.status != RolloutAttemptStatus.SEALED + ] + group_id = recovery_group.group_id + start_version = recovery_group.start_weight_version + rollout_ids = tuple(recovery_group.gate_rollout_ids) + attempt_input_sample = copy.deepcopy(input_sample) + attempt_extra_env_info = attempt_input_sample.get("extra_env_info") + if isinstance(attempt_extra_env_info, dict): + attempt_extra_env_info[NEMO_GYM_GROUP_ID_KEY] = group_id + attempt_extra_env_info[NEMO_GYM_GROUP_ATTEMPT_KEY] = max( + len(sibling.attempts) for sibling in recovery_group.siblings + ) - 1 self._tq_buffer.reserve( weight_version=start_version, - target_step=target_step, + target_step=recovery_group.target_step, group_id=group_id, rollout_ids=list(rollout_ids), ) + + async def _record_streamed_completion( + generation_index: int, completion: Completion + ) -> None: + env_extras = completion.env_extras + if env_extras is None: + raise ValueError( + "token-capture completion must contain environment extras" + ) + receipt = env_extras.get("ng_receipt") + gate_rollout_id = env_extras.get("ng_rollout_id") + if not isinstance(receipt, dict): + raise ValueError( + "token-capture completion must contain a receipt mapping" + ) + if not isinstance(gate_rollout_id, str): + raise ValueError( + "token-capture completion must contain its Gate rollout ID" + ) + self._recovery_ledger.mark_sibling_sealed( + group_id, + generation_index=generation_index, + gate_rollout_id=gate_rollout_id, + receipt=receipt, + reward=completion.reward, + ) + + mask_sample_by_index: dict[int, bool] = {} + + async def _record_completion( + generation_index: int, completion: Completion + ) -> None: + mask_sample_by_index[generation_index] = bool( + (((completion.env_extras or {}).get("instance_config") or {}).get( + MASK_SAMPLE, False + )) + ) + await _record_streamed_completion(generation_index, completion) + try: if inflight_registry is not None: current_task = asyncio.current_task() assert current_task is not None inflight_registry[group_id] = (current_task, start_version) try: - record = await self.run_rollout( - input_sample, - rollout_ids=list(rollout_ids), - ) + if pending_indices: + self._recovery_ledger.mark_group_dispatched( + group_id, generation_indices=pending_indices + ) + await self.run_rollout( + attempt_input_sample, + rollout_ids=list(rollout_ids), + generation_indices=pending_indices, + on_completion=_record_completion, + ) finally: if inflight_registry is not None: inflight_registry.pop(group_id, None) - receipts = tuple(c.env_extras.get("ng_receipt") for c in record.completions) - rewards = tuple(float(c.reward) for c in record.completions) - # Same read as the token path's ``_mask_sample_flags``; the impl - # already applied the ``mask_env_flagged_samples`` gate by popping - # the flag from ``instance_config`` when masking is off. - mask_sample = tuple( - bool( - ((c.env_extras or {}).get("instance_config") or {}).get( - MASK_SAMPLE, False - ) - ) - for c in record.completions - ) + ( + physical_rollout_ids, + canonical_sample_ids, + receipts, + rewards, + ) = self._recovery_ledger.finalization_inputs(group_id) request = ReassemblyRequest( group_id=group_id, - rollout_ids=rollout_ids, - receipts=receipts, - rewards=rewards, + rollout_ids=tuple(physical_rollout_ids), + canonical_sample_ids=tuple(canonical_sample_ids), + receipts=tuple(receipts), + rewards=tuple(rewards), fallback_weight_version=start_version, - prompt_idx=record.prompt_idx, - mask_sample=mask_sample, - loss_multiplier=record.loss_multiplier, + prompt_idx=int(recovery_group.prompt_id), + mask_sample=tuple( + mask_sample_by_index.get(index, False) + for index in range(len(receipts)) + ), + loss_multiplier=float(input_sample.get("loss_multiplier", 1.0)), ) from nemo_rl.experience.rollout_reassembler_actor import ( assert_metadata_only, @@ -1952,4 +2030,20 @@ async def _generate_for_finalization_attempt( # yet). Their ledger files are inert — failure rows or missing # terminal rows keep any later read fail-closed. self._tq_buffer.abort(group_id) + self._recovery_ledger.abandon_unsealed(group_id) + # The capture ledger has no per-rollout fail endpoint. Rows from + # abandoned attempts are unreferenced and are swept with the + # staging partition at run teardown. raise + + async def _discard_recovery_group(self, group_id: str) -> None: + """Clean known staged rows before intentionally dropping lineage.""" + assert self._tq_buffer is not None + group = self._recovery_ledger.get_group(group_id) + staging_keys = [ + key + for sibling in group.siblings + for key in sibling.current_attempt.staging_keys + ] + await self._tq_buffer.clear_staging_keys(staging_keys) + self._recovery_ledger.discard_group(group_id) diff --git a/nemo_rl/experience/rollout_reassembler.py b/nemo_rl/experience/rollout_reassembler.py index 2b5818c8a2d..06b32e7bc99 100644 --- a/nemo_rl/experience/rollout_reassembler.py +++ b/nemo_rl/experience/rollout_reassembler.py @@ -365,6 +365,7 @@ def finalize_group( fallback_weight_version: int, prompt_idx: int, loss_multiplier: float = 1.0, + canonical_sample_ids: Optional[list[str]] = None, ) -> FinalizedGroup: """Publish exactly N canonical rows for one prompt group. @@ -385,6 +386,11 @@ def finalize_group( assert len(rollout_ids) == len(receipts) == len(rewards) == len(mask_sample), ( "rollout_ids, receipts, rewards, and mask_sample must be parallel" ) + if canonical_sample_ids is None: + canonical_sample_ids = rollout_ids + assert len(canonical_sample_ids) == len(rollout_ids), ( + "canonical_sample_ids must be one per rollout" + ) _group_t0 = time.perf_counter() rows = [ self.finalize_rollout(rollout_id, receipt, reward=reward) @@ -594,9 +600,9 @@ def finalize_group( metrics["finalize/routed_experts_row_coverage"] = ( valid_route_rows / len(valid_rows) ) - assert sample_ids == rollout_ids, ( - "canonical sample ids must equal the ledger-registered rollout ids: " - f"{sample_ids} != {rollout_ids}" + assert sample_ids == canonical_sample_ids, ( + "canonical sample ids must equal the stable logical rollout ids: " + f"{sample_ids} != {canonical_sample_ids}" ) _tensorize_ms = (time.perf_counter() - _tensorize_t0) * 1000.0 _put_t0 = time.perf_counter() diff --git a/nemo_rl/experience/rollout_reassembler_actor.py b/nemo_rl/experience/rollout_reassembler_actor.py index 57d91953d29..d22923f5caf 100644 --- a/nemo_rl/experience/rollout_reassembler_actor.py +++ b/nemo_rl/experience/rollout_reassembler_actor.py @@ -53,6 +53,7 @@ class ReassemblyRequest: group_id: str rollout_ids: tuple[str, ...] + canonical_sample_ids: tuple[str, ...] receipts: tuple[Optional[dict[str, Any]], ...] rewards: tuple[float, ...] fallback_weight_version: int @@ -138,13 +139,14 @@ def finalize(self, request: ReassemblyRequest) -> FinalizedGroup: assert_metadata_only(request) if not ( len(request.rollout_ids) + == len(request.canonical_sample_ids) == len(request.receipts) == len(request.rewards) == len(request.mask_sample) ): raise ValueError( - "finalizer request rollout_ids, receipts, rewards, and " - "mask_sample must be parallel" + "finalizer request rollout_ids, canonical_sample_ids, receipts, " + "rewards, and mask_sample must be parallel" ) result = self._finalizer.finalize_group( request.group_id, @@ -155,6 +157,7 @@ def finalize(self, request: ReassemblyRequest) -> FinalizedGroup: fallback_weight_version=request.fallback_weight_version, prompt_idx=request.prompt_idx, loss_multiplier=request.loss_multiplier, + canonical_sample_ids=list(request.canonical_sample_ids), ) assert_metadata_only(result) return result diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index 98afc450ecb..3bc68e55b81 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -12,90 +12,141 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Versioned ownership state for unfinished SingleController prompt groups.""" +"""Controller-owned lineage for recoverable token-capture prompt groups. + +The ledger deliberately contains control-plane metadata only. Token tensors and +router-replay payloads remain in TQ. Persistence is added by a later change; the +versioned ``state_dict`` boundary lives here so that change does not have to +invent a second lifecycle model. +""" from __future__ import annotations import copy import uuid -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import StrEnum -from typing import TYPE_CHECKING, Any, NotRequired, TypedDict +from typing import TYPE_CHECKING, Any, Optional, Self, TypeAlias + +from nemo_rl.data_plane import KVBatchMeta if TYPE_CHECKING: from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneMutationCut from nemo_rl.data.interfaces import DatumSpec -ROLLOUT_RECOVERY_SCHEMA_VERSION = 1 +ROLLOUT_RECOVERY_SCHEMA_VERSION = 2 ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" +RolloutRecoveryState: TypeAlias = dict[str, Any] class PromptGroupPhase(StrEnum): - """Durable admission phase for an unfinished prompt group.""" + """Durable sampler-admission phase for an unfinished prompt group.""" RESERVED = "reserved" ADMITTED = "admitted" -class PromptRefState(TypedDict): - """Serializable locator for rebuilding one prompt from the dataset.""" - - sample_id: str - task_name: str | None - - -class PromptGroupRecoveryState(TypedDict): - """Serializable ownership state for one unfinished prompt group.""" +class RolloutAttemptStatus(StrEnum): + """Lifecycle of one physical Gate execution attempt.""" - group_id: str - admission_id: str - prompt_id: str - prompt_ref: PromptRefState - expected_generations: int - target_step: int | None - start_weight_version: int - phase: str + RESERVED = "reserved" + DISPATCHED = "dispatched" + SEALED = "sealed" + FAILED = "failed" + ABANDONED = "abandoned" -class RolloutRecoveryLedgerState(TypedDict): - """Versioned prompt-group ownership state managed by the ledger.""" +class PromptGroupStatus(StrEnum): + """Ownership lifecycle of one logical prompt group.""" - schema_version: int - groups: list[PromptGroupRecoveryState] + GENERATING = "generating" + READY_TO_FINALIZE = "ready_to_finalize" + FINALIZING = "finalizing" + FINALIZATION_UNKNOWN = "finalization_unknown" + FINALIZED = "finalized" + CLAIMED_FOR_TRAINING = "claimed_for_training" + APPLIED_UNCHECKPOINTED = "applied_uncheckpointed" -class RolloutRecoveryState(RolloutRecoveryLedgerState): - """Complete checkpoint sidecar for unfinished rollout scheduling state.""" +class TrainStepStatus(StrEnum): + """State of the one optimizer step the SingleController may have open.""" - batch_shortfall: NotRequired[dict[int, int]] - sampler_stamps_target_steps: NotRequired[bool] + OPEN = "open" + APPLIED_UNCHECKPOINTED = "applied_uncheckpointed" @dataclass(frozen=True) class PromptRef: - """Stable dataset identity for rebuilding one prompt.""" + """Small durable locator for a prompt owned by the input dataset. + + The persistence layer will resolve this reference and validate the dataset + identity before redispatch. The full ``DatumSpec`` is runtime-only state and + is deliberately excluded from the serialized ledger. + """ sample_id: str - task_name: str | None + task_name: Optional[str] = None + def __post_init__(self) -> None: + if not self.sample_id: + raise ValueError("prompt sample_id must not be empty") -@dataclass(frozen=True) + +@dataclass +class RolloutAttemptRecord: + """One physical attempt for a stable logical sibling.""" + + attempt_uuid: uuid.UUID + status: RolloutAttemptStatus + receipt: Optional[dict[str, Any]] = None + reward: Optional[float] = None + staging_keys: list[str] = field(default_factory=list) + + @property + def attempt_id(self) -> str: + """Return the compact external representation of this attempt UUID.""" + return self.attempt_uuid.hex + + +@dataclass +class RolloutSiblingRecord: + """One stable GRPO generation slot and its physical attempts.""" + + generation_index: int + attempts: list[RolloutAttemptRecord] + + @property + def current_attempt(self) -> RolloutAttemptRecord: + if not self.attempts: + raise RuntimeError( + f"generation index {self.generation_index} has no attempts" + ) + return self.attempts[-1] + + +@dataclass class PromptGroupRecoveryRecord: - """In-memory ownership record for one prompt group.""" + """Lineage and ownership for one logical prompt group.""" group_id: str admission_id: str prompt_id: str prompt_ref: PromptRef - runtime_prompt_payload: DatumSpec | None + runtime_prompt_payload: Optional[DatumSpec] expected_generations: int - target_step: int | None + target_step: Optional[int] start_weight_version: int + siblings: list[RolloutSiblingRecord] phase: PromptGroupPhase + status: PromptGroupStatus = PromptGroupStatus.GENERATING + canonical_meta: Optional[KVBatchMeta] = None + group_min_weight_version: Optional[int] = None + group_max_weight_version: Optional[int] = None + claimed_train_step: Optional[int] = None @property def prompt_payload(self) -> DatumSpec: - """Return the rehydrated prompt required for rollout redispatch.""" + """Return the runtime prompt required to redispatch unfinished work.""" if self.runtime_prompt_payload is None: raise RuntimeError( f"recovery group {self.group_id!r} has not rehydrated prompt " @@ -103,66 +154,91 @@ def prompt_payload(self) -> DatumSpec: ) return self.runtime_prompt_payload + @property + def logical_rollout_ids(self) -> list[str]: + return [ + self.logical_rollout_id(sibling.generation_index) + for sibling in self.siblings + ] -@dataclass(frozen=True) -class ParsedRolloutRecoveryState: - """Validated controller and ledger state loaded from one checkpoint sidecar.""" + @property + def gate_rollout_ids(self) -> list[str]: + return [ + self.gate_rollout_id(sibling.generation_index) for sibling in self.siblings + ] + + def logical_rollout_id(self, generation_index: int) -> str: + """Derive the stable sibling ID instead of storing another UUID string.""" + return f"{self.group_id}_g{generation_index}" + + def gate_rollout_id(self, generation_index: int) -> str: + """Derive the physical Gate ID from group, sibling, and attempt UUID.""" + sibling = self.siblings[generation_index] + return ( + f"{self.logical_rollout_id(generation_index)}" + f"_a{sibling.current_attempt.attempt_id}" + ) - ledger_state: RolloutRecoveryLedgerState - batch_shortfall: dict[int, int] - sampler_stamps_target_steps: bool | None + @property + def sealed_generation_indices(self) -> list[int]: + return [ + sibling.generation_index + for sibling in self.siblings + if sibling.current_attempt.status == RolloutAttemptStatus.SEALED + ] -def _require_int(value: Any, *, field: str, minimum: int) -> int: - """Validate one integer field without accepting booleans.""" - if isinstance(value, bool) or not isinstance(value, int) or value < minimum: - raise ValueError(f"{field} must be an integer >= {minimum}, got {value!r}") - return value +@dataclass +class OpenTrainStepRecord: + """Groups contributing to the current, not-yet-durable optimizer step.""" + train_step: int + trainer_version: int + expected_group_count: int + group_ids: list[str] = field(default_factory=list) + status: TrainStepStatus = TrainStepStatus.OPEN -def _prompt_task_name(prompt_payload: DatumSpec) -> str | None: - task_name = prompt_payload.get("task_name") - if task_name is not None and not isinstance(task_name, str): - raise TypeError( - "prompt_payload.task_name must be a string or None, got " - f"{type(task_name).__name__}" - ) - return task_name - - -def _validate_prompt_identity( - prompt_ref: PromptRef, - prompt_payload: DatumSpec, - *, - group_id: str, -) -> None: - sample_id = prompt_payload.get("idx") - if isinstance(sample_id, bool) or not isinstance(sample_id, int): - raise ValueError( - f"recovery group {group_id!r} prompt payload must contain an integer idx" - ) - if str(sample_id) != prompt_ref.sample_id: - raise ValueError( - f"recovery group {group_id!r} resolved sample_id={sample_id!r}; " - f"expected {prompt_ref.sample_id!r}" - ) - task_name = _prompt_task_name(prompt_payload) - if task_name != prompt_ref.task_name: - raise ValueError( - f"recovery group {group_id!r} resolved task_name={task_name!r}; " - f"expected {prompt_ref.task_name!r}" - ) + +def _new_attempt() -> RolloutAttemptRecord: + return RolloutAttemptRecord( + attempt_uuid=uuid.uuid4(), + status=RolloutAttemptStatus.RESERVED, + ) + + +def _receipt_staging_keys(receipt: dict[str, Any]) -> list[str]: + """Validate a sealed Gate receipt and return its ordered staging keys.""" + manifest = receipt.get("manifest") + if not isinstance(manifest, list): + raise ValueError("sealed rollout receipt must contain a manifest list") + staging_keys: list[str] = [] + for entry in manifest: + if not isinstance(entry, dict) or not isinstance(entry.get("staging_key"), str): + raise ValueError( + "sealed rollout receipt manifest entries must contain string " + "staging_key values" + ) + staging_keys.append(entry["staging_key"]) + return staging_keys class RolloutRecoveryLedger: - """Own prompts after dataloader advance and before canonical TQ commit. + """In-memory source of truth for token-capture rollout ownership. - Every mutating operation requires a live data-plane cut so ownership cannot + Controller-owned mutations require a live data-plane cut so lineage cannot change outside the checkpoint barrier's consistent snapshot boundary. """ def __init__(self) -> None: self._groups: dict[str, PromptGroupRecoveryRecord] = {} + self._open_train_step: Optional[OpenTrainStepRecord] = None + + @property + def open_train_step(self) -> Optional[OpenTrainStepRecord]: + return copy.deepcopy(self._open_train_step) + + def groups(self) -> list[PromptGroupRecoveryRecord]: + return [self._copy_group(group) for group in self._groups.values()] def reserve_group( self, @@ -171,165 +247,411 @@ def reserve_group( prompt_id: str, prompt_payload: DatumSpec, expected_generations: int, - target_step: int | None, + target_step: Optional[int], start_weight_version: int, - admitted: bool, - group_id: str | None = None, - admission_id: str | None = None, + admitted: bool = True, + group_id: Optional[str] = None, + admission_id: Optional[str] = None, + prompt_ref: Optional[PromptRef] = None, ) -> PromptGroupRecoveryRecord: - """Record ownership before the prompt can disappear from the dataloader. - - Args: - cut: Live capability yielded by the shared data-plane barrier. - prompt_id: Dataset-level prompt identity used for diagnostics. - prompt_payload: Runtime prompt used for whole-group regeneration. Only - its stable dataset reference is checkpointed. - expected_generations: Number of GRPO siblings in the prompt group. - target_step: Original gated training step, when the sampler stamps one. - start_weight_version: Policy version visible at reservation time. - admitted: Whether sampler admission already completed. This is explicit - because ``target_step=None`` is also valid for admitted ungated groups. - group_id: Stable logical and canonical TQ group ID. Generated when absent. - admission_id: Stable identity shared by every prompt in one sampler - admission. Defaults to ``group_id`` for single-prompt direct callers. - - Returns: - A defensive copy of the new record. - """ + """Create one logical group and its first physical sibling attempts.""" cut.require_live() if not prompt_id: raise ValueError("prompt_id must not be empty") - sample_id = prompt_payload.get("idx") - if isinstance(sample_id, bool) or not isinstance(sample_id, int): - raise ValueError("prompt_payload must contain an integer idx") - if prompt_id != str(sample_id): + if prompt_ref is None: + task_name = prompt_payload.get("task_name") + if task_name is not None and not isinstance(task_name, str): + raise TypeError("prompt task_name must be a string or None") + prompt_ref = PromptRef(sample_id=prompt_id, task_name=task_name) + if prompt_ref.sample_id != prompt_id: raise ValueError( - f"prompt_id={prompt_id!r} does not match prompt_payload idx={sample_id!r}" + "dataset prompt reference must match prompt_id: " + f"{prompt_ref.sample_id!r} != {prompt_id!r}" ) - _require_int( - expected_generations, - field="expected_generations", - minimum=1, - ) - _require_int( - start_weight_version, - field="start_weight_version", - minimum=0, - ) - if target_step is not None: - _require_int(target_step, field="target_step", minimum=0) + if expected_generations < 1: + raise ValueError("expected_generations must be at least one") group_id = group_id or str(uuid.uuid4()) - if not group_id: - raise ValueError("group_id must not be empty") if group_id in self._groups: raise ValueError(f"duplicate recovery group_id={group_id!r}") admission_id = admission_id or group_id if not admission_id: raise ValueError("admission_id must not be empty") + siblings = [] + for generation_index in range(expected_generations): + siblings.append( + RolloutSiblingRecord( + generation_index=generation_index, + attempts=[_new_attempt()], + ) + ) record = PromptGroupRecoveryRecord( group_id=group_id, admission_id=admission_id, prompt_id=prompt_id, - # The rollout path treats the dataloader sample as immutable and builds - # mutable environment inputs from copies. Retaining that sample by - # reference avoids cloning a potentially very long prompt on every - # dispatch; state_dict() persists only its dataset locator. - prompt_ref=PromptRef( - sample_id=prompt_id, - task_name=_prompt_task_name(prompt_payload), - ), + prompt_ref=prompt_ref, + # Retain the immutable dataloader sample by reference instead of copying + # a potentially 131k-token payload. This cache is never serialized and + # is released as soon as canonical rows take over recovery ownership. runtime_prompt_payload=prompt_payload, expected_generations=expected_generations, target_step=target_step, start_weight_version=start_weight_version, + siblings=siblings, phase=( PromptGroupPhase.ADMITTED if admitted else PromptGroupPhase.RESERVED ), ) self._groups[group_id] = record - return copy.copy(record) + return self._copy_group(record) def mark_group_admitted( self, cut: DataPlaneMutationCut, group_id: str, *, - target_step: int | None, + target_step: Optional[int], start_weight_version: int, ) -> None: - """Attach the sampler result to a previously reserved prompt group.""" + """Commit sampler admission without replacing sibling lineage.""" cut.require_live() record = self._require_group(group_id) if record.phase is not PromptGroupPhase.RESERVED: raise ValueError( f"recovery group {group_id!r} is already {record.phase.value}" ) - if target_step is not None: - _require_int(target_step, field="target_step", minimum=0) - _require_int( - start_weight_version, - field="start_weight_version", - minimum=0, - ) - self._groups[group_id] = PromptGroupRecoveryRecord( - group_id=record.group_id, - admission_id=record.admission_id, - prompt_id=record.prompt_id, - prompt_ref=record.prompt_ref, - runtime_prompt_payload=record.runtime_prompt_payload, - expected_generations=record.expected_generations, - target_step=target_step, - start_weight_version=start_weight_version, - phase=PromptGroupPhase.ADMITTED, - ) + record.target_step = target_step + record.start_weight_version = start_weight_version + record.phase = PromptGroupPhase.ADMITTED + + def bind_runtime_prompt(self, group_id: str, prompt_payload: DatumSpec) -> None: + """Attach a dataset-reconstructed prompt after identity validation.""" + record = self._require_group(group_id) + sample_id = prompt_payload.get("idx") + if str(sample_id) != record.prompt_ref.sample_id: + raise ValueError( + f"recovery group {group_id!r} resolved sample_id={sample_id!r}; " + f"expected {record.prompt_ref.sample_id!r}" + ) + task_name = prompt_payload.get("task_name") + if task_name != record.prompt_ref.task_name: + raise ValueError( + f"recovery group {group_id!r} resolved task_name={task_name!r}; " + f"expected {record.prompt_ref.task_name!r}" + ) + record.runtime_prompt_payload = prompt_payload - def bind_runtime_prompt( + def prepare_incomplete_retry(self, group_id: str) -> PromptGroupRecoveryRecord: + """Mint fresh physical attempts only for siblings that are not sealed.""" + record = self._require_group(group_id) + if record.status != PromptGroupStatus.GENERATING: + raise ValueError( + f"cannot retry group {group_id!r} from status {record.status.value!r}" + ) + for sibling in record.siblings: + attempt = sibling.current_attempt + if attempt.status == RolloutAttemptStatus.SEALED: + continue + if attempt.status == RolloutAttemptStatus.RESERVED: + continue + if attempt.status not in { + RolloutAttemptStatus.ABANDONED, + RolloutAttemptStatus.FAILED, + }: + raise ValueError( + "cannot retry logical rollout " + f"{record.logical_rollout_id(sibling.generation_index)!r} " + f"from status {attempt.status.value!r}" + ) + sibling.attempts.append(_new_attempt()) + return self._copy_group(record) + + def mark_group_dispatched( + self, group_id: str, *, generation_indices: Optional[list[int]] = None + ) -> None: + """Move the selected current sibling attempts to dispatched.""" + record = self._require_group(group_id) + if record.phase is not PromptGroupPhase.ADMITTED: + raise ValueError( + f"cannot dispatch unadmitted recovery group {group_id!r}" + ) + if record.status != PromptGroupStatus.GENERATING: + raise ValueError( + f"cannot dispatch group {group_id!r} from {record.status.value!r}" + ) + indices = ( + generation_indices + if generation_indices is not None + else list(range(record.expected_generations)) + ) + attempts = [ + self._require_sibling(record, index).current_attempt for index in indices + ] + if any(attempt.status != RolloutAttemptStatus.RESERVED for attempt in attempts): + raise ValueError("only reserved rollout attempts may be dispatched") + for attempt in attempts: + attempt.status = RolloutAttemptStatus.DISPATCHED + + def mark_sibling_sealed( self, cut: DataPlaneMutationCut, group_id: str, - prompt_payload: DatumSpec, + *, + generation_index: int, + gate_rollout_id: str, + receipt: dict[str, Any], + reward: float, ) -> None: - """Attach a dataset-rehydrated prompt after identity validation. - - The current reference is a positional index into a map-style dataset. - Recovery therefore requires dataset ordering to remain unchanged between - checkpoint and restart. - """ + """Record one streamed sibling receipt as soon as the row arrives.""" cut.require_live() record = self._require_group(group_id) - _validate_prompt_identity( - record.prompt_ref, - prompt_payload, - group_id=group_id, + sibling = self._require_sibling(record, generation_index) + attempt = sibling.current_attempt + expected_gate_rollout_id = record.gate_rollout_id(generation_index) + staging_keys = _receipt_staging_keys(receipt) + if gate_rollout_id != expected_gate_rollout_id: + raise ValueError( + "streamed rollout identity mismatch: " + f"result={gate_rollout_id!r}, expected={expected_gate_rollout_id!r}" + ) + if receipt.get("rollout_id") != gate_rollout_id: + raise ValueError( + "receipt rollout identity mismatch: " + f"receipt={receipt.get('rollout_id')!r}, expected={gate_rollout_id!r}" + ) + if attempt.status == RolloutAttemptStatus.SEALED: + if ( + attempt.receipt == receipt + and attempt.reward == float(reward) + and attempt.staging_keys == staging_keys + ): + return + raise ValueError( + "conflicting duplicate seal for " + f"{record.logical_rollout_id(generation_index)!r}" + ) + if attempt.status != RolloutAttemptStatus.DISPATCHED: + raise ValueError( + "cannot seal logical rollout " + f"{record.logical_rollout_id(generation_index)!r} " + f"from status {attempt.status.value!r}" + ) + + attempt.receipt = copy.deepcopy(receipt) + attempt.reward = float(reward) + attempt.staging_keys = staging_keys + attempt.status = RolloutAttemptStatus.SEALED + if all( + item.current_attempt.status == RolloutAttemptStatus.SEALED + for item in record.siblings + ): + record.status = PromptGroupStatus.READY_TO_FINALIZE + + def abandon_unsealed(self, group_id: str) -> None: + """Abandon failed attempts without destroying reusable sealed receipts.""" + record = self._require_group(group_id) + if record.status not in { + PromptGroupStatus.GENERATING, + PromptGroupStatus.READY_TO_FINALIZE, + }: + raise ValueError( + f"cannot abandon group {group_id!r} from {record.status.value!r}" + ) + for sibling in record.siblings: + attempt = sibling.current_attempt + if attempt.status == RolloutAttemptStatus.SEALED: + continue + attempt.status = RolloutAttemptStatus.ABANDONED + record.status = ( + PromptGroupStatus.READY_TO_FINALIZE + if all( + sibling.current_attempt.status == RolloutAttemptStatus.SEALED + for sibling in record.siblings + ) + else PromptGroupStatus.GENERATING ) - self._groups[group_id] = PromptGroupRecoveryRecord( - group_id=record.group_id, - admission_id=record.admission_id, - prompt_id=record.prompt_id, - prompt_ref=PromptRef( - sample_id=record.prompt_ref.sample_id, - task_name=record.prompt_ref.task_name, - ), - runtime_prompt_payload=prompt_payload, - expected_generations=record.expected_generations, - target_step=record.target_step, - start_weight_version=record.start_weight_version, - phase=record.phase, + + def finalization_inputs( + self, group_id: str + ) -> tuple[list[str], list[str], list[dict[str, Any]], list[float]]: + """Return physical IDs, canonical IDs, receipts and rewards in sibling order.""" + record = self._require_group(group_id) + if record.status != PromptGroupStatus.READY_TO_FINALIZE: + raise ValueError( + f"group {group_id!r} is not ready to finalize: {record.status.value!r}" + ) + receipts: list[dict[str, Any]] = [] + rewards: list[float] = [] + for sibling in record.siblings: + attempt = sibling.current_attempt + if ( + attempt.status != RolloutAttemptStatus.SEALED + or attempt.receipt is None + or attempt.reward is None + ): + raise ValueError( + "logical rollout " + f"{record.logical_rollout_id(sibling.generation_index)!r} " + "is not sealed" + ) + receipts.append(copy.deepcopy(attempt.receipt)) + rewards.append(attempt.reward) + return ( + record.gate_rollout_ids, + record.logical_rollout_ids, + receipts, + rewards, ) - def get_group(self, group_id: str) -> PromptGroupRecoveryRecord: - """Return a record copy while sharing its immutable runtime prompt.""" - return copy.copy(self._require_group(group_id)) + def mark_finalization_started(self, group_id: str) -> None: + record = self._require_group(group_id) + self._require_group_status( + record, + allowed={PromptGroupStatus.READY_TO_FINALIZE}, + transition="start finalization", + ) + record.status = PromptGroupStatus.FINALIZING - def groups(self) -> list[PromptGroupRecoveryRecord]: - """Return record copies in reservation order without cloning prompts.""" - return [copy.copy(record) for record in self._groups.values()] + def mark_finalization_unknown(self, group_id: str) -> None: + record = self._require_group(group_id) + self._require_group_status( + record, + allowed={PromptGroupStatus.FINALIZING}, + transition="mark finalization unknown", + ) + record.status = PromptGroupStatus.FINALIZATION_UNKNOWN + + def mark_group_finalized( + self, + group_id: str, + *, + meta: KVBatchMeta, + group_min_weight_version: int, + group_max_weight_version: int, + ) -> None: + """Transfer recovery ownership from staged receipts to canonical TQ rows.""" + record = self._require_group(group_id) + self._require_group_status( + record, + allowed={PromptGroupStatus.FINALIZING}, + transition="finalize", + ) + if list(meta.sample_ids) != record.logical_rollout_ids: + raise ValueError( + "finalized sample IDs do not match stable logical rollout IDs: " + f"{meta.sample_ids!r} != {record.logical_rollout_ids!r}" + ) + record.canonical_meta = copy.deepcopy(meta) + record.group_min_weight_version = int(group_min_weight_version) + record.group_max_weight_version = int(group_max_weight_version) + record.runtime_prompt_payload = None + record.status = PromptGroupStatus.FINALIZED + + def claim_groups_for_training( + self, + group_ids: list[str], + *, + train_step: int, + trainer_version: int, + expected_group_count: int, + ) -> None: + """Move finalized groups into the controller's one open optimizer step.""" + if not group_ids: + raise ValueError("training claim must contain at least one group") + if len(group_ids) != len(set(group_ids)): + raise ValueError("training claim contains duplicate group IDs") + if self._open_train_step is None: + self._open_train_step = OpenTrainStepRecord( + train_step=train_step, + trainer_version=trainer_version, + expected_group_count=expected_group_count, + ) + open_step = self._open_train_step + if ( + open_step.train_step != train_step + or open_step.trainer_version != trainer_version + or open_step.expected_group_count != expected_group_count + or open_step.status != TrainStepStatus.OPEN + ): + raise ValueError( + "training claim does not match the existing open train step" + ) + records = [self._require_group(group_id) for group_id in group_ids] + for record in records: + if record.status != PromptGroupStatus.FINALIZED: + raise ValueError( + f"cannot claim group {record.group_id!r} from " + f"{record.status.value!r}" + ) + if len(open_step.group_ids) + len(group_ids) > expected_group_count: + raise ValueError("training claim exceeds the step's expected group count") + for record in records: + record.status = PromptGroupStatus.CLAIMED_FOR_TRAINING + record.claimed_train_step = train_step + open_step.group_ids.append(record.group_id) + + def mark_train_step_applied(self, train_step: int) -> None: + """Record optimizer success while rows are still not checkpoint-covered.""" + open_step = self._require_open_train_step(train_step) + if open_step.status != TrainStepStatus.OPEN: + raise ValueError( + f"train step {train_step} is already {open_step.status.value!r}" + ) + if len(open_step.group_ids) != open_step.expected_group_count: + raise ValueError( + f"train step {train_step} has {len(open_step.group_ids)} claimed " + f"groups; expected {open_step.expected_group_count}" + ) + for group_id in open_step.group_ids: + record = self._require_group(group_id) + if record.status != PromptGroupStatus.CLAIMED_FOR_TRAINING: + raise ValueError( + f"train step {train_step} owns group {group_id!r} in state " + f"{record.status.value!r}" + ) + for group_id in open_step.group_ids: + self._groups[group_id].status = PromptGroupStatus.APPLIED_UNCHECKPOINTED + open_step.status = TrainStepStatus.APPLIED_UNCHECKPOINTED + + def release_applied_train_step(self, train_step: int) -> None: + """Drop group metadata after the caller has cleared all owned TQ rows. + + The current controller clears immediately after optimizer success. A later + persistence change will delay this call until a trainer checkpoint covers + the applied update. + """ + open_step = self._require_open_train_step(train_step) + if open_step.status != TrainStepStatus.APPLIED_UNCHECKPOINTED: + raise ValueError( + f"cannot release train step {train_step} from " + f"{open_step.status.value!r}" + ) + for group_id in open_step.group_ids: + del self._groups[group_id] + self._open_train_step = None + + def rollback_open_train_step(self, train_step: int) -> None: + """Return an uncheckpointed step's groups to finalized ownership.""" + open_step = self._require_open_train_step(train_step) + for group_id in open_step.group_ids: + record = self._require_group(group_id) + if record.status not in { + PromptGroupStatus.CLAIMED_FOR_TRAINING, + PromptGroupStatus.APPLIED_UNCHECKPOINTED, + }: + raise ValueError( + f"cannot roll back group {group_id!r} from {record.status.value!r}" + ) + record.status = PromptGroupStatus.FINALIZED + record.claimed_train_step = None + self._open_train_step = None def discard_group(self, cut: DataPlaneMutationCut, group_id: str) -> None: - """Release ownership after canonical commit or intentional discard.""" + """Drop a group only after its external TQ/Gate ownership is cleaned.""" cut.require_live() - self._require_group(group_id) + record = self._require_group(group_id) + if record.claimed_train_step is not None: + raise ValueError(f"cannot discard training-owned group {group_id!r}") del self._groups[group_id] def discard_canonical_groups( @@ -337,7 +659,7 @@ def discard_canonical_groups( cut: DataPlaneMutationCut, group_ids: set[str], ) -> int: - """Drop ledger copies already owned by canonical replay metadata.""" + """Prefer canonical TQ ownership over a stale unfinished sidecar row.""" cut.require_live() discarded = 0 for group_id in list(self._groups): @@ -346,9 +668,18 @@ def discard_canonical_groups( discarded += 1 return discarded - def state_dict(self) -> RolloutRecoveryLedgerState: - """Return versioned references without serializing full prompt payloads.""" - groups: list[PromptGroupRecoveryState] = [] + def get_group(self, group_id: str) -> PromptGroupRecoveryRecord: + return self._copy_group(self._require_group(group_id)) + + def __len__(self) -> int: + return len(self._groups) + + def __contains__(self, group_id: object) -> bool: + return isinstance(group_id, str) and group_id in self._groups + + def state_dict(self) -> dict[str, Any]: + """Return the versioned metadata envelope used by later persistence.""" + groups = [] for record in self._groups.values(): prompt_payload = record.runtime_prompt_payload if prompt_payload is None: @@ -361,9 +692,6 @@ def state_dict(self) -> RolloutRecoveryLedgerState: prompt_payload, group_id=record.group_id, ) - # sample_id is currently a positional index into a map-style dataset, - # not a dataset-independent identity. The checkpoint is recoverable only - # when that dataset's ordering remains unchanged across the restart. groups.append( { "group_id": record.group_id, @@ -376,25 +704,50 @@ def state_dict(self) -> RolloutRecoveryLedgerState: "expected_generations": record.expected_generations, "target_step": record.target_step, "start_weight_version": record.start_weight_version, + "status": record.status.value, "phase": record.phase.value, + "canonical_meta": copy.deepcopy(record.canonical_meta), + "group_min_weight_version": record.group_min_weight_version, + "group_max_weight_version": record.group_max_weight_version, + "claimed_train_step": record.claimed_train_step, + "siblings": [ + { + "generation_index": sibling.generation_index, + "attempts": [ + { + "attempt_uuid": attempt.attempt_uuid.bytes, + "status": attempt.status.value, + "receipt": copy.deepcopy(attempt.receipt), + "reward": attempt.reward, + "staging_keys": list(attempt.staging_keys), + } + for attempt in sibling.attempts + ], + } + for sibling in record.siblings + ], } ) + open_step = self._open_train_step return { "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, "groups": groups, + "open_train_step": ( + None + if open_step is None + else { + "train_step": open_step.train_step, + "trainer_version": open_step.trainer_version, + "expected_group_count": open_step.expected_group_count, + "group_ids": list(open_step.group_ids), + "status": open_step.status.value, + } + ), } - def load_state_dict( - self, - cut: DataPlaneMutationCut, - state: RolloutRecoveryLedgerState, - ) -> None: - """Replace this empty ledger from a validated checkpoint payload.""" - cut.require_live() - if self._groups: - raise RuntimeError( - "cannot restore into a non-empty rollout recovery ledger" - ) + @classmethod + def from_state_dict(cls, state: dict[str, Any]) -> Self: + """Restore and validate a ledger metadata envelope.""" if not isinstance(state, dict): raise TypeError( "rollout recovery state must be a dictionary, got " @@ -402,194 +755,345 @@ def load_state_dict( ) if state.get("schema_version") != ROLLOUT_RECOVERY_SCHEMA_VERSION: raise ValueError( - "unsupported rollout recovery schema_version=" - f"{state.get('schema_version')!r}; expected " - f"{ROLLOUT_RECOVERY_SCHEMA_VERSION}" - ) - groups = state.get("groups") - if not isinstance(groups, list): - raise TypeError("rollout recovery groups must be a list") - - restored: dict[str, PromptGroupRecoveryRecord] = {} - for index, raw_group in enumerate(groups): - if not isinstance(raw_group, dict): - raise TypeError( - f"rollout recovery groups[{index}] must be a dictionary" - ) - group_id = raw_group.get("group_id") - prompt_id = raw_group.get("prompt_id") - admission_id = raw_group.get("admission_id") - if not isinstance(group_id, str) or not group_id: - raise ValueError( - f"rollout recovery groups[{index}].group_id must be non-empty" - ) - if group_id in restored: - raise ValueError(f"duplicate recovery group_id={group_id!r}") - if not isinstance(admission_id, str) or not admission_id: - raise ValueError( - f"rollout recovery groups[{index}].admission_id must be non-empty" - ) - if not isinstance(prompt_id, str) or not prompt_id: - raise ValueError( - f"rollout recovery groups[{index}].prompt_id must be non-empty" - ) - expected_generations = _require_int( - raw_group.get("expected_generations"), - field=f"groups[{index}].expected_generations", - minimum=1, - ) - start_weight_version = _require_int( - raw_group.get("start_weight_version"), - field=f"groups[{index}].start_weight_version", - minimum=0, - ) - target_step = raw_group.get("target_step") - if target_step is not None: - target_step = _require_int( - target_step, - field=f"groups[{index}].target_step", - minimum=0, - ) - raw_phase = raw_group.get("phase") - if not isinstance(raw_phase, str): - raise ValueError( - f"rollout recovery groups[{index}].phase is invalid: {raw_phase!r}" - ) + "Unsupported rollout-recovery schema version: " + f"{state.get('schema_version')!r}" + ) + raw_groups = state.get("groups") + if not isinstance(raw_groups, list): + raise ValueError("rollout-recovery state must contain a groups list") + + ledger = cls() + seen_attempt_uuids: set[uuid.UUID] = set() + for raw_group in raw_groups: + record = cls._group_from_state( + raw_group, + seen_attempt_uuids=seen_attempt_uuids, + ) + if record.group_id in ledger._groups: + raise ValueError(f"duplicate recovery group_id={record.group_id!r}") + ledger._groups[record.group_id] = record + + raw_open_step = state.get("open_train_step") + if raw_open_step is not None: + if not isinstance(raw_open_step, dict): + raise ValueError("open_train_step must be a mapping or None") try: - phase = PromptGroupPhase(raw_phase) - except ValueError as error: - raise ValueError( - f"rollout recovery groups[{index}].phase is invalid: {raw_phase!r}" - ) from error - raw_prompt_ref = raw_group.get("prompt_ref") - if not isinstance(raw_prompt_ref, dict): - raise TypeError( - f"rollout recovery groups[{index}].prompt_ref must be a dictionary" + open_step = OpenTrainStepRecord( + train_step=int(raw_open_step["train_step"]), + trainer_version=int(raw_open_step["trainer_version"]), + expected_group_count=int(raw_open_step["expected_group_count"]), + group_ids=list(raw_open_step["group_ids"]), + status=TrainStepStatus(raw_open_step["status"]), ) - sample_id = raw_prompt_ref.get("sample_id") - task_name = raw_prompt_ref.get("task_name") - if not isinstance(sample_id, str) or not sample_id: - raise ValueError( - f"rollout recovery groups[{index}].prompt_ref.sample_id " - "must be non-empty" - ) - if sample_id != prompt_id: + except (KeyError, TypeError, ValueError) as error: + raise ValueError("invalid open_train_step state") from error + if not all(isinstance(group_id, str) for group_id in open_step.group_ids): + raise ValueError("open_train_step group_ids must be strings") + if len(open_step.group_ids) != len(set(open_step.group_ids)): + raise ValueError("open_train_step contains duplicate group IDs") + if len(open_step.group_ids) > open_step.expected_group_count: + raise ValueError("open_train_step exceeds expected_group_count") + expected_group_status = ( + PromptGroupStatus.CLAIMED_FOR_TRAINING + if open_step.status == TrainStepStatus.OPEN + else PromptGroupStatus.APPLIED_UNCHECKPOINTED + ) + for group_id in open_step.group_ids: + record = ledger._require_group(group_id) + if ( + record.claimed_train_step != open_step.train_step + or record.status != expected_group_status + ): + raise ValueError( + f"open_train_step ownership mismatch for group {group_id!r}" + ) + claimed_group_ids = { + record.group_id + for record in ledger._groups.values() + if record.claimed_train_step is not None + } + if claimed_group_ids != set(open_step.group_ids): raise ValueError( - f"rollout recovery groups[{index}] prompt_id and " - "prompt_ref.sample_id must match" - ) - if task_name is not None and not isinstance(task_name, str): - raise TypeError( - f"rollout recovery groups[{index}].prompt_ref.task_name " - "must be a string or None" + "open_train_step does not list every training-owned group" ) - restored[group_id] = PromptGroupRecoveryRecord( - group_id=group_id, - admission_id=admission_id, - prompt_id=prompt_id, - prompt_ref=PromptRef( - sample_id=sample_id, - task_name=task_name, - ), - runtime_prompt_payload=None, - expected_generations=expected_generations, - target_step=target_step, - start_weight_version=start_weight_version, - phase=phase, - ) - - admission_states: dict[str, tuple[PromptGroupPhase, int | None]] = {} - for record in restored.values(): + ledger._open_train_step = open_step + elif any( + record.claimed_train_step is not None for record in ledger._groups.values() + ): + raise ValueError("claimed groups require an open_train_step record") + admission_states: dict[ + str, tuple[PromptGroupPhase, Optional[int]] + ] = {} + for record in ledger._groups.values(): signature = (record.phase, record.target_step) - prior = admission_states.setdefault(record.admission_id, signature) - if prior != signature: + previous = admission_states.setdefault(record.admission_id, signature) + if previous != signature: raise ValueError( "rollout recovery groups sharing admission_id=" f"{record.admission_id!r} disagree on phase or target_step" ) - self._groups = restored + return ledger + + def load_state_dict(self, state: RolloutRecoveryState) -> None: + """Replace this empty ledger from a validated checkpoint envelope.""" + if self._groups or self._open_train_step is not None: + raise RuntimeError("cannot restore into a non-empty rollout recovery ledger") + restored = self.from_state_dict(state) + self._groups = restored._groups + self._open_train_step = restored._open_train_step + + @classmethod + def _group_from_state( + cls, + raw_group: Any, + *, + seen_attempt_uuids: set[uuid.UUID], + ) -> PromptGroupRecoveryRecord: + if not isinstance(raw_group, dict): + raise ValueError("rollout-recovery group must be a mapping") + group_id = raw_group.get("group_id") + admission_id = raw_group.get("admission_id") + prompt_id = raw_group.get("prompt_id") + expected_generations = raw_group.get("expected_generations") + siblings_state = raw_group.get("siblings") + if not isinstance(group_id, str) or not group_id: + raise ValueError("group_id must be a non-empty string") + if not isinstance(admission_id, str) or not admission_id: + raise ValueError("admission_id must be a non-empty string") + if not isinstance(prompt_id, str) or not prompt_id: + raise ValueError("prompt_id must be a non-empty string") + if not isinstance(expected_generations, int) or expected_generations < 1: + raise ValueError("expected_generations must be a positive integer") + if ( + not isinstance(siblings_state, list) + or len(siblings_state) != expected_generations + ): + raise ValueError( + f"recovery group {group_id!r} must contain " + f"{expected_generations} siblings" + ) + raw_status = raw_group.get("status") + if not isinstance(raw_status, str): + raise ValueError(f"invalid prompt group status={raw_status!r}") + try: + status = PromptGroupStatus(raw_status) + except ValueError as error: + raise ValueError(f"invalid prompt group status={raw_status!r}") from error + raw_phase = raw_group.get("phase") + if not isinstance(raw_phase, str): + raise ValueError(f"invalid prompt group phase={raw_phase!r}") + try: + phase = PromptGroupPhase(raw_phase) + except ValueError as error: + raise ValueError(f"invalid prompt group phase={raw_phase!r}") from error + + siblings: list[RolloutSiblingRecord] = [] + for generation_index, sibling_state in enumerate(siblings_state): + if not isinstance(sibling_state, dict): + raise ValueError("rollout-recovery sibling must be a mapping") + if sibling_state.get("generation_index") != generation_index: + raise ValueError("generation indices must be contiguous") + logical_id = f"{group_id}_g{generation_index}" + attempts_state = sibling_state.get("attempts") + if not isinstance(attempts_state, list) or not attempts_state: + raise ValueError(f"logical rollout {logical_id!r} has no attempts") + attempts: list[RolloutAttemptRecord] = [] + for attempt_state in attempts_state: + if not isinstance(attempt_state, dict): + raise ValueError("rollout-recovery attempt must be a mapping") + raw_attempt_uuid = attempt_state.get("attempt_uuid") + if ( + not isinstance(raw_attempt_uuid, bytes) + or len(raw_attempt_uuid) != 16 + ): + raise ValueError("attempt_uuid must contain exactly 16 bytes") + attempt_uuid = uuid.UUID(bytes=raw_attempt_uuid) + if attempt_uuid in seen_attempt_uuids: + raise ValueError("duplicate rollout attempt identity") + seen_attempt_uuids.add(attempt_uuid) + gate_id = f"{logical_id}_a{attempt_uuid.hex}" + raw_attempt_status = attempt_state.get("status") + if not isinstance(raw_attempt_status, str): + raise ValueError( + f"invalid rollout attempt status={raw_attempt_status!r}" + ) + try: + attempt_status = RolloutAttemptStatus(raw_attempt_status) + except ValueError as error: + raise ValueError( + f"invalid rollout attempt status={raw_attempt_status!r}" + ) from error + receipt = attempt_state.get("receipt") + reward = attempt_state.get("reward") + staging_keys = attempt_state.get("staging_keys") + if not isinstance(staging_keys, list) or not all( + isinstance(key, str) for key in staging_keys + ): + raise ValueError("staging_keys must be a list of strings") + if attempt_status == RolloutAttemptStatus.SEALED: + if not isinstance(receipt, dict) or not isinstance( + reward, (int, float) + ): + raise ValueError("sealed attempts require receipt and reward") + if receipt.get("rollout_id") != gate_id: + raise ValueError("sealed receipt identity mismatch") + if _receipt_staging_keys(receipt) != staging_keys: + raise ValueError("sealed receipt staging manifest mismatch") + elif receipt is not None or reward is not None or staging_keys: + raise ValueError("only sealed attempts may retain receipt data") + attempts.append( + RolloutAttemptRecord( + attempt_uuid=attempt_uuid, + status=attempt_status, + receipt=copy.deepcopy(receipt), + reward=float(reward) if reward is not None else None, + staging_keys=list(staging_keys), + ) + ) + siblings.append( + RolloutSiblingRecord( + generation_index=generation_index, + attempts=attempts, + ) + ) + + raw_prompt_ref = raw_group.get("prompt_ref") + if not isinstance(raw_prompt_ref, dict): + raise ValueError("prompt_ref must be a mapping") + sample_id = raw_prompt_ref.get("sample_id") + task_name = raw_prompt_ref.get("task_name") + if not isinstance(sample_id, str) or not sample_id: + raise ValueError("prompt_ref sample_id must be a non-empty string") + if task_name is not None and not isinstance(task_name, str): + raise ValueError("prompt_ref task_name must be a string or None") + if sample_id != prompt_id: + raise ValueError("prompt_ref sample_id must match prompt_id") + canonical_meta = raw_group.get("canonical_meta") + if canonical_meta is not None and not isinstance(canonical_meta, KVBatchMeta): + raise ValueError("canonical_meta must be KVBatchMeta or None") + target_step = raw_group.get("target_step") + claimed_train_step = raw_group.get("claimed_train_step") + if target_step is not None and not isinstance(target_step, int): + raise ValueError("target_step must be an integer or None") + if claimed_train_step is not None and not isinstance(claimed_train_step, int): + raise ValueError("claimed_train_step must be an integer or None") + start_weight = raw_group.get("start_weight_version") + if not isinstance(start_weight, int): + raise ValueError("start_weight_version must be an integer") + min_weight = raw_group.get("group_min_weight_version") + max_weight = raw_group.get("group_max_weight_version") + if min_weight is not None and not isinstance(min_weight, int): + raise ValueError("group_min_weight_version must be an integer or None") + if max_weight is not None and not isinstance(max_weight, int): + raise ValueError("group_max_weight_version must be an integer or None") + + finalized_states = { + PromptGroupStatus.FINALIZED, + PromptGroupStatus.CLAIMED_FOR_TRAINING, + PromptGroupStatus.APPLIED_UNCHECKPOINTED, + } + prefinalization_sealed_states = { + PromptGroupStatus.READY_TO_FINALIZE, + PromptGroupStatus.FINALIZING, + PromptGroupStatus.FINALIZATION_UNKNOWN, + } + all_current_attempts_sealed = all( + sibling.current_attempt.status == RolloutAttemptStatus.SEALED + for sibling in siblings + ) + if status in prefinalization_sealed_states and not all_current_attempts_sealed: + raise ValueError( + f"group state {status.value!r} requires every sibling to be sealed" + ) + if status in finalized_states: + if canonical_meta is None or min_weight is None or max_weight is None: + raise ValueError("finalized group state requires canonical metadata") + if list(canonical_meta.sample_ids) != [ + f"{group_id}_g{s.generation_index}" for s in siblings + ]: + raise ValueError("canonical sample IDs do not match logical lineage") + elif canonical_meta is not None: + raise ValueError("unfinished group cannot contain canonical metadata") + claimed_states = { + PromptGroupStatus.CLAIMED_FOR_TRAINING, + PromptGroupStatus.APPLIED_UNCHECKPOINTED, + } + if (status in claimed_states) != (claimed_train_step is not None): + raise ValueError( + "claimed_train_step must be present exactly for training-owned groups" + ) + + return PromptGroupRecoveryRecord( + group_id=group_id, + admission_id=admission_id, + prompt_id=prompt_id, + prompt_ref=PromptRef(sample_id=sample_id, task_name=task_name), + runtime_prompt_payload=None, + expected_generations=expected_generations, + target_step=target_step, + start_weight_version=start_weight, + siblings=siblings, + phase=phase, + status=status, + canonical_meta=copy.deepcopy(canonical_meta), + group_min_weight_version=min_weight, + group_max_weight_version=max_weight, + claimed_train_step=claimed_train_step, + ) def _require_group(self, group_id: str) -> PromptGroupRecoveryRecord: try: return self._groups[group_id] except KeyError as error: - raise KeyError(f"unknown recovery group_id={group_id!r}") from error + raise ValueError(f"unknown recovery group_id={group_id!r}") from error - def __len__(self) -> int: - return len(self._groups) + @staticmethod + def _copy_group(record: PromptGroupRecoveryRecord) -> PromptGroupRecoveryRecord: + """Copy mutable lineage metadata without duplicating the prompt payload.""" + return PromptGroupRecoveryRecord( + group_id=record.group_id, + admission_id=record.admission_id, + prompt_id=record.prompt_id, + prompt_ref=record.prompt_ref, + runtime_prompt_payload=record.runtime_prompt_payload, + expected_generations=record.expected_generations, + target_step=record.target_step, + start_weight_version=record.start_weight_version, + siblings=copy.deepcopy(record.siblings), + phase=record.phase, + status=record.status, + canonical_meta=copy.deepcopy(record.canonical_meta), + group_min_weight_version=record.group_min_weight_version, + group_max_weight_version=record.group_max_weight_version, + claimed_train_step=record.claimed_train_step, + ) + @staticmethod + def _require_sibling( + record: PromptGroupRecoveryRecord, generation_index: int + ) -> RolloutSiblingRecord: + if not 0 <= generation_index < len(record.siblings): + raise ValueError( + f"generation_index={generation_index} is outside group " + f"{record.group_id!r}" + ) + return record.siblings[generation_index] -def _validate_batch_shortfall(value: object) -> dict[int, int]: - """Return a defensive copy of per-step permanent rollout losses.""" - if not isinstance(value, dict): - raise TypeError("rollout recovery batch_shortfall must be a dictionary") - batch_shortfall: dict[int, int] = {} - for step, count in value.items(): - if ( - isinstance(step, bool) - or not isinstance(step, int) - or step < 0 - or isinstance(count, bool) - or not isinstance(count, int) - or count < 0 - ): + @staticmethod + def _require_group_status( + record: PromptGroupRecoveryRecord, + *, + allowed: set[PromptGroupStatus], + transition: str, + ) -> None: + if record.status not in allowed: raise ValueError( - "rollout recovery batch_shortfall entries must contain " - f"non-negative integer steps and counts, got {step!r}: {count!r}" - ) - batch_shortfall[step] = count - return batch_shortfall - - -def build_rollout_recovery_state( - ledger: RolloutRecoveryLedger, - *, - batch_shortfall: dict[int, int], - sampler_stamps_target_steps: bool, -) -> RolloutRecoveryState: - """Build the complete versioned sidecar from ledger and controller state.""" - if not isinstance(sampler_stamps_target_steps, bool): - raise TypeError( - "rollout recovery sampler_stamps_target_steps must be a boolean" - ) - ledger_state = ledger.state_dict() - return { - "schema_version": ledger_state["schema_version"], - "groups": ledger_state["groups"], - "batch_shortfall": _validate_batch_shortfall(batch_shortfall), - "sampler_stamps_target_steps": sampler_stamps_target_steps, - } - - -def parse_rollout_recovery_state(state: object) -> ParsedRolloutRecoveryState: - """Validate and split a complete checkpoint sidecar by runtime owner.""" - if not isinstance(state, dict): - raise TypeError( - "rollout recovery sidecar must contain a dictionary, got " - f"{type(state).__name__}" - ) - if state.get("schema_version") != ROLLOUT_RECOVERY_SCHEMA_VERSION: - raise ValueError( - "unsupported rollout recovery schema_version=" - f"{state.get('schema_version')!r}; expected " - f"{ROLLOUT_RECOVERY_SCHEMA_VERSION}" - ) - groups = state.get("groups") - if not isinstance(groups, list): - raise TypeError("rollout recovery groups must be a list") - - raw_sampler_stamps = state.get("sampler_stamps_target_steps") - if raw_sampler_stamps is not None and not isinstance(raw_sampler_stamps, bool): - raise TypeError( - "rollout recovery sampler_stamps_target_steps must be a boolean" - ) + f"cannot {transition} group {record.group_id!r} from " + f"{record.status.value!r}" + ) - ledger_state: RolloutRecoveryLedgerState = { - "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, - "groups": groups, - } - return ParsedRolloutRecoveryState( - ledger_state=ledger_state, - batch_shortfall=_validate_batch_shortfall(state.get("batch_shortfall", {})), - sampler_stamps_target_steps=raw_sampler_stamps, - ) + def _require_open_train_step(self, train_step: int) -> OpenTrainStepRecord: + open_step = self._open_train_step + if open_step is None or open_step.train_step != train_step: + raise ValueError(f"train step {train_step} is not open") + return open_step diff --git a/tests/unit/data_plane/test_rollout_reassembler.py b/tests/unit/data_plane/test_rollout_reassembler.py index 27816aad807..0173c4773e9 100644 --- a/tests/unit/data_plane/test_rollout_reassembler.py +++ b/tests/unit/data_plane/test_rollout_reassembler.py @@ -255,11 +255,37 @@ def test_finalize_group_publishes_n_rows_with_placeholder(tq_client, partitions) finalizer._source.fetch([receipt["manifest"][0]["staging_key"]]) +def test_finalize_group_maps_physical_attempt_to_stable_canonical_id( + tq_client, partitions +): + group_id = "stable" + physical_id = f"{group_id}_g0_aattempt" + canonical_id = f"{group_id}_g0" + receipt, _ = _stage_fixture( + tq_client, + "worked_example", + rollout_id=physical_id, + ) + receipt["rollout_id"] = physical_id + + finalized = _finalizer(tq_client).finalize_group( + group_id, + [physical_id], + [receipt], + [1.0], + mask_sample=[False], + fallback_weight_version=4, + prompt_idx=0, + canonical_sample_ids=[canonical_id], + ) + + assert finalized.meta is not None + assert finalized.meta.sample_ids == [canonical_id] + assert _fetch_rows(tq_client, [canonical_id])["input_ids"] is not None + + def test_finalize_group_reports_valid_and_total_row_counts(tq_client, partitions): - """The finalizer no longer drops a low-valid-fraction group itself -- - only the controller can source a replacement, so it reports the counts - and always finalizes (see min_valid_fraction_per_group's removal from - RolloutReassembler; the threshold now lives in single_controller.py).""" + """The finalizer reports validity; the controller owns replacement policy.""" group_id = "grp2" rollout_ids = [f"{group_id}_g0", f"{group_id}_g1"] finalizer = _finalizer(tq_client) diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index ea7f78c9e29..56ef5d25b0c 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -44,6 +44,7 @@ from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data.processors import nemo_gym_data_processor from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.experience.failures import GenerationUnavailable from nemo_rl.experience.interfaces import ( NEMO_GYM_GROUP_ATTEMPT_KEY, NEMO_GYM_GROUP_ID_KEY, @@ -1516,6 +1517,9 @@ def reserve( rollout_ids=rollout_ids, ) + async def clear_staging_keys(self, staging_keys): + del staging_keys + def _receipt_record( rollout_ids, receipts, instance_configs=None, *, loss_multiplier=1.0 @@ -1558,7 +1562,6 @@ def _make_capture_manager( mgr._tokenizer = None mgr._num_generations_per_prompt = num_generations mgr._tq_buffer = buf - mgr._env_handles = {} mgr._weight_version = 7 mgr._retry_policy = ( retry_policy @@ -1568,21 +1571,47 @@ def _make_capture_manager( mgr._stats = RolloutStats() mgr._skipped_prompts = 0 mgr._consecutive_infra_drops = 0 + mgr._recovery_ledger = RolloutRecoveryLedger() class _CaptureImpl: def __init__(self): self.seen_rollout_ids = None - async def run_rollout(self, _sample, *, rollout_ids=None): + async def run_rollout( + self, + _sample, + *, + rollout_ids=None, + generation_indices=None, + on_completion=None, + ): self.seen_rollout_ids = rollout_ids if on_run is not None: await on_run(_sample) - return _receipt_record( - rollout_ids, - [{"rollout_id": rid} for rid in rollout_ids], - instance_configs=instance_configs, + indices = generation_indices or list(range(len(rollout_ids))) + selected_ids = [rollout_ids[index] for index in indices] + selected_configs = ( + [instance_configs[index] for index in indices] + if instance_configs is not None + else None + ) + receipts = [ + { + "rollout_id": rollout_id, + "manifest": [{"staging_key": f"{rollout_id}/call"}], + } + for rollout_id in selected_ids + ] + record = _receipt_record( + selected_ids, + receipts, + instance_configs=selected_configs, loss_multiplier=float(_sample.get("loss_multiplier", 1.0)), ) + if on_completion is not None: + for generation_index, completion in zip(indices, record.completions): + await on_completion(generation_index, completion) + return record mgr._impl = _CaptureImpl() return mgr @@ -1609,19 +1638,26 @@ def test_mints_ids_and_returns_metadata_request(self): request = _run( mgr.generate_for_finalization( - {"prompt": "p", "loss_multiplier": 0.25}, target_step=5 + {"prompt": "p", "idx": 0, "loss_multiplier": 0.25}, target_step=5 ) ) + assert request is not None # Rollout ids were minted from the reserved group id and threaded # end to end: reserve -> impl -> metadata-only actor request. (group_id,) = buf._slots - expected_ids = [f"{group_id}_g0", f"{group_id}_g1"] - assert buf.reserve_rollout_ids == [expected_ids] - assert mgr._impl.seen_rollout_ids == expected_ids + canonical_ids = [f"{group_id}_g0", f"{group_id}_g1"] + attempt_ids = buf.reserve_rollout_ids[0] + assert attempt_ids is not None + assert all( + attempt_id.startswith(f"{canonical_id}_a") + for attempt_id, canonical_id in zip(attempt_ids, canonical_ids) + ) + assert mgr._impl.seen_rollout_ids == attempt_ids assert request.group_id == group_id - assert request.rollout_ids == tuple(expected_ids) - assert [r["rollout_id"] for r in request.receipts] == expected_ids + assert request.rollout_ids == tuple(attempt_ids) + assert request.canonical_sample_ids == tuple(canonical_ids) + assert [r["rollout_id"] for r in request.receipts] == attempt_ids assert request.rewards == (0.5, 0.5) assert request.mask_sample == (False, False) assert request.loss_multiplier == 0.25 @@ -1638,7 +1674,91 @@ async def _boom(_sample): mgr = _make_capture_manager(buf, on_run=_boom) with pytest.raises(RuntimeError, match="rollout exploded"): - _run(mgr.generate_for_finalization({"prompt": "p"})) - # The slot is released; abandoned staged rows are swept with the - # staging partition at run end (no per-rollout control-plane call). + _run(mgr.generate_for_finalization({"prompt": "p", "idx": 0})) assert len(buf.abort_calls) == 1 + + def test_retries_infrastructure_failure_with_stable_logical_ids(self): + buf = _FakeCaptureBuffer() + attempts = 0 + + async def _fail_once(_sample): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise GenerationUnavailable("worker disappeared") + + mgr = _make_capture_manager(buf, on_run=_fail_once) + mgr._retry_policy = RolloutRetryPolicy.single_attempt( + max_infra_attempts=2, + backoff_base_s=0.0, + ) + + request = _run(mgr.generate_for_finalization({"prompt": "p", "idx": 4})) + + assert request is not None + assert attempts == 2 + assert len(buf.reserve_rollout_ids) == 2 + assert buf.reserve_rollout_ids[0] != buf.reserve_rollout_ids[1] + assert len(buf.abort_calls) == 1 + assert buf.abort_calls[0] == request.group_id + assert request.canonical_sample_ids == ( + f"{request.group_id}_g0", + f"{request.group_id}_g1", + ) + assert mgr.stats.as_metrics()["rollout/redispatch_total"] == 1.0 + + def test_reuses_sealed_sibling_and_redispatches_only_incomplete_one(self): + buf = _FakeCaptureBuffer() + mgr = _make_capture_manager(buf) + mgr._retry_policy = RolloutRetryPolicy.single_attempt( + max_infra_attempts=2, + backoff_base_s=0.0, + ) + + class _PartialCaptureImpl: + def __init__(self): + self.generation_indices: list[list[int]] = [] + + async def run_rollout( + self, + _sample, + *, + rollout_ids=None, + generation_indices=None, + on_completion=None, + ): + indices = list(generation_indices) + self.generation_indices.append(indices) + completions = [] + for generation_index in indices: + rollout_id = rollout_ids[generation_index] + receipt = { + "rollout_id": rollout_id, + "manifest": [{"staging_key": f"{rollout_id}/call"}], + } + completion = _receipt_record([rollout_id], [receipt]).completions[0] + completions.append(completion) + await on_completion(generation_index, completion) + if len(self.generation_indices) == 1: + raise GenerationUnavailable("worker disappeared") + return PromptGroupRecord( + prompt_idx=0, + prompt=[], + extra_env_info={}, + metadata={"task_name": "nemo_gym"}, + completions=completions, + rollout_metrics={}, + ) + + impl = _PartialCaptureImpl() + mgr._impl = impl + + request = _run(mgr.generate_for_finalization({"prompt": "p", "idx": 9})) + + assert request is not None + assert impl.generation_indices == [[0, 1], [1]] + first_ids, second_ids = buf.reserve_rollout_ids + assert first_ids is not None and second_ids is not None + assert second_ids[0] == first_ids[0] + assert second_ids[1] != first_ids[1] + assert request.rollout_ids == (second_ids[0], second_ids[1]) diff --git a/tests/unit/experience/test_rollout_reassembler_actor.py b/tests/unit/experience/test_rollout_reassembler_actor.py index 4a21747e228..f8705f207e5 100644 --- a/tests/unit/experience/test_rollout_reassembler_actor.py +++ b/tests/unit/experience/test_rollout_reassembler_actor.py @@ -35,6 +35,7 @@ def _request() -> ReassemblyRequest: return ReassemblyRequest( group_id="group", rollout_ids=("group_g0",), + canonical_sample_ids=("group_g0",), receipts=( { "rollout_id": "group_g0", diff --git a/tests/unit/experience/test_rollouts.py b/tests/unit/experience/test_rollouts.py index 3891e2564d4..e75dba5332e 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -2226,6 +2226,7 @@ def remote(self, inputs, timer_prefix): return _Stream() manager = object.__new__(AsyncNemoGymRolloutImpl) + manager._num_generations_per_prompt = 2 # These tests cover stream ordering/dedup, not deadlines or re-dispatch. manager._timeouts = RolloutTimeouts() manager._max_gym_row_attempts = 1 @@ -2342,6 +2343,7 @@ def remote(self, inputs, timer_prefix): return _DuplicateStream() manager = object.__new__(AsyncNemoGymRolloutImpl) + manager._num_generations_per_prompt = 2 # These tests cover stream ordering/dedup, not deadlines or re-dispatch. manager._timeouts = RolloutTimeouts() manager._max_gym_row_attempts = 1 diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 7400cdad49b..32fdf367e57 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -416,6 +416,10 @@ def __init__( self.load_calls: list[dict[str, Any]] = [] self.checkpoint_barrier: Optional[DataPlaneCheckpointBarrier] = None + @property + def group_ids(self) -> tuple[str, ...]: + return () + def set_data_plane_checkpoint_barrier( self, barrier: DataPlaneCheckpointBarrier ) -> None: diff --git a/tests/unit/single_controller/test_finalizer_lifecycle.py b/tests/unit/single_controller/test_finalizer_lifecycle.py index bf34cee73ad..47f417bfd64 100644 --- a/tests/unit/single_controller/test_finalizer_lifecycle.py +++ b/tests/unit/single_controller/test_finalizer_lifecycle.py @@ -77,6 +77,7 @@ def _request() -> ReassemblyRequest: return ReassemblyRequest( group_id="group", rollout_ids=("group_g0",), + canonical_sample_ids=("group_g0",), receipts=( { "manifest": [ @@ -101,14 +102,20 @@ def _controller(actor: object) -> Any: ctrl._finalizer_waiters = 0 ctrl._finalizer_unknown_outcomes = 0 ctrl._finalizer_metrics_by_group = {} + ctrl._rollout_recovery_ledger = MagicMock() + ctrl._rollout_recovery_ledger.__contains__.return_value = False + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._buffer = MagicMock() ctrl._buffer.commit_finalized = AsyncMock() ctrl._dp_client = _DataPlaneClient() ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._partition_id = "canonical" ctrl._master_config = SimpleNamespace( - token_capture=SimpleNamespace(staging_partition="staging") + token_capture=SimpleNamespace(staging_partition="staging"), + grpo=SimpleNamespace(num_prompts_per_step=1), ) + ctrl._trainer_version = 3 + ctrl._train_steps = 3 return ctrl From 3947302b12e98aaad4239a00895cec4d8f3a3de0 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 19 Aug 2026 13:38:11 -0400 Subject: [PATCH 02/28] fix(data-plane): make checkpoint barrier reentrant Signed-off-by: Anish Mahishi (cherry picked from commit c1f1225ba99fa570521f95c86d8c9a1ec41ac5a5) --- .../algorithms/async_utils/replay_buffer.py | 28 +++++++++++++-- .../test_tq_replay_buffer.py | 34 +++++++++++++++++++ 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 529a27f9f43..bf479b8de14 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -230,18 +230,42 @@ def __init__(self) -> None: self._condition = asyncio.Condition() self._checkpoint_active = False self._active_mutations = 0 + self._mutation_depth_by_task: dict[asyncio.Task[Any], int] = {} + self._mutation_cut_by_task: dict[ + asyncio.Task[Any], DataPlaneMutationCut + ] = {} @asynccontextmanager async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: - """Yield a live mutation capability after any active checkpoint exits.""" + """Yield one task-local live cut after any active checkpoint exits.""" + task = asyncio.current_task() + if task is None: + raise RuntimeError("data-plane mutation must run inside an asyncio task") + depth = self._mutation_depth_by_task.get(task, 0) + if depth: + # Replay-buffer helpers may join a controller-owned mutation. Do not + # wait behind a checkpoint that is already waiting for this outer + # mutation, or the two tasks deadlock. + self._mutation_depth_by_task[task] = depth + 1 + cut = self._mutation_cut_by_task[task] + cut.require_live() + try: + yield cut + finally: + self._mutation_depth_by_task[task] -= 1 + return async with self._condition: await self._condition.wait_for(lambda: not self._checkpoint_active) self._active_mutations += 1 - cut = DataPlaneMutationCut(self) + self._mutation_depth_by_task[task] = 1 + cut = DataPlaneMutationCut(self) + self._mutation_cut_by_task[task] = cut try: yield cut finally: cut._invalidate() + del self._mutation_cut_by_task[task] + del self._mutation_depth_by_task[task] async with self._condition: self._active_mutations -= 1 if self._active_mutations == 0: diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index c4191fce21a..eeec744e3da 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -298,6 +298,40 @@ async def checkpoint() -> None: asyncio.run(exercise()) + def test_nested_mutation_does_not_deadlock_with_waiting_checkpoint(self): + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + outer_entered = asyncio.Event() + allow_nested = asyncio.Event() + nested_entered = asyncio.Event() + checkpoint_entered = asyncio.Event() + + async def mutate() -> None: + async with barrier.mutation(): + outer_entered.set() + await allow_nested.wait() + async with barrier.mutation(): + nested_entered.set() + + async def checkpoint() -> None: + async with barrier.checkpoint(): + checkpoint_entered.set() + + mutation_task = asyncio.create_task(mutate()) + await outer_entered.wait() + checkpoint_task = asyncio.create_task(checkpoint()) + await asyncio.sleep(0) + allow_nested.set() + + await asyncio.wait_for(nested_entered.wait(), timeout=5.0) + assert not checkpoint_entered.is_set() + await asyncio.wait_for( + asyncio.gather(mutation_task, checkpoint_task), timeout=5.0 + ) + assert checkpoint_entered.is_set() + + asyncio.run(exercise()) + def test_two_checkpoints_serialize_without_deadlock(self): async def exercise() -> None: barrier = DataPlaneCheckpointBarrier() From e09aab41a24e5d8e8e8908149046cde0c58e7794 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 23:58:19 -0400 Subject: [PATCH 03/28] feat(rollout): persist sibling recovery state Signed-off-by: Anish Mahishi --- .../algorithms/async_utils/replay_buffer.py | 22 ++++- nemo_rl/algorithms/single_controller.py | 88 ++++++++++++++++-- nemo_rl/experience/rollout_manager.py | 90 ++++++++++++------ nemo_rl/experience/rollout_recovery.py | 92 +++++++++++++++++-- .../unit/experience/test_rollout_recovery.py | 51 ++++++++++ .../single_controller/test_checkpointing.py | 3 + .../test_single_controller_actor.py | 19 +++- .../test_tq_replay_buffer.py | 33 +++++++ 8 files changed, 347 insertions(+), 51 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index bf479b8de14..1686edb819c 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -1683,16 +1683,30 @@ async def load_state_dict( for group in groups: meta = group["meta"] + staging_keys: list[str] = [] + for tag in meta.tags or []: + encoded_plan = tag.get(ROUTE_PLAN_TAG) + if encoded_plan is None: + continue + from nemo_rl.experience.route_plan import decode_route_plan + + staging_keys.extend( + decode_route_plan(encoded_plan).cleanup_staging_keys + ) self.meta_list.append(meta) self.start_weight_list.append(group["start_weight"]) self.end_weight_list.append(group["end_weight"]) self.target_step_list.append(group["target_step"]) self.ready_list.append(True) self._group_ids.append(group["group_id"]) - # Token-capture bookkeeping is not checkpointed: restored groups - # are already finalized, so there are no staged rows to own. - self._rollout_ids_list.append(None) - self._staging_keys_list.append(None) + # Live token-capture reservations retain physical rollout IDs. Once a + # group is canonical, only stable sample IDs are durable and sufficient + # for replay ownership; staging cleanup is reconstructed from the route + # plans stored in canonical row tags. + self._rollout_ids_list.append(list(meta.sample_ids)) + self._staging_keys_list.append( + list(dict.fromkeys(staging_keys)) if staging_keys else None + ) print( f"📦 Restored {len(groups)} replay group(s) from checkpoint", diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index e161178fee7..a44fdaf8ea8 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -396,16 +396,15 @@ def __init__( # makes the native snapshot match the controller's metadata-only replay # index exactly. Generation may continue, but completed rollouts wait at # commit; _buffer_capacity bounds reservations and eventually stalls - # dispatch instead of allowing unbounded TQ growth. The finalizer commit, - # pre-publication cleanup, and post-train cleanup paths take the mutation - # side too. Authoritative (replay-index) snapshots are still rejected in - # token-capture mode by validate_single_controller_config: the capture - # dispatch path does not thread lineage group ids into the recovery - # ledger, so a resume could not reap it. Shadow snapshots still run. + # dispatch instead of allowing unbounded TQ growth. Finalizer commits, + # cleanup, and streamed recovery-ledger transitions use the mutation side. self._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() self._buffer.set_data_plane_checkpoint_barrier( self._data_plane_checkpoint_barrier ) + self._rollout_manager.set_data_plane_checkpoint_barrier( + self._data_plane_checkpoint_barrier + ) # Gate: cleared during _sync_weights, set when generation may proceed self._rollout_permitted: asyncio.Event = asyncio.Event() @@ -737,6 +736,7 @@ async def _maybe_restore_rollout_recovery( recovery_ledger = self._rollout_manager.recovery_ledger async with self._data_plane_checkpoint_barrier.mutation() as cut: recovery_ledger.load_state_dict(cut, parsed_state.ledger_state) + recovery_ledger.prepare_for_restart(cut) self._batch_shortfall = parsed_state.batch_shortfall canonical_state = self._buffer.metadata_state_dict( saved_capacity=self._async_cfg.max_buffered_rollouts @@ -745,6 +745,11 @@ async def _maybe_restore_rollout_recovery( group["group_id"] for group in canonical_state["groups"] } recovery_ledger.discard_canonical_groups(cut, canonical_group_ids) + if self._master_config.token_capture.enabled: + await self._validate_rollout_recovery_inventory( + replay_metadata=canonical_state, + clear_unreferenced=True, + ) await self._rehydrate_rollout_recovery_prompts(cut) self._sampler_stamps_target_steps = ( parsed_state.sampler_stamps_target_steps @@ -1026,6 +1031,57 @@ async def _validate_replay_inventory( flush=True, ) + async def _validate_rollout_recovery_inventory( + self, + *, + replay_metadata: Optional[TQReplayMetadataState], + clear_unreferenced: bool, + ) -> None: + """Require every unfinished receipt or deferred route to retain staging.""" + expected_staging_keys = ( + self._rollout_recovery_ledger.expected_staging_keys() + ) + if replay_metadata is not None: + for group in replay_metadata["groups"]: + for tag in group["meta"].tags or []: + encoded_plan = tag.get(ROUTE_PLAN_TAG) + if encoded_plan is not None: + expected_staging_keys.update( + decode_route_plan(encoded_plan).cleanup_staging_keys + ) + + staging_partition = self._master_config.token_capture.staging_partition + actual_staging_keys = set( + await self._call_dp( + "list_sample_ids", + partition_id=staging_partition, + ) + ) + missing = sorted(expected_staging_keys - actual_staging_keys) + if missing: + raise RuntimeError( + "rollout-recovery ownership references staging rows missing " + f"from live TQ state: missing={missing[:10]!r} " + f"(total={len(missing)})" + ) + unreferenced = sorted(actual_staging_keys - expected_staging_keys) + if clear_unreferenced and unreferenced: + await self._call_dp( + "clear_samples", + sample_ids=unreferenced, + partition_id=staging_partition, + ) + print( + "rollout recovery cleared unreferenced staging rows: " + f"count={len(unreferenced)}", + flush=True, + ) + print( + "📦 Rollout-recovery staging inventory validated: " + f"referenced={len(expected_staging_keys)}", + flush=True, + ) + async def _maybe_restore_replacement_reserve(self) -> None: """Restore spare prompts diverted before the previous run's checkpoint. @@ -1264,13 +1320,13 @@ async def _finalize_with_actor( # one mutation cut so a TQ snapshot sees all of them or none of them. async with self._data_plane_checkpoint_barrier.mutation() as cut: ledger = self._rollout_recovery_ledger - ledger.mark_finalization_started(request.group_id) + ledger.mark_finalization_started(cut, request.group_id) try: rpc_submitted = True finalized = await actor.finalize.remote(request) except BaseException: self._finalizer_unknown_outcomes += 1 - ledger.mark_finalization_unknown(request.group_id) + ledger.mark_finalization_unknown(cut, request.group_id) print( "FATAL: finalizer actor RPC failed after submission; canonical " f"publication outcome is unknown for group {request.group_id}. " @@ -1484,6 +1540,16 @@ async def _dispatch_one_prompt( # committed, so the train pump will never release # this permit, and the step it was stamped for # must be allowed to close short. + if ( + self._rollout_recovery_enabled + and lineage_group_id is not None + ): + async with ( + self._data_plane_checkpoint_barrier.mutation() + ) as cut: + await self._rollout_manager.discard_recovery_group( + cut, lineage_group_id + ) self._buffer_capacity.release() self._credit_shortfall(target_step) return @@ -3311,6 +3377,12 @@ async def _save_checkpoint( rollout_recovery_payload ).hexdigest() + if self._master_config.token_capture.enabled: + await self._validate_rollout_recovery_inventory( + replay_metadata=replay_metadata, + clear_unreferenced=False, + ) + await self._save_data_plane_checkpoint( checkpoint_path, replay_metadata=replay_metadata, diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 7b5443b119f..6d975c7a16a 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -17,7 +17,8 @@ import enum import json import uuid -from collections.abc import Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable +from contextlib import asynccontextmanager from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Optional @@ -27,6 +28,7 @@ from wandb import Table from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneCheckpointBarrier, DataPlaneMutationCut, PostWriteEnrichmentError, TQReplayBuffer, @@ -1445,6 +1447,7 @@ def __init__( self._num_generations_per_prompt = num_generations_per_prompt self._tq_buffer = tq_buffer self._recovery_ledger = RolloutRecoveryLedger() + self._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() self._env_handles = task_to_env self._weight_version: int = 0 # Run-wide, shared across concurrent generate_and_push calls. Safe as a plain @@ -1481,6 +1484,18 @@ def record_finalizer_dropped_prompt(self) -> None: ) self._stats.skipped += 1 + def set_data_plane_checkpoint_barrier( + self, barrier: DataPlaneCheckpointBarrier + ) -> None: + """Join streamed sibling transitions to the SC snapshot barrier.""" + self._data_plane_checkpoint_barrier = barrier + + @asynccontextmanager + async def _recovery_mutation(self) -> AsyncIterator[DataPlaneMutationCut]: + """Serialize short lineage transitions with native TQ snapshots.""" + async with self._data_plane_checkpoint_barrier.mutation() as cut: + yield cut + def reserve_prompt_group( self, cut: DataPlaneMutationCut, @@ -1838,11 +1853,13 @@ async def generate_for_finalization( ) recovery_group_id = lineage_group_id if recovery_group_id is None: - recovery_group_id = self.reserve_prompt_group( - input_sample, - target_step=target_step, - admitted=True, - ) + async with self._recovery_mutation() as cut: + recovery_group_id = self.reserve_prompt_group( + cut, + input_sample, + target_step=target_step, + admitted=True, + ) recovery_group = self._recovery_ledger.get_group(recovery_group_id) if recovery_group.phase is not PromptGroupPhase.ADMITTED: raise RuntimeError( @@ -1913,11 +1930,12 @@ async def _generate_for_finalization_attempt( from nemo_rl.experience.rollout_reassembler_actor import ReassemblyRequest assert self._tq_buffer is not None - recovery_group = self._recovery_ledger.get_group(recovery_group_id) - if recovery_group.status == PromptGroupStatus.GENERATING: - recovery_group = self._recovery_ledger.prepare_incomplete_retry( - recovery_group_id - ) + async with self._recovery_mutation() as cut: + recovery_group = self._recovery_ledger.get_group(recovery_group_id) + if recovery_group.status == PromptGroupStatus.GENERATING: + recovery_group = self._recovery_ledger.prepare_incomplete_retry( + cut, recovery_group_id + ) pending_indices = [ sibling.generation_index for sibling in recovery_group.siblings @@ -1958,13 +1976,15 @@ async def _record_streamed_completion( raise ValueError( "token-capture completion must contain its Gate rollout ID" ) - self._recovery_ledger.mark_sibling_sealed( - group_id, - generation_index=generation_index, - gate_rollout_id=gate_rollout_id, - receipt=receipt, - reward=completion.reward, - ) + async with self._recovery_mutation() as cut: + self._recovery_ledger.mark_sibling_sealed( + cut, + group_id, + generation_index=generation_index, + gate_rollout_id=gate_rollout_id, + receipt=receipt, + reward=completion.reward, + ) mask_sample_by_index: dict[int, bool] = {} @@ -1985,9 +2005,12 @@ async def _record_completion( inflight_registry[group_id] = (current_task, start_version) try: if pending_indices: - self._recovery_ledger.mark_group_dispatched( - group_id, generation_indices=pending_indices - ) + async with self._recovery_mutation() as cut: + self._recovery_ledger.mark_group_dispatched( + cut, + group_id, + generation_indices=pending_indices, + ) await self.run_rollout( attempt_input_sample, rollout_ids=list(rollout_ids), @@ -1997,12 +2020,13 @@ async def _record_completion( finally: if inflight_registry is not None: inflight_registry.pop(group_id, None) - ( - physical_rollout_ids, - canonical_sample_ids, - receipts, - rewards, - ) = self._recovery_ledger.finalization_inputs(group_id) + async with self._recovery_mutation(): + ( + physical_rollout_ids, + canonical_sample_ids, + receipts, + rewards, + ) = self._recovery_ledger.finalization_inputs(group_id) request = ReassemblyRequest( group_id=group_id, rollout_ids=tuple(physical_rollout_ids), @@ -2030,14 +2054,20 @@ async def _record_completion( # yet). Their ledger files are inert — failure rows or missing # terminal rows keep any later read fail-closed. self._tq_buffer.abort(group_id) - self._recovery_ledger.abandon_unsealed(group_id) + async with self._recovery_mutation() as cut: + self._recovery_ledger.abandon_unsealed(cut, group_id) # The capture ledger has no per-rollout fail endpoint. Rows from # abandoned attempts are unreferenced and are swept with the # staging partition at run teardown. raise - async def _discard_recovery_group(self, group_id: str) -> None: + async def discard_recovery_group( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> None: """Clean known staged rows before intentionally dropping lineage.""" + cut.require_live() assert self._tq_buffer is not None group = self._recovery_ledger.get_group(group_id) staging_keys = [ @@ -2046,4 +2076,4 @@ async def _discard_recovery_group(self, group_id: str) -> None: for key in sibling.current_attempt.staging_keys ] await self._tq_buffer.clear_staging_keys(staging_keys) - self._recovery_ledger.discard_group(group_id) + self._recovery_ledger.discard_group(cut, group_id) diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index 3bc68e55b81..deba66258e9 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -324,8 +324,14 @@ def mark_group_admitted( record.start_weight_version = start_weight_version record.phase = PromptGroupPhase.ADMITTED - def bind_runtime_prompt(self, group_id: str, prompt_payload: DatumSpec) -> None: + def bind_runtime_prompt( + self, + cut: DataPlaneMutationCut, + group_id: str, + prompt_payload: DatumSpec, + ) -> None: """Attach a dataset-reconstructed prompt after identity validation.""" + cut.require_live() record = self._require_group(group_id) sample_id = prompt_payload.get("idx") if str(sample_id) != record.prompt_ref.sample_id: @@ -341,8 +347,56 @@ def bind_runtime_prompt(self, group_id: str, prompt_payload: DatumSpec) -> None: ) record.runtime_prompt_payload = prompt_payload - def prepare_incomplete_retry(self, group_id: str) -> PromptGroupRecoveryRecord: + def prepare_for_restart(self, cut: DataPlaneMutationCut) -> None: + """Turn crash-interrupted physical attempts into retryable state.""" + cut.require_live() + self.assert_checkpoint_safe() + for record in self._groups.values(): + if record.status is PromptGroupStatus.GENERATING: + self.abandon_unsealed(cut, record.group_id) + + def assert_checkpoint_safe(self) -> None: + """Reject states whose publication or optimizer outcome is ambiguous.""" + if self._open_train_step is not None: + raise RuntimeError( + "rollout recovery contains an open optimizer step; restoring " + "mid-step training ownership is not supported" + ) + unsafe = [ + record.group_id + for record in self._groups.values() + if record.status + in { + PromptGroupStatus.FINALIZING, + PromptGroupStatus.FINALIZATION_UNKNOWN, + PromptGroupStatus.CLAIMED_FOR_TRAINING, + PromptGroupStatus.APPLIED_UNCHECKPOINTED, + } + ] + if unsafe: + raise RuntimeError( + "rollout recovery contains checkpoint-unsafe group states: " + f"groups={unsafe!r}" + ) + + def expected_staging_keys(self) -> set[str]: + """Return staged token rows still owned by sealed sibling attempts.""" + return { + staging_key + for record in self._groups.values() + for sibling in record.siblings + for attempt in sibling.attempts[-1:] + if attempt.status is RolloutAttemptStatus.SEALED + for staging_key in attempt.staging_keys + } + + def prepare_incomplete_retry( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> PromptGroupRecoveryRecord: """Mint fresh physical attempts only for siblings that are not sealed.""" + cut.require_live() record = self._require_group(group_id) if record.status != PromptGroupStatus.GENERATING: raise ValueError( @@ -367,9 +421,14 @@ def prepare_incomplete_retry(self, group_id: str) -> PromptGroupRecoveryRecord: return self._copy_group(record) def mark_group_dispatched( - self, group_id: str, *, generation_indices: Optional[list[int]] = None + self, + cut: DataPlaneMutationCut, + group_id: str, + *, + generation_indices: Optional[list[int]] = None, ) -> None: """Move the selected current sibling attempts to dispatched.""" + cut.require_live() record = self._require_group(group_id) if record.phase is not PromptGroupPhase.ADMITTED: raise ValueError( @@ -447,8 +506,9 @@ def mark_sibling_sealed( ): record.status = PromptGroupStatus.READY_TO_FINALIZE - def abandon_unsealed(self, group_id: str) -> None: + def abandon_unsealed(self, cut: DataPlaneMutationCut, group_id: str) -> None: """Abandon failed attempts without destroying reusable sealed receipts.""" + cut.require_live() record = self._require_group(group_id) if record.status not in { PromptGroupStatus.GENERATING, @@ -503,7 +563,12 @@ def finalization_inputs( rewards, ) - def mark_finalization_started(self, group_id: str) -> None: + def mark_finalization_started( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> None: + cut.require_live() record = self._require_group(group_id) self._require_group_status( record, @@ -512,7 +577,12 @@ def mark_finalization_started(self, group_id: str) -> None: ) record.status = PromptGroupStatus.FINALIZING - def mark_finalization_unknown(self, group_id: str) -> None: + def mark_finalization_unknown( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> None: + cut.require_live() record = self._require_group(group_id) self._require_group_status( record, @@ -679,6 +749,7 @@ def __contains__(self, group_id: object) -> bool: def state_dict(self) -> dict[str, Any]: """Return the versioned metadata envelope used by later persistence.""" + self.assert_checkpoint_safe() groups = [] for record in self._groups.values(): prompt_payload = record.runtime_prompt_payload @@ -834,8 +905,13 @@ def from_state_dict(cls, state: dict[str, Any]) -> Self: ) return ledger - def load_state_dict(self, state: RolloutRecoveryState) -> None: + def load_state_dict( + self, + cut: DataPlaneMutationCut, + state: RolloutRecoveryState, + ) -> None: """Replace this empty ledger from a validated checkpoint envelope.""" + cut.require_live() if self._groups or self._open_train_step is not None: raise RuntimeError("cannot restore into a non-empty rollout recovery ledger") restored = self.from_state_dict(state) @@ -1001,6 +1077,8 @@ def _group_from_state( sibling.current_attempt.status == RolloutAttemptStatus.SEALED for sibling in siblings ) + if status is PromptGroupStatus.GENERATING and all_current_attempts_sealed: + raise ValueError("generating group must retain an unfinished sibling") if status in prefinalization_sealed_states and not all_current_attempts_sealed: raise ValueError( f"group state {status.value!r} requires every sibling to be sealed" diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index ef96bbf0b93..8927995e50a 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -30,6 +30,7 @@ from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, PromptGroupPhase, + RolloutAttemptStatus, RolloutRecoveryLedger, build_rollout_recovery_state, parse_rollout_recovery_state, @@ -416,6 +417,56 @@ def test_prompt_ref_rehydrates_through_a_restored_shuffled_dataloader() -> None: assert restored_ledger.get_group("unfinished").prompt_payload == owned_prompt +def test_restart_preserves_sealed_sibling_and_retries_only_interrupted_one() -> None: + ledger = RolloutRecoveryLedger() + group = _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=6, + admitted=True, + ) + _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) + sealed_attempt_id = group.siblings[0].current_attempt.attempt_id + sealed_id = group.gate_rollout_id(0) + _mutate( + lambda cut: ledger.mark_sibling_sealed( + cut, + "g7", + generation_index=0, + gate_rollout_id=sealed_id, + receipt={ + "rollout_id": sealed_id, + "manifest": [{"staging_key": "g7/sibling-0/call-0"}], + }, + reward=1.0, + ) + ) + + restored = RolloutRecoveryLedger.from_state_dict(ledger.state_dict()) + _mutate(lambda cut: restored.prepare_for_restart(cut)) + recovered_group = restored.get_group("g7") + + assert ( + recovered_group.siblings[0].current_attempt.status + is RolloutAttemptStatus.SEALED + ) + assert ( + recovered_group.siblings[1].current_attempt.status + is RolloutAttemptStatus.ABANDONED + ) + assert restored.expected_staging_keys() == {"g7/sibling-0/call-0"} + + retry = _mutate(lambda cut: restored.prepare_incomplete_retry(cut, "g7")) + assert retry.siblings[0].current_attempt.attempt_id == sealed_attempt_id + assert retry.siblings[0].current_attempt.status is RolloutAttemptStatus.SEALED + assert retry.siblings[1].current_attempt.status is RolloutAttemptStatus.RESERVED + + @pytest.mark.parametrize( "state", [ diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 32fdf367e57..2837b454527 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -389,6 +389,9 @@ def __init__(self) -> None: self._tq_buffer = None self.recovery_ledger = RolloutRecoveryLedger() + def set_data_plane_checkpoint_barrier(self, barrier: Any) -> None: + self.data_plane_checkpoint_barrier = barrier + def set_weight_version(self, version: int) -> None: self.weight_versions.append(version) diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 7e8f3b56cd7..9416e9989c1 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -43,6 +43,7 @@ from nemo_rl.data_plane import KVBatchMeta from nemo_rl.data_plane.schema import ROLLOUT_METRICS from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.experience.rollout_recovery import RolloutRecoveryLedger from nemo_rl.utils.timer import TimeoutChecker, Timer @@ -62,6 +63,20 @@ def set_data_plane_checkpoint_barrier( self.checkpoint_barrier = barrier +class _InitRolloutManager: + """Minimal rollout-manager contract for actor-init tests.""" + + def __init__(self, tq_buffer: _InitBuffer) -> None: + self._tq_buffer = tq_buffer + self.recovery_ledger = RolloutRecoveryLedger() + self.checkpoint_barrier: DataPlaneCheckpointBarrier | None = None + + def set_data_plane_checkpoint_barrier( + self, barrier: DataPlaneCheckpointBarrier + ) -> None: + self.checkpoint_barrier = barrier + + def _checkpointing_config(tmp_path) -> dict: """Minimal checkpointing block for actors built through __init__.""" return { @@ -111,7 +126,7 @@ def _actor_args_for_init(**overrides) -> SimpleNamespace: advantage_estimator=None, loss_fn=None, tq_buffer=tq_buffer, - rollout_manager=SimpleNamespace(_tq_buffer=tq_buffer), + rollout_manager=_InitRolloutManager(tq_buffer), env_handles={}, fleet_monitor=None, generation_router=None, @@ -161,7 +176,7 @@ def test_rejects_multiple_optimizer_steps_per_rl_step(monkeypatch) -> None: advantage_estimator=None, loss_fn=None, tq_buffer=tq_buffer, - rollout_manager=SimpleNamespace(_tq_buffer=tq_buffer), + rollout_manager=_InitRolloutManager(tq_buffer), env_handles={}, fleet_monitor=None, generation_router=None, diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index eeec744e3da..7dc379d8a5f 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -1043,6 +1043,10 @@ def test_native_tq_round_trip_restores_index_without_reputting_rows(self): "per_worker_token_counts": {0: 7}, } ] + assert restored_buf._rollout_ids_list == [ + list(meta.sample_ids) for meta in metas + ] + assert restored_buf._staging_keys_list == [None, None] assert restored_dp.put_calls == [] def test_checkpoint_serialization_preserves_full_result_table(self): @@ -1093,6 +1097,35 @@ def test_checkpoint_serialization_preserves_full_result_table(self): assert restored_table.columns == ["Full result"] assert restored_table.data == [['{"reward":1.0,"status":"completed"}']] + def test_token_capture_round_trip_restores_staging_cleanup_ownership(self): + plan = encode_route_plan( + RouteAssemblyPlan( + schema_version=ROUTE_PLAN_SCHEMA_VERSION, + staging_partition="rollout_staging", + spans=(RouteSpan("r0/on_chain", 0, 2, 2, 1, "0" * 64),), + cleanup_staging_keys=("r0/on_chain", "r0/off_chain"), + expected_token_length=2, + ) + ) + group = _make_group_entry("g0", weight=1) + group["meta"].tags = [ + {ROUTE_PLAN_TAG: plan} for _ in group["meta"].sample_ids + ] + state = _make_metadata_envelope([group]) + + restored = TQReplayBuffer( + MultiPartitionFakeDataPlaneClient(), + partition_id="rollout_data", + pad_value_dict={"token_ids": 0}, + staging_partition_id="rollout_staging", + ) + assert _load(restored, state) == 1 + + assert restored._rollout_ids_list == [list(group["meta"].sample_ids)] + assert restored._staging_keys_list == [ + ["r0/on_chain", "r0/off_chain"] + ] + def test_round_trip_preserves_end_weight_and_target_step(self): # start != end and a non-None target_step must survive the round-trip: # a load that swapped start/end or dropped target_step (the From ee3906b9c09289ae207e798c1b60d8854ed6628e Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 28 Aug 2026 17:05:23 -0400 Subject: [PATCH 04/28] feat(rollout): add configurable recovery granularity Signed-off-by: Anish Mahishi --- ...po_math_1B_megatron_single_controller.yaml | 9 + .../algorithms/async_utils/replay_buffer.py | 4 +- nemo_rl/algorithms/single_controller.py | 8 +- .../single_controller_utils/config.py | 57 +++++ .../single_controller_utils/setup.py | 1 + nemo_rl/experience/rollout_manager.py | 8 + nemo_rl/experience/rollout_recovery.py | 194 ++++++++++++++++-- tests/unit/experience/test_rollout_manager.py | 113 +++++++++- .../unit/experience/test_rollout_recovery.py | 149 ++++++++++++++ .../test_checkpoint_dispatch_races.py | 3 + .../single_controller/test_rollout_pump.py | 6 +- .../test_tq_replay_buffer.py | 8 +- tests/unit/test_effort_shaping.py | 2 + 13 files changed, 520 insertions(+), 42 deletions(-) diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index 90b10ba69ce..f54daaa1f7c 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -180,6 +180,15 @@ token_capture: enabled: false staging_partition: rollout_staging +# Restore-only policy for unfinished token-capture groups. "sibling" reuses +# already sealed generations; "prompt_group" regenerates every sibling when +# any sibling was unfinished at the checkpoint. Agent overrides win over task +# overrides. Live in-process retries remain sibling-level. +rollout_recovery: + default_granularity: sibling + agent_granularity_overrides: {} + task_granularity_overrides: {} + cluster: # Master ports inherit the shared 1400-1999 band from grpo_math_1B.yaml. gpus_per_node: 2 diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 1686edb819c..de33e76e64d 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -231,9 +231,7 @@ def __init__(self) -> None: self._checkpoint_active = False self._active_mutations = 0 self._mutation_depth_by_task: dict[asyncio.Task[Any], int] = {} - self._mutation_cut_by_task: dict[ - asyncio.Task[Any], DataPlaneMutationCut - ] = {} + self._mutation_cut_by_task: dict[asyncio.Task[Any], DataPlaneMutationCut] = {} @asynccontextmanager async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index a44fdaf8ea8..66ec66395b5 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1038,9 +1038,7 @@ async def _validate_rollout_recovery_inventory( clear_unreferenced: bool, ) -> None: """Require every unfinished receipt or deferred route to retain staging.""" - expected_staging_keys = ( - self._rollout_recovery_ledger.expected_staging_keys() - ) + expected_staging_keys = self._rollout_recovery_ledger.expected_staging_keys() if replay_metadata is not None: for group in replay_metadata["groups"]: for tag in group["meta"].tags or []: @@ -1410,9 +1408,7 @@ async def _finalize_with_actor( self._finalizer_metrics_by_group[request.group_id] = dict(finalized.metrics) return finalized - async def _cleanup_consumed_metas_unlocked( - self, metas: list[KVBatchMeta] - ) -> None: + async def _cleanup_consumed_metas_unlocked(self, metas: list[KVBatchMeta]) -> None: """Clear consumed ownership while holding a barrier mutation slot.""" canonical_by_partition: dict[str, list[str]] = {} staging_by_partition: dict[str, list[str]] = {} diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 3c0506498bc..a305e71b6a0 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -16,6 +16,7 @@ import math import warnings +from collections.abc import Mapping from dataclasses import dataclass, field from typing import Annotated, Any, Literal, Optional @@ -58,6 +59,7 @@ ClusterConfig, ) from nemo_rl.environments.nemo_gym import should_use_nemo_gym +from nemo_rl.experience.rollout_recovery import RecoveryGranularity from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.value import ValueConfig from nemo_rl.utils.checkpoint import CheckpointingConfig @@ -621,6 +623,58 @@ class TokenCaptureConfig(BaseModel, extra="allow"): num_reassembler_workers: PositiveInt = 2 +@dataclass(frozen=True) +class ResolvedRolloutRecovery: + """Policy coordinates stamped on one newly reserved ledger group.""" + + agent_name: Optional[str] + granularity: RecoveryGranularity + + +class RolloutRecoveryConfig(BaseModel, extra="allow"): + """Restore policy for unfinished token-capture prompt groups. + + The resolved value is persisted on each ledger group, so restoring a saved + group does not reinterpret it using a newer configuration. + """ + + default_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING + # NeMo-Gym agent_ref.name takes precedence over task_name when both match. + agent_granularity_overrides: dict[str, RecoveryGranularity] = Field( + default_factory=dict + ) + task_granularity_overrides: dict[str, RecoveryGranularity] = Field( + default_factory=dict + ) + + def resolve_for_prompt(self, prompt: Mapping[str, Any]) -> ResolvedRolloutRecovery: + """Resolve one new group using agent, then task, then the global default.""" + extra_env_info = prompt.get("extra_env_info") + agent_name: Optional[str] = None + if isinstance(extra_env_info, Mapping): + agent_ref = extra_env_info.get("agent_ref") + if agent_ref is not None and not isinstance(agent_ref, Mapping): + raise TypeError("prompt agent_ref must be a mapping or None") + if isinstance(agent_ref, Mapping): + raw_agent_name = agent_ref.get("name") + if raw_agent_name is not None and not isinstance(raw_agent_name, str): + raise TypeError("prompt agent_ref.name must be a string or None") + agent_name = raw_agent_name + if agent_name is not None: + override = self.agent_granularity_overrides.get(agent_name) + if override is not None: + return ResolvedRolloutRecovery(agent_name, override) + + task_name = prompt.get("task_name") + if task_name is not None and not isinstance(task_name, str): + raise TypeError("prompt task_name must be a string or None") + if task_name is not None: + override = self.task_granularity_overrides.get(task_name) + if override is not None: + return ResolvedRolloutRecovery(agent_name, override) + return ResolvedRolloutRecovery(agent_name, self.default_granularity) + + class MasterConfig(BaseModel, extra="allow"): # algo configs grpo: Optional[GRPOConfig] = None @@ -638,6 +692,9 @@ class MasterConfig(BaseModel, extra="allow"): reward_penalties: RewardPenaltyConfig = Field(default_factory=RewardPenaltyConfig) data_plane: DataPlaneConfig async_rl: AsyncRLConfig + rollout_recovery: RolloutRecoveryConfig = Field( + default_factory=RolloutRecoveryConfig + ) on_policy_distillation: Optional[OnPolicyDistillationConfig] = None token_capture: TokenCaptureConfig = Field(default_factory=TokenCaptureConfig) diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index a3411f11bf0..41f53c89d3d 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -1562,6 +1562,7 @@ def _build_generation_then_trainer( task_to_env=env_handles, num_generations_per_prompt=algo_cfg.num_generations_per_prompt, max_seq_len=_generation_max_seq_len(generation_config), + rollout_recovery_config=master_config.rollout_recovery, max_rollout_turns=algo_cfg.max_rollout_turns, policy_generation=generation, generation_config=generation_config, diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 6d975c7a16a..f92534237c6 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import asyncio import copy import enum @@ -86,6 +88,7 @@ RolloutCompletionCallback = Callable[[int, Completion], Awaitable[None]] if TYPE_CHECKING: + from nemo_rl.algorithms.single_controller_utils.config import RolloutRecoveryConfig from nemo_rl.experience.rollout_reassembler_actor import ReassemblyRequest @@ -1387,6 +1390,7 @@ def __init__( task_to_env: dict[str, EnvironmentInterface], num_generations_per_prompt: int, max_seq_len: int, + rollout_recovery_config: RolloutRecoveryConfig, max_rollout_turns: int = 1, policy_generation: Optional[GenerationInterface] = None, generation_config: Optional[GenerationConfig] = None, @@ -1445,6 +1449,7 @@ def __init__( ) self._tokenizer = tokenizer self._num_generations_per_prompt = num_generations_per_prompt + self._rollout_recovery_config = rollout_recovery_config self._tq_buffer = tq_buffer self._recovery_ledger = RolloutRecoveryLedger() self._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() @@ -1512,6 +1517,7 @@ def reserve_prompt_group( "rollout recovery requires every dataloader sample to contain " f"a stable integer idx, got {prompt_idx!r}" ) + recovery_policy = self._rollout_recovery_config.resolve_for_prompt(input_sample) record = self._recovery_ledger.reserve_group( cut, prompt_id=str(prompt_idx), @@ -1519,6 +1525,8 @@ def reserve_prompt_group( expected_generations=self._num_generations_per_prompt, target_step=target_step, start_weight_version=self._weight_version, + agent_name=recovery_policy.agent_name, + recovery_granularity=recovery_policy.granularity, admitted=admitted, admission_id=admission_id, ) diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index deba66258e9..f691e6e5056 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -34,7 +34,7 @@ from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneMutationCut from nemo_rl.data.interfaces import DatumSpec -ROLLOUT_RECOVERY_SCHEMA_VERSION = 2 +ROLLOUT_RECOVERY_SCHEMA_VERSION = 3 ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" RolloutRecoveryState: TypeAlias = dict[str, Any] @@ -46,6 +46,13 @@ class PromptGroupPhase(StrEnum): ADMITTED = "admitted" +class RecoveryGranularity(StrEnum): + """Unit of completed work reused after restoring an unfinished group.""" + + SIBLING = "sibling" + PROMPT_GROUP = "prompt_group" + + class RolloutAttemptStatus(StrEnum): """Lifecycle of one physical Gate execution attempt.""" @@ -92,6 +99,33 @@ def __post_init__(self) -> None: raise ValueError("prompt sample_id must not be empty") +def _validate_prompt_identity( + prompt_ref: PromptRef, + prompt_payload: DatumSpec, + *, + group_id: str, +) -> None: + """Require a runtime prompt to resolve the ledger's durable dataset key.""" + sample_id = prompt_payload.get("idx") + if isinstance(sample_id, bool) or not isinstance(sample_id, int): + raise ValueError( + f"recovery group {group_id!r} prompt payload must contain an integer idx" + ) + if str(sample_id) != prompt_ref.sample_id: + raise ValueError( + f"recovery group {group_id!r} resolved sample_id={sample_id!r}; " + f"expected {prompt_ref.sample_id!r}" + ) + task_name = prompt_payload.get("task_name") + if task_name is not None and not isinstance(task_name, str): + raise TypeError("prompt task_name must be a string or None") + if task_name != prompt_ref.task_name: + raise ValueError( + f"recovery group {group_id!r} resolved task_name={task_name!r}; " + f"expected {prompt_ref.task_name!r}" + ) + + @dataclass class RolloutAttemptRecord: """One physical attempt for a stable logical sibling.""" @@ -132,6 +166,8 @@ class PromptGroupRecoveryRecord: admission_id: str prompt_id: str prompt_ref: PromptRef + agent_name: Optional[str] + recovery_granularity: RecoveryGranularity runtime_prompt_payload: Optional[DatumSpec] expected_generations: int target_step: Optional[int] @@ -199,6 +235,15 @@ class OpenTrainStepRecord: status: TrainStepStatus = TrainStepStatus.OPEN +@dataclass(frozen=True) +class ParsedRolloutRecoveryState: + """Validated controller and ledger state loaded from one checkpoint sidecar.""" + + ledger_state: RolloutRecoveryState + batch_shortfall: dict[int, int] + sampler_stamps_target_steps: Optional[bool] + + def _new_attempt() -> RolloutAttemptRecord: return RolloutAttemptRecord( attempt_uuid=uuid.uuid4(), @@ -249,6 +294,8 @@ def reserve_group( expected_generations: int, target_step: Optional[int], start_weight_version: int, + agent_name: Optional[str] = None, + recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, admitted: bool = True, group_id: Optional[str] = None, admission_id: Optional[str] = None, @@ -290,6 +337,8 @@ def reserve_group( admission_id=admission_id, prompt_id=prompt_id, prompt_ref=prompt_ref, + agent_name=agent_name, + recovery_granularity=recovery_granularity, # Retain the immutable dataloader sample by reference instead of copying # a potentially 131k-token payload. This cache is never serialized and # is released as soon as canonical rows take over recovery ownership. @@ -333,27 +382,38 @@ def bind_runtime_prompt( """Attach a dataset-reconstructed prompt after identity validation.""" cut.require_live() record = self._require_group(group_id) - sample_id = prompt_payload.get("idx") - if str(sample_id) != record.prompt_ref.sample_id: - raise ValueError( - f"recovery group {group_id!r} resolved sample_id={sample_id!r}; " - f"expected {record.prompt_ref.sample_id!r}" - ) - task_name = prompt_payload.get("task_name") - if task_name != record.prompt_ref.task_name: - raise ValueError( - f"recovery group {group_id!r} resolved task_name={task_name!r}; " - f"expected {record.prompt_ref.task_name!r}" - ) + _validate_prompt_identity( + record.prompt_ref, + prompt_payload, + group_id=group_id, + ) record.runtime_prompt_payload = prompt_payload def prepare_for_restart(self, cut: DataPlaneMutationCut) -> None: - """Turn crash-interrupted physical attempts into retryable state.""" + """Apply each group's persisted restore policy to interrupted attempts.""" cut.require_live() self.assert_checkpoint_safe() for record in self._groups.values(): if record.status is PromptGroupStatus.GENERATING: - self.abandon_unsealed(cut, record.group_id) + if record.recovery_granularity is RecoveryGranularity.PROMPT_GROUP: + self._abandon_entire_group(record) + else: + self.abandon_unsealed(cut, record.group_id) + + @staticmethod + def _abandon_entire_group(record: PromptGroupRecoveryRecord) -> None: + """Discard every current sibling when an incomplete group is atomic. + + Sealed staging rows become unreferenced here. The controller's restore + inventory pass removes those rows from TQ before redispatch. + """ + for sibling in record.siblings: + attempt = sibling.current_attempt + attempt.status = RolloutAttemptStatus.ABANDONED + attempt.receipt = None + attempt.reward = None + attempt.staging_keys.clear() + record.status = PromptGroupStatus.GENERATING def assert_checkpoint_safe(self) -> None: """Reject states whose publication or optimizer outcome is ambiguous.""" @@ -431,9 +491,7 @@ def mark_group_dispatched( cut.require_live() record = self._require_group(group_id) if record.phase is not PromptGroupPhase.ADMITTED: - raise ValueError( - f"cannot dispatch unadmitted recovery group {group_id!r}" - ) + raise ValueError(f"cannot dispatch unadmitted recovery group {group_id!r}") if record.status != PromptGroupStatus.GENERATING: raise ValueError( f"cannot dispatch group {group_id!r} from {record.status.value!r}" @@ -772,6 +830,8 @@ def state_dict(self) -> dict[str, Any]: "sample_id": record.prompt_ref.sample_id, "task_name": record.prompt_ref.task_name, }, + "agent_name": record.agent_name, + "recovery_granularity": record.recovery_granularity.value, "expected_generations": record.expected_generations, "target_step": record.target_step, "start_weight_version": record.start_weight_version, @@ -892,9 +952,7 @@ def from_state_dict(cls, state: dict[str, Any]) -> Self: record.claimed_train_step is not None for record in ledger._groups.values() ): raise ValueError("claimed groups require an open_train_step record") - admission_states: dict[ - str, tuple[PromptGroupPhase, Optional[int]] - ] = {} + admission_states: dict[str, tuple[PromptGroupPhase, Optional[int]]] = {} for record in ledger._groups.values(): signature = (record.phase, record.target_step) previous = admission_states.setdefault(record.admission_id, signature) @@ -913,7 +971,9 @@ def load_state_dict( """Replace this empty ledger from a validated checkpoint envelope.""" cut.require_live() if self._groups or self._open_train_step is not None: - raise RuntimeError("cannot restore into a non-empty rollout recovery ledger") + raise RuntimeError( + "cannot restore into a non-empty rollout recovery ledger" + ) restored = self.from_state_dict(state) self._groups = restored._groups self._open_train_step = restored._open_train_step @@ -930,6 +990,8 @@ def _group_from_state( group_id = raw_group.get("group_id") admission_id = raw_group.get("admission_id") prompt_id = raw_group.get("prompt_id") + agent_name = raw_group.get("agent_name") + raw_recovery_granularity = raw_group.get("recovery_granularity") expected_generations = raw_group.get("expected_generations") siblings_state = raw_group.get("siblings") if not isinstance(group_id, str) or not group_id: @@ -938,6 +1000,16 @@ def _group_from_state( raise ValueError("admission_id must be a non-empty string") if not isinstance(prompt_id, str) or not prompt_id: raise ValueError("prompt_id must be a non-empty string") + if agent_name is not None and not isinstance(agent_name, str): + raise ValueError("agent_name must be a string or None") + if not isinstance(raw_recovery_granularity, str): + raise ValueError("recovery_granularity must be a string") + try: + recovery_granularity = RecoveryGranularity(raw_recovery_granularity) + except ValueError as error: + raise ValueError( + f"invalid recovery_granularity={raw_recovery_granularity!r}" + ) from error if not isinstance(expected_generations, int) or expected_generations < 1: raise ValueError("expected_generations must be a positive integer") if ( @@ -1106,6 +1178,8 @@ def _group_from_state( admission_id=admission_id, prompt_id=prompt_id, prompt_ref=PromptRef(sample_id=sample_id, task_name=task_name), + agent_name=agent_name, + recovery_granularity=recovery_granularity, runtime_prompt_payload=None, expected_generations=expected_generations, target_step=target_step, @@ -1133,6 +1207,8 @@ def _copy_group(record: PromptGroupRecoveryRecord) -> PromptGroupRecoveryRecord: admission_id=record.admission_id, prompt_id=record.prompt_id, prompt_ref=record.prompt_ref, + agent_name=record.agent_name, + recovery_granularity=record.recovery_granularity, runtime_prompt_payload=record.runtime_prompt_payload, expected_generations=record.expected_generations, target_step=record.target_step, @@ -1175,3 +1251,77 @@ def _require_open_train_step(self, train_step: int) -> OpenTrainStepRecord: if open_step is None or open_step.train_step != train_step: raise ValueError(f"train step {train_step} is not open") return open_step + + +def _validate_batch_shortfall(value: object) -> dict[int, int]: + """Return a defensive copy of per-step permanent rollout losses.""" + if not isinstance(value, dict): + raise TypeError("rollout recovery batch_shortfall must be a dictionary") + batch_shortfall: dict[int, int] = {} + for step, count in value.items(): + if ( + isinstance(step, bool) + or not isinstance(step, int) + or step < 0 + or isinstance(count, bool) + or not isinstance(count, int) + or count < 0 + ): + raise ValueError( + "rollout recovery batch_shortfall entries must contain " + f"non-negative integer steps and counts, got {step!r}: {count!r}" + ) + batch_shortfall[step] = count + return batch_shortfall + + +def build_rollout_recovery_state( + ledger: RolloutRecoveryLedger, + *, + batch_shortfall: dict[int, int], + sampler_stamps_target_steps: bool, +) -> RolloutRecoveryState: + """Build the complete versioned sidecar from ledger and controller state.""" + if not isinstance(sampler_stamps_target_steps, bool): + raise TypeError( + "rollout recovery sampler_stamps_target_steps must be a boolean" + ) + state = ledger.state_dict() + state["batch_shortfall"] = _validate_batch_shortfall(batch_shortfall) + state["sampler_stamps_target_steps"] = sampler_stamps_target_steps + return state + + +def parse_rollout_recovery_state(state: object) -> ParsedRolloutRecoveryState: + """Validate and split a complete checkpoint sidecar by runtime owner.""" + if not isinstance(state, dict): + raise TypeError( + "rollout recovery sidecar must contain a dictionary, got " + f"{type(state).__name__}" + ) + if state.get("schema_version") != ROLLOUT_RECOVERY_SCHEMA_VERSION: + raise ValueError( + "unsupported rollout recovery schema_version=" + f"{state.get('schema_version')!r}; expected " + f"{ROLLOUT_RECOVERY_SCHEMA_VERSION}" + ) + groups = state.get("groups") + if not isinstance(groups, list): + raise TypeError("rollout recovery groups must be a list") + + raw_sampler_stamps = state.get("sampler_stamps_target_steps") + if raw_sampler_stamps is not None and not isinstance(raw_sampler_stamps, bool): + raise TypeError( + "rollout recovery sampler_stamps_target_steps must be a boolean" + ) + + ledger_state: RolloutRecoveryState = { + "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "groups": groups, + "open_train_step": state.get("open_train_step"), + } + return ParsedRolloutRecoveryState( + ledger_state=ledger_state, + batch_shortfall=_validate_batch_shortfall(state.get("batch_shortfall", {})), + sampler_stamps_target_steps=raw_sampler_stamps, + ) diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 56ef5d25b0c..e611dde88da 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -38,6 +38,7 @@ DataPlaneCheckpointBarrier, PostWriteEnrichmentError, ) +from nemo_rl.algorithms.single_controller_utils.config import RolloutRecoveryConfig from nemo_rl.data.collate_fn import rl_collate_fn from nemo_rl.data.datasets.response_datasets import NemoGymDataset from nemo_rl.data.interfaces import DatumSpec @@ -60,7 +61,10 @@ RolloutRetryPolicy, RolloutStats, ) -from nemo_rl.experience.rollout_recovery import RolloutRecoveryLedger +from nemo_rl.experience.rollout_recovery import ( + RecoveryGranularity, + RolloutRecoveryLedger, +) from nemo_rl.experience.rollouts import ( run_async_multi_turn_rollout, run_async_nemo_gym_rollout, @@ -239,8 +243,10 @@ def _make_manager( mgr._impl = impl mgr._tokenizer = None mgr._num_generations_per_prompt = 1 + mgr._rollout_recovery_config = RolloutRecoveryConfig() mgr._tq_buffer = buffer mgr._recovery_ledger = RolloutRecoveryLedger() + mgr._data_plane_checkpoint_barrier = buffer.data_plane_checkpoint_barrier mgr._env_handles = {} mgr._weight_version = 0 mgr._retry_policy = ( @@ -471,6 +477,26 @@ async def _assert_ledger_owns_inflight_prompt(_sample): assert buf._slots == [group_id] assert buf.commit_calls[0][0] == group_id + def test_reservation_persists_the_resolved_agent_recovery_policy(self): + mgr = _make_manager(_FakeBuffer(), _FakeImpl()) + mgr._rollout_recovery_config = RolloutRecoveryConfig( + agent_granularity_overrides={ + "genrm_agent": RecoveryGranularity.PROMPT_GROUP + } + ) + prompt = { + "idx": 0, + "message_log": [], + "task_name": "nemo_gym", + "extra_env_info": {"agent_ref": {"name": "genrm_agent"}}, + } + + group_id = mgr.reserve_prompt_group(prompt, target_step=0) + group = mgr.recovery_ledger.get_group(group_id) + + assert group.agent_name == "genrm_agent" + assert group.recovery_granularity is RecoveryGranularity.PROMPT_GROUP + def test_skipped_tracked_prompt_remains_owned_for_controller_handoff(self): async def _fail_rollout(_sample): raise RuntimeError("bad prompt") @@ -514,6 +540,8 @@ def test_tracked_dispatch_rejects_changed_generations_per_prompt(self): expected_generations=2, target_step=0, start_weight_version=0, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ), ) @@ -660,6 +688,7 @@ def test_rollout_manager_raises_without_impl_params(): "task_to_env": {}, "num_generations_per_prompt": 1, "max_seq_len": 1, + "rollout_recovery_config": RolloutRecoveryConfig(), } with pytest.raises(AssertionError, match="num_generations_per_prompt must be >= 1"): @@ -681,6 +710,7 @@ def test_rollout_manager_forwards_mask_env_flagged_samples(): "task_to_env": {}, "num_generations_per_prompt": 1, "max_seq_len": 1, + "rollout_recovery_config": RolloutRecoveryConfig(), "generation_config": { "stop_strings": None, "stop_token_ids": None, @@ -1048,6 +1078,7 @@ def test_async_rollout_manager( task_to_env=task_to_env, num_generations_per_prompt=num_generations, max_seq_len=max_seq_len, + rollout_recovery_config=RolloutRecoveryConfig(), max_rollout_turns=max_rollout_turns, policy_generation=vllm_generation, ) @@ -1108,6 +1139,7 @@ def test_async_rollout_manager_truncation( task_to_env=task_to_env, num_generations_per_prompt=num_generations, max_seq_len=max_seq_len, + rollout_recovery_config=RolloutRecoveryConfig(), max_rollout_turns=max_rollout_turns, policy_generation=vllm_generation, ) @@ -1173,6 +1205,7 @@ def test_async_rollout_manager_matches_original( task_to_env=task_to_env, num_generations_per_prompt=num_generations, max_seq_len=max_seq_len, + rollout_recovery_config=RolloutRecoveryConfig(), max_rollout_turns=max_rollout_turns, policy_generation=vllm_generation, ) @@ -1307,6 +1340,7 @@ def test_async_nemo_gym_rollout_manager( task_to_env={"nemo_gym": nemo_gym}, num_generations_per_prompt=num_generations, max_seq_len=nemo_gym_vllm_generation.cfg["vllm_cfg"]["max_model_len"], + rollout_recovery_config=RolloutRecoveryConfig(), generation_config=nemo_gym_vllm_generation.cfg, ) record = asyncio.run(manager.run_rollout(single_prompt)) @@ -1429,6 +1463,7 @@ async def _collect_original_results(): task_to_env={"nemo_gym": nemo_gym}, num_generations_per_prompt=num_generations, max_seq_len=nemo_gym_vllm_generation.cfg["vllm_cfg"]["max_model_len"], + rollout_recovery_config=RolloutRecoveryConfig(), generation_config=nemo_gym_vllm_generation.cfg, ) record = asyncio.run(manager.run_rollout(single_prompt)) @@ -1557,10 +1592,12 @@ def _make_capture_manager( num_generations=2, retry_policy: RolloutRetryPolicy | None = None, instance_configs=None, + recovery_config: RolloutRecoveryConfig | None = None, ): mgr = object.__new__(RolloutManager) mgr._tokenizer = None mgr._num_generations_per_prompt = num_generations + mgr._rollout_recovery_config = recovery_config or RolloutRecoveryConfig() mgr._tq_buffer = buf mgr._weight_version = 7 mgr._retry_policy = ( @@ -1572,10 +1609,12 @@ def _make_capture_manager( mgr._skipped_prompts = 0 mgr._consecutive_infra_drops = 0 mgr._recovery_ledger = RolloutRecoveryLedger() + mgr._data_plane_checkpoint_barrier = buf.data_plane_checkpoint_barrier class _CaptureImpl: def __init__(self): self.seen_rollout_ids = None + self.seen_generation_indices = None async def run_rollout( self, @@ -1586,6 +1625,7 @@ async def run_rollout( on_completion=None, ): self.seen_rollout_ids = rollout_ids + self.seen_generation_indices = list(generation_indices or []) if on_run is not None: await on_run(_sample) indices = generation_indices or list(range(len(rollout_ids))) @@ -1707,9 +1747,14 @@ async def _fail_once(_sample): ) assert mgr.stats.as_metrics()["rollout/redispatch_total"] == 1.0 - def test_reuses_sealed_sibling_and_redispatches_only_incomplete_one(self): + def test_prompt_group_policy_does_not_change_live_infrastructure_retry(self): buf = _FakeCaptureBuffer() - mgr = _make_capture_manager(buf) + mgr = _make_capture_manager( + buf, + recovery_config=RolloutRecoveryConfig( + default_granularity=RecoveryGranularity.PROMPT_GROUP + ), + ) mgr._retry_policy = RolloutRetryPolicy.single_attempt( max_infra_attempts=2, backoff_base_s=0.0, @@ -1762,3 +1807,65 @@ async def run_rollout( assert second_ids[0] == first_ids[0] assert second_ids[1] != first_ids[1] assert request.rollout_ids == (second_ids[0], second_ids[1]) + + def test_prompt_group_restore_redispatches_every_sibling(self): + recovery_config = RolloutRecoveryConfig( + default_granularity=RecoveryGranularity.PROMPT_GROUP + ) + first = _make_capture_manager( + _FakeCaptureBuffer(), recovery_config=recovery_config + ) + prompt = {"prompt": "p", "idx": 9} + group_id = _with_cut( + first._tq_buffer, + lambda cut: first.reserve_prompt_group(cut, prompt, target_step=7), + ) + _with_cut( + first._tq_buffer, + lambda cut: first.recovery_ledger.mark_group_dispatched(cut, group_id), + ) + group = first.recovery_ledger.get_group(group_id) + gate_id = group.gate_rollout_id(0) + _with_cut( + first._tq_buffer, + lambda cut: first.recovery_ledger.mark_sibling_sealed( + cut, + group_id, + generation_index=0, + gate_rollout_id=gate_id, + receipt={ + "rollout_id": gate_id, + "manifest": [{"staging_key": f"{gate_id}/call"}], + }, + reward=0.5, + ), + ) + + restored = _make_capture_manager( + _FakeCaptureBuffer(), + # The saved group policy wins over the new process configuration. + recovery_config=RolloutRecoveryConfig( + default_granularity=RecoveryGranularity.SIBLING + ), + ) + _with_cut( + restored._tq_buffer, + lambda cut: restored.recovery_ledger.load_state_dict( + cut, first.recovery_ledger.state_dict() + ), + ) + _with_cut( + restored._tq_buffer, + lambda cut: restored.recovery_ledger.prepare_for_restart(cut), + ) + + request = _run( + restored.generate_for_finalization( + prompt, + target_step=7, + lineage_group_id=group_id, + ) + ) + + assert request is not None + assert restored._impl.seen_generation_indices == [0, 1] diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 8927995e50a..6f7f7b46672 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -26,10 +26,12 @@ DataPlaneCheckpointBarrier, DataPlaneMutationCut, ) +from nemo_rl.algorithms.single_controller_utils.config import RolloutRecoveryConfig from nemo_rl.data.interfaces import DatumSpec from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, PromptGroupPhase, + RecoveryGranularity, RolloutAttemptStatus, RolloutRecoveryLedger, build_rollout_recovery_state, @@ -103,6 +105,8 @@ def _group_state( "sample_id": str(idx), "task_name": None, }, + "agent_name": None, + "recovery_granularity": "sibling", "expected_generations": 2, "target_step": target_step, "start_weight_version": 7, @@ -121,6 +125,8 @@ def test_ledger_round_trip_preserves_group_ownership() -> None: expected_generations=2, target_step=7, start_weight_version=6, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -236,6 +242,35 @@ async def exercise() -> None: asyncio.run(exercise()) +def test_recovery_config_resolves_agent_then_task_then_default() -> None: + config = RolloutRecoveryConfig( + default_granularity=RecoveryGranularity.SIBLING, + agent_granularity_overrides={"genrm_agent": RecoveryGranularity.PROMPT_GROUP}, + task_granularity_overrides={ + "math": RecoveryGranularity.PROMPT_GROUP, + "agent_wins": RecoveryGranularity.SIBLING, + }, + ) + + agent_policy = config.resolve_for_prompt( + { + "task_name": "agent_wins", + "extra_env_info": {"agent_ref": {"name": "genrm_agent"}}, + } + ) + task_policy = config.resolve_for_prompt( + {"task_name": "math", "extra_env_info": None} + ) + default_policy = config.resolve_for_prompt( + {"task_name": "other", "extra_env_info": None} + ) + + assert agent_policy.agent_name == "genrm_agent" + assert agent_policy.granularity is RecoveryGranularity.PROMPT_GROUP + assert task_policy.granularity is RecoveryGranularity.PROMPT_GROUP + assert default_policy.granularity is RecoveryGranularity.SIBLING + + def test_target_step_none_does_not_mean_unadmitted() -> None: ledger = RolloutRecoveryLedger() record = _reserve( @@ -247,6 +282,8 @@ def test_target_step_none_does_not_mean_unadmitted() -> None: expected_generations=2, target_step=None, start_weight_version=6, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -265,6 +302,8 @@ def test_reserved_group_can_be_admitted_exactly_once() -> None: expected_generations=2, target_step=None, start_weight_version=6, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, admitted=False, ) @@ -300,6 +339,8 @@ def test_canonical_groups_are_discarded_without_touching_unfinished_groups() -> expected_generations=2, target_step=7, start_weight_version=7, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -319,6 +360,8 @@ def test_state_dict_stores_a_prompt_ref_without_the_full_payload() -> None: expected_generations=2, target_step=7, start_weight_version=7, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -346,6 +389,8 @@ def test_bind_runtime_prompt_accepts_changed_content_with_the_same_identity() -> expected_generations=2, target_step=7, start_weight_version=7, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) restored = RolloutRecoveryLedger() @@ -369,6 +414,8 @@ def test_bind_runtime_prompt_rejects_the_wrong_dataset_sample() -> None: expected_generations=2, target_step=7, start_weight_version=7, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) restored = RolloutRecoveryLedger() @@ -396,6 +443,8 @@ def test_prompt_ref_rehydrates_through_a_restored_shuffled_dataloader() -> None: expected_generations=2, target_step=1, start_weight_version=0, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) ledger_state = ledger.state_dict() @@ -428,6 +477,8 @@ def test_restart_preserves_sealed_sibling_and_retries_only_interrupted_one() -> expected_generations=2, target_step=7, start_weight_version=6, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) @@ -467,6 +518,104 @@ def test_restart_preserves_sealed_sibling_and_retries_only_interrupted_one() -> assert retry.siblings[1].current_attempt.status is RolloutAttemptStatus.RESERVED +def test_prompt_group_restart_retries_every_sibling_when_one_is_unfinished() -> None: + ledger = RolloutRecoveryLedger() + group = _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=6, + agent_name="genrm_agent", + recovery_granularity=RecoveryGranularity.PROMPT_GROUP, + admitted=True, + ) + _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) + sealed_id = group.gate_rollout_id(0) + _mutate( + lambda cut: ledger.mark_sibling_sealed( + cut, + "g7", + generation_index=0, + gate_rollout_id=sealed_id, + receipt={ + "rollout_id": sealed_id, + "manifest": [{"staging_key": "g7/sibling-0/call-0"}], + }, + reward=1.0, + ) + ) + + state = ledger.state_dict() + assert state["groups"][0]["agent_name"] == "genrm_agent" + assert state["groups"][0]["recovery_granularity"] == "prompt_group" + + restored = RolloutRecoveryLedger.from_state_dict(state) + _mutate(lambda cut: restored.prepare_for_restart(cut)) + recovered = restored.get_group("g7") + + assert [sibling.current_attempt.status for sibling in recovered.siblings] == [ + RolloutAttemptStatus.ABANDONED, + RolloutAttemptStatus.ABANDONED, + ] + assert restored.expected_staging_keys() == set() + + retry = _mutate(lambda cut: restored.prepare_incomplete_retry(cut, "g7")) + assert [sibling.current_attempt.status for sibling in retry.siblings] == [ + RolloutAttemptStatus.RESERVED, + RolloutAttemptStatus.RESERVED, + ] + + +def test_prompt_group_restart_keeps_a_fully_sealed_group() -> None: + ledger = RolloutRecoveryLedger() + group = _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=6, + agent_name="genrm_agent", + recovery_granularity=RecoveryGranularity.PROMPT_GROUP, + admitted=True, + ) + _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) + for generation_index in range(2): + gate_id = group.gate_rollout_id(generation_index) + _mutate( + lambda cut, generation_index=generation_index, gate_id=gate_id: ( + ledger.mark_sibling_sealed( + cut, + "g7", + generation_index=generation_index, + gate_rollout_id=gate_id, + receipt={ + "rollout_id": gate_id, + "manifest": [ + {"staging_key": (f"g7/sibling-{generation_index}/call-0")} + ], + }, + reward=1.0, + ) + ) + ) + + restored = RolloutRecoveryLedger.from_state_dict(ledger.state_dict()) + _mutate(lambda cut: restored.prepare_for_restart(cut)) + + assert restored.get_group("g7").sealed_generation_indices == [0, 1] + assert restored.expected_staging_keys() == { + "g7/sibling-0/call-0", + "g7/sibling-1/call-0", + } + + @pytest.mark.parametrize( "state", [ diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index eb80541c928..37b97e3bd1d 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -60,6 +60,7 @@ ROLLOUT_RECOVERY_SCHEMA_VERSION, ROLLOUT_RECOVERY_STATE_FILENAME, PromptGroupPhase, + RecoveryGranularity, RolloutRecoveryLedger, build_rollout_recovery_state, ) @@ -354,6 +355,8 @@ def reserve_prompt_group( expected_generations=2, target_step=target_step, start_weight_version=7, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, admitted=admitted, admission_id=admission_id, ) diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 4061c2553f5..de7720a7c7a 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -45,13 +45,14 @@ from nemo_rl.algorithms.single_controller_utils.config import ( AsyncRLConfig, MasterConfig, + RolloutRecoveryConfig, ) from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.rollout_manager import RolloutManager, RolloutOutcome from nemo_rl.experience.rollout_recovery import ( RolloutRecoveryLedger, - RolloutRecoveryLedgerState, + RolloutRecoveryState, ) # Reuse fixtures from the experience tests; same shape as test_async_rollout_manager. @@ -931,7 +932,7 @@ async def _main() -> None: checkpoint_entered = asyncio.Event() - async def checkpoint_snapshot() -> RolloutRecoveryLedgerState: + async def checkpoint_snapshot() -> RolloutRecoveryState: async with barrier.checkpoint(): checkpoint_entered.set() return ledger.state_dict() @@ -1299,6 +1300,7 @@ def test_rollout_pump_writes_expected_tq_data( task_to_env=task_to_env, num_generations_per_prompt=num_generations, max_seq_len=max_seq_len, + rollout_recovery_config=RolloutRecoveryConfig(), max_rollout_turns=max_rollout_turns, policy_generation=vllm_generation, use_nemo_gym=False, diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 7dc379d8a5f..19693b8cbaa 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -1108,9 +1108,7 @@ def test_token_capture_round_trip_restores_staging_cleanup_ownership(self): ) ) group = _make_group_entry("g0", weight=1) - group["meta"].tags = [ - {ROUTE_PLAN_TAG: plan} for _ in group["meta"].sample_ids - ] + group["meta"].tags = [{ROUTE_PLAN_TAG: plan} for _ in group["meta"].sample_ids] state = _make_metadata_envelope([group]) restored = TQReplayBuffer( @@ -1122,9 +1120,7 @@ def test_token_capture_round_trip_restores_staging_cleanup_ownership(self): assert _load(restored, state) == 1 assert restored._rollout_ids_list == [list(group["meta"].sample_ids)] - assert restored._staging_keys_list == [ - ["r0/on_chain", "r0/off_chain"] - ] + assert restored._staging_keys_list == [["r0/on_chain", "r0/off_chain"]] def test_round_trip_preserves_end_weight_and_target_step(self): # start != end and a non-None target_step must survive the round-trip: diff --git a/tests/unit/test_effort_shaping.py b/tests/unit/test_effort_shaping.py index 3be2efe74f7..e3f2e9f8cc2 100644 --- a/tests/unit/test_effort_shaping.py +++ b/tests/unit/test_effort_shaping.py @@ -17,6 +17,7 @@ import pytest +from nemo_rl.algorithms.single_controller_utils.config import RolloutRecoveryConfig from nemo_rl.experience.rollout_manager import ( AsyncNemoGymRolloutImpl, RolloutManager, @@ -331,6 +332,7 @@ def test_rollout_manager_forwards_effort_config(): "task_to_env": {}, "num_generations_per_prompt": 1, "max_seq_len": 1, + "rollout_recovery_config": RolloutRecoveryConfig(), "generation_config": _GENERATION_CONFIG, "use_nemo_gym": True, } From 6e33523bd685390f527ede4a5ed4a81291d3bfb0 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 28 Aug 2026 18:51:53 -0400 Subject: [PATCH 05/28] fix(rollout): reconcile sibling recovery tests Signed-off-by: Anish Mahishi --- nemo_rl/utils/checkpoint.py | 4 +- .../test_checkpoint_dispatch_races.py | 15 ++ .../single_controller/test_rollout_pump.py | 132 +++++++++++++++++- tests/unit/utils/test_checkpoint.py | 6 +- 4 files changed, 152 insertions(+), 5 deletions(-) diff --git a/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py index 6e4db00bde5..56f95dedef7 100644 --- a/nemo_rl/utils/checkpoint.py +++ b/nemo_rl/utils/checkpoint.py @@ -322,7 +322,9 @@ def init_tmp_checkpoint( # save config if run_config is not None: with open(save_dir / "config.yaml", "w") as f: - yaml.safe_dump(run_config.model_dump(), f) + # JSON mode converts enums and other Pydantic-supported scalar + # types to the primitive values expected by safe YAML. + yaml.safe_dump(run_config.model_dump(mode="json"), f) return Path(os.path.abspath(save_dir)) diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 37b97e3bd1d..7d1c4c20b93 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -252,6 +252,12 @@ def __init__(self, ledger: _PendingLedger) -> None: def set_weight_version(self, version: int) -> None: self.weight_version = version + def set_data_plane_checkpoint_barrier( + self, barrier: DataPlaneCheckpointBarrier + ) -> None: + """Accept the controller-owned barrier used by the production manager.""" + del barrier + def reserve_prompt_group( self, cut: DataPlaneMutationCut | None, @@ -909,6 +915,9 @@ async def exercise() -> None: controller = object.__new__(controller_cls) controller._sampler = sampler controller._rollout_manager = rollout_manager + controller._master_config = SimpleNamespace( + token_capture=SimpleNamespace(enabled=False) + ) controller._last_checkpoint_path = str(tmp_path) controller._data_plane_checkpoint_metadata = { "rollout_recovery_schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, @@ -1016,6 +1025,9 @@ async def exercise() -> None: controller = object.__new__(controller_cls) controller._sampler = sampler controller._rollout_manager = rollout_manager + controller._master_config = SimpleNamespace( + token_capture=SimpleNamespace(enabled=False) + ) controller._last_checkpoint_path = str(tmp_path) controller._data_plane_checkpoint_metadata = { "rollout_recovery_schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, @@ -1169,6 +1181,9 @@ async def exercise() -> None: controller = object.__new__(controller_cls) controller._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() controller._rollout_manager = rollout_manager + controller._master_config = SimpleNamespace( + token_capture=SimpleNamespace(enabled=False) + ) controller._last_checkpoint_path = str(tmp_path) controller._data_plane_checkpoint_metadata = { "rollout_recovery_schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index de7720a7c7a..2353bc44cff 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -1162,9 +1162,13 @@ async def _main() -> None: release_finalizers = asyncio.Event() finalizers_started = 0 - async def _delayed_finalize(request: Any) -> bool: + async def _delayed_finalize( + request: Any, + *, + target_step: int | None = None, + ) -> bool: nonlocal finalizers_started - del request + del request, target_step finalizers_started += 1 await release_finalizers.wait() return True @@ -1221,6 +1225,130 @@ async def _delayed_finalize(request: Any) -> bool: asyncio.run(_main()) +@pytest.mark.parametrize("committed", [False, True]) +def test_actor_finalization_discards_recovery_ledger_ownership( + committed: bool, +) -> None: + class _RecoveryCaptureManager: + def __init__(self) -> None: + self.recovery_ledger = RolloutRecoveryLedger() + self.stats = SimpleNamespace(committed=0) + + def reserve_prompt_group( + self, + cut: DataPlaneMutationCut, + prompt: Any, + *, + target_step: int | None, + admitted: bool, + admission_id: str, + ) -> str: + return self.recovery_ledger.reserve_group( + cut, + prompt_id=str(prompt["idx"]), + prompt_payload=prompt, + expected_generations=1, + target_step=target_step, + start_weight_version=0, + admitted=admitted, + admission_id=admission_id, + ).group_id + + def mark_prompt_group_admitted( + self, + cut: DataPlaneMutationCut, + group_id: str, + *, + target_step: int | None, + ) -> None: + self.recovery_ledger.mark_group_admitted( + cut, + group_id, + target_step=target_step, + start_weight_version=0, + ) + + def discard_prompt_group( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> None: + self.recovery_ledger.discard_group(cut, group_id) + + async def generate_for_finalization( + self, + prompt: Any, + *, + target_step: int | None, + inflight_registry: dict[str, tuple[asyncio.Task[None], int]], + lineage_group_id: str, + ) -> Any: + del prompt, target_step, inflight_registry + assert self.recovery_ledger.get_group(lineage_group_id) + return SimpleNamespace(group_id=lineage_group_id) + + async def _main() -> None: + manager = _RecoveryCaptureManager() + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._async_cfg = SimpleNamespace( + max_inflight_prompts=1, + diagnostics=False, + rollout_failure=_failure_cfg(), + ) + ctrl._master_config = SimpleNamespace( + grpo=GRPOConfig.model_construct(max_num_epochs=1), + token_capture=SimpleNamespace(min_valid_fraction_per_group=None), + ) + ctrl._algo_cfg = ctrl._master_config.grpo + ctrl._buffer = _RecordingBuffer() + ctrl._rollout_manager = manager + ctrl._sampler = WindowedSampler(None, max_staleness_versions=1) + ctrl._dataloader = [ + BatchedDataDict( + { + "idx": [7], + "message_log": [[{"role": "user", "content": "prompt"}]], + } + ) + ] + ctrl._rollout_permitted = asyncio.Event() + ctrl._rollout_permitted.set() + ctrl._rollout_exhausted = asyncio.Event() + ctrl._buffer_capacity = asyncio.Semaphore(1) + ctrl._inflight_rollouts = 0 + ctrl._inflight_by_group_id = {} + ctrl._dispatched_rollouts = set() + ctrl._trainer_version = 0 + ctrl._current_epoch = 0 + ctrl._sampler_stamps_target_steps = False + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + _init_pump_ledgers(ctrl) + ctrl._finalizer_actors = [object()] + ctrl._rollout_recovery_enabled = True + + async def _finalize( + request: Any, + ) -> Any: + async with ctrl._data_plane_checkpoint_barrier.mutation() as cut: + manager.recovery_ledger.discard_group(cut, request.group_id) + if not committed: + return None + return SimpleNamespace(valid_row_count=1, total_row_count=1) + + ctrl._finalize_with_actor = _finalize + + await ctrl._rollout_pump() + + assert manager.recovery_ledger.groups() == [] + # A committed group transfers its permit to the train pump; a dropped + # group returns it immediately because no canonical replay row owns it. + assert ctrl._buffer_capacity._value == (0 if committed else 1) + assert ctrl._rollout_exhausted.is_set() + + asyncio.run(_main()) + + @pytest.mark.vllm def test_rollout_pump_writes_expected_tq_data( multi_step_setup_vllm_async, # noqa: F811 diff --git a/tests/unit/utils/test_checkpoint.py b/tests/unit/utils/test_checkpoint.py index 885412c1eb6..e2b342b3ff5 100644 --- a/tests/unit/utils/test_checkpoint.py +++ b/tests/unit/utils/test_checkpoint.py @@ -55,7 +55,8 @@ def test_init_tmp_checkpoint(checkpoint_manager, checkpoint_dir): step = 1 training_info = {"loss": 0.5, "tensor": torch.tensor(0.5), "numpy": np.array(0.5)} run_config = MagicMock() - run_config.model_dump.return_value = {"model": "test"} + expected_config = {"model": "test"} + run_config.model_dump.return_value = expected_config save_dir = checkpoint_manager.init_tmp_checkpoint(step, training_info, run_config) @@ -73,7 +74,8 @@ def test_init_tmp_checkpoint(checkpoint_manager, checkpoint_dir): # Check if config was saved with open(save_dir / "config.yaml", "r") as f: saved_config = yaml.safe_load(f) - assert saved_config == run_config.model_dump() + assert saved_config == expected_config + run_config.model_dump.assert_called_once_with(mode="json") def test_finalize_checkpoint(checkpoint_manager, checkpoint_dir): From b7fc4dc0cd45dc2aaa39aa60645e53003572ace9 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 28 Aug 2026 19:16:14 -0400 Subject: [PATCH 06/28] test(rollout): cover sibling recovery across restart Signed-off-by: Anish Mahishi --- .../L1_Functional_Tests_SingleController.sh | 3 + ...single_controller_sibling_recovery_hook.py | 193 ++++++++++++++++++ .../grpo_async_gym_single_controller.sh | 27 +-- ..._gym_single_controller_sibling_recovery.sh | 80 ++++++++ 4 files changed, 291 insertions(+), 12 deletions(-) create mode 100644 tests/functional/_single_controller_sibling_recovery_hook.py create mode 100755 tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index 5922fa9016d..ded20da8876 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -182,6 +182,9 @@ run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller # Token-capture (gate-authoritative) path: same SC+Gym smoke with the gate # custodying token lineage and the finalizer publishing training rows. run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh ++token_capture.enabled=true +# Two-process token-capture recovery: preserve one sealed sibling in TQ and +# redispatch only its unfinished peer after restoring the step checkpoint. +run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/_single_controller_sibling_recovery_hook.py b/tests/functional/_single_controller_sibling_recovery_hook.py new file mode 100644 index 00000000000..a57f3cf9b77 --- /dev/null +++ b/tests/functional/_single_controller_sibling_recovery_hook.py @@ -0,0 +1,193 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Test-only SC entrypoint for deterministic sibling-level recovery. + +The first process lets one sibling become ledger-sealed, then parks the next +completion before the ledger records it. Earlier-step completions wait for that +cut, guaranteeing that the step-1 checkpoint contains the selected partial +group. The second process records which generation indices are redispatched. +""" + +from __future__ import annotations + +import asyncio +import json +import os +from pathlib import Path +from typing import Any + +from examples import run_grpo_single_controller +from nemo_rl.experience.rollout_manager import RolloutCompletionCallback + + +class _InstrumentedNemoGymRolloutImpl: + """Delegate Gym rollouts while controlling one streamed-completion cut.""" + + def __init__( + self, + delegate: Any, + *, + recovery_ledger: Any, + events_path: Path, + block_target_step: int | None, + ) -> None: + self._delegate = delegate + self._recovery_ledger = recovery_ledger + self._events_path = events_path + self._block_target_step = block_target_step + self._selected_group_id: str | None = None + # Construct this lazily inside the Ray actor's event loop. The rollout + # manager and this test wrapper are built driver-side and serialized into + # that actor, so an eagerly created asyncio primitive could bind to the + # wrong loop. + self._selected_sibling_sealed: asyncio.Event | None = None + + def __getattr__(self, name: str) -> Any: + delegate = self.__dict__.get("_delegate") + if delegate is None: + raise AttributeError(name) + return getattr(delegate, name) + + def _append_event(self, event: str, **fields: Any) -> None: + self._events_path.parent.mkdir(parents=True, exist_ok=True) + with self._events_path.open("a", encoding="utf-8") as stream: + stream.write(json.dumps({"event": event, **fields}, sort_keys=True) + "\n") + + def _find_group(self, rollout_ids: list[str]) -> Any: + rollout_id_set = set(rollout_ids) + matches = [ + group + for group in self._recovery_ledger.groups() + if rollout_id_set.intersection(group.gate_rollout_ids) + ] + if len(matches) != 1: + raise RuntimeError( + "sibling recovery hook could not uniquely resolve rollout IDs " + f"to one ledger group: ids={rollout_ids!r}, matches=" + f"{[group.group_id for group in matches]!r}" + ) + return matches[0] + + def _sibling_sealed_event(self) -> asyncio.Event: + if self._selected_sibling_sealed is None: + self._selected_sibling_sealed = asyncio.Event() + return self._selected_sibling_sealed + + async def run_rollout( + self, + input_sample: Any, + *, + rollout_ids: list[str] | None = None, + generation_indices: list[int] | None = None, + on_completion: RolloutCompletionCallback | None = None, + ) -> Any: + if rollout_ids is None or generation_indices is None or on_completion is None: + raise RuntimeError( + "sibling recovery hook requires the token-capture rollout path" + ) + + group = self._find_group(rollout_ids) + indices = list(generation_indices) + fields = { + "group_id": group.group_id, + "prompt_idx": int(input_sample["idx"]), + "target_step": group.target_step, + "generation_indices": indices, + "rollout_ids": list(rollout_ids), + } + self._append_event("dispatch", **fields) + + selected = False + if ( + self._block_target_step is not None + and group.target_step == self._block_target_step + and self._selected_group_id is None + ): + # Selection occurs before the first await, so concurrent rollout tasks + # cannot select two groups on this event loop. + self._selected_group_id = group.group_id + selected = True + + sealed_in_selected_call = False + + async def _instrumented_completion( + generation_index: int, completion: Any + ) -> None: + nonlocal sealed_in_selected_call + completion_fields = { + **fields, + "generation_index": generation_index, + "rollout_id": rollout_ids[generation_index], + } + if selected: + if not sealed_in_selected_call: + await on_completion(generation_index, completion) + sealed_in_selected_call = True + self._append_event("sibling_sealed", **completion_fields) + self._sibling_sealed_event().set() + return + + self._append_event("blocked_before_ledger_seal", **completion_fields) + print( + "sibling recovery functional hook: blocked group=" + f"{group.group_id} generation_index={generation_index}", + flush=True, + ) + await asyncio.Event().wait() + + # Do not allow the preceding train step to complete until the selected + # lookahead group has one durable sibling. This removes checkpoint timing + # from the test: the step-1 save cannot occur before the intended cut. + if ( + self._block_target_step is not None + and group.target_step is not None + and group.target_step < self._block_target_step + ): + await self._sibling_sealed_event().wait() + await on_completion(generation_index, completion) + + result = await self._delegate.run_rollout( + input_sample, + rollout_ids=rollout_ids, + generation_indices=indices, + on_completion=_instrumented_completion, + ) + self._append_event("capture_complete", **fields) + return result + + +_original_setup_single_controller = run_grpo_single_controller.setup_single_controller + + +def _setup_with_sibling_recovery_hook(*args: Any, **kwargs: Any) -> Any: + actor_args, timing_metrics = _original_setup_single_controller(*args, **kwargs) + events_path = Path(os.environ["SC_SIBLING_RECOVERY_TEST_EVENTS"]) + raw_target_step = os.environ.get("SC_SIBLING_RECOVERY_BLOCK_TARGET_STEP") + block_target_step = int(raw_target_step) if raw_target_step is not None else None + manager = actor_args.rollout_manager + manager._impl = _InstrumentedNemoGymRolloutImpl( + manager._impl, + recovery_ledger=manager.recovery_ledger, + events_path=events_path, + block_target_step=block_target_step, + ) + return actor_args, timing_metrics + + +run_grpo_single_controller.setup_single_controller = _setup_with_sibling_recovery_hook + + +if __name__ == "__main__": + run_grpo_single_controller.main() diff --git a/tests/functional/grpo_async_gym_single_controller.sh b/tests/functional/grpo_async_gym_single_controller.sh index b9c5734c98d..323cd96de95 100755 --- a/tests/functional/grpo_async_gym_single_controller.sh +++ b/tests/functional/grpo_async_gym_single_controller.sh @@ -17,6 +17,7 @@ JSON_METRICS=$EXP_DIR/metrics.json RUN_LOG=$EXP_DIR/run.log CHECKPOINT_DIR=$EXP_DIR/checkpoints DATA_DIR=$EXP_DIR/data +SC_ENTRYPOINT=${SC_TEST_ENTRYPOINT:-$PROJECT_ROOT/examples/run_grpo_single_controller.py} export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} rm -rf $EXP_DIR $LOG_DIR @@ -56,7 +57,7 @@ jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/da jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ - $PROJECT_ROOT/examples/run_grpo_single_controller.py \ + $SC_ENTRYPOINT \ --config $PROJECT_ROOT/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml \ policy.model_name=Qwen/Qwen3-0.6B \ policy.dtensor_cfg.enabled=false \ @@ -105,16 +106,18 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE $@ \ 2>&1 | tee $RUN_LOG -uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS +if [[ "${RUN_CONVERGENCE_CHECKS:-1}" == "1" ]]; then + uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS -EXTRA_CHECKS=() -if [[ "$*" == *token_capture.enabled=true* ]]; then - # Nonzero only when a finalizer actor ran, i.e. capture really was on. - EXTRA_CHECKS+=('max(data["train/finalize/total_ms"]) > 0') -fi + EXTRA_CHECKS=() + if [[ "$*" == *token_capture.enabled=true* ]]; then + # Nonzero only when a finalizer actor ran, i.e. capture really was on. + EXTRA_CHECKS+=('max(data["train/finalize/total_ms"]) > 0') + fi -# Observed to be between 0.8-1.3 -uv run tests/check_metrics.py $JSON_METRICS \ - 'median(data["train/gen_kl_error"]) < 1.3' \ - 'max(data["train/reward"]) > 0' \ - "${EXTRA_CHECKS[@]}" + # Observed to be between 0.8-1.3 + uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' \ + 'max(data["train/reward"]) > 0' \ + "${EXTRA_CHECKS[@]}" +fi diff --git a/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh b/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh new file mode 100755 index 00000000000..3a51ab8657e --- /dev/null +++ b/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh @@ -0,0 +1,80 @@ +#!/bin/bash +# Two-process functional test for sibling-level token-capture recovery. + +set -eou pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd) +PROJECT_ROOT=$(realpath "$SCRIPT_DIR/../..") +BASE_TEST=$SCRIPT_DIR/grpo_async_gym_single_controller.sh +TEST_DIR=$SCRIPT_DIR/grpo_async_gym_single_controller_sibling_recovery +CHECKPOINT_DIR=$TEST_DIR/checkpoints +BASE_RUN_LOG=$SCRIPT_DIR/grpo_async_gym_single_controller/run.log +PHASE1_LOG=$TEST_DIR/phase1.log +PHASE2_LOG=$TEST_DIR/phase2.log +PHASE1_EVENTS=$TEST_DIR/phase1-events.jsonl +PHASE2_EVENTS=$TEST_DIR/phase2-events.jsonl +RECOVERY_HOOK=$SCRIPT_DIR/_single_controller_sibling_recovery_hook.py + +rm -rf "$TEST_DIR" +mkdir -p "$TEST_DIR" + +COMMON_OVERRIDES=( + checkpointing.enabled=true + checkpointing.checkpoint_dir="$CHECKPOINT_DIR" + checkpointing.save_period=1 + checkpointing.save_data_plane=true + ++token_capture.enabled=true + ++rollout_recovery.default_granularity=sibling + async_rl.sampler.name=in_order + async_rl.sampler.max_lookahead_versions=1 + async_rl.max_inflight_prompts=8 + async_rl.max_buffered_rollouts=8 + ++async_rl.rollout_failure.native.generation_timeout_s=120 + ++async_rl.stall_watchdog.interval_s=10 + ++async_rl.stall_watchdog.stall_timeout_s=300 + ++async_rl.stall_watchdog.stall_action=abort + grpo.max_num_steps=2 +) + +echo "=== Phase 1: checkpoint one sealed and one unfinished sibling ===" +# Target step 1 is the lookahead batch while trainer step 1 consumes target +# step 0. The hook holds target-step-0 completions until the selected group has +# sealed one sibling, then parks its next completion before the ledger update. +SC_TEST_ENTRYPOINT="$RECOVERY_HOOK" \ +SC_SIBLING_RECOVERY_TEST_EVENTS="$PHASE1_EVENTS" \ +SC_SIBLING_RECOVERY_BLOCK_TARGET_STEP=1 \ +RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" \ + "${COMMON_OVERRIDES[@]}" \ + checkpointing.checkpoint_must_save_by=0:0:0:1 +cp "$BASE_RUN_LOG" "$PHASE1_LOG" + +STEP1=$CHECKPOINT_DIR/step_1 +test -d "$STEP1/data_plane" +test -f "$STEP1/replay_buffer_metadata.pt" +test -f "$STEP1/rollout_recovery.pt" +PARTIAL_GROUP_ID=$(uv run --directory "$PROJECT_ROOT" --no-sync python -c \ + 'import json, sys; events = [json.loads(line) for line in open(sys.argv[1])]; sealed = [event for event in events if event["event"] == "sibling_sealed"]; blocked = [event for event in events if event["event"] == "blocked_before_ledger_seal"]; assert len(sealed) == 1, sealed; assert len(blocked) == 1, blocked; assert sealed[0]["group_id"] == blocked[0]["group_id"], (sealed, blocked); assert sealed[0]["generation_index"] != blocked[0]["generation_index"], (sealed, blocked); print(sealed[0]["group_id"])' \ + "$PHASE1_EVENTS") +uv run --directory "$PROJECT_ROOT" --no-sync python -c \ + 'import sys, torch; state = torch.load(sys.argv[1], weights_only=True); group_id = sys.argv[2]; groups = [group for group in state["groups"] if group["group_id"] == group_id]; assert len(groups) == 1, state; group = groups[0]; statuses = [sibling["attempts"][-1]["status"] for sibling in group["siblings"]]; assert group["recovery_granularity"] == "sibling", group; assert sorted(statuses) == ["dispatched", "sealed"], statuses' \ + "$STEP1/rollout_recovery.pt" "$PARTIAL_GROUP_ID" + +echo "=== Phase 2: reuse the sealed sibling and regenerate only its peer ===" +SC_TEST_ENTRYPOINT="$RECOVERY_HOOK" \ +SC_SIBLING_RECOVERY_TEST_EVENTS="$PHASE2_EVENTS" \ +RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" \ + "${COMMON_OVERRIDES[@]}" +cp "$BASE_RUN_LOG" "$PHASE2_LOG" + +grep -q "Native TQ checkpoint restored and validated" "$PHASE2_LOG" +grep -q "Loaded .* unfinished rollout group(s)" "$PHASE2_LOG" +test -d "$CHECKPOINT_DIR/step_2/data_plane" +test -f "$CHECKPOINT_DIR/step_2/rollout_recovery.pt" +uv run --directory "$PROJECT_ROOT" --no-sync python -c \ + 'import json, sys, torch, uuid; phase1 = torch.load(sys.argv[1], weights_only=True); group_id = sys.argv[2]; events = [json.loads(line) for line in open(sys.argv[3])]; group = next(group for group in phase1["groups"] if group["group_id"] == group_id); old_attempts = ["{}_g{}_a{}".format(group_id, i, uuid.UUID(bytes=sibling["attempts"][-1]["attempt_uuid"]).hex) for i, sibling in enumerate(group["siblings"])]; sealed_index = next(i for i, sibling in enumerate(group["siblings"]) if sibling["attempts"][-1]["status"] == "sealed"); unfinished_index = 1 - sealed_index; dispatches = [event for event in events if event["event"] == "dispatch" and event["group_id"] == group_id]; assert len(dispatches) == 1, dispatches; dispatch = dispatches[0]; assert dispatch["generation_indices"] == [unfinished_index], dispatch; assert dispatch["rollout_ids"][sealed_index] == old_attempts[sealed_index], dispatch; assert dispatch["rollout_ids"][unfinished_index] != old_attempts[unfinished_index], dispatch' \ + "$STEP1/rollout_recovery.pt" "$PARTIAL_GROUP_ID" "$PHASE2_EVENTS" +uv run --directory "$PROJECT_ROOT" --no-sync python -c \ + 'import sys, torch; state = torch.load(sys.argv[1], weights_only=True); group_id = sys.argv[2]; assert group_id not in {group["group_id"] for group in state["groups"]}, state' \ + "$CHECKPOINT_DIR/step_2/rollout_recovery.pt" "$PARTIAL_GROUP_ID" + +echo "Sibling-level token-capture recovery functional test passed." From 402fc4d903062f49cab38f3186c47c43f10e14f3 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 28 Aug 2026 22:36:15 -0400 Subject: [PATCH 07/28] test(rollout): add data-plane checkpoint override Signed-off-by: Anish Mahishi --- .../grpo_async_gym_single_controller_sibling_recovery.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh b/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh index 3a51ab8657e..7c7f01e0c31 100755 --- a/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh @@ -22,7 +22,7 @@ COMMON_OVERRIDES=( checkpointing.enabled=true checkpointing.checkpoint_dir="$CHECKPOINT_DIR" checkpointing.save_period=1 - checkpointing.save_data_plane=true + +checkpointing.save_data_plane=true ++token_capture.enabled=true ++rollout_recovery.default_granularity=sibling async_rl.sampler.name=in_order From d06318d72cf1cf58e469c126de4936151e83093c Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 28 Aug 2026 22:46:21 -0400 Subject: [PATCH 08/28] test(rollout): validate sibling recovery config Signed-off-by: Anish Mahishi --- ..._gym_single_controller_sibling_recovery.sh | 3 +- tests/unit/single_controller/test_setup.py | 78 ++++++++++++++++++- 2 files changed, 79 insertions(+), 2 deletions(-) diff --git a/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh b/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh index 7c7f01e0c31..1d2182f25c2 100755 --- a/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh @@ -21,6 +21,7 @@ mkdir -p "$TEST_DIR" COMMON_OVERRIDES=( checkpointing.enabled=true checkpointing.checkpoint_dir="$CHECKPOINT_DIR" + checkpointing.metric_name=null checkpointing.save_period=1 +checkpointing.save_data_plane=true ++token_capture.enabled=true @@ -29,7 +30,7 @@ COMMON_OVERRIDES=( async_rl.sampler.max_lookahead_versions=1 async_rl.max_inflight_prompts=8 async_rl.max_buffered_rollouts=8 - ++async_rl.rollout_failure.native.generation_timeout_s=120 + ++async_rl.rollout_failure.nemo_gym.rollout_timeout_s=120 ++async_rl.stall_watchdog.interval_s=10 ++async_rl.stall_watchdog.stall_timeout_s=300 ++async_rl.stall_watchdog.stall_action=abort diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 25d695f5ad0..a50d4143e79 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -64,7 +64,11 @@ from nemo_rl.data_plane.schema import SC_ROLLOUT_SCHEMA_FIELDS from nemo_rl.experience.rollouts import EffortLevelsConfig from nemo_rl.models.generation.megatron.megatron_generation import MegatronGeneration -from nemo_rl.utils.config import load_config, register_omegaconf_resolvers +from nemo_rl.utils.config import ( + load_config, + parse_hydra_overrides, + register_omegaconf_resolvers, +) # Captured at import, before the patched_factories fixture swaps it for a mock. _REAL_BUILD_GENERATION = sc_setup_mod._build_generation @@ -485,6 +489,78 @@ def test_build_trainer_initializes_reference_model_only_for_nonzero_kl( ) +def test_sibling_recovery_functional_config_resolves_to_runtime_contract(): + """The two-phase Gym recovery fixture must pass SC config validation.""" + register_omegaconf_resolvers() + repo_root = Path(__file__).resolve().parents[3] + config = load_config( + repo_root / "examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml" + ) + overrides = [ + "policy.model_name=Qwen/Qwen3-0.6B", + "policy.dtensor_cfg.enabled=false", + "policy.megatron_cfg.enabled=true", + "policy.megatron_cfg.tensor_model_parallel_size=1", + "policy.megatron_cfg.pipeline_model_parallel_size=1", + "policy.megatron_cfg.expert_model_parallel_size=1", + "policy.megatron_cfg.context_parallel_size=1", + "policy.megatron_cfg.sequence_parallel=false", + "policy.generation.vllm_cfg.tensor_parallel_size=1", + "policy.generation.vllm_cfg.async_engine=true", + "policy.max_total_sequence_length=512", + "policy.generation.colocated.enabled=false", + "policy.generation.colocated.resources.num_nodes=1", + "policy.generation.colocated.resources.gpus_per_node=1", + "grpo.num_prompts_per_step=4", + "grpo.num_generations_per_prompt=2", + "grpo.max_num_steps=2", + "grpo.val_period=-1", + "grpo.val_at_start=false", + "grpo.async_grpo=null", + "policy.train_global_batch_size=8", + "policy.train_micro_batch_size=1", + "cluster.gpus_per_node=2", + "loss_fn.reference_policy_kl_penalty=0.01", + "grpo.skip_reference_policy_logprobs_calculation=false", + "loss_fn.use_importance_sampling_correction=true", + "checkpointing.enabled=true", + "checkpointing.checkpoint_dir=/tmp/sibling-recovery-checkpoints", + "checkpointing.metric_name=null", + "checkpointing.save_period=1", + "+checkpointing.save_data_plane=true", + "++data_plane.enabled=true", + "++data_plane.impl=transfer_queue", + "++data_plane.backend=simple", + "++data_plane.simple.storage_capacity=1000000", + "++data_plane.simple.num_storage_units=2", + "++data_plane.claim_meta_poll_interval_s=0.5", + "++token_capture.enabled=true", + "++rollout_recovery.default_granularity=sibling", + "++async_rl.sampler.name=in_order", + "++async_rl.sampler.max_lookahead_versions=1", + "++async_rl.min_groups_for_streaming_train=4", + "++async_rl.max_inflight_prompts=8", + "++async_rl.max_buffered_rollouts=8", + "++async_rl.rollout_failure.nemo_gym.rollout_timeout_s=120", + "++async_rl.stall_watchdog.interval_s=10", + "++async_rl.stall_watchdog.stall_timeout_s=300", + "++async_rl.stall_watchdog.stall_action=abort", + ] + + resolved = OmegaConf.to_container( + parse_hydra_overrides(config, overrides), + resolve=True, + ) + + assert isinstance(resolved, dict) + master_config = MasterConfig.model_validate(resolved) + validate_single_controller_config(master_config) + assert master_config.checkpointing["metric_name"] is None + assert master_config.checkpointing["save_data_plane"] is True + assert master_config.async_rl.rollout_failure.native.generation_timeout_s is None + assert master_config.async_rl.rollout_failure.nemo_gym.rollout_timeout_s == 120 + + class TestSetup: """setup arg validation + actor_args assembly.""" From d966d3aaa8f923812394f2249d1ade6b72e43a33 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 28 Aug 2026 23:10:27 -0400 Subject: [PATCH 09/28] fix(rollout): propagate prompt index through finalizer Signed-off-by: Anish Mahishi --- nemo_rl/experience/rollout_manager.py | 1 + nemo_rl/experience/rollout_reassembler_actor.py | 1 + tests/unit/data_plane/test_rollout_reassembler.py | 13 +++++++------ tests/unit/experience/test_rollout_manager.py | 4 ++++ .../experience/test_rollout_reassembler_actor.py | 1 + .../single_controller/test_finalizer_lifecycle.py | 1 + 6 files changed, 15 insertions(+), 6 deletions(-) diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index f92534237c6..2208cdebb46 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -2037,6 +2037,7 @@ async def _record_completion( ) = self._recovery_ledger.finalization_inputs(group_id) request = ReassemblyRequest( group_id=group_id, + prompt_idx=int(recovery_group.prompt_id), rollout_ids=tuple(physical_rollout_ids), canonical_sample_ids=tuple(canonical_sample_ids), receipts=tuple(receipts), diff --git a/nemo_rl/experience/rollout_reassembler_actor.py b/nemo_rl/experience/rollout_reassembler_actor.py index d22923f5caf..aefd80f3337 100644 --- a/nemo_rl/experience/rollout_reassembler_actor.py +++ b/nemo_rl/experience/rollout_reassembler_actor.py @@ -52,6 +52,7 @@ class ReassemblyRequest: """Metadata-only input for one prompt group's finalization.""" group_id: str + prompt_idx: int rollout_ids: tuple[str, ...] canonical_sample_ids: tuple[str, ...] receipts: tuple[Optional[dict[str, Any]], ...] diff --git a/tests/unit/data_plane/test_rollout_reassembler.py b/tests/unit/data_plane/test_rollout_reassembler.py index 0173c4773e9..e3cc8de2d58 100644 --- a/tests/unit/data_plane/test_rollout_reassembler.py +++ b/tests/unit/data_plane/test_rollout_reassembler.py @@ -214,12 +214,13 @@ def test_finalize_group_publishes_n_rows_with_placeholder(tq_client, partitions) [1.0, 0.0], mask_sample=[True, False], fallback_weight_version=9, - prompt_idx=0, + prompt_idx=17, loss_multiplier=0.25, ) assert not finalized.dropped assert finalized.meta is not None assert finalized.meta.sample_ids == rollout_ids + assert [tag["prompt_idx"] for tag in finalized.meta.tags] == [17, 17] # Group staleness comes from the valid rollout's calls (wv 4), not the fallback. assert (finalized.group_min_wv, finalized.group_max_wv) == (4, 4) assert finalized.metrics["finalize/invalid_row_rate"] == 0.5 @@ -275,7 +276,7 @@ def test_finalize_group_maps_physical_attempt_to_stable_canonical_id( [1.0], mask_sample=[False], fallback_weight_version=4, - prompt_idx=0, + prompt_idx=17, canonical_sample_ids=[canonical_id], ) @@ -296,7 +297,7 @@ def test_finalize_group_reports_valid_and_total_row_counts(tq_client, partitions [0.0, 0.0], mask_sample=[False] * 2, fallback_weight_version=3, - prompt_idx=0, + prompt_idx=17, ) assert not finalized.dropped assert finalized.meta is not None @@ -442,7 +443,7 @@ def test_finalize_group_publishes_routed_experts(tq_client, r3_partitions): [1.0, 0.0], mask_sample=[False] * 2, fallback_weight_version=9, - prompt_idx=0, + prompt_idx=17, ) assert not finalized.dropped assert "routed_experts" in finalized.meta.fields @@ -498,7 +499,7 @@ def test_finalize_group_router_replay_without_routes_fails_loudly( [1.0], mask_sample=[False], fallback_weight_version=9, - prompt_idx=0, + prompt_idx=17, ) @@ -592,7 +593,7 @@ def test_deferred_finalizer_publishes_plans_and_worker_replays_routes( [1.0, 0.0], mask_sample=[False] * 2, fallback_weight_version=9, - prompt_idx=0, + prompt_idx=17, ) assert finalized.meta is not None diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index e611dde88da..83a5ddbc774 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -1695,6 +1695,7 @@ def test_mints_ids_and_returns_metadata_request(self): ) assert mgr._impl.seen_rollout_ids == attempt_ids assert request.group_id == group_id + assert request.prompt_idx == 0 assert request.rollout_ids == tuple(attempt_ids) assert request.canonical_sample_ids == tuple(canonical_ids) assert [r["rollout_id"] for r in request.receipts] == attempt_ids @@ -1736,6 +1737,7 @@ async def _fail_once(_sample): request = _run(mgr.generate_for_finalization({"prompt": "p", "idx": 4})) assert request is not None + assert request.prompt_idx == 4 assert attempts == 2 assert len(buf.reserve_rollout_ids) == 2 assert buf.reserve_rollout_ids[0] != buf.reserve_rollout_ids[1] @@ -1801,6 +1803,7 @@ async def run_rollout( request = _run(mgr.generate_for_finalization({"prompt": "p", "idx": 9})) assert request is not None + assert request.prompt_idx == 9 assert impl.generation_indices == [[0, 1], [1]] first_ids, second_ids = buf.reserve_rollout_ids assert first_ids is not None and second_ids is not None @@ -1868,4 +1871,5 @@ def test_prompt_group_restore_redispatches_every_sibling(self): ) assert request is not None + assert request.prompt_idx == 9 assert restored._impl.seen_generation_indices == [0, 1] diff --git a/tests/unit/experience/test_rollout_reassembler_actor.py b/tests/unit/experience/test_rollout_reassembler_actor.py index f8705f207e5..931076c9034 100644 --- a/tests/unit/experience/test_rollout_reassembler_actor.py +++ b/tests/unit/experience/test_rollout_reassembler_actor.py @@ -34,6 +34,7 @@ def _request() -> ReassemblyRequest: return ReassemblyRequest( group_id="group", + prompt_idx=17, rollout_ids=("group_g0",), canonical_sample_ids=("group_g0",), receipts=( diff --git a/tests/unit/single_controller/test_finalizer_lifecycle.py b/tests/unit/single_controller/test_finalizer_lifecycle.py index 47f417bfd64..4da6f92fec9 100644 --- a/tests/unit/single_controller/test_finalizer_lifecycle.py +++ b/tests/unit/single_controller/test_finalizer_lifecycle.py @@ -76,6 +76,7 @@ async def clear_samples(self, *, sample_ids: list[str], partition_id: str) -> No def _request() -> ReassemblyRequest: return ReassemblyRequest( group_id="group", + prompt_idx=17, rollout_ids=("group_g0",), canonical_sample_ids=("group_g0",), receipts=( From 079f6f677bcc82fa8edebaa8374c3b432a1bc8f0 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 29 Aug 2026 14:45:57 -0400 Subject: [PATCH 10/28] fix(rollout): harden sibling recovery checkpointing Signed-off-by: Anish Mahishi --- ...po_math_1B_megatron_single_controller.yaml | 9 + nemo_rl/experience/rollout_manager.py | 28 ++-- tests/unit/experience/test_rollout_manager.py | 22 ++- .../unit/experience/test_rollout_recovery.py | 154 ++++++++++++++++++ .../single_controller/test_checkpointing.py | 122 ++++++++++++++ 5 files changed, 323 insertions(+), 12 deletions(-) diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 6c4fd99c7c9..1f711bc0c30 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -218,6 +218,15 @@ token_capture: enabled: false staging_partition: rollout_staging +# Restore-only policy for unfinished token-capture groups. "sibling" reuses +# already sealed generations; "prompt_group" regenerates every sibling when +# any sibling was unfinished at the checkpoint. Agent overrides win over task +# overrides. Live in-process retries remain sibling-level. +rollout_recovery: + default_granularity: sibling + agent_granularity_overrides: {} + task_granularity_overrides: {} + cluster: # Master ports inherit the shared 1400-1999 band from ppo_math_1B.yaml. gpus_per_node: 2 diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 2208cdebb46..2b53d3b528e 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -1452,7 +1452,7 @@ def __init__( self._rollout_recovery_config = rollout_recovery_config self._tq_buffer = tq_buffer self._recovery_ledger = RolloutRecoveryLedger() - self._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + self._data_plane_checkpoint_barrier: Optional[DataPlaneCheckpointBarrier] = None self._env_handles = task_to_env self._weight_version: int = 0 # Run-wide, shared across concurrent generate_and_push calls. Safe as a plain @@ -1493,12 +1493,22 @@ def set_data_plane_checkpoint_barrier( self, barrier: DataPlaneCheckpointBarrier ) -> None: """Join streamed sibling transitions to the SC snapshot barrier.""" + if self._data_plane_checkpoint_barrier is not None: + raise RuntimeError( + "RolloutManager data-plane checkpoint barrier is already bound" + ) self._data_plane_checkpoint_barrier = barrier @asynccontextmanager async def _recovery_mutation(self) -> AsyncIterator[DataPlaneMutationCut]: """Serialize short lineage transitions with native TQ snapshots.""" - async with self._data_plane_checkpoint_barrier.mutation() as cut: + barrier = self._data_plane_checkpoint_barrier + if barrier is None: + raise RuntimeError( + "RolloutManager must be bound to the SingleController data-plane " + "checkpoint barrier before mutating rollout recovery state" + ) + async with barrier.mutation() as cut: yield cut def reserve_prompt_group( @@ -2028,16 +2038,14 @@ async def _record_completion( finally: if inflight_registry is not None: inflight_registry.pop(group_id, None) - async with self._recovery_mutation(): - ( - physical_rollout_ids, - canonical_sample_ids, - receipts, - rewards, - ) = self._recovery_ledger.finalization_inputs(group_id) + ( + physical_rollout_ids, + canonical_sample_ids, + receipts, + rewards, + ) = self._recovery_ledger.finalization_inputs(group_id) request = ReassemblyRequest( group_id=group_id, - prompt_idx=int(recovery_group.prompt_id), rollout_ids=tuple(physical_rollout_ids), canonical_sample_ids=tuple(canonical_sample_ids), receipts=tuple(receipts), diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 83a5ddbc774..39982589751 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -478,7 +478,8 @@ async def _assert_ledger_owns_inflight_prompt(_sample): assert buf.commit_calls[0][0] == group_id def test_reservation_persists_the_resolved_agent_recovery_policy(self): - mgr = _make_manager(_FakeBuffer(), _FakeImpl()) + buf = _FakeBuffer() + mgr = _make_manager(buf, _FakeImpl()) mgr._rollout_recovery_config = RolloutRecoveryConfig( agent_granularity_overrides={ "genrm_agent": RecoveryGranularity.PROMPT_GROUP @@ -491,12 +492,29 @@ def test_reservation_persists_the_resolved_agent_recovery_policy(self): "extra_env_info": {"agent_ref": {"name": "genrm_agent"}}, } - group_id = mgr.reserve_prompt_group(prompt, target_step=0) + group_id = _with_cut( + buf, + lambda cut: mgr.reserve_prompt_group(cut, prompt, target_step=0), + ) group = mgr.recovery_ledger.get_group(group_id) assert group.agent_name == "genrm_agent" assert group.recovery_granularity is RecoveryGranularity.PROMPT_GROUP + def test_recovery_mutation_requires_the_controller_barrier(self): + mgr = _make_manager(_FakeBuffer(), _FakeImpl()) + mgr._data_plane_checkpoint_barrier = None + + async def mutate() -> None: + async with mgr._recovery_mutation(): + pass + + with pytest.raises( + RuntimeError, + match="must be bound to the SingleController data-plane checkpoint barrier", + ): + _run(mutate()) + def test_skipped_tracked_prompt_remains_owned_for_controller_handoff(self): async def _fail_rollout(_sample): raise RuntimeError("bad prompt") diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 6f7f7b46672..86d9030664f 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -28,6 +28,7 @@ ) from nemo_rl.algorithms.single_controller_utils.config import RolloutRecoveryConfig from nemo_rl.data.interfaces import DatumSpec +from nemo_rl.data_plane import KVBatchMeta from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, PromptGroupPhase, @@ -271,6 +272,30 @@ def test_recovery_config_resolves_agent_then_task_then_default() -> None: assert default_policy.granularity is RecoveryGranularity.SIBLING +@pytest.mark.parametrize( + ("prompt", "error_fragment"), + [ + ( + {"extra_env_info": {"agent_ref": "genrm_agent"}}, + "agent_ref must be a mapping or None", + ), + ( + {"extra_env_info": {"agent_ref": {"name": 7}}}, + "agent_ref.name must be a string or None", + ), + ( + {"task_name": 7}, + "task_name must be a string or None", + ), + ], +) +def test_recovery_config_rejects_malformed_prompt_identity( + prompt: dict[str, Any], error_fragment: str +) -> None: + with pytest.raises(TypeError, match=error_fragment): + RolloutRecoveryConfig().resolve_for_prompt(prompt) + + def test_target_step_none_does_not_mean_unadmitted() -> None: ledger = RolloutRecoveryLedger() record = _reserve( @@ -616,6 +641,135 @@ def test_prompt_group_restart_keeps_a_fully_sealed_group() -> None: } +@pytest.mark.parametrize("unknown_outcome", [False, True]) +def test_checkpoint_rejects_ambiguous_finalization_state( + unknown_outcome: bool, +) -> None: + ledger = RolloutRecoveryLedger() + group = _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=1, + target_step=7, + start_weight_version=6, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, + admitted=True, + ) + _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) + gate_id = group.gate_rollout_id(0) + _mutate( + lambda cut: ledger.mark_sibling_sealed( + cut, + "g7", + generation_index=0, + gate_rollout_id=gate_id, + receipt={ + "rollout_id": gate_id, + "manifest": [{"staging_key": "g7/sibling-0/call-0"}], + }, + reward=1.0, + ) + ) + _mutate(lambda cut: ledger.mark_finalization_started(cut, "g7")) + if unknown_outcome: + _mutate(lambda cut: ledger.mark_finalization_unknown(cut, "g7")) + + with pytest.raises(RuntimeError, match="checkpoint-unsafe group states"): + ledger.state_dict() + + +def test_checkpoint_rejects_an_open_optimizer_step() -> None: + ledger = RolloutRecoveryLedger() + group = _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=1, + target_step=7, + start_weight_version=6, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, + admitted=True, + ) + _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) + gate_id = group.gate_rollout_id(0) + _mutate( + lambda cut: ledger.mark_sibling_sealed( + cut, + "g7", + generation_index=0, + gate_rollout_id=gate_id, + receipt={ + "rollout_id": gate_id, + "manifest": [{"staging_key": "g7/sibling-0/call-0"}], + }, + reward=1.0, + ) + ) + _mutate(lambda cut: ledger.mark_finalization_started(cut, "g7")) + ledger.mark_group_finalized( + "g7", + meta=KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=["g7_g0"], + ), + group_min_weight_version=6, + group_max_weight_version=6, + ) + ledger.claim_groups_for_training( + ["g7"], + train_step=7, + trainer_version=7, + expected_group_count=1, + ) + + with pytest.raises(RuntimeError, match="open optimizer step"): + ledger.state_dict() + + +@pytest.mark.parametrize( + ("field", "value", "error_fragment"), + [ + ("recovery_granularity", "banana", "invalid recovery_granularity"), + ("recovery_granularity", None, "recovery_granularity must be a string"), + ("agent_name", 123, "agent_name must be a string or None"), + ], +) +def test_restore_rejects_malformed_recovery_policy_fields( + field: str, value: Any, error_fragment: str +) -> None: + ledger = RolloutRecoveryLedger() + _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=1, + target_step=7, + start_weight_version=6, + agent_name=None, + recovery_granularity=RecoveryGranularity.SIBLING, + admitted=True, + ) + state = ledger.state_dict() + group_state = state["groups"][0] + if value is None: + del group_state[field] + else: + group_state[field] = value + + with pytest.raises(ValueError, match=error_fragment): + _load(RolloutRecoveryLedger(), state) + + @pytest.mark.parametrize( "state", [ diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 2837b454527..d9aa64136af 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -80,6 +80,12 @@ from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs from nemo_rl.data.utils import load_dataloader_state from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, KVBatchMeta +from nemo_rl.data_plane.schema import ROUTE_PLAN_TAG +from nemo_rl.experience.route_plan import ( + ROUTE_PLAN_SCHEMA_VERSION, + RouteAssemblyPlan, + encode_route_plan, +) from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, ROLLOUT_RECOVERY_STATE_FILENAME, @@ -349,6 +355,25 @@ def save_checkpoint( json.dump({"user_metadata": metadata or {}}, f) +class _StagingInventoryDPClient: + """Partition-scoped fake for rollout-recovery inventory validation.""" + + def __init__(self, sample_ids: list[str], *, partition_id: str) -> None: + self.sample_ids = list(sample_ids) + self.partition_id = partition_id + self.clear_calls: list[tuple[list[str], str]] = [] + + def list_sample_ids(self, partition_id: str) -> list[str]: + assert partition_id == self.partition_id + return sorted(self.sample_ids) + + def clear_samples(self, sample_ids: list[str], partition_id: str) -> None: + assert partition_id == self.partition_id + self.clear_calls.append((list(sample_ids), partition_id)) + cleared = set(sample_ids) + self.sample_ids = [key for key in self.sample_ids if key not in cleared] + + class _BlockingDPClient(_FakeDPClient): def __init__(self) -> None: super().__init__() @@ -621,6 +646,41 @@ def _data_plane_checkpoint_metadata( } +def _sealed_recovery_ledger(staging_key: str) -> RolloutRecoveryLedger: + """Build one ledger whose only sibling owns a sealed staging row.""" + ledger = RolloutRecoveryLedger() + + async def seed() -> None: + async with DataPlaneCheckpointBarrier().mutation() as cut: + group = ledger.reserve_group( + cut, + group_id="recovery-group", + admission_id="recovery-batch", + prompt_id="7", + prompt_payload={"idx": 7, "message_log": []}, + expected_generations=1, + target_step=7, + start_weight_version=6, + admitted=True, + ) + ledger.mark_group_dispatched(cut, group.group_id) + gate_id = group.gate_rollout_id(0) + ledger.mark_sibling_sealed( + cut, + group.group_id, + generation_index=0, + gate_rollout_id=gate_id, + receipt={ + "rollout_id": gate_id, + "manifest": [{"staging_key": staging_key}], + }, + reward=1.0, + ) + + asyncio.run(seed()) + return ledger + + def _run_train_pump( mc: MasterConfig, actor_args: SingleControllerActorArgs, @@ -1161,6 +1221,68 @@ def test_tq_save_rejects_inventory_mismatch( assert not (tmp_path / "checkpoints" / "step_1").exists() + def test_rollout_recovery_inventory_rejects_missing_staging_rows(self): + staging_partition = "rollout_staging" + actor = object.__new__(_ACTOR_CLS) + actor._rollout_recovery_ledger = _sealed_recovery_ledger("sealed-key") + actor._master_config = SimpleNamespace( + token_capture=SimpleNamespace(staging_partition=staging_partition) + ) + actor._dp_client = _StagingInventoryDPClient([], partition_id=staging_partition) + + with pytest.raises(RuntimeError, match=r"missing=\['sealed-key'\]"): + asyncio.run( + actor._validate_rollout_recovery_inventory( + replay_metadata=None, + clear_unreferenced=False, + ) + ) + + def test_rollout_recovery_inventory_merges_routes_and_clears_orphans(self): + staging_partition = "rollout_staging" + route_key = "canonical-route-key" + route_plan = encode_route_plan( + RouteAssemblyPlan( + schema_version=ROUTE_PLAN_SCHEMA_VERSION, + staging_partition=staging_partition, + spans=(), + cleanup_staging_keys=(route_key,), + expected_token_length=0, + ) + ) + replay_metadata = { + "groups": [ + { + "meta": KVBatchMeta( + partition_id=_PARTITION_ID, + task_name="train", + sample_ids=["canonical-sample"], + tags=[{ROUTE_PLAN_TAG: route_plan}], + ) + } + ] + } + dp_client = _StagingInventoryDPClient( + ["sealed-key", route_key, "orphan-key"], + partition_id=staging_partition, + ) + actor = object.__new__(_ACTOR_CLS) + actor._rollout_recovery_ledger = _sealed_recovery_ledger("sealed-key") + actor._master_config = SimpleNamespace( + token_capture=SimpleNamespace(staging_partition=staging_partition) + ) + actor._dp_client = dp_client + + asyncio.run( + actor._validate_rollout_recovery_inventory( + replay_metadata=replay_metadata, # type: ignore[arg-type] + clear_unreferenced=True, + ) + ) + + assert dp_client.clear_calls == [(["orphan-key"], staging_partition)] + assert sorted(dp_client.sample_ids) == [route_key, "sealed-key"] + def test_gated_sampler_writes_authoritative_tq_checkpoint(self, tmp_path): mc = _actor_master_config( tmp_path, From 43871ff93a6e3510ee0dbe68e5d87af36caf9656 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 29 Aug 2026 21:09:13 -0400 Subject: [PATCH 11/28] fix(rollout): remove replay groups by stable id Signed-off-by: Anish Mahishi --- .../algorithms/async_utils/replay_buffer.py | 92 ++++++++++++------- .../test_tq_replay_buffer.py | 79 ++++++++++++++++ 2 files changed, 139 insertions(+), 32 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index de33e76e64d..fc5cacfc0d9 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -1104,6 +1104,8 @@ def reserve( """ if group_id is None: group_id = str(uuid.uuid4()) + if group_id in self._group_ids: + raise ValueError(f"duplicate live group_id={group_id!r}") self.meta_list.append(None) self.start_weight_list.append(weight_version) self.end_weight_list.append(-1) @@ -1238,7 +1240,8 @@ async def remove_group(self, group_id: str, *, remove_in_dp: bool = False) -> in remove_in_dp: Whether to clear rows referenced by a committed slot. Returns: - Number of removed slots (always one on success). + One when this call removes the slot, or zero if another concurrent + mutation removed it while DataPlane cleanup was awaiting. Raises: ValueError: ``group_id`` has no live slot. @@ -1249,11 +1252,9 @@ async def remove_group(self, group_id: str, *, remove_in_dp: bool = False) -> in "checkpoint barrier before removing a group" ) async with self._data_plane_checkpoint_barrier.mutation(): - try: - idx = self._group_ids.index(group_id) - except ValueError as error: - raise ValueError(f"unknown group_id={group_id!r}") from error - return await self._remove_unlocked([idx], clear_data_plane=remove_in_dp) + return await self._remove_groups_unlocked( + [group_id], clear_data_plane=remove_in_dp + ) async def clear_staging_keys(self, staging_keys: list[str]) -> None: """Clear known token-capture staging rows under the checkpoint barrier.""" @@ -1423,28 +1424,46 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: "TQReplayBuffer must be bound to the controller data-plane " "checkpoint barrier before removing groups" ) + if len(idxs) != len(set(idxs)): + raise ValueError("replay removal contains duplicate indices") + if min(idxs) < 0: + raise IndexError("replay removal indices must be non-negative") + drop_idxs = sorted(idxs, reverse=True) + if drop_idxs[0] >= len(self.meta_list): + raise IndexError( + f"TQReplayBuffer.remove: indices out of range: {drop_idxs[0]}; " + f"size={len(self.meta_list)}" + ) + # Convert the caller's transient list coordinates into durable ownership + # coordinates before the first await. Mutations are concurrent, so another + # removal may shift every list index while this task waits for the barrier + # or for DataPlane cleanup. + drop_group_ids = [self._group_ids[i] for i in drop_idxs] async with self._data_plane_checkpoint_barrier.mutation(): - drop_idxs = sorted(idxs, reverse=True) - if drop_idxs[0] >= len(self.meta_list): - raise IndexError( - f"TQReplayBuffer.remove: indices out of range: {drop_idxs[0]}; " - f"size={len(self.meta_list)}" - ) - return await self._remove_unlocked(drop_idxs, clear_data_plane=remove_in_dp) + return await self._remove_groups_unlocked( + drop_group_ids, clear_data_plane=remove_in_dp + ) - async def _remove_unlocked( - self, drop_idxs: list[int], *, clear_data_plane: bool + async def _remove_groups_unlocked( + self, group_ids: list[str], *, clear_data_plane: bool ) -> int: - """Remove validated indices while the caller owns any required lock. - - Slots are deleted by group id, not position: the clears below yield the - event loop, and abort() or a concurrent remove() can renumber the - parallel lists while they are in flight. - """ + """Remove stable groups while the caller owns a mutation slot.""" + if len(group_ids) != len(set(group_ids)): + raise ValueError("replay removal contains duplicate group IDs") + index_by_group_id = { + group_id: i for i, group_id in enumerate(self._group_ids) + } + if len(index_by_group_id) != len(self._group_ids): + raise RuntimeError("replay buffer contains duplicate live group IDs") + missing_group_ids = [ + group_id for group_id in group_ids if group_id not in index_by_group_id + ] + if missing_group_ids: + raise ValueError(f"unknown group_ids={missing_group_ids!r}") dropped_sample_ids: list[str] = [] dropped_staging_keys: list[str] = [] - drop_group_ids = [self._group_ids[i] for i in drop_idxs] - for i in drop_idxs: + for group_id in group_ids: + i = index_by_group_id[group_id] meta = self.meta_list[i] if meta is not None: dropped_sample_ids.extend(meta.sample_ids) @@ -1481,16 +1500,25 @@ async def _remove_unlocked( "may already be cleared" ) from error - removed = 0 - for group_id in drop_group_ids: - try: - idx = self._group_ids.index(group_id) - except ValueError: - continue # already dropped by a concurrent abort() or remove() - self._delete_slot(idx) - removed += 1 + # A different mutation may have removed a lower list slot while the + # DataPlane calls were awaiting. Resolve the original stable IDs again; + # never apply pre-await indices to the now-shifted parallel lists. A group + # already removed concurrently needs no second local deletion. + current_index_by_group_id = { + group_id: i for i, group_id in enumerate(self._group_ids) + } + current_drop_idxs = sorted( + ( + current_index_by_group_id[group_id] + for group_id in group_ids + if group_id in current_index_by_group_id + ), + reverse=True, + ) + for i in current_drop_idxs: + self._delete_slot(i) - return removed + return len(current_drop_idxs) def metadata_state_dict(self, *, saved_capacity: int) -> TQReplayMetadataState: """Capture the controller index for ready groups without tensor payloads. diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 19693b8cbaa..47832a85876 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -170,6 +170,24 @@ def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None raise OSError("injected rollback failure") +class BlockFirstClearDataPlaneClient(FakeDataPlaneClient): + """Pause one clear so another structural mutation can shift list indices.""" + + def __init__(self) -> None: + super().__init__() + self.clear_started = threading.Event() + self.release_clear = threading.Event() + self._blocked = False + + def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None: + if not self._blocked: + self._blocked = True + self.clear_started.set() + if not self.release_clear.wait(timeout=5): + raise TimeoutError("timed out waiting to release the first clear") + super().clear_samples(sample_ids, partition_id) + + def _run(coro): return asyncio.run(coro) @@ -363,6 +381,13 @@ async def checkpoint(tag: str) -> None: class TestTQReplayBufferReserveCommit: + def test_reserve_rejects_duplicate_live_group_id(self): + buf = _make_buffer(FakeDataPlaneClient()) + buf.reserve(weight_version=0, group_id="group-0") + + with pytest.raises(ValueError, match="duplicate live group_id"): + buf.reserve(weight_version=1, group_id="group-0") + def test_commit_waits_for_active_checkpoint(self): async def exercise() -> None: dp = FakeDataPlaneClient() @@ -691,6 +716,38 @@ async def exercise() -> tuple[FakeDataPlaneClient, int]: assert dp.clear_thread_ids assert dp.clear_thread_ids[0] != event_loop_thread_id + def test_concurrent_removal_re_resolves_stable_group_ids_after_dp_await(self): + async def exercise() -> None: + dp = BlockFirstClearDataPlaneClient() + buf = _make_buffer(dp) + group_ids = [buf.reserve(weight_version=i) for i in range(3)] + for i, group_id in enumerate(group_ids): + await buf.commit( + group_id, + _make_record(), + start_weight_version=i, + end_weight_version=i, + ) + + # Removing the final slot pauses in DataPlane. Removing the first slot + # concurrently shifts the final slot from index 2 to index 1. + remove_last = asyncio.create_task(buf.remove([2], remove_in_dp=True)) + try: + clear_started = await asyncio.to_thread(dp.clear_started.wait, 2) + assert clear_started + assert await buf.remove([0], remove_in_dp=True) == 1 + finally: + dp.release_clear.set() + assert await remove_last == 1 + + assert buf.group_ids == (group_ids[1],) + assert buf.start_weight_list == [1] + assert buf.end_weight_list == [1] + assert buf.ready_list == [True] + assert set(dp._rows) == set(buf.meta_list[0].sample_ids) + + asyncio.run(exercise()) + def test_remove_drops_indices_and_clears_dp_when_requested(self): dp = FakeDataPlaneClient() buf = _make_buffer(dp) @@ -737,6 +794,28 @@ def test_remove_rejects_out_of_range_before_mutating(self): assert dp.depth() == 2 * _N_GENS assert dp.clear_calls == [] + def test_remove_rejects_duplicate_indices_before_mutating(self): + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + _add_group(buf, weight=0) + + with pytest.raises(ValueError, match="duplicate indices"): + _run(buf.remove([0, 0], remove_in_dp=True)) + + assert buf.size() == 1 + assert dp.clear_calls == [] + + def test_remove_rejects_negative_indices_before_mutating(self): + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + _add_group(buf, weight=0) + + with pytest.raises(IndexError, match="must be non-negative"): + _run(buf.remove([-1], remove_in_dp=True)) + + assert buf.size() == 1 + assert dp.clear_calls == [] + def test_remove_empty_is_noop(self): dp = FakeDataPlaneClient() buf = _make_buffer(dp) From ff7fb980b069a8e9eec13d9193bac53efb41c358 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 29 Aug 2026 23:40:47 -0400 Subject: [PATCH 12/28] fix(rollout): apply prompt-group recovery to live retries Signed-off-by: Anish Mahishi --- ...po_math_1B_megatron_single_controller.yaml | 8 +- ...po_math_1B_megatron_single_controller.yaml | 8 +- .../single_controller_utils/config.py | 5 +- nemo_rl/experience/rollout_manager.py | 86 +++++++++++-- nemo_rl/experience/rollout_recovery.py | 120 +++++++++++++++++- .../test_rollout_generation_failures.py | 18 +++ tests/unit/experience/test_rollout_manager.py | 32 ++--- .../unit/experience/test_rollout_recovery.py | 81 +++++++----- 8 files changed, 283 insertions(+), 75 deletions(-) diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index f54daaa1f7c..8b0d71d6216 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -180,10 +180,10 @@ token_capture: enabled: false staging_partition: rollout_staging -# Restore-only policy for unfinished token-capture groups. "sibling" reuses -# already sealed generations; "prompt_group" regenerates every sibling when -# any sibling was unfinished at the checkpoint. Agent overrides win over task -# overrides. Live in-process retries remain sibling-level. +# 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. Agent overrides win +# over task overrides. rollout_recovery: default_granularity: sibling agent_granularity_overrides: {} diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 1f711bc0c30..7f548486e87 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -218,10 +218,10 @@ token_capture: enabled: false staging_partition: rollout_staging -# Restore-only policy for unfinished token-capture groups. "sibling" reuses -# already sealed generations; "prompt_group" regenerates every sibling when -# any sibling was unfinished at the checkpoint. Agent overrides win over task -# overrides. Live in-process retries remain sibling-level. +# 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. Agent overrides win +# over task overrides. rollout_recovery: default_granularity: sibling agent_granularity_overrides: {} diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index a305e71b6a0..0f7d016c86a 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -632,10 +632,11 @@ class ResolvedRolloutRecovery: class RolloutRecoveryConfig(BaseModel, extra="allow"): - """Restore policy for unfinished token-capture prompt groups. + """Retry and restore policy for unfinished token-capture prompt groups. The resolved value is persisted on each ledger group, so restoring a saved - group does not reinterpret it using a newer configuration. + group does not reinterpret it using a newer configuration. The same + granularity governs failures handled in-process and after a process restart. """ default_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 2b53d3b528e..047d9aecf0a 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -61,8 +61,10 @@ from nemo_rl.experience.rollout_recovery import ( PromptGroupPhase, PromptGroupStatus, + RecoveryGranularity, RolloutAttemptStatus, RolloutRecoveryLedger, + SiblingSealResult, ) from nemo_rl.experience.rollouts import ( EffortLevelsConfig, @@ -427,6 +429,7 @@ async def run_rollout( rollout_ids: Optional[list[str]] = None, generation_indices: Optional[list[int]] = None, on_completion: Optional[RolloutCompletionCallback] = None, + recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, ) -> PromptGroupRecord: """Run num_generations_per_prompt rollouts for one prompt. @@ -446,6 +449,9 @@ async def run_rollout( assert on_completion is None, ( "streamed completion callbacks are only supported on the NeMo-Gym path" ) + assert recovery_granularity is RecoveryGranularity.SIBLING, ( + "recovery granularity is only supported on the NeMo-Gym path" + ) timer = Timer() timer_prefix = "timing/rollout" timer.start(f"{timer_prefix}/total") @@ -854,6 +860,7 @@ async def run_rollout( rollout_ids: Optional[list[str]] = None, generation_indices: Optional[list[int]] = None, on_completion: Optional[RolloutCompletionCallback] = None, + recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, ) -> PromptGroupRecord: """Run num_generations_per_prompt rollouts for one prompt. @@ -881,6 +888,7 @@ async def run_rollout( timer, timer_prefix, on_completion=on_completion, + recovery_granularity=recovery_granularity, ) # Token-capture receipt rows carry empty message logs by design — the # canonical row (and any media it needs) is rebuilt by the finalizer @@ -1064,14 +1072,13 @@ async def _run_rollouts( timer_prefix: str, *, on_completion: Optional[RolloutCompletionCallback] = None, + recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, ) -> tuple[list[Completion], LLMMessageLogType, dict[str, Any]]: """Dispatch rows to NeMo-Gym; return completions, prompt, and metrics. - Rows that never arrive are re-dispatched on their own rather than by redoing the - whole group. NeMo-Gym's stream dies on the first failing row, so one bad row - takes every later row with it; at num_generations_per_prompt=16 a naive whole - group retry pays 16 generations to recover one. Completed rows are kept across - attempts, which is the same shape as the legacy collector's pending-group retry. + Sibling recovery re-dispatches only rows that never arrive. Prompt-group + recovery performs one physical Gym dispatch here and delegates a complete + cohort replacement to the outer recovery loop. """ nemo_gym_env = self._task_to_env["nemo_gym"] if not inputs: @@ -1109,15 +1116,20 @@ async def _run_rollouts( # below, and a wider annotation makes the `raise ... from last_error` at the # end unverifiable. last_error: Optional[Exception] = None + max_row_attempts = ( + 1 + if recovery_granularity is RecoveryGranularity.PROMPT_GROUP + else self._max_gym_row_attempts + ) async with _Deadline(self._timeouts.rollout_s, "NeMo-Gym prompt group"): - for attempt in range(1, self._max_gym_row_attempts + 1): + for attempt in range(1, max_row_attempts + 1): pending = [row for row in inputs if results[row["_rowidx"]] is None] if not pending: break if attempt > 1: print( f"NeMo-Gym: re-dispatching {len(pending)}/{total_rows} " - f"row(s) (attempt {attempt}/{self._max_gym_row_attempts})", + f"row(s) (attempt {attempt}/{max_row_attempts})", flush=True, ) # Row re-dispatches are invisible in redispatch_total -- they @@ -1140,7 +1152,7 @@ async def _run_rollouts( # prompt NeMo-Gym cannot serve fails the same way every time. if ( classify_rollout_failure(error) is not FailureClass.INFRA - or attempt == self._max_gym_row_attempts + or attempt == max_row_attempts ): raise else: @@ -1152,7 +1164,7 @@ async def _run_rollouts( failure = GymTransportError( "NeMo-Gym rollout stream ended before all rows arrived; missing " f"rows {missing} of {total_rows} after " - f"{self._max_gym_row_attempts} attempt(s)" + f"{max_row_attempts} attempt(s)" ) # Narrowed before the raise: pyrefly rejects an Optional in a `from` # clause, even though `raise ... from None` is legal at runtime. @@ -1580,10 +1592,12 @@ async def run_rollout( rollout_ids: Optional[list[str]] = None, generation_indices: Optional[list[int]] = None, on_completion: Optional[RolloutCompletionCallback] = None, + recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, ) -> PromptGroupRecord: if rollout_ids is None: assert generation_indices is None assert on_completion is None + assert recovery_granularity is RecoveryGranularity.SIBLING # Legacy path: keep the impl call signature byte-identical. return await self._impl.run_rollout(input_sample) return await self._impl.run_rollout( @@ -1591,6 +1605,7 @@ async def run_rollout( rollout_ids=rollout_ids, generation_indices=generation_indices, on_completion=on_completion, + recovery_granularity=recovery_granularity, ) async def generate_and_push( @@ -1865,7 +1880,7 @@ async def generate_for_finalization( inflight_registry: Optional[dict[str, tuple[asyncio.Task[None], int]]] = None, lineage_group_id: Optional[str] = None, ) -> Optional["ReassemblyRequest"]: - """Capture siblings with stable lineage and retry only unfinished work.""" + """Capture siblings with stable lineage and configured retry granularity.""" assert self._tq_buffer is not None, ( "generate_for_finalization requires tq_buffer to be set at __init__" ) @@ -1944,7 +1959,7 @@ async def _generate_for_finalization_attempt( recovery_group_id: str, inflight_registry: Optional[dict[str, tuple[asyncio.Task[None], int]]], ) -> "ReassemblyRequest": - """Dispatch only unfinished siblings and leave one reserved slot unready.""" + """Dispatch the current sibling cohort and leave one slot unready.""" from nemo_rl.experience.rollout_reassembler_actor import ReassemblyRequest assert self._tq_buffer is not None @@ -1975,6 +1990,7 @@ async def _generate_for_finalization_attempt( group_id=group_id, rollout_ids=list(rollout_ids), ) + pending_group_results: dict[int, SiblingSealResult] = {} async def _record_streamed_completion( generation_index: int, completion: Completion @@ -1994,6 +2010,53 @@ async def _record_streamed_completion( raise ValueError( "token-capture completion must contain its Gate rollout ID" ) + if not 0 <= generation_index < len(rollout_ids): + raise ValueError( + f"streamed generation index {generation_index} is outside " + f"prompt group {group_id!r}" + ) + expected_gate_rollout_id = rollout_ids[generation_index] + if gate_rollout_id != expected_gate_rollout_id: + raise ValueError( + "streamed rollout identity mismatch: " + f"result={gate_rollout_id!r}, " + f"expected={expected_gate_rollout_id!r}" + ) + if receipt.get("rollout_id") != gate_rollout_id: + raise ValueError( + "receipt rollout identity mismatch: " + f"receipt={receipt.get('rollout_id')!r}, " + f"expected={gate_rollout_id!r}" + ) + + if ( + recovery_group.recovery_granularity + is RecoveryGranularity.PROMPT_GROUP + ): + result = SiblingSealResult( + gate_rollout_id=gate_rollout_id, + receipt=receipt, + reward=completion.reward, + ) + previous = pending_group_results.get(generation_index) + if previous is not None: + if previous != result: + raise ValueError( + "conflicting duplicate prompt-group completion for " + f"generation_index={generation_index}" + ) + return + pending_group_results[generation_index] = result + if len(pending_group_results) < recovery_group.expected_generations: + return + async with self._recovery_mutation() as cut: + self._recovery_ledger.mark_group_sealed( + cut, + group_id, + pending_group_results, + ) + return + async with self._recovery_mutation() as cut: self._recovery_ledger.mark_sibling_sealed( cut, @@ -2034,6 +2097,7 @@ async def _record_completion( rollout_ids=list(rollout_ids), generation_indices=pending_indices, on_completion=_record_completion, + recovery_granularity=recovery_group.recovery_granularity, ) finally: if inflight_registry is not None: diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index f691e6e5056..e5f779be704 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -47,7 +47,7 @@ class PromptGroupPhase(StrEnum): class RecoveryGranularity(StrEnum): - """Unit of completed work reused after restoring an unfinished group.""" + """Unit of completed work reused after a live failure or process restart.""" SIBLING = "sibling" PROMPT_GROUP = "prompt_group" @@ -244,6 +244,15 @@ class ParsedRolloutRecoveryState: sampler_stamps_target_steps: Optional[bool] +@dataclass(frozen=True) +class SiblingSealResult: + """One terminal sibling result waiting for an atomic prompt-group seal.""" + + gate_rollout_id: str + receipt: dict[str, Any] + reward: float + + def _new_attempt() -> RolloutAttemptRecord: return RolloutAttemptRecord( attempt_uuid=uuid.uuid4(), @@ -455,16 +464,46 @@ def prepare_incomplete_retry( cut: DataPlaneMutationCut, group_id: str, ) -> PromptGroupRecoveryRecord: - """Mint fresh physical attempts only for siblings that are not sealed.""" + """Mint fresh physical attempts according to the persisted granularity.""" cut.require_live() record = self._require_group(group_id) if record.status != PromptGroupStatus.GENERATING: raise ValueError( f"cannot retry group {group_id!r} from status {record.status.value!r}" ) + current_statuses = [ + sibling.current_attempt.status for sibling in record.siblings + ] + retry_prompt_group = ( + record.recovery_granularity is RecoveryGranularity.PROMPT_GROUP + and any( + status + in { + RolloutAttemptStatus.ABANDONED, + RolloutAttemptStatus.FAILED, + } + for status in current_statuses + ) + ) + if retry_prompt_group and any( + status + not in { + RolloutAttemptStatus.ABANDONED, + RolloutAttemptStatus.FAILED, + } + for status in current_statuses + ): + raise ValueError( + "prompt-group retry requires every sibling attempt to be abandoned " + "or failed together" + ) + for sibling in record.siblings: attempt = sibling.current_attempt - if attempt.status == RolloutAttemptStatus.SEALED: + if ( + record.recovery_granularity is RecoveryGranularity.SIBLING + and attempt.status == RolloutAttemptStatus.SEALED + ): continue if attempt.status == RolloutAttemptStatus.RESERVED: continue @@ -522,6 +561,10 @@ def mark_sibling_sealed( """Record one streamed sibling receipt as soon as the row arrives.""" cut.require_live() record = self._require_group(group_id) + if record.recovery_granularity is RecoveryGranularity.PROMPT_GROUP: + raise ValueError( + "prompt-group recovery must seal every sibling atomically" + ) sibling = self._require_sibling(record, generation_index) attempt = sibling.current_attempt expected_gate_rollout_id = record.gate_rollout_id(generation_index) @@ -564,8 +607,71 @@ def mark_sibling_sealed( ): record.status = PromptGroupStatus.READY_TO_FINALIZE + def mark_group_sealed( + self, + cut: DataPlaneMutationCut, + group_id: str, + results: dict[int, SiblingSealResult], + ) -> None: + """Atomically seal one complete prompt-group-scoped physical cohort.""" + cut.require_live() + record = self._require_group(group_id) + if record.recovery_granularity is not RecoveryGranularity.PROMPT_GROUP: + raise ValueError( + "atomic group sealing requires prompt-group recovery granularity" + ) + if record.status is not PromptGroupStatus.GENERATING: + raise ValueError( + f"cannot seal group {group_id!r} from status {record.status.value!r}" + ) + expected_indices = set(range(record.expected_generations)) + if set(results) != expected_indices: + raise ValueError( + "prompt-group seal requires every logical sibling exactly once: " + f"expected={sorted(expected_indices)}, actual={sorted(results)}" + ) + + validated: list[ + tuple[RolloutAttemptRecord, SiblingSealResult, list[str]] + ] = [] + for generation_index in range(record.expected_generations): + result = results[generation_index] + sibling = self._require_sibling(record, generation_index) + attempt = sibling.current_attempt + expected_gate_rollout_id = record.gate_rollout_id(generation_index) + if attempt.status is not RolloutAttemptStatus.DISPATCHED: + raise ValueError( + "cannot seal logical rollout " + f"{record.logical_rollout_id(generation_index)!r} " + f"from status {attempt.status.value!r}" + ) + if result.gate_rollout_id != expected_gate_rollout_id: + raise ValueError( + "streamed rollout identity mismatch: " + f"result={result.gate_rollout_id!r}, " + f"expected={expected_gate_rollout_id!r}" + ) + if result.receipt.get("rollout_id") != expected_gate_rollout_id: + raise ValueError( + "receipt rollout identity mismatch: " + f"receipt={result.receipt.get('rollout_id')!r}, " + f"expected={expected_gate_rollout_id!r}" + ) + validated.append( + (attempt, result, _receipt_staging_keys(result.receipt)) + ) + + # Validate the complete cohort before changing any sibling. A checkpoint + # therefore observes either no committed siblings or the complete group. + for attempt, result, staging_keys in validated: + attempt.receipt = copy.deepcopy(result.receipt) + attempt.reward = float(result.reward) + attempt.staging_keys = staging_keys + attempt.status = RolloutAttemptStatus.SEALED + record.status = PromptGroupStatus.READY_TO_FINALIZE + def abandon_unsealed(self, cut: DataPlaneMutationCut, group_id: str) -> None: - """Abandon failed attempts without destroying reusable sealed receipts.""" + """Abandon failed work at the group's persisted recovery granularity.""" cut.require_live() record = self._require_group(group_id) if record.status not in { @@ -575,6 +681,12 @@ def abandon_unsealed(self, cut: DataPlaneMutationCut, group_id: str) -> None: raise ValueError( f"cannot abandon group {group_id!r} from {record.status.value!r}" ) + if ( + record.recovery_granularity is RecoveryGranularity.PROMPT_GROUP + and record.status is PromptGroupStatus.GENERATING + ): + self._abandon_entire_group(record) + return for sibling in record.siblings: attempt = sibling.current_attempt if attempt.status == RolloutAttemptStatus.SEALED: diff --git a/tests/unit/experience/test_rollout_generation_failures.py b/tests/unit/experience/test_rollout_generation_failures.py index 3b191494959..54136ca8708 100644 --- a/tests/unit/experience/test_rollout_generation_failures.py +++ b/tests/unit/experience/test_rollout_generation_failures.py @@ -54,6 +54,7 @@ _Deadline, _gather_cancelling_siblings, ) +from nemo_rl.experience.rollout_recovery import RecoveryGranularity from nemo_rl.utils.timer import Timer @@ -633,6 +634,23 @@ def test_completed_rows_survive_across_attempts(self): assert len(completions) == 5 assert sum(len(d) for d in method.dispatched) == 7 + def test_prompt_group_defers_complete_retry_to_the_outer_manager(self): + method = _PartialGymMethod(fail_after_rows=2, failures_before_success=1) + impl = _make_gym_impl(method, num_generations=4, row_attempts=3) + + with pytest.raises(ConnectionResetError, match="gym stream died"): + asyncio.run( + impl._run_rollouts( + _gym_rows(4), + Timer(), + "timing/rollout", + recovery_granularity=RecoveryGranularity.PROMPT_GROUP, + ) + ) + + assert method.dispatched == [[0, 1, 2, 3]] + assert impl._stats.gym_row_redispatches == 0 + def test_a_stale_echo_of_a_landed_row_is_rejected(self): """Re-dispatch narrows the stream; an echo of an already-landed row must not win. diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 39982589751..bd216562333 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -1633,6 +1633,7 @@ class _CaptureImpl: def __init__(self): self.seen_rollout_ids = None self.seen_generation_indices = None + self.seen_recovery_granularity = None async def run_rollout( self, @@ -1641,9 +1642,11 @@ async def run_rollout( rollout_ids=None, generation_indices=None, on_completion=None, + recovery_granularity=RecoveryGranularity.SIBLING, ): self.seen_rollout_ids = rollout_ids self.seen_generation_indices = list(generation_indices or []) + self.seen_recovery_granularity = recovery_granularity if on_run is not None: await on_run(_sample) indices = generation_indices or list(range(len(rollout_ids))) @@ -1767,7 +1770,7 @@ async def _fail_once(_sample): ) assert mgr.stats.as_metrics()["rollout/redispatch_total"] == 1.0 - def test_prompt_group_policy_does_not_change_live_infrastructure_retry(self): + def test_prompt_group_policy_retries_the_complete_live_cohort(self): buf = _FakeCaptureBuffer() mgr = _make_capture_manager( buf, @@ -1783,6 +1786,7 @@ def test_prompt_group_policy_does_not_change_live_infrastructure_retry(self): class _PartialCaptureImpl: def __init__(self): self.generation_indices: list[list[int]] = [] + self.recovery_granularities: list[RecoveryGranularity] = [] async def run_rollout( self, @@ -1791,9 +1795,11 @@ async def run_rollout( rollout_ids=None, generation_indices=None, on_completion=None, + recovery_granularity=RecoveryGranularity.SIBLING, ): indices = list(generation_indices) self.generation_indices.append(indices) + self.recovery_granularities.append(recovery_granularity) completions = [] for generation_index in indices: rollout_id = rollout_ids[generation_index] @@ -1822,10 +1828,14 @@ async def run_rollout( assert request is not None assert request.prompt_idx == 9 - assert impl.generation_indices == [[0, 1], [1]] + assert impl.generation_indices == [[0, 1], [0, 1]] + assert impl.recovery_granularities == [ + RecoveryGranularity.PROMPT_GROUP, + RecoveryGranularity.PROMPT_GROUP, + ] first_ids, second_ids = buf.reserve_rollout_ids assert first_ids is not None and second_ids is not None - assert second_ids[0] == first_ids[0] + assert second_ids[0] != first_ids[0] assert second_ids[1] != first_ids[1] assert request.rollout_ids == (second_ids[0], second_ids[1]) @@ -1845,22 +1855,6 @@ def test_prompt_group_restore_redispatches_every_sibling(self): first._tq_buffer, lambda cut: first.recovery_ledger.mark_group_dispatched(cut, group_id), ) - group = first.recovery_ledger.get_group(group_id) - gate_id = group.gate_rollout_id(0) - _with_cut( - first._tq_buffer, - lambda cut: first.recovery_ledger.mark_sibling_sealed( - cut, - group_id, - generation_index=0, - gate_rollout_id=gate_id, - receipt={ - "rollout_id": gate_id, - "manifest": [{"staging_key": f"{gate_id}/call"}], - }, - reward=0.5, - ), - ) restored = _make_capture_manager( _FakeCaptureBuffer(), diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 86d9030664f..a5ccd4d9c42 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -35,6 +35,7 @@ RecoveryGranularity, RolloutAttemptStatus, RolloutRecoveryLedger, + SiblingSealResult, build_rollout_recovery_state, parse_rollout_recovery_state, ) @@ -545,7 +546,7 @@ def test_restart_preserves_sealed_sibling_and_retries_only_interrupted_one() -> def test_prompt_group_restart_retries_every_sibling_when_one_is_unfinished() -> None: ledger = RolloutRecoveryLedger() - group = _reserve( + _reserve( ledger, group_id="g7", admission_id="batch-7", @@ -559,20 +560,6 @@ def test_prompt_group_restart_retries_every_sibling_when_one_is_unfinished() -> admitted=True, ) _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) - sealed_id = group.gate_rollout_id(0) - _mutate( - lambda cut: ledger.mark_sibling_sealed( - cut, - "g7", - generation_index=0, - gate_rollout_id=sealed_id, - receipt={ - "rollout_id": sealed_id, - "manifest": [{"staging_key": "g7/sibling-0/call-0"}], - }, - reward=1.0, - ) - ) state = ledger.state_dict() assert state["groups"][0]["agent_name"] == "genrm_agent" @@ -611,25 +598,20 @@ def test_prompt_group_restart_keeps_a_fully_sealed_group() -> None: admitted=True, ) _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) + results = {} for generation_index in range(2): gate_id = group.gate_rollout_id(generation_index) - _mutate( - lambda cut, generation_index=generation_index, gate_id=gate_id: ( - ledger.mark_sibling_sealed( - cut, - "g7", - generation_index=generation_index, - gate_rollout_id=gate_id, - receipt={ - "rollout_id": gate_id, - "manifest": [ - {"staging_key": (f"g7/sibling-{generation_index}/call-0")} - ], - }, - reward=1.0, - ) - ) + results[generation_index] = SiblingSealResult( + gate_rollout_id=gate_id, + receipt={ + "rollout_id": gate_id, + "manifest": [ + {"staging_key": f"g7/sibling-{generation_index}/call-0"} + ], + }, + reward=1.0, ) + _mutate(lambda cut: ledger.mark_group_sealed(cut, "g7", results)) restored = RolloutRecoveryLedger.from_state_dict(ledger.state_dict()) _mutate(lambda cut: restored.prepare_for_restart(cut)) @@ -641,6 +623,43 @@ def test_prompt_group_restart_keeps_a_fully_sealed_group() -> None: } +def test_prompt_group_seal_is_atomic() -> None: + ledger = RolloutRecoveryLedger() + group = _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=6, + agent_name="genrm_agent", + recovery_granularity=RecoveryGranularity.PROMPT_GROUP, + admitted=True, + ) + _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) + gate_id = group.gate_rollout_id(0) + partial = { + 0: SiblingSealResult( + gate_rollout_id=gate_id, + receipt={ + "rollout_id": gate_id, + "manifest": [{"staging_key": "g7/sibling-0/call-0"}], + }, + reward=1.0, + ) + } + + with pytest.raises(ValueError, match="requires every logical sibling"): + _mutate(lambda cut: ledger.mark_group_sealed(cut, "g7", partial)) + + assert [ + sibling.current_attempt.status for sibling in ledger.get_group("g7").siblings + ] == [RolloutAttemptStatus.DISPATCHED, RolloutAttemptStatus.DISPATCHED] + assert ledger.expected_staging_keys() == set() + + @pytest.mark.parametrize("unknown_outcome", [False, True]) def test_checkpoint_rejects_ambiguous_finalization_state( unknown_outcome: bool, From e852891c0227218469722cd0689b81e6a24e90aa Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 30 Aug 2026 00:00:41 -0400 Subject: [PATCH 13/28] fix(rollout): preserve missing receipts across restart Signed-off-by: Anish Mahishi --- nemo_rl/experience/rollout_manager.py | 12 ++- nemo_rl/experience/rollout_recovery.py | 86 +++++++++++++------ .../unit/experience/test_rollout_recovery.py | 75 ++++++++++++++++ 3 files changed, 145 insertions(+), 28 deletions(-) diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 047d9aecf0a..b93db5496d8 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -2000,11 +2000,15 @@ async def _record_streamed_completion( raise ValueError( "token-capture completion must contain environment extras" ) - receipt = env_extras.get("ng_receipt") + if "ng_receipt" not in env_extras: + raise ValueError( + "token-capture completion must contain ng_receipt" + ) + receipt = env_extras["ng_receipt"] gate_rollout_id = env_extras.get("ng_rollout_id") - if not isinstance(receipt, dict): + if receipt is not None and not isinstance(receipt, dict): raise ValueError( - "token-capture completion must contain a receipt mapping" + "token-capture completion ng_receipt must be a mapping or None" ) if not isinstance(gate_rollout_id, str): raise ValueError( @@ -2022,7 +2026,7 @@ async def _record_streamed_completion( f"result={gate_rollout_id!r}, " f"expected={expected_gate_rollout_id!r}" ) - if receipt.get("rollout_id") != gate_rollout_id: + if receipt is not None and receipt.get("rollout_id") != gate_rollout_id: raise ValueError( "receipt rollout identity mismatch: " f"receipt={receipt.get('rollout_id')!r}, " diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index e5f779be704..d60f9094b68 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -34,7 +34,11 @@ from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneMutationCut from nemo_rl.data.interfaces import DatumSpec -ROLLOUT_RECOVERY_SCHEMA_VERSION = 3 +ROLLOUT_RECOVERY_SCHEMA_VERSION = 4 +_SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS = { + 3, + ROLLOUT_RECOVERY_SCHEMA_VERSION, +} ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" RolloutRecoveryState: TypeAlias = dict[str, Any] @@ -249,7 +253,9 @@ class SiblingSealResult: """One terminal sibling result waiting for an atomic prompt-group seal.""" gate_rollout_id: str - receipt: dict[str, Any] + # None is an explicit terminal capture failure. The finalizer turns it + # into a masked placeholder, matching the base token-capture contract. + receipt: Optional[dict[str, Any]] reward: float @@ -260,8 +266,10 @@ def _new_attempt() -> RolloutAttemptRecord: ) -def _receipt_staging_keys(receipt: dict[str, Any]) -> list[str]: - """Validate a sealed Gate receipt and return its ordered staging keys.""" +def _receipt_staging_keys(receipt: Optional[dict[str, Any]]) -> list[str]: + """Validate a terminal Gate receipt and return its ordered staging keys.""" + if receipt is None: + return [] manifest = receipt.get("manifest") if not isinstance(manifest, list): raise ValueError("sealed rollout receipt must contain a manifest list") @@ -555,7 +563,7 @@ def mark_sibling_sealed( *, generation_index: int, gate_rollout_id: str, - receipt: dict[str, Any], + receipt: Optional[dict[str, Any]], reward: float, ) -> None: """Record one streamed sibling receipt as soon as the row arrives.""" @@ -574,7 +582,7 @@ def mark_sibling_sealed( "streamed rollout identity mismatch: " f"result={gate_rollout_id!r}, expected={expected_gate_rollout_id!r}" ) - if receipt.get("rollout_id") != gate_rollout_id: + if receipt is not None and receipt.get("rollout_id") != gate_rollout_id: raise ValueError( "receipt rollout identity mismatch: " f"receipt={receipt.get('rollout_id')!r}, expected={gate_rollout_id!r}" @@ -651,7 +659,10 @@ def mark_group_sealed( f"result={result.gate_rollout_id!r}, " f"expected={expected_gate_rollout_id!r}" ) - if result.receipt.get("rollout_id") != expected_gate_rollout_id: + if ( + result.receipt is not None + and result.receipt.get("rollout_id") != expected_gate_rollout_id + ): raise ValueError( "receipt rollout identity mismatch: " f"receipt={result.receipt.get('rollout_id')!r}, " @@ -703,20 +714,21 @@ def abandon_unsealed(self, cut: DataPlaneMutationCut, group_id: str) -> None: def finalization_inputs( self, group_id: str - ) -> tuple[list[str], list[str], list[dict[str, Any]], list[float]]: + ) -> tuple[ + list[str], list[str], list[Optional[dict[str, Any]]], list[float] + ]: """Return physical IDs, canonical IDs, receipts and rewards in sibling order.""" record = self._require_group(group_id) if record.status != PromptGroupStatus.READY_TO_FINALIZE: raise ValueError( f"group {group_id!r} is not ready to finalize: {record.status.value!r}" ) - receipts: list[dict[str, Any]] = [] + receipts: list[Optional[dict[str, Any]]] = [] rewards: list[float] = [] for sibling in record.siblings: attempt = sibling.current_attempt if ( attempt.status != RolloutAttemptStatus.SEALED - or attempt.receipt is None or attempt.reward is None ): raise ValueError( @@ -996,10 +1008,15 @@ def from_state_dict(cls, state: dict[str, Any]) -> Self: "rollout recovery state must be a dictionary, got " f"{type(state).__name__}" ) - if state.get("schema_version") != ROLLOUT_RECOVERY_SCHEMA_VERSION: + schema_version = state.get("schema_version") + if ( + isinstance(schema_version, bool) + or not isinstance(schema_version, int) + or schema_version not in _SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS + ): raise ValueError( "Unsupported rollout-recovery schema version: " - f"{state.get('schema_version')!r}" + f"{schema_version!r}" ) raw_groups = state.get("groups") if not isinstance(raw_groups, list): @@ -1011,6 +1028,7 @@ def from_state_dict(cls, state: dict[str, Any]) -> Self: record = cls._group_from_state( raw_group, seen_attempt_uuids=seen_attempt_uuids, + schema_version=schema_version, ) if record.group_id in ledger._groups: raise ValueError(f"duplicate recovery group_id={record.group_id!r}") @@ -1096,6 +1114,7 @@ def _group_from_state( raw_group: Any, *, seen_attempt_uuids: set[uuid.UUID], + schema_version: int, ) -> PromptGroupRecoveryRecord: if not isinstance(raw_group, dict): raise ValueError("rollout-recovery group must be a mapping") @@ -1191,14 +1210,28 @@ def _group_from_state( ): raise ValueError("staging_keys must be a list of strings") if attempt_status == RolloutAttemptStatus.SEALED: - if not isinstance(receipt, dict) or not isinstance( - reward, (int, float) - ): - raise ValueError("sealed attempts require receipt and reward") - if receipt.get("rollout_id") != gate_id: - raise ValueError("sealed receipt identity mismatch") - if _receipt_staging_keys(receipt) != staging_keys: - raise ValueError("sealed receipt staging manifest mismatch") + if not isinstance(reward, (int, float)): + raise ValueError("sealed attempts require a reward") + if receipt is None: + if schema_version < 4: + raise ValueError( + "sealed attempts require a receipt before schema v4" + ) + if staging_keys: + raise ValueError( + "sealed missing-receipt attempt cannot own staging keys" + ) + elif isinstance(receipt, dict): + if receipt.get("rollout_id") != gate_id: + raise ValueError("sealed receipt identity mismatch") + if _receipt_staging_keys(receipt) != staging_keys: + raise ValueError( + "sealed receipt staging manifest mismatch" + ) + else: + raise ValueError( + "sealed attempt receipt must be a mapping or None" + ) elif receipt is not None or reward is not None or staging_keys: raise ValueError("only sealed attempts may retain receipt data") attempts.append( @@ -1411,11 +1444,16 @@ def parse_rollout_recovery_state(state: object) -> ParsedRolloutRecoveryState: "rollout recovery sidecar must contain a dictionary, got " f"{type(state).__name__}" ) - if state.get("schema_version") != ROLLOUT_RECOVERY_SCHEMA_VERSION: + schema_version = state.get("schema_version") + if ( + isinstance(schema_version, bool) + or not isinstance(schema_version, int) + or schema_version not in _SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS + ): raise ValueError( "unsupported rollout recovery schema_version=" - f"{state.get('schema_version')!r}; expected " - f"{ROLLOUT_RECOVERY_SCHEMA_VERSION}" + f"{schema_version!r}; supported versions are " + f"{sorted(_SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS)}" ) groups = state.get("groups") if not isinstance(groups, list): @@ -1428,7 +1466,7 @@ def parse_rollout_recovery_state(state: object) -> ParsedRolloutRecoveryState: ) ledger_state: RolloutRecoveryState = { - "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "schema_version": schema_version, "groups": groups, "open_train_step": state.get("open_train_step"), } diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index a5ccd4d9c42..b0151936ab8 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -544,6 +544,81 @@ def test_restart_preserves_sealed_sibling_and_retries_only_interrupted_one() -> assert retry.siblings[1].current_attempt.status is RolloutAttemptStatus.RESERVED +@pytest.mark.parametrize( + "recovery_granularity", + [RecoveryGranularity.SIBLING, RecoveryGranularity.PROMPT_GROUP], +) +def test_missing_receipt_is_a_restart_safe_sealed_placeholder( + recovery_granularity: RecoveryGranularity, +) -> None: + ledger = RolloutRecoveryLedger() + group = _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=6, + agent_name=None, + recovery_granularity=recovery_granularity, + admitted=True, + ) + _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) + gate_ids = group.gate_rollout_ids + receipts = [ + None, + { + "rollout_id": gate_ids[1], + "manifest": [{"staging_key": f"{gate_ids[1]}/call"}], + }, + ] + + if recovery_granularity is RecoveryGranularity.SIBLING: + for generation_index, receipt in enumerate(receipts): + _mutate( + lambda cut, generation_index=generation_index, receipt=receipt: ( + ledger.mark_sibling_sealed( + cut, + "g7", + generation_index=generation_index, + gate_rollout_id=gate_ids[generation_index], + receipt=receipt, + reward=float(generation_index), + ) + ) + ) + else: + _mutate( + lambda cut: ledger.mark_group_sealed( + cut, + "g7", + { + generation_index: SiblingSealResult( + gate_rollout_id=gate_ids[generation_index], + receipt=receipt, + reward=float(generation_index), + ) + for generation_index, receipt in enumerate(receipts) + }, + ) + ) + + state = ledger.state_dict() + restored = RolloutRecoveryLedger.from_state_dict(state) + physical_ids, _, restored_receipts, rewards = restored.finalization_inputs("g7") + + assert physical_ids == gate_ids + assert restored_receipts[0] is None + assert restored_receipts[1] == receipts[1] + assert rewards == [0.0, 1.0] + + state["schema_version"] = 3 + with pytest.raises(ValueError, match="before schema v4"): + RolloutRecoveryLedger.from_state_dict(state) + + def test_prompt_group_restart_retries_every_sibling_when_one_is_unfinished() -> None: ledger = RolloutRecoveryLedger() _reserve( From 4bfc4ae8fe1d8fb1723b714b1716d5ca156eed38 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 30 Aug 2026 15:29:32 -0400 Subject: [PATCH 14/28] test(rollout): forward recovery granularity in sibling hook Signed-off-by: Anish Mahishi --- .../functional/_single_controller_sibling_recovery_hook.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/functional/_single_controller_sibling_recovery_hook.py b/tests/functional/_single_controller_sibling_recovery_hook.py index a57f3cf9b77..7213e2551d7 100644 --- a/tests/functional/_single_controller_sibling_recovery_hook.py +++ b/tests/functional/_single_controller_sibling_recovery_hook.py @@ -30,6 +30,7 @@ from examples import run_grpo_single_controller from nemo_rl.experience.rollout_manager import RolloutCompletionCallback +from nemo_rl.experience.rollout_recovery import RecoveryGranularity class _InstrumentedNemoGymRolloutImpl: @@ -92,11 +93,16 @@ async def run_rollout( rollout_ids: list[str] | None = None, generation_indices: list[int] | None = None, on_completion: RolloutCompletionCallback | None = None, + recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, ) -> Any: if rollout_ids is None or generation_indices is None or on_completion is None: raise RuntimeError( "sibling recovery hook requires the token-capture rollout path" ) + if recovery_granularity is not RecoveryGranularity.SIBLING: + raise RuntimeError( + "sibling recovery hook requires sibling recovery granularity" + ) group = self._find_group(rollout_ids) indices = list(generation_indices) @@ -163,6 +169,7 @@ async def _instrumented_completion( rollout_ids=rollout_ids, generation_indices=indices, on_completion=_instrumented_completion, + recovery_granularity=recovery_granularity, ) self._append_event("capture_complete", **fields) return result From 472ccd498b8f434015438791516abca2863bd154 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 31 Aug 2026 23:05:53 -0400 Subject: [PATCH 15/28] fix(rollout): address sibling recovery review findings Signed-off-by: Anish Mahishi --- docs/guides/single-controller.md | 15 +- nemo_rl/algorithms/single_controller.py | 4 + .../single_controller_utils/config.py | 11 + nemo_rl/experience/rollout_manager.py | 64 +++- nemo_rl/experience/rollout_recovery.py | 326 +----------------- .../test_rollout_generation_failures.py | 20 +- tests/unit/experience/test_rollout_manager.py | 29 ++ .../unit/experience/test_rollout_recovery.py | 171 ++++----- .../test_finalizer_lifecycle.py | 3 - .../single_controller/test_rollout_pump.py | 1 + tests/unit/single_controller/test_setup.py | 28 +- tests/unit/test_effort_shaping.py | 41 ++- 12 files changed, 274 insertions(+), 439 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 391ff704df8..3447a21b85f 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -104,7 +104,20 @@ On resume, Single-Controller validates the TQ snapshot against the trainer check Replay recovery is supported by all built-in samplers: `in_order`, `weight_fifo`, `ready_first`, and `windowed`. Custom samplers must explicitly declare `supports_buffer_checkpoint = True`. Otherwise, setup emits a warning and completed buffered groups are not restored. :::{note} -Completed groups are restored directly from the TQ snapshot. Prompt groups whose generations were still in flight at the checkpoint boundary are recovered by ownership: `rollout_recovery.pt` records them, and on resume they are redispatched and regenerated from the same dataset rows. Only rows already committed to TQ preserve their exact generated tokens; redispatched groups produce new samples from the same prompts. +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: + +- `sibling` preserves each sealed sibling and redispatches only unfinished ones. +- `prompt_group` retries every sibling in the group when any sibling is unfinished. + +`agent_granularity_overrides` and `task_granularity_overrides` can select the +policy per Gym agent or dataset task; agent overrides take precedence. These +non-default policies require `token_capture.enabled: true`. The resolved policy +is persisted in `rollout_recovery.pt`, so recovery does not reinterpret an +existing group using changed configuration. Only sealed TQ rows preserve their +exact generated tokens; redispatched siblings produce new samples from the same +prompt. ::: When a sampler does not support replay recovery, a requested data-plane checkpoint is written in `shadow` mode. The TQ snapshot is retained, but no authoritative replay index is written and its rows are not restored into the training replay buffer. diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 66ec66395b5..071ee59c1f9 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1316,6 +1316,10 @@ async def _finalize_with_actor( # The actor publishes canonical rows before returning metadata. Keep # the remote write, local replay-index update, and lineage hand-off in # one mutation cut so a TQ snapshot sees all of them or none of them. + # This deliberately makes checkpoint acquisition wait for the tail of + # every in-flight finalizer RPC. Releasing the cut across the await + # would let a snapshot preserve canonical rows without the matching + # replay index and lineage transition, which is not recoverable. async with self._data_plane_checkpoint_barrier.mutation() as cut: ledger = self._rollout_recovery_ledger ledger.mark_finalization_started(cut, request.group_id) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 0f7d016c86a..1875e754091 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -1108,6 +1108,17 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: ) token_capture_config = master_config.token_capture + recovery_config = master_config.rollout_recovery + if not token_capture_config.enabled and ( + recovery_config.default_granularity is not RecoveryGranularity.SIBLING + or recovery_config.agent_granularity_overrides + or recovery_config.task_granularity_overrides + ): + raise ValueError( + "non-default rollout_recovery policies require " + "token_capture.enabled=true; without token capture, unfinished Gym " + "siblings have no durable receipts to recover" + ) if token_capture_config.defer_routed_experts_to_policy and not ( token_capture_config.enabled ): diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index b93db5496d8..c08284aa74c 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -72,6 +72,7 @@ _attach_routed_experts_to_message_log_prefix, _dummy_routed_experts_for_tokens, _effort_shaping_metrics, + _EffortShapingMetrics, _find_routed_experts_template, _tensorize_by_key, apply_reward_penalties, @@ -1016,6 +1017,7 @@ async def _stream_rows( nemo_gym_env: Any, pending: list[dict], results: list[Optional[dict]], + shaping_by_rowidx: list[Optional[_EffortShapingMetrics]], total_rows: int, timer_prefix: str, on_completion: Optional[RolloutCompletionCallback] = None, @@ -1026,6 +1028,8 @@ async def _stream_rows( nemo_gym_env: The NeMo-Gym environment actor handle. pending: Rows still awaiting a result; each carries its original ``_rowidx``. results: Full-length result list, mutated in place. + shaping_by_rowidx: Per-row shaping metrics, populated before a + completion can be published to the recovery ledger. total_rows: Size of the original prompt group, used to validate row indices. timer_prefix: Timer namespace forwarded to the environment. @@ -1033,7 +1037,7 @@ async def _stream_rows( The environment's timing metrics, or None if the stream ended without them. """ dispatched = {row["_rowidx"] for row in pending} - pending_by_rowidx = {row["_rowidx"]: row for row in pending} + inputs_by_rowidx = {row["_rowidx"]: row for row in pending} received: set[int] = set() env_timing_metrics: Optional[dict[str, Any]] = None @@ -1056,7 +1060,17 @@ async def _stream_rows( if rowidx in received: raise ValueError(f"NeMo-Gym returned duplicate row index {rowidx}") received.add(rowidx) - pending_by_rowidx[rowidx]["agent_ref"] = resolved_agent_ref + inputs_by_rowidx[rowidx]["agent_ref"] = resolved_agent_ref + # A streamed completion may become durable recovery ownership before + # the rest of its prompt group finishes. Shape its reward first so a + # checkpoint never preserves a raw reward that finalization will later + # train on. The shaping rule is row-local; aggregation below is metrics + # only. + shaping_by_rowidx[rowidx] = _apply_effort_shaping( + [result], + [inputs_by_rowidx[rowidx]], + self._effort_config, + ) results[rowidx] = result if on_completion is not None: await on_completion(rowidx, self._result_to_completion(result)) @@ -1101,6 +1115,9 @@ async def _run_rollouts( # Run generation and restore input order as results stream back. with timer.time(f"{timer_prefix}/run_rollouts"): results: list[dict | None] = [None for _ in range(total_rows)] + shaping_by_rowidx: list[Optional[_EffortShapingMetrics]] = [ + None for _ in range(total_rows) + ] env_timing_metrics: dict[str, Any] = {} # One deadline for the whole prompt group, re-dispatches included -- it is # the group that has a budget, not each attempt. It also spans the stream @@ -1142,6 +1159,7 @@ async def _run_rollouts( nemo_gym_env, pending, results, + shaping_by_rowidx, total_rows, timer_prefix, on_completion=on_completion, @@ -1173,9 +1191,30 @@ async def _run_rollouts( raise failure from last_error completed_results = [result for result in results if result is not None] - # Shape rewards for low-effort prompts before completions are built. - shaping = _apply_effort_shaping( - completed_results, inputs, self._effort_config + completed_shaping = [ + metrics for metrics in shaping_by_rowidx if metrics is not None + ] + shaping = _EffortShapingMetrics( + length_rewards_low=[ + value + for metrics in completed_shaping + for value in metrics.length_rewards_low + ], + rewards_low=[ + value + for metrics in completed_shaping + for value in metrics.rewards_low + ], + low_lengths=[ + value + for metrics in completed_shaping + for value in metrics.low_lengths + ], + high_lengths=[ + value + for metrics in completed_shaping + for value in metrics.high_lengths + ], ) # All N rollouts share the same input prompt; tensorize one copy. prompt_message_log = completed_results[0]["input_message_log"] @@ -2001,9 +2040,7 @@ async def _record_streamed_completion( "token-capture completion must contain environment extras" ) if "ng_receipt" not in env_extras: - raise ValueError( - "token-capture completion must contain ng_receipt" - ) + raise ValueError("token-capture completion must contain ng_receipt") receipt = env_extras["ng_receipt"] gate_rollout_id = env_extras.get("ng_rollout_id") if receipt is not None and not isinstance(receipt, dict): @@ -2033,10 +2070,7 @@ async def _record_streamed_completion( f"expected={gate_rollout_id!r}" ) - if ( - recovery_group.recovery_granularity - is RecoveryGranularity.PROMPT_GROUP - ): + if recovery_group.recovery_granularity is RecoveryGranularity.PROMPT_GROUP: result = SiblingSealResult( gate_rollout_id=gate_rollout_id, receipt=receipt, @@ -2140,7 +2174,11 @@ async def _record_completion( # terminal rows keep any later read fail-closed. self._tq_buffer.abort(group_id) async with self._recovery_mutation() as cut: - self._recovery_ledger.abandon_unsealed(cut, group_id) + # Intentional staleness aborts discard the ledger owner before + # cancelling this task. Preserve the original cancellation rather + # than replacing it with "unknown group" during cleanup. + if group_id in self._recovery_ledger: + self._recovery_ledger.abandon_unsealed(cut, group_id) # The capture ledger has no per-rollout fail endpoint. Rows from # abandoned attempts are unreferenced and are swept with the # staging partition at run teardown. diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index d60f9094b68..f64afff1c35 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -28,17 +28,12 @@ from enum import StrEnum from typing import TYPE_CHECKING, Any, Optional, Self, TypeAlias -from nemo_rl.data_plane import KVBatchMeta - if TYPE_CHECKING: from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneMutationCut from nemo_rl.data.interfaces import DatumSpec ROLLOUT_RECOVERY_SCHEMA_VERSION = 4 -_SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS = { - 3, - ROLLOUT_RECOVERY_SCHEMA_VERSION, -} +_SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS = {ROLLOUT_RECOVERY_SCHEMA_VERSION} ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" RolloutRecoveryState: TypeAlias = dict[str, Any] @@ -74,16 +69,6 @@ class PromptGroupStatus(StrEnum): READY_TO_FINALIZE = "ready_to_finalize" FINALIZING = "finalizing" FINALIZATION_UNKNOWN = "finalization_unknown" - FINALIZED = "finalized" - CLAIMED_FOR_TRAINING = "claimed_for_training" - APPLIED_UNCHECKPOINTED = "applied_uncheckpointed" - - -class TrainStepStatus(StrEnum): - """State of the one optimizer step the SingleController may have open.""" - - OPEN = "open" - APPLIED_UNCHECKPOINTED = "applied_uncheckpointed" @dataclass(frozen=True) @@ -179,10 +164,6 @@ class PromptGroupRecoveryRecord: siblings: list[RolloutSiblingRecord] phase: PromptGroupPhase status: PromptGroupStatus = PromptGroupStatus.GENERATING - canonical_meta: Optional[KVBatchMeta] = None - group_min_weight_version: Optional[int] = None - group_max_weight_version: Optional[int] = None - claimed_train_step: Optional[int] = None @property def prompt_payload(self) -> DatumSpec: @@ -228,17 +209,6 @@ def sealed_generation_indices(self) -> list[int]: ] -@dataclass -class OpenTrainStepRecord: - """Groups contributing to the current, not-yet-durable optimizer step.""" - - train_step: int - trainer_version: int - expected_group_count: int - group_ids: list[str] = field(default_factory=list) - status: TrainStepStatus = TrainStepStatus.OPEN - - @dataclass(frozen=True) class ParsedRolloutRecoveryState: """Validated controller and ledger state loaded from one checkpoint sidecar.""" @@ -293,11 +263,6 @@ class RolloutRecoveryLedger: def __init__(self) -> None: self._groups: dict[str, PromptGroupRecoveryRecord] = {} - self._open_train_step: Optional[OpenTrainStepRecord] = None - - @property - def open_train_step(self) -> Optional[OpenTrainStepRecord]: - return copy.deepcopy(self._open_train_step) def groups(self) -> list[PromptGroupRecoveryRecord]: return [self._copy_group(group) for group in self._groups.values()] @@ -433,12 +398,7 @@ def _abandon_entire_group(record: PromptGroupRecoveryRecord) -> None: record.status = PromptGroupStatus.GENERATING def assert_checkpoint_safe(self) -> None: - """Reject states whose publication or optimizer outcome is ambiguous.""" - if self._open_train_step is not None: - raise RuntimeError( - "rollout recovery contains an open optimizer step; restoring " - "mid-step training ownership is not supported" - ) + """Reject states whose canonical publication outcome is ambiguous.""" unsafe = [ record.group_id for record in self._groups.values() @@ -446,8 +406,6 @@ def assert_checkpoint_safe(self) -> None: in { PromptGroupStatus.FINALIZING, PromptGroupStatus.FINALIZATION_UNKNOWN, - PromptGroupStatus.CLAIMED_FOR_TRAINING, - PromptGroupStatus.APPLIED_UNCHECKPOINTED, } ] if unsafe: @@ -570,9 +528,7 @@ def mark_sibling_sealed( cut.require_live() record = self._require_group(group_id) if record.recovery_granularity is RecoveryGranularity.PROMPT_GROUP: - raise ValueError( - "prompt-group recovery must seal every sibling atomically" - ) + raise ValueError("prompt-group recovery must seal every sibling atomically") sibling = self._require_sibling(record, generation_index) attempt = sibling.current_attempt expected_gate_rollout_id = record.gate_rollout_id(generation_index) @@ -639,9 +595,7 @@ def mark_group_sealed( f"expected={sorted(expected_indices)}, actual={sorted(results)}" ) - validated: list[ - tuple[RolloutAttemptRecord, SiblingSealResult, list[str]] - ] = [] + validated: list[tuple[RolloutAttemptRecord, SiblingSealResult, list[str]]] = [] for generation_index in range(record.expected_generations): result = results[generation_index] sibling = self._require_sibling(record, generation_index) @@ -668,9 +622,7 @@ def mark_group_sealed( f"receipt={result.receipt.get('rollout_id')!r}, " f"expected={expected_gate_rollout_id!r}" ) - validated.append( - (attempt, result, _receipt_staging_keys(result.receipt)) - ) + validated.append((attempt, result, _receipt_staging_keys(result.receipt))) # Validate the complete cohort before changing any sibling. A checkpoint # therefore observes either no committed siblings or the complete group. @@ -714,9 +666,7 @@ def abandon_unsealed(self, cut: DataPlaneMutationCut, group_id: str) -> None: def finalization_inputs( self, group_id: str - ) -> tuple[ - list[str], list[str], list[Optional[dict[str, Any]]], list[float] - ]: + ) -> tuple[list[str], list[str], list[Optional[dict[str, Any]]], list[float]]: """Return physical IDs, canonical IDs, receipts and rewards in sibling order.""" record = self._require_group(group_id) if record.status != PromptGroupStatus.READY_TO_FINALIZE: @@ -727,10 +677,7 @@ def finalization_inputs( rewards: list[float] = [] for sibling in record.siblings: attempt = sibling.current_attempt - if ( - attempt.status != RolloutAttemptStatus.SEALED - or attempt.reward is None - ): + if attempt.status != RolloutAttemptStatus.SEALED or attempt.reward is None: raise ValueError( "logical rollout " f"{record.logical_rollout_id(sibling.generation_index)!r} " @@ -773,137 +720,10 @@ def mark_finalization_unknown( ) record.status = PromptGroupStatus.FINALIZATION_UNKNOWN - def mark_group_finalized( - self, - group_id: str, - *, - meta: KVBatchMeta, - group_min_weight_version: int, - group_max_weight_version: int, - ) -> None: - """Transfer recovery ownership from staged receipts to canonical TQ rows.""" - record = self._require_group(group_id) - self._require_group_status( - record, - allowed={PromptGroupStatus.FINALIZING}, - transition="finalize", - ) - if list(meta.sample_ids) != record.logical_rollout_ids: - raise ValueError( - "finalized sample IDs do not match stable logical rollout IDs: " - f"{meta.sample_ids!r} != {record.logical_rollout_ids!r}" - ) - record.canonical_meta = copy.deepcopy(meta) - record.group_min_weight_version = int(group_min_weight_version) - record.group_max_weight_version = int(group_max_weight_version) - record.runtime_prompt_payload = None - record.status = PromptGroupStatus.FINALIZED - - def claim_groups_for_training( - self, - group_ids: list[str], - *, - train_step: int, - trainer_version: int, - expected_group_count: int, - ) -> None: - """Move finalized groups into the controller's one open optimizer step.""" - if not group_ids: - raise ValueError("training claim must contain at least one group") - if len(group_ids) != len(set(group_ids)): - raise ValueError("training claim contains duplicate group IDs") - if self._open_train_step is None: - self._open_train_step = OpenTrainStepRecord( - train_step=train_step, - trainer_version=trainer_version, - expected_group_count=expected_group_count, - ) - open_step = self._open_train_step - if ( - open_step.train_step != train_step - or open_step.trainer_version != trainer_version - or open_step.expected_group_count != expected_group_count - or open_step.status != TrainStepStatus.OPEN - ): - raise ValueError( - "training claim does not match the existing open train step" - ) - records = [self._require_group(group_id) for group_id in group_ids] - for record in records: - if record.status != PromptGroupStatus.FINALIZED: - raise ValueError( - f"cannot claim group {record.group_id!r} from " - f"{record.status.value!r}" - ) - if len(open_step.group_ids) + len(group_ids) > expected_group_count: - raise ValueError("training claim exceeds the step's expected group count") - for record in records: - record.status = PromptGroupStatus.CLAIMED_FOR_TRAINING - record.claimed_train_step = train_step - open_step.group_ids.append(record.group_id) - - def mark_train_step_applied(self, train_step: int) -> None: - """Record optimizer success while rows are still not checkpoint-covered.""" - open_step = self._require_open_train_step(train_step) - if open_step.status != TrainStepStatus.OPEN: - raise ValueError( - f"train step {train_step} is already {open_step.status.value!r}" - ) - if len(open_step.group_ids) != open_step.expected_group_count: - raise ValueError( - f"train step {train_step} has {len(open_step.group_ids)} claimed " - f"groups; expected {open_step.expected_group_count}" - ) - for group_id in open_step.group_ids: - record = self._require_group(group_id) - if record.status != PromptGroupStatus.CLAIMED_FOR_TRAINING: - raise ValueError( - f"train step {train_step} owns group {group_id!r} in state " - f"{record.status.value!r}" - ) - for group_id in open_step.group_ids: - self._groups[group_id].status = PromptGroupStatus.APPLIED_UNCHECKPOINTED - open_step.status = TrainStepStatus.APPLIED_UNCHECKPOINTED - - def release_applied_train_step(self, train_step: int) -> None: - """Drop group metadata after the caller has cleared all owned TQ rows. - - The current controller clears immediately after optimizer success. A later - persistence change will delay this call until a trainer checkpoint covers - the applied update. - """ - open_step = self._require_open_train_step(train_step) - if open_step.status != TrainStepStatus.APPLIED_UNCHECKPOINTED: - raise ValueError( - f"cannot release train step {train_step} from " - f"{open_step.status.value!r}" - ) - for group_id in open_step.group_ids: - del self._groups[group_id] - self._open_train_step = None - - def rollback_open_train_step(self, train_step: int) -> None: - """Return an uncheckpointed step's groups to finalized ownership.""" - open_step = self._require_open_train_step(train_step) - for group_id in open_step.group_ids: - record = self._require_group(group_id) - if record.status not in { - PromptGroupStatus.CLAIMED_FOR_TRAINING, - PromptGroupStatus.APPLIED_UNCHECKPOINTED, - }: - raise ValueError( - f"cannot roll back group {group_id!r} from {record.status.value!r}" - ) - record.status = PromptGroupStatus.FINALIZED - record.claimed_train_step = None - self._open_train_step = None - def discard_group(self, cut: DataPlaneMutationCut, group_id: str) -> None: """Drop a group only after its external TQ/Gate ownership is cleaned.""" cut.require_live() - record = self._require_group(group_id) - if record.claimed_train_step is not None: - raise ValueError(f"cannot discard training-owned group {group_id!r}") + self._require_group(group_id) del self._groups[group_id] def discard_canonical_groups( @@ -961,10 +781,6 @@ def state_dict(self) -> dict[str, Any]: "start_weight_version": record.start_weight_version, "status": record.status.value, "phase": record.phase.value, - "canonical_meta": copy.deepcopy(record.canonical_meta), - "group_min_weight_version": record.group_min_weight_version, - "group_max_weight_version": record.group_max_weight_version, - "claimed_train_step": record.claimed_train_step, "siblings": [ { "generation_index": sibling.generation_index, @@ -983,21 +799,9 @@ def state_dict(self) -> dict[str, Any]: ], } ) - open_step = self._open_train_step return { "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, "groups": groups, - "open_train_step": ( - None - if open_step is None - else { - "train_step": open_step.train_step, - "trainer_version": open_step.trainer_version, - "expected_group_count": open_step.expected_group_count, - "group_ids": list(open_step.group_ids), - "status": open_step.status.value, - } - ), } @classmethod @@ -1015,8 +819,7 @@ def from_state_dict(cls, state: dict[str, Any]) -> Self: or schema_version not in _SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS ): raise ValueError( - "Unsupported rollout-recovery schema version: " - f"{schema_version!r}" + f"Unsupported rollout-recovery schema version: {schema_version!r}" ) raw_groups = state.get("groups") if not isinstance(raw_groups, list): @@ -1028,60 +831,11 @@ def from_state_dict(cls, state: dict[str, Any]) -> Self: record = cls._group_from_state( raw_group, seen_attempt_uuids=seen_attempt_uuids, - schema_version=schema_version, ) if record.group_id in ledger._groups: raise ValueError(f"duplicate recovery group_id={record.group_id!r}") ledger._groups[record.group_id] = record - raw_open_step = state.get("open_train_step") - if raw_open_step is not None: - if not isinstance(raw_open_step, dict): - raise ValueError("open_train_step must be a mapping or None") - try: - open_step = OpenTrainStepRecord( - train_step=int(raw_open_step["train_step"]), - trainer_version=int(raw_open_step["trainer_version"]), - expected_group_count=int(raw_open_step["expected_group_count"]), - group_ids=list(raw_open_step["group_ids"]), - status=TrainStepStatus(raw_open_step["status"]), - ) - except (KeyError, TypeError, ValueError) as error: - raise ValueError("invalid open_train_step state") from error - if not all(isinstance(group_id, str) for group_id in open_step.group_ids): - raise ValueError("open_train_step group_ids must be strings") - if len(open_step.group_ids) != len(set(open_step.group_ids)): - raise ValueError("open_train_step contains duplicate group IDs") - if len(open_step.group_ids) > open_step.expected_group_count: - raise ValueError("open_train_step exceeds expected_group_count") - expected_group_status = ( - PromptGroupStatus.CLAIMED_FOR_TRAINING - if open_step.status == TrainStepStatus.OPEN - else PromptGroupStatus.APPLIED_UNCHECKPOINTED - ) - for group_id in open_step.group_ids: - record = ledger._require_group(group_id) - if ( - record.claimed_train_step != open_step.train_step - or record.status != expected_group_status - ): - raise ValueError( - f"open_train_step ownership mismatch for group {group_id!r}" - ) - claimed_group_ids = { - record.group_id - for record in ledger._groups.values() - if record.claimed_train_step is not None - } - if claimed_group_ids != set(open_step.group_ids): - raise ValueError( - "open_train_step does not list every training-owned group" - ) - ledger._open_train_step = open_step - elif any( - record.claimed_train_step is not None for record in ledger._groups.values() - ): - raise ValueError("claimed groups require an open_train_step record") admission_states: dict[str, tuple[PromptGroupPhase, Optional[int]]] = {} for record in ledger._groups.values(): signature = (record.phase, record.target_step) @@ -1100,13 +854,12 @@ def load_state_dict( ) -> None: """Replace this empty ledger from a validated checkpoint envelope.""" cut.require_live() - if self._groups or self._open_train_step is not None: + if self._groups: raise RuntimeError( "cannot restore into a non-empty rollout recovery ledger" ) restored = self.from_state_dict(state) self._groups = restored._groups - self._open_train_step = restored._open_train_step @classmethod def _group_from_state( @@ -1114,7 +867,6 @@ def _group_from_state( raw_group: Any, *, seen_attempt_uuids: set[uuid.UUID], - schema_version: int, ) -> PromptGroupRecoveryRecord: if not isinstance(raw_group, dict): raise ValueError("rollout-recovery group must be a mapping") @@ -1213,10 +965,6 @@ def _group_from_state( if not isinstance(reward, (int, float)): raise ValueError("sealed attempts require a reward") if receipt is None: - if schema_version < 4: - raise ValueError( - "sealed attempts require a receipt before schema v4" - ) if staging_keys: raise ValueError( "sealed missing-receipt attempt cannot own staging keys" @@ -1225,9 +973,7 @@ def _group_from_state( if receipt.get("rollout_id") != gate_id: raise ValueError("sealed receipt identity mismatch") if _receipt_staging_keys(receipt) != staging_keys: - raise ValueError( - "sealed receipt staging manifest mismatch" - ) + raise ValueError("sealed receipt staging manifest mismatch") else: raise ValueError( "sealed attempt receipt must be a mapping or None" @@ -1261,30 +1007,12 @@ def _group_from_state( raise ValueError("prompt_ref task_name must be a string or None") if sample_id != prompt_id: raise ValueError("prompt_ref sample_id must match prompt_id") - canonical_meta = raw_group.get("canonical_meta") - if canonical_meta is not None and not isinstance(canonical_meta, KVBatchMeta): - raise ValueError("canonical_meta must be KVBatchMeta or None") target_step = raw_group.get("target_step") - claimed_train_step = raw_group.get("claimed_train_step") if target_step is not None and not isinstance(target_step, int): raise ValueError("target_step must be an integer or None") - if claimed_train_step is not None and not isinstance(claimed_train_step, int): - raise ValueError("claimed_train_step must be an integer or None") start_weight = raw_group.get("start_weight_version") if not isinstance(start_weight, int): raise ValueError("start_weight_version must be an integer") - min_weight = raw_group.get("group_min_weight_version") - max_weight = raw_group.get("group_max_weight_version") - if min_weight is not None and not isinstance(min_weight, int): - raise ValueError("group_min_weight_version must be an integer or None") - if max_weight is not None and not isinstance(max_weight, int): - raise ValueError("group_max_weight_version must be an integer or None") - - finalized_states = { - PromptGroupStatus.FINALIZED, - PromptGroupStatus.CLAIMED_FOR_TRAINING, - PromptGroupStatus.APPLIED_UNCHECKPOINTED, - } prefinalization_sealed_states = { PromptGroupStatus.READY_TO_FINALIZE, PromptGroupStatus.FINALIZING, @@ -1300,23 +1028,6 @@ def _group_from_state( raise ValueError( f"group state {status.value!r} requires every sibling to be sealed" ) - if status in finalized_states: - if canonical_meta is None or min_weight is None or max_weight is None: - raise ValueError("finalized group state requires canonical metadata") - if list(canonical_meta.sample_ids) != [ - f"{group_id}_g{s.generation_index}" for s in siblings - ]: - raise ValueError("canonical sample IDs do not match logical lineage") - elif canonical_meta is not None: - raise ValueError("unfinished group cannot contain canonical metadata") - claimed_states = { - PromptGroupStatus.CLAIMED_FOR_TRAINING, - PromptGroupStatus.APPLIED_UNCHECKPOINTED, - } - if (status in claimed_states) != (claimed_train_step is not None): - raise ValueError( - "claimed_train_step must be present exactly for training-owned groups" - ) return PromptGroupRecoveryRecord( group_id=group_id, @@ -1332,10 +1043,6 @@ def _group_from_state( siblings=siblings, phase=phase, status=status, - canonical_meta=copy.deepcopy(canonical_meta), - group_min_weight_version=min_weight, - group_max_weight_version=max_weight, - claimed_train_step=claimed_train_step, ) def _require_group(self, group_id: str) -> PromptGroupRecoveryRecord: @@ -1361,10 +1068,6 @@ def _copy_group(record: PromptGroupRecoveryRecord) -> PromptGroupRecoveryRecord: siblings=copy.deepcopy(record.siblings), phase=record.phase, status=record.status, - canonical_meta=copy.deepcopy(record.canonical_meta), - group_min_weight_version=record.group_min_weight_version, - group_max_weight_version=record.group_max_weight_version, - claimed_train_step=record.claimed_train_step, ) @staticmethod @@ -1391,12 +1094,6 @@ def _require_group_status( f"{record.status.value!r}" ) - def _require_open_train_step(self, train_step: int) -> OpenTrainStepRecord: - open_step = self._open_train_step - if open_step is None or open_step.train_step != train_step: - raise ValueError(f"train step {train_step} is not open") - return open_step - def _validate_batch_shortfall(value: object) -> dict[int, int]: """Return a defensive copy of per-step permanent rollout losses.""" @@ -1468,7 +1165,6 @@ def parse_rollout_recovery_state(state: object) -> ParsedRolloutRecoveryState: ledger_state: RolloutRecoveryState = { "schema_version": schema_version, "groups": groups, - "open_train_step": state.get("open_train_step"), } return ParsedRolloutRecoveryState( ledger_state=ledger_state, diff --git a/tests/unit/experience/test_rollout_generation_failures.py b/tests/unit/experience/test_rollout_generation_failures.py index 54136ca8708..56dd6ff4799 100644 --- a/tests/unit/experience/test_rollout_generation_failures.py +++ b/tests/unit/experience/test_rollout_generation_failures.py @@ -753,13 +753,31 @@ def test_rows_must_carry_their_own_index(self): method = _PartialGymMethod(fail_after_rows=99, failures_before_success=0) impl = _make_gym_impl(method, num_generations=2, row_attempts=2) - with pytest.raises(ValueError, match="must be stamped with their own position"): + with pytest.raises(ValueError, match="carries invalid _rowidx"): asyncio.run( impl._run_rollouts( [{"agent_ref": {"name": "a"}}], Timer(), "timing/rollout" ) ) + def test_row_indices_must_fit_within_the_prompt_group(self): + method = _PartialGymMethod(fail_after_rows=99, failures_before_success=0) + impl = _make_gym_impl(method, num_generations=2, row_attempts=2) + rows = _gym_rows(2) + rows[1]["_rowidx"] = 2 + + with pytest.raises(ValueError, match="carries invalid _rowidx=2"): + asyncio.run(impl._run_rollouts(rows, Timer(), "timing/rollout")) + + def test_row_indices_must_be_unique(self): + method = _PartialGymMethod(fail_after_rows=99, failures_before_success=0) + impl = _make_gym_impl(method, num_generations=2, row_attempts=2) + rows = _gym_rows(2) + rows[1]["_rowidx"] = 0 + + with pytest.raises(ValueError, match="duplicate _rowidx values"): + asyncio.run(impl._run_rollouts(rows, Timer(), "timing/rollout")) + def test_the_group_deadline_spans_re_dispatches(self): """The budget belongs to the prompt group, not to each attempt. diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index bd216562333..0b79937764d 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -1739,6 +1739,35 @@ async def _boom(_sample): _run(mgr.generate_for_finalization({"prompt": "p", "idx": 0})) assert len(buf.abort_calls) == 1 + def test_cancel_after_controller_discard_preserves_cancelled_error(self): + """A stale abort may delete lineage before rollout cleanup runs.""" + + async def _scenario() -> None: + started = asyncio.Event() + + async def _block(_sample: object) -> None: + started.set() + await asyncio.Event().wait() + + buf = _FakeCaptureBuffer() + mgr = _make_capture_manager(buf, on_run=_block) + task = asyncio.create_task( + mgr.generate_for_finalization({"prompt": "p", "idx": 0}) + ) + await asyncio.wait_for(started.wait(), timeout=1.0) + (group_id,) = buf._slots + async with buf.data_plane_checkpoint_barrier.mutation() as cut: + mgr.discard_prompt_group(cut, group_id) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert group_id not in mgr.recovery_ledger + assert buf.abort_calls == [group_id] + + asyncio.run(_scenario()) + def test_retries_infrastructure_failure_with_stable_logical_ids(self): buf = _FakeCaptureBuffer() attempts = 0 diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index b0151936ab8..c61be7af2f7 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -28,7 +28,6 @@ ) from nemo_rl.algorithms.single_controller_utils.config import RolloutRecoveryConfig from nemo_rl.data.interfaces import DatumSpec -from nemo_rl.data_plane import KVBatchMeta from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, PromptGroupPhase, @@ -93,29 +92,6 @@ def _shuffled_prompt_loader(seed: int = 123) -> StatefulDataLoader: ) -def _group_state( - idx: int = 7, - *, - target_step: int | None = 7, - phase: str = "admitted", -) -> dict: - return { - "group_id": f"g{idx}", - "admission_id": "batch-7", - "prompt_id": str(idx), - "prompt_ref": { - "sample_id": str(idx), - "task_name": None, - }, - "agent_name": None, - "recovery_granularity": "sibling", - "expected_generations": 2, - "target_step": target_step, - "start_weight_version": 7, - "phase": phase, - } - - def test_ledger_round_trip_preserves_group_ownership() -> None: ledger = RolloutRecoveryLedger() _reserve( @@ -133,6 +109,13 @@ def test_ledger_round_trip_preserves_group_ownership() -> None: ) state = ledger.state_dict() + assert "open_train_step" not in state + assert { + "canonical_meta", + "group_min_weight_version", + "group_max_weight_version", + "claimed_train_step", + }.isdisjoint(state["groups"][0]) restored = RolloutRecoveryLedger() _load(restored, state) @@ -615,7 +598,7 @@ def test_missing_receipt_is_a_restart_safe_sealed_placeholder( assert rewards == [0.0, 1.0] state["schema_version"] = 3 - with pytest.raises(ValueError, match="before schema v4"): + with pytest.raises(ValueError, match="Unsupported rollout-recovery schema version"): RolloutRecoveryLedger.from_state_dict(state) @@ -680,9 +663,7 @@ def test_prompt_group_restart_keeps_a_fully_sealed_group() -> None: gate_rollout_id=gate_id, receipt={ "rollout_id": gate_id, - "manifest": [ - {"staging_key": f"g7/sibling-{generation_index}/call-0"} - ], + "manifest": [{"staging_key": f"g7/sibling-{generation_index}/call-0"}], }, reward=1.0, ) @@ -776,58 +757,6 @@ def test_checkpoint_rejects_ambiguous_finalization_state( ledger.state_dict() -def test_checkpoint_rejects_an_open_optimizer_step() -> None: - ledger = RolloutRecoveryLedger() - group = _reserve( - ledger, - group_id="g7", - admission_id="batch-7", - prompt_id="7", - prompt_payload=_prompt(), - expected_generations=1, - target_step=7, - start_weight_version=6, - agent_name=None, - recovery_granularity=RecoveryGranularity.SIBLING, - admitted=True, - ) - _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) - gate_id = group.gate_rollout_id(0) - _mutate( - lambda cut: ledger.mark_sibling_sealed( - cut, - "g7", - generation_index=0, - gate_rollout_id=gate_id, - receipt={ - "rollout_id": gate_id, - "manifest": [{"staging_key": "g7/sibling-0/call-0"}], - }, - reward=1.0, - ) - ) - _mutate(lambda cut: ledger.mark_finalization_started(cut, "g7")) - ledger.mark_group_finalized( - "g7", - meta=KVBatchMeta( - partition_id="rollout_data", - task_name="train", - sample_ids=["g7_g0"], - ), - group_min_weight_version=6, - group_max_weight_version=6, - ) - ledger.claim_groups_for_training( - ["g7"], - train_step=7, - trainer_version=7, - expected_group_count=1, - ) - - with pytest.raises(RuntimeError, match="open optimizer step"): - ledger.state_dict() - - @pytest.mark.parametrize( ("field", "value", "error_fragment"), [ @@ -864,27 +793,63 @@ def test_restore_rejects_malformed_recovery_policy_fields( _load(RolloutRecoveryLedger(), state) -@pytest.mark.parametrize( - "state", - [ - {"schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION + 1, "groups": []}, - {"schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, "groups": {}}, - { - "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, - "groups": [_group_state(phase="unknown")], - }, - { - "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, - "groups": [ - _group_state(idx, target_step=target_step, phase=phase) - for idx, target_step, phase in ( - (7, None, "reserved"), - (8, 7, "admitted"), - ) - ], - }, - ], -) -def test_restore_rejects_incompatible_or_malformed_state(state: dict) -> None: - with pytest.raises((TypeError, ValueError)): +def test_restore_rejects_unsupported_schema_version() -> None: + state = { + "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION + 1, + "groups": [], + } + + with pytest.raises(ValueError, match="Unsupported rollout-recovery schema"): _load(RolloutRecoveryLedger(), state) # type: ignore[arg-type] + + +def test_restore_rejects_non_list_groups() -> None: + state = { + "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "groups": {}, + } + + with pytest.raises(ValueError, match="must contain a groups list"): + _load(RolloutRecoveryLedger(), state) # type: ignore[arg-type] + + +def test_restore_rejects_invalid_prompt_group_phase() -> None: + ledger = RolloutRecoveryLedger() + _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=1, + target_step=7, + start_weight_version=6, + admitted=True, + ) + state = ledger.state_dict() + state["groups"][0]["phase"] = "unknown" + + with pytest.raises(ValueError, match="invalid prompt group phase"): + _load(RolloutRecoveryLedger(), state) + + +def test_restore_rejects_inconsistent_shared_admission_state() -> None: + ledger = RolloutRecoveryLedger() + for idx in (7, 8): + _reserve( + ledger, + group_id=f"g{idx}", + admission_id="batch-7", + prompt_id=str(idx), + prompt_payload=_prompt(idx), + expected_generations=1, + target_step=7, + start_weight_version=6, + admitted=True, + ) + state = ledger.state_dict() + state["groups"][0]["target_step"] = None + state["groups"][0]["phase"] = "reserved" + + with pytest.raises(ValueError, match="disagree on phase or target_step"): + _load(RolloutRecoveryLedger(), state) diff --git a/tests/unit/single_controller/test_finalizer_lifecycle.py b/tests/unit/single_controller/test_finalizer_lifecycle.py index 4da6f92fec9..b508a2db1cf 100644 --- a/tests/unit/single_controller/test_finalizer_lifecycle.py +++ b/tests/unit/single_controller/test_finalizer_lifecycle.py @@ -90,7 +90,6 @@ def _request() -> ReassemblyRequest: rewards=(1.0,), mask_sample=(False,), fallback_weight_version=3, - prompt_idx=0, ) @@ -272,8 +271,6 @@ def test_post_train_cleanup_clears_canonical_rows_and_route_plan_staging_keys() "partition_id": "staging", }, ] - - class _SyncDataPlaneClient: """Synchronous client like the production TQ adapter; records caller threads.""" diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 2353bc44cff..63ea1f0a22e 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -1341,6 +1341,7 @@ async def _finalize( await ctrl._rollout_pump() assert manager.recovery_ledger.groups() == [] + assert manager.stats.committed == int(committed) # A committed group transfers its permit to the train pump; a dropped # group returns it immediately because no canonical replay row owns it. assert ctrl._buffer_capacity._value == (0 if committed else 1) diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index a50d4143e79..86df2696ec2 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -62,6 +62,7 @@ from nemo_rl.data.multimodal_utils import WIRE_MULTIMODAL_FIELDS from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION from nemo_rl.data_plane.schema import SC_ROLLOUT_SCHEMA_FIELDS +from nemo_rl.experience.rollout_recovery import RecoveryGranularity from nemo_rl.experience.rollouts import EffortLevelsConfig from nemo_rl.models.generation.megatron.megatron_generation import MegatronGeneration from nemo_rl.utils.config import ( @@ -489,7 +490,9 @@ def test_build_trainer_initializes_reference_model_only_for_nonzero_kl( ) -def test_sibling_recovery_functional_config_resolves_to_runtime_contract(): +def test_sibling_recovery_functional_config_resolves_to_runtime_contract( + tmp_path: Path, +) -> None: """The two-phase Gym recovery fixture must pass SC config validation.""" register_omegaconf_resolvers() repo_root = Path(__file__).resolve().parents[3] @@ -524,7 +527,7 @@ def test_sibling_recovery_functional_config_resolves_to_runtime_contract(): "grpo.skip_reference_policy_logprobs_calculation=false", "loss_fn.use_importance_sampling_correction=true", "checkpointing.enabled=true", - "checkpointing.checkpoint_dir=/tmp/sibling-recovery-checkpoints", + f"checkpointing.checkpoint_dir={tmp_path / 'sibling-recovery-checkpoints'}", "checkpointing.metric_name=null", "checkpointing.save_period=1", "+checkpointing.save_data_plane=true", @@ -557,6 +560,11 @@ def test_sibling_recovery_functional_config_resolves_to_runtime_contract(): validate_single_controller_config(master_config) assert master_config.checkpointing["metric_name"] is None assert master_config.checkpointing["save_data_plane"] is True + assert master_config.token_capture.enabled is True + assert ( + master_config.rollout_recovery.default_granularity + is RecoveryGranularity.SIBLING + ) assert master_config.async_rl.rollout_failure.native.generation_timeout_s is None assert master_config.async_rl.rollout_failure.nemo_gym.rollout_timeout_s == 120 @@ -881,6 +889,16 @@ def create_teachers(*args, **kwargs): ValueError, "defer_routed_experts_to_policy requires", ), + ( + "prompt_group_recovery_without_capture", + ValueError, + "non-default rollout_recovery policies require", + ), + ( + "recovery_override_without_capture", + ValueError, + "non-default rollout_recovery policies require", + ), ], ) def test_invalid_config_fails_before_setup_factories( @@ -920,6 +938,12 @@ def test_invalid_config_fails_before_setup_factories( mc.async_rl.generation_fleet_health.enabled = True elif invalid_case == "gym_on_sglang": mc = _make_master_config(colocated=False, backend="sglang") + elif invalid_case == "prompt_group_recovery_without_capture": + mc.rollout_recovery.default_granularity = RecoveryGranularity.PROMPT_GROUP + elif invalid_case == "recovery_override_without_capture": + mc.rollout_recovery.task_granularity_overrides = { + "genrm": RecoveryGranularity.PROMPT_GROUP + } else: # pragma: no cover raise AssertionError(f"unknown test case {invalid_case}") diff --git a/tests/unit/test_effort_shaping.py b/tests/unit/test_effort_shaping.py index e3f2e9f8cc2..306b591e979 100644 --- a/tests/unit/test_effort_shaping.py +++ b/tests/unit/test_effort_shaping.py @@ -13,11 +13,13 @@ # limitations under the License. import asyncio +from collections.abc import Awaitable, Callable from typing import Optional import pytest from nemo_rl.algorithms.single_controller_utils.config import RolloutRecoveryConfig +from nemo_rl.experience.interfaces import Completion from nemo_rl.experience.rollout_manager import ( AsyncNemoGymRolloutImpl, RolloutManager, @@ -305,6 +307,8 @@ def _run_gym_rollouts( effort_config: Optional[EffortLevelsConfig], prompt: str, results: list[dict], + *, + on_completion: Optional[Callable[[int, Completion], Awaitable[None]]] = None, ): """Drive the real _run_rollouts against a fake NeMo-Gym stream.""" impl = AsyncNemoGymRolloutImpl( @@ -322,7 +326,14 @@ def _run_gym_rollouts( )() } inputs = [_gym_input(i, prompt) for i in range(len(results))] - return asyncio.run(impl._run_rollouts(inputs, Timer(), "timing/test")) + return asyncio.run( + impl._run_rollouts( + inputs, + Timer(), + "timing/test", + on_completion=on_completion, + ) + ) def test_rollout_manager_forwards_effort_config(): @@ -362,6 +373,34 @@ def test_run_rollouts_shapes_completion_reward_and_emits_low_metrics(): assert "median_length_high" not in metrics +def test_streamed_completion_is_shaped_before_recovery_callback() -> None: + """Recovery ownership must never seal the environment's raw reward.""" + observed_rewards: list[float] = [] + + async def record_completion(_rowidx: int, completion: Completion) -> None: + observed_rewards.append(completion.reward) + + result = { + "input_message_log": [{"role": "user", "token_ids": [1, 2]}], + "message_log": [], + "full_result": { + "reward": 1.0, + "response": {"usage": {"output_tokens": 100}}, + }, + "receipt": {"rollout_id": "rollout-0", "manifest": []}, + "rollout_id": "rollout-0", + } + + _run_gym_rollouts( + _LOW_EFFORT_CONFIG, + " be concise", + [result], + on_completion=record_completion, + ) + + assert observed_rewards == pytest.approx([1.9]) + + def test_run_rollouts_leaves_high_effort_prompt_reward_untouched(): """A prompt without low_string is only counted, never re-scored.""" completions, _, metrics = _run_gym_rollouts( From ff3626749076fcdd31f84d9db036dc6dc645d656 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 3 Sep 2026 12:00:44 -0400 Subject: [PATCH 16/28] fix(rollout): address token capture review findings Signed-off-by: Anish Mahishi --- docs/guides/single-controller.md | 11 ++- ...po_math_1B_megatron_single_controller.yaml | 10 ++ ...po_math_1B_megatron_single_controller.yaml | 10 ++ .../algorithms/async_utils/replay_buffer.py | 96 ++++++------------- nemo_rl/algorithms/single_controller.py | 9 +- .../single_controller_utils/config.py | 9 ++ nemo_rl/data_plane/tq_token_sink.py | 8 +- nemo_rl/experience/rollout_manager.py | 28 +++++- nemo_rl/experience/rollout_recovery.py | 86 +++++++++++++++-- tests/unit/experience/test_rollout_manager.py | 49 +++++++++- .../unit/experience/test_rollout_recovery.py | 59 +++++++++++- tests/unit/experience/test_rollouts.py | 1 + .../single_controller/test_checkpointing.py | 36 ++++--- .../test_finalizer_lifecycle.py | 3 +- tests/unit/single_controller/test_setup.py | 6 +- .../test_tq_replay_buffer.py | 93 ++++++++++++++---- 16 files changed, 388 insertions(+), 126 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 3447a21b85f..751dd9385b5 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -111,13 +111,18 @@ failure and restart behavior: - `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. + `agent_granularity_overrides` and `task_granularity_overrides` can select the policy per Gym agent or dataset task; agent overrides take precedence. These non-default policies require `token_capture.enabled: true`. The resolved policy is persisted in `rollout_recovery.pt`, so recovery does not reinterpret an -existing group using changed configuration. Only sealed TQ rows preserve their -exact generated tokens; redispatched siblings produce new samples from the same -prompt. +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 sample from the same prompt. +Neither becomes a training row until every generation in the group has finished. ::: When a sampler does not support replay recovery, a requested data-plane checkpoint is written in `shadow` mode. The TQ snapshot is retained, but no authoritative replay index is written and its rows are not restored into the training replay buffer. diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index 8b0d71d6216..9444628eac1 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -186,9 +186,19 @@ token_capture: # over task overrides. rollout_recovery: default_granularity: sibling + # Per-Gym-agent override, keyed on extra_env_info.agent_ref.name. Wins over + # the task map below. + # agent_granularity_overrides: {genrm_agent: prompt_group} agent_granularity_overrides: {} + # Per-dataset-task override, keyed on the prompt's task_name. + # task_granularity_overrides: {math: prompt_group} task_granularity_overrides: {} +# Leaving either map above non-empty requires token capture, which is a separate +# top-level section (not in this file, and off by default): +# token_capture: +# enabled: true + cluster: # Master ports inherit the shared 1400-1999 band from grpo_math_1B.yaml. gpus_per_node: 2 diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 7f548486e87..7fdd5a6f8f8 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -224,9 +224,19 @@ token_capture: # over task overrides. rollout_recovery: default_granularity: sibling + # Per-Gym-agent override, keyed on extra_env_info.agent_ref.name. Wins over + # the task map below. + # agent_granularity_overrides: {genrm_agent: prompt_group} agent_granularity_overrides: {} + # Per-dataset-task override, keyed on the prompt's task_name. + # task_granularity_overrides: {math: prompt_group} task_granularity_overrides: {} +# Leaving either map above non-empty requires token capture, which is a separate +# top-level section (not in this file, and off by default): +# token_capture: +# enabled: true + cluster: # Master ports inherit the shared 1400-1999 band from ppo_math_1B.yaml. gpus_per_node: 2 diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index fc5cacfc0d9..9427fc06afb 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -230,40 +230,18 @@ def __init__(self) -> None: self._condition = asyncio.Condition() self._checkpoint_active = False self._active_mutations = 0 - self._mutation_depth_by_task: dict[asyncio.Task[Any], int] = {} - self._mutation_cut_by_task: dict[asyncio.Task[Any], DataPlaneMutationCut] = {} @asynccontextmanager async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: - """Yield one task-local live cut after any active checkpoint exits.""" - task = asyncio.current_task() - if task is None: - raise RuntimeError("data-plane mutation must run inside an asyncio task") - depth = self._mutation_depth_by_task.get(task, 0) - if depth: - # Replay-buffer helpers may join a controller-owned mutation. Do not - # wait behind a checkpoint that is already waiting for this outer - # mutation, or the two tasks deadlock. - self._mutation_depth_by_task[task] = depth + 1 - cut = self._mutation_cut_by_task[task] - cut.require_live() - try: - yield cut - finally: - self._mutation_depth_by_task[task] -= 1 - return + """Yield a live cut after any active checkpoint exits.""" async with self._condition: await self._condition.wait_for(lambda: not self._checkpoint_active) self._active_mutations += 1 - self._mutation_depth_by_task[task] = 1 cut = DataPlaneMutationCut(self) - self._mutation_cut_by_task[task] = cut try: yield cut finally: cut._invalidate() - del self._mutation_cut_by_task[task] - del self._mutation_depth_by_task[task] async with self._condition: self._active_mutations -= 1 if self._active_mutations == 0: @@ -1251,13 +1229,18 @@ async def remove_group(self, group_id: str, *, remove_in_dp: bool = False) -> in "TQReplayBuffer must be bound to the controller data-plane " "checkpoint barrier before removing a group" ) - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation() as cut: return await self._remove_groups_unlocked( - [group_id], clear_data_plane=remove_in_dp + cut, [group_id], clear_data_plane=remove_in_dp ) - async def clear_staging_keys(self, staging_keys: list[str]) -> None: - """Clear known token-capture staging rows under the checkpoint barrier.""" + async def clear_staging_keys( + self, + cut: DataPlaneMutationCut, + staging_keys: list[str], + ) -> None: + """Clear known token-capture staging rows under a caller-owned cut.""" + cut.require_live() if not staging_keys: return if self._staging_partition_id is None: @@ -1270,17 +1253,17 @@ async def clear_staging_keys(self, staging_keys: list[str]) -> None: "checkpoint barrier before clearing staging samples" ) unique_keys = list(dict.fromkeys(staging_keys)) - async with self._data_plane_checkpoint_barrier.mutation(): - await call_data_plane( - self._dp_client, - "clear_samples", - offload_sync=True, - sample_ids=unique_keys, - partition_id=self._staging_partition_id, - ) + await call_data_plane( + self._dp_client, + "clear_samples", + offload_sync=True, + sample_ids=unique_keys, + partition_id=self._staging_partition_id, + ) async def commit_finalized( self, + cut: DataPlaneMutationCut, group_id: str, meta: KVBatchMeta, group_min_wv: int, @@ -1297,6 +1280,7 @@ async def commit_finalized( rollout straddles a refit. Args: + cut: Live cut acquired by the owner coordinating finalization. group_id: group_id returned by the matching reserve call. meta: KVBatchMeta the finalizer built over its published rows. group_min_wv: Oldest weight version any call in the group used. @@ -1307,30 +1291,7 @@ async def commit_finalized( Raises: ValueError: group_id has no live slot (removed or never reserved). """ - if self._data_plane_checkpoint_barrier is None: - raise RuntimeError( - "TQReplayBuffer must be bound to the controller data-plane " - "checkpoint barrier before committing finalized groups" - ) - async with self._data_plane_checkpoint_barrier.mutation(): - return self._commit_finalized_unlocked( - group_id, - meta, - group_min_wv, - group_max_wv, - staging_keys=staging_keys, - ) - - def _commit_finalized_unlocked( - self, - group_id: str, - meta: KVBatchMeta, - group_min_wv: int, - group_max_wv: int, - *, - staging_keys: Optional[list[str]] = None, - ) -> KVBatchMeta: - """Fill the slot while the caller holds a barrier mutation slot.""" + cut.require_live() try: idx = self._group_ids.index(group_id) except ValueError: @@ -1439,20 +1400,23 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: # removal may shift every list index while this task waits for the barrier # or for DataPlane cleanup. drop_group_ids = [self._group_ids[i] for i in drop_idxs] - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation() as cut: return await self._remove_groups_unlocked( - drop_group_ids, clear_data_plane=remove_in_dp + cut, drop_group_ids, clear_data_plane=remove_in_dp ) async def _remove_groups_unlocked( - self, group_ids: list[str], *, clear_data_plane: bool + self, + cut: DataPlaneMutationCut, + group_ids: list[str], + *, + clear_data_plane: bool, ) -> int: - """Remove stable groups while the caller owns a mutation slot.""" + """Remove stable groups while the caller owns a live mutation cut.""" + cut.require_live() if len(group_ids) != len(set(group_ids)): raise ValueError("replay removal contains duplicate group IDs") - index_by_group_id = { - group_id: i for i, group_id in enumerate(self._group_ids) - } + index_by_group_id = {group_id: i for i, group_id in enumerate(self._group_ids)} if len(index_by_group_id) != len(self._group_ids): raise RuntimeError("replay buffer contains duplicate live group IDs") missing_group_ids = [ diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 071ee59c1f9..c6d2c51feee 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -747,6 +747,7 @@ async def _maybe_restore_rollout_recovery( recovery_ledger.discard_canonical_groups(cut, canonical_group_ids) if self._master_config.token_capture.enabled: await self._validate_rollout_recovery_inventory( + cut, replay_metadata=canonical_state, clear_unreferenced=True, ) @@ -1033,11 +1034,13 @@ async def _validate_replay_inventory( async def _validate_rollout_recovery_inventory( self, + cut: DataPlaneMutationCut, *, replay_metadata: Optional[TQReplayMetadataState], clear_unreferenced: bool, ) -> None: - """Require every unfinished receipt or deferred route to retain staging.""" + """Validate staging ownership while the caller holds a stable cut.""" + cut.require_live() expected_staging_keys = self._rollout_recovery_ledger.expected_staging_keys() if replay_metadata is not None: for group in replay_metadata["groups"]: @@ -1372,6 +1375,7 @@ async def _finalize_with_actor( else: try: await self._buffer.commit_finalized( + cut, request.group_id, finalized.meta, finalized.group_min_wv, @@ -3322,7 +3326,7 @@ async def _save_checkpoint( # controller artifact under the exclusive side so the checkpoint cannot # contain a cursor without its prompt owner, or two durable owners for one # canonical group. - async with self._data_plane_checkpoint_barrier.checkpoint(): + async with self._data_plane_checkpoint_barrier.checkpoint() as cut: save_state.current_step = self._train_steps save_state.total_steps = self._train_steps save_state.trainer_version = self._trainer_version @@ -3379,6 +3383,7 @@ async def _save_checkpoint( if self._master_config.token_capture.enabled: await self._validate_rollout_recovery_inventory( + cut, replay_metadata=replay_metadata, clear_unreferenced=False, ) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 1875e754091..394f59c34e1 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -634,6 +634,15 @@ class ResolvedRolloutRecovery: class RolloutRecoveryConfig(BaseModel, extra="allow"): """Retry and restore policy for unfinished token-capture prompt groups. + ``sibling`` (the default) preserves completed generations and retries only + the missing ones. Prefer it when reusing work and avoiding repeated long-tail + generations matters more than keeping a group on one policy version. + + ``prompt_group`` discards and regenerates every sibling when any generation + is unfinished. It costs a full group per recovery, but keeps the regenerated + group on the policy weights live at redispatch instead of mixing those results + with older sealed siblings. + The resolved value is persisted on each ledger group, so restoring a saved group does not reinterpret it using a newer configuration. The same granularity governs failures handled in-process and after a process restart. diff --git a/nemo_rl/data_plane/tq_token_sink.py b/nemo_rl/data_plane/tq_token_sink.py index 54ef6b96596..83a7974e6ea 100644 --- a/nemo_rl/data_plane/tq_token_sink.py +++ b/nemo_rl/data_plane/tq_token_sink.py @@ -464,8 +464,12 @@ def _select_row(rows: TensorDict, index: int) -> dict[str, torch.Tensor]: """ row: dict[str, torch.Tensor] = {} for field in rows.keys(): - leaf = cast(torch.Tensor, rows.get(field)) - row[str(field)] = leaf[index].unsqueeze(0) + value = rows.get(field) + if not isinstance(value, torch.Tensor): + raise TypeError( + f"staging field {field!r} must be a tensor, got {type(value).__name__}" + ) + row[str(field)] = value[index].unsqueeze(0) return row diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index c08284aa74c..11ca9d9a595 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -1332,7 +1332,7 @@ def _compute_rollout_metrics( # (cum_len of the deepest chain; delta sums as the generation # proxy) instead of a message_log walk. manifests = [ - ((c.env_extras.get("ng_receipt") or {}).get("manifest") or []) + (((c.env_extras or {}).get("ng_receipt") or {}).get("manifest") or []) for c in completions ] # .get with 0: _assemble_receipt ships raw ledger rows unvalidated @@ -1406,7 +1406,7 @@ def _compute_rollout_metrics( # Agent-level metrics. Receipts are lineage records, not agent # results — keep them (and their manifests) out of the logged table. agent_extras = [ - {k: v for k, v in c.env_extras.items() if k not in ("ng_receipt",)} + {k: v for k, v in (c.env_extras or {}).items() if k not in ("ng_receipt",)} for c in completions ] for key in agent_extras[0].keys(): @@ -1919,10 +1919,16 @@ async def generate_for_finalization( inflight_registry: Optional[dict[str, tuple[asyncio.Task[None], int]]] = None, lineage_group_id: Optional[str] = None, ) -> Optional["ReassemblyRequest"]: - """Capture siblings with stable lineage and configured retry granularity.""" + """Capture siblings with stable lineage and configured retry granularity. + + Returns ``None`` when infrastructure retries are exhausted within the + configured drop budget. The caller then owns the backpressure permit and + target-step shortfall. + """ assert self._tq_buffer is not None, ( "generate_for_finalization requires tq_buffer to be set at __init__" ) + owns_recovery_group = lineage_group_id is None recovery_group_id = lineage_group_id if recovery_group_id is None: async with self._recovery_mutation() as cut: @@ -1978,6 +1984,12 @@ async def generate_for_finalization( assert last_infra_error is not None reason = type(last_infra_error).__name__ self._consecutive_infra_drops += 1 + if owns_recovery_group: + # Without controller-owned recovery lineage, nobody above this method + # knows the temporary group ID. Clean its known staging ownership before + # dropping the only record that names those rows. + async with self._recovery_mutation() as cut: + await self.discard_recovery_group(cut, recovery_group_id) if self._consecutive_infra_drops > policy.max_consecutive_dropped_prompts: raise RolloutRedispatchExhausted( f"prompt idx={input_sample['idx']} exhausted its infrastructure " @@ -1988,6 +2000,14 @@ async def generate_for_finalization( f"{reason}: {last_infra_error}" ) from last_infra_error self._stats.record_infra_drop(reason, self._consecutive_infra_drops) + print( + f"dropping capture prompt idx={input_sample['idx']} after " + f"{infra_attempts} infrastructure failure(s) ({reason}: " + f"{last_infra_error}) [consecutive drop " + f"{self._consecutive_infra_drops}/" + f"{policy.max_consecutive_dropped_prompts}]", + flush=True, + ) self._stats.skipped += 1 return None @@ -2198,5 +2218,5 @@ async def discard_recovery_group( for sibling in group.siblings for key in sibling.current_attempt.staging_keys ] - await self._tq_buffer.clear_staging_keys(staging_keys) + await self._tq_buffer.clear_staging_keys(cut, staging_keys) self._recovery_ledger.discard_group(cut, group_id) diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index f64afff1c35..56b14bf0599 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -15,15 +15,16 @@ """Controller-owned lineage for recoverable token-capture prompt groups. The ledger deliberately contains control-plane metadata only. Token tensors and -router-replay payloads remain in TQ. Persistence is added by a later change; the -versioned ``state_dict`` boundary lives here so that change does not have to -invent a second lifecycle model. +router-replay payloads remain in TQ. The versioned ``state_dict`` boundary here is +what the controller writes into ``rollout_recovery.pt`` at each checkpoint and +reads back on restore. """ from __future__ import annotations import copy import uuid +from collections.abc import Mapping from dataclasses import dataclass, field from enum import StrEnum from typing import TYPE_CHECKING, Any, Optional, Self, TypeAlias @@ -37,6 +38,50 @@ ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" RolloutRecoveryState: TypeAlias = dict[str, Any] +_LEDGER_STATE_FIELDS = frozenset({"schema_version", "groups"}) +_SIDECAR_STATE_FIELDS = frozenset( + { + *_LEDGER_STATE_FIELDS, + "batch_shortfall", + "sampler_stamps_target_steps", + } +) +_GROUP_STATE_FIELDS = frozenset( + { + "group_id", + "admission_id", + "prompt_id", + "prompt_ref", + "agent_name", + "recovery_granularity", + "expected_generations", + "target_step", + "start_weight_version", + "status", + "phase", + "siblings", + } +) +_PROMPT_REF_STATE_FIELDS = frozenset({"sample_id", "task_name"}) +_SIBLING_STATE_FIELDS = frozenset({"generation_index", "attempts"}) +_ATTEMPT_STATE_FIELDS = frozenset( + {"attempt_uuid", "status", "receipt", "reward", "staging_keys"} +) + + +def _reject_unknown_fields( + mapping: Mapping[Any, Any], + *, + expected: frozenset[str], + context: str, +) -> None: + """Reject fields that are not part of the current versioned schema.""" + unknown = set(mapping) - expected + if unknown: + raise ValueError( + f"{context} contains unknown fields: {sorted(unknown, key=repr)!r}" + ) + class PromptGroupPhase(StrEnum): """Durable sampler-admission phase for an unfinished prompt group.""" @@ -382,8 +427,7 @@ def prepare_for_restart(self, cut: DataPlaneMutationCut) -> None: else: self.abandon_unsealed(cut, record.group_id) - @staticmethod - def _abandon_entire_group(record: PromptGroupRecoveryRecord) -> None: + def _abandon_entire_group(self, record: PromptGroupRecoveryRecord) -> None: """Discard every current sibling when an incomplete group is atomic. Sealed staging rows become unreferenced here. The controller's restore @@ -750,7 +794,7 @@ def __contains__(self, group_id: object) -> bool: return isinstance(group_id, str) and group_id in self._groups def state_dict(self) -> dict[str, Any]: - """Return the versioned metadata envelope used by later persistence.""" + """Return the versioned metadata persisted in ``rollout_recovery.pt``.""" self.assert_checkpoint_safe() groups = [] for record in self._groups.values(): @@ -812,6 +856,11 @@ def from_state_dict(cls, state: dict[str, Any]) -> Self: "rollout recovery state must be a dictionary, got " f"{type(state).__name__}" ) + _reject_unknown_fields( + state, + expected=_LEDGER_STATE_FIELDS, + context="rollout recovery state", + ) schema_version = state.get("schema_version") if ( isinstance(schema_version, bool) @@ -870,6 +919,11 @@ def _group_from_state( ) -> PromptGroupRecoveryRecord: if not isinstance(raw_group, dict): raise ValueError("rollout-recovery group must be a mapping") + _reject_unknown_fields( + raw_group, + expected=_GROUP_STATE_FIELDS, + context="rollout-recovery group", + ) group_id = raw_group.get("group_id") admission_id = raw_group.get("admission_id") prompt_id = raw_group.get("prompt_id") @@ -922,6 +976,11 @@ def _group_from_state( for generation_index, sibling_state in enumerate(siblings_state): if not isinstance(sibling_state, dict): raise ValueError("rollout-recovery sibling must be a mapping") + _reject_unknown_fields( + sibling_state, + expected=_SIBLING_STATE_FIELDS, + context="rollout-recovery sibling", + ) if sibling_state.get("generation_index") != generation_index: raise ValueError("generation indices must be contiguous") logical_id = f"{group_id}_g{generation_index}" @@ -932,6 +991,11 @@ def _group_from_state( for attempt_state in attempts_state: if not isinstance(attempt_state, dict): raise ValueError("rollout-recovery attempt must be a mapping") + _reject_unknown_fields( + attempt_state, + expected=_ATTEMPT_STATE_FIELDS, + context="rollout-recovery attempt", + ) raw_attempt_uuid = attempt_state.get("attempt_uuid") if ( not isinstance(raw_attempt_uuid, bytes) @@ -999,6 +1063,11 @@ def _group_from_state( raw_prompt_ref = raw_group.get("prompt_ref") if not isinstance(raw_prompt_ref, dict): raise ValueError("prompt_ref must be a mapping") + _reject_unknown_fields( + raw_prompt_ref, + expected=_PROMPT_REF_STATE_FIELDS, + context="rollout-recovery prompt_ref", + ) sample_id = raw_prompt_ref.get("sample_id") task_name = raw_prompt_ref.get("task_name") if not isinstance(sample_id, str) or not sample_id: @@ -1141,6 +1210,11 @@ def parse_rollout_recovery_state(state: object) -> ParsedRolloutRecoveryState: "rollout recovery sidecar must contain a dictionary, got " f"{type(state).__name__}" ) + _reject_unknown_fields( + state, + expected=_SIDECAR_STATE_FIELDS, + context="rollout recovery sidecar", + ) schema_version = state.get("schema_version") if ( isinstance(schema_version, bool) diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 0b79937764d..b3e058fe1c3 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -1558,6 +1558,7 @@ class _FakeCaptureBuffer(_FakeBuffer): def __init__(self): super().__init__() self.reserve_rollout_ids: list[list[str] | None] = [] + self.cleared_staging_key_batches: list[list[str]] = [] def reserve( self, *, weight_version, target_step=None, group_id=None, rollout_ids=None @@ -1570,8 +1571,9 @@ def reserve( rollout_ids=rollout_ids, ) - async def clear_staging_keys(self, staging_keys): - del staging_keys + async def clear_staging_keys(self, cut, staging_keys): + cut.require_live() + self.cleared_staging_key_batches.append(list(staging_keys)) def _receipt_record( @@ -1739,6 +1741,49 @@ async def _boom(_sample): _run(mgr.generate_for_finalization({"prompt": "p", "idx": 0})) assert len(buf.abort_calls) == 1 + def test_exhausted_capture_cleans_internally_owned_recovery_group(self, capsys): + buf = _FakeCaptureBuffer() + mgr = _make_capture_manager(buf) + mgr._retry_policy = RolloutRetryPolicy.single_attempt( + max_consecutive_dropped_prompts=1 + ) + + class _PartialCaptureImpl: + async def run_rollout( + self, + _sample, + *, + rollout_ids=None, + generation_indices=None, + on_completion=None, + recovery_granularity=RecoveryGranularity.SIBLING, + ): + del _sample, recovery_granularity + generation_index = generation_indices[0] + rollout_id = rollout_ids[generation_index] + receipt = { + "rollout_id": rollout_id, + "manifest": [{"staging_key": f"{rollout_id}/call"}], + } + completion = _receipt_record([rollout_id], [receipt]).completions[0] + await on_completion(generation_index, completion) + raise GenerationUnavailable("worker disappeared") + + mgr._impl = _PartialCaptureImpl() + + request = _run(mgr.generate_for_finalization({"prompt": "p", "idx": 0})) + + assert request is None + assert len(mgr.recovery_ledger) == 0 + first_rollout_ids = buf.reserve_rollout_ids[0] + assert first_rollout_ids is not None + assert buf.cleared_staging_key_batches == [[f"{first_rollout_ids[0]}/call"]] + assert ( + "dropping capture prompt idx=0 after 1 infrastructure failure(s) " + "(GenerationUnavailable: worker disappeared) [consecutive drop 1/1]" + in capsys.readouterr().out + ) + def test_cancel_after_controller_discard_preserves_cancelled_error(self): """A stale abort may delete lineage before rollout cleanup runs.""" diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index c61be7af2f7..afc03e6392e 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -127,6 +127,42 @@ def test_ledger_round_trip_preserves_group_ownership() -> None: assert restored.get_group("g7").phase is PromptGroupPhase.ADMITTED +@pytest.mark.parametrize( + ("path", "context"), + [ + ((), "rollout recovery state"), + (("groups", 0), "rollout-recovery group"), + (("groups", 0, "prompt_ref"), "rollout-recovery prompt_ref"), + (("groups", 0, "siblings", 0), "rollout-recovery sibling"), + (("groups", 0, "siblings", 0, "attempts", 0), "rollout-recovery attempt"), + ], +) +def test_ledger_restore_rejects_unknown_fields( + path: tuple[object, ...], context: str +) -> None: + ledger = RolloutRecoveryLedger() + _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=1, + target_step=7, + start_weight_version=6, + admitted=True, + ) + state = ledger.state_dict() + target: Any = state + for component in path: + target = target[component] + assert isinstance(target, dict) + target["unexpected"] = True + + with pytest.raises(ValueError, match=rf"{context} contains unknown fields"): + RolloutRecoveryLedger.from_state_dict(state) + + def test_checkpoint_state_round_trip_preserves_controller_and_ledger_state() -> None: ledger = RolloutRecoveryLedger() _reserve( @@ -155,6 +191,20 @@ def test_checkpoint_state_round_trip_preserves_controller_and_ledger_state() -> assert parsed.sampler_stamps_target_steps is True +def test_checkpoint_parser_rejects_unknown_sidecar_fields() -> None: + state = build_rollout_recovery_state( + RolloutRecoveryLedger(), + batch_shortfall={}, + sampler_stamps_target_steps=True, + ) + state["unexpected"] = True + + with pytest.raises( + ValueError, match="rollout recovery sidecar contains unknown fields" + ): + parse_rollout_recovery_state(state) + + def test_checkpoint_parser_defaults_fields_absent_from_older_state() -> None: parsed = parse_rollout_recovery_state(RolloutRecoveryLedger().state_dict()) @@ -627,6 +677,8 @@ def test_prompt_group_restart_retries_every_sibling_when_one_is_unfinished() -> _mutate(lambda cut: restored.prepare_for_restart(cut)) recovered = restored.get_group("g7") + assert recovered.agent_name == "genrm_agent" + assert recovered.recovery_granularity is RecoveryGranularity.PROMPT_GROUP assert [sibling.current_attempt.status for sibling in recovered.siblings] == [ RolloutAttemptStatus.ABANDONED, RolloutAttemptStatus.ABANDONED, @@ -669,7 +721,12 @@ def test_prompt_group_restart_keeps_a_fully_sealed_group() -> None: ) _mutate(lambda cut: ledger.mark_group_sealed(cut, "g7", results)) - restored = RolloutRecoveryLedger.from_state_dict(ledger.state_dict()) + state = ledger.state_dict() + restored = RolloutRecoveryLedger.from_state_dict(state) + _bind(restored, "g7", _prompt()) + + assert restored.state_dict() == state + _mutate(lambda cut: restored.prepare_for_restart(cut)) assert restored.get_group("g7").sealed_generation_indices == [0, 1] diff --git a/tests/unit/experience/test_rollouts.py b/tests/unit/experience/test_rollouts.py index e75dba5332e..32cac116ae9 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -2351,6 +2351,7 @@ def remote(self, inputs, timer_prefix): "nemo_gym": type("_Environment", (), {"run_rollouts": _RunRolloutsRemote()})() } manager._tokenizer = None + manager._effort_config = None with pytest.raises(ValueError, match="duplicate row index 0"): asyncio.run( diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index d9aa64136af..64546e2433c 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -81,16 +81,16 @@ from nemo_rl.data.utils import load_dataloader_state from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, KVBatchMeta from nemo_rl.data_plane.schema import ROUTE_PLAN_TAG -from nemo_rl.experience.route_plan import ( - ROUTE_PLAN_SCHEMA_VERSION, - RouteAssemblyPlan, - encode_route_plan, -) from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, ROLLOUT_RECOVERY_STATE_FILENAME, RolloutRecoveryLedger, ) +from nemo_rl.experience.route_plan import ( + ROUTE_PLAN_SCHEMA_VERSION, + RouteAssemblyPlan, + encode_route_plan, +) from nemo_rl.utils.checkpoint import CheckpointManager # Reuse the factory patches from the setup tests (same cross-module fixture @@ -1230,13 +1230,16 @@ def test_rollout_recovery_inventory_rejects_missing_staging_rows(self): ) actor._dp_client = _StagingInventoryDPClient([], partition_id=staging_partition) - with pytest.raises(RuntimeError, match=r"missing=\['sealed-key'\]"): - asyncio.run( - actor._validate_rollout_recovery_inventory( + async def validate_inventory() -> None: + async with DataPlaneCheckpointBarrier().mutation() as cut: + await actor._validate_rollout_recovery_inventory( + cut, replay_metadata=None, clear_unreferenced=False, ) - ) + + with pytest.raises(RuntimeError, match=r"missing=\['sealed-key'\]"): + asyncio.run(validate_inventory()) def test_rollout_recovery_inventory_merges_routes_and_clears_orphans(self): staging_partition = "rollout_staging" @@ -1273,12 +1276,15 @@ def test_rollout_recovery_inventory_merges_routes_and_clears_orphans(self): ) actor._dp_client = dp_client - asyncio.run( - actor._validate_rollout_recovery_inventory( - replay_metadata=replay_metadata, # type: ignore[arg-type] - clear_unreferenced=True, - ) - ) + async def validate_inventory() -> None: + async with DataPlaneCheckpointBarrier().mutation() as cut: + await actor._validate_rollout_recovery_inventory( + cut, + replay_metadata=replay_metadata, # type: ignore[arg-type] + clear_unreferenced=True, + ) + + asyncio.run(validate_inventory()) assert dp_client.clear_calls == [(["orphan-key"], staging_partition)] assert sorted(dp_client.sample_ids) == [route_key, "sealed-key"] diff --git a/tests/unit/single_controller/test_finalizer_lifecycle.py b/tests/unit/single_controller/test_finalizer_lifecycle.py index b508a2db1cf..c8fa52cea84 100644 --- a/tests/unit/single_controller/test_finalizer_lifecycle.py +++ b/tests/unit/single_controller/test_finalizer_lifecycle.py @@ -20,7 +20,7 @@ import threading from types import SimpleNamespace from typing import Any -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import ANY, AsyncMock, MagicMock import pytest @@ -148,6 +148,7 @@ def test_successful_actor_finalization_returns_actor_and_transfers_ownership() - assert ctrl._active_finalizers == 0 assert ctrl._finalizer_unknown_outcomes == 0 ctrl._buffer.commit_finalized.assert_awaited_once_with( + ANY, "group", meta, 3, diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 86df2696ec2..06e4db07d24 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -490,7 +490,7 @@ def test_build_trainer_initializes_reference_model_only_for_nonzero_kl( ) -def test_sibling_recovery_functional_config_resolves_to_runtime_contract( +def test_rollout_recovery_functional_config_resolves_to_runtime_contract( tmp_path: Path, ) -> None: """The two-phase Gym recovery fixture must pass SC config validation.""" @@ -538,7 +538,7 @@ def test_sibling_recovery_functional_config_resolves_to_runtime_contract( "++data_plane.simple.num_storage_units=2", "++data_plane.claim_meta_poll_interval_s=0.5", "++token_capture.enabled=true", - "++rollout_recovery.default_granularity=sibling", + "++rollout_recovery.default_granularity=prompt_group", "++async_rl.sampler.name=in_order", "++async_rl.sampler.max_lookahead_versions=1", "++async_rl.min_groups_for_streaming_train=4", @@ -563,7 +563,7 @@ def test_sibling_recovery_functional_config_resolves_to_runtime_contract( assert master_config.token_capture.enabled is True assert ( master_config.rollout_recovery.default_granularity - is RecoveryGranularity.SIBLING + is RecoveryGranularity.PROMPT_GROUP ) assert master_config.async_rl.rollout_failure.native.generation_timeout_s is None assert master_config.async_rl.rollout_failure.nemo_gym.rollout_timeout_s == 120 diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 47832a85876..6b0306d1964 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -192,6 +192,28 @@ def _run(coro): return asyncio.run(coro) +async def _commit_finalized( + buffer: TQReplayBuffer, + group_id: str, + meta: KVBatchMeta, + group_min_wv: int, + group_max_wv: int, + *, + staging_keys: list[str] | None = None, +) -> KVBatchMeta: + barrier = buffer._data_plane_checkpoint_barrier + assert barrier is not None + async with barrier.mutation() as cut: + return await buffer.commit_finalized( + cut, + group_id, + meta, + group_min_wv, + group_max_wv, + staging_keys=staging_keys, + ) + + def _make_record( rollout_metrics: dict[str, Any] | None = None, *, @@ -316,36 +338,31 @@ async def checkpoint() -> None: asyncio.run(exercise()) - def test_nested_mutation_does_not_deadlock_with_waiting_checkpoint(self): + def test_cancelled_mutation_releases_waiting_checkpoint(self): async def exercise() -> None: barrier = DataPlaneCheckpointBarrier() - outer_entered = asyncio.Event() - allow_nested = asyncio.Event() - nested_entered = asyncio.Event() + mutation_entered = asyncio.Event() checkpoint_entered = asyncio.Event() async def mutate() -> None: async with barrier.mutation(): - outer_entered.set() - await allow_nested.wait() - async with barrier.mutation(): - nested_entered.set() + mutation_entered.set() + await asyncio.Event().wait() async def checkpoint() -> None: async with barrier.checkpoint(): checkpoint_entered.set() mutation_task = asyncio.create_task(mutate()) - await outer_entered.wait() + await mutation_entered.wait() checkpoint_task = asyncio.create_task(checkpoint()) await asyncio.sleep(0) - allow_nested.set() - - await asyncio.wait_for(nested_entered.wait(), timeout=5.0) assert not checkpoint_entered.is_set() - await asyncio.wait_for( - asyncio.gather(mutation_task, checkpoint_task), timeout=5.0 - ) + + mutation_task.cancel() + with pytest.raises(asyncio.CancelledError): + await mutation_task + await asyncio.wait_for(checkpoint_task, timeout=5.0) assert checkpoint_entered.is_set() asyncio.run(exercise()) @@ -1431,6 +1448,36 @@ def test_reserve_records_rollout_ids(self): buf.reserve(weight_version=1) assert buf._rollout_ids_list[1] is None + def test_mutating_helpers_reject_expired_cut(self): + buf = self._make_capture_buffer(MultiPartitionFakeDataPlaneClient()) + group_id = buf.reserve(weight_version=1, rollout_ids=["r0"]) + meta = KVBatchMeta( + partition_id="rollout_data", + task_name=None, + sample_ids=[f"{group_id}_g0"], + fields=None, + ) + + async def get_expired_cut(): + barrier = buf._data_plane_checkpoint_barrier + assert barrier is not None + async with barrier.mutation() as cut: + return cut + + cut = _run(get_expired_cut()) + with pytest.raises(RuntimeError, match="no longer active"): + _run(buf.clear_staging_keys(cut, [])) + with pytest.raises(RuntimeError, match="no longer active"): + _run(buf.commit_finalized(cut, group_id, meta, 1, 1)) + with pytest.raises(RuntimeError, match="no longer active"): + _run( + buf._remove_groups_unlocked( + cut, + [group_id], + clear_data_plane=False, + ) + ) + def test_commit_finalized_fills_slot_with_group_min_wv(self): dp = MultiPartitionFakeDataPlaneClient() buf = self._make_capture_buffer(dp) @@ -1443,7 +1490,8 @@ def test_commit_finalized_fills_slot_with_group_min_wv(self): fields=None, ) _run( - buf.commit_finalized( + _commit_finalized( + buf, group_id, meta, group_min_wv=3, @@ -1465,7 +1513,7 @@ def test_commit_finalized_raises_for_evicted_slot(self): partition_id="rollout_data", task_name=None, sample_ids=[], fields=None ) with pytest.raises(ValueError, match="no live slot"): - _run(buf.commit_finalized("ghost", meta, group_min_wv=0, group_max_wv=0)) + _run(_commit_finalized(buf, "ghost", meta, 0, 0)) def test_commit_finalized_verifies_full_plan_manifest_ownership(self): dp = MultiPartitionFakeDataPlaneClient() @@ -1490,7 +1538,8 @@ def test_commit_finalized_verifies_full_plan_manifest_ownership(self): with pytest.raises(ValueError, match="ownership does not match"): _run( - buf.commit_finalized( + _commit_finalized( + buf, group_id, meta, group_min_wv=1, @@ -1533,7 +1582,8 @@ def test_remove_clears_staging_rows_alongside_canonical(self): fields=None, ) _run( - buf.commit_finalized( + _commit_finalized( + buf, group_id, meta, group_min_wv=1, @@ -1562,7 +1612,7 @@ def test_remove_without_staging_partition_skips_staging_clear(self): sample_ids=[f"{group_id}_g0"], fields=None, ) - _run(buf.commit_finalized(group_id, meta, group_min_wv=1, group_max_wv=1)) + _run(_commit_finalized(buf, group_id, meta, 1, 1)) _run(buf.remove([0], remove_in_dp=True)) partitions_cleared = {p for p, _ in dp.clear_calls_by_partition} assert partitions_cleared == {"rollout_data"} @@ -1584,7 +1634,8 @@ def clear_samples(self, sample_ids, partition_id): fields=None, ) _run( - buf.commit_finalized( + _commit_finalized( + buf, group_id, meta, group_min_wv=1, From 4d31868c79a99108b11d9bcfa436aef5d4dca5bc Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 3 Sep 2026 18:54:18 -0400 Subject: [PATCH 17/28] refactor(rollout): rename resolved recovery granularity Signed-off-by: Anish Mahishi --- .../single_controller_utils/config.py | 18 ++++++++++++------ tests/unit/experience/test_rollout_recovery.py | 2 ++ 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 394f59c34e1..487f61173a3 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -624,8 +624,12 @@ class TokenCaptureConfig(BaseModel, extra="allow"): @dataclass(frozen=True) -class ResolvedRolloutRecovery: - """Policy coordinates stamped on one newly reserved ledger group.""" +class RecoveryGranularityResolution: + """Recovery granularity selected for a prompt-group reservation. + + ``agent_name`` is copied from the prompt when present. ``granularity`` is + selected from an agent override, task override, or the global default. + """ agent_name: Optional[str] granularity: RecoveryGranularity @@ -657,7 +661,9 @@ class RolloutRecoveryConfig(BaseModel, extra="allow"): default_factory=dict ) - def resolve_for_prompt(self, prompt: Mapping[str, Any]) -> ResolvedRolloutRecovery: + def resolve_for_prompt( + self, prompt: Mapping[str, Any] + ) -> RecoveryGranularityResolution: """Resolve one new group using agent, then task, then the global default.""" extra_env_info = prompt.get("extra_env_info") agent_name: Optional[str] = None @@ -673,7 +679,7 @@ def resolve_for_prompt(self, prompt: Mapping[str, Any]) -> ResolvedRolloutRecove if agent_name is not None: override = self.agent_granularity_overrides.get(agent_name) if override is not None: - return ResolvedRolloutRecovery(agent_name, override) + return RecoveryGranularityResolution(agent_name, override) task_name = prompt.get("task_name") if task_name is not None and not isinstance(task_name, str): @@ -681,8 +687,8 @@ def resolve_for_prompt(self, prompt: Mapping[str, Any]) -> ResolvedRolloutRecove if task_name is not None: override = self.task_granularity_overrides.get(task_name) if override is not None: - return ResolvedRolloutRecovery(agent_name, override) - return ResolvedRolloutRecovery(agent_name, self.default_granularity) + return RecoveryGranularityResolution(agent_name, override) + return RecoveryGranularityResolution(agent_name, self.default_granularity) class MasterConfig(BaseModel, extra="allow"): diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index afc03e6392e..4b79dfbb866 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -302,7 +302,9 @@ def test_recovery_config_resolves_agent_then_task_then_default() -> None: assert agent_policy.agent_name == "genrm_agent" assert agent_policy.granularity is RecoveryGranularity.PROMPT_GROUP + assert task_policy.agent_name is None assert task_policy.granularity is RecoveryGranularity.PROMPT_GROUP + assert default_policy.agent_name is None assert default_policy.granularity is RecoveryGranularity.SIBLING From 2755a266f1141f0b10d217977f20a4d42a45f453 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 3 Sep 2026 23:11:14 -0400 Subject: [PATCH 18/28] fix(sc): enforce data-plane barrier ownership Signed-off-by: Anish Mahishi --- .../algorithms/async_utils/replay_buffer.py | 29 ++++++++++++-- nemo_rl/algorithms/single_controller.py | 9 +++-- .../test_tq_replay_buffer.py | 38 +++++++++++++++++++ 3 files changed, 70 insertions(+), 6 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 9427fc06afb..adebeec8fd0 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -230,19 +230,35 @@ def __init__(self) -> None: self._condition = asyncio.Condition() self._checkpoint_active = False self._active_mutations = 0 + self._section_holders: set[asyncio.Task[Any]] = set() + + def _current_task(self) -> asyncio.Task[Any]: + """Return the task entering a barrier section and reject reentrancy.""" + task = asyncio.current_task() + if task is None: + raise RuntimeError("data-plane barrier sections require an asyncio task") + if task in self._section_holders: + raise RuntimeError( + "this task already holds a data-plane barrier section; pass the " + "DataPlaneMutationCut you already have instead of opening another" + ) + return task @asynccontextmanager async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: """Yield a live cut after any active checkpoint exits.""" async with self._condition: + task = self._current_task() await self._condition.wait_for(lambda: not self._checkpoint_active) self._active_mutations += 1 + self._section_holders.add(task) cut = DataPlaneMutationCut(self) try: yield cut finally: cut._invalidate() async with self._condition: + self._section_holders.discard(task) self._active_mutations -= 1 if self._active_mutations == 0: self._condition.notify_all() @@ -251,6 +267,7 @@ async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: async def checkpoint(self) -> AsyncIterator[DataPlaneMutationCut]: """Yield a live capability after blocking and draining all mutations.""" async with self._condition: + task = self._current_task() await self._condition.wait_for(lambda: not self._checkpoint_active) self._checkpoint_active = True try: @@ -259,12 +276,14 @@ async def checkpoint(self) -> AsyncIterator[DataPlaneMutationCut]: self._checkpoint_active = False self._condition.notify_all() raise + self._section_holders.add(task) cut = DataPlaneMutationCut(self) try: yield cut finally: cut._invalidate() async with self._condition: + self._section_holders.discard(task) self._checkpoint_active = False self._condition.notify_all() @@ -1150,7 +1169,7 @@ async def commit( "the async message-log flattening path." ) trace_rollout_payload(keys=sample_ids, data=train_batch) - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation() as cut: try: await call_data_plane( self._dp_client, @@ -1199,6 +1218,7 @@ async def commit( # deterministic IDs while retaining the barrier mutation slot. try: await self._clear_samples_unlocked( + cut, sample_ids=list(sample_ids), ) except BaseException as rollback_error: @@ -1439,7 +1459,7 @@ async def _remove_groups_unlocked( if dropped_sample_ids: try: await self._clear_samples_unlocked( - sample_ids=dropped_sample_ids, + cut, sample_ids=dropped_sample_ids ) except Exception as error: raise RuntimeError( @@ -1760,8 +1780,11 @@ def size(self) -> int: def __len__(self) -> int: return len(self.meta_list) - async def _clear_samples_unlocked(self, *, sample_ids: list[str]) -> None: + async def _clear_samples_unlocked( + self, cut: DataPlaneMutationCut, *, sample_ids: list[str] + ) -> None: """Clear rows while the caller holds a barrier mutation slot.""" + cut.require_live() await call_data_plane( self._dp_client, "clear_samples", diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index c6d2c51feee..20ce8baca79 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1416,8 +1416,11 @@ async def _finalize_with_actor( self._finalizer_metrics_by_group[request.group_id] = dict(finalized.metrics) return finalized - async def _cleanup_consumed_metas_unlocked(self, metas: list[KVBatchMeta]) -> None: + async def _cleanup_consumed_metas_unlocked( + self, cut: DataPlaneMutationCut, metas: list[KVBatchMeta] + ) -> None: """Clear consumed ownership while holding a barrier mutation slot.""" + cut.require_live() canonical_by_partition: dict[str, list[str]] = {} staging_by_partition: dict[str, list[str]] = {} for meta in metas: @@ -1458,8 +1461,8 @@ async def _cleanup_consumed_metas_unlocked(self, metas: list[KVBatchMeta]) -> No async def _cleanup_consumed_metas(self, metas: list[KVBatchMeta]) -> None: """Clear consumed rows without racing a native TQ checkpoint.""" - async with self._data_plane_checkpoint_barrier.mutation(): - await self._cleanup_consumed_metas_unlocked(metas) + async with self._data_plane_checkpoint_barrier.mutation() as cut: + await self._cleanup_consumed_metas_unlocked(cut, metas) # ── the three pumps + the inline advantage stage ─────────────────────── diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 6b0306d1964..9967ace22cf 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -396,6 +396,44 @@ async def checkpoint(tag: str) -> None: asyncio.run(exercise()) + @pytest.mark.parametrize( + ("outer_section", "inner_section"), + [ + ("mutation", "mutation"), + ("mutation", "checkpoint"), + ("checkpoint", "mutation"), + ("checkpoint", "checkpoint"), + ], + ) + def test_same_task_cannot_nest_barrier_sections( + self, outer_section: str, inner_section: str + ): + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + outer = ( + barrier.mutation() + if outer_section == "mutation" + else barrier.checkpoint() + ) + inner = ( + barrier.mutation() + if inner_section == "mutation" + else barrier.checkpoint() + ) + + async with outer: + with pytest.raises( + RuntimeError, match="already holds a data-plane barrier section" + ): + async with inner: + pytest.fail("nested barrier section unexpectedly opened") + + # A rejected nested section must not poison later acquisitions. + async with barrier.mutation() as cut: + cut.require_live() + + asyncio.run(exercise()) + class TestTQReplayBufferReserveCommit: def test_reserve_rejects_duplicate_live_group_id(self): From 839f38d16281af15fcdfd50d38f2bee77134e30f Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 5 Sep 2026 19:22:16 -0400 Subject: [PATCH 19/28] fix(rollout): preserve mask state across recovery Signed-off-by: Anish Mahishi --- .../single_controller_utils/config.py | 17 ------ nemo_rl/data_plane/tq_token_sink.py | 2 +- nemo_rl/experience/rollout_manager.py | 27 ++++------ nemo_rl/experience/rollout_recovery.py | 52 ++++++++++++++++--- .../test_rollout_reassembler_actor.py | 1 - .../unit/experience/test_rollout_recovery.py | 8 ++- 6 files changed, 64 insertions(+), 43 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 487f61173a3..e6ab7ad0e91 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -35,7 +35,6 @@ ReadyFirstSamplerConfig, SamplerConfig, required_buffer_capacity_for_config, - sampler_supports_buffer_checkpoint, ) from nemo_rl.algorithms.grpo import ( _REWARD_PENALTY_FLAGS, @@ -1151,22 +1150,6 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: "async_rl.max_buffered_rollouts; excess finalizer actors cannot be busy", stacklevel=2, ) - if ( - token_capture_config.enabled - and master_config.checkpointing["enabled"] - and master_config.checkpointing.get("save_data_plane") - and sampler_supports_buffer_checkpoint(async_config.sampler) - ): - raise NotImplementedError( - "token_capture.enabled does not support rollout-recovery " - "checkpointing (checkpointing.save_data_plane=true with a " - "replay-checkpoint-capable sampler): the capture dispatch path " - "does not thread lineage group ids, so the recovery ledger would " - "never be reaped and a resume would re-dispatch every restored " - "prompt on top of the recovered replay groups. Set " - "checkpointing.enabled=false, or use a sampler without buffer " - "checkpointing." - ) if token_capture_config.enabled and reward_penalties_enabled: warnings.warn( "reward_penalties are enabled but token-capture receipt rollouts " diff --git a/nemo_rl/data_plane/tq_token_sink.py b/nemo_rl/data_plane/tq_token_sink.py index 83a7974e6ea..ed8568a2718 100644 --- a/nemo_rl/data_plane/tq_token_sink.py +++ b/nemo_rl/data_plane/tq_token_sink.py @@ -37,7 +37,7 @@ import json import logging from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any import ray import torch diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 11ca9d9a595..09b506303e6 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -2089,12 +2089,18 @@ async def _record_streamed_completion( f"receipt={receipt.get('rollout_id')!r}, " f"expected={gate_rollout_id!r}" ) + mask_sample = bool( + (((completion.env_extras or {}).get("instance_config") or {}).get( + MASK_SAMPLE, False + )) + ) if recovery_group.recovery_granularity is RecoveryGranularity.PROMPT_GROUP: result = SiblingSealResult( gate_rollout_id=gate_rollout_id, receipt=receipt, reward=completion.reward, + mask_sample=mask_sample, ) previous = pending_group_results.get(generation_index) if previous is not None: @@ -2123,20 +2129,9 @@ async def _record_streamed_completion( gate_rollout_id=gate_rollout_id, receipt=receipt, reward=completion.reward, + mask_sample=mask_sample, ) - mask_sample_by_index: dict[int, bool] = {} - - async def _record_completion( - generation_index: int, completion: Completion - ) -> None: - mask_sample_by_index[generation_index] = bool( - (((completion.env_extras or {}).get("instance_config") or {}).get( - MASK_SAMPLE, False - )) - ) - await _record_streamed_completion(generation_index, completion) - try: if inflight_registry is not None: current_task = asyncio.current_task() @@ -2154,7 +2149,7 @@ async def _record_completion( attempt_input_sample, rollout_ids=list(rollout_ids), generation_indices=pending_indices, - on_completion=_record_completion, + on_completion=_record_streamed_completion, recovery_granularity=recovery_group.recovery_granularity, ) finally: @@ -2165,6 +2160,7 @@ async def _record_completion( canonical_sample_ids, receipts, rewards, + mask_sample, ) = self._recovery_ledger.finalization_inputs(group_id) request = ReassemblyRequest( group_id=group_id, @@ -2174,10 +2170,7 @@ async def _record_completion( rewards=tuple(rewards), fallback_weight_version=start_version, prompt_idx=int(recovery_group.prompt_id), - mask_sample=tuple( - mask_sample_by_index.get(index, False) - for index in range(len(receipts)) - ), + mask_sample=tuple(mask_sample), loss_multiplier=float(input_sample.get("loss_multiplier", 1.0)), ) from nemo_rl.experience.rollout_reassembler_actor import ( diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index 56b14bf0599..0c9f31ad552 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -33,7 +33,7 @@ from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneMutationCut from nemo_rl.data.interfaces import DatumSpec -ROLLOUT_RECOVERY_SCHEMA_VERSION = 4 +ROLLOUT_RECOVERY_SCHEMA_VERSION = 5 _SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS = {ROLLOUT_RECOVERY_SCHEMA_VERSION} ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" RolloutRecoveryState: TypeAlias = dict[str, Any] @@ -65,7 +65,14 @@ _PROMPT_REF_STATE_FIELDS = frozenset({"sample_id", "task_name"}) _SIBLING_STATE_FIELDS = frozenset({"generation_index", "attempts"}) _ATTEMPT_STATE_FIELDS = frozenset( - {"attempt_uuid", "status", "receipt", "reward", "staging_keys"} + { + "attempt_uuid", + "status", + "receipt", + "reward", + "mask_sample", + "staging_keys", + } ) @@ -168,6 +175,7 @@ class RolloutAttemptRecord: status: RolloutAttemptStatus receipt: Optional[dict[str, Any]] = None reward: Optional[float] = None + mask_sample: Optional[bool] = None staging_keys: list[str] = field(default_factory=list) @property @@ -272,6 +280,7 @@ class SiblingSealResult: # into a masked placeholder, matching the base token-capture contract. receipt: Optional[dict[str, Any]] reward: float + mask_sample: bool = False def _new_attempt() -> RolloutAttemptRecord: @@ -567,6 +576,7 @@ def mark_sibling_sealed( gate_rollout_id: str, receipt: Optional[dict[str, Any]], reward: float, + mask_sample: bool = False, ) -> None: """Record one streamed sibling receipt as soon as the row arrives.""" cut.require_live() @@ -577,6 +587,8 @@ def mark_sibling_sealed( attempt = sibling.current_attempt expected_gate_rollout_id = record.gate_rollout_id(generation_index) staging_keys = _receipt_staging_keys(receipt) + if not isinstance(mask_sample, bool): + raise TypeError("mask_sample must be a bool") if gate_rollout_id != expected_gate_rollout_id: raise ValueError( "streamed rollout identity mismatch: " @@ -591,6 +603,7 @@ def mark_sibling_sealed( if ( attempt.receipt == receipt and attempt.reward == float(reward) + and attempt.mask_sample is mask_sample and attempt.staging_keys == staging_keys ): return @@ -607,6 +620,7 @@ def mark_sibling_sealed( attempt.receipt = copy.deepcopy(receipt) attempt.reward = float(reward) + attempt.mask_sample = mask_sample attempt.staging_keys = staging_keys attempt.status = RolloutAttemptStatus.SEALED if all( @@ -666,6 +680,8 @@ def mark_group_sealed( f"receipt={result.receipt.get('rollout_id')!r}, " f"expected={expected_gate_rollout_id!r}" ) + if not isinstance(result.mask_sample, bool): + raise TypeError("mask_sample must be a bool") validated.append((attempt, result, _receipt_staging_keys(result.receipt))) # Validate the complete cohort before changing any sibling. A checkpoint @@ -673,6 +689,7 @@ def mark_group_sealed( for attempt, result, staging_keys in validated: attempt.receipt = copy.deepcopy(result.receipt) attempt.reward = float(result.reward) + attempt.mask_sample = result.mask_sample attempt.staging_keys = staging_keys attempt.status = RolloutAttemptStatus.SEALED record.status = PromptGroupStatus.READY_TO_FINALIZE @@ -710,8 +727,14 @@ def abandon_unsealed(self, cut: DataPlaneMutationCut, group_id: str) -> None: def finalization_inputs( self, group_id: str - ) -> tuple[list[str], list[str], list[Optional[dict[str, Any]]], list[float]]: - """Return physical IDs, canonical IDs, receipts and rewards in sibling order.""" + ) -> tuple[ + list[str], + list[str], + list[Optional[dict[str, Any]]], + list[float], + list[bool], + ]: + """Return sealed finalization inputs in stable sibling order.""" record = self._require_group(group_id) if record.status != PromptGroupStatus.READY_TO_FINALIZE: raise ValueError( @@ -719,9 +742,14 @@ def finalization_inputs( ) receipts: list[Optional[dict[str, Any]]] = [] rewards: list[float] = [] + mask_sample: list[bool] = [] for sibling in record.siblings: attempt = sibling.current_attempt - if attempt.status != RolloutAttemptStatus.SEALED or attempt.reward is None: + if ( + attempt.status != RolloutAttemptStatus.SEALED + or attempt.reward is None + or attempt.mask_sample is None + ): raise ValueError( "logical rollout " f"{record.logical_rollout_id(sibling.generation_index)!r} " @@ -729,11 +757,13 @@ def finalization_inputs( ) receipts.append(copy.deepcopy(attempt.receipt)) rewards.append(attempt.reward) + mask_sample.append(attempt.mask_sample) return ( record.gate_rollout_ids, record.logical_rollout_ids, receipts, rewards, + mask_sample, ) def mark_finalization_started( @@ -834,6 +864,7 @@ def state_dict(self) -> dict[str, Any]: "status": attempt.status.value, "receipt": copy.deepcopy(attempt.receipt), "reward": attempt.reward, + "mask_sample": attempt.mask_sample, "staging_keys": list(attempt.staging_keys), } for attempt in sibling.attempts @@ -1020,6 +1051,7 @@ def _group_from_state( ) from error receipt = attempt_state.get("receipt") reward = attempt_state.get("reward") + mask_sample = attempt_state.get("mask_sample") staging_keys = attempt_state.get("staging_keys") if not isinstance(staging_keys, list) or not all( isinstance(key, str) for key in staging_keys @@ -1028,6 +1060,8 @@ def _group_from_state( if attempt_status == RolloutAttemptStatus.SEALED: if not isinstance(reward, (int, float)): raise ValueError("sealed attempts require a reward") + if not isinstance(mask_sample, bool): + raise ValueError("sealed attempts require a boolean mask_sample") if receipt is None: if staging_keys: raise ValueError( @@ -1042,7 +1076,12 @@ def _group_from_state( raise ValueError( "sealed attempt receipt must be a mapping or None" ) - elif receipt is not None or reward is not None or staging_keys: + elif ( + receipt is not None + or reward is not None + or mask_sample is not None + or staging_keys + ): raise ValueError("only sealed attempts may retain receipt data") attempts.append( RolloutAttemptRecord( @@ -1050,6 +1089,7 @@ def _group_from_state( status=attempt_status, receipt=copy.deepcopy(receipt), reward=float(reward) if reward is not None else None, + mask_sample=mask_sample, staging_keys=list(staging_keys), ) ) diff --git a/tests/unit/experience/test_rollout_reassembler_actor.py b/tests/unit/experience/test_rollout_reassembler_actor.py index 931076c9034..e381abd619d 100644 --- a/tests/unit/experience/test_rollout_reassembler_actor.py +++ b/tests/unit/experience/test_rollout_reassembler_actor.py @@ -51,7 +51,6 @@ def _request() -> ReassemblyRequest: ), rewards=(1.0,), mask_sample=(False,), - prompt_idx=0, fallback_weight_version=4, ) diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 4b79dfbb866..5303b9dfff7 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -556,6 +556,7 @@ def test_restart_preserves_sealed_sibling_and_retries_only_interrupted_one() -> "manifest": [{"staging_key": "g7/sibling-0/call-0"}], }, reward=1.0, + mask_sample=True, ) ) @@ -621,6 +622,7 @@ def test_missing_receipt_is_a_restart_safe_sealed_placeholder( gate_rollout_id=gate_ids[generation_index], receipt=receipt, reward=float(generation_index), + mask_sample=generation_index == 0, ) ) ) @@ -634,6 +636,7 @@ def test_missing_receipt_is_a_restart_safe_sealed_placeholder( gate_rollout_id=gate_ids[generation_index], receipt=receipt, reward=float(generation_index), + mask_sample=generation_index == 0, ) for generation_index, receipt in enumerate(receipts) }, @@ -642,12 +645,15 @@ def test_missing_receipt_is_a_restart_safe_sealed_placeholder( state = ledger.state_dict() restored = RolloutRecoveryLedger.from_state_dict(state) - physical_ids, _, restored_receipts, rewards = restored.finalization_inputs("g7") + physical_ids, _, restored_receipts, rewards, mask_sample = ( + restored.finalization_inputs("g7") + ) assert physical_ids == gate_ids assert restored_receipts[0] is None assert restored_receipts[1] == receipts[1] assert rewards == [0.0, 1.0] + assert mask_sample == [True, False] state["schema_version"] = 3 with pytest.raises(ValueError, match="Unsupported rollout-recovery schema version"): From 284d4044ca12091ed6644162758409ff912d9e4a Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 5 Sep 2026 19:48:32 -0400 Subject: [PATCH 20/28] fix(rollout): repair recovery test integration Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 14 ++++++-------- nemo_rl/experience/rollout_reassembler_actor.py | 1 - .../functional/grpo_async_gym_single_controller.sh | 10 ++++++++-- tests/unit/experience/test_rollout_manager.py | 3 ++- .../experience/test_rollout_reassembler_actor.py | 1 + tests/unit/single_controller/test_rollout_pump.py | 1 + tests/unit/single_controller/test_setup.py | 2 ++ .../single_controller/test_tq_replay_buffer.py | 8 ++++++-- 8 files changed, 26 insertions(+), 14 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 20ce8baca79..eda20e0e259 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1125,14 +1125,12 @@ async def _ray_get(self, obj_ref: Any) -> Any: async def _call_dp(self, method_name: str, **kwargs) -> Any: """Call a DataPlaneClient method or a Ray actor exposing that method.""" - method = getattr(self._dp_client, method_name) - remote = getattr(method, "remote", None) - if remote is not None: - return await self._ray_get(remote(**kwargs)) - result = method(**kwargs) - if asyncio.iscoroutine(result): - return await result - return result + return await call_data_plane( + self._dp_client, + method_name, + offload_sync=True, + **kwargs, + ) async def _save_data_plane_checkpoint( self, diff --git a/nemo_rl/experience/rollout_reassembler_actor.py b/nemo_rl/experience/rollout_reassembler_actor.py index aefd80f3337..d22923f5caf 100644 --- a/nemo_rl/experience/rollout_reassembler_actor.py +++ b/nemo_rl/experience/rollout_reassembler_actor.py @@ -52,7 +52,6 @@ class ReassemblyRequest: """Metadata-only input for one prompt group's finalization.""" group_id: str - prompt_idx: int rollout_ids: tuple[str, ...] canonical_sample_ids: tuple[str, ...] receipts: tuple[Optional[dict[str, Any]], ...] diff --git a/tests/functional/grpo_async_gym_single_controller.sh b/tests/functional/grpo_async_gym_single_controller.sh index 323cd96de95..a6eb6d7f200 100755 --- a/tests/functional/grpo_async_gym_single_controller.sh +++ b/tests/functional/grpo_async_gym_single_controller.sh @@ -53,8 +53,14 @@ cd - # smoke test, we trim all but the first tool TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl -jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH -jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH +jq -c ' + .agent_ref //= {"name": "workplace_assistant_simple_agent"} + | .responses_create_params.tools |= (.[0:1]) +' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c ' + .agent_ref //= {"name": "workplace_assistant_simple_agent"} + | .responses_create_params.tools |= (.[0:1]) +' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ $SC_ENTRYPOINT \ diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index b3e058fe1c3..fa631539dff 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -751,6 +751,7 @@ def test_rollout_manager_forwards_log_full_result_tables(): "task_to_env": {}, "num_generations_per_prompt": 1, "max_seq_len": 1, + "rollout_recovery_config": RolloutRecoveryConfig(), "generation_config": { "stop_strings": None, "stop_token_ids": None, @@ -1687,7 +1688,7 @@ def test_request_carries_env_mask_flags(self): buf, instance_configs=[{"mask_sample": True}, {"other": 1}] ) - request = _run(mgr.generate_for_finalization({"prompt": "p"})) + request = _run(mgr.generate_for_finalization({"prompt": "p", "idx": 0})) # The gym mask flag is read from env_extras exactly like the token # path's _mask_sample_flags. truncated is not part of this request -- diff --git a/tests/unit/experience/test_rollout_reassembler_actor.py b/tests/unit/experience/test_rollout_reassembler_actor.py index e381abd619d..85a6896bef8 100644 --- a/tests/unit/experience/test_rollout_reassembler_actor.py +++ b/tests/unit/experience/test_rollout_reassembler_actor.py @@ -126,6 +126,7 @@ def test_rpc_dataclass_fields_are_classified() -> None: assert {f.name for f in fields(ReassemblyRequest)} == { "group_id", "rollout_ids", + "canonical_sample_ids", "receipts", "rewards", "fallback_weight_version", diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 63ea1f0a22e..34db681229f 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -1143,6 +1143,7 @@ class _SplitRolloutManager: def __init__(self) -> None: self.generated = 0 self.two_generated = asyncio.Event() + self.stats = SimpleNamespace(committed=0) async def generate_for_finalization( self, diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 06e4db07d24..77fab5315cc 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -939,8 +939,10 @@ def test_invalid_config_fails_before_setup_factories( elif invalid_case == "gym_on_sglang": mc = _make_master_config(colocated=False, backend="sglang") elif invalid_case == "prompt_group_recovery_without_capture": + mc = _make_master_config() mc.rollout_recovery.default_granularity = RecoveryGranularity.PROMPT_GROUP elif invalid_case == "recovery_override_without_capture": + mc = _make_master_config() mc.rollout_recovery.task_granularity_overrides = { "genrm": RecoveryGranularity.PROMPT_GROUP } diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 9967ace22cf..1d2f8fa6ee1 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -1250,6 +1250,7 @@ def test_token_capture_round_trip_restores_staging_cleanup_ownership(self): partition_id="rollout_data", pad_value_dict={"token_ids": 0}, staging_partition_id="rollout_staging", + include_message_violation_fields=False, ) assert _load(restored, state) == 1 @@ -1723,8 +1724,11 @@ async def put_samples( result = FakeDataPlaneClient.put_samples( self, sample_ids, partition_id, fields=fields, tags=tags ) - # Simulate the sampler evicting the slot mid-write. - await self.buf.remove([0], remove_in_dp=False) + # Simulate another task evicting the slot mid-write. Barrier + # sections are deliberately non-reentrant within one task, but + # independent mutation tasks may overlap when no checkpoint is + # active. + await asyncio.create_task(self.buf.remove([0], remove_in_dp=False)) return result dp = EvictDuringPut() From cb78d634f9ba7a1009babdee10fa0dcafcb50f14 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 5 Sep 2026 19:59:16 -0400 Subject: [PATCH 21/28] fix(rollout): convert streamed capture receipts Signed-off-by: Anish Mahishi --- nemo_rl/experience/rollout_manager.py | 7 ++- tests/unit/experience/test_rollout_manager.py | 45 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 09b506303e6..a8acf1ac17e 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -1073,7 +1073,12 @@ async def _stream_rows( ) results[rowidx] = result if on_completion is not None: - await on_completion(rowidx, self._result_to_completion(result)) + # Use the same conversion path as completed groups so streamed + # recovery records inherit the current mask and reward semantics. + # Completion callbacks are token-capture receipt-only, making this + # conversion lightweight and safe to repeat during group metrics. + row_completions, _ = self._results_to_completions([result]) + await on_completion(rowidx, row_completions[0]) if timing_metrics is not None: env_timing_metrics = timing_metrics diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index fa631539dff..fc3174a2e9b 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -850,6 +850,51 @@ def test_receipt_completion_drops_mask_flag_when_gate_off(): assert completion.env_extras["instance_config"]["other_key"] == "kept" +def test_streamed_receipt_callback_uses_current_completion_conversion(): + class _RunRolloutsRemote: + def options(self, *, num_returns): + assert num_returns == "streaming" + return self + + def remote(self, pending, timer_prefix): + del pending, timer_prefix + + async def result_ref(): + return 0, _mask_gate_receipt_result(), None + + async def stream(): + yield result_ref() + + return stream() + + impl = _nemo_gym_impl(False) + env = type("_Environment", (), {"run_rollouts": _RunRolloutsRemote()})() + results = [None] + shaping = [None] + streamed = [] + + async def on_completion(generation_index, completion): + streamed.append((generation_index, completion)) + + _run( + impl._stream_rows( + env, + [{"_rowidx": 0}], + results, + shaping, + 1, + "timing/test", + on_completion=on_completion, + ) + ) + + assert len(streamed) == 1 + generation_index, completion = streamed[0] + assert generation_index == 0 + assert completion.env_extras["ng_rollout_id"] == "r0" + assert "mask_sample" not in completion.env_extras["instance_config"] + + @pytest.mark.parametrize("log_full_result_tables", [False, True]) def test_nemo_gym_full_result_tables_are_opt_in(log_full_result_tables): impl = _nemo_gym_impl(True, log_full_result_tables=log_full_result_tables) From 9d2b3f6be2cf8f3f2587bdf6b39fdf796f5b45d3 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 5 Sep 2026 21:25:47 -0400 Subject: [PATCH 22/28] fix(rollout): harden recovery schema coverage Signed-off-by: Anish Mahishi --- ...po_math_1B_megatron_single_controller.yaml | 6 +- ...po_math_1B_megatron_single_controller.yaml | 6 +- .../algorithms/async_utils/replay_buffer.py | 2 +- nemo_rl/algorithms/single_controller.py | 1 + nemo_rl/experience/rollout_recovery.py | 16 +- tests/unit/experience/test_rollout_manager.py | 4 + .../unit/experience/test_rollout_recovery.py | 145 ++++++++++++++++++ 7 files changed, 158 insertions(+), 22 deletions(-) diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index 9444628eac1..3689a031b04 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -194,10 +194,8 @@ rollout_recovery: # task_granularity_overrides: {math: prompt_group} task_granularity_overrides: {} -# Leaving either map above non-empty requires token capture, which is a separate -# top-level section (not in this file, and off by default): -# token_capture: -# enabled: true +# Leaving either map above non-empty requires the top-level token_capture section +# above to set enabled: true. cluster: # Master ports inherit the shared 1400-1999 band from grpo_math_1B.yaml. diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 7fdd5a6f8f8..62628dd67eb 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -232,10 +232,8 @@ rollout_recovery: # task_granularity_overrides: {math: prompt_group} task_granularity_overrides: {} -# Leaving either map above non-empty requires token capture, which is a separate -# top-level section (not in this file, and off by default): -# token_capture: -# enabled: true +# Leaving either map above non-empty requires the top-level token_capture section +# above to set enabled: true. cluster: # Master ports inherit the shared 1400-1999 band from ppo_math_1B.yaml. diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index adebeec8fd0..60f94efd3eb 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -1783,7 +1783,7 @@ def __len__(self) -> int: async def _clear_samples_unlocked( self, cut: DataPlaneMutationCut, *, sample_ids: list[str] ) -> None: - """Clear rows while the caller holds a barrier mutation slot.""" + """Clear rows while the caller owns the provided live mutation cut.""" cut.require_live() await call_data_plane( self._dp_client, diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index eda20e0e259..dabc88c978a 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1233,6 +1233,7 @@ async def _cleanup_known_finalization_request_unlocked( request: "ReassemblyRequest", ) -> None: """Clear known request ownership while holding a barrier mutation slot.""" + cut.require_live() errors: list[BaseException] = [] try: await self._call_dp( diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index 0c9f31ad552..b0473d40ca3 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -23,6 +23,7 @@ from __future__ import annotations import copy +import dataclasses import uuid from collections.abc import Mapping from dataclasses import dataclass, field @@ -1163,20 +1164,9 @@ def _require_group(self, group_id: str) -> PromptGroupRecoveryRecord: @staticmethod def _copy_group(record: PromptGroupRecoveryRecord) -> PromptGroupRecoveryRecord: """Copy mutable lineage metadata without duplicating the prompt payload.""" - return PromptGroupRecoveryRecord( - group_id=record.group_id, - admission_id=record.admission_id, - prompt_id=record.prompt_id, - prompt_ref=record.prompt_ref, - agent_name=record.agent_name, - recovery_granularity=record.recovery_granularity, - runtime_prompt_payload=record.runtime_prompt_payload, - expected_generations=record.expected_generations, - target_step=record.target_step, - start_weight_version=record.start_weight_version, + return dataclasses.replace( + record, siblings=copy.deepcopy(record.siblings), - phase=record.phase, - status=record.status, ) @staticmethod diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index fc3174a2e9b..29440a2bfa0 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -2005,3 +2005,7 @@ def test_prompt_group_restore_redispatches_every_sibling(self): assert request is not None assert request.prompt_idx == 9 assert restored._impl.seen_generation_indices == [0, 1] + assert ( + restored._impl.seen_recovery_granularity + is RecoveryGranularity.PROMPT_GROUP + ) diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 5303b9dfff7..1d445a5d00a 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -15,6 +15,7 @@ from __future__ import annotations import asyncio +import dataclasses from collections.abc import Callable from typing import Any, TypeVar @@ -30,10 +31,18 @@ from nemo_rl.data.interfaces import DatumSpec from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, + _ATTEMPT_STATE_FIELDS, + _GROUP_STATE_FIELDS, + _PROMPT_REF_STATE_FIELDS, + _SIBLING_STATE_FIELDS, PromptGroupPhase, + PromptGroupRecoveryRecord, + PromptRef, RecoveryGranularity, + RolloutAttemptRecord, RolloutAttemptStatus, RolloutRecoveryLedger, + RolloutSiblingRecord, SiblingSealResult, build_rollout_recovery_state, parse_rollout_recovery_state, @@ -127,6 +136,22 @@ def test_ledger_round_trip_preserves_group_ownership() -> None: assert restored.get_group("g7").phase is PromptGroupPhase.ADMITTED +def test_serialized_state_fields_match_recovery_dataclasses() -> None: + """Require every durable dataclass field to be classified explicitly.""" + assert _PROMPT_REF_STATE_FIELDS == { + field.name for field in dataclasses.fields(PromptRef) + } + assert _SIBLING_STATE_FIELDS == { + field.name for field in dataclasses.fields(RolloutSiblingRecord) + } + assert _ATTEMPT_STATE_FIELDS == { + field.name for field in dataclasses.fields(RolloutAttemptRecord) + } + assert _GROUP_STATE_FIELDS == { + field.name for field in dataclasses.fields(PromptGroupRecoveryRecord) + } - {"runtime_prompt_payload"} + + @pytest.mark.parametrize( ("path", "context"), [ @@ -163,6 +188,126 @@ def test_ledger_restore_rejects_unknown_fields( RolloutRecoveryLedger.from_state_dict(state) +def _sealed_attempt_state() -> dict[str, Any]: + ledger = RolloutRecoveryLedger() + group = _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=1, + target_step=7, + start_weight_version=6, + admitted=True, + ) + _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) + gate_id = group.gate_rollout_id(0) + _mutate( + lambda cut: ledger.mark_sibling_sealed( + cut, + "g7", + generation_index=0, + gate_rollout_id=gate_id, + receipt={ + "rollout_id": gate_id, + "manifest": [{"staging_key": "g7/sibling-0/call-0"}], + }, + reward=1.0, + mask_sample=True, + ) + ) + return ledger.state_dict() + + +@pytest.mark.parametrize( + ("case", "error_fragment"), + [ + ("attempt_uuid", "attempt_uuid must contain exactly 16 bytes"), + ("status_type", "invalid rollout attempt status"), + ("status_value", "invalid rollout attempt status"), + ("staging_keys", "staging_keys must be a list of strings"), + ("reward", "sealed attempts require a reward"), + ("mask_sample", "sealed attempts require a boolean mask_sample"), + ( + "missing_receipt_staging", + "sealed missing-receipt attempt cannot own staging keys", + ), + ("receipt_type", "sealed attempt receipt must be a mapping or None"), + ("receipt_manifest_type", "receipt must contain a manifest list"), + ("receipt_identity", "sealed receipt identity mismatch"), + ("receipt_manifest", "sealed receipt staging manifest mismatch"), + ("unsealed_payload", "only sealed attempts may retain receipt data"), + ], +) +def test_restore_rejects_malformed_attempt_fields( + case: str, error_fragment: str +) -> None: + state = _sealed_attempt_state() + attempt = state["groups"][0]["siblings"][0]["attempts"][0] + + if case == "attempt_uuid": + attempt["attempt_uuid"] = b"short" + elif case == "status_type": + attempt["status"] = None + elif case == "status_value": + attempt["status"] = "unknown" + elif case == "staging_keys": + attempt["staging_keys"] = ["valid", 7] + elif case == "reward": + attempt["reward"] = None + elif case == "mask_sample": + attempt["mask_sample"] = "yes" + elif case == "missing_receipt_staging": + attempt["receipt"] = None + elif case == "receipt_type": + attempt["receipt"] = [] + elif case == "receipt_manifest_type": + attempt["receipt"]["manifest"] = None + elif case == "receipt_identity": + attempt["receipt"]["rollout_id"] = "wrong" + elif case == "receipt_manifest": + attempt["receipt"]["manifest"] = [{"staging_key": "wrong"}] + elif case == "unsealed_payload": + attempt["status"] = RolloutAttemptStatus.DISPATCHED.value + else: # pragma: no cover - the parameter table above owns the cases. + raise AssertionError(f"unknown malformed-attempt case={case!r}") + + with pytest.raises(ValueError, match=error_fragment): + RolloutRecoveryLedger.from_state_dict(state) + + +def test_restore_rejects_non_mapping_attempt() -> None: + state = _sealed_attempt_state() + state["groups"][0]["siblings"][0]["attempts"][0] = None + + with pytest.raises(ValueError, match="rollout-recovery attempt must be a mapping"): + RolloutRecoveryLedger.from_state_dict(state) + + +def test_restore_rejects_duplicate_attempt_identity() -> None: + ledger = RolloutRecoveryLedger() + _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=6, + admitted=True, + ) + state = ledger.state_dict() + siblings = state["groups"][0]["siblings"] + siblings[1]["attempts"][0]["attempt_uuid"] = siblings[0]["attempts"][0][ + "attempt_uuid" + ] + + with pytest.raises(ValueError, match="duplicate rollout attempt identity"): + RolloutRecoveryLedger.from_state_dict(state) + + def test_checkpoint_state_round_trip_preserves_controller_and_ledger_state() -> None: ledger = RolloutRecoveryLedger() _reserve( From 6549009a041a0cc681561c7641edbe416fffa3d4 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 5 Sep 2026 18:30:44 -0700 Subject: [PATCH 23/28] fix: lint errors Signed-off-by: Anish Mahishi --- nemo_rl/experience/rollout_manager.py | 14 ++++++++------ nemo_rl/experience/rollout_recovery.py | 4 +++- tests/unit/experience/test_rollout_manager.py | 3 +-- tests/unit/experience/test_rollout_recovery.py | 2 +- .../single_controller/test_finalizer_lifecycle.py | 2 ++ 5 files changed, 15 insertions(+), 10 deletions(-) diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index a8acf1ac17e..7d68361b87e 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -2045,9 +2045,9 @@ async def _generate_for_finalization_attempt( attempt_extra_env_info = attempt_input_sample.get("extra_env_info") if isinstance(attempt_extra_env_info, dict): attempt_extra_env_info[NEMO_GYM_GROUP_ID_KEY] = group_id - attempt_extra_env_info[NEMO_GYM_GROUP_ATTEMPT_KEY] = max( - len(sibling.attempts) for sibling in recovery_group.siblings - ) - 1 + attempt_extra_env_info[NEMO_GYM_GROUP_ATTEMPT_KEY] = ( + max(len(sibling.attempts) for sibling in recovery_group.siblings) - 1 + ) self._tq_buffer.reserve( weight_version=start_version, target_step=recovery_group.target_step, @@ -2095,9 +2095,11 @@ async def _record_streamed_completion( f"expected={gate_rollout_id!r}" ) mask_sample = bool( - (((completion.env_extras or {}).get("instance_config") or {}).get( - MASK_SAMPLE, False - )) + ( + ((completion.env_extras or {}).get("instance_config") or {}).get( + MASK_SAMPLE, False + ) + ) ) if recovery_group.recovery_granularity is RecoveryGranularity.PROMPT_GROUP: diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index b0473d40ca3..cbbd7b9ee43 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -1062,7 +1062,9 @@ def _group_from_state( if not isinstance(reward, (int, float)): raise ValueError("sealed attempts require a reward") if not isinstance(mask_sample, bool): - raise ValueError("sealed attempts require a boolean mask_sample") + raise ValueError( + "sealed attempts require a boolean mask_sample" + ) if receipt is None: if staging_keys: raise ValueError( diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 29440a2bfa0..e83dc9d3a65 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -2006,6 +2006,5 @@ def test_prompt_group_restore_redispatches_every_sibling(self): assert request.prompt_idx == 9 assert restored._impl.seen_generation_indices == [0, 1] assert ( - restored._impl.seen_recovery_granularity - is RecoveryGranularity.PROMPT_GROUP + restored._impl.seen_recovery_granularity is RecoveryGranularity.PROMPT_GROUP ) diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 1d445a5d00a..4563a500c00 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -30,11 +30,11 @@ from nemo_rl.algorithms.single_controller_utils.config import RolloutRecoveryConfig from nemo_rl.data.interfaces import DatumSpec from nemo_rl.experience.rollout_recovery import ( - ROLLOUT_RECOVERY_SCHEMA_VERSION, _ATTEMPT_STATE_FIELDS, _GROUP_STATE_FIELDS, _PROMPT_REF_STATE_FIELDS, _SIBLING_STATE_FIELDS, + ROLLOUT_RECOVERY_SCHEMA_VERSION, PromptGroupPhase, PromptGroupRecoveryRecord, PromptRef, diff --git a/tests/unit/single_controller/test_finalizer_lifecycle.py b/tests/unit/single_controller/test_finalizer_lifecycle.py index c8fa52cea84..9cbfb2f661b 100644 --- a/tests/unit/single_controller/test_finalizer_lifecycle.py +++ b/tests/unit/single_controller/test_finalizer_lifecycle.py @@ -272,6 +272,8 @@ def test_post_train_cleanup_clears_canonical_rows_and_route_plan_staging_keys() "partition_id": "staging", }, ] + + class _SyncDataPlaneClient: """Synchronous client like the production TQ adapter; records caller threads.""" From fa7f2bd7a4322dda377f3c01489afbf15456b53b Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 6 Sep 2026 11:36:27 -0400 Subject: [PATCH 24/28] fix(rollout): harden sibling recovery invariants Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 4 + nemo_rl/experience/rollout_recovery.py | 4 +- .../unit/experience/test_rollout_recovery.py | 3 + .../_checkpoint_scenarios.py | 116 +++++++++++++++++- .../test_checkpoint_recovery_matrix.py | 33 ++++- .../single_controller/test_checkpointing.py | 1 + 6 files changed, 152 insertions(+), 9 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index dabc88c978a..af55ab9b493 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -3921,6 +3921,10 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: fields_to_put[adv_cfg.returns_field] = returns new_fields.append(adv_cfg.returns_field) + # Trainer-step checkpointing runs later in this same train-pump task, so + # this publication cannot race a checkpoint save. If advantage staging + # moves to another task, the write must participate in the data-plane + # mutation barrier. await self._call_dp( "put_samples", sample_ids=meta.sample_ids, diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index cbbd7b9ee43..b9f3e63f7a2 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -281,7 +281,7 @@ class SiblingSealResult: # into a masked placeholder, matching the base token-capture contract. receipt: Optional[dict[str, Any]] reward: float - mask_sample: bool = False + mask_sample: bool def _new_attempt() -> RolloutAttemptRecord: @@ -577,7 +577,7 @@ def mark_sibling_sealed( gate_rollout_id: str, receipt: Optional[dict[str, Any]], reward: float, - mask_sample: bool = False, + mask_sample: bool, ) -> None: """Record one streamed sibling receipt as soon as the row arrives.""" cut.require_live() diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 4563a500c00..f133126c61f 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -871,6 +871,7 @@ def test_prompt_group_restart_keeps_a_fully_sealed_group() -> None: "manifest": [{"staging_key": f"g7/sibling-{generation_index}/call-0"}], }, reward=1.0, + mask_sample=False, ) _mutate(lambda cut: ledger.mark_group_sealed(cut, "g7", results)) @@ -914,6 +915,7 @@ def test_prompt_group_seal_is_atomic() -> None: "manifest": [{"staging_key": "g7/sibling-0/call-0"}], }, reward=1.0, + mask_sample=False, ) } @@ -957,6 +959,7 @@ def test_checkpoint_rejects_ambiguous_finalization_state( "manifest": [{"staging_key": "g7/sibling-0/call-0"}], }, reward=1.0, + mask_sample=False, ) ) _mutate(lambda cut: ledger.mark_finalization_started(cut, "g7")) diff --git a/tests/unit/single_controller/_checkpoint_scenarios.py b/tests/unit/single_controller/_checkpoint_scenarios.py index 2fbe41696a6..17b0a4816c9 100644 --- a/tests/unit/single_controller/_checkpoint_scenarios.py +++ b/tests/unit/single_controller/_checkpoint_scenarios.py @@ -40,6 +40,7 @@ from typing import Any, Optional import torch +from tensordict import TensorDict from nemo_rl.algorithms.async_utils import replay_buffer as _rb from nemo_rl.algorithms.async_utils.replay_buffer import ( @@ -56,9 +57,13 @@ from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.interfaces import PromptGroupRecord -from nemo_rl.experience.rollout_recovery import RolloutRecoveryLedger +from nemo_rl.experience.rollout_recovery import ( + RolloutAttemptStatus, + RolloutRecoveryLedger, +) PARTITION = "rollout_data" +STAGING_PARTITION = "rollout_staging" ROLLOUTS_PER_GROUP = 2 # rollouts_per_prompt_group GROUPS_PER_STEP = 3 # prompt_groups per training step CAPACITY = 64 # max_buffered_rollouts @@ -91,7 +96,8 @@ class Group: gid: Prompt-group number, matching the order the dataloader served it. done: How many of its ``ROLLOUTS_PER_GROUP`` rollouts have finished. ``ROLLOUTS_PER_GROUP`` means the group committed; anything less - means it is still in flight. + means it is still in flight. Finished siblings of an incomplete + group are sealed in the recovery ledger and backed by staging rows. weight: Weight version the group was dispatched at. target: ``target_step`` stamp, used by the gated samplers. evicted: The sampler deliberately dropped it (too stale). An evicted @@ -164,6 +170,10 @@ def _gid(n: int) -> str: return f"g{n:02d}" +def _staging_key(group_id: str, generation_index: int) -> str: + return f"{group_id}/sibling-{generation_index}/call-0" + + @dataclass(frozen=True) class Case: """One row of the test matrix: a scenario run under one sampler. @@ -227,6 +237,12 @@ def _fresh_client(register: bool) -> NoOpDataPlaneClient: num_samples=CAPACITY * ROLLOUTS_PER_GROUP, consumer_tasks=["train"], ) + dp.register_partition( + partition_id=STAGING_PARTITION, + fields=["token_ids"], + num_samples=CAPACITY * ROLLOUTS_PER_GROUP, + consumer_tasks=[], + ) return dp @@ -279,7 +295,10 @@ class RoundTrip: nothing asserts on them. ``stamps`` records each restored group's ``target_step`` and start weight. ``selected`` and ``selected_count`` report the optional restore-then-select result used to verify each sampler's - recovery key at multiple gate lags. + recovery key at multiple gate lags. The sealed-sibling and redispatch maps + verify that an unfinished group keeps completed work and retries only its + missing generation indices. The staging-row sets verify that the matching + token-capture payload survived the data-plane checkpoint. """ recovered: set[str] @@ -291,6 +310,11 @@ class RoundTrip: stamps: dict[str, tuple[int | None, int]] selected: set[str] selected_count: int + sealed_before: dict[str, tuple[int, ...]] + sealed_after: dict[str, tuple[int, ...]] + redispatched: dict[str, tuple[int, ...]] + staging_rows_before: set[str] + staging_rows_after_restore: set[str] async def _round_trip( @@ -321,7 +345,7 @@ async def _round_trip( or group.done == ROLLOUTS_PER_GROUP ): continue - recovery_ledger_a.reserve_group( + recovery_group = recovery_ledger_a.reserve_group( cut, group_id=_gid(group.gid), admission_id=f"batch-{group.target}", @@ -332,8 +356,37 @@ async def _round_trip( start_weight_version=group.weight, admitted=True, ) + recovery_ledger_a.mark_group_dispatched(cut, recovery_group.group_id) + for generation_index in range(group.done): + gate_rollout_id = recovery_group.gate_rollout_id(generation_index) + staging_key = _staging_key(recovery_group.group_id, generation_index) + dp_a.put_samples( + sample_ids=[staging_key], + partition_id=STAGING_PARTITION, + fields=TensorDict( + {"token_ids": torch.tensor([[generation_index]])}, + batch_size=[1], + ), + ) + recovery_ledger_a.mark_sibling_sealed( + cut, + recovery_group.group_id, + generation_index=generation_index, + gate_rollout_id=gate_rollout_id, + receipt={ + "rollout_id": gate_rollout_id, + "manifest": [{"staging_key": staging_key}], + }, + reward=float(generation_index), + mask_sample=False, + ) + sealed_before = { + group.group_id: tuple(group.sealed_generation_indices) + for group in recovery_ledger_a.groups() + } recovery_sidecar = recovery_ledger_a.state_dict() rows_before = set(dp_a.list_sample_ids(PARTITION)) + staging_rows_before = set(dp_a.list_sample_ids(STAGING_PARTITION)) dp_a.save_checkpoint(tmp_path / "data_plane") # ---- restart: brand new process, nothing carried over in memory ---- @@ -354,8 +407,30 @@ async def _round_trip( recovery_ledger_b = RolloutRecoveryLedger() async with buf_b.data_plane_checkpoint_barrier.mutation() as cut: recovery_ledger_b.load_state_dict(cut, recovery_sidecar) + recovery_ledger_b.prepare_for_restart(cut) recovery_ledger_b.discard_canonical_groups(cut, set(buf_b._group_ids)) + staging_rows_after_restore = set(dp_b.list_sample_ids(STAGING_PARTITION)) + sealed_after = { + group.group_id: tuple(group.sealed_generation_indices) + for group in recovery_ledger_b.groups() + } + redispatched: dict[str, tuple[int, ...]] = {} for group in recovery_ledger_b.groups(): + async with buf_b.data_plane_checkpoint_barrier.mutation() as cut: + retry_group = recovery_ledger_b.prepare_incomplete_retry( + cut, group.group_id + ) + generation_indices = tuple( + sibling.generation_index + for sibling in retry_group.siblings + if sibling.current_attempt.status is RolloutAttemptStatus.RESERVED + ) + recovery_ledger_b.mark_group_dispatched( + cut, + group.group_id, + generation_indices=list(generation_indices), + ) + redispatched[group.group_id] = generation_indices group_id = buf_b.reserve( weight_version=group.start_weight_version, target_step=group.target_step, @@ -369,6 +444,10 @@ async def _round_trip( ) async with buf_b.data_plane_checkpoint_barrier.mutation() as cut: recovery_ledger_b.discard_group(cut, group_id) + else: + sealed_after = {} + redispatched = {} + staging_rows_after_restore = set(dp_b.list_sample_ids(STAGING_PARTITION)) ready = { gid for gid, is_ready in zip(buf_b._group_ids, buf_b.ready_list) if is_ready @@ -405,6 +484,11 @@ async def _round_trip( stamps=stamps, selected=selected, selected_count=selected_count, + sealed_before=sealed_before, + sealed_after=sealed_after, + redispatched=redispatched, + staging_rows_before=staging_rows_before, + staging_rows_after_restore=staging_rows_after_restore, ) @@ -506,6 +590,21 @@ def assert_completed_groups_survive( lag=1, ) +S_ZERO_LAG_PARTIAL = Scenario( + name="lag0-current-step-partly-generated", + groups=( + Group(9, 2, weight=4, target=4), + Group(10, 2, weight=4, target=4), + Group(11, 2, weight=4, target=4), + Group(12, 1, weight=5, target=5), + Group(13, 2, weight=5, target=5), + Group(14, 0, weight=5, target=5), + ), + cursor=15, + trained=frozenset({9, 10, 11}), + lag=0, +) + S_LAG2 = Scenario( name="lag2-two-batches-in-flight", groups=( @@ -568,5 +667,12 @@ def assert_completed_groups_survive( # Everything fully generated -- the case this PR set out to recover. FULLY_GENERATED = (S_ZERO_LAG_ALL_COMPLETE, S_ALL_COMPLETE, S_STALE_ONLY) # At least one group still generating when the snapshot was taken. -WITH_IN_FLIGHT = (S_PARTIAL, S_LAG2, S_EVICTED, S_TRAINED_OUT_OF_ORDER) +WITH_IN_FLIGHT = ( + S_ZERO_LAG_PARTIAL, + S_PARTIAL, + S_LAG2, + S_EVICTED, + S_TRAINED_OUT_OF_ORDER, +) +WITH_SEALED_SIBLINGS = (S_ZERO_LAG_PARTIAL, S_PARTIAL, S_LAG2) ALL_SCENARIOS = FULLY_GENERATED + WITH_IN_FLIGHT diff --git a/tests/unit/single_controller/test_checkpoint_recovery_matrix.py b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py index c99e9cb15a4..fd04cc48933 100644 --- a/tests/unit/single_controller/test_checkpoint_recovery_matrix.py +++ b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py @@ -18,8 +18,9 @@ in_order; a zero-lag completed row and ready_first are included here to cover the complete built-in recovery contract. -Unfinished rows are regenerated as whole prompt groups from the group-level -ledger. Sibling-level continuation remains outside this recovery foundation. +Unfinished prompt groups are recovered from the rollout ledger. Under sibling +recovery, completed siblings retain their captured staging rows and only their +missing siblings are redispatched. """ from __future__ import annotations @@ -30,11 +31,13 @@ ALL_SCENARIOS, FULLY_GENERATED, GROUPS_PER_STEP, + ROLLOUTS_PER_GROUP, S_ALL_COMPLETE, S_LAG2, S_ZERO_LAG_ALL_COMPLETE, SAMPLERS, WITH_IN_FLIGHT, + WITH_SEALED_SIBLINGS, Case, assert_completed_groups_survive, assert_no_data_loss, @@ -56,6 +59,9 @@ for sampler in SAMPLERS for scenario in (S_ZERO_LAG_ALL_COMPLETE, S_ALL_COMPLETE, S_LAG2) ] +SEALED_SIBLING_CASES = [ + Case(scenario, sampler) for sampler in SAMPLERS for scenario in WITH_SEALED_SIBLINGS +] @pytest.fixture(autouse=True) @@ -86,6 +92,29 @@ def test_unfinished_groups_are_owned_across_restart(case, tmp_path): assert_no_data_loss(case.scenario, case.sampler, tmp_path) +@pytest.mark.parametrize("case", SEALED_SIBLING_CASES, ids=lambda case: case.id) +def test_sealed_siblings_survive_and_only_missing_siblings_redispatch(case, tmp_path): + """Preserve completed sibling work across samplers and gate lags 0, 1, and 2.""" + result = round_trip(case.scenario, case.sampler, tmp_path) + partial_groups = { + f"g{group.gid:02d}": group + for group in case.scenario.groups + if 0 < group.done < ROLLOUTS_PER_GROUP + and not group.evicted + and group.gid not in case.scenario.trained + } + + assert partial_groups + for group_id, group in partial_groups.items(): + expected_sealed = tuple(range(group.done)) + expected_missing = tuple(range(group.done, ROLLOUTS_PER_GROUP)) + assert result.sealed_before[group_id] == expected_sealed + assert result.sealed_after[group_id] == expected_sealed + assert result.redispatched[group_id] == expected_missing + assert result.staging_rows_before + assert result.staging_rows_after_restore == result.staging_rows_before + + @pytest.mark.parametrize("case", ALL_CASES, ids=lambda case: case.id) def test_restore_preserves_sampler_stamps(case, tmp_path): """Every restored group retains the keys its sampler uses for selection.""" diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 64546e2433c..a813610185f 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -675,6 +675,7 @@ async def seed() -> None: "manifest": [{"staging_key": staging_key}], }, reward=1.0, + mask_sample=False, ) asyncio.run(seed()) From efcf43b48b53c4049783a18ddda6a75d24b045b0 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 6 Sep 2026 11:50:53 -0400 Subject: [PATCH 25/28] refactor(rollout): rename recovery granularity result Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller_utils/config.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index e6ab7ad0e91..79559e4e078 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -623,7 +623,7 @@ class TokenCaptureConfig(BaseModel, extra="allow"): @dataclass(frozen=True) -class RecoveryGranularityResolution: +class AgentRecoveryGranularity: """Recovery granularity selected for a prompt-group reservation. ``agent_name`` is copied from the prompt when present. ``granularity`` is @@ -660,9 +660,7 @@ class RolloutRecoveryConfig(BaseModel, extra="allow"): default_factory=dict ) - def resolve_for_prompt( - self, prompt: Mapping[str, Any] - ) -> RecoveryGranularityResolution: + def resolve_for_prompt(self, prompt: Mapping[str, Any]) -> AgentRecoveryGranularity: """Resolve one new group using agent, then task, then the global default.""" extra_env_info = prompt.get("extra_env_info") agent_name: Optional[str] = None @@ -678,7 +676,7 @@ def resolve_for_prompt( if agent_name is not None: override = self.agent_granularity_overrides.get(agent_name) if override is not None: - return RecoveryGranularityResolution(agent_name, override) + return AgentRecoveryGranularity(agent_name, override) task_name = prompt.get("task_name") if task_name is not None and not isinstance(task_name, str): @@ -686,8 +684,8 @@ def resolve_for_prompt( if task_name is not None: override = self.task_granularity_overrides.get(task_name) if override is not None: - return RecoveryGranularityResolution(agent_name, override) - return RecoveryGranularityResolution(agent_name, self.default_granularity) + return AgentRecoveryGranularity(agent_name, override) + return AgentRecoveryGranularity(agent_name, self.default_granularity) class MasterConfig(BaseModel, extra="allow"): From 34bcf9b36480ef47d2a9f697e22fa6328cb4156c Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 6 Sep 2026 14:06:53 -0400 Subject: [PATCH 26/28] fix(rollout): align recovery policy with Gym routing Signed-off-by: Anish Mahishi --- docs/guides/single-controller.md | 23 ++-- ...po_math_1B_megatron_single_controller.yaml | 17 +-- ...po_math_1B_megatron_single_controller.yaml | 17 +-- .../single_controller_utils/config.py | 64 +++++++--- nemo_rl/environments/nemo_gym.py | 37 ++++-- nemo_rl/experience/rollout_manager.py | 20 +++- nemo_rl/experience/rollout_recovery.py | 20 ++-- .../grpo_async_gym_single_controller.sh | 10 +- tests/unit/environments/test_nemo_gym.py | 86 +++++++++++++- tests/unit/experience/test_rollout_manager.py | 37 +++++- .../unit/experience/test_rollout_recovery.py | 109 +++++++++++------- .../test_checkpoint_dispatch_races.py | 2 +- tests/unit/single_controller/test_setup.py | 12 +- 13 files changed, 331 insertions(+), 123 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 751dd9385b5..30262c4abf5 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -115,14 +115,21 @@ failure and restart behavior: `prompt_group` when every generation in a recovered group must come from the policy weights live at redispatch. -`agent_granularity_overrides` and `task_granularity_overrides` can select the -policy per Gym agent or dataset task; agent overrides take precedence. These -non-default policies require `token_capture.enabled: true`. The resolved policy -is 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 sample from the same prompt. -Neither becomes a training row until every generation in the group has finished. +`task_source_granularity_overrides` can select the policy 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, +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 +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 +sample from the same prompt. Neither becomes a training row until every +generation in the group has finished. ::: When a sampler does not support replay recovery, a requested data-plane checkpoint is written in `shadow` mode. The TQ snapshot is retained, but no authoritative replay index is written and its rows are not restored into the training replay buffer. diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index 3689a031b04..96f2d1ea69d 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -182,20 +182,21 @@ token_capture: # 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. Agent overrides win -# over task overrides. +# sibling fails in-process or was unfinished at a checkpoint. rollout_recovery: default_granularity: sibling - # Per-Gym-agent override, keyed on extra_env_info.agent_ref.name. Wins over - # the task map below. + # Per-Gym-task-source override, keyed on extra_env_info.task_source. + # task_source_granularity_overrides: + # genrm_compare_resources_server: prompt_group + task_source_granularity_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: {} - # Per-dataset-task override, keyed on the prompt's task_name. - # task_granularity_overrides: {math: prompt_group} - task_granularity_overrides: {} # Leaving either map above non-empty requires the top-level token_capture section -# above to set enabled: true. +# to set enabled: true. cluster: # Master ports inherit the shared 1400-1999 band from grpo_math_1B.yaml. diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 62628dd67eb..0c414a80216 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -220,20 +220,21 @@ token_capture: # 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. Agent overrides win -# over task overrides. +# sibling fails in-process or was unfinished at a checkpoint. rollout_recovery: default_granularity: sibling - # Per-Gym-agent override, keyed on extra_env_info.agent_ref.name. Wins over - # the task map below. + # Per-Gym-task-source override, keyed on extra_env_info.task_source. + # task_source_granularity_overrides: + # genrm_compare_resources_server: prompt_group + task_source_granularity_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: {} - # Per-dataset-task override, keyed on the prompt's task_name. - # task_granularity_overrides: {math: prompt_group} - task_granularity_overrides: {} # Leaving either map above non-empty requires the top-level token_capture section -# above to set enabled: true. +# to set enabled: true. cluster: # Master ports inherit the shared 1400-1999 band from ppo_math_1B.yaml. diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 79559e4e078..72419d0e009 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -623,14 +623,15 @@ class TokenCaptureConfig(BaseModel, extra="allow"): @dataclass(frozen=True) -class AgentRecoveryGranularity: +class TaskSourceRecoveryGranularity: """Recovery granularity selected for a prompt-group reservation. - ``agent_name`` is copied from the prompt when present. ``granularity`` is - selected from an agent override, task override, or the global default. + ``task_source`` is copied from the raw Gym row when present. ``granularity`` + is selected from an explicit agent override, a task-source override, or the + global default. """ - agent_name: Optional[str] + task_source: Optional[str] granularity: RecoveryGranularity @@ -652,19 +653,40 @@ class RolloutRecoveryConfig(BaseModel, extra="allow"): """ default_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING - # NeMo-Gym agent_ref.name takes precedence over task_name when both match. - agent_granularity_overrides: dict[str, RecoveryGranularity] = Field( + # Keyed by ``extra_env_info.task_source``, which is available before Gym + # resolves the concrete agent used to execute the row. + task_source_granularity_overrides: dict[str, RecoveryGranularity] = Field( default_factory=dict ) - task_granularity_overrides: dict[str, RecoveryGranularity] = Field( + # Keyed by ``extra_env_info.agent_ref.name`` when the input row already has + # a concrete Gym route. A matching agent override wins over task_source. + agent_granularity_overrides: dict[str, RecoveryGranularity] = Field( default_factory=dict ) - def resolve_for_prompt(self, prompt: Mapping[str, Any]) -> AgentRecoveryGranularity: - """Resolve one new group using agent, then task, then the global default.""" + @model_validator(mode="after") + def _reject_removed_override_keys(self) -> "RolloutRecoveryConfig": + """Reject the removed task-name map instead of silently ignoring it.""" + removed = {"task_granularity_overrides"}.intersection(self.model_extra or {}) + if removed: + raise ValueError( + f"rollout_recovery fields {sorted(removed)!r} were replaced by " + "task_source_granularity_overrides" + ) + return self + + def resolve_for_prompt( + self, prompt: Mapping[str, Any] + ) -> TaskSourceRecoveryGranularity: + """Resolve using matching agent, matching task source, then default.""" extra_env_info = prompt.get("extra_env_info") + task_source: Optional[str] = None agent_name: Optional[str] = None if isinstance(extra_env_info, Mapping): + raw_task_source = extra_env_info.get("task_source") + if raw_task_source is not None and not isinstance(raw_task_source, str): + raise TypeError("prompt task_source must be a string or None") + task_source = raw_task_source agent_ref = extra_env_info.get("agent_ref") if agent_ref is not None and not isinstance(agent_ref, Mapping): raise TypeError("prompt agent_ref must be a mapping or None") @@ -674,18 +696,22 @@ def resolve_for_prompt(self, prompt: Mapping[str, Any]) -> AgentRecoveryGranular raise TypeError("prompt agent_ref.name must be a string or None") agent_name = raw_agent_name if agent_name is not None: + if task_source is None: + warnings.warn( + "rollout recovery is using legacy agent_ref because " + "task_source is missing; re-collate the dataset with " + "the current NeMo Gym", + FutureWarning, + stacklevel=2, + ) override = self.agent_granularity_overrides.get(agent_name) if override is not None: - return AgentRecoveryGranularity(agent_name, override) - - task_name = prompt.get("task_name") - if task_name is not None and not isinstance(task_name, str): - raise TypeError("prompt task_name must be a string or None") - if task_name is not None: - override = self.task_granularity_overrides.get(task_name) + return TaskSourceRecoveryGranularity(task_source, override) + if task_source is not None: + override = self.task_source_granularity_overrides.get(task_source) if override is not None: - return AgentRecoveryGranularity(agent_name, override) - return AgentRecoveryGranularity(agent_name, self.default_granularity) + return TaskSourceRecoveryGranularity(task_source, override) + return TaskSourceRecoveryGranularity(task_source, self.default_granularity) class MasterConfig(BaseModel, extra="allow"): @@ -1123,8 +1149,8 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: recovery_config = master_config.rollout_recovery if not token_capture_config.enabled and ( recovery_config.default_granularity is not RecoveryGranularity.SIBLING + or recovery_config.task_source_granularity_overrides or recovery_config.agent_granularity_overrides - or recovery_config.task_granularity_overrides ): raise ValueError( "non-default rollout_recovery policies require " diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 7415aaa985c..fc63620057b 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -78,6 +78,20 @@ DEFAULT_THINKING_TAGS = ["", ""] +def _rollout_progress_identity(row: Mapping[str, Any]) -> str: + """Return the concrete agent route, falling back to task provenance.""" + agent_ref = row.get("agent_ref") + if isinstance(agent_ref, Mapping): + agent_name = agent_ref.get("name") + if isinstance(agent_name, str) and agent_name: + return f"agent:{agent_name}" + + task_source = row.get("task_source") + if isinstance(task_source, str) and task_source: + return f"task-source:{task_source}" + return "" + + class NemoGymCompatibleConfig(Protocol): """Configuration fields required to select the NeMo Gym rollout path.""" @@ -615,10 +629,11 @@ async def run_rollouts( nemo_gym_result_iterator = self.rch.run_examples( examples=nemo_gym_examples, head_server_config=self.head_server_config ) - # Current Gym collates data with ``task_source`` rather than a baked-in - # ``agent_ref``. ``run_examples`` resolves that routing synchronously and - # stamps each input row before returning its result iterator. - counts_left = Counter(row["agent_ref"]["name"] for row in nemo_gym_examples) + # Gym resolves task_source to agent_ref synchronously in run_examples(). + # Build the counter afterward so completion rows use the same identity. + routing_counts_left = Counter( + _rollout_progress_identity(row) for row in nemo_gym_examples + ) num_results = 0 for task in nemo_gym_result_iterator: @@ -670,18 +685,18 @@ async def run_rollouts( / total_time ) - agent_name = nemo_gym_row["agent_ref"]["name"] - counts_left[agent_name] -= 1 - if counts_left[agent_name] <= 0: - counts_left.pop(agent_name) - if num_results % 10 == 0 and counts_left: - top_left = counts_left.most_common(5) + routing_identity = _rollout_progress_identity(nemo_gym_row) + routing_counts_left[routing_identity] -= 1 + if routing_counts_left[routing_identity] <= 0: + routing_counts_left.pop(routing_identity) + if num_results % 10 == 0 and routing_counts_left: + top_left = routing_counts_left.most_common(5) top_left_str = "\n".join( f"{index + 1}. {name}: {count}" for index, (name, count) in enumerate(top_left) ) print( - "Top 5 NeMo Gym agent refs left in this rollout batch: " + "Top 5 NeMo Gym routing identities left in this rollout batch: " f"{top_left_str}", file=sys.stderr, ) diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 7d68361b87e..a368acbc68e 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -19,7 +19,7 @@ import enum import json import uuid -from collections.abc import AsyncIterator, Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from contextlib import asynccontextmanager from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Optional @@ -106,6 +106,20 @@ def _contains_post_write_enrichment_error(error: BaseException) -> bool: return False +def _nemo_gym_metric_namespace(row: Mapping[str, Any]) -> str: + """Return the best available namespace for NeMo-Gym rollout metrics.""" + agent_ref = row.get("agent_ref") + if isinstance(agent_ref, Mapping): + agent_name = agent_ref.get("name") + if isinstance(agent_name, str) and agent_name: + return agent_name + + task_source = row.get("task_source") + if isinstance(task_source, str) and task_source: + return f"task-source:{task_source}" + return "nemo_gym" + + class RolloutOutcome(str, enum.Enum): """How :meth:`RolloutManager.generate_and_push` finished for one prompt.""" @@ -1233,7 +1247,7 @@ async def _run_rollouts( # Compute rollout metrics. with timer.time(f"{timer_prefix}/compute_metrics"): rollout_metrics = self._compute_rollout_metrics( - completions, inputs[0]["agent_ref"]["name"] + completions, _nemo_gym_metric_namespace(inputs[0]) ) # Same helper the batched path uses, so the two cannot drift apart. rollout_metrics.update(_effort_shaping_metrics(shaping)) @@ -1591,7 +1605,7 @@ def reserve_prompt_group( expected_generations=self._num_generations_per_prompt, target_step=target_step, start_weight_version=self._weight_version, - agent_name=recovery_policy.agent_name, + task_source=recovery_policy.task_source, recovery_granularity=recovery_policy.granularity, admitted=admitted, admission_id=admission_id, diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index b9f3e63f7a2..0694652606d 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -34,7 +34,7 @@ from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneMutationCut from nemo_rl.data.interfaces import DatumSpec -ROLLOUT_RECOVERY_SCHEMA_VERSION = 5 +ROLLOUT_RECOVERY_SCHEMA_VERSION = 2 _SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS = {ROLLOUT_RECOVERY_SCHEMA_VERSION} ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" RolloutRecoveryState: TypeAlias = dict[str, Any] @@ -53,7 +53,7 @@ "admission_id", "prompt_id", "prompt_ref", - "agent_name", + "task_source", "recovery_granularity", "expected_generations", "target_step", @@ -209,7 +209,7 @@ class PromptGroupRecoveryRecord: admission_id: str prompt_id: str prompt_ref: PromptRef - agent_name: Optional[str] + task_source: Optional[str] recovery_granularity: RecoveryGranularity runtime_prompt_payload: Optional[DatumSpec] expected_generations: int @@ -331,7 +331,7 @@ def reserve_group( expected_generations: int, target_step: Optional[int], start_weight_version: int, - agent_name: Optional[str] = None, + task_source: Optional[str] = None, recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, admitted: bool = True, group_id: Optional[str] = None, @@ -374,7 +374,7 @@ def reserve_group( admission_id=admission_id, prompt_id=prompt_id, prompt_ref=prompt_ref, - agent_name=agent_name, + task_source=task_source, recovery_granularity=recovery_granularity, # Retain the immutable dataloader sample by reference instead of copying # a potentially 131k-token payload. This cache is never serialized and @@ -849,7 +849,7 @@ def state_dict(self) -> dict[str, Any]: "sample_id": record.prompt_ref.sample_id, "task_name": record.prompt_ref.task_name, }, - "agent_name": record.agent_name, + "task_source": record.task_source, "recovery_granularity": record.recovery_granularity.value, "expected_generations": record.expected_generations, "target_step": record.target_step, @@ -959,7 +959,7 @@ def _group_from_state( group_id = raw_group.get("group_id") admission_id = raw_group.get("admission_id") prompt_id = raw_group.get("prompt_id") - agent_name = raw_group.get("agent_name") + task_source = raw_group.get("task_source") raw_recovery_granularity = raw_group.get("recovery_granularity") expected_generations = raw_group.get("expected_generations") siblings_state = raw_group.get("siblings") @@ -969,8 +969,8 @@ def _group_from_state( raise ValueError("admission_id must be a non-empty string") if not isinstance(prompt_id, str) or not prompt_id: raise ValueError("prompt_id must be a non-empty string") - if agent_name is not None and not isinstance(agent_name, str): - raise ValueError("agent_name must be a string or None") + if task_source is not None and not isinstance(task_source, str): + raise ValueError("task_source must be a string or None") if not isinstance(raw_recovery_granularity, str): raise ValueError("recovery_granularity must be a string") try: @@ -1146,7 +1146,7 @@ def _group_from_state( admission_id=admission_id, prompt_id=prompt_id, prompt_ref=PromptRef(sample_id=sample_id, task_name=task_name), - agent_name=agent_name, + task_source=task_source, recovery_granularity=recovery_granularity, runtime_prompt_payload=None, expected_generations=expected_generations, diff --git a/tests/functional/grpo_async_gym_single_controller.sh b/tests/functional/grpo_async_gym_single_controller.sh index a6eb6d7f200..323cd96de95 100755 --- a/tests/functional/grpo_async_gym_single_controller.sh +++ b/tests/functional/grpo_async_gym_single_controller.sh @@ -53,14 +53,8 @@ cd - # smoke test, we trim all but the first tool TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl -jq -c ' - .agent_ref //= {"name": "workplace_assistant_simple_agent"} - | .responses_create_params.tools |= (.[0:1]) -' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH -jq -c ' - .agent_ref //= {"name": "workplace_assistant_simple_agent"} - | .responses_create_params.tools |= (.[0:1]) -' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ $SC_ENTRYPOINT \ diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index 0c57afd293f..47d5f900792 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -41,6 +41,7 @@ from nemo_rl.environments.nemo_gym import ( NemoGym, NemoGymConfig, + _rollout_progress_identity, build_reward_component_columns, extract_reward_components, setup_nemo_gym_config, @@ -77,6 +78,87 @@ ) +@pytest.mark.parametrize( + ("row", "expected"), + [ + ( + { + "task_source": "math_resources_server", + "agent_ref": {"name": "routed_agent"}, + }, + "agent:routed_agent", + ), + ({"task_source": "math_resources_server"}, "task-source:math_resources_server"), + ({"agent_ref": {"name": "legacy_agent"}}, "agent:legacy_agent"), + ({}, ""), + ], +) +def test_rollout_progress_identity_matches_gym_routing_precedence( + row: dict, expected: str +) -> None: + assert _rollout_progress_identity(row) == expected + + +def test_rollout_progress_counter_is_built_after_gym_resolves_task_source( + capsys, +) -> None: + async def _run() -> None: + rows = [ + { + "_rowidx": index, + "task_source": "test_resources_server", + "responses_create_params": {"input": []}, + } + for index in range(11) + ] + + class _RolloutCollectionHelper: + def run_examples(self, examples, head_server_config): + del head_server_config + for row in examples: + row["agent_ref"] = {"name": "resolved_agent"} + + async def _completed_result(row): + return row, {"response": {"output": []}} + + return [_completed_result(row) for row in examples] + + class _MockSelf: + cfg = {} + rch = _RolloutCollectionHelper() + head_server_config = object() + _token_capture_enabled = False + _tokenizer = object() + + def _require_spinup(self): + pass + + def _postprocess_nemo_gym_to_nemo_rl_result( + self, + row, + result, + result_tokenizer, + *, + include_initial_multimodal_data, + ): + del self, row, result, result_tokenizer, include_initial_multimodal_data + return {"message_log": []} + + streamed = [] + async for result in NemoGym.__ray_metadata__.modified_class.run_rollouts( + _MockSelf(), rows, "test" + ): + streamed.append(result) + + assert len(streamed) == len(rows) + + asyncio.run(_run()) + + captured = capsys.readouterr() + assert "1. agent:resolved_agent: 1" in captured.err + assert "task-source:test_resources_server" not in captured.err + + def test_multimodal_content_types_cover_responses_media_aliases(): assert { "input_image", @@ -1528,7 +1610,7 @@ def test_nemo_gym_run_rollouts_normalizes_mixed_media_before_dispatch(tmp_path): async def _run(): nemo_gym_row = { "_rowidx": 7, - "agent_ref": {"name": "test_agent"}, + "agent_ref": {"name": "legacy_test_agent"}, "responses_create_params": { "input": [ { @@ -1625,6 +1707,7 @@ async def _run(): row = { "_rowidx": 3, + "task_source": "test_resources_server", "agent_ref": {"name": "mock-megatron-agent"}, "responses_create_params": { "input": [ @@ -1801,6 +1884,7 @@ def test_nemo_gym_sanity( "temperature" ] example["responses_create_params"]["top_p"] = generation_config["top_p"] + example["task_source"] = "example_multi_step_resources_server" example["_rowidx"] = idx actual_result = [None] * len(nemo_gym_sanity_test_data["input"]) diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index e83dc9d3a65..5429548c780 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -60,6 +60,7 @@ RolloutOutcome, RolloutRetryPolicy, RolloutStats, + _nemo_gym_metric_namespace, ) from nemo_rl.experience.rollout_recovery import ( RecoveryGranularity, @@ -477,19 +478,19 @@ async def _assert_ledger_owns_inflight_prompt(_sample): assert buf._slots == [group_id] assert buf.commit_calls[0][0] == group_id - def test_reservation_persists_the_resolved_agent_recovery_policy(self): + def test_reservation_persists_the_resolved_task_source_recovery_policy(self): buf = _FakeBuffer() mgr = _make_manager(buf, _FakeImpl()) mgr._rollout_recovery_config = RolloutRecoveryConfig( - agent_granularity_overrides={ - "genrm_agent": RecoveryGranularity.PROMPT_GROUP + task_source_granularity_overrides={ + "genrm_compare": RecoveryGranularity.PROMPT_GROUP } ) prompt = { "idx": 0, "message_log": [], "task_name": "nemo_gym", - "extra_env_info": {"agent_ref": {"name": "genrm_agent"}}, + "extra_env_info": {"task_source": "genrm_compare"}, } group_id = _with_cut( @@ -498,7 +499,7 @@ def test_reservation_persists_the_resolved_agent_recovery_policy(self): ) group = mgr.recovery_ledger.get_group(group_id) - assert group.agent_name == "genrm_agent" + assert group.task_source == "genrm_compare" assert group.recovery_granularity is RecoveryGranularity.PROMPT_GROUP def test_recovery_mutation_requires_the_controller_barrier(self): @@ -558,7 +559,7 @@ def test_tracked_dispatch_rejects_changed_generations_per_prompt(self): expected_generations=2, target_step=0, start_weight_version=0, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ), @@ -791,6 +792,30 @@ def _nemo_gym_impl( ) +@pytest.mark.parametrize( + ("row", "expected"), + [ + ( + { + "task_source": "shared_resources_server", + "agent_ref": {"name": "resolved_agent"}, + }, + "resolved_agent", + ), + ( + {"task_source": "shared_resources_server"}, + "task-source:shared_resources_server", + ), + ({"agent_ref": {"name": "legacy_agent"}}, "legacy_agent"), + ({}, "nemo_gym"), + ], +) +def test_nemo_gym_metric_namespace_supports_task_source_only_rows( + row: dict, expected: str +) -> None: + assert _nemo_gym_metric_namespace(row) == expected + + def _mask_gate_result(): return { "message_log": [ diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index f133126c61f..96b8a3fd691 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -112,7 +112,7 @@ def test_ledger_round_trip_preserves_group_ownership() -> None: expected_generations=2, target_step=7, start_weight_version=6, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -422,52 +422,72 @@ async def exercise() -> None: asyncio.run(exercise()) -def test_recovery_config_resolves_agent_then_task_then_default() -> None: +def test_recovery_config_resolves_agent_then_task_source_then_default() -> None: config = RolloutRecoveryConfig( default_granularity=RecoveryGranularity.SIBLING, - agent_granularity_overrides={"genrm_agent": RecoveryGranularity.PROMPT_GROUP}, - task_granularity_overrides={ - "math": RecoveryGranularity.PROMPT_GROUP, - "agent_wins": RecoveryGranularity.SIBLING, + task_source_granularity_overrides={ + "genrm_compare": RecoveryGranularity.PROMPT_GROUP, + }, + agent_granularity_overrides={ + "legacy_genrm_agent": RecoveryGranularity.PROMPT_GROUP, + "sibling_agent": RecoveryGranularity.SIBLING, }, ) - agent_policy = config.resolve_for_prompt( + source_policy = config.resolve_for_prompt( { - "task_name": "agent_wins", - "extra_env_info": {"agent_ref": {"name": "genrm_agent"}}, + "extra_env_info": { + "task_source": "genrm_compare", + "agent_ref": {"name": "unmapped_agent"}, + } } ) - task_policy = config.resolve_for_prompt( - {"task_name": "math", "extra_env_info": None} + agent_policy = config.resolve_for_prompt( + { + "extra_env_info": { + "task_source": "genrm_compare", + "agent_ref": {"name": "sibling_agent"}, + } + } ) default_policy = config.resolve_for_prompt( - {"task_name": "other", "extra_env_info": None} + { + "extra_env_info": { + "task_source": "other", + "agent_ref": {"name": "unmapped_agent"}, + } + } ) + with pytest.warns(FutureWarning, match="legacy agent_ref"): + legacy_policy = config.resolve_for_prompt( + {"extra_env_info": {"agent_ref": {"name": "legacy_genrm_agent"}}} + ) - assert agent_policy.agent_name == "genrm_agent" - assert agent_policy.granularity is RecoveryGranularity.PROMPT_GROUP - assert task_policy.agent_name is None - assert task_policy.granularity is RecoveryGranularity.PROMPT_GROUP - assert default_policy.agent_name is None + assert source_policy.task_source == "genrm_compare" + assert source_policy.granularity is RecoveryGranularity.PROMPT_GROUP + assert agent_policy.task_source == "genrm_compare" + assert agent_policy.granularity is RecoveryGranularity.SIBLING + assert default_policy.task_source == "other" assert default_policy.granularity is RecoveryGranularity.SIBLING + assert legacy_policy.task_source is None + assert legacy_policy.granularity is RecoveryGranularity.PROMPT_GROUP @pytest.mark.parametrize( ("prompt", "error_fragment"), [ ( - {"extra_env_info": {"agent_ref": "genrm_agent"}}, + {"extra_env_info": {"task_source": 7}}, + "task_source must be a string or None", + ), + ( + {"extra_env_info": {"agent_ref": "legacy_agent"}}, "agent_ref must be a mapping or None", ), ( {"extra_env_info": {"agent_ref": {"name": 7}}}, "agent_ref.name must be a string or None", ), - ( - {"task_name": 7}, - "task_name must be a string or None", - ), ], ) def test_recovery_config_rejects_malformed_prompt_identity( @@ -477,6 +497,17 @@ def test_recovery_config_rejects_malformed_prompt_identity( RolloutRecoveryConfig().resolve_for_prompt(prompt) +def test_recovery_config_rejects_removed_task_name_override() -> None: + with pytest.raises(ValueError, match="task_source_granularity_overrides"): + RolloutRecoveryConfig( + **{ + "task_granularity_overrides": { + "legacy": RecoveryGranularity.PROMPT_GROUP + } + } + ) + + def test_target_step_none_does_not_mean_unadmitted() -> None: ledger = RolloutRecoveryLedger() record = _reserve( @@ -488,7 +519,7 @@ def test_target_step_none_does_not_mean_unadmitted() -> None: expected_generations=2, target_step=None, start_weight_version=6, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -508,7 +539,7 @@ def test_reserved_group_can_be_admitted_exactly_once() -> None: expected_generations=2, target_step=None, start_weight_version=6, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=False, ) @@ -545,7 +576,7 @@ def test_canonical_groups_are_discarded_without_touching_unfinished_groups() -> expected_generations=2, target_step=7, start_weight_version=7, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -566,7 +597,7 @@ def test_state_dict_stores_a_prompt_ref_without_the_full_payload() -> None: expected_generations=2, target_step=7, start_weight_version=7, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -595,7 +626,7 @@ def test_bind_runtime_prompt_accepts_changed_content_with_the_same_identity() -> expected_generations=2, target_step=7, start_weight_version=7, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -620,7 +651,7 @@ def test_bind_runtime_prompt_rejects_the_wrong_dataset_sample() -> None: expected_generations=2, target_step=7, start_weight_version=7, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -649,7 +680,7 @@ def test_prompt_ref_rehydrates_through_a_restored_shuffled_dataloader() -> None: expected_generations=2, target_step=1, start_weight_version=0, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -683,7 +714,7 @@ def test_restart_preserves_sealed_sibling_and_retries_only_interrupted_one() -> expected_generations=2, target_step=7, start_weight_version=6, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -742,7 +773,7 @@ def test_missing_receipt_is_a_restart_safe_sealed_placeholder( expected_generations=2, target_step=7, start_weight_version=6, - agent_name=None, + task_source=None, recovery_granularity=recovery_granularity, admitted=True, ) @@ -816,21 +847,21 @@ def test_prompt_group_restart_retries_every_sibling_when_one_is_unfinished() -> expected_generations=2, target_step=7, start_weight_version=6, - agent_name="genrm_agent", + task_source="genrm_compare", recovery_granularity=RecoveryGranularity.PROMPT_GROUP, admitted=True, ) _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) state = ledger.state_dict() - assert state["groups"][0]["agent_name"] == "genrm_agent" + assert state["groups"][0]["task_source"] == "genrm_compare" assert state["groups"][0]["recovery_granularity"] == "prompt_group" restored = RolloutRecoveryLedger.from_state_dict(state) _mutate(lambda cut: restored.prepare_for_restart(cut)) recovered = restored.get_group("g7") - assert recovered.agent_name == "genrm_agent" + assert recovered.task_source == "genrm_compare" assert recovered.recovery_granularity is RecoveryGranularity.PROMPT_GROUP assert [sibling.current_attempt.status for sibling in recovered.siblings] == [ RolloutAttemptStatus.ABANDONED, @@ -856,7 +887,7 @@ def test_prompt_group_restart_keeps_a_fully_sealed_group() -> None: expected_generations=2, target_step=7, start_weight_version=6, - agent_name="genrm_agent", + task_source="genrm_compare", recovery_granularity=RecoveryGranularity.PROMPT_GROUP, admitted=True, ) @@ -901,7 +932,7 @@ def test_prompt_group_seal_is_atomic() -> None: expected_generations=2, target_step=7, start_weight_version=6, - agent_name="genrm_agent", + task_source="genrm_compare", recovery_granularity=RecoveryGranularity.PROMPT_GROUP, admitted=True, ) @@ -942,7 +973,7 @@ def test_checkpoint_rejects_ambiguous_finalization_state( expected_generations=1, target_step=7, start_weight_version=6, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) @@ -975,7 +1006,7 @@ def test_checkpoint_rejects_ambiguous_finalization_state( [ ("recovery_granularity", "banana", "invalid recovery_granularity"), ("recovery_granularity", None, "recovery_granularity must be a string"), - ("agent_name", 123, "agent_name must be a string or None"), + ("task_source", 123, "task_source must be a string or None"), ], ) def test_restore_rejects_malformed_recovery_policy_fields( @@ -991,7 +1022,7 @@ def test_restore_rejects_malformed_recovery_policy_fields( expected_generations=1, target_step=7, start_weight_version=6, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=True, ) diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 7d1c4c20b93..9eb3f3d7ab8 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -361,7 +361,7 @@ def reserve_prompt_group( expected_generations=2, target_step=target_step, start_weight_version=7, - agent_name=None, + task_source=None, recovery_granularity=RecoveryGranularity.SIBLING, admitted=admitted, admission_id=admission_id, diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 77fab5315cc..b370bf2b8d6 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -899,6 +899,11 @@ def create_teachers(*args, **kwargs): ValueError, "non-default rollout_recovery policies require", ), + ( + "legacy_agent_recovery_override_without_capture", + ValueError, + "non-default rollout_recovery policies require", + ), ], ) def test_invalid_config_fails_before_setup_factories( @@ -943,9 +948,14 @@ def test_invalid_config_fails_before_setup_factories( mc.rollout_recovery.default_granularity = RecoveryGranularity.PROMPT_GROUP elif invalid_case == "recovery_override_without_capture": mc = _make_master_config() - mc.rollout_recovery.task_granularity_overrides = { + mc.rollout_recovery.task_source_granularity_overrides = { "genrm": RecoveryGranularity.PROMPT_GROUP } + elif invalid_case == "legacy_agent_recovery_override_without_capture": + mc = _make_master_config() + mc.rollout_recovery.agent_granularity_overrides = { + "genrm_agent": RecoveryGranularity.PROMPT_GROUP + } else: # pragma: no cover raise AssertionError(f"unknown test case {invalid_case}") From 5a74a6afc7f2df3be485f95dfbbf513e117505bb Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 7 Sep 2026 17:56:40 -0400 Subject: [PATCH 27/28] fix(nemo-gym): count resolved rollout agents Signed-off-by: Anish Mahishi --- nemo_rl/environments/nemo_gym.py | 34 +++++++----------------- tests/unit/environments/test_nemo_gym.py | 24 +---------------- 2 files changed, 10 insertions(+), 48 deletions(-) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index fc63620057b..a9726d8aac5 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -78,20 +78,6 @@ DEFAULT_THINKING_TAGS = ["", ""] -def _rollout_progress_identity(row: Mapping[str, Any]) -> str: - """Return the concrete agent route, falling back to task provenance.""" - agent_ref = row.get("agent_ref") - if isinstance(agent_ref, Mapping): - agent_name = agent_ref.get("name") - if isinstance(agent_name, str) and agent_name: - return f"agent:{agent_name}" - - task_source = row.get("task_source") - if isinstance(task_source, str) and task_source: - return f"task-source:{task_source}" - return "" - - class NemoGymCompatibleConfig(Protocol): """Configuration fields required to select the NeMo Gym rollout path.""" @@ -630,10 +616,8 @@ async def run_rollouts( examples=nemo_gym_examples, head_server_config=self.head_server_config ) # Gym resolves task_source to agent_ref synchronously in run_examples(). - # Build the counter afterward so completion rows use the same identity. - routing_counts_left = Counter( - _rollout_progress_identity(row) for row in nemo_gym_examples - ) + # Build the counter afterward so completion rows use the resolved identity. + counts_left = Counter(row["agent_ref"]["name"] for row in nemo_gym_examples) num_results = 0 for task in nemo_gym_result_iterator: @@ -685,18 +669,18 @@ async def run_rollouts( / total_time ) - routing_identity = _rollout_progress_identity(nemo_gym_row) - routing_counts_left[routing_identity] -= 1 - if routing_counts_left[routing_identity] <= 0: - routing_counts_left.pop(routing_identity) - if num_results % 10 == 0 and routing_counts_left: - top_left = routing_counts_left.most_common(5) + agent_name = nemo_gym_row["agent_ref"]["name"] + counts_left[agent_name] -= 1 + if counts_left[agent_name] <= 0: + counts_left.pop(agent_name) + if num_results % 10 == 0 and counts_left: + top_left = counts_left.most_common(5) top_left_str = "\n".join( f"{index + 1}. {name}: {count}" for index, (name, count) in enumerate(top_left) ) print( - "Top 5 NeMo Gym routing identities left in this rollout batch: " + "Top 5 NeMo Gym agent refs left in this rollout batch: " f"{top_left_str}", file=sys.stderr, ) diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index 47d5f900792..fbe0bee4097 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -41,7 +41,6 @@ from nemo_rl.environments.nemo_gym import ( NemoGym, NemoGymConfig, - _rollout_progress_identity, build_reward_component_columns, extract_reward_components, setup_nemo_gym_config, @@ -78,27 +77,6 @@ ) -@pytest.mark.parametrize( - ("row", "expected"), - [ - ( - { - "task_source": "math_resources_server", - "agent_ref": {"name": "routed_agent"}, - }, - "agent:routed_agent", - ), - ({"task_source": "math_resources_server"}, "task-source:math_resources_server"), - ({"agent_ref": {"name": "legacy_agent"}}, "agent:legacy_agent"), - ({}, ""), - ], -) -def test_rollout_progress_identity_matches_gym_routing_precedence( - row: dict, expected: str -) -> None: - assert _rollout_progress_identity(row) == expected - - def test_rollout_progress_counter_is_built_after_gym_resolves_task_source( capsys, ) -> None: @@ -155,7 +133,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( asyncio.run(_run()) captured = capsys.readouterr() - assert "1. agent:resolved_agent: 1" in captured.err + assert "1. resolved_agent: 1" in captured.err assert "task-source:test_resources_server" not in captured.err From e7447f7631fd34267f605f3d136a6abd206488b4 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 8 Sep 2026 11:15:21 -0400 Subject: [PATCH 28/28] test(sc): update rollout recovery fixtures Signed-off-by: Anish Mahishi --- tests/unit/experience/test_rollout_manager.py | 7 ++++++- .../experience/test_rollout_reassembler_actor.py | 3 ++- tests/unit/experience/test_rollouts.py | 12 ++++++++++-- tests/unit/single_controller/test_train_pump_e2e.py | 5 ++++- 4 files changed, 22 insertions(+), 5 deletions(-) diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 5429548c780..fce44c92ef3 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -885,7 +885,12 @@ def remote(self, pending, timer_prefix): del pending, timer_prefix async def result_ref(): - return 0, _mask_gate_receipt_result(), None + return ( + 0, + {"name": "resolved-agent"}, + _mask_gate_receipt_result(), + None, + ) async def stream(): yield result_ref() diff --git a/tests/unit/experience/test_rollout_reassembler_actor.py b/tests/unit/experience/test_rollout_reassembler_actor.py index 85a6896bef8..a31c22777d9 100644 --- a/tests/unit/experience/test_rollout_reassembler_actor.py +++ b/tests/unit/experience/test_rollout_reassembler_actor.py @@ -97,8 +97,9 @@ def test_finalize_forwards_loss_multiplier_to_reassembler() -> None: [1.0], mask_sample=[False], fallback_weight_version=4, - prompt_idx=0, + prompt_idx=17, loss_multiplier=0.25, + canonical_sample_ids=["group_g0"], ) diff --git a/tests/unit/experience/test_rollouts.py b/tests/unit/experience/test_rollouts.py index 32cac116ae9..4c6f1533771 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -56,6 +56,7 @@ AsyncNemoGymRolloutImpl, RolloutTimeouts, ) +from nemo_rl.experience.rollout_recovery import RecoveryGranularity from nemo_rl.experience.rollouts import ( _add_multimodal_generation_payload, _reattach_original_multimodal_payloads, @@ -2280,8 +2281,15 @@ def test_nemo_gym_rollout_record_persists_runtime_resolved_agent_ref(): "name": "workplace_assistant_simple_agent", } - async def _run_rollouts(inputs, timer, timer_prefix): - del timer, timer_prefix + async def _run_rollouts( + inputs, + timer, + timer_prefix, + *, + on_completion=None, + recovery_granularity=RecoveryGranularity.SIBLING, + ): + del timer, timer_prefix, on_completion, recovery_granularity for row in inputs: row["agent_ref"] = resolved_agent_ref receipt_completion = SimpleNamespace(env_extras={"ng_receipt": {}}) diff --git a/tests/unit/single_controller/test_train_pump_e2e.py b/tests/unit/single_controller/test_train_pump_e2e.py index 6423b2612bf..cdff8beeca6 100644 --- a/tests/unit/single_controller/test_train_pump_e2e.py +++ b/tests/unit/single_controller/test_train_pump_e2e.py @@ -306,9 +306,12 @@ def test_train_pump_drives_mcore_training_step( sync_weights=lambda *, kv_scales=None: None, ) adv_est = _FakeAdvEstimator() - # Rollout manager stub — SC.__init__ only touches ._tq_buffer. + # Rollout manager stub — recovery is disabled for this native rollout test, + # but SC still binds the shared data-plane checkpoint barrier at startup. rollout_manager = SimpleNamespace( _tq_buffer=None, + recovery_ledger=None, + set_data_plane_checkpoint_barrier=lambda _barrier: None, set_weight_version=lambda v: ray.get( log.record.remote("set_weight_version", {"version": int(v)}) ),