From dd2b9d4d84a17b55eabef27e5174e4aa2d00729f Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 10 Aug 2026 14:57:58 -0400 Subject: [PATCH 1/6] feat(rollout): add partial recovery benchmark telemetry Signed-off-by: Anish Mahishi (cherry picked from commit 52aa3159ed0b9c86b2985443a3927a2e444b49ed) --- nemo_rl/utils/logger.py | 15 +++++++++++++++ tests/unit/utils/test_logger.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index 75df1c9499a..fc26e86c2a7 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -1052,6 +1052,21 @@ def log_metrics( for logger in self.loggers: logger.log_metrics(metrics, step, prefix, step_metric, step_finished) + def define_metric( + self, + name: str, + *, + step_metric: Optional[str] = None, + ) -> None: + """Define a W&B metric series without affecting other backends. + + TensorBoard and MLflow maintain independent steps per metric key, while + W&B needs an explicit custom step metric for event streams that advance + independently from the trainer step. + """ + if self.wandb_logger is not None: + self.wandb_logger.define_metric(name, step_metric=step_metric) + def log_hyperparams(self, params: Mapping[str, Any]) -> None: """Log hyperparameters to all enabled backends. diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py index 5b90f482981..0c82d539c1e 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -1663,6 +1663,34 @@ def test_log_metrics(self, mock_tb_logger, mock_wandb_logger, temp_dir): metrics, step, "", None, False ) + @patch("nemo_rl.utils.logger.WandbLogger") + @patch("nemo_rl.utils.logger.TensorboardLogger") + def test_define_metric_only_targets_wandb( + self, mock_tb_logger, mock_wandb_logger, temp_dir + ): + cfg = { + "wandb_enabled": True, + "tensorboard_enabled": True, + "mlflow_enabled": False, + "swanlab_enabled": False, + "monitor_gpus": False, + "wandb": {"project": "test-project"}, + "tensorboard": {"log_dir": "test_logs"}, + "log_dir": temp_dir, + } + logger = Logger(cfg) + + logger.define_metric( + "rollout/throughput/*", + step_metric="telemetry/wall_time_seconds", + ) + + mock_wandb_logger.return_value.define_metric.assert_called_once_with( + "rollout/throughput/*", + step_metric="telemetry/wall_time_seconds", + ) + assert not mock_tb_logger.return_value.define_metric.called + @patch("nemo_rl.utils.logger.WandbLogger") @patch("nemo_rl.utils.logger.TensorboardLogger") def test_log_hyperparams(self, mock_tb_logger, mock_wandb_logger, temp_dir): From b9955751e4f5177e81bcb2da2a312747d960fa2a Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 11 Aug 2026 16:50:52 -0400 Subject: [PATCH 2/6] fix(rollout): expose checkpoint backpressure telemetry Signed-off-by: Anish Mahishi (cherry picked from commit 815f3d94f44727edb56f18dd5219ccb80a1d4fbf) --- nemo_rl/models/generation/interfaces.py | 4 +++ .../models/generation/vllm/vllm_generation.py | 18 ++++++++-- .../generation/vllm/vllm_worker_async.py | 22 ++++++++++++ nemo_rl/utils/logger.py | 6 ++-- .../models/generation/test_vllm_generation.py | 24 +++++++++++++ .../test_tq_replay_buffer.py | 36 +++++++++++++++++++ tests/unit/utils/test_logger.py | 10 +++--- 7 files changed, 108 insertions(+), 12 deletions(-) diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 4afd8c9938f..9e5f3fbb979 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -417,3 +417,7 @@ def get_logger_metrics(self) -> dict[str, Any]: Dictionary of metrics. Format may vary by backend. """ return {} + + def get_latest_logger_metrics(self) -> dict[str, Any]: + """Get a bounded latest-value snapshot for frequent telemetry polls.""" + return self.get_logger_metrics() diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index f71b329475e..8fa1967c034 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1120,8 +1120,8 @@ def stop_gpu_profiling(self) -> None: futures = self.worker_group.run_all_workers_single_data("stop_gpu_profiling") ray.get(futures) - def get_vllm_logger_metrics(self) -> dict[str, Any]: - """Collect vLLM logger metrics from vLLM workers (model-owner actors only).""" + def _collect_vllm_logger_metrics(self, worker_method_name: str) -> dict[str, Any]: + """Collect one logger payload from every model-owner vLLM worker.""" if not self.cfg["vllm_cfg"].get("enable_vllm_metrics_logger", False): return {} if not self.cfg["vllm_cfg"].get("async_engine", False): @@ -1132,7 +1132,7 @@ def get_vllm_logger_metrics(self) -> dict[str, Any]: for dp_idx in range(self.worker_group.dp_size): worker_idx = self.worker_group.get_dp_leader_worker_idx(dp_idx) future = self.worker_group.run_single_worker_single_data( - "get_vllm_logger_metrics", + worker_method_name, worker_idx=worker_idx, ) futures.append(future) @@ -1166,6 +1166,14 @@ def get_vllm_logger_metrics(self) -> dict[str, Any]: return vllm_logger_metrics + def get_vllm_logger_metrics(self) -> dict[str, Any]: + """Collect vLLM metric histories for step-level performance reports.""" + return self._collect_vllm_logger_metrics("get_vllm_logger_metrics") + + def get_latest_vllm_logger_metrics(self) -> dict[str, Any]: + """Collect bounded latest-value snapshots for frequent telemetry polls.""" + return self._collect_vllm_logger_metrics("get_latest_vllm_logger_metrics") + def clear_vllm_logger_metrics(self) -> None: if not self.cfg["vllm_cfg"].get("enable_vllm_metrics_logger", False): return @@ -1185,6 +1193,10 @@ def get_logger_metrics(self) -> dict[str, Any]: """Get logger metrics for performance reporting.""" return self.get_vllm_logger_metrics() + def get_latest_logger_metrics(self) -> dict[str, Any]: + """Get latest logger values without transferring full worker histories.""" + return self.get_latest_vllm_logger_metrics() + def __del__(self) -> None: """Shuts down the worker groups when the object is deleted or is garbage collected. diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 2a1c3545fad..886af15a5f7 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -326,6 +326,28 @@ def get_vllm_logger_metrics(self) -> dict[str, Any]: } return metric + def get_latest_vllm_logger_metrics(self) -> dict[str, Any]: + """Return latest samples and prune histories after a telemetry poll.""" + if not self.cfg["vllm_cfg"].get("enable_vllm_metrics_logger", False): + return {} + + with self._vllm_metrics_lock: + histories = { + "inflight_batch_sizes": self.inflight_batch_sizes, + "num_pending_samples": self.num_pending_samples, + "kv_cache_usage_perc": self.kv_cache_usage_perc, + "generation_tokens": self.generation_tokens, + } + latest = { + name: [values[-1]] if values else [] + for name, values in histories.items() + } + self.inflight_batch_sizes = latest["inflight_batch_sizes"] + self.num_pending_samples = latest["num_pending_samples"] + self.kv_cache_usage_perc = latest["kv_cache_usage_perc"] + self.generation_tokens = latest["generation_tokens"] + return latest + def clear_vllm_logger_metrics(self) -> None: if not self.cfg["vllm_cfg"].get("enable_vllm_metrics_logger", False): return diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index fc26e86c2a7..a565d8c3565 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -378,10 +378,10 @@ def log_metrics( for k, v in metrics.items() } - # If step_metric is provided, use the corresponding value from metrics as step + # A defined custom axis selects the plotted x-value; each call still + # needs to commit its own W&B history row. if step_metric and step_metric in metrics: - # commit=False so the step does not get incremented - self.run.log(metrics, commit=False) + self.run.log(metrics) elif step_finished: # Commit param defaults to None. By default if step is set, then commit defaults to False # Here, we have an explicit fork for commit in case W&B ever decides to change their default logic. diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 33f0507750e..0d5bb9c9174 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -16,6 +16,7 @@ import json import os import sys +import threading import types from copy import deepcopy from pathlib import Path @@ -189,6 +190,29 @@ def test_sampling_params_preserve_bad_words(): assert sampling_params["bad_words"] == ["", ""] +def test_vllm_latest_metric_snapshot_prunes_worker_histories(): + worker = object.__new__(VllmAsyncGenerationWorkerImpl) + worker.cfg = {"vllm_cfg": {"enable_vllm_metrics_logger": True}} + worker._vllm_metrics_lock = threading.Lock() + worker.inflight_batch_sizes = [1, 2] + worker.num_pending_samples = [3, 4] + worker.kv_cache_usage_perc = [0.2, 0.6] + worker.generation_tokens = [10, 30] + + latest = worker.get_latest_vllm_logger_metrics() + + assert latest == { + "inflight_batch_sizes": [2], + "num_pending_samples": [4], + "kv_cache_usage_perc": [0.6], + "generation_tokens": [30], + } + assert worker.inflight_batch_sizes == [2] + assert worker.num_pending_samples == [4] + assert worker.kv_cache_usage_perc == [0.6] + assert worker.generation_tokens == [30] + + def test_resolve_enable_prefix_caching_respects_explicit_config(monkeypatch): def raise_if_called(): raise AssertionError("CUDA capability should not be queried") diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 36f8eb5367d..67708a05ae7 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -26,6 +26,7 @@ from nemo_rl.algorithms.async_utils.replay_buffer import ( REPLAY_BUFFER_METADATA_SCHEMA_VERSION, REPLAY_BUFFER_METADATA_STORAGE, + CheckpointMutationKind, DataPlaneCheckpointBarrier, TQReplayBuffer, replay_manifest_digest, @@ -338,6 +339,41 @@ async def checkpoint(tag: str) -> None: asyncio.run(exercise()) + def test_reports_mutations_blocked_by_checkpoint(self): + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + + async def mutate(kind: CheckpointMutationKind) -> None: + async with barrier.mutation(kind): + pass + + async with barrier.checkpoint(): + commit_task = asyncio.create_task(mutate("group_commits")) + seal_task = asyncio.create_task(mutate("sibling_seals")) + await asyncio.sleep(0) + active = await barrier.drain_telemetry() + assert active.checkpoint_active + assert active.waiting_mutations == 2 + assert active.max_waiting_mutations == 2 + assert sum(active.blocked_by_kind.values()) == 0 + + await asyncio.gather(commit_task, seal_task) + completed = await barrier.drain_telemetry() + assert not completed.checkpoint_active + assert completed.waiting_mutations == 0 + assert completed.max_waiting_mutations == 2 + assert completed.blocked_by_kind["group_commits"] == 1 + assert completed.blocked_by_kind["sibling_seals"] == 1 + assert len(completed.wait_durations_s) == 2 + assert all(duration >= 0 for duration in completed.wait_durations_s) + + drained = await barrier.drain_telemetry() + assert sum(drained.blocked_by_kind.values()) == 0 + assert drained.wait_durations_s == () + assert drained.max_waiting_mutations == 0 + + asyncio.run(exercise()) + class TestTQReplayBufferReserveCommit: def test_commit_waits_for_active_checkpoint(self): diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py index 0c82d539c1e..2f441870898 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -386,10 +386,9 @@ def test_log_metrics_with_step_metric(self, mock_wandb): logger.log_metrics(metrics, step, step_metric=step_metric) - # Check that log was called with metrics and commit=False - # When using step_metric, step should be ignored and commit=False should be used + # The custom metric supplies the x-axis, while each event commits a row. mock_run = mock_wandb.init.return_value - mock_run.log.assert_called_once_with(metrics, commit=False) + mock_run.log.assert_called_once_with(metrics) @patch("nemo_rl.utils.logger.wandb") def test_log_metrics_with_prefix_and_step_metric(self, mock_wandb): @@ -407,15 +406,14 @@ def test_log_metrics_with_prefix_and_step_metric(self, mock_wandb): logger.log_metrics(metrics, step, prefix=prefix, step_metric=step_metric) - # Check that log was called with prefixed metrics and commit=False - # The step_metric key gets prefixed based on the current implementation + # The step_metric key gets prefixed based on the current implementation. mock_run = mock_wandb.init.return_value expected_metrics = { "train/loss": 0.5, "train/accuracy": 0.8, "train/iteration": 15, } - mock_run.log.assert_called_once_with(expected_metrics, commit=False) + mock_run.log.assert_called_once_with(expected_metrics) @patch("nemo_rl.utils.logger.wandb") def test_define_metric(self, mock_wandb): From ecb02d3e84a5ad2efb81b4f17f3035e70a01b7e6 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 11 Aug 2026 18:36:12 -0400 Subject: [PATCH 3/6] fix(rollout): address telemetry review feedback Signed-off-by: Anish Mahishi (cherry picked from commit 9853cf487046fea367995ed1d68a8e4a210495a2) --- .../nemo_gym/grpo_qwen3_30ba3b_instruct.yaml | 4 ++ .../single_controller_utils/setup.py | 14 +++++ nemo_rl/models/generation/interfaces.py | 4 +- nemo_rl/models/generation/vllm/config.py | 5 ++ .../models/generation/vllm/vllm_generation.py | 12 ++-- .../generation/vllm/vllm_worker_async.py | 16 +++--- nemo_rl/utils/logger.py | 54 ++++++++++++------ .../models/generation/test_vllm_generation.py | 5 +- .../test_single_controller_setup.py | 8 +++ tests/unit/utils/test_logger.py | 55 +++++++++++++++---- 10 files changed, 132 insertions(+), 45 deletions(-) diff --git a/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml b/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml index 68ff288e2a4..397d52023be 100644 --- a/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml +++ b/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml @@ -53,6 +53,10 @@ policy: generation: vllm_cfg: tensor_parallel_size: 4 + # Required for rollout-throughput checkpoint A/B telemetry. Keep this + # identical in baseline and checkpoint-enabled runs. + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 # This is a very low GPU mem utilization. We GPU OOM in two places: # Refit after train, refit before validation. gpu_memory_utilization: 0.7 diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index a3706e1a8c0..a2f62237bed 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -25,6 +25,7 @@ import io import os import time +import warnings from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from functools import partial @@ -718,6 +719,19 @@ def setup_single_controller( "single_controller_utils.setup requires policy.generation in master_config" ) + telemetry_interval_s = master_config.rollout_checkpointing.telemetry_interval_s + vllm_cfg = generation_config.get("vllm_cfg", {}) + if telemetry_interval_s is not None and not vllm_cfg.get( + "enable_vllm_metrics_logger", False + ): + warnings.warn( + "rollout_checkpointing.telemetry_interval_s is enabled, but " + "policy.generation.vllm_cfg.enable_vllm_metrics_logger is false. " + "Canonical rollout telemetry will be recorded, but vLLM token, " + "request, and KV-cache signals will be absent.", + stacklevel=2, + ) + if data_config["use_multiple_dataloader"]: raise NotImplementedError( "single_controller_utils does not support " diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 9e5f3fbb979..6e8e539b979 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -418,6 +418,6 @@ def get_logger_metrics(self) -> dict[str, Any]: """ return {} - def get_latest_logger_metrics(self) -> dict[str, Any]: - """Get a bounded latest-value snapshot for frequent telemetry polls.""" + def drain_latest_logger_metrics(self) -> dict[str, Any]: + """Consume a bounded latest-value snapshot for frequent telemetry polls.""" return self.get_logger_metrics() diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index 6281bb2af18..7672a4e8291 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -48,6 +48,11 @@ class VllmSpecificArgs(TypedDict): kv_cache_dtype: Literal["auto", "fp8", "fp8_e4m3"] enforce_eager: NotRequired[bool] enable_return_routed_experts: NotRequired[bool] + # Collect vLLM request, cache, and cumulative token counters in a model-owner + # background thread for performance diagnostics. + enable_vllm_metrics_logger: NotRequired[bool] + # Sampling cadence for the optional vLLM metrics logger. + vllm_metrics_logger_interval: NotRequired[float] # Whether to show a tqdm progress bar during generation. Defaults to vLLM's own default (True) when absent. Only applies when async_engine is False. use_tqdm: NotRequired[bool] # By default, NeMo RL only has a Python handle to the vllm.LLM generation engine. The expose_http_server flag here will expose that generation engine as an HTTP server. diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 8fa1967c034..b41f657349d 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1170,9 +1170,9 @@ def get_vllm_logger_metrics(self) -> dict[str, Any]: """Collect vLLM metric histories for step-level performance reports.""" return self._collect_vllm_logger_metrics("get_vllm_logger_metrics") - def get_latest_vllm_logger_metrics(self) -> dict[str, Any]: - """Collect bounded latest-value snapshots for frequent telemetry polls.""" - return self._collect_vllm_logger_metrics("get_latest_vllm_logger_metrics") + def drain_latest_vllm_logger_metrics(self) -> dict[str, Any]: + """Consume bounded latest-value snapshots for frequent telemetry polls.""" + return self._collect_vllm_logger_metrics("drain_latest_vllm_logger_metrics") def clear_vllm_logger_metrics(self) -> None: if not self.cfg["vllm_cfg"].get("enable_vllm_metrics_logger", False): @@ -1193,9 +1193,9 @@ def get_logger_metrics(self) -> dict[str, Any]: """Get logger metrics for performance reporting.""" return self.get_vllm_logger_metrics() - def get_latest_logger_metrics(self) -> dict[str, Any]: - """Get latest logger values without transferring full worker histories.""" - return self.get_latest_vllm_logger_metrics() + def drain_latest_logger_metrics(self) -> dict[str, Any]: + """Consume latest values without transferring full worker histories.""" + return self.drain_latest_vllm_logger_metrics() def __del__(self) -> None: """Shuts down the worker groups when the object is deleted or is garbage collected. diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 886af15a5f7..b8935ce2199 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -246,7 +246,7 @@ def locked_send_multipart(*args: Any, **kwargs: Any) -> Any: def _start_vllm_metrics_logger(self) -> None: """Start a background thread that periodically collects vLLM logger metrics. - Controlled by vllm_metrics_logger_interval (default: 0.5) in vllm_cfg. + Controlled by the required vllm_metrics_logger_interval in vllm_cfg. Runs only on the model-owner actor. """ from vllm.v1.metrics.reader import Gauge, Counter, get_metrics_snapshot @@ -326,7 +326,7 @@ def get_vllm_logger_metrics(self) -> dict[str, Any]: } return metric - def get_latest_vllm_logger_metrics(self) -> dict[str, Any]: + def drain_latest_vllm_logger_metrics(self) -> dict[str, Any]: """Return latest samples and prune histories after a telemetry poll.""" if not self.cfg["vllm_cfg"].get("enable_vllm_metrics_logger", False): return {} @@ -342,11 +342,13 @@ def get_latest_vllm_logger_metrics(self) -> dict[str, Any]: name: [values[-1]] if values else [] for name, values in histories.items() } - self.inflight_batch_sizes = latest["inflight_batch_sizes"] - self.num_pending_samples = latest["num_pending_samples"] - self.kv_cache_usage_perc = latest["kv_cache_usage_perc"] - self.generation_tokens = latest["generation_tokens"] - return latest + # Keep worker-owned histories distinct from the lists handed to Ray; + # the sampling thread may append immediately after this lock exits. + self.inflight_batch_sizes = list(latest["inflight_batch_sizes"]) + self.num_pending_samples = list(latest["num_pending_samples"]) + self.kv_cache_usage_perc = list(latest["kv_cache_usage_perc"]) + self.generation_tokens = list(latest["generation_tokens"]) + return {name: list(values) for name, values in latest.items()} def clear_vllm_logger_metrics(self) -> None: if not self.cfg["vllm_cfg"].get("enable_vllm_metrics_logger", False): diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index a565d8c3565..27930c2b976 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -46,6 +46,11 @@ # Flag to track if rich logging has been configured _rich_logging_configured = False +# W&B owns one monotonically increasing internal history row counter per run. +# Keep the NeMo-RL caller step as a custom axis so independently committed +# telemetry rows cannot advance past, and thereby invalidate, trainer steps. +_WANDB_CALLER_STEP_METRIC = "nemo_rl/step" + class WandbConfig(TypedDict): project: NotRequired[str] @@ -220,6 +225,9 @@ def __init__(self, cfg: WandbConfig, log_dir: Optional[str] = None): wandb_init_config = dict(cfg) wandb_init_config.pop("log_nemo_gym_full_result_tables", None) self.run = wandb.init(**wandb_init_config, dir=log_dir) + self._log_lock = threading.Lock() + self.run.define_metric(_WANDB_CALLER_STEP_METRIC, hidden=True) + self.run.define_metric("*", step_metric=_WANDB_CALLER_STEP_METRIC) if os.environ.get("RAY_BACKEND_LOG_LEVEL", "").lower() == "debug": print( @@ -353,7 +361,8 @@ def define_metric( name: Name of the metric or pattern (e.g. 'ray/*') step_metric: Optional name of the step metric to use """ - self.run.define_metric(name, step_metric=step_metric) + with self._log_lock: + self.run.define_metric(name, step_metric=step_metric) def log_metrics( self, @@ -378,16 +387,13 @@ def log_metrics( for k, v in metrics.items() } - # A defined custom axis selects the plotted x-value; each call still - # needs to commit its own W&B history row. - if step_metric and step_metric in metrics: - self.run.log(metrics) - elif step_finished: - # Commit param defaults to None. By default if step is set, then commit defaults to False - # Here, we have an explicit fork for commit in case W&B ever decides to change their default logic. - self.run.log(metrics, step=step, commit=True) - else: - self.run.log(metrics, step=step) + # Every call is an independent W&B event row. The caller's logical step + # is a custom axis rather than W&B's internal row number, so asynchronous + # telemetry cannot make a later trainer step look stale and get dropped. + event_metrics = dict(metrics) + event_metrics[_WANDB_CALLER_STEP_METRIC] = step + with self._log_lock: + self.run.log(event_metrics) def log_hyperparams(self, params: Mapping[str, Any]) -> None: """Log hyperparameters to wandb. @@ -404,7 +410,8 @@ def log_plot(self, figure: plt.Figure, step: int, name: str) -> None: figure: Matplotlib figure to log step: Global step value """ - self.run.log({name: figure}, step=step) + with self._log_lock: + self.run.log({name: figure, _WANDB_CALLER_STEP_METRIC: step}) def finish(self) -> None: """Flush queued metrics and close the wandb service. @@ -412,9 +419,10 @@ def finish(self) -> None: Required when the run lives inside a Ray actor: Ray tears the worker down before wandb's atexit hook can drain the IPC queue to the service. """ - if self.run is not None: - self.run.finish() - self.run = None + with self._log_lock: + if self.run is not None: + self.run.finish() + self.run = None def log_histogram(self, histogram: list[Any], step: int, name: str) -> None: """Log histogram metrics to wandb. @@ -425,11 +433,23 @@ def log_histogram(self, histogram: list[Any], step: int, name: str) -> None: name: Name of the metric """ try: - self.run.log({name: wandb.Histogram(histogram)}, step=step) + with self._log_lock: + self.run.log( + { + name: wandb.Histogram(histogram), + _WANDB_CALLER_STEP_METRIC: step, + } + ) except ValueError: # When all values are identical, numpy cannot create finite-sized bins. # Log the scalar value instead. - self.run.log({name: histogram[0] if len(histogram) > 0 else 0}, step=step) + with self._log_lock: + self.run.log( + { + name: histogram[0] if len(histogram) > 0 else 0, + _WANDB_CALLER_STEP_METRIC: step, + } + ) class SwanlabLogger(LoggerInterface): diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 0d5bb9c9174..815b85bc8eb 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -190,7 +190,7 @@ def test_sampling_params_preserve_bad_words(): assert sampling_params["bad_words"] == ["", ""] -def test_vllm_latest_metric_snapshot_prunes_worker_histories(): +def test_vllm_latest_metric_drain_prunes_worker_histories(): worker = object.__new__(VllmAsyncGenerationWorkerImpl) worker.cfg = {"vllm_cfg": {"enable_vllm_metrics_logger": True}} worker._vllm_metrics_lock = threading.Lock() @@ -199,7 +199,7 @@ def test_vllm_latest_metric_snapshot_prunes_worker_histories(): worker.kv_cache_usage_perc = [0.2, 0.6] worker.generation_tokens = [10, 30] - latest = worker.get_latest_vllm_logger_metrics() + latest = worker.drain_latest_vllm_logger_metrics() assert latest == { "inflight_batch_sizes": [2], @@ -211,6 +211,7 @@ def test_vllm_latest_metric_snapshot_prunes_worker_histories(): assert worker.num_pending_samples == [4] assert worker.kv_cache_usage_perc == [0.6] assert worker.generation_tokens == [30] + assert latest["generation_tokens"] is not worker.generation_tokens def test_resolve_enable_prefix_caching_respects_explicit_config(monkeypatch): diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index f1246cd5e73..f7238786b13 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -286,6 +286,14 @@ def test_raises_when_data_plane_disabled(self): with pytest.raises(ValueError, match="data_plane.enabled=True"): setup_single_controller(mc, MagicMock()) + def test_warns_when_rollout_telemetry_lacks_vllm_metrics(self, patched_factories): + mc = _make_master_config() + mc.rollout_checkpointing.telemetry_interval_s = 30.0 + mc.policy["generation"]["vllm_cfg"] = {} + + with pytest.warns(UserWarning, match="vLLM token, request, and KV-cache"): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_rejects_mooncake_data_plane_checkpointing(self): mc = _make_master_config() mc.data_plane.update( diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py index 2f441870898..e209311579c 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -351,9 +351,9 @@ def test_log_metrics(self, mock_wandb): step = 10 logger.log_metrics(metrics, step) - # Check that log was called with metrics and step + # W&B's internal row step is implicit; the caller step is a custom axis. mock_run = mock_wandb.init.return_value - mock_run.log.assert_called_once_with(metrics, step=step) + mock_run.log.assert_called_once_with({**metrics, "nemo_rl/step": step}) @patch("nemo_rl.utils.logger.wandb") def test_log_metrics_with_prefix(self, mock_wandb): @@ -366,10 +366,14 @@ def test_log_metrics_with_prefix(self, mock_wandb): prefix = "train" logger.log_metrics(metrics, step, prefix) - # Check that log was called with prefixed metrics and step + # Check that prefixed metrics retain the caller step as a custom axis. mock_run = mock_wandb.init.return_value - expected_metrics = {"train/loss": 0.5, "train/accuracy": 0.8} - mock_run.log.assert_called_once_with(expected_metrics, step=step) + expected_metrics = { + "train/loss": 0.5, + "train/accuracy": 0.8, + "nemo_rl/step": step, + } + mock_run.log.assert_called_once_with(expected_metrics) @patch("nemo_rl.utils.logger.wandb") def test_log_metrics_with_step_metric(self, mock_wandb): @@ -386,9 +390,10 @@ def test_log_metrics_with_step_metric(self, mock_wandb): logger.log_metrics(metrics, step, step_metric=step_metric) - # The custom metric supplies the x-axis, while each event commits a row. + # The requested custom metric supplies the series x-axis, while the + # caller step remains available to all other metrics in the row. mock_run = mock_wandb.init.return_value - mock_run.log.assert_called_once_with(metrics) + mock_run.log.assert_called_once_with({**metrics, "nemo_rl/step": step}) @patch("nemo_rl.utils.logger.wandb") def test_log_metrics_with_prefix_and_step_metric(self, mock_wandb): @@ -412,9 +417,38 @@ def test_log_metrics_with_prefix_and_step_metric(self, mock_wandb): "train/loss": 0.5, "train/accuracy": 0.8, "train/iteration": 15, + "nemo_rl/step": step, } mock_run.log.assert_called_once_with(expected_metrics) + @patch("nemo_rl.utils.logger.wandb") + def test_independent_events_do_not_reuse_wandb_internal_step(self, mock_wandb): + """Telemetry commits must not make a later trainer step stale.""" + logger = WandbLogger({}) + + logger.log_metrics({"loss": 1.0}, step=1, prefix="train") + logger.log_metrics( + {"telemetry/wall_time_seconds": 30.0, "tokens_per_second": 10.0}, + step=1, + prefix="rollout/throughput", + step_metric="telemetry/wall_time_seconds", + ) + logger.log_metrics({"loss": 0.5}, step=2, prefix="train") + + mock_run = mock_wandb.init.return_value + assert mock_run.log.call_args_list == [ + call({"train/loss": 1.0, "nemo_rl/step": 1}), + call( + { + "telemetry/wall_time_seconds": 30.0, + "rollout/throughput/tokens_per_second": 10.0, + "nemo_rl/step": 1, + } + ), + call({"train/loss": 0.5, "nemo_rl/step": 2}), + ] + assert all("step" not in kwargs for _, kwargs in mock_run.log.call_args_list) + @patch("nemo_rl.utils.logger.wandb") def test_define_metric(self, mock_wandb): """Test defining a metric with a custom step metric.""" @@ -424,11 +458,10 @@ def test_define_metric(self, mock_wandb): # Define metric pattern and step metric logger.define_metric("ray/*", step_metric="ray/ray_step") - # Check that define_metric was called + # Check that the caller's custom series definition is preserved in + # addition to the logger-wide logical-step axis. mock_run = mock_wandb.init.return_value - mock_run.define_metric.assert_called_once_with( - "ray/*", step_metric="ray/ray_step" - ) + mock_run.define_metric.assert_any_call("ray/*", step_metric="ray/ray_step") @patch("nemo_rl.utils.logger.wandb") def test_log_hyperparams(self, mock_wandb): From 1f44e1d47984646aa6d5006477122b612304521b Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 11 Aug 2026 19:03:00 -0400 Subject: [PATCH 4/6] fix(rollout): make telemetry axes deterministic Signed-off-by: Anish Mahishi (cherry picked from commit 67f8278ea5f56e21bfc7820f2d3395ec8b7808e1) --- .../single_controller_utils/setup.py | 41 +++- nemo_rl/utils/logger.py | 210 +++++++++++++++--- .../test_single_controller_setup.py | 26 ++- tests/unit/utils/test_logger.py | 86 ++++++- 4 files changed, 315 insertions(+), 48 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index a2f62237bed..7762b58543c 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -720,17 +720,36 @@ def setup_single_controller( ) telemetry_interval_s = master_config.rollout_checkpointing.telemetry_interval_s - vllm_cfg = generation_config.get("vllm_cfg", {}) - if telemetry_interval_s is not None and not vllm_cfg.get( - "enable_vllm_metrics_logger", False - ): - warnings.warn( - "rollout_checkpointing.telemetry_interval_s is enabled, but " - "policy.generation.vllm_cfg.enable_vllm_metrics_logger is false. " - "Canonical rollout telemetry will be recorded, but vLLM token, " - "request, and KV-cache signals will be absent.", - stacklevel=2, - ) + if telemetry_interval_s is not None: + generation_backend = generation_config["backend"] + if generation_backend != "vllm": + warnings.warn( + "rollout_checkpointing.telemetry_interval_s is enabled with " + f"policy.generation.backend={generation_backend!r}. Canonical " + "rollout telemetry will be recorded, but vLLM token, request, " + "and KV-cache signals are unavailable for this backend.", + stacklevel=2, + ) + else: + vllm_cfg = generation_config["vllm_cfg"] + if not vllm_cfg.get("enable_vllm_metrics_logger"): + warnings.warn( + "rollout_checkpointing.telemetry_interval_s is enabled, but " + "policy.generation.vllm_cfg.enable_vllm_metrics_logger is " + "false. Canonical rollout telemetry will be recorded, but " + "vLLM token, request, and KV-cache signals will be absent.", + stacklevel=2, + ) + elif not vllm_cfg["async_engine"]: + warnings.warn( + "rollout_checkpointing.telemetry_interval_s and " + "policy.generation.vllm_cfg.enable_vllm_metrics_logger are " + "enabled, but vLLM metric collection requires " + "policy.generation.vllm_cfg.async_engine=true. Canonical " + "rollout telemetry will be recorded, but vLLM token, request, " + "and KV-cache signals will be absent.", + stacklevel=2, + ) if data_config["use_multiple_dataloader"]: raise NotImplementedError( diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index 27930c2b976..4c392d64cd8 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -49,6 +49,9 @@ # W&B owns one monotonically increasing internal history row counter per run. # Keep the NeMo-RL caller step as a custom axis so independently committed # telemetry rows cannot advance past, and thereby invalidate, trainer steps. +# Metric axes are registered by exact name when a key is first logged. A +# catch-all ``*`` definition is intentionally avoided because wandb-core does +# not deterministically resolve overlapping glob definitions. _WANDB_CALLER_STEP_METRIC = "nemo_rl/step" @@ -226,8 +229,13 @@ def __init__(self, cfg: WandbConfig, log_dir: Optional[str] = None): wandb_init_config.pop("log_nemo_gym_full_result_tables", None) self.run = wandb.init(**wandb_init_config, dir=log_dir) self._log_lock = threading.Lock() + self._metric_step_patterns: dict[str, Optional[str]] = {} + self._defined_metric_axes: dict[str, Optional[str]] = { + _WANDB_CALLER_STEP_METRIC: None + } + self._pending_step: Optional[int] = None + self._pending_metrics: dict[str, Any] = {} self.run.define_metric(_WANDB_CALLER_STEP_METRIC, hidden=True) - self.run.define_metric("*", step_metric=_WANDB_CALLER_STEP_METRIC) if os.environ.get("RAY_BACKEND_LOG_LEVEL", "").lower() == "debug": print( @@ -355,14 +363,151 @@ def define_metric( name: str, step_metric: Optional[str] = None, ) -> None: - """Define a metric with custom step metric. + """Register the axis to apply lazily to matching exact metric names. Args: name: Name of the metric or pattern (e.g. 'ray/*') step_metric: Optional name of the step metric to use """ with self._log_lock: - self.run.define_metric(name, step_metric=step_metric) + existing_step_metric = self._metric_step_patterns.get(name) + if name in self._metric_step_patterns and ( + existing_step_metric != step_metric + ): + raise ValueError( + f"W&B metric pattern {name!r} is already registered with " + f"step metric {existing_step_metric!r}, not {step_metric!r}" + ) + self._metric_step_patterns[name] = step_metric + if step_metric is not None: + self._define_exact_metric_locked( + step_metric, + step_metric=None, + hidden=True, + ) + if "*" not in name and name != step_metric: + self._define_exact_metric_locked(name, step_metric=step_metric) + + def _define_exact_metric_locked( + self, + name: str, + *, + step_metric: Optional[str], + hidden: bool = False, + ) -> None: + """Define one exact W&B key once and reject conflicting axes.""" + if name in self._defined_metric_axes: + existing_step_metric = self._defined_metric_axes[name] + if existing_step_metric != step_metric: + raise ValueError( + f"W&B metric {name!r} is already defined with step metric " + f"{existing_step_metric!r}, not {step_metric!r}" + ) + return + + define_kwargs: dict[str, Any] = {} + if step_metric is not None: + define_kwargs["step_metric"] = step_metric + if hidden: + define_kwargs["hidden"] = True + self.run.define_metric(name, **define_kwargs) + self._defined_metric_axes[name] = step_metric + + def _matching_step_metric(self, metric_name: str) -> Optional[str]: + """Resolve a locally registered pattern without sending globs to W&B.""" + matches: list[tuple[int, Optional[str]]] = [] + for pattern, step_metric in self._metric_step_patterns.items(): + if pattern.endswith("*"): + matches_pattern = metric_name.startswith(pattern[:-1]) + else: + matches_pattern = metric_name == pattern + if matches_pattern: + matches.append((len(pattern), step_metric)) + + if not matches: + return None + longest_match = max(length for length, _ in matches) + matching_axes = { + step_metric for length, step_metric in matches if length == longest_match + } + if len(matching_axes) != 1: + raise ValueError( + f"W&B metric {metric_name!r} matches conflicting step metrics: " + f"{matching_axes!r}" + ) + return matching_axes.pop() + + def _resolve_event_step_metric( + self, + metrics: Mapping[str, Any], + explicit_step_metric: Optional[str], + ) -> Optional[str]: + """Choose the single custom axis for an event, if it has one.""" + if explicit_step_metric is not None: + return explicit_step_metric + + matching_axes = { + step_metric + for metric_name in metrics + if (step_metric := self._matching_step_metric(metric_name)) is not None + } + if len(matching_axes) > 1: + raise ValueError( + "One W&B event cannot contain metrics with different custom " + f"step metrics: {matching_axes!r}" + ) + return next(iter(matching_axes), None) + + def _define_event_metrics_locked( + self, + metrics: Mapping[str, Any], + *, + step_metric: str, + ) -> None: + """Define the event axis and each data key by exact name.""" + self._define_exact_metric_locked( + step_metric, + step_metric=None, + hidden=True, + ) + for metric_name in metrics: + if metric_name != step_metric: + self._define_exact_metric_locked( + metric_name, + step_metric=step_metric, + ) + + def _flush_pending_metrics_locked(self) -> None: + """Commit the accumulated metrics for one trainer step.""" + if self._pending_step is None: + return + event_metrics = dict(self._pending_metrics) + event_metrics[_WANDB_CALLER_STEP_METRIC] = self._pending_step + self.run.log(event_metrics) + self._pending_step = None + self._pending_metrics = {} + + def _buffer_step_metrics_locked( + self, + metrics: Mapping[str, Any], + *, + step: int, + step_finished: bool, + ) -> None: + """Accumulate correlated trainer metrics into one W&B history row.""" + if self._pending_step is not None and self._pending_step != step: + self._flush_pending_metrics_locked() + if self._pending_step is None: + self._pending_step = step + + for metric_name in metrics: + self._define_exact_metric_locked( + metric_name, + step_metric=_WANDB_CALLER_STEP_METRIC, + ) + self._pending_metrics.update(metrics) + if step_finished: + self._flush_pending_metrics_locked() def log_metrics( self, @@ -387,13 +532,28 @@ def log_metrics( for k, v in metrics.items() } - # Every call is an independent W&B event row. The caller's logical step - # is a custom axis rather than W&B's internal row number, so asynchronous - # telemetry cannot make a later trainer step look stale and get dropped. - event_metrics = dict(metrics) - event_metrics[_WANDB_CALLER_STEP_METRIC] = step with self._log_lock: - self.run.log(event_metrics) + event_step_metric = self._resolve_event_step_metric(metrics, step_metric) + if event_step_metric is not None: + if event_step_metric not in metrics: + raise ValueError( + f"Custom W&B step metric {event_step_metric!r} is missing " + "from the logged event" + ) + self._define_event_metrics_locked( + metrics, + step_metric=event_step_metric, + ) + # Custom-axis streams (rollout telemetry, GPU monitoring) are + # independent events and must not carry nemo_rl/step. + self.run.log(dict(metrics)) + return + + self._buffer_step_metrics_locked( + metrics, + step=step, + step_finished=step_finished, + ) def log_hyperparams(self, params: Mapping[str, Any]) -> None: """Log hyperparameters to wandb. @@ -401,7 +561,8 @@ def log_hyperparams(self, params: Mapping[str, Any]) -> None: Args: params: Dict of hyperparameters to log """ - self.run.config.update(params, allow_val_change=True) + with self._log_lock: + self.run.config.update(params, allow_val_change=True) def log_plot(self, figure: plt.Figure, step: int, name: str) -> None: """Log a plot to wandb. @@ -411,7 +572,11 @@ def log_plot(self, figure: plt.Figure, step: int, name: str) -> None: step: Global step value """ with self._log_lock: - self.run.log({name: figure, _WANDB_CALLER_STEP_METRIC: step}) + self._buffer_step_metrics_locked( + {name: figure}, + step=step, + step_finished=False, + ) def finish(self) -> None: """Flush queued metrics and close the wandb service. @@ -421,6 +586,7 @@ def finish(self) -> None: """ with self._log_lock: if self.run is not None: + self._flush_pending_metrics_locked() self.run.finish() self.run = None @@ -433,23 +599,17 @@ def log_histogram(self, histogram: list[Any], step: int, name: str) -> None: name: Name of the metric """ try: - with self._log_lock: - self.run.log( - { - name: wandb.Histogram(histogram), - _WANDB_CALLER_STEP_METRIC: step, - } - ) + value: Any = wandb.Histogram(histogram) except ValueError: # When all values are identical, numpy cannot create finite-sized bins. # Log the scalar value instead. - with self._log_lock: - self.run.log( - { - name: histogram[0] if len(histogram) > 0 else 0, - _WANDB_CALLER_STEP_METRIC: step, - } - ) + value = histogram[0] if len(histogram) > 0 else 0 + with self._log_lock: + self._buffer_step_metrics_locked( + {name: value}, + step=step, + step_finished=False, + ) class SwanlabLogger(LoggerInterface): diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index f7238786b13..8b739cdb182 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -289,11 +289,35 @@ def test_raises_when_data_plane_disabled(self): def test_warns_when_rollout_telemetry_lacks_vllm_metrics(self, patched_factories): mc = _make_master_config() mc.rollout_checkpointing.telemetry_interval_s = 30.0 - mc.policy["generation"]["vllm_cfg"] = {} + mc.policy["generation"]["vllm_cfg"] = {"async_engine": True} with pytest.warns(UserWarning, match="vLLM token, request, and KV-cache"): setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_warns_when_vllm_telemetry_uses_sync_engine(self, patched_factories): + mc = _make_master_config() + mc.rollout_checkpointing.telemetry_interval_s = 30.0 + mc.policy["generation"]["vllm_cfg"] = { + "async_engine": False, + "enable_vllm_metrics_logger": True, + } + + with pytest.warns(UserWarning, match="async_engine=true"): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + def test_non_vllm_telemetry_warning_has_no_vllm_config_guidance( + self, patched_factories + ): + mc = _make_master_config(backend="sglang") + mc.rollout_checkpointing.telemetry_interval_s = 30.0 + + with pytest.warns(UserWarning) as warning_records: + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + messages = [str(record.message) for record in warning_records] + assert any("backend='sglang'" in message for message in messages) + assert all("enable_vllm_metrics_logger" not in message for message in messages) + def test_rejects_mooncake_data_plane_checkpointing(self): mc = _make_master_config() mc.data_plane.update( diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py index e209311579c..766fc3a282a 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -349,7 +349,7 @@ def test_log_metrics(self, mock_wandb): metrics = {"loss": 0.5, "accuracy": 0.8} step = 10 - logger.log_metrics(metrics, step) + logger.log_metrics(metrics, step, step_finished=True) # W&B's internal row step is implicit; the caller step is a custom axis. mock_run = mock_wandb.init.return_value @@ -364,7 +364,7 @@ def test_log_metrics_with_prefix(self, mock_wandb): metrics = {"loss": 0.5, "accuracy": 0.8} step = 10 prefix = "train" - logger.log_metrics(metrics, step, prefix) + logger.log_metrics(metrics, step, prefix, step_finished=True) # Check that prefixed metrics retain the caller step as a custom axis. mock_run = mock_wandb.init.return_value @@ -390,10 +390,12 @@ def test_log_metrics_with_step_metric(self, mock_wandb): logger.log_metrics(metrics, step, step_metric=step_metric) - # The requested custom metric supplies the series x-axis, while the - # caller step remains available to all other metrics in the row. + # The requested custom metric supplies the series x-axis. Independent + # event streams must not also carry the trainer-step axis. mock_run = mock_wandb.init.return_value - mock_run.log.assert_called_once_with({**metrics, "nemo_rl/step": step}) + mock_run.log.assert_called_once_with(metrics) + mock_run.define_metric.assert_any_call("loss", step_metric="iteration") + mock_run.define_metric.assert_any_call("accuracy", step_metric="iteration") @patch("nemo_rl.utils.logger.wandb") def test_log_metrics_with_prefix_and_step_metric(self, mock_wandb): @@ -417,9 +419,11 @@ def test_log_metrics_with_prefix_and_step_metric(self, mock_wandb): "train/loss": 0.5, "train/accuracy": 0.8, "train/iteration": 15, - "nemo_rl/step": step, } mock_run.log.assert_called_once_with(expected_metrics) + mock_run.define_metric.assert_any_call( + "train/loss", step_metric="train/iteration" + ) @patch("nemo_rl.utils.logger.wandb") def test_independent_events_do_not_reuse_wandb_internal_step(self, mock_wandb): @@ -427,21 +431,34 @@ def test_independent_events_do_not_reuse_wandb_internal_step(self, mock_wandb): logger = WandbLogger({}) logger.log_metrics({"loss": 1.0}, step=1, prefix="train") + logger.log_metrics({"seconds": 5.0}, step=1, prefix="timing/train") logger.log_metrics( {"telemetry/wall_time_seconds": 30.0, "tokens_per_second": 10.0}, step=1, prefix="rollout/throughput", step_metric="telemetry/wall_time_seconds", ) - logger.log_metrics({"loss": 0.5}, step=2, prefix="train") + logger.log_metrics( + {"tokens_per_second": 20.0}, + step=1, + prefix="performance", + step_finished=True, + ) + logger.log_metrics({"loss": 0.5}, step=2, prefix="train", step_finished=True) mock_run = mock_wandb.init.return_value assert mock_run.log.call_args_list == [ - call({"train/loss": 1.0, "nemo_rl/step": 1}), call( { "telemetry/wall_time_seconds": 30.0, "rollout/throughput/tokens_per_second": 10.0, + } + ), + call( + { + "train/loss": 1.0, + "timing/train/seconds": 5.0, + "performance/tokens_per_second": 20.0, "nemo_rl/step": 1, } ), @@ -458,10 +475,57 @@ def test_define_metric(self, mock_wandb): # Define metric pattern and step metric logger.define_metric("ray/*", step_metric="ray/ray_step") - # Check that the caller's custom series definition is preserved in - # addition to the logger-wide logical-step axis. + logger.log_metrics( + {"ray/ray_step": 15.0, "gpu_utilization": 80.0}, + step=10, + prefix="ray", + step_metric="ray/ray_step", + ) + logger.log_metrics( + {"ray/ray_step": 16.0, "gpu_utilization": 81.0}, + step=11, + prefix="ray", + step_metric="ray/ray_step", + ) + + # Patterns stay local; W&B receives deterministic exact-name rules. mock_run = mock_wandb.init.return_value - mock_run.define_metric.assert_any_call("ray/*", step_metric="ray/ray_step") + assert call("ray/*", step_metric="ray/ray_step") not in ( + mock_run.define_metric.call_args_list + ) + mock_run.define_metric.assert_any_call("ray/ray_step", hidden=True) + mock_run.define_metric.assert_any_call( + "ray/gpu_utilization", step_metric="ray/ray_step" + ) + assert ( + mock_run.define_metric.call_args_list.count( + call("ray/gpu_utilization", step_metric="ray/ray_step") + ) + == 1 + ) + + @patch("nemo_rl.utils.logger.wandb") + def test_does_not_define_catch_all_metric(self, mock_wandb): + """Overlapping W&B globs must not choose axes nondeterministically.""" + WandbLogger({}) + + mock_run = mock_wandb.init.return_value + assert call("*", step_metric="nemo_rl/step") not in ( + mock_run.define_metric.call_args_list + ) + + @patch("nemo_rl.utils.logger.wandb") + def test_finish_flushes_pending_trainer_row(self, mock_wandb): + """A final incomplete step is not lost during logger teardown.""" + logger = WandbLogger({}) + logger.log_metrics({"loss": 0.5}, step=7, prefix="train") + + mock_run = mock_wandb.init.return_value + mock_run.log.assert_not_called() + logger.finish() + + mock_run.log.assert_called_once_with({"train/loss": 0.5, "nemo_rl/step": 7}) + mock_run.finish.assert_called_once_with() @patch("nemo_rl.utils.logger.wandb") def test_log_hyperparams(self, mock_wandb): From f26f8a4314e80d9183c1814e20d682725fca15ca Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 11 Aug 2026 19:30:16 -0400 Subject: [PATCH 5/6] fix(logger): flush final training metrics Signed-off-by: Anish Mahishi (cherry picked from commit eafe0a56cf8d93d0bf34f7c6c1d4563507800da3) --- nemo_rl/algorithms/distillation.py | 7 ++++++- nemo_rl/algorithms/dpo.py | 7 ++++++- nemo_rl/algorithms/rm.py | 7 ++++++- nemo_rl/algorithms/sft.py | 7 ++++++- nemo_rl/utils/logger.py | 6 ++++++ tests/unit/algorithms/test_distillation.py | 6 ++++++ tests/unit/algorithms/test_dpo.py | 6 ++++++ tests/unit/algorithms/test_rm.py | 6 ++++++ tests/unit/algorithms/test_sft.py | 6 ++++++ .../single_controller/test_sc_checkpointing.py | 4 ++++ tests/unit/utils/test_logger.py | 15 +++++++++++++++ 11 files changed, 73 insertions(+), 4 deletions(-) diff --git a/nemo_rl/algorithms/distillation.py b/nemo_rl/algorithms/distillation.py index 985cb9fd05d..30e093daf69 100644 --- a/nemo_rl/algorithms/distillation.py +++ b/nemo_rl/algorithms/distillation.py @@ -1145,7 +1145,12 @@ def distillation_train( metrics["global_valid_toks"] / total_time / total_num_gpus ) logger.log_metrics(metrics, total_steps + 1, prefix="train") - logger.log_metrics(timing_metrics, total_steps + 1, prefix="timing/train") + logger.log_metrics( + timing_metrics, + total_steps + 1, + prefix="timing/train", + step_finished=True, + ) timer.reset() current_step += 1 diff --git a/nemo_rl/algorithms/dpo.py b/nemo_rl/algorithms/dpo.py index de0a60af265..fc221db6f60 100644 --- a/nemo_rl/algorithms/dpo.py +++ b/nemo_rl/algorithms/dpo.py @@ -792,7 +792,12 @@ def dpo_train( metrics["global_valid_toks"] / total_time / total_num_gpus ) logger.log_metrics(metrics, total_steps + 1, prefix="train") - logger.log_metrics(timing_metrics, total_steps + 1, prefix="timing/train") + logger.log_metrics( + timing_metrics, + total_steps + 1, + prefix="timing/train", + step_finished=True, + ) timer.reset() current_step += 1 diff --git a/nemo_rl/algorithms/rm.py b/nemo_rl/algorithms/rm.py index 0cbe3fc02dc..986ae8e76c7 100644 --- a/nemo_rl/algorithms/rm.py +++ b/nemo_rl/algorithms/rm.py @@ -715,7 +715,12 @@ def rm_train( metrics["global_valid_toks"] / total_time / total_num_gpus ) logger.log_metrics(metrics, total_steps + 1, prefix="train") - logger.log_metrics(timing_metrics, total_steps + 1, prefix="timing/train") + logger.log_metrics( + timing_metrics, + total_steps + 1, + prefix="timing/train", + step_finished=True, + ) timer.reset() current_step += 1 diff --git a/nemo_rl/algorithms/sft.py b/nemo_rl/algorithms/sft.py index b6d96778ef0..bfe50739eeb 100644 --- a/nemo_rl/algorithms/sft.py +++ b/nemo_rl/algorithms/sft.py @@ -651,7 +651,12 @@ def sft_train( else: timing_metrics["valid_tokens_per_sec_per_gpu"] = 0.0 logger.log_metrics(metrics, total_steps + 1, prefix="train") - logger.log_metrics(timing_metrics, total_steps + 1, prefix="timing/train") + logger.log_metrics( + timing_metrics, + total_steps + 1, + prefix="timing/train", + step_finished=True, + ) timer.reset() current_step += 1 diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index 4c392d64cd8..b03513178e3 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -13,6 +13,7 @@ # limitations under the License. +import atexit import glob import json import logging @@ -236,6 +237,11 @@ def __init__(self, cfg: WandbConfig, log_dir: Optional[str] = None): self._pending_step: Optional[int] = None self._pending_metrics: dict[str, Any] = {} self.run.define_metric(_WANDB_CALLER_STEP_METRIC, hidden=True) + # Most training entrypoints run W&B in the driver process and do not + # own an explicit logger teardown today. Register after wandb.init so + # our Python-side row buffer drains before W&B's own atexit handlers. + # finish() is idempotent for actors that already close explicitly. + atexit.register(self.finish) if os.environ.get("RAY_BACKEND_LOG_LEVEL", "").lower() == "debug": print( diff --git a/tests/unit/algorithms/test_distillation.py b/tests/unit/algorithms/test_distillation.py index df9188c02b2..2ad705c8a02 100644 --- a/tests/unit/algorithms/test_distillation.py +++ b/tests/unit/algorithms/test_distillation.py @@ -274,6 +274,12 @@ def test_distillation_train_max_steps(mock_components): ) assert mock_components["student_policy"].train.call_count == 5 + final_timing_call = [ + call + for call in mock_components["logger"].log_metrics.call_args_list + if call.kwargs.get("prefix") == "timing/train" + ][-1] + assert final_timing_call.kwargs["step_finished"] is True def test_ft_save_period_triggers_periodic_saves(mock_components): diff --git a/tests/unit/algorithms/test_dpo.py b/tests/unit/algorithms/test_dpo.py index ffb5bcf6ac5..90705b474b7 100644 --- a/tests/unit/algorithms/test_dpo.py +++ b/tests/unit/algorithms/test_dpo.py @@ -269,6 +269,12 @@ def test_exit_on_max_steps(mock_dpo_components): # Verify we only trained for 12 steps. assert mock_dpo_components["policy"].train.call_count == 12 + final_timing_call = [ + call + for call in mock_dpo_components["logger"].log_metrics.call_args_list + if call.kwargs.get("prefix") == "timing/train" + ][-1] + assert final_timing_call.kwargs["step_finished"] is True def test_exit_on_max_epochs(mock_dpo_components): diff --git a/tests/unit/algorithms/test_rm.py b/tests/unit/algorithms/test_rm.py index 9695b9a7790..7857db825cd 100644 --- a/tests/unit/algorithms/test_rm.py +++ b/tests/unit/algorithms/test_rm.py @@ -240,6 +240,12 @@ def test_exit_on_max_steps(mock_components): # Verify we only trained for 12 steps. assert mock_components["policy"].train.call_count == 12 + final_timing_call = [ + call + for call in mock_components["logger"].log_metrics.call_args_list + if call.kwargs.get("prefix") == "timing/train" + ][-1] + assert final_timing_call.kwargs["step_finished"] is True def test_exit_on_max_epochs(mock_components): diff --git a/tests/unit/algorithms/test_sft.py b/tests/unit/algorithms/test_sft.py index 6d38da3fd47..e7707dd9e00 100644 --- a/tests/unit/algorithms/test_sft.py +++ b/tests/unit/algorithms/test_sft.py @@ -153,6 +153,12 @@ def test_exit_on_max_steps(mock_components): # Verify we only trained for 12 steps. assert mock_components["policy"].train.call_count == 12 + final_timing_call = [ + call + for call in mock_components["logger"].log_metrics.call_args_list + if call.kwargs.get("prefix") == "timing/train" + ][-1] + assert final_timing_call.kwargs["step_finished"] is True def test_exit_on_max_epochs(mock_components): diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py index 21981269454..808cee75e83 100644 --- a/tests/unit/single_controller/test_sc_checkpointing.py +++ b/tests/unit/single_controller/test_sc_checkpointing.py @@ -448,6 +448,10 @@ def group_ids(self) -> tuple[str, ...]: def target_step_list(self) -> tuple[Optional[int], ...]: return tuple(self._target_step_list) + def __len__(self) -> int: + """Match the production TQReplayBuffer occupancy contract.""" + return len(self.target_step_list) + def set_data_plane_checkpoint_barrier( self, barrier: DataPlaneCheckpointBarrier ) -> None: diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py index 766fc3a282a..b1a7bd5b54b 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -514,6 +514,21 @@ def test_does_not_define_catch_all_metric(self, mock_wandb): mock_run.define_metric.call_args_list ) + @patch("nemo_rl.utils.logger.atexit.register") + @patch("nemo_rl.utils.logger.wandb") + def test_registers_teardown_flush(self, mock_wandb, mock_atexit_register): + """Driver entrypoints flush the final pending row at process exit.""" + logger = WandbLogger({}) + logger.log_metrics({"loss": 0.5}, step=7, prefix="train") + + mock_atexit_register.assert_called_once_with(logger.finish) + callback = mock_atexit_register.call_args.args[0] + callback() + + mock_run = mock_wandb.init.return_value + mock_run.log.assert_called_once_with({"train/loss": 0.5, "nemo_rl/step": 7}) + mock_run.finish.assert_called_once_with() + @patch("nemo_rl.utils.logger.wandb") def test_finish_flushes_pending_trainer_row(self, mock_wandb): """A final incomplete step is not lost during logger teardown.""" From a7d417d8bfc0750babbf46f240f42b6a09bc260f Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 19 Aug 2026 17:47:49 -0400 Subject: [PATCH 6/6] feat(rollout): refresh recovery benchmark telemetry (cherry picked from commit af43edd0033ab14b91f7740c33c73101d445f728) --- .../nemo_gym/grpo_qwen3_30ba3b_instruct.yaml | 5 + .../algorithms/async_utils/replay_buffer.py | 173 +++++-- nemo_rl/algorithms/single_controller.py | 483 +++++++++++++++++- .../single_controller_utils/config.py | 5 + .../single_controller_utils/setup.py | 27 + nemo_rl/experience/blackbox_finalizer.py | 6 + nemo_rl/experience/rollout_manager.py | 69 ++- .../data_plane/test_blackbox_finalizer.py | 2 + tests/unit/experience/test_rollout_manager.py | 14 + .../test_finalizer_lifecycle.py | 1 + .../test_sc_checkpointing.py | 119 ++++- 11 files changed, 830 insertions(+), 74 deletions(-) diff --git a/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml b/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml index 397d52023be..93d8aedd5e9 100644 --- a/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml +++ b/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml @@ -176,3 +176,8 @@ checkpointing: # 4. The checkpoint time for this model is around 10 mins. checkpoint_must_save_by: "00:03:30:00" save_period: 1 + +# Optional wall-clock benchmark sampling. This is independent of the rollout +# checkpoint save interval and may be enabled for baseline runs as well. +rollout_checkpointing: + telemetry_interval_s: null diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index ef028bee0f1..c1518c36cb9 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -18,10 +18,12 @@ import math import statistics import threading as _threading +import time import uuid -from collections import Counter +from collections import Counter, deque from collections.abc import AsyncIterator, Mapping from contextlib import asynccontextmanager +from dataclasses import dataclass from numbers import Integral, Real from typing import Any, Iterable, Literal, Optional, TypedDict @@ -42,6 +44,38 @@ REPLAY_BUFFER_METADATA_SCHEMA_VERSION = 1 REPLAY_BUFFER_METADATA_STORAGE: Literal["tq_checkpoint"] = "tq_checkpoint" +CheckpointMutationKind = Literal[ + "group_commits", + "group_removals", + "other", + "prompt_reservations", + "recovery_retries", + "sample_clears", + "sibling_seals", +] +CHECKPOINT_MUTATION_KINDS: tuple[CheckpointMutationKind, ...] = ( + "group_commits", + "group_removals", + "prompt_reservations", + "recovery_retries", + "sample_clears", + "sibling_seals", + "other", +) + + +@dataclass(frozen=True) +class DataPlaneCheckpointBarrierTelemetry: + """Bounded interval telemetry for checkpoint-induced mutation waits.""" + + blocked_by_kind: dict[CheckpointMutationKind, int] + wait_durations_s: tuple[float, ...] + active_mutations: int + waiting_mutations: int + max_waiting_mutations: int + checkpoint_active: bool + + # These TypedDicts describe the versioned, plain-mapping checkpoint wire # format. They are intentionally not dataclass instances: persisting a # dataclass would couple recovery to its Python import path and class layout. @@ -160,6 +194,10 @@ def __init__(self) -> None: self._active_mutations = 0 self._mutation_depth_by_task: dict[asyncio.Task[Any], int] = {} self._mutation_version = 0 + self._waiting_mutations = 0 + self._max_waiting_mutations = 0 + self._blocked_by_kind: Counter[CheckpointMutationKind] = Counter() + self._wait_durations_s: deque[float] = deque(maxlen=10_000) @property def mutation_version(self) -> int: @@ -167,7 +205,9 @@ def mutation_version(self) -> int: return self._mutation_version @asynccontextmanager - async def mutation(self) -> AsyncIterator[None]: + async def mutation( + self, kind: CheckpointMutationKind = "other" + ) -> AsyncIterator[None]: """Enter a commit/clear section, waiting only for an active checkpoint.""" task = asyncio.current_task() if task is None: @@ -184,7 +224,20 @@ async def mutation(self) -> AsyncIterator[None]: self._mutation_depth_by_task[task] -= 1 return async with self._condition: - await self._condition.wait_for(lambda: not self._checkpoint_active) + wait_started: Optional[float] = None + if self._checkpoint_active: + wait_started = time.monotonic() + self._waiting_mutations += 1 + self._max_waiting_mutations = max( + self._max_waiting_mutations, self._waiting_mutations + ) + try: + await self._condition.wait_for(lambda: not self._checkpoint_active) + finally: + if wait_started is not None: + self._waiting_mutations -= 1 + self._blocked_by_kind[kind] += 1 + self._wait_durations_s.append(time.monotonic() - wait_started) self._active_mutations += 1 self._mutation_depth_by_task[task] = 1 try: @@ -200,6 +253,25 @@ async def mutation(self) -> AsyncIterator[None]: if self._active_mutations == 0: self._condition.notify_all() + async def drain_telemetry(self) -> DataPlaneCheckpointBarrierTelemetry: + """Return and reset interval waits while preserving current state.""" + async with self._condition: + telemetry = DataPlaneCheckpointBarrierTelemetry( + blocked_by_kind={ + kind: self._blocked_by_kind.get(kind, 0) + for kind in CHECKPOINT_MUTATION_KINDS + }, + wait_durations_s=tuple(self._wait_durations_s), + active_mutations=self._active_mutations, + waiting_mutations=self._waiting_mutations, + max_waiting_mutations=self._max_waiting_mutations, + checkpoint_active=self._checkpoint_active, + ) + self._blocked_by_kind.clear() + self._wait_durations_s.clear() + self._max_waiting_mutations = self._waiting_mutations + return telemetry + @asynccontextmanager async def checkpoint(self) -> AsyncIterator[None]: """Block new mutations and wait for active ones before snapshotting.""" @@ -970,7 +1042,7 @@ async def commit( weight_version=start_weight_version, ) trace_rollout_payload(keys=sample_ids, data=train_batch) - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation("group_commits"): try: await call_data_plane( self._dp_client, @@ -1037,7 +1109,7 @@ async def remove_group(self, group_id: str, *, remove_in_dp: bool = False) -> in "TQReplayBuffer must be bound to the controller data-plane " "checkpoint barrier before removing a group" ) - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation("group_removals"): try: idx = self._group_ids.index(group_id) except ValueError as error: @@ -1058,7 +1130,7 @@ async def clear_staging_keys(self, staging_keys: list[str]) -> None: "checkpoint barrier before clearing staging samples" ) unique_keys = list(dict.fromkeys(staging_keys)) - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation("sample_clears"): await call_data_plane( self._dp_client, "clear_samples", @@ -1095,45 +1167,56 @@ async def commit_finalized( Raises: ValueError: group_id has no live slot (removed or never reserved). """ - try: - idx = self._group_ids.index(group_id) - except ValueError: - raise ValueError( - f"TQReplayBuffer.commit_finalized: group {group_id} has no " - "live slot (evicted or never reserved)" - ) from None - tagged_plans = [ - tag[ROUTE_PLAN_TAG] for tag in (meta.tags or []) if ROUTE_PLAN_TAG in tag - ] - if tagged_plans: - if len(tagged_plans) != len(meta.sample_ids): - raise ValueError( - "commit_finalized received mixed deferred/direct route plans" - ) - from nemo_rl.experience.route_plan import decode_route_plan + if self._data_plane_checkpoint_barrier is None: + raise RuntimeError( + "TQReplayBuffer must be bound to the controller data-plane " + "checkpoint barrier before committing finalized samples" + ) - plan_cleanup_keys = { - key - for encoded in tagged_plans - for key in decode_route_plan(encoded).cleanup_staging_keys - } - provided_staging_keys = list(staging_keys or []) - if len(provided_staging_keys) != len(set(provided_staging_keys)): - raise ValueError("commit_finalized staging_keys contains duplicates") - if set(provided_staging_keys) != plan_cleanup_keys: + async with self._data_plane_checkpoint_barrier.mutation("group_commits"): + try: + idx = self._group_ids.index(group_id) + except ValueError: raise ValueError( - "commit_finalized staging ownership does not match route plans: " - f"provided={sorted(provided_staging_keys)!r}, " - f"planned={sorted(plan_cleanup_keys)!r}" - ) - self.meta_list[idx] = meta - self.start_weight_list[idx] = group_min_wv - self.end_weight_list[idx] = group_max_wv - self.ready_list[idx] = True - self._staging_keys_list[idx] = ( - list(staging_keys) if staging_keys is not None else None - ) - return meta + f"TQReplayBuffer.commit_finalized: group {group_id} has no " + "live slot (evicted or never reserved)" + ) from None + tagged_plans = [ + tag[ROUTE_PLAN_TAG] + for tag in (meta.tags or []) + if ROUTE_PLAN_TAG in tag + ] + if tagged_plans: + if len(tagged_plans) != len(meta.sample_ids): + raise ValueError( + "commit_finalized received mixed deferred/direct route plans" + ) + from nemo_rl.experience.route_plan import decode_route_plan + + plan_cleanup_keys = { + key + for encoded in tagged_plans + for key in decode_route_plan(encoded).cleanup_staging_keys + } + provided_staging_keys = list(staging_keys or []) + if len(provided_staging_keys) != len(set(provided_staging_keys)): + raise ValueError( + "commit_finalized staging_keys contains duplicates" + ) + if set(provided_staging_keys) != plan_cleanup_keys: + raise ValueError( + "commit_finalized staging ownership does not match route plans: " + f"provided={sorted(provided_staging_keys)!r}, " + f"planned={sorted(plan_cleanup_keys)!r}" + ) + self.meta_list[idx] = meta + self.start_weight_list[idx] = group_min_wv + self.end_weight_list[idx] = group_max_wv + self.ready_list[idx] = True + self._staging_keys_list[idx] = ( + list(staging_keys) if staging_keys is not None else None + ) + return meta def abort(self, group_id: str) -> bool: """Drop an unready slot whose dispatch failed or was cancelled. @@ -1188,7 +1271,7 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: "TQReplayBuffer must be bound to the controller data-plane " "checkpoint barrier before removing groups" ) - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation("group_removals"): drop_idxs = sorted(idxs, reverse=True) if drop_idxs[0] >= len(self.meta_list): raise IndexError( @@ -1478,7 +1561,7 @@ async def _clear_samples(self, *, sample_ids: list[str]) -> None: "TQReplayBuffer must be bound to the controller data-plane " "checkpoint barrier before clearing samples" ) - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation("sample_clears"): await self._clear_samples_unlocked(sample_ids=sample_ids) async def _clear_samples_unlocked(self, *, sample_ids: list[str]) -> None: diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 61d604c78ce..70f8c0cfaba 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -38,11 +38,13 @@ import hashlib import io import json +import math import os import shutil import statistics import time import warnings +from collections import deque from dataclasses import dataclass from functools import partial from pathlib import Path @@ -57,6 +59,7 @@ LEGACY_REPLAY_BUFFER_FILENAME, REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + CHECKPOINT_MUTATION_KINDS, DataPlaneCheckpointBarrier, TQReplayGroupMetadata, TQReplayMetadataState, @@ -119,6 +122,14 @@ Generation = Union[VllmGeneration, SGLangGeneration] +_TELEMETRY_WALL_TIME_METRIC = "telemetry/wall_time_seconds" +_TELEMETRY_PREFIXES = ( + "timing/rollout_checkpoint", + "timing/rollout_recovery", + "rollout/checkpoint_outcome", + "rollout/throughput", +) + @dataclass(frozen=True) class _RolloutWorkItem: @@ -139,6 +150,31 @@ class _RolloutCheckpointCut: rollout_recovery_group_count: Optional[int] rolled_back_train_group_count: int mutation_version: int + tq_save_seconds: float + + +def _latest_vllm_values(metrics: dict[str, Any], metric_name: str) -> dict[int, float]: + """Return the latest value keyed by vLLM data-parallel worker.""" + per_worker = metrics.get(metric_name) + if not isinstance(per_worker, dict): + return {} + return { + worker_id: float(values[-1]) + for worker_id, values in per_worker.items() + if isinstance(worker_id, int) + and isinstance(values, list) + and values + and isinstance(values[-1], (int, float)) + } + + +def _percentile(values: list[float], quantile: float) -> float: + """Return a deterministic nearest-rank percentile for telemetry.""" + if not values: + raise ValueError("percentile requires at least one value") + ordered = sorted(values) + index = max(0, min(len(ordered) - 1, math.ceil(quantile * len(ordered)) - 1)) + return ordered[index] @ray.remote(num_cpus=1, num_gpus=0) # pragma: no cover @@ -225,6 +261,26 @@ def __init__( setup_timing_metrics.to_metrics_dict(), step=0, prefix="timing/setup" ) self._timer = Timer() + self._telemetry_started_at = time.monotonic() + self._telemetry_sample_index = 0 + try: + for prefix in _TELEMETRY_PREFIXES: + self._logger.define_metric( + f"{prefix}/*", step_metric=_TELEMETRY_WALL_TIME_METRIC + ) + except Exception as error: + warnings.warn( + "Failed to define rollout benchmark telemetry series: " + f"{type(error).__name__}: {error}", + stacklevel=2, + ) + self._throughput_sample_time: Optional[float] = None + self._throughput_generation_tokens_by_worker: Optional[dict[int, int]] = None + self._throughput_rollout_counters: Optional[dict[str, int]] = None + self._rollout_telemetry_lock = asyncio.Lock() + self._rollout_completion_durations_s: deque[float] = deque(maxlen=10_000) + self._rollout_queue_wait_durations_s: deque[float] = deque(maxlen=10_000) + self._last_successful_rollout_checkpoint_time: Optional[float] = None # Also built here, not on the driver: TimeoutChecker must capture # wall-clock start times inside the actor, not at driver setup time. @@ -244,6 +300,9 @@ def __init__( self._data_plane_checkpoint_metadata: Optional[dict[str, Any]] = ( actor_args.data_plane_checkpoint_metadata ) + self._rollout_checkpoint_load_metrics = ( + actor_args.rollout_checkpoint_load_metrics + ) self._consumed_samples: int = actor_args.save_state.consumed_samples self._total_valid_tokens: int = actor_args.save_state.total_valid_tokens @@ -320,6 +379,9 @@ def __init__( # Count of in-flight generate_and_push calls self._inflight_rollouts: int = 0 + self._rollout_slot_waiters: int = 0 + self._rollout_permitted_waiters: int = 0 + self._buffer_capacity_waiters: int = 0 # Cancellation handles for in-flight rollout dispatches. self._dispatched_rollouts: set[asyncio.Task[None]] = set() @@ -364,10 +426,32 @@ async def run(self) -> dict[str, Any]: # Synchronize weights before starting the pumps await self._sync_weights() + replay_restore_started = time.monotonic() restored_replay_groups = await self._maybe_restore_replay_buffer() + replay_restore_seconds = time.monotonic() - replay_restore_started + recovery_prepare_started = time.monotonic() recovery_groups = await self._prepare_rollout_recovery( restored_replay_groups=restored_replay_groups ) + recovery_prepare_seconds = time.monotonic() - recovery_prepare_started + if self._rollout_checkpoint_load_metrics is not None: + load_metrics = dict(self._rollout_checkpoint_load_metrics) + load_metrics["replay_metadata_load_seconds"] = replay_restore_seconds + load_metrics["recovery_prepare_seconds"] = recovery_prepare_seconds + load_metrics["total_load_seconds"] = sum(load_metrics.values()) + self._log_telemetry_metrics( + load_metrics, + step=self._train_steps, + prefix="timing/rollout_checkpoint", + ) + + telemetry_interval_s = ( + self._master_config.rollout_checkpointing.telemetry_interval_s + ) + if telemetry_interval_s is not None: + # Establish a process-local rate baseline before restored and fresh + # work begin competing for the same generation capacity. + await self._log_rollout_throughput_metrics(emit=False) # Start restored work and normal pumps together. Preparation above has # already reconstructed sampler state and reserved restored capacity, @@ -385,11 +469,18 @@ async def run(self) -> dict[str, Any]: if self._master_config.rollout_checkpointing.interval_s is not None else None ) + rollout_telemetry_task = ( + asyncio.create_task(self._rollout_telemetry_pump()) + if telemetry_interval_s is not None + else None + ) tasks = {rollout_task, train_task, watchdog_task} if recovery_task is not None: tasks.add(recovery_task) if rollout_checkpoint_task is not None: tasks.add(rollout_checkpoint_task) + if rollout_telemetry_task is not None: + tasks.add(rollout_telemetry_task) try: pending = set(tasks) while pending: @@ -417,6 +508,11 @@ async def run(self) -> dict[str, Any]: raise RuntimeError( "rollout checkpoint pump exited without requesting stop" ) + if ( + rollout_telemetry_task is not None + and rollout_telemetry_task in done + ): + raise RuntimeError("rollout telemetry pump exited unexpectedly") finally: for task in tasks: task.cancel() @@ -455,6 +551,41 @@ async def ping(self) -> dict[str, Any]: "finalizer_unknown_outcomes": self._finalizer_unknown_outcomes, } + def _log_telemetry_metrics( + self, metrics: dict[str, float], *, step: int, prefix: str + ) -> None: + """Log benchmark telemetry on an axis independent of trainer steps.""" + try: + self._telemetry_sample_index += 1 + event_metrics = dict(metrics) + event_metrics["train_step"] = float(step) + event_metrics["sample_index"] = float(self._telemetry_sample_index) + event_metrics[_TELEMETRY_WALL_TIME_METRIC] = ( + time.monotonic() - self._telemetry_started_at + ) + self._logger.log_metrics( + event_metrics, + step=self._telemetry_sample_index, + prefix=prefix, + step_metric=_TELEMETRY_WALL_TIME_METRIC, + ) + except Exception as error: + warnings.warn( + f"Failed to log {prefix} telemetry: {type(error).__name__}: {error}", + stacklevel=2, + ) + + def _record_rollout_timing( + self, *, work_started: float, dispatch_started: float + ) -> None: + """Record one committed group's queue and execution durations.""" + if not hasattr(self, "_rollout_queue_wait_durations_s"): + self._rollout_queue_wait_durations_s = deque(maxlen=10_000) + if not hasattr(self, "_rollout_completion_durations_s"): + self._rollout_completion_durations_s = deque(maxlen=10_000) + self._rollout_queue_wait_durations_s.append(dispatch_started - work_started) + self._rollout_completion_durations_s.append(time.monotonic() - dispatch_started) + # ── internal helpers ─────────────────────────────────────────────────── async def _maybe_restore_replay_buffer(self) -> int: @@ -648,7 +779,7 @@ async def _cleanup_known_finalization_request( self, request: "FinalizationRequest" ) -> None: """Clear known request ownership without racing a TQ checkpoint.""" - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation("group_removals"): await self._cleanup_known_finalization_request_unlocked(request) async def _finalize_with_actor(self, request: "FinalizationRequest") -> None: @@ -678,7 +809,7 @@ async def _finalize_with_actor(self, request: "FinalizationRequest") -> None: # the matching replay-index transition. A checkpoint therefore # observes either the complete group or neither half of it while # tensor payloads remain outside the SingleController process. - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation("group_commits"): ledger = self._rollout_recovery_ledger assert ledger is not None ledger.mark_finalization_started(request.group_id) @@ -746,6 +877,9 @@ async def _finalize_with_actor(self, request: "FinalizationRequest") -> None: group_min_weight_version=finalized.group_min_wv, group_max_weight_version=finalized.group_max_wv, ) + self._rollout_manager.record_canonical_publication( + finalized.canonical_output_tokens + ) finally: self._active_finalizers -= 1 # Cancellation before RPC submission leaves the actor untouched; @@ -818,7 +952,7 @@ async def _cleanup_consumed_metas_unlocked(self, metas: list[KVBatchMeta]) -> No async def _cleanup_consumed_metas(self, metas: list[KVBatchMeta]) -> None: """Clear consumed rows without racing a native TQ checkpoint.""" - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation("sample_clears"): await self._cleanup_consumed_metas_unlocked(metas) async def _validate_replay_inventory( @@ -1034,10 +1168,22 @@ async def _recover_one_prepared_group( recovery_record = ledger.get_group(group.group_id) reused_siblings = len(recovery_record.sealed_generation_indices) redispatched_siblings = recovery_record.expected_generations - reused_siblings + work_started = time.monotonic() try: - await self._rollout_slots.acquire() + self._rollout_slot_waiters = getattr(self, "_rollout_slot_waiters", 0) + 1 + try: + await self._rollout_slots.acquire() + finally: + self._rollout_slot_waiters -= 1 rollout_slot_owned = True - await self._rollout_permitted.wait() + self._rollout_permitted_waiters = ( + getattr(self, "_rollout_permitted_waiters", 0) + 1 + ) + try: + await self._rollout_permitted.wait() + finally: + self._rollout_permitted_waiters -= 1 + dispatch_started = time.monotonic() self._inflight_rollouts += 1 counted_inflight = True @@ -1056,6 +1202,14 @@ async def _recover_one_prepared_group( if request is None: return False await self._finalize_with_actor(request) + self._rollout_manager.record_recovery_siblings( + reused=reused_siblings, + redispatched=redispatched_siblings, + ) + self._record_rollout_timing( + work_started=work_started, + dispatch_started=dispatch_started, + ) print( "rollout recovery finalized group: " f"group={group.group_id} reused={reused_siblings} " @@ -1076,6 +1230,8 @@ async def _recover_prepared_rollout_groups( self, groups: list[PromptGroupRecoverySummary] ) -> None: """Recover restored groups concurrently beside fresh rollout dispatch.""" + recovery_started = time.monotonic() + counters_before = self._rollout_manager.telemetry_snapshot() recovery_tasks: list[asyncio.Task[bool]] = [] try: async with asyncio.TaskGroup() as task_group: @@ -1103,6 +1259,25 @@ async def _recover_prepared_rollout_groups( f"dropped={len(groups) - completed}", flush=True, ) + counters_after = self._rollout_manager.telemetry_snapshot() + self._log_telemetry_metrics( + { + "total_recovery_seconds": time.monotonic() - recovery_started, + "groups_considered": float(len(groups)), + "groups_finalized": float(completed), + "groups_dropped": float(len(groups) - completed), + "siblings_reused": float( + counters_after["recovery_siblings_reused"] + - counters_before["recovery_siblings_reused"] + ), + "siblings_redispatched": float( + counters_after["recovery_siblings_redispatched"] + - counters_before["recovery_siblings_redispatched"] + ), + }, + step=self._train_steps, + prefix="timing/rollout_recovery", + ) finally: self._rollout_recovery_complete.set() @@ -1154,7 +1329,7 @@ def _restore_sampler_dispatch_state( async def _clear_data_plane_samples(self, sample_ids: list[str]) -> None: """Clear consumed rows without overlapping a data-plane checkpoint.""" - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation("sample_clears"): await call_data_plane( self._dp_client, "clear_samples", @@ -1314,12 +1489,14 @@ async def _capture_rollout_checkpoint_cut( await self._validate_replay_inventory(replay_metadata) if rollout_recovery_payload is not None: await self._validate_rollout_recovery_inventory(clear_unreferenced=False) + tq_save_started = time.monotonic() await self._save_data_plane_checkpoint( checkpoint_path, replay_metadata=replay_metadata, rollout_recovery_payload_sha256=rollout_recovery_digest, rollout_recovery_group_count=rollout_recovery_group_count, ) + tq_save_seconds = time.monotonic() - tq_save_started return _RolloutCheckpointCut( dataloader_state=dataloader_state, replay_metadata=replay_metadata, @@ -1327,6 +1504,7 @@ async def _capture_rollout_checkpoint_cut( rollout_recovery_group_count=rollout_recovery_group_count, rolled_back_train_group_count=len(additional_groups), mutation_version=self._data_plane_checkpoint_barrier.mutation_version, + tq_save_seconds=tq_save_seconds, ) async def _write_rollout_checkpoint_sidecars( @@ -1375,6 +1553,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: ): return False + save_started = time.monotonic() await asyncio.to_thread(self._checkpointer.finalize_pending) if self._train_steps == 0: if self._trainer_version != 0: @@ -1425,11 +1604,13 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: expected_train_step = self._train_steps expected_trainer_version = self._trainer_version - tmp_path, final_path, _ = await asyncio.to_thread( + tmp_path, final_path, snapshot_sequence = await asyncio.to_thread( prepare_snapshot_paths, anchor ) try: + barrier_requested = time.monotonic() async with self._data_plane_checkpoint_barrier.checkpoint(): + barrier_acquired = time.monotonic() if ( self._optimizer_commit_in_progress or self._train_steps != expected_train_step @@ -1442,6 +1623,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: tmp_path, periodic=True, ) + barrier_released = time.monotonic() await self._write_rollout_checkpoint_sidecars( tmp_path, @@ -1476,15 +1658,65 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: self._last_rollout_snapshot_mutation_version = cut.mutation_version self._last_missing_rollout_snapshot_anchor = None + save_completed = time.monotonic() + checkpoint_metrics = { + "snapshot_sequence": float(snapshot_sequence), + "total_save_seconds": save_completed - save_started, + "tq_save_seconds": cut.tq_save_seconds, + "barrier_wait_seconds": barrier_acquired - barrier_requested, + "exclusive_hold_seconds": barrier_released - barrier_acquired, + "replay_groups": float( + len(cut.replay_metadata["groups"]) + if cut.replay_metadata is not None + else 0 + ), + "ledger_groups": float(cut.rollout_recovery_group_count or 0), + "rolled_back_train_groups": float(cut.rolled_back_train_group_count), + } + self._log_telemetry_metrics( + checkpoint_metrics, + step=expected_train_step, + prefix="timing/rollout_checkpoint", + ) print( "rollout checkpoint save completed: " f"{final_path} (step={expected_train_step}, " f"trainer_version={expected_trainer_version}, " - f"ledger_groups={cut.rollout_recovery_group_count or 0})", + f"ledger_groups={cut.rollout_recovery_group_count or 0}, " + f"total_save_seconds={checkpoint_metrics['total_save_seconds']:.2f}, " + "exclusive_hold_seconds=" + f"{checkpoint_metrics['exclusive_hold_seconds']:.2f})", flush=True, ) return True + def _log_rollout_checkpoint_outcome( + self, *, outcome: str, attempt_duration_seconds: float + ) -> None: + """Record every scheduled checkpoint attempt, including no-op cuts.""" + metrics = { + "attempt": 1.0, + "attempt_duration_seconds": attempt_duration_seconds, + "completed": float(outcome == "completed"), + "skipped": float(outcome == "skipped"), + "failed": float(outcome == "failed"), + "configured_interval_seconds": float( + self._master_config.rollout_checkpointing.interval_s or 0.0 + ), + } + if outcome == "completed": + completed_at = time.monotonic() + if self._last_successful_rollout_checkpoint_time is not None: + metrics["seconds_since_previous_success"] = ( + completed_at - self._last_successful_rollout_checkpoint_time + ) + self._last_successful_rollout_checkpoint_time = completed_at + self._log_telemetry_metrics( + metrics, + step=self._train_steps, + prefix="rollout/checkpoint_outcome", + ) + async def _rollout_checkpoint_pump(self) -> None: """Persist rollout state periodically, including during streamed train.""" interval_s = self._master_config.rollout_checkpointing.interval_s @@ -1492,10 +1724,15 @@ async def _rollout_checkpoint_pump(self) -> None: raise RuntimeError("rollout checkpoint pump started while disabled") while True: await asyncio.sleep(interval_s) + attempt_started = time.monotonic() deadline_due = self._train_steps == 0 and self._timeout.would_save() try: saved = await self._save_rollout_checkpoint(force=deadline_due) except Exception as error: + self._log_rollout_checkpoint_outcome( + outcome="failed", + attempt_duration_seconds=time.monotonic() - attempt_started, + ) if deadline_due: raise RuntimeError( "failed to save the required pre-step rollout checkpoint" @@ -1506,6 +1743,10 @@ async def _rollout_checkpoint_pump(self) -> None: stacklevel=2, ) continue + self._log_rollout_checkpoint_outcome( + outcome="completed" if saved else "skipped", + attempt_duration_seconds=time.monotonic() - attempt_started, + ) if deadline_due: if not saved: continue @@ -1519,6 +1760,15 @@ async def _rollout_checkpoint_pump(self) -> None: self._rollout_checkpoint_stop_requested.set() return + async def _rollout_telemetry_pump(self) -> None: + """Sample generation and publication throughput at a fixed cadence.""" + interval_s = self._master_config.rollout_checkpointing.telemetry_interval_s + if interval_s is None: + raise RuntimeError("rollout telemetry pump started while disabled") + while True: + await asyncio.sleep(interval_s) + await self._log_rollout_throughput_metrics() + @staticmethod def _group_ids_from_meta(meta: KVBatchMeta) -> list[str]: """Return stable prompt-group IDs in their canonical sample order.""" @@ -1594,15 +1844,29 @@ async def _dispatch_one_prompt( dispatch_admitted_event: asyncio.Event, ) -> None: task_started_event.set() + work_started = time.monotonic() generation_permit_released = False inflight_count_released = False counted_inflight = False ownership_transferred = False rollout_slot_owned = False try: - await self._rollout_slots.acquire() + self._rollout_slot_waiters = ( + getattr(self, "_rollout_slot_waiters", 0) + 1 + ) + try: + await self._rollout_slots.acquire() + finally: + self._rollout_slot_waiters -= 1 rollout_slot_owned = True - await self._rollout_permitted.wait() + self._rollout_permitted_waiters = ( + getattr(self, "_rollout_permitted_waiters", 0) + 1 + ) + try: + await self._rollout_permitted.wait() + finally: + self._rollout_permitted_waiters -= 1 + dispatch_started = time.monotonic() dispatch_admitted_event.set() self._inflight_rollouts += 1 counted_inflight = True @@ -1656,6 +1920,11 @@ async def _dispatch_one_prompt( self._buffer_capacity.release() return + self._record_rollout_timing( + work_started=work_started, + dispatch_started=dispatch_started, + ) + if self._async_cfg.diagnostics: content = "" for i in range(len(prompt["message_log"])): @@ -1682,7 +1951,13 @@ async def _acquire_capacity(permits: int) -> None: acquired = 0 try: for _ in range(permits): - await self._buffer_capacity.acquire() + self._buffer_capacity_waiters = ( + getattr(self, "_buffer_capacity_waiters", 0) + 1 + ) + try: + await self._buffer_capacity.acquire() + finally: + self._buffer_capacity_waiters -= 1 acquired += 1 except BaseException: _release_capacity(acquired) @@ -1729,7 +2004,9 @@ def _num_prompts_to_dispatch( # The cursor and all durable prompt reservations move # together. A snapshot therefore cannot skip a batch # whose lineage was not yet made recoverable. - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation( + "prompt_reservations" + ): try: prompt_batch = next(dataloader_iterator) except StopIteration: @@ -1808,7 +2085,13 @@ def _num_prompts_to_dispatch( for work_item in rollout_work_items: capacity_acquired_here = False if not token_capture_enabled: - await self._buffer_capacity.acquire() + self._buffer_capacity_waiters = ( + getattr(self, "_buffer_capacity_waiters", 0) + 1 + ) + try: + await self._buffer_capacity.acquire() + finally: + self._buffer_capacity_waiters -= 1 capacity_acquired_here = True task_started_event = asyncio.Event() dispatch_admitted_event = asyncio.Event() @@ -1890,7 +2173,9 @@ async def _train_pump(self) -> None: current_train_weight=self._trainer_version, ) else: - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation( + "group_removals" + ): group_ids_before_evict = set(self._buffer.group_ids) evicted = await self._sampler.evict( current_train_weight=self._trainer_version, @@ -1930,7 +2215,9 @@ async def _train_pump(self) -> None: # Selection removes local replay ownership. Keep that # removal and the matching ledger claim in one mutation # boundary so a later periodic checkpoint cannot split them. - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation( + "group_removals" + ): train_meta, num_groups = await self._sampler.select( current_train_weight=self._trainer_version, min_prompt_groups=min_prompt_groups, @@ -2053,7 +2340,9 @@ async def _train_pump(self) -> None: result = await asyncio.to_thread(self._trainer.finish_train_step) self._optimizer_commit_in_progress = True - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation( + "sample_clears" + ): if self._rollout_recovery_ledger is not None: self._rollout_recovery_ledger.mark_train_step_applied( self._train_steps @@ -2164,6 +2453,11 @@ async def _train_pump(self) -> None: self._logger.log_metrics( timing_metrics, step=self._train_steps, prefix="timing/train" ) + if ( + self._master_config.rollout_checkpointing.telemetry_interval_s + is not None + ): + await self._log_rollout_throughput_metrics() self._timer.reset() # min sample version refers to the version each consumed sample was @@ -2540,6 +2834,163 @@ async def _sync_weights( self._rollout_permitted.set() return aborted_stale_inflight_groups + async def _log_rollout_throughput_metrics(self, *, emit: bool = True) -> None: + """Serialize backend sampling across telemetry and train-step paths.""" + async with self._rollout_telemetry_lock: + await self._collect_and_log_rollout_throughput_metrics(emit=emit) + + async def _collect_and_log_rollout_throughput_metrics( + self, *, emit: bool = True + ) -> None: + """Compare generation-engine output with canonical TQ publication.""" + now = time.monotonic() + counters = self._rollout_manager.telemetry_snapshot() + barrier = await self._data_plane_checkpoint_barrier.drain_telemetry() + generation_metrics: dict[str, Any] = {} + try: + generation_metrics = await asyncio.to_thread( + self._gen.drain_latest_logger_metrics + ) + except Exception as error: + warnings.warn( + "Failed to collect generation throughput telemetry: " + f"{type(error).__name__}: {error}", + stacklevel=2, + ) + + generated_values = _latest_vllm_values(generation_metrics, "generation_tokens") + generated_by_worker = ( + {worker_id: int(value) for worker_id, value in generated_values.items()} + if generated_values + else None + ) + generated_tokens = ( + sum(generated_by_worker.values()) + if generated_by_worker is not None + else None + ) + metrics: dict[str, float] = { + key: float(value) for key, value in counters.items() + } + metrics.update( + { + "controller_dispatched_tasks": float(len(self._dispatched_rollouts)), + "controller_inflight_rollouts": float(self._inflight_rollouts), + "controller_rollout_capacity": float( + self._async_cfg.max_inflight_prompts + ), + "controller_rollout_slot_waiters": float(self._rollout_slot_waiters), + "controller_rollout_permitted_waiters": float( + self._rollout_permitted_waiters + ), + "controller_buffer_capacity_waiters": float( + self._buffer_capacity_waiters + ), + "checkpoint_barrier_active": float(barrier.checkpoint_active), + "checkpoint_active_mutations": float(barrier.active_mutations), + "checkpoint_waiting_mutations": float(barrier.waiting_mutations), + "checkpoint_max_simultaneous_waiters": float( + barrier.max_waiting_mutations + ), + "checkpoint_blocked_mutations": float( + sum(barrier.blocked_by_kind.values()) + ), + "buffer_occupancy_groups": float(len(self._buffer)), + } + ) + for kind in CHECKPOINT_MUTATION_KINDS: + metrics[f"checkpoint_blocked_{kind}"] = float(barrier.blocked_by_kind[kind]) + mutation_waits = list(barrier.wait_durations_s) + if mutation_waits: + metrics["checkpoint_mutation_wait_seconds_total"] = sum(mutation_waits) + metrics["checkpoint_mutation_wait_seconds_mean"] = sum( + mutation_waits + ) / len(mutation_waits) + metrics["checkpoint_mutation_wait_seconds_p95"] = _percentile( + mutation_waits, 0.95 + ) + metrics["checkpoint_mutation_wait_seconds_max"] = max(mutation_waits) + + running = _latest_vllm_values(generation_metrics, "inflight_batch_sizes") + waiting = _latest_vllm_values(generation_metrics, "num_pending_samples") + kv_usage = _latest_vllm_values(generation_metrics, "kv_cache_usage_perc") + if running: + metrics["vllm_requests_running"] = sum(running.values()) + if waiting: + metrics["vllm_requests_waiting"] = sum(waiting.values()) + if kv_usage: + metrics["vllm_kv_cache_usage_mean"] = sum(kv_usage.values()) / len(kv_usage) + if generated_tokens is not None: + metrics["vllm_output_tokens"] = float(generated_tokens) + + completion_durations = list(self._rollout_completion_durations_s) + queue_wait_durations = list(self._rollout_queue_wait_durations_s) + self._rollout_completion_durations_s.clear() + self._rollout_queue_wait_durations_s.clear() + if completion_durations: + metrics["group_completion_samples"] = float(len(completion_durations)) + metrics["group_completion_seconds_mean"] = sum(completion_durations) / len( + completion_durations + ) + metrics["group_completion_seconds_p50"] = _percentile( + completion_durations, 0.50 + ) + metrics["group_completion_seconds_p95"] = _percentile( + completion_durations, 0.95 + ) + if queue_wait_durations: + metrics["group_queue_wait_seconds_mean"] = sum(queue_wait_durations) / len( + queue_wait_durations + ) + metrics["group_queue_wait_seconds_p95"] = _percentile( + queue_wait_durations, 0.95 + ) + + previous_time = self._throughput_sample_time + previous_counters = self._throughput_rollout_counters + previous_generated = self._throughput_generation_tokens_by_worker + if previous_time is not None and previous_counters is not None: + elapsed = now - previous_time + if elapsed > 0: + metrics["sample_elapsed_seconds"] = elapsed + metrics["canonical_groups_per_second"] = ( + counters["canonical_groups_finalized"] + - previous_counters["canonical_groups_finalized"] + ) / elapsed + metrics["canonical_output_tokens_per_second"] = ( + counters["canonical_output_tokens"] + - previous_counters["canonical_output_tokens"] + ) / elapsed + if generated_by_worker is not None and previous_generated is not None: + current_workers = set(generated_by_worker) + previous_workers = set(previous_generated) + monotonic = all( + generated_by_worker[worker_id] >= previous_generated[worker_id] + for worker_id in current_workers & previous_workers + ) + if current_workers == previous_workers and monotonic: + generated_delta = sum( + generated_by_worker[worker_id] + - previous_generated[worker_id] + for worker_id in current_workers + ) + metrics["vllm_output_tokens_per_second"] = ( + generated_delta / elapsed + ) + else: + metrics["vllm_counter_discontinuity"] = 1.0 + + self._throughput_sample_time = now + self._throughput_generation_tokens_by_worker = generated_by_worker + self._throughput_rollout_counters = counters + if emit: + self._log_telemetry_metrics( + metrics, + step=self._train_steps, + prefix="rollout/throughput", + ) + print(f"rollout_throughput_metrics={metrics}", flush=True) + async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: """Fetch advantage inputs, compute advantages, and write them back. diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index db9e777763d..a766e8675fb 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -315,9 +315,14 @@ class RolloutCheckpointConfig(BaseModel, extra="allow"): SingleController has no validation loop, so checkpoint selection must use ``checkpointing.metric_name=None`` or a ``train:`` metric. Inherited ``val:`` settings are rejected during setup. + + ``telemetry_interval_s=None`` disables the independent wall-clock sampler + for rollout/checkpoint benchmark metrics. It does not enable checkpointing + and may be configured without ``interval_s``. """ interval_s: Optional[float] = Field(default=None, gt=0) + telemetry_interval_s: Optional[float] = Field(default=None, gt=0) keep_latest_k: int = Field(default=2, ge=1) restore_mode: Literal["latest", "trainer_checkpoint", "none"] = "latest" diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 7762b58543c..90ff8eae493 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -137,6 +137,7 @@ class SingleControllerActorArgs: save_state: GRPOSaveState last_checkpoint_path: Optional[str] data_plane_checkpoint_metadata: Optional[dict[str, Any]] = None + rollout_checkpoint_load_metrics: Optional[dict[str, float]] = None bootstrap_fingerprint: Optional[str] = None @@ -924,6 +925,16 @@ def setup_single_controller( "dataloader recovery.", flush=True, ) + recovery_path = ( + Path(recovery_checkpoint_path) if recovery_checkpoint_path is not None else None + ) + has_rollout_checkpoint_payload = recovery_path is not None and ( + (recovery_path / REPLAY_BUFFER_METADATA_FILENAME).is_file() + or (recovery_path / ROLLOUT_RECOVERY_STATE_FILENAME).is_file() + ) + rollout_checkpoint_load_metrics: Optional[dict[str, float]] = ( + {} if has_rollout_checkpoint_payload else None + ) # ========================== # Setup Dataset & Environments @@ -960,7 +971,12 @@ def setup_single_controller( print( f"📦 Restoring dataloader state from checkpoint: {recovery_checkpoint_path}" ) + dataloader_load_started = time.monotonic() load_dataloader_state(dataloader, recovery_checkpoint_path, data_config) + if rollout_checkpoint_load_metrics is not None: + rollout_checkpoint_load_metrics["dataloader_load_seconds"] = ( + time.monotonic() - dataloader_load_started + ) _clamp_max_num_steps(master_config, dataloader) _maybe_inject_megatron_train_iters(master_config) @@ -1082,6 +1098,7 @@ def _build_generation_then_trainer( # Native TQ restore must run through the trainer's bootstrap client before # the normal SC data-plane client is created or any rollout/train data-plane # operation starts. + data_plane_load_started = time.monotonic() data_plane_checkpoint_metadata = _maybe_restore_native_data_plane_checkpoint( trainer, last_checkpoint_path=recovery_checkpoint_path, @@ -1089,12 +1106,21 @@ def _build_generation_then_trainer( partition_id=partition_id, sampler_name=master_config.async_rl.sampler.name, ) + if rollout_checkpoint_load_metrics is not None: + rollout_checkpoint_load_metrics["tq_load_seconds"] = ( + time.monotonic() - data_plane_load_started + ) + ledger_load_started = time.monotonic() recovery_ledger = _maybe_restore_rollout_recovery_ledger( last_checkpoint_path=recovery_checkpoint_path, data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, token_capture_enabled=token_capture_cfg.enabled, ) _rehydrate_rollout_recovery_prompts(recovery_ledger, dataset) + if rollout_checkpoint_load_metrics is not None: + rollout_checkpoint_load_metrics["ledger_load_seconds"] = ( + time.monotonic() - ledger_load_started + ) if use_nemo_gym: env_handles["nemo_gym"], gym_time = results["nemo_gym"] @@ -1274,6 +1300,7 @@ def _build_generation_then_trainer( save_state=save_state, last_checkpoint_path=recovery_checkpoint_path, data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, + rollout_checkpoint_load_metrics=rollout_checkpoint_load_metrics, bootstrap_fingerprint=bootstrap_digest, ) return actor_args, setup_timing_metrics diff --git a/nemo_rl/experience/blackbox_finalizer.py b/nemo_rl/experience/blackbox_finalizer.py index bf96573a74b..ba5bf846160 100644 --- a/nemo_rl/experience/blackbox_finalizer.py +++ b/nemo_rl/experience/blackbox_finalizer.py @@ -88,6 +88,7 @@ class FinalizedGroup: group_min_wv: int group_max_wv: int staging_keys: list[str] + canonical_output_tokens: int = 0 metrics: dict[str, float] = field(default_factory=dict) # True when min_valid_fraction_per_group rejected the whole group; the # caller aborts the slot instead of committing it. @@ -603,6 +604,7 @@ def finalize_group( group_min_wv=group_min_wv, group_max_wv=group_max_wv, staging_keys=[], + canonical_output_tokens=0, metrics=metrics, dropped=True, ) @@ -663,6 +665,7 @@ def finalize_group( group_min_wv=group_min_wv, group_max_wv=group_max_wv, staging_keys=[], + canonical_output_tokens=0, metrics=metrics, dropped=True, ) @@ -744,6 +747,9 @@ def finalize_group( group_min_wv=group_min_wv, group_max_wv=group_max_wv, staging_keys=(staging_keys if self._defer_routed_experts_to_policy else []), + canonical_output_tokens=sum( + int(mask) for row in valid_rows for mask in row.token_mask + ), metrics=metrics, ) diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 4456c733ac0..4c178092cc5 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -16,6 +16,7 @@ import copy import enum import json +import math from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager from dataclasses import dataclass, field @@ -27,6 +28,7 @@ from wandb import Table from nemo_rl.algorithms.async_utils.replay_buffer import ( + CheckpointMutationKind, DataPlaneCheckpointBarrier, TQReplayBuffer, ) @@ -1277,6 +1279,12 @@ def __init__( self._recovery_ledger = recovery_ledger self._data_plane_checkpoint_barrier: Optional[DataPlaneCheckpointBarrier] = None self._weight_version: int = 0 + # Process-local benchmark counters. Restored work is counted only when + # it is republished in this process, so interval rates remain coherent. + self._canonical_groups_finalized = 0 + self._canonical_output_tokens = 0 + self._recovery_siblings_reused = 0 + self._recovery_siblings_redispatched = 0 # Run-wide, shared across concurrent generate_and_push calls. Safe as a plain # int: every caller runs on the SingleController's single event loop. self._skipped_prompts: int = 0 @@ -1298,13 +1306,15 @@ def set_data_plane_checkpoint_barrier( self._data_plane_checkpoint_barrier = barrier @asynccontextmanager - async def _recovery_mutation(self) -> AsyncIterator[None]: + async def _recovery_mutation( + self, kind: CheckpointMutationKind = "recovery_retries" + ) -> AsyncIterator[None]: """Serialize short lineage transitions with native TQ snapshots.""" barrier = self._data_plane_checkpoint_barrier if barrier is None: yield return - async with barrier.mutation(): + async with barrier.mutation(kind): yield def set_weight_version(self, version: int) -> None: @@ -1315,6 +1325,37 @@ def set_weight_version(self, version: int) -> None: """ self._weight_version = int(version) + def telemetry_snapshot(self) -> dict[str, int]: + """Return cumulative canonical-publication and recovery counters.""" + return { + "canonical_groups_finalized": getattr( + self, "_canonical_groups_finalized", 0 + ), + "canonical_output_tokens": getattr(self, "_canonical_output_tokens", 0), + "recovery_siblings_reused": getattr(self, "_recovery_siblings_reused", 0), + "recovery_siblings_redispatched": ( + getattr(self, "_recovery_siblings_redispatched", 0) + ), + } + + def record_canonical_publication(self, output_tokens: int) -> None: + """Count one prompt group after its canonical TQ commit succeeds.""" + self._canonical_groups_finalized = ( + getattr(self, "_canonical_groups_finalized", 0) + 1 + ) + self._canonical_output_tokens = getattr( + self, "_canonical_output_tokens", 0 + ) + max(0, int(output_tokens)) + + def record_recovery_siblings(self, *, reused: int, redispatched: int) -> None: + """Count the sibling work avoided and repeated during recovery.""" + self._recovery_siblings_reused = getattr( + self, "_recovery_siblings_reused", 0 + ) + max(0, int(reused)) + self._recovery_siblings_redispatched = getattr( + self, "_recovery_siblings_redispatched", 0 + ) + max(0, int(redispatched)) + def reserve_prompt_group( self, input_sample: DatumSpec, *, target_step: Optional[int] = None ) -> str: @@ -1499,6 +1540,16 @@ async def generate_and_push( raise self._stats.committed += 1 + rollout_metrics = getattr(record, "rollout_metrics", {}) + mean_output_tokens = rollout_metrics.get("mean_gen_tokens_per_sample", 0) + output_tokens = 0 + if isinstance(mean_output_tokens, (int, float)): + total_output_tokens = float(mean_output_tokens) * len( + getattr(record, "completions", ()) + ) + if math.isfinite(total_output_tokens): + output_tokens = max(0, round(total_output_tokens)) + self.record_canonical_publication(output_tokens) return RolloutOutcome.COMMITTED # The infrastructure budget ran out. The same failure followed the prompt across @@ -1544,7 +1595,7 @@ async def generate_for_finalization( "generate_for_finalization requires a rollout recovery ledger" ) if recovery_group_id is None: - async with self._recovery_mutation(): + async with self._recovery_mutation("prompt_reservations"): recovery_group_id = self.reserve_prompt_group( input_sample, target_step=target_step ) @@ -1685,7 +1736,7 @@ async def _generate_for_finalization_attempt( assert self._tq_buffer is not None assert self._recovery_ledger is not None - async with self._recovery_mutation(): + async with self._recovery_mutation("recovery_retries"): recovery_group = self._recovery_ledger.get_group(recovery_group_id) if recovery_group.status == PromptGroupStatus.GENERATING: recovery_group = self._recovery_ledger.prepare_incomplete_retry( @@ -1734,7 +1785,7 @@ async def _record_streamed_completion( reward=completion.reward, ) else: - async with barrier.mutation(): + async with barrier.mutation("sibling_seals"): self._recovery_ledger.mark_sibling_sealed( group_id, generation_index=generation_index, @@ -1755,7 +1806,7 @@ async def _record_streamed_completion( f"recovery group {group_id!r} needs a prompt for " "unfinished sibling generation" ) - async with self._recovery_mutation(): + async with self._recovery_mutation("recovery_retries"): self._recovery_ledger.mark_group_dispatched( group_id, generation_indices=pending_indices ) @@ -1768,7 +1819,7 @@ async def _record_streamed_completion( finally: if inflight_registry is not None: inflight_registry.pop(group_id, None) - async with self._recovery_mutation(): + async with self._recovery_mutation("recovery_retries"): ( physical_rollout_ids, canonical_sample_ids, @@ -1789,7 +1840,7 @@ async def _record_streamed_completion( return request except BaseException: self._tq_buffer.abort(group_id) - async with self._recovery_mutation(): + async with self._recovery_mutation("recovery_retries"): self._recovery_ledger.abandon_unsealed(group_id) # A failure after the last sibling sealed leaves the group ready to # finalize; an earlier failure leaves it generating. The outer retry @@ -1809,6 +1860,6 @@ async def _discard_recovery_group(self, group_id: str) -> None: for sibling in group.siblings for key in sibling.current_attempt.staging_keys ] - async with self._recovery_mutation(): + async with self._recovery_mutation("group_removals"): await self._tq_buffer.clear_staging_keys(staging_keys) self._recovery_ledger.discard_group(group_id) diff --git a/tests/unit/data_plane/test_blackbox_finalizer.py b/tests/unit/data_plane/test_blackbox_finalizer.py index 42f9a6b4f94..25732e800fd 100644 --- a/tests/unit/data_plane/test_blackbox_finalizer.py +++ b/tests/unit/data_plane/test_blackbox_finalizer.py @@ -222,6 +222,7 @@ def test_finalize_group_publishes_n_rows_with_placeholder(tq_client, partitions) assert finalized.metrics["finalize/invalid_row_rate"] == 0.5 assert finalized.metrics["finalize/heuristic_terminal_count"] == 1.0 assert finalized.metrics["finalize/heuristic_terminal_fraction"] == 0.5 + assert finalized.canonical_output_tokens == sum(expected.token_mask) rows = _fetch_rows(tq_client, rollout_ids) sample_mask = torch.as_tensor(rows["sample_mask"]).flatten() @@ -285,6 +286,7 @@ def test_finalize_group_min_valid_fraction_drops(tq_client, partitions): ) assert finalized.dropped assert finalized.meta is None + assert finalized.canonical_output_tokens == 0 assert (finalized.group_min_wv, finalized.group_max_wv) == (3, 3) with pytest.raises((KeyError, RuntimeError, ValueError)): rows = _fetch_rows(tq_client, rollout_ids) diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 98a60596232..763f6c3b695 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -255,6 +255,20 @@ async def _logged_commit(*args, **kwargs): assert record == "r0" assert start_v == 0 assert end_v == 0 + assert mgr.telemetry_snapshot()["canonical_groups_finalized"] == 1 + + def test_publication_and_recovery_telemetry_are_cumulative(self): + mgr = _make_manager(_FakeBuffer(), _FakeImpl()) + + mgr.record_canonical_publication(42) + mgr.record_recovery_siblings(reused=3, redispatched=1) + + assert mgr.telemetry_snapshot() == { + "canonical_groups_finalized": 1, + "canonical_output_tokens": 42, + "recovery_siblings_reused": 3, + "recovery_siblings_redispatched": 1, + } def test_start_weight_version_pinned_at_reserve_time(self): """If set_weight_version is called mid-rollout, start != end.""" diff --git a/tests/unit/single_controller/test_finalizer_lifecycle.py b/tests/unit/single_controller/test_finalizer_lifecycle.py index c642f9df3b0..2ca2499338a 100644 --- a/tests/unit/single_controller/test_finalizer_lifecycle.py +++ b/tests/unit/single_controller/test_finalizer_lifecycle.py @@ -122,6 +122,7 @@ def _controller(actor: object) -> Any: ctrl._finalizer_metrics_by_group = {} ctrl._rollout_recovery_ledger = MagicMock() ctrl._rollout_recovery_ledger.__contains__.return_value = False + ctrl._rollout_manager = MagicMock() ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._buffer = MagicMock() ctrl._buffer.commit_finalized = AsyncMock() diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py index 808cee75e83..f1c7862ee7b 100644 --- a/tests/unit/single_controller/test_sc_checkpointing.py +++ b/tests/unit/single_controller/test_sc_checkpointing.py @@ -102,6 +102,19 @@ _ACTOR_CLS = SingleControllerActor.__ray_metadata__.modified_class _PARTITION_ID = "rollout_data" + + +class _SteppingClock: + def __init__(self, *, start: float = 0.0, step: float = 1.0) -> None: + self._next = start + self._step = step + + def __call__(self) -> float: + value = self._next + self._next += self._step + return value + + _STAGING_PARTITION_ID = "rollout_staging" @@ -371,6 +384,9 @@ def __init__(self) -> None: def set_rollout_weight_version(self, version: int) -> None: self.rollout_weight_versions.append(version) + def drain_latest_logger_metrics(self) -> dict[str, Any]: + return {} + class _FakeRolloutManager: def __init__(self, recovery_ledger: Optional[RolloutRecoveryLedger] = None) -> None: @@ -380,6 +396,12 @@ def __init__(self, recovery_ledger: Optional[RolloutRecoveryLedger] = None) -> N self.checkpoint_barrier: Optional[DataPlaneCheckpointBarrier] = None self.recovery_calls: list[str] = [] self.discard_calls: list[str] = [] + self.telemetry = { + "canonical_groups_finalized": 0, + "canonical_output_tokens": 0, + "recovery_siblings_reused": 0, + "recovery_siblings_redispatched": 0, + } def set_weight_version(self, version: int) -> None: self.weight_versions.append(version) @@ -399,6 +421,17 @@ async def discard_recovery_group(self, group_id: str) -> None: assert self.recovery_ledger is not None self.recovery_ledger.discard_group(group_id) + def telemetry_snapshot(self) -> dict[str, int]: + return dict(self.telemetry) + + def record_canonical_publication(self, output_tokens: int) -> None: + self.telemetry["canonical_groups_finalized"] += 1 + self.telemetry["canonical_output_tokens"] += output_tokens + + def record_recovery_siblings(self, *, reused: int, redispatched: int) -> None: + self.telemetry["recovery_siblings_reused"] += reused + self.telemetry["recovery_siblings_redispatched"] += redispatched + def _unfinished_recovery_ledger(group_count: int) -> RolloutRecoveryLedger: ledger = RolloutRecoveryLedger() @@ -448,10 +481,6 @@ def group_ids(self) -> tuple[str, ...]: def target_step_list(self) -> tuple[Optional[int], ...]: return tuple(self._target_step_list) - def __len__(self) -> int: - """Match the production TQReplayBuffer occupancy contract.""" - return len(self.target_step_list) - def set_data_plane_checkpoint_barrier( self, barrier: DataPlaneCheckpointBarrier ) -> None: @@ -616,6 +645,7 @@ def _make_actor_args( data_plane_checkpoint_metadata: Optional[dict[str, Any]] = None, rollout_manager: Optional[_FakeRolloutManager] = None, bootstrap_fingerprint: Optional[str] = None, + rollout_checkpoint_load_metrics: Optional[dict[str, float]] = None, ) -> SingleControllerActorArgs: return SingleControllerActorArgs( gen_handle=_FakeGeneration(), @@ -640,6 +670,7 @@ def _make_actor_args( last_checkpoint_path=last_checkpoint_path, data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, bootstrap_fingerprint=bootstrap_fingerprint, + rollout_checkpoint_load_metrics=rollout_checkpoint_load_metrics, ) @@ -1019,6 +1050,86 @@ def test_pre_step_snapshot_omits_trainer_payload(self, tmp_path: Path) -> None: assert (snapshot / "config.yaml").is_file() assert not (snapshot / "policy").exists() + def test_logs_snapshot_phase_durations(self, tmp_path: Path) -> None: + actor = self._make_actor(tmp_path) + actor._logger = MagicMock() + actor._telemetry_started_at = 0.0 + + with patch( + "nemo_rl.algorithms.single_controller.time.monotonic", + new=_SteppingClock(), + ): + assert asyncio.run(actor._save_rollout_checkpoint(force=True)) + + logged = actor._logger.log_metrics.call_args.args[0] + assert logged["snapshot_sequence"] == 1.0 + assert logged["total_save_seconds"] > 0 + assert 0 <= logged["barrier_wait_seconds"] <= logged["total_save_seconds"] + assert 0 <= logged["tq_save_seconds"] <= logged["exclusive_hold_seconds"] + assert logged["rolled_back_train_groups"] == 0.0 + assert actor._logger.log_metrics.call_args.kwargs == { + "step": 1, + "prefix": "timing/rollout_checkpoint", + "step_metric": "telemetry/wall_time_seconds", + } + + def test_logs_raw_and_canonical_rollout_throughput(self, tmp_path: Path) -> None: + actor = self._make_actor(tmp_path) + actor._logger = MagicMock() + actor._telemetry_started_at = 0.0 + snapshots = iter( + [ + { + "canonical_groups_finalized": 0, + "canonical_output_tokens": 0, + "recovery_siblings_reused": 0, + "recovery_siblings_redispatched": 0, + }, + { + "canonical_groups_finalized": 2, + "canonical_output_tokens": 40, + "recovery_siblings_reused": 1, + "recovery_siblings_redispatched": 3, + }, + ] + ) + actor._rollout_manager.telemetry_snapshot = lambda: next(snapshots) + generation_snapshots = iter( + [ + {"generation_tokens": {0: [60], 1: [40]}}, + { + "generation_tokens": {0: [180], 1: [120]}, + "inflight_batch_sizes": {0: [2], 1: [1]}, + "num_pending_samples": {0: [1], 1: [3]}, + "kv_cache_usage_perc": {0: [0.5], 1: [0.7]}, + }, + ] + ) + actor._gen.drain_latest_logger_metrics = lambda: next(generation_snapshots) + + async def _sample_twice() -> None: + await actor._log_rollout_throughput_metrics(emit=False) + actor._rollout_completion_durations_s.extend([2.0, 4.0]) + actor._rollout_queue_wait_durations_s.extend([1.0, 3.0]) + await actor._log_rollout_throughput_metrics() + + with patch( + "nemo_rl.algorithms.single_controller.time.monotonic", + new=_SteppingClock(start=10.0, step=10.0), + ): + asyncio.run(_sample_twice()) + + logged = actor._logger.log_metrics.call_args.args[0] + assert logged["vllm_output_tokens_per_second"] == pytest.approx(20.0) + assert logged["canonical_output_tokens_per_second"] == pytest.approx(4.0) + assert logged["canonical_groups_per_second"] == pytest.approx(0.2) + assert logged["vllm_requests_running"] == 3 + assert logged["vllm_requests_waiting"] == 4 + assert logged["vllm_kv_cache_usage_mean"] == pytest.approx(0.6) + assert logged["group_completion_seconds_p50"] == pytest.approx(2.0) + assert logged["group_completion_seconds_p95"] == pytest.approx(4.0) + assert logged["group_queue_wait_seconds_p95"] == pytest.approx(3.0) + def test_unchanged_state_is_not_saved_twice(self, tmp_path: Path) -> None: actor = self._make_actor(tmp_path)