diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index a40aee9ce5f..f04bf3225b4 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -184,6 +184,8 @@ rollout_checkpointing: # reaches this many tokens, rather than only when a checkpoint cut fires. # Bounds crash exposure independently of checkpoint cadence; 0 disables it. generation_chunk_flush_tokens: 0 + generation_prefix_batch_size: 256 # Checkpoint PUT rows per generation owner; 1 disables batching. + generation_prefix_batch_max_tokens: 4194304 # Aggregate tokens; one oversized row goes alone. prepare_timeout_s: 300 policy: diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index dc82b05cb9f..e45925933d7 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -196,6 +196,8 @@ rollout_checkpointing: # reaches this many tokens, rather than only when a checkpoint cut fires. # Bounds crash exposure independently of checkpoint cadence; 0 disables it. generation_chunk_flush_tokens: 0 + generation_prefix_batch_size: 256 # Checkpoint PUT rows per generation owner; 1 disables batching. + generation_prefix_batch_max_tokens: 4194304 # Aggregate tokens; one oversized row goes alone. prepare_timeout_s: 300 policy: diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index aa4dd67ff10..7cda7b1769d 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -759,6 +759,13 @@ class GymRolloutCheckpointConfig(BaseModel, extra="forbid"): # Stage an in-flight call's accumulated tokens once its unstaged segment # reaches this many tokens. 0 disables periodic chunk flushing. generation_chunk_flush_tokens: Annotated[int, Field(ge=0)] = 0 + # Checkpoint-triggered prefix PUTs per generation owner. 1 retains the + # single-row path. Periodic chunk flushing is unchanged. + generation_prefix_batch_size: Annotated[int, Field(gt=0)] = 256 + # Aggregate staged tokens per batch (~64 MiB of base tensors at default). + # Not a total-memory bound: metadata/copies add overhead; one oversized + # record is allowed alone, since a record cannot be split here. + generation_prefix_batch_max_tokens: Annotated[int, Field(gt=0)] = 4_194_304 prepare_timeout_s: Annotated[float, Field(gt=0)] = 300.0 @property diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 9d32c6261fe..aa3c91b3c96 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -2172,6 +2172,8 @@ def _build_trainer_then_megatron_generation() -> tuple[ rollout_checkpoint_cfg.gym.generation_prefix_cuts_enabled ), generation_cut_control_token=token_capture_cfg.control_auth_token, + generation_prefix_batch_size=rollout_checkpoint_cfg.gym.generation_prefix_batch_size, + generation_prefix_batch_max_tokens=rollout_checkpoint_cfg.gym.generation_prefix_batch_max_tokens, generation_chunk_flush_tokens=( rollout_checkpoint_cfg.gym.generation_chunk_flush_tokens ), diff --git a/nemo_rl/data_plane/tq_token_sink.py b/nemo_rl/data_plane/tq_token_sink.py index 8e0c626a369..1c266ff8b53 100644 --- a/nemo_rl/data_plane/tq_token_sink.py +++ b/nemo_rl/data_plane/tq_token_sink.py @@ -36,6 +36,7 @@ import json import logging +from collections.abc import Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -53,6 +54,7 @@ StageResult, ) +from nemo_rl.data_plane.codec import stack_or_nest from nemo_rl.data_plane.schema import ( ROUTE_ENCODING_ENVELOPE, ROUTE_ENCODING_LIST, @@ -240,124 +242,7 @@ def _stage_at_key(self, record: StagedCallRecord, key: str) -> StageResult: from nemo_gym.token_id_capture.staging.records import StageResult try: - field_dict = { - "token_ids_delta": torch.tensor( - [record.token_ids_delta], dtype=torch.int64 - ), - "token_mask_delta": torch.tensor( - [record.token_mask_delta], dtype=torch.float32 - ), - "generation_logprobs_delta": torch.tensor( - [record.generation_log_probs_delta], dtype=torch.float32 - ), - "schema_version": torch.tensor( - [record.schema_version], dtype=torch.int64 - ), - "digest_version": torch.tensor( - [record.digest_version], dtype=torch.int64 - ), - "extras_digest_version": torch.tensor( - [record.extras_digest_version], dtype=torch.int64 - ), - "rollout_id_utf8": _bytes_tensor(record.rollout_id.encode("utf-8")), - "model_call_id_utf8": _bytes_tensor( - record.model_call_id.encode("utf-8") - ), - "parent_call_id_utf8": _bytes_tensor( - (record.parent_call_id or "\0").encode("utf-8") - ), - "parent_call_id_present": torch.tensor( - [record.parent_call_id is not None], dtype=torch.bool - ), - "capture_mode": torch.tensor( - [_MODE_TO_CODE[record.mode]], dtype=torch.int64 - ), - "prev_len": torch.tensor([record.prev_len], dtype=torch.int64), - "delta_len": torch.tensor([record.delta_len], dtype=torch.int64), - "cum_len": torch.tensor([record.cum_len], dtype=torch.int64), - "weight_version": torch.tensor( - [record.weight_version], dtype=torch.int64 - ), - "digest_bytes": _bytes_tensor(bytes.fromhex(record.digest)), - "extras_digest_bytes": _bytes_tensor( - bytes.fromhex(record.extras_digest) - ), - } - chain_hash, chain_hash_present = _optional_digest_fields(record.chain_hash) - cumulative_hash, cumulative_hash_present = _optional_digest_fields( - record.cumulative_hash - ) - field_dict.update( - { - "chain_hash_bytes": chain_hash, - "chain_hash_present": chain_hash_present, - "cumulative_hash_bytes": cumulative_hash, - "cumulative_hash_present": cumulative_hash_present, - } - ) - extras_metadata = dict(record.extras) if record.extras is not None else None - routed = ( - extras_metadata.pop("routed_experts", None) - if extras_metadata is not None - else None - ) - field_dict[ROUTED_EXTRAS_METADATA_FIELD] = _bytes_tensor( - json.dumps( - extras_metadata, - sort_keys=True, - separators=(",", ":"), - ensure_ascii=False, - ).encode("utf-8") - ) - routed_len = 0 - routed_encoding = ROUTE_ENCODING_NONE - if routed is not None: - delta_len = len(record.token_ids_delta) - if isinstance(routed, str): - from nemo_rl.utils.routed_experts_codec import ( - decode_routed_experts, - ) - - 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}" - ) - experts = decode_routed_experts(routed, dtype) - routed_encoding = ROUTE_ENCODING_ENVELOPE - else: - experts = torch.tensor(routed, dtype=torch.int16) - routed_encoding = ROUTE_ENCODING_LIST - if experts.dim() != 3 or experts.shape[0] != delta_len: - raise ValueError( - "routed_experts must already be delta-aligned: " - f"got shape {tuple(experts.shape)} for delta_len={delta_len}" - ) - field_dict[ROUTED_EXPERTS_FIELD] = experts.unsqueeze(0) - routed_len = int(experts.shape[0]) - field_dict[ROUTED_EXPERTS_ENCODING_FIELD] = torch.tensor( - [routed_encoding], dtype=torch.int64 - ) - field_dict[ROUTED_LEN_FIELD] = torch.tensor([routed_len], dtype=torch.int64) - fields = TensorDict(field_dict, batch_size=[1]) - tags = [ - { - "rollout_id": record.rollout_id, - "model_call_id": record.model_call_id, - "parent_call_id": record.parent_call_id, - "prev_len": record.prev_len, - "delta_len": record.delta_len, - "cum_len": record.cum_len, - "weight_version": record.weight_version, - "digest": record.digest, - "schema_version": record.schema_version, - } - ] + fields, tags = self._encode_record(record) _call_dp( self._dp_client, "put_samples", @@ -381,6 +266,208 @@ def _stage_at_key(self, record: StagedCallRecord, key: str) -> StageResult: ) return StageResult(ok=True, staging_key=key) + @staticmethod + def _encode_record( + record: StagedCallRecord, + ) -> tuple[TensorDict, list[dict[str, Any]]]: + """Encode one row without publishing it.""" + field_dict = { + "token_ids_delta": torch.tensor( + [record.token_ids_delta], dtype=torch.int64 + ), + "token_mask_delta": torch.tensor( + [record.token_mask_delta], dtype=torch.float32 + ), + "generation_logprobs_delta": torch.tensor( + [record.generation_log_probs_delta], dtype=torch.float32 + ), + "schema_version": torch.tensor([record.schema_version], dtype=torch.int64), + "digest_version": torch.tensor([record.digest_version], dtype=torch.int64), + "extras_digest_version": torch.tensor( + [record.extras_digest_version], dtype=torch.int64 + ), + "rollout_id_utf8": _bytes_tensor(record.rollout_id.encode("utf-8")), + "model_call_id_utf8": _bytes_tensor(record.model_call_id.encode("utf-8")), + "parent_call_id_utf8": _bytes_tensor( + (record.parent_call_id or "\0").encode("utf-8") + ), + "parent_call_id_present": torch.tensor( + [record.parent_call_id is not None], dtype=torch.bool + ), + "capture_mode": torch.tensor( + [_MODE_TO_CODE[record.mode]], dtype=torch.int64 + ), + "prev_len": torch.tensor([record.prev_len], dtype=torch.int64), + "delta_len": torch.tensor([record.delta_len], dtype=torch.int64), + "cum_len": torch.tensor([record.cum_len], dtype=torch.int64), + "weight_version": torch.tensor([record.weight_version], dtype=torch.int64), + "digest_bytes": _bytes_tensor(bytes.fromhex(record.digest)), + "extras_digest_bytes": _bytes_tensor(bytes.fromhex(record.extras_digest)), + } + chain_hash, chain_hash_present = _optional_digest_fields(record.chain_hash) + cumulative_hash, cumulative_hash_present = _optional_digest_fields( + record.cumulative_hash + ) + field_dict.update( + { + "chain_hash_bytes": chain_hash, + "chain_hash_present": chain_hash_present, + "cumulative_hash_bytes": cumulative_hash, + "cumulative_hash_present": cumulative_hash_present, + } + ) + extras_metadata = dict(record.extras) if record.extras is not None else None + routed = ( + extras_metadata.pop("routed_experts", None) + if extras_metadata is not None + else None + ) + field_dict[ROUTED_EXTRAS_METADATA_FIELD] = _bytes_tensor( + json.dumps( + extras_metadata, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + ) + routed_len = 0 + routed_encoding = ROUTE_ENCODING_NONE + if routed is not None: + delta_len = len(record.token_ids_delta) + if isinstance(routed, str): + from nemo_rl.utils.routed_experts_codec import ( + decode_routed_experts, + ) + + 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}") + experts = decode_routed_experts(routed, dtype) + routed_encoding = ROUTE_ENCODING_ENVELOPE + else: + experts = torch.tensor(routed, dtype=torch.int16) + routed_encoding = ROUTE_ENCODING_LIST + if experts.dim() != 3 or experts.shape[0] != delta_len: + raise ValueError( + "routed_experts must already be delta-aligned: " + f"got shape {tuple(experts.shape)} for delta_len={delta_len}" + ) + field_dict[ROUTED_EXPERTS_FIELD] = experts.unsqueeze(0) + routed_len = int(experts.shape[0]) + field_dict[ROUTED_EXPERTS_ENCODING_FIELD] = torch.tensor( + [routed_encoding], dtype=torch.int64 + ) + field_dict[ROUTED_LEN_FIELD] = torch.tensor([routed_len], dtype=torch.int64) + fields = TensorDict(field_dict, batch_size=[1]) + tags = [ + { + "rollout_id": record.rollout_id, + "model_call_id": record.model_call_id, + "parent_call_id": record.parent_call_id, + "prev_len": record.prev_len, + "delta_len": record.delta_len, + "cum_len": record.cum_len, + "weight_version": record.weight_version, + "digest": record.digest, + "schema_version": record.schema_version, + } + ] + return fields, tags + + def stage_generation_prefix_batch( + self, + records: Sequence[StagedCallRecord], + *, + checkpoint_id: str, + chunk_sequences: Sequence[int], + ) -> list[StageResult]: + """Synchronously publish prefixes, preserving one result per input row. + + Success is returned only after TQ acknowledges the whole compatible + batch. A failed write reports failure for every member of that batch; + callers must not publish its lineage. Optional route columns and route + shapes are grouped separately so missing routes never become fabricated + payloads. Variable token and identity lengths remain jagged. + """ + # Gym is an optional dependency outside capture-enabled runs. + from nemo_gym.token_id_capture.staging.records import StageResult + + if len(records) != len(chunk_sequences): + raise ValueError("records and chunk_sequences must have equal lengths") + keys = [ + generation_cut_staging_key( + checkpoint_id, + record.rollout_id, + record.model_call_id, + chunk_sequence=sequence, + ) + for record, sequence in zip(records, chunk_sequences, strict=True) + ] + if len(set(keys)) != len(keys): + raise ValueError("generation-prefix batch contains duplicate staging keys") + if not records: + return [] + + results = [ + StageResult(ok=False, staging_key=key, error="not written") for key in keys + ] + groups: dict[tuple[Any, ...], list[tuple[int, TensorDict, dict[str, Any]]]] = {} + for index, record in enumerate(records): + try: + fields, tags = self._encode_record(record) + signature = tuple( + (name, fields[name].dtype, tuple(fields[name].shape[2:])) + for name in sorted(fields.keys()) + ) + groups.setdefault(signature, []).append((index, fields, tags[0])) + except Exception as error: # Encoding failures must fail closed. + results[index] = StageResult( + ok=False, + staging_key=keys[index], + error=f"{type(error).__name__}: {error}", + ) + + for group in groups.values(): + indices = [item[0] for item in group] + try: + fields = TensorDict( + { + name: stack_or_nest([item[1][name][0] for item in group]) + for name in group[0][1].keys() + }, + batch_size=[len(group)], + ) + _call_dp( + self._dp_client, + "put_samples", + sample_ids=[keys[index] for index in indices], + partition_id=self._staging_partition, + fields=fields, + tags=[item[2] for item in group], + ) + except Exception as error: # Batch failures must fail closed. + logging.getLogger(__name__).warning( + "TQ generation-prefix batch failed for %d rows: %s: %s", + len(group), + type(error).__name__, + error, + ) + for index in indices: + results[index] = StageResult( + ok=False, + staging_key=keys[index], + error=f"{type(error).__name__}: {error}", + ) + else: + for index in indices: + results[index] = StageResult(ok=True, staging_key=keys[index]) + return results + def clear(self, staging_keys: list[str]) -> None: """Drop staged rows (finalizer / eviction cleanup).""" if not staging_keys: diff --git a/nemo_rl/models/generation/generation_cut_capture.py b/nemo_rl/models/generation/generation_cut_capture.py index a53d81fea33..47b226ccbc8 100644 --- a/nemo_rl/models/generation/generation_cut_capture.py +++ b/nemo_rl/models/generation/generation_cut_capture.py @@ -18,7 +18,7 @@ import logging import threading import time -from collections.abc import Callable +from collections.abc import Callable, Generator from dataclasses import dataclass, field from typing import Any, Literal @@ -31,6 +31,20 @@ _COMPLETED_CAPTURE_RETENTION_S = 60.0 * 60.0 +@dataclass(frozen=True) +class GenerationPrefixBatchLimits: + """Backend-neutral limits for one worker's generation-prefix TQ writes.""" + + max_rows: int + max_tokens: int + + def __post_init__(self) -> None: + if self.max_rows < 1 or self.max_tokens < 1: + raise ValueError( + "generation-prefix batch row and token limits must be positive" + ) + + @dataclass class _RequestCaptureBuffer: """One append-only generation segment owned by an in-flight request.""" @@ -243,6 +257,27 @@ class GenerationCutCaptureMixin: and against ``_RequestCaptureState``; nothing here touches an engine. """ + _generation_prefix_batch_limits: GenerationPrefixBatchLimits | None = None + + def _configure_generation_prefix_batching( + self, *, max_rows: int, max_tokens: int + ) -> None: + """Install validated per-worker limits for backend-neutral cut batching.""" + self._generation_prefix_batch_limits = GenerationPrefixBatchLimits( + max_rows=max_rows, + max_tokens=max_tokens, + ) + + def _require_generation_prefix_batch_limits( + self, + ) -> GenerationPrefixBatchLimits: + limits = self._generation_prefix_batch_limits + if limits is None: + raise RuntimeError( + "generation-prefix batch limits require token capture setup" + ) + return limits + def _pop_request_capture(self, request: Any) -> _RequestCaptureState | None: with self._capture_registry_lock: state = self._capture_calls.pop(id(request), None) @@ -702,7 +737,34 @@ def _checkpoint_active_generation_cut( state: _RequestCaptureState, checkpoint_id: str, ) -> Any | None: - """Cut one live call, or return ``None`` if abort won the race.""" + """Single-row compatibility path using the freeze/seal transaction.""" + transaction = self._generation_cut_transaction(prefix, state, checkpoint_id) + try: + try: + record, sequence = next(transaction) + except StopIteration as finished: + return finished.value + result = self._capture_sink.stage_generation_prefix( + record, checkpoint_id=checkpoint_id, chunk_sequence=sequence + ) + acknowledgement = transaction.send(result) + try: + transaction.send(None) + except StopIteration: + return acknowledgement + raise RuntimeError("generation-cut transaction did not seal") + finally: + transaction.close() + + def _generation_cut_transaction( + self, prefix: Any, state: _RequestCaptureState, checkpoint_id: str + ) -> Generator[Any, Any, Any]: + """Own lifecycle until a staged row and acknowledgement are validated. + + Yield ``(record, sequence)``, receive a ``StageResult``, then yield a + validated acknowledgement. The caller resumes once more to seal, or + closes to roll back. Calls needing no new row return immediately. + """ from nemo_gym._checkpoint.model_control_contracts import ( GenerationCutPrefixAck, ) @@ -772,6 +834,11 @@ def _checkpoint_active_generation_cut( disposition="durable_failure", ) if effective_output_limit is None: + with state.lock: + state.rollback_frozen_buffer(checkpoint_id) + state.refresh_periodic_flush_due( + self._generation_chunk_flush_tokens + ) raise RuntimeError( f"cannot cut model call {prefix.model_call_id!r}: " "vLLM did not resolve its effective output limit" @@ -795,11 +862,7 @@ def _checkpoint_active_generation_cut( generated_logprobs=chunk_logprobs, ) ) - result = sink.stage_generation_prefix( - chunk_record, - checkpoint_id=checkpoint_id, - chunk_sequence=chunk_sequence, - ) + result = yield chunk_record, chunk_sequence staged_key = result.staging_key if not result.ok: raise RuntimeError( @@ -827,6 +890,10 @@ def _checkpoint_active_generation_cut( terminal_finish_reason=terminal_finish_reason, terminal_stop_reason=terminal_stop_reason, ) + if chunk_token_ids: + # Every row in a batch must validate before any row seals. + yield acknowledgement + with state.lock: # Validate the complete acknowledgement before adopting the # staged chunk. A failed acknowledgement must leave the # frozen buffer available for rollback. @@ -836,7 +903,9 @@ def _checkpoint_active_generation_cut( state.refresh_periodic_flush_due( self._generation_chunk_flush_tokens ) - except Exception: + except BaseException: + # Generator.close() injects GeneratorExit when a batch fails. + # Roll back the detached state without hiding the root failure. with state.lock: frozen = state.frozen_buffer if frozen is not None and frozen.checkpoint_id == checkpoint_id: @@ -855,6 +924,131 @@ def _checkpoint_active_generation_cut( raise return acknowledgement + def _checkpoint_generation_cut_batched(self, inventory: Any) -> list[Any]: + """Batch independently per owner, retaining lifecycle locks until seal.""" + from nemo_rl.data_plane.tq_token_sink import generation_cut_staging_key + + batch_limits = self._require_generation_prefix_batch_limits() + call_ids = [prefix.model_call_id for prefix in inventory.active_prefixes] + if len(set(call_ids)) != len(call_ids): + raise ValueError( + "generation-cut inventory contains duplicate model-call identities" + ) + + sink = self._capture_sink + acknowledgements: list[Any] = [] + pending: list[tuple[Generator[Any, Any, Any], Any, int]] = [] + pending_tokens = 0 + + def flush() -> None: + nonlocal pending_tokens + if not pending: + return + keys = [ + generation_cut_staging_key( + inventory.checkpoint_id, + record.rollout_id, + record.model_call_id, + chunk_sequence=sequence, + ) + for _, record, sequence in pending + ] + sealed_keys: set[str] = set() + try: + results = sink.stage_generation_prefix_batch( + [record for _, record, _ in pending], + checkpoint_id=inventory.checkpoint_id, + chunk_sequences=[sequence for _, _, sequence in pending], + ) + if len(results) != len(pending): + raise RuntimeError( + "generation-prefix batch returned an incorrect result count" + ) + for result, key in zip(results, keys, strict=True): + if not result.ok or result.staging_key != key: + raise RuntimeError( + "generation-prefix batch staging failed: " + f"{result.error}; key={key}" + ) + batch_acks = [ + transaction.send(result) + for (transaction, _, _), result in zip( + pending, results, strict=True + ) + ] + for (transaction, _, _), key in zip(pending, keys, strict=True): + try: + transaction.send(None) + except StopIteration: + sealed_keys.add(key) + else: + raise RuntimeError("generation-cut transaction did not seal") + acknowledgements.extend(batch_acks) + except BaseException: + # A transport error may happen after TQ accepted some rows. + # These are new chunk keys, never earlier sealed prefix keys. + rejected = [key for key in keys if key not in sealed_keys] + if rejected: + try: + sink.clear(rejected) + except Exception: + LOGGER.exception( + "Failed to clear rejected generation-prefix batch" + ) + raise + finally: + for transaction, _, _ in pending: + transaction.close() + pending.clear() + pending_tokens = 0 + + try: + for prefix in inventory.active_prefixes: + with self._capture_registry_lock: + state = self._capture_calls_by_model_call_id.get( + prefix.model_call_id + ) + if state is None: + acknowledgement = self._completed_generation_cut_ack( + prefix, inventory.checkpoint_id + ) + acknowledgements.append(acknowledgement) + continue + transaction = self._generation_cut_transaction( + prefix, state, inventory.checkpoint_id + ) + try: + try: + record, sequence = next(transaction) + except StopIteration as finished: + acknowledgement = finished.value + if acknowledgement is None: + acknowledgement = self._completed_generation_cut_ack( + prefix, inventory.checkpoint_id + ) + acknowledgements.append(acknowledgement) + continue + tokens = len(record.token_ids_delta) + if pending and pending_tokens + tokens > batch_limits.max_tokens: + flush() + pending.append((transaction, record, sequence)) + except BaseException: + transaction.close() + raise + pending_tokens += tokens + if ( + len(pending) >= batch_limits.max_rows + or pending_tokens >= batch_limits.max_tokens + ): + # An oversized single row is written alone, never split. + flush() + flush() + finally: + for transaction, _, _ in pending: + transaction.close() + by_ticket = {ack.ticket_id: ack for ack in acknowledgements} + return [by_ticket[prefix.ticket_id] for prefix in inventory.active_prefixes] + def _checkpoint_generation_cut(self, inventory: Any) -> Any: """Stage a stable prefix for every call named by Gym's frozen inventory.""" from nemo_gym._checkpoint.model_control_contracts import ( @@ -871,7 +1065,12 @@ def _checkpoint_generation_cut(self, inventory: Any) -> Any: if cached_receipt is not None: return cached_receipt acknowledgements = [] - for prefix in inventory.active_prefixes: + single_prefixes = inventory.active_prefixes + batch_limits = self._require_generation_prefix_batch_limits() + if batch_limits.max_rows > 1: + acknowledgements = self._checkpoint_generation_cut_batched(inventory) + single_prefixes = () + for prefix in single_prefixes: with self._capture_registry_lock: state = self._capture_calls_by_model_call_id.get(prefix.model_call_id) if state is None: diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 4c24003c867..507bdde5ef7 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -630,6 +630,8 @@ def setup_token_capture( generation_prefix_cuts_enabled: bool = False, generation_cut_control_token: str | None = None, generation_chunk_flush_tokens: int = 0, + generation_prefix_batch_size: int, + generation_prefix_batch_max_tokens: int, ) -> None: """Install ledger-authoritative token capture in every DP-leader worker. @@ -648,6 +650,8 @@ def setup_token_capture( generation_prefix_cuts_enabled=generation_prefix_cuts_enabled, generation_cut_control_token=generation_cut_control_token, generation_chunk_flush_tokens=generation_chunk_flush_tokens, + generation_prefix_batch_size=generation_prefix_batch_size, + generation_prefix_batch_max_tokens=generation_prefix_batch_max_tokens, run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], ) ray.get(futures) diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 24645ad1ffe..b070c20f6a0 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -767,6 +767,8 @@ async def setup_token_capture( generation_prefix_cuts_enabled: bool = False, generation_cut_control_token: str | None = None, generation_chunk_flush_tokens: int = 0, + generation_prefix_batch_size: int, + generation_prefix_batch_max_tokens: int, ) -> bool: """Host ledger-authoritative token capture in this worker. @@ -793,6 +795,10 @@ async def setup_token_capture( ) if generation_chunk_flush_tokens < 0: raise ValueError("generation_chunk_flush_tokens must not be negative") + self._configure_generation_prefix_batching( + max_rows=generation_prefix_batch_size, + max_tokens=generation_prefix_batch_max_tokens, + ) if generation_chunk_flush_tokens and not generation_prefix_cuts_enabled: raise ValueError( "generation_chunk_flush_tokens requires generation-prefix cuts; " diff --git a/pyrefly.toml b/pyrefly.toml index 79334aba1ce..7f55397ef02 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -328,6 +328,8 @@ project-includes = [ "nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py", "nemo_rl/weight_sync/xferdtensor.py", "nemo_rl/weight_sync/xferdtensor_python.py", + "tools/benchmark_tq_prefix_writes.py", + "tools/benchmark_prefix_checkpoint_proxy.py", "tools/external_gym_vllm/serve_vllm_on_ray.py", "tools/external_gym_vllm/vllm_pool_lb.py", "tools/model_diagnostics/1.max_model_len_respected.py", diff --git a/tests/unit/data_plane/test_prefix_checkpoint_proxy.py b/tests/unit/data_plane/test_prefix_checkpoint_proxy.py new file mode 100644 index 00000000000..51fe1fd6a3c --- /dev/null +++ b/tests/unit/data_plane/test_prefix_checkpoint_proxy.py @@ -0,0 +1,346 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dependency-light tests; run directly with unittest to avoid GPU conftest.""" + +import asyncio +import importlib.util +import tempfile +import threading +import time +import unittest +from collections import Counter +from pathlib import Path +from types import SimpleNamespace + +from tools.benchmark_prefix_checkpoint_proxy import ( + BoundedStager, + CachedRecordBuilder, + PerWorkerStager, + RealTQCutBackend, + TimedPutClient, + latency_summary, + load_checkpoint_driver, + worker_for_row, +) + + +class ProxyTests(unittest.TestCase): + def test_per_worker_routes_id_only_batches_without_shared_gate(self) -> None: + args = SimpleNamespace(rows=8, generation_workers=2) + measured = TimedPutClient(None) + stager = PerWorkerStager(measured, args) + received: list[tuple[int, list[int]]] = [] + + def actor(owner: int) -> SimpleNamespace: + async def remote(indices: list[int], checkpoint_id: str) -> dict: + self.assertEqual(checkpoint_id, "checkpoint") + self.assertTrue(all(isinstance(index, int) for index in indices)) + received.append((owner, indices)) + await asyncio.sleep(0.01) + return { + "results": [dict(ok=True, staging_key=f"key-{i}") for i in indices], + "stage_seconds": 0.005, + "put_latencies": [0.004], + "batch_rows": [len(indices)], + } + + return SimpleNamespace(stage=SimpleNamespace(remote=remote)) + + stager.actors = [actor(0), actor(1)] + + def record(index: int) -> SimpleNamespace: + return SimpleNamespace(rollout_id=f"rollout-{index:09d}") + + async def run() -> None: + results = await asyncio.gather( + stager.stage([record(0), record(1)], "checkpoint"), + stager.stage([record(4)], "checkpoint"), + ) + self.assertEqual(results[0][0].staging_key, "key-0") + with self.assertRaisesRegex(ValueError, "crosses"): + await stager.stage([record(0), record(4)], "checkpoint") + + asyncio.run(run()) + self.assertEqual(received, [(0, [0, 1]), (1, [4])]) + self.assertEqual(stager.peak_requests, 2) + self.assertEqual(stager.active_requests, 0) + self.assertEqual(measured.batch_rows, [2, 1]) + self.assertEqual(len(measured.latencies), 2) + self.assertEqual(stager.waits, []) + + def test_per_worker_remote_failure_propagates(self) -> None: + async def remote(indices: list[int], checkpoint_id: str) -> None: + raise RuntimeError("injected writer failure") + + stager = PerWorkerStager( + TimedPutClient(None), SimpleNamespace(rows=1, generation_workers=1) + ) + stager.actors = [SimpleNamespace(stage=SimpleNamespace(remote=remote))] + + async def run() -> None: + with self.assertRaisesRegex(RuntimeError, "injected writer failure"): + await stager.stage( + [SimpleNamespace(rollout_id="rollout-000000000")], "checkpoint" + ) + + asyncio.run(run()) + self.assertEqual(stager.active_requests, 0) + self.assertEqual(stager.stages, []) + + @unittest.skipUnless( + importlib.util.find_spec("nemo_gym"), "requires Gym dependencies" + ) + def test_cached_records_match_production_encoder(self) -> None: + # Optional Gym dependency, not needed for scheduler-only tests. + from nemo_gym.token_id_capture.staging.digest import ( + compute_chain_hash, + compute_extras_digest, + compute_staging_digest, + hash_token_ids, + ) + from nemo_gym.token_id_capture.staging.records import StagedCallRecord + + for length in (1, 17, 128, 16384): + for child in (False, True): + tokens = [i % 32000 for i in range(length)] + masks = [float(i % 2) for i in range(length)] + logprobs = [-0.125 if mask else 0.0 for mask in masks] + extras = {"test": [1, "π"]} if child else None + values = dict( + schema_version=2, + digest_version=2, + extras_digest_version=1, + rollout_id="original", + model_call_id="call-0", + parent_call_id="parent" if child else None, + mode="token_in" if child else "text", + prev_len=3 if child else 0, + delta_len=length, + cum_len=length + (3 if child else 0), + weight_version=7, + token_ids_delta=tokens, + token_mask_delta=masks, + generation_log_probs_delta=logprobs, + extras_digest=compute_extras_digest(extras), + chain_hash=compute_chain_hash("a" * 64 if child else None, tokens), + cumulative_hash=hash_token_ids( + ([1, 2, 3] if child else []) + tokens + ), + ) + template = StagedCallRecord( + **values, extras=extras, digest=compute_staging_digest(**values) + ) + builder = CachedRecordBuilder(template) + for rollout_id, call_id in ( + ("a", "b"), + ("rollout-000131071", "call-000131071"), + ("π/日本語", "call/é"), + ): + changed = values | { + "rollout_id": rollout_id, + "model_call_id": call_id, + } + expected = StagedCallRecord( + **changed, + extras=extras, + digest=compute_staging_digest(**changed), + ) + actual = builder.build(rollout_id=rollout_id, model_call_id=call_id) + self.assertEqual(actual.model_dump(), expected.model_dump()) + StagedCallRecord.model_validate(actual.model_dump()) + another = builder.build(rollout_id="next", model_call_id="next") + self.assertIs(actual.token_ids_delta, another.token_ids_delta) + self.assertIsNot(template.token_ids_delta, actual.token_ids_delta) + with self.assertRaises(ValueError): + builder.build(rollout_id="", model_call_id="call") + with self.assertRaises(ValueError): + CachedRecordBuilder( + template.model_copy(update={"digest": "0" * 64}) + ) + + @unittest.skipUnless( + importlib.util.find_spec("nemo_gym"), "requires Gym dependencies" + ) + def test_real_coordinator_commit_and_coordinate_restore(self) -> None: + """Real Gym checkpoint lifecycle with a fake acknowledged storage sink.""" + driver = load_checkpoint_driver() + args = SimpleNamespace( + rows=16, generation_workers=4, batch_size=3, prefix_tokens=8 + ) + records = [ + SimpleNamespace( + rollout_id=f"rollout-{i:09d}", + model_call_id=f"call-{i:09d}", + digest="a" * 64, + ) + for i in range(args.rows) + ] + + def key(record: SimpleNamespace, checkpoint_id: str) -> str: + return f"__generation_cut__/{checkpoint_id}/{record.rollout_id}/{record.model_call_id}/0" + + class Sink: + def stage_generation_prefix_batch( + self, + batch: list[SimpleNamespace], + *, + checkpoint_id: str, + chunk_sequences: list[int], + ) -> list[SimpleNamespace]: + return [ + SimpleNamespace(ok=True, staging_key=key(record, checkpoint_id)) + for record in batch + ] + + async def run(root: Path) -> None: + stager = BoundedStager(Sink(), concurrency=2) + try: + result = await driver.run_checkpoint_benchmark( + root, + workers=2, + cuts=args.rows, + hot_worker_fraction=0.5, + prefix_tokens=8, + staging_key_bytes=0, + timeout_s=30, + backend_factory=lambda owner: RealTQCutBackend( + owner, records, stager, args + ), + staging_keys_for_ticket=lambda ticket, checkpoint_id: ( + key(ticket, checkpoint_id), + ), + ) + self.assertEqual(result["generation_cuts_restored"], args.rows) + self.assertTrue(result["restored_prefix_coordinates_match"]) + timings = result["restore_timings"] + for phase in ( + "metadata_validation_seconds", + "archive_validation_seconds", + "tar_inventory_seconds", + "member_read_validation_seconds", + "lineage_parse_seconds", + "file_materialization_seconds", + "namespace_inventory_seconds", + "directory_publish_seconds", + "receipt_reconstruction_seconds", + "result_serialization_seconds", + ): + self.assertGreater(timings[phase], 0, phase) + self.assertAlmostEqual( + sum( + value + for name, value in timings.items() + if name != "total_seconds" + ), + timings["total_seconds"], + ) + self.assertLessEqual( + timings["total_seconds"], result["restore_seconds"] + ) + self.assertEqual(len(stager.stages), 8) + finally: + stager.close() + + with tempfile.TemporaryDirectory(prefix="proxy-unit-") as root: + asyncio.run(run(Path(root))) + + def test_ownership_and_expected_batch_counts(self) -> None: + for rows, expected in ( + (8192, 128), + (16384, 128), + (65536, 256), + (131072, 512), + (262144, 1024), + ): + counts = Counter( + worker_for_row(index, rows=rows, workers=128) for index in range(rows) + ) + self.assertEqual(len(counts), 128) + self.assertEqual( + sum((count + 255) // 256 for count in counts.values()), expected + ) + self.assertEqual( + [worker_for_row(i, rows=3, workers=8) for i in range(3)], [0, 1, 2] + ) + self.assertEqual( + [worker_for_row(i, rows=5, workers=2) for i in range(5)], [0, 0, 0, 1, 1] + ) + + def test_percentiles(self) -> None: + summary = latency_summary(list(range(1, 101))) + self.assertEqual(summary["p50_s"], 50) + self.assertEqual(summary["p95_s"], 95) + self.assertEqual(summary["p99_s"], 99) + self.assertEqual(latency_summary([]), {"count": 0}) + + def test_put_failure_is_measured_and_propagated(self) -> None: + class BrokenClient: + def put_samples(self, **kwargs: object) -> None: + raise RuntimeError("injected") + + measured = TimedPutClient(BrokenClient()) + with self.assertRaisesRegex(RuntimeError, "injected"): + measured.put_samples(sample_ids=["a", "b"]) + self.assertEqual(measured.failures, 1) + self.assertEqual(measured.batch_rows, [2]) + self.assertEqual(measured.active, 0) + + def test_concurrency_is_bounded_and_all_batches_finish(self) -> None: + class Sink: + def __init__(self) -> None: + self.lock = threading.Lock() + self.active = 0 + self.peak = 0 + self.rows = 0 + + def stage_generation_prefix_batch( + self, records: list[int], **kwargs: object + ) -> list[SimpleNamespace]: + with self.lock: + self.active += 1 + self.peak = max(self.peak, self.active) + time.sleep(0.01) + with self.lock: + self.active -= 1 + self.rows += len(records) + return [SimpleNamespace(ok=True) for _ in records] + + sink = Sink() + + async def run() -> None: + stager = BoundedStager(sink, concurrency=2) + try: + await asyncio.gather( + *(stager.stage([index], "checkpoint") for index in range(20)) + ) + self.assertEqual(len(stager.stages), 20) + self.assertEqual(len(stager.waits), 20) + self.assertGreater(max(stager.waits), 0.01) + finally: + stager.close() + + asyncio.run(run()) + self.assertEqual(sink.peak, 2) + self.assertEqual(sink.active, 0) + self.assertEqual(sink.rows, 20) + + def test_failed_batch_cannot_become_evidence(self) -> None: + class Sink: + def stage_generation_prefix_batch( + self, records: list[int], **kwargs: object + ) -> list[SimpleNamespace]: + return [SimpleNamespace(ok=False)] + + async def run() -> None: + stager = BoundedStager(Sink(), concurrency=1) + try: + with self.assertRaisesRegex(RuntimeError, "Prefix staging failed"): + await stager.stage([0], "checkpoint") + finally: + stager.close() + + asyncio.run(run()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/data_plane/test_tq_token_sink.py b/tests/unit/data_plane/test_tq_token_sink.py index 60967821e8d..aa9ed87c28d 100644 --- a/tests/unit/data_plane/test_tq_token_sink.py +++ b/tests/unit/data_plane/test_tq_token_sink.py @@ -28,6 +28,10 @@ nemo_gym = pytest.importorskip("nemo_gym.token_id_capture.staging") +from nemo_gym.token_id_capture.staging.digest import ( # noqa: E402 + compute_extras_digest, + compute_staging_digest, +) from nemo_gym.token_id_capture.staging.protocols import ( # noqa: E402 StagingSink as TokenSinkProtocol, ) @@ -45,6 +49,12 @@ build_fixture_artifacts, fixture_names, ) +from tools.benchmark_tq_prefix_writes import ( # noqa: E402 + BenchmarkConfig, + _data_plane_config, + _write_ranges, + run_prefix_write_benchmark, +) STAGING_PARTITION = "rollout_staging_test" @@ -258,3 +268,222 @@ def test_fetch_prefix_token_ids_rejects_duplicates(tq_client, staging_partition) source = TQTokenSource(tq_client, staging_partition=staging_partition) with pytest.raises(KeyError, match="duplicates"): source.fetch_prefix_token_ids(["r/c", "r/c"]) + + +def test_tq_prefix_write_benchmark_exercises_live_unbatched_path(tq_client): + partition_id = "rollout_staging_prefix_write_benchmark_test" + result = run_prefix_write_benchmark( + tq_client, + BenchmarkConfig( + rows=8, + prefix_tokens=4, + writers_per_client=2, + num_storage_units=1, + verify_rows=3, + verify_batch_size=2, + partition_id=partition_id, + checkpoint_id="benchmark-test", + ), + ) + + assert result["put_calls"] == 8 + assert result["stored_keys"] == 8 + assert result["verified_rows"] == 3 + assert result["writers_used"] == 2 + assert result["rows_per_second"] > 0 + assert result["cleanup_seconds"] is not None + assert tq_client.list_sample_ids(partition_id) == [] + + +@pytest.mark.parametrize("batch_size,expected_calls", [(1, 8), (3, 4)]) +def test_tq_prefix_write_benchmark_uses_process_isolated_clients( + tq_client, batch_size, expected_calls +): + partition_id = "rollout_staging_prefix_write_multi_client_test" + config = BenchmarkConfig( + rows=8, + prefix_tokens=4, + clients=2, + writers_per_client=1, + batch_size=batch_size, + num_storage_units=1, + verify_rows=3, + verify_batch_size=2, + partition_id=partition_id, + checkpoint_id="multi-client-benchmark-test", + ) + result = run_prefix_write_benchmark( + tq_client, + config, + dp_config=_data_plane_config(config), + ) + + assert result["put_calls"] == expected_calls + assert result["stored_keys"] == 8 + assert result["verified_rows"] == 3 + assert result["clients_used"] == 2 + assert result["writers_used"] == 2 + assert result["client_setup_seconds"] > 0 + assert result["rows_per_second"] > 0 + assert result["cleanup_seconds"] is not None + assert tq_client.list_sample_ids(partition_id) == [] + + +def test_generation_prefix_batch_round_trips_ragged_rows(tq_client, staging_partition): + class RecordingClient: + def __init__(self): + self.calls = [] + + def put_samples(self, **kwargs): + self.calls.append(kwargs) + return tq_client.put_samples(**kwargs) + + client = RecordingClient() + sink = TQTokenSink(client, staging_partition=staging_partition) + records, _, _ = build_fixture_artifacts("worked_example") + more_records, _, _ = build_fixture_artifacts( + "single_call", rollout_id="a-different-length-rollout-id" + ) + records.extend(more_records) + results = sink.stage_generation_prefix_batch( + records, checkpoint_id="batch-checkpoint", chunk_sequences=[0, 2, 3] + ) + assert all(result.ok for result in results) + assert len(client.calls) == 1 + assert client.calls[0]["sample_ids"] == [result.staging_key for result in results] + assert [tag["digest"] for tag in client.calls[0]["tags"]] == [ + record.digest for record in records + ] + source = TQTokenSource(tq_client, staging_partition=staging_partition) + restored = source.fetch([result.staging_key for result in results]) + assert [row.model_dump() for row in restored] == [ + record.model_dump(exclude={"extras"}) for record in records + ] + + +def test_generation_prefix_batch_rejects_bad_inventory_before_writing(): + class UnexpectedClient: + def put_samples(self, **kwargs): + pytest.fail("invalid inventory must not reach TQ") + + sink = TQTokenSink(UnexpectedClient(), staging_partition="test") + records, _, _ = build_fixture_artifacts("single_call") + assert ( + sink.stage_generation_prefix_batch( + [], checkpoint_id="checkpoint", chunk_sequences=[] + ) + == [] + ) + with pytest.raises(ValueError, match="equal lengths"): + sink.stage_generation_prefix_batch( + records, checkpoint_id="checkpoint", chunk_sequences=[] + ) + with pytest.raises(ValueError, match="duplicate"): + sink.stage_generation_prefix_batch( + records * 2, checkpoint_id="checkpoint", chunk_sequences=[0, 0] + ) + + +def test_generation_prefix_batch_preserves_optional_routes( + tq_client, staging_partition +): + class RecordingClient: + def __init__(self): + self.calls = [] + + def put_samples(self, **kwargs): + self.calls.append(kwargs) + return tq_client.put_samples(**kwargs) + + records, _, _ = build_fixture_artifacts("worked_example") + values = records[0].model_dump(exclude={"digest"}) + values["extras"] = {"routed_experts": [[[1, 2]]] * records[0].delta_len} + values["extras_digest"] = compute_extras_digest(values["extras"]) + values["digest"] = compute_staging_digest( + **{key: value for key, value in values.items() if key != "extras"} + ) + records[0] = type(records[0]).model_validate(values) + client = RecordingClient() + sink = TQTokenSink(client, staging_partition=staging_partition) + results = sink.stage_generation_prefix_batch( + records, checkpoint_id="routes", chunk_sequences=[0, 0] + ) + assert all(result.ok for result in results) + assert len(client.calls) == 2 + assert "routed_experts" in client.calls[0]["fields"] + assert "routed_experts" not in client.calls[1]["fields"] + source = TQTokenSource(tq_client, staging_partition=staging_partition) + routed = source.fetch_for_finalization( + [results[0].staging_key], include_route_fragments=True + )[0] + assert routed.fragment is not None + assert routed.routed_len == records[0].delta_len + assert routed.snapshot.model_dump() == records[0].model_dump(exclude={"extras"}) + plain = source.fetch([results[1].staging_key])[0] + assert plain.model_dump() == records[1].model_dump(exclude={"extras"}) + + +def test_generation_prefix_batch_failure_does_not_advertise_success(): + class FailingClient: + def __init__(self): + self.calls = 0 + + def put_samples(self, **kwargs): + self.calls += 1 + raise RuntimeError("injected partial batch write") + + client = FailingClient() + sink = TQTokenSink(client, staging_partition="test") + records, _, _ = build_fixture_artifacts("worked_example") + results = sink.stage_generation_prefix_batch( + records, checkpoint_id="checkpoint", chunk_sequences=[0, 1] + ) + assert client.calls == 1 + assert len(results) == len(records) + assert all(not result.ok for result in results) + assert all(result.staging_key for result in results) + assert all("injected partial batch write" in result.error for result in results) + + +def test_prefix_batch_benchmark_1000_rows_issue_four_puts(): + class RecordingClient: + def __init__(self): + self.batches = [] + + def put_samples(self, **kwargs): + self.batches.append(list(kwargs["sample_ids"])) + + client = RecordingClient() + sink = TQTokenSink(client, staging_partition="test") + timings = _write_ranges( + sink, + range(1000), + checkpoint_id="checkpoint", + prefix_tokens=4, + writers=1, + batch_size=256, + ) + assert [len(batch) for batch in client.batches] == [256, 256, 256, 232] + assert len({key for batch in client.batches for key in batch}) == 1000 + assert sum(timing.put_calls for timing in timings) == 4 + + +@pytest.mark.parametrize("batch_size,expected_calls", [(1, 11), (4, 3), (256, 1)]) +def test_tq_prefix_write_benchmark_batches_and_flushes_tail( + tq_client, batch_size, expected_calls +): + partition_id = f"prefix_batch_benchmark_{batch_size}" + result = run_prefix_write_benchmark( + tq_client, + BenchmarkConfig( + rows=11, + prefix_tokens=4, + batch_size=batch_size, + partition_id=partition_id, + verify_rows=11, + ), + ) + assert result["put_calls"] == expected_calls + assert result["stored_keys"] == 11 + assert result["verified_rows"] == 11 + assert tq_client.list_sample_ids(partition_id) == [] 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 dd4297fc110..23f6748bf13 100644 --- a/tests/unit/models/generation/test_vllm_token_capture_hosting.py +++ b/tests/unit/models/generation/test_vllm_token_capture_hosting.py @@ -49,6 +49,10 @@ StageResult, ) +from nemo_rl.models.generation.generation_cut_capture import ( # noqa: E402 + GenerationCutCaptureMixin, + GenerationPrefixBatchLimits, +) from nemo_rl.models.generation.vllm.vllm_generation import VllmGeneration # noqa: E402 from nemo_rl.models.generation.vllm.vllm_worker_async import ( # noqa: E402 VllmAsyncGenerationWorkerImpl, @@ -176,9 +180,13 @@ def stage_generation_prefix( ) -def _fake_worker(*, is_model_owner: bool = True) -> SimpleNamespace: +class _FakeWorker(GenerationCutCaptureMixin, SimpleNamespace): + """Minimal host for shared capture logic plus vLLM setup seams.""" + + +def _fake_worker(*, is_model_owner: bool = True) -> _FakeWorker: """The attribute surface setup_token_capture touches, minus the engine.""" - worker = SimpleNamespace( + worker = _FakeWorker( is_model_owner=is_model_owner, token_capture=None, _rollout_weight_version=0, @@ -191,6 +199,7 @@ def _fake_worker(*, is_model_owner: bool = True) -> SimpleNamespace: _generation_prefix_cuts_enabled=False, _generation_cut_control_token=None, _generation_chunk_flush_tokens=0, + _start_generation_chunk_flusher=MagicMock(), _generation_checkpoint_gate=_CheckpointCaptureGate(), _staging_source=None, _prefix_cache={}, @@ -216,11 +225,20 @@ def test_setup_token_capture_installs_capture_with_vllm_adapter(monkeypatch): installed = asyncio.run( VllmAsyncGenerationWorkerImpl.setup_token_capture( - worker, dp_cfg={"backend": "simple"}, staging_partition="rollout_staging" + worker, + dp_cfg={"backend": "simple"}, + staging_partition="rollout_staging", + generation_prefix_batch_size=256, + generation_prefix_batch_max_tokens=4_194_304, ) ) assert installed is True + worker._start_generation_chunk_flusher.assert_called_once_with() + assert worker._generation_prefix_batch_limits == GenerationPrefixBatchLimits( + max_rows=256, + max_tokens=4_194_304, + ) assert isinstance(worker.token_capture, RolloutTokenCapture) assert worker.token_capture.adapter is not None # The adapter is the vLLM one (prefix ids enter via the worker's field). @@ -232,13 +250,30 @@ def test_setup_token_capture_skips_non_model_owners(monkeypatch): worker = _fake_worker(is_model_owner=False) installed = asyncio.run( VllmAsyncGenerationWorkerImpl.setup_token_capture( - worker, dp_cfg={}, staging_partition="rollout_staging" + worker, + dp_cfg={}, + staging_partition="rollout_staging", + generation_prefix_batch_size=256, + generation_prefix_batch_max_tokens=4_194_304, ) ) assert installed is False + worker._start_generation_chunk_flusher.assert_not_called() assert worker.token_capture is None +@pytest.mark.parametrize("max_rows,max_tokens", [(0, 1), (1, 0), (-1, 1), (1, -1)]) +def test_shared_generation_prefix_batch_limits_require_positive_values( + max_rows: int, max_tokens: int +) -> None: + worker = _fake_worker() + with pytest.raises(ValueError, match="row and token limits must be positive"): + worker._configure_generation_prefix_batching( + max_rows=max_rows, + max_tokens=max_tokens, + ) + + def test_weight_version_is_stamped_from_worker_state(monkeypatch): """The install closure reads _rollout_weight_version live: a set_rollout_weight_version between calls changes the stamp.""" @@ -254,7 +289,11 @@ def test_weight_version_is_stamped_from_worker_state(monkeypatch): worker = _fake_worker() asyncio.run( VllmAsyncGenerationWorkerImpl.setup_token_capture( - worker, dp_cfg={}, staging_partition="rollout_staging" + worker, + dp_cfg={}, + staging_partition="rollout_staging", + generation_prefix_batch_size=256, + generation_prefix_batch_max_tokens=4_194_304, ) ) @@ -290,13 +329,21 @@ def test_generation_setup_token_capture_fans_out(monkeypatch): "nemo_rl.models.generation.vllm.vllm_generation.ray.get", lambda futures, timeout=None: futures, ) - gen.setup_token_capture({"backend": "simple"}, "rollout_staging") + gen.setup_token_capture( + {"backend": "simple"}, + "rollout_staging", + generation_prefix_batch_size=128, + generation_prefix_batch_max_tokens=2048, + ) gen.worker_group.run_all_workers_single_data.assert_called_once_with( "setup_token_capture", dp_cfg={"backend": "simple"}, staging_partition="rollout_staging", generation_prefix_cuts_enabled=False, generation_cut_control_token=None, + generation_chunk_flush_tokens=0, + generation_prefix_batch_size=128, + generation_prefix_batch_max_tokens=2048, run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], ) @@ -304,7 +351,12 @@ def test_generation_setup_token_capture_fans_out(monkeypatch): def test_generation_setup_token_capture_requires_async_engine(): gen = _generation_with_mock_group(async_engine=False) with pytest.raises(AssertionError, match="async vLLM engine"): - gen.setup_token_capture({}, "rollout_staging") + gen.setup_token_capture( + {}, + "rollout_staging", + generation_prefix_batch_size=256, + generation_prefix_batch_max_tokens=4_194_304, + ) @pytest.mark.parametrize( @@ -657,6 +709,10 @@ def _worker_with_capture(sink: _MemorySink): from nemo_gym.token_id_capture.adapters.vllm import VLLMCaptureAdapter worker = _fake_worker() + worker._configure_generation_prefix_batching( + max_rows=1, + max_tokens=4_194_304, + ) worker._capture_calls = {} worker._capture_calls_by_model_call_id = {} worker._completed_capture_calls = {} @@ -674,13 +730,8 @@ def _worker_with_capture(sink: _MemorySink): "_fetch_chain_prefix", "_capture_admission", "_resolve_admission_prefix", - "_resolve_generation_cut", "_enter_request_prefix", - "_get_request_capture", - "_pop_request_capture", - "_remember_completed_capture", - "_completed_generation_cut_ack", - "_checkpoint_active_generation_cut", + "_observe_request_capture", "_finish_request_capture_after_checkpoint_gate", "_finish_request_capture_with_lifecycle_owned", ): @@ -695,6 +746,259 @@ def _worker_with_capture(sink: _MemorySink): return worker +class _BatchPrefixSink(_MemorySink): + def __init__(self) -> None: + super().__init__() + self.batches: list[list[int]] = [] + self.failure: str | None = None + self.before_write = lambda: None + + def stage_generation_prefix_batch(self, records, *, checkpoint_id, chunk_sequences): + self.batches.append([len(record.token_ids_delta) for record in records]) + self.before_write() + results = [ + self.stage_generation_prefix( + record, checkpoint_id=checkpoint_id, chunk_sequence=sequence + ) + for record, sequence in zip(records, chunk_sequences, strict=True) + ] + if self.failure == "raise": + raise RuntimeError("injected batch transport failure after writes") + if self.failure == "short": + return results[:-1] + if self.failure == "partial": + results[-1] = StageResult( + ok=False, + staging_key=results[-1].staging_key, + error="injected partial failure", + ) + if self.failure == "wrong-key": + results[-1] = StageResult(ok=True, staging_key="wrong-key") + return results + + +def _batch_worker_fixture(lengths, *, batch_size=256, max_tokens=4_194_304): + sink = _BatchPrefixSink() + worker = _worker_with_capture(sink) + worker._configure_generation_prefix_batching( + max_rows=batch_size, + max_tokens=max_tokens, + ) + requests, prefixes = [], [] + for i, length in enumerate(lengths): + request = _FakeRequest( + ng_capture={ + "rollout_id": f"r{i}", + "model_call_id": f"c{i}", + "parent_call_id": None, + "prev_len": 0, + "mode": "text", + }, + stream=False, + ) + VllmAsyncGenerationWorkerImpl._begin_request_capture(worker, request, [10]) + state = worker._capture_calls[id(request)] + state.effective_output_limit = 128 + with state.lock: + state.observe(list(range(20, 20 + length)), [-0.1] * length) + requests.append(request) + prefixes.append( + GenerationCutPrefix( + ticket_id=f"t{i}", + rollout_id=f"r{i}", + attempt_index=0, + model_call_id=f"c{i}", + admitted_at=1.0, + ) + ) + inventory = GenerationCutInventory.build( + checkpoint_id="batch-cut", server_name="policy", active_prefixes=prefixes + ) + return worker, sink, requests, inventory + + +def test_worker_batches_1000_prefixes_and_caches_receipt(): + worker, sink, requests, inventory = _batch_worker_fixture([2] * 1000) + receipt = worker._checkpoint_generation_cut(inventory) + assert [len(batch) for batch in sink.batches] == [256, 256, 256, 232] + assert [ack.ticket_id for ack in receipt.prefixes] == [ + p.ticket_id for p in inventory.active_prefixes + ] + assert all(ack.prefix_token_count == 2 for ack in receipt.prefixes) + assert all(worker._capture_calls[id(r)].frozen_buffer is None for r in requests) + assert worker._checkpoint_generation_cut(inventory) is receipt + assert len(sink.batches) == 4 + + +def test_worker_batch_token_limit_ragged_and_oversized_rows(): + worker, sink, _, inventory = _batch_worker_fixture([1, 2, 8, 1], max_tokens=5) + worker._checkpoint_generation_cut(inventory) + assert sink.batches == [[2, 3], [9], [2]] # First chunk includes prompt. + + +@pytest.mark.parametrize("failure", ["raise", "short", "partial", "wrong-key"]) +def test_worker_batch_failure_rolls_back_every_unsealed_row(failure): + worker, sink, requests, inventory = _batch_worker_fixture([2, 3, 4]) + sink.failure = failure + with pytest.raises(RuntimeError): + worker._checkpoint_generation_cut(inventory) + assert not worker._generation_cut_receipts + assert set(sink.generation_prefix_keys) <= set(sink.cleared_generation_prefix_keys) + for request, length in zip(requests, [2, 3, 4], strict=True): + state = worker._capture_calls[id(request)] + assert state.frozen_buffer is None + assert not state.generation_cut_staging_keys + assert not state.sealed_generated_token_ids + assert len(state.active_buffer.generated_token_ids) == length + assert state.lifecycle_lock.acquire(blocking=False) + state.lifecycle_lock.release() + sink.failure = None + receipt = worker._checkpoint_generation_cut(inventory) + assert [ack.prefix_token_count for ack in receipt.prefixes] == [2, 3, 4] + + +def test_worker_batch_ack_validation_is_all_before_seal(monkeypatch): + from nemo_gym._checkpoint import model_control_contracts + + worker, sink, requests, inventory = _batch_worker_fixture([2, 3]) + original = model_control_contracts.GenerationCutPrefixAck + + def reject_second(**kwargs): + if kwargs["model_call_id"] == "c1": + raise ValueError("injected ack failure") + return original(**kwargs) + + monkeypatch.setattr( + model_control_contracts, "GenerationCutPrefixAck", reject_second + ) + with pytest.raises(ValueError, match="injected ack failure"): + worker._checkpoint_generation_cut(inventory) + for request in requests: + state = worker._capture_calls[id(request)] + assert state.frozen_buffer is None and not state.sealed_generated_token_ids + assert not state.generation_cut_staging_keys + assert set(sink.generation_prefix_keys) <= set(sink.cleared_generation_prefix_keys) + + +def test_worker_batch_retains_earlier_success_on_later_batch_failure(): + worker, sink, requests, inventory = _batch_worker_fixture([2] * 3, batch_size=2) + + def fail_second(): + if len(sink.batches) == 2: + sink.failure = "partial" + + sink.before_write = fail_second + with pytest.raises(RuntimeError): + worker._checkpoint_generation_cut(inventory) + earlier_keys = set(sink.generation_prefix_keys[:2]) + assert earlier_keys.isdisjoint(sink.cleared_generation_prefix_keys) + sink.failure = None + sink.before_write = lambda: None + for request in requests: + state = worker._capture_calls[id(request)] + with state.lock: + state.observe([99], [-0.2]) + receipt = worker._checkpoint_generation_cut(inventory) + assert [ack.prefix_token_count for ack in receipt.prefixes] == [3, 3, 3] + assert [len(ack.staging_keys) for ack in receipt.prefixes] == [2, 2, 1] + assert earlier_keys.isdisjoint(sink.cleared_generation_prefix_keys) + + +def test_worker_batch_observation_continues_but_lifecycle_waits(): + worker, sink, requests, inventory = _batch_worker_fixture([2, 3]) + state = worker._capture_calls[id(requests[0])] + write_started, release_write, lifecycle_done = ( + threading.Event(), + threading.Event(), + threading.Event(), + ) + + def block_write(): + write_started.set() + assert release_write.wait(5) + + sink.before_write = block_write + with ThreadPoolExecutor(max_workers=2) as executor: + cut = executor.submit(worker._checkpoint_generation_cut, inventory) + assert write_started.wait(5) + + def terminal_operation(): + with state.lifecycle_lock: + lifecycle_done.set() + + terminal = executor.submit(terminal_operation) + try: + assert not lifecycle_done.wait(0.05) + with state.lock: + state.observe([99], [-0.2]) + finally: + release_write.set() + receipt = cut.result(timeout=5) + terminal.result(timeout=5) + assert receipt.prefixes[0].prefix_token_count == 2 + assert state.active_buffer.generated_token_ids == [99] + assert state.sealed_generated_token_ids == [20, 21] + + +@pytest.mark.parametrize("batch_size", [1, 256]) +def test_worker_batch_consecutive_cuts_reuse_then_append(batch_size): + worker, sink, requests, inventory = _batch_worker_fixture( + [2, 3], batch_size=batch_size + ) + first = worker._checkpoint_generation_cut(inventory) + second_inventory = GenerationCutInventory.build( + checkpoint_id="second-cut", + server_name="policy", + active_prefixes=inventory.active_prefixes, + ) + second = worker._checkpoint_generation_cut(second_inventory) + assert [ack.staging_keys for ack in second.prefixes] == [ + ack.staging_keys for ack in first.prefixes + ] + assert len(sink.generation_prefix_keys) == 2 + for request in requests: + state = worker._capture_calls[id(request)] + with state.lock: + state.observe([99], [-0.2]) + third_inventory = GenerationCutInventory.build( + checkpoint_id="third-cut", + server_name="policy", + active_prefixes=inventory.active_prefixes, + ) + third = worker._checkpoint_generation_cut(third_inventory) + assert [ack.prefix_token_count for ack in third.prefixes] == [3, 4] + assert all(len(ack.staging_keys) == 2 for ack in third.prefixes) + assert [ + list(record.token_ids_delta) + for _, record in sink.generation_prefix_records[-2:] + ] == [[99], [99]] + + +def test_worker_batch_preparation_failure_releases_pending_calls(): + worker, sink, requests, inventory = _batch_worker_fixture([2, 3]) + worker._capture_calls[id(requests[1])].effective_output_limit = None + with pytest.raises(RuntimeError): + worker._checkpoint_generation_cut(inventory) + assert not sink.batches + for request in requests: + state = worker._capture_calls[id(request)] + assert state.frozen_buffer is None + assert not state.sealed_generated_token_ids + assert state.lifecycle_lock.acquire(blocking=False) + state.lifecycle_lock.release() + + +def test_worker_batch_mixed_zero_token_and_active_calls(): + worker, sink, _, inventory = _batch_worker_fixture([2, 0, 3]) + receipt = worker._checkpoint_generation_cut(inventory) + assert [len(batch) for batch in sink.batches] == [2] + assert [ack.disposition for ack in receipt.prefixes] == [ + "durable_prefix", + "durable_failure", + "durable_prefix", + ] + + class _MemoryPrefixSource: def __init__( self, @@ -790,7 +1094,12 @@ def test_request_capture_round_trip_stages_and_rides_coords( assert coords["cumulative_hash"] == sink.records[0].cumulative_hash # Only the ordinary response fields and coords transit worker -> gate. assert content["choices"] == [ - {"index": 0, "message": {"role": "assistant", "content": "x"}} + { + "index": 0, + "finish_reason": "stop", + "stop_reason": "", + "message": {"role": "assistant", "content": "x"}, + } ] assert worker._capture_calls == {} @@ -1289,7 +1598,7 @@ def test_abort_waits_for_generation_cut_before_failing_the_call(): assert not abort_thread.is_alive() assert receipt[0].prefixes[0].disposition == "durable_prefix" assert len(sink.generation_prefix_records) == 1 - assert worker._completed_capture_calls["c1"].coords.disposition == "failed" + assert worker._completed_capture_calls["c1"].coords.disposition == "capture_failed" def test_terminal_completion_waits_for_concurrent_generation_cut_and_publishes_once(): diff --git a/tools/benchmark_prefix_checkpoint_proxy.py b/tools/benchmark_prefix_checkpoint_proxy.py new file mode 100644 index 00000000000..3f54f73d6e8 --- /dev/null +++ b/tools/benchmark_prefix_checkpoint_proxy.py @@ -0,0 +1,749 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU proxy: real TQ prefix writes inside Gym model prepare. + +Includes production sink encoding, bounded concurrent PUTs, worker journals, +model-ledger commit, and post-timing lineage/coordinate restore verification. +Excludes vLLM cuts, agent state/ACKs, tools, TQ snapshot save/load, and networking +between physical nodes. Shared mode uses logical partitions in one process; +per-worker mode uses one synchronous Ray actor/client per generation worker. +Synthetic records are built before prepare; their immutable payload lists are +shared to avoid holding a second Python-list copy of the entire TQ dataset. +""" + +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import importlib.util +import json +import math +import os +import resource +import struct +import sys +import tempfile +import threading +import time +from collections import defaultdict +from collections.abc import Awaitable +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any + + +class CachedRecordBuilder: + """Benchmark-only v2 digest encoder for a fixed, validated payload. + + Owns a validated copy of the template. Returned records share its lists; + callers must treat them as immutable. Only IDs and digest vary. No TQ + tensors are precomputed: sink encoding remains inside prepare timing. + """ + + def __init__(self, template: Any) -> None: + # Gym is optional for the dependency-light scheduling tests. + from nemo_gym.token_id_capture.staging import digest as wire + from nemo_gym.token_id_capture.staging.records import StagedCallRecord + + self._template = StagedCallRecord.model_validate(template.model_dump()) + record = self._template + if ( + record.schema_version, + record.digest_version, + record.extras_digest_version, + ) != (2, 2, 1): + raise ValueError("Cached benchmark builder requires staging digest v2") + self._header = wire._CALL_DIGEST_DOMAIN + struct.pack( + ">BBB", + record.schema_version, + record.digest_version, + record.extras_digest_version, + ) + # This is exactly the suffix following the two IDs in the production + # v2 wire layout. Reuse its encoders, including length/presence markers. + self._suffix = b"".join( + ( + wire._encode_optional_text(record.parent_call_id), + wire._encode_text(record.mode), + wire._encode_uint(record.prev_len, field="prev_len"), + wire._encode_uint(record.delta_len, field="delta_len"), + wire._encode_uint(record.cum_len, field="cum_len"), + wire._encode_uint(record.weight_version, field="weight_version"), + wire._encode_bytes(wire.encode_token_ids(record.token_ids_delta)), + wire._encode_bytes( + wire._encode_float32_values( + record.token_mask_delta, field="token_mask_delta" + ) + ), + wire._encode_bytes( + wire._encode_float32_values( + record.generation_log_probs_delta, + field="generation_log_probs_delta", + ) + ), + bytes.fromhex(record.extras_digest), + wire._encode_present_digest(record.chain_hash, field="chain_hash"), + wire._encode_present_digest( + record.cumulative_hash, field="cumulative_hash" + ), + ) + ) + # Fail closed on wire-layout drift before generating a large inventory. + if ( + self.build( + rollout_id=record.rollout_id, model_call_id=record.model_call_id + ).digest + != record.digest + ): + raise ValueError( + "Cached benchmark digest disagrees with production encoder" + ) + probe = self.build(rollout_id="benchmark-parity-π", model_call_id="call-parity") + StagedCallRecord.model_validate(probe.model_dump()) + + def build(self, *, rollout_id: str, model_call_id: str) -> Any: + """Reuse validated payload, hashing the exact ID-specific v2 bytes.""" + digest = hashlib.sha256(self._header) + for value in (rollout_id, model_call_id): + if not isinstance(value, str) or not value: + raise ValueError( + "rollout_id and model_call_id must be non-empty strings" + ) + encoded = value.encode("utf-8") + digest.update(struct.pack(">Q", len(encoded))) + digest.update(encoded) + digest.update(self._suffix) + return self._template.model_copy( + update={ + "rollout_id": rollout_id, + "model_call_id": model_call_id, + "digest": digest.hexdigest(), + } + ) + + +def worker_for_row(index: int, *, rows: int, workers: int) -> int: + """Contiguous balanced ownership, including uneven and tiny workloads.""" + base, extra = divmod(rows, workers) + boundary = (base + 1) * extra + if index < boundary: + return index // (base + 1) + return extra + (index - boundary) // base + + +def latency_summary(values: list[float]) -> dict[str, float | int]: + """Nearest-rank latency percentiles, in seconds.""" + ordered = sorted(values) + if not ordered: + return {"count": 0} + return { + "count": len(ordered), + "sum_s": sum(ordered), + "mean_s": sum(ordered) / len(ordered), + "p50_s": ordered[math.ceil(0.50 * len(ordered)) - 1], + "p95_s": ordered[math.ceil(0.95 * len(ordered)) - 1], + "p99_s": ordered[math.ceil(0.99 * len(ordered)) - 1], + "max_s": ordered[-1], + } + + +class TimedPutClient: + """Instrument the exact synchronous client API boundary used by the sink.""" + + def __init__(self, client: Any) -> None: + self.client = client + self.lock = threading.Lock() + self.latencies: list[float] = [] + self.batch_rows: list[int] = [] + self.failures = 0 + self.active = 0 + self.peak_active = 0 + + def put_samples(self, **kwargs: Any) -> Any: + with self.lock: + self.active += 1 + self.peak_active = max(self.peak_active, self.active) + started = time.perf_counter() + succeeded = False + try: + result = self.client.put_samples(**kwargs) + succeeded = True + return result + finally: + elapsed = time.perf_counter() - started + with self.lock: + self.active -= 1 + self.latencies.append(elapsed) + self.batch_rows.append(len(kwargs["sample_ids"])) + self.failures += int(not succeeded) + + +class BoundedStager: + """Bound whole sink calls (encoding plus PUT), independently of owners.""" + + def __init__(self, sink: Any, *, concurrency: int) -> None: + if concurrency < 1: + raise ValueError("concurrency must be positive") + self.sink = sink + self.semaphore = asyncio.Semaphore(concurrency) + self.pool = ThreadPoolExecutor(max_workers=concurrency) + self.waits: list[float] = [] + self.stages: list[float] = [] + self.first_started: float | None = None + self.last_finished: float | None = None + + async def stage(self, records: list[Any], checkpoint_id: str) -> list[Any]: + queued = time.perf_counter() + if self.first_started is None: + self.first_started = queued + async with self.semaphore: + started = time.perf_counter() + self.waits.append(started - queued) + results = await asyncio.get_running_loop().run_in_executor( + self.pool, + lambda: self.sink.stage_generation_prefix_batch( + records, + checkpoint_id=checkpoint_id, + chunk_sequences=[0] * len(records), + ), + ) + self.last_finished = time.perf_counter() + self.stages.append(self.last_finished - started) + if len(results) != len(records) or any(not result.ok for result in results): + raise RuntimeError( + "Prefix staging failed; no successful cut receipt will be published" + ) + return results + + def close(self) -> None: + """Join outstanding writes before the client/storage can be closed.""" + self.pool.shutdown(wait=True, cancel_futures=True) + + +class RealTQCutBackend: + """Replace synthetic coordinates with acknowledged TQ prefix coordinates.""" + + def __init__( + self, + owner: str, + records: list[Any], + stager: BoundedStager | PerWorkerStager, + args: argparse.Namespace, + ) -> None: + self.owner = owner + self.records = records + self.stager = stager + self.args = args + self.keys: dict[int, str] = {} + + async def checkpoint_generation_cut(self, inventory: Any) -> Any: + # Gym is optional; only the integrated runtime needs these contracts. + from nemo_gym._checkpoint.model_control_contracts import ( + GenerationCutPrefixAck, + GenerationCutReceipt, + ) + + groups: dict[int, list[Any]] = defaultdict(list) + for prefix in inventory.active_prefixes: + index = int(prefix.rollout_id.rsplit("-", 1)[1]) + groups[ + worker_for_row( + index, rows=self.args.rows, workers=self.args.generation_workers + ) + ].append(prefix) + + async def write_owner(prefixes: list[Any]) -> list[Any]: + acknowledgements = [] + for offset in range(0, len(prefixes), self.args.batch_size): + batch = prefixes[offset : offset + self.args.batch_size] + indices = [int(prefix.rollout_id.rsplit("-", 1)[1]) for prefix in batch] + records = [self.records[index] for index in indices] + results = await self.stager.stage(records, inventory.checkpoint_id) + for prefix, index, record, result in zip( + batch, indices, records, results, strict=True + ): + if ( + prefix.model_call_id != record.model_call_id + or prefix.rollout_id != record.rollout_id + ): + raise RuntimeError( + "Cut inventory identity does not match staged record" + ) + self.keys[index] = result.staging_key + acknowledgements.append( + GenerationCutPrefixAck( + **prefix.model_dump(mode="json"), + disposition="durable_prefix", + cut_kind="active_prefix", + frozen_buffer_id=f"active/{inventory.checkpoint_id}/{prefix.ticket_id}", + staging_keys=(result.staging_key,), + prefix_token_count=self.args.prefix_tokens, + prefix_digest=record.digest, + effective_output_limit=self.args.prefix_tokens + 1, + ) + ) + return acknowledgements + + tasks = [asyncio.create_task(write_owner(group)) for group in groups.values()] + try: + chunks = await asyncio.gather(*tasks) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + return GenerationCutReceipt( + checkpoint_id=inventory.checkpoint_id, + cut_id=f"cut-{self.owner}-{inventory.inventory_digest[:16]}", + inventory_digest=inventory.inventory_digest, + inventory=inventory, + backend_snapshot_id=f"snapshot-{self.owner}-{inventory.inventory_digest[:16]}", + prefixes=tuple(prefix for chunk in chunks for prefix in chunk), + ) + + +class GenerationWriter: + """Synchronous Ray actor implementation: one local client, one batch at a time.""" + + def __init__( + self, dp_config: Any, partition: str, start: int, stop: int, prefix_tokens: int + ) -> None: + # Ray/TQ dependencies are needed only inside actual writer processes. + import ray + from nemo_rl.data_plane import build_data_plane_client + from nemo_rl.data_plane.tq_token_sink import TQTokenSink + from tools.benchmark_tq_prefix_writes import _build_record, _record_payload + + self.start = start + self.stop = stop + self.measured = TimedPutClient( + build_data_plane_client(dp_config, bootstrap=False) + ) + self.sink = TQTokenSink(self.measured, staging_partition=partition) + started = time.perf_counter() + builder = CachedRecordBuilder(_build_record(0, *_record_payload(prefix_tokens))) + self.records = { + i: builder.build( + rollout_id=f"rollout-{i:09d}", model_call_id=f"call-{i:09d}" + ) + for i in range(start, stop) + } + self.placement = { + "pid": os.getpid(), + "node_id": str(ray.get_runtime_context().get_node_id()), + "start": start, + "stop": stop, + "record_setup_seconds": time.perf_counter() - started, + } + + def ready(self) -> dict[str, Any]: + return self.placement + + def stage(self, indices: list[int], checkpoint_id: str) -> dict[str, Any]: + if ( + not indices + or len(set(indices)) != len(indices) + or any(i not in self.records for i in indices) + ): + raise ValueError( + "Batch must contain unique rows owned by this generation worker" + ) + before = len(self.measured.latencies) + started = time.perf_counter() + results = self.sink.stage_generation_prefix_batch( + [self.records[i] for i in indices], + checkpoint_id=checkpoint_id, + chunk_sequences=[0] * len(indices), + ) + stage_seconds = time.perf_counter() - started + if len(results) != len(indices) or any(not result.ok for result in results): + raise RuntimeError( + "Worker prefix staging failed; no cut receipt may be published" + ) + return { + "results": [result.model_dump() for result in results], + "stage_seconds": stage_seconds, + "put_latencies": self.measured.latencies[before:], + "batch_rows": self.measured.batch_rows[before:], + } + + +class PerWorkerStager: + """Route ID-only batch requests to independent, pre-initialized clients. + + No shared eight-slot gate: each actor serializes its own batches while + separate actors run concurrently. Ray RPC time is measured separately. + """ + + def __init__(self, measured: TimedPutClient, args: argparse.Namespace) -> None: + self.measured = measured + self.args = args + self.actors: list[Any] = [] + self.placements: list[dict[str, Any]] = [] + self.setup_seconds = 0.0 + self.waits: list[float] = [] + self.stages: list[float] = [] + self.rpc_latencies: list[float] = [] + self.first_started: float | None = None + self.last_finished: float | None = None + self.active_requests = 0 + self.peak_requests = 0 + + def setup(self, dp_config: Any, partition: str) -> None: + # Optional Ray dependency only needed for the process-isolated mode. + import ray + from tools.benchmark_tq_prefix_writes import _wait_for_actor_results + + started = time.perf_counter() + writer_class = ray.remote( + num_cpus=self.args.worker_cpus, + max_concurrency=1, + max_restarts=0, + max_task_retries=0, + )(GenerationWriter) + width = self.args.rows // self.args.generation_workers + for owner in range(self.args.generation_workers): + self.actors.append( + writer_class.remote( + dp_config, + partition, + owner * width, + (owner + 1) * width, + self.args.prefix_tokens, + ) + ) + self.placements = _wait_for_actor_results( + [actor.ready.remote() for actor in self.actors], + stage="proxy_worker_setup", + timeout_s=self.args.client_setup_timeout_s, + progress_interval_s=10, + ) + self.setup_seconds = time.perf_counter() - started + + async def stage(self, records: list[Any], checkpoint_id: str) -> list[Any]: + indices = [int(record.rollout_id.rsplit("-", 1)[1]) for record in records] + owners = { + worker_for_row(i, rows=self.args.rows, workers=self.args.generation_workers) + for i in indices + } + if len(owners) != 1: + raise ValueError("Batch crosses generation-worker ownership") + started = time.perf_counter() + if self.first_started is None: + self.first_started = started + self.active_requests += 1 + self.peak_requests = max(self.peak_requests, self.active_requests) + try: + payload = await self.actors[owners.pop()].stage.remote( + indices, checkpoint_id + ) + finally: + self.active_requests -= 1 + self.last_finished = time.perf_counter() + self.rpc_latencies.append(self.last_finished - started) + self.stages.append(payload["stage_seconds"]) + self.measured.latencies.extend(payload["put_latencies"]) + self.measured.batch_rows.extend(payload["batch_rows"]) + return [SimpleNamespace(**result) for result in payload["results"]] + + def close(self) -> None: + # Attached clients must not close the shared TQ deployment. Kill only + # actors created by this harness, as in the existing multi-client test. + import ray + + for actor in self.actors: + ray.kill(actor, no_restart=True) + self.actors.clear() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + for name, default in ( + ("rows", 512), + ("prefix-tokens", 16384), + ("generation-workers", 128), + ("gym-workers", 8), + ("batch-size", 256), + ("num-storage-units", 8), + ("verify-rows", 256), + ("timeout-s", 1800), + ): + parser.add_argument(f"--{name}", type=int, default=default) + parser.add_argument( + "--client-mode", choices=("shared", "per-worker"), default="shared" + ) + parser.add_argument( + "--put-concurrency", + type=int, + default=None, + help="Shared-client mode only; defaults to 8", + ) + parser.add_argument( + "--worker-cpus", + type=float, + default=0.25, + help="Ray scheduling CPU reservation per writer actor; not a CPU thread cap", + ) + parser.add_argument("--client-setup-timeout-s", type=int, default=300) + parser.add_argument( + "--root", + type=Path, + required=True, + help="Artifact filesystem; use shared storage for representative commit timing", + ) + parser.add_argument("--output-json", type=Path, required=True) + args = parser.parse_args() + if args.client_mode == "per-worker" and args.put_concurrency is not None: + parser.error( + "--put-concurrency applies only to shared mode; per-worker writes are sequential per actor" + ) + if args.client_mode == "shared" and args.put_concurrency is None: + args.put_concurrency = 8 + if not math.isfinite(args.worker_cpus) or args.worker_cpus <= 0: + parser.error("worker-cpus must be finite and positive") + for name, value in vars(args).items(): + if isinstance(value, int) and value < 1: + parser.error(f"{name} must be positive") + # Aligned contiguous policy inventories keep each generation owner within + # exactly one policy participant, avoiding artificial cross-inventory batches. + if ( + args.generation_workers % args.gym_workers + or args.rows % args.generation_workers + ): + parser.error( + "generation-workers must divide rows and be divisible by gym-workers" + ) + return args + + +async def with_progress( + operation: Awaitable[dict[str, Any]], + measured: TimedPutClient, + stager: BoundedStager | PerWorkerStager | None = None, +) -> dict[str, Any]: + """Report progress through prepare, commit and post-timing ledger restore.""" + + async def report() -> None: + started = time.perf_counter() + while True: + await asyncio.sleep(10) + with measured.lock: + calls = len(measured.latencies) + rows = sum(measured.batch_rows) + active = measured.active + failures = measured.failures + activity = ( + f"remote_stage_requests={stager.active_requests}" + if isinstance(stager, PerWorkerStager) + else f"active_puts={active}" + ) + print( + f"proxy scope=prepare_commit_restore_verify puts_completed={calls} attempted_rows={rows} {activity} failed_puts={failures} pipeline_elapsed_s={time.perf_counter() - started:.1f}", + flush=True, + ) + + reporter = asyncio.create_task(report()) + try: + return await operation + finally: + reporter.cancel() + await asyncio.gather(reporter, return_exceptions=True) + + +def load_checkpoint_driver() -> ModuleType: + """Load the adjacent Gym benchmark without relying on a scripts package.""" + gym_script = ( + Path(__file__).resolve().parents[1] + / "3rdparty/Gym-workspace/Gym/scripts/benchmark_checkpoint_cpu.py" + ) + spec = importlib.util.spec_from_file_location( + "gym_checkpoint_proxy_driver", gym_script + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"Cannot load {gym_script}") + gym = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = gym + spec.loader.exec_module(gym) + return gym + + +def main() -> None: + # Heavy Ray/Torch/Gym imports are optional for helper unit tests. + from tools.benchmark_tq_prefix_writes import ( + BenchmarkConfig, + _build_record, + _data_plane_config, + _record_payload, + _verification_indices, + ) + from nemo_rl.data_plane import build_data_plane_client + from nemo_rl.data_plane.tq_token_sink import ( + STAGING_FIELDS, + TQTokenSink, + TQTokenSource, + generation_cut_staging_key, + ) + + args = parse_args() + gym = load_checkpoint_driver() + args.root.mkdir(parents=True, exist_ok=True) + run_root = Path(tempfile.mkdtemp(prefix="prefix-proxy-", dir=args.root)) + config = BenchmarkConfig( + rows=args.rows, + prefix_tokens=args.prefix_tokens, + num_storage_units=args.num_storage_units, + ) + started = time.perf_counter() + payload = _record_payload(args.prefix_tokens) + records = [] + print("proxy stage=record_setup started (excluded from prepare)", flush=True) + print( + f"proxy token_tensor_bytes_lower_bound={args.rows * args.prefix_tokens * 16} (excludes metadata, staging copies and Ray overhead)", + flush=True, + ) + builder = CachedRecordBuilder(_build_record(0, *payload)) + for index in range(args.rows): + record = builder.build( + rollout_id=f"rollout-{index:09d}", + model_call_id=f"call-{index:09d}", + ) + records.append(record) + if (index + 1) % 1024 == 0: + print(f"proxy stage=record_setup rows={index + 1}/{args.rows}", flush=True) + record_setup_seconds = time.perf_counter() - started + client = build_data_plane_client(_data_plane_config(config), bootstrap=True) + partition = config.partition_id + measured = TimedPutClient(client) + stager = ( + PerWorkerStager(measured, args) + if args.client_mode == "per-worker" + else BoundedStager( + TQTokenSink(measured, staging_partition=partition), + concurrency=args.put_concurrency, + ) + ) + backends: list[RealTQCutBackend] = [] + + def backend_factory(owner: str) -> RealTQCutBackend: + backend = RealTQCutBackend(owner, records, stager, args) + backends.append(backend) + return backend + + try: + client.register_partition( + partition_id=partition, + fields=list(STAGING_FIELDS), + num_samples=args.rows, + consumer_tasks=["benchmark"], + ) + if isinstance(stager, PerWorkerStager): + stager.setup(_data_plane_config(config), partition) + print("proxy stage=checkpoint started", flush=True) + checkpoint = asyncio.run( + with_progress( + gym.run_checkpoint_benchmark( + run_root, + workers=args.gym_workers, + cuts=args.rows, + hot_worker_fraction=1 / args.gym_workers, + prefix_tokens=args.prefix_tokens, + staging_key_bytes=0, + timeout_s=args.timeout_s, + backend_factory=backend_factory, + staging_keys_for_ticket=lambda ticket, checkpoint_id: ( + generation_cut_staging_key( + checkpoint_id, + ticket.rollout_id, + ticket.model_call_id, + chunk_sequence=0, + ), + ), + ), + measured, + stager, + ) + ) + keys = { + index: key for backend in backends for index, key in backend.keys.items() + } + if len(keys) != args.rows or set(client.list_sample_ids(partition)) != set( + keys.values() + ): + raise RuntimeError( + "TQ key inventory does not match successful cut receipts" + ) + source = TQTokenSource(client, staging_partition=partition) + indices = _verification_indices(args.rows, args.verify_rows) + for offset in range(0, len(indices), 16): + batch = indices[offset : offset + 16] + if source.fetch_prefix_token_ids( + [keys[index] for index in batch] + ) != payload[0] * len(batch): + raise RuntimeError("Stored prefix token verification failed") + if stager.first_started is None or stager.last_finished is None: + raise RuntimeError("No prefix staging was measured") + write_window = stager.last_finished - stager.first_started + result = { + "config": { + key: str(value) if isinstance(value, Path) else value + for key, value in vars(args).items() + }, + "scope": "model prepare + real prefix PUTs + model-ledger commit; excludes agent commit and TQ snapshot", + "run_root": str(run_root), + "clients": args.generation_workers + if args.client_mode == "per-worker" + else 1, + "driver_bootstrap_client": True, + "client_setup_seconds": stager.setup_seconds + if isinstance(stager, PerWorkerStager) + else 0.0, + "writer_placements": stager.placements + if isinstance(stager, PerWorkerStager) + else [], + "record_setup_seconds": record_setup_seconds, + "record_builder": "cached_v2_payload", + "checkpoint": checkpoint, + "proxy_prepare_commit_seconds": checkpoint["prepare_commit_seconds"], + "prefix_write_window_seconds": write_window, + "prefix_rows_per_second": args.rows / write_window, + "prefix_token_bytes_per_second": args.rows + * args.prefix_tokens + * 16 + / write_window, + "put_samples_latency": latency_summary(measured.latencies), + "batch_slot_wait": latency_summary(stager.waits), + "sink_stage_latency": latency_summary(stager.stages), + "put_calls": len(measured.latencies), + "put_batch_rows": measured.batch_rows, + "put_failures": measured.failures, + "peak_concurrent_puts": None + if isinstance(stager, PerWorkerStager) + else measured.peak_active, + "peak_remote_stage_requests": stager.peak_requests + if isinstance(stager, PerWorkerStager) + else None, + "remote_stage_roundtrip_latency": latency_summary(stager.rpc_latencies) + if isinstance(stager, PerWorkerStager) + else None, + "stored_keys": len(keys), + "verified_rows": len(indices), + "driver_max_rss_kib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, + "token_tensor_bytes_lower_bound": args.rows * args.prefix_tokens * 16, + } + finally: + try: + stager.close() + finally: + try: + client.clear_samples(sample_ids=None, partition_id=partition) + finally: + client.close() + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/benchmark_tq_prefix_writes.py b/tools/benchmark_tq_prefix_writes.py new file mode 100644 index 00000000000..567864696f1 --- /dev/null +++ b/tools/benchmark_tq_prefix_writes.py @@ -0,0 +1,781 @@ +# 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. + +"""Benchmark generation-prefix writes through the production TQ token sink. + +Batch size one exercises the original single-prefix path. Larger batches call +``TQTokenSink.stage_generation_prefix_batch()``, issuing one synchronous +``DataPlaneClient.put_samples()`` per batch. ``--clients`` creates separate +Ray actor processes that attach to one shared TQ controller/storage deployment, +matching production's one process-local TQ client per generation model owner. +``--writers-per-client`` optionally adds concurrent callers inside each client +with independently bounded batches in each writer. + +``stage_seconds_mean`` is the mean time a row waits for its batch write; +``stage_seconds_per_row_amortized`` divides total PUT time by row count. +``put_calls`` includes each writer's final partial batch. + +Example: + uv run --no-sync python tools/benchmark_tq_prefix_writes.py \ + --rows 10000 --prefix-tokens 128 --clients 32 \ + --writers-per-client 1 \ + --num-storage-units 8 --output-json /tmp/tq-prefix-131k.json +""" + +from __future__ import annotations + +import argparse +import json +import resource +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, cast + +import ray + +from nemo_gym.token_id_capture.staging.digest import ( + compute_chain_hash, + compute_extras_digest, + compute_staging_digest, + hash_token_ids, +) +from nemo_gym.token_id_capture.staging.records import StagedCallRecord + +from nemo_rl.data_plane import DataPlaneConfig, build_data_plane_client +from nemo_rl.data_plane.interfaces import DataPlaneClient +from nemo_rl.data_plane.tq_token_sink import ( + STAGING_FIELDS, + TQTokenSink, + TQTokenSource, + generation_cut_staging_key, +) + +DEFAULT_PARTITION_ID = "tq_prefix_write_benchmark" +DEFAULT_CHECKPOINT_ID = "benchmark-checkpoint" +CONSUMER_TASK = "benchmark" + + +@dataclass(frozen=True) +class BenchmarkConfig: + """Configuration for one TQ generation-prefix write benchmark.""" + + rows: int = 10_000 + prefix_tokens: int = 128 + clients: int = 1 + writers_per_client: int = 1 + batch_size: int = 1 + num_storage_units: int = 8 + storage_capacity: int | None = None + verify_rows: int = 16 + verify_batch_size: int = 16 + partition_id: str = DEFAULT_PARTITION_ID + checkpoint_id: str = DEFAULT_CHECKPOINT_ID + keep_rows: bool = False + client_setup_timeout_s: float = 180.0 + write_timeout_s: float = 1_800.0 + progress_interval_s: float = 5.0 + + +@dataclass(frozen=True) +class _WriterResult: + rows: int + record_build_seconds: float + stage_seconds: float + put_calls: int + row_wait_seconds: float + + +def _validate_config(config: BenchmarkConfig) -> None: + """Reject invalid or misleading benchmark configurations.""" + positive_fields = { + "rows": config.rows, + "prefix_tokens": config.prefix_tokens, + "clients": config.clients, + "writers_per_client": config.writers_per_client, + "batch_size": config.batch_size, + "num_storage_units": config.num_storage_units, + "verify_batch_size": config.verify_batch_size, + } + for name, value in positive_fields.items(): + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + if config.storage_capacity is not None and config.storage_capacity < config.rows: + raise ValueError( + "storage_capacity must be at least rows: " + f"{config.storage_capacity} < {config.rows}" + ) + if config.verify_rows < 0: + raise ValueError(f"verify_rows must be non-negative, got {config.verify_rows}") + positive_float_fields = { + "client_setup_timeout_s": config.client_setup_timeout_s, + "write_timeout_s": config.write_timeout_s, + "progress_interval_s": config.progress_interval_s, + } + for name, value in positive_float_fields.items(): + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + if not config.partition_id: + raise ValueError("partition_id must be non-empty") + if not config.checkpoint_id: + raise ValueError("checkpoint_id must be non-empty") + + +def _data_plane_config(config: BenchmarkConfig) -> DataPlaneConfig: + """Build a simple-backend TQ configuration sized for this run.""" + storage_capacity = config.storage_capacity or max(1_024, config.rows) + return cast( + DataPlaneConfig, + { + "enabled": True, + "impl": "transfer_queue", + "backend": "simple", + "claim_meta_poll_interval_s": 0.05, + "simple": { + "storage_capacity": storage_capacity, + "num_storage_units": config.num_storage_units, + }, + }, + ) + + +def _record_payload(prefix_tokens: int) -> tuple[list[int], list[float], list[float]]: + """Build deterministic token payloads shared by every benchmark row.""" + token_ids = [index % 32_000 for index in range(prefix_tokens)] + return token_ids, [1.0] * prefix_tokens, [-0.5] * prefix_tokens + + +def _build_record( + row_index: int, + token_ids: list[int], + token_mask: list[float], + logprobs: list[float], + *, + rollout_id: str | None = None, + model_call_id: str = "call-0", +) -> StagedCallRecord: + """Build one valid production-shaped generation-prefix record.""" + rollout_id = rollout_id or f"benchmark-rollout-{row_index:09d}" + extras_digest = compute_extras_digest(None) + chain_hash = compute_chain_hash(None, token_ids) + cumulative_hash = hash_token_ids(token_ids) + values: dict[str, Any] = { + "rollout_id": rollout_id, + "model_call_id": model_call_id, + "parent_call_id": None, + "mode": "text", + "prev_len": 0, + "delta_len": len(token_ids), + "cum_len": len(token_ids), + "weight_version": 0, + "token_ids_delta": token_ids, + "token_mask_delta": token_mask, + "generation_log_probs_delta": logprobs, + "extras": None, + "extras_digest": extras_digest, + "chain_hash": chain_hash, + "cumulative_hash": cumulative_hash, + } + digest_values = {key: value for key, value in values.items() if key != "extras"} + return StagedCallRecord( + **values, + digest=compute_staging_digest( + schema_version=2, + digest_version=2, + extras_digest_version=1, + **digest_values, + ), + ) + + +def _split_range(row_range: range, parts: int) -> list[range]: + """Split one unit-step range into balanced, non-overlapping ranges.""" + if row_range.step != 1: + raise ValueError("benchmark row ranges must have a step of one") + part_count = min(len(row_range), parts) + base, remainder = divmod(len(row_range), part_count) + ranges: list[range] = [] + start = row_range.start + for part_index in range(part_count): + size = base + (1 if part_index < remainder else 0) + ranges.append(range(start, start + size)) + start += size + return ranges + + +def _row_ranges(rows: int, writers: int) -> list[range]: + """Split zero-based rows into balanced, non-overlapping writer ranges.""" + return _split_range(range(rows), writers) + + +def _write_range( + sink: TQTokenSink, + row_range: range, + *, + checkpoint_id: str, + token_ids: list[int], + token_mask: list[float], + logprobs: list[float], + batch_size: int = 1, +) -> _WriterResult: + """Build and synchronously stage one writer's prefix rows.""" + build_seconds = 0.0 + stage_seconds = 0.0 + row_wait_seconds = 0.0 + put_calls = 0 + for offset in range(row_range.start, row_range.stop, batch_size): + indices = range(offset, min(offset + batch_size, row_range.stop)) + started = time.perf_counter() + records = [ + _build_record(index, token_ids, token_mask, logprobs) for index in indices + ] + build_seconds += time.perf_counter() - started + + started = time.perf_counter() + if batch_size == 1: + results = [ + sink.stage_generation_prefix( + records[0], + checkpoint_id=checkpoint_id, + chunk_sequence=0, + ) + ] + else: + results = sink.stage_generation_prefix_batch( + records, + checkpoint_id=checkpoint_id, + chunk_sequences=[0] * len(records), + ) + elapsed = time.perf_counter() - started + stage_seconds += elapsed + row_wait_seconds += elapsed * len(records) + put_calls += 1 + if len(results) != len(records): + raise RuntimeError("TQ prefix batch returned an unexpected result count") + for index, result in zip(indices, results, strict=True): + if not result.ok: + raise RuntimeError( + f"TQ prefix write failed for row {index}: {result.error}" + ) + return _WriterResult( + rows=len(row_range), + record_build_seconds=build_seconds, + stage_seconds=stage_seconds, + put_calls=put_calls, + row_wait_seconds=row_wait_seconds, + ) + + +def _write_ranges( + sink: TQTokenSink, + row_range: range, + *, + checkpoint_id: str, + prefix_tokens: int, + writers: int, + batch_size: int = 1, +) -> list[_WriterResult]: + """Stage one client's assigned rows with optional local concurrency.""" + token_ids, token_mask, logprobs = _record_payload(prefix_tokens) + ranges = _split_range(row_range, writers) + with ThreadPoolExecutor(max_workers=len(ranges)) as executor: + futures = [ + executor.submit( + _write_range, + sink, + writer_range, + checkpoint_id=checkpoint_id, + token_ids=token_ids, + token_mask=token_mask, + logprobs=logprobs, + batch_size=batch_size, + ) + for writer_range in ranges + ] + return [future.result() for future in futures] + + +class _AttachedTQClientWriter: + """One production-shaped process-local client attached to shared TQ.""" + + def __init__( + self, + dp_config: DataPlaneConfig, + partition_id: str, + checkpoint_id: str, + ) -> None: + self._client = build_data_plane_client(dp_config, bootstrap=False) + self._sink = TQTokenSink(self._client, staging_partition=partition_id) + self._checkpoint_id = checkpoint_id + + def ready(self) -> bool: + """Confirm that this process has attached its TQ client.""" + return True + + def write( + self, + start: int, + stop: int, + prefix_tokens: int, + writers: int, + batch_size: int, + ) -> list[dict[str, Any]]: + """Write this client's assigned rows and return per-writer timings.""" + results = _write_ranges( + self._sink, + range(start, stop), + checkpoint_id=self._checkpoint_id, + prefix_tokens=prefix_tokens, + writers=writers, + batch_size=batch_size, + ) + return [asdict(result) for result in results] + + +_AttachedTQClientWriterActor = ray.remote(num_cpus=0)(_AttachedTQClientWriter) + + +def _wait_for_actor_results( + references: list[Any], + *, + stage: str, + timeout_s: float, + progress_interval_s: float, +) -> list[Any]: + """Collect Ray results with bounded, periodic progress reporting.""" + pending = list(references) + completed = 0 + results: list[Any] = [] + started = time.monotonic() + print( + f"tq_prefix_benchmark stage={stage} started total={len(pending)}", + file=sys.stderr, + flush=True, + ) + while pending: + elapsed = time.monotonic() - started + remaining = timeout_s - elapsed + if remaining <= 0: + raise TimeoutError( + f"Timed out during {stage}: completed={completed}, " + f"pending={len(pending)}, timeout_s={timeout_s}" + ) + ready, pending = ray.wait( + pending, + num_returns=len(pending), + timeout=min(progress_interval_s, remaining), + ) + if ready: + results.extend(ray.get(ready)) + completed += len(ready) + print( + f"tq_prefix_benchmark stage={stage} completed={completed} " + f"pending={len(pending)} elapsed_s={time.monotonic() - started:.1f}", + file=sys.stderr, + flush=True, + ) + return results + + +def _run_single_client_writes( + dp_client: DataPlaneClient, + config: BenchmarkConfig, +) -> tuple[list[_WriterResult], float, int, int]: + """Run the original one-client benchmark path.""" + sink = TQTokenSink(dp_client, staging_partition=config.partition_id) + started = time.perf_counter() + writer_results = _write_ranges( + sink, + range(config.rows), + checkpoint_id=config.checkpoint_id, + prefix_tokens=config.prefix_tokens, + writers=config.writers_per_client, + batch_size=config.batch_size, + ) + return ( + writer_results, + time.perf_counter() - started, + 1, + len(writer_results), + ) + + +def _run_multi_client_writes( + dp_config: DataPlaneConfig, + config: BenchmarkConfig, +) -> tuple[list[_WriterResult], float, float, int, int]: + """Write through process-isolated clients attached to one TQ deployment.""" + client_ranges = _row_ranges(config.rows, config.clients) + setup_started = time.perf_counter() + actors = [ + _AttachedTQClientWriterActor.remote( + dp_config, + config.partition_id, + config.checkpoint_id, + ) + for _ in client_ranges + ] + try: + # Constructors finish before timing begins, just as production TQ + # clients are already attached when checkpoint prepare starts. + _wait_for_actor_results( + [actor.ready.remote() for actor in actors], + stage="client_setup", + timeout_s=config.client_setup_timeout_s, + progress_interval_s=config.progress_interval_s, + ) + client_setup_seconds = time.perf_counter() - setup_started + started = time.perf_counter() + payloads = _wait_for_actor_results( + [ + actor.write.remote( + client_range.start, + client_range.stop, + config.prefix_tokens, + config.writers_per_client, + config.batch_size, + ) + for actor, client_range in zip(actors, client_ranges) + ], + stage="prefix_writes", + timeout_s=config.write_timeout_s, + progress_interval_s=config.progress_interval_s, + ) + write_seconds = time.perf_counter() - started + finally: + # Do not call DataPlaneClient.close() in attached actors: TQ's close + # operation has deployment-level semantics and can contend with the + # driver-owned controller/storage that must remain live for verification. + # Terminating the client processes drops their local attachments without + # asking them to close the shared deployment. + print( + f"tq_prefix_benchmark stage=client_teardown started total={len(actors)}", + file=sys.stderr, + flush=True, + ) + for actor in actors: + try: + ray.kill(actor, no_restart=True) + except Exception: + pass + print( + f"tq_prefix_benchmark stage=client_teardown completed={len(actors)}", + file=sys.stderr, + flush=True, + ) + + writer_results = [ + _WriterResult(**payload) + for client_payloads in payloads + for payload in client_payloads + ] + return ( + writer_results, + write_seconds, + client_setup_seconds, + len(actors), + len(writer_results), + ) + + +def _benchmark_keys(config: BenchmarkConfig, row_indices: list[int]) -> list[str]: + """Return deterministic TQ keys for selected benchmark rows.""" + return [ + generation_cut_staging_key( + config.checkpoint_id, + f"benchmark-rollout-{row_index:09d}", + "call-0", + chunk_sequence=0, + ) + for row_index in row_indices + ] + + +def _verification_indices(rows: int, verify_rows: int) -> list[int]: + """Select evenly spread rows, including both ends when possible.""" + count = min(rows, verify_rows) + if count == 0: + return [] + if count == 1: + return [0] + return [round(index * (rows - 1) / (count - 1)) for index in range(count)] + + +def _verify_prefix_tokens( + source: TQTokenSource, + config: BenchmarkConfig, + expected_token_ids: list[int], +) -> tuple[int, float]: + """Fetch a bounded sample of prefixes and verify token contents.""" + indices = _verification_indices(config.rows, config.verify_rows) + keys = _benchmark_keys(config, indices) + started = time.perf_counter() + for offset in range(0, len(keys), config.verify_batch_size): + batch_keys = keys[offset : offset + config.verify_batch_size] + actual = source.fetch_prefix_token_ids(batch_keys) + expected = expected_token_ids * len(batch_keys) + if actual != expected: + raise AssertionError( + "TQ prefix token verification failed for keys " + f"{batch_keys[0]!r} through {batch_keys[-1]!r}" + ) + return len(keys), time.perf_counter() - started + + +def run_prefix_write_benchmark( + dp_client: DataPlaneClient, + config: BenchmarkConfig, + *, + dp_config: DataPlaneConfig | None = None, +) -> dict[str, Any]: + """Run the benchmark against an initialized data-plane client.""" + _validate_config(config) + token_ids, token_mask, logprobs = _record_payload(config.prefix_tokens) + dp_client.register_partition( + partition_id=config.partition_id, + fields=list(STAGING_FIELDS), + num_samples=config.rows, + consumer_tasks=[CONSUMER_TASK], + ) + source = TQTokenSource(dp_client, staging_partition=config.partition_id) + + cleanup_seconds: float | None = None + try: + client_setup_seconds = 0.0 + if config.clients == 1: + ( + writer_results, + write_seconds, + clients_used, + writers_used, + ) = _run_single_client_writes(dp_client, config) + else: + if dp_config is None: + raise ValueError( + "dp_config is required when clients is greater than one" + ) + ( + writer_results, + write_seconds, + client_setup_seconds, + clients_used, + writers_used, + ) = _run_multi_client_writes(dp_config, config) + + print( + "tq_prefix_benchmark stage=list_stored_keys started", + file=sys.stderr, + flush=True, + ) + started = time.perf_counter() + stored_keys = dp_client.list_sample_ids(config.partition_id) + list_seconds = time.perf_counter() - started + print( + f"tq_prefix_benchmark stage=list_stored_keys completed " + f"rows={len(stored_keys)} elapsed_s={list_seconds:.1f}", + file=sys.stderr, + flush=True, + ) + if len(stored_keys) != config.rows: + raise AssertionError( + f"Expected {config.rows} stored prefixes, found {len(stored_keys)}" + ) + + print( + "tq_prefix_benchmark stage=verify_prefixes started", + file=sys.stderr, + flush=True, + ) + verified_rows, verify_seconds = _verify_prefix_tokens( + source, + config, + token_ids, + ) + print( + f"tq_prefix_benchmark stage=verify_prefixes completed " + f"rows={verified_rows} elapsed_s={verify_seconds:.1f}", + file=sys.stderr, + flush=True, + ) + result: dict[str, Any] = { + "backend": "simple", + "rows": config.rows, + "prefix_tokens": config.prefix_tokens, + "clients_requested": config.clients, + "clients_used": clients_used, + "writers_per_client_requested": config.writers_per_client, + "writers_requested_total": config.clients * config.writers_per_client, + "writers_used": writers_used, + "client_setup_seconds": client_setup_seconds, + "num_storage_units": config.num_storage_units, + "storage_capacity": config.storage_capacity or max(1_024, config.rows), + "batch_size": config.batch_size, + "put_calls": sum(item.put_calls for item in writer_results), + "stage_batch_seconds_mean": sum( + item.stage_seconds for item in writer_results + ) + / sum(item.put_calls for item in writer_results), + "write_seconds": write_seconds, + "rows_per_second": config.rows / write_seconds, + "tokens_per_second": config.rows * config.prefix_tokens / write_seconds, + "record_build_seconds_sum": sum( + item.record_build_seconds for item in writer_results + ), + "stage_seconds_sum": sum(item.stage_seconds for item in writer_results), + "stage_seconds_mean": sum(item.row_wait_seconds for item in writer_results) + / config.rows, + "stage_seconds_per_row_amortized": sum( + item.stage_seconds for item in writer_results + ) + / config.rows, + "stored_keys": len(stored_keys), + "list_seconds": list_seconds, + "verified_rows": verified_rows, + "verify_seconds": verify_seconds, + "max_rss_kib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, + "rows_kept": config.keep_rows, + } + finally: + if not config.keep_rows: + print( + "tq_prefix_benchmark stage=cleanup started", + file=sys.stderr, + flush=True, + ) + started = time.perf_counter() + dp_client.clear_samples( + sample_ids=None, + partition_id=config.partition_id, + ) + cleanup_seconds = time.perf_counter() - started + print( + f"tq_prefix_benchmark stage=cleanup completed " + f"elapsed_s={cleanup_seconds:.1f}", + file=sys.stderr, + flush=True, + ) + result["cleanup_seconds"] = cleanup_seconds + return result + + +def _parse_args() -> tuple[BenchmarkConfig, Path | None]: + """Parse CLI arguments into a typed benchmark configuration.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, default=10_000) + parser.add_argument("--prefix-tokens", type=int, default=128) + parser.add_argument( + "--clients", + type=int, + default=1, + help=( + "Process-isolated TQ clients attached to one shared deployment; " + "use more than one to model separate generation model owners." + ), + ) + parser.add_argument( + "--writers-per-client", + "--writers", + dest="writers_per_client", + type=int, + default=1, + help=( + "Concurrent writer threads inside each client. --writers is kept " + "as a backward-compatible alias." + ), + ) + parser.add_argument( + "--batch-size", + type=int, + default=1, + help="Prefix rows per PUT in each writer; 1 preserves the baseline.", + ) + parser.add_argument("--num-storage-units", type=int, default=8) + parser.add_argument( + "--storage-capacity", + type=int, + default=None, + help="Maximum retained rows; defaults to max(1024, --rows).", + ) + parser.add_argument("--verify-rows", type=int, default=16) + parser.add_argument("--verify-batch-size", type=int, default=16) + parser.add_argument("--partition-id", default=DEFAULT_PARTITION_ID) + parser.add_argument("--checkpoint-id", default=DEFAULT_CHECKPOINT_ID) + parser.add_argument("--client-setup-timeout-s", type=float, default=180.0) + parser.add_argument("--write-timeout-s", type=float, default=1_800.0) + parser.add_argument("--progress-interval-s", type=float, default=5.0) + parser.add_argument( + "--keep-rows", + action="store_true", + help="Do not clear benchmark rows before closing the TQ client.", + ) + parser.add_argument( + "--output-json", + type=Path, + default=None, + help="Also write the final JSON result to this path.", + ) + args = parser.parse_args() + config = BenchmarkConfig( + rows=args.rows, + prefix_tokens=args.prefix_tokens, + clients=args.clients, + writers_per_client=args.writers_per_client, + batch_size=args.batch_size, + num_storage_units=args.num_storage_units, + storage_capacity=args.storage_capacity, + verify_rows=args.verify_rows, + verify_batch_size=args.verify_batch_size, + partition_id=args.partition_id, + checkpoint_id=args.checkpoint_id, + keep_rows=args.keep_rows, + client_setup_timeout_s=args.client_setup_timeout_s, + write_timeout_s=args.write_timeout_s, + progress_interval_s=args.progress_interval_s, + ) + try: + _validate_config(config) + except ValueError as error: + parser.error(str(error)) + return config, args.output_json + + +def main() -> None: + """Run the CLI benchmark and print one machine-readable result.""" + config, output_json = _parse_args() + dp_config = _data_plane_config(config) + dp_client = build_data_plane_client( + dp_config, + bootstrap=True, + ) + try: + result = run_prefix_write_benchmark( + dp_client, + config, + dp_config=dp_config, + ) + finally: + dp_client.close() + + payload = {"config": asdict(config), "result": result} + rendered = json.dumps(payload, indent=2, sort_keys=True) + print(rendered) + if output_json is not None: + output_json.parent.mkdir(parents=True, exist_ok=True) + output_json.write_text(f"{rendered}\n", encoding="utf-8") + + +if __name__ == "__main__": + main()