diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index 17abf2ddc76..120d516f09c 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit 17abf2ddc76f74e40d0ba5a15b0401710a46db26 +Subproject commit 120d516f09cfd0377dbed6828247ab1606e9e686 diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index c7ce87acb74..ed03bdc5a84 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -113,6 +113,10 @@ With `checkpointing.save_data_plane: true`, each Single-Controller checkpoint co - A native TQ snapshot containing rollout tensor payloads and TQ state. - A metadata-only replay index describing the completed rollout groups stored in TQ. - A `rollout_recovery.pt` ownership ledger describing unfinished prompt groups that must be redispatched after a restart. +- With `rollout_recovery.target_level: turn` or `prefix`, a + `gym_checkpoint.json` manifest and per-Gym-instance participant snapshots for + unfinished episodes. `prefix` additionally stages decoded tokens from active + model calls at each coordinated checkpoint. - A `replacement_reserve.pt` sidecar containing prompts held for dropped-rollout replacement, when applicable. - The sampler dispatch position needed to continue scheduling from the correct point. @@ -237,9 +241,10 @@ checkpoint exists. :::{note} Completed groups are restored directly from the TQ snapshot. For unfinished -token-capture groups, `rollout_recovery.default_granularity` controls both live -failure and restart behavior: +token-capture groups, `rollout_recovery.target_level` selects the requested +restart boundary: +- `turn` requests coordinated Gym participant recovery from the last saved turn. - `sibling` preserves each sealed sibling and redispatches only unfinished ones. - `prompt_group` retries every sibling in the group when any sibling is unfinished. @@ -247,16 +252,48 @@ failure and restart behavior: `prompt_group` when every generation in a recovered group must come from the policy weights live at redispatch. -`task_source_granularity_overrides` can select the policy using the Gym +Turn recovery coordinates Gym and RL as one checkpoint cut. Each prompt group +is pinned to the Gym shard replica that accepted it. SC first closes its narrow +`/run` dispatch admission gate; already-submitted requests, completion callbacks, +finalization, and TQ writes remain live. Gym parks its participants and commits +the subset of candidate episodes it still owns into the checkpoint's +`gym-instances/` tree. Candidate replies already on the wire drain through the +ordinary RL/TQ completion path. SC then acquires the exclusive data-plane +barrier, verifies that the remaining Gym-owned episodes and their TQ staging +keys match the Gym commit, and captures TQ plus the rollout ledger. If a +candidate is neither exported by Gym nor drained before the deadline, the +checkpoint aborts. Gym leases are renewed throughout and all deployments are +resumed after the cut succeeds or aborts. + +On restart, SC validates that `gym_checkpoint.json` names exactly the episodes +owned by the rollout ledger and that the saved Gym actor topology matches the +live one. It restores Gym before starting rollout dispatch, advances each +restored episode to its next attempt, and routes it back to the same Gym actor. +If one shard replica fails restore, every replica conservatively retires its +replacement attempts before the outer restore fails. This also covers a lost +Ray reply where the remote restore may actually have succeeded, so a partially +restored topology is never released for execution. A second checkpoint taken +after restore but before redispatch includes the unclaimed restored attempt, so +another crash does not lose the saved turn boundary. + +The `turn` target currently requires vLLM generation, NeMo Gym, token capture, +periodic rollout snapshots, `restore_mode: latest`, and the same named Gym +shards and replica counts on restart. Only checkpoints using the current rollout +recovery schema are accepted. + +`task_source_target_level_overrides` can coarsen the target using the Gym `task_source` embedded in the raw rollout row. Unlike `agent_ref`, this identity is available before Gym resolves the concrete agent and SC reserves the recovery group. When a row already carries an `agent_ref`, a matching -`agent_granularity_overrides` entry wins over a matching task-source entry, +`agent_target_level_overrides` entry wins over a matching task-source entry, mirroring Gym's concrete-route precedence. Otherwise the task-source override, then the global default, applies. The agent map also keeps datasets collated before Gym recorded `task_source` working, although re-collating them is -recommended. Non-default policies require `token_capture.enabled: true`. The -task source and resolved policy are persisted in `rollout_recovery.pt`, so +recommended. Overrides may only select a coarser level than the global target. +Until Gym advertises group-scoring capabilities, configure group-scored routes +such as GenRM explicitly as `prompt_group`. Non-default policies require +`token_capture.enabled: true`. The task source and resolved policy are persisted +in `rollout_recovery.pt`, so recovery does not reinterpret an existing group using changed configuration. A generation that already finished keeps its tokens in the token-capture staging area, so `sibling` reuses them unchanged; a redispatched sibling produces a new diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index ecd7b66ad9f..13ccea271e3 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -219,20 +219,22 @@ token_capture: enabled: false staging_partition: rollout_staging -# Retry and restore policy for unfinished token-capture groups. "sibling" reuses -# already sealed generations; "prompt_group" regenerates every sibling when any -# sibling fails in-process or was unfinished at a checkpoint. +# Target recovery boundary for unfinished token-capture groups. "prefix" also +# preserves decoded tokens from an active model call; "turn" enables coordinated +# Gym participant snapshots; "sibling" preserves sealed siblings; "prompt_group" +# regenerates every sibling in an unfinished group. rollout_recovery: - default_granularity: sibling - # Per-Gym-task-source override, keyed on extra_env_info.task_source. - # task_source_granularity_overrides: + target_level: sibling + # Per-Gym-task-source coarsening override, keyed on extra_env_info.task_source. + # Until Gym advertises group-scoring capability, configure GenRM explicitly: + # task_source_target_level_overrides: # genrm_compare_resources_server: prompt_group - task_source_granularity_overrides: {} + task_source_target_level_overrides: {} # Agent-specific override for rows already carrying extra_env_info.agent_ref. # A matching agent entry wins over a matching task-source entry and also # provides compatibility for legacy datasets without task_source. - # agent_granularity_overrides: {genrm_agent: prompt_group} - agent_granularity_overrides: {} + # agent_target_level_overrides: {genrm_agent: prompt_group} + agent_target_level_overrides: {} # Leaving either map above non-empty requires the top-level token_capture section # to set enabled: true. diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index a4b32fdf879..8c165f884ac 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -267,20 +267,22 @@ token_capture: enabled: false staging_partition: rollout_staging -# Retry and restore policy for unfinished token-capture groups. "sibling" reuses -# already sealed generations; "prompt_group" regenerates every sibling when any -# sibling fails in-process or was unfinished at a checkpoint. +# Target recovery boundary for unfinished token-capture groups. "prefix" also +# preserves decoded tokens from an active model call; "turn" enables coordinated +# Gym participant snapshots; "sibling" preserves sealed siblings; "prompt_group" +# regenerates every sibling in an unfinished group. rollout_recovery: - default_granularity: sibling - # Per-Gym-task-source override, keyed on extra_env_info.task_source. - # task_source_granularity_overrides: + target_level: sibling + # Per-Gym-task-source coarsening override, keyed on extra_env_info.task_source. + # Until Gym advertises group-scoring capability, configure GenRM explicitly: + # task_source_target_level_overrides: # genrm_compare_resources_server: prompt_group - task_source_granularity_overrides: {} + task_source_target_level_overrides: {} # Agent-specific override for rows already carrying extra_env_info.agent_ref. # A matching agent entry wins over a matching task-source entry and also # provides compatibility for legacy datasets without task_source. - # agent_granularity_overrides: {genrm_agent: prompt_group} - agent_granularity_overrides: {} + # agent_target_level_overrides: {genrm_agent: prompt_group} + agent_target_level_overrides: {} # Leaving either map above non-empty requires the top-level token_capture section # to set enabled: true. diff --git a/examples/nemo_gym/grpo_workplace_assistant_two_replicas.yaml b/examples/nemo_gym/grpo_workplace_assistant_two_replicas.yaml new file mode 100644 index 00000000000..9a1e6f5c2cc --- /dev/null +++ b/examples/nemo_gym/grpo_workplace_assistant_two_replicas.yaml @@ -0,0 +1,25 @@ +# Two Workplace Assistant Gym replicas for checkpoint/recovery tests. Replicas +# share one port range, so NeMo-RL intentionally places them on separate Ray +# nodes with the default STRICT_SPREAD strategy. The training shape and model +# are supplied by the functional launcher. + +defaults: "grpo_qwen3_30ba3b_instruct.yaml" + +env: + nemo_gym: + # Shards replace the inherited unsharded server list. Policy-model config is + # an RL-owned overlay duplicated into every otherwise-identical replica. + config_paths: null + code_gen: null + shards: + - name: workplace + replicas: 2 + port_range_low: 5000 + port_range_high: 5999 + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - responses_api_agents/checkpoint_test_agent/configs/workplace_assistant.yaml + common_inherited_overlays: + - policy_model + allowed_duplicate_entries: + - policy_model diff --git a/examples/nemo_gym/prefetch_venvs.py b/examples/nemo_gym/prefetch_venvs.py index c0ab0c8c570..1fbe45918c1 100644 --- a/examples/nemo_gym/prefetch_venvs.py +++ b/examples/nemo_gym/prefetch_venvs.py @@ -29,6 +29,7 @@ from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import init_ray +from nemo_rl.environments.gym_checkpoint_adapter import GymCheckpointInstance from nemo_rl.environments.nemo_gym import ( NemoGym, NemoGymConfig, @@ -86,6 +87,8 @@ def prefetch_nemo_gym_venvs(config_paths: list[str]) -> None: ray_namespace=None, initial_global_config_dict=nemo_gym_dict, invalid_tool_call_patterns=None, + checkpoint_instance=GymCheckpointInstance("nemo_gym", 0), + turn_checkpointing_enabled=False, ) nemo_gym_opts = { diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 49d3f5186eb..2e6428fb6c3 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -263,7 +263,11 @@ class DataPlaneCheckpointBarrier: """ def __init__(self) -> None: - self._condition = asyncio.Condition() + lock = asyncio.Lock() + self._condition = asyncio.Condition(lock) + # Separate from _condition so a finished mutation wakes only waiters for + # it, not every mutation queued behind an active checkpoint. + self._mutation_finished = asyncio.Condition(lock) self._checkpoint_active = False self._active_mutations = 0 self._section_holders: set[asyncio.Task[Any]] = set() @@ -329,9 +333,25 @@ async def mutation( # Count the section even when its body raised. A redundant # snapshot is safe; skipping a partially applied mutation is not. self._mutation_version += 1 + self._mutation_finished.notify_all() if self._active_mutations == 0: self._condition.notify_all() + async def wait_for_mutation(self, since_version: int, timeout_s: float) -> None: + """Return once a mutation finishes after ``since_version``, or on timeout. + + Read :attr:`mutation_version` before inspecting state, then pass it here: + a mutation that lands in between returns at once instead of being missed. + """ + async with self._mutation_finished: + try: + async with asyncio.timeout(max(0.0, timeout_s)): + await self._mutation_finished.wait_for( + lambda: self._mutation_version != since_version + ) + except TimeoutError: + return + async def drain_telemetry(self) -> DataPlaneCheckpointBarrierTelemetry: """Return and reset interval waits while preserving current state.""" async with self._condition: diff --git a/nemo_rl/algorithms/distillation.py b/nemo_rl/algorithms/distillation.py index d4a67911555..e530dddb6c7 100644 --- a/nemo_rl/algorithms/distillation.py +++ b/nemo_rl/algorithms/distillation.py @@ -540,6 +540,7 @@ def init_nemo_gym(): # Distillation does not configure vLLM for router replay. enable_router_replay=False, use_fastokens=bool(policy_config["tokenizer"].get("use_fastokens")), + turn_checkpointing_enabled=False, ) try: validate_dataset_agent_coverage( diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index d7482dde68d..a331e5d35f6 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -893,6 +893,7 @@ def _spinup_nemo_gym(base_urls, model_name): tokenizer=tokenizer, enable_router_replay=router_replay_enabled(policy_config), use_fastokens=bool(policy_config["tokenizer"].get("use_fastokens")), + turn_checkpointing_enabled=False, ) train_splits = ( {f"train[{name}]": split for name, split in dataset.items()} diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index eb63392701b..b4c0c32e645 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -54,8 +54,8 @@ import uuid import warnings from collections import deque -from collections.abc import Iterator -from dataclasses import dataclass +from collections.abc import AsyncIterator, Iterable, Iterator, Mapping +from dataclasses import dataclass, replace from functools import partial from pathlib import Path from typing import ( @@ -68,6 +68,7 @@ cast, ) +import numpy as np import ray import torch from ray.exceptions import RayActorError @@ -160,10 +161,26 @@ ) from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.refit_watchdog import RefitAborted, is_refit_context_lost -from nemo_rl.environments.nemo_gym import should_use_nemo_gym +from nemo_rl.environments.gym_checkpoint_adapter import GymCheckpointEpisode +from nemo_rl.environments.gym_checkpoint_coordinator import ( + GYM_CHECKPOINT_MANIFEST_FILENAME, + GymCheckpointCommitResult, + GymCheckpointCoordinator, + is_transient_gym_checkpoint_failure, + load_gym_checkpoint_manifest, + staging_keys_of, + write_gym_checkpoint_manifest, +) +from nemo_rl.environments.nemo_gym import ( + as_nemo_gym_shard_set, + should_use_nemo_gym, +) from nemo_rl.experience.failures import RolloutStall from nemo_rl.experience.payload import VIOLATION_TAG_KEYS -from nemo_rl.experience.rollout_manager import RolloutOutcome +from nemo_rl.experience.rollout_manager import ( + RolloutDispatchAdmissionGate, + RolloutOutcome, +) from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, ROLLOUT_RECOVERY_STATE_FILENAME, @@ -219,6 +236,28 @@ # cleaner exit. Not configurable for the same reason. _SUPERVISOR_DRAIN_TIMEOUT_S = 30.0 +# Rollout-deadline pause owned by a Gym checkpoint while it holds rollouts parked. +_GYM_CHECKPOINT_DEADLINE_HOLDER = "gym_checkpoint" + +# Episodes named in a Gym checkpoint error; the rest are only counted. +_GYM_EPISODE_SAMPLE_SIZE = 5 + + +def _summarize_gym_episodes(episodes_by_instance: Mapping[str, Iterable[Any]]) -> str: + """Count episodes per instance and name a few, however many there are.""" + parts = [] + for instance_id, episodes in sorted(episodes_by_instance.items()): + ordered = sorted(episodes, key=repr) + if not ordered: + continue + sample = ordered[:_GYM_EPISODE_SAMPLE_SIZE] + more = len(ordered) - len(sample) + parts.append( + f"{instance_id}: {len(ordered)} episode(s), e.g. {sample!r}" + + (f" and {more} more" if more else "") + ) + return "; ".join(parts) or "none" + @dataclass(frozen=True) class _RolloutCheckpointSaveResult: @@ -245,6 +284,20 @@ class _RolloutCheckpointCut: tq_save_seconds: float +def _write_training_info( + checkpoint_path: PathLike, training_info: dict[str, Any] +) -> None: + """Atomically refresh training metadata after the final checkpoint cut.""" + serializable = dict(training_info) + for key, value in serializable.items(): + if isinstance(value, (torch.Tensor, np.ndarray)): + serializable[key] = value.item() + path = Path(checkpoint_path) / "training_info.json" + tmp_path = path.with_suffix(".json.tmp") + tmp_path.write_text(json.dumps(serializable)) + os.replace(tmp_path, path) + + def _latest_generation_values( metrics: dict[str, Any], metric_name: str ) -> dict[int, float]: @@ -476,6 +529,22 @@ def __init__( # exists to remove. A missing field should break loudly at construction, where # it costs five minutes, not quietly at hour three of a run. self._env_handles = actor_args.env_handles + self._gym_checkpoint_coordinator: Optional[GymCheckpointCoordinator] = None + self._generation_prefix_cuts_enabled = ( + self._master_config.rollout_recovery.generation_prefix_cuts_enabled + ) + self._rollout_dispatch_admission_gate: Optional[ + RolloutDispatchAdmissionGate + ] = None + if self._master_config.rollout_recovery.turn_checkpointing_enabled: + self._gym_checkpoint_coordinator = GymCheckpointCoordinator( + as_nemo_gym_shard_set(self._env_handles["nemo_gym"]), + control_timeout_s=(self._master_config.token_capture.control_timeout_s), + ) + self._rollout_dispatch_admission_gate = RolloutDispatchAdmissionGate() + self._rollout_manager.set_rollout_dispatch_admission_gate( + self._rollout_dispatch_admission_gate + ) # These two keep the getattr for a genuinely different reason: None is a # meaningful value meaning "feature off", and it is also their default. Absence # therefore degrades to the documented off state rather than to a broken one. @@ -1073,12 +1142,9 @@ async def _maybe_restore_rollout_recovery( "TQ checkpoint metadata" ) expected_schema_version = metadata.get("rollout_recovery_schema_version") - if ( - isinstance(expected_schema_version, bool) - or expected_schema_version != ROLLOUT_RECOVERY_SCHEMA_VERSION - ): + if expected_schema_version != ROLLOUT_RECOVERY_SCHEMA_VERSION: raise ValueError( - "native TQ checkpoint rollout recovery schema mismatch: " + "native TQ checkpoint has an unsupported rollout recovery schema: " f"checkpoint={expected_schema_version!r}, " f"expected={ROLLOUT_RECOVERY_SCHEMA_VERSION}" ) @@ -1112,6 +1178,12 @@ async def _maybe_restore_rollout_recovery( weights_only=True, ) parsed_state = parse_rollout_recovery_state(state) + if parsed_state.ledger_state["schema_version"] != expected_schema_version: + raise ValueError( + "rollout recovery sidecar schema does not match native TQ " + f"metadata: sidecar={parsed_state.ledger_state['schema_version']!r}, " + f"metadata={expected_schema_version!r}" + ) if len(parsed_state.ledger_state["groups"]) != expected_group_count: raise ValueError( "rollout recovery sidecar group count does not match native " @@ -1123,22 +1195,115 @@ async def _maybe_restore_rollout_recovery( "recovery_restore" ) 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 - ) - canonical_group_ids = { - 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( + gym_state_groups = [ + group + for group in recovery_ledger.groups() + if group.restore_level.requires_gym_checkpoint + ] + restore_id: Optional[str] = None + restore_coordinator: Optional[GymCheckpointCoordinator] = None + restore_manifest = None + gym_restore_completed = False + restored_gym_episodes: set[tuple[str, str, int]] = set() + restored_gym_staging_keys: set[str] = set() + try: + if gym_state_groups: + coordinator = self._gym_checkpoint_coordinator + if coordinator is None: + raise RuntimeError( + "rollout checkpoint contains turn/prefix Gym state, but " + "Gym checkpoint coordination is disabled" + ) + manifest_path = ( + Path(self._last_checkpoint_path) + / GYM_CHECKPOINT_MANIFEST_FILENAME + ) + if not manifest_path.is_file(): + raise FileNotFoundError( + "turn/prefix rollout recovery requires the matching Gym " + f"manifest at {manifest_path}" + ) + manifest = await asyncio.to_thread( + load_gym_checkpoint_manifest, + self._last_checkpoint_path, + ) + expected_raw_inventory = recovery_ledger.gym_checkpoint_inventory( + coordinator.instance_ids + ) + expected_inventory = { + instance_id: tuple( + GymCheckpointEpisode(rollout_id, attempt) + for rollout_id, attempt in episodes + ) + for instance_id, episodes in expected_raw_inventory.items() + } + if manifest.instances != expected_inventory: + raise ValueError( + "Gym checkpoint manifest does not match the rollout " + "recovery ledger" + ) + restore_id = f"restore-{uuid.uuid4().hex}" + restore_coordinator = coordinator + restore_manifest = manifest + await coordinator.restore( + restore_id, + self._last_checkpoint_path, + manifest, + ) + gym_restore_completed = True + restored_gym_episodes = { + (instance_id, episode.rollout_id, episode.attempt) + for instance_id, episodes in manifest.instances.items() + for episode in episodes + } + restored_gym_staging_keys = { + key for keys in manifest.staging_keys.values() for key in keys + } + + recovery_ledger.prepare_for_restart( cut, - replay_metadata=canonical_state, - clear_unreferenced=True, + restored_gym_episodes=restored_gym_episodes, ) - await self._rehydrate_rollout_recovery_prompts(cut) + self._batch_shortfall = parsed_state.batch_shortfall + canonical_state = self._buffer.metadata_state_dict( + saved_capacity=self._async_cfg.max_buffered_rollouts + ) + canonical_group_ids = { + 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( + cut, + replay_metadata=canonical_state, + clear_unreferenced=True, + additional_expected_staging_keys=restored_gym_staging_keys, + ) + await self._rehydrate_rollout_recovery_prompts(cut) + except BaseException as error: + if restore_id is not None and restore_coordinator is not None: + if gym_restore_completed and restore_manifest is not None: + try: + await restore_coordinator.retire_restored( + restore_id, + restore_manifest, + ) + except Exception as retire_error: + error.add_note( + "Gym retire of restored attempts also failed: " + f"{type(retire_error).__name__}: {retire_error}" + ) + try: + await restore_coordinator.resume(restore_id) + except Exception as resume_error: + error.add_note( + "Gym restore cleanup also failed: " + f"{type(resume_error).__name__}: {resume_error}" + ) + raise + else: + if restore_id is not None and restore_coordinator is not None: + await restore_coordinator.resume(restore_id) self._sampler_stamps_target_steps = ( parsed_state.sampler_stamps_target_steps if parsed_state.sampler_stamps_target_steps is not None @@ -1475,10 +1640,12 @@ async def _validate_rollout_recovery_inventory( *, replay_metadata: Optional[TQReplayMetadataState], clear_unreferenced: bool, + additional_expected_staging_keys: Optional[set[str]] = None, ) -> int: """Validate staging ownership while the caller holds a stable cut.""" cut.require_live() expected_staging_keys = self._rollout_recovery_ledger.expected_staging_keys() + expected_staging_keys.update(additional_expected_staging_keys or ()) if replay_metadata is not None: for group in replay_metadata["groups"]: for tag in group["meta"].tags or []: @@ -4108,6 +4275,8 @@ async def _capture_rollout_checkpoint_cut( self, cut: DataPlaneMutationCut, checkpoint_path: PathLike, + *, + additional_expected_staging_keys: Optional[set[str]] = None, ) -> _RolloutCheckpointCut: """Save TQ and capture matching restart state under the barrier. @@ -4150,6 +4319,7 @@ async def _capture_rollout_checkpoint_cut( cut, replay_metadata=replay_metadata, clear_unreferenced=False, + additional_expected_staging_keys=additional_expected_staging_keys, ) tq_save_started = time.monotonic() await self._save_data_plane_checkpoint( @@ -4176,6 +4346,327 @@ async def _capture_rollout_checkpoint_cut( tq_save_seconds=tq_save_seconds, ) + async def _capture_coordinated_rollout_checkpoint_cut( + self, + cut: DataPlaneMutationCut, + checkpoint_path: PathLike, + *, + gym_commit: GymCheckpointCommitResult | None, + ) -> _RolloutCheckpointCut: + """Validate the prepared Gym cut, then capture matching TQ state.""" + gym_staging_keys: set[str] = set() + if gym_commit is not None: + gym_staging_keys = await self._validate_gym_checkpoint_cut( + cut, gym_commit, checkpoint_path=checkpoint_path + ) + return await self._capture_rollout_checkpoint_cut( + cut, + checkpoint_path, + additional_expected_staging_keys=gym_staging_keys, + ) + + def _gym_checkpoint_inventory( + self, coordinator: GymCheckpointCoordinator + ) -> dict[str, tuple[GymCheckpointEpisode, ...]]: + """Return Gym-owned episode attempts using the coordinator wire type.""" + raw_inventory = self._rollout_recovery_ledger.gym_checkpoint_inventory( + coordinator.instance_ids + ) + return { + instance_id: tuple( + GymCheckpointEpisode(rollout_id, attempt) + for rollout_id, attempt in episodes + ) + for instance_id, episodes in raw_inventory.items() + } + + def _queued_gym_retirements( + self, coordinator: GymCheckpointCoordinator + ) -> dict[str, frozenset[GymCheckpointEpisode]]: + """Return dropped episodes Gym may still hold, awaiting their retire.""" + raw = self._rollout_recovery_ledger.gym_checkpoint_retirements( + coordinator.instance_ids + ) + return { + instance_id: frozenset( + GymCheckpointEpisode(rollout_id, attempt) + for rollout_id, attempt in episodes + ) + for instance_id, episodes in raw.items() + } + + async def _retire_dropped_gym_episodes( + self, + coordinator: GymCheckpointCoordinator, + checkpoint_id: str, + ) -> None: + """Retire episodes RL dropped while Gym may still be running them. + + Abandoning or discarding a group leaves an admitted episode running in + Gym, which would export it at commit, outside the candidate set. Gym + accepts a retire only while idle, so this runs before prepare. A drop + recorded after it stays queued: this checkpoint commits it in scope but + leaves it out of the saved manifest, and a later retire stops it. + """ + ledger = self._rollout_recovery_ledger + raw = ledger.gym_checkpoint_retirements(coordinator.instance_ids) + if not any(raw.values()): + return + await coordinator.retire( + checkpoint_id, + { + instance_id: tuple( + GymCheckpointEpisode(rollout_id, attempt) + for rollout_id, attempt in episodes + ) + for instance_id, episodes in raw.items() + }, + ) + # Only what was retired: a drop recorded meanwhile stays for next time. + ledger.mark_gym_retired(raw) + + async def _drain_non_exported_gym_candidates( + self, + coordinator: GymCheckpointCoordinator, + candidates: dict[str, tuple[GymCheckpointEpisode, ...]], + commit: GymCheckpointCommitResult, + ) -> None: + """Wait for on-wire replies that Gym did not retain in its checkpoint.""" + pending = { + instance_id: set(candidates[instance_id]) + - set(commit.manifest.instances[instance_id]) + for instance_id in coordinator.instance_ids + } + if not any(pending.values()): + return + # Only an episode whose reply is already on the wire can resolve here. A + # session Gym still holds is parked by prepare, so waiting would only time out. + stranded = { + instance_id: pending[instance_id].intersection( + commit.instances[instance_id].unowned_live_episodes + ) + for instance_id in coordinator.instance_ids + } + if any(stranded.values()): + raise RuntimeError( + "Gym checkpoint holds parked sessions that no environment or legacy " + "agent participant exported, so they can neither be restored nor " + "finish while the deployment is prepared: " + f"stranded={_summarize_gym_episodes(stranded)}" + ) + + # Only the pending episodes are checked, and only after a ledger mutation + # could have resolved one: rescanning every live rollout on a timer would + # block the event loop that delivers the replies this waits for. + barrier = self._data_plane_checkpoint_barrier + ledger = self._rollout_recovery_ledger + waiting = { + instance_id: {(episode.rollout_id, episode.attempt) for episode in episodes} + for instance_id, episodes in pending.items() + } + deadline = time.monotonic() + coordinator.control_timeout_s + while True: + # Read before checking, so a mutation landing in between is not missed. + version = barrier.mutation_version + waiting = ledger.gym_episodes_still_held(waiting) + if not any(waiting.values()): + return + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + "Gym checkpoint could not classify every candidate before the " + f"drain deadline: unresolved={_summarize_gym_episodes(waiting)}" + ) + await barrier.wait_for_mutation(version, remaining) + + async def _validate_gym_checkpoint_cut( + self, + cut: DataPlaneMutationCut, + commit: GymCheckpointCommitResult, + *, + checkpoint_path: PathLike, + ) -> set[str]: + """Require Gym ownership and referenced staging rows to match this cut. + + Gym may have exported episodes RL dropped while it was prepared: it + refuses their retire until resume, so they are still queued for one. + They and their staged rows are left out of the saved manifest, since + restore requires it to match RL's inventory and RL no longer continues + them; a dropped group may already have cleaned those rows up. + + Returns the staging rows the saved manifest keeps. ``commit`` still + names the dropped episodes' rows, so callers must use these instead. + """ + cut.require_live() + coordinator = self._gym_checkpoint_coordinator + if coordinator is None: + raise RuntimeError("Gym commit exists without a live Gym coordinator") + + actual = self._gym_checkpoint_inventory(coordinator) + expected = commit.manifest.instances + dropped = self._queued_gym_retirements(coordinator) + mismatches = {} + kept: dict[str, tuple[GymCheckpointEpisode, ...]] = {} + for instance_id in coordinator.instance_ids: + owned = set(actual[instance_id]) + exported = set(expected[instance_id]) + if owned - exported or not exported - owned <= dropped[instance_id]: + mismatches[instance_id] = { + "expected": sorted(exported, key=repr), + "actual": sorted(owned, key=repr), + } + kept[instance_id] = tuple( + episode for episode in expected[instance_id] if episode in owned + ) + if mismatches: + raise RuntimeError( + "Gym episode ownership changed after commit and before the TQ cut: " + f"{mismatches!r}" + ) + staging_keys = { + instance_id: staging_keys_of( + commit.instances[instance_id], kept[instance_id] + ) + for instance_id in coordinator.instance_ids + } + if kept != expected: + await asyncio.to_thread( + write_gym_checkpoint_manifest, + Path(checkpoint_path), + replace(commit.manifest, instances=kept, staging_keys=staging_keys), + ) + + expected_staging_keys = {key for keys in staging_keys.values() for key in keys} + if not expected_staging_keys: + return expected_staging_keys + actual_staging_keys = set( + await self._call_dp( + "list_sample_ids", + partition_id=self._master_config.token_capture.staging_partition, + ) + ) + missing = sorted(expected_staging_keys - actual_staging_keys) + if missing: + raise RuntimeError( + "Gym checkpoint references staging rows missing from live TQ state: " + f"missing={missing[:10]!r} (total={len(missing)})" + ) + return expected_staging_keys + + @contextlib.asynccontextmanager + async def _prepared_gym_checkpoint( + self, + checkpoint_path: PathLike, + *, + checkpoint_id: str, + enabled: bool = True, + ) -> AsyncIterator[GymCheckpointCommitResult | None]: + """Close Gym admission, commit owned episodes, and drain on-wire replies.""" + coordinator = self._gym_checkpoint_coordinator + if not enabled or coordinator is None: + yield None + return + gate = self._rollout_dispatch_admission_gate + if gate is None: + raise RuntimeError( + "Gym checkpoint coordinator has no dispatch admission gate" + ) + + fence_timeout_s: float | None = None + fence_started = False + if self._generation_prefix_cuts_enabled: + fence_timeout_s = self._master_config.token_capture.control_timeout_s + try: + await asyncio.to_thread( + self._gen.begin_token_capture_snapshot_fence, + timeout_s=fence_timeout_s, + ) + fence_started = True + except BaseException as fence_error: + # Fan-out may have fenced a subset of workers before raising. + # Always attempt the idempotent release on every worker. + try: + await asyncio.to_thread( + self._gen.end_token_capture_snapshot_fence, + timeout_s=fence_timeout_s, + ) + except BaseException as release_error: + raise BaseExceptionGroup( + "token-capture snapshot fencing failed and partially fenced " + "workers could not be released", + [fence_error, release_error], + ) + raise + + checkpoint_error: BaseException | None = None + try: + async with gate.closed(): + # Gym parks every live rollout until resume, so the wait must not + # count against their deadlines. A separate holder keeps a + # colocated engine's own pause intact when this one ends. + self._rollout_manager.suspend_request_deadlines( + _GYM_CHECKPOINT_DEADLINE_HOLDER + ) + try: + # Gym refuses a retire once a checkpoint is open, so retire + # while it is still idle; dispatch is already closed, so + # nothing new races in. + await self._retire_dropped_gym_episodes(coordinator, checkpoint_id) + async with coordinator.prepared(checkpoint_id): + candidates = self._gym_checkpoint_inventory(coordinator) + # A row dropped since that retire is still running in + # Gym, which exports it; naming it keeps the export in + # scope, and the cut leaves it out of the saved manifest. + dropped = self._queued_gym_retirements(coordinator) + scope = { + instance_id: episodes + + tuple( + sorted( + dropped[instance_id] - set(episodes), + key=lambda episode: ( + episode.rollout_id, + episode.attempt, + ), + ) + ) + for instance_id, episodes in candidates.items() + } + commit = await coordinator.commit( + checkpoint_id, + Path(checkpoint_path), + scope, + ) + await self._drain_non_exported_gym_candidates( + coordinator, + candidates, + commit, + ) + yield commit + finally: + self._rollout_manager.resume_request_deadlines( + _GYM_CHECKPOINT_DEADLINE_HOLDER + ) + except BaseException as error: + checkpoint_error = error + finally: + if fence_started: + assert fence_timeout_s is not None + try: + await asyncio.to_thread( + self._gen.end_token_capture_snapshot_fence, + timeout_s=fence_timeout_s, + ) + except BaseException as release_error: + if checkpoint_error is not None: + raise BaseExceptionGroup( + "Gym checkpoint failed and generation workers could " + "not be released", + [checkpoint_error, release_error], + ) + raise + if checkpoint_error is not None: + raise checkpoint_error + async def _write_rollout_checkpoint_sidecars( self, checkpoint_path: Path, @@ -4243,6 +4734,12 @@ async def _save_rollout_checkpoint( not force and self._last_rollout_snapshot_mutation_version == self._data_plane_checkpoint_barrier.mutation_version + # Decoding grows a live call's recoverable prefix without + # mutating the data plane, so keep cutting while one is active. + and not ( + self._generation_prefix_cuts_enabled + and self._rollout_recovery_ledger.has_dispatched_generation_prefix() + ) ): return _RolloutCheckpointSaveResult( saved=False, @@ -4307,24 +4804,32 @@ async def _save_rollout_checkpoint( prepare_snapshot_paths, anchor ) try: - barrier_requested = time.monotonic() - async with self._data_plane_checkpoint_barrier.checkpoint() as cut: - barrier_acquired = time.monotonic() - if ( - self._optimizer_commit_in_progress - or self._train_steps != expected_train_step - or self._trainer_version != expected_trainer_version - ): - await asyncio.to_thread(partial(shutil.rmtree, tmp_path)) - return _RolloutCheckpointSaveResult( - saved=False, - reason="trainer_state_changed", + async with self._prepared_gym_checkpoint( + tmp_path, + checkpoint_id=f"rollout-{uuid.uuid4().hex}", + ) as gym_commit: + barrier_requested = time.monotonic() + async with self._data_plane_checkpoint_barrier.checkpoint() as cut: + barrier_acquired = time.monotonic() + if ( + self._optimizer_commit_in_progress + or self._train_steps != expected_train_step + or self._trainer_version != expected_trainer_version + ): + await asyncio.to_thread(partial(shutil.rmtree, tmp_path)) + return _RolloutCheckpointSaveResult( + saved=False, + reason="trainer_state_changed", + ) + snapshot_epoch = self._current_epoch + snapshot_cut = ( + await self._capture_coordinated_rollout_checkpoint_cut( + cut, + tmp_path, + gym_commit=gym_commit, + ) ) - snapshot_epoch = self._current_epoch - snapshot_cut = await self._capture_rollout_checkpoint_cut( - cut, tmp_path - ) - barrier_released = time.monotonic() + barrier_released = time.monotonic() sidecar_save_started = time.monotonic() controller_sidecar_bytes = ( @@ -4470,6 +4975,10 @@ async def _rollout_checkpoint_pump(self) -> None: failure_reason: RolloutCheckpointAttemptReason = "timeout" elif isinstance(error, OSError): failure_reason = "io_error" + elif is_transient_gym_checkpoint_failure(error): + # Gym was resumed, or its lease reopens it; nothing RL + # holds changed, so the next attempt can simply retry. + failure_reason = "gym_unavailable" else: failure_reason = "invariant_error" self._log_rollout_checkpoint_outcome( @@ -4477,7 +4986,7 @@ async def _rollout_checkpoint_pump(self) -> None: reason=failure_reason, attempt_duration_seconds=time.monotonic() - attempt_started, ) - if not isinstance(error, (OSError, TimeoutError)): + if failure_reason == "invariant_error": raise if deadline_due: raise RuntimeError( @@ -4747,95 +5256,123 @@ async def _save_checkpoint_impl( rollout_recovery_payload: Optional[bytes] = None rollout_recovery_payload_sha256: Optional[str] = None - # Admission, dataloader movement, replay mutations, and canonical TQ writes - # all take the mutation side of this barrier. Capture every restart-facing - # 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() as cut: - save_state.current_step = self._train_steps - save_state.total_steps = self._train_steps - save_state.trainer_version = self._trainer_version - save_state.current_epoch = self._current_epoch - save_state.consumed_samples = self._consumed_samples - save_state.total_valid_tokens = self._total_valid_tokens - save_state.sampler_name = self._async_cfg.sampler.name - save_state.sampler_dispatch_index = self._sampler.dispatch_index - dataloader_state = self._dataloader.state_dict() - # The spare pool and dataloader advance together under the same mutation - # cut in recovery-enabled dispatch, so preserve them in this cut too. - reserve_state = list(self._replacement_reserve) - - checkpoint_path: PathLike = await asyncio.to_thread( # pyrefly: ignore[bad-assignment] the PathLike alias resolves inconsistently under pyrefly's import-cycle breaking - self._checkpointer.init_tmp_checkpoint, - self._train_steps, - vars(save_state), - self._master_config, - ) + # Gym commit needs a destination before the final data-plane cut. Seed + # the temporary trainer directory now, then atomically refresh its + # training_info.json with the exact controller values captured below. + save_state.current_step = self._train_steps + save_state.total_steps = self._train_steps + save_state.trainer_version = self._trainer_version + save_state.current_epoch = self._current_epoch + save_state.consumed_samples = self._consumed_samples + save_state.total_valid_tokens = self._total_valid_tokens + save_state.sampler_name = self._async_cfg.sampler.name + save_state.sampler_dispatch_index = self._sampler.dispatch_index + checkpoint_path: PathLike = await asyncio.to_thread( # pyrefly: ignore[bad-assignment] the PathLike alias resolves inconsistently under pyrefly's import-cycle breaking + self._checkpointer.init_tmp_checkpoint, + self._train_steps, + vars(save_state), + self._master_config, + ) - if self._master_config.checkpointing.get("save_data_plane"): - training_owned_groups = self._buffer.training_owned_replay_groups() - if training_owned_groups: - raise RuntimeError( - "full trainer checkpoint still owns streamed training rows: " - f"groups={[group['group_id'] for group in training_owned_groups]!r}" - ) - if self._sampler.supports_buffer_checkpoint: - replay_metadata = self._buffer.metadata_state_dict( - saved_capacity=self._async_cfg.max_buffered_rollouts - ) - if replay_metadata is not None: - await self._validate_replay_inventory(replay_metadata) - - if self._rollout_recovery_enabled: - rollout_recovery_state = build_rollout_recovery_state( - self._rollout_manager.recovery_ledger, - batch_shortfall=self._batch_shortfall, - sampler_stamps_target_steps=(self._sampler_stamps_target_steps), - ) + save_data_plane = bool(self._master_config.checkpointing.get("save_data_plane")) + async with self._prepared_gym_checkpoint( + checkpoint_path, + checkpoint_id=f"trainer-{uuid.uuid4().hex}", + enabled=save_data_plane, + ) as gym_commit: + # Admission, dataloader movement, replay mutations, and canonical TQ + # writes all take the mutation side of this barrier. Gym is already + # parked and on-wire replies have drained before exclusivity starts. + 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 + save_state.current_epoch = self._current_epoch + save_state.consumed_samples = self._consumed_samples + save_state.total_valid_tokens = self._total_valid_tokens + save_state.sampler_name = self._async_cfg.sampler.name + save_state.sampler_dispatch_index = self._sampler.dispatch_index + await asyncio.to_thread( + _write_training_info, + checkpoint_path, + vars(save_state), + ) + dataloader_state = self._dataloader.state_dict() + # The spare pool and dataloader advance together under the same mutation + # cut in recovery-enabled dispatch, so preserve them in this cut too. + reserve_state = list(self._replacement_reserve) + + if save_data_plane: + if gym_commit is not None: + await self._validate_gym_checkpoint_cut( + cut, gym_commit, checkpoint_path=checkpoint_path + ) + training_owned_groups = self._buffer.training_owned_replay_groups() + if training_owned_groups: + raise RuntimeError( + "full trainer checkpoint still owns streamed training rows: " + f"groups={[group['group_id'] for group in training_owned_groups]!r}" + ) + if self._sampler.supports_buffer_checkpoint: + replay_metadata = self._buffer.metadata_state_dict( + saved_capacity=self._async_cfg.max_buffered_rollouts + ) if replay_metadata is not None: - canonical_group_ids = { - group["group_id"] for group in replay_metadata["groups"] - } - rollout_recovery_state["groups"] = [ - group - for group in rollout_recovery_state["groups"] - if group["group_id"] not in canonical_group_ids - ] - payload_buffer = io.BytesIO() - await asyncio.to_thread( - torch.save, - rollout_recovery_state, - payload_buffer, - ) - rollout_recovery_payload = payload_buffer.getvalue() - rollout_recovery_payload_sha256 = hashlib.sha256( - rollout_recovery_payload - ).hexdigest() + await self._validate_replay_inventory(replay_metadata) + + if self._rollout_recovery_enabled: + rollout_recovery_state = build_rollout_recovery_state( + self._rollout_manager.recovery_ledger, + batch_shortfall=self._batch_shortfall, + sampler_stamps_target_steps=( + self._sampler_stamps_target_steps + ), + ) + if replay_metadata is not None: + canonical_group_ids = { + group["group_id"] for group in replay_metadata["groups"] + } + rollout_recovery_state["groups"] = [ + group + for group in rollout_recovery_state["groups"] + if group["group_id"] not in canonical_group_ids + ] + payload_buffer = io.BytesIO() + await asyncio.to_thread( + torch.save, + rollout_recovery_state, + payload_buffer, + ) + rollout_recovery_payload = payload_buffer.getvalue() + rollout_recovery_payload_sha256 = hashlib.sha256( + rollout_recovery_payload + ).hexdigest() - if self._master_config.token_capture.enabled: - await self._validate_rollout_recovery_inventory( - cut, + if self._master_config.token_capture.enabled: + await self._validate_rollout_recovery_inventory( + cut, + replay_metadata=replay_metadata, + clear_unreferenced=False, + ) + + await self._save_data_plane_checkpoint( + checkpoint_path, + train_steps=save_state.current_step, + trainer_version=self._trainer_version, + current_epoch=save_state.current_epoch, replay_metadata=replay_metadata, - clear_unreferenced=False, + rollout_recovery_payload_sha256=( + rollout_recovery_payload_sha256 + ), + rollout_recovery_group_count=( + len(rollout_recovery_state["groups"]) + if rollout_recovery_state is not None + else None + ), + ) + self._last_rollout_snapshot_mutation_version = ( + self._data_plane_checkpoint_barrier.mutation_version ) - - await self._save_data_plane_checkpoint( - checkpoint_path, - train_steps=save_state.current_step, - trainer_version=self._trainer_version, - current_epoch=save_state.current_epoch, - replay_metadata=replay_metadata, - rollout_recovery_payload_sha256=(rollout_recovery_payload_sha256), - rollout_recovery_group_count=( - len(rollout_recovery_state["groups"]) - if rollout_recovery_state is not None - else None - ), - ) - self._last_rollout_snapshot_mutation_version = ( - self._data_plane_checkpoint_barrier.mutation_version - ) # Save value model if self._is_ppo: diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 45299dd5c29..997cb646510 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -57,7 +57,10 @@ ClusterConfig, ) from nemo_rl.environments.nemo_gym import should_use_nemo_gym -from nemo_rl.experience.rollout_recovery import RecoveryGranularity +from nemo_rl.experience.rollout_recovery import ( + RecoveryGranularity, + RecoveryTargetLevel, +) from nemo_rl.models.generation.vllm.config import ( VllmConfig, parse_nvfp4_pertoken_rollout, @@ -656,57 +659,95 @@ class TokenCaptureConfig(BaseModel, extra="allow"): @dataclass(frozen=True) class TaskSourceRecoveryGranularity: - """Recovery granularity selected for a prompt-group reservation. + """Resolved retry and restore policy for a prompt-group reservation. - ``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. + ``task_source`` is copied from the raw Gym row when present. ``restore_level`` + is the configured restart target. Until Gym advertises group-scoring + capabilities, ``granularity`` is derived from that target and explicit + task-source or agent overrides. """ task_source: Optional[str] granularity: RecoveryGranularity + restore_level: RecoveryTargetLevel -class RolloutRecoveryConfig(BaseModel, extra="allow"): - """Retry and restore policy for unfinished token-capture prompt groups. +_RECOVERY_TARGET_RANK = { + RecoveryTargetLevel.PREFIX: 0, + RecoveryTargetLevel.TURN: 1, + RecoveryTargetLevel.SIBLING: 2, + RecoveryTargetLevel.PROMPT_GROUP: 3, +} - ``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. +class RolloutRecoveryConfig(BaseModel, extra="allow"): + """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. The same - granularity governs failures handled in-process and after a process restart. + ``prefix`` restores Gym participant state and decoded tokens from an active + model call, ``turn`` restarts an active call from its preceding durable turn, + ``sibling`` preserves completed generations and restarts unfinished siblings, + and ``prompt_group`` restarts every sibling in an unfinished group. Overrides + may only make a particular task source or agent coarser than the global target. """ - default_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING + target_level: RecoveryTargetLevel = RecoveryTargetLevel.SIBLING # 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( + task_source_target_level_overrides: dict[str, RecoveryTargetLevel] = Field( default_factory=dict ) # 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( + agent_target_level_overrides: dict[str, RecoveryTargetLevel] = Field( default_factory=dict ) @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 {}) + def _validate_target_overrides(self) -> "RolloutRecoveryConfig": + """Reject removed keys and overrides finer than the global target.""" + removed = { + "default_granularity", + "task_granularity_overrides", + "task_source_granularity_overrides", + "agent_granularity_overrides", + "preserve_generation_prefixes", + "generation_chunk_flush_tokens", + }.intersection(self.model_extra or {}) if removed: raise ValueError( f"rollout_recovery fields {sorted(removed)!r} were replaced by " - "task_source_granularity_overrides" + "target_level and *_target_level_overrides" ) + for mapping_name, overrides in ( + ( + "task_source_target_level_overrides", + self.task_source_target_level_overrides, + ), + ("agent_target_level_overrides", self.agent_target_level_overrides), + ): + finer = sorted( + name + for name, level in overrides.items() + if _RECOVERY_TARGET_RANK[level] + < _RECOVERY_TARGET_RANK[self.target_level] + ) + if finer: + raise ValueError( + f"{mapping_name} may only force a coarser recovery level than " + f"target_level={self.target_level.value!r}; invalid={finer!r}" + ) return self + @property + def turn_checkpointing_enabled(self) -> bool: + """Whether Gym participant state is required for the requested target.""" + return self.target_level.requires_gym_checkpoint + + @property + def generation_prefix_cuts_enabled(self) -> bool: + """Whether active model calls participate in Gym's checkpoint cut.""" + return self.target_level.preserves_generation_prefix + def resolve_for_prompt( self, prompt: Mapping[str, Any] ) -> TaskSourceRecoveryGranularity: @@ -727,6 +768,7 @@ def resolve_for_prompt( 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 + target_level = self.target_level if agent_name is not None: if task_source is None: warnings.warn( @@ -736,14 +778,27 @@ def resolve_for_prompt( FutureWarning, stacklevel=2, ) - override = self.agent_granularity_overrides.get(agent_name) + override = self.agent_target_level_overrides.get(agent_name) if override is not None: - 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 TaskSourceRecoveryGranularity(task_source, override) - return TaskSourceRecoveryGranularity(task_source, self.default_granularity) + target_level = override + elif task_source is not None: + target_level = self.task_source_target_level_overrides.get( + task_source, target_level + ) + elif task_source is not None: + target_level = self.task_source_target_level_overrides.get( + task_source, target_level + ) + granularity = ( + RecoveryGranularity.PROMPT_GROUP + if target_level is RecoveryTargetLevel.PROMPT_GROUP + else RecoveryGranularity.SIBLING + ) + return TaskSourceRecoveryGranularity( + task_source=task_source, + granularity=granularity, + restore_level=target_level, + ) class RolloutCheckpointConfig(BaseModel, extra="forbid"): @@ -1350,10 +1405,14 @@ 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.task_source_granularity_overrides - or recovery_config.agent_granularity_overrides + if ( + not token_capture_config.enabled + and not recovery_config.turn_checkpointing_enabled + and ( + recovery_config.target_level is not RecoveryTargetLevel.SIBLING + or recovery_config.task_source_target_level_overrides + or recovery_config.agent_target_level_overrides + ) ): raise ValueError( "non-default rollout_recovery policies require " diff --git a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py index c79163d9e6d..ca95b9bad5f 100644 --- a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py +++ b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py @@ -42,6 +42,7 @@ RolloutCheckpointAttemptReason = Literal[ "completed", + "gym_unavailable", "invariant_error", "io_error", "missing_trainer_anchor", diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 494c21529a5..435f525ed5d 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -770,6 +770,15 @@ def _spinup_gym( policy_config = master_config.policy generation_config = policy_config["generation"] enable_router_replay = router_replay_enabled(policy_config) + token_capture = ( + master_config.token_capture.model_dump() + if master_config.token_capture.enabled + else None + ) + if token_capture is not None: + token_capture["generation_prefix_cuts_enabled"] = ( + master_config.rollout_recovery.generation_prefix_cuts_enabled + ) shard_set = build_nemo_gym_actors( master_config.env, base_urls=base_urls, @@ -777,12 +786,11 @@ def _spinup_gym( tokenizer=tokenizer, enable_router_replay=enable_router_replay, use_fastokens=bool(policy_config["tokenizer"].get("use_fastokens")), - # Ledger config rides into Gym's policy model server. - token_capture=( - master_config.token_capture.model_dump() - if master_config.token_capture.enabled - else None + turn_checkpointing_enabled=( + master_config.rollout_recovery.turn_checkpointing_enabled ), + # Ledger config rides into Gym's policy model server. + token_capture=token_capture, ) return shard_set, time.perf_counter() - t0 @@ -1092,6 +1100,36 @@ def setup_single_controller( ) data_plane_checkpointing_supported = data_plane_supports_checkpointing(dp_config) rollout_checkpoint_cfg = master_config.rollout_checkpointing + if master_config.rollout_recovery.turn_checkpointing_enabled: + recovery_target = master_config.rollout_recovery.target_level.value + if generation_config["backend"] != "vllm": + raise NotImplementedError( + f"rollout_recovery.target_level={recovery_target!r} currently " + "supports only " + "the vllm generation backend; " + f"got {generation_config['backend']!r}" + ) + if not should_use_nemo_gym(master_config): + raise ValueError( + f"rollout_recovery.target_level={recovery_target!r} requires " + "the NeMo-Gym " + "rollout path (env.should_use_nemo_gym=true)" + ) + if rollout_checkpoint_cfg.snapshot_attempt_interval_s is None: + raise ValueError( + f"rollout_recovery.target_level={recovery_target!r} requires " + "rollout_checkpointing.snapshot_attempt_interval_s" + ) + if not master_config.token_capture.enabled: + raise ValueError( + f"rollout_recovery.target_level={recovery_target!r} requires " + "token_capture.enabled=true" + ) + if rollout_checkpoint_cfg.restore_mode != "latest": + raise ValueError( + f"rollout_recovery.target_level={recovery_target!r} requires " + "rollout_checkpointing.restore_mode='latest'" + ) if ( master_config.checkpointing.get("save_data_plane") or rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None @@ -1203,6 +1241,23 @@ def setup_single_controller( # change the worker's environment. token_capture_cfg = master_config.token_capture capture_media = token_capture_cfg.enabled and processor is not None + generation_prefix_cuts_enabled = ( + master_config.rollout_recovery.generation_prefix_cuts_enabled + ) + if generation_prefix_cuts_enabled: + if generation_config["backend"] != "vllm": + raise NotImplementedError( + "generation-prefix recovery currently supports only the vLLM " + f"generation backend; got {generation_config['backend']!r}" + ) + if capture_media: + raise NotImplementedError( + "generation-prefix recovery does not yet support multimodal capture" + ) + if router_replay_enabled(master_config.policy): + raise NotImplementedError( + "generation-prefix recovery does not yet support router replay" + ) if capture_media: if generation_config["backend"] != "vllm": raise NotImplementedError( @@ -1938,6 +1993,9 @@ def _build_trainer_then_megatron_generation() -> tuple[ dp_config, token_capture_cfg.staging_partition, capture_media=capture_media, + generation_prefix_cuts_enabled=generation_prefix_cuts_enabled, + generation_cut_control_token=token_capture_cfg.control_auth_token, + generation_cut_control_timeout_s=token_capture_cfg.control_timeout_s, ) generation.set_rollout_weight_version(0) diff --git a/nemo_rl/data_plane/tq_token_sink.py b/nemo_rl/data_plane/tq_token_sink.py index 3ce8af7653e..3cbae968499 100644 --- a/nemo_rl/data_plane/tq_token_sink.py +++ b/nemo_rl/data_plane/tq_token_sink.py @@ -135,6 +135,54 @@ _MODE_TO_CODE = {"text": 0, "token_in": 1} _CODE_TO_MODE = {code: mode for mode, code in _MODE_TO_CODE.items()} +GENERATION_CUT_STAGING_PREFIX = "__generation_cut__/" + + +def generation_cut_staging_key( + checkpoint_id: str, + rollout_id: str, + model_call_id: str, + *, + chunk_sequence: int, +) -> str: + """Return an immutable key for one checkpointed generation chunk.""" + if not checkpoint_id or not rollout_id or not model_call_id: + raise ValueError("generation-cut key components must be non-empty") + if chunk_sequence < 0: + raise ValueError("generation-cut chunk sequence must be non-negative") + return ( + f"{GENERATION_CUT_STAGING_PREFIX}{checkpoint_id}/{rollout_id}/" + f"{model_call_id}/{chunk_sequence}" + ) + + +def _staging_key_matches_snapshot(key: str, snapshot: StagedCallBaseSnapshot) -> bool: + """Match a physical TQ key to the row's logical call identity.""" + if key == snapshot.staging_key: + return True + if not key.startswith(GENERATION_CUT_STAGING_PREFIX): + return False + checkpoint_and_identity = key.removeprefix(GENERATION_CUT_STAGING_PREFIX) + checkpoint_id, checkpoint_separator, identity_and_sequence = ( + checkpoint_and_identity.partition("/") + ) + if ( + checkpoint_id + and checkpoint_separator + and identity_and_sequence == snapshot.staging_key + ): + # Compatibility with cuts written before chunks had a sequence suffix. + return True + logical_key, sequence_separator, chunk_sequence = identity_and_sequence.rpartition( + "/" + ) + return bool( + checkpoint_id + and checkpoint_separator + and sequence_separator + and chunk_sequence.isdecimal() + and logical_key == snapshot.staging_key + ) class MediaMetadataIntegrityError(ValueError): @@ -402,6 +450,39 @@ def stage( record: StagedCallRecord, *, attachments: Mapping[str, Any] | None = None, + ) -> StageResult: + return self._stage_at_key( + record, + record.staging_key, + attachments=attachments, + ) + + def stage_generation_prefix( + self, + record: StagedCallRecord, + *, + checkpoint_id: str, + chunk_sequence: int, + attachments: Mapping[str, Any] | None = None, + ) -> StageResult: + """Stage an immutable active-call prefix under a checkpoint-scoped key.""" + return self._stage_at_key( + record, + generation_cut_staging_key( + checkpoint_id, + record.rollout_id, + record.model_call_id, + chunk_sequence=chunk_sequence, + ), + attachments=attachments, + ) + + def _stage_at_key( + self, + record: StagedCallRecord, + key: str, + *, + attachments: Mapping[str, Any] | None = None, ) -> StageResult: """Write the token row and its media attachments in one ``put``. @@ -414,7 +495,6 @@ def stage( # Deferred: nemo_gym is an optional extra absent in non-gym runs. from nemo_gym.token_id_capture.staging.records import StageResult - key = record.staging_key write_started = False try: if attachments is not None and not self._capture_media: @@ -877,7 +957,7 @@ def fetch_for_finalization( for index, key in enumerate(staging_keys): row = _select_row(rows, index) snapshot = _row_to_base_snapshot(row) - if snapshot.staging_key != key: + if not _staging_key_matches_snapshot(key, snapshot): raise KeyError( f"staged row identity mismatch: requested {key!r}, got {snapshot.staging_key!r}" ) diff --git a/nemo_rl/environments/gym_checkpoint_adapter.py b/nemo_rl/environments/gym_checkpoint_adapter.py new file mode 100644 index 00000000000..69ff58d23e7 --- /dev/null +++ b/nemo_rl/environments/gym_checkpoint_adapter.py @@ -0,0 +1,705 @@ +# 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. +"""Thin NeMo RL adapter over one NeMo Gym v2 checkpoint coordinator.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Iterable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, get_args + +if TYPE_CHECKING: + from nemo_gym._checkpoint.coordination import Participants, PrepareResult + from nemo_gym.server_utils import ServerClient + + +@dataclass(frozen=True) +class GymCheckpointInstance: + """Stable identity and filesystem namespace for one logical Gym actor.""" + + shard_name: str + replica_index: int + + def __post_init__(self) -> None: + if ( + not self.shard_name + or self.shard_name in {".", ".."} + or "/" in self.shard_name + or "\\" in self.shard_name + ): + raise ValueError( + "Gym checkpoint shard_name must be one safe path component" + ) + if self.replica_index < 0: + raise ValueError("Gym checkpoint replica_index must be non-negative") + + @property + def instance_id(self) -> str: + return f"{self.shard_name}/replica-{self.replica_index}" + + def live_capture_dir(self, capture_root: str | Path) -> Path: + return Path(capture_root) / "gym-instances" / self.instance_id + + def checkpoint_dir(self, checkpoint_root: str | Path) -> Path: + return Path(checkpoint_root) / "gym-instances" / self.instance_id + + +@dataclass(frozen=True) +class GymCheckpointParticipantSummary: + """Serializable discovery result; the live Gym client remains actor-local.""" + + instance_id: str + members: tuple[tuple[str, str], ...] + + +@dataclass(frozen=True) +class GymCheckpointParticipantManifest: + """Integrity metadata for one Gym participant's committed state.""" + + schema_version: int + kind: str + instance: str + checkpoint_id: str + records_file: str + records_sha256: str + record_count: int + + +GymCheckpointParticipantKind = Literal[ + "environment", + "model", + "agent", + "resources", +] + + +GYM_CHECKPOINT_PARTICIPANT_KINDS: frozenset[str] = frozenset( + get_args(GymCheckpointParticipantKind) +) + + +@dataclass(frozen=True) +class GymCheckpointParticipantRecord: + """Durable identity of one Gym participant whose state a checkpoint holds.""" + + server_name: str + kind: GymCheckpointParticipantKind + # The participant's state directory under its kind, as Gym wrote it. + instance: str + records_sha256: str + + +@dataclass(frozen=True) +class GymCheckpointParticipantCommitSummary: + """Validated commit reply from one discovered Gym participant.""" + + server_name: str + kind: GymCheckpointParticipantKind + phase: str + capture_keys: tuple[str, ...] + manifest: GymCheckpointParticipantManifest + staging_keys: tuple[str, ...] + + +@dataclass(frozen=True) +class GymCheckpointCommitSummary: + """Actor-local Gym state selected by one checkpoint commit.""" + + exported_episodes: tuple[GymCheckpointEpisode, ...] + staging_keys: tuple[str, ...] + participants: tuple[GymCheckpointParticipantCommitSummary, ...] + # Each committed episode's staged rows, as Gym's records attribute them. A + # row carried over from an earlier attempt belongs to both attempts. + staging_keys_by_episode: Mapping[GymCheckpointEpisode, tuple[str, ...]] = field( + default_factory=dict + ) + + def __post_init__(self) -> None: + attributed = { + key for keys in self.staging_keys_by_episode.values() for key in keys + } + if attributed != set(self.staging_keys): + raise ValueError( + "Gym checkpoint staging_keys_by_episode must attribute exactly the " + "committed staging keys: " + f"unattributed={sorted(set(self.staging_keys) - attributed)[:10]!r}, " + f"unlisted={sorted(attributed - set(self.staging_keys))[:10]!r}" + ) + + # Parked agent sessions that no environment or legacy agent participant owns. + # They cannot be restored and cannot finish while the deployment is prepared. + unowned_live_episodes: tuple[GymCheckpointEpisode, ...] = () + + +@dataclass(frozen=True) +class GymCheckpointEpisode: + """Serializable identity of one physical Gym episode attempt.""" + + rollout_id: str + attempt: int + + def __post_init__(self) -> None: + if not isinstance(self.rollout_id, str) or not self.rollout_id: + raise ValueError("Gym checkpoint rollout_id must not be empty") + if ( + not isinstance(self.attempt, int) + or isinstance(self.attempt, bool) + or self.attempt < 0 + ): + raise ValueError("Gym checkpoint attempt must be non-negative") + + def next_attempt(self) -> GymCheckpointEpisode: + """Return the physical episode Gym creates when this one is restored.""" + return GymCheckpointEpisode(self.rollout_id, self.attempt + 1) + + +@dataclass(frozen=True) +class GymCheckpointPrepareSummary: + """Serializable result of preparing one actor's Gym deployment.""" + + prepared: bool + blockers: dict[str, tuple[str, ...]] + + +class GymCheckpointUnavailable(RuntimeError): + """Gym's control plane could not complete one live checkpoint call. + + Raised in place of Gym's ``CoordinationError``, which cannot be pickled and + so reaches the driver as an untyped ``UnserializableException``. The + participants' failures arrive as text, so this says only that the call + failed, not why; nothing RL holds changed, and a later checkpoint may + succeed. + """ + + +@contextmanager +def _live_checkpoint_call() -> Iterator[None]: + """Type a Gym coordination failure before it crosses the Ray boundary.""" + from nemo_gym._checkpoint.coordination import CoordinationError + + try: + yield + except CoordinationError as error: + raise GymCheckpointUnavailable(str(error)) from error + + +class GymCheckpointAdapter: + """Delegate checkpoint operations for exactly one Gym deployment to Gym.""" + + def __init__( + self, + *, + instance: GymCheckpointInstance, + client: ServerClient, + auth_token: str, + ) -> None: + self._instance = instance + self._client = client + self._auth_token = auth_token + self._participants: Participants | None = None + + async def discover(self) -> GymCheckpointParticipantSummary: + from nemo_gym._checkpoint.coordination import discover + + if self._participants is None: + self._participants = await discover( + self._client, auth_token=self._auth_token + ) + participants = self._participants + return GymCheckpointParticipantSummary( + instance_id=self._instance.instance_id, + members=tuple( + (member.server_name, member.kind) for member in participants.members + ), + ) + + def _require_participants(self) -> Participants: + if self._participants is None: + raise RuntimeError( + f"Gym checkpoint instance {self._instance.instance_id!r} has not " + "discovered its deployment" + ) + return self._participants + + async def prepare(self, checkpoint_id: str, *, deadline_ts: float) -> PrepareResult: + from nemo_gym._checkpoint.coordination import prepare + + with _live_checkpoint_call(): + return await prepare( + self._require_participants(), checkpoint_id, deadline_ts=deadline_ts + ) + + async def renew(self, checkpoint_id: str, *, deadline_ts: float) -> None: + from nemo_gym._checkpoint.coordination import renew + + with _live_checkpoint_call(): + await renew( + self._require_participants(), checkpoint_id, deadline_ts=deadline_ts + ) + + @staticmethod + def _episode_ids(episodes: Iterable[GymCheckpointEpisode]) -> list[Any]: + """Convert the RL wire identity to Gym's validated API model.""" + from nemo_gym.episode_types import EpisodeId + + return [ + EpisodeId(rollout_id=episode.rollout_id, attempt=episode.attempt) + for episode in episodes + ] + + @staticmethod + def _participant_manifest( + raw: object, + *, + participant: str, + kind: GymCheckpointParticipantKind, + checkpoint_id: str, + ) -> GymCheckpointParticipantManifest: + """Validate Gym's participant manifest before it leaves the actor.""" + if not isinstance(raw, Mapping): + raise RuntimeError( + "Gym checkpoint commit returned no manifest for participant " + f"{participant!r}" + ) + schema_version = raw.get("schema_version") + manifest_kind = raw.get("kind") + instance = raw.get("instance") + manifest_checkpoint_id = raw.get("checkpoint_id") + records_file = raw.get("records_file") + records_sha256 = raw.get("records_sha256") + record_count = raw.get("record_count") + if not isinstance(schema_version, int) or isinstance(schema_version, bool): + raise RuntimeError( + "Gym checkpoint participant manifest has invalid schema_version: " + f"participant={participant!r}, schema_version={schema_version!r}" + ) + if manifest_kind != kind: + raise RuntimeError( + "Gym checkpoint participant manifest has the wrong kind: " + f"participant={participant!r}, expected={kind!r}, " + f"actual={manifest_kind!r}" + ) + if not isinstance(instance, str) or not instance: + raise RuntimeError( + "Gym checkpoint participant manifest has invalid instance: " + f"participant={participant!r}, instance={instance!r}" + ) + if ( + not isinstance(manifest_checkpoint_id, str) + or manifest_checkpoint_id != checkpoint_id + ): + raise RuntimeError( + "Gym checkpoint participant manifest has the wrong checkpoint ID: " + f"participant={participant!r}, expected={checkpoint_id!r}, " + f"actual={manifest_checkpoint_id!r}" + ) + if not isinstance(records_file, str) or not records_file: + raise RuntimeError( + "Gym checkpoint participant manifest has invalid records_file: " + f"participant={participant!r}, records_file={records_file!r}" + ) + if not isinstance(records_sha256, str) or not records_sha256: + raise RuntimeError( + "Gym checkpoint participant manifest has invalid records_sha256: " + f"participant={participant!r}, records_sha256={records_sha256!r}" + ) + if ( + not isinstance(record_count, int) + or isinstance(record_count, bool) + or record_count < 0 + ): + raise RuntimeError( + "Gym checkpoint participant manifest has invalid record_count: " + f"participant={participant!r}, record_count={record_count!r}" + ) + return GymCheckpointParticipantManifest( + schema_version=schema_version, + kind=kind, + instance=instance, + checkpoint_id=manifest_checkpoint_id, + records_file=records_file, + records_sha256=records_sha256, + record_count=record_count, + ) + + @staticmethod + def _legacy_agent_capture_keys( + checkpoint_dir: Path, + *, + participant: str, + manifest: GymCheckpointParticipantManifest, + ) -> frozenset[str]: + """Return the legacy ``/run`` episodes one agent participant committed. + + An agent exports every parked session, but only a legacy ``/run`` session + (a record with a non-null ``episode`` step boundary) owns its episode. A + native session runs under an environment participant, which owns it. + """ + from nemo_gym._checkpoint.store import read_participant_state + from nemo_gym.episode_types import EpisodeId + + stored, records = read_participant_state( + checkpoint_dir, kind="agent", instance=manifest.instance + ) + if stored.get("records_sha256") != manifest.records_sha256: + raise RuntimeError( + "Gym checkpoint agent records on disk do not match the commit " + f"reply: participant={participant!r}, " + f"expected={manifest.records_sha256!r}, " + f"actual={stored.get('records_sha256')!r}" + ) + return frozenset( + EpisodeId.model_validate(record["episode_id"]).capture_key + for record in records + if record.get("episode") is not None + ) + + def _commit_summary( + self, + checkpoint_id: str, + episode_ids: list[Any], + replies: dict[str, dict[str, Any]], + *, + checkpoint_dir: Path, + ) -> GymCheckpointCommitSummary: + """Convert participant replies into an actor-local ownership summary.""" + participants = self._require_participants() + expected_participants = { + participant.server_name for participant in participants.members + } + if set(replies) != expected_participants: + raise RuntimeError( + "Gym checkpoint commit returned an unexpected participant set: " + f"expected={sorted(expected_participants)!r}, " + f"actual={sorted(replies)!r}" + ) + + requested_keys = {episode_id.capture_key for episode_id in episode_ids} + exported_environment_keys: set[str] = set() + # A legacy_agent environment server only relays /run, so it never holds + # the episode; the agent's legacy session owns it instead. + legacy_agent_keys: set[str] = set() + agent_session_keys: set[str] = set() + staging_keys: set[str] = set() + keys_by_capture_key: dict[str, set[str]] = {} + participant_summaries: list[GymCheckpointParticipantCommitSummary] = [] + for member in participants.members: + reply = replies.get(member.server_name) + if not isinstance(reply, Mapping): + raise RuntimeError( + "Gym checkpoint commit returned an invalid reply for participant " + f"{member.server_name!r}" + ) + phase = reply.get("phase") + if phase != "committed": + raise RuntimeError( + "Gym checkpoint participant did not enter committed phase: " + f"participant={member.server_name!r}, phase={phase!r}" + ) + exported = reply.get("episode_ids") + if not isinstance(exported, list) or not all( + isinstance(capture_key, str) for capture_key in exported + ): + raise RuntimeError( + "Gym checkpoint commit returned invalid episode_ids for " + f"participant {member.server_name!r}" + ) + capture_keys = tuple(sorted(set(exported))) + if len(capture_keys) != len(exported): + raise RuntimeError( + "Gym checkpoint commit returned duplicate episode IDs for " + f"participant {member.server_name!r}" + ) + unexpected = set(capture_keys) - requested_keys + if unexpected: + raise RuntimeError( + "Gym checkpoint participant exported episodes outside the " + f"requested scope: participant={member.server_name!r}, " + f"unexpected={sorted(unexpected)!r}" + ) + manifest = self._participant_manifest( + reply.get("manifest"), + participant=member.server_name, + kind=member.kind, + checkpoint_id=checkpoint_id, + ) + if member.kind == "agent" and capture_keys: + agent_session_keys.update(capture_keys) + legacy_keys = self._legacy_agent_capture_keys( + checkpoint_dir, + participant=member.server_name, + manifest=manifest, + ) + unreported = legacy_keys - set(capture_keys) + if unreported: + raise RuntimeError( + "Gym checkpoint agent records name episodes absent from its " + f"commit reply: participant={member.server_name!r}, " + f"unreported={sorted(unreported)!r}" + ) + duplicate_owners = legacy_agent_keys.intersection(legacy_keys) + if duplicate_owners: + raise RuntimeError( + "Gym checkpoint legacy episode was exported by multiple agent " + f"participants: {sorted(duplicate_owners)!r}" + ) + legacy_agent_keys.update(legacy_keys) + if member.kind == "environment": + duplicate_owners = exported_environment_keys.intersection(capture_keys) + if duplicate_owners: + raise RuntimeError( + "Gym checkpoint episode was exported by multiple environment " + f"participants: {sorted(duplicate_owners)!r}" + ) + exported_environment_keys.update(capture_keys) + + participant_keys: tuple[str, ...] = () + if member.kind == "model": + raw_staging_keys = reply.get("staging_keys") + if not isinstance(raw_staging_keys, list) or not all( + isinstance(key, str) for key in raw_staging_keys + ): + raise RuntimeError( + "Gym checkpoint policy model returned invalid staging_keys: " + f"participant={member.server_name!r}, " + f"staging_keys={raw_staging_keys!r}" + ) + participant_keys = tuple(sorted(set(raw_staging_keys))) + if len(participant_keys) != len(raw_staging_keys): + raise RuntimeError( + "Gym checkpoint policy model returned duplicate staging_keys: " + f"participant={member.server_name!r}" + ) + staging_keys.update(participant_keys) + raw_by_episode = reply.get("staging_keys_by_episode") + if not isinstance(raw_by_episode, dict) or not all( + isinstance(capture_key, str) + and isinstance(keys, list) + and all(isinstance(key, str) for key in keys) + for capture_key, keys in raw_by_episode.items() + ): + raise RuntimeError( + "Gym checkpoint policy model returned invalid " + f"staging_keys_by_episode: participant={member.server_name!r}" + ) + attributed = {key for keys in raw_by_episode.values() for key in keys} + if attributed != set(participant_keys): + raise RuntimeError( + "Gym checkpoint policy model staging_keys_by_episode does not " + "attribute exactly its staging_keys: " + f"participant={member.server_name!r}, " + f"unattributed={sorted(set(participant_keys) - attributed)[:10]!r}, " + f"unlisted={sorted(attributed - set(participant_keys))[:10]!r}" + ) + for capture_key, keys in raw_by_episode.items(): + keys_by_capture_key.setdefault(capture_key, set()).update(keys) + + participant_summaries.append( + GymCheckpointParticipantCommitSummary( + server_name=member.server_name, + kind=member.kind, + phase=phase, + capture_keys=capture_keys, + manifest=manifest, + staging_keys=participant_keys, + ) + ) + + from nemo_gym.episode_types import EpisodeId + + def _episodes(capture_keys: set[str]) -> tuple[GymCheckpointEpisode, ...]: + return tuple( + GymCheckpointEpisode(parsed.rollout_id, parsed.attempt) + for parsed in ( + EpisodeId.from_capture_key(capture_key) + for capture_key in sorted(capture_keys) + ) + ) + + # An environment participant owns every episode it exported; a legacy + # agent session owns its episode only when no environment does. + owned_keys = exported_environment_keys | legacy_agent_keys + return GymCheckpointCommitSummary( + exported_episodes=_episodes(owned_keys), + staging_keys=tuple(sorted(staging_keys)), + participants=tuple(participant_summaries), + staging_keys_by_episode={ + GymCheckpointEpisode(parsed.rollout_id, parsed.attempt): tuple( + sorted(keys) + ) + for parsed, keys in ( + (EpisodeId.from_capture_key(capture_key), keys) + for capture_key, keys in keys_by_capture_key.items() + ) + }, + unowned_live_episodes=_episodes(agent_session_keys - owned_keys), + ) + + async def retire( + self, + checkpoint_id: str, + episodes: Iterable[GymCheckpointEpisode], + *, + deadline_ts: float, + ) -> None: + from nemo_gym._checkpoint.coordination import retire + + with _live_checkpoint_call(): + await retire( + self._require_participants(), + checkpoint_id, + self._episode_ids(episodes), + deadline_ts=deadline_ts, + ) + + async def commit( + self, + checkpoint_id: str, + checkpoint_root: str | Path, + episodes: Iterable[GymCheckpointEpisode], + *, + deadline_ts: float, + ) -> GymCheckpointCommitSummary: + from nemo_gym._checkpoint.coordination import commit + + episode_ids = self._episode_ids(episodes) + checkpoint_dir = self._instance.checkpoint_dir(checkpoint_root) + with _live_checkpoint_call(): + replies = await commit( + self._require_participants(), + checkpoint_id, + str(checkpoint_dir), + episode_ids, + deadline_ts=deadline_ts, + ) + # Reads committed agent records from disk, so keep it off the event loop. + return await asyncio.to_thread( + self._commit_summary, + checkpoint_id, + episode_ids, + replies, + checkpoint_dir=checkpoint_dir, + ) + + def _require_restorable_topology( + self, + recorded: Iterable[GymCheckpointParticipantRecord], + checkpoint_dir: Path, + ) -> None: + """Require this deployment to be exactly the one the checkpoint recorded. + + Gym restores each live participant from its own state, so a participant + missing from this deployment would have its state silently dropped, + and one the checkpoint never saw would start empty. Each recorded + participant's state on disk must also be the state that was committed. + """ + from nemo_gym._checkpoint.errors import CheckpointStateError + from nemo_gym._checkpoint.store import read_participant_state + + expected = {record.server_name: record for record in recorded} + live = { + member.server_name: member.kind + for member in self._require_participants().members + } + missing = sorted(set(expected) - set(live)) + unexpected = sorted(set(live) - set(expected)) + changed_kind = sorted( + name + for name in set(expected) & set(live) + if expected[name].kind != live[name] + ) + if missing or unexpected or changed_kind: + raise RuntimeError( + "Gym deployment does not match the participants this checkpoint " + f"recorded for instance {self._instance.instance_id!r}: " + f"missing={missing!r}, unexpected={unexpected!r}, " + f"changed_kind={changed_kind!r}" + ) + for record in expected.values(): + try: + stored, _ = read_participant_state( + checkpoint_dir, kind=record.kind, instance=record.instance + ) + except (CheckpointStateError, OSError, ValueError) as error: + raise RuntimeError( + "Gym checkpoint state is unreadable for participant " + f"{record.server_name!r}: {error}" + ) from error + if stored.get("records_sha256") != record.records_sha256: + raise RuntimeError( + "Gym checkpoint state for participant " + f"{record.server_name!r} is not the state that was committed: " + f"expected records_sha256={record.records_sha256!r}, " + f"found={stored.get('records_sha256')!r}" + ) + + async def restore( + self, + checkpoint_id: str, + checkpoint_root: str | Path, + episodes: Iterable[GymCheckpointEpisode], + *, + participants: Iterable[GymCheckpointParticipantRecord], + source_checkpoint_id: str, + deadline_ts: float, + ) -> None: + from nemo_gym._checkpoint.coordination import restore + + checkpoint_dir = self._instance.checkpoint_dir(checkpoint_root) + # Checked before Gym installs anything, so a mismatch leaves it untouched. + await asyncio.to_thread( + self._require_restorable_topology, tuple(participants), checkpoint_dir + ) + replies = await restore( + self._require_participants(), + checkpoint_id, + str(checkpoint_dir), + self._episode_ids(episodes), + deadline_ts=deadline_ts, + ) + expected_participants = { + participant.server_name + for participant in self._require_participants().members + } + if set(replies) != expected_participants: + raise RuntimeError( + "Gym checkpoint restore returned an unexpected participant set: " + f"expected={sorted(expected_participants)!r}, " + f"actual={sorted(replies)!r}" + ) + mismatched_sources: dict[str, object] = {} + for participant, reply in replies.items(): + if not isinstance(reply, Mapping): + mismatched_sources[participant] = type(reply).__name__ + continue + actual_source = reply.get("source_checkpoint_id") + if actual_source != source_checkpoint_id: + mismatched_sources[participant] = actual_source + if mismatched_sources: + raise RuntimeError( + "Gym checkpoint restore loaded participant state from the wrong " + f"checkpoint: expected={source_checkpoint_id!r}, " + f"actual={mismatched_sources!r}" + ) + + async def resume(self, checkpoint_id: str, *, deadline_ts: float) -> None: + from nemo_gym._checkpoint.coordination import resume + + with _live_checkpoint_call(): + await resume( + self._require_participants(), checkpoint_id, deadline_ts=deadline_ts + ) diff --git a/nemo_rl/environments/gym_checkpoint_coordinator.py b/nemo_rl/environments/gym_checkpoint_coordinator.py new file mode 100644 index 00000000000..967d3f14a78 --- /dev/null +++ b/nemo_rl/environments/gym_checkpoint_coordinator.py @@ -0,0 +1,740 @@ +# 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. + +"""NeMo RL coordination across independent Gym checkpoint deployments.""" + +from __future__ import annotations + +import asyncio +import json +import os +import re +import time +from collections.abc import AsyncIterator, Iterable, Mapping +from contextlib import asynccontextmanager +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import ray + +from nemo_rl.environments.gym_checkpoint_adapter import ( + GYM_CHECKPOINT_PARTICIPANT_KINDS, + GymCheckpointCommitSummary, + GymCheckpointEpisode, + GymCheckpointParticipantRecord, + GymCheckpointPrepareSummary, + GymCheckpointUnavailable, +) +from nemo_rl.environments.nemo_gym import NemoGymShardSet +from nemo_rl.experience.failures import FailureClass, classify_rollout_failure + +GYM_CHECKPOINT_MANIFEST_FILENAME = "gym_checkpoint.json" +GYM_CHECKPOINT_SCHEMA_VERSION = 2 +# Upper bound on how long past a call's deadline RL waits for Gym's reply. Gym +# answers by the deadline once it runs the call, so this only covers the reply +# crossing Ray; a call still unanswered after it was never run. +_MAX_REPLY_GRACE_S = 5.0 +_SHA256_PATTERN = re.compile(r"[0-9a-f]{64}") + + +class GymCheckpointOperationError(RuntimeError): + """One or more independent Gym deployments failed one operation.""" + + def __init__(self, operation: str, failures: Mapping[str, BaseException]) -> None: + self.operation = operation + self.failures = dict(failures) + detail = "; ".join( + f"{instance}: {type(error).__name__}: {error}" + for instance, error in sorted(self.failures.items()) + ) + super().__init__( + f"Gym checkpoint {operation} failed on " + f"{len(self.failures)} instance(s): {detail}" + ) + + +class GymCheckpointNotReady(RuntimeError): + """Gym participants could not all park before the checkpoint deadline.""" + + def __init__( + self, + blockers: Mapping[str, Mapping[str, tuple[str, ...]]], + ) -> None: + self.blockers = { + instance: dict(instance_blockers) + for instance, instance_blockers in blockers.items() + } + super().__init__( + f"Gym checkpoint participants are not ready: {self.blockers!r}" + ) + + +def is_transient_gym_checkpoint_failure(error: BaseException) -> bool: + """Whether a failed Gym checkpoint may succeed if it is simply tried again. + + Participants that could not park in time, calls that went unanswered, and + actors or control planes that failed a call leave nothing behind: the + coordinator resumes Gym, and Gym's lease reopens it if that resume is lost + too. Anything else, such as a malformed reply or an ownership or topology + mismatch, is a broken invariant that a retry would only hide. + """ + if isinstance(error, GymCheckpointNotReady): + return True + if isinstance(error, GymCheckpointOperationError): + return bool(error.failures) and all( + isinstance(failure, GymCheckpointUnavailable) + or classify_rollout_failure(failure) is FailureClass.INFRA + for failure in error.failures.values() + ) + return False + + +@dataclass(frozen=True) +class GymCheckpointManifest: + """Durable mapping from Gym actor instances to continued episodes.""" + + schema_version: int + checkpoint_id: str + instances: dict[str, tuple[GymCheckpointEpisode, ...]] + staging_keys: dict[str, tuple[str, ...]] + # The Gym participants whose state each instance committed. Restore requires + # the live deployment to match, so no participant's state is silently lost. + participants: dict[str, tuple[GymCheckpointParticipantRecord, ...]] + + def __post_init__(self) -> None: + if self.schema_version != GYM_CHECKPOINT_SCHEMA_VERSION: + raise ValueError( + f"unsupported Gym checkpoint schema version {self.schema_version!r}" + ) + if not self.checkpoint_id: + raise ValueError("Gym checkpoint_id must not be empty") + if not self.instances: + raise ValueError( + "Gym checkpoint manifest must contain at least one instance" + ) + for instance_id, episodes in self.instances.items(): + if not instance_id: + raise ValueError("Gym checkpoint instance ID must not be empty") + episode_keys = [ + (episode.rollout_id, episode.attempt) for episode in episodes + ] + if len(episode_keys) != len(set(episode_keys)): + raise ValueError( + "Gym checkpoint instance " + f"{instance_id!r} contains duplicate episodes" + ) + if set(self.staging_keys) != set(self.instances): + raise ValueError( + "Gym checkpoint staging-key topology must match its episode topology" + ) + for instance_id, keys in self.staging_keys.items(): + if any(not key for key in keys): + raise ValueError( + "Gym checkpoint instance " + f"{instance_id!r} contains an empty staging key" + ) + # pyrefly widens the tuple ``keys`` to Iterable inside this loop. + if len(keys) != len(set(keys)): # pyrefly: ignore[bad-argument-type] + raise ValueError( + "Gym checkpoint instance " + f"{instance_id!r} contains duplicate staging keys" + ) + + if set(self.participants) != set(self.instances): + raise ValueError( + "Gym checkpoint participant topology must match its episode topology" + ) + for instance_id, records in self.participants.items(): + names = [record.server_name for record in records] + if len(names) != len(set(names)): + raise ValueError( + "Gym checkpoint instance " + f"{instance_id!r} contains duplicate participants" + ) + for record in records: + if not record.server_name or not record.instance: + raise ValueError( + "Gym checkpoint participant " + f"{record!r} must name its server and state directory" + ) + if record.kind not in GYM_CHECKPOINT_PARTICIPANT_KINDS: + raise ValueError( + "Gym checkpoint participant " + f"{record.server_name!r} has unknown kind {record.kind!r}" + ) + if not _SHA256_PATTERN.fullmatch(record.records_sha256): + raise ValueError( + "Gym checkpoint participant " + f"{record.server_name!r} has an invalid records_sha256" + ) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "checkpoint_id": self.checkpoint_id, + "instances": { + instance_id: [asdict(episode) for episode in episodes] + for instance_id, episodes in sorted(self.instances.items()) + }, + "staging_keys": { + instance_id: list(keys) + for instance_id, keys in sorted(self.staging_keys.items()) + }, + "participants": { + instance_id: [ + asdict(record) + for record in sorted(records, key=lambda item: item.server_name) + ] + for instance_id, records in sorted(self.participants.items()) + }, + } + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any]) -> GymCheckpointManifest: + unknown = set(raw) - { + "schema_version", + "checkpoint_id", + "instances", + "staging_keys", + "participants", + } + if unknown: + raise ValueError( + f"Gym checkpoint manifest contains unknown fields: {sorted(unknown)!r}" + ) + schema_version = raw.get("schema_version") + checkpoint_id = raw.get("checkpoint_id") + raw_instances = raw.get("instances") + raw_staging_keys = raw.get("staging_keys") + if not isinstance(schema_version, int) or isinstance(schema_version, bool): + raise ValueError("Gym checkpoint schema_version must be an integer") + if not isinstance(checkpoint_id, str): + raise ValueError("Gym checkpoint checkpoint_id must be a string") + if not isinstance(raw_instances, dict): + raise ValueError("Gym checkpoint instances must be a mapping") + if not isinstance(raw_staging_keys, dict): + raise ValueError("Gym checkpoint staging_keys must be a mapping") + + instances: dict[str, tuple[GymCheckpointEpisode, ...]] = {} + for instance_id, raw_episodes in raw_instances.items(): + if not isinstance(instance_id, str): + raise ValueError("Gym checkpoint instance IDs must be strings") + if not isinstance(raw_episodes, list): + raise ValueError( + f"Gym checkpoint episodes for {instance_id!r} must be a list" + ) + episodes = [] + for raw_episode in raw_episodes: + if not isinstance(raw_episode, dict): + raise ValueError("Gym checkpoint episode must be a mapping") + if set(raw_episode) != {"rollout_id", "attempt"}: + raise ValueError( + "Gym checkpoint episode must contain exactly rollout_id " + "and attempt" + ) + rollout_id = raw_episode["rollout_id"] + attempt = raw_episode["attempt"] + if not isinstance(rollout_id, str): + raise ValueError( + "Gym checkpoint episode rollout_id must be a string" + ) + if not isinstance(attempt, int) or isinstance(attempt, bool): + raise ValueError( + "Gym checkpoint episode attempt must be an integer" + ) + episodes.append(GymCheckpointEpisode(rollout_id, attempt)) + instances[instance_id] = tuple(episodes) + staging_keys: dict[str, tuple[str, ...]] = {} + for instance_id, raw_keys in raw_staging_keys.items(): + if not isinstance(instance_id, str): + raise ValueError( + "Gym checkpoint staging-key instance IDs must be strings" + ) + if not isinstance(raw_keys, list) or not all( + isinstance(key, str) for key in raw_keys + ): + raise ValueError( + f"Gym checkpoint staging keys for {instance_id!r} must be strings" + ) + staging_keys[instance_id] = tuple(raw_keys) + raw_participants = raw.get("participants") + if not isinstance(raw_participants, dict): + raise ValueError("Gym checkpoint participants must be a mapping") + participants: dict[str, tuple[GymCheckpointParticipantRecord, ...]] = {} + record_fields = {"server_name", "kind", "instance", "records_sha256"} + for instance_id, raw_records in raw_participants.items(): + if not isinstance(instance_id, str): + raise ValueError( + "Gym checkpoint participant instance IDs must be strings" + ) + if not isinstance(raw_records, list): + raise ValueError( + f"Gym checkpoint participants for {instance_id!r} must be a list" + ) + records = [] + for raw_record in raw_records: + if not isinstance(raw_record, dict) or set(raw_record) != record_fields: + raise ValueError( + "Gym checkpoint participant must contain exactly " + f"{sorted(record_fields)!r}" + ) + if not all(isinstance(value, str) for value in raw_record.values()): + raise ValueError( + "Gym checkpoint participant fields must be strings" + ) + records.append(GymCheckpointParticipantRecord(**raw_record)) + participants[instance_id] = tuple(records) + return cls( + schema_version=schema_version, + checkpoint_id=checkpoint_id, + instances=instances, + staging_keys=staging_keys, + participants=participants, + ) + + +@dataclass(frozen=True) +class GymCheckpointCommitResult: + """Outer manifest plus each Gym actor's validated commit summary.""" + + manifest: GymCheckpointManifest + instances: dict[str, GymCheckpointCommitSummary] + + +def staging_keys_of( + summary: GymCheckpointCommitSummary, episodes: Iterable[GymCheckpointEpisode] +) -> tuple[str, ...]: + """Return the staged rows of ``episodes``, the only rows a restore continues. + + Gym attributes each row to the episodes whose records hold it, so a + generation-cut row needs no parsing and a carried-over row stays while any + attempt that uses it is kept. + """ + by_episode = summary.staging_keys_by_episode + return tuple( + sorted({key for episode in episodes for key in by_episode.get(episode, ())}) + ) + + +def load_gym_checkpoint_manifest(checkpoint_root: str | Path) -> GymCheckpointManifest: + """Load and validate the Gym manifest from one committed RL snapshot.""" + manifest_path = Path(checkpoint_root) / GYM_CHECKPOINT_MANIFEST_FILENAME + raw = json.loads(manifest_path.read_text()) + if not isinstance(raw, Mapping): + raise ValueError("Gym checkpoint manifest must contain a mapping") + return GymCheckpointManifest.from_mapping(raw) + + +def write_gym_checkpoint_manifest( + checkpoint_root: str | Path, + manifest: GymCheckpointManifest, +) -> None: + """Atomically write a Gym manifest inside an unpublished RL snapshot.""" + root = Path(checkpoint_root) + path = root / GYM_CHECKPOINT_MANIFEST_FILENAME + tmp_path = path.with_suffix(".json.tmp") + tmp_path.write_text(json.dumps(manifest.to_dict(), sort_keys=True, indent=2) + "\n") + os.replace(tmp_path, path) + + +class GymCheckpointCoordinator: + """Drive one Gym checkpoint operation across every NeMo-Gym actor.""" + + def __init__( + self, + shard_set: NemoGymShardSet, + *, + control_timeout_s: float, + ) -> None: + if control_timeout_s <= 0: + raise ValueError("Gym checkpoint control_timeout_s must be positive") + self._handles = shard_set.checkpoint_handles + if not self._handles: + raise ValueError("Gym checkpoint coordinator requires at least one actor") + self._control_timeout_s = control_timeout_s + + @property + def instance_ids(self) -> frozenset[str]: + return frozenset(self._handles) + + @property + def control_timeout_s(self) -> float: + """Maximum wall-clock budget for one coordinated control phase.""" + return self._control_timeout_s + + def _deadline(self) -> float: + return time.time() + self._control_timeout_s + + async def _collect( + self, + operation: str, + references: Mapping[str, Any], + *, + deadline_ts: float, + ) -> dict[str, Any]: + """Wait for every actor's reply, but never past the call's deadline. + + Gym enforces ``deadline_ts`` only once it runs the call. An actor whose + event loop is stuck, or a call Ray never schedules, would otherwise be + waited on forever. Giving up is safe: a Gym deployment left parked + reopens by itself when its checkpoint lease expires. + """ + tasks = { + instance_id: asyncio.ensure_future(reference) + for instance_id, reference in references.items() + } + grace_s = min(self._control_timeout_s, _MAX_REPLY_GRACE_S) + timeout_s = max(0.0, deadline_ts - time.time()) + grace_s + try: + await asyncio.wait(tasks.values(), timeout=timeout_s) + except BaseException: + for task in tasks.values(): + task.cancel() + raise + results: dict[str, Any] = {} + failures: dict[str, BaseException] = {} + for instance_id, task in tasks.items(): + if not task.done(): + task.cancel() + _cancel_remote(references[instance_id]) + failures[instance_id] = TimeoutError( + f"no reply within {timeout_s:.1f}s of the {operation} call" + ) + elif task.cancelled(): + failures[instance_id] = asyncio.CancelledError() + elif task.exception() is not None: + failures[instance_id] = task.exception() + else: + results[instance_id] = task.result() + if failures: + raise GymCheckpointOperationError(operation, failures) + return results + + async def prepare( + self, + checkpoint_id: str, + ) -> dict[str, GymCheckpointPrepareSummary]: + deadline_ts = self._deadline() + results = await self._collect( + "prepare", + { + instance_id: handle.checkpoint_prepare.remote( + checkpoint_id, + deadline_ts=deadline_ts, + ) + for instance_id, handle in self._handles.items() + }, + deadline_ts=deadline_ts, + ) + summaries = {instance_id: result for instance_id, result in results.items()} + blockers = { + instance_id: summary.blockers + for instance_id, summary in summaries.items() + if not summary.prepared + } + if blockers: + raise GymCheckpointNotReady(blockers) + return summaries + + async def renew(self, checkpoint_id: str) -> None: + deadline_ts = self._deadline() + await self._collect( + "renew", + { + instance_id: handle.checkpoint_renew.remote( + checkpoint_id, + deadline_ts=deadline_ts, + ) + for instance_id, handle in self._handles.items() + }, + deadline_ts=deadline_ts, + ) + + def _validate_inventory( + self, + episodes_by_instance: Mapping[str, tuple[GymCheckpointEpisode, ...]], + ) -> dict[str, tuple[GymCheckpointEpisode, ...]]: + actual = set(episodes_by_instance) + expected = set(self._handles) + if actual != expected: + raise ValueError( + "Gym checkpoint inventory does not match the live actor topology: " + f"missing={sorted(expected - actual)!r}, " + f"unknown={sorted(actual - expected)!r}" + ) + return { + instance_id: tuple(episodes_by_instance[instance_id]) + for instance_id in self._handles + } + + async def commit( + self, + checkpoint_id: str, + checkpoint_root: str | Path, + episodes_by_instance: Mapping[str, tuple[GymCheckpointEpisode, ...]], + ) -> GymCheckpointCommitResult: + inventory = self._validate_inventory(episodes_by_instance) + deadline_ts = self._deadline() + raw_summaries = await self._collect( + "commit", + { + instance_id: self._handles[instance_id].checkpoint_commit.remote( + checkpoint_id, + str(checkpoint_root), + episodes, + deadline_ts=deadline_ts, + ) + for instance_id, episodes in inventory.items() + }, + deadline_ts=deadline_ts, + ) + summaries: dict[str, GymCheckpointCommitSummary] = {} + for instance_id, summary in raw_summaries.items(): + if not isinstance(summary, GymCheckpointCommitSummary): + raise TypeError( + "Gym checkpoint actor returned an invalid commit summary: " + f"instance={instance_id!r}, type={type(summary).__name__}" + ) + summaries[instance_id] = summary + expected = set(inventory[instance_id]) + actual = set(summary.exported_episodes) + unexpected = actual - expected + if unexpected: + raise RuntimeError( + "Gym checkpoint actor exported an episode outside the " + "candidate set: " + f"instance={instance_id!r}, " + f"unexpected={sorted(unexpected, key=repr)!r}" + ) + manifest = GymCheckpointManifest( + schema_version=GYM_CHECKPOINT_SCHEMA_VERSION, + checkpoint_id=checkpoint_id, + instances={ + instance_id: summary.exported_episodes + for instance_id, summary in summaries.items() + }, + # Gym's model participant reports the rows of every candidate in + # scope, including one whose reply was already on the wire. Only + # saved episodes are restored; RL's ledger owns the rest once sealed. + staging_keys={ + instance_id: staging_keys_of(summary, summary.exported_episodes) + for instance_id, summary in summaries.items() + }, + participants={ + instance_id: tuple( + GymCheckpointParticipantRecord( + server_name=participant.server_name, + kind=participant.kind, + instance=participant.manifest.instance, + records_sha256=participant.manifest.records_sha256, + ) + for participant in sorted( + summary.participants, key=lambda item: item.server_name + ) + ) + for instance_id, summary in summaries.items() + }, + ) + await asyncio.to_thread( + write_gym_checkpoint_manifest, + checkpoint_root, + manifest, + ) + return GymCheckpointCommitResult( + manifest=manifest, + instances=summaries, + ) + + async def retire( + self, + checkpoint_id: str, + episodes_by_instance: Mapping[str, tuple[GymCheckpointEpisode, ...]], + ) -> None: + """Retire selected physical attempts from every Gym deployment.""" + inventory = self._validate_inventory(episodes_by_instance) + deadline_ts = self._deadline() + await self._collect( + "retire", + { + instance_id: self._handles[instance_id].checkpoint_retire.remote( + checkpoint_id, + episodes, + deadline_ts=deadline_ts, + ) + for instance_id, episodes in inventory.items() + }, + deadline_ts=deadline_ts, + ) + + async def retire_restored( + self, + restore_id: str, + manifest: GymCheckpointManifest, + ) -> None: + """Retire the replacement attempts installed by a successful restore.""" + await self.retire( + restore_id, + { + instance_id: tuple(episode.next_attempt() for episode in episodes) + for instance_id, episodes in manifest.instances.items() + }, + ) + + async def restore( + self, + restore_id: str, + checkpoint_root: str | Path, + manifest: GymCheckpointManifest, + ) -> None: + inventory = self._validate_inventory(manifest.instances) + deadline_ts = self._deadline() + try: + await self._collect( + "restore", + { + instance_id: self._handles[instance_id].checkpoint_restore.remote( + restore_id, + str(checkpoint_root), + episodes, + participants=manifest.participants[instance_id], + source_checkpoint_id=manifest.checkpoint_id, + deadline_ts=deadline_ts, + ) + for instance_id, episodes in inventory.items() + }, + deadline_ts=deadline_ts, + ) + except GymCheckpointOperationError as error: + # Gym makes one deployment's restore atomic. RL must extend that + # guarantee across independent Gym actors. Retire replacement + # attempts on every actor before reopening them: a failed Ray reply + # or actor-local validation error is an uncertain outcome, so the + # remote restore may have succeeded before the failure reached this + # coordinator. + cleanup_errors: list[BaseException] = [] + try: + cleanup_deadline_ts = self._deadline() + await self._collect( + "restore_retire", + { + instance_id: self._handles[ + instance_id + ].checkpoint_retire.remote( + restore_id, + tuple( + episode.next_attempt() + for episode in inventory[instance_id] + ), + deadline_ts=cleanup_deadline_ts, + ) + for instance_id in inventory + }, + deadline_ts=cleanup_deadline_ts, + ) + except BaseException as cleanup_error: + cleanup_errors.append(cleanup_error) + try: + await self.resume(restore_id) + except BaseException as cleanup_error: + cleanup_errors.append(cleanup_error) + for cleanup_error in cleanup_errors: + error.add_note( + "Gym cross-instance restore cleanup also failed: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + raise + + async def resume(self, checkpoint_id: str) -> None: + deadline_ts = self._deadline() + await self._collect( + "resume", + { + instance_id: handle.checkpoint_resume.remote( + checkpoint_id, + deadline_ts=deadline_ts, + ) + for instance_id, handle in self._handles.items() + }, + deadline_ts=deadline_ts, + ) + + async def _renew_until_stopped( + self, + checkpoint_id: str, + stop: asyncio.Event, + ) -> None: + """Keep a prepared Gym deployment parked during a slow outer save.""" + interval_s = max(0.01, min(self._control_timeout_s / 3.0, 10.0)) + while True: + try: + await asyncio.wait_for(stop.wait(), timeout=interval_s) + except TimeoutError: + await self.renew(checkpoint_id) + else: + return + + @asynccontextmanager + async def prepared(self, checkpoint_id: str) -> AsyncIterator[None]: + """Prepare every deployment, renew its lease, and always release it.""" + try: + await self.prepare(checkpoint_id) + # A prepare may use most of its original deadline. Refresh once before + # the outer controller begins disk I/O, then keep the lease alive. + await self.renew(checkpoint_id) + stop = asyncio.Event() + renew_task = asyncio.create_task( + self._renew_until_stopped(checkpoint_id, stop) + ) + body_error: BaseException | None = None + try: + yield + except BaseException as error: + body_error = error + raise + finally: + stop.set() + (renew_result,) = await asyncio.gather( + renew_task, + return_exceptions=True, + ) + if isinstance(renew_result, BaseException): + if body_error is None: + raise renew_result + body_error.add_note( + "Gym checkpoint lease renewal also failed: " + f"{type(renew_result).__name__}: {renew_result}" + ) + except BaseException as error: + try: + await self.resume(checkpoint_id) + except Exception as resume_error: + error.add_note( + "Gym checkpoint cleanup also failed: " + f"{type(resume_error).__name__}: {resume_error}" + ) + raise + else: + await self.resume(checkpoint_id) + + +def _cancel_remote(reference: Any) -> None: + """Ask Ray to stop a control call RL stopped waiting for; best effort.""" + if not isinstance(reference, ray.ObjectRef): + return + try: + ray.cancel(reference) + except Exception: + # The lease still reopens Gym; a failed cancel only leaves the call queued. + pass diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index ddd5723d324..f904597d79d 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -52,6 +52,15 @@ _get_free_port_local, _get_node_ip_local, ) +from nemo_rl.environments.gym_checkpoint_adapter import ( + GymCheckpointAdapter, + GymCheckpointCommitSummary, + GymCheckpointEpisode, + GymCheckpointInstance, + GymCheckpointParticipantRecord, + GymCheckpointParticipantSummary, + GymCheckpointPrepareSummary, +) from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.environments.nemo_gym_multimodal import ( _index_per_turn_images, @@ -73,6 +82,7 @@ ) from nemo_rl.environments.utils import shutdown_environments from nemo_rl.experience.failures import ( + GymAdmissionClosed, GymTransportError, RolloutDataFailure, http_status_is_infra, @@ -281,6 +291,75 @@ def _typed_gym_failure(error: Exception) -> Optional[Exception]: return RolloutDataFailure(detail) +def _gym_error_code(response_body: object) -> Optional[str]: + """Return the ``error.code`` of a Gym control-plane error body, if it has one.""" + if not isinstance(response_body, str): + return None + try: + payload = json.loads(response_body) + except json.JSONDecodeError: + return None + error = payload.get("error") if isinstance(payload, Mapping) else None + code = error.get("code") if isinstance(error, Mapping) else None + return code if isinstance(code, str) else None + + +def _gym_failure_row_error( + row: Mapping[str, Any], result: Mapping[str, Any] +) -> Optional[Exception]: + """Map one Gym sidecar failure row onto the typed failure its ``/run`` raised. + + ``run_examples(route_failures_to_sidecar=True)`` resolves a failed ``/run`` to its + row plus ``_ng_failure_*`` fields instead of an exception that carries no row. That + is what lets an ``admission_closed`` refusal name the row it refused while the rest + of the batch keeps streaming. Every other failure maps to what + :func:`_typed_gym_failure` returns for the same status, and a failure with no HTTP + reply is transport-shaped unless the reply arrived but was not JSON. + + Only Gym's no-rollout classes are failures. A row that ran can carry a class + too, such as the ``judge_failed`` row Gym's judge failsafe scores zero and masks; + it is post-processed like any other. + + Returns None for a row that ran. + """ + # Checked before the imports: every streamed row lands here, and NeMo-Gym is an + # optional extra that rows which ran must not need. + if "_ng_failure_class" not in result: + return None + from nemo_gym._checkpoint.errors import AdmissionClosedError + from nemo_gym.rollout_collection import ( + AGENT_REQUEST_FAILED_FAILURE_CLASS, + AGENT_RUN_ERROR_FAILURE_CLASS, + ENVIRONMENT_SERVER_FAILURE_CLASS, + NG_FAILURE_CLASS_KEY, + ) + + if result[NG_FAILURE_CLASS_KEY] not in { + AGENT_REQUEST_FAILED_FAILURE_CLASS, + AGENT_RUN_ERROR_FAILURE_CLASS, + ENVIRONMENT_SERVER_FAILURE_CLASS, + }: + return None + status = result.get("_ng_failure_http_status") + body = result.get("_ng_failure_response_body") + message = result.get("_ng_failure_message") + if status == 409 and _gym_error_code(body) == AdmissionClosedError.code: + return GymAdmissionClosed( + f"NeMo-Gym refused /run while a checkpoint has admission closed: {body}", + row["_rowidx"], + ) + if isinstance(status, int): + detail = f"NeMo-Gym /run failed with HTTP {status}: {message}" + if http_status_is_infra(status): + return GymTransportError(detail) + return RolloutDataFailure(detail) + failure_type = result.get("_ng_failure_type") + detail = f"NeMo-Gym /run failed without an HTTP reply ({failure_type}): {message}" + if failure_type == "JSONDecodeError": + return RolloutDataFailure(detail) + return GymTransportError(detail) + + def get_nemo_gym_uv_cache_dir() -> str | None: """Return the uv cache directory inside a container, or None outside one. @@ -341,6 +420,10 @@ class NemoGymConfig(TypedDict): # server, switches run_rollouts to receipt mode, and assembles receipts # from the manifest control route. None/absent = legacy token-echo path. token_capture: NotRequired[Dict[str, Any] | None] + # Internal checkpoint topology. These are derived by the actor builder, + # never read from user YAML. + checkpoint_instance: GymCheckpointInstance + turn_checkpointing_enabled: bool # Gym control-plane server name (the model server hosting the ledger) and the @@ -474,6 +557,10 @@ def __init__(self, cfg: NemoGymConfig): self.head_server_config: Any = None self.node_ip: Optional[str] = None self.head_server_port: Optional[int] = None + self._checkpoint_adapter: Optional[GymCheckpointAdapter] = None + self._server_client: Any = None + self._control_headers: Dict[str, str] = {} + self._control_timeout_s = 60.0 # Installed by set_tokenizer at spinup, not passed per rollout call. Declared # here rather than in _spinup so a second spinup cannot wipe an installed # tokenizer and then report that set_tokenizer was never called. @@ -612,7 +699,8 @@ def _spinup(self) -> None: token_capture and token_capture.get("enabled") ) self._server_client = None - self._control_headers: Dict[str, str] = {} + self._checkpoint_adapter = None + self._control_headers = {} self._control_timeout_s = 60.0 if self._token_capture_enabled: assert token_capture is not None @@ -702,6 +790,125 @@ def _control_client(self): ) return self._server_client + async def initialize_checkpoint_adapter( + self, + ) -> GymCheckpointParticipantSummary: + """Discover this actor's Gym v2 checkpoint participants once.""" + self._require_spinup() + if not self.cfg["turn_checkpointing_enabled"]: + raise RuntimeError("Gym checkpoint coordination is not enabled") + if self._checkpoint_adapter is None: + token_capture = self.cfg.get("token_capture") + if token_capture is None: + raise RuntimeError( + "Gym checkpoint coordination has no token-capture credentials" + ) + self._checkpoint_adapter = GymCheckpointAdapter( + instance=self.cfg["checkpoint_instance"], + client=self._control_client(), + auth_token=token_capture["control_auth_token"], + ) + return await self._checkpoint_adapter.discover() + + def _require_checkpoint_adapter(self) -> GymCheckpointAdapter: + """Return the discovered actor-local adapter or fail before coordination.""" + if self._checkpoint_adapter is None: + raise RuntimeError("Gym checkpoint adapter has not been initialized") + return self._checkpoint_adapter + + async def checkpoint_prepare( + self, + checkpoint_id: str, + *, + deadline_ts: float, + ) -> GymCheckpointPrepareSummary: + """Park this actor's Gym deployment at one recoverable boundary.""" + result = await self._require_checkpoint_adapter().prepare( + checkpoint_id, + deadline_ts=deadline_ts, + ) + return GymCheckpointPrepareSummary( + prepared=result.prepared, + blockers={ + server: tuple(blockers) + for server, blockers in result.blockers().items() + }, + ) + + async def checkpoint_renew( + self, + checkpoint_id: str, + *, + deadline_ts: float, + ) -> None: + """Extend this actor's active Gym checkpoint lease.""" + await self._require_checkpoint_adapter().renew( + checkpoint_id, + deadline_ts=deadline_ts, + ) + + async def checkpoint_retire( + self, + checkpoint_id: str, + episodes: tuple[GymCheckpointEpisode, ...], + *, + deadline_ts: float, + ) -> None: + """Retire selected episode attempts from this Gym deployment.""" + await self._require_checkpoint_adapter().retire( + checkpoint_id, + episodes, + deadline_ts=deadline_ts, + ) + + async def checkpoint_commit( + self, + checkpoint_id: str, + checkpoint_root: str, + episodes: tuple[GymCheckpointEpisode, ...], + *, + deadline_ts: float, + ) -> GymCheckpointCommitSummary: + """Persist this actor's selected episode state under ``checkpoint_root``.""" + return await self._require_checkpoint_adapter().commit( + checkpoint_id, + checkpoint_root, + episodes, + deadline_ts=deadline_ts, + ) + + async def checkpoint_restore( + self, + checkpoint_id: str, + checkpoint_root: str, + episodes: tuple[GymCheckpointEpisode, ...], + *, + participants: tuple[GymCheckpointParticipantRecord, ...], + source_checkpoint_id: str, + deadline_ts: float, + ) -> None: + """Restore and validate this actor's state from ``checkpoint_root``.""" + await self._require_checkpoint_adapter().restore( + checkpoint_id, + checkpoint_root, + episodes, + participants=participants, + source_checkpoint_id=source_checkpoint_id, + deadline_ts=deadline_ts, + ) + + async def checkpoint_resume( + self, + checkpoint_id: str, + *, + deadline_ts: float, + ) -> None: + """Release this actor's Gym participants after commit or abort.""" + await self._require_checkpoint_adapter().resume( + checkpoint_id, + deadline_ts=deadline_ts, + ) + async def _control(self, method: str, path: str, **kwargs: Any) -> dict: headers = {**kwargs.pop("headers", {}), **self._control_headers} try: @@ -779,7 +986,7 @@ async def run_rollouts( timer_prefix: str, deduplicate_multimodal_data: bool = False, per_prompt: bool = False, - ) -> AsyncGenerator[tuple[int, dict, dict, dict | None], None]: + ) -> AsyncGenerator[tuple[int, dict, dict, dict | None] | GymAdmissionClosed, None]: """Stream postprocessed rollouts as NeMo-Gym tasks complete. A thin span-opening wrapper over :meth:`_stream_rollouts`, which holds @@ -810,7 +1017,9 @@ async def run_rollouts( ``rowidx`` echoes back the ``_rowidx`` the caller stamped on the example, which is how the caller maps a result to its slot. ``timing_metrics`` is ``None`` on every tuple but the last, which - carries the batch totals. + carries the batch totals. A row Gym refused because a checkpoint + had admission closed is yielded as a ``GymAdmissionClosed`` + naming its ``rowidx`` instead of a tuple. """ attributes = {"rl.gym.batch_size": len(nemo_gym_examples)} # Two branches rather than a group variable, so the drift test can read @@ -845,7 +1054,7 @@ async def _stream_rollouts( nemo_gym_examples: list[dict], timer_prefix: str, deduplicate_multimodal_data: bool = False, - ) -> AsyncGenerator[tuple[int, dict, dict, dict | None], None]: + ) -> AsyncGenerator[tuple[int, dict, dict, dict | None] | GymAdmissionClosed, None]: """Body of :meth:`run_rollouts`; see there for the tracing wrapper.""" self._require_spinup() if not nemo_gym_examples: @@ -867,8 +1076,12 @@ async def _stream_rollouts( timer = Timer() timer.start("_run_rollouts_total") + # Sidecar mode resolves a failed /run to its row rather than an exception + # that names no row; see _gym_failure_row_error. nemo_gym_result_iterator = self.rch.run_examples( - examples=nemo_gym_examples, head_server_config=self.head_server_config + examples=nemo_gym_examples, + head_server_config=self.head_server_config, + route_failures_to_sidecar=True, ) # Gym resolves task_source to agent_ref synchronously in run_examples(). # Build the counter afterward so completion rows use the resolved identity. @@ -894,6 +1107,24 @@ async def _stream_rollouts( # the whole point. The status and message are already in `detail`. raise typed from None raise + failure = _gym_failure_row_error(nemo_gym_row, nemo_gym_result) + # A refusal is expected while a checkpoint is open; only a real + # failure is worth its body on stderr. + if failure is not None and not isinstance(failure, GymAdmissionClosed): + print( + "EXCEPTION RESULT", + nemo_gym_result.get("_ng_failure_response_body"), + file=sys.stderr, + ) + raise failure + + if isinstance(failure, GymAdmissionClosed): + # Yielded, not raised: rows Gym admitted before the checkpoint + # closed admission keep streaming, and the caller unwinds this one + # row before the checkpoint drains. + num_results += 1 + yield failure + continue with timer.time(label=f"{timer_prefix}/postprocess_results"): if self._token_capture_enabled: @@ -1548,6 +1779,7 @@ def build_nemo_gym_config( model_name: str, enable_router_replay: bool, use_fastokens: bool, + turn_checkpointing_enabled: bool, token_capture: Optional[dict[str, Any]] = None, ) -> NemoGymConfig: """Build the ``NemoGymConfig`` for a single, unsharded NeMo-Gym actor. @@ -1566,6 +1798,8 @@ def build_nemo_gym_config( routed-experts carry dtype ("int8"/"int16"/"int32") for the model. use_fastokens: Forwarded from ``policy.tokenizer.use_fastokens`` so the actor patches its tokenizer the same way the driver does. + turn_checkpointing_enabled: Install Gym v2 checkpoint participants and + validate their discovery during actor startup. Returns: A ``NemoGymConfig`` with NeMo-RL fields at the top level and the @@ -1592,6 +1826,11 @@ def build_nemo_gym_config( model_name=model_name, enable_router_replay=enable_router_replay, use_fastokens=use_fastokens, + checkpoint_instance=GymCheckpointInstance( + shard_name=DEFAULT_SHARD_NAME, + replica_index=0, + ), + turn_checkpointing_enabled=turn_checkpointing_enabled, token_capture=token_capture, ) @@ -1603,6 +1842,8 @@ def _build_gym_actor_config( model_name: str, enable_router_replay: bool, use_fastokens: bool, + checkpoint_instance: GymCheckpointInstance, + turn_checkpointing_enabled: bool, token_capture: Optional[dict[str, Any]] = None, ) -> NemoGymConfig: """Turn one already-resolved Gym config mapping into a ``NemoGymConfig``. @@ -1652,6 +1893,43 @@ def _build_gym_actor_config( else "int16" ) + instance_token_capture = None + if token_capture is not None: + instance_token_capture = dict(token_capture) + if instance_token_capture.get("enabled"): + capture_root = instance_token_capture.get("capture_dir") + if not isinstance(capture_root, str) or not capture_root: + raise ValueError( + "enabled token capture requires a non-empty capture_dir before " + "building NeMo-Gym actors" + ) + instance_token_capture["capture_dir"] = str( + checkpoint_instance.live_capture_dir(capture_root) + ) + + if turn_checkpointing_enabled: + if instance_token_capture is None or not instance_token_capture.get("enabled"): + raise ValueError( + "Gym checkpoint coordination requires token capture to be enabled" + ) + control_auth_token = instance_token_capture.get("control_auth_token") + if not isinstance(control_auth_token, str) or not control_auth_token: + raise ValueError( + "Gym checkpoint coordination requires a token-capture control token" + ) + nemo_gym_dict["checkpoint"] = { + "enabled": True, + "control_auth_token": control_auth_token, + } + policy_overrides = ( + nemo_gym_dict.setdefault("policy_model", {}) + .setdefault("responses_api_models", {}) + .setdefault("vllm_model", {}) + ) + policy_overrides["checkpoint_policy"] = True + if instance_token_capture.get("generation_prefix_cuts_enabled", False): + policy_overrides["checkpoint_generation_cuts"] = True + return NemoGymConfig( model_name=model_name, base_urls=base_urls, @@ -1662,7 +1940,9 @@ def _build_gym_actor_config( routed_experts_dtype=routed_experts_dtype, use_fastokens=use_fastokens, initial_global_config_dict=nemo_gym_dict, - token_capture=token_capture, + token_capture=instance_token_capture, + checkpoint_instance=checkpoint_instance, + turn_checkpointing_enabled=turn_checkpointing_enabled, **port_range, **multimodal_flags, ) @@ -1727,6 +2007,15 @@ def is_sharded(self) -> bool: def all_handles(self) -> List[Any]: return [handle for replicas in self.handles.values() for handle in replicas] + @property + def checkpoint_handles(self) -> Dict[str, Any]: + """Map every stable checkpoint instance ID to its actor handle.""" + return { + GymCheckpointInstance(shard_name, replica).instance_id: handle + for shard_name, replicas in self.handles.items() + for replica, handle in enumerate(replicas) + } + @property def hosted_routes(self) -> frozenset[str]: """Agent and task-source entry names this set can route to.""" @@ -1786,6 +2075,24 @@ def instance_label(self, handle: Any) -> str: return shard_name if len(replicas) == 1 else f"{shard_name}/{index}" raise ShardSetupError("Handle does not belong to this NeMo-Gym shard set") + def checkpoint_instance_for_handle(self, handle: Any) -> GymCheckpointInstance: + """Return the durable checkpoint identity of an actor handle.""" + for shard_name, replicas in self.handles.items(): + for replica, candidate in enumerate(replicas): + if candidate is handle: + return GymCheckpointInstance(shard_name, replica) + raise ShardSetupError("Handle does not belong to this NeMo-Gym shard set") + + def handle_for_checkpoint_instance(self, instance_id: str) -> Any: + """Resolve a persisted checkpoint instance ID back to its actor handle.""" + try: + return self.checkpoint_handles[instance_id] + except KeyError: + raise ShardSetupError( + f"No NeMo-Gym actor has checkpoint instance {instance_id!r}; " + f"available={sorted(self.checkpoint_handles)!r}" + ) from None + def sole_handle(self) -> Any: """The only actor, for callers that predate routing. @@ -1863,6 +2170,7 @@ def build_nemo_gym_actors( tokenizer: PreTrainedTokenizerBase, enable_router_replay: bool, use_fastokens: bool, + turn_checkpointing_enabled: bool, token_capture: Optional[dict[str, Any]] = None, pg_ready_timeout: float = DEFAULT_SHARD_PG_READY_TIMEOUT_SECONDS, spinup_timeout: float = DEFAULT_SHARD_SPINUP_TIMEOUT_SECONDS, @@ -1878,6 +2186,8 @@ def build_nemo_gym_actors( Args: tokenizer: Installed on every actor once it is up, rather than passed per rollout call. See ``NemoGym.set_tokenizer`` for why. + turn_checkpointing_enabled: Install and discover one Gym v2 checkpoint + coordinator per actor without triggering checkpoint saves. Returns: A :class:`NemoGymShardSet` whose actors are all running and validated. @@ -1898,6 +2208,7 @@ def build_nemo_gym_actors( tokenizer=tokenizer, enable_router_replay=enable_router_replay, use_fastokens=use_fastokens, + turn_checkpointing_enabled=turn_checkpointing_enabled, token_capture=token_capture, ) @@ -1909,6 +2220,7 @@ def build_nemo_gym_actors( tokenizer=tokenizer, enable_router_replay=enable_router_replay, use_fastokens=use_fastokens, + turn_checkpointing_enabled=turn_checkpointing_enabled, token_capture=token_capture, pg_ready_timeout=pg_ready_timeout, spinup_timeout=spinup_timeout, @@ -1923,12 +2235,13 @@ def _build_single_gym_actor( tokenizer: PreTrainedTokenizerBase, enable_router_replay: bool, use_fastokens: bool, + turn_checkpointing_enabled: bool, token_capture: Optional[dict[str, Any]], ) -> NemoGymShardSet: - """The pre-sharding path: one actor, no placement group, no discovery. + """The pre-sharding path: one actor and no placement group. - Discovery is skipped rather than merely unused. Its checks compare entry - names *between* shards, so with one shard there is nothing they could find. + Cross-shard route discovery is skipped because there is only one shard. + Gym checkpoint-participant discovery still runs when turn recovery is enabled. """ actor_config = _build_gym_actor_config( nemo_gym_dict, @@ -1936,6 +2249,11 @@ def _build_single_gym_actor( model_name=model_name, enable_router_replay=enable_router_replay, use_fastokens=use_fastokens, + checkpoint_instance=GymCheckpointInstance( + shard_name=DEFAULT_SHARD_NAME, + replica_index=0, + ), + turn_checkpointing_enabled=turn_checkpointing_enabled, token_capture=token_capture, ) @@ -1953,6 +2271,8 @@ def _build_single_gym_actor( try: ray.get(actor._spinup.remote()) ray.get(actor.set_tokenizer.remote(tokenizer)) + if turn_checkpointing_enabled: + ray.get(actor.initialize_checkpoint_adapter.remote()) except BaseException: shard_set.shutdown( timeout=NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S, @@ -1971,6 +2291,7 @@ def _build_sharded_gym_actors( tokenizer: PreTrainedTokenizerBase, enable_router_replay: bool, use_fastokens: bool, + turn_checkpointing_enabled: bool, token_capture: Optional[dict[str, Any]], pg_ready_timeout: float, spinup_timeout: float, @@ -2054,12 +2375,22 @@ def _build_sharded_gym_actors( model_name=model_name, enable_router_replay=enable_router_replay, use_fastokens=use_fastokens, + checkpoint_instance=GymCheckpointInstance( + shard_name=shard.name, + replica_index=replica, + ), + turn_checkpointing_enabled=turn_checkpointing_enabled, token_capture=token_capture, ) ) shard_set.handles.setdefault(shard.name, []).append(actor) _spinup_shards_concurrently(shard_set, spinup_timeout, tokenizer=tokenizer) + if turn_checkpointing_enabled: + _initialize_checkpoint_adapters_concurrently( + shard_set, + timeout=spinup_timeout, + ) shard_set.route_to_shard = _discover_route_shard_map(shard_set, plan) except BaseException: # A ray.get timeout does not cancel the actor-side work, so a @@ -2146,6 +2477,28 @@ def _spinup_shards_concurrently( ) from error +def _initialize_checkpoint_adapters_concurrently( + shard_set: NemoGymShardSet, + *, + timeout: float, +) -> None: + """Discover each actor's independent Gym checkpoint deployment.""" + deadline = monotonic() + timeout + pending = [ + (shard_name, replica, handle.initialize_checkpoint_adapter.remote()) + for shard_name, replicas in shard_set.handles.items() + for replica, handle in enumerate(replicas) + ] + for shard_name, replica, reference in pending: + try: + ray.get(reference, timeout=max(0.0, deadline - monotonic())) + except BaseException as error: + raise ShardSetupError( + f"NeMo-Gym shard '{shard_name}' (replica {replica}) could not " + f"discover its Gym v2 checkpoint participants: {error}" + ) from error + + def _discover_route_shard_map( shard_set: NemoGymShardSet, plan: ShardPlan ) -> Dict[str, str]: @@ -2169,6 +2522,7 @@ def spinup_nemo_gym_actor( tokenizer: PreTrainedTokenizerBase, enable_router_replay: bool, use_fastokens: bool, + turn_checkpointing_enabled: bool, token_capture: Optional[dict[str, Any]] = None, ) -> Any: """Spin up a single NeMo-Gym actor against the given generation server URLs. @@ -2185,6 +2539,8 @@ def spinup_nemo_gym_actor( token_capture: Dumped ``TokenCaptureConfig`` when ledger-authoritative token capture is enabled, else ``None``. Forwarded to ``build_nemo_gym_config``. + turn_checkpointing_enabled: Install and discover Gym v2 checkpoint + participants. This helper still returns one unsharded actor. Returns: The spun-up ``NemoGym`` Ray actor handle (``_spinup`` already awaited). @@ -2209,6 +2565,7 @@ def spinup_nemo_gym_actor( tokenizer=tokenizer, enable_router_replay=enable_router_replay, use_fastokens=use_fastokens, + turn_checkpointing_enabled=turn_checkpointing_enabled, token_capture=token_capture, ).sole_handle() diff --git a/nemo_rl/experience/failures.py b/nemo_rl/experience/failures.py index 60a400196df..58a1e70a2d9 100644 --- a/nemo_rl/experience/failures.py +++ b/nemo_rl/experience/failures.py @@ -109,6 +109,26 @@ class GymTransportError(RolloutInfraFailure): """NeMo-Gym failed at the transport layer rather than returning a rollout.""" +class GymAdmissionClosed(RolloutInfraFailure): + """NeMo-Gym refused a new ``/run`` because a checkpoint has admission closed. + + Gym refuses before it records anything for the episode, so the attempt is + unwound and re-dispatched once dispatch admission reopens. The NemoGym actor + yields it per row instead of raising, so rows Gym did admit keep streaming. + It is INFRA only as a fallback: a caller without a dispatch recorder that + raises it retries the row instead of charging the prompt a data failure. + """ + + def __init__(self, detail: str, rowidx: int) -> None: + # Both values ride in ``args`` so the exception pickles across Ray. + super().__init__(detail, rowidx) + self.detail = detail + self.rowidx = rowidx + + def __str__(self) -> str: + return self.detail + + class RolloutDataFailure(RolloutFailure): """This prompt cannot be rolled out, and another shard would fail identically. diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 489bacf8a12..e266703127b 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -19,11 +19,12 @@ import enum import json import math +import time import uuid -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping -from contextlib import asynccontextmanager +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, asynccontextmanager, nullcontext from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any, Optional, Protocol import ray.exceptions import torch @@ -42,6 +43,7 @@ from nemo_rl.data.multimodal_utils import VLLM_CONTENT_KEY, VLLM_PROMPT_KEYS from nemo_rl.data_plane.schema import MASK_SAMPLE from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.environments.gym_checkpoint_adapter import GymCheckpointEpisode from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.environments.nemo_gym import ( as_nemo_gym_shard_set, @@ -50,6 +52,7 @@ from nemo_rl.experience.failures import ( FailureClass, GenerationUnavailable, + GymAdmissionClosed, GymTransportError, RolloutDataFailure, RolloutFailure, @@ -117,6 +120,113 @@ def _contains_post_write_enrichment_error(error: BaseException) -> bool: return False +class RolloutDispatchAdmissionGate: + """Pause new Gym ``/run`` submissions without blocking live completions. + + An admission is held only while the Ray streaming RPC is submitted. The + returned stream is consumed after the admission is released, so closing + this gate does not wait for generation, completion callbacks, finalization, + or TQ writes. + """ + + def __init__(self) -> None: + self._condition = asyncio.Condition() + self._closed = False + self._active_submissions = 0 + + @asynccontextmanager + async def admission(self) -> AsyncIterator[None]: + """Admit one Gym RPC submission after any checkpoint pause exits.""" + async with self._condition: + await self._condition.wait_for(lambda: not self._closed) + self._active_submissions += 1 + try: + yield + finally: + async with self._condition: + self._active_submissions -= 1 + if self._active_submissions == 0: + self._condition.notify_all() + + @asynccontextmanager + async def closed(self) -> AsyncIterator[None]: + """Block new submissions after every admitted submission has crossed.""" + async with self._condition: + await self._condition.wait_for(lambda: not self._closed) + self._closed = True + try: + await self._condition.wait_for(lambda: self._active_submissions == 0) + except BaseException: + self._closed = False + self._condition.notify_all() + raise + try: + yield + finally: + async with self._condition: + self._closed = False + self._condition.notify_all() + + +class RolloutDispatchRecorder(Protocol): + """Recovery-ledger hooks around one physical Gym ``/run`` submission. + + A ledger attempt is ``DISPATCHED`` only while a ``/run`` is actually on its + way to Gym. Marking it before the dispatch gate opened left an attempt that + a checkpoint counted as Gym-owned, yet that could not reach Gym until the + checkpoint released the gate, so the checkpoint drain could only time out. + """ + + def submitting( + self, generation_indices: Sequence[int] + ) -> AbstractAsyncContextManager[Callable[[], None]]: + """Hold the ledger for one submission; call the yielded hook right after it. + + Entered inside the dispatch admission, so a checkpoint cannot close the + gate between the submission and the mark. The yielded hook is + synchronous so nothing can interleave once the RPC is in flight. + """ + ... + + async def refused(self, generation_index: int) -> None: + """Unwind one attempt Gym refused because a checkpoint closed admission.""" + ... + + async def dropped( + self, generation_indices: Sequence[int], *, replace: bool + ) -> DroppedDispatch: + """Abandon attempts a broken stream left running in Gym. + + With ``replace``, each abandoned row gets a fresh attempt for its re-send. + """ + ... + + def retired(self, episodes: Sequence[GymCheckpointEpisode]) -> None: + """Forget dropped episodes once Gym has retired them.""" + ... + + +@dataclass(frozen=True) +class DroppedDispatch: + """Rows a broken Gym stream never returned, as the recovery ledger left them.""" + + # Gym-owned episodes to retire before anything replaces them. + episodes: tuple[GymCheckpointEpisode, ...] + # Replacement Gate rollout IDs by row index; empty unless replaced. + rollout_ids: Mapping[int, str] + + +# Bound on retiring the rows a broken stream left running. A failed retire is +# not fatal: the episodes stay queued and the next checkpoint retires them. +_DROPPED_EPISODE_RETIRE_TIMEOUT_S = 30.0 + + +# First pause before re-sending rows Gym refused at checkpoint admission. Doubles +# per consecutive refusal; the prompt group's rollout deadline bounds the total. +_PARKED_REDISPATCH_BACKOFF_S = 0.5 +_MAX_PARKED_REDISPATCH_BACKOFF_S = 5.0 + + 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") @@ -242,6 +352,9 @@ class RolloutStats: # redoing the whole thing, so they never reached the counters above and gym could # retry rows all run with redispatch_total sitting flat. gym_row_redispatches: int = 0 + # Rows Gym refused because a checkpoint had admission closed. Re-sent without + # spending any retry budget, so they are counted apart from every family above. + gym_admission_closed_rows: int = 0 def record_redispatch(self, reason: str) -> None: self.redispatches_by_reason[reason] = ( @@ -269,6 +382,9 @@ def record_infra_drop(self, reason: str, consecutive: int) -> None: def record_gym_row_redispatch(self, rows: int = 1) -> None: self.gym_row_redispatches += rows + def record_gym_admission_closed(self, rows: int = 1) -> None: + self.gym_admission_closed_rows += rows + def as_metrics(self) -> dict[str, float]: """Flatten into a metric dict for the SingleController logger.""" # Every family gets an aggregate, not just per-exception series: alerting on @@ -286,6 +402,9 @@ def as_metrics(self) -> dict[str, float]: sum(self.data_failures_by_reason.values()) ), "rollout/gym_row_redispatch_total": float(self.gym_row_redispatches), + "rollout/gym_checkpoint_parked_rows_total": float( + self.gym_admission_closed_rows + ), "rollout/infra_drops_total": float( sum(self.infra_drops_by_reason.values()) ), @@ -390,11 +509,18 @@ async def _gather_cancelling_siblings(coros: list[Any]) -> list[Any]: class RequestDeadlineRegistry: - """Live request deadlines, pausable while a colocated engine has switched to training.""" + """Live request deadlines shared by independently pausable subsystems.""" + + _DEFAULT_HOLDER = "generation_stand_down" def __init__(self) -> None: self._live: set["_Deadline"] = set() - self.suspended = False + self._suspension_holders: set[str] = set() + + @property + def suspended(self) -> bool: + """Whether at least one subsystem still owns the deadline pause.""" + return bool(self._suspension_holders) def add(self, deadline: "_Deadline") -> None: self._live.add(deadline) @@ -404,17 +530,24 @@ def add(self, deadline: "_Deadline") -> None: def discard(self, deadline: "_Deadline") -> None: self._live.discard(deadline) - def suspend(self) -> None: - if self.suspended: + def suspend(self, holder: str = _DEFAULT_HOLDER) -> None: + """Pause all deadline clocks for one idempotent owner.""" + if holder in self._suspension_holders: + return + already_suspended = self.suspended + self._suspension_holders.add(holder) + if already_suspended: return - self.suspended = True for deadline in self._live: deadline.suspend() - def resume(self) -> None: - if not self.suspended: + def resume(self, holder: str = _DEFAULT_HOLDER) -> None: + """Release one owner and resume clocks after the final owner leaves.""" + if holder not in self._suspension_holders: + return + self._suspension_holders.remove(holder) + if self.suspended: return - self.suspended = False for deadline in self._live: deadline.resume() @@ -522,6 +655,7 @@ async def run_rollout( generation_indices: Optional[list[int]] = None, on_completion: Optional[RolloutCompletionCallback] = None, recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, + gym_instance_id: Optional[str] = None, ) -> PromptGroupRecord: """Run num_generations_per_prompt rollouts for one prompt. @@ -544,6 +678,9 @@ async def run_rollout( assert recovery_granularity is RecoveryGranularity.SIBLING, ( "recovery granularity is only supported on the NeMo-Gym path" ) + assert gym_instance_id is None, ( + "Gym checkpoint ownership is only supported on the NeMo-Gym path" + ) timer = Timer() timer_prefix = "timing/rollout" timer.start(f"{timer_prefix}/total") @@ -968,9 +1105,18 @@ def __init__( ).max_gym_row_attempts self._stats = stats self._effort_config = effort_config + self._dispatch_admission_gate: RolloutDispatchAdmissionGate | None = None self._validate_init_params() + def set_rollout_dispatch_admission_gate( + self, gate: RolloutDispatchAdmissionGate + ) -> None: + """Bind the controller-owned gate used by Gym checkpoint coordination.""" + if self._dispatch_admission_gate is not None: + raise RuntimeError("Gym rollout dispatch admission gate is already bound") + self._dispatch_admission_gate = gate + async def run_rollout( self, input_sample: DatumSpec, @@ -979,6 +1125,8 @@ async def run_rollout( generation_indices: Optional[list[int]] = None, on_completion: Optional[RolloutCompletionCallback] = None, recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, + gym_instance_id: Optional[str] = None, + dispatch_recorder: Optional[RolloutDispatchRecorder] = None, ) -> PromptGroupRecord: """Run num_generations_per_prompt rollouts for one prompt. @@ -988,6 +1136,8 @@ async def run_rollout( per generation, riding each row's run body as the opaque ``_ng_rollout_id`` key (agents stamp /ng-rollout/ from it; zero agent changes). + dispatch_recorder: Recovery-ledger hooks that mark each submission + dispatched and unwind rows Gym refuses at checkpoint admission. Returns: PromptGroupRecord with num_generations_per_prompt completions. @@ -1007,6 +1157,8 @@ async def run_rollout( timer_prefix, on_completion=on_completion, recovery_granularity=recovery_granularity, + gym_instance_id=gym_instance_id, + dispatch_recorder=dispatch_recorder, ) # Token-capture receipt rows carry empty message logs by design — the # canonical row (and any media it needs) is rebuilt by the finalizer @@ -1138,6 +1290,8 @@ async def _stream_rows( total_rows: int, timer_prefix: str, on_completion: Optional[RolloutCompletionCallback] = None, + dispatch_recorder: Optional[RolloutDispatchRecorder] = None, + admission_closed_rows: Optional[set[int]] = None, ) -> Optional[dict[str, Any]]: """Dispatch ``pending`` rows and fill their slots in ``results`` as they land. @@ -1149,10 +1303,18 @@ async def _stream_rows( 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. + dispatch_recorder: Marks the submission dispatched in the recovery + ledger and unwinds rows Gym refuses at checkpoint admission. + admission_closed_rows: Collects the indices of rows Gym refused because a + checkpoint had admission closed; required with ``dispatch_recorder``. Returns: The environment's timing metrics, or None if the stream ended without them. """ + if dispatch_recorder is not None and admission_closed_rows is None: + raise ValueError( + "a dispatch recorder requires a admission_closed_rows collector" + ) dispatched = {row["_rowidx"] for row in pending} inputs_by_rowidx = {row["_rowidx"]: row for row in pending} received: set[int] = set() @@ -1160,13 +1322,53 @@ async def _stream_rows( # Read here, not in the actor: the scope is a ContextVar in this # process, and NemoGym is a separate Ray actor. - async for result_ref in dispatch_with_trace_context( - nemo_gym_env.run_rollouts.options(num_returns="streaming"), - pending, - timer_prefix, - per_prompt=in_per_prompt_scope(), - ): - rowidx, resolved_agent_ref, result, timing_metrics = await result_ref + remote_method = nemo_gym_env.run_rollouts.options(num_returns="streaming") + gate = self._dispatch_admission_gate + # ``dispatch_with_trace_context`` calls ``.remote`` synchronously. + # Release admission immediately after Ray accepts the streaming RPC; + # consuming the returned stream must remain live during Gym prepare. + async with gate.admission() if gate is not None else nullcontext(): + if dispatch_recorder is None: + result_refs = dispatch_with_trace_context( + remote_method, + pending, + timer_prefix, + per_prompt=in_per_prompt_scope(), + ) + else: + # The ledger turns DISPATCHED only once Ray has accepted the + # RPC, and still inside the admission, so a checkpoint never + # counts an attempt that is waiting at the closed gate. + async with dispatch_recorder.submitting( + sorted(dispatched) + ) as mark_dispatched: + result_refs = dispatch_with_trace_context( + remote_method, + pending, + timer_prefix, + per_prompt=in_per_prompt_scope(), + ) + mark_dispatched() + + async for result_ref in result_refs: + item = await result_ref + if isinstance(item, GymAdmissionClosed): + rowidx = item.rowidx + if rowidx not in dispatched or rowidx in received: + raise ValueError( + f"NeMo-Gym refused row index {rowidx}, which is not an " + f"outstanding row of this attempt ({sorted(dispatched)})" + ) + if dispatch_recorder is None: + raise item + received.add(rowidx) + # Unwound before the next row is read, so a checkpoint draining + # right now stops waiting on an episode Gym never admitted. + await dispatch_recorder.refused(rowidx) + assert admission_closed_rows is not None + admission_closed_rows.add(rowidx) + continue + rowidx, resolved_agent_ref, result, timing_metrics = item # Validated against the original group, not the pending subset: on a # re-dispatch the row keeps its original index so results stay ordered. if not isinstance(rowidx, int) or not 0 <= rowidx < total_rows: @@ -1206,6 +1408,67 @@ async def _stream_rows( return env_timing_metrics + async def _drop_unreturned_rows( + self, + nemo_gym_env: Any, + pending: list[dict], + results: list[Optional[dict]], + admission_closed_rows: set[int], + *, + dispatch_recorder: Optional[RolloutDispatchRecorder], + replace: bool, + ) -> None: + """Retire the rows a stream ended without, before anything replaces them. + + Gym runs every row of a ``/run`` as its own task, and an ended stream + cancels none of them. A re-send under the same identity would run two + copies of one episode, and a prompt retry would leave the old ones for + the next checkpoint to export. With ``replace``, the rows are re-stamped + with fresh rollout IDs for their re-send. + """ + if dispatch_recorder is None: + return + unreturned = [ + row + for row in pending + if results[row["_rowidx"]] is None + and row["_rowidx"] not in admission_closed_rows + ] + if not unreturned: + return + dropped = await dispatch_recorder.dropped( + [row["_rowidx"] for row in unreturned], replace=replace + ) + for row in unreturned: + rollout_id = dropped.rollout_ids.get(row["_rowidx"]) + if rollout_id is not None: + row["_ng_rollout_id"] = rollout_id + if not dropped.episodes: + return + gate = self._dispatch_admission_gate + try: + # Admission keeps any checkpoint out, so Gym is idle and accepts the + # retire under a checkpoint ID of its own. + async with gate.admission() if gate is not None else nullcontext(): + await asyncio.wait_for( + nemo_gym_env.checkpoint_retire.remote( + f"rollout-drop-{uuid.uuid4().hex}", + dropped.episodes, + deadline_ts=time.time() + _DROPPED_EPISODE_RETIRE_TIMEOUT_S, + ), + timeout=_DROPPED_EPISODE_RETIRE_TIMEOUT_S, + ) + dispatch_recorder.retired(dropped.episodes) + except Exception as error: + # The ledger still queues them, so the next checkpoint retires them + # before it commits. + print( + f"NeMo-Gym: could not retire {len(dropped.episodes)} episode(s) " + f"the rollout stream dropped; the next checkpoint retires them: " + f"{error!r}", + flush=True, + ) + async def _run_rollouts( self, inputs: list[dict], @@ -1214,6 +1477,8 @@ async def _run_rollouts( *, on_completion: Optional[RolloutCompletionCallback] = None, recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, + gym_instance_id: Optional[str] = None, + dispatch_recorder: Optional[RolloutDispatchRecorder] = None, ) -> tuple[list[Completion], LLMMessageLogType, dict[str, Any]]: """Dispatch rows to NeMo-Gym; return completions, prompt, and metrics. @@ -1226,7 +1491,11 @@ async def _run_rollouts( # These rows are all one prompt's generations. # They share one Gym route and must stay on one instance. shard_set = as_nemo_gym_shard_set(self._task_to_env["nemo_gym"]) - nemo_gym_env = shard_set.pick_handle(get_nemo_gym_route_name(inputs[0])) + nemo_gym_env = ( + shard_set.pick_handle(get_nemo_gym_route_name(inputs[0])) + if gym_instance_id is None + else shard_set.handle_for_checkpoint_instance(gym_instance_id) + ) instance_label = shard_set.instance_label(nemo_gym_env) instance_timer_prefix = f"{timer_prefix}/shard/{instance_label}" total_rows = self._num_generations_per_prompt @@ -1275,11 +1544,15 @@ async def _run_rollouts( "NeMo-Gym prompt group", registry=self._deadline_registry, ): - for attempt in range(1, max_row_attempts + 1): + # Rows Gym refused at checkpoint admission never ran, so re-sending + # them spends no row attempt; admission_closed_streak only paces the re-sends. + attempt = 1 + admission_closed_streak = 0 + while attempt <= max_row_attempts: pending = [row for row in inputs if results[row["_rowidx"]] is None] if not pending: break - if attempt > 1: + if attempt > 1 and admission_closed_streak == 0: print( f"NeMo-Gym: re-dispatching {len(pending)}/{total_rows} " f"row(s) (attempt {attempt}/{max_row_attempts})", @@ -1290,6 +1563,8 @@ async def _run_rollouts( # gym could retry rows all run with every counter flat. if self._stats is not None: self._stats.record_gym_row_redispatch(len(pending)) + admission_closed_rows: set[int] = set() + stream_failed = False try: timing_metrics = await self._stream_rows( nemo_gym_env, @@ -1299,22 +1574,67 @@ async def _run_rollouts( total_rows, instance_timer_prefix, on_completion=on_completion, + dispatch_recorder=dispatch_recorder, + admission_closed_rows=admission_closed_rows, ) except Exception as error: last_error = error # Only transport-shaped failures are worth another dispatch; a # prompt NeMo-Gym cannot serve fails the same way every time. - if ( - classify_rollout_failure(error) is not FailureClass.INFRA - or attempt == max_row_attempts - ): + resend = ( + classify_rollout_failure(error) is FailureClass.INFRA + and attempt < max_row_attempts + ) + await self._drop_unreturned_rows( + nemo_gym_env, + pending, + results, + admission_closed_rows, + dispatch_recorder=dispatch_recorder, + replace=resend, + ) + if not resend: error.add_note( f"NeMo-Gym instance '{instance_label}' failed during rollout collection" ) raise + stream_failed = True else: if timing_metrics is not None: env_timing_metrics = timing_metrics + await self._drop_unreturned_rows( + nemo_gym_env, + pending, + results, + admission_closed_rows, + dispatch_recorder=dispatch_recorder, + replace=attempt < max_row_attempts, + ) + if admission_closed_rows: + if self._stats is not None: + self._stats.record_gym_admission_closed( + len(admission_closed_rows) + ) + if not stream_failed: + admission_closed_streak += 1 + print( + f"NeMo-Gym: {len(admission_closed_rows)} row(s) refused at " + "checkpoint admission; re-sending once dispatch " + "admission reopens", + flush=True, + ) + # The re-send also waits on the dispatch gate; the pause + # only stops a hot loop if Gym stays closed after it opens. + await asyncio.sleep( + min( + _PARKED_REDISPATCH_BACKOFF_S + * 2 ** (admission_closed_streak - 1), + _MAX_PARKED_REDISPATCH_BACKOFF_S, + ) + ) + continue + admission_closed_streak = 0 + attempt += 1 missing = [index for index in expected_indices if results[index] is None] if missing: @@ -1386,6 +1706,15 @@ async def _run_rollouts( return completions, prompt_message_log, rollout_metrics + def select_checkpoint_instance(self, input_sample: DatumSpec) -> str: + """Choose and return the durable Gym actor identity for one prompt group.""" + extra_env_info = input_sample.get("extra_env_info") + if not isinstance(extra_env_info, Mapping): + raise ValueError("NeMo-Gym input must contain extra_env_info") + shard_set = as_nemo_gym_shard_set(self._task_to_env["nemo_gym"]) + handle = shard_set.pick_handle(get_nemo_gym_route_name(extra_env_info)) + return shard_set.checkpoint_instance_for_handle(handle).instance_id + def _results_to_completions( self, results: list[dict] ) -> tuple[list[Completion], dict[str, int]]: @@ -1576,6 +1905,90 @@ def _compute_rollout_metrics( return rollout_metrics +@dataclass(frozen=True) +class _LedgerDispatchRecorder: + """Bind one prompt group's ledger attempts to its physical Gym submissions.""" + + manager: RolloutManager + group_id: str + gym_instance_id: Optional[str] + + @asynccontextmanager + async def submitting( + self, generation_indices: Sequence[int] + ) -> AsyncIterator[Callable[[], None]]: + ledger = self.manager._recovery_ledger + async with self.manager._recovery_mutation() as cut: + record = ledger.get_group(self.group_id) + # Rows a broken stream dropped were abandoned and re-minted, so a + # re-send, like a first send, only moves reserved attempts. + reserved = [ + index + for index in generation_indices + if record.siblings[index].current_attempt.status + is RolloutAttemptStatus.RESERVED + ] + + def mark_dispatched() -> None: + if reserved: + ledger.mark_group_dispatched( + cut, + self.group_id, + generation_indices=reserved, + gym_instance_id=self.gym_instance_id, + ) + + yield mark_dispatched + + async def refused(self, generation_index: int) -> None: + async with self.manager._recovery_mutation() as cut: + self.manager._recovery_ledger.release_refused_dispatch( + cut, + self.group_id, + generation_index=generation_index, + ) + + async def dropped( + self, generation_indices: Sequence[int], *, replace: bool + ) -> DroppedDispatch: + ledger = self.manager._recovery_ledger + indices = list(generation_indices) + async with self.manager._recovery_mutation() as cut: + if self.group_id not in ledger: + # Discarding the group already queued its Gym episodes. + return DroppedDispatch(episodes=(), rollout_ids={}) + episodes = ledger.abandon_dropped_dispatch( + cut, + self.group_id, + generation_indices=indices, + replace=replace, + ) + record = ledger.get_group(self.group_id) + return DroppedDispatch( + episodes=tuple( + GymCheckpointEpisode(rollout_id, attempt) + for rollout_id, attempt in episodes + ), + rollout_ids=( + {index: record.gate_rollout_id(index) for index in indices} + if replace + else {} + ), + ) + + def retired(self, episodes: Sequence[GymCheckpointEpisode]) -> None: + # Every attempt of one dispatch shares the recorder's Gym owner. + if self.gym_instance_id is None: + return + self.manager._recovery_ledger.mark_gym_retired( + { + self.gym_instance_id: [ + (episode.rollout_id, episode.attempt) for episode in episodes + ] + } + ) + + class RolloutManager: """Routes to AsyncRolloutImpl (native async) or AsyncNemoGymRolloutImpl (NeMo-Gym), and pushes results to a TQReplayBuffer.""" @@ -1671,13 +2084,19 @@ def stats(self) -> RolloutStats: """Counters describing retry/skip activity so far.""" return self._stats - def suspend_request_deadlines(self) -> None: - """Pause live request-deadline clocks while a colocated engine is in training mode.""" - self._request_deadlines.suspend() + def suspend_request_deadlines( + self, + holder: str = RequestDeadlineRegistry._DEFAULT_HOLDER, + ) -> None: + """Pause live request-deadline clocks for one idempotent owner.""" + self._request_deadlines.suspend(holder) - def resume_request_deadlines(self) -> None: - """Resume live request-deadline clocks when a colocated engine exits training mode.""" - self._request_deadlines.resume() + def resume_request_deadlines( + self, + holder: str = RequestDeadlineRegistry._DEFAULT_HOLDER, + ) -> None: + """Resume live clocks after one owner releases its pause.""" + self._request_deadlines.resume(holder) @property def recovery_ledger(self) -> RolloutRecoveryLedger: @@ -1710,6 +2129,16 @@ def set_data_plane_checkpoint_barrier( ) self._data_plane_checkpoint_barrier = barrier + def set_rollout_dispatch_admission_gate( + self, gate: RolloutDispatchAdmissionGate + ) -> None: + """Bind the narrow Gym RPC gate used while a checkpoint is prepared.""" + if not isinstance(self._impl, AsyncNemoGymRolloutImpl): + raise RuntimeError( + "rollout dispatch admission is only supported on the NeMo-Gym path" + ) + self._impl.set_rollout_dispatch_admission_gate(gate) + @asynccontextmanager async def _recovery_mutation( self, kind: CheckpointMutationKind = "recovery_retries" @@ -1769,6 +2198,7 @@ def reserve_prompt_group( start_weight_version=self._weight_version, task_source=recovery_policy.task_source, recovery_granularity=recovery_policy.granularity, + restore_level=recovery_policy.restore_level, admitted=admitted, admission_id=admission_id, ) @@ -1813,19 +2243,29 @@ async def run_rollout( generation_indices: Optional[list[int]] = None, on_completion: Optional[RolloutCompletionCallback] = None, recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, + gym_instance_id: Optional[str] = None, + dispatch_recorder: Optional[RolloutDispatchRecorder] = None, ) -> PromptGroupRecord: if rollout_ids is None: assert generation_indices is None assert on_completion is None assert recovery_granularity is RecoveryGranularity.SIBLING + assert gym_instance_id is None + assert dispatch_recorder is None # Legacy path: keep the impl call signature byte-identical. return await self._impl.run_rollout(input_sample) + checkpoint_ownership: dict[str, Any] = ( + {"gym_instance_id": gym_instance_id} if gym_instance_id is not None else {} + ) + if dispatch_recorder is not None: + checkpoint_ownership["dispatch_recorder"] = dispatch_recorder return await self._impl.run_rollout( input_sample, rollout_ids=rollout_ids, generation_indices=generation_indices, on_completion=on_completion, recovery_granularity=recovery_granularity, + **checkpoint_ownership, ) async def generate_and_push( @@ -2224,6 +2664,18 @@ async def _generate_for_finalization_attempt( for sibling in recovery_group.siblings if sibling.current_attempt.status != RolloutAttemptStatus.SEALED ] + pending_owners: set[str] = set() + for generation_index in pending_indices: + owner = recovery_group.siblings[ + generation_index + ].current_attempt.gym_instance_id + if owner is not None: + pending_owners.add(owner) + if len(pending_owners) > 1: + raise RuntimeError( + f"recovery group {recovery_group.group_id!r} spans multiple Gym " + f"checkpoint instances: {sorted(pending_owners)!r}" + ) group_id = recovery_group.group_id start_version = recovery_group.start_weight_version rollout_ids = tuple(recovery_group.gate_rollout_ids) @@ -2234,6 +2686,17 @@ async def _generate_for_finalization_attempt( attempt_extra_env_info[NEMO_GYM_GROUP_ATTEMPT_KEY] = ( max(len(sibling.attempts) for sibling in recovery_group.siblings) - 1 ) + gym_instance_id: Optional[str] = None + if pending_owners: + gym_instance_id = next(iter(pending_owners)) + elif pending_indices and recovery_group.restore_level.requires_gym_checkpoint: + if not isinstance(self._impl, AsyncNemoGymRolloutImpl): + raise RuntimeError( + "turn/prefix recovery requires the NeMo-Gym rollout path" + ) + gym_instance_id = self._impl.select_checkpoint_instance( + attempt_input_sample + ) self._tq_buffer.reserve( weight_version=start_version, target_step=recovery_group.target_step, @@ -2267,7 +2730,11 @@ async def _record_streamed_completion( f"streamed generation index {generation_index} is outside " f"prompt group {group_id!r}" ) - expected_gate_rollout_id = rollout_ids[generation_index] + # Read from the ledger: a row re-sent after a broken stream runs + # under a fresh attempt, not the ID this dispatch started with. + expected_gate_rollout_id = self._recovery_ledger.get_group( + group_id + ).gate_rollout_id(generation_index) if gate_rollout_id != expected_gate_rollout_id: raise ValueError( "streamed rollout identity mismatch: " @@ -2332,18 +2799,30 @@ async def _record_streamed_completion( inflight_registry[group_id] = (current_task, start_version) try: if pending_indices: - async with self._recovery_mutation() as cut: - self._recovery_ledger.mark_group_dispatched( - cut, - group_id, - generation_indices=pending_indices, + dispatch_recorder: Optional[RolloutDispatchRecorder] = None + if isinstance(self._impl, AsyncNemoGymRolloutImpl): + # The Gym impl marks each submission once Ray accepts it. + dispatch_recorder = _LedgerDispatchRecorder( + manager=self, + group_id=group_id, + gym_instance_id=gym_instance_id, ) + else: + async with self._recovery_mutation() as cut: + self._recovery_ledger.mark_group_dispatched( + cut, + group_id, + generation_indices=pending_indices, + gym_instance_id=gym_instance_id, + ) await self.run_rollout( attempt_input_sample, rollout_ids=list(rollout_ids), generation_indices=pending_indices, on_completion=_record_streamed_completion, recovery_granularity=recovery_group.recovery_granularity, + gym_instance_id=gym_instance_id, + dispatch_recorder=dispatch_recorder, ) finally: if inflight_registry is not None: diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index 0694652606d..d9e9c0c7e36 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -25,7 +25,7 @@ import copy import dataclasses import uuid -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from dataclasses import dataclass, field from enum import StrEnum from typing import TYPE_CHECKING, Any, Optional, Self, TypeAlias @@ -34,8 +34,7 @@ from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneMutationCut from nemo_rl.data.interfaces import DatumSpec -ROLLOUT_RECOVERY_SCHEMA_VERSION = 2 -_SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS = {ROLLOUT_RECOVERY_SCHEMA_VERSION} +ROLLOUT_RECOVERY_SCHEMA_VERSION = 4 ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" RolloutRecoveryState: TypeAlias = dict[str, Any] @@ -55,6 +54,7 @@ "prompt_ref", "task_source", "recovery_granularity", + "restore_level", "expected_generations", "target_step", "start_weight_version", @@ -73,6 +73,8 @@ "reward", "mask_sample", "staging_keys", + "gym_instance_id", + "gym_attempt", } ) @@ -99,12 +101,31 @@ class PromptGroupPhase(StrEnum): class RecoveryGranularity(StrEnum): - """Unit of completed work reused after a live failure or process restart.""" + """Unit retried after a live rollout failure.""" SIBLING = "sibling" PROMPT_GROUP = "prompt_group" +class RecoveryTargetLevel(StrEnum): + """Finest requested restart boundary for an unfinished rollout.""" + + PREFIX = "prefix" + TURN = "turn" + SIBLING = "sibling" + PROMPT_GROUP = "prompt_group" + + @property + def requires_gym_checkpoint(self) -> bool: + """Whether recovery needs Gym participant state for the active episode.""" + return self in {RecoveryTargetLevel.PREFIX, RecoveryTargetLevel.TURN} + + @property + def preserves_generation_prefix(self) -> bool: + """Whether recovery continues an active model call from decoded tokens.""" + return self is RecoveryTargetLevel.PREFIX + + class RolloutAttemptStatus(StrEnum): """Lifecycle of one physical Gate execution attempt.""" @@ -178,6 +199,8 @@ class RolloutAttemptRecord: reward: Optional[float] = None mask_sample: Optional[bool] = None staging_keys: list[str] = field(default_factory=list) + gym_instance_id: Optional[str] = None + gym_attempt: int = 0 @property def attempt_id(self) -> str: @@ -211,6 +234,7 @@ class PromptGroupRecoveryRecord: prompt_ref: PromptRef task_source: Optional[str] recovery_granularity: RecoveryGranularity + restore_level: RecoveryTargetLevel runtime_prompt_payload: Optional[DatumSpec] expected_generations: int target_step: Optional[int] @@ -247,12 +271,30 @@ def logical_rollout_id(self, generation_index: int) -> str: 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.""" + """Return the legacy string carrier for one structured Gym episode. + + Durable checkpoint state uses :meth:`gym_episode`. Keep this transport + encoding synchronized with Gym's ``EpisodeId`` contract; the NeMo-Gym + test suite checks the two implementations against each other. + """ sibling = self.siblings[generation_index] - return ( + base_rollout_id = ( f"{self.logical_rollout_id(generation_index)}" f"_a{sibling.current_attempt.attempt_id}" ) + gym_attempt = sibling.current_attempt.gym_attempt + return ( + base_rollout_id if gym_attempt == 0 else f"{base_rollout_id}-a{gym_attempt}" + ) + + def gym_episode(self, generation_index: int) -> tuple[str, int]: + """Return Gym's base rollout ID and attempt number for one sibling.""" + sibling = self.siblings[generation_index] + return ( + f"{self.logical_rollout_id(generation_index)}" + f"_a{sibling.current_attempt.attempt_id}", + sibling.current_attempt.gym_attempt, + ) @property def sealed_generation_indices(self) -> list[int]: @@ -284,6 +326,41 @@ class SiblingSealResult: mask_sample: bool +def _holds_gym_episode( + record: PromptGroupRecoveryRecord, sibling: RolloutSiblingRecord +) -> bool: + """Whether this sibling's current attempt is live turn/prefix Gym state. + + A dispatched attempt may have an admitted episode, and a reserved restored + attempt (``gym_attempt > 0``) carries committed Gym state forward. + """ + if ( + record.status is not PromptGroupStatus.GENERATING + or not record.restore_level.requires_gym_checkpoint + ): + return False + attempt = sibling.current_attempt + return attempt.status is RolloutAttemptStatus.DISPATCHED or ( + attempt.status is RolloutAttemptStatus.RESERVED and attempt.gym_attempt > 0 + ) + + +def _parse_gym_rollout_id(rollout_id: str) -> Optional[tuple[str, int]]: + """Invert :meth:`PromptGroupRecoveryRecord.gym_episode`'s base rollout ID. + + The ID is ``{group_id}_g{generation_index}_a{attempt_id}``. The index and + the hex attempt ID never contain ``_``, so splitting from the right is + unambiguous; callers still confirm the match against the live record. + """ + logical_rollout_id, separator, _ = rollout_id.rpartition("_a") + if not separator: + return None + group_id, separator, generation_index = logical_rollout_id.rpartition("_g") + if not separator or not generation_index.isdigit(): + return None + return group_id, int(generation_index) + + def _new_attempt() -> RolloutAttemptRecord: return RolloutAttemptRecord( attempt_uuid=uuid.uuid4(), @@ -318,6 +395,10 @@ class RolloutRecoveryLedger: def __init__(self) -> None: self._groups: dict[str, PromptGroupRecoveryRecord] = {} + # Gym episodes RL dropped after Gym may have admitted them, by owning + # Gym instance. In memory only: they belong to this process's Gym + # executions and must be retired before the next checkpoint commit. + self._gym_retirements: dict[str, set[tuple[str, int]]] = {} def groups(self) -> list[PromptGroupRecoveryRecord]: return [self._copy_group(group) for group in self._groups.values()] @@ -333,6 +414,7 @@ def reserve_group( start_weight_version: int, task_source: Optional[str] = None, recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, + restore_level: Optional[RecoveryTargetLevel] = None, admitted: bool = True, group_id: Optional[str] = None, admission_id: Optional[str] = None, @@ -360,6 +442,12 @@ def reserve_group( admission_id = admission_id or group_id if not admission_id: raise ValueError("admission_id must not be empty") + if restore_level is None: + restore_level = ( + RecoveryTargetLevel.PROMPT_GROUP + if recovery_granularity is RecoveryGranularity.PROMPT_GROUP + else RecoveryTargetLevel.SIBLING + ) siblings = [] for generation_index in range(expected_generations): @@ -376,6 +464,7 @@ def reserve_group( prompt_ref=prompt_ref, task_source=task_source, recovery_granularity=recovery_granularity, + restore_level=restore_level, # 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. @@ -426,16 +515,62 @@ def bind_runtime_prompt( ) record.runtime_prompt_payload = prompt_payload - def prepare_for_restart(self, cut: DataPlaneMutationCut) -> None: - """Apply each group's persisted restore policy to interrupted attempts.""" + def prepare_for_restart( + self, + cut: DataPlaneMutationCut, + *, + restored_gym_episodes: Optional[set[tuple[str, str, int]]] = None, + ) -> None: + """Apply each group's persisted restore policy to interrupted attempts. + + ``restored_gym_episodes`` contains ``(instance_id, rollout_id, attempt)`` + triples successfully installed by Gym. Those attempts are re-issued as + Gym's next attempt while keeping the same RL attempt and TQ lineage. + """ cut.require_live() self.assert_checkpoint_safe() + restored = restored_gym_episodes or set() + consumed_restores: set[tuple[str, str, int]] = set() for record in self._groups.values(): if record.status is PromptGroupStatus.GENERATING: if record.recovery_granularity is RecoveryGranularity.PROMPT_GROUP: self._abandon_entire_group(record) else: - self.abandon_unsealed(cut, record.group_id) + for sibling in record.siblings: + attempt = sibling.current_attempt + if attempt.status is RolloutAttemptStatus.SEALED: + continue + rollout_id, gym_attempt = record.gym_episode( + sibling.generation_index + ) + restore_key = ( + attempt.gym_instance_id or "", + rollout_id, + gym_attempt, + ) + if ( + record.restore_level.requires_gym_checkpoint + and attempt.status + in { + RolloutAttemptStatus.DISPATCHED, + RolloutAttemptStatus.RESERVED, + } + and attempt.gym_instance_id is not None + and restore_key in restored + ): + attempt.gym_attempt += 1 + attempt.status = RolloutAttemptStatus.RESERVED + consumed_restores.add(restore_key) + continue + if attempt.status is RolloutAttemptStatus.RESERVED: + continue + attempt.status = RolloutAttemptStatus.ABANDONED + unused_restores = restored - consumed_restores + if unused_restores: + raise ValueError( + "Gym restored episodes are not owned by dispatched Gym-state " + f"rollouts: {sorted(unused_restores)!r}" + ) def _abandon_entire_group(self, record: PromptGroupRecoveryRecord) -> None: """Discard every current sibling when an incomplete group is atomic. @@ -479,6 +614,128 @@ def expected_staging_keys(self) -> set[str]: for staging_key in attempt.staging_keys } + def gym_checkpoint_inventory( + self, + instance_ids: set[str] | frozenset[str], + ) -> dict[str, tuple[tuple[str, int], ...]]: + """Return active turn/prefix Gym episodes grouped by owning actor.""" + inventory: dict[str, list[tuple[str, int]]] = { + instance_id: [] for instance_id in instance_ids + } + for record in self._groups.values(): + for sibling in record.siblings: + if not _holds_gym_episode(record, sibling): + continue + attempt = sibling.current_attempt + instance_id = attempt.gym_instance_id + if instance_id is None: + raise RuntimeError( + "cannot checkpoint a dispatched Gym-state rollout without " + f"Gym ownership: group={record.group_id!r}, " + f"generation_index={sibling.generation_index}" + ) + if instance_id not in inventory: + raise RuntimeError( + f"Gym-state rollout references unavailable Gym instance " + f"{instance_id!r}" + ) + inventory[instance_id].append( + record.gym_episode(sibling.generation_index) + ) + return { + instance_id: tuple(sorted(episodes)) + for instance_id, episodes in inventory.items() + } + + def gym_episodes_still_held( + self, + episodes_by_instance: Mapping[str, Iterable[tuple[str, int]]], + ) -> dict[str, set[tuple[str, int]]]: + """Return which of these episodes the checkpoint inventory would still list. + + Looks up only the given episodes, so checking a few outstanding replies + does not cost a scan of every live rollout. + """ + held: dict[str, set[tuple[str, int]]] = {} + for instance_id, episodes in episodes_by_instance.items(): + held[instance_id] = set() + for episode in episodes: + owner = _parse_gym_rollout_id(episode[0]) + if owner is None: + continue + group_id, generation_index = owner + record = self._groups.get(group_id) + if record is None or not 0 <= generation_index < len(record.siblings): + continue + sibling = record.siblings[generation_index] + if ( + _holds_gym_episode(record, sibling) + and sibling.current_attempt.gym_instance_id == instance_id + and record.gym_episode(generation_index) == tuple(episode) + ): + held[instance_id].add(tuple(episode)) + return held + + def has_dispatched_generation_prefix(self) -> bool: + """Whether a dispatched attempt may be growing a recoverable prefix. + + Decoding advances such an attempt without mutating the data plane, so a + snapshot cadence keyed only on data-plane mutations would never cut it. + """ + return any( + record.status is PromptGroupStatus.GENERATING + and record.restore_level.preserves_generation_prefix + and sibling.current_attempt.status is RolloutAttemptStatus.DISPATCHED + for record in self._groups.values() + for sibling in record.siblings + ) + + def _record_gym_retirements(self, record: PromptGroupRecoveryRecord) -> None: + """Remember each current attempt Gym may still be running for this group. + + A dispatched attempt can have an admitted episode, and a reserved + restored attempt has restored Gym state. Dropping either from the ledger + does not stop it in Gym, which would export it at the next commit. + """ + for sibling in record.siblings: + attempt = sibling.current_attempt + holds_gym_state = attempt.status is RolloutAttemptStatus.DISPATCHED or ( + attempt.status is RolloutAttemptStatus.RESERVED + and attempt.gym_attempt > 0 + ) + if not holds_gym_state or attempt.gym_instance_id is None: + continue + self._gym_retirements.setdefault(attempt.gym_instance_id, set()).add( + record.gym_episode(sibling.generation_index) + ) + + def gym_checkpoint_retirements( + self, + instance_ids: set[str] | frozenset[str], + ) -> dict[str, tuple[tuple[str, int], ...]]: + """Return dropped Gym episodes to retire before the next commit.""" + unknown = set(self._gym_retirements) - set(instance_ids) + if unknown: + raise RuntimeError( + f"dropped rollouts reference unavailable Gym instances {sorted(unknown)!r}" + ) + return { + instance_id: tuple(sorted(self._gym_retirements.get(instance_id, ()))) + for instance_id in instance_ids + } + + def mark_gym_retired( + self, retired: Mapping[str, Iterable[tuple[str, int]]] + ) -> None: + """Forget episodes Gym has retired; later drops stay pending.""" + for instance_id, episodes in retired.items(): + pending = self._gym_retirements.get(instance_id) + if pending is None: + continue + pending.difference_update(episodes) + if not pending: + del self._gym_retirements[instance_id] + def prepare_incomplete_retry( self, cut: DataPlaneMutationCut, @@ -545,6 +802,7 @@ def mark_group_dispatched( group_id: str, *, generation_indices: Optional[list[int]] = None, + gym_instance_id: Optional[str] = None, ) -> None: """Move the selected current sibling attempts to dispatched.""" cut.require_live() @@ -566,8 +824,101 @@ def mark_group_dispatched( if any(attempt.status != RolloutAttemptStatus.RESERVED for attempt in attempts): raise ValueError("only reserved rollout attempts may be dispatched") for attempt in attempts: + if ( + record.restore_level.requires_gym_checkpoint + and attempt.gym_instance_id is None + and gym_instance_id is None + ): + raise ValueError( + "turn/prefix rollout dispatch requires a Gym checkpoint owner" + ) + if attempt.gym_instance_id is not None: + if gym_instance_id != attempt.gym_instance_id: + raise ValueError( + "restored Gym rollout must return to its checkpoint owner: " + f"expected={attempt.gym_instance_id!r}, " + f"actual={gym_instance_id!r}" + ) + elif gym_instance_id is not None: + attempt.gym_instance_id = gym_instance_id attempt.status = RolloutAttemptStatus.DISPATCHED + def release_refused_dispatch( + self, + cut: DataPlaneMutationCut, + group_id: str, + *, + generation_index: int, + ) -> None: + """Return one attempt Gym refused at admission to reserved. + + Gym refuses a ``/run`` before recording anything for its episode, so the + attempt keeps its identity and is re-dispatched as-is. A fresh attempt + leaves the Gym checkpoint inventory until then. A restored turn + (``gym_attempt > 0``) stays Gym-owned while reserved, so the next + checkpoint carries its committed continuation forward. + """ + cut.require_live() + record = self._require_group(group_id) + if record.status != PromptGroupStatus.GENERATING: + raise ValueError( + f"cannot release a refused dispatch of group {group_id!r} from " + f"{record.status.value!r}" + ) + attempt = self._require_sibling(record, generation_index).current_attempt + if attempt.status != RolloutAttemptStatus.DISPATCHED: + raise ValueError( + "only a dispatched rollout attempt can be refused: " + f"{record.logical_rollout_id(generation_index)!r} is " + f"{attempt.status.value!r}" + ) + attempt.status = RolloutAttemptStatus.RESERVED + + def abandon_dropped_dispatch( + self, + cut: DataPlaneMutationCut, + group_id: str, + *, + generation_indices: Iterable[int], + replace: bool, + ) -> tuple[tuple[str, int], ...]: + """Abandon dispatched attempts whose Gym stream ended before they returned. + + Gym runs every row of a ``/run`` as its own task, so a broken stream + stops none of them. Each Gym-owned attempt is queued for retirement and + returned so the caller can retire it right away. With ``replace``, each + abandoned sibling gets a fresh attempt, so its re-send runs under a new + rollout ID and never overlaps the old episode. Attempts that never + reached Gym keep their identity. + """ + cut.require_live() + record = self._require_group(group_id) + if record.status != PromptGroupStatus.GENERATING: + raise ValueError( + f"cannot drop dispatched attempts of group {group_id!r} from " + f"{record.status.value!r}" + ) + if replace and record.recovery_granularity is RecoveryGranularity.PROMPT_GROUP: + raise ValueError( + "prompt-group recovery replaces the whole cohort, not single siblings" + ) + dropped: list[tuple[str, int]] = [] + for generation_index in generation_indices: + sibling = self._require_sibling(record, generation_index) + attempt = sibling.current_attempt + if attempt.status is not RolloutAttemptStatus.DISPATCHED: + continue + if attempt.gym_instance_id is not None: + episode = record.gym_episode(generation_index) + self._gym_retirements.setdefault(attempt.gym_instance_id, set()).add( + episode + ) + dropped.append(episode) + attempt.status = RolloutAttemptStatus.ABANDONED + if replace: + sibling.attempts.append(_new_attempt()) + return tuple(dropped) + def mark_sibling_sealed( self, cut: DataPlaneMutationCut, @@ -706,6 +1057,7 @@ def abandon_unsealed(self, cut: DataPlaneMutationCut, group_id: str) -> None: raise ValueError( f"cannot abandon group {group_id!r} from {record.status.value!r}" ) + self._record_gym_retirements(record) if ( record.recovery_granularity is RecoveryGranularity.PROMPT_GROUP and record.status is PromptGroupStatus.GENERATING @@ -798,7 +1150,7 @@ def mark_finalization_unknown( 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() - self._require_group(group_id) + self._record_gym_retirements(self._require_group(group_id)) del self._groups[group_id] def discard_canonical_groups( @@ -851,6 +1203,7 @@ def state_dict(self) -> dict[str, Any]: }, "task_source": record.task_source, "recovery_granularity": record.recovery_granularity.value, + "restore_level": record.restore_level.value, "expected_generations": record.expected_generations, "target_step": record.target_step, "start_weight_version": record.start_weight_version, @@ -867,6 +1220,8 @@ def state_dict(self) -> dict[str, Any]: "reward": attempt.reward, "mask_sample": attempt.mask_sample, "staging_keys": list(attempt.staging_keys), + "gym_instance_id": attempt.gym_instance_id, + "gym_attempt": attempt.gym_attempt, } for attempt in sibling.attempts ], @@ -894,13 +1249,10 @@ def from_state_dict(cls, state: dict[str, Any]) -> Self: context="rollout recovery state", ) 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 - ): + if schema_version != ROLLOUT_RECOVERY_SCHEMA_VERSION: raise ValueError( - f"Unsupported rollout-recovery schema version: {schema_version!r}" + f"Unsupported rollout-recovery schema version: {schema_version!r}; " + f"expected {ROLLOUT_RECOVERY_SCHEMA_VERSION}" ) raw_groups = state.get("groups") if not isinstance(raw_groups, list): @@ -961,6 +1313,7 @@ def _group_from_state( prompt_id = raw_group.get("prompt_id") task_source = raw_group.get("task_source") raw_recovery_granularity = raw_group.get("recovery_granularity") + raw_restore_level = raw_group.get("restore_level") expected_generations = raw_group.get("expected_generations") siblings_state = raw_group.get("siblings") if not isinstance(group_id, str) or not group_id: @@ -979,6 +1332,12 @@ def _group_from_state( raise ValueError( f"invalid recovery_granularity={raw_recovery_granularity!r}" ) from error + if not isinstance(raw_restore_level, str): + raise ValueError("restore_level must be a string") + try: + restore_level = RecoveryTargetLevel(raw_restore_level) + except ValueError as error: + raise ValueError(f"invalid restore_level={raw_restore_level!r}") from error if not isinstance(expected_generations, int) or expected_generations < 1: raise ValueError("expected_generations must be a positive integer") if ( @@ -1038,7 +1397,7 @@ def _group_from_state( 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}" + base_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( @@ -1054,6 +1413,27 @@ def _group_from_state( reward = attempt_state.get("reward") mask_sample = attempt_state.get("mask_sample") staging_keys = attempt_state.get("staging_keys") + gym_instance_id = attempt_state.get("gym_instance_id") + gym_attempt = attempt_state.get("gym_attempt", 0) + if gym_instance_id is not None and ( + not isinstance(gym_instance_id, str) or not gym_instance_id + ): + raise ValueError( + "gym_instance_id must be a non-empty string or None" + ) + if ( + not isinstance(gym_attempt, int) + or isinstance(gym_attempt, bool) + or gym_attempt < 0 + ): + raise ValueError("gym_attempt must be a non-negative integer") + if gym_attempt > 0 and gym_instance_id is None: + raise ValueError("resumed Gym attempts require gym_instance_id") + gate_id = ( + base_gate_id + if gym_attempt == 0 + else f"{base_gate_id}-a{gym_attempt}" + ) if not isinstance(staging_keys, list) or not all( isinstance(key, str) for key in staging_keys ): @@ -1094,6 +1474,8 @@ def _group_from_state( reward=float(reward) if reward is not None else None, mask_sample=mask_sample, staging_keys=list(staging_keys), + gym_instance_id=gym_instance_id, + gym_attempt=gym_attempt, ) ) siblings.append( @@ -1148,6 +1530,7 @@ def _group_from_state( prompt_ref=PromptRef(sample_id=sample_id, task_name=task_name), task_source=task_source, recovery_granularity=recovery_granularity, + restore_level=restore_level, runtime_prompt_payload=None, expected_generations=expected_generations, target_step=target_step, @@ -1248,15 +1631,10 @@ def parse_rollout_recovery_state(state: object) -> ParsedRolloutRecoveryState: context="rollout recovery sidecar", ) 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 - ): + if schema_version != ROLLOUT_RECOVERY_SCHEMA_VERSION: raise ValueError( "unsupported rollout recovery schema_version=" - f"{schema_version!r}; supported versions are " - f"{sorted(_SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS)}" + f"{schema_version!r}; expected {ROLLOUT_RECOVERY_SCHEMA_VERSION}" ) groups = state.get("groups") if not isinstance(groups, list): diff --git a/nemo_rl/models/generation/generation_cut_capture.py b/nemo_rl/models/generation/generation_cut_capture.py new file mode 100644 index 00000000000..0e2b4f0f44c --- /dev/null +++ b/nemo_rl/models/generation/generation_cut_capture.py @@ -0,0 +1,821 @@ +# 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. + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any, Literal + +LOGGER = logging.getLogger(__name__) + +# Completion evidence only bridges the short race between a terminal response +# and Gym freezing its admitted-call inventory. Time-based retention avoids +# evicting fresh evidence merely because a large deployment completed more +# than 100k calls, while still bounding idle-job memory growth. +_COMPLETED_CAPTURE_RETENTION_S = 60.0 * 60.0 + + +@dataclass +class _RequestCaptureBuffer: + """One append-only generation segment owned by an in-flight request.""" + + sequence: int + generated_token_ids: list[int] = field(default_factory=list) + generated_logprobs: list[float] = field(default_factory=list) + + +@dataclass +class _FrozenRequestCaptureBuffer: + """An active buffer detached for a checkpoint write.""" + + checkpoint_id: str + buffer: _RequestCaptureBuffer + + +@dataclass +class _RequestCaptureState: + """Append-only token buffers for one in-flight captured request. + + vLLM publishes delta progress. ``observe`` appends each delta to the current + active buffer exactly once. A checkpoint swaps the active buffer under + ``lock`` and performs its blocking TQ write after releasing the lock, so + later observations can continue in a fresh buffer. + + Each successful cut stages only the detached buffer. The ordered TQ keys + remain in Gym lineage, while the sealed token arrays are retained locally + only to compute the cumulative digest advertised by the latest cut. + """ + + call: Any + prompt_token_ids: list[int] + # Backend-owned opaque state that must survive until the terminal row is + # staged. vLLM uses this for its processed-media snapshot; prefix recovery + # is currently setup-guarded to text-only, but ordinary terminal capture + # still shares this lifecycle container. + media: Any | None = None + resumed_generation_token_ids: list[int] = field(default_factory=list) + effective_output_limit: int | None = None + terminal_finish_reason: Literal["stop", "length"] | None = None + terminal_stop_reason: str | int | None = None + observation_error: str | None = None + lock: threading.Lock = field(default_factory=threading.Lock) + # Serializes lifecycle-changing TQ operations for this call. Token + # observation deliberately uses ``lock`` instead so decoding can continue + # while a detached generation chunk is staged. + lifecycle_lock: threading.Lock = field(default_factory=threading.Lock) + terminal_started: bool = False + sealed_generated_token_ids: list[int] = field(default_factory=list) + sealed_generated_logprobs: list[float] = field(default_factory=list) + generation_cut_staging_keys: list[str] = field(default_factory=list) + active_buffer: _RequestCaptureBuffer = field( + default_factory=lambda: _RequestCaptureBuffer(sequence=0) + ) + frozen_buffer: _FrozenRequestCaptureBuffer | None = None + observed_generation_token_count: int = 0 + next_buffer_sequence: int = 1 + + def observe( + self, + generated_token_ids: list[int], + generated_logprobs: list[float], + ) -> None: + """Append one immutable vLLM output delta.""" + if len(generated_token_ids) != len(generated_logprobs): + self.observation_error = ( + "generated token IDs and log probabilities must have equal lengths" + ) + return + self.active_buffer.generated_token_ids.extend(generated_token_ids) + self.active_buffer.generated_logprobs.extend(generated_logprobs) + self.observed_generation_token_count += len(generated_token_ids) + + def freeze_for_checkpoint( + self, checkpoint_id: str + ) -> tuple[str, int, list[int], list[float], list[int], list[float]]: + """Swap the active buffer and return its delta plus the stable prefix.""" + if self.frozen_buffer is not None: + raise RuntimeError( + "cannot start a generation cut while another cut is in progress" + ) + buffer = self.active_buffer + self.active_buffer = _RequestCaptureBuffer(sequence=self.next_buffer_sequence) + self.next_buffer_sequence += 1 + self.frozen_buffer = _FrozenRequestCaptureBuffer( + checkpoint_id=checkpoint_id, + buffer=buffer, + ) + return ( + f"active/{checkpoint_id}/{buffer.sequence}", + buffer.sequence, + list(buffer.generated_token_ids), + list(buffer.generated_logprobs), + self.sealed_generated_token_ids + buffer.generated_token_ids, + self.sealed_generated_logprobs + buffer.generated_logprobs, + ) + + def seal_frozen_buffer(self, checkpoint_id: str) -> None: + """Adopt a successfully staged frozen buffer into the live prefix.""" + frozen = self._require_frozen_buffer(checkpoint_id) + self.sealed_generated_token_ids.extend(frozen.generated_token_ids) + self.sealed_generated_logprobs.extend(frozen.generated_logprobs) + self.frozen_buffer = None + + def rollback_frozen_buffer(self, checkpoint_id: str) -> None: + """Restore a failed cut ahead of progress collected after its swap.""" + frozen = self._require_frozen_buffer(checkpoint_id) + self.active_buffer.generated_token_ids[:0] = frozen.generated_token_ids + self.active_buffer.generated_logprobs[:0] = frozen.generated_logprobs + self.frozen_buffer = None + + def _require_frozen_buffer(self, checkpoint_id: str) -> _RequestCaptureBuffer: + frozen = self.frozen_buffer + if frozen is None or frozen.checkpoint_id != checkpoint_id: + raise RuntimeError( + f"generation cut {checkpoint_id!r} does not own the frozen buffer" + ) + return frozen.buffer + + +@dataclass(frozen=True) +class _CompletedCaptureState: + """Time-bounded terminal evidence retained across the response/cut race.""" + + coords: Any + generation_token_count: int + effective_output_limit: int | None + terminal_finish_reason: Literal["stop", "length"] | None + terminal_stop_reason: str | int | None + completed_at_monotonic: float + + +def _remaining_generation_limits_after_prefix( + *, + max_tokens: int | None, + min_tokens: int | None, + prefix_token_count: int, +) -> tuple[int | None, int | None]: + """Return output limits for the suffix after restoring generated tokens.""" + if prefix_token_count < 0: + raise ValueError("prefix_token_count must be non-negative") + return ( + None if max_tokens is None else max_tokens - prefix_token_count, + None if min_tokens is None else max(0, min_tokens - prefix_token_count), + ) + + +# One character spans at most four UTF-8 bytes, hence at most four byte-level +# tokens; decoding a slightly longer tail is enough to see whether it is whole. +_PREFIX_BOUNDARY_TAIL_TOKENS = 8 + + +def _ends_inside_character(tokenizer: Any, generated_token_ids: list[int]) -> bool: + """Whether the generated text stops partway through a multi-byte character. + + vLLM detokenizes a resumed call's tail with a stream primed on the prompt, + which ends with the restored prefix. A prefix ending inside a character + makes that stream emit the whole decoded prompt as the tail's first text. + """ + tail = generated_token_ids[-_PREFIX_BOUNDARY_TAIL_TOKENS:] + return tokenizer.decode(tail, skip_special_tokens=False).endswith("\ufffd") + + +class _TokenCaptureSnapshotGate: + """Drain active terminal writes and block new ones across a TQ snapshot. + + Each fence takes an epoch before its close is queued. ``reopen`` releases + every epoch issued so far, so a close that runs after the driver already + timed out and released its fence is a no-op instead of a stuck gate. + """ + + def __init__(self) -> None: + self._condition = threading.Condition() + self._open = True + self._active = 0 + self._epoch = 0 + self._released_epoch = 0 + + def enter(self) -> None: + with self._condition: + while not self._open: + self._condition.wait() + self._active += 1 + + def exit(self) -> None: + with self._condition: + self._active -= 1 + if self._active == 0: + self._condition.notify_all() + + def begin_epoch(self) -> int: + with self._condition: + self._epoch += 1 + return self._epoch + + def close_and_wait(self, epoch: int) -> None: + with self._condition: + if epoch <= self._released_epoch: + return + self._open = False + while self._active and self._released_epoch < epoch: + self._condition.wait() + + def reopen(self) -> None: + with self._condition: + self._released_epoch = self._epoch + self._open = True + self._condition.notify_all() + + +class GenerationCutCaptureMixin: + """Registry, checkpoint-cut, and resume logic shared by generation backends. + + Methods use the capture-host state initialized by the vLLM worker and + ``_RequestCaptureState``; nothing here touches a generation engine. + """ + + # Initialized by the vLLM worker; declared so type checkers can see the + # insertion-ordered dicts this mixin evicts from. + _completed_capture_calls: dict[str, _CompletedCaptureState] + _generation_cut_receipts: dict[tuple[str, str], Any] + _generation_cut_tokenizer: Any + + def _pop_request_capture(self, request: Any) -> _RequestCaptureState | None: + with self._capture_registry_lock: + state = self._capture_calls.pop(id(request), None) + if state is not None: + self._capture_calls_by_model_call_id.pop(state.call.model_call_id, None) + return state + + def _get_request_capture(self, request: Any) -> _RequestCaptureState | None: + with self._capture_registry_lock: + return self._capture_calls.get(id(request)) + + def _remember_completed_capture( + self, + model_call_id: str, + coords: Any, + generation_token_count: int, + *, + effective_output_limit: int | None = None, + terminal_finish_reason: Literal["stop", "length"] | None = None, + terminal_stop_reason: str | int | None = None, + ) -> None: + now = time.monotonic() + with self._capture_registry_lock: + # Dict insertion order is completion order. Remove only entries old + # enough that the response/inventory race has certainly elapsed; + # never discard fresh evidence because the deployment crossed an + # arbitrary call-count threshold. + while self._completed_capture_calls: + oldest_id = next(iter(self._completed_capture_calls)) + oldest = self._completed_capture_calls[oldest_id] + if ( + now - oldest.completed_at_monotonic + <= _COMPLETED_CAPTURE_RETENTION_S + ): + break + self._completed_capture_calls.pop(oldest_id) + self._completed_capture_calls.pop(model_call_id, None) + self._completed_capture_calls[model_call_id] = _CompletedCaptureState( + coords=coords, + generation_token_count=generation_token_count, + effective_output_limit=effective_output_limit, + terminal_finish_reason=terminal_finish_reason, + terminal_stop_reason=terminal_stop_reason, + completed_at_monotonic=now, + ) + + def _fetch_generation_cut_chunks(self, continuation: Any) -> list[Any]: + """Fetch the staged chunks named by a ``GenerationCutContinuation``, in order.""" + if self._staging_source is None: + raise RuntimeError( + "_staging_source not initialized; call setup_token_capture() first" + ) + staging_keys = list(continuation.staging_keys) + snapshots = self._staging_source.fetch(staging_keys) + if len(snapshots) != len(staging_keys): + raise RuntimeError("generation-cut fetch did not return every staged chunk") + if not snapshots: + raise RuntimeError("generation-cut continuation has no staged chunks") + return snapshots + + def _rebuild_generation_cut_snapshot( + self, admission: Any, prefix_token_ids: list[int], snapshots: list[Any] + ) -> Any: + """Rebuild and verify the cumulative staged prefix from its fetched chunks.""" + continuation = admission.generation_cut + weight_versions = [snapshot.weight_version for snapshot in snapshots] + if weight_versions != sorted(weight_versions): + raise RuntimeError( + "generation-cut chunk policy versions are not monotonically non-decreasing" + ) + current_weight_version = self._rollout_weight_version + if weight_versions[-1] > current_weight_version: + raise RuntimeError( + "generation-cut prefix contains policy version " + f"{weight_versions[-1]}, newer than the current rollout version " + f"{current_weight_version}" + ) + token_ids_delta = [ + token_id for snapshot in snapshots for token_id in snapshot.token_ids_delta + ] + token_mask_delta = [ + mask for snapshot in snapshots for mask in snapshot.token_mask_delta + ] + generation_logprobs_delta = [ + logprob + for snapshot in snapshots + for logprob in snapshot.generation_log_probs_delta + ] + if any( + mask == 0.0 + for snapshot in snapshots[1:] + for mask in snapshot.token_mask_delta + ): + raise RuntimeError( + "only the first generation-cut chunk may contain prompt tokens" + ) + + # Deferred: nemo_gym is an optional extra absent in non-Gym runs. + from nemo_gym.token_id_capture.staging.digest import ( + EXTRAS_DIGEST_VERSION, + STAGING_DIGEST_VERSION, + compute_chain_hash, + compute_extras_digest, + compute_staging_digest, + hash_token_ids, + ) + from nemo_gym.token_id_capture.staging.records import StagedCallBaseSnapshot + + delta_len = len(token_ids_delta) + cum_len = admission.prev_len + delta_len + # A cut chain may span retries performed after later weight updates. + # Its token-level behavior logprobs remain attached to their original + # chunks. The cumulative snapshot uses the oldest contributing version + # so downstream replay-buffer staleness checks remain conservative. + weight_version = weight_versions[0] + extras_digest = compute_extras_digest(None) + chain_hash = compute_chain_hash(admission.parent_chain_hash, token_ids_delta) + cumulative_hash = hash_token_ids(prefix_token_ids + token_ids_delta) + digest = compute_staging_digest( + schema_version=admission.schema_version, + digest_version=STAGING_DIGEST_VERSION, + extras_digest_version=EXTRAS_DIGEST_VERSION, + rollout_id=continuation.source_capture_key, + model_call_id=continuation.source_model_call_id, + parent_call_id=admission.parent_call_id, + mode=admission.mode, + prev_len=admission.prev_len, + delta_len=delta_len, + cum_len=cum_len, + weight_version=weight_version, + token_ids_delta=token_ids_delta, + token_mask_delta=token_mask_delta, + generation_log_probs_delta=generation_logprobs_delta, + extras_digest=extras_digest, + chain_hash=chain_hash, + cumulative_hash=cumulative_hash, + ) + prefix_token_count = sum(mask == 1.0 for mask in token_mask_delta) + if prefix_token_count != continuation.prefix_token_count: + message = ( + "generation-cut token count mismatch: " + f"expected={continuation.prefix_token_count} " + f"actual={prefix_token_count} " + f"source_capture_key={continuation.source_capture_key!r} " + f"source_model_call_id={continuation.source_model_call_id!r}" + ) + LOGGER.error("%s", message) + raise RuntimeError(message) + if digest != continuation.prefix_digest: + parent_chain_hash_prefix = ( + admission.parent_chain_hash[:12] + if admission.parent_chain_hash is not None + else None + ) + message = ( + "generation-cut digest mismatch: " + f"expected={continuation.prefix_digest} actual={digest} " + f"source_capture_key={continuation.source_capture_key!r} " + f"source_model_call_id={continuation.source_model_call_id!r} " + f"parent_call_id={admission.parent_call_id!r} " + f"mode={admission.mode!r} prev_len={admission.prev_len} " + f"delta_len={delta_len} cum_len={cum_len} " + f"weight_version={weight_version} " + f"schema_version={admission.schema_version} " + f"parent_chain_hash_prefix={parent_chain_hash_prefix!r} " + f"chain_hash_prefix={chain_hash[:12]!r} " + f"cumulative_hash_prefix={cumulative_hash[:12]!r}" + ) + LOGGER.error("%s", message) + raise RuntimeError(message) + snapshot = StagedCallBaseSnapshot( + rollout_id=continuation.source_capture_key, + model_call_id=continuation.source_model_call_id, + parent_call_id=admission.parent_call_id, + mode=admission.mode, + prev_len=admission.prev_len, + delta_len=delta_len, + cum_len=cum_len, + weight_version=weight_version, + digest=digest, + token_ids_delta=token_ids_delta, + token_mask_delta=token_mask_delta, + generation_log_probs_delta=generation_logprobs_delta, + extras_digest=extras_digest, + chain_hash=chain_hash, + cumulative_hash=cumulative_hash, + ) + # Lets a consumer confirm the exact generated tokens it later receives + # (e.g. a trained row) without access to the staging store. + generated_prefix_ids = [ + token_id + for token_id, mask in zip(token_ids_delta, token_mask_delta, strict=True) + if mask == 1.0 + ] + prefix_ids_sha256 = hashlib.sha256( + ",".join(map(str, generated_prefix_ids)).encode() + ).hexdigest() + LOGGER.info( + "generation prefix restored: rollout_id=%s model_call_id=%s " + "source_model_call_id=%s prefix_tokens=%d prefix_digest=%s " + "weight_version_span=[%d,%d] prefix_ids_sha256=%s", + admission.rollout_id, + admission.model_call_id, + continuation.source_model_call_id, + continuation.prefix_token_count, + continuation.prefix_digest, + weight_versions[0], + weight_versions[-1], + prefix_ids_sha256, + ) + return snapshot + + def _resolve_generation_cut( + self, admission: Any, prefix_token_ids: list[int] + ) -> Any | None: + """Fetch and rebuild the cumulative staged prefix named by an admission.""" + continuation = admission.generation_cut + if continuation is None: + return None + # Class-qualified so a host that borrows this method alone (tests bind the + # worker's methods onto plain namespaces) needs nothing else. + snapshots = GenerationCutCaptureMixin._fetch_generation_cut_chunks( + self, continuation + ) + return GenerationCutCaptureMixin._rebuild_generation_cut_snapshot( + self, admission, prefix_token_ids, snapshots + ) + + def _abort_request_capture(self, request: Any, *, reason: str) -> None: + """Drop the in-flight capture state for a request that errored.""" + state = self._get_request_capture(request) + if state is None: + return + with state.lifecycle_lock: + if state.terminal_started: + return + state.terminal_started = True + popped = self._pop_request_capture(request) + if popped is not state or self.token_capture is None: + return + coords = self.token_capture.fail_call(state.call, reason=reason) + self._remember_completed_capture(state.call.model_call_id, coords, 0) + + async def _run_generation_cut_control( + self, + operation: Callable[..., Any], + *args: Any, + ) -> Any: + """Run cut control independently of blocked terminal capture writes.""" + return await asyncio.get_running_loop().run_in_executor( + self._generation_cut_control_executor, + operation, + *args, + ) + + def _completed_generation_cut_ack(self, prefix: Any, checkpoint_id: str) -> Any: + """Describe a call whose terminal path won the cut/abort race.""" + from nemo_gym._checkpoint.generation_cut import ( + GenerationCutPrefixAck, + ) + from nemo_gym.token_id_capture.staging.records import staging_key + + with self._capture_registry_lock: + completed = self._completed_capture_calls.get(prefix.model_call_id) + expected_capture_key = ( + prefix.rollout_id + if prefix.attempt == 0 + else f"{prefix.rollout_id}-a{prefix.attempt}" + ) + if ( + completed is not None + and completed.coords.rollout_id != expected_capture_key + ): + raise RuntimeError( + "generation-prefix inventory identity does not match the " + f"completed call: model_call_id={prefix.model_call_id!r}" + ) + if ( + completed is not None + and completed.coords.disposition == "staged" + and completed.effective_output_limit is not None + and completed.terminal_finish_reason is not None + ): + return GenerationCutPrefixAck( + **prefix.model_dump(mode="json"), + disposition="durable_prefix", + cut_kind="terminal_completion", + frozen_buffer_id=f"terminal/{checkpoint_id}", + staging_keys=(completed.coords.staging_key,), + prefix_token_count=completed.generation_token_count, + prefix_digest=completed.coords.digest, + effective_output_limit=completed.effective_output_limit, + terminal_finish_reason=completed.terminal_finish_reason, + terminal_stop_reason=completed.terminal_stop_reason, + ) + # The in-memory cache is only an optimization. Its durable canonical + # row remains authoritative if the race evidence aged out before a + # delayed checkpoint inventory arrived. + if completed is None and self._staging_source is not None: + canonical_key = staging_key(expected_capture_key, prefix.model_call_id) + try: + snapshots = self._staging_source.fetch([canonical_key]) + except KeyError: + snapshots = [] + if snapshots: + if len(snapshots) != 1: + raise RuntimeError( + "terminal completion lookup returned an unexpected row count" + ) + snapshot = snapshots[0] + if ( + snapshot.rollout_id != expected_capture_key + or snapshot.model_call_id != prefix.model_call_id + ): + raise RuntimeError( + "generation-prefix inventory identity does not match the " + f"durable terminal row: model_call_id={prefix.model_call_id!r}" + ) + # The canonical token row proves completion, but it predates + # the recovery contract's effective budget and finish reason. + # Without those fields we cannot reproduce the same terminal + # API result, so restart this attempt from its prior boundary. + LOGGER.warning( + "terminal completion row lacks recovery metadata: model_call_id=%s", + prefix.model_call_id, + ) + return GenerationCutPrefixAck( + **prefix.model_dump(mode="json"), + disposition="durable_failure", + ) + + def _checkpoint_active_generation_cut( + self, + prefix: Any, + state: _RequestCaptureState, + checkpoint_id: str, + ) -> Any | None: + """Cut one live call, or return ``None`` if abort won the race.""" + from nemo_gym._checkpoint.generation_cut import ( + GenerationCutPrefixAck, + ) + + capture = self.token_capture + sink = self._capture_sink + if capture is None or sink is None: + raise RuntimeError("generation-prefix cuts require token capture setup") + + # Abort and terminal completion use the same lifecycle lock. Once the + # cut owns it, the call cannot be removed until its detached buffer is + # staged or rolled back. If abort won first, use its terminal evidence. + with state.lifecycle_lock: + with self._capture_registry_lock: + current = self._capture_calls_by_model_call_id.get(prefix.model_call_id) + if current is not state or state.terminal_started: + return None + + expected_capture_key = ( + prefix.rollout_id + if prefix.attempt == 0 + else f"{prefix.rollout_id}-a{prefix.attempt}" + ) + if state.call.rollout_id != expected_capture_key: + raise RuntimeError( + "generation-prefix inventory identity does not match the " + f"active call: model_call_id={prefix.model_call_id!r}, " + f"inventory_rollout_id={prefix.rollout_id!r}, " + f"worker_rollout_id={state.call.rollout_id!r}" + ) + + with state.lock: + observation_error = state.observation_error + effective_output_limit = state.effective_output_limit + terminal_finish_reason = state.terminal_finish_reason + terminal_stop_reason = state.terminal_stop_reason + if observation_error is None: + ( + frozen_buffer_id, + chunk_sequence, + chunk_token_ids, + chunk_logprobs, + generated_token_ids, + generated_logprobs, + ) = state.freeze_for_checkpoint(checkpoint_id) + if observation_error is not None: + # A dropped delta leaves the buffer short of what vLLM decoded. + # Only this call regenerates; the rest of the inventory is cut. + LOGGER.warning( + "cannot cut model call %s: %s", + prefix.model_call_id, + observation_error, + ) + return GenerationCutPrefixAck( + **prefix.model_dump(mode="json"), + disposition="durable_failure", + ) + inherited_prefix_token_count = ( + state.call.admission.generation_cut.prefix_token_count + if state.call.admission.generation_cut is not None + else 0 + ) + total_generation_token_count = inherited_prefix_token_count + len( + generated_token_ids + ) + if total_generation_token_count == 0 or ( + generated_token_ids + and _ends_inside_character( + self._generation_cut_tokenizer, + state.resumed_generation_token_ids + generated_token_ids, + ) + ): + with state.lock: + state.rollback_frozen_buffer(checkpoint_id) + return GenerationCutPrefixAck( + **prefix.model_dump(mode="json"), + disposition="durable_failure", + ) + if effective_output_limit is None: + raise RuntimeError( + f"cannot cut model call {prefix.model_call_id!r}: " + "vLLM did not resolve its effective output limit" + ) + + staged_key = None + try: + cumulative_record = capture.build_prefix_record( + state.call, + prompt_token_ids=state.prompt_token_ids, + generated_token_ids=generated_token_ids, + generated_logprobs=generated_logprobs, + ) + if chunk_token_ids: + chunk_record = ( + cumulative_record + if not state.generation_cut_staging_keys + else capture.build_generation_chunk_record( + state.call, + generated_token_ids=chunk_token_ids, + generated_logprobs=chunk_logprobs, + ) + ) + result = sink.stage_generation_prefix( + chunk_record, + checkpoint_id=checkpoint_id, + chunk_sequence=chunk_sequence, + ) + staged_key = result.staging_key + if not result.ok: + raise RuntimeError( + "generation-prefix staging failed for " + f"{prefix.model_call_id!r}: {result.error}" + ) + if staged_key is None: + raise RuntimeError( + "generation-prefix staging returned no staging key for " + f"{prefix.model_call_id!r}" + ) + with state.lock: + candidate_staging_keys = tuple(state.generation_cut_staging_keys) + if staged_key is not None: + candidate_staging_keys += (staged_key,) + acknowledgement = GenerationCutPrefixAck( + **prefix.model_dump(mode="json"), + disposition="durable_prefix", + cut_kind="active_prefix", + frozen_buffer_id=frozen_buffer_id, + staging_keys=candidate_staging_keys, + prefix_token_count=total_generation_token_count, + prefix_digest=cumulative_record.digest, + effective_output_limit=effective_output_limit, + terminal_finish_reason=terminal_finish_reason, + terminal_stop_reason=terminal_stop_reason, + ) + # Validate the complete acknowledgement before adopting the + # staged chunk. A failed acknowledgement must leave the + # frozen buffer available for rollback. + state.seal_frozen_buffer(checkpoint_id) + if staged_key is not None: + state.generation_cut_staging_keys.append(staged_key) + except Exception: + with state.lock: + frozen = state.frozen_buffer + if frozen is not None and frozen.checkpoint_id == checkpoint_id: + state.rollback_frozen_buffer(checkpoint_id) + if staged_key is not None: + try: + sink.clear([staged_key]) + except Exception: # noqa: BLE001 — preserve the cut failure + LOGGER.exception( + "Failed to clear rejected generation-prefix row %s", + staged_key, + ) + raise + return acknowledgement + + def _checkpoint_generation_cut( + self, inventory: Any, *, deadline: float | None = None + ) -> Any: + """Stage a stable prefix for every call named by Gym's frozen inventory. + + ``deadline`` is a ``time.monotonic()`` bound. Gym stops waiting for the + receipt at its own control deadline, so calls not reached by then are + acknowledged as failures (they regenerate after restore) instead of + holding the single cut-control thread for the next checkpoint. + """ + from nemo_gym._checkpoint.generation_cut import ( + GenerationCutPrefixAck, + GenerationCutReceipt, + ) + + capture = self.token_capture + sink = self._capture_sink + if capture is None or sink is None: + raise RuntimeError("generation-prefix cuts require token capture setup") + receipt_key = (inventory.checkpoint_id, inventory.inventory_digest) + with self._capture_registry_lock: + cached_receipt = self._generation_cut_receipts.get(receipt_key) + if cached_receipt is not None: + return cached_receipt + acknowledgements = [] + for index, prefix in enumerate(inventory.active_prefixes): + if deadline is not None and time.monotonic() >= deadline: + remaining = inventory.active_prefixes[index:] + LOGGER.warning( + "generation cut passed its control deadline: checkpoint_id=%s " + "uncut=%d of %d", + inventory.checkpoint_id, + len(remaining), + len(inventory.active_prefixes), + ) + acknowledgements.extend( + GenerationCutPrefixAck.failure(uncut) for uncut in remaining + ) + break + with self._capture_registry_lock: + state = self._capture_calls_by_model_call_id.get(prefix.model_call_id) + if state is None: + acknowledgements.append( + self._completed_generation_cut_ack(prefix, inventory.checkpoint_id) + ) + continue + acknowledgement = self._checkpoint_active_generation_cut( + prefix, state, inventory.checkpoint_id + ) + if acknowledgement is None: + acknowledgement = self._completed_generation_cut_ack( + prefix, inventory.checkpoint_id + ) + acknowledgements.append(acknowledgement) + receipt = GenerationCutReceipt( + checkpoint_id=inventory.checkpoint_id, + cut_id=f"worker-{inventory.inventory_digest}", + inventory_digest=inventory.inventory_digest, + inventory=inventory, + backend_snapshot_id=f"tq-{inventory.inventory_digest}", + prefixes=tuple(acknowledgements), + ) + with self._capture_registry_lock: + self._generation_cut_receipts[receipt_key] = receipt + if len(self._generation_cut_receipts) > 256: + self._generation_cut_receipts.pop( + next(iter(self._generation_cut_receipts)) + ) + return receipt diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 3a87e7d945f..b0d2d12ef6c 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -618,6 +618,9 @@ def setup_token_capture( staging_partition: str, *, capture_media: bool = False, + generation_prefix_cuts_enabled: bool = False, + generation_cut_control_token: str | None = None, + generation_cut_control_timeout_s: float | None = None, ) -> None: """Install token capture in the serving workers (``token_capture.enabled``). @@ -630,6 +633,12 @@ def setup_token_capture( staging_partition: Data-plane partition that captured rows are staged in. capture_media: Also stage the processed VLM media each call ran on beside its token delta (vLLM only; see ``MEDIA_STAGING_FIELDS``). + generation_prefix_cuts_enabled: Allow checkpoint-time cuts of active + model calls for prefix recovery. + generation_cut_control_token: Bearer token used by Gym's checkpoint + coordinator to authorize generation-cut requests. + generation_cut_control_timeout_s: Bound on one generation-cut + request; calls not cut by then are acknowledged as failures. """ raise NotImplementedError( f"token_capture.enabled is not supported for {type(self).__name__}" @@ -685,6 +694,24 @@ def resume_generation_after_refit(self) -> bool: _warn_unsupported_in_flight_refit_pause_once(type(self).__name__) return False + def begin_token_capture_snapshot_fence( + self, *, timeout_s: Optional[float] = None + ) -> bool: + """Fence terminal token writes while a coordinated TQ snapshot is taken.""" + raise NotImplementedError( + "generation-prefix checkpointing is not supported for " + f"{type(self).__name__}" + ) + + def end_token_capture_snapshot_fence( + self, *, timeout_s: Optional[float] = None + ) -> bool: + """Release terminal token writes after the coordinated TQ snapshot.""" + raise NotImplementedError( + "generation-prefix checkpointing is not supported for " + f"{type(self).__name__}" + ) + def blocks_training(self) -> bool: """Whether this engine must stand down before a training step. diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index 6476db8a689..61b82c93e1b 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -640,8 +640,15 @@ def setup_token_capture( staging_partition: str, *, capture_media: bool = False, + generation_prefix_cuts_enabled: bool = False, + generation_cut_control_token: str | None = None, + generation_cut_control_timeout_s: float | None = None, ) -> None: """Install MInf's canonical prompt and completion capture hooks.""" + if generation_prefix_cuts_enabled: + raise NotImplementedError( + "generation-prefix recovery is not implemented for Megatron inference" + ) if capture_media: raise NotImplementedError( "Media token capture is only implemented for the vLLM generation " diff --git a/nemo_rl/models/generation/vllm/utils.py b/nemo_rl/models/generation/vllm/utils.py index 84624afad27..614fc25d666 100644 --- a/nemo_rl/models/generation/vllm/utils.py +++ b/nemo_rl/models/generation/vllm/utils.py @@ -51,6 +51,56 @@ # The expert-id range vs carry dtype is model-constant, so it is verified on the # first non-empty routed-experts tensor per process and skipped afterwards. G_ROUTED_EXPERTS_RANGE_CHECKED = False + + +def _extract_selected_token_logprobs( + generation_token_ids: list[int], generation_logprob_details: list[Any] +) -> list[float]: + """Extract the sampled token's log probability at every output position.""" + selected_logprobs: list[float] = [] + for token_id, position_logprobs in zip( + generation_token_ids, generation_logprob_details, strict=True + ): + selected = position_logprobs.get(token_id) + if selected is None: + raise RuntimeError( + "vLLM generation log probabilities did not include the selected " + f"token: token_id={token_id}" + ) + selected_logprobs.append(max(float(selected.logprob), VLLM_LOGPROB_FLOOR)) + return selected_logprobs + + +def _snapshot_generation_logprob_details( + generation_logprob_details: Any, +) -> list[Any]: + """Copy only fully published positions from an append-only container.""" + position_count = len(generation_logprob_details) + end_indices = getattr(generation_logprob_details, "end_indices", None) + if end_indices is not None: + position_count = min(position_count, len(end_indices)) + return [generation_logprob_details[index] for index in range(position_count)] + + +def extract_selected_token_logprobs(generation_details: Any) -> list[float]: + """Return the sampled log probability for each published output token.""" + generation_token_ids = list(getattr(generation_details, "token_ids", ()) or ()) + details = getattr(generation_details, "logprobs", None) + if details is None: + if generation_token_ids: + raise RuntimeError( + "vLLM generation output with token IDs did not include logprobs" + ) + return [] + details = _snapshot_generation_logprob_details(details) + if len(generation_token_ids) != len(details): + raise RuntimeError( + "vLLM returned mismatched generation token IDs and log probabilities: " + f"token_count={len(generation_token_ids)}, logprob_count={len(details)}" + ) + return _extract_selected_token_logprobs(generation_token_ids, details) + + GROUPED_MOE_MXFP8_REFIT_ERROR = ( "MXFP8 refit does not support grouped MoE expert weights." ) diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 250542d1b77..b2f38e941c4 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -640,6 +640,9 @@ def setup_token_capture( staging_partition: str, *, capture_media: bool = False, + generation_prefix_cuts_enabled: bool = False, + generation_cut_control_token: str | None = None, + generation_cut_control_timeout_s: float | None = None, ) -> None: """Install ledger-authoritative token capture in every DP-leader worker. @@ -656,6 +659,9 @@ def setup_token_capture( dp_cfg=dp_cfg, staging_partition=staging_partition, capture_media=capture_media, + generation_prefix_cuts_enabled=generation_prefix_cuts_enabled, + generation_cut_control_token=generation_cut_control_token, + generation_cut_control_timeout_s=generation_cut_control_timeout_s, run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], ) ray.get(futures) @@ -1715,6 +1721,46 @@ def resume_generation_after_refit(self) -> bool: raise RuntimeError("Failed to resume every async vLLM engine") return True + def begin_token_capture_snapshot_fence( + self, *, timeout_s: Optional[float] = None + ) -> bool: + """Fence terminal staging while vLLM keeps decoding into live buffers.""" + if not self.cfg["vllm_cfg"]["async_engine"]: + raise RuntimeError( + "begin_token_capture_snapshot_fence requires async_engine=True" + ) + if not self.worker_group or not self.worker_group.workers: + raise RuntimeError("Worker group is not initialized") + futures = self.worker_group.run_all_workers_single_data( + "begin_token_capture_snapshot_fence_async", + run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], + ) + if not all(ray.get(futures, timeout=timeout_s)): + raise RuntimeError( + "Failed to fence terminal staging on every async vLLM engine" + ) + return True + + def end_token_capture_snapshot_fence( + self, *, timeout_s: Optional[float] = None + ) -> bool: + """Release terminal capture writes after the coordinated TQ snapshot.""" + if not self.cfg["vllm_cfg"]["async_engine"]: + raise RuntimeError( + "end_token_capture_snapshot_fence requires async_engine=True" + ) + if not self.worker_group or not self.worker_group.workers: + raise RuntimeError("Worker group is not initialized") + futures = self.worker_group.run_all_workers_single_data( + "end_token_capture_snapshot_fence_async", + run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], + ) + if not all(ray.get(futures, timeout=timeout_s)): + raise RuntimeError( + "Failed to release every async vLLM token-capture snapshot fence" + ) + return True + @property def requires_kv_scale_sync(self) -> bool: """Check if KV cache scales should be synchronized during refit. diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index c47d1a3802c..d9a5a9555e4 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -14,15 +14,17 @@ import asyncio import copy +import functools import gc import logging import threading import time import uuid import warnings -from collections.abc import Awaitable, Callable -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, cast +from collections.abc import Awaitable, Callable, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Any, AsyncGenerator, Literal, Optional, cast import ray import torch @@ -45,6 +47,13 @@ _get_node_ip_local, ) from nemo_rl.distributed.worker_group_utils import get_nsight_config_if_pattern_matches +from nemo_rl.models.generation.generation_cut_capture import ( + GenerationCutCaptureMixin, + _CompletedCaptureState, + _remaining_generation_limits_after_prefix, + _RequestCaptureState, + _TokenCaptureSnapshotGate, +) from nemo_rl.models.generation.interfaces import ( GenerationDatumSpec, GenerationOutputSpec, @@ -60,6 +69,7 @@ from nemo_rl.models.generation.vllm.utils import ( attach_routed_experts_to_chat_response_choices, attach_token_information_to_chat_response_choices, + extract_selected_token_logprobs, format_prompt_for_vllm_generation, validate_rollout_prompt, model_dump_chat_response_with_dynamic_message_fields, @@ -73,18 +83,233 @@ from nemo_rl.telemetry.setup import shutdown_telemetry LOGGER = logging.getLogger(__name__) +_RESTORED_PREFIX_TERMINAL_PROMPT_KEY = "__nemo_rl_restored_prefix_terminal__" + -if TYPE_CHECKING: - from nemo_gym.token_id_capture.staging.capture import ActiveCall +@dataclass(frozen=True) +class _RestoredPrefixTerminal: + """A restored call whose durable output already reached a terminal limit.""" + + prompt_token_ids: tuple[int, ...] + prefix_token_count: int + reason: str + finish_reason: Literal["stop", "length"] + stop_reason: str | int | None = None + + @property + def original_prompt_token_count(self) -> int: + return len(self.prompt_token_ids) - self.prefix_token_count + + +def _classify_restored_prefix_terminal( + *, + prompt_token_ids: list[int], + prefix_token_count: int, + requested_output_tokens: int | None, + model_max_tokens: int, + terminal_finish_reason: Literal["stop", "length"] | None = None, + terminal_stop_reason: str | int | None = None, +) -> _RestoredPrefixTerminal | None: + """Classify an exact terminal prefix or reject an incompatible restore.""" + prompt_token_count = len(prompt_token_ids) + if prefix_token_count < 0: + raise ValueError("prefix_token_count must be non-negative") + if prefix_token_count > prompt_token_count: + raise ValueError( + "Durable generation prefix contains more generated tokens than the " + "restored engine prompt." + ) + if ( + requested_output_tokens is not None + and prefix_token_count > requested_output_tokens + ): + raise ValueError( + "Durable generation prefix token count " + f"({prefix_token_count}) exceeds the restored request output " + f"budget ({requested_output_tokens})." + ) + if prompt_token_count > model_max_tokens: + raise ValueError( + "Durable generation prefix prompt length " + f"({prompt_token_count}) exceeds restored model capacity " + f"({model_max_tokens})." + ) + if terminal_stop_reason is not None and terminal_finish_reason is None: + raise ValueError("terminal_stop_reason requires terminal_finish_reason") + if terminal_finish_reason is not None: + return _RestoredPrefixTerminal( + prompt_token_ids=tuple(prompt_token_ids), + prefix_token_count=prefix_token_count, + reason=f"observed_{terminal_finish_reason}", + finish_reason=terminal_finish_reason, + stop_reason=terminal_stop_reason, + ) + reached_output_limit = ( + requested_output_tokens is not None + and prefix_token_count == requested_output_tokens + ) + reached_model_capacity = prompt_token_count == model_max_tokens + if not reached_output_limit and not reached_model_capacity: + return None + return _RestoredPrefixTerminal( + prompt_token_ids=tuple(prompt_token_ids), + prefix_token_count=prefix_token_count, + reason=( + "output_and_model_limit" + if reached_output_limit and reached_model_capacity + else "output_limit" + if reached_output_limit + else "model_limit" + ), + finish_reason="length", + ) + + +def _build_restored_prefix_terminal_output( + request_id: str, terminal: _RestoredPrefixTerminal +) -> Any: + """Build the final vLLM output for a prefix that needs no more decoding.""" + from vllm.outputs import CompletionOutput, RequestOutput + + return RequestOutput( + request_id=request_id, + prompt=None, + prompt_token_ids=list(terminal.prompt_token_ids), + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=0.0, + logprobs=[], + finish_reason=terminal.finish_reason, + stop_reason=terminal.stop_reason, + ) + ], + finished=True, + ) + + +@dataclass(frozen=True) +class _RestoredPrefixParser: + """Give vLLM's parser the complete output while capture keeps tail deltas.""" + + delegate: Any + prefix_token_ids: tuple[int, ...] + + def parse( + self, + model_output: str, + request: Any, + *, + enable_auto_tools: bool, + model_output_token_ids: list[int], + ) -> Any: + return self.delegate.parse( + model_output, + request, + enable_auto_tools=enable_auto_tools, + model_output_token_ids=[ + *self.prefix_token_ids, + *model_output_token_ids, + ], + ) + + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + return self.delegate.count_reasoning_tokens( + [*self.prefix_token_ids, *token_ids] + ) + + +@dataclass +class _CompletionOutputDeltaAccumulator: + """Linear-time accumulator for one vLLM completion index.""" + + template: Any | None = None + text_parts: list[str] = field(default_factory=list) + token_ids: list[int] = field(default_factory=list) + logprobs: list[Any] | None = None + routed_expert_chunks: list[Any] = field(default_factory=list) + + def append(self, output: Any) -> None: + if self.template is None: + self.template = copy.copy(output) + else: + for name in ( + "cumulative_logprob", + "finish_reason", + "stop_reason", + "lora_request", + ): + if hasattr(output, name): + setattr(self.template, name, getattr(output, name)) + self.text_parts.append(str(getattr(output, "text", ""))) + self.token_ids.extend(list(getattr(output, "token_ids", ()) or ())) + delta_logprobs = getattr(output, "logprobs", None) + if delta_logprobs is not None: + if self.logprobs is None: + self.logprobs = [] + self.logprobs.extend(list(delta_logprobs)) + routed = getattr(output, "routed_experts", None) + if routed is not None: + self.routed_expert_chunks.append(copy.deepcopy(routed)) + + def build(self) -> Any: + if self.template is None: + raise RuntimeError("cannot build an empty completion output") + self.template.text = "".join(self.text_parts) + self.template.token_ids = self.token_ids + self.template.logprobs = self.logprobs + if self.routed_expert_chunks: + self.template.routed_experts = torch.cat( + [torch.as_tensor(chunk) for chunk in self.routed_expert_chunks], + dim=0, + ) + return self.template @dataclass -class CapturedRequest: - """Request-local ownership of tokens and processed-media snapshots.""" +class _RequestOutputDeltaAccumulator: + """Reconstruct one final RequestOutput from immutable engine deltas.""" - call: "ActiveCall" - prompt_token_ids: list[int] - media: CapturedMedia | None + template: Any | None = None + completions: dict[int, _CompletionOutputDeltaAccumulator] = field( + default_factory=dict + ) + + def append(self, delta: Any) -> None: + if self.template is None: + self.template = copy.copy(delta) + else: + previous = self.template + self.template = copy.copy(delta) + for name in ( + "prompt", + "prompt_token_ids", + "prompt_logprobs", + "encoder_prompt", + "encoder_prompt_token_ids", + "lora_request", + "num_cached_tokens", + "prompt_routed_experts", + ): + if getattr(self.template, name, None) is None and hasattr( + previous, name + ): + setattr(self.template, name, getattr(previous, name)) + for output in getattr(delta, "outputs", ()): + self.completions.setdefault( + output.index, _CompletionOutputDeltaAccumulator() + ).append(output) + + def build(self) -> Any: + if self.template is None: + raise RuntimeError("cannot build an empty request output") + self.template.outputs = [ + self.completions[index].build() for index in sorted(self.completions) + ] + return self.template from nemo_rl.distributed.refit_watchdog import RefitAborted, is_refit_abort @@ -132,6 +357,15 @@ async def _generate( request_id: str, kwargs: dict[str, Any], ) -> AsyncGenerator[Any, None]: + terminal = ( + prompt.get(_RESTORED_PREFIX_TERMINAL_PROMPT_KEY) + if isinstance(prompt, dict) + else None + ) + if isinstance(terminal, _RestoredPrefixTerminal): + yield _build_restored_prefix_terminal_output(request_id, terminal) + return + iterator = None completed = False @@ -186,7 +420,9 @@ async def is_tracing_enabled(self) -> bool: class VllmAsyncGenerationWorkerImpl( - VllmAsyncCheckpointEngineRpcMixin, BaseVllmGenerationWorker + GenerationCutCaptureMixin, + VllmAsyncCheckpointEngineRpcMixin, + BaseVllmGenerationWorker, ): def __init__( self, @@ -234,10 +470,32 @@ def __init__( # the set_rollout_weight_version fan-out from the SC's _sync_weights. self.token_capture = None self._rollout_weight_version = 0 - # In-flight captured calls keyed by id(request): the ActiveCall, the - # exact engine prompt ids recorded at preprocess time, and the media - # snapshot the engine ran on (see CapturedRequest). - self._capture_calls: dict[int, CapturedRequest] = {} + # In-flight calls are indexed by both the request object and Gym's + # stable model-call ID. The second index is what a checkpoint inventory + # uses while the original HTTP request remains active. + self._capture_calls: dict[int, _RequestCaptureState] = {} + self._capture_calls_by_model_call_id: dict[str, _RequestCaptureState] = {} + self._completed_capture_calls: dict[str, _CompletedCaptureState] = {} + self._generation_cut_receipts: dict[tuple[str, str], Any] = {} + self._capture_registry_lock = threading.Lock() + self._capture_sink: Any | None = None + self._generation_prefix_cuts_enabled = False + self._generation_cut_control_token: str | None = None + self._generation_cut_control_timeout_s: float | None = None + self._generation_cut_tokenizer: Any | None = None + self._token_capture_snapshot_gate = _TokenCaptureSnapshotGate() + # A dedicated executor prevents completions waiting on the closed gate + # from consuming all threads needed to create the checkpoint cut. + self._generation_cut_control_executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="nrl-generation-cut-control", + ) + # Fence closes get their own thread so a stale cut still staging rows + # cannot delay the next checkpoint's fence past the driver's timeout. + self._token_capture_fence_executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="nrl-token-capture-fence", + ) self._capture_media = False self._capture_patch_size: int | None = None self._capture_image_token_id: int | None = None @@ -529,6 +787,9 @@ async def setup_token_capture( staging_partition: str, *, capture_media: bool = False, + generation_prefix_cuts_enabled: bool = False, + generation_cut_control_token: str | None = None, + generation_cut_control_timeout_s: float | None = None, ) -> bool: """Host ledger-authoritative token capture in this worker. @@ -558,6 +819,14 @@ async def setup_token_capture( capture_media=capture_media, media_pixel_dtype=pixel_dtype, ) + if generation_prefix_cuts_enabled and not generation_cut_control_token: + raise ValueError( + "generation-prefix cuts require a non-empty control bearer token" + ) + if generation_prefix_cuts_enabled and capture_media: + raise ValueError( + "generation-prefix recovery does not yet support multimodal capture" + ) if capture_media: # Omni-only: a new processor family must also change setup.py (driver # checks), captured_media.py (_processed_omni_tensors, pack_images, @@ -589,6 +858,13 @@ async def setup_token_capture( self._capture_image_token_id = int(context_ids[0]) self._capture_patch_size = int(info.get_hf_config().patch_size) self._capture_media = capture_media + self._capture_sink = sink + self._generation_prefix_cuts_enabled = generation_prefix_cuts_enabled + self._generation_cut_control_token = generation_cut_control_token + self._generation_cut_control_timeout_s = generation_cut_control_timeout_s + if generation_prefix_cuts_enabled: + # Cuts must land on character boundaries; see _ends_inside_character. + self._generation_cut_tokenizer = self.llm.renderer.get_tokenizer() source = TQTokenSource( dp_client, staging_partition=staging_partition, capture_media=capture_media ) @@ -633,6 +909,8 @@ def _begin_request_capture( admission: Any | None = None, prefix_token_ids: list[int] | None = None, media: CapturedMedia | None = None, + generation_cut: Any | None = None, + resumed_generation_token_ids: list[int] | None = None, ) -> None: """Admit one ledger-forwarded call into the capture layer. @@ -656,11 +934,116 @@ def _begin_request_capture( call = capture.begin_call( admission, prefix_token_ids=prefix_token_ids, + generation_cut=generation_cut, + generation_cut_staging_keys=( + admission.generation_cut.staging_keys + if admission.generation_cut is not None + else None + ), stream=bool(getattr(request, "stream", False)), ) - self._capture_calls[id(request)] = CapturedRequest( - call, list(prompt_token_ids), media + state = _RequestCaptureState( + call=call, + prompt_token_ids=list(prompt_token_ids), + media=media, + resumed_generation_token_ids=list(resumed_generation_token_ids or ()), + effective_output_limit=( + admission.generation_cut.effective_output_limit + if admission.generation_cut is not None + else None + ), + generation_cut_staging_keys=( + list(admission.generation_cut.staging_keys) + if admission.generation_cut is not None + else [] + ), ) + with self._capture_registry_lock: + if call.model_call_id in self._capture_calls_by_model_call_id: + raise RuntimeError( + f"model call {call.model_call_id!r} is already active in token capture" + ) + self._capture_calls[id(request)] = state + self._capture_calls_by_model_call_id[call.model_call_id] = state + + def _observe_request_capture(self, request: Any, request_output: Any) -> None: + """Append one request's immutable vLLM output delta.""" + state = self._get_request_capture(request) + if state is None: + return + outputs = getattr(request_output, "outputs", None) + if not outputs: + return + try: + output = outputs[0] + generation_token_ids = list(getattr(output, "token_ids", ()) or ()) + generation_logprobs = extract_selected_token_logprobs(output) + finish_reason = getattr(output, "finish_reason", None) + stop_reason = getattr(output, "stop_reason", None) + except (RuntimeError, TypeError, ValueError) as error: + with state.lock: + state.observation_error = f"{type(error).__name__}: {error}" + return + with state.lock: + state.observe(generation_token_ids, generation_logprobs) + if finish_reason is not None: + if finish_reason not in ("stop", "length"): + state.observation_error = ( + f"unsupported terminal finish_reason {finish_reason!r}" + ) + elif ( + state.terminal_finish_reason is not None + and state.terminal_finish_reason != finish_reason + ): + state.observation_error = ( + "conflicting terminal finish reasons: " + f"{state.terminal_finish_reason!r} and {finish_reason!r}" + ) + elif not isinstance(stop_reason, (str, int, type(None))): + state.observation_error = ( + "terminal stop_reason must be a string, integer, or None" + ) + else: + state.terminal_finish_reason = finish_reason + state.terminal_stop_reason = stop_reason + + def _restore_response_prefix( + self, request: Any, request_output: Any, *, tokenizer: Any + ) -> None: + """Prepend the durable assistant prefix before vLLM parses the response.""" + state = self._get_request_capture(request) + if state is None or not state.resumed_generation_token_ids: + return + outputs = getattr(request_output, "outputs", None) + if not outputs: + return + output = outputs[0] + # vLLM already detokenized the tail with the request's settings and + # stripped EOS and any matched stop string from it. Re-decoding the + # whole sequence would restore both, so decode only the prefix, with + # the same settings, and prepend it to the processed tail. + with state.lock: + prefix_text = tokenizer.decode( + state.resumed_generation_token_ids, + skip_special_tokens=getattr(request, "skip_special_tokens", True), + spaces_between_special_tokens=getattr( + request, "spaces_between_special_tokens", True + ), + ) + output.text = prefix_text + (getattr(output, "text", None) or "") + + def _record_request_effective_output_limit( + self, request: Any, effective_output_limit: int + ) -> None: + """Remember vLLM's resolved total output budget for future cuts.""" + if effective_output_limit <= 0: + raise ValueError("effective output limit must be positive") + state = self._get_request_capture(request) + if state is None: + return + with state.lock: + if state.effective_output_limit is None: + state.effective_output_limit = effective_output_limit def _capture_request_media( self, @@ -828,17 +1211,47 @@ def _delta_align_routed_experts( payload["choices"] = [choice] def _finish_request_capture(self, request: Any, content: dict) -> dict: - """Stage the finished call and ride its coords on the response. + """Run terminal token staging outside an active checkpoint cut.""" + self._token_capture_snapshot_gate.enter() + try: + return self._finish_request_capture_after_snapshot_fence(request, content) + finally: + self._token_capture_snapshot_gate.exit() - Tokens and the call's new media tensors go to TQ in one write (the - media ride as opaque attachments beside the record), so ``staged`` - coords vouch for both and any failure is ``capture_failed`` at call - time. Token ids, logprobs, and routes are stripped after staging, so - the worker->gate hop carries the completion and coords only. - """ - state = self._capture_calls.pop(id(request), None) + def _finish_request_capture_after_snapshot_fence( + self, request: Any, content: dict + ) -> dict: + """Stage one canonical terminal row and retire its prefix chunks.""" + state = self._get_request_capture(request) if state is None: return content + with state.lifecycle_lock: + if state.terminal_started: + raise RuntimeError( + f"model call {state.call.model_call_id!r} already started " + "terminal token capture" + ) + state.terminal_started = True + content, obsolete_staging_keys = ( + self._finish_request_capture_with_lifecycle_owned( + state, request, content + ) + ) + sink = self._capture_sink + if obsolete_staging_keys and sink is not None: + try: + sink.clear(list(obsolete_staging_keys)) + except Exception: # noqa: BLE001 - completion is already durable + LOGGER.exception( + "failed to clear obsolete generation chunks for model call %s", + state.call.model_call_id, + ) + return content + + def _finish_request_capture_with_lifecycle_owned( + self, state: _RequestCaptureState, request: Any, content: dict + ) -> tuple[dict, tuple[str, ...]]: + """Stage a terminal call while owning ``state.lifecycle_lock``.""" call, prompt_token_ids = state.call, state.prompt_token_ids payload = dict(content) # vLLM's OpenAI response carries no prompt ids; the adapter reads the @@ -853,6 +1266,7 @@ def _finish_request_capture(self, request: Any, content: dict) -> dict: # digest-covered extras; the pixels themselves are attachments. payload[MEDIA_SPANS_FIELD] = [item.to_dict() for item in state.media.items] adapter = self.token_capture.adapter + generated_token_ids: list[int] = [] if adapter is not None: try: generated_token_ids, _ = adapter.extract_generation(payload) @@ -869,6 +1283,47 @@ def _finish_request_capture(self, request: Any, content: dict) -> dict: payload, attachments=state.media.tensors if state.media is not None else None, ) + total_generation_token_count = len(generated_token_ids) + if call.generation_cut is not None: + prefix_tokens = sum( + mask == 1.0 for mask in call.generation_cut.token_mask_delta + ) + total_generation_token_count += prefix_tokens + LOGGER.info( + "generation prefix completed: rollout_id=%s model_call_id=%s " + "source_model_call_id=%s prefix_tokens=%d tail_tokens=%d " + "total_generation_tokens=%d", + call.rollout_id, + call.model_call_id, + call.generation_cut.model_call_id, + prefix_tokens, + len(generated_token_ids), + total_generation_token_count, + ) + with state.lock: + effective_output_limit = state.effective_output_limit + terminal_finish_reason = state.terminal_finish_reason + terminal_stop_reason = state.terminal_stop_reason + if terminal_finish_reason is None: + choices = content.get("choices") or [] + finish_reason = choices[0].get("finish_reason") if choices else None + if finish_reason in ("stop", "length"): + terminal_finish_reason = finish_reason + terminal_stop_reason = choices[0].get("stop_reason") + self._remember_completed_capture( + call.model_call_id, + coords, + total_generation_token_count, + effective_output_limit=effective_output_limit, + terminal_finish_reason=terminal_finish_reason, + terminal_stop_reason=terminal_stop_reason, + ) + self._pop_request_capture(request) + obsolete_staging_keys = ( + tuple(state.generation_cut_staging_keys) + if coords.disposition == "staged" + else () + ) for choice in content.get("choices") or []: choice.pop("logprobs", None) # Token arrays and delta-aligned routes were staged to TQ above; @@ -883,13 +1338,7 @@ def _finish_request_capture(self, request: Any, content: dict) -> dict: ): message.pop(field, None) content["ng_commit_coords"] = coords.model_dump() - return content - - def _abort_request_capture(self, request: Any, *, reason: str) -> None: - """Drop the in-flight capture state for a request that errored.""" - state = self._capture_calls.pop(id(request), None) - if state is not None and self.token_capture is not None: - self.token_capture.fail_call(state.call, reason=reason) + return content, obsolete_staging_keys # ruff: noqa def _setup_vllm_openai_api_server(self, app: FastAPI) -> FastAPI: @@ -899,8 +1348,9 @@ def _setup_vllm_openai_api_server(self, app: FastAPI) -> FastAPI: from logging import LogRecord from typing import List, Optional, Union - from fastapi import Request + from fastapi import Header, HTTPException, Request from fastapi.responses import JSONResponse, StreamingResponse + from pydantic import PrivateAttr from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, @@ -923,6 +1373,7 @@ def _setup_vllm_openai_api_server(self, app: FastAPI) -> FastAPI: ServingTokenization, ) from vllm.renderers.online_renderer import OnlineRenderer + from vllm.sampling_params import RequestOutputKind from vllm.exceptions import VLLMValidationError from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager from vllm.tool_parsers.abstract_tool_parser import ToolParserManager @@ -1080,11 +1531,56 @@ async def preprocess_chat( # the single splice path for staged and inline prefixes. admission = worker_self._capture_admission(request) capture_prefix_token_ids: list[int] | None = None - if admission is not None and admission.mode == "token_in": + generation_cut = None + resumed_generation_token_ids: list[int] = [] + restored_request_output_tokens: int | None = None + if admission is not None: capture_prefix_token_ids = await asyncio.to_thread( worker_self._resolve_admission_prefix, admission ) - worker_self._enter_request_prefix(request, capture_prefix_token_ids) + generation_cut = await asyncio.to_thread( + worker_self._resolve_generation_cut, + admission, + capture_prefix_token_ids, + ) + engine_prefix_token_ids = list(capture_prefix_token_ids) + if generation_cut is not None: + restored_request_output_tokens = ( + admission.generation_cut.effective_output_limit + ) + engine_prefix_token_ids.extend(generation_cut.token_ids_delta) + resumed_generation_token_ids = [ + token_id + for token_id, mask in zip( + generation_cut.token_ids_delta, + generation_cut.token_mask_delta, + ) + if mask == 1.0 + ] + request._restored_generation_token_ids = tuple( + resumed_generation_token_ids + ) + ( + remaining_output_tokens, + remaining_min_tokens, + ) = _remaining_generation_limits_after_prefix( + max_tokens=restored_request_output_tokens, + min_tokens=getattr(request, "min_tokens", None), + prefix_token_count=( + admission.generation_cut.prefix_token_count + ), + ) + if ( + remaining_output_tokens is not None + and remaining_output_tokens > 0 + ): + actual_request_max_tokens = remaining_output_tokens + if remaining_min_tokens is not None: + request.min_tokens = remaining_min_tokens + if engine_prefix_token_ids: + worker_self._enter_request_prefix( + request, engine_prefix_token_ids + ) if ( not hasattr(request, "required_prefix_token_ids") @@ -1108,7 +1604,10 @@ async def preprocess_chat( request, res[1][0]["prompt_token_ids"], admission=admission, + prefix_token_ids=capture_prefix_token_ids, media=media, + generation_cut=generation_cut, + resumed_generation_token_ids=resumed_generation_token_ids, ) return res @@ -1152,23 +1651,57 @@ async def preprocess_chat( engine_prompt = res[1][0] - splice = splice_prefix_tokens( - tokenizer=self.renderer.tokenizer, - model_prefix_token_ids=model_prefix_token_ids, - template_prefix_token_ids=actual_corresponding_token_ids, - template_token_ids=engine_prompt["prompt_token_ids"], - ) - final_prompt_token_ids = splice.token_ids - media = await asyncio.to_thread( - worker_self._capture_request_media, - engine_prompt, - admission=admission, - splice=splice, - ) + splice = None + if generation_cut is not None: + # The durable prefix is the exact engine prompt at the cut; + # rendering/splicing it again could add or drop tokens. + final_prompt_token_ids = model_prefix_token_ids + media = None + else: + splice = splice_prefix_tokens( + tokenizer=self.renderer.tokenizer, + model_prefix_token_ids=model_prefix_token_ids, + template_prefix_token_ids=actual_corresponding_token_ids, + template_token_ids=engine_prompt["prompt_token_ids"], + ) + final_prompt_token_ids = splice.token_ids + media = await asyncio.to_thread( + worker_self._capture_request_media, + engine_prompt, + admission=admission, + splice=splice, + ) engine_prompt["prompt_token_ids"] = final_prompt_token_ids + restored_prefix_terminal = None + if generation_cut is not None: + try: + restored_prefix_terminal = _classify_restored_prefix_terminal( + prompt_token_ids=final_prompt_token_ids, + prefix_token_count=( + admission.generation_cut.prefix_token_count + ), + requested_output_tokens=restored_request_output_tokens, + model_max_tokens=self.model_config.max_model_len, + terminal_finish_reason=( + admission.generation_cut.terminal_finish_reason + ), + terminal_stop_reason=( + admission.generation_cut.terminal_stop_reason + ), + ) + except ValueError as error: + raise VLLMValidationError( + str(error), + parameter="generation_prefix", + value=admission.generation_cut.prefix_token_count, + ) from error + # Clamp after prefix replacement since the prompt length may have changed. - if actual_request_max_tokens is not None: + if ( + restored_prefix_terminal is None + and actual_request_max_tokens is not None + ): self._clamp_max_tokens( request, actual_request_max_tokens, @@ -1184,8 +1717,28 @@ async def preprocess_chat( admission=admission, prefix_token_ids=capture_prefix_token_ids, media=media, + generation_cut=generation_cut, + resumed_generation_token_ids=resumed_generation_token_ids, ) + if restored_prefix_terminal is not None: + request.min_tokens = 0 + self._set_max_tokens(request, 1) + engine_prompt[_RESTORED_PREFIX_TERMINAL_PROMPT_KEY] = ( + restored_prefix_terminal + ) + if len(final_prompt_token_ids) == self.model_config.max_model_len: + engine_prompt["prompt_token_ids"] = final_prompt_token_ids[:-1] + request._restored_prefix_terminal = restored_prefix_terminal + LOGGER.info( + "generation prefix already terminal: " + "rollout_id=%s model_call_id=%s prefix_tokens=%d reason=%s", + admission.rollout_id, + admission.model_call_id, + restored_prefix_terminal.prefix_token_count, + restored_prefix_terminal.reason, + ) + return res ######################################## @@ -1200,12 +1753,63 @@ class NeMoRLChatCompletionRequest( # Ledger-authoritative token capture: the call identity the ledger # attaches (rollout_id, call_id, parent_call_id, prev_len, mode). ng_capture: Optional[dict[str, Any]] = None + _restored_prefix_terminal: _RestoredPrefixTerminal | None = PrivateAttr( + default=None + ) + _restored_generation_token_ids: tuple[int, ...] = PrivateAttr(default=()) + + def to_sampling_params(self, *args, **kwargs): + sampling_params = super().to_sampling_params(*args, **kwargs) + if ( + worker_self._generation_prefix_cuts_enabled + and self.ng_capture is not None + ): + sampling_params.output_kind = RequestOutputKind.DELTA + worker_self._record_request_effective_output_limit( + self, int(sampling_params.max_tokens) + ) + return sampling_params # vLLM 0.25 routes both /v1/chat/completions and /tokenize through # OnlineRenderer.preprocess_chat, so the prefix-token override # belongs on the renderer subclass. worker_self = self + @app.post("/ng-control/v1/generation-cut") + async def checkpoint_generation_cut( + inventory: dict[str, Any], + authorization: str | None = Header(default=None), + ): + """Persist one checkpoint's frozen active-call prefixes to TQ.""" + import secrets + + from nemo_gym._checkpoint.generation_cut import GenerationCutInventory + + expected = worker_self._generation_cut_control_token + supplied = ( + authorization.removeprefix("Bearer ") + if authorization is not None and authorization.startswith("Bearer ") + else "" + ) + if not worker_self._generation_prefix_cuts_enabled or expected is None: + raise HTTPException( + status_code=404, detail="generation-prefix cuts are disabled" + ) + if not secrets.compare_digest(supplied, expected): + raise HTTPException(status_code=401, detail="invalid control bearer") + typed_inventory = GenerationCutInventory.model_validate(inventory) + # Measured from arrival so time queued behind a stale cut counts. + timeout_s = worker_self._generation_cut_control_timeout_s + deadline = None if timeout_s is None else time.monotonic() + timeout_s + receipt = await worker_self._run_generation_cut_control( + functools.partial( + worker_self._checkpoint_generation_cut, + typed_inventory, + deadline=deadline, + ) + ) + return receipt.model_dump(mode="json") + class NeMoRLOpenAIServingChatMixin: async def chat_completion_full_generator( self, @@ -1230,13 +1834,53 @@ async def chat_completion_full_generator( parameter="top_logprobs", ) + aggregate_deltas = bool( + worker_self._generation_prefix_cuts_enabled + and request.ng_capture is not None + ) final_res = None + delta_accumulator = _RequestOutputDeltaAccumulator() async def capture_result_generator(): nonlocal final_res async for res in result_generator: - final_res = res - yield res + worker_self._observe_request_capture(request, res) + if not aggregate_deltas: + final_res = res + worker_self._restore_response_prefix( + request, + res, + tokenizer=self.renderer.tokenizer, + ) + yield res + continue + delta_accumulator.append(res) + + if not aggregate_deltas or delta_accumulator.template is None: + return + final_res = delta_accumulator.build() + worker_self._restore_response_prefix( + request, + final_res, + tokenizer=self.renderer.tokenizer, + ) + yield final_res + + restored_parser_prefix = request._restored_generation_token_ids + if restored_parser_prefix: + if len(args) >= 6 and args[5] is not None: + mutable_args = list(args) + mutable_args[5] = _RestoredPrefixParser( + delegate=args[5], + prefix_token_ids=restored_parser_prefix, + ) + args = tuple(mutable_args) + elif kwargs.get("parser") is not None: + kwargs = dict(kwargs) + kwargs["parser"] = _RestoredPrefixParser( + delegate=kwargs["parser"], + prefix_token_ids=restored_parser_prefix, + ) response = await super().chat_completion_full_generator( request, @@ -1250,6 +1894,32 @@ async def capture_result_generator(): ): return response + restored_prefix_terminal = request._restored_prefix_terminal + if ( + isinstance(restored_prefix_terminal, _RestoredPrefixTerminal) + and response.usage is not None + ): + response.usage.prompt_tokens = ( + restored_prefix_terminal.original_prompt_token_count + ) + response.usage.completion_tokens = ( + restored_prefix_terminal.prefix_token_count + ) + response.usage.total_tokens = ( + response.usage.prompt_tokens + response.usage.completion_tokens + ) + if response.prompt_token_ids is not None: + response.prompt_token_ids = list( + restored_prefix_terminal.prompt_token_ids[ + : restored_prefix_terminal.original_prompt_token_count + ] + ) + request_metadata = ( + args[4] if len(args) >= 5 else kwargs.get("request_metadata") + ) + if request_metadata is not None: + request_metadata.final_usage_info = response.usage + if request.logprobs and return_as_token_id: response = attach_token_information_to_chat_response_choices( response, @@ -2270,6 +2940,21 @@ async def pause_generation_async(self, *, clear_cache: bool) -> bool: await self.llm.pause_generation(mode="keep", clear_cache=clear_cache) return True + async def begin_token_capture_snapshot_fence_async(self) -> bool: + """Fence terminal staging without pausing vLLM decoding.""" + if not self.cfg["vllm_cfg"]["async_engine"]: + raise RuntimeError( + "begin_token_capture_snapshot_fence_async requires async_engine=True" + ) + gate = self._token_capture_snapshot_gate + # Take the epoch on the loop, in call order with the matching end: if + # the driver times out and releases first, the late close is a no-op. + epoch = gate.begin_epoch() + await asyncio.get_running_loop().run_in_executor( + self._token_capture_fence_executor, gate.close_and_wait, epoch + ) + return True + async def resume_generation_async(self) -> bool: """Resume vLLM generation after an in-flight weight update.""" assert self.llm is not None, ( @@ -2284,6 +2969,15 @@ async def resume_generation_async(self) -> bool: await self.llm.resume_generation() return True + async def end_token_capture_snapshot_fence_async(self) -> bool: + """Release terminal token staging after the TQ snapshot is durable.""" + if not self.cfg["vllm_cfg"]["async_engine"]: + raise RuntimeError( + "end_token_capture_snapshot_fence_async requires async_engine=True" + ) + self._token_capture_snapshot_gate.reopen() + return True + async def sleep_async(self): """Async version of sleep.""" assert self.llm is not None, ( @@ -2330,6 +3024,22 @@ async def wake_up_async(self, **kwargs): async def shutdown(self) -> bool: """Clean up vLLM resources.""" try: + # A terminal write may be waiting at the TQ snapshot fence when the + # actor is asked to stop. Release it before retiring the isolated + # cut-control and fence executors; waiting for a wedged TQ write + # here would otherwise make actor shutdown hang indefinitely. + token_capture_snapshot_gate = getattr( + self, "_token_capture_snapshot_gate", None + ) + if token_capture_snapshot_gate is not None: + token_capture_snapshot_gate.reopen() + for executor_name in ( + "_generation_cut_control_executor", + "_token_capture_fence_executor", + ): + executor = getattr(self, executor_name, None) + if executor is not None: + executor.shutdown(wait=False, cancel_futures=True) if self.server_thread is not None: self.http_server.should_exit = True await asyncio.to_thread(self.server_thread.join) diff --git a/pyrefly.toml b/pyrefly.toml index 60783954201..6a5265abc93 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -174,6 +174,8 @@ project-includes = [ "nemo_rl/environments/__init__.py", "nemo_rl/environments/code_jaccard_environment.py", "nemo_rl/environments/games/sliding_puzzle.py", + "nemo_rl/environments/gym_checkpoint_adapter.py", + "nemo_rl/environments/gym_checkpoint_coordinator.py", "nemo_rl/environments/interfaces.py", "nemo_rl/environments/math_environment.py", "nemo_rl/environments/metrics.py", @@ -227,6 +229,7 @@ project-includes = [ "nemo_rl/models/generation/dynamo/worker_pool.py", "nemo_rl/models/generation/engine_supervisor.py", "nemo_rl/models/generation/fleet_health.py", + "nemo_rl/models/generation/generation_cut_capture.py", "nemo_rl/models/generation/generation_router.py", "nemo_rl/models/generation/interfaces.py", "nemo_rl/models/generation/megatron/__init__.py", diff --git a/tests/functional/L1_Functional_Tests_SingleController_3.sh b/tests/functional/L1_Functional_Tests_SingleController_3.sh index b056970d426..4f4416d6e68 100755 --- a/tests/functional/L1_Functional_Tests_SingleController_3.sh +++ b/tests/functional/L1_Functional_Tests_SingleController_3.sh @@ -55,6 +55,13 @@ run_test fast env SC_SIBLING_RECOVERY_GENERATION_BACKEND=megatron uv run --no-sy # Periodic native-TQ snapshot while a streamed step owns only part of its # rollout batch, followed by SIGKILL and rollback to the durable trainer anchor. run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh +# Two-process Gym turn-level recovery against a stateful Workplace Assistant +# server: a restored episode resumes on its Gym instance as its next attempt. +run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_turn_recovery.sh +# Gym-v2 active-decode recovery: checkpoint a nonterminal vLLM prefix, restart +# the process, restore that exact prefix, and generate only its remaining tail. +# This is intentionally full-tier because it performs two long-generation runs. +run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_prefix_recovery.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/_gym_prefix_recovery_snapshot.py b/tests/functional/_gym_prefix_recovery_snapshot.py new file mode 100644 index 00000000000..4a2fffc675f --- /dev/null +++ b/tests/functional/_gym_prefix_recovery_snapshot.py @@ -0,0 +1,363 @@ +# 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. + +"""Select and verify the Gym-v2 snapshot used by the prefix recovery test.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import shutil +import time +import uuid +from pathlib import Path +from typing import Any + + +def _read_json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text()) + if not isinstance(value, dict): + raise TypeError(f"expected a JSON object in {path}") + return value + + +def _ledger_episode_index(path: Path) -> dict[tuple[str, int], dict[str, Any]]: + import torch + + state = torch.load(path, weights_only=True) + result: dict[tuple[str, int], dict[str, Any]] = {} + for group in state["groups"]: + for sibling in group["siblings"]: + attempt = sibling["attempts"][-1] + generation_index = sibling["generation_index"] + rollout_id = ( + f"{group['group_id']}_g{generation_index}_a" + f"{uuid.UUID(bytes=attempt['attempt_uuid']).hex}" + ) + result[(rollout_id, attempt["gym_attempt"])] = { + "group_id": group["group_id"], + "generation_index": generation_index, + "target_step": group["target_step"], + "restore_level": group["restore_level"], + "status": attempt["status"], + "gym_instance_id": attempt["gym_instance_id"], + } + return result + + +def _participant_records( + snapshot: Path, kind: str +) -> list[tuple[Path, dict[str, Any]]]: + records: list[tuple[Path, dict[str, Any]]] = [] + pattern = f"gym-instances/*/replica-*/gym/{kind}/*/manifest.json" + for manifest_path in snapshot.glob(pattern): + manifest = _read_json(manifest_path) + if manifest.get("kind") != kind or manifest.get("record_count", 0) <= 0: + continue + records_path = manifest_path.parent / manifest["records_file"] + for line in records_path.read_text().splitlines(): + if line: + record = json.loads(line) + if not isinstance(record, dict): + raise TypeError( + f"participant row must be an object in {records_path}" + ) + records.append((manifest_path, record)) + return records + + +def _episode_key(record: dict[str, Any]) -> tuple[str, int]: + episode = record["episode_id"] + return episode["rollout_id"], episode["attempt"] + + +def _process_alive(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + return True + + +def _log_tail(path: Path, lines: int = 80) -> str: + if not path.is_file(): + return "" + return "\n".join(path.read_text(errors="replace").splitlines()[-lines:]) + + +def _recoverable_cut(record: dict[str, Any]) -> dict[str, Any] | None: + for cut in record.get("generation_cuts", []): + continuation = cut.get("continuation") + if not isinstance(continuation, dict): + continue + token_count = continuation.get("prefix_token_count") + output_limit = continuation.get("effective_output_limit") + staging_keys = continuation.get("staging_keys") + if ( + isinstance(token_count, int) + and isinstance(output_limit, int) + and 0 < token_count < output_limit + and isinstance(staging_keys, list) + and bool(staging_keys) + and continuation.get("terminal_finish_reason") is None + and continuation.get("terminal_stop_reason") is None + ): + return { + "model_call_id": cut["model_call_id"], + "request_digest": cut["request_digest"], + **continuation, + } + return None + + +def inspect_snapshot(snapshot: Path) -> dict[str, Any]: + """Return one fully cross-checked nonterminal prefix in ``snapshot``.""" + + required = [ + snapshot / "manifest.json", + snapshot / "gym_checkpoint.json", + snapshot / "rollout_recovery.pt", + snapshot / "replay_buffer_metadata.pt", + ] + if ( + not all(path.is_file() for path in required) + or not (snapshot / "data_plane").is_dir() + ): + raise FileNotFoundError("snapshot is not fully published") + manifest = _read_json(required[0]) + if manifest.get("base_train_step") != 0 or manifest.get("trainer_version") != 0: + raise AssertionError("prefix fault injection must use a bootstrap snapshot") + + gym = _read_json(required[1]) + ledger = _ledger_episode_index(required[2]) + agent_keys = { + _episode_key(record) for _, record in _participant_records(snapshot, "agent") + } + resources_keys = { + _episode_key(record) + for _, record in _participant_records(snapshot, "resources") + } + model_records = _participant_records(snapshot, "model") + + candidates: list[dict[str, Any]] = [] + for model_manifest, record in model_records: + key = _episode_key(record) + recovery = ledger.get(key) + cut = _recoverable_cut(record) + if recovery is None or cut is None: + continue + instance_id = recovery["gym_instance_id"] + checkpointed_episodes = { + (episode["rollout_id"], episode["attempt"]) + for episode in gym.get("instances", {}).get(instance_id, []) + } + retained_keys = set(gym.get("staging_keys", {}).get(instance_id, [])) + if ( + key not in checkpointed_episodes + or key not in agent_keys + or key not in resources_keys + or recovery["restore_level"] != "prefix" + or recovery["status"] != "dispatched" + or not set(cut["staging_keys"]).issubset(retained_keys) + ): + continue + candidates.append( + { + "snapshot": str(snapshot), + "snapshot_name": snapshot.name, + "gym_instance_id": instance_id, + "gym_rollout_id": key[0], + "gym_attempt": key[1], + "group_id": recovery["group_id"], + "generation_index": recovery["generation_index"], + "target_step": recovery["target_step"], + "model_manifest": str(model_manifest.relative_to(snapshot)), + **cut, + } + ) + if not candidates: + raise AssertionError( + "snapshot has no recoverable nonterminal active generation prefix" + ) + return max(candidates, key=lambda item: item["prefix_token_count"]) + + +def select_snapshot(args: argparse.Namespace) -> None: + root = Path(args.snapshot_root) + selection_path = Path(args.selection_file) + backup_path = Path(args.backup_path) + phase_log = Path(args.phase_log) + deadline = time.monotonic() + args.timeout_s + last_error = "no snapshot observed" + + while time.monotonic() < deadline: + for snapshot in sorted(root.glob("snapshot_*"), reverse=True): + try: + selection = inspect_snapshot(snapshot) + except (AssertionError, FileNotFoundError, json.JSONDecodeError) as error: + last_error = f"{snapshot}: {error}" + continue + try: + shutil.copytree(snapshot, backup_path) + except FileExistsError: + shutil.rmtree(backup_path) + continue + except FileNotFoundError: + shutil.rmtree(backup_path, ignore_errors=True) + continue + selection_path.write_text( + json.dumps(selection, sort_keys=True, indent=2) + "\n" + ) + return + if not _process_alive(args.phase_pid): + raise RuntimeError( + "phase one exited before publishing a recoverable generation prefix:\n" + + _log_tail(phase_log) + ) + time.sleep(0.2) + raise TimeoutError(f"no recoverable prefix snapshot found ({last_error})") + + +def _events(path: Path) -> list[dict[str, Any]]: + return [json.loads(line) for line in path.read_text().splitlines() if line] + + +def _matching_siblings( + events: list[dict[str, Any]], selection: dict[str, Any], event_name: str +) -> list[dict[str, Any]]: + found: list[dict[str, Any]] = [] + for event in events: + if event["event"] != event_name or event["group_id"] != selection["group_id"]: + continue + siblings = event.get("siblings") or [event.get("sibling")] + found.extend( + sibling + for sibling in siblings + if sibling is not None + and sibling["gym_rollout_id"] == selection["gym_rollout_id"] + ) + return found + + +def verify_restore(args: argparse.Namespace) -> None: + selection = _read_json(Path(args.selection_file)) + phase1 = _events(Path(args.phase1_events)) + phase2 = _events(Path(args.phase2_events)) + log = Path(args.phase2_log).read_text() + training_info = _read_json(Path(args.training_info)) + + phase1_dispatch = _matching_siblings(phase1, selection, "dispatch") + if not phase1_dispatch: + raise AssertionError( + "selected Gym episode was not observed in phase-one dispatches" + ) + if not all( + sibling["gym_instance_id"] == selection["gym_instance_id"] + and sibling["gym_attempt"] == selection["gym_attempt"] + for sibling in phase1_dispatch + ): + raise AssertionError( + f"phase-one dispatch identity mismatch: {phase1_dispatch!r}" + ) + + phase2_dispatch = _matching_siblings(phase2, selection, "dispatch") + phase2_refused = _matching_siblings(phase2, selection, "refused") + if len(phase2_dispatch) - len(phase2_refused) != 1: + raise AssertionError( + f"expected one accepted restored dispatch: {phase2_dispatch!r}, {phase2_refused!r}" + ) + restored = phase2_dispatch[-1] + expected_identity = ( + selection["gym_instance_id"], + selection["gym_attempt"] + 1, + selection["generation_index"], + ) + actual_identity = ( + restored["gym_instance_id"], + restored["gym_attempt"], + restored["generation_index"], + ) + if actual_identity != expected_identity: + raise AssertionError((actual_identity, expected_identity)) + + completions = _matching_siblings(phase2, selection, "completion") + if not any( + sibling["gym_instance_id"] == selection["gym_instance_id"] + and sibling["gym_attempt"] == selection["gym_attempt"] + 1 + and sibling["status"] == "sealed" + for sibling in completions + ): + raise AssertionError(f"restored Gym episode did not seal: {completions!r}") + + restored_prefixes = re.findall( + r"generation prefix restored: rollout_id=(\S+) model_call_id=(\S+) " + r"source_model_call_id=(\S+) prefix_tokens=(\d+) " + r"prefix_digest=([0-9a-f]{64})", + log, + ) + expected_prefix = ( + selection["model_call_id"], + selection["prefix_token_count"], + selection["prefix_digest"], + ) + if not any( + source_model_call_id == expected_prefix[0] + and int(prefix_tokens) == expected_prefix[1] + and prefix_digest == expected_prefix[2] + for _, _, source_model_call_id, prefix_tokens, prefix_digest in restored_prefixes + ): + raise AssertionError( + f"selected prefix was not restored: expected={expected_prefix!r}, " + f"observed={restored_prefixes!r}" + ) + + if training_info["current_step"] != args.max_steps: + raise AssertionError(training_info) + if training_info["trainer_version"] != args.max_steps: + raise AssertionError(training_info) + for step in range(1, args.max_steps + 1): + matches = re.findall(rf"train step {step}/{args.max_steps}(?:\s|$)", log) + if len(matches) != 1: + raise AssertionError(f"train step {step} appeared {len(matches)} times") + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser() + commands = parser.add_subparsers(dest="command", required=True) + + select = commands.add_parser("select") + select.add_argument("snapshot_root") + select.add_argument("selection_file") + select.add_argument("phase_pid", type=int) + select.add_argument("phase_log") + select.add_argument("timeout_s", type=float) + select.add_argument("backup_path") + select.set_defaults(func=select_snapshot) + + verify = commands.add_parser("verify-restore") + verify.add_argument("selection_file") + verify.add_argument("phase1_events") + verify.add_argument("phase2_events") + verify.add_argument("phase2_log") + verify.add_argument("training_info") + verify.add_argument("max_steps", type=int) + verify.set_defaults(func=verify_restore) + return parser + + +if __name__ == "__main__": + arguments = _parser().parse_args() + arguments.func(arguments) diff --git a/tests/functional/_single_controller_turn_recovery_hook.py b/tests/functional/_single_controller_turn_recovery_hook.py new file mode 100644 index 00000000000..eb324f8e473 --- /dev/null +++ b/tests/functional/_single_controller_turn_recovery_hook.py @@ -0,0 +1,224 @@ +# 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 that records Gym-recovery ownership. + +The wrapper does not alter scheduling or checkpoint timing. It records the +structured recovery identity once each Gym submission is marked dispatched, +when Gym refuses one at checkpoint admission, and after each completion has +been sealed in the RL ledger. A two-process functional test can +therefore prove that a restored episode returns to the same Gym instance under +``gym_attempt + 1`` without depending on log wording. +""" + +from __future__ import annotations + +import json +import os +from collections.abc import AsyncIterator, Callable, Sequence +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any + +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: + """Delegate Gym rollouts while recording durable identity transitions.""" + + def __init__( + self, + delegate: Any, + *, + recovery_ledger: Any, + events_path: Path, + ) -> None: + self._delegate = delegate + # Captured before the hook patches ``delegate.run_rollout`` to point here. + self._delegate_run_rollout = delegate.run_rollout + self._recovery_ledger = recovery_ledger + self._events_path = events_path + + 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( + "Gym 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] + + @staticmethod + def _sibling_fields(group: Any, generation_index: int) -> dict[str, Any]: + sibling = group.siblings[generation_index] + attempt = sibling.current_attempt + gym_rollout_id, gym_attempt = group.gym_episode(generation_index) + return { + "generation_index": generation_index, + "logical_rollout_id": group.logical_rollout_id(generation_index), + "gate_rollout_id": group.gate_rollout_id(generation_index), + "gym_rollout_id": gym_rollout_id, + "gym_instance_id": attempt.gym_instance_id, + "gym_attempt": gym_attempt, + "attempt_id": attempt.attempt_id, + "status": attempt.status.value, + } + + 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, + recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, + gym_instance_id: str | None = None, + dispatch_recorder: Any = None, + ) -> Any: + if rollout_ids is None or generation_indices is None or on_completion is None: + raise RuntimeError("Gym recovery hook requires token capture") + if gym_instance_id is None: + raise RuntimeError("Gym recovery hook requires a Gym checkpoint owner") + if dispatch_recorder is None: + raise RuntimeError("Gym recovery hook requires a dispatch recorder") + + group = self._find_group(rollout_ids) + indices = list(generation_indices) + common = { + "group_id": group.group_id, + "prompt_idx": int(input_sample["idx"]), + "target_step": group.target_step, + "generation_indices": indices, + "gym_instance_id": gym_instance_id, + } + + def _record_submission(event: str, submitted: list[int]) -> None: + # By group ID: a row re-sent after a broken stream has a new rollout ID. + refreshed = self._recovery_ledger.get_group(group.group_id) + self._append_event( + event, + **common, + siblings=[ + self._sibling_fields(refreshed, index) for index in submitted + ], + ) + + async def _record_completion(generation_index: int, completion: Any) -> None: + await on_completion(generation_index, completion) + refreshed = self._recovery_ledger.get_group(group.group_id) + self._append_event( + "completion", + **common, + sibling=self._sibling_fields(refreshed, generation_index), + ) + + result = await self._delegate_run_rollout( + input_sample, + rollout_ids=rollout_ids, + generation_indices=indices, + on_completion=_record_completion, + recovery_granularity=recovery_granularity, + gym_instance_id=gym_instance_id, + dispatch_recorder=_RecordingDispatchRecorder( + dispatch_recorder, record=_record_submission + ), + ) + self._append_event("rollout_returned", **common) + return result + + +class _RecordingDispatchRecorder: + """Record each physical Gym submission once the ledger marks it dispatched. + + A ``dispatch`` event is written right after Ray accepts the submission, so it + sees the attempt's real owner and status. A ``refused`` event marks a + submission Gym turned away at checkpoint admission; it never ran. + """ + + def __init__( + self, delegate: Any, *, record: Callable[[str, list[int]], None] + ) -> None: + self._delegate = delegate + self._record = record + + @asynccontextmanager + async def submitting( + self, generation_indices: Sequence[int] + ) -> AsyncIterator[Callable[[], None]]: + indices = list(generation_indices) + async with self._delegate.submitting(indices) as mark_dispatched: + + def mark_and_record() -> None: + mark_dispatched() + self._record("dispatch", indices) + + yield mark_and_record + + async def refused(self, generation_index: int) -> None: + await self._delegate.refused(generation_index) + self._record("refused", [generation_index]) + + async def dropped(self, generation_indices: Sequence[int], *, replace: bool) -> Any: + return await self._delegate.dropped(generation_indices, replace=replace) + + def retired(self, episodes: Sequence[Any]) -> None: + self._delegate.retired(episodes) + + +_original_setup_single_controller = run_grpo_single_controller.setup_single_controller + + +def _setup_with_gym_recovery_hook(*args: Any, **kwargs: Any) -> Any: + actor_args, timing_metrics = _original_setup_single_controller(*args, **kwargs) + manager = actor_args.rollout_manager + impl = manager._impl + events_path = os.environ.get("SC_GYM_RECOVERY_TEST_EVENTS") + if events_path is None: + events_path = os.environ["SC_TURN_RECOVERY_TEST_EVENTS"] + instrumented = _InstrumentedNemoGymRolloutImpl( + impl, + recovery_ledger=manager.recovery_ledger, + events_path=Path(events_path), + ) + # Patch the instance method rather than replacing ``_impl``: RolloutManager + # gates checkpoint admission and Gym recovery on + # ``isinstance(self._impl, AsyncNemoGymRolloutImpl)``. + impl.run_rollout = instrumented.run_rollout + return actor_args, timing_metrics + + +run_grpo_single_controller.setup_single_controller = _setup_with_gym_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 323cd96de95..d59d1dad7df 100755 --- a/tests/functional/grpo_async_gym_single_controller.sh +++ b/tests/functional/grpo_async_gym_single_controller.sh @@ -18,6 +18,7 @@ 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} +SC_CONFIG=${SC_TEST_CONFIG:-$PROJECT_ROOT/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml} export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} rm -rf $EXP_DIR $LOG_DIR @@ -58,7 +59,7 @@ jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/da uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ $SC_ENTRYPOINT \ - --config $PROJECT_ROOT/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml \ + --config "$SC_CONFIG" \ policy.model_name=Qwen/Qwen3-0.6B \ policy.dtensor_cfg.enabled=false \ policy.megatron_cfg.enabled=true \ diff --git a/tests/functional/grpo_async_gym_single_controller_prefix_recovery.sh b/tests/functional/grpo_async_gym_single_controller_prefix_recovery.sh new file mode 100755 index 00000000000..a2321630b39 --- /dev/null +++ b/tests/functional/grpo_async_gym_single_controller_prefix_recovery.sh @@ -0,0 +1,202 @@ +#!/bin/bash +# Two-process functional coverage for a Gym-v2 checkpoint taken during active +# vLLM decode. The restored attempt must reuse the exact durable prefix and +# generate only the remaining suffix before the rollout is finalized/trained. + +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 +BASE_RUN_LOG=$SCRIPT_DIR/grpo_async_gym_single_controller/run.log +TEST_CONFIG=${SC_PREFIX_RECOVERY_TEST_CONFIG:-$PROJECT_ROOT/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml} +RECOVERY_HOOK=$SCRIPT_DIR/_single_controller_turn_recovery_hook.py +SNAPSHOT_HELPER=$SCRIPT_DIR/_gym_prefix_recovery_snapshot.py +TEST_DIR=$SCRIPT_DIR/grpo_async_gym_single_controller_prefix_recovery +CHECKPOINT_DIR=$TEST_DIR/checkpoints +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 +SELECTION_FILE=$TEST_DIR/selected-snapshot.json +SELECTED_BACKUP=$TEST_DIR/selected-snapshot +TEST_DATA=$TEST_DIR/test-data.jsonl +GYM_ROOT=$PROJECT_ROOT/3rdparty/Gym-workspace/Gym +PHASE1_PID="" + +NUM_PROMPTS=${SC_PREFIX_RECOVERY_NUM_PROMPTS:-2} +NUM_GENERATIONS=${SC_PREFIX_RECOVERY_NUM_GENERATIONS:-2} +MIN_GENERATION_TOKENS=${SC_PREFIX_RECOVERY_MIN_TOKENS:-4096} +MAX_TOTAL_SEQUENCE_LENGTH=${SC_PREFIX_RECOVERY_MAX_TOTAL_SEQUENCE_LENGTH:-8192} +SNAPSHOT_INTERVAL_S=${SC_PREFIX_RECOVERY_INTERVAL_S:-0.05} +PHASE2_SNAPSHOT_INTERVAL_S=${SC_PREFIX_RECOVERY_PHASE2_INTERVAL_S:-600} +SNAPSHOT_TIMEOUT_S=${SC_PREFIX_RECOVERY_TIMEOUT_S:-2400} +PHASE2_TIMEOUT_S=${SC_PREFIX_RECOVERY_PHASE2_TIMEOUT_S:-2400} +MAX_STEPS=${SC_PREFIX_RECOVERY_MAX_STEPS:-1} +TRAIN_GLOBAL_BATCH_SIZE=$((NUM_PROMPTS * NUM_GENERATIONS)) + +rm -rf "$TEST_DIR" +mkdir -p "$TEST_DIR" + +# A long, tool-free first policy call isolates active decode recovery. min_tokens +# keeps the request alive until a periodic checkpoint can cut a non-empty, +# nonterminal prefix. +jq -c -s \ + --argjson count "$NUM_PROMPTS" \ + --argjson min_tokens "$MIN_GENERATION_TOKENS" ' + limit($count; .[]) + | .task_source = "example_session_state_mgmt_simple_agent" + | .responses_create_params.input = [{ + "role": "user", + "content": "Write a long numbered list. Continue until the output limit and do not call tools." + }] + | .responses_create_params.tools = [] + | .responses_create_params.tool_choice = "none" + | .responses_create_params.max_output_tokens = $min_tokens + | .responses_create_params.metadata = ((.responses_create_params.metadata // {}) + { + "extra_body": ({"min_tokens": $min_tokens} | tojson) + }) + ' "$GYM_ROOT/resources_servers/example_session_state_mgmt/data/example.jsonl" \ + > "$TEST_DATA" + +stop_phase1() { + if [[ -z "$PHASE1_PID" ]]; then + return + fi + kill -KILL -- "-$PHASE1_PID" 2>/dev/null || true + wait "$PHASE1_PID" 2>/dev/null || true + PHASE1_PID="" + # Ray actors outlive the driver until their raylet notices the crash. Do + # not start phase two while the old controller can still write snapshots. + local deadline=$((SECONDS + 120)) + while pgrep -f "^ray::SingleControllerActor" >/dev/null; do + if ((SECONDS >= deadline)); then + echo "[ERROR] phase-one SingleControllerActor outlived its driver" + return 1 + fi + sleep 1 + done +} + +cleanup() { + local status=$? + stop_phase1 + if [[ "$status" -eq 0 && "${SC_PREFIX_RECOVERY_KEEP_CHECKPOINTS:-0}" != "1" ]]; then + rm -rf "$CHECKPOINT_DIR" "$SELECTED_BACKUP" + else + echo "Preserving prefix-recovery artifacts for inspection: $TEST_DIR" + fi + return "$status" +} +trap cleanup EXIT + +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 + ++rollout_recovery.target_level=prefix + ++rollout_checkpointing.snapshot_attempt_interval_s="$SNAPSHOT_INTERVAL_S" + ++rollout_checkpointing.keep_latest_k=16 + ++rollout_checkpointing.restore_mode=latest + async_rl.sampler.name=in_order + async_rl.sampler.max_lookahead_versions=0 + async_rl.min_groups_for_streaming_train="$NUM_PROMPTS" + async_rl.max_inflight_prompts="$NUM_PROMPTS" + async_rl.max_buffered_rollouts="$NUM_PROMPTS" + ++async_rl.rollout_failure.nemo_gym.rollout_timeout_s=300 + ++async_rl.stall_watchdog.interval_s=10 + ++async_rl.stall_watchdog.stall_timeout_s=600 + ++async_rl.stall_watchdog.stall_action=abort + grpo.num_prompts_per_step="$NUM_PROMPTS" + grpo.num_generations_per_prompt="$NUM_GENERATIONS" + grpo.max_num_steps="$MAX_STEPS" + policy.max_total_sequence_length="$MAX_TOTAL_SEQUENCE_LENGTH" + policy.generation.max_new_tokens="$MIN_GENERATION_TOKENS" + policy.train_global_batch_size="$TRAIN_GLOBAL_BATCH_SIZE" + policy.generation.temperature=1.0 + '~env.nemo_gym.code_gen' + "env.nemo_gym.config_paths=[responses_api_models/vllm_model/configs/vllm_model_for_training.yaml,resources_servers/example_session_state_mgmt/configs/example_session_state_mgmt.yaml]" + data.train.data_path="$TEST_DATA" + data.validation.data_path="$TEST_DATA" +) + +echo "=== Phase 1: publish a checkpoint containing an active generation prefix ===" +command -v setsid >/dev/null +setsid env \ + SC_TEST_ENTRYPOINT="$RECOVERY_HOOK" \ + SC_TEST_CONFIG="$TEST_CONFIG" \ + SC_GYM_RECOVERY_TEST_EVENTS="$PHASE1_EVENTS" \ + RUN_CONVERGENCE_CHECKS=0 \ + bash "$BASE_TEST" "${COMMON_OVERRIDES[@]}" "$@" & +PHASE1_PID=$! + +uv run --directory "$PROJECT_ROOT" --no-sync python "$SNAPSHOT_HELPER" select \ + "$CHECKPOINT_DIR/bootstrap/rollout_snapshots" \ + "$SELECTION_FILE" \ + "$PHASE1_PID" \ + "$BASE_RUN_LOG" \ + "$SNAPSHOT_TIMEOUT_S" \ + "$SELECTED_BACKUP" + +stop_phase1 +cp "$BASE_RUN_LOG" "$PHASE1_LOG" +SNAPSHOT_DIR=$(uv run --directory "$PROJECT_ROOT" --no-sync python -c \ + 'import json, sys; print(json.load(open(sys.argv[1]))["snapshot"])' \ + "$SELECTION_FILE") +SNAPSHOT_ROOT=$(dirname "$SNAPSHOT_DIR") + +# Recover exactly the prefix selected above. A newer periodic snapshot or a +# trainer checkpoint would represent work performed after the simulated crash. +rm -rf "$SNAPSHOT_ROOT" +mkdir -p "$SNAPSHOT_ROOT" +cp -a "$SELECTED_BACKUP" "$SNAPSHOT_DIR" +for trainer_checkpoint in "$CHECKPOINT_DIR"/step_*; do + if [[ -d "$trainer_checkpoint" ]]; then + rm -rf "$trainer_checkpoint" + fi +done + +# Avoid cutting the replacement request again while verifying that it consumes +# the selected prefix and reaches normal completion. +for index in "${!COMMON_OVERRIDES[@]}"; do + if [[ "${COMMON_OVERRIDES[$index]}" == ++rollout_checkpointing.snapshot_attempt_interval_s=* ]]; then + COMMON_OVERRIDES[$index]="++rollout_checkpointing.snapshot_attempt_interval_s=$PHASE2_SNAPSHOT_INTERVAL_S" + fi +done + +echo "=== Phase 2: restore the prefix and generate its remaining suffix ===" +timeout --signal=TERM --kill-after=30s "${PHASE2_TIMEOUT_S}s" \ + env \ + SC_TEST_ENTRYPOINT="$RECOVERY_HOOK" \ + SC_TEST_CONFIG="$TEST_CONFIG" \ + SC_GYM_RECOVERY_TEST_EVENTS="$PHASE2_EVENTS" \ + RUN_CONVERGENCE_CHECKS=0 \ + bash "$BASE_TEST" "${COMMON_OVERRIDES[@]}" "$@" +cp "$BASE_RUN_LOG" "$PHASE2_LOG" + +grep -Fq "Selected rollout recovery snapshot: $SNAPSHOT_DIR" "$PHASE2_LOG" +grep -q "Native TQ checkpoint restored and validated" "$PHASE2_LOG" +grep -q "Loaded .* unfinished rollout group(s)" "$PHASE2_LOG" +grep -q "generation prefix restored:" "$PHASE2_LOG" +grep -q "train step $MAX_STEPS/$MAX_STEPS" "$PHASE2_LOG" + +PHASE2_METRICS=$TEST_DIR/phase2-metrics.json +uv run --directory "$PROJECT_ROOT" --no-sync tests/json_dump_tb_logs.py \ + "$SCRIPT_DIR/grpo_async_gym_single_controller/logs" \ + --output_path "$PHASE2_METRICS" +uv run --directory "$PROJECT_ROOT" --no-sync tests/check_metrics.py "$PHASE2_METRICS" \ + 'max(data["train/finalize/invalid_row_rate"]) == 0' \ + 'max(data["train/finalize/capture_poisoned_rollouts"]) == 0' + +uv run --directory "$PROJECT_ROOT" --no-sync python "$SNAPSHOT_HELPER" verify-restore \ + "$SELECTION_FILE" \ + "$PHASE1_EVENTS" \ + "$PHASE2_EVENTS" \ + "$PHASE2_LOG" \ + "$CHECKPOINT_DIR/step_$MAX_STEPS/training_info.json" \ + "$MAX_STEPS" + +echo "Single-controller Gym generation-prefix recovery functional test passed" 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 ddd1eafd7bc..8bf22aa812d 100755 --- a/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh @@ -43,7 +43,7 @@ COMMON_OVERRIDES=( checkpointing.save_period=1 ++checkpointing.save_data_plane=true ++token_capture.enabled=true - ++rollout_recovery.default_granularity=sibling + ++rollout_recovery.target_level=sibling async_rl.sampler.name=in_order async_rl.sampler.max_lookahead_versions=1 async_rl.max_inflight_prompts=8 diff --git a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh index 238a88d5475..c8b06e02cbc 100755 --- a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh @@ -54,7 +54,7 @@ COMMON_OVERRIDES=( checkpointing.metric_name=null ++checkpointing.save_data_plane=true ++token_capture.enabled=true - ++rollout_recovery.default_granularity=sibling + ++rollout_recovery.target_level=sibling ++rollout_checkpointing.snapshot_attempt_interval_s="$SNAPSHOT_INTERVAL_S" ++rollout_checkpointing.keep_latest_k=8 ++rollout_checkpointing.restore_mode=latest diff --git a/tests/functional/grpo_async_gym_single_controller_turn_recovery.sh b/tests/functional/grpo_async_gym_single_controller_turn_recovery.sh new file mode 100755 index 00000000000..f2948cf25ab --- /dev/null +++ b/tests/functional/grpo_async_gym_single_controller_turn_recovery.sh @@ -0,0 +1,442 @@ +#!/bin/bash +# Two-process functional coverage for Gym turn-level recovery with a real, +# stateful Workplace Assistant resources server. By default this uses one Gym +# actor and runs on one Ray node. Set SC_TURN_RECOVERY_TEST_CONFIG to the +# two-replica config and SC_TURN_RECOVERY_EXPECTED_GYM_INSTANCES=2 to exercise +# cross-node replica ownership. Policy training and vLLM generation still use +# the small model. + +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 +BASE_RUN_LOG=$SCRIPT_DIR/grpo_async_gym_single_controller/run.log +TEST_CONFIG=${SC_TURN_RECOVERY_TEST_CONFIG:-$PROJECT_ROOT/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml} +RECOVERY_HOOK=$SCRIPT_DIR/_single_controller_turn_recovery_hook.py +TEST_DIR=$SCRIPT_DIR/grpo_async_gym_single_controller_turn_recovery +CHECKPOINT_DIR=$TEST_DIR/checkpoints +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 +SELECTION_FILE=$TEST_DIR/selected-snapshot.json +# Bootstrap snapshots are pruned once step 1's trainer checkpoint lands, so the +# selector copies the chosen cut here before phase one can delete it. +SELECTED_BACKUP=$TEST_DIR/selected-snapshot +TEST_DATA=$TEST_DIR/test_data.jsonl +GYM_ROOT=$PROJECT_ROOT/3rdparty/Gym-workspace/Gym +SENTINEL_EVENT="NeMo RL checkpoint recovery sentinel" +PHASE1_PID="" + +EXPECTED_GYM_INSTANCES=${SC_TURN_RECOVERY_EXPECTED_GYM_INSTANCES:-1} +NUM_PROMPTS=${SC_TURN_RECOVERY_NUM_PROMPTS:-4} +NUM_GENERATIONS=${SC_TURN_RECOVERY_NUM_GENERATIONS:-2} +MAX_STEPS=${SC_TURN_RECOVERY_MAX_STEPS:-3} +SNAPSHOT_INTERVAL_S=${SC_TURN_RECOVERY_INTERVAL_S:-0.2} +SNAPSHOT_TIMEOUT_S=${SC_TURN_RECOVERY_TIMEOUT_S:-2400} +PHASE2_TIMEOUT_S=${SC_TURN_RECOVERY_PHASE2_TIMEOUT_S:-2400} +TRAIN_GLOBAL_BATCH_SIZE=$((NUM_PROMPTS * NUM_GENERATIONS)) + +rm -rf "$TEST_DIR" +mkdir -p "$TEST_DIR" + +# Force one real Workplace mutation per rollout. The checkpoint test agent holds +# the first boundary after it in phase one, so a published checkpoint always +# holds a turn whose calendar write has already been applied. +jq -c -s --argjson count "$NUM_PROMPTS" --arg event "$SENTINEL_EVENT" ' + limit($count; .[]) + | del(.agent_ref) + | .task_source = "workplace_assistant_checkpoint_test_agent" + | .responses_create_params.input = [{ + "role": "user", + "content": ("Call calendar_create_event exactly once with event_name " + $event + + ", participant_email checkpoint-recovery@example.com, event_start " + + "2025-01-15 10:00:00, and duration 30.") + }] + | .responses_create_params.tools = [ + .responses_create_params.tools[] | select(.name == "calendar_create_event") + ] + | .responses_create_params.tool_choice = {"type": "function", "name": "calendar_create_event"} + | .responses_create_params.parallel_tool_calls = false + | .ground_truth = [{ + "name": "calendar_create_event", + "arguments": ({ + "event_name": $event, + "participant_email": "checkpoint-recovery@example.com", + "event_start": "2025-01-15 10:00:00", + "duration": "30" + } | tojson) + }] + | .category = "workplace_assistant_calendar" + | .environment_name = "workplace_assistant" +' "$GYM_ROOT/resources_servers/workplace_assistant/data/example.jsonl" > "$TEST_DATA" + +stop_phase1() { + if [[ -z "$PHASE1_PID" ]]; then + return + fi + kill -KILL -- "-$PHASE1_PID" 2>/dev/null || true + wait "$PHASE1_PID" 2>/dev/null || true + PHASE1_PID="" + # Ray actors run in their own process groups and outlive the driver until + # their raylet notices; the controller may still be writing a snapshot. + local deadline=$((SECONDS + 120)) + while pgrep -f "^ray::SingleControllerActor" >/dev/null; do + if ((SECONDS >= deadline)); then + echo "[ERROR] phase-one SingleControllerActor outlived its driver" + return 1 + fi + sleep 1 + done +} + +cleanup() { + local status=$? + stop_phase1 + if [[ "$status" -eq 0 && "${SC_TURN_RECOVERY_KEEP_CHECKPOINTS:-0}" != "1" ]]; then + rm -rf "$CHECKPOINT_DIR" "$SELECTED_BACKUP" + else + echo "Preserving recovery artifacts for inspection: $TEST_DIR" + fi + return "$status" +} +trap cleanup EXIT + +# A sharded config names its servers per shard, rejects a top-level list, and +# already drops the base config's code_gen server. +if [[ "$EXPECTED_GYM_INSTANCES" -eq 1 ]]; then + GYM_CONFIG_OVERRIDES=( + # The base config overrides a code_gen server it never launches; Gym + # checkpoint discovery would probe it as a participant. + '~env.nemo_gym.code_gen' + "env.nemo_gym.config_paths=[responses_api_models/vllm_model/configs/vllm_model_for_training.yaml,responses_api_agents/checkpoint_test_agent/configs/workplace_assistant.yaml]" + ) +else + GYM_CONFIG_OVERRIDES=() +fi + +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 + ++rollout_recovery.target_level=turn + ++rollout_checkpointing.snapshot_attempt_interval_s="$SNAPSHOT_INTERVAL_S" + ++rollout_checkpointing.keep_latest_k=8 + ++rollout_checkpointing.restore_mode=latest + async_rl.sampler.name=in_order + async_rl.sampler.max_lookahead_versions=1 + async_rl.min_groups_for_streaming_train=2 + async_rl.max_inflight_prompts=4 + async_rl.max_buffered_rollouts=8 + ++async_rl.rollout_failure.nemo_gym.rollout_timeout_s=180 + ++async_rl.stall_watchdog.interval_s=10 + ++async_rl.stall_watchdog.stall_timeout_s=300 + ++async_rl.stall_watchdog.stall_action=abort + grpo.num_prompts_per_step="$NUM_PROMPTS" + grpo.num_generations_per_prompt="$NUM_GENERATIONS" + grpo.max_num_steps="$MAX_STEPS" + policy.train_global_batch_size="$TRAIN_GLOBAL_BATCH_SIZE" + data.train.data_path="$TEST_DATA" + data.validation.data_path="$TEST_DATA" + policy.generation.temperature=0.1 + "${GYM_CONFIG_OVERRIDES[@]}" +) + +echo "=== Phase 1: publish a Workplace turn checkpoint, then crash ===" +command -v setsid >/dev/null +setsid env \ + SC_TEST_ENTRYPOINT="$RECOVERY_HOOK" \ + SC_TEST_CONFIG="$TEST_CONFIG" \ + SC_TURN_RECOVERY_TEST_EVENTS="$PHASE1_EVENTS" \ + NEMO_GYM_TEST_HOLD_FIRST_MUTATED_BOUNDARY=1 \ + RUN_CONVERGENCE_CHECKS=0 \ + bash "$BASE_TEST" "${COMMON_OVERRIDES[@]}" "$@" & +PHASE1_PID=$! + +# Select only a fully published bootstrap (pre-step-1) snapshot whose Gym +# manifest, RL ledger, Workplace resource export and TQ references all describe +# the same episode. +uv run --directory "$PROJECT_ROOT" --no-sync python - \ + "$CHECKPOINT_DIR/bootstrap/rollout_snapshots" \ + "$SELECTION_FILE" \ + "$PHASE1_PID" \ + "$BASE_RUN_LOG" \ + "$SNAPSHOT_TIMEOUT_S" \ + "$EXPECTED_GYM_INSTANCES" \ + "$SELECTED_BACKUP" \ + "$SENTINEL_EVENT" <<'PY' +import json +import os +import shutil +import sys +import time +import uuid +from pathlib import Path + +import torch + +root = Path(sys.argv[1]) +selection_path = Path(sys.argv[2]) +phase_pid = int(sys.argv[3]) +phase_log = Path(sys.argv[4]) +deadline = time.monotonic() + float(sys.argv[5]) +expected_gym_instances = int(sys.argv[6]) +backup_path = Path(sys.argv[7]) +sentinel_event = sys.argv[8] +if expected_gym_instances < 1: + raise ValueError("expected Gym instance count must be positive") + + +def ledger_episode_index(path: Path) -> dict[tuple[str, int], dict]: + state = torch.load(path, weights_only=True) + result = {} + for group in state["groups"]: + for sibling in group["siblings"]: + attempt = sibling["attempts"][-1] + generation_index = sibling["generation_index"] + rollout_id = ( + f"{group['group_id']}_g{generation_index}_a" + f"{uuid.UUID(bytes=attempt['attempt_uuid']).hex}" + ) + result[(rollout_id, attempt["gym_attempt"])] = { + "group_id": group["group_id"], + "generation_index": generation_index, + "target_step": group["target_step"], + "restore_level": group["restore_level"], + "status": attempt["status"], + "gym_instance_id": attempt["gym_instance_id"], + } + return result + + +def sentinel_count(state: dict) -> int: + """Count calendar rows the held turn wrote before the crash.""" + columns = state["calendar"]["_calendar_events"]["columns"] + names = next(column["values"] for column in columns if column["name"] == "event_name") + return sum(str(name).lower() == sentinel_event.lower() for name in names) + + +def participant_records(snapshot: Path, kind: str) -> list[tuple[Path, dict]]: + records = [] + pattern = f"gym-instances/*/replica-*/gym/{kind}/*/manifest.json" + for manifest_path in snapshot.glob(pattern): + manifest = json.loads(manifest_path.read_text()) + if manifest.get("kind") != kind or manifest.get("record_count", 0) <= 0: + continue + records_path = manifest_path.parent / manifest["records_file"] + for line in records_path.read_text().splitlines(): + if line: + records.append((manifest_path, json.loads(line))) + return records + + +while time.monotonic() < deadline: + for snapshot in sorted(root.glob("snapshot_*"), reverse=True): + required = [ + snapshot / "manifest.json", + snapshot / "gym_checkpoint.json", + snapshot / "rollout_recovery.pt", + snapshot / "replay_buffer_metadata.pt", + ] + if not all(path.is_file() for path in required) or not (snapshot / "data_plane").is_dir(): + continue + manifest = json.loads(required[0].read_text()) + if manifest.get("base_train_step") != 0 or manifest.get("trainer_version") != 0: + continue + try: + gym = json.loads(required[1].read_text()) + ledger = ledger_episode_index(required[2]) + resources = participant_records(snapshot, "resources") + model = participant_records(snapshot, "model") + except FileNotFoundError: + # Pruned by phase one's first trainer checkpoint mid-inspection. + continue + if len(gym["instances"]) != expected_gym_instances: + continue + resource_index = { + (record["episode_id"]["rollout_id"], record["episode_id"]["attempt"]): ( + path, + record["state"], + ) + for path, record in resources + if record.get("state") + } + model_index = { + (record["episode_id"]["rollout_id"], record["episode_id"]["attempt"]): path + for path, record in model + if record.get("rows") + } + for instance_id, episodes in gym["instances"].items(): + if not gym["staging_keys"].get(instance_id): + continue + for episode in episodes: + key = (episode["rollout_id"], episode["attempt"]) + recovery = ledger.get(key) + resources_entry = resource_index.get(key) + model_manifest = model_index.get(key) + if ( + recovery is None + or resources_entry is None + or model_manifest is None + ): + continue + resources_manifest, resources_state = resources_entry + # The held turn applied its calendar write exactly once. + if sentinel_count(resources_state) != 1: + continue + if ( + recovery["restore_level"] != "turn" + or recovery["status"] != "dispatched" + or recovery["gym_instance_id"] != instance_id + ): + continue + try: + shutil.copytree(snapshot, backup_path) + except FileNotFoundError: + shutil.rmtree(backup_path, ignore_errors=True) + continue + selection_path.write_text( + json.dumps( + { + "snapshot": str(snapshot), + "snapshot_name": snapshot.name, + "gym_instance_id": instance_id, + "gym_rollout_id": episode["rollout_id"], + "gym_attempt": episode["attempt"], + "group_id": recovery["group_id"], + "generation_index": recovery["generation_index"], + "target_step": recovery["target_step"], + "resources_manifest": str(resources_manifest.relative_to(snapshot)), + "model_manifest": str(model_manifest.relative_to(snapshot)), + }, + sort_keys=True, + indent=2, + ) + + "\n" + ) + raise SystemExit(0) + try: + os.kill(phase_pid, 0) + except ProcessLookupError as error: + tail = "" + if phase_log.is_file(): + tail = "\n".join(phase_log.read_text(errors="replace").splitlines()[-60:]) + raise RuntimeError( + "phase one exited before publishing a recoverable Workplace turn:\n" + tail + ) from error + time.sleep(0.2) +raise TimeoutError("no bootstrap snapshot contained a recoverable Workplace turn") +PY + +stop_phase1 +cp "$BASE_RUN_LOG" "$PHASE1_LOG" +SNAPSHOT_DIR=$(uv run --directory "$PROJECT_ROOT" --no-sync python -c \ + 'import json, sys; print(json.load(open(sys.argv[1]))["snapshot"])' \ + "$SELECTION_FILE") +SNAPSHOT_ROOT=$(dirname "$SNAPSHOT_DIR") + +# Make recovery select exactly the validated cut: phase one may have published +# newer snapshots or pruned this one, and any trainer checkpoint it wrote would +# take precedence over the bootstrap anchor. +rm -rf "$SNAPSHOT_ROOT" +mkdir -p "$SNAPSHOT_ROOT" +cp -a "$SELECTED_BACKUP" "$SNAPSHOT_DIR" +for trainer_checkpoint in "$CHECKPOINT_DIR"/step_*; do + if [[ -d "$trainer_checkpoint" ]]; then + rm -rf "$trainer_checkpoint" + fi +done + +echo "=== Phase 2: restore the turn on its original Gym replica ===" +timeout --signal=TERM --kill-after=30s "${PHASE2_TIMEOUT_S}s" \ + env \ + SC_TEST_ENTRYPOINT="$RECOVERY_HOOK" \ + SC_TEST_CONFIG="$TEST_CONFIG" \ + SC_TURN_RECOVERY_TEST_EVENTS="$PHASE2_EVENTS" \ + RUN_CONVERGENCE_CHECKS=0 \ + bash "$BASE_TEST" "${COMMON_OVERRIDES[@]}" "$@" +cp "$BASE_RUN_LOG" "$PHASE2_LOG" + +grep -Fq "Selected rollout recovery snapshot: $SNAPSHOT_DIR" "$PHASE2_LOG" +grep -q "Native TQ checkpoint restored and validated" "$PHASE2_LOG" +grep -q "Loaded .* unfinished rollout group(s)" "$PHASE2_LOG" +grep -q "train step $MAX_STEPS/$MAX_STEPS" "$PHASE2_LOG" + +PHASE2_METRICS=$TEST_DIR/phase2-metrics.json +uv run --directory "$PROJECT_ROOT" --no-sync tests/json_dump_tb_logs.py \ + "$SCRIPT_DIR/grpo_async_gym_single_controller/logs" \ + --output_path "$PHASE2_METRICS" +uv run --directory "$PROJECT_ROOT" --no-sync tests/check_metrics.py "$PHASE2_METRICS" \ + 'max(data["train/finalize/invalid_row_rate"]) == 0' \ + 'max(data["train/finalize/capture_poisoned_rollouts"]) == 0' + +uv run --directory "$PROJECT_ROOT" --no-sync python - \ + "$SELECTION_FILE" "$PHASE1_EVENTS" "$PHASE2_EVENTS" "$PHASE2_LOG" \ + "$CHECKPOINT_DIR/step_$MAX_STEPS/training_info.json" "$MAX_STEPS" <<'PY' +import json +import re +import sys +from pathlib import Path + +selection = json.loads(Path(sys.argv[1]).read_text()) +phase1 = [json.loads(line) for line in Path(sys.argv[2]).read_text().splitlines()] +phase2 = [json.loads(line) for line in Path(sys.argv[3]).read_text().splitlines()] +log = Path(sys.argv[4]).read_text() +training_info = json.loads(Path(sys.argv[5]).read_text()) +max_steps = int(sys.argv[6]) + + +def matching_siblings(events: list[dict], event_name: str) -> list[dict]: + found = [] + for event in events: + if event["event"] != event_name or event["group_id"] != selection["group_id"]: + continue + siblings = event.get("siblings") or [event.get("sibling")] + found.extend( + sibling + for sibling in siblings + if sibling is not None + and sibling["gym_rollout_id"] == selection["gym_rollout_id"] + ) + return found + + +phase1_dispatch = matching_siblings(phase1, "dispatch") +assert phase1_dispatch, "selected Gym episode was not observed in phase-one dispatches" +assert all( + sibling["gym_instance_id"] == selection["gym_instance_id"] + and sibling["gym_attempt"] == selection["gym_attempt"] + for sibling in phase1_dispatch +), phase1_dispatch + +phase2_dispatch = matching_siblings(phase2, "dispatch") +# A submission Gym refused at checkpoint admission never ran; it is re-sent. +phase2_refused = matching_siblings(phase2, "refused") +assert len(phase2_dispatch) - len(phase2_refused) == 1, (phase2_dispatch, phase2_refused) +restored = phase2_dispatch[-1] +assert restored["gym_instance_id"] == selection["gym_instance_id"], restored +assert restored["gym_attempt"] == selection["gym_attempt"] + 1, restored +assert restored["generation_index"] == selection["generation_index"], restored + +phase2_completion = matching_siblings(phase2, "completion") +assert phase2_completion, "restored Gym episode never completed" +assert any( + sibling["gym_instance_id"] == selection["gym_instance_id"] + and sibling["gym_attempt"] == selection["gym_attempt"] + 1 + and sibling["status"] == "sealed" + for sibling in phase2_completion +), phase2_completion + +assert training_info["current_step"] == max_steps, training_info +assert training_info["trainer_version"] == max_steps, training_info +# Restored from the bootstrap anchor, so every step trains exactly once. +for step in range(1, max_steps + 1): + matches = re.findall(rf"train step {step}/{max_steps}(?:\s|$)", log) + assert len(matches) == 1, (step, len(matches)) +PY + +echo "Workplace turn-level recovery functional test passed." diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index c3f6b84f3ee..537ac252523 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -3306,6 +3306,7 @@ def __init__(self, *, cluster, config, tokenizer, tokenizer_config): tokenizer=tokenizer, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) @@ -3687,6 +3688,7 @@ def get_refit_payload_mode(self): tokenizer=tokenizer, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) diff --git a/tests/unit/data/test_captured_media.py b/tests/unit/data/test_captured_media.py index a00b2702b7d..b07a137725e 100644 --- a/tests/unit/data/test_captured_media.py +++ b/tests/unit/data/test_captured_media.py @@ -15,6 +15,7 @@ import inspect import json +import threading from dataclasses import dataclass, replace from types import SimpleNamespace from unittest.mock import Mock @@ -674,15 +675,25 @@ def test_worker_restart_recovers_retained_geometry_without_fetching_pixels( def test_worker_completion_stages_pixels_and_only_returns_capture_coordinates(dp): from nemo_gym.token_id_capture.adapters.vllm import VLLMCaptureAdapter + from nemo_rl.models.generation.generation_cut_capture import ( + _TokenCaptureSnapshotGate, + ) + worker = object.__new__(VllmAsyncGenerationWorkerImpl) worker._capture_calls = {} + worker._capture_calls_by_model_call_id = {} + worker._completed_capture_calls = {} + worker._capture_registry_lock = threading.Lock() + worker._token_capture_snapshot_gate = _TokenCaptureSnapshotGate() + sink = TQTokenSink( + dp, + staging_partition="staging", + capture_media=True, + media_pixel_dtype=torch.float32, + ) + worker._capture_sink = sink worker.token_capture = RolloutTokenCapture( - sink=TQTokenSink( - dp, - staging_partition="staging", - capture_media=True, - media_pixel_dtype=torch.float32, - ), + sink=sink, weight_version_fn=lambda: 0, adapter=VLLMCaptureAdapter(), ) @@ -1273,6 +1284,9 @@ async def test_retained_media_rejection_precedes_inference_and_survives_gym( ) worker.token_capture = MagicMock(adapter=VLLMCaptureAdapter()) worker._capture_calls = {} + worker._capture_calls_by_model_call_id = {} + worker._completed_capture_calls = {} + worker._capture_registry_lock = threading.Lock() worker._http_engine_client = MagicMock() worker.llm_async_engine_args = MagicMock() worker.llm_async_engine_args.create_model_config.return_value = SimpleNamespace( diff --git a/tests/unit/data_plane/test_tq_token_sink.py b/tests/unit/data_plane/test_tq_token_sink.py index d9ad74d408c..69fab4b4412 100644 --- a/tests/unit/data_plane/test_tq_token_sink.py +++ b/tests/unit/data_plane/test_tq_token_sink.py @@ -53,6 +53,7 @@ ChainPrefixCache, TQTokenSink, TQTokenSource, + generation_cut_staging_key, resolve_admission_prefix, ) from nemo_rl.models.generation.megatron.token_capture import ( # noqa: E402 @@ -471,6 +472,38 @@ def test_sink_clear_drops_rows(tq_client, staging_partition): source.fetch(keys) +def test_generation_prefix_round_trips_under_checkpoint_scoped_keys( + tq_client, staging_partition +): + sink = TQTokenSink(tq_client, staging_partition=staging_partition) + source = TQTokenSource(tq_client, staging_partition=staging_partition) + records, _, _ = build_fixture_artifacts("single_call") + record = records[0] + + first = sink.stage_generation_prefix( + record, checkpoint_id="checkpoint-1", chunk_sequence=0 + ) + second = sink.stage_generation_prefix( + record, checkpoint_id="checkpoint-1", chunk_sequence=1 + ) + keys = [ + generation_cut_staging_key( + "checkpoint-1", + record.rollout_id, + record.model_call_id, + chunk_sequence=sequence, + ) + for sequence in (0, 1) + ] + + assert first.ok and second.ok + assert [first.staging_key, second.staging_key] == keys + assert [snapshot.model_dump() for snapshot in source.fetch(keys)] == [ + record.model_dump(exclude={"extras"}), + record.model_dump(exclude={"extras"}), + ] + + def test_fetch_prefix_token_ids_empty(tq_client, staging_partition): source = TQTokenSource(tq_client, staging_partition=staging_partition) assert source.fetch_prefix_token_ids([]) == [] diff --git a/tests/unit/environments/test_gym_checkpoint_adapter.py b/tests/unit/environments/test_gym_checkpoint_adapter.py new file mode 100644 index 00000000000..86a0186aac9 --- /dev/null +++ b/tests/unit/environments/test_gym_checkpoint_adapter.py @@ -0,0 +1,952 @@ +# 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. + +from __future__ import annotations + +import asyncio +import pickle +import uuid +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from nemo_rl.environments.gym_checkpoint_adapter import ( + GymCheckpointAdapter, + GymCheckpointEpisode, + GymCheckpointInstance, + GymCheckpointParticipantManifest, + GymCheckpointParticipantRecord, + GymCheckpointUnavailable, +) +from nemo_rl.experience.rollout_recovery import ( + PromptGroupPhase, + PromptGroupRecoveryRecord, + PromptGroupStatus, + PromptRef, + RecoveryGranularity, + RecoveryTargetLevel, + RolloutAttemptRecord, + RolloutAttemptStatus, + RolloutSiblingRecord, +) + +pytestmark = pytest.mark.nemo_gym + + +@pytest.mark.parametrize("gym_attempt", [0, 1, 17]) +def test_rollout_capture_key_matches_gym_episode_identity(gym_attempt: int) -> None: + """Fail loudly if RL's legacy string carrier drifts from Gym's codec.""" + from nemo_gym.episode_types import EpisodeId + + group = PromptGroupRecoveryRecord( + group_id="group-7", + admission_id="batch-7", + prompt_id="7", + prompt_ref=PromptRef(sample_id="7", task_name=None), + task_source=None, + recovery_granularity=RecoveryGranularity.SIBLING, + restore_level=RecoveryTargetLevel.TURN, + runtime_prompt_payload=None, + expected_generations=1, + target_step=7, + start_weight_version=6, + siblings=[ + RolloutSiblingRecord( + generation_index=0, + attempts=[ + RolloutAttemptRecord( + attempt_uuid=uuid.UUID(int=1), + status=RolloutAttemptStatus.DISPATCHED, + gym_instance_id="tools/replica-0", + gym_attempt=gym_attempt, + ) + ], + ) + ], + phase=PromptGroupPhase.ADMITTED, + status=PromptGroupStatus.GENERATING, + ) + + rollout_id, attempt = group.gym_episode(0) + gym_episode = EpisodeId(rollout_id=rollout_id, attempt=attempt) + capture_key = group.gate_rollout_id(0) + + assert capture_key == gym_episode.capture_key + assert EpisodeId.from_capture_key(capture_key) == gym_episode + + +def _participants(client: object, label: str) -> SimpleNamespace: + return SimpleNamespace( + client=client, + members=( + SimpleNamespace(server_name=f"{label}-environment", kind="environment"), + SimpleNamespace(server_name=f"{label}-model", kind="model"), + ), + ) + + +def _committed_participants( + checkpoint_root: Path, label: str, *, checkpoint_id: str = "save-1" +) -> tuple[GymCheckpointParticipantRecord, ...]: + """Write each participant's state as Gym commits it; return what RL records.""" + from nemo_gym._checkpoint.store import write_participant_state + + checkpoint_dir = GymCheckpointInstance(label, 0).checkpoint_dir(checkpoint_root) + records = [] + for kind in ("environment", "model"): + server_name = f"{label}-{kind}" + stored = write_participant_state( + checkpoint_dir, + kind=kind, + instance=server_name, + checkpoint_id=checkpoint_id, + records=[{"owner": server_name}], + ) + records.append( + GymCheckpointParticipantRecord( + server_name=server_name, + kind=kind, + instance=server_name, + records_sha256=stored["records_sha256"], + ) + ) + return tuple(records) + + +def _participant_manifest( + label: str, + kind: str, + *, + checkpoint_id: str = "save-1", + record_count: int = 1, +) -> dict[str, object]: + return { + "schema_version": 1, + "kind": kind, + "instance": f"{label}-{kind}", + "checkpoint_id": checkpoint_id, + "records_file": "records.jsonl", + "records_sha256": "0" * 64, + "record_count": record_count, + } + + +def _commit_reply( + label: str, + kind: str, + episode_ids: list[str], + *, + staging_keys: list[str] | None = None, + staging_keys_by_episode: dict[str, list[str]] | None = None, +) -> dict[str, object]: + reply: dict[str, object] = { + "phase": "committed", + "manifest": _participant_manifest( + label, + kind, + record_count=len(episode_ids), + ), + "episode_ids": episode_ids, + } + if staging_keys is not None: + reply["staging_keys"] = staging_keys + if staging_keys_by_episode is not None: + reply["staging_keys_by_episode"] = staging_keys_by_episode + return reply + + +def test_checkpoint_instance_paths_are_stable_and_isolated() -> None: + first = GymCheckpointInstance(shard_name="tools", replica_index=0) + replacement = GymCheckpointInstance(shard_name="tools", replica_index=0) + sibling = GymCheckpointInstance(shard_name="tools", replica_index=1) + + assert first.instance_id == replacement.instance_id == "tools/replica-0" + assert first.live_capture_dir("/capture") == replacement.live_capture_dir( + "/capture" + ) + assert first.live_capture_dir("/capture") != sibling.live_capture_dir("/capture") + assert first.checkpoint_dir("/checkpoint") == ( + replacement.checkpoint_dir("/checkpoint") + ) + + +@pytest.mark.parametrize( + ("shard_name", "replica_index"), + [("", 0), (".", 0), ("..", 0), ("a/b", 0), ("a\\b", 0), ("a", -1)], +) +def test_checkpoint_instance_rejects_unsafe_identity( + shard_name: str, replica_index: int +) -> None: + with pytest.raises(ValueError): + GymCheckpointInstance( + shard_name=shard_name, + replica_index=replica_index, + ) + + +def test_adapter_delegates_the_complete_gym_v2_lifecycle( + monkeypatch, tmp_path: Path +) -> None: + from nemo_gym._checkpoint import coordination + + client = object() + participants = _participants(client, "actor-a") + prepare_result = object() + commit_result = { + "actor-a-environment": _commit_reply( + "actor-a", "environment", ["rollout-1-a2"] + ), + "actor-a-model": _commit_reply( + "actor-a", + "model", + ["rollout-1-a2"], + staging_keys=["stage-2", "stage-1"], + staging_keys_by_episode={"rollout-1-a2": ["stage-2", "stage-1"]}, + ), + } + restore_result = { + "actor-a-environment": {"source_checkpoint_id": "save-1"}, + "actor-a-model": {"source_checkpoint_id": "save-1"}, + } + + discover = AsyncMock(return_value=participants) + prepare = AsyncMock(return_value=prepare_result) + renew = AsyncMock() + retire = AsyncMock() + commit = AsyncMock(return_value=commit_result) + restore = AsyncMock(return_value=restore_result) + resume = AsyncMock() + for name, operation in ( + ("discover", discover), + ("prepare", prepare), + ("renew", renew), + ("retire", retire), + ("commit", commit), + ("restore", restore), + ("resume", resume), + ): + monkeypatch.setattr(coordination, name, operation) + + instance = GymCheckpointInstance(shard_name="actor-a", replica_index=0) + adapter = GymCheckpointAdapter( + instance=instance, + client=client, + auth_token="secret", + ) + episodes = [GymCheckpointEpisode("rollout-1", 2)] + + async def exercise() -> None: + summary = await adapter.discover() + assert summary.instance_id == "actor-a/replica-0" + assert summary.members == ( + ("actor-a-environment", "environment"), + ("actor-a-model", "model"), + ) + # Discovery is cached inside the actor-local adapter. + assert await adapter.discover() == summary + assert await adapter.prepare("save-1", deadline_ts=10.0) is prepare_result + await adapter.renew("save-1", deadline_ts=11.0) + await adapter.retire("save-1", episodes, deadline_ts=12.0) + commit_summary = await adapter.commit( + "save-1", + "/checkpoints/step-1", + episodes, + deadline_ts=13.0, + ) + assert commit_summary.exported_episodes == ( + GymCheckpointEpisode("rollout-1", 2), + ) + assert commit_summary.staging_keys == ("stage-1", "stage-2") + assert commit_summary.staging_keys_by_episode == { + GymCheckpointEpisode("rollout-1", 2): ("stage-1", "stage-2") + } + assert [ + participant.server_name for participant in commit_summary.participants + ] == ["actor-a-environment", "actor-a-model"] + assert commit_summary.participants[0].manifest == ( + GymCheckpointParticipantManifest( + schema_version=1, + kind="environment", + instance="actor-a-environment", + checkpoint_id="save-1", + records_file="records.jsonl", + records_sha256="0" * 64, + record_count=1, + ) + ) + await adapter.restore( + "restore-1", + tmp_path, + episodes, + participants=_committed_participants(tmp_path, "actor-a"), + source_checkpoint_id="save-1", + deadline_ts=14.0, + ) + await adapter.resume("save-1", deadline_ts=15.0) + + asyncio.run(exercise()) + + discover.assert_awaited_once_with(client, auth_token="secret") + prepare.assert_awaited_once_with(participants, "save-1", deadline_ts=10.0) + renew.assert_awaited_once_with(participants, "save-1", deadline_ts=11.0) + (retire_episode,) = retire.await_args.args[2] + assert retire_episode.rollout_id == "rollout-1" + assert retire_episode.attempt == 2 + retire.assert_awaited_once_with( + participants, + "save-1", + retire.await_args.args[2], + deadline_ts=12.0, + ) + instance_dir = "/checkpoints/step-1/gym-instances/actor-a/replica-0" + (commit_episode,) = commit.await_args.args[3] + assert commit_episode == retire_episode + commit.assert_awaited_once_with( + participants, + "save-1", + instance_dir, + commit.await_args.args[3], + deadline_ts=13.0, + ) + (restore_episode,) = restore.await_args.args[3] + assert restore_episode == retire_episode + restore.assert_awaited_once_with( + participants, + "restore-1", + str(tmp_path / "gym-instances" / "actor-a" / "replica-0"), + restore.await_args.args[3], + deadline_ts=14.0, + ) + resume.assert_awaited_once_with(participants, "save-1", deadline_ts=15.0) + + +def test_two_adapters_keep_their_clients_and_participants_isolated(monkeypatch) -> None: + from nemo_gym._checkpoint import coordination + + client_a = object() + client_b = object() + participants_a = _participants(client_a, "actor-a") + participants_b = _participants(client_b, "actor-b") + + async def discover(client, *, auth_token): + assert auth_token == "secret" + return participants_a if client is client_a else participants_b + + prepare = AsyncMock(return_value=object()) + monkeypatch.setattr(coordination, "discover", discover) + monkeypatch.setattr(coordination, "prepare", prepare) + + adapter_a = GymCheckpointAdapter( + instance=GymCheckpointInstance("actor-a", 0), + client=client_a, + auth_token="secret", + ) + adapter_b = GymCheckpointAdapter( + instance=GymCheckpointInstance("actor-b", 0), + client=client_b, + auth_token="secret", + ) + + async def exercise() -> None: + await asyncio.gather(adapter_a.discover(), adapter_b.discover()) + await asyncio.gather( + adapter_a.prepare("save-1", deadline_ts=10.0), + adapter_b.prepare("save-1", deadline_ts=10.0), + ) + + asyncio.run(exercise()) + + assert prepare.await_args_list[0].args[0] is participants_a + assert prepare.await_args_list[1].args[0] is participants_b + + +def test_adapter_reports_candidates_not_exported_by_the_environment( + monkeypatch, +) -> None: + from nemo_gym._checkpoint import coordination + + client = object() + participants = _participants(client, "actor-a") + monkeypatch.setattr( + coordination, + "discover", + AsyncMock(return_value=participants), + ) + monkeypatch.setattr( + coordination, + "commit", + AsyncMock( + return_value={ + "actor-a-environment": _commit_reply("actor-a", "environment", []), + "actor-a-model": _commit_reply( + "actor-a", "model", [], staging_keys=[], staging_keys_by_episode={} + ), + } + ), + ) + adapter = GymCheckpointAdapter( + instance=GymCheckpointInstance("actor-a", 0), + client=client, + auth_token="secret", + ) + + async def exercise() -> None: + await adapter.discover() + summary = await adapter.commit( + "save-1", + "/checkpoints/step-1", + [GymCheckpointEpisode("rollout-1", 0)], + deadline_ts=10.0, + ) + assert summary.exported_episodes == () + assert summary.staging_keys == () + + asyncio.run(exercise()) + + +def test_adapter_does_not_require_optional_agent_or_resource_episode_state( + monkeypatch, +) -> None: + from nemo_gym._checkpoint import coordination + + client = object() + participants = SimpleNamespace( + client=client, + members=( + SimpleNamespace(server_name="environment", kind="environment"), + SimpleNamespace(server_name="model", kind="model"), + SimpleNamespace(server_name="agent", kind="agent"), + SimpleNamespace(server_name="resources", kind="resources"), + ), + ) + monkeypatch.setattr( + coordination, + "discover", + AsyncMock(return_value=participants), + ) + monkeypatch.setattr( + coordination, + "commit", + AsyncMock( + return_value={ + "environment": _commit_reply( + "deployment", "environment", ["rollout-1"] + ), + "model": _commit_reply( + "deployment", + "model", + ["rollout-1"], + staging_keys=["rollout-1/call-1"], + staging_keys_by_episode={"rollout-1": ["rollout-1/call-1"]}, + ), + "agent": _commit_reply("deployment", "agent", []), + "resources": _commit_reply("deployment", "resources", []), + } + ), + ) + adapter = GymCheckpointAdapter( + instance=GymCheckpointInstance("deployment", 0), + client=client, + auth_token="secret", + ) + + async def exercise() -> None: + await adapter.discover() + summary = await adapter.commit( + "save-1", + "/checkpoints/step-1", + [GymCheckpointEpisode("rollout-1", 0)], + deadline_ts=10.0, + ) + assert summary.exported_episodes == (GymCheckpointEpisode("rollout-1", 0),) + assert summary.staging_keys == ("rollout-1/call-1",) + assert { + participant.kind: participant.capture_keys + for participant in summary.participants + } == { + "environment": ("rollout-1",), + "model": ("rollout-1",), + "agent": (), + "resources": (), + } + + asyncio.run(exercise()) + + +def _commit_model_reply(monkeypatch, model_reply: dict[str, object]): + """Commit two episodes with an environment and the given policy-model reply.""" + from nemo_gym._checkpoint import coordination + + client = object() + participants = SimpleNamespace( + client=client, + members=( + SimpleNamespace(server_name="environment", kind="environment"), + SimpleNamespace(server_name="model", kind="model"), + ), + ) + monkeypatch.setattr(coordination, "discover", AsyncMock(return_value=participants)) + monkeypatch.setattr( + coordination, + "commit", + AsyncMock( + return_value={ + "environment": _commit_reply( + "deployment", "environment", ["rollout-1", "rollout-2-a1"] + ), + "model": model_reply, + } + ), + ) + adapter = GymCheckpointAdapter( + instance=GymCheckpointInstance("deployment", 0), + client=client, + auth_token="secret", + ) + + async def exercise(): + await adapter.discover() + return await adapter.commit( + "save-1", + "/checkpoints/step-1", + [ + GymCheckpointEpisode("rollout-1", 0), + GymCheckpointEpisode("rollout-2", 1), + ], + deadline_ts=10.0, + ) + + return asyncio.run(exercise()) + + +def test_staging_keys_are_attributed_to_their_episodes_by_gym(monkeypatch) -> None: + """A cut key names no rollout and a carried-over row serves two attempts.""" + cut_key = "__generation_cut__/save-1/c9" + summary = _commit_model_reply( + monkeypatch, + _commit_reply( + "deployment", + "model", + ["rollout-1", "rollout-2-a1"], + staging_keys=["rollout-1/c1", "rollout-2/c1", cut_key], + staging_keys_by_episode={ + "rollout-1": ["rollout-1/c1"], + "rollout-2-a1": ["rollout-2/c1", cut_key], + }, + ), + ) + + assert summary.staging_keys_by_episode == { + GymCheckpointEpisode("rollout-1", 0): ("rollout-1/c1",), + GymCheckpointEpisode("rollout-2", 1): (cut_key, "rollout-2/c1"), + } + + +@pytest.mark.parametrize( + "staging_keys_by_episode", + [ + pytest.param(None, id="missing"), + pytest.param({"rollout-1": ["rollout-1/c1"]}, id="a-key-with-no-episode"), + pytest.param( + {"rollout-1": ["rollout-1/c1"], "rollout-2-a1": ["rollout-2/c1", "extra"]}, + id="a-key-outside-the-flat-list", + ), + ], +) +def test_staging_keys_gym_does_not_attribute_fail_the_commit( + monkeypatch, staging_keys_by_episode +) -> None: + """Keeping or dropping a row whose owner is unknown would both be a guess.""" + with pytest.raises(RuntimeError, match="staging_keys_by_episode"): + _commit_model_reply( + monkeypatch, + _commit_reply( + "deployment", + "model", + ["rollout-1", "rollout-2-a1"], + staging_keys=["rollout-1/c1", "rollout-2/c1"], + staging_keys_by_episode=staging_keys_by_episode, + ), + ) + + +def _commit_with_agent_records( + monkeypatch, + tmp_path, + *, + environment_keys: list[str], + agent_records: list[dict[str, object]], + tamper_agent_digest: bool = False, +): + """Commit one episode against real agent records written by Gym's store.""" + from nemo_gym._checkpoint import coordination + from nemo_gym._checkpoint.store import write_participant_state + from nemo_gym.episode_types import EpisodeId + + instance = GymCheckpointInstance("deployment", 0) + checkpoint_root = tmp_path / "step-1" + agent_manifest = write_participant_state( + instance.checkpoint_dir(checkpoint_root), + kind="agent", + instance="workplace-agent", + checkpoint_id="save-1", + records=agent_records, + ) + if tamper_agent_digest: + agent_manifest = {**agent_manifest, "records_sha256": "f" * 64} + agent_keys = sorted( + { + EpisodeId.model_validate(record["episode_id"]).capture_key + for record in agent_records + } + ) + client = object() + participants = SimpleNamespace( + client=client, + members=( + SimpleNamespace(server_name="environment", kind="environment"), + SimpleNamespace(server_name="model", kind="model"), + SimpleNamespace(server_name="agent", kind="agent"), + ), + ) + monkeypatch.setattr(coordination, "discover", AsyncMock(return_value=participants)) + monkeypatch.setattr( + coordination, + "commit", + AsyncMock( + return_value={ + "environment": _commit_reply( + "deployment", "environment", environment_keys + ), + "model": _commit_reply( + "deployment", + "model", + [], + staging_keys=[], + staging_keys_by_episode={}, + ), + "agent": { + "phase": "committed", + "manifest": agent_manifest, + "episode_ids": agent_keys, + }, + } + ), + ) + adapter = GymCheckpointAdapter( + instance=instance, client=client, auth_token="secret" + ) + + async def exercise(): + await adapter.discover() + return await adapter.commit( + "save-1", + checkpoint_root, + [GymCheckpointEpisode("rollout-1", 0)], + deadline_ts=10.0, + ) + + return asyncio.run(exercise()) + + +def _agent_record(*, legacy: bool) -> dict[str, object]: + return { + "episode_id": {"rollout_id": "rollout-1", "attempt": 0}, + "session_key": "session-1", + "session": {"messages": []}, + "boundary": None, + "episode": {"next_step": 2} if legacy else None, + } + + +def test_adapter_assigns_legacy_relay_episode_to_its_agent( + monkeypatch, tmp_path +) -> None: + """A legacy_agent environment relays /run, so the agent session owns it.""" + summary = _commit_with_agent_records( + monkeypatch, + tmp_path, + environment_keys=[], + agent_records=[_agent_record(legacy=True)], + ) + + assert summary.exported_episodes == (GymCheckpointEpisode("rollout-1", 0),) + assert summary.unowned_live_episodes == () + + +def test_adapter_keeps_environment_ownership_of_native_agent_session( + monkeypatch, tmp_path +) -> None: + summary = _commit_with_agent_records( + monkeypatch, + tmp_path, + environment_keys=["rollout-1"], + agent_records=[_agent_record(legacy=False)], + ) + + assert summary.exported_episodes == (GymCheckpointEpisode("rollout-1", 0),) + assert summary.unowned_live_episodes == () + + +def test_adapter_reports_native_agent_session_without_an_owner( + monkeypatch, tmp_path +) -> None: + """A parked session nobody can restore must not be mistaken for an export.""" + summary = _commit_with_agent_records( + monkeypatch, + tmp_path, + environment_keys=[], + agent_records=[_agent_record(legacy=False)], + ) + + assert summary.exported_episodes == () + assert summary.unowned_live_episodes == (GymCheckpointEpisode("rollout-1", 0),) + + +def test_adapter_rejects_agent_records_that_differ_from_the_commit_reply( + monkeypatch, tmp_path +) -> None: + with pytest.raises(RuntimeError, match="agent records on disk do not match"): + _commit_with_agent_records( + monkeypatch, + tmp_path, + environment_keys=[], + agent_records=[_agent_record(legacy=True)], + tamper_agent_digest=True, + ) + + +def test_adapter_rejects_restore_from_another_checkpoint( + monkeypatch, tmp_path: Path +) -> None: + from nemo_gym._checkpoint import coordination + + client = object() + participants = _participants(client, "actor-a") + monkeypatch.setattr( + coordination, + "discover", + AsyncMock(return_value=participants), + ) + monkeypatch.setattr( + coordination, + "restore", + AsyncMock( + return_value={ + "actor-a-environment": {"source_checkpoint_id": "save-other"}, + "actor-a-model": {"source_checkpoint_id": "save-other"}, + } + ), + ) + adapter = GymCheckpointAdapter( + instance=GymCheckpointInstance("actor-a", 0), + client=client, + auth_token="secret", + ) + + async def exercise() -> None: + await adapter.discover() + with pytest.raises(RuntimeError, match="wrong checkpoint"): + await adapter.restore( + "restore-1", + tmp_path, + [GymCheckpointEpisode("rollout-1", 0)], + participants=_committed_participants(tmp_path, "actor-a"), + source_checkpoint_id="save-1", + deadline_ts=10.0, + ) + + asyncio.run(exercise()) + + +def test_adapter_rejects_operations_before_discovery() -> None: + adapter = GymCheckpointAdapter( + instance=GymCheckpointInstance("actor-a", 0), + client=object(), + auth_token="secret", + ) + + with pytest.raises(RuntimeError, match="has not discovered"): + asyncio.run(adapter.prepare("save-1", deadline_ts=10.0)) + + +@pytest.mark.parametrize( + "operation", ["prepare", "renew", "retire", "commit", "resume"] +) +def test_a_failed_live_checkpoint_call_crosses_ray_typed( + monkeypatch, operation: str +) -> None: + """Gym's CoordinationError cannot be pickled and reaches RL untyped.""" + from nemo_gym._checkpoint import coordination + + failure = coordination.CoordinationError( + operation, {"actor-a-agent": "HTTP 409 deadline_exceeded"} + ) + monkeypatch.setattr(coordination, operation, AsyncMock(side_effect=failure)) + adapter = GymCheckpointAdapter( + instance=GymCheckpointInstance(shard_name="actor-a", replica_index=0), + client=object(), + auth_token="secret", + ) + adapter._participants = _participants(object(), "actor-a") + episodes = [GymCheckpointEpisode("rollout-1", 0)] + calls = { + "prepare": lambda: adapter.prepare("save-1", deadline_ts=1.0), + "renew": lambda: adapter.renew("save-1", deadline_ts=1.0), + "retire": lambda: adapter.retire("save-1", episodes, deadline_ts=1.0), + "commit": lambda: adapter.commit( + "save-1", "/unused", episodes, deadline_ts=1.0 + ), + "resume": lambda: adapter.resume("save-1", deadline_ts=1.0), + } + + with pytest.raises(GymCheckpointUnavailable, match="deadline_exceeded") as raised: + asyncio.run(calls[operation]()) + + restored = pickle.loads(pickle.dumps(raised.value)) + assert type(restored) is GymCheckpointUnavailable + assert str(restored) == str(raised.value) + + +class TestRestoreRequiresTheCommittedTopology: + """A participant missing on restore would have its state silently dropped.""" + + @staticmethod + def _adapter(monkeypatch, members: tuple[tuple[str, str], ...]): + from nemo_gym._checkpoint import coordination + + restore = AsyncMock( + return_value={ + name: {"source_checkpoint_id": "save-1"} for name, _ in members + } + ) + monkeypatch.setattr(coordination, "restore", restore) + adapter = GymCheckpointAdapter( + instance=GymCheckpointInstance("actor-a", 0), + client=object(), + auth_token="secret", + ) + adapter._participants = SimpleNamespace( + client=object(), + members=tuple( + SimpleNamespace(server_name=name, kind=kind) for name, kind in members + ), + ) + return adapter, restore + + @staticmethod + def _restore(adapter, checkpoint_root: Path, participants) -> None: + asyncio.run( + adapter.restore( + "restore-1", + checkpoint_root, + [GymCheckpointEpisode("rollout-1", 0)], + participants=participants, + source_checkpoint_id="save-1", + deadline_ts=10.0, + ) + ) + + def test_a_deployment_without_a_committed_participant_is_refused( + self, monkeypatch, tmp_path: Path + ) -> None: + recorded = _committed_participants(tmp_path, "actor-a") + adapter, restore = self._adapter( + monkeypatch, (("actor-a-environment", "environment"),) + ) + + with pytest.raises(RuntimeError, match=r"missing=\['actor-a-model'\]"): + self._restore(adapter, tmp_path, recorded) + restore.assert_not_awaited() + + def test_a_participant_the_checkpoint_never_saw_is_refused( + self, monkeypatch, tmp_path: Path + ) -> None: + recorded = _committed_participants(tmp_path, "actor-a") + adapter, restore = self._adapter( + monkeypatch, + ( + ("actor-a-environment", "environment"), + ("actor-a-model", "model"), + ("actor-a-resources", "resources"), + ), + ) + + with pytest.raises(RuntimeError, match=r"unexpected=\['actor-a-resources'\]"): + self._restore(adapter, tmp_path, recorded) + restore.assert_not_awaited() + + def test_a_participant_whose_kind_changed_is_refused( + self, monkeypatch, tmp_path: Path + ) -> None: + recorded = _committed_participants(tmp_path, "actor-a") + adapter, restore = self._adapter( + monkeypatch, + (("actor-a-environment", "environment"), ("actor-a-model", "agent")), + ) + + with pytest.raises(RuntimeError, match=r"changed_kind=\['actor-a-model'\]"): + self._restore(adapter, tmp_path, recorded) + restore.assert_not_awaited() + + def test_state_that_is_not_what_was_committed_is_refused( + self, monkeypatch, tmp_path: Path + ) -> None: + recorded = _committed_participants(tmp_path, "actor-a") + environment, model = recorded + swapped = ( + environment, + GymCheckpointParticipantRecord( + server_name=model.server_name, + kind=model.kind, + instance=model.instance, + records_sha256="f" * 64, + ), + ) + adapter, restore = self._adapter( + monkeypatch, + (("actor-a-environment", "environment"), ("actor-a-model", "model")), + ) + + with pytest.raises(RuntimeError, match="not the state that was committed"): + self._restore(adapter, tmp_path, swapped) + restore.assert_not_awaited() + + def test_missing_participant_state_is_refused( + self, monkeypatch, tmp_path: Path + ) -> None: + recorded = _committed_participants(tmp_path, "actor-a") + adapter, restore = self._adapter( + monkeypatch, + (("actor-a-environment", "environment"), ("actor-a-model", "model")), + ) + + with pytest.raises(RuntimeError, match="unreadable for participant"): + self._restore(adapter, tmp_path / "elsewhere", recorded) + restore.assert_not_awaited() + + def test_the_committed_deployment_restores( + self, monkeypatch, tmp_path: Path + ) -> None: + recorded = _committed_participants(tmp_path, "actor-a") + adapter, restore = self._adapter( + monkeypatch, + (("actor-a-model", "model"), ("actor-a-environment", "environment")), + ) + + self._restore(adapter, tmp_path, recorded) + + restore.assert_awaited_once() diff --git a/tests/unit/environments/test_gym_checkpoint_coordinator.py b/tests/unit/environments/test_gym_checkpoint_coordinator.py new file mode 100644 index 00000000000..e1d5018a342 --- /dev/null +++ b/tests/unit/environments/test_gym_checkpoint_coordinator.py @@ -0,0 +1,658 @@ +# 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. + +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from typing import Any + +import pytest +import ray.exceptions + +from nemo_rl.environments.gym_checkpoint_adapter import ( + GymCheckpointCommitSummary, + GymCheckpointEpisode, + GymCheckpointParticipantCommitSummary, + GymCheckpointParticipantManifest, + GymCheckpointParticipantRecord, + GymCheckpointPrepareSummary, + GymCheckpointUnavailable, +) +from nemo_rl.environments.gym_checkpoint_coordinator import ( + GYM_CHECKPOINT_SCHEMA_VERSION, + GymCheckpointCoordinator, + GymCheckpointManifest, + GymCheckpointNotReady, + GymCheckpointOperationError, + is_transient_gym_checkpoint_failure, + load_gym_checkpoint_manifest, +) +from nemo_rl.environments.nemo_gym import NemoGymShardSet + +pytestmark = pytest.mark.nemo_gym + + +class _RemoteMethod: + def __init__(self, function) -> None: + self._function = function + + def remote(self, *args, **kwargs): + return self._function(*args, **kwargs) + + +class _FakeGymActor: + def __init__( + self, + *, + prepared: bool = True, + restore_error: BaseException | None = None, + exported_episodes: tuple[GymCheckpointEpisode, ...] | None = None, + staging_keys_by_episode: dict[GymCheckpointEpisode, tuple[str, ...]] + | None = None, + hang: frozenset[str] = frozenset(), + participants: tuple[GymCheckpointParticipantCommitSummary, ...] = (), + ) -> None: + self.prepared = prepared + self.participants = participants + # Operations whose call is accepted but never answered, like an actor + # whose event loop is stuck. + self.hang = hang + self.restore_error = restore_error + self.exported_episodes = exported_episodes + self.staging_keys_by_episode = staging_keys_by_episode or {} + self.calls: list[tuple[str, tuple[Any, ...], dict[str, Any]]] = [] + self.checkpoint_prepare = _RemoteMethod(self._prepare) + self.checkpoint_renew = _RemoteMethod(self._renew) + self.checkpoint_retire = _RemoteMethod(self._retire) + self.checkpoint_commit = _RemoteMethod(self._commit) + self.checkpoint_restore = _RemoteMethod(self._restore) + self.checkpoint_resume = _RemoteMethod(self._resume) + + async def _record(self, operation: str, args, kwargs) -> None: + self.calls.append((operation, args, kwargs)) + if operation in self.hang: + await asyncio.Event().wait() + + async def _prepare(self, *args, **kwargs) -> GymCheckpointPrepareSummary: + await self._record("prepare", args, kwargs) + return GymCheckpointPrepareSummary( + prepared=self.prepared, + blockers={} if self.prepared else {"agent": ("rollout-1",)}, + ) + + async def _renew(self, *args, **kwargs) -> None: + await self._record("renew", args, kwargs) + + async def _commit(self, *args, **kwargs) -> GymCheckpointCommitSummary: + await self._record("commit", args, kwargs) + episodes = args[2] if self.exported_episodes is None else self.exported_episodes + return GymCheckpointCommitSummary( + exported_episodes=episodes, + staging_keys=tuple( + sorted( + {k for keys in self.staging_keys_by_episode.values() for k in keys} + ) + ), + participants=self.participants, + staging_keys_by_episode=self.staging_keys_by_episode, + ) + + async def _retire(self, *args, **kwargs) -> None: + await self._record("retire", args, kwargs) + + async def _restore(self, *args, **kwargs) -> None: + await self._record("restore", args, kwargs) + if self.restore_error is not None: + raise self.restore_error + + async def _resume(self, *args, **kwargs) -> None: + await self._record("resume", args, kwargs) + + +def _coordinator(*actors: _FakeGymActor) -> GymCheckpointCoordinator: + return GymCheckpointCoordinator( + NemoGymShardSet(handles={"tools": list(actors)}), + control_timeout_s=30.0, + ) + + +def test_manifest_rejects_duplicate_episode_identity() -> None: + episode = GymCheckpointEpisode("rollout-0", 2) + + with pytest.raises(ValueError, match="duplicate episodes"): + GymCheckpointManifest( + schema_version=GYM_CHECKPOINT_SCHEMA_VERSION, + checkpoint_id="save-1", + instances={"tools/replica-0": (episode, episode)}, + staging_keys={"tools/replica-0": ()}, + participants={"tools/replica-0": ()}, + ) + + +def test_coordinator_commits_and_restores_each_instance_inventory( + tmp_path: Path, +) -> None: + first = _FakeGymActor( + staging_keys_by_episode={ + GymCheckpointEpisode("rollout-0", 0): ("rollout-0/call-0",) + } + ) + second = _FakeGymActor() + coordinator = _coordinator(first, second) + inventory = { + "tools/replica-0": (GymCheckpointEpisode("rollout-0", 0),), + "tools/replica-1": (GymCheckpointEpisode("rollout-1", 2),), + } + + async def exercise() -> None: + async with coordinator.prepared("save-1"): + result = await coordinator.commit("save-1", tmp_path, inventory) + assert result.manifest.instances == inventory + assert set(result.instances) == set(inventory) + assert result.instances["tools/replica-0"].exported_episodes == ( + GymCheckpointEpisode("rollout-0", 0), + ) + assert result.manifest.staging_keys == { + "tools/replica-0": ("rollout-0/call-0",), + "tools/replica-1": (), + } + await coordinator.renew("save-1") + manifest = load_gym_checkpoint_manifest(tmp_path) + assert manifest.checkpoint_id == "save-1" + assert manifest.instances == inventory + assert manifest.staging_keys == { + "tools/replica-0": ("rollout-0/call-0",), + "tools/replica-1": (), + } + await coordinator.restore("restore-1", tmp_path, manifest) + await coordinator.retire_restored("restore-1", manifest) + await coordinator.resume("restore-1") + + asyncio.run(exercise()) + + for replica, actor in enumerate((first, second)): + operations = [call[0] for call in actor.calls] + assert operations == [ + "prepare", + "renew", + "commit", + "renew", + "resume", + "restore", + "retire", + "resume", + ] + commit_args = actor.calls[2][1] + assert commit_args[0] == "save-1" + assert commit_args[2] == inventory[f"tools/replica-{replica}"] + retire_args = actor.calls[-2][1] + assert retire_args[0] == "restore-1" + assert retire_args[1] == tuple( + episode.next_attempt() for episode in inventory[f"tools/replica-{replica}"] + ) + assert actor.calls[5][2]["source_checkpoint_id"] == "save-1" + + +def test_coordinator_resumes_every_actor_when_prepare_is_not_ready() -> None: + ready = _FakeGymActor() + blocked = _FakeGymActor(prepared=False) + coordinator = _coordinator(ready, blocked) + + async def exercise() -> None: + with pytest.raises(GymCheckpointNotReady, match="rollout-1"): + async with coordinator.prepared("save-1"): + raise AssertionError("unreachable") + + asyncio.run(exercise()) + + assert [call[0] for call in ready.calls] == ["prepare", "resume"] + assert [call[0] for call in blocked.calls] == ["prepare", "resume"] + + +def test_coordinator_renews_lease_while_outer_checkpoint_is_slow() -> None: + actor = _FakeGymActor() + coordinator = GymCheckpointCoordinator( + NemoGymShardSet(handles={"tools": [actor]}), + control_timeout_s=0.03, + ) + + async def exercise() -> None: + async with coordinator.prepared("save-1"): + await asyncio.sleep(0.04) + + asyncio.run(exercise()) + + operations = [call[0] for call in actor.calls] + assert operations[0:2] == ["prepare", "renew"] + assert operations.count("renew") >= 2 + assert operations[-1] == "resume" + + +def test_coordinator_rejects_inventory_for_the_wrong_actor_topology( + tmp_path: Path, +) -> None: + coordinator = _coordinator(_FakeGymActor()) + + with pytest.raises(ValueError, match="live actor topology"): + asyncio.run( + coordinator.commit( + "save-1", + tmp_path, + {"other/replica-0": ()}, + ) + ) + + +def test_coordinator_uses_environment_exported_subset_as_manifest( + tmp_path: Path, +) -> None: + coordinator = _coordinator(_FakeGymActor(exported_episodes=())) + + result = asyncio.run( + coordinator.commit( + "save-1", + tmp_path, + {"tools/replica-0": (GymCheckpointEpisode("rollout-on-wire", 0),)}, + ) + ) + + assert result.manifest.instances == {"tools/replica-0": ()} + assert load_gym_checkpoint_manifest(tmp_path).instances == {"tools/replica-0": ()} + + +def test_manifest_keeps_only_the_staging_keys_of_episodes_gym_saved( + tmp_path: Path, +) -> None: + """A candidate whose reply is on the wire is not saved, so neither are its rows. + + Gym's model participant reports the staged rows of every candidate in scope. + An on-wire candidate seals and its group can clean those rows up before the + cut, so keeping its keys would fail the checkpoint, or a later restore, + over rows nothing needs. Gym says whose each key is: a cut key names no + rollout, and a row carried over from an earlier attempt serves both. + """ + saved = (GymCheckpointEpisode("rollout-0", 0), GymCheckpointEpisode("rollout-1", 2)) + on_wire = GymCheckpointEpisode("rollout-on-wire", 0) + coordinator = _coordinator( + _FakeGymActor( + exported_episodes=saved, + staging_keys_by_episode={ + saved[0]: ("rollout-0/call-0",), + saved[1]: ("__generation_cut__/save-1/call-1", "carried/call-0"), + on_wire: ("carried/call-0", "rollout-on-wire/call-0"), + }, + ) + ) + + result = asyncio.run( + coordinator.commit("save-1", tmp_path, {"tools/replica-0": (*saved, on_wire)}) + ) + + kept = { + "tools/replica-0": ( + "__generation_cut__/save-1/call-1", + "carried/call-0", + "rollout-0/call-0", + ) + } + assert result.manifest.staging_keys == kept + assert load_gym_checkpoint_manifest(tmp_path).staging_keys == kept + + +def test_a_summary_must_attribute_every_staging_key() -> None: + """A key no episode owns could be neither kept nor dropped safely.""" + with pytest.raises(ValueError, match="staging_keys_by_episode"): + GymCheckpointCommitSummary( + exported_episodes=(), + staging_keys=("orphan/call-0",), + participants=(), + ) + + +def test_coordinator_rejects_episode_outside_candidate_set(tmp_path: Path) -> None: + coordinator = _coordinator( + _FakeGymActor( + exported_episodes=(GymCheckpointEpisode("unexpected-rollout", 0),) + ) + ) + + with pytest.raises(RuntimeError, match="outside the candidate set"): + asyncio.run( + coordinator.commit( + "save-1", + tmp_path, + {"tools/replica-0": (GymCheckpointEpisode("rollout-on-wire", 0),)}, + ) + ) + + +def test_restore_failure_discards_every_actor_under_uncertain_outcome() -> None: + restored = _FakeGymActor() + failed = _FakeGymActor(restore_error=RuntimeError("restore failed")) + coordinator = _coordinator(restored, failed) + manifest = GymCheckpointManifest( + schema_version=GYM_CHECKPOINT_SCHEMA_VERSION, + checkpoint_id="save-1", + instances={ + "tools/replica-0": (GymCheckpointEpisode("rollout-0", 0),), + "tools/replica-1": (GymCheckpointEpisode("rollout-1", 0),), + }, + staging_keys={ + "tools/replica-0": (), + "tools/replica-1": (), + }, + participants={ + "tools/replica-0": (), + "tools/replica-1": (), + }, + ) + + with pytest.raises(GymCheckpointOperationError, match="restore failed"): + asyncio.run(coordinator.restore("restore-1", "/unused", manifest)) + + assert [call[0] for call in restored.calls] == ["restore", "retire", "resume"] + assert [call[0] for call in failed.calls] == ["restore", "retire", "resume"] + + +# Bound on how long a test waits for the coordinator before calling it hung. +_HANG_GUARD_S = 5.0 + + +def _run_guarded(exercise) -> None: + async def guarded() -> None: + try: + await asyncio.wait_for(exercise(), timeout=_HANG_GUARD_S) + except TimeoutError: + pytest.fail( + f"Gym checkpoint control call still waiting after {_HANG_GUARD_S}s" + ) + + asyncio.run(guarded()) + + +def _fast_coordinator(*actors: _FakeGymActor) -> GymCheckpointCoordinator: + return GymCheckpointCoordinator( + NemoGymShardSet(handles={"tools": list(actors)}), + control_timeout_s=0.05, + ) + + +def test_an_actor_that_never_answers_prepare_fails_the_checkpoint() -> None: + """Gym only enforces deadline_ts once it runs the call; a stuck actor never does.""" + healthy = _FakeGymActor() + stuck = _FakeGymActor(hang=frozenset({"prepare"})) + coordinator = _fast_coordinator(healthy, stuck) + + async def exercise() -> None: + with pytest.raises(GymCheckpointOperationError) as raised: + async with coordinator.prepared("save-1"): + raise AssertionError("unreachable") + assert set(raised.value.failures) == {"tools/replica-1"} + assert isinstance(raised.value.failures["tools/replica-1"], TimeoutError) + + _run_guarded(exercise) + + # Admission was closed on the healthy actor, so cleanup must still reopen it. + assert [call[0] for call in healthy.calls] == ["prepare", "resume"] + assert [call[0] for call in stuck.calls] == ["prepare", "resume"] + + +def test_a_renewal_that_never_answers_does_not_hang_the_checkpoint() -> None: + actor = _FakeGymActor(hang=frozenset({"renew"})) + coordinator = _fast_coordinator(actor) + + async def exercise() -> None: + with pytest.raises(GymCheckpointOperationError) as raised: + async with coordinator.prepared("save-1"): + await asyncio.sleep(0.2) + assert raised.value.operation == "renew" + assert isinstance(raised.value.failures["tools/replica-0"], TimeoutError) + + _run_guarded(exercise) + + assert [call[0] for call in actor.calls][-1] == "resume" + + +def test_a_resume_that_never_answers_keeps_the_original_failure() -> None: + """Cleanup is best effort: Gym's lease reopens admission if resume is lost.""" + actor = _FakeGymActor(prepared=False, hang=frozenset({"resume"})) + coordinator = _fast_coordinator(actor) + + async def exercise() -> None: + with pytest.raises(GymCheckpointNotReady) as raised: + async with coordinator.prepared("save-1"): + raise AssertionError("unreachable") + assert any("cleanup also failed" in note for note in raised.value.__notes__) + + _run_guarded(exercise) + + assert [call[0] for call in actor.calls] == ["prepare", "resume"] + + +@pytest.mark.parametrize("operation", ["commit", "restore", "retire"]) +def test_every_control_operation_is_bounded(operation: str, tmp_path: Path) -> None: + actor = _FakeGymActor(hang=frozenset({operation})) + coordinator = _fast_coordinator(actor) + episodes = {"tools/replica-0": (GymCheckpointEpisode("rollout-0", 0),)} + manifest = GymCheckpointManifest( + schema_version=GYM_CHECKPOINT_SCHEMA_VERSION, + checkpoint_id="save-1", + instances=episodes, + staging_keys={"tools/replica-0": ()}, + participants={"tools/replica-0": ()}, + ) + calls = { + "commit": lambda: coordinator.commit("save-1", tmp_path, episodes), + "restore": lambda: coordinator.restore("restore-1", tmp_path, manifest), + "retire": lambda: coordinator.retire("save-1", episodes), + } + + async def exercise() -> None: + with pytest.raises(GymCheckpointOperationError) as raised: + await calls[operation]() + assert any( + isinstance(failure, TimeoutError) + for failure in raised.value.failures.values() + ) + + _run_guarded(exercise) + + +@pytest.mark.parametrize( + "error", + [ + GymCheckpointNotReady({"tools/replica-0": {"agent": ("rollout-1",)}}), + GymCheckpointOperationError("prepare", {"tools/replica-0": TimeoutError()}), + GymCheckpointOperationError( + "commit", + {"tools/replica-0": GymCheckpointUnavailable("HTTP 500 participant")}, + ), + GymCheckpointOperationError( + "renew", {"tools/replica-0": ray.exceptions.RayActorError()} + ), + ], + ids=["not-ready", "unanswered", "control-plane", "dead-actor"], +) +def test_failures_that_leave_nothing_behind_are_transient(error) -> None: + assert is_transient_gym_checkpoint_failure(error) + + +@pytest.mark.parametrize( + "error", + [ + RuntimeError("Gym episode ownership changed after commit"), + ValueError("inventory does not match the live actor topology"), + # An actor-side validation failure, such as an unexpected participant set. + GymCheckpointOperationError( + "commit", {"tools/replica-0": RuntimeError("unexpected participant set")} + ), + # One broken instance is not excused by another that merely timed out. + GymCheckpointOperationError( + "commit", + { + "tools/replica-0": TimeoutError(), + "tools/replica-1": TypeError("invalid commit summary"), + }, + ), + ], + ids=["controller-invariant", "topology", "actor-invariant", "mixed"], +) +def test_broken_invariants_are_not_transient(error) -> None: + assert not is_transient_gym_checkpoint_failure(error) + + +def _committed_participant( + server_name: str, kind: str, digest: str +) -> GymCheckpointParticipantCommitSummary: + return GymCheckpointParticipantCommitSummary( + server_name=server_name, + kind=kind, + phase="committed", + capture_keys=(), + manifest=GymCheckpointParticipantManifest( + schema_version=1, + kind=kind, + instance=f"{server_name}-state", + checkpoint_id="save-1", + records_file="records.jsonl", + records_sha256=digest, + record_count=0, + ), + staging_keys=(), + ) + + +def test_commit_records_each_instance_participant_topology(tmp_path: Path) -> None: + """Restore needs it: a participant missing later must not be silently skipped.""" + first = _FakeGymActor( + participants=( + _committed_participant("tools-environment", "environment", "a" * 64), + _committed_participant("tools-resources", "resources", "b" * 64), + ) + ) + second = _FakeGymActor( + participants=(_committed_participant("tools-env", "environment", "c" * 64),) + ) + coordinator = _coordinator(first, second) + inventory = {"tools/replica-0": (), "tools/replica-1": ()} + + result = asyncio.run(coordinator.commit("save-1", tmp_path, inventory)) + + expected = { + "tools/replica-0": ( + GymCheckpointParticipantRecord( + "tools-environment", "environment", "tools-environment-state", "a" * 64 + ), + GymCheckpointParticipantRecord( + "tools-resources", "resources", "tools-resources-state", "b" * 64 + ), + ), + "tools/replica-1": ( + GymCheckpointParticipantRecord( + "tools-env", "environment", "tools-env-state", "c" * 64 + ), + ), + } + assert result.manifest.participants == expected + assert load_gym_checkpoint_manifest(tmp_path).participants == expected + + +def test_restore_hands_each_instance_its_recorded_participants() -> None: + first = _FakeGymActor() + second = _FakeGymActor() + coordinator = _coordinator(first, second) + recorded = { + "tools/replica-0": ( + GymCheckpointParticipantRecord("a-env", "environment", "a-env", "a" * 64), + ), + "tools/replica-1": ( + GymCheckpointParticipantRecord("b-env", "environment", "b-env", "b" * 64), + ), + } + manifest = GymCheckpointManifest( + schema_version=GYM_CHECKPOINT_SCHEMA_VERSION, + checkpoint_id="save-1", + instances={"tools/replica-0": (), "tools/replica-1": ()}, + staging_keys={"tools/replica-0": (), "tools/replica-1": ()}, + participants=recorded, + ) + + asyncio.run(coordinator.restore("restore-1", "/unused", manifest)) + + for actor, instance_id in ((first, "tools/replica-0"), (second, "tools/replica-1")): + (restore_call,) = [call for call in actor.calls if call[0] == "restore"] + assert restore_call[2]["participants"] == recorded[instance_id] + + +def test_a_manifest_without_participant_topology_is_rejected(tmp_path: Path) -> None: + manifest = GymCheckpointManifest( + schema_version=GYM_CHECKPOINT_SCHEMA_VERSION, + checkpoint_id="save-1", + instances={"tools/replica-0": ()}, + staging_keys={"tools/replica-0": ()}, + participants={"tools/replica-0": ()}, + ) + raw = manifest.to_dict() + del raw["participants"] + (tmp_path / "gym_checkpoint.json").write_text(json.dumps(raw)) + + with pytest.raises(ValueError, match="participants must be a mapping"): + load_gym_checkpoint_manifest(tmp_path) + + +@pytest.mark.parametrize( + ("participants", "match"), + [ + ({}, "participant topology must match"), + ( + { + "tools/replica-0": ( + GymCheckpointParticipantRecord("a", "environment", "a", "a" * 64), + GymCheckpointParticipantRecord("a", "model", "a2", "b" * 64), + ) + }, + "duplicate participants", + ), + ( + { + "tools/replica-0": ( + GymCheckpointParticipantRecord("a", "judge", "a", "a" * 64), + ) + }, + "unknown kind", + ), + ( + { + "tools/replica-0": ( + GymCheckpointParticipantRecord("a", "environment", "a", "nope"), + ) + }, + "invalid records_sha256", + ), + ], + ids=["missing-instance", "duplicate", "unknown-kind", "bad-digest"], +) +def test_manifest_rejects_a_malformed_participant_topology( + participants, match: str +) -> None: + with pytest.raises(ValueError, match=match): + GymCheckpointManifest( + schema_version=GYM_CHECKPOINT_SCHEMA_VERSION, + checkpoint_id="save-1", + instances={"tools/replica-0": ()}, + staging_keys={"tools/replica-0": ()}, + participants=participants, + ) diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index 37699264d2d..d34901182ae 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -39,6 +39,7 @@ from nemo_rl.distributed.ray_actor_environment_registry import ( get_actor_python_env, ) +from nemo_rl.environments.gym_checkpoint_adapter import GymCheckpointInstance from nemo_rl.environments.nemo_gym import ( NemoGym, NemoGymConfig, @@ -93,8 +94,10 @@ async def _run() -> None: ] class _RolloutCollectionHelper: - def run_examples(self, examples, head_server_config): - del head_server_config + def run_examples( + self, examples, head_server_config, route_failures_to_sidecar=False + ): + del head_server_config, route_failures_to_sidecar for row in examples: row["agent_ref"] = {"name": "resolved_agent"} @@ -1161,6 +1164,8 @@ def nemo_gym(nemo_gym_vllm_generation, nemo_gym_tokenizer): # noqa: F811 model_name=nemo_gym_vllm_generation.cfg["model_name"], base_urls=nemo_gym_vllm_generation.dp_openai_server_base_urls, initial_global_config_dict=safe_load(yaml_str), + checkpoint_instance=GymCheckpointInstance("nemo_gym", 0), + turn_checkpointing_enabled=False, ) env = NemoGym.options( runtime_env={ @@ -1650,8 +1655,10 @@ async def _run(): postprocess_calls = [] class _RolloutCollectionHelper: - def run_examples(self, examples, head_server_config): - del head_server_config + def run_examples( + self, examples, head_server_config, route_failures_to_sidecar=False + ): + del head_server_config, route_failures_to_sidecar content = examples[0]["responses_create_params"]["input"][0]["content"] assert content[0]["video_url"].startswith("data:video/mp4;base64,") assert content[1]["image_url"].startswith("data:image/png;base64,") @@ -1747,7 +1754,10 @@ def batch_decode(self, batches): return [" ".join(map(str, token_ids)) for token_ids in batches] class _RolloutCollectionHelper: - def run_examples(self, examples, head_server_config): + def run_examples( + self, examples, head_server_config, route_failures_to_sidecar=False + ): + del route_failures_to_sidecar assert head_server_config.backend == "megatron" dispatched_row = examples[0] dispatched_part = dispatched_row["responses_create_params"]["input"][0][ diff --git a/tests/unit/environments/test_nemo_gym_health.py b/tests/unit/environments/test_nemo_gym_health.py index c35e125078a..b02fc212eb7 100644 --- a/tests/unit/environments/test_nemo_gym_health.py +++ b/tests/unit/environments/test_nemo_gym_health.py @@ -30,6 +30,7 @@ import pytest +from nemo_rl.environments.gym_checkpoint_adapter import GymCheckpointInstance from nemo_rl.environments.nemo_gym import NemoGym # NemoGym is a Ray actor; grab the plain class so these run without a cluster. @@ -39,7 +40,13 @@ def _unspun() -> NemoGymClass: """A NemoGym exactly as Ray would recreate it after a restart.""" return NemoGymClass( - {"model_name": "m", "base_urls": [], "initial_global_config_dict": {}} + { + "model_name": "m", + "base_urls": [], + "initial_global_config_dict": {}, + "checkpoint_instance": GymCheckpointInstance("nemo_gym", 0), + "turn_checkpointing_enabled": False, + } ) @@ -62,9 +69,14 @@ class _TaskSourceResolvingRolloutHelper: """Mimic Gym's synchronous task_source-to-agent_ref resolution.""" def run_examples( - self, examples: list[dict[str, Any]], head_server_config: str + self, + examples: list[dict[str, Any]], + head_server_config: str, + route_failures_to_sidecar: bool = False, ) -> list[Any]: assert head_server_config == "head-server" + # A failed /run must come back as its row so a refusal can name it. + assert route_failures_to_sidecar is True assert all("agent_ref" not in example for example in examples) for example in examples: example["agent_ref"] = { @@ -76,9 +88,16 @@ def run_examples( class _TaskSourceResolvingRolloutHelperWithResult(_TaskSourceResolvingRolloutHelper): def run_examples( - self, examples: list[dict[str, Any]], head_server_config: str + self, + examples: list[dict[str, Any]], + head_server_config: str, + route_failures_to_sidecar: bool = False, ) -> list[Any]: - super().run_examples(examples, head_server_config) + super().run_examples( + examples, + head_server_config, + route_failures_to_sidecar=route_failures_to_sidecar, + ) async def completed(example): return example, {} diff --git a/tests/unit/environments/test_nemo_gym_rollout.py b/tests/unit/environments/test_nemo_gym_rollout.py index 7610f531291..bd394cc1575 100644 --- a/tests/unit/environments/test_nemo_gym_rollout.py +++ b/tests/unit/environments/test_nemo_gym_rollout.py @@ -539,6 +539,7 @@ def l0_nemo_gym(scripted_openai_base_url): tokenizer=tokenizer, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) try: yield env diff --git a/tests/unit/environments/test_nemo_gym_utils.py b/tests/unit/environments/test_nemo_gym_utils.py index 07b3a61a7d5..b6eeca057ea 100644 --- a/tests/unit/environments/test_nemo_gym_utils.py +++ b/tests/unit/environments/test_nemo_gym_utils.py @@ -171,6 +171,7 @@ def test_build_nemo_gym_config_splits_nemo_rl_keys(detected_uv_dirs): model_name="test-model", enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) assert cfg["model_name"] == "test-model" @@ -207,6 +208,7 @@ def test_build_nemo_gym_config_uv_dirs(detected_uv_dirs, configured, expected): model_name="test-model", enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) global_config = cfg["initial_global_config_dict"] assert (global_config["uv_cache_dir"], global_config["uv_venv_dir"]) == expected @@ -219,6 +221,7 @@ def test_build_nemo_gym_config_moves_port_range_to_actor_fields(detected_uv_dirs model_name="test-model", enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) assert (cfg["port_range_low"], cfg["port_range_high"]) == (6000, 6999) @@ -233,6 +236,7 @@ def test_build_nemo_gym_config_router_replay_off_uses_default_dtype(detected_uv_ model_name="test-model", enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) assert cfg["require_routed_experts"] is False assert cfg["routed_experts_dtype"] == "int16" @@ -250,6 +254,7 @@ def test_build_nemo_gym_config_router_replay_resolves_dtype(detected_uv_dirs): model_name="test-model", enable_router_replay=True, use_fastokens=False, + turn_checkpointing_enabled=False, ) mock_resolve.assert_called_once_with("test-model") @@ -285,6 +290,7 @@ def test_an_unsharded_job_gets_the_registry_runtime_env( tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=True, + turn_checkpointing_enabled=False, token_capture=token_capture, ) @@ -307,7 +313,12 @@ def test_an_unsharded_job_gets_the_registry_runtime_env( assert cfg["use_fastokens"] is True # The ledger config must ride through to the actor; a refactor of this # wrapper once dropped it without any type or test catching it. - assert cfg["token_capture"] == token_capture + assert cfg["token_capture"] == { + **token_capture, + "capture_dir": "/tmp/cap/gym-instances/nemo_gym/replica-0", + } + assert token_capture == {"enabled": True, "capture_dir": "/tmp/cap"} + assert cfg["checkpoint_instance"].instance_id == "nemo_gym/replica-0" # Spinup is deferred from __init__, so the factory must await it. The # tokenizer install has to follow it, not race it. @@ -349,6 +360,7 @@ def get_or_fail(ref, **_kwargs): tokenizer=MagicMock(), enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) shutdown.assert_called_once() @@ -394,6 +406,7 @@ def get_or_fail(ref, **_kwargs): tokenizer=MagicMock(), enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) assert exc_info.value is startup_error @@ -573,6 +586,7 @@ def __init__( spinup_timeouts=None, wedged_spinups=None, tokenizer_timeouts=None, + checkpoint_failures=None, pg_ready_error=None, ): self.entries_by_index = entries_by_index or {} @@ -580,6 +594,7 @@ def __init__( self.spinup_timeouts = set(spinup_timeouts or ()) self.wedged_spinups = set(wedged_spinups or ()) self.tokenizer_timeouts = set(tokenizer_timeouts or ()) + self.checkpoint_failures = checkpoint_failures or {} self.timed_out_spinups = set() self.pg_ready_error = pg_ready_error self.actors = [] @@ -615,6 +630,10 @@ def remote(config): actor = MagicMock(name=f"gym-actor-{index}") actor._spinup.remote.return_value = ("spinup", index) actor.set_tokenizer.remote.return_value = ("tokenizer", index) + actor.initialize_checkpoint_adapter.remote.return_value = ( + "checkpoint", + index, + ) actor.list_entries.remote.return_value = ("entries", index) self.actors.append(actor) self.actor_options.append(option_kwargs) @@ -652,6 +671,10 @@ def ray_get(self, reference, timeout=None): if index in self.tokenizer_timeouts and timeout is not None: raise TimeoutError("startup budget expired") return None + if kind == "checkpoint": + if index in self.checkpoint_failures: + raise self.checkpoint_failures[index] + return None return self.entries_by_index.get(index, {}) @@ -708,6 +731,7 @@ def test_build_nemo_gym_actors_unsharded_makes_exactly_one_actor(detected_uv_dir tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) assert not shard_set.is_sharded @@ -720,6 +744,83 @@ def test_build_nemo_gym_actors_unsharded_makes_exactly_one_actor(detected_uv_dir ) +def test_unsharded_checkpointing_discovers_the_actor_at_startup(detected_uv_dirs): + cluster = _FakeGymCluster() + token_capture = { + "enabled": True, + "capture_dir": "/capture", + "control_auth_token": "secret", + } + + with _patched_cluster(cluster): + shard_set = nemo_gym_mod.build_nemo_gym_actors( + _env_configs(), + base_urls=["http://vllm-0"], + model_name="test-model", + tokenizer=_TOKENIZER, + enable_router_replay=False, + use_fastokens=False, + turn_checkpointing_enabled=True, + token_capture=token_capture, + ) + + actor = shard_set.sole_handle() + actor.initialize_checkpoint_adapter.remote.assert_called_once_with() + config = cluster.actor_configs[0] + assert config["checkpoint_instance"].instance_id == "nemo_gym/replica-0" + assert config["token_capture"]["capture_dir"] == ( + "/capture/gym-instances/nemo_gym/replica-0" + ) + assert config["initial_global_config_dict"]["checkpoint"] == { + "enabled": True, + "control_auth_token": "secret", + } + assert ( + config["initial_global_config_dict"]["policy_model"]["responses_api_models"][ + "vllm_model" + ]["checkpoint_policy"] + is True + ) + assert token_capture["capture_dir"] == "/capture" + + +def test_prefix_checkpointing_enables_generation_cuts_on_policy_model( + detected_uv_dirs, +): + config = build_nemo_gym_config( + _env_configs(), + base_urls=["http://vllm-0"], + model_name="test-model", + enable_router_replay=False, + use_fastokens=False, + turn_checkpointing_enabled=True, + token_capture={ + "enabled": True, + "capture_dir": "/capture", + "control_auth_token": "secret", + "generation_prefix_cuts_enabled": True, + }, + ) + + model = config["initial_global_config_dict"]["policy_model"][ + "responses_api_models" + ]["vllm_model"] + assert model["checkpoint_policy"] is True + assert model["checkpoint_generation_cuts"] is True + + +def test_checkpointing_requires_token_capture(detected_uv_dirs): + with pytest.raises(ValueError, match="requires token capture"): + build_nemo_gym_config( + _env_configs(), + base_urls=["http://vllm-0"], + model_name="test-model", + enable_router_replay=False, + use_fastokens=False, + turn_checkpointing_enabled=True, + ) + + def test_build_nemo_gym_actors_spreads_every_replica_onto_its_own_node( detected_uv_dirs, ): @@ -738,6 +839,7 @@ def test_build_nemo_gym_actors_spreads_every_replica_onto_its_own_node( tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) # Two shards, one with replicas: 2, so three actors on three nodes. @@ -771,6 +873,7 @@ def test_shards_get_their_own_config_paths_and_log_directories(detected_uv_dirs) tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) gym_configs = [ @@ -794,6 +897,72 @@ def test_shards_get_their_own_config_paths_and_log_directories(detected_uv_dirs) assert "allowed_duplicate_entries" not in config +def test_shards_get_stable_isolated_checkpoint_namespaces(detected_uv_dirs): + cluster = _FakeGymCluster() + token_capture = { + "enabled": True, + "capture_dir": "/capture", + "control_auth_token": "secret", + } + + with _patched_cluster(cluster): + nemo_gym_mod.build_nemo_gym_actors( + _shard_env_configs(), + base_urls=["http://vllm-0"], + model_name="test-model", + tokenizer=_TOKENIZER, + enable_router_replay=False, + use_fastokens=False, + turn_checkpointing_enabled=True, + token_capture=token_capture, + ) + + assert [ + config["checkpoint_instance"].instance_id for config in cluster.actor_configs + ] == ["judged/replica-0", "tools/replica-0", "tools/replica-1"] + assert [ + config["token_capture"]["capture_dir"] for config in cluster.actor_configs + ] == [ + "/capture/gym-instances/judged/replica-0", + "/capture/gym-instances/tools/replica-0", + "/capture/gym-instances/tools/replica-1", + ] + for actor in cluster.actors: + actor.initialize_checkpoint_adapter.remote.assert_called_once_with() + + +def test_checkpoint_discovery_failure_tears_down_every_shard(detected_uv_dirs): + cluster = _FakeGymCluster( + checkpoint_failures={1: RuntimeError("participant unavailable")} + ) + token_capture = { + "enabled": True, + "capture_dir": "/capture", + "control_auth_token": "secret", + } + + with ( + _patched_cluster(cluster), + pytest.raises( + nemo_gym_mod.ShardSetupError, + match="tools.*replica 0.*participant unavailable", + ), + ): + nemo_gym_mod.build_nemo_gym_actors( + _shard_env_configs(), + base_urls=["http://vllm-0"], + model_name="test-model", + tokenizer=_TOKENIZER, + enable_router_replay=False, + use_fastokens=False, + turn_checkpointing_enabled=True, + token_capture=token_capture, + ) + + cluster.shutdown_environments.assert_called_once() + assert len(cluster.removed_placement_groups) == 1 + + def test_every_replica_gets_the_tokenizer_installed(detected_uv_dirs): """A replica missing the tokenizer fails on its first rollout, not at setup. @@ -810,6 +979,7 @@ def test_every_replica_gets_the_tokenizer_installed(detected_uv_dirs): tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) assert len(cluster.actors) == 3 @@ -834,6 +1004,7 @@ def test_actor_cpus_override_sizes_that_shards_bundle(detected_uv_dirs): tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) (pg_kwargs,) = cluster.placement_group_calls @@ -905,6 +1076,7 @@ def test_a_shard_that_fails_to_start_names_itself_and_tears_everything_down( tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) # Gym names the offending entry; we add the shard it belongs to. @@ -928,6 +1100,7 @@ def test_unsharded_startup_failure_tears_down_the_actor(detected_uv_dirs): tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) cluster.shutdown_environments.assert_called_once() @@ -947,6 +1120,7 @@ def test_timed_out_shard_startup_is_drained_before_teardown(detected_uv_dirs): tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) spinup_zero_calls = [ @@ -975,6 +1149,7 @@ def test_wedged_shard_startup_cannot_block_forced_teardown(detected_uv_dirs): tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) cluster.shutdown_environments.assert_called_once() @@ -1003,6 +1178,7 @@ def test_tokenizer_timeout_reports_exhausted_shared_startup_budget( tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, spinup_timeout=1.0, ) @@ -1021,6 +1197,7 @@ def test_unplaceable_bundles_fail_fast_and_release_the_group(detected_uv_dirs): tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) assert cluster.actors == [] @@ -1044,6 +1221,7 @@ def test_duplicate_agent_across_shards_tears_the_set_down(detected_uv_dirs): tokenizer=_TOKENIZER, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, ) cluster.shutdown_environments.assert_called_once() @@ -1080,6 +1258,24 @@ def test_a_replicated_shard_labels_each_instance_apart(): shard_set.instance_label(MagicMock()) +def test_checkpoint_instance_identity_round_trips_to_the_same_replica(): + first, second = MagicMock(), MagicMock() + shard_set = nemo_gym_mod.NemoGymShardSet(handles={"tools": [first, second]}) + + assert shard_set.checkpoint_instance_for_handle(first).instance_id == ( + "tools/replica-0" + ) + assert shard_set.checkpoint_instance_for_handle(second).instance_id == ( + "tools/replica-1" + ) + assert shard_set.handle_for_checkpoint_instance("tools/replica-0") is first + assert shard_set.handle_for_checkpoint_instance("tools/replica-1") is second + assert shard_set.checkpoint_handles == { + "tools/replica-0": first, + "tools/replica-1": second, + } + + def test_an_unreplicated_shard_is_labelled_by_its_name_alone(): """Names a metric and its log directory the same way.""" handle = MagicMock() diff --git a/tests/unit/experience/test_failures.py b/tests/unit/experience/test_failures.py index 5a111bd5d2f..c69742f8e4e 100644 --- a/tests/unit/experience/test_failures.py +++ b/tests/unit/experience/test_failures.py @@ -20,6 +20,7 @@ """ import asyncio +import sys import aiohttp import pytest @@ -28,10 +29,11 @@ from ray import cloudpickle as ray_cloudpickle from yarl import URL -from nemo_rl.environments.nemo_gym import _typed_gym_failure +from nemo_rl.environments.nemo_gym import _gym_failure_row_error, _typed_gym_failure from nemo_rl.experience.failures import ( FailureClass, GenerationUnavailable, + GymAdmissionClosed, GymTransportError, NoHealthyShards, RolloutDataFailure, @@ -282,3 +284,114 @@ def test_both_sides_of_the_boundary_share_one_status_policy(self): GymTransportError, ) assert at_source is http_status_is_infra(status) + + +def _gym_failure_row( + *, + status: int | None, + body: str | None = None, + failure_type: str = "ClientResponseError", +) -> dict[str, object]: + """The ``result`` Gym's ``run_examples`` builds for a failed /run in sidecar mode.""" + return { + "_ng_failure_class": "agent_run_error", + "_ng_failure_type": failure_type, + "_ng_failure_message": f"synthetic {status}", + "_ng_failure_http_status": status, + "_ng_failure_response_body": body, + } + + +_ADMISSION_CLOSED_BODY = ( + '{"error":{"code":"admission_closed",' + '"detail":"agent admission is closed for a checkpoint"}}' +) + + +@pytest.mark.nemo_gym +class TestGymFailureRowsKeepTheirRow: + """Sidecar failure rows map to the same typed failures, and a refusal names its row.""" + + def test_admission_closed_names_its_row_and_survives_the_boundary(self): + failure = _gym_failure_row_error( + {"_rowidx": 3}, _gym_failure_row(status=409, body=_ADMISSION_CLOSED_BODY) + ) + + assert isinstance(failure, GymAdmissionClosed) + restored = ray_cloudpickle.loads(ray_cloudpickle.dumps(failure)) + assert isinstance(restored, GymAdmissionClosed) + assert restored.rowidx == 3 + assert "admission_closed" in str(restored) + # A raise without a dispatch recorder is retried, never charged as data. + assert classify_rollout_failure(restored) is FailureClass.INFRA + + def test_another_409_stays_a_data_failure(self): + body = '{"error":{"code":"stale_attempt","detail":"older attempt"}}' + failure = _gym_failure_row_error( + {"_rowidx": 0}, _gym_failure_row(status=409, body=body) + ) + assert type(failure) is RolloutDataFailure + + @pytest.mark.parametrize("status", [400, 404, 408, 429, 500, 503]) + def test_http_failures_match_the_exception_path(self, status): + """Sidecar mode must not change how an ordinary /run failure is budgeted.""" + failure = _gym_failure_row_error( + {"_rowidx": 0}, _gym_failure_row(status=status) + ) + assert isinstance(failure, GymTransportError) is http_status_is_infra(status) + assert str(status) in str(failure) + + def test_no_http_reply_is_transport_unless_the_reply_was_not_json(self): + transport = _gym_failure_row_error( + {"_rowidx": 0}, + _gym_failure_row(status=None, failure_type="ServerDisconnectedError"), + ) + undecodable = _gym_failure_row_error( + {"_rowidx": 0}, + _gym_failure_row(status=None, failure_type="JSONDecodeError"), + ) + assert isinstance(transport, GymTransportError) + assert type(undecodable) is RolloutDataFailure + + def test_a_row_that_ran_is_not_a_failure(self): + assert _gym_failure_row_error({"_rowidx": 0}, {"reward": 1.0}) is None + + def test_a_judge_failure_is_a_scored_row_not_a_transport_failure(self): + """Gym's judge failsafe answers 200 with a masked zero-reward row. + + It carries ``_ng_failure_class`` but ran to completion, so re-sending it + would hit the same judge and drop the prompt instead of masking one row. + """ + judge_failed = { + "reward": 0.0, + "mask_sample": True, + "failure_kind": "judge_failed", + "_ng_failure_class": "judge_failed", + "_ng_failure_judge_error": "judge returned no verdict", + } + assert _gym_failure_row_error({"_rowidx": 0}, judge_failed) is None + + def test_an_agent_run_error_without_an_http_reply_stays_retryable(self): + """An agent may answer 200 with its own run error; it holds no rollout.""" + failure = _gym_failure_row_error( + {"_rowidx": 0}, + { + "reward": None, + "response": None, + "_ng_failure_class": "agent_run_error", + "_ng_failure_message": "Hermes response status: failed", + }, + ) + assert isinstance(failure, GymTransportError) + + +def test_a_row_that_ran_needs_no_nemo_gym_import(monkeypatch): + """Every streamed row is checked, including in environments without the extra.""" + for module in ( + "nemo_gym", + "nemo_gym._checkpoint.errors", + "nemo_gym.rollout_collection", + ): + monkeypatch.setitem(sys.modules, module, None) + + assert _gym_failure_row_error({"_rowidx": 0}, {"reward": 1.0}) is None diff --git a/tests/unit/experience/test_rollout_generation_failures.py b/tests/unit/experience/test_rollout_generation_failures.py index 6ef0095f085..7245bfc3a26 100644 --- a/tests/unit/experience/test_rollout_generation_failures.py +++ b/tests/unit/experience/test_rollout_generation_failures.py @@ -27,16 +27,19 @@ """ import asyncio +import contextlib import uuid import pytest import ray.exceptions import torch +from nemo_rl.environments.gym_checkpoint_adapter import GymCheckpointEpisode from nemo_rl.environments.interfaces import EnvironmentReturn from nemo_rl.experience.failures import ( FailureClass, GenerationUnavailable, + GymAdmissionClosed, GymTransportError, RolloutDataFailure, RolloutFailure, @@ -47,6 +50,7 @@ from nemo_rl.experience.rollout_manager import ( AsyncRolloutImpl, RequestDeadlineRegistry, + RolloutDispatchAdmissionGate, RolloutManager, RolloutRetryPolicy, RolloutStats, @@ -54,6 +58,7 @@ _classify_generation_failure, _Deadline, _gather_cancelling_siblings, + _LedgerDispatchRecorder, ) from nemo_rl.experience.rollout_recovery import RecoveryGranularity from nemo_rl.utils.timer import Timer @@ -540,6 +545,7 @@ def _make_gym_impl( from nemo_rl.experience.rollout_manager import AsyncNemoGymRolloutImpl impl = object.__new__(AsyncNemoGymRolloutImpl) + impl._dispatch_admission_gate = None impl._tokenizer = _FakeTokenizer() impl._task_to_env = {"nemo_gym": type("Env", (), {"run_rollouts": gym_method})()} impl._num_generations_per_prompt = num_generations @@ -877,10 +883,11 @@ async def credit_flow() -> None: loop = asyncio.get_running_loop() async with _Deadline(1.0, "generation turn", registry=registry) as live: await asyncio.sleep(0.2) - registry.suspend() + registry.suspend("weight_sync") assert live._timeout is not None and live._timeout.when() is None banked = live._remaining assert banked is not None and 0.0 < banked <= 0.8 + registry.suspend("gym_checkpoint") async with _Deadline(0.2, "generation turn", registry=registry) as late: # Entered while suspended: disarmed, full budget banked. assert late._timeout is not None and late._timeout.when() is None @@ -889,7 +896,13 @@ async def credit_flow() -> None: assert off._remaining is None # disabled: nothing to bank # Wall clock far past both budgets: nothing fires while frozen. await asyncio.sleep(1.2) - registry.resume() + # Clocks stay frozen until every holder has released its pause. + registry.resume("weight_sync") + assert registry.suspended + assert live._timeout.when() is None + assert late._timeout.when() is None + registry.resume("gym_checkpoint") + assert not registry.suspended assert live._timeout.when() == pytest.approx( loop.time() + banked, abs=0.05 ) @@ -906,3 +919,447 @@ async def expiry_still_fires() -> None: with pytest.raises(RolloutTimeout, match="generation turn exceeded"): asyncio.run(expiry_still_fires()) + + +class _RecordingDispatchRecorder: + """RolloutDispatchRecorder stand-in that logs every ledger hook in order.""" + + def __init__(self, events: list[tuple]) -> None: + self.events = events + + @contextlib.asynccontextmanager + async def submitting(self, generation_indices): + indices = list(generation_indices) + self.events.append(("submitting", indices)) + yield lambda: self.events.append(("mark", indices)) + + async def refused(self, generation_index: int) -> None: + self.events.append(("refused", generation_index)) + + +class _AdmissionClosedGymMethod: + """Gym stream whose first dispatch refuses ``parked`` rows at checkpoint admission.""" + + def __init__(self, parked: set[int], events: list[tuple]) -> None: + self._parked = set(parked) + self.events = events + self.dispatched: list[list[int]] = [] + self.release_admitted = asyncio.Event() + self.release_admitted.set() + + def options(self, **kwargs): + del kwargs + return self + + def remote(self, inputs, timer_prefix, per_prompt=False): + del timer_prefix, per_prompt + rows = [row["_rowidx"] for row in inputs] + self.dispatched.append(rows) + self.events.append(("remote", rows)) + first = len(self.dispatched) == 1 + return self._stream(rows, first) + + async def _stream(self, rows, first): + refused = [row for row in rows if first and row in self._parked] + admitted = [row for row in rows if row not in refused] + # Gym answers a refusal immediately; admitted rows finish later. + for rowidx in refused: + yield self._parked_item(rowidx) + for rowidx in admitted: + yield self._admitted_item(rowidx) + + @staticmethod + async def _parked_item(rowidx): + return GymAdmissionClosed("admission_closed", rowidx) + + async def _admitted_item(self, rowidx): + await self.release_admitted.wait() + return await _row_result(rowidx) + + +class TestGymAdmissionClosedRows: + """A row Gym refuses at checkpoint admission never ran: unwind it, re-send it.""" + + def test_refused_row_is_unwound_and_resent_without_spending_a_row_attempt( + self, monkeypatch + ): + from nemo_rl.experience import rollout_manager + + monkeypatch.setattr(rollout_manager, "_PARKED_REDISPATCH_BACKOFF_S", 0.0) + events: list[tuple] = [] + method = _AdmissionClosedGymMethod(parked={1}, events=events) + stats = RolloutStats() + # One row attempt: a refusal must not consume it. + impl = _make_gym_impl(method, num_generations=2, row_attempts=1, stats=stats) + impl._dispatch_admission_gate = RolloutDispatchAdmissionGate() + + completions, _, _ = asyncio.run( + impl._run_rollouts( + _gym_rows(2), + Timer(), + "timing/rollout", + dispatch_recorder=_RecordingDispatchRecorder(events), + ) + ) + + assert len(completions) == 2 + assert method.dispatched == [[0, 1], [1]] + # Marked only after Ray accepted each submission; the refusal is unwound + # before the row is re-sent. + assert events == [ + ("submitting", [0, 1]), + ("remote", [0, 1]), + ("mark", [0, 1]), + ("refused", 1), + ("submitting", [1]), + ("remote", [1]), + ("mark", [1]), + ] + assert stats.gym_admission_closed_rows == 1 + assert stats.gym_row_redispatches == 0 + assert not stats.data_retries_by_reason + + def test_nothing_is_submitted_or_marked_while_the_gate_is_closed(self): + events: list[tuple] = [] + method = _AdmissionClosedGymMethod(parked=set(), events=events) + impl = _make_gym_impl(method, num_generations=2) + gate = RolloutDispatchAdmissionGate() + impl._dispatch_admission_gate = gate + + async def exercise() -> None: + async with gate.closed(): + task = asyncio.create_task( + impl._run_rollouts( + _gym_rows(2), + Timer(), + "timing/rollout", + dispatch_recorder=_RecordingDispatchRecorder(events), + ) + ) + await asyncio.sleep(0.05) + assert events == [] + await asyncio.wait_for(task, timeout=5.0) + + asyncio.run(exercise()) + assert events[:3] == [ + ("submitting", [0, 1]), + ("remote", [0, 1]), + ("mark", [0, 1]), + ] + + def test_checkpoint_excludes_a_refused_row_and_does_not_deadlock(self, monkeypatch): + """The real recorder, ledger and barrier, with a checkpoint mid-flight. + + Row 1 is refused while row 0 is still running in Gym. A checkpoint that + closes the gate must then see only row 0 as Gym-owned, and row 1's + re-send must wait at the gate instead of becoming DISPATCHED under it. + """ + from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneCheckpointBarrier, + ) + from nemo_rl.experience import rollout_manager + from nemo_rl.experience.rollout_recovery import ( + RecoveryTargetLevel, + RolloutAttemptStatus, + RolloutRecoveryLedger, + ) + + # Long enough for the checkpoint below to close the gate before the re-send. + monkeypatch.setattr(rollout_manager, "_PARKED_REDISPATCH_BACKOFF_S", 0.2) + events: list[tuple] = [] + method = _AdmissionClosedGymMethod(parked={1}, events=events) + method.release_admitted.clear() + impl = _make_gym_impl(method, num_generations=2, row_attempts=1) + gate = RolloutDispatchAdmissionGate() + impl._dispatch_admission_gate = gate + + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + ledger = RolloutRecoveryLedger() + async with barrier.mutation() as cut: + group = ledger.reserve_group( + cut, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload={"idx": 7, "message_log": []}, + expected_generations=2, + target_step=7, + start_weight_version=7, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + manager = object.__new__(RolloutManager) + manager._recovery_ledger = ledger + manager._data_plane_checkpoint_barrier = barrier + recorder = _LedgerDispatchRecorder( + manager=manager, group_id="g7", gym_instance_id="tools/replica-0" + ) + + def status(index: int) -> RolloutAttemptStatus: + return ledger.get_group("g7").siblings[index].current_attempt.status + + task = asyncio.create_task( + impl._run_rollouts( + _gym_rows(2), + Timer(), + "timing/rollout", + dispatch_recorder=recorder, + ) + ) + # Both rows start RESERVED and are marked together, so this only + # holds once Gym's refusal of row 1 has been unwound. + while not ( + status(0) is RolloutAttemptStatus.DISPATCHED + and status(1) is RolloutAttemptStatus.RESERVED + ): + await asyncio.sleep(0.001) + + async def checkpoint() -> tuple: + async with gate.closed(): + async with barrier.checkpoint(): + inventory = ledger.gym_checkpoint_inventory( + {"tools/replica-0"} + )["tools/replica-0"] + await asyncio.sleep(0.05) + return inventory, status(1) + + inventory, held_status = await asyncio.wait_for(checkpoint(), timeout=5.0) + assert inventory == (group.gym_episode(0),) + assert held_status is RolloutAttemptStatus.RESERVED + + method.release_admitted.set() + await asyncio.wait_for(task, timeout=5.0) + assert status(0) is RolloutAttemptStatus.DISPATCHED + assert status(1) is RolloutAttemptStatus.DISPATCHED + assert method.dispatched == [[0, 1], [1]] + + asyncio.run(exercise()) + + +class _EarlyEndingGym: + """NemoGym actor whose first stream returns row 0, then dies with rows still running. + + Gym runs every row of a dispatch as its own task, so the rows that never + came back are still executing there after the stream breaks. + """ + + def __init__(self) -> None: + self.events: list[tuple] = [] + self.run_rollouts = self._RunRollouts(self) + self.checkpoint_retire = self._Retire(self) + + class _RunRollouts: + def __init__(self, gym: "_EarlyEndingGym") -> None: + self._gym = gym + + def options(self, **kwargs): + del kwargs + return self + + def remote(self, inputs, timer_prefix, per_prompt=False): + del timer_prefix, per_prompt + first = not any(event[0] == "dispatch" for event in self._gym.events) + self._gym.events.append( + ("dispatch", {row["_rowidx"]: row["_ng_rollout_id"] for row in inputs}) + ) + return self._stream([row["_rowidx"] for row in inputs], first) + + async def _stream(self, rows, first): + if first: + yield _row_result(rows[0]) + raise ConnectionResetError("gym stream died mid-flight") + for rowidx in rows: + yield _row_result(rowidx) + + class _Retire: + def __init__(self, gym: "_EarlyEndingGym") -> None: + self._gym = gym + + def remote(self, checkpoint_id, episodes, *, deadline_ts): + del checkpoint_id, deadline_ts + self._gym.events.append(("retire", frozenset(episodes))) + + async def _done() -> None: + return None + + return _done() + + +class _CuedEarlyEndingGym(_EarlyEndingGym): + """Like ``_EarlyEndingGym``, but the first stream breaks only on cue.""" + + def __init__(self) -> None: + super().__init__() + self.break_stream = asyncio.Event() + stream = self.run_rollouts._stream + + async def cued(rows, first): + async for result in stream(rows, first): + yield result + if first: + await self.break_stream.wait() + + self.run_rollouts._stream = cued + + +class TestGymEarlyStreamEnd: + """Rows still running in Gym when a stream breaks must not overlap their replacements.""" + + @staticmethod + def _dispatch( + gym: _EarlyEndingGym, + granularity: RecoveryGranularity, + row_attempts: int, + gate: RolloutDispatchAdmissionGate | None = None, + ): + from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneCheckpointBarrier, + ) + from nemo_rl.experience.rollout_recovery import ( + RecoveryTargetLevel, + RolloutRecoveryLedger, + ) + + impl = _make_gym_impl(None, num_generations=3, row_attempts=row_attempts) + impl._task_to_env = {"nemo_gym": gym} + impl._dispatch_admission_gate = gate or RolloutDispatchAdmissionGate() + barrier = DataPlaneCheckpointBarrier() + ledger = RolloutRecoveryLedger() + manager = object.__new__(RolloutManager) + manager._recovery_ledger = ledger + manager._data_plane_checkpoint_barrier = barrier + + async def reserve(): + async with barrier.mutation() as cut: + return ledger.reserve_group( + cut, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload={"idx": 7, "message_log": []}, + expected_generations=3, + target_step=7, + start_weight_version=7, + recovery_granularity=granularity, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + + group = asyncio.run(reserve()) + rows = [ + { + "_rowidx": index, + "task_source": "workplace_assistant", + "_ng_rollout_id": group.gate_rollout_id(index), + } + for index in range(3) + ] + recorder = _LedgerDispatchRecorder( + manager=manager, group_id="g7", gym_instance_id="tools/replica-0" + ) + episodes = { + index: GymCheckpointEpisode(*group.gym_episode(index)) for index in range(3) + } + + async def run(): + return await impl._run_rollouts( + rows, + Timer(), + "timing/rollout", + recovery_granularity=granularity, + dispatch_recorder=recorder, + ) + + return run, episodes, ledger + + def test_rows_still_running_are_retired_before_they_are_resent_under_new_ids(self): + from nemo_rl.experience.rollout_recovery import RolloutAttemptStatus + + gym = _EarlyEndingGym() + run, episodes, ledger = self._dispatch( + gym, RecoveryGranularity.SIBLING, row_attempts=3 + ) + + completions, _, _ = asyncio.run(run()) + + assert len(completions) == 3 + kinds = [event[0] for event in gym.events] + assert kinds == ["dispatch", "retire", "dispatch"], gym.events + first, retired, resent = (event[1] for event in gym.events) + # Rows 1 and 2 never came back, so Gym may still be running them. + assert retired == frozenset({episodes[1], episodes[2]}) + assert set(resent) == {1, 2} + assert set(resent.values()).isdisjoint(first.values()) + group = ledger.get_group("g7") + for index in (1, 2): + sibling = group.siblings[index] + assert sibling.attempts[0].status is RolloutAttemptStatus.ABANDONED + assert sibling.current_attempt.status is RolloutAttemptStatus.DISPATCHED + assert group.gate_rollout_id(index) == resent[index] + assert ledger.gym_checkpoint_retirements({"tools/replica-0"}) == { + "tools/replica-0": () + } + + def test_a_prompt_group_retires_its_running_siblings_before_giving_up(self): + gym = _EarlyEndingGym() + run, episodes, _ = self._dispatch( + gym, RecoveryGranularity.PROMPT_GROUP, row_attempts=1 + ) + + with pytest.raises(ConnectionResetError): + asyncio.run(run()) + + assert [event[0] for event in gym.events] == ["dispatch", "retire"] + retired = gym.events[1][1] + # The prompt retry dispatches fresh attempts next; none may overlap these. + assert {episodes[1], episodes[2]} <= retired <= set(episodes.values()) + + def test_a_failed_retire_leaves_the_dropped_rows_for_the_next_checkpoint(self): + gym = _EarlyEndingGym() + + def unreachable(checkpoint_id, episodes, *, deadline_ts): + del checkpoint_id, deadline_ts + gym.events.append(("retire", frozenset(episodes))) + raise ConnectionRefusedError("gym control plane is down") + + gym.checkpoint_retire.remote = unreachable + run, episodes, ledger = self._dispatch( + gym, RecoveryGranularity.SIBLING, row_attempts=3 + ) + + completions, _, _ = asyncio.run(run()) + + assert len(completions) == 3 + assert [event[0] for event in gym.events] == ["dispatch", "retire", "dispatch"] + pending = ledger.gym_checkpoint_retirements({"tools/replica-0"}) + assert set(pending["tools/replica-0"]) == { + (episodes[index].rollout_id, episodes[index].attempt) for index in (1, 2) + } + + def test_a_dropped_row_is_not_retired_while_a_checkpoint_holds_dispatch_closed( + self, + ): + """Gym refuses a retire once a checkpoint is open, so it waits for resume.""" + gym = _CuedEarlyEndingGym() + gate = RolloutDispatchAdmissionGate() + run, episodes, _ = self._dispatch( + gym, RecoveryGranularity.SIBLING, row_attempts=3, gate=gate + ) + + async def exercise() -> None: + rollout = asyncio.create_task(run()) + while not gym.events: + await asyncio.sleep(0) + async with gate.closed(): + # The stream breaks while the checkpoint holds Gym prepared. + gym.break_stream.set() + await asyncio.sleep(0.05) + assert [event[0] for event in gym.events] == ["dispatch"] + completions, _, _ = await asyncio.wait_for(rollout, timeout=5.0) + assert len(completions) == 3 + + asyncio.run(exercise()) + + assert [event[0] for event in gym.events] == ["dispatch", "retire", "dispatch"] + assert gym.events[1][1] == frozenset({episodes[1], episodes[2]}) diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 283e90d91ca..bd6bf147f12 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -65,6 +65,7 @@ ) from nemo_rl.experience.rollout_recovery import ( RecoveryGranularity, + RecoveryTargetLevel, RolloutRecoveryLedger, ) from nemo_rl.experience.rollouts import ( @@ -648,8 +649,8 @@ def test_reservation_persists_the_resolved_task_source_recovery_policy(self): buf = _FakeBuffer() mgr = _make_manager(buf, _FakeImpl()) mgr._rollout_recovery_config = RolloutRecoveryConfig( - task_source_granularity_overrides={ - "genrm_compare": RecoveryGranularity.PROMPT_GROUP + task_source_target_level_overrides={ + "genrm_compare": RecoveryTargetLevel.PROMPT_GROUP } ) prompt = { @@ -667,6 +668,7 @@ def test_reservation_persists_the_resolved_task_source_recovery_policy(self): assert group.task_source == "genrm_compare" assert group.recovery_granularity is RecoveryGranularity.PROMPT_GROUP + assert group.restore_level is RecoveryTargetLevel.PROMPT_GROUP def test_recovery_mutation_requires_the_controller_barrier(self): mgr = _make_manager(_FakeBuffer(), _FakeImpl()) @@ -2113,7 +2115,7 @@ def test_prompt_group_policy_retries_the_complete_live_cohort(self): mgr = _make_capture_manager( buf, recovery_config=RolloutRecoveryConfig( - default_granularity=RecoveryGranularity.PROMPT_GROUP + target_level=RecoveryTargetLevel.PROMPT_GROUP ), ) mgr._retry_policy = RolloutRetryPolicy.single_attempt( @@ -2179,7 +2181,7 @@ async def run_rollout( def test_prompt_group_restore_redispatches_every_sibling(self): recovery_config = RolloutRecoveryConfig( - default_granularity=RecoveryGranularity.PROMPT_GROUP + target_level=RecoveryTargetLevel.PROMPT_GROUP ) first = _make_capture_manager( _FakeCaptureBuffer(), recovery_config=recovery_config @@ -2198,7 +2200,7 @@ def test_prompt_group_restore_redispatches_every_sibling(self): _FakeCaptureBuffer(), # The saved group policy wins over the new process configuration. recovery_config=RolloutRecoveryConfig( - default_granularity=RecoveryGranularity.SIBLING + target_level=RecoveryTargetLevel.SIBLING ), ) _with_cut( diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 96b8a3fd691..584359ae4de 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -39,6 +39,7 @@ PromptGroupRecoveryRecord, PromptRef, RecoveryGranularity, + RecoveryTargetLevel, RolloutAttemptRecord, RolloutAttemptStatus, RolloutRecoveryLedger, @@ -424,13 +425,13 @@ async def exercise() -> None: def test_recovery_config_resolves_agent_then_task_source_then_default() -> None: config = RolloutRecoveryConfig( - default_granularity=RecoveryGranularity.SIBLING, - task_source_granularity_overrides={ - "genrm_compare": RecoveryGranularity.PROMPT_GROUP, + target_level=RecoveryTargetLevel.TURN, + task_source_target_level_overrides={ + "genrm_compare": RecoveryTargetLevel.PROMPT_GROUP, }, - agent_granularity_overrides={ - "legacy_genrm_agent": RecoveryGranularity.PROMPT_GROUP, - "sibling_agent": RecoveryGranularity.SIBLING, + agent_target_level_overrides={ + "legacy_genrm_agent": RecoveryTargetLevel.PROMPT_GROUP, + "sibling_agent": RecoveryTargetLevel.SIBLING, }, ) @@ -465,12 +466,37 @@ def test_recovery_config_resolves_agent_then_task_source_then_default() -> None: assert source_policy.task_source == "genrm_compare" assert source_policy.granularity is RecoveryGranularity.PROMPT_GROUP + assert source_policy.restore_level is RecoveryTargetLevel.PROMPT_GROUP assert agent_policy.task_source == "genrm_compare" assert agent_policy.granularity is RecoveryGranularity.SIBLING + assert agent_policy.restore_level is RecoveryTargetLevel.SIBLING assert default_policy.task_source == "other" assert default_policy.granularity is RecoveryGranularity.SIBLING + assert default_policy.restore_level is RecoveryTargetLevel.TURN assert legacy_policy.task_source is None assert legacy_policy.granularity is RecoveryGranularity.PROMPT_GROUP + assert legacy_policy.restore_level is RecoveryTargetLevel.PROMPT_GROUP + + +def test_recovery_overrides_may_only_make_the_target_coarser() -> None: + with pytest.raises(ValueError, match="may only force a coarser recovery level"): + RolloutRecoveryConfig( + target_level=RecoveryTargetLevel.SIBLING, + agent_target_level_overrides={ + "too-fine": RecoveryTargetLevel.TURN, + }, + ) + + prefix_config = RolloutRecoveryConfig( + target_level=RecoveryTargetLevel.PREFIX, + task_source_target_level_overrides={ + "restart-call": RecoveryTargetLevel.TURN, + }, + ) + policy = prefix_config.resolve_for_prompt( + {"extra_env_info": {"task_source": "restart-call"}} + ) + assert policy.restore_level is RecoveryTargetLevel.TURN @pytest.mark.parametrize( @@ -498,7 +524,7 @@ def test_recovery_config_rejects_malformed_prompt_identity( def test_recovery_config_rejects_removed_task_name_override() -> None: - with pytest.raises(ValueError, match="task_source_granularity_overrides"): + with pytest.raises(ValueError, match="target_level"): RolloutRecoveryConfig( **{ "task_granularity_overrides": { @@ -756,6 +782,77 @@ def test_restart_preserves_sealed_sibling_and_retries_only_interrupted_one() -> assert retry.siblings[1].current_attempt.status is RolloutAttemptStatus.RESERVED +@pytest.mark.parametrize( + "restore_level", [RecoveryTargetLevel.PREFIX, RecoveryTargetLevel.TURN] +) +def test_gym_state_restart_preserves_sealed_sibling_and_restores_other_from_gym( + restore_level: RecoveryTargetLevel, +) -> 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, + recovery_granularity=RecoveryGranularity.SIBLING, + restore_level=restore_level, + admitted=True, + ) + _mutate( + lambda cut: ledger.mark_group_dispatched( + cut, + "g7", + gym_instance_id="tools/replica-0", + ) + ) + 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, + mask_sample=False, + ) + ) + + restored = RolloutRecoveryLedger.from_state_dict(ledger.state_dict()) + inventory = restored.gym_checkpoint_inventory({"tools/replica-0"}) + (unfinished_episode,) = inventory["tools/replica-0"] + + _mutate( + lambda cut: restored.prepare_for_restart( + cut, + restored_gym_episodes={ + ("tools/replica-0", *unfinished_episode), + }, + ) + ) + recovered = restored.get_group("g7") + + assert recovered.siblings[0].current_attempt.attempt_id == sealed_attempt_id + assert recovered.siblings[0].current_attempt.status is RolloutAttemptStatus.SEALED + resumed = recovered.siblings[1].current_attempt + assert resumed.status is RolloutAttemptStatus.RESERVED + assert resumed.gym_instance_id == "tools/replica-0" + assert resumed.gym_attempt == 1 + 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.status is RolloutAttemptStatus.SEALED + assert retry.siblings[1].current_attempt.status is RolloutAttemptStatus.RESERVED + + @pytest.mark.parametrize( "recovery_granularity", [RecoveryGranularity.SIBLING, RecoveryGranularity.PROMPT_GROUP], @@ -831,7 +928,7 @@ def test_missing_receipt_is_a_restart_safe_sealed_placeholder( assert rewards == [0.0, 1.0] assert mask_sample == [True, False] - state["schema_version"] = 3 + state["schema_version"] = ROLLOUT_RECOVERY_SCHEMA_VERSION + 1 with pytest.raises(ValueError, match="Unsupported rollout-recovery schema version"): RolloutRecoveryLedger.from_state_dict(state) @@ -856,6 +953,7 @@ def test_prompt_group_restart_retries_every_sibling_when_one_is_unfinished() -> state = ledger.state_dict() assert state["groups"][0]["task_source"] == "genrm_compare" assert state["groups"][0]["recovery_granularity"] == "prompt_group" + assert state["groups"][0]["restore_level"] == "prompt_group" restored = RolloutRecoveryLedger.from_state_dict(state) _mutate(lambda cut: restored.prepare_for_restart(cut)) @@ -863,6 +961,7 @@ def test_prompt_group_restart_retries_every_sibling_when_one_is_unfinished() -> assert recovered.task_source == "genrm_compare" assert recovered.recovery_granularity is RecoveryGranularity.PROMPT_GROUP + assert recovered.restore_level is RecoveryTargetLevel.PROMPT_GROUP assert [sibling.current_attempt.status for sibling in recovered.siblings] == [ RolloutAttemptStatus.ABANDONED, RolloutAttemptStatus.ABANDONED, @@ -876,6 +975,185 @@ def test_prompt_group_restart_retries_every_sibling_when_one_is_unfinished() -> ] +def test_turn_restore_target_is_persisted_separately_from_live_retry() -> 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, + recovery_granularity=RecoveryGranularity.SIBLING, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + + state = ledger.state_dict() + assert state["groups"][0]["recovery_granularity"] == "sibling" + assert state["groups"][0]["restore_level"] == "turn" + + restored = RolloutRecoveryLedger.from_state_dict(state).get_group("g7") + assert restored.recovery_granularity is RecoveryGranularity.SIBLING + assert restored.restore_level is RecoveryTargetLevel.TURN + + +@pytest.mark.parametrize("schema_version", [2, 3]) +def test_old_rollout_recovery_schemas_are_rejected( + schema_version: int, +) -> 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, + recovery_granularity=RecoveryGranularity.SIBLING, + admitted=True, + ) + _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) + + state = ledger.state_dict() + state["schema_version"] = schema_version + + with pytest.raises(ValueError, match="Unsupported rollout-recovery schema"): + RolloutRecoveryLedger.from_state_dict(state) + + +def test_turn_restart_preserves_gym_owner_and_advances_attempt() -> 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, + recovery_granularity=RecoveryGranularity.SIBLING, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + _mutate( + lambda cut: ledger.mark_group_dispatched( + cut, + "g7", + gym_instance_id="tools/replica-1", + ) + ) + + state = ledger.state_dict() + attempt_state = state["groups"][0]["siblings"][0]["attempts"][0] + assert attempt_state["gym_instance_id"] == "tools/replica-1" + assert attempt_state["gym_attempt"] == 0 + + restored = RolloutRecoveryLedger.from_state_dict(state) + inventory = restored.gym_checkpoint_inventory( + {"judge/replica-0", "tools/replica-1"} + ) + ((rollout_id, attempt),) = inventory["tools/replica-1"] + assert inventory["judge/replica-0"] == () + assert attempt == 0 + + _mutate( + lambda cut: restored.prepare_for_restart( + cut, + restored_gym_episodes={ + ("tools/replica-1", rollout_id, attempt), + }, + ) + ) + recovered = restored.get_group("g7") + current = recovered.siblings[0].current_attempt + assert current.status is RolloutAttemptStatus.RESERVED + assert current.gym_instance_id == "tools/replica-1" + assert current.gym_attempt == 1 + assert recovered.gate_rollout_id(0) == f"{rollout_id}-a1" + + repeated_inventory = restored.gym_checkpoint_inventory( + {"judge/replica-0", "tools/replica-1"} + ) + assert repeated_inventory["tools/replica-1"] == ((rollout_id, 1),) + + _mutate( + lambda cut: restored.prepare_for_restart( + cut, + restored_gym_episodes={ + ("tools/replica-1", rollout_id, 1), + }, + ) + ) + recovered = restored.get_group("g7") + current = recovered.siblings[0].current_attempt + assert current.status is RolloutAttemptStatus.RESERVED + assert current.gym_attempt == 2 + assert recovered.gate_rollout_id(0) == f"{rollout_id}-a2" + + with pytest.raises(ValueError, match="checkpoint owner"): + _mutate( + lambda cut: restored.mark_group_dispatched( + cut, + "g7", + gym_instance_id="judge/replica-0", + ) + ) + + +def test_turn_restart_falls_back_to_sibling_when_gym_episode_is_absent() -> 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, + recovery_granularity=RecoveryGranularity.SIBLING, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + _mutate( + lambda cut: ledger.mark_group_dispatched( + cut, + "g7", + gym_instance_id="tools/replica-0", + ) + ) + + _mutate(lambda cut: ledger.prepare_for_restart(cut)) + + current = ledger.get_group("g7").siblings[0].current_attempt + assert current.status is RolloutAttemptStatus.ABANDONED + + +def test_turn_dispatch_requires_a_gym_owner() -> 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, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + + with pytest.raises(ValueError, match="requires a Gym checkpoint owner"): + _mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7")) + + def test_prompt_group_restart_keeps_a_fully_sealed_group() -> None: ledger = RolloutRecoveryLedger() group = _reserve( @@ -1006,6 +1284,8 @@ def test_checkpoint_rejects_ambiguous_finalization_state( [ ("recovery_granularity", "banana", "invalid recovery_granularity"), ("recovery_granularity", None, "recovery_granularity must be a string"), + ("restore_level", "banana", "invalid restore_level"), + ("restore_level", None, "restore_level must be a string"), ("task_source", 123, "task_source must be a string or None"), ], ) @@ -1097,3 +1377,371 @@ def test_restore_rejects_inconsistent_shared_admission_state() -> None: with pytest.raises(ValueError, match="disagree on phase or target_step"): _load(RolloutRecoveryLedger(), state) + + +def _dispatched_turn_group(ledger: RolloutRecoveryLedger) -> PromptGroupRecoveryRecord: + 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, + recovery_granularity=RecoveryGranularity.SIBLING, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + _mutate( + lambda cut: ledger.mark_group_dispatched( + cut, "g7", gym_instance_id="tools/replica-0" + ) + ) + return group + + +@pytest.mark.parametrize( + ("restore_level", "expected"), + [(RecoveryTargetLevel.PREFIX, True), (RecoveryTargetLevel.TURN, False)], +) +def test_only_a_dispatched_prefix_attempt_grows_a_recoverable_prefix( + restore_level: RecoveryTargetLevel, expected: bool +) -> None: + """Decoding advances only dispatched prefix-level attempts between mutations.""" + 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, + restore_level=restore_level, + admitted=True, + ) + assert not ledger.has_dispatched_generation_prefix() + + _mutate( + lambda cut: ledger.mark_group_dispatched( + cut, "g7", gym_instance_id="tools/replica-0" + ) + ) + assert ledger.has_dispatched_generation_prefix() is expected + + _mutate(lambda cut: ledger.abandon_unsealed(cut, "g7")) + assert not ledger.has_dispatched_generation_prefix() + + +def test_refused_fresh_dispatch_keeps_identity_and_leaves_gym_inventory() -> None: + """Gym recorded nothing for a refused /run, so a checkpoint must not wait on it.""" + ledger = RolloutRecoveryLedger() + group = _dispatched_turn_group(ledger) + refused_attempt_id = group.siblings[1].current_attempt.attempt_id + refused_rollout_id = group.gate_rollout_id(1) + assert ( + len(ledger.gym_checkpoint_inventory({"tools/replica-0"})["tools/replica-0"]) + == 2 + ) + + _mutate(lambda cut: ledger.release_refused_dispatch(cut, "g7", generation_index=1)) + + released = ledger.get_group("g7") + attempt = released.siblings[1].current_attempt + assert attempt.status is RolloutAttemptStatus.RESERVED + assert attempt.attempt_id == refused_attempt_id + assert released.gate_rollout_id(1) == refused_rollout_id + assert ( + released.siblings[0].current_attempt.status is RolloutAttemptStatus.DISPATCHED + ) + (remaining,) = ledger.gym_checkpoint_inventory({"tools/replica-0"})[ + "tools/replica-0" + ] + assert remaining == group.gym_episode(0) + + # The same attempt is re-sent once admission reopens. + _mutate( + lambda cut: ledger.mark_group_dispatched( + cut, "g7", generation_indices=[1], gym_instance_id="tools/replica-0" + ) + ) + assert ( + ledger.get_group("g7").siblings[1].current_attempt.status + is RolloutAttemptStatus.DISPATCHED + ) + + +def test_refused_restored_turn_stays_gym_owned() -> None: + """A restored turn waiting to be re-sent must keep its committed continuation.""" + ledger = RolloutRecoveryLedger() + _dispatched_turn_group(ledger) + restored = RolloutRecoveryLedger.from_state_dict(ledger.state_dict()) + episodes = restored.gym_checkpoint_inventory({"tools/replica-0"})["tools/replica-0"] + _mutate( + lambda cut: restored.prepare_for_restart( + cut, + restored_gym_episodes={ + ("tools/replica-0", *episode) for episode in episodes + }, + ) + ) + _mutate( + lambda cut: restored.mark_group_dispatched( + cut, "g7", gym_instance_id="tools/replica-0" + ) + ) + + _mutate( + lambda cut: restored.release_refused_dispatch(cut, "g7", generation_index=1) + ) + + attempt = restored.get_group("g7").siblings[1].current_attempt + assert attempt.status is RolloutAttemptStatus.RESERVED + assert attempt.gym_attempt == 1 + assert ( + restored.get_group("g7").gym_episode(1) + in restored.gym_checkpoint_inventory({"tools/replica-0"})["tools/replica-0"] + ) + + +def test_refusal_requires_a_dispatched_attempt() -> None: + ledger = RolloutRecoveryLedger() + _dispatched_turn_group(ledger) + _mutate(lambda cut: ledger.release_refused_dispatch(cut, "g7", generation_index=0)) + + with pytest.raises(ValueError, match="only a dispatched rollout attempt"): + _mutate( + lambda cut: ledger.release_refused_dispatch(cut, "g7", generation_index=0) + ) + + +def test_dropping_a_dispatched_group_records_its_gym_episodes_for_retirement() -> None: + """Gym keeps running what it admitted, so the next checkpoint must retire it.""" + abandoned_ledger = RolloutRecoveryLedger() + abandoned = _dispatched_turn_group(abandoned_ledger) + _mutate(lambda cut: abandoned_ledger.abandon_unsealed(cut, "g7")) + + discarded_ledger = RolloutRecoveryLedger() + discarded = _dispatched_turn_group(discarded_ledger) + _mutate(lambda cut: discarded_ledger.discard_group(cut, "g7")) + + for ledger, group in ((abandoned_ledger, abandoned), (discarded_ledger, discarded)): + assert ledger.gym_checkpoint_retirements({"tools/replica-0"}) == { + "tools/replica-0": tuple( + sorted(group.gym_episode(index) for index in (0, 1)) + ) + } + # Bookkeeping for this process's Gym executions, never persisted. + reloaded = RolloutRecoveryLedger.from_state_dict(abandoned_ledger.state_dict()) + assert reloaded.gym_checkpoint_retirements({"tools/replica-0"}) == { + "tools/replica-0": () + } + + +def test_attempts_that_never_reached_gym_are_not_retired() -> 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, + recovery_granularity=RecoveryGranularity.SIBLING, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + _mutate(lambda cut: ledger.abandon_unsealed(cut, "g7")) + + assert ledger.gym_checkpoint_retirements({"tools/replica-0"}) == { + "tools/replica-0": () + } + + +def test_dropping_a_restored_reserved_attempt_retires_its_restored_gym_state() -> None: + ledger = RolloutRecoveryLedger() + _dispatched_turn_group(ledger) + restored = RolloutRecoveryLedger.from_state_dict(ledger.state_dict()) + episodes = restored.gym_checkpoint_inventory({"tools/replica-0"})["tools/replica-0"] + _mutate( + lambda cut: restored.prepare_for_restart( + cut, + restored_gym_episodes={ + ("tools/replica-0", *episode) for episode in episodes + }, + ) + ) + group = restored.get_group("g7") + + _mutate(lambda cut: restored.discard_group(cut, "g7")) + + retired = restored.gym_checkpoint_retirements({"tools/replica-0"})[ + "tools/replica-0" + ] + assert set(retired) == {group.gym_episode(index) for index in (0, 1)} + assert {attempt for _, attempt in retired} == {1} + + +def test_mark_gym_retired_keeps_drops_recorded_after_the_retire_was_sent() -> None: + ledger = RolloutRecoveryLedger() + first = _dispatched_turn_group(ledger) + _mutate(lambda cut: ledger.abandon_unsealed(cut, "g7")) + sent = ledger.gym_checkpoint_retirements({"tools/replica-0"}) + + later = _mutate( + lambda cut: ledger.reserve_group( + cut, + group_id="g8", + admission_id="batch-8", + prompt_id="8", + prompt_payload=_prompt(8), + expected_generations=1, + target_step=8, + start_weight_version=7, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + ) + _mutate( + lambda cut: ledger.mark_group_dispatched( + cut, "g8", gym_instance_id="tools/replica-0" + ) + ) + later = ledger.get_group("g8") + _mutate(lambda cut: ledger.abandon_unsealed(cut, "g8")) + + ledger.mark_gym_retired(sent) + + assert ledger.gym_checkpoint_retirements({"tools/replica-0"}) == { + "tools/replica-0": (later.gym_episode(0),) + } + assert ( + first.gym_episode(0) + not in ledger.gym_checkpoint_retirements({"tools/replica-0"})["tools/replica-0"] + ) + + +def test_a_dropped_dispatch_is_queued_for_retirement_and_replaced() -> None: + """A broken stream leaves the row running in Gym; its re-send needs a new ID.""" + ledger = RolloutRecoveryLedger() + group = _dispatched_turn_group(ledger) + + dropped = _mutate( + lambda cut: ledger.abandon_dropped_dispatch( + cut, "g7", generation_indices=[1], replace=True + ) + ) + + assert dropped == (group.gym_episode(1),) + assert ledger.gym_checkpoint_retirements({"tools/replica-0"}) == { + "tools/replica-0": (group.gym_episode(1),) + } + replaced = ledger.get_group("g7") + assert replaced.siblings[1].attempts[0].status is RolloutAttemptStatus.ABANDONED + assert replaced.siblings[1].current_attempt.status is RolloutAttemptStatus.RESERVED + assert replaced.gate_rollout_id(1) != group.gate_rollout_id(1) + # The sibling that returned keeps its attempt and is not Gym-retired. + assert replaced.gate_rollout_id(0) == group.gate_rollout_id(0) + assert ledger.gym_checkpoint_inventory({"tools/replica-0"}) == { + "tools/replica-0": (group.gym_episode(0),) + } + + +def test_a_dropped_attempt_that_never_reached_gym_keeps_its_identity() -> None: + ledger = RolloutRecoveryLedger() + group = _dispatched_turn_group(ledger) + _mutate(lambda cut: ledger.release_refused_dispatch(cut, "g7", generation_index=1)) + + dropped = _mutate( + lambda cut: ledger.abandon_dropped_dispatch( + cut, "g7", generation_indices=[1], replace=True + ) + ) + + assert dropped == () + assert ledger.get_group("g7").gate_rollout_id(1) == group.gate_rollout_id(1) + + +def test_a_prompt_group_cannot_replace_one_dropped_sibling() -> 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, + recovery_granularity=RecoveryGranularity.PROMPT_GROUP, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + _mutate( + lambda cut: ledger.mark_group_dispatched( + cut, "g7", gym_instance_id="tools/replica-0" + ) + ) + + with pytest.raises(ValueError, match="replaces the whole cohort"): + _mutate( + lambda cut: ledger.abandon_dropped_dispatch( + cut, "g7", generation_indices=[1], replace=True + ) + ) + + +def test_gym_episodes_still_held_looks_up_only_live_current_attempts() -> None: + ledger = RolloutRecoveryLedger() + group = _dispatched_turn_group(ledger) + first, second = group.gym_episode(0), group.gym_episode(1) + # Sibling 1's stream broke: its old attempt is abandoned and re-minted. + _mutate( + lambda cut: ledger.abandon_dropped_dispatch( + cut, "g7", generation_indices=[1], replace=True + ) + ) + + held = ledger.gym_episodes_still_held( + { + "tools/replica-0": [first, second, ("not-a-rollout-id", 0)], + "tools/replica-1": [first], + } + ) + + assert held == {"tools/replica-0": {first}, "tools/replica-1": set()} + # It agrees with the full inventory it stands in for. + inventory = ledger.gym_checkpoint_inventory({"tools/replica-0", "tools/replica-1"}) + assert held["tools/replica-0"] == set(inventory["tools/replica-0"]) & { + first, + second, + } + + +def test_wait_for_mutation_wakes_on_the_next_finished_mutation() -> None: + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + version = barrier.mutation_version + + waiter = asyncio.create_task(barrier.wait_for_mutation(version, 5.0)) + for _ in range(3): + await asyncio.sleep(0) + assert not waiter.done() + async with barrier.mutation(): + pass + for _ in range(3): + await asyncio.sleep(0) + assert waiter.done() + + # A mutation that finished before the wait began is not missed. + await asyncio.wait_for(barrier.wait_for_mutation(version, 5.0), 0.5) + # With nothing happening, the wait simply ends at its timeout. + await asyncio.wait_for( + barrier.wait_for_mutation(barrier.mutation_version, 0.05), 0.5 + ) + + asyncio.run(exercise()) diff --git a/tests/unit/experience/test_rollouts.py b/tests/unit/experience/test_rollouts.py index 502a8c6c3af..21c7f5b0542 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -55,6 +55,7 @@ from nemo_rl.experience.metric_utils import calculate_single_metric, pct from nemo_rl.experience.rollout_manager import ( AsyncNemoGymRolloutImpl, + RolloutDispatchAdmissionGate, RolloutTimeouts, ) from nemo_rl.experience.rollout_recovery import RecoveryGranularity @@ -2300,6 +2301,7 @@ def remote(self, inputs, timer_prefix, per_prompt=False): return _Stream() manager = object.__new__(AsyncNemoGymRolloutImpl) + manager._dispatch_admission_gate = RolloutDispatchAdmissionGate() manager._num_generations_per_prompt = 2 # These tests cover stream ordering/dedup, not deadlines or re-dispatch. manager._timeouts = RolloutTimeouts() @@ -2352,6 +2354,7 @@ def remote(self, inputs, timer_prefix, per_prompt=False): def test_nemo_gym_rollout_record_persists_runtime_resolved_agent_ref(): manager = object.__new__(AsyncNemoGymRolloutImpl) + manager._dispatch_admission_gate = None manager._num_generations_per_prompt = 2 manager._generation_config = { "temperature": 1.0, @@ -2371,8 +2374,11 @@ async def _run_rollouts( *, on_completion=None, recovery_granularity=RecoveryGranularity.SIBLING, + gym_instance_id=None, + dispatch_recorder=None, ): - del timer, timer_prefix, on_completion, recovery_granularity + del timer, timer_prefix, on_completion, recovery_granularity, gym_instance_id + del dispatch_recorder for row in inputs: row["agent_ref"] = resolved_agent_ref receipt_completion = SimpleNamespace(env_extras={"ng_receipt": {}}) @@ -2434,6 +2440,7 @@ def remote(self, inputs, timer_prefix, per_prompt=False): return _DuplicateStream() manager = object.__new__(AsyncNemoGymRolloutImpl) + manager._dispatch_admission_gate = None manager._num_generations_per_prompt = 2 # These tests cover stream ordering/dedup, not deadlines or re-dispatch. manager._timeouts = RolloutTimeouts() @@ -2480,6 +2487,7 @@ def test_rollout_manager_rotates_replicas_and_reports_group_share(): first, second = object(), object() selected = [] manager = object.__new__(AsyncNemoGymRolloutImpl) + manager._dispatch_admission_gate = None manager._timeouts = RolloutTimeouts() manager._max_gym_row_attempts = 1 manager._deadline_registry = None @@ -2503,8 +2511,16 @@ async def fake_stream_rows( timer_prefix, *, on_completion, + dispatch_recorder=None, + admission_closed_rows=None, ): - del total_rows, timer_prefix, on_completion + del ( + total_rows, + timer_prefix, + on_completion, + dispatch_recorder, + admission_closed_rows, + ) selected.append(environment) for row in pending: rowidx = row["_rowidx"] @@ -2570,6 +2586,7 @@ def remote(self, inputs, timer_prefix, per_prompt=False): return _FailedStream() manager = object.__new__(AsyncNemoGymRolloutImpl) + manager._dispatch_admission_gate = None manager._timeouts = RolloutTimeouts() manager._max_gym_row_attempts = 1 manager._deadline_registry = None diff --git a/tests/unit/models/generation/test_vllm_chat_template_wiring.py b/tests/unit/models/generation/test_vllm_chat_template_wiring.py index 82f065fceef..02a6440ffca 100644 --- a/tests/unit/models/generation/test_vllm_chat_template_wiring.py +++ b/tests/unit/models/generation/test_vllm_chat_template_wiring.py @@ -147,6 +147,10 @@ def placeholder(name): ServingTokenization=_ServingTokenization, ) module("vllm.renderers.online_renderer", OnlineRenderer=_OnlineRenderer) + module( + "vllm.sampling_params", + RequestOutputKind=type("RequestOutputKind", (), {"DELTA": "delta"}), + ) module( "vllm.exceptions", VLLMValidationError=type("VLLMValidationError", (Exception,), {}), diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 594412e7005..8938a4be39b 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -858,6 +858,10 @@ class ReasoningParserManager: ServingTokenization=ServingTokenization, ) make_module("vllm.exceptions", VLLMValidationError=VLLMValidationError) + make_module( + "vllm.sampling_params", + RequestOutputKind=type("RequestOutputKind", (), {"DELTA": "delta"}), + ) make_module( "vllm.reasoning.abs_reasoning_parsers", ReasoningParserManager=ReasoningParserManager, @@ -946,6 +950,7 @@ async def test_context_overflow_returns_http_400_for_nemo_gym(monkeypatch): worker.llm = MagicMock(model_config="model-config", renderer="renderer") worker._http_engine_client = worker.llm worker._capture_calls = {} + worker._capture_registry_lock = threading.Lock() worker.token_capture = None worker.llm_async_engine_args = MagicMock() worker.llm_async_engine_args.create_model_config.return_value = MagicMock( diff --git a/tests/unit/models/generation/test_vllm_token_capture_hosting.py b/tests/unit/models/generation/test_vllm_token_capture_hosting.py index 7337dae094b..093bc18cda4 100644 --- a/tests/unit/models/generation/test_vllm_token_capture_hosting.py +++ b/tests/unit/models/generation/test_vllm_token_capture_hosting.py @@ -23,6 +23,11 @@ from __future__ import annotations import asyncio +import hashlib +import logging +import threading +import time +from concurrent.futures import ThreadPoolExecutor from copy import deepcopy from types import SimpleNamespace from unittest.mock import MagicMock @@ -32,6 +37,10 @@ nemo_gym = pytest.importorskip("nemo_gym.token_id_capture.staging") +from nemo_gym._checkpoint.generation_cut import ( # noqa: E402 + GenerationCutInventory, + GenerationCutPrefix, +) from nemo_gym.token_id_capture.staging.capture import ( # noqa: E402 CaptureError, RolloutTokenCapture, @@ -43,9 +52,15 @@ ) from nemo_rl.data_plane.tq_token_sink import ChainPrefixCache # noqa: E402 +from nemo_rl.models.generation.generation_cut_capture import ( # noqa: E402 + _remaining_generation_limits_after_prefix, + _TokenCaptureSnapshotGate, +) from nemo_rl.models.generation.vllm.vllm_generation import VllmGeneration # noqa: E402 from nemo_rl.models.generation.vllm.vllm_worker_async import ( # noqa: E402 VllmAsyncGenerationWorkerImpl, + _classify_restored_prefix_terminal, + _RequestOutputDeltaAccumulator, ) pytestmark = pytest.mark.nemo_gym @@ -55,6 +70,9 @@ class _MemorySink: def __init__(self) -> None: self.records: list[StagedCallRecord] = [] self.attachments: list[dict | None] = [] + self.generation_prefix_records: list[tuple[str, StagedCallRecord]] = [] + self.generation_prefix_keys: list[str] = [] + self.cleared_generation_prefix_keys: list[str] = [] def stage( self, record: StagedCallRecord, *, attachments: dict | None = None @@ -63,6 +81,33 @@ def stage( self.attachments.append(attachments) return StageResult(ok=True, staging_key=record.staging_key) + def stage_generation_prefix( + self, + record: StagedCallRecord, + *, + checkpoint_id: str, + chunk_sequence: int, + attachments: dict | None = None, + ) -> StageResult: + assert attachments is None + self.generation_prefix_records.append((checkpoint_id, record)) + key = ( + f"__generation_cut__/{checkpoint_id}/{record.rollout_id}/" + f"{record.model_call_id}/{chunk_sequence}" + ) + self.generation_prefix_keys.append(key) + return StageResult(ok=True, staging_key=key) + + def clear(self, staging_keys: list[str]) -> None: + self.cleared_generation_prefix_keys.extend(staging_keys) + + +class _ByteTokenizer: + """Byte-level stand-in: each token ID is one UTF-8 byte of the output.""" + + def decode(self, token_ids: list[int], **kwargs) -> str: + return bytes(token_ids).decode("utf-8", errors="replace") + def _fake_worker(*, is_model_owner: bool = True) -> SimpleNamespace: """The attribute surface setup_token_capture touches, minus the engine.""" @@ -71,6 +116,19 @@ def _fake_worker(*, is_model_owner: bool = True) -> SimpleNamespace: token_capture=None, _rollout_weight_version=0, _chain_prefix=ChainPrefixCache(), + _capture_calls={}, + _capture_calls_by_model_call_id={}, + _completed_capture_calls={}, + _generation_cut_receipts={}, + _capture_registry_lock=threading.Lock(), + _capture_sink=None, + _generation_prefix_cuts_enabled=False, + _generation_cut_control_token=None, + _generation_cut_control_timeout_s=None, + _generation_cut_tokenizer=_ByteTokenizer(), + _token_capture_snapshot_gate=_TokenCaptureSnapshotGate(), + _staging_source=None, + _capture_media=False, ) worker.install_token_capture = lambda capture: setattr( worker, "token_capture", capture @@ -86,11 +144,9 @@ def test_setup_token_capture_installs_capture_with_vllm_adapter(monkeypatch): ) monkeypatch.setattr( "nemo_rl.data_plane.tq_token_sink.TQTokenSink", - lambda dp_client, - *, - staging_partition, - capture_media, - media_pixel_dtype=None: sink, + lambda dp_client, *, staging_partition, capture_media, media_pixel_dtype=None: ( + sink + ), ) worker = _fake_worker() @@ -129,11 +185,9 @@ def test_weight_version_is_stamped_from_worker_state(monkeypatch): ) monkeypatch.setattr( "nemo_rl.data_plane.tq_token_sink.TQTokenSink", - lambda dp_client, - *, - staging_partition, - capture_media, - media_pixel_dtype=None: sink, + lambda dp_client, *, staging_partition, capture_media, media_pixel_dtype=None: ( + sink + ), ) worker = _fake_worker() asyncio.run( @@ -174,12 +228,19 @@ def test_generation_setup_token_capture_fans_out(monkeypatch): "nemo_rl.models.generation.vllm.vllm_generation.ray.get", lambda futures: futures, ) - gen.setup_token_capture({"backend": "simple"}, "rollout_staging") + gen.setup_token_capture( + {"backend": "simple"}, + "rollout_staging", + generation_cut_control_timeout_s=60.0, + ) gen.worker_group.run_all_workers_single_data.assert_called_once_with( "setup_token_capture", dp_cfg={"backend": "simple"}, staging_partition="rollout_staging", capture_media=False, + generation_prefix_cuts_enabled=False, + generation_cut_control_token=None, + generation_cut_control_timeout_s=60.0, run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], ) @@ -190,6 +251,273 @@ def test_generation_setup_token_capture_requires_async_engine(): gen.setup_token_capture({}, "rollout_staging") +@pytest.mark.parametrize( + ("operation", "worker_method"), + [ + ( + "begin_token_capture_snapshot_fence", + "begin_token_capture_snapshot_fence_async", + ), + ( + "end_token_capture_snapshot_fence", + "end_token_capture_snapshot_fence_async", + ), + ], +) +def test_token_capture_snapshot_fence_control_fans_out( + monkeypatch, operation: str, worker_method: str +): + gen = _generation_with_mock_group() + gen.worker_group.workers = [object()] + gen.worker_group.run_all_workers_single_data.return_value = [True] + monkeypatch.setattr( + "nemo_rl.models.generation.vllm.vllm_generation.ray.get", + lambda futures, timeout=None: futures, + ) + + assert getattr(gen, operation)(timeout_s=17.0) + gen.worker_group.run_all_workers_single_data.assert_called_once_with( + worker_method, + run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], + ) + + +def test_prefix_capture_setup_requires_control_token_and_rejects_media(monkeypatch): + monkeypatch.setattr( + "nemo_rl.data_plane.build_data_plane_client", + lambda dp_cfg, bootstrap: MagicMock(), + ) + worker = _fake_worker() + with pytest.raises(ValueError, match="control bearer token"): + asyncio.run( + VllmAsyncGenerationWorkerImpl.setup_token_capture( + worker, + {}, + "rollout_staging", + generation_prefix_cuts_enabled=True, + ) + ) + + worker = _fake_worker() + worker.llm = SimpleNamespace(model_config=SimpleNamespace(dtype=torch.bfloat16)) + with pytest.raises(ValueError, match="does not yet support multimodal"): + asyncio.run( + VllmAsyncGenerationWorkerImpl.setup_token_capture( + worker, + {}, + "rollout_staging", + capture_media=True, + generation_prefix_cuts_enabled=True, + generation_cut_control_token="secret", + ) + ) + + +@pytest.mark.parametrize( + ("max_tokens", "min_tokens", "prefix_token_count", "expected"), + [ + (10, 6, 3, (7, 3)), + (None, 6, 3, (None, 3)), + (10, None, 3, (7, None)), + (10, 2, 3, (7, 0)), + ], +) +def test_restored_prefix_reduces_remaining_output_limits( + max_tokens, min_tokens, prefix_token_count, expected +): + assert ( + _remaining_generation_limits_after_prefix( + max_tokens=max_tokens, + min_tokens=min_tokens, + prefix_token_count=prefix_token_count, + ) + == expected + ) + + +def test_restored_prefix_at_output_limit_is_terminal(): + terminal = _classify_restored_prefix_terminal( + prompt_token_ids=[1, 2, 3, 4], + prefix_token_count=2, + requested_output_tokens=2, + model_max_tokens=8, + ) + assert terminal is not None + assert terminal.reason == "output_limit" + assert terminal.original_prompt_token_count == 2 + + +def test_request_output_deltas_are_assembled_without_mutating_inputs(): + def output(text, token_ids): + return SimpleNamespace( + outputs=[ + SimpleNamespace( + index=0, + text=text, + token_ids=list(token_ids), + logprobs=[f"lp-{token_id}" for token_id in token_ids], + ) + ] + ) + + first = output("a", [1]) + second = output("bc", [2, 3]) + accumulator = _RequestOutputDeltaAccumulator() + accumulator.append(first) + accumulator.append(second) + accumulated = accumulator.build() + assert accumulated.outputs[0].text == "abc" + assert accumulated.outputs[0].token_ids == [1, 2, 3] + assert first.outputs[0].token_ids == [1] + assert second.outputs[0].token_ids == [2, 3] + + +def test_token_capture_snapshot_fence_does_not_pause_decoding(): + class FakeLLM: + pause_calls = 0 + + async def pause_generation(self, **kwargs): + self.pause_calls += 1 + + async def scenario(): + worker = object.__new__(VllmAsyncGenerationWorkerImpl) + worker.cfg = {"vllm_cfg": {"async_engine": True}} + worker.llm = FakeLLM() + worker._token_capture_snapshot_gate = _TokenCaptureSnapshotGate() + worker._token_capture_fence_executor = ThreadPoolExecutor(max_workers=1) + try: + assert await worker.begin_token_capture_snapshot_fence_async() + assert worker.llm.pause_calls == 0 + entered = threading.Event() + released = threading.Event() + + def terminal_write(): + entered.set() + worker._token_capture_snapshot_gate.enter() + released.set() + worker._token_capture_snapshot_gate.exit() + + thread = threading.Thread(target=terminal_write) + thread.start() + assert entered.wait(timeout=5) + assert not released.wait(timeout=0.05) + assert await worker.end_token_capture_snapshot_fence_async() + assert released.wait(timeout=5) + thread.join(timeout=5) + finally: + worker._token_capture_snapshot_gate.reopen() + worker._token_capture_fence_executor.shutdown() + + asyncio.run(scenario()) + + +def _gate_admits_write(gate: _TokenCaptureSnapshotGate, *, timeout: float) -> bool: + """Whether a terminal write passes the gate within ``timeout`` seconds.""" + admitted = threading.Event() + + def terminal_write(): + gate.enter() + admitted.set() + gate.exit() + + threading.Thread(target=terminal_write, daemon=True).start() + return admitted.wait(timeout=timeout) + + +def test_snapshot_gate_close_after_its_release_is_a_no_op(): + gate = _TokenCaptureSnapshotGate() + epoch = gate.begin_epoch() + # The driver timed out and released before the worker ran the close. + gate.reopen() + gate.close_and_wait(epoch) + assert _gate_admits_write(gate, timeout=5) + + # A later fence still closes the gate normally. + gate.close_and_wait(gate.begin_epoch()) + assert not _gate_admits_write(gate, timeout=0.05) + gate.reopen() + + +def test_snapshot_gate_release_unblocks_a_close_still_draining(): + gate = _TokenCaptureSnapshotGate() + gate.enter() + closer = threading.Thread( + target=gate.close_and_wait, args=(gate.begin_epoch(),), daemon=True + ) + closer.start() + closer.join(timeout=0.05) + assert closer.is_alive() + gate.reopen() + closer.join(timeout=5) + assert not closer.is_alive() + gate.exit() + assert _gate_admits_write(gate, timeout=5) + + +def _fence_worker() -> VllmAsyncGenerationWorkerImpl: + worker = object.__new__(VllmAsyncGenerationWorkerImpl) + worker.cfg = {"vllm_cfg": {"async_engine": True}} + worker._token_capture_snapshot_gate = _TokenCaptureSnapshotGate() + worker._generation_cut_control_executor = ThreadPoolExecutor(max_workers=1) + worker._token_capture_fence_executor = ThreadPoolExecutor(max_workers=1) + return worker + + +def _shutdown_fence_worker(worker: VllmAsyncGenerationWorkerImpl) -> None: + worker._token_capture_snapshot_gate.reopen() + worker._generation_cut_control_executor.shutdown() + worker._token_capture_fence_executor.shutdown() + + +def test_snapshot_fence_is_not_queued_behind_a_running_cut(): + async def scenario(): + worker = _fence_worker() + stale_cut_running = threading.Event() + release_stale_cut = threading.Event() + try: + stale_cut = asyncio.ensure_future( + worker._run_generation_cut_control( + lambda: (stale_cut_running.set(), release_stale_cut.wait(10)) + ) + ) + assert await asyncio.to_thread(stale_cut_running.wait, 5) + await asyncio.wait_for( + worker.begin_token_capture_snapshot_fence_async(), timeout=5 + ) + assert not stale_cut.done() + release_stale_cut.set() + await stale_cut + finally: + release_stale_cut.set() + _shutdown_fence_worker(worker) + + asyncio.run(scenario()) + + +def test_snapshot_fence_closed_after_its_release_leaves_the_gate_open(): + async def scenario(): + worker = _fence_worker() + release_fence_thread = threading.Event() + try: + # Hold the fence thread so the close is still queued when the + # driver gives up and releases the fence. + worker._token_capture_fence_executor.submit(release_fence_thread.wait, 10) + begin = asyncio.ensure_future( + worker.begin_token_capture_snapshot_fence_async() + ) + await asyncio.sleep(0) + assert not begin.done() + assert await worker.end_token_capture_snapshot_fence_async() + release_fence_thread.set() + assert await asyncio.wait_for(begin, timeout=5) + assert _gate_admits_write(worker._token_capture_snapshot_gate, timeout=5) + finally: + release_fence_thread.set() + _shutdown_fence_worker(worker) + + asyncio.run(scenario()) + + def test_generation_set_rollout_weight_version_fans_out(monkeypatch): gen = _generation_with_mock_group() monkeypatch.setattr( @@ -217,8 +545,8 @@ def _worker_with_capture(sink: _MemorySink): from nemo_gym.token_id_capture.adapters.vllm import VLLMCaptureAdapter worker = _fake_worker() - worker._capture_calls = {} worker._chain_prefix = ChainPrefixCache() + worker._capture_sink = sink worker._delta_align_routed_experts = ( VllmAsyncGenerationWorkerImpl._delta_align_routed_experts ) @@ -226,7 +554,20 @@ def _worker_with_capture(sink: _MemorySink): "_fetch_chain_prefix", "_capture_admission", "_resolve_admission_prefix", + "_fetch_generation_cut_chunks", + "_rebuild_generation_cut_snapshot", + "_resolve_generation_cut", "_enter_request_prefix", + "_pop_request_capture", + "_get_request_capture", + "_remember_completed_capture", + "_abort_request_capture", + "_finish_request_capture_after_snapshot_fence", + "_finish_request_capture_with_lifecycle_owned", + "_restore_response_prefix", + "_checkpoint_active_generation_cut", + "_completed_generation_cut_ack", + "_checkpoint_generation_cut", ): setattr( worker, name, getattr(VllmAsyncGenerationWorkerImpl, name).__get__(worker) @@ -240,14 +581,24 @@ def _worker_with_capture(sink: _MemorySink): class _MemoryPrefixSource: - def __init__(self, deltas: dict[str, list[int]]) -> None: + def __init__( + self, + deltas: dict[str, list[int]], + *, + records: dict[str, StagedCallRecord] | None = None, + ) -> None: self.deltas = deltas + self.records = records or {} self.calls: list[list[str]] = [] def fetch_prefix_token_ids(self, staging_keys: list[str]) -> list[int]: self.calls.append(list(staging_keys)) return [token for key in staging_keys for token in self.deltas[key]] + def fetch(self, staging_keys: list[str]) -> list[StagedCallRecord]: + self.calls.append(list(staging_keys)) + return [self.records[key] for key in staging_keys] + def _served_content(gen_ids, logprobs): return { @@ -327,6 +678,419 @@ def test_request_capture_round_trip_stages_and_rides_coords( assert worker._capture_calls == {} +def test_generation_cut_stages_prefix_then_terminal_row_replaces_it(): + sink = _MemorySink() + worker = _worker_with_capture(sink) + worker._generation_prefix_cuts_enabled = True + request = _FakeRequest( + ng_capture={ + "rollout_id": "r0", + "model_call_id": "c1", + "parent_call_id": None, + "prev_len": 0, + "mode": "text", + }, + stream=False, + ) + VllmAsyncGenerationWorkerImpl._begin_request_capture(worker, request, [10, 11]) + state = worker._capture_calls[id(request)] + with state.lock: + state.effective_output_limit = 8 + state.observe([20, 21], [-0.1, -0.2]) + + inventory = GenerationCutInventory.build( + checkpoint_id="checkpoint-1", + server_name="policy_model", + active_prefixes=[ + GenerationCutPrefix( + ticket_id="ticket-1", + rollout_id="r0", + attempt=0, + model_call_id="c1", + admitted_at=1.0, + ) + ], + ) + receipt = worker._checkpoint_generation_cut(inventory) + replayed = worker._checkpoint_generation_cut(inventory) + + assert replayed is receipt + assert len(sink.generation_prefix_records) == 1 + assert receipt.prefixes[0].disposition == "durable_prefix" + assert receipt.prefixes[0].prefix_token_count == 2 + assert sink.generation_prefix_records[0][1].token_ids_delta == [10, 11, 20, 21] + assert state.generation_cut_staging_keys == sink.generation_prefix_keys + + content = VllmAsyncGenerationWorkerImpl._finish_request_capture( + worker, + request, + _served_content([20, 21, 22], [-0.1, -0.2, -0.3]), + ) + assert content["ng_commit_coords"]["disposition"] == "staged" + assert sink.records[-1].token_ids_delta == [10, 11, 20, 21, 22] + assert sink.cleared_generation_prefix_keys == sink.generation_prefix_keys + assert worker._capture_calls == {} + assert worker._capture_calls_by_model_call_id == {} + + +def test_generation_cut_past_its_deadline_fails_uncut_calls_without_staging(): + sink = _MemorySink() + worker = _worker_with_capture(sink) + worker._generation_prefix_cuts_enabled = True + request = _FakeRequest( + ng_capture={"rollout_id": "r0", "model_call_id": "c1", "mode": "text"}, + stream=False, + ) + VllmAsyncGenerationWorkerImpl._begin_request_capture(worker, request, [10]) + state = worker._capture_calls[id(request)] + with state.lock: + state.effective_output_limit = 8 + state.observe([11, 12], [-0.1, -0.2]) + inventory = GenerationCutInventory.build( + checkpoint_id="checkpoint-1", + server_name="policy_model", + active_prefixes=[ + GenerationCutPrefix( + ticket_id="ticket-1", + rollout_id="r0", + attempt=0, + model_call_id="c1", + admitted_at=1.0, + ) + ], + ) + + # Gym has already stopped waiting: the queued cut must not stage rows. + receipt = worker._checkpoint_generation_cut( + inventory, deadline=time.monotonic() - 1.0 + ) + + receipt.validate_for(inventory) + assert [ack.disposition for ack in receipt.prefixes] == ["durable_failure"] + assert sink.generation_prefix_records == [] + assert state.frozen_buffer is None + assert state.active_buffer.generated_token_ids == [11, 12] + + +def _active_call_inventory( + worker, *calls: tuple[str, list[int]] +) -> tuple[GenerationCutInventory, list]: + """Admit one live call per (model_call_id, generated IDs) and inventory them.""" + states = [] + prefixes = [] + for index, (model_call_id, generated) in enumerate(calls): + request = _FakeRequest( + ng_capture={ + "rollout_id": f"r{index}", + "model_call_id": model_call_id, + "mode": "text", + }, + stream=False, + ) + VllmAsyncGenerationWorkerImpl._begin_request_capture(worker, request, [10]) + state = worker._capture_calls[id(request)] + with state.lock: + state.effective_output_limit = 16 + state.observe(generated, [-0.1] * len(generated)) + states.append(state) + prefixes.append( + GenerationCutPrefix( + ticket_id=f"ticket-{index}", + rollout_id=f"r{index}", + attempt=0, + model_call_id=model_call_id, + admitted_at=1.0, + ) + ) + inventory = GenerationCutInventory.build( + checkpoint_id="checkpoint-1", + server_name="policy_model", + active_prefixes=prefixes, + ) + return inventory, states + + +def test_generation_cut_refuses_a_prefix_ending_inside_a_character(): + sink = _MemorySink() + worker = _worker_with_capture(sink) + worker._generation_prefix_cuts_enabled = True + dinosaur = list("🦖".encode()) + inventory, (split, whole) = _active_call_inventory( + worker, + ("split", list(b"Sure ") + dinosaur[:2]), + ("whole", list(b"Sure ") + dinosaur), + ) + + receipt = worker._checkpoint_generation_cut(inventory) + + # vLLM would prime the resumed tail on half a character and echo the + # prompt into the response, so that call regenerates instead. + assert [ack.disposition for ack in receipt.prefixes] == [ + "durable_failure", + "durable_prefix", + ] + assert len(sink.generation_prefix_records) == 1 + assert split.frozen_buffer is None + assert split.active_buffer.generated_token_ids == list(b"Sure ") + dinosaur[:2] + assert whole.generation_cut_staging_keys == sink.generation_prefix_keys + + +def test_dropped_output_delta_keeps_its_call_uncuttable(): + sink = _MemorySink() + worker = _worker_with_capture(sink) + worker._generation_prefix_cuts_enabled = True + inventory, (poisoned, healthy) = _active_call_inventory( + worker, ("poisoned", [11, 12]), ("healthy", [13, 14]) + ) + with poisoned.lock: + poisoned.observe([15], []) + # A later valid delta must not hide the tokens the bad one dropped. + poisoned.observe([16], [-0.1]) + assert poisoned.observation_error is not None + + receipt = worker._checkpoint_generation_cut(inventory) + + assert [ack.disposition for ack in receipt.prefixes] == [ + "durable_failure", + "durable_prefix", + ] + assert poisoned.frozen_buffer is None + assert healthy.generation_cut_staging_keys == sink.generation_prefix_keys + + +class _FailOncePrefixSink(_MemorySink): + def __init__(self) -> None: + super().__init__() + self._fail_next_prefix = True + + def stage_generation_prefix(self, *args, **kwargs) -> StageResult: + if self._fail_next_prefix: + self._fail_next_prefix = False + raise RuntimeError("injected generation-prefix staging failure") + return super().stage_generation_prefix(*args, **kwargs) + + +def test_generation_cut_rolls_back_tokens_after_staging_failure(): + sink = _FailOncePrefixSink() + worker = _worker_with_capture(sink) + request = _FakeRequest( + ng_capture={ + "rollout_id": "r0", + "model_call_id": "c1", + "mode": "text", + }, + stream=False, + ) + VllmAsyncGenerationWorkerImpl._begin_request_capture(worker, request, [10]) + state = worker._capture_calls[id(request)] + with state.lock: + state.effective_output_limit = 8 + state.observe([11, 12], [-0.1, -0.2]) + inventory = GenerationCutInventory.build( + checkpoint_id="checkpoint-1", + server_name="policy_model", + active_prefixes=[ + GenerationCutPrefix( + ticket_id="ticket-1", + rollout_id="r0", + attempt=0, + model_call_id="c1", + admitted_at=1.0, + ) + ], + ) + + with pytest.raises(RuntimeError, match="injected generation-prefix"): + worker._checkpoint_generation_cut(inventory) + assert state.frozen_buffer is None + assert state.active_buffer.generated_token_ids == [11, 12] + assert state.sealed_generated_token_ids == [] + + receipt = worker._checkpoint_generation_cut(inventory) + assert receipt.prefixes[0].prefix_token_count == 2 + assert sink.generation_prefix_records[-1][1].token_ids_delta == [10, 11, 12] + + +def test_restored_generation_cut_is_extended_and_retired_on_completion(caplog): + caplog.set_level(logging.INFO) + sink = _MemorySink() + original_worker = _worker_with_capture(sink) + original_worker._rollout_weight_version = 7 + original_request = _FakeRequest( + ng_capture={ + "rollout_id": "r0", + "model_call_id": "c1", + "mode": "text", + }, + stream=False, + ) + VllmAsyncGenerationWorkerImpl._begin_request_capture( + original_worker, original_request, [10, 11] + ) + original_worker._capture_calls[id(original_request)].effective_output_limit = 128 + VllmAsyncGenerationWorkerImpl._observe_request_capture( + original_worker, + original_request, + SimpleNamespace( + outputs=[ + SimpleNamespace( + token_ids=[12, 13], + logprobs=[ + {12: SimpleNamespace(logprob=-0.1)}, + {13: SimpleNamespace(logprob=-0.2)}, + ], + ) + ] + ), + ) + inventory = GenerationCutInventory.build( + checkpoint_id="checkpoint-1", + server_name="policy_model", + active_prefixes=[ + GenerationCutPrefix( + ticket_id="ticket-1", + rollout_id="r0", + attempt=0, + model_call_id="c1", + admitted_at=1.0, + ) + ], + ) + receipt = original_worker._checkpoint_generation_cut(inventory) + (cut_key,) = receipt.prefixes[0].staging_keys + cut_record = sink.generation_prefix_records[-1][1] + + resumed_worker = _worker_with_capture(sink) + resumed_worker._rollout_weight_version = 9 + resumed_worker._staging_source = _MemoryPrefixSource( + {}, records={cut_key: cut_record} + ) + request = _FakeRequest( + ng_capture={ + "rollout_id": "r0-a1", + "model_call_id": "c2", + "mode": "text", + "generation_cut": { + "source_capture_key": "r0", + "source_model_call_id": "c1", + "staging_keys": [cut_key], + "prefix_token_count": 2, + "prefix_digest": cut_record.digest, + "effective_output_limit": receipt.prefixes[0].effective_output_limit, + }, + }, + stream=False, + ) + admission = resumed_worker._capture_admission(request) + bad_count_admission = admission.model_copy( + update={ + "generation_cut": admission.generation_cut.model_copy( + update={"prefix_token_count": 3} + ) + } + ) + with pytest.raises( + RuntimeError, + match="generation-cut token count mismatch: expected=3 actual=2", + ): + resumed_worker._resolve_generation_cut(bad_count_admission, []) + bad_digest_admission = admission.model_copy( + update={ + "generation_cut": admission.generation_cut.model_copy( + update={"prefix_digest": "0" * 64} + ) + } + ) + with pytest.raises(RuntimeError, match="generation-cut digest mismatch"): + resumed_worker._resolve_generation_cut(bad_digest_admission, []) + + with caplog.at_level(logging.INFO): + cut = resumed_worker._resolve_generation_cut(admission, []) + expected_sha = hashlib.sha256(b"12,13").hexdigest() + assert f"prefix_ids_sha256={expected_sha}" in caplog.text + VllmAsyncGenerationWorkerImpl._begin_request_capture( + resumed_worker, + request, + [10, 11, 12, 13], + admission=admission, + prefix_token_ids=[], + generation_cut=cut, + resumed_generation_token_ids=[12, 13], + ) + tokenizer = MagicMock() + tokenizer.decode.return_value = "partial" + # vLLM's processed tail: EOS and any matched stop string already stripped. + output = SimpleNamespace(token_ids=[14, 151645], text=" tail") + resumed_worker._restore_response_prefix( + request, + SimpleNamespace(outputs=[output]), + tokenizer=tokenizer, + ) + tokenizer.decode.assert_called_once_with( + [12, 13], skip_special_tokens=True, spaces_between_special_tokens=True + ) + assert output.text == "partial tail" + + # A second checkpoint taken before the resumed call finishes must retain + # the original cut key and append the newly generated tail. Gym v2 stores + # this complete key list in the replacement attempt's model record. + VllmAsyncGenerationWorkerImpl._observe_request_capture( + resumed_worker, + request, + SimpleNamespace( + outputs=[ + SimpleNamespace( + token_ids=[14], + logprobs=[{14: SimpleNamespace(logprob=-0.3)}], + ) + ] + ), + ) + second_receipt = resumed_worker._checkpoint_generation_cut( + GenerationCutInventory.build( + checkpoint_id="checkpoint-2", + server_name="policy_model", + active_prefixes=[ + GenerationCutPrefix( + ticket_id="ticket-2", + rollout_id="r0", + attempt=1, + model_call_id="c2", + admitted_at=2.0, + ) + ], + ) + ) + assert second_receipt.prefixes[0].prefix_token_count == 3 + assert second_receipt.prefixes[0].staging_keys[0] == cut_key + assert len(second_receipt.prefixes[0].staging_keys) == 2 + + content = VllmAsyncGenerationWorkerImpl._finish_request_capture( + resumed_worker, + request, + _served_content([14], [-0.3]), + ) + + final_record = sink.records[-1] + assert content["ng_commit_coords"]["rollout_id"] == "r0-a1" + assert final_record.token_ids_delta == [10, 11, 12, 13, 14] + assert final_record.token_mask_delta == [0.0, 0.0, 1.0, 1.0, 1.0] + assert final_record.generation_log_probs_delta == [ + 0.0, + 0.0, + -0.1, + -0.2, + -0.3, + ] + assert final_record.weight_version == 7 + assert content["ng_commit_coords"]["weight_version"] == 7 + assert resumed_worker._completed_capture_calls["c2"].generation_token_count == 3 + assert sink.cleared_generation_prefix_keys == list( + second_receipt.prefixes[0].staging_keys + ) + assert "generation prefix restored:" in caplog.text + + def test_request_capture_token_in_prev_len_chains(): sink = _MemorySink() worker = _worker_with_capture(sink) diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index a1de4318507..cb34fe1a57e 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -33,6 +33,7 @@ import threading from collections import deque from collections.abc import Callable +from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace @@ -56,12 +57,33 @@ from nemo_rl.data.interfaces import DatumSpec from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.experience.rollout_manager import RolloutOutcome +from nemo_rl.environments.gym_checkpoint_adapter import ( + GymCheckpointCommitSummary, + GymCheckpointEpisode, +) +from nemo_rl.environments.gym_checkpoint_coordinator import ( + GYM_CHECKPOINT_SCHEMA_VERSION, + GymCheckpointCommitResult, + GymCheckpointManifest, + GymCheckpointNotReady, + load_gym_checkpoint_manifest, + staging_keys_of, + write_gym_checkpoint_manifest, +) +from nemo_rl.experience.rollout_manager import ( + RequestDeadlineRegistry, + RolloutDispatchAdmissionGate, + RolloutOutcome, + _Deadline, +) from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, ROLLOUT_RECOVERY_STATE_FILENAME, PromptGroupPhase, + PromptGroupRecoveryRecord, RecoveryGranularity, + RecoveryTargetLevel, + RolloutAttemptStatus, RolloutRecoveryLedger, build_rollout_recovery_state, ) @@ -250,6 +272,172 @@ def record_recovery_siblings(self, *, reused: int, redispatched: int) -> None: del reused, redispatched +class _RestoringGymCoordinator: + def __init__(self) -> None: + self.instance_ids = frozenset({"tools/replica-0"}) + self.restores: list[tuple[str, Path, GymCheckpointManifest]] = [] + self.discarded: list[tuple[str, GymCheckpointManifest]] = [] + self.resumes: list[str] = [] + + async def restore( + self, + restore_id: str, + checkpoint_root: str | Path, + manifest: GymCheckpointManifest, + ) -> None: + self.restores.append((restore_id, Path(checkpoint_root), manifest)) + + async def resume(self, checkpoint_id: str) -> None: + self.resumes.append(checkpoint_id) + + async def retire_restored( + self, + restore_id: str, + manifest: GymCheckpointManifest, + ) -> None: + self.discarded.append((restore_id, manifest)) + + +def _deadline_pausing_manager() -> SimpleNamespace: + """Rollout-manager stand-in for the deadline pause a Gym checkpoint takes.""" + registry = RequestDeadlineRegistry() + return SimpleNamespace( + suspend_request_deadlines=registry.suspend, + resume_request_deadlines=registry.resume, + ) + + +class _SavingGymCoordinator: + def __init__( + self, + *, + exported_episodes: tuple[GymCheckpointEpisode, ...] | None = None, + staging_keys_by_episode: dict[GymCheckpointEpisode, tuple[str, ...]] + | None = None, + unowned_live_episodes: tuple[GymCheckpointEpisode, ...] = (), + ) -> None: + self.instance_ids = frozenset({"tools/replica-0"}) + self.control_timeout_s = 1.0 + self.exported_episodes = exported_episodes + self.staging_keys_by_episode = staging_keys_by_episode or {} + self.unowned_live_episodes = unowned_live_episodes + self.events: list[object] = [] + self.checkpoint_open = False + + @asynccontextmanager + async def prepared(self, checkpoint_id: str): + self.events.append(("prepare", checkpoint_id)) + self.checkpoint_open = True + try: + yield + finally: + self.checkpoint_open = False + self.events.append(("resume", checkpoint_id)) + + async def commit( + self, + checkpoint_id: str, + checkpoint_root: Path, + inventory: dict[str, tuple[GymCheckpointEpisode, ...]], + ) -> GymCheckpointCommitResult: + self.events.append(("commit", checkpoint_id, checkpoint_root, inventory)) + exported = ( + inventory["tools/replica-0"] + if self.exported_episodes is None + else self.exported_episodes + ) + # Like Gym's model participant: the rows of every candidate in scope. + summary = GymCheckpointCommitSummary( + exported_episodes=exported, + staging_keys=tuple( + sorted( + {k for keys in self.staging_keys_by_episode.values() for k in keys} + ) + ), + participants=(), + unowned_live_episodes=self.unowned_live_episodes, + staging_keys_by_episode=self.staging_keys_by_episode, + ) + return GymCheckpointCommitResult( + manifest=GymCheckpointManifest( + schema_version=GYM_CHECKPOINT_SCHEMA_VERSION, + checkpoint_id=checkpoint_id, + instances={"tools/replica-0": summary.exported_episodes}, + staging_keys={"tools/replica-0": staging_keys_of(summary, exported)}, + participants={"tools/replica-0": ()}, + ), + instances={"tools/replica-0": summary}, + ) + + +class _BlockingCommitGymCoordinator(_SavingGymCoordinator): + """Hold commit open so a completion callback can mutate the ledger.""" + + def __init__(self) -> None: + super().__init__(exported_episodes=()) + self.commit_started = asyncio.Event() + self.release_commit = asyncio.Event() + + async def commit( + self, + checkpoint_id: str, + checkpoint_root: Path, + inventory: dict[str, tuple[GymCheckpointEpisode, ...]], + ) -> GymCheckpointCommitResult: + self.commit_started.set() + await self.release_commit.wait() + return await super().commit(checkpoint_id, checkpoint_root, inventory) + + +class _LiveSessionGymCoordinator(_SavingGymCoordinator): + """Gym deployment that keeps every episode it admitted until RL retires it. + + Like the real deployment, a commit exports every live session, and the + coordinator rejects one that is outside RL's candidate set. + """ + + def __init__(self, live: set[GymCheckpointEpisode]) -> None: + super().__init__() + self.live = set(live) + + async def retire( + self, + checkpoint_id: str, + episodes_by_instance: dict[str, tuple[GymCheckpointEpisode, ...]], + ) -> None: + if self.checkpoint_open: + # Gym stops a retired episode by waiting for its final cleanup, + # whose calls wait for resume, so it refuses a retire mid-checkpoint. + raise RuntimeError("cannot retire while a checkpoint is open; resume first") + self.events.append( + ( + "retire", + checkpoint_id, + { + instance_id: frozenset(episodes) + for instance_id, episodes in episodes_by_instance.items() + }, + ) + ) + for episodes in episodes_by_instance.values(): + self.live.difference_update(episodes) + + async def commit( + self, + checkpoint_id: str, + checkpoint_root: Path, + inventory: dict[str, tuple[GymCheckpointEpisode, ...]], + ) -> GymCheckpointCommitResult: + unexpected = self.live - set(inventory["tools/replica-0"]) + if unexpected: + raise RuntimeError( + "Gym checkpoint actor exported an episode outside the " + f"candidate set: unexpected={sorted(unexpected, key=repr)!r}" + ) + self.exported_episodes = tuple(sorted(self.live, key=repr)) + return await super().commit(checkpoint_id, checkpoint_root, inventory) + + class _BlockingRolloutManager: """Hold one admitted rollout unfinished while the checkpoint is written.""" @@ -611,6 +799,493 @@ async def exercise() -> None: asyncio.run(exercise()) +def test_turn_checkpoint_drains_on_wire_reply_before_tq_cut(tmp_path: Path) -> None: + async def exercise() -> None: + ledger = RolloutRecoveryLedger() + barrier = DataPlaneCheckpointBarrier() + async with barrier.mutation() as cut: + group = ledger.reserve_group( + cut, + group_id="batch-7-prompt-0", + admission_id="batch-7", + prompt_id="70", + prompt_payload={"idx": 70, "message_log": []}, + expected_generations=1, + target_step=7, + start_weight_version=7, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + ledger.mark_group_dispatched( + cut, + group.group_id, + gym_instance_id="tools/replica-0", + ) + + coordinator = _SavingGymCoordinator(exported_episodes=()) + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._gym_checkpoint_coordinator = coordinator + controller._rollout_dispatch_admission_gate = RolloutDispatchAdmissionGate() + controller._rollout_manager = _deadline_pausing_manager() + controller._rollout_recovery_ledger = ledger + controller._data_plane_checkpoint_barrier = barrier + controller._generation_prefix_cuts_enabled = False + + async def complete_on_wire_reply() -> None: + while not any( + isinstance(event, tuple) and event[0] == "commit" + for event in coordinator.events + ): + await asyncio.sleep(0) + async with barrier.mutation("sibling_seals") as cut: + ledger.mark_sibling_sealed( + cut, + group.group_id, + generation_index=0, + gate_rollout_id=group.gate_rollout_id(0), + receipt={ + "rollout_id": group.gate_rollout_id(0), + "manifest": [{"staging_key": "on-wire/call-0"}], + }, + reward=1.0, + mask_sample=False, + ) + + completion = asyncio.create_task(complete_on_wire_reply()) + async with controller._prepared_gym_checkpoint( + tmp_path, + checkpoint_id="save-1", + ) as gym_commit: + assert gym_commit is not None + assert gym_commit.manifest.instances == {"tools/replica-0": ()} + async with barrier.checkpoint() as cut: + await controller._validate_gym_checkpoint_cut( + cut, gym_commit, checkpoint_path=tmp_path + ) + assert ledger.expected_staging_keys() == {"on-wire/call-0"} + coordinator.events.append("tq-cut") + await completion + + assert coordinator.events == [ + ("prepare", "save-1"), + ( + "commit", + "save-1", + tmp_path, + {"tools/replica-0": (GymCheckpointEpisode(*group.gym_episode(0)),)}, + ), + "tq-cut", + ("resume", "save-1"), + ] + + asyncio.run(exercise()) + + +def test_completion_callback_can_run_while_gym_commit_is_pending( + tmp_path: Path, +) -> None: + async def exercise() -> None: + ledger = RolloutRecoveryLedger() + barrier = DataPlaneCheckpointBarrier() + async with barrier.mutation() as cut: + group = ledger.reserve_group( + cut, + group_id="batch-7-prompt-0", + admission_id="batch-7", + prompt_id="70", + prompt_payload={"idx": 70, "message_log": []}, + expected_generations=1, + target_step=7, + start_weight_version=7, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + ledger.mark_group_dispatched( + cut, + group.group_id, + gym_instance_id="tools/replica-0", + ) + + coordinator = _BlockingCommitGymCoordinator() + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._gym_checkpoint_coordinator = coordinator + controller._rollout_dispatch_admission_gate = RolloutDispatchAdmissionGate() + controller._rollout_manager = _deadline_pausing_manager() + controller._rollout_recovery_ledger = ledger + controller._data_plane_checkpoint_barrier = barrier + controller._generation_prefix_cuts_enabled = False + + async def checkpoint() -> GymCheckpointCommitResult: + async with controller._prepared_gym_checkpoint( + tmp_path, + checkpoint_id="save-1", + ) as gym_commit: + assert gym_commit is not None + return gym_commit + + checkpoint_task = asyncio.create_task(checkpoint()) + await asyncio.wait_for( + coordinator.commit_started.wait(), + timeout=_ASYNC_TEST_TIMEOUT_S, + ) + + async def complete_rollout() -> None: + async with barrier.mutation("sibling_seals") as cut: + ledger.mark_sibling_sealed( + cut, + group.group_id, + generation_index=0, + gate_rollout_id=group.gate_rollout_id(0), + receipt={ + "rollout_id": group.gate_rollout_id(0), + "manifest": [{"staging_key": "on-wire/call-0"}], + }, + reward=1.0, + mask_sample=False, + ) + + await asyncio.wait_for( + complete_rollout(), + timeout=_ASYNC_TEST_TIMEOUT_S, + ) + + coordinator.release_commit.set() + commit = await asyncio.wait_for( + checkpoint_task, + timeout=_ASYNC_TEST_TIMEOUT_S, + ) + + assert commit.manifest.instances == {"tools/replica-0": ()} + assert ledger.expected_staging_keys() == {"on-wire/call-0"} + assert coordinator.events[-1] == ("resume", "save-1") + + asyncio.run(exercise()) + + +def test_rollout_dispatch_gate_blocks_only_new_rpc_submissions() -> None: + async def exercise() -> None: + gate = RolloutDispatchAdmissionGate() + first_entered = asyncio.Event() + release_first = asyncio.Event() + checkpoint_closed = asyncio.Event() + release_checkpoint = asyncio.Event() + second_entered = asyncio.Event() + + async def first_submission() -> None: + async with gate.admission(): + first_entered.set() + await release_first.wait() + + async def checkpoint() -> None: + async with gate.closed(): + checkpoint_closed.set() + await release_checkpoint.wait() + + async def second_submission() -> None: + async with gate.admission(): + second_entered.set() + + first = asyncio.create_task(first_submission()) + await first_entered.wait() + closing = asyncio.create_task(checkpoint()) + await asyncio.sleep(0) + assert not checkpoint_closed.is_set() + + release_first.set() + await checkpoint_closed.wait() + second = asyncio.create_task(second_submission()) + await asyncio.sleep(0) + assert not second_entered.is_set() + + release_checkpoint.set() + await asyncio.gather(first, closing, second) + assert second_entered.is_set() + + asyncio.run(exercise()) + + +def test_turn_checkpoint_aborts_when_candidate_is_neither_exported_nor_drained( + tmp_path: Path, +) -> None: + async def exercise() -> None: + ledger = RolloutRecoveryLedger() + async with DataPlaneCheckpointBarrier().mutation() as cut: + group = ledger.reserve_group( + cut, + group_id="batch-7-prompt-0", + admission_id="batch-7", + prompt_id="70", + prompt_payload={"idx": 70, "message_log": []}, + expected_generations=1, + target_step=7, + start_weight_version=7, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + ledger.mark_group_dispatched( + cut, + group.group_id, + gym_instance_id="tools/replica-0", + ) + + coordinator = _SavingGymCoordinator(exported_episodes=()) + coordinator.control_timeout_s = 0.01 + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._gym_checkpoint_coordinator = coordinator + controller._rollout_dispatch_admission_gate = RolloutDispatchAdmissionGate() + controller._rollout_manager = _deadline_pausing_manager() + controller._rollout_recovery_ledger = ledger + controller._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + controller._generation_prefix_cuts_enabled = False + + with pytest.raises(TimeoutError, match="neither|classify every candidate"): + async with controller._prepared_gym_checkpoint( + tmp_path, + checkpoint_id="save-1", + ): + raise AssertionError("unreachable") + + assert coordinator.events[-1] == ("resume", "save-1") + + asyncio.run(exercise()) + + +def test_turn_checkpoint_fails_fast_on_parked_session_without_an_owner( + tmp_path: Path, +) -> None: + """A session Gym parks but cannot export never drains; do not wait for it.""" + + async def exercise() -> None: + ledger = RolloutRecoveryLedger() + async with DataPlaneCheckpointBarrier().mutation() as cut: + group = ledger.reserve_group( + cut, + group_id="batch-7-prompt-0", + admission_id="batch-7", + prompt_id="70", + prompt_payload={"idx": 70, "message_log": []}, + expected_generations=1, + target_step=7, + start_weight_version=7, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + ledger.mark_group_dispatched( + cut, + group.group_id, + gym_instance_id="tools/replica-0", + ) + + (stranded,) = ledger.gym_checkpoint_inventory(frozenset({"tools/replica-0"}))[ + "tools/replica-0" + ] + coordinator = _SavingGymCoordinator( + exported_episodes=(), + unowned_live_episodes=(GymCheckpointEpisode(*stranded),), + ) + coordinator.control_timeout_s = 60.0 + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._gym_checkpoint_coordinator = coordinator + controller._rollout_dispatch_admission_gate = RolloutDispatchAdmissionGate() + controller._rollout_manager = _deadline_pausing_manager() + controller._rollout_recovery_ledger = ledger + controller._generation_prefix_cuts_enabled = False + + async def checkpoint() -> None: + async with controller._prepared_gym_checkpoint( + tmp_path, + checkpoint_id="save-1", + ): + raise AssertionError("unreachable") + + with pytest.raises(RuntimeError, match="parked sessions"): + await asyncio.wait_for(checkpoint(), timeout=5.0) + + assert coordinator.events[-1] == ("resume", "save-1") + + asyncio.run(exercise()) + + +def _dispatched_two_sibling_group( + ledger: RolloutRecoveryLedger, cut: DataPlaneMutationCut +) -> PromptGroupRecoveryRecord: + group = ledger.reserve_group( + cut, + group_id="batch-7-prompt-0", + admission_id="batch-7", + prompt_id="70", + prompt_payload={"idx": 70, "message_log": []}, + expected_generations=2, + target_step=7, + start_weight_version=7, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + ledger.mark_group_dispatched(cut, group.group_id, gym_instance_id="tools/replica-0") + return ledger.get_group(group.group_id) + + +def _gym_episodes(group: PromptGroupRecoveryRecord) -> set[GymCheckpointEpisode]: + return { + GymCheckpointEpisode(*group.gym_episode(sibling.generation_index)) + for sibling in group.siblings + } + + +def _controller_for( + coordinator: _SavingGymCoordinator, ledger: RolloutRecoveryLedger +) -> Any: + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._gym_checkpoint_coordinator = coordinator + controller._rollout_dispatch_admission_gate = RolloutDispatchAdmissionGate() + controller._rollout_manager = _deadline_pausing_manager() + controller._rollout_recovery_ledger = ledger + controller._generation_prefix_cuts_enabled = False + return controller + + +def test_turn_checkpoint_retires_abandoned_attempts_gym_is_still_running( + tmp_path: Path, +) -> None: + """A cancelled group's /run keeps running in Gym; RL must retire it. + + Abandoning drops the attempts from RL's inventory, but nothing stops the + episodes Gym already admitted. Unretired, Gym exports them at the next + commit and every checkpoint fails until they finish on their own. + """ + + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + ledger = RolloutRecoveryLedger() + async with barrier.mutation() as cut: + group = _dispatched_two_sibling_group(ledger, cut) + abandoned = _gym_episodes(group) + # The group task was cancelled after Ray accepted its submission. + async with barrier.mutation() as cut: + ledger.abandon_unsealed(cut, group.group_id) + + coordinator = _LiveSessionGymCoordinator(live=abandoned) + controller = _controller_for(coordinator, ledger) + + async with controller._prepared_gym_checkpoint( + tmp_path, + checkpoint_id="save-1", + ) as gym_commit: + assert gym_commit is not None + assert gym_commit.manifest.instances == {"tools/replica-0": ()} + + # Retired while Gym is idle, before prepare parks anything. + assert [event[0] for event in coordinator.events] == [ + "retire", + "prepare", + "commit", + "resume", + ] + assert coordinator.events[0][2] == {"tools/replica-0": frozenset(abandoned)} + assert coordinator.live == set() + + asyncio.run(exercise()) + + +def test_turn_checkpoint_retires_only_the_abandoned_attempts_of_a_retried_group( + tmp_path: Path, +) -> None: + """The retry's fresh attempts stay live and are committed normally.""" + + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + ledger = RolloutRecoveryLedger() + async with barrier.mutation() as cut: + group = _dispatched_two_sibling_group(ledger, cut) + abandoned = _gym_episodes(group) + async with barrier.mutation() as cut: + ledger.abandon_unsealed(cut, group.group_id) + ledger.prepare_incomplete_retry(cut, group.group_id) + ledger.mark_group_dispatched( + cut, group.group_id, gym_instance_id="tools/replica-0" + ) + replacements = _gym_episodes(ledger.get_group(group.group_id)) + assert replacements.isdisjoint(abandoned) + + coordinator = _LiveSessionGymCoordinator(live=abandoned | replacements) + controller = _controller_for(coordinator, ledger) + + async with controller._prepared_gym_checkpoint( + tmp_path, + checkpoint_id="save-1", + ) as gym_commit: + assert gym_commit is not None + assert set(gym_commit.manifest.instances["tools/replica-0"]) == ( + replacements + ) + + assert coordinator.events[0][0] == "retire" + assert coordinator.events[0][2] == {"tools/replica-0": frozenset(abandoned)} + assert coordinator.live == replacements + + asyncio.run(exercise()) + + +def test_turn_checkpoint_rejects_missing_gym_staging_key(tmp_path: Path) -> None: + async def exercise() -> None: + ledger = RolloutRecoveryLedger() + barrier = DataPlaneCheckpointBarrier() + async with barrier.mutation() as cut: + group = ledger.reserve_group( + cut, + group_id="batch-7-prompt-0", + admission_id="batch-7", + prompt_id="70", + prompt_payload={"idx": 70, "message_log": []}, + expected_generations=1, + target_step=7, + start_weight_version=7, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + ledger.mark_group_dispatched( + cut, + group.group_id, + gym_instance_id="tools/replica-0", + ) + + episode = GymCheckpointEpisode(*group.gym_episode(0)) + coordinator = _SavingGymCoordinator( + staging_keys_by_episode={episode: ("missing-call-row",)} + ) + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._gym_checkpoint_coordinator = coordinator + controller._rollout_dispatch_admission_gate = RolloutDispatchAdmissionGate() + controller._rollout_manager = _deadline_pausing_manager() + controller._rollout_recovery_ledger = ledger + controller._generation_prefix_cuts_enabled = False + controller._dp_client = NoOpDataPlaneClient() + controller._master_config = SimpleNamespace( + token_capture=SimpleNamespace(staging_partition="rollout_staging") + ) + + with pytest.raises(RuntimeError, match="staging rows missing"): + async with controller._prepared_gym_checkpoint( + tmp_path, + checkpoint_id="save-1", + ) as gym_commit: + assert gym_commit is not None + async with barrier.checkpoint() as cut: + await controller._validate_gym_checkpoint_cut( + cut, gym_commit, checkpoint_path=tmp_path + ) + + assert coordinator.events[-1] == ("resume", "save-1") + + asyncio.run(exercise()) + + def test_checkpoint_owns_batch_7_while_its_rollout_is_unfinished(tmp_path) -> None: """A finalized checkpoint cannot contain a cursor hole for target step 7.""" @@ -984,6 +1659,219 @@ async def _recover( asyncio.run(exercise()) +def test_turn_recovery_restores_gym_before_redispatch(tmp_path: Path) -> None: + async def exercise() -> None: + dataset_prompt: DatumSpec = { + "idx": 70, + "message_log": [], + "length": 1, + "extra_env_info": None, + "loss_multiplier": 1.0, + } + prompt_batch = rl_collate_fn([dataset_prompt]) + dispatched_prompt = cast( + DatumSpec, + {key: value[0] for key, value in prompt_batch.items()}, + ) + saved_ledger = RolloutRecoveryLedger() + async with DataPlaneCheckpointBarrier().mutation() as cut: + saved_group = saved_ledger.reserve_group( + cut, + group_id="batch-7-prompt-0", + admission_id="batch-7", + prompt_id="70", + prompt_payload=dispatched_prompt, + expected_generations=1, + target_step=7, + start_weight_version=7, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + saved_ledger.mark_group_dispatched( + cut, + saved_group.group_id, + gym_instance_id="tools/replica-0", + ) + rollout_id, attempt = saved_group.gym_episode(0) + saved_state = build_rollout_recovery_state( + saved_ledger, + batch_shortfall={}, + sampler_stamps_target_steps=True, + ) + recovery_path = tmp_path / ROLLOUT_RECOVERY_STATE_FILENAME + torch.save(saved_state, recovery_path) + payload_sha256 = hashlib.sha256(recovery_path.read_bytes()).hexdigest() + manifest = GymCheckpointManifest( + schema_version=GYM_CHECKPOINT_SCHEMA_VERSION, + checkpoint_id="save-1", + instances={ + "tools/replica-0": (GymCheckpointEpisode(rollout_id, attempt),), + }, + staging_keys={"tools/replica-0": ()}, + participants={"tools/replica-0": ()}, + ) + write_gym_checkpoint_manifest(tmp_path, manifest) + + rollout_manager = _RecoveryRolloutManager(RolloutRecoveryLedger()) + gym_coordinator = _RestoringGymCoordinator() + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + _init_recovery_telemetry(controller, train_steps=7) + controller._sampler = _CountingInOrderSampler() + controller._rollout_manager = rollout_manager + controller._gym_checkpoint_coordinator = gym_coordinator + 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, + "rollout_recovery_payload_sha256": payload_sha256, + "rollout_recovery_group_count": 1, + } + controller._async_cfg = SimpleNamespace( + max_buffered_rollouts=4, + max_inflight_prompts=2, + ) + controller._buffer_capacity = asyncio.Semaphore(4) + controller._trainer_version = 7 + controller._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + controller._dataloader = SimpleNamespace( + dataset={70: dataset_prompt}, + collate_fn=rl_collate_fn, + ) + controller._buffer = SimpleNamespace( + count_for_target_step=lambda _target_step: 0, + metadata_state_dict=lambda *, saved_capacity: { + "groups": [], + "saved_capacity": saved_capacity, + }, + ) + + await controller._maybe_restore_rollout_recovery(restored_replay_groups=0) + + restored_group = rollout_manager.recovery_ledger.get_group("batch-7-prompt-0") + restored_attempt = restored_group.siblings[0].current_attempt + assert restored_attempt.status is RolloutAttemptStatus.RESERVED + assert restored_attempt.gym_instance_id == "tools/replica-0" + assert restored_attempt.gym_attempt == 1 + assert len(gym_coordinator.restores) == 1 + restore_id, checkpoint_root, restored_manifest = gym_coordinator.restores[0] + assert restore_id.startswith("restore-") + assert checkpoint_root == tmp_path + assert restored_manifest == manifest + assert gym_coordinator.resumes == [restore_id] + + asyncio.run(exercise()) + + +def test_turn_checkpoint_commits_owned_gym_episode_before_tq_cut( + tmp_path: Path, +) -> None: + async def exercise() -> None: + ledger = RolloutRecoveryLedger() + async with DataPlaneCheckpointBarrier().mutation() as cut: + group = ledger.reserve_group( + cut, + group_id="batch-7-prompt-0", + admission_id="batch-7", + prompt_id="70", + prompt_payload={"idx": 70, "message_log": []}, + expected_generations=1, + target_step=7, + start_weight_version=7, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + ledger.mark_group_dispatched( + cut, + group.group_id, + gym_instance_id="tools/replica-0", + ) + + coordinator = _SavingGymCoordinator() + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._gym_checkpoint_coordinator = coordinator + controller._rollout_dispatch_admission_gate = RolloutDispatchAdmissionGate() + controller._rollout_manager = _deadline_pausing_manager() + controller._rollout_recovery_ledger = ledger + controller._generation_prefix_cuts_enabled = False + + async with controller._prepared_gym_checkpoint( + tmp_path, + checkpoint_id="save-1", + ): + async with DataPlaneCheckpointBarrier().checkpoint(): + coordinator.events.append("tq-cut") + + rollout_id, attempt = group.gym_episode(0) + assert coordinator.events == [ + ("prepare", "save-1"), + ( + "commit", + "save-1", + tmp_path, + {"tools/replica-0": (GymCheckpointEpisode(rollout_id, attempt),)}, + ), + "tq-cut", + ("resume", "save-1"), + ] + + asyncio.run(exercise()) + + +def test_prefix_checkpoint_fences_terminal_writes_through_tq_cut( + tmp_path: Path, +) -> None: + async def exercise() -> None: + ledger = RolloutRecoveryLedger() + coordinator = _SavingGymCoordinator() + + class FakeGeneration: + def begin_token_capture_snapshot_fence(self, *, timeout_s=None): + coordinator.events.append(("generation-fence", timeout_s)) + return True + + def end_token_capture_snapshot_fence(self, *, timeout_s=None): + coordinator.events.append(("generation-release", timeout_s)) + return True + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._gym_checkpoint_coordinator = coordinator + controller._rollout_dispatch_admission_gate = RolloutDispatchAdmissionGate() + controller._rollout_manager = _deadline_pausing_manager() + controller._rollout_recovery_ledger = ledger + controller._generation_prefix_cuts_enabled = True + controller._gen = FakeGeneration() + controller._master_config = SimpleNamespace( + token_capture=SimpleNamespace(control_timeout_s=17.0) + ) + + async with controller._prepared_gym_checkpoint( + tmp_path, + checkpoint_id="save-prefix-1", + ): + coordinator.events.append("tq-cut") + + assert coordinator.events == [ + ("generation-fence", 17.0), + ("prepare", "save-prefix-1"), + ( + "commit", + "save-prefix-1", + tmp_path, + {"tools/replica-0": ()}, + ), + "tq-cut", + ("resume", "save-prefix-1"), + ("generation-release", 17.0), + ] + + asyncio.run(exercise()) + + def test_recovery_rejects_an_unhandled_phase_before_redispatch() -> None: """A future phase must fail loudly instead of remaining owned forever.""" @@ -1426,3 +2314,555 @@ def test_recovery_rejects_a_missing_advertised_ledger_sidecar(tmp_path) -> None: asyncio.run( controller._maybe_restore_rollout_recovery(restored_replay_groups=0) ) + + +class _NotReadyGymCoordinator(_SavingGymCoordinator): + @asynccontextmanager + async def prepared(self, checkpoint_id: str): + self.events.append(("prepare", checkpoint_id)) + raise GymCheckpointNotReady({"tools/replica-0": {"agent": ("rollout-1",)}}) + yield + + +def _gym_checkpoint_controller( + coordinator: _SavingGymCoordinator, registry: RequestDeadlineRegistry +) -> Any: + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._gym_checkpoint_coordinator = coordinator + controller._rollout_dispatch_admission_gate = RolloutDispatchAdmissionGate() + controller._rollout_recovery_ledger = RolloutRecoveryLedger() + controller._rollout_manager = SimpleNamespace( + suspend_request_deadlines=registry.suspend, + resume_request_deadlines=registry.resume, + ) + controller._generation_prefix_cuts_enabled = False + return controller + + +class TestGymCheckpointPausesRolloutDeadlines: + """Parked rollouts cannot progress, so a checkpoint must not use up their budget.""" + + def test_a_parked_rollout_outlives_a_checkpoint_longer_than_its_budget( + self, tmp_path: Path + ) -> None: + registry = RequestDeadlineRegistry() + controller = _gym_checkpoint_controller( + _SavingGymCoordinator(exported_episodes=()), registry + ) + + async def exercise() -> None: + entered = asyncio.Event() + checkpoint_done = asyncio.Event() + + async def parked_rollout() -> None: + async with _Deadline(0.2, "NeMo-Gym prompt group", registry=registry): + entered.set() + await checkpoint_done.wait() + await asyncio.sleep(0.05) + + rollout = asyncio.create_task(parked_rollout()) + await entered.wait() + async with controller._prepared_gym_checkpoint( + tmp_path, checkpoint_id="save-1" + ): + assert registry.suspended + # Twice the rollout's whole budget. + await asyncio.sleep(0.4) + checkpoint_done.set() + # Raises RolloutTimeout if the checkpoint was charged to the rollout. + await rollout + assert not registry.suspended + + asyncio.run(exercise()) + + def test_ending_a_checkpoint_keeps_a_generation_stand_down_pause( + self, tmp_path: Path + ) -> None: + registry = RequestDeadlineRegistry() + controller = _gym_checkpoint_controller( + _SavingGymCoordinator(exported_episodes=()), registry + ) + + async def exercise() -> None: + # The colocated engine is training; its own pause must outlast ours. + registry.suspend() + async with controller._prepared_gym_checkpoint( + tmp_path, checkpoint_id="save-1" + ): + pass + assert registry.suspended + registry.resume() + assert not registry.suspended + + asyncio.run(exercise()) + + def test_a_checkpoint_that_cannot_park_releases_its_pause( + self, tmp_path: Path + ) -> None: + registry = RequestDeadlineRegistry() + controller = _gym_checkpoint_controller(_NotReadyGymCoordinator(), registry) + + async def exercise() -> None: + with pytest.raises(GymCheckpointNotReady): + async with controller._prepared_gym_checkpoint( + tmp_path, checkpoint_id="save-1" + ): + raise AssertionError("unreachable") + assert not registry.suspended + + asyncio.run(exercise()) + + +class TestDrainScalesWithPendingRepliesNotTheLedger: + """The drain waits on a few on-wire replies; its cost must not track every rollout. + + At hundreds of thousands of live rollouts, rescanning the whole ledger on + each check blocks the event loop that delivers the very replies it waits + for, with every rollout parked behind the checkpoint. + """ + + @staticmethod + async def _dispatched_groups( + ledger: RolloutRecoveryLedger, + barrier: DataPlaneCheckpointBarrier, + count: int, + ) -> list[PromptGroupRecoveryRecord]: + groups = [] + async with barrier.mutation() as cut: + for index in range(count): + group = ledger.reserve_group( + cut, + group_id=f"batch-7-prompt-{index}", + admission_id="batch-7", + prompt_id=str(index), + prompt_payload={"idx": index, "message_log": []}, + expected_generations=1, + target_step=7, + start_weight_version=7, + restore_level=RecoveryTargetLevel.TURN, + admitted=True, + ) + ledger.mark_group_dispatched( + cut, group.group_id, gym_instance_id="tools/replica-0" + ) + groups.append(ledger.get_group(group.group_id)) + return groups + + @staticmethod + async def _drain_setup( + tmp_path: Path, + *, + rollouts: int, + on_wire: int, + control_timeout_s: float = 5.0, + ): + """Commit a cut in which Gym exported all but the first ``on_wire`` episodes.""" + ledger = RolloutRecoveryLedger() + barrier = DataPlaneCheckpointBarrier() + groups = await TestDrainScalesWithPendingRepliesNotTheLedger._dispatched_groups( + ledger, barrier, rollouts + ) + episodes = [GymCheckpointEpisode(*group.gym_episode(0)) for group in groups] + coordinator = _SavingGymCoordinator(exported_episodes=tuple(episodes[on_wire:])) + coordinator.control_timeout_s = control_timeout_s + candidates = {"tools/replica-0": tuple(episodes)} + commit = await coordinator.commit("save-1", tmp_path, candidates) + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._rollout_recovery_ledger = ledger + controller._data_plane_checkpoint_barrier = barrier + return controller, coordinator, candidates, commit, ledger, barrier, groups + + @staticmethod + async def _resolve( + ledger: RolloutRecoveryLedger, + barrier: DataPlaneCheckpointBarrier, + group: PromptGroupRecoveryRecord, + how: str, + ) -> None: + async with barrier.mutation("sibling_seals") as cut: + if how == "sealed": + ledger.mark_sibling_sealed( + cut, + group.group_id, + generation_index=0, + gate_rollout_id=group.gate_rollout_id(0), + receipt=None, + reward=1.0, + mask_sample=False, + ) + else: + ledger.abandon_unsealed(cut, group.group_id) + + def test_waiting_never_rescans_every_rollout(self, tmp_path: Path) -> None: + async def exercise() -> None: + ( + controller, + coordinator, + candidates, + commit, + ledger, + barrier, + groups, + ) = await self._drain_setup(tmp_path, rollouts=50, on_wire=1) + scans = 0 + full_scan = ledger.gym_checkpoint_inventory + + def counting_scan(instance_ids): + nonlocal scans + scans += 1 + return full_scan(instance_ids) + + ledger.gym_checkpoint_inventory = counting_scan + drain = asyncio.create_task( + controller._drain_non_exported_gym_candidates( + coordinator, candidates, commit + ) + ) + # Long enough for a 50 ms poll to have rescanned several times. + await asyncio.sleep(0.3) + await self._resolve(ledger, barrier, groups[0], "sealed") + await asyncio.wait_for(drain, timeout=5.0) + + assert scans == 0 + + asyncio.run(exercise()) + + @pytest.mark.parametrize("how", ["sealed", "abandoned"]) + def test_it_returns_as_soon_as_the_last_reply_resolves( + self, tmp_path: Path, how: str + ) -> None: + async def exercise() -> None: + ( + controller, + coordinator, + candidates, + commit, + ledger, + barrier, + groups, + ) = await self._drain_setup(tmp_path, rollouts=3, on_wire=2) + drain = asyncio.create_task( + controller._drain_non_exported_gym_candidates( + coordinator, candidates, commit + ) + ) + await asyncio.sleep(0.1) + await self._resolve(ledger, barrier, groups[0], how) + for _ in range(5): + await asyncio.sleep(0) + assert not drain.done(), "one on-wire reply is still outstanding" + + await self._resolve(ledger, barrier, groups[1], how) + # Woken by the change itself, not by the next timed poll. + for _ in range(5): + await asyncio.sleep(0) + assert drain.done() + await drain + + asyncio.run(exercise()) + + def test_a_timeout_reports_a_bounded_summary(self, tmp_path: Path) -> None: + async def exercise() -> None: + controller, coordinator, candidates, commit, *_ = await self._drain_setup( + tmp_path, rollouts=500, on_wire=500, control_timeout_s=0.05 + ) + with pytest.raises(TimeoutError) as raised: + await controller._drain_non_exported_gym_candidates( + coordinator, candidates, commit + ) + message = str(raised.value) + assert "500" in message + # A count and a short sample, not every unresolved identity. + assert len(message) < 2000, len(message) + + asyncio.run(exercise()) + + +class _DiskLiveSessionGymCoordinator(_LiveSessionGymCoordinator): + """Live-session Gym that also writes its manifest, and can act mid-checkpoint.""" + + def __init__( + self, + live: set[GymCheckpointEpisode], + *, + while_prepared: Callable[[], Any] | None = None, + ) -> None: + super().__init__(live) + self.while_prepared = while_prepared + + @asynccontextmanager + async def prepared(self, checkpoint_id: str): + async with super().prepared(checkpoint_id): + if self.while_prepared is not None: + await self.while_prepared() + yield + + async def commit( + self, + checkpoint_id: str, + checkpoint_root: Path, + inventory: dict[str, tuple[GymCheckpointEpisode, ...]], + ) -> GymCheckpointCommitResult: + result = await super().commit(checkpoint_id, checkpoint_root, inventory) + write_gym_checkpoint_manifest(checkpoint_root, result.manifest) + return result + + +class TestARowDroppedDuringACheckpointDoesNotFailIt: + """A broken stream can drop rows while Gym is prepared; the checkpoint must hold. + + The drop marks the rows abandoned at once, but Gym refuses a retire until it + resumes, so it still holds and exports them. Restore requires the saved Gym + manifest to equal RL's inventory, so the saved checkpoint must leave the + dropped episodes out, and they stay queued for the retire after resume. + """ + + @staticmethod + async def _drop( + ledger: RolloutRecoveryLedger, + barrier: DataPlaneCheckpointBarrier, + group_id: str, + generation_indices: list[int], + ) -> None: + # What _drop_unreturned_rows records before its retire waits for resume. + async with barrier.mutation() as cut: + ledger.abandon_dropped_dispatch( + cut, group_id, generation_indices=generation_indices, replace=True + ) + + @staticmethod + async def _save( + controller: Any, barrier: DataPlaneCheckpointBarrier, tmp_path: Path + ) -> None: + async with controller._prepared_gym_checkpoint( + tmp_path, checkpoint_id="save-1" + ) as gym_commit: + assert gym_commit is not None + async with barrier.checkpoint() as cut: + await controller._validate_gym_checkpoint_cut( + cut, gym_commit, checkpoint_path=tmp_path + ) + + @staticmethod + def _assert_restorable( + controller: Any, + coordinator: _SavingGymCoordinator, + tmp_path: Path, + *, + kept: set[GymCheckpointEpisode], + ) -> None: + saved = load_gym_checkpoint_manifest(tmp_path).instances + assert saved == controller._gym_checkpoint_inventory(coordinator) + assert set(saved["tools/replica-0"]) == kept + + def test_a_drop_after_the_pre_prepare_retire(self, tmp_path: Path) -> None: + """Gym exports the dropped row it still holds, outside RL's candidates.""" + + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + ledger = RolloutRecoveryLedger() + async with barrier.mutation() as cut: + group = _dispatched_two_sibling_group(ledger, cut) + row_0, row_1 = ( + GymCheckpointEpisode(*group.gym_episode(index)) for index in (0, 1) + ) + + coordinator = _DiskLiveSessionGymCoordinator( + live={row_0, row_1}, + while_prepared=lambda: self._drop(ledger, barrier, group.group_id, [1]), + ) + controller = _controller_for(coordinator, ledger) + controller._data_plane_checkpoint_barrier = barrier + + await asyncio.wait_for( + self._save(controller, barrier, tmp_path), _ASYNC_TEST_TIMEOUT_S + ) + + self._assert_restorable(controller, coordinator, tmp_path, kept={row_0}) + assert ledger.gym_checkpoint_retirements({"tools/replica-0"}) == { + "tools/replica-0": ((row_1.rollout_id, row_1.attempt),) + } + + asyncio.run(exercise()) + + def test_a_drop_while_the_commit_drains(self, tmp_path: Path) -> None: + """A failed on-wire reply ends the stream, dropping a row Gym exported.""" + + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + ledger = RolloutRecoveryLedger() + async with barrier.mutation() as cut: + group = _dispatched_two_sibling_group(ledger, cut) + row_0, row_1 = ( + GymCheckpointEpisode(*group.gym_episode(index)) for index in (0, 1) + ) + # Row 0 is parked in Gym and exported; row 1's reply is on the wire. + coordinator = _DiskLiveSessionGymCoordinator(live={row_0}) + controller = _controller_for(coordinator, ledger) + controller._data_plane_checkpoint_barrier = barrier + + async def fail_the_on_wire_reply() -> None: + while not any( + isinstance(event, tuple) and event[0] == "commit" + for event in coordinator.events + ): + await asyncio.sleep(0) + await self._drop(ledger, barrier, group.group_id, [0, 1]) + + failing = asyncio.create_task(fail_the_on_wire_reply()) + await asyncio.wait_for( + self._save(controller, barrier, tmp_path), _ASYNC_TEST_TIMEOUT_S + ) + await failing + + self._assert_restorable(controller, coordinator, tmp_path, kept=set()) + assert ledger.gym_checkpoint_retirements({"tools/replica-0"}) == { + "tools/replica-0": tuple( + sorted( + (episode.rollout_id, episode.attempt) + for episode in (row_0, row_1) + ) + ) + } + + asyncio.run(exercise()) + + def test_a_dropped_episode_takes_its_staging_keys_with_it( + self, tmp_path: Path + ) -> None: + """The cut needs, and the saved manifest keeps, only the kept episodes' rows.""" + + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + ledger = RolloutRecoveryLedger() + async with barrier.mutation() as cut: + group = _dispatched_two_sibling_group(ledger, cut) + row_0, row_1 = ( + GymCheckpointEpisode(*group.gym_episode(index)) for index in (0, 1) + ) + key_0 = f"{row_0.rollout_id}/call-0" + key_1 = f"{row_1.rollout_id}/call-0" + coordinator = _DiskLiveSessionGymCoordinator( + live={row_0, row_1}, + while_prepared=lambda: self._drop(ledger, barrier, group.group_id, [1]), + ) + coordinator.staging_keys_by_episode = {row_0: (key_0,), row_1: (key_1,)} + controller = _controller_for(coordinator, ledger) + controller._data_plane_checkpoint_barrier = barrier + controller._master_config = SimpleNamespace( + token_capture=SimpleNamespace(staging_partition="rollout_staging") + ) + + async def list_sample_ids(method: str, *, partition_id: str) -> list[str]: + assert (method, partition_id) == ("list_sample_ids", "rollout_staging") + # The dropped row's staging is gone by the time the cut looks. + return [key_0] + + controller._call_dp = list_sample_ids + + await asyncio.wait_for( + self._save(controller, barrier, tmp_path), _ASYNC_TEST_TIMEOUT_S + ) + + saved = load_gym_checkpoint_manifest(tmp_path) + assert set(saved.instances["tools/replica-0"]) == {row_0} + assert saved.staging_keys == {"tools/replica-0": (key_0,)} + + asyncio.run(exercise()) + + def test_the_tq_capture_expects_only_the_kept_episodes_rows( + self, tmp_path: Path + ) -> None: + """The TQ capture must not expect the rows of an episode RL dropped.""" + + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + ledger = RolloutRecoveryLedger() + async with barrier.mutation() as cut: + group = _dispatched_two_sibling_group(ledger, cut) + row_0, row_1 = ( + GymCheckpointEpisode(*group.gym_episode(index)) for index in (0, 1) + ) + key_0 = f"{row_0.rollout_id}/call-0" + key_1 = f"{row_1.rollout_id}/call-0" + coordinator = _DiskLiveSessionGymCoordinator( + live={row_0, row_1}, + while_prepared=lambda: self._drop(ledger, barrier, group.group_id, [1]), + ) + coordinator.staging_keys_by_episode = {row_0: (key_0,), row_1: (key_1,)} + controller = _controller_for(coordinator, ledger) + controller._data_plane_checkpoint_barrier = barrier + controller._master_config = SimpleNamespace( + token_capture=SimpleNamespace(staging_partition="rollout_staging") + ) + + async def list_sample_ids(method: str, *, partition_id: str) -> list[str]: + assert (method, partition_id) == ("list_sample_ids", "rollout_staging") + return [key_0] + + controller._call_dp = list_sample_ids + expected_by_capture: list[set[str]] = [] + + async def capture( + cut, checkpoint_path, *, additional_expected_staging_keys + ): + expected_by_capture.append(set(additional_expected_staging_keys)) + return "captured" + + controller._capture_rollout_checkpoint_cut = capture + + async def save() -> None: + async with controller._prepared_gym_checkpoint( + tmp_path, checkpoint_id="save-1" + ) as gym_commit: + assert gym_commit is not None + async with barrier.checkpoint() as cut: + assert ( + await controller._capture_coordinated_rollout_checkpoint_cut( + cut, tmp_path, gym_commit=gym_commit + ) + == "captured" + ) + + await asyncio.wait_for(save(), _ASYNC_TEST_TIMEOUT_S) + + # The dropped row's staging is gone; expecting it would fail the cut. + assert expected_by_capture == [{key_0}] + + asyncio.run(exercise()) + + def test_only_a_queued_drop_may_differ_from_the_saved_manifest( + self, tmp_path: Path + ) -> None: + """An exported episode RL dropped without queuing a retire still fails.""" + + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + ledger = RolloutRecoveryLedger() + async with barrier.mutation() as cut: + group = _dispatched_two_sibling_group(ledger, cut) + row_0, row_1 = ( + GymCheckpointEpisode(*group.gym_episode(index)) for index in (0, 1) + ) + coordinator = _DiskLiveSessionGymCoordinator(live={row_0, row_1}) + controller = _controller_for(coordinator, ledger) + controller._data_plane_checkpoint_barrier = barrier + + async with controller._prepared_gym_checkpoint( + tmp_path, checkpoint_id="save-1" + ) as gym_commit: + assert gym_commit is not None + await self._drop(ledger, barrier, group.group_id, [0]) + # As if the drop's retire had already been acknowledged. + ledger.mark_gym_retired( + {"tools/replica-0": [(row_0.rollout_id, row_0.attempt)]} + ) + async with barrier.checkpoint() as cut: + with pytest.raises(RuntimeError, match="ownership changed"): + await controller._validate_gym_checkpoint_cut( + cut, gym_commit, checkpoint_path=tmp_path + ) + + asyncio.run(exercise()) diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 865325586f1..cd1985d963d 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -99,9 +99,14 @@ 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.distributed.virtual_cluster import ClusterConfig +from nemo_rl.environments.gym_checkpoint_coordinator import ( + GymCheckpointNotReady, + GymCheckpointOperationError, +) from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, ROLLOUT_RECOVERY_STATE_FILENAME, + RecoveryTargetLevel, RolloutRecoveryLedger, ) from nemo_rl.experience.route_plan import ( @@ -517,11 +522,11 @@ def set_data_plane_checkpoint_barrier(self, barrier: Any) -> None: def set_weight_version(self, version: int) -> None: self.weight_versions.append(version) - def suspend_request_deadlines(self) -> None: + def suspend_request_deadlines(self, holder: str = "generation_stand_down") -> None: if self._events is not None: self._events.append("suspend_deadlines") - def resume_request_deadlines(self) -> None: + def resume_request_deadlines(self, holder: str = "generation_stand_down") -> None: if self._events is not None: self._events.append("resume_deadlines") @@ -1496,6 +1501,7 @@ def test_logs_checkpoint_outcome_reason_and_effective_cadence( outcome_keys = {"completed", "failed", "skipped"} reason_keys = { "reason_completed", + "reason_gym_unavailable", "reason_invariant_error", "reason_io_error", "reason_missing_trainer_anchor", @@ -1744,6 +1750,70 @@ def test_snapshot_reports_no_new_mutations(self, tmp_path: Path) -> None: assert result.reason == "no_data_plane_mutations" assert actor._dp_client.save_calls == [] + @pytest.mark.parametrize( + ("prefix_cuts_enabled", "restore_level", "expect_save"), + [ + (True, RecoveryTargetLevel.PREFIX, True), + (False, RecoveryTargetLevel.PREFIX, False), + (True, RecoveryTargetLevel.TURN, False), + ], + ) + def test_unchanged_data_plane_still_cuts_a_decoding_prefix( + self, + tmp_path: Path, + prefix_cuts_enabled: bool, + restore_level: RecoveryTargetLevel, + expect_save: bool, + ) -> None: + """Decoding grows a recoverable prefix without mutating the data plane.""" + + class _SaveStarted(Exception): + pass + + actor = self._actor(tmp_path) + actor._generation_prefix_cuts_enabled = prefix_cuts_enabled + ledger = actor._rollout_recovery_ledger + + async def dispatch() -> None: + async with actor._data_plane_checkpoint_barrier.mutation() as cut: + ledger.reserve_group( + cut, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload={"idx": 7, "message_log": []}, + expected_generations=2, + target_step=0, + start_weight_version=0, + restore_level=restore_level, + admitted=True, + ) + ledger.mark_group_dispatched( + cut, "g7", gym_instance_id="tools/replica-0" + ) + + asyncio.run(dispatch()) + actor._last_rollout_snapshot_mutation_version = ( + actor._data_plane_checkpoint_barrier.mutation_version + ) + result = None + try: + with patch.object( + actor._checkpointer, "finalize_pending", side_effect=_SaveStarted + ): + try: + result = asyncio.run(actor._save_rollout_checkpoint()) + except _SaveStarted: + pass + finally: + actor._checkpointer.shutdown() + + if expect_save: + assert result is None + else: + assert result is not None + assert result.reason == "no_data_plane_mutations" + def test_periodic_pump_reports_each_consecutive_failure( self, tmp_path: Path, @@ -1814,6 +1884,89 @@ async def _failing_save(*, force: bool = False) -> bool: output = capsys.readouterr().out assert output.count("Periodic rollout checkpoint failed") == 2 + @pytest.mark.parametrize( + "error", + [ + GymCheckpointNotReady({"tools/replica-0": {"agent": ("rollout-1",)}}), + GymCheckpointOperationError( + "prepare", {"tools/replica-0": TimeoutError("no reply")} + ), + ], + ids=["not-ready", "unanswered"], + ) + def test_periodic_pump_retries_a_gym_checkpoint_that_could_not_run( + self, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + error: Exception, + ): + actor = self._actor(tmp_path) + actor._logger = MagicMock() + actor._master_config.rollout_checkpointing.snapshot_attempt_interval_s = 0.001 + actor._train_steps = 1 + + async def _main() -> None: + two_failures = asyncio.Event() + calls = 0 + + async def _failing_save(*, force: bool = False) -> bool: + nonlocal calls + del force + calls += 1 + if calls == 2: + two_failures.set() + raise error + + actor._save_rollout_checkpoint = _failing_save + pump = asyncio.create_task(actor._rollout_checkpoint_pump()) + await asyncio.wait_for(two_failures.wait(), timeout=1.0) + pump.cancel() + await asyncio.gather(pump, return_exceptions=True) + + try: + asyncio.run(_main()) + finally: + actor._checkpointer.shutdown() + + output = capsys.readouterr().out + assert "retaining the previous committed snapshot" in output + assert "consecutive_failures=2" in output + logged = actor._logger.log_metrics.call_args.args[0] + assert logged["reason_gym_unavailable"] == 1.0 + + def test_periodic_pump_does_not_retry_a_broken_gym_reply(self, tmp_path: Path): + actor = self._actor(tmp_path) + actor._logger = MagicMock() + actor._master_config.rollout_checkpointing.snapshot_attempt_interval_s = 0.001 + actor._train_steps = 1 + calls = 0 + + async def _main() -> None: + async def _failing_save(*, force: bool = False) -> bool: + nonlocal calls + del force + calls += 1 + raise GymCheckpointOperationError( + "commit", + {"tools/replica-0": RuntimeError("unexpected participant set")}, + ) + + actor._save_rollout_checkpoint = _failing_save + with pytest.raises(GymCheckpointOperationError, match="participant set"): + await asyncio.wait_for( + actor._rollout_checkpoint_pump(), + timeout=1.0, + ) + + try: + asyncio.run(_main()) + finally: + actor._checkpointer.shutdown() + + assert calls == 1 + logged = actor._logger.log_metrics.call_args.args[0] + assert logged["reason_invariant_error"] == 1.0 + def test_periodic_pump_does_not_retry_invariant_failure(self, tmp_path: Path): actor = self._actor(tmp_path) actor._logger = MagicMock() @@ -2091,6 +2244,60 @@ async def validate_inventory() -> int: assert dp_client.clear_calls == [(["orphan-key"], staging_partition)] assert sorted(dp_client.sample_ids) == [route_key, "sealed-key"] + def test_rollout_recovery_inventory_preserves_gym_manifest_staging_keys(self): + staging_partition = "rollout_staging" + dp_client = _StagingInventoryDPClient( + ["sealed-key", "gym-owned-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 + + async def validate_inventory() -> int: + async with DataPlaneCheckpointBarrier().mutation() as cut: + return await actor._validate_rollout_recovery_inventory( + cut, + replay_metadata=None, + clear_unreferenced=True, + additional_expected_staging_keys={"gym-owned-key"}, + ) + + assert asyncio.run(validate_inventory()) == 2 + assert dp_client.clear_calls == [(["orphan-key"], staging_partition)] + assert sorted(dp_client.sample_ids) == ["gym-owned-key", "sealed-key"] + + def test_rollout_recovery_inventory_rejects_missing_gym_staging_key(self): + staging_partition = "rollout_staging" + dp_client = _StagingInventoryDPClient( + ["sealed-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 + + async def validate_inventory() -> None: + async with DataPlaneCheckpointBarrier().mutation() as cut: + await actor._validate_rollout_recovery_inventory( + cut, + replay_metadata=None, + clear_unreferenced=True, + additional_expected_staging_keys={"missing-gym-key"}, + ) + + with pytest.raises(RuntimeError, match="staging rows missing"): + asyncio.run(validate_inventory()) + + assert dp_client.clear_calls == [] + assert dp_client.sample_ids == ["sealed-key"] + def test_gated_sampler_writes_authoritative_tq_checkpoint(self, tmp_path): mc = _actor_master_config( tmp_path, @@ -2331,6 +2538,8 @@ def _ppo_save_actor(tmp_path: Path, calls: list[str]): actor._checkpointer = MagicMock() actor._checkpointer.save_optimizer = True actor._checkpointer.init_tmp_checkpoint.return_value = str(checkpoint_path) + actor._gym_checkpoint_coordinator = None + actor._rollout_dispatch_admission_gate = None actor._is_ppo = True actor._trainer = _OrderRecordingPolicy(calls) actor._value = _OrderRecordingCritic(calls) diff --git a/tests/unit/single_controller/test_rollout_checkpoint.py b/tests/unit/single_controller/test_rollout_checkpoint.py index e1e97bd2c13..7201b2475e7 100644 --- a/tests/unit/single_controller/test_rollout_checkpoint.py +++ b/tests/unit/single_controller/test_rollout_checkpoint.py @@ -378,7 +378,7 @@ def _fields(schema: Any) -> set[str]: "async_rl", {"sampler": {"name": "windowed", "max_staleness_versions": 2}}, ), - ("rollout_recovery", {"default_granularity": "prompt_group"}), + ("rollout_recovery", {"target_level": "prompt_group"}), ], ) def test_bootstrap_fingerprint_rejects_rollout_semantic_changes( @@ -398,7 +398,7 @@ def test_bootstrap_fingerprint_rejects_rollout_semantic_changes( "mixed_weight_version_policy": "allow", }, "async_rl": {"sampler": {"name": "windowed", "max_staleness_versions": 1}}, - "rollout_recovery": {"default_granularity": "sibling"}, + "rollout_recovery": {"target_level": "sibling"}, } modified = {**base, section: {**base[section], **changed}} diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index b4ceb836625..648d07d6738 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -67,6 +67,7 @@ ) from nemo_rl.algorithms.single_controller_utils.config import ( RolloutCheckpointConfig, + RolloutRecoveryConfig, TokenCaptureConfig, _validate_opd_full_config, validate_single_controller_config, @@ -83,7 +84,7 @@ ) from nemo_rl.data_plane.tq_token_sink import MEDIA_STAGING_FIELDS from nemo_rl.distributed.virtual_cluster import ClusterConfig -from nemo_rl.experience.rollout_recovery import RecoveryGranularity +from nemo_rl.experience.rollout_recovery import RecoveryTargetLevel from nemo_rl.experience.rollouts import EffortLevelsConfig from nemo_rl.models.generation.megatron.megatron_generation import MegatronGeneration from nemo_rl.utils.config import ( @@ -654,7 +655,7 @@ def test_rollout_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=prompt_group", + "++rollout_recovery.target_level=prompt_group", "++async_rl.sampler.name=in_order", "++async_rl.sampler.max_lookahead_versions=1", "++async_rl.min_groups_for_streaming_train=4", @@ -678,13 +679,41 @@ def test_rollout_recovery_functional_config_resolves_to_runtime_contract( 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.PROMPT_GROUP + master_config.rollout_recovery.target_level is RecoveryTargetLevel.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 +@pytest.mark.parametrize( + ("target_level", "turn_checkpointing", "generation_prefix_cuts"), + [ + (RecoveryTargetLevel.PREFIX, True, True), + (RecoveryTargetLevel.TURN, True, False), + (RecoveryTargetLevel.SIBLING, False, False), + (RecoveryTargetLevel.PROMPT_GROUP, False, False), + ], +) +def test_recovery_target_enables_required_checkpoint_capabilities( + target_level: RecoveryTargetLevel, + turn_checkpointing: bool, + generation_prefix_cuts: bool, +): + config = RolloutRecoveryConfig(target_level=target_level) + + assert config.turn_checkpointing_enabled is turn_checkpointing + assert config.generation_prefix_cuts_enabled is generation_prefix_cuts + + +@pytest.mark.parametrize( + "removed_field", + ["preserve_generation_prefixes", "generation_chunk_flush_tokens"], +) +def test_recovery_config_rejects_removed_prefix_fields(removed_field: str): + with pytest.raises(ValueError, match=removed_field): + RolloutRecoveryConfig.model_validate({removed_field: True}) + + class TestSetup: """setup arg validation + actor_args assembly.""" @@ -1023,6 +1052,49 @@ def test_periodic_checkpointing_requires_claim_aware_custom_sampler(self): with pytest.raises(ValueError, match="supports training-claim ownership"): setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_periodic_checkpointing_does_not_enable_gym_participants(self): + mc = _make_master_config(env={"should_use_nemo_gym": True}) + mc.policy["generation"]["model_name"] = "test-model" + mc.rollout_checkpointing = RolloutCheckpointConfig( + snapshot_attempt_interval_s=1.0 + ) + + with patch.object( + sc_setup_mod, + "build_nemo_gym_actors", + return_value=MagicMock(), + ) as mock_spinup: + sc_setup_mod._spinup_gym( + mc, + ["http://generation"], + MagicMock(name="tokenizer"), + ) + + assert mock_spinup.call_args.kwargs["turn_checkpointing_enabled"] is False + + @pytest.mark.parametrize( + "target_level", [RecoveryTargetLevel.TURN, RecoveryTargetLevel.PREFIX] + ) + def test_gym_state_recovery_requires_periodic_snapshots( + self, target_level: RecoveryTargetLevel + ): + mc = _make_master_config(env={"should_use_nemo_gym": True}) + mc.policy["generation"]["vllm_cfg"] = { + "async_engine": True, + "expose_http_server": True, + } + mc.token_capture.enabled = True + mc.rollout_recovery.target_level = target_level + + with pytest.raises( + ValueError, + match=( + f"target_level='{target_level.value}' requires.*" + "snapshot_attempt_interval_s" + ), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_periodic_checkpointing_warns_without_per_step_trainer_anchors( self, tmp_path: Path, @@ -1047,6 +1119,7 @@ def test_periodic_checkpointing_warns_without_per_step_trainer_anchors( } ) mc.token_capture.enabled = True + mc.rollout_recovery.target_level = RecoveryTargetLevel.TURN mc.rollout_checkpointing = RolloutCheckpointConfig( snapshot_attempt_interval_s=1.0 ) @@ -1061,7 +1134,7 @@ def test_periodic_checkpointing_warns_without_per_step_trainer_anchors( patch.object(sc_setup_mod, "should_use_nemo_gym", return_value=True), patch.object( sc_setup_mod, "build_nemo_gym_actors", return_value=MagicMock() - ), + ) as mock_spinup, patch.object(sc_setup_mod, "router_replay_enabled", return_value=False), patch( "nemo_rl.experience.rollout_reassembler_actor." @@ -1075,6 +1148,7 @@ def test_periodic_checkpointing_warns_without_per_step_trainer_anchors( ) assert actor_args.finalizer_actors == fake_finalizers + assert mock_spinup.call_args.kwargs["turn_checkpointing_enabled"] is True def test_disabled_periodic_checkpointing_ignores_existing_snapshots( self, @@ -1488,16 +1562,16 @@ def test_invalid_config_fails_before_setup_factories( 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 + mc.rollout_recovery.target_level = RecoveryTargetLevel.PROMPT_GROUP elif invalid_case == "recovery_override_without_capture": mc = _make_master_config() - mc.rollout_recovery.task_source_granularity_overrides = { - "genrm": RecoveryGranularity.PROMPT_GROUP + mc.rollout_recovery.task_source_target_level_overrides = { + "genrm": RecoveryTargetLevel.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 + mc.rollout_recovery.agent_target_level_overrides = { + "genrm_agent": RecoveryTargetLevel.PROMPT_GROUP } else: # pragma: no cover raise AssertionError(f"unknown test case {invalid_case}") @@ -1873,6 +1947,7 @@ def test_nemo_gym_wires_env_handle(self, patched_factories): tokenizer=tokenizer, enable_router_replay=False, use_fastokens=False, + turn_checkpointing_enabled=False, token_capture=None, ) mock_validate.assert_called_once_with( @@ -1954,6 +2029,11 @@ def test_token_capture_always_creates_finalizer_actor_pool( generation, _ = patched_factories["_build_generation"].return_value _, setup_kwargs = generation.setup_token_capture.call_args assert setup_kwargs["capture_media"] is with_processor + # Workers bound a generation cut by Gym's control timeout, counted + # from when the cut request reaches them. + assert setup_kwargs["generation_cut_control_timeout_s"] == ( + mc.token_capture.control_timeout_s + ) def test_nemo_gym_coverage_failure_shuts_down_shards(self, patched_factories): mc = _make_master_config(colocated=False, backend="vllm") diff --git a/uv.lock b/uv.lock index 8141a20efc8..06f782ec6ae 100644 --- a/uv.lock +++ b/uv.lock @@ -4152,37 +4152,35 @@ source = { editable = "3rdparty/Gym-workspace/Gym" } dependencies = [ { name = "aiohttp" }, { name = "anthropic", version = "0.109.2", source = { registry = "https://pypi.org/simple" } }, + { name = "anyio" }, { name = "datasets" }, { name = "devtools" }, { name = "fastapi", extra = ["standard"], marker = "extra == 'extra-7-nemo-rl-automodel' or extra == 'extra-7-nemo-rl-mcore' or extra == 'extra-7-nemo-rl-nemo-gym' or (extra == 'extra-7-nemo-rl-modelopt' and extra != 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nvrx' and extra != 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (extra != 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-trtllm')" }, { name = "fonttools" }, - { name = "gitpython" }, { name = "httptools" }, { name = "hydra-core" }, { name = "itsdangerous" }, + { name = "jinja2" }, + { name = "jsonschema" }, { name = "mcp" }, - { name = "mlflow" }, - { name = "mlflow-skinny" }, { name = "omegaconf" }, { name = "openai", version = "2.44.0", source = { registry = "https://pypi.org/simple" } }, { name = "orjson" }, { name = "packaging" }, { name = "psutil" }, - { name = "pyarrow" }, { name = "pyasn1" }, { name = "pydantic" }, { name = "pydantic-core" }, { name = "pydot" }, + { name = "pyjwt" }, { name = "python-multipart" }, { name = "ray", extra = ["default"], marker = "extra == 'extra-7-nemo-rl-automodel' or extra == 'extra-7-nemo-rl-mcore' or extra == 'extra-7-nemo-rl-nemo-gym' or (extra == 'extra-7-nemo-rl-modelopt' and extra != 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nvrx' and extra != 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (extra != 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-trtllm')" }, { name = "rich" }, { name = "scipy" }, - { name = "sqlparse" }, { name = "tqdm" }, { name = "urllib3" }, { name = "uvicorn" }, { name = "uvloop" }, - { name = "wandb" }, { name = "yappi" }, ] @@ -4190,6 +4188,7 @@ dependencies = [ requires-dist = [ { name = "aiohttp", specifier = ">=3.14.1" }, { name = "anthropic", specifier = "<=0.109.2" }, + { name = "anyio", specifier = ">=4.15.1" }, { name = "boto3", marker = "extra == 'sandbox'", specifier = ">=1.34" }, { name = "coverage", extras = ["toml"], marker = "extra == 'dev'" }, { name = "datasets" }, @@ -4198,55 +4197,58 @@ requires-dist = [ { name = "devtools" }, { name = "e2b", marker = "extra == 'sandbox'", specifier = ">=2.36.0,<3.0.0" }, { name = "fastapi" }, - { name = "flashinfer-python", marker = "extra == 'vllm'", specifier = "==0.6.13" }, + { name = "flashinfer-python", marker = "extra == 'vllm'", specifier = "==0.6.18" }, { name = "fonttools", specifier = ">=4.60.2" }, - { name = "gitpython", specifier = ">=3.1.61" }, + { name = "gitpython", marker = "extra == 'wandb'", specifier = ">=3.1.61" }, { name = "gprof2dot", marker = "extra == 'dev'" }, { name = "httptools" }, { name = "hydra-core" }, { name = "itsdangerous" }, + { name = "jinja2", specifier = ">=3.1" }, + { name = "jsonschema", specifier = ">=4.18,<5" }, { name = "mcp", specifier = ">=1.28.1,<2" }, - { name = "mlflow", specifier = ">=3.16.0" }, - { name = "mlflow-skinny", specifier = ">=3.16.0" }, + { name = "mlflow", marker = "extra == 'mlflow'", specifier = ">=3.16.1" }, + { name = "mlflow-skinny", marker = "extra == 'mlflow'", specifier = ">=3.16.1" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.8.0" }, - { name = "nemo-gym", extras = ["dev", "sandbox"], marker = "extra == 'all'" }, - { name = "nemo-lens", extras = ["sdk"], marker = "extra == 'telemetry'", git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d" }, + { name = "nemo-gym", extras = ["dev", "mlflow", "sandbox", "wandb"], marker = "extra == 'all'" }, + { name = "nemo-lens", extras = ["sdk"], marker = "extra == 'telemetry'", git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b0f977d414b2f89938604a0b7eaa78ee08bc8700" }, { name = "omegaconf" }, { name = "openai", specifier = "==2.44.0" }, { name = "opensandbox", marker = "extra == 'sandbox'", specifier = ">=0.1.15" }, { name = "openshell", marker = "extra == 'openshell'", specifier = ">=0.0.92,<0.1" }, { name = "opentelemetry-instrumentation-aiohttp-client", marker = "extra == 'telemetry'", specifier = ">=0.60b0" }, { name = "opentelemetry-instrumentation-fastapi", marker = "extra == 'telemetry'", specifier = ">=0.60b0" }, - { name = "orjson" }, + { name = "orjson", specifier = ">=3.11.6" }, { name = "packaging", specifier = ">=24.2" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.6.0" }, { name = "psutil" }, - { name = "pyarrow", specifier = ">=23.0.1" }, + { name = "pyarrow", marker = "extra == 'mlflow'", specifier = ">=23.0.1" }, { name = "pyasn1", specifier = ">=0.6.4" }, { name = "pydantic" }, { name = "pydantic-core" }, { name = "pydot" }, + { name = "pyjwt", specifier = ">=2.15.1" }, { name = "pytest", marker = "extra == 'dev'" }, { name = "pytest-asyncio", marker = "extra == 'dev'" }, { name = "pytest-cov", marker = "extra == 'dev'" }, { name = "pytest-xdist", marker = "extra == 'dev'" }, - { name = "python-multipart", specifier = ">=0.0.22" }, - { name = "ray", extras = ["default"], specifier = ">=2.56.1" }, + { name = "python-multipart", specifier = ">=0.0.27" }, + { name = "ray", extras = ["default"], specifier = ">=2.58.0" }, { name = "requests-mock", marker = "extra == 'dev'" }, { name = "rich" }, { name = "ruff", marker = "extra == 'dev'" }, { name = "scipy", specifier = ">=1.14.1" }, - { name = "sqlparse", specifier = ">=0.6.0" }, + { name = "sqlparse", marker = "extra == 'mlflow'", specifier = ">=0.6.0" }, { name = "tenacity", marker = "extra == 'sandbox'", specifier = ">=9.1.4" }, { name = "tqdm" }, - { name = "urllib3", specifier = ">=2.7.0" }, + { name = "urllib3", specifier = ">=2.8.0" }, { name = "uvicorn" }, { name = "uvloop" }, - { name = "vllm", marker = "extra == 'vllm'", specifier = "==0.25.1" }, - { name = "wandb" }, + { name = "vllm", marker = "extra == 'vllm'", specifier = "==0.29.0" }, + { name = "wandb", marker = "extra == 'wandb'", specifier = ">=0.30.0" }, { name = "yappi", specifier = ">=1.7.6" }, ] -provides-extras = ["all", "vllm", "sandbox", "openshell", "telemetry", "dev"] +provides-extras = ["all", "mlflow", "wandb", "vllm", "sandbox", "openshell", "telemetry", "dev"] [package.metadata.requires-dev] codec = [ @@ -6523,11 +6525,11 @@ wheels = [ [[package]] name = "pyjwt" -version = "2.14.0" +version = "2.15.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/af/c3/8a3b59c25070cc61dc517fbdfa5dc0904670c96f605cc69759dc09166b99/pyjwt-2.14.0.tar.gz", hash = "sha256:77283c83fb56ecf566a886c757a714bc83668e38156de2cce8263302f42e0b86", size = 113177, upload-time = "2026-09-11T13:11:54.638Z" } +sdist = { url = "https://files.pythonhosted.org/packages/43/ea/5194e52748b0da83d71e082d75496eaec6e58f419f5e184786ded517e6a9/pyjwt-2.15.1.tar.gz", hash = "sha256:4f259e80cdfb6b3fc18a7de51fd1ef9ec79652f25019bae68975ca2468a34df8", size = 121252, upload-time = "2026-09-28T18:40:42.598Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9c/97/672cb32ce0dfea44b740cb7b4f97038463b9cf7c0ead1aacf595572851d6/pyjwt-2.14.0-py3-none-any.whl", hash = "sha256:ad0cef71c756a56e74863c2919cf0985f72decbcfcb550ee2f422e7c62b5eedc", size = 32896, upload-time = "2026-09-11T13:11:53.409Z" }, + { url = "https://files.pythonhosted.org/packages/50/ca/44de4e75f8aadc457f0634be3b542815078ded46dca30efb960edeecad6e/pyjwt-2.15.1-py3-none-any.whl", hash = "sha256:42d59d631f7768a1028a64c7ff581a9bf7519804daf91fc5b6c56e30eec5e193", size = 33860, upload-time = "2026-09-28T18:40:41.429Z" }, ] [package.optional-dependencies]