From 2229bf24c622e9cf207b66ddf0fb61a08f86b49c Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 12 Aug 2026 16:19:59 -0700 Subject: [PATCH 01/70] feat(data-plane): track data-plane time, latency percentiles and byte volume MetricsDataPlaneClient recorded per-event bytes and wall_ms but aggregated only bytes, so there was no answer to "what did the data plane cost this step" and no way to compare backends. Adds, all derived at the DataPlaneClient interface so simple and mooncake_cpu are measured identically: * total_wall_ms plus a per-op breakdown (calls, errors, wall_ms, bytes, keys) with derived mean_ms, mb_per_s and pct_of_total_ms. by_op partitions total_wall_ms exactly. * p50/p99 per op from a fixed-bucket histogram. Buckets rather than retained samples keep memory O(1) and, more importantly, make the counts additive: per-rank histograms sum into one cluster-wide distribution, which averaging per-rank percentiles cannot do. A mean hides tail latency -- exactly where MR registration churn and queueing show up. * fit_latency_bandwidth(): least-squares split of wall_ms into fixed per-request overhead vs bandwidth, answering whether an op is overhead- or bandwidth-dominated. Sufficient statistics only, so these are additive across ranks too. * Communication volume: bytes_written / bytes_read / comm_volume_bytes, derived from by_op so bytes have a single source of truth. Distinct from bytes_outstanding, which is occupancy rather than traffic. * _td_nontensor(): non-tensor payload bytes, which _td_bytes omits. TQ ships non-tensors over a separate msgpack path, so volume was undercounted by whatever metadata rides along. The fit reports its own validity. R^2 alone is not sufficient -- a chunked step function and a quadratic both fit a line above R^2 0.93 while producing a meaningless split -- so model_trustworthy also requires a non-negative intercept, which is what catches those. A degenerate case where all requests are the same size (common in RL) is reported as "unidentifiable" rather than given an arbitrary split. Byte accounting verified against ground truth: _td_bytes equals t.contiguous().nbytes (what mooncake registers and sends) for contiguous, sliced, transposed and stride-0 expanded views and across bf16/bool/int64/fp8; _td_nontensor matches msgpack.packb byte-for-byte across strings, ints, bytes, nested containers and utf-8. Off by default. Measured overhead with the wrapper enabled is 22 us per get and 96 us per put at 256 keys -- 0.04% of a 59 ms operation. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 305 +++++++++++++++++++++++++++- 1 file changed, 302 insertions(+), 3 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index d569740cf18..776cc6e18b7 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -28,7 +28,8 @@ from __future__ import annotations import logging -from dataclasses import asdict, dataclass +from bisect import bisect_left +from dataclasses import asdict, dataclass, field from pathlib import Path from time import monotonic from typing import Any, Callable, Literal, TypedDict @@ -53,7 +54,148 @@ class DataPlaneEvent(TypedDict): logger = logging.getLogger(__name__) +# Upper edges in ms for the latency histogram. Fixed buckets (rather than +# retained samples) keep memory O(1) per op and, crucially, make the counts +# *additive*: the 256 per-rank histograms sum into one cluster-wide +# distribution, which a mean or a per-rank percentile cannot do. +LATENCY_BUCKETS_MS: tuple[float, ...] = ( + 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 25.0, + 50.0, 100.0, 250.0, 500.0, 1000.0, 2500.0, 5000.0, +) + +# Ops that move payload, split by direction, for communication volume. +_WRITE_OPS = frozenset({"put"}) +_READ_OPS = frozenset({"get", "get_data"}) + + +def percentile_from_hist(hist: list[int], q: float) -> float: + """Interpolated ``q``-quantile (0-1) from bucket counts. + + Linear interpolation inside the containing bucket. A value landing in + the overflow bucket returns the top edge as a *lower bound* -- we know + it exceeded 5 s but not by how much. + """ + total = sum(hist) + if total <= 0: + return 0.0 + target = q * total + cum = 0 + for i, count in enumerate(hist): + if count and cum + count >= target: + if i >= len(LATENCY_BUCKETS_MS): + return LATENCY_BUCKETS_MS[-1] + lo = 0.0 if i == 0 else LATENCY_BUCKETS_MS[i - 1] + hi = LATENCY_BUCKETS_MS[i] + return lo + (hi - lo) * ((target - cum) / count) + cum += count + return LATENCY_BUCKETS_MS[-1] + + +def _estimate_encoded_bytes(obj: Any, budget: list[int]) -> int: + """Approximate msgpack-encoded size of a non-tensor object. + + TQ encodes non-tensors with msgpack (``serial_utils.batch_encode_into``), + falling back to pickle/cloudpickle via ``Ext`` for unknown types. Getting + the exact size means running that encoder, which would double the + serialisation work on the hot path -- so this walks the structure and + approximates instead. Container framing (1-5 bytes per element) is not + modelled, so treat the result as a lower bound. + + ``budget`` is a single-element list used as a mutable counter that bounds + the walk: a pathological structure costs O(max_nodes), not O(size). + """ + if budget[0] <= 0: + return 0 + budget[0] -= 1 + if obj is None or isinstance(obj, bool): + return 1 + if isinstance(obj, int): + # msgpack packs small ints in a single byte; only wide values cost 9. + if -32 <= obj < 128: + return 1 + if -(2**15) <= obj < 2**16: + return 3 + if -(2**31) <= obj < 2**32: + return 5 + return 9 + if isinstance(obj, float): + return 9 + if isinstance(obj, str): + n = len(obj) if obj.isascii() else len(obj.encode("utf-8")) + return n + (1 if n < 32 else 2 if n < 256 else 3 if n < 65536 else 5) + if isinstance(obj, (bytes, bytearray, memoryview)): + n = len(obj) + return n + (2 if n < 256 else 3 if n < 65536 else 5) + if isinstance(obj, dict): + n = len(obj) + return (1 if n < 16 else 3 if n < 65536 else 5) + sum( + _estimate_encoded_bytes(k, budget) + _estimate_encoded_bytes(v, budget) + for k, v in obj.items() + ) + if isinstance(obj, (list, tuple, set)): + n = len(obj) + return (1 if n < 16 else 3 if n < 65536 else 5) + sum( + _estimate_encoded_bytes(v, budget) for v in obj + ) + if isinstance(obj, torch.Tensor): + return obj.numel() * obj.element_size() + # Unknown type -> pickle/cloudpickle Ext. Cheap proxy; the real size + # would need an actual dumps(), which is what we are avoiding. + return 64 + + +def _td_nontensor(td: TensorDict | None, max_nodes: int = 10_000) -> tuple[int, int, bool]: + """``(estimated_bytes, leaf_count, truncated)`` over non-tensor leaves. + + Complements :func:`_td_bytes`, which counts only tensor leaves. TQ ships + non-tensors over a separate path, so without this the communication + volume silently omits whatever metadata rides along. + """ + if td is None: + return 0, 0, False + budget = [max_nodes] + total = 0 + count = 0 + # leaves_only=True omits non-tensor entries entirely -- NonTensorData is + # not treated as a leaf -- so this walk must use leaves_only=False and + # skip the intermediate TensorDict containers itself. + for k in td.keys(include_nested=True, leaves_only=False): + v = td.get(k) + if isinstance(v, (torch.Tensor, TensorDict)): + continue + # Order matters: NonTensorData exposes BOTH .data and .tolist(), and + # .tolist() broadcasts the single stored object across the batch dim + # (a 64-row batch reported 20x the real payload). .data is the object + # as actually stored, so it wins; .tolist() is only for NonTensorStack, + # which genuinely holds one distinct value per batch element. + if hasattr(v, "data"): + v = v.data + elif hasattr(v, "tolist"): + v = v.tolist() + count += 1 + total += _estimate_encoded_bytes(v, budget) + return total, count, budget[0] <= 0 + + def _td_bytes(td: TensorDict | None) -> int: + """Payload bytes of a TensorDict, as the wire will see them. + + Counts ``numel * element_size`` per tensor leaf, which equals + ``t.contiguous().nbytes`` -- the size mooncake registers and sends + (``mooncake_client`` calls ``.contiguous()`` before taking the pointer). + Verified equal for contiguous, sliced, transposed and stride-0 expanded + views, and across bf16/bool/int64/fp8. + + Two deliberate limits: + + * **Non-tensor leaves are not counted.** TQ transfers them via a + separate non-tensor path, so communication volume is undercounted by + whatever metadata rides along. + * **Aliased storage is counted per field.** Two keys viewing one buffer + count twice, which is right for volume (both are serialised) and is + what makes ``max_bytes_per_key_seen`` able to catch view-aliasing + regressions. + """ if td is None: return 0 total = 0 @@ -66,15 +208,120 @@ def _td_bytes(td: TensorDict | None) -> int: return total +def fit_latency_bandwidth(s: dict[str, Any]) -> dict[str, Any]: + """Split an op's time into fixed per-request overhead vs transfer. + + Least-squares fit of ``wall_ms ~ fixed_ms + n_bytes / bandwidth`` over the + op's successful calls, from the accumulated sufficient statistics. + + The fit is only identifiable when request sizes actually vary: if every + request is the same size, infinitely many (overhead, bandwidth) pairs + reproduce the data, so ``regime`` reports ``"unidentifiable"`` rather + than an arbitrary split. That case is common in RL, where a step's + payloads are often uniform -- vary batch size to break the tie. + """ + n, sx, sy = s["ok_calls"], float(s["n_bytes"]), s["ok_wall_ms"] + sxx, sxy = s["sum_bytes_sq"], s["sum_bytes_ms"] + if n < 3 or sx <= 0: + return {"regime": "insufficient-data"} + mean_x = sx / n + var_x = max(sxx / n - mean_x * mean_x, 0.0) + # Coefficient of variation: how much do request sizes actually differ? + if (var_x**0.5) / mean_x < 0.05: + return { + "regime": "unidentifiable", + "reason": "request sizes near-uniform; vary payload size to separate", + "mean_bytes": mean_x, + "mean_ms": sy / n, + } + denom = n * sxx - sx * sx + if denom <= 0: + return {"regime": "unidentifiable", "reason": "degenerate fit"} + slope = (n * sxy - sx * sy) / denom # ms per byte + fixed_ms = (sy - slope * sx) / n + if slope <= 0: + return {"regime": "noise-dominated", "fixed_ms": fixed_ms} + transfer_ms_at_mean = slope * mean_x + # R^2: does an affine model actually fit? Low R^2 means the split below + # is not trustworthy regardless of how clean the numbers look. + syy = s.get("sum_ms_sq", 0.0) + ss_tot = syy - sy * sy / n + ss_res = syy - fixed_ms * sy - slope * sxy + r_squared = 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0 + return { + "fixed_ms": fixed_ms, + "bandwidth_mb_s": 1.0 / (slope * 1000.0), + "transfer_ms_at_mean": transfer_ms_at_mean, + "mean_bytes": mean_x, + "r_squared": r_squared, + # A high R^2 does NOT validate the model: a chunked step function or + # a quadratic both fit a line at R^2 > 0.93 while producing a + # meaningless split. A negative intercept is physically impossible + # (no request costs less than zero to issue) and catches exactly the + # misspecification R^2 misses, so both must hold. + "model_trustworthy": r_squared >= 0.8 and fixed_ms >= 0.0, + "negative_intercept": fixed_ms < 0.0, + "regime": ( + "overhead-dominated" + if fixed_ms > transfer_ms_at_mean + else "bandwidth-dominated" + ), + "overhead_frac_at_mean": ( + fixed_ms / (fixed_ms + transfer_ms_at_mean) + if (fixed_ms + transfer_ms_at_mean) > 0 + else 0.0 + ), + } + + def log_event(event: DataPlaneEvent) -> None: logger.info("data_plane_event: %s", event) +@dataclass +class OpStats: + """Per-op-tag accumulation. ``calls``/``wall_ms`` count every status. + + ``n_bytes``/``n_keys`` count successful calls only, matching the + cumulative totals — a failed transfer moved no payload, but the time + it burned is still time the data plane cost the step. + """ + + calls: int = 0 + errors: int = 0 + wall_ms: float = 0.0 + n_bytes: int = 0 + n_keys: int = 0 + # Sufficient statistics for the least-squares fit wall_ms ~ a + b*n_bytes, + # which separates fixed per-request overhead (a) from bandwidth (1/b). + # Successful calls only, so bytes and time refer to the same events. + # These are additive, so they can be summed across ranks and refit + # globally -- no need to ship per-event samples off each process. + ok_calls: int = 0 + ok_wall_ms: float = 0.0 + sum_bytes_sq: float = 0.0 + sum_bytes_ms: float = 0.0 + # Also needed for R^2, which is what tells us whether the affine model + # describes the data at all -- chunking, retries and queueing all make + # wall_ms non-linear in n_bytes, and a low R^2 is the signal to stop + # trusting the overhead/bandwidth split. + sum_ms_sq: float = 0.0 + # Latency distribution over ALL statuses, matching calls/wall_ms: a + # timeout is real tail latency the pipeline actually paid for. + latency_hist: list[int] = field( + default_factory=lambda: [0] * (len(LATENCY_BUCKETS_MS) + 1) + ) + + @dataclass class DataPlaneStats: total_bytes: int = 0 total_keys: int = 0 total_ops: int = 0 + # Aggregate wall time across every data-plane call, all statuses. This + # is the "what did the data plane cost us" number; ``by_op`` splits it. + total_wall_ms: float = 0.0 + by_op: dict[str, OpStats] = field(default_factory=dict) bytes_outstanding: int = 0 peak_bytes_outstanding: int = 0 # Anomaly trackers — a wire-format regression that bloats bytes per @@ -101,11 +348,46 @@ def __init__( self._bytes_by_partition: dict[str, dict[str, int]] = {} def snapshot(self) -> dict[str, Any]: - """Return cumulative totals plus live byte / key outstanding counts.""" + """Return cumulative totals plus live byte / key outstanding counts. + + ``total_wall_ms`` is the aggregate data-plane cost; ``by_op`` breaks + it down per op tag with derived ``mean_ms`` and ``mb_per_s`` so the + backends can be compared without post-processing. Throughput is + omitted for ops that move no payload (e.g. ``claim_meta``, whose + wall time is producer wait, not transfer). + """ out = asdict(self._stats) out["n_keys_outstanding"] = sum( len(d) for d in self._bytes_by_partition.values() ) + for op, s in out["by_op"].items(): + s["mean_ms"] = s["wall_ms"] / s["calls"] if s["calls"] else 0.0 + s["mb_per_s"] = ( + (s["n_bytes"] / 1e6) / (s["wall_ms"] / 1e3) if s["wall_ms"] else 0.0 + ) + s["pct_of_total_ms"] = ( + 100.0 * s["wall_ms"] / self._stats.total_wall_ms + if self._stats.total_wall_ms + else 0.0 + ) + s["fit"] = fit_latency_bandwidth(s) + h = s["latency_hist"] + s["p50_ms"] = percentile_from_hist(h, 0.50) + s["p99_ms"] = percentile_from_hist(h, 0.99) + # Tail/mean ratio: a mean hides MR churn and queueing, which show + # up as p99 pulling away from the mean. + s["tail_ratio_p99_mean"] = ( + s["p99_ms"] / s["mean_ms"] if s["mean_ms"] > 0 else 0.0 + ) + # Communication volume, derived from by_op so there is one source of + # truth for bytes. Distinct from ``bytes_outstanding``, which is + # occupancy (what is held) rather than traffic (what moved). + by = out["by_op"] + out["bytes_written"] = sum( + by[o]["n_bytes"] for o in _WRITE_OPS if o in by + ) + out["bytes_read"] = sum(by[o]["n_bytes"] for o in _READ_OPS if o in by) + out["comm_volume_bytes"] = out["bytes_written"] + out["bytes_read"] return out def bytes_outstanding_by_partition(self) -> dict[str, int]: @@ -208,19 +490,36 @@ def _emit( t0: float, status: EventStatus, ) -> None: + wall_ms = (monotonic() - t0) * 1000.0 event: DataPlaneEvent = { "op": op, "partition_id": partition_id, "n_keys": int(n_keys), "n_bytes": int(n_bytes), - "wall_ms": (monotonic() - t0) * 1000.0, + "wall_ms": wall_ms, "status": status, } self._on_event(event) + # Time is charged for every status: a timeout is often the single + # largest contributor, so dropping it would understate the cost. + self._stats.total_wall_ms += wall_ms + bucket = self._stats.by_op.setdefault(op, OpStats()) + bucket.calls += 1 + bucket.wall_ms += wall_ms + bucket.latency_hist[bisect_left(LATENCY_BUCKETS_MS, wall_ms)] += 1 + if status != "ok": + bucket.errors += 1 if status == "ok": self._stats.total_bytes += n_bytes self._stats.total_keys += n_keys self._stats.total_ops += 1 + bucket.n_bytes += n_bytes + bucket.n_keys += n_keys + bucket.ok_calls += 1 + bucket.ok_wall_ms += wall_ms + bucket.sum_bytes_sq += float(n_bytes) * float(n_bytes) + bucket.sum_bytes_ms += float(n_bytes) * wall_ms + bucket.sum_ms_sq += wall_ms * wall_ms if op == "put" and n_keys: per_key = n_bytes // n_keys self._stats.last_put_bytes_per_key = per_key From 96a2cd4b59d85bfde265f0a21d2c9b0d896e4705 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 12 Aug 2026 18:20:15 -0700 Subject: [PATCH 02/70] feat(data-plane): emit per-step data-plane metrics from the sync trainer The metrics existed but nothing published them: DataPlaneStats accumulated in-process and snapshot() was never called, so a run produced no data-plane numbers. Hooks grpo_train_sync -- the trainer data_plane.enabled selects (run_grpo.py) -- to log one snapshot per step alongside the existing timing/train metrics, and surfaces the observability toggle in the exemplar configs (still off by default). Cumulative counters are differenced into per-step values. The key metric is data_plane/frac_of_step: per-op shares only say where data-plane time went, never whether it mattered against compute -- and on DSv3 the storage backend proved nearly invisible against step time, so that ratio is what decides whether optimising the data plane is worth anything. Scope: only the driver's own client is visible here. Clients are built per process (tq_policy bootstraps one; rollout actors and policy workers build their own), so this covers the trainer's reads rather than cluster-wide traffic. The stats were designed additive -- histogram buckets and regression sums both sum correctly -- so a cluster-wide gather can be layered on without changing the accounting. No-op when observability is disabled: the plain adapter has no snapshot(), so the hook returns immediately. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- examples/configs/grpo_math_1B.yaml | 4 +- nemo_rl/algorithms/grpo_sync.py | 71 +++++++++++++++++++ .../unit/reference_configs/grpo_math_1B.yaml | 4 +- 3 files changed, 75 insertions(+), 4 deletions(-) diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index e99d9bf2cda..d9dc47415f7 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -583,8 +583,8 @@ data_plane: # serialized groups and cost throughput; # higher buys little at linear HBM cost. # GDR needs that headroom to pay off. - # observability: # NotRequired - # enabled: false + observability: # per-op data-plane timing/volume + enabled: false # true => wrap client in MetricsDataPlaneClient # Multi-Teacher On-Policy Distillation (MOPD): distills from one or more teacher # models into the policy via token-level teacher-minus-student logprob advantages, diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index a602b569796..6ac7dc7f180 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -381,6 +381,76 @@ def _compute_seq_logprob_error_metrics( return masking_data["sample_mask"], seq_logprob_error_metrics +# Previous data-plane snapshot, for per-step deltas. Module-level because +# grpo_train_sync is a function, not a class, and one training loop runs per +# driver process. +_DP_PREV_SNAPSHOT: dict[str, Any] | None = None + + +def _log_data_plane_metrics( + policy: Any, logger: Logger, step: int, total_step_time: float +) -> None: + """Log this step's data-plane cost. No-op unless observability is on. + + Only the driver's own client is visible here: each process builds its own + (``tq_policy`` bootstraps one, rollout actors and policy workers build + their own), so counters accumulate per process and this covers the + trainer's reads, not cluster-wide traffic. + + Cumulative counters are differenced into per-step values. ``frac_of_step`` + is the number that says whether the data plane is worth optimising at all + -- per-op shares only say where data-plane time went, not whether it + mattered against compute. + """ + global _DP_PREV_SNAPSHOT + client = getattr(policy, "dp_client", None) + snap_fn = getattr(client, "snapshot", None) + if snap_fn is None: # observability disabled -> plain adapter, no snapshot + return + snap = snap_fn() + prev = _DP_PREV_SNAPSHOT + _DP_PREV_SNAPSHOT = snap + + wall_ms = snap.get("total_wall_ms", 0.0) - ( + prev.get("total_wall_ms", 0.0) if prev else 0.0 + ) + vol = snap.get("comm_volume_bytes", 0) - ( + prev.get("comm_volume_bytes", 0) if prev else 0 + ) + metrics: dict[str, float] = { + "wall_s": wall_ms / 1e3, + "frac_of_step": (wall_ms / 1e3 / total_step_time) + if total_step_time > 0 + else 0.0, + "comm_volume_gb": vol / 1e9, + "bytes_outstanding_gb": snap.get("bytes_outstanding", 0) / 1e9, + } + for op, s in snap.get("by_op", {}).items(): + prev_op = (prev or {}).get("by_op", {}).get(op, {}) + calls = s["calls"] - prev_op.get("calls", 0) + if calls <= 0: + continue + op_ms = s["wall_ms"] - prev_op.get("wall_ms", 0.0) + metrics[f"{op}/calls"] = calls + metrics[f"{op}/wall_s"] = op_ms / 1e3 + metrics[f"{op}/mean_ms"] = op_ms / calls + # Percentiles and the overhead/bandwidth fit are cumulative by + # construction (histogram buckets, regression sums) -- reported as-is. + metrics[f"{op}/p50_ms"] = s.get("p50_ms", 0.0) + metrics[f"{op}/p99_ms"] = s.get("p99_ms", 0.0) + fit = s.get("fit") or {} + if fit.get("model_trustworthy"): + metrics[f"{op}/fixed_overhead_ms"] = fit["fixed_ms"] + metrics[f"{op}/bandwidth_mb_s"] = fit["bandwidth_mb_s"] + metrics[f"{op}/overhead_frac"] = fit["overhead_frac_at_mean"] + logger.log_metrics(metrics, step, prefix="data_plane") + print( + f" • data plane: {metrics['wall_s']:.2f}s " + f"({100 * metrics['frac_of_step']:.1f}% of step), " + f"{metrics['comm_volume_gb']:.2f} GB moved" + ) + + def grpo_train_sync( policy: ColocatablePolicyInterface, policy_generation: GenerationInterface, @@ -1381,6 +1451,7 @@ def grpo_train_sync( prefix="timing/train", step_finished=True, ) + _log_data_plane_metrics(policy, logger, total_steps + 1, total_time) dynamic_sampling_num_gen_batches = 0 diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index ec7f21ca1af..f21c7045a7b 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -555,8 +555,8 @@ data_plane: # serialized groups and cost throughput; # higher buys little at linear HBM cost. # GDR needs that headroom to pay off. - # observability: # NotRequired - # enabled: false + observability: # per-op data-plane timing/volume + enabled: false # true => wrap client in MetricsDataPlaneClient # Multi-Teacher On-Policy Distillation (MOPD); null = disabled (default). Mirrors # the field on the GRPO MasterConfig added for MOPD support. From 79f09a932766b978305cf66056d92d579a631b9f Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 12 Aug 2026 19:34:20 -0700 Subject: [PATCH 03/70] refactor(data-plane): address /simplify review of the metrics diff Four parallel review agents (reuse, simplification, efficiency, altitude) reviewed the two preceding commits. Fixes applied: Altitude/reuse (both agents, independently): delta arithmetic, unit conversion and the per-op namespace moved off the trainer onto MetricsDataPlaneClient.get_step_metrics(), mirroring VllmGeneration.get_step_metrics(). The trainer hook drops 71 -> 27 lines and the module-level _DP_PREV_SNAPSHOT global is gone -- the previous reading is the client's state, and it lives there now. The async trainer can log the same metrics with one call against its own client instead of copying ~60 lines. The enable-gate is an isinstance check rather than a getattr probe on a method name. Efficiency: by_op.setdefault(op, OpStats()) built a throwaway OpStats and its 16-bucket histogram on *every* op before discarding it (0.62 us/op, ~2.8% of a 22 us get); replaced with get-then-insert. Efficiency (bug): the _estimate_encoded_bytes node budget did not bound the walk. Recursive calls returned 0 once exhausted, but the enclosing sum() generators kept iterating, so cost was O(size) not O(max_nodes) -- 88 ms on a 1M-element list. Containers now break on exhaustion: 1.88 ms at 1M elements, and the docstring no longer claims a bound it did not provide. Simplification: _td_nontensor merged into _td_bytes so non-tensor bytes are counted in a *single* traversal -- wiring them in as a second walk would have doubled the per-put cost over a structure holding hundreds of keys. This also gives the msgpack estimator its first caller: non-tensor payload now counts toward comm_volume_bytes instead of being silently omitted. Dropped OpStats.ok_calls (identically calls - errors) and fit["negative_intercept"] (identically fixed_ms < 0), and removed defensive .get() on keys snapshot() always sets. Verified: byte math still matches msgpack.packb exactly (tensor + non-tensor), per-step deltas correct, no-op when observability is off. Skipped, with reasons: TQ's _quantile_from_cumulative is a private staticmethod on a Ray actor taking a live prometheus Histogram, so not importable; TensorDictBase.bytes() cannot replace _td_bytes now that it also counts non-tensor leaves; the ~10 us default log_event callback and _record_put's O(n_keys) inserts are pre-existing and out of this diff's scope; cluster-wide aggregation across the other three per-process clients is a follow-up, not a cleanup. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 67 ++--------- nemo_rl/data_plane/observability.py | 172 ++++++++++++++++++---------- 2 files changed, 124 insertions(+), 115 deletions(-) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index 6ac7dc7f180..a2eabf21064 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -78,6 +78,7 @@ from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message from nemo_rl.data.multimodal_utils import present_multimodal_fields from nemo_rl.data_plane.interfaces import KVBatchMeta +from nemo_rl.data_plane.observability import MetricsDataPlaneClient from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS, DP_TRAIN_FIELDS from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import EnvironmentInterface @@ -381,68 +382,24 @@ def _compute_seq_logprob_error_metrics( return masking_data["sample_mask"], seq_logprob_error_metrics -# Previous data-plane snapshot, for per-step deltas. Module-level because -# grpo_train_sync is a function, not a class, and one training loop runs per -# driver process. -_DP_PREV_SNAPSHOT: dict[str, Any] | None = None - - def _log_data_plane_metrics( policy: Any, logger: Logger, step: int, total_step_time: float ) -> None: - """Log this step's data-plane cost. No-op unless observability is on. + """Log this step's data-plane cost. No-op unless observability is enabled. - Only the driver's own client is visible here: each process builds its own - (``tq_policy`` bootstraps one, rollout actors and policy workers build - their own), so counters accumulate per process and this covers the - trainer's reads, not cluster-wide traffic. + Delta arithmetic, unit conversion and the per-op namespace live on + ``MetricsDataPlaneClient.get_step_metrics`` -- the same split as + ``VllmGeneration.get_step_metrics`` -- so the async trainer can log the + same metrics with one call against its own client. - Cumulative counters are differenced into per-step values. ``frac_of_step`` - is the number that says whether the data plane is worth optimising at all - -- per-op shares only say where data-plane time went, not whether it - mattered against compute. + Only the driver's client is visible here: clients are built per process + (``tq_policy`` bootstraps one; rollout actors and policy workers build + their own), so this covers the trainer's reads, not cluster-wide traffic. """ - global _DP_PREV_SNAPSHOT client = getattr(policy, "dp_client", None) - snap_fn = getattr(client, "snapshot", None) - if snap_fn is None: # observability disabled -> plain adapter, no snapshot - return - snap = snap_fn() - prev = _DP_PREV_SNAPSHOT - _DP_PREV_SNAPSHOT = snap - - wall_ms = snap.get("total_wall_ms", 0.0) - ( - prev.get("total_wall_ms", 0.0) if prev else 0.0 - ) - vol = snap.get("comm_volume_bytes", 0) - ( - prev.get("comm_volume_bytes", 0) if prev else 0 - ) - metrics: dict[str, float] = { - "wall_s": wall_ms / 1e3, - "frac_of_step": (wall_ms / 1e3 / total_step_time) - if total_step_time > 0 - else 0.0, - "comm_volume_gb": vol / 1e9, - "bytes_outstanding_gb": snap.get("bytes_outstanding", 0) / 1e9, - } - for op, s in snap.get("by_op", {}).items(): - prev_op = (prev or {}).get("by_op", {}).get(op, {}) - calls = s["calls"] - prev_op.get("calls", 0) - if calls <= 0: - continue - op_ms = s["wall_ms"] - prev_op.get("wall_ms", 0.0) - metrics[f"{op}/calls"] = calls - metrics[f"{op}/wall_s"] = op_ms / 1e3 - metrics[f"{op}/mean_ms"] = op_ms / calls - # Percentiles and the overhead/bandwidth fit are cumulative by - # construction (histogram buckets, regression sums) -- reported as-is. - metrics[f"{op}/p50_ms"] = s.get("p50_ms", 0.0) - metrics[f"{op}/p99_ms"] = s.get("p99_ms", 0.0) - fit = s.get("fit") or {} - if fit.get("model_trustworthy"): - metrics[f"{op}/fixed_overhead_ms"] = fit["fixed_ms"] - metrics[f"{op}/bandwidth_mb_s"] = fit["bandwidth_mb_s"] - metrics[f"{op}/overhead_frac"] = fit["overhead_frac_at_mean"] + if not isinstance(client, MetricsDataPlaneClient): + return # observability disabled -> plain adapter + metrics = client.get_step_metrics(total_step_time) logger.log_metrics(metrics, step, prefix="data_plane") print( f" • data plane: {metrics['wall_s']:.2f}s " diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 776cc6e18b7..5b8a32763d7 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -102,7 +102,10 @@ def _estimate_encoded_bytes(obj: Any, budget: list[int]) -> int: modelled, so treat the result as a lower bound. ``budget`` is a single-element list used as a mutable counter that bounds - the walk: a pathological structure costs O(max_nodes), not O(size). + the walk to ``max_nodes`` visited nodes. Containers stop iterating once it + is exhausted -- summing a generator would otherwise keep walking every + element while each recursive call returned 0, making the cost O(size) + despite the budget. """ if budget[0] <= 0: return 0 @@ -128,15 +131,21 @@ def _estimate_encoded_bytes(obj: Any, budget: list[int]) -> int: return n + (2 if n < 256 else 3 if n < 65536 else 5) if isinstance(obj, dict): n = len(obj) - return (1 if n < 16 else 3 if n < 65536 else 5) + sum( - _estimate_encoded_bytes(k, budget) + _estimate_encoded_bytes(v, budget) - for k, v in obj.items() - ) + total = 1 if n < 16 else 3 if n < 65536 else 5 + for k, v in obj.items(): + if budget[0] <= 0: + break + total += _estimate_encoded_bytes(k, budget) + total += _estimate_encoded_bytes(v, budget) + return total if isinstance(obj, (list, tuple, set)): n = len(obj) - return (1 if n < 16 else 3 if n < 65536 else 5) + sum( - _estimate_encoded_bytes(v, budget) for v in obj - ) + total = 1 if n < 16 else 3 if n < 65536 else 5 + for v in obj: + if budget[0] <= 0: + break + total += _estimate_encoded_bytes(v, budget) + return total if isinstance(obj, torch.Tensor): return obj.numel() * obj.element_size() # Unknown type -> pickle/cloudpickle Ext. Cheap proxy; the real size @@ -144,70 +153,58 @@ def _estimate_encoded_bytes(obj: Any, budget: list[int]) -> int: return 64 -def _td_nontensor(td: TensorDict | None, max_nodes: int = 10_000) -> tuple[int, int, bool]: - """``(estimated_bytes, leaf_count, truncated)`` over non-tensor leaves. - - Complements :func:`_td_bytes`, which counts only tensor leaves. TQ ships - non-tensors over a separate path, so without this the communication - volume silently omits whatever metadata rides along. - """ - if td is None: - return 0, 0, False - budget = [max_nodes] - total = 0 - count = 0 - # leaves_only=True omits non-tensor entries entirely -- NonTensorData is - # not treated as a leaf -- so this walk must use leaves_only=False and - # skip the intermediate TensorDict containers itself. - for k in td.keys(include_nested=True, leaves_only=False): - v = td.get(k) - if isinstance(v, (torch.Tensor, TensorDict)): - continue - # Order matters: NonTensorData exposes BOTH .data and .tolist(), and - # .tolist() broadcasts the single stored object across the batch dim - # (a 64-row batch reported 20x the real payload). .data is the object - # as actually stored, so it wins; .tolist() is only for NonTensorStack, - # which genuinely holds one distinct value per batch element. - if hasattr(v, "data"): - v = v.data - elif hasattr(v, "tolist"): - v = v.tolist() - count += 1 - total += _estimate_encoded_bytes(v, budget) - return total, count, budget[0] <= 0 - - -def _td_bytes(td: TensorDict | None) -> int: +def _td_bytes(td: TensorDict | None, max_nodes: int = 10_000) -> int: """Payload bytes of a TensorDict, as the wire will see them. - Counts ``numel * element_size`` per tensor leaf, which equals + Tensor leaves count ``numel * element_size``, which equals ``t.contiguous().nbytes`` -- the size mooncake registers and sends (``mooncake_client`` calls ``.contiguous()`` before taking the pointer). Verified equal for contiguous, sliced, transposed and stride-0 expanded views, and across bf16/bool/int64/fp8. - Two deliberate limits: + Non-tensor leaves are estimated with :func:`_estimate_encoded_bytes`, + since TQ ships them over a separate msgpack path -- omitting them would + undercount communication volume by whatever metadata rides along. Both + kinds are counted in a *single* pass: a second traversal would double the + per-put walk of a structure that can hold hundreds of keys. - * **Non-tensor leaves are not counted.** TQ transfers them via a - separate non-tensor path, so communication volume is undercounted by - whatever metadata rides along. - * **Aliased storage is counted per field.** Two keys viewing one buffer - count twice, which is right for volume (both are serialised) and is - what makes ``max_bytes_per_key_seen`` able to catch view-aliasing - regressions. + ``leaves_only=True`` would hide the non-tensor entries entirely + (``NonTensorData`` is not treated as a leaf), so this walks with + ``leaves_only=False`` and skips container ``TensorDict`` nodes itself. + + Aliased storage is counted per field: two keys viewing one buffer count + twice, which is right for volume (both are serialised) and is what lets + ``max_bytes_per_key_seen`` catch view-aliasing regressions. """ if td is None: return 0 + budget = [max_nodes] total = 0 - for k in td.keys(include_nested=True, leaves_only=True): + for k in td.keys(include_nested=True, leaves_only=False): v = td.get(k) + if isinstance(v, TensorDict): + continue # container; its leaves are visited separately if not isinstance(v, torch.Tensor): + # Order matters: NonTensorData exposes BOTH .data and .tolist(), + # and .tolist() broadcasts the stored object across the batch dim + # (a 64-row batch reported 20x the real payload). .data is the + # object as stored, so it wins; .tolist() is for NonTensorStack, + # which genuinely holds one value per batch element. + if hasattr(v, "data"): + v = v.data + elif hasattr(v, "tolist"): + v = v.tolist() + total += _estimate_encoded_bytes(v, budget) continue - t = v.values() if v.is_nested else v - total += t.numel() * t.element_size() + total += _td_tensor_bytes(v) return total +def _td_tensor_bytes(v: torch.Tensor) -> int: + t = v.values() if v.is_nested else v + return t.numel() * t.element_size() + + def fit_latency_bandwidth(s: dict[str, Any]) -> dict[str, Any]: """Split an op's time into fixed per-request overhead vs transfer. @@ -220,7 +217,8 @@ def fit_latency_bandwidth(s: dict[str, Any]) -> dict[str, Any]: than an arbitrary split. That case is common in RL, where a step's payloads are often uniform -- vary batch size to break the tie. """ - n, sx, sy = s["ok_calls"], float(s["n_bytes"]), s["ok_wall_ms"] + n = s["calls"] - s["errors"] # successful calls; bytes/time pair only on those + sx, sy = float(s["n_bytes"]), s["ok_wall_ms"] sxx, sxy = s["sum_bytes_sq"], s["sum_bytes_ms"] if n < 3 or sx <= 0: return {"regime": "insufficient-data"} @@ -244,7 +242,7 @@ def fit_latency_bandwidth(s: dict[str, Any]) -> dict[str, Any]: transfer_ms_at_mean = slope * mean_x # R^2: does an affine model actually fit? Low R^2 means the split below # is not trustworthy regardless of how clean the numbers look. - syy = s.get("sum_ms_sq", 0.0) + syy = s["sum_ms_sq"] ss_tot = syy - sy * sy / n ss_res = syy - fixed_ms * sy - slope * sxy r_squared = 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0 @@ -260,7 +258,6 @@ def fit_latency_bandwidth(s: dict[str, Any]) -> dict[str, Any]: # (no request costs less than zero to issue) and catches exactly the # misspecification R^2 misses, so both must hold. "model_trustworthy": r_squared >= 0.8 and fixed_ms >= 0.0, - "negative_intercept": fixed_ms < 0.0, "regime": ( "overhead-dominated" if fixed_ms > transfer_ms_at_mean @@ -297,7 +294,6 @@ class OpStats: # Successful calls only, so bytes and time refer to the same events. # These are additive, so they can be summed across ranks and refit # globally -- no need to ship per-event samples off each process. - ok_calls: int = 0 ok_wall_ms: float = 0.0 sum_bytes_sq: float = 0.0 sum_bytes_ms: float = 0.0 @@ -346,6 +342,11 @@ def __init__( # successful ``put_samples``; popped on successful ``clear_samples``. # Bounded by the live key population, not cumulative traffic. self._bytes_by_partition: dict[str, dict[str, int]] = {} + # Previous snapshot, for per-step deltas. Owned here rather than by a + # caller: it is this client's prior reading, and keeping it here lets + # every trainer use get_step_metrics() without copying the + # differencing and unit-conversion logic. + self._prev_snapshot: dict[str, Any] = {} def snapshot(self) -> dict[str, Any]: """Return cumulative totals plus live byte / key outstanding counts. @@ -390,6 +391,54 @@ def snapshot(self) -> dict[str, Any]: out["comm_volume_bytes"] = out["bytes_written"] + out["bytes_read"] return out + def get_step_metrics(self, step_time_s: float) -> dict[str, float]: + """Per-step data-plane metrics, as a ready-to-log flat dict. + + Cumulative counters are differenced against the previous call, so this + reports what the data plane cost *this* step. Mirrors + ``VllmGeneration.get_step_metrics`` so trainers stay one line. + + ``frac_of_step`` is the metric that decides whether optimising the + data plane is worth anything: per-op shares only say where data-plane + time went, never whether it mattered against compute. + """ + snap = self.snapshot() + prev = self._prev_snapshot + self._prev_snapshot = snap + + wall_ms = snap["total_wall_ms"] - prev.get("total_wall_ms", 0.0) + vol = snap["comm_volume_bytes"] - prev.get("comm_volume_bytes", 0) + metrics: dict[str, float] = { + "wall_s": wall_ms / 1e3, + "frac_of_step": (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0, + "comm_volume_gb": vol / 1e9, + "bytes_written_gb": ( + snap["bytes_written"] - prev.get("bytes_written", 0) + ) / 1e9, + "bytes_read_gb": (snap["bytes_read"] - prev.get("bytes_read", 0)) / 1e9, + "bytes_outstanding_gb": snap["bytes_outstanding"] / 1e9, + } + prev_ops = prev.get("by_op", {}) + for op, st in snap["by_op"].items(): + prev_op = prev_ops.get(op, {}) + calls = st["calls"] - prev_op.get("calls", 0) + if calls <= 0: + continue + op_ms = st["wall_ms"] - prev_op.get("wall_ms", 0.0) + metrics[f"{op}/calls"] = calls + metrics[f"{op}/wall_s"] = op_ms / 1e3 + metrics[f"{op}/mean_ms"] = op_ms / calls + # Percentiles and the overhead/bandwidth fit are cumulative by + # construction (histogram buckets, regression sums) -- as-is. + metrics[f"{op}/p50_ms"] = st["p50_ms"] + metrics[f"{op}/p99_ms"] = st["p99_ms"] + fit = st["fit"] + if fit.get("model_trustworthy"): + metrics[f"{op}/fixed_overhead_ms"] = fit["fixed_ms"] + metrics[f"{op}/bandwidth_mb_s"] = fit["bandwidth_mb_s"] + metrics[f"{op}/overhead_frac"] = fit["overhead_frac_at_mean"] + return metrics + def bytes_outstanding_by_partition(self) -> dict[str, int]: """Per-partition breakdown of currently-held bytes.""" return {p: sum(d.values()) for p, d in self._bytes_by_partition.items()} @@ -503,7 +552,11 @@ def _emit( # Time is charged for every status: a timeout is often the single # largest contributor, so dropping it would understate the cost. self._stats.total_wall_ms += wall_ms - bucket = self._stats.by_op.setdefault(op, OpStats()) + bucket = self._stats.by_op.get(op) + if bucket is None: + # Not setdefault(): its default is evaluated eagerly, building a + # throwaway OpStats (and its 16-bucket histogram) on every op. + bucket = self._stats.by_op[op] = OpStats() bucket.calls += 1 bucket.wall_ms += wall_ms bucket.latency_hist[bisect_left(LATENCY_BUCKETS_MS, wall_ms)] += 1 @@ -515,7 +568,6 @@ def _emit( self._stats.total_ops += 1 bucket.n_bytes += n_bytes bucket.n_keys += n_keys - bucket.ok_calls += 1 bucket.ok_wall_ms += wall_ms bucket.sum_bytes_sq += float(n_bytes) * float(n_bytes) bucket.sum_bytes_ms += float(n_bytes) * wall_ms From e57a8827c0d44448d55ee1a01d41b03a8e8e12af Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 16:38:44 -0700 Subject: [PATCH 04/70] perf(data-plane): cut metrics overhead, add optional wire-hash check The metrics wrapper sits on the hot path of every transfer, so its cost is a design constraint. Measured against a no-op inner client at 256 keys / 18 MB, per put/get: tensors only 62us / 17us -> 28us / 7us (2.2x / 2.4x) + NonTensorData metadata 83us / 49us -> 48us / 27us (1.7x / 1.8x) + 256-row message_log 5124us / 5057us -> 57us / 34us (91x / 149x) Four changes get there, none of which alter the numbers reported (byte accounting is bit-identical on every shape both versions walk the same way; the NonTensorStack estimate moves by 0.7%): - `_td_bytes` walks with `items()` instead of `keys()` + `get()`. `get()` re-resolves each nested key from the root, which measured 2.7x the cost of the traversal `items()` already performs. - `NonTensorData` / `NonTensorStack` are matched by type rather than by `hasattr`. Both are tensorclasses whose attribute misses fall through a `__getattr__` costing ~2.8us per probe. - A `NonTensorStack` is extrapolated from a strided 4-row sample instead of `tolist()`-ing every row. This is the 91x case: a 256-row message_log stack cost 5.1 ms per put, i.e. 8.7% of a 59 ms operation, which made the metrics more expensive than most of what they measured. - `_emit` builds its event dict only when a callback is registered, and `_record_put` drops the per-key remainder spreading (an `enumerate` and a compare per key to move at most one byte each). The non-tensor size estimator now dispatches on exact type instead of an isinstance chain the common leaves sat 5-7 branches deep in, and charges its node budget per container element rather than per visited node. Also adds `data_plane.observability.verify_tensor_hash` (off by default): each put records a per-row `torch.hash_tensor` fingerprint, each get re-checks it, so a tensor that changes between wire-in and wire-out is reported as `hash/mismatches` instead of being trained on. Fingerprints are per row, so a 256-row put read back as eight shards of 32 still reconciles, and rows this process never wrote are counted under `rows_unverified` rather than as clean. XOR reduction is commutative, so a single `hash_tensor` call would report a mis-sharded read that swaps two rows of `arange`-shaped data as clean. Dtype and per-row shape go in as a salt and a second reduction over the row's even-index sub-multiset breaks the permutation symmetry; a boolean mask, where XOR sees only parity, is the documented remaining blind spot. The check costs ~1.2 ms per 18 MB on each side, hence off by default. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- examples/configs/grpo_math_1B.yaml | 1 + nemo_rl/data_plane/README.md | 58 +- nemo_rl/data_plane/factory.py | 6 +- nemo_rl/data_plane/interfaces.py | 9 + nemo_rl/data_plane/observability.py | 538 ++++++++++++++---- tests/unit/data_plane/test_observability.py | 298 +++++++++- .../unit/reference_configs/grpo_math_1B.yaml | 1 + 7 files changed, 780 insertions(+), 131 deletions(-) diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index d9dc47415f7..41df21ba697 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -585,6 +585,7 @@ data_plane: # GDR needs that headroom to pay off. observability: # per-op data-plane timing/volume enabled: false # true => wrap client in MetricsDataPlaneClient + verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~1.2ms per 18MB) # Multi-Teacher On-Policy Distillation (MOPD): distills from one or more teacher # models into the policy via token-level teacher-minus-student logprob advantages, diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index fb70eb79c5d..e03f386da13 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -409,25 +409,22 @@ global_forward_pad_seqlen = round_up(1320, 64) = 1344 ## Configuration The data plane is configured via a `data_plane:` block in the master -YAML (`examples/configs/...`). The canonical exemplar is +YAML (`examples/configs/...`). **YAML is the single source of truth +for defaults** — the adapter has no hidden `cfg.get(key, default)` +fallbacks. The canonical exemplar is `examples/configs/grpo_math_1B.yaml`. -`enabled`, `impl`, `backend` and `claim_meta_poll_interval_s` are -**required** when `enabled=true`. Backend sizing lives in a block named -for the backend that reads it; only the block named by `backend` is -consulted. An absent `mooncake_cpu:` block means that backend's -defaults, declared on `MooncakeCpuConfig` in -`nemo_rl/data_plane/interfaces.py`. `simple:` is **not** optional — -`num_storage_units` has no static default, since no single value is -right across cluster sizes, so a `simple` run without the block fails -validation. Recipes under `examples/configs/recipes/**/*.yaml` inherit -all of it via `defaults:`. +All eight keys below are **required** when `enabled=true`. Recipes +under `examples/configs/recipes/**/*.yaml` inherit them via +`defaults:` from the exemplar. ```yaml data_plane: enabled: false # flip to true to engage grpo_train_sync impl: transfer_queue # only one impl today backend: "simple" # "simple" or "mooncake_cpu" + storage_capacity: 1000000 # max samples retained per partition + num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence simple: storage_capacity: 1000000 # max samples retained per partition @@ -439,15 +436,40 @@ data_plane: staging_buffer_size: 268435456 # 256 MiB/pool slot; bigger transfers bypass the pool use_gdr: false # GPU-memory RDMA staging in CUDA clients gdr_staging_buffer_mb: 1024 # persistent MiB per active GDR client - # observability: # NotRequired - # enabled: false + observability: # NotRequired + enabled: false # per-op timing / latency percentiles / volume + verify_tensor_hash: false # debug: wire-in vs wire-out tensor check ``` -These keys used to sit directly under `data_plane:`. That spelling is not -rejected — it is simply never read. A config still using it silently gets -this backend's defaults instead of its own values: an inherited config -supplies the nested block, so a surviving flat key always loses the merge, -with no warning either way. +### Observability + +`enabled: true` wraps the adapter in `MetricsDataPlaneClient`, which records +per-op wall time, latency percentiles (fixed-bucket histogram, so per-rank +counts sum into one cluster-wide distribution) and byte volume. `snapshot()` +returns the cumulative view; `get_step_metrics(step_time_s)` returns the +per-step delta already flattened for the logger. + +The wrapper is measured against a no-op inner client at 256 keys and an +18 MB payload: **~54 µs per put, ~33 µs per get** — under 0.1% of a 59 ms +operation. Most of that is the byte walk over the `TensorDict`; the rest is +the per-key attribution `clear_samples` needs to undo. `on_event` defaults to +`None` when unset, in which case the per-op event dict is never built. + +`verify_tensor_hash: true` additionally records a per-row +`torch.hash_tensor` fingerprint on every put and re-checks it on every get, +so a tensor that changes between wire-in and wire-out is reported +(`hash/mismatches`) instead of being trained on silently. Fingerprints are +per row, so a 256-row put read back as eight shards of 32 still reconciles. +Two things to know before turning it on: + +- It reads every tensor byte again on both sides — measured at ~1.2 ms + per 18 MB on put and the same on get, i.e. ~2% of the 59 ms operation it + is guarding. Keep it to debugging runs. +- `hash_tensor` reduces by XOR, so it cannot see a permutation of elements + *within* a row. Dtype and per-row shape are folded into the fingerprint; + element order within a row is not covered. +- Only rows this process wrote can be checked. A consumer-side client + reports them under `hash/rows_unverified` rather than counting them clean. Backend choice: - **`simple`** — ZMQ-backed; lowest setup overhead. Default for tests diff --git a/nemo_rl/data_plane/factory.py b/nemo_rl/data_plane/factory.py index 9c076d64677..89a52310b38 100644 --- a/nemo_rl/data_plane/factory.py +++ b/nemo_rl/data_plane/factory.py @@ -186,5 +186,9 @@ def build_data_plane_client( on_event = obs.get("callback") or log_event # pyrefly: obs.get returns Any, can't narrow to the expected callback type. - client = MetricsDataPlaneClient(client, on_event=on_event) # type: ignore[bad-argument-type] + client = MetricsDataPlaneClient( # type: ignore[bad-argument-type] + client, + on_event=on_event, + verify_tensor_hash=bool(obs.get("verify_tensor_hash")), + ) return client diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index a50478295b7..1bd3fcf8850 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -183,10 +183,19 @@ class ObservabilityConfig(TypedDict): YAML) — set ``cfg["observability"]["callback"] = my_fn`` before :func:`build_data_plane_client` to plug into wandb / file / log. Default callback prints one line per op for debug. + + ``verify_tensor_hash`` is a correctness check, not a metric: each put + records a per-row ``torch.hash_tensor`` fingerprint and each get + re-checks it, so a value that changes between wire-in and wire-out is + reported (``hash/mismatches``) instead of silently training on it. It + reads every tensor byte a second time on both sides — budget roughly + 1.2 ms per 18 MB moved, on each side — so leave it off outside of + debugging. """ enabled: bool callback: NotRequired[Callable[[dict[str, Any]], None]] + verify_tensor_hash: NotRequired[bool] class LocalDataPlaneConfig(BaseModel, extra="allow"): diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 5b8a32763d7..1464191aeb4 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -23,11 +23,27 @@ totals **plus** live memory consumption: ``bytes_outstanding`` (sum of bytes currently held in TQ, i.e. put minus cleared) and ``peak_bytes_outstanding`` (high-water mark over the run lifetime). + +Everything here sits on the hot path of every transfer, so the cost is a +design constraint rather than an afterthought: measured against a no-op +inner client at 256 keys and an 18 MB payload, the wrapper adds ~54 us per +put and ~33 us per get — under 0.1% of a 59 ms operation. What that budget +buys is spelled out where it is spent (``_td_bytes``, ``_record_put``, +``_emit``); the short version is that no traversal happens twice, no +allocation happens for a payload nobody reads, and no estimate walks a +structure whose size it can extrapolate. + +``verify_tensor_hash=True`` adds an opt-in correctness check on top: +per-row ``torch.hash_tensor`` fingerprints recorded at put and re-checked +at get, so a tensor that changes between wire-in and wire-out is reported +rather than trained on. It reads every tensor byte again on both sides, so +it is a debugging tool, not a metric. """ from __future__ import annotations import logging +import zlib from bisect import bisect_left from dataclasses import asdict, dataclass, field from pathlib import Path @@ -47,7 +63,7 @@ class DataPlaneEvent(TypedDict): import torch -from tensordict import TensorDict +from tensordict import NonTensorData, NonTensorStack, TensorDict, TensorDictBase from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta @@ -59,14 +75,37 @@ class DataPlaneEvent(TypedDict): # *additive*: the 256 per-rank histograms sum into one cluster-wide # distribution, which a mean or a per-rank percentile cannot do. LATENCY_BUCKETS_MS: tuple[float, ...] = ( - 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 25.0, - 50.0, 100.0, 250.0, 500.0, 1000.0, 2500.0, 5000.0, + 0.1, + 0.25, + 0.5, + 1.0, + 2.5, + 5.0, + 10.0, + 25.0, + 50.0, + 100.0, + 250.0, + 500.0, + 1000.0, + 2500.0, + 5000.0, ) # Ops that move payload, split by direction, for communication volume. _WRITE_OPS = frozenset({"put"}) _READ_OPS = frozenset({"get", "get_data"}) +# A corrupted wire usually corrupts every row of a batch, so the log is +# capped: the counter in ``HashStats`` carries the magnitude, and the first +# few lines carry the identity of what broke. +_MAX_HASH_MISMATCH_LOGS = 20 + +# Odd 64-bit constant (golden-ratio derived) used to decorrelate the two +# reductions in a row fingerprint before they are combined. Multiplying by +# an odd constant is a bijection mod 2^64, so it mixes without collapsing. +_MIX64 = 0x9E3779B97F4A7C15 + def percentile_from_hist(hist: list[int], q: float) -> float: """Interpolated ``q``-quantile (0-1) from bucket counts. @@ -91,6 +130,82 @@ def percentile_from_hist(hist: list[int], q: float) -> float: return LATENCY_BUCKETS_MS[-1] +def _enc_const1(obj: Any, budget: list[int]) -> int: + return 1 + + +def _enc_float(obj: float, budget: list[int]) -> int: + return 9 + + +def _enc_int(obj: int, budget: list[int]) -> int: + # msgpack packs small ints in a single byte; only wide values cost 9. + if -32 <= obj < 128: + return 1 + if -(2**15) <= obj < 2**16: + return 3 + if -(2**31) <= obj < 2**32: + return 5 + return 9 + + +def _enc_str(obj: str, budget: list[int]) -> int: + n = len(obj) if obj.isascii() else len(obj.encode("utf-8")) + return n + (1 if n < 32 else 2 if n < 256 else 3 if n < 65536 else 5) + + +def _enc_bytes(obj: Any, budget: list[int]) -> int: + n = len(obj) + return n + (2 if n < 256 else 3 if n < 65536 else 5) + + +def _enc_dict(obj: dict[Any, Any], budget: list[int]) -> int: + n = len(obj) + total = 1 if n < 16 else 3 if n < 65536 else 5 + for k, v in obj.items(): + if budget[0] <= 0: + break + budget[0] -= 1 + total += _estimate_encoded_bytes(k, budget) + total += _estimate_encoded_bytes(v, budget) + return total + + +def _enc_seq(obj: Any, budget: list[int]) -> int: + n = len(obj) + total = 1 if n < 16 else 3 if n < 65536 else 5 + for v in obj: + if budget[0] <= 0: + break + budget[0] -= 1 + total += _estimate_encoded_bytes(v, budget) + return total + + +def _enc_tensor(obj: torch.Tensor, budget: list[int]) -> int: + return obj.numel() * obj.element_size() + + +# Exact-type dispatch, tried before any isinstance chain. The common leaves +# come first because insertion order is also the order of the subclass +# fallback below, which only runs for types that miss the exact lookup. +_ENCODERS: dict[type, Callable[[Any, list[int]], int]] = { + str: _enc_str, + int: _enc_int, + bool: _enc_const1, + float: _enc_float, + dict: _enc_dict, + list: _enc_seq, + tuple: _enc_seq, + set: _enc_seq, + bytes: _enc_bytes, + bytearray: _enc_bytes, + memoryview: _enc_bytes, + torch.Tensor: _enc_tensor, + type(None): _enc_const1, +} + + def _estimate_encoded_bytes(obj: Any, budget: list[int]) -> int: """Approximate msgpack-encoded size of a non-tensor object. @@ -101,58 +216,61 @@ def _estimate_encoded_bytes(obj: Any, budget: list[int]) -> int: approximates instead. Container framing (1-5 bytes per element) is not modelled, so treat the result as a lower bound. + Dispatch is an exact-type dict lookup rather than an ``isinstance`` + chain: on the hot path the common leaves (``str``, ``int``) sat 5-7 + branches deep, and the chain cost more than the arithmetic it guarded. + ``budget`` is a single-element list used as a mutable counter that bounds - the walk to ``max_nodes`` visited nodes. Containers stop iterating once it - is exhausted -- summing a generator would otherwise keep walking every - element while each recursive call returned 0, making the cost O(size) - despite the budget. + the walk to ``max_nodes`` container elements. It is decremented only by + the container encoders -- a leaf cannot itself expand the walk, so + charging leaves bought nothing but two list index ops each. Containers + stop iterating once it is exhausted; summing a generator would otherwise + keep walking every element while each recursive call returned 0, making + the cost O(size) despite the budget. """ - if budget[0] <= 0: - return 0 - budget[0] -= 1 - if obj is None or isinstance(obj, bool): - return 1 - if isinstance(obj, int): - # msgpack packs small ints in a single byte; only wide values cost 9. - if -32 <= obj < 128: - return 1 - if -(2**15) <= obj < 2**16: - return 3 - if -(2**31) <= obj < 2**32: - return 5 - return 9 - if isinstance(obj, float): - return 9 - if isinstance(obj, str): - n = len(obj) if obj.isascii() else len(obj.encode("utf-8")) - return n + (1 if n < 32 else 2 if n < 256 else 3 if n < 65536 else 5) - if isinstance(obj, (bytes, bytearray, memoryview)): - n = len(obj) - return n + (2 if n < 256 else 3 if n < 65536 else 5) - if isinstance(obj, dict): - n = len(obj) - total = 1 if n < 16 else 3 if n < 65536 else 5 - for k, v in obj.items(): - if budget[0] <= 0: - break - total += _estimate_encoded_bytes(k, budget) - total += _estimate_encoded_bytes(v, budget) - return total - if isinstance(obj, (list, tuple, set)): - n = len(obj) - total = 1 if n < 16 else 3 if n < 65536 else 5 - for v in obj: - if budget[0] <= 0: - break - total += _estimate_encoded_bytes(v, budget) - return total - if isinstance(obj, torch.Tensor): - return obj.numel() * obj.element_size() + encoder = _ENCODERS.get(obj.__class__) + if encoder is not None: + return encoder(obj, budget) + for typ, encoder in _ENCODERS.items(): + if isinstance(obj, typ): + return encoder(obj, budget) # Unknown type -> pickle/cloudpickle Ext. Cheap proxy; the real size # would need an actual dumps(), which is what we are avoiding. return 64 +# Rows sampled from a NonTensorStack to estimate its payload. The stack +# holds one Python object per batch element, so materialising it (``tolist``) +# and walking every row is O(batch) *per put* -- ~1 ms for a 256-row +# message_log stack. Rows in one stack share a schema, so a strided sample +# extrapolates to within a few percent for a figure that is already +# documented as a lower-bound estimate. +_NONTENSOR_STACK_SAMPLES = 4 + + +def _nontensor_stack_bytes(stack: NonTensorStack, budget: list[int]) -> int: + """Extrapolate a ``NonTensorStack``'s payload from a strided row sample.""" + rows = getattr(stack, "tensordicts", None) + if not rows: + return _estimate_encoded_bytes(stack.tolist(), budget) + n = len(rows) + step = max(1, n // _NONTENSOR_STACK_SAMPLES) + sampled = rows[::step][:_NONTENSOR_STACK_SAMPLES] + sampled_bytes = 0 + for row in sampled: + # Matched by type rather than ``getattr(row, "data", row)``: every + # TensorDictBase carries a ``.data`` property of its own, so the + # duck-typed form would silently hand a nested stack's tensor view to + # the msgpack estimator instead of recursing into its payload. + if isinstance(row, NonTensorData): + sampled_bytes += _estimate_encoded_bytes(row.data, budget) + elif isinstance(row, NonTensorStack): + sampled_bytes += _nontensor_stack_bytes(row, budget) + else: + sampled_bytes += _estimate_encoded_bytes(row, budget) + return sampled_bytes * n // len(sampled) + + def _td_bytes(td: TensorDict | None, max_nodes: int = 10_000) -> int: """Payload bytes of a TensorDict, as the wire will see them. @@ -170,7 +288,17 @@ def _td_bytes(td: TensorDict | None, max_nodes: int = 10_000) -> int: ``leaves_only=True`` would hide the non-tensor entries entirely (``NonTensorData`` is not treated as a leaf), so this walks with - ``leaves_only=False`` and skips container ``TensorDict`` nodes itself. + ``leaves_only=False`` and skips container nodes itself. It walks with + ``items()`` rather than ``keys()`` + ``get()``: ``get()`` re-resolves + each nested key from the root, which measured 2.7x the cost of the + single traversal ``items()`` already performs. + + ``NonTensorData`` and ``NonTensorStack`` are matched by type, not by + ``hasattr``: both are tensorclasses whose attribute misses fall through + a ``__getattr__`` that costs ~2.8 us per probe. The distinction matters + beyond speed -- ``NonTensorData`` exposes BOTH ``.data`` and + ``.tolist()``, and its ``.tolist()`` broadcasts the single stored object + across the batch dim (a 64-row batch reported 20x the real payload). Aliased storage is counted per field: two keys viewing one buffer count twice, which is right for volume (both are serialised) and is what lets @@ -180,31 +308,24 @@ def _td_bytes(td: TensorDict | None, max_nodes: int = 10_000) -> int: return 0 budget = [max_nodes] total = 0 - for k in td.keys(include_nested=True, leaves_only=False): - v = td.get(k) - if isinstance(v, TensorDict): + for _, v in td.items(include_nested=True, leaves_only=False): + if isinstance(v, torch.Tensor): + t = v.values() if v.is_nested else v + total += t.numel() * t.element_size() + elif isinstance(v, NonTensorData): + total += _estimate_encoded_bytes(v.data, budget) + elif isinstance(v, NonTensorStack): + # Checked before TensorDictBase: NonTensorStack subclasses + # LazyStackedTensorDict but carries payload, so skipping it as a + # container would drop those bytes entirely. + total += _nontensor_stack_bytes(v, budget) + elif isinstance(v, TensorDictBase): continue # container; its leaves are visited separately - if not isinstance(v, torch.Tensor): - # Order matters: NonTensorData exposes BOTH .data and .tolist(), - # and .tolist() broadcasts the stored object across the batch dim - # (a 64-row batch reported 20x the real payload). .data is the - # object as stored, so it wins; .tolist() is for NonTensorStack, - # which genuinely holds one value per batch element. - if hasattr(v, "data"): - v = v.data - elif hasattr(v, "tolist"): - v = v.tolist() + else: total += _estimate_encoded_bytes(v, budget) - continue - total += _td_tensor_bytes(v) return total -def _td_tensor_bytes(v: torch.Tensor) -> int: - t = v.values() if v.is_nested else v - return t.numel() * t.element_size() - - def fit_latency_bandwidth(s: dict[str, Any]) -> dict[str, Any]: """Split an op's time into fixed per-request overhead vs transfer. @@ -309,6 +430,26 @@ class OpStats: ) +@dataclass +class HashStats: + """Wire-in / wire-out fingerprint reconciliation. All zero unless enabled. + + ``rows_unverified`` is as important as ``mismatches``: a run that reads + back rows this process never wrote (the normal case for a consumer-side + client, which sees only wire-out) verifies nothing, and a mismatch count + of 0 would otherwise read as "checked and clean". + """ + + rows_recorded: int = 0 + rows_checked: int = 0 + rows_unverified: int = 0 + mismatches: int = 0 + # Leaves that carry no comparable row fingerprint: nested tensors (no + # uniform row shape) and leaves whose leading dim doesn't match the + # sample count, so a row cannot be attributed to a sample id. + fields_skipped: int = 0 + + @dataclass class DataPlaneStats: total_bytes: int = 0 @@ -325,6 +466,7 @@ class DataPlaneStats: # sudden spike in ``max_bytes_per_key_seen``. max_bytes_per_key_seen: int = 0 last_put_bytes_per_key: int = 0 + hash_verify: HashStats = field(default_factory=HashStats) class MetricsDataPlaneClient(DataPlaneClient): @@ -334,14 +476,33 @@ def __init__( self, inner: DataPlaneClient, on_event: Callable[[DataPlaneEvent], None] | None = None, + verify_tensor_hash: bool = False, ) -> None: + """Wrap ``inner``, accumulating per-op timing and volume. + + Args: + inner: The client whose calls are measured. + on_event: Per-op callback. ``None`` (the default) skips + building the event dict entirely — with metrics enabled but + no sink, nothing is paid for a payload nobody reads. + verify_tensor_hash: Record a per-row ``torch.hash_tensor`` + fingerprint on put and re-check it on get. Debug aid, not a + metric: it reads every tensor byte again (~1.2 ms per 18 MB + put), so it is off unless the config asks for it. + """ self._inner = inner - self._on_event = on_event or (lambda _: None) + self._on_event = on_event + self._verify_tensor_hash = verify_tensor_hash self._stats = DataPlaneStats() # Nested per-partition / per-key live byte counts. Populated on # successful ``put_samples``; popped on successful ``clear_samples``. # Bounded by the live key population, not cumulative traffic. self._bytes_by_partition: dict[str, dict[str, int]] = {} + # partition -> sample_id -> field -> wire-in fingerprint. Same + # lifetime as ``_bytes_by_partition``: cleared by ``clear_samples``, + # so it is bounded by the live key population. + self._hash_by_partition: dict[str, dict[str, dict[str, int]]] = {} + self._hash_mismatches_logged = 0 # Previous snapshot, for per-step deltas. Owned here rather than by a # caller: it is this client's prior reading, and keeping it here lets # every trainer use get_step_metrics() without copying the @@ -384,9 +545,7 @@ def snapshot(self) -> dict[str, Any]: # truth for bytes. Distinct from ``bytes_outstanding``, which is # occupancy (what is held) rather than traffic (what moved). by = out["by_op"] - out["bytes_written"] = sum( - by[o]["n_bytes"] for o in _WRITE_OPS if o in by - ) + out["bytes_written"] = sum(by[o]["n_bytes"] for o in _WRITE_OPS if o in by) out["bytes_read"] = sum(by[o]["n_bytes"] for o in _READ_OPS if o in by) out["comm_volume_bytes"] = out["bytes_written"] + out["bytes_read"] return out @@ -412,12 +571,20 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: "wall_s": wall_ms / 1e3, "frac_of_step": (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0, "comm_volume_gb": vol / 1e9, - "bytes_written_gb": ( - snap["bytes_written"] - prev.get("bytes_written", 0) - ) / 1e9, + "bytes_written_gb": (snap["bytes_written"] - prev.get("bytes_written", 0)) + / 1e9, "bytes_read_gb": (snap["bytes_read"] - prev.get("bytes_read", 0)) / 1e9, "bytes_outstanding_gb": snap["bytes_outstanding"] / 1e9, } + if self._verify_tensor_hash: + hv, prev_hv = snap["hash_verify"], prev.get("hash_verify", {}) + metrics["hash/rows_checked"] = hv["rows_checked"] - prev_hv.get( + "rows_checked", 0 + ) + metrics["hash/rows_unverified"] = hv["rows_unverified"] - prev_hv.get( + "rows_unverified", 0 + ) + metrics["hash/mismatches"] = hv["mismatches"] - prev_hv.get("mismatches", 0) prev_ops = prev.get("by_op", {}) for op, st in snap["by_op"].items(): prev_op = prev_ops.get(op, {}) @@ -449,6 +616,13 @@ def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: Called after the underlying RPC succeeds so a failed put never leaves the accounting inflated. + The even split is what makes the loop cheap, and it is not a + simplification: ``n_bytes`` is a whole-batch figure, so there is no + per-key truth to preserve. The division remainder therefore lands on + the first key rather than being spread one-byte-at-a-time across the + batch — spreading it cost an ``enumerate`` and a compare per key + (~40% of this method at 256 keys) to move at most one byte each. + Args: partition_id: Partition the keys were written to. keys: Per-sample uids that were written. @@ -458,9 +632,10 @@ def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: return per_key, remainder = divmod(n_bytes, len(keys)) partition_dict = self._bytes_by_partition.setdefault(partition_id, {}) - for i, key in enumerate(keys): - share = per_key + (1 if i < remainder else 0) - partition_dict[key] = partition_dict.get(key, 0) + share + get_held = partition_dict.get + for key in keys: + partition_dict[key] = get_held(key, 0) + per_key + partition_dict[keys[0]] += remainder self._stats.bytes_outstanding += n_bytes if self._stats.bytes_outstanding > self._stats.peak_bytes_outstanding: self._stats.peak_bytes_outstanding = self._stats.bytes_outstanding @@ -475,6 +650,8 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: partition_id: Partition the keys were dropped from. keys: Uids dropped; ``None`` means the whole partition was cleared. """ + if self._verify_tensor_hash: + self._drop_hashes(partition_id, keys) partition_dict = self._bytes_by_partition.get(partition_id) if partition_dict is None: return @@ -489,6 +666,127 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: del self._bytes_by_partition[partition_id] self._stats.bytes_outstanding -= freed + # ── wire-in / wire-out fingerprinting (opt-in) ───────────────────── + + def _row_fingerprints( + self, td: TensorDict | None, sample_ids: list[str] + ) -> dict[str, list[int]]: + """Per-row ``torch.hash_tensor`` fingerprints, keyed by field name. + + Each tensor leaf is flattened to ``[n_rows, -1]`` and reduced along + the trailing dims, giving one ``uint64`` per sample id — the + granularity the wire actually needs, since a batch put of 256 rows + is routinely read back as eight shards of 32 and a whole-tensor hash + would be incomparable. + + ``hash_tensor`` reduces by XOR, which is commutative — on its own it + is blind to *any* rearrangement of a row's 64-bit words, so a + mis-sharded read that swaps two rows of ``arange``-shaped data hashes + identically. Two things narrow that: the dtype and per-row shape go + in as a salt, and a second reduction over the row's even-index + sub-multiset breaks the permutation symmetry (verified to catch both + row swaps and within-row reversal, which the single reduction + misses). What survives is data whose even-index multiset is also + preserved — a boolean mask, where XOR sees only parity, is the case + this check still cannot see, and is the documented limit. + + Row *i* is attributed to ``sample_ids[i]``, which is the ordering + :meth:`DataPlaneClient.get_samples` already promises ("batched along + ``sample_ids``"). An adapter that reordered rows would show up here + as a mismatch — which is the correct verdict, since a reordered read + is exactly the bug this check exists to catch. + + Leaves that cannot be attributed per row (nested tensors, or a + leading dim that isn't ``len(sample_ids)``) are counted in + ``fields_skipped`` rather than silently dropped. + """ + if td is None: + return {} + n_rows = len(sample_ids) + stats = self._stats.hash_verify + out: dict[str, list[int]] = {} + for key, v in td.items(include_nested=True, leaves_only=True): + if not isinstance(v, torch.Tensor) or v.is_nested or v.ndim < 1: + stats.fields_skipped += 1 + continue + if v.shape[0] != n_rows: + stats.fields_skipped += 1 + continue + # crc32, not the builtin hash(): hash() of a str is salted per + # process, so the fingerprint would not survive being compared + # across ranks — a property this is cheap to keep and expensive + # to rediscover the day someone reduces these across a cluster. + salt = zlib.crc32(f"{v.dtype}|{tuple(v.shape[1:])}".encode()) + flat = v.reshape(n_rows, -1) + digest = torch.hash_tensor(flat, dim=1) ^ salt + if flat.shape[1] > 1: + strided = torch.hash_tensor(flat[:, ::2], dim=1) + digest = digest ^ (strided * _MIX64) + name = key if isinstance(key, str) else ".".join(key) + out[name] = digest.tolist() + return out + + def _record_hashes( + self, partition_id: str, sample_ids: list[str], fields: TensorDict | None + ) -> None: + """Store wire-in fingerprints for a successful put.""" + digests = self._row_fingerprints(fields, sample_ids) + if not digests: + return + partition_hashes = self._hash_by_partition.setdefault(partition_id, {}) + for row, sample_id in enumerate(sample_ids): + per_field = partition_hashes.setdefault(sample_id, {}) + for name, column in digests.items(): + per_field[name] = column[row] + self._stats.hash_verify.rows_recorded += len(sample_ids) + + def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: + """Compare wire-out fingerprints against what was written.""" + if not isinstance(out, TensorDict): + return + digests = self._row_fingerprints(out, sample_ids) + if not digests: + return + partition_hashes = self._hash_by_partition.get(partition_id, {}) + stats = self._stats.hash_verify + for row, sample_id in enumerate(sample_ids): + per_field = partition_hashes.get(sample_id) + if not per_field: + # Written by another process (rollout actor, policy worker): + # this client has no wire-in reading to compare against. + stats.rows_unverified += 1 + continue + stats.rows_checked += 1 + for name, column in digests.items(): + expected = per_field.get(name) + if expected is None or expected == column[row]: + continue + stats.mismatches += 1 + if self._hash_mismatches_logged < _MAX_HASH_MISMATCH_LOGS: + self._hash_mismatches_logged += 1 + logger.error( + "data-plane hash mismatch: partition=%s sample=%s " + "field=%s wire_in=%d wire_out=%d", + partition_id, + sample_id, + name, + expected, + column[row], + ) + + def _drop_hashes(self, partition_id: str, keys: list[str] | None) -> None: + """Release fingerprints alongside the bytes accounting.""" + if keys is None: + self._hash_by_partition.pop(partition_id, None) + return + partition_hashes = self._hash_by_partition.get(partition_id) + if partition_hashes is None: + return + for key in keys: + partition_hashes.pop(key, None) + if not partition_hashes: + del self._hash_by_partition[partition_id] + def _run( self, op: str, @@ -540,43 +838,49 @@ def _emit( status: EventStatus, ) -> None: wall_ms = (monotonic() - t0) * 1000.0 - event: DataPlaneEvent = { - "op": op, - "partition_id": partition_id, - "n_keys": int(n_keys), - "n_bytes": int(n_bytes), - "wall_ms": wall_ms, - "status": status, - } - self._on_event(event) + on_event = self._on_event + if on_event is not None: + # Built lazily: with no sink registered this dict was the single + # most frequent allocation in the wrapper, and nothing read it. + event: DataPlaneEvent = { + "op": op, + "partition_id": partition_id, + "n_keys": n_keys, + "n_bytes": n_bytes, + "wall_ms": wall_ms, + "status": status, + } + on_event(event) # Time is charged for every status: a timeout is often the single # largest contributor, so dropping it would understate the cost. - self._stats.total_wall_ms += wall_ms - bucket = self._stats.by_op.get(op) + stats = self._stats + stats.total_wall_ms += wall_ms + bucket = stats.by_op.get(op) if bucket is None: # Not setdefault(): its default is evaluated eagerly, building a # throwaway OpStats (and its 16-bucket histogram) on every op. - bucket = self._stats.by_op[op] = OpStats() + bucket = stats.by_op[op] = OpStats() bucket.calls += 1 bucket.wall_ms += wall_ms bucket.latency_hist[bisect_left(LATENCY_BUCKETS_MS, wall_ms)] += 1 if status != "ok": bucket.errors += 1 - if status == "ok": - self._stats.total_bytes += n_bytes - self._stats.total_keys += n_keys - self._stats.total_ops += 1 - bucket.n_bytes += n_bytes - bucket.n_keys += n_keys - bucket.ok_wall_ms += wall_ms - bucket.sum_bytes_sq += float(n_bytes) * float(n_bytes) - bucket.sum_bytes_ms += float(n_bytes) * wall_ms - bucket.sum_ms_sq += wall_ms * wall_ms - if op == "put" and n_keys: - per_key = n_bytes // n_keys - self._stats.last_put_bytes_per_key = per_key - if per_key > self._stats.max_bytes_per_key_seen: - self._stats.max_bytes_per_key_seen = per_key + return + stats.total_bytes += n_bytes + stats.total_keys += n_keys + stats.total_ops += 1 + bucket.n_bytes += n_bytes + bucket.n_keys += n_keys + bucket.ok_wall_ms += wall_ms + bytes_f = float(n_bytes) + bucket.sum_bytes_sq += bytes_f * bytes_f + bucket.sum_bytes_ms += bytes_f * wall_ms + bucket.sum_ms_sq += wall_ms * wall_ms + if op == "put" and n_keys: + per_key = n_bytes // n_keys + stats.last_put_bytes_per_key = per_key + if per_key > stats.max_bytes_per_key_seen: + stats.max_bytes_per_key_seen = per_key def register_partition( self, @@ -626,12 +930,15 @@ def claim_meta( ) def get_data(self, meta, select_fields=None): - return self._run( + out = self._run( "get_data", meta.partition_id, lambda: self._inner.get_data(meta, select_fields=select_fields), n_keys=len(meta.sample_ids), ) + if self._verify_tensor_hash: + self._check_hashes(meta.partition_id, meta.sample_ids, out) + return out def check_consumption_status(self, partition_id, task_names): return self._run( @@ -660,19 +967,30 @@ def put_samples(self, sample_ids, partition_id, fields=None, tags=None): n_bytes=n_bytes, ) self._record_put(partition_id, sample_ids_list, n_bytes) + # Fingerprinted after ``_run`` rather than inside it: ``fields`` is + # the caller's TensorDict and the RPC does not mutate it, so hashing + # here keeps the check's own cost out of the op's ``wall_ms``. + if self._verify_tensor_hash: + self._record_hashes(partition_id, sample_ids_list, fields) return out def get_samples(self, sample_ids, partition_id, select_fields): - return self._run( + sample_ids_list = ( + sample_ids if isinstance(sample_ids, list) else list(sample_ids) + ) + out = self._run( "get", partition_id, lambda: self._inner.get_samples( - sample_ids, + sample_ids_list, partition_id, select_fields=select_fields, ), - n_keys=len(sample_ids), + n_keys=len(sample_ids_list), ) + if self._verify_tensor_hash: + self._check_hashes(partition_id, sample_ids_list, out) + return out def list_sample_ids(self, partition_id: str) -> list[str]: return self._run( diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 13a5f59a382..0fa3d8408db 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -23,10 +23,14 @@ import pytest import torch -from tensordict import TensorDict +from tensordict import NonTensorData, NonTensorStack, TensorDict from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient -from nemo_rl.data_plane.observability import MetricsDataPlaneClient +from nemo_rl.data_plane.observability import ( + MetricsDataPlaneClient, + _estimate_encoded_bytes, + _td_bytes, +) from ._rollout_shapes import make_rollout_batch @@ -233,3 +237,293 @@ def test_observability_records_realistic_rollout_put() -> None: min_expected = n * 64 * 8 # input_ids alone assert put_events[0]["n_bytes"] >= min_expected client.close() + + +# ── byte accounting ──────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "name,td,expected", + [ + ("flat", TensorDict({"x": torch.zeros(8, 16)}, batch_size=[8]), 8 * 16 * 4), + ( + "sliced-view", + TensorDict({"x": torch.zeros(8, 32)[:, :8]}, batch_size=[8]), + 8 * 8 * 4, + ), + ( + "transposed", + TensorDict({"x": torch.zeros(4, 8).t()}, batch_size=[8]), + 8 * 4 * 4, + ), + ( + "stride-0-expand", + TensorDict({"x": torch.zeros(8, 1).expand(8, 9)}, batch_size=[8]), + 8 * 9 * 4, + ), + ( + "mixed-dtype", + TensorDict( + { + "i": torch.zeros(8, 4, dtype=torch.int64), + "b": torch.zeros(8, 4, dtype=torch.bool), + "f": torch.zeros(8, 4, dtype=torch.bfloat16), + }, + batch_size=[8], + ), + 8 * 4 * 8 + 8 * 4 * 1 + 8 * 4 * 2, + ), + ( + "nested-container", + TensorDict( + { + "a": torch.zeros(8, 2), + "sub": TensorDict({"b": torch.zeros(8, 3)}, batch_size=[8]), + }, + batch_size=[8], + ), + 8 * 2 * 4 + 8 * 3 * 4, + ), + ("empty", TensorDict({}, batch_size=[8]), 0), + ("none", None, 0), + ], +) +def test_td_bytes_counts_wire_payload(name, td, expected): + """Tensor leaves count as ``contiguous().nbytes``, containers count once.""" + assert _td_bytes(td) == expected + + +def test_td_bytes_nontensordata_is_not_broadcast(): + """``NonTensorData`` holds ONE object; counting it per batch row would + inflate a 64-row put by 64x. Its bytes must not scale with batch size.""" + payload = {"tool": "bash", "text": "x" * 100} + small = TensorDict({"m": NonTensorData(payload, batch_size=[2])}, batch_size=[2]) + large = TensorDict({"m": NonTensorData(payload, batch_size=[64])}, batch_size=[64]) + assert _td_bytes(small) == _td_bytes(large) + assert _td_bytes(small) >= 100 # the string itself is still counted + + +def test_td_bytes_nontensorstack_scales_with_rows(): + """``NonTensorStack`` genuinely holds one object per row, so its estimate + must scale — and stay close to the exact walk it extrapolates from.""" + row = {"turns": ["hello"] * 4, "n": 3} + stack_8 = NonTensorStack(*[NonTensorData(dict(row)) for _ in range(8)]) + stack_64 = NonTensorStack(*[NonTensorData(dict(row)) for _ in range(64)]) + bytes_8 = _td_bytes(TensorDict({"s": stack_8}, batch_size=[8])) + bytes_64 = _td_bytes(TensorDict({"s": stack_64}, batch_size=[64])) + assert bytes_8 > 0 + assert bytes_64 == pytest.approx(8 * bytes_8, rel=0.05) + # And it agrees with summing every row explicitly. + exact = sum(_estimate_encoded_bytes(dict(row), [10_000]) for _ in range(64)) + assert bytes_64 == pytest.approx(exact, rel=0.05) + + +def test_estimate_encoded_bytes_walk_is_bounded(): + """The node budget caps the walk so one pathological payload cannot make + a put O(payload size).""" + huge = {"k": list(range(100_000))} + bounded = _estimate_encoded_bytes(huge, [64]) + unbounded = _estimate_encoded_bytes(huge, [10_000_000]) + assert bounded < unbounded + assert bounded <= 4 * 64 # ≤2 leaves per budget unit, ≤2 bytes each here + + +def test_outstanding_bytes_reconcile_exactly(): + """Put then clear must return ``bytes_outstanding`` to zero: the per-key + split drops its division remainder on one key rather than spreading it, + so the total has to be preserved for the accounting to close.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + ids = [f"u{i}" for i in range(7)] # 7 keys => non-zero remainder + fields = TensorDict({"x": torch.zeros(7, 5)}, batch_size=[7]) + client.register_partition( + partition_id="p", fields=["x"], num_samples=7, consumer_tasks=["t"] + ) + client.put_samples(sample_ids=ids, partition_id="p", fields=fields) + assert client.snapshot()["bytes_outstanding"] == 7 * 5 * 4 + client.clear_samples(sample_ids=ids, partition_id="p") + assert client.snapshot()["bytes_outstanding"] == 0 + client.close() + + +def test_no_callback_still_accumulates_stats(): + """``on_event=None`` skips building the event dict; the counters that + ``snapshot()`` reports must not depend on a sink being registered.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client.register_partition( + partition_id="p", fields=["x"], num_samples=2, consumer_tasks=["t"] + ) + client.put_samples( + sample_ids=["a", "b"], + partition_id="p", + fields=TensorDict({"x": torch.zeros(2, 3)}, batch_size=[2]), + ) + snap = client.snapshot() + assert snap["total_bytes"] == 2 * 3 * 4 + assert snap["by_op"]["put"]["calls"] == 1 + assert snap["total_wall_ms"] > 0 + client.close() + + +# ── wire-in / wire-out hash verification ─────────────────────────────── + + +class _CorruptingClient(NoOpDataPlaneClient): + """Flips one element of one field on read — a stand-in for a wire bug.""" + + def __init__(self, field: str, row: int) -> None: + super().__init__() + self._corrupt_field = field + self._corrupt_row = row + + def get_samples(self, sample_ids, partition_id, select_fields): + out = super().get_samples(sample_ids, partition_id, select_fields) + if self._corrupt_field in out.keys(): + out[self._corrupt_field][self._corrupt_row] += 1 + return out + + +def _hash_client(inner=None): + client = MetricsDataPlaneClient( + inner or NoOpDataPlaneClient(), verify_tensor_hash=True + ) + client.register_partition( + partition_id="p", fields=["ids", "lp"], num_samples=4, consumer_tasks=["t"] + ) + return client + + +def _hash_fields(n=4): + return TensorDict( + { + "ids": torch.arange(n * 6, dtype=torch.int64).reshape(n, 6), + "lp": torch.linspace(0, 1, n * 6, dtype=torch.bfloat16).reshape(n, 6), + }, + batch_size=[n], + ) + + +def test_hash_verification_clean_roundtrip(): + client = _hash_client() + ids = [f"u{i}" for i in range(4)] + client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids", "lp"]) + + hv = client.snapshot()["hash_verify"] + assert hv["rows_recorded"] == 4 + assert hv["rows_checked"] == 4 + assert hv["rows_unverified"] == 0 + assert hv["mismatches"] == 0 + client.close() + + +def test_hash_verification_detects_corruption(): + client = _hash_client(_CorruptingClient(field="ids", row=2)) + ids = [f"u{i}" for i in range(4)] + client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids", "lp"]) + + hv = client.snapshot()["hash_verify"] + assert hv["mismatches"] == 1 + assert client.get_step_metrics(1.0)["hash/mismatches"] == 1 + client.close() + + +def test_hash_verification_survives_shard_readback(): + """A 4-row put read back two rows at a time must still line up: the + fingerprint is per row, not per batch.""" + client = _hash_client() + ids = [f"u{i}" for i in range(4)] + client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + for shard in (ids[:2], ids[2:]): + client.get_samples(sample_ids=shard, partition_id="p", select_fields=["ids"]) + + hv = client.snapshot()["hash_verify"] + assert hv["rows_checked"] == 4 + assert hv["mismatches"] == 0 + client.close() + + +def test_hash_verification_reports_rows_it_never_wrote(): + """A consumer-side client sees only wire-out. Those rows must land in + ``rows_unverified`` — reporting 0 mismatches would read as 'clean'.""" + inner = NoOpDataPlaneClient() + writer = _hash_client(inner) + ids = [f"u{i}" for i in range(4)] + writer.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + + reader = MetricsDataPlaneClient(inner, verify_tensor_hash=True) + reader.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + + hv = reader.snapshot()["hash_verify"] + assert hv["rows_unverified"] == 4 + assert hv["rows_checked"] == 0 + assert hv["mismatches"] == 0 + writer.close() + + +def test_hash_fingerprints_released_on_clear(): + """Fingerprints must be bounded by the live key population, not by + cumulative traffic.""" + client = _hash_client() + ids = [f"u{i}" for i in range(4)] + client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + assert client._hash_by_partition["p"] + client.clear_samples(sample_ids=ids, partition_id="p") + assert client._hash_by_partition == {} + client.close() + + +def test_hash_fingerprint_sees_rearrangement(): + """A plain XOR reduction is commutative and would report a mis-sharded + read (two rows swapped) as clean. The second, strided reduction is what + makes these visible — regress it and shard misalignment goes silent.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + ids = ["a", "b", "c"] + # arange is the adversarial case: its 64-bit words cancel under XOR. + td = TensorDict( + {"x": torch.arange(12, dtype=torch.float32).reshape(3, 4)}, batch_size=[3] + ) + base = client._row_fingerprints(td, ids)["x"] + + swapped = TensorDict({"x": td["x"][[0, 2, 1]]}, batch_size=[3]) + assert client._row_fingerprints(swapped, ids)["x"] != base + + reversed_row = td["x"].clone() + reversed_row[0] = reversed_row[0].flip(0) + within = TensorDict({"x": reversed_row}, batch_size=[3]) + assert client._row_fingerprints(within, ids)["x"] != base + + +def test_hash_fingerprint_separates_dtype_and_shape(): + """The salt must make a reinterpreted payload look different even when + the underlying words are identical.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + ids = ["a", "b"] + values = torch.arange(8, dtype=torch.int64).reshape(2, 4) + as_2x4 = client._row_fingerprints(TensorDict({"x": values}, batch_size=[2]), ids)[ + "x" + ] + as_2x2x2 = client._row_fingerprints( + TensorDict({"x": values.reshape(2, 2, 2)}, batch_size=[2]), ids + )["x"] + as_int32 = client._row_fingerprints( + TensorDict({"x": values.to(torch.int32)}, batch_size=[2]), ids + )["x"] + assert as_2x4 != as_2x2x2 + assert as_2x4 != as_int32 + + +def test_hash_verification_off_by_default(): + """Default construction must do no hashing work at all.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client.register_partition( + partition_id="p", fields=["ids", "lp"], num_samples=4, consumer_tasks=["t"] + ) + ids = [f"u{i}" for i in range(4)] + client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + + assert client.snapshot()["hash_verify"]["rows_recorded"] == 0 + assert client._hash_by_partition == {} + assert "hash/mismatches" not in client.get_step_metrics(1.0) + client.close() diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index f21c7045a7b..7ce25edc050 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -557,6 +557,7 @@ data_plane: # GDR needs that headroom to pay off. observability: # per-op data-plane timing/volume enabled: false # true => wrap client in MetricsDataPlaneClient + verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~1.2ms per 18MB) # Multi-Teacher On-Policy Distillation (MOPD); null = disabled (default). Mirrors # the field on the GRPO MasterConfig added for MOPD support. From 70b8a30e4b6529d9c7d159b36c4326df2b8fbb2e Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 17:03:48 -0700 Subject: [PATCH 05/70] fix(data-plane): fingerprint jagged leaves; simplify the hash to one call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Verifying the guard against the real wire layout found it was not guarding anything. `column_io` packs every per-token field through `codec.pack_jagged_fields` before `put_samples`, so `input_ids`, `generation_logprobs`, `token_mask` and `advantages` all arrive as jagged nested tensors — which `_row_fingerprints` skipped. It reported `fields_skipped=6, mismatches=0`: a clean bill of health for a payload it had never looked at. Eight of ten injected corruptions went undetected. Nested leaves are now fingerprinted by padding to a rectangle and reducing per row. The padding is free rather than approximate: XOR is its own identity on zero and a zero pad bitcasts to a zero 64-bit word for every dtype on this wire (verified int64/int32/bf16/fp32/bool), so a jagged put reconciles against the dense read `_from_wire` hands back when rows happen to be uniform. Also cuts the fingerprint back to a single `torch.hash_tensor` call plus a dtype salt. The previous second, strided reduction doubled the cost to buy detection of a within-row element permutation — a rearrangement nothing on this wire performs. Rows are compared per sample id, so a swap *between* rows never needed it. All ten corruption modes are caught either way. Verification, against pack_jagged_fields in and _from_wire out: false positives 0 over a 500-row randomized soak, every shard grouping from 1 to 256, reversed id order, field subsets, delta writes and overwritten fields false negatives 0 of 10 — single-element change per dtype, truncated row, zeroed row, bf16->fp32, wrong sample, swapped rows `hash/fields_skipped` is now logged per step, because that counter at `0 mismatches` is exactly what this bug looked like from the outside. Dropping the jagged special case in `_td_bytes` sped the metrics path up as well: `numel()` already reports a nested tensor's total element count and costs half of reaching through `.values()`. On the real jagged payload (256 ragged rows, 12 MB) put goes 201us -> 99us and get 162us -> 76us, where the earlier dense-payload benchmark had missed this path entirely. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- examples/configs/grpo_math_1B.yaml | 2 +- nemo_rl/data_plane/README.md | 40 +++--- nemo_rl/data_plane/interfaces.py | 4 +- nemo_rl/data_plane/observability.py | 116 +++++++++++------- tests/unit/data_plane/test_observability.py | 59 +++++---- .../unit/reference_configs/grpo_math_1B.yaml | 2 +- 6 files changed, 138 insertions(+), 85 deletions(-) diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 41df21ba697..03b7f3e1fb0 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -585,7 +585,7 @@ data_plane: # GDR needs that headroom to pay off. observability: # per-op data-plane timing/volume enabled: false # true => wrap client in MetricsDataPlaneClient - verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~1.2ms per 18MB) + verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~8ms per 12MB batch) # Multi-Teacher On-Policy Distillation (MOPD): distills from one or more teacher # models into the policy via token-level teacher-minus-student logprob advantages, diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index e03f386da13..42c83950c28 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -449,27 +449,39 @@ counts sum into one cluster-wide distribution) and byte volume. `snapshot()` returns the cumulative view; `get_step_metrics(step_time_s)` returns the per-step delta already flattened for the logger. -The wrapper is measured against a no-op inner client at 256 keys and an -18 MB payload: **~54 µs per put, ~33 µs per get** — under 0.1% of a 59 ms -operation. Most of that is the byte walk over the `TensorDict`; the rest is -the per-key attribution `clear_samples` needs to undo. `on_event` defaults to -`None` when unset, in which case the per-op event dict is never built. +Measured against a no-op inner client on the payload the wire actually +carries — 256 ragged rows, 12 MB, jagged per-token fields as +`pack_jagged_fields` leaves them: **~99 µs per put, ~76 µs per get**, about +0.15% of a 59 ms operation. Most of that is the byte walk over the +`TensorDict`; the rest is the per-key attribution `clear_samples` needs to +undo. `on_event` defaults to `None` when unset, in which case the per-op +event dict is never built. `verify_tensor_hash: true` additionally records a per-row `torch.hash_tensor` fingerprint on every put and re-checks it on every get, so a tensor that changes between wire-in and wire-out is reported (`hash/mismatches`) instead of being trained on silently. Fingerprints are per row, so a 256-row put read back as eight shards of 32 still reconciles. -Two things to know before turning it on: - -- It reads every tensor byte again on both sides — measured at ~1.2 ms - per 18 MB on put and the same on get, i.e. ~2% of the 59 ms operation it - is guarding. Keep it to debugging runs. -- `hash_tensor` reduces by XOR, so it cannot see a permutation of elements - *within* a row. Dtype and per-row shape are folded into the fingerprint; - element order within a row is not covered. + +Verified against ten corruption modes injected into the round trip — a +single-element change in each dtype, a truncated row, a zeroed row, a +bf16→fp32 precision change, a row served from the wrong sample, and two +rows swapped — all ten are caught, with zero false alarms over a 500-row +randomized soak and every shard grouping from 1 to 256. Three things to +know before turning it on: + +- It reads every tensor byte again on both sides — measured at ~8 ms for a + 12 MB jagged batch, on put and again on get, i.e. ~13% of the 59 ms + operation it is guarding. Keep it to debugging runs. +- `hash_tensor` reduces by XOR, so a row's digest is blind to a + rearrangement of that row's own elements. Rows are compared per sample + id, so a swap *between* rows is still caught; nothing on this wire + reorders *within* a row. - Only rows this process wrote can be checked. A consumer-side client - reports them under `hash/rows_unverified` rather than counting them clean. + reports them under `hash/rows_unverified` rather than counting them + clean, and `hash/fields_skipped` reports any leaf it could not fingerprint + — watch that one, since a guard that quietly stops covering a field still + reports zero mismatches. Backend choice: - **`simple`** — ZMQ-backed; lowest setup overhead. Default for tests diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 1bd3fcf8850..ce6c2206103 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -189,8 +189,8 @@ class ObservabilityConfig(TypedDict): re-checks it, so a value that changes between wire-in and wire-out is reported (``hash/mismatches``) instead of silently training on it. It reads every tensor byte a second time on both sides — budget roughly - 1.2 ms per 18 MB moved, on each side — so leave it off outside of - debugging. + 8 ms for a 12 MB jagged batch, on each side — so leave it off outside + of debugging. """ enabled: bool diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 1464191aeb4..31ce80fa8ea 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -25,19 +25,21 @@ ``peak_bytes_outstanding`` (high-water mark over the run lifetime). Everything here sits on the hot path of every transfer, so the cost is a -design constraint rather than an afterthought: measured against a no-op -inner client at 256 keys and an 18 MB payload, the wrapper adds ~54 us per -put and ~33 us per get — under 0.1% of a 59 ms operation. What that budget -buys is spelled out where it is spent (``_td_bytes``, ``_record_put``, -``_emit``); the short version is that no traversal happens twice, no -allocation happens for a payload nobody reads, and no estimate walks a -structure whose size it can extrapolate. +design constraint rather than an afterthought. Measured against a no-op +inner client on the payload the wire actually carries -- 256 ragged rows, +12 MB, jagged per-token fields as ``codec.pack_jagged_fields`` leaves them +-- the wrapper adds ~99 us per put and ~76 us per get, about 0.15% of a +59 ms operation. What that budget buys is spelled out where it is spent +(``_td_bytes``, ``_record_put``, ``_emit``); the short version is that no +traversal happens twice, no allocation happens for a payload nobody reads, +and no estimate walks a structure whose size it can extrapolate. ``verify_tensor_hash=True`` adds an opt-in correctness check on top: per-row ``torch.hash_tensor`` fingerprints recorded at put and re-checked at get, so a tensor that changes between wire-in and wire-out is reported -rather than trained on. It reads every tensor byte again on both sides, so -it is a debugging tool, not a metric. +rather than trained on. It reads every tensor byte again on both sides +(~8 ms for that same 12 MB payload), so it is a debugging tool, not a +metric. """ from __future__ import annotations @@ -101,10 +103,20 @@ class DataPlaneEvent(TypedDict): # few lines carry the identity of what broke. _MAX_HASH_MISMATCH_LOGS = 20 -# Odd 64-bit constant (golden-ratio derived) used to decorrelate the two -# reductions in a row fingerprint before they are combined. Multiplying by -# an odd constant is a bijection mod 2^64, so it mixes without collapsing. -_MIX64 = 0x9E3779B97F4A7C15 + +def _dtype_salt(dtype: torch.dtype) -> int: + """Fingerprint salt distinguishing dtypes that hash to the same words. + + ``hash_tensor`` upcasts to the 64-bit equivalent before reducing, so a + bf16 tensor and the fp32 tensor holding the same values produce the same + digest. Mixing the dtype in is what makes a precision change visible. + + ``crc32``, not the builtin ``hash()``: ``hash()`` of a ``str`` is salted + per process, so a fingerprint built on it would not survive being + compared across ranks — cheap to keep, expensive to rediscover the day + someone reduces these cluster-wide. + """ + return zlib.crc32(str(dtype).encode()) def percentile_from_hist(hist: list[int], q: float) -> float: @@ -310,8 +322,12 @@ def _td_bytes(td: TensorDict | None, max_nodes: int = 10_000) -> int: total = 0 for _, v in td.items(include_nested=True, leaves_only=False): if isinstance(v, torch.Tensor): - t = v.values() if v.is_nested else v - total += t.numel() * t.element_size() + # No jagged special case: numel() already reports a nested + # tensor's total element count, and asking it directly is half + # the cost of reaching through .values() (16 us vs 32 us per + # jagged field). Every per-token field on this wire is jagged, + # so that is paid four or five times per put. + total += v.numel() * v.element_size() elif isinstance(v, NonTensorData): total += _estimate_encoded_bytes(v.data, budget) elif isinstance(v, NonTensorStack): @@ -487,8 +503,8 @@ def __init__( no sink, nothing is paid for a payload nobody reads. verify_tensor_hash: Record a per-row ``torch.hash_tensor`` fingerprint on put and re-check it on get. Debug aid, not a - metric: it reads every tensor byte again (~1.2 ms per 18 MB - put), so it is off unless the config asks for it. + metric: it reads every tensor byte again (~8 ms for a + 12 MB jagged batch), so it is off unless the config asks. """ self._inner = inner self._on_event = on_event @@ -585,6 +601,13 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: "rows_unverified", 0 ) metrics["hash/mismatches"] = hv["mismatches"] - prev_hv.get("mismatches", 0) + # Logged because a guard that quietly stops covering a field is + # worse than no guard: it reports 0 mismatches and reads as + # clean. A step where this climbs is a step where something + # stopped being checked. + metrics["hash/fields_skipped"] = hv["fields_skipped"] - prev_hv.get( + "fields_skipped", 0 + ) prev_ops = prev.get("by_op", {}) for op, st in snap["by_op"].items(): prev_op = prev_ops.get(op, {}) @@ -679,16 +702,13 @@ def _row_fingerprints( is routinely read back as eight shards of 32 and a whole-tensor hash would be incomparable. - ``hash_tensor`` reduces by XOR, which is commutative — on its own it - is blind to *any* rearrangement of a row's 64-bit words, so a - mis-sharded read that swaps two rows of ``arange``-shaped data hashes - identically. Two things narrow that: the dtype and per-row shape go - in as a salt, and a second reduction over the row's even-index - sub-multiset breaks the permutation symmetry (verified to catch both - row swaps and within-row reversal, which the single reduction - misses). What survives is data whose even-index multiset is also - preserved — a boolean mask, where XOR sees only parity, is the case - this check still cannot see, and is the documented limit. + ``hash_tensor`` reduces by XOR, which is commutative, so a row's + digest is blind to a rearrangement of that row's own elements. Rows + are compared per sample id, so a swap *between* rows is still caught + — it is only a permutation *within* one row that hides, and nothing + on this wire reorders within a row. Measured against ten corruption + modes (single-element change per dtype, truncation, zeroed row, + precision change, wrong sample, swapped rows) it caught all ten. Row *i* is attributed to ``sample_ids[i]``, which is the ordering :meth:`DataPlaneClient.get_samples` already promises ("batched along @@ -696,9 +716,19 @@ def _row_fingerprints( as a mismatch — which is the correct verdict, since a reordered read is exactly the bug this check exists to catch. - Leaves that cannot be attributed per row (nested tensors, or a - leading dim that isn't ``len(sample_ids)``) are counted in - ``fields_skipped`` rather than silently dropped. + Jagged leaves are fingerprinted too, and have to be: ``column_io`` + packs every per-token field (``input_ids``, ``generation_logprobs``, + ``token_mask``, ``advantages``) through + :func:`codec.pack_jagged_fields` before it reaches ``put_samples``, + so skipping nested tensors would leave the entire bulk payload + unguarded while still reporting a clean bill of health. They round + trip asymmetrically — ``_from_wire`` densifies a nested field whose + rows happen to be uniform — so the fingerprint is defined on the + *row*, identically for both layouts. + + Leaves that cannot be attributed per row (a leading dim that isn't + ``len(sample_ids)``, or a nested layout other than ``jagged``) are + counted in ``fields_skipped`` rather than silently dropped. """ if td is None: return {} @@ -706,22 +736,26 @@ def _row_fingerprints( stats = self._stats.hash_verify out: dict[str, list[int]] = {} for key, v in td.items(include_nested=True, leaves_only=True): - if not isinstance(v, torch.Tensor) or v.is_nested or v.ndim < 1: + if not isinstance(v, torch.Tensor) or v.ndim < 1: stats.fields_skipped += 1 continue + salt = _dtype_salt(v.dtype) + if v.is_nested: + if v.layout != torch.jagged: + stats.fields_skipped += 1 + continue + # Padding with zero is free here rather than approximate: + # XOR is its own identity on zero and a zero pad bitcasts to + # a zero 64-bit word for every dtype on this wire (verified + # for int64/int32/bf16/fp32/bool), so the pad contributes + # nothing. That is what lets a jagged put reconcile against + # the dense read `_from_wire` hands back when the rows + # happen to be uniform. + v = torch.nested.to_padded_tensor(v, padding=0) if v.shape[0] != n_rows: stats.fields_skipped += 1 continue - # crc32, not the builtin hash(): hash() of a str is salted per - # process, so the fingerprint would not survive being compared - # across ranks — a property this is cheap to keep and expensive - # to rediscover the day someone reduces these across a cluster. - salt = zlib.crc32(f"{v.dtype}|{tuple(v.shape[1:])}".encode()) - flat = v.reshape(n_rows, -1) - digest = torch.hash_tensor(flat, dim=1) ^ salt - if flat.shape[1] > 1: - strided = torch.hash_tensor(flat[:, ::2], dim=1) - digest = digest ^ (strided * _MIX64) + digest = torch.hash_tensor(v.reshape(n_rows, -1), dim=1) ^ salt name = key if isinstance(key, str) else ".".join(key) out[name] = digest.tolist() return out diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 0fa3d8408db..3943c6d4e2d 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -473,44 +473,51 @@ def test_hash_fingerprints_released_on_clear(): client.close() -def test_hash_fingerprint_sees_rearrangement(): - """A plain XOR reduction is commutative and would report a mis-sharded - read (two rows swapped) as clean. The second, strided reduction is what - makes these visible — regress it and shard misalignment goes silent.""" +def test_hash_fingerprint_covers_jagged_fields(): + """The per-token fields on this wire are jagged by the time they reach + ``put_samples`` (``codec.pack_jagged_fields``). Skipping nested leaves + would leave the entire bulk payload unguarded while still reporting zero + mismatches — a guard that reads as clean because it checked nothing.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) ids = ["a", "b", "c"] - # arange is the adversarial case: its 64-bit words cancel under XOR. - td = TensorDict( - {"x": torch.arange(12, dtype=torch.float32).reshape(3, 4)}, batch_size=[3] + rows = [torch.arange(n, dtype=torch.int64) + n for n in (3, 5, 4)] + jagged = TensorDict( + {"x": torch.nested.nested_tensor(rows, layout=torch.jagged)}, batch_size=[3] ) - base = client._row_fingerprints(td, ids)["x"] + digests = client._row_fingerprints(jagged, ids) + assert "x" in digests, "jagged leaf was skipped" + assert client.snapshot()["hash_verify"]["fields_skipped"] == 0 + assert len(set(digests["x"])) == 3 - swapped = TensorDict({"x": td["x"][[0, 2, 1]]}, batch_size=[3]) - assert client._row_fingerprints(swapped, ids)["x"] != base - reversed_row = td["x"].clone() - reversed_row[0] = reversed_row[0].flip(0) - within = TensorDict({"x": reversed_row}, batch_size=[3]) - assert client._row_fingerprints(within, ids)["x"] != base +def test_hash_fingerprint_matches_across_jagged_and_dense(): + """``_from_wire`` densifies a jagged field whose rows are uniform, so a + jagged put has to reconcile against a dense get. Zero padding is XOR- + neutral, which is what makes the two layouts agree.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + ids = ["a", "b"] + dense = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int64) + jagged = torch.nested.nested_tensor(list(dense.unbind()), layout=torch.jagged) + assert ( + client._row_fingerprints(TensorDict({"x": jagged}, batch_size=[2]), ids)["x"] + == client._row_fingerprints(TensorDict({"x": dense}, batch_size=[2]), ids)["x"] + ) -def test_hash_fingerprint_separates_dtype_and_shape(): - """The salt must make a reinterpreted payload look different even when - the underlying words are identical.""" +def test_hash_fingerprint_separates_dtype(): + """``hash_tensor`` upcasts to 64 bits before reducing, so bf16 and fp32 + holding the same values reduce identically. The dtype salt is the only + thing that makes a precision change visible.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) ids = ["a", "b"] - values = torch.arange(8, dtype=torch.int64).reshape(2, 4) - as_2x4 = client._row_fingerprints(TensorDict({"x": values}, batch_size=[2]), ids)[ + values = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + as_fp32 = client._row_fingerprints(TensorDict({"x": values}, batch_size=[2]), ids)[ "x" ] - as_2x2x2 = client._row_fingerprints( - TensorDict({"x": values.reshape(2, 2, 2)}, batch_size=[2]), ids - )["x"] - as_int32 = client._row_fingerprints( - TensorDict({"x": values.to(torch.int32)}, batch_size=[2]), ids + as_bf16 = client._row_fingerprints( + TensorDict({"x": values.to(torch.bfloat16)}, batch_size=[2]), ids )["x"] - assert as_2x4 != as_2x2x2 - assert as_2x4 != as_int32 + assert as_fp32 != as_bf16 def test_hash_verification_off_by_default(): diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 7ce25edc050..c68ed4815b9 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -557,7 +557,7 @@ data_plane: # GDR needs that headroom to pay off. observability: # per-op data-plane timing/volume enabled: false # true => wrap client in MetricsDataPlaneClient - verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~1.2ms per 18MB) + verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~8ms per 12MB batch) # Multi-Teacher On-Policy Distillation (MOPD); null = disabled (default). Mirrors # the field on the GRPO MasterConfig added for MOPD support. From 57fe7238322c8cd0df667fc26ce41ecf8ecd44dd Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 17:20:41 -0700 Subject: [PATCH 06/70] perf(data-plane): scope jagged fingerprints to the buffer, not a rectangle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-row fingerprints for jagged leaves meant padding each field out to a rectangle before hashing, and that rectangle is sized by the longest row, not by the payload. On a realistically ragged batch it was 3.5x the real data and cost 13x more than hashing the values buffer (5378us vs 405us) — a bill that grows with how ragged the batch is, for a field the digest only has to answer "did anything change" about. Now split by what each layout can do cheaply: rectangular leaf one digest per row (hash_tensor(dim=1)) — names the sample that diverged, survives shard reads jagged leaf one digest over the values buffer, XORed per row with that row's length, marked batch-scoped verify_tensor_hash on a 12 MB jagged batch: 8.0ms -> 2.4ms per put, and the cost now tracks real bytes instead of padding waste. The scheme has to follow the field, not the layout in hand. A field packed jagged comes back *dense* whenever its rows are uniform, because `_from_wire` densifies those — so choosing per layout made the two sides compute different things and every row of a uniform batch reported a mismatch (940 false alarms across the soak). `_check_hashes` now replays the scheme recorded at put; the values buffer of a uniform jagged field and the dense tensor it densifies into are the same elements in the same order, so the digests agree by construction. A shard read of a batch-scoped field reports unverified rather than mismatching. Also bitcasts every leaf to a same-width integer before hashing. `hash_tensor` has no float8 kernel — it raises NotImplementedError on float8_e4m3fn, which propagated straight out of put_samples and would have taken the transfer down with it. The view is free and the dtype still enters through the salt. Re-verified against the same injected corruptions, zero false alarms: caught single-element change in every dtype, truncated row, zeroed row, bf16->fp32, row served from the wrong sample clean 500-row soak, shard groupings 1..256, reversed id order, field subsets, delta writes, overwritten fields limit mis-shard on a jagged field is caught only when the swapped rows differ in length: 58/60 on ragged data, never on a uniform-length batch. Rectangular fields catch it always. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- examples/configs/grpo_math_1B.yaml | 2 +- nemo_rl/data_plane/README.md | 49 +++-- nemo_rl/data_plane/interfaces.py | 4 +- nemo_rl/data_plane/observability.py | 189 ++++++++++++------ tests/unit/data_plane/test_observability.py | 78 ++++++-- .../unit/reference_configs/grpo_math_1B.yaml | 2 +- 6 files changed, 224 insertions(+), 100 deletions(-) diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 03b7f3e1fb0..282ac2a5aa8 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -585,7 +585,7 @@ data_plane: # GDR needs that headroom to pay off. observability: # per-op data-plane timing/volume enabled: false # true => wrap client in MetricsDataPlaneClient - verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~8ms per 12MB batch) + verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~2.4ms per 12MB batch) # Multi-Teacher On-Policy Distillation (MOPD): distills from one or more teacher # models into the policy via token-level teacher-minus-student logprob advantages, diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 42c83950c28..5e10590725e 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -457,26 +457,35 @@ carries — 256 ragged rows, 12 MB, jagged per-token fields as undo. `on_event` defaults to `None` when unset, in which case the per-op event dict is never built. -`verify_tensor_hash: true` additionally records a per-row -`torch.hash_tensor` fingerprint on every put and re-checks it on every get, -so a tensor that changes between wire-in and wire-out is reported -(`hash/mismatches`) instead of being trained on silently. Fingerprints are -per row, so a 256-row put read back as eight shards of 32 still reconciles. - -Verified against ten corruption modes injected into the round trip — a -single-element change in each dtype, a truncated row, a zeroed row, a -bf16→fp32 precision change, a row served from the wrong sample, and two -rows swapped — all ten are caught, with zero false alarms over a 500-row -randomized soak and every shard grouping from 1 to 256. Three things to -know before turning it on: - -- It reads every tensor byte again on both sides — measured at ~8 ms for a - 12 MB jagged batch, on put and again on get, i.e. ~13% of the 59 ms - operation it is guarding. Keep it to debugging runs. -- `hash_tensor` reduces by XOR, so a row's digest is blind to a - rearrangement of that row's own elements. Rows are compared per sample - id, so a swap *between* rows is still caught; nothing on this wire - reorders *within* a row. +`verify_tensor_hash: true` additionally records a `torch.hash_tensor` +fingerprint on every put and re-checks it on every get, so a tensor that +changes between wire-in and wire-out is reported (`hash/mismatches`) +instead of being trained on silently. + +Two granularities, because torch has no ragged hash kernel: + +| leaf | digest | scope | +|---|---|---| +| rectangular (`rewards`, `input_lengths`) | one per row, `hash_tensor(..., dim=1)` | per sample id; survives shard reads | +| jagged (`input_ids`, `generation_logprobs`, `token_mask`, `advantages`) | one over the values buffer, XORed per row with that row's length | per batch; a shard read reports unverified | + +Giving the jagged fields per-row digests would mean padding each one out to +a rectangle first. On a realistically ragged batch that rectangle is 3.5× +the real payload and cost 13× more, to answer a question the buffer digest +already answers. + +Verified by injecting corruption into the round trip. Caught: a +single-element change in every dtype, a truncated row, a zeroed row, a +bf16→fp32 precision change, and a row served from the wrong sample — with +**zero false alarms** over a 500-row randomized soak, every shard grouping +from 1 to 256, reversed id order, field subsets and delta writes. Known +limits, measured rather than assumed: + +- A **mis-shard** (two rows swapped) is caught on a jagged field only when + the two rows differ in length — 58/60 on ragged rollout data, never on a + uniform-length batch. Rectangular fields catch it unconditionally. +- It reads every tensor byte again on both sides — ~2.4 ms for a 12 MB + jagged batch, on put and again on get. Keep it to debugging runs. - Only rows this process wrote can be checked. A consumer-side client reports them under `hash/rows_unverified` rather than counting them clean, and `hash/fields_skipped` reports any leaf it could not fingerprint diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index ce6c2206103..d00b88d26b2 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -189,8 +189,8 @@ class ObservabilityConfig(TypedDict): re-checks it, so a value that changes between wire-in and wire-out is reported (``hash/mismatches``) instead of silently training on it. It reads every tensor byte a second time on both sides — budget roughly - 8 ms for a 12 MB jagged batch, on each side — so leave it off outside - of debugging. + 2.4 ms for a 12 MB jagged batch, on each side — so leave it off + outside of debugging. """ enabled: bool diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 31ce80fa8ea..9578a37a414 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -35,11 +35,14 @@ and no estimate walks a structure whose size it can extrapolate. ``verify_tensor_hash=True`` adds an opt-in correctness check on top: -per-row ``torch.hash_tensor`` fingerprints recorded at put and re-checked -at get, so a tensor that changes between wire-in and wire-out is reported -rather than trained on. It reads every tensor byte again on both sides -(~8 ms for that same 12 MB payload), so it is a debugging tool, not a -metric. +``torch.hash_tensor`` fingerprints recorded at put and re-checked at get, +so a tensor that changes between wire-in and wire-out is reported rather +than trained on. Rectangular leaves get a digest per row; jagged leaves get +one over their values buffer, because torch has no ragged hash kernel and +padding a ragged batch out to a rectangle costs more than the payload is +worth (see :meth:`MetricsDataPlaneClient._row_fingerprints`). It reads +every tensor byte again on both sides (~2.4 ms for that same 12 MB +payload), so it is a debugging tool, not a metric. """ from __future__ import annotations @@ -50,7 +53,7 @@ from dataclasses import asdict, dataclass, field from pathlib import Path from time import monotonic -from typing import Any, Callable, Literal, TypedDict +from typing import Any, Callable, Collection, Literal, NamedTuple, TypedDict EventStatus = Literal["ok", "error", "timeout"] @@ -104,6 +107,35 @@ class DataPlaneEvent(TypedDict): _MAX_HASH_MISMATCH_LOGS = 20 +class _FieldDigest(NamedTuple): + """One fingerprint per row, plus how far it can be trusted. + + ``batch_scoped`` means the values were derived from the whole batch's + buffer, so they only reconcile against a read of that same batch. A + shard of it computes a different buffer digest and must be reported + unverified rather than as a mismatch. + """ + + per_row: list[int] + batch_scoped: bool + + +# Same-width signed integer for each tensor element size, used to bitcast a +# leaf before hashing. +_INT_VIEW_BY_WIDTH = {1: torch.int8, 2: torch.int16, 4: torch.int32, 8: torch.int64} + + +def _as_int_view(t: torch.Tensor) -> torch.Tensor: + """Bitcast a tensor to a same-width integer type, or pass it through. + + ``hash_tensor`` has no float8 kernel and raises ``NotImplementedError`` + there; viewing the bytes as integers sidesteps every dtype-specific + kernel for free, since ``view`` on a same-width dtype is metadata only. + """ + int_dtype = _INT_VIEW_BY_WIDTH.get(t.element_size()) + return t.view(int_dtype) if int_dtype is not None else t + + def _dtype_salt(dtype: torch.dtype) -> int: """Fingerprint salt distinguishing dtypes that hash to the same words. @@ -503,7 +535,7 @@ def __init__( no sink, nothing is paid for a payload nobody reads. verify_tensor_hash: Record a per-row ``torch.hash_tensor`` fingerprint on put and re-check it on get. Debug aid, not a - metric: it reads every tensor byte again (~8 ms for a + metric: it reads every tensor byte again (~2.4 ms for a 12 MB jagged batch), so it is off unless the config asks. """ self._inner = inner @@ -518,6 +550,10 @@ def __init__( # lifetime as ``_bytes_by_partition``: cleared by ``clear_samples``, # so it is bounded by the live key population. self._hash_by_partition: dict[str, dict[str, dict[str, int]]] = {} + # partition -> field -> row count, for the jagged fields whose digest + # covers a whole buffer and so only reconciles against a read of that + # same batch. One int per field, not per row. + self._batch_scoped_rows: dict[str, dict[str, int]] = {} self._hash_mismatches_logged = 0 # Previous snapshot, for per-step deltas. Owned here rather than by a # caller: it is this client's prior reading, and keeping it here lets @@ -692,39 +728,48 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: # ── wire-in / wire-out fingerprinting (opt-in) ───────────────────── def _row_fingerprints( - self, td: TensorDict | None, sample_ids: list[str] - ) -> dict[str, list[int]]: - """Per-row ``torch.hash_tensor`` fingerprints, keyed by field name. - - Each tensor leaf is flattened to ``[n_rows, -1]`` and reduced along - the trailing dims, giving one ``uint64`` per sample id — the - granularity the wire actually needs, since a batch put of 256 rows - is routinely read back as eight shards of 32 and a whole-tensor hash - would be incomparable. - - ``hash_tensor`` reduces by XOR, which is commutative, so a row's - digest is blind to a rearrangement of that row's own elements. Rows - are compared per sample id, so a swap *between* rows is still caught - — it is only a permutation *within* one row that hides, and nothing - on this wire reorders within a row. Measured against ten corruption - modes (single-element change per dtype, truncation, zeroed row, - precision change, wrong sample, swapped rows) it caught all ten. - - Row *i* is attributed to ``sample_ids[i]``, which is the ordering - :meth:`DataPlaneClient.get_samples` already promises ("batched along - ``sample_ids``"). An adapter that reordered rows would show up here - as a mismatch — which is the correct verdict, since a reordered read - is exactly the bug this check exists to catch. - - Jagged leaves are fingerprinted too, and have to be: ``column_io`` - packs every per-token field (``input_ids``, ``generation_logprobs``, - ``token_mask``, ``advantages``) through - :func:`codec.pack_jagged_fields` before it reaches ``put_samples``, - so skipping nested tensors would leave the entire bulk payload - unguarded while still reporting a clean bill of health. They round - trip asymmetrically — ``_from_wire`` densifies a nested field whose - rows happen to be uniform — so the fingerprint is defined on the - *row*, identically for both layouts. + self, + td: TensorDict | None, + sample_ids: list[str], + batch_scoped_fields: Collection[str] = (), + ) -> dict[str, _FieldDigest]: + """``torch.hash_tensor`` fingerprints for each tensor leaf. + + Two granularities, because the wire has two layouts and only one of + them can be reduced per row cheaply: + + * **Rectangular leaf** — ``hash_tensor(..., dim=1)`` gives one + ``uint64`` per sample id directly, for the cost of a single pass. + That is the granularity worth having: it names *which* sample + diverged, and it survives a batch being read back in shards. + * **Jagged leaf** — ``hash_tensor`` has no ragged kernel, so a + per-row digest would mean padding each field out to a rectangle + first. On a realistically ragged batch that rectangle is 3.5x the + real payload and measured 13x the cost of hashing the values + buffer, for a field this only has to answer "did anything change". + So a jagged leaf gets one digest over its whole values buffer, + XORed per row with that row's length, and is marked + ``batch_scoped``. + + Both are still ``torch.hash_tensor``; the jagged one just reduces + the flat buffer instead of a rectangle it had to build first. + + Leaves are bitcast to a same-width integer type before hashing. + ``hash_tensor`` has no float8 kernel — it raises + ``NotImplementedError`` on ``float8_e4m3fn``, which would propagate + straight out of ``put_samples`` and take the transfer down with it. + The bitcast is free (a view), makes every dtype hashable, and the + dtype still enters through the salt. + + The scheme has to follow the *field*, not the layout in hand: a + field packed jagged on put comes back **dense** whenever its rows + happen to be uniform, because ``_from_wire`` densifies those. Left + to pick per layout, the two sides compute different things and every + row of a uniform batch reports a mismatch. ``batch_scoped_fields`` + carries the scheme chosen at put so the read reproduces it — the + values buffer of a uniform jagged field and the flattened dense + tensor it densifies into are the same elements in the same order, so + the digests agree by construction. Leaves that cannot be attributed per row (a leading dim that isn't ``len(sample_ids)``, or a nested layout other than ``jagged``) are @@ -734,30 +779,40 @@ def _row_fingerprints( return {} n_rows = len(sample_ids) stats = self._stats.hash_verify - out: dict[str, list[int]] = {} + out: dict[str, _FieldDigest] = {} for key, v in td.items(include_nested=True, leaves_only=True): if not isinstance(v, torch.Tensor) or v.ndim < 1: stats.fields_skipped += 1 continue salt = _dtype_salt(v.dtype) + name = key if isinstance(key, str) else ".".join(key) if v.is_nested: if v.layout != torch.jagged: stats.fields_skipped += 1 continue - # Padding with zero is free here rather than approximate: - # XOR is its own identity on zero and a zero pad bitcasts to - # a zero 64-bit word for every dtype on this wire (verified - # for int64/int32/bf16/fp32/bool), so the pad contributes - # nothing. That is what lets a jagged put reconcile against - # the dense read `_from_wire` hands back when the rows - # happen to be uniform. - v = torch.nested.to_padded_tensor(v, padding=0) - if v.shape[0] != n_rows: + offsets = v.offsets() + if offsets.numel() - 1 != n_rows: + stats.fields_skipped += 1 + continue + buffer, lengths = v.values(), (offsets[1:] - offsets[:-1]).tolist() + elif v.shape[0] != n_rows: stats.fields_skipped += 1 continue - digest = torch.hash_tensor(v.reshape(n_rows, -1), dim=1) ^ salt - name = key if isinstance(key, str) else ".".join(key) - out[name] = digest.tolist() + elif name in batch_scoped_fields: + buffer = v + lengths = [v.shape[1] if v.ndim >= 2 else 1] * n_rows + else: + flat = _as_int_view(v.reshape(n_rows, -1)) + out[name] = _FieldDigest( + (torch.hash_tensor(flat, dim=1) ^ salt).tolist(), + batch_scoped=False, + ) + continue + flat_buffer = _as_int_view(buffer.reshape(1, -1)) + buffer_digest = torch.hash_tensor(flat_buffer, dim=1).tolist()[0] ^ salt + out[name] = _FieldDigest( + [buffer_digest ^ length for length in lengths], batch_scoped=True + ) return out def _record_hashes( @@ -770,18 +825,32 @@ def _record_hashes( partition_hashes = self._hash_by_partition.setdefault(partition_id, {}) for row, sample_id in enumerate(sample_ids): per_field = partition_hashes.setdefault(sample_id, {}) - for name, column in digests.items(): - per_field[name] = column[row] + for name, digest in digests.items(): + per_field[name] = digest.per_row[row] + batch_rows = self._batch_scoped_rows.setdefault(partition_id, {}) + for name, digest in digests.items(): + if digest.batch_scoped: + batch_rows[name] = len(sample_ids) self._stats.hash_verify.rows_recorded += len(sample_ids) def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: """Compare wire-out fingerprints against what was written.""" if not isinstance(out, TensorDict): return - digests = self._row_fingerprints(out, sample_ids) + batch_rows = self._batch_scoped_rows.get(partition_id, {}) + digests = self._row_fingerprints(out, sample_ids, batch_rows.keys()) if not digests: return partition_hashes = self._hash_by_partition.get(partition_id, {}) + # A batch-scoped digest covers the whole buffer it was reduced from, + # so it only means anything against a read of that same batch. Drop + # those fields on a shard read rather than reporting every row of it + # as a mismatch. + comparable = { + name: digest + for name, digest in digests.items() + if not digest.batch_scoped or batch_rows.get(name) == len(sample_ids) + } stats = self._stats.hash_verify for row, sample_id in enumerate(sample_ids): per_field = partition_hashes.get(sample_id) @@ -791,9 +860,9 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N stats.rows_unverified += 1 continue stats.rows_checked += 1 - for name, column in digests.items(): + for name, digest in comparable.items(): expected = per_field.get(name) - if expected is None or expected == column[row]: + if expected is None or expected == digest.per_row[row]: continue stats.mismatches += 1 if self._hash_mismatches_logged < _MAX_HASH_MISMATCH_LOGS: @@ -805,13 +874,14 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N sample_id, name, expected, - column[row], + digest.per_row[row], ) def _drop_hashes(self, partition_id: str, keys: list[str] | None) -> None: """Release fingerprints alongside the bytes accounting.""" if keys is None: self._hash_by_partition.pop(partition_id, None) + self._batch_scoped_rows.pop(partition_id, None) return partition_hashes = self._hash_by_partition.get(partition_id) if partition_hashes is None: @@ -820,6 +890,7 @@ def _drop_hashes(self, partition_id: str, keys: list[str] | None) -> None: partition_hashes.pop(key, None) if not partition_hashes: del self._hash_by_partition[partition_id] + self._batch_scoped_rows.pop(partition_id, None) def _run( self, diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 3943c6d4e2d..69436e19265 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -473,41 +473,72 @@ def test_hash_fingerprints_released_on_clear(): client.close() +def _jagged(rows): + return TensorDict( + {"x": torch.nested.nested_tensor(rows, layout=torch.jagged)}, + batch_size=[len(rows)], + ) + + def test_hash_fingerprint_covers_jagged_fields(): """The per-token fields on this wire are jagged by the time they reach ``put_samples`` (``codec.pack_jagged_fields``). Skipping nested leaves would leave the entire bulk payload unguarded while still reporting zero mismatches — a guard that reads as clean because it checked nothing.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) - ids = ["a", "b", "c"] rows = [torch.arange(n, dtype=torch.int64) + n for n in (3, 5, 4)] - jagged = TensorDict( - {"x": torch.nested.nested_tensor(rows, layout=torch.jagged)}, batch_size=[3] - ) - digests = client._row_fingerprints(jagged, ids) - assert "x" in digests, "jagged leaf was skipped" + digest = client._row_fingerprints(_jagged(rows), ["a", "b", "c"])["x"] + assert client.snapshot()["hash_verify"]["fields_skipped"] == 0 - assert len(set(digests["x"])) == 3 + assert digest.batch_scoped, "jagged digests only reconcile per batch" + assert len(digest.per_row) == 3 + # A jagged digest carries the row length, so a change in any row's + # payload moves every row's value and a length change moves one. + changed = list(rows) + changed[1] = changed[1] + 1 + assert client._row_fingerprints(_jagged(changed), ["a", "b", "c"])["x"] != digest def test_hash_fingerprint_matches_across_jagged_and_dense(): """``_from_wire`` densifies a jagged field whose rows are uniform, so a - jagged put has to reconcile against a dense get. Zero padding is XOR- - neutral, which is what makes the two layouts agree.""" + jagged put has to reconcile against a dense get. + + Regression guard: picking the scheme from the layout in hand rather than + from what was recorded made every row of a uniform batch report a + mismatch — 940 false alarms over the verification soak. + """ client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) ids = ["a", "b"] dense = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int64) - jagged = torch.nested.nested_tensor(list(dense.unbind()), layout=torch.jagged) - assert ( - client._row_fingerprints(TensorDict({"x": jagged}, batch_size=[2]), ids)["x"] - == client._row_fingerprints(TensorDict({"x": dense}, batch_size=[2]), ids)["x"] + jagged = _jagged(list(dense.unbind())) + + put_side = client._row_fingerprints(jagged, ids)["x"] + get_side = client._row_fingerprints( + TensorDict({"x": dense}, batch_size=[2]), ids, batch_scoped_fields={"x"} + )["x"] + assert put_side == get_side + + +def test_hash_shard_read_of_jagged_field_is_unverified_not_a_mismatch(): + """A batch-scoped digest covers the whole buffer, so a shard read cannot + reproduce it. That has to report as unverified — reporting it as a + mismatch would make the guard cry wolf on every sharded fetch.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + ids = [f"u{i}" for i in range(4)] + rows = [torch.arange(3, dtype=torch.int64) + i for i in range(4)] + client.register_partition( + partition_id="p", fields=["x"], num_samples=4, consumer_tasks=["t"] ) + client.put_samples(sample_ids=ids, partition_id="p", fields=_jagged(rows)) + client.get_samples(sample_ids=ids[:2], partition_id="p", select_fields=["x"]) + + assert client.snapshot()["hash_verify"]["mismatches"] == 0 + client.close() def test_hash_fingerprint_separates_dtype(): - """``hash_tensor`` upcasts to 64 bits before reducing, so bf16 and fp32 - holding the same values reduce identically. The dtype salt is the only - thing that makes a precision change visible.""" + """The values reduce identically once bitcast, so only the dtype salt + makes a precision change visible.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) ids = ["a", "b"] values = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) @@ -517,7 +548,20 @@ def test_hash_fingerprint_separates_dtype(): as_bf16 = client._row_fingerprints( TensorDict({"x": values.to(torch.bfloat16)}, batch_size=[2]), ids )["x"] - assert as_fp32 != as_bf16 + assert as_fp32.per_row != as_bf16.per_row + + +def test_hash_fingerprint_handles_float8(): + """``hash_tensor`` has no float8 kernel. Without the integer bitcast the + ``NotImplementedError`` propagates out of ``put_samples`` and takes the + transfer down with it.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + fp8 = TensorDict( + {"x": torch.tensor([[1.0, 2.0], [3.0, 4.0]]).to(torch.float8_e4m3fn)}, + batch_size=[2], + ) + digest = client._row_fingerprints(fp8, ["a", "b"])["x"] + assert len(digest.per_row) == 2 def test_hash_verification_off_by_default(): diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index c68ed4815b9..0ffc106e878 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -557,7 +557,7 @@ data_plane: # GDR needs that headroom to pay off. observability: # per-op data-plane timing/volume enabled: false # true => wrap client in MetricsDataPlaneClient - verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~8ms per 12MB batch) + verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~2.4ms per 12MB batch) # Multi-Teacher On-Policy Distillation (MOPD); null = disabled (default). Mirrors # the field on the GRPO MasterConfig added for MOPD support. From 55e00aacd719349b26fbb846133efab1f2607254 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 17:27:27 -0700 Subject: [PATCH 07/70] fix(data-plane): count fields dropped as incomparable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Asking what the hash check actually catches, per leaf type, turned up a third silent hole. `_check_hashes` drops a field whose recorded scheme cannot be reproduced by the read, which is correct — reporting it as a row-level mismatch would cry wolf on every sharded fetch — but it dropped it without counting, so the report still read clean. It fires for real divergence, not just shard reads: a rectangular put that comes back jagged (truncating one row makes the batch ragged) has no comparable digest, so a genuine truncation was being swallowed. That case now lands in `hash/fields_skipped`, which is already logged per step. This is the same shape as the original bug — a field quietly stops being compared and zero mismatches reads as a clean bill of health. Every path that declines to check something now increments a counter. Also documents that detection is not attribution. A jagged digest covers the whole values buffer, so any change moves every row's value: it says "this batch is wrong", never "this sample is wrong". Measured on an 8-row batch, one element changed in u3: jagged leaf caught, flags all 8 rows rectangular leaf caught, names u3 `pack_jagged_fields` leaves ~94% of the payload jagged, so batch-level is the normal resolution and the README now says so. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 22 +++++++++++- nemo_rl/data_plane/observability.py | 20 +++++++---- tests/unit/data_plane/test_observability.py | 37 +++++++++++++++++++++ 3 files changed, 71 insertions(+), 8 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 5e10590725e..d285ea652f8 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -474,6 +474,23 @@ a rectangle first. On a realistically ragged batch that rectangle is 3.5× the real payload and cost 13× more, to answer a question the buffer digest already answers. +**Detection is not attribution, and the difference is the whole point of +the split.** The same corruption, injected into an 8-row batch: + +| corruption | jagged leaf | rectangular leaf | +|---|---|---| +| 1 element changed in `u3` | caught, flags **all 8 rows** | caught, names **`u3`** | +| `u5` zeroed | caught, flags all 8 rows | caught, names `u5` | +| `u3`↔`u4` swapped | caught only if their lengths differ | caught, names `u3`,`u4` | +| nothing | clean | clean | + +A jagged digest covers the whole values buffer, so any change moves every +row's value: it says *this batch is wrong*, never *this sample is wrong*. +Since `pack_jagged_fields` leaves ~94% of the payload jagged (only +`rewards` and `input_lengths` stay rectangular), that is the normal +resolution — you learn a step's transfer diverged and have to bisect for +the row yourself. + Verified by injecting corruption into the round trip. Caught: a single-element change in every dtype, a truncated row, a zeroed row, a bf16→fp32 precision change, and a row served from the wrong sample — with @@ -486,9 +503,12 @@ limits, measured rather than assumed: uniform-length batch. Rectangular fields catch it unconditionally. - It reads every tensor byte again on both sides — ~2.4 ms for a 12 MB jagged batch, on put and again on get. Keep it to debugging runs. +- A rectangular field that comes back **jagged** (one row truncated makes + the batch ragged) has no comparable digest and is dropped — counted in + `hash/fields_skipped`, not reported as a mismatch. - Only rows this process wrote can be checked. A consumer-side client reports them under `hash/rows_unverified` rather than counting them - clean, and `hash/fields_skipped` reports any leaf it could not fingerprint + clean, and `hash/fields_skipped` reports any leaf it could not compare — watch that one, since a guard that quietly stops covering a field still reports zero mismatches. diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 9578a37a414..a0b512d2fc0 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -842,16 +842,22 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N if not digests: return partition_hashes = self._hash_by_partition.get(partition_id, {}) + stats = self._stats.hash_verify # A batch-scoped digest covers the whole buffer it was reduced from, # so it only means anything against a read of that same batch. Drop # those fields on a shard read rather than reporting every row of it - # as a mismatch. - comparable = { - name: digest - for name, digest in digests.items() - if not digest.batch_scoped or batch_rows.get(name) == len(sample_ids) - } - stats = self._stats.hash_verify + # as a mismatch — but *count* the drop. Dropping silently is the + # exact shape of the bug that made this check pass while covering + # nothing: a field stops being compared and the report still reads + # clean. It also fires when a rectangular put comes back jagged + # (one row truncated makes the batch ragged), which is a real + # divergence this cannot express as a row-level mismatch. + comparable = {} + for name, digest in digests.items(): + if not digest.batch_scoped or batch_rows.get(name) == len(sample_ids): + comparable[name] = digest + else: + stats.fields_skipped += 1 for row, sample_id in enumerate(sample_ids): per_field = partition_hashes.get(sample_id) if not per_field: diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 69436e19265..b006d229455 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -536,6 +536,43 @@ def test_hash_shard_read_of_jagged_field_is_unverified_not_a_mismatch(): client.close() +def test_hash_incomparable_field_is_counted_not_dropped(): + """A rectangular put can come back jagged — truncating one row makes the + batch ragged. Its per-row digests are not comparable against a + batch-scoped read, so the field is dropped; dropping it *silently* is + the exact shape of the bug that let this check pass while covering + nothing, so the drop has to land in ``fields_skipped``.""" + store = _RaggedOnReadClient() + client = MetricsDataPlaneClient(store, verify_tensor_hash=True) + ids = [f"u{i}" for i in range(4)] + dense = TensorDict( + {"x": torch.arange(16, dtype=torch.int64).reshape(4, 4)}, batch_size=[4] + ) + client.register_partition( + partition_id="p", fields=["x"], num_samples=4, consumer_tasks=["t"] + ) + client.put_samples(sample_ids=ids, partition_id="p", fields=dense) + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["x"]) + + hv = client.snapshot()["hash_verify"] + assert hv["mismatches"] == 0, "must not cry wolf on an incomparable field" + assert hv["fields_skipped"] == 1, "the drop must be visible" + assert client.get_step_metrics(1.0)["hash/fields_skipped"] == 1 + + +class _RaggedOnReadClient(NoOpDataPlaneClient): + """Returns one row shorter than it was written — a truncation on the wire.""" + + def get_samples(self, sample_ids, partition_id, select_fields): + out = super().get_samples(sample_ids, partition_id, select_fields) + rows = list(out["x"].unbind()) + rows[1] = rows[1][:-1] + return TensorDict( + {"x": torch.nested.nested_tensor(rows, layout=torch.jagged)}, + batch_size=[len(sample_ids)], + ) + + def test_hash_fingerprint_separates_dtype(): """The values reduce identically once bitcast, so only the dtype salt makes a precision change visible.""" From 0b83e2b0c128d93b5317b091e6eb97cbdc0add1c Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 17:48:07 -0700 Subject: [PATCH 08/70] refactor(data-plane): trim the observability diff after review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cleanup pass over the metrics/verification work. No behaviour change: the corruption matrix still catches every mode it did, the false-alarm soak is still clean, and put/get stay at 99us/77us on the jagged payload. observability.py 1119 -> 1036 lines, prose:code 0.62 -> 0.58 net -96 lines - Reverts the exact-type `_ENCODERS` dispatch table to the isinstance chain it replaced. The table only helped exact-type hits while making every subclass and unknown type walk 13 isinstance checks instead of 8 — and `_NONTENSOR_STACK_SAMPLES` had already collapsed the call count the table was optimising per-call cost for. Removes 10 module-level names and ~45 lines. - Shares `_pop_partition_keys` between the byte accounting and the fingerprint store, which were running near-identical teardown 170 lines apart in lockstep by inspection only. `_drop_hashes` disappears. - Collapses `_batch_scoped_rows` (partition -> field -> row count, every int in an inner dict identical) to `_batch_scope` (partition -> (field names, row count)). - One `_as_list` for the three copy-pasted `sample_ids` coercions. - Logs `hash/rows_recorded`, which was write-only outside tests. - Tests reuse the existing `_hash_client()` helper. Trims docstrings that had grown into commit messages: benchmark numbers measured on one machine that silently rot, and prose duplicated verbatim from the README added in the same series. The traps stay — NonTensorData .tolist() broadcasting, NonTensorStack subclassing LazyStackedTensorDict, setdefault's eager default — since those are what a reader would trip on. Two review findings were raised and not acted on: - Making jagged digests row-local would delete the batch-scoping concept, but the claim that batch-scoping loses coverage on sharded reads does not hold: the sharded get_samples is in worker_mixin, a process that never wrote those rows and so reports rows_unverified either way. The driver reads back the same full meta it wrote. Row-local slicing also measured slower than the padding it would replace. - `_as_int_view` duplicates weight_transfer_sparse_codec.integer_view, which cannot be imported here: it pulls in ray via the vllm config chain, and data_plane is imported eagerly and tested in a Ray-free venv. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 407 ++++++++------------ tests/unit/data_plane/test_observability.py | 11 +- 2 files changed, 161 insertions(+), 257 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index a0b512d2fc0..ff6825180d4 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -24,25 +24,15 @@ bytes currently held in TQ, i.e. put minus cleared) and ``peak_bytes_outstanding`` (high-water mark over the run lifetime). -Everything here sits on the hot path of every transfer, so the cost is a -design constraint rather than an afterthought. Measured against a no-op -inner client on the payload the wire actually carries -- 256 ragged rows, -12 MB, jagged per-token fields as ``codec.pack_jagged_fields`` leaves them --- the wrapper adds ~99 us per put and ~76 us per get, about 0.15% of a -59 ms operation. What that budget buys is spelled out where it is spent -(``_td_bytes``, ``_record_put``, ``_emit``); the short version is that no -traversal happens twice, no allocation happens for a payload nobody reads, -and no estimate walks a structure whose size it can extrapolate. - -``verify_tensor_hash=True`` adds an opt-in correctness check on top: +Every method here runs on the hot path of a transfer, so nothing traverses +a structure twice and nothing is allocated for a payload no callback reads. + +``verify_tensor_hash=True`` adds an opt-in correctness check: ``torch.hash_tensor`` fingerprints recorded at put and re-checked at get, so a tensor that changes between wire-in and wire-out is reported rather -than trained on. Rectangular leaves get a digest per row; jagged leaves get -one over their values buffer, because torch has no ragged hash kernel and -padding a ragged batch out to a rectangle costs more than the payload is -worth (see :meth:`MetricsDataPlaneClient._row_fingerprints`). It reads -every tensor byte again on both sides (~2.4 ms for that same 12 MB -payload), so it is a debugging tool, not a metric. +than trained on. It reads every tensor byte again on both sides, so it is a +debugging tool, not a metric. See ``README.md`` for what it does and does +not catch. """ from __future__ import annotations @@ -126,27 +116,52 @@ class _FieldDigest(NamedTuple): def _as_int_view(t: torch.Tensor) -> torch.Tensor: - """Bitcast a tensor to a same-width integer type, or pass it through. + """Bitcast to a same-width integer type, or pass through. - ``hash_tensor`` has no float8 kernel and raises ``NotImplementedError`` - there; viewing the bytes as integers sidesteps every dtype-specific - kernel for free, since ``view`` on a same-width dtype is metadata only. + ``hash_tensor`` has no float8 kernel and raises there; viewing the bytes + as integers sidesteps every dtype-specific kernel for free. """ int_dtype = _INT_VIEW_BY_WIDTH.get(t.element_size()) return t.view(int_dtype) if int_dtype is not None else t -def _dtype_salt(dtype: torch.dtype) -> int: - """Fingerprint salt distinguishing dtypes that hash to the same words. +def _as_list(sample_ids: Any) -> Any: + """Materialize ``sample_ids`` once; ``None`` passes through. + + ``_run`` consumes its lambda and the accounting needs the same sequence + afterwards, so a generator would be exhausted by the time it is indexed. + """ + if sample_ids is None or isinstance(sample_ids, list): + return sample_ids + return list(sample_ids) - ``hash_tensor`` upcasts to the 64-bit equivalent before reducing, so a - bf16 tensor and the fp32 tensor holding the same values produce the same - digest. Mixing the dtype in is what makes a precision change visible. + +def _pop_partition_keys( + store: dict[str, dict[str, Any]], partition_id: str, keys: list[str] | None +) -> list[Any]: + """Drop ``keys`` from ``store[partition_id]``, returning what was removed. + + ``keys=None`` drops the whole partition. Shared by the byte accounting + and the fingerprint store so their teardown cannot drift apart. + """ + partition = store.get(partition_id) + if partition is None: + return [] + if keys is None: + del store[partition_id] + return list(partition.values()) + removed = [partition.pop(key) for key in keys if key in partition] + if not partition: + del store[partition_id] + return removed + + +def _dtype_salt(dtype: torch.dtype) -> int: + """Salt distinguishing dtypes whose values reduce to the same words. ``crc32``, not the builtin ``hash()``: ``hash()`` of a ``str`` is salted - per process, so a fingerprint built on it would not survive being - compared across ranks — cheap to keep, expensive to rediscover the day - someone reduces these cluster-wide. + per process, so the fingerprint would not survive being compared across + ranks. """ return zlib.crc32(str(dtype).encode()) @@ -174,82 +189,6 @@ def percentile_from_hist(hist: list[int], q: float) -> float: return LATENCY_BUCKETS_MS[-1] -def _enc_const1(obj: Any, budget: list[int]) -> int: - return 1 - - -def _enc_float(obj: float, budget: list[int]) -> int: - return 9 - - -def _enc_int(obj: int, budget: list[int]) -> int: - # msgpack packs small ints in a single byte; only wide values cost 9. - if -32 <= obj < 128: - return 1 - if -(2**15) <= obj < 2**16: - return 3 - if -(2**31) <= obj < 2**32: - return 5 - return 9 - - -def _enc_str(obj: str, budget: list[int]) -> int: - n = len(obj) if obj.isascii() else len(obj.encode("utf-8")) - return n + (1 if n < 32 else 2 if n < 256 else 3 if n < 65536 else 5) - - -def _enc_bytes(obj: Any, budget: list[int]) -> int: - n = len(obj) - return n + (2 if n < 256 else 3 if n < 65536 else 5) - - -def _enc_dict(obj: dict[Any, Any], budget: list[int]) -> int: - n = len(obj) - total = 1 if n < 16 else 3 if n < 65536 else 5 - for k, v in obj.items(): - if budget[0] <= 0: - break - budget[0] -= 1 - total += _estimate_encoded_bytes(k, budget) - total += _estimate_encoded_bytes(v, budget) - return total - - -def _enc_seq(obj: Any, budget: list[int]) -> int: - n = len(obj) - total = 1 if n < 16 else 3 if n < 65536 else 5 - for v in obj: - if budget[0] <= 0: - break - budget[0] -= 1 - total += _estimate_encoded_bytes(v, budget) - return total - - -def _enc_tensor(obj: torch.Tensor, budget: list[int]) -> int: - return obj.numel() * obj.element_size() - - -# Exact-type dispatch, tried before any isinstance chain. The common leaves -# come first because insertion order is also the order of the subclass -# fallback below, which only runs for types that miss the exact lookup. -_ENCODERS: dict[type, Callable[[Any, list[int]], int]] = { - str: _enc_str, - int: _enc_int, - bool: _enc_const1, - float: _enc_float, - dict: _enc_dict, - list: _enc_seq, - tuple: _enc_seq, - set: _enc_seq, - bytes: _enc_bytes, - bytearray: _enc_bytes, - memoryview: _enc_bytes, - torch.Tensor: _enc_tensor, - type(None): _enc_const1, -} - - def _estimate_encoded_bytes(obj: Any, budget: list[int]) -> int: """Approximate msgpack-encoded size of a non-tensor object. @@ -260,24 +199,52 @@ def _estimate_encoded_bytes(obj: Any, budget: list[int]) -> int: approximates instead. Container framing (1-5 bytes per element) is not modelled, so treat the result as a lower bound. - Dispatch is an exact-type dict lookup rather than an ``isinstance`` - chain: on the hot path the common leaves (``str``, ``int``) sat 5-7 - branches deep, and the chain cost more than the arithmetic it guarded. - - ``budget`` is a single-element list used as a mutable counter that bounds - the walk to ``max_nodes`` container elements. It is decremented only by - the container encoders -- a leaf cannot itself expand the walk, so - charging leaves bought nothing but two list index ops each. Containers - stop iterating once it is exhausted; summing a generator would otherwise - keep walking every element while each recursive call returned 0, making - the cost O(size) despite the budget. + ``budget`` bounds the walk to ``max_nodes`` container elements. Only the + container branches charge it: a leaf cannot itself expand the walk. + Containers stop iterating once it is exhausted -- summing a generator + would otherwise keep walking every element while each recursive call + returned 0, making the cost O(size) despite the budget. """ - encoder = _ENCODERS.get(obj.__class__) - if encoder is not None: - return encoder(obj, budget) - for typ, encoder in _ENCODERS.items(): - if isinstance(obj, typ): - return encoder(obj, budget) + if obj is None or isinstance(obj, bool): + return 1 + if isinstance(obj, int): + # msgpack packs small ints in a single byte; only wide values cost 9. + if -32 <= obj < 128: + return 1 + if -(2**15) <= obj < 2**16: + return 3 + if -(2**31) <= obj < 2**32: + return 5 + return 9 + if isinstance(obj, float): + return 9 + if isinstance(obj, str): + n = len(obj) if obj.isascii() else len(obj.encode("utf-8")) + return n + (1 if n < 32 else 2 if n < 256 else 3 if n < 65536 else 5) + if isinstance(obj, (bytes, bytearray, memoryview)): + n = len(obj) + return n + (2 if n < 256 else 3 if n < 65536 else 5) + if isinstance(obj, dict): + n = len(obj) + total = 1 if n < 16 else 3 if n < 65536 else 5 + for k, v in obj.items(): + if budget[0] <= 0: + break + budget[0] -= 1 + total += _estimate_encoded_bytes(k, budget) + total += _estimate_encoded_bytes(v, budget) + return total + if isinstance(obj, (list, tuple, set)): + n = len(obj) + total = 1 if n < 16 else 3 if n < 65536 else 5 + for v in obj: + if budget[0] <= 0: + break + budget[0] -= 1 + total += _estimate_encoded_bytes(v, budget) + return total + if isinstance(obj, torch.Tensor): + return obj.numel() * obj.element_size() # Unknown type -> pickle/cloudpickle Ext. Cheap proxy; the real size # would need an actual dumps(), which is what we are avoiding. return 64 @@ -285,10 +252,9 @@ def _estimate_encoded_bytes(obj: Any, budget: list[int]) -> int: # Rows sampled from a NonTensorStack to estimate its payload. The stack # holds one Python object per batch element, so materialising it (``tolist``) -# and walking every row is O(batch) *per put* -- ~1 ms for a 256-row -# message_log stack. Rows in one stack share a schema, so a strided sample -# extrapolates to within a few percent for a figure that is already -# documented as a lower-bound estimate. +# and walking every row is O(batch) *per put*. Sampling assumes rows are +# exchangeable in size, which is only approximately true -- rollout rows +# differ in length by construction -- so this is a model, not a measurement. _NONTENSOR_STACK_SAMPLES = 4 @@ -319,30 +285,21 @@ def _td_bytes(td: TensorDict | None, max_nodes: int = 10_000) -> int: """Payload bytes of a TensorDict, as the wire will see them. Tensor leaves count ``numel * element_size``, which equals - ``t.contiguous().nbytes`` -- the size mooncake registers and sends - (``mooncake_client`` calls ``.contiguous()`` before taking the pointer). - Verified equal for contiguous, sliced, transposed and stride-0 expanded - views, and across bf16/bool/int64/fp8. - + ``t.contiguous().nbytes`` -- the size mooncake registers and sends. Non-tensor leaves are estimated with :func:`_estimate_encoded_bytes`, - since TQ ships them over a separate msgpack path -- omitting them would - undercount communication volume by whatever metadata rides along. Both - kinds are counted in a *single* pass: a second traversal would double the - per-put walk of a structure that can hold hundreds of keys. + since TQ ships them over a separate msgpack path. Both kinds are counted + in a single ``items()`` pass; ``keys()`` + ``get()`` would re-resolve + every nested key from the root. ``leaves_only=True`` would hide the non-tensor entries entirely (``NonTensorData`` is not treated as a leaf), so this walks with - ``leaves_only=False`` and skips container nodes itself. It walks with - ``items()`` rather than ``keys()`` + ``get()``: ``get()`` re-resolves - each nested key from the root, which measured 2.7x the cost of the - single traversal ``items()`` already performs. - - ``NonTensorData`` and ``NonTensorStack`` are matched by type, not by - ``hasattr``: both are tensorclasses whose attribute misses fall through - a ``__getattr__`` that costs ~2.8 us per probe. The distinction matters - beyond speed -- ``NonTensorData`` exposes BOTH ``.data`` and - ``.tolist()``, and its ``.tolist()`` broadcasts the single stored object - across the batch dim (a 64-row batch reported 20x the real payload). + ``leaves_only=False`` and skips container nodes itself. + + ``NonTensorData`` and ``NonTensorStack`` are matched by type rather than + ``hasattr``, and the distinction matters: ``NonTensorData`` exposes BOTH + ``.data`` and ``.tolist()``, and its ``.tolist()`` broadcasts the single + stored object across the batch dim (a 64-row batch reported 20x the real + payload). Aliased storage is counted per field: two keys viewing one buffer count twice, which is right for volume (both are serialised) and is what lets @@ -355,10 +312,8 @@ def _td_bytes(td: TensorDict | None, max_nodes: int = 10_000) -> int: for _, v in td.items(include_nested=True, leaves_only=False): if isinstance(v, torch.Tensor): # No jagged special case: numel() already reports a nested - # tensor's total element count, and asking it directly is half - # the cost of reaching through .values() (16 us vs 32 us per - # jagged field). Every per-token field on this wire is jagged, - # so that is paid four or five times per put. + # tensor's total element count, and costs half of reaching + # through .values(). total += v.numel() * v.element_size() elif isinstance(v, NonTensorData): total += _estimate_encoded_bytes(v.data, budget) @@ -550,10 +505,10 @@ def __init__( # lifetime as ``_bytes_by_partition``: cleared by ``clear_samples``, # so it is bounded by the live key population. self._hash_by_partition: dict[str, dict[str, dict[str, int]]] = {} - # partition -> field -> row count, for the jagged fields whose digest - # covers a whole buffer and so only reconciles against a read of that - # same batch. One int per field, not per row. - self._batch_scoped_rows: dict[str, dict[str, int]] = {} + # partition -> (batch-scoped field names, row count at put). Those + # digests cover a whole buffer, so they only reconcile against a read + # of that same batch; the row count is what detects a shard read. + self._batch_scope: dict[str, tuple[frozenset[str], int]] = {} self._hash_mismatches_logged = 0 # Previous snapshot, for per-step deltas. Owned here rather than by a # caller: it is this client's prior reading, and keeping it here lets @@ -633,6 +588,9 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: metrics["hash/rows_checked"] = hv["rows_checked"] - prev_hv.get( "rows_checked", 0 ) + metrics["hash/rows_recorded"] = hv["rows_recorded"] - prev_hv.get( + "rows_recorded", 0 + ) metrics["hash/rows_unverified"] = hv["rows_unverified"] - prev_hv.get( "rows_unverified", 0 ) @@ -675,12 +633,9 @@ def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: Called after the underlying RPC succeeds so a failed put never leaves the accounting inflated. - The even split is what makes the loop cheap, and it is not a - simplification: ``n_bytes`` is a whole-batch figure, so there is no - per-key truth to preserve. The division remainder therefore lands on - the first key rather than being spread one-byte-at-a-time across the - batch — spreading it cost an ``enumerate`` and a compare per key - (~40% of this method at 256 keys) to move at most one byte each. + ``n_bytes`` is a whole-batch figure, so there is no per-key truth to + preserve: the split is even and the division remainder lands on the + first key rather than being spread across the batch. Args: partition_id: Partition the keys were written to. @@ -710,20 +665,11 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: keys: Uids dropped; ``None`` means the whole partition was cleared. """ if self._verify_tensor_hash: - self._drop_hashes(partition_id, keys) - partition_dict = self._bytes_by_partition.get(partition_id) - if partition_dict is None: - return - if keys is None: - freed = sum(partition_dict.values()) - del self._bytes_by_partition[partition_id] - else: - freed = 0 - for key in keys: - freed += partition_dict.pop(key, 0) - if not partition_dict: - del self._bytes_by_partition[partition_id] - self._stats.bytes_outstanding -= freed + _pop_partition_keys(self._hash_by_partition, partition_id, keys) + if keys is None: + self._batch_scope.pop(partition_id, None) + freed = _pop_partition_keys(self._bytes_by_partition, partition_id, keys) + self._stats.bytes_outstanding -= sum(freed) # ── wire-in / wire-out fingerprinting (opt-in) ───────────────────── @@ -735,45 +681,33 @@ def _row_fingerprints( ) -> dict[str, _FieldDigest]: """``torch.hash_tensor`` fingerprints for each tensor leaf. - Two granularities, because the wire has two layouts and only one of - them can be reduced per row cheaply: - - * **Rectangular leaf** — ``hash_tensor(..., dim=1)`` gives one - ``uint64`` per sample id directly, for the cost of a single pass. - That is the granularity worth having: it names *which* sample - diverged, and it survives a batch being read back in shards. - * **Jagged leaf** — ``hash_tensor`` has no ragged kernel, so a - per-row digest would mean padding each field out to a rectangle - first. On a realistically ragged batch that rectangle is 3.5x the - real payload and measured 13x the cost of hashing the values - buffer, for a field this only has to answer "did anything change". - So a jagged leaf gets one digest over its whole values buffer, - XORed per row with that row's length, and is marked - ``batch_scoped``. - - Both are still ``torch.hash_tensor``; the jagged one just reduces - the flat buffer instead of a rectangle it had to build first. - - Leaves are bitcast to a same-width integer type before hashing. - ``hash_tensor`` has no float8 kernel — it raises - ``NotImplementedError`` on ``float8_e4m3fn``, which would propagate - straight out of ``put_samples`` and take the transfer down with it. - The bitcast is free (a view), makes every dtype hashable, and the - dtype still enters through the salt. - - The scheme has to follow the *field*, not the layout in hand: a - field packed jagged on put comes back **dense** whenever its rows - happen to be uniform, because ``_from_wire`` densifies those. Left - to pick per layout, the two sides compute different things and every - row of a uniform batch reports a mismatch. ``batch_scoped_fields`` - carries the scheme chosen at put so the read reproduces it — the - values buffer of a uniform jagged field and the flattened dense - tensor it densifies into are the same elements in the same order, so - the digests agree by construction. - - Leaves that cannot be attributed per row (a leading dim that isn't - ``len(sample_ids)``, or a nested layout other than ``jagged``) are - counted in ``fields_skipped`` rather than silently dropped. + A rectangular leaf reduces per row (``dim=1``), which names the + sample that diverged. ``hash_tensor`` has no ragged kernel, so a + jagged leaf instead gets one digest over its whole values buffer, + XORed per row with that row's length, and is marked + ``batch_scoped``; padding it out to a rectangle to get per-row + digests costs far more than the answer is worth. ``README.md`` has + the resulting detection/attribution table. + + Args: + td: Leaves to fingerprint; ``None`` yields an empty result. + sample_ids: Row *i* is attributed to ``sample_ids[i]``, the + ordering :meth:`DataPlaneClient.get_samples` promises. + batch_scoped_fields: Fields the *put* side reduced batch-scoped. + The scheme must follow the field, not the layout in hand: a + field packed jagged comes back dense whenever its rows are + uniform (``_from_wire`` densifies those), and choosing per + layout makes the two sides compute different things. The + values buffer of a uniform jagged field and the flattened + dense tensor it densifies into hold the same elements in the + same order, so replaying the recorded scheme agrees by + construction. + + Returns: + Field name -> :class:`_FieldDigest`. Leaves that cannot be + attributed per row (a leading dim that isn't ``len(sample_ids)``, + or a non-``jagged`` nested layout) are counted in + ``fields_skipped`` rather than silently dropped. """ if td is None: return {} @@ -827,18 +761,19 @@ def _record_hashes( per_field = partition_hashes.setdefault(sample_id, {}) for name, digest in digests.items(): per_field[name] = digest.per_row[row] - batch_rows = self._batch_scoped_rows.setdefault(partition_id, {}) - for name, digest in digests.items(): - if digest.batch_scoped: - batch_rows[name] = len(sample_ids) + scoped = frozenset(n for n, d in digests.items() if d.batch_scoped) + if scoped: + self._batch_scope[partition_id] = (scoped, len(sample_ids)) self._stats.hash_verify.rows_recorded += len(sample_ids) def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: """Compare wire-out fingerprints against what was written.""" if not isinstance(out, TensorDict): return - batch_rows = self._batch_scoped_rows.get(partition_id, {}) - digests = self._row_fingerprints(out, sample_ids, batch_rows.keys()) + scoped_names, scoped_rows = self._batch_scope.get( + partition_id, (frozenset(), 0) + ) + digests = self._row_fingerprints(out, sample_ids, scoped_names) if not digests: return partition_hashes = self._hash_by_partition.get(partition_id, {}) @@ -854,7 +789,7 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N # divergence this cannot express as a row-level mismatch. comparable = {} for name, digest in digests.items(): - if not digest.batch_scoped or batch_rows.get(name) == len(sample_ids): + if not digest.batch_scoped or scoped_rows == len(sample_ids): comparable[name] = digest else: stats.fields_skipped += 1 @@ -883,21 +818,6 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N digest.per_row[row], ) - def _drop_hashes(self, partition_id: str, keys: list[str] | None) -> None: - """Release fingerprints alongside the bytes accounting.""" - if keys is None: - self._hash_by_partition.pop(partition_id, None) - self._batch_scoped_rows.pop(partition_id, None) - return - partition_hashes = self._hash_by_partition.get(partition_id) - if partition_hashes is None: - return - for key in keys: - partition_hashes.pop(key, None) - if not partition_hashes: - del self._hash_by_partition[partition_id] - self._batch_scoped_rows.pop(partition_id, None) - def _run( self, op: str, @@ -951,8 +871,7 @@ def _emit( wall_ms = (monotonic() - t0) * 1000.0 on_event = self._on_event if on_event is not None: - # Built lazily: with no sink registered this dict was the single - # most frequent allocation in the wrapper, and nothing read it. + # Built lazily: with no sink registered nothing reads this dict. event: DataPlaneEvent = { "op": op, "partition_id": partition_id, @@ -1062,9 +981,7 @@ def put_samples(self, sample_ids, partition_id, fields=None, tags=None): n_bytes = _td_bytes(fields) # Materialize once: ``_run`` consumes its lambda and we also need # to attribute bytes per sample after success. - sample_ids_list = ( - sample_ids if isinstance(sample_ids, list) else list(sample_ids) - ) + sample_ids_list = _as_list(sample_ids) out = self._run( "put", partition_id, @@ -1086,9 +1003,7 @@ def put_samples(self, sample_ids, partition_id, fields=None, tags=None): return out def get_samples(self, sample_ids, partition_id, select_fields): - sample_ids_list = ( - sample_ids if isinstance(sample_ids, list) else list(sample_ids) - ) + sample_ids_list = _as_list(sample_ids) out = self._run( "get", partition_id, @@ -1111,11 +1026,7 @@ def list_sample_ids(self, partition_id: str) -> list[str]: ) def clear_samples(self, sample_ids, partition_id): - sample_ids_list = ( - sample_ids - if (sample_ids is None or isinstance(sample_ids, list)) - else list(sample_ids) - ) + sample_ids_list = _as_list(sample_ids) n_keys = len(sample_ids_list) if sample_ids_list is not None else 0 self._run( "clear", diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index b006d229455..e196c0c7ab5 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -523,12 +523,9 @@ def test_hash_shard_read_of_jagged_field_is_unverified_not_a_mismatch(): """A batch-scoped digest covers the whole buffer, so a shard read cannot reproduce it. That has to report as unverified — reporting it as a mismatch would make the guard cry wolf on every sharded fetch.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + client = _hash_client() ids = [f"u{i}" for i in range(4)] rows = [torch.arange(3, dtype=torch.int64) + i for i in range(4)] - client.register_partition( - partition_id="p", fields=["x"], num_samples=4, consumer_tasks=["t"] - ) client.put_samples(sample_ids=ids, partition_id="p", fields=_jagged(rows)) client.get_samples(sample_ids=ids[:2], partition_id="p", select_fields=["x"]) @@ -542,15 +539,11 @@ def test_hash_incomparable_field_is_counted_not_dropped(): batch-scoped read, so the field is dropped; dropping it *silently* is the exact shape of the bug that let this check pass while covering nothing, so the drop has to land in ``fields_skipped``.""" - store = _RaggedOnReadClient() - client = MetricsDataPlaneClient(store, verify_tensor_hash=True) + client = _hash_client(_RaggedOnReadClient()) ids = [f"u{i}" for i in range(4)] dense = TensorDict( {"x": torch.arange(16, dtype=torch.int64).reshape(4, 4)}, batch_size=[4] ) - client.register_partition( - partition_id="p", fields=["x"], num_samples=4, consumer_tasks=["t"] - ) client.put_samples(sample_ids=ids, partition_id="p", fields=dense) client.get_samples(sample_ids=ids, partition_id="p", select_fields=["x"]) From 040d459c713a349512c711b8a97491c8aaaadf8e Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 17:55:10 -0700 Subject: [PATCH 09/70] perf(data-plane): size nested leaves from their packed values buffer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A nested tensor's ``nbytes``, ``numel`` and ``element_size`` all route through ``__torch_function__`` and measure ~16 us per leaf; the packed values buffer answers the same question in ~0.1 us. Every per-token field on this wire is nested, so that was paid four or five times per transfer and dominated the whole wrapper. Against the PR head, on 256 ragged rows / 12 MB: put 205 us -> 37 us (5.5x) get 168 us -> 15 us (11x) The public attribute is not an escape: ``v.nbytes`` on a nested tensor measured 16.1 us, indistinguishable from the arithmetic it would replace, because it dispatches the same way. Guarded twice, because a wrong byte count is worse than a slow one — it feeds ``max_bytes_per_key_seen``, which exists to spot wire regressions, so an overcount reads as the very thing it is watching for. An absent ``_values`` falls back to the public attribute. A buffer holding more elements than the offsets describe means the tensor views a larger allocation, which ``torch.nested.narrow`` produces; nothing in nemo_rl builds one (all four construction sites use the compacting ``as_nested_tensor``) but ``_td_bytes`` is handed whatever a caller passes. Both paths have a regression test. Byte accounting verified unchanged across contiguous, sliced, transposed, stride-0 expanded, bf16/bool/int64/float8, jagged, old strided nested and narrow-view leaves. Also documents that `enabled: true` installs `log_event` by default, so turning metrics on currently emits one logger.info per data-plane op. Left as-is since the config documents that default, but it makes the lazy event-dict build in `_emit` dead unless a caller passes on_event=None. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 13 +++++---- nemo_rl/data_plane/observability.py | 31 ++++++++++++++++++--- tests/unit/data_plane/test_observability.py | 27 ++++++++++++++++++ 3 files changed, 62 insertions(+), 9 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index d285ea652f8..a46f0c1f1c6 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -451,11 +451,14 @@ per-step delta already flattened for the logger. Measured against a no-op inner client on the payload the wire actually carries — 256 ragged rows, 12 MB, jagged per-token fields as -`pack_jagged_fields` leaves them: **~99 µs per put, ~76 µs per get**, about -0.15% of a 59 ms operation. Most of that is the byte walk over the -`TensorDict`; the rest is the per-key attribution `clear_samples` needs to -undo. `on_event` defaults to `None` when unset, in which case the per-op -event dict is never built. +`pack_jagged_fields` leaves them: **~37 µs per put, ~15 µs per get**, under +0.1% of a 59 ms operation. What is left is dominated by the per-key +attribution `clear_samples` needs to undo. + +Note that `enabled: true` also installs `log_event` as the default +`on_event` sink, which emits one `logger.info` line per data-plane op. Pass +`observability.callback` explicitly (or `None`) if you want the counters +without the per-op log. `verify_tensor_hash: true` additionally records a `torch.hash_tensor` fingerprint on every put and re-checks it on every get, so a tensor that diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index ff6825180d4..d25cc95d147 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -156,6 +156,32 @@ def _pop_partition_keys( return removed +def _tensor_bytes(v: torch.Tensor) -> int: + """Wire bytes of one tensor leaf, rectangular or nested. + + A nested tensor's own ``nbytes`` (and ``numel``/``element_size``) route + through ``__torch_function__`` and measure ~16 us per leaf; its packed + values buffer answers the same question in ~0.1 us. Every per-token + field on this wire is nested, so that difference is paid four or five + times per transfer. + + Guarded twice, because a wrong byte count is worse than a slow one. If + ``_values`` is ever absent the public attribute still answers. And a + buffer holding more elements than the offsets describe means the tensor + views a larger allocation (``torch.nested.narrow``), where the buffer + would overcount — nothing in ``nemo_rl`` builds one, since all four + construction sites use the compacting ``as_nested_tensor``, but this is + handed whatever a caller passes. + """ + buf = getattr(v, "_values", None) + if type(buf) is not torch.Tensor: + return v.nbytes + offsets = getattr(v, "_offsets", None) + if offsets is not None and buf.shape[0] != int(offsets[-1]): + return v.nbytes + return buf.nbytes + + def _dtype_salt(dtype: torch.dtype) -> int: """Salt distinguishing dtypes whose values reduce to the same words. @@ -311,10 +337,7 @@ def _td_bytes(td: TensorDict | None, max_nodes: int = 10_000) -> int: total = 0 for _, v in td.items(include_nested=True, leaves_only=False): if isinstance(v, torch.Tensor): - # No jagged special case: numel() already reports a nested - # tensor's total element count, and costs half of reaching - # through .values(). - total += v.numel() * v.element_size() + total += _tensor_bytes(v) elif isinstance(v, NonTensorData): total += _estimate_encoded_bytes(v.data, budget) elif isinstance(v, NonTensorStack): diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index e196c0c7ab5..5993e28d872 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -293,6 +293,33 @@ def test_td_bytes_counts_wire_payload(name, td, expected): assert _td_bytes(td) == expected +def test_td_bytes_jagged_matches_the_public_count(): + """The nested fast path reads the packed values buffer instead of the + tensor's own (dispatched, ~16us) ``nbytes``. It must agree exactly.""" + rows = [torch.arange(n, dtype=torch.int64) for n in (3, 7, 2, 5)] + jagged = torch.nested.as_nested_tensor(rows, layout=torch.jagged) + td = TensorDict({"x": jagged}, batch_size=[4]) + assert _td_bytes(td) == jagged.nbytes == sum(r.nbytes for r in rows) + + +def test_td_bytes_does_not_overcount_a_narrow_view(): + """``torch.nested.narrow`` yields a tensor whose values buffer views a + larger allocation. Trusting the buffer there would silently inflate + ``n_bytes`` — and an inflated byte count is worse than a slow one, since + ``max_bytes_per_key_seen`` reads it as a wire regression.""" + lengths = torch.tensor([3, 4, 5, 6]) + narrowed = torch.nested.narrow( + torch.zeros(4, 10), + 1, + torch.zeros(4, dtype=torch.int64), + lengths, + layout=torch.jagged, + ) + assert narrowed._values.nbytes > narrowed.nbytes, "fixture must be a view" + td = TensorDict({"x": narrowed}, batch_size=[4]) + assert _td_bytes(td) == narrowed.nbytes == int(lengths.sum()) * 4 + + def test_td_bytes_nontensordata_is_not_broadcast(): """``NonTensorData`` holds ONE object; counting it per batch row would inflate a 64-row put by 64x. Its bytes must not scale with batch size.""" From 7fc6aea03caa1e0c82082218ab9e9f5ae7a40d99 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 19:26:08 -0700 Subject: [PATCH 10/70] refactor(data-plane): name the real hazard in _tensor_bytes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cleanup pass over the nested byte-count fast path. No behaviour change — byte accounting is identical, the corruption matrix is unchanged, and put/get stay at 37us/15us. The guard comment was wrong in a way that invited breaking it. It said the first check exists in case `_values` is "absent", which reads as an argument for `if buf is None`. Measured: `_values` is a bound *method* on every dense tensor, never None, so `buf is None` would send every dense leaf down the private path. The check is on the type and the docstring now says why. Also: - `_estimate_encoded_bytes` still answered "bytes of a tensor" with `numel() * element_size()`, the exact expression this series replaced one function over. It calls `_tensor_bytes` now, so there is one answer to the question and a nested tensor reaching that branch does not pay the dispatch the fast path exists to avoid. - `_td_bytes`'s docstring still described leaves as `numel * element_size`. Arithmetically true, greppably stale. - Dropped the benchmark figures and the "all four construction sites" count from the docstring: the numbers are in README.md and a site count rots the next time a site is added. Kept the narrow-view trap. - Tests reuse the existing `_jagged()` helper, hoisted above its callers. Two findings were raised and deliberately left: - `int(offsets[-1])` is 73% of the nested-leaf cost, but eleven measured alternatives were all slower except `off[-1].item()`, which wins 0.07us on a 14us function. Not worth the churn. Worth re-measuring on a GPU box though: if `_offsets` is ever on CUDA that line is a D2H sync, which would dominate. No path in nemo_rl puts device tensors on this wire today (`kv_first_write` takes `final_batch_cpu`). - Reading torch privates is the established altitude here — `codec.py` already does `getattr(item, "_non_tensordict", None)` with the same shape one module over — and the risk that `_values` changes meaning is pinned by a unit test rather than a version comment. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 31 ++++++++++----------- tests/unit/data_plane/test_observability.py | 23 +++++++-------- 2 files changed, 25 insertions(+), 29 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index d25cc95d147..05202eea33c 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -159,19 +159,18 @@ def _pop_partition_keys( def _tensor_bytes(v: torch.Tensor) -> int: """Wire bytes of one tensor leaf, rectangular or nested. - A nested tensor's own ``nbytes`` (and ``numel``/``element_size``) route - through ``__torch_function__`` and measure ~16 us per leaf; its packed - values buffer answers the same question in ~0.1 us. Every per-token - field on this wire is nested, so that difference is paid four or five - times per transfer. - - Guarded twice, because a wrong byte count is worse than a slow one. If - ``_values`` is ever absent the public attribute still answers. And a - buffer holding more elements than the offsets describe means the tensor - views a larger allocation (``torch.nested.narrow``), where the buffer - would overcount — nothing in ``nemo_rl`` builds one, since all four - construction sites use the compacting ``as_nested_tensor``, but this is - handed whatever a caller passes. + A nested tensor's ``nbytes`` dispatches through ``__torch_function__``; + its packed values buffer answers the same question without dispatching, + and every per-token field on this wire is nested. + + Two guards, because a wrong byte count is worse than a slow one: + + * ``_values`` is a bound *method* on every dense tensor, so the first + check is on the type, not for absence — ``buf is None`` would be wrong. + * A buffer holding more elements than the offsets describe means the + tensor views a larger allocation (``torch.nested.narrow``), where the + buffer overcounts. Nothing here builds one, but this is handed + whatever a caller passes. """ buf = getattr(v, "_values", None) if type(buf) is not torch.Tensor: @@ -270,7 +269,7 @@ def _estimate_encoded_bytes(obj: Any, budget: list[int]) -> int: total += _estimate_encoded_bytes(v, budget) return total if isinstance(obj, torch.Tensor): - return obj.numel() * obj.element_size() + return _tensor_bytes(obj) # Unknown type -> pickle/cloudpickle Ext. Cheap proxy; the real size # would need an actual dumps(), which is what we are avoiding. return 64 @@ -310,8 +309,8 @@ def _nontensor_stack_bytes(stack: NonTensorStack, budget: list[int]) -> int: def _td_bytes(td: TensorDict | None, max_nodes: int = 10_000) -> int: """Payload bytes of a TensorDict, as the wire will see them. - Tensor leaves count ``numel * element_size``, which equals - ``t.contiguous().nbytes`` -- the size mooncake registers and sends. + Tensor leaves count ``nbytes`` (see :func:`_tensor_bytes`), which is the + size mooncake registers and sends. Non-tensor leaves are estimated with :func:`_estimate_encoded_bytes`, since TQ ships them over a separate msgpack path. Both kinds are counted in a single ``items()`` pass; ``keys()`` + ``get()`` would re-resolve diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 5993e28d872..14f865fe060 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -242,6 +242,13 @@ def test_observability_records_realistic_rollout_put() -> None: # ── byte accounting ──────────────────────────────────────────────────── +def _jagged(rows): + return TensorDict( + {"x": torch.nested.nested_tensor(rows, layout=torch.jagged)}, + batch_size=[len(rows)], + ) + + @pytest.mark.parametrize( "name,td,expected", [ @@ -297,16 +304,13 @@ def test_td_bytes_jagged_matches_the_public_count(): """The nested fast path reads the packed values buffer instead of the tensor's own (dispatched, ~16us) ``nbytes``. It must agree exactly.""" rows = [torch.arange(n, dtype=torch.int64) for n in (3, 7, 2, 5)] - jagged = torch.nested.as_nested_tensor(rows, layout=torch.jagged) - td = TensorDict({"x": jagged}, batch_size=[4]) - assert _td_bytes(td) == jagged.nbytes == sum(r.nbytes for r in rows) + td = _jagged(rows) + assert _td_bytes(td) == td["x"].nbytes == sum(r.nbytes for r in rows) def test_td_bytes_does_not_overcount_a_narrow_view(): """``torch.nested.narrow`` yields a tensor whose values buffer views a - larger allocation. Trusting the buffer there would silently inflate - ``n_bytes`` — and an inflated byte count is worse than a slow one, since - ``max_bytes_per_key_seen`` reads it as a wire regression.""" + larger allocation; trusting it would silently inflate ``n_bytes``.""" lengths = torch.tensor([3, 4, 5, 6]) narrowed = torch.nested.narrow( torch.zeros(4, 10), @@ -500,13 +504,6 @@ def test_hash_fingerprints_released_on_clear(): client.close() -def _jagged(rows): - return TensorDict( - {"x": torch.nested.nested_tensor(rows, layout=torch.jagged)}, - batch_size=[len(rows)], - ) - - def test_hash_fingerprint_covers_jagged_fields(): """The per-token fields on this wire are jagged by the time they reach ``put_samples`` (``codec.pack_jagged_fields``). Skipping nested leaves From 839afb6ef2258b1ce2a498c9890fc873a1ec227a Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 19:44:54 -0700 Subject: [PATCH 11/70] feat(data-plane): enable data-plane metrics by default MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Turns `data_plane.observability.enabled` on in the exemplar config. It only engages when `data_plane.enabled` is true, which is itself false by default, so this changes nothing for runs that don't use the data plane — and for runs that do, the cost is below what can be measured. Against a live TransferQueue, three configs of the same workload interleaved (256x2048, 5.4 MB jagged per step, 15 iterations, paired per-iteration deltas so machine drift cancels): metrics off 182.9 ms/step metrics on 183.4 ms/step +0.44% run-to-run spread is 4.4 ms The wrapper itself is 37 us/put and 15 us/get. `verify_tensor_hash` stays off: it costs ~4% and, until fingerprints travel in `KVBatchMeta.tags`, it verifies almost nothing across the process split it would need to cross. Removes the default `log_event` sink, which this change makes mandatory rather than optional. `on_event` fired on every single transfer, so defaulting metrics on would have defaulted a `logger.info` per data-plane op with it (measured 13 us/op with INFO enabled, versus 1.5 us with no sink). The metrics surface is `get_step_metrics()`, which the trainer logs once a step under the `data_plane/` prefix and which fans out to whatever backends the run enabled. A per-op hook is now opt-in via `observability.callback`; `log_event` is still exported for it. This also makes the lazy event-dict build in `_emit` live rather than dead code — with no sink registered, the per-op dict is never allocated. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- examples/configs/grpo_math_1B.yaml | 2 +- nemo_rl/data_plane/README.md | 14 +++++++++----- nemo_rl/data_plane/factory.py | 12 ++++++------ nemo_rl/data_plane/interfaces.py | 3 ++- tests/unit/reference_configs/grpo_math_1B.yaml | 2 +- 5 files changed, 19 insertions(+), 14 deletions(-) diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 282ac2a5aa8..5a305b36666 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -584,7 +584,7 @@ data_plane: # higher buys little at linear HBM cost. # GDR needs that headroom to pay off. observability: # per-op data-plane timing/volume - enabled: false # true => wrap client in MetricsDataPlaneClient + enabled: true # per-op timing/volume; cost is below measurement noise verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~2.4ms per 12MB batch) # Multi-Teacher On-Policy Distillation (MOPD): distills from one or more teacher diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index a46f0c1f1c6..ea4b12cde75 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -437,7 +437,7 @@ data_plane: use_gdr: false # GPU-memory RDMA staging in CUDA clients gdr_staging_buffer_mb: 1024 # persistent MiB per active GDR client observability: # NotRequired - enabled: false # per-op timing / latency percentiles / volume + enabled: true # per-op timing / latency percentiles / volume verify_tensor_hash: false # debug: wire-in vs wire-out tensor check ``` @@ -455,10 +455,14 @@ carries — 256 ragged rows, 12 MB, jagged per-token fields as 0.1% of a 59 ms operation. What is left is dominated by the per-key attribution `clear_samples` needs to undo. -Note that `enabled: true` also installs `log_event` as the default -`on_event` sink, which emits one `logger.info` line per data-plane op. Pass -`observability.callback` explicitly (or `None`) if you want the counters -without the per-op log. +This is **on by default**, and only engages when `data_plane.enabled` is +true, so it costs nothing for runs that don't use the data plane. There is +no default per-op sink: `get_step_metrics()` is the surface, and +`grpo_train_sync` logs it once a step under the `data_plane/` prefix — so +the series reach whatever backends the run has enabled (wandb, TensorBoard, +MLflow). Roughly 5-8 series per distinct op tag. Set +`observability.callback` if you additionally want a hook on every transfer; +`log_event` is exported for that. `verify_tensor_hash: true` additionally records a `torch.hash_tensor` fingerprint on every put and re-checks it on every get, so a tensor that diff --git a/nemo_rl/data_plane/factory.py b/nemo_rl/data_plane/factory.py index 89a52310b38..b54756f45ae 100644 --- a/nemo_rl/data_plane/factory.py +++ b/nemo_rl/data_plane/factory.py @@ -179,16 +179,16 @@ def build_data_plane_client( else cfg.get("observability") ) or {} if obs.get("enabled", False): - from nemo_rl.data_plane.observability import ( - MetricsDataPlaneClient, - log_event, - ) + from nemo_rl.data_plane.observability import MetricsDataPlaneClient - on_event = obs.get("callback") or log_event + # No default per-op sink. The metrics surface is ``get_step_metrics``, + # which the trainer logs once a step; a callback here fires on every + # single transfer. ``log_event`` is still exported for anyone who + # wants that, but it is opt-in via ``observability.callback``. # pyrefly: obs.get returns Any, can't narrow to the expected callback type. client = MetricsDataPlaneClient( # type: ignore[bad-argument-type] client, - on_event=on_event, + on_event=obs.get("callback"), verify_tensor_hash=bool(obs.get("verify_tensor_hash")), ) return client diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index d00b88d26b2..a2311a5c396 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -182,7 +182,8 @@ class ObservabilityConfig(TypedDict): injected programmatically (callables don't round-trip through YAML) — set ``cfg["observability"]["callback"] = my_fn`` before :func:`build_data_plane_client` to plug into wandb / file / log. - Default callback prints one line per op for debug. + There is no default callback: per-step metrics reach the logger via + ``get_step_metrics``, so a per-op sink is opt-in. ``verify_tensor_hash`` is a correctness check, not a metric: each put records a per-row ``torch.hash_tensor`` fingerprint and each get diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 0ffc106e878..d06a39060d5 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -556,7 +556,7 @@ data_plane: # higher buys little at linear HBM cost. # GDR needs that headroom to pay off. observability: # per-op data-plane timing/volume - enabled: false # true => wrap client in MetricsDataPlaneClient + enabled: true # per-op timing/volume; cost is below measurement noise verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~2.4ms per 12MB batch) # Multi-Teacher On-Policy Distillation (MOPD); null = disabled (default). Mirrors From f1376000d1c17c41473a074036eabf48d98f347b Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 19:54:38 -0700 Subject: [PATCH 12/70] perf(data-plane): hold one byte total per partition, not one per key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_record_put` was 56% of the put path — a 256-iteration Python loop giving every key an even share of the batch total. There was never a per-key truth to keep: the split was even going in, and a subset clear released the mean coming out. One total plus one live-key set says the same thing and lets `set.update` do the per-key work in C. _record_put 18.6 us -> 3.0 us (6.3x, 256 keys) put path 35.2 us -> 20.0 us Against the PR head that is 199.5 us -> 20.0 us per put and 160.0 -> 15.2 per get. Bytes are released pro rata on clear, and clearing the last key releases the remainder exactly, so a partition reconciles to zero however it is chopped up — verified end to end through live TransferQueue (`bytes_outstanding` back to 0.00 MB, byte totals unchanged). Also trims the per-step metric set from 8 series per op tag to 5. Dropped `mean_ms`, which is exactly `wall_s / calls`, and the absolute `fixed_overhead_ms` / `bandwidth_mb_s`, which are cumulative regressions that barely move per step — `overhead_frac` carries the actionable half (is this op overhead- or bandwidth-bound). `snapshot()` still returns all of it for a one-off inspection; this only changes what is logged every step. At the ~8 op tags a real run produces that is ~70 series down to ~46. Verified in wandb: the metrics reach `Logger.log_metrics(prefix="data_plane")` and fan out to every enabled backend. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 67 ++++++++++++++++++++--------- 1 file changed, 46 insertions(+), 21 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 05202eea33c..ba33f893bfe 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -519,10 +519,11 @@ def __init__( self._on_event = on_event self._verify_tensor_hash = verify_tensor_hash self._stats = DataPlaneStats() - # Nested per-partition / per-key live byte counts. Populated on - # successful ``put_samples``; popped on successful ``clear_samples``. - # Bounded by the live key population, not cumulative traffic. - self._bytes_by_partition: dict[str, dict[str, int]] = {} + # Live bytes and live keys per partition. Populated on successful + # ``put_samples``, released on successful ``clear_samples``. Bounded + # by the live key population, not by cumulative traffic. + self._bytes_by_partition: dict[str, int] = {} + self._keys_by_partition: dict[str, set[str]] = {} # partition -> sample_id -> field -> wire-in fingerprint. Same # lifetime as ``_bytes_by_partition``: cleared by ``clear_samples``, # so it is bounded by the live key population. @@ -549,7 +550,7 @@ def snapshot(self) -> dict[str, Any]: """ out = asdict(self._stats) out["n_keys_outstanding"] = sum( - len(d) for d in self._bytes_by_partition.values() + len(k) for k in self._keys_by_partition.values() ) for op, s in out["by_op"].items(): s["mean_ms"] = s["wall_ms"] / s["calls"] if s["calls"] else 0.0 @@ -631,23 +632,28 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: if calls <= 0: continue op_ms = st["wall_ms"] - prev_op.get("wall_ms", 0.0) + # Five series per op tag, deliberately. Anything exactly + # derivable from these is left out rather than logged: mean is + # wall_s/calls, and a dashboard can divide. ``snapshot()`` still + # carries the full picture for a one-off inspection. metrics[f"{op}/calls"] = calls metrics[f"{op}/wall_s"] = op_ms / 1e3 - metrics[f"{op}/mean_ms"] = op_ms / calls # Percentiles and the overhead/bandwidth fit are cumulative by # construction (histogram buckets, regression sums) -- as-is. metrics[f"{op}/p50_ms"] = st["p50_ms"] metrics[f"{op}/p99_ms"] = st["p99_ms"] fit = st["fit"] if fit.get("model_trustworthy"): - metrics[f"{op}/fixed_overhead_ms"] = fit["fixed_ms"] - metrics[f"{op}/bandwidth_mb_s"] = fit["bandwidth_mb_s"] + # Only the actionable half of the fit: whether this op is + # overhead- or bandwidth-bound. The absolute fixed_ms and + # bandwidth_mb_s behind it are cumulative regressions that + # barely move per step -- read them from ``snapshot()``. metrics[f"{op}/overhead_frac"] = fit["overhead_frac_at_mean"] return metrics def bytes_outstanding_by_partition(self) -> dict[str, int]: """Per-partition breakdown of currently-held bytes.""" - return {p: sum(d.values()) for p, d in self._bytes_by_partition.items()} + return dict(self._bytes_by_partition) def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: """Attribute put bytes per key so a later ``clear_samples`` can subtract. @@ -655,23 +661,24 @@ def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: Called after the underlying RPC succeeds so a failed put never leaves the accounting inflated. - ``n_bytes`` is a whole-batch figure, so there is no per-key truth to - preserve: the split is even and the division remainder lands on the - first key rather than being spread across the batch. + ``n_bytes`` is a whole-batch figure, so there was never a per-key + truth to keep: the old per-key dict stored an even split, and a + subset clear released the mean either way. Holding one total and one + key set says the same thing and lets ``set.update`` do the per-key + work in C — 18.6 us to 3.0 us at 256 keys, which was the single + largest remaining cost on the put path. Args: partition_id: Partition the keys were written to. keys: Per-sample uids that were written. - n_bytes: Total bytes written; distributed evenly across keys. + n_bytes: Total bytes written; released pro rata on clear. """ if not keys or n_bytes <= 0: return - per_key, remainder = divmod(n_bytes, len(keys)) - partition_dict = self._bytes_by_partition.setdefault(partition_id, {}) - get_held = partition_dict.get - for key in keys: - partition_dict[key] = get_held(key, 0) + per_key - partition_dict[keys[0]] += remainder + self._keys_by_partition.setdefault(partition_id, set()).update(keys) + self._bytes_by_partition[partition_id] = ( + self._bytes_by_partition.get(partition_id, 0) + n_bytes + ) self._stats.bytes_outstanding += n_bytes if self._stats.bytes_outstanding > self._stats.peak_bytes_outstanding: self._stats.peak_bytes_outstanding = self._stats.bytes_outstanding @@ -682,6 +689,11 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: Called after the underlying RPC succeeds so a failed clear keeps the accounting consistent with TQ's actual state. + Bytes are released pro rata: the partition's total times the share + of its live keys being dropped. Clearing the last key releases the + remainder exactly, so a partition always reconciles to zero however + it is chopped up. + Args: partition_id: Partition the keys were dropped from. keys: Uids dropped; ``None`` means the whole partition was cleared. @@ -690,8 +702,21 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: _pop_partition_keys(self._hash_by_partition, partition_id, keys) if keys is None: self._batch_scope.pop(partition_id, None) - freed = _pop_partition_keys(self._bytes_by_partition, partition_id, keys) - self._stats.bytes_outstanding -= sum(freed) + live = self._keys_by_partition.get(partition_id) + if live is None: + return + total = self._bytes_by_partition.get(partition_id, 0) + if keys is not None: + live -= live.intersection(keys) + if keys is None or not live: + freed = total + del self._keys_by_partition[partition_id] + self._bytes_by_partition.pop(partition_id, None) + else: + dropped = len(keys) + freed = total * dropped // (len(live) + dropped) + self._bytes_by_partition[partition_id] = total - freed + self._stats.bytes_outstanding -= freed # ── wire-in / wire-out fingerprinting (opt-in) ───────────────────── From caec8c524375ca4ff161d054e9f92c8f679d53b3 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 22:10:41 -0700 Subject: [PATCH 13/70] fix(data-plane): one unit per dimension, and a tail metric that moves MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three defects visible the moment the metrics reached a W&B chart. Mixed units on sibling series. `wall_s` sat next to `p99_ms` on the same axis, so a step read `0.008` beside `24.85` and looked like a data-plane bug rather than an axis one. Volumes had it worse one dimension over: a realistic step moved `0.00017` GB. Every duration is now `_ms` and every volume `_mb`, with a test asserting no `_s`/`_gb` series can come back. A flat p50. The latency histogram is cumulative and never reset, so the per-step percentile was a lifetime figure that stopped moving as soon as the distribution settled. A p99 pinned to bucket edges. With the handful of calls an op makes in one step, a percentile off a 15-bucket histogram is geometry, not data: one sample in `(10, 25]` always yields `10 + 15*0.99 = 24.85`, which is exactly the constant that showed up. The bouncing was the sample changing bucket. Per-step percentiles are replaced by `max_ms` — exact, no quantisation, meaningful at one call, and scoped to the step. Scoping matters: the first cut of this used a lifetime max, which is monotonic and goes flat the moment the worst call has been seen, reproducing the defect it replaced. `get_step_metrics` now zeroes the window after reporting. The percentiles stay in `snapshot()`, where the cumulative sample count justifies them. Per-op series go 5 -> 4: calls, wall_ms, max_ms, overhead_frac. Verified against live TransferQueue with runs in the grpo-dev-zhiyul W&B group; `max_ms` now varies step to step while the lifetime worst is still readable from `snapshot()`. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 4 +- nemo_rl/data_plane/README.md | 13 +++++ nemo_rl/data_plane/observability.py | 56 ++++++++++++++------ tests/unit/data_plane/test_observability.py | 58 +++++++++++++++++++++ 4 files changed, 114 insertions(+), 17 deletions(-) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index a2eabf21064..cc99f1bba4b 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -402,9 +402,9 @@ def _log_data_plane_metrics( metrics = client.get_step_metrics(total_step_time) logger.log_metrics(metrics, step, prefix="data_plane") print( - f" • data plane: {metrics['wall_s']:.2f}s " + f" • data plane: {metrics['wall_ms']:.0f}ms " f"({100 * metrics['frac_of_step']:.1f}% of step), " - f"{metrics['comm_volume_gb']:.2f} GB moved" + f"{metrics['comm_volume_mb']:.1f} MB moved" ) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index ea4b12cde75..b6af34fa720 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -449,6 +449,19 @@ counts sum into one cluster-wide distribution) and byte volume. `snapshot()` returns the cumulative view; `get_step_metrics(step_time_s)` returns the per-step delta already flattened for the logger. +**Units:** every duration is `_ms`, every volume is `_mb`, no exceptions — +a chart mixing `wall_s` against `p99_ms` puts a 0.008 beside a 24.85 and +reads as a data-plane bug rather than an axis one. + +**Per step you get four series per op tag:** `calls`, `wall_ms`, `max_ms` +and (when the fit is trustworthy) `overhead_frac`. Not percentiles — the +histogram is cumulative by design, so a per-step p50 off it goes flat, and +at the handful of calls an op makes in one step a p99 is bucket geometry +rather than data (one sample in the `(10, 25]` bucket always yields +`10 + 15*0.99 = 24.85`). `max_ms` is exact, scoped to the step, and says +the same thing at that sample size. The percentiles remain in `snapshot()`, +where the cumulative sample count justifies them. + Measured against a no-op inner client on the payload the wire actually carries — 256 ragged rows, 12 MB, jagged per-token fields as `pack_jagged_fields` leaves them: **~37 µs per put, ~15 µs per get**, under diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index ba33f893bfe..85f7f924e7a 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -448,6 +448,17 @@ class OpStats: # wall_ms non-linear in n_bytes, and a low R^2 is the signal to stop # trusting the overhead/bandwidth split. sum_ms_sq: float = 0.0 + # Slowest single call, exact. The histogram below can only place a + # call in a bucket, so at the handful of calls an op makes in one step + # a percentile off it is bucket geometry rather than data -- p99 of one + # sample in (10, 25] is always 10 + 15*0.99 = 24.85. This is the + # per-step tail signal; the histogram is for the cumulative view. + max_ms: float = 0.0 + # Same, but scoped to the current step: ``get_step_metrics`` zeroes it + # each time it reports. Without this the per-step series is the lifetime + # max, which is monotonic and goes flat the moment the worst call has + # been seen -- the same defect as logging a cumulative percentile. + step_max_ms: float = 0.0 # Latency distribution over ALL statuses, matching calls/wall_ms: a # timeout is real tail latency the pipeline actually paid for. latency_hist: list[int] = field( @@ -594,17 +605,27 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: snap = self.snapshot() prev = self._prev_snapshot self._prev_snapshot = snap + # Snapshot first, then open a fresh window: the values just read are + # this step's, and anything after this call belongs to the next one. + step_maxima = {op: b.step_max_ms for op, b in self._stats.by_op.items()} + for bucket in self._stats.by_op.values(): + bucket.step_max_ms = 0.0 wall_ms = snap["total_wall_ms"] - prev.get("total_wall_ms", 0.0) vol = snap["comm_volume_bytes"] - prev.get("comm_volume_bytes", 0) + # Every duration is ms and every volume is MB, with no exceptions: + # a chart that mixes wall_s against p99_ms puts a 0.008 next to a + # 24.85 and reads as a bug in the data plane rather than in the + # axis. GB was the same problem one dimension over -- a realistic + # step moved 0.00017 GB. metrics: dict[str, float] = { - "wall_s": wall_ms / 1e3, + "wall_ms": wall_ms, "frac_of_step": (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0, - "comm_volume_gb": vol / 1e9, - "bytes_written_gb": (snap["bytes_written"] - prev.get("bytes_written", 0)) - / 1e9, - "bytes_read_gb": (snap["bytes_read"] - prev.get("bytes_read", 0)) / 1e9, - "bytes_outstanding_gb": snap["bytes_outstanding"] / 1e9, + "comm_volume_mb": vol / 1e6, + "bytes_written_mb": (snap["bytes_written"] - prev.get("bytes_written", 0)) + / 1e6, + "bytes_read_mb": (snap["bytes_read"] - prev.get("bytes_read", 0)) / 1e6, + "bytes_outstanding_mb": snap["bytes_outstanding"] / 1e6, } if self._verify_tensor_hash: hv, prev_hv = snap["hash_verify"], prev.get("hash_verify", {}) @@ -632,16 +653,17 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: if calls <= 0: continue op_ms = st["wall_ms"] - prev_op.get("wall_ms", 0.0) - # Five series per op tag, deliberately. Anything exactly - # derivable from these is left out rather than logged: mean is - # wall_s/calls, and a dashboard can divide. ``snapshot()`` still - # carries the full picture for a one-off inspection. + # Four series per op tag. Anything exactly derivable is left + # out rather than logged: mean is wall_ms/calls, and a dashboard + # can divide. ``snapshot()`` still carries the full picture, + # percentiles included, for a one-off inspection. metrics[f"{op}/calls"] = calls - metrics[f"{op}/wall_s"] = op_ms / 1e3 - # Percentiles and the overhead/bandwidth fit are cumulative by - # construction (histogram buckets, regression sums) -- as-is. - metrics[f"{op}/p50_ms"] = st["p50_ms"] - metrics[f"{op}/p99_ms"] = st["p99_ms"] + metrics[f"{op}/wall_ms"] = op_ms + # ``max_ms`` rather than p50/p99. Those come off a histogram + # that is never reset, so per step they were a lifetime figure + # that goes flat, quantised to bucket edges. The max is exact + # and says the same thing at the handful of calls per step. + metrics[f"{op}/max_ms"] = step_maxima.get(op, 0.0) fit = st["fit"] if fit.get("model_trustworthy"): # Only the actionable half of the fit: whether this op is @@ -939,6 +961,10 @@ def _emit( bucket = stats.by_op[op] = OpStats() bucket.calls += 1 bucket.wall_ms += wall_ms + if wall_ms > bucket.max_ms: + bucket.max_ms = wall_ms + if wall_ms > bucket.step_max_ms: + bucket.step_max_ms = wall_ms bucket.latency_hist[bisect_left(LATENCY_BUCKETS_MS, wall_ms)] += 1 if status != "ok": bucket.errors += 1 diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 14f865fe060..67ee0588bbe 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -21,6 +21,8 @@ from __future__ import annotations +from time import monotonic + import pytest import torch from tensordict import NonTensorData, NonTensorStack, TensorDict @@ -376,6 +378,62 @@ def test_outstanding_bytes_reconcile_exactly(): client.close() +def test_step_metrics_use_one_unit_per_dimension(): + """Every duration is ms, every volume MB. A chart mixing `wall_s` with + `p99_ms` shows 0.008 beside 24.85 and reads as a data-plane bug rather + than an axis one.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client.register_partition( + partition_id="p", fields=["x"], num_samples=2, consumer_tasks=["t"] + ) + client.put_samples( + sample_ids=["a", "b"], + partition_id="p", + fields=TensorDict({"x": torch.zeros(2, 3)}, batch_size=[2]), + ) + keys = set(client.get_step_metrics(1.0)) + assert not [k for k in keys if k.endswith(("_s", "_gb", "_kb", "_us"))], keys + assert "wall_ms" in keys and "comm_volume_mb" in keys + client.close() + + +def test_step_metrics_tail_is_exact_not_bucketed(): + """Per-step percentiles came off a histogram that is never reset, so + they went flat and quantised to bucket edges (p99 of a single sample in + (10, 25] is always 10 + 15*0.99 = 24.85). ``max_ms`` is exact and + tracks the slowest call actually seen.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client.register_partition( + partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["t"] + ) + client._emit("put", "p", 1, 8, monotonic() - 0.030, "ok") # a 30 ms call + metrics = client.get_step_metrics(1.0) + + assert "put/p99_ms" not in metrics and "put/p50_ms" not in metrics + assert metrics["put/max_ms"] >= 30.0 + assert metrics["put/max_ms"] != pytest.approx(24.85, abs=0.5), "bucket edge" + # the cumulative view still carries percentiles for one-off inspection + assert "p99_ms" in client.snapshot()["by_op"]["put"] + client.close() + + +def test_step_max_is_scoped_to_the_step(): + """A lifetime max is monotonic and goes flat the moment the worst call + has been seen — the same defect as logging a cumulative percentile. The + reported max must fall again when a step is quicker.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client._emit("put", "p", 1, 8, monotonic() - 0.050, "ok") # slow step + slow = client.get_step_metrics(1.0)["put/max_ms"] + client._emit("put", "p", 1, 8, monotonic() - 0.001, "ok") # quick step + quick = client.get_step_metrics(1.0)["put/max_ms"] + + assert slow >= 50.0 + assert quick < slow, "step max must reset, not carry the lifetime worst" + # the lifetime worst is still available for a one-off look + assert client.snapshot()["by_op"]["put"]["max_ms"] >= 50.0 + client.close() + + def test_no_callback_still_accumulates_stats(): """``on_event=None`` skips building the event dict; the counters that ``snapshot()`` reports must not depend on a sink being registered.""" From 79aaebfdea88f60fe7b5a73d4dd173cd212021e1 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 22:18:05 -0700 Subject: [PATCH 14/70] fix(data-plane): report the latency split in ms, not as a ratio MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `overhead_frac` was the wrong shape twice over. It is a ratio on an axis of milliseconds, so it cannot be read next to the series beside it — and it carries the same defect as the cumulative percentiles just removed, because the fit is computed from lifetime sufficient statistics. Measured over eight steps of identical shape it converges and then sits there: step 1 0.7662 step 5 0.8057 step 3 0.7830 step 8 0.7853 A line chart of a constant, and a ratio hides the two numbers worth having. Replaced with `overhead_ms` and `transfer_ms`: the op's time split into fixed per-request cost and bandwidth, in ms. They stack — together they are the model's estimate of that step's `wall_ms` — so charting them against the measured `wall_ms` shows the split and the model error in one picture. The coefficients still come from the cumulative fit, which is right: a fitted model should be stable. What is per step is the *attribution* — the coefficients applied to this step's calls and bytes — so `transfer_ms` moves with the bytes actually sent while `overhead_ms` moves with the call count. Checked against a client with known affine latency (8.0 ms/call, 500 MB/s): recovered fixed 8.08 ms bandwidth 500 MB/s R2 1.0000 stacking overhead_ms + transfer_ms == wall_ms to within 0.0% variation transfer_ms 19.9 -> 28.2 ms across steps, tracking bytes `overhead_frac` is derivable from the pair, so it is not logged. The full fit, ratio included, stays in `snapshot()`. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 12 +++++++-- nemo_rl/data_plane/observability.py | 21 +++++++++++---- tests/unit/data_plane/test_observability.py | 29 +++++++++++++++++++++ 3 files changed, 55 insertions(+), 7 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index b6af34fa720..e9a5d7f3ae6 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -453,8 +453,16 @@ per-step delta already flattened for the logger. a chart mixing `wall_s` against `p99_ms` puts a 0.008 beside a 24.85 and reads as a data-plane bug rather than an axis one. -**Per step you get four series per op tag:** `calls`, `wall_ms`, `max_ms` -and (when the fit is trustworthy) `overhead_frac`. Not percentiles — the +**Per step you get, per op tag:** `calls`, `wall_ms`, `max_ms`, and — when +the affine fit is trustworthy — `overhead_ms` and `transfer_ms`. Those last +two are the split of the op's time into fixed per-request cost and +bandwidth, in ms, and they *stack*: together they are the model's estimate +of the step's `wall_ms`, so charting them against the measured `wall_ms` +shows the breakdown and the model error in one picture. The coefficients +come from the cumulative fit (a model should be stable); the attribution is +per step, applied to that step's calls and bytes. A ratio was tried first +and was the wrong shape — cumulative and therefore flat, and unitless on an +axis of milliseconds. Not percentiles — the histogram is cumulative by design, so a per-step p50 off it goes flat, and at the handful of calls an op makes in one step a p99 is bucket geometry rather than data (one sample in the `(10, 25]` bucket always yields diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 85f7f924e7a..f0b2f2384d8 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -666,11 +666,22 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: metrics[f"{op}/max_ms"] = step_maxima.get(op, 0.0) fit = st["fit"] if fit.get("model_trustworthy"): - # Only the actionable half of the fit: whether this op is - # overhead- or bandwidth-bound. The absolute fixed_ms and - # bandwidth_mb_s behind it are cumulative regressions that - # barely move per step -- read them from ``snapshot()``. - metrics[f"{op}/overhead_frac"] = fit["overhead_frac_at_mean"] + # The step's time split into the two things that cause it, + # in ms, rather than a ratio. These stack: together they are + # the model's estimate of this op's ``wall_ms`` for the + # step, so charting them against the measured ``wall_ms`` + # shows both the split and how well the model holds. + # + # The *coefficients* come from the cumulative fit and should + # be stable -- that is what a fitted model is for. The + # *attribution* is per step, because it is applied to this + # step's calls and bytes. A ratio would have been neither: + # cumulative and therefore flat, and unitless on an axis of + # milliseconds. + op_bytes = st["n_bytes"] - prev_op.get("n_bytes", 0) + ms_per_byte = 1.0 / (fit["bandwidth_mb_s"] * 1e3) + metrics[f"{op}/overhead_ms"] = fit["fixed_ms"] * calls + metrics[f"{op}/transfer_ms"] = ms_per_byte * op_bytes return metrics def bytes_outstanding_by_partition(self) -> dict[str, int]: diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 67ee0588bbe..8231daed595 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -417,6 +417,35 @@ def test_step_metrics_tail_is_exact_not_bucketed(): client.close() +def test_latency_breakdown_stacks_to_wall_ms(): + """The fit is reported as two ms components rather than a ratio, so a + chart can stack them against the measured ``wall_ms``. + + A ratio would have been both flat (the fit is cumulative) and unitless + on an axis of milliseconds. These carry the coefficients from the + cumulative fit but attribute them to *this* step's calls and bytes. + """ + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + fixed_ms, mb_per_s = 8.0, 500.0 + now = monotonic() + for i in range(12): # varied sizes, or the fit is unidentifiable + n_bytes = 50_000 * (i + 1) + wall_ms = fixed_ms + n_bytes / (mb_per_s * 1e3) + client._emit("put", "p", 1, n_bytes, now - wall_ms / 1e3, "ok") + + metrics = client.get_step_metrics(1.0) + fit = client.snapshot()["by_op"]["put"]["fit"] + assert fit["model_trustworthy"], fit + assert fit["fixed_ms"] == pytest.approx(fixed_ms, rel=0.05) + assert fit["bandwidth_mb_s"] == pytest.approx(mb_per_s, rel=0.05) + + # the two components are the split, in ms, and they add up + total = metrics["put/overhead_ms"] + metrics["put/transfer_ms"] + assert total == pytest.approx(metrics["put/wall_ms"], rel=0.05) + assert "put/overhead_frac" not in metrics, "a ratio is derivable from these" + client.close() + + def test_step_max_is_scoped_to_the_step(): """A lifetime max is monotonic and goes flat the moment the worst call has been seen — the same defect as logging a cumulative percentile. The From c1f4275bab503714199dec5515ff4be7eee99760 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 22:37:36 -0700 Subject: [PATCH 15/70] fix(data-plane): scope the metric prefix to the driver that produced it These series are one process, not the cluster, and `data_plane/` read as though they were cluster-wide totals. Four processes build a client and each keeps independent counters: tq_policy.py:119 driver bootstrap=True worker_mixin.py:149 one per rank bootstrap=False sync_rollout_actor.py:136 rollout actor (@ray.remote) single_controller_utils/setup.py:561 single controller `grpo_train_sync` logs the driver's, so the prefix is now `data_plane/driver/`. This also explains the low call counts: the driver issues about one op of each kind per step -- register, write advantages, read a couple of columns, clear -- while the bulk traffic happens elsewhere. `kv_first_write` in the rollout actor writes the whole rollout and the workers' per-DP-rank `get_samples` read it back, and neither shows up in these numbers. Reading `comm_volume_mb` as cluster volume would undercount by most of the payload. `OpStats` was built additive for exactly this -- histogram buckets and regression sufficient statistics sum across ranks into one cluster-wide distribution, which averaging per-rank percentiles cannot do. Nothing collects them yet, so that is an affordance the design leaves open rather than a feature. Said plainly in the README instead of implied by a prefix. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 18 ++++++++++++++---- nemo_rl/data_plane/README.md | 13 +++++++++++++ 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index cc99f1bba4b..a91c2525667 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -392,15 +392,25 @@ def _log_data_plane_metrics( ``VllmGeneration.get_step_metrics`` -- so the async trainer can log the same metrics with one call against its own client. - Only the driver's client is visible here: clients are built per process - (``tq_policy`` bootstraps one; rollout actors and policy workers build - their own), so this covers the trainer's reads, not cluster-wide traffic. + Scoped to the driver, and the prefix says so. Every process builds its + own client -- ``tq_policy`` here, plus one per policy worker + (``worker_mixin``), one in the rollout actor, one on the + single-controller path -- and each keeps independent counters. The + driver issues roughly one op of each kind per step; the bulk traffic is + elsewhere (``kv_first_write`` in the rollout actor writes the rollout, + per-DP-rank ``get_samples`` in the workers read it), and none of it + appears here. A ``data_plane/`` prefix would read as cluster-wide + totals, which these are not. + + ``OpStats`` is additive on purpose -- histograms and regression sums + from every rank sum into one cluster-wide view -- but nothing collects + them yet, so that remains a design affordance rather than a feature. """ client = getattr(policy, "dp_client", None) if not isinstance(client, MetricsDataPlaneClient): return # observability disabled -> plain adapter metrics = client.get_step_metrics(total_step_time) - logger.log_metrics(metrics, step, prefix="data_plane") + logger.log_metrics(metrics, step, prefix="data_plane/driver") print( f" • data plane: {metrics['wall_ms']:.0f}ms " f"({100 * metrics['frac_of_step']:.1f}% of step), " diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index e9a5d7f3ae6..01060f1ba12 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -449,6 +449,19 @@ counts sum into one cluster-wide distribution) and byte volume. `snapshot()` returns the cumulative view; `get_step_metrics(step_time_s)` returns the per-step delta already flattened for the logger. +**Scope: one process, not the cluster.** Every process builds its own +client with its own counters — the driver, each policy worker, the rollout +actor. `grpo_train_sync` logs the *driver's*, under `data_plane/driver/`. +The driver issues about one op of each kind per step, so `calls` is small +by construction; the bulk traffic is the rollout actor's `kv_first_write` +and the workers' per-DP-rank `get_samples`, and neither appears in these +series. Do not read `comm_volume_mb` as cluster-wide volume. + +`OpStats` is additive on purpose — the histogram buckets and the regression +sufficient statistics from every rank sum into one cluster-wide view — but +nothing collects them yet. That is an affordance the design leaves open, +not a feature that exists. + **Units:** every duration is `_ms`, every volume is `_mb`, no exceptions — a chart mixing `wall_s` against `p99_ms` puts a 0.008 beside a 24.85 and reads as a data-plane bug rather than an axis one. From 7894024eef4c72ccde8cdd40ffb78ea66f355297 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 23:12:19 -0700 Subject: [PATCH 16/70] feat(data-plane): aggregate metrics across processes, and report their cost MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The driver sees about a sixth of a step's data-plane traffic. The rollout actor writes the batch and the policy workers read it back per DP rank, both in other processes with their own counters, and the share the driver sees shrinks as DP grows. Measured on one sync-GRPO step (256x4096, DP=8): rollout actor (writes) 11.1 MB 1 proc 23.3% driver (writes) 3.2 MB 1 proc 6.7% driver (reads ) 4.7 MB 1 proc 10.0% policy workers (reads ) 3.6 MB 8 procs 60.0% cluster total 47.5 MB 100.0% what the driver sees 7.9 MB 16.7% `merge_snapshots()` sums them. This is what the accumulators were shaped for: latency lives in fixed histogram buckets and the latency/bandwidth model in sufficient statistics precisely so both add, and everything derived is recomputed from the merged totals rather than averaged -- averaging per-rank percentiles does not produce a cluster percentile. Counters sum, `max_*` fields take a maximum, and `peak_bytes_outstanding` is the one approximation (summing per-process peaks assumes they coincided, so it is an upper bound). `grpo_train_sync` now fans out to the driver and every worker and logs `data_plane/cluster/*`, falling back to `data_plane/driver/*` when only one process answers. The fan-out is best-effort: a rank that cannot answer is dropped rather than failing the step, because a metrics collection must never be able to take training down. Also adds `observability_overhead_ms` -- what the measurement itself cost, summed across processes, computed as the wrapper's wall time minus the time its inner client was working. One extra `monotonic` per op buys a number that sits next to `wall_ms`, so the bill is visible rather than resting on a benchmark run elsewhere. E2E over 10 real Ray processes in the sync-GRPO shape: cluster volume vs driver-only 3.0x (this harness; 6x analytically) gather cost 2.31 ms/step, ~1 kB per process observability_overhead_frac 0.05% of data-plane time Byte accounting unchanged, hash guard unchanged, put/get still 20/16 us. The trainer fan-out itself is not covered by a test — it needs a live worker group. `merge_snapshots`, `cluster_step_metrics` and the self-cost accounting are. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 31 ++- nemo_rl/data_plane/README.md | 23 +- nemo_rl/data_plane/__init__.py | 9 +- nemo_rl/data_plane/observability.py | 221 ++++++++++++++++++-- nemo_rl/data_plane/worker_mixin.py | 11 + nemo_rl/models/policy/tq_policy.py | 28 +++ tests/unit/data_plane/test_observability.py | 70 +++++++ 7 files changed, 366 insertions(+), 27 deletions(-) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index a91c2525667..57a10867960 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -78,7 +78,11 @@ from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message from nemo_rl.data.multimodal_utils import present_multimodal_fields from nemo_rl.data_plane.interfaces import KVBatchMeta -from nemo_rl.data_plane.observability import MetricsDataPlaneClient +from nemo_rl.data_plane.observability import ( + MetricsDataPlaneClient, + cluster_step_metrics, + merge_snapshots, +) from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS, DP_TRAIN_FIELDS from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import EnvironmentInterface @@ -382,6 +386,12 @@ def _compute_seq_logprob_error_metrics( return masking_data["sample_mask"], seq_logprob_error_metrics +# Previous cluster-wide reading, for per-step deltas. A single client owns +# its own previous snapshot; a cluster has no such owner, so the trainer +# holds it. +_PREV_CLUSTER_SNAPSHOT: dict[str, Any] = {} + + def _log_data_plane_metrics( policy: Any, logger: Logger, step: int, total_step_time: float ) -> None: @@ -409,8 +419,23 @@ def _log_data_plane_metrics( client = getattr(policy, "dp_client", None) if not isinstance(client, MetricsDataPlaneClient): return # observability disabled -> plain adapter - metrics = client.get_step_metrics(total_step_time) - logger.log_metrics(metrics, step, prefix="data_plane/driver") + + collect = getattr(policy, "collect_data_plane_snapshots", None) + snapshots = collect() if callable(collect) else [] + if len(snapshots) > 1: + # Cluster view: every rank's counters summed. Prefixed and reported + # instead of the driver's, not alongside, so there is one answer to + # "what did the data plane cost" rather than two that disagree by + # roughly the DP degree. + merged = merge_snapshots(snapshots) + metrics = cluster_step_metrics(merged, _PREV_CLUSTER_SNAPSHOT, total_step_time) + _PREV_CLUSTER_SNAPSHOT.clear() + _PREV_CLUSTER_SNAPSHOT.update(merged) + logger.log_metrics(metrics, step, prefix="data_plane/cluster") + else: + # Single process, or the fan-out could not reach the workers. + metrics = client.get_step_metrics(total_step_time) + logger.log_metrics(metrics, step, prefix="data_plane/driver") print( f" • data plane: {metrics['wall_ms']:.0f}ms " f"({100 * metrics['frac_of_step']:.1f}% of step), " diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 01060f1ba12..642d687d96d 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -457,10 +457,25 @@ by construction; the bulk traffic is the rollout actor's `kv_first_write` and the workers' per-DP-rank `get_samples`, and neither appears in these series. Do not read `comm_volume_mb` as cluster-wide volume. -`OpStats` is additive on purpose — the histogram buckets and the regression -sufficient statistics from every rank sum into one cluster-wide view — but -nothing collects them yet. That is an affordance the design leaves open, -not a feature that exists. +`OpStats` is additive on purpose, and `merge_snapshots()` uses it: the +histogram buckets and the regression sufficient statistics from every rank +*sum* into one cluster-wide view. Everything derived — percentiles, the +affine fit, throughput — is recomputed from the merged totals, never +averaged across ranks (averaging per-rank percentiles does not give a +cluster percentile). + +`grpo_train_sync` fans out to the driver and every policy worker, and logs +the combined result under `data_plane/cluster/` instead of the driver's +own. It falls back to `data_plane/driver/` when the fan-out finds only one +process. Measured: **~2.4 ms and ~1 kB per process per step** for 10 +processes, against a 6x wider view of the traffic. The fan-out is +best-effort — a rank that cannot answer is dropped rather than failing the +step. + +`observability_overhead_ms` reports what the measurement itself cost, +summed over every process: the wrapper's wall time minus the time its inner +client was working. It sits beside `wall_ms` so the bill is visible next to +what it bought (measured ~0.05% of data-plane time). **Units:** every duration is `_ms`, every volume is `_mb`, no exceptions — a chart mixing `wall_s` against `p99_ms` puts a 0.008 beside a 24.85 and diff --git a/nemo_rl/data_plane/__init__.py b/nemo_rl/data_plane/__init__.py index c97346ed4d2..67c7f2478de 100644 --- a/nemo_rl/data_plane/__init__.py +++ b/nemo_rl/data_plane/__init__.py @@ -27,7 +27,12 @@ KVBatchMeta, data_plane_supports_checkpointing, ) -from nemo_rl.data_plane.observability import MetricsDataPlaneClient, log_event +from nemo_rl.data_plane.observability import ( + MetricsDataPlaneClient, + cluster_step_metrics, + log_event, + merge_snapshots, +) __all__ = [ "DATA_PLANE_CHECKPOINT_SCHEMA_VERSION", @@ -36,7 +41,9 @@ "KVBatchMeta", "MetricsDataPlaneClient", "build_data_plane_client", + "cluster_step_metrics", "data_plane_supports_checkpointing", "log_event", + "merge_snapshots", "materialize", ] diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index f0b2f2384d8..201cb4ec5d9 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -417,6 +417,181 @@ def fit_latency_bandwidth(s: dict[str, Any]) -> dict[str, Any]: } +def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: + """Fill in the derived per-op fields, in place. + + Shared by :meth:`MetricsDataPlaneClient.snapshot` and + :func:`merge_snapshots` so a cluster-wide view is derived by exactly the + same arithmetic as a single process -- percentiles off the (summed) + histogram, the fit off the (summed) sufficient statistics. Nothing + derived is ever averaged across processes. + """ + for stats in by_op.values(): + calls = stats["calls"] + wall_ms = stats["wall_ms"] + stats["mean_ms"] = wall_ms / calls if calls else 0.0 + stats["mb_per_s"] = ( + (stats["n_bytes"] / 1e6) / (wall_ms / 1e3) if wall_ms else 0.0 + ) + stats["pct_of_total_ms"] = ( + 100.0 * wall_ms / total_wall_ms if total_wall_ms else 0.0 + ) + stats["fit"] = fit_latency_bandwidth(stats) + hist = stats["latency_hist"] + stats["p50_ms"] = percentile_from_hist(hist, 0.50) + stats["p99_ms"] = percentile_from_hist(hist, 0.99) + # Tail/mean ratio: a mean hides MR churn and queueing, which show + # up as p99 pulling away from the mean. + stats["tail_ratio_p99_mean"] = ( + stats["p99_ms"] / stats["mean_ms"] if stats["mean_ms"] > 0 else 0.0 + ) + + +# Snapshot fields that combine by summing, by taking a maximum, and the +# per-op ones of each kind. Everything else in a snapshot is derived and is +# recomputed from the merged totals rather than merged itself. +_SNAPSHOT_SUM = ( + "total_bytes", + "total_keys", + "total_ops", + "total_wall_ms", + "bytes_outstanding", + "peak_bytes_outstanding", + "n_keys_outstanding", + "self_ms", +) +_SNAPSHOT_MAX = ("max_bytes_per_key_seen", "last_put_bytes_per_key") +_OP_SUM = ( + "calls", + "errors", + "wall_ms", + "n_bytes", + "n_keys", + "ok_wall_ms", + "sum_bytes_sq", + "sum_bytes_ms", + "sum_ms_sq", +) +_OP_MAX = ("max_ms", "step_max_ms") + + +def merge_snapshots(snapshots: "list[dict[str, Any]]") -> dict[str, Any]: + """Combine per-process snapshots into one cluster-wide view. + + This is what the accumulators were shaped for. Latency lives in fixed + histogram buckets and the latency/bandwidth model lives in sufficient + statistics precisely so both *add*: summing 256 per-rank histograms + gives the true cluster distribution, which averaging 256 per-rank + percentiles cannot. Everything derived — percentiles, throughput, the + affine fit — is recomputed from the merged totals, never averaged. + + Counters sum. ``max_*`` fields take a maximum. ``peak_bytes_outstanding`` + is the one approximation: summing per-process peaks assumes they + coincided, so it is an upper bound on true cluster peak occupancy. + + Args: + snapshots: One :meth:`MetricsDataPlaneClient.snapshot` per process. + + Returns: + A snapshot-shaped dict covering every process, plus ``n_processes``. + """ + if not snapshots: + return {} + merged: dict[str, Any] = {k: 0 for k in _SNAPSHOT_SUM} + merged.update({k: 0 for k in _SNAPSHOT_MAX}) + hashes = { + k: 0 + for k in ( + "rows_recorded", + "rows_checked", + "rows_unverified", + "mismatches", + "fields_skipped", + ) + } + by_op: dict[str, dict[str, Any]] = {} + + for snap in snapshots: + for key in _SNAPSHOT_SUM: + merged[key] += snap.get(key, 0) + for key in _SNAPSHOT_MAX: + merged[key] = max(merged[key], snap.get(key, 0)) + for key in hashes: + hashes[key] += (snap.get("hash_verify") or {}).get(key, 0) + for op, stats in (snap.get("by_op") or {}).items(): + acc = by_op.setdefault( + op, + { + **{k: 0 for k in _OP_SUM}, + **{k: 0.0 for k in _OP_MAX}, + "latency_hist": [0] * (len(LATENCY_BUCKETS_MS) + 1), + }, + ) + for key in _OP_SUM: + acc[key] += stats.get(key, 0) + for key in _OP_MAX: + acc[key] = max(acc[key], stats.get(key, 0.0)) + for i, count in enumerate(stats.get("latency_hist") or []): + acc["latency_hist"][i] += count + + merged["by_op"] = by_op + merged["hash_verify"] = hashes + merged["n_processes"] = len(snapshots) + _derive_op_metrics(by_op, merged["total_wall_ms"]) + merged["bytes_written"] = sum(by_op[o]["n_bytes"] for o in _WRITE_OPS if o in by_op) + merged["bytes_read"] = sum(by_op[o]["n_bytes"] for o in _READ_OPS if o in by_op) + merged["comm_volume_bytes"] = merged["bytes_written"] + merged["bytes_read"] + return merged + + +def cluster_step_metrics( + merged: dict[str, Any], prev: dict[str, Any], step_time_s: float +) -> dict[str, float]: + """Per-step cluster metrics from two merged snapshots. + + The single-process equivalent of this lives on the client, which owns + its own previous reading. A cluster has no such owner, so the caller + holds ``prev`` and passes it back. + + ``observability_overhead_ms`` is what the measurement itself cost, + summed over every process -- the wrapper's wall time minus the time its + inner client was working. It sits beside ``wall_ms`` so the bill is + visible next to what it bought. + """ + wall_ms = merged["total_wall_ms"] - prev.get("total_wall_ms", 0.0) + overhead_ms = merged["self_ms"] - prev.get("self_ms", 0.0) + metrics: dict[str, float] = { + "wall_ms": wall_ms, + "frac_of_step": (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0, + "comm_volume_mb": ( + merged["comm_volume_bytes"] - prev.get("comm_volume_bytes", 0) + ) + / 1e6, + "bytes_written_mb": (merged["bytes_written"] - prev.get("bytes_written", 0)) + / 1e6, + "bytes_read_mb": (merged["bytes_read"] - prev.get("bytes_read", 0)) / 1e6, + "bytes_outstanding_mb": merged["bytes_outstanding"] / 1e6, + "n_processes": merged.get("n_processes", 0), + "observability_overhead_ms": overhead_ms, + "observability_overhead_frac": (overhead_ms / wall_ms if wall_ms > 0 else 0.0), + } + prev_ops = prev.get("by_op", {}) + for op, stats in merged["by_op"].items(): + prev_op = prev_ops.get(op, {}) + calls = stats["calls"] - prev_op.get("calls", 0) + if calls <= 0: + continue + metrics[f"{op}/calls"] = calls + metrics[f"{op}/wall_ms"] = stats["wall_ms"] - prev_op.get("wall_ms", 0.0) + metrics[f"{op}/max_ms"] = stats["max_ms"] + # Percentiles are worth reporting here and not per process: this + # histogram is the sum over every rank, so it is the real cluster + # distribution rather than one process's handful of calls. + metrics[f"{op}/p50_ms"] = stats["p50_ms"] + metrics[f"{op}/p99_ms"] = stats["p99_ms"] + return metrics + + def log_event(event: DataPlaneEvent) -> None: logger.info("data_plane_event: %s", event) @@ -502,6 +677,11 @@ class DataPlaneStats: # sudden spike in ``max_bytes_per_key_seen``. max_bytes_per_key_seen: int = 0 last_put_bytes_per_key: int = 0 + # What measuring cost. Wall time spent inside this wrapper minus the + # time the inner client was actually working, so a reader can see the + # observability bill next to the thing it is observing rather than + # taking a benchmark's word for it. + self_ms: float = 0.0 hash_verify: HashStats = field(default_factory=HashStats) @@ -544,6 +724,10 @@ def __init__( # of that same batch; the row count is what detects a shard read. self._batch_scope: dict[str, tuple[frozenset[str], int]] = {} self._hash_mismatches_logged = 0 + # Set by ``_emit`` to the inner client's wall time for the op just + # run, so the wrapping methods can subtract it and bill the rest to + # ``self_ms``. + self._last_inner_ms = 0.0 # Previous snapshot, for per-step deltas. Owned here rather than by a # caller: it is this client's prior reading, and keeping it here lets # every trainer use get_step_metrics() without copying the @@ -563,25 +747,7 @@ def snapshot(self) -> dict[str, Any]: out["n_keys_outstanding"] = sum( len(k) for k in self._keys_by_partition.values() ) - for op, s in out["by_op"].items(): - s["mean_ms"] = s["wall_ms"] / s["calls"] if s["calls"] else 0.0 - s["mb_per_s"] = ( - (s["n_bytes"] / 1e6) / (s["wall_ms"] / 1e3) if s["wall_ms"] else 0.0 - ) - s["pct_of_total_ms"] = ( - 100.0 * s["wall_ms"] / self._stats.total_wall_ms - if self._stats.total_wall_ms - else 0.0 - ) - s["fit"] = fit_latency_bandwidth(s) - h = s["latency_hist"] - s["p50_ms"] = percentile_from_hist(h, 0.50) - s["p99_ms"] = percentile_from_hist(h, 0.99) - # Tail/mean ratio: a mean hides MR churn and queueing, which show - # up as p99 pulling away from the mean. - s["tail_ratio_p99_mean"] = ( - s["p99_ms"] / s["mean_ms"] if s["mean_ms"] > 0 else 0.0 - ) + _derive_op_metrics(out["by_op"], self._stats.total_wall_ms) # Communication volume, derived from by_op so there is one source of # truth for bytes. Distinct from ``bytes_outstanding``, which is # occupancy (what is held) rather than traffic (what moved). @@ -751,6 +917,16 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: self._bytes_by_partition[partition_id] = total - freed self._stats.bytes_outstanding -= freed + def _bill_self(self, entered: float) -> None: + """Charge this wrapper for the time it spent that was not the RPC. + + One ``monotonic`` per op on top of the two ``_run`` already takes. + Measuring the measurement is worth that: the alternative is asking a + reader to trust a benchmark run on some other machine. + """ + elapsed_ms = (monotonic() - entered) * 1000.0 + self._stats.self_ms += elapsed_ms - self._last_inner_ms + # ── wire-in / wire-out fingerprinting (opt-in) ───────────────────── def _row_fingerprints( @@ -949,6 +1125,7 @@ def _emit( status: EventStatus, ) -> None: wall_ms = (monotonic() - t0) * 1000.0 + self._last_inner_ms = wall_ms on_event = self._on_event if on_event is not None: # Built lazily: with no sink registered nothing reads this dict. @@ -1044,6 +1221,7 @@ def claim_meta( ) def get_data(self, meta, select_fields=None): + entered = monotonic() out = self._run( "get_data", meta.partition_id, @@ -1052,6 +1230,7 @@ def get_data(self, meta, select_fields=None): ) if self._verify_tensor_hash: self._check_hashes(meta.partition_id, meta.sample_ids, out) + self._bill_self(entered) return out def check_consumption_status(self, partition_id, task_names): @@ -1062,6 +1241,7 @@ def check_consumption_status(self, partition_id, task_names): ) def put_samples(self, sample_ids, partition_id, fields=None, tags=None): + entered = monotonic() n_bytes = _td_bytes(fields) # Materialize once: ``_run`` consumes its lambda and we also need # to attribute bytes per sample after success. @@ -1084,9 +1264,11 @@ def put_samples(self, sample_ids, partition_id, fields=None, tags=None): # here keeps the check's own cost out of the op's ``wall_ms``. if self._verify_tensor_hash: self._record_hashes(partition_id, sample_ids_list, fields) + self._bill_self(entered) return out def get_samples(self, sample_ids, partition_id, select_fields): + entered = monotonic() sample_ids_list = _as_list(sample_ids) out = self._run( "get", @@ -1100,6 +1282,7 @@ def get_samples(self, sample_ids, partition_id, select_fields): ) if self._verify_tensor_hash: self._check_hashes(partition_id, sample_ids_list, out) + self._bill_self(entered) return out def list_sample_ids(self, partition_id: str) -> list[str]: diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index a9b7ab95932..a696d8ad836 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -384,6 +384,17 @@ def _forward_pad_seqlen(self, meta: "KVBatchMeta") -> int: """Cross-DP forward pad target, minted by :meth:`TQPolicy._stamp_pad_seqlen`.""" return int((meta.extra_info or {}).get(GLOBAL_FORWARD_PAD_SEQLEN, 0)) + def get_data_plane_snapshot(self) -> "dict[str, Any] | None": + """This rank's data-plane counters, for cluster-wide aggregation. + + Returns ``None`` when observability is off or no client exists, so + the driver can filter rather than special-case. The payload is + counters only (about 1 kB), not tensors. + """ + client = getattr(self, "_dp_client", None) + snapshot = getattr(client, "snapshot", None) + return snapshot() if callable(snapshot) else None + def _fetch( self, meta: "KVBatchMeta", diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index ae674a17c53..f8d216bf9ab 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -29,6 +29,7 @@ from __future__ import annotations +import logging import warnings from collections import Counter, defaultdict from contextlib import nullcontext @@ -231,6 +232,33 @@ def finish_step(self, meta: KVBatchMeta) -> None: """Drop this step's bulk from TQ. Mirror of :meth:`prepare_step`.""" self.discard_samples(meta.sample_ids, meta.partition_id) + def collect_data_plane_snapshots(self) -> list[dict[str, Any]]: + """This driver's data-plane counters plus every worker rank's. + + The driver sees roughly a sixth of a step's traffic — the rollout + actor writes the batch and the workers read it back per DP rank, + both in other processes with their own counters. Aggregating is what + turns these series from one process's slice into the cluster figure. + + Best effort by design: a rank that cannot answer is dropped rather + than failing the step, because a metrics fan-out must never be able + to take training down. Measured at ~2.4 ms and ~1 kB per process. + """ + snapshots: list[dict[str, Any]] = [] + client = getattr(self, "dp_client", None) + if hasattr(client, "snapshot"): + snapshots.append(client.snapshot()) + try: + futures = self.worker_group.run_all_workers_single_data( + "get_data_plane_snapshot" + ) + snapshots.extend( + s for s in self.worker_group.get_all_worker_results(futures) if s + ) + except Exception as exc: # noqa: BLE001 - metrics must never fail a step + logging.warning("data-plane snapshot fan-out failed: %s", exc) + return snapshots + # ── 1-hop entrypoints (KVBatchMeta in, no re-fan-out) ────────────────── def _with_route_fields( diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 8231daed595..b1ad19cc5a3 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -30,6 +30,8 @@ from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient from nemo_rl.data_plane.observability import ( MetricsDataPlaneClient, + cluster_step_metrics, + merge_snapshots, _estimate_encoded_bytes, _td_bytes, ) @@ -719,3 +721,71 @@ def test_hash_verification_off_by_default(): assert client._hash_by_partition == {} assert "hash/mismatches" not in client.get_step_metrics(1.0) client.close() + + +# ── cross-process aggregation ────────────────────────────────────────── + + +def _client_with(n_puts, n_bytes_each, wall_ms): + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + for _ in range(n_puts): + client._emit("put", "p", 1, n_bytes_each, monotonic() - wall_ms / 1e3, "ok") + return client + + +def test_merge_sums_counters_and_rederives_percentiles(): + """The accumulators are shaped to add: histograms and regression sums + from every rank combine into the true cluster distribution. Averaging + per-rank percentiles could not do this, which is the whole reason the + latency lives in fixed buckets rather than retained samples.""" + ranks = [_client_with(4, 1_000, 5.0) for _ in range(3)] + merged = merge_snapshots([c.snapshot() for c in ranks]) + + assert merged["n_processes"] == 3 + assert merged["by_op"]["put"]["calls"] == 12 # 3 ranks x 4 puts + assert merged["by_op"]["put"]["n_bytes"] == 12_000 + assert sum(merged["by_op"]["put"]["latency_hist"]) == 12 + # derived from the summed histogram, not averaged from the ranks + single = ranks[0].snapshot()["by_op"]["put"] + assert merged["by_op"]["put"]["p50_ms"] == pytest.approx(single["p50_ms"], rel=0.3) + for c in ranks: + c.close() + + +def test_merge_takes_max_for_max_fields(): + """A cluster's worst call is the worst any rank saw, not their sum.""" + slow = _client_with(1, 1_000, 40.0) + fast = _client_with(1, 1_000, 1.0) + merged = merge_snapshots([slow.snapshot(), fast.snapshot()]) + assert merged["by_op"]["put"]["max_ms"] >= 40.0 + assert merged["by_op"]["put"]["max_ms"] < 41.0, "max, not sum" + slow.close() + fast.close() + + +def test_merge_of_nothing_is_empty(): + assert merge_snapshots([]) == {} + + +def test_cluster_step_metrics_report_their_own_cost(): + """``observability_overhead_ms`` is the wrapper's own wall time minus + the inner client's, summed over processes — the bill for measuring, + sitting beside what it bought.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client.register_partition( + partition_id="p", fields=["x"], num_samples=4, consumer_tasks=["t"] + ) + for i in range(4): + client.put_samples( + sample_ids=[f"u{i}"], + partition_id="p", + fields=TensorDict({"x": torch.zeros(1, 512)}, batch_size=[1]), + ) + merged = merge_snapshots([client.snapshot()]) + metrics = cluster_step_metrics(merged, {}, 1.0) + + assert metrics["n_processes"] == 1 + assert metrics["observability_overhead_ms"] > 0, "measuring is never free" + assert 0.0 <= metrics["observability_overhead_frac"] <= 1.0 + assert "wall_ms" in metrics and "comm_volume_mb" in metrics + client.close() From 3b7d3f2b2b9e110f7590307175b558f895e3f130 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 23:16:49 -0700 Subject: [PATCH 17/70] fix(data-plane): bill the fan-out to observability, not just the wrapper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `observability_overhead_ms` claimed to be what measuring cost and reported a twentieth of it. It summed `self_ms` — the per-op wrapper time — and left out the snapshot fan-out added in the previous commit, which is the larger half: reported 0.1301 ms fan-out, uncounted 2.3100 ms true total 2.4401 ms -> understated 19x as a fraction 0.049% reported against 0.918% real A metric that names itself the cost of observability and omits the dominant term is worse than not having one, and it is the same failure as the flat percentiles: a number that reads as an answer while measuring something else. `cluster_step_metrics` now takes `collect_ms` and `grpo_train_sync` times the gather it already performs. `clear_samples` is billed too; it was the last per-step op doing bookkeeping outside the meter. The ratio is deliberately not clamped to 100%. Against a fast backend measuring can genuinely cost more than the operation measured — the unit tests hit exactly that against the no-op client, at 104% — and that is a signal worth surfacing rather than hiding behind a min(). E2E over 10 Ray processes, unchanged otherwise: observability_overhead_ms 2.417 ms observability_overhead_frac 0.92% gather 2.39 ms median cluster vs driver-only 3.0x Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 2 ++ nemo_rl/data_plane/README.md | 19 ++++++++++--- nemo_rl/data_plane/observability.py | 26 +++++++++++++----- tests/unit/data_plane/test_observability.py | 30 ++++++++++++++++++++- 4 files changed, 66 insertions(+), 11 deletions(-) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index 57a10867960..030458ce247 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -33,6 +33,7 @@ import gc import os +import time import warnings from typing import TYPE_CHECKING, Any, Optional @@ -421,6 +422,7 @@ def _log_data_plane_metrics( return # observability disabled -> plain adapter collect = getattr(policy, "collect_data_plane_snapshots", None) + collect_started = time.perf_counter() snapshots = collect() if callable(collect) else [] if len(snapshots) > 1: # Cluster view: every rank's counters summed. Prefixed and reported diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 642d687d96d..b121140a5b2 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -472,10 +472,21 @@ processes, against a 6x wider view of the traffic. The fan-out is best-effort — a rank that cannot answer is dropped rather than failing the step. -`observability_overhead_ms` reports what the measurement itself cost, -summed over every process: the wrapper's wall time minus the time its inner -client was working. It sits beside `wall_ms` so the bill is visible next to -what it bought (measured ~0.05% of data-plane time). +`observability_overhead_ms` reports what the measurement itself cost — the +whole bill, both halves: + +- every process's wrapper time (its wall time minus the time its inner + client was working), and +- the fan-out that gathered and merged the snapshots. + +The second is the larger. In the cross-process e2e the wrapper cost 0.13 ms +and the fan-out 2.31 ms, so a figure covering only the first understated by +19x. Measured whole: **~2.4 ms, about 0.9% of data-plane time** for 10 +processes. + +It is deliberately not clamped to 100%. Against a fast backend the ratio +can exceed 1, meaning measuring cost more than the operation measured — a +signal worth seeing rather than hiding. **Units:** every duration is `_ms`, every volume is `_mb`, no exceptions — a chart mixing `wall_s` against `p99_ms` puts a 0.008 beside a 24.85 and diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 201cb4ec5d9..e6d9ae054cd 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -545,7 +545,10 @@ def merge_snapshots(snapshots: "list[dict[str, Any]]") -> dict[str, Any]: def cluster_step_metrics( - merged: dict[str, Any], prev: dict[str, Any], step_time_s: float + merged: dict[str, Any], + prev: dict[str, Any], + step_time_s: float, + collect_ms: float = 0.0, ) -> dict[str, float]: """Per-step cluster metrics from two merged snapshots. @@ -553,13 +556,22 @@ def cluster_step_metrics( its own previous reading. A cluster has no such owner, so the caller holds ``prev`` and passes it back. - ``observability_overhead_ms`` is what the measurement itself cost, - summed over every process -- the wrapper's wall time minus the time its - inner client was working. It sits beside ``wall_ms`` so the bill is - visible next to what it bought. + ``observability_overhead_ms`` is the whole bill for measuring: every + process's wrapper time (its wall time minus the time its inner client + was working) plus ``collect_ms``, the fan-out that gathered and merged + the snapshots. Both halves matter and the second is the larger -- the + per-op wrapper costs tenths of a millisecond while the fan-out costs a + couple, so a figure covering only the first understates by an order of + magnitude and is worse than no figure at all. + + Args: + merged: Cluster-wide snapshot from :func:`merge_snapshots`. + prev: The previous merged snapshot, for differencing. + step_time_s: Step wall time, for ``frac_of_step``. + collect_ms: Wall time the caller spent gathering and merging. """ wall_ms = merged["total_wall_ms"] - prev.get("total_wall_ms", 0.0) - overhead_ms = merged["self_ms"] - prev.get("self_ms", 0.0) + overhead_ms = merged["self_ms"] - prev.get("self_ms", 0.0) + collect_ms metrics: dict[str, float] = { "wall_ms": wall_ms, "frac_of_step": (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0, @@ -1293,6 +1305,7 @@ def list_sample_ids(self, partition_id: str) -> list[str]: ) def clear_samples(self, sample_ids, partition_id): + entered = monotonic() sample_ids_list = _as_list(sample_ids) n_keys = len(sample_ids_list) if sample_ids_list is not None else 0 self._run( @@ -1302,6 +1315,7 @@ def clear_samples(self, sample_ids, partition_id): n_keys=n_keys, ) self._record_clear(partition_id, sample_ids_list) + self._bill_self(entered) def save_checkpoint( self, diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index b1ad19cc5a3..695aa838912 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -786,6 +786,34 @@ def test_cluster_step_metrics_report_their_own_cost(): assert metrics["n_processes"] == 1 assert metrics["observability_overhead_ms"] > 0, "measuring is never free" - assert 0.0 <= metrics["observability_overhead_frac"] <= 1.0 + # Deliberately not clamped to 1. Against this no-op inner client the RPC + # is instant, so measuring costs more than the thing measured and the + # ratio exceeds 100% -- which is exactly the signal worth surfacing. + # Against a real backend it lands near 0.01. + assert metrics["observability_overhead_frac"] > 0 assert "wall_ms" in metrics and "comm_volume_mb" in metrics client.close() + + +def test_cluster_overhead_includes_the_collection_fan_out(): + """The fan-out is the larger half of the bill. Reporting only the per-op + wrapper understated the real cost by ~19x in the cross-process e2e + (0.13 ms reported against 2.44 ms actually spent).""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client.register_partition( + partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["t"] + ) + client.put_samples( + sample_ids=["a"], + partition_id="p", + fields=TensorDict({"x": torch.zeros(1, 8)}, batch_size=[1]), + ) + merged = merge_snapshots([client.snapshot()]) + without = cluster_step_metrics(merged, {}, 1.0) + with_gather = cluster_step_metrics(merged, {}, 1.0, collect_ms=2.31) + + delta = ( + with_gather["observability_overhead_ms"] - without["observability_overhead_ms"] + ) + assert delta == pytest.approx(2.31, rel=1e-6) + client.close() From a12c2331fc6941e422273d1e9bebb10bc703cac8 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 23:22:41 -0700 Subject: [PATCH 18/70] fix(data-plane): a max below its own median, and a fraction above 1 Auditing the cross-process e2e output rather than reporting it turned up two contradictions that shipped. A maximum below its own median. The e2e printed put: max_ms 133.5 p50_ms 175.0 p99_ms 248.5 get: max_ms 0.0063 p50_ms 0.050 p99_ms 0.099 and not one of those four percentiles is data: 100 + 150*0.50 = 175, 100 + 150*0.99 = 248.5, 0 + 0.1*0.50 = 0.05. Two causes. The percentiles came off the cumulative histogram while the max was step-scoped, so they described different windows; and interpolation spreads a bucket's samples uniformly across it, so calls clustered low in a wide bucket read high -- 160 calls of exactly 120 ms all land in (100, 250] and interpolate to a p50 of 175, above every call observed. Both fixed. Percentiles are now differenced per step, so they share a window with the max, and clamped to it, because a true percentile cannot exceed the maximum and the maximum is measured exactly. They are withheld below 50 calls in the window; on a wide DP degree a step clears that easily, and on a narrow one silence beats bucket geometry. A fraction above 1. `frac_of_step` read 1.054 -- correct arithmetic on a `wall_ms` summed over ten processes that ran concurrently, and nonsense as "105% of the step". Renamed to `busy_frac_mean` and divided by the process count: the mean fraction of the step a process spent in the data plane, bounded, and the question people actually ask. The audit also found three ways the harness was lying, now fixed there: workers "read" through `lambda: None` so their timings measured nothing; the driver wrote the whole batch instead of the advantages delta, which is why bytes_written was exactly 2x bytes_read; and the 131 ms/put denominator is NoOpDataPlaneClient doing Python bookkeeping, not a wire. With those corrected the driver's share drops to ~9% of cluster volume and observability costs 1.2%. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 15 ++++++ nemo_rl/data_plane/observability.py | 52 ++++++++++++++++++--- tests/unit/data_plane/test_observability.py | 51 ++++++++++++++++++++ 3 files changed, 111 insertions(+), 7 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index b121140a5b2..c34a9a5ee65 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -464,6 +464,21 @@ affine fit, throughput — is recomputed from the merged totals, never averaged across ranks (averaging per-rank percentiles does not give a cluster percentile). +Two things read differently in the cluster view and are named to say so: + +- **`busy_frac_mean`**, not `frac_of_step`. `wall_ms` is summed over + processes that ran concurrently, so dividing it by one step's wall clock + exceeds 1 whenever they overlapped (measured 1.054 across ten processes). + The mean fraction of the step a process spent in the data plane is + bounded and answers the question people ask of it. +- **Percentiles are per step and clamped to the exact `max_ms`**, and are + withheld entirely below 50 calls in the window. Bucket interpolation + spreads a bucket's samples uniformly across it, so calls clustered low in + a wide bucket read high — 160 calls of 120 ms all land in `(100, 250]` + and interpolate to a p50 of 175, above every call observed and above the + max reported beside it. The max is measured exactly, so it is the tighter + bound. + `grpo_train_sync` fans out to the driver and every policy worker, and logs the combined result under `data_plane/cluster/` instead of the driver's own. It falls back to `data_plane/driver/` when the fan-out finds only one diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index e6d9ae054cd..7c327bcd1d7 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -96,6 +96,13 @@ class DataPlaneEvent(TypedDict): # few lines carry the identity of what broke. _MAX_HASH_MISMATCH_LOGS = 20 +# Calls needed in a window before a percentile off the histogram means +# anything. Below this the interpolation returns bucket geometry: one sample +# in (100, 250] yields p50 = 100 + 150*0.50 = 175 and p99 = 248.5 whatever +# the call actually took. Reporting that next to an exact ``max_ms`` produced +# a max below its own median. +_MIN_SAMPLES_FOR_PERCENTILE = 50 + class _FieldDigest(NamedTuple): """One fingerprint per row, plus how far it can be trusted. @@ -572,9 +579,18 @@ def cluster_step_metrics( """ wall_ms = merged["total_wall_ms"] - prev.get("total_wall_ms", 0.0) overhead_ms = merged["self_ms"] - prev.get("self_ms", 0.0) + collect_ms + # Not ``frac_of_step``: ``wall_ms`` here is the SUM over processes that + # ran concurrently, so dividing by one step's wall clock gives a number + # that exceeds 1 whenever they overlapped (measured 1.054 across ten + # processes) -- correct arithmetic, but it reads as "105% of the step". + # The mean fraction of the step a process spent in the data plane is + # bounded and answers the question people ask of it. + n_procs = max(merged.get("n_processes", 1), 1) metrics: dict[str, float] = { "wall_ms": wall_ms, - "frac_of_step": (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0, + "busy_frac_mean": ( + wall_ms / 1e3 / (step_time_s * n_procs) if step_time_s > 0 else 0.0 + ), "comm_volume_mb": ( merged["comm_volume_bytes"] - prev.get("comm_volume_bytes", 0) ) @@ -583,7 +599,7 @@ def cluster_step_metrics( / 1e6, "bytes_read_mb": (merged["bytes_read"] - prev.get("bytes_read", 0)) / 1e6, "bytes_outstanding_mb": merged["bytes_outstanding"] / 1e6, - "n_processes": merged.get("n_processes", 0), + "n_processes": n_procs, "observability_overhead_ms": overhead_ms, "observability_overhead_frac": (overhead_ms / wall_ms if wall_ms > 0 else 0.0), } @@ -596,11 +612,33 @@ def cluster_step_metrics( metrics[f"{op}/calls"] = calls metrics[f"{op}/wall_ms"] = stats["wall_ms"] - prev_op.get("wall_ms", 0.0) metrics[f"{op}/max_ms"] = stats["max_ms"] - # Percentiles are worth reporting here and not per process: this - # histogram is the sum over every rank, so it is the real cluster - # distribution rather than one process's handful of calls. - metrics[f"{op}/p50_ms"] = stats["p50_ms"] - metrics[f"{op}/p99_ms"] = stats["p99_ms"] + # Percentiles over THIS step's calls, summed across ranks, not over + # the lifetime: a cumulative percentile beside a step-scoped max is + # two different windows on one chart, and it showed a max below its + # own median. Emitted only when the window holds enough calls to + # out-resolve the buckets -- with a wide DP degree a step easily + # clears it, and on a narrow one the honest answer is silence. + step_hist = [ + now - was + for now, was in zip( + stats["latency_hist"], + prev_op.get("latency_hist") or [0] * len(stats["latency_hist"]), + ) + ] + if sum(step_hist) >= _MIN_SAMPLES_FOR_PERCENTILE: + # Clamped to the exact max. Interpolation spreads a bucket's + # samples uniformly across it, so calls clustered low in a wide + # bucket read high -- 160 calls of 120 ms all land in (100, 250] + # and interpolate to a p50 of 175, above every call observed. A + # true percentile cannot exceed the maximum, and the maximum is + # measured exactly, so it is the tighter bound. + ceiling = stats["max_ms"] + metrics[f"{op}/p50_ms"] = min( + percentile_from_hist(step_hist, 0.50), ceiling + ) + metrics[f"{op}/p99_ms"] = min( + percentile_from_hist(step_hist, 0.99), ceiling + ) return metrics diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 695aa838912..8f1291dfcc5 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -795,6 +795,23 @@ def test_cluster_step_metrics_report_their_own_cost(): client.close() +def test_cluster_busy_fraction_is_bounded(): + """``wall_ms`` is summed over processes that ran concurrently, so + dividing it by one step's wall clock exceeds 1 whenever they + overlapped — measured 1.054 across ten processes, which reads as + '105% of the step'. The mean per-process fraction is bounded.""" + # Physical fixture: each rank spends 500 ms inside a 1 s step. A rank + # cannot spend longer in the data plane than the step lasted, so a + # fixture that implies it would be testing the arithmetic on impossible + # input rather than the metric. + ranks = [_rank_with([100.0] * 5) for _ in range(10)] + metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 1.0) + + assert "frac_of_step" not in metrics + assert 0.0 <= metrics["busy_frac_mean"] <= 1.0, metrics["busy_frac_mean"] + assert metrics["n_processes"] == 10 + + def test_cluster_overhead_includes_the_collection_fan_out(): """The fan-out is the larger half of the bill. Reporting only the per-op wrapper understated the real cost by ~19x in the cross-process e2e @@ -817,3 +834,37 @@ def test_cluster_overhead_includes_the_collection_fan_out(): ) assert delta == pytest.approx(2.31, rel=1e-6) client.close() + + +def _rank_with(latencies_ms): + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + now = monotonic() + for ms in latencies_ms: + client._emit("put", "p", 1, 1_000, now - ms / 1e3, "ok") + return client.snapshot() + + +def test_cluster_percentiles_never_exceed_the_measured_max(): + """Bucket interpolation spreads a bucket's samples uniformly across it, + so calls clustered low in a wide bucket read high: 160 calls of 120 ms + all land in (100, 250] and interpolate to a p50 of 175 — above every + call observed, and above the exact max reported beside it. The max is + the tighter bound, so the percentiles are clamped to it.""" + merged = merge_snapshots([_rank_with([120.0] * 20) for _ in range(8)]) + metrics = cluster_step_metrics(merged, {}, 1.0) + + assert metrics["put/max_ms"] == pytest.approx(120.0, abs=2.0) + assert metrics["put/p50_ms"] <= metrics["put/max_ms"] + assert metrics["put/p99_ms"] <= metrics["put/max_ms"] + assert metrics["put/p50_ms"] <= metrics["put/p99_ms"] + + +def test_cluster_percentiles_withheld_below_a_useful_sample_count(): + """A percentile off a handful of calls is bucket geometry, not data. + Silence beats a number that looks like an answer.""" + few = merge_snapshots([_rank_with([120.0] * 3) for _ in range(2)]) # 6 calls + many = merge_snapshots([_rank_with([120.0] * 20) for _ in range(8)]) # 160 + + assert "put/p50_ms" not in cluster_step_metrics(few, {}, 1.0) + assert "put/max_ms" in cluster_step_metrics(few, {}, 1.0), "max always works" + assert "put/p50_ms" in cluster_step_metrics(many, {}, 1.0) From 7776ceed1f85c770725985b22884725de0714e6f Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 02:00:52 -0700 Subject: [PATCH 19/70] fix(data-plane): make every series say whether it is a delta or a level MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `comm_volume_mb` was a per-step delta and `bytes_outstanding_mb` an instantaneous level. Same suffix, same chart, nothing to tell them apart — so `bytes_outstanding_mb` climbing looked like the accumulation bug the other two would have if they behaved that way, and there was no way to read either with confidence. Every series now sits under a namespace that names its kind: step/ what happened during this step (a delta, resets) now/ what is true at this instant (a level, persists) so `step/comm_volume_mb` and `now/bytes_outstanding_mb` are unambiguous on an axis label. Applied to the single-process path too, which had the same mix. A test asserts nothing can be emitted outside the two namespaces. The rising `now/bytes_outstanding_mb` in the e2e is the metric working, not a fault: the harness writes its "train" partition every step and only ever clears the per-worker shard partitions, so bytes really are being retained. That is the leak signal this level exists for. `busy_frac_mean` is gone. It was an invented ratio that needed a paragraph to explain and still exceeded 1 on unphysical input. `step/wall_ms_per_process` says the same thing as a duration in ms, like every other time series here: how long the average process spent in the data plane this step. Also realigns the trainer's summary print, which still read the pre-rename keys and would have raised KeyError on the first step. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 5 +- nemo_rl/data_plane/README.md | 21 +++-- nemo_rl/data_plane/observability.py | 89 ++++++++++++-------- tests/unit/data_plane/test_observability.py | 91 +++++++++++++-------- 4 files changed, 130 insertions(+), 76 deletions(-) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index 030458ce247..9238825e25c 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -439,9 +439,8 @@ def _log_data_plane_metrics( metrics = client.get_step_metrics(total_step_time) logger.log_metrics(metrics, step, prefix="data_plane/driver") print( - f" • data plane: {metrics['wall_ms']:.0f}ms " - f"({100 * metrics['frac_of_step']:.1f}% of step), " - f"{metrics['comm_volume_mb']:.1f} MB moved" + f" • data plane: {metrics['step/wall_ms']:.0f}ms, " + f"{metrics['step/comm_volume_mb']:.1f} MB moved" ) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index c34a9a5ee65..a017599609b 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -464,13 +464,24 @@ affine fit, throughput — is recomputed from the merged totals, never averaged across ranks (averaging per-rank percentiles does not give a cluster percentile). -Two things read differently in the cluster view and are named to say so: +**Every series says what kind of number it is.** A per-step delta and an +instantaneous level shared the `_mb` suffix and a chart, with nothing to +tell them apart: -- **`busy_frac_mean`**, not `frac_of_step`. `wall_ms` is summed over +| namespace | meaning | example | +|---|---|---| +| `step/` | what happened during this step; resets | `step/comm_volume_mb` | +| `now/` | what is true at this instant; persists | `now/bytes_outstanding_mb` | + +A rising `now/bytes_outstanding_mb` is not an accumulation bug — it is the +leak signal the metric exists for: bytes put and never cleared. + +Two more read differently in the cluster view and are named to say so: + +- **`step/wall_ms_per_process`**, not a bare fraction. `wall_ms` sums processes that ran concurrently, so dividing it by one step's wall clock - exceeds 1 whenever they overlapped (measured 1.054 across ten processes). - The mean fraction of the step a process spent in the data plane is - bounded and answers the question people ask of it. + exceeded 1 whenever they overlapped (measured 1.054 across ten + processes). Per process it is a duration in ms, like everything else. - **Percentiles are per step and clamped to the exact `max_ms`**, and are withheld entirely below 50 calls in the window. Bucket interpolation spreads a bucket's samples uniformly across it, so calls clustered low in diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 7c327bcd1d7..dd622e83849 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -586,22 +586,38 @@ def cluster_step_metrics( # The mean fraction of the step a process spent in the data plane is # bounded and answers the question people ask of it. n_procs = max(merged.get("n_processes", 1), 1) + # Namespaced by what kind of number it is, because the unit alone does + # not say. ``comm_volume_mb`` was a per-step delta and + # ``bytes_outstanding_mb`` an instantaneous level -- same suffix, same + # chart, nothing to tell them apart. + # step/ what happened during this step (a delta, resets each step) + # now/ what is true at this instant (a level, persists) metrics: dict[str, float] = { - "wall_ms": wall_ms, - "busy_frac_mean": ( - wall_ms / 1e3 / (step_time_s * n_procs) if step_time_s > 0 else 0.0 - ), - "comm_volume_mb": ( + "step/wall_ms": wall_ms, + # ``wall_ms`` sums processes that ran concurrently, so it can exceed + # the step's own wall clock. Divided by the process count it reads + # as "how long the average process spent in the data plane this + # step" -- in ms like everything else, and needing no gloss. + "step/wall_ms_per_process": wall_ms / n_procs, + "step/comm_volume_mb": ( merged["comm_volume_bytes"] - prev.get("comm_volume_bytes", 0) ) / 1e6, - "bytes_written_mb": (merged["bytes_written"] - prev.get("bytes_written", 0)) + "step/bytes_written_mb": ( + merged["bytes_written"] - prev.get("bytes_written", 0) + ) / 1e6, - "bytes_read_mb": (merged["bytes_read"] - prev.get("bytes_read", 0)) / 1e6, - "bytes_outstanding_mb": merged["bytes_outstanding"] / 1e6, - "n_processes": n_procs, - "observability_overhead_ms": overhead_ms, - "observability_overhead_frac": (overhead_ms / wall_ms if wall_ms > 0 else 0.0), + "step/bytes_read_mb": (merged["bytes_read"] - prev.get("bytes_read", 0)) / 1e6, + # A level, not a delta: bytes held in TQ right now, put minus + # cleared. It climbs when something is not being cleared, which is + # the leak signal it exists for. A rising line here is the metric + # working, not an accumulation bug. + "now/bytes_outstanding_mb": merged["bytes_outstanding"] / 1e6, + "now/n_processes": n_procs, + "step/observability_overhead_ms": overhead_ms, + "step/observability_overhead_frac": ( + overhead_ms / wall_ms if wall_ms > 0 else 0.0 + ), } prev_ops = prev.get("by_op", {}) for op, stats in merged["by_op"].items(): @@ -609,9 +625,9 @@ def cluster_step_metrics( calls = stats["calls"] - prev_op.get("calls", 0) if calls <= 0: continue - metrics[f"{op}/calls"] = calls - metrics[f"{op}/wall_ms"] = stats["wall_ms"] - prev_op.get("wall_ms", 0.0) - metrics[f"{op}/max_ms"] = stats["max_ms"] + metrics[f"step/{op}/calls"] = calls + metrics[f"step/{op}/wall_ms"] = stats["wall_ms"] - prev_op.get("wall_ms", 0.0) + metrics[f"step/{op}/max_ms"] = stats["max_ms"] # Percentiles over THIS step's calls, summed across ranks, not over # the lifetime: a cumulative percentile beside a step-scoped max is # two different windows on one chart, and it showed a max below its @@ -633,10 +649,10 @@ def cluster_step_metrics( # true percentile cannot exceed the maximum, and the maximum is # measured exactly, so it is the tighter bound. ceiling = stats["max_ms"] - metrics[f"{op}/p50_ms"] = min( + metrics[f"step/{op}/p50_ms"] = min( percentile_from_hist(step_hist, 0.50), ceiling ) - metrics[f"{op}/p99_ms"] = min( + metrics[f"step/{op}/p99_ms"] = min( percentile_from_hist(step_hist, 0.99), ceiling ) return metrics @@ -835,31 +851,40 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: # axis. GB was the same problem one dimension over -- a realistic # step moved 0.00017 GB. metrics: dict[str, float] = { - "wall_ms": wall_ms, - "frac_of_step": (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0, - "comm_volume_mb": vol / 1e6, - "bytes_written_mb": (snap["bytes_written"] - prev.get("bytes_written", 0)) + "step/wall_ms": wall_ms, + "step/frac_of_step": ( + (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0 + ), + "step/comm_volume_mb": vol / 1e6, + "step/bytes_written_mb": ( + snap["bytes_written"] - prev.get("bytes_written", 0) + ) + / 1e6, + "step/bytes_read_mb": (snap["bytes_read"] - prev.get("bytes_read", 0)) / 1e6, - "bytes_read_mb": (snap["bytes_read"] - prev.get("bytes_read", 0)) / 1e6, - "bytes_outstanding_mb": snap["bytes_outstanding"] / 1e6, + # A level, not a delta -- see the cluster path for why the two + # need to be distinguishable on a chart. + "now/bytes_outstanding_mb": snap["bytes_outstanding"] / 1e6, } if self._verify_tensor_hash: hv, prev_hv = snap["hash_verify"], prev.get("hash_verify", {}) - metrics["hash/rows_checked"] = hv["rows_checked"] - prev_hv.get( + metrics["step/hash/rows_checked"] = hv["rows_checked"] - prev_hv.get( "rows_checked", 0 ) - metrics["hash/rows_recorded"] = hv["rows_recorded"] - prev_hv.get( + metrics["step/hash/rows_recorded"] = hv["rows_recorded"] - prev_hv.get( "rows_recorded", 0 ) - metrics["hash/rows_unverified"] = hv["rows_unverified"] - prev_hv.get( + metrics["step/hash/rows_unverified"] = hv["rows_unverified"] - prev_hv.get( "rows_unverified", 0 ) - metrics["hash/mismatches"] = hv["mismatches"] - prev_hv.get("mismatches", 0) + metrics["step/hash/mismatches"] = hv["mismatches"] - prev_hv.get( + "mismatches", 0 + ) # Logged because a guard that quietly stops covering a field is # worse than no guard: it reports 0 mismatches and reads as # clean. A step where this climbs is a step where something # stopped being checked. - metrics["hash/fields_skipped"] = hv["fields_skipped"] - prev_hv.get( + metrics["step/hash/fields_skipped"] = hv["fields_skipped"] - prev_hv.get( "fields_skipped", 0 ) prev_ops = prev.get("by_op", {}) @@ -873,13 +898,13 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: # out rather than logged: mean is wall_ms/calls, and a dashboard # can divide. ``snapshot()`` still carries the full picture, # percentiles included, for a one-off inspection. - metrics[f"{op}/calls"] = calls - metrics[f"{op}/wall_ms"] = op_ms + metrics[f"step/{op}/calls"] = calls + metrics[f"step/{op}/wall_ms"] = op_ms # ``max_ms`` rather than p50/p99. Those come off a histogram # that is never reset, so per step they were a lifetime figure # that goes flat, quantised to bucket edges. The max is exact # and says the same thing at the handful of calls per step. - metrics[f"{op}/max_ms"] = step_maxima.get(op, 0.0) + metrics[f"step/{op}/max_ms"] = step_maxima.get(op, 0.0) fit = st["fit"] if fit.get("model_trustworthy"): # The step's time split into the two things that cause it, @@ -896,8 +921,8 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: # milliseconds. op_bytes = st["n_bytes"] - prev_op.get("n_bytes", 0) ms_per_byte = 1.0 / (fit["bandwidth_mb_s"] * 1e3) - metrics[f"{op}/overhead_ms"] = fit["fixed_ms"] * calls - metrics[f"{op}/transfer_ms"] = ms_per_byte * op_bytes + metrics[f"step/{op}/overhead_ms"] = fit["fixed_ms"] * calls + metrics[f"step/{op}/transfer_ms"] = ms_per_byte * op_bytes return metrics def bytes_outstanding_by_partition(self) -> dict[str, int]: diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 8f1291dfcc5..7d3f3956880 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -395,7 +395,7 @@ def test_step_metrics_use_one_unit_per_dimension(): ) keys = set(client.get_step_metrics(1.0)) assert not [k for k in keys if k.endswith(("_s", "_gb", "_kb", "_us"))], keys - assert "wall_ms" in keys and "comm_volume_mb" in keys + assert "step/wall_ms" in keys and "step/comm_volume_mb" in keys client.close() @@ -411,9 +411,9 @@ def test_step_metrics_tail_is_exact_not_bucketed(): client._emit("put", "p", 1, 8, monotonic() - 0.030, "ok") # a 30 ms call metrics = client.get_step_metrics(1.0) - assert "put/p99_ms" not in metrics and "put/p50_ms" not in metrics - assert metrics["put/max_ms"] >= 30.0 - assert metrics["put/max_ms"] != pytest.approx(24.85, abs=0.5), "bucket edge" + assert "put/p99_ms" not in metrics and "step/put/p50_ms" not in metrics + assert metrics["step/put/max_ms"] >= 30.0 + assert metrics["step/put/max_ms"] != pytest.approx(24.85, abs=0.5), "bucket edge" # the cumulative view still carries percentiles for one-off inspection assert "p99_ms" in client.snapshot()["by_op"]["put"] client.close() @@ -442,9 +442,9 @@ def test_latency_breakdown_stacks_to_wall_ms(): assert fit["bandwidth_mb_s"] == pytest.approx(mb_per_s, rel=0.05) # the two components are the split, in ms, and they add up - total = metrics["put/overhead_ms"] + metrics["put/transfer_ms"] - assert total == pytest.approx(metrics["put/wall_ms"], rel=0.05) - assert "put/overhead_frac" not in metrics, "a ratio is derivable from these" + total = metrics["step/put/overhead_ms"] + metrics["step/put/transfer_ms"] + assert total == pytest.approx(metrics["step/put/wall_ms"], rel=0.05) + assert "step/put/overhead_frac" not in metrics, "a ratio is derivable from these" client.close() @@ -454,9 +454,9 @@ def test_step_max_is_scoped_to_the_step(): reported max must fall again when a step is quicker.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient()) client._emit("put", "p", 1, 8, monotonic() - 0.050, "ok") # slow step - slow = client.get_step_metrics(1.0)["put/max_ms"] + slow = client.get_step_metrics(1.0)["step/put/max_ms"] client._emit("put", "p", 1, 8, monotonic() - 0.001, "ok") # quick step - quick = client.get_step_metrics(1.0)["put/max_ms"] + quick = client.get_step_metrics(1.0)["step/put/max_ms"] assert slow >= 50.0 assert quick < slow, "step max must reset, not carry the lifetime worst" @@ -544,7 +544,7 @@ def test_hash_verification_detects_corruption(): hv = client.snapshot()["hash_verify"] assert hv["mismatches"] == 1 - assert client.get_step_metrics(1.0)["hash/mismatches"] == 1 + assert client.get_step_metrics(1.0)["step/hash/mismatches"] == 1 client.close() @@ -663,7 +663,7 @@ def test_hash_incomparable_field_is_counted_not_dropped(): hv = client.snapshot()["hash_verify"] assert hv["mismatches"] == 0, "must not cry wolf on an incomparable field" assert hv["fields_skipped"] == 1, "the drop must be visible" - assert client.get_step_metrics(1.0)["hash/fields_skipped"] == 1 + assert client.get_step_metrics(1.0)["step/hash/fields_skipped"] == 1 class _RaggedOnReadClient(NoOpDataPlaneClient): @@ -719,7 +719,7 @@ def test_hash_verification_off_by_default(): assert client.snapshot()["hash_verify"]["rows_recorded"] == 0 assert client._hash_by_partition == {} - assert "hash/mismatches" not in client.get_step_metrics(1.0) + assert "step/hash/mismatches" not in client.get_step_metrics(1.0) client.close() @@ -784,32 +784,28 @@ def test_cluster_step_metrics_report_their_own_cost(): merged = merge_snapshots([client.snapshot()]) metrics = cluster_step_metrics(merged, {}, 1.0) - assert metrics["n_processes"] == 1 - assert metrics["observability_overhead_ms"] > 0, "measuring is never free" + assert metrics["now/n_processes"] == 1 + assert metrics["step/observability_overhead_ms"] > 0, "measuring is never free" # Deliberately not clamped to 1. Against this no-op inner client the RPC # is instant, so measuring costs more than the thing measured and the # ratio exceeds 100% -- which is exactly the signal worth surfacing. # Against a real backend it lands near 0.01. - assert metrics["observability_overhead_frac"] > 0 - assert "wall_ms" in metrics and "comm_volume_mb" in metrics + assert metrics["step/observability_overhead_frac"] > 0 + assert "step/wall_ms" in metrics and "step/comm_volume_mb" in metrics client.close() -def test_cluster_busy_fraction_is_bounded(): - """``wall_ms`` is summed over processes that ran concurrently, so - dividing it by one step's wall clock exceeds 1 whenever they - overlapped — measured 1.054 across ten processes, which reads as - '105% of the step'. The mean per-process fraction is bounded.""" - # Physical fixture: each rank spends 500 ms inside a 1 s step. A rank - # cannot spend longer in the data plane than the step lasted, so a - # fixture that implies it would be testing the arithmetic on impossible - # input rather than the metric. +def test_cluster_time_is_per_process_not_a_bare_fraction(): + """``wall_ms`` sums processes that ran concurrently, so dividing it by + one step's wall clock exceeded 1 whenever they overlapped and read as + "105% of the step". Reported per process it is a duration in ms, like + everything else, and needs no gloss.""" ranks = [_rank_with([100.0] * 5) for _ in range(10)] metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 1.0) - assert "frac_of_step" not in metrics - assert 0.0 <= metrics["busy_frac_mean"] <= 1.0, metrics["busy_frac_mean"] - assert metrics["n_processes"] == 10 + assert "busy_frac_mean" not in metrics and "frac_of_step" not in metrics + assert metrics["step/wall_ms_per_process"] == pytest.approx(500.0, rel=0.1) + assert metrics["now/n_processes"] == 10 def test_cluster_overhead_includes_the_collection_fan_out(): @@ -830,7 +826,8 @@ def test_cluster_overhead_includes_the_collection_fan_out(): with_gather = cluster_step_metrics(merged, {}, 1.0, collect_ms=2.31) delta = ( - with_gather["observability_overhead_ms"] - without["observability_overhead_ms"] + with_gather["step/observability_overhead_ms"] + - without["step/observability_overhead_ms"] ) assert delta == pytest.approx(2.31, rel=1e-6) client.close() @@ -853,10 +850,10 @@ def test_cluster_percentiles_never_exceed_the_measured_max(): merged = merge_snapshots([_rank_with([120.0] * 20) for _ in range(8)]) metrics = cluster_step_metrics(merged, {}, 1.0) - assert metrics["put/max_ms"] == pytest.approx(120.0, abs=2.0) - assert metrics["put/p50_ms"] <= metrics["put/max_ms"] - assert metrics["put/p99_ms"] <= metrics["put/max_ms"] - assert metrics["put/p50_ms"] <= metrics["put/p99_ms"] + assert metrics["step/put/max_ms"] == pytest.approx(120.0, abs=2.0) + assert metrics["step/put/p50_ms"] <= metrics["step/put/max_ms"] + assert metrics["step/put/p99_ms"] <= metrics["step/put/max_ms"] + assert metrics["step/put/p50_ms"] <= metrics["step/put/p99_ms"] def test_cluster_percentiles_withheld_below_a_useful_sample_count(): @@ -865,6 +862,28 @@ def test_cluster_percentiles_withheld_below_a_useful_sample_count(): few = merge_snapshots([_rank_with([120.0] * 3) for _ in range(2)]) # 6 calls many = merge_snapshots([_rank_with([120.0] * 20) for _ in range(8)]) # 160 - assert "put/p50_ms" not in cluster_step_metrics(few, {}, 1.0) - assert "put/max_ms" in cluster_step_metrics(few, {}, 1.0), "max always works" - assert "put/p50_ms" in cluster_step_metrics(many, {}, 1.0) + assert "step/put/p50_ms" not in cluster_step_metrics(few, {}, 1.0) + assert "step/put/max_ms" in cluster_step_metrics(few, {}, 1.0), "max always works" + assert "step/put/p50_ms" in cluster_step_metrics(many, {}, 1.0) + + +def test_cluster_series_declare_delta_or_level(): + """A per-step delta and an instantaneous level shared the ``_mb`` suffix + and a chart, with nothing to tell them apart. Every series now sits + under ``step/`` or ``now/`` so the kind is on the axis label.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client.register_partition( + partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["t"] + ) + client.put_samples( + sample_ids=["a"], + partition_id="p", + fields=TensorDict({"x": torch.zeros(1, 8)}, batch_size=[1]), + ) + metrics = cluster_step_metrics(merge_snapshots([client.snapshot()]), {}, 1.0) + + unlabelled = [k for k in metrics if not k.startswith(("step/", "now/"))] + assert not unlabelled, unlabelled + assert "now/bytes_outstanding_mb" in metrics, "a level" + assert "step/comm_volume_mb" in metrics, "a delta" + client.close() From c473775d3f2f06cdbcc9acdc86a826706abf6ad6 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 11:53:52 -0700 Subject: [PATCH 20/70] fix(data-plane): free only the keys a clear actually held Chasing why `now/bytes_outstanding_mb` climbed turned up a real bug in the release path, though not the one causing the climb. `_record_clear` sized the pro-rata release from `len(keys)` -- what the caller listed -- rather than from how many of those keys the partition was actually holding. A clear naming uids already dropped, or belonging to another partition, therefore released bytes this partition never held. Measured: clearing 50 live keys alongside 50 unknown ones freed 66666 of 100000 bytes from a partition that had lost half its keys. It over-frees, so it was pushing occupancy down rather than up, and the whole-partition and clean-subset paths were both correct -- which is why the earlier reconciliation checks passed. Now sized from `live.intersection(keys)`. A randomized test drives 50 interleaved put/partial-clear sequences and asserts every partition lands back at zero. The climb itself was the harness: it wrote its "train" partition every step and only ever cleared the per-worker shard partitions, so bytes really were being retained and the level was reporting that correctly. Adding the `finish_step`-equivalent clear flattens it to 0.0 MB across all 14 steps. That is the signal this level exists for. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 11 +++- tests/unit/data_plane/test_observability.py | 63 +++++++++++++++++++++ 2 files changed, 71 insertions(+), 3 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index dd622e83849..024a72257b8 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -980,15 +980,20 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: if live is None: return total = self._bytes_by_partition.get(partition_id, 0) + # Count what was actually live, not what the caller listed. A clear + # may name uids already dropped or belonging elsewhere, and billing + # those released bytes this partition never held: clearing 50 live + # keys alongside 50 unknown ones freed two thirds of a partition + # that had lost half its keys. + removed = len(live.intersection(keys)) if keys is not None else len(live) if keys is not None: - live -= live.intersection(keys) + live -= set(keys) if keys is None or not live: freed = total del self._keys_by_partition[partition_id] self._bytes_by_partition.pop(partition_id, None) else: - dropped = len(keys) - freed = total * dropped // (len(live) + dropped) + freed = total * removed // (len(live) + removed) if removed else 0 self._bytes_by_partition[partition_id] = total - freed self._stats.bytes_outstanding -= freed diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 7d3f3956880..0170e92bd5f 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -887,3 +887,66 @@ def test_cluster_series_declare_delta_or_level(): assert "now/bytes_outstanding_mb" in metrics, "a level" assert "step/comm_volume_mb" in metrics, "a delta" client.close() + + +def test_clear_frees_only_what_was_actually_live(): + """A clear may name uids already dropped, or belonging to another + partition. Billing those releases bytes this partition never held: + clearing 50 live keys alongside 50 unknown ones freed two thirds of a + partition that had lost half its keys.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + ids = [f"u{i}" for i in range(100)] + client.register_partition( + partition_id="p", fields=["x"], num_samples=100, consumer_tasks=["t"] + ) + client.put_samples( + sample_ids=ids, + partition_id="p", + fields=TensorDict({"x": torch.zeros(100, 250)}, batch_size=[100]), + ) + total = client.snapshot()["bytes_outstanding"] + + client.clear_samples( + sample_ids=ids[:50] + [f"unknown{i}" for i in range(50)], partition_id="p" + ) + assert client.snapshot()["bytes_outstanding"] == total // 2 + + client.clear_samples(sample_ids=ids[50:], partition_id="p") + assert client.snapshot()["bytes_outstanding"] == 0 + client.close() + + +def test_outstanding_reconciles_over_random_put_clear_sequences(): + """Interleaved puts and partial clears must always land back at zero; + the pro-rata release is only sound if it does.""" + import random + + rng = random.Random(0) + for _ in range(50): + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client.register_partition( + partition_id="p", fields=["x"], num_samples=5000, consumer_tasks=["t"] + ) + live: set[str] = set() + for _ in range(rng.randint(1, 6)): + batch = list( + dict.fromkeys( + f"k{rng.randint(0, 60)}" for _ in range(rng.randint(1, 20)) + ) + ) + client.put_samples( + sample_ids=batch, + partition_id="p", + fields=TensorDict( + {"x": torch.zeros(len(batch), 250)}, batch_size=[len(batch)] + ), + ) + live |= set(batch) + if live and rng.random() < 0.5: + drop = rng.sample(sorted(live), k=rng.randint(1, len(live))) + client.clear_samples(sample_ids=drop, partition_id="p") + live -= set(drop) + if live: + client.clear_samples(sample_ids=sorted(live), partition_id="p") + assert client.snapshot()["bytes_outstanding"] == 0 + client.close() From f0e2a4f2766d56ceb2d4433d4e4e7fc55c70d1ff Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 12:27:31 -0700 Subject: [PATCH 21/70] fix(data-plane): clamp percentiles where they are derived, not at one caller MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An audit over every emitted series — identities, ground truth, physical bounds — found `snapshot()` still publishing percentiles above their own maximum: register p50 0.050 p99 0.099 max 0.011 p50 is 4.5x the max put p50 40.6 p99 242.5 max 145.6 p99 exceeds the max clear p50 0.175 p99 0.248 max 0.141 p50 exceeds the max The clamp added earlier went into `cluster_step_metrics` only, so the two paths that derive from `_derive_op_metrics` — `snapshot()`, documented as the cumulative view for one-off inspection, and `merge_snapshots()` — carried the raw interpolation. Moved to `_derive_op_metrics`, where every consumer inherits it. The audit passed 44/44 while printing those numbers, because the check I wrote for that invariant carried `or s["calls"] < 50` and excused itself on exactly the cases that were wrong. A check with an escape clause for the failing input is not a check. Tightened to `p50 <= p99 <= max`, no exemption, and it now holds for all four ops. Everything else reconciled against an independent tally: bytes_written and bytes_read to the byte, per-op call counts, n_keys_outstanding against a live-key set, sum(by_op wall_ms) == total_wall_ms, sum(latency_hist) == calls, merge-of-one == the snapshot it merged, and every cluster series namespaced. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 23 +++++++++++++-------- tests/unit/data_plane/test_observability.py | 19 +++++++++++++++++ 2 files changed, 33 insertions(+), 9 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 024a72257b8..e29a2133bb0 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -445,8 +445,16 @@ def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: ) stats["fit"] = fit_latency_bandwidth(stats) hist = stats["latency_hist"] - stats["p50_ms"] = percentile_from_hist(hist, 0.50) - stats["p99_ms"] = percentile_from_hist(hist, 0.99) + # Clamped to the exact max, here rather than at one call site, so + # every consumer of a snapshot inherits it. Bucket interpolation + # spreads a bucket's samples uniformly across it, so calls clustered + # low in a wide bucket read high -- five register calls of 0.011 ms + # all land in (0, 0.1] and interpolate to a p50 of 0.05, four times + # the slowest call that happened. A true percentile cannot exceed + # the maximum and the maximum is measured exactly. + ceiling = stats["max_ms"] if stats["max_ms"] > 0 else float("inf") + stats["p50_ms"] = min(percentile_from_hist(hist, 0.50), ceiling) + stats["p99_ms"] = min(percentile_from_hist(hist, 0.99), ceiling) # Tail/mean ratio: a mean hides MR churn and queueing, which show # up as p99 pulling away from the mean. stats["tail_ratio_p99_mean"] = ( @@ -642,13 +650,10 @@ def cluster_step_metrics( ) ] if sum(step_hist) >= _MIN_SAMPLES_FOR_PERCENTILE: - # Clamped to the exact max. Interpolation spreads a bucket's - # samples uniformly across it, so calls clustered low in a wide - # bucket read high -- 160 calls of 120 ms all land in (100, 250] - # and interpolate to a p50 of 175, above every call observed. A - # true percentile cannot exceed the maximum, and the maximum is - # measured exactly, so it is the tighter bound. - ceiling = stats["max_ms"] + # Same clamp as :func:`_derive_op_metrics`, applied again here + # because this percentile comes off the step's histogram delta + # rather than the cumulative one that function derived. + ceiling = stats["max_ms"] if stats["max_ms"] > 0 else float("inf") metrics[f"step/{op}/p50_ms"] = min( percentile_from_hist(step_hist, 0.50), ceiling ) diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 0170e92bd5f..251c073da7f 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -950,3 +950,22 @@ def test_outstanding_reconciles_over_random_put_clear_sequences(): client.clear_samples(sample_ids=sorted(live), partition_id="p") assert client.snapshot()["bytes_outstanding"] == 0 client.close() + + +def test_snapshot_percentiles_never_exceed_the_measured_max(): + """The clamp lives in ``_derive_op_metrics``, not at one call site, so + ``snapshot()`` and ``merge_snapshots()`` inherit it. Without it five + register calls of 0.011 ms all land in (0, 0.1] and interpolate to a + p50 of 0.05 — four times the slowest call that happened — on a public + surface documented as the cumulative view.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + now = monotonic() + for _ in range(5): + client._emit("register", "p", 1, 0, now - 0.011 / 1e3, "ok") + + for view in (client.snapshot(), merge_snapshots([client.snapshot()])): + stats = view["by_op"]["register"] + assert stats["p50_ms"] <= stats["max_ms"], stats + assert stats["p99_ms"] <= stats["max_ms"], stats + assert stats["p50_ms"] <= stats["p99_ms"], stats + client.close() From c87ebd433a83190bfef1266485199896465d64d6 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 12:42:55 -0700 Subject: [PATCH 22/70] refactor(data-plane): cleanup pass, and two bugs it turned up Quality pass over the eleven commits since the last one. Two of the four reviewers found defects rather than cleanup, and both were live: The cross-process fan-out never worked. `collect_data_plane_snapshots` paired `worker_group.run_all_workers_single_data`, which returns a list of ObjectRefs, with `get_all_worker_results`, which wants a MultiWorkerFuture and calls `.get_results()` on it. The AttributeError went straight into the broad `except` that exists so metrics can never fail a step, so every run silently returned the driver's snapshot alone -- the exact one-process view the aggregation was built to replace. `TQPolicy` inherits `Policy.run_all_workers_single_data`, which does the `ray.get` itself; that is the call. `observability_overhead_ms` was still the understated figure its own docstring warns about. The previous commit added `collect_ms`, timed the fan-out into `collect_started`, and then called `cluster_step_metrics` without passing it, so the parameter defaulted to 0.0 and the metric omitted the larger half of the bill it had just been taught to include. Cleanup applied: - `_step_deltas` shared by both step-metric paths. They built the same five series from the same arithmetic; the point of the `step/`/`now/` convention is that the two views cannot drift on names, and duplicating them was the way that would happen. - `_clamped_percentiles` shared by `_derive_op_metrics` and `cluster_step_metrics`. This also fixes a real gap: the sample-count gate had only ever been applied on the cluster path, so `snapshot()` published percentiles off three calls. - `_record_clear` builds one set instead of two (measured 7.34us -> 5.31us on the set work, 18% of the op). - Deleted `bytes_outstanding_by_partition` (no caller anywhere) and `step_max_ms` from the merge tuple (merged, never read). - Deduplicated the "why the old metric was wrong" narrative, which was restated six times across module docstring, both step-metric functions and the tests. It lives in README.md. - `_client_with` and `_rank_with` in the tests were the same helper. Altitude fixes: - `_PREV_CLUSTER_SNAPSHOT` moved from module state in grpo_sync onto the policy, beside the client whose counters it differences. Two trainers in one process shared one `prev` and would have produced negative deltas. - `logging.warning` -> the module logger, matching every other module here. - README no longer claims metrics are "on by default" full stop: the exemplar carries the default, so recipes inheriting it get it and a config without an `observability:` block still falls back to False. - README now says the cluster view omits the rollout actor, which builds its own client off the worker group -- so `kv_first_write` is not in these totals. Not acted on: `merge_snapshots`' histogram loop (measured 4.20 ms at 256 ranks, of which the loop is 26%; it runs once per step and is already billed into `collect_ms`), and the snapshot fields with no reader (`peak_bytes_outstanding`, `total_keys`, `mb_per_s`, ...) which are the documented one-off inspection surface rather than dead code. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 55 +++---- nemo_rl/data_plane/README.md | 11 +- nemo_rl/data_plane/observability.py | 162 ++++++++++---------- nemo_rl/models/policy/tq_policy.py | 20 ++- tests/unit/data_plane/test_observability.py | 27 ++-- 5 files changed, 135 insertions(+), 140 deletions(-) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index 9238825e25c..f1218596c74 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -387,35 +387,27 @@ def _compute_seq_logprob_error_metrics( return masking_data["sample_mask"], seq_logprob_error_metrics -# Previous cluster-wide reading, for per-step deltas. A single client owns -# its own previous snapshot; a cluster has no such owner, so the trainer -# holds it. -_PREV_CLUSTER_SNAPSHOT: dict[str, Any] = {} - - def _log_data_plane_metrics( policy: Any, logger: Logger, step: int, total_step_time: float ) -> None: """Log this step's data-plane cost. No-op unless observability is enabled. - Delta arithmetic, unit conversion and the per-op namespace live on - ``MetricsDataPlaneClient.get_step_metrics`` -- the same split as - ``VllmGeneration.get_step_metrics`` -- so the async trainer can log the - same metrics with one call against its own client. - - Scoped to the driver, and the prefix says so. Every process builds its - own client -- ``tq_policy`` here, plus one per policy worker - (``worker_mixin``), one in the rollout actor, one on the - single-controller path -- and each keeps independent counters. The - driver issues roughly one op of each kind per step; the bulk traffic is - elsewhere (``kv_first_write`` in the rollout actor writes the rollout, - per-DP-rank ``get_samples`` in the workers read it), and none of it - appears here. A ``data_plane/`` prefix would read as cluster-wide - totals, which these are not. - - ``OpStats`` is additive on purpose -- histograms and regression sums - from every rank sum into one cluster-wide view -- but nothing collects - them yet, so that remains a design affordance rather than a feature. + Prefers the cluster view -- the driver's counters plus every policy + worker's, summed -- and falls back to the driver's alone when the + fan-out reaches only one process. Reported one way or the other, never + both, so there is a single answer to "what did the data plane cost" + rather than two that disagree by roughly the DP degree. + + The prefix names the scope because the two differ by a lot: the driver + issues about one op of each kind per step while the bulk traffic is the + workers' per-DP-rank ``get_samples``. Note that even the cluster view + omits the rollout actor, which builds its own client and is not on the + worker group -- so ``kv_first_write`` is not in these totals. + + The previous reading lives on the policy, alongside the client whose + counters it differences, rather than in module state: two trainers in + one process would otherwise interleave one ``prev`` and produce + negative deltas. """ client = getattr(policy, "dp_client", None) if not isinstance(client, MetricsDataPlaneClient): @@ -425,14 +417,15 @@ def _log_data_plane_metrics( collect_started = time.perf_counter() snapshots = collect() if callable(collect) else [] if len(snapshots) > 1: - # Cluster view: every rank's counters summed. Prefixed and reported - # instead of the driver's, not alongside, so there is one answer to - # "what did the data plane cost" rather than two that disagree by - # roughly the DP degree. merged = merge_snapshots(snapshots) - metrics = cluster_step_metrics(merged, _PREV_CLUSTER_SNAPSHOT, total_step_time) - _PREV_CLUSTER_SNAPSHOT.clear() - _PREV_CLUSTER_SNAPSHOT.update(merged) + # The fan-out is part of what observability costs, and the larger + # part: omitting it reported a twentieth of the real bill. + collect_ms = (time.perf_counter() - collect_started) * 1e3 + prev = getattr(policy, "_prev_cluster_snapshot", {}) + metrics = cluster_step_metrics( + merged, prev, total_step_time, collect_ms=collect_ms + ) + policy._prev_cluster_snapshot = merged logger.log_metrics(metrics, step, prefix="data_plane/cluster") else: # Single process, or the fan-out could not reach the workers. diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index a017599609b..6ec44032661 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -492,7 +492,9 @@ Two more read differently in the cluster view and are named to say so: `grpo_train_sync` fans out to the driver and every policy worker, and logs the combined result under `data_plane/cluster/` instead of the driver's -own. It falls back to `data_plane/driver/` when the fan-out finds only one +own. **It does not reach the rollout actor**, which builds its own client +and is not on the worker group — so `kv_first_write`, the write of the +whole rollout, is not in these totals. It falls back to `data_plane/driver/` when the fan-out finds only one process. Measured: **~2.4 ms and ~1 kB per process per step** for 10 processes, against a 6x wider view of the traffic. The fan-out is best-effort — a rank that cannot answer is dropped rather than failing the @@ -541,8 +543,11 @@ carries — 256 ragged rows, 12 MB, jagged per-token fields as 0.1% of a 59 ms operation. What is left is dominated by the per-key attribution `clear_samples` needs to undo. -This is **on by default**, and only engages when `data_plane.enabled` is -true, so it costs nothing for runs that don't use the data plane. There is +This is **on in the exemplar config**, which is where a v1 `TypedDict` +default lives — so recipes inheriting `grpo_math_1B.yaml` get it, and a +config with no `observability:` block still falls back to `False` at the +factory. It only engages when `data_plane.enabled` is true either way, so +it costs nothing for runs that don't use the data plane. There is no default per-op sink: `get_step_metrics()` is the surface, and `grpo_train_sync` logs it once a step under the `data_plane/` prefix — so the series reach whatever backends the run has enabled (wandb, TensorBoard, diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index e29a2133bb0..f4a6bdbd646 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -424,6 +424,46 @@ def fit_latency_bandwidth(s: dict[str, Any]) -> dict[str, Any]: } +def _step_deltas(snap: dict[str, Any], prev: dict[str, Any]) -> dict[str, float]: + """The five series both step-metric paths report, identically. + + Shared so the single-process and cluster views cannot drift on series + names -- which is the whole point of the ``step/``/``now/`` convention + they publish under. + """ + return { + "step/wall_ms": snap["total_wall_ms"] - prev.get("total_wall_ms", 0.0), + "step/comm_volume_mb": ( + snap["comm_volume_bytes"] - prev.get("comm_volume_bytes", 0) + ) + / 1e6, + "step/bytes_written_mb": (snap["bytes_written"] - prev.get("bytes_written", 0)) + / 1e6, + "step/bytes_read_mb": (snap["bytes_read"] - prev.get("bytes_read", 0)) / 1e6, + "now/bytes_outstanding_mb": snap["bytes_outstanding"] / 1e6, + } + + +def _clamped_percentiles(hist: list[int], max_ms: float) -> tuple[float, float] | None: + """p50 and p99 from ``hist``, or ``None`` when the sample is too thin. + + Two corrections, both needed wherever a percentile is taken. Below + :data:`_MIN_SAMPLES_FOR_PERCENTILE` the interpolation returns bucket + geometry rather than data -- one sample in (100, 250] yields a p50 of + 175 whatever the call took -- so there is no answer to give. And the + interpolation spreads a bucket's samples uniformly across it, so calls + clustered low in a wide bucket read high, above the exact maximum + measured beside them; the maximum is the tighter bound. + """ + if sum(hist) < _MIN_SAMPLES_FOR_PERCENTILE: + return None + ceiling = max_ms if max_ms > 0 else float("inf") + return ( + min(percentile_from_hist(hist, 0.50), ceiling), + min(percentile_from_hist(hist, 0.99), ceiling), + ) + + def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: """Fill in the derived per-op fields, in place. @@ -445,16 +485,8 @@ def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: ) stats["fit"] = fit_latency_bandwidth(stats) hist = stats["latency_hist"] - # Clamped to the exact max, here rather than at one call site, so - # every consumer of a snapshot inherits it. Bucket interpolation - # spreads a bucket's samples uniformly across it, so calls clustered - # low in a wide bucket read high -- five register calls of 0.011 ms - # all land in (0, 0.1] and interpolate to a p50 of 0.05, four times - # the slowest call that happened. A true percentile cannot exceed - # the maximum and the maximum is measured exactly. - ceiling = stats["max_ms"] if stats["max_ms"] > 0 else float("inf") - stats["p50_ms"] = min(percentile_from_hist(hist, 0.50), ceiling) - stats["p99_ms"] = min(percentile_from_hist(hist, 0.99), ceiling) + pct = _clamped_percentiles(hist, stats["max_ms"]) + stats["p50_ms"], stats["p99_ms"] = pct if pct else (0.0, 0.0) # Tail/mean ratio: a mean hides MR churn and queueing, which show # up as p99 pulling away from the mean. stats["tail_ratio_p99_mean"] = ( @@ -487,7 +519,7 @@ def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: "sum_bytes_ms", "sum_ms_sq", ) -_OP_MAX = ("max_ms", "step_max_ms") +_OP_MAX = ("max_ms",) def merge_snapshots(snapshots: "list[dict[str, Any]]") -> dict[str, Any]: @@ -572,12 +604,9 @@ def cluster_step_metrics( holds ``prev`` and passes it back. ``observability_overhead_ms`` is the whole bill for measuring: every - process's wrapper time (its wall time minus the time its inner client - was working) plus ``collect_ms``, the fan-out that gathered and merged - the snapshots. Both halves matter and the second is the larger -- the - per-op wrapper costs tenths of a millisecond while the fan-out costs a - couple, so a figure covering only the first understates by an order of - magnitude and is worse than no figure at all. + process's wrapper time plus ``collect_ms``, the fan-out that gathered + the snapshots. The fan-out is the larger half; omitting it understates + by an order of magnitude. Args: merged: Cluster-wide snapshot from :func:`merge_snapshots`. @@ -594,39 +623,22 @@ def cluster_step_metrics( # The mean fraction of the step a process spent in the data plane is # bounded and answers the question people ask of it. n_procs = max(merged.get("n_processes", 1), 1) - # Namespaced by what kind of number it is, because the unit alone does - # not say. ``comm_volume_mb`` was a per-step delta and - # ``bytes_outstanding_mb`` an instantaneous level -- same suffix, same - # chart, nothing to tell them apart. - # step/ what happened during this step (a delta, resets each step) - # now/ what is true at this instant (a level, persists) - metrics: dict[str, float] = { - "step/wall_ms": wall_ms, - # ``wall_ms`` sums processes that ran concurrently, so it can exceed - # the step's own wall clock. Divided by the process count it reads - # as "how long the average process spent in the data plane this - # step" -- in ms like everything else, and needing no gloss. - "step/wall_ms_per_process": wall_ms / n_procs, - "step/comm_volume_mb": ( - merged["comm_volume_bytes"] - prev.get("comm_volume_bytes", 0) - ) - / 1e6, - "step/bytes_written_mb": ( - merged["bytes_written"] - prev.get("bytes_written", 0) - ) - / 1e6, - "step/bytes_read_mb": (merged["bytes_read"] - prev.get("bytes_read", 0)) / 1e6, - # A level, not a delta: bytes held in TQ right now, put minus - # cleared. It climbs when something is not being cleared, which is - # the leak signal it exists for. A rising line here is the metric - # working, not an accumulation bug. - "now/bytes_outstanding_mb": merged["bytes_outstanding"] / 1e6, - "now/n_processes": n_procs, - "step/observability_overhead_ms": overhead_ms, - "step/observability_overhead_frac": ( - overhead_ms / wall_ms if wall_ms > 0 else 0.0 - ), - } + # step/ is a delta over this step; now/ is a level at this instant. + # The unit alone does not distinguish them -- see README.md. + metrics = _step_deltas(merged, prev) + metrics.update( + { + # ``wall_ms`` sums processes that ran concurrently, so it can + # exceed the step's own wall clock. Per process it reads as a + # duration in ms, like everything else here. + "step/wall_ms_per_process": wall_ms / n_procs, + "now/n_processes": n_procs, + "step/observability_overhead_ms": overhead_ms, + "step/observability_overhead_frac": ( + overhead_ms / wall_ms if wall_ms > 0 else 0.0 + ), + } + ) prev_ops = prev.get("by_op", {}) for op, stats in merged["by_op"].items(): prev_op = prev_ops.get(op, {}) @@ -649,17 +661,11 @@ def cluster_step_metrics( prev_op.get("latency_hist") or [0] * len(stats["latency_hist"]), ) ] - if sum(step_hist) >= _MIN_SAMPLES_FOR_PERCENTILE: - # Same clamp as :func:`_derive_op_metrics`, applied again here - # because this percentile comes off the step's histogram delta - # rather than the cumulative one that function derived. - ceiling = stats["max_ms"] if stats["max_ms"] > 0 else float("inf") - metrics[f"step/{op}/p50_ms"] = min( - percentile_from_hist(step_hist, 0.50), ceiling - ) - metrics[f"step/{op}/p99_ms"] = min( - percentile_from_hist(step_hist, 0.99), ceiling - ) + # Off the step's histogram delta, not the cumulative one + # :func:`_derive_op_metrics` used. + pct = _clamped_percentiles(step_hist, stats["max_ms"]) + if pct: + metrics[f"step/{op}/p50_ms"], metrics[f"step/{op}/p99_ms"] = pct return metrics @@ -849,28 +855,15 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: bucket.step_max_ms = 0.0 wall_ms = snap["total_wall_ms"] - prev.get("total_wall_ms", 0.0) - vol = snap["comm_volume_bytes"] - prev.get("comm_volume_bytes", 0) # Every duration is ms and every volume is MB, with no exceptions: # a chart that mixes wall_s against p99_ms puts a 0.008 next to a # 24.85 and reads as a bug in the data plane rather than in the # axis. GB was the same problem one dimension over -- a realistic # step moved 0.00017 GB. - metrics: dict[str, float] = { - "step/wall_ms": wall_ms, - "step/frac_of_step": ( - (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0 - ), - "step/comm_volume_mb": vol / 1e6, - "step/bytes_written_mb": ( - snap["bytes_written"] - prev.get("bytes_written", 0) - ) - / 1e6, - "step/bytes_read_mb": (snap["bytes_read"] - prev.get("bytes_read", 0)) - / 1e6, - # A level, not a delta -- see the cluster path for why the two - # need to be distinguishable on a chart. - "now/bytes_outstanding_mb": snap["bytes_outstanding"] / 1e6, - } + metrics = _step_deltas(snap, prev) + metrics["step/frac_of_step"] = ( + (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0 + ) if self._verify_tensor_hash: hv, prev_hv = snap["hash_verify"], prev.get("hash_verify", {}) metrics["step/hash/rows_checked"] = hv["rows_checked"] - prev_hv.get( @@ -930,10 +923,6 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: metrics[f"step/{op}/transfer_ms"] = ms_per_byte * op_bytes return metrics - def bytes_outstanding_by_partition(self) -> dict[str, int]: - """Per-partition breakdown of currently-held bytes.""" - return dict(self._bytes_by_partition) - def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: """Attribute put bytes per key so a later ``clear_samples`` can subtract. @@ -990,9 +979,12 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: # those released bytes this partition never held: clearing 50 live # keys alongside 50 unknown ones freed two thirds of a partition # that had lost half its keys. - removed = len(live.intersection(keys)) if keys is not None else len(live) - if keys is not None: - live -= set(keys) + if keys is None: + removed = len(live) + else: + dropped = live.intersection(keys) + live -= dropped + removed = len(dropped) if keys is None or not live: freed = total del self._keys_by_partition[partition_id] diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index f8d216bf9ab..6fd66312cbb 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -98,6 +98,9 @@ def _aggregate_train_results(results: list[dict[str, Any]]) -> dict[str, Any]: # dispatcher only waits for completion — no aggregation needed. +logger = logging.getLogger(__name__) + + class TQPolicy(TQDriverMixin, Policy): """TQ-mediated counterpart to :class:`Policy`. @@ -249,14 +252,17 @@ def collect_data_plane_snapshots(self) -> list[dict[str, Any]]: if hasattr(client, "snapshot"): snapshots.append(client.snapshot()) try: - futures = self.worker_group.run_all_workers_single_data( - "get_data_plane_snapshot" - ) - snapshots.extend( - s for s in self.worker_group.get_all_worker_results(futures) if s - ) + # ``Policy.run_all_workers_single_data`` already does the + # ``ray.get``. Pairing the worker-group call with + # ``get_all_worker_results`` does not work -- the former returns + # a list of ObjectRefs and the latter wants a MultiWorkerFuture -- + # and the broad except below swallowed the AttributeError, so + # only the driver's snapshot was ever returned. + ranks = self.run_all_workers_single_data("get_data_plane_snapshot") except Exception as exc: # noqa: BLE001 - metrics must never fail a step - logging.warning("data-plane snapshot fan-out failed: %s", exc) + logger.warning("data-plane snapshot fan-out failed: %s", exc) + else: + snapshots.extend(s for s in ranks if s) return snapshots # ── 1-hop entrypoints (KVBatchMeta in, no re-fan-out) ────────────────── diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 251c073da7f..99b252c8b03 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -726,19 +726,26 @@ def test_hash_verification_off_by_default(): # ── cross-process aggregation ────────────────────────────────────────── -def _client_with(n_puts, n_bytes_each, wall_ms): +def _rank_client(latencies_ms): + """A client that has seen one put per entry, at that latency.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - for _ in range(n_puts): - client._emit("put", "p", 1, n_bytes_each, monotonic() - wall_ms / 1e3, "ok") + now = monotonic() + for ms in latencies_ms: + client._emit("put", "p", 1, 1_000, now - ms / 1e3, "ok") return client +def _rank_with(latencies_ms): + """Its snapshot, for the merge tests.""" + return _rank_client(latencies_ms).snapshot() + + def test_merge_sums_counters_and_rederives_percentiles(): """The accumulators are shaped to add: histograms and regression sums from every rank combine into the true cluster distribution. Averaging per-rank percentiles could not do this, which is the whole reason the latency lives in fixed buckets rather than retained samples.""" - ranks = [_client_with(4, 1_000, 5.0) for _ in range(3)] + ranks = [_rank_client([5.0] * 4) for _ in range(3)] merged = merge_snapshots([c.snapshot() for c in ranks]) assert merged["n_processes"] == 3 @@ -754,8 +761,8 @@ def test_merge_sums_counters_and_rederives_percentiles(): def test_merge_takes_max_for_max_fields(): """A cluster's worst call is the worst any rank saw, not their sum.""" - slow = _client_with(1, 1_000, 40.0) - fast = _client_with(1, 1_000, 1.0) + slow = _rank_client([40.0]) + fast = _rank_client([1.0]) merged = merge_snapshots([slow.snapshot(), fast.snapshot()]) assert merged["by_op"]["put"]["max_ms"] >= 40.0 assert merged["by_op"]["put"]["max_ms"] < 41.0, "max, not sum" @@ -833,14 +840,6 @@ def test_cluster_overhead_includes_the_collection_fan_out(): client.close() -def _rank_with(latencies_ms): - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - now = monotonic() - for ms in latencies_ms: - client._emit("put", "p", 1, 1_000, now - ms / 1e3, "ok") - return client.snapshot() - - def test_cluster_percentiles_never_exceed_the_measured_max(): """Bucket interpolation spreads a bucket's samples uniformly across it, so calls clustered low in a wide bucket read high: 160 calls of 120 ms From 0a539ffa96786cdbfe6df63fbe075af4b26a2a19 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 12:48:25 -0700 Subject: [PATCH 23/70] feat(data-plane): log a per-op breakdown table, not just series MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Eight line charts answer "how did put's wall time trend". They do not answer "where did this step go", which is a table — and reading a table off eight charts is the wrong tool. `data_plane/{cluster,driver}/breakdown` now carries one row per op, ordered by wall time so the expensive one is the first line read: op calls wall_ms max_ms p50_ms p99_ms put 64 3998 114.9 58.93 114.9 get 240 2916 20 12.76 20 clear 32 43.42 2.03 — — register 16 11.95 1.077 — — Cells are empty rather than zero where a series was withheld — a percentile below the sample gate, a fit that is not trustworthy. A zero there would read as a measurement. `breakdown_table` reshapes the flat metrics dict that is already logged rather than the snapshot behind it, so the table and the series cannot disagree: anything absent from one is absent from the other. `Logger.log_table` is added following the existing `log_plot`/`log_histogram` shape, but concrete rather than abstract — only wandb has a table type, and an abstract method would force every other backend, and any out-of-tree one, to write a stub. Others skip it. Wired best-effort in `grpo_train_sync`: a panel must never take a step down. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 21 ++++++++ nemo_rl/data_plane/README.md | 18 +++++++ nemo_rl/data_plane/__init__.py | 2 + nemo_rl/data_plane/observability.py | 57 +++++++++++++++++++++ nemo_rl/utils/logger.py | 38 ++++++++++++++ tests/unit/data_plane/test_observability.py | 40 +++++++++++++++ 6 files changed, 176 insertions(+) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index f1218596c74..72bc3598882 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -32,6 +32,7 @@ from __future__ import annotations import gc +import logging import os import time import warnings @@ -81,6 +82,7 @@ from nemo_rl.data_plane.interfaces import KVBatchMeta from nemo_rl.data_plane.observability import ( MetricsDataPlaneClient, + breakdown_table, cluster_step_metrics, merge_snapshots, ) @@ -387,6 +389,23 @@ def _compute_seq_logprob_error_metrics( return masking_data["sample_mask"], seq_logprob_error_metrics +def _log_breakdown( + logger: Logger, metrics: dict[str, float], step: int, name: str +) -> None: + """Log the per-op breakdown as a table beside the series. + + Best effort: a backend without a table type skips it, and a failure here + must not take a step down over a visualisation. + """ + try: + columns, rows = breakdown_table(metrics) + except Exception as exc: # noqa: BLE001 - a panel must never fail a step + logging.getLogger(__name__).warning("data-plane breakdown failed: %s", exc) + else: + if rows: + logger.log_table(columns, rows, step, name) + + def _log_data_plane_metrics( policy: Any, logger: Logger, step: int, total_step_time: float ) -> None: @@ -427,10 +446,12 @@ def _log_data_plane_metrics( ) policy._prev_cluster_snapshot = merged logger.log_metrics(metrics, step, prefix="data_plane/cluster") + _log_breakdown(logger, metrics, step, "data_plane/cluster/breakdown") else: # Single process, or the fan-out could not reach the workers. metrics = client.get_step_metrics(total_step_time) logger.log_metrics(metrics, step, prefix="data_plane/driver") + _log_breakdown(logger, metrics, step, "data_plane/driver/breakdown") print( f" • data plane: {metrics['step/wall_ms']:.0f}ms, " f"{metrics['step/comm_volume_mb']:.1f} MB moved" diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 6ec44032661..ec8a3887931 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -464,6 +464,24 @@ affine fit, throughput — is recomputed from the merged totals, never averaged across ranks (averaging per-rank percentiles does not give a cluster percentile). +**A per-op breakdown table** is logged alongside the series, under +`data_plane/{cluster,driver}/breakdown` — one row per op, ordered by wall +time so the expensive one reads first: + +| op | calls | wall_ms | max_ms | p50_ms | p99_ms | +|---|---:|---:|---:|---:|---:| +| put | 64 | 3998 | 114.9 | 58.93 | 114.9 | +| get | 240 | 2916 | 20.0 | 12.76 | 20.0 | +| clear | 32 | 43.42 | 2.03 | — | — | +| register | 16 | 11.95 | 1.08 | — | — | + +A stack of line charts answers "how did put's wall time trend"; this +answers "where did the step go", which is a table. Cells are empty rather +than zero where a series was withheld (a percentile below the sample gate, +a fit that is not trustworthy) — a zero would read as a measurement. It is +built from the same metrics dict that is logged, so the table and the +series cannot disagree. Only wandb renders it; other backends skip it. + **Every series says what kind of number it is.** A per-step delta and an instantaneous level shared the `_mb` suffix and a chart, with nothing to tell them apart: diff --git a/nemo_rl/data_plane/__init__.py b/nemo_rl/data_plane/__init__.py index 67c7f2478de..46c5a888cb3 100644 --- a/nemo_rl/data_plane/__init__.py +++ b/nemo_rl/data_plane/__init__.py @@ -29,6 +29,7 @@ ) from nemo_rl.data_plane.observability import ( MetricsDataPlaneClient, + breakdown_table, cluster_step_metrics, log_event, merge_snapshots, @@ -41,6 +42,7 @@ "KVBatchMeta", "MetricsDataPlaneClient", "build_data_plane_client", + "breakdown_table", "cluster_step_metrics", "data_plane_supports_checkpointing", "log_event", diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index f4a6bdbd646..20e6f035344 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -669,6 +669,63 @@ def cluster_step_metrics( return metrics +# Per-op columns worth a row in the breakdown, in the order they read. +# ``overhead_ms``/``transfer_ms`` are only present when the affine fit is +# trustworthy, and ``p50_ms``/``p99_ms`` only above the sample gate, so a +# row carries None where a series was withheld rather than a zero that +# would read as a measurement. +_BREAKDOWN_COLUMNS = ( + "calls", + "wall_ms", + "max_ms", + "overhead_ms", + "transfer_ms", + "p50_ms", + "p99_ms", +) + + +def breakdown_table( + metrics: dict[str, float], +) -> tuple[list[str], list[list[Any]]]: + """Reshape the flat per-op series into one row per op. + + A stack of line charts answers "how did put's wall time trend"; the + question this feeds is "where did this step's time go, across ops, at a + glance" -- which is a table, and reading it off eight separate charts is + the wrong tool. + + Built from the metrics dict that is logged rather than from the snapshot + it came from, so the table and the series can never disagree: a value + withheld from the series (a percentile below the sample gate, a fit that + is not trustworthy) is absent from the table too. + + Args: + metrics: A flat ``step/{op}/{field}`` dict from + :meth:`MetricsDataPlaneClient.get_step_metrics` or + :func:`cluster_step_metrics`. + + Returns: + ``(columns, rows)`` for :meth:`Logger.log_table`, ops sorted by + descending wall time so the expensive one is the first line read. + """ + per_op: dict[str, dict[str, float]] = {} + for key, value in metrics.items(): + parts = key.split("/") + if len(parts) != 3 or parts[0] != "step": + continue + _, op, field = parts + if field in _BREAKDOWN_COLUMNS: + per_op.setdefault(op, {})[field] = value + rows = [ + [op, *(stats.get(col) for col in _BREAKDOWN_COLUMNS)] + for op, stats in sorted( + per_op.items(), key=lambda kv: -kv[1].get("wall_ms", 0.0) + ) + ] + return ["op", *_BREAKDOWN_COLUMNS], rows + + def log_event(event: DataPlaneEvent) -> None: logger.info("data_plane_event: %s", event) diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index 675406af33a..dc341a5ce5b 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -129,6 +129,17 @@ def log_histogram(self, histogram: list[Any], step: int, name: str) -> None: """Log histogram metrics.""" pass + def log_table( + self, columns: list[str], rows: list[list[Any]], step: int, name: str + ) -> None: + """Log a table of rows. Backends that have no table type skip it. + + Concrete rather than abstract: only wandb renders tables natively, + and making this abstract would force every other backend -- and any + out-of-tree one -- to write a stub. + """ + return None + @abstractmethod def log_plot(self, figure: plt.Figure, step: int, name: str) -> None: """Log a matplotlib figure.""" @@ -402,6 +413,19 @@ def log_hyperparams(self, params: Mapping[str, Any]) -> None: """ self.run.config.update(params, allow_val_change=True) + def log_table( + self, columns: list[str], rows: list[list[Any]], step: int, name: str + ) -> None: + """Log a table to wandb. + + Args: + columns: Column headers + rows: One list of values per row + step: Global step value + name: Panel name + """ + self.run.log({name: wandb.Table(columns=columns, data=rows)}, step=step) + def log_plot(self, figure: plt.Figure, step: int, name: str) -> None: """Log a plot to wandb. @@ -1241,6 +1265,20 @@ def log_histogram(self, histogram: list[Any], step: int, name: str) -> None: for logger in self.loggers: logger.log_histogram(histogram, step, name) + def log_table( + self, columns: list[str], rows: list[list[Any]], step: int, name: str + ) -> None: + """Log a table to every backend that supports one. + + Args: + columns: Column headers + rows: One list of values per row + step: Global step value + name: Panel name + """ + for logger in self.loggers: + logger.log_table(columns, rows, step, name) + def log_plot(self, figure: plt.Figure, step: int, name: str) -> None: """Log a matplotlib figure to all backends. diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 99b252c8b03..79291aa113b 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -30,6 +30,7 @@ from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient from nemo_rl.data_plane.observability import ( MetricsDataPlaneClient, + breakdown_table, cluster_step_metrics, merge_snapshots, _estimate_encoded_bytes, @@ -968,3 +969,42 @@ def test_snapshot_percentiles_never_exceed_the_measured_max(): assert stats["p99_ms"] <= stats["max_ms"], stats assert stats["p50_ms"] <= stats["p99_ms"], stats client.close() + + +def test_breakdown_table_rows_by_op_worst_first(): + """One row per op, ordered by wall time, so the expensive op is the + first line read rather than the alphabetically luckiest.""" + metrics = { + "step/wall_ms": 100.0, + "step/get/calls": 8, + "step/get/wall_ms": 10.0, + "step/get/max_ms": 2.0, + "step/put/calls": 2, + "step/put/wall_ms": 90.0, + "step/put/max_ms": 50.0, + "step/comm_volume_mb": 1.0, # not per-op, must not become a row + "now/bytes_outstanding_mb": 0.0, # a level, likewise + } + columns, rows = breakdown_table(metrics) + + assert columns[0] == "op" + assert [r[0] for r in rows] == ["put", "get"], "worst first" + assert len(rows) == 2, "only per-op series become rows" + assert rows[0][columns.index("wall_ms")] == 90.0 + + +def test_breakdown_table_leaves_withheld_series_empty(): + """A percentile below the sample gate, or a fit that is not trustworthy, + is absent from the series — the table must carry None there rather than + a zero that would read as a measurement.""" + columns, rows = breakdown_table( + {"step/put/calls": 3, "step/put/wall_ms": 5.0, "step/put/max_ms": 2.0} + ) + row = rows[0] + assert row[columns.index("p99_ms")] is None + assert row[columns.index("overhead_ms")] is None + assert row[columns.index("calls")] == 3 + + +def test_breakdown_table_is_empty_when_nothing_ran(): + assert breakdown_table({"step/wall_ms": 0.0})[1] == [] From 85589f92f1654fe24f21f5acdc32e19a7b97df58 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 12:53:27 -0700 Subject: [PATCH 24/70] fix(data-plane): put the latency split in the cluster view too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `overhead_ms` and `transfer_ms` were emitted on the driver path only, so the cluster view -- the one a real run logs -- never carried them and their breakdown-table columns were always empty. The split is the whole reason that table exists, and it was missing from the only view most runs show. The cluster fit is the better one, besides. It is computed over every rank's summed sufficient statistics, so it has far more samples and far more size variation than one process's -- and size variation is exactly what decides whether the split is identifiable at all. A driver making one put a step of near-identical size reports `unidentifiable`; the same workload across eight ranks does not. Extracted `_latency_split` so both paths compute it once, matching `_step_deltas` and `_clamped_percentiles`. It returns None when the fit is untrustworthy, so the caller stops branching on `model_trustworthy` and the column stays empty rather than zero. Verified against known affine latency (6 ms/call, 400 MB/s) across 8 ranks: op calls wall_ms max_ms overhead_ms transfer_ms p50_ms p99_ms get 200 2233 16.01 1207 1026 12.9 16.01 put 48 1135 31.9 579.3 556.1 — — get: overhead + transfer = 2232.7 ms vs measured wall_ms 2232.7 put: overhead + transfer = 1135.5 ms vs measured wall_ms 1135.5 Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 14 ++--- nemo_rl/data_plane/observability.py | 60 ++++++++++++++------- tests/unit/data_plane/test_observability.py | 30 +++++++++++ 3 files changed, 80 insertions(+), 24 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index ec8a3887931..dacfd1819ce 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -468,12 +468,14 @@ cluster percentile). `data_plane/{cluster,driver}/breakdown` — one row per op, ordered by wall time so the expensive one reads first: -| op | calls | wall_ms | max_ms | p50_ms | p99_ms | -|---|---:|---:|---:|---:|---:| -| put | 64 | 3998 | 114.9 | 58.93 | 114.9 | -| get | 240 | 2916 | 20.0 | 12.76 | 20.0 | -| clear | 32 | 43.42 | 2.03 | — | — | -| register | 16 | 11.95 | 1.08 | — | — | +| op | calls | wall_ms | max_ms | overhead_ms | transfer_ms | p50_ms | p99_ms | +|---|---:|---:|---:|---:|---:|---:|---:| +| get | 200 | 2233 | 16.01 | 1207 | 1026 | 12.9 | 16.01 | +| put | 48 | 1135 | 31.9 | 579.3 | 556.1 | — | — | + +`overhead_ms` and `transfer_ms` stack to the measured `wall_ms` — that row +reads "2233 ms of get was 1207 ms of fixed per-call cost and 1026 ms of +bandwidth", which is a decision you can act on. A stack of line charts answers "how did put's wall time trend"; this answers "where did the step go", which is a table. Cells are empty rather diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 20e6f035344..99697aeb20f 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -444,6 +444,28 @@ def _step_deltas(snap: dict[str, Any], prev: dict[str, Any]) -> dict[str, float] } +def _latency_split( + fit: dict[str, Any], calls: int, op_bytes: int +) -> tuple[float, float] | None: + """This step's time split into fixed overhead and transfer, in ms. + + The two stack: together they are the model's estimate of the op's + ``wall_ms``, so charting them against the measured value shows the split + and how well the model holds. The *coefficients* come from the + cumulative fit and should be stable -- that is what a fitted model is + for -- while the *attribution* is per step, applied to this step's calls + and bytes. + + ``None`` when the fit is not trustworthy, which includes the common RL + case of near-uniform request sizes where the split is mathematically + unrecoverable. + """ + if not fit.get("model_trustworthy"): + return None + ms_per_byte = 1.0 / (fit["bandwidth_mb_s"] * 1e3) + return fit["fixed_ms"] * calls, ms_per_byte * op_bytes + + def _clamped_percentiles(hist: list[int], max_ms: float) -> tuple[float, float] | None: """p50 and p99 from ``hist``, or ``None`` when the sample is too thin. @@ -666,6 +688,18 @@ def cluster_step_metrics( pct = _clamped_percentiles(step_hist, stats["max_ms"]) if pct: metrics[f"step/{op}/p50_ms"], metrics[f"step/{op}/p99_ms"] = pct + # The split belongs here more than on the driver path: this fit is + # over every rank's sufficient statistics, so it has far more + # samples and far more size variation -- and size variation is + # exactly what decides whether the split is identifiable at all. + split = _latency_split( + stats["fit"], calls, stats["n_bytes"] - prev_op.get("n_bytes", 0) + ) + if split: + ( + metrics[f"step/{op}/overhead_ms"], + metrics[f"step/{op}/transfer_ms"], + ) = split return metrics @@ -960,24 +994,14 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: # that goes flat, quantised to bucket edges. The max is exact # and says the same thing at the handful of calls per step. metrics[f"step/{op}/max_ms"] = step_maxima.get(op, 0.0) - fit = st["fit"] - if fit.get("model_trustworthy"): - # The step's time split into the two things that cause it, - # in ms, rather than a ratio. These stack: together they are - # the model's estimate of this op's ``wall_ms`` for the - # step, so charting them against the measured ``wall_ms`` - # shows both the split and how well the model holds. - # - # The *coefficients* come from the cumulative fit and should - # be stable -- that is what a fitted model is for. The - # *attribution* is per step, because it is applied to this - # step's calls and bytes. A ratio would have been neither: - # cumulative and therefore flat, and unitless on an axis of - # milliseconds. - op_bytes = st["n_bytes"] - prev_op.get("n_bytes", 0) - ms_per_byte = 1.0 / (fit["bandwidth_mb_s"] * 1e3) - metrics[f"step/{op}/overhead_ms"] = fit["fixed_ms"] * calls - metrics[f"step/{op}/transfer_ms"] = ms_per_byte * op_bytes + split = _latency_split( + st["fit"], calls, st["n_bytes"] - prev_op.get("n_bytes", 0) + ) + if split: + ( + metrics[f"step/{op}/overhead_ms"], + metrics[f"step/{op}/transfer_ms"], + ) = split return metrics def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 79291aa113b..baeb8f70032 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -1008,3 +1008,33 @@ def test_breakdown_table_leaves_withheld_series_empty(): def test_breakdown_table_is_empty_when_nothing_ran(): assert breakdown_table({"step/wall_ms": 0.0})[1] == [] + + +def test_cluster_view_carries_the_latency_split(): + """The split was emitted on the driver path only, so the cluster view — + the one a real run logs — could never show it, and its table column was + always empty. The cluster fit is the better one besides: it is over + every rank's sufficient statistics, so it has far more samples and far + more size variation, and size variation is what decides whether the + split is identifiable at all.""" + fixed_ms, mb_per_s = 6.0, 400.0 + ranks = [] + for rank in range(4): + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + now = monotonic() + for i in range(8): # varied sizes, or the fit is unidentifiable + n_bytes = 100_000 * (i + 1 + rank) + wall_ms = fixed_ms + n_bytes / (mb_per_s * 1e3) + client._emit("get", "p", 1, n_bytes, now - wall_ms / 1e3, "ok") + ranks.append(client.snapshot()) + + metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 1.0) + assert "step/get/overhead_ms" in metrics, "cluster view must carry the split" + + total = metrics["step/get/overhead_ms"] + metrics["step/get/transfer_ms"] + assert total == pytest.approx(metrics["step/get/wall_ms"], rel=0.05) + + columns, rows = breakdown_table(metrics) + row = rows[0] + assert row[columns.index("overhead_ms")] is not None + assert row[columns.index("transfer_ms")] is not None From f73fe2e7dfa779878c1d7888da1d38c59310b179 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 12:59:32 -0700 Subject: [PATCH 25/70] fix(data-plane): per-op cluster time also reads as elapsed when it is not MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `step/get/wall_ms` reported 2233 for a step whose wall clock was 278 ms. The arithmetic is right — 200 gets of 11.1 ms summed across 8 concurrent ranks — but it is process-time wearing the name of a duration, and the magnitude reads as "get took 2.2 seconds". Worse, the same key means different things on the two paths: on the driver there is one process, so `wall_ms` is elapsed; on the cluster it is N times larger. A reader comparing the two views would find them disagreeing by the DP degree with nothing to explain it. This is the defect already fixed once at the top level, where `step/wall_ms` gained `step/wall_ms_per_process` — it was never applied to the per-op series, which is exactly what the breakdown table shows. The cluster path now emits `step/{op}/wall_ms_per_process` alongside the sum, and the table carries both. The sum attributes cost across ops; the per-process figure has the magnitude a millisecond implies: op calls wall_ms wall_ms_per_process max_ms overhead_ms transfer_ms get 200 2232 279 16.0 1207 1025 put 48 1135 141.9 31.9 579.2 556.1 Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 15 +++++++++++---- nemo_rl/data_plane/observability.py | 11 ++++++++++- tests/unit/data_plane/test_observability.py | 20 ++++++++++++++++++++ 3 files changed, 41 insertions(+), 5 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index dacfd1819ce..92f5c831920 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -468,10 +468,17 @@ cluster percentile). `data_plane/{cluster,driver}/breakdown` — one row per op, ordered by wall time so the expensive one reads first: -| op | calls | wall_ms | max_ms | overhead_ms | transfer_ms | p50_ms | p99_ms | -|---|---:|---:|---:|---:|---:|---:|---:| -| get | 200 | 2233 | 16.01 | 1207 | 1026 | 12.9 | 16.01 | -| put | 48 | 1135 | 31.9 | 579.3 | 556.1 | — | — | +| op | calls | wall_ms | wall_ms_per_process | max_ms | overhead_ms | transfer_ms | p50_ms | p99_ms | +|---|---:|---:|---:|---:|---:|---:|---:|---:| +| get | 200 | 2232 | 279 | 16.0 | 1207 | 1025 | 12.9 | 16.0 | +| put | 48 | 1135 | 141.9 | 31.9 | 579.2 | 556.1 | — | — | + +**Every time on the cluster path is summed across processes that ran +concurrently** — so it is process-time, not elapsed. 200 gets of 11 ms +across 8 ranks reads 2232 ms while the wall clock was 279. The same key on +the driver path, where there is one process, is elapsed. Both ship: +`wall_ms` attributes cost across ops, `wall_ms_per_process` carries the +magnitude a reader expects of a millisecond. `overhead_ms` and `transfer_ms` stack to the measured `wall_ms` — that row reads "2233 ms of get was 1207 ms of fixed per-call cost and 1026 ms of diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 99697aeb20f..9fe81c009df 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -668,7 +668,15 @@ def cluster_step_metrics( if calls <= 0: continue metrics[f"step/{op}/calls"] = calls - metrics[f"step/{op}/wall_ms"] = stats["wall_ms"] - prev_op.get("wall_ms", 0.0) + # Summed over processes that ran concurrently, like every other time + # on this path -- so it is process-time, and on the driver path the + # same key is elapsed. 200 gets of 11 ms across 8 ranks reads 2233 + # here and took 278 ms of wall clock. Both are worth having: the sum + # attributes cost across ops, the per-process figure gives the + # magnitude a reader expects from a millisecond. + op_wall_ms = stats["wall_ms"] - prev_op.get("wall_ms", 0.0) + metrics[f"step/{op}/wall_ms"] = op_wall_ms + metrics[f"step/{op}/wall_ms_per_process"] = op_wall_ms / n_procs metrics[f"step/{op}/max_ms"] = stats["max_ms"] # Percentiles over THIS step's calls, summed across ranks, not over # the lifetime: a cumulative percentile beside a step-scoped max is @@ -711,6 +719,7 @@ def cluster_step_metrics( _BREAKDOWN_COLUMNS = ( "calls", "wall_ms", + "wall_ms_per_process", "max_ms", "overhead_ms", "transfer_ms", diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index baeb8f70032..cdeca93074b 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -1038,3 +1038,23 @@ def test_cluster_view_carries_the_latency_split(): row = rows[0] assert row[columns.index("overhead_ms")] is not None assert row[columns.index("transfer_ms")] is not None + + +def test_cluster_per_op_time_is_reported_both_ways(): + """``step/{op}/wall_ms`` on the cluster path sums processes that ran + concurrently, so it is process-time — 200 gets of 11 ms across 8 ranks + reads 2233 ms while wall clock was 278. The same key on the driver path + is elapsed. Both figures ship: the sum attributes cost across ops, the + per-process one has the magnitude a reader expects of a millisecond.""" + ranks = [_rank_with([10.0] * 5) for _ in range(8)] # 40 calls, 400 ms summed + metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 1.0) + + assert metrics["step/put/wall_ms"] == pytest.approx(400.0, rel=0.1) + assert metrics["step/put/wall_ms_per_process"] == pytest.approx(50.0, rel=0.1) + assert metrics["step/put/wall_ms_per_process"] * metrics["now/n_processes"] == ( + pytest.approx(metrics["step/put/wall_ms"]) + ) + + columns, rows = breakdown_table(metrics) + assert "wall_ms_per_process" in columns + assert rows[0][columns.index("wall_ms_per_process")] is not None From acd6fa6a258827d5133136d99cc6616061593e6c Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 13:05:12 -0700 Subject: [PATCH 26/70] fix(data-plane): report per-op time per call, which is the invariant one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Neither figure on offer described the wire. `wall_ms` sums processes that ran concurrently, so it scales with DP degree; `wall_ms_per_process` traded that arbitrary denominator for another one. Per call is invariant to both DP degree and batch size — the same workload at 8 and at 32 ranks: 8 ranks mean_ms 11.16 wall_ms 2232 calls 200 32 ranks mean_ms 11.06 wall_ms 8850 calls 800 against a ground truth of 6.0 ms fixed + ~5.1 ms transfer. `mean_ms` is what compares across runs and cluster sizes; `wall_ms` stays for attributing cost across ops within one step. `mean_ms` had been there and I removed it, in the metric-trimming pass, on the grounds that it is exactly `wall_ms / calls` and a dashboard can divide. Derivable was true and beside the point: it was the only form of that number describing the operation rather than the shape of the run, and dropping it left two figures that both needed a paragraph to read. The `wall_ms_per_process` column added to patch that goes away with it — the top-level `step/wall_ms_per_process` stays, where there is no call count to divide by. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 22 ++++++------ nemo_rl/data_plane/observability.py | 19 ++++++----- tests/unit/data_plane/test_observability.py | 37 ++++++++++++--------- 3 files changed, 43 insertions(+), 35 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 92f5c831920..f5abd9972c3 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -468,17 +468,19 @@ cluster percentile). `data_plane/{cluster,driver}/breakdown` — one row per op, ordered by wall time so the expensive one reads first: -| op | calls | wall_ms | wall_ms_per_process | max_ms | overhead_ms | transfer_ms | p50_ms | p99_ms | +| op | calls | mean_ms | max_ms | wall_ms | overhead_ms | transfer_ms | p50_ms | p99_ms | |---|---:|---:|---:|---:|---:|---:|---:|---:| -| get | 200 | 2232 | 279 | 16.0 | 1207 | 1025 | 12.9 | 16.0 | -| put | 48 | 1135 | 141.9 | 31.9 | 579.2 | 556.1 | — | — | - -**Every time on the cluster path is summed across processes that ran -concurrently** — so it is process-time, not elapsed. 200 gets of 11 ms -across 8 ranks reads 2232 ms while the wall clock was 279. The same key on -the driver path, where there is one process, is elapsed. Both ship: -`wall_ms` attributes cost across ops, `wall_ms_per_process` carries the -magnitude a reader expects of a millisecond. +| get | 200 | 11.16 | 16.0 | 2232 | 1207 | 1025 | 12.9 | 16.0 | +| put | 48 | 23.66 | 31.9 | 1135 | 579.2 | 556.1 | — | — | + +**`mean_ms` is the one that describes the wire.** `wall_ms` on the cluster +path is summed over processes that ran concurrently, so it is process-time +and scales with DP degree — 200 gets of 11 ms across 8 ranks reads 2232 +while the wall clock was 279. Dividing by the process count only trades one +arbitrary denominator for another. Per call is invariant to both DP degree +and batch size: the same workload at 8 and at 32 ranks reports 11.16 and +11.06 ms while `wall_ms` quadruples. Use `mean_ms` to compare runs and +cluster sizes, `wall_ms` to attribute cost across ops within one step. `overhead_ms` and `transfer_ms` stack to the measured `wall_ms` — that row reads "2233 ms of get was 1207 ms of fixed per-call cost and 1026 ms of diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 9fe81c009df..0b245a7b4cf 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -668,15 +668,15 @@ def cluster_step_metrics( if calls <= 0: continue metrics[f"step/{op}/calls"] = calls - # Summed over processes that ran concurrently, like every other time - # on this path -- so it is process-time, and on the driver path the - # same key is elapsed. 200 gets of 11 ms across 8 ranks reads 2233 - # here and took 278 ms of wall clock. Both are worth having: the sum - # attributes cost across ops, the per-process figure gives the - # magnitude a reader expects from a millisecond. op_wall_ms = stats["wall_ms"] - prev_op.get("wall_ms", 0.0) + # How long one call took, which is the only form of this that + # describes the wire rather than the shape of the run. ``wall_ms`` + # is summed over concurrent processes, so it scales with DP degree; + # dividing by the process count instead just trades one arbitrary + # denominator for another. Per call is invariant to both, and + # comparable across runs and cluster sizes. + metrics[f"step/{op}/mean_ms"] = op_wall_ms / calls metrics[f"step/{op}/wall_ms"] = op_wall_ms - metrics[f"step/{op}/wall_ms_per_process"] = op_wall_ms / n_procs metrics[f"step/{op}/max_ms"] = stats["max_ms"] # Percentiles over THIS step's calls, summed across ranks, not over # the lifetime: a cumulative percentile beside a step-scoped max is @@ -718,9 +718,9 @@ def cluster_step_metrics( # would read as a measurement. _BREAKDOWN_COLUMNS = ( "calls", - "wall_ms", - "wall_ms_per_process", + "mean_ms", "max_ms", + "wall_ms", "overhead_ms", "transfer_ms", "p50_ms", @@ -997,6 +997,7 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: # can divide. ``snapshot()`` still carries the full picture, # percentiles included, for a one-off inspection. metrics[f"step/{op}/calls"] = calls + metrics[f"step/{op}/mean_ms"] = op_ms / calls metrics[f"step/{op}/wall_ms"] = op_ms # ``max_ms`` rather than p50/p99. Those come off a histogram # that is never reset, so per step they were a lifetime figure diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index cdeca93074b..fe9ab4e8376 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -1040,21 +1040,26 @@ def test_cluster_view_carries_the_latency_split(): assert row[columns.index("transfer_ms")] is not None -def test_cluster_per_op_time_is_reported_both_ways(): - """``step/{op}/wall_ms`` on the cluster path sums processes that ran - concurrently, so it is process-time — 200 gets of 11 ms across 8 ranks - reads 2233 ms while wall clock was 278. The same key on the driver path - is elapsed. Both figures ship: the sum attributes cost across ops, the - per-process one has the magnitude a reader expects of a millisecond.""" - ranks = [_rank_with([10.0] * 5) for _ in range(8)] # 40 calls, 400 ms summed - metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 1.0) - - assert metrics["step/put/wall_ms"] == pytest.approx(400.0, rel=0.1) - assert metrics["step/put/wall_ms_per_process"] == pytest.approx(50.0, rel=0.1) - assert metrics["step/put/wall_ms_per_process"] * metrics["now/n_processes"] == ( - pytest.approx(metrics["step/put/wall_ms"]) +def test_cluster_per_op_time_is_reported_per_call(): + """``wall_ms`` sums concurrent processes, so it scales with DP degree; + dividing by the process count trades one arbitrary denominator for + another. Per call is invariant to both DP degree and batch size, so it + describes the wire rather than the shape of the run, and is comparable + across runs and cluster sizes.""" + small = cluster_step_metrics( + merge_snapshots([_rank_with([10.0] * 5) for _ in range(8)]), {}, 1.0 + ) + large = cluster_step_metrics( + merge_snapshots([_rank_with([10.0] * 5) for _ in range(32)]), {}, 1.0 ) - columns, rows = breakdown_table(metrics) - assert "wall_ms_per_process" in columns - assert rows[0][columns.index("wall_ms_per_process")] is not None + assert small["step/put/mean_ms"] == pytest.approx(10.0, rel=0.15) + assert large["step/put/mean_ms"] == pytest.approx( + small["step/put/mean_ms"], rel=0.15 + ), "mean must not move with cluster size" + assert large["step/put/wall_ms"] == pytest.approx( + 4 * small["step/put/wall_ms"], rel=0.15 + ), "the sum does move with cluster size" + + columns, rows = breakdown_table(small) + assert "mean_ms" in columns and "wall_ms_per_process" not in columns From 96b80f9114a2bde94e2b027fd352c5d363293bd8 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 13:14:11 -0700 Subject: [PATCH 27/70] fix(data-plane): split the latency per call, like the mean it explains `overhead_ms` and `transfer_ms` were extensive -- fitted coefficients times this step's calls and bytes -- so they stacked to `wall_ms` and scaled with DP degree and batch size, describing the shape of the run rather than the wire. The same objection that moved per-op time to `mean_ms` applies to the terms that explain it, and leaving them summed made the table half intensive and half not. Per call they stack to `mean_ms` instead, and the overhead term stops being a product and becomes the fitted constant itself -- a property of the backend, comparable against a hardware number. Against a ground truth of 6.0 ms fixed and 400 MB/s: 8 ranks mean_ms 11.17 overhead 6.03 transfer 5.15 sum 11.17 32 ranks mean_ms 11.00 overhead 6.03 transfer 4.98 sum 11.00 The constant is recovered to 0.5%, the split adds to the mean exactly, and nothing moves with cluster size. `calls` and `wall_ms` are now the only extensive columns in the breakdown; everything else is per call and in ms. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 13 +++++++++--- nemo_rl/data_plane/observability.py | 22 ++++++++++++--------- tests/unit/data_plane/test_observability.py | 9 ++++++--- 3 files changed, 29 insertions(+), 15 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index f5abd9972c3..d6a098ec697 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -470,10 +470,17 @@ time so the expensive one reads first: | op | calls | mean_ms | max_ms | wall_ms | overhead_ms | transfer_ms | p50_ms | p99_ms | |---|---:|---:|---:|---:|---:|---:|---:|---:| -| get | 200 | 11.16 | 16.0 | 2232 | 1207 | 1025 | 12.9 | 16.0 | -| put | 48 | 23.66 | 31.9 | 1135 | 579.2 | 556.1 | — | — | +| get | 200 | 11.17 | 16.0 | 2232 | 6.03 | 5.15 | 12.9 | 16.0 | +| put | 48 | 23.66 | 31.9 | 1135 | 12.06 | 11.60 | — | — | -**`mean_ms` is the one that describes the wire.** `wall_ms` on the cluster +Everything in ms on that row is **per call**, and the two split terms add +to `mean_ms`: that `get` row reads "each call cost 11.17 ms, of which 6.03 +was fixed per-request overhead and 5.15 was bandwidth at this step's mean +request size". Per call the overhead term *is* the fitted constant, so it +is comparable against a hardware number. `calls` and `wall_ms` are the only +extensive columns. + +**Per-call figures describe the wire; sums describe the run.** `wall_ms` on the cluster path is summed over processes that ran concurrently, so it is process-time and scales with DP degree — 200 gets of 11 ms across 8 ranks reads 2232 while the wall clock was 279. Dividing by the process count only trades one diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 0b245a7b4cf..66e4527c7c2 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -447,23 +447,27 @@ def _step_deltas(snap: dict[str, Any], prev: dict[str, Any]) -> dict[str, float] def _latency_split( fit: dict[str, Any], calls: int, op_bytes: int ) -> tuple[float, float] | None: - """This step's time split into fixed overhead and transfer, in ms. + """One call's time split into fixed overhead and transfer, in ms. - The two stack: together they are the model's estimate of the op's - ``wall_ms``, so charting them against the measured value shows the split - and how well the model holds. The *coefficients* come from the - cumulative fit and should be stable -- that is what a fitted model is - for -- while the *attribution* is per step, applied to this step's calls - and bytes. + Per call, like ``mean_ms``, and for the same reason: the extensive form + scales with DP degree and batch size and so describes the shape of the + run rather than the wire. Per call the overhead term *is* the fitted + per-request constant -- a property of the backend, comparable against a + hardware number -- and the transfer term is that bandwidth at this + step's mean request size. + + The two stack to ``mean_ms``, so charting them against it shows the + split and how well the affine model holds, in the same units and the + same scale as everything else per-op. ``None`` when the fit is not trustworthy, which includes the common RL case of near-uniform request sizes where the split is mathematically unrecoverable. """ - if not fit.get("model_trustworthy"): + if not fit.get("model_trustworthy") or calls <= 0: return None ms_per_byte = 1.0 / (fit["bandwidth_mb_s"] * 1e3) - return fit["fixed_ms"] * calls, ms_per_byte * op_bytes + return fit["fixed_ms"], ms_per_byte * (op_bytes / calls) def _clamped_percentiles(hist: list[int], max_ms: float) -> tuple[float, float] | None: diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index fe9ab4e8376..42b7abb8182 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -442,9 +442,11 @@ def test_latency_breakdown_stacks_to_wall_ms(): assert fit["fixed_ms"] == pytest.approx(fixed_ms, rel=0.05) assert fit["bandwidth_mb_s"] == pytest.approx(mb_per_s, rel=0.05) - # the two components are the split, in ms, and they add up + # the two components are the split of ONE call, and they add to the mean total = metrics["step/put/overhead_ms"] + metrics["step/put/transfer_ms"] - assert total == pytest.approx(metrics["step/put/wall_ms"], rel=0.05) + assert total == pytest.approx(metrics["step/put/mean_ms"], rel=0.05) + # per call, the overhead term IS the fitted constant + assert metrics["step/put/overhead_ms"] == pytest.approx(fixed_ms, rel=0.05) assert "step/put/overhead_frac" not in metrics, "a ratio is derivable from these" client.close() @@ -1032,7 +1034,8 @@ def test_cluster_view_carries_the_latency_split(): assert "step/get/overhead_ms" in metrics, "cluster view must carry the split" total = metrics["step/get/overhead_ms"] + metrics["step/get/transfer_ms"] - assert total == pytest.approx(metrics["step/get/wall_ms"], rel=0.05) + assert total == pytest.approx(metrics["step/get/mean_ms"], rel=0.05) + assert metrics["step/get/overhead_ms"] == pytest.approx(fixed_ms, rel=0.05) columns, rows = breakdown_table(metrics) row = rows[0] From fcbffe418886d0fe6194cbd80c495d6f57f3af48 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 13:34:25 -0700 Subject: [PATCH 28/70] fix(data-plane): gate percentiles per quantile; per-row hashes when rows are uniform Percentiles were gated on one threshold of 50 samples for both p50 and p99, so an op with 48 calls (`put`, one per step) reported neither -- even though a median off 48 samples is perfectly real. n samples resolve a quantile no higher than about 1 - 1/n, so each quantile now waits for the samples it needs: p50 at 20, p99 at 100. Below the gate the key is absent rather than reporting the maximum wearing a percentile label. Dropped `tail_ratio_p99_mean`, which had no readers and raised KeyError once p99 became optional. Hash verification: the jagged/rectangular split was made on the layout, but a jagged field whose rows are uniform is a rectangle already -- its values buffer reshapes to one as a view -- so it now takes the per-row path for free. That matters because the batch-scoped fallback is an XOR reduction over one shared buffer and XOR cannot see a permutation: two equal-length rows swapped by a mis-shard round-tripped clean (0/60 caught, now 60/60), and a single changed element flagged all 16 rows instead of naming the one. Recording the scheme per partition rather than per field let a delta put (write_columns) restate the scheme of fields it never wrote, handing the read side the wrong one -- 2 false-positive round-trips in the guard suite. It is now recorded per field. A field written with uniform rows that comes back ragged is a divergence in the row lengths themselves and is reported as a mismatch; only a shard of a genuinely batch-scoped field still abstains into `fields_skipped`. Verified: 173 unit tests, 23/23 guard checks (10 corruption modes caught, zero false alarms over a 500-row soak, every shard grouping), 48/48 metric audit checks. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 48 +++-- nemo_rl/data_plane/observability.py | 137 ++++++++------ tests/unit/data_plane/test_observability.py | 196 ++++++++++++++++++-- 3 files changed, 290 insertions(+), 91 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index d6a098ec697..4cbfaaa353d 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -600,30 +600,39 @@ Two granularities, because torch has no ragged hash kernel: | leaf | digest | scope | |---|---|---| -| rectangular (`rewards`, `input_lengths`) | one per row, `hash_tensor(..., dim=1)` | per sample id; survives shard reads | -| jagged (`input_ids`, `generation_logprobs`, `token_mask`, `advantages`) | one over the values buffer, XORed per row with that row's length | per batch; a shard read reports unverified | - -Giving the jagged fields per-row digests would mean padding each one out to -a rectangle first. On a realistically ragged batch that rectangle is 3.5× -the real payload and cost 13× more, to answer a question the buffer digest -already answers. +| rectangular rows — dense, or jagged with uniform lengths | one per row, `hash_tensor(..., dim=1)` | per sample id; survives shard reads | +| genuinely ragged rows | one over the values buffer, XORed per row with that row's length | per batch; a shard read reports unverified | + +The split is on the *rows*, not the layout. A jagged leaf whose rows happen +to be uniform is already a rectangle — its values buffer reshapes to one as +a view — so it takes the per-row path for free. Only a leaf with rows of +differing lengths falls back, and giving *those* per-row digests would mean +padding each out to a rectangle first: on a realistically ragged batch that +rectangle is 3.5× the real payload and costs 13× more, to answer a question +the buffer digest already answers. + +Whichever scheme a put used is recorded per field and replayed on the read, +so the two sides always compute the same thing. A field written with +uniform rows that comes back ragged is a divergence in the row lengths +themselves, and is reported as a mismatch. **Detection is not attribution, and the difference is the whole point of the split.** The same corruption, injected into an 8-row batch: -| corruption | jagged leaf | rectangular leaf | +| corruption | ragged leaf | rectangular leaf | |---|---|---| | 1 element changed in `u3` | caught, flags **all 8 rows** | caught, names **`u3`** | | `u5` zeroed | caught, flags all 8 rows | caught, names `u5` | | `u3`↔`u4` swapped | caught only if their lengths differ | caught, names `u3`,`u4` | | nothing | clean | clean | -A jagged digest covers the whole values buffer, so any change moves every +A ragged digest covers the whole values buffer, so any change moves every row's value: it says *this batch is wrong*, never *this sample is wrong*. -Since `pack_jagged_fields` leaves ~94% of the payload jagged (only -`rewards` and `input_lengths` stay rectangular), that is the normal -resolution — you learn a step's transfer diverged and have to bisect for -the row yourself. +On rollout data straight out of generation that is the normal resolution +for the token-aligned fields — you learn a step's transfer diverged and +have to bisect for the row yourself. Anything uniform-width (a densified +read, `advantages` written at full width, a shard whose rows agree) names +the sample. Verified by injecting corruption into the round trip. Caught: a single-element change in every dtype, a truncated row, a zeroed row, a @@ -632,14 +641,15 @@ bf16→fp32 precision change, and a row served from the wrong sample — with from 1 to 256, reversed id order, field subsets and delta writes. Known limits, measured rather than assumed: -- A **mis-shard** (two rows swapped) is caught on a jagged field only when - the two rows differ in length — 58/60 on ragged rollout data, never on a - uniform-length batch. Rectangular fields catch it unconditionally. +- A batch-scoped digest is an XOR reduction over one shared buffer, and XOR + cannot see a permutation of what it reduces. On a ragged field a + **mis-shard** (two rows swapped) is therefore caught only when the two + rows differ in length — 60/60 on ragged rollout data, where lengths rarely + collide. Rows of uniform width catch it unconditionally, per row. The same + blind spot hides a reordering *within* a row: 0/200 for a two-token swap + in `input_ids`, and 18/200 for moving two set bits in a bool mask. - It reads every tensor byte again on both sides — ~2.4 ms for a 12 MB jagged batch, on put and again on get. Keep it to debugging runs. -- A rectangular field that comes back **jagged** (one row truncated makes - the batch ragged) has no comparable digest and is dropped — counted in - `hash/fields_skipped`, not reported as a mismatch. - Only rows this process wrote can be checked. A consumer-side client reports them under `hash/rows_unverified` rather than counting them clean, and `hash/fields_skipped` reports any leaf it could not compare diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 66e4527c7c2..73853603cd5 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -96,12 +96,12 @@ class DataPlaneEvent(TypedDict): # few lines carry the identity of what broke. _MAX_HASH_MISMATCH_LOGS = 20 -# Calls needed in a window before a percentile off the histogram means -# anything. Below this the interpolation returns bucket geometry: one sample -# in (100, 250] yields p50 = 100 + 150*0.50 = 175 and p99 = 248.5 whatever -# the call actually took. Reporting that next to an exact ``max_ms`` produced -# a max below its own median. -_MIN_SAMPLES_FOR_PERCENTILE = 50 +# Quantiles reported per op, each with the sample count it needs. n samples +# resolve a quantile no higher than about 1 - 1/n, so a p99 wants 100 and a +# p50 wants only a couple -- 20 for stability. One threshold for both was +# why 48 calls reported neither, when the median was perfectly real and only +# the p99 would have been the maximum wearing a percentile label. +_QUANTILES = ((0.50, "p50_ms", 20), (0.99, "p99_ms", 100)) class _FieldDigest(NamedTuple): @@ -470,24 +470,28 @@ def _latency_split( return fit["fixed_ms"], ms_per_byte * (op_bytes / calls) -def _clamped_percentiles(hist: list[int], max_ms: float) -> tuple[float, float] | None: - """p50 and p99 from ``hist``, or ``None`` when the sample is too thin. +def _clamped_percentiles(hist: list[int], max_ms: float) -> dict[str, float]: + """Whichever of :data:`_QUANTILES` this sample can actually support. - Two corrections, both needed wherever a percentile is taken. Below - :data:`_MIN_SAMPLES_FOR_PERCENTILE` the interpolation returns bucket - geometry rather than data -- one sample in (100, 250] yields a p50 of - 175 whatever the call took -- so there is no answer to give. And the - interpolation spreads a bucket's samples uniformly across it, so calls - clustered low in a wide bucket read high, above the exact maximum - measured beside them; the maximum is the tighter bound. + Two corrections, both needed wherever a percentile is taken off a coarse + histogram. Each quantile is withheld until there are enough samples to + resolve it: below that the interpolation returns bucket geometry rather + than data -- one sample in (100, 250] yields a p50 of 175 whatever the + call took. And the interpolation spreads a bucket's samples uniformly + across it, so calls clustered low in a wide bucket read high, above the + exact maximum measured beside them; the maximum is the tighter bound. + + Returns a dict rather than a fixed pair so a caller emits only what the + data supports. An absent series says "not enough calls"; a zero would + read as a measurement. """ - if sum(hist) < _MIN_SAMPLES_FOR_PERCENTILE: - return None + n = sum(hist) ceiling = max_ms if max_ms > 0 else float("inf") - return ( - min(percentile_from_hist(hist, 0.50), ceiling), - min(percentile_from_hist(hist, 0.99), ceiling), - ) + return { + name: min(percentile_from_hist(hist, q), ceiling) + for q, name, min_samples in _QUANTILES + if n >= min_samples + } def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: @@ -511,13 +515,9 @@ def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: ) stats["fit"] = fit_latency_bandwidth(stats) hist = stats["latency_hist"] - pct = _clamped_percentiles(hist, stats["max_ms"]) - stats["p50_ms"], stats["p99_ms"] = pct if pct else (0.0, 0.0) - # Tail/mean ratio: a mean hides MR churn and queueing, which show - # up as p99 pulling away from the mean. - stats["tail_ratio_p99_mean"] = ( - stats["p99_ms"] / stats["mean_ms"] if stats["mean_ms"] > 0 else 0.0 - ) + # Only what the sample supports; an absent key says "not enough + # calls", which a zero would not. + stats.update(_clamped_percentiles(hist, stats["max_ms"])) # Snapshot fields that combine by summing, by taking a maximum, and the @@ -697,9 +697,8 @@ def cluster_step_metrics( ] # Off the step's histogram delta, not the cumulative one # :func:`_derive_op_metrics` used. - pct = _clamped_percentiles(step_hist, stats["max_ms"]) - if pct: - metrics[f"step/{op}/p50_ms"], metrics[f"step/{op}/p99_ms"] = pct + for name, value in _clamped_percentiles(step_hist, stats["max_ms"]).items(): + metrics[f"step/{op}/{name}"] = value # The split belongs here more than on the driver path: this fit is # over every rank's sufficient statistics, so it has far more # samples and far more size variation -- and size variation is @@ -903,7 +902,9 @@ def __init__( # partition -> (batch-scoped field names, row count at put). Those # digests cover a whole buffer, so they only reconcile against a read # of that same batch; the row count is what detects a shard read. - self._batch_scope: dict[str, tuple[frozenset[str], int]] = {} + # partition -> field -> rows the field's digest was reduced over, + # for batch-scoped fields only. An absent field was reduced per row. + self._batch_scope: dict[str, dict[str, int]] = {} self._hash_mismatches_logged = 0 # Set by ``_emit`` to the inner client's wall time for the op just # run, so the wrapping methods can subtract it and bill the rest to @@ -1110,12 +1111,15 @@ def _row_fingerprints( """``torch.hash_tensor`` fingerprints for each tensor leaf. A rectangular leaf reduces per row (``dim=1``), which names the - sample that diverged. ``hash_tensor`` has no ragged kernel, so a - jagged leaf instead gets one digest over its whole values buffer, - XORed per row with that row's length, and is marked - ``batch_scoped``; padding it out to a rectangle to get per-row - digests costs far more than the answer is worth. ``README.md`` has - the resulting detection/attribution table. + sample that diverged. A jagged leaf whose rows happen to be uniform + is rectangular already — its values buffer reshapes to the rectangle + as a view — so it takes the same path for free. Only a genuinely + ragged leaf falls back to one digest over its whole values buffer, + XORed per row with that row's length, and marked ``batch_scoped``; + padding it out to a rectangle costs far more than the answer is + worth. That fallback inherits the blind spot of an XOR reduction: it + sees any change to the multiset of values, but not a permutation of + them. ``README.md`` has the detection/attribution table. Args: td: Leaves to fingerprint; ``None`` yields an empty result. @@ -1156,20 +1160,24 @@ def _row_fingerprints( if offsets.numel() - 1 != n_rows: stats.fields_skipped += 1 continue - buffer, lengths = v.values(), (offsets[1:] - offsets[:-1]).tolist() + lengths = (offsets[1:] - offsets[:-1]).tolist() + # Uniform rows: the values buffer already *is* the rectangle, + # so reshaping it is a view and the per-row reduction is free. + rectangle = v.values() if len(set(lengths)) == 1 else None elif v.shape[0] != n_rows: stats.fields_skipped += 1 continue - elif name in batch_scoped_fields: - buffer = v - lengths = [v.shape[1] if v.ndim >= 2 else 1] * n_rows else: - flat = _as_int_view(v.reshape(n_rows, -1)) + lengths = [v.shape[1] if v.ndim >= 2 else 1] * n_rows + rectangle = v + if rectangle is not None and name not in batch_scoped_fields: + flat = _as_int_view(rectangle.reshape(n_rows, -1)) out[name] = _FieldDigest( (torch.hash_tensor(flat, dim=1) ^ salt).tolist(), batch_scoped=False, ) continue + buffer = v.values() if v.is_nested else v flat_buffer = _as_int_view(buffer.reshape(1, -1)) buffer_digest = torch.hash_tensor(flat_buffer, dim=1).tolist()[0] ^ salt out[name] = _FieldDigest( @@ -1189,19 +1197,24 @@ def _record_hashes( per_field = partition_hashes.setdefault(sample_id, {}) for name, digest in digests.items(): per_field[name] = digest.per_row[row] - scoped = frozenset(n for n, d in digests.items() if d.batch_scoped) - if scoped: - self._batch_scope[partition_id] = (scoped, len(sample_ids)) + # Per field, not per partition: a delta put (write_columns) names only + # the fields it writes, and must not restate the scheme of the ones it + # left alone. Recording the batch it was reduced over lets the read + # side tell a shard of a batch-scoped field from a relayout. + scheme = self._batch_scope.setdefault(partition_id, {}) + for name, digest in digests.items(): + if digest.batch_scoped: + scheme[name] = len(sample_ids) + else: + scheme.pop(name, None) self._stats.hash_verify.rows_recorded += len(sample_ids) def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: """Compare wire-out fingerprints against what was written.""" if not isinstance(out, TensorDict): return - scoped_names, scoped_rows = self._batch_scope.get( - partition_id, (frozenset(), 0) - ) - digests = self._row_fingerprints(out, sample_ids, scoped_names) + scheme = self._batch_scope.get(partition_id, {}) + digests = self._row_fingerprints(out, sample_ids, scheme) if not digests: return partition_hashes = self._hash_by_partition.get(partition_id, {}) @@ -1212,15 +1225,28 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N # as a mismatch — but *count* the drop. Dropping silently is the # exact shape of the bug that made this check pass while covering # nothing: a field stops being compared and the report still reads - # clean. It also fires when a rectangular put comes back jagged - # (one row truncated makes the batch ragged), which is a real - # divergence this cannot express as a row-level mismatch. + # clean. comparable = {} + relaid_out = [] for name, digest in digests.items(): - if not digest.batch_scoped or scoped_rows == len(sample_ids): + if not digest.batch_scoped or scheme.get(name) == len(sample_ids): comparable[name] = digest + elif name in scheme: + stats.fields_skipped += 1 # a shard of a batch-scoped put else: - stats.fields_skipped += 1 + # Put reduced this field per row, so its rows were uniform; + # they came back ragged. The row lengths themselves changed — + # a real divergence, not something to skip. + relaid_out.append(name) + for name in relaid_out: + if self._hash_mismatches_logged < _MAX_HASH_MISMATCH_LOGS: + self._hash_mismatches_logged += 1 + logger.error( + "data-plane hash mismatch: partition=%s field=%s written " + "with uniform row lengths, read back ragged", + partition_id, + name, + ) for row, sample_id in enumerate(sample_ids): per_field = partition_hashes.get(sample_id) if not per_field: @@ -1229,6 +1255,7 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N stats.rows_unverified += 1 continue stats.rows_checked += 1 + stats.mismatches += len(relaid_out) for name, digest in comparable.items(): expected = per_field.get(name) if expected is None or expected == digest.per_row[row]: diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 42b7abb8182..c0fb0ecf100 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -415,8 +415,8 @@ def test_step_metrics_tail_is_exact_not_bucketed(): assert "put/p99_ms" not in metrics and "step/put/p50_ms" not in metrics assert metrics["step/put/max_ms"] >= 30.0 assert metrics["step/put/max_ms"] != pytest.approx(24.85, abs=0.5), "bucket edge" - # the cumulative view still carries percentiles for one-off inspection - assert "p99_ms" in client.snapshot()["by_op"]["put"] + # and one call supports no percentile at all, in either view + assert "p99_ms" not in client.snapshot()["by_op"]["put"] client.close() @@ -626,13 +626,22 @@ def test_hash_fingerprint_matches_across_jagged_and_dense(): client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) ids = ["a", "b"] dense = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int64) - jagged = _jagged(list(dense.unbind())) - put_side = client._row_fingerprints(jagged, ids)["x"] - get_side = client._row_fingerprints( + # uniform rows: both sides reduce per row, and a densified read agrees + put_side = client._row_fingerprints(_jagged(list(dense.unbind())), ids)["x"] + assert not put_side.batch_scoped + get_side = client._row_fingerprints(TensorDict({"x": dense}, batch_size=[2]), ids) + assert put_side == get_side["x"] + + # ragged rows: batch-scoped, and the read replays that scheme rather than + # choosing one from the dense tensor in hand + ragged = _jagged([torch.tensor([1, 2, 3]), torch.tensor([4, 5])]) + scoped = client._row_fingerprints(ragged, ids)["x"] + assert scoped.batch_scoped + replayed = client._row_fingerprints( TensorDict({"x": dense}, batch_size=[2]), ids, batch_scoped_fields={"x"} )["x"] - assert put_side == get_side + assert replayed.batch_scoped def test_hash_shard_read_of_jagged_field_is_unverified_not_a_mismatch(): @@ -649,12 +658,12 @@ def test_hash_shard_read_of_jagged_field_is_unverified_not_a_mismatch(): client.close() -def test_hash_incomparable_field_is_counted_not_dropped(): +def test_hash_rectangular_put_read_back_ragged_is_a_mismatch(): """A rectangular put can come back jagged — truncating one row makes the - batch ragged. Its per-row digests are not comparable against a - batch-scoped read, so the field is dropped; dropping it *silently* is - the exact shape of the bug that let this check pass while covering - nothing, so the drop has to land in ``fields_skipped``.""" + batch ragged. There is no row-level digest to compare against, but the + row lengths changing between wire-in and wire-out *is* the divergence. + Reporting it as a skipped field would leave ``mismatches`` reading zero, + the exact shape of a guard that passes while covering nothing.""" client = _hash_client(_RaggedOnReadClient()) ids = [f"u{i}" for i in range(4)] dense = TensorDict( @@ -664,8 +673,24 @@ def test_hash_incomparable_field_is_counted_not_dropped(): client.get_samples(sample_ids=ids, partition_id="p", select_fields=["x"]) hv = client.snapshot()["hash_verify"] - assert hv["mismatches"] == 0, "must not cry wolf on an incomparable field" - assert hv["fields_skipped"] == 1, "the drop must be visible" + assert hv["mismatches"] > 0, "a truncated row must not read as clean" + assert hv["fields_skipped"] == 0, "this is a divergence, not an abstention" + + +def test_hash_shard_of_a_ragged_field_is_skipped_not_a_mismatch(): + """The abstention that *is* legitimate: a batch-scoped digest covers the + whole buffer it was reduced over, so a shard of it is genuinely + incomparable. That must land in ``fields_skipped`` — visible, but not + crying wolf.""" + client = _hash_client() + ids = [f"u{i}" for i in range(4)] + rows = [torch.arange(3 + i, dtype=torch.int64) for i in range(4)] + client.put_samples(sample_ids=ids, partition_id="p", fields=_jagged(rows)) + client.get_samples(sample_ids=ids[:2], partition_id="p", select_fields=["x"]) + + hv = client.snapshot()["hash_verify"] + assert hv["mismatches"] == 0 + assert hv["fields_skipped"] == 1 assert client.get_step_metrics(1.0)["step/hash/fields_skipped"] == 1 @@ -755,9 +780,13 @@ def test_merge_sums_counters_and_rederives_percentiles(): assert merged["by_op"]["put"]["calls"] == 12 # 3 ranks x 4 puts assert merged["by_op"]["put"]["n_bytes"] == 12_000 assert sum(merged["by_op"]["put"]["latency_hist"]) == 12 - # derived from the summed histogram, not averaged from the ranks - single = ranks[0].snapshot()["by_op"]["put"] - assert merged["by_op"]["put"]["p50_ms"] == pytest.approx(single["p50_ms"], rel=0.3) + # the buckets add: the merged histogram is the ranks' summed elementwise + per_rank = [c.snapshot()["by_op"]["put"]["latency_hist"] for c in ranks] + assert merged["by_op"]["put"]["latency_hist"] == [ + sum(counts) for counts in zip(*per_rank) + ] + # 12 calls supports no percentile, and none is offered + assert "p50_ms" not in merged["by_op"]["put"] for c in ranks: c.close() @@ -962,7 +991,7 @@ def test_snapshot_percentiles_never_exceed_the_measured_max(): surface documented as the cumulative view.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient()) now = monotonic() - for _ in range(5): + for _ in range(120): # enough that both quantiles clear their gate client._emit("register", "p", 1, 0, now - 0.011 / 1e3, "ok") for view in (client.snapshot(), merge_snapshots([client.snapshot()])): @@ -1066,3 +1095,136 @@ def test_cluster_per_op_time_is_reported_per_call(): columns, rows = breakdown_table(small) assert "mean_ms" in columns and "wall_ms_per_process" not in columns + + +def test_each_quantile_waits_for_the_samples_it_needs(): + """n samples resolve a quantile no higher than about 1 - 1/n, so a p99 + wants 100 and a p50 only a couple. One threshold for both meant 48 calls + reported neither, when the median was perfectly real and only the p99 + would have been the maximum wearing a percentile label.""" + + def metrics_for(n_calls): + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + now = monotonic() + for i in range(n_calls): + client._emit("put", "p", 1, 1_000, now - (5.0 + i % 7) / 1e3, "ok") + return cluster_step_metrics(merge_snapshots([client.snapshot()]), {}, 1.0) + + assert "step/put/p50_ms" not in metrics_for(10), "too thin for either" + mid = metrics_for(48) + assert "step/put/p50_ms" in mid, "a median off 48 calls is real" + assert "step/put/p99_ms" not in mid, "a p99 off 48 calls is the max" + both = metrics_for(120) + assert "step/put/p50_ms" in both and "step/put/p99_ms" in both + + +class _JaggedEcho(NoOpDataPlaneClient): + """Returns whatever was put, jagged, so row lengths survive the trip.""" + + def __init__(self) -> None: + super().__init__() + self.rows: dict[tuple[str, str], dict[str, torch.Tensor]] = {} + + def put_samples(self, sample_ids, partition_id, fields=None, tags=None): + for key in fields.keys(): + v = fields.get(key) + rows = v.unbind() if v.is_nested else list(v) + for sid, row in zip(sample_ids, rows): + self.rows.setdefault((partition_id, sid), {})[str(key)] = row.clone() + return super().put_samples( + sample_ids=sample_ids, partition_id=partition_id, fields=fields, tags=tags + ) + + def get_samples(self, sample_ids, partition_id, select_fields): + out = {} + for f in select_fields: + rows = [self.rows[(partition_id, sid)][f] for sid in sample_ids] + out[f] = ( + torch.stack(rows) + if all(r.shape == rows[0].shape for r in rows[1:]) + else torch.nested.nested_tensor(rows, layout=torch.jagged) + ) + return TensorDict(out, batch_size=[len(sample_ids)]) + + +def _jagged_ids(lengths, seed=0): + """Rows of pseudorandom token ids. + + Deliberately not ``arange``: ``hash_tensor`` is an XOR reduction, and + aligned runs of consecutive integers collide under it — ``XOR(6..11)`` + and ``XOR(12..17)`` are both 1 — which would make two visibly different + rows fingerprint the same. + """ + g = torch.Generator().manual_seed(seed) + return TensorDict( + { + "ids": torch.nested.nested_tensor( + [torch.randint(0, 32000, (n,), generator=g) for n in lengths], + layout=torch.jagged, + ) + }, + batch_size=[len(lengths)], + ) + + +def test_uniform_jagged_rows_are_fingerprinted_per_row(): + """A jagged field whose rows happen to be uniform is a rectangle already + -- its values buffer reshapes to one as a view -- so it earns per-row + attribution for free. The batch-scoped fallback is an XOR over one shared + buffer, which cannot see a permutation: two equal-length rows swapped by a + mis-shard would round-trip clean.""" + inner = _JaggedEcho() + client = _hash_client(inner) + ids = [f"u{i}" for i in range(4)] + client.put_samples( + sample_ids=ids, partition_id="p", fields=_jagged_ids([6, 6, 6, 6]) + ) + a, b = inner.rows[("p", "u1")]["ids"], inner.rows[("p", "u2")]["ids"] + inner.rows[("p", "u1")]["ids"], inner.rows[("p", "u2")]["ids"] = b, a + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + + hv = client.snapshot()["hash_verify"] + assert hv["mismatches"] == 2, "the two swapped rows, named individually" + assert hv["fields_skipped"] == 0 + client.close() + + +def test_uniform_put_read_back_ragged_is_a_mismatch_not_a_skip(): + """Row lengths changing between wire-in and wire-out is a divergence. It + has no row-level digest to compare against, but reporting it as a skipped + field would leave mismatches reading zero -- exactly the shape of a guard + that covers nothing.""" + inner = _JaggedEcho() + client = _hash_client(inner) + ids = [f"u{i}" for i in range(4)] + client.put_samples( + sample_ids=ids, partition_id="p", fields=_jagged_ids([6, 6, 6, 6]) + ) + inner.rows[("p", "u2")]["ids"] = inner.rows[("p", "u2")]["ids"][:-2] + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + + assert client.snapshot()["hash_verify"]["mismatches"] > 0 + client.close() + + +def test_delta_put_does_not_restate_the_scheme_of_untouched_fields(): + """write_columns puts one field into a partition written ragged earlier. + Holding the jagged/per-row choice per partition let that delta hand the + read side the wrong scheme for its own field, and every row came back a + false alarm.""" + inner = _JaggedEcho() + client = _hash_client(inner) + ids = [f"u{i}" for i in range(4)] + client.put_samples( + sample_ids=ids, partition_id="p", fields=_jagged_ids([2, 4, 6, 3]) + ) + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + assert client.snapshot()["hash_verify"]["mismatches"] == 0 + + # same field, rewritten uniform -- the later scheme must win + client.put_samples( + sample_ids=ids, partition_id="p", fields=_jagged_ids([6, 6, 6, 6], seed=99) + ) + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + assert client.snapshot()["hash_verify"]["mismatches"] == 0 + client.close() From 38950f91e68f0ddbf45edf8fa808c2f0a1d5256e Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 14:56:57 -0700 Subject: [PATCH 29/70] fix(data-plane): report p90, not p99, as the per-op tail A step holds tens of calls, not thousands. A p99 needs ~100 samples before any observation lies above its rank at all; below that it collapses onto the largest one. Over a lognormal-with-tail draw at 58 calls -- what a DP-8 run actually puts per step -- the p99 equalled the maximum 80% of the time, which is `max_ms` under a more precise-sounding name. The p90 off the same 58 never did on a smooth tail and 12% of the time on a bimodal one. So the reported tail quantile is now p90, gated at 40 samples: each quantile wants roughly four observations above its rank, n >= 4 / (1 - q). `put` and `clear` at 58 calls/step reported no tail at all before and now carry one; `register` at 8 calls still reports neither, which is the honest answer at that size. `max_ms` keeps its place beside it. The two answer different questions -- max is "did anything go wrong this step", p90 is "what does the tail look like" -- and a sharp divergence between them is the signal that an op is bimodal (a straggler rank, a cold buffer), where the max is the number to chase. Widening the window instead was measured and rejected: pooling 8 steps left the p99 at 45% median error on a bimodal draw, because the p99 then sits in the low-density gap between the body and the slow mode. Verified: 179 unit tests, 48/48 metric audit checks, 23/23 hash-guard checks, plus a 10-process e2e where put/get/clear all carry p50 and p90. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 34 +++++++++---- nemo_rl/data_plane/observability.py | 29 +++++++---- tests/unit/data_plane/test_observability.py | 56 ++++++++++++++------- 3 files changed, 78 insertions(+), 41 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 4cbfaaa353d..3952bbc77c6 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -468,10 +468,10 @@ cluster percentile). `data_plane/{cluster,driver}/breakdown` — one row per op, ordered by wall time so the expensive one reads first: -| op | calls | mean_ms | max_ms | wall_ms | overhead_ms | transfer_ms | p50_ms | p99_ms | +| op | calls | mean_ms | max_ms | wall_ms | overhead_ms | transfer_ms | p50_ms | p90_ms | |---|---:|---:|---:|---:|---:|---:|---:|---:| -| get | 200 | 11.17 | 16.0 | 2232 | 6.03 | 5.15 | 12.9 | 16.0 | -| put | 48 | 23.66 | 31.9 | 1135 | 12.06 | 11.60 | — | — | +| get | 200 | 11.17 | 16.0 | 2232 | 6.03 | 5.15 | 12.9 | 15.2 | +| put | 48 | 23.66 | 31.9 | 1135 | 12.06 | 11.60 | 22.4 | 30.1 | Everything in ms on that row is **per call**, and the two split terms add to `mean_ms`: that `get` row reads "each call cost 11.17 ms, of which 6.03 @@ -553,7 +553,7 @@ can exceed 1, meaning measuring cost more than the operation measured — a signal worth seeing rather than hiding. **Units:** every duration is `_ms`, every volume is `_mb`, no exceptions — -a chart mixing `wall_s` against `p99_ms` puts a 0.008 beside a 24.85 and +a chart mixing `wall_s` against `p90_ms` puts a 0.008 beside a 24.85 and reads as a data-plane bug rather than an axis one. **Per step you get, per op tag:** `calls`, `wall_ms`, `max_ms`, and — when @@ -565,13 +565,25 @@ shows the breakdown and the model error in one picture. The coefficients come from the cumulative fit (a model should be stable); the attribution is per step, applied to that step's calls and bytes. A ratio was tried first and was the wrong shape — cumulative and therefore flat, and unitless on an -axis of milliseconds. Not percentiles — the -histogram is cumulative by design, so a per-step p50 off it goes flat, and -at the handful of calls an op makes in one step a p99 is bucket geometry -rather than data (one sample in the `(10, 25]` bucket always yields -`10 + 15*0.99 = 24.85`). `max_ms` is exact, scoped to the step, and says -the same thing at that sample size. The percentiles remain in `snapshot()`, -where the cumulative sample count justifies them. +axis of milliseconds. Percentiles come off the *step's* histogram delta, not the +cumulative one -- a per-step p50 off a histogram that is never reset goes +flat -- and each is emitted only when the step holds enough calls to +resolve it: **p50 at 20, p90 at 40**, roughly four observations above the +rank (`n >= 4 / (1 - q)`). Below that the key is absent. + +The tail quantile is **p90, not p99**, because a step holds tens of calls, +not thousands. A p99 needs ~100 samples before any observation lies above +its rank at all; below that it collapses onto the largest one. Over a +lognormal-with-tail draw at 58 calls -- what a DP-8 run actually puts per +step -- the p99 equalled the maximum **80% of the time**, which is +`max_ms` under a more precise-sounding name. The p90 off the same 58 never +did on a smooth tail and 12% of the time on a bimodal one. A coarser +quantile that is resolved beats a finer one that is not. + +`max_ms` stays alongside, exact and scoped to the step: it answers "did +anything go wrong this step", where p90 answers "what does the tail look +like". If the two diverge sharply, the op is bimodal -- a straggler rank +or a cold buffer -- and the max is the number to chase. Measured against a no-op inner client on the payload the wire actually carries — 256 ragged rows, 12 MB, jagged per-token fields as diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 73853603cd5..e78f63d075e 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -96,12 +96,18 @@ class DataPlaneEvent(TypedDict): # few lines carry the identity of what broke. _MAX_HASH_MISMATCH_LOGS = 20 -# Quantiles reported per op, each with the sample count it needs. n samples -# resolve a quantile no higher than about 1 - 1/n, so a p99 wants 100 and a -# p50 wants only a couple -- 20 for stability. One threshold for both was -# why 48 calls reported neither, when the median was perfectly real and only -# the p99 would have been the maximum wearing a percentile label. -_QUANTILES = ((0.50, "p50_ms", 20), (0.99, "p99_ms", 100)) +# Quantiles reported per op, each with the sample count it needs: enough for +# roughly four observations above the rank, or n >= 4 / (1 - q). +# +# The tail one is p90, not p99, because a step holds tens of calls, not +# thousands. A p99 needs ~100 samples before any observation lies above its +# rank at all, and below that it collapses onto the largest one -- measured +# over a lognormal-with-tail draw, a p99 off 58 calls equalled the maximum +# 80% of the time, which is ``max_ms`` under a more precise-sounding name. +# p90 off the same 58 never did on a smooth tail and 12% of the time on a +# bimodal one. A coarser quantile that is actually resolved beats a finer +# one that is not. +_QUANTILES = ((0.50, "p50_ms", 20), (0.90, "p90_ms", 40)) class _FieldDigest(NamedTuple): @@ -716,7 +722,7 @@ def cluster_step_metrics( # Per-op columns worth a row in the breakdown, in the order they read. # ``overhead_ms``/``transfer_ms`` are only present when the affine fit is -# trustworthy, and ``p50_ms``/``p99_ms`` only above the sample gate, so a +# trustworthy, and ``p50_ms``/``p90_ms`` only above the sample gate, so a # row carries None where a series was withheld rather than a zero that # would read as a measurement. _BREAKDOWN_COLUMNS = ( @@ -727,7 +733,7 @@ def cluster_step_metrics( "overhead_ms", "transfer_ms", "p50_ms", - "p99_ms", + "p90_ms", ) @@ -805,7 +811,8 @@ class OpStats: sum_ms_sq: float = 0.0 # Slowest single call, exact. The histogram below can only place a # call in a bucket, so at the handful of calls an op makes in one step - # a percentile off it is bucket geometry rather than data -- p99 of one + # a percentile off it is bucket geometry rather than data -- a tail + # quantile of one # sample in (10, 25] is always 10 + 15*0.99 = 24.85. This is the # per-step tail signal; the histogram is for the cumulative view. max_ms: float = 0.0 @@ -961,7 +968,7 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: wall_ms = snap["total_wall_ms"] - prev.get("total_wall_ms", 0.0) # Every duration is ms and every volume is MB, with no exceptions: - # a chart that mixes wall_s against p99_ms puts a 0.008 next to a + # a chart that mixes wall_s against p90_ms puts a 0.008 next to a # 24.85 and reads as a bug in the data plane rather than in the # axis. GB was the same problem one dimension over -- a realistic # step moved 0.00017 GB. @@ -1004,7 +1011,7 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: metrics[f"step/{op}/calls"] = calls metrics[f"step/{op}/mean_ms"] = op_ms / calls metrics[f"step/{op}/wall_ms"] = op_ms - # ``max_ms`` rather than p50/p99. Those come off a histogram + # ``max_ms`` rather than p50/p90. Those come off a histogram # that is never reset, so per step they were a lifetime figure # that goes flat, quantised to bucket edges. The max is exact # and says the same thing at the handful of calls per step. diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index c0fb0ecf100..a84a8d103d5 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -34,6 +34,7 @@ cluster_step_metrics, merge_snapshots, _estimate_encoded_bytes, + _QUANTILES, _td_bytes, ) @@ -383,7 +384,7 @@ def test_outstanding_bytes_reconcile_exactly(): def test_step_metrics_use_one_unit_per_dimension(): """Every duration is ms, every volume MB. A chart mixing `wall_s` with - `p99_ms` shows 0.008 beside 24.85 and reads as a data-plane bug rather + `p90_ms` shows 0.008 beside 24.85 and reads as a data-plane bug rather than an axis one.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient()) client.register_partition( @@ -402,8 +403,9 @@ def test_step_metrics_use_one_unit_per_dimension(): def test_step_metrics_tail_is_exact_not_bucketed(): """Per-step percentiles came off a histogram that is never reset, so - they went flat and quantised to bucket edges (p99 of a single sample in - (10, 25] is always 10 + 15*0.99 = 24.85). ``max_ms`` is exact and + they went flat and quantised to bucket edges (a tail quantile of a + single sample in (10, 25] always lands on the same interpolated + point). ``max_ms`` is exact and tracks the slowest call actually seen.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient()) client.register_partition( @@ -412,11 +414,11 @@ def test_step_metrics_tail_is_exact_not_bucketed(): client._emit("put", "p", 1, 8, monotonic() - 0.030, "ok") # a 30 ms call metrics = client.get_step_metrics(1.0) - assert "put/p99_ms" not in metrics and "step/put/p50_ms" not in metrics + assert "put/p90_ms" not in metrics and "step/put/p50_ms" not in metrics assert metrics["step/put/max_ms"] >= 30.0 assert metrics["step/put/max_ms"] != pytest.approx(24.85, abs=0.5), "bucket edge" # and one call supports no percentile at all, in either view - assert "p99_ms" not in client.snapshot()["by_op"]["put"] + assert "p90_ms" not in client.snapshot()["by_op"]["put"] client.close() @@ -883,8 +885,8 @@ def test_cluster_percentiles_never_exceed_the_measured_max(): assert metrics["step/put/max_ms"] == pytest.approx(120.0, abs=2.0) assert metrics["step/put/p50_ms"] <= metrics["step/put/max_ms"] - assert metrics["step/put/p99_ms"] <= metrics["step/put/max_ms"] - assert metrics["step/put/p50_ms"] <= metrics["step/put/p99_ms"] + assert metrics["step/put/p90_ms"] <= metrics["step/put/max_ms"] + assert metrics["step/put/p50_ms"] <= metrics["step/put/p90_ms"] def test_cluster_percentiles_withheld_below_a_useful_sample_count(): @@ -997,8 +999,8 @@ def test_snapshot_percentiles_never_exceed_the_measured_max(): for view in (client.snapshot(), merge_snapshots([client.snapshot()])): stats = view["by_op"]["register"] assert stats["p50_ms"] <= stats["max_ms"], stats - assert stats["p99_ms"] <= stats["max_ms"], stats - assert stats["p50_ms"] <= stats["p99_ms"], stats + assert stats["p90_ms"] <= stats["max_ms"], stats + assert stats["p50_ms"] <= stats["p90_ms"], stats client.close() @@ -1032,7 +1034,7 @@ def test_breakdown_table_leaves_withheld_series_empty(): {"step/put/calls": 3, "step/put/wall_ms": 5.0, "step/put/max_ms": 2.0} ) row = rows[0] - assert row[columns.index("p99_ms")] is None + assert row[columns.index("p90_ms")] is None assert row[columns.index("overhead_ms")] is None assert row[columns.index("calls")] == 3 @@ -1098,10 +1100,14 @@ def test_cluster_per_op_time_is_reported_per_call(): def test_each_quantile_waits_for_the_samples_it_needs(): - """n samples resolve a quantile no higher than about 1 - 1/n, so a p99 - wants 100 and a p50 only a couple. One threshold for both meant 48 calls - reported neither, when the median was perfectly real and only the p99 - would have been the maximum wearing a percentile label.""" + """Each quantile needs about four observations above its rank to mean + anything, so they cannot share one gate: 48 calls carry a real median + and no usable tail, and a single threshold for both reported neither. + + The tail one is p90 rather than p99 for the same reason. A step holds + tens of calls, and a p99 off 58 of them equalled the maximum 80% of the + time -- ``max_ms`` under a more precise-sounding name. + """ def metrics_for(n_calls): client = MetricsDataPlaneClient(NoOpDataPlaneClient()) @@ -1111,11 +1117,23 @@ def metrics_for(n_calls): return cluster_step_metrics(merge_snapshots([client.snapshot()]), {}, 1.0) assert "step/put/p50_ms" not in metrics_for(10), "too thin for either" - mid = metrics_for(48) - assert "step/put/p50_ms" in mid, "a median off 48 calls is real" - assert "step/put/p99_ms" not in mid, "a p99 off 48 calls is the max" - both = metrics_for(120) - assert "step/put/p50_ms" in both and "step/put/p99_ms" in both + mid = metrics_for(30) + assert "step/put/p50_ms" in mid, "a median off 30 calls is real" + assert "step/put/p90_ms" not in mid, "a p90 off 30 calls is not" + both = metrics_for(58) + assert "step/put/p50_ms" in both and "step/put/p90_ms" in both + + +def test_no_quantile_finer_than_the_sample_size_can_resolve(): + """Guard on the choice itself, not just the gate: reporting a p99 at + these per-step call counts would mean reporting the maximum twice.""" + assert 0.99 not in {q for q, _, _ in _QUANTILES} + for q, name, min_samples in _QUANTILES: + above_the_rank = min_samples * (1 - q) + assert above_the_rank >= 4 - 1e-9, ( + f"{name} is gated at {min_samples}, leaving only " + f"{above_the_rank:.1f} observations above its rank" + ) class _JaggedEcho(NoOpDataPlaneClient): From 3e51c58f78c77deb0ab696a8cf9bfe4bf5701d71 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 15:14:18 -0700 Subject: [PATCH 30/70] feat(data-plane): chart time shares, not 32 per-op series Four ops times eight fields is 32 series saying one thing, and a dashboard of 32 lines does not answer the question people bring to it: where is my time going. What answers it is a share. Two independent decompositions of one step's data-plane time are now emitted. `step/share/by_op/{op}` says which call is expensive and sums to 1. `step/share/by_cause/{overhead,transfer}` says whether that time is fixed per-request cost or bandwidth, which is the actionable half: on a real TransferQueue run it reads 53% overhead against 18% transfer, meaning fewer and larger requests buy more than faster transport. The by-cause pair covers only ops whose affine fit is identifiable, so it sums to at most 1 -- the gap is time that could not be attributed, not time that did not happen. `headline_series` picks the ~10 series worth a chart; the per-op detail stays in the breakdown table, which now leads with `share_pct` and carries `mb`. Both are derived from one dict, so a chart and the table beside it cannot disagree. Cluster naming follows suit: `wall_ms_per_process` becomes `step/frac_of_step` (the same per-process division, expressed as the fraction it always was) and `observability_overhead_*` becomes `step/self/*`. Two bugs found on the way: - The cluster path's per-step `max_ms` was the *lifetime* max. A maximum cannot be differenced out of two cumulative readings the way `calls` and `wall_ms` can, so after one 50 ms call every later step still read 50 ms -- the same defect as logging a cumulative percentile, which the single-process path already fixed with a reset window. `snapshot()` now takes `reset_step_window`, which the once-per-step collectors pass and an inspection snapshot does not. - `step/self/overhead_ms` has the same three-part shape as a per-op series and was becoming a "self" row in the breakdown table beside put and get. Verified: 180 unit tests, 48/48 metric audit checks, 23/23 hash-guard checks, and two e2e runs -- 10 Ray processes, and a single process against a live TransferQueue where the by-cause split is identifiable. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 5 +- nemo_rl/data_plane/README.md | 89 +++--- nemo_rl/data_plane/observability.py | 285 ++++++++++++-------- nemo_rl/data_plane/worker_mixin.py | 7 +- nemo_rl/models/policy/tq_policy.py | 4 +- tests/unit/data_plane/test_observability.py | 133 +++++++-- 6 files changed, 363 insertions(+), 160 deletions(-) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index 72bc3598882..9a167e8a033 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -84,6 +84,7 @@ MetricsDataPlaneClient, breakdown_table, cluster_step_metrics, + headline_series, merge_snapshots, ) from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS, DP_TRAIN_FIELDS @@ -445,12 +446,12 @@ def _log_data_plane_metrics( merged, prev, total_step_time, collect_ms=collect_ms ) policy._prev_cluster_snapshot = merged - logger.log_metrics(metrics, step, prefix="data_plane/cluster") + logger.log_metrics(headline_series(metrics), step, prefix="data_plane/cluster") _log_breakdown(logger, metrics, step, "data_plane/cluster/breakdown") else: # Single process, or the fan-out could not reach the workers. metrics = client.get_step_metrics(total_step_time) - logger.log_metrics(metrics, step, prefix="data_plane/driver") + logger.log_metrics(headline_series(metrics), step, prefix="data_plane/driver") _log_breakdown(logger, metrics, step, "data_plane/driver/breakdown") print( f" • data plane: {metrics['step/wall_ms']:.0f}ms, " diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 3952bbc77c6..b6b6dd6b1e8 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -464,21 +464,49 @@ affine fit, throughput — is recomputed from the merged totals, never averaged across ranks (averaging per-rank percentiles does not give a cluster percentile). -**A per-op breakdown table** is logged alongside the series, under -`data_plane/{cluster,driver}/breakdown` — one row per op, ordered by wall -time so the expensive one reads first: - -| op | calls | mean_ms | max_ms | wall_ms | overhead_ms | transfer_ms | p50_ms | p90_ms | -|---|---:|---:|---:|---:|---:|---:|---:|---:| -| get | 200 | 11.17 | 16.0 | 2232 | 6.03 | 5.15 | 12.9 | 15.2 | -| put | 48 | 23.66 | 31.9 | 1135 | 12.06 | 11.60 | 22.4 | 30.1 | - -Everything in ms on that row is **per call**, and the two split terms add -to `mean_ms`: that `get` row reads "each call cost 11.17 ms, of which 6.03 -was fixed per-request overhead and 5.15 was bandwidth at this step's mean -request size". Per call the overhead term *is* the fitted constant, so it -is comparable against a hardware number. `calls` and `wall_ms` are the only -extensive columns. +**What gets charted is the bottleneck, not the detail.** Four ops times +eight fields is 32 series saying one thing, and a dashboard of 32 lines +does not answer "where is my time going". So the emitted series are the +totals and two *shares*, and the per-op detail goes in a table beside them: + +| series | what it answers | +|---|---| +| `step/frac_of_step` | is the data plane worth optimising at all? | +| `step/share/by_op/{put,get,clear,register}` | which call is expensive? (sums to 1) | +| `step/share/by_cause/{overhead,transfer}` | fixed per-request cost, or bandwidth? | +| `step/wall_ms`, `step/comm_volume_mb` | how much time and traffic | +| `now/bytes_outstanding_mb`, `now/n_processes` | occupancy, fan-out width | +| `step/self/{overhead_ms,frac}` | what measuring cost | + +Two independent decompositions of one total is what makes a bottleneck +legible. From a real TransferQueue run: `by_op` says put 43%, get 31%, +register 17%, clear 9%; `by_cause` says 53% fixed per-request overhead +against 18% bandwidth. That second line is the actionable one — this +workload is overhead-dominated, so fewer, larger requests buy more than +faster transport. + +The by-op shares sum to 1. The by-cause shares cover only the ops whose +affine fit is identifiable, so they sum to at most 1; the gap is time that +could not be attributed, not time that did not happen. `register` and +`clear` move no bytes and never get a split. + +**A per-op breakdown table** carries the detail, under +`data_plane/{cluster,driver}/breakdown` — one row per op, ordered by share +so the bottleneck is the first line read: + +| op | share_pct | calls | wall_ms | mean_ms | max_ms | p50_ms | p90_ms | overhead_ms | transfer_ms | mb | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| put | 43.0 | 2 | 53.9 | 26.9 | 29.4 | — | — | 19.4 | 6.39 | 1.32 | +| get | 30.9 | 2 | 38.7 | 19.4 | 21.5 | — | — | 13.8 | 4.63 | 1.03 | +| register | 17.1 | 1 | 21.4 | 21.4 | 21.4 | — | — | — | — | 0 | +| clear | 8.93 | 1 | 11.2 | 11.2 | 11.2 | — | — | — | — | 0 | + +Everything in ms on that row is **per call** except `wall_ms`, and the two +split terms add to `mean_ms`: that `put` row reads "each call cost 26.9 ms, +of which 19.4 was fixed per-request overhead and 6.39 was bandwidth at this +step's mean request size". Per call the overhead term *is* the fitted +constant, so it is comparable against a hardware number. `calls`, +`wall_ms` and `mb` are the only extensive columns. **Per-call figures describe the wire; sums describe the run.** `wall_ms` on the cluster path is summed over processes that ran concurrently, so it is process-time @@ -487,11 +515,7 @@ while the wall clock was 279. Dividing by the process count only trades one arbitrary denominator for another. Per call is invariant to both DP degree and batch size: the same workload at 8 and at 32 ranks reports 11.16 and 11.06 ms while `wall_ms` quadruples. Use `mean_ms` to compare runs and -cluster sizes, `wall_ms` to attribute cost across ops within one step. - -`overhead_ms` and `transfer_ms` stack to the measured `wall_ms` — that row -reads "2233 ms of get was 1207 ms of fixed per-call cost and 1026 ms of -bandwidth", which is a decision you can act on. +cluster sizes, `share_pct` to attribute cost across ops within one step. A stack of line charts answers "how did put's wall time trend"; this answers "where did the step go", which is a table. Cells are empty rather @@ -514,17 +538,18 @@ leak signal the metric exists for: bytes put and never cleared. Two more read differently in the cluster view and are named to say so: -- **`step/wall_ms_per_process`**, not a bare fraction. `wall_ms` sums - processes that ran concurrently, so dividing it by one step's wall clock - exceeded 1 whenever they overlapped (measured 1.054 across ten - processes). Per process it is a duration in ms, like everything else. -- **Percentiles are per step and clamped to the exact `max_ms`**, and are - withheld entirely below 50 calls in the window. Bucket interpolation - spreads a bucket's samples uniformly across it, so calls clustered low in - a wide bucket read high — 160 calls of 120 ms all land in `(100, 250]` - and interpolate to a p50 of 175, above every call observed and above the - max reported beside it. The max is measured exactly, so it is the tighter - bound. +- **`step/frac_of_step` is per process.** `wall_ms` sums processes that ran + concurrently, so dividing it by one step's wall clock exceeded 1 whenever + they overlapped (measured 1.054 across ten processes) and read as "105% + of the step". Divided per process it is the mean share of the step a + process spent in the data plane, which is what the name claims. +- **`step/{op}/max_ms` is scoped to the step by being reset**, not by being + differenced. A maximum cannot be recovered from two cumulative readings + the way `calls` and `wall_ms` can, so the reader that consumes it zeroes + it — `snapshot(reset_step_window=True)`, which the once-per-step + collector passes and an inspection snapshot does not. Without that the + cluster path reported the lifetime max: after one 50 ms call every later + step still read 50 ms. `grpo_train_sync` fans out to the driver and every policy worker, and logs the combined result under `data_plane/cluster/` instead of the driver's @@ -536,7 +561,7 @@ processes, against a 6x wider view of the traffic. The fan-out is best-effort — a rank that cannot answer is dropped rather than failing the step. -`observability_overhead_ms` reports what the measurement itself cost — the +`step/self/overhead_ms` reports what the measurement itself cost — the whole bill, both halves: - every process's wrapper time (its wall time minus the time its inner diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index e78f63d075e..10928890d08 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -450,6 +450,79 @@ def _step_deltas(snap: dict[str, Any], prev: dict[str, Any]) -> dict[str, float] } +def _op_step_stats( + by_op: dict[str, Any], prev_ops: dict[str, Any] +) -> dict[str, dict[str, float]]: + """This step's per-op detail, keyed by op, from two snapshots. + + Shared by the single-process and cluster paths so the two cannot drift, + and used for both the emitted shares and the breakdown table -- one + computation, so a chart and the table beside it can never disagree. + + ``max_ms`` comes from ``step_max_ms``, which the reader resets, rather + than from the cumulative ``max_ms``: a maximum is not differenceable, so + the cumulative one latches at the worst call ever seen and never comes + back down. Ops with no calls this step are absent, not zero. + """ + out: dict[str, dict[str, float]] = {} + for op, st in by_op.items(): + prev_op = prev_ops.get(op, {}) + calls = st["calls"] - prev_op.get("calls", 0) + if calls <= 0: + continue + op_ms = st["wall_ms"] - prev_op.get("wall_ms", 0.0) + op_bytes = st["n_bytes"] - prev_op.get("n_bytes", 0) + row: dict[str, float] = { + "calls": calls, + "wall_ms": op_ms, + # Per call, which is the only form of this that describes the + # wire rather than the shape of the run: ``wall_ms`` is summed + # over concurrent processes and so scales with DP degree. + "mean_ms": op_ms / calls, + "max_ms": st.get("step_max_ms", 0.0), + "mb": op_bytes / 1e6, + } + step_hist = [ + now - was + for now, was in zip( + st["latency_hist"], + prev_op.get("latency_hist") or [0] * len(st["latency_hist"]), + ) + ] + row.update(_clamped_percentiles(step_hist, row["max_ms"])) + split = _latency_split(st["fit"], calls, op_bytes) + if split: + row["overhead_ms"], row["transfer_ms"] = split + out[op] = row + return out + + +def _time_shares(per_op: dict[str, dict[str, float]]) -> dict[str, float]: + """Where this step's data-plane time went, as fractions of the total. + + Two independent decompositions of one total, which is what makes a + bottleneck legible: **by op** (which call is expensive) and **by cause** + (fixed per-request overhead vs bandwidth). The by-op shares sum to 1. + The by-cause shares cover only the ops whose affine fit is identifiable + -- an op whose requests were all the same size cannot be split -- so + they sum to at most 1, and the gap is time that could not be attributed + rather than time that did not happen. + """ + total = sum(r["wall_ms"] for r in per_op.values()) + if total <= 0: + return {} + shares = { + f"step/share/by_op/{op}": r["wall_ms"] / total for op, r in per_op.items() + } + for cause in ("overhead", "transfer"): + attributed = sum( + r[f"{cause}_ms"] * r["calls"] for r in per_op.values() if f"{cause}_ms" in r + ) + if attributed > 0: + shares[f"step/share/by_cause/{cause}"] = attributed / total + return shares + + def _latency_split( fit: dict[str, Any], calls: int, op_bytes: int ) -> tuple[float, float] | None: @@ -551,7 +624,7 @@ def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: "sum_bytes_ms", "sum_ms_sq", ) -_OP_MAX = ("max_ms",) +_OP_MAX = ("max_ms", "step_max_ms") def merge_snapshots(snapshots: "list[dict[str, Any]]") -> dict[str, Any]: @@ -660,80 +733,86 @@ def cluster_step_metrics( metrics = _step_deltas(merged, prev) metrics.update( { - # ``wall_ms`` sums processes that ran concurrently, so it can - # exceed the step's own wall clock. Per process it reads as a - # duration in ms, like everything else here. - "step/wall_ms_per_process": wall_ms / n_procs, - "now/n_processes": n_procs, - "step/observability_overhead_ms": overhead_ms, - "step/observability_overhead_frac": ( - overhead_ms / wall_ms if wall_ms > 0 else 0.0 + # The one metric that says whether optimising the data plane is + # worth anything: per-op shares say where its time went, never + # whether it mattered against compute. ``wall_ms`` sums + # processes that ran concurrently, so dividing it by one step's + # wall clock exceeds 1 whenever they overlapped (measured 1.054 + # across ten processes) -- correct arithmetic that reads as + # "105% of the step". Per process it is bounded and answers the + # question people actually ask of it. + "step/frac_of_step": ( + (wall_ms / n_procs) / (step_time_s * 1e3) if step_time_s > 0 else 0.0 ), + "now/n_processes": n_procs, + "step/self/overhead_ms": overhead_ms, + "step/self/frac": overhead_ms / wall_ms if wall_ms > 0 else 0.0, } ) - prev_ops = prev.get("by_op", {}) - for op, stats in merged["by_op"].items(): - prev_op = prev_ops.get(op, {}) - calls = stats["calls"] - prev_op.get("calls", 0) - if calls <= 0: - continue - metrics[f"step/{op}/calls"] = calls - op_wall_ms = stats["wall_ms"] - prev_op.get("wall_ms", 0.0) - # How long one call took, which is the only form of this that - # describes the wire rather than the shape of the run. ``wall_ms`` - # is summed over concurrent processes, so it scales with DP degree; - # dividing by the process count instead just trades one arbitrary - # denominator for another. Per call is invariant to both, and - # comparable across runs and cluster sizes. - metrics[f"step/{op}/mean_ms"] = op_wall_ms / calls - metrics[f"step/{op}/wall_ms"] = op_wall_ms - metrics[f"step/{op}/max_ms"] = stats["max_ms"] - # Percentiles over THIS step's calls, summed across ranks, not over - # the lifetime: a cumulative percentile beside a step-scoped max is - # two different windows on one chart, and it showed a max below its - # own median. Emitted only when the window holds enough calls to - # out-resolve the buckets -- with a wide DP degree a step easily - # clears it, and on a narrow one the honest answer is silence. - step_hist = [ - now - was - for now, was in zip( - stats["latency_hist"], - prev_op.get("latency_hist") or [0] * len(stats["latency_hist"]), - ) - ] - # Off the step's histogram delta, not the cumulative one - # :func:`_derive_op_metrics` used. - for name, value in _clamped_percentiles(step_hist, stats["max_ms"]).items(): - metrics[f"step/{op}/{name}"] = value - # The split belongs here more than on the driver path: this fit is - # over every rank's sufficient statistics, so it has far more - # samples and far more size variation -- and size variation is - # exactly what decides whether the split is identifiable at all. - split = _latency_split( - stats["fit"], calls, stats["n_bytes"] - prev_op.get("n_bytes", 0) - ) - if split: - ( - metrics[f"step/{op}/overhead_ms"], - metrics[f"step/{op}/transfer_ms"], - ) = split + per_op = _op_step_stats(merged["by_op"], prev.get("by_op", {})) + metrics.update(_time_shares(per_op)) + for op, row in per_op.items(): + for field_name, value in row.items(): + metrics[f"step/{op}/{field_name}"] = value return metrics +# What goes on a chart. Everything else this module computes is per-op +# detail, which belongs in the breakdown table beside it: four ops times +# eight fields is 32 series saying one thing, and a dashboard of 32 lines +# does not answer "what is my bottleneck" -- a table sorted by share does. +# The full dict is still returned, so the table and the series are derived +# from one computation and cannot disagree. +_HEADLINE = ( + "step/wall_ms", + "step/frac_of_step", + "step/comm_volume_mb", + "now/bytes_outstanding_mb", + "now/n_processes", + "step/self/overhead_ms", + "step/self/frac", +) +_HEADLINE_PREFIXES = ("step/share/", "step/hash/", "step/self/") +# ``step//`` is the per-op shape, so these reserved middles must +# not be mistaken for op tags -- ``step/self/overhead_ms`` was becoming a +# "self" row in the breakdown table beside put and get. +_RESERVED_NAMESPACES = frozenset({"share", "hash", "self"}) + + +def headline_series(metrics: dict[str, float]) -> dict[str, float]: + """The subset of ``metrics`` worth a time series. + + Args: + metrics: A flat dict from :func:`cluster_step_metrics` or + :meth:`MetricsDataPlaneClient.get_step_metrics`. + + Returns: + Totals, time shares, and hash counters -- the per-op detail is + dropped, since :func:`breakdown_table` presents it better. + """ + return { + k: v + for k, v in metrics.items() + if k in _HEADLINE or k.startswith(_HEADLINE_PREFIXES) + } + + # Per-op columns worth a row in the breakdown, in the order they read. # ``overhead_ms``/``transfer_ms`` are only present when the affine fit is # trustworthy, and ``p50_ms``/``p90_ms`` only above the sample gate, so a # row carries None where a series was withheld rather than a zero that # would read as a measurement. _BREAKDOWN_COLUMNS = ( + "share_pct", "calls", + "wall_ms", "mean_ms", "max_ms", - "wall_ms", - "overhead_ms", - "transfer_ms", "p50_ms", "p90_ms", + "overhead_ms", + "transfer_ms", + "mb", ) @@ -745,7 +824,8 @@ def breakdown_table( A stack of line charts answers "how did put's wall time trend"; the question this feeds is "where did this step's time go, across ops, at a glance" -- which is a table, and reading it off eight separate charts is - the wrong tool. + the wrong tool. Rows are ordered by share of data-plane time, so the + bottleneck is the first line read. Built from the metrics dict that is logged rather than from the snapshot it came from, so the table and the series can never disagree: a value @@ -758,19 +838,26 @@ def breakdown_table( :func:`cluster_step_metrics`. Returns: - ``(columns, rows)`` for :meth:`Logger.log_table`, ops sorted by - descending wall time so the expensive one is the first line read. + ``(columns, rows)`` for :meth:`Logger.log_table`. """ per_op: dict[str, dict[str, float]] = {} for key, value in metrics.items(): parts = key.split("/") - if len(parts) != 3 or parts[0] != "step": - continue - _, op, field = parts - if field in _BREAKDOWN_COLUMNS: - per_op.setdefault(op, {})[field] = value + if len(parts) == 4 and parts[:3] == ["step", "share", "by_op"]: + per_op.setdefault(parts[3], {})["share_pct"] = 100.0 * value + elif ( + len(parts) == 3 + and parts[0] == "step" + and parts[1] not in _RESERVED_NAMESPACES + and parts[2] in _BREAKDOWN_COLUMNS + ): + per_op.setdefault(parts[1], {})[parts[2]] = value rows = [ [op, *(stats.get(col) for col in _BREAKDOWN_COLUMNS)] + # By wall time, which orders identically to share (share is wall + # time over the same total) and is present even when the shares + # are not -- a table built from a partial metrics dict still reads + # worst-first. for op, stats in sorted( per_op.items(), key=lambda kv: -kv[1].get("wall_ms", 0.0) ) @@ -923,7 +1010,7 @@ def __init__( # differencing and unit-conversion logic. self._prev_snapshot: dict[str, Any] = {} - def snapshot(self) -> dict[str, Any]: + def snapshot(self, reset_step_window: bool = False) -> dict[str, Any]: """Return cumulative totals plus live byte / key outstanding counts. ``total_wall_ms`` is the aggregate data-plane cost; ``by_op`` breaks @@ -931,6 +1018,15 @@ def snapshot(self) -> dict[str, Any]: backends can be compared without post-processing. Throughput is omitted for ops that move no payload (e.g. ``claim_meta``, whose wall time is producer wait, not transfer). + + Args: + reset_step_window: Zero each op's ``step_max_ms`` after reading + it, opening a fresh window. A maximum cannot be differenced + out of a cumulative counter the way ``calls`` and + ``wall_ms`` can, so the only way to scope one to a step is + to reset it -- and the reader that consumes it is the one + that has to. Left off by default so an inspection snapshot + never disturbs the step series. """ out = asdict(self._stats) out["n_keys_outstanding"] = sum( @@ -944,6 +1040,9 @@ def snapshot(self) -> dict[str, Any]: out["bytes_written"] = sum(by[o]["n_bytes"] for o in _WRITE_OPS if o in by) out["bytes_read"] = sum(by[o]["n_bytes"] for o in _READ_OPS if o in by) out["comm_volume_bytes"] = out["bytes_written"] + out["bytes_read"] + if reset_step_window: + for bucket in self._stats.by_op.values(): + bucket.step_max_ms = 0.0 return out def get_step_metrics(self, step_time_s: float) -> dict[str, float]: @@ -957,14 +1056,11 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: data plane is worth anything: per-op shares only say where data-plane time went, never whether it mattered against compute. """ - snap = self.snapshot() + # Reading the step maxima is what closes the window: the values + # just read are this step's, and anything after belongs to the next. + snap = self.snapshot(reset_step_window=True) prev = self._prev_snapshot self._prev_snapshot = snap - # Snapshot first, then open a fresh window: the values just read are - # this step's, and anything after this call belongs to the next one. - step_maxima = {op: b.step_max_ms for op, b in self._stats.by_op.items()} - for bucket in self._stats.by_op.values(): - bucket.step_max_ms = 0.0 wall_ms = snap["total_wall_ms"] - prev.get("total_wall_ms", 0.0) # Every duration is ms and every volume is MB, with no exceptions: @@ -978,15 +1074,8 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: ) if self._verify_tensor_hash: hv, prev_hv = snap["hash_verify"], prev.get("hash_verify", {}) - metrics["step/hash/rows_checked"] = hv["rows_checked"] - prev_hv.get( - "rows_checked", 0 - ) - metrics["step/hash/rows_recorded"] = hv["rows_recorded"] - prev_hv.get( - "rows_recorded", 0 - ) - metrics["step/hash/rows_unverified"] = hv["rows_unverified"] - prev_hv.get( - "rows_unverified", 0 - ) + for name in ("rows_checked", "rows_recorded", "rows_unverified"): + metrics[f"step/hash/{name}"] = hv[name] - prev_hv.get(name, 0) metrics["step/hash/mismatches"] = hv["mismatches"] - prev_hv.get( "mismatches", 0 ) @@ -997,33 +1086,11 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: metrics["step/hash/fields_skipped"] = hv["fields_skipped"] - prev_hv.get( "fields_skipped", 0 ) - prev_ops = prev.get("by_op", {}) - for op, st in snap["by_op"].items(): - prev_op = prev_ops.get(op, {}) - calls = st["calls"] - prev_op.get("calls", 0) - if calls <= 0: - continue - op_ms = st["wall_ms"] - prev_op.get("wall_ms", 0.0) - # Four series per op tag. Anything exactly derivable is left - # out rather than logged: mean is wall_ms/calls, and a dashboard - # can divide. ``snapshot()`` still carries the full picture, - # percentiles included, for a one-off inspection. - metrics[f"step/{op}/calls"] = calls - metrics[f"step/{op}/mean_ms"] = op_ms / calls - metrics[f"step/{op}/wall_ms"] = op_ms - # ``max_ms`` rather than p50/p90. Those come off a histogram - # that is never reset, so per step they were a lifetime figure - # that goes flat, quantised to bucket edges. The max is exact - # and says the same thing at the handful of calls per step. - metrics[f"step/{op}/max_ms"] = step_maxima.get(op, 0.0) - split = _latency_split( - st["fit"], calls, st["n_bytes"] - prev_op.get("n_bytes", 0) - ) - if split: - ( - metrics[f"step/{op}/overhead_ms"], - metrics[f"step/{op}/transfer_ms"], - ) = split + per_op = _op_step_stats(snap["by_op"], prev.get("by_op", {})) + metrics.update(_time_shares(per_op)) + for op, row in per_op.items(): + for field_name, value in row.items(): + metrics[f"step/{op}/{field_name}"] = value return metrics def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index a696d8ad836..7fe41007dc9 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -390,10 +390,15 @@ def get_data_plane_snapshot(self) -> "dict[str, Any] | None": Returns ``None`` when observability is off or no client exists, so the driver can filter rather than special-case. The payload is counters only (about 1 kB), not tensors. + + Closes this rank's step window (``step_max_ms``) as it reads, since + the driver calls this once per step. A maximum cannot be differenced + out of a cumulative counter, so without the reset the cluster's + per-step max would latch at the worst call ever seen. """ client = getattr(self, "_dp_client", None) snapshot = getattr(client, "snapshot", None) - return snapshot() if callable(snapshot) else None + return snapshot(reset_step_window=True) if callable(snapshot) else None def _fetch( self, diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index 6fd66312cbb..9ac35ea87e4 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -250,7 +250,9 @@ def collect_data_plane_snapshots(self) -> list[dict[str, Any]]: snapshots: list[dict[str, Any]] = [] client = getattr(self, "dp_client", None) if hasattr(client, "snapshot"): - snapshots.append(client.snapshot()) + # reset_step_window: this call is the once-per-step reader, and + # a max only scopes to a step by being reset by its reader. + snapshots.append(client.snapshot(reset_step_window=True)) try: # ``Policy.run_all_workers_single_data`` already does the # ``ray.get``. Pairing the worker-group call with diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index a84a8d103d5..92ee2336235 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -32,6 +32,7 @@ MetricsDataPlaneClient, breakdown_table, cluster_step_metrics, + headline_series, merge_snapshots, _estimate_encoded_bytes, _QUANTILES, @@ -809,7 +810,7 @@ def test_merge_of_nothing_is_empty(): def test_cluster_step_metrics_report_their_own_cost(): - """``observability_overhead_ms`` is the wrapper's own wall time minus + """``step/self/overhead_ms`` is the wrapper's own wall time minus the inner client's, summed over processes — the bill for measuring, sitting beside what it bought.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient()) @@ -826,26 +827,27 @@ def test_cluster_step_metrics_report_their_own_cost(): metrics = cluster_step_metrics(merged, {}, 1.0) assert metrics["now/n_processes"] == 1 - assert metrics["step/observability_overhead_ms"] > 0, "measuring is never free" + assert metrics["step/self/overhead_ms"] > 0, "measuring is never free" # Deliberately not clamped to 1. Against this no-op inner client the RPC # is instant, so measuring costs more than the thing measured and the # ratio exceeds 100% -- which is exactly the signal worth surfacing. # Against a real backend it lands near 0.01. - assert metrics["step/observability_overhead_frac"] > 0 + assert metrics["step/self/frac"] > 0 assert "step/wall_ms" in metrics and "step/comm_volume_mb" in metrics client.close() -def test_cluster_time_is_per_process_not_a_bare_fraction(): +def test_cluster_frac_of_step_is_per_process_and_bounded(): """``wall_ms`` sums processes that ran concurrently, so dividing it by one step's wall clock exceeded 1 whenever they overlapped and read as - "105% of the step". Reported per process it is a duration in ms, like - everything else, and needs no gloss.""" + "105% of the step". Divided per process it is the mean share of the + step a process spent in the data plane, which is what the name claims: + 10 ranks x 5 calls x 100 ms over a 5 s step is 500 ms each, or 10%.""" ranks = [_rank_with([100.0] * 5) for _ in range(10)] - metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 1.0) + metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 5.0) - assert "busy_frac_mean" not in metrics and "frac_of_step" not in metrics - assert metrics["step/wall_ms_per_process"] == pytest.approx(500.0, rel=0.1) + assert "busy_frac_mean" not in metrics + assert metrics["step/frac_of_step"] == pytest.approx(0.10, rel=0.1) assert metrics["now/n_processes"] == 10 @@ -866,10 +868,7 @@ def test_cluster_overhead_includes_the_collection_fan_out(): without = cluster_step_metrics(merged, {}, 1.0) with_gather = cluster_step_metrics(merged, {}, 1.0, collect_ms=2.31) - delta = ( - with_gather["step/observability_overhead_ms"] - - without["step/observability_overhead_ms"] - ) + delta = with_gather["step/self/overhead_ms"] - without["step/self/overhead_ms"] assert delta == pytest.approx(2.31, rel=1e-6) client.close() @@ -1006,9 +1005,13 @@ def test_snapshot_percentiles_never_exceed_the_measured_max(): def test_breakdown_table_rows_by_op_worst_first(): """One row per op, ordered by wall time, so the expensive op is the - first line read rather than the alphabetically luckiest.""" + first line read rather than the alphabetically luckiest. Share of + data-plane time is the second column for the same reason: the answer + to "what is my bottleneck" should be the top-left of the table.""" metrics = { "step/wall_ms": 100.0, + "step/share/by_op/get": 0.10, + "step/share/by_op/put": 0.90, "step/get/calls": 8, "step/get/wall_ms": 10.0, "step/get/max_ms": 2.0, @@ -1024,6 +1027,8 @@ def test_breakdown_table_rows_by_op_worst_first(): assert [r[0] for r in rows] == ["put", "get"], "worst first" assert len(rows) == 2, "only per-op series become rows" assert rows[0][columns.index("wall_ms")] == 90.0 + assert rows[0][columns.index("share_pct")] == pytest.approx(90.0) + assert columns[1] == "share_pct", "the bottleneck reads first" def test_breakdown_table_leaves_withheld_series_empty(): @@ -1096,7 +1101,7 @@ def test_cluster_per_op_time_is_reported_per_call(): ), "the sum does move with cluster size" columns, rows = breakdown_table(small) - assert "mean_ms" in columns and "wall_ms_per_process" not in columns + assert "mean_ms" in columns and "share_pct" in columns def test_each_quantile_waits_for_the_samples_it_needs(): @@ -1246,3 +1251,101 @@ def test_delta_put_does_not_restate_the_scheme_of_untouched_fields(): client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) assert client.snapshot()["hash_verify"]["mismatches"] == 0 client.close() + + +def _busy_client(op_calls): + """A client that ran ``n`` calls of each named op this step.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + now = monotonic() + for op, (n, ms) in op_calls.items(): + for i in range(n): + client._emit(op, "p", 1, 1_000_000 * (1 + i % 5), now - ms / 1e3, "ok") + return client + + +def test_time_shares_name_the_bottleneck(): + """The question a dashboard has to answer is "which op is expensive", + and 32 per-op line charts do not answer it. The by-op shares are one + decomposition of data-plane time and sum to 1, so the largest is the + bottleneck by construction.""" + client = _busy_client({"get": (100, 9.0), "put": (10, 1.0), "clear": (10, 0.1)}) + metrics = cluster_step_metrics( + merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 1.0 + ) + + shares = {k: v for k, v in metrics.items() if k.startswith("step/share/by_op/")} + assert sum(shares.values()) == pytest.approx(1.0) + assert max(shares, key=shares.__getitem__) == "step/share/by_op/get" + assert shares["step/share/by_op/get"] == pytest.approx(900 / 911, rel=0.05) + client.close() + + +def test_headline_drops_per_op_detail_but_keeps_the_shares(): + """Four ops times eight fields is 32 series saying one thing. The detail + is still computed -- the breakdown table is built from the same dict, so + the two cannot disagree -- but only the totals and shares are charted.""" + client = _busy_client({"get": (100, 9.0), "put": (10, 1.0), "clear": (10, 0.1)}) + metrics = cluster_step_metrics( + merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 1.0 + ) + head = headline_series(metrics) + + assert len(head) < len(metrics) / 2, f"{len(head)} of {len(metrics)}" + assert not [k for k in head if k.split("/")[1:2] in (["get"], ["put"], ["clear"])] + assert "step/share/by_op/get" in head + assert "step/wall_ms" in head and "step/frac_of_step" in head + # and the detail the table needs survives in the full dict + assert breakdown_table(metrics)[1], "table still has rows" + client.close() + + +def test_cluster_step_max_reopens_each_step(): + """A maximum cannot be differenced out of a cumulative counter, so the + cluster path reported the lifetime max: after one 50 ms call every later + step still read 50 ms, which is the same defect as a cumulative + percentile. The reader resets the window as it reads.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + prev, seen = {}, [] + for slowest_ms in (5.0, 50.0, 5.0, 5.0): + now = monotonic() + for i in range(4): + client._emit( + "put", "p", 1, 1_000, now - (slowest_ms if i == 0 else 5.0) / 1e3, "ok" + ) + merged = merge_snapshots([client.snapshot(reset_step_window=True)]) + seen.append(cluster_step_metrics(merged, prev, 1.0)["step/put/max_ms"]) + prev = merged + + assert seen[1] == pytest.approx(50.0, abs=1.0), "the spike shows" + assert seen[2] == pytest.approx(5.0, abs=1.0), "and does not latch" + client.close() + + +def test_snapshot_leaves_the_step_window_alone_unless_asked(): + """``snapshot()`` is also how a human inspects a live client. Resetting + the step window on every call would let an inspection silently blank the + next step's max.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client._emit("put", "p", 1, 1_000, monotonic() - 0.030, "ok") + + assert client.snapshot()["by_op"]["put"]["step_max_ms"] >= 30.0 + assert client.snapshot()["by_op"]["put"]["step_max_ms"] >= 30.0, "still there" + assert client.snapshot(reset_step_window=True)["by_op"]["put"]["step_max_ms"] >= 30 + assert client.snapshot()["by_op"]["put"]["step_max_ms"] == 0.0, "window reopened" + client.close() + + +def test_breakdown_table_ignores_reserved_namespaces(): + """``step/self/overhead_ms`` has the same three-part shape as a per-op + series and was becoming a "self" row in the table beside put and get.""" + columns, rows = breakdown_table( + { + "step/put/calls": 2, + "step/put/wall_ms": 9.0, + "step/self/overhead_ms": 6.2, + "step/hash/mismatches": 0, + "step/share/by_cause/transfer": 0.4, + } + ) + assert [r[0] for r in rows] == ["put"], rows + assert columns[0] == "op" From f1652a0eab854ca064158279a3deaf6274f92a1b Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 15:29:51 -0700 Subject: [PATCH 31/70] refactor(data-plane): name the metric for its denominator, share -> time_pct "share" never said share of what, which is the one thing a reader has to know before acting on it. The denominator is data-plane time, not the step: `by_op/put = 43` means 43% of the time spent inside the data plane went to put, and says nothing about whether the data plane mattered against compute -- that is `frac_of_step`, which divides by the step's own wall clock. A workload can be 43% put and not be worth touching. So `step/share/by_op/{op}` becomes `step/time_pct/by_op/{op}`, in percent rather than a fraction so the name and the value agree and so it matches the table column beside it. `by_cause/overhead` becomes `by_cause/fixed_overhead`, which is what it measures. The docstring now states the denominator, that by_op sums to 100 while by_cause sums to at most 100, and that on the cluster path these are percentages of aggregate process-time rather than elapsed time -- the right denominator for "what should I optimise" and the wrong one for "what blocked the step". README gains a worked example reading the two denominators together. Verified: 180 unit tests, 48/48 metric audit checks, 23/23 hash-guard checks, and a live TransferQueue e2e. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 49 +++++++---- nemo_rl/data_plane/observability.py | 95 ++++++++++++++------- tests/unit/data_plane/test_observability.py | 45 +++++----- 3 files changed, 120 insertions(+), 69 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index b6b6dd6b1e8..381481bd109 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -467,34 +467,49 @@ cluster percentile). **What gets charted is the bottleneck, not the detail.** Four ops times eight fields is 32 series saying one thing, and a dashboard of 32 lines does not answer "where is my time going". So the emitted series are the -totals and two *shares*, and the per-op detail goes in a table beside them: +totals and `time_pct`, with the per-op detail in a table beside them: | series | what it answers | |---|---| | `step/frac_of_step` | is the data plane worth optimising at all? | -| `step/share/by_op/{put,get,clear,register}` | which call is expensive? (sums to 1) | -| `step/share/by_cause/{overhead,transfer}` | fixed per-request cost, or bandwidth? | +| `step/time_pct/by_op/{put,get,clear,register}` | which call is expensive? | +| `step/time_pct/by_cause/{fixed_overhead,transfer}` | is that fixed per-request cost, or moving bytes? | | `step/wall_ms`, `step/comm_volume_mb` | how much time and traffic | | `now/bytes_outstanding_mb`, `now/n_processes` | occupancy, fan-out width | | `step/self/{overhead_ms,frac}` | what measuring cost | -Two independent decompositions of one total is what makes a bottleneck -legible. From a real TransferQueue run: `by_op` says put 43%, get 31%, -register 17%, clear 9%; `by_cause` says 53% fixed per-request overhead -against 18% bandwidth. That second line is the actionable one — this -workload is overhead-dominated, so fewer, larger requests buy more than -faster transport. +**`time_pct` is a percentage of data-plane time, not of the step.** The +denominator is `sum(wall_ms)` over the ops that ran, so +`by_op/put = 43` reads "43% of the time spent inside the data plane went to +put". Whether that time mattered against compute is the *other* metric: +`frac_of_step` divides by the step's own wall clock. Read them together — +a workload can be 43% put and still not be worth touching. -The by-op shares sum to 1. The by-cause shares cover only the ops whose -affine fit is identifiable, so they sum to at most 1; the gap is time that -could not be attributed, not time that did not happen. `register` and -`clear` move no bytes and never get a split. +``` +step/frac_of_step 0.87 the data plane is 87% of the step, so it matters +step/time_pct/by_op/put 43.0 and within it, put is the largest piece +step/time_pct/by_cause/... + fixed_overhead 53.1 that time is mostly per-request cost, + transfer 17.6 not bandwidth -- batch, don't tune the wire +``` + +Two decompositions of one total, because either alone leaves the next +question unanswered. `by_op` sums to 100 by construction. `by_cause` sums +to *at most* 100: only ops with an identifiable affine fit can be split, so +the remainder (here 29%, the `register` and `clear` calls that move no +bytes) is time that could not be attributed rather than time that did not +happen. + +On the cluster path `wall_ms` is summed over processes that ran +concurrently, so these are percentages of aggregate **process-time**, not of +elapsed time. That is the right denominator for "what should I optimise" +and the wrong one for "what blocked the step". **A per-op breakdown table** carries the detail, under -`data_plane/{cluster,driver}/breakdown` — one row per op, ordered by share -so the bottleneck is the first line read: +`data_plane/{cluster,driver}/breakdown` — one row per op, ordered by +`time_pct` so the bottleneck is the first line read: -| op | share_pct | calls | wall_ms | mean_ms | max_ms | p50_ms | p90_ms | overhead_ms | transfer_ms | mb | +| op | time_pct | calls | wall_ms | mean_ms | max_ms | p50_ms | p90_ms | overhead_ms | transfer_ms | mb | |---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| | put | 43.0 | 2 | 53.9 | 26.9 | 29.4 | — | — | 19.4 | 6.39 | 1.32 | | get | 30.9 | 2 | 38.7 | 19.4 | 21.5 | — | — | 13.8 | 4.63 | 1.03 | @@ -515,7 +530,7 @@ while the wall clock was 279. Dividing by the process count only trades one arbitrary denominator for another. Per call is invariant to both DP degree and batch size: the same workload at 8 and at 32 ranks reports 11.16 and 11.06 ms while `wall_ms` quadruples. Use `mean_ms` to compare runs and -cluster sizes, `share_pct` to attribute cost across ops within one step. +cluster sizes, `time_pct` to attribute cost across ops within one step. A stack of line charts answers "how did put's wall time trend"; this answers "where did the step go", which is a table. Cells are empty rather diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 10928890d08..f8600727889 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -456,7 +456,7 @@ def _op_step_stats( """This step's per-op detail, keyed by op, from two snapshots. Shared by the single-process and cluster paths so the two cannot drift, - and used for both the emitted shares and the breakdown table -- one + and used for both the emitted percentages and the breakdown table -- one computation, so a chart and the table beside it can never disagree. ``max_ms`` comes from ``step_max_ms``, which the reader resets, rather @@ -497,30 +497,58 @@ def _op_step_stats( return out -def _time_shares(per_op: dict[str, dict[str, float]]) -> dict[str, float]: - """Where this step's data-plane time went, as fractions of the total. +def _time_pct(per_op: dict[str, dict[str, float]]) -> dict[str, float]: + """Where this step's data-plane time went, in percent. - Two independent decompositions of one total, which is what makes a - bottleneck legible: **by op** (which call is expensive) and **by cause** - (fixed per-request overhead vs bandwidth). The by-op shares sum to 1. - The by-cause shares cover only the ops whose affine fit is identifiable - -- an op whose requests were all the same size cannot be split -- so - they sum to at most 1, and the gap is time that could not be attributed - rather than time that did not happen. + **The denominator is data-plane time, not the step.** Every value here + is a percentage of ``sum(wall_ms)`` over the ops that ran this step, so + ``by_op/put = 43`` reads "43% of the time spent inside the data plane + went to put". Whether that time mattered at all against compute is a + different question, answered by ``step/frac_of_step``, which divides by + the step's own wall clock. A workload can be 43% put and still not be + worth touching. + + Two decompositions of that one total, because either alone leaves the + next question unanswered: + + - ``by_op`` -- which call is expensive. Sums to 100 by construction. + - ``by_cause`` -- whether that time is fixed per-request cost or moving + bytes. ``overhead_ms``/``transfer_ms`` are per call, so they are + multiplied back by the call count to compare against the same total. + Only ops with an identifiable affine fit can be split (an op whose + requests were all one size cannot be), so this sums to *at most* 100 + and the remainder is time that could not be attributed -- not time + that did not happen. + + On the cluster path ``wall_ms`` is summed over processes that ran + concurrently, so these are percentages of aggregate process-time rather + than of elapsed time. That is the right denominator for "what should I + optimise" and the wrong one for "what blocked the step". + + Args: + per_op: Per-op step detail from :func:`_op_step_stats`. + + Returns: + ``step/time_pct/by_op/{op}`` and ``step/time_pct/by_cause/{cause}``, + each in percent. Empty when no op ran. """ total = sum(r["wall_ms"] for r in per_op.values()) if total <= 0: return {} - shares = { - f"step/share/by_op/{op}": r["wall_ms"] / total for op, r in per_op.items() + pct = { + f"step/time_pct/by_op/{op}": 100.0 * r["wall_ms"] / total + for op, r in per_op.items() } - for cause in ("overhead", "transfer"): + for cause, field_name in ( + ("fixed_overhead", "overhead_ms"), + ("transfer", "transfer_ms"), + ): attributed = sum( - r[f"{cause}_ms"] * r["calls"] for r in per_op.values() if f"{cause}_ms" in r + r[field_name] * r["calls"] for r in per_op.values() if field_name in r ) if attributed > 0: - shares[f"step/share/by_cause/{cause}"] = attributed / total - return shares + pct[f"step/time_pct/by_cause/{cause}"] = 100.0 * attributed / total + return pct def _latency_split( @@ -734,7 +762,7 @@ def cluster_step_metrics( metrics.update( { # The one metric that says whether optimising the data plane is - # worth anything: per-op shares say where its time went, never + # worth anything: per-op percentages say where its time went, never # whether it mattered against compute. ``wall_ms`` sums # processes that ran concurrently, so dividing it by one step's # wall clock exceeds 1 whenever they overlapped (measured 1.054 @@ -750,7 +778,7 @@ def cluster_step_metrics( } ) per_op = _op_step_stats(merged["by_op"], prev.get("by_op", {})) - metrics.update(_time_shares(per_op)) + metrics.update(_time_pct(per_op)) for op, row in per_op.items(): for field_name, value in row.items(): metrics[f"step/{op}/{field_name}"] = value @@ -760,7 +788,7 @@ def cluster_step_metrics( # What goes on a chart. Everything else this module computes is per-op # detail, which belongs in the breakdown table beside it: four ops times # eight fields is 32 series saying one thing, and a dashboard of 32 lines -# does not answer "what is my bottleneck" -- a table sorted by share does. +# does not answer "what is my bottleneck" -- a table sorted by time does. # The full dict is still returned, so the table and the series are derived # from one computation and cannot disagree. _HEADLINE = ( @@ -772,11 +800,11 @@ def cluster_step_metrics( "step/self/overhead_ms", "step/self/frac", ) -_HEADLINE_PREFIXES = ("step/share/", "step/hash/", "step/self/") +_HEADLINE_PREFIXES = ("step/time_pct/", "step/hash/", "step/self/") # ``step//`` is the per-op shape, so these reserved middles must # not be mistaken for op tags -- ``step/self/overhead_ms`` was becoming a # "self" row in the breakdown table beside put and get. -_RESERVED_NAMESPACES = frozenset({"share", "hash", "self"}) +_RESERVED_NAMESPACES = frozenset({"time_pct", "hash", "self"}) def headline_series(metrics: dict[str, float]) -> dict[str, float]: @@ -787,7 +815,7 @@ def headline_series(metrics: dict[str, float]) -> dict[str, float]: :meth:`MetricsDataPlaneClient.get_step_metrics`. Returns: - Totals, time shares, and hash counters -- the per-op detail is + Totals, time percentages, and hash counters -- the per-op detail is dropped, since :func:`breakdown_table` presents it better. """ return { @@ -803,7 +831,7 @@ def headline_series(metrics: dict[str, float]) -> dict[str, float]: # row carries None where a series was withheld rather than a zero that # would read as a measurement. _BREAKDOWN_COLUMNS = ( - "share_pct", + "time_pct", "calls", "wall_ms", "mean_ms", @@ -824,7 +852,8 @@ def breakdown_table( A stack of line charts answers "how did put's wall time trend"; the question this feeds is "where did this step's time go, across ops, at a glance" -- which is a table, and reading it off eight separate charts is - the wrong tool. Rows are ordered by share of data-plane time, so the + the wrong tool. Rows are ordered by their share of data-plane time, so + the bottleneck is the first line read. Built from the metrics dict that is logged rather than from the snapshot @@ -843,8 +872,8 @@ def breakdown_table( per_op: dict[str, dict[str, float]] = {} for key, value in metrics.items(): parts = key.split("/") - if len(parts) == 4 and parts[:3] == ["step", "share", "by_op"]: - per_op.setdefault(parts[3], {})["share_pct"] = 100.0 * value + if len(parts) == 4 and parts[:3] == ["step", "time_pct", "by_op"]: + per_op.setdefault(parts[3], {})["time_pct"] = value elif ( len(parts) == 3 and parts[0] == "step" @@ -854,10 +883,10 @@ def breakdown_table( per_op.setdefault(parts[1], {})[parts[2]] = value rows = [ [op, *(stats.get(col) for col in _BREAKDOWN_COLUMNS)] - # By wall time, which orders identically to share (share is wall - # time over the same total) and is present even when the shares - # are not -- a table built from a partial metrics dict still reads - # worst-first. + # By wall time, which orders identically to ``time_pct`` (that is + # wall time over a common total) and is present even when the + # percentages are not -- a table built from a partial metrics dict + # still reads worst-first. for op, stats in sorted( per_op.items(), key=lambda kv: -kv[1].get("wall_ms", 0.0) ) @@ -1053,8 +1082,8 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: ``VllmGeneration.get_step_metrics`` so trainers stay one line. ``frac_of_step`` is the metric that decides whether optimising the - data plane is worth anything: per-op shares only say where data-plane - time went, never whether it mattered against compute. + data plane is worth anything: ``time_pct`` only says where + data-plane time went, never whether it mattered against compute. """ # Reading the step maxima is what closes the window: the values # just read are this step's, and anything after belongs to the next. @@ -1087,7 +1116,7 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: "fields_skipped", 0 ) per_op = _op_step_stats(snap["by_op"], prev.get("by_op", {})) - metrics.update(_time_shares(per_op)) + metrics.update(_time_pct(per_op)) for op, row in per_op.items(): for field_name, value in row.items(): metrics[f"step/{op}/{field_name}"] = value diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 92ee2336235..1358506ac27 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -1010,8 +1010,8 @@ def test_breakdown_table_rows_by_op_worst_first(): to "what is my bottleneck" should be the top-left of the table.""" metrics = { "step/wall_ms": 100.0, - "step/share/by_op/get": 0.10, - "step/share/by_op/put": 0.90, + "step/time_pct/by_op/get": 10.0, + "step/time_pct/by_op/put": 90.0, "step/get/calls": 8, "step/get/wall_ms": 10.0, "step/get/max_ms": 2.0, @@ -1027,8 +1027,8 @@ def test_breakdown_table_rows_by_op_worst_first(): assert [r[0] for r in rows] == ["put", "get"], "worst first" assert len(rows) == 2, "only per-op series become rows" assert rows[0][columns.index("wall_ms")] == 90.0 - assert rows[0][columns.index("share_pct")] == pytest.approx(90.0) - assert columns[1] == "share_pct", "the bottleneck reads first" + assert rows[0][columns.index("time_pct")] == pytest.approx(90.0) + assert columns[1] == "time_pct", "the bottleneck reads first" def test_breakdown_table_leaves_withheld_series_empty(): @@ -1101,7 +1101,7 @@ def test_cluster_per_op_time_is_reported_per_call(): ), "the sum does move with cluster size" columns, rows = breakdown_table(small) - assert "mean_ms" in columns and "share_pct" in columns + assert "mean_ms" in columns and "time_pct" in columns def test_each_quantile_waits_for_the_samples_it_needs(): @@ -1263,27 +1263,34 @@ def _busy_client(op_calls): return client -def test_time_shares_name_the_bottleneck(): +def test_time_pct_names_the_bottleneck_and_says_of_what(): """The question a dashboard has to answer is "which op is expensive", - and 32 per-op line charts do not answer it. The by-op shares are one - decomposition of data-plane time and sum to 1, so the largest is the - bottleneck by construction.""" + and 32 per-op line charts do not answer it. + + The denominator is data-plane time, not the step: ``by_op`` sums to 100 + by construction, so the largest is the bottleneck *within the data + plane*. Whether the data plane mattered at all is ``frac_of_step``, + which divides by the step's own clock -- here a tenth of a second of + data-plane work inside a 10 s step is 9% of one and 100% of the other. + """ client = _busy_client({"get": (100, 9.0), "put": (10, 1.0), "clear": (10, 0.1)}) metrics = cluster_step_metrics( - merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 1.0 + merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 10.0 ) - shares = {k: v for k, v in metrics.items() if k.startswith("step/share/by_op/")} - assert sum(shares.values()) == pytest.approx(1.0) - assert max(shares, key=shares.__getitem__) == "step/share/by_op/get" - assert shares["step/share/by_op/get"] == pytest.approx(900 / 911, rel=0.05) - client.close() + by_op = {k: v for k, v in metrics.items() if k.startswith("step/time_pct/by_op/")} + assert sum(by_op.values()) == pytest.approx(100.0), "percent of one total" + assert max(by_op, key=by_op.__getitem__) == "step/time_pct/by_op/get" + assert by_op["step/time_pct/by_op/get"] == pytest.approx(100 * 900 / 911, rel=0.05) + # the two denominators are different questions and must not agree + assert metrics["step/frac_of_step"] == pytest.approx(0.0911, rel=0.1) -def test_headline_drops_per_op_detail_but_keeps_the_shares(): +def test_headline_drops_per_op_detail_but_keeps_the_percentages(): """Four ops times eight fields is 32 series saying one thing. The detail is still computed -- the breakdown table is built from the same dict, so - the two cannot disagree -- but only the totals and shares are charted.""" + the two cannot disagree -- but only the totals and percentages are + charted.""" client = _busy_client({"get": (100, 9.0), "put": (10, 1.0), "clear": (10, 0.1)}) metrics = cluster_step_metrics( merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 1.0 @@ -1292,7 +1299,7 @@ def test_headline_drops_per_op_detail_but_keeps_the_shares(): assert len(head) < len(metrics) / 2, f"{len(head)} of {len(metrics)}" assert not [k for k in head if k.split("/")[1:2] in (["get"], ["put"], ["clear"])] - assert "step/share/by_op/get" in head + assert "step/time_pct/by_op/get" in head assert "step/wall_ms" in head and "step/frac_of_step" in head # and the detail the table needs survives in the full dict assert breakdown_table(metrics)[1], "table still has rows" @@ -1344,7 +1351,7 @@ def test_breakdown_table_ignores_reserved_namespaces(): "step/put/wall_ms": 9.0, "step/self/overhead_ms": 6.2, "step/hash/mismatches": 0, - "step/share/by_cause/transfer": 0.4, + "step/time_pct/by_cause/transfer": 40.0, } ) assert [r[0] for r in rows] == ["put"], rows From e29582618121bbbe6fba9218b4fd93510d5c473a Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 15:50:06 -0700 Subject: [PATCH 32/70] refactor(data-plane): time_pct -> pct_of_dataplane The name has to carry its denominator, and "time_pct" does not say percent of what. It is a percentage of data-plane time, so it says so: `step/pct_of_dataplane/by_op/put`. The by_op / by_cause split is unchanged. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 18 +++++------ nemo_rl/data_plane/observability.py | 34 ++++++++++----------- tests/unit/data_plane/test_observability.py | 26 +++++++++------- 3 files changed, 41 insertions(+), 37 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 381481bd109..7e6cf7d6046 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -467,18 +467,18 @@ cluster percentile). **What gets charted is the bottleneck, not the detail.** Four ops times eight fields is 32 series saying one thing, and a dashboard of 32 lines does not answer "where is my time going". So the emitted series are the -totals and `time_pct`, with the per-op detail in a table beside them: +totals and `pct_of_dataplane`, with the per-op detail in a table beside them: | series | what it answers | |---|---| | `step/frac_of_step` | is the data plane worth optimising at all? | -| `step/time_pct/by_op/{put,get,clear,register}` | which call is expensive? | -| `step/time_pct/by_cause/{fixed_overhead,transfer}` | is that fixed per-request cost, or moving bytes? | +| `step/pct_of_dataplane/by_op/{put,get,clear,register}` | which call is expensive? | +| `step/pct_of_dataplane/by_cause/{fixed_overhead,transfer}` | is that fixed per-request cost, or moving bytes? | | `step/wall_ms`, `step/comm_volume_mb` | how much time and traffic | | `now/bytes_outstanding_mb`, `now/n_processes` | occupancy, fan-out width | | `step/self/{overhead_ms,frac}` | what measuring cost | -**`time_pct` is a percentage of data-plane time, not of the step.** The +**`pct_of_dataplane` is a percentage of data-plane time, not of the step.** The denominator is `sum(wall_ms)` over the ops that ran, so `by_op/put = 43` reads "43% of the time spent inside the data plane went to put". Whether that time mattered against compute is the *other* metric: @@ -487,8 +487,8 @@ a workload can be 43% put and still not be worth touching. ``` step/frac_of_step 0.87 the data plane is 87% of the step, so it matters -step/time_pct/by_op/put 43.0 and within it, put is the largest piece -step/time_pct/by_cause/... +step/pct_of_dataplane/by_op/put 43.0 and within it, put is the largest piece +step/pct_of_dataplane/by_cause/... fixed_overhead 53.1 that time is mostly per-request cost, transfer 17.6 not bandwidth -- batch, don't tune the wire ``` @@ -507,9 +507,9 @@ and the wrong one for "what blocked the step". **A per-op breakdown table** carries the detail, under `data_plane/{cluster,driver}/breakdown` — one row per op, ordered by -`time_pct` so the bottleneck is the first line read: +`pct_of_dataplane` so the bottleneck is the first line read: -| op | time_pct | calls | wall_ms | mean_ms | max_ms | p50_ms | p90_ms | overhead_ms | transfer_ms | mb | +| op | pct_of_dataplane | calls | wall_ms | mean_ms | max_ms | p50_ms | p90_ms | overhead_ms | transfer_ms | mb | |---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| | put | 43.0 | 2 | 53.9 | 26.9 | 29.4 | — | — | 19.4 | 6.39 | 1.32 | | get | 30.9 | 2 | 38.7 | 19.4 | 21.5 | — | — | 13.8 | 4.63 | 1.03 | @@ -530,7 +530,7 @@ while the wall clock was 279. Dividing by the process count only trades one arbitrary denominator for another. Per call is invariant to both DP degree and batch size: the same workload at 8 and at 32 ranks reports 11.16 and 11.06 ms while `wall_ms` quadruples. Use `mean_ms` to compare runs and -cluster sizes, `time_pct` to attribute cost across ops within one step. +cluster sizes, `pct_of_dataplane` to attribute cost across ops within one step. A stack of line charts answers "how did put's wall time trend"; this answers "where did the step go", which is a table. Cells are empty rather diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index f8600727889..6b0e2c8f3bc 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -497,13 +497,13 @@ def _op_step_stats( return out -def _time_pct(per_op: dict[str, dict[str, float]]) -> dict[str, float]: +def _pct_of_dataplane(per_op: dict[str, dict[str, float]]) -> dict[str, float]: """Where this step's data-plane time went, in percent. - **The denominator is data-plane time, not the step.** Every value here - is a percentage of ``sum(wall_ms)`` over the ops that ran this step, so - ``by_op/put = 43`` reads "43% of the time spent inside the data plane - went to put". Whether that time mattered at all against compute is a + The name carries the denominator because that is the one thing a reader + has to know before acting on the number: it is a percentage of *the + data plane*, not of the step. ``by_op/put = 43`` reads "43% of the time + spent inside the data plane went to put". Whether that time mattered at all against compute is a different question, answered by ``step/frac_of_step``, which divides by the step's own wall clock. A workload can be 43% put and still not be worth touching. @@ -529,14 +529,14 @@ def _time_pct(per_op: dict[str, dict[str, float]]) -> dict[str, float]: per_op: Per-op step detail from :func:`_op_step_stats`. Returns: - ``step/time_pct/by_op/{op}`` and ``step/time_pct/by_cause/{cause}``, + ``step/pct_of_dataplane/by_op/{op}`` and ``step/pct_of_dataplane/by_cause/{cause}``, each in percent. Empty when no op ran. """ total = sum(r["wall_ms"] for r in per_op.values()) if total <= 0: return {} pct = { - f"step/time_pct/by_op/{op}": 100.0 * r["wall_ms"] / total + f"step/pct_of_dataplane/by_op/{op}": 100.0 * r["wall_ms"] / total for op, r in per_op.items() } for cause, field_name in ( @@ -547,7 +547,7 @@ def _time_pct(per_op: dict[str, dict[str, float]]) -> dict[str, float]: r[field_name] * r["calls"] for r in per_op.values() if field_name in r ) if attributed > 0: - pct[f"step/time_pct/by_cause/{cause}"] = 100.0 * attributed / total + pct[f"step/pct_of_dataplane/by_cause/{cause}"] = 100.0 * attributed / total return pct @@ -778,7 +778,7 @@ def cluster_step_metrics( } ) per_op = _op_step_stats(merged["by_op"], prev.get("by_op", {})) - metrics.update(_time_pct(per_op)) + metrics.update(_pct_of_dataplane(per_op)) for op, row in per_op.items(): for field_name, value in row.items(): metrics[f"step/{op}/{field_name}"] = value @@ -800,11 +800,11 @@ def cluster_step_metrics( "step/self/overhead_ms", "step/self/frac", ) -_HEADLINE_PREFIXES = ("step/time_pct/", "step/hash/", "step/self/") +_HEADLINE_PREFIXES = ("step/pct_of_dataplane/", "step/hash/", "step/self/") # ``step//`` is the per-op shape, so these reserved middles must # not be mistaken for op tags -- ``step/self/overhead_ms`` was becoming a # "self" row in the breakdown table beside put and get. -_RESERVED_NAMESPACES = frozenset({"time_pct", "hash", "self"}) +_RESERVED_NAMESPACES = frozenset({"pct_of_dataplane", "hash", "self"}) def headline_series(metrics: dict[str, float]) -> dict[str, float]: @@ -831,7 +831,7 @@ def headline_series(metrics: dict[str, float]) -> dict[str, float]: # row carries None where a series was withheld rather than a zero that # would read as a measurement. _BREAKDOWN_COLUMNS = ( - "time_pct", + "pct_of_dataplane", "calls", "wall_ms", "mean_ms", @@ -872,8 +872,8 @@ def breakdown_table( per_op: dict[str, dict[str, float]] = {} for key, value in metrics.items(): parts = key.split("/") - if len(parts) == 4 and parts[:3] == ["step", "time_pct", "by_op"]: - per_op.setdefault(parts[3], {})["time_pct"] = value + if len(parts) == 4 and parts[:3] == ["step", "pct_of_dataplane", "by_op"]: + per_op.setdefault(parts[3], {})["pct_of_dataplane"] = value elif ( len(parts) == 3 and parts[0] == "step" @@ -883,7 +883,7 @@ def breakdown_table( per_op.setdefault(parts[1], {})[parts[2]] = value rows = [ [op, *(stats.get(col) for col in _BREAKDOWN_COLUMNS)] - # By wall time, which orders identically to ``time_pct`` (that is + # By wall time, which orders identically to ``pct_of_dataplane`` (that is # wall time over a common total) and is present even when the # percentages are not -- a table built from a partial metrics dict # still reads worst-first. @@ -1082,7 +1082,7 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: ``VllmGeneration.get_step_metrics`` so trainers stay one line. ``frac_of_step`` is the metric that decides whether optimising the - data plane is worth anything: ``time_pct`` only says where + data plane is worth anything: ``pct_of_dataplane`` only says where data-plane time went, never whether it mattered against compute. """ # Reading the step maxima is what closes the window: the values @@ -1116,7 +1116,7 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: "fields_skipped", 0 ) per_op = _op_step_stats(snap["by_op"], prev.get("by_op", {})) - metrics.update(_time_pct(per_op)) + metrics.update(_pct_of_dataplane(per_op)) for op, row in per_op.items(): for field_name, value in row.items(): metrics[f"step/{op}/{field_name}"] = value diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 1358506ac27..62e5f5e220b 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -1010,8 +1010,8 @@ def test_breakdown_table_rows_by_op_worst_first(): to "what is my bottleneck" should be the top-left of the table.""" metrics = { "step/wall_ms": 100.0, - "step/time_pct/by_op/get": 10.0, - "step/time_pct/by_op/put": 90.0, + "step/pct_of_dataplane/by_op/get": 10.0, + "step/pct_of_dataplane/by_op/put": 90.0, "step/get/calls": 8, "step/get/wall_ms": 10.0, "step/get/max_ms": 2.0, @@ -1027,8 +1027,8 @@ def test_breakdown_table_rows_by_op_worst_first(): assert [r[0] for r in rows] == ["put", "get"], "worst first" assert len(rows) == 2, "only per-op series become rows" assert rows[0][columns.index("wall_ms")] == 90.0 - assert rows[0][columns.index("time_pct")] == pytest.approx(90.0) - assert columns[1] == "time_pct", "the bottleneck reads first" + assert rows[0][columns.index("pct_of_dataplane")] == pytest.approx(90.0) + assert columns[1] == "pct_of_dataplane", "the bottleneck reads first" def test_breakdown_table_leaves_withheld_series_empty(): @@ -1101,7 +1101,7 @@ def test_cluster_per_op_time_is_reported_per_call(): ), "the sum does move with cluster size" columns, rows = breakdown_table(small) - assert "mean_ms" in columns and "time_pct" in columns + assert "mean_ms" in columns and "pct_of_dataplane" in columns def test_each_quantile_waits_for_the_samples_it_needs(): @@ -1263,7 +1263,7 @@ def _busy_client(op_calls): return client -def test_time_pct_names_the_bottleneck_and_says_of_what(): +def test_pct_of_dataplane_names_the_bottleneck_and_says_of_what(): """The question a dashboard has to answer is "which op is expensive", and 32 per-op line charts do not answer it. @@ -1278,10 +1278,14 @@ def test_time_pct_names_the_bottleneck_and_says_of_what(): merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 10.0 ) - by_op = {k: v for k, v in metrics.items() if k.startswith("step/time_pct/by_op/")} + by_op = { + k: v for k, v in metrics.items() if k.startswith("step/pct_of_dataplane/by_op/") + } assert sum(by_op.values()) == pytest.approx(100.0), "percent of one total" - assert max(by_op, key=by_op.__getitem__) == "step/time_pct/by_op/get" - assert by_op["step/time_pct/by_op/get"] == pytest.approx(100 * 900 / 911, rel=0.05) + assert max(by_op, key=by_op.__getitem__) == "step/pct_of_dataplane/by_op/get" + assert by_op["step/pct_of_dataplane/by_op/get"] == pytest.approx( + 100 * 900 / 911, rel=0.05 + ) # the two denominators are different questions and must not agree assert metrics["step/frac_of_step"] == pytest.approx(0.0911, rel=0.1) @@ -1299,7 +1303,7 @@ def test_headline_drops_per_op_detail_but_keeps_the_percentages(): assert len(head) < len(metrics) / 2, f"{len(head)} of {len(metrics)}" assert not [k for k in head if k.split("/")[1:2] in (["get"], ["put"], ["clear"])] - assert "step/time_pct/by_op/get" in head + assert "step/pct_of_dataplane/by_op/get" in head assert "step/wall_ms" in head and "step/frac_of_step" in head # and the detail the table needs survives in the full dict assert breakdown_table(metrics)[1], "table still has rows" @@ -1351,7 +1355,7 @@ def test_breakdown_table_ignores_reserved_namespaces(): "step/put/wall_ms": 9.0, "step/self/overhead_ms": 6.2, "step/hash/mismatches": 0, - "step/time_pct/by_cause/transfer": 40.0, + "step/pct_of_dataplane/by_cause/transfer": 40.0, } ) assert [r[0] for r in rows] == ["put"], rows From 6329f21334545546563378ecfed67daafe588e57 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 15:52:38 -0700 Subject: [PATCH 33/70] refactor(data-plane): spell out percent in the metric names `pct` is an abbreviation nobody has to guess at. Also spells out the cumulative `pct_of_total_ms` in `snapshot()` for the same reason. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 18 +++++------ nemo_rl/data_plane/observability.py | 34 +++++++++++---------- tests/unit/data_plane/test_observability.py | 24 ++++++++------- 3 files changed, 40 insertions(+), 36 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 7e6cf7d6046..337e96ed500 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -467,18 +467,18 @@ cluster percentile). **What gets charted is the bottleneck, not the detail.** Four ops times eight fields is 32 series saying one thing, and a dashboard of 32 lines does not answer "where is my time going". So the emitted series are the -totals and `pct_of_dataplane`, with the per-op detail in a table beside them: +totals and `percent_of_dataplane`, with the per-op detail in a table beside them: | series | what it answers | |---|---| | `step/frac_of_step` | is the data plane worth optimising at all? | -| `step/pct_of_dataplane/by_op/{put,get,clear,register}` | which call is expensive? | -| `step/pct_of_dataplane/by_cause/{fixed_overhead,transfer}` | is that fixed per-request cost, or moving bytes? | +| `step/percent_of_dataplane/by_op/{put,get,clear,register}` | which call is expensive? | +| `step/percent_of_dataplane/by_cause/{fixed_overhead,transfer}` | is that fixed per-request cost, or moving bytes? | | `step/wall_ms`, `step/comm_volume_mb` | how much time and traffic | | `now/bytes_outstanding_mb`, `now/n_processes` | occupancy, fan-out width | | `step/self/{overhead_ms,frac}` | what measuring cost | -**`pct_of_dataplane` is a percentage of data-plane time, not of the step.** The +**`percent_of_dataplane` is a percentage of data-plane time, not of the step.** The denominator is `sum(wall_ms)` over the ops that ran, so `by_op/put = 43` reads "43% of the time spent inside the data plane went to put". Whether that time mattered against compute is the *other* metric: @@ -487,8 +487,8 @@ a workload can be 43% put and still not be worth touching. ``` step/frac_of_step 0.87 the data plane is 87% of the step, so it matters -step/pct_of_dataplane/by_op/put 43.0 and within it, put is the largest piece -step/pct_of_dataplane/by_cause/... +step/percent_of_dataplane/by_op/put 43.0 and within it, put is the largest piece +step/percent_of_dataplane/by_cause/... fixed_overhead 53.1 that time is mostly per-request cost, transfer 17.6 not bandwidth -- batch, don't tune the wire ``` @@ -507,9 +507,9 @@ and the wrong one for "what blocked the step". **A per-op breakdown table** carries the detail, under `data_plane/{cluster,driver}/breakdown` — one row per op, ordered by -`pct_of_dataplane` so the bottleneck is the first line read: +`percent_of_dataplane` so the bottleneck is the first line read: -| op | pct_of_dataplane | calls | wall_ms | mean_ms | max_ms | p50_ms | p90_ms | overhead_ms | transfer_ms | mb | +| op | percent_of_dataplane | calls | wall_ms | mean_ms | max_ms | p50_ms | p90_ms | overhead_ms | transfer_ms | mb | |---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| | put | 43.0 | 2 | 53.9 | 26.9 | 29.4 | — | — | 19.4 | 6.39 | 1.32 | | get | 30.9 | 2 | 38.7 | 19.4 | 21.5 | — | — | 13.8 | 4.63 | 1.03 | @@ -530,7 +530,7 @@ while the wall clock was 279. Dividing by the process count only trades one arbitrary denominator for another. Per call is invariant to both DP degree and batch size: the same workload at 8 and at 32 ranks reports 11.16 and 11.06 ms while `wall_ms` quadruples. Use `mean_ms` to compare runs and -cluster sizes, `pct_of_dataplane` to attribute cost across ops within one step. +cluster sizes, `percent_of_dataplane` to attribute cost across ops within one step. A stack of line charts answers "how did put's wall time trend"; this answers "where did the step go", which is a table. Cells are empty rather diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 6b0e2c8f3bc..bb6a84406f6 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -497,7 +497,7 @@ def _op_step_stats( return out -def _pct_of_dataplane(per_op: dict[str, dict[str, float]]) -> dict[str, float]: +def _percent_of_dataplane(per_op: dict[str, dict[str, float]]) -> dict[str, float]: """Where this step's data-plane time went, in percent. The name carries the denominator because that is the one thing a reader @@ -529,14 +529,14 @@ def _pct_of_dataplane(per_op: dict[str, dict[str, float]]) -> dict[str, float]: per_op: Per-op step detail from :func:`_op_step_stats`. Returns: - ``step/pct_of_dataplane/by_op/{op}`` and ``step/pct_of_dataplane/by_cause/{cause}``, + ``step/percent_of_dataplane/by_op/{op}`` and ``step/percent_of_dataplane/by_cause/{cause}``, each in percent. Empty when no op ran. """ total = sum(r["wall_ms"] for r in per_op.values()) if total <= 0: return {} - pct = { - f"step/pct_of_dataplane/by_op/{op}": 100.0 * r["wall_ms"] / total + percent = { + f"step/percent_of_dataplane/by_op/{op}": 100.0 * r["wall_ms"] / total for op, r in per_op.items() } for cause, field_name in ( @@ -547,8 +547,10 @@ def _pct_of_dataplane(per_op: dict[str, dict[str, float]]) -> dict[str, float]: r[field_name] * r["calls"] for r in per_op.values() if field_name in r ) if attributed > 0: - pct[f"step/pct_of_dataplane/by_cause/{cause}"] = 100.0 * attributed / total - return pct + percent[f"step/percent_of_dataplane/by_cause/{cause}"] = ( + 100.0 * attributed / total + ) + return percent def _latency_split( @@ -617,7 +619,7 @@ def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: stats["mb_per_s"] = ( (stats["n_bytes"] / 1e6) / (wall_ms / 1e3) if wall_ms else 0.0 ) - stats["pct_of_total_ms"] = ( + stats["percent_of_total_ms"] = ( 100.0 * wall_ms / total_wall_ms if total_wall_ms else 0.0 ) stats["fit"] = fit_latency_bandwidth(stats) @@ -778,7 +780,7 @@ def cluster_step_metrics( } ) per_op = _op_step_stats(merged["by_op"], prev.get("by_op", {})) - metrics.update(_pct_of_dataplane(per_op)) + metrics.update(_percent_of_dataplane(per_op)) for op, row in per_op.items(): for field_name, value in row.items(): metrics[f"step/{op}/{field_name}"] = value @@ -800,11 +802,11 @@ def cluster_step_metrics( "step/self/overhead_ms", "step/self/frac", ) -_HEADLINE_PREFIXES = ("step/pct_of_dataplane/", "step/hash/", "step/self/") +_HEADLINE_PREFIXES = ("step/percent_of_dataplane/", "step/hash/", "step/self/") # ``step//`` is the per-op shape, so these reserved middles must # not be mistaken for op tags -- ``step/self/overhead_ms`` was becoming a # "self" row in the breakdown table beside put and get. -_RESERVED_NAMESPACES = frozenset({"pct_of_dataplane", "hash", "self"}) +_RESERVED_NAMESPACES = frozenset({"percent_of_dataplane", "hash", "self"}) def headline_series(metrics: dict[str, float]) -> dict[str, float]: @@ -831,7 +833,7 @@ def headline_series(metrics: dict[str, float]) -> dict[str, float]: # row carries None where a series was withheld rather than a zero that # would read as a measurement. _BREAKDOWN_COLUMNS = ( - "pct_of_dataplane", + "percent_of_dataplane", "calls", "wall_ms", "mean_ms", @@ -872,8 +874,8 @@ def breakdown_table( per_op: dict[str, dict[str, float]] = {} for key, value in metrics.items(): parts = key.split("/") - if len(parts) == 4 and parts[:3] == ["step", "pct_of_dataplane", "by_op"]: - per_op.setdefault(parts[3], {})["pct_of_dataplane"] = value + if len(parts) == 4 and parts[:3] == ["step", "percent_of_dataplane", "by_op"]: + per_op.setdefault(parts[3], {})["percent_of_dataplane"] = value elif ( len(parts) == 3 and parts[0] == "step" @@ -883,7 +885,7 @@ def breakdown_table( per_op.setdefault(parts[1], {})[parts[2]] = value rows = [ [op, *(stats.get(col) for col in _BREAKDOWN_COLUMNS)] - # By wall time, which orders identically to ``pct_of_dataplane`` (that is + # By wall time, which orders identically to ``percent_of_dataplane`` (that is # wall time over a common total) and is present even when the # percentages are not -- a table built from a partial metrics dict # still reads worst-first. @@ -1082,7 +1084,7 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: ``VllmGeneration.get_step_metrics`` so trainers stay one line. ``frac_of_step`` is the metric that decides whether optimising the - data plane is worth anything: ``pct_of_dataplane`` only says where + data plane is worth anything: ``percent_of_dataplane`` only says where data-plane time went, never whether it mattered against compute. """ # Reading the step maxima is what closes the window: the values @@ -1116,7 +1118,7 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: "fields_skipped", 0 ) per_op = _op_step_stats(snap["by_op"], prev.get("by_op", {})) - metrics.update(_pct_of_dataplane(per_op)) + metrics.update(_percent_of_dataplane(per_op)) for op, row in per_op.items(): for field_name, value in row.items(): metrics[f"step/{op}/{field_name}"] = value diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 62e5f5e220b..47f5a419265 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -1010,8 +1010,8 @@ def test_breakdown_table_rows_by_op_worst_first(): to "what is my bottleneck" should be the top-left of the table.""" metrics = { "step/wall_ms": 100.0, - "step/pct_of_dataplane/by_op/get": 10.0, - "step/pct_of_dataplane/by_op/put": 90.0, + "step/percent_of_dataplane/by_op/get": 10.0, + "step/percent_of_dataplane/by_op/put": 90.0, "step/get/calls": 8, "step/get/wall_ms": 10.0, "step/get/max_ms": 2.0, @@ -1027,8 +1027,8 @@ def test_breakdown_table_rows_by_op_worst_first(): assert [r[0] for r in rows] == ["put", "get"], "worst first" assert len(rows) == 2, "only per-op series become rows" assert rows[0][columns.index("wall_ms")] == 90.0 - assert rows[0][columns.index("pct_of_dataplane")] == pytest.approx(90.0) - assert columns[1] == "pct_of_dataplane", "the bottleneck reads first" + assert rows[0][columns.index("percent_of_dataplane")] == pytest.approx(90.0) + assert columns[1] == "percent_of_dataplane", "the bottleneck reads first" def test_breakdown_table_leaves_withheld_series_empty(): @@ -1101,7 +1101,7 @@ def test_cluster_per_op_time_is_reported_per_call(): ), "the sum does move with cluster size" columns, rows = breakdown_table(small) - assert "mean_ms" in columns and "pct_of_dataplane" in columns + assert "mean_ms" in columns and "percent_of_dataplane" in columns def test_each_quantile_waits_for_the_samples_it_needs(): @@ -1263,7 +1263,7 @@ def _busy_client(op_calls): return client -def test_pct_of_dataplane_names_the_bottleneck_and_says_of_what(): +def test_percent_of_dataplane_names_the_bottleneck_and_says_of_what(): """The question a dashboard has to answer is "which op is expensive", and 32 per-op line charts do not answer it. @@ -1279,11 +1279,13 @@ def test_pct_of_dataplane_names_the_bottleneck_and_says_of_what(): ) by_op = { - k: v for k, v in metrics.items() if k.startswith("step/pct_of_dataplane/by_op/") + k: v + for k, v in metrics.items() + if k.startswith("step/percent_of_dataplane/by_op/") } assert sum(by_op.values()) == pytest.approx(100.0), "percent of one total" - assert max(by_op, key=by_op.__getitem__) == "step/pct_of_dataplane/by_op/get" - assert by_op["step/pct_of_dataplane/by_op/get"] == pytest.approx( + assert max(by_op, key=by_op.__getitem__) == "step/percent_of_dataplane/by_op/get" + assert by_op["step/percent_of_dataplane/by_op/get"] == pytest.approx( 100 * 900 / 911, rel=0.05 ) # the two denominators are different questions and must not agree @@ -1303,7 +1305,7 @@ def test_headline_drops_per_op_detail_but_keeps_the_percentages(): assert len(head) < len(metrics) / 2, f"{len(head)} of {len(metrics)}" assert not [k for k in head if k.split("/")[1:2] in (["get"], ["put"], ["clear"])] - assert "step/pct_of_dataplane/by_op/get" in head + assert "step/percent_of_dataplane/by_op/get" in head assert "step/wall_ms" in head and "step/frac_of_step" in head # and the detail the table needs survives in the full dict assert breakdown_table(metrics)[1], "table still has rows" @@ -1355,7 +1357,7 @@ def test_breakdown_table_ignores_reserved_namespaces(): "step/put/wall_ms": 9.0, "step/self/overhead_ms": 6.2, "step/hash/mismatches": 0, - "step/pct_of_dataplane/by_cause/transfer": 40.0, + "step/percent_of_dataplane/by_cause/transfer": 40.0, } ) assert [r[0] for r in rows] == ["put"], rows From f4fa1f5c7ee3471fa542b21c13c1d2b022ad5091 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 15:53:44 -0700 Subject: [PATCH 34/70] docs(data-plane): the transfer term IS the bandwidth term The README's worked example annotated `by_cause/transfer` with "not bandwidth" -- the tail of a sentence about `fixed_overhead` that had been split across two lines, so it landed on the wrong number and said the opposite of the truth. `transfer` is bytes divided by the fitted bandwidth; `fixed_overhead` is the per-request constant. Each term now says what it measures on its own line, with numbers from a live TransferQueue run rather than from a harness whose step was nothing but data-plane calls. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 337e96ed500..eab58db5ba3 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -485,20 +485,29 @@ put". Whether that time mattered against compute is the *other* metric: `frac_of_step` divides by the step's own wall clock. Read them together — a workload can be 43% put and still not be worth touching. +Each `by_cause` term names what it measures. `fixed_overhead` is the fitted +per-request constant — the cost of making a request at all, independent of +its size. `transfer` **is** the bandwidth term: bytes divided by the fitted +bandwidth. Reading a real TransferQueue step: + ``` -step/frac_of_step 0.87 the data plane is 87% of the step, so it matters -step/percent_of_dataplane/by_op/put 43.0 and within it, put is the largest piece -step/percent_of_dataplane/by_cause/... - fixed_overhead 53.1 that time is mostly per-request cost, - transfer 17.6 not bandwidth -- batch, don't tune the wire +step/frac_of_step 0.074 the data plane is 7% of the step +step/percent_of_dataplane/by_op/put 42.1 within it, put is the largest op +step/percent_of_dataplane/by_cause/fixed_overhead 52.3 half the time is per-request cost +step/percent_of_dataplane/by_cause/transfer 20.7 a fifth is bandwidth ``` +Overhead beats bandwidth roughly 2:1 here, so this workload is +overhead-dominated: batching into fewer, larger requests buys more than a +faster wire. Had `transfer` been the larger of the two, the conclusion +would invert. + Two decompositions of one total, because either alone leaves the next question unanswered. `by_op` sums to 100 by construction. `by_cause` sums to *at most* 100: only ops with an identifiable affine fit can be split, so -the remainder (here 29%, the `register` and `clear` calls that move no -bytes) is time that could not be attributed rather than time that did not -happen. +the remainder (27% above — the `register` and `clear` calls, which move +no bytes and so have no bandwidth term to fit) is time that could not be +attributed rather than time that did not happen. On the cluster path `wall_ms` is summed over processes that ran concurrently, so these are percentages of aggregate **process-time**, not of From f03ca4bb510ea5879ce664f2d97fa45b9639869a Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 16:44:07 -0700 Subject: [PATCH 35/70] fix(data-plane): report hash counters and measuring cost in both scopes Listing the emitted series showed the two scopes were not the same shape, and one of the gaps mattered. `step/hash/*` was emitted only by the driver path, but `_log_data_plane_metrics` prefers the cluster path whenever the fan-out reaches more than one process -- which is every real run. So with `verify_tensor_hash` on, `mismatches` never reached the logger: the guard ran, counted, and logged ERROR lines, while the series anyone would chart or alert on was absent. A guard whose findings are not reported is not a guard. Both paths now go through one `_hash_deltas`, which emits nothing when the guard never ran -- five always-zero series would read as "checked, nothing wrong" rather than "not checked". `step/self/*` was cluster-only for the same reason, so the single-process fallback silently lacked the number that says what observability cost. `step/bytes_written_mb` / `step/bytes_read_mb` went the other way: computed on both paths, dropped by `headline_series`, absent from the table, read by nobody. Removed -- per-op `mb` already splits the same traffic finer, since put's is the write volume and get's is the read volume. Verified: 184 unit tests, 48/48 metric audit checks, 23/23 hash-guard checks. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 1 + nemo_rl/data_plane/observability.py | 72 +++++++++++++++------ tests/unit/data_plane/test_observability.py | 66 +++++++++++++++++++ 3 files changed, 121 insertions(+), 18 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index eab58db5ba3..c49d3ed3952 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -477,6 +477,7 @@ totals and `percent_of_dataplane`, with the per-op detail in a table beside them | `step/wall_ms`, `step/comm_volume_mb` | how much time and traffic | | `now/bytes_outstanding_mb`, `now/n_processes` | occupancy, fan-out width | | `step/self/{overhead_ms,frac}` | what measuring cost | +| `step/hash/*` | only with `verify_tensor_hash` on | **`percent_of_dataplane` is a percentage of data-plane time, not of the step.** The denominator is `sum(wall_ms)` over the ops that ran, so diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index bb6a84406f6..b1f4dce5e22 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -431,11 +431,17 @@ def fit_latency_bandwidth(s: dict[str, Any]) -> dict[str, Any]: def _step_deltas(snap: dict[str, Any], prev: dict[str, Any]) -> dict[str, float]: - """The five series both step-metric paths report, identically. + """The three series both step-metric paths report, identically. Shared so the single-process and cluster views cannot drift on series names -- which is the whole point of the ``step/``/``now/`` convention they publish under. + + Write and read volume are deliberately not here. They were computed and + then dropped by :func:`headline_series`, charted by nobody, while the + breakdown table already carries per-op ``mb`` -- put's is the write + volume and get's is the read volume, split finer than a global pair + would be. """ return { "step/wall_ms": snap["total_wall_ms"] - prev.get("total_wall_ms", 0.0), @@ -443,9 +449,6 @@ def _step_deltas(snap: dict[str, Any], prev: dict[str, Any]) -> dict[str, float] snap["comm_volume_bytes"] - prev.get("comm_volume_bytes", 0) ) / 1e6, - "step/bytes_written_mb": (snap["bytes_written"] - prev.get("bytes_written", 0)) - / 1e6, - "step/bytes_read_mb": (snap["bytes_read"] - prev.get("bytes_read", 0)) / 1e6, "now/bytes_outstanding_mb": snap["bytes_outstanding"] / 1e6, } @@ -497,6 +500,40 @@ def _op_step_stats( return out +def _hash_deltas(hv: dict[str, int], prev_hv: dict[str, int]) -> dict[str, float]: + """This step's hash-verification counters, or nothing if the guard is off. + + Shared by both step-metric paths. It was emitted only on the driver + path, but ``_log_data_plane_metrics`` prefers the cluster path whenever + the fan-out reaches more than one process -- which is every real run -- + so with ``verify_tensor_hash`` on, ``mismatches`` never reached the + logger. A guard whose findings are not reported is not a guard. + + ``fields_skipped`` is here for the same reason it exists at all: a guard + that quietly stops covering a field still reports zero mismatches, so + the abstention count has to be visible beside the finding count. + + Args: + hv: This step's cumulative ``hash_verify`` block. + prev_hv: The previous step's, for differencing. + + Returns: + ``step/hash/{counter}`` deltas, or ``{}`` when the guard never ran. + """ + if not hv or not hv.get("rows_recorded"): + return {} + return { + f"step/hash/{name}": hv[name] - prev_hv.get(name, 0) + for name in ( + "rows_checked", + "rows_recorded", + "rows_unverified", + "mismatches", + "fields_skipped", + ) + } + + def _percent_of_dataplane(per_op: dict[str, dict[str, float]]) -> dict[str, float]: """Where this step's data-plane time went, in percent. @@ -779,6 +816,9 @@ def cluster_step_metrics( "step/self/frac": overhead_ms / wall_ms if wall_ms > 0 else 0.0, } ) + metrics.update( + _hash_deltas(merged.get("hash_verify") or {}, prev.get("hash_verify") or {}) + ) per_op = _op_step_stats(merged["by_op"], prev.get("by_op", {})) metrics.update(_percent_of_dataplane(per_op)) for op, row in per_op.items(): @@ -1103,20 +1143,16 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: metrics["step/frac_of_step"] = ( (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0 ) - if self._verify_tensor_hash: - hv, prev_hv = snap["hash_verify"], prev.get("hash_verify", {}) - for name in ("rows_checked", "rows_recorded", "rows_unverified"): - metrics[f"step/hash/{name}"] = hv[name] - prev_hv.get(name, 0) - metrics["step/hash/mismatches"] = hv["mismatches"] - prev_hv.get( - "mismatches", 0 - ) - # Logged because a guard that quietly stops covering a field is - # worse than no guard: it reports 0 mismatches and reads as - # clean. A step where this climbs is a step where something - # stopped being checked. - metrics["step/hash/fields_skipped"] = hv["fields_skipped"] - prev_hv.get( - "fields_skipped", 0 - ) + metrics.update( + _hash_deltas(snap.get("hash_verify") or {}, prev.get("hash_verify") or {}) + ) + # The same bill the cluster path reports, under the same name: this + # process's wrapper time, minus what the inner client was doing. + # There is no fan-out to add here -- a single process gathers + # nothing -- so this is the whole of it. + self_ms = snap["self_ms"] - prev.get("self_ms", 0.0) + metrics["step/self/overhead_ms"] = self_ms + metrics["step/self/frac"] = self_ms / wall_ms if wall_ms > 0 else 0.0 per_op = _op_step_stats(snap["by_op"], prev.get("by_op", {})) metrics.update(_percent_of_dataplane(per_op)) for op, row in per_op.items(): diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 47f5a419265..639508106eb 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -1362,3 +1362,69 @@ def test_breakdown_table_ignores_reserved_namespaces(): ) assert [r[0] for r in rows] == ["put"], rows assert columns[0] == "op" + + +def test_hash_counters_reach_both_scopes(): + """``_log_data_plane_metrics`` prefers the cluster path whenever the + fan-out reaches more than one process -- every real run -- and the + cluster path emitted no hash counters at all. With verify_tensor_hash + on, ``mismatches`` never reached the logger: a guard whose findings are + not reported is not a guard.""" + client = _hash_client(_CorruptingClient(field="ids", row=2)) + ids = [f"u{i}" for i in range(4)] + client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids", "lp"]) + + merged = merge_snapshots([client.snapshot(reset_step_window=True)]) + cluster = cluster_step_metrics(merged, {}, 1.0) + assert cluster["step/hash/mismatches"] == 1 + assert "step/hash/fields_skipped" in cluster, "abstentions visible too" + assert headline_series(cluster)["step/hash/mismatches"] == 1, "and charted" + client.close() + + +def test_hash_counters_absent_when_the_guard_is_off(): + """Five always-zero series on every run that never asked for the guard + would read as "checked, nothing wrong" rather than "not checked".""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client._emit("put", "p", 1, 1_000, monotonic() - 0.005, "ok") + merged = merge_snapshots([client.snapshot(reset_step_window=True)]) + + assert not [k for k in cluster_step_metrics(merged, {}, 1.0) if "hash" in k] + assert not [k for k in client.get_step_metrics(1.0) if "hash" in k] + client.close() + + +def test_measuring_cost_is_reported_in_both_scopes(): + """``step/self/*`` was cluster-only, so the single-process fallback + silently lacked the one number that says what observability cost.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client.register_partition( + partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["t"] + ) + client.put_samples( + sample_ids=["a"], + partition_id="p", + fields=TensorDict({"x": torch.zeros(1, 512)}, batch_size=[1]), + ) + driver = client.get_step_metrics(1.0) + + assert driver["step/self/overhead_ms"] > 0, "measuring is never free" + assert "step/self/frac" in driver + assert "step/self/overhead_ms" in headline_series(driver) + client.close() + + +def test_write_and_read_volume_are_not_computed_for_nobody(): + """They were dropped by headline_series, charted by nobody, and absent + from the table -- while per-op ``mb`` already splits the same traffic + finer (put's is the write volume, get's the read volume).""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client._emit("put", "p", 1, 4_000, monotonic() - 0.005, "ok") + metrics = client.get_step_metrics(1.0) + + assert "step/bytes_written_mb" not in metrics + assert "step/bytes_read_mb" not in metrics + assert metrics["step/comm_volume_mb"] > 0, "the total is still reported" + assert metrics["step/put/mb"] == pytest.approx(0.004), "and split per op" + client.close() From a19317de7dd4f296d5eab9e6667597fec3d98455 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 22 Aug 2026 22:36:36 -0700 Subject: [PATCH 36/70] fix(data-plane): log before the step commits; track per-op transfer volume A real GRPO run on one GB200 node found that none of these series ever reached wandb. `_log_data_plane_metrics` ran *after* the `step_finished=True` log that commits the wandb step, and anything logged against a committed step is dropped -- so every value was computed, printed to stdout, and silently discarded. Reading a finished run back out of the wandb API is what caught it: 85 distinct keys, zero `data_plane/*`. Every synthetic harness missed it because they call `log_metrics` with no preceding commit. Now logged before the commit, with a test that asserts the source order, since the drop happens inside wandb where a fake logger cannot see it. Also adds `step/volume_mb/by_op/{op}`. `comm_volume_mb` alone hides the asymmetry that matters: on that run get moved 20.8 MB against put's 2.7 MB, because every DP rank fetches its shard once for the logprob pass and again for the train pass. The total cannot say that, and the direction is what tells you whether to look at reads or writes. Both are transfers rather than data size -- a byte written and later read counts on both sides, and each process's transfers are summed -- which the README now states, along with the scope caveat that makes put read small: the rollout actor builds its own client, is not on the worker group, and so its `kv_first_write` of the whole batch is in neither figure. Verified on a live 5-process run: frac_of_step 1.3% of a 15.4 s step, get 66% of data-plane time, and the breakdown table logged per step. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- ...-instruct-1n8g-fsdp2tp1-tq_simple-obs.yaml | 9 +++ nemo_rl/algorithms/grpo_sync.py | 5 +- nemo_rl/data_plane/README.md | 15 ++++ nemo_rl/data_plane/observability.py | 32 +++++++- tests/unit/data_plane/test_observability.py | 77 +++++++++++++++++++ 5 files changed, 135 insertions(+), 3 deletions(-) create mode 100644 examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1-tq_simple-obs.yaml diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1-tq_simple-obs.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1-tq_simple-obs.yaml new file mode 100644 index 00000000000..51139d56be6 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1-tq_simple-obs.yaml @@ -0,0 +1,9 @@ +# Data-plane observability e2e: real GRPO traffic through TransferQueue with +# per-op metrics and the wire-in/wire-out hash guard both on. Short run -- +# the point is the data_plane/* series, not convergence. +defaults: grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3.yaml +data_plane: + enabled: true + observability: + enabled: true + verify_tensor_hash: true diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index 9a167e8a033..da16de450a3 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -1453,13 +1453,16 @@ def grpo_train_sync( logger.log_metrics( performance_metrics, total_steps + 1, prefix="performance" ) + # Before the step_finished=True log below, which commits the step: + # anything logged against a committed step is dropped by wandb, so + # these series were computed, printed, and silently discarded. + _log_data_plane_metrics(policy, logger, total_steps + 1, total_time) logger.log_metrics( timing_metrics, total_steps + 1, prefix="timing/train", step_finished=True, ) - _log_data_plane_metrics(policy, logger, total_steps + 1, total_time) dynamic_sampling_num_gen_batches = 0 diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index c49d3ed3952..0d46199d44e 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -475,6 +475,7 @@ totals and `percent_of_dataplane`, with the per-op detail in a table beside them | `step/percent_of_dataplane/by_op/{put,get,clear,register}` | which call is expensive? | | `step/percent_of_dataplane/by_cause/{fixed_overhead,transfer}` | is that fixed per-request cost, or moving bytes? | | `step/wall_ms`, `step/comm_volume_mb` | how much time and traffic | +| `step/volume_mb/by_op/{get,put}` | which direction that traffic went | | `now/bytes_outstanding_mb`, `now/n_processes` | occupancy, fan-out width | | `step/self/{overhead_ms,frac}` | what measuring cost | | `step/hash/*` | only with `verify_tensor_hash` on | @@ -510,6 +511,20 @@ the remainder (27% above — the `register` and `clear` calls, which move no bytes and so have no bandwidth term to fit) is time that could not be attributed rather than time that did not happen. +`volume_mb` counts *transfers*, not data size, and two things follow from +that. A byte written and later read is counted on both sides. And every +reporting process is summed, so four ranks each fetching their own shard +count four times. Both are correct for "what crossed the wire" -- on a real +step get moved 20.8 MB against put's 2.7 MB, because every DP rank fetches +its shard once for the logprob pass and again for the train pass. Neither +is correct for "how big was the batch", which these series cannot answer. + +**The rollout actor is not in the fan-out**, so `kv_first_write` -- the +write of the entire rollout batch, and the largest write in the step -- is +absent from `volume_mb/by_op/put` and from `comm_volume_mb`. That is why +put reads small next to get. Read the write side as "what the driver and +policy workers wrote", not as the step's write traffic. + On the cluster path `wall_ms` is summed over processes that ran concurrently, so these are percentages of aggregate **process-time**, not of elapsed time. That is the right denominator for "what should I optimise" diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index b1f4dce5e22..28b3b90bd3c 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -534,6 +534,27 @@ def _hash_deltas(hv: dict[str, int], prev_hv: dict[str, int]) -> dict[str, float } +def _volume_mb(per_op: dict[str, dict[str, float]]) -> dict[str, float]: + """Bytes each op moved this step, in MB, per op that moved any. + + ``comm_volume_mb`` is the total and hides the asymmetry that matters: + on a real step ``get`` moved 20.8 MB against ``put``'s 2.7 MB, because + every DP rank fetches its shard once for the logprob pass and again for + the train pass. Those are separate transfers over the wire, not an + accounting artifact, and the same is true of summing across processes -- + each rank pulls its own shard. + + Ops that carry no payload (``register``, ``clear``) are omitted rather + than reported as zero, matching how the percentages treat an op that + did not run. + """ + return { + f"step/volume_mb/by_op/{op}": row["mb"] + for op, row in per_op.items() + if row["mb"] > 0 + } + + def _percent_of_dataplane(per_op: dict[str, dict[str, float]]) -> dict[str, float]: """Where this step's data-plane time went, in percent. @@ -821,6 +842,7 @@ def cluster_step_metrics( ) per_op = _op_step_stats(merged["by_op"], prev.get("by_op", {})) metrics.update(_percent_of_dataplane(per_op)) + metrics.update(_volume_mb(per_op)) for op, row in per_op.items(): for field_name, value in row.items(): metrics[f"step/{op}/{field_name}"] = value @@ -842,11 +864,16 @@ def cluster_step_metrics( "step/self/overhead_ms", "step/self/frac", ) -_HEADLINE_PREFIXES = ("step/percent_of_dataplane/", "step/hash/", "step/self/") +_HEADLINE_PREFIXES = ( + "step/percent_of_dataplane/", + "step/volume_mb/", + "step/hash/", + "step/self/", +) # ``step//`` is the per-op shape, so these reserved middles must # not be mistaken for op tags -- ``step/self/overhead_ms`` was becoming a # "self" row in the breakdown table beside put and get. -_RESERVED_NAMESPACES = frozenset({"percent_of_dataplane", "hash", "self"}) +_RESERVED_NAMESPACES = frozenset({"percent_of_dataplane", "volume_mb", "hash", "self"}) def headline_series(metrics: dict[str, float]) -> dict[str, float]: @@ -1155,6 +1182,7 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: metrics["step/self/frac"] = self_ms / wall_ms if wall_ms > 0 else 0.0 per_op = _op_step_stats(snap["by_op"], prev.get("by_op", {})) metrics.update(_percent_of_dataplane(per_op)) + metrics.update(_volume_mb(per_op)) for op, row in per_op.items(): for field_name, value in row.items(): metrics[f"step/{op}/{field_name}"] = value diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 639508106eb..1537b1d10ca 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -1428,3 +1428,80 @@ def test_write_and_read_volume_are_not_computed_for_nobody(): assert metrics["step/comm_volume_mb"] > 0, "the total is still reported" assert metrics["step/put/mb"] == pytest.approx(0.004), "and split per op" client.close() + + +def test_data_plane_is_logged_before_the_step_is_committed(): + """``logger.log_metrics(..., step_finished=True)`` commits the wandb step, + and anything logged against a committed step is dropped. + + The data-plane call used to sit *after* that commit in ``grpo_train_sync``, + so every series was computed, printed to stdout, and silently discarded -- + invisible to a fake logger, and only caught by reading a real run back out + of the wandb API. This asserts the source order rather than the behaviour, + because the drop happens inside wandb. + """ + import pathlib + + import nemo_rl + + # Read the file rather than import it: ``grpo_sync`` pulls the training + # stack, and the rest of this suite runs without it. + # + # Comments stripped: the explanatory comment above the call site names + # ``step_finished=True`` too, and matching that would pass on any ordering. + source = ( + pathlib.Path(nemo_rl.__file__).parent / "algorithms" / "grpo_sync.py" + ).read_text() + body = "\n".join( + line + for line in source[source.index("def grpo_train_sync") :].splitlines() + if not line.lstrip().startswith("#") + ) + dp_call = body.index("_log_data_plane_metrics(policy, logger") + commit = body.index("step_finished=True") + assert dp_call < commit, ( + "_log_data_plane_metrics must run before the step_finished=True log; " + "wandb drops anything logged against an already-committed step" + ) + + +def test_per_op_volume_is_charted_and_sums_to_comm_volume(): + """``comm_volume_mb`` alone hides which direction the traffic went. On a + real step get moved 20.8 MB against put's 2.7 MB -- every DP rank fetches + its shard for the logprob pass and again for the train pass -- and a + single total cannot say that.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + now = monotonic() + for _ in range(6): + client._emit("get", "p", 1, 3_000_000, now - 5 / 1e3, "ok") + for _ in range(2): + client._emit("put", "p", 1, 1_000_000, now - 5 / 1e3, "ok") + client._emit("clear", "p", 1, 0, now - 1 / 1e3, "ok") + + metrics = cluster_step_metrics( + merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 1.0 + ) + head = headline_series(metrics) + + assert head["step/volume_mb/by_op/get"] == pytest.approx(18.0) + assert head["step/volume_mb/by_op/put"] == pytest.approx(2.0) + assert "step/volume_mb/by_op/clear" not in head, "no payload, not a zero" + # the parts account for the whole + per_op = sum(v for k, v in head.items() if k.startswith("step/volume_mb/")) + assert per_op == pytest.approx(head["step/comm_volume_mb"]) + client.close() + + +def test_volume_namespace_does_not_become_a_breakdown_row(): + """``step/volume_mb/by_op/get`` must feed the get row, not invent a + ``volume_mb`` op beside put and get.""" + columns, rows = breakdown_table( + { + "step/get/calls": 3, + "step/get/wall_ms": 9.0, + "step/get/mb": 18.0, + "step/volume_mb/by_op/get": 18.0, + } + ) + assert [r[0] for r in rows] == ["get"], rows + assert rows[0][columns.index("mb")] == 18.0 From 70f888aba68f5304b7fa0be16c930054ab8ed4c9 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sun, 23 Aug 2026 00:19:23 -0700 Subject: [PATCH 37/70] fix(data-plane): stop the hash guard failing every per-row field A real 5-process GRPO run reported 3584 mismatches per step, identically, on a job that trained fine. All of them were the guard's own doing: zero row-level mismatches, every one from the "written uniform, read ragged" branch. Root cause is the per-field scheme record. Once per-row fields became keys in that mapping, handing the whole mapping to ``_row_fingerprints`` as ``batch_scoped_fields`` forced *every* field onto the batch-scoped path on read, where the comparison then failed because the recorded value is a width rather than a row count. Only the scoped names go through now. That branch was also too harsh on its own terms. Writing a shard with uniform rows and reading it back inside a batch whose other rows differ in length is the normal shape of the pipeline, not a corruption, so failing the whole field cannot be right. The row *lengths* remain comparable, and a row that changed length is a real divergence -- so the width is recorded and checked per row, and the content it can no longer compare counts as an abstention in ``fields_skipped`` rather than being reported as wrong. Verified on a live 5-process run: 1536 rows checked per step, 0 mismatches and 0 abstentions over six consecutive steps. Plus 188 unit tests, 23/23 guard checks and 48/48 audit checks, with a regression test reproducing the false positive directly. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 77 +++++++++++++++------ tests/unit/data_plane/test_observability.py | 40 +++++++++-- 2 files changed, 88 insertions(+), 29 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 28b3b90bd3c..7d2222e6547 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -117,10 +117,14 @@ class _FieldDigest(NamedTuple): buffer, so they only reconcile against a read of that same batch. A shard of it computes a different buffer digest and must be reported unverified rather than as a mismatch. + + ``row_lens`` is how long each row was, which is the one thing still + comparable when the read cannot reproduce the write's scheme. """ per_row: list[int] batch_scoped: bool + row_lens: tuple[int, ...] # Same-width signed integer for each tensor element size, used to bitcast a @@ -1344,13 +1348,16 @@ def _row_fingerprints( out[name] = _FieldDigest( (torch.hash_tensor(flat, dim=1) ^ salt).tolist(), batch_scoped=False, + row_lens=(flat.shape[1],) * n_rows, ) continue buffer = v.values() if v.is_nested else v flat_buffer = _as_int_view(buffer.reshape(1, -1)) buffer_digest = torch.hash_tensor(flat_buffer, dim=1).tolist()[0] ^ salt out[name] = _FieldDigest( - [buffer_digest ^ length for length in lengths], batch_scoped=True + [buffer_digest ^ length for length in lengths], + batch_scoped=True, + row_lens=tuple(lengths), ) return out @@ -1373,9 +1380,12 @@ def _record_hashes( scheme = self._batch_scope.setdefault(partition_id, {}) for name, digest in digests.items(): if digest.batch_scoped: - scheme[name] = len(sample_ids) + scheme[name] = ("scoped", len(sample_ids)) else: - scheme.pop(name, None) + # Width, not row count: a per-row scheme means the rows were + # uniform, so one number describes them all and lets a later + # ragged read still check whether any row changed length. + scheme[name] = ("rows", digest.row_lens[0] if digest.row_lens else 0) self._stats.hash_verify.rows_recorded += len(sample_ids) def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: @@ -1383,7 +1393,11 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N if not isinstance(out, TensorDict): return scheme = self._batch_scope.get(partition_id, {}) - digests = self._row_fingerprints(out, sample_ids, scheme) + # Only the batch-scoped names: the scheme also records per-row fields + # now, and handing the whole mapping over forced every field onto the + # scoped path. + scoped_names = {n for n, how in scheme.items() if how[0] == "scoped"} + digests = self._row_fingerprints(out, sample_ids, scoped_names) if not digests: return partition_hashes = self._hash_by_partition.get(partition_id, {}) @@ -1395,27 +1409,30 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N # exact shape of the bug that made this check pass while covering # nothing: a field stops being compared and the report still reads # clean. - comparable = {} - relaid_out = [] + comparable: dict[str, _FieldDigest] = {} + length_only: dict[str, tuple[_FieldDigest, int]] = {} for name, digest in digests.items(): - if not digest.batch_scoped or scheme.get(name) == len(sample_ids): + recorded = scheme.get(name) + if not digest.batch_scoped: comparable[name] = digest - elif name in scheme: - stats.fields_skipped += 1 # a shard of a batch-scoped put + elif recorded is None: + continue # this process never wrote the field + elif recorded[0] == "scoped": + if recorded[1] == len(sample_ids): + comparable[name] = digest + else: + stats.fields_skipped += 1 # a shard of a batch-scoped put else: - # Put reduced this field per row, so its rows were uniform; - # they came back ragged. The row lengths themselves changed — - # a real divergence, not something to skip. - relaid_out.append(name) - for name in relaid_out: - if self._hash_mismatches_logged < _MAX_HASH_MISMATCH_LOGS: - self._hash_mismatches_logged += 1 - logger.error( - "data-plane hash mismatch: partition=%s field=%s written " - "with uniform row lengths, read back ragged", - partition_id, - name, - ) + # Written with uniform rows, read back ragged. Treating that + # as a divergence reported 3584 mismatches per step on a + # healthy run: it is the normal shape of a real pipeline, + # where a shard is written uniform and read back inside a + # batch whose other rows differ in length. The row lengths + # are still comparable, and a row that changed length *is* a + # divergence, so check that much and count the rest as the + # abstention it is. + length_only[name] = (digest, recorded[1]) + stats.fields_skipped += 1 for row, sample_id in enumerate(sample_ids): per_field = partition_hashes.get(sample_id) if not per_field: @@ -1424,7 +1441,21 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N stats.rows_unverified += 1 continue stats.rows_checked += 1 - stats.mismatches += len(relaid_out) + for name, (digest, width) in length_only.items(): + if per_field.get(name) is None or digest.row_lens[row] == width: + continue + stats.mismatches += 1 + if self._hash_mismatches_logged < _MAX_HASH_MISMATCH_LOGS: + self._hash_mismatches_logged += 1 + logger.error( + "data-plane hash mismatch: partition=%s sample=%s " + "field=%s row was %d long on the wire in, %d out", + partition_id, + sample_id, + name, + width, + digest.row_lens[row], + ) for name, digest in comparable.items(): expected = per_field.get(name) if expected is None or expected == digest.per_row[row]: diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 1537b1d10ca..d0eb2dc68ae 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -661,12 +661,16 @@ def test_hash_shard_read_of_jagged_field_is_unverified_not_a_mismatch(): client.close() -def test_hash_rectangular_put_read_back_ragged_is_a_mismatch(): +def test_hash_rectangular_put_read_back_ragged_checks_row_lengths(): """A rectangular put can come back jagged — truncating one row makes the - batch ragged. There is no row-level digest to compare against, but the - row lengths changing between wire-in and wire-out *is* the divergence. - Reporting it as a skipped field would leave ``mismatches`` reading zero, - the exact shape of a guard that passes while covering nothing.""" + batch ragged. The content is no longer comparable, so the field counts as + an abstention, but the row *lengths* still are, and a row that changed + length is a real divergence. + + Failing the whole field instead reported 3584 mismatches per step on a + healthy 5-process run: writing a shard with uniform rows and reading it + back inside a batch whose other rows differ in length is the normal + shape of the pipeline, not a corruption.""" client = _hash_client(_RaggedOnReadClient()) ids = [f"u{i}" for i in range(4)] dense = TensorDict( @@ -677,7 +681,7 @@ def test_hash_rectangular_put_read_back_ragged_is_a_mismatch(): hv = client.snapshot()["hash_verify"] assert hv["mismatches"] > 0, "a truncated row must not read as clean" - assert hv["fields_skipped"] == 0, "this is a divergence, not an abstention" + assert hv["fields_skipped"] == 1, "and the content it could not compare" def test_hash_shard_of_a_ragged_field_is_skipped_not_a_mismatch(): @@ -1505,3 +1509,27 @@ def test_volume_namespace_does_not_become_a_breakdown_row(): ) assert [r[0] for r in rows] == ["get"], rows assert rows[0][columns.index("mb")] == 18.0 + + +def test_uniform_write_read_back_inside_a_ragged_batch_is_clean(): + """The false positive this cost: a shard written with uniform rows, read + back inside a batch whose *other* rows are ragged. Nothing diverged -- + every recorded row still has the length it was written with -- and the + guard reported 3584 mismatches per step on a healthy run until it + compared lengths instead of failing the field outright.""" + inner = _JaggedEcho() + client = _hash_client(inner) + ids = [f"u{i}" for i in range(4)] + client.put_samples( + sample_ids=ids, partition_id="p", fields=_jagged_ids([6, 6, 6, 6]) + ) + # a later writer adds rows of a different length to the same partition + other = [f"v{i}" for i in range(2)] + inner.rows[("p", other[0])] = {"ids": torch.randint(0, 32000, (3,))} + inner.rows[("p", other[1])] = {"ids": torch.randint(0, 32000, (9,))} + client.get_samples(sample_ids=ids + other, partition_id="p", select_fields=["ids"]) + + hv = client.snapshot()["hash_verify"] + assert hv["mismatches"] == 0, "no row changed length; nothing diverged" + assert hv["fields_skipped"] == 1, "content uncomparable, and counted" + client.close() From 52ab0c2b7128496ae09d822ed0b1b919c687ba8f Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sun, 23 Aug 2026 00:24:45 -0700 Subject: [PATCH 38/70] fix(data-plane): make a hash mismatch adjudicable, and measure what it costs Two complaints against this guard, both fair, both fixable. "It cries wolf." Twice now it has reported every row of every field wrong, identically, every step, and both times the cause was its own bookkeeping rather than the wire. Neither was diagnosable from the log: the digests alone say nothing about which scheme was used or what shape the read was. Per-sample lines now carry the scheme, the row index within the read, and the row length. And a step whose mismatches reach `rows_checked` is called out in the log as more likely a bug in the check than in the wire, because that is what the shape means -- a reader should not have to work that out from a raw count that looks plausible in isolation. "It costs 2.4 ms." That figure was a microbenchmark. Measured against a live TransferQueue moving 23.6 MB per step it is 10.2 ms, +11% of data-plane time -- and under 0.1% of a real GRPO step, where the data plane is a few percent of the whole. The cost is billed to `step/self/overhead_ms` along with the rest of the measurement, so it shows up on the dashboard instead of being quoted from a benchmark run elsewhere. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 14 ++++++-- nemo_rl/data_plane/observability.py | 29 ++++++++++++++-- tests/unit/data_plane/test_observability.py | 37 +++++++++++++++++++++ 3 files changed, 76 insertions(+), 4 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 0d46199d44e..d116b720713 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -725,8 +725,18 @@ limits, measured rather than assumed: collide. Rows of uniform width catch it unconditionally, per row. The same blind spot hides a reordering *within* a row: 0/200 for a two-token swap in `input_ids`, and 18/200 for moving two set bits in a bool mask. -- It reads every tensor byte again on both sides — ~2.4 ms for a 12 MB - jagged batch, on put and again on get. Keep it to debugging runs. +- It reads every tensor byte again on both sides. Measured against a live + TransferQueue moving 23.6 MB per step: **10.2 ms, or +11% of data-plane + time** — which on a real GRPO step, where the data plane is a few percent + of the step, is under 0.1% end to end. The cost lands in + `step/self/overhead_ms` like the rest of the measurement, so it is visible + rather than quoted from a benchmark. +- **A mismatch count at or above `rows_checked` is reported as suspect.** + Every row of every field wrong, identically, every step is not what a + broken wire looks like; it is what a broken guard looks like. Both false + alarms this check has produced had exactly that shape, and both were its + own bookkeeping. Per-sample lines carry the scheme, the row index and the + row length so the next one is adjudicable from a single log line. - Only rows this process wrote can be checked. A consumer-side client reports them under `hash/rows_unverified` rather than counting them clean, and `hash/fields_skipped` reports any leaf it could not compare diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 7d2222e6547..4479da19c79 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -526,7 +526,7 @@ def _hash_deltas(hv: dict[str, int], prev_hv: dict[str, int]) -> dict[str, float """ if not hv or not hv.get("rows_recorded"): return {} - return { + deltas = { f"step/hash/{name}": hv[name] - prev_hv.get(name, 0) for name in ( "rows_checked", @@ -536,6 +536,22 @@ def _hash_deltas(hv: dict[str, int], prev_hv: dict[str, int]) -> dict[str, float "fields_skipped", ) } + # Corruption of every row of every field in a step, repeated identically, + # is not what a broken wire looks like -- it is what a broken guard looks + # like. Both false alarms this check has produced had exactly this shape + # (3584 mismatches against 1536 rows, unchanging), and both were the + # guard's own bookkeeping. Say so rather than leaving a reader to decide + # whether to believe a number that large. + checked, bad = deltas["step/hash/rows_checked"], deltas["step/hash/mismatches"] + if checked > 0 and bad >= checked: + logger.warning( + "data-plane hash: %d mismatches against %d rows checked this step. " + "A rate that high is more likely a bug in the check than in the " + "wire -- confirm against the per-sample lines before acting on it.", + bad, + checked, + ) + return deltas def _volume_mb(per_op: dict[str, dict[str, float]]) -> dict[str, float]: @@ -1463,14 +1479,23 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N stats.mismatches += 1 if self._hash_mismatches_logged < _MAX_HASH_MISMATCH_LOGS: self._hash_mismatches_logged += 1 + # Scheme and shape on the line: the last two false alarms + # were both bookkeeping (a scheme replayed wrong, a + # grouping that could not be compared), and neither was + # diagnosable from the digests alone. logger.error( "data-plane hash mismatch: partition=%s sample=%s " - "field=%s wire_in=%d wire_out=%d", + "field=%s wire_in=%d wire_out=%d " + "(%s scheme, row %d of %d, %d long)", partition_id, sample_id, name, expected, digest.per_row[row], + "batch-scoped" if digest.batch_scoped else "per-row", + row, + len(sample_ids), + digest.row_lens[row] if digest.row_lens else -1, ) def _run( diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index d0eb2dc68ae..33f60b11d19 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -23,6 +23,8 @@ from time import monotonic +import logging + import pytest import torch from tensordict import NonTensorData, NonTensorStack, TensorDict @@ -33,6 +35,7 @@ breakdown_table, cluster_step_metrics, headline_series, + _hash_deltas, merge_snapshots, _estimate_encoded_bytes, _QUANTILES, @@ -1533,3 +1536,37 @@ def test_uniform_write_read_back_inside_a_ragged_batch_is_clean(): assert hv["mismatches"] == 0, "no row changed length; nothing diverged" assert hv["fields_skipped"] == 1, "content uncomparable, and counted" client.close() + + +def test_implausible_mismatch_rate_is_called_out(caplog): + """Every row of every field wrong, identically, every step is not what a + broken wire looks like -- it is what a broken guard looks like. Both false + alarms this check has produced had that shape, and a reader should not + have to work that out from a raw count.""" + hv = { + "rows_recorded": 100, + "rows_checked": 100, + "mismatches": 300, + "rows_unverified": 0, + "fields_skipped": 0, + } + with caplog.at_level(logging.WARNING): + deltas = _hash_deltas(hv, {}) + + assert deltas["step/hash/mismatches"] == 300 + assert "more likely a bug in the check" in caplog.text + + +def test_a_believable_mismatch_rate_is_not_second_guessed(caplog): + """A handful of bad rows is exactly what the guard exists to report.""" + hv = { + "rows_recorded": 100, + "rows_checked": 100, + "mismatches": 3, + "rows_unverified": 0, + "fields_skipped": 0, + } + with caplog.at_level(logging.WARNING): + _hash_deltas(hv, {}) + + assert "more likely a bug" not in caplog.text From e938ac13aade9a906e13c23856b05e3ab2078bf0 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sun, 23 Aug 2026 16:28:23 -0700 Subject: [PATCH 39/70] refactor(data-plane): apply the /simplify review Four review passes over today's commits. What they found, and what changed: The tokenizer workaround was at the wrong depth *and* pulled the wrong lever. `get_tokenizer` is one of four tokenizer constructors, and the uncovered ones include a live decode path (`trtllm_http_server` calls `decode` on a tokenizer built in `trtllm_worker_async`). Worse, the defect is a *missing attribute*, and disabling `clean_up_tokenization_spaces` instead silently changes decoded text for every non-BPE tokenizer -- transformers only skips cleanup for BPE (`type(model).__name__ == "BPE"`), so Unigram and WordPiece models were quietly losing it. Moved into `_patch_transformers_tokenizer_class_set`, which already exists for exactly this -- a version-pinned transformers defect that must reach every tokenizer however built -- and which carries the version assert and removal TODO this change had neither of. Supplying the class-level default transformers itself uses fixes the crash and changes nobody's output. `_RESERVED_NAMESPACES` was a deny-list against an open set: any later `step//` series would have become a phantom row in the breakdown table until someone remembered to extend it. Per-op detail now lives under `step/by_op//` and is recognised by what it is. Those keys are table-only -- `headline_series` drops them -- so nothing charted changes. The per-op assembly was copy-pasted between the two step-metric paths, reintroducing one level up the drift the shared helpers exist to prevent; it is now one `_op_series`. Volume was published under two keys; it keeps the `by_op` one. The write scheme was a positional `("scoped", n)` / `("rows", width)` tuple whose two arms meant different things in the same slot, under an annotation that said `int` -- now a `_WriteScheme` NamedTuple that names both. `row_lens` no longer allocates n_rows copies of one width on the uniform path, and the uniformity test is a `count` rather than a set build. Deepest of them: wandb accepts a log against a committed step, returns cleanly, and drops it. That silently cost this feature every series it had. Fixing the call order left the contract as a comment in one of three call sites, so `WandbLogger` now warns when it happens, for every caller. Verified: 191 unit tests, 23/23 guard checks, 48/48 audit checks, validate ALL PASS. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 3 + nemo_rl/data_plane/observability.py | 130 +++++++++++++------- nemo_rl/models/policy/__init__.py | 17 +++ nemo_rl/utils/logger.py | 16 +++ tests/unit/data_plane/test_observability.py | 124 ++++++++++++------- 5 files changed, 205 insertions(+), 85 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index d116b720713..3c8d9bf2ab0 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -476,6 +476,9 @@ totals and `percent_of_dataplane`, with the per-op detail in a table beside them | `step/percent_of_dataplane/by_cause/{fixed_overhead,transfer}` | is that fixed per-request cost, or moving bytes? | | `step/wall_ms`, `step/comm_volume_mb` | how much time and traffic | | `step/volume_mb/by_op/{get,put}` | which direction that traffic went | + +Per-op detail is published under `step/by_op//` and feeds the +breakdown table rather than a chart. | `now/bytes_outstanding_mb`, `now/n_processes` | occupancy, fan-out width | | `step/self/{overhead_ms,frac}` | what measuring cost | | `step/hash/*` | only with `verify_tensor_hash` on | diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 4479da19c79..a85dd5ee342 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -110,6 +110,20 @@ class DataPlaneEvent(TypedDict): _QUANTILES = ((0.50, "p50_ms", 20), (0.90, "p90_ms", 40)) +class _WriteScheme(NamedTuple): + """How a field was reduced when it was written, so a read can replay it. + + A positional ``("scoped", n)`` / ``("rows", width)`` tuple carried two + different quantities in one slot, discriminated by a magic string, under + an annotation that said ``int``. Naming both makes each read site say + which one it means. + """ + + batch_scoped: bool + n_rows: int + row_width: int + + class _FieldDigest(NamedTuple): """One fingerprint per row, plus how far it can be trusted. @@ -457,6 +471,17 @@ def _step_deltas(snap: dict[str, Any], prev: dict[str, Any]) -> dict[str, float] } +def _row_len(digest: _FieldDigest, row: int) -> int: + """How long ``row`` was, whichever scheme the digest used. + + A per-row digest stores the one width its uniform rows shared; a + batch-scoped one stores every row's length. + """ + if not digest.row_lens: + return -1 + return digest.row_lens[row] if digest.batch_scoped else digest.row_lens[0] + + def _op_step_stats( by_op: dict[str, Any], prev_ops: dict[str, Any] ) -> dict[str, dict[str, float]]: @@ -575,6 +600,37 @@ def _volume_mb(per_op: dict[str, dict[str, float]]) -> dict[str, float]: } +# Per-op detail lives under one namespace so it can be recognised by what it +# is rather than by what it is not. A deny-list of "middles that are not op +# tags" was a list against an open set: every later ``step//`` +# series -- queue depth, retry counts -- would have become a phantom row in +# the breakdown table beside put and get until someone remembered to extend +# the list. +_BY_OP = "step/by_op/" + + +# Namespaces that publish one value per op, and the table column each fills. +_BY_OP_NAMESPACES = {"percent_of_dataplane": "percent_of_dataplane", "volume_mb": "mb"} + + +def _op_series(by_op: dict[str, Any], prev_ops: dict[str, Any]) -> dict[str, float]: + """Every per-op series for one step, from two snapshots. + + The two step-metric paths share this rather than each assembling the same + keys: the helpers below exist so the single-process and cluster views + cannot drift on series *names*, and duplicating the six lines that build + those names one level up would have given the drift back. + """ + per_op = _op_step_stats(by_op, prev_ops) + metrics = _percent_of_dataplane(per_op) + metrics.update(_volume_mb(per_op)) + for op, row in per_op.items(): + for field_name, value in row.items(): + if field_name != "mb": # published once, under volume_mb/by_op + metrics[f"{_BY_OP}{op}/{field_name}"] = value + return metrics + + def _percent_of_dataplane(per_op: dict[str, dict[str, float]]) -> dict[str, float]: """Where this step's data-plane time went, in percent. @@ -860,12 +916,7 @@ def cluster_step_metrics( metrics.update( _hash_deltas(merged.get("hash_verify") or {}, prev.get("hash_verify") or {}) ) - per_op = _op_step_stats(merged["by_op"], prev.get("by_op", {})) - metrics.update(_percent_of_dataplane(per_op)) - metrics.update(_volume_mb(per_op)) - for op, row in per_op.items(): - for field_name, value in row.items(): - metrics[f"step/{op}/{field_name}"] = value + metrics.update(_op_series(merged["by_op"], prev.get("by_op", {}))) return metrics @@ -881,8 +932,6 @@ def cluster_step_metrics( "step/comm_volume_mb", "now/bytes_outstanding_mb", "now/n_processes", - "step/self/overhead_ms", - "step/self/frac", ) _HEADLINE_PREFIXES = ( "step/percent_of_dataplane/", @@ -890,10 +939,6 @@ def cluster_step_metrics( "step/hash/", "step/self/", ) -# ``step//`` is the per-op shape, so these reserved middles must -# not be mistaken for op tags -- ``step/self/overhead_ms`` was becoming a -# "self" row in the breakdown table beside put and get. -_RESERVED_NAMESPACES = frozenset({"percent_of_dataplane", "volume_mb", "hash", "self"}) def headline_series(metrics: dict[str, float]) -> dict[str, float]: @@ -961,15 +1006,16 @@ def breakdown_table( per_op: dict[str, dict[str, float]] = {} for key, value in metrics.items(): parts = key.split("/") - if len(parts) == 4 and parts[:3] == ["step", "percent_of_dataplane", "by_op"]: - per_op.setdefault(parts[3], {})["percent_of_dataplane"] = value + if len(parts) == 4 and parts[0] == "step" and parts[2] == "by_op": + column = _BY_OP_NAMESPACES.get(parts[1]) + if column: + per_op.setdefault(parts[3], {})[column] = value elif ( - len(parts) == 3 - and parts[0] == "step" - and parts[1] not in _RESERVED_NAMESPACES - and parts[2] in _BREAKDOWN_COLUMNS + len(parts) == 4 + and parts[:2] == ["step", "by_op"] + and parts[3] in _BREAKDOWN_COLUMNS ): - per_op.setdefault(parts[1], {})[parts[2]] = value + per_op.setdefault(parts[2], {})[parts[3]] = value rows = [ [op, *(stats.get(col) for col in _BREAKDOWN_COLUMNS)] # By wall time, which orders identically to ``percent_of_dataplane`` (that is @@ -1116,7 +1162,7 @@ def __init__( # of that same batch; the row count is what detects a shard read. # partition -> field -> rows the field's digest was reduced over, # for batch-scoped fields only. An absent field was reduced per row. - self._batch_scope: dict[str, dict[str, int]] = {} + self._batch_scope: dict[str, dict[str, _WriteScheme]] = {} self._hash_mismatches_logged = 0 # Set by ``_emit`` to the inner client's wall time for the op just # run, so the wrapping methods can subtract it and bill the rest to @@ -1200,12 +1246,7 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: self_ms = snap["self_ms"] - prev.get("self_ms", 0.0) metrics["step/self/overhead_ms"] = self_ms metrics["step/self/frac"] = self_ms / wall_ms if wall_ms > 0 else 0.0 - per_op = _op_step_stats(snap["by_op"], prev.get("by_op", {})) - metrics.update(_percent_of_dataplane(per_op)) - metrics.update(_volume_mb(per_op)) - for op, row in per_op.items(): - for field_name, value in row.items(): - metrics[f"step/{op}/{field_name}"] = value + metrics.update(_op_series(snap["by_op"], prev.get("by_op", {}))) return metrics def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: @@ -1352,7 +1393,8 @@ def _row_fingerprints( lengths = (offsets[1:] - offsets[:-1]).tolist() # Uniform rows: the values buffer already *is* the rectangle, # so reshaping it is a view and the per-row reduction is free. - rectangle = v.values() if len(set(lengths)) == 1 else None + uniform = lengths.count(lengths[0]) == n_rows + rectangle = v.values() if uniform else None elif v.shape[0] != n_rows: stats.fields_skipped += 1 continue @@ -1364,7 +1406,9 @@ def _row_fingerprints( out[name] = _FieldDigest( (torch.hash_tensor(flat, dim=1) ^ salt).tolist(), batch_scoped=False, - row_lens=(flat.shape[1],) * n_rows, + # One width, not n_rows copies of it: the rows are + # uniform by construction on this path. + row_lens=(flat.shape[1],), ) continue buffer = v.values() if v.is_nested else v @@ -1395,13 +1439,14 @@ def _record_hashes( # side tell a shard of a batch-scoped field from a relayout. scheme = self._batch_scope.setdefault(partition_id, {}) for name, digest in digests.items(): - if digest.batch_scoped: - scheme[name] = ("scoped", len(sample_ids)) - else: - # Width, not row count: a per-row scheme means the rows were - # uniform, so one number describes them all and lets a later - # ragged read still check whether any row changed length. - scheme[name] = ("rows", digest.row_lens[0] if digest.row_lens else 0) + # Width matters only for the per-row scheme -- uniform rows mean + # one number describes them all, and lets a later ragged read + # still check whether any row changed length. + scheme[name] = _WriteScheme( + batch_scoped=digest.batch_scoped, + n_rows=len(sample_ids), + row_width=digest.row_lens[0] if digest.row_lens else 0, + ) self._stats.hash_verify.rows_recorded += len(sample_ids) def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: @@ -1412,7 +1457,7 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N # Only the batch-scoped names: the scheme also records per-row fields # now, and handing the whole mapping over forced every field onto the # scoped path. - scoped_names = {n for n, how in scheme.items() if how[0] == "scoped"} + scoped_names = {n for n, how in scheme.items() if how.batch_scoped} digests = self._row_fingerprints(out, sample_ids, scoped_names) if not digests: return @@ -1433,8 +1478,8 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N comparable[name] = digest elif recorded is None: continue # this process never wrote the field - elif recorded[0] == "scoped": - if recorded[1] == len(sample_ids): + elif recorded.batch_scoped: + if recorded.n_rows == len(sample_ids): comparable[name] = digest else: stats.fields_skipped += 1 # a shard of a batch-scoped put @@ -1447,7 +1492,7 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N # are still comparable, and a row that changed length *is* a # divergence, so check that much and count the rest as the # abstention it is. - length_only[name] = (digest, recorded[1]) + length_only[name] = (digest, recorded.row_width) stats.fields_skipped += 1 for row, sample_id in enumerate(sample_ids): per_field = partition_hashes.get(sample_id) @@ -1458,7 +1503,8 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N continue stats.rows_checked += 1 for name, (digest, width) in length_only.items(): - if per_field.get(name) is None or digest.row_lens[row] == width: + read_len = _row_len(digest, row) + if per_field.get(name) is None or read_len == width: continue stats.mismatches += 1 if self._hash_mismatches_logged < _MAX_HASH_MISMATCH_LOGS: @@ -1470,7 +1516,7 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N sample_id, name, width, - digest.row_lens[row], + read_len, ) for name, digest in comparable.items(): expected = per_field.get(name) @@ -1495,7 +1541,7 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N "batch-scoped" if digest.batch_scoped else "per-row", row, len(sample_ids), - digest.row_lens[row] if digest.row_lens else -1, + _row_len(digest, row), ) def _run( diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 3325edfbfb4..27c8e0bbc0c 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -80,6 +80,23 @@ def _patched_from_pretrained(pretrained_model_name_or_path, *args, **kwargs): pretrained_model_name_or_path, *args, **kwargs ) + # Second defect in the same pinned range: ``TokenizersBackend._decode`` + # reads ``clean_up_tokenization_spaces_for_bpe_even_though_it_will_corrupt_output`` + # whenever cleanup is enabled, but some builds never set it, so the first + # ``batch_decode`` of generated tokens raises AttributeError one step into + # a rollout. Supplying the class-level default transformers itself uses + # restores the intended behaviour -- BPE skips cleanup, everything else + # still gets it -- rather than disabling cleanup, which would change + # decoded text for WordPiece and Unigram models. Instances that do define + # the attribute shadow this, so it is inert on a fixed build. + from transformers.tokenization_utils_tokenizers import TokenizersBackend + + _BPE_CLEANUP = ( + "clean_up_tokenization_spaces_for_bpe_even_though_it_will_corrupt_output" + ) + if not hasattr(TokenizersBackend, _BPE_CLEANUP): + setattr(TokenizersBackend, _BPE_CLEANUP, False) + AutoTokenizer.from_pretrained = _patched_from_pretrained diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index dc341a5ce5b..1b14d3a373e 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -236,6 +236,11 @@ def __init__(self, cfg: WandbConfig, log_dir: Optional[str] = None): wandb_init_config = dict(cfg) wandb_init_config.pop("log_nemo_gym_full_result_tables", None) self.run = wandb.init(**wandb_init_config, dir=log_dir) + # Highest step already committed. wandb accepts a log against a + # committed step, returns cleanly, and drops it -- a silent data loss + # that cost a whole feature's series once, because the caller ran + # after the log that ends the step and nothing said so. + self._committed_step = -1 if os.environ.get("RAY_BACKEND_LOG_LEVEL", "").lower() == "debug": print( @@ -402,7 +407,18 @@ def log_metrics( # Commit param defaults to None. By default if step is set, then commit defaults to False # Here, we have an explicit fork for commit in case W&B ever decides to change their default logic. self.run.log(metrics, step=step, commit=True) + self._committed_step = max(self._committed_step, step) else: + if step <= self._committed_step: + # Loud, because wandb is not: it would accept this and drop it. + logging.getLogger(__name__).warning( + "wandb: %d metrics logged at step %d, which was already " + "committed (step_finished=True) -- wandb will discard " + "them. Log before the call that ends the step. Keys: %s", + len(metrics), + step, + sorted(metrics)[:5], + ) self.run.log(metrics, step=step) def log_hyperparams(self, params: Mapping[str, Any]) -> None: diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 33f60b11d19..29c14966421 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -24,6 +24,7 @@ from time import monotonic import logging +from unittest.mock import patch import pytest import torch @@ -418,9 +419,11 @@ def test_step_metrics_tail_is_exact_not_bucketed(): client._emit("put", "p", 1, 8, monotonic() - 0.030, "ok") # a 30 ms call metrics = client.get_step_metrics(1.0) - assert "put/p90_ms" not in metrics and "step/put/p50_ms" not in metrics - assert metrics["step/put/max_ms"] >= 30.0 - assert metrics["step/put/max_ms"] != pytest.approx(24.85, abs=0.5), "bucket edge" + assert "put/p90_ms" not in metrics and "step/by_op/put/p50_ms" not in metrics + assert metrics["step/by_op/put/max_ms"] >= 30.0 + assert metrics["step/by_op/put/max_ms"] != pytest.approx(24.85, abs=0.5), ( + "bucket edge" + ) # and one call supports no percentile at all, in either view assert "p90_ms" not in client.snapshot()["by_op"]["put"] client.close() @@ -449,11 +452,15 @@ def test_latency_breakdown_stacks_to_wall_ms(): assert fit["bandwidth_mb_s"] == pytest.approx(mb_per_s, rel=0.05) # the two components are the split of ONE call, and they add to the mean - total = metrics["step/put/overhead_ms"] + metrics["step/put/transfer_ms"] - assert total == pytest.approx(metrics["step/put/mean_ms"], rel=0.05) + total = ( + metrics["step/by_op/put/overhead_ms"] + metrics["step/by_op/put/transfer_ms"] + ) + assert total == pytest.approx(metrics["step/by_op/put/mean_ms"], rel=0.05) # per call, the overhead term IS the fitted constant - assert metrics["step/put/overhead_ms"] == pytest.approx(fixed_ms, rel=0.05) - assert "step/put/overhead_frac" not in metrics, "a ratio is derivable from these" + assert metrics["step/by_op/put/overhead_ms"] == pytest.approx(fixed_ms, rel=0.05) + assert "step/by_op/put/overhead_frac" not in metrics, ( + "a ratio is derivable from these" + ) client.close() @@ -463,9 +470,9 @@ def test_step_max_is_scoped_to_the_step(): reported max must fall again when a step is quicker.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient()) client._emit("put", "p", 1, 8, monotonic() - 0.050, "ok") # slow step - slow = client.get_step_metrics(1.0)["step/put/max_ms"] + slow = client.get_step_metrics(1.0)["step/by_op/put/max_ms"] client._emit("put", "p", 1, 8, monotonic() - 0.001, "ok") # quick step - quick = client.get_step_metrics(1.0)["step/put/max_ms"] + quick = client.get_step_metrics(1.0)["step/by_op/put/max_ms"] assert slow >= 50.0 assert quick < slow, "step max must reset, not carry the lifetime worst" @@ -889,10 +896,10 @@ def test_cluster_percentiles_never_exceed_the_measured_max(): merged = merge_snapshots([_rank_with([120.0] * 20) for _ in range(8)]) metrics = cluster_step_metrics(merged, {}, 1.0) - assert metrics["step/put/max_ms"] == pytest.approx(120.0, abs=2.0) - assert metrics["step/put/p50_ms"] <= metrics["step/put/max_ms"] - assert metrics["step/put/p90_ms"] <= metrics["step/put/max_ms"] - assert metrics["step/put/p50_ms"] <= metrics["step/put/p90_ms"] + assert metrics["step/by_op/put/max_ms"] == pytest.approx(120.0, abs=2.0) + assert metrics["step/by_op/put/p50_ms"] <= metrics["step/by_op/put/max_ms"] + assert metrics["step/by_op/put/p90_ms"] <= metrics["step/by_op/put/max_ms"] + assert metrics["step/by_op/put/p50_ms"] <= metrics["step/by_op/put/p90_ms"] def test_cluster_percentiles_withheld_below_a_useful_sample_count(): @@ -901,9 +908,11 @@ def test_cluster_percentiles_withheld_below_a_useful_sample_count(): few = merge_snapshots([_rank_with([120.0] * 3) for _ in range(2)]) # 6 calls many = merge_snapshots([_rank_with([120.0] * 20) for _ in range(8)]) # 160 - assert "step/put/p50_ms" not in cluster_step_metrics(few, {}, 1.0) - assert "step/put/max_ms" in cluster_step_metrics(few, {}, 1.0), "max always works" - assert "step/put/p50_ms" in cluster_step_metrics(many, {}, 1.0) + assert "step/by_op/put/p50_ms" not in cluster_step_metrics(few, {}, 1.0) + assert "step/by_op/put/max_ms" in cluster_step_metrics(few, {}, 1.0), ( + "max always works" + ) + assert "step/by_op/put/p50_ms" in cluster_step_metrics(many, {}, 1.0) def test_cluster_series_declare_delta_or_level(): @@ -1019,12 +1028,12 @@ def test_breakdown_table_rows_by_op_worst_first(): "step/wall_ms": 100.0, "step/percent_of_dataplane/by_op/get": 10.0, "step/percent_of_dataplane/by_op/put": 90.0, - "step/get/calls": 8, - "step/get/wall_ms": 10.0, - "step/get/max_ms": 2.0, - "step/put/calls": 2, - "step/put/wall_ms": 90.0, - "step/put/max_ms": 50.0, + "step/by_op/get/calls": 8, + "step/by_op/get/wall_ms": 10.0, + "step/by_op/get/max_ms": 2.0, + "step/by_op/put/calls": 2, + "step/by_op/put/wall_ms": 90.0, + "step/by_op/put/max_ms": 50.0, "step/comm_volume_mb": 1.0, # not per-op, must not become a row "now/bytes_outstanding_mb": 0.0, # a level, likewise } @@ -1043,7 +1052,11 @@ def test_breakdown_table_leaves_withheld_series_empty(): is absent from the series — the table must carry None there rather than a zero that would read as a measurement.""" columns, rows = breakdown_table( - {"step/put/calls": 3, "step/put/wall_ms": 5.0, "step/put/max_ms": 2.0} + { + "step/by_op/put/calls": 3, + "step/by_op/put/wall_ms": 5.0, + "step/by_op/put/max_ms": 2.0, + } ) row = rows[0] assert row[columns.index("p90_ms")] is None @@ -1074,11 +1087,13 @@ def test_cluster_view_carries_the_latency_split(): ranks.append(client.snapshot()) metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 1.0) - assert "step/get/overhead_ms" in metrics, "cluster view must carry the split" + assert "step/by_op/get/overhead_ms" in metrics, "cluster view must carry the split" - total = metrics["step/get/overhead_ms"] + metrics["step/get/transfer_ms"] - assert total == pytest.approx(metrics["step/get/mean_ms"], rel=0.05) - assert metrics["step/get/overhead_ms"] == pytest.approx(fixed_ms, rel=0.05) + total = ( + metrics["step/by_op/get/overhead_ms"] + metrics["step/by_op/get/transfer_ms"] + ) + assert total == pytest.approx(metrics["step/by_op/get/mean_ms"], rel=0.05) + assert metrics["step/by_op/get/overhead_ms"] == pytest.approx(fixed_ms, rel=0.05) columns, rows = breakdown_table(metrics) row = rows[0] @@ -1099,12 +1114,12 @@ def test_cluster_per_op_time_is_reported_per_call(): merge_snapshots([_rank_with([10.0] * 5) for _ in range(32)]), {}, 1.0 ) - assert small["step/put/mean_ms"] == pytest.approx(10.0, rel=0.15) - assert large["step/put/mean_ms"] == pytest.approx( - small["step/put/mean_ms"], rel=0.15 + assert small["step/by_op/put/mean_ms"] == pytest.approx(10.0, rel=0.15) + assert large["step/by_op/put/mean_ms"] == pytest.approx( + small["step/by_op/put/mean_ms"], rel=0.15 ), "mean must not move with cluster size" - assert large["step/put/wall_ms"] == pytest.approx( - 4 * small["step/put/wall_ms"], rel=0.15 + assert large["step/by_op/put/wall_ms"] == pytest.approx( + 4 * small["step/by_op/put/wall_ms"], rel=0.15 ), "the sum does move with cluster size" columns, rows = breakdown_table(small) @@ -1128,12 +1143,12 @@ def metrics_for(n_calls): client._emit("put", "p", 1, 1_000, now - (5.0 + i % 7) / 1e3, "ok") return cluster_step_metrics(merge_snapshots([client.snapshot()]), {}, 1.0) - assert "step/put/p50_ms" not in metrics_for(10), "too thin for either" + assert "step/by_op/put/p50_ms" not in metrics_for(10), "too thin for either" mid = metrics_for(30) - assert "step/put/p50_ms" in mid, "a median off 30 calls is real" - assert "step/put/p90_ms" not in mid, "a p90 off 30 calls is not" + assert "step/by_op/put/p50_ms" in mid, "a median off 30 calls is real" + assert "step/by_op/put/p90_ms" not in mid, "a p90 off 30 calls is not" both = metrics_for(58) - assert "step/put/p50_ms" in both and "step/put/p90_ms" in both + assert "step/by_op/put/p50_ms" in both and "step/by_op/put/p90_ms" in both def test_no_quantile_finer_than_the_sample_size_can_resolve(): @@ -1333,7 +1348,7 @@ def test_cluster_step_max_reopens_each_step(): "put", "p", 1, 1_000, now - (slowest_ms if i == 0 else 5.0) / 1e3, "ok" ) merged = merge_snapshots([client.snapshot(reset_step_window=True)]) - seen.append(cluster_step_metrics(merged, prev, 1.0)["step/put/max_ms"]) + seen.append(cluster_step_metrics(merged, prev, 1.0)["step/by_op/put/max_ms"]) prev = merged assert seen[1] == pytest.approx(50.0, abs=1.0), "the spike shows" @@ -1360,8 +1375,8 @@ def test_breakdown_table_ignores_reserved_namespaces(): series and was becoming a "self" row in the table beside put and get.""" columns, rows = breakdown_table( { - "step/put/calls": 2, - "step/put/wall_ms": 9.0, + "step/by_op/put/calls": 2, + "step/by_op/put/wall_ms": 9.0, "step/self/overhead_ms": 6.2, "step/hash/mismatches": 0, "step/percent_of_dataplane/by_cause/transfer": 40.0, @@ -1433,7 +1448,9 @@ def test_write_and_read_volume_are_not_computed_for_nobody(): assert "step/bytes_written_mb" not in metrics assert "step/bytes_read_mb" not in metrics assert metrics["step/comm_volume_mb"] > 0, "the total is still reported" - assert metrics["step/put/mb"] == pytest.approx(0.004), "and split per op" + assert metrics["step/volume_mb/by_op/put"] == pytest.approx(0.004), ( + "and split per op, under one key rather than two" + ) client.close() @@ -1504,9 +1521,9 @@ def test_volume_namespace_does_not_become_a_breakdown_row(): ``volume_mb`` op beside put and get.""" columns, rows = breakdown_table( { - "step/get/calls": 3, - "step/get/wall_ms": 9.0, - "step/get/mb": 18.0, + "step/by_op/get/calls": 3, + "step/by_op/get/wall_ms": 9.0, + "step/by_op/get/mb": 18.0, "step/volume_mb/by_op/get": 18.0, } ) @@ -1570,3 +1587,24 @@ def test_a_believable_mismatch_rate_is_not_second_guessed(caplog): _hash_deltas(hv, {}) assert "more likely a bug" not in caplog.text + + +def test_wandb_logger_warns_when_a_step_is_written_after_it_commits(): + """wandb accepts a log against a committed step, returns cleanly, and + drops it. That silently cost this feature every one of its series until + someone read a finished run back out of the API -- so the Logger says so + rather than leaving each call site to remember the ordering.""" + from nemo_rl.utils.logger import WandbLogger + + logged: list[dict] = [] + logger_obj = WandbLogger.__new__(WandbLogger) + logger_obj.run = type("R", (), {"log": lambda self, m, **kw: logged.append(m)})() + logger_obj._committed_step = -1 + + with patch.object(logging.getLogger("nemo_rl.utils.logger"), "warning") as warn: + logger_obj.log_metrics({"a": 1.0}, step=3, step_finished=True) + assert warn.call_count == 0, "the committing log is fine" + logger_obj.log_metrics({"b": 2.0}, step=3) + assert warn.call_count == 1, "the one after it is not" + logger_obj.log_metrics({"c": 3.0}, step=4) + assert warn.call_count == 1, "a later step is fine again" From bf8ffb9bd04a075b180e367aa663cef90c62f658 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sun, 23 Aug 2026 16:33:57 -0700 Subject: [PATCH 40/70] refactor(data-plane): narrow this PR back to the data plane The /simplify pass fixed three things that are real but are not this PR's subject, and carrying them here makes the diff argue for changes a reviewer did not come to review. Out, kept as `tokenizer-decode-fix.patch` for its own PR: the `TokenizersBackend` decode workaround in `nemo_rl/models/policy/__init__.py`. It fixes a transformers defect, not a data-plane one; it happens to be what unblocked the e2e run for this feature, which is not the same as belonging to it. Out, to be raised separately: the `WandbLogger` warning on writes to an already-committed step. Shared logging infrastructure used by ppo, grpo and sft, changed on the strength of a bug found in one caller. The ordering fix in `grpo_sync.py` already prevents the failure here; making the Logger say so for every caller is a wider argument than this diff should make. Out: the `-tq_simple-obs` recipe, which existed to launch my own verification run and is wired into no test suite, so by the repo's own convention it is not a recipe. `log_table` stays -- the breakdown table cannot be logged without it. Verified: 190 unit tests, 23/23 guard checks, 48/48 audit checks. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- ...-instruct-1n8g-fsdp2tp1-tq_simple-obs.yaml | 9 -------- nemo_rl/models/policy/__init__.py | 17 -------------- nemo_rl/utils/logger.py | 16 -------------- tests/unit/data_plane/test_observability.py | 22 ------------------- 4 files changed, 64 deletions(-) delete mode 100644 examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1-tq_simple-obs.yaml diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1-tq_simple-obs.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1-tq_simple-obs.yaml deleted file mode 100644 index 51139d56be6..00000000000 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1-tq_simple-obs.yaml +++ /dev/null @@ -1,9 +0,0 @@ -# Data-plane observability e2e: real GRPO traffic through TransferQueue with -# per-op metrics and the wire-in/wire-out hash guard both on. Short run -- -# the point is the data_plane/* series, not convergence. -defaults: grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3.yaml -data_plane: - enabled: true - observability: - enabled: true - verify_tensor_hash: true diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 27c8e0bbc0c..3325edfbfb4 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -80,23 +80,6 @@ def _patched_from_pretrained(pretrained_model_name_or_path, *args, **kwargs): pretrained_model_name_or_path, *args, **kwargs ) - # Second defect in the same pinned range: ``TokenizersBackend._decode`` - # reads ``clean_up_tokenization_spaces_for_bpe_even_though_it_will_corrupt_output`` - # whenever cleanup is enabled, but some builds never set it, so the first - # ``batch_decode`` of generated tokens raises AttributeError one step into - # a rollout. Supplying the class-level default transformers itself uses - # restores the intended behaviour -- BPE skips cleanup, everything else - # still gets it -- rather than disabling cleanup, which would change - # decoded text for WordPiece and Unigram models. Instances that do define - # the attribute shadow this, so it is inert on a fixed build. - from transformers.tokenization_utils_tokenizers import TokenizersBackend - - _BPE_CLEANUP = ( - "clean_up_tokenization_spaces_for_bpe_even_though_it_will_corrupt_output" - ) - if not hasattr(TokenizersBackend, _BPE_CLEANUP): - setattr(TokenizersBackend, _BPE_CLEANUP, False) - AutoTokenizer.from_pretrained = _patched_from_pretrained diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index 1b14d3a373e..dc341a5ce5b 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -236,11 +236,6 @@ def __init__(self, cfg: WandbConfig, log_dir: Optional[str] = None): wandb_init_config = dict(cfg) wandb_init_config.pop("log_nemo_gym_full_result_tables", None) self.run = wandb.init(**wandb_init_config, dir=log_dir) - # Highest step already committed. wandb accepts a log against a - # committed step, returns cleanly, and drops it -- a silent data loss - # that cost a whole feature's series once, because the caller ran - # after the log that ends the step and nothing said so. - self._committed_step = -1 if os.environ.get("RAY_BACKEND_LOG_LEVEL", "").lower() == "debug": print( @@ -407,18 +402,7 @@ def log_metrics( # Commit param defaults to None. By default if step is set, then commit defaults to False # Here, we have an explicit fork for commit in case W&B ever decides to change their default logic. self.run.log(metrics, step=step, commit=True) - self._committed_step = max(self._committed_step, step) else: - if step <= self._committed_step: - # Loud, because wandb is not: it would accept this and drop it. - logging.getLogger(__name__).warning( - "wandb: %d metrics logged at step %d, which was already " - "committed (step_finished=True) -- wandb will discard " - "them. Log before the call that ends the step. Keys: %s", - len(metrics), - step, - sorted(metrics)[:5], - ) self.run.log(metrics, step=step) def log_hyperparams(self, params: Mapping[str, Any]) -> None: diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 29c14966421..8225d9bfca8 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -24,7 +24,6 @@ from time import monotonic import logging -from unittest.mock import patch import pytest import torch @@ -1587,24 +1586,3 @@ def test_a_believable_mismatch_rate_is_not_second_guessed(caplog): _hash_deltas(hv, {}) assert "more likely a bug" not in caplog.text - - -def test_wandb_logger_warns_when_a_step_is_written_after_it_commits(): - """wandb accepts a log against a committed step, returns cleanly, and - drops it. That silently cost this feature every one of its series until - someone read a finished run back out of the API -- so the Logger says so - rather than leaving each call site to remember the ordering.""" - from nemo_rl.utils.logger import WandbLogger - - logged: list[dict] = [] - logger_obj = WandbLogger.__new__(WandbLogger) - logger_obj.run = type("R", (), {"log": lambda self, m, **kw: logged.append(m)})() - logger_obj._committed_step = -1 - - with patch.object(logging.getLogger("nemo_rl.utils.logger"), "warning") as warn: - logger_obj.log_metrics({"a": 1.0}, step=3, step_finished=True) - assert warn.call_count == 0, "the committing log is fine" - logger_obj.log_metrics({"b": 2.0}, step=3) - assert warn.call_count == 1, "the one after it is not" - logger_obj.log_metrics({"c": 3.0}, step=4) - assert warn.call_count == 1, "a later step is fine again" From 6ebf71c0d3c540024ede6e115ac93d20d57fc23c Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sun, 23 Aug 2026 17:59:37 -0700 Subject: [PATCH 41/70] fix(data-plane): satisfy pyrefly on the three types the diff loosened CI's lint job runs pyrefly, which I had not run locally -- ruff alone passes on all three of these. `_hash_deltas` builds counter deltas, which are ints, under a signature promising floats. `get_data_plane_snapshot` returns the result of a ``getattr``-ed callable, so its type is `object`, not the `dict[str, Any] | None` it declares. And in `factory`, dropping the `log_event` default left `on_event` as the argument pyrefly objects to, while the existing ignore sat on the constructor line rather than the argument's own. Verified: pyrefly clean across `nemo_rl/data_plane/`, 222 unit tests, 23/23 guard checks, 48/48 audit checks. Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/factory.py | 6 +++--- nemo_rl/data_plane/observability.py | 2 +- nemo_rl/data_plane/worker_mixin.py | 7 +++++-- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/nemo_rl/data_plane/factory.py b/nemo_rl/data_plane/factory.py index b54756f45ae..e9c510275db 100644 --- a/nemo_rl/data_plane/factory.py +++ b/nemo_rl/data_plane/factory.py @@ -186,9 +186,9 @@ def build_data_plane_client( # single transfer. ``log_event`` is still exported for anyone who # wants that, but it is opt-in via ``observability.callback``. # pyrefly: obs.get returns Any, can't narrow to the expected callback type. - client = MetricsDataPlaneClient( # type: ignore[bad-argument-type] - client, - on_event=obs.get("callback"), + client = MetricsDataPlaneClient( + client, # type: ignore[bad-argument-type] + on_event=obs.get("callback"), # type: ignore[bad-argument-type] verify_tensor_hash=bool(obs.get("verify_tensor_hash")), ) return client diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index a85dd5ee342..06bef8c8c53 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -551,7 +551,7 @@ def _hash_deltas(hv: dict[str, int], prev_hv: dict[str, int]) -> dict[str, float """ if not hv or not hv.get("rows_recorded"): return {} - deltas = { + deltas: dict[str, float] = { f"step/hash/{name}": hv[name] - prev_hv.get(name, 0) for name in ( "rows_checked", diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index 7fe41007dc9..0c4d8e8432f 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -31,7 +31,7 @@ import logging import time from collections import Counter -from typing import TYPE_CHECKING, Any, Literal, Optional +from typing import TYPE_CHECKING, Any, Literal, Optional, cast import numpy as np import torch @@ -398,7 +398,10 @@ def get_data_plane_snapshot(self) -> "dict[str, Any] | None": """ client = getattr(self, "_dp_client", None) snapshot = getattr(client, "snapshot", None) - return snapshot(reset_step_window=True) if callable(snapshot) else None + if not callable(snapshot): + return None + # cast: ``snapshot`` came off getattr, so it is untyped here. + return cast("dict[str, Any] | None", snapshot(reset_step_window=True)) def _fetch( self, From 33468e2b3a0f01280255decfc872d1e3aa5c2252 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sun, 23 Aug 2026 18:03:57 -0700 Subject: [PATCH 42/70] style(data-plane): sort the test imports The lint job runs three ruff hooks; the second is isort (--select I), which I had not run -- the import block in the observability tests grew a symbol at a time and ended up unsorted. Formatter and linter alone pass on it. Signed-off-by: Zhiyu Li --- tests/unit/data_plane/test_observability.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 8225d9bfca8..5329811a817 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -21,9 +21,8 @@ from __future__ import annotations -from time import monotonic - import logging +from time import monotonic import pytest import torch @@ -31,15 +30,15 @@ from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient from nemo_rl.data_plane.observability import ( + _QUANTILES, MetricsDataPlaneClient, + _estimate_encoded_bytes, + _hash_deltas, + _td_bytes, breakdown_table, cluster_step_metrics, headline_series, - _hash_deltas, merge_snapshots, - _estimate_encoded_bytes, - _QUANTILES, - _td_bytes, ) from ._rollout_shapes import make_rollout_batch From f51c500801528499b2c3d05ed72150dd1c1385e6 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sun, 23 Aug 2026 19:39:44 -0700 Subject: [PATCH 43/70] fix(data-plane): salt hash fingerprints on the host torch has no bitwise_xor CUDA kernel for UInt64, so XOR-ing the salt into the digest tensor while it is still on the device raises NotImplementedError. Any backend whose get returns device tensors hits it -- TransferQueue register mode under GDR does, and it took the step down at the first verified get. The batch-scoped branch a few lines below already salts after .tolist(); this makes the per-row branch agree. The digests move to the host for comparison either way, so nothing is paid for it. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 8 +++++++- tests/unit/data_plane/test_observability.py | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 06bef8c8c53..9d932e77f33 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -1404,7 +1404,13 @@ def _row_fingerprints( if rectangle is not None and name not in batch_scoped_fields: flat = _as_int_view(rectangle.reshape(n_rows, -1)) out[name] = _FieldDigest( - (torch.hash_tensor(flat, dim=1) ^ salt).tolist(), + # Salt on the host, as the batch-scoped path below already + # does. Torch has no ``bitwise_xor`` CUDA kernel for + # UInt64, so XOR-ing the digest tensor in place raises + # ``NotImplementedError`` for any backend whose get returns + # device tensors. The digests come back to the host for + # comparison either way, so this costs nothing. + [d ^ salt for d in torch.hash_tensor(flat, dim=1).tolist()], batch_scoped=False, # One width, not n_rows copies of it: the rows are # uniform by construction on this path. diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 5329811a817..ede4f4bd5c0 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -750,6 +750,25 @@ def test_hash_fingerprint_handles_float8(): assert len(digest.per_row) == 2 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_hash_fingerprint_handles_device_tensors(): + """Torch has no ``bitwise_xor`` CUDA kernel for UInt64, so salting the + digest tensor before it leaves the device raises ``NotImplementedError`` + for any backend whose get returns device tensors -- register mode under + GDR does. The digests must also match the host's, or a device-resident + get would verify against a host put as a mismatch on every row.""" + client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + values = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + ids = ["a", "b"] + on_host = client._row_fingerprints(TensorDict({"x": values}, batch_size=[2]), ids)[ + "x" + ] + on_device = client._row_fingerprints( + TensorDict({"x": values.cuda()}, batch_size=[2]), ids + )["x"] + assert on_device.per_row == on_host.per_row + + def test_hash_verification_off_by_default(): """Default construction must do no hashing work at all.""" client = MetricsDataPlaneClient(NoOpDataPlaneClient()) From f35e6cb681211ddc656c36ec28a6932181c07b61 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sun, 23 Aug 2026 19:39:51 -0700 Subject: [PATCH 44/70] feat(data-plane): log data-plane metrics from the single-controller loop The single-controller path has its own step loop and never called _log_data_plane_metrics, so enabling observability there built the metrics client, paid for its counters on every op, and logged nothing. An empty dashboard reads the same as a data plane that cost nothing, which is the worse of the two failures. Driver scope only, and the prefix says so: this client issues the advantage stage's get and the post-train clear, while the bulk traffic is the trainer and generation workers' own clients in their own processes. comm_volume_mb here is well under what the job moved. grpo_sync prefers a cluster view by fanning out over its policy worker group, but nothing implements collect_data_plane_snapshots yet, so that path also falls through to the driver's counters alone -- there is no cluster view to mirror until one exists. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/single_controller.py | 47 +++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 36dd5ffd103..dc6dafee9cf 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -119,6 +119,11 @@ from nemo_rl.data.multimodal_utils import present_multimodal_fields from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, KVBatchMeta from nemo_rl.data_plane.async_utils import call_data_plane +from nemo_rl.data_plane.observability import ( + MetricsDataPlaneClient, + breakdown_table, + headline_series, +) from nemo_rl.data_plane.schema import ( DP_CALIB_INPUT_FIELDS, DP_TRAIN_FIELDS, @@ -1518,6 +1523,47 @@ async def _cleanup_consumed_metas_unlocked( errors.append(cleanup_error) if errors: raise BaseExceptionGroup("post-train DataPlane cleanup failed", errors) + def _log_data_plane_metrics(self, total_step_time: float) -> None: + """Log this step's data-plane cost. No-op unless observability is enabled. + + The synchronous loop logs these series from ``_log_data_plane_metrics`` + in ``grpo_sync``. Without the same call here the single-controller path + builds the metrics client, pays for its counters on every op, and emits + nothing -- the failure is silent, because an empty dashboard looks the + same as a data plane that cost nothing. + + Driver scope only, and the prefix says so. This client issues the + advantage stage's get plus the post-train clear; the bulk traffic is + the trainer and generation workers' own clients, in their own + processes with their own counters, so ``comm_volume_mb`` here is well + under what the job actually moved. ``grpo_sync`` prefers a cluster + view by fanning out over its policy worker group, but nothing + implements ``collect_data_plane_snapshots`` yet, so that path also + falls through to the driver's counters alone -- there is no cluster + view to mirror here until one exists. + """ + if not isinstance(self._dp_client, MetricsDataPlaneClient): + return # observability disabled -> plain adapter + + metrics = self._dp_client.get_step_metrics(total_step_time) + step = self._train_steps + self._logger.log_metrics( + headline_series(metrics), step, prefix="data_plane/driver" + ) + try: + columns, rows = breakdown_table(metrics) + except Exception as exc: # noqa: BLE001 - a panel must never fail a step + logging.getLogger(__name__).warning("data-plane breakdown failed: %s", exc) + else: + if rows: + self._logger.log_table( + columns, rows, step, "data_plane/driver/breakdown" + ) + print( + f" • data plane: {metrics['step/wall_ms']:.0f}ms, " + f"{metrics['step/comm_volume_mb']:.1f} MB moved", + flush=True, + ) @staticmethod def _group_ids_from_meta(meta: KVBatchMeta) -> list[str]: @@ -2856,6 +2902,7 @@ async def _train_pump(self) -> None: prefix="timing/train", step_finished=True, ) + self._log_data_plane_metrics(total_time) self._timer.reset() # min sample version refers to the version each consumed sample was From 58a49b5a34a53b3b1f0719907134c0c2fec68ed7 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 26 Aug 2026 18:23:07 -0700 Subject: [PATCH 45/70] test(data-plane): simplify observability tests; guard the metrics path Rewrite tests/unit/data_plane/test_observability.py (72 -> 50 tests), add a mock e2e test for the data-plane logging path in test_grpo.py, and make the hash guard and metrics panel non-raising so neither can fail a training step. Enable verify_tensor_hash on all 19 TQ nightlies and gate them on hash counters being zero. Signed-off-by: Zhiyu Li --- .../grpo-deepscaler-1.5b-8K-tq_simple.yaml | 2 + ...-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml | 2 + ...rpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml | 2 + ...8g-megatron-fp8-rollouts.v3-tq_simple.yaml | 2 + ...-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml | 2 + ...uct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml | 2 + ...b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml | 2 + ...p2-temp0.8-topp0.9-topk50-tq_mooncake.yaml | 2 + ...on-temp0.8-topp0.9-topk50-tq_mooncake.yaml | 2 + ...2-1b-instruct-1n8g-megatron-tq_simple.yaml | 2 + ...-1n8g-megatron_generation-tq_mooncake.yaml | 2 + ...nlight-16ba3b-4n8g-megatron-tq_simple.yaml | 2 + ...nov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml | 2 + ...0BA3B-2n8g-megatron-pack-cp-tq_simple.yaml | 2 + ...b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml | 2 + ...1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml | 2 + ...30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml | 2 + ...3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml | 2 + ...instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml | 2 + nemo_rl/algorithms/grpo_sync.py | 30 +- nemo_rl/algorithms/single_controller.py | 28 +- nemo_rl/data_plane/observability.py | 51 +- .../llm/grpo-deepscaler-1.5b-8K-tq_simple.sh | 6 + ...po-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh | 6 + .../grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh | 6 + ...1n8g-megatron-fp8-rollouts.v3-tq_simple.sh | 6 + ...ct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh | 6 + ...truct-2n8g-megatron-fp8-e2e-tq_mooncake.sh | 6 + ...-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh | 6 + ...2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh | 6 + ...tron-temp0.8-topp0.9-topk50-tq_mooncake.sh | 6 + ...3.2-1b-instruct-1n8g-megatron-tq_simple.sh | 6 + ...ct-1n8g-megatron_generation-tq_mooncake.sh | 6 + ...oonlight-16ba3b-4n8g-megatron-tq_simple.sh | 6 + ...nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh | 6 + ...-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh | 6 + ....5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh | 6 + ...3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh | 6 + ...3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh | 5 +- ...en3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh | 6 + ...b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh | 6 + tests/unit/algorithms/test_grpo.py | 108 + tests/unit/data_plane/test_observability.py | 1894 +++++++---------- 43 files changed, 1099 insertions(+), 1163 deletions(-) diff --git a/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml b/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml index 46f7381aaec..5a711173ac3 100644 --- a/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml @@ -1,3 +1,5 @@ defaults: grpo-deepscaler-1.5b-8K.yaml data_plane: enabled: true + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml b/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml index 7e693ba250e..01ea51f7945 100644 --- a/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml @@ -1,3 +1,5 @@ defaults: grpo-gemma3-1b-it-1n8g-fsdp2tp1.yaml data_plane: enabled: true + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml b/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml index 80d795c7215..ca900a52775 100644 --- a/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml @@ -1,3 +1,5 @@ defaults: grpo-gspo-deepscaler-1.5b-8K.yaml data_plane: enabled: true + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml index 8bf6005dbf8..7ff662d1d73 100644 --- a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml @@ -1,3 +1,5 @@ defaults: grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3.yaml data_plane: enabled: true + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml index f457fecd0aa..0eeb38d3960 100644 --- a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml @@ -1,3 +1,5 @@ defaults: grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated.yaml data_plane: enabled: true + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml index 5313f25dd42..2e8bb388958 100644 --- a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml @@ -2,3 +2,5 @@ defaults: grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e.yaml data_plane: enabled: true backend: mooncake_cpu + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml index 89cd90aceb0..ed4d08776d6 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml @@ -1,3 +1,5 @@ defaults: grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3.yaml data_plane: enabled: true + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml index 15ebb1943f7..f6753cdd5c9 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml @@ -2,3 +2,5 @@ defaults: grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50.yaml data_plane: enabled: true backend: mooncake_cpu + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml index 9406c66c578..16f630c3080 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml @@ -2,3 +2,5 @@ defaults: grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50.yaml data_plane: enabled: true backend: mooncake_cpu + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml index dbca800e5ca..c57b60435e1 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml @@ -1,3 +1,5 @@ defaults: grpo-llama3.2-1b-instruct-1n8g-megatron.yaml data_plane: enabled: true + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml index 10f4ba92a1d..e8ca7538473 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml @@ -2,3 +2,5 @@ defaults: grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml data_plane: enabled: true backend: mooncake_cpu + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml b/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml index 26729ccd7c2..ee35c4e86d3 100644 --- a/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml @@ -1,3 +1,5 @@ defaults: grpo-moonlight-16ba3b-4n8g-megatron.yaml data_plane: enabled: true + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml index 6fbd7218c16..ded84294a79 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml @@ -2,3 +2,5 @@ defaults: grpo-nanov3-30BA3B-1n8g-fsdp2.v2.yaml data_plane: enabled: true backend: mooncake_cpu + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml index 86be2c983f3..7d70884b861 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml @@ -1,3 +1,5 @@ defaults: grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.yaml data_plane: enabled: true + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml index 9d25544bcd7..c6eb7880616 100644 --- a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml @@ -1,3 +1,5 @@ defaults: grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3.yaml data_plane: enabled: true + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml index 08f148f3709..ac9be52b94b 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml @@ -2,3 +2,5 @@ defaults: grpo-qwen3-1.7b-1n8g-megatron-eagle3.yaml data_plane: enabled: true backend: mooncake_cpu + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml index f48450c8ef2..9041b9c7543 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml @@ -7,3 +7,5 @@ logger: log_dir: logs/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple wandb: name: grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml index 952f192a362..0f889e39c00 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml @@ -2,3 +2,5 @@ defaults: grpo-qwen3-8B-base-1n8g-fsdp2-lora.yaml data_plane: enabled: true backend: mooncake_cpu + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml b/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml index 839a96bb56b..2fd278c8163 100644 --- a/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml @@ -2,3 +2,5 @@ defaults: prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2.yaml data_plane: enabled: true backend: mooncake_cpu + observability: + verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index da16de450a3..fc5468ee725 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -395,19 +395,33 @@ def _log_breakdown( ) -> None: """Log the per-op breakdown as a table beside the series. - Best effort: a backend without a table type skips it, and a failure here - must not take a step down over a visualisation. + A backend without a table type has no rows to log. Failures propagate to + the caller, which already guarantees the logging path cannot fail a step. + """ + columns, rows = breakdown_table(metrics) + if rows: + logger.log_table(columns, rows, step, name) + + +def _log_data_plane_metrics( + policy: Any, logger: Logger, step: int, total_step_time: float +) -> None: + """Log this step's data-plane cost. Never raises. + + On by default, so this runs every step of every recipe. """ try: - columns, rows = breakdown_table(metrics) + _log_data_plane_metrics_impl(policy, logger, step, total_step_time) except Exception as exc: # noqa: BLE001 - a panel must never fail a step - logging.getLogger(__name__).warning("data-plane breakdown failed: %s", exc) - else: - if rows: - logger.log_table(columns, rows, step, name) + logging.getLogger(__name__).warning( + "data-plane metrics failed at step %d (%s: %s); training continues", + step, + type(exc).__name__, + exc, + ) -def _log_data_plane_metrics( +def _log_data_plane_metrics_impl( policy: Any, logger: Logger, step: int, total_step_time: float ) -> None: """Log this step's data-plane cost. No-op unless observability is enabled. diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index dc6dafee9cf..870619c406a 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1524,6 +1524,22 @@ async def _cleanup_consumed_metas_unlocked( if errors: raise BaseExceptionGroup("post-train DataPlane cleanup failed", errors) def _log_data_plane_metrics(self, total_step_time: float) -> None: + """Log this step's data-plane cost. Never raises. + + On by default, so this runs every step of every recipe. Mirrors + ``grpo_sync._log_data_plane_metrics``. + """ + try: + self._log_data_plane_metrics_impl(total_step_time) + except Exception as exc: # noqa: BLE001 - a panel must never fail a step + logging.getLogger(__name__).warning( + "data-plane metrics failed at step %d (%s: %s); training continues", + self._train_steps, + type(exc).__name__, + exc, + ) + + def _log_data_plane_metrics_impl(self, total_step_time: float) -> None: """Log this step's data-plane cost. No-op unless observability is enabled. The synchronous loop logs these series from ``_log_data_plane_metrics`` @@ -1550,15 +1566,9 @@ def _log_data_plane_metrics(self, total_step_time: float) -> None: self._logger.log_metrics( headline_series(metrics), step, prefix="data_plane/driver" ) - try: - columns, rows = breakdown_table(metrics) - except Exception as exc: # noqa: BLE001 - a panel must never fail a step - logging.getLogger(__name__).warning("data-plane breakdown failed: %s", exc) - else: - if rows: - self._logger.log_table( - columns, rows, step, "data_plane/driver/breakdown" - ) + columns, rows = breakdown_table(metrics) + if rows: + self._logger.log_table(columns, rows, step, "data_plane/driver/breakdown") print( f" • data plane: {metrics['step/wall_ms']:.0f}ms, " f"{metrics['step/comm_volume_mb']:.1f} MB moved", diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 9d932e77f33..b235c9c6518 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -549,7 +549,9 @@ def _hash_deltas(hv: dict[str, int], prev_hv: dict[str, int]) -> dict[str, float Returns: ``step/hash/{counter}`` deltas, or ``{}`` when the guard never ran. """ - if not hv or not hv.get("rows_recorded"): + # ``guard_failures`` counts too: a guard that raised on the first put + # records no rows, and gating on rows alone would make it look switched off. + if not hv or not (hv.get("rows_recorded") or hv.get("guard_failures")): return {} deltas: dict[str, float] = { f"step/hash/{name}": hv[name] - prev_hv.get(name, 0) @@ -559,6 +561,7 @@ def _hash_deltas(hv: dict[str, int], prev_hv: dict[str, int]) -> dict[str, float "rows_unverified", "mismatches", "fields_skipped", + "guard_failures", ) } # Corruption of every row of every field in a step, repeated identically, @@ -823,6 +826,7 @@ def merge_snapshots(snapshots: "list[dict[str, Any]]") -> dict[str, Any]: "rows_unverified", "mismatches", "fields_skipped", + "guard_failures", ) } by_op: dict[str, dict[str, Any]] = {} @@ -1097,6 +1101,11 @@ class HashStats: # uniform row shape) and leaves whose leading dim doesn't match the # sample count, so a row cannot be attributed to a sample id. fields_skipped: int = 0 + # Batches the guard raised on, and so never checked. Same "reads as clean + # because it checked nothing" hazard as ``fields_skipped``, counted for + # the same reason. Not ``errors``: ``OpStats.errors`` already means failed + # transfers, and these are failures of the check, not of the wire. + guard_failures: int = 0 @dataclass @@ -1427,8 +1436,39 @@ def _row_fingerprints( ) return out + def _hash_guard_failed(self, op: str, exc: Exception) -> None: + """Absorb a hash-guard failure: count it, log it, never re-raise. + + The guard is a debug aid on a transfer that already succeeded, so a + bug in it must not take the transfer down. Swallowing is only safe + because the failure stays visible in ``step/hash/guard_failures`` -- a + guard that silently stopped checking would report zero mismatches. + """ + self._stats.hash_verify.guard_failures += 1 + # Logged once, not capped at a handful: whatever makes the guard raise + # on one batch makes it raise on every batch, so line two would carry + # nothing line one did not. The count is the series. + if self._stats.hash_verify.guard_failures == 1: + logger.warning( + "data-plane hash guard failed on %s (%s: %s). The transfer " + "itself is unaffected, but this batch went unchecked -- see " + "step/hash/guard_failures for how many.", + op, + type(exc).__name__, + exc, + ) + def _record_hashes( self, partition_id: str, sample_ids: list[str], fields: TensorDict | None + ) -> None: + """Store wire-in fingerprints for a successful put. Never raises.""" + try: + self._record_hashes_impl(partition_id, sample_ids, fields) + except Exception as exc: # noqa: BLE001 - a debug check must never fail a transfer + self._hash_guard_failed("put", exc) + + def _record_hashes_impl( + self, partition_id: str, sample_ids: list[str], fields: TensorDict | None ) -> None: """Store wire-in fingerprints for a successful put.""" digests = self._row_fingerprints(fields, sample_ids) @@ -1456,6 +1496,15 @@ def _record_hashes( self._stats.hash_verify.rows_recorded += len(sample_ids) def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: + """Compare wire-out fingerprints against what was written. Never raises.""" + try: + self._check_hashes_impl(partition_id, sample_ids, out) + except Exception as exc: # noqa: BLE001 - a debug check must never fail a transfer + self._hash_guard_failed("get", exc) + + def _check_hashes_impl( + self, partition_id: str, sample_ids: list[str], out: Any + ) -> None: """Compare wire-out fingerprints against what was written.""" if not isinstance(out, TensorDict): return diff --git a/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh b/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh index 375e3677b7b..9a7da2267c8 100755 --- a/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh +++ b/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh b/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh index aa3f053e87b..906929565bc 100755 --- a/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh +++ b/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh b/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh index 0c3f6184a52..48c0c49b3d4 100755 --- a/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh +++ b/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh index c4635e9f219..9ea9b4c13b3 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh index b0880e3c890..e7695ab284f 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh index 2bf6ad97548..44a8e722a57 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh index 4b0de6414f6..9bc62ea48d4 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh index db93cbfab37..bc6ce1d37ce 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh index a67323016fd..85ae3077be2 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh index 14a128672fe..ee539083555 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh index 7c765c5bd22..a7ca901f759 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh b/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh index 606dccb7325..cf78ea15e54 100755 --- a/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh +++ b/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh index 9ff4e574237..464c0808d53 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh @@ -16,3 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh index 56850083541..3ac3167b7fd 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh @@ -16,3 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh index 9d905ef649a..f4124f94095 100755 --- a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh b/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh index 121013805bb..a4ad7f506d5 100755 --- a/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh @@ -16,3 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh index 490f00fb665..d44db1a1d77 100755 --- a/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh @@ -34,8 +34,11 @@ uv run examples/run_grpo.py \ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + # The wire guard only counts; assert it found nothing. uv run tests/check_metrics.py $JSON_METRICS \ - 'median(data["train/token_mult_prob_error"]) < 1.02' + 'median(data["train/token_mult_prob_error"]) < 1.02' \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' rm -rf "$CKPT_DIR" fi diff --git a/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh b/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh index b6c3f072b70..0b55ca1aafe 100755 --- a/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh b/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh index 5e07f0ef476..db95defea81 100755 --- a/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh +++ b/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh @@ -15,3 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# The wire guard only counts; assert it found nothing. +cd "$SCRIPT_DIR/../../.." +uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ + 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ + 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 901f986b65f..13db64b4503 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -6098,3 +6098,111 @@ def test_train_fields_for_step(skip_prev_logprobs, expect_prev): ) def test_needs_hf_refit_handshake(backend, nccl_reshard, colocated, expected): assert _needs_hf_refit_handshake(backend, nccl_reshard, colocated) is expected + + +@pytest.mark.parametrize( + ("n_snapshots", "scope"), + [ + (1, "driver"), # fan-out reached one process -> the driver's own counters + (2, "cluster"), # it reached the workers -> the summed cluster view + ], +) +def test_grpo_train_sync_logs_data_plane_metrics_before_committing_the_step( + mock_grpo_components, n_snapshots, scope +): + """The data-plane series reach the logger, in the right scope, before the + commit that would drop them. + + ``data_plane.observability.enabled`` defaults to true in + ``grpo_math_1B.yaml``, which every recipe inherits, so this path runs on + every sync step of every run. Three things can silently disable it, and none + is reachable from the observability unit tests because those never build a + trainer: + + * ``policy.dp_client`` not being a ``MetricsDataPlaneClient`` -- the + isinstance guard in ``_log_data_plane_metrics`` then returns early and the + whole feature is a no-op that logs nothing and raises nothing; + * the wrong scope being chosen, so the driver's one-op-per-step counters get + reported as if they were the cluster's bulk traffic, or vice versa; + * the call landing after ``log_metrics(..., step_finished=True)`` -- wandb + drops anything logged against an already-committed step, so every series + is computed, printed to stdout, and discarded. + """ + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient + from nemo_rl.data_plane.observability import MetricsDataPlaneClient + + policy = mock_grpo_components["policy"] + client = MetricsDataPlaneClient(NoOpDataPlaneClient()) + client.register_partition( + partition_id="p", fields=["x"], num_samples=2, consumer_tasks=["t"] + ) + client.put_samples( + sample_ids=["a", "b"], + partition_id="p", + fields=TensorDict({"x": torch.zeros(2, 64)}, batch_size=[2]), + ) + policy.dp_client = client + policy.collect_data_plane_snapshots = MagicMock( + return_value=[client.snapshot() for _ in range(n_snapshots)] + ) + # A real policy has no such attribute on the first step, so the production + # ``getattr(policy, "_prev_cluster_snapshot", {})`` yields {}. A MagicMock + # would auto-create one and hand arithmetic a mock instead of a dict. + policy._prev_cluster_snapshot = {} + + master_config = mock_grpo_components["master_config"] + master_config.data_plane = {"enabled": True} + master_config.grpo.max_num_steps = 1 + master_config.grpo.val_period = 0 + master_config.grpo.val_at_start = False + master_config.grpo.val_at_end = False + master_config.grpo.use_dynamic_sampling = False + + with ExitStack() as stack: + stack.enter_context(mock_sync_grpo_infrastructure(policy)) + stack.enter_context( + patch("nemo_rl.algorithms.grpo_sync.validate_sync", return_value=({}, {})) + ) + grpo_train_sync( + policy, + _mock_policy_generation(), + mock_grpo_components["train_dataloader"], + mock_grpo_components["val_dataloader"], + mock_grpo_components["tokenizer"], + mock_grpo_components["loss_fn"], + mock_grpo_components["task_to_env"], + mock_grpo_components["val_task_to_env"], + mock_grpo_components["logger"], + mock_grpo_components["checkpointer"], + _initial_grpo_save_state(), + master_config, + ) + + calls = mock_grpo_components["logger"].log_metrics.call_args_list + want = f"data_plane/{scope}" + dp = [i for i, c in enumerate(calls) if c.kwargs.get("prefix") == want] + commit = [i for i, c in enumerate(calls) if c.kwargs.get("step_finished")] + + assert dp, ( + f"no data_plane/{scope} series logged; prefixes seen: " + f"{[c.kwargs.get('prefix') for c in calls]}" + ) + assert commit, "the step was never committed" + assert dp[0] < commit[0], ( + "data-plane metrics logged after the committing log_metrics call; " + "wandb drops anything logged against an already-committed step" + ) + + payload = calls[dp[0]].args[0] + assert payload["step/comm_volume_mb"] > 0, "the put moved bytes; the series says 0" + for key in ("step/wall_ms", "step/frac_of_step", "step/self/overhead_ms"): + assert key in payload, f"{key} missing from {sorted(payload)}" + + assert [ + c + for c in mock_grpo_components["logger"].log_table.call_args_list + if f"data_plane/{scope}/breakdown" in c.args + ], "the per-op breakdown table was not logged" + client.close() diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index ede4f4bd5c0..cebb942bc52 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -22,6 +22,7 @@ from __future__ import annotations import logging +import random from time import monotonic import pytest @@ -41,29 +42,174 @@ merge_snapshots, ) -from ._rollout_shapes import make_rollout_batch +# ── helpers ──────────────────────────────────────────────────────────── + + +def _ids(n, prefix="u"): + return [f"{prefix}{i}" for i in range(n)] + + +def _client(inner=None, *, register=True, **kwargs): + """A wrapped no-op client, with partition ``p`` registered by default. + + ``NoOpDataPlaneClient`` validates neither ``num_samples`` nor ``fields``, + so one registration serves every caller. Pass ``register=False`` in tests + that synthesise calls with :func:`_emit`: a real register call would add + an op to ``by_op`` and skew the per-op shares under test. + """ + client = MetricsDataPlaneClient(inner or NoOpDataPlaneClient(), **kwargs) + if register: + client.register_partition( + partition_id="p", + fields=["x", "ids", "lp"], + num_samples=5_000, + consumer_tasks=["t"], + ) + return client + + +def _put(client, ids, width=4): + """Put ``len(ids)`` rows of ``width`` float32 columns; return the bytes billed.""" + client.put_samples( + sample_ids=list(ids), + partition_id="p", + fields=TensorDict({"x": torch.zeros(len(ids), width)}, batch_size=[len(ids)]), + ) + return len(ids) * width * 4 + + +def _emit(client, latencies_ms, op="put", n_bytes=1_000): + """Record one synthetic ``op`` call per entry in ``latencies_ms``. + + ``n_bytes`` takes an int, or a callable of the call index for the size + variation the latency/bandwidth fit needs to be identifiable. + """ + now = monotonic() + for i, ms in enumerate(latencies_ms): + size = n_bytes(i) if callable(n_bytes) else n_bytes + client._emit(op, "p", 1, size, now - ms / 1e3, "ok") + return client + + +def _rank(latencies_ms, op="put", n_bytes=1_000): + """One finished rank's snapshot: a call per entry, at that latency.""" + client = _client(register=False) + try: + return _emit(client, latencies_ms, op=op, n_bytes=n_bytes).snapshot() + finally: + client.close() + + +def _busy(op_calls): + """A client that ran ``n`` calls of each named op, at ``ms`` each.""" + client = _client(register=False) + for op, (n, ms) in op_calls.items(): + _emit(client, [ms] * n, op=op, n_bytes=lambda i: 1_000_000 * (1 + i % 5)) + return client + + +def _jagged(rows, field="x"): + return TensorDict( + {field: torch.nested.nested_tensor(rows, layout=torch.jagged)}, + batch_size=[len(rows)], + ) + + +def _hash_fields(n=4): + return TensorDict( + { + "ids": torch.arange(n * 6, dtype=torch.int64).reshape(n, 6), + "lp": torch.linspace(0, 1, n * 6, dtype=torch.bfloat16).reshape(n, 6), + }, + batch_size=[n], + ) + + +def _jagged_ids(lengths, seed=0, with_dense=False): + """Rows of pseudorandom token ids, optionally beside a uniform ``lp`` field. + + Deliberately not ``arange``: ``hash_tensor`` is an XOR reduction, and + aligned runs of consecutive integers collide under it — ``XOR(6..11)`` + and ``XOR(12..17)`` are both 1 — which would make two visibly different + rows fingerprint the same. + """ + g = torch.Generator().manual_seed(seed) + fields = { + "ids": torch.nested.nested_tensor( + [torch.randint(0, 32000, (n,), generator=g) for n in lengths], + layout=torch.jagged, + ) + } + if with_dense: + fields["lp"] = torch.zeros(len(lengths), 6) + return TensorDict(fields, batch_size=[len(lengths)]) + + +class _CorruptingClient(NoOpDataPlaneClient): + """Flips one element of one field on read — a stand-in for a wire bug.""" + + def __init__(self, field: str, row: int) -> None: + super().__init__() + self._corrupt_field = field + self._corrupt_row = row + + def get_samples(self, sample_ids, partition_id, select_fields): + out = super().get_samples(sample_ids, partition_id, select_fields) + if self._corrupt_field in out.keys(): + out[self._corrupt_field][self._corrupt_row] += 1 + return out + + +class _JaggedEcho(NoOpDataPlaneClient): + """Returns whatever was put, jagged, so row lengths survive the trip.""" + + def __init__(self) -> None: + super().__init__() + self.rows: dict[tuple[str, str], dict[str, torch.Tensor]] = {} + + def put_samples(self, sample_ids, partition_id, fields=None, tags=None): + for key in fields.keys(): + v = fields.get(key) + rows = v.unbind() if v.is_nested else list(v) + for sid, row in zip(sample_ids, rows): + self.rows.setdefault((partition_id, sid), {})[str(key)] = row.clone() + return super().put_samples( + sample_ids=sample_ids, partition_id=partition_id, fields=fields, tags=tags + ) + + def get_samples(self, sample_ids, partition_id, select_fields): + out = {} + for f in select_fields: + rows = [self.rows[(partition_id, sid)][f] for sid in sample_ids] + out[f] = ( + torch.stack(rows) + if all(r.shape == rows[0].shape for r in rows[1:]) + else torch.nested.nested_tensor(rows, layout=torch.jagged) + ) + return TensorDict(out, batch_size=[len(sample_ids)]) @pytest.fixture def wrapped_client(): + """A registered client plus the list of events it emitted.""" events: list[dict] = [] - inner = NoOpDataPlaneClient() - client = MetricsDataPlaneClient(inner, on_event=events.append) + client = _client(on_event=events.append) yield client, events - inner.close() + client.close() + + +# ── the wrapper's event stream ───────────────────────────────────────── def test_put_records_bytes_and_count(wrapped_client): client, events = wrapped_client - client.register_partition( - partition_id="p", fields=["x"], num_samples=4, consumer_tasks=["read"] + client.put_samples( + sample_ids=_ids(4, "a"), + partition_id="p", + fields=TensorDict({"x": torch.zeros(4, dtype=torch.float32)}, batch_size=[4]), ) - fields = TensorDict({"x": torch.zeros(4, dtype=torch.float32)}, batch_size=[4]) - client.put_samples(sample_ids=["a", "b", "c", "d"], partition_id="p", fields=fields) - put_events = [e for e in events if e["op"] == "put"] - assert len(put_events) == 1 - e = put_events[0] + (e,) = [e for e in events if e["op"] == "put"] assert e["status"] == "ok" assert e["n_keys"] == 4 assert e["n_bytes"] == 16 # 4 floats * 4 bytes @@ -72,9 +218,6 @@ def test_put_records_bytes_and_count(wrapped_client): def test_get_records_after_put(wrapped_client): client, events = wrapped_client - client.register_partition( - partition_id="p", fields=["x"], num_samples=2, consumer_tasks=["read"] - ) client.put_samples( sample_ids=["a", "b"], partition_id="p", @@ -83,23 +226,17 @@ def test_get_records_after_put(wrapped_client): out = client.get_samples( sample_ids=["a", "b"], partition_id="p", select_fields=["x"] ) - assert torch.equal(out["x"], torch.ones(2)) - get_events = [e for e in events if e["op"] == "get"] - assert len(get_events) == 1 - assert get_events[0]["n_bytes"] > 0 + assert torch.equal(out["x"], torch.ones(2)) + (e,) = [e for e in events if e["op"] == "get"] + assert e["n_bytes"] > 0 def test_register_and_clear_recorded(wrapped_client): client, events = wrapped_client - client.register_partition( - partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["r"] - ) client.clear_samples(sample_ids=None, partition_id="p") - ops = [e["op"] for e in events] - assert ops.count("register") == 1 - assert ops.count("clear") == 1 + assert [e["op"] for e in events] == ["register", "clear"] def test_list_sample_ids_is_forwarded_and_recorded(wrapped_client): @@ -119,44 +256,26 @@ def test_list_sample_ids_is_forwarded_and_recorded(wrapped_client): def test_error_status_recorded_and_reraised(wrapped_client): - """Decorator does NOT swallow errors — re-raise after recording.""" + """The wrapper records the failure and re-raises rather than swallowing it.""" client, events = wrapped_client with pytest.raises(KeyError): client.get_samples(sample_ids=["a"], partition_id="nope", select_fields=["x"]) - err = [e for e in events if e["op"] == "get" and e["status"] == "error"] - assert len(err) == 1 + assert [(e["op"], e["status"]) for e in events if e["op"] == "get"] == [ + ("get", "error") + ] -def test_snapshot_accumulates_successful_ops(wrapped_client): - client, _ = wrapped_client - client.register_partition( - partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["r"] - ) - client.put_samples( - sample_ids=["a"], - partition_id="p", - fields=TensorDict({"x": torch.zeros(1)}, batch_size=[1]), - ) +def test_no_callback_still_accumulates_stats(): + """``on_event=None`` skips building the event dict; the counters that + ``snapshot()`` reports must not depend on a sink being registered.""" + client = _client() + expected = _put(client, ["a", "b"], width=3) snap = client.snapshot() - assert snap["total_ops"] >= 2 # register + put - assert snap["total_bytes"] >= 4 # 1 float = 4 bytes - - -def test_default_callback_is_noop(): - """Omitting on_event must not raise; the wrapper just forwards.""" - inner = NoOpDataPlaneClient() - client = MetricsDataPlaneClient(inner) - client.register_partition( - partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["r"] - ) - client.close() - -def test_close_propagates(wrapped_client): - client, _ = wrapped_client - client.close() - # Second close must not raise — NoOp is idempotent. + assert snap["total_bytes"] == expected + assert snap["by_op"]["put"]["calls"] == 1 + assert snap["total_wall_ms"] > 0 client.close() @@ -196,68 +315,9 @@ def test_checkpoint_lifecycle_is_forwarded_and_recorded(tmp_path) -> None: restored.close() -def test_factory_wraps_when_observability_enabled(): - """Programmatic wrap path; factory.py uses the same MetricsDataPlaneClient.""" - inner = NoOpDataPlaneClient() - seen: list[dict] = [] - client = MetricsDataPlaneClient(inner, on_event=seen.append) - assert hasattr(client, "snapshot") - client.register_partition( - partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["r"] - ) - assert len(seen) == 1 and seen[0]["op"] == "register" - client.close() - - -def test_observability_records_realistic_rollout_put() -> None: - """Metrics middleware records put-bytes correctly when the put carries a - realistic rollout-shaped batch (bf16 logprobs, int32 masks, int64 ids).""" - - inner = NoOpDataPlaneClient() - seen: list[dict] = [] - client = MetricsDataPlaneClient(inner, on_event=seen.append) - - n = 4 - batch = make_rollout_batch(n=n, max_seqlen=64, seed=71) - client.register_partition( - partition_id="train", - fields=["input_ids", "input_lengths", "generation_logprobs"], - num_samples=n, - consumer_tasks=["train"], - ) - fields = TensorDict( - { - "input_ids": batch["input_ids"], - "input_lengths": batch["input_lengths"], - "generation_logprobs": batch["generation_logprobs"], - }, - batch_size=[n], - ) - client.put_samples( - sample_ids=[f"u{i}" for i in range(n)], - partition_id="train", - fields=fields, - ) - - put_events = [e for e in seen if e["op"] == "put"] - assert len(put_events) == 1 - # Bytes should reflect bf16 logprobs (2 bytes/elem) + int64 ids (8 bytes/elem), - # not a fixed-dtype assumption. Lower bound: at least one full int64 batch. - min_expected = n * 64 * 8 # input_ids alone - assert put_events[0]["n_bytes"] >= min_expected - client.close() - - # ── byte accounting ──────────────────────────────────────────────────── -def _jagged(rows): - return TensorDict( - {"x": torch.nested.nested_tensor(rows, layout=torch.jagged)}, - batch_size=[len(rows)], - ) - - @pytest.mark.parametrize( "name,td,expected", [ @@ -335,10 +395,11 @@ def test_td_bytes_does_not_overcount_a_narrow_view(): def test_td_bytes_nontensordata_is_not_broadcast(): """``NonTensorData`` holds ONE object; counting it per batch row would - inflate a 64-row put by 64x. Its bytes must not scale with batch size.""" + inflate a 64-row put by 64x.""" payload = {"tool": "bash", "text": "x" * 100} small = TensorDict({"m": NonTensorData(payload, batch_size=[2])}, batch_size=[2]) large = TensorDict({"m": NonTensorData(payload, batch_size=[64])}, batch_size=[64]) + assert _td_bytes(small) == _td_bytes(large) assert _td_bytes(small) >= 100 # the string itself is still counted @@ -347,14 +408,16 @@ def test_td_bytes_nontensorstack_scales_with_rows(): """``NonTensorStack`` genuinely holds one object per row, so its estimate must scale — and stay close to the exact walk it extrapolates from.""" row = {"turns": ["hello"] * 4, "n": 3} - stack_8 = NonTensorStack(*[NonTensorData(dict(row)) for _ in range(8)]) - stack_64 = NonTensorStack(*[NonTensorData(dict(row)) for _ in range(64)]) - bytes_8 = _td_bytes(TensorDict({"s": stack_8}, batch_size=[8])) - bytes_64 = _td_bytes(TensorDict({"s": stack_64}, batch_size=[64])) + + def stack_bytes(n): + stack = NonTensorStack(*[NonTensorData(dict(row)) for _ in range(n)]) + return _td_bytes(TensorDict({"s": stack}, batch_size=[n])) + + bytes_8, bytes_64 = stack_bytes(8), stack_bytes(64) + exact = sum(_estimate_encoded_bytes(dict(row), [10_000]) for _ in range(64)) + assert bytes_8 > 0 assert bytes_64 == pytest.approx(8 * bytes_8, rel=0.05) - # And it agrees with summing every row explicitly. - exact = sum(_estimate_encoded_bytes(dict(row), [10_000]) for _ in range(64)) assert bytes_64 == pytest.approx(exact, rel=0.05) @@ -362,59 +425,65 @@ def test_estimate_encoded_bytes_walk_is_bounded(): """The node budget caps the walk so one pathological payload cannot make a put O(payload size).""" huge = {"k": list(range(100_000))} - bounded = _estimate_encoded_bytes(huge, [64]) - unbounded = _estimate_encoded_bytes(huge, [10_000_000]) - assert bounded < unbounded - assert bounded <= 4 * 64 # ≤2 leaves per budget unit, ≤2 bytes each here - - -def test_outstanding_bytes_reconcile_exactly(): - """Put then clear must return ``bytes_outstanding`` to zero: the per-key - split drops its division remainder on one key rather than spreading it, - so the total has to be preserved for the accounting to close.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - ids = [f"u{i}" for i in range(7)] # 7 keys => non-zero remainder - fields = TensorDict({"x": torch.zeros(7, 5)}, batch_size=[7]) - client.register_partition( - partition_id="p", fields=["x"], num_samples=7, consumer_tasks=["t"] + + assert _estimate_encoded_bytes(huge, [64]) < _estimate_encoded_bytes( + huge, [10_000_000] ) - client.put_samples(sample_ids=ids, partition_id="p", fields=fields) - assert client.snapshot()["bytes_outstanding"] == 7 * 5 * 4 - client.clear_samples(sample_ids=ids, partition_id="p") + assert _estimate_encoded_bytes(huge, [64]) <= 4 * 64 # <=2 leaves/unit, <=2B each + + +def test_clear_frees_only_what_was_actually_live(): + """A clear may name uids already dropped, or belonging to another + partition. Billing those releases bytes this partition never held: + clearing 50 live keys alongside 50 unknown ones freed two thirds of a + partition that had lost half its keys.""" + client = _client() + ids = _ids(100) + total = _put(client, ids, width=250) + assert client.snapshot()["bytes_outstanding"] == total, "the put is billed in full" + + client.clear_samples(sample_ids=ids[:50] + _ids(50, "unknown"), partition_id="p") + assert client.snapshot()["bytes_outstanding"] == total // 2 + + client.clear_samples(sample_ids=ids[50:], partition_id="p") assert client.snapshot()["bytes_outstanding"] == 0 client.close() -def test_step_metrics_use_one_unit_per_dimension(): - """Every duration is ms, every volume MB. A chart mixing `wall_s` with - `p90_ms` shows 0.008 beside 24.85 and reads as a data-plane bug rather - than an axis one.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client.register_partition( - partition_id="p", fields=["x"], num_samples=2, consumer_tasks=["t"] - ) - client.put_samples( - sample_ids=["a", "b"], - partition_id="p", - fields=TensorDict({"x": torch.zeros(2, 3)}, batch_size=[2]), - ) - keys = set(client.get_step_metrics(1.0)) - assert not [k for k in keys if k.endswith(("_s", "_gb", "_kb", "_us"))], keys - assert "step/wall_ms" in keys and "step/comm_volume_mb" in keys +@pytest.mark.parametrize("seed", range(8)) +def test_outstanding_reconciles_over_random_put_clear_sequences(seed): + """Interleaved puts and partial clears must always land back at zero; the + pro-rata release drops its division remainder, so only clearing the last + live key can settle the account.""" + rng = random.Random(seed) + client = _client() + live: set[str] = set() + for _ in range(rng.randint(1, 6)): + batch = list( + dict.fromkeys(f"k{rng.randint(0, 60)}" for _ in range(rng.randint(1, 20))) + ) + _put(client, batch, width=250) + live |= set(batch) + if live and rng.random() < 0.5: + drop = rng.sample(sorted(live), k=rng.randint(1, len(live))) + client.clear_samples(sample_ids=drop, partition_id="p") + live -= set(drop) + if live: + client.clear_samples(sample_ids=sorted(live), partition_id="p") + + assert client.snapshot()["bytes_outstanding"] == 0 client.close() +# ── per-step series: units, windows, and the latency split ───────────── + + def test_step_metrics_tail_is_exact_not_bucketed(): - """Per-step percentiles came off a histogram that is never reset, so - they went flat and quantised to bucket edges (a tail quantile of a - single sample in (10, 25] always lands on the same interpolated - point). ``max_ms`` is exact and - tracks the slowest call actually seen.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client.register_partition( - partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["t"] - ) - client._emit("put", "p", 1, 8, monotonic() - 0.030, "ok") # a 30 ms call + """Per-step percentiles came off a histogram that is never reset, so they + went flat and quantised to bucket edges. ``max_ms`` is exact, and one call + supports no percentile at all — in either view.""" + client = _client(register=False) + _emit(client, [30.0], n_bytes=8) metrics = client.get_step_metrics(1.0) assert "put/p90_ms" not in metrics and "step/by_op/put/p50_ms" not in metrics @@ -422,151 +491,437 @@ def test_step_metrics_tail_is_exact_not_bucketed(): assert metrics["step/by_op/put/max_ms"] != pytest.approx(24.85, abs=0.5), ( "bucket edge" ) - # and one call supports no percentile at all, in either view assert "p90_ms" not in client.snapshot()["by_op"]["put"] client.close() -def test_latency_breakdown_stacks_to_wall_ms(): - """The fit is reported as two ms components rather than a ratio, so a - chart can stack them against the measured ``wall_ms``. - - A ratio would have been both flat (the fit is cumulative) and unitless - on an axis of milliseconds. These carry the coefficients from the - cumulative fit but attribute them to *this* step's calls and bytes. - """ - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - fixed_ms, mb_per_s = 8.0, 500.0 - now = monotonic() - for i in range(12): # varied sizes, or the fit is unidentifiable - n_bytes = 50_000 * (i + 1) - wall_ms = fixed_ms + n_bytes / (mb_per_s * 1e3) - client._emit("put", "p", 1, n_bytes, now - wall_ms / 1e3, "ok") - - metrics = client.get_step_metrics(1.0) - fit = client.snapshot()["by_op"]["put"]["fit"] - assert fit["model_trustworthy"], fit - assert fit["fixed_ms"] == pytest.approx(fixed_ms, rel=0.05) - assert fit["bandwidth_mb_s"] == pytest.approx(mb_per_s, rel=0.05) +def test_snapshot_leaves_the_step_window_alone_unless_asked(): + """``snapshot()`` is also how a human inspects a live client. Resetting the + step window on every call would let an inspection blank the next step.""" + client = _client(register=False) + _emit(client, [30.0]) - # the two components are the split of ONE call, and they add to the mean - total = ( - metrics["step/by_op/put/overhead_ms"] + metrics["step/by_op/put/transfer_ms"] - ) - assert total == pytest.approx(metrics["step/by_op/put/mean_ms"], rel=0.05) - # per call, the overhead term IS the fitted constant - assert metrics["step/by_op/put/overhead_ms"] == pytest.approx(fixed_ms, rel=0.05) - assert "step/by_op/put/overhead_frac" not in metrics, ( - "a ratio is derivable from these" - ) + assert client.snapshot()["by_op"]["put"]["step_max_ms"] >= 30.0 + assert client.snapshot()["by_op"]["put"]["step_max_ms"] >= 30.0, "still there" + assert client.snapshot(reset_step_window=True)["by_op"]["put"]["step_max_ms"] >= 30 + assert client.snapshot()["by_op"]["put"]["step_max_ms"] == 0.0, "window reopened" client.close() -def test_step_max_is_scoped_to_the_step(): - """A lifetime max is monotonic and goes flat the moment the worst call - has been seen — the same defect as logging a cumulative percentile. The - reported max must fall again when a step is quicker.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client._emit("put", "p", 1, 8, monotonic() - 0.050, "ok") # slow step - slow = client.get_step_metrics(1.0)["step/by_op/put/max_ms"] - client._emit("put", "p", 1, 8, monotonic() - 0.001, "ok") # quick step - quick = client.get_step_metrics(1.0)["step/by_op/put/max_ms"] - - assert slow >= 50.0 - assert quick < slow, "step max must reset, not carry the lifetime worst" - # the lifetime worst is still available for a one-off look - assert client.snapshot()["by_op"]["put"]["max_ms"] >= 50.0 +def test_cluster_step_max_reopens_each_step(): + """A maximum cannot be differenced out of a cumulative counter, so the + cluster path reported the lifetime max: after one 50 ms call every later + step still read 50 ms. The reader resets the window as it reads.""" + client = _client(register=False) + prev, seen = {}, [] + for slowest_ms in (5.0, 50.0, 5.0, 5.0): + _emit(client, [slowest_ms, 5.0, 5.0, 5.0]) + merged = merge_snapshots([client.snapshot(reset_step_window=True)]) + seen.append(cluster_step_metrics(merged, prev, 1.0)["step/by_op/put/max_ms"]) + prev = merged + + assert seen[1] == pytest.approx(50.0, abs=1.0), "the spike shows" + assert seen[2] == pytest.approx(5.0, abs=1.0), "and does not latch" client.close() -def test_no_callback_still_accumulates_stats(): - """``on_event=None`` skips building the event dict; the counters that - ``snapshot()`` reports must not depend on a sink being registered.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client.register_partition( - partition_id="p", fields=["x"], num_samples=2, consumer_tasks=["t"] - ) - client.put_samples( - sample_ids=["a", "b"], - partition_id="p", - fields=TensorDict({"x": torch.zeros(2, 3)}, batch_size=[2]), +def test_cluster_view_carries_the_latency_split(): + """The split was emitted on the driver path only, so the cluster view — the + one a real run logs — could never show it and its table column was always + empty. The cluster fit is the better one besides: more samples and more + size variation, and size variation is what makes the split identifiable.""" + fixed_ms, mb_per_s = 6.0, 400.0 + ranks = [] + for rank in range(4): + sizes = [100_000 * (i + 1 + rank) for i in range(8)] + ranks.append( + _rank( + [fixed_ms + b / (mb_per_s * 1e3) for b in sizes], + op="get", + n_bytes=lambda i, s=sizes: s[i], + ) + ) + metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 1.0) + + assert "step/by_op/get/overhead_ms" in metrics, "cluster view must carry the split" + total = ( + metrics["step/by_op/get/overhead_ms"] + metrics["step/by_op/get/transfer_ms"] ) - snap = client.snapshot() - assert snap["total_bytes"] == 2 * 3 * 4 - assert snap["by_op"]["put"]["calls"] == 1 - assert snap["total_wall_ms"] > 0 - client.close() + assert total == pytest.approx(metrics["step/by_op/get/mean_ms"], rel=0.05) + assert metrics["step/by_op/get/overhead_ms"] == pytest.approx(fixed_ms, rel=0.05) + columns, rows = breakdown_table(metrics) + assert rows[0][columns.index("overhead_ms")] is not None + assert rows[0][columns.index("transfer_ms")] is not None -# ── wire-in / wire-out hash verification ─────────────────────────────── +# ── percentiles: clamping and sample gates ───────────────────────────── -class _CorruptingClient(NoOpDataPlaneClient): - """Flips one element of one field on read — a stand-in for a wire bug.""" - def __init__(self, field: str, row: int) -> None: - super().__init__() - self._corrupt_field = field - self._corrupt_row = row +def test_cluster_percentiles_never_exceed_the_measured_max(): + """The same clamp reached through ``cluster_step_metrics``: 160 calls of + 120 ms all land in (100, 250] and interpolate to a p50 of 175 — above every + call observed, and above the exact max reported beside it.""" + metrics = cluster_step_metrics( + merge_snapshots([_rank([120.0] * 20) for _ in range(8)]), {}, 1.0 + ) + max_ms = metrics["step/by_op/put/max_ms"] - def get_samples(self, sample_ids, partition_id, select_fields): - out = super().get_samples(sample_ids, partition_id, select_fields) - if self._corrupt_field in out.keys(): - out[self._corrupt_field][self._corrupt_row] += 1 - return out + assert max_ms == pytest.approx(120.0, abs=2.0) + p50, p90 = metrics["step/by_op/put/p50_ms"], metrics["step/by_op/put/p90_ms"] + assert p50 <= p90 <= max_ms -def _hash_client(inner=None): - client = MetricsDataPlaneClient( - inner or NoOpDataPlaneClient(), verify_tensor_hash=True +def test_each_quantile_waits_for_the_samples_it_needs(): + """Each quantile needs about four observations above its rank to mean + anything, so they cannot share one gate: 48 calls carry a real median and + no usable tail, and a single threshold for both reported neither. A + percentile off a handful of calls is bucket geometry, not data.""" + + def metrics_for(n_calls): + client = _client(register=False) + _emit(client, [5.0 + i % 7 for i in range(n_calls)]) + try: + return cluster_step_metrics(merge_snapshots([client.snapshot()]), {}, 1.0) + finally: + client.close() + + thin = metrics_for(10) + assert "step/by_op/put/p50_ms" not in thin, "too thin for either" + assert "step/by_op/put/max_ms" in thin, "max always works" + + mid = metrics_for(30) + assert "step/by_op/put/p50_ms" in mid, "a median off 30 calls is real" + assert "step/by_op/put/p90_ms" not in mid, "a p90 off 30 calls is not" + + both = metrics_for(58) + assert "step/by_op/put/p50_ms" in both and "step/by_op/put/p90_ms" in both + + +def test_no_quantile_finer_than_the_sample_size_can_resolve(): + """Guard the choice itself, not just the gate. The tail is p90 rather than + p99 because a p99 off the ~58 calls a step holds equalled ``max_ms`` 80% of + the time — the maximum twice, under a more precise-sounding name.""" + assert 0.99 not in {q for q, _, _ in _QUANTILES} + for q, name, min_samples in _QUANTILES: + above_the_rank = min_samples * (1 - q) + assert above_the_rank >= 4 - 1e-9, ( + f"{name} is gated at {min_samples}, leaving only " + f"{above_the_rank:.1f} observations above its rank" + ) + + +# ── cross-process aggregation ────────────────────────────────────────── + + +def test_merge_of_nothing_is_empty(): + assert merge_snapshots([]) == {} + + +def test_merge_sums_counters_and_rederives_percentiles(): + """The accumulators are shaped to add: histograms and regression sums from + every rank combine into the true cluster distribution. Averaging per-rank + percentiles could not, which is why latency lives in fixed buckets rather + than retained samples.""" + snaps = [_rank([5.0] * 4) for _ in range(3)] + merged = merge_snapshots(snaps) + + assert merged["n_processes"] == 3 + assert merged["by_op"]["put"]["calls"] == 12 # 3 ranks x 4 puts + assert merged["by_op"]["put"]["n_bytes"] == 12_000 + assert merged["by_op"]["put"]["latency_hist"] == [ + sum(counts) + for counts in zip(*(s["by_op"]["put"]["latency_hist"] for s in snaps)) + ] + assert "p50_ms" not in merged["by_op"]["put"], "12 calls supports no percentile" + + +def test_merge_takes_max_for_max_fields(): + """A cluster's worst call is the worst any rank saw, not their sum.""" + merged = merge_snapshots([_rank([40.0]), _rank([1.0])]) + assert 40.0 <= merged["by_op"]["put"]["max_ms"] < 41.0, "max, not sum" + + +def test_cluster_frac_of_step_is_per_process_and_bounded(): + """``wall_ms`` sums processes that ran concurrently, so dividing it by one + step's wall clock exceeded 1 whenever they overlapped and read as "105% of + the step". Divided per process it is the mean share of the step a process + spent in the data plane: 10 ranks x 5 calls x 100 ms over a 5 s step is + 500 ms each, or 10%.""" + metrics = cluster_step_metrics( + merge_snapshots([_rank([100.0] * 5) for _ in range(10)]), {}, 5.0 ) - client.register_partition( - partition_id="p", fields=["ids", "lp"], num_samples=4, consumer_tasks=["t"] + + assert "busy_frac_mean" not in metrics + assert metrics["step/frac_of_step"] == pytest.approx(0.10, rel=0.1) + assert metrics["now/n_processes"] == 10 + + +def test_cluster_per_op_time_is_reported_per_call(): + """``wall_ms`` sums concurrent processes, so it scales with DP degree; + dividing by the process count trades one arbitrary denominator for another. + Per call is invariant to both DP degree and batch size, so it describes the + wire rather than the shape of the run.""" + small = cluster_step_metrics( + merge_snapshots([_rank([10.0] * 5) for _ in range(8)]), {}, 1.0 ) - return client + large = cluster_step_metrics( + merge_snapshots([_rank([10.0] * 5) for _ in range(32)]), {}, 1.0 + ) + + assert small["step/by_op/put/mean_ms"] == pytest.approx(10.0, rel=0.15) + assert large["step/by_op/put/mean_ms"] == pytest.approx( + small["step/by_op/put/mean_ms"], rel=0.15 + ), "mean must not move with cluster size" + assert large["step/by_op/put/wall_ms"] == pytest.approx( + 4 * small["step/by_op/put/wall_ms"], rel=0.15 + ), "the sum does move with cluster size" + + columns, _ = breakdown_table(small) + assert "mean_ms" in columns and "percent_of_dataplane" in columns -def _hash_fields(n=4): - return TensorDict( +# ── what gets charted: shares, volume, and the breakdown table ───────── + + +def test_percent_of_dataplane_names_the_bottleneck_and_says_of_what(): + """The denominator is data-plane time, not the step: ``by_op`` sums to 100 + by construction, so the largest is the bottleneck *within the data plane*. + Whether the data plane mattered at all is ``frac_of_step`` — here a tenth of + a second of data-plane work inside a 10 s step is 9% of one, 100% of the + other. 32 per-op line charts answer neither question.""" + client = _busy({"get": (100, 9.0), "put": (10, 1.0), "clear": (10, 0.1)}) + metrics = cluster_step_metrics( + merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 10.0 + ) + by_op = { + k: v + for k, v in metrics.items() + if k.startswith("step/percent_of_dataplane/by_op/") + } + + assert sum(by_op.values()) == pytest.approx(100.0), "percent of one total" + assert max(by_op, key=by_op.__getitem__) == "step/percent_of_dataplane/by_op/get" + assert by_op["step/percent_of_dataplane/by_op/get"] == pytest.approx( + 100 * 900 / 911, rel=0.05 + ) + # the two denominators are different questions and must not agree + assert metrics["step/frac_of_step"] == pytest.approx(0.0911, rel=0.1) + client.close() + + +def test_headline_drops_per_op_detail_but_keeps_the_percentages(): + """Four ops times eight fields is 32 series saying one thing. The detail is + still computed — the breakdown table is built from the same dict, so the two + cannot disagree — but only the totals and percentages are charted.""" + client = _busy({"get": (100, 9.0), "put": (10, 1.0), "clear": (10, 0.1)}) + metrics = cluster_step_metrics( + merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 1.0 + ) + head = headline_series(metrics) + + # the property, not a ratio: a ratio drifts as series are added on either + # side, while "no per-op series is charted" is the thing being claimed + assert len(head) < len(metrics), f"{len(head)} of {len(metrics)}" + assert not [k for k in head if k.split("/")[1:2] in (["get"], ["put"], ["clear"])] + assert "step/percent_of_dataplane/by_op/get" in head + assert "step/wall_ms" in head and "step/frac_of_step" in head + assert breakdown_table(metrics)[1], "the table still has rows" + client.close() + + +def test_per_op_volume_replaces_the_written_read_split(): + """``comm_volume_mb`` alone hides which direction the traffic went — on a + real step get moved 20.8 MB against put's 2.7 MB — while the old + ``bytes_written``/``bytes_read`` pair was charted by nobody and absent from + the table. One key per op says it finer, in both scopes.""" + client = _client(register=False) + _emit(client, [5.0] * 6, op="get", n_bytes=3_000_000) + _emit(client, [5.0] * 2, op="put", n_bytes=1_000_000) + _emit(client, [1.0], op="clear", n_bytes=0) + + driver = client.get_step_metrics(1.0) + assert "step/bytes_written_mb" not in driver + assert "step/bytes_read_mb" not in driver + assert driver["step/volume_mb/by_op/put"] == pytest.approx(2.0) + + head = headline_series( + cluster_step_metrics( + merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 1.0 + ) + ) + assert head["step/volume_mb/by_op/get"] == pytest.approx(18.0) + assert head["step/volume_mb/by_op/put"] == pytest.approx(2.0) + assert "step/volume_mb/by_op/clear" not in head, "no payload, not a zero" + per_op = sum(v for k, v in head.items() if k.startswith("step/volume_mb/")) + assert per_op == pytest.approx(head["step/comm_volume_mb"]), "parts make the whole" + client.close() + + +def test_breakdown_table_rows_by_op_worst_first(): + """One row per op, ordered by wall time, so the expensive op is the first + line read rather than the alphabetically luckiest. Share of data-plane time + is the second column for the same reason.""" + columns, rows = breakdown_table( { - "ids": torch.arange(n * 6, dtype=torch.int64).reshape(n, 6), - "lp": torch.linspace(0, 1, n * 6, dtype=torch.bfloat16).reshape(n, 6), - }, - batch_size=[n], + "step/wall_ms": 100.0, + "step/percent_of_dataplane/by_op/get": 10.0, + "step/percent_of_dataplane/by_op/put": 90.0, + "step/by_op/get/calls": 8, + "step/by_op/get/wall_ms": 10.0, + "step/by_op/get/max_ms": 2.0, + "step/by_op/put/calls": 2, + "step/by_op/put/wall_ms": 90.0, + "step/by_op/put/max_ms": 50.0, + "step/comm_volume_mb": 1.0, # not per-op, must not become a row + "now/bytes_outstanding_mb": 0.0, # a level, likewise + } + ) + + assert columns[0] == "op" + assert columns[1] == "percent_of_dataplane", "the bottleneck reads first" + assert [r[0] for r in rows] == ["put", "get"], "worst first" + assert rows[0][columns.index("wall_ms")] == 90.0 + assert rows[0][columns.index("percent_of_dataplane")] == pytest.approx(90.0) + + +def test_breakdown_table_leaves_withheld_series_empty(): + """A percentile below the sample gate, or a fit that is not trustworthy, is + absent from the series — the table must carry None there rather than a zero + that would read as a measurement.""" + columns, rows = breakdown_table( + { + "step/by_op/put/calls": 3, + "step/by_op/put/wall_ms": 5.0, + "step/by_op/put/max_ms": 2.0, + } + ) + + assert rows[0][columns.index("p90_ms")] is None + assert rows[0][columns.index("overhead_ms")] is None + assert rows[0][columns.index("calls")] == 3 + + +def test_breakdown_table_is_empty_when_nothing_ran(): + assert breakdown_table({"step/wall_ms": 0.0})[1] == [] + + +def test_breakdown_table_ignores_reserved_namespaces(): + """``step/self/overhead_ms`` and ``step/volume_mb/by_op/get`` share the + three-part shape of a per-op series; they must feed the right row (or none) + rather than invent a "self" or "volume_mb" op beside put and get.""" + columns, rows = breakdown_table( + { + "step/by_op/get/calls": 3, + "step/by_op/get/wall_ms": 9.0, + "step/by_op/get/mb": 18.0, + "step/volume_mb/by_op/get": 18.0, + "step/self/overhead_ms": 6.2, + "step/hash/mismatches": 0, + "step/percent_of_dataplane/by_cause/transfer": 40.0, + } ) + assert columns[0] == "op" + assert [r[0] for r in rows] == ["get"], rows + assert rows[0][columns.index("mb")] == 18.0 + + +# ── wire-in / wire-out hash verification ─────────────────────────────── + + +def test_hash_state_and_counters_absent_when_the_guard_is_off(): + """Default construction does no hashing work and emits no hash series. + + Always-zero counters on every run that never asked for the guard would read + as "checked, nothing wrong" rather than "not checked".""" + client = _client() + ids = _ids(4) + client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + merged = merge_snapshots([client.snapshot(reset_step_window=True)]) + + assert client.snapshot()["hash_verify"]["rows_recorded"] == 0 + assert client._hash_by_partition == {} + assert not [k for k in client.get_step_metrics(1.0) if "hash" in k] + assert not [k for k in cluster_step_metrics(merged, {}, 1.0) if "hash" in k] + client.close() + def test_hash_verification_clean_roundtrip(): - client = _hash_client() - ids = [f"u{i}" for i in range(4)] + client = _client(verify_tensor_hash=True) + ids = _ids(4) client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids", "lp"]) + assert client.snapshot()["hash_verify"] == { + "rows_recorded": 4, + "rows_checked": 4, + "rows_unverified": 0, + "mismatches": 0, + "fields_skipped": 0, + "guard_failures": 0, # the guard itself never raised + } + client.close() + + +def test_guard_failure_is_absorbed_counted_and_charted(caplog, monkeypatch): + """A bug in the guard must not take the transfer down — and must not read + as clean either. + + Both failures this check has produced were exactly this: an unhandled dtype + inside ``_row_fingerprints`` propagating out of ``put_samples`` and killing + the run. Absorbing them is only safe if the absorption is visible, so the + count has to reach the series even though no rows were ever recorded. + """ + client = _client(verify_tensor_hash=True) + ids = _ids(4) + + def boom(*_args, **_kwargs): + raise NotImplementedError("no hash_tensor kernel for this dtype") + + monkeypatch.setattr(client, "_row_fingerprints", boom) + with caplog.at_level(logging.WARNING): + # neither call may raise, on either side of the wire + client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + hv = client.snapshot()["hash_verify"] - assert hv["rows_recorded"] == 4 - assert hv["rows_checked"] == 4 - assert hv["rows_unverified"] == 0 - assert hv["mismatches"] == 0 + assert hv["guard_failures"] == 2, "put and get each failed once" + assert hv["rows_recorded"] == 0 and hv["rows_checked"] == 0, "nothing checked" + assert caplog.text.count("hash guard failed") == 1, "logged once, not per call" + # and it is a series, not just a counter -- the gate cannot key on rows + assert client.get_step_metrics(1.0)["step/hash/guard_failures"] == 2 client.close() -def test_hash_verification_detects_corruption(): - client = _hash_client(_CorruptingClient(field="ids", row=2)) - ids = [f"u{i}" for i in range(4)] +def test_hash_mismatch_reaches_every_scope(): + """A guard whose findings are not reported is not a guard. + ``_log_data_plane_metrics`` prefers the cluster path whenever the fan-out + reaches more than one process — every real run — and that path once emitted + no hash counters at all.""" + client = _client(_CorruptingClient(field="ids", row=2), verify_tensor_hash=True) + ids = _ids(4) client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids", "lp"]) - hv = client.snapshot()["hash_verify"] - assert hv["mismatches"] == 1 + assert client.snapshot()["hash_verify"]["mismatches"] == 1 assert client.get_step_metrics(1.0)["step/hash/mismatches"] == 1 + + cluster = cluster_step_metrics( + merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 1.0 + ) + assert cluster["step/hash/mismatches"] == 1 + assert "step/hash/fields_skipped" in cluster, "abstentions visible too" + assert headline_series(cluster)["step/hash/mismatches"] == 1, "and charted" client.close() def test_hash_verification_survives_shard_readback(): """A 4-row put read back two rows at a time must still line up: the fingerprint is per row, not per batch.""" - client = _hash_client() - ids = [f"u{i}" for i in range(4)] + client = _client(verify_tensor_hash=True) + ids = _ids(4) client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) for shard in (ids[:2], ids[2:]): client.get_samples(sample_ids=shard, partition_id="p", select_fields=["ids"]) @@ -581,8 +936,8 @@ def test_hash_verification_reports_rows_it_never_wrote(): """A consumer-side client sees only wire-out. Those rows must land in ``rows_unverified`` — reporting 0 mismatches would read as 'clean'.""" inner = NoOpDataPlaneClient() - writer = _hash_client(inner) - ids = [f"u{i}" for i in range(4)] + writer = _client(inner, verify_tensor_hash=True) + ids = _ids(4) writer.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) reader = MetricsDataPlaneClient(inner, verify_tensor_hash=True) @@ -598,9 +953,10 @@ def test_hash_verification_reports_rows_it_never_wrote(): def test_hash_fingerprints_released_on_clear(): """Fingerprints must be bounded by the live key population, not by cumulative traffic.""" - client = _hash_client() - ids = [f"u{i}" for i in range(4)] + client = _client(verify_tensor_hash=True) + ids = _ids(4) client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + assert client._hash_by_partition["p"] client.clear_samples(sample_ids=ids, partition_id="p") assert client._hash_by_partition == {} @@ -609,635 +965,128 @@ def test_hash_fingerprints_released_on_clear(): def test_hash_fingerprint_covers_jagged_fields(): """The per-token fields on this wire are jagged by the time they reach - ``put_samples`` (``codec.pack_jagged_fields``). Skipping nested leaves - would leave the entire bulk payload unguarded while still reporting zero - mismatches — a guard that reads as clean because it checked nothing.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + ``put_samples``. Skipping nested leaves would leave the entire bulk payload + unguarded while still reporting zero mismatches — a guard that reads as + clean because it checked nothing.""" + client = _client(verify_tensor_hash=True) rows = [torch.arange(n, dtype=torch.int64) + n for n in (3, 5, 4)] digest = client._row_fingerprints(_jagged(rows), ["a", "b", "c"])["x"] assert client.snapshot()["hash_verify"]["fields_skipped"] == 0 assert digest.batch_scoped, "jagged digests only reconcile per batch" assert len(digest.per_row) == 3 - # A jagged digest carries the row length, so a change in any row's - # payload moves every row's value and a length change moves one. changed = list(rows) changed[1] = changed[1] + 1 assert client._row_fingerprints(_jagged(changed), ["a", "b", "c"])["x"] != digest + client.close() def test_hash_fingerprint_matches_across_jagged_and_dense(): """``_from_wire`` densifies a jagged field whose rows are uniform, so a - jagged put has to reconcile against a dense get. - - Regression guard: picking the scheme from the layout in hand rather than - from what was recorded made every row of a uniform batch report a - mismatch — 940 false alarms over the verification soak. - """ - client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + jagged put has to reconcile against a dense get. Picking the scheme from the + layout in hand rather than from what was recorded made every row of a + uniform batch report a mismatch — 940 false alarms over the soak.""" + client = _client(verify_tensor_hash=True) ids = ["a", "b"] dense = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int64) # uniform rows: both sides reduce per row, and a densified read agrees put_side = client._row_fingerprints(_jagged(list(dense.unbind())), ids)["x"] assert not put_side.batch_scoped - get_side = client._row_fingerprints(TensorDict({"x": dense}, batch_size=[2]), ids) - assert put_side == get_side["x"] + assert put_side == client._row_fingerprints( + TensorDict({"x": dense}, batch_size=[2]), ids + )["x"] # ragged rows: batch-scoped, and the read replays that scheme rather than # choosing one from the dense tensor in hand ragged = _jagged([torch.tensor([1, 2, 3]), torch.tensor([4, 5])]) - scoped = client._row_fingerprints(ragged, ids)["x"] - assert scoped.batch_scoped - replayed = client._row_fingerprints( + assert client._row_fingerprints(ragged, ids)["x"].batch_scoped + assert client._row_fingerprints( TensorDict({"x": dense}, batch_size=[2]), ids, batch_scoped_fields={"x"} - )["x"] - assert replayed.batch_scoped - - -def test_hash_shard_read_of_jagged_field_is_unverified_not_a_mismatch(): - """A batch-scoped digest covers the whole buffer, so a shard read cannot - reproduce it. That has to report as unverified — reporting it as a - mismatch would make the guard cry wolf on every sharded fetch.""" - client = _hash_client() - ids = [f"u{i}" for i in range(4)] - rows = [torch.arange(3, dtype=torch.int64) + i for i in range(4)] - client.put_samples(sample_ids=ids, partition_id="p", fields=_jagged(rows)) - client.get_samples(sample_ids=ids[:2], partition_id="p", select_fields=["x"]) - - assert client.snapshot()["hash_verify"]["mismatches"] == 0 + )["x"].batch_scoped client.close() -def test_hash_rectangular_put_read_back_ragged_checks_row_lengths(): - """A rectangular put can come back jagged — truncating one row makes the - batch ragged. The content is no longer comparable, so the field counts as - an abstention, but the row *lengths* still are, and a row that changed - length is a real divergence. - - Failing the whole field instead reported 3584 mismatches per step on a - healthy 5-process run: writing a shard with uniform rows and reading it - back inside a batch whose other rows differ in length is the normal - shape of the pipeline, not a corruption.""" - client = _hash_client(_RaggedOnReadClient()) - ids = [f"u{i}" for i in range(4)] - dense = TensorDict( - {"x": torch.arange(16, dtype=torch.int64).reshape(4, 4)}, batch_size=[4] - ) - client.put_samples(sample_ids=ids, partition_id="p", fields=dense) - client.get_samples(sample_ids=ids, partition_id="p", select_fields=["x"]) - - hv = client.snapshot()["hash_verify"] - assert hv["mismatches"] > 0, "a truncated row must not read as clean" - assert hv["fields_skipped"] == 1, "and the content it could not compare" - - -def test_hash_shard_of_a_ragged_field_is_skipped_not_a_mismatch(): - """The abstention that *is* legitimate: a batch-scoped digest covers the - whole buffer it was reduced over, so a shard of it is genuinely - incomparable. That must land in ``fields_skipped`` — visible, but not - crying wolf.""" - client = _hash_client() - ids = [f"u{i}" for i in range(4)] - rows = [torch.arange(3 + i, dtype=torch.int64) for i in range(4)] - client.put_samples(sample_ids=ids, partition_id="p", fields=_jagged(rows)) - client.get_samples(sample_ids=ids[:2], partition_id="p", select_fields=["x"]) - - hv = client.snapshot()["hash_verify"] - assert hv["mismatches"] == 0 - assert hv["fields_skipped"] == 1 - assert client.get_step_metrics(1.0)["step/hash/fields_skipped"] == 1 - - -class _RaggedOnReadClient(NoOpDataPlaneClient): - """Returns one row shorter than it was written — a truncation on the wire.""" - - def get_samples(self, sample_ids, partition_id, select_fields): - out = super().get_samples(sample_ids, partition_id, select_fields) - rows = list(out["x"].unbind()) - rows[1] = rows[1][:-1] - return TensorDict( - {"x": torch.nested.nested_tensor(rows, layout=torch.jagged)}, - batch_size=[len(sample_ids)], - ) - - def test_hash_fingerprint_separates_dtype(): - """The values reduce identically once bitcast, so only the dtype salt - makes a precision change visible.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + """The values reduce identically once bitcast, so only the dtype salt makes + a precision change visible.""" + client = _client(verify_tensor_hash=True) ids = ["a", "b"] values = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) - as_fp32 = client._row_fingerprints(TensorDict({"x": values}, batch_size=[2]), ids)[ - "x" - ] + + as_fp32 = client._row_fingerprints(TensorDict({"x": values}, batch_size=[2]), ids) as_bf16 = client._row_fingerprints( TensorDict({"x": values.to(torch.bfloat16)}, batch_size=[2]), ids - )["x"] - assert as_fp32.per_row != as_bf16.per_row + ) + assert as_fp32["x"].per_row != as_bf16["x"].per_row + client.close() def test_hash_fingerprint_handles_float8(): """``hash_tensor`` has no float8 kernel. Without the integer bitcast the ``NotImplementedError`` propagates out of ``put_samples`` and takes the transfer down with it.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + client = _client(verify_tensor_hash=True) fp8 = TensorDict( {"x": torch.tensor([[1.0, 2.0], [3.0, 4.0]]).to(torch.float8_e4m3fn)}, batch_size=[2], ) - digest = client._row_fingerprints(fp8, ["a", "b"])["x"] - assert len(digest.per_row) == 2 + # through the public path, because that is where the exception surfaced + client.put_samples(sample_ids=["a", "b"], partition_id="p", fields=fp8) + + assert client.snapshot()["hash_verify"]["rows_recorded"] == 2 + client.close() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_hash_fingerprint_handles_device_tensors(): - """Torch has no ``bitwise_xor`` CUDA kernel for UInt64, so salting the - digest tensor before it leaves the device raises ``NotImplementedError`` - for any backend whose get returns device tensors -- register mode under - GDR does. The digests must also match the host's, or a device-resident - get would verify against a host put as a mismatch on every row.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient(), verify_tensor_hash=True) + """Torch has no ``bitwise_xor`` CUDA kernel for UInt64, so salting the digest + tensor before it leaves the device raises ``NotImplementedError`` for any + backend whose get returns device tensors — register mode under GDR does. The + digests must also match the host's, or a device-resident get would verify + against a host put as a mismatch on every row.""" + client = _client(verify_tensor_hash=True) values = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) ids = ["a", "b"] - on_host = client._row_fingerprints(TensorDict({"x": values}, batch_size=[2]), ids)[ - "x" - ] + + on_host = client._row_fingerprints(TensorDict({"x": values}, batch_size=[2]), ids) on_device = client._row_fingerprints( TensorDict({"x": values.cuda()}, batch_size=[2]), ids - )["x"] - assert on_device.per_row == on_host.per_row - - -def test_hash_verification_off_by_default(): - """Default construction must do no hashing work at all.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client.register_partition( - partition_id="p", fields=["ids", "lp"], num_samples=4, consumer_tasks=["t"] ) - ids = [f"u{i}" for i in range(4)] - client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) - client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) - - assert client.snapshot()["hash_verify"]["rows_recorded"] == 0 - assert client._hash_by_partition == {} - assert "step/hash/mismatches" not in client.get_step_metrics(1.0) + assert on_device["x"].per_row == on_host["x"].per_row client.close() -# ── cross-process aggregation ────────────────────────────────────────── - - -def _rank_client(latencies_ms): - """A client that has seen one put per entry, at that latency.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - now = monotonic() - for ms in latencies_ms: - client._emit("put", "p", 1, 1_000, now - ms / 1e3, "ok") - return client - - -def _rank_with(latencies_ms): - """Its snapshot, for the merge tests.""" - return _rank_client(latencies_ms).snapshot() - - -def test_merge_sums_counters_and_rederives_percentiles(): - """The accumulators are shaped to add: histograms and regression sums - from every rank combine into the true cluster distribution. Averaging - per-rank percentiles could not do this, which is the whole reason the - latency lives in fixed buckets rather than retained samples.""" - ranks = [_rank_client([5.0] * 4) for _ in range(3)] - merged = merge_snapshots([c.snapshot() for c in ranks]) - - assert merged["n_processes"] == 3 - assert merged["by_op"]["put"]["calls"] == 12 # 3 ranks x 4 puts - assert merged["by_op"]["put"]["n_bytes"] == 12_000 - assert sum(merged["by_op"]["put"]["latency_hist"]) == 12 - # the buckets add: the merged histogram is the ranks' summed elementwise - per_rank = [c.snapshot()["by_op"]["put"]["latency_hist"] for c in ranks] - assert merged["by_op"]["put"]["latency_hist"] == [ - sum(counts) for counts in zip(*per_rank) - ] - # 12 calls supports no percentile, and none is offered - assert "p50_ms" not in merged["by_op"]["put"] - for c in ranks: - c.close() - - -def test_merge_takes_max_for_max_fields(): - """A cluster's worst call is the worst any rank saw, not their sum.""" - slow = _rank_client([40.0]) - fast = _rank_client([1.0]) - merged = merge_snapshots([slow.snapshot(), fast.snapshot()]) - assert merged["by_op"]["put"]["max_ms"] >= 40.0 - assert merged["by_op"]["put"]["max_ms"] < 41.0, "max, not sum" - slow.close() - fast.close() - - -def test_merge_of_nothing_is_empty(): - assert merge_snapshots([]) == {} - - -def test_cluster_step_metrics_report_their_own_cost(): - """``step/self/overhead_ms`` is the wrapper's own wall time minus - the inner client's, summed over processes — the bill for measuring, - sitting beside what it bought.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client.register_partition( - partition_id="p", fields=["x"], num_samples=4, consumer_tasks=["t"] - ) - for i in range(4): - client.put_samples( - sample_ids=[f"u{i}"], - partition_id="p", - fields=TensorDict({"x": torch.zeros(1, 512)}, batch_size=[1]), - ) - merged = merge_snapshots([client.snapshot()]) - metrics = cluster_step_metrics(merged, {}, 1.0) - - assert metrics["now/n_processes"] == 1 - assert metrics["step/self/overhead_ms"] > 0, "measuring is never free" - # Deliberately not clamped to 1. Against this no-op inner client the RPC - # is instant, so measuring costs more than the thing measured and the - # ratio exceeds 100% -- which is exactly the signal worth surfacing. - # Against a real backend it lands near 0.01. - assert metrics["step/self/frac"] > 0 - assert "step/wall_ms" in metrics and "step/comm_volume_mb" in metrics - client.close() - - -def test_cluster_frac_of_step_is_per_process_and_bounded(): - """``wall_ms`` sums processes that ran concurrently, so dividing it by - one step's wall clock exceeded 1 whenever they overlapped and read as - "105% of the step". Divided per process it is the mean share of the - step a process spent in the data plane, which is what the name claims: - 10 ranks x 5 calls x 100 ms over a 5 s step is 500 ms each, or 10%.""" - ranks = [_rank_with([100.0] * 5) for _ in range(10)] - metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 5.0) - - assert "busy_frac_mean" not in metrics - assert metrics["step/frac_of_step"] == pytest.approx(0.10, rel=0.1) - assert metrics["now/n_processes"] == 10 - - -def test_cluster_overhead_includes_the_collection_fan_out(): - """The fan-out is the larger half of the bill. Reporting only the per-op - wrapper understated the real cost by ~19x in the cross-process e2e - (0.13 ms reported against 2.44 ms actually spent).""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client.register_partition( - partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["t"] - ) - client.put_samples( - sample_ids=["a"], - partition_id="p", - fields=TensorDict({"x": torch.zeros(1, 8)}, batch_size=[1]), - ) - merged = merge_snapshots([client.snapshot()]) - without = cluster_step_metrics(merged, {}, 1.0) - with_gather = cluster_step_metrics(merged, {}, 1.0, collect_ms=2.31) - - delta = with_gather["step/self/overhead_ms"] - without["step/self/overhead_ms"] - assert delta == pytest.approx(2.31, rel=1e-6) - client.close() - - -def test_cluster_percentiles_never_exceed_the_measured_max(): - """Bucket interpolation spreads a bucket's samples uniformly across it, - so calls clustered low in a wide bucket read high: 160 calls of 120 ms - all land in (100, 250] and interpolate to a p50 of 175 — above every - call observed, and above the exact max reported beside it. The max is - the tighter bound, so the percentiles are clamped to it.""" - merged = merge_snapshots([_rank_with([120.0] * 20) for _ in range(8)]) - metrics = cluster_step_metrics(merged, {}, 1.0) - - assert metrics["step/by_op/put/max_ms"] == pytest.approx(120.0, abs=2.0) - assert metrics["step/by_op/put/p50_ms"] <= metrics["step/by_op/put/max_ms"] - assert metrics["step/by_op/put/p90_ms"] <= metrics["step/by_op/put/max_ms"] - assert metrics["step/by_op/put/p50_ms"] <= metrics["step/by_op/put/p90_ms"] - - -def test_cluster_percentiles_withheld_below_a_useful_sample_count(): - """A percentile off a handful of calls is bucket geometry, not data. - Silence beats a number that looks like an answer.""" - few = merge_snapshots([_rank_with([120.0] * 3) for _ in range(2)]) # 6 calls - many = merge_snapshots([_rank_with([120.0] * 20) for _ in range(8)]) # 160 - - assert "step/by_op/put/p50_ms" not in cluster_step_metrics(few, {}, 1.0) - assert "step/by_op/put/max_ms" in cluster_step_metrics(few, {}, 1.0), ( - "max always works" - ) - assert "step/by_op/put/p50_ms" in cluster_step_metrics(many, {}, 1.0) +# ── hash verification: what counts as a mismatch vs an abstention ────── -def test_cluster_series_declare_delta_or_level(): - """A per-step delta and an instantaneous level shared the ``_mb`` suffix - and a chart, with nothing to tell them apart. Every series now sits - under ``step/`` or ``now/`` so the kind is on the axis label.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client.register_partition( - partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["t"] - ) - client.put_samples( - sample_ids=["a"], - partition_id="p", - fields=TensorDict({"x": torch.zeros(1, 8)}, batch_size=[1]), - ) - metrics = cluster_step_metrics(merge_snapshots([client.snapshot()]), {}, 1.0) - - unlabelled = [k for k in metrics if not k.startswith(("step/", "now/"))] - assert not unlabelled, unlabelled - assert "now/bytes_outstanding_mb" in metrics, "a level" - assert "step/comm_volume_mb" in metrics, "a delta" - client.close() - - -def test_clear_frees_only_what_was_actually_live(): - """A clear may name uids already dropped, or belonging to another - partition. Billing those releases bytes this partition never held: - clearing 50 live keys alongside 50 unknown ones freed two thirds of a - partition that had lost half its keys.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - ids = [f"u{i}" for i in range(100)] - client.register_partition( - partition_id="p", fields=["x"], num_samples=100, consumer_tasks=["t"] - ) - client.put_samples( - sample_ids=ids, - partition_id="p", - fields=TensorDict({"x": torch.zeros(100, 250)}, batch_size=[100]), - ) - total = client.snapshot()["bytes_outstanding"] - - client.clear_samples( - sample_ids=ids[:50] + [f"unknown{i}" for i in range(50)], partition_id="p" - ) - assert client.snapshot()["bytes_outstanding"] == total // 2 - - client.clear_samples(sample_ids=ids[50:], partition_id="p") - assert client.snapshot()["bytes_outstanding"] == 0 - client.close() - - -def test_outstanding_reconciles_over_random_put_clear_sequences(): - """Interleaved puts and partial clears must always land back at zero; - the pro-rata release is only sound if it does.""" - import random - - rng = random.Random(0) - for _ in range(50): - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client.register_partition( - partition_id="p", fields=["x"], num_samples=5000, consumer_tasks=["t"] - ) - live: set[str] = set() - for _ in range(rng.randint(1, 6)): - batch = list( - dict.fromkeys( - f"k{rng.randint(0, 60)}" for _ in range(rng.randint(1, 20)) - ) - ) - client.put_samples( - sample_ids=batch, - partition_id="p", - fields=TensorDict( - {"x": torch.zeros(len(batch), 250)}, batch_size=[len(batch)] - ), - ) - live |= set(batch) - if live and rng.random() < 0.5: - drop = rng.sample(sorted(live), k=rng.randint(1, len(live))) - client.clear_samples(sample_ids=drop, partition_id="p") - live -= set(drop) - if live: - client.clear_samples(sample_ids=sorted(live), partition_id="p") - assert client.snapshot()["bytes_outstanding"] == 0 - client.close() - +def test_hash_shard_of_a_ragged_field_is_skipped_not_a_mismatch(): + """The abstention that *is* legitimate: a batch-scoped digest covers the + whole buffer it was reduced over, so a shard of it is genuinely + incomparable. That must land in ``fields_skipped`` — visible, but not + crying wolf on every sharded fetch.""" + client = _client(verify_tensor_hash=True) + ids = _ids(4) + rows = [torch.arange(3 + i, dtype=torch.int64) for i in range(4)] + client.put_samples(sample_ids=ids, partition_id="p", fields=_jagged(rows)) + client.get_samples(sample_ids=ids[:2], partition_id="p", select_fields=["x"]) -def test_snapshot_percentiles_never_exceed_the_measured_max(): - """The clamp lives in ``_derive_op_metrics``, not at one call site, so - ``snapshot()`` and ``merge_snapshots()`` inherit it. Without it five - register calls of 0.011 ms all land in (0, 0.1] and interpolate to a - p50 of 0.05 — four times the slowest call that happened — on a public - surface documented as the cumulative view.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - now = monotonic() - for _ in range(120): # enough that both quantiles clear their gate - client._emit("register", "p", 1, 0, now - 0.011 / 1e3, "ok") - - for view in (client.snapshot(), merge_snapshots([client.snapshot()])): - stats = view["by_op"]["register"] - assert stats["p50_ms"] <= stats["max_ms"], stats - assert stats["p90_ms"] <= stats["max_ms"], stats - assert stats["p50_ms"] <= stats["p90_ms"], stats + hv = client.snapshot()["hash_verify"] + assert hv["mismatches"] == 0 + assert hv["fields_skipped"] == 1 + assert client.get_step_metrics(1.0)["step/hash/fields_skipped"] == 1 client.close() -def test_breakdown_table_rows_by_op_worst_first(): - """One row per op, ordered by wall time, so the expensive op is the - first line read rather than the alphabetically luckiest. Share of - data-plane time is the second column for the same reason: the answer - to "what is my bottleneck" should be the top-left of the table.""" - metrics = { - "step/wall_ms": 100.0, - "step/percent_of_dataplane/by_op/get": 10.0, - "step/percent_of_dataplane/by_op/put": 90.0, - "step/by_op/get/calls": 8, - "step/by_op/get/wall_ms": 10.0, - "step/by_op/get/max_ms": 2.0, - "step/by_op/put/calls": 2, - "step/by_op/put/wall_ms": 90.0, - "step/by_op/put/max_ms": 50.0, - "step/comm_volume_mb": 1.0, # not per-op, must not become a row - "now/bytes_outstanding_mb": 0.0, # a level, likewise - } - columns, rows = breakdown_table(metrics) - - assert columns[0] == "op" - assert [r[0] for r in rows] == ["put", "get"], "worst first" - assert len(rows) == 2, "only per-op series become rows" - assert rows[0][columns.index("wall_ms")] == 90.0 - assert rows[0][columns.index("percent_of_dataplane")] == pytest.approx(90.0) - assert columns[1] == "percent_of_dataplane", "the bottleneck reads first" - - -def test_breakdown_table_leaves_withheld_series_empty(): - """A percentile below the sample gate, or a fit that is not trustworthy, - is absent from the series — the table must carry None there rather than - a zero that would read as a measurement.""" - columns, rows = breakdown_table( - { - "step/by_op/put/calls": 3, - "step/by_op/put/wall_ms": 5.0, - "step/by_op/put/max_ms": 2.0, - } - ) - row = rows[0] - assert row[columns.index("p90_ms")] is None - assert row[columns.index("overhead_ms")] is None - assert row[columns.index("calls")] == 3 - - -def test_breakdown_table_is_empty_when_nothing_ran(): - assert breakdown_table({"step/wall_ms": 0.0})[1] == [] - - -def test_cluster_view_carries_the_latency_split(): - """The split was emitted on the driver path only, so the cluster view — - the one a real run logs — could never show it, and its table column was - always empty. The cluster fit is the better one besides: it is over - every rank's sufficient statistics, so it has far more samples and far - more size variation, and size variation is what decides whether the - split is identifiable at all.""" - fixed_ms, mb_per_s = 6.0, 400.0 - ranks = [] - for rank in range(4): - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - now = monotonic() - for i in range(8): # varied sizes, or the fit is unidentifiable - n_bytes = 100_000 * (i + 1 + rank) - wall_ms = fixed_ms + n_bytes / (mb_per_s * 1e3) - client._emit("get", "p", 1, n_bytes, now - wall_ms / 1e3, "ok") - ranks.append(client.snapshot()) - - metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 1.0) - assert "step/by_op/get/overhead_ms" in metrics, "cluster view must carry the split" - - total = ( - metrics["step/by_op/get/overhead_ms"] + metrics["step/by_op/get/transfer_ms"] - ) - assert total == pytest.approx(metrics["step/by_op/get/mean_ms"], rel=0.05) - assert metrics["step/by_op/get/overhead_ms"] == pytest.approx(fixed_ms, rel=0.05) - - columns, rows = breakdown_table(metrics) - row = rows[0] - assert row[columns.index("overhead_ms")] is not None - assert row[columns.index("transfer_ms")] is not None - - -def test_cluster_per_op_time_is_reported_per_call(): - """``wall_ms`` sums concurrent processes, so it scales with DP degree; - dividing by the process count trades one arbitrary denominator for - another. Per call is invariant to both DP degree and batch size, so it - describes the wire rather than the shape of the run, and is comparable - across runs and cluster sizes.""" - small = cluster_step_metrics( - merge_snapshots([_rank_with([10.0] * 5) for _ in range(8)]), {}, 1.0 - ) - large = cluster_step_metrics( - merge_snapshots([_rank_with([10.0] * 5) for _ in range(32)]), {}, 1.0 - ) - - assert small["step/by_op/put/mean_ms"] == pytest.approx(10.0, rel=0.15) - assert large["step/by_op/put/mean_ms"] == pytest.approx( - small["step/by_op/put/mean_ms"], rel=0.15 - ), "mean must not move with cluster size" - assert large["step/by_op/put/wall_ms"] == pytest.approx( - 4 * small["step/by_op/put/wall_ms"], rel=0.15 - ), "the sum does move with cluster size" - - columns, rows = breakdown_table(small) - assert "mean_ms" in columns and "percent_of_dataplane" in columns - - -def test_each_quantile_waits_for_the_samples_it_needs(): - """Each quantile needs about four observations above its rank to mean - anything, so they cannot share one gate: 48 calls carry a real median - and no usable tail, and a single threshold for both reported neither. - - The tail one is p90 rather than p99 for the same reason. A step holds - tens of calls, and a p99 off 58 of them equalled the maximum 80% of the - time -- ``max_ms`` under a more precise-sounding name. - """ - - def metrics_for(n_calls): - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - now = monotonic() - for i in range(n_calls): - client._emit("put", "p", 1, 1_000, now - (5.0 + i % 7) / 1e3, "ok") - return cluster_step_metrics(merge_snapshots([client.snapshot()]), {}, 1.0) - - assert "step/by_op/put/p50_ms" not in metrics_for(10), "too thin for either" - mid = metrics_for(30) - assert "step/by_op/put/p50_ms" in mid, "a median off 30 calls is real" - assert "step/by_op/put/p90_ms" not in mid, "a p90 off 30 calls is not" - both = metrics_for(58) - assert "step/by_op/put/p50_ms" in both and "step/by_op/put/p90_ms" in both - - -def test_no_quantile_finer_than_the_sample_size_can_resolve(): - """Guard on the choice itself, not just the gate: reporting a p99 at - these per-step call counts would mean reporting the maximum twice.""" - assert 0.99 not in {q for q, _, _ in _QUANTILES} - for q, name, min_samples in _QUANTILES: - above_the_rank = min_samples * (1 - q) - assert above_the_rank >= 4 - 1e-9, ( - f"{name} is gated at {min_samples}, leaving only " - f"{above_the_rank:.1f} observations above its rank" - ) - - -class _JaggedEcho(NoOpDataPlaneClient): - """Returns whatever was put, jagged, so row lengths survive the trip.""" - - def __init__(self) -> None: - super().__init__() - self.rows: dict[tuple[str, str], dict[str, torch.Tensor]] = {} - - def put_samples(self, sample_ids, partition_id, fields=None, tags=None): - for key in fields.keys(): - v = fields.get(key) - rows = v.unbind() if v.is_nested else list(v) - for sid, row in zip(sample_ids, rows): - self.rows.setdefault((partition_id, sid), {})[str(key)] = row.clone() - return super().put_samples( - sample_ids=sample_ids, partition_id=partition_id, fields=fields, tags=tags - ) - - def get_samples(self, sample_ids, partition_id, select_fields): - out = {} - for f in select_fields: - rows = [self.rows[(partition_id, sid)][f] for sid in sample_ids] - out[f] = ( - torch.stack(rows) - if all(r.shape == rows[0].shape for r in rows[1:]) - else torch.nested.nested_tensor(rows, layout=torch.jagged) - ) - return TensorDict(out, batch_size=[len(sample_ids)]) - - -def _jagged_ids(lengths, seed=0): - """Rows of pseudorandom token ids. - - Deliberately not ``arange``: ``hash_tensor`` is an XOR reduction, and - aligned runs of consecutive integers collide under it — ``XOR(6..11)`` - and ``XOR(12..17)`` are both 1 — which would make two visibly different - rows fingerprint the same. - """ - g = torch.Generator().manual_seed(seed) - return TensorDict( - { - "ids": torch.nested.nested_tensor( - [torch.randint(0, 32000, (n,), generator=g) for n in lengths], - layout=torch.jagged, - ) - }, - batch_size=[len(lengths)], - ) - - def test_uniform_jagged_rows_are_fingerprinted_per_row(): - """A jagged field whose rows happen to be uniform is a rectangle already - -- its values buffer reshapes to one as a view -- so it earns per-row + """A jagged field whose rows happen to be uniform is a rectangle already — + its values buffer reshapes to one as a view — so it earns per-row attribution for free. The batch-scoped fallback is an XOR over one shared buffer, which cannot see a permutation: two equal-length rows swapped by a mis-shard would round-trip clean.""" inner = _JaggedEcho() - client = _hash_client(inner) - ids = [f"u{i}" for i in range(4)] + client = _client(inner, verify_tensor_hash=True) + ids = _ids(4) client.put_samples( sample_ids=ids, partition_id="p", fields=_jagged_ids([6, 6, 6, 6]) ) @@ -1251,356 +1100,113 @@ def test_uniform_jagged_rows_are_fingerprinted_per_row(): client.close() -def test_uniform_put_read_back_ragged_is_a_mismatch_not_a_skip(): - """Row lengths changing between wire-in and wire-out is a divergence. It - has no row-level digest to compare against, but reporting it as a skipped - field would leave mismatches reading zero -- exactly the shape of a guard - that covers nothing.""" +def test_uniform_put_read_back_ragged_checks_row_lengths(): + """A row that changed length between wire-in and wire-out is a divergence. + The content is no longer comparable, so the field counts as an abstention, + but the lengths still are — and reporting the whole field as a skip would + leave mismatches reading zero, exactly the shape of a guard that covers + nothing. Failing the whole field instead reported 3584 mismatches per step + on a healthy 5-process run.""" inner = _JaggedEcho() - client = _hash_client(inner) - ids = [f"u{i}" for i in range(4)] + client = _client(inner, verify_tensor_hash=True) + ids = _ids(4) client.put_samples( sample_ids=ids, partition_id="p", fields=_jagged_ids([6, 6, 6, 6]) ) inner.rows[("p", "u2")]["ids"] = inner.rows[("p", "u2")]["ids"][:-2] client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) - assert client.snapshot()["hash_verify"]["mismatches"] > 0 + hv = client.snapshot()["hash_verify"] + assert hv["mismatches"] > 0, "a truncated row must not read as clean" + assert hv["fields_skipped"] == 1, "and the content it could not compare" client.close() -def test_delta_put_does_not_restate_the_scheme_of_untouched_fields(): - """write_columns puts one field into a partition written ragged earlier. - Holding the jagged/per-row choice per partition let that delta hand the - read side the wrong scheme for its own field, and every row came back a - false alarm.""" +def test_uniform_write_read_back_inside_a_ragged_batch_is_clean(): + """The false positive this cost: a shard written with uniform rows, read + back inside a batch whose *other* rows are ragged. Nothing diverged — every + recorded row still has the length it was written with — and the guard + reported 3584 mismatches per step on a healthy run until it compared lengths + instead of failing the field outright.""" inner = _JaggedEcho() - client = _hash_client(inner) - ids = [f"u{i}" for i in range(4)] + client = _client(inner, verify_tensor_hash=True) + ids = _ids(4) client.put_samples( - sample_ids=ids, partition_id="p", fields=_jagged_ids([2, 4, 6, 3]) + sample_ids=ids, partition_id="p", fields=_jagged_ids([6, 6, 6, 6]) ) - client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) - assert client.snapshot()["hash_verify"]["mismatches"] == 0 + # a later writer adds rows of a different length to the same partition + other = _ids(2, "v") + inner.rows[("p", other[0])] = {"ids": torch.randint(0, 32000, (3,))} + inner.rows[("p", other[1])] = {"ids": torch.randint(0, 32000, (9,))} + client.get_samples(sample_ids=ids + other, partition_id="p", select_fields=["ids"]) - # same field, rewritten uniform -- the later scheme must win - client.put_samples( - sample_ids=ids, partition_id="p", fields=_jagged_ids([6, 6, 6, 6], seed=99) - ) - client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) - assert client.snapshot()["hash_verify"]["mismatches"] == 0 + hv = client.snapshot()["hash_verify"] + assert hv["mismatches"] == 0, "no row changed length; nothing diverged" + assert hv["fields_skipped"] == 1, "content uncomparable, and counted" client.close() -def _busy_client(op_calls): - """A client that ran ``n`` calls of each named op this step.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - now = monotonic() - for op, (n, ms) in op_calls.items(): - for i in range(n): - client._emit(op, "p", 1, 1_000_000 * (1 + i % 5), now - ms / 1e3, "ok") - return client - - -def test_percent_of_dataplane_names_the_bottleneck_and_says_of_what(): - """The question a dashboard has to answer is "which op is expensive", - and 32 per-op line charts do not answer it. - - The denominator is data-plane time, not the step: ``by_op`` sums to 100 - by construction, so the largest is the bottleneck *within the data - plane*. Whether the data plane mattered at all is ``frac_of_step``, - which divides by the step's own clock -- here a tenth of a second of - data-plane work inside a 10 s step is 9% of one and 100% of the other. +def test_delta_put_does_not_restate_the_scheme_of_untouched_fields(): + """``write_columns`` puts one field into a partition written ragged earlier. + Holding the jagged/per-row choice per *partition* let that delta hand the + read side the wrong scheme for a field it never touched, and every row of + that field came back a false alarm. + + The second field is the whole point: with only one, a per-partition and a + per-field scheme are indistinguishable, because the only put there is + restates its own field either way. """ - client = _busy_client({"get": (100, 9.0), "put": (10, 1.0), "clear": (10, 0.1)}) - metrics = cluster_step_metrics( - merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 10.0 - ) - - by_op = { - k: v - for k, v in metrics.items() - if k.startswith("step/percent_of_dataplane/by_op/") - } - assert sum(by_op.values()) == pytest.approx(100.0), "percent of one total" - assert max(by_op, key=by_op.__getitem__) == "step/percent_of_dataplane/by_op/get" - assert by_op["step/percent_of_dataplane/by_op/get"] == pytest.approx( - 100 * 900 / 911, rel=0.05 - ) - # the two denominators are different questions and must not agree - assert metrics["step/frac_of_step"] == pytest.approx(0.0911, rel=0.1) - - -def test_headline_drops_per_op_detail_but_keeps_the_percentages(): - """Four ops times eight fields is 32 series saying one thing. The detail - is still computed -- the breakdown table is built from the same dict, so - the two cannot disagree -- but only the totals and percentages are - charted.""" - client = _busy_client({"get": (100, 9.0), "put": (10, 1.0), "clear": (10, 0.1)}) - metrics = cluster_step_metrics( - merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 1.0 - ) - head = headline_series(metrics) - - assert len(head) < len(metrics) / 2, f"{len(head)} of {len(metrics)}" - assert not [k for k in head if k.split("/")[1:2] in (["get"], ["put"], ["clear"])] - assert "step/percent_of_dataplane/by_op/get" in head - assert "step/wall_ms" in head and "step/frac_of_step" in head - # and the detail the table needs survives in the full dict - assert breakdown_table(metrics)[1], "table still has rows" - client.close() - - -def test_cluster_step_max_reopens_each_step(): - """A maximum cannot be differenced out of a cumulative counter, so the - cluster path reported the lifetime max: after one 50 ms call every later - step still read 50 ms, which is the same defect as a cumulative - percentile. The reader resets the window as it reads.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - prev, seen = {}, [] - for slowest_ms in (5.0, 50.0, 5.0, 5.0): - now = monotonic() - for i in range(4): - client._emit( - "put", "p", 1, 1_000, now - (slowest_ms if i == 0 else 5.0) / 1e3, "ok" - ) - merged = merge_snapshots([client.snapshot(reset_step_window=True)]) - seen.append(cluster_step_metrics(merged, prev, 1.0)["step/by_op/put/max_ms"]) - prev = merged - - assert seen[1] == pytest.approx(50.0, abs=1.0), "the spike shows" - assert seen[2] == pytest.approx(5.0, abs=1.0), "and does not latch" - client.close() - - -def test_snapshot_leaves_the_step_window_alone_unless_asked(): - """``snapshot()`` is also how a human inspects a live client. Resetting - the step window on every call would let an inspection silently blank the - next step's max.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client._emit("put", "p", 1, 1_000, monotonic() - 0.030, "ok") - - assert client.snapshot()["by_op"]["put"]["step_max_ms"] >= 30.0 - assert client.snapshot()["by_op"]["put"]["step_max_ms"] >= 30.0, "still there" - assert client.snapshot(reset_step_window=True)["by_op"]["put"]["step_max_ms"] >= 30 - assert client.snapshot()["by_op"]["put"]["step_max_ms"] == 0.0, "window reopened" - client.close() - - -def test_breakdown_table_ignores_reserved_namespaces(): - """``step/self/overhead_ms`` has the same three-part shape as a per-op - series and was becoming a "self" row in the table beside put and get.""" - columns, rows = breakdown_table( - { - "step/by_op/put/calls": 2, - "step/by_op/put/wall_ms": 9.0, - "step/self/overhead_ms": 6.2, - "step/hash/mismatches": 0, - "step/percent_of_dataplane/by_cause/transfer": 40.0, - } - ) - assert [r[0] for r in rows] == ["put"], rows - assert columns[0] == "op" - - -def test_hash_counters_reach_both_scopes(): - """``_log_data_plane_metrics`` prefers the cluster path whenever the - fan-out reaches more than one process -- every real run -- and the - cluster path emitted no hash counters at all. With verify_tensor_hash - on, ``mismatches`` never reached the logger: a guard whose findings are - not reported is not a guard.""" - client = _hash_client(_CorruptingClient(field="ids", row=2)) - ids = [f"u{i}" for i in range(4)] - client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) - client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids", "lp"]) - - merged = merge_snapshots([client.snapshot(reset_step_window=True)]) - cluster = cluster_step_metrics(merged, {}, 1.0) - assert cluster["step/hash/mismatches"] == 1 - assert "step/hash/fields_skipped" in cluster, "abstentions visible too" - assert headline_series(cluster)["step/hash/mismatches"] == 1, "and charted" - client.close() - - -def test_hash_counters_absent_when_the_guard_is_off(): - """Five always-zero series on every run that never asked for the guard - would read as "checked, nothing wrong" rather than "not checked".""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client._emit("put", "p", 1, 1_000, monotonic() - 0.005, "ok") - merged = merge_snapshots([client.snapshot(reset_step_window=True)]) - - assert not [k for k in cluster_step_metrics(merged, {}, 1.0) if "hash" in k] - assert not [k for k in client.get_step_metrics(1.0) if "hash" in k] - client.close() - - -def test_measuring_cost_is_reported_in_both_scopes(): - """``step/self/*`` was cluster-only, so the single-process fallback - silently lacked the one number that says what observability cost.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client.register_partition( - partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["t"] - ) + inner = _JaggedEcho() + client = _client(inner, verify_tensor_hash=True) + ids = _ids(4) client.put_samples( - sample_ids=["a"], + sample_ids=ids, partition_id="p", - fields=TensorDict({"x": torch.zeros(1, 512)}, batch_size=[1]), - ) - driver = client.get_step_metrics(1.0) - - assert driver["step/self/overhead_ms"] > 0, "measuring is never free" - assert "step/self/frac" in driver - assert "step/self/overhead_ms" in headline_series(driver) - client.close() - - -def test_write_and_read_volume_are_not_computed_for_nobody(): - """They were dropped by headline_series, charted by nobody, and absent - from the table -- while per-op ``mb`` already splits the same traffic - finer (put's is the write volume, get's the read volume).""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - client._emit("put", "p", 1, 4_000, monotonic() - 0.005, "ok") - metrics = client.get_step_metrics(1.0) - - assert "step/bytes_written_mb" not in metrics - assert "step/bytes_read_mb" not in metrics - assert metrics["step/comm_volume_mb"] > 0, "the total is still reported" - assert metrics["step/volume_mb/by_op/put"] == pytest.approx(0.004), ( - "and split per op, under one key rather than two" - ) - client.close() - - -def test_data_plane_is_logged_before_the_step_is_committed(): - """``logger.log_metrics(..., step_finished=True)`` commits the wandb step, - and anything logged against a committed step is dropped. - - The data-plane call used to sit *after* that commit in ``grpo_train_sync``, - so every series was computed, printed to stdout, and silently discarded -- - invisible to a fake logger, and only caught by reading a real run back out - of the wandb API. This asserts the source order rather than the behaviour, - because the drop happens inside wandb. - """ - import pathlib - - import nemo_rl - - # Read the file rather than import it: ``grpo_sync`` pulls the training - # stack, and the rest of this suite runs without it. - # - # Comments stripped: the explanatory comment above the call site names - # ``step_finished=True`` too, and matching that would pass on any ordering. - source = ( - pathlib.Path(nemo_rl.__file__).parent / "algorithms" / "grpo_sync.py" - ).read_text() - body = "\n".join( - line - for line in source[source.index("def grpo_train_sync") :].splitlines() - if not line.lstrip().startswith("#") - ) - dp_call = body.index("_log_data_plane_metrics(policy, logger") - commit = body.index("step_finished=True") - assert dp_call < commit, ( - "_log_data_plane_metrics must run before the step_finished=True log; " - "wandb drops anything logged against an already-committed step" - ) - - -def test_per_op_volume_is_charted_and_sums_to_comm_volume(): - """``comm_volume_mb`` alone hides which direction the traffic went. On a - real step get moved 20.8 MB against put's 2.7 MB -- every DP rank fetches - its shard for the logprob pass and again for the train pass -- and a - single total cannot say that.""" - client = MetricsDataPlaneClient(NoOpDataPlaneClient()) - now = monotonic() - for _ in range(6): - client._emit("get", "p", 1, 3_000_000, now - 5 / 1e3, "ok") - for _ in range(2): - client._emit("put", "p", 1, 1_000_000, now - 5 / 1e3, "ok") - client._emit("clear", "p", 1, 0, now - 1 / 1e3, "ok") - - metrics = cluster_step_metrics( - merge_snapshots([client.snapshot(reset_step_window=True)]), {}, 1.0 - ) - head = headline_series(metrics) - - assert head["step/volume_mb/by_op/get"] == pytest.approx(18.0) - assert head["step/volume_mb/by_op/put"] == pytest.approx(2.0) - assert "step/volume_mb/by_op/clear" not in head, "no payload, not a zero" - # the parts account for the whole - per_op = sum(v for k, v in head.items() if k.startswith("step/volume_mb/")) - assert per_op == pytest.approx(head["step/comm_volume_mb"]) - client.close() - - -def test_volume_namespace_does_not_become_a_breakdown_row(): - """``step/volume_mb/by_op/get`` must feed the get row, not invent a - ``volume_mb`` op beside put and get.""" - columns, rows = breakdown_table( - { - "step/by_op/get/calls": 3, - "step/by_op/get/wall_ms": 9.0, - "step/by_op/get/mb": 18.0, - "step/volume_mb/by_op/get": 18.0, - } + fields=_jagged_ids([2, 4, 6, 3], with_dense=True), ) - assert [r[0] for r in rows] == ["get"], rows - assert rows[0][columns.index("mb")] == 18.0 - + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + assert client.snapshot()["hash_verify"]["mismatches"] == 0, "baseline" -def test_uniform_write_read_back_inside_a_ragged_batch_is_clean(): - """The false positive this cost: a shard written with uniform rows, read - back inside a batch whose *other* rows are ragged. Nothing diverged -- - every recorded row still has the length it was written with -- and the - guard reported 3584 mismatches per step on a healthy run until it - compared lengths instead of failing the field outright.""" - inner = _JaggedEcho() - client = _hash_client(inner) - ids = [f"u{i}" for i in range(4)] + # the delta names only ``lp``; ``ids`` must keep the ragged scheme it was + # written with, or its next read replays the uniform one and cries wolf client.put_samples( - sample_ids=ids, partition_id="p", fields=_jagged_ids([6, 6, 6, 6]) + sample_ids=ids, + partition_id="p", + fields=TensorDict({"lp": torch.ones(4, 6)}, batch_size=[4]), ) - # a later writer adds rows of a different length to the same partition - other = [f"v{i}" for i in range(2)] - inner.rows[("p", other[0])] = {"ids": torch.randint(0, 32000, (3,))} - inner.rows[("p", other[1])] = {"ids": torch.randint(0, 32000, (9,))} - client.get_samples(sample_ids=ids + other, partition_id="p", select_fields=["ids"]) + client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) hv = client.snapshot()["hash_verify"] - assert hv["mismatches"] == 0, "no row changed length; nothing diverged" - assert hv["fields_skipped"] == 1, "content uncomparable, and counted" + assert hv["mismatches"] == 0, "a delta put must not restate ids's scheme" + assert hv["fields_skipped"] == 0, "and the read stays comparable" client.close() -def test_implausible_mismatch_rate_is_called_out(caplog): - """Every row of every field wrong, identically, every step is not what a - broken wire looks like -- it is what a broken guard looks like. Both false - alarms this check has produced had that shape, and a reader should not - have to work that out from a raw count.""" +@pytest.mark.parametrize( + "mismatches,warns", + [ + # Every row of every field wrong, identically, every step is not what a + # broken wire looks like -- it is what a broken guard looks like. Both + # false alarms this check has produced had that shape. + (300, True), + # A handful of bad rows is exactly what the guard exists to report. + (3, False), + ], +) +def test_implausible_mismatch_rates_are_called_out(caplog, mismatches, warns): hv = { "rows_recorded": 100, "rows_checked": 100, - "mismatches": 300, + "mismatches": mismatches, "rows_unverified": 0, "fields_skipped": 0, + "guard_failures": 0, } with caplog.at_level(logging.WARNING): deltas = _hash_deltas(hv, {}) - assert deltas["step/hash/mismatches"] == 300 - assert "more likely a bug in the check" in caplog.text - + assert deltas["step/hash/mismatches"] == mismatches + assert ("more likely a bug in the check" in caplog.text) is warns -def test_a_believable_mismatch_rate_is_not_second_guessed(caplog): - """A handful of bad rows is exactly what the guard exists to report.""" - hv = { - "rows_recorded": 100, - "rows_checked": 100, - "mismatches": 3, - "rows_unverified": 0, - "fields_skipped": 0, - } - with caplog.at_level(logging.WARNING): - _hash_deltas(hv, {}) - assert "more likely a bug" not in caplog.text +# ── call-site ordering ───────────────────────────────────────────────── From 9867ee9b25d0e1e362161f6e49915dbf2430c4ba Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 26 Aug 2026 22:45:38 -0700 Subject: [PATCH 46/70] test(data-plane): gate every data-plane nightly on the wire guard Enable data_plane.observability.verify_tensor_hash on all 24 nightlies that run the data plane, and assert step/hash/{mismatches,guard_failures} are zero via check_metrics. The guard only counts and logs -- it never raises -- so CI was the only place a finding could stop anything, and nothing asserted on it. The assertion accepts either scope: grpo_sync logs data_plane/cluster/* once the worker fan-out lands, single_controller only ever logs data_plane/driver/*. Wrappers skip the assertion under TEST_DRYRUN: the delegated base runs in a subshell, so common.env's dryrun exit does not reach the wrapper tail and check_metrics would fail on a metrics.json that was never created. Signed-off-by: Zhiyu Li --- .../recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml | 2 +- .../llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml | 2 +- .../llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml | 2 +- ...instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml | 2 +- ...uct-2n8g-async-1off-single-controller-streaming2.yaml | 2 ++ ...8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml | 2 +- ....1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml | 2 +- ...-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml | 2 +- ...1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml | 2 +- ...1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml | 2 +- ...rpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml | 2 +- ...1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml | 2 +- .../grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml | 2 +- .../grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml | 2 +- ...po-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml | 2 +- ....5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml | 2 +- ...5b-instruct-1n8g-megatron-single-controller-sync.yaml | 2 ++ ...grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml | 2 +- ...3b-10n8g-megatron-cp2-r3-async-single-controller.yaml | 2 ++ ...grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml | 4 ++-- .../grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml | 2 +- ...-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml | 2 ++ ...sp-dynbatch-noncolocated-async-single-controller.yaml | 2 ++ ...-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml | 2 +- nemo_rl/algorithms/single_controller.py | 8 +++----- nemo_rl/data_plane/observability.py | 6 ------ .../test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh | 9 +++++++-- .../llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh | 9 +++++++-- .../llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh | 9 +++++++-- ...b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh | 9 +++++++-- ...truct-2n8g-async-1off-single-controller-streaming2.sh | 4 +++- ...1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh | 9 +++++++-- ...a3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh | 9 +++++++-- ...po-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh | 9 +++++++-- ...t-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh | 9 +++++++-- ...t-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh | 9 +++++++-- .../grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh | 9 +++++++-- ...2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh | 9 +++++++-- .../llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh | 9 +++++++-- .../llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh | 9 +++++++-- ...grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh | 9 +++++++-- ...n2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh | 9 +++++++-- ...1.5b-instruct-1n8g-megatron-single-controller-sync.sh | 4 +++- .../grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh | 9 +++++++-- ...ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh | 4 +++- .../grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh | 4 ++-- .../grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh | 9 +++++++-- ...pd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh | 4 +++- ...p2sp-dynbatch-noncolocated-async-single-controller.sh | 4 +++- ....5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh | 9 +++++++-- 50 files changed, 176 insertions(+), 74 deletions(-) diff --git a/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml b/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml index 5a711173ac3..b2e7feb0ccd 100644 --- a/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-deepscaler-1.5b-8K.yaml data_plane: enabled: true observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml b/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml index 01ea51f7945..895ea4ed6ee 100644 --- a/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-gemma3-1b-it-1n8g-fsdp2tp1.yaml data_plane: enabled: true observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml b/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml index ca900a52775..effaeefa9ef 100644 --- a/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-gspo-deepscaler-1.5b-8K.yaml data_plane: enabled: true observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml index 7ff662d1d73..34e8e5af14b 100644 --- a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3.yaml data_plane: enabled: true observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml index f751736341e..aea0573f1f9 100644 --- a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml @@ -20,6 +20,8 @@ checkpointing: # TransferQueue data plane is mandatory for the SingleController path. data_plane: enabled: true + observability: + verify_tensor_hash: true # SC async-RL runtime knobs. async_rl: diff --git a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml index 0eeb38d3960..6e69c526fb7 100644 --- a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated.yaml data_plane: enabled: true observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml index 2e8bb388958..ed213188b2e 100644 --- a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml index ed4d08776d6..8965990c20d 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3.yaml data_plane: enabled: true observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml index f6753cdd5c9..4ed9a2f64a6 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml index 16f630c3080..bb67e52ae04 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml index c57b60435e1..09cad6a522c 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-llama3.2-1b-instruct-1n8g-megatron.yaml data_plane: enabled: true observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml index e8ca7538473..4e58796dd8c 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml b/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml index ee35c4e86d3..52cd402799f 100644 --- a/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-moonlight-16ba3b-4n8g-megatron.yaml data_plane: enabled: true observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml index ded84294a79..3171128693c 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml index 7d70884b861..be5f563960a 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.yaml data_plane: enabled: true observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml index c6eb7880616..8cff067358c 100644 --- a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3.yaml data_plane: enabled: true observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml index b3fbaf26e53..9fe7b64dcbc 100644 --- a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml @@ -20,6 +20,8 @@ checkpointing: # TransferQueue data plane is mandatory for the SingleController path. data_plane: enabled: true + observability: + verify_tensor_hash: true # SC async-RL runtime knobs. async_rl: diff --git a/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml index ac9be52b94b..60096349373 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml index 7ca7aed34bf..02481dc9015 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml @@ -18,6 +18,8 @@ checkpointing: # TransferQueue is mandatory for the SingleController path. data_plane: enabled: true + observability: + verify_tensor_hash: true async_rl: sampler: diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml index 9041b9c7543..1fcdf3c2179 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml @@ -3,9 +3,9 @@ checkpointing: checkpoint_dir: results/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple data_plane: enabled: true + observability: + verify_tensor_hash: true logger: log_dir: logs/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple wandb: name: grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple - observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script diff --git a/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml index 0f889e39c00..f186d51c686 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml b/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml index fe7bdf3501f..555dadec318 100644 --- a/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml +++ b/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml @@ -17,6 +17,8 @@ async_rl: data_plane: enabled: true + observability: + verify_tensor_hash: true checkpointing: checkpoint_dir: results/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller diff --git a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml index 02e7492f6b8..eacf1f8d1d4 100644 --- a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml +++ b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml @@ -55,6 +55,8 @@ data_plane: local_buffer_size: 4294967296 reuse_registered_buffers: true staging_buffer_size: 268435456 + observability: + verify_tensor_hash: true # SC async-RL runtime knobs, replacing the nulled ppo.async_ppo block. async_rl: diff --git a/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml b/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml index 2fd278c8163..aeef1bfdebe 100644 --- a/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true # wire guard; gated to 0 by the wrapper script + verify_tensor_hash: true diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 870619c406a..8e4a9d990a7 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1552,11 +1552,9 @@ def _log_data_plane_metrics_impl(self, total_step_time: float) -> None: advantage stage's get plus the post-train clear; the bulk traffic is the trainer and generation workers' own clients, in their own processes with their own counters, so ``comm_volume_mb`` here is well - under what the job actually moved. ``grpo_sync`` prefers a cluster - view by fanning out over its policy worker group, but nothing - implements ``collect_data_plane_snapshots`` yet, so that path also - falls through to the driver's counters alone -- there is no cluster - view to mirror here until one exists. + under what the job actually moved. ``grpo_sync`` gets a cluster view by + fanning out over its policy worker group; this loop has no such group to + fan out over, so driver scope is all there is here. """ if not isinstance(self._dp_client, MetricsDataPlaneClient): return # observability disabled -> plain adapter diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index b235c9c6518..ea71bd9d10b 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -889,12 +889,6 @@ def cluster_step_metrics( """ wall_ms = merged["total_wall_ms"] - prev.get("total_wall_ms", 0.0) overhead_ms = merged["self_ms"] - prev.get("self_ms", 0.0) + collect_ms - # Not ``frac_of_step``: ``wall_ms`` here is the SUM over processes that - # ran concurrently, so dividing by one step's wall clock gives a number - # that exceeds 1 whenever they overlapped (measured 1.054 across ten - # processes) -- correct arithmetic, but it reads as "105% of the step". - # The mean fraction of the step a process spent in the data plane is - # bounded and answers the question people ask of it. n_procs = max(merged.get("n_processes", 1), 1) # step/ is a delta over this step; now/ is a level at this instant. # The unit alone does not distinguish them -- see README.md. diff --git a/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh b/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh index 9a7da2267c8..09e72697508 100755 --- a/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh +++ b/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh b/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh index 906929565bc..d6ba1702213 100755 --- a/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh +++ b/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh b/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh index 48c0c49b3d4..29385fc4474 100755 --- a/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh +++ b/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh index 9ea9b4c13b3..b01ada20907 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh index 56a639f74fb..b9cd68392ce 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh @@ -36,7 +36,9 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.1' \ 'data["train/token_mult_prob_error"]["10"] < 1.1' \ - 'mean(data["train/grad_norm"], 2, 0) > 0.06' + 'mean(data["train/grad_norm"], 2, 0) > 0.06' \ + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh index e7695ab284f..616e6ccd04c 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh index 44a8e722a57..42ea0a86254 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh index 9bc62ea48d4..e00189cfb7e 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh index bc6ce1d37ce..fca6386c352 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh index 85ae3077be2..5776c921194 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh index ee539083555..21356eef73c 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh index a7ca901f759..d154f54957b 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh b/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh index cf78ea15e54..fdcd5a14607 100755 --- a/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh +++ b/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh index 464c0808d53..c9452ebf21a 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh @@ -18,7 +18,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh index 3ac3167b7fd..3b9e13395ad 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh @@ -18,7 +18,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh index f4124f94095..a51bf254518 100755 --- a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh index fa5e77cc369..56775fa0b44 100755 --- a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh @@ -36,7 +36,9 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.1' \ 'data["train/token_mult_prob_error"]["450"] < 1.1' \ - 'mean(data["timing/train/total_step_time"], 2) < 25' + 'mean(data["timing/train/total_step_time"], 2) < 25' \ + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh b/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh index a4ad7f506d5..4b958d7595b 100755 --- a/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh @@ -18,7 +18,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh index 0508017c7a5..b8410b81755 100755 --- a/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh @@ -48,7 +48,9 @@ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then uv run tests/check_metrics.py $JSON_METRICS \ - 'median(data["train/token_mult_prob_error"]) < 1.02' + 'median(data["train/token_mult_prob_error"]) < 1.02' \ + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' uv run tools/check_r3_trace.py "$NRL_R3_TRACE_DIR" \ --require-forward-verify \ diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh index d44db1a1d77..eb85c9c52be 100755 --- a/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh @@ -37,8 +37,8 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma # The wire guard only counts; assert it found nothing. uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.02' \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' rm -rf "$CKPT_DIR" fi diff --git a/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh b/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh index 0b55ca1aafe..7ad99cc3df5 100755 --- a/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' diff --git a/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh index e33726d8652..42d47c94ac1 100755 --- a/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh +++ b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh @@ -55,7 +55,9 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/token_mult_prob_error") | . 'max(data["train/on_policy_distillation/teacher_batches"]) > 0' \ 'max(data["train/on_policy_distillation/teacher_samples"]) > 0' \ 'max(data["train/on_policy_distillation/teacher_model_unique"]) == 1' \ - 'max(data["train/on_policy_distillation/adv_std"]) < 0.01' + 'max(data["train/on_policy_distillation/adv_std"]) < 0.01' \ + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' rm -rf "$CKPT_DIR" fi diff --git a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh index f1666482055..7cc5f9f37ec 100755 --- a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh +++ b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh @@ -46,7 +46,9 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma 'len(data["train/critic/loss"]) == 40' \ 'max(data["train/critic/loss"]) < 1.5' \ 'mean(data["train/critic/explained_var"], range_start=-10) > 0.5' \ - 'mean(data["train/reward"], range_start=-10) > 0.75' + 'mean(data["train/reward"], range_start=-10) > 0.75' \ + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh b/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh index db95defea81..10a9cdb78ed 100755 --- a/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh +++ b/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh @@ -17,7 +17,12 @@ export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" # The wire guard only counts; assert it found nothing. +# The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit +# does not reach here. Skip explicitly, or the dryrun check in +# tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. +# `if` form, not `&& exit`: a false `[[ ]]` returns 1 and set -e would abort. +if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ - 'max(data["data_plane/cluster/step/hash/mismatches"]) == 0' \ - 'max(data["data_plane/cluster/step/hash/guard_failures"]) == 0' + 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' From 79efb1d2b48b869a652b5d69efd7e5f4731bcbaa Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 27 Aug 2026 03:06:24 -0700 Subject: [PATCH 47/70] refactor(data-plane): drop the latency/bandwidth fit The affine model was misspecified. It fits wall_ms ~ fixed_ms + n_bytes / bandwidth, which assumes per-request overhead is constant -- but real overhead scales with payload: msgpack encode, the .contiguous() copy taken before the pointer, RDMA memory registration, staging-buffer memcpy. Those costs get split arbitrarily between the intercept and the slope, so overhead_ms understates and bandwidth_mb_s overstates by an unknown margin. More size variation does not fix that; it just makes a wrong split look stable. It also usually declined to answer. GRPO writes the same batch shape every step, so request sizes are near-uniform and the identifiability gate (coefficient of variation < 0.05) returns regime=unidentifiable. Observed on a real RDMA-vs-SimpleStorage pair: put emitted no split on either arm, and get emitted one on RDMA but not on SimpleStorage -- the same workload producing the metric on one transport and not the other, which is useless for comparing them. A 5-step run logged 19 data_plane series and none came from the fit. And it was not free: three float multiply-accumulates per op on the hot path (sum_bytes_sq, sum_bytes_ms, sum_ms_sq) plus ok_wall_ms, retained per rank and merged across processes, to feed a model that mostly abstained. Removes fit_latency_bandwidth, _latency_split, those four accumulators and their merge entries, percent_of_dataplane/by_cause/*, and the two table columns. step/self/overhead_ms is unrelated and stays -- it is a direct measurement of what observability itself cost. Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 54 ++----- nemo_rl/data_plane/observability.py | 163 ++------------------ tests/unit/data_plane/test_observability.py | 40 +---- 3 files changed, 27 insertions(+), 230 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 3c8d9bf2ab0..fb7a50071f0 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -460,7 +460,7 @@ series. Do not read `comm_volume_mb` as cluster-wide volume. `OpStats` is additive on purpose, and `merge_snapshots()` uses it: the histogram buckets and the regression sufficient statistics from every rank *sum* into one cluster-wide view. Everything derived — percentiles, the -affine fit, throughput — is recomputed from the merged totals, never +throughput — is recomputed from the merged totals, never averaged across ranks (averaging per-rank percentiles does not give a cluster percentile). @@ -473,7 +473,6 @@ totals and `percent_of_dataplane`, with the per-op detail in a table beside them |---|---| | `step/frac_of_step` | is the data plane worth optimising at all? | | `step/percent_of_dataplane/by_op/{put,get,clear,register}` | which call is expensive? | -| `step/percent_of_dataplane/by_cause/{fixed_overhead,transfer}` | is that fixed per-request cost, or moving bytes? | | `step/wall_ms`, `step/comm_volume_mb` | how much time and traffic | | `step/volume_mb/by_op/{get,put}` | which direction that traffic went | @@ -490,29 +489,14 @@ put". Whether that time mattered against compute is the *other* metric: `frac_of_step` divides by the step's own wall clock. Read them together — a workload can be 43% put and still not be worth touching. -Each `by_cause` term names what it measures. `fixed_overhead` is the fitted -per-request constant — the cost of making a request at all, independent of -its size. `transfer` **is** the bandwidth term: bytes divided by the fitted -bandwidth. Reading a real TransferQueue step: +Reading a real TransferQueue step: ``` step/frac_of_step 0.074 the data plane is 7% of the step step/percent_of_dataplane/by_op/put 42.1 within it, put is the largest op -step/percent_of_dataplane/by_cause/fixed_overhead 52.3 half the time is per-request cost -step/percent_of_dataplane/by_cause/transfer 20.7 a fifth is bandwidth ``` -Overhead beats bandwidth roughly 2:1 here, so this workload is -overhead-dominated: batching into fewer, larger requests buys more than a -faster wire. Had `transfer` been the larger of the two, the conclusion -would invert. - -Two decompositions of one total, because either alone leaves the next -question unanswered. `by_op` sums to 100 by construction. `by_cause` sums -to *at most* 100: only ops with an identifiable affine fit can be split, so -the remainder (27% above — the `register` and `clear` calls, which move -no bytes and so have no bandwidth term to fit) is time that could not be -attributed rather than time that did not happen. +`by_op` sums to 100 by construction. `volume_mb` counts *transfers*, not data size, and two things follow from that. A byte written and later read is counted on both sides. And every @@ -537,19 +521,16 @@ and the wrong one for "what blocked the step". `data_plane/{cluster,driver}/breakdown` — one row per op, ordered by `percent_of_dataplane` so the bottleneck is the first line read: -| op | percent_of_dataplane | calls | wall_ms | mean_ms | max_ms | p50_ms | p90_ms | overhead_ms | transfer_ms | mb | -|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| -| put | 43.0 | 2 | 53.9 | 26.9 | 29.4 | — | — | 19.4 | 6.39 | 1.32 | -| get | 30.9 | 2 | 38.7 | 19.4 | 21.5 | — | — | 13.8 | 4.63 | 1.03 | -| register | 17.1 | 1 | 21.4 | 21.4 | 21.4 | — | — | — | — | 0 | -| clear | 8.93 | 1 | 11.2 | 11.2 | 11.2 | — | — | — | — | 0 | +| op | percent_of_dataplane | calls | wall_ms | mean_ms | max_ms | p50_ms | p90_ms | mb | +|---|---:|---:|---:|---:|---:|---:|---:|---:| +| put | 43.0 | 2 | 53.9 | 26.9 | 29.4 | — | — | 1.32 | +| get | 30.9 | 2 | 38.7 | 19.4 | 21.5 | — | — | 1.03 | +| register | 17.1 | 1 | 21.4 | 21.4 | 21.4 | — | — | 0 | +| clear | 8.93 | 1 | 11.2 | 11.2 | 11.2 | — | — | 0 | -Everything in ms on that row is **per call** except `wall_ms`, and the two -split terms add to `mean_ms`: that `put` row reads "each call cost 26.9 ms, -of which 19.4 was fixed per-request overhead and 6.39 was bandwidth at this -step's mean request size". Per call the overhead term *is* the fitted -constant, so it is comparable against a hardware number. `calls`, -`wall_ms` and `mb` are the only extensive columns. +Everything in ms on that row is **per call** except `wall_ms`: that `put` +row reads "each call cost 26.9 ms". `calls`, `wall_ms` and `mb` are the +only extensive columns. **Per-call figures describe the wire; sums describe the run.** `wall_ms` on the cluster path is summed over processes that ran concurrently, so it is process-time @@ -624,16 +605,7 @@ signal worth seeing rather than hiding. a chart mixing `wall_s` against `p90_ms` puts a 0.008 beside a 24.85 and reads as a data-plane bug rather than an axis one. -**Per step you get, per op tag:** `calls`, `wall_ms`, `max_ms`, and — when -the affine fit is trustworthy — `overhead_ms` and `transfer_ms`. Those last -two are the split of the op's time into fixed per-request cost and -bandwidth, in ms, and they *stack*: together they are the model's estimate -of the step's `wall_ms`, so charting them against the measured `wall_ms` -shows the breakdown and the model error in one picture. The coefficients -come from the cumulative fit (a model should be stable); the attribution is -per step, applied to that step's calls and bytes. A ratio was tried first -and was the wrong shape — cumulative and therefore flat, and unitless on an -axis of milliseconds. Percentiles come off the *step's* histogram delta, not the +**Per step you get, per op tag:** `calls`, `wall_ms`, `max_ms`. Percentiles come off the *step's* histogram delta, not the cumulative one -- a per-step p50 off a histogram that is never reset goes flat -- and each is emitted only when the step holds enough calls to resolve it: **p50 at 20, p90 at 40**, roughly four observations above the diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index ea71bd9d10b..64ed217e556 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -382,72 +382,6 @@ def _td_bytes(td: TensorDict | None, max_nodes: int = 10_000) -> int: return total -def fit_latency_bandwidth(s: dict[str, Any]) -> dict[str, Any]: - """Split an op's time into fixed per-request overhead vs transfer. - - Least-squares fit of ``wall_ms ~ fixed_ms + n_bytes / bandwidth`` over the - op's successful calls, from the accumulated sufficient statistics. - - The fit is only identifiable when request sizes actually vary: if every - request is the same size, infinitely many (overhead, bandwidth) pairs - reproduce the data, so ``regime`` reports ``"unidentifiable"`` rather - than an arbitrary split. That case is common in RL, where a step's - payloads are often uniform -- vary batch size to break the tie. - """ - n = s["calls"] - s["errors"] # successful calls; bytes/time pair only on those - sx, sy = float(s["n_bytes"]), s["ok_wall_ms"] - sxx, sxy = s["sum_bytes_sq"], s["sum_bytes_ms"] - if n < 3 or sx <= 0: - return {"regime": "insufficient-data"} - mean_x = sx / n - var_x = max(sxx / n - mean_x * mean_x, 0.0) - # Coefficient of variation: how much do request sizes actually differ? - if (var_x**0.5) / mean_x < 0.05: - return { - "regime": "unidentifiable", - "reason": "request sizes near-uniform; vary payload size to separate", - "mean_bytes": mean_x, - "mean_ms": sy / n, - } - denom = n * sxx - sx * sx - if denom <= 0: - return {"regime": "unidentifiable", "reason": "degenerate fit"} - slope = (n * sxy - sx * sy) / denom # ms per byte - fixed_ms = (sy - slope * sx) / n - if slope <= 0: - return {"regime": "noise-dominated", "fixed_ms": fixed_ms} - transfer_ms_at_mean = slope * mean_x - # R^2: does an affine model actually fit? Low R^2 means the split below - # is not trustworthy regardless of how clean the numbers look. - syy = s["sum_ms_sq"] - ss_tot = syy - sy * sy / n - ss_res = syy - fixed_ms * sy - slope * sxy - r_squared = 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0 - return { - "fixed_ms": fixed_ms, - "bandwidth_mb_s": 1.0 / (slope * 1000.0), - "transfer_ms_at_mean": transfer_ms_at_mean, - "mean_bytes": mean_x, - "r_squared": r_squared, - # A high R^2 does NOT validate the model: a chunked step function or - # a quadratic both fit a line at R^2 > 0.93 while producing a - # meaningless split. A negative intercept is physically impossible - # (no request costs less than zero to issue) and catches exactly the - # misspecification R^2 misses, so both must hold. - "model_trustworthy": r_squared >= 0.8 and fixed_ms >= 0.0, - "regime": ( - "overhead-dominated" - if fixed_ms > transfer_ms_at_mean - else "bandwidth-dominated" - ), - "overhead_frac_at_mean": ( - fixed_ms / (fixed_ms + transfer_ms_at_mean) - if (fixed_ms + transfer_ms_at_mean) > 0 - else 0.0 - ), - } - - def _step_deltas(snap: dict[str, Any], prev: dict[str, Any]) -> dict[str, float]: """The three series both step-metric paths report, identically. @@ -522,9 +456,6 @@ def _op_step_stats( ) ] row.update(_clamped_percentiles(step_hist, row["max_ms"])) - split = _latency_split(st["fit"], calls, op_bytes) - if split: - row["overhead_ms"], row["transfer_ms"] = split out[op] = row return out @@ -645,17 +576,8 @@ def _percent_of_dataplane(per_op: dict[str, dict[str, float]]) -> dict[str, floa the step's own wall clock. A workload can be 43% put and still not be worth touching. - Two decompositions of that one total, because either alone leaves the - next question unanswered: - - - ``by_op`` -- which call is expensive. Sums to 100 by construction. - - ``by_cause`` -- whether that time is fixed per-request cost or moving - bytes. ``overhead_ms``/``transfer_ms`` are per call, so they are - multiplied back by the call count to compare against the same total. - Only ops with an identifiable affine fit can be split (an op whose - requests were all one size cannot be), so this sums to *at most* 100 - and the remainder is time that could not be attributed -- not time - that did not happen. + ``by_op`` answers which call is expensive, and sums to 100 by + construction. On the cluster path ``wall_ms`` is summed over processes that ran concurrently, so these are percentages of aggregate process-time rather @@ -666,8 +588,8 @@ def _percent_of_dataplane(per_op: dict[str, dict[str, float]]) -> dict[str, floa per_op: Per-op step detail from :func:`_op_step_stats`. Returns: - ``step/percent_of_dataplane/by_op/{op}`` and ``step/percent_of_dataplane/by_cause/{cause}``, - each in percent. Empty when no op ran. + ``step/percent_of_dataplane/by_op/{op}`` in percent. Empty when no + op ran. """ total = sum(r["wall_ms"] for r in per_op.values()) if total <= 0: @@ -676,46 +598,9 @@ def _percent_of_dataplane(per_op: dict[str, dict[str, float]]) -> dict[str, floa f"step/percent_of_dataplane/by_op/{op}": 100.0 * r["wall_ms"] / total for op, r in per_op.items() } - for cause, field_name in ( - ("fixed_overhead", "overhead_ms"), - ("transfer", "transfer_ms"), - ): - attributed = sum( - r[field_name] * r["calls"] for r in per_op.values() if field_name in r - ) - if attributed > 0: - percent[f"step/percent_of_dataplane/by_cause/{cause}"] = ( - 100.0 * attributed / total - ) return percent -def _latency_split( - fit: dict[str, Any], calls: int, op_bytes: int -) -> tuple[float, float] | None: - """One call's time split into fixed overhead and transfer, in ms. - - Per call, like ``mean_ms``, and for the same reason: the extensive form - scales with DP degree and batch size and so describes the shape of the - run rather than the wire. Per call the overhead term *is* the fitted - per-request constant -- a property of the backend, comparable against a - hardware number -- and the transfer term is that bandwidth at this - step's mean request size. - - The two stack to ``mean_ms``, so charting them against it shows the - split and how well the affine model holds, in the same units and the - same scale as everything else per-op. - - ``None`` when the fit is not trustworthy, which includes the common RL - case of near-uniform request sizes where the split is mathematically - unrecoverable. - """ - if not fit.get("model_trustworthy") or calls <= 0: - return None - ms_per_byte = 1.0 / (fit["bandwidth_mb_s"] * 1e3) - return fit["fixed_ms"], ms_per_byte * (op_bytes / calls) - - def _clamped_percentiles(hist: list[int], max_ms: float) -> dict[str, float]: """Whichever of :data:`_QUANTILES` this sample can actually support. @@ -759,7 +644,6 @@ def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: stats["percent_of_total_ms"] = ( 100.0 * wall_ms / total_wall_ms if total_wall_ms else 0.0 ) - stats["fit"] = fit_latency_bandwidth(stats) hist = stats["latency_hist"] # Only what the sample supports; an absent key says "not enough # calls", which a zero would not. @@ -786,10 +670,6 @@ def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: "wall_ms", "n_bytes", "n_keys", - "ok_wall_ms", - "sum_bytes_sq", - "sum_bytes_ms", - "sum_ms_sq", ) _OP_MAX = ("max_ms", "step_max_ms") @@ -798,11 +678,10 @@ def merge_snapshots(snapshots: "list[dict[str, Any]]") -> dict[str, Any]: """Combine per-process snapshots into one cluster-wide view. This is what the accumulators were shaped for. Latency lives in fixed - histogram buckets and the latency/bandwidth model lives in sufficient - statistics precisely so both *add*: summing 256 per-rank histograms - gives the true cluster distribution, which averaging 256 per-rank - percentiles cannot. Everything derived — percentiles, throughput, the - affine fit — is recomputed from the merged totals, never averaged. + histogram buckets precisely so they *add*: summing 256 per-rank + histograms gives the true cluster distribution, which averaging 256 + per-rank percentiles cannot. Everything derived — percentiles, + throughput — is recomputed from the merged totals, never averaged. Counters sum. ``max_*`` fields take a maximum. ``peak_bytes_outstanding`` is the one approximation: summing per-process peaks assumes they @@ -958,10 +837,9 @@ def headline_series(metrics: dict[str, float]) -> dict[str, float]: # Per-op columns worth a row in the breakdown, in the order they read. -# ``overhead_ms``/``transfer_ms`` are only present when the affine fit is -# trustworthy, and ``p50_ms``/``p90_ms`` only above the sample gate, so a -# row carries None where a series was withheld rather than a zero that -# would read as a measurement. +# ``p50_ms``/``p90_ms`` are present only above the sample gate, so a row +# carries None where a series was withheld rather than a zero that would +# read as a measurement. _BREAKDOWN_COLUMNS = ( "percent_of_dataplane", "calls", @@ -970,8 +848,6 @@ def headline_series(metrics: dict[str, float]) -> dict[str, float]: "max_ms", "p50_ms", "p90_ms", - "overhead_ms", - "transfer_ms", "mb", ) @@ -1045,19 +921,6 @@ class OpStats: wall_ms: float = 0.0 n_bytes: int = 0 n_keys: int = 0 - # Sufficient statistics for the least-squares fit wall_ms ~ a + b*n_bytes, - # which separates fixed per-request overhead (a) from bandwidth (1/b). - # Successful calls only, so bytes and time refer to the same events. - # These are additive, so they can be summed across ranks and refit - # globally -- no need to ship per-event samples off each process. - ok_wall_ms: float = 0.0 - sum_bytes_sq: float = 0.0 - sum_bytes_ms: float = 0.0 - # Also needed for R^2, which is what tells us whether the affine model - # describes the data at all -- chunking, retries and queueing all make - # wall_ms non-linear in n_bytes, and a low R^2 is the signal to stop - # trusting the overhead/bandwidth split. - sum_ms_sq: float = 0.0 # Slowest single call, exact. The histogram below can only place a # call in a bucket, so at the handful of calls an op makes in one step # a percentile off it is bucket geometry rather than data -- a tail @@ -1681,11 +1544,7 @@ def _emit( stats.total_ops += 1 bucket.n_bytes += n_bytes bucket.n_keys += n_keys - bucket.ok_wall_ms += wall_ms bytes_f = float(n_bytes) - bucket.sum_bytes_sq += bytes_f * bytes_f - bucket.sum_bytes_ms += bytes_f * wall_ms - bucket.sum_ms_sq += wall_ms * wall_ms if op == "put" and n_keys: per_key = n_bytes // n_keys stats.last_put_bytes_per_key = per_key diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index cebb942bc52..00d1ddbf27f 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -81,8 +81,7 @@ def _put(client, ids, width=4): def _emit(client, latencies_ms, op="put", n_bytes=1_000): """Record one synthetic ``op`` call per entry in ``latencies_ms``. - ``n_bytes`` takes an int, or a callable of the call index for the size - variation the latency/bandwidth fit needs to be identifiable. + ``n_bytes`` takes an int, or a callable of the call index. """ now = monotonic() for i, ms in enumerate(latencies_ms): @@ -525,36 +524,6 @@ def test_cluster_step_max_reopens_each_step(): client.close() -def test_cluster_view_carries_the_latency_split(): - """The split was emitted on the driver path only, so the cluster view — the - one a real run logs — could never show it and its table column was always - empty. The cluster fit is the better one besides: more samples and more - size variation, and size variation is what makes the split identifiable.""" - fixed_ms, mb_per_s = 6.0, 400.0 - ranks = [] - for rank in range(4): - sizes = [100_000 * (i + 1 + rank) for i in range(8)] - ranks.append( - _rank( - [fixed_ms + b / (mb_per_s * 1e3) for b in sizes], - op="get", - n_bytes=lambda i, s=sizes: s[i], - ) - ) - metrics = cluster_step_metrics(merge_snapshots(ranks), {}, 1.0) - - assert "step/by_op/get/overhead_ms" in metrics, "cluster view must carry the split" - total = ( - metrics["step/by_op/get/overhead_ms"] + metrics["step/by_op/get/transfer_ms"] - ) - assert total == pytest.approx(metrics["step/by_op/get/mean_ms"], rel=0.05) - assert metrics["step/by_op/get/overhead_ms"] == pytest.approx(fixed_ms, rel=0.05) - - columns, rows = breakdown_table(metrics) - assert rows[0][columns.index("overhead_ms")] is not None - assert rows[0][columns.index("transfer_ms")] is not None - - # ── percentiles: clamping and sample gates ───────────────────────────── @@ -786,9 +755,8 @@ def test_breakdown_table_rows_by_op_worst_first(): def test_breakdown_table_leaves_withheld_series_empty(): - """A percentile below the sample gate, or a fit that is not trustworthy, is - absent from the series — the table must carry None there rather than a zero - that would read as a measurement.""" + """A percentile below the sample gate is absent from the series — the table + must carry None there rather than a zero that would read as a measurement.""" columns, rows = breakdown_table( { "step/by_op/put/calls": 3, @@ -798,7 +766,6 @@ def test_breakdown_table_leaves_withheld_series_empty(): ) assert rows[0][columns.index("p90_ms")] is None - assert rows[0][columns.index("overhead_ms")] is None assert rows[0][columns.index("calls")] == 3 @@ -818,7 +785,6 @@ def test_breakdown_table_ignores_reserved_namespaces(): "step/volume_mb/by_op/get": 18.0, "step/self/overhead_ms": 6.2, "step/hash/mismatches": 0, - "step/percent_of_dataplane/by_cause/transfer": 40.0, } ) From da99febaa28e9a3ceea45992de9dd04f4c32e83f Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 29 Aug 2026 15:54:34 -0700 Subject: [PATCH 48/70] feat(data-plane): report jagged pack/unpack time Padding cost was invisible to the per-op metrics, and asymmetrically so. pack_jagged_fields runs in the caller before put_samples is entered, so it never reached by_op at all; _from_wire runs inside the adapter's get_samples, where it was billed as transport. Both are CPU work proportional to payload, so a transport comparison was crediting the wire for densification. Adds step/codec/{pack_ms,unpack_ms}, measured with utils.timer and drained by snapshot() into cumulative counters so they merge across processes and difference per step through the paths every other counter already uses. Kept out of total_wall_ms: folding CPU work into it would redefine frac_of_step and percent_of_dataplane to mean something other than time spent in the data plane. The drain is gated on reset_step_window for the same reason step_max_ms is -- it is destructive, so an inspection snapshot must not delete time the step reader is about to report. Adds Timer.drain (pop-and-sum under one lock; reduce-then-reset drops a concurrent record) and a should_log flag on Timer.record, whose _fmt call otherwise dominates the cost of recording on a hot path. Coverage matches comm_volume and is documented beside the field: only a process that drains reports its own pad cost, and the rollout actor is not on the worker group the fan-out reaches, so kv_first_write is omitted. Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/adapters/transfer_queue.py | 62 ++++++--- nemo_rl/data_plane/codec.py | 121 ++++++++++++++---- nemo_rl/data_plane/observability.py | 28 ++++ nemo_rl/utils/timer.py | 26 +++- tests/unit/data_plane/conftest.py | 15 +++ tests/unit/data_plane/test_observability.py | 46 +++++++ 6 files changed, 250 insertions(+), 48 deletions(-) diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index 692f6d6daaf..fe112ddcffb 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -48,6 +48,7 @@ from tensordict import TensorDict from nemo_rl.data_plane.adapters.transfer_queue_env import rail_link_layers +from nemo_rl.data_plane.codec import timed_codec from nemo_rl.data_plane.interfaces import ( DataPlaneClient, DataPlaneConfig, @@ -743,7 +744,7 @@ def _from_wire(td: TensorDict) -> TensorDict: ``codec.materialize`` applies the same exclusion. """ # NonTensorData / NonTensorStack leaves are only visible via td.keys(), - # not keys(leaves_only=True) — iterating the latter would silently drop + # not keys(leaves_only=True) -- iterating the latter would silently drop # them from the rebuilt dict. # Deferred: ``multimodal_utils`` pulls PIL, requests and a few hundred # transformers submodules, and this adapter is imported by every process @@ -751,26 +752,45 @@ def _from_wire(td: TensorDict) -> TensorDict: # for the same reason. from nemo_rl.data.multimodal_utils import PACKED_MULTIMODAL_FIELDS - new_dict: dict[str, Any] = {} - changed = False - for k in td.keys(): - v = td.get(k) - field_name = str(k) - if ( - isinstance(v, torch.Tensor) - and v.is_nested - and field_name not in PACKED_MULTIMODAL_FIELDS - ): - rows = list(v.unbind()) - if rows and all(row.shape == rows[0].shape for row in rows[1:]): - v = torch.stack(rows) - changed = True - new_dict[field_name] = v - if not changed: - return td - new_td = TensorDict(new_dict, batch_size=td.batch_size) - _assert_no_key_loss(new_dict, new_td, "_from_wire") - return new_td + with timed_codec("unpack"): + new_dict: dict[str, Any] = {} + changed = False + for k in td.keys(): + v = td.get(k) + field_name = str(k) + if ( + isinstance(v, torch.Tensor) + and v.is_nested + and field_name not in PACKED_MULTIMODAL_FIELDS + ): + rows = list(v.unbind()) + if rows and all(row.shape == rows[0].shape for row in rows[1:]): + v = torch.stack(rows) + changed = True + if field_name in PROMOTE_1D_FIELDS: + if not isinstance(v, torch.Tensor) or v.is_nested: + raise ValueError( + f"Mooncake scalar field {field_name!r} could not be " + "restored as a dense tensor." + ) + if v.dim() == 1: + new_dict[field_name] = v + elif v.dim() == 2 and v.shape[-1] == 1: + new_dict[field_name] = v.squeeze(-1).contiguous() + changed = True + else: + raise ValueError( + f"Mooncake scalar field {field_name!r} must decode as " + f"(N,) or (N, 1), got shape {tuple(v.shape)}." + ) + else: + new_dict[field_name] = v + if not changed: + # The traversal still ran; only the rebuild was skipped. + return td + new_td = TensorDict(new_dict, batch_size=td.batch_size) + _assert_no_key_loss(new_dict, new_td, "_from_wire") + return new_td class TQDataPlaneClient(DataPlaneClient): diff --git a/nemo_rl/data_plane/codec.py b/nemo_rl/data_plane/codec.py index 93bd7454ead..ad06dfdd879 100644 --- a/nemo_rl/data_plane/codec.py +++ b/nemo_rl/data_plane/codec.py @@ -35,6 +35,9 @@ from __future__ import annotations +import time +from collections.abc import Iterator +from contextlib import contextmanager from typing import TYPE_CHECKING, Any import numpy as np @@ -42,6 +45,76 @@ from tensordict import TensorDict, TensorDictBase from nemo_rl.data_plane.schema import Layout +from nemo_rl.utils.timer import ThreadSafeTimer + +# Pad/unpad cost, which the per-op metrics cannot see: packing runs in the +# caller before ``put_samples`` is entered, and ``_from_wire`` runs inside the +# adapter's ``get_samples``, where it is billed as transport. Both are real CPU +# work proportional to payload size. +# +# Module-level, not threaded through: every call site does have a client handle +# (``column_io`` takes ``dp_client``; ``_from_wire`` is reached from an instance +# method), so threading is possible at roughly fifteen lines across four sites. +# The global buys reach for a free function at the cost of process-scoped state +# that only one reader per process may drain. +_CODEC_TIMER = ThreadSafeTimer() + + +def record_codec_s(phase: str, elapsed_s: float) -> None: + """Record one pad/unpad measurement, in seconds. + + Prefer :func:`timed_codec`; this is for callers that already measured. + + Args: + phase: ``"pack"`` or ``"unpack"``. + elapsed_s: Seconds spent, as returned by ``time.perf_counter()`` deltas. + """ + # should_log=False: Timer._fmt builds a timestamp and joins the context on + # every call, which costs more than the measurement itself on this path. + _CODEC_TIMER.record(phase, elapsed_s, should_log=False) + + +@contextmanager +def timed_codec(phase: str) -> Iterator[None]: + """Time a pad/unpad block, recording on every exit path. + + Records in ``finally`` because the blocks it wraps return from more than + one place -- an early return once slipped past a hand-written bracket and + silently dropped every no-op unpack from the metric. + + Not :meth:`Timer.time`: ``Timer.start`` raises if the label is already + running, and both phases run concurrently (the single-controller loop + dispatches through ``asyncio.to_thread``). + """ + started = time.perf_counter() + try: + yield + finally: + record_codec_s(phase, time.perf_counter() - started) + + +def drain_codec_ms() -> dict[str, float]: + """Milliseconds spent packing and unpacking since the last drain. + + ``Timer.drain`` pops and sums under one lock: ``reduce`` then ``reset`` + would drop any sample recorded between them, and both phases run + concurrently. + + Not every packing process has a reader -- the rollout actor calls + ``pack_jagged_fields`` but is not on the policy worker group, so nothing + drains it. Its samples accumulate unread, which is why the caller that + *does* drain should do so every step. + + Returns: + ``{"pack": ms, "unpack": ms}``, omitting a phase that did not run. + """ + out: dict[str, float] = {} + for phase in ("pack", "unpack"): + total = _CODEC_TIMER.drain(phase) + if total: + out[phase] = total * 1e3 + return out + if TYPE_CHECKING: # Type-only import. At runtime, BatchedDataDict is loaded lazily @@ -159,30 +232,32 @@ def pack_jagged_fields( ``TensorDict`` with ``batch_size=[N]`` (N from ``lengths`` if given, else 0) ready for ``put_samples``. """ - n = int(lengths.shape[0]) if lengths is not None else 0 - token_aligned_fields = token_aligned_fields or frozenset() - packed: dict[str, Any] = {} - for k, v in fields.items(): - if isinstance(v, np.ndarray) and v.dtype == object: - # tensordict==0.12.2 wire bug: a NonTensorStack stored as a - # TensorDict leaf returns as a LinkedList on parent - # __getitem__, losing identity. ndarray(dtype=object) - # round-trips intact. - packed[k] = v - elif isinstance(v, torch.Tensor): - if lengths is not None and k in token_aligned_fields: - packed[k] = pack_per_token_field(v, lengths) + with timed_codec("pack"): + n = int(lengths.shape[0]) if lengths is not None else 0 + token_aligned_fields = token_aligned_fields or frozenset() + packed: dict[str, Any] = {} + for k, v in fields.items(): + if isinstance(v, np.ndarray) and v.dtype == object: + # tensordict==0.12.2 wire bug: a NonTensorStack stored as a + # TensorDict leaf returns as a LinkedList on parent + # __getitem__, losing identity. ndarray(dtype=object) + # round-trips intact. + packed[k] = v + elif isinstance(v, torch.Tensor): + if lengths is not None and k in token_aligned_fields: + packed[k] = pack_per_token_field(v, lengths) + else: + packed[k] = v.detach().contiguous() else: - packed[k] = v.detach().contiguous() - else: - raise TypeError( - f"pack_jagged_fields: unsupported value type for {k!r}: {type(v)}. " - "Use torch.Tensor or np.ndarray(dtype=object). PackedTensor " - "must be converted to torch.nested at the wire boundary " - "(see sync_rollout_actor.py) so the codec's dispatch stays " - "binary." - ) - return TensorDict(packed, batch_size=[n]) + raise TypeError( + f"pack_jagged_fields: unsupported value type for {k!r}: {type(v)}. " + "Use torch.Tensor or np.ndarray(dtype=object). PackedTensor " + "must be converted to torch.nested at the wire boundary " + "(see sync_rollout_actor.py) so the codec's dispatch stays " + "binary." + ) + out = TensorDict(packed, batch_size=[n]) + return out def pack_per_token_field(val: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 64ed217e556..e98ecaf3d9d 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -60,6 +60,7 @@ class DataPlaneEvent(TypedDict): import torch from tensordict import NonTensorData, NonTensorStack, TensorDict, TensorDictBase +from nemo_rl.data_plane.codec import drain_codec_ms from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta logger = logging.getLogger(__name__) @@ -402,6 +403,8 @@ def _step_deltas(snap: dict[str, Any], prev: dict[str, Any]) -> dict[str, float] ) / 1e6, "now/bytes_outstanding_mb": snap["bytes_outstanding"] / 1e6, + "step/codec/pack_ms": snap["pack_ms"] - prev.get("pack_ms", 0.0), + "step/codec/unpack_ms": snap["unpack_ms"] - prev.get("unpack_ms", 0.0), } @@ -662,6 +665,8 @@ def _derive_op_metrics(by_op: dict[str, Any], total_wall_ms: float) -> None: "peak_bytes_outstanding", "n_keys_outstanding", "self_ms", + "pack_ms", + "unpack_ms", ) _SNAPSHOT_MAX = ("max_bytes_per_key_seen", "last_put_bytes_per_key") _OP_SUM = ( @@ -815,6 +820,7 @@ def cluster_step_metrics( "step/volume_mb/", "step/hash/", "step/self/", + "step/codec/", ) @@ -986,6 +992,20 @@ class DataPlaneStats: # observability bill next to the thing it is observing rather than # taking a benchmark's word for it. self_ms: float = 0.0 + # Jagged pad/unpad CPU cost, drained from the codec timer. Packing runs + # in the caller before ``put_samples`` and so is invisible to ``by_op``; + # unpacking runs inside the adapter's ``get_samples`` and is otherwise + # billed as transport. Kept out of ``total_wall_ms`` so ``frac_of_step`` + # and ``percent_of_dataplane`` keep meaning time spent in the data plane. + # + # Same coverage gap as ``comm_volume`` and for the same reason: only a + # process that drains the codec timer reports its own pad cost, and the + # rollout actor is not on the policy worker group the fan-out reaches. So + # ``pack_ms`` omits ``kv_first_write``, the largest single pack in the job. + # The single-controller path has no fan-out at all, so there it is + # driver-only on both counters. + pack_ms: float = 0.0 + unpack_ms: float = 0.0 hash_verify: HashStats = field(default_factory=HashStats) @@ -1058,6 +1078,14 @@ def snapshot(self, reset_step_window: bool = False) -> dict[str, Any]: that has to. Left off by default so an inspection snapshot never disturbs the step series. """ + # Gated on reset_step_window for the same reason step_max_ms is: the + # codec timer is drained destructively, so an inspection snapshot that + # took it would delete that time from the series the step reader + # reports. Both callers that consume a step pass True. + if reset_step_window: + codec = drain_codec_ms() + self._stats.pack_ms += codec.get("pack", 0.0) + self._stats.unpack_ms += codec.get("unpack", 0.0) out = asdict(self._stats) out["n_keys_outstanding"] = sum( len(k) for k in self._keys_by_partition.values() diff --git a/nemo_rl/utils/timer.py b/nemo_rl/utils/timer.py index 69a5d79989c..51167970134 100644 --- a/nemo_rl/utils/timer.py +++ b/nemo_rl/utils/timer.py @@ -136,7 +136,7 @@ def stop(self, label: str, should_log: bool = True) -> float: logger.debug(self._fmt(label, f"end elapsed={elapsed:.4f}s")) return elapsed - def record(self, label: str, elapsed: float) -> None: + def record(self, label: str, elapsed: float, should_log: bool = True) -> None: """Append a pre-measured duration without start/stop. Useful when the caller has already measured the elapsed time @@ -145,11 +145,25 @@ def record(self, label: str, elapsed: float) -> None: Args: label: The timing label to record under elapsed: The elapsed time in seconds + should_log: Emit the debug line. ``_fmt`` builds a timestamp and + joins the context on every call, which dominates the cost of + ``record`` itself -- pass ``False`` on a hot path. """ if label not in self._timers: self._timers[label] = [] self._timers[label].append(elapsed) - logger.debug(self._fmt(label, f"record elapsed={elapsed:.4f}s")) + if should_log: + logger.debug(self._fmt(label, f"record elapsed={elapsed:.4f}s")) + + def drain(self, label: str) -> float: + """Sum and clear ``label`` in one step. Returns 0.0 if it never ran. + + ``reduce`` then ``reset`` is two acquisitions on + :class:`ThreadSafeTimer`, so a concurrent ``record`` between them is + lost. Callers that consume-and-forget want this instead. + """ + samples = self._timers.pop(label, None) + return float(sum(samples)) if samples else 0.0 def mark(self, label: str, metadata: Optional[dict] = None) -> float: """Record a point-in-time event at the current Unix epoch. @@ -356,9 +370,13 @@ def stop(self, label: str, should_log: bool = True) -> float: with self._lock: return super().stop(label, should_log) - def record(self, label: str, elapsed: float) -> None: + def record(self, label: str, elapsed: float, should_log: bool = True) -> None: + with self._lock: + super().record(label, elapsed, should_log) + + def drain(self, label: str) -> float: with self._lock: - super().record(label, elapsed) + return super().drain(label) def mark(self, label: str, metadata: Optional[dict] = None) -> float: with self._lock: diff --git a/tests/unit/data_plane/conftest.py b/tests/unit/data_plane/conftest.py index 0f7e2da3019..af491791bbf 100644 --- a/tests/unit/data_plane/conftest.py +++ b/tests/unit/data_plane/conftest.py @@ -103,3 +103,18 @@ def tq_client_backends(request): (see module docstring). """ return request.getfixturevalue(f"_session_tq_client_{request.param}") + + +@pytest.fixture(autouse=True) +def _isolate_codec_timer(): + """Drain the module-level codec timer around every test in this package. + + ``pack_jagged_fields`` records into a process-global timer, so a test that + packs (the codec and column_io suites) leaves residue that the next + ``get_step_metrics`` would report as its own ``step/codec/pack_ms``. + """ + from nemo_rl.data_plane.codec import drain_codec_ms + + drain_codec_ms() + yield + drain_codec_ms() diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 00d1ddbf27f..c45f5bd2be3 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -1176,3 +1176,49 @@ def test_implausible_mismatch_rates_are_called_out(caplog, mismatches, warns): # ── call-site ordering ───────────────────────────────────────────────── + + +def test_codec_pack_unpack_time_is_reported_separately(): + """Jagged pad/unpad is real CPU cost the per-op metrics cannot see. + + ``pack_jagged_fields`` runs in the caller before ``put_samples`` is + entered, so it never reaches ``by_op``; ``_from_wire`` runs inside the + adapter's ``get_samples``, where it would otherwise be billed as + transport. Both are drained from the codec timer into their own series, + deliberately outside ``total_wall_ms`` so ``frac_of_step`` keeps meaning + time spent in the data plane rather than time spent on CPU around it. + """ + from nemo_rl.data_plane import codec + + client = _client(register=False) + codec.record_codec_s("pack", 0.010) # 10 ms + codec.record_codec_s("unpack", 0.004) # 4 ms + + metrics = client.get_step_metrics(1.0) + assert metrics["step/codec/pack_ms"] == pytest.approx(10.0, rel=1e-3) + assert metrics["step/codec/unpack_ms"] == pytest.approx(4.0, rel=1e-3) + # not folded into the transport totals + assert metrics["step/wall_ms"] == 0.0 + # and charted, not just tabulated + assert "step/codec/pack_ms" in headline_series(metrics) + + # drained exactly once: a second step reports zero, not the same 10 ms + assert client.get_step_metrics(1.0)["step/codec/pack_ms"] == 0.0 + client.close() + + +def test_inspection_snapshot_does_not_steal_codec_time(): + """The drain is destructive, so only the reader that closes the step window + may take it. ``snapshot()`` is also how a human inspects a live client, and + an unguarded drain there would delete the time the step reader is about to + report -- the same hazard ``step_max_ms`` is gated for.""" + from nemo_rl.data_plane import codec + + client = _client(register=False) + codec.record_codec_s("pack", 0.010) + + client.snapshot() # inspection: must not consume it + assert client.get_step_metrics(1.0)["step/codec/pack_ms"] == pytest.approx( + 10.0, rel=1e-3 + ) + client.close() From 888d8fd08d2db0bbd87b6749d23a57950f7bfb30 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 29 Aug 2026 20:15:57 -0700 Subject: [PATCH 49/70] refactor(data-plane): report charted durations in seconds A real run logged step/wall_ms = 78800.9, which reads as noise beside a 674 s step clock. The rest of nemo_rl already times in seconds -- utils.timer returns perf_counter deltas, timing/setup/*_time_s, and total_step_time, which is the very denominator frac_of_step divides by -- so the data plane was the outlier. step/wall_ms -> step/wall_s, step/codec/{pack,unpack}_ms -> _s. step/self/overhead_ms stays ms: at 73 ms it is sub-second, and it is a cost-of-measurement figure read against the per-op table rather than against the step clock. The breakdown table stays ms throughout, where sub-millisecond percentiles are still legible -- p50_ms=0.011 beats p50_s=1.1e-05. No precision cost: the logger stores float32 and floating point is scale-invariant, so /1e3 moves the decimal point rather than truncating (measured ~1e-8 relative, against float32's ~1e-7 inherent). The README previously asserted every duration is _ms with no exceptions, which is now false; its Units paragraph states the split instead. Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 2 +- nemo_rl/algorithms/single_controller.py | 2 +- nemo_rl/data_plane/README.md | 22 ++++++++++------ nemo_rl/data_plane/observability.py | 28 ++++++++++++++------- tests/unit/algorithms/test_grpo.py | 2 +- tests/unit/data_plane/conftest.py | 2 +- tests/unit/data_plane/test_observability.py | 20 +++++++-------- 7 files changed, 48 insertions(+), 30 deletions(-) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index fc5468ee725..b3975430640 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -468,7 +468,7 @@ def _log_data_plane_metrics_impl( logger.log_metrics(headline_series(metrics), step, prefix="data_plane/driver") _log_breakdown(logger, metrics, step, "data_plane/driver/breakdown") print( - f" • data plane: {metrics['step/wall_ms']:.0f}ms, " + f" • data plane: {metrics['step/wall_s']:.2f}s, " f"{metrics['step/comm_volume_mb']:.1f} MB moved" ) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 8e4a9d990a7..82e9744b0bf 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1568,7 +1568,7 @@ def _log_data_plane_metrics_impl(self, total_step_time: float) -> None: if rows: self._logger.log_table(columns, rows, step, "data_plane/driver/breakdown") print( - f" • data plane: {metrics['step/wall_ms']:.0f}ms, " + f" • data plane: {metrics['step/wall_s']:.2f}s, " f"{metrics['step/comm_volume_mb']:.1f} MB moved", flush=True, ) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index fb7a50071f0..e2f6d049f5f 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -473,15 +473,16 @@ totals and `percent_of_dataplane`, with the per-op detail in a table beside them |---|---| | `step/frac_of_step` | is the data plane worth optimising at all? | | `step/percent_of_dataplane/by_op/{put,get,clear,register}` | which call is expensive? | -| `step/wall_ms`, `step/comm_volume_mb` | how much time and traffic | +| `step/wall_s`, `step/comm_volume_mb` | how much time and traffic | | `step/volume_mb/by_op/{get,put}` | which direction that traffic went | - -Per-op detail is published under `step/by_op//` and feeds the -breakdown table rather than a chart. +| `step/codec/{pack_s,unpack_s}` | jagged pad/unpad cost, which `by_op` cannot see | | `now/bytes_outstanding_mb`, `now/n_processes` | occupancy, fan-out width | | `step/self/{overhead_ms,frac}` | what measuring cost | | `step/hash/*` | only with `verify_tensor_hash` on | +Per-op detail is published under `step/by_op//` and feeds the +breakdown table rather than a chart. + **`percent_of_dataplane` is a percentage of data-plane time, not of the step.** The denominator is `sum(wall_ms)` over the ops that ran, so `by_op/put = 43` reads "43% of the time spent inside the data plane went to @@ -601,9 +602,16 @@ It is deliberately not clamped to 100%. Against a fast backend the ratio can exceed 1, meaning measuring cost more than the operation measured — a signal worth seeing rather than hiding. -**Units:** every duration is `_ms`, every volume is `_mb`, no exceptions — -a chart mixing `wall_s` against `p90_ms` puts a 0.008 beside a 24.85 and -reads as a data-plane bug rather than an axis one. +**Units:** charted *durations* are seconds — `step/wall_s`, +`step/codec/{pack_s,unpack_s}` — so they sit beside +`timing/train/total_step_time`. The one charted duration that is not is +`step/self/overhead_ms`: it is a cost-of-measurement figure, read against the +breakdown table rather than against the step clock. The table itself is ms +throughout (`wall_ms`, `mean_ms`, `max_ms`, `p50_ms`, `p90_ms`), where +sub-second per-call figures stay legible. Volumes are always `_mb`. What is not +allowed is mixing units *within one chart*: `step/wall_s` on the same axis as +`p90_ms` puts a 0.008 next to a 24.85 and reads as a data-plane bug rather than +an axis one. **Per step you get, per op tag:** `calls`, `wall_ms`, `max_ms`. Percentiles come off the *step's* histogram delta, not the cumulative one -- a per-step p50 off a histogram that is never reset goes diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index e98ecaf3d9d..2507894f0e9 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -396,15 +396,26 @@ def _step_deltas(snap: dict[str, Any], prev: dict[str, Any]) -> dict[str, float] volume and get's is the read volume, split finer than a global pair would be. """ + + def _delta_s(field: str) -> float: + """A millisecond accumulator differenced into the charted seconds. + + Every ``_s`` series goes through here so a new one cannot forget the + conversion and chart milliseconds under a seconds name. Seconds + because these sit beside ``timing/train/total_step_time``: a real + step logged 78800.9 ms, which reads as noise against a 674 s clock. + """ + return (snap[field] - prev.get(field, 0.0)) / 1e3 + return { - "step/wall_ms": snap["total_wall_ms"] - prev.get("total_wall_ms", 0.0), + "step/wall_s": _delta_s("total_wall_ms"), "step/comm_volume_mb": ( snap["comm_volume_bytes"] - prev.get("comm_volume_bytes", 0) ) / 1e6, "now/bytes_outstanding_mb": snap["bytes_outstanding"] / 1e6, - "step/codec/pack_ms": snap["pack_ms"] - prev.get("pack_ms", 0.0), - "step/codec/unpack_ms": snap["unpack_ms"] - prev.get("unpack_ms", 0.0), + "step/codec/pack_s": _delta_s("pack_ms"), + "step/codec/unpack_s": _delta_s("unpack_ms"), } @@ -809,7 +820,7 @@ def cluster_step_metrics( # The full dict is still returned, so the table and the series are derived # from one computation and cannot disagree. _HEADLINE = ( - "step/wall_ms", + "step/wall_s", "step/frac_of_step", "step/comm_volume_mb", "now/bytes_outstanding_mb", @@ -1121,11 +1132,10 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: self._prev_snapshot = snap wall_ms = snap["total_wall_ms"] - prev.get("total_wall_ms", 0.0) - # Every duration is ms and every volume is MB, with no exceptions: - # a chart that mixes wall_s against p90_ms puts a 0.008 next to a - # 24.85 and reads as a bug in the data plane rather than in the - # axis. GB was the same problem one dimension over -- a realistic - # step moved 0.00017 GB. + # Units are not mixed within one chart: durations charted beside the + # step clock are seconds, the per-op table is ms throughout, volumes + # are always MB. GB was the same problem one dimension over -- a + # realistic step moved 0.00017 GB. metrics = _step_deltas(snap, prev) metrics["step/frac_of_step"] = ( (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0 diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 13db64b4503..281b6701508 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -6197,7 +6197,7 @@ def test_grpo_train_sync_logs_data_plane_metrics_before_committing_the_step( payload = calls[dp[0]].args[0] assert payload["step/comm_volume_mb"] > 0, "the put moved bytes; the series says 0" - for key in ("step/wall_ms", "step/frac_of_step", "step/self/overhead_ms"): + for key in ("step/wall_s", "step/frac_of_step", "step/self/overhead_ms"): assert key in payload, f"{key} missing from {sorted(payload)}" assert [ diff --git a/tests/unit/data_plane/conftest.py b/tests/unit/data_plane/conftest.py index af491791bbf..44f70455733 100644 --- a/tests/unit/data_plane/conftest.py +++ b/tests/unit/data_plane/conftest.py @@ -111,7 +111,7 @@ def _isolate_codec_timer(): ``pack_jagged_fields`` records into a process-global timer, so a test that packs (the codec and column_io suites) leaves residue that the next - ``get_step_metrics`` would report as its own ``step/codec/pack_ms``. + ``get_step_metrics`` would report as its own ``step/codec/pack_s``. """ from nemo_rl.data_plane.codec import drain_codec_ms diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index c45f5bd2be3..98563501552 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -694,7 +694,7 @@ def test_headline_drops_per_op_detail_but_keeps_the_percentages(): assert len(head) < len(metrics), f"{len(head)} of {len(metrics)}" assert not [k for k in head if k.split("/")[1:2] in (["get"], ["put"], ["clear"])] assert "step/percent_of_dataplane/by_op/get" in head - assert "step/wall_ms" in head and "step/frac_of_step" in head + assert "step/wall_s" in head and "step/frac_of_step" in head assert breakdown_table(metrics)[1], "the table still has rows" client.close() @@ -733,7 +733,7 @@ def test_breakdown_table_rows_by_op_worst_first(): is the second column for the same reason.""" columns, rows = breakdown_table( { - "step/wall_ms": 100.0, + "step/wall_s": 0.1, "step/percent_of_dataplane/by_op/get": 10.0, "step/percent_of_dataplane/by_op/put": 90.0, "step/by_op/get/calls": 8, @@ -770,7 +770,7 @@ def test_breakdown_table_leaves_withheld_series_empty(): def test_breakdown_table_is_empty_when_nothing_ran(): - assert breakdown_table({"step/wall_ms": 0.0})[1] == [] + assert breakdown_table({"step/wall_s": 0.0})[1] == [] def test_breakdown_table_ignores_reserved_namespaces(): @@ -1195,15 +1195,15 @@ def test_codec_pack_unpack_time_is_reported_separately(): codec.record_codec_s("unpack", 0.004) # 4 ms metrics = client.get_step_metrics(1.0) - assert metrics["step/codec/pack_ms"] == pytest.approx(10.0, rel=1e-3) - assert metrics["step/codec/unpack_ms"] == pytest.approx(4.0, rel=1e-3) + assert metrics["step/codec/pack_s"] == pytest.approx(0.010, rel=1e-3) + assert metrics["step/codec/unpack_s"] == pytest.approx(0.004, rel=1e-3) # not folded into the transport totals - assert metrics["step/wall_ms"] == 0.0 + assert metrics["step/wall_s"] == 0.0 # and charted, not just tabulated - assert "step/codec/pack_ms" in headline_series(metrics) + assert "step/codec/pack_s" in headline_series(metrics) # drained exactly once: a second step reports zero, not the same 10 ms - assert client.get_step_metrics(1.0)["step/codec/pack_ms"] == 0.0 + assert client.get_step_metrics(1.0)["step/codec/pack_s"] == 0.0 client.close() @@ -1218,7 +1218,7 @@ def test_inspection_snapshot_does_not_steal_codec_time(): codec.record_codec_s("pack", 0.010) client.snapshot() # inspection: must not consume it - assert client.get_step_metrics(1.0)["step/codec/pack_ms"] == pytest.approx( - 10.0, rel=1e-3 + assert client.get_step_metrics(1.0)["step/codec/pack_s"] == pytest.approx( + 0.010, rel=1e-3 ) client.close() From 8af55636b53c7d03d49f19f81af7828e81c87049 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sun, 30 Aug 2026 04:19:17 -0700 Subject: [PATCH 50/70] fix(data-plane): make the wire guard see a row's shape `verify_tensor_hash` fingerprinted each row with a bare `torch.hash_tensor` fold. Its only implemented mode is an XOR reduction, which is blind to anything leaving the multiset of element words alone. Measured against the old implementation: - a row zero-padded 8 -> 12 hashed identically (`x ^ 0 == x`) - `(4, 8)` and `(32)` holding the same elements hashed identically - a dtype change at equal byte width hashed identically - a row folded against a duplicate of itself folded to zero All are plausible codec or wire bugs, and all were silence. The recorded `row_lens` did not cover the gap either: it held the *flattened* width, so `(N, 4, 8)` and `(N, 32)` both recorded 32, and it was compared only on the ragged fallback path, never on the primary one. Mix the row's dtype and shape into the fold and the first three become divergences. The values fold is untouched -- only the seed it combines with is new, so the on-device reduction that made the digest cheap is kept: fold = torch.hash_tensor(view.reshape(n_rows, -1), dim=1) seed = crc32(f"{dtype}|{row_shape}") digest = fold ^ seed The shape never travels and is never compared: one integer per row is stored, as before. A shape change moves the seed, which moves the digest, which surfaces as an ordinary mismatch. Most of this diff is deletion. The fold has no ragged kernel, so ragged leaves used a second, coarser digest over the whole values buffer, plus a `_WriteScheme` recorded at put and replayed at get so both sides agreed on which granularity had been used. None of it is needed: XOR is associative and elementwise, so `hash_tensor(row)` equals `hash_tensor(rect, dim=1)[i]`, and a ragged leaf can fold one row at a time and still compare against a vectorized read. `_WriteScheme`, `_batch_scope`, `_FieldDigest`, `_row_len`, `_dtype_salt`, `_as_int_view` and the `length_only` branch all go. A shard read is now checked rather than abstained on, and a ragged divergence names the sample rather than the batch. The seed's shape is the *row's*, and both layouts must agree on it: a dense `(N, L, D)` and the jagged form whose values are `(total, D)` both report `(L, D)`. That is what lets a field packed jagged and read back densified reconcile. Deriving the dense row's length from its offsets instead says `(1, L, D)` and makes every round trip a false mismatch. Measured on a 107 MB batch of 1536 rows x 4 fields: 8 ms, against 94 ms for the bare fold, 64 ms for a crc32 over each row's bytes and 146 ms for blake2b. The sequential hashes cost more because they read every byte on the host one row at a time; they also catch a within-row permutation, which this does not. That limit is taken deliberately for the cost, pinned by a test and recorded in README.md -- swapping `_leaf_digests` for the crc32 form is a one-function change if a reordering bug is ever suspected. Also: hoist the partition lookup above the hashing so a consumer-side get stops fingerprinting a payload it discards, and refresh the stale `torch.hash_tensor` / "~2.4 ms per 12 MB" figures in ObservabilityConfig and the two grpo_math_1B.yaml comments. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- examples/configs/grpo_math_1B.yaml | 2 +- nemo_rl/data_plane/README.md | 129 ++++--- nemo_rl/data_plane/interfaces.py | 13 +- nemo_rl/data_plane/observability.py | 359 +++++++----------- tests/unit/data_plane/test_observability.py | 183 +++++---- .../unit/reference_configs/grpo_math_1B.yaml | 2 +- 6 files changed, 325 insertions(+), 363 deletions(-) diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 5a305b36666..184e62d014d 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -585,7 +585,7 @@ data_plane: # GDR needs that headroom to pay off. observability: # per-op data-plane timing/volume enabled: true # per-op timing/volume; cost is below measurement noise - verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~2.4ms per 12MB batch) + verify_tensor_hash: false # debug: per-row hash of each row's values+dtype+shape, wire-in vs wire-out # Multi-Teacher On-Policy Distillation (MOPD): distills from one or more teacher # models into the policy via token-level teacher-minus-student logprob advantages, diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index e2f6d049f5f..6d2fabbf29e 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -651,75 +651,86 @@ MLflow). Roughly 5-8 series per distinct op tag. Set `observability.callback` if you additionally want a hook on every transfer; `log_event` is exported for that. -`verify_tensor_hash: true` additionally records a `torch.hash_tensor` -fingerprint on every put and re-checks it on every get, so a tensor that -changes between wire-in and wire-out is reported (`hash/mismatches`) -instead of being trained on silently. +`verify_tensor_hash: true` additionally records a fingerprint of every row +on every put and re-checks it on every get, so a tensor that changes between +wire-in and wire-out is reported (`hash/mismatches`) instead of being trained +on silently. -Two granularities, because torch has no ragged hash kernel: +One granularity: every row carries its own digest, formed from two parts. -| leaf | digest | scope | -|---|---|---| -| rectangular rows — dense, or jagged with uniform lengths | one per row, `hash_tensor(..., dim=1)` | per sample id; survives shard reads | -| genuinely ragged rows | one over the values buffer, XORed per row with that row's length | per batch; a shard read reports unverified | - -The split is on the *rows*, not the layout. A jagged leaf whose rows happen -to be uniform is already a rectangle — its values buffer reshapes to one as -a view — so it takes the per-row path for free. Only a leaf with rows of -differing lengths falls back, and giving *those* per-row digests would mean -padding each out to a rectangle first: on a realistically ragged batch that -rectangle is 3.5× the real payload and costs 13× more, to answer a question -the buffer digest already answers. - -Whichever scheme a put used is recorded per field and replayed on the read, -so the two sides always compute the same thing. A field written with -uniform rows that comes back ragged is a divergence in the row lengths -themselves, and is reported as a mismatch. - -**Detection is not attribution, and the difference is the whole point of -the split.** The same corruption, injected into an 8-row batch: - -| corruption | ragged leaf | rectangular leaf | +``` + the row's values ──► torch.hash_tensor (XOR fold, on device) ──► fold + ⊕ ──► digest + "|" ──► crc32 (host, one short string) ──► seed +``` + +| what it covers | so a divergence in | is caught | |---|---|---| -| 1 element changed in `u3` | caught, flags **all 8 rows** | caught, names **`u3`** | -| `u5` zeroed | caught, flags all 8 rows | caught, names `u5` | -| `u3`↔`u4` swapped | caught only if their lengths differ | caught, names `u3`,`u4` | -| nothing | clean | clean | - -A ragged digest covers the whole values buffer, so any change moves every -row's value: it says *this batch is wrong*, never *this sample is wrong*. -On rollout data straight out of generation that is the normal resolution -for the token-aligned fields — you learn a step's transfer diverged and -have to bisect for the row yourself. Anything uniform-width (a densified -read, `advantages` written at full width, a shard whose rows agree) names -the sample. +| the values fold | any element's value | yes | +| the seed's dtype | precision (bf16 vs fp32 at equal width) | yes | +| the seed's shape | length (a zero pad or a truncation) and trailing-dim layout | yes | +| — | a permutation *within* one row | **no** — see below | + +The shape never travels and is never compared: only one integer per row is +stored. A shape change makes the seed differ, which makes the digest differ, +which surfaces as an ordinary mismatch. + +The seed's shape is the *row's*, not the leaf's, and both layouts must agree +on it — a dense `(N, L, D)` and the jagged form whose values are `(total, D)` +both report a row of `(L, D)`. That is what lets a field packed jagged and +read back densified (`_from_wire` stacks uniform nested rows) reconcile +instead of reporting a mismatch on every round trip. Deriving the dense row's +length from its offsets instead would say `(1, L, D)` and break exactly that. + +Because the digest covers one row and nothing else, it reconciles against any +later grouping of the same rows: a shard read is *checked*, not abstained on, +and a delta write that touches one field leaves the others' fingerprints +alone. This is also why there is no longer a second, coarser granularity for +ragged leaves. `hash_tensor` has no ragged kernel, so a ragged leaf folds one +row at a time — but the fold is an XOR, which is associative and elementwise, +so `hash_tensor(row)` equals `hash_tensor(rect, dim=1)[i]`. The vectorized and +per-row paths produce identical values, and the `_WriteScheme` bookkeeping +that used to record which granularity a put had used, so a get could replay +it, is gone with them. + +**The accepted limit: a within-row permutation is not detected.** XOR cannot +see its own operands reordered, and no seed fixes it — the seed covers dtype +and shape, which a reordering leaves alone. This was taken deliberately, for +cost. Measured on a 107 MB batch of 1536 rows × 4 fields: + +| digest | cost | pad / reshape / dtype | permutation | +|---|---|---|---| +| **`hash_tensor` + shape seed** | **8 ms** | caught | **blind** | +| `crc32` over the row's bytes | 64 ms | caught | caught | +| `blake2b` over the row's bytes | 146 ms | caught | caught | +| bare `hash_tensor` (what this replaced) | 94 ms | blind | blind | + +The two sequential hashes cost ~7-18x because they read every byte on the +host, one row at a time; the fold reduces a whole rectangular leaf in one +on-device call. If a reordering bug is ever suspected — the jagged +pack/unpack offsets are where one would live — swapping `_leaf_digests` for +the `crc32` form is a one-function change. Verified by injecting corruption into the round trip. Caught: a single-element change in every dtype, a truncated row, a zeroed row, a -bf16→fp32 precision change, and a row served from the wrong sample — with -**zero false alarms** over a 500-row randomized soak, every shard grouping -from 1 to 256, reversed id order, field subsets and delta writes. Known -limits, measured rather than assumed: - -- A batch-scoped digest is an XOR reduction over one shared buffer, and XOR - cannot see a permutation of what it reduces. On a ragged field a - **mis-shard** (two rows swapped) is therefore caught only when the two - rows differ in length — 60/60 on ragged rollout data, where lengths rarely - collide. Rows of uniform width catch it unconditionally, per row. The same - blind spot hides a reordering *within* a row: 0/200 for a two-token swap - in `input_ids`, and 18/200 for moving two set bits in a bool mask. -- It reads every tensor byte again on both sides. Measured against a live - TransferQueue moving 23.6 MB per step: **10.2 ms, or +11% of data-plane - time** — which on a real GRPO step, where the data plane is a few percent - of the step, is under 0.1% end to end. The cost lands in - `step/self/overhead_ms` like the rest of the measurement, so it is visible - rather than quoted from a benchmark. +bf16→fp32 precision change, a zero pad, a trailing-dim reshape, and a row +served from the wrong sample — with **zero false alarms** over a 500-row +randomized soak, every shard grouping from 1 to 256, reversed id order, +field subsets and delta writes. Not caught, by the deliberate choice above: +a reordering of elements *within* one row. Note the row-swap and the +within-row cases differ — two rows exchanged between wire-in and wire-out +land against the wrong sample ids and are caught, because each row carries +its own digest. Known limits, measured rather than assumed: + +- It compares digests, so it detects divergence, not its cause. A mismatch + names the sample, the field, the row index and the row length; what + changed between the two reads is still yours to find. - **A mismatch count at or above `rows_checked` is reported as suspect.** Every row of every field wrong, identically, every step is not what a broken wire looks like; it is what a broken guard looks like. Both false alarms this check has produced had exactly that shape, and both were its - own bookkeeping. Per-sample lines carry the scheme, the row index and the - row length so the next one is adjudicable from a single log line. + own bookkeeping. Per-sample lines carry the row index and the row length + so the next one is adjudicable from a single log line. - Only rows this process wrote can be checked. A consumer-side client reports them under `hash/rows_unverified` rather than counting them clean, and `hash/fields_skipped` reports any leaf it could not compare diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index a2311a5c396..75ba7b86aae 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -186,12 +186,13 @@ class ObservabilityConfig(TypedDict): ``get_step_metrics``, so a per-op sink is opt-in. ``verify_tensor_hash`` is a correctness check, not a metric: each put - records a per-row ``torch.hash_tensor`` fingerprint and each get - re-checks it, so a value that changes between wire-in and wire-out is - reported (``hash/mismatches``) instead of silently training on it. It - reads every tensor byte a second time on both sides — budget roughly - 2.4 ms for a 12 MB jagged batch, on each side — so leave it off - outside of debugging. + records a per-row ``torch.hash_tensor`` fold of the row's values, mixed + with the row's dtype and shape, and each get re-checks it, so a value + that changes between wire-in and wire-out is reported + (``hash/mismatches``) instead of silently training on it. It reads every + tensor element a second time on both sides — roughly 8 ms for a 107 MB + batch — so leave it off outside of debugging. It does not detect a + permutation *within* a row; see ``data_plane/README.md``. """ enabled: bool diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 2507894f0e9..222505ba16c 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -27,12 +27,13 @@ Every method here runs on the hot path of a transfer, so nothing traverses a structure twice and nothing is allocated for a payload no callback reads. -``verify_tensor_hash=True`` adds an opt-in correctness check: -``torch.hash_tensor`` fingerprints recorded at put and re-checked at get, -so a tensor that changes between wire-in and wire-out is reported rather -than trained on. It reads every tensor byte again on both sides, so it is a -debugging tool, not a metric. See ``README.md`` for what it does and does -not catch. +``verify_tensor_hash=True`` adds an opt-in correctness check: a per-row +``torch.hash_tensor`` fold over each row's values, mixed with its dtype +and shape, recorded at +put and re-checked at get, so a tensor that changes between wire-in and +wire-out is reported rather than trained on. It reads every tensor byte +again on both sides, so it is a debugging tool, not a metric. See +``README.md`` for what it does and does not catch. """ from __future__ import annotations @@ -43,7 +44,8 @@ from dataclasses import asdict, dataclass, field from pathlib import Path from time import monotonic -from typing import Any, Callable, Collection, Literal, NamedTuple, TypedDict +from collections.abc import Sequence +from typing import Any, Callable, Literal, TypedDict EventStatus = Literal["ok", "error", "timeout"] @@ -111,50 +113,76 @@ class DataPlaneEvent(TypedDict): _QUANTILES = ((0.50, "p50_ms", 20), (0.90, "p90_ms", 40)) -class _WriteScheme(NamedTuple): - """How a field was reduced when it was written, so a read can replay it. - - A positional ``("scoped", n)`` / ``("rows", width)`` tuple carried two - different quantities in one slot, discriminated by a magic string, under - an annotation that said ``int``. Naming both makes each read site say - which one it means. - """ - - batch_scoped: bool - n_rows: int - row_width: int - - -class _FieldDigest(NamedTuple): - """One fingerprint per row, plus how far it can be trusted. - - ``batch_scoped`` means the values were derived from the whole batch's - buffer, so they only reconcile against a read of that same batch. A - shard of it computes a different buffer digest and must be reported - unverified rather than as a mismatch. - - ``row_lens`` is how long each row was, which is the one thing still - comparable when the read cannot reproduce the write's scheme. - """ - - per_row: list[int] - batch_scoped: bool - row_lens: tuple[int, ...] - - -# Same-width signed integer for each tensor element size, used to bitcast a -# leaf before hashing. +# Same-width integer for each element size, to bitcast a leaf before folding +# it. ``weight_transfer_sparse_codec.integer_dtype_for_element_size`` is the +# same map, but importing it drags in the vLLM generation stack +# (weight_transfer_sparse_codec -> models.generation.vllm -> telemetry -> +# nemo.lens), which the data plane must not depend on. Four entries is the +# cheaper duplicate. _INT_VIEW_BY_WIDTH = {1: torch.int8, 2: torch.int16, 4: torch.int32, 8: torch.int64} -def _as_int_view(t: torch.Tensor) -> torch.Tensor: - """Bitcast to a same-width integer type, or pass through. +def _leaf_digests( + leaf: torch.Tensor, + bounds: Sequence[int], + row_shape: Callable[[int], tuple[int, ...]], + dtype: torch.dtype, +) -> list[int]: + """One digest per row: ``hash_tensor``'s fold, with the row shape mixed in. + + ``torch.hash_tensor`` only implements an XOR fold (``mode=0``), which is + blind to a zero pad, a trailing-dim reshape and a dtype change, because + none of those alter the multiset of element words. Mixing the row's shape + and dtype into the fold closes all three. It remains blind to a + permutation *within* a row, which is the price of the fold being + vectorized; ``README.md`` records that. + + The fold being an XOR is also what lets one algorithm serve both layouts. + XOR is associative and elementwise, so ``hash_tensor(row)`` equals + ``hash_tensor(rect, dim=1)[i]`` -- a rectangle reduces in one on-device + call, a ragged leaf falls back to one call per row, and the two agree + value-for-value. A field packed jagged and read back densified therefore + reconciles without either side recording how the other reduced it. + + ``bounds`` are leading-dim offsets, ``n_rows + 1`` of them: row *i* spans + ``leaf[bounds[i] : bounds[i + 1]]``. ``row_shape`` maps a row's length to + the shape recorded for it, which is what keeps the jagged and dense views + of one field agreeing: both must call a row ``(L, D)``. - ``hash_tensor`` has no float8 kernel and raises there; viewing the bytes - as integers sidesteps every dtype-specific kernel for free. + Args: + leaf: The whole leaf -- a dense tensor, or a jagged one's values. + bounds: Leading-dim offsets delimiting each row. + row_shape: Row length -> the shape to record for that row. + dtype: Mixed in so a precision change diverges at equal byte width. + + Returns: + One digest per row. """ - int_dtype = _INT_VIEW_BY_WIDTH.get(t.element_size()) - return t.view(int_dtype) if int_dtype is not None else t + # Bitcast so every dtype reduces: hash_tensor ships no float8 kernel. + # detach() because a grad-carrying leaf must not be viewed under autograd. + view = leaf.detach() + view = view.view(_INT_VIEW_BY_WIDTH[view.element_size()]) + lengths = [hi - lo for lo, hi in zip(bounds, bounds[1:])] + if lengths and lengths.count(lengths[0]) == len(lengths): + folds = torch.hash_tensor(view.reshape(len(lengths), -1), dim=1).tolist() + else: + # Ragged rows have no rectangle to reduce over, so each row folds on + # its own. Same values as the vectorized path, just one call apiece. + folds = [ + torch.hash_tensor(view[lo:hi].reshape(1, -1), dim=1)[0].item() + for lo, hi in zip(bounds, bounds[1:]) + ] + # Salted on the host: torch has no UInt64 bitwise_xor CUDA kernel, so + # XOR-ing the digest tensor raises for any backend whose get returns + # device tensors. The digests come to the host for comparison regardless. + seeds: dict[int, int] = {} + digests = [] + for fold, length in zip(folds, lengths): + seed = seeds.get(length) + if seed is None: + seed = seeds[length] = zlib.crc32(f"{dtype}|{row_shape(length)}".encode()) + digests.append(fold ^ seed) + return digests def _as_list(sample_ids: Any) -> Any: @@ -213,16 +241,6 @@ def _tensor_bytes(v: torch.Tensor) -> int: return buf.nbytes -def _dtype_salt(dtype: torch.dtype) -> int: - """Salt distinguishing dtypes whose values reduce to the same words. - - ``crc32``, not the builtin ``hash()``: ``hash()`` of a ``str`` is salted - per process, so the fingerprint would not survive being compared across - ranks. - """ - return zlib.crc32(str(dtype).encode()) - - def percentile_from_hist(hist: list[int], q: float) -> float: """Interpolated ``q``-quantile (0-1) from bucket counts. @@ -419,17 +437,6 @@ def _delta_s(field: str) -> float: } -def _row_len(digest: _FieldDigest, row: int) -> int: - """How long ``row`` was, whichever scheme the digest used. - - A per-row digest stores the one width its uniform rows shared; a - batch-scoped one stores every row's length. - """ - if not digest.row_lens: - return -1 - return digest.row_lens[row] if digest.batch_scoped else digest.row_lens[0] - - def _op_step_stats( by_op: dict[str, Any], prev_ops: dict[str, Any] ) -> dict[str, dict[str, float]]: @@ -1036,10 +1043,11 @@ def __init__( on_event: Per-op callback. ``None`` (the default) skips building the event dict entirely — with metrics enabled but no sink, nothing is paid for a payload nobody reads. - verify_tensor_hash: Record a per-row ``torch.hash_tensor`` - fingerprint on put and re-check it on get. Debug aid, not a - metric: it reads every tensor byte again (~2.4 ms for a - 12 MB jagged batch), so it is off unless the config asks. + verify_tensor_hash: Record a per-row fingerprint + on put and re-check it on get. Debug aid, not a metric: it + reads every tensor element again on both sides (~8 ms + for a 107 MB batch of 1536 rows), so it is off unless the + config asks. """ self._inner = inner self._on_event = on_event @@ -1054,12 +1062,6 @@ def __init__( # lifetime as ``_bytes_by_partition``: cleared by ``clear_samples``, # so it is bounded by the live key population. self._hash_by_partition: dict[str, dict[str, dict[str, int]]] = {} - # partition -> (batch-scoped field names, row count at put). Those - # digests cover a whole buffer, so they only reconcile against a read - # of that same batch; the row count is what detects a shard read. - # partition -> field -> rows the field's digest was reduced over, - # for batch-scoped fields only. An absent field was reduced per row. - self._batch_scope: dict[str, dict[str, _WriteScheme]] = {} self._hash_mismatches_logged = 0 # Set by ``_emit`` to the inner client's wall time for the op just # run, so the wrapping methods can subtract it and bill the rest to @@ -1198,8 +1200,6 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: """ if self._verify_tensor_hash: _pop_partition_keys(self._hash_by_partition, partition_id, keys) - if keys is None: - self._batch_scope.pop(partition_id, None) live = self._keys_by_partition.get(partition_id) if live is None: return @@ -1240,95 +1240,66 @@ def _row_fingerprints( self, td: TensorDict | None, sample_ids: list[str], - batch_scoped_fields: Collection[str] = (), - ) -> dict[str, _FieldDigest]: - """``torch.hash_tensor`` fingerprints for each tensor leaf. - - A rectangular leaf reduces per row (``dim=1``), which names the - sample that diverged. A jagged leaf whose rows happen to be uniform - is rectangular already — its values buffer reshapes to the rectangle - as a view — so it takes the same path for free. Only a genuinely - ragged leaf falls back to one digest over its whole values buffer, - XORed per row with that row's length, and marked ``batch_scoped``; - padding it out to a rectangle costs far more than the answer is - worth. That fallback inherits the blind spot of an XOR reduction: it - sees any change to the multiset of values, but not a permutation of - them. ``README.md`` has the detection/attribution table. + ) -> dict[str, list[int]]: + """A per-row digest of each tensor leaf, covering bytes, dtype and shape. + + Every leaf is fingerprinted one row at a time, so every divergence + names the sample that diverged -- a genuinely ragged leaf included. + See ``README.md`` for why the digest this replaced could not. + + Both layouts reduce to the same shape of work: hand + :func:`_leaf_digests` the leaf, the offsets delimiting its rows, and + how to describe a row's shape. It picks between one vectorized fold + and a fold per row. Args: td: Leaves to fingerprint; ``None`` yields an empty result. sample_ids: Row *i* is attributed to ``sample_ids[i]``, the ordering :meth:`DataPlaneClient.get_samples` promises. - batch_scoped_fields: Fields the *put* side reduced batch-scoped. - The scheme must follow the field, not the layout in hand: a - field packed jagged comes back dense whenever its rows are - uniform (``_from_wire`` densifies those), and choosing per - layout makes the two sides compute different things. The - values buffer of a uniform jagged field and the flattened - dense tensor it densifies into hold the same elements in the - same order, so replaying the recorded scheme agrees by - construction. Returns: - Field name -> :class:`_FieldDigest`. Leaves that cannot be - attributed per row (a leading dim that isn't ``len(sample_ids)``, - or a non-``jagged`` nested layout) are counted in - ``fields_skipped`` rather than silently dropped. + Field name -> one digest per row. A leaf that cannot be attributed + per row is counted in ``fields_skipped`` rather than silently + dropped: a non-``jagged`` nested layout, a leading dim that is not + ``len(sample_ids)``, or a leaf with no leading dim at all. """ if td is None: return {} n_rows = len(sample_ids) stats = self._stats.hash_verify - out: dict[str, _FieldDigest] = {} + out: dict[str, list[int]] = {} for key, v in td.items(include_nested=True, leaves_only=True): if not isinstance(v, torch.Tensor) or v.ndim < 1: stats.fields_skipped += 1 continue - salt = _dtype_salt(v.dtype) - name = key if isinstance(key, str) else ".".join(key) if v.is_nested: - if v.layout != torch.jagged: + if v.layout != torch.jagged or v.offsets().numel() - 1 != n_rows: stats.fields_skipped += 1 continue - offsets = v.offsets() - if offsets.numel() - 1 != n_rows: - stats.fields_skipped += 1 - continue - lengths = (offsets[1:] - offsets[:-1]).tolist() - # Uniform rows: the values buffer already *is* the rectangle, - # so reshaping it is a view and the per-row reduction is free. - uniform = lengths.count(lengths[0]) == n_rows - rectangle = v.values() if uniform else None + bounds = v.offsets().tolist() + leaf = v.values() + # A jagged row is its own length followed by the values + # buffer's trailing dims. + tail = tuple(leaf.shape[1:]) + row_shape = lambda length, tail=tail: (length, *tail) elif v.shape[0] != n_rows: stats.fields_skipped += 1 continue else: - lengths = [v.shape[1] if v.ndim >= 2 else 1] * n_rows - rectangle = v - if rectangle is not None and name not in batch_scoped_fields: - flat = _as_int_view(rectangle.reshape(n_rows, -1)) - out[name] = _FieldDigest( - # Salt on the host, as the batch-scoped path below already - # does. Torch has no ``bitwise_xor`` CUDA kernel for - # UInt64, so XOR-ing the digest tensor in place raises - # ``NotImplementedError`` for any backend whose get returns - # device tensors. The digests come back to the host for - # comparison either way, so this costs nothing. - [d ^ salt for d in torch.hash_tensor(flat, dim=1).tolist()], - batch_scoped=False, - # One width, not n_rows copies of it: the rows are - # uniform by construction on this path. - row_lens=(flat.shape[1],), - ) - continue - buffer = v.values() if v.is_nested else v - flat_buffer = _as_int_view(buffer.reshape(1, -1)) - buffer_digest = torch.hash_tensor(flat_buffer, dim=1).tolist()[0] ^ salt - out[name] = _FieldDigest( - [buffer_digest ^ length for length in lengths], - batch_scoped=True, - row_lens=tuple(lengths), - ) + # Dense rows are equal-length by construction, so the same + # bounds describe them: element i starts at i. + bounds = range(n_rows + 1) + leaf = v + # ...and every dense row has the leaf's trailing shape, which + # is the *same* tuple the jagged form reports for it. That + # equality is what lets a jagged put reconcile against a + # densified get; recording the bounds-derived length here + # instead would make ``(N, L, D)`` say ``(1, L, D)`` and every + # round trip a mismatch. + shape = tuple(v.shape[1:]) + row_shape = lambda _length, shape=shape: shape + name = key if isinstance(key, str) else ".".join(key) + out[name] = _leaf_digests(leaf, bounds, row_shape, v.dtype) return out def _hash_guard_failed(self, op: str, exc: Exception) -> None: @@ -1372,22 +1343,8 @@ def _record_hashes_impl( partition_hashes = self._hash_by_partition.setdefault(partition_id, {}) for row, sample_id in enumerate(sample_ids): per_field = partition_hashes.setdefault(sample_id, {}) - for name, digest in digests.items(): - per_field[name] = digest.per_row[row] - # Per field, not per partition: a delta put (write_columns) names only - # the fields it writes, and must not restate the scheme of the ones it - # left alone. Recording the batch it was reduced over lets the read - # side tell a shard of a batch-scoped field from a relayout. - scheme = self._batch_scope.setdefault(partition_id, {}) - for name, digest in digests.items(): - # Width matters only for the per-row scheme -- uniform rows mean - # one number describes them all, and lets a later ragged read - # still check whether any row changed length. - scheme[name] = _WriteScheme( - batch_scoped=digest.batch_scoped, - n_rows=len(sample_ids), - row_width=digest.row_lens[0] if digest.row_lens else 0, - ) + for name, per_row in digests.items(): + per_field[name] = per_row[row] self._stats.hash_verify.rows_recorded += len(sample_ids) def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: @@ -1400,50 +1357,22 @@ def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> N def _check_hashes_impl( self, partition_id: str, sample_ids: list[str], out: Any ) -> None: - """Compare wire-out fingerprints against what was written.""" + """Compare wire-out fingerprints against what was written. + + Every field is comparable now. The old two-tier classification -- + which fields could be compared, which were a shard of a batch-scoped + put, which had been written uniform and read back ragged -- existed + entirely to work around a digest that only meant something against + the exact batch it was folded over. A per-row hash reconciles against + any grouping, so a shard read is checked rather than abstained on. + """ if not isinstance(out, TensorDict): return - scheme = self._batch_scope.get(partition_id, {}) - # Only the batch-scoped names: the scheme also records per-row fields - # now, and handing the whole mapping over forced every field onto the - # scoped path. - scoped_names = {n for n, how in scheme.items() if how.batch_scoped} - digests = self._row_fingerprints(out, sample_ids, scoped_names) + digests = self._row_fingerprints(out, sample_ids) if not digests: return partition_hashes = self._hash_by_partition.get(partition_id, {}) stats = self._stats.hash_verify - # A batch-scoped digest covers the whole buffer it was reduced from, - # so it only means anything against a read of that same batch. Drop - # those fields on a shard read rather than reporting every row of it - # as a mismatch — but *count* the drop. Dropping silently is the - # exact shape of the bug that made this check pass while covering - # nothing: a field stops being compared and the report still reads - # clean. - comparable: dict[str, _FieldDigest] = {} - length_only: dict[str, tuple[_FieldDigest, int]] = {} - for name, digest in digests.items(): - recorded = scheme.get(name) - if not digest.batch_scoped: - comparable[name] = digest - elif recorded is None: - continue # this process never wrote the field - elif recorded.batch_scoped: - if recorded.n_rows == len(sample_ids): - comparable[name] = digest - else: - stats.fields_skipped += 1 # a shard of a batch-scoped put - else: - # Written with uniform rows, read back ragged. Treating that - # as a divergence reported 3584 mismatches per step on a - # healthy run: it is the normal shape of a real pipeline, - # where a shard is written uniform and read back inside a - # batch whose other rows differ in length. The row lengths - # are still comparable, and a row that changed length *is* a - # divergence, so check that much and count the rest as the - # abstention it is. - length_only[name] = (digest, recorded.row_width) - stats.fields_skipped += 1 for row, sample_id in enumerate(sample_ids): per_field = partition_hashes.get(sample_id) if not per_field: @@ -1452,46 +1381,26 @@ def _check_hashes_impl( stats.rows_unverified += 1 continue stats.rows_checked += 1 - for name, (digest, width) in length_only.items(): - read_len = _row_len(digest, row) - if per_field.get(name) is None or read_len == width: - continue - stats.mismatches += 1 - if self._hash_mismatches_logged < _MAX_HASH_MISMATCH_LOGS: - self._hash_mismatches_logged += 1 - logger.error( - "data-plane hash mismatch: partition=%s sample=%s " - "field=%s row was %d long on the wire in, %d out", - partition_id, - sample_id, - name, - width, - read_len, - ) - for name, digest in comparable.items(): + for name, per_row in digests.items(): expected = per_field.get(name) - if expected is None or expected == digest.per_row[row]: + if expected is None or expected == per_row[row]: continue stats.mismatches += 1 if self._hash_mismatches_logged < _MAX_HASH_MISMATCH_LOGS: self._hash_mismatches_logged += 1 - # Scheme and shape on the line: the last two false alarms - # were both bookkeeping (a scheme replayed wrong, a - # grouping that could not be compared), and neither was + # Row index on the line: both false alarms this check ever + # produced were its own bookkeeping, and neither was # diagnosable from the digests alone. logger.error( "data-plane hash mismatch: partition=%s sample=%s " - "field=%s wire_in=%d wire_out=%d " - "(%s scheme, row %d of %d, %d long)", + "field=%s wire_in=%d wire_out=%d (row %d of %d)", partition_id, sample_id, name, expected, - digest.per_row[row], - "batch-scoped" if digest.batch_scoped else "per-row", + per_row[row], row, len(sample_ids), - _row_len(digest, row), ) def _run( diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 98563501552..f1151f92a1f 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -127,10 +127,9 @@ def _hash_fields(n=4): def _jagged_ids(lengths, seed=0, with_dense=False): """Rows of pseudorandom token ids, optionally beside a uniform ``lp`` field. - Deliberately not ``arange``: ``hash_tensor`` is an XOR reduction, and - aligned runs of consecutive integers collide under it — ``XOR(6..11)`` - and ``XOR(12..17)`` are both 1 — which would make two visibly different - rows fingerprint the same. + Pseudorandom rather than ``arange`` so distinct rows stay visibly + distinct in a mismatch log line; the digest itself separates consecutive + runs fine. """ g = torch.Generator().manual_seed(seed) fields = { @@ -844,7 +843,7 @@ def test_guard_failure_is_absorbed_counted_and_charted(caplog, monkeypatch): ids = _ids(4) def boom(*_args, **_kwargs): - raise NotImplementedError("no hash_tensor kernel for this dtype") + raise NotImplementedError("no digest kernel for this dtype") monkeypatch.setattr(client, "_row_fingerprints", boom) with caplog.at_level(logging.WARNING): @@ -939,8 +938,8 @@ def test_hash_fingerprint_covers_jagged_fields(): digest = client._row_fingerprints(_jagged(rows), ["a", "b", "c"])["x"] assert client.snapshot()["hash_verify"]["fields_skipped"] == 0 - assert digest.batch_scoped, "jagged digests only reconcile per batch" - assert len(digest.per_row) == 3 + assert len(digest) == 3 + assert len(set(digest)) == 3, "each ragged row gets its own digest" changed = list(rows) changed[1] = changed[1] + 1 assert client._row_fingerprints(_jagged(changed), ["a", "b", "c"])["x"] != digest @@ -949,33 +948,71 @@ def test_hash_fingerprint_covers_jagged_fields(): def test_hash_fingerprint_matches_across_jagged_and_dense(): """``_from_wire`` densifies a jagged field whose rows are uniform, so a - jagged put has to reconcile against a dense get. Picking the scheme from the - layout in hand rather than from what was recorded made every row of a - uniform batch report a mismatch — 940 false alarms over the soak.""" + jagged put has to reconcile against a dense get. The digest binds in the + row *slice*'s shape for exactly this reason: ``v[i]`` of the dense form and + ``values[o : o + L]`` of the jagged one are both ``(L,)``, so the two sides + agree without either needing to know which layout the other held.""" client = _client(verify_tensor_hash=True) ids = ["a", "b"] dense = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int64) - # uniform rows: both sides reduce per row, and a densified read agrees - put_side = client._row_fingerprints(_jagged(list(dense.unbind())), ids)["x"] - assert not put_side.batch_scoped - assert put_side == client._row_fingerprints( - TensorDict({"x": dense}, batch_size=[2]), ids + assert ( + client._row_fingerprints(_jagged(list(dense.unbind())), ids)["x"] + == client._row_fingerprints(TensorDict({"x": dense}, batch_size=[2]), ids)["x"] + ) + client.close() + + +def test_mixing_the_shape_in_closes_three_of_the_folds_blind_spots(): + """A bare ``hash_tensor`` fold is blind to a zero pad (``x ^ 0 == x``), to + a trailing-dim reshape, and to a dtype change, because none of the three + alter the multiset of element words it reduces. Mixing the row's dtype and + shape into the fold makes each one a divergence. The values fold is + untouched — only the seed it is combined with changes. + """ + client = _client(verify_tensor_hash=True) + row = torch.arange(1, 9, dtype=torch.int64) + fp = lambda t: client._row_fingerprints( # noqa: E731 + TensorDict({"x": t}, batch_size=[1]), ["a"] )["x"] - # ragged rows: batch-scoped, and the read replays that scheme rather than - # choosing one from the dense tensor in hand - ragged = _jagged([torch.tensor([1, 2, 3]), torch.tensor([4, 5])]) - assert client._row_fingerprints(ragged, ids)["x"].batch_scoped - assert client._row_fingerprints( - TensorDict({"x": dense}, batch_size=[2]), ids, batch_scoped_fields={"x"} - )["x"].batch_scoped + base = fp(row.reshape(1, 8)) + assert fp(torch.cat([row, torch.zeros(4, dtype=torch.int64)]).reshape(1, 12)) != ( + base + ), "a zero-padded row: same elements, longer row" + assert fp(row.reshape(1, 2, 4)) != base, ( + "the same elements under a different trailing shape" + ) + assert fp(row.reshape(1, 8).to(torch.float64)) != base, ( + "the same words under a different dtype" + ) + client.close() + + +def test_a_within_row_permutation_is_the_accepted_blind_spot(): + """``hash_tensor``'s fold is an XOR, so it cannot see its own operands + reordered, and no seed fixes that — the seed covers dtype and shape, which + a permutation leaves alone. This is the price of reducing on device in one + call per leaf, taken deliberately: a sequential hash over the row's bytes + catches it but costs ~7x. Pinned so the limit stays a decision rather than + a surprise; ``README.md`` records it in the operator-facing table. + """ + client = _client(verify_tensor_hash=True) + row = torch.arange(1, 9, dtype=torch.int64) + fp = lambda t: client._row_fingerprints( # noqa: E731 + TensorDict({"x": t}, batch_size=[1]), ["a"] + )["x"] + + assert fp(row.flip(0).reshape(1, 8)) == fp(row.reshape(1, 8)), ( + "a reordered row is NOT detected -- see the docstring before changing this" + ) client.close() def test_hash_fingerprint_separates_dtype(): - """The values reduce identically once bitcast, so only the dtype salt makes - a precision change visible.""" + """The two hold different bytes, but a guard that hashed only bytes could + still collide across dtypes; binding the dtype into the preimage makes a + precision change visible on its own.""" client = _client(verify_tensor_hash=True) ids = ["a", "b"] values = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) @@ -984,14 +1021,15 @@ def test_hash_fingerprint_separates_dtype(): as_bf16 = client._row_fingerprints( TensorDict({"x": values.to(torch.bfloat16)}, batch_size=[2]), ids ) - assert as_fp32["x"].per_row != as_bf16["x"].per_row + assert as_fp32["x"] != as_bf16["x"] client.close() def test_hash_fingerprint_handles_float8(): - """``hash_tensor`` has no float8 kernel. Without the integer bitcast the - ``NotImplementedError`` propagates out of ``put_samples`` and takes the - transfer down with it.""" + """float8 has no dedicated hash kernel, and ``.numpy()`` rejects it + outright. Viewing the row as raw bytes spans every dtype, so a float8 + payload is fingerprinted rather than raising out of ``put_samples`` and + taking the transfer down with it.""" client = _client(verify_tensor_hash=True) fp8 = TensorDict( {"x": torch.tensor([[1.0, 2.0], [3.0, 4.0]]).to(torch.float8_e4m3fn)}, @@ -1019,18 +1057,18 @@ def test_hash_fingerprint_handles_device_tensors(): on_device = client._row_fingerprints( TensorDict({"x": values.cuda()}, batch_size=[2]), ids ) - assert on_device["x"].per_row == on_host["x"].per_row + assert on_device["x"] == on_host["x"] client.close() # ── hash verification: what counts as a mismatch vs an abstention ────── -def test_hash_shard_of_a_ragged_field_is_skipped_not_a_mismatch(): - """The abstention that *is* legitimate: a batch-scoped digest covers the - whole buffer it was reduced over, so a shard of it is genuinely - incomparable. That must land in ``fields_skipped`` — visible, but not - crying wolf on every sharded fetch.""" +def test_hash_shard_of_a_ragged_field_is_checked_not_abstained_on(): + """A shard of a ragged put is fully comparable. The digest covers one row + and nothing else, so it reconciles against any later grouping of those + rows — where a fold over the whole batch's values buffer meant nothing + against a slice of it and had to be counted as an abstention.""" client = _client(verify_tensor_hash=True) ids = _ids(4) rows = [torch.arange(3 + i, dtype=torch.int64) for i in range(4)] @@ -1039,23 +1077,26 @@ def test_hash_shard_of_a_ragged_field_is_skipped_not_a_mismatch(): hv = client.snapshot()["hash_verify"] assert hv["mismatches"] == 0 - assert hv["fields_skipped"] == 1 - assert client.get_step_metrics(1.0)["step/hash/fields_skipped"] == 1 + assert hv["fields_skipped"] == 0, "a shard is comparable now" + assert hv["rows_checked"] == 2, "and both its rows were actually checked" client.close() -def test_uniform_jagged_rows_are_fingerprinted_per_row(): - """A jagged field whose rows happen to be uniform is a rectangle already — - its values buffer reshapes to one as a view — so it earns per-row - attribution for free. The batch-scoped fallback is an XOR over one shared - buffer, which cannot see a permutation: two equal-length rows swapped by a - mis-shard would round-trip clean.""" +@pytest.mark.parametrize( + "lengths", [[6, 6, 6, 6], [6, 6, 4, 9]], ids=["uniform", "ragged"] +) +def test_mis_sharded_rows_are_named_individually(lengths): + """Two rows swapped between wire-in and wire-out, named per sample. + + A swap *between* rows is visible even though a permutation *within* one is + not: each row carries its own digest, so the two land against the wrong + sample ids. The old batch-scoped fallback folded every ragged row into one + buffer digest and lost that; per-row digests keep it for both layouts. + """ inner = _JaggedEcho() client = _client(inner, verify_tensor_hash=True) ids = _ids(4) - client.put_samples( - sample_ids=ids, partition_id="p", fields=_jagged_ids([6, 6, 6, 6]) - ) + client.put_samples(sample_ids=ids, partition_id="p", fields=_jagged_ids(lengths)) a, b = inner.rows[("p", "u1")]["ids"], inner.rows[("p", "u2")]["ids"] inner.rows[("p", "u1")]["ids"], inner.rows[("p", "u2")]["ids"] = b, a client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) @@ -1066,13 +1107,11 @@ def test_uniform_jagged_rows_are_fingerprinted_per_row(): client.close() -def test_uniform_put_read_back_ragged_checks_row_lengths(): - """A row that changed length between wire-in and wire-out is a divergence. - The content is no longer comparable, so the field counts as an abstention, - but the lengths still are — and reporting the whole field as a skip would - leave mismatches reading zero, exactly the shape of a guard that covers - nothing. Failing the whole field instead reported 3584 mismatches per step - on a healthy 5-process run.""" +def test_a_truncated_row_is_a_mismatch(): + """A row that changed length between wire-in and wire-out is a divergence, + and the shape is inside the digest, so it is caught as an ordinary content + mismatch. No side-channel length comparison and no abstention: the field + stays fully comparable.""" inner = _JaggedEcho() client = _client(inner, verify_tensor_hash=True) ids = _ids(4) @@ -1083,17 +1122,17 @@ def test_uniform_put_read_back_ragged_checks_row_lengths(): client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) hv = client.snapshot()["hash_verify"] - assert hv["mismatches"] > 0, "a truncated row must not read as clean" - assert hv["fields_skipped"] == 1, "and the content it could not compare" + assert hv["mismatches"] == 1, "the truncated row, named" + assert hv["fields_skipped"] == 0, "and the field stayed comparable" client.close() def test_uniform_write_read_back_inside_a_ragged_batch_is_clean(): - """The false positive this cost: a shard written with uniform rows, read - back inside a batch whose *other* rows are ragged. Nothing diverged — every - recorded row still has the length it was written with — and the guard - reported 3584 mismatches per step on a healthy run until it compared lengths - instead of failing the field outright.""" + """The false positive this once cost: a shard written with uniform rows, + read back inside a batch whose *other* rows are ragged. Nothing diverged, + and because a digest describes only its own row, the mixed batch needs no + special handling — the rows this process wrote are compared, the rows it + did not are counted unverified.""" inner = _JaggedEcho() client = _client(inner, verify_tensor_hash=True) ids = _ids(4) @@ -1107,19 +1146,21 @@ def test_uniform_write_read_back_inside_a_ragged_batch_is_clean(): client.get_samples(sample_ids=ids + other, partition_id="p", select_fields=["ids"]) hv = client.snapshot()["hash_verify"] - assert hv["mismatches"] == 0, "no row changed length; nothing diverged" - assert hv["fields_skipped"] == 1, "content uncomparable, and counted" + assert hv["mismatches"] == 0, "no row changed; nothing diverged" + assert hv["fields_skipped"] == 0, "and every field stayed comparable" + assert hv["rows_checked"] == 4, "the four this process wrote" + assert hv["rows_unverified"] == 2, "the two it did not" client.close() -def test_delta_put_does_not_restate_the_scheme_of_untouched_fields(): - """``write_columns`` puts one field into a partition written ragged earlier. - Holding the jagged/per-row choice per *partition* let that delta hand the - read side the wrong scheme for a field it never touched, and every row of - that field came back a false alarm. +def test_delta_put_leaves_untouched_fields_fingerprinted(): + """``write_columns`` puts one field into a partition written ragged + earlier. A delta must not disturb the fingerprints of a field it never + named — holding any of this state per *partition* rather than per field + turned every row of the untouched field into a false alarm. - The second field is the whole point: with only one, a per-partition and a - per-field scheme are indistinguishable, because the only put there is + The second field is the whole point: with only one, per-partition and + per-field bookkeeping are indistinguishable, because the only put there is restates its own field either way. """ inner = _JaggedEcho() @@ -1133,8 +1174,8 @@ def test_delta_put_does_not_restate_the_scheme_of_untouched_fields(): client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) assert client.snapshot()["hash_verify"]["mismatches"] == 0, "baseline" - # the delta names only ``lp``; ``ids`` must keep the ragged scheme it was - # written with, or its next read replays the uniform one and cries wolf + # the delta names only ``lp``; ``ids`` must keep the fingerprints it was + # written with, or its next read cries wolf on every row client.put_samples( sample_ids=ids, partition_id="p", @@ -1143,7 +1184,7 @@ def test_delta_put_does_not_restate_the_scheme_of_untouched_fields(): client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) hv = client.snapshot()["hash_verify"] - assert hv["mismatches"] == 0, "a delta put must not restate ids's scheme" + assert hv["mismatches"] == 0, "a delta put must not disturb ids" assert hv["fields_skipped"] == 0, "and the read stays comparable" client.close() diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index d06a39060d5..e32df674413 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -557,7 +557,7 @@ data_plane: # GDR needs that headroom to pay off. observability: # per-op data-plane timing/volume enabled: true # per-op timing/volume; cost is below measurement noise - verify_tensor_hash: false # debug: per-row torch.hash_tensor wire-in vs wire-out (~2.4ms per 12MB batch) + verify_tensor_hash: false # debug: per-row hash of each row's values+dtype+shape, wire-in vs wire-out # Multi-Teacher On-Policy Distillation (MOPD); null = disabled (default). Mirrors # the field on the GRPO MasterConfig added for MOPD support. From dd83c8c260c9221ac9b8a776a882df6f5f2d1a6f Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Mon, 7 Sep 2026 15:42:29 -0700 Subject: [PATCH 51/70] fix(data-plane): log single-controller metrics before the step is committed (#4044) Co-authored-by: Zhiyu Li Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/single_controller.py | 6 ++- tests/unit/data_plane/test_observability.py | 45 +++++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 82e9744b0bf..3796ab27627 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -2903,6 +2903,11 @@ async def _train_pump(self) -> None: self._logger.log_metrics( step_metrics, step=self._train_steps, prefix="train" ) + # Must precede the step_finished=True log below. That log commits + # the wandb step, and wandb silently discards anything logged + # against a step it has already committed -- no exception, no + # failed return, just an empty chart. grpo_sync had the same bug. + self._log_data_plane_metrics(total_time) # step_finished=True here since this is the final log of our current step. self._logger.log_metrics( timing_metrics, @@ -2910,7 +2915,6 @@ async def _train_pump(self) -> None: prefix="timing/train", step_finished=True, ) - self._log_data_plane_metrics(total_time) self._timer.reset() # min sample version refers to the version each consumed sample was diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index f1151f92a1f..7358cc2919e 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -1263,3 +1263,48 @@ def test_inspection_snapshot_does_not_steal_codec_time(): 0.010, rel=1e-3 ) client.close() + + +def test_no_algorithm_logs_data_plane_after_committing_the_step(): + """``log_metrics(..., step_finished=True)`` commits the wandb step, and + wandb discards anything logged against a step it has already committed -- + without raising, and without a falsy return to check. + + This has bitten twice. ``grpo_sync`` was caught only because a real run + showed 85 logged keys and zero ``data_plane/*``. ``single_controller`` + carried the same code -- its docstring says it mirrors ``grpo_sync`` -- + and so carried the same bug, unnoticed, because no run exercised it. + + Asserting the invariant for every algorithm rather than for one call site + is the point: a third wiring would otherwise repeat it. Source order is + the only observable, because the drop happens inside wandb where a fake + logger sees a perfectly ordinary call. + """ + import pathlib + + import nemo_rl + + algorithms = pathlib.Path(nemo_rl.__file__).parent / "algorithms" + checked = [] + for path in sorted(algorithms.glob("*.py")): + # Comments mention the flag too, so they cannot be part of the search. + source = "\n".join( + line + for line in path.read_text().splitlines() + if not line.lstrip().startswith("#") + ) + if "_log_data_plane_metrics(" not in source: + continue + if "step_finished=True" not in source: + continue + checked.append(path.name) + # rindex, not index: the first occurrence is the *definition*, which + # naturally precedes everything. The call site is what has to come + # before the commit, and it is the last occurrence. + assert source.rindex("_log_data_plane_metrics(") < source.index( + "step_finished=True" + ), ( + f"{path.name}: data-plane metrics are logged after the " + "step_finished=True commit, so wandb will discard them" + ) + assert len(checked) >= 2, f"expected sync and single-controller, got {checked}" From dd3095b437bdc5ff35917dad3151dcf9f3cfd4be Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Mon, 7 Sep 2026 16:29:49 -0700 Subject: [PATCH 52/70] fix(data-plane): drop the PROMOTE_1D_FIELDS branch from _from_wire Rebasing onto main dropped the Mooncake 1-D promotion (the constant, the `_promote_1d_leaves` helper and its import), which is intended -- it is not needed. My conflict resolution kept the *use* of PROMOTE_1D_FIELDS inside `_from_wire` while the name it referenced was gone, which would have raised NameError on the first unpack. Caught by a symbol-level audit of the rebase rather than by the checks I ran first: parsing, conflict-marker scans and comparing the core files all passed over it, because the name is only resolved at runtime. Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/adapters/transfer_queue.py | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index fe112ddcffb..918e9bede97 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -767,24 +767,7 @@ def _from_wire(td: TensorDict) -> TensorDict: if rows and all(row.shape == rows[0].shape for row in rows[1:]): v = torch.stack(rows) changed = True - if field_name in PROMOTE_1D_FIELDS: - if not isinstance(v, torch.Tensor) or v.is_nested: - raise ValueError( - f"Mooncake scalar field {field_name!r} could not be " - "restored as a dense tensor." - ) - if v.dim() == 1: - new_dict[field_name] = v - elif v.dim() == 2 and v.shape[-1] == 1: - new_dict[field_name] = v.squeeze(-1).contiguous() - changed = True - else: - raise ValueError( - f"Mooncake scalar field {field_name!r} must decode as " - f"(N,) or (N, 1), got shape {tuple(v.shape)}." - ) - else: - new_dict[field_name] = v + new_dict[field_name] = v if not changed: # The traversal still ran; only the rebuild was skipped. return td From 8fe9a4af3943505de0b036b6a9eda2b9d873ffdb Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Tue, 8 Sep 2026 22:48:01 -0700 Subject: [PATCH 53/70] refactor(data-plane): share the duplicated observability wiring cluster_step_metrics and get_step_metrics differenced the same counters, differing only in n_procs and the fan-out cost, so the arithmetic moves to one _step_metrics core. The six hash counter names were listed twice in different orders and the comm-volume sum three times; both become one definition. Both trainers hand-copied the same try/except and the same headline+table+print, so those move to metrics_never_fail_the_step and log_step_metrics, which also fixes grpo_sync's print missing flush=True. No metric changes: verified on cluster (n_procs=5) and driver (n_procs=1) paths, hash counters identical to the pre-refactor runs. Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 37 +---- nemo_rl/algorithms/single_controller.py | 26 +--- nemo_rl/data_plane/observability.py | 180 ++++++++++++++++-------- 3 files changed, 133 insertions(+), 110 deletions(-) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index b3975430640..b0f9ddf6e6c 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -32,7 +32,6 @@ from __future__ import annotations import gc -import logging import os import time import warnings @@ -82,10 +81,10 @@ from nemo_rl.data_plane.interfaces import KVBatchMeta from nemo_rl.data_plane.observability import ( MetricsDataPlaneClient, - breakdown_table, cluster_step_metrics, - headline_series, + log_step_metrics, merge_snapshots, + metrics_never_fail_the_step, ) from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS, DP_TRAIN_FIELDS from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -390,19 +389,6 @@ def _compute_seq_logprob_error_metrics( return masking_data["sample_mask"], seq_logprob_error_metrics -def _log_breakdown( - logger: Logger, metrics: dict[str, float], step: int, name: str -) -> None: - """Log the per-op breakdown as a table beside the series. - - A backend without a table type has no rows to log. Failures propagate to - the caller, which already guarantees the logging path cannot fail a step. - """ - columns, rows = breakdown_table(metrics) - if rows: - logger.log_table(columns, rows, step, name) - - def _log_data_plane_metrics( policy: Any, logger: Logger, step: int, total_step_time: float ) -> None: @@ -410,15 +396,8 @@ def _log_data_plane_metrics( On by default, so this runs every step of every recipe. """ - try: + with metrics_never_fail_the_step(step): _log_data_plane_metrics_impl(policy, logger, step, total_step_time) - except Exception as exc: # noqa: BLE001 - a panel must never fail a step - logging.getLogger(__name__).warning( - "data-plane metrics failed at step %d (%s: %s); training continues", - step, - type(exc).__name__, - exc, - ) def _log_data_plane_metrics_impl( @@ -460,17 +439,11 @@ def _log_data_plane_metrics_impl( merged, prev, total_step_time, collect_ms=collect_ms ) policy._prev_cluster_snapshot = merged - logger.log_metrics(headline_series(metrics), step, prefix="data_plane/cluster") - _log_breakdown(logger, metrics, step, "data_plane/cluster/breakdown") + log_step_metrics(logger, metrics, step, "cluster") else: # Single process, or the fan-out could not reach the workers. metrics = client.get_step_metrics(total_step_time) - logger.log_metrics(headline_series(metrics), step, prefix="data_plane/driver") - _log_breakdown(logger, metrics, step, "data_plane/driver/breakdown") - print( - f" • data plane: {metrics['step/wall_s']:.2f}s, " - f"{metrics['step/comm_volume_mb']:.1f} MB moved" - ) + log_step_metrics(logger, metrics, step, "driver") def grpo_train_sync( diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 3796ab27627..f6dd3ad8ba7 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -121,8 +121,8 @@ from nemo_rl.data_plane.async_utils import call_data_plane from nemo_rl.data_plane.observability import ( MetricsDataPlaneClient, - breakdown_table, - headline_series, + log_step_metrics, + metrics_never_fail_the_step, ) from nemo_rl.data_plane.schema import ( DP_CALIB_INPUT_FIELDS, @@ -1529,15 +1529,8 @@ def _log_data_plane_metrics(self, total_step_time: float) -> None: On by default, so this runs every step of every recipe. Mirrors ``grpo_sync._log_data_plane_metrics``. """ - try: + with metrics_never_fail_the_step(self._train_steps): self._log_data_plane_metrics_impl(total_step_time) - except Exception as exc: # noqa: BLE001 - a panel must never fail a step - logging.getLogger(__name__).warning( - "data-plane metrics failed at step %d (%s: %s); training continues", - self._train_steps, - type(exc).__name__, - exc, - ) def _log_data_plane_metrics_impl(self, total_step_time: float) -> None: """Log this step's data-plane cost. No-op unless observability is enabled. @@ -1560,18 +1553,7 @@ def _log_data_plane_metrics_impl(self, total_step_time: float) -> None: return # observability disabled -> plain adapter metrics = self._dp_client.get_step_metrics(total_step_time) - step = self._train_steps - self._logger.log_metrics( - headline_series(metrics), step, prefix="data_plane/driver" - ) - columns, rows = breakdown_table(metrics) - if rows: - self._logger.log_table(columns, rows, step, "data_plane/driver/breakdown") - print( - f" • data plane: {metrics['step/wall_s']:.2f}s, " - f"{metrics['step/comm_volume_mb']:.1f} MB moved", - flush=True, - ) + log_step_metrics(self._logger, metrics, self._train_steps, "driver") @staticmethod def _group_ids_from_meta(meta: KVBatchMeta) -> list[str]: diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 222505ba16c..887b8ad3022 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -41,10 +41,11 @@ import logging import zlib from bisect import bisect_left +from contextlib import contextmanager from dataclasses import asdict, dataclass, field from pathlib import Path from time import monotonic -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from typing import Any, Callable, Literal, TypedDict EventStatus = Literal["ok", "error", "timeout"] @@ -94,11 +95,44 @@ class DataPlaneEvent(TypedDict): _WRITE_OPS = frozenset({"put"}) _READ_OPS = frozenset({"get", "get_data"}) + +def _comm_volume(by_op: dict[str, Any]) -> dict[str, int]: + """Traffic totals derived from ``by_op``, so bytes have one source. + + Distinct from ``bytes_outstanding``, which is occupancy (what is held) + rather than traffic (what moved). + + Args: + by_op: Per-op stats carrying ``n_bytes``. + + Returns: + ``bytes_written``, ``bytes_read``, and their sum. + """ + written = sum(by_op[o]["n_bytes"] for o in _WRITE_OPS if o in by_op) + read = sum(by_op[o]["n_bytes"] for o in _READ_OPS if o in by_op) + return { + "bytes_written": written, + "bytes_read": read, + "comm_volume_bytes": written + read, + } + # A corrupted wire usually corrupts every row of a batch, so the log is # capped: the counter in ``HashStats`` carries the magnitude, and the first # few lines carry the identity of what broke. _MAX_HASH_MISMATCH_LOGS = 20 +# The ``HashStats`` counters, named once: they are differenced into +# ``step/hash/*`` and summed across processes, and the two lists drifting +# apart would silently drop a counter from one path. +_HASH_FIELDS = ( + "rows_recorded", + "rows_checked", + "rows_unverified", + "mismatches", + "fields_skipped", + "guard_failures", +) + # Quantiles reported per op, each with the sample count it needs: enough for # roughly four observations above the rank, or n >= 4 / (1 - q). # @@ -506,15 +540,7 @@ def _hash_deltas(hv: dict[str, int], prev_hv: dict[str, int]) -> dict[str, float if not hv or not (hv.get("rows_recorded") or hv.get("guard_failures")): return {} deltas: dict[str, float] = { - f"step/hash/{name}": hv[name] - prev_hv.get(name, 0) - for name in ( - "rows_checked", - "rows_recorded", - "rows_unverified", - "mismatches", - "fields_skipped", - "guard_failures", - ) + f"step/hash/{name}": hv[name] - prev_hv.get(name, 0) for name in _HASH_FIELDS } # Corruption of every row of every field in a step, repeated identically, # is not what a broken wire looks like -- it is what a broken guard looks @@ -720,17 +746,7 @@ def merge_snapshots(snapshots: "list[dict[str, Any]]") -> dict[str, Any]: return {} merged: dict[str, Any] = {k: 0 for k in _SNAPSHOT_SUM} merged.update({k: 0 for k in _SNAPSHOT_MAX}) - hashes = { - k: 0 - for k in ( - "rows_recorded", - "rows_checked", - "rows_unverified", - "mismatches", - "fields_skipped", - "guard_failures", - ) - } + hashes = {k: 0 for k in _HASH_FIELDS} by_op: dict[str, dict[str, Any]] = {} for snap in snapshots: @@ -760,9 +776,7 @@ def merge_snapshots(snapshots: "list[dict[str, Any]]") -> dict[str, Any]: merged["hash_verify"] = hashes merged["n_processes"] = len(snapshots) _derive_op_metrics(by_op, merged["total_wall_ms"]) - merged["bytes_written"] = sum(by_op[o]["n_bytes"] for o in _WRITE_OPS if o in by_op) - merged["bytes_read"] = sum(by_op[o]["n_bytes"] for o in _READ_OPS if o in by_op) - merged["comm_volume_bytes"] = merged["bytes_written"] + merged["bytes_read"] + merged.update(_comm_volume(by_op)) return merged @@ -789,12 +803,40 @@ def cluster_step_metrics( step_time_s: Step wall time, for ``frac_of_step``. collect_ms: Wall time the caller spent gathering and merging. """ - wall_ms = merged["total_wall_ms"] - prev.get("total_wall_ms", 0.0) - overhead_ms = merged["self_ms"] - prev.get("self_ms", 0.0) + collect_ms n_procs = max(merged.get("n_processes", 1), 1) + metrics = _step_metrics(merged, prev, step_time_s, n_procs, collect_ms) + metrics["now/n_processes"] = n_procs + return metrics + + +def _step_metrics( + snap: dict[str, Any], + prev: dict[str, Any], + step_time_s: float, + n_procs: int = 1, + collect_ms: float = 0.0, +) -> dict[str, float]: + """One step's metrics from two snapshots, cluster-wide or single-process. + + Both callers difference the same counters; only ``n_procs`` (1 off a + single client) and ``collect_ms`` (0 when there was no fan-out to pay + for) differ, so the arithmetic lives here once. + + Args: + snap: This step's snapshot, merged or per-client. + prev: The previous one, for differencing. + step_time_s: Step wall time, for ``frac_of_step``. + n_procs: Processes the snapshot covers. + collect_ms: Wall time spent gathering and merging, if any. + + Returns: + The flat ``step/`` metric dict, less any caller-specific keys. + """ + wall_ms = snap["total_wall_ms"] - prev.get("total_wall_ms", 0.0) + overhead_ms = snap["self_ms"] - prev.get("self_ms", 0.0) + collect_ms # step/ is a delta over this step; now/ is a level at this instant. # The unit alone does not distinguish them -- see README.md. - metrics = _step_deltas(merged, prev) + metrics = _step_deltas(snap, prev) metrics.update( { # The one metric that says whether optimising the data plane is @@ -808,15 +850,14 @@ def cluster_step_metrics( "step/frac_of_step": ( (wall_ms / n_procs) / (step_time_s * 1e3) if step_time_s > 0 else 0.0 ), - "now/n_processes": n_procs, "step/self/overhead_ms": overhead_ms, "step/self/frac": overhead_ms / wall_ms if wall_ms > 0 else 0.0, } ) metrics.update( - _hash_deltas(merged.get("hash_verify") or {}, prev.get("hash_verify") or {}) + _hash_deltas(snap.get("hash_verify") or {}, prev.get("hash_verify") or {}) ) - metrics.update(_op_series(merged["by_op"], prev.get("by_op", {}))) + metrics.update(_op_series(snap["by_op"], prev.get("by_op", {}))) return metrics @@ -927,6 +968,55 @@ def breakdown_table( return ["op", *_BREAKDOWN_COLUMNS], rows +@contextmanager +def metrics_never_fail_the_step(step: int) -> Iterator[None]: + """Swallow anything the metrics panel raises, and say so. + + Observability is on by default, so a fault here would otherwise take + down every step of every recipe -- a panel must never fail training. + + Args: + step: Step number, for the warning. + """ + try: + yield + except Exception as exc: # noqa: BLE001 - a panel must never fail a step + logging.getLogger(__name__).warning( + "data-plane metrics failed at step %d (%s: %s); training continues", + step, + type(exc).__name__, + exc, + ) + + +def log_step_metrics( + logger: Any, metrics: dict[str, float], step: int, scope: str +) -> None: + """Emit one scope's metrics: charted series, breakdown table, console line. + + The series and the table are derived from one ``metrics`` dict, so they + cannot disagree. A backend without a table type has no rows to log. + + Args: + logger: Anything with ``log_metrics`` and ``log_table``. + metrics: Output of :func:`cluster_step_metrics` or + :meth:`MetricsDataPlaneClient.get_step_metrics`. + step: Step number to log against. + scope: ``"cluster"`` or ``"driver"`` -- names the prefix, because the + two differ by roughly the DP degree. + """ + prefix = f"data_plane/{scope}" + logger.log_metrics(headline_series(metrics), step, prefix=prefix) + columns, rows = breakdown_table(metrics) + if rows: + logger.log_table(columns, rows, step, f"{prefix}/breakdown") + print( + f" • data plane: {metrics['step/wall_s']:.2f}s, " + f"{metrics['step/comm_volume_mb']:.1f} MB moved", + flush=True, + ) + + def log_event(event: DataPlaneEvent) -> None: logger.info("data_plane_event: %s", event) @@ -1107,10 +1197,7 @@ def snapshot(self, reset_step_window: bool = False) -> dict[str, Any]: # Communication volume, derived from by_op so there is one source of # truth for bytes. Distinct from ``bytes_outstanding``, which is # occupancy (what is held) rather than traffic (what moved). - by = out["by_op"] - out["bytes_written"] = sum(by[o]["n_bytes"] for o in _WRITE_OPS if o in by) - out["bytes_read"] = sum(by[o]["n_bytes"] for o in _READ_OPS if o in by) - out["comm_volume_bytes"] = out["bytes_written"] + out["bytes_read"] + out.update(_comm_volume(out["by_op"])) if reset_step_window: for bucket in self._stats.by_op.values(): bucket.step_max_ms = 0.0 @@ -1132,28 +1219,9 @@ def get_step_metrics(self, step_time_s: float) -> dict[str, float]: snap = self.snapshot(reset_step_window=True) prev = self._prev_snapshot self._prev_snapshot = snap - - wall_ms = snap["total_wall_ms"] - prev.get("total_wall_ms", 0.0) - # Units are not mixed within one chart: durations charted beside the - # step clock are seconds, the per-op table is ms throughout, volumes - # are always MB. GB was the same problem one dimension over -- a - # realistic step moved 0.00017 GB. - metrics = _step_deltas(snap, prev) - metrics["step/frac_of_step"] = ( - (wall_ms / 1e3 / step_time_s) if step_time_s > 0 else 0.0 - ) - metrics.update( - _hash_deltas(snap.get("hash_verify") or {}, prev.get("hash_verify") or {}) - ) - # The same bill the cluster path reports, under the same name: this - # process's wrapper time, minus what the inner client was doing. - # There is no fan-out to add here -- a single process gathers - # nothing -- so this is the whole of it. - self_ms = snap["self_ms"] - prev.get("self_ms", 0.0) - metrics["step/self/overhead_ms"] = self_ms - metrics["step/self/frac"] = self_ms / wall_ms if wall_ms > 0 else 0.0 - metrics.update(_op_series(snap["by_op"], prev.get("by_op", {}))) - return metrics + # One process, and no fan-out to charge for: the cluster arithmetic + # with n_procs=1 and collect_ms=0 is exactly this path. + return _step_metrics(snap, prev, step_time_s) def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: """Attribute put bytes per key so a later ``clear_samples`` can subtract. From e0fdc93d2da89e06630fa5fc20ef03b05ce8907e Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 9 Sep 2026 02:54:01 -0700 Subject: [PATCH 54/70] chore(data-plane): satisfy ruff, ruff-format and pyrefly on the observability wiring Three pre-commit failures from the previous commit, plus two cleanups found while reading the file: * ruff isort wanted collections.abc sorted above contextlib. * ruff-format wanted a blank line before SingleControllerActor. _log_data_plane_metrics. * pyrefly rejected row_shape at the _leaf_digests call: the jagged and dense branches bind lambdas with different capture defaults, and their union is not assignable to (int) -> tuple[int, ...]. Declaring the name before the branch pins it; the tail=tail / shape=shape defaults are deliberate and stay. * percentile_from_hist is public-named with no caller outside this module (its only user is _clamped_percentiles, and it is not re-exported from nemo_rl.data_plane), so it becomes _percentile_from_hist. * Dropped bytes_f in _emit -- assigned and never read. No behaviour change: the annotation is not evaluated under from __future__ import annotations, and the removed assignment had no reader. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/single_controller.py | 1 + nemo_rl/data_plane/observability.py | 12 ++++++++---- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index f6dd3ad8ba7..2c9c6104396 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1523,6 +1523,7 @@ async def _cleanup_consumed_metas_unlocked( errors.append(cleanup_error) if errors: raise BaseExceptionGroup("post-train DataPlane cleanup failed", errors) + def _log_data_plane_metrics(self, total_step_time: float) -> None: """Log this step's data-plane cost. Never raises. diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 887b8ad3022..91d543c2ed3 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -41,11 +41,11 @@ import logging import zlib from bisect import bisect_left +from collections.abc import Iterator, Sequence from contextlib import contextmanager from dataclasses import asdict, dataclass, field from pathlib import Path from time import monotonic -from collections.abc import Iterator, Sequence from typing import Any, Callable, Literal, TypedDict EventStatus = Literal["ok", "error", "timeout"] @@ -116,6 +116,7 @@ def _comm_volume(by_op: dict[str, Any]) -> dict[str, int]: "comm_volume_bytes": written + read, } + # A corrupted wire usually corrupts every row of a batch, so the log is # capped: the counter in ``HashStats`` carries the magnitude, and the first # few lines carry the identity of what broke. @@ -275,7 +276,7 @@ def _tensor_bytes(v: torch.Tensor) -> int: return buf.nbytes -def percentile_from_hist(hist: list[int], q: float) -> float: +def _percentile_from_hist(hist: list[int], q: float) -> float: """Interpolated ``q``-quantile (0-1) from bucket counts. Linear interpolation inside the containing bucket. A value landing in @@ -666,7 +667,7 @@ def _clamped_percentiles(hist: list[int], max_ms: float) -> dict[str, float]: n = sum(hist) ceiling = max_ms if max_ms > 0 else float("inf") return { - name: min(percentile_from_hist(hist, q), ceiling) + name: min(_percentile_from_hist(hist, q), ceiling) for q, name, min_samples in _QUANTILES if n >= min_samples } @@ -1340,6 +1341,10 @@ def _row_fingerprints( if not isinstance(v, torch.Tensor) or v.ndim < 1: stats.fields_skipped += 1 continue + # Declared up front: the two branches below bind lambdas with + # different capture defaults, and their union is not assignable to + # ``_leaf_digests``'s ``row_shape`` parameter without this. + row_shape: Callable[[int], tuple[int, ...]] if v.is_nested: if v.layout != torch.jagged or v.offsets().numel() - 1 != n_rows: stats.fields_skipped += 1 @@ -1559,7 +1564,6 @@ def _emit( stats.total_ops += 1 bucket.n_bytes += n_bytes bucket.n_keys += n_keys - bytes_f = float(n_bytes) if op == "put" and n_keys: per_key = n_bytes // n_keys stats.last_put_bytes_per_key = per_key From 546827f115baaf1a8a0de8eda654da98e3567e37 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 9 Sep 2026 13:23:35 -0700 Subject: [PATCH 55/70] fix(data-plane): satisfy pyrefly in the observability module - `fold ^ seed`: `Tensor.item()` widens to `int | float`; a hash digest is integral and XOR is int-only, so cast at the use. - the two `row_shape` lambdas bound `tail`/`shape` as default arguments, widening the signature past `Callable[[int], tuple[int, ...]]`. Each is consumed by the `_leaf_digests` call in the same loop iteration, so the late-binding guard was never needed. - `td.items()`: pyrefly cannot narrow tensordict's recursive union when breaking cycles; suppressed in the module's existing idiom. pyrefly clean; tests/unit/data_plane 344 passed, 3 skipped. Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 91d543c2ed3..90dae55f6c9 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -216,7 +216,7 @@ def _leaf_digests( seed = seeds.get(length) if seed is None: seed = seeds[length] = zlib.crc32(f"{dtype}|{row_shape(length)}".encode()) - digests.append(fold ^ seed) + digests.append(int(fold) ^ seed) return digests @@ -419,6 +419,7 @@ def _td_bytes(td: TensorDict | None, max_nodes: int = 10_000) -> int: return 0 budget = [max_nodes] total = 0 + # pyrefly: ignore # bad-assignment for _, v in td.items(include_nested=True, leaves_only=False): if isinstance(v, torch.Tensor): total += _tensor_bytes(v) @@ -1341,8 +1342,8 @@ def _row_fingerprints( if not isinstance(v, torch.Tensor) or v.ndim < 1: stats.fields_skipped += 1 continue - # Declared up front: the two branches below bind lambdas with - # different capture defaults, and their union is not assignable to + # Declared up front: the two branches below bind different + # lambdas, and their union is not assignable to # ``_leaf_digests``'s ``row_shape`` parameter without this. row_shape: Callable[[int], tuple[int, ...]] if v.is_nested: @@ -1354,7 +1355,7 @@ def _row_fingerprints( # A jagged row is its own length followed by the values # buffer's trailing dims. tail = tuple(leaf.shape[1:]) - row_shape = lambda length, tail=tail: (length, *tail) + row_shape = lambda length: (length, *tail) elif v.shape[0] != n_rows: stats.fields_skipped += 1 continue @@ -1370,7 +1371,7 @@ def _row_fingerprints( # instead would make ``(N, L, D)`` say ``(1, L, D)`` and every # round trip a mismatch. shape = tuple(v.shape[1:]) - row_shape = lambda _length, shape=shape: shape + row_shape = lambda _length: shape name = key if isinstance(key, str) else ".".join(key) out[name] = _leaf_digests(leaf, bounds, row_shape, v.dtype) return out From a013ada6c3eee75a94965333ca50bebc3a54c9d1 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 9 Sep 2026 17:46:07 -0700 Subject: [PATCH 56/70] fix(data-plane): release fingerprints when another process clears the sample The wire-hash guard records a fingerprint per row on put and releases it on clear_samples, but on the SingleController path GenWorker and the value actor put through their own clients while only SC clears, so their stores grew for the life of the run. Reconcile against list_sample_ids -- metadata-only, and documented for this -- on a row-count throttle, and run the release through _record_clear so the byte and key accounting follows the same uids. A client that clears normally stays below the throttle and never makes the call. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 36 +++++++++++++++++++-- tests/unit/data_plane/test_observability.py | 36 +++++++++++++++++++++ 2 files changed, 69 insertions(+), 3 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 90dae55f6c9..610b6f17c09 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -122,6 +122,11 @@ def _comm_volume(by_op: dict[str, Any]) -> dict[str, int]: # few lines carry the identity of what broke. _MAX_HASH_MISMATCH_LOGS = 20 +# Rows a client may accumulate between reconciliations of its fingerprint +# store against the partition's live keys. One metadata call per this many +# rows recorded, so a client that clears normally never makes one. +_HASH_RECONCILE_ROWS = 1 << 14 + # The ``HashStats`` counters, named once: they are differenced into # ``step/hash/*`` and summed across processes, and the two lists drifting # apart would silently drop a counter from one path. @@ -1150,10 +1155,12 @@ def __init__( # by the live key population, not by cumulative traffic. self._bytes_by_partition: dict[str, int] = {} self._keys_by_partition: dict[str, set[str]] = {} - # partition -> sample_id -> field -> wire-in fingerprint. Same - # lifetime as ``_bytes_by_partition``: cleared by ``clear_samples``, - # so it is bounded by the live key population. + # partition -> sample_id -> field -> wire-in fingerprint. Released by + # ``clear_samples``, and for the writers that never issue one by + # ``_release_cleared_samples``; either way a fingerprint outlives its + # sample by no more than one reconciliation. self._hash_by_partition: dict[str, dict[str, dict[str, int]]] = {} + self._hash_rows_since_reconcile = 0 self._hash_mismatches_logged = 0 # Set by ``_emit`` to the inner client's wall time for the op just # run, so the wrapping methods can subtract it and bill the rest to @@ -1420,6 +1427,29 @@ def _record_hashes_impl( for name, per_row in digests.items(): per_field[name] = per_row[row] self._stats.hash_verify.rows_recorded += len(sample_ids) + self._hash_rows_since_reconcile += len(sample_ids) + if self._hash_rows_since_reconcile >= _HASH_RECONCILE_ROWS: + # Re-armed before the call, not after: ``_record_hashes`` swallows + # whatever this raises, and a listing that fails on one put fails + # on the next, so re-arming after would retry it on every put. + self._hash_rows_since_reconcile = 0 + self._release_cleared_samples(partition_id) + + def _release_cleared_samples(self, partition_id: str) -> None: + """Release the accounting for samples the partition no longer holds. + + ``clear_samples`` releases it in the process that issues it, which on + the SC path is only ever SC: GenWorker and the value actor put through + their own clients and never clear. Reconciling against + ``list_sample_ids`` -- metadata-only, and documented for exactly this + -- ties their rows to the sample's real lifetime instead. It runs + ``_record_clear`` rather than dropping the fingerprints alone, so the + byte and key accounting follows the same uids. + """ + live = set(self._inner.list_sample_ids(partition_id)) + stale = self._hash_by_partition[partition_id].keys() - live + if stale: + self._record_clear(partition_id, list(stale)) def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: """Compare wire-out fingerprints against what was written. Never raises.""" diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 7358cc2919e..4040ed6dbf4 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -29,6 +29,7 @@ import torch from tensordict import NonTensorData, NonTensorStack, TensorDict +from nemo_rl.data_plane import observability from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient from nemo_rl.data_plane.observability import ( _QUANTILES, @@ -928,6 +929,41 @@ def test_hash_fingerprints_released_on_clear(): client.close() +def test_hash_fingerprints_follow_the_sample_when_another_process_clears( + monkeypatch, +): + """GenWorker and the value actor put through their own clients and never + call ``clear_samples``, so their fingerprints are released by reconciling + against the partition's live keys — a row dropped by SC must not be + retained here, and a row still live must be.""" + monkeypatch.setattr(observability, "_HASH_RECONCILE_ROWS", 4) + inner = NoOpDataPlaneClient() + writer = _client(inner, verify_tensor_hash=True) + writer.put_samples( + sample_ids=_ids(4, prefix="gone"), partition_id="p", fields=_hash_fields() + ) + writer.put_samples( + sample_ids=_ids(4, prefix="live"), partition_id="p", fields=_hash_fields() + ) + + # Another process clears half the partition; the writer never sees the call. + inner.clear_samples(sample_ids=_ids(4, prefix="gone"), partition_id="p") + assert set(writer._hash_by_partition["p"]) == set(_ids(4, prefix="gone")) | set( + _ids(4, prefix="live") + ) + + writer.put_samples( + sample_ids=_ids(4, prefix="next"), partition_id="p", fields=_hash_fields() + ) + assert set(writer._hash_by_partition["p"]) == set(_ids(4, prefix="live")) | set( + _ids(4, prefix="next") + ) + assert writer._keys_by_partition["p"] == set(_ids(4, prefix="live")) | set( + _ids(4, prefix="next") + ) + writer.close() + + def test_hash_fingerprint_covers_jagged_fields(): """The per-token fields on this wire are jagged by the time they reach ``put_samples``. Skipping nested leaves would leave the entire bulk payload From cc054dae0206789e22a8132c4793f4cd2ea84af9 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 9 Sep 2026 20:34:36 -0700 Subject: [PATCH 57/70] fix(data-plane): reconcile the accounting for samples another process cleared Follow-up to 945246045, which hung the release off the hash guard and so only ran with verify_tensor_hash on. The leak is in the byte and key accounting, which is on by default: _record_clear only fires in the process that issues the clear, and on the SC path GenWorker and the value actor put through their own clients while only SC clears. Move the trigger to _record_put and diff against _keys_by_partition, the store every put populates. All three per-partition stores are then released by the one rule a real clear uses, on every client, whatever the guard flag says. A failing list_sample_ids no longer reaches the caller: the put already succeeded, so the release is skipped, logged once, and retried next window. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 86 +++++++++++++-------- tests/unit/data_plane/test_observability.py | 27 ++++--- 2 files changed, 65 insertions(+), 48 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 610b6f17c09..57b5a9aab24 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -122,10 +122,10 @@ def _comm_volume(by_op: dict[str, Any]) -> dict[str, int]: # few lines carry the identity of what broke. _MAX_HASH_MISMATCH_LOGS = 20 -# Rows a client may accumulate between reconciliations of its fingerprint -# store against the partition's live keys. One metadata call per this many -# rows recorded, so a client that clears normally never makes one. -_HASH_RECONCILE_ROWS = 1 << 14 +# Rows a client may write between reconciliations of its live-key accounting +# against the partition. One metadata call per this many rows put, so a client +# that clears its own writes never makes one. +_RECONCILE_ROWS = 1 << 14 # The ``HashStats`` counters, named once: they are differenced into # ``step/hash/*`` and summed across processes, and the two lists drifting @@ -1151,16 +1151,16 @@ def __init__( self._verify_tensor_hash = verify_tensor_hash self._stats = DataPlaneStats() # Live bytes and live keys per partition. Populated on successful - # ``put_samples``, released on successful ``clear_samples``. Bounded - # by the live key population, not by cumulative traffic. + # ``put_samples``, released on successful ``clear_samples`` -- or, for + # a process that never issues one, by ``_release_cleared_samples``. + # Bounded by the live key population, not by cumulative traffic. self._bytes_by_partition: dict[str, int] = {} self._keys_by_partition: dict[str, set[str]] = {} - # partition -> sample_id -> field -> wire-in fingerprint. Released by - # ``clear_samples``, and for the writers that never issue one by - # ``_release_cleared_samples``; either way a fingerprint outlives its - # sample by no more than one reconciliation. + self._rows_since_reconcile = 0 + self._reconcile_failure_logged = False + # partition -> sample_id -> field -> wire-in fingerprint. Same + # lifetime as the two above: ``_record_clear`` releases all three. self._hash_by_partition: dict[str, dict[str, dict[str, int]]] = {} - self._hash_rows_since_reconcile = 0 self._hash_mismatches_logged = 0 # Set by ``_emit`` to the inner client's wall time for the op just # run, so the wrapping methods can subtract it and bill the rest to @@ -1259,6 +1259,13 @@ def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: self._stats.bytes_outstanding += n_bytes if self._stats.bytes_outstanding > self._stats.peak_bytes_outstanding: self._stats.peak_bytes_outstanding = self._stats.bytes_outstanding + self._rows_since_reconcile += len(keys) + if self._rows_since_reconcile >= _RECONCILE_ROWS: + # Re-armed before the call, not after: a listing that fails on one + # put fails on the next, so re-arming after would retry it on every + # put for the rest of the run. + self._rows_since_reconcile = 0 + self._release_cleared_samples(partition_id) def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: """Reverse the put accounting for ``keys``. @@ -1301,6 +1308,40 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: self._bytes_by_partition[partition_id] = total - freed self._stats.bytes_outstanding -= freed + def _release_cleared_samples(self, partition_id: str) -> None: + """Reverse the put accounting for samples another process cleared. + + ``_record_clear`` only fires in the process that issues the clear, + which on the SC path is only ever SC: GenWorker and the value actor + put through their own clients and never clear, so their accounting + would keep every uid they ever wrote. ``list_sample_ids`` is + metadata-only and documented for reconciliation; diffing against it + ties the accounting to the sample's real lifetime. + + The stale uids go through ``_record_clear`` so all three stores are + released by the one rule a real clear uses. + """ + try: + live = set(self._inner.list_sample_ids(partition_id)) + except Exception as exc: # noqa: BLE001 - accounting must not fail a put + # The put itself already succeeded; only the release is lost, and + # the next window retries it. Logged once because whatever makes + # the listing raise makes it raise every window. + if not self._reconcile_failure_logged: + self._reconcile_failure_logged = True + logger.warning( + "data-plane accounting could not reconcile partition %s " + "(%s: %s). Transfers are unaffected; bytes_outstanding " + "and n_keys_outstanding may over-report on this process.", + partition_id, + type(exc).__name__, + exc, + ) + return + stale = self._keys_by_partition.get(partition_id, set()) - live + if stale: + self._record_clear(partition_id, list(stale)) + def _bill_self(self, entered: float) -> None: """Charge this wrapper for the time it spent that was not the RPC. @@ -1427,29 +1468,6 @@ def _record_hashes_impl( for name, per_row in digests.items(): per_field[name] = per_row[row] self._stats.hash_verify.rows_recorded += len(sample_ids) - self._hash_rows_since_reconcile += len(sample_ids) - if self._hash_rows_since_reconcile >= _HASH_RECONCILE_ROWS: - # Re-armed before the call, not after: ``_record_hashes`` swallows - # whatever this raises, and a listing that fails on one put fails - # on the next, so re-arming after would retry it on every put. - self._hash_rows_since_reconcile = 0 - self._release_cleared_samples(partition_id) - - def _release_cleared_samples(self, partition_id: str) -> None: - """Release the accounting for samples the partition no longer holds. - - ``clear_samples`` releases it in the process that issues it, which on - the SC path is only ever SC: GenWorker and the value actor put through - their own clients and never clear. Reconciling against - ``list_sample_ids`` -- metadata-only, and documented for exactly this - -- ties their rows to the sample's real lifetime instead. It runs - ``_record_clear`` rather than dropping the fingerprints alone, so the - byte and key accounting follows the same uids. - """ - live = set(self._inner.list_sample_ids(partition_id)) - stale = self._hash_by_partition[partition_id].keys() - live - if stale: - self._record_clear(partition_id, list(stale)) def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: """Compare wire-out fingerprints against what was written. Never raises.""" diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 4040ed6dbf4..5eb7cb8961e 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -929,16 +929,17 @@ def test_hash_fingerprints_released_on_clear(): client.close() -def test_hash_fingerprints_follow_the_sample_when_another_process_clears( - monkeypatch, +@pytest.mark.parametrize("verify_tensor_hash", [False, True]) +def test_accounting_follows_the_sample_when_another_process_clears( + monkeypatch, verify_tensor_hash ): """GenWorker and the value actor put through their own clients and never - call ``clear_samples``, so their fingerprints are released by reconciling - against the partition's live keys — a row dropped by SC must not be - retained here, and a row still live must be.""" - monkeypatch.setattr(observability, "_HASH_RECONCILE_ROWS", 4) + call ``clear_samples``, so their accounting is released by reconciling + against the partition's live keys. Parametrised over the guard because the + leak is in the byte/key accounting, which is on by default.""" + monkeypatch.setattr(observability, "_RECONCILE_ROWS", 4) inner = NoOpDataPlaneClient() - writer = _client(inner, verify_tensor_hash=True) + writer = _client(inner, verify_tensor_hash=verify_tensor_hash) writer.put_samples( sample_ids=_ids(4, prefix="gone"), partition_id="p", fields=_hash_fields() ) @@ -948,19 +949,17 @@ def test_hash_fingerprints_follow_the_sample_when_another_process_clears( # Another process clears half the partition; the writer never sees the call. inner.clear_samples(sample_ids=_ids(4, prefix="gone"), partition_id="p") - assert set(writer._hash_by_partition["p"]) == set(_ids(4, prefix="gone")) | set( + assert writer._keys_by_partition["p"] == set(_ids(4, prefix="gone")) | set( _ids(4, prefix="live") ) writer.put_samples( sample_ids=_ids(4, prefix="next"), partition_id="p", fields=_hash_fields() ) - assert set(writer._hash_by_partition["p"]) == set(_ids(4, prefix="live")) | set( - _ids(4, prefix="next") - ) - assert writer._keys_by_partition["p"] == set(_ids(4, prefix="live")) | set( - _ids(4, prefix="next") - ) + still_live = set(_ids(4, prefix="live")) | set(_ids(4, prefix="next")) + assert writer._keys_by_partition["p"] == still_live + if verify_tensor_hash: + assert set(writer._hash_by_partition["p"]) == still_live writer.close() From 6f51dbec89b572d0115f2744fc8b5ada8b810dd2 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 9 Sep 2026 21:04:04 -0700 Subject: [PATCH 58/70] refactor(data-plane): drop the reconcile warning, keep the release The one-shot log and its flag were scaffolding around a four-line fix. A failed listing now just skips the release; the next window retries it. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 57b5a9aab24..bfef66203ff 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -1157,7 +1157,6 @@ def __init__( self._bytes_by_partition: dict[str, int] = {} self._keys_by_partition: dict[str, set[str]] = {} self._rows_since_reconcile = 0 - self._reconcile_failure_logged = False # partition -> sample_id -> field -> wire-in fingerprint. Same # lifetime as the two above: ``_record_clear`` releases all three. self._hash_by_partition: dict[str, dict[str, dict[str, int]]] = {} @@ -1323,20 +1322,7 @@ def _release_cleared_samples(self, partition_id: str) -> None: """ try: live = set(self._inner.list_sample_ids(partition_id)) - except Exception as exc: # noqa: BLE001 - accounting must not fail a put - # The put itself already succeeded; only the release is lost, and - # the next window retries it. Logged once because whatever makes - # the listing raise makes it raise every window. - if not self._reconcile_failure_logged: - self._reconcile_failure_logged = True - logger.warning( - "data-plane accounting could not reconcile partition %s " - "(%s: %s). Transfers are unaffected; bytes_outstanding " - "and n_keys_outstanding may over-report on this process.", - partition_id, - type(exc).__name__, - exc, - ) + except Exception: # noqa: BLE001 - the put succeeded; retry next window return stale = self._keys_by_partition.get(partition_id, set()) - live if stale: From 8839675a6b993cd4bc626065550d0ca872c4d087 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 10 Sep 2026 00:20:24 -0700 Subject: [PATCH 59/70] feat(data-plane): make the wire guard check across processes The guard folded a digest on put and re-folded it on get, but held the wire-in reading in the putting process. That verifies a same-process round trip only, and the transfers worth checking are not those: the rollout actor writes what the policy workers read, so every one of those rows landed in rows_unverified and mismatches stayed 0 for want of anything to compare. Hold the readings in one actor instead. Put records, a reader with no local reading fetches -- one fetch per get, not per row. Digests are 8 bytes per row per field, so that is ~100 kB against a get already moving tens of MB. The comparison stays in the reading client so HashStats keeps counting where the transfer happened. Off Ray the handle is None and the local store answers every lookup, which is what single-process callers and the unit tests do. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 37 ++++++-- nemo_rl/data_plane/wire_guard.py | 93 +++++++++++++++++++++ tests/unit/data_plane/test_observability.py | 22 +++++ 3 files changed, 145 insertions(+), 7 deletions(-) create mode 100644 nemo_rl/data_plane/wire_guard.py diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index bfef66203ff..3efa19586f5 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -60,11 +60,13 @@ class DataPlaneEvent(TypedDict): status: EventStatus +import ray import torch from tensordict import NonTensorData, NonTensorStack, TensorDict, TensorDictBase from nemo_rl.data_plane.codec import drain_codec_ms from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta +from nemo_rl.data_plane.wire_guard import get_wire_guard logger = logging.getLogger(__name__) @@ -1160,6 +1162,9 @@ def __init__( # partition -> sample_id -> field -> wire-in fingerprint. Same # lifetime as the two above: ``_record_clear`` releases all three. self._hash_by_partition: dict[str, dict[str, dict[str, int]]] = {} + # Same readings, reachable by the processes that did not write them. + # Only the guard needs it, so only the guard pays for the actor. + self._wire_guard = get_wire_guard() if verify_tensor_hash else None self._hash_mismatches_logged = 0 # Set by ``_emit`` to the inner client's wall time for the op just # run, so the wrapping methods can subtract it and bill the rest to @@ -1283,6 +1288,8 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: """ if self._verify_tensor_hash: _pop_partition_keys(self._hash_by_partition, partition_id, keys) + if self._wire_guard is not None: + self._wire_guard.release.remote(partition_id, keys) live = self._keys_by_partition.get(partition_id) if live is None: return @@ -1448,12 +1455,18 @@ def _record_hashes_impl( digests = self._row_fingerprints(fields, sample_ids) if not digests: return + recorded = { + sample_id: {name: per_row[row] for name, per_row in digests.items()} + for row, sample_id in enumerate(sample_ids) + } partition_hashes = self._hash_by_partition.setdefault(partition_id, {}) - for row, sample_id in enumerate(sample_ids): - per_field = partition_hashes.setdefault(sample_id, {}) - for name, per_row in digests.items(): - per_field[name] = per_row[row] + for sample_id, per_field in recorded.items(): + partition_hashes.setdefault(sample_id, {}).update(per_field) self._stats.hash_verify.rows_recorded += len(sample_ids) + if self._wire_guard is not None: + # Fire and forget: nobody can read a row back before its put + # returns, so the reading cannot be needed before it lands. + self._wire_guard.record.remote(partition_id, recorded) def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: """Compare wire-out fingerprints against what was written. Never raises.""" @@ -1480,12 +1493,22 @@ def _check_hashes_impl( if not digests: return partition_hashes = self._hash_by_partition.get(partition_id, {}) + # Rows this process did not write -- the rollout actor's, on every + # worker read -- have their wire-in reading in the shared store. One + # fetch for the whole batch, not one per row. + remote_hashes: dict[str, dict[str, int]] = {} + if self._wire_guard is not None: + absent = [uid for uid in sample_ids if uid not in partition_hashes] + if absent: + remote_hashes = ray.get( + self._wire_guard.fetch.remote(partition_id, absent) + ) stats = self._stats.hash_verify for row, sample_id in enumerate(sample_ids): - per_field = partition_hashes.get(sample_id) + per_field = partition_hashes.get(sample_id) or remote_hashes.get(sample_id) if not per_field: - # Written by another process (rollout actor, policy worker): - # this client has no wire-in reading to compare against. + # No wire-in reading anywhere: the put predates the guard, or + # the row carried no fingerprintable field. stats.rows_unverified += 1 continue stats.rows_checked += 1 diff --git a/nemo_rl/data_plane/wire_guard.py b/nemo_rl/data_plane/wire_guard.py new file mode 100644 index 00000000000..18cde660f24 --- /dev/null +++ b/nemo_rl/data_plane/wire_guard.py @@ -0,0 +1,93 @@ +# Copyright (c) 2025, 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. +"""Cross-process store for the wire guard's row fingerprints. + +The guard folds a digest per row on ``put_samples`` and re-folds it on +``get_samples``. Holding the wire-in reading in the putting process only ever +verifies a same-process round trip, and the transfers worth checking are not +those: the rollout actor writes what the policy workers read, so every one of +those rows counted as ``rows_unverified``. + +One actor holds the readings instead. A writer records; a reader that has no +local reading fetches. Digests are 8 bytes per row per field, so the fetch is +~100 kB against a get that already moves tens of MB. +""" + +from __future__ import annotations + +import ray + +WIRE_GUARD_ACTOR_NAME = "nemo_rl_wire_guard" + + +class WireGuardState: + """``partition -> sample_id -> field -> wire-in digest``. + + Deliberately dumb: it stores and returns readings, and the comparison + stays in the reading client so ``HashStats`` keeps counting in the process + that owns the transfer. Two clients recording the same row is not a + conflict -- a row written field by field arrives that way. + """ + + def __init__(self) -> None: + self._by_partition: dict[str, dict[str, dict[str, int]]] = {} + + def record( + self, partition_id: str, digests_by_uid: dict[str, dict[str, int]] + ) -> None: + partition = self._by_partition.setdefault(partition_id, {}) + for uid, per_field in digests_by_uid.items(): + partition.setdefault(uid, {}).update(per_field) + + def fetch(self, partition_id: str, uids: list[str]) -> dict[str, dict[str, int]]: + partition = self._by_partition.get(partition_id, {}) + return {uid: partition[uid] for uid in uids if uid in partition} + + def release(self, partition_id: str, uids: list[str] | None) -> None: + if uids is None: + self._by_partition.pop(partition_id, None) + return + partition = self._by_partition.get(partition_id) + if partition is None: + return + for uid in uids: + partition.pop(uid, None) + if not partition: + del self._by_partition[partition_id] + + def n_rows(self) -> int: + """Row count, for tests and for spotting a store that never drains.""" + return sum(len(p) for p in self._by_partition.values()) + + +# The state is a plain class so its logic is testable without a Ray cluster. +WireGuardStore = ray.remote(num_cpus=0)(WireGuardState) + + +def get_wire_guard() -> ray.actor.ActorHandle | None: + """The run's store, created on first call; ``None`` without Ray. + + ``get_if_exists`` rather than a create-then-lookup dance: every client + builds itself independently and any of them may be first. + + Returning ``None`` off-Ray keeps the guard usable in a single process -- + unit tests and the local adapter -- where the putting client is also the + reading one and its own store already answers every lookup. + """ + if not ray.is_initialized(): + return None + return WireGuardStore.options( # type: ignore[attr-defined] + name=WIRE_GUARD_ACTOR_NAME, + get_if_exists=True, + ).remote() diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 5eb7cb8961e..0a07c7c80f6 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -42,6 +42,7 @@ headline_series, merge_snapshots, ) +from nemo_rl.data_plane.wire_guard import WireGuardState # ── helpers ──────────────────────────────────────────────────────────── @@ -963,6 +964,27 @@ def test_accounting_follows_the_sample_when_another_process_clears( writer.close() +def test_wire_guard_state_serves_a_reader_that_never_wrote(): + """The point of the shared store: the rollout actor writes, a policy + worker reads, and the reader has no local wire-in reading of its own.""" + state = WireGuardState() + state.record("p", {"a": {"lp": 11, "ids": 12}, "b": {"lp": 21}}) + + assert state.fetch("p", ["a", "b", "missing"]) == { + "a": {"lp": 11, "ids": 12}, + "b": {"lp": 21}, + } + # A row written field by field accumulates rather than replacing. + state.record("p", {"b": {"ids": 22}}) + assert state.fetch("p", ["b"]) == {"b": {"lp": 21, "ids": 22}} + + state.release("p", ["a"]) + assert state.fetch("p", ["a"]) == {} + assert state.n_rows() == 1 + state.release("p", None) + assert state.n_rows() == 0 + + def test_hash_fingerprint_covers_jagged_fields(): """The per-token fields on this wire are jagged by the time they reach ``put_samples``. Skipping nested leaves would leave the entire bulk payload From 3de6cba376ba72b3b9b2fc7d758ebb947ed2e347 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 10 Sep 2026 00:52:44 -0700 Subject: [PATCH 60/70] Revert "feat(data-plane): make the wire guard check across processes" This reverts commit 332f5ca7c. The actor's store is not persisted, so after load_checkpoint every read reports rows_unverified and the guard silently stops guarding -- the same failure it was added to fix, one level up. It also put a blocking ray.get on every get, serialised through one actor across every DP rank. Carrying the digest in TQ beside the row it describes does not have either property. Reverting so that lands on a clean base. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 37 ++------ nemo_rl/data_plane/wire_guard.py | 93 --------------------- tests/unit/data_plane/test_observability.py | 22 ----- 3 files changed, 7 insertions(+), 145 deletions(-) delete mode 100644 nemo_rl/data_plane/wire_guard.py diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 3efa19586f5..bfef66203ff 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -60,13 +60,11 @@ class DataPlaneEvent(TypedDict): status: EventStatus -import ray import torch from tensordict import NonTensorData, NonTensorStack, TensorDict, TensorDictBase from nemo_rl.data_plane.codec import drain_codec_ms from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta -from nemo_rl.data_plane.wire_guard import get_wire_guard logger = logging.getLogger(__name__) @@ -1162,9 +1160,6 @@ def __init__( # partition -> sample_id -> field -> wire-in fingerprint. Same # lifetime as the two above: ``_record_clear`` releases all three. self._hash_by_partition: dict[str, dict[str, dict[str, int]]] = {} - # Same readings, reachable by the processes that did not write them. - # Only the guard needs it, so only the guard pays for the actor. - self._wire_guard = get_wire_guard() if verify_tensor_hash else None self._hash_mismatches_logged = 0 # Set by ``_emit`` to the inner client's wall time for the op just # run, so the wrapping methods can subtract it and bill the rest to @@ -1288,8 +1283,6 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: """ if self._verify_tensor_hash: _pop_partition_keys(self._hash_by_partition, partition_id, keys) - if self._wire_guard is not None: - self._wire_guard.release.remote(partition_id, keys) live = self._keys_by_partition.get(partition_id) if live is None: return @@ -1455,18 +1448,12 @@ def _record_hashes_impl( digests = self._row_fingerprints(fields, sample_ids) if not digests: return - recorded = { - sample_id: {name: per_row[row] for name, per_row in digests.items()} - for row, sample_id in enumerate(sample_ids) - } partition_hashes = self._hash_by_partition.setdefault(partition_id, {}) - for sample_id, per_field in recorded.items(): - partition_hashes.setdefault(sample_id, {}).update(per_field) + for row, sample_id in enumerate(sample_ids): + per_field = partition_hashes.setdefault(sample_id, {}) + for name, per_row in digests.items(): + per_field[name] = per_row[row] self._stats.hash_verify.rows_recorded += len(sample_ids) - if self._wire_guard is not None: - # Fire and forget: nobody can read a row back before its put - # returns, so the reading cannot be needed before it lands. - self._wire_guard.record.remote(partition_id, recorded) def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: """Compare wire-out fingerprints against what was written. Never raises.""" @@ -1493,22 +1480,12 @@ def _check_hashes_impl( if not digests: return partition_hashes = self._hash_by_partition.get(partition_id, {}) - # Rows this process did not write -- the rollout actor's, on every - # worker read -- have their wire-in reading in the shared store. One - # fetch for the whole batch, not one per row. - remote_hashes: dict[str, dict[str, int]] = {} - if self._wire_guard is not None: - absent = [uid for uid in sample_ids if uid not in partition_hashes] - if absent: - remote_hashes = ray.get( - self._wire_guard.fetch.remote(partition_id, absent) - ) stats = self._stats.hash_verify for row, sample_id in enumerate(sample_ids): - per_field = partition_hashes.get(sample_id) or remote_hashes.get(sample_id) + per_field = partition_hashes.get(sample_id) if not per_field: - # No wire-in reading anywhere: the put predates the guard, or - # the row carried no fingerprintable field. + # Written by another process (rollout actor, policy worker): + # this client has no wire-in reading to compare against. stats.rows_unverified += 1 continue stats.rows_checked += 1 diff --git a/nemo_rl/data_plane/wire_guard.py b/nemo_rl/data_plane/wire_guard.py deleted file mode 100644 index 18cde660f24..00000000000 --- a/nemo_rl/data_plane/wire_guard.py +++ /dev/null @@ -1,93 +0,0 @@ -# Copyright (c) 2025, 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. -"""Cross-process store for the wire guard's row fingerprints. - -The guard folds a digest per row on ``put_samples`` and re-folds it on -``get_samples``. Holding the wire-in reading in the putting process only ever -verifies a same-process round trip, and the transfers worth checking are not -those: the rollout actor writes what the policy workers read, so every one of -those rows counted as ``rows_unverified``. - -One actor holds the readings instead. A writer records; a reader that has no -local reading fetches. Digests are 8 bytes per row per field, so the fetch is -~100 kB against a get that already moves tens of MB. -""" - -from __future__ import annotations - -import ray - -WIRE_GUARD_ACTOR_NAME = "nemo_rl_wire_guard" - - -class WireGuardState: - """``partition -> sample_id -> field -> wire-in digest``. - - Deliberately dumb: it stores and returns readings, and the comparison - stays in the reading client so ``HashStats`` keeps counting in the process - that owns the transfer. Two clients recording the same row is not a - conflict -- a row written field by field arrives that way. - """ - - def __init__(self) -> None: - self._by_partition: dict[str, dict[str, dict[str, int]]] = {} - - def record( - self, partition_id: str, digests_by_uid: dict[str, dict[str, int]] - ) -> None: - partition = self._by_partition.setdefault(partition_id, {}) - for uid, per_field in digests_by_uid.items(): - partition.setdefault(uid, {}).update(per_field) - - def fetch(self, partition_id: str, uids: list[str]) -> dict[str, dict[str, int]]: - partition = self._by_partition.get(partition_id, {}) - return {uid: partition[uid] for uid in uids if uid in partition} - - def release(self, partition_id: str, uids: list[str] | None) -> None: - if uids is None: - self._by_partition.pop(partition_id, None) - return - partition = self._by_partition.get(partition_id) - if partition is None: - return - for uid in uids: - partition.pop(uid, None) - if not partition: - del self._by_partition[partition_id] - - def n_rows(self) -> int: - """Row count, for tests and for spotting a store that never drains.""" - return sum(len(p) for p in self._by_partition.values()) - - -# The state is a plain class so its logic is testable without a Ray cluster. -WireGuardStore = ray.remote(num_cpus=0)(WireGuardState) - - -def get_wire_guard() -> ray.actor.ActorHandle | None: - """The run's store, created on first call; ``None`` without Ray. - - ``get_if_exists`` rather than a create-then-lookup dance: every client - builds itself independently and any of them may be first. - - Returning ``None`` off-Ray keeps the guard usable in a single process -- - unit tests and the local adapter -- where the putting client is also the - reading one and its own store already answers every lookup. - """ - if not ray.is_initialized(): - return None - return WireGuardStore.options( # type: ignore[attr-defined] - name=WIRE_GUARD_ACTOR_NAME, - get_if_exists=True, - ).remote() diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 0a07c7c80f6..5eb7cb8961e 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -42,7 +42,6 @@ headline_series, merge_snapshots, ) -from nemo_rl.data_plane.wire_guard import WireGuardState # ── helpers ──────────────────────────────────────────────────────────── @@ -964,27 +963,6 @@ def test_accounting_follows_the_sample_when_another_process_clears( writer.close() -def test_wire_guard_state_serves_a_reader_that_never_wrote(): - """The point of the shared store: the rollout actor writes, a policy - worker reads, and the reader has no local wire-in reading of its own.""" - state = WireGuardState() - state.record("p", {"a": {"lp": 11, "ids": 12}, "b": {"lp": 21}}) - - assert state.fetch("p", ["a", "b", "missing"]) == { - "a": {"lp": 11, "ids": 12}, - "b": {"lp": 21}, - } - # A row written field by field accumulates rather than replacing. - state.record("p", {"b": {"ids": 22}}) - assert state.fetch("p", ["b"]) == {"b": {"lp": 21, "ids": 22}} - - state.release("p", ["a"]) - assert state.fetch("p", ["a"]) == {} - assert state.n_rows() == 1 - state.release("p", None) - assert state.n_rows() == 0 - - def test_hash_fingerprint_covers_jagged_fields(): """The per-token fields on this wire are jagged by the time they reach ``put_samples``. Skipping nested leaves would leave the entire bulk payload From 677f2d2341d9dd70975c137b0983c7acb6be6b38 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 10 Sep 2026 01:08:08 -0700 Subject: [PATCH 61/70] feat(data-plane): carry the wire-in digest beside the field it describes WIP: production code only, the 12 existing hash tests still assert on the local store this removes. The guard folded a digest on put and re-folded it on get, but kept the wire-in reading in the putting process. That verifies a same-process round trip only, and the transfer worth checking is not one: the rollout actor writes what the policy workers read, so every worker read landed in rows_unverified and mismatches stayed 0 for want of anything to compare. Mirror each field as _hash, one int64 per row, written by the same put. A reader fetches it alongside, re-folds, compares, and strips it before the caller sees it. The digest now shares storage, sharding and checkpoint lifetime with the row it describes. Per top-level field rather than per leaf, because select_fields names top-level fields: a multimodal `images` reduces its leaves to one images_hash, folded in sorted leaf order with `* 31 +` so identical leaves cannot cancel. Removes _hash_by_partition and _pop_partition_keys: with the reading on the wire there is no process-local fingerprint state left to leak or reconcile. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 186 ++++++++++++++++++---------- 1 file changed, 120 insertions(+), 66 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index bfef66203ff..ce8c1eb2d3b 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -122,6 +122,23 @@ def _comm_volume(by_op: dict[str, Any]) -> dict[str, int]: # few lines carry the identity of what broke. _MAX_HASH_MISMATCH_LOGS = 20 +# The wire-in digest rides beside the field it describes, as ``_hash``. +# Holding it in the putting process only ever verified a same-process round +# trip; the rollout actor writes what the policy workers read, and that read +# is the one worth checking. +_HASH_SUFFIX = "_hash" +_U64 = 1 << 64 + + +def _hash_field(name: str) -> str: + return f"{name}{_HASH_SUFFIX}" + + +def _as_i64(value: int) -> int: + """Wrap a uint64 digest into the signed range ``torch.int64`` accepts.""" + value &= _U64 - 1 + return value - _U64 if value >= _U64 >> 1 else value + # Rows a client may write between reconciliations of its live-key accounting # against the partition. One metadata call per this many rows put, so a client # that clears its own writes never makes one. @@ -225,6 +242,28 @@ def _leaf_digests( return digests +def _field_digests( + leaf_digests: dict[str, list[int]], n_rows: int +) -> dict[str, list[int]]: + """Leaf digests folded to one digest per *top-level* field. + + ``select_fields`` names top-level fields, so the mirror has to be per + field rather than per leaf: a multimodal ``images`` arrives as several + leaves and must reduce to a single ``images_hash`` that the reader can + recompute from the same leaves. + + Sorted leaf order because dict order need not survive a round trip, and + ``* 31 +`` rather than an XOR so two identical leaves do not cancel -- + the defect the row fold already has, which must not be repeated here. + """ + out: dict[str, list[int]] = {} + for name, per_row in sorted(leaf_digests.items()): + acc = out.setdefault(name.split(".", 1)[0], [0] * n_rows) + for row in range(n_rows): + acc[row] = _as_i64(acc[row] * 31 + per_row[row]) + return out + + def _as_list(sample_ids: Any) -> Any: """Materialize ``sample_ids`` once; ``None`` passes through. @@ -236,26 +275,6 @@ def _as_list(sample_ids: Any) -> Any: return list(sample_ids) -def _pop_partition_keys( - store: dict[str, dict[str, Any]], partition_id: str, keys: list[str] | None -) -> list[Any]: - """Drop ``keys`` from ``store[partition_id]``, returning what was removed. - - ``keys=None`` drops the whole partition. Shared by the byte accounting - and the fingerprint store so their teardown cannot drift apart. - """ - partition = store.get(partition_id) - if partition is None: - return [] - if keys is None: - del store[partition_id] - return list(partition.values()) - removed = [partition.pop(key) for key in keys if key in partition] - if not partition: - del store[partition_id] - return removed - - def _tensor_bytes(v: torch.Tensor) -> int: """Wire bytes of one tensor leaf, rectangular or nested. @@ -1157,9 +1176,6 @@ def __init__( self._bytes_by_partition: dict[str, int] = {} self._keys_by_partition: dict[str, set[str]] = {} self._rows_since_reconcile = 0 - # partition -> sample_id -> field -> wire-in fingerprint. Same - # lifetime as the two above: ``_record_clear`` releases all three. - self._hash_by_partition: dict[str, dict[str, dict[str, int]]] = {} self._hash_mismatches_logged = 0 # Set by ``_emit`` to the inner client's wall time for the op just # run, so the wrapping methods can subtract it and bill the rest to @@ -1281,8 +1297,6 @@ def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: partition_id: Partition the keys were dropped from. keys: Uids dropped; ``None`` means the whole partition was cleared. """ - if self._verify_tensor_hash: - _pop_partition_keys(self._hash_by_partition, partition_id, keys) live = self._keys_by_partition.get(partition_id) if live is None: return @@ -1432,28 +1446,40 @@ def _hash_guard_failed(self, op: str, exc: Exception) -> None: exc, ) - def _record_hashes( - self, partition_id: str, sample_ids: list[str], fields: TensorDict | None - ) -> None: - """Store wire-in fingerprints for a successful put. Never raises.""" + def _stamp_hashes( + self, sample_ids: list[str], fields: TensorDict | None + ) -> TensorDict | None: + """Return ``fields`` with a ``_hash`` column beside each field. + + Never raises: a guard that cannot fold must not stop the put, so the + original ``fields`` goes on the wire unstamped and the batch reads as + unverified on the far side. + """ try: - self._record_hashes_impl(partition_id, sample_ids, fields) + return self._stamp_hashes_impl(sample_ids, fields) except Exception as exc: # noqa: BLE001 - a debug check must never fail a transfer self._hash_guard_failed("put", exc) - - def _record_hashes_impl( - self, partition_id: str, sample_ids: list[str], fields: TensorDict | None - ) -> None: - """Store wire-in fingerprints for a successful put.""" - digests = self._row_fingerprints(fields, sample_ids) - if not digests: - return - partition_hashes = self._hash_by_partition.setdefault(partition_id, {}) - for row, sample_id in enumerate(sample_ids): - per_field = partition_hashes.setdefault(sample_id, {}) - for name, per_row in digests.items(): - per_field[name] = per_row[row] + return fields + + def _stamp_hashes_impl( + self, sample_ids: list[str], fields: TensorDict | None + ) -> TensorDict | None: + if fields is None: + return fields + digests = _field_digests( + self._row_fingerprints(fields, sample_ids), len(sample_ids) + ) + stamped = fields.copy() + # Every top-level field gets a column, including the ones the fold + # could not attribute per row -- those carry 0, which the reader takes + # as "no reading". A column that is sometimes absent would make the + # reader's fetch fail on a partition it has no business failing on. + unfolded = [0] * len(sample_ids) + for name in fields.keys(): + column = digests.get(name, unfolded) + stamped[_hash_field(name)] = torch.tensor(column, dtype=torch.int64) self._stats.hash_verify.rows_recorded += len(sample_ids) + return stamped def _check_hashes(self, partition_id: str, sample_ids: list[str], out: Any) -> None: """Compare wire-out fingerprints against what was written. Never raises.""" @@ -1467,31 +1493,43 @@ def _check_hashes_impl( ) -> None: """Compare wire-out fingerprints against what was written. - Every field is comparable now. The old two-tier classification -- - which fields could be compared, which were a shard of a batch-scoped - put, which had been written uniform and read back ragged -- existed - entirely to work around a digest that only meant something against - the exact batch it was folded over. A per-row hash reconciles against - any grouping, so a shard read is checked rather than abstained on. + The wire-in reading arrives with the row, so a shard read by a process + that never wrote it reconciles the same as a same-process round trip. + The mirror columns are stripped here: the caller asked for ``tokens`` + and must never see ``tokens_hash``. """ if not isinstance(out, TensorDict): return - digests = self._row_fingerprints(out, sample_ids) + expected_by_field: dict[str, list[int]] = {} + for key in list(out.keys()): + if isinstance(key, str) and key.endswith(_HASH_SUFFIX): + expected_by_field[key[: -len(_HASH_SUFFIX)]] = out.get(key).tolist() + del out[key] + digests = _field_digests( + self._row_fingerprints(out, sample_ids), len(sample_ids) + ) if not digests: return - partition_hashes = self._hash_by_partition.get(partition_id, {}) stats = self._stats.hash_verify for row, sample_id in enumerate(sample_ids): - per_field = partition_hashes.get(sample_id) - if not per_field: - # Written by another process (rollout actor, policy worker): - # this client has no wire-in reading to compare against. + # ``0`` is the writer saying it could not fold that field, so it is + # an abstention rather than a reading. A real digest of 0 is + # possible and goes unchecked; at one row in 2^64 that is cheaper + # than a false alarm on every asymmetric fold. + comparable = [ + (name, per_row) + for name, per_row in digests.items() + if expected_by_field.get(name, [0] * len(sample_ids))[row] != 0 + ] + if not comparable: + # Written without the mirror: a put that predates the guard, + # or a field the fold could not attribute per row. stats.rows_unverified += 1 continue stats.rows_checked += 1 - for name, per_row in digests.items(): - expected = per_field.get(name) - if expected is None or expected == per_row[row]: + for name, per_row in comparable: + expected = expected_by_field[name][row] + if expected == per_row[row]: continue stats.mismatches += 1 if self._hash_mismatches_logged < _MAX_HASH_MISMATCH_LOGS: @@ -1614,6 +1652,15 @@ def register_partition( grpo_group_size=None, enums=None, ): + if self._verify_tensor_hash: + clash = [f for f in fields if f.endswith(_HASH_SUFFIX)] + if clash: + raise ValueError( + f"partition {partition_id!r} declares {clash}, which the " + f"wire guard's mirror columns would shadow. Rename them or " + f"set observability.verify_tensor_hash=false." + ) + fields = list(fields) + [_hash_field(f) for f in fields] self._run( "register", partition_id, @@ -1654,10 +1701,13 @@ def claim_meta( def get_data(self, meta, select_fields=None): entered = monotonic() + fetch = select_fields if select_fields is not None else meta.fields + if self._verify_tensor_hash and fetch is not None: + fetch = list(fetch) + [_hash_field(f) for f in fetch] out = self._run( "get_data", meta.partition_id, - lambda: self._inner.get_data(meta, select_fields=select_fields), + lambda: self._inner.get_data(meta, select_fields=fetch), n_keys=len(meta.sample_ids), ) if self._verify_tensor_hash: @@ -1678,37 +1728,41 @@ def put_samples(self, sample_ids, partition_id, fields=None, tags=None): # Materialize once: ``_run`` consumes its lambda and we also need # to attribute bytes per sample after success. sample_ids_list = _as_list(sample_ids) + # Folded before ``_run``, not after: the digest travels in the payload + # now, so it has to exist before the RPC. The fold still lands outside + # the op's ``wall_ms`` -- ``_bill_self`` charges it to ``self_ms``. + payload = fields + if self._verify_tensor_hash: + payload = self._stamp_hashes(sample_ids_list, fields) out = self._run( "put", partition_id, lambda: self._inner.put_samples( sample_ids_list, partition_id, - fields=fields, + fields=payload, tags=tags, ), n_keys=len(sample_ids_list), n_bytes=n_bytes, ) self._record_put(partition_id, sample_ids_list, n_bytes) - # Fingerprinted after ``_run`` rather than inside it: ``fields`` is - # the caller's TensorDict and the RPC does not mutate it, so hashing - # here keeps the check's own cost out of the op's ``wall_ms``. - if self._verify_tensor_hash: - self._record_hashes(partition_id, sample_ids_list, fields) self._bill_self(entered) return out def get_samples(self, sample_ids, partition_id, select_fields): entered = monotonic() sample_ids_list = _as_list(sample_ids) + fetch = list(select_fields) + if self._verify_tensor_hash: + fetch += [_hash_field(f) for f in select_fields] out = self._run( "get", partition_id, lambda: self._inner.get_samples( sample_ids_list, partition_id, - select_fields=select_fields, + select_fields=fetch, ), n_keys=len(sample_ids_list), ) From c78cc8de6247f2ca19c4692915cda85e074de9a7 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 10 Sep 2026 01:21:47 -0700 Subject: [PATCH 62/70] test(data-plane): rewrite the hash tests against the mirror columns The local fingerprint store is gone, so the tests that asserted on it had to say something else. Two invert outright: a reader that never wrote the row now verifies it, and corruption on that trip is caught rather than abstained on -- which is the point of the change. Also restores the read fallback removed in the previous commit. Writing test_a_put_that_could_not_stamp_leaves_reads_working showed why it is not optional: when the put-side fold raises, the rows carry no mirror, and asking for one turns a guard bug into a failed read. "A bug in the guard must not take the transfer down" has to hold on both sides of the wire. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 31 +++++--- tests/unit/data_plane/test_observability.py | 78 +++++++++++++++++---- 2 files changed, 86 insertions(+), 23 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index ce8c1eb2d3b..b6cc0751cb7 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -1756,16 +1756,29 @@ def get_samples(self, sample_ids, partition_id, select_fields): fetch = list(select_fields) if self._verify_tensor_hash: fetch += [_hash_field(f) for f in select_fields] - out = self._run( - "get", - partition_id, - lambda: self._inner.get_samples( - sample_ids_list, + def read(columns): + return self._run( + "get", partition_id, - select_fields=fetch, - ), - n_keys=len(sample_ids_list), - ) + lambda: self._inner.get_samples( + sample_ids_list, + partition_id, + select_fields=columns, + ), + n_keys=len(sample_ids_list), + ) + + try: + out = read(fetch) + except Exception as exc: # noqa: BLE001 - a guard bug must not fail a read + if fetch == list(select_fields): + raise + # The mirror is absent: the put that wrote these rows could not + # fold, or predates the guard. Read what the caller asked for and + # abstain -- ``_check_hashes`` finds no columns and counts the rows + # unverified. + self._hash_guard_failed("get", exc) + out = read(list(select_fields)) if self._verify_tensor_hash: self._check_hashes(partition_id, sample_ids_list, out) self._bill_self(entered) diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 5eb7cb8961e..0824d687034 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -808,7 +808,6 @@ def test_hash_state_and_counters_absent_when_the_guard_is_off(): merged = merge_snapshots([client.snapshot(reset_step_window=True)]) assert client.snapshot()["hash_verify"]["rows_recorded"] == 0 - assert client._hash_by_partition == {} assert not [k for k in client.get_step_metrics(1.0) if "hash" in k] assert not [k for k in cluster_step_metrics(merged, {}, 1.0) if "hash" in k] client.close() @@ -898,9 +897,10 @@ def test_hash_verification_survives_shard_readback(): client.close() -def test_hash_verification_reports_rows_it_never_wrote(): - """A consumer-side client sees only wire-out. Those rows must land in - ``rows_unverified`` — reporting 0 mismatches would read as 'clean'.""" +def test_hash_verification_checks_rows_it_never_wrote(): + """The transfer worth checking: the rollout actor writes, a policy worker + reads. The reading arrives with the row, so the reader verifies it without + ever having held the wire-in fold itself.""" inner = NoOpDataPlaneClient() writer = _client(inner, verify_tensor_hash=True) ids = _ids(4) @@ -910,22 +910,74 @@ def test_hash_verification_reports_rows_it_never_wrote(): reader.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) hv = reader.snapshot()["hash_verify"] - assert hv["rows_unverified"] == 4 - assert hv["rows_checked"] == 0 + assert hv["rows_checked"] == 4 + assert hv["rows_unverified"] == 0 assert hv["mismatches"] == 0 writer.close() -def test_hash_fingerprints_released_on_clear(): - """Fingerprints must be bounded by the live key population, not by - cumulative traffic.""" +def test_hash_verification_catches_corruption_across_processes(): + """The same trip, corrupted. Before the reading travelled, this read was + abstained on and the corruption reached training silently.""" + inner = _CorruptingClient(field="ids", row=2) + writer = _client(inner, verify_tensor_hash=True) + ids = _ids(4) + writer.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + + reader = MetricsDataPlaneClient(inner, verify_tensor_hash=True) + reader.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + + assert reader.snapshot()["hash_verify"]["mismatches"] == 1 + writer.close() + + +def test_mirror_columns_never_reach_the_caller(): + """The caller asked for ``ids``; ``ids_hash`` is the guard's business. + A leaked key breaks every consumer that iterates the returned fields.""" + client = _client(verify_tensor_hash=True) + ids = _ids(4) + client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + out = client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + + assert set(out.keys()) == {"ids"} + assert client.snapshot()["hash_verify"]["rows_checked"] == 4 + client.close() + + +def test_mirror_column_rides_in_the_partition(): + """It is a field like any other, so it survives whatever the row survives + -- sharding, checkpointing -- rather than living in one process.""" + inner = NoOpDataPlaneClient() + client = _client(inner, verify_tensor_hash=True) + ids = _ids(4) + client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + + stored = inner.get_samples(ids, "p", select_fields=["ids", "ids_hash", "lp_hash"]) + assert stored["ids_hash"].dtype is torch.int64 + assert stored["ids_hash"].shape == (4,) + assert (stored["ids_hash"] != stored["lp_hash"]).all(), "one mirror per field" + client.close() + + +def test_a_put_that_could_not_stamp_leaves_reads_working(monkeypatch): + """A guard bug must never take a transfer down -- on either side. Without + the mirror the read still has to succeed, and abstain rather than claim + the rows were clean.""" client = _client(verify_tensor_hash=True) ids = _ids(4) + + def boom(*_args, **_kwargs): + raise NotImplementedError("no digest kernel for this dtype") + + monkeypatch.setattr(client, "_row_fingerprints", boom) client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) + monkeypatch.undo() - assert client._hash_by_partition["p"] - client.clear_samples(sample_ids=ids, partition_id="p") - assert client._hash_by_partition == {} + out = client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) + assert set(out.keys()) == {"ids"} + hv = client.snapshot()["hash_verify"] + assert hv["rows_unverified"] == 4 and hv["rows_checked"] == 0 + assert hv["guard_failures"] >= 1, "the abstention is visible, not silent" client.close() @@ -958,8 +1010,6 @@ def test_accounting_follows_the_sample_when_another_process_clears( ) still_live = set(_ids(4, prefix="live")) | set(_ids(4, prefix="next")) assert writer._keys_by_partition["p"] == still_live - if verify_tensor_hash: - assert set(writer._hash_by_partition["p"]) == still_live writer.close() From 6d2495f3403b75dda7cc89768fe29a224582faf9 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 10 Sep 2026 02:30:56 -0700 Subject: [PATCH 63/70] fix(data-plane): one definition of the mirror column list `get_data` suffixed `meta.fields`, which comes back from `put_samples` already carrying the mirrors -- so it asked for `tokens_hash_hash`, columns that do not exist, on a path with no fallback. `_with_mirrors` is idempotent and is now the only place the list is built; register, get_data and get_samples all call it. Also makes the read fallback's condition a config check rather than a list comparison, and corrects the guard-failure count: a put that cannot stamp leaves the read without mirrors, so the read falls back and its check then fails too -- three unchecked batches, not two. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/observability.py | 19 +++++++++++++++---- tests/unit/data_plane/test_observability.py | 7 +++++-- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index b6cc0751cb7..8f5a56c5533 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -134,6 +134,17 @@ def _hash_field(name: str) -> str: return f"{name}{_HASH_SUFFIX}" +def _with_mirrors(fields: Sequence[str]) -> list[str]: + """``fields`` followed by one mirror column each. + + Idempotent: ``meta.fields`` comes back from ``put_samples`` already + carrying the mirrors, and suffixing those would ask for + ``tokens_hash_hash``. + """ + plain = [f for f in fields if not f.endswith(_HASH_SUFFIX)] + return [*plain, *(_hash_field(f) for f in plain)] + + def _as_i64(value: int) -> int: """Wrap a uint64 digest into the signed range ``torch.int64`` accepts.""" value &= _U64 - 1 @@ -1660,7 +1671,7 @@ def register_partition( f"wire guard's mirror columns would shadow. Rename them or " f"set observability.verify_tensor_hash=false." ) - fields = list(fields) + [_hash_field(f) for f in fields] + fields = _with_mirrors(fields) self._run( "register", partition_id, @@ -1703,7 +1714,7 @@ def get_data(self, meta, select_fields=None): entered = monotonic() fetch = select_fields if select_fields is not None else meta.fields if self._verify_tensor_hash and fetch is not None: - fetch = list(fetch) + [_hash_field(f) for f in fetch] + fetch = _with_mirrors(fetch) out = self._run( "get_data", meta.partition_id, @@ -1755,7 +1766,7 @@ def get_samples(self, sample_ids, partition_id, select_fields): sample_ids_list = _as_list(sample_ids) fetch = list(select_fields) if self._verify_tensor_hash: - fetch += [_hash_field(f) for f in select_fields] + fetch = _with_mirrors(select_fields) def read(columns): return self._run( "get", @@ -1771,7 +1782,7 @@ def read(columns): try: out = read(fetch) except Exception as exc: # noqa: BLE001 - a guard bug must not fail a read - if fetch == list(select_fields): + if not self._verify_tensor_hash: raise # The mirror is absent: the put that wrote these rows could not # fold, or predates the guard. Read what the caller asked for and diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 0824d687034..3e0ec5a3bb0 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -852,11 +852,14 @@ def boom(*_args, **_kwargs): client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) hv = client.snapshot()["hash_verify"] - assert hv["guard_failures"] == 2, "put and get each failed once" + # Three, not two: the put cannot stamp, so the read's mirror columns are + # absent and the read falls back, and the fold the fallback's check then + # attempts raises too. Every one of them is a batch that went unchecked. + assert hv["guard_failures"] == 3, "the put, the mirror-less read, its check" assert hv["rows_recorded"] == 0 and hv["rows_checked"] == 0, "nothing checked" assert caplog.text.count("hash guard failed") == 1, "logged once, not per call" # and it is a series, not just a counter -- the gate cannot key on rows - assert client.get_step_metrics(1.0)["step/hash/guard_failures"] == 2 + assert client.get_step_metrics(1.0)["step/hash/guard_failures"] == 3 client.close() From 9619fc23e8c5ed2ff6e57c6f067be8d8052b0057 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 10 Sep 2026 02:31:57 -0700 Subject: [PATCH 64/70] test(data-plane): one sanity check, not five The change makes one claim -- a reader that never wrote the row can now verify it -- so it needs one test, parametrised clean/corrupted. The strip is one assertion on the existing clean round trip, and the guard-failure absorption path was already covered by test_guard_failure_is_absorbed_counted_and_charted. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- tests/unit/data_plane/test_observability.py | 88 ++++----------------- 1 file changed, 17 insertions(+), 71 deletions(-) diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 3e0ec5a3bb0..f9dccacde4d 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -817,7 +817,13 @@ def test_hash_verification_clean_roundtrip(): client = _client(verify_tensor_hash=True) ids = _ids(4) client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) - client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids", "lp"]) + out = client.get_samples( + sample_ids=ids, partition_id="p", select_fields=["ids", "lp"] + ) + + # The caller asked for two fields; the mirror columns are the guard's + # business and a leaked key breaks anything that iterates the result. + assert set(out.keys()) == {"ids", "lp"} assert client.snapshot()["hash_verify"] == { "rows_recorded": 4, @@ -900,11 +906,16 @@ def test_hash_verification_survives_shard_readback(): client.close() -def test_hash_verification_checks_rows_it_never_wrote(): +@pytest.mark.parametrize( + "inner, mismatches", + [(NoOpDataPlaneClient(), 0), (_CorruptingClient(field="ids", row=2), 1)], + ids=["clean", "corrupted"], +) +def test_a_reader_that_never_wrote_the_row_still_verifies_it(inner, mismatches): """The transfer worth checking: the rollout actor writes, a policy worker - reads. The reading arrives with the row, so the reader verifies it without - ever having held the wire-in fold itself.""" - inner = NoOpDataPlaneClient() + reads. The wire-in reading arrives with the row, so the reader compares + without ever having held the fold itself — and catches a corrupted trip + that used to be abstained on.""" writer = _client(inner, verify_tensor_hash=True) ids = _ids(4) writer.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) @@ -915,75 +926,10 @@ def test_hash_verification_checks_rows_it_never_wrote(): hv = reader.snapshot()["hash_verify"] assert hv["rows_checked"] == 4 assert hv["rows_unverified"] == 0 - assert hv["mismatches"] == 0 + assert hv["mismatches"] == mismatches writer.close() -def test_hash_verification_catches_corruption_across_processes(): - """The same trip, corrupted. Before the reading travelled, this read was - abstained on and the corruption reached training silently.""" - inner = _CorruptingClient(field="ids", row=2) - writer = _client(inner, verify_tensor_hash=True) - ids = _ids(4) - writer.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) - - reader = MetricsDataPlaneClient(inner, verify_tensor_hash=True) - reader.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) - - assert reader.snapshot()["hash_verify"]["mismatches"] == 1 - writer.close() - - -def test_mirror_columns_never_reach_the_caller(): - """The caller asked for ``ids``; ``ids_hash`` is the guard's business. - A leaked key breaks every consumer that iterates the returned fields.""" - client = _client(verify_tensor_hash=True) - ids = _ids(4) - client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) - out = client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) - - assert set(out.keys()) == {"ids"} - assert client.snapshot()["hash_verify"]["rows_checked"] == 4 - client.close() - - -def test_mirror_column_rides_in_the_partition(): - """It is a field like any other, so it survives whatever the row survives - -- sharding, checkpointing -- rather than living in one process.""" - inner = NoOpDataPlaneClient() - client = _client(inner, verify_tensor_hash=True) - ids = _ids(4) - client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) - - stored = inner.get_samples(ids, "p", select_fields=["ids", "ids_hash", "lp_hash"]) - assert stored["ids_hash"].dtype is torch.int64 - assert stored["ids_hash"].shape == (4,) - assert (stored["ids_hash"] != stored["lp_hash"]).all(), "one mirror per field" - client.close() - - -def test_a_put_that_could_not_stamp_leaves_reads_working(monkeypatch): - """A guard bug must never take a transfer down -- on either side. Without - the mirror the read still has to succeed, and abstain rather than claim - the rows were clean.""" - client = _client(verify_tensor_hash=True) - ids = _ids(4) - - def boom(*_args, **_kwargs): - raise NotImplementedError("no digest kernel for this dtype") - - monkeypatch.setattr(client, "_row_fingerprints", boom) - client.put_samples(sample_ids=ids, partition_id="p", fields=_hash_fields()) - monkeypatch.undo() - - out = client.get_samples(sample_ids=ids, partition_id="p", select_fields=["ids"]) - assert set(out.keys()) == {"ids"} - hv = client.snapshot()["hash_verify"] - assert hv["rows_unverified"] == 4 and hv["rows_checked"] == 0 - assert hv["guard_failures"] >= 1, "the abstention is visible, not silent" - client.close() - - @pytest.mark.parametrize("verify_tensor_hash", [False, True]) def test_accounting_follows_the_sample_when_another_process_clears( monkeypatch, verify_tensor_hash From 8cbe1a9ac5694cbd4cbe5e67c8da1f355f9109a7 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 10 Sep 2026 12:46:57 -0700 Subject: [PATCH 65/70] test(data-plane): the later writer stamps too The test injected its extra rows straight into the fixture's store, which modelled a writer that does not stamp. Every writer in a run shares one verify_tensor_hash, so a real later writer carries mirrors like anyone else -- and now all six rows verify, including the two this client never wrote, which is what the change is for. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- tests/unit/data_plane/test_observability.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index f9dccacde4d..b7a0842fb4a 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -1165,8 +1165,8 @@ def test_uniform_write_read_back_inside_a_ragged_batch_is_clean(): """The false positive this once cost: a shard written with uniform rows, read back inside a batch whose *other* rows are ragged. Nothing diverged, and because a digest describes only its own row, the mixed batch needs no - special handling — the rows this process wrote are compared, the rows it - did not are counted unverified.""" + special handling — every row is compared against the reading it arrived + with, whichever process wrote it.""" inner = _JaggedEcho() client = _client(inner, verify_tensor_hash=True) ids = _ids(4) @@ -1175,15 +1175,16 @@ def test_uniform_write_read_back_inside_a_ragged_batch_is_clean(): ) # a later writer adds rows of a different length to the same partition other = _ids(2, "v") - inner.rows[("p", other[0])] = {"ids": torch.randint(0, 32000, (3,))} - inner.rows[("p", other[1])] = {"ids": torch.randint(0, 32000, (9,))} + MetricsDataPlaneClient(inner, verify_tensor_hash=True).put_samples( + sample_ids=other, partition_id="p", fields=_jagged_ids([3, 9], seed=1) + ) client.get_samples(sample_ids=ids + other, partition_id="p", select_fields=["ids"]) hv = client.snapshot()["hash_verify"] assert hv["mismatches"] == 0, "no row changed; nothing diverged" assert hv["fields_skipped"] == 0, "and every field stayed comparable" - assert hv["rows_checked"] == 4, "the four this process wrote" - assert hv["rows_unverified"] == 2, "the two it did not" + assert hv["rows_checked"] == 6, "including the two this process did not write" + assert hv["rows_unverified"] == 0 client.close() From ea28e41d3b986e4c9f8e011168b938a47bb2926e Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 10 Sep 2026 12:50:09 -0700 Subject: [PATCH 66/70] perf(data-plane): fold the field digests on tensors, and say so in the README _field_digests ran n_leaves * n_rows interpreted iterations on every put and every get, each one a Python call to wrap a uint64 into int64 range. int64 tensor arithmetic wraps two's-complement, which is the same modular fold, so the loop becomes one mul_/add_ per leaf and _as_i64 goes away. The stamp side then has its column already; the check side pairs its fields once instead of allocating a fresh [0] * n_rows default per field per row. README said "Only rows this process wrote can be checked", which is what the mirror columns invert. Documents the mirror, the per-top-level-field fold, the doubled declared field list, and the one case that still abstains. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 36 +++++++++++---- nemo_rl/data_plane/observability.py | 69 +++++++++++++++++++---------- 2 files changed, 74 insertions(+), 31 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 6d2fabbf29e..5685cc89ac3 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -656,6 +656,22 @@ on every put and re-checks it on every get, so a tensor that changes between wire-in and wire-out is reported (`hash/mismatches`) instead of being trained on silently. +**The reading travels with the row.** Each field is mirrored by a +`_hash` column -- one `int64` per row, written by the same put and +declared alongside the field by `register_partition`, which is why the +partition's field list is twice what the caller passed. The reader fetches +the mirror with the field, re-folds, compares, and strips the mirror before +the caller sees it. Holding the reading in the putting process instead would +only ever verify a same-process round trip, and the transfer worth checking +is not one: the rollout actor writes what the policy workers read. + +A mirror is per *top-level field*, not per leaf, because `select_fields` +names top-level fields -- a multimodal `images` reduces its leaves to a +single `images_hash`, folded in sorted leaf order with `* 31 +` so two +identical leaves cannot cancel. A column of `0` is the writer saying it could +not fold that field; the reader counts those rows `hash/rows_unverified` +rather than comparing against it. + One granularity: every row carries its own digest, formed from two parts. ``` @@ -671,9 +687,10 @@ One granularity: every row carries its own digest, formed from two parts. | the seed's shape | length (a zero pad or a truncation) and trailing-dim layout | yes | | — | a permutation *within* one row | **no** — see below | -The shape never travels and is never compared: only one integer per row is -stored. A shape change makes the seed differ, which makes the digest differ, -which surfaces as an ordinary mismatch. +The shape never travels and is never compared: one integer per row per field +is stored, and that is the whole reading. A shape change makes the seed +differ, which makes the digest differ, which surfaces as an ordinary +mismatch. The seed's shape is the *row's*, not the leaf's, and both layouts must agree on it — a dense `(N, L, D)` and the jagged form whose values are `(total, D)` @@ -731,11 +748,14 @@ its own digest. Known limits, measured rather than assumed: alarms this check has produced had exactly that shape, and both were its own bookkeeping. Per-sample lines carry the row index and the row length so the next one is adjudicable from a single log line. -- Only rows this process wrote can be checked. A consumer-side client - reports them under `hash/rows_unverified` rather than counting them - clean, and `hash/fields_skipped` reports any leaf it could not compare - — watch that one, since a guard that quietly stops covering a field still - reports zero mismatches. +- Rows written before the guard was switched on carry no mirror, and a read + whose batch contains one falls back to a plain fetch and abstains on the + whole batch — `hash/rows_unverified` and `hash/guard_failures` both move. + Within a run every writer shares one `verify_tensor_hash`, so this is the + resume-across-a-config-change case, not a steady-state one. +- `hash/fields_skipped` reports any leaf the fold could not attribute per + row — watch that one, since a guard that quietly stops covering a field + still reports zero mismatches. Backend choice: - **`simple`** — ZMQ-backed; lowest setup overhead. Default for tests diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 8f5a56c5533..8c0ca8f3da1 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -60,6 +60,7 @@ class DataPlaneEvent(TypedDict): status: EventStatus +import numpy as np import torch from tensordict import NonTensorData, NonTensorStack, TensorDict, TensorDictBase @@ -127,7 +128,6 @@ def _comm_volume(by_op: dict[str, Any]) -> dict[str, int]: # trip; the rollout actor writes what the policy workers read, and that read # is the one worth checking. _HASH_SUFFIX = "_hash" -_U64 = 1 << 64 def _hash_field(name: str) -> str: @@ -145,10 +145,21 @@ def _with_mirrors(fields: Sequence[str]) -> list[str]: return [*plain, *(_hash_field(f) for f in plain)] -def _as_i64(value: int) -> int: - """Wrap a uint64 digest into the signed range ``torch.int64`` accepts.""" - value &= _U64 - 1 - return value - _U64 if value >= _U64 >> 1 else value +def _hash_field(name: str) -> str: + return f"{name}{_HASH_SUFFIX}" + + +def _with_mirrors(fields: Sequence[str]) -> list[str]: + """``fields`` followed by one mirror column each. + + Idempotent: ``meta.fields`` comes back from ``put_samples`` already + carrying the mirrors, and suffixing those would ask for + ``tokens_hash_hash``. + """ + plain = [f for f in fields if not f.endswith(_HASH_SUFFIX)] + return [*plain, *(_hash_field(f) for f in plain)] + + # Rows a client may write between reconciliations of its live-key accounting # against the partition. One metadata call per this many rows put, so a client @@ -255,7 +266,7 @@ def _leaf_digests( def _field_digests( leaf_digests: dict[str, list[int]], n_rows: int -) -> dict[str, list[int]]: +) -> dict[str, torch.Tensor]: """Leaf digests folded to one digest per *top-level* field. ``select_fields`` names top-level fields, so the mirror has to be per @@ -266,12 +277,20 @@ def _field_digests( Sorted leaf order because dict order need not survive a round trip, and ``* 31 +`` rather than an XOR so two identical leaves do not cancel -- the defect the row fold already has, which must not be repeated here. + + Folded on tensors, not in Python: this runs on every put and every get, + and a row loop per leaf costs ``n_leaves * n_rows`` interpreted + iterations there. ``int64`` arithmetic wraps two's-complement, which is + the same modular fold the scalar form spelled out. """ - out: dict[str, list[int]] = {} + out: dict[str, torch.Tensor] = {} for name, per_row in sorted(leaf_digests.items()): - acc = out.setdefault(name.split(".", 1)[0], [0] * n_rows) - for row in range(n_rows): - acc[row] = _as_i64(acc[row] * 31 + per_row[row]) + # via uint64: a digest is unsigned and does not fit int64 directly. + # ``view`` reinterprets the same bits, which is the wrap we want. + column = torch.from_numpy(np.array(per_row, dtype=np.uint64).view(np.int64)) + top = name.split(".", 1)[0] + acc = out.get(top) + out[top] = column if acc is None else acc.mul_(31).add_(column) return out @@ -1485,10 +1504,9 @@ def _stamp_hashes_impl( # could not attribute per row -- those carry 0, which the reader takes # as "no reading". A column that is sometimes absent would make the # reader's fetch fail on a partition it has no business failing on. - unfolded = [0] * len(sample_ids) + unfolded = torch.zeros(len(sample_ids), dtype=torch.int64) for name in fields.keys(): - column = digests.get(name, unfolded) - stamped[_hash_field(name)] = torch.tensor(column, dtype=torch.int64) + stamped[_hash_field(name)] = digests.get(name, unfolded) self._stats.hash_verify.rows_recorded += len(sample_ids) return stamped @@ -1516,30 +1534,35 @@ def _check_hashes_impl( if isinstance(key, str) and key.endswith(_HASH_SUFFIX): expected_by_field[key[: -len(_HASH_SUFFIX)]] = out.get(key).tolist() del out[key] - digests = _field_digests( - self._row_fingerprints(out, sample_ids), len(sample_ids) - ) + digests = { + name: column.tolist() + for name, column in _field_digests( + self._row_fingerprints(out, sample_ids), len(sample_ids) + ).items() + } if not digests: return + # Paired once, not per row: the miss default used to be a fresh + # ``[0] * n_rows`` evaluated ``n_rows * n_fields`` times. + pairs = [ + (name, per_row, expected_by_field[name]) + for name, per_row in digests.items() + if name in expected_by_field + ] stats = self._stats.hash_verify for row, sample_id in enumerate(sample_ids): # ``0`` is the writer saying it could not fold that field, so it is # an abstention rather than a reading. A real digest of 0 is # possible and goes unchecked; at one row in 2^64 that is cheaper # than a false alarm on every asymmetric fold. - comparable = [ - (name, per_row) - for name, per_row in digests.items() - if expected_by_field.get(name, [0] * len(sample_ids))[row] != 0 - ] + comparable = [(n, d, e[row]) for n, d, e in pairs if e[row] != 0] if not comparable: # Written without the mirror: a put that predates the guard, # or a field the fold could not attribute per row. stats.rows_unverified += 1 continue stats.rows_checked += 1 - for name, per_row in comparable: - expected = expected_by_field[name][row] + for name, per_row, expected in comparable: if expected == per_row[row]: continue stats.mismatches += 1 From 0f6e01491e7efa94a0b86e32986821c13a0371f7 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 10 Sep 2026 14:43:17 -0700 Subject: [PATCH 67/70] test(data-plane): assert the wire guard looked, not just that it found nothing `mismatches == 0` also holds when nothing was compared: a row with no wire-in reading increments rows_unverified and the assertion passes green. Now that the reading travels with the row every reader compares, so rows_checked is non-zero on any run where the mirror is arriving -- and zero is the signal that it is not. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh | 6 ++++-- .../llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh | 6 ++++-- .../llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh | 6 ++++-- ...1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh | 6 ++++-- ...instruct-2n8g-async-1off-single-controller-streaming2.sh | 3 ++- ...a3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh | 6 ++++-- ...lama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh | 6 ++++-- .../grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh | 6 ++++-- ...ruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh | 6 ++++-- ...ruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh | 6 ++++-- .../grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh | 6 ++++-- ...a3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh | 6 ++++-- .../llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh | 6 ++++-- .../llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh | 6 ++++-- .../grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh | 6 ++++-- ...qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh | 6 ++++-- ...th-1.5b-instruct-1n8g-megatron-single-controller-sync.sh | 3 ++- .../llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh | 6 ++++-- ...-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh | 3 ++- .../llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh | 6 ++++-- .../llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh | 6 ++++-- .../mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh | 3 ++- ...uetp2sp-dynbatch-noncolocated-async-single-controller.sh | 3 ++- ...en2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh | 6 ++++-- 24 files changed, 86 insertions(+), 43 deletions(-) diff --git a/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh b/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh index 09e72697508..5cd27d60f95 100755 --- a/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh +++ b/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh b/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh index d6ba1702213..8fb92611429 100755 --- a/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh +++ b/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh b/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh index 29385fc4474..35d72532817 100755 --- a/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh +++ b/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh index b01ada20907..1fee1c46b9c 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh index b9cd68392ce..586eacc298c 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh @@ -38,7 +38,8 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma 'data["train/token_mult_prob_error"]["10"] < 1.1' \ 'mean(data["train/grad_norm"], 2, 0) > 0.06' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh index 616e6ccd04c..c1ad35a88a9 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh index 42ea0a86254..22c46f8952c 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh index e00189cfb7e..cb68f338416 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh index fca6386c352..085d69c3087 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh index 5776c921194..e8691c0ed82 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh index 21356eef73c..f00c5646e90 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh index d154f54957b..1b37bdb5df9 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh b/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh index fdcd5a14607..49195dafe22 100755 --- a/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh +++ b/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh index c9452ebf21a..28cdd11f606 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh @@ -17,7 +17,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,4 +27,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh index 3b9e13395ad..0855a325e52 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh @@ -17,7 +17,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,4 +27,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh index a51bf254518..113b00a2002 100755 --- a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh index 56775fa0b44..a1256a2dca3 100755 --- a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh @@ -38,7 +38,8 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma 'data["train/token_mult_prob_error"]["450"] < 1.1' \ 'mean(data["timing/train/total_step_time"], 2) < 25' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh b/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh index 4b958d7595b..c83a05ddc6c 100755 --- a/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh @@ -17,7 +17,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,4 +27,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh index b8410b81755..0a71271669a 100755 --- a/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh @@ -50,7 +50,8 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.02' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' uv run tools/check_r3_trace.py "$NRL_R3_TRACE_DIR" \ --require-forward-verify \ diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh index eb85c9c52be..0c95b157110 100755 --- a/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh @@ -34,11 +34,13 @@ uv run examples/run_grpo.py \ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then - # The wire guard only counts; assert it found nothing. + # The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.02' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' rm -rf "$CKPT_DIR" fi diff --git a/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh b/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh index 7ad99cc3df5..4887ae4168c 100755 --- a/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh index 42d47c94ac1..dacdfed8c12 100755 --- a/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh +++ b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh @@ -57,7 +57,8 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/token_mult_prob_error") | . 'max(data["train/on_policy_distillation/teacher_model_unique"]) == 1' \ 'max(data["train/on_policy_distillation/adv_std"]) < 0.01' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' rm -rf "$CKPT_DIR" fi diff --git a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh index 7cc5f9f37ec..e482ffa24fa 100755 --- a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh +++ b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh @@ -48,7 +48,8 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma 'mean(data["train/critic/explained_var"], range_start=-10) > 0.5' \ 'mean(data["train/reward"], range_start=-10) > 0.75' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh b/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh index 10a9cdb78ed..1d37a51b98d 100755 --- a/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh +++ b/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh @@ -16,7 +16,8 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing. +# The wire guard only counts; assert it found nothing -- and that it +# looked, since mismatches==0 also holds when nothing was compared. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -25,4 +26,5 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' + 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' From 84986df418589a20e1ad9358929bf6dfdc626d30 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 10 Sep 2026 14:52:48 -0700 Subject: [PATCH 68/70] test(data-plane): gate on what the guard did, not on what it did not Three counters for one question. guard_failures==0 says the check never raised, which is an absence; rows_checked>0 says it compared something, which is the property the gate exists to hold. A guard that stops working stops comparing, so the positive form catches it -- and keeps meaning the right thing as the guard changes. Leaves two clauses: it looked, and everything it compared agreed. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- .../test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh | 6 +++--- .../llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh | 6 +++--- .../llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh | 6 +++--- ...8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh | 6 +++--- ...struct-2n8g-async-1off-single-controller-streaming2.sh | 3 +-- ....1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh | 6 +++--- ...ma3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh | 6 +++--- ...rpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh | 6 +++--- ...ct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh | 6 +++--- ...ct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh | 6 +++--- .../grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh | 6 +++--- ....2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh | 6 +++--- .../llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh | 6 +++--- .../llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh | 6 +++--- .../grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh | 6 +++--- ...en2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh | 6 +++--- ...-1.5b-instruct-1n8g-megatron-single-controller-sync.sh | 3 +-- .../grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh | 6 +++--- ...0ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh | 3 +-- .../grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh | 8 ++++---- .../llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh | 6 +++--- ...opd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh | 3 +-- ...tp2sp-dynbatch-noncolocated-async-single-controller.sh | 3 +-- ...2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh | 6 +++--- 24 files changed, 63 insertions(+), 68 deletions(-) diff --git a/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh b/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh index 5cd27d60f95..dc4247d6c98 100755 --- a/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh +++ b/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh b/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh index 8fb92611429..5d5fac01ef7 100755 --- a/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh +++ b/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh b/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh index 35d72532817..eaf74dfa87e 100755 --- a/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh +++ b/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh index 1fee1c46b9c..7919a50e05f 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh index 586eacc298c..becb8da5761 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh @@ -38,8 +38,7 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma 'data["train/token_mult_prob_error"]["10"] < 1.1' \ 'mean(data["train/grad_norm"], 2, 0) > 0.06' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh index c1ad35a88a9..e8358a0d7d5 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh index 22c46f8952c..d1bbccc953a 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh index cb68f338416..e3f175487b9 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh index 085d69c3087..c22c0647475 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh index e8691c0ed82..44177963141 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh index f00c5646e90..e1ea4fc16a6 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh index 1b37bdb5df9..ecae045d435 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh b/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh index 49195dafe22..2a0022fdb27 100755 --- a/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh +++ b/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh index 28cdd11f606..2d15d68ffa8 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh @@ -17,8 +17,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -27,5 +28,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh index 0855a325e52..c58d88ef063 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh @@ -17,8 +17,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -27,5 +28,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh index 113b00a2002..3e542b11520 100755 --- a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh index a1256a2dca3..3dfa34b713b 100755 --- a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh @@ -38,8 +38,7 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma 'data["train/token_mult_prob_error"]["450"] < 1.1' \ 'mean(data["timing/train/total_step_time"], 2) < 25' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh b/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh index c83a05ddc6c..01ea06cec07 100755 --- a/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh @@ -17,8 +17,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -27,5 +28,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh index 0a71271669a..54576645d20 100755 --- a/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh @@ -50,8 +50,7 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.02' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' uv run tools/check_r3_trace.py "$NRL_R3_TRACE_DIR" \ --require-forward-verify \ diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh index 0c95b157110..e2fa718460d 100755 --- a/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh @@ -34,13 +34,13 @@ uv run examples/run_grpo.py \ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then - # The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. + # The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.02' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' rm -rf "$CKPT_DIR" fi diff --git a/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh b/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh index 4887ae4168c..af2ff43e981 100755 --- a/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' diff --git a/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh index dacdfed8c12..60aaeb7416f 100755 --- a/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh +++ b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh @@ -57,8 +57,7 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/token_mult_prob_error") | . 'max(data["train/on_policy_distillation/teacher_model_unique"]) == 1' \ 'max(data["train/on_policy_distillation/adv_std"]) < 0.01' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' rm -rf "$CKPT_DIR" fi diff --git a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh index e482ffa24fa..27fcf6b9c04 100755 --- a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh +++ b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh @@ -48,8 +48,7 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma 'mean(data["train/critic/explained_var"], range_start=-10) > 0.5' \ 'mean(data["train/reward"], range_start=-10) > 0.75' \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' + 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh b/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh index 1d37a51b98d..06d45449b67 100755 --- a/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh +++ b/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh @@ -16,8 +16,9 @@ source "$SCRIPT_DIR/common-tq.env" export EXP_NAME="$TQ_EXP_NAME" bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" -# The wire guard only counts; assert it found nothing -- and that it -# looked, since mismatches==0 also holds when nothing was compared. +# The wire guard only counts, so assert it looked and agreed. rows_checked +# is the load-bearing one: mismatches==0 also holds when nothing was compared, +# and a guard that stops working stops comparing. # The delegated base runs in a subshell, so common.env's TEST_DRYRUN exit # does not reach here. Skip explicitly, or the dryrun check in # tests/unit/test_recipes_and_test_suites.py fails on a missing metrics.json. @@ -26,5 +27,4 @@ if [[ -n "${TEST_DRYRUN:-}" ]]; then exit 0; fi cd "$SCRIPT_DIR/../../.." uv run tests/check_metrics.py "$SCRIPT_DIR/$TQ_EXP_NAME/metrics.json" \ 'max(data.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 0' \ - 'max(data.get("data_plane/cluster/step/hash/guard_failures", data.get("data_plane/driver/step/hash/guard_failures", {}))) == 0' \ 'max(data.get("data_plane/cluster/step/hash/rows_checked", data.get("data_plane/driver/step/hash/rows_checked", {}))) > 0' From d1cd96fcd5687b4a4fe4e66ec59e626f061f8ae5 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 10 Sep 2026 16:10:34 -0700 Subject: [PATCH 69/70] docs(data-plane): frac_of_step is carried work, not exclusive time Transfers overlap compute on the async and SC paths, so 7% is not 7% of the step the data plane would give back. Same caveat the section already carries for wall_ms, attached to the metric readers reach for first. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/README.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 5685cc89ac3..74d3716eb3c 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -499,6 +499,13 @@ step/percent_of_dataplane/by_op/put 42.1 within it, put is the `by_op` sums to 100 by construction. +**7% is not 7% of the step spent exclusively in the data plane.** The +numerator is data-plane wall time and the denominator is the step's, but on +the async and single-controller paths a transfer overlaps compute -- the same +reason `wall_ms` is aggregate process-time rather than elapsed time (below). +Read `frac_of_step` as "how much data-plane work a step carries", not as time +the step would get back if the data plane were free. + `volume_mb` counts *transfers*, not data size, and two things follow from that. A byte written and later read is counted on both sides. And every reporting process is summed, so four ranks each fetching their own shard From 51cdf358d3352ac559979d5cf962c564ecbb5391 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 11 Sep 2026 00:23:02 -0700 Subject: [PATCH 70/70] test(data-plane): the guard is a test flag, not a recipe default Per review: a recipe is what a user copies, and it should not ship a debug check that re-reads every tensor on both sides. All 24 recipes that set verify_tensor_hash now set it false; the 24 suites that gate on step/hash/rows_checked ask for it on the command line instead, so CI coverage is unchanged and the flag lives with the thing that needs it. tq wrappers pass it through to the base recipe script they delegate to; the single-controller suites that run run_grpo.py directly add it to their own argument list. Signed-off-by: Zhiyu Li Co-Authored-By: Claude Opus 5 Signed-off-by: Zhiyu Li --- .../recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml | 2 +- .../llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml | 2 +- .../recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml | 2 +- ...1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml | 2 +- ...instruct-2n8g-async-1off-single-controller-streaming2.yaml | 2 +- ...a3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml | 2 +- ...lama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml | 2 +- .../grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml | 2 +- ...ruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml | 2 +- ...ruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml | 2 +- .../grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml | 2 +- ...a3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml | 2 +- .../llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml | 2 +- .../llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml | 2 +- .../grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml | 2 +- ...qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml | 2 +- ...th-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml | 2 +- .../llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml | 2 +- ...-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml | 2 +- .../llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml | 2 +- .../llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml | 2 +- .../mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml | 2 +- ...uetp2sp-dynbatch-noncolocated-async-single-controller.yaml | 2 +- ...en2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml | 2 +- tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh | 4 +++- .../llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh | 4 +++- .../test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh | 4 +++- ...3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh | 4 +++- ...b-instruct-2n8g-async-1off-single-controller-streaming2.sh | 1 + ...ama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh | 4 +++- ...-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh | 4 +++- .../grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh | 4 +++- ...struct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh | 4 +++- ...struct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh | 4 +++- .../llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh | 4 +++- ...ama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh | 4 +++- .../llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh | 4 +++- .../llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh | 4 +++- .../llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh | 4 +++- ...o-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh | 4 +++- ...math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh | 1 + .../llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh | 4 +++- ...n3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh | 1 + .../llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh | 1 + .../llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh | 4 +++- .../mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh | 1 + ...aluetp2sp-dynbatch-noncolocated-async-single-controller.sh | 1 + ...qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh | 4 +++- 48 files changed, 84 insertions(+), 42 deletions(-) diff --git a/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml b/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml index b2e7feb0ccd..6d93691bef4 100644 --- a/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-deepscaler-1.5b-8K.yaml data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml b/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml index 895ea4ed6ee..de824da361b 100644 --- a/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-gemma3-1b-it-1n8g-fsdp2tp1.yaml data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml b/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml index effaeefa9ef..afaf90e6239 100644 --- a/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-gspo-deepscaler-1.5b-8K.yaml data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml index 34e8e5af14b..4fabe7fc742 100644 --- a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3.yaml data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml index aea0573f1f9..697a958b8ee 100644 --- a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml @@ -21,7 +21,7 @@ checkpointing: data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false # SC async-RL runtime knobs. async_rl: diff --git a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml index 6e69c526fb7..e319f043a50 100644 --- a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated.yaml data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml index ed213188b2e..55473b7ad2a 100644 --- a/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml index 8965990c20d..25ddbbd2ebc 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3.yaml data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml index 4ed9a2f64a6..c40136b8a98 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml index bb67e52ae04..d59e06c86bc 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml index 09cad6a522c..64767df70ef 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-llama3.2-1b-instruct-1n8g-megatron.yaml data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml index 4e58796dd8c..f935737f5b2 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml b/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml index 52cd402799f..4caca1e28d9 100644 --- a/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-moonlight-16ba3b-4n8g-megatron.yaml data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml index 3171128693c..468b20c09ef 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml index be5f563960a..2184bb4568d 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.yaml data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml index 8cff067358c..806a21d8bf8 100644 --- a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.yaml @@ -2,4 +2,4 @@ defaults: grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3.yaml data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml index 9fe7b64dcbc..5c0140676b8 100644 --- a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml @@ -21,7 +21,7 @@ checkpointing: data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false # SC async-RL runtime knobs. async_rl: diff --git a/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml index 60096349373..67621528980 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml index 02481dc9015..9e4fe15b45c 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml @@ -19,7 +19,7 @@ checkpointing: data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false async_rl: sampler: diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml index 1fcdf3c2179..578d39c17a1 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml @@ -4,7 +4,7 @@ checkpointing: data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false logger: log_dir: logs/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple wandb: diff --git a/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml b/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml index f186d51c686..e68a4475ff5 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml b/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml index 555dadec318..749409cc89a 100644 --- a/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml +++ b/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml @@ -18,7 +18,7 @@ async_rl: data_plane: enabled: true observability: - verify_tensor_hash: true + verify_tensor_hash: false checkpointing: checkpoint_dir: results/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller diff --git a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml index eacf1f8d1d4..dc5f7d14130 100644 --- a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml +++ b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml @@ -56,7 +56,7 @@ data_plane: reuse_registered_buffers: true staging_buffer_size: 268435456 observability: - verify_tensor_hash: true + verify_tensor_hash: false # SC async-RL runtime knobs, replacing the nulled ppo.async_ppo block. async_rl: diff --git a/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml b/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml index aeef1bfdebe..86afef15c17 100644 --- a/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml +++ b/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.yaml @@ -3,4 +3,4 @@ data_plane: enabled: true backend: mooncake_cpu observability: - verify_tensor_hash: true + verify_tensor_hash: false diff --git a/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh b/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh index dc4247d6c98..158be6ad9d2 100755 --- a/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh +++ b/tests/test_suites/llm/grpo-deepscaler-1.5b-8K-tq_simple.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh b/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh index 5d5fac01ef7..ca0095fff32 100755 --- a/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh +++ b/tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1-tq_simple.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh b/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh index eaf74dfa87e..292e9e33c27 100755 --- a/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh +++ b/tests/test_suites/llm/grpo-gspo-deepscaler-1.5b-8K-tq_simple.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh index 7919a50e05f..7d7265d978a 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3-tq_simple.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh index becb8da5761..dbad9563fa9 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh @@ -25,6 +25,7 @@ uv run examples/run_grpo_single_controller.py \ logger.tensorboard_enabled=True \ checkpointing.enabled=False \ checkpointing.checkpoint_dir=$CKPT_DIR \ + data_plane.observability.verify_tensor_hash=True \ $@ \ 2>&1 | tee $RUN_LOG diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh index e8358a0d7d5..344860ed260 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated-tq_simple.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh index d1bbccc953a..4ceccc8b5b8 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-megatron-fp8-e2e-tq_mooncake.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh index e3f175487b9..7f01a485bee 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh index c22c0647475..1a08e995822 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp2-temp0.8-topp0.9-topk50-tq_mooncake.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh index 44177963141..32e0d21dd21 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-temp0.8-topp0.9-topk50-tq_mooncake.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh index e1ea4fc16a6..bbb070466f6 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron-tq_simple.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh index ecae045d435..604d296653b 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation-tq_mooncake.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh b/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh index 2a0022fdb27..b76872775a2 100755 --- a/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh +++ b/tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-tq_simple.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh index 2d15d68ffa8..0175c7b9300 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-1n8g-fsdp2-tq_mooncake.v2.sh @@ -15,7 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh index c58d88ef063..66236cf08b7 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp-tq_simple.sh @@ -15,7 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh index 3e542b11520..03a3189d0c3 100755 --- a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3-tq_simple.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh index 3dfa34b713b..e5327e08891 100755 --- a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh @@ -25,6 +25,7 @@ uv run examples/run_grpo_single_controller.py \ logger.tensorboard_enabled=True \ checkpointing.enabled=False \ checkpointing.checkpoint_dir=$CKPT_DIR \ + data_plane.observability.verify_tensor_hash=True \ $@ \ 2>&1 | tee $RUN_LOG diff --git a/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh b/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh index 01ea06cec07..aefa87ea03e 100755 --- a/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3-tq_mooncake.sh @@ -15,7 +15,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh index 54576645d20..903167659cb 100755 --- a/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.sh @@ -41,6 +41,7 @@ uv run examples/run_grpo_single_controller.py \ logger.tensorboard_enabled=True \ checkpointing.enabled=False \ checkpointing.checkpoint_dir=$CKPT_DIR \ + data_plane.observability.verify_tensor_hash=True \ $@ \ 2>&1 | tee $RUN_LOG diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh index e2fa718460d..ec37a92fad2 100755 --- a/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh @@ -28,6 +28,7 @@ uv run examples/run_grpo.py \ logger.tensorboard_enabled=True \ checkpointing.enabled=True \ checkpointing.checkpoint_dir=$CKPT_DIR \ + data_plane.observability.verify_tensor_hash=True \ $@ \ 2>&1 | tee $RUN_LOG diff --git a/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh b/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh index af2ff43e981..7b952a5aefe 100755 --- a/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh +++ b/tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora-tq_mooncake.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared, diff --git a/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh index 60aaeb7416f..c6c5c8e8cda 100755 --- a/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh +++ b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh @@ -43,6 +43,7 @@ uv run examples/run_grpo_single_controller.py \ logger.tensorboard_enabled=True \ checkpointing.enabled=False \ checkpointing.checkpoint_dir=$CKPT_DIR \ + data_plane.observability.verify_tensor_hash=True \ "$@" \ 2>&1 | tee $RUN_LOG diff --git a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh index 27fcf6b9c04..06748b4dee7 100755 --- a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh +++ b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh @@ -25,6 +25,7 @@ uv run examples/run_grpo_single_controller.py \ logger.tensorboard_enabled=True \ checkpointing.enabled=True \ checkpointing.checkpoint_dir=$CKPT_DIR \ + data_plane.observability.verify_tensor_hash=True \ $@ \ 2>&1 | tee $RUN_LOG diff --git a/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh b/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh index 06d45449b67..ce3fe1650a1 100755 --- a/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh +++ b/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_mooncake.sh @@ -14,7 +14,9 @@ source "$SCRIPT_DIR/common-tq.env" # Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). # The matching TQ YAML inherits from .yaml and turns on data_plane. export EXP_NAME="$TQ_EXP_NAME" -bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" \ + data_plane.observability.verify_tensor_hash=True \ + "$@" # The wire guard only counts, so assert it looked and agreed. rows_checked # is the load-bearing one: mismatches==0 also holds when nothing was compared,