From 8732bec0cbc440602860ffa49d26213ecf309f20 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 10 Aug 2026 14:57:58 -0400 Subject: [PATCH 01/16] 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 675406af33a..c66d37d53cc 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -1088,6 +1088,21 @@ def log_metrics( tee_rl_metrics_to_otel(metrics, prefix) + 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 cecc49769c5..b8a2e97a7f0 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -1838,6 +1838,34 @@ def test_log_metrics_preserves_non_train_generation_histogram_prefix( "generation_metrics/validation/histogram/gen_tokens_length", ) + @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 ef9bf4cdc314f84e5fd1132600550070646fa2f9 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 11 Aug 2026 16:50:52 -0400 Subject: [PATCH 02/16] 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 | 23 +++++++++++++ .../test_tq_replay_buffer.py | 34 +++++++++++++++++++ tests/unit/utils/test_logger.py | 10 +++--- 7 files changed, 105 insertions(+), 12 deletions(-) diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 00f95740e7d..d485da43009 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -652,3 +652,7 @@ def get_step_metrics(self) -> dict[str, float]: metrics return an empty dictionary. """ 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 99b5f0793aa..83643ff7ee3 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1410,8 +1410,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): @@ -1422,7 +1422,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) @@ -1456,6 +1456,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 @@ -1475,6 +1483,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 bb77973fedb..e8f3fa96987 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -392,6 +392,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 c66d37d53cc..45a40184f66 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -383,10 +383,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 e33a6630900..85eb4f8ea72 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -561,6 +561,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 46504442116..77ff988f4a9 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -29,6 +29,7 @@ from nemo_rl.algorithms.async_utils.replay_buffer import ( REPLAY_BUFFER_METADATA_SCHEMA_VERSION, REPLAY_BUFFER_METADATA_STORAGE, + CheckpointMutationKind, DataPlaneCheckpointBarrier, PostWriteEnrichmentError, TQReplayBuffer, @@ -446,6 +447,39 @@ async def exercise() -> None: async with barrier.mutation() as cut: cut.require_live() + 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()) diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py index b8a2e97a7f0..01bc8cec8fc 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -459,10 +459,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): @@ -480,15 +479,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 10615d1463b256f2a12a7fd9c7bcaa3df079ff69 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 11 Aug 2026 18:36:12 -0400 Subject: [PATCH 03/16] 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 | 13 +++++ 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 +- tests/unit/single_controller/test_setup.py | 8 +++ tests/unit/utils/test_logger.py | 55 +++++++++++++++---- 10 files changed, 131 insertions(+), 45 deletions(-) diff --git a/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml b/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml index 5997a6d433e..e1bf7e72aa4 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 d16674c1eb8..ad96de03424 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -1039,6 +1039,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 d485da43009..e53141f81a5 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -653,6 +653,6 @@ def get_step_metrics(self) -> dict[str, float]: """ 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 e1735064a45..48d3b968951 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -74,6 +74,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 83643ff7ee3..4db0ceef279 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1460,9 +1460,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): @@ -1483,9 +1483,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 e8f3fa96987..9994326c236 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -312,7 +312,7 @@ def _create_engine(self, llm_kwargs: dict[str, Any]) -> None: 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 @@ -392,7 +392,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 {} @@ -408,11 +408,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 45a40184f66..2cd31790b9a 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -48,6 +48,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] @@ -225,6 +230,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( @@ -358,7 +366,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, @@ -383,16 +392,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. @@ -409,7 +415,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. @@ -417,9 +424,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. @@ -430,11 +438,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 85eb4f8ea72..5142f917179 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -561,7 +561,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() @@ -570,7 +570,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], @@ -582,6 +582,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_setup.py b/tests/unit/single_controller/test_setup.py index 0872bb6f759..70c360ee0e7 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -694,6 +694,14 @@ def test_reward_kl_rejects_skipping_policy_logprobs(self, patched_factories): patched_factories["_build_clusters"].assert_not_called() patched_factories["_build_trainer"].assert_not_called() + 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["backend"] = "mooncake_cpu" diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py index 01bc8cec8fc..f0d8098b0df 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -424,9 +424,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): @@ -439,10 +439,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): @@ -459,9 +463,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): @@ -485,9 +490,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.""" @@ -497,11 +531,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 cca3a661634a5028ddce34bc5ecea6f00397ebff Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 11 Aug 2026 19:03:00 -0400 Subject: [PATCH 04/16] 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 +++++++++++++++--- tests/unit/single_controller/test_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 ad96de03424..22132feee42 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -1040,17 +1040,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 2cd31790b9a..3ed4fb9cf1a 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -51,6 +51,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" @@ -231,8 +234,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( @@ -360,14 +368,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, @@ -392,13 +537,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. @@ -406,7 +566,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. @@ -416,7 +577,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. @@ -426,6 +591,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 @@ -438,23 +604,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_setup.py b/tests/unit/single_controller/test_setup.py index 70c360ee0e7..d8552b4446c 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -697,11 +697,35 @@ def test_reward_kl_rejects_skipping_policy_logprobs(self, patched_factories): 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["backend"] = "mooncake_cpu" diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py index f0d8098b0df..895c3a98785 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -422,7 +422,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 @@ -437,7 +437,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 @@ -463,10 +463,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): @@ -490,9 +492,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): @@ -500,21 +504,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, } ), @@ -531,10 +548,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 4e67b5ff767163305cee43471a4b80449cd3f496 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 11 Aug 2026 19:30:16 -0400 Subject: [PATCH 05/16] 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 ++++++ .../unit/single_controller/test_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 aec53201dde..b77a7c8837d 100644 --- a/nemo_rl/algorithms/distillation.py +++ b/nemo_rl/algorithms/distillation.py @@ -1154,7 +1154,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 ac775f16340..4b0d8d193ab 100644 --- a/nemo_rl/algorithms/dpo.py +++ b/nemo_rl/algorithms/dpo.py @@ -889,7 +889,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 5b2e7b9216e..4d21d387b82 100644 --- a/nemo_rl/algorithms/rm.py +++ b/nemo_rl/algorithms/rm.py @@ -747,7 +747,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 2ebdb90250b..3ba477645d2 100644 --- a/nemo_rl/algorithms/sft.py +++ b/nemo_rl/algorithms/sft.py @@ -795,7 +795,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 3ed4fb9cf1a..164f30c1108 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 @@ -241,6 +242,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 a642e47a1f4..f1a08e69da4 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 7d2a05a237f..d627a5a5e1f 100644 --- a/tests/unit/algorithms/test_dpo.py +++ b/tests/unit/algorithms/test_dpo.py @@ -305,6 +305,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_configured_stop_forces_checkpoint_without_shortening_run( 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 d5fe1082735..288da063e09 100644 --- a/tests/unit/algorithms/test_sft.py +++ b/tests/unit/algorithms/test_sft.py @@ -186,6 +186,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_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 668ec6d52de..022754b0523 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -525,6 +525,10 @@ def __init__( def group_ids(self) -> tuple[str, ...]: return () + 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 895c3a98785..07bfb444be8 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -587,6 +587,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 55e0f772ac3e88d934e7ca8d6b859c51b1f2fbcf Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 19 Aug 2026 17:47:49 -0400 Subject: [PATCH 06/16] feat(rollout): refresh recovery benchmark telemetry (cherry picked from commit af43edd0033ab14b91f7740c33c73101d445f728) Signed-off-by: Anish Mahishi --- .../nemo_gym/grpo_qwen3_30ba3b_instruct.yaml | 5 + .../algorithms/async_utils/replay_buffer.py | 95 +++- nemo_rl/algorithms/single_controller.py | 471 +++++++++++++++++- .../single_controller_utils/config.py | 5 + .../single_controller_utils/setup.py | 23 + nemo_rl/experience/rollout_manager.py | 63 ++- nemo_rl/experience/rollout_reassembler.py | 5 + .../data_plane/test_rollout_reassembler.py | 1 + tests/unit/experience/test_rollout_manager.py | 14 + .../single_controller/test_checkpointing.py | 118 ++++- .../test_finalizer_lifecycle.py | 1 + 11 files changed, 768 insertions(+), 33 deletions(-) diff --git a/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml b/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml index e1bf7e72aa4..1669e1ae890 100644 --- a/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml +++ b/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml @@ -178,3 +178,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 608da389533..baff24e1979 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -20,10 +20,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, @@ -63,6 +65,37 @@ 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. @@ -233,6 +266,10 @@ def __init__(self) -> None: self._active_mutations = 0 self._section_holders: set[asyncio.Task[Any]] = set() 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) def _current_task(self) -> asyncio.Task[Any]: """Return the task entering a barrier section and reject reentrancy.""" @@ -252,11 +289,26 @@ def mutation_version(self) -> int: return self._mutation_version @asynccontextmanager - async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: - """Yield a live cut after any active checkpoint exits.""" + async def mutation( + self, kind: CheckpointMutationKind = "other" + ) -> AsyncIterator[DataPlaneMutationCut]: + """Yield one task-local live cut after any active checkpoint exits.""" async with self._condition: task = self._current_task() - 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._section_holders.add(task) cut = DataPlaneMutationCut(self) @@ -273,6 +325,25 @@ async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: 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[DataPlaneMutationCut]: """Yield a live capability after blocking and draining all mutations.""" @@ -1184,7 +1255,9 @@ async def commit( "the async message-log flattening path." ) trace_rollout_payload(keys=sample_ids, data=train_batch) - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "group_commits" + ) as cut: try: await call_data_plane( self._dp_client, @@ -1264,7 +1337,9 @@ 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() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "group_removals" + ) as cut: return await self._remove_groups_unlocked( cut, [group_id], clear_data_plane=remove_in_dp ) @@ -1435,7 +1510,9 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: # removal may shift every list index while this task waits for the barrier # or for DataPlane cleanup. drop_group_ids = [self._group_ids[i] for i in drop_idxs] - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "group_removals" + ) as cut: return await self._remove_groups_unlocked( cut, drop_group_ids, clear_data_plane=remove_in_dp ) @@ -1465,7 +1542,9 @@ async def claim_for_training(self, idxs: list[int]) -> int: f"{claim_idxs[0]}; size={len(self.meta_list)}" ) claim_group_ids = [self._group_ids[i] for i in claim_idxs] - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "group_removals" + ) as cut: return await self._remove_groups_unlocked( cut, claim_group_ids, diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 36dd5ffd103..387ab56c71c 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -66,6 +66,7 @@ from nemo_rl.algorithms import opd as opd_module from nemo_rl.algorithms.async_utils.replay_buffer import ( + CHECKPOINT_MUTATION_KINDS, DATA_PLANE_CHECKPOINT_DIR, LEGACY_REPLAY_BUFFER_FILENAME, REPLACEMENT_RESERVE_FILENAME, @@ -135,6 +136,7 @@ ROLLOUT_RECOVERY_SCHEMA_VERSION, ROLLOUT_RECOVERY_STATE_FILENAME, PromptGroupPhase, + PromptGroupRecoveryRecord, RolloutRecoveryState, build_rollout_recovery_state, parse_rollout_recovery_state, @@ -156,6 +158,14 @@ Generation = Union[VllmGeneration, SGLangGeneration, MegatronGeneration] +_TELEMETRY_WALL_TIME_METRIC = "telemetry/wall_time_seconds" +_TELEMETRY_PREFIXES = ( + "timing/rollout_checkpoint", + "timing/rollout_recovery", + "rollout/checkpoint_outcome", + "rollout/throughput", +) + # Named `log` rather than `logger` to keep it distinct from the experiment # Logger this module also uses as `self._logger`. log = logging.getLogger(__name__) @@ -175,6 +185,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] def _pooled_opd_metrics( @@ -344,6 +379,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. @@ -363,6 +418,9 @@ def __init__( self._data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = ( actor_args.data_plane_checkpoint_metadata ) + self._rollout_checkpoint_load_metrics = getattr( + actor_args, "rollout_checkpoint_load_metrics", None + ) self._consumed_samples: int = actor_args.save_state.consumed_samples self._total_valid_tokens: int = actor_args.save_state.total_valid_tokens @@ -453,6 +511,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() @@ -533,17 +594,39 @@ async def run(self) -> dict[str, Any]: await self._sync_weights() self._rollout_manager.set_weight_version(self._trainer_version) + 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() await self._maybe_restore_rollout_recovery( restored_replay_groups=restored_replay_groups ) await self._maybe_restore_replacement_reserve() + 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", + ) # Start the rollout and train pumps, plus the watchdog rollout_task = asyncio.create_task(self._rollout_pump()) train_task = asyncio.create_task(self._train_pump()) watchdog_task = asyncio.create_task(self._stall_watchdog_pump()) tasks = [rollout_task, train_task, watchdog_task] + rollout_checkpoint_cfg = getattr( + self._master_config, "rollout_checkpointing", None + ) + telemetry_interval_s = getattr( + rollout_checkpoint_cfg, "telemetry_interval_s", None + ) + if telemetry_interval_s is not None: + await self._log_rollout_throughput_metrics(emit=False) rollout_checkpoint_task = ( asyncio.create_task(self._rollout_checkpoint_pump()) if self._master_config.rollout_checkpointing.snapshot_attempt_interval_s @@ -552,6 +635,13 @@ async def run(self) -> dict[str, Any]: ) if rollout_checkpoint_task is not None: tasks.append(rollout_checkpoint_task) + rollout_telemetry_task = ( + asyncio.create_task(self._rollout_telemetry_pump()) + if telemetry_interval_s is not None + else None + ) + if rollout_telemetry_task is not None: + tasks.append(rollout_telemetry_task) # Only with fleet health on. Created unconditionally it would be a timer firing # every probe_interval_s for every run that does not use the feature, which is # the default. @@ -588,6 +678,9 @@ async def run(self) -> dict[str, Any]: ): # Loops forever like the watchdog, so finishing at all means it raised. await probe_task + if rollout_telemetry_task is not None and rollout_telemetry_task in done: + await rollout_telemetry_task + raise RuntimeError("rollout telemetry pump exited unexpectedly") if not stop_after_rollout_checkpoint and watchdog_task in done: # The watchdog loops forever, so finishing at all means it raised -- # a stall or an unhealthy environment. Surface that ahead of the @@ -636,6 +729,37 @@ 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.""" + 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: @@ -968,7 +1092,9 @@ def _commit( await self._sampler.wait_until_admissible( trainer_version_fn=lambda: self._trainer_version ) - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "prompt_reservations" + ) as cut: target_step = self._sampler.commit_admission(cut) return _commit(cut, target_step) @@ -978,7 +1104,9 @@ def _commit( # a checkpoint drains mutation slots while blocking the train pump, so a # longer wait deadlocks the run. Implement TransactionalAdmissionSampler # to keep the gate wait outside the mutation cut entirely. - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "prompt_reservations" + ) as cut: target_step = await self._sampler.admit( trainer_version_fn=lambda: self._trainer_version ) @@ -1001,6 +1129,15 @@ async def _redispatch_restored_rollouts( groups_to_recover = recovery_ledger.groups() if not groups_to_recover: return + recovery_started = time.monotonic() + reused_siblings = 0 + redispatched_siblings = 0 + + def _record_sibling_work(group: PromptGroupRecoveryRecord) -> None: + nonlocal reused_siblings, redispatched_siblings + reused = len(group.sealed_generation_indices) + reused_siblings += reused + redispatched_siblings += group.expected_generations - reused recognized_phases = ( PromptGroupPhase.ADMITTED, @@ -1027,6 +1164,7 @@ async def _redispatch_restored_rollouts( group.target_step, group.group_id, ) + _record_sibling_work(group) redispatched += 1 # A checkpoint may land after dataloader ownership is recorded but before @@ -1049,8 +1187,25 @@ async def _redispatch_restored_rollouts( group.target_step, group.group_id, ) + _record_sibling_work(group) redispatched += 1 + self._rollout_manager.record_recovery_siblings( + reused=reused_siblings, + redispatched=redispatched_siblings, + ) + self._log_telemetry_metrics( + { + "redispatch_schedule_seconds": time.monotonic() - recovery_started, + "groups_considered": float(len(groups_to_recover)), + "groups_redispatched": float(redispatched), + "siblings_reused": float(reused_siblings), + "siblings_redispatched": float(redispatched_siblings), + }, + step=self._train_steps, + prefix="timing/rollout_recovery", + ) + print( f"📦 Redispatched {redispatched} unfinished rollout " "group(s) before new dataloader work", @@ -1340,7 +1495,9 @@ async def _cleanup_known_finalization_request( self, request: "ReassemblyRequest" ) -> None: """Clear a known request outcome without racing a native TQ snapshot.""" - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "group_removals" + ) as cut: await self._cleanup_known_finalization_request_unlocked(cut, request) async def _finalize_with_actor( @@ -1383,7 +1540,9 @@ async def _finalize_with_actor( # every in-flight finalizer RPC. Releasing the cut across the await # would let a snapshot preserve canonical rows without the matching # replay index and lineage transition, which is not recoverable. - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "group_commits" + ) as cut: ledger = self._rollout_recovery_ledger ledger.mark_finalization_started(cut, request.group_id) try: @@ -1457,6 +1616,9 @@ async def _finalize_with_actor( # Canonical TQ rows plus replay metadata now own the completed # group; keep only unfinished work in the lineage sidecar. ledger.discard_group(cut, request.group_id) + self._rollout_manager.record_canonical_publication( + finalized.canonical_output_tokens + ) committed = True finally: self._active_finalizers -= 1 @@ -1574,6 +1736,9 @@ async def _dispatch_one_prompt( target_step: Optional[int], lineage_group_id: Optional[str], task_started_event: asyncio.Event, + *, + work_started: float, + dispatch_started: float, ) -> None: task_started_event.set() self._inflight_rollouts += 1 @@ -1623,7 +1788,9 @@ async def _dispatch_one_prompt( and lineage_group_id is not None ): async with ( - self._data_plane_checkpoint_barrier.mutation() + self._data_plane_checkpoint_barrier.mutation( + "group_removals" + ) ) as cut: await self._rollout_manager.discard_recovery_group( cut, lineage_group_id @@ -1737,7 +1904,9 @@ async def _dispatch_one_prompt( if self._rollout_recovery_enabled: assert lineage_group_id is not None async with ( - self._data_plane_checkpoint_barrier.mutation() + self._data_plane_checkpoint_barrier.mutation( + "prompt_reservations" + ) ) as cut: replacement = self._take_replacement( target_step, replacements @@ -1809,6 +1978,11 @@ async def _dispatch_one_prompt( self._batch_replacements.get(target_step, 0) + 1 ) + 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"])): @@ -1836,12 +2010,26 @@ async def _launch( "recovery-enabled rollout dispatch requires a pre-reserved " "prompt-group ID" ) + work_started = time.monotonic() # check if buffer is full - await self._buffer_capacity.acquire() + self._buffer_capacity_waiters += 1 + try: + await self._buffer_capacity.acquire() + finally: + self._buffer_capacity_waiters -= 1 # check if inflight rollouts is full - await sem.acquire() + self._rollout_slot_waiters += 1 + try: + await sem.acquire() + finally: + self._rollout_slot_waiters -= 1 # wait for rollout to be permitted - await self._rollout_permitted.wait() + self._rollout_permitted_waiters += 1 + try: + await self._rollout_permitted.wait() + finally: + self._rollout_permitted_waiters -= 1 + dispatch_started = time.monotonic() task_started_event = asyncio.Event() # dispatch rollout @@ -1851,6 +2039,8 @@ async def _launch( target_step, lineage_group_id, task_started_event, + work_started=work_started, + dispatch_started=dispatch_started, ) ) self._dispatched_rollouts.add(task) @@ -1898,7 +2088,9 @@ async def _launch( dataloader_iterator = iter(self._dataloader) while True: prompt_dispatches: list[tuple[DatumSpec, str]] = [] - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "prompt_reservations" + ) as cut: try: prompt_batch = next(dataloader_iterator) except StopIteration: @@ -2033,7 +2225,9 @@ async def _drain_reserve_into_steps( while len(self._replacement_reserve) >= num_prompts_per_step: if self._rollout_recovery_enabled: prompt_dispatches: list[tuple[DatumSpec, str]] = [] - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "prompt_reservations" + ) as cut: step_prompts = [ self._replacement_reserve.popleft() for _ in range(num_prompts_per_step) @@ -2293,9 +2487,12 @@ async def _train_pump(self) -> None: await asyncio.sleep(0) # Evict stale groups - evicted = await self._sampler.evict( - current_train_weight=self._trainer_version, - ) + async with self._data_plane_checkpoint_barrier.mutation( + "group_removals" + ): + evicted = await self._sampler.evict( + current_train_weight=self._trainer_version, + ) evicted_stale_prompt_groups += evicted if evicted: print( @@ -2630,7 +2827,9 @@ async def _train_pump(self) -> None: step_metrics.update(aggregate_step_metrics(policy_result)) if value_result is not None: step_metrics.update(_compute_critic_metrics(value_result)) - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "sample_clears" + ) as cut: await self._cleanup_consumed_metas_unlocked(cut, consumed_metas) self._buffer.release_training_claims(consumed_training_claim_ids) for _ in range(consumed_group_count): @@ -3388,7 +3587,9 @@ def _stale_groups() -> list[tuple[str, asyncio.Task[None]]]: return stale_groups if self._rollout_recovery_enabled: - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "group_removals" + ) as cut: # Re-evaluate after acquiring the cut: a rollout may have completed # while a checkpoint holder delayed this mutation. stale_groups = _stale_groups() @@ -3470,6 +3671,7 @@ async def _capture_rollout_checkpoint_cut( replay_metadata=replay_metadata, clear_unreferenced=False, ) + tq_save_started = time.monotonic() await self._save_data_plane_checkpoint( checkpoint_path, train_steps=self._train_steps, @@ -3479,6 +3681,7 @@ async def _capture_rollout_checkpoint_cut( rollout_recovery_payload_sha256=recovery_digest, rollout_recovery_group_count=len(recovery_state["groups"]), ) + tq_save_seconds = time.monotonic() - tq_save_started return _RolloutCheckpointCut( dataloader_state=dataloader_state, sampler_dispatch_index=self._sampler.dispatch_index, @@ -3488,6 +3691,7 @@ async def _capture_rollout_checkpoint_cut( rollout_recovery_group_count=len(recovery_state["groups"]), rolled_back_train_group_count=len(training_owned_groups), mutation_version=self._data_plane_checkpoint_barrier.mutation_version, + tq_save_seconds=tq_save_seconds, ) async def _write_rollout_checkpoint_sidecars( @@ -3540,6 +3744,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: @@ -3590,11 +3795,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() as cut: + barrier_acquired = time.monotonic() if ( self._optimizer_commit_in_progress or self._train_steps != expected_train_step @@ -3606,6 +3813,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: snapshot_cut = await self._capture_rollout_checkpoint_cut( cut, tmp_path ) + barrier_released = time.monotonic() await self._write_rollout_checkpoint_sidecars(tmp_path, snapshot_cut) manifest = RolloutSnapshotManifest( @@ -3639,15 +3847,67 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: self._last_rollout_snapshot_mutation_version = snapshot_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": snapshot_cut.tq_save_seconds, + "barrier_wait_seconds": barrier_acquired - barrier_requested, + "exclusive_hold_seconds": barrier_released - barrier_acquired, + "replay_groups": float( + len(snapshot_cut.replay_metadata["groups"]) + if snapshot_cut.replay_metadata is not None + else 0 + ), + "ledger_groups": float( + snapshot_cut.rollout_recovery_group_count or 0 + ), + } + 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={snapshot_cut.rollout_recovery_group_count or 0})", + f"ledger_groups={snapshot_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.snapshot_attempt_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.""" snapshot_attempt_interval_s = ( @@ -3658,10 +3918,15 @@ async def _rollout_checkpoint_pump(self) -> None: consecutive_failures = 0 while True: await asyncio.sleep(snapshot_attempt_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 (OSError, TimeoutError) 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" @@ -3681,6 +3946,10 @@ async def _rollout_checkpoint_pump(self) -> None: ) from error continue consecutive_failures = 0 + self._log_rollout_checkpoint_outcome( + outcome="completed" if saved else "skipped", + attempt_duration_seconds=time.monotonic() - attempt_started, + ) if deadline_due and saved and self._timeout.check_save(): print( "Checkpoint deadline reached before the first train step; " @@ -3690,6 +3959,172 @@ 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() + + async def _log_rollout_throughput_metrics(self, *, emit: bool = True) -> None: + """Serialize backend sampling so cumulative counters have one baseline.""" + 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 _save_checkpoint( self, step_metrics: dict[str, Any], diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 32f187b7e6b..8e959db8d1e 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -741,9 +741,14 @@ class RolloutCheckpointConfig(BaseModel, extra="forbid"): ``val:`` settings are rejected during setup. Unknown keys are forbidden because a misspelled interval, retention, or restore option can silently disable the durability behavior the operator intended. + + ``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 ``snapshot_attempt_interval_s``. """ snapshot_attempt_interval_s: Annotated[Optional[float], Field(gt=0)] = None + telemetry_interval_s: Annotated[Optional[float], Field(gt=0)] = None keep_latest_k: Annotated[int, Field(ge=1)] = 2 restore_mode: Literal["latest", "trainer_checkpoint"] = "latest" extra_fingerprint_excluded_paths: list[str] = Field(default_factory=list) diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 22132feee42..b52382741b7 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -103,6 +103,7 @@ RolloutRetryPolicy, RolloutTimeouts, ) +from nemo_rl.experience.rollout_recovery import ROLLOUT_RECOVERY_STATE_FILENAME from nemo_rl.experience.rollouts import ( get_nemo_gym_thinking_tags, resolve_reward_penalty_config, @@ -165,6 +166,7 @@ class SingleControllerActorArgs: # Defaulted fields must follow the required ones above, so these stay last. data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = None bootstrap_identity: Optional[BootstrapCompatibilityIdentity] = None + rollout_checkpoint_load_metrics: Optional[dict[str, float]] = None # None when async_rl.generation_fleet_health is disabled; the SingleController # drives the probe loop when it is present. fleet_monitor: Optional[GenerationFleetHealth] = None @@ -1281,6 +1283,16 @@ def setup_single_controller( f"without considering newer periodic snapshots: {trainer_checkpoint_path}", 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 @@ -1326,7 +1338,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) @@ -1634,6 +1651,7 @@ def _build_trainer_then_megatron_generation() -> tuple[ # 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, @@ -1641,6 +1659,10 @@ def _build_trainer_then_megatron_generation() -> tuple[ 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 + ) if use_nemo_gym: # the two fields are only meaningful when use_nemo_gym enabled @@ -1886,6 +1908,7 @@ def _build_trainer_then_megatron_generation() -> tuple[ last_checkpoint_path=recovery_checkpoint_path, data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, bootstrap_identity=bootstrap_identity, + rollout_checkpoint_load_metrics=rollout_checkpoint_load_metrics, finalizer_actors=finalizer_actors, fleet_monitor=fleet_monitor, generation_router=generation_router, diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 21f8cc68fc0..41621d86310 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -18,6 +18,7 @@ import copy import enum import json +import math import uuid from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from contextlib import asynccontextmanager @@ -30,6 +31,7 @@ from wandb import Table from nemo_rl.algorithms.async_utils.replay_buffer import ( + CheckpointMutationKind, DataPlaneCheckpointBarrier, DataPlaneMutationCut, PostWriteEnrichmentError, @@ -1605,6 +1607,10 @@ def __init__( self._data_plane_checkpoint_barrier: Optional[DataPlaneCheckpointBarrier] = None self._env_handles = task_to_env self._weight_version: int = 0 + 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 @@ -1658,7 +1664,9 @@ def set_data_plane_checkpoint_barrier( self._data_plane_checkpoint_barrier = barrier @asynccontextmanager - async def _recovery_mutation(self) -> AsyncIterator[DataPlaneMutationCut]: + async def _recovery_mutation( + self, kind: CheckpointMutationKind = "recovery_retries" + ) -> AsyncIterator[DataPlaneMutationCut]: """Serialize short lineage transitions with native TQ snapshots.""" barrier = self._data_plane_checkpoint_barrier if barrier is None: @@ -1666,9 +1674,40 @@ async def _recovery_mutation(self) -> AsyncIterator[DataPlaneMutationCut]: "RolloutManager must be bound to the SingleController data-plane " "checkpoint barrier before mutating rollout recovery state" ) - async with barrier.mutation() as cut: + async with barrier.mutation(kind) as cut: yield cut + 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 sibling work avoided and repeated after a process restart.""" + 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, cut: DataPlaneMutationCut, @@ -1969,6 +2008,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) # A commit proves the fleet is answering, which is exactly the claim the # consecutive budget is testing, so it clears the run of drops. Placed on # the success path rather than in the infra handler so that a prompt which @@ -1976,7 +2025,9 @@ async def generate_and_push( self._consecutive_infra_drops = 0 if lineage_group_id is not None: async with ( - self._tq_buffer.data_plane_checkpoint_barrier.mutation() + self._tq_buffer.data_plane_checkpoint_barrier.mutation( + "group_removals" + ) ) as cut: self._recovery_ledger.discard_group(cut, lineage_group_id) return RolloutOutcome.COMMITTED @@ -2038,7 +2089,7 @@ async def generate_for_finalization( owns_recovery_group = lineage_group_id is None recovery_group_id = lineage_group_id if recovery_group_id is None: - async with self._recovery_mutation() as cut: + async with self._recovery_mutation("prompt_reservations") as cut: recovery_group_id = self.reserve_prompt_group( cut, input_sample, @@ -2222,7 +2273,7 @@ async def _record_streamed_completion( pending_group_results[generation_index] = result if len(pending_group_results) < recovery_group.expected_generations: return - async with self._recovery_mutation() as cut: + async with self._recovery_mutation("sibling_seals") as cut: self._recovery_ledger.mark_group_sealed( cut, group_id, @@ -2230,7 +2281,7 @@ async def _record_streamed_completion( ) return - async with self._recovery_mutation() as cut: + async with self._recovery_mutation("sibling_seals") as cut: self._recovery_ledger.mark_sibling_sealed( cut, group_id, diff --git a/nemo_rl/experience/rollout_reassembler.py b/nemo_rl/experience/rollout_reassembler.py index 06b32e7bc99..c108bbdec8b 100644 --- a/nemo_rl/experience/rollout_reassembler.py +++ b/nemo_rl/experience/rollout_reassembler.py @@ -89,6 +89,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 the finalizer rejected the whole group as a structural outcome # (see drop_reason); the caller aborts the slot instead of committing it. @@ -556,6 +557,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, drop_reason=( @@ -639,6 +641,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, valid_row_count=len(valid_rows), total_row_count=len(rows), diff --git a/tests/unit/data_plane/test_rollout_reassembler.py b/tests/unit/data_plane/test_rollout_reassembler.py index e3cc8de2d58..025ee9acb32 100644 --- a/tests/unit/data_plane/test_rollout_reassembler.py +++ b/tests/unit/data_plane/test_rollout_reassembler.py @@ -228,6 +228,7 @@ def test_finalize_group_publishes_n_rows_with_placeholder(tq_client, partitions) assert finalized.metrics["finalize/terminal_selection_heuristic_fraction"] == 0.5 assert finalized.metrics["finalize/terminal_selection_declared_count"] == 0.0 assert finalized.metrics["finalize/terminal_witness_disagreement_count"] == 0.0 + 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() diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index e6c191df657..12e4af88520 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -445,6 +445,20 @@ async def _logged_commit(*args, **kwargs): assert start_v == 0 assert end_v == 0 assert len(mgr.recovery_ledger) == 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_ledger_hands_ownership_to_canonical_buffer_on_commit(self): buf = _FakeBuffer() diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 022754b0523..536d99b452f 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -134,6 +134,17 @@ def _consumed_meta(*sample_ids: str) -> KVBatchMeta: ) +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 + + # ── fakes ──────────────────────────────────────────────────────────────────── @@ -481,6 +492,12 @@ def __init__(self, events: Optional[list[str]] = None) -> None: self._tq_buffer = None self.recovery_ledger = RolloutRecoveryLedger() self._events = events + self.telemetry = { + "canonical_groups_finalized": 0, + "canonical_output_tokens": 0, + "recovery_siblings_reused": 0, + "recovery_siblings_redispatched": 0, + } def set_data_plane_checkpoint_barrier(self, barrier: Any) -> None: self.data_plane_checkpoint_barrier = barrier @@ -496,6 +513,17 @@ def resume_request_deadlines(self) -> None: if self._events is not None: self._events.append("resume_deadlines") + 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 + class _FakeTQBuffer: """TQReplayBuffer stand-in for the SC save/restore integration tests.""" @@ -527,7 +555,7 @@ def group_ids(self) -> tuple[str, ...]: def __len__(self) -> int: """Match the production TQReplayBuffer occupancy contract.""" - return len(self.target_step_list) + return self._num_groups def set_data_plane_checkpoint_barrier( self, barrier: DataPlaneCheckpointBarrier @@ -709,6 +737,7 @@ def _make_actor_args( last_checkpoint_path: Optional[str] = None, data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = None, bootstrap_identity: Optional[BootstrapCompatibilityIdentity] = None, + rollout_checkpoint_load_metrics: Optional[dict[str, float]] = None, ) -> SingleControllerActorArgs: return SingleControllerActorArgs( gen_handle=gen if gen is not None else _FakeGeneration(), @@ -737,6 +766,7 @@ def _make_actor_args( finalizer_actors=[], data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, bootstrap_identity=bootstrap_identity, + rollout_checkpoint_load_metrics=rollout_checkpoint_load_metrics, ) @@ -1326,6 +1356,92 @@ def test_pre_step_snapshot_contains_only_rollout_state(self, tmp_path: Path): assert (snapshot / ROLLOUT_RECOVERY_STATE_FILENAME).is_file() assert not (snapshot / "policy").exists() + def test_logs_snapshot_phase_durations(self, tmp_path: Path) -> None: + actor = self._actor(tmp_path) + actor._logger = MagicMock() + actor._telemetry_started_at = 0.0 + try: + with patch( + "nemo_rl.algorithms.single_controller.time.monotonic", + new=_SteppingClock(), + ): + assert asyncio.run(actor._save_rollout_checkpoint(force=True)) + finally: + actor._checkpointer.shutdown() + + 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 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._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 = SimpleNamespace( + 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() + + try: + with patch( + "nemo_rl.algorithms.single_controller.time.monotonic", + new=_SteppingClock(start=10.0, step=10.0), + ): + asyncio.run(sample_twice()) + finally: + actor._checkpointer.shutdown() + + 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_snapshot_reindexes_rows_owned_by_active_streamed_step( self, tmp_path: Path ): diff --git a/tests/unit/single_controller/test_finalizer_lifecycle.py b/tests/unit/single_controller/test_finalizer_lifecycle.py index 049e68855d4..eee55ac8a92 100644 --- a/tests/unit/single_controller/test_finalizer_lifecycle.py +++ b/tests/unit/single_controller/test_finalizer_lifecycle.py @@ -104,6 +104,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() From c522831b4ee0552814fe6ba1378e3b250a0ed278 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 29 Aug 2026 22:32:00 -0400 Subject: [PATCH 07/16] fix(rollout): address telemetry review findings Signed-off-by: Anish Mahishi --- ...po_math_1B_megatron_single_controller.yaml | 1 + .../algorithms/async_utils/replay_buffer.py | 9 ++-- nemo_rl/algorithms/single_controller.py | 19 ++++--- nemo_rl/experience/rollout_manager.py | 6 +-- nemo_rl/models/generation/interfaces.py | 7 ++- nemo_rl/utils/logger.py | 4 ++ .../test_checkpoint_dispatch_races.py | 3 ++ .../single_controller/test_checkpointing.py | 3 -- .../single_controller/test_rollout_pump.py | 3 ++ tests/unit/utils/test_logger.py | 52 +++++++++++++++++++ 10 files changed, 86 insertions(+), 21 deletions(-) diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index 38cf8d4ece6..d51e6bff8c4 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -156,6 +156,7 @@ checkpointing: # checkpoint support. save_period=1 provides an anchor after every train step. rollout_checkpointing: snapshot_attempt_interval_s: null + telemetry_interval_s: null keep_latest_k: 2 restore_mode: latest # Advanced escape hatch for runtime-only fields from external integrations. diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index baff24e1979..50ce965f93a 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -66,18 +66,22 @@ REPLAY_BUFFER_METADATA_STORAGE: Literal["tq_checkpoint"] = "tq_checkpoint" CheckpointMutationKind = Literal[ + "advantage_writeback", "group_commits", "group_removals", "other", "prompt_reservations", + "recovery_restore", "recovery_retries", "sample_clears", "sibling_seals", ] CHECKPOINT_MUTATION_KINDS: tuple[CheckpointMutationKind, ...] = ( + "advantage_writeback", "group_commits", "group_removals", "prompt_reservations", + "recovery_restore", "recovery_retries", "sample_clears", "sibling_seals", @@ -96,6 +100,7 @@ class DataPlaneCheckpointBarrierTelemetry: 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. @@ -1255,9 +1260,7 @@ async def commit( "the async message-log flattening path." ) trace_rollout_payload(keys=sample_ids, data=train_batch) - async with self._data_plane_checkpoint_barrier.mutation( - "group_commits" - ) as cut: + async with self._data_plane_checkpoint_barrier.mutation("group_commits") as cut: try: await call_data_plane( self._dp_client, diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 387ab56c71c..725df1a4d05 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -679,8 +679,9 @@ async def run(self) -> dict[str, Any]: # Loops forever like the watchdog, so finishing at all means it raised. await probe_task if rollout_telemetry_task is not None and rollout_telemetry_task in done: + # This pump has no normal return path. Awaiting it propagates the + # exception, including one concurrent with an orderly checkpoint stop. await rollout_telemetry_task - raise RuntimeError("rollout telemetry pump exited unexpectedly") if not stop_after_rollout_checkpoint and watchdog_task in done: # The watchdog loops forever, so finishing at all means it raised -- # a stall or an unhealthy environment. Surface that ahead of the @@ -922,7 +923,9 @@ async def _maybe_restore_rollout_recovery( ) recovery_ledger = self._rollout_manager.recovery_ledger - async with self._data_plane_checkpoint_barrier.mutation() as cut: + async with self._data_plane_checkpoint_barrier.mutation( + "recovery_restore" + ) as cut: recovery_ledger.load_state_dict(cut, parsed_state.ledger_state) recovery_ledger.prepare_for_restart(cut) self._batch_shortfall = parsed_state.batch_shortfall @@ -1903,10 +1906,8 @@ async def _dispatch_one_prompt( if self._rollout_recovery_enabled: assert lineage_group_id is not None - async with ( - self._data_plane_checkpoint_barrier.mutation( - "prompt_reservations" - ) + async with self._data_plane_checkpoint_barrier.mutation( + "prompt_reservations" ) as cut: replacement = self._take_replacement( target_step, replacements @@ -3859,9 +3860,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: if snapshot_cut.replay_metadata is not None else 0 ), - "ledger_groups": float( - snapshot_cut.rollout_recovery_group_count or 0 - ), + "ledger_groups": float(snapshot_cut.rollout_recovery_group_count or 0), } self._log_telemetry_metrics( checkpoint_metrics, @@ -4797,7 +4796,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: fields_to_put[adv_cfg.returns_field] = returns new_fields.append(adv_cfg.returns_field) - async with self._data_plane_checkpoint_barrier.mutation(): + async with self._data_plane_checkpoint_barrier.mutation("advantage_writeback"): await self._call_dp( "put_samples", sample_ids=meta.sample_ids, diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 41621d86310..9c85a4aa71d 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -2024,10 +2024,8 @@ async def generate_and_push( # succeeded on a retry also counts -- the fleet recovered either way. self._consecutive_infra_drops = 0 if lineage_group_id is not None: - async with ( - self._tq_buffer.data_plane_checkpoint_barrier.mutation( - "group_removals" - ) + async with self._tq_buffer.data_plane_checkpoint_barrier.mutation( + "group_removals" ) as cut: self._recovery_ledger.discard_group(cut, lineage_group_id) return RolloutOutcome.COMMITTED diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index e53141f81a5..d8dfb9cdcca 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -654,5 +654,10 @@ def get_step_metrics(self) -> dict[str, float]: return {} def drain_latest_logger_metrics(self) -> dict[str, Any]: - """Consume a bounded latest-value snapshot for frequent telemetry polls.""" + """Consume a bounded latest-value snapshot for frequent telemetry polls. + + Implementations may clear or compact their accumulated metric histories. + Callers must not assume that a later ``get_logger_metrics`` includes values + observed before this drain. + """ return self.get_logger_metrics() diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index 164f30c1108..77f2e5623bf 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -380,6 +380,10 @@ def define_metric( name: Name of the metric or pattern (e.g. 'ray/*') step_metric: Optional name of the step metric to use """ + if "*" in name and (not name.endswith("*") or name.count("*") != 1): + raise ValueError( + f"W&B metric patterns support exactly one trailing '*': {name!r}" + ) with self._log_lock: existing_step_metric = self._metric_step_patterns.get(name) if name in self._metric_step_patterns and ( diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 9eb3f3d7ab8..198b0250cdb 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -239,6 +239,9 @@ def discard_prompt_group( ) -> None: self.recovery_ledger.discard_group(cut, group_id) + def record_recovery_siblings(self, *, reused: int, redispatched: int) -> None: + del reused, redispatched + class _BlockingRolloutManager: """Hold one admitted rollout unfinished while the checkpoint is written.""" diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 536d99b452f..4b731b930ce 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -615,9 +615,6 @@ async def load_state_dict( ) return self.load_return - def __len__(self) -> int: - return self._num_groups - # Default position sentinel the fake dataloader reports via state_dict(). _SENTINEL_DL_STATE = {"fake_position": 42} diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 34db681229f..5fcd6ca1c99 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -105,6 +105,9 @@ def _init_pump_ledgers(ctrl: Any) -> None: ctrl._finalizer_actors = [] ctrl._replacement_reserve = deque() ctrl._rollout_recovery_enabled = False + ctrl._rollout_slot_waiters = 0 + ctrl._rollout_permitted_waiters = 0 + ctrl._buffer_capacity_waiters = 0 class _PausingMutationBarrier(DataPlaneCheckpointBarrier): diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py index 07bfb444be8..852b7e26bf5 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -577,6 +577,58 @@ def test_define_metric(self, mock_wandb): == 1 ) + @patch("nemo_rl.utils.logger.wandb") + def test_define_metric_rejects_conflicting_registration(self, mock_wandb): + logger = WandbLogger({}) + logger.define_metric("ray/*", step_metric="ray/ray_step") + + with pytest.raises(ValueError, match="already registered"): + logger.define_metric("ray/*", step_metric="other/step") + + @patch("nemo_rl.utils.logger.wandb") + def test_define_metric_rejects_non_terminal_wildcard(self, mock_wandb): + logger = WandbLogger({}) + + with pytest.raises(ValueError, match="exactly one trailing"): + logger.define_metric("ray/*/util", step_metric="ray/ray_step") + + @patch("nemo_rl.utils.logger.wandb") + def test_log_metrics_requires_registered_step_metric(self, mock_wandb): + logger = WandbLogger({}) + logger.define_metric( + "rollout/throughput/*", + step_metric="telemetry/wall_time_seconds", + ) + + with pytest.raises(ValueError, match="is missing from the logged event"): + logger.log_metrics( + {"rollout/throughput/tokens_per_second": 10.0}, + step=0, + ) + + @patch("nemo_rl.utils.logger.wandb") + def test_define_metric_uses_longest_matching_prefix(self, mock_wandb): + logger = WandbLogger({}) + logger.define_metric("rollout/*", step_metric="rollout/step") + logger.define_metric( + "rollout/throughput/*", + step_metric="telemetry/wall_time_seconds", + ) + + logger.log_metrics( + { + "telemetry/wall_time_seconds": 30.0, + "rollout/throughput/tokens_per_second": 10.0, + }, + step=0, + ) + + mock_run = mock_wandb.init.return_value + mock_run.define_metric.assert_any_call( + "rollout/throughput/tokens_per_second", + step_metric="telemetry/wall_time_seconds", + ) + @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.""" From e0a8b7570010e89e6d3a0e5f6af2711470252371 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 7 Sep 2026 12:01:47 -0400 Subject: [PATCH 08/16] fix(rollout): address telemetry review follow-ups Signed-off-by: Anish Mahishi --- docs/observability/metrics.md | 15 ++++++++ nemo_rl/algorithms/single_controller.py | 26 +++++++------- .../single_controller_utils/setup.py | 2 +- nemo_rl/experience/rollout_manager.py | 34 ++++++------------- .../test_rollout_generation_failures.py | 4 +++ tests/unit/experience/test_rollout_manager.py | 26 ++++++++++++-- .../test_rollout_reassembler_actor.py | 1 + .../experience/test_rollout_redispatch.py | 14 +++++++- .../test_checkpoint_dispatch_races.py | 13 +++++++ .../single_controller/test_rollout_pump.py | 9 +++-- .../test_tq_replay_buffer.py | 2 ++ 11 files changed, 101 insertions(+), 45 deletions(-) diff --git a/docs/observability/metrics.md b/docs/observability/metrics.md index 00ba9609e10..ea73500cdfe 100644 --- a/docs/observability/metrics.md +++ b/docs/observability/metrics.md @@ -6,6 +6,21 @@ Metrics are emitted **only when telemetry is exporting** — the driver always e Training scalars — reward, loss, KL, grad norm, learning rate, throughput — are **not** mirrored to OTel. nemo-lens declares `record_rl_metrics` gauges for most of them, plus `rl.generation.duration_ms` and `rl.rollout.duration_ms` histograms, but NeMo-RL emits none of them: mapping its logger keys onto lens's fixed fields is still being settled with the lens owners. Read those scalars from W&B / TensorBoard, and phase durations from the spans. +## W&B history step axes + +W&B's internal `_step` is a monotonically increasing history-row number, not a +NeMo-RL trainer step. Trainer-correlated metrics are buffered into one history +row and carry the explicit `nemo_rl/step` field. Use `nemo_rl/step` when plotting +or joining training metrics by optimizer step; relying on `_step` also counts +independently committed telemetry rows and can misalign the series. + +Independent event streams use their own custom axes and do not carry +`nemo_rl/step`. In particular, Single-Controller rollout benchmark series under +`rollout/throughput/*`, `timing/rollout_checkpoint/*`, +`timing/rollout_recovery/*`, and `rollout/checkpoint_outcome/*` use +`telemetry/wall_time_seconds`. This lets those series continue through a long +or paused trainer step without changing the meaning of the trainer-step axis. + ## Async efficiency metrics (`rl.efficiency.*`) Async GRPO measures where wall time goes with a `Timer` and logs the result as `efficiency/*` scalars (`print_efficiency_summary` in `nemo_rl/algorithms/utils.py`). Those same values are teed to OTel as one **dimensioned** gauge rather than one instrument per category, so adding a category needs no instrument change. diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 725df1a4d05..84580a86c32 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1136,12 +1136,6 @@ async def _redispatch_restored_rollouts( reused_siblings = 0 redispatched_siblings = 0 - def _record_sibling_work(group: PromptGroupRecoveryRecord) -> None: - nonlocal reused_siblings, redispatched_siblings - reused = len(group.sealed_generation_indices) - reused_siblings += reused - redispatched_siblings += group.expected_generations - reused - recognized_phases = ( PromptGroupPhase.ADMITTED, PromptGroupPhase.RESERVED, @@ -1167,7 +1161,9 @@ def _record_sibling_work(group: PromptGroupRecoveryRecord) -> None: group.target_step, group.group_id, ) - _record_sibling_work(group) + reused = len(group.sealed_generation_indices) + reused_siblings += reused + redispatched_siblings += group.expected_generations - reused redispatched += 1 # A checkpoint may land after dataloader ownership is recorded but before @@ -1190,7 +1186,9 @@ def _record_sibling_work(group: PromptGroupRecoveryRecord) -> None: group.target_step, group.group_id, ) - _record_sibling_work(group) + reused = len(group.sealed_generation_indices) + reused_siblings += reused + redispatched_siblings += group.expected_generations - reused redispatched += 1 self._rollout_manager.record_recovery_siblings( @@ -2488,12 +2486,12 @@ async def _train_pump(self) -> None: await asyncio.sleep(0) # Evict stale groups - async with self._data_plane_checkpoint_barrier.mutation( - "group_removals" - ): - evicted = await self._sampler.evict( - current_train_weight=self._trainer_version, - ) + # TQReplayBuffer.remove() owns the group-removal mutation + # cut. Acquiring another cut here would nest the same + # non-reentrant barrier section in this task. + evicted = await self._sampler.evict( + current_train_weight=self._trainer_version, + ) evicted_stale_prompt_groups += evicted if evicted: print( diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index b52382741b7..614daec6370 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -1053,7 +1053,7 @@ def setup_single_controller( stacklevel=2, ) else: - vllm_cfg = generation_config["vllm_cfg"] + vllm_cfg = cast(dict[str, Any], generation_config)["vllm_cfg"] if not vllm_cfg.get("enable_vllm_metrics_logger"): warnings.warn( "rollout_checkpointing.telemetry_interval_s is enabled, but " diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 9c85a4aa71d..59fdf257139 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -1680,33 +1680,21 @@ async def _recovery_mutation( 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 - ), + "canonical_groups_finalized": self._canonical_groups_finalized, + "canonical_output_tokens": self._canonical_output_tokens, + "recovery_siblings_reused": self._recovery_siblings_reused, + "recovery_siblings_redispatched": self._recovery_siblings_redispatched, } 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)) + self._canonical_groups_finalized += 1 + self._canonical_output_tokens += max(0, int(output_tokens)) def record_recovery_siblings(self, *, reused: int, redispatched: int) -> None: """Count sibling work avoided and repeated after a process restart.""" - 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)) + self._recovery_siblings_reused += max(0, int(reused)) + self._recovery_siblings_redispatched += max(0, int(redispatched)) def reserve_prompt_group( self, @@ -2008,13 +1996,11 @@ async def generate_and_push( raise self._stats.committed += 1 - rollout_metrics = getattr(record, "rollout_metrics", {}) + rollout_metrics = 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", ()) - ) + total_output_tokens = float(mean_output_tokens) * len(record.completions) if math.isfinite(total_output_tokens): output_tokens = max(0, round(total_output_tokens)) self.record_canonical_publication(output_tokens) diff --git a/tests/unit/experience/test_rollout_generation_failures.py b/tests/unit/experience/test_rollout_generation_failures.py index 4ae7944fabe..29aa66bbb1b 100644 --- a/tests/unit/experience/test_rollout_generation_failures.py +++ b/tests/unit/experience/test_rollout_generation_failures.py @@ -181,6 +181,10 @@ def _make_manager(buffer, impl, retry_policy=None) -> RolloutManager: else RolloutRetryPolicy.single_attempt() ) manager._stats = RolloutStats() + manager._canonical_groups_finalized = 0 + manager._canonical_output_tokens = 0 + manager._recovery_siblings_reused = 0 + manager._recovery_siblings_redispatched = 0 manager._skipped_prompts = 0 manager._consecutive_infra_drops = 0 return manager diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 12e4af88520..2f7496e1974 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -221,10 +221,21 @@ async def remove_group(self, group_id: str, *, remove_in_dp: bool = False) -> in class _FakeImpl: - """Stand-in for AsyncRolloutImpl that returns a sentinel record.""" + """Stand-in for AsyncRolloutImpl that returns a typed sentinel record.""" def __init__(self, record="sentinel-record", on_run=None) -> None: - self._record = record + self._record = ( + record + if isinstance(record, PromptGroupRecord) + else PromptGroupRecord( + prompt_idx=0, + prompt=[], + extra_env_info=None, + metadata={"sentinel": record}, + completions=[], + rollout_metrics={}, + ) + ) self._on_run = on_run async def run_rollout(self, input_sample): @@ -257,6 +268,10 @@ def _make_manager( else RolloutRetryPolicy.single_attempt() ) mgr._stats = RolloutStats() + mgr._canonical_groups_finalized = 0 + mgr._canonical_output_tokens = 0 + mgr._recovery_siblings_reused = 0 + mgr._recovery_siblings_redispatched = 0 mgr._skipped_prompts = 0 mgr._consecutive_infra_drops = 0 return mgr @@ -441,7 +456,8 @@ async def _logged_commit(*args, **kwargs): assert len(buf.commit_calls) == 1 gid, record, start_v, end_v = buf.commit_calls[0] assert gid in buf._slots - assert record == "r0" + assert isinstance(record, PromptGroupRecord) + assert record.metadata["sentinel"] == "r0" assert start_v == 0 assert end_v == 0 assert len(mgr.recovery_ledger) == 0 @@ -1735,6 +1751,10 @@ def _make_capture_manager( else RolloutRetryPolicy.single_attempt() ) mgr._stats = RolloutStats() + mgr._canonical_groups_finalized = 0 + mgr._canonical_output_tokens = 0 + mgr._recovery_siblings_reused = 0 + mgr._recovery_siblings_redispatched = 0 mgr._skipped_prompts = 0 mgr._consecutive_infra_drops = 0 mgr._recovery_ledger = RolloutRecoveryLedger() diff --git a/tests/unit/experience/test_rollout_reassembler_actor.py b/tests/unit/experience/test_rollout_reassembler_actor.py index a31c22777d9..4461ba0e5c1 100644 --- a/tests/unit/experience/test_rollout_reassembler_actor.py +++ b/tests/unit/experience/test_rollout_reassembler_actor.py @@ -140,6 +140,7 @@ def test_rpc_dataclass_fields_are_classified() -> None: "group_min_wv", "group_max_wv", "staging_keys", + "canonical_output_tokens", "metrics", "dropped", "drop_reason", diff --git a/tests/unit/experience/test_rollout_redispatch.py b/tests/unit/experience/test_rollout_redispatch.py index 81b215f3575..b1fe0909f34 100644 --- a/tests/unit/experience/test_rollout_redispatch.py +++ b/tests/unit/experience/test_rollout_redispatch.py @@ -49,6 +49,7 @@ from nemo_rl.experience.interfaces import ( NEMO_GYM_GROUP_ATTEMPT_KEY, NEMO_GYM_GROUP_ID_KEY, + PromptGroupRecord, ) from nemo_rl.experience.rollout_manager import ( RolloutManager, @@ -99,7 +100,14 @@ async def run_rollout(self, input_sample): failure = self._failures[index] if failure is not None: raise failure - return f"record-{index}" + return PromptGroupRecord( + prompt_idx=0, + prompt=[], + extra_env_info=None, + metadata={"sentinel": f"record-{index}"}, + completions=[], + rollout_metrics={}, + ) def _make_manager(buffer, impl, policy) -> RolloutManager: @@ -111,6 +119,10 @@ def _make_manager(buffer, impl, policy) -> RolloutManager: manager._weight_version = 0 manager._retry_policy = policy manager._stats = RolloutStats() + manager._canonical_groups_finalized = 0 + manager._canonical_output_tokens = 0 + manager._recovery_siblings_reused = 0 + manager._recovery_siblings_redispatched = 0 manager._skipped_prompts = 0 manager._consecutive_infra_drops = 0 return manager diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 198b0250cdb..1ba41a4680b 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -37,6 +37,7 @@ from pathlib import Path from types import SimpleNamespace from typing import Any, TypeVar, cast +from unittest.mock import MagicMock import pytest import torch @@ -86,6 +87,14 @@ async def apply() -> _T: return asyncio.run(apply()) +def _init_recovery_telemetry(controller: Any, *, train_steps: int = 0) -> None: + """Initialize constructor-owned telemetry state for hand-built controllers.""" + controller._train_steps = train_steps + controller._telemetry_sample_index = 0 + controller._telemetry_started_at = 0.0 + controller._logger = MagicMock() + + async def _wait_for_event_or_pump( event: asyncio.Event, pump: asyncio.Task[None], @@ -916,6 +925,7 @@ async def exercise() -> None: controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller = object.__new__(controller_cls) + _init_recovery_telemetry(controller, train_steps=7) controller._sampler = sampler controller._rollout_manager = rollout_manager controller._master_config = SimpleNamespace( @@ -978,6 +988,7 @@ async def exercise() -> None: ) controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller = object.__new__(controller_cls) + _init_recovery_telemetry(controller) controller._rollout_manager = SimpleNamespace(recovery_ledger=recovery_ledger) launched = False @@ -1026,6 +1037,7 @@ async def exercise() -> None: rollout_manager = _RecoveryRolloutManager(RolloutRecoveryLedger()) controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller = object.__new__(controller_cls) + _init_recovery_telemetry(controller, train_steps=7) controller._sampler = sampler controller._rollout_manager = rollout_manager controller._master_config = SimpleNamespace( @@ -1117,6 +1129,7 @@ async def exercise() -> None: controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller = object.__new__(controller_cls) + _init_recovery_telemetry(controller, train_steps=6) controller._sampler = sampler controller._rollout_manager = rollout_manager controller._trainer_version = 6 diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 5fcd6ca1c99..994117806aa 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -28,6 +28,7 @@ import torch from nemo_rl.algorithms.async_utils.replay_buffer import ( + CheckpointMutationKind, DataPlaneCheckpointBarrier, DataPlaneMutationCut, TQReplayBuffer, @@ -108,6 +109,8 @@ def _init_pump_ledgers(ctrl: Any) -> None: ctrl._rollout_slot_waiters = 0 ctrl._rollout_permitted_waiters = 0 ctrl._buffer_capacity_waiters = 0 + ctrl._rollout_completion_durations_s = deque(maxlen=10_000) + ctrl._rollout_queue_wait_durations_s = deque(maxlen=10_000) class _PausingMutationBarrier(DataPlaneCheckpointBarrier): @@ -119,8 +122,10 @@ def __init__(self) -> None: self.release_mutation = asyncio.Event() @asynccontextmanager - async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: - async with super().mutation() as cut: + async def mutation( + self, kind: CheckpointMutationKind = "other" + ) -> AsyncIterator[DataPlaneMutationCut]: + async with super().mutation(kind) as cut: yield cut self.mutation_applied.set() await self.release_mutation.wait() diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 77ff988f4a9..6e5538b8d06 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -447,6 +447,8 @@ async def exercise() -> None: async with barrier.mutation() as cut: cut.require_live() + asyncio.run(exercise()) + def test_reports_mutations_blocked_by_checkpoint(self): async def exercise() -> None: barrier = DataPlaneCheckpointBarrier() From 7c123822f4eb6e2a3a6d1eaf2572aa0b787e1974 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 7 Sep 2026 12:25:39 -0400 Subject: [PATCH 09/16] feat(rollout): focus checkpoint recovery telemetry Signed-off-by: Anish Mahishi --- docs/observability/metrics.md | 30 +++ nemo_rl/algorithms/single_controller.py | 223 ++++++++++++++---- .../single_controller_utils/setup.py | 6 +- nemo_rl/models/generation/interfaces.py | 5 +- .../single_controller/test_checkpointing.py | 114 ++++++++- 5 files changed, 324 insertions(+), 54 deletions(-) diff --git a/docs/observability/metrics.md b/docs/observability/metrics.md index ea73500cdfe..3ad5bad1620 100644 --- a/docs/observability/metrics.md +++ b/docs/observability/metrics.md @@ -21,6 +21,36 @@ Independent event streams use their own custom axes and do not carry `telemetry/wall_time_seconds`. This lets those series continue through a long or paused trainer step without changing the meaning of the trainer-step axis. +## Single-Controller rollout recovery metrics + +The rollout checkpoint benchmark focuses on the following questions: + +1. Did checkpointing reduce raw generation or usable rollout throughput? +2. Which save or restore phase took the time? +3. How much live data-plane work waited behind the checkpoint barrier? +4. Did scheduled checkpoint attempts succeed at the intended cadence? +5. How much completed rollout work was reused after restart? + +| Prefix | Important fields | Meaning | +|---|---|---| +| `rollout/throughput` | `generation_output_tokens_per_second`, `canonical_output_tokens_per_second`, `canonical_groups_per_second` | Raw backend decoding throughput compared with finalized token and group throughput available for training. The raw metric is absent when the backend does not expose compatible cumulative counters. | +| `rollout/throughput` | `checkpoint_blocked_mutations`, `checkpoint_mutation_wait_seconds_p95`, `checkpoint_mutation_wait_seconds_max` | Number and latency of live data-plane mutations delayed by an exclusive checkpoint. | +| `timing/rollout_checkpoint` | `total_save_seconds`, `tq_save_seconds`, `barrier_wait_seconds`, `exclusive_hold_seconds`, `sidecar_save_seconds`, `snapshot_commit_seconds` | End-to-end save latency and its storage, fencing, controller-sidecar, and atomic-publication components. | +| `timing/rollout_checkpoint` | `snapshot_rows`, `replay_rows`, `staging_rows`, `replay_groups`, `ledger_groups`, `controller_sidecar_bytes` | Logical volume captured by the snapshot. `controller_sidecar_bytes` excludes the native TQ payload because the current TQ checkpoint API does not report bytes written. NeMo-RL deliberately does not recursively scan the shared checkpoint directory because that scan would perturb the benchmark. | +| `rollout/checkpoint_outcome` | `completed`, `skipped`, `failed`, `reason_*`, `seconds_since_previous_success`, `seconds_since_last_success` | Result, actionable reason, and effective cadence of every scheduled checkpoint attempt. | +| `timing/rollout_recovery` | `snapshot_resolution_seconds`, `dataloader_load_seconds`, `tq_load_seconds`, `replay_metadata_load_seconds`, `recovery_prepare_seconds`, `total_load_seconds` | Rollout-state restore latency. `total_load_seconds` is the sum of these non-overlapping restore phases. | +| `timing/rollout_recovery` | `groups_reused`, `groups_redispatched`, `siblings_reused`, `siblings_redispatched`, `redispatch_schedule_seconds` | Completed groups restored without generation and unfinished sibling work preserved or repeated after restart. | + +`barrier_wait_seconds` and mutation wait latency measure opposite sides of the +same fence. The former is how long the checkpoint waits for already-running +mutations; the latter is how long rollout or training mutations wait for the +checkpoint to release its exclusive cut. + +The vLLM request, KV-cache, controller-waiter, buffer-occupancy, and per-mutation +kind fields are diagnostic drill-down signals. They are useful when one of the +primary throughput or barrier metrics regresses, but do not need to appear on +the primary rollout checkpoint dashboard. + ## Async efficiency metrics (`rl.efficiency.*`) Async GRPO measures where wall time goes with a `Timer` and logs the result as `efficiency/*` scalars (`print_efficiency_summary` in `nemo_rl/algorithms/utils.py`). Those same values are teed to OTel as one **dimensioned** gauge rather than one instrument per category, so adding a category needs no instrument change. diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 84580a86c32..d944cceb381 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -58,7 +58,16 @@ from dataclasses import dataclass from functools import partial from pathlib import Path -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Literal, + Optional, + Union, + cast, +) import ray import torch @@ -136,7 +145,6 @@ ROLLOUT_RECOVERY_SCHEMA_VERSION, ROLLOUT_RECOVERY_STATE_FILENAME, PromptGroupPhase, - PromptGroupRecoveryRecord, RolloutRecoveryState, build_rollout_recovery_state, parse_rollout_recovery_state, @@ -172,6 +180,25 @@ _MAX_CONSECUTIVE_ROLLOUT_CHECKPOINT_FAILURES = 3 +RolloutCheckpointAttemptReason = Literal[ + "completed", + "invariant_error", + "io_error", + "missing_trainer_anchor", + "no_data_plane_mutations", + "optimizer_commit_in_progress", + "timeout", + "trainer_state_changed", +] + + +@dataclass(frozen=True) +class _RolloutCheckpointSaveResult: + """Outcome returned by one rollout checkpoint save attempt.""" + + saved: bool + reason: RolloutCheckpointAttemptReason + @dataclass(frozen=True) class _RolloutCheckpointCut: @@ -183,13 +210,17 @@ class _RolloutCheckpointCut: replay_metadata: Optional[TQReplayMetadataState] rollout_recovery_payload: Optional[bytes] rollout_recovery_group_count: Optional[int] + replay_row_count: int + staging_row_count: 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.""" +def _latest_generation_values( + metrics: dict[str, Any], metric_name: str +) -> dict[int, float]: + """Return the latest value keyed by generation data-parallel worker.""" per_worker = metrics.get(metric_name) if not isinstance(per_worker, dict): return {} @@ -603,16 +634,11 @@ async def run(self) -> dict[str, Any]: ) await self._maybe_restore_replacement_reserve() 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", - ) + self._log_rollout_restore_metrics( + replay_metadata_load_seconds=replay_restore_seconds, + recovery_prepare_seconds=recovery_prepare_seconds, + restored_replay_groups=restored_replay_groups, + ) # Start the rollout and train pumps, plus the watchdog rollout_task = asyncio.create_task(self._rollout_pump()) @@ -761,6 +787,37 @@ def _record_rollout_timing( self._rollout_queue_wait_durations_s.append(dispatch_started - work_started) self._rollout_completion_durations_s.append(time.monotonic() - dispatch_started) + def _log_rollout_restore_metrics( + self, + *, + replay_metadata_load_seconds: float, + recovery_prepare_seconds: float, + restored_replay_groups: int, + ) -> None: + """Log rollout restore phases after the controller is ready to dispatch.""" + if self._rollout_checkpoint_load_metrics is None: + return + load_metrics = dict(self._rollout_checkpoint_load_metrics) + load_metrics["replay_metadata_load_seconds"] = replay_metadata_load_seconds + load_metrics["recovery_prepare_seconds"] = recovery_prepare_seconds + load_metrics["total_load_seconds"] = sum( + load_metrics[key] + for key in ( + "snapshot_resolution_seconds", + "dataloader_load_seconds", + "tq_load_seconds", + "replay_metadata_load_seconds", + "recovery_prepare_seconds", + ) + if key in load_metrics + ) + load_metrics["groups_reused"] = float(restored_replay_groups) + self._log_telemetry_metrics( + load_metrics, + step=self._train_steps, + prefix="timing/rollout_recovery", + ) + # ── internal helpers ─────────────────────────────────────────────────── async def _maybe_restore_replay_buffer(self) -> int: @@ -1258,7 +1315,7 @@ async def _validate_rollout_recovery_inventory( *, replay_metadata: Optional[TQReplayMetadataState], clear_unreferenced: bool, - ) -> None: + ) -> int: """Validate staging ownership while the caller holds a stable cut.""" cut.require_live() expected_staging_keys = self._rollout_recovery_ledger.expected_staging_keys() @@ -1302,6 +1359,7 @@ async def _validate_rollout_recovery_inventory( f"referenced={len(expected_staging_keys)}", flush=True, ) + return len(expected_staging_keys) async def _maybe_restore_replacement_reserve(self) -> None: """Restore spare prompts diverted before the previous run's checkpoint. @@ -3647,6 +3705,9 @@ async def _capture_rollout_checkpoint_cut( additional_groups=training_owned_groups, ) await self._validate_replay_inventory(replay_metadata) + replay_row_count = sum( + len(group["meta"].sample_ids) for group in replay_metadata["groups"] + ) recovery_state = self._rollout_manager.recovery_ledger.state_dict() recovery_state["batch_shortfall"] = self._batch_shortfall.copy() @@ -3664,8 +3725,9 @@ async def _capture_rollout_checkpoint_cut( recovery_payload = payload_buffer.getvalue() recovery_digest = hashlib.sha256(recovery_payload).hexdigest() + staging_row_count = 0 if self._master_config.token_capture.enabled: - await self._validate_rollout_recovery_inventory( + staging_row_count = await self._validate_rollout_recovery_inventory( cut, replay_metadata=replay_metadata, clear_unreferenced=False, @@ -3688,6 +3750,8 @@ async def _capture_rollout_checkpoint_cut( replay_metadata=replay_metadata, rollout_recovery_payload=recovery_payload, rollout_recovery_group_count=len(recovery_state["groups"]), + replay_row_count=replay_row_count, + staging_row_count=staging_row_count, rolled_back_train_group_count=len(training_owned_groups), mutation_version=self._data_plane_checkpoint_barrier.mutation_version, tq_save_seconds=tq_save_seconds, @@ -3697,51 +3761,74 @@ async def _write_rollout_checkpoint_sidecars( self, checkpoint_path: Path, cut: _RolloutCheckpointCut, - ) -> None: - """Write metadata-only controller state beside a native TQ snapshot.""" + ) -> int: + """Write controller state beside TQ and return its on-disk byte size.""" + written_paths: list[Path] = [] + dataloader_path = checkpoint_path / "train_dataloader.pt" await asyncio.to_thread( torch.save, cut.dataloader_state, - checkpoint_path / "train_dataloader.pt", + dataloader_path, ) + written_paths.append(dataloader_path) if cut.replacement_reserve: + replacement_reserve_path = checkpoint_path / REPLACEMENT_RESERVE_FILENAME await asyncio.to_thread( torch.save, cut.replacement_reserve, - checkpoint_path / "replacement_reserve.pt", + replacement_reserve_path, ) + written_paths.append(replacement_reserve_path) if cut.replay_metadata is not None: + replay_metadata_path = checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME await asyncio.to_thread( torch.save, cut.replay_metadata, - checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME, + replay_metadata_path, ) + written_paths.append(replay_metadata_path) if cut.rollout_recovery_payload is not None: + rollout_recovery_path = checkpoint_path / ROLLOUT_RECOVERY_STATE_FILENAME await asyncio.to_thread( - (checkpoint_path / ROLLOUT_RECOVERY_STATE_FILENAME).write_bytes, + rollout_recovery_path.write_bytes, cut.rollout_recovery_payload, ) + written_paths.append(rollout_recovery_path) + + config_path = checkpoint_path / "config.yaml" def _write_config() -> None: import yaml dumped = self._master_config.model_dump(mode="json") - with (checkpoint_path / "config.yaml").open("w") as config_file: + with config_path.open("w") as config_file: yaml.safe_dump(dumped, config_file) await asyncio.to_thread(_write_config) + written_paths.append(config_path) + return await asyncio.to_thread( + lambda: sum(path.stat().st_size for path in written_paths) + ) - async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: + async def _save_rollout_checkpoint( + self, *, force: bool = False + ) -> _RolloutCheckpointSaveResult: """Publish one rollout-only snapshot anchored to durable trainer state.""" async with self._checkpoint_save_lock: if self._optimizer_commit_in_progress: - return False + return _RolloutCheckpointSaveResult( + saved=False, + reason="optimizer_commit_in_progress", + ) if ( not force and self._last_rollout_snapshot_mutation_version == self._data_plane_checkpoint_barrier.mutation_version ): - return False + return _RolloutCheckpointSaveResult( + saved=False, + reason="no_data_plane_mutations", + ) save_started = time.monotonic() await asyncio.to_thread(self._checkpointer.finalize_pending) @@ -3777,7 +3864,10 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: flush=True, ) self._last_missing_rollout_snapshot_anchor = skip_key - return False + return _RolloutCheckpointSaveResult( + saved=False, + reason="missing_trainer_anchor", + ) try: await asyncio.to_thread( prune_bootstrap_snapshots, @@ -3807,14 +3897,23 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: or self._trainer_version != expected_trainer_version ): await asyncio.to_thread(partial(shutil.rmtree, tmp_path)) - return False + return _RolloutCheckpointSaveResult( + saved=False, + reason="trainer_state_changed", + ) snapshot_epoch = self._current_epoch snapshot_cut = await self._capture_rollout_checkpoint_cut( cut, tmp_path ) barrier_released = time.monotonic() - await self._write_rollout_checkpoint_sidecars(tmp_path, snapshot_cut) + sidecar_save_started = time.monotonic() + controller_sidecar_bytes = ( + await self._write_rollout_checkpoint_sidecars( + tmp_path, + snapshot_cut, + ) + ) manifest = RolloutSnapshotManifest( schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, base_train_step=expected_train_step, @@ -3827,10 +3926,16 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: ), bootstrap_fingerprint=snapshot_fingerprint, ) + manifest_text = ( + json.dumps(manifest.to_dict(), sort_keys=True, indent=2) + "\n" + ) await asyncio.to_thread( (tmp_path / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text, - json.dumps(manifest.to_dict(), sort_keys=True, indent=2) + "\n", + manifest_text, ) + controller_sidecar_bytes += len(manifest_text.encode()) + sidecar_save_seconds = time.monotonic() - sidecar_save_started + snapshot_commit_started = time.monotonic() await asyncio.to_thread( commit_snapshot, tmp_path, @@ -3839,6 +3944,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: self._master_config.rollout_checkpointing.keep_latest_k ), ) + snapshot_commit_seconds = time.monotonic() - snapshot_commit_started except BaseException: if tmp_path.exists(): await asyncio.to_thread(partial(shutil.rmtree, tmp_path)) @@ -3853,12 +3959,20 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: "tq_save_seconds": snapshot_cut.tq_save_seconds, "barrier_wait_seconds": barrier_acquired - barrier_requested, "exclusive_hold_seconds": barrier_released - barrier_acquired, + "sidecar_save_seconds": sidecar_save_seconds, + "snapshot_commit_seconds": snapshot_commit_seconds, "replay_groups": float( len(snapshot_cut.replay_metadata["groups"]) if snapshot_cut.replay_metadata is not None else 0 ), "ledger_groups": float(snapshot_cut.rollout_recovery_group_count or 0), + "replay_rows": float(snapshot_cut.replay_row_count), + "staging_rows": float(snapshot_cut.staging_row_count), + "snapshot_rows": float( + snapshot_cut.replay_row_count + snapshot_cut.staging_row_count + ), + "controller_sidecar_bytes": float(controller_sidecar_bytes), } self._log_telemetry_metrics( checkpoint_metrics, @@ -3875,10 +3989,14 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: f"{checkpoint_metrics['exclusive_hold_seconds']:.2f})", flush=True, ) - return True + return _RolloutCheckpointSaveResult(saved=True, reason="completed") def _log_rollout_checkpoint_outcome( - self, *, outcome: str, attempt_duration_seconds: float + self, + *, + outcome: Literal["completed", "failed", "skipped"], + reason: RolloutCheckpointAttemptReason, + attempt_duration_seconds: float, ) -> None: """Record every scheduled checkpoint attempt, including no-op cuts.""" metrics = { @@ -3891,14 +4009,19 @@ def _log_rollout_checkpoint_outcome( self._master_config.rollout_checkpointing.snapshot_attempt_interval_s or 0.0 ), + f"reason_{reason}": 1.0, } + completed_at = time.monotonic() 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 + elif self._last_successful_rollout_checkpoint_time is not None: + metrics["seconds_since_last_success"] = ( + completed_at - self._last_successful_rollout_checkpoint_time + ) self._log_telemetry_metrics( metrics, step=self._train_steps, @@ -3918,12 +4041,21 @@ async def _rollout_checkpoint_pump(self) -> None: 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 (OSError, TimeoutError) as error: + result = await self._save_rollout_checkpoint(force=deadline_due) + except Exception as error: + if isinstance(error, TimeoutError): + failure_reason: RolloutCheckpointAttemptReason = "timeout" + elif isinstance(error, OSError): + failure_reason = "io_error" + else: + failure_reason = "invariant_error" self._log_rollout_checkpoint_outcome( outcome="failed", + reason=failure_reason, attempt_duration_seconds=time.monotonic() - attempt_started, ) + if not isinstance(error, (OSError, TimeoutError)): + raise if deadline_due: raise RuntimeError( "failed to save the required pre-step rollout checkpoint" @@ -3944,10 +4076,11 @@ async def _rollout_checkpoint_pump(self) -> None: continue consecutive_failures = 0 self._log_rollout_checkpoint_outcome( - outcome="completed" if saved else "skipped", + outcome="completed" if result.saved else "skipped", + reason=result.reason, attempt_duration_seconds=time.monotonic() - attempt_started, ) - if deadline_due and saved and self._timeout.check_save(): + if deadline_due and result.saved and self._timeout.check_save(): print( "Checkpoint deadline reached before the first train step; " "stopping after a durable rollout snapshot", @@ -3989,7 +4122,9 @@ async def _collect_and_log_rollout_throughput_metrics( stacklevel=2, ) - generated_values = _latest_vllm_values(generation_metrics, "generation_tokens") + generated_values = _latest_generation_values( + generation_metrics, "generation_tokens" + ) generated_by_worker = ( {worker_id: int(value) for worker_id, value in generated_values.items()} if generated_values @@ -4042,9 +4177,9 @@ async def _collect_and_log_rollout_throughput_metrics( ) 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") + running = _latest_generation_values(generation_metrics, "inflight_batch_sizes") + waiting = _latest_generation_values(generation_metrics, "num_pending_samples") + kv_usage = _latest_generation_values(generation_metrics, "kv_cache_usage_perc") if running: metrics["vllm_requests_running"] = sum(running.values()) if waiting: @@ -4052,7 +4187,7 @@ async def _collect_and_log_rollout_throughput_metrics( 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) + metrics["generation_output_tokens"] = float(generated_tokens) completion_durations = list(self._rollout_completion_durations_s) queue_wait_durations = list(self._rollout_queue_wait_durations_s) @@ -4105,11 +4240,11 @@ async def _collect_and_log_rollout_throughput_metrics( - previous_generated[worker_id] for worker_id in current_workers ) - metrics["vllm_output_tokens_per_second"] = ( + metrics["generation_output_tokens_per_second"] = ( generated_delta / elapsed ) else: - metrics["vllm_counter_discontinuity"] = 1.0 + metrics["generation_counter_discontinuity"] = 1.0 self._throughput_sample_time = now self._throughput_generation_tokens_by_worker = generated_by_worker diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 614daec6370..5c8a2f3bf92 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -1228,6 +1228,7 @@ def setup_single_controller( if save_state.trainer_version is not None else save_state.current_step ) + snapshot_resolution_started = time.monotonic() if ( trainer_checkpoint_path is not None and rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None @@ -1267,6 +1268,7 @@ def setup_single_controller( expected_trainer_version=0, expected_bootstrap_fingerprint=bootstrap_digest, ) + snapshot_resolution_seconds = time.monotonic() - snapshot_resolution_started if resolved_snapshot is not None: recovery_checkpoint_path = str(resolved_snapshot.path) save_state.current_epoch = resolved_snapshot.manifest.current_epoch @@ -1291,7 +1293,9 @@ def setup_single_controller( or (recovery_path / ROLLOUT_RECOVERY_STATE_FILENAME).is_file() ) rollout_checkpoint_load_metrics: Optional[dict[str, float]] = ( - {} if has_rollout_checkpoint_payload else None + {"snapshot_resolution_seconds": snapshot_resolution_seconds} + if has_rollout_checkpoint_payload + else None ) # ========================== diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index d8dfb9cdcca..4b35a4f8197 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -658,6 +658,9 @@ def drain_latest_logger_metrics(self) -> dict[str, Any]: Implementations may clear or compact their accumulated metric histories. Callers must not assume that a later ``get_logger_metrics`` includes values - observed before this drain. + observed before this drain. Backends supporting raw rollout throughput + should return cumulative sampled-token counters under ``generation_tokens`` + as ``data_parallel_worker_id -> list[counter]``. The controller computes + per-worker deltas before summing them, so counter resets are detectable. """ return self.get_logger_metrics() diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 4b731b930ce..97294723297 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -39,6 +39,7 @@ import json import os import threading +import time from collections.abc import Callable from pathlib import Path from types import SimpleNamespace @@ -1331,7 +1332,9 @@ def test_pre_step_snapshot_contains_only_rollout_state(self, tmp_path: Path): actor = self._actor(tmp_path) try: actor._sampler.restore_dispatch_index(5) - assert asyncio.run(actor._save_rollout_checkpoint(force=True)) + result = asyncio.run(actor._save_rollout_checkpoint(force=True)) + assert result.saved + assert result.reason == "completed" finally: actor._checkpointer.shutdown() @@ -1362,7 +1365,8 @@ def test_logs_snapshot_phase_durations(self, tmp_path: Path) -> None: "nemo_rl.algorithms.single_controller.time.monotonic", new=_SteppingClock(), ): - assert asyncio.run(actor._save_rollout_checkpoint(force=True)) + result = asyncio.run(actor._save_rollout_checkpoint(force=True)) + assert result.saved finally: actor._checkpointer.shutdown() @@ -1371,12 +1375,73 @@ def test_logs_snapshot_phase_durations(self, tmp_path: Path) -> None: 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 0 <= logged["sidecar_save_seconds"] <= logged["total_save_seconds"] + assert 0 <= logged["snapshot_commit_seconds"] <= logged["total_save_seconds"] + assert logged["controller_sidecar_bytes"] > 0 + assert logged["snapshot_rows"] == ( + logged["replay_rows"] + logged["staging_rows"] + ) assert actor._logger.log_metrics.call_args.kwargs == { "step": 1, "prefix": "timing/rollout_checkpoint", "step_metric": "telemetry/wall_time_seconds", } + def test_logs_checkpoint_outcome_reason_and_effective_cadence( + self, tmp_path: Path + ) -> None: + actor = self._actor(tmp_path) + actor._logger = MagicMock() + actor._telemetry_started_at = 0.0 + try: + with patch( + "nemo_rl.algorithms.single_controller.time.monotonic", + new=_SteppingClock(start=10.0), + ): + actor._log_rollout_checkpoint_outcome( + outcome="completed", + reason="completed", + attempt_duration_seconds=2.0, + ) + actor._log_rollout_checkpoint_outcome( + outcome="skipped", + reason="no_data_plane_mutations", + attempt_duration_seconds=0.1, + ) + finally: + actor._checkpointer.shutdown() + + first, second = actor._logger.log_metrics.call_args_list + assert first.args[0]["reason_completed"] == 1.0 + assert "seconds_since_last_success" not in first.args[0] + assert second.args[0]["reason_no_data_plane_mutations"] == 1.0 + assert second.args[0]["seconds_since_last_success"] == pytest.approx(2.0) + + def test_logs_restore_phase_total_and_reused_groups(self, tmp_path: Path) -> None: + actor = self._actor(tmp_path) + actor._logger = MagicMock() + actor._telemetry_started_at = time.monotonic() + actor._rollout_checkpoint_load_metrics = { + "snapshot_resolution_seconds": 0.5, + "dataloader_load_seconds": 1.0, + "tq_load_seconds": 2.0, + } + try: + actor._log_rollout_restore_metrics( + replay_metadata_load_seconds=3.0, + recovery_prepare_seconds=4.0, + restored_replay_groups=5, + ) + finally: + actor._checkpointer.shutdown() + + logged = actor._logger.log_metrics.call_args.args[0] + assert logged["total_load_seconds"] == 10.5 + assert logged["groups_reused"] == 5.0 + assert actor._logger.log_metrics.call_args.kwargs["prefix"] == ( + "timing/rollout_recovery" + ) + def test_logs_raw_and_canonical_rollout_throughput(self, tmp_path: Path) -> None: actor = self._actor(tmp_path) actor._logger = MagicMock() @@ -1429,7 +1494,7 @@ async def sample_twice() -> None: actor._checkpointer.shutdown() logged = actor._logger.log_metrics.call_args.args[0] - assert logged["vllm_output_tokens_per_second"] == pytest.approx(20.0) + assert logged["generation_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 @@ -1443,6 +1508,8 @@ def test_snapshot_reindexes_rows_owned_by_active_streamed_step( self, tmp_path: Path ): actor = self._actor(tmp_path) + actor._logger = MagicMock() + actor._telemetry_started_at = time.monotonic() claimed_meta = KVBatchMeta( partition_id=_PARTITION_ID, task_name=None, @@ -1462,7 +1529,8 @@ def test_snapshot_reindexes_rows_owned_by_active_streamed_step( actor._dp_client.sample_ids = list(claimed_meta.sample_ids) try: - assert asyncio.run(actor._save_rollout_checkpoint(force=True)) + result = asyncio.run(actor._save_rollout_checkpoint(force=True)) + assert result.saved finally: actor._checkpointer.shutdown() @@ -1484,16 +1552,35 @@ def test_snapshot_reindexes_rows_owned_by_active_streamed_step( assert [group["group_id"] for group in replay_state["groups"]] == [ "claimed-group" ] + logged = actor._logger.log_metrics.call_args.args[0] + assert logged["replay_rows"] == 1.0 + assert logged["snapshot_rows"] == 1.0 def test_snapshot_skips_optimizer_commit_window(self, tmp_path: Path): actor = self._actor(tmp_path) actor._optimizer_commit_in_progress = True try: - assert not asyncio.run(actor._save_rollout_checkpoint(force=True)) + result = asyncio.run(actor._save_rollout_checkpoint(force=True)) + assert not result.saved + assert result.reason == "optimizer_commit_in_progress" assert actor._dp_client.save_calls == [] finally: actor._checkpointer.shutdown() + def test_snapshot_reports_no_new_mutations(self, tmp_path: Path) -> None: + actor = self._actor(tmp_path) + actor._last_rollout_snapshot_mutation_version = ( + actor._data_plane_checkpoint_barrier.mutation_version + ) + try: + result = asyncio.run(actor._save_rollout_checkpoint()) + finally: + actor._checkpointer.shutdown() + + assert not result.saved + assert result.reason == "no_data_plane_mutations" + assert actor._dp_client.save_calls == [] + def test_periodic_pump_reports_each_consecutive_failure( self, tmp_path: Path, @@ -1565,6 +1652,8 @@ async def _failing_save(*, force: bool = False) -> bool: def test_periodic_pump_does_not_retry_invariant_failure(self, tmp_path: Path): actor = self._actor(tmp_path) + actor._logger = MagicMock() + actor._telemetry_started_at = time.monotonic() actor._master_config.rollout_checkpointing.snapshot_attempt_interval_s = 0.001 calls = 0 @@ -1588,6 +1677,9 @@ async def _failing_save(*, force: bool = False) -> bool: actor._checkpointer.shutdown() assert calls == 1 + logged = actor._logger.log_metrics.call_args.args[0] + assert logged["failed"] == 1.0 + assert logged["reason_invariant_error"] == 1.0 class TestDataPlaneCheckpoint: @@ -1823,15 +1915,15 @@ def test_rollout_recovery_inventory_merges_routes_and_clears_orphans(self): ) actor._dp_client = dp_client - async def validate_inventory() -> None: + async def validate_inventory() -> int: async with DataPlaneCheckpointBarrier().mutation() as cut: - await actor._validate_rollout_recovery_inventory( + return await actor._validate_rollout_recovery_inventory( cut, replay_metadata=replay_metadata, # type: ignore[arg-type] clear_unreferenced=True, ) - asyncio.run(validate_inventory()) + assert asyncio.run(validate_inventory()) == 2 assert dp_client.clear_calls == [(["orphan-key"], staging_partition)] assert sorted(dp_client.sample_ids) == [route_key, "sealed-key"] @@ -2430,6 +2522,12 @@ def test_periodic_snapshot_restores_exact_dispatch_cursor( assert actor_args.save_state.current_epoch == 4 assert actor_args.save_state.sampler_dispatch_index == 6 assert actor_args.last_checkpoint_path == str(final_snapshot) + assert actor_args.rollout_checkpoint_load_metrics is not None + assert { + "snapshot_resolution_seconds", + "dataloader_load_seconds", + "tq_load_seconds", + } <= actor_args.rollout_checkpoint_load_metrics.keys() def test_disabled_periodic_checkpointing_uses_trainer_anchor( self, From ab0f1936404e2ef44591d14ef5faede06627f5a3 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 7 Sep 2026 13:22:11 -0400 Subject: [PATCH 10/16] fix(rollout): complete telemetry restore coverage Signed-off-by: Anish Mahishi --- examples/configs/ppo_math_1B_megatron_single_controller.yaml | 1 + tests/unit/single_controller/test_checkpointing.py | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 3a13053c7f7..a0cc210bd7f 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -161,6 +161,7 @@ checkpointing: # checkpoint support. save_period=1 provides an anchor after every train step. rollout_checkpointing: snapshot_attempt_interval_s: null + telemetry_interval_s: null keep_latest_k: 2 restore_mode: latest # Advanced escape hatch for runtime-only fields from external integrations. diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 97294723297..ded36172b3f 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -2332,6 +2332,10 @@ def _write_periodic_snapshot(step_dir: Path) -> Path: {"fake_position": 7}, tmp_snapshot / "train_dataloader.pt", ) + # Production periodic snapshots include replay metadata alongside the + # dataloader and manifest. Keep this fixture representative so setup + # exercises rollout-payload restore timing as well as cursor restoration. + torch.save({"groups": []}, tmp_snapshot / REPLAY_BUFFER_METADATA_FILENAME) manifest = RolloutSnapshotManifest( schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, base_train_step=3, From 398574e56757f183d2e15e76de62d6900d21e40c Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 7 Sep 2026 14:53:59 -0400 Subject: [PATCH 11/16] fix(rollout): repair telemetry recovery validation Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 15 ++++++--------- nemo_rl/experience/rollout_manager.py | 4 +++- .../unit/single_controller/test_checkpointing.py | 11 ++++++++++- 3 files changed, 19 insertions(+), 11 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index d944cceb381..b634df3695b 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -800,16 +800,13 @@ def _log_rollout_restore_metrics( load_metrics = dict(self._rollout_checkpoint_load_metrics) load_metrics["replay_metadata_load_seconds"] = replay_metadata_load_seconds load_metrics["recovery_prepare_seconds"] = recovery_prepare_seconds + # At this point, load_metrics contains restore-phase timers only. Sum by + # the metric contract instead of maintaining a second hard-coded phase + # inventory that can silently omit a newly added restore timer. load_metrics["total_load_seconds"] = sum( - load_metrics[key] - for key in ( - "snapshot_resolution_seconds", - "dataloader_load_seconds", - "tq_load_seconds", - "replay_metadata_load_seconds", - "recovery_prepare_seconds", - ) - if key in load_metrics + value + for key, value in load_metrics.items() + if key.endswith("_seconds") and key != "total_load_seconds" ) load_metrics["groups_reused"] = float(restored_replay_groups) self._log_telemetry_metrics( diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 59fdf257139..bddc5f8be60 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -2000,7 +2000,9 @@ async def generate_and_push( 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(record.completions) + total_output_tokens = float(mean_output_tokens) * len( + record.completions + ) if math.isfinite(total_output_tokens): output_tokens = max(0, round(total_output_tokens)) self.record_canonical_publication(output_tokens) diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index ded36172b3f..c382d860946 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -53,6 +53,7 @@ from torchdata.stateful_dataloader import StatefulDataLoader from nemo_rl.algorithms.async_utils.replay_buffer import ( + DATA_PLANE_CHECKPOINT_DIR, LEGACY_REPLAY_BUFFER_FILENAME, REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, @@ -113,6 +114,7 @@ # Reuse the factory patches from the setup tests (same cross-module fixture # import pattern as test_rollout_pump.py). from tests.unit.single_controller.test_setup import ( + _native_tq_metadata, patched_factories, # noqa: F401 ) @@ -1425,6 +1427,7 @@ def test_logs_restore_phase_total_and_reused_groups(self, tmp_path: Path) -> Non "snapshot_resolution_seconds": 0.5, "dataloader_load_seconds": 1.0, "tq_load_seconds": 2.0, + "future_restore_phase_seconds": 5.0, } try: actor._log_rollout_restore_metrics( @@ -1436,7 +1439,7 @@ def test_logs_restore_phase_total_and_reused_groups(self, tmp_path: Path) -> Non actor._checkpointer.shutdown() logged = actor._logger.log_metrics.call_args.args[0] - assert logged["total_load_seconds"] == 10.5 + assert logged["total_load_seconds"] == 15.5 assert logged["groups_reused"] == 5.0 assert actor._logger.log_metrics.call_args.kwargs["prefix"] == ( "timing/rollout_recovery" @@ -2336,6 +2339,7 @@ def _write_periodic_snapshot(step_dir: Path) -> Path: # dataloader and manifest. Keep this fixture representative so setup # exercises rollout-payload restore timing as well as cursor restoration. torch.save({"groups": []}, tmp_snapshot / REPLAY_BUFFER_METADATA_FILENAME) + (tmp_snapshot / DATA_PLANE_CHECKPOINT_DIR).mkdir() manifest = RolloutSnapshotManifest( schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, base_train_step=3, @@ -2501,6 +2505,11 @@ def test_periodic_snapshot_restores_exact_dispatch_cursor( list(range(8)), None, ) + tq_metadata = _native_tq_metadata(step=3, trainer_version=3, epoch=4) + tq_metadata["replay_group_count"] = 0 + patched_factories[ + "fake_policy" + ].load_data_plane_checkpoint.return_value = tq_metadata with ( patch( From 44005707e634498b50c4eac0792ec2b4120aa029 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 8 Sep 2026 13:26:05 -0400 Subject: [PATCH 12/16] test(sc): cover long-stalled sibling recovery Signed-off-by: Anish Mahishi --- .../_checkpoint_scenarios.py | 49 ++++++++++++++++--- .../test_checkpoint_recovery_matrix.py | 31 ++++++++++++ 2 files changed, 74 insertions(+), 6 deletions(-) diff --git a/tests/unit/single_controller/_checkpoint_scenarios.py b/tests/unit/single_controller/_checkpoint_scenarios.py index 17b0a4816c9..a03d655cb8c 100644 --- a/tests/unit/single_controller/_checkpoint_scenarios.py +++ b/tests/unit/single_controller/_checkpoint_scenarios.py @@ -293,9 +293,10 @@ class RoundTrip: run has not lost the prompt, and either way these tests notice. ``ready`` and ``pending`` are reported separately for diagnosis only; nothing asserts on them. ``stamps`` records each restored group's - ``target_step`` and start weight. ``selected`` and ``selected_count`` report - the optional restore-then-select result used to verify each sampler's - recovery key at multiple gate lags. The sealed-sibling and redispatch maps + ``target_step`` and start weight. ``selected``, ``selected_count``, + ``evicted_after_restore``, and ``evicted_count_after_restore`` report the + optional restore-then-evict-and-select result used to verify each sampler's + recovery key and staleness behavior. The sealed-sibling and redispatch maps verify that an unfinished group keeps completed work and retries only its missing generation indices. The staging-row sets verify that the matching token-capture payload survived the data-plane checkpoint. @@ -310,6 +311,8 @@ class RoundTrip: stamps: dict[str, tuple[int | None, int]] selected: set[str] selected_count: int + evicted_after_restore: set[str] + evicted_count_after_restore: int sealed_before: dict[str, tuple[int, ...]] sealed_after: dict[str, tuple[int, ...]] redispatched: dict[str, tuple[int, ...]] @@ -323,6 +326,8 @@ async def _round_trip( tmp_path: Path, *, select_current_train_weight: int | None = None, + select_min_prompt_groups: int = GROUPS_PER_STEP, + select_max_prompt_groups: int = GROUPS_PER_STEP, ) -> RoundTrip: dp_a = _fresh_client(register=True) buf_a = _new_buffer(dp_a) @@ -464,11 +469,18 @@ async def _round_trip( } selected: set[str] = set() selected_count = 0 + evicted_after_restore: set[str] = set() + evicted_count_after_restore = 0 if select_current_train_weight is not None: + groups_before_evict = set(buf_b._group_ids) + evicted_count_after_restore = await sampler_b.evict( + current_train_weight=select_current_train_weight + ) + evicted_after_restore = groups_before_evict - set(buf_b._group_ids) selected_meta, selected_count = await sampler_b.select( current_train_weight=select_current_train_weight, - min_prompt_groups=GROUPS_PER_STEP, - max_prompt_groups=GROUPS_PER_STEP, + min_prompt_groups=select_min_prompt_groups, + max_prompt_groups=select_max_prompt_groups, ) if selected_meta is not None: selected = { @@ -484,6 +496,8 @@ async def _round_trip( stamps=stamps, selected=selected, selected_count=selected_count, + evicted_after_restore=evicted_after_restore, + evicted_count_after_restore=evicted_count_after_restore, sealed_before=sealed_before, sealed_after=sealed_after, redispatched=redispatched, @@ -498,6 +512,8 @@ def round_trip( tmp_path: Path, *, select_current_train_weight: int | None = None, + select_min_prompt_groups: int = GROUPS_PER_STEP, + select_max_prompt_groups: int = GROUPS_PER_STEP, ) -> RoundTrip: """Save the scenario, restore it into a fresh buffer, report what came back.""" return asyncio.run( @@ -506,6 +522,8 @@ def round_trip( sampler_name, tmp_path, select_current_train_weight=select_current_train_weight, + select_min_prompt_groups=select_min_prompt_groups, + select_max_prompt_groups=select_max_prompt_groups, ) ) @@ -664,6 +682,19 @@ def assert_completed_groups_survive( lag=1, ) +S_LONG_STALLED_PARTIAL = Scenario( + name="long-stalled-partly-generated", + groups=( + Group(12, 1, weight=1, target=1), + Group(13, ROLLOUTS_PER_GROUP, weight=8, target=8), + Group(14, ROLLOUTS_PER_GROUP, weight=8, target=8), + Group(15, ROLLOUTS_PER_GROUP, weight=8, target=8), + ), + cursor=16, + trained=frozenset(), + lag=1, +) + # Everything fully generated -- the case this PR set out to recover. FULLY_GENERATED = (S_ZERO_LAG_ALL_COMPLETE, S_ALL_COMPLETE, S_STALE_ONLY) # At least one group still generating when the snapshot was taken. @@ -673,6 +704,12 @@ def assert_completed_groups_survive( S_LAG2, S_EVICTED, S_TRAINED_OUT_OF_ORDER, + S_LONG_STALLED_PARTIAL, +) +WITH_SEALED_SIBLINGS = ( + S_ZERO_LAG_PARTIAL, + S_PARTIAL, + S_LAG2, + S_LONG_STALLED_PARTIAL, ) -WITH_SEALED_SIBLINGS = (S_ZERO_LAG_PARTIAL, S_PARTIAL, S_LAG2) ALL_SCENARIOS = FULLY_GENERATED + WITH_IN_FLIGHT diff --git a/tests/unit/single_controller/test_checkpoint_recovery_matrix.py b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py index fd04cc48933..d7ea5b11f10 100644 --- a/tests/unit/single_controller/test_checkpoint_recovery_matrix.py +++ b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py @@ -34,6 +34,7 @@ ROLLOUTS_PER_GROUP, S_ALL_COMPLETE, S_LAG2, + S_LONG_STALLED_PARTIAL, S_ZERO_LAG_ALL_COMPLETE, SAMPLERS, WITH_IN_FLIGHT, @@ -62,6 +63,7 @@ SEALED_SIBLING_CASES = [ Case(scenario, sampler) for sampler in SAMPLERS for scenario in WITH_SEALED_SIBLINGS ] +LONG_STALLED_CASES = [Case(S_LONG_STALLED_PARTIAL, sampler) for sampler in SAMPLERS] @pytest.fixture(autouse=True) @@ -115,6 +117,35 @@ def test_sealed_siblings_survive_and_only_missing_siblings_redispatch(case, tmp_ assert result.staging_rows_after_restore == result.staging_rows_before +@pytest.mark.parametrize("case", LONG_STALLED_CASES, ids=lambda case: case.id) +def test_long_stalled_partial_group_is_selected_or_deliberately_evicted(case, tmp_path): + """A restored straggler follows the configured sampler's stale-data policy.""" + result = round_trip( + case.scenario, + case.sampler, + tmp_path, + select_current_train_weight=8, + select_min_prompt_groups=1, + select_max_prompt_groups=len(case.scenario.groups), + ) + + stalled_group = "g12" + fresh_groups = {"g13", "g14", "g15"} + assert stalled_group in result.recovered + assert result.sealed_before[stalled_group] == (0,) + assert result.sealed_after[stalled_group] == (0,) + assert result.redispatched[stalled_group] == (1,) + + if case.sampler == "ready_first": + assert result.evicted_after_restore == set() + assert result.evicted_count_after_restore == 0 + assert result.selected == fresh_groups | {stalled_group} + else: + assert result.evicted_after_restore == {stalled_group} + assert result.evicted_count_after_restore == 1 + assert result.selected == fresh_groups + + @pytest.mark.parametrize("case", ALL_CASES, ids=lambda case: case.id) def test_restore_preserves_sampler_stamps(case, tmp_path): """Every restored group retains the keys its sampler uses for selection.""" From 195df040781dae2588c2e1fe9c9e2017fe316b0e Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 10 Sep 2026 11:48:23 -0400 Subject: [PATCH 13/16] fix(sc): address rollout telemetry review feedback Signed-off-by: Anish Mahishi --- docs/guides/single-controller.md | 13 + docs/observability/metrics.md | 29 +- ...po_math_1B_megatron_single_controller.yaml | 1 + ...po_math_1B_megatron_single_controller.yaml | 1 + .../algorithms/async_utils/replay_buffer.py | 20 +- nemo_rl/algorithms/single_controller.py | 227 ++++------ .../single_controller_utils/config.py | 6 + .../rollout_checkpoint.py | 21 +- nemo_rl/experience/rollout_manager.py | 8 +- nemo_rl/utils/logger.py | 11 +- tests/unit/experience/test_rollout_manager.py | 45 +- .../test_checkpoint_borrow_restore.py | 424 ++++++++++++++++++ .../test_checkpoint_dispatch_races.py | 11 +- .../single_controller/test_checkpointing.py | 197 ++++++-- .../single_controller/test_rollout_pump.py | 71 +-- .../test_single_controller_actor.py | 1 + .../test_tq_replay_buffer.py | 79 +++- tests/unit/utils/test_logger.py | 68 ++- 18 files changed, 958 insertions(+), 275 deletions(-) create mode 100644 tests/unit/single_controller/test_checkpoint_borrow_restore.py diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 10f3a8a8fad..b61b07e8713 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -119,6 +119,8 @@ checkpointing: rollout_checkpointing: snapshot_attempt_interval_s: 120 + telemetry_interval_s: null + max_consecutive_failures: 3 keep_latest_k: 2 restore_mode: latest extra_fingerprint_excluded_paths: [] @@ -136,6 +138,17 @@ are skipped until the matching trainer checkpoint exists. Before the first training step, snapshots are anchored to the initial model and a fingerprint of the rollout-semantic configuration. +`telemetry_interval_s` controls an independent wall-clock sampler for rollout +throughput and checkpoint pressure. It is `null` (disabled) by default; set it +to a positive number such as `30` to emit one sample every 30 seconds. This does +not change the checkpoint cadence. See the +[Single-Controller rollout recovery metrics](../observability/metrics.md#single-controller-rollout-recovery-metrics) +for the emitted fields. + +`max_consecutive_failures` is the number of consecutive retryable periodic-save +failures tolerated before training aborts. A successful or skipped checkpoint +attempt resets the count; checkpoint invariant failures still fail immediately. + The bootstrap fingerprint is fail-closed: every configuration value affects compatibility unless NeMo-RL's built-in denylist identifies it as operational, such as logging, cluster placement, checkpoint location, runtime ports, or diff --git a/docs/observability/metrics.md b/docs/observability/metrics.md index 3ad5bad1620..ff3f99b0dfa 100644 --- a/docs/observability/metrics.md +++ b/docs/observability/metrics.md @@ -18,11 +18,26 @@ Independent event streams use their own custom axes and do not carry `nemo_rl/step`. In particular, Single-Controller rollout benchmark series under `rollout/throughput/*`, `timing/rollout_checkpoint/*`, `timing/rollout_recovery/*`, and `rollout/checkpoint_outcome/*` use -`telemetry/wall_time_seconds`. This lets those series continue through a long -or paused trainer step without changing the meaning of the trainer-step axis. +`telemetry/wall_time_seconds`, recorded as Unix wall-clock time. This lets those +series continue through a long or paused trainer step and across process restart +without changing the meaning of the trainer-step axis. ## Single-Controller rollout recovery metrics +Wall-clock rollout telemetry sampling is disabled by default. Enable it without +changing the checkpoint cadence by setting, for example: + +```yaml +rollout_checkpointing: + telemetry_interval_s: 30.0 +``` + +This records one throughput and controller-pressure sample every 30 seconds. +Set the value back to `null` to disable the sampler. Detailed per-interval +metric dictionaries are printed to the controller log only when +`async_rl.diagnostics: true`; configured metric backends receive them either +way. + The rollout checkpoint benchmark focuses on the following questions: 1. Did checkpointing reduce raw generation or usable rollout throughput? @@ -33,13 +48,19 @@ The rollout checkpoint benchmark focuses on the following questions: | Prefix | Important fields | Meaning | |---|---|---| -| `rollout/throughput` | `generation_output_tokens_per_second`, `canonical_output_tokens_per_second`, `canonical_groups_per_second` | Raw backend decoding throughput compared with finalized token and group throughput available for training. The raw metric is absent when the backend does not expose compatible cumulative counters. | +| `rollout/throughput` | `generation_output_tokens_per_second`, `committed_output_tokens_per_second`, `committed_groups_per_second` | Raw backend decoding throughput compared with output committed for training. The raw metric is absent when the backend does not expose compatible cumulative counters. With token capture, committed tokens are counted exactly from valid staged rows; without token capture, they are estimated by multiplying the reported per-sample mean by the number of completions, so compare like-for-like runs. | | `rollout/throughput` | `checkpoint_blocked_mutations`, `checkpoint_mutation_wait_seconds_p95`, `checkpoint_mutation_wait_seconds_max` | Number and latency of live data-plane mutations delayed by an exclusive checkpoint. | | `timing/rollout_checkpoint` | `total_save_seconds`, `tq_save_seconds`, `barrier_wait_seconds`, `exclusive_hold_seconds`, `sidecar_save_seconds`, `snapshot_commit_seconds` | End-to-end save latency and its storage, fencing, controller-sidecar, and atomic-publication components. | | `timing/rollout_checkpoint` | `snapshot_rows`, `replay_rows`, `staging_rows`, `replay_groups`, `ledger_groups`, `controller_sidecar_bytes` | Logical volume captured by the snapshot. `controller_sidecar_bytes` excludes the native TQ payload because the current TQ checkpoint API does not report bytes written. NeMo-RL deliberately does not recursively scan the shared checkpoint directory because that scan would perturb the benchmark. | | `rollout/checkpoint_outcome` | `completed`, `skipped`, `failed`, `reason_*`, `seconds_since_previous_success`, `seconds_since_last_success` | Result, actionable reason, and effective cadence of every scheduled checkpoint attempt. | | `timing/rollout_recovery` | `snapshot_resolution_seconds`, `dataloader_load_seconds`, `tq_load_seconds`, `replay_metadata_load_seconds`, `recovery_prepare_seconds`, `total_load_seconds` | Rollout-state restore latency. `total_load_seconds` is the sum of these non-overlapping restore phases. | -| `timing/rollout_recovery` | `groups_reused`, `groups_redispatched`, `siblings_reused`, `siblings_redispatched`, `redispatch_schedule_seconds` | Completed groups restored without generation and unfinished sibling work preserved or repeated after restart. | +| `timing/rollout_recovery` | `groups_complete_restored`, `groups_unfinished_found`, `siblings_reused`, `siblings_rerun`, `redispatch_schedule_seconds` | Training-ready groups restored without generation, unfinished groups found for redispatch, and sibling work preserved or rerun after restart. | + +Recovery counters describe work reconstructed at restore time, not a promise +that every restored group will eventually train. After recovery, the configured +sampler may still evict a restored group under its normal staleness rules (for +example, when a `windowed` sampler finds that the group's policy version has +fallen outside its valid window). `barrier_wait_seconds` and mutation wait latency measure opposite sides of the same fence. The former is how long the checkpoint waits for already-running diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index d51e6bff8c4..3c839839ed9 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -157,6 +157,7 @@ checkpointing: rollout_checkpointing: snapshot_attempt_interval_s: null telemetry_interval_s: null + max_consecutive_failures: 3 keep_latest_k: 2 restore_mode: latest # Advanced escape hatch for runtime-only fields from external integrations. diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index a0cc210bd7f..e82b640e4cb 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -162,6 +162,7 @@ checkpointing: rollout_checkpointing: snapshot_attempt_interval_s: null telemetry_interval_s: null + max_consecutive_failures: 3 keep_latest_k: 2 restore_mode: latest # Advanced escape hatch for runtime-only fields from external integrations. diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 50ce965f93a..6e0d7b53d21 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -36,6 +36,8 @@ NotRequired, Optional, TypedDict, + cast, + get_args, ) import ray @@ -76,16 +78,9 @@ "sample_clears", "sibling_seals", ] -CHECKPOINT_MUTATION_KINDS: tuple[CheckpointMutationKind, ...] = ( - "advantage_writeback", - "group_commits", - "group_removals", - "prompt_reservations", - "recovery_restore", - "recovery_retries", - "sample_clears", - "sibling_seals", - "other", +CHECKPOINT_MUTATION_KINDS = cast( + tuple[CheckpointMutationKind, ...], + get_args(CheckpointMutationKind), ) @@ -298,6 +293,11 @@ async def mutation( self, kind: CheckpointMutationKind = "other" ) -> AsyncIterator[DataPlaneMutationCut]: """Yield one task-local live cut after any active checkpoint exits.""" + if kind not in CHECKPOINT_MUTATION_KINDS: + raise ValueError( + f"unknown checkpoint mutation kind {kind!r}; expected one of " + f"{CHECKPOINT_MUTATION_KINDS!r}" + ) async with self._condition: task = self._current_task() wait_started: Optional[float] = None diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index b634df3695b..746c4ccc2f0 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -63,7 +63,6 @@ Any, Awaitable, Callable, - Literal, Optional, Union, cast, @@ -107,9 +106,13 @@ is_ppo_run, ) from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + ROLLOUT_CHECKPOINT_ATTEMPT_OUTCOMES, + ROLLOUT_CHECKPOINT_ATTEMPT_REASONS, ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, ROLLOUT_SNAPSHOT_SCHEMA_VERSION, BootstrapCompatibilityIdentity, + RolloutCheckpointAttemptOutcome, + RolloutCheckpointAttemptReason, RolloutSnapshotManifest, commit_snapshot, ensure_bootstrap_anchor, @@ -157,7 +160,7 @@ from nemo_rl.models.policy.tq_policy import TQPolicy from nemo_rl.models.value.tq_value import TQValue from nemo_rl.utils.checkpoint import CheckpointManager, PathLike -from nemo_rl.utils.logger import Logger +from nemo_rl.utils.logger import Logger, TELEMETRY_WALL_TIME_METRIC from nemo_rl.utils.timer import TimeoutChecker, Timer if TYPE_CHECKING: @@ -166,31 +169,10 @@ Generation = Union[VllmGeneration, SGLangGeneration, MegatronGeneration] -_TELEMETRY_WALL_TIME_METRIC = "telemetry/wall_time_seconds" -_TELEMETRY_PREFIXES = ( - "timing/rollout_checkpoint", - "timing/rollout_recovery", - "rollout/checkpoint_outcome", - "rollout/throughput", -) - # Named `log` rather than `logger` to keep it distinct from the experiment # Logger this module also uses as `self._logger`. log = logging.getLogger(__name__) -_MAX_CONSECUTIVE_ROLLOUT_CHECKPOINT_FAILURES = 3 - -RolloutCheckpointAttemptReason = Literal[ - "completed", - "invariant_error", - "io_error", - "missing_trainer_anchor", - "no_data_plane_mutations", - "optimizer_commit_in_progress", - "timeout", - "trainer_state_changed", -] - @dataclass(frozen=True) class _RolloutCheckpointSaveResult: @@ -410,19 +392,6 @@ 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 @@ -449,8 +418,8 @@ def __init__( self._data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = ( actor_args.data_plane_checkpoint_metadata ) - self._rollout_checkpoint_load_metrics = getattr( - actor_args, "rollout_checkpoint_load_metrics", None + 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 @@ -633,6 +602,7 @@ async def run(self) -> dict[str, Any]: restored_replay_groups=restored_replay_groups ) await self._maybe_restore_replacement_reserve() + self._validate_restored_sampler_cursor() recovery_prepare_seconds = time.monotonic() - recovery_prepare_started self._log_rollout_restore_metrics( replay_metadata_load_seconds=replay_restore_seconds, @@ -645,11 +615,8 @@ async def run(self) -> dict[str, Any]: train_task = asyncio.create_task(self._train_pump()) watchdog_task = asyncio.create_task(self._stall_watchdog_pump()) tasks = [rollout_task, train_task, watchdog_task] - rollout_checkpoint_cfg = getattr( - self._master_config, "rollout_checkpointing", None - ) - telemetry_interval_s = getattr( - rollout_checkpoint_cfg, "telemetry_interval_s", None + telemetry_interval_s = ( + self._master_config.rollout_checkpointing.telemetry_interval_s ) if telemetry_interval_s is not None: await self._log_rollout_throughput_metrics(emit=False) @@ -761,18 +728,17 @@ def _log_telemetry_metrics( ) -> None: """Log benchmark telemetry on an axis independent of trainer steps.""" try: - self._telemetry_sample_index += 1 + wall_time_ns = time.time_ns() 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 - ) + event_metrics[TELEMETRY_WALL_TIME_METRIC] = wall_time_ns / 1_000_000_000 self._logger.log_metrics( event_metrics, - step=self._telemetry_sample_index, + # W&B uses the custom wall-time axis and ignores this value. + # Other backends still need a restart-stable integer event step. + step=wall_time_ns, prefix=prefix, - step_metric=_TELEMETRY_WALL_TIME_METRIC, + step_metric=TELEMETRY_WALL_TIME_METRIC, ) except Exception as error: warnings.warn( @@ -808,7 +774,7 @@ def _log_rollout_restore_metrics( for key, value in load_metrics.items() if key.endswith("_seconds") and key != "total_load_seconds" ) - load_metrics["groups_reused"] = float(restored_replay_groups) + load_metrics["groups_complete_restored"] = float(restored_replay_groups) self._log_telemetry_metrics( load_metrics, step=self._train_steps, @@ -1018,6 +984,43 @@ async def _maybe_restore_rollout_recovery( flush=True, ) + def _validate_restored_sampler_cursor(self) -> None: + """Require the sampler cursor to cover every restored target step.""" + restored_target_steps = [ + target_step + for target_step in self._buffer.target_step_list + if target_step is not None + ] + restored_target_steps.extend( + group.target_step + for group in self._rollout_manager.recovery_ledger.groups() + if group.target_step is not None + ) + if not restored_target_steps: + return + + max_target_step = max(restored_target_steps) + if self._sampler.dispatch_index < max_target_step: + raise RuntimeError( + "restored sampler cursor is older than restored rollout work: " + f"dispatch_index={self._sampler.dispatch_index}, " + f"max_target_step={max_target_step}. Refusing to continue because " + "a new admission could collide with restored work and discard prompts." + ) + + def _require_unoccupied_target_step(self, target_step: Optional[int]) -> None: + """Reject a sampler admission that collides with existing buffer ownership.""" + if target_step is None: + return + buffered = self._buffer.count_for_target_step(target_step) + if buffered: + raise RuntimeError( + f"sampler admitted target_step={target_step}, but the replay buffer " + f"already contains {buffered} group(s) for that step. The sampler " + "cursor and rollout ownership are inconsistent; refusing to discard " + "reserved prompts." + ) + async def _rehydrate_rollout_recovery_prompts( self, cut: DataPlaneMutationCut, @@ -1110,13 +1113,11 @@ async def _rehydrate_rollout_recovery_prompts( async def _admit_reserved_prompt_groups( self, group_ids: list[str], - ) -> tuple[Optional[int], list[str], int]: - """Commit one admission and atomically reconcile restored canonical groups. + ) -> Optional[int]: + """Commit one admission and mark every reserved group as admitted. Returns: - The target-step stamp, IDs that still require rollout dispatch, and the - number of already-canonical groups that replaced reservations in this - admission. + The target-step stamp assigned to every group. """ if not group_ids: raise ValueError("sampler admission requires at least one prompt group") @@ -1124,9 +1125,10 @@ async def _admit_reserved_prompt_groups( def _commit( cut: DataPlaneMutationCut, target_step: Optional[int], - ) -> tuple[Optional[int], list[str], int]: + ) -> Optional[int]: if target_step is not None: self._sampler_stamps_target_steps = True + self._require_unoccupied_target_step(target_step) for group_id in group_ids: self._rollout_manager.mark_prompt_group_admitted( cut, @@ -1134,16 +1136,7 @@ def _commit( target_step=target_step, ) - buffered = 0 - dispatch_group_ids = group_ids - if target_step is not None: - buffered = self._buffer.count_for_target_step(target_step) - if buffered: - dispatch_count = max(0, len(group_ids) - buffered) - dispatch_group_ids = group_ids[:dispatch_count] - for group_id in group_ids[dispatch_count:]: - self._rollout_manager.discard_prompt_group(cut, group_id) - return target_step, dispatch_group_ids, buffered + return target_step if isinstance(self._sampler, TransactionalAdmissionSampler): await self._sampler.wait_until_admissible( @@ -1207,7 +1200,6 @@ async def _redispatch_restored_rollouts( # opening the sampler gate. Launch them before waiting to re-admit RESERVED # groups, or restore can deadlock with the trainer waiting for recovered work # that this method has not launched yet. - redispatched = 0 for group in groups_to_recover: if group.phase is PromptGroupPhase.ADMITTED: await launch( @@ -1218,7 +1210,6 @@ async def _redispatch_restored_rollouts( reused = len(group.sealed_generation_indices) reused_siblings += reused redispatched_siblings += group.expected_generations - reused - redispatched += 1 # A checkpoint may land after dataloader ownership is recorded but before # sampler admission commits. Re-admit each original dataloader batch once and @@ -1230,10 +1221,8 @@ async def _redispatch_restored_rollouts( group.group_id ) for group_ids in reserved_admissions.values(): - _, dispatch_group_ids, _ = await self._admit_reserved_prompt_groups( - group_ids - ) - for group_id in dispatch_group_ids: + await self._admit_reserved_prompt_groups(group_ids) + for group_id in group_ids: group = recovery_ledger.get_group(group_id) await launch( group.prompt_payload, @@ -1243,7 +1232,6 @@ async def _redispatch_restored_rollouts( reused = len(group.sealed_generation_indices) reused_siblings += reused redispatched_siblings += group.expected_generations - reused - redispatched += 1 self._rollout_manager.record_recovery_siblings( reused=reused_siblings, @@ -1252,17 +1240,16 @@ async def _redispatch_restored_rollouts( self._log_telemetry_metrics( { "redispatch_schedule_seconds": time.monotonic() - recovery_started, - "groups_considered": float(len(groups_to_recover)), - "groups_redispatched": float(redispatched), + "groups_unfinished_found": float(len(groups_to_recover)), "siblings_reused": float(reused_siblings), - "siblings_redispatched": float(redispatched_siblings), + "siblings_rerun": float(redispatched_siblings), }, step=self._train_steps, prefix="timing/rollout_recovery", ) print( - f"📦 Redispatched {redispatched} unfinished rollout " + f"📦 Redispatched {len(groups_to_recover)} unfinished rollout " "group(s) before new dataloader work", flush=True, ) @@ -2120,18 +2107,8 @@ async def _launch( ) if target_step is not None: self._sampler_stamps_target_steps = True - num_prompts = prompt_batch.size - if target_step is not None: - buffered = self._buffer.count_for_target_step(target_step) - if buffered: - num_prompts = max(0, prompt_batch.size - buffered) - print( - f" target_step={target_step}: {buffered} group(s) " - f"already buffered; dispatching {num_prompts} of " - f"{prompt_batch.size} prompt(s), dropping the rest", - flush=True, - ) - for prompt_idx in range(num_prompts): + self._require_unoccupied_target_step(target_step) + for prompt_idx in range(prompt_batch.size): prompt: DatumSpec = { # type: ignore k: v[prompt_idx] for k, v in prompt_batch.items() } @@ -2166,30 +2143,10 @@ async def _launch( ) prompt_dispatches.append((prompt, group_id)) - ( - target_step, - dispatch_group_ids, - buffered, - ) = await self._admit_reserved_prompt_groups( + target_step = await self._admit_reserved_prompt_groups( [group_id for _, group_id in prompt_dispatches] ) - if target_step is not None: - if buffered: - print( - f" target_step={target_step}: {buffered} group(s) " - f"already buffered; dispatching " - f"{len(dispatch_group_ids)} of " - f"{len(prompt_dispatches)} prompt(s), dropping the rest", - flush=True, - ) - dispatch_group_id_set = set(dispatch_group_ids) - prompt_dispatches = [ - (prompt, group_id) - for prompt, group_id in prompt_dispatches - if group_id in dispatch_group_id_set - ] - for prompt, group_id in prompt_dispatches: await _launch(prompt, target_step, group_id) @@ -2296,23 +2253,12 @@ async def _drain_reserve_into_steps( admission_id=admission_id, ) prompt_dispatches.append((prompt, group_id)) - ( - target_step, - dispatch_group_ids, - buffered, - ) = await self._admit_reserved_prompt_groups( + target_step = await self._admit_reserved_prompt_groups( [group_id for _, group_id in prompt_dispatches] ) - dispatch_group_id_set = set(dispatch_group_ids) - prompt_dispatches = [ - (prompt, group_id) - for prompt, group_id in prompt_dispatches - if group_id in dispatch_group_id_set - ] print( f" dataloader exhausted; training on {len(prompt_dispatches)} " - f"pooled spare(s) as target_step={target_step}" - + (f" ({buffered} group(s) already buffered)" if buffered else ""), + f"pooled spare(s) as target_step={target_step}", flush=True, ) for prompt, group_id in prompt_dispatches: @@ -3991,7 +3937,7 @@ async def _save_rollout_checkpoint( def _log_rollout_checkpoint_outcome( self, *, - outcome: Literal["completed", "failed", "skipped"], + outcome: RolloutCheckpointAttemptOutcome, reason: RolloutCheckpointAttemptReason, attempt_duration_seconds: float, ) -> None: @@ -3999,14 +3945,18 @@ def _log_rollout_checkpoint_outcome( 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.snapshot_attempt_interval_s or 0.0 ), - f"reason_{reason}": 1.0, + **{ + f"reason_{candidate}": float(reason == candidate) + for candidate in ROLLOUT_CHECKPOINT_ATTEMPT_REASONS + }, + **{ + candidate: float(outcome == candidate) + for candidate in ROLLOUT_CHECKPOINT_ATTEMPT_OUTCOMES + }, } completed_at = time.monotonic() if outcome == "completed": @@ -4065,7 +4015,10 @@ async def _rollout_checkpoint_pump(self) -> None: f"{type(error).__name__}: {error}", flush=True, ) - if consecutive_failures >= _MAX_CONSECUTIVE_ROLLOUT_CHECKPOINT_FAILURES: + if ( + consecutive_failures + >= self._master_config.rollout_checkpointing.max_consecutive_failures + ): raise RuntimeError( "periodic rollout checkpoint failed " f"{consecutive_failures} consecutive times" @@ -4216,13 +4169,12 @@ async def _collect_and_log_rollout_throughput_metrics( 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"] + metrics["committed_groups_per_second"] = ( + counters["committed_groups"] - previous_counters["committed_groups"] ) / elapsed - metrics["canonical_output_tokens_per_second"] = ( - counters["canonical_output_tokens"] - - previous_counters["canonical_output_tokens"] + metrics["committed_output_tokens_per_second"] = ( + counters["committed_output_tokens"] + - previous_counters["committed_output_tokens"] ) / elapsed if generated_by_worker is not None and previous_generated is not None: current_workers = set(generated_by_worker) @@ -4252,7 +4204,8 @@ async def _collect_and_log_rollout_throughput_metrics( step=self._train_steps, prefix="rollout/throughput", ) - print(f"rollout_throughput_metrics={metrics}", flush=True) + if self._async_cfg.diagnostics: + print(f"rollout_throughput_metrics={metrics}", flush=True) async def _save_checkpoint( self, diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 8e959db8d1e..91a4f482d4d 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -745,10 +745,16 @@ class RolloutCheckpointConfig(BaseModel, extra="forbid"): ``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 ``snapshot_attempt_interval_s``. + + ``max_consecutive_failures`` controls how many consecutive retryable + periodic-checkpoint failures are tolerated before the controller aborts the + run. A successful or skipped attempt resets the counter; checkpoint + invariant failures still fail immediately. """ snapshot_attempt_interval_s: Annotated[Optional[float], Field(gt=0)] = None telemetry_interval_s: Annotated[Optional[float], Field(gt=0)] = None + max_consecutive_failures: Annotated[int, Field(ge=1)] = 3 keep_latest_k: Annotated[int, Field(ge=1)] = 2 restore_mode: Literal["latest", "trainer_checkpoint"] = "latest" extra_fingerprint_excluded_paths: list[str] = Field(default_factory=list) diff --git a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py index 539aa7d86c4..c79163d9e6d 100644 --- a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py +++ b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py @@ -24,7 +24,7 @@ from dataclasses import asdict, dataclass from fnmatch import fnmatchcase from pathlib import Path -from typing import Any, Mapping, Optional +from typing import Any, Literal, Mapping, Optional, get_args from nemo_rl.algorithms.single_controller_utils.config import MasterConfig @@ -35,6 +35,25 @@ ROLLOUT_SNAPSHOTS_DIRNAME = "rollout_snapshots" ROLLOUT_SNAPSHOT_MANIFEST_FILENAME = "manifest.json" +RolloutCheckpointAttemptOutcome = Literal["completed", "failed", "skipped"] +ROLLOUT_CHECKPOINT_ATTEMPT_OUTCOMES: tuple[RolloutCheckpointAttemptOutcome, ...] = ( + get_args(RolloutCheckpointAttemptOutcome) +) + +RolloutCheckpointAttemptReason = Literal[ + "completed", + "invariant_error", + "io_error", + "missing_trainer_anchor", + "no_data_plane_mutations", + "optimizer_commit_in_progress", + "timeout", + "trainer_state_changed", +] +ROLLOUT_CHECKPOINT_ATTEMPT_REASONS: tuple[RolloutCheckpointAttemptReason, ...] = ( + get_args(RolloutCheckpointAttemptReason) +) + _SNAPSHOT_RE = re.compile(r"snapshot_(\d+)") _TMP_SNAPSHOT_RE = re.compile(r"tmp_snapshot_(\d+)") _TRASH_SNAPSHOT_RE = re.compile(r"trash_snapshot_(\d+)") diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index bddc5f8be60..5c448caee8d 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -1678,12 +1678,12 @@ async def _recovery_mutation( yield cut def telemetry_snapshot(self) -> dict[str, int]: - """Return cumulative canonical-publication and recovery counters.""" + """Return cumulative committed-publication and recovery counters.""" return { - "canonical_groups_finalized": self._canonical_groups_finalized, - "canonical_output_tokens": self._canonical_output_tokens, + "committed_groups": self._canonical_groups_finalized, + "committed_output_tokens": self._canonical_output_tokens, "recovery_siblings_reused": self._recovery_siblings_reused, - "recovery_siblings_redispatched": self._recovery_siblings_redispatched, + "recovery_siblings_rerun": self._recovery_siblings_redispatched, } def record_canonical_publication(self, output_tokens: int) -> None: diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index 77f2e5623bf..9b00271e56a 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -55,7 +55,8 @@ # 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" +WANDB_CALLER_STEP_METRIC = "nemo_rl/step" +TELEMETRY_WALL_TIME_METRIC = "telemetry/wall_time_seconds" class WandbConfig(TypedDict): @@ -237,11 +238,11 @@ def __init__(self, cfg: WandbConfig, log_dir: Optional[str] = None): 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 + 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(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. @@ -497,7 +498,7 @@ def _flush_pending_metrics_locked(self) -> None: if self._pending_step is None: return event_metrics = dict(self._pending_metrics) - event_metrics[_WANDB_CALLER_STEP_METRIC] = self._pending_step + event_metrics[WANDB_CALLER_STEP_METRIC] = self._pending_step self.run.log(event_metrics) self._pending_step = None self._pending_metrics = {} @@ -518,7 +519,7 @@ def _buffer_step_metrics_locked( for metric_name in metrics: self._define_exact_metric_locked( metric_name, - step_metric=_WANDB_CALLER_STEP_METRIC, + step_metric=WANDB_CALLER_STEP_METRIC, ) self._pending_metrics.update(metrics) if step_finished: diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 2f7496e1974..c0a4bb5bd42 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -461,7 +461,7 @@ async def _logged_commit(*args, **kwargs): assert start_v == 0 assert end_v == 0 assert len(mgr.recovery_ledger) == 0 - assert mgr.telemetry_snapshot()["canonical_groups_finalized"] == 1 + assert mgr.telemetry_snapshot()["committed_groups"] == 1 def test_publication_and_recovery_telemetry_are_cumulative(self): mgr = _make_manager(_FakeBuffer(), _FakeImpl()) @@ -470,12 +470,49 @@ def test_publication_and_recovery_telemetry_are_cumulative(self): mgr.record_recovery_siblings(reused=3, redispatched=1) assert mgr.telemetry_snapshot() == { - "canonical_groups_finalized": 1, - "canonical_output_tokens": 42, + "committed_groups": 1, + "committed_output_tokens": 42, "recovery_siblings_reused": 3, - "recovery_siblings_redispatched": 1, + "recovery_siblings_rerun": 1, } + @pytest.mark.parametrize( + ("mean_output_tokens", "expected_output_tokens"), + [(3.75, 8), (-2.0, 0)], + ids=["rounded-group-total", "negative-total-clamped"], + ) + def test_non_capture_commit_estimates_committed_output_tokens( + self, + mean_output_tokens: float, + expected_output_tokens: int, + ) -> None: + """The legacy path rounds its per-sample mean and clamps bad totals.""" + completions = [ + Completion( + message_log=[], + env_extras=None, + truncated=False, + reward=0.0, + ) + for _ in range(2) + ] + record = PromptGroupRecord( + prompt_idx=0, + prompt=[], + extra_env_info=None, + metadata={}, + completions=completions, + rollout_metrics={"mean_gen_tokens_per_sample": mean_output_tokens}, + ) + mgr = _make_manager(_FakeBuffer(), _FakeImpl(record=record)) + + _run(mgr.generate_and_push({"prompt": "p"})) + + assert ( + mgr.telemetry_snapshot()["committed_output_tokens"] + == expected_output_tokens + ) + def test_ledger_hands_ownership_to_canonical_buffer_on_commit(self): buf = _FakeBuffer() diff --git a/tests/unit/single_controller/test_checkpoint_borrow_restore.py b/tests/unit/single_controller/test_checkpoint_borrow_restore.py new file mode 100644 index 00000000000..a2f554a3eeb --- /dev/null +++ b/tests/unit/single_controller/test_checkpoint_borrow_restore.py @@ -0,0 +1,424 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Round-trip a borrowed and repaid rollout through an SC checkpoint cut.""" + +from __future__ import annotations + +import asyncio +from collections import deque +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import pytest +import torch + +from nemo_rl.algorithms.async_utils.replay_buffer import ( + REPLAY_BUFFER_METADATA_FILENAME, + DataPlaneCheckpointBarrier, + TQReplayBuffer, +) +from nemo_rl.algorithms.async_utils.staleness_sampler import InOrderSampler +from nemo_rl.algorithms.grpo import GRPOConfig +from nemo_rl.algorithms.single_controller import ( + DATA_PLANE_CHECKPOINT_DIR, + SingleControllerActor, +) +from nemo_rl.algorithms.single_controller_utils.config import RolloutRecoveryConfig +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.experience.failures import RolloutDataFailure +from nemo_rl.experience.interfaces import PromptGroupRecord +from nemo_rl.experience.rollout_manager import ( + RolloutManager, + RolloutRetryPolicy, + RolloutStats, +) +from nemo_rl.experience.rollout_recovery import ( + ROLLOUT_RECOVERY_STATE_FILENAME, + PromptGroupPhase, + RolloutRecoveryLedger, +) +from tests.unit.single_controller import _checkpoint_scenarios as scenarios + +_GROUP_SIZE = 2 +_PROMPTS_PER_STEP = 3 +_CAPACITY = 64 +_TIMEOUT_S = 20.0 + + +class _Generation: + """Finish immediately unless a prompt is deliberately held or failed.""" + + def __init__(self) -> None: + self.hold: dict[int, asyncio.Event] = {} + self.fail: set[int] = set() + + async def run_rollout(self, input_sample: dict[str, Any]) -> PromptGroupRecord: + prompt_idx = input_sample["idx"] + gate = self.hold.get(prompt_idx) + if gate is not None: + await gate.wait() + if prompt_idx in self.fail: + raise RolloutDataFailure(f"prompt {prompt_idx} is bad on purpose") + return PromptGroupRecord( + prompt_idx=prompt_idx, + prompt=[], + extra_env_info=None, + metadata={}, + completions=[], + rollout_metrics={}, + ) + + +class _Loader: + def __init__( + self, + batches: list[BatchedDataDict], + dataset: dict[int, dict[str, Any]], + ) -> None: + self._batches = batches + self.dataset = dataset + + def __iter__(self): + return iter(self._batches) + + def __len__(self) -> int: + return len(self._batches) + + def state_dict(self) -> dict[str, bool]: + return {"fake": True} + + +class _TrackingRolloutManager(RolloutManager): + """Associate stable recovery group IDs with their source prompt indexes.""" + + def reserve_prompt_group(self, cut, input_sample, **kwargs) -> str: + group_id = super().reserve_prompt_group(cut, input_sample, **kwargs) + self.prompt_by_group_id[group_id] = int(input_sample["idx"]) + return group_id + + +def _prompt(prompt_idx: int) -> dict[str, Any]: + return { + "idx": prompt_idx, + "message_log": [{"role": "user", "content": f"p{prompt_idx}"}], + } + + +def _batch(prompt_indices: list[int]) -> BatchedDataDict: + return BatchedDataDict( + { + "idx": prompt_indices, + "message_log": [ + [{"role": "user", "content": f"p{prompt_idx}"}] + for prompt_idx in prompt_indices + ], + } + ) + + +def _client(*, register: bool) -> NoOpDataPlaneClient: + client = NoOpDataPlaneClient() + if register: + client.register_partition( + partition_id=scenarios.PARTITION, + fields=list(scenarios._FIELDS), + num_samples=_CAPACITY * _GROUP_SIZE, + consumer_tasks=["train"], + ) + return client + + +def _manager( + buffer: TQReplayBuffer, + barrier: DataPlaneCheckpointBarrier, + generation: _Generation, +) -> _TrackingRolloutManager: + manager = object.__new__(_TrackingRolloutManager) + manager._impl = generation + manager._tokenizer = None + manager._num_generations_per_prompt = _GROUP_SIZE + manager._rollout_recovery_config = RolloutRecoveryConfig() + manager._tq_buffer = buffer + manager._recovery_ledger = RolloutRecoveryLedger() + manager._data_plane_checkpoint_barrier = barrier + manager._env_handles = {} + manager._weight_version = 0 + manager._retry_policy = RolloutRetryPolicy( + max_infra_attempts=1, + max_data_attempts=1, + max_gym_row_attempts=1, + max_skipped_prompts=8, + ) + manager._stats = RolloutStats() + manager._canonical_groups_finalized = 0 + manager._canonical_output_tokens = 0 + manager._recovery_siblings_reused = 0 + manager._recovery_siblings_redispatched = 0 + manager._skipped_prompts = 0 + manager._consecutive_infra_drops = 0 + manager.prompt_by_group_id: dict[str, int] = {} + return manager + + +def _controller( + *, + client: NoOpDataPlaneClient, + loader: _Loader, + generation: _Generation, + dispatch_index: int | None = None, + replacement_reserve: list[dict[str, Any]] | None = None, + checkpoint_path: Path | None = None, + checkpoint_metadata: dict[str, Any] | None = None, +) -> Any: + barrier = DataPlaneCheckpointBarrier() + buffer = TQReplayBuffer( + client, + partition_id=scenarios.PARTITION, + pad_value_dict={"input_ids": 0}, + include_message_violation_fields=False, + require_routed_experts=False, + ) + buffer.set_data_plane_checkpoint_barrier(barrier) + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._data_plane_checkpoint_barrier = barrier + controller._buffer = buffer + controller._dp_client = client + controller._partition_id = scenarios.PARTITION + controller._rollout_manager = _manager(buffer, barrier, generation) + controller._sampler = InOrderSampler(buffer, max_lookahead_versions=2) + if dispatch_index is not None: + controller._sampler.restore_dispatch_index(dispatch_index) + elif checkpoint_path is not None: + controller._sampler.set_dispatch_index(0) + controller._async_cfg = SimpleNamespace( + max_inflight_prompts=16, + max_buffered_rollouts=_CAPACITY, + diagnostics=False, + sampler=SimpleNamespace(name="in_order"), + rollout_failure=SimpleNamespace( + on_dropped_prompt="replace", + max_replacement_attempts=1, + replacement_reserve_prompts=_PROMPTS_PER_STEP, + min_step_batch_fraction=0.9, + ), + ) + controller._algo_cfg = GRPOConfig.model_construct( + max_num_epochs=1, + num_prompts_per_step=_PROMPTS_PER_STEP, + num_generations_per_prompt=_GROUP_SIZE, + ) + controller._master_config = SimpleNamespace( + grpo=controller._algo_cfg, + token_capture=SimpleNamespace(enabled=False), + ) + controller._dataloader = loader + controller._rollout_permitted = asyncio.Event() + controller._rollout_permitted.set() + controller._rollout_exhausted = asyncio.Event() + controller._buffer_capacity = asyncio.Semaphore(_CAPACITY) + controller._inflight_rollouts = 0 + controller._inflight_by_group_id = {} + controller._dispatched_rollouts = set() + controller._trainer_version = 0 + controller._train_steps = 0 + controller._current_epoch = 0 + controller._sampler_stamps_target_steps = False + controller._rollout_recovery_enabled = True + controller._batch_shortfall = {} + controller._batch_replacements = {} + controller._batch_promotions = {} + controller._finalizer_actors = [] + controller._replacement_reserve = deque(replacement_reserve or []) + controller._rollout_slot_waiters = 0 + controller._rollout_permitted_waiters = 0 + controller._buffer_capacity_waiters = 0 + controller._rollout_completion_durations_s = deque(maxlen=10_000) + controller._rollout_queue_wait_durations_s = deque(maxlen=10_000) + controller._logger = MagicMock() + controller._last_checkpoint_path = ( + str(checkpoint_path) if checkpoint_path is not None else None + ) + controller._data_plane_checkpoint_metadata = checkpoint_metadata + return controller + + +async def _wait_for(predicate, description: str, pump: asyncio.Task[None]) -> None: + deadline = asyncio.get_running_loop().time() + _TIMEOUT_S + while not predicate(): + if pump.done(): + pump.result() + raise AssertionError(f"rollout pump exited before {description}") + if asyncio.get_running_loop().time() > deadline: + raise TimeoutError(f"timed out waiting for {description}") + await asyncio.sleep(0.005) + + +def _stamps( + buffer: TQReplayBuffer, + prompt_by_group_id: dict[str, int], +) -> dict[int, list[int]]: + result: dict[int, list[int]] = {} + for group_id, target_step in zip(buffer._group_ids, buffer.target_step_list): + assert target_step is not None + result.setdefault(target_step, []).append(prompt_by_group_id[group_id]) + return { + target_step: sorted(prompt_indices) + for target_step, prompt_indices in sorted(result.items()) + } + + +async def _exercise_round_trip( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + scenarios._rb, "record_to_train_batch", scenarios._stub_converter + ) + + dataset = {prompt_idx: _prompt(prompt_idx) for prompt_idx in range(15)} + batches = [ + _batch([0, 1, 2]), + _batch([3, 4, 5]), + _batch([6, 7, 8]), + _batch([9, 10, 11]), + _batch([12, 13, 14]), + ] + generation = _Generation() + generation.hold[1] = asyncio.Event() + generation.fail.add(1) + generation.hold[3] = asyncio.Event() + + first_client = _client(register=True) + first = _controller( + client=first_client, + loader=_Loader(batches, dataset), + generation=generation, + ) + first_pump = asyncio.create_task(first._rollout_pump()) + first_ledger = first._rollout_manager.recovery_ledger + await _wait_for( + lambda: sum(first._buffer.ready_list) == 8 + and sum( + group.phase is PromptGroupPhase.RESERVED for group in first_ledger.groups() + ) + == 3, + "lookahead work and the reserved batch", + first_pump, + ) + + generation.hold[1].set() + await _wait_for( + lambda: first._batch_promotions == {0: 1} + and any( + first._rollout_manager.prompt_by_group_id.get(group_id) == 3 + for group_id in first._buffer._group_ids + ), + "the borrow and its repayment dispatch", + first_pump, + ) + + checkpoint_path = tmp_path / "step_0" + checkpoint_path.mkdir() + async with first._data_plane_checkpoint_barrier.checkpoint() as cut: + snapshot = await first._capture_rollout_checkpoint_cut(cut, checkpoint_path) + assert snapshot.replay_metadata is not None + assert snapshot.rollout_recovery_payload is not None + torch.save( + snapshot.replay_metadata, + checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME, + ) + (checkpoint_path / ROLLOUT_RECOVERY_STATE_FILENAME).write_bytes( + snapshot.rollout_recovery_payload + ) + first_pump.cancel() + await asyncio.gather(first_pump, return_exceptions=True) + + second_client = _client(register=False) + checkpoint_metadata = second_client.load_checkpoint( + checkpoint_path / DATA_PLANE_CHECKPOINT_DIR + ) + restored = _controller( + client=second_client, + loader=_Loader([], dataset), + generation=_Generation(), + dispatch_index=snapshot.sampler_dispatch_index, + replacement_reserve=snapshot.replacement_reserve, + checkpoint_path=checkpoint_path, + checkpoint_metadata=checkpoint_metadata, + ) + restored_group_count = await restored._maybe_restore_replay_buffer() + await restored._maybe_restore_rollout_recovery( + restored_replay_groups=restored_group_count + ) + restored._validate_restored_sampler_cursor() + assert restored._sampler.dispatch_index == snapshot.sampler_dispatch_index + restored._rollout_manager.prompt_by_group_id.update( + first._rollout_manager.prompt_by_group_id + ) + + restored_pump = asyncio.create_task(restored._rollout_pump()) + await _wait_for( + lambda: 3 in restored._rollout_manager.prompt_by_group_id.values() + and any( + restored._rollout_manager.prompt_by_group_id.get(group_id) == 3 + for group_id in restored._buffer._group_ids + ), + "the restored repayment rollout", + restored_pump, + ) + restored._trainer_version = 1 + await asyncio.wait_for(restored_pump, timeout=_TIMEOUT_S) + + present = { + prompt_idx + for prompt_indices in _stamps( + restored._buffer, + restored._rollout_manager.prompt_by_group_id, + ).values() + for prompt_idx in prompt_indices + } + spares = {prompt["idx"] for prompt in restored._replacement_reserve} + assert present | spares | {1} == set(range(15)) + + recovery_metrics: dict[str, float] = {} + for call in restored._logger.log_metrics.call_args_list: + recovery_metrics.update( + { + key: value + for key, value in call.args[0].items() + if key + in { + "groups_unfinished_found", + "siblings_reused", + "siblings_rerun", + } + } + ) + assert recovery_metrics["groups_unfinished_found"] == 4.0 + assert recovery_metrics["siblings_reused"] == 0.0 + assert recovery_metrics["siblings_rerun"] == 8.0 + + +def test_borrow_and_repayment_survive_controller_checkpoint_restore( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A saved sampler cursor preserves every prompt after borrow and repayment.""" + asyncio.run(_exercise_round_trip(tmp_path, monkeypatch)) diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 1ba41a4680b..d05743586fa 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -90,8 +90,6 @@ async def apply() -> _T: def _init_recovery_telemetry(controller: Any, *, train_steps: int = 0) -> None: """Initialize constructor-owned telemetry state for hand-built controllers.""" controller._train_steps = train_steps - controller._telemetry_sample_index = 0 - controller._telemetry_started_at = 0.0 controller._logger = MagicMock() @@ -973,6 +971,11 @@ async def _recover( assert sampler.dispatch_index == 7 assert controller._batch_shortfall == {6: 1} assert controller._sampler_stamps_target_steps is True + recovery_metrics = controller._logger.log_metrics.call_args.args[0] + assert recovery_metrics["groups_unfinished_found"] == 1.0 + assert recovery_metrics["siblings_reused"] == 0.0 + assert recovery_metrics["siblings_rerun"] == 2.0 + assert "groups_redispatched" not in recovery_metrics asyncio.run(exercise()) @@ -1344,7 +1347,7 @@ async def exercise() -> None: async def block_admission( group_ids: list[str], - ) -> tuple[int, list[str], int]: + ) -> int: admission_started.set() await release_admission.wait() async with controller._data_plane_checkpoint_barrier.mutation() as cut: @@ -1354,7 +1357,7 @@ async def block_admission( group_id, target_step=7, ) - return 7, group_ids, 0 + return 7 async def launch( prompt: DatumSpec, diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index c382d860946..3a2d89b4e1a 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -39,7 +39,6 @@ import json import os import threading -import time from collections.abc import Callable from pathlib import Path from types import SimpleNamespace @@ -110,6 +109,7 @@ encode_route_plan, ) from nemo_rl.utils.checkpoint import CheckpointManager +from nemo_rl.utils.logger import TELEMETRY_WALL_TIME_METRIC # Reuse the factory patches from the setup tests (same cross-module fixture # import pattern as test_rollout_pump.py). @@ -496,10 +496,10 @@ def __init__(self, events: Optional[list[str]] = None) -> None: self.recovery_ledger = RolloutRecoveryLedger() self._events = events self.telemetry = { - "canonical_groups_finalized": 0, - "canonical_output_tokens": 0, + "committed_groups": 0, + "committed_output_tokens": 0, "recovery_siblings_reused": 0, - "recovery_siblings_redispatched": 0, + "recovery_siblings_rerun": 0, } def set_data_plane_checkpoint_barrier(self, barrier: Any) -> None: @@ -520,12 +520,12 @@ 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 + self.telemetry["committed_groups"] += 1 + self.telemetry["committed_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 + self.telemetry["recovery_siblings_rerun"] += redispatched class _FakeTQBuffer: @@ -943,6 +943,45 @@ def _training_info(ckpt_dir: Path, step: int) -> dict[str, Any]: class TestCounterRestore: + @pytest.mark.parametrize( + ("buffer_target_steps", "recovery_target_steps"), + [([8], []), ([], [8])], + ) + def test_rejects_sampler_cursor_older_than_restored_work( + self, + buffer_target_steps: list[int], + recovery_target_steps: list[int], + ) -> None: + actor = object.__new__(_ACTOR_CLS) + actor._buffer = SimpleNamespace(target_step_list=buffer_target_steps) + actor._rollout_manager = SimpleNamespace( + recovery_ledger=SimpleNamespace( + groups=lambda: [ + SimpleNamespace(target_step=target_step) + for target_step in recovery_target_steps + ] + ) + ) + actor._sampler = SimpleNamespace(dispatch_index=7) + + with pytest.raises( + RuntimeError, + match=r"dispatch_index=7, max_target_step=8", + ): + actor._validate_restored_sampler_cursor() + + def test_accepts_sampler_cursor_covering_restored_work(self) -> None: + actor = object.__new__(_ACTOR_CLS) + actor._buffer = SimpleNamespace(target_step_list=[None, 7]) + actor._rollout_manager = SimpleNamespace( + recovery_ledger=SimpleNamespace( + groups=lambda: [SimpleNamespace(target_step=8)] + ) + ) + actor._sampler = SimpleNamespace(dispatch_index=8) + + actor._validate_restored_sampler_cursor() + def test_restore_from_step_n(self, tmp_path): save_state = _initial_grpo_save_state() save_state.current_step = 7 @@ -1308,6 +1347,7 @@ def test_restore_mode_rejects_removed_none_value(self): {"snapshot_attempt_interval_s": 0}, {"interval_s": 1}, {"keep_latest_k": 0}, + {"max_consecutive_failures": 0}, {"unknown_option": True}, ], ) @@ -1361,11 +1401,17 @@ def test_pre_step_snapshot_contains_only_rollout_state(self, tmp_path: Path): def test_logs_snapshot_phase_durations(self, tmp_path: Path) -> None: actor = self._actor(tmp_path) actor._logger = MagicMock() - actor._telemetry_started_at = 0.0 + wall_time_ns = 1_750_000_000_000_000_000 try: - with patch( - "nemo_rl.algorithms.single_controller.time.monotonic", - new=_SteppingClock(), + with ( + patch( + "nemo_rl.algorithms.single_controller.time.monotonic", + new=_SteppingClock(), + ), + patch( + "nemo_rl.algorithms.single_controller.time.time_ns", + return_value=wall_time_ns, + ), ): result = asyncio.run(actor._save_rollout_checkpoint(force=True)) assert result.saved @@ -1383,10 +1429,12 @@ def test_logs_snapshot_phase_durations(self, tmp_path: Path) -> None: assert logged["snapshot_rows"] == ( logged["replay_rows"] + logged["staging_rows"] ) + assert logged[TELEMETRY_WALL_TIME_METRIC] == 1_750_000_000.0 + assert "sample_index" not in logged assert actor._logger.log_metrics.call_args.kwargs == { - "step": 1, + "step": wall_time_ns, "prefix": "timing/rollout_checkpoint", - "step_metric": "telemetry/wall_time_seconds", + "step_metric": TELEMETRY_WALL_TIME_METRIC, } def test_logs_checkpoint_outcome_reason_and_effective_cadence( @@ -1394,7 +1442,6 @@ def test_logs_checkpoint_outcome_reason_and_effective_cadence( ) -> None: actor = self._actor(tmp_path) actor._logger = MagicMock() - actor._telemetry_started_at = 0.0 try: with patch( "nemo_rl.algorithms.single_controller.time.monotonic", @@ -1410,19 +1457,53 @@ def test_logs_checkpoint_outcome_reason_and_effective_cadence( reason="no_data_plane_mutations", attempt_duration_seconds=0.1, ) + actor._log_rollout_checkpoint_outcome( + outcome="completed", + reason="completed", + attempt_duration_seconds=1.0, + ) finally: actor._checkpointer.shutdown() - first, second = actor._logger.log_metrics.call_args_list + first, second, third = actor._logger.log_metrics.call_args_list + outcome_keys = {"completed", "failed", "skipped"} + reason_keys = { + "reason_completed", + "reason_invariant_error", + "reason_io_error", + "reason_missing_trainer_anchor", + "reason_no_data_plane_mutations", + "reason_optimizer_commit_in_progress", + "reason_timeout", + "reason_trainer_state_changed", + } + assert { + key for key in first.args[0] if key.startswith("reason_") + } == reason_keys + assert { + key for key in second.args[0] if key.startswith("reason_") + } == reason_keys + assert { + key for key in third.args[0] if key.startswith("reason_") + } == reason_keys + assert sum(first.args[0][key] for key in reason_keys) == 1.0 + assert sum(second.args[0][key] for key in reason_keys) == 1.0 + assert sum(third.args[0][key] for key in reason_keys) == 1.0 + assert sum(first.args[0][key] for key in outcome_keys) == 1.0 + assert sum(second.args[0][key] for key in outcome_keys) == 1.0 + assert sum(third.args[0][key] for key in outcome_keys) == 1.0 assert first.args[0]["reason_completed"] == 1.0 + assert first.args[0]["reason_no_data_plane_mutations"] == 0.0 assert "seconds_since_last_success" not in first.args[0] + assert second.args[0]["reason_completed"] == 0.0 assert second.args[0]["reason_no_data_plane_mutations"] == 1.0 - assert second.args[0]["seconds_since_last_success"] == pytest.approx(2.0) + assert second.args[0]["seconds_since_last_success"] == pytest.approx(1.0) + assert third.args[0]["seconds_since_previous_success"] == pytest.approx(2.0) + assert "seconds_since_last_success" not in third.args[0] def test_logs_restore_phase_total_and_reused_groups(self, tmp_path: Path) -> None: actor = self._actor(tmp_path) actor._logger = MagicMock() - actor._telemetry_started_at = time.monotonic() actor._rollout_checkpoint_load_metrics = { "snapshot_resolution_seconds": 0.5, "dataloader_load_seconds": 1.0, @@ -1440,28 +1521,32 @@ def test_logs_restore_phase_total_and_reused_groups(self, tmp_path: Path) -> Non logged = actor._logger.log_metrics.call_args.args[0] assert logged["total_load_seconds"] == 15.5 - assert logged["groups_reused"] == 5.0 + assert logged["groups_complete_restored"] == 5.0 assert actor._logger.log_metrics.call_args.kwargs["prefix"] == ( "timing/rollout_recovery" ) - def test_logs_raw_and_canonical_rollout_throughput(self, tmp_path: Path) -> None: + def test_logs_raw_and_canonical_rollout_throughput( + self, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + ) -> None: actor = self._actor(tmp_path) actor._logger = MagicMock() - actor._telemetry_started_at = 0.0 + capsys.readouterr() snapshots = iter( [ { - "canonical_groups_finalized": 0, - "canonical_output_tokens": 0, + "committed_groups": 0, + "committed_output_tokens": 0, "recovery_siblings_reused": 0, - "recovery_siblings_redispatched": 0, + "recovery_siblings_rerun": 0, }, { - "canonical_groups_finalized": 2, - "canonical_output_tokens": 40, + "committed_groups": 2, + "committed_output_tokens": 40, "recovery_siblings_reused": 1, - "recovery_siblings_redispatched": 3, + "recovery_siblings_rerun": 3, }, ] ) @@ -1498,21 +1583,69 @@ async def sample_twice() -> None: logged = actor._logger.log_metrics.call_args.args[0] assert logged["generation_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["committed_output_tokens_per_second"] == pytest.approx(4.0) + assert logged["committed_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) + assert "rollout_throughput_metrics=" not in capsys.readouterr().out + + @pytest.mark.parametrize( + "second_generation_tokens", + [ + {0: [90], 1: [150]}, + {0: [150], 2: [50]}, + ], + ids=["counter-decreased", "worker-set-changed"], + ) + def test_generation_counter_discontinuity_suppresses_invalid_rate( + self, + tmp_path: Path, + second_generation_tokens: dict[int, list[int]], + ) -> None: + actor = self._actor(tmp_path) + actor._logger = MagicMock() + actor._rollout_manager.telemetry_snapshot = lambda: { + "committed_groups": 0, + "committed_output_tokens": 0, + "recovery_siblings_reused": 0, + "recovery_siblings_rerun": 0, + } + generation_snapshots = iter( + [ + {"generation_tokens": {0: [100], 1: [100]}}, + {"generation_tokens": second_generation_tokens}, + ] + ) + actor._gen = SimpleNamespace( + drain_latest_logger_metrics=lambda: next(generation_snapshots) + ) + + async def sample_twice() -> None: + await actor._log_rollout_throughput_metrics(emit=False) + await actor._log_rollout_throughput_metrics() + + try: + with patch( + "nemo_rl.algorithms.single_controller.time.monotonic", + new=_SteppingClock(start=10.0, step=10.0), + ): + asyncio.run(sample_twice()) + finally: + actor._checkpointer.shutdown() + + logged = actor._logger.log_metrics.call_args.args[0] + assert logged["generation_counter_discontinuity"] == 1.0 + assert "generation_output_tokens_per_second" not in logged def test_snapshot_reindexes_rows_owned_by_active_streamed_step( self, tmp_path: Path ): actor = self._actor(tmp_path) actor._logger = MagicMock() - actor._telemetry_started_at = time.monotonic() claimed_meta = KVBatchMeta( partition_id=_PARTITION_ID, task_name=None, @@ -1628,6 +1761,7 @@ def test_periodic_pump_aborts_after_repeated_failures( ): actor = self._actor(tmp_path) actor._master_config.rollout_checkpointing.snapshot_attempt_interval_s = 0.001 + actor._master_config.rollout_checkpointing.max_consecutive_failures = 2 actor._train_steps = 1 async def _main() -> None: @@ -1638,7 +1772,7 @@ async def _failing_save(*, force: bool = False) -> bool: actor._save_rollout_checkpoint = _failing_save with pytest.raises( RuntimeError, - match="periodic rollout checkpoint failed 3 consecutive times", + match="periodic rollout checkpoint failed 2 consecutive times", ): await asyncio.wait_for( actor._rollout_checkpoint_pump(), @@ -1651,12 +1785,11 @@ async def _failing_save(*, force: bool = False) -> bool: actor._checkpointer.shutdown() output = capsys.readouterr().out - assert output.count("Periodic rollout checkpoint failed") == 3 + assert output.count("Periodic rollout checkpoint failed") == 2 def test_periodic_pump_does_not_retry_invariant_failure(self, tmp_path: Path): actor = self._actor(tmp_path) actor._logger = MagicMock() - actor._telemetry_started_at = time.monotonic() actor._master_config.rollout_checkpointing.snapshot_attempt_interval_s = 0.001 calls = 0 diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 994117806aa..f4136d7662f 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -22,6 +22,7 @@ from contextlib import asynccontextmanager from types import SimpleNamespace from typing import Any +from unittest.mock import MagicMock import pytest import ray @@ -288,71 +289,27 @@ async def generate_and_push( assert ctrl._inflight_rollouts == 0 -@pytest.mark.parametrize( - ("restored", "expected_new_dispatches"), - [ - # Room left for a partial top-up. - (1, 1), - # Target step already full: the whole batch is dropped. - (2, 0), - # More restored than a batch: still zero, never negative. - (3, 0), - ], -) -def test_rollout_pump_tops_up_restored_target_step( - restored: int, - expected_new_dispatches: int, -) -> None: - # On resume the buffer holds groups still stamped for the next target - # step. In-order selection consumes a target step as one fixed-size batch, - # so the pump must dispatch only the shortfall — a full batch on top would - # leave surplus groups that are never selected and whose capacity permits - # are held until evict. - buffer = _RecordingBuffer([0] * restored) +def test_reserved_admission_rejects_an_occupied_target_step() -> None: + """A stale sampler cursor must not silently discard newly yielded prompts.""" + buffer = _RecordingBuffer([0]) controller_cls = SingleControllerActor.__ray_metadata__.modified_class ctrl = object.__new__(controller_cls) ctrl._buffer = buffer - ctrl._async_cfg = SimpleNamespace( - max_inflight_prompts=2, diagnostics=False, rollout_failure=_failure_cfg() + ctrl._rollout_manager = SimpleNamespace( + mark_prompt_group_admitted=MagicMock(), + discard_prompt_group=MagicMock(), ) - ctrl._master_config = SimpleNamespace( - grpo=GRPOConfig.model_construct(max_num_epochs=1) - ) - ctrl._algo_cfg = ctrl._master_config.grpo - ctrl._rollout_manager = _RecordingRolloutManager(buffer) - # lookahead=0 keeps the single batch on target_step 0. ctrl._sampler = InOrderSampler(buffer, max_lookahead_versions=0) - ctrl._dataloader = [ - BatchedDataDict( - { - "message_log": [ - [{"role": "user", "content": "p0"}], - [{"role": "user", "content": "p1"}], - ] - } - ) - ] - ctrl._rollout_permitted = asyncio.Event() - ctrl._rollout_permitted.set() - ctrl._rollout_exhausted = asyncio.Event() - ctrl._buffer_capacity = asyncio.Semaphore(4) - ctrl._inflight_rollouts = 0 - ctrl._inflight_by_group_id = {} - ctrl._dispatched_rollouts = set() ctrl._trainer_version = 0 - ctrl._current_epoch = 0 - _init_pump_ledgers(ctrl) + ctrl._sampler_stamps_target_steps = False + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() - asyncio.run(ctrl._rollout_pump()) + with pytest.raises(RuntimeError, match="already contains 1 group"): + asyncio.run(ctrl._admit_reserved_prompt_groups(["new-group"])) - # Only the shortfall was dispatched on top of the restored groups. - assert buffer.target_step_list == [0] * (restored + expected_new_dispatches) - # A dispatched prompt keeps its permit (the train pump releases it after - # consuming the group), so exactly one permit per dispatch is held and the - # dropped prompts consume none. - assert ctrl._buffer_capacity._value == 4 - expected_new_dispatches - assert ctrl._inflight_rollouts == 0 - assert ctrl._rollout_exhausted.is_set() + assert buffer.target_step_list == [0] + ctrl._rollout_manager.mark_prompt_group_admitted.assert_not_called() + ctrl._rollout_manager.discard_prompt_group.assert_not_called() class _SkippingRolloutManager: diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 470ae99a3fb..b4cdf923071 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -137,6 +137,7 @@ def _actor_args_for_init(**overrides) -> SimpleNamespace: finalizer_actors=[], data_plane_checkpoint_metadata=None, bootstrap_identity=None, + rollout_checkpoint_load_metrics=None, ) args.update(overrides) return SimpleNamespace(**args) diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 6e5538b8d06..5f1e7d04f85 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -19,7 +19,7 @@ import asyncio import io import threading -from typing import Any +from typing import Any, cast import pytest import torch @@ -35,6 +35,7 @@ TQReplayBuffer, replay_manifest_digest, ) +from nemo_rl.algorithms.async_utils.staleness_sampler import InOrderSampler from nemo_rl.data_plane import KVBatchMeta from nemo_rl.data_plane.schema import ROLLOUT_METRICS, ROUTE_PLAN_TAG from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -484,6 +485,19 @@ async def mutate(kind: CheckpointMutationKind) -> None: asyncio.run(exercise()) + def test_rejects_unknown_mutation_kind(self) -> None: + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + unknown = cast(CheckpointMutationKind, "group_commit") + + with pytest.raises( + ValueError, match="unknown checkpoint mutation kind 'group_commit'" + ): + async with barrier.mutation(unknown): + pytest.fail("unknown mutation kind unexpectedly admitted") + + asyncio.run(exercise()) + class TestTQReplayBufferReserveCommit: def test_reserve_rejects_duplicate_live_group_id(self): @@ -1171,6 +1185,69 @@ def test_ignores_rollout_metrics_logging_sidecar(self): class TestTQReplayBufferStateDict: + def test_borrow_and_repayment_remain_selectable_after_restore(self) -> None: + async def exercise() -> None: + dp = FakeDataPlaneClient() + original = _make_buffer(dp) + lender_id = original.reserve( + weight_version=0, + target_step=1, + group_id="lender", + ) + await original.commit( + lender_id, + _make_record(), + start_weight_version=0, + end_weight_version=0, + ) + + assert original.promote_ready_group(to_target_step=0) == 1 + + repayment_id = original.reserve( + weight_version=0, + target_step=1, + group_id="repayment", + ) + await original.commit( + repayment_id, + _make_record(), + start_weight_version=0, + end_weight_version=0, + ) + state = original.metadata_state_dict(saved_capacity=8) + + restored = _make_buffer(dp) + await restored.load_state_dict( + state, + max_groups=8, + expected_partition_id="rollout_data", + expected_group_size=_N_GENS, + expected_manifest_digest=state["manifest_digest"], + ) + sampler = InOrderSampler(restored, max_lookahead_versions=1) + sampler.restore_dispatch_index(1) + + borrowed, borrowed_count = await sampler.select( + current_train_weight=0, + min_prompt_groups=1, + max_prompt_groups=1, + ) + repaid, repaid_count = await sampler.select( + current_train_weight=1, + min_prompt_groups=1, + max_prompt_groups=1, + ) + + assert borrowed is not None + assert borrowed.sample_ids == ["lender_g0", "lender_g1"] + assert borrowed_count == 1 + assert repaid is not None + assert repaid.sample_ids == ["repayment_g0", "repayment_g1"] + assert repaid_count == 1 + assert restored.group_ids == () + + asyncio.run(exercise()) + def test_training_claim_is_reindexed_only_for_periodic_snapshot(self): dp = FakeDataPlaneClient() buf = _make_buffer(dp) diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py index 852b7e26bf5..0b9ffb74657 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -30,7 +30,9 @@ MLflowLogger, RayGpuMonitorLogger, SwanlabLogger, + TELEMETRY_WALL_TIME_METRIC, TensorboardLogger, + WANDB_CALLER_STEP_METRIC, WandbLogger, flatten_dict, log_container_init_timing, @@ -426,7 +428,9 @@ def test_log_metrics(self, mock_wandb): # 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, "nemo_rl/step": step}) + mock_run.log.assert_called_once_with( + {**metrics, WANDB_CALLER_STEP_METRIC: step} + ) @patch("nemo_rl.utils.logger.wandb") def test_log_metrics_with_prefix(self, mock_wandb): @@ -444,7 +448,7 @@ def test_log_metrics_with_prefix(self, mock_wandb): expected_metrics = { "train/loss": 0.5, "train/accuracy": 0.8, - "nemo_rl/step": step, + WANDB_CALLER_STEP_METRIC: step, } mock_run.log.assert_called_once_with(expected_metrics) @@ -506,10 +510,10 @@ def test_independent_events_do_not_reuse_wandb_internal_step(self, mock_wandb): 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}, + {TELEMETRY_WALL_TIME_METRIC: 30.0, "tokens_per_second": 10.0}, step=1, prefix="rollout/throughput", - step_metric="telemetry/wall_time_seconds", + step_metric=TELEMETRY_WALL_TIME_METRIC, ) logger.log_metrics( {"tokens_per_second": 20.0}, @@ -523,7 +527,7 @@ def test_independent_events_do_not_reuse_wandb_internal_step(self, mock_wandb): assert mock_run.log.call_args_list == [ call( { - "telemetry/wall_time_seconds": 30.0, + TELEMETRY_WALL_TIME_METRIC: 30.0, "rollout/throughput/tokens_per_second": 10.0, } ), @@ -532,10 +536,10 @@ def test_independent_events_do_not_reuse_wandb_internal_step(self, mock_wandb): "train/loss": 1.0, "timing/train/seconds": 5.0, "performance/tokens_per_second": 20.0, - "nemo_rl/step": 1, + WANDB_CALLER_STEP_METRIC: 1, } ), - call({"train/loss": 0.5, "nemo_rl/step": 2}), + call({"train/loss": 0.5, WANDB_CALLER_STEP_METRIC: 2}), ] assert all("step" not in kwargs for _, kwargs in mock_run.log.call_args_list) @@ -597,7 +601,7 @@ def test_log_metrics_requires_registered_step_metric(self, mock_wandb): logger = WandbLogger({}) logger.define_metric( "rollout/throughput/*", - step_metric="telemetry/wall_time_seconds", + step_metric=TELEMETRY_WALL_TIME_METRIC, ) with pytest.raises(ValueError, match="is missing from the logged event"): @@ -612,12 +616,12 @@ def test_define_metric_uses_longest_matching_prefix(self, mock_wandb): logger.define_metric("rollout/*", step_metric="rollout/step") logger.define_metric( "rollout/throughput/*", - step_metric="telemetry/wall_time_seconds", + step_metric=TELEMETRY_WALL_TIME_METRIC, ) logger.log_metrics( { - "telemetry/wall_time_seconds": 30.0, + TELEMETRY_WALL_TIME_METRIC: 30.0, "rollout/throughput/tokens_per_second": 10.0, }, step=0, @@ -626,7 +630,7 @@ def test_define_metric_uses_longest_matching_prefix(self, mock_wandb): mock_run = mock_wandb.init.return_value mock_run.define_metric.assert_any_call( "rollout/throughput/tokens_per_second", - step_metric="telemetry/wall_time_seconds", + step_metric=TELEMETRY_WALL_TIME_METRIC, ) @patch("nemo_rl.utils.logger.wandb") @@ -635,7 +639,7 @@ def test_does_not_define_catch_all_metric(self, mock_wandb): WandbLogger({}) mock_run = mock_wandb.init.return_value - assert call("*", step_metric="nemo_rl/step") not in ( + assert call("*", step_metric=WANDB_CALLER_STEP_METRIC) not in ( mock_run.define_metric.call_args_list ) @@ -651,7 +655,9 @@ def test_registers_teardown_flush(self, mock_wandb, mock_atexit_register): 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.log.assert_called_once_with( + {"train/loss": 0.5, WANDB_CALLER_STEP_METRIC: 7} + ) mock_run.finish.assert_called_once_with() @patch("nemo_rl.utils.logger.wandb") @@ -664,7 +670,37 @@ def test_finish_flushes_pending_trainer_row(self, mock_wandb): 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.log.assert_called_once_with( + {"train/loss": 0.5, WANDB_CALLER_STEP_METRIC: 7} + ) + mock_run.finish.assert_called_once_with() + + @patch("nemo_rl.utils.logger.wandb") + def test_finish_flushes_histogram_and_plot_with_pending_trainer_row( + self, mock_wandb + ) -> None: + """Typed W&B values join the final scalar row instead of being lost.""" + histogram_value = MagicMock(name="histogram_value") + plot_value = MagicMock(name="plot_value") + mock_wandb.Histogram.return_value = histogram_value + logger = WandbLogger({}) + + logger.log_metrics({"loss": 0.5}, step=7, prefix="train") + logger.log_histogram([1.0, 2.0], step=7, name="train/reward_histogram") + logger.log_plot(plot_value, step=7, name="train/reward_plot") + + 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, + "train/reward_histogram": histogram_value, + "train/reward_plot": plot_value, + WANDB_CALLER_STEP_METRIC: 7, + } + ) mock_run.finish.assert_called_once_with() @patch("nemo_rl.utils.logger.wandb") @@ -2019,12 +2055,12 @@ def test_define_metric_only_targets_wandb( logger.define_metric( "rollout/throughput/*", - step_metric="telemetry/wall_time_seconds", + step_metric=TELEMETRY_WALL_TIME_METRIC, ) mock_wandb_logger.return_value.define_metric.assert_called_once_with( "rollout/throughput/*", - step_metric="telemetry/wall_time_seconds", + step_metric=TELEMETRY_WALL_TIME_METRIC, ) assert not mock_tb_logger.return_value.define_metric.called From a3dfa6ca8c96f34d674ee2e854bc1e09c066ccc0 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 10 Sep 2026 10:02:01 -0700 Subject: [PATCH 14/16] fix: lint Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 2 +- tests/unit/utils/test_logger.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 746c4ccc2f0..998cf33105e 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -160,7 +160,7 @@ from nemo_rl.models.policy.tq_policy import TQPolicy from nemo_rl.models.value.tq_value import TQValue from nemo_rl.utils.checkpoint import CheckpointManager, PathLike -from nemo_rl.utils.logger import Logger, TELEMETRY_WALL_TIME_METRIC +from nemo_rl.utils.logger import TELEMETRY_WALL_TIME_METRIC, Logger from nemo_rl.utils.timer import TimeoutChecker, Timer if TYPE_CHECKING: diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py index 0b9ffb74657..50329e02de3 100644 --- a/tests/unit/utils/test_logger.py +++ b/tests/unit/utils/test_logger.py @@ -26,13 +26,13 @@ import torch from nemo_rl.utils.logger import ( + TELEMETRY_WALL_TIME_METRIC, + WANDB_CALLER_STEP_METRIC, Logger, MLflowLogger, RayGpuMonitorLogger, SwanlabLogger, - TELEMETRY_WALL_TIME_METRIC, TensorboardLogger, - WANDB_CALLER_STEP_METRIC, WandbLogger, flatten_dict, log_container_init_timing, From e48366dc446d5265016bc29ccf1b2bf65f2d3b19 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 10 Sep 2026 19:35:27 -0400 Subject: [PATCH 15/16] fix(sc): validate restored target ownership Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 1 + .../single_controller/test_checkpointing.py | 10 +++++++++ .../single_controller/test_rollout_pump.py | 22 +++++++++++++++++++ 3 files changed, 33 insertions(+) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 998cf33105e..ef1cac6bc14 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -2274,6 +2274,7 @@ async def _drain_reserve_into_steps( target_step = await self._sampler.admit( trainer_version_fn=lambda: self._trainer_version ) + self._require_unoccupied_target_step(target_step) print( f" dataloader exhausted; training on {len(step_prompts)} pooled " f"spare(s) as target_step={target_step}", diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 3a2d89b4e1a..daaa098237a 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -546,6 +546,9 @@ def __init__( "manifest_digest": "fake-manifest-digest", "groups": [], } + self.target_step_list = [ + group["target_step"] for group in self._metadata_state["groups"] + ] self.load_return = load_return self.metadata_state_dict_calls: list[int] = [] self.load_calls: list[dict[str, Any]] = [] @@ -616,6 +619,10 @@ async def load_state_dict( "expected_manifest_digest": expected_manifest_digest, } ) + self._metadata_state = state + self.target_step_list = [ + group["target_step"] for group in self._metadata_state["groups"] + ] return self.load_return @@ -2933,6 +2940,8 @@ def test_run_restores_native_tq_replay_metadata_without_payload_reput( data_plane_checkpoint=True, ) buffer = _FakeTQBuffer(load_return=2) + save_state = _matching_save_state() + save_state.sampler_dispatch_index = 1 actor, result = _run_actor_run( mc, @@ -2941,6 +2950,7 @@ def test_run_restores_native_tq_replay_metadata_without_payload_reput( dp_client=_FakeDPClient(sample_ids=sample_ids), last_checkpoint_path=str(ckpt_dir), data_plane_checkpoint_metadata=tq_metadata, + save_state=save_state, ), ) diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index f4136d7662f..4e8bd1bd4f0 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -312,6 +312,28 @@ def test_reserved_admission_rejects_an_occupied_target_step() -> None: ctrl._rollout_manager.discard_prompt_group.assert_not_called() +def test_non_recovery_reserve_drain_rejects_an_occupied_target_step() -> None: + """The ordinary spare-pool admission enforces the same cursor invariant.""" + buffer = _RecordingBuffer([0]) + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._algo_cfg = SimpleNamespace(num_prompts_per_step=1) + ctrl._buffer = buffer + ctrl._replacement_reserve = deque([{"idx": 0}]) + ctrl._rollout_recovery_enabled = False + ctrl._sampler = InOrderSampler(buffer, max_lookahead_versions=0) + ctrl._trainer_version = 0 + launched: list[Any] = [] + + async def launch(*args: Any) -> None: + launched.append(args) + + with pytest.raises(RuntimeError, match="already contains 1 group"): + asyncio.run(ctrl._drain_reserve_into_steps(launch)) + + assert launched == [] + + class _SkippingRolloutManager: """Every prompt is given up on within budget, so nothing is ever committed.""" From 054bb6f0b292a06f30bc99cc744c5f0c2a3e73fd Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 11 Sep 2026 13:42:56 -0700 Subject: [PATCH 16/16] fix(test): fix regressed test Signed-off-by: Anish Mahishi --- .../single_controller/test_checkpointing.py | 28 +++++++++++++++++-- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index daaa098237a..9151293f1ba 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -300,9 +300,10 @@ async def select(self, **kwargs) -> tuple[Optional[KVBatchMeta], int]: class _RestoredGroupsSampler(_FakeSampler): """Drain the exact groups represented by a restored replay metadata file.""" - def __init__(self, groups: list[dict[str, Any]]) -> None: + def __init__(self, groups: list[dict[str, Any]], buffer: "_FakeTQBuffer") -> None: super().__init__() self._groups = list(groups) + self._buffer = buffer async def select( self, @@ -316,6 +317,10 @@ async def select( if len(selected) < min_prompt_groups: return None, 0 del self._groups[: len(selected)] + # Legacy local-removal contract: a sampler without training claims drops + # the rows from the replay index at selection, so a checkpoint taken + # after the step cannot list groups whose canonical rows are gone. + self._buffer.drop_groups([group["group_id"] for group in selected]) metas = [group["meta"] for group in selected] return ( @@ -582,6 +587,21 @@ def metadata_state_dict( ] return state + def drop_groups(self, group_ids: list[str]) -> None: + """Remove groups from the replay index, as a selection or eviction does.""" + dropped = set(group_ids) + remaining = [ + group + for group in self._metadata_state["groups"] + if group["group_id"] not in dropped + ] + unknown = dropped - { + group["group_id"] for group in self._metadata_state["groups"] + } + assert not unknown, f"unknown group_ids={sorted(unknown)!r}" + self._metadata_state = {**self._metadata_state, "groups": remaining} + self.target_step_list = [group["target_step"] for group in remaining] + def training_owned_replay_groups(self) -> list[dict[str, Any]]: return list(self.training_claims) @@ -619,7 +639,9 @@ async def load_state_dict( "expected_manifest_digest": expected_manifest_digest, } ) - self._metadata_state = state + # A load repopulates rows; the envelope fields are the live buffer's own, + # so a later save still emits a complete state dict. + self._metadata_state = {**self._metadata_state, "groups": state["groups"]} self.target_step_list = [ group["target_step"] for group in self._metadata_state["groups"] ] @@ -924,7 +946,7 @@ def _run_restore_then_train_pump( async def _main(): actor = _ACTOR_CLS(mc, actor_args, SetupTimingMetrics()) await actor._maybe_restore_replay_buffer() - actor._sampler = _RestoredGroupsSampler(restored_groups) + actor._sampler = _RestoredGroupsSampler(restored_groups, actor._buffer) with patch("ray.cluster_resources", return_value={"GPU": 0}): await asyncio.wait_for(actor._train_pump(), timeout=60.0) actor._checkpointer.shutdown()