diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index e99d9bf2cda..184e62d014d 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -583,8 +583,9 @@ 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: true # per-op timing/volume; cost is below measurement noise + 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/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..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 @@ -1,3 +1,5 @@ defaults: grpo-deepscaler-1.5b-8K.yaml data_plane: enabled: true + observability: + 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 7e693ba250e..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 @@ -1,3 +1,5 @@ defaults: grpo-gemma3-1b-it-1n8g-fsdp2tp1.yaml data_plane: enabled: true + observability: + 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 80d795c7215..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 @@ -1,3 +1,5 @@ defaults: grpo-gspo-deepscaler-1.5b-8K.yaml data_plane: enabled: true + observability: + 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 8bf6005dbf8..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 @@ -1,3 +1,5 @@ defaults: grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3.yaml data_plane: enabled: true + observability: + 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 f751736341e..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 @@ -20,6 +20,8 @@ checkpointing: # TransferQueue data plane is mandatory for the SingleController path. data_plane: enabled: true + observability: + 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 f457fecd0aa..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 @@ -1,3 +1,5 @@ defaults: grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated.yaml data_plane: enabled: true + observability: + 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 5313f25dd42..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 @@ -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: 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 89cd90aceb0..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 @@ -1,3 +1,5 @@ defaults: grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3.yaml data_plane: enabled: true + observability: + 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 15ebb1943f7..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 @@ -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: 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 9406c66c578..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 @@ -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: 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 dbca800e5ca..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 @@ -1,3 +1,5 @@ defaults: grpo-llama3.2-1b-instruct-1n8g-megatron.yaml data_plane: enabled: true + observability: + 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 10f4ba92a1d..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 @@ -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: 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 26729ccd7c2..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 @@ -1,3 +1,5 @@ defaults: grpo-moonlight-16ba3b-4n8g-megatron.yaml data_plane: enabled: true + observability: + 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 6fbd7218c16..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 @@ -2,3 +2,5 @@ defaults: grpo-nanov3-30BA3B-1n8g-fsdp2.v2.yaml data_plane: enabled: true backend: mooncake_cpu + observability: + 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 86be2c983f3..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 @@ -1,3 +1,5 @@ defaults: grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.yaml data_plane: enabled: true + observability: + 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 9d25544bcd7..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 @@ -1,3 +1,5 @@ defaults: grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3.yaml data_plane: enabled: true + observability: + 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 b3fbaf26e53..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 @@ -20,6 +20,8 @@ checkpointing: # TransferQueue data plane is mandatory for the SingleController path. data_plane: enabled: true + observability: + 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 08f148f3709..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 @@ -2,3 +2,5 @@ defaults: grpo-qwen3-1.7b-1n8g-megatron-eagle3.yaml data_plane: enabled: true backend: mooncake_cpu + observability: + 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 7ca7aed34bf..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 @@ -18,6 +18,8 @@ checkpointing: # TransferQueue is mandatory for the SingleController path. data_plane: enabled: true + observability: + 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 f48450c8ef2..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 @@ -3,6 +3,8 @@ checkpointing: checkpoint_dir: results/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple data_plane: enabled: true + observability: + 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 952f192a362..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 @@ -2,3 +2,5 @@ defaults: grpo-qwen3-8B-base-1n8g-fsdp2-lora.yaml data_plane: enabled: true backend: mooncake_cpu + observability: + 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 fe7bdf3501f..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 @@ -17,6 +17,8 @@ async_rl: data_plane: enabled: true + observability: + 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 02e7492f6b8..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 @@ -55,6 +55,8 @@ data_plane: local_buffer_size: 4294967296 reuse_registered_buffers: true staging_buffer_size: 268435456 + observability: + 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 839a96bb56b..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 @@ -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: false diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index a602b569796..b0f9ddf6e6c 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 @@ -78,6 +79,13 @@ 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, + cluster_step_metrics, + 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 from nemo_rl.environments.interfaces import EnvironmentInterface @@ -381,6 +389,63 @@ def _compute_seq_logprob_error_metrics( return masking_data["sample_mask"], seq_logprob_error_metrics +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. + """ + with metrics_never_fail_the_step(step): + _log_data_plane_metrics_impl(policy, logger, step, total_step_time) + + +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. + + 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): + 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: + merged = merge_snapshots(snapshots) + # 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 + 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) + log_step_metrics(logger, metrics, step, "driver") + + def grpo_train_sync( policy: ColocatablePolicyInterface, policy_generation: GenerationInterface, @@ -1375,6 +1440,10 @@ 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, diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 36dd5ffd103..2c9c6104396 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, + log_step_metrics, + metrics_never_fail_the_step, +) from nemo_rl.data_plane.schema import ( DP_CALIB_INPUT_FIELDS, DP_TRAIN_FIELDS, @@ -1519,6 +1524,38 @@ 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``. + """ + with metrics_never_fail_the_step(self._train_steps): + self._log_data_plane_metrics_impl(total_step_time) + + 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`` + 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`` 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 + + metrics = self._dp_client.get_step_metrics(total_step_time) + log_step_metrics(self._logger, metrics, self._train_steps, "driver") + @staticmethod def _group_ids_from_meta(meta: KVBatchMeta) -> list[str]: """Return stable prompt-group IDs in canonical sample order.""" @@ -2849,6 +2886,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, diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index fb70eb79c5d..87bf96a61e8 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,359 @@ 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: true # 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. + +**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, 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 +throughput — is recomputed from the merged totals, never +averaged across ranks (averaging per-rank percentiles does not give a +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 `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/percent_of_dataplane/by_op/{put,get,clear,register}` | which call is expensive? | +| `step/wall_s`, `step/comm_volume_mb` | how much time and traffic | +| `step/volume_mb/by_op/{get,put}` | which direction that traffic went | +| `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 +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. + +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 +``` + +`by_op` sums to 100 by construction. + +**`frac_of_step` and `wall_s` report the busiest process, not the sum.** +The denominator is one step's wall clock, so the numerator has to be wall time +too. These processes run concurrently inside that window: summing them +exceeds the step itself (measured 1.054 across ten processes), and averaging +them dilutes the busiest process with idle ones. The step waits on one +process, so the max is what was exposed. It is a lower bound -- if the +driver's ops are serial with the workers' fetches, the real figure is +`driver + max(workers)` -- but it is a bound in the right direction, where the +sum was neither. `by_op` percentages still sum, because "where did the time +go" is a process-time question. + +`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 the per-op `wall_ms` is summed over processes that ran +concurrently, so the **`by_op` percentages** are shares 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" -- which is +what `frac_of_step` answers, and why it takes a max instead. + +**A per-op breakdown table** carries the detail, under +`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 | 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`: 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 +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, `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 +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: + +| 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/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 +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 +step. + +`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 + 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:** 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 +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 +`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. + +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, +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 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. + +**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. + +``` + 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 | +|---|---|---| +| 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: 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)` +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. + +Measured cost, from a 15-step A/B on one recipe (Llama-3.2-1B, 1 node x 4 +GB300, TQ `simple`, 24 MB/step) differing only in `verify_tensor_hash`: + +| | guard on | guard off | +|---|---|---| +| `step/self/overhead_ms` | ~123 ms | ~4.4 ms | +| `step/wall_s` | ~1.35 s | ~1.05 s | +| `total_step_time` (mean, steps 8-15) | 16.8 s | 15.1 s (see below) | +| `step/hash/rows_checked` | 2560 | — | + +The wrapper's own accounting is ~4 ms; the guard is essentially all of the +~119 ms difference, and the data plane's own wall time rose 297 ms. + +The end-to-end step moved 1.67 s, which those two do **not** explain, and it +should not be read as the guard's cost. `comm_volume_mb` was unchanged +(24.42 vs 24.58 MB), step time varied 14.3-20.3 s *within* each run, and at +step 10 the guard-off run was the slower of the two. The two runs were also on +different nodes. Against that spread a 1.67 s mean difference is not +attributable; a same-node A/B is needed before any end-to-end figure is +quoted. + +**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, 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 row index and the row length + so the next one is adjudicable from a single log line. +- 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/__init__.py b/nemo_rl/data_plane/__init__.py index c97346ed4d2..46c5a888cb3 100644 --- a/nemo_rl/data_plane/__init__.py +++ b/nemo_rl/data_plane/__init__.py @@ -27,7 +27,13 @@ KVBatchMeta, data_plane_supports_checkpointing, ) -from nemo_rl.data_plane.observability import MetricsDataPlaneClient, log_event +from nemo_rl.data_plane.observability import ( + MetricsDataPlaneClient, + breakdown_table, + cluster_step_metrics, + log_event, + merge_snapshots, +) __all__ = [ "DATA_PLANE_CHECKPOINT_SCHEMA_VERSION", @@ -36,7 +42,10 @@ "KVBatchMeta", "MetricsDataPlaneClient", "build_data_plane_client", + "breakdown_table", + "cluster_step_metrics", "data_plane_supports_checkpointing", "log_event", + "merge_snapshots", "materialize", ] diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index 692f6d6daaf..918e9bede97 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,28 @@ 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 + 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/factory.py b/nemo_rl/data_plane/factory.py index 9c076d64677..e9c510275db 100644 --- a/nemo_rl/data_plane/factory.py +++ b/nemo_rl/data_plane/factory.py @@ -179,12 +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(client, on_event=on_event) # type: ignore[bad-argument-type] + 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/interfaces.py b/nemo_rl/data_plane/interfaces.py index a50478295b7..75ba7b86aae 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -182,11 +182,22 @@ 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`` 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 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 d569740cf18..9b968c94646 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -23,12 +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). + +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: 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 import logging -from dataclasses import asdict, dataclass +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 typing import Any, Callable, Literal, TypedDict @@ -45,36 +60,1104 @@ class DataPlaneEvent(TypedDict): status: EventStatus +import numpy as np import torch -from tensordict import TensorDict +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__) -def _td_bytes(td: TensorDict | None) -> int: +# 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 _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 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" + + +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 _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 +# 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 +# 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). +# +# 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)) + + +# 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 _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)``. + + 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. + """ + # 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(int(fold) ^ seed) + return digests + + +def _field_digests( + leaf_digests: dict[str, list[int]], n_rows: 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 + 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. + + 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, torch.Tensor] = {} + for name, per_row in sorted(leaf_digests.items()): + # 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 + + +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) + + +def _tensor_bytes(v: torch.Tensor) -> int: + """Wire bytes of one tensor leaf, rectangular or nested. + + 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: + 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 _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`` 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. + """ + 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 _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 + + +# 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*. 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 + + +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. + + 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 + 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. + + ``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 + ``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): - v = td.get(k) - if not isinstance(v, torch.Tensor): - continue - t = v.values() if v.is_nested else v - total += t.numel() * t.element_size() + # pyrefly: ignore # bad-assignment + for _, v in td.items(include_nested=True, leaves_only=False): + if isinstance(v, torch.Tensor): + total += _tensor_bytes(v) + 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 + else: + total += _estimate_encoded_bytes(v, budget) return total +def _step_deltas(snap: dict[str, Any], prev: dict[str, Any]) -> dict[str, float]: + """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. + """ + + 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_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_s": _delta_s("pack_ms"), + "step/codec/unpack_s": _delta_s("unpack_ms"), + } + + +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 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 + 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"])) + out[op] = row + 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. + """ + # ``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) 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 + # 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]: + """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 + } + + +# 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. + + 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. + + ``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 + 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/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: + return {} + percent = { + f"step/percent_of_dataplane/by_op/{op}": 100.0 * r["wall_ms"] / total + for op, r in per_op.items() + } + return percent + + +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 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. + """ + n = sum(hist) + ceiling = max_ms if max_ms > 0 else float("inf") + 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: + """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["percent_of_total_ms"] = ( + 100.0 * wall_ms / total_wall_ms if total_wall_ms else 0.0 + ) + hist = stats["latency_hist"] + # 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 +# 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", + "pack_ms", + "unpack_ms", +) +_SNAPSHOT_MAX = ("max_bytes_per_key_seen", "last_put_bytes_per_key") +_OP_SUM = ( + "calls", + "errors", + "wall_ms", + "n_bytes", + "n_keys", +) +_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 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 + 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 _HASH_FIELDS} + 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) + # The busiest single process, kept beside the sum: these ran concurrently + # inside one step, so the sum is process-time and only the max is wall + # time the step could have waited on. ``frac_of_step`` differences it. + merged["max_process_wall_ms"] = max( + (s.get("total_wall_ms", 0.0) for s in snapshots), default=0.0 + ) + _derive_op_metrics(by_op, merged["total_wall_ms"]) + merged.update(_comm_volume(by_op)) + return merged + + +def cluster_step_metrics( + 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. + + 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 the whole bill for measuring: every + 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`. + 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. + """ + 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 + exposed_ms = snap.get("max_process_wall_ms", wall_ms / n_procs) - prev.get( + "max_process_wall_ms", 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(snap, prev) + metrics.update( + { + # The one metric that says whether optimising the data plane is + # worth anything: per-op percentages say where its time went, never + # whether it mattered against compute. The denominator is one + # step's wall clock, so the numerator has to be wall time too: + # processes run concurrently inside that window, so summing them + # exceeds 1 whenever they overlapped (measured 1.054 across ten + # processes) and averaging them dilutes the busiest process with + # idle ones. The busiest process is the one the step waits on, so + # the max is what was exposed -- a lower bound on the step's own + # serial cost, where the sum was an upper bound on nothing. + "step/frac_of_step": ( + exposed_ms / (step_time_s * 1e3) if step_time_s > 0 else 0.0 + ), + # Same reduction, same reason: the step waited on one process, + # not on all of them added together. ``_step_deltas`` summed it. + "step/wall_s": exposed_ms / 1e3, + "step/self/overhead_ms": overhead_ms, + "step/self/frac": overhead_ms / wall_ms if wall_ms > 0 else 0.0, + } + ) + metrics.update( + _hash_deltas(snap.get("hash_verify") or {}, prev.get("hash_verify") or {}) + ) + metrics.update(_op_series(snap["by_op"], prev.get("by_op", {}))) + 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 time does. +# The full dict is still returned, so the table and the series are derived +# from one computation and cannot disagree. +_HEADLINE = ( + "step/wall_s", + "step/frac_of_step", + "step/comm_volume_mb", + "now/bytes_outstanding_mb", + "now/n_processes", +) +_HEADLINE_PREFIXES = ( + "step/percent_of_dataplane/", + "step/volume_mb/", + "step/hash/", + "step/self/", + "step/codec/", +) + + +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 percentages, 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. +# ``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", + "wall_ms", + "mean_ms", + "max_ms", + "p50_ms", + "p90_ms", + "mb", +) + + +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. 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 + 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`. + """ + per_op: dict[str, dict[str, float]] = {} + for key, value in metrics.items(): + parts = key.split("/") + 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) == 4 + and parts[:2] == ["step", "by_op"] + and parts[3] in _BREAKDOWN_COLUMNS + ): + 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 + # 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) + ) + ] + 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) +@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 + # 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 + # 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 + # 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( + default_factory=lambda: [0] * (len(LATENCY_BUCKETS_MS) + 1) + ) + + +@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 + # 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 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 @@ -82,6 +1165,26 @@ 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 + # 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) class MetricsDataPlaneClient(DataPlaneClient): @@ -91,26 +1194,102 @@ 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 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 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]] = {} + # Live bytes and live keys per partition. Populated on successful + # ``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]] = {} + self._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 + # ``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 + # 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.""" + 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 + 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). + + 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. + """ + # 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(d) for d in self._bytes_by_partition.values() + len(k) for k in self._keys_by_partition.values() ) + _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). + 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 return out - 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()} + 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: ``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 + # 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 + # 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. @@ -118,21 +1297,34 @@ 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 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, {}) - for i, key in enumerate(keys): - share = per_key + (1 if i < remainder else 0) - partition_dict[key] = partition_dict.get(key, 0) + share + 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 + 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``. @@ -140,24 +1332,271 @@ 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. """ - partition_dict = self._bytes_by_partition.get(partition_id) - if partition_dict is None: + live = self._keys_by_partition.get(partition_id) + 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. if keys is None: - freed = sum(partition_dict.values()) - del self._bytes_by_partition[partition_id] + 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] + self._bytes_by_partition.pop(partition_id, None) else: - freed = 0 - for key in keys: - freed += partition_dict.pop(key, 0) - if not partition_dict: - del self._bytes_by_partition[partition_id] + freed = total * removed // (len(live) + removed) if removed else 0 + 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: # noqa: BLE001 - the put succeeded; retry next window + 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. + + 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( + self, + td: TensorDict | None, + sample_ids: list[str], + ) -> 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. + + Returns: + 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, 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 + # 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: + if v.layout != torch.jagged or v.offsets().numel() - 1 != n_rows: + stats.fields_skipped += 1 + continue + 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: (length, *tail) + elif v.shape[0] != n_rows: + stats.fields_skipped += 1 + continue + else: + # 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 + 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: + """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 _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: + 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) + 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 = torch.zeros(len(sample_ids), dtype=torch.int64) + for name in fields.keys(): + stamped[_hash_field(name)] = digests.get(name, unfolded) + 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.""" + 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. + + 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 + 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 = { + 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 = [(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, expected in comparable: + if expected == per_row[row]: + continue + stats.mismatches += 1 + if self._hash_mismatches_logged < _MAX_HASH_MISMATCH_LOGS: + self._hash_mismatches_logged += 1 + # 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 (row %d of %d)", + partition_id, + sample_id, + name, + expected, + per_row[row], + row, + len(sample_ids), + ) + def _run( self, op: str, @@ -208,24 +1647,49 @@ def _emit( t0: float, status: EventStatus, ) -> None: - event: DataPlaneEvent = { - "op": op, - "partition_id": partition_id, - "n_keys": int(n_keys), - "n_bytes": int(n_bytes), - "wall_ms": (monotonic() - t0) * 1000.0, - "status": status, - } - self._on_event(event) - if status == "ok": - self._stats.total_bytes += n_bytes - self._stats.total_keys += n_keys - self._stats.total_ops += 1 - 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 + 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. + 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. + 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 = 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 + 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 + 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, @@ -236,6 +1700,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 = _with_mirrors(fields) self._run( "register", partition_id, @@ -275,12 +1748,20 @@ def claim_meta( ) def get_data(self, meta, select_fields=None): - return self._run( + 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 = _with_mirrors(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: + 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): return self._run( @@ -290,38 +1771,66 @@ 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. - sample_ids_list = ( - sample_ids if isinstance(sample_ids, list) else list(sample_ids) - ) + 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) + self._bill_self(entered) return out def get_samples(self, sample_ids, partition_id, select_fields): - return self._run( - "get", - partition_id, - lambda: self._inner.get_samples( - sample_ids, + entered = monotonic() + sample_ids_list = _as_list(sample_ids) + fetch = list(select_fields) + if self._verify_tensor_hash: + fetch = _with_mirrors(select_fields) + def read(columns): + return self._run( + "get", partition_id, - select_fields=select_fields, - ), - n_keys=len(sample_ids), - ) + 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 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 + # 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) + return out def list_sample_ids(self, partition_id: str) -> list[str]: return self._run( @@ -331,11 +1840,8 @@ 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) - ) + 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( "clear", @@ -344,6 +1850,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/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index a9b7ab95932..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 @@ -384,6 +384,25 @@ 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. + + 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) + 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, meta: "KVBatchMeta", diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index ae674a17c53..9ac35ea87e4 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 @@ -97,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`. @@ -231,6 +235,38 @@ 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"): + # 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 + # ``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 + 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) ────────────────── def _with_route_fields( 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/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/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..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 aa3f053e87b..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 0c3f6184a52..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 c4635e9f219..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 56a639f74fb..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 @@ -36,7 +37,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/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 b0880e3c890..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 2bf6ad97548..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 4b0de6414f6..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 db93cbfab37..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 a67323016fd..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 14a128672fe..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 7c765c5bd22..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 606dccb7325..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 9ff4e574237..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,4 +15,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 56850083541..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,4 +15,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 9d905ef649a..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 fa5e77cc369..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 @@ -36,7 +37,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/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 121013805bb..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,4 +15,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 0508017c7a5..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 @@ -48,7 +49,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/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 490f00fb665..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,14 +28,20 @@ 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 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, 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' + '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/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 b6c3f072b70..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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 e33726d8652..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 @@ -55,7 +56,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/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 f1666482055..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 @@ -46,7 +47,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/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 5e07f0ef476..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,4 +14,19 @@ 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, +# 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. +# `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.get("data_plane/cluster/step/hash/mismatches", data.get("data_plane/driver/step/hash/mismatches", {}))) == 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/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 901f986b65f..281b6701508 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_s", "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/conftest.py b/tests/unit/data_plane/conftest.py index 0f7e2da3019..44f70455733 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_s``. + """ + 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 13a5f59a382..b7a0842fb4a 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -21,36 +21,194 @@ from __future__ import annotations +import logging +import random +from time import monotonic + import pytest import torch -from tensordict import TensorDict +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 MetricsDataPlaneClient +from nemo_rl.data_plane.observability import ( + _QUANTILES, + MetricsDataPlaneClient, + _estimate_encoded_bytes, + _hash_deltas, + _td_bytes, + breakdown_table, + cluster_step_metrics, + headline_series, + merge_snapshots, +) + +# ── 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. + """ + 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. + + 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 = { + "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 + ) -from ._rollout_shapes import make_rollout_batch + 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 @@ -59,9 +217,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", @@ -70,23 +225,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): @@ -106,44 +255,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() @@ -183,53 +314,1032 @@ 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"] +# ── 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_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)] + 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 it would silently inflate ``n_bytes``.""" + 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.""" + 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} + + 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) + 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))} + + assert _estimate_encoded_bytes(huge, [64]) < _estimate_encoded_bytes( + huge, [10_000_000] ) - assert len(seen) == 1 and seen[0]["op"] == "register" + 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_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).""" +@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") - inner = NoOpDataPlaneClient() - seen: list[dict] = [] - client = MetricsDataPlaneClient(inner, on_event=seen.append) + assert client.snapshot()["bytes_outstanding"] == 0 + client.close() - 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"], + +# ── 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. ``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 + 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" + ) + assert "p90_ms" not in client.snapshot()["by_op"]["put"] + 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 blank the next step.""" + client = _client(register=False) + _emit(client, [30.0]) + + 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_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() + + +# ── percentiles: clamping and sample gates ───────────────────────────── + + +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"] + + 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 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 + ) + + 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 ) - fields = TensorDict( + 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 + + +# ── 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_s" 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( { - "input_ids": batch["input_ids"], - "input_lengths": batch["input_lengths"], - "generation_logprobs": batch["generation_logprobs"], - }, - batch_size=[n], + "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, + "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 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("calls")] == 3 + + +def test_breakdown_table_is_empty_when_nothing_ran(): + assert breakdown_table({"step/wall_s": 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, + } + ) + + 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 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 = _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", "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, + "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 digest 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"] + # 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"] == 3 + client.close() + + +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"]) + + 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 = _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"]) + + hv = client.snapshot()["hash_verify"] + assert hv["rows_checked"] == 4 + assert hv["mismatches"] == 0 + client.close() + + +@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 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()) + + 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_checked"] == 4 + assert hv["rows_unverified"] == 0 + assert hv["mismatches"] == mismatches + writer.close() + + +@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 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=verify_tensor_hash) + 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 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() + ) + still_live = set(_ids(4, prefix="live")) | set(_ids(4, prefix="next")) + assert writer._keys_by_partition["p"] == still_live + 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 + 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 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 + 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. 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) + + 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"] + + 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 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]]) + + 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 + ) + assert as_fp32["x"] != as_bf16["x"] + client.close() + + +def test_hash_fingerprint_handles_float8(): + """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)}, + batch_size=[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 = _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) + on_device = client._row_fingerprints( + TensorDict({"x": values.cuda()}, batch_size=[2]), ids + ) + 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_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)] + 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"] == 0, "a shard is comparable now" + assert hv["rows_checked"] == 2, "and both its rows were actually checked" + client.close() + + +@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(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"]) + + hv = client.snapshot()["hash_verify"] + assert hv["mismatches"] == 2, "the two swapped rows, named individually" + assert hv["fields_skipped"] == 0 + client.close() + + +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) client.put_samples( - sample_ids=[f"u{i}" for i in range(n)], - partition_id="train", - fields=fields, + 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"]) - 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 + hv = client.snapshot()["hash_verify"] + 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 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 — 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) + 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 = _ids(2, "v") + 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"] == 6, "including the two this process did not write" + assert hv["rows_unverified"] == 0 + client.close() + + +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, per-partition and + per-field bookkeeping are indistinguishable, because the only put there is + restates its own field either way. + """ + inner = _JaggedEcho() + 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], with_dense=True), + ) + 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 fingerprints it was + # written with, or its next read cries wolf on every row + client.put_samples( + sample_ids=ids, + partition_id="p", + fields=TensorDict({"lp": torch.ones(4, 6)}, batch_size=[4]), + ) + 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 disturb ids" + assert hv["fields_skipped"] == 0, "and the read stays comparable" + client.close() + + +@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": 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"] == mismatches + assert ("more likely a bug in the check" in caplog.text) is 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_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_s"] == 0.0 + # and charted, not just tabulated + 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_s"] == 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_s"] == pytest.approx( + 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}" diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index ec7f21ca1af..e32df674413 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -555,8 +555,9 @@ 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: true # per-op timing/volume; cost is below measurement noise + 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.