diff --git a/nemo_rl/data_plane/tq_token_sink.py b/nemo_rl/data_plane/tq_token_sink.py index ed8568a271..b62102c320 100644 --- a/nemo_rl/data_plane/tq_token_sink.py +++ b/nemo_rl/data_plane/tq_token_sink.py @@ -36,7 +36,8 @@ import json import logging -from dataclasses import dataclass +import time +from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any import ray @@ -130,6 +131,16 @@ class FetchedStagedCall: fragment: RouteFragment | None = None +@dataclass +class StagedGroupFetch: + """Per-key results and time spent waiting on data-plane reads only.""" + + calls: dict[str, FetchedStagedCall | KeyError | TypeError | ValueError] = field( + default_factory=dict + ) + fetch_ms: float = 0.0 + + def _call_dp(dp_client: Any, method_name: str, **kwargs: Any) -> Any: """Call a DataPlaneClient method on a local client or a Ray actor handle.""" method = getattr(dp_client, method_name) @@ -319,15 +330,10 @@ def clear(self, staging_keys: list[str]) -> None: class TQTokenSource: """Gym ``StagingSource`` over ``DataPlaneClient.get_samples``. - All requested rows are fetched in a single batched ``get_samples`` call - (TQ returns jagged delta columns as nested tensors; ``_from_wire`` - preserves the raggedness), in the order requested. A missing or - unreadable row raises ``KeyError`` per the protocol — the finalizer maps - that to a placeholder, never a silent skip. TQ's field-readiness check - is all-or-nothing across a batch, so the extras fallback is batch-level: - extras-free runs land in the base schema exactly like the old per-key - probe, but a batch with *mixed* extras presence degrades every row to - the base schema (worker feature-gating makes presence uniform per run). + Healthy requests use one base-column read and, in direct mode, one + payload read restricted to rows whose encoding declares routes. The Gym + interface raises on missing or unreadable rows. Group finalization uses + ``fetch_for_group`` to isolate failures by staging key. """ def __init__(self, dp_client: Any, *, staging_partition: str) -> None: @@ -377,45 +383,47 @@ def fetch_for_finalization( ) -> list[FetchedStagedCall]: """Fetch digest-covered base columns, plus route payloads when requested. - Deferred mode (the default) never selects ``routed_experts`` — route - bytes stay in TQ for the policy worker. Direct mode passes - ``include_route_fragments=True`` to pull the payloads in the same - batched read and receives them as ``RouteFragment`` values beside the - base snapshots, never inside them. + Deferred mode never selects ``routed_experts``. Direct mode first + reads the base columns, then selects payloads only for keys whose + encoding declares routes, so mixed route presence is supported. """ + return self._fetch_for_finalization( + staging_keys, include_route_fragments=include_route_fragments + ) + + def _read_finalization_rows( + self, + staging_keys: list[str], + select_fields: list[str], + timings: StagedGroupFetch | None, + ) -> TensorDict: + started = time.perf_counter() + try: + rows = _call_dp( + self._dp_client, + "get_samples", + sample_ids=staging_keys, + partition_id=self._staging_partition, + select_fields=select_fields, + ) + finally: + if timings is not None: + timings.fetch_ms += (time.perf_counter() - started) * 1000.0 + return rows + + def _fetch_for_finalization( + self, + staging_keys: list[str], + *, + include_route_fragments: bool, + timings: StagedGroupFetch | None = None, + ) -> list[FetchedStagedCall]: if not staging_keys: return [] if len(set(staging_keys)) != len(staging_keys): raise KeyError("finalization staging request contains duplicate keys") try: - if include_route_fragments: - # Route payloads are optional per run (feature-gated at the - # worker); fall back to the base schema so extras-free rows - # keep fetching. - try: - rows = _call_dp( - self._dp_client, - "get_samples", - sample_ids=list(staging_keys), - partition_id=self._staging_partition, - select_fields=STAGING_FIELDS + [ROUTED_EXPERTS_FIELD], - ) - except Exception: # noqa: BLE001 — field-not-present probe - rows = _call_dp( - self._dp_client, - "get_samples", - sample_ids=list(staging_keys), - partition_id=self._staging_partition, - select_fields=STAGING_FIELDS, - ) - else: - rows = _call_dp( - self._dp_client, - "get_samples", - sample_ids=list(staging_keys), - partition_id=self._staging_partition, - select_fields=STAGING_FIELDS, - ) + rows = self._read_finalization_rows(staging_keys, STAGING_FIELDS, timings) except Exception as error: # noqa: BLE001 — protocol maps misses to KeyError raise KeyError( f"staged rows for {len(staging_keys)} keys could not be " @@ -446,13 +454,86 @@ def fetch_for_finalization( staging_key=key, snapshot=snapshot, routed_len=_row_scalar_int(row, ROUTED_LEN_FIELD), - fragment=( - _row_to_route_fragment(row) if include_route_fragments else None - ), ) ) + if include_route_fragments: + route_indices = [ + index + for index in range(n_rows) + if _row_scalar_int( + _select_row(rows, index), ROUTED_EXPERTS_ENCODING_FIELD + ) + != ROUTE_ENCODING_NONE + ] + if route_indices: + route_keys = [staging_keys[index] for index in route_indices] + try: + route_rows = self._read_finalization_rows( + route_keys, + ["rollout_id_utf8", "model_call_id_utf8", ROUTED_EXPERTS_FIELD], + timings, + ) + except Exception as error: # noqa: BLE001 — Gym source maps misses to KeyError + raise KeyError( + f"staged route payloads could not be fetched: {error}" + ) from error + if not route_rows.batch_size or int(route_rows.batch_size[0]) != len( + route_keys + ): + raise KeyError("staged route payloads missing") + for route_index, base_index in enumerate(route_indices): + route_row = _select_row(route_rows, route_index) + base_row = _select_row(rows, base_index) + if any( + _row_text(route_row, name) != _row_text(base_row, name) + for name in ("rollout_id_utf8", "model_call_id_utf8") + ): + raise KeyError("staged route payload identity mismatch") + base_row[ROUTED_EXPERTS_FIELD] = route_row[ROUTED_EXPERTS_FIELD] + fetched[base_index] = replace( + fetched[base_index], fragment=_row_to_route_fragment(base_row) + ) return fetched + def fetch_for_group( + self, + staging_keys: list[str], + *, + include_route_fragments: bool = False, + ) -> StagedGroupFetch: + """Fetch a key union while isolating unreadable rows from siblings. + + TQ omits missing keys and rejects an entire read for unready fields. + Its TensorDict does not carry the resolved key list. Strict identity + and cardinality checks therefore precede mapping results to keys; + failed batches are bisected until each failure has one owner. Extra + reads occur only on the failure path, with at most 2*N-1 batch + attempts (up to two data-plane reads per attempt in direct mode). + """ + result = StagedGroupFetch() + + def fetch_batch(keys: list[str]) -> None: + if not keys: + return + try: + fetched = self._fetch_for_finalization( + keys, + include_route_fragments=include_route_fragments, + timings=result, + ) + except (KeyError, TypeError, ValueError) as error: + if len(keys) == 1: + result.calls[keys[0]] = error + else: + middle = len(keys) // 2 + fetch_batch(keys[:middle]) + fetch_batch(keys[middle:]) + else: + result.calls.update((item.staging_key, item) for item in fetched) + + fetch_batch(list(dict.fromkeys(staging_keys))) + return result + def _select_row(rows: TensorDict, index: int) -> dict[str, torch.Tensor]: """Slice one row out of a batched fetch, restoring single-row shapes. @@ -463,13 +544,13 @@ def _select_row(rows: TensorDict, index: int) -> dict[str, torch.Tensor]: dense component, which is exactly the jagged-row payload. """ row: dict[str, torch.Tensor] = {} - for field in rows.keys(): - value = rows.get(field) + for column in rows.keys(): + value = rows.get(column) if not isinstance(value, torch.Tensor): raise TypeError( - f"staging field {field!r} must be a tensor, got {type(value).__name__}" + f"staging field {column!r} must be a tensor, got {type(value).__name__}" ) - row[str(field)] = value[index].unsqueeze(0) + row[str(column)] = value[index].unsqueeze(0) return row diff --git a/nemo_rl/experience/rollout_reassembler.py b/nemo_rl/experience/rollout_reassembler.py index 06b32e7bc9..7fd4a38723 100644 --- a/nemo_rl/experience/rollout_reassembler.py +++ b/nemo_rl/experience/rollout_reassembler.py @@ -14,9 +14,9 @@ """Blackbox finalization: token-free receipts + staged deltas -> canonical rows. Orchestration only: -per rollout, apply the rollout-level receipt guards, fetch the staged base -rows the receipt manifest names through the ``TokenSource`` (normally -validated ``StagedCallBaseSnapshot`` values), and delegate all token, digest, +apply per-rollout receipt guards, fetch the group's union of staged keys +through the ``TokenSource`` (normally validated ``StagedCallBaseSnapshot`` +values), and delegate all token, digest, lineage, and terminal-chain semantics to Gym's ``verify_and_linearize``. Any rejection becomes a masked placeholder row — the group always publishes exactly N rows so GRPO group shape survives; validity folds into @@ -28,22 +28,27 @@ ``RouteAssemblyPlan`` from Gym's link spans and extras commitments. Deferred mode publishes the encoded plan beside the canonical row and leaves staged route fragments live until policy consumption; direct mode executes the plan -eagerly with fragments fetched in the same batch — any executor failure is a -pre-publication ``route_assembly:`` rejection. +eagerly with fragments fetched in one targeted group read — any executor +failure is a pre-publication ``route_assembly:`` rejection. """ from __future__ import annotations import time from collections import Counter +from collections.abc import Sequence from dataclasses import dataclass, field -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional import torch from nemo_rl.data_plane import KVBatchMeta from nemo_rl.data_plane.schema import MASK_SAMPLE, ROUTE_PLAN_TAG, TRUNCATED -from nemo_rl.data_plane.tq_token_sink import TQTokenSink, TQTokenSource +from nemo_rl.data_plane.tq_token_sink import ( + FetchedStagedCall, + TQTokenSink, + TQTokenSource, +) from nemo_rl.experience.payload import pack_payload from nemo_rl.experience.route_assembly import ( ROUTE_MISSING_SENTINEL, @@ -59,6 +64,9 @@ validate_route_plan, ) +if TYPE_CHECKING: + from nemo_gym.token_id_capture.staging.records import RolloutReceipt + @dataclass(frozen=True) class FinalizedRollout: @@ -149,75 +157,139 @@ def finalize_rollout( invalid row whose reason feeds the metrics; the group publisher substitutes a placeholder. """ - # Deferred: nemo_gym is an optional extra absent in non-gym runs. - from nemo_gym.token_id_capture.staging.rebuild import ( - RebuildError, - ReceiptVerificationError, - verify_and_linearize, + parsed = self._prepare_rollout(rollout_id, receipt, reward=reward) + if isinstance(parsed, FinalizedRollout): + return parsed + staging_keys = [record.staging_key for record in parsed.manifest] + try: + fetched = self._source.fetch_for_finalization( + staging_keys, + include_route_fragments=( + self._router_replay_enabled + and not self._defer_routed_experts_to_policy + ), + ) + except KeyError as error: + return self._rejected( + rollout_id, reward, f"missing_staging_row:{error}", staging_keys + ) + except (TypeError, ValueError) as error: + return self._rejected( + rollout_id, reward, f"invalid_staging_row:{error}", staging_keys + ) + return self._verify_rollout(rollout_id, parsed, fetched, reward=reward) + + @staticmethod + def _rejected( + rollout_id: str, reward: float, reason: str, staging_keys: list[str] + ) -> FinalizedRollout: + return FinalizedRollout( + rollout_id=rollout_id, + valid=False, + rejection_reason=reason, + token_ids=[], + token_mask=[], + logprobs=[], + prompt_len=0, + reward=reward, + staging_keys=staging_keys, ) - from nemo_gym.token_id_capture.staging.records import RolloutReceipt - def rejected(reason: str, staging_keys: list[str]) -> FinalizedRollout: - return FinalizedRollout( - rollout_id=rollout_id, - valid=False, - rejection_reason=reason, - token_ids=[], - token_mask=[], - logprobs=[], - prompt_len=0, - reward=reward, - staging_keys=staging_keys, - ) + def _prepare_rollout( + self, rollout_id: str, receipt: Optional[dict[str, Any]], *, reward: float + ) -> RolloutReceipt | FinalizedRollout: + # Deferred: nemo_gym is an optional extra absent in non-gym runs. + from nemo_gym.token_id_capture.staging.records import RolloutReceipt if receipt is None: - return rejected("missing_receipt", []) + return self._rejected(rollout_id, reward, "missing_receipt", []) try: parsed = RolloutReceipt.model_validate(receipt) except ValueError as error: - return rejected(f"invalid_receipt:{error}", []) + return self._rejected(rollout_id, reward, f"invalid_receipt:{error}", []) staging_keys = [record.staging_key for record in parsed.manifest] if parsed.rollout_id != rollout_id: - return rejected(f"identity_mismatch:{parsed.rollout_id}", staging_keys) + return self._rejected( + rollout_id, + reward, + f"identity_mismatch:{parsed.rollout_id}", + staging_keys, + ) if parsed.failure_reason is not None: - return rejected(f"rollout_failed:{parsed.failure_reason}", staging_keys) + return self._rejected( + rollout_id, + reward, + f"rollout_failed:{parsed.failure_reason}", + staging_keys, + ) if parsed.capture_poisoned: - return rejected("capture_poisoned", staging_keys) + return self._rejected(rollout_id, reward, "capture_poisoned", staging_keys) if not parsed.manifest: - return rejected("empty_manifest", staging_keys) + return self._rejected(rollout_id, reward, "empty_manifest", staging_keys) if len(set(staging_keys)) != len(staging_keys): - return rejected( + return self._rejected( + rollout_id, + reward, "duplicate_staging_key", list(dict.fromkeys(staging_keys)), ) records_by_call = {record.model_call_id: record for record in parsed.manifest} if len(records_by_call) != len(parsed.manifest): - return rejected("duplicate_manifest_call_id", staging_keys) + return self._rejected( + rollout_id, reward, "duplicate_manifest_call_id", staging_keys + ) + + return parsed - fetch_fragments = ( - self._router_replay_enabled and not self._defer_routed_experts_to_policy + def _verify_rollout( + self, + rollout_id: str, + parsed: RolloutReceipt, + fetch_results: Sequence[FetchedStagedCall | KeyError | TypeError | ValueError], + *, + reward: float, + ) -> FinalizedRollout: + # Deferred: nemo_gym is an optional extra absent in non-gym runs. + from nemo_gym.token_id_capture.staging.rebuild import ( + RebuildError, + ReceiptVerificationError, + verify_and_linearize, ) - try: - fetched = self._source.fetch_for_finalization( - staging_keys, include_route_fragments=fetch_fragments - ) - except KeyError as error: - return rejected(f"missing_staging_row:{error}", staging_keys) - except (TypeError, ValueError) as error: - return rejected(f"invalid_staging_row:{error}", staging_keys) + + staging_keys = [record.staging_key for record in parsed.manifest] + records_by_call = {record.model_call_id: record for record in parsed.manifest} + fetched: list[FetchedStagedCall] = [] + for item in fetch_results: + if isinstance(item, KeyError): + return self._rejected( + rollout_id, reward, f"missing_staging_row:{item}", staging_keys + ) + if isinstance(item, (TypeError, ValueError)): + return self._rejected( + rollout_id, reward, f"invalid_staging_row:{item}", staging_keys + ) + fetched.append(item) fetched_by_call = {} for record, item in zip(parsed.manifest, fetched): if item.staging_key != record.staging_key: - return rejected( - f"staging_key_mismatch:{record.model_call_id}", staging_keys + return self._rejected( + rollout_id, + reward, + f"staging_key_mismatch:{record.model_call_id}", + staging_keys, ) if item.snapshot.model_call_id != record.model_call_id: - return rejected( - f"call_id_mismatch:{record.model_call_id}", staging_keys + return self._rejected( + rollout_id, + reward, + f"call_id_mismatch:{record.model_call_id}", + staging_keys, ) fetched_by_call[record.model_call_id] = item if len(fetched_by_call) != len(fetched): - return rejected("duplicate_fetched_call_id", staging_keys) + return self._rejected( + rollout_id, reward, "duplicate_fetched_call_id", staging_keys + ) # All base token/digest/lineage/terminal semantics belong to Gym; the # finalizer never re-verifies them. @@ -231,7 +303,9 @@ def rejected(reason: str, staging_keys: list[str]) -> FinalizedRollout: RebuildError, NotImplementedError, ) as error: - return rejected(f"rebuild_failed:{error}", staging_keys) + return self._rejected( + rollout_id, reward, f"rebuild_failed:{error}", staging_keys + ) weight_versions = [record.weight_version for record in parsed.manifest] min_wv, max_wv = min(weight_versions), max(weight_versions) @@ -250,22 +324,43 @@ def rejected(reason: str, staging_keys: list[str]) -> FinalizedRollout: seen_span_call_ids: set[str] = set() for call_id, carry_len, generation_len in row.link_spans: if call_id in seen_span_call_ids: - return rejected(f"duplicate_route_span:{call_id}", staging_keys) + return self._rejected( + rollout_id, + reward, + f"duplicate_route_span:{call_id}", + staging_keys, + ) seen_span_call_ids.add(call_id) record = records_by_call.get(call_id) item = fetched_by_call.get(call_id) commitment = commitments_by_call.get(call_id) if record is None or item is None or commitment is None: - return rejected(f"route_span_identity:{call_id}", staging_keys) + return self._rejected( + rollout_id, + reward, + f"route_span_identity:{call_id}", + staging_keys, + ) if item.routed_len not in (0, record.delta_len): - return rejected(f"routed_len_mismatch:{call_id}", staging_keys) + return self._rejected( + rollout_id, + reward, + f"routed_len_mismatch:{call_id}", + staging_keys, + ) if generation_len < 0 or generation_len > record.delta_len: - return rejected( - f"route_generation_span_mismatch:{call_id}", staging_keys + return self._rejected( + rollout_id, + reward, + f"route_generation_span_mismatch:{call_id}", + staging_keys, ) if carry_len < 0: - return rejected( - f"route_carry_span_mismatch:{call_id}", staging_keys + return self._rejected( + rollout_id, + reward, + f"route_carry_span_mismatch:{call_id}", + staging_keys, ) route_spans.append( RouteSpan( @@ -280,7 +375,9 @@ def rejected(reason: str, staging_keys: list[str]) -> FinalizedRollout: if sum(span.carry_len + span.generation_len for span in route_spans) != len( row.token_ids ): - return rejected("route_span_length_mismatch", staging_keys) + return self._rejected( + rollout_id, reward, "route_span_length_mismatch", staging_keys + ) plan = RouteAssemblyPlan( schema_version=ROUTE_PLAN_SCHEMA_VERSION, staging_partition=self._staging_partition, @@ -291,7 +388,9 @@ def rejected(reason: str, staging_keys: list[str]) -> FinalizedRollout: try: validate_route_plan(plan) except (TypeError, ValueError) as error: - return rejected(f"invalid_route_plan:{error}", staging_keys) + return self._rejected( + rollout_id, reward, f"invalid_route_plan:{error}", staging_keys + ) # Both modes carry the constructed plan on the rollout; only # deferred mode publishes it (direct mode executes it eagerly and # the published row carries the assembled tensor instead). @@ -299,7 +398,9 @@ def rejected(reason: str, staging_keys: list[str]) -> FinalizedRollout: if not self._defer_routed_experts_to_policy: routed_experts, failure = self._execute_direct_plan(plan, fetched) if failure is not None: - return rejected(f"route_assembly:{failure}", staging_keys) + return self._rejected( + rollout_id, reward, f"route_assembly:{failure}", staging_keys + ) return FinalizedRollout( rollout_id=rollout_id, @@ -392,14 +493,49 @@ def finalize_group( "canonical_sample_ids must be one per rollout" ) _group_t0 = time.perf_counter() - rows = [ - self.finalize_rollout(rollout_id, receipt, reward=reward) + prepared = [ + self._prepare_rollout(rollout_id, receipt, reward=reward) for rollout_id, receipt, reward in zip(rollout_ids, receipts, rewards) ] + staging_union = list( + dict.fromkeys( + record.staging_key + for parsed in prepared + if not isinstance(parsed, FinalizedRollout) + for record in parsed.manifest + ) + ) + # fetch_ms measures only data-plane reads, including failed attempts. + # verify_ms covers CPU work: receipt guards, wire decoding, Gym + # verification and direct route execution. Their sum is rollouts_ms. + fetched_by_key = self._source.fetch_for_group( + staging_union, + include_route_fragments=( + self._router_replay_enabled and not self._defer_routed_experts_to_policy + ), + ) + _fetch_ms = fetched_by_key.fetch_ms + rows = [ + parsed + if isinstance(parsed, FinalizedRollout) + else self._verify_rollout( + rollout_id, + parsed, + [ + fetched_by_key.calls[record.staging_key] + for record in parsed.manifest + ], + reward=reward, + ) + for rollout_id, parsed, reward in zip(rollout_ids, prepared, rewards) + ] _rollouts_ms = (time.perf_counter() - _group_t0) * 1000.0 valid_rows = [row for row in rows if row.valid] staging_keys = [key for row in rows for key in row.staging_keys] metrics = { + "row_assembly/fetch_ms": _fetch_ms, + "row_assembly/verify_ms": _rollouts_ms - _fetch_ms, + "row_assembly/rollouts_ms": _rollouts_ms, "finalize/invalid_row_rate": 1.0 - len(valid_rows) / len(rows), "finalize/calls_per_rollout": ( sum(len(row.staging_keys) for row in rows) / len(rows) @@ -621,7 +757,6 @@ def finalize_group( _clear_ms = (time.perf_counter() - _clear_t0) * 1000.0 # Per-step W&B breakdown of training-row assembly (capture arm) rides # FinalizedGroup.metrics into the controller's rollout metrics. - metrics["row_assembly/rollouts_ms"] = _rollouts_ms metrics["row_assembly/tensorize_ms"] = _tensorize_ms metrics["row_assembly/tq_put_ms"] = _put_ms if not self._defer_routed_experts_to_policy: diff --git a/nemo_rl/experience/rollout_reassembler_actor.py b/nemo_rl/experience/rollout_reassembler_actor.py index d22923f5ca..cbe967e291 100644 --- a/nemo_rl/experience/rollout_reassembler_actor.py +++ b/nemo_rl/experience/rollout_reassembler_actor.py @@ -20,6 +20,7 @@ import ray import torch +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from nemo_rl.data_plane import DataPlaneConfig, build_data_plane_client from nemo_rl.experience.rollout_reassembler import FinalizedGroup, RolloutReassembler @@ -172,6 +173,14 @@ def create_rollout_reassembler_actors( """Construct the fixed validation pool after TQ partitions are registered.""" if num_workers <= 0: raise ValueError(f"num_reassembler_workers must be positive, got {num_workers}") - return [ - RolloutReassemblerActor.remote(dp_config, config) for _ in range(num_workers) - ] + # Prefer the controller node to avoid competing with policy workers, but + # permit spillover when its CPUs are busy. soft=True alone only spills + # when the node is dead or infeasible, not when its resources are occupied. + actor = RolloutReassemblerActor.options( + scheduling_strategy=NodeAffinitySchedulingStrategy( + node_id=ray.get_runtime_context().get_node_id(), + soft=True, + _spill_on_unavailable=True, + ) + ) + return [actor.remote(dp_config, config) for _ in range(num_workers)] diff --git a/tests/unit/data_plane/test_rollout_reassembler.py b/tests/unit/data_plane/test_rollout_reassembler.py index e3cc8de2d5..e9f4162206 100644 --- a/tests/unit/data_plane/test_rollout_reassembler.py +++ b/tests/unit/data_plane/test_rollout_reassembler.py @@ -824,3 +824,175 @@ def test_deferred_chain_hash_corruption_rejects_the_row( ) assert not row.valid assert (row.rejection_reason or "").startswith("rebuild_failed:chain_hash_mismatch") + + +@pytest.mark.parametrize("bad_second", [False, True]) +def test_group_batches_reads_and_preserves_rollout_failure_isolation( + tq_client, partitions, monkeypatch, bad_second +): + group_id = "batched" + ids = [f"{group_id}_g{i}" for i in range(3)] + receipts = [ + _stage_fixture(tq_client, "single_call", rollout_id=rollout_id)[0] + for rollout_id in ids[:2] + ] + if bad_second: + tq_client.clear_samples( + sample_ids=[receipts[1]["manifest"][0]["staging_key"]], + partition_id=STAGING_PARTITION, + ) + calls = [] + read = tq_client.get_samples + + def recording_read(**kwargs): + calls.append(kwargs) + return read(**kwargs) + + monkeypatch.setattr(tq_client, "get_samples", recording_read) + result = _finalizer(tq_client).finalize_group( + group_id, + ids, + receipts + [None], + [1.0, 2.0, 0.0], + mask_sample=[False] * 3, + fallback_weight_version=9, + prompt_idx=17, + ) + assert result.valid_row_count == (1 if bad_second else 2) + assert result.total_row_count == 3 + assert len(calls) == (3 if bad_second else 1) + assert calls[0]["sample_ids"] == [r["manifest"][0]["staging_key"] for r in receipts] + assert all(call["select_fields"] == STAGING_FIELDS for call in calls) + assert result.metrics["row_assembly/fetch_ms"] >= 0 + assert result.metrics["row_assembly/verify_ms"] >= 0 + assert result.metrics["row_assembly/rollouts_ms"] == pytest.approx( + result.metrics["row_assembly/fetch_ms"] + + result.metrics["row_assembly/verify_ms"] + ) + if bad_second: + assert ( + result.metrics["finalize/capture_failure_reason_missing_staging_row_count"] + == 1 + ) + published = read( + sample_ids=ids, partition_id=CANONICAL_PARTITION, select_fields=["sample_mask"] + ) + assert published["sample_mask"].tolist() == [1.0, 0.0 if bad_second else 1.0, 0.0] + + +@pytest.mark.parametrize("router_replay", [False, True]) +def test_placeholder_only_groups_have_zero_fetch_time( + tq_client, partitions, monkeypatch, router_replay +): + def unexpected_read(**kwargs): + raise AssertionError("placeholder groups must not read staging") + + monkeypatch.setattr(tq_client, "get_samples", unexpected_read) + result = _finalizer(tq_client, router_replay_enabled=router_replay).finalize_group( + "empty", + ["empty_g0"], + [None], + [0.0], + mask_sample=[False], + fallback_weight_version=9, + prompt_idx=17, + ) + assert result.metrics["row_assembly/fetch_ms"] == 0 + assert result.metrics["row_assembly/verify_ms"] >= 0 + assert result.dropped == router_replay + + +@pytest.mark.parametrize("damage", [None, "missing", "malformed", "reordered_partial"]) +def test_direct_group_fetch_targets_only_route_keys_and_isolates_payload_failures( + tq_client, r3_deferred_partitions, damage +): + from nemo_rl.data_plane.tq_token_sink import ROUTED_EXPERTS_FIELD, FetchedStagedCall + + records = [ + build_fixture_artifacts("single_call", rollout_id=f"mixed_{i}")[0][0] + for i in range(3) + ] + records[0] = _record_with_routes( + records[0], _routes_for_delta(0, records[0].delta_len) + ) + records[2] = _record_with_routes( + records[2], _routes_for_delta(2, records[2].delta_len) + ) + sink = TQTokenSink(tq_client, staging_partition=_R3_DEFERRED_STAGING) + for record in records: + assert sink.stage(record).ok + keys = [r.staging_key for r in records] + calls = [] + + class RouteClient: + def get_samples(self, **kwargs): + calls.append((list(kwargs["sample_ids"]), list(kwargs["select_fields"]))) + route_read = ROUTED_EXPERTS_FIELD in kwargs["select_fields"] + if route_read and keys[2] in kwargs["sample_ids"]: + if damage == "missing": + raise ValueError("route field is not ready") + if damage == "reordered_partial": + kwargs["sample_ids"] = list(reversed(kwargs["sample_ids"][:-1])) + rows = tq_client.get_samples(**kwargs) + if route_read and damage == "malformed" and keys[2] in kwargs["sample_ids"]: + rows[ROUTED_EXPERTS_FIELD] = "invalid non-tensor payload" + return rows + + result = TQTokenSource( + RouteClient(), staging_partition=_R3_DEFERRED_STAGING + ).fetch_for_group(keys, include_route_fragments=True) + assert isinstance(result.calls[keys[0]], FetchedStagedCall) + assert result.calls[keys[0]].fragment is not None + assert isinstance(result.calls[keys[1]], FetchedStagedCall) + assert result.calls[keys[1]].fragment is None + if damage is None: + assert isinstance(result.calls[keys[2]], FetchedStagedCall) + assert result.calls[keys[2]].fragment is not None + assert len(calls) == 2 + else: + assert isinstance(result.calls[keys[2]], (KeyError, TypeError, ValueError)) + assert calls[0] == (keys, STAGING_FIELDS) + for selected_keys, fields in calls: + if ROUTED_EXPERTS_FIELD in fields: + assert keys[1] not in selected_keys + assert set(fields) == { + "rollout_id_utf8", + "model_call_id_utf8", + ROUTED_EXPERTS_FIELD, + } + + +def test_group_timing_separates_reads_from_verification( + tq_client, partitions, monkeypatch +): + import nemo_rl.experience.rollout_reassembler as module + + receipt, _ = _stage_fixture(tq_client, "single_call", rollout_id="timed_g0") + finalizer = _finalizer(tq_client) + clock = [0.0] + read = tq_client.get_samples + verify = finalizer._verify_rollout + + def timed_read(**kwargs): + clock[0] += 0.010 + return read(**kwargs) + + def timed_verify(*args, **kwargs): + clock[0] += 0.025 + return verify(*args, **kwargs) + + monkeypatch.setattr(module.time, "perf_counter", lambda: clock[0]) + monkeypatch.setattr(tq_client, "get_samples", timed_read) + monkeypatch.setattr(finalizer, "_verify_rollout", timed_verify) + result = finalizer.finalize_group( + "timed", + ["timed_g0"], + [receipt], + [1.0], + mask_sample=[False], + fallback_weight_version=9, + prompt_idx=17, + ) + assert result.metrics["row_assembly/fetch_ms"] == pytest.approx(10) + assert result.metrics["row_assembly/verify_ms"] == pytest.approx(25) + assert result.metrics["row_assembly/rollouts_ms"] == pytest.approx(35) diff --git a/tests/unit/data_plane/test_tq_token_sink.py b/tests/unit/data_plane/test_tq_token_sink.py index 8438cb8093..3ae634b4d7 100644 --- a/tests/unit/data_plane/test_tq_token_sink.py +++ b/tests/unit/data_plane/test_tq_token_sink.py @@ -218,3 +218,90 @@ 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"]) + + +@pytest.mark.parametrize( + "damage", ["missing", "malformed", "unready", "reordered", "reordered_partial"] +) +def test_group_fetch_isolates_bad_rows(tq_client, staging_partition, damage): + """Partial results and all-or-nothing readiness cannot poison siblings.""" + import torch + + from nemo_rl.data_plane.tq_token_sink import FetchedStagedCall + + sink = TQTokenSink(tq_client, staging_partition=staging_partition) + records = [ + build_fixture_artifacts("single_call", rollout_id=f"isolated_{i}")[0][0] + for i in range(3) + ] + keys = [record.staging_key for record in records] + for index, record in enumerate(records): + if damage not in ("missing", "reordered_partial") or index != 1: + assert sink.stage(record).ok + calls = [] + + class DamagedClient: + def get_samples(self, **kwargs): + requested = kwargs["sample_ids"] + calls.append(list(requested)) + if damage == "unready" and keys[1] in requested: + raise ValueError("Some fields are not ready in all the requested keys!") + if damage in ("reordered", "reordered_partial"): + kwargs["sample_ids"] = list(reversed(requested)) + rows = tq_client.get_samples(**kwargs) + if damage == "malformed" and keys[1] in requested: + values = rows["capture_mode"].clone() + values[requested.index(keys[1])] = torch.tensor(99) + rows["capture_mode"] = values + return rows + + result = TQTokenSource( + DamagedClient(), staging_partition=staging_partition + ).fetch_for_group(keys) + assert set(result.calls) == set(keys) + assert isinstance(result.calls[keys[0]], FetchedStagedCall) + assert isinstance(result.calls[keys[2]], FetchedStagedCall) + if damage == "reordered": + assert isinstance(result.calls[keys[1]], FetchedStagedCall) + else: + assert isinstance(result.calls[keys[1]], (KeyError, ValueError)) + assert len(calls) <= 2 * len(keys) - 1 + for key, item in result.calls.items(): + if isinstance(item, FetchedStagedCall): + assert item.snapshot.staging_key == key + + +def test_group_fetch_deduplicates_keys_and_times_only_reads( + tq_client, staging_partition, monkeypatch +): + import nemo_rl.data_plane.tq_token_sink as module + + sink = TQTokenSink(tq_client, staging_partition=staging_partition) + record = build_fixture_artifacts("single_call")[0][0] + assert sink.stage(record).ok + clock = [0.0] + calls = [] + decode = module._row_to_base_snapshot + + class TimedClient: + def get_samples(self, **kwargs): + calls.append(kwargs) + clock[0] += 0.010 + return tq_client.get_samples(**kwargs) + + def timed_decode(row): + clock[0] += 0.050 + return decode(row) + + monkeypatch.setattr(module.time, "perf_counter", lambda: clock[0]) + monkeypatch.setattr(module, "_row_to_base_snapshot", timed_decode) + source = TQTokenSource(TimedClient(), staging_partition=staging_partition) + result = source.fetch_for_group([record.staging_key] * 2) + assert len(calls) == 1 + assert calls[0]["sample_ids"] == [record.staging_key] + assert result.fetch_ms == pytest.approx(10) + assert clock[0] == pytest.approx(0.060) + empty = source.fetch_for_group([]) + assert empty.calls == {} + assert empty.fetch_ms == 0 + assert len(calls) == 1 diff --git a/tests/unit/experience/test_rollout_reassembler_actor.py b/tests/unit/experience/test_rollout_reassembler_actor.py index a31c22777d..a289bc21de 100644 --- a/tests/unit/experience/test_rollout_reassembler_actor.py +++ b/tests/unit/experience/test_rollout_reassembler_actor.py @@ -153,3 +153,38 @@ def test_every_forbidden_key_is_rejected(key) -> None: """Removing an entry from the denylist should fail loudly.""" with pytest.raises(TypeError, match="forbidden heavy field"): assert_metadata_only({key: [1, 2, 3]}) + + +@pytest.mark.parametrize("num_workers", [1, 3]) +def test_reassembler_pool_prefers_local_node_with_soft_affinity( + monkeypatch, num_workers +): + import nemo_rl.experience.rollout_reassembler_actor as module + + actor = MagicMock() + context = MagicMock() + context.get_node_id.return_value = "a" * 56 + monkeypatch.setattr(module, "RolloutReassemblerActor", actor) + monkeypatch.setattr(module.ray, "get_runtime_context", lambda: context) + dp_config, config = MagicMock(), MagicMock() + result = module.create_rollout_reassembler_actors( + dp_config, config, num_workers=num_workers + ) + strategy = actor.options.call_args.kwargs["scheduling_strategy"] + assert strategy.node_id == "a" * 56 + assert strategy.soft is True + assert strategy._spill_on_unavailable is True + assert len(result) == num_workers + assert actor.options.return_value.remote.call_count == num_workers + actor.options.return_value.remote.assert_called_with(dp_config, config) + + +def test_reassembler_pool_rejects_invalid_size_before_ray_initialization(monkeypatch): + import nemo_rl.experience.rollout_reassembler_actor as module + + context = MagicMock(side_effect=AssertionError("must validate size first")) + monkeypatch.setattr(module.ray, "get_runtime_context", context) + with pytest.raises(ValueError, match="must be positive"): + module.create_rollout_reassembler_actors( + MagicMock(), MagicMock(), num_workers=0 + )