diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 2c7c2235853..e55dc91abeb 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -50,6 +50,12 @@ RolloutDataFailure, http_status_is_infra, ) +from nemo_rl.experience.interfaces import ( + NEMO_GYM_ATTEMPT_INDEX_KEY, + NEMO_GYM_CAPTURE_ID_KEY, + NEMO_GYM_ROLLOUT_ID_KEY, + nemo_gym_capture_key, +) from nemo_rl.models.generation.interfaces import ( resolve_routed_experts_dtype_name_for_model, should_use_async_rollouts, @@ -231,12 +237,10 @@ class NemoGymConfig(TypedDict): token_capture: NotRequired[Dict[str, Any] | None] -# Gym control-plane server name (the model server hosting the ledger) and the -# opaque run-body key rollout ids ride on (Gym's ROLLOUT_ID_KEY_NAME): the -# agent derives the id from the run body and stamps /ng-rollout/ on every -# model call, so the TQ sample id IS the capture key end to end. +# Gym control-plane server name for the model server hosting the capture ledger. +# Gym derives its physical capture key from the logical run-body rollout ID and +# numeric attempt index. _POLICY_SERVER_NAME = "policy_model" -_NG_ROLLOUT_ID_BODY_KEY = "_ng_rollout_id" _TOKEN_CAPTURE_CONTROL_PREFIX = "/training-token-capture/control" _TOKEN_CAPTURE_CONTROL_ENV = "NEMO_GYM_TOKEN_CAPTURE_CONTROL_TOKEN" @@ -602,6 +606,29 @@ async def run_rollouts( ) tokenizer = self._tokenizer + for row in nemo_gym_examples: + logical_rollout_id = row.get(NEMO_GYM_ROLLOUT_ID_KEY) + attempt_index = row.get(NEMO_GYM_ATTEMPT_INDEX_KEY) + if not isinstance(logical_rollout_id, str) or not logical_rollout_id: + raise ValueError( + f"{NEMO_GYM_ROLLOUT_ID_KEY} must be a non-empty string" + ) + if ( + isinstance(attempt_index, bool) + or not isinstance(attempt_index, int) + or attempt_index < 0 + ): + raise ValueError( + f"{NEMO_GYM_ATTEMPT_INDEX_KEY} must be a non-negative integer" + ) + expected_capture_id = nemo_gym_capture_key( + logical_rollout_id, attempt_index + ) + if row.get(NEMO_GYM_CAPTURE_ID_KEY) != expected_capture_id: + raise ValueError( + f"{NEMO_GYM_CAPTURE_ID_KEY} must equal {expected_capture_id!r}" + ) + from nemo_rl.utils.fastokens import maybe_patch_fastokens maybe_patch_fastokens(bool(self.cfg.get("use_fastokens"))) @@ -710,7 +737,7 @@ async def _postprocess_receipt_mode( assert isinstance(nemo_gym_result, dict), ( f"Hit a non-successful response when querying NeMo Gym for rollouts: {nemo_gym_result}" ) - rollout_id = nemo_gym_row[_NG_ROLLOUT_ID_BODY_KEY] + rollout_id = nemo_gym_row[NEMO_GYM_CAPTURE_ID_KEY] # Gym's TERMINAL_RESPONSE_ID_KEY: the served response envelope id the # harness kept (``response.id``), not the logical-request header. terminal_response_id = nemo_gym_result.get("terminal_response_id") diff --git a/nemo_rl/experience/interfaces.py b/nemo_rl/experience/interfaces.py index 29320e93f2b..09abe3002c3 100644 --- a/nemo_rl/experience/interfaces.py +++ b/nemo_rl/experience/interfaces.py @@ -21,6 +21,9 @@ NEMO_GYM_GROUP_ID_KEY = "_ng_group_id" NEMO_GYM_GROUP_ATTEMPT_KEY = "_ng_group_attempt" NEMO_GYM_ROLLOUT_INDEX_KEY = "_ng_rollout_index" +NEMO_GYM_ROLLOUT_ID_KEY = "_ng_rollout_id" +NEMO_GYM_ATTEMPT_INDEX_KEY = "_ng_attempt_index" +NEMO_GYM_CAPTURE_ID_KEY = "_ng_capture_id" NEXT_NEMO_GYM_TASK_INDEX_KEY = "next_ng_task_index" # Unconsumed suffix of a gap-fill dataloader batch, carried in the async # collector's rollouts state so a checkpoint cannot strand yielded prompts. @@ -41,6 +44,15 @@ TRAINED_TASK_INDICES_KEY = "trained_task_indices" +def nemo_gym_capture_key(logical_rollout_id: str, attempt_index: int) -> str: + """Return Gym's physical capture key for one logical rollout attempt.""" + if attempt_index < 0: + raise ValueError("attempt_index must be non-negative") + if attempt_index == 0: + return logical_rollout_id + return f"{logical_rollout_id}-a{attempt_index}" + + @dataclass class Completion: """A single generated completion for one prompt.""" diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 21f8cc68fc0..9ba2a1779d5 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -22,7 +22,7 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from contextlib import asynccontextmanager from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any, Optional, cast import ray.exceptions import torch @@ -51,11 +51,15 @@ classify_rollout_failure, ) from nemo_rl.experience.interfaces import ( + NEMO_GYM_ATTEMPT_INDEX_KEY, + NEMO_GYM_CAPTURE_ID_KEY, NEMO_GYM_GROUP_ATTEMPT_KEY, NEMO_GYM_GROUP_ID_KEY, + NEMO_GYM_ROLLOUT_ID_KEY, NEMO_GYM_ROLLOUT_INDEX_KEY, Completion, PromptGroupRecord, + nemo_gym_capture_key, ) from nemo_rl.experience.metric_utils import calculate_single_metric, pct from nemo_rl.experience.rollout_recovery import ( @@ -89,12 +93,39 @@ TokenizerType = PreTrainedTokenizerBase RolloutCompletionCallback = Callable[[int, Completion], Awaitable[None]] +RolloutAttemptAllocator = Callable[[list[int]], Awaitable[dict[int, int]]] if TYPE_CHECKING: from nemo_rl.algorithms.single_controller_utils.config import RolloutRecoveryConfig from nemo_rl.experience.rollout_reassembler_actor import ReassemblyRequest +@dataclass +class _LocalSiblingAttemptState: + """Run-local attempt ownership for ordinary NeMo-Gym rollouts.""" + + logical_rollout_ids: tuple[str, ...] + next_attempt_indices: list[int] + + @classmethod + def create(cls, group_id: str, sibling_count: int) -> "_LocalSiblingAttemptState": + return cls( + logical_rollout_ids=tuple( + f"{group_id}_g{generation_index}" + for generation_index in range(sibling_count) + ), + next_attempt_indices=[0] * sibling_count, + ) + + async def allocate(self, generation_indices: list[int]) -> dict[int, int]: + allocated: dict[int, int] = {} + for generation_index in generation_indices: + attempt_index = self.next_attempt_indices[generation_index] + allocated[generation_index] = attempt_index + self.next_attempt_indices[generation_index] = attempt_index + 1 + return allocated + + def _contains_post_write_enrichment_error(error: BaseException) -> bool: """Whether an error, including a rollback ExceptionGroup, is post-write.""" if isinstance(error, PostWriteEnrichmentError): @@ -944,19 +975,17 @@ async def run_rollout( self, input_sample: DatumSpec, *, - rollout_ids: Optional[list[str]] = None, + logical_rollout_ids: Optional[list[str]] = None, generation_indices: Optional[list[int]] = None, on_completion: Optional[RolloutCompletionCallback] = None, + attempt_allocator: Optional[RolloutAttemptAllocator] = None, recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, ) -> PromptGroupRecord: """Run num_generations_per_prompt rollouts for one prompt. Args: input_sample: A single prompt (one DatumSpec entry). - rollout_ids: Token-capture mode: gate-registered rollout ids, one - per generation, riding each row's run body as the opaque - ``_ng_rollout_id`` key (agents stamp /ng-rollout/ from it; - zero agent changes). + logical_rollout_ids: Stable sibling identities, one per generation. Returns: PromptGroupRecord with num_generations_per_prompt completions. @@ -967,14 +996,23 @@ async def run_rollout( rollout_inputs = self._build_inputs( input_sample, - rollout_ids=rollout_ids, + logical_rollout_ids=logical_rollout_ids, generation_indices=generation_indices, ) + if attempt_allocator is None: + local_state = _LocalSiblingAttemptState( + logical_rollout_ids=tuple( + row[NEMO_GYM_ROLLOUT_ID_KEY] for row in rollout_inputs + ), + next_attempt_indices=[0] * self._num_generations_per_prompt, + ) + attempt_allocator = local_state.allocate completions, prompt_message_log, rollout_metrics = await self._run_rollouts( rollout_inputs, timer, timer_prefix, on_completion=on_completion, + attempt_allocator=attempt_allocator, recovery_granularity=recovery_granularity, ) # Token-capture receipt rows carry empty message logs by design — the @@ -1035,7 +1073,7 @@ def _build_inputs( self, input_sample: DatumSpec, *, - rollout_ids: Optional[list[str]] = None, + logical_rollout_ids: Optional[list[str]] = None, generation_indices: Optional[list[int]] = None, ) -> list[dict]: """Build N row dicts from input_sample, applying generation config params.""" @@ -1058,9 +1096,9 @@ def _build_inputs( ) # Build N rows with distinct rowidxs so run_rollouts can sort them correctly. - if rollout_ids is not None: - assert len(rollout_ids) == self._num_generations_per_prompt, ( - "token-capture rollout ids must be one per generation" + if logical_rollout_ids is not None: + assert len(logical_rollout_ids) == self._num_generations_per_prompt, ( + "logical rollout ids must be one per generation" ) group_id = template_row.get(NEMO_GYM_GROUP_ID_KEY) or uuid.uuid4().hex group_attempt = template_row.get(NEMO_GYM_GROUP_ATTEMPT_KEY, 0) @@ -1090,11 +1128,12 @@ def _build_inputs( row[NEMO_GYM_GROUP_ID_KEY] = group_id row[NEMO_GYM_GROUP_ATTEMPT_KEY] = group_attempt row[NEMO_GYM_ROLLOUT_INDEX_KEY] = i - if rollout_ids is not None: - # Opaque run-body carrier (Gym's _ng_rollout_id key): the agent - # derives the id from the run body and stamps /ng-rollout/ - # on every model call, so the TQ sample id IS the capture key. - row["_ng_rollout_id"] = rollout_ids[i] + row[NEMO_GYM_ROLLOUT_ID_KEY] = ( + logical_rollout_ids[i] + if logical_rollout_ids is not None + else f"{group_id}_g{i}" + ) + row[NEMO_GYM_ATTEMPT_INDEX_KEY] = 0 rows.append(row) return rows @@ -1177,6 +1216,7 @@ async def _run_rollouts( timer_prefix: str, *, on_completion: Optional[RolloutCompletionCallback] = None, + attempt_allocator: Optional[RolloutAttemptAllocator] = None, recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, ) -> tuple[list[Completion], LLMMessageLogType, dict[str, Any]]: """Dispatch rows to NeMo-Gym; return completions, prompt, and metrics. @@ -1189,6 +1229,22 @@ async def _run_rollouts( if not inputs: raise ValueError("NeMo-Gym rollout dispatch requires at least one row") total_rows = self._num_generations_per_prompt + if attempt_allocator is None: + fallback_group_id = uuid.uuid4().hex + for row in inputs: + rowidx = row.get("_rowidx") + if isinstance(rowidx, int): + row.setdefault( + NEMO_GYM_ROLLOUT_ID_KEY, + f"{fallback_group_id}_g{rowidx}", + ) + local_state = _LocalSiblingAttemptState( + logical_rollout_ids=tuple( + f"{fallback_group_id}_g{index}" for index in range(total_rows) + ), + next_attempt_indices=[0] * total_rows, + ) + attempt_allocator = local_state.allocate # Re-dispatch maps NeMo-Gym's echoed _rowidx back onto the original group, so # the rows must carry the index _build_inputs stamped on them. Checked here # because the alternative is a KeyError several frames deeper. @@ -1238,6 +1294,28 @@ async def _run_rollouts( pending = [row for row in inputs if results[row["_rowidx"]] is None] if not pending: break + pending_indices = [row["_rowidx"] for row in pending] + attempt_indices = await attempt_allocator(pending_indices) + if set(attempt_indices) != set(pending_indices): + raise ValueError( + "attempt allocator must return every dispatched sibling" + ) + for row in pending: + rowidx = row["_rowidx"] + attempt_index = attempt_indices[rowidx] + if ( + isinstance(attempt_index, bool) + or not isinstance(attempt_index, int) + or attempt_index < 0 + ): + raise ValueError( + "attempt allocator returned an invalid attempt index" + ) + logical_rollout_id = row[NEMO_GYM_ROLLOUT_ID_KEY] + row[NEMO_GYM_ATTEMPT_INDEX_KEY] = attempt_index + row[NEMO_GYM_CAPTURE_ID_KEY] = nemo_gym_capture_key( + logical_rollout_id, attempt_index + ) if attempt > 1: print( f"NeMo-Gym: re-dispatching {len(pending)}/{total_rows} " @@ -1598,6 +1676,7 @@ def __init__( effort_config=effort_config, ) self._tokenizer = tokenizer + self._use_nemo_gym = use_nemo_gym self._num_generations_per_prompt = num_generations_per_prompt self._rollout_recovery_config = rollout_recovery_config self._tq_buffer = tq_buffer @@ -1735,22 +1814,25 @@ async def run_rollout( self, input_sample: DatumSpec, *, - rollout_ids: Optional[list[str]] = None, + logical_rollout_ids: Optional[list[str]] = None, generation_indices: Optional[list[int]] = None, on_completion: Optional[RolloutCompletionCallback] = None, + attempt_allocator: Optional[RolloutAttemptAllocator] = None, recovery_granularity: RecoveryGranularity = RecoveryGranularity.SIBLING, ) -> PromptGroupRecord: - if rollout_ids is None: + if logical_rollout_ids is None: assert generation_indices is None assert on_completion is None + assert attempt_allocator is None assert recovery_granularity is RecoveryGranularity.SIBLING - # Legacy path: keep the impl call signature byte-identical. return await self._impl.run_rollout(input_sample) - return await self._impl.run_rollout( + nemo_gym_impl = cast(AsyncNemoGymRolloutImpl, self._impl) + return await nemo_gym_impl.run_rollout( input_sample, - rollout_ids=rollout_ids, + logical_rollout_ids=logical_rollout_ids, generation_indices=generation_indices, on_completion=on_completion, + attempt_allocator=attempt_allocator, recovery_granularity=recovery_granularity, ) @@ -1821,6 +1903,7 @@ async def generate_and_push( data_attempts = 0 last_infra_error: Optional[Exception] = None logical_group_id: Optional[str] = None + local_attempt_state: Optional[_LocalSiblingAttemptState] = None group_attempt = 0 extra_env_info = input_sample.get("extra_env_info") if isinstance(extra_env_info, dict): @@ -1830,6 +1913,11 @@ async def generate_and_push( ): raise ValueError(f"{NEMO_GYM_GROUP_ID_KEY} must be a non-empty string") logical_group_id = configured_group_id or uuid.uuid4().hex + if self._use_nemo_gym: + local_attempt_state = _LocalSiblingAttemptState.create( + logical_group_id, + self._num_generations_per_prompt, + ) configured_group_attempt = extra_env_info.get(NEMO_GYM_GROUP_ATTEMPT_KEY, 0) if ( not isinstance(configured_group_attempt, int) @@ -1873,7 +1961,16 @@ async def generate_and_push( attempt_extra_env_info[NEMO_GYM_GROUP_ATTEMPT_KEY] = ( group_attempt ) - record = await self.run_rollout(attempt_input_sample) + if local_attempt_state is None: + record = await self.run_rollout(attempt_input_sample) + else: + record = await self.run_rollout( + attempt_input_sample, + logical_rollout_ids=list( + local_attempt_state.logical_rollout_ids + ), + attempt_allocator=local_attempt_state.allocate, + ) finally: if inflight_registry is not None: inflight_registry.pop(tq_group_id, None) @@ -2129,12 +2226,12 @@ async def _generate_for_finalization_attempt( from nemo_rl.experience.rollout_reassembler_actor import ReassemblyRequest assert self._tq_buffer is not None - async with self._recovery_mutation() as cut: - recovery_group = self._recovery_ledger.get_group(recovery_group_id) - if recovery_group.status == PromptGroupStatus.GENERATING: - recovery_group = self._recovery_ledger.prepare_incomplete_retry( - cut, recovery_group_id - ) + recovery_group = self._recovery_ledger.get_group(recovery_group_id) + if recovery_group.status is not PromptGroupStatus.GENERATING: + raise ValueError( + f"cannot generate recovery group {recovery_group_id!r} from " + f"{recovery_group.status.value!r}" + ) pending_indices = [ sibling.generation_index for sibling in recovery_group.siblings @@ -2142,7 +2239,22 @@ async def _generate_for_finalization_attempt( ] group_id = recovery_group.group_id start_version = recovery_group.start_weight_version - rollout_ids = tuple(recovery_group.gate_rollout_ids) + logical_rollout_ids = tuple(recovery_group.logical_rollout_ids) + planned_rollout_ids = [ + nemo_gym_capture_key( + logical_rollout_ids[sibling.generation_index], + ( + sibling.current_attempt.attempt_index + if sibling.current_attempt.status + in { + RolloutAttemptStatus.RESERVED, + RolloutAttemptStatus.SEALED, + } + else len(sibling.attempts) + ), + ) + for sibling in recovery_group.siblings + ] attempt_input_sample = copy.deepcopy(input_sample) attempt_extra_env_info = attempt_input_sample.get("extra_env_info") if isinstance(attempt_extra_env_info, dict): @@ -2154,10 +2266,20 @@ async def _generate_for_finalization_attempt( weight_version=start_version, target_step=recovery_group.target_step, group_id=group_id, - rollout_ids=list(rollout_ids), + rollout_ids=planned_rollout_ids, ) pending_group_results: dict[int, SiblingSealResult] = {} + async def _allocate_recovery_attempts( + generation_indices: list[int], + ) -> dict[int, int]: + async with self._recovery_mutation() as cut: + return self._recovery_ledger.allocate_dispatch_attempts( + cut, + group_id, + generation_indices=generation_indices, + ) + async def _record_streamed_completion( generation_index: int, completion: Completion ) -> None: @@ -2178,12 +2300,14 @@ async def _record_streamed_completion( raise ValueError( "token-capture completion must contain its Gate rollout ID" ) - if not 0 <= generation_index < len(rollout_ids): + if not 0 <= generation_index < len(logical_rollout_ids): raise ValueError( f"streamed generation index {generation_index} is outside " f"prompt group {group_id!r}" ) - expected_gate_rollout_id = rollout_ids[generation_index] + expected_gate_rollout_id = self._recovery_ledger.get_group( + group_id + ).gate_rollout_id(generation_index) if gate_rollout_id != expected_gate_rollout_id: raise ValueError( "streamed rollout identity mismatch: " @@ -2248,17 +2372,12 @@ async def _record_streamed_completion( inflight_registry[group_id] = (current_task, start_version) try: if pending_indices: - async with self._recovery_mutation() as cut: - self._recovery_ledger.mark_group_dispatched( - cut, - group_id, - generation_indices=pending_indices, - ) await self.run_rollout( attempt_input_sample, - rollout_ids=list(rollout_ids), + logical_rollout_ids=list(logical_rollout_ids), generation_indices=pending_indices, on_completion=_record_streamed_completion, + attempt_allocator=_allocate_recovery_attempts, recovery_granularity=recovery_group.recovery_granularity, ) finally: diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index 0694652606d..0abf2965032 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -30,11 +30,13 @@ from enum import StrEnum from typing import TYPE_CHECKING, Any, Optional, Self, TypeAlias +from nemo_rl.experience.interfaces import nemo_gym_capture_key + if TYPE_CHECKING: from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneMutationCut from nemo_rl.data.interfaces import DatumSpec -ROLLOUT_RECOVERY_SCHEMA_VERSION = 2 +ROLLOUT_RECOVERY_SCHEMA_VERSION = 3 _SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS = {ROLLOUT_RECOVERY_SCHEMA_VERSION} ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" RolloutRecoveryState: TypeAlias = dict[str, Any] @@ -67,7 +69,7 @@ _SIBLING_STATE_FIELDS = frozenset({"generation_index", "attempts"}) _ATTEMPT_STATE_FIELDS = frozenset( { - "attempt_uuid", + "attempt_index", "status", "receipt", "reward", @@ -172,7 +174,7 @@ def _validate_prompt_identity( class RolloutAttemptRecord: """One physical attempt for a stable logical sibling.""" - attempt_uuid: uuid.UUID + attempt_index: int status: RolloutAttemptStatus receipt: Optional[dict[str, Any]] = None reward: Optional[float] = None @@ -181,8 +183,8 @@ class RolloutAttemptRecord: @property def attempt_id(self) -> str: - """Return the compact external representation of this attempt UUID.""" - return self.attempt_uuid.hex + """Return the external representation of this numeric attempt.""" + return str(self.attempt_index) @dataclass @@ -247,11 +249,11 @@ def logical_rollout_id(self, generation_index: int) -> str: return f"{self.group_id}_g{generation_index}" def gate_rollout_id(self, generation_index: int) -> str: - """Derive the physical Gate ID from group, sibling, and attempt UUID.""" + """Derive the physical Gate ID from logical identity and attempt index.""" sibling = self.siblings[generation_index] - return ( - f"{self.logical_rollout_id(generation_index)}" - f"_a{sibling.current_attempt.attempt_id}" + return nemo_gym_capture_key( + self.logical_rollout_id(generation_index), + sibling.current_attempt.attempt_index, ) @property @@ -284,9 +286,9 @@ class SiblingSealResult: mask_sample: bool -def _new_attempt() -> RolloutAttemptRecord: +def _new_attempt(attempt_index: int) -> RolloutAttemptRecord: return RolloutAttemptRecord( - attempt_uuid=uuid.uuid4(), + attempt_index=attempt_index, status=RolloutAttemptStatus.RESERVED, ) @@ -366,7 +368,7 @@ def reserve_group( siblings.append( RolloutSiblingRecord( generation_index=generation_index, - attempts=[_new_attempt()], + attempts=[_new_attempt(0)], ) ) record = PromptGroupRecoveryRecord( @@ -536,9 +538,55 @@ def prepare_incomplete_retry( f"{record.logical_rollout_id(sibling.generation_index)!r} " f"from status {attempt.status.value!r}" ) - sibling.attempts.append(_new_attempt()) + sibling.attempts.append(_new_attempt(len(sibling.attempts))) return self._copy_group(record) + def allocate_dispatch_attempts( + self, + cut: DataPlaneMutationCut, + group_id: str, + *, + generation_indices: list[int], + ) -> dict[int, int]: + """Allocate one numeric attempt for each sibling actually dispatched.""" + cut.require_live() + record = self._require_group(group_id) + if record.phase is not PromptGroupPhase.ADMITTED: + raise ValueError(f"cannot dispatch unadmitted recovery group {group_id!r}") + if record.status is not PromptGroupStatus.GENERATING: + raise ValueError( + f"cannot dispatch group {group_id!r} from {record.status.value!r}" + ) + if len(generation_indices) != len(set(generation_indices)): + raise ValueError("generation_indices must be unique") + + allocated: dict[int, int] = {} + for generation_index in generation_indices: + sibling = self._require_sibling(record, generation_index) + attempt = sibling.current_attempt + if attempt.status is RolloutAttemptStatus.SEALED: + raise ValueError( + "cannot redispatch sealed logical rollout " + f"{record.logical_rollout_id(generation_index)!r}" + ) + if attempt.status is RolloutAttemptStatus.DISPATCHED: + attempt.status = RolloutAttemptStatus.ABANDONED + if attempt.status in { + RolloutAttemptStatus.ABANDONED, + RolloutAttemptStatus.FAILED, + }: + attempt = _new_attempt(len(sibling.attempts)) + sibling.attempts.append(attempt) + if attempt.status is not RolloutAttemptStatus.RESERVED: + raise ValueError( + "cannot dispatch logical rollout " + f"{record.logical_rollout_id(generation_index)!r} " + f"from status {attempt.status.value!r}" + ) + attempt.status = RolloutAttemptStatus.DISPATCHED + allocated[generation_index] = attempt.attempt_index + return allocated + def mark_group_dispatched( self, cut: DataPlaneMutationCut, @@ -560,13 +608,11 @@ def mark_group_dispatched( if generation_indices is not None else list(range(record.expected_generations)) ) - attempts = [ - self._require_sibling(record, index).current_attempt for index in indices - ] - if any(attempt.status != RolloutAttemptStatus.RESERVED for attempt in attempts): - raise ValueError("only reserved rollout attempts may be dispatched") - for attempt in attempts: - attempt.status = RolloutAttemptStatus.DISPATCHED + self.allocate_dispatch_attempts( + cut, + group_id, + generation_indices=indices, + ) def mark_sibling_sealed( self, @@ -861,7 +907,7 @@ def state_dict(self) -> dict[str, Any]: "generation_index": sibling.generation_index, "attempts": [ { - "attempt_uuid": attempt.attempt_uuid.bytes, + "attempt_index": attempt.attempt_index, "status": attempt.status.value, "receipt": copy.deepcopy(attempt.receipt), "reward": attempt.reward, @@ -894,6 +940,11 @@ def from_state_dict(cls, state: dict[str, Any]) -> Self: context="rollout recovery state", ) schema_version = state.get("schema_version") + if schema_version == 2: + raise ValueError( + "Unsupported rollout-recovery schema version 2: UUID attempt " + "identities cannot be safely migrated to numeric attempt indices" + ) if ( isinstance(schema_version, bool) or not isinstance(schema_version, int) @@ -907,12 +958,8 @@ def from_state_dict(cls, state: dict[str, Any]) -> Self: raise ValueError("rollout-recovery state must contain a groups list") ledger = cls() - seen_attempt_uuids: set[uuid.UUID] = set() for raw_group in raw_groups: - record = cls._group_from_state( - raw_group, - seen_attempt_uuids=seen_attempt_uuids, - ) + record = cls._group_from_state(raw_group) if record.group_id in ledger._groups: raise ValueError(f"duplicate recovery group_id={record.group_id!r}") ledger._groups[record.group_id] = record @@ -946,8 +993,6 @@ def load_state_dict( def _group_from_state( cls, raw_group: Any, - *, - seen_attempt_uuids: set[uuid.UUID], ) -> PromptGroupRecoveryRecord: if not isinstance(raw_group, dict): raise ValueError("rollout-recovery group must be a mapping") @@ -1020,7 +1065,7 @@ def _group_from_state( if not isinstance(attempts_state, list) or not attempts_state: raise ValueError(f"logical rollout {logical_id!r} has no attempts") attempts: list[RolloutAttemptRecord] = [] - for attempt_state in attempts_state: + for expected_attempt_index, attempt_state in enumerate(attempts_state): if not isinstance(attempt_state, dict): raise ValueError("rollout-recovery attempt must be a mapping") _reject_unknown_fields( @@ -1028,17 +1073,16 @@ def _group_from_state( expected=_ATTEMPT_STATE_FIELDS, context="rollout-recovery attempt", ) - raw_attempt_uuid = attempt_state.get("attempt_uuid") + attempt_index = attempt_state.get("attempt_index") if ( - not isinstance(raw_attempt_uuid, bytes) - or len(raw_attempt_uuid) != 16 + isinstance(attempt_index, bool) + or not isinstance(attempt_index, int) + or attempt_index != expected_attempt_index ): - raise ValueError("attempt_uuid must contain exactly 16 bytes") - attempt_uuid = uuid.UUID(bytes=raw_attempt_uuid) - if attempt_uuid in seen_attempt_uuids: - raise ValueError("duplicate rollout attempt identity") - seen_attempt_uuids.add(attempt_uuid) - gate_id = f"{logical_id}_a{attempt_uuid.hex}" + raise ValueError( + "attempt indices must be contiguous non-negative integers" + ) + gate_id = nemo_gym_capture_key(logical_id, attempt_index) raw_attempt_status = attempt_state.get("status") if not isinstance(raw_attempt_status, str): raise ValueError( @@ -1088,7 +1132,7 @@ def _group_from_state( raise ValueError("only sealed attempts may retain receipt data") attempts.append( RolloutAttemptRecord( - attempt_uuid=attempt_uuid, + attempt_index=attempt_index, status=attempt_status, receipt=copy.deepcopy(receipt), reward=float(reward) if reward is not None else None, diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index fbe0bee4097..1804f1baae5 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -84,6 +84,9 @@ async def _run() -> None: rows = [ { "_rowidx": index, + "_ng_rollout_id": f"group_g{index}", + "_ng_attempt_index": 0, + "_ng_capture_id": f"group_g{index}", "task_source": "test_resources_server", "responses_create_params": {"input": []}, } @@ -137,6 +140,107 @@ def _postprocess_nemo_gym_to_nemo_rl_result( assert "task-source:test_resources_server" not in captured.err +def test_token_capture_forwards_logical_id_and_attempt_unchanged() -> None: + row = { + "_rowidx": 0, + "_ng_rollout_id": "group_g0", + "_ng_attempt_index": 2, + "_ng_capture_id": "group_g0-a2", + "agent_ref": {"name": "agent"}, + "responses_create_params": {"input": []}, + } + + class _RolloutCollectionHelper: + def run_examples(self, examples, head_server_config): + del head_server_config + forwarded = examples[0] + assert forwarded["_ng_rollout_id"] == "group_g0" + assert forwarded["_ng_attempt_index"] == 2 + assert forwarded["_ng_capture_id"] == "group_g0-a2" + + async def _completed_result(): + return forwarded, {"response": {"output": []}} + + return [_completed_result()] + + class _MockSelf: + cfg = {} + rch = _RolloutCollectionHelper() + head_server_config = object() + _token_capture_enabled = True + _tokenizer = object() + + def _require_spinup(self): + pass + + async def _postprocess_receipt_mode(self, forwarded, result): + del result + assert forwarded["_ng_rollout_id"] == "group_g0" + assert forwarded["_ng_attempt_index"] == 2 + assert forwarded["_ng_capture_id"] == "group_g0-a2" + return {"message_log": []} + + async def _run() -> tuple: + stream = NemoGym.__ray_metadata__.modified_class.run_rollouts( + _MockSelf(), [row], "test" + ) + return await stream.__anext__() + + rowidx, _, _, _ = asyncio.run(_run()) + assert rowidx == 0 + assert row["_ng_rollout_id"] == "group_g0" + assert row["_ng_attempt_index"] == 2 + + +@pytest.mark.parametrize( + ("attempt_index", "capture_key"), + [(0, "logical-rollout"), (1, "logical-rollout-a1")], +) +def test_receipt_lookup_uses_attempt_qualified_capture_key( + attempt_index: int, capture_key: str +) -> None: + class _MockSelf: + def __init__(self) -> None: + self.control_paths: list[str] = [] + + async def _control(self, method, path): + assert method == "GET" + self.control_paths.append(path) + return {} + + def _assemble_receipt( + self, + rollout_id, + manifest, + *, + terminal_response_id, + scored_response, + reward, + ): + del manifest, terminal_response_id, scored_response, reward + return {"rollout_id": rollout_id} + + mock_self = _MockSelf() + row = { + "_ng_rollout_id": "logical-rollout", + "_ng_attempt_index": attempt_index, + "_ng_capture_id": capture_key, + } + result = asyncio.run( + NemoGym.__ray_metadata__.modified_class._postprocess_receipt_mode( + mock_self, + row, + {"response": {}, "reward": 1.0}, + ) + ) + + assert row["_ng_rollout_id"] == "logical-rollout" + assert result["rollout_id"] == capture_key + assert mock_self.control_paths == [ + f"/training-token-capture/control/rollouts/{capture_key}/manifest" + ] + + def test_multimodal_content_types_cover_responses_media_aliases(): assert { "input_image", @@ -1588,6 +1692,9 @@ def test_nemo_gym_run_rollouts_normalizes_mixed_media_before_dispatch(tmp_path): async def _run(): nemo_gym_row = { "_rowidx": 7, + "_ng_rollout_id": "group_g7", + "_ng_attempt_index": 0, + "_ng_capture_id": "group_g7", "agent_ref": {"name": "legacy_test_agent"}, "responses_create_params": { "input": [ @@ -1685,6 +1792,9 @@ async def _run(): row = { "_rowidx": 3, + "_ng_rollout_id": "group_g3", + "_ng_attempt_index": 0, + "_ng_capture_id": "group_g3", "task_source": "test_resources_server", "agent_ref": {"name": "mock-megatron-agent"}, "responses_create_params": { diff --git a/tests/unit/experience/test_rollout_generation_failures.py b/tests/unit/experience/test_rollout_generation_failures.py index 4ae7944fabe..73aef35ee49 100644 --- a/tests/unit/experience/test_rollout_generation_failures.py +++ b/tests/unit/experience/test_rollout_generation_failures.py @@ -172,6 +172,7 @@ def _make_manager(buffer, impl, retry_policy=None) -> RolloutManager: manager = object.__new__(RolloutManager) manager._impl = impl manager._tokenizer = None + manager._use_nemo_gym = False manager._num_generations_per_prompt = 1 manager._tq_buffer = buffer manager._weight_version = 0 @@ -461,6 +462,7 @@ def __init__(self, fail_after_rows: int, failures_before_success: int) -> None: self._failures_before_success = failures_before_success self.attempts = 0 self.dispatched: list[list[int]] = [] + self.dispatch_identities: list[list[tuple[str, int, str]]] = [] def options(self, **kwargs): del kwargs @@ -469,6 +471,16 @@ def options(self, **kwargs): def remote(self, inputs, timer_prefix): del timer_prefix self.dispatched.append([row["_rowidx"] for row in inputs]) + self.dispatch_identities.append( + [ + ( + row["_ng_rollout_id"], + row["_ng_attempt_index"], + row["_ng_capture_id"], + ) + for row in inputs + ] + ) attempt = self.attempts self.attempts += 1 return self._stream(inputs, attempt) @@ -617,6 +629,17 @@ def test_only_the_missing_rows_are_re_dispatched(self): ) assert method.dispatched == [[0, 1, 2, 3], [2, 3]] + first_logical_ids = [identity[0] for identity in method.dispatch_identities[0]] + assert method.dispatch_identities == [ + [ + (first_logical_ids[index], 0, first_logical_ids[index]) + for index in range(4) + ], + [ + (first_logical_ids[index], 1, f"{first_logical_ids[index]}-a1") + for index in (2, 3) + ], + ] assert len(completions) == 4 def test_completed_rows_survive_across_attempts(self): diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index e6c191df657..d3e7475eea2 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -47,11 +47,14 @@ from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.failures import GenerationUnavailable from nemo_rl.experience.interfaces import ( + NEMO_GYM_ATTEMPT_INDEX_KEY, NEMO_GYM_GROUP_ATTEMPT_KEY, NEMO_GYM_GROUP_ID_KEY, + NEMO_GYM_ROLLOUT_ID_KEY, NEMO_GYM_ROLLOUT_INDEX_KEY, Completion, PromptGroupRecord, + nemo_gym_capture_key, ) from nemo_rl.experience.rollout_manager import ( AsyncNemoGymRolloutImpl, @@ -244,6 +247,7 @@ def _make_manager( mgr = object.__new__(RolloutManager) mgr._impl = impl mgr._tokenizer = None + mgr._use_nemo_gym = False mgr._num_generations_per_prompt = 1 mgr._rollout_recovery_config = RolloutRecoveryConfig() mgr._tq_buffer = buffer @@ -263,6 +267,62 @@ def _make_manager( class TestGenerateAndPushFlow: + def test_outer_retry_continues_after_inner_sibling_attempts(self): + class _AttemptingImpl: + def __init__(self) -> None: + self.dispatches: list[list[tuple[str, int]]] = [] + self.calls = 0 + + async def run_rollout( + self, + _sample, + *, + logical_rollout_ids, + generation_indices=None, + on_completion=None, + attempt_allocator=None, + recovery_granularity=RecoveryGranularity.SIBLING, + ): + del generation_indices, on_completion, recovery_granularity + self.calls += 1 + indices = [0, 1] + first = await attempt_allocator(indices) + self.dispatches.append( + [(logical_rollout_ids[index], first[index]) for index in indices] + ) + if self.calls == 1: + second = await attempt_allocator([1]) + self.dispatches.append([(logical_rollout_ids[1], second[1])]) + raise GenerationUnavailable("outer retry") + return "record" + + impl = _AttemptingImpl() + buf = _FakeBuffer() + mgr = _make_manager( + buf, + impl, + retry_policy=RolloutRetryPolicy( + max_infra_attempts=2, + max_data_attempts=1, + max_gym_row_attempts=2, + backoff_base_s=0.0, + ), + ) + mgr._num_generations_per_prompt = 2 + mgr._use_nemo_gym = True + sample = { + "idx": 3, + "extra_env_info": {"responses_create_params": {}}, + } + + assert _run(mgr.generate_and_push(sample)) is RolloutOutcome.COMMITTED + logical_ids = [dispatch[0] for dispatch in impl.dispatches[0]] + assert impl.dispatches == [ + [(logical_ids[0], 0), (logical_ids[1], 0)], + [(logical_ids[1], 1)], + [(logical_ids[0], 1), (logical_ids[1], 2)], + ] + def test_post_write_failure_does_not_regenerate_the_rollout(self): class _EnrichmentFailBuffer(_FakeBuffer): async def commit( @@ -1080,6 +1140,12 @@ def test_nemo_gym_build_inputs_stamps_logical_group_coordinates(): assert len({row[NEMO_GYM_GROUP_ID_KEY] for row in rows}) == 1 assert [row[NEMO_GYM_GROUP_ATTEMPT_KEY] for row in rows] == [0, 0, 0] assert [row[NEMO_GYM_ROLLOUT_INDEX_KEY] for row in rows] == [0, 1, 2] + assert [row[NEMO_GYM_ATTEMPT_INDEX_KEY] for row in rows] == [0, 0, 0] + assert [row[NEMO_GYM_ROLLOUT_ID_KEY] for row in rows] == [ + f"{rows[0][NEMO_GYM_GROUP_ID_KEY]}_g0", + f"{rows[0][NEMO_GYM_GROUP_ID_KEY]}_g1", + f"{rows[0][NEMO_GYM_GROUP_ID_KEY]}_g2", + ] assert [row["_rowidx"] for row in rows] == [0, 1, 2] @@ -1102,6 +1168,23 @@ def test_nemo_gym_build_inputs_preserves_explicit_group_identity(): ] assert [row[NEMO_GYM_GROUP_ATTEMPT_KEY] for row in rows] == [2, 2] assert [row[NEMO_GYM_ROLLOUT_INDEX_KEY] for row in rows] == [0, 1] + assert [row[NEMO_GYM_ROLLOUT_ID_KEY] for row in rows] == [ + "stable-group_g0", + "stable-group_g1", + ] + + +@pytest.mark.parametrize( + ("attempt_index", "expected_capture_key"), + [(0, "logical-rollout"), (1, "logical-rollout-a1")], +) +def test_rollout_identity_uses_attempt_qualified_capture_key( + attempt_index: int, + expected_capture_key: str, +) -> None: + assert ( + nemo_gym_capture_key("logical-rollout", attempt_index) == expected_capture_key + ) # --------------------------------------------------------------------------- @@ -1711,6 +1794,7 @@ def _make_capture_manager( ): mgr = object.__new__(RolloutManager) mgr._tokenizer = None + mgr._use_nemo_gym = True mgr._num_generations_per_prompt = num_generations mgr._rollout_recovery_config = recovery_config or RolloutRecoveryConfig() mgr._tq_buffer = buf @@ -1736,17 +1820,24 @@ async def run_rollout( self, _sample, *, - rollout_ids=None, + logical_rollout_ids=None, generation_indices=None, on_completion=None, + attempt_allocator=None, recovery_granularity=RecoveryGranularity.SIBLING, ): + indices = generation_indices or list(range(len(logical_rollout_ids))) + attempt_indices = await attempt_allocator(indices) + rollout_ids = list(logical_rollout_ids) + for index in indices: + rollout_ids[index] = nemo_gym_capture_key( + logical_rollout_ids[index], attempt_indices[index] + ) self.seen_rollout_ids = rollout_ids self.seen_generation_indices = list(generation_indices or []) self.seen_recovery_granularity = recovery_granularity if on_run is not None: await on_run(_sample) - indices = generation_indices or list(range(len(rollout_ids))) selected_ids = [rollout_ids[index] for index in indices] selected_configs = ( [instance_configs[index] for index in indices] @@ -1807,10 +1898,7 @@ def test_mints_ids_and_returns_metadata_request(self): canonical_ids = [f"{group_id}_g0", f"{group_id}_g1"] attempt_ids = buf.reserve_rollout_ids[0] assert attempt_ids is not None - assert all( - attempt_id.startswith(f"{canonical_id}_a") - for attempt_id, canonical_id in zip(attempt_ids, canonical_ids) - ) + assert attempt_ids == canonical_ids assert mgr._impl.seen_rollout_ids == attempt_ids assert request.group_id == group_id assert request.prompt_idx == 0 @@ -1848,12 +1936,19 @@ async def run_rollout( self, _sample, *, - rollout_ids=None, + logical_rollout_ids=None, generation_indices=None, on_completion=None, + attempt_allocator=None, recovery_granularity=RecoveryGranularity.SIBLING, ): del _sample, recovery_granularity + attempt_indices = await attempt_allocator(generation_indices) + rollout_ids = list(logical_rollout_ids) + for index in generation_indices: + rollout_ids[index] = nemo_gym_capture_key( + logical_rollout_ids[index], attempt_indices[index] + ) generation_index = generation_indices[0] rollout_id = rollout_ids[generation_index] receipt = { @@ -1961,12 +2056,19 @@ async def run_rollout( self, _sample, *, - rollout_ids=None, + logical_rollout_ids=None, generation_indices=None, on_completion=None, + attempt_allocator=None, recovery_granularity=RecoveryGranularity.SIBLING, ): indices = list(generation_indices) + attempt_indices = await attempt_allocator(indices) + rollout_ids = list(logical_rollout_ids) + for index in indices: + rollout_ids[index] = nemo_gym_capture_key( + logical_rollout_ids[index], attempt_indices[index] + ) self.generation_indices.append(indices) self.recovery_granularities.append(recovery_granularity) completions = [] diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 96b8a3fd691..0485e0cb9bc 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -223,7 +223,7 @@ def _sealed_attempt_state() -> dict[str, Any]: @pytest.mark.parametrize( ("case", "error_fragment"), [ - ("attempt_uuid", "attempt_uuid must contain exactly 16 bytes"), + ("attempt_index", "attempt indices must be contiguous"), ("status_type", "invalid rollout attempt status"), ("status_value", "invalid rollout attempt status"), ("staging_keys", "staging_keys must be a list of strings"), @@ -246,8 +246,8 @@ def test_restore_rejects_malformed_attempt_fields( state = _sealed_attempt_state() attempt = state["groups"][0]["siblings"][0]["attempts"][0] - if case == "attempt_uuid": - attempt["attempt_uuid"] = b"short" + if case == "attempt_index": + attempt["attempt_index"] = 2 elif case == "status_type": attempt["status"] = None elif case == "status_value": @@ -285,7 +285,7 @@ def test_restore_rejects_non_mapping_attempt() -> None: RolloutRecoveryLedger.from_state_dict(state) -def test_restore_rejects_duplicate_attempt_identity() -> None: +def test_restore_rejects_noncontiguous_attempt_identity() -> None: ledger = RolloutRecoveryLedger() _reserve( ledger, @@ -293,18 +293,16 @@ def test_restore_rejects_duplicate_attempt_identity() -> None: admission_id="batch-7", prompt_id="7", prompt_payload=_prompt(), - expected_generations=2, + expected_generations=1, target_step=7, start_weight_version=6, admitted=True, ) state = ledger.state_dict() - siblings = state["groups"][0]["siblings"] - siblings[1]["attempts"][0]["attempt_uuid"] = siblings[0]["attempts"][0][ - "attempt_uuid" - ] + attempt = state["groups"][0]["siblings"][0]["attempts"][0] + attempt["attempt_index"] = 1 - with pytest.raises(ValueError, match="duplicate rollout attempt identity"): + with pytest.raises(ValueError, match="attempt indices must be contiguous"): RolloutRecoveryLedger.from_state_dict(state) @@ -756,6 +754,67 @@ def test_restart_preserves_sealed_sibling_and_retries_only_interrupted_one() -> assert retry.siblings[1].current_attempt.status is RolloutAttemptStatus.RESERVED +def test_dispatch_attempts_are_numeric_persisted_and_sibling_local() -> None: + ledger = RolloutRecoveryLedger() + group = _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=6, + admitted=True, + ) + + first = _mutate( + lambda cut: ledger.allocate_dispatch_attempts( + cut, "g7", generation_indices=[0, 1] + ) + ) + sealed_id = group.gate_rollout_id(0) + _mutate( + lambda cut: ledger.mark_sibling_sealed( + cut, + "g7", + generation_index=0, + gate_rollout_id=sealed_id, + receipt={ + "rollout_id": sealed_id, + "manifest": [{"staging_key": f"{sealed_id}/call"}], + }, + reward=1.0, + mask_sample=False, + ) + ) + second = _mutate( + lambda cut: ledger.allocate_dispatch_attempts(cut, "g7", generation_indices=[1]) + ) + restored = RolloutRecoveryLedger.from_state_dict(ledger.state_dict()) + third = _mutate( + lambda cut: restored.allocate_dispatch_attempts( + cut, "g7", generation_indices=[1] + ) + ) + + assert first == {0: 0, 1: 0} + assert second == {1: 1} + assert third == {1: 2} + restored_group = restored.get_group("g7") + assert [ + attempt.attempt_index for attempt in restored_group.siblings[1].attempts + ] == [ + 0, + 1, + 2, + ] + assert ( + restored_group.siblings[0].current_attempt.status is RolloutAttemptStatus.SEALED + ) + assert restored_group.gate_rollout_id(1) == "g7_g1-a2" + + @pytest.mark.parametrize( "recovery_granularity", [RecoveryGranularity.SIBLING, RecoveryGranularity.PROMPT_GROUP], @@ -831,7 +890,7 @@ def test_missing_receipt_is_a_restart_safe_sealed_placeholder( assert rewards == [0.0, 1.0] assert mask_sample == [True, False] - state["schema_version"] = 3 + state["schema_version"] = ROLLOUT_RECOVERY_SCHEMA_VERSION + 1 with pytest.raises(ValueError, match="Unsupported rollout-recovery schema version"): RolloutRecoveryLedger.from_state_dict(state) @@ -1047,6 +1106,16 @@ def test_restore_rejects_unsupported_schema_version() -> None: _load(RolloutRecoveryLedger(), state) # type: ignore[arg-type] +def test_restore_rejects_uuid_attempt_schema_with_migration_guidance() -> None: + state = { + "schema_version": 2, + "groups": [], + } + + with pytest.raises(ValueError, match="cannot be safely migrated"): + _load(RolloutRecoveryLedger(), state) # type: ignore[arg-type] + + def test_restore_rejects_non_list_groups() -> None: state = { "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, diff --git a/tests/unit/experience/test_rollout_redispatch.py b/tests/unit/experience/test_rollout_redispatch.py index 81b215f3575..ea02d475c76 100644 --- a/tests/unit/experience/test_rollout_redispatch.py +++ b/tests/unit/experience/test_rollout_redispatch.py @@ -106,6 +106,7 @@ def _make_manager(buffer, impl, policy) -> RolloutManager: manager = object.__new__(RolloutManager) manager._impl = impl manager._tokenizer = None + manager._use_nemo_gym = False manager._num_generations_per_prompt = 1 manager._tq_buffer = buffer manager._weight_version = 0