diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 10f3a8a8fad..9cc1d7b98e4 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -360,6 +360,7 @@ The SC path is still under active development. Feature gaps are tracked in [issu - Train backend: only Megatron is supported and validated; the AutoModel training path has not been tested on SC. - Generation backend: vLLM and Megatron generation are supported (Megatron in both non-colocated and colocated modes); SGLang and TRT-LLM have not been tested on SC. - Validation is not yet supported (setup raises on `val_period > 0`, `val_at_start`, or `val_at_end`); checkpointing is. +- Rollout retry budgets apply to both ordinary generation and token capture. After `max_data_attempts`, a deterministic failure skips the prompt only within the run-wide `max_skipped_prompts` budget. Its default of `0` propagates the original error. Skipped capture prompts retain controller-owned lineage until replacement or step-shortfall accounting completes; standalone capture callers clean their own recovery group. - (PPO) Rollout drop budgets — `async_rl.rollout_failure.max_skipped_prompts` and `max_consecutive_dropped_prompts` must both be `0`. A drop shortens the step, and the critic shards it against the configured `value.train_global_batch_size` rather than its actual size, so setup rejects a non-zero budget. The resiliency layer stays available on GRPO. - Reward shaping and sample filtering — `reward_shaping`, `reward_scaling`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. Environment-flagged sample masking and `overlong_filtering` are supported; truncated completions are excluded from the loss through `sample_mask`, and a step in which every completion is filtered is rejected rather than skipped. - The `windowed` sampler has no `over_sampling_ratio` cap — over-produced groups aged past the window are evicted, wasting rollout compute. diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 36dd5ffd103..2bbed4c9bad 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -131,6 +131,7 @@ 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_reassembler_pool import RolloutReassemblerPool from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, ROLLOUT_RECOVERY_STATE_FILENAME, @@ -308,13 +309,7 @@ def __init__( # therefore degrades to the documented off state rather than to a broken one. self._gen_fleet = getattr(actor_args, "fleet_monitor", None) self._generation_router = getattr(actor_args, "generation_router", None) - self._finalizer_actors = list(actor_args.finalizer_actors) - self._available_finalizers: asyncio.Queue[Any] = asyncio.Queue() - for actor in self._finalizer_actors: - self._available_finalizers.put_nowait(actor) - self._active_finalizers = 0 - self._finalizer_waiters = 0 - self._finalizer_unknown_outcomes = 0 + self._reassembler_pool = RolloutReassemblerPool(actor_args.reassembler_actors) self._finalizer_metrics_by_group: dict[str, dict[str, float]] = {} teacher_worker_groups = getattr(actor_args, "teacher_worker_groups", None) or {} if teacher_worker_groups: @@ -603,11 +598,7 @@ async def run(self) -> dict[str, Any]: for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) - for actor in self._finalizer_actors: - try: - ray.kill(actor, no_restart=True) - except Exception as error: - print(f"finalizer actor termination failed: {error}", flush=True) + self._reassembler_pool.shutdown() try: self._weight_synchronizer.shutdown() except Exception as e: # teardown must not mask the original failure @@ -630,10 +621,10 @@ async def ping(self) -> dict[str, Any]: "inflight_rollouts": self._inflight_rollouts, "rollout_permitted": self._rollout_permitted.is_set(), "epoch": self._current_epoch, - "active_finalizers": self._active_finalizers, - "finalizer_waiters": self._finalizer_waiters, - "finalizer_queue_depth": self._available_finalizers.qsize(), - "finalizer_unknown_outcomes": self._finalizer_unknown_outcomes, + "active_finalizers": self._reassembler_pool.active, + "finalizer_waiters": self._reassembler_pool.waiters, + "finalizer_queue_depth": self._reassembler_pool.available_count, + "finalizer_unknown_outcomes": self._reassembler_pool.unknown_outcomes, } # ── internal helpers ─────────────────────────────────────────────────── @@ -1356,25 +1347,12 @@ async def _finalize_with_actor( valid-row fraction is no longer a finalizer-side drop -- the caller decides that, since only the caller can source a replacement. """ - self._finalizer_waiters += 1 - queue_depth = max( - 0, - self._finalizer_waiters - self._available_finalizers.qsize(), - ) - queue_start = time.perf_counter() try: - actor = await self._available_finalizers.get() + lease = await self._reassembler_pool.acquire() except asyncio.CancelledError: await self._cleanup_known_finalization_request(request) raise - finally: - self._finalizer_waiters -= 1 - queue_wait_ms = (time.perf_counter() - queue_start) * 1000.0 - self._active_finalizers += 1 - active_actor_count = self._active_finalizers finalize_start = time.perf_counter() - rpc_submitted = False - actor_reusable = False try: # The actor publishes canonical rows before returning metadata. Keep # the remote write, local replay-index update, and lineage hand-off in @@ -1387,10 +1365,9 @@ async def _finalize_with_actor( ledger = self._rollout_recovery_ledger ledger.mark_finalization_started(cut, request.group_id) try: - rpc_submitted = True - finalized = await actor.finalize.remote(request) + lease.rpc_submitted = True + finalized = await lease.actor.finalize.remote(request) except BaseException: - self._finalizer_unknown_outcomes += 1 ledger.mark_finalization_unknown(cut, request.group_id) print( "FATAL: finalizer actor RPC failed after submission; canonical " @@ -1400,7 +1377,7 @@ async def _finalize_with_actor( ) raise else: - actor_reusable = True + lease.outcome_known = True if finalized.dropped: try: @@ -1459,18 +1436,16 @@ async def _finalize_with_actor( ledger.discard_group(cut, request.group_id) committed = True finally: - self._active_finalizers -= 1 - if actor_reusable or not rpc_submitted: - self._available_finalizers.put_nowait(actor) + self._reassembler_pool.release(lease) finalize_total_ms = (time.perf_counter() - finalize_start) * 1000.0 if not committed: return None finalized.metrics.update( { - "finalize/queue_wait_ms": queue_wait_ms, + "finalize/queue_wait_ms": lease.queue_wait_ms, "finalize/total_ms": finalize_total_ms, - "finalize/queue_depth": float(queue_depth), - "finalize/active_actor_count": float(active_actor_count), + "finalize/queue_depth": float(lease.queue_depth), + "finalize/active_actor_count": float(lease.active_actor_count), } ) self._finalizer_metrics_by_group[request.group_id] = dict(finalized.metrics) @@ -1587,7 +1562,7 @@ async def _dispatch_one_prompt( generation_permit_released = False inflight_count_released = False try: - if self._finalizer_actors: + if self._reassembler_pool: # Token-capture path: run capture generation, release the # generation permits as soon as the tokens are staged, then # hand the metadata-only request to the finalizer actor pool. @@ -1614,23 +1589,54 @@ async def _dispatch_one_prompt( sem.release() generation_permit_released = True if request is None: - # Dropped within the infra budget: nothing was - # committed, so the train pump will never release - # this permit, and the step it was stamped for - # must be allowed to close short. - if ( - self._rollout_recovery_enabled - and lineage_group_id is not None - ): + # Capture exhausted a tolerated retry budget. Transfer + # the retained lineage to a spare or credit shortfall + # in one checkpoint-atomic controller decision. + if self._rollout_recovery_enabled: + assert lineage_group_id is not None async with ( - self._data_plane_checkpoint_barrier.mutation() - ) as cut: + self._data_plane_checkpoint_barrier.mutation() as cut + ): + replacement = self._take_replacement( + target_step, replacements + ) await self._rollout_manager.discard_recovery_group( cut, lineage_group_id ) - self._buffer_capacity.release() - self._credit_shortfall(target_step) - return + if replacement is not None: + lender_step = self._promote_into_step( + target_step + ) + if lender_step is not None: + target_step = lender_step + lineage_group_id = self._rollout_manager.reserve_prompt_group( + cut, + replacement, + target_step=target_step, + ) + else: + self._credit_shortfall(target_step) + else: + replacement = self._take_replacement( + target_step, replacements + ) + if replacement is None: + self._buffer_capacity.release() + if not self._rollout_recovery_enabled: + self._credit_shortfall(target_step) + return + replacements += 1 + prompt = replacement + if not self._rollout_recovery_enabled: + lender_step = self._promote_into_step(target_step) + if lender_step is not None: + target_step = lender_step + await self._rollout_permitted.wait() + await sem.acquire() + self._inflight_rollouts += 1 + inflight_count_released = False + generation_permit_released = False + continue finalized = await self._finalize_with_actor(request) if finalized is None: # Finalizer dropped the group as a structural diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 39084e004fb..b3cacfc1456 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -1249,17 +1249,6 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: "train with penalized rewards.", stacklevel=2, ) - if ( - token_capture_config.enabled - and async_config.rollout_failure.max_skipped_prompts - ): - warnings.warn( - "async_rl.rollout_failure.max_skipped_prompts does nothing with " - "token_capture.enabled=true: the capture dispatch path re-raises a " - "deterministic failure instead of skipping the prompt, so the run " - "ends on the first prompt that exhausts max_data_attempts.", - stacklevel=2, - ) # A non-zero reference-policy KL penalty makes the loss read # ``reference_policy_logprobs``, but the SC train pump only computes them diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 05872ba63bb..50b5b94a783 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -160,7 +160,7 @@ class SingleControllerActorArgs: partition_id: str save_state: GRPOSaveState last_checkpoint_path: Optional[str] - finalizer_actors: list[Any] + reassembler_actors: list[Any] # Defaulted fields must follow the required ones above, so these stay last. data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = None bootstrap_identity: Optional[BootstrapCompatibilityIdentity] = None @@ -1715,14 +1715,14 @@ def _build_trainer_then_megatron_generation() -> tuple[ token_capture_cfg.staging_partition if token_capture_cfg.enabled else None ), ) - finalizer_actors: list[Any] = [] + reassembler_actors: list[Any] = [] if token_capture_cfg.enabled: from nemo_rl.experience.rollout_reassembler_actor import ( RolloutReassemblerActorConfig, create_rollout_reassembler_actors, ) - finalizer_actors = create_rollout_reassembler_actors( + reassembler_actors = create_rollout_reassembler_actors( dp_config, RolloutReassemblerActorConfig( partition_id=partition_id, @@ -1785,7 +1785,7 @@ def _build_trainer_then_megatron_generation() -> tuple[ last_checkpoint_path=recovery_checkpoint_path, data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, bootstrap_identity=bootstrap_identity, - finalizer_actors=finalizer_actors, + reassembler_actors=reassembler_actors, fleet_monitor=fleet_monitor, generation_router=generation_router, teacher_worker_groups=teacher_worker_groups, diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 21f8cc68fc0..9356d551989 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -198,6 +198,19 @@ def backoff_for(self, attempt: int) -> float: return min(self.backoff_base_s * 2 ** (attempt - 1), self.max_backoff_s) +class _CaptureCleanupError(RuntimeError): + """A capture attempt could not release its reservation safely.""" + + +@dataclass +class _RetryState: + """Per-prompt attempt accounting shared by both dispatch paths.""" + + infra_attempts: int = 0 + data_attempts: int = 0 + group_attempt: int = 0 + + @dataclass class RolloutStats: """Counters describing what the retry policy has been doing. @@ -1816,12 +1829,8 @@ async def generate_and_push( "the resumed configuration requests " f"{self._num_generations_per_prompt}" ) - policy = self._retry_policy - infra_attempts = 0 - data_attempts = 0 - last_infra_error: Optional[Exception] = None + retry = _RetryState() logical_group_id: Optional[str] = None - group_attempt = 0 extra_env_info = input_sample.get("extra_env_info") if isinstance(extra_env_info, dict): configured_group_id = extra_env_info.get(NEMO_GYM_GROUP_ID_KEY) @@ -1839,13 +1848,10 @@ async def generate_and_push( raise ValueError( f"{NEMO_GYM_GROUP_ATTEMPT_KEY} must be a non-negative integer" ) - group_attempt = configured_group_attempt + retry.group_attempt = configured_group_attempt - # The loop condition is the infrastructure budget, so running out of it exits - # here rather than raising from inside the handler. The data budget is tracked - # separately and terminates from within, since exhausting it is a statement - # about the prompt rather than about the fleet. - while infra_attempts < policy.max_infra_attempts: + # Each failure class has an independent budget in the shared retry state. + while True: start_version = self._weight_version # A lineage-tracked prompt reuses its durable logical ID only after the # prior attempt's buffer slot was removed successfully. Ordinary callers @@ -1871,7 +1877,7 @@ async def generate_and_push( attempt_extra_env_info = attempt_input_sample["extra_env_info"] attempt_extra_env_info[NEMO_GYM_GROUP_ID_KEY] = logical_group_id attempt_extra_env_info[NEMO_GYM_GROUP_ATTEMPT_KEY] = ( - group_attempt + retry.group_attempt ) record = await self.run_rollout(attempt_input_sample) finally: @@ -1912,51 +1918,9 @@ async def generate_and_push( # and would spend the rollout retry budget on the wrong subsystem. if _contains_post_write_enrichment_error(error): raise - reason = type(error).__name__ - - if classify_rollout_failure(error) is FailureClass.INFRA: - infra_attempts += 1 - last_infra_error = error - if infra_attempts >= policy.max_infra_attempts: - break - self._stats.record_redispatch(reason) - # The backpressure permit is held across this sleep, so the wait is - # capped by max_backoff_s rather than growing without bound. - await asyncio.sleep(policy.backoff_for(infra_attempts)) - group_attempt += 1 + if await self._retry_after_failure(retry, error, input_sample): continue - - data_attempts += 1 - if data_attempts >= policy.max_data_attempts: - self._stats.record_data_failure(reason) - if self._skipped_prompts >= policy.max_skipped_prompts: - # At the default of 0 this fires on the first exhaustion and the - # original failure propagates unchanged -- one knob, and its - # zero value is the old "fail_fast" without a second key to - # contradict it. - if policy.max_skipped_prompts == 0: - raise - raise RolloutDataFailure( - f"skipped {self._skipped_prompts} prompts and this one also " - f"exhausted its data budget, exceeding max_skipped_prompts=" - f"{policy.max_skipped_prompts}; the dataset or " - "sequence-length configuration is likely wrong" - ) from error - self._skipped_prompts += 1 - print( - f"skipping prompt idx={input_sample['idx']} after " - f"{data_attempts} deterministic failure(s) ({reason}: {error})", - flush=True, - ) - self._stats.skipped += 1 - return RolloutOutcome.SKIPPED - # A data retry, NOT a re-dispatch: the fleet is fine, this prompt is - # suspect. Recording it as a re-dispatch made rollout/redispatch_total -- - # documented above as the sign the fleet is degrading -- climb for bad - # data, which is the one distinction the two budgets exist to draw. - self._stats.record_data_retry(reason) - group_attempt += 1 - continue + return RolloutOutcome.SKIPPED except BaseException: # Cancellation and other non-Exception exits: clean up, never retry. try: @@ -1968,55 +1932,82 @@ async def generate_and_push( ) raise - self._stats.committed += 1 - # A commit proves the fleet is answering, which is exactly the claim the - # consecutive budget is testing, so it clears the run of drops. Placed on - # the success path rather than in the infra handler so that a prompt which - # succeeded on a retry also counts -- the fleet recovered either way. - self._consecutive_infra_drops = 0 - if lineage_group_id is not None: - async with ( - self._tq_buffer.data_plane_checkpoint_barrier.mutation() - ) as cut: - self._recovery_ledger.discard_group(cut, lineage_group_id) - return RolloutOutcome.COMMITTED - - # The infrastructure budget ran out. The same failure followed the prompt across - # repeated shard selections, which says the fleet is broken rather than the - # prompt. - # - # The budget is >= 1 (enforced in RolloutRetryPolicy), so the loop ran at least - # once and can only have exited through the infra branch's break. - assert last_infra_error is not None - reason = type(last_infra_error).__name__ - self._consecutive_infra_drops += 1 - if self._consecutive_infra_drops > policy.max_consecutive_dropped_prompts: - raise RolloutRedispatchExhausted( - f"prompt idx={input_sample['idx']} exhausted its infrastructure retry " - f"budget after {infra_attempts} attempt(s) " - f"(max_infra_attempts_per_prompt=" - f"{policy.max_infra_attempts}), and this was drop " - f"{self._consecutive_infra_drops} with no rollout committed in between, " - f"exceeding max_consecutive_dropped_prompts=" - f"{policy.max_consecutive_dropped_prompts}; the generation fleet is not " - f"recovering. Last failure was {reason}: {last_infra_error}" - ) from last_infra_error - - # Under the budget: give up on this prompt and let the run continue. The caller - # owns the backpressure permit for a SKIPPED outcome, and -- because the prompt - # may have been stamped for a specific training step that will now never fill -- - # owns atomically replacing its retained ledger entry or crediting the shortfall - # so the train pump can close that step short. - self._stats.record_infra_drop(reason, self._consecutive_infra_drops) - print( - f"dropping prompt idx={input_sample['idx']} after {infra_attempts} " - f"infrastructure failure(s) ({reason}: {last_infra_error}) " - f"[consecutive drop {self._consecutive_infra_drops}/" - f"{policy.max_consecutive_dropped_prompts}]", - flush=True, - ) + break + + self._stats.committed += 1 + # A commit proves the fleet is answering, which is exactly the claim the + # consecutive budget is testing, so it clears the run of drops. Placed on + # the success path rather than in the infra handler so that a prompt which + # succeeded on a retry also counts -- the fleet recovered either way. + self._consecutive_infra_drops = 0 + if lineage_group_id is not None: + async with self._tq_buffer.data_plane_checkpoint_barrier.mutation() as cut: + self._recovery_ledger.discard_group(cut, lineage_group_id) + return RolloutOutcome.COMMITTED + + async def _retry_after_failure( + self, + state: _RetryState, + error: Exception, + input_sample: DatumSpec, + ) -> bool: + """Account for a cleaned-up attempt; return whether to try again. + + Terminal budget failures propagate; False means the caller must transfer + its retained ownership to the replacement/shortfall path. + """ + policy = self._retry_policy + reason = type(error).__name__ + if classify_rollout_failure(error) is FailureClass.INFRA: + state.infra_attempts += 1 + if state.infra_attempts < policy.max_infra_attempts: + self._stats.record_redispatch(reason) + await asyncio.sleep(policy.backoff_for(state.infra_attempts)) + state.group_attempt += 1 + return True + self._consecutive_infra_drops += 1 + if self._consecutive_infra_drops > policy.max_consecutive_dropped_prompts: + raise RolloutRedispatchExhausted( + f"prompt idx={input_sample['idx']} exhausted its infrastructure retry " + f"budget after {state.infra_attempts} attempt(s) " + f"(max_infra_attempts_per_prompt={policy.max_infra_attempts}), " + f"and this was drop {self._consecutive_infra_drops} with no rollout " + "committed in between, exceeding max_consecutive_dropped_prompts=" + f"{policy.max_consecutive_dropped_prompts}; the generation fleet " + f"is not recovering. Last failure was {reason}: {error}" + ) from error + self._stats.record_infra_drop(reason, self._consecutive_infra_drops) + print( + f"dropping prompt idx={input_sample['idx']} after " + f"{state.infra_attempts} infrastructure failure(s) ({reason}: {error}) " + f"[consecutive drop {self._consecutive_infra_drops}/" + f"{policy.max_consecutive_dropped_prompts}]", + flush=True, + ) + else: + state.data_attempts += 1 + if state.data_attempts < policy.max_data_attempts: + self._stats.record_data_retry(reason) + state.group_attempt += 1 + return True + self._stats.record_data_failure(reason) + if policy.max_skipped_prompts == 0: + raise error + if self._skipped_prompts >= policy.max_skipped_prompts: + raise RolloutDataFailure( + f"skipped {self._skipped_prompts} prompts and this one also " + "exhausted its data budget, exceeding max_skipped_prompts=" + f"{policy.max_skipped_prompts}; the dataset or " + "sequence-length configuration is likely wrong" + ) from error + self._skipped_prompts += 1 + print( + f"skipping prompt idx={input_sample['idx']} after " + f"{state.data_attempts} deterministic failure(s) ({reason}: {error})", + flush=True, + ) self._stats.skipped += 1 - return RolloutOutcome.SKIPPED + return False async def generate_for_finalization( self, @@ -2028,9 +2019,9 @@ async def generate_for_finalization( ) -> Optional["ReassemblyRequest"]: """Capture siblings with stable lineage and configured retry granularity. - Returns ``None`` when infrastructure retries are exhausted within the - configured drop budget. The caller then owns the backpressure permit and - target-step shortfall. + Returns ``None`` when infrastructure or data retries are exhausted within + their configured drop/skip budget. The controller owns replacement and + the backpressure permit; standalone calls clean their own recovery group. """ assert self._tq_buffer is not None, ( "generate_for_finalization requires tq_buffer to be set at __init__" @@ -2056,11 +2047,8 @@ async def generate_for_finalization( f"{recovery_group.expected_generations} generation(s), but the " f"resumed configuration requests {self._num_generations_per_prompt}" ) - policy = self._retry_policy - infra_attempts = 0 - data_attempts = 0 - last_infra_error: Optional[Exception] = None - while infra_attempts < policy.max_infra_attempts: + retry = _RetryState() + while True: try: request = await self._generate_for_finalization_attempt( input_sample, @@ -2068,56 +2056,33 @@ async def generate_for_finalization( inflight_registry=inflight_registry, ) except Exception as error: - reason = type(error).__name__ - if classify_rollout_failure(error) is FailureClass.INFRA: - infra_attempts += 1 - last_infra_error = error - if infra_attempts >= policy.max_infra_attempts: - break - self._stats.record_redispatch(reason) - await asyncio.sleep(policy.backoff_for(infra_attempts)) - continue - - data_attempts += 1 - if data_attempts >= policy.max_data_attempts: - self._stats.record_data_failure(reason) + if isinstance( + error, _CaptureCleanupError + ) or _contains_post_write_enrichment_error(error): raise - self._stats.record_data_retry(reason) - continue + + if owns_recovery_group and ( + ( + classify_rollout_failure(error) is FailureClass.INFRA + and retry.infra_attempts + 1 + >= self._retry_policy.max_infra_attempts + ) + or ( + classify_rollout_failure(error) is FailureClass.DATA + and retry.data_attempts + 1 + >= self._retry_policy.max_data_attempts + ) + ): + # An independent caller has no controller to release its lineage. + async with self._recovery_mutation() as cut: + await self.discard_recovery_group(cut, recovery_group_id) + if await self._retry_after_failure(retry, error, input_sample): + continue + return None self._consecutive_infra_drops = 0 return request - assert last_infra_error is not None - reason = type(last_infra_error).__name__ - self._consecutive_infra_drops += 1 - if owns_recovery_group: - # Without controller-owned recovery lineage, nobody above this method - # knows the temporary group ID. Clean its known staging ownership before - # dropping the only record that names those rows. - async with self._recovery_mutation() as cut: - await self.discard_recovery_group(cut, recovery_group_id) - if self._consecutive_infra_drops > policy.max_consecutive_dropped_prompts: - raise RolloutRedispatchExhausted( - f"prompt idx={input_sample['idx']} exhausted its infrastructure " - f"retry budget after {infra_attempts} capture attempt(s) and this " - f"was drop {self._consecutive_infra_drops}, exceeding " - f"max_consecutive_dropped_prompts=" - f"{policy.max_consecutive_dropped_prompts}; last failure was " - f"{reason}: {last_infra_error}" - ) from last_infra_error - self._stats.record_infra_drop(reason, self._consecutive_infra_drops) - print( - f"dropping capture prompt idx={input_sample['idx']} after " - f"{infra_attempts} infrastructure failure(s) ({reason}: " - f"{last_infra_error}) [consecutive drop " - f"{self._consecutive_infra_drops}/" - f"{policy.max_consecutive_dropped_prompts}]", - flush=True, - ) - self._stats.skipped += 1 - return None - async def _generate_for_finalization_attempt( self, input_sample: DatumSpec, @@ -2294,13 +2259,18 @@ async def _record_streamed_completion( # run end (there is no prefix-clear primitive in the data plane # yet). Their ledger files are inert — failure rows or missing # terminal rows keep any later read fail-closed. - self._tq_buffer.abort(group_id) - async with self._recovery_mutation() as cut: - # Intentional staleness aborts discard the ledger owner before - # cancelling this task. Preserve the original cancellation rather - # than replacing it with "unknown group" during cleanup. - if group_id in self._recovery_ledger: - self._recovery_ledger.abandon_unsealed(cut, group_id) + try: + self._tq_buffer.abort(group_id) + async with self._recovery_mutation() as cut: + # Intentional staleness aborts discard the ledger owner before + # cancelling this task. Preserve the original cancellation rather + # than replacing it with "unknown group" during cleanup. + if group_id in self._recovery_ledger: + self._recovery_ledger.abandon_unsealed(cut, group_id) + except Exception as cleanup_error: + raise _CaptureCleanupError( + f"capture cleanup failed for group {group_id}; refusing to retry" + ) from cleanup_error # The capture ledger has no per-rollout fail endpoint. Rows from # abandoned attempts are unreferenced and are swept with the # staging partition at run teardown. diff --git a/nemo_rl/experience/rollout_reassembler_pool.py b/nemo_rl/experience/rollout_reassembler_pool.py new file mode 100644 index 00000000000..23bff2cce16 --- /dev/null +++ b/nemo_rl/experience/rollout_reassembler_pool.py @@ -0,0 +1,104 @@ +# 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 time +from dataclasses import dataclass +from typing import Any + +import ray + + +@dataclass +class ReassemblerLease: + """One exclusive actor checkout, including its publication outcome.""" + + actor: Any + queue_wait_ms: float + queue_depth: int + active_actor_count: int + rpc_submitted: bool = False + outcome_known: bool = False + + +class RolloutReassemblerPool: + """Own actor availability while the controller owns publication and commits. + + An actor with an unknown RPC outcome is quarantined until shutdown. It must + never be replaced or retried: its canonical rows may already be published. + """ + + def __init__(self, actors: list[Any]) -> None: + self._actors = list(actors) + self._available: asyncio.Queue[Any] = asyncio.Queue() + for actor in actors: + self._available.put_nowait(actor) + self.active = 0 + self.waiters = 0 + self.unknown_outcomes = 0 + self._closed = False + + def __bool__(self) -> bool: + return bool(self._actors) + + @property + def available_count(self) -> int: + return 0 if self._closed else self._available.qsize() + + async def acquire(self) -> ReassemblerLease: + """Wait for an actor; cancellation never consumes a checkout.""" + if self._closed: + raise RuntimeError("reassembler pool is closed") + self.waiters += 1 + queue_depth = max(0, self.waiters - self.available_count) + start = time.perf_counter() + try: + actor = await self._available.get() + finally: + self.waiters -= 1 + if actor is None: + raise RuntimeError("reassembler pool is closed") + self.active += 1 + return ReassemblerLease( + actor=actor, + queue_wait_ms=(time.perf_counter() - start) * 1000.0, + queue_depth=queue_depth, + active_actor_count=self.active, + ) + + def release(self, lease: ReassemblerLease) -> None: + """Return a safe actor or quarantine an unresolved publication.""" + self.active -= 1 + if lease.rpc_submitted and not lease.outcome_known: + self.unknown_outcomes += 1 + elif not self._closed: + self._available.put_nowait(lease.actor) + + def shutdown(self) -> None: + """Terminate all actors after controller tasks have drained.""" + if self._closed: + return + self._closed = True + while not self._available.empty(): + self._available.get_nowait() + for _ in range(self.waiters): + self._available.put_nowait(None) + for actor in self._actors: + try: + ray.kill(actor, no_restart=True) + except Exception as error: + # Teardown must not mask the controller's original failure. + print(f"reassembler actor termination failed: {error}", flush=True) diff --git a/nemo_rl/models/generation/vllm/token_capture_host.py b/nemo_rl/models/generation/vllm/token_capture_host.py new file mode 100644 index 00000000000..619f9b8987e --- /dev/null +++ b/nemo_rl/models/generation/vllm/token_capture_host.py @@ -0,0 +1,286 @@ +# 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. + +"""Worker-local token capture lifecycle, independent of the vLLM engine.""" + +from __future__ import annotations + +import logging +import threading +from typing import TYPE_CHECKING, Any, Protocol + +import torch + +from nemo_rl.data_plane.interfaces import DataPlaneRuntimeConfig + +if TYPE_CHECKING: + from nemo_gym.token_id_capture.staging.capture import ( + ActiveCall, + RolloutTokenCapture, + ) + from nemo_gym.token_id_capture.staging.records import CaptureAdmission + +LOGGER = logging.getLogger(__name__) + + +class _PrefixSource(Protocol): + def fetch_prefix_token_ids(self, staging_keys: list[str]) -> list[int]: ... + + +class TokenCaptureHost: + """Own capture state and staging prefixes for one serving worker. + + Gym imports remain deferred so constructing a disabled host does not require + the optional Gym dependency. Engine interception belongs to the worker. + """ + + def __init__(self) -> None: + self.token_capture: RolloutTokenCapture | None = None + self._rollout_weight_version = 0 + self._capture_calls: dict[int, tuple[ActiveCall, list[int]]] = {} + self._staging_source: _PrefixSource | None = None + self._prefix_cache: dict[str, list[int]] = {} + self._prefix_cache_lock = threading.Lock() + + def install_token_capture(self, capture: RolloutTokenCapture) -> None: + """Gym's ``install_capture`` seam (the ``CaptureHost`` contract).""" + self.token_capture = capture + + def setup(self, dp_cfg: DataPlaneRuntimeConfig, staging_partition: str) -> bool: + """Install Gym capture with a worker-local staging client and adapter.""" + # Deferred: nemo_gym is an optional extra absent in non-gym runs. + from nemo_gym.token_id_capture.adapters.vllm import VLLMCaptureAdapter + from nemo_gym.token_id_capture.staging import install_capture + + from nemo_rl.data_plane import build_data_plane_client + from nemo_rl.data_plane.tq_token_sink import TQTokenSink, TQTokenSource + + dp_client = build_data_plane_client(dp_cfg, bootstrap=False) + sink = TQTokenSink(dp_client, staging_partition=staging_partition) + self._staging_source = TQTokenSource( + dp_client, staging_partition=staging_partition + ) + self._prefix_cache.clear() + install_capture( + self, + sink=sink, + weight_version_fn=lambda: self._rollout_weight_version, + adapter=VLLMCaptureAdapter(), + ) + return True + + def set_weight_version(self, version: int) -> None: + """Rotate the weight version stamped on subsequent captured calls.""" + self._rollout_weight_version = int(version) + + def admission(self, request: Any) -> CaptureAdmission | None: + """Parse the ledger's ``ng_capture`` context into a ``CaptureAdmission``. + + Returns None unless capture is installed and the request carries the + context. The dict itself is never mutated: the admission is the typed, + read-only contract that the prefix resolution and ``begin_call`` share. + """ + context = getattr(request, "ng_capture", None) + if self.token_capture is None or not context: + return None + # Deferred: nemo_gym is an optional extra absent in non-gym runs. + from nemo_gym.token_id_capture.staging.records import CaptureAdmission + + return CaptureAdmission.model_validate(context) + + def begin_request( + self, + request: Any, + prompt_token_ids: list[int], + *, + admission: CaptureAdmission | None = None, + prefix_token_ids: list[int] | None = None, + ) -> None: + """Admit one ledger-forwarded call into the capture layer. + + Called from preprocess_chat once the exact engine prompt is known + (post-splice in token-in mode, full render in text mode). No-op + unless capture is installed and the request carries the ledger's + ``ng_capture`` context. + + ``prefix_token_ids`` is the prefix resolved by + :meth:`resolve_prefix`; Gym's ``begin_call`` checks it + against the admission (length == ``prev_len``, equal to an inline + prefix) and requires it for a ``staging_chain`` admission. + """ + capture = self.token_capture + if capture is None: + return + if admission is None: + admission = self.admission(request) + if admission is None: + return + call = capture.begin_call( + admission, + prefix_token_ids=prefix_token_ids, + stream=bool(getattr(request, "stream", False)), + ) + self._capture_calls[id(request)] = (call, list(prompt_token_ids)) + + def fetch_chain_prefix(self, staging_chain: list[str]) -> list[int]: + """Assemble prefix token ids from staging_chain, with a worker-local LRU cache.""" + cache = self._prefix_cache + with self._prefix_cache_lock: + cached_ids: list[int] = [] + miss_start = 0 + for i, key in enumerate(staging_chain): + if key in cache: + cached_ids = cache[key] + miss_start = i + 1 + miss_keys = staging_chain[miss_start:] + if not miss_keys: + return list(cached_ids) + if self._staging_source is None: + raise RuntimeError( + "_staging_source not initialized; call setup_token_capture() first" + ) + # TQ read stays outside the lock so concurrent fetches overlap. + fetched = self._staging_source.fetch_prefix_token_ids(miss_keys) + result = cached_ids + fetched + last_key = staging_chain[-1] + with self._prefix_cache_lock: + cache[last_key] = result + if len(cache) > 256: + del cache[next(iter(cache))] + return result + + def resolve_prefix(self, admission: CaptureAdmission) -> list[int]: + """Resolve a ``CaptureAdmission`` to the flat prefix the engine prompt starts with. + + A ``staging_chain`` is fetched through the cached TransferQueue read; + an inline ``required_prefix_token_ids`` is used as is; a text root has + no prefix. Length checks are Gym's: ``begin_call`` rejects a prefix + that does not match ``prev_len``. + """ + if admission.mode == "text": + return [] + if admission.staging_chain: + return self.fetch_chain_prefix(list(admission.staging_chain)) + return list(admission.required_prefix_token_ids) + + def enter_prefix(self, request: Any, prefix_token_ids: list[int]) -> None: + """Attach the resolved prefix to the request through the capture adapter. + + ``VLLMCaptureAdapter.enter_prefix`` writes the engine-native field + (``required_prefix_token_ids``) into a payload; the same fields are + applied to the pydantic request so the existing prefix-splice branch + of preprocess_chat handles staged and inline prefixes alike. + """ + capture = self.token_capture + assert capture is not None + adapter = capture.adapter + assert adapter is not None + for field_name, value in adapter.enter_prefix({}, prefix_token_ids).items(): + setattr(request, field_name, value) + + @staticmethod + def _delta_align_routed_experts( + payload: dict[str, Any], *, prev_len: int, prompt_len: int, generated_len: int + ) -> None: + """Normalize optional vLLM routes to the exact staged token delta.""" + choices = payload.get("choices") or [] + if len(choices) != 1 or not isinstance(choices[0], dict): + return + choice = dict(choices[0]) + message = dict(choice.get("message") or {}) + routed = message.get("routed_experts") + if routed is None: + return + try: + from nemo_rl.utils.routed_experts_codec import ( + decode_routed_experts, + encode_routed_experts, + ) + + if isinstance(routed, str): + dtype_name = routed.split(":", 3)[1] + dtype = { + "int8": torch.int8, + "int16": torch.int16, + "int32": torch.int32, + }.get(dtype_name) + if dtype is None: + raise ValueError(f"unsupported routed_experts dtype {dtype_name!r}") + else: + dtype = torch.int16 + experts = decode_routed_experts(routed, dtype) + expected_full_len = prompt_len + generated_len + if experts.dim() != 3 or experts.shape[0] != expected_full_len: + raise ValueError( + f"route length {experts.shape[0]} does not match engine sequence " + f"length {expected_full_len}" + ) + message["routed_experts"] = encode_routed_experts(experts[prev_len:]) + except (IndexError, TypeError, ValueError) as error: + LOGGER.warning( + "dropping invalid routed_experts from staged capture: %s", error + ) + message.pop("routed_experts", None) + choice["message"] = message + payload["choices"] = [choice] + + def finish_request(self, request: Any, content: dict) -> dict: + """Stage the finished call and ride its coords on the response. + + Fail-closed: the sink write happens inside complete_call — + the coords exist only after the bytes are durable, and any capture + failure degrades to capture_failed coords without breaking the + completion. Token ids and logprobs are stripped: the staged delta is + the only token store on this path, so the worker->gate hop carries + text + delta ids + coords only. + """ + state = self._capture_calls.pop(id(request), None) + if state is None: + return content + call, prompt_token_ids = state + payload = dict(content) + # vLLM's OpenAI response carries no prompt ids; the adapter reads the + # preprocess-time engine prompt off the payload (see + # nemo_gym.token_id_capture.adapters.vllm.extract_prompt_ids). + payload["prompt_token_ids"] = prompt_token_ids + capture = self.token_capture + assert capture is not None + adapter = capture.adapter + if adapter is not None: + try: + generated_token_ids, _ = adapter.extract_generation(payload) + except Exception: # capture core will report the authoritative failure + generated_token_ids = [] + self._delta_align_routed_experts( + payload, + prev_len=call.admission.prev_len, + prompt_len=len(prompt_token_ids), + generated_len=len(generated_token_ids), + ) + coords = capture.complete_call_from_response(call, payload) + for choice in content.get("choices") or []: + choice.pop("logprobs", None) + # The delta-aligned routes were staged to TQ above; the served + # full-length copy is dead weight the gate strips on arrival. + message = choice.get("message") + if isinstance(message, dict): + message.pop("routed_experts", None) + content["ng_commit_coords"] = coords.model_dump() + return content + + def abort_request(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[0], reason=reason) diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index bb77973fedb..c5289f7f2bd 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -25,6 +25,8 @@ import ray import torch + +from nemo_rl.data_plane.interfaces import DataPlaneRuntimeConfig import uvicorn from fastapi import FastAPI @@ -44,6 +46,7 @@ from nemo_rl.models.generation.vllm.checkpoint_engine import ( VllmAsyncCheckpointEngineRpcMixin, ) +from nemo_rl.models.generation.vllm.token_capture_host import TokenCaptureHost from nemo_rl.models.generation.vllm.utils import ( attach_routed_experts_to_chat_response_choices, attach_token_information_to_chat_response_choices, @@ -189,20 +192,7 @@ def __init__( self._engine_loop = None self._http_engine_client = None - # Ledger-authoritative token capture (dormant until the - # setup_token_capture fan-out runs). The weight - # version is stamped per model call at begin_call time and rotated by - # 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): (ActiveCall, the - # exact engine prompt ids recorded at preprocess time). - self._capture_calls: dict[int, tuple[Any, list[int]]] = {} - self._staging_source: Any | None = None - # Guarded by _prefix_cache_lock: _fetch_chain_prefix runs on executor - # threads (asyncio.to_thread), so lookups/evictions can be concurrent. - self._prefix_cache: dict[str, list[int]] = {} - self._prefix_cache_lock = threading.Lock() + self.capture_host = TokenCaptureHost() super().__init__( config, @@ -436,243 +426,17 @@ async def get_reserved_url(self) -> Optional[str]: async def report_dp_openai_server_base_url(self) -> Optional[str]: return self.base_url - def install_token_capture(self, capture: Any) -> None: - """Gym's ``install_capture`` seam (the ``CaptureHost`` contract).""" - self.token_capture = capture - async def setup_token_capture( - self, dp_cfg: dict[str, Any], staging_partition: str + self, dp_cfg: DataPlaneRuntimeConfig, staging_partition: str ) -> bool: - """Host ledger-authoritative token capture in this worker. - - Fan-out target (token_capture.enabled only): builds the in-worker - data-plane client and TQTokenSink, then makes the single - ``install_capture`` call wiring Gym's engine-blind capture core + - vLLM adapter into this worker. Returns whether capture was installed - (False on non-model-owner ranks, which serve no HTTP). - """ + """Install capture on model-owner ranks that serve HTTP requests.""" if not self.is_model_owner: return False - # Deferred: nemo_gym is an optional extra absent in non-gym runs. - from nemo_gym.token_id_capture.adapters.vllm import VLLMCaptureAdapter - from nemo_gym.token_id_capture.staging import install_capture - - from nemo_rl.data_plane import build_data_plane_client - from nemo_rl.data_plane.tq_token_sink import TQTokenSink, TQTokenSource - - dp_client = build_data_plane_client(dp_cfg, bootstrap=False) - sink = TQTokenSink(dp_client, staging_partition=staging_partition) - self._staging_source = TQTokenSource( - dp_client, staging_partition=staging_partition - ) - self._prefix_cache.clear() - install_capture( - self, - sink=sink, - weight_version_fn=lambda: self._rollout_weight_version, - adapter=VLLMCaptureAdapter(), - ) - return True + return self.capture_host.setup(dp_cfg, staging_partition) async def set_rollout_weight_version(self, version: int) -> None: - """Rotate the weight version stamped on subsequent captured calls.""" - self._rollout_weight_version = int(version) - - def _capture_admission(self, request: Any) -> Any | None: - """Parse the ledger's ``ng_capture`` context into a ``CaptureAdmission``. - - Returns None unless capture is installed and the request carries the - context. The dict itself is never mutated: the admission is the typed, - read-only contract that the prefix resolution and ``begin_call`` share. - """ - context = getattr(request, "ng_capture", None) - if self.token_capture is None or not context: - return None - # Deferred: nemo_gym is an optional extra absent in non-gym runs. - from nemo_gym.token_id_capture.staging.records import CaptureAdmission - - return CaptureAdmission.model_validate(context) - - def _begin_request_capture( - self, - request: Any, - prompt_token_ids: list[int], - *, - admission: Any | None = None, - prefix_token_ids: list[int] | None = None, - ) -> None: - """Admit one ledger-forwarded call into the capture layer. - - Called from preprocess_chat once the exact engine prompt is known - (post-splice in token-in mode, full render in text mode). No-op - unless capture is installed and the request carries the ledger's - ``ng_capture`` context. - - ``prefix_token_ids`` is the prefix resolved by - :meth:`_resolve_admission_prefix`; Gym's ``begin_call`` checks it - against the admission (length == ``prev_len``, equal to an inline - prefix) and requires it for a ``staging_chain`` admission. - """ - capture = self.token_capture - if capture is None: - return - if admission is None: - admission = self._capture_admission(request) - if admission is None: - return - call = capture.begin_call( - admission, - prefix_token_ids=prefix_token_ids, - stream=bool(getattr(request, "stream", False)), - ) - self._capture_calls[id(request)] = (call, list(prompt_token_ids)) - - def _fetch_chain_prefix(self, staging_chain: list[str]) -> list[int]: - """Assemble prefix token ids from staging_chain, with a worker-local LRU cache.""" - cache = self._prefix_cache - with self._prefix_cache_lock: - cached_ids: list[int] = [] - miss_start = 0 - for i, key in enumerate(staging_chain): - if key in cache: - cached_ids = cache[key] - miss_start = i + 1 - miss_keys = staging_chain[miss_start:] - if not miss_keys: - return list(cached_ids) - if self._staging_source is None: - raise RuntimeError( - "_staging_source not initialized; call setup_token_capture() first" - ) - # TQ read stays outside the lock so concurrent fetches overlap. - fetched = self._staging_source.fetch_prefix_token_ids(miss_keys) - result = cached_ids + fetched - last_key = staging_chain[-1] - with self._prefix_cache_lock: - cache[last_key] = result - if len(cache) > 256: - del cache[next(iter(cache))] - return result - - def _resolve_admission_prefix(self, admission: Any) -> list[int]: - """Resolve a ``CaptureAdmission`` to the flat prefix the engine prompt starts with. - - A ``staging_chain`` is fetched through the cached TransferQueue read; - an inline ``required_prefix_token_ids`` is used as is; a text root has - no prefix. Length checks are Gym's: ``begin_call`` rejects a prefix - that does not match ``prev_len``. - """ - if admission.mode == "text": - return [] - if admission.staging_chain: - return self._fetch_chain_prefix(list(admission.staging_chain)) - return list(admission.required_prefix_token_ids) - - def _enter_request_prefix(self, request: Any, prefix_token_ids: list[int]) -> None: - """Attach the resolved prefix to the request through the capture adapter. - - ``VLLMCaptureAdapter.enter_prefix`` writes the engine-native field - (``required_prefix_token_ids``) into a payload; the same fields are - applied to the pydantic request so the existing prefix-splice branch - of preprocess_chat handles staged and inline prefixes alike. - """ - adapter = self.token_capture.adapter - for field_name, value in adapter.enter_prefix({}, prefix_token_ids).items(): - setattr(request, field_name, value) - - @staticmethod - def _delta_align_routed_experts( - payload: dict[str, Any], *, prev_len: int, prompt_len: int, generated_len: int - ) -> None: - """Normalize optional vLLM routes to the exact staged token delta.""" - choices = payload.get("choices") or [] - if len(choices) != 1 or not isinstance(choices[0], dict): - return - choice = dict(choices[0]) - message = dict(choice.get("message") or {}) - routed = message.get("routed_experts") - if routed is None: - return - try: - from nemo_rl.utils.routed_experts_codec import ( - decode_routed_experts, - encode_routed_experts, - ) - - if isinstance(routed, str): - dtype_name = routed.split(":", 3)[1] - dtype = { - "int8": torch.int8, - "int16": torch.int16, - "int32": torch.int32, - }.get(dtype_name) - if dtype is None: - raise ValueError(f"unsupported routed_experts dtype {dtype_name!r}") - else: - dtype = torch.int16 - experts = decode_routed_experts(routed, dtype) - expected_full_len = prompt_len + generated_len - if experts.dim() != 3 or experts.shape[0] != expected_full_len: - raise ValueError( - f"route length {experts.shape[0]} does not match engine sequence " - f"length {expected_full_len}" - ) - message["routed_experts"] = encode_routed_experts(experts[prev_len:]) - except (IndexError, TypeError, ValueError) as error: - LOGGER.warning( - "dropping invalid routed_experts from staged capture: %s", error - ) - message.pop("routed_experts", None) - choice["message"] = message - payload["choices"] = [choice] - - def _finish_request_capture(self, request: Any, content: dict) -> dict: - """Stage the finished call and ride its coords on the response. - - Fail-closed: the sink write happens inside complete_call — - the coords exist only after the bytes are durable, and any capture - failure degrades to capture_failed coords without breaking the - completion. Token ids and logprobs are stripped: the staged delta is - the only token store on this path, so the worker->gate hop carries - text + delta ids + coords only. - """ - state = self._capture_calls.pop(id(request), None) - if state is None: - return content - call, prompt_token_ids = state - payload = dict(content) - # vLLM's OpenAI response carries no prompt ids; the adapter reads the - # preprocess-time engine prompt off the payload (see - # nemo_gym.token_id_capture.adapters.vllm.extract_prompt_ids). - payload["prompt_token_ids"] = prompt_token_ids - adapter = self.token_capture.adapter - if adapter is not None: - try: - generated_token_ids, _ = adapter.extract_generation(payload) - except Exception: # capture core will report the authoritative failure - generated_token_ids = [] - self._delta_align_routed_experts( - payload, - prev_len=call.admission.prev_len, - prompt_len=len(prompt_token_ids), - generated_len=len(generated_token_ids), - ) - coords = self.token_capture.complete_call_from_response(call, payload) - for choice in content.get("choices") or []: - choice.pop("logprobs", None) - # The delta-aligned routes were staged to TQ above; the served - # full-length copy is dead weight the gate strips on arrival. - message = choice.get("message") - if isinstance(message, dict): - message.pop("routed_experts", 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[0], reason=reason) + """Rotate the version stamped on subsequent captured calls.""" + self.capture_host.set_weight_version(version) # ruff: noqa def _setup_vllm_openai_api_server(self, app: FastAPI) -> FastAPI: @@ -852,13 +616,15 @@ async def preprocess_chat( # its event loop explicitly. The adapter then attaches the # prefix to the request, so the inline-prefix branch below is # the single splice path for staged and inline prefixes. - admission = worker_self._capture_admission(request) + admission = worker_self.capture_host.admission(request) capture_prefix_token_ids: list[int] | None = None if admission is not None and admission.mode == "token_in": capture_prefix_token_ids = await asyncio.to_thread( - worker_self._resolve_admission_prefix, admission + worker_self.capture_host.resolve_prefix, admission + ) + worker_self.capture_host.enter_prefix( + request, capture_prefix_token_ids ) - worker_self._enter_request_prefix(request, capture_prefix_token_ids) if ( not hasattr(request, "required_prefix_token_ids") @@ -873,7 +639,7 @@ async def preprocess_chat( ) # Token capture, text mode: the full render is the exact # engine prompt. - worker_self._begin_request_capture( + worker_self.capture_host.begin_request( request, res[1][0]["prompt_token_ids"], admission=admission ) return res @@ -938,7 +704,7 @@ async def preprocess_chat( # Token capture, token-in mode: the spliced prompt is the # exact engine prompt; begin_call re-checks the prefix it # was spliced from against the admission. - worker_self._begin_request_capture( + worker_self.capture_host.begin_request( request, final_prompt_token_ids, admission=admission, @@ -1159,7 +925,7 @@ async def create_chat_completion( # max_model_len during tokenization, instead of returning an # ErrorResponse. Convert to HTTP 400 so the Gym proxy can # detect context-length overflow and handle it gracefully. - worker_self._abort_request_capture(request, reason="context_length") + worker_self.capture_host.abort_request(request, reason="context_length") return JSONResponse( content={ "error": { @@ -1172,11 +938,11 @@ async def create_chat_completion( status_code=400, ) except BaseException: - worker_self._abort_request_capture(request, reason="engine_error") + worker_self.capture_host.abort_request(request, reason="engine_error") raise if isinstance(generator, ErrorResponse): - worker_self._abort_request_capture(request, reason="error_response") + worker_self.capture_host.abort_request(request, reason="error_response") return JSONResponse( content=generator.model_dump(), status_code=generator.error.code ) @@ -1190,11 +956,11 @@ async def create_chat_completion( # Off-loop: the sink write inside complete_call is a blocking # TQ round trip (see the staging protocol's serving-host rule). content = await asyncio.to_thread( - worker_self._finish_request_capture, request, content + worker_self.capture_host.finish_request, request, content ) return JSONResponse(content=content) - worker_self._abort_request_capture(request, reason="streaming_response") + worker_self.capture_host.abort_request(request, reason="streaming_response") return StreamingResponse(content=generator, media_type="text/event-stream") ######################################## diff --git a/pyrefly.toml b/pyrefly.toml index 01788d1fb52..3133469edcb 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -24,6 +24,8 @@ replace-imports-with-any = [ "zstandard.*", ] project-includes = [ + "nemo_rl/experience/rollout_reassembler_pool.py", + "nemo_rl/models/generation/vllm/token_capture_host.py", # TODO: enable these once we have 100 correctness #"nemo_rl/**/*.py", #"examples/**/*.py", diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index e6c191df657..72317eb912e 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -45,7 +45,7 @@ from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data.processors import nemo_gym_data_processor from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.experience.failures import GenerationUnavailable +from nemo_rl.experience.failures import GenerationUnavailable, RolloutDataFailure from nemo_rl.experience.interfaces import ( NEMO_GYM_GROUP_ATTEMPT_KEY, NEMO_GYM_GROUP_ID_KEY, @@ -1776,6 +1776,110 @@ async def run_rollout( class TestGenerateForFinalizationFlow: + @pytest.mark.parametrize("lineage_owned", [False, True]) + def test_data_skip_releases_only_internally_owned_lineage(self, lineage_owned): + buf = _FakeCaptureBuffer() + calls = 0 + + async def fail(_sample): + nonlocal calls + calls += 1 + raise ValueError("bad prompt") + + mgr = _make_capture_manager( + buf, + on_run=fail, + retry_policy=RolloutRetryPolicy.single_attempt( + max_data_attempts=2, + max_skipped_prompts=1, + ), + ) + + async def scenario(): + sample = {"prompt": "p", "idx": 0} + group_id = None + if lineage_owned: + async with mgr._recovery_mutation() as cut: + group_id = mgr.reserve_prompt_group( + cut, + sample, + target_step=3, + admitted=True, + ) + assert ( + await mgr.generate_for_finalization( + sample, + lineage_group_id=group_id, + ) + is None + ) + assert calls == 2 + assert not buf._slots + assert len(mgr.recovery_ledger) == int(lineage_owned) + assert mgr.stats.skipped == 1 + assert mgr.stats.data_retries_by_reason == {"ValueError": 1} + assert mgr.stats.data_failures_by_reason == {"ValueError": 1} + with pytest.raises(RolloutDataFailure, match="max_skipped_prompts=1"): + await mgr.generate_for_finalization(sample) + assert calls == 4 + assert mgr.stats.skipped == 1 + assert len(mgr.recovery_ledger) == int(lineage_owned) + + _run(scenario()) + + def test_mixed_retry_classes_have_independent_capture_budgets(self): + errors = iter([ValueError("bad data"), GenerationUnavailable("lost"), None]) + + async def attempt(_sample): + error = next(errors) + if error is not None: + raise error + + mgr = _make_capture_manager( + _FakeCaptureBuffer(), + on_run=attempt, + retry_policy=RolloutRetryPolicy.single_attempt( + max_data_attempts=2, + max_infra_attempts=2, + backoff_base_s=0, + ), + ) + assert ( + _run(mgr.generate_for_finalization({"prompt": "p", "idx": 0})) is not None + ) + assert mgr.stats.data_retries_by_reason == {"ValueError": 1} + assert mgr.stats.redispatches_by_reason == {"GenerationUnavailable": 1} + assert mgr.stats.data_failures_by_reason == {} + assert mgr.stats.skipped == 0 + + def test_capture_cleanup_failure_is_never_retried(self): + class BrokenBuffer(_FakeCaptureBuffer): + def abort(self, group_id): + raise RuntimeError("abort failed") + + calls = 0 + + async def fail(_sample): + nonlocal calls + calls += 1 + raise GenerationUnavailable("lost") + + mgr = _make_capture_manager( + BrokenBuffer(), + on_run=fail, + retry_policy=RolloutRetryPolicy.single_attempt( + max_data_attempts=3, + max_infra_attempts=3, + backoff_base_s=0, + max_skipped_prompts=5, + ), + ) + with pytest.raises(RuntimeError, match="capture cleanup failed"): + _run(mgr.generate_for_finalization({"prompt": "p", "idx": 0})) + assert calls == 1 + assert mgr.stats.redispatches_by_reason == {} + assert mgr.stats.data_retries_by_reason == {} + def test_request_carries_env_mask_flags(self): buf = _FakeCaptureBuffer() mgr = _make_capture_manager( @@ -1874,7 +1978,7 @@ async def run_rollout( assert first_rollout_ids is not None assert buf.cleared_staging_key_batches == [[f"{first_rollout_ids[0]}/call"]] assert ( - "dropping capture prompt idx=0 after 1 infrastructure failure(s) " + "dropping prompt idx=0 after 1 infrastructure failure(s) " "(GenerationUnavailable: worker disappeared) [consecutive drop 1/1]" in capsys.readouterr().out ) diff --git a/tests/unit/experience/test_rollout_reassembler_pool.py b/tests/unit/experience/test_rollout_reassembler_pool.py new file mode 100644 index 00000000000..be6d9de8eaf --- /dev/null +++ b/tests/unit/experience/test_rollout_reassembler_pool.py @@ -0,0 +1,91 @@ +# 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. + +import asyncio +from unittest.mock import Mock + +import pytest + +from nemo_rl.experience.rollout_reassembler_pool import RolloutReassemblerPool + + +def test_contention_cancellation_and_unknown_outcome_quarantine(): + async def scenario(): + actor = object() + pool = RolloutReassemblerPool([actor]) + lease = await pool.acquire() + waiter = asyncio.create_task(pool.acquire()) + await asyncio.sleep(0) + assert pool.active == 1 + assert pool.waiters == 1 + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + assert pool.waiters == 0 + pool.release(lease) + reused = await pool.acquire() + assert reused.actor is actor + reused.rpc_submitted = True + pool.release(reused) + assert pool.active == 0 + assert pool.available_count == 0 + assert pool.unknown_outcomes == 1 + + asyncio.run(scenario()) + + +def test_known_outcome_releases_waiter_and_records_queue_depth(): + async def scenario(): + pool = RolloutReassemblerPool([object()]) + lease = await pool.acquire() + waiter = asyncio.create_task(pool.acquire()) + await asyncio.sleep(0) + lease.rpc_submitted = True + lease.outcome_known = True + pool.release(lease) + next_lease = await waiter + assert next_lease.actor is lease.actor + assert next_lease.queue_depth == 1 + assert next_lease.queue_wait_ms >= 0 + assert next_lease.active_actor_count == 1 + pool.release(next_lease) + assert pool.unknown_outcomes == 0 + assert pool.available_count == 1 + + asyncio.run(scenario()) + + +def test_shutdown_wakes_waiters_and_terminates_quarantined_actors(monkeypatch): + kill = Mock() + monkeypatch.setattr("nemo_rl.experience.rollout_reassembler_pool.ray.kill", kill) + + async def scenario(): + actor = object() + pool = RolloutReassemblerPool([actor]) + lease = await pool.acquire() + lease.rpc_submitted = True + pool.release(lease) + waiter = asyncio.create_task(pool.acquire()) + await asyncio.sleep(0) + pool.shutdown() + pool.shutdown() + with pytest.raises(RuntimeError, match="closed"): + await waiter + with pytest.raises(RuntimeError, match="closed"): + await pool.acquire() + assert pool.waiters == 0 + assert pool.available_count == 0 + kill.assert_called_once_with(actor, no_restart=True) + + asyncio.run(scenario()) diff --git a/tests/unit/models/generation/test_token_capture_host.py b/tests/unit/models/generation/test_token_capture_host.py new file mode 100644 index 00000000000..b781a6ec649 --- /dev/null +++ b/tests/unit/models/generation/test_token_capture_host.py @@ -0,0 +1,74 @@ +# 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. + +"""Capture host isolation and disabled-path behavior without Gym installed.""" + +import builtins +from types import SimpleNamespace +from unittest.mock import MagicMock + +from nemo_rl.models.generation.vllm.token_capture_host import TokenCaptureHost + + +def test_disabled_host_never_imports_gym(monkeypatch): + original_import = builtins.__import__ + + def reject_gym(name, *args, **kwargs): + if name.startswith("nemo_gym"): + raise AssertionError("disabled capture imported Gym") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", reject_gym) + host = TokenCaptureHost() + request = SimpleNamespace(ng_capture={"rollout_id": "r"}, stream=False) + host.set_weight_version(7) + assert host.admission(request) is None + host.begin_request(request, [1, 2]) + response = {"choices": []} + assert host.finish_request(request, response) is response + host.abort_request(request, reason="engine_error") + + +def test_active_calls_are_isolated_and_abort_only_once(): + first, second = TokenCaptureHost(), TokenCaptureHost() + capture = MagicMock() + first.install_token_capture(capture) + request = SimpleNamespace(stream=False) + first.begin_request(request, [1, 2], admission=SimpleNamespace()) + + # An unrelated host cannot finish or abort another worker's active call. + response = {"choices": []} + assert second.finish_request(request, response) is response + second.abort_request(request, reason="engine_error") + capture.fail_call.assert_not_called() + first.abort_request(request, reason="engine_error") + first.abort_request(request, reason="engine_error") + capture.fail_call.assert_called_once_with( + capture.begin_call.return_value, reason="engine_error" + ) + + +def test_prefix_cache_is_local_and_returned_values_do_not_mutate_cache(): + first, second = TokenCaptureHost(), TokenCaptureHost() + first._staging_source = MagicMock() + second._staging_source = MagicMock() + first._staging_source.fetch_prefix_token_ids.return_value = [1, 2] + second._staging_source.fetch_prefix_token_ids.return_value = [3, 4] + assert first.fetch_chain_prefix(["same/key"]) == [1, 2] + cached = first.fetch_chain_prefix(["same/key"]) + cached.append(99) + assert first.fetch_chain_prefix(["same/key"]) == [1, 2] + assert second.fetch_chain_prefix(["same/key"]) == [3, 4] + first._staging_source.fetch_prefix_token_ids.assert_called_once() + second._staging_source.fetch_prefix_token_ids.assert_called_once() diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index e33a6630900..2b04da85d22 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -813,8 +813,9 @@ 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.token_capture = None + from nemo_rl.models.generation.vllm.token_capture_host import TokenCaptureHost + + worker.capture_host = TokenCaptureHost() worker.llm_async_engine_args = MagicMock() worker.llm_async_engine_args.create_model_config.return_value = MagicMock( served_model_name="served-model", model="model-path" 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 4f59c59c413..681193f5791 100644 --- a/tests/unit/models/generation/test_vllm_token_capture_hosting.py +++ b/tests/unit/models/generation/test_vllm_token_capture_hosting.py @@ -23,7 +23,6 @@ from __future__ import annotations import asyncio -import threading from types import SimpleNamespace from unittest.mock import MagicMock @@ -46,6 +45,8 @@ VllmAsyncGenerationWorkerImpl, ) +from nemo_rl.models.generation.vllm.token_capture_host import TokenCaptureHost + pytestmark = pytest.mark.nemo_gym @@ -60,18 +61,10 @@ def stage(self, record: StagedCallRecord) -> StageResult: def _fake_worker(*, is_model_owner: bool = True) -> SimpleNamespace: """The attribute surface setup_token_capture touches, minus the engine.""" - worker = SimpleNamespace( + return SimpleNamespace( is_model_owner=is_model_owner, - token_capture=None, - _rollout_weight_version=0, - _staging_source=None, - _prefix_cache={}, - _prefix_cache_lock=threading.Lock(), - ) - worker.install_token_capture = lambda capture: setattr( - worker, "token_capture", capture + capture_host=TokenCaptureHost(), ) - return worker def test_setup_token_capture_installs_capture_with_vllm_adapter(monkeypatch): @@ -93,10 +86,10 @@ def test_setup_token_capture_installs_capture_with_vllm_adapter(monkeypatch): ) assert installed is True - assert isinstance(worker.token_capture, RolloutTokenCapture) - assert worker.token_capture.adapter is not None + assert isinstance(worker.capture_host.token_capture, RolloutTokenCapture) + assert worker.capture_host.token_capture.adapter is not None # The adapter is the vLLM one (prefix ids enter via the worker's field). - payload = worker.token_capture.adapter.enter_prefix({}, [1, 2]) + payload = worker.capture_host.token_capture.adapter.enter_prefix({}, [1, 2]) assert payload["required_prefix_token_ids"] == [1, 2] @@ -108,7 +101,7 @@ def test_setup_token_capture_skips_non_model_owners(monkeypatch): ) ) assert installed is False - assert worker.token_capture is None + assert worker.capture_host.token_capture is None def test_weight_version_is_stamped_from_worker_state(monkeypatch): @@ -131,17 +124,17 @@ def test_weight_version_is_stamped_from_worker_state(monkeypatch): ) asyncio.run(VllmAsyncGenerationWorkerImpl.set_rollout_weight_version(worker, 4)) - first = worker.token_capture.begin_call( + first = worker.capture_host.token_capture.begin_call( CaptureAdmission(rollout_id="r", model_call_id="c1", mode="text") ) asyncio.run(VllmAsyncGenerationWorkerImpl.set_rollout_weight_version(worker, 5)) - second = worker.token_capture.begin_call( + second = worker.capture_host.token_capture.begin_call( CaptureAdmission(rollout_id="r", model_call_id="c2", mode="text") ) assert (first.weight_version, second.weight_version) == (4, 5) - coords = worker.token_capture.complete_call( + coords = worker.capture_host.token_capture.complete_call( first, prompt_token_ids=[1], generated_token_ids=[2], generated_logprobs=[-0.1] ) assert coords.weight_version == 4 @@ -203,23 +196,7 @@ class _FakeRequest(SimpleNamespace): def _worker_with_capture(sink: _MemorySink): from nemo_gym.token_id_capture.adapters.vllm import VLLMCaptureAdapter - worker = _fake_worker() - worker._capture_calls = {} - worker._prefix_cache = {} - worker._prefix_cache_lock = threading.Lock() - worker._staging_source = None - worker._delta_align_routed_experts = ( - VllmAsyncGenerationWorkerImpl._delta_align_routed_experts - ) - for name in ( - "_fetch_chain_prefix", - "_capture_admission", - "_resolve_admission_prefix", - "_enter_request_prefix", - ): - setattr( - worker, name, getattr(VllmAsyncGenerationWorkerImpl, name).__get__(worker) - ) + worker = TokenCaptureHost() worker.token_capture = RolloutTokenCapture( sink=sink, weight_version_fn=lambda: worker._rollout_weight_version, @@ -268,13 +245,11 @@ def test_request_capture_round_trip_stages_and_rides_coords(): }, stream=False, ) - VllmAsyncGenerationWorkerImpl._begin_request_capture(worker, request, [10, 11, 12]) + TokenCaptureHost.begin_request(worker, request, [10, 11, 12]) content = _served_content([13, 14], [-0.1, -0.2]) # Full-length routes on the served response must not survive the strip. content["choices"][0]["message"]["routed_experts"] = [[[0]]] * 5 - content = VllmAsyncGenerationWorkerImpl._finish_request_capture( - worker, request, content - ) + content = TokenCaptureHost.finish_request(worker, request, content) # Bytes were staged before the coords existed (fail-closed ordering). assert len(sink.records) == 1 assert sink.records[0].token_ids_delta == [10, 11, 12, 13, 14] @@ -310,10 +285,8 @@ def test_request_capture_token_in_prev_len_chains(): stream=False, ) spliced_prompt = [10, 11, 12, 20, 21] # exact prefix + fresh suffix - VllmAsyncGenerationWorkerImpl._begin_request_capture( - worker, request, spliced_prompt - ) - content = VllmAsyncGenerationWorkerImpl._finish_request_capture( + TokenCaptureHost.begin_request(worker, request, spliced_prompt) + content = TokenCaptureHost.finish_request( worker, request, _served_content([22], [-0.5]) ) coords = content["ng_commit_coords"] @@ -347,10 +320,10 @@ def test_staging_chain_prefix_flows_through_adapter_and_begin_call(): request = _staging_chain_request() context_before = dict(request.ng_capture) - admission = worker._capture_admission(request) - prefix = worker._resolve_admission_prefix(admission) - worker._enter_request_prefix(request, prefix) - VllmAsyncGenerationWorkerImpl._begin_request_capture( + admission = worker.admission(request) + prefix = worker.resolve_prefix(admission) + worker.enter_prefix(request, prefix) + TokenCaptureHost.begin_request( worker, request, prefix + [20], @@ -385,15 +358,15 @@ def test_inline_prefix_admission_resolves_without_a_fetch(): }, stream=False, ) - admission = worker._capture_admission(request) - assert worker._resolve_admission_prefix(admission) == [10, 11] + admission = worker.admission(request) + assert worker.resolve_prefix(admission) == [10, 11] assert worker._staging_source.calls == [] - text_root = worker._capture_admission( + text_root = worker.admission( _FakeRequest( ng_capture={"rollout_id": "r0", "model_call_id": "c1", "mode": "text"} ) ) - assert worker._resolve_admission_prefix(text_root) == [] + assert worker.resolve_prefix(text_root) == [] def test_staging_chain_cache_fetches_only_uncached_suffix(): @@ -401,10 +374,8 @@ def test_staging_chain_cache_fetches_only_uncached_suffix(): source = _MemoryPrefixSource({"r0/c1": [10, 11], "r0/c2": [12]}) worker._staging_source = source - first = VllmAsyncGenerationWorkerImpl._fetch_chain_prefix(worker, ["r0/c1"]) - second = VllmAsyncGenerationWorkerImpl._fetch_chain_prefix( - worker, ["r0/c1", "r0/c2"] - ) + first = TokenCaptureHost.fetch_chain_prefix(worker, ["r0/c1"]) + second = TokenCaptureHost.fetch_chain_prefix(worker, ["r0/c1", "r0/c2"]) assert first == [10, 11] assert second == [10, 11, 12] @@ -418,11 +389,11 @@ def test_staging_chain_prefix_length_mismatch_is_rejected_by_begin_call(): request = _staging_chain_request(prev_len=3) context_before = dict(request.ng_capture) - admission = worker._capture_admission(request) - prefix = worker._resolve_admission_prefix(admission) + admission = worker.admission(request) + prefix = worker.resolve_prefix(admission) assert prefix == [10, 11] with pytest.raises(CaptureError, match="does not equal prev_len 3"): - VllmAsyncGenerationWorkerImpl._begin_request_capture( + TokenCaptureHost.begin_request( worker, request, prefix + [20], admission=admission, prefix_token_ids=prefix ) @@ -435,9 +406,7 @@ def test_staging_chain_admission_requires_the_resolved_prefix_keyword(): request = _staging_chain_request() with pytest.raises(CaptureError, match="pass the resolved prefix_token_ids"): - VllmAsyncGenerationWorkerImpl._begin_request_capture( - worker, request, [10, 11, 12, 20] - ) + TokenCaptureHost.begin_request(worker, request, [10, 11, 12, 20]) assert worker._capture_calls == {} @@ -446,13 +415,11 @@ def test_request_capture_is_a_noop_without_context_or_capture(): sink = _MemorySink() worker = _worker_with_capture(sink) plain = _FakeRequest(stream=False) # no ng_capture attribute - VllmAsyncGenerationWorkerImpl._begin_request_capture(worker, plain, [1, 2]) + TokenCaptureHost.begin_request(worker, plain, [1, 2]) content = { "choices": [{"message": {"role": "assistant"}, "logprobs": {"content": []}}] } - out = VllmAsyncGenerationWorkerImpl._finish_request_capture( - worker, plain, dict(content) - ) + out = TokenCaptureHost.finish_request(worker, plain, dict(content)) assert "ng_commit_coords" not in out assert out["choices"][0]["logprobs"] is not None # untouched off the capture path assert sink.records == [] @@ -471,14 +438,10 @@ def test_request_capture_abort_fails_the_call_and_drains_state(): }, stream=False, ) - VllmAsyncGenerationWorkerImpl._begin_request_capture(worker, request, [1, 2]) - VllmAsyncGenerationWorkerImpl._abort_request_capture( - worker, request, reason="engine_error" - ) + TokenCaptureHost.begin_request(worker, request, [1, 2]) + TokenCaptureHost.abort_request(worker, request, reason="engine_error") assert worker._capture_calls == {} assert sink.records == [] # A late finish after abort is a no-op (state already drained). - out = VllmAsyncGenerationWorkerImpl._finish_request_capture( - worker, request, _served_content([3], [-0.1]) - ) + out = TokenCaptureHost.finish_request(worker, request, _served_content([3], [-0.1])) assert "ng_commit_coords" not in out diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 668ec6d52de..68f939f9f4f 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -730,7 +730,7 @@ def _make_actor_args( save_state if save_state is not None else _initial_grpo_save_state() ), last_checkpoint_path=last_checkpoint_path, - finalizer_actors=[], + reassembler_actors=[], data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, bootstrap_identity=bootstrap_identity, ) diff --git a/tests/unit/single_controller/test_finalizer_lifecycle.py b/tests/unit/single_controller/test_finalizer_lifecycle.py index 049e68855d4..c1f00c70db5 100644 --- a/tests/unit/single_controller/test_finalizer_lifecycle.py +++ b/tests/unit/single_controller/test_finalizer_lifecycle.py @@ -32,6 +32,7 @@ from nemo_rl.data_plane.schema import ROUTE_PLAN_TAG from nemo_rl.experience.rollout_reassembler import FinalizedGroup from nemo_rl.experience.rollout_reassembler_actor import ReassemblyRequest +from nemo_rl.experience.rollout_reassembler_pool import RolloutReassemblerPool from nemo_rl.experience.route_plan import ( ROUTE_PLAN_SCHEMA_VERSION, RouteAssemblyPlan, @@ -96,11 +97,7 @@ def _request() -> ReassemblyRequest: def _controller(actor: object) -> Any: controller_cls = SingleControllerActor.__ray_metadata__.modified_class ctrl = object.__new__(controller_cls) - ctrl._available_finalizers = asyncio.Queue() - ctrl._available_finalizers.put_nowait(actor) - ctrl._active_finalizers = 0 - ctrl._finalizer_waiters = 0 - ctrl._finalizer_unknown_outcomes = 0 + ctrl._reassembler_pool = RolloutReassemblerPool([actor]) ctrl._finalizer_metrics_by_group = {} ctrl._rollout_recovery_ledger = MagicMock() ctrl._rollout_recovery_ledger.__contains__.return_value = False @@ -144,9 +141,9 @@ def test_successful_actor_finalization_returns_actor_and_transfers_ownership() - assert committed is result assert finalize.calls == [request] - assert ctrl._available_finalizers.get_nowait() is actor - assert ctrl._active_finalizers == 0 - assert ctrl._finalizer_unknown_outcomes == 0 + assert ctrl._reassembler_pool.available_count == 1 + assert ctrl._reassembler_pool.active == 0 + assert ctrl._reassembler_pool.unknown_outcomes == 0 ctrl._buffer.commit_finalized.assert_awaited_once_with( ANY, "group", @@ -168,9 +165,9 @@ def test_actor_rpc_failure_is_fatal_and_does_not_retry_or_requeue_actor() -> Non asyncio.run(ctrl._finalize_with_actor(request)) assert finalize.calls == [request] - assert ctrl._available_finalizers.empty() - assert ctrl._active_finalizers == 0 - assert ctrl._finalizer_unknown_outcomes == 1 + assert ctrl._reassembler_pool.available_count == 0 + assert ctrl._reassembler_pool.active == 0 + assert ctrl._reassembler_pool.unknown_outcomes == 1 ctrl._buffer.commit_finalized.assert_not_awaited() ctrl._buffer.abort.assert_not_called() assert ctrl._dp_client.clear_calls == [] @@ -195,7 +192,7 @@ def test_missing_actor_metadata_cleans_known_canonical_and_staging_ownership() - {"sample_ids": ["group_g0/call"], "partition_id": "staging"}, ] ctrl._buffer.abort.assert_called_once_with("group") - assert ctrl._available_finalizers.get_nowait() is actor + assert ctrl._reassembler_pool.available_count == 1 def test_dropped_actor_group_cleans_ownership_and_returns_uncommitted() -> None: @@ -228,7 +225,7 @@ def test_dropped_actor_group_cleans_ownership_and_returns_uncommitted() -> None: ctrl._buffer.abort.assert_called_once_with("group") ctrl._buffer.commit_finalized.assert_not_awaited() assert "group" not in ctrl._finalizer_metrics_by_group - assert ctrl._available_finalizers.get_nowait() is actor + assert ctrl._reassembler_pool.available_count == 1 def test_post_train_cleanup_clears_canonical_rows_and_route_plan_staging_keys() -> None: @@ -309,3 +306,78 @@ async def _main() -> int: assert ctrl._dp_client.clear_thread_ids assert all(tid != loop_thread_id for tid in ctrl._dp_client.clear_thread_ids) ctrl._buffer.abort.assert_called_once_with("group") + + +def test_checkpoint_waits_for_actor_publication_and_replay_commit() -> None: + async def scenario() -> None: + rpc_started = asyncio.Event() + release_rpc = asyncio.Event() + commit_started = asyncio.Event() + release_commit = asyncio.Event() + checkpoint_entered = asyncio.Event() + result = FinalizedGroup( + meta=KVBatchMeta( + partition_id="canonical", + task_name="train", + sample_ids=["group_g0"], + fields=["input_ids"], + sequence_lengths=[3], + tags=[{"weight_version": 3}], + ), + group_min_wv=3, + group_max_wv=3, + staging_keys=["group_g0/call"], + metrics={}, + ) + + async def remote(_request): + rpc_started.set() + await release_rpc.wait() + return result + + ctrl = _controller(SimpleNamespace(finalize=SimpleNamespace(remote=remote))) + + async def commit(*args, **kwargs): + commit_started.set() + await release_commit.wait() + + ctrl._buffer.commit_finalized.side_effect = commit + + async def checkpoint(): + async with ctrl._data_plane_checkpoint_barrier.checkpoint(): + ctrl._rollout_recovery_ledger.discard_group.assert_called_once() + checkpoint_entered.set() + + finalize_task = asyncio.create_task(ctrl._finalize_with_actor(_request())) + await rpc_started.wait() + checkpoint_task = asyncio.create_task(checkpoint()) + await asyncio.sleep(0) + assert not checkpoint_entered.is_set() + release_rpc.set() + await commit_started.wait() + assert not checkpoint_entered.is_set() + release_commit.set() + await asyncio.wait_for(asyncio.gather(finalize_task, checkpoint_task), 1) + assert checkpoint_entered.is_set() + assert ctrl._reassembler_pool.available_count == 1 + + asyncio.run(scenario()) + + +def test_cancel_waiting_for_actor_cleans_known_request() -> None: + async def scenario() -> None: + ctrl = _controller(object()) + lease = await ctrl._reassembler_pool.acquire() + task = asyncio.create_task(ctrl._finalize_with_actor(_request())) + await asyncio.sleep(0) + assert ctrl._reassembler_pool.waiters == 1 + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert ctrl._reassembler_pool.waiters == 0 + assert ctrl._reassembler_pool.active == 1 + ctrl._buffer.abort.assert_called_once_with("group") + ctrl._buffer.commit_finalized.assert_not_awaited() + ctrl._reassembler_pool.release(lease) + + asyncio.run(scenario()) diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 34db681229f..3891a0931c1 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -50,6 +50,7 @@ from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.rollout_manager import RolloutManager, RolloutOutcome +from nemo_rl.experience.rollout_reassembler_pool import RolloutReassemblerPool from nemo_rl.experience.rollout_recovery import ( RolloutRecoveryLedger, RolloutRecoveryState, @@ -102,7 +103,7 @@ def _init_pump_ledgers(ctrl: Any) -> None: ctrl._batch_replacements = {} ctrl._batch_promotions = {} # Empty means the legacy (non-token-capture) dispatch path. - ctrl._finalizer_actors = [] + ctrl._reassembler_pool = RolloutReassemblerPool([]) ctrl._replacement_reserve = deque() ctrl._rollout_recovery_enabled = False @@ -196,7 +197,7 @@ def test_rollout_pump_stamps_target_steps( ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = _RecordingRolloutManager(buffer) - ctrl._finalizer_actors = [] + ctrl._reassembler_pool = RolloutReassemblerPool([]) # The sampler owns admission + target_step stamping (the dispatch counter # lives on the sampler, not the actor). ctrl._sampler = make_sampler(buffer) @@ -1019,7 +1020,7 @@ async def _main() -> None: ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = manager - ctrl._finalizer_actors = [] + ctrl._reassembler_pool = RolloutReassemblerPool([]) # Over-sampled windowed policy: admit never gates (buffer unused here). ctrl._sampler = WindowedSampler(None, max_staleness_versions=1) ctrl._dataloader = [ @@ -1109,7 +1110,7 @@ async def _main() -> None: ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = _NeverCalledRolloutManager() - ctrl._finalizer_actors = [] + ctrl._reassembler_pool = RolloutReassemblerPool([]) # Over-sampled windowed policy: admit never gates (buffer unused here). ctrl._sampler = WindowedSampler(None, max_staleness_versions=1) ctrl._dataloader = [ @@ -1188,7 +1189,7 @@ async def _delayed_finalize( ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = manager _init_pump_ledgers(ctrl) - ctrl._finalizer_actors = [object()] + ctrl._reassembler_pool = RolloutReassemblerPool([object()]) ctrl._finalize_with_actor = _delayed_finalize ctrl._sampler = WindowedSampler(None, max_staleness_versions=1) ctrl._dataloader = [ @@ -1325,7 +1326,7 @@ async def _main() -> None: ctrl._sampler_stamps_target_steps = False ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() _init_pump_ledgers(ctrl) - ctrl._finalizer_actors = [object()] + ctrl._reassembler_pool = RolloutReassemblerPool([object()]) ctrl._rollout_recovery_enabled = True async def _finalize( @@ -1452,7 +1453,7 @@ def test_rollout_pump_writes_expected_tq_data( partition_id=_PARTITION_ID, save_state=_initial_grpo_save_state(), last_checkpoint_path=None, - finalizer_actors=[], + reassembler_actors=[], ) ctrl = SingleControllerActor.remote( master_config=master_config, @@ -1534,3 +1535,86 @@ def test_rollout_pump_writes_expected_tq_data( "num_assistant_messages", "num_routed_experts_backfilled", } + + +@pytest.mark.parametrize("recovery_enabled", [False, True]) +@pytest.mark.parametrize("replace", [False, True]) +def test_capture_skip_uses_replacement_policy_and_releases_ownership( + recovery_enabled, replace +): + class CaptureManager: + def __init__(self): + self.recovery_ledger = RolloutRecoveryLedger() + self.stats = SimpleNamespace(committed=0) + self.prompts_seen = [] + + def reserve_prompt_group( + self, cut, prompt, *, target_step, admitted=True, admission_id=None + ): + return self.recovery_ledger.reserve_group( + cut, + prompt_id=prompt["message_log"][0]["content"], + prompt_payload=prompt, + expected_generations=1, + target_step=target_step, + start_weight_version=0, + admitted=admitted, + admission_id=admission_id, + ).group_id + + def mark_prompt_group_admitted(self, cut, group_id, *, target_step): + self.recovery_ledger.mark_group_admitted( + cut, + group_id, + target_step=target_step, + start_weight_version=0, + ) + + def discard_prompt_group(self, cut, group_id): + self.recovery_ledger.discard_group(cut, group_id) + + async def discard_recovery_group(self, cut, group_id): + self.discard_prompt_group(cut, group_id) + + async def generate_for_finalization( + self, prompt, *, target_step, inflight_registry, lineage_group_id=None + ): + content = prompt["message_log"][0]["content"] + self.prompts_seen.append(content) + if lineage_group_id is not None: + assert self.recovery_ledger.get_group(lineage_group_id) + if content == "bad": + return None + return SimpleNamespace(group_id=lineage_group_id) + + async def scenario(): + manager = CaptureManager() + ctrl = _pump_controller( + manager, + [_batch("bad"), _batch("spare")], + on_dropped_prompt="replace" if replace else "shrink", + ) + ctrl._master_config.token_capture = SimpleNamespace( + min_valid_fraction_per_group=None, + ) + ctrl._reassembler_pool = RolloutReassemblerPool([object()]) + ctrl._rollout_recovery_enabled = recovery_enabled + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + + async def finalize(request): + if request.group_id is not None: + async with ctrl._data_plane_checkpoint_barrier.mutation() as cut: + manager.discard_prompt_group(cut, request.group_id) + return SimpleNamespace(valid_row_count=1, total_row_count=1) + + ctrl._finalize_with_actor = finalize + await ctrl._rollout_pump() + assert manager.prompts_seen == ["bad", "spare"] + assert ctrl._batch_shortfall == ({} if replace else {0: 1}) + assert ctrl._batch_replacements == ({0: 1} if replace else {}) + assert len(manager.recovery_ledger) == 0 + assert manager.stats.committed == 1 + assert ctrl._inflight_rollouts == 0 + assert ctrl._buffer_capacity._value == 3 + + asyncio.run(scenario()) diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index a8a180cba2a..769c174dae5 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -828,7 +828,7 @@ def test_periodic_checkpointing_warns_without_per_step_trainer_anchors( MagicMock(pad_token_id=0), ) - assert actor_args.finalizer_actors == fake_finalizers + assert actor_args.reassembler_actors == fake_finalizers def test_disabled_periodic_checkpointing_ignores_existing_snapshots( self, @@ -1291,7 +1291,7 @@ def test_returns_actor_args(self, patched_factories): assert actor_args.partition_id == "rollout_data" assert actor_args.tq_buffer._partition_id == "rollout_data" assert actor_args.tq_buffer._require_routed_experts is False - assert actor_args.finalizer_actors == [] + assert actor_args.reassembler_actors == [] actor_args.dp_client.register_partition.assert_called_once() warmup = actor_args.dp_client.register_partition.call_args.kwargs assert warmup["partition_id"] == "rollout_data" @@ -1614,19 +1614,19 @@ def test_token_capture_always_creates_finalizer_actor_pool(self, patched_factori patch( "nemo_rl.experience.rollout_reassembler_actor.create_rollout_reassembler_actors", return_value=fake_actors, - ) as mock_create_finalizer_actors, + ) as mock_create_reassembler_actors, ): actor_args, _ = setup_single_controller(mc, tokenizer, processor=processor) (actor_dp_config, actor_config), actor_kwargs = ( - mock_create_finalizer_actors.call_args + mock_create_reassembler_actors.call_args ) assert actor_dp_config == mc.data_plane assert actor_config.partition_id == "rollout_data" assert actor_config.staging_partition == mc.token_capture.staging_partition assert actor_config.pad_token_id == 9 assert actor_kwargs == {"num_workers": 3} - assert actor_args.finalizer_actors == fake_actors + assert actor_args.reassembler_actors == fake_actors assert not hasattr(actor_args.rollout_manager, "_finalizer") partition_calls = actor_args.dp_client.register_partition.call_args_list assert WIRE_MULTIMODAL_FIELDS <= set(partition_calls[0].kwargs["fields"]) diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 470ae99a3fb..798a3dbd00c 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -134,7 +134,7 @@ def _actor_args_for_init(**overrides) -> SimpleNamespace: inference_cluster=None, save_state=_initial_grpo_save_state(), last_checkpoint_path=None, - finalizer_actors=[], + reassembler_actors=[], data_plane_checkpoint_metadata=None, bootstrap_identity=None, ) diff --git a/tests/unit/single_controller/test_train_pump_e2e.py b/tests/unit/single_controller/test_train_pump_e2e.py index b994dce577a..4c5d5e10846 100644 --- a/tests/unit/single_controller/test_train_pump_e2e.py +++ b/tests/unit/single_controller/test_train_pump_e2e.py @@ -385,7 +385,7 @@ def test_train_pump_drives_mcore_training_step( partition_id=_PARTITION_ID, save_state=_initial_grpo_save_state(), last_checkpoint_path=None, - finalizer_actors=[], + reassembler_actors=[], ) ctrl = _RecordingSingleControllerActor.remote( metric_log_handle=log,