From d3a404df1ddc5321438e478be99cd65a5eacf0c8 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 12 Aug 2026 15:30:56 -0400 Subject: [PATCH 01/10] fix(rl): report the GRPO metrics that say whether training is working The NeMo-RL logger was written against DPO and forwards almost nothing from a GRPO run. Three problems, all in the translation layer rather than upstream. Validation was reported for DPO only. The branch gated on a `loss` key, but GRPO's validation dict is {accuracy, avg_length} and carries no loss at all (nemo_rl/algorithms/grpo.py builds it in validate()). Every GRPO validation pass was therefore dropped silently -- no report, no best-metric update, no log line saying anything had been skipped. `val_loss` is now optional: report whenever anything usable arrived, and omit the key rather than sending null, which would chart as a real zero. The metric key list was DPO-shaped. `preference_loss` and `rewards_rejected_mean` do not exist under GRPO, so a GRPO run forwarded only loss/lr/grad_norm and three token counts -- no reward, no advantages, no KL, no truncation rate. Reward is the metric that says whether RL is converging, and it never left the pod. The list is now a union across algorithms, selected by presence, covering the reward/advantage block, policy-optimization health, and NeMo-Gym's rollout metrics. Widening the list meant `has_metric_value` had to stop assuming numbers. NeMo-RL's metric dicts interleave non-scalars with the scalars: calculate_single_metric emits a `/histogram` Histogram object, NeMo-Gym adds a per-agent `full_result` Table, and `generation_logger_metrics` is a nested dict. `math.isnan` raises TypeError on all three, so the old None-check would have crashed mid-training on the first widened key rather than dropping the value. Separately, RL never accumulated a metric time series. `report_running` REPLACES the task's status_details blob, so reporting only the current step leaves no history -- the loss curve was unrecoverable and fetch_current_metrics always came back empty, so resume could not seed either. The shared customization callback has done this since it was written; the RL copy simply had not. train_loss/val_loss now accumulate in the same {step, epoch, value} shape Studio already renders as CustomizationMetricValue[], and every report path carries the payload, including report_training_start -- which would otherwise blank a resumed job's seeded series. Only those two series accumulate. The wider metric set rides along as current-step scalars: every series is resent in full on every update, so the payload grows with series count times step count. Putting the whole GRPO surface on that wire needs the Jobs metrics transport reworked, which is deliberately out of scope here. Both files were byte-identical to main before this change. Signed-off-by: Albert Cui --- .../training/backends/nemo_rl/callbacks.py | 70 +++- .../backends/nemo_rl/nemo_rl_logger.py | 106 ++++-- services/rl/tests/test_nemo_rl_callbacks.py | 178 ++++++++++ services/rl/tests/test_nemo_rl_logger.py | 311 ++++++++++++++++++ 4 files changed, 622 insertions(+), 43 deletions(-) create mode 100644 services/rl/tests/test_nemo_rl_callbacks.py create mode 100644 services/rl/tests/test_nemo_rl_logger.py diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py index 8ac0760844..b9a2ad1048 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py @@ -22,15 +22,50 @@ class TrainingProgressCallback: This class composes JobsServiceProgressReporter and provides training-specific methods for reporting detailed metrics during training. + + ``train_loss`` and ``val_loss`` are accumulated as time-series lists and resent + under a ``metrics`` key on every update, matching + ``nmp.customization_common.training.callbacks.TrainingProgressCallback``. This + is what makes a loss curve recoverable at all: ``report_running`` REPLACES the + task's ``status_details`` blob, so a report carrying only the current step + leaves no history behind. Studio reads exactly this shape + (``CustomizationMetricValue[]``). + + Only these two series accumulate. The wider RL metric set rides along as + current-step scalars, because every series is resent in full on every update + and the payload grows with the product of series count and step count. """ def __init__(self, reporter: JobsServiceProgressReporter): self._reporter = reporter + prior = reporter.fetch_current_metrics() + self._train_metrics: list[dict[str, float | int]] = prior.get("train_loss", []) + self._val_metrics: list[dict[str, float | int]] = prior.get("val_loss", []) + if self._train_metrics or self._val_metrics: + logger.info( + "Seeded metrics from server: %d train_loss, %d val_loss entries", + len(self._train_metrics), + len(self._val_metrics), + ) + + def _build_metrics_summary(self) -> dict[str, list[dict[str, float | int]]]: + """Build the accumulated metrics payload for inclusion in status_details.""" + return { + "train_loss": list(self._train_metrics), + "val_loss": list(self._val_metrics), + } + def report_training_start(self, max_steps: int, num_epochs: int) -> None: """Report that training has started with schedule information.""" self._reporter.configure_progress_tracking(max_steps, num_epochs) - self._reporter.report_running(phase="training", step=0, max_steps=max_steps, num_epochs=num_epochs) + self._reporter.report_running( + phase="training", + step=0, + max_steps=max_steps, + num_epochs=num_epochs, + metrics=self._build_metrics_summary(), + ) def report_train_step( self, @@ -49,9 +84,12 @@ def report_train_step( loss: Training loss value lr: Learning rate (optional) grad_norm: Gradient norm (optional) - **additional_metrics: Additional training metrics to report (e.g., num_valid_samples, - preference_loss, rewards_rejected_mean, global_valid_seqs, global_valid_toks) + **additional_metrics: Additional training metrics to report as current-step + scalars — DPO's preference_loss/rewards_rejected_mean, GRPO's + reward/advantages/kl_penalty, or the shared token counts. These are + not accumulated into the series; see the class docstring. """ + self._train_metrics.append({"step": step, "epoch": epoch, "value": loss}) self._reporter.report_running( phase="training", step=step, @@ -59,6 +97,7 @@ def report_train_step( train_loss=loss, lr=lr, grad_norm=grad_norm, + metrics=self._build_metrics_summary(), **additional_metrics, ) @@ -66,7 +105,7 @@ def report_validation( self, step: int, epoch: int, - val_loss: float, + val_loss: float | None = None, **additional_metrics: Any, ) -> None: """Report validation results. @@ -74,21 +113,34 @@ def report_validation( Args: step: Training step number epoch: Current epoch number - val_loss: Validation loss value + val_loss: Validation loss value, or None for algorithms that do not + produce one. GRPO validates on ``accuracy``/``avg_length`` and + reports no loss at all, so the key is omitted rather than sent as + null and charted as zero. **additional_metrics: Additional validation metrics to report (e.g., accuracy, num_valid_samples, or any other validation-specific metrics) """ + details: dict[str, Any] = {"step": step, "epoch": epoch} + if val_loss is not None: + self._val_metrics.append({"step": step, "epoch": epoch, "value": val_loss}) + details["val_loss"] = val_loss + self._reporter.report_running( phase="validation", - step=step, - epoch=epoch, - val_loss=val_loss, + metrics=self._build_metrics_summary(), + **details, **additional_metrics, ) def report_checkpoint_saved(self, step: int, epoch: int, checkpoint_path: str | None = None) -> None: """Report that a checkpoint was saved.""" - self._reporter.report_running(phase="checkpoint_saved", step=step, epoch=epoch, checkpoint_path=checkpoint_path) + self._reporter.report_running( + phase="checkpoint_saved", + step=step, + epoch=epoch, + checkpoint_path=checkpoint_path, + metrics=self._build_metrics_summary(), + ) def close(self) -> None: """Clean up resources.""" diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py index f7534bc2c7..b08bc7d055 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py @@ -10,6 +10,7 @@ import logging import math +import numbers from typing import Any, Mapping, Optional from nemo_rl.utils.logger import LoggerInterface @@ -20,12 +21,58 @@ _logger = logging.getLogger(__name__) +# Metric keys forwarded to Jobs Service, in addition to loss/lr/grad_norm. +# +# A union across algorithms: selection is by presence, so a DPO run simply has no +# GRPO keys and vice versa. GRPO's `train` dict is the merge of the policy-loss +# metrics, the reward/advantage block and NeMo-Gym's rollout metrics +# (nemo_rl/algorithms/grpo.py builds it, nemo_rl/experience/rollouts.py supplies +# the rollout half), so all three families are represented here. +_TRAIN_METRIC_KEYS = ( + # Shared + "num_valid_samples", + "global_valid_seqs", + "global_valid_toks", + # DPO + "preference_loss", + "rewards_rejected_mean", + # GRPO: reward and advantages — the signal that says whether RL is working + "reward", + "total_reward/mean", + "advantages/mean", + "advantages/min", + "advantages/max", + # GRPO: policy-optimization health + "kl_penalty", + "approx_entropy", + "token_mult_prob_error", + # GRPO: rollout shape + "truncation_rate", + "natural_termination_rate", + "turns_per_sample/mean", + "mean_gen_tokens_per_sample", +) + +# GRPO reports `accuracy`/`avg_length` here and no loss at all; DPO reports loss. +_VALIDATION_METRIC_KEYS = _TRAIN_METRIC_KEYS + ("accuracy", "avg_length") + def has_metric_value(metric: Any) -> bool: - """Check if a metric has a valid value.""" - if metric is not None and not math.isnan(metric): - return True - return False + """Whether ``metric`` is a finite-enough scalar to forward to Jobs Service. + + The type check is load-bearing, not defensive. NeMo-RL's metric dicts carry + non-scalars alongside the numbers: ``calculate_single_metric`` emits a + ``/histogram`` holding a ``Histogram`` object, NeMo-Gym adds a + per-agent ``full_result`` ``Table``, and ``generation_logger_metrics`` is a + nested dict. ``math.isnan`` raises ``TypeError`` on all of those, so a bare + None-check would turn a widened key list into a crash mid-training. + + ``bool`` is rejected despite being an ``int`` subclass: no metric here is a + flag, and silently charting one as 0/1 is worse than dropping it. + """ + if isinstance(metric, bool) or not isinstance(metric, numbers.Real): + return False + return not math.isnan(float(metric)) class NemoRLLogger(LoggerInterface): @@ -110,7 +157,11 @@ def log_metrics( # Calculate epoch from step (epochs start from 1) epoch = ((step - 1) // self._steps_per_epoch) + 1 - # Handle training loss + # Handle training loss. + # + # The loss gate also de-duplicates: GRPO logs under `train` twice per step — + # once mid-step with the rollout metrics alone (no loss), then again with the + # full merged dict. Requiring loss keeps the second, complete one. if prefix == "train" and has_metric_value(metrics.get("loss")): # Only report at log_interval to reduce output if step % self._log_interval == 0: @@ -119,44 +170,26 @@ def log_metrics( lr = metrics.get("lr") grad_norm = metrics.get("grad_norm") - # Extract additional training metrics (whitelisted only) - additional_metrics = {} - for key in [ - "num_valid_samples", - "preference_loss", - "rewards_rejected_mean", - "global_valid_seqs", - "global_valid_toks", - ]: - if has_metric_value(metrics.get(key)): - additional_metrics[key] = metrics[key] - self._callback.report_train_step( step=step, epoch=epoch, loss=loss, lr=lr, grad_norm=grad_norm, - **additional_metrics, + **self._select_metrics(metrics, _TRAIN_METRIC_KEYS), ) - # Handle validation metrics + # Handle validation metrics. + # + # `val_loss` is optional because GRPO has none: its validation dict is + # {accuracy, avg_length}. Gating on loss here dropped every GRPO validation + # report silently, so instead report whenever anything usable came through. elif prefix and prefix.startswith("validation"): - if has_metric_value(metrics.get("loss")): - val_loss = metrics["loss"] - - # Extract additional validation metrics (whitelisted only) - additional_metrics = {} - for key in [ - "num_valid_samples", - "preference_loss", - "rewards_rejected_mean", - "global_valid_seqs", - "global_valid_toks", - ]: - if has_metric_value(metrics.get(key)): - additional_metrics[key] = metrics[key] + raw_val_loss = metrics.get("loss") + val_loss = raw_val_loss if has_metric_value(raw_val_loss) else None + additional_metrics = self._select_metrics(metrics, _VALIDATION_METRIC_KEYS) + if val_loss is not None or additional_metrics: self._callback.report_validation( step=step, epoch=epoch, @@ -164,12 +197,17 @@ def log_metrics( **additional_metrics, ) # Track best validation loss - if val_loss < self._best_metric_value: + if val_loss is not None and val_loss < self._best_metric_value: self._best_metric_value = val_loss self._best_epoch = epoch _logger.debug(f"log_metrics: step={step}, prefix={prefix}, metrics={metrics}") + @staticmethod + def _select_metrics(metrics: dict[str, Any], keys: tuple[str, ...]) -> dict[str, Any]: + """Pick the whitelisted keys that carry a forwardable scalar.""" + return {key: metrics[key] for key in keys if has_metric_value(metrics.get(key))} + def log_hyperparams(self, params: Mapping[str, Any]) -> None: """Log hyperparameters and report training start. diff --git a/services/rl/tests/test_nemo_rl_callbacks.py b/services/rl/tests/test_nemo_rl_callbacks.py new file mode 100644 index 0000000000..98bc9ec698 --- /dev/null +++ b/services/rl/tests/test_nemo_rl_callbacks.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the RL TrainingProgressCallback's metric accumulation. + +``JobsServiceProgressReporter.report_running`` REPLACES the task's ``status_details`` +blob rather than merging into it, so a report that carries only the current step +erases everything before it. These tests pin the consequence: every report must carry +the full accumulated ``metrics`` payload, in the ``{step, epoch, value}`` shape Studio +reads as ``CustomizationMetricValue[]``. +""" + +from __future__ import annotations + +from typing import Any, cast + +import pytest +from nmp.customization_common.training.progress import JobsServiceProgressReporter +from nmp.rl.tasks.training.backends.nemo_rl.callbacks import TrainingProgressCallback + + +class _RecordingReporter: + """Stands in for JobsServiceProgressReporter, capturing each report payload.""" + + def __init__(self, prior: dict[str, list[dict[str, Any]]] | None = None) -> None: + self._prior = prior or {"train_loss": [], "val_loss": []} + self.reports: list[dict[str, Any]] = [] + self.tracking: tuple[int, int] | None = None + self.closed = False + + def fetch_current_metrics(self) -> dict[str, list[dict[str, Any]]]: + return self._prior + + def configure_progress_tracking(self, max_steps: int, num_epochs: int) -> None: + self.tracking = (max_steps, num_epochs) + + def report_running(self, phase: str, **details: Any) -> None: + self.reports.append({"phase": phase, **details}) + + def close(self) -> None: + self.closed = True + + +@pytest.fixture +def reporter() -> _RecordingReporter: + return _RecordingReporter() + + +def _make_callback(reporter: _RecordingReporter) -> TrainingProgressCallback: + """Build the callback over a duck-typed reporter, narrowing the type once here.""" + return TrainingProgressCallback(cast(JobsServiceProgressReporter, reporter)) + + +# --------------------------------------------------------------------------- # +# Accumulation +# --------------------------------------------------------------------------- # + + +def test_train_loss_accumulates_across_steps(reporter: _RecordingReporter) -> None: + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5) + callback.report_train_step(step=2, epoch=1, loss=0.4) + callback.report_train_step(step=3, epoch=1, loss=0.3) + + # The final report carries the whole curve, not just the last point. + series = reporter.reports[-1]["metrics"]["train_loss"] + assert series == [ + {"step": 1, "epoch": 1, "value": 0.5}, + {"step": 2, "epoch": 1, "value": 0.4}, + {"step": 3, "epoch": 1, "value": 0.3}, + ] + + +def test_every_report_carries_the_full_series(reporter: _RecordingReporter) -> None: + """report_running replaces status_details, so an omission is a data loss.""" + callback = _make_callback(reporter) + callback.report_training_start(max_steps=10, num_epochs=1) + callback.report_train_step(step=1, epoch=1, loss=0.5) + callback.report_validation(step=1, epoch=1, val_loss=0.45) + callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") + + for report in reporter.reports: + assert "metrics" in report, report["phase"] + assert "train_loss" in report["metrics"] + assert "val_loss" in report["metrics"] + + +def test_val_loss_accumulates_separately(reporter: _RecordingReporter) -> None: + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5) + callback.report_validation(step=1, epoch=1, val_loss=0.45) + callback.report_validation(step=2, epoch=1, val_loss=0.40) + + metrics = reporter.reports[-1]["metrics"] + assert len(metrics["train_loss"]) == 1 + assert metrics["val_loss"] == [ + {"step": 1, "epoch": 1, "value": 0.45}, + {"step": 2, "epoch": 1, "value": 0.40}, + ] + + +def test_series_are_copies_not_live_references(reporter: _RecordingReporter) -> None: + """Each payload must snapshot the series; a shared list would mutate old reports.""" + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5) + first_payload = reporter.reports[-1]["metrics"]["train_loss"] + callback.report_train_step(step=2, epoch=1, loss=0.4) + + assert len(first_payload) == 1 + + +# --------------------------------------------------------------------------- # +# Resume seeding +# --------------------------------------------------------------------------- # + + +def test_prior_metrics_seed_the_series() -> None: + """A resumed job continues the curve instead of restarting it.""" + prior = { + "train_loss": [{"step": 1, "epoch": 1, "value": 0.9}], + "val_loss": [{"step": 1, "epoch": 1, "value": 0.8}], + } + reporter = _RecordingReporter(prior) + callback = _make_callback(reporter) + callback.report_train_step(step=2, epoch=1, loss=0.5) + + series = reporter.reports[-1]["metrics"]["train_loss"] + assert [entry["step"] for entry in series] == [1, 2] + + +def test_training_start_does_not_erase_seeded_metrics() -> None: + """report_training_start fires before the first step; it must not blank the blob.""" + prior = {"train_loss": [{"step": 1, "epoch": 1, "value": 0.9}], "val_loss": []} + reporter = _RecordingReporter(prior) + callback = _make_callback(reporter) + callback.report_training_start(max_steps=10, num_epochs=1) + + assert reporter.reports[0]["metrics"]["train_loss"] == prior["train_loss"] + + +# --------------------------------------------------------------------------- # +# Optional val_loss (GRPO) +# --------------------------------------------------------------------------- # + + +def test_validation_without_loss_omits_the_key(reporter: _RecordingReporter) -> None: + """GRPO validates on accuracy; a null val_loss would chart as zero.""" + callback = _make_callback(reporter) + callback.report_validation(step=1, epoch=1, val_loss=None, accuracy=0.75) + + report = reporter.reports[-1] + assert "val_loss" not in report + assert report["accuracy"] == 0.75 + assert report["phase"] == "validation" + + +def test_validation_without_loss_leaves_the_series_empty(reporter: _RecordingReporter) -> None: + callback = _make_callback(reporter) + callback.report_validation(step=1, epoch=1, val_loss=None, accuracy=0.75) + + assert reporter.reports[-1]["metrics"]["val_loss"] == [] + + +def test_additional_metrics_ride_along_as_scalars(reporter: _RecordingReporter) -> None: + """The wide RL metric set is current-step only; it must not enter the series.""" + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5, reward=0.62, kl_penalty=0.008) + + report = reporter.reports[-1] + assert report["reward"] == 0.62 + assert report["kl_penalty"] == 0.008 + assert report["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] + + +def test_close_delegates_to_the_reporter(reporter: _RecordingReporter) -> None: + _make_callback(reporter).close() + + assert reporter.closed diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py new file mode 100644 index 0000000000..02e9659e70 --- /dev/null +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -0,0 +1,311 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for NemoRLLogger's translation of NeMo-RL metrics into Jobs Service reports. + +The metric dicts here mirror what NeMo-RL actually hands the logger, including the +non-scalar entries (``Histogram``, ``Table``, nested dicts) that share the dict with +the numbers. Those are the reason ``has_metric_value`` type-checks rather than just +None-checks, so they are exercised rather than sanitised away. +""" + +from __future__ import annotations + +import importlib.util +import math +import sys +import types +from typing import Any + +import pytest + +# NeMo-RL is only installed inside the training image, so the LoggerInterface import +# at nemo_rl_logger module scope fails in a plain repo checkout. Stub just enough to +# import the module under test; when the real package IS present (the in-image smoke +# run) this is skipped and the genuine base class is used. +if importlib.util.find_spec("nemo_rl") is None: # pragma: no cover - env dependent + _nemo_rl = types.ModuleType("nemo_rl") + _utils = types.ModuleType("nemo_rl.utils") + _logger_mod = types.ModuleType("nemo_rl.utils.logger") + + class LoggerInterface: # minimal stand-in for the abstract base + pass + + setattr(_logger_mod, "LoggerInterface", LoggerInterface) + sys.modules.setdefault("nemo_rl", _nemo_rl) + sys.modules.setdefault("nemo_rl.utils", _utils) + sys.modules.setdefault("nemo_rl.utils.logger", _logger_mod) + +from nmp.rl.tasks.training.backends.nemo_rl import nemo_rl_logger # noqa: E402 +from nmp.rl.tasks.training.backends.nemo_rl.nemo_rl_logger import ( # noqa: E402 + NemoRLLogger, + has_metric_value, +) + + +class _RecordingCallback: + """Stands in for TrainingProgressCallback, capturing what the logger forwards.""" + + def __init__(self) -> None: + self.train_steps: list[dict[str, Any]] = [] + self.validations: list[dict[str, Any]] = [] + self.training_starts: list[dict[str, Any]] = [] + self.closed = False + + def report_training_start(self, max_steps: int, num_epochs: int) -> None: + self.training_starts.append({"max_steps": max_steps, "num_epochs": num_epochs}) + + def report_train_step(self, step, epoch, loss, lr=None, grad_norm=None, **additional): + self.train_steps.append( + {"step": step, "epoch": epoch, "loss": loss, "lr": lr, "grad_norm": grad_norm, **additional} + ) + + def report_validation(self, step, epoch, val_loss=None, **additional): + self.validations.append({"step": step, "epoch": epoch, "val_loss": val_loss, **additional}) + + def close(self) -> None: + self.closed = True + + +@pytest.fixture +def callback(monkeypatch: pytest.MonkeyPatch) -> _RecordingCallback: + """Build a NemoRLLogger whose reporter/callback are inert local objects.""" + recorder = _RecordingCallback() + monkeypatch.setattr(nemo_rl_logger, "JobsServiceProgressReporter", lambda *a, **k: object()) + monkeypatch.setattr(nemo_rl_logger, "TrainingProgressCallback", lambda _reporter: recorder) + return recorder + + +def _make_logger(**kwargs: Any) -> NemoRLLogger: + params: dict[str, Any] = {"steps_per_epoch": 10, "log_interval": 1} + params.update(kwargs) + return NemoRLLogger(**params) + + +class _Histogram: + """Stand-in for nemo_rl's wandb Histogram — non-numeric, and NaN-hostile.""" + + +# Trimmed but faithful shape of GRPO's `train` dict: policy-loss metrics, the +# reward/advantage block, NeMo-Gym rollout metrics, and the non-scalars that ride +# along with them. +GRPO_TRAIN_METRICS: dict[str, Any] = { + "loss": 0.31, + "lr": 5e-6, + "grad_norm": 1.7, + "reward": 0.62, + "total_reward/mean": 0.62, + "total_reward/histogram": _Histogram(), + "advantages/mean": 0.04, + "advantages/min": -1.2, + "advantages/max": 1.4, + "kl_penalty": 0.008, + "approx_entropy": 0.55, + "token_mult_prob_error": 1.02, + "truncation_rate": 0.125, + "natural_termination_rate": 0.875, + "turns_per_sample/mean": 2.5, + "mean_gen_tokens_per_sample": 148.0, + "num_valid_samples": 64, + "global_valid_seqs": 64.0, + "global_valid_toks": 9472.0, + "generation_logger_metrics": {"inflight": [1, 2, 3]}, + "per_worker_token_counts": [{0: 100, 1: 120}], + "ascii_tree_agent/full_result": object(), +} + +# GRPO validation, verbatim in shape: no `loss` key anywhere. +GRPO_VALIDATION_METRICS: dict[str, Any] = {"accuracy": 0.75, "avg_length": 143.2} + +DPO_TRAIN_METRICS: dict[str, Any] = { + "loss": 0.5, + "lr": 1e-5, + "grad_norm": 0.9, + "preference_loss": 0.42, + "rewards_rejected_mean": -0.3, + "num_valid_samples": 8, + "global_valid_seqs": 8.0, + "global_valid_toks": 1024.0, +} + + +# --------------------------------------------------------------------------- # +# has_metric_value +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "value,expected", + [ + (0.5, True), + (0, True), + (-1.5, True), + (float("nan"), False), + (None, False), + # Non-scalars that genuinely appear in NeMo-RL metric dicts. Each of these + # raises TypeError under a bare math.isnan, which is the regression guarded here. + (_Histogram(), False), + ({"inflight": [1, 2]}, False), + ([1, 2, 3], False), + ("0.5", False), + # bool is an int subclass; charting a flag as 0/1 is not wanted. + (True, False), + (False, False), + ], +) +def test_has_metric_value(value: Any, expected: bool) -> None: + assert has_metric_value(value) is expected + + +def test_has_metric_value_does_not_raise_on_any_grpo_metric() -> None: + """Every value in a real GRPO dict must be classifiable without raising.""" + for key, value in GRPO_TRAIN_METRICS.items(): + assert isinstance(has_metric_value(value), bool), key + + +def test_has_metric_value_accepts_numpy_scalars() -> None: + np = pytest.importorskip("numpy") + assert has_metric_value(np.float32(0.5)) is True + assert has_metric_value(np.float64(0.5)) is True + assert has_metric_value(np.int64(3)) is True + assert has_metric_value(np.float32("nan")) is False + + +# --------------------------------------------------------------------------- # +# GRPO train metrics +# --------------------------------------------------------------------------- # + + +def test_grpo_train_step_forwards_reward_and_rollout_metrics(callback: _RecordingCallback) -> None: + """The reward signal is the point of GRPO; it must reach Jobs Service.""" + _make_logger().log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + + assert len(callback.train_steps) == 1 + reported = callback.train_steps[0] + assert reported["loss"] == 0.31 + assert reported["reward"] == 0.62 + assert reported["total_reward/mean"] == 0.62 + assert reported["advantages/mean"] == 0.04 + assert reported["kl_penalty"] == 0.008 + assert reported["approx_entropy"] == 0.55 + assert reported["truncation_rate"] == 0.125 + assert reported["turns_per_sample/mean"] == 2.5 + + +def test_grpo_train_step_drops_non_scalar_metrics(callback: _RecordingCallback) -> None: + """Histograms/Tables/nested dicts must not be forwarded, and must not raise.""" + _make_logger().log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + + reported = callback.train_steps[0] + for key in ( + "total_reward/histogram", + "generation_logger_metrics", + "per_worker_token_counts", + "ascii_tree_agent/full_result", + ): + assert key not in reported + + +def test_rollout_only_train_call_is_ignored(callback: _RecordingCallback) -> None: + """GRPO logs `train` twice per step; the mid-step call has no loss and is a partial.""" + rollout_only = {k: v for k, v in GRPO_TRAIN_METRICS.items() if k != "loss"} + _make_logger().log_metrics(rollout_only, step=0, prefix="train") + + assert callback.train_steps == [] + + +def test_dpo_train_metrics_still_forwarded(callback: _RecordingCallback) -> None: + """Widening the key list for GRPO must not drop DPO's existing metrics.""" + _make_logger().log_metrics(DPO_TRAIN_METRICS, step=0, prefix="train") + + reported = callback.train_steps[0] + assert reported["loss"] == 0.5 + assert reported["preference_loss"] == 0.42 + assert reported["rewards_rejected_mean"] == -0.3 + assert reported["global_valid_toks"] == 1024.0 + + +def test_log_interval_throttles_train_reports(callback: _RecordingCallback) -> None: + logger = _make_logger(log_interval=5) + for step in range(10): + logger.log_metrics(GRPO_TRAIN_METRICS, step=step, prefix="train") + + # log_metrics increments step by 1, so steps 5 and 10 report. + assert [r["step"] for r in callback.train_steps] == [5, 10] + + +# --------------------------------------------------------------------------- # +# Validation — the branch GRPO never reached +# --------------------------------------------------------------------------- # + + +def test_grpo_validation_is_reported_without_a_loss(callback: _RecordingCallback) -> None: + """GRPO validates on accuracy/avg_length and reports no loss. + + Gating this branch on `loss` silently dropped every GRPO validation report. + """ + _make_logger().log_metrics(GRPO_VALIDATION_METRICS, step=9, prefix="validation") + + assert len(callback.validations) == 1 + reported = callback.validations[0] + assert reported["val_loss"] is None + assert reported["accuracy"] == 0.75 + assert reported["avg_length"] == 143.2 + assert reported["step"] == 10 + assert reported["epoch"] == 1 + + +def test_dpo_validation_still_reports_loss(callback: _RecordingCallback) -> None: + _make_logger().log_metrics({"loss": 0.25, "num_valid_samples": 8}, step=9, prefix="validation") + + reported = callback.validations[0] + assert reported["val_loss"] == 0.25 + assert reported["num_valid_samples"] == 8 + + +def test_validation_with_nothing_usable_is_not_reported(callback: _RecordingCallback) -> None: + """An empty or all-non-scalar dict must not produce a hollow report.""" + _make_logger().log_metrics({}, step=9, prefix="validation") + _make_logger().log_metrics({"histogram/x": _Histogram()}, step=9, prefix="validation") + + assert callback.validations == [] + + +def test_best_validation_loss_tracks_minimum(callback: _RecordingCallback) -> None: + logger = _make_logger() + logger.log_metrics({"loss": 0.5}, step=9, prefix="validation") + logger.log_metrics({"loss": 0.2}, step=19, prefix="validation") + logger.log_metrics({"loss": 0.7}, step=29, prefix="validation") + + assert logger._best_metric_value == 0.2 + assert logger._best_epoch == 2 + + +def test_grpo_validation_leaves_best_loss_untouched(callback: _RecordingCallback) -> None: + """No loss means no best-loss update — and no crash comparing None.""" + logger = _make_logger() + logger.log_metrics(GRPO_VALIDATION_METRICS, step=9, prefix="validation") + + assert math.isinf(logger._best_metric_value) + assert logger._best_epoch is None + + +@pytest.mark.parametrize("prefix", ["validation", "validation-0", "validation/nemo_gym"]) +def test_all_validation_prefixes_are_handled(callback: _RecordingCallback, prefix: str) -> None: + """NeMo-RL suffixes the prefix per dataloader; all must route to validation.""" + _make_logger().log_metrics(GRPO_VALIDATION_METRICS, step=9, prefix=prefix) + + assert len(callback.validations) == 1 + + +# --------------------------------------------------------------------------- # +# Prefixes we intentionally ignore +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("prefix", ["timing/train", "timing/validation", "timing/setup", "performance", "refit", ""]) +def test_unhandled_prefixes_produce_no_reports(callback: _RecordingCallback, prefix: str) -> None: + _make_logger().log_metrics({"loss": 0.1, "total_step_time": 12.0}, step=0, prefix=prefix) + + assert callback.train_steps == [] + assert callback.validations == [] From 9961fba19091ce90f20c19ce338271fb00068674 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 12 Aug 2026 16:51:06 -0400 Subject: [PATCH 02/10] test(rl): stop the chat_template test depending on ambient transformers test_tokenizer_omits_chat_template_when_none asserted that a model with no chat template omits the key, and established the premise indirectly: the fixture model dir has no tokenizer, so resolve_chat_template was expected to fall through to None. It does not, once another suite has run in the same session. services/automodel/tests/tasks/training/backends/test_config.py installs `sys.modules.setdefault("transformers", MagicMock())` at module scope and never removes it, so resolve_chat_template's AutoTokenizer.from_pretrained(...) call returns a Mock whose `.chat_template` is truthy. The key is then present and the assertion fails -- but only when the automodel tests are collected first, which is why it passed per-service and failed in a combined run. Patch resolve_chat_template directly instead, mirroring the neighbouring test_tokenizer_keeps_chat_template_when_present. That states the actual premise rather than arranging for it, and is immune to what else is in sys.modules. Does not address the leak itself; two unsloth hf_trainer_callback tests fail the same way (`from transformers import TrainerCallback` yields a Mock base class) and are untouched here. Signed-off-by: Albert Cui --- services/rl/tests/test_grpo_config.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/services/rl/tests/test_grpo_config.py b/services/rl/tests/test_grpo_config.py index 666d0eb11a..017b5da51c 100644 --- a/services/rl/tests/test_grpo_config.py +++ b/services/rl/tests/test_grpo_config.py @@ -261,8 +261,15 @@ def test_tokenizer_omits_chat_template_when_none( one. Qwen3 has one, so a single-model GPU run never sees this. """ monkeypatch.setenv("NMP_JOB_STORAGE_PVC_CLAIM", "nmp-job-storage") + # Patched rather than relying on the fixture model dir having no tokenizer: + # resolution falls through to transformers.AutoTokenizer, and another suite in + # the same session installs a module-scope `transformers` MagicMock into + # sys.modules, whose truthy `.chat_template` silently invalidates the premise. + monkeypatch.setattr( + "nmp.rl.tasks.training.backends.nemo_rl.grpo_config.resolve_chat_template", + lambda **_: None, + ) step, _ = _prepared_step(tmp_path) - # The fixture model dir has no tokenizer, so resolution falls through to None. tokenizer = compile_grpo_config(step, job_ctx)["policy"]["tokenizer"] assert "chat_template" not in tokenizer From 5a39483156721cb2ed4e776e51ea659a7bf8fa44 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 12 Aug 2026 16:51:25 -0400 Subject: [PATCH 03/10] refactor(rl): fold the RL progress callback onto the shared implementation RL carried a standalone TrainingProgressCallback rather than subclassing the shared one, and the previous commit made that worse by copying ~25 lines of accumulation logic into it. Every other customization service (unsloth, automodel) subclasses the shared class; RL was the outlier, and the one that had drifted ahead in features. The three RL-only behaviours are additive, so they move up into the shared class rather than justifying a fork: `**additional_metrics` for backend-specific current-step scalars, and an optional `val_loss` for algorithms that do not produce one. RL's copy becomes a two-line subclass, leaving `_default_backend` as None so no `backend` key is added and its status-detail shape is unchanged. Promoting the accumulation exposed that the shared class had the same data-loss hole this branch just closed for RL. report_training_start, report_epoch_end and report_checkpoint_saved all omitted the `metrics` payload, and report_running REPLACES status_details rather than merging, so any of them erased the accumulated series from stored status until the next train step resent it -- and lost it outright if the job died in that window. automodel calls both report_epoch_end and report_checkpoint_saved mid-training, so this was reachable in practice, not theoretical. Two automodel tests and one unsloth test pinned the buggy payload with exact-kwargs assertions; they now assert the series survives instead. Also adds the missing services/rl/.../training/progress.py, matching the unsloth and automodel modules that bind SERVICE_NAME, so the two RL construction sites stop passing it by hand. Two defects fixed while in here: The final training step was never reported. The throttle is `step % log_interval == 0`, so when max_steps is not a multiple of log_interval the last steps are dropped -- at 23 steps and an interval of 10 the run's last recorded loss was step 20's. A withheld step is now held as pending and flushed by close(), the only end-of-run hook available (`step_finished` is per-step). The flush is reachable from __del__, so it swallows and logs rather than raising. The two drivers derived the same two parameters with different formulas: at val_period=100 DPO computed a log_interval of 11 and GRPO 10, and steps_per_epoch was read from config in one and derived in the other. Both now call NemoRLLogger.for_schedule, which owns the arithmetic and still prefers an explicit steps_per_epoch when the algorithm config carries one (DPO does). DPO's reporting cadence changes slightly as a result -- its `+1` was a divide-by-zero guard that also skewed every value. Behaviour-preserving elsewhere: no wire-shape change for any backend beyond the `metrics` payload now being present on reports that previously dropped it. Signed-off-by: Albert Cui --- .../training/callbacks.py | 74 ++++++-- .../tests/training/test_callbacks.py | 157 +++++++++++++++ .../tasks/training/backends/test_callbacks.py | 38 +++- .../training/backends/nemo_rl/callbacks.py | 144 ++------------ .../training/backends/nemo_rl/dpo_driver.py | 16 +- .../training/backends/nemo_rl/grpo_driver.py | 13 +- .../backends/nemo_rl/nemo_rl_logger.py | 100 ++++++++-- .../rl/src/nmp/rl/tasks/training/progress.py | 33 ++++ .../rl/src/nmp/rl/tasks/training/runner.py | 6 +- services/rl/tests/test_nemo_rl_callbacks.py | 179 ++---------------- services/rl/tests/test_nemo_rl_logger.py | 105 ++++++++++ services/unsloth/tests/test_callbacks.py | 1 + 12 files changed, 514 insertions(+), 352 deletions(-) create mode 100644 packages/nmp_customization_common/tests/training/test_callbacks.py create mode 100644 services/rl/src/nmp/rl/tasks/training/progress.py diff --git a/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py b/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py index 870b64d6ae..2b57105d63 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py @@ -5,14 +5,25 @@ Composes a :class:`nmp.customization_common.training.progress.JobsServiceProgressReporter` and provides training-specific methods. Metric accumulation: ``train_loss`` and -``val_loss`` are accumulated as time-series lists and included in every +``val_loss`` are accumulated as time-series lists and included in EVERY ``status_details`` update under a ``metrics`` key, enabling loss-curve reconstruction from job status. +Every update matters because ``report_running`` REPLACES the task's +``status_details`` blob rather than merging into it. A report that omits +``metrics`` therefore erases the accumulated series from stored status until the +next train step resends it -- and if the job dies inside that window, the curve +is gone. Checkpoint and epoch-end reports fire mid-training, so they carry the +payload too. + +Only these two series accumulate. Anything passed as ``**additional_metrics`` +rides along as a current-step scalar: the full series set is resent on every +update, so the payload grows with series count times step count. + Backends subclass this and set :attr:`_default_backend`: unsloth stamps a -``backend`` field on each report (``"unsloth"``); automodel leaves it ``None`` so -no ``backend`` key is added (preserving its status-detail shape). Callers may also -pass ``backend`` per call (e.g. unsloth's HF trainer callback). +``backend`` field on each report (``"unsloth"``); automodel and NeMo-RL leave it +``None`` so no ``backend`` key is added (preserving their status-detail shape). +Callers may also pass ``backend`` per call (e.g. unsloth's HF trainer callback). """ import logging @@ -56,7 +67,12 @@ def _build_metrics_summary(self) -> dict[str, list[dict[str, float | int]]]: def report_training_start(self, max_steps: int, num_epochs: int, *, backend: str | None = None) -> None: """Report that training has started with schedule information.""" self._reporter.configure_progress_tracking(max_steps, num_epochs) - details: dict[str, object] = {"step": 0, "max_steps": max_steps, "num_epochs": num_epochs} + details: dict[str, object] = { + "step": 0, + "max_steps": max_steps, + "num_epochs": num_epochs, + "metrics": self._build_metrics_summary(), + } resolved = self._resolve_backend(backend) if resolved is not None: details["backend"] = resolved @@ -71,8 +87,14 @@ def report_train_step( grad_norm: float | None = None, *, backend: str | None = None, + **additional_metrics: object, ) -> None: - """Report training step with metrics.""" + """Report training step with metrics. + + ``additional_metrics`` are backend-specific current-step scalars (DPO's + ``preference_loss``, GRPO's ``reward``/``kl_penalty``, ...). They are not + accumulated into the series; see the module docstring. + """ self._train_metrics.append({"step": step, "epoch": epoch, "value": loss}) details: dict[str, object] = { "step": step, @@ -81,21 +103,38 @@ def report_train_step( "lr": lr, "grad_norm": grad_norm, "metrics": self._build_metrics_summary(), + **additional_metrics, } resolved = self._resolve_backend(backend) if resolved is not None: details["backend"] = resolved self._reporter.report_running(phase="training", **details) - def report_validation(self, step: int, epoch: int, val_loss: float, *, backend: str | None = None) -> None: - """Report validation results.""" - self._val_metrics.append({"step": step, "epoch": epoch, "value": val_loss}) + def report_validation( + self, + step: int, + epoch: int, + val_loss: float | None = None, + *, + backend: str | None = None, + **additional_metrics: object, + ) -> None: + """Report validation results. + + ``val_loss`` is optional because not every algorithm produces one: GRPO + validates on ``accuracy``/``avg_length`` and reports no loss at all. The + key is omitted rather than sent as null, which would chart as a real zero. + """ details: dict[str, object] = { "step": step, "epoch": epoch, - "val_loss": val_loss, - "metrics": self._build_metrics_summary(), + **additional_metrics, } + if val_loss is not None: + self._val_metrics.append({"step": step, "epoch": epoch, "value": val_loss}) + details["val_loss"] = val_loss + details["metrics"] = self._build_metrics_summary() + resolved = self._resolve_backend(backend) if resolved is not None: details["backend"] = resolved @@ -110,7 +149,12 @@ def report_checkpoint_saved( backend: str | None = None, ) -> None: """Report that a checkpoint was saved.""" - details: dict[str, object] = {"step": step, "epoch": epoch, "checkpoint_path": checkpoint_path} + details: dict[str, object] = { + "step": step, + "epoch": epoch, + "checkpoint_path": checkpoint_path, + "metrics": self._build_metrics_summary(), + } resolved = self._resolve_backend(backend) if resolved is not None: details["backend"] = resolved @@ -118,7 +162,11 @@ def report_checkpoint_saved( def report_epoch_end(self, step: int, epoch: int, *, backend: str | None = None) -> None: """Report that an epoch has completed.""" - details: dict[str, object] = {"step": step, "epoch": epoch} + details: dict[str, object] = { + "step": step, + "epoch": epoch, + "metrics": self._build_metrics_summary(), + } resolved = self._resolve_backend(backend) if resolved is not None: details["backend"] = resolved diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py new file mode 100644 index 0000000000..abe5a149d0 --- /dev/null +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the shared TrainingProgressCallback. + +Focused on the contract that every backend depends on: ``report_running`` REPLACES +the task's ``status_details``, so the accumulated series must ride on every report +or it is erased from stored status. Basic accumulation and resume-seeding are +covered by the per-backend suites; this file covers the shared surface itself. +""" + +from __future__ import annotations + +from typing import Any, ClassVar, cast + +import pytest +from nmp.customization_common.training.callbacks import TrainingProgressCallback +from nmp.customization_common.training.progress import JobsServiceProgressReporter + + +class _RecordingReporter: + """Stands in for JobsServiceProgressReporter, capturing each report payload.""" + + def __init__(self, prior: dict[str, list[dict[str, Any]]] | None = None) -> None: + self._prior = prior or {"train_loss": [], "val_loss": []} + self.reports: list[dict[str, Any]] = [] + self.tracking: tuple[int, int] | None = None + self.closed = False + + def fetch_current_metrics(self) -> dict[str, list[dict[str, Any]]]: + return self._prior + + def configure_progress_tracking(self, max_steps: int, num_epochs: int) -> None: + self.tracking = (max_steps, num_epochs) + + def report_running(self, phase: str, **details: Any) -> None: + self.reports.append({"phase": phase, **details}) + + def close(self) -> None: + self.closed = True + + +@pytest.fixture +def reporter() -> _RecordingReporter: + return _RecordingReporter() + + +def _make_callback(reporter: _RecordingReporter) -> TrainingProgressCallback: + """Build the callback over a duck-typed reporter, narrowing the type once here.""" + return TrainingProgressCallback(cast(JobsServiceProgressReporter, reporter)) + + +# --------------------------------------------------------------------------- # +# Every report path carries the series +# --------------------------------------------------------------------------- # + + +def test_every_report_path_carries_the_series(reporter: _RecordingReporter) -> None: + """An omitted payload erases the curve from stored status_details.""" + callback = _make_callback(reporter) + callback.report_training_start(max_steps=10, num_epochs=1) + callback.report_train_step(step=1, epoch=1, loss=0.5) + callback.report_validation(step=1, epoch=1, val_loss=0.45) + callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") + callback.report_epoch_end(step=1, epoch=1) + + assert len(reporter.reports) == 5 + for report in reporter.reports: + assert "metrics" in report, report["phase"] + assert set(report["metrics"]) == {"train_loss", "val_loss"} + + +def test_training_start_does_not_erase_seeded_metrics() -> None: + """report_training_start fires before the first step; it must not blank the blob.""" + prior = {"train_loss": [{"step": 1, "epoch": 1, "value": 0.9}], "val_loss": []} + reporter = _RecordingReporter(prior) + _make_callback(reporter).report_training_start(max_steps=10, num_epochs=1) + + assert reporter.reports[0]["metrics"]["train_loss"] == prior["train_loss"] + + +def test_series_are_snapshots_not_live_references(reporter: _RecordingReporter) -> None: + """A shared list would retroactively mutate already-sent payloads.""" + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5) + first_payload = reporter.reports[-1]["metrics"]["train_loss"] + callback.report_train_step(step=2, epoch=1, loss=0.4) + + assert len(first_payload) == 1 + + +# --------------------------------------------------------------------------- # +# Optional val_loss +# --------------------------------------------------------------------------- # + + +def test_validation_without_loss_omits_the_key(reporter: _RecordingReporter) -> None: + """GRPO validates on accuracy; a null val_loss would chart as a real zero.""" + _make_callback(reporter).report_validation(step=1, epoch=1, val_loss=None, accuracy=0.75) + + report = reporter.reports[-1] + assert "val_loss" not in report + assert report["accuracy"] == 0.75 + assert report["metrics"]["val_loss"] == [] + + +def test_validation_with_loss_records_both_key_and_series(reporter: _RecordingReporter) -> None: + _make_callback(reporter).report_validation(step=1, epoch=1, val_loss=0.25) + + report = reporter.reports[-1] + assert report["val_loss"] == 0.25 + assert report["metrics"]["val_loss"] == [{"step": 1, "epoch": 1, "value": 0.25}] + + +# --------------------------------------------------------------------------- # +# additional_metrics +# --------------------------------------------------------------------------- # + + +def test_additional_train_metrics_ride_along_without_entering_the_series( + reporter: _RecordingReporter, +) -> None: + """The wide backend metric set is current-step only; series stay bounded.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, reward=0.62, kl_penalty=0.008) + + report = reporter.reports[-1] + assert report["reward"] == 0.62 + assert report["kl_penalty"] == 0.008 + assert report["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] + + +def test_additional_validation_metrics_ride_along(reporter: _RecordingReporter) -> None: + _make_callback(reporter).report_validation(step=1, epoch=1, val_loss=0.25, accuracy=0.9) + + assert reporter.reports[-1]["accuracy"] == 0.9 + + +def test_additional_metrics_do_not_collide_with_backend_stamping( + reporter: _RecordingReporter, +) -> None: + """`backend` is keyword-only, so **additional_metrics can never capture it.""" + + class _Stamped(TrainingProgressCallback): + _default_backend: ClassVar[str | None] = "test-backend" + + callback = _Stamped(cast(JobsServiceProgressReporter, reporter)) + callback.report_train_step(step=1, epoch=1, loss=0.5, reward=0.62) + + report = reporter.reports[-1] + assert report["backend"] == "test-backend" + assert report["reward"] == 0.62 + + +def test_close_delegates_to_the_reporter(reporter: _RecordingReporter) -> None: + _make_callback(reporter).close() + + assert reporter.closed diff --git a/services/automodel/tests/tasks/training/backends/test_callbacks.py b/services/automodel/tests/tasks/training/backends/test_callbacks.py index 71dcb278ca..46132c3903 100644 --- a/services/automodel/tests/tasks/training/backends/test_callbacks.py +++ b/services/automodel/tests/tasks/training/backends/test_callbacks.py @@ -142,7 +142,13 @@ def test_report_training_start_delegates(self): callback.report_training_start(max_steps=500, num_epochs=2) reporter.configure_progress_tracking.assert_called_once_with(500, 2) - reporter.report_running.assert_called_once_with(phase="training", step=0, max_steps=500, num_epochs=2) + reporter.report_running.assert_called_once_with( + phase="training", + step=0, + max_steps=500, + num_epochs=2, + metrics={"train_loss": [], "val_loss": []}, + ) def test_report_checkpoint_saved_delegates(self): callback, reporter = self._make_callback() @@ -150,9 +156,37 @@ def test_report_checkpoint_saved_delegates(self): callback.report_checkpoint_saved(step=100, epoch=1, checkpoint_path="/tmp/ckpt") reporter.report_running.assert_called_once_with( - phase="checkpoint_saved", step=100, epoch=1, checkpoint_path="/tmp/ckpt" + phase="checkpoint_saved", + step=100, + epoch=1, + checkpoint_path="/tmp/ckpt", + metrics={"train_loss": [], "val_loss": []}, ) + def test_checkpoint_report_preserves_accumulated_series(self): + """report_running REPLACES status_details, so an omitted payload erases the curve. + + Checkpoint saves fire mid-training (finetune.py calls this from the save + hook), so a report without `metrics` would drop the series from stored + status until the next train step -- and lose it entirely if the job then died. + """ + callback, reporter = self._make_callback() + + callback.report_train_step(step=1, epoch=1, loss=3.21) + callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/tmp/ckpt") + + kwargs = self._last_report_kwargs(reporter) + assert kwargs["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 3.21}] + + def test_epoch_end_report_preserves_accumulated_series(self): + callback, reporter = self._make_callback() + + callback.report_train_step(step=1, epoch=1, loss=3.21) + callback.report_epoch_end(step=1, epoch=1) + + kwargs = self._last_report_kwargs(reporter) + assert kwargs["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 3.21}] + def test_close_delegates(self): callback, reporter = self._make_callback() callback.close() diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py index b9a2ad1048..f86444ae8d 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py @@ -8,140 +8,20 @@ # without an express license agreement from NVIDIA CORPORATION or # its affiliates is strictly prohibited. -import logging -from typing import Any +"""Training progress callbacks for NeMo-RL Jobs-service reporting. -from nmp.customization_common.training.progress import JobsServiceProgressReporter +Thin subclass of the shared +:class:`nmp.customization_common.training.callbacks.TrainingProgressCallback`, +matching the unsloth and automodel pattern. ``_default_backend`` stays ``None`` +so no ``backend`` key is added and the existing status-detail shape is preserved. +""" -logger = logging.getLogger(__name__) +from nmp.customization_common.training.callbacks import ( + TrainingProgressCallback as _BaseTrainingProgressCallback, +) +__all__ = ["TrainingProgressCallback"] -class TrainingProgressCallback: - """ - Callback for reporting NeMo RL training progress to the Jobs service. - This class composes JobsServiceProgressReporter and provides training-specific - methods for reporting detailed metrics during training. - - ``train_loss`` and ``val_loss`` are accumulated as time-series lists and resent - under a ``metrics`` key on every update, matching - ``nmp.customization_common.training.callbacks.TrainingProgressCallback``. This - is what makes a loss curve recoverable at all: ``report_running`` REPLACES the - task's ``status_details`` blob, so a report carrying only the current step - leaves no history behind. Studio reads exactly this shape - (``CustomizationMetricValue[]``). - - Only these two series accumulate. The wider RL metric set rides along as - current-step scalars, because every series is resent in full on every update - and the payload grows with the product of series count and step count. - """ - - def __init__(self, reporter: JobsServiceProgressReporter): - self._reporter = reporter - - prior = reporter.fetch_current_metrics() - self._train_metrics: list[dict[str, float | int]] = prior.get("train_loss", []) - self._val_metrics: list[dict[str, float | int]] = prior.get("val_loss", []) - if self._train_metrics or self._val_metrics: - logger.info( - "Seeded metrics from server: %d train_loss, %d val_loss entries", - len(self._train_metrics), - len(self._val_metrics), - ) - - def _build_metrics_summary(self) -> dict[str, list[dict[str, float | int]]]: - """Build the accumulated metrics payload for inclusion in status_details.""" - return { - "train_loss": list(self._train_metrics), - "val_loss": list(self._val_metrics), - } - - def report_training_start(self, max_steps: int, num_epochs: int) -> None: - """Report that training has started with schedule information.""" - self._reporter.configure_progress_tracking(max_steps, num_epochs) - self._reporter.report_running( - phase="training", - step=0, - max_steps=max_steps, - num_epochs=num_epochs, - metrics=self._build_metrics_summary(), - ) - - def report_train_step( - self, - step: int, - epoch: int, - loss: float, - lr: float | None = None, - grad_norm: float | None = None, - **additional_metrics: Any, - ) -> None: - """Report training step with metrics. - - Args: - step: Training step number - epoch: Current epoch number - loss: Training loss value - lr: Learning rate (optional) - grad_norm: Gradient norm (optional) - **additional_metrics: Additional training metrics to report as current-step - scalars — DPO's preference_loss/rewards_rejected_mean, GRPO's - reward/advantages/kl_penalty, or the shared token counts. These are - not accumulated into the series; see the class docstring. - """ - self._train_metrics.append({"step": step, "epoch": epoch, "value": loss}) - self._reporter.report_running( - phase="training", - step=step, - epoch=epoch, - train_loss=loss, - lr=lr, - grad_norm=grad_norm, - metrics=self._build_metrics_summary(), - **additional_metrics, - ) - - def report_validation( - self, - step: int, - epoch: int, - val_loss: float | None = None, - **additional_metrics: Any, - ) -> None: - """Report validation results. - - Args: - step: Training step number - epoch: Current epoch number - val_loss: Validation loss value, or None for algorithms that do not - produce one. GRPO validates on ``accuracy``/``avg_length`` and - reports no loss at all, so the key is omitted rather than sent as - null and charted as zero. - **additional_metrics: Additional validation metrics to report (e.g., accuracy, - num_valid_samples, or any other validation-specific metrics) - """ - details: dict[str, Any] = {"step": step, "epoch": epoch} - if val_loss is not None: - self._val_metrics.append({"step": step, "epoch": epoch, "value": val_loss}) - details["val_loss"] = val_loss - - self._reporter.report_running( - phase="validation", - metrics=self._build_metrics_summary(), - **details, - **additional_metrics, - ) - - def report_checkpoint_saved(self, step: int, epoch: int, checkpoint_path: str | None = None) -> None: - """Report that a checkpoint was saved.""" - self._reporter.report_running( - phase="checkpoint_saved", - step=step, - epoch=epoch, - checkpoint_path=checkpoint_path, - metrics=self._build_metrics_summary(), - ) - - def close(self) -> None: - """Clean up resources.""" - self._reporter.close() +class TrainingProgressCallback(_BaseTrainingProgressCallback): + """Report NeMo-RL training progress to the Jobs service.""" diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py index 897a14a10b..e6d51b237d 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py @@ -109,18 +109,12 @@ def main(): # and identifiers that should not be dumped to stdout. print(f"Job context loaded (job_id={job_ctx.job_id})") if job_ctx.jobs_url: - # Extract training parameters for progress reporting - max_steps = config.dpo.max_num_steps - num_epochs = config.dpo.max_num_epochs - steps_per_epoch = config.dpo.steps_per_epoch # type: ignore[attr-defined] - extra (undeclared) DPOConfig field, allowed via extra="allow" - log_interval = (config.dpo.val_period // 10) + 1 - - customizer_logger = NemoRLLogger( - steps_per_epoch=steps_per_epoch, + customizer_logger = NemoRLLogger.for_schedule( job_ctx=job_ctx, - log_interval=log_interval, - max_steps=max_steps, - num_epochs=num_epochs, + max_steps=config.dpo.max_num_steps, + num_epochs=config.dpo.max_num_epochs, + val_period=config.dpo.val_period, + steps_per_epoch=config.dpo.steps_per_epoch, # type: ignore[attr-defined] - extra (undeclared) DPOConfig field, allowed via extra="allow" ) # The setup() logger is a composite with a `.loggers` list; guard in case # that internal shape changes. diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py index 0e3ca33ccb..48c561867d 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py @@ -140,16 +140,11 @@ def main() -> None: job_ctx = NMPJobContext.from_env() print(f"Job context loaded (job_id={job_ctx.job_id})") if job_ctx.jobs_url: - max_steps = config.grpo.max_num_steps - num_epochs = config.grpo.max_num_epochs - val_period = config.grpo.val_period or 1 - log_interval = max(val_period // 10, 1) - customizer_logger = NemoRLLogger( - steps_per_epoch=max(max_steps // max(num_epochs, 1), 1), + customizer_logger = NemoRLLogger.for_schedule( job_ctx=job_ctx, - log_interval=log_interval, - max_steps=max_steps, - num_epochs=num_epochs, + max_steps=config.grpo.max_num_steps, + num_epochs=config.grpo.max_num_epochs, + val_period=config.grpo.val_period, ) if hasattr(logger_inst, "loggers"): logger_inst.loggers.append(customizer_logger) diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py index b08bc7d055..5b767f6a4b 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py @@ -15,12 +15,15 @@ from nemo_rl.utils.logger import LoggerInterface from nmp.customization_common.service.context import NMPJobContext -from nmp.customization_common.training.progress import JobsServiceProgressReporter -from nmp.rl.app.constants import SERVICE_NAME from nmp.rl.tasks.training.backends.nemo_rl.callbacks import TrainingProgressCallback +from nmp.rl.tasks.training.progress import JobsServiceProgressReporter _logger = logging.getLogger(__name__) +# How many progress reports to aim for across one validation period. NeMo-RL has no +# notion of a reporting cadence, so it is derived from val_period. +_REPORTS_PER_VAL_PERIOD = 10 + # Metric keys forwarded to Jobs Service, in addition to loss/lr/grad_norm. # # A union across algorithms: selection is by presence, so a DPO run simply has no @@ -75,6 +78,18 @@ def has_metric_value(metric: Any) -> bool: return not math.isnan(float(metric)) +def resolve_log_interval(val_period: int | None) -> int: + """Steps between progress reports, targeting ~10 reports per validation period.""" + return max((val_period or 0) // _REPORTS_PER_VAL_PERIOD, 1) + + +def resolve_steps_per_epoch(max_steps: int, num_epochs: int | None, explicit: int | None = None) -> int: + """Steps per epoch, preferring an explicit value from the algorithm config.""" + if explicit is not None and explicit >= 1: + return explicit + return max(max_steps // max(num_epochs or 1, 1), 1) + + class NemoRLLogger(LoggerInterface): """ NemoRLLogger is a logger implementation that reports training updates to Jobs Service. @@ -121,7 +136,7 @@ def __init__( self._steps_per_epoch = steps_per_epoch # Create the callback for progress reporting - self._reporter = JobsServiceProgressReporter(self._job_ctx, SERVICE_NAME) + self._reporter = JobsServiceProgressReporter(self._job_ctx) self._callback = TrainingProgressCallback(self._reporter) # Track best metrics for monitoring @@ -129,12 +144,46 @@ def __init__( self._best_epoch: int | None = None self._closed = False + # Last train step built but withheld by the log_interval throttle. Flushed on + # close() so the final step is reported even when max_steps is not a multiple + # of log_interval -- otherwise the run's last recorded loss is stale. + self._pending_train_report: dict[str, Any] | None = None + _logger.info( f"Initialized NemoRLLogger with jobs_url={self._job_ctx.jobs_url}, " f"log_interval={log_interval}, max_steps={max_steps}, num_epochs={num_epochs}, " f"steps_per_epoch={steps_per_epoch}" ) + @classmethod + def for_schedule( + cls, + *, + max_steps: int, + num_epochs: int | None, + val_period: int | None, + steps_per_epoch: int | None = None, + job_ctx: NMPJobContext | None = None, + ) -> "NemoRLLogger": + """Build a logger from a NeMo-RL training schedule. + + The DPO and GRPO drivers previously derived ``log_interval`` and + ``steps_per_epoch`` with different formulas for the same intent -- at + ``val_period=100`` one produced 11 and the other 10 -- so the arithmetic + lives here instead of being restated per algorithm. + + Args: + steps_per_epoch: Authoritative value when the algorithm config carries + one (DPO does); otherwise derived from max_steps and num_epochs. + """ + return cls( + steps_per_epoch=resolve_steps_per_epoch(max_steps, num_epochs, steps_per_epoch), + job_ctx=job_ctx, + log_interval=resolve_log_interval(val_period), + max_steps=max_steps, + num_epochs=num_epochs, + ) + def log_metrics( self, metrics: dict[str, Any], @@ -163,21 +212,21 @@ def log_metrics( # once mid-step with the rollout metrics alone (no loss), then again with the # full merged dict. Requiring loss keeps the second, complete one. if prefix == "train" and has_metric_value(metrics.get("loss")): - # Only report at log_interval to reduce output + report = { + "step": step, + "epoch": epoch, + "loss": metrics["loss"], + "lr": metrics.get("lr"), + "grad_norm": metrics.get("grad_norm"), + **self._select_metrics(metrics, _TRAIN_METRIC_KEYS), + } + # Throttled to log_interval to reduce output. A withheld step is held as + # pending rather than dropped, so close() can flush the last one. if step % self._log_interval == 0: - # Extract core metrics - loss = metrics["loss"] - lr = metrics.get("lr") - grad_norm = metrics.get("grad_norm") - - self._callback.report_train_step( - step=step, - epoch=epoch, - loss=loss, - lr=lr, - grad_norm=grad_norm, - **self._select_metrics(metrics, _TRAIN_METRIC_KEYS), - ) + self._callback.report_train_step(**report) + self._pending_train_report = None + else: + self._pending_train_report = report # Handle validation metrics. # @@ -245,13 +294,28 @@ def log_plot(self, figure: Any, step: int, name: str) -> None: return None def close(self) -> None: - """Clean up resources.""" + """Flush any withheld final step, then clean up resources.""" if self._closed: return self._closed = True + self._flush_pending_train_report() _logger.info("NemoRLLogger closing") self._callback.close() + def _flush_pending_train_report(self) -> None: + """Report the last step if the log_interval throttle withheld it. + + Reachable from ``__del__``, so failures must not propagate; the reporter + already swallows and logs transport errors, and this guards the rest. + """ + if self._pending_train_report is None: + return + report, self._pending_train_report = self._pending_train_report, None + try: + self._callback.report_train_step(**report) + except Exception as exc: # pragma: no cover - defensive, shutdown path + _logger.warning(f"Failed to flush final train step: {exc}") + def __del__(self): """Cleanup when the logger is destroyed.""" try: diff --git a/services/rl/src/nmp/rl/tasks/training/progress.py b/services/rl/src/nmp/rl/tasks/training/progress.py new file mode 100644 index 0000000000..f998b1adfa --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/progress.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""Progress reporting for RL training tasks. + +Thin subclass of the shared +:class:`nmp.customization_common.training.progress.JobsServiceProgressReporter` +that bakes in the RL ``SERVICE_NAME`` so callers keep the +``JobsServiceProgressReporter(job_ctx)`` constructor. Mirrors the equivalent +modules in the unsloth and automodel services. +""" + +from nmp.customization_common.service.context import NMPJobContext +from nmp.customization_common.training.progress import ( + JobsServiceProgressReporter as _BaseJobsServiceProgressReporter, +) +from nmp.rl.app.constants import SERVICE_NAME + +__all__ = ["JobsServiceProgressReporter"] + + +class JobsServiceProgressReporter(_BaseJobsServiceProgressReporter): + """RL training progress reporter (binds the RL service name).""" + + def __init__(self, job_ctx: NMPJobContext): + super().__init__(job_ctx, service_name=SERVICE_NAME) diff --git a/services/rl/src/nmp/rl/tasks/training/runner.py b/services/rl/src/nmp/rl/tasks/training/runner.py index e9573dc68d..fe51f41150 100644 --- a/services/rl/src/nmp/rl/tasks/training/runner.py +++ b/services/rl/src/nmp/rl/tasks/training/runner.py @@ -19,8 +19,7 @@ import yaml from nmp.customization_common.service.context import NMPJobContext -from nmp.customization_common.training.progress import JobsServiceProgressReporter -from nmp.rl.app.constants import DEFAULT_TRAINING_RESULT_FILE_NAME, SERVICE_NAME +from nmp.rl.app.constants import DEFAULT_TRAINING_RESULT_FILE_NAME from nmp.rl.app.jobs.training.schemas import ( GPUInfo, TrainingMetrics, @@ -28,6 +27,7 @@ TrainingStepConfig, ) from nmp.rl.app.jobs.training.schemas import TrainingBackend as TrainingBackendEnum +from nmp.rl.tasks.training.progress import JobsServiceProgressReporter from .distributed import DistributedContext from .errors.converter import create_error_details @@ -72,7 +72,7 @@ def __init__(self, backend: TrainingBackend | None = None) -> None: self._job_ctx = NMPJobContext.from_env() self._config = self._load_config(self._job_ctx.config_path) - self._progress = JobsServiceProgressReporter(self._job_ctx, SERVICE_NAME) + self._progress = JobsServiceProgressReporter(self._job_ctx) self._dist_ctx = DistributedContext.from_env(self._get_barrier_dir()) self._backend = backend or self._load_backend(self._config.backend) # workspace_path and output_path are absolute paths from the config diff --git a/services/rl/tests/test_nemo_rl_callbacks.py b/services/rl/tests/test_nemo_rl_callbacks.py index 98bc9ec698..df13896eb1 100644 --- a/services/rl/tests/test_nemo_rl_callbacks.py +++ b/services/rl/tests/test_nemo_rl_callbacks.py @@ -1,178 +1,29 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for the RL TrainingProgressCallback's metric accumulation. +"""Wiring tests for the RL TrainingProgressCallback subclass. -``JobsServiceProgressReporter.report_running`` REPLACES the task's ``status_details`` -blob rather than merging into it, so a report that carries only the current step -erases everything before it. These tests pin the consequence: every report must carry -the full accumulated ``metrics`` payload, in the ``{step, epoch, value}`` shape Studio -reads as ``CustomizationMetricValue[]``. +The behaviour lives in the shared base +(``packages/nmp_customization_common/tests/training/test_callbacks.py``); what is +RL-specific is only which base it inherits and that it stays unstamped. """ from __future__ import annotations -from typing import Any, cast - -import pytest -from nmp.customization_common.training.progress import JobsServiceProgressReporter +from nmp.customization_common.training.callbacks import ( + TrainingProgressCallback as SharedTrainingProgressCallback, +) from nmp.rl.tasks.training.backends.nemo_rl.callbacks import TrainingProgressCallback -class _RecordingReporter: - """Stands in for JobsServiceProgressReporter, capturing each report payload.""" - - def __init__(self, prior: dict[str, list[dict[str, Any]]] | None = None) -> None: - self._prior = prior or {"train_loss": [], "val_loss": []} - self.reports: list[dict[str, Any]] = [] - self.tracking: tuple[int, int] | None = None - self.closed = False - - def fetch_current_metrics(self) -> dict[str, list[dict[str, Any]]]: - return self._prior - - def configure_progress_tracking(self, max_steps: int, num_epochs: int) -> None: - self.tracking = (max_steps, num_epochs) - - def report_running(self, phase: str, **details: Any) -> None: - self.reports.append({"phase": phase, **details}) - - def close(self) -> None: - self.closed = True - - -@pytest.fixture -def reporter() -> _RecordingReporter: - return _RecordingReporter() - - -def _make_callback(reporter: _RecordingReporter) -> TrainingProgressCallback: - """Build the callback over a duck-typed reporter, narrowing the type once here.""" - return TrainingProgressCallback(cast(JobsServiceProgressReporter, reporter)) - - -# --------------------------------------------------------------------------- # -# Accumulation -# --------------------------------------------------------------------------- # - - -def test_train_loss_accumulates_across_steps(reporter: _RecordingReporter) -> None: - callback = _make_callback(reporter) - callback.report_train_step(step=1, epoch=1, loss=0.5) - callback.report_train_step(step=2, epoch=1, loss=0.4) - callback.report_train_step(step=3, epoch=1, loss=0.3) - - # The final report carries the whole curve, not just the last point. - series = reporter.reports[-1]["metrics"]["train_loss"] - assert series == [ - {"step": 1, "epoch": 1, "value": 0.5}, - {"step": 2, "epoch": 1, "value": 0.4}, - {"step": 3, "epoch": 1, "value": 0.3}, - ] - - -def test_every_report_carries_the_full_series(reporter: _RecordingReporter) -> None: - """report_running replaces status_details, so an omission is a data loss.""" - callback = _make_callback(reporter) - callback.report_training_start(max_steps=10, num_epochs=1) - callback.report_train_step(step=1, epoch=1, loss=0.5) - callback.report_validation(step=1, epoch=1, val_loss=0.45) - callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") - - for report in reporter.reports: - assert "metrics" in report, report["phase"] - assert "train_loss" in report["metrics"] - assert "val_loss" in report["metrics"] - - -def test_val_loss_accumulates_separately(reporter: _RecordingReporter) -> None: - callback = _make_callback(reporter) - callback.report_train_step(step=1, epoch=1, loss=0.5) - callback.report_validation(step=1, epoch=1, val_loss=0.45) - callback.report_validation(step=2, epoch=1, val_loss=0.40) - - metrics = reporter.reports[-1]["metrics"] - assert len(metrics["train_loss"]) == 1 - assert metrics["val_loss"] == [ - {"step": 1, "epoch": 1, "value": 0.45}, - {"step": 2, "epoch": 1, "value": 0.40}, - ] - - -def test_series_are_copies_not_live_references(reporter: _RecordingReporter) -> None: - """Each payload must snapshot the series; a shared list would mutate old reports.""" - callback = _make_callback(reporter) - callback.report_train_step(step=1, epoch=1, loss=0.5) - first_payload = reporter.reports[-1]["metrics"]["train_loss"] - callback.report_train_step(step=2, epoch=1, loss=0.4) - - assert len(first_payload) == 1 - - -# --------------------------------------------------------------------------- # -# Resume seeding -# --------------------------------------------------------------------------- # - - -def test_prior_metrics_seed_the_series() -> None: - """A resumed job continues the curve instead of restarting it.""" - prior = { - "train_loss": [{"step": 1, "epoch": 1, "value": 0.9}], - "val_loss": [{"step": 1, "epoch": 1, "value": 0.8}], - } - reporter = _RecordingReporter(prior) - callback = _make_callback(reporter) - callback.report_train_step(step=2, epoch=1, loss=0.5) - - series = reporter.reports[-1]["metrics"]["train_loss"] - assert [entry["step"] for entry in series] == [1, 2] - - -def test_training_start_does_not_erase_seeded_metrics() -> None: - """report_training_start fires before the first step; it must not blank the blob.""" - prior = {"train_loss": [{"step": 1, "epoch": 1, "value": 0.9}], "val_loss": []} - reporter = _RecordingReporter(prior) - callback = _make_callback(reporter) - callback.report_training_start(max_steps=10, num_epochs=1) - - assert reporter.reports[0]["metrics"]["train_loss"] == prior["train_loss"] - - -# --------------------------------------------------------------------------- # -# Optional val_loss (GRPO) -# --------------------------------------------------------------------------- # - - -def test_validation_without_loss_omits_the_key(reporter: _RecordingReporter) -> None: - """GRPO validates on accuracy; a null val_loss would chart as zero.""" - callback = _make_callback(reporter) - callback.report_validation(step=1, epoch=1, val_loss=None, accuracy=0.75) - - report = reporter.reports[-1] - assert "val_loss" not in report - assert report["accuracy"] == 0.75 - assert report["phase"] == "validation" - - -def test_validation_without_loss_leaves_the_series_empty(reporter: _RecordingReporter) -> None: - callback = _make_callback(reporter) - callback.report_validation(step=1, epoch=1, val_loss=None, accuracy=0.75) - - assert reporter.reports[-1]["metrics"]["val_loss"] == [] - - -def test_additional_metrics_ride_along_as_scalars(reporter: _RecordingReporter) -> None: - """The wide RL metric set is current-step only; it must not enter the series.""" - callback = _make_callback(reporter) - callback.report_train_step(step=1, epoch=1, loss=0.5, reward=0.62, kl_penalty=0.008) - - report = reporter.reports[-1] - assert report["reward"] == 0.62 - assert report["kl_penalty"] == 0.008 - assert report["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] +def test_rl_callback_subclasses_the_shared_one() -> None: + """RL used to carry a standalone copy; accumulation fixes must land once.""" + assert issubclass(TrainingProgressCallback, SharedTrainingProgressCallback) -def test_close_delegates_to_the_reporter(reporter: _RecordingReporter) -> None: - _make_callback(reporter).close() +def test_rl_callback_adds_no_backend_field() -> None: + """Stamping `backend` would change RL's status-detail shape on the wire. - assert reporter.closed + unsloth opts in; automodel and RL deliberately do not. + """ + assert TrainingProgressCallback._default_backend is None diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py index 02e9659e70..a441695e64 100644 --- a/services/rl/tests/test_nemo_rl_logger.py +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -40,6 +40,8 @@ class LoggerInterface: # minimal stand-in for the abstract base from nmp.rl.tasks.training.backends.nemo_rl.nemo_rl_logger import ( # noqa: E402 NemoRLLogger, has_metric_value, + resolve_log_interval, + resolve_steps_per_epoch, ) @@ -234,6 +236,109 @@ def test_log_interval_throttles_train_reports(callback: _RecordingCallback) -> N assert [r["step"] for r in callback.train_steps] == [5, 10] +# --------------------------------------------------------------------------- # +# Final-step flush +# --------------------------------------------------------------------------- # + + +def test_close_flushes_the_withheld_final_step(callback: _RecordingCallback) -> None: + """When max_steps is not a multiple of log_interval the last step is throttled out. + + Without a flush the run's last recorded loss is stale — for 23 steps at an + interval of 10 it would be step 20's, and steps 21-23 would never be seen. + """ + logger = _make_logger(log_interval=10) + for step in range(23): + logger.log_metrics(GRPO_TRAIN_METRICS, step=step, prefix="train") + + assert [r["step"] for r in callback.train_steps] == [10, 20] + + logger.close() + + assert [r["step"] for r in callback.train_steps] == [10, 20, 23] + + +def test_close_does_not_duplicate_an_already_reported_step(callback: _RecordingCallback) -> None: + logger = _make_logger(log_interval=10) + for step in range(20): + logger.log_metrics(GRPO_TRAIN_METRICS, step=step, prefix="train") + + logger.close() + + assert [r["step"] for r in callback.train_steps] == [10, 20] + + +def test_flushed_step_carries_the_full_metric_payload(callback: _RecordingCallback) -> None: + logger = _make_logger(log_interval=10) + logger.log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + logger.close() + + flushed = callback.train_steps[-1] + assert flushed["step"] == 1 + assert flushed["reward"] == 0.62 + assert flushed["loss"] == 0.31 + + +def test_double_close_flushes_once(callback: _RecordingCallback) -> None: + logger = _make_logger(log_interval=10) + logger.log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + + logger.close() + logger.close() + + assert len(callback.train_steps) == 1 + + +def test_close_with_nothing_pending_reports_nothing(callback: _RecordingCallback) -> None: + _make_logger().close() + + assert callback.train_steps == [] + + +# --------------------------------------------------------------------------- # +# Schedule resolution — one formula for both drivers +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "val_period,expected", + [ + (100, 10), + (10, 1), + (5, 1), # floors to 0 -> clamped + (1, 1), + (0, 1), + (None, 1), # GRPO's val_period is Optional + ], +) +def test_resolve_log_interval(val_period: int | None, expected: int) -> None: + assert resolve_log_interval(val_period) == expected + + +@pytest.mark.parametrize( + "max_steps,num_epochs,explicit,expected", + [ + (100, 4, None, 25), + (100, None, None, 100), + (100, 0, None, 100), # guard against a zero divisor + (3, 10, None, 1), # floors to 0 -> clamped + (100, 4, 40, 40), # explicit wins (DPO carries steps_per_epoch) + (100, 4, 0, 25), # ...unless it is unusable + ], +) +def test_resolve_steps_per_epoch(max_steps: int, num_epochs: int | None, explicit: int | None, expected: int) -> None: + assert resolve_steps_per_epoch(max_steps, num_epochs, explicit) == expected + + +def test_for_schedule_builds_a_consistent_logger(callback: _RecordingCallback) -> None: + """Both drivers now share this path; DPO used to derive a different interval.""" + logger = NemoRLLogger.for_schedule(max_steps=100, num_epochs=4, val_period=100) + + assert logger._log_interval == 10 + assert logger._steps_per_epoch == 25 + assert logger._max_steps == 100 + + # --------------------------------------------------------------------------- # # Validation — the branch GRPO never reached # --------------------------------------------------------------------------- # diff --git a/services/unsloth/tests/test_callbacks.py b/services/unsloth/tests/test_callbacks.py index bdd09282c7..15a83a6bd0 100644 --- a/services/unsloth/tests/test_callbacks.py +++ b/services/unsloth/tests/test_callbacks.py @@ -65,6 +65,7 @@ def test_report_training_start_delegates(self): step=0, max_steps=500, num_epochs=2, + metrics={"train_loss": [], "val_loss": []}, backend="unsloth", ) From 78d8463d7b18a01990a5b217bb012211256aa520 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 12 Aug 2026 17:41:22 -0400 Subject: [PATCH 04/10] fix(customization): stop every job from erasing its own loss curve The metric series survived mid-training reports as of the previous commit, but not the end of the run. `status_details` is REPLACED by the Jobs service, and the last three writes of any training job come from the runner process, not the training driver: runner report_running("processing_checkpoint") <- no metrics runner report_completed("Training completed") <- no metrics runner report_error(...) <- blanks the whole blob TrainingProgressCallback resends the series on every report it makes, which is why the fix for report_epoch_end/report_checkpoint_saved worked. It cannot help here: the runner is a *different process* from the driver that accumulated the series (backend.execute_training spawns the driver as a subprocess), so it holds nothing to resend. Studio reads the curve straight out of status_details (CustomizationDetailsPanel), so in practice the chart was populated while a job ran and empty the moment it stopped -- and emptiest on the failure path, where the partial curve is worth the most. Preservation therefore has to live below the callback, at update_task, the one choke point every write path shares. When an update does not carry `metrics`, read the stored series back and re-attach it. The server is already the source of truth for the series -- that is how resume seeding works -- so this reuses fetch_current_metrics rather than introducing a second notion of "current". Only `metrics` is carried over. The rest of the blob is deliberately a current-state snapshot (phase, step, lr, ...); merging that would leave a completed task advertising a mid-training step. Costs one GET per update that omits `metrics`, which is the handful the runner makes per job. Per-step training reports always carry their own series and skip the fetch, so the hot path is unchanged. Applies to all three customization services, since they share this reporter. Adds the first test coverage for progress.py. Signed-off-by: Albert Cui --- .../customization_common/training/progress.py | 40 +++- .../tests/training/test_progress.py | 172 ++++++++++++++++++ 2 files changed, 210 insertions(+), 2 deletions(-) create mode 100644 packages/nmp_customization_common/tests/training/test_progress.py diff --git a/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py b/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py index f56221da11..e80c0d0252 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py @@ -8,6 +8,10 @@ training runner; backends subclass it (or instantiate it directly) supplying their own ``service_name`` so the task SDK resolves the right credentials. +Every update REPLACES the task's ``status_details``. The accumulated metric +series is the one cumulative field in that blob, so ``update_task`` carries it +across updates that don't supply their own -- see :meth:`_preserve_metrics`. + For training-specific metrics (loss, validation, checkpoints) see the ``TrainingProgressCallback`` which composes this reporter. """ @@ -53,6 +57,33 @@ def _calculate_percentage_done(self, step: int | None) -> int: # downstream progress consumers expect a bounded percentage. return min(100, int((step / self._max_steps) * 100)) + def _preserve_metrics(self, status_details: dict[str, Any] | None) -> dict[str, Any]: + """Re-attach the stored metric series to an update that doesn't carry one. + + ``status_details`` is REPLACED by the Jobs service, not merged, so any + update omitting ``metrics`` blanks the accumulated loss curve. + ``TrainingProgressCallback`` resends the series on every report it makes, + but the surrounding runner cannot: it reports ``processing_checkpoint``, + completion and failure from a *different process* than the training + driver that accumulated the series, so it holds nothing to resend. Every + job would otherwise end by erasing its own curve -- including, and most + expensively, on the failure path. + + Only ``metrics`` is carried over. The rest of the blob is deliberately a + current-state snapshot (``phase``, ``step``, ``lr``, ...); merging that + would leave a completed task advertising a mid-training step. + + Costs a GET only on updates that omit ``metrics``, which is the handful + the runner makes per job -- never the per-step training reports. + """ + details = dict(status_details or {}) + if "metrics" in details: + return details + stored = self.fetch_current_metrics() + if any(stored.values()): + details["metrics"] = stored + return details + def update_task( self, status: str = "active", @@ -65,6 +96,8 @@ def update_task( if not self._is_main_rank: return + details = self._preserve_metrics(status_details) + try: jobs = client_from_platform(self._sdk, JobsClient) jobs.update_job_step_task( @@ -74,7 +107,7 @@ def update_task( step=self._job_ctx.step, body=PlatformJobTaskUpdate( status=PlatformJobStatus(status), - status_details=status_details or {}, + status_details=details, error_details=error_details or {}, ), ) @@ -99,7 +132,10 @@ def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: "val_loss": metrics.get("val_loss", []), } except Exception as e: - logger.info(f"No prior metrics to seed (expected on first run): {e}") + # Expected on a first run, where the task has no stored details yet. + # Serves both resume seeding and update_task's metric preservation, + # so the message stays neutral about which caller hit it. + logger.info(f"No stored metrics available: {e}") return {"train_loss": [], "val_loss": []} def report_running(self, phase: str, **details: Any) -> None: diff --git a/packages/nmp_customization_common/tests/training/test_progress.py b/packages/nmp_customization_common/tests/training/test_progress.py new file mode 100644 index 0000000000..7dae541bf8 --- /dev/null +++ b/packages/nmp_customization_common/tests/training/test_progress.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for JobsServiceProgressReporter's status_details handling. + +Focused on metric preservation: the Jobs service REPLACES ``status_details``, and +the runner reports checkpoint processing, completion and failure from a different +process than the training driver that accumulated the loss curve. Without the +carry-over in ``update_task`` every job ends by erasing its own metrics. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from nmp.customization_common.training.progress import JobsServiceProgressReporter + +SERIES: dict[str, list[dict[str, float | int]]] = { + "train_loss": [{"step": 10, "epoch": 1, "value": 0.5}], + "val_loss": [{"step": 10, "epoch": 1, "value": 0.45}], +} +EMPTY: dict[str, list[dict[str, float | int]]] = {"train_loss": [], "val_loss": []} + + +class _JobCtx: + """The four identifiers update_task reads off the job context.""" + + normalized_task = "training" + workspace = "default" + job_id = "job-1" + step = "train" + + +class _Reporter(JobsServiceProgressReporter): + """Reporter with the SDK and job context stubbed out. + + Bypasses ``__init__`` rather than mocking the SDK factory: what is under test + is the status_details logic, and the real constructor calls ``get_task_sdk``, + which wants credentials. Every attribute ``update_task`` touches is set here. + """ + + def __init__(self, stored: dict[str, list[dict[str, float | int]]]) -> None: + self._job_ctx = _JobCtx() # type: ignore[assignment] - duck-typed stand-in + self._sdk = object() # type: ignore[assignment] - never dereferenced; the client is patched + self._is_main_rank = True + self._enabled = True + self._max_steps = 0 + self._num_epochs = 0 + self._stored = stored + self.fetch_calls = 0 + + def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: + self.fetch_calls += 1 + return self._stored + + +class _StubJobsClient: + """Captures task updates instead of issuing them.""" + + def __init__(self, sink: list[dict[str, Any]]) -> None: + self._sink = sink + + def update_job_step_task(self, **kwargs: Any) -> None: + self._sink.append(kwargs) + + +@pytest.fixture +def sent(monkeypatch: pytest.MonkeyPatch) -> list[dict[str, Any]]: + """Capture what ``update_task`` would send, running its real body. + + ``update_task`` swallows exceptions, so anything wrong with the stubs would + surface as an empty capture; the helpers below index into it eagerly so that + reads as a failure rather than a pass. + """ + sink: list[dict[str, Any]] = [] + monkeypatch.setattr( + "nmp.customization_common.training.progress.client_from_platform", + lambda _sdk, _cls: _StubJobsClient(sink), + ) + return sink + + +def _details(sent: list[dict[str, Any]]) -> dict[str, Any]: + assert len(sent) == 1, f"expected exactly one task update, got {len(sent)}" + return dict(sent[0]["body"].status_details or {}) + + +# --------------------------------------------------------------------------- # +# Preservation +# --------------------------------------------------------------------------- # + + +def test_completion_preserves_the_accumulated_series(sent: list[dict[str, Any]]) -> None: + """The last write of a successful job must not blank the loss curve.""" + _Reporter(SERIES).report_completed("Training completed") + + details = _details(sent) + assert details["metrics"] == SERIES + assert details["phase"] == "completed" + + +def test_failure_preserves_the_accumulated_series(sent: list[dict[str, Any]]) -> None: + """A failed run is exactly when the partial curve is most worth keeping.""" + _Reporter(SERIES).report_error("boom") + + assert _details(sent)["metrics"] == SERIES + + +def test_intermediate_phase_preserves_the_accumulated_series(sent: list[dict[str, Any]]) -> None: + """processing_checkpoint fires after the driver exits, before completion.""" + _Reporter(SERIES).report_running("processing_checkpoint") + + details = _details(sent) + assert details["metrics"] == SERIES + assert details["phase"] == "processing_checkpoint" + + +def test_caller_supplied_metrics_win_and_skip_the_fetch(sent: list[dict[str, Any]]) -> None: + """Per-step training reports carry their own series; no round-trip for them.""" + fresher = {"train_loss": [{"step": 20, "epoch": 2, "value": 0.1}], "val_loss": []} + reporter = _Reporter(SERIES) + + reporter.report_running("training", step=20, metrics=fresher) + + assert _details(sent)["metrics"] == fresher + assert reporter.fetch_calls == 0 + + +def test_no_stored_metrics_adds_no_key(sent: list[dict[str, Any]]) -> None: + """Before training starts there is nothing to preserve; don't invent a key.""" + _Reporter(EMPTY).report_running("compiling_config") + + assert "metrics" not in _details(sent) + + +def test_preservation_does_not_resurrect_stale_current_state(sent: list[dict[str, Any]]) -> None: + """Only the cumulative field carries over, not the step/lr snapshot.""" + _Reporter(SERIES).report_completed("Training completed") + + assert set(_details(sent)) == {"message", "phase", "metrics"} + + +def test_error_details_still_ride_along(sent: list[dict[str, Any]]) -> None: + """Preserving metrics must not displace the error payload.""" + _Reporter(SERIES).report_error({"message": "oom", "code": "OOM"}) + + assert sent[0]["body"].error_details == {"message": "oom", "code": "OOM"} + + +# --------------------------------------------------------------------------- # +# Gating +# --------------------------------------------------------------------------- # + + +def test_disabled_reporter_sends_nothing_and_does_not_fetch(sent: list[dict[str, Any]]) -> None: + reporter = _Reporter(SERIES) + reporter._enabled = False + + reporter.report_completed("Training completed") + + assert sent == [] + assert reporter.fetch_calls == 0 + + +def test_non_main_rank_sends_nothing(sent: list[dict[str, Any]]) -> None: + reporter = _Reporter(SERIES) + reporter._is_main_rank = False + + reporter.report_completed("Training completed") + + assert sent == [] From c6341f807333042d18b6d38193fbae06ae9343d0 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 12 Aug 2026 17:43:10 -0400 Subject: [PATCH 05/10] fix(rl): actually invoke the teardown that flushes the final train step The previous commit added a pending-report flush on NemoRLLogger.close() to stop the last training step being dropped by the log_interval throttle. Nothing calls close(). Both drivers append the logger to `logger_inst.loggers` and never tear it down. nemo_rl.utils.logger.Logger has no close() at all -- its only teardown hook is finish(), dispatched as `getattr(logger, "finish", None)`, which skipped us because NemoRLLogger did not define one. And grpo_train/dpo_train never call finish() either; the only caller upstream is the single-controller path, which these drivers do not use. So the flush ran only from __del__, at GC or interpreter shutdown, where both _flush_pending_train_report and the reporter's update_task swallow exceptions. On a clean return refcounting probably got there. On SIGTERM, an unhandled exception or a cancelled job it did not -- which is the case the flush exists for. Two hooks, because neither alone is sufficient: finish() aliases close() under the name the composite dispatches, so the flush happens even if a driver forgets. drivers call close() explicitly from a finally, because dpo_train and grpo_train never trigger the composite's finish() at all. DPO's driver had no try/finally around dpo_train; it has one now. GRPO's runs before the existing environment teardown, which is slow and can itself raise. Tested at the seam that broke: one test performs the composite's exact `getattr(logger, "finish", None)` lookup, so renaming the method fails loudly. The driver calls are asserted against the AST -- the drivers cannot be imported outside the training image -- with the detector's own negative cases pinned, since a tripwire that cannot trip is worse than none. Signed-off-by: Albert Cui --- .../training/backends/nemo_rl/dpo_driver.py | 30 +++++--- .../training/backends/nemo_rl/grpo_driver.py | 9 +++ .../backends/nemo_rl/nemo_rl_logger.py | 13 ++++ services/rl/tests/test_nemo_rl_drivers.py | 68 +++++++++++++++++++ services/rl/tests/test_nemo_rl_logger.py | 39 +++++++++++ 5 files changed, 148 insertions(+), 11 deletions(-) create mode 100644 services/rl/tests/test_nemo_rl_drivers.py diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py index e6d51b237d..53e093ad56 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py @@ -108,6 +108,7 @@ def main(): # Log only the non-sensitive job id; the full context carries service URLs # and identifiers that should not be dumped to stdout. print(f"Job context loaded (job_id={job_ctx.job_id})") + customizer_logger: NemoRLLogger | None = None if job_ctx.jobs_url: customizer_logger = NemoRLLogger.for_schedule( job_ctx=job_ctx, @@ -125,17 +126,24 @@ def main(): logger.log_hyperparams(config.model_dump()) - dpo_train( - policy, - train_dataloader, - val_dataloader, - tokenizer, - loss_fn, - master_config, - logger, - checkpointer, - dpo_save_state, - ) + try: + dpo_train( + policy, + train_dataloader, + val_dataloader, + tokenizer, + loss_fn, + master_config, + logger, + checkpointer, + dpo_save_state, + ) + finally: + # Flushes the final training step. NeMo-RL never closes the loggers it is + # handed, so without this the only fallback is NemoRLLogger.__del__ at + # interpreter shutdown, which does not run at all on an abnormal exit. + if customizer_logger is not None: + customizer_logger.close() if __name__ == "__main__": diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py index 48c561867d..5fdc16199a 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py @@ -139,6 +139,7 @@ def main() -> None: job_ctx = NMPJobContext.from_env() print(f"Job context loaded (job_id={job_ctx.job_id})") + customizer_logger: NemoRLLogger | None = None if job_ctx.jobs_url: customizer_logger = NemoRLLogger.for_schedule( job_ctx=job_ctx, @@ -167,6 +168,14 @@ def main() -> None: master_config, ) finally: + # Before the env teardown below, which can be slow and can itself raise: + # this flushes the final training step, and it is the only deterministic + # chance to do so. NeMo-RL never closes the loggers it is handed, so the + # fallback is NemoRLLogger.__del__ at interpreter shutdown, which does not + # run at all on an abnormal exit. + if customizer_logger is not None: + customizer_logger.close() + for task_name, env in task_to_env.items(): try: import ray diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py index 5b767f6a4b..38963526df 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py @@ -293,6 +293,19 @@ def log_plot(self, figure: Any, step: int, name: str) -> None: """ return None + def finish(self) -> None: + """Alias for :meth:`close` under the name NeMo-RL's composite fans out. + + ``nemo_rl.utils.logger.Logger`` has no ``close()`` at all; its only + teardown hook is ``finish()``, dispatched via + ``getattr(logger, "finish", None)``. Without this method the composite + silently skips us and the withheld final step is never flushed. The + drivers also call ``close()`` directly, because ``grpo_train``/ + ``dpo_train`` never invoke ``finish()`` either -- only the + single-controller path does. + """ + self.close() + def close(self) -> None: """Flush any withheld final step, then clean up resources.""" if self._closed: diff --git a/services/rl/tests/test_nemo_rl_drivers.py b/services/rl/tests/test_nemo_rl_drivers.py new file mode 100644 index 0000000000..4fc2d28882 --- /dev/null +++ b/services/rl/tests/test_nemo_rl_drivers.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Source-level checks that the drivers tear the progress logger down. + +These are tripwires, not behaviour tests. The drivers cannot be imported outside +the training image -- they pull in nemo_rl, ray and omegaconf at module scope -- +so the wiring is asserted against the AST instead. + +It is worth asserting at all because the failure is silent: NeMo-RL never closes +the loggers it is handed, so if these calls are dropped the final training step +stops being reported and every unit test still passes. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +DRIVERS = Path(__file__).resolve().parents[1] / "src/nmp/rl/tasks/training/backends/nemo_rl" + + +def _closes_logger_in_finally(source: str) -> bool: + """Whether some `try/finally` closes `customizer_logger` in its finally body.""" + for node in ast.walk(ast.parse(source)): + if not isinstance(node, ast.Try): + continue + for stmt in node.finalbody: + for inner in ast.walk(stmt): + if ( + isinstance(inner, ast.Call) + and isinstance(inner.func, ast.Attribute) + and inner.func.attr == "close" + and isinstance(inner.func.value, ast.Name) + and inner.func.value.id == "customizer_logger" + ): + return True + return False + + +@pytest.mark.parametrize( + "source,expected", + [ + ("try:\n train()\nfinally:\n customizer_logger.close()\n", True), + # Guarded call — how the drivers actually write it. + ("try:\n train()\nfinally:\n if customizer_logger:\n customizer_logger.close()\n", True), + # Present, but not on the abnormal-exit path. + ("try:\n train()\nfinally:\n pass\ncustomizer_logger.close()\n", False), + ("try:\n train()\nfinally:\n other_logger.close()\n", False), + ("try:\n train()\nfinally:\n customizer_logger.flush()\n", False), + ], +) +def test_detector_discriminates(source: str, expected: bool) -> None: + """The tripwire is only worth having if it can actually trip.""" + assert _closes_logger_in_finally(source) is expected + + +@pytest.mark.parametrize("driver", ["grpo_driver.py", "dpo_driver.py"]) +def test_driver_closes_the_progress_logger_in_a_finally(driver: str) -> None: + """`finally`, not the happy path: an aborted run is when the flush matters.""" + source = (DRIVERS / driver).read_text() + + assert _closes_logger_in_finally(source), ( + f"{driver} must close customizer_logger from a finally block; " + "NeMo-RL does not close loggers, and __del__ does not run on abnormal exit" + ) diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py index a441695e64..b378f17a24 100644 --- a/services/rl/tests/test_nemo_rl_logger.py +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -289,6 +289,45 @@ def test_double_close_flushes_once(callback: _RecordingCallback) -> None: assert len(callback.train_steps) == 1 +def test_finish_flushes_like_close(callback: _RecordingCallback) -> None: + """`finish` is the name NeMo-RL's composite Logger actually dispatches. + + nemo_rl.utils.logger.Logger has no close(); its teardown fan-out is + `getattr(logger, "finish", None)`. Without this alias the composite skips us + entirely and the withheld final step is never flushed. + """ + logger = _make_logger(log_interval=10) + logger.log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + + logger.finish() + + assert [r["step"] for r in callback.train_steps] == [1] + assert callback.closed + + +def test_finish_is_reachable_through_the_composite_dispatch(callback: _RecordingCallback) -> None: + """Mirrors Logger.finish()'s exact lookup, so a rename here fails loudly.""" + logger = _make_logger(log_interval=10) + logger.log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + + finish = getattr(logger, "finish", None) + assert callable(finish) + finish() + + assert [r["step"] for r in callback.train_steps] == [1] + + +def test_finish_then_close_flushes_once(callback: _RecordingCallback) -> None: + """Both the composite and the driver may call in; the step reports once.""" + logger = _make_logger(log_interval=10) + logger.log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + + logger.finish() + logger.close() + + assert len(callback.train_steps) == 1 + + def test_close_with_nothing_pending_reports_nothing(callback: _RecordingCallback) -> None: _make_logger().close() From 5366a3ac72d8256d546faed7e7b2f0f92910117e Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 12 Aug 2026 17:45:05 -0400 Subject: [PATCH 06/10] fix(rl): stop double-counting the training step `log_metrics` opened with `step = step + 1 # ...we start counting from 1`, but both callers already count from 1: grpo.py and dpo.py log `total_steps + 1`, where total_steps is 0-based and incremented *after* the log. The logger added a second increment on top. A 23-step run therefore recorded steps 2..24 against max_steps=23. Until the previous commit that only skewed a `step` field; now that the series is the x-axis of a rendered loss curve, the whole curve sat one step right of the truth and percentage_done saturated a step early. It also interacted with the log_interval throttle. `step % log_interval == 0` was evaluated on the inflated step, so reports landed on true steps 9, 19, 29 -- and the final step was withheld even when max_steps *was* a multiple of the interval (at max_steps=20, interval=10 the run's last reported loss was step 19's). The pending-report flush was covering for this; it now handles only the case it was written for. Epoch derivation is fixed by the same change, since it was reading the same inflated step and flipping an epoch early at the boundary. It now clamps at zero: step 0 does arrive, from the validate-at-start path both algorithms run before training, and it belongs to epoch 1 rather than epoch 0. Rides with this branch because it is the same wire values, and this branch already changes DPO's reporting cadence -- correcting the labels separately would mean reviewing that blast radius twice. The tests drove on 0-indexed steps, which is why this survived them. They now generate the sequence a real N-step run produces, via a helper that says so, and pin both ends of the range plus the epoch boundaries. Signed-off-by: Albert Cui --- .../backends/nemo_rl/nemo_rl_logger.py | 14 +++- services/rl/tests/test_nemo_rl_logger.py | 83 +++++++++++++++---- 2 files changed, 78 insertions(+), 19 deletions(-) diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py index 38963526df..cba9ed3dc0 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py @@ -201,10 +201,16 @@ def log_metrics( step_metric: Optional step metric name (ignored in this implementation) step_finished: Whether the step is finished (part of NeMo-RL's LoggerInterface; ignored here) """ - step = step + 1 # Increment step since we start counting from 1 - - # Calculate epoch from step (epochs start from 1) - epoch = ((step - 1) // self._steps_per_epoch) + 1 + # `step` arrives 1-indexed and is used as-is. Both callers pass + # `total_steps + 1`, where total_steps is 0-based and incremented *after* + # logging (nemo_rl/algorithms/grpo.py, .../dpo.py), so it is already the + # 1-indexed step number. Incrementing again put the last step of an + # N-step run at N+1 and shifted the whole series one to the right of the + # axis Studio draws it against. + # + # Step 0 does arrive, from the validate-at-start path only; it belongs to + # epoch 1, hence the clamp rather than a bare `step - 1`. + epoch = (max(step - 1, 0) // self._steps_per_epoch) + 1 # Handle training loss. # diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py index b378f17a24..ae8af3ae29 100644 --- a/services/rl/tests/test_nemo_rl_logger.py +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -84,6 +84,17 @@ def _make_logger(**kwargs: Any) -> NemoRLLogger: return NemoRLLogger(**params) +def _driver_steps(max_steps: int) -> range: + """The step sequence an N-step run actually produces. + + grpo.py and dpo.py both log `total_steps + 1` with total_steps 0-based and + incremented after the log, so an N-step run emits 1..N -- not 0..N-1. Tests + that use range(N) directly would validate the throttle against a convention + no caller uses. + """ + return range(1, max_steps + 1) + + class _Histogram: """Stand-in for nemo_rl's wandb Histogram — non-numeric, and NaN-hostile.""" @@ -229,10 +240,9 @@ def test_dpo_train_metrics_still_forwarded(callback: _RecordingCallback) -> None def test_log_interval_throttles_train_reports(callback: _RecordingCallback) -> None: logger = _make_logger(log_interval=5) - for step in range(10): + for step in _driver_steps(10): logger.log_metrics(GRPO_TRAIN_METRICS, step=step, prefix="train") - # log_metrics increments step by 1, so steps 5 and 10 report. assert [r["step"] for r in callback.train_steps] == [5, 10] @@ -248,7 +258,7 @@ def test_close_flushes_the_withheld_final_step(callback: _RecordingCallback) -> interval of 10 it would be step 20's, and steps 21-23 would never be seen. """ logger = _make_logger(log_interval=10) - for step in range(23): + for step in _driver_steps(23): logger.log_metrics(GRPO_TRAIN_METRICS, step=step, prefix="train") assert [r["step"] for r in callback.train_steps] == [10, 20] @@ -260,7 +270,7 @@ def test_close_flushes_the_withheld_final_step(callback: _RecordingCallback) -> def test_close_does_not_duplicate_an_already_reported_step(callback: _RecordingCallback) -> None: logger = _make_logger(log_interval=10) - for step in range(20): + for step in _driver_steps(20): logger.log_metrics(GRPO_TRAIN_METRICS, step=step, prefix="train") logger.close() @@ -270,7 +280,7 @@ def test_close_does_not_duplicate_an_already_reported_step(callback: _RecordingC def test_flushed_step_carries_the_full_metric_payload(callback: _RecordingCallback) -> None: logger = _make_logger(log_interval=10) - logger.log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + logger.log_metrics(GRPO_TRAIN_METRICS, step=1, prefix="train") logger.close() flushed = callback.train_steps[-1] @@ -281,7 +291,7 @@ def test_flushed_step_carries_the_full_metric_payload(callback: _RecordingCallba def test_double_close_flushes_once(callback: _RecordingCallback) -> None: logger = _make_logger(log_interval=10) - logger.log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + logger.log_metrics(GRPO_TRAIN_METRICS, step=1, prefix="train") logger.close() logger.close() @@ -297,7 +307,7 @@ def test_finish_flushes_like_close(callback: _RecordingCallback) -> None: entirely and the withheld final step is never flushed. """ logger = _make_logger(log_interval=10) - logger.log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + logger.log_metrics(GRPO_TRAIN_METRICS, step=1, prefix="train") logger.finish() @@ -308,7 +318,7 @@ def test_finish_flushes_like_close(callback: _RecordingCallback) -> None: def test_finish_is_reachable_through_the_composite_dispatch(callback: _RecordingCallback) -> None: """Mirrors Logger.finish()'s exact lookup, so a rename here fails loudly.""" logger = _make_logger(log_interval=10) - logger.log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + logger.log_metrics(GRPO_TRAIN_METRICS, step=1, prefix="train") finish = getattr(logger, "finish", None) assert callable(finish) @@ -320,7 +330,7 @@ def test_finish_is_reachable_through_the_composite_dispatch(callback: _Recording def test_finish_then_close_flushes_once(callback: _RecordingCallback) -> None: """Both the composite and the driver may call in; the step reports once.""" logger = _make_logger(log_interval=10) - logger.log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + logger.log_metrics(GRPO_TRAIN_METRICS, step=1, prefix="train") logger.finish() logger.close() @@ -378,6 +388,49 @@ def test_for_schedule_builds_a_consistent_logger(callback: _RecordingCallback) - assert logger._max_steps == 100 +# --------------------------------------------------------------------------- # +# Step and epoch arithmetic +# --------------------------------------------------------------------------- # + + +def test_step_is_reported_as_the_caller_numbered_it(callback: _RecordingCallback) -> None: + """The caller's step is already 1-indexed; re-incrementing shifted the curve.""" + logger = _make_logger(max_steps=23) + for step in _driver_steps(23): + logger.log_metrics(GRPO_TRAIN_METRICS, step=step, prefix="train") + + reported = [r["step"] for r in callback.train_steps] + assert reported[0] == 1, "an N-step run starts at 1" + assert reported[-1] == 23, "...and ends at N, not N+1" + + +@pytest.mark.parametrize( + "step,expected_epoch", + [ + (0, 1), # validate-at-start, before any training + (1, 1), + (10, 1), # last step of epoch 1 at steps_per_epoch=10 + (11, 2), # first of epoch 2 + (20, 2), + (21, 3), + ], +) +def test_epoch_boundaries(callback: _RecordingCallback, step: int, expected_epoch: int) -> None: + """Epoch flips on the step after a full epoch, not the last step of one.""" + _make_logger().log_metrics({"loss": 0.5}, step=step, prefix="train") + + assert callback.train_steps[0]["epoch"] == expected_epoch + + +def test_validate_at_start_reports_step_zero(callback: _RecordingCallback) -> None: + """Both algorithms run an optional validation pass at step 0 before training.""" + _make_logger().log_metrics(GRPO_VALIDATION_METRICS, step=0, prefix="validation") + + reported = callback.validations[0] + assert reported["step"] == 0 + assert reported["epoch"] == 1 + + # --------------------------------------------------------------------------- # # Validation — the branch GRPO never reached # --------------------------------------------------------------------------- # @@ -388,7 +441,7 @@ def test_grpo_validation_is_reported_without_a_loss(callback: _RecordingCallback Gating this branch on `loss` silently dropped every GRPO validation report. """ - _make_logger().log_metrics(GRPO_VALIDATION_METRICS, step=9, prefix="validation") + _make_logger().log_metrics(GRPO_VALIDATION_METRICS, step=10, prefix="validation") assert len(callback.validations) == 1 reported = callback.validations[0] @@ -400,7 +453,7 @@ def test_grpo_validation_is_reported_without_a_loss(callback: _RecordingCallback def test_dpo_validation_still_reports_loss(callback: _RecordingCallback) -> None: - _make_logger().log_metrics({"loss": 0.25, "num_valid_samples": 8}, step=9, prefix="validation") + _make_logger().log_metrics({"loss": 0.25, "num_valid_samples": 8}, step=10, prefix="validation") reported = callback.validations[0] assert reported["val_loss"] == 0.25 @@ -417,9 +470,9 @@ def test_validation_with_nothing_usable_is_not_reported(callback: _RecordingCall def test_best_validation_loss_tracks_minimum(callback: _RecordingCallback) -> None: logger = _make_logger() - logger.log_metrics({"loss": 0.5}, step=9, prefix="validation") - logger.log_metrics({"loss": 0.2}, step=19, prefix="validation") - logger.log_metrics({"loss": 0.7}, step=29, prefix="validation") + logger.log_metrics({"loss": 0.5}, step=10, prefix="validation") + logger.log_metrics({"loss": 0.2}, step=20, prefix="validation") + logger.log_metrics({"loss": 0.7}, step=30, prefix="validation") assert logger._best_metric_value == 0.2 assert logger._best_epoch == 2 @@ -428,7 +481,7 @@ def test_best_validation_loss_tracks_minimum(callback: _RecordingCallback) -> No def test_grpo_validation_leaves_best_loss_untouched(callback: _RecordingCallback) -> None: """No loss means no best-loss update — and no crash comparing None.""" logger = _make_logger() - logger.log_metrics(GRPO_VALIDATION_METRICS, step=9, prefix="validation") + logger.log_metrics(GRPO_VALIDATION_METRICS, step=10, prefix="validation") assert math.isinf(logger._best_metric_value) assert logger._best_epoch is None From 11ee9ec71d8525f217917290ee6f16c82cdc94f1 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 12 Aug 2026 17:46:36 -0400 Subject: [PATCH 07/10] chore(rl): tighten four loose ends from review Independent one-liners, grouped so they are easy to drop; none change behaviour on a path exercised today. callbacks: splat `**additional_metrics` first in report_train_step, so a backend metric named `metrics` or `train_loss` cannot silently replace the accumulated series or the step's own loss. report_validation already had this order; the two now agree. Every other colliding name is a real parameter and so already errors at the call site. nemo_rl_logger: add `rewards_chosen_mean` alongside `rewards_rejected_mean`. Forwarding one half of DPO's reward pair makes it hard to read. test_nemo_rl_logger: give the nemo_rl module stubs a real `__spec__`. find_spec consults sys.modules before the finders and raises on a `__spec__` of None, so the bare ModuleType turned any later `find_spec("nemo_rl")` in the same session into a ValueError. The stub is installed at import time and never torn down -- exactly the shape of leak the chat_template test on this branch had to be rewritten around, so it should not be left as a trap for the next one. test_grpo_config: `lambda *_, **__` for the resolve_chat_template patch. The call site passes keywords today; the keyword-only stub breaks silently if that ever changes. Signed-off-by: Albert Cui --- .../training/callbacks.py | 5 ++- .../tests/training/test_callbacks.py | 18 ++++++++++ .../backends/nemo_rl/nemo_rl_logger.py | 3 +- services/rl/tests/test_grpo_config.py | 2 +- services/rl/tests/test_nemo_rl_logger.py | 33 ++++++++++++++++--- 5 files changed, 53 insertions(+), 8 deletions(-) diff --git a/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py b/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py index 2b57105d63..9e5b6c88c6 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py @@ -96,14 +96,17 @@ def report_train_step( accumulated into the series; see the module docstring. """ self._train_metrics.append({"step": step, "epoch": epoch, "value": loss}) + # `**additional_metrics` is splatted first, matching report_validation, so a + # backend metric cannot shadow the accumulated series or the step's own loss. + # `step`/`epoch`/`lr`/`grad_norm` are named parameters and so already safe. details: dict[str, object] = { + **additional_metrics, "step": step, "epoch": epoch, "train_loss": loss, "lr": lr, "grad_norm": grad_norm, "metrics": self._build_metrics_summary(), - **additional_metrics, } resolved = self._resolve_backend(backend) if resolved is not None: diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py index abe5a149d0..b0873381e4 100644 --- a/packages/nmp_customization_common/tests/training/test_callbacks.py +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -135,6 +135,24 @@ def test_additional_validation_metrics_ride_along(reporter: _RecordingReporter) assert reporter.reports[-1]["accuracy"] == 0.9 +def test_additional_metrics_cannot_shadow_the_series(reporter: _RecordingReporter) -> None: + """`metrics` is not a parameter, so only splat order stops a silent override. + + `step`/`epoch`/`lr`/`grad_norm`/`backend` are named parameters -- passing one + is a TypeError at the call site. `metrics` and `train_loss` would just win. + """ + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, metrics="clobbered") + + report = reporter.reports[-1] + assert report["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] + + +def test_additional_metrics_cannot_shadow_the_step_loss(reporter: _RecordingReporter) -> None: + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, train_loss="clobbered") + + assert reporter.reports[-1]["train_loss"] == 0.5 + + def test_additional_metrics_do_not_collide_with_backend_stamping( reporter: _RecordingReporter, ) -> None: diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py index cba9ed3dc0..1ecec2fbc5 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py @@ -36,8 +36,9 @@ "num_valid_samples", "global_valid_seqs", "global_valid_toks", - # DPO + # DPO — both halves of the reward pair; either alone is hard to read. "preference_loss", + "rewards_chosen_mean", "rewards_rejected_mean", # GRPO: reward and advantages — the signal that says whether RL is working "reward", diff --git a/services/rl/tests/test_grpo_config.py b/services/rl/tests/test_grpo_config.py index 017b5da51c..2b91440fb3 100644 --- a/services/rl/tests/test_grpo_config.py +++ b/services/rl/tests/test_grpo_config.py @@ -267,7 +267,7 @@ def test_tokenizer_omits_chat_template_when_none( # sys.modules, whose truthy `.chat_template` silently invalidates the premise. monkeypatch.setattr( "nmp.rl.tasks.training.backends.nemo_rl.grpo_config.resolve_chat_template", - lambda **_: None, + lambda *_, **__: None, ) step, _ = _prepared_step(tmp_path) tokenizer = compile_grpo_config(step, job_ctx)["policy"]["tokenizer"] diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py index ae8af3ae29..62113418c4 100644 --- a/services/rl/tests/test_nemo_rl_logger.py +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -11,6 +11,7 @@ from __future__ import annotations +import importlib.machinery import importlib.util import math import sys @@ -24,16 +25,27 @@ # import the module under test; when the real package IS present (the in-image smoke # run) this is skipped and the genuine base class is used. if importlib.util.find_spec("nemo_rl") is None: # pragma: no cover - env dependent - _nemo_rl = types.ModuleType("nemo_rl") - _utils = types.ModuleType("nemo_rl.utils") - _logger_mod = types.ModuleType("nemo_rl.utils.logger") class LoggerInterface: # minimal stand-in for the abstract base pass + def _stub(name: str) -> types.ModuleType: + """Build a stub module that survives a later importlib.util.find_spec. + + A bare ModuleType has ``__spec__ = None``, and find_spec consults + sys.modules first -- so leaving it unset makes a later + ``find_spec("nemo_rl")`` raise ValueError rather than return None. The + stub outlives this module (nothing tears it down), so it must not booby + trap whatever runs next in the session. + """ + module = types.ModuleType(name) + module.__spec__ = importlib.machinery.ModuleSpec(name, loader=None) + return module + + _logger_mod = _stub("nemo_rl.utils.logger") setattr(_logger_mod, "LoggerInterface", LoggerInterface) - sys.modules.setdefault("nemo_rl", _nemo_rl) - sys.modules.setdefault("nemo_rl.utils", _utils) + sys.modules.setdefault("nemo_rl", _stub("nemo_rl")) + sys.modules.setdefault("nemo_rl.utils", _stub("nemo_rl.utils")) sys.modules.setdefault("nemo_rl.utils.logger", _logger_mod) from nmp.rl.tasks.training.backends.nemo_rl import nemo_rl_logger # noqa: E402 @@ -176,6 +188,17 @@ def test_has_metric_value_does_not_raise_on_any_grpo_metric() -> None: assert isinstance(has_metric_value(value), bool), key +def test_module_stub_does_not_break_find_spec() -> None: + """The stub installed at import time outlives this module; it must be inert. + + find_spec consults sys.modules first and raises on a `__spec__` of None, so a + bare ModuleType here would turn an unrelated later `find_spec("nemo_rl")` + into a ValueError -- the same kind of cross-suite leak this file's sibling + test_grpo_config had to be rewritten around. + """ + assert importlib.util.find_spec("nemo_rl") is not None + + def test_has_metric_value_accepts_numpy_scalars() -> None: np = pytest.importorskip("numpy") assert has_metric_value(np.float32(0.5)) is True From a933ccda797a919404973c685821de291d216011 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 11:20:12 -0400 Subject: [PATCH 08/10] feat(customization): accumulate a time series for every reported metric Only train_loss and val_loss were series; every other metric rode as a current-step scalar that the next update overwrote. For GRPO that left exactly one plottable curve -- and not a useful one, since val_loss is permanently empty for an algorithm with no validation loss. Reward, accuracy, KL and truncation rate, the metrics you actually watch to decide whether a run is working, existed only as "latest value". Now every numeric metric a backend reports accumulates into its own series in the same {step, epoch, value} shape Studio already renders. A GRPO run produces 22 of them. The current-step scalars stay on the blob alongside, so consumers can read either the curve or the latest value. Series are namespaced by phase: train_ / val_. The prefix is load-bearing, not cosmetic -- GRPO reports truncation_rate in both its train and validation dicts and DPO reports accuracy in both, so unprefixed names would interleave two different quantities into one curve. train_loss and val_loss keep their bare names, so the existing Studio chart is unaffected. lr and grad_norm accumulate too; they are curves people read, and they were only excluded because they happen to be named parameters rather than **additional_metrics. fetch_current_metrics had to stop hardcoding the two names, or a resumed job would silently restart all 20 other curves from empty. It now returns whatever list-valued series are stored. The numeric guard moves into the shared module as is_chartable(), and RL's has_metric_value delegates to it. They were about to be two copies of the same rule, and a metric the logger forwards must be one the callback can chart -- letting those drift is how a Histogram ends up in a series. Size, measured for GRPO's 22 series rather than estimated. The cost scales with reports, not steps, and backends throttle reporting: 500 steps, log_interval 10 -> 42 KB final blob, 1.1 MB uploaded 500 steps, log_interval 1 -> 413 KB final blob, 101.3 MB uploaded Accepted for batch training jobs. A backend reporting every step of a long run pays quadratically; if that becomes a real configuration the fix is delta appends in the transport, not trimming the series here. Signed-off-by: Albert Cui --- .../training/callbacks.py | 128 ++++++++++++++---- .../customization_common/training/progress.py | 16 ++- .../tests/training/test_callbacks.py | 82 +++++++++-- .../tasks/training/backends/test_callbacks.py | 4 +- .../backends/nemo_rl/nemo_rl_logger.py | 12 +- 5 files changed, 193 insertions(+), 49 deletions(-) diff --git a/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py b/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py index 9e5b6c88c6..608032a434 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py @@ -4,10 +4,9 @@ """Training progress callback shared by the customization backends. Composes a :class:`nmp.customization_common.training.progress.JobsServiceProgressReporter` -and provides training-specific methods. Metric accumulation: ``train_loss`` and -``val_loss`` are accumulated as time-series lists and included in EVERY -``status_details`` update under a ``metrics`` key, enabling loss-curve -reconstruction from job status. +and provides training-specific methods. Every numeric metric a backend reports is +accumulated as a time series and included in EVERY ``status_details`` update under +a ``metrics`` key, so any of them can be charted from job status alone. Every update matters because ``report_running`` REPLACES the task's ``status_details`` blob rather than merging into it. A report that omits @@ -16,9 +15,35 @@ is gone. Checkpoint and epoch-end reports fire mid-training, so they carry the payload too. -Only these two series accumulate. Anything passed as ``**additional_metrics`` -rides along as a current-step scalar: the full series set is resent on every -update, so the payload grows with series count times step count. +Series naming +------------- +Series are namespaced by the phase that produced them: ``train_`` and +``val_``, which is what the long-standing ``train_loss``/``val_loss`` pair +already did. The prefix is load-bearing rather than cosmetic -- GRPO reports +``truncation_rate`` in both its train and validation dicts, and DPO reports +``accuracy`` in both, so unprefixed names would interleave two different +quantities into one series. + +``train_loss`` and ``val_loss`` keep those exact names, so existing consumers +(the Studio loss chart) are unaffected. + +Payload size +------------ +Every series is resent in full on every update, so the stored blob grows as +``series x reports`` and total upload as the square of it. The driver of that +cost is the number of *reports*, not training steps -- backends throttle +reporting, so a 500-step GRPO run at ``log_interval=10`` accumulates 50 points +per series, not 500. + +Measured, for GRPO's ~22 series: + + 500 steps, log_interval 10 -> 42 KB final blob, 1.1 MB uploaded + 500 steps, log_interval 1 -> 413 KB final blob, 101.3 MB uploaded + +Deliberately accepted for batch training jobs. It does mean a backend that +reports every step of a long run pays quadratically, so if that becomes a real +configuration the transport should move to delta appends rather than the series +being trimmed here. Backends subclass this and set :attr:`_default_backend`: unsloth stamps a ``backend`` field on each report (``"unsloth"``); automodel and NeMo-RL leave it @@ -27,12 +52,34 @@ """ import logging -from typing import ClassVar +import math +import numbers +from typing import Any, ClassVar, cast from nmp.customization_common.training.progress import JobsServiceProgressReporter logger = logging.getLogger(__name__) +#: Series that keep their bare name instead of taking a phase prefix, because +#: they predate the prefixing scheme and are read by name downstream. +_UNPREFIXED = frozenset({"train_loss", "val_loss"}) + + +def is_chartable(value: Any) -> bool: + """Whether ``value`` is a finite scalar that can enter a metric series. + + Backends hand us whatever their framework produced, which is not always a + number: NeMo-RL's metric dicts interleave ``Histogram`` objects, tables and + nested dicts with the scalars, and ``math.isnan`` raises ``TypeError`` on all + of those rather than returning False. + + ``bool`` is rejected despite being an ``int`` subclass: no metric here is a + flag, and silently charting one as 0/1 is worse than dropping it. + """ + if isinstance(value, bool) or not isinstance(value, numbers.Real): + return False + return not math.isnan(float(value)) + class TrainingProgressCallback: """Report training progress to the Jobs service.""" @@ -44,25 +91,46 @@ class TrainingProgressCallback: def __init__(self, reporter: JobsServiceProgressReporter): self._reporter = reporter - prior = reporter.fetch_current_metrics() - self._train_metrics: list[dict[str, float | int]] = prior.get("train_loss", []) - self._val_metrics: list[dict[str, float | int]] = prior.get("val_loss", []) - if self._train_metrics or self._val_metrics: + #: series name -> [{step, epoch, value}], seeded from the server so a + #: resumed job continues its curves instead of restarting them. + self._series: dict[str, list[dict[str, float | int]]] = dict(reporter.fetch_current_metrics()) + if any(self._series.values()): logger.info( - "Seeded metrics from server: %d train_loss, %d val_loss entries", - len(self._train_metrics), - len(self._val_metrics), + "Seeded %d metric series from server (%d points): %s", + len(self._series), + sum(len(points) for points in self._series.values()), + ", ".join(sorted(self._series)), ) def _resolve_backend(self, backend: str | None) -> str | None: return backend if backend is not None else self._default_backend + def _record(self, phase: str, name: str, step: int, epoch: int, value: object) -> None: + """Append one point to the ``_`` series, if it is chartable. + + Silently drops non-numeric values rather than raising: a backend adding a + metric that turns out to be a histogram should lose that one series, not + fail the training run's progress reporting. + """ + if not is_chartable(value): + return + # Coerce to a built-in: numpy scalars satisfy numbers.Real but are not + # JSON-serializable. Counts stay ints rather than becoming 64.0. + real = cast(numbers.Real, value) + numeric: float | int = int(real) if isinstance(real, numbers.Integral) else float(real) + series = name if name in _UNPREFIXED else f"{phase}_{name}" + self._series.setdefault(series, []).append({"step": step, "epoch": epoch, "value": numeric}) + def _build_metrics_summary(self) -> dict[str, list[dict[str, float | int]]]: - """Build the accumulated metrics payload for inclusion in status_details.""" - return { - "train_loss": list(self._train_metrics), - "val_loss": list(self._val_metrics), - } + """Build the accumulated metrics payload for inclusion in status_details. + + ``train_loss``/``val_loss`` are always present, even when empty, so the + shape stays stable for consumers that index them directly. Lists are + copied: the payload must not mutate after it is handed over. + """ + summary: dict[str, list[dict[str, float | int]]] = {"train_loss": [], "val_loss": []} + summary.update({name: list(points) for name, points in self._series.items()}) + return summary def report_training_start(self, max_steps: int, num_epochs: int, *, backend: str | None = None) -> None: """Report that training has started with schedule information.""" @@ -91,11 +159,16 @@ def report_train_step( ) -> None: """Report training step with metrics. - ``additional_metrics`` are backend-specific current-step scalars (DPO's - ``preference_loss``, GRPO's ``reward``/``kl_penalty``, ...). They are not - accumulated into the series; see the module docstring. + ``additional_metrics`` are backend-specific (DPO's ``preference_loss``, + GRPO's ``reward``/``kl_penalty``, ...). Each numeric one accumulates into + its own ``train_`` series *and* rides along as a current-step + scalar, so consumers can read either the curve or the latest value. """ - self._train_metrics.append({"step": step, "epoch": epoch, "value": loss}) + self._record("train", "train_loss", step, epoch, loss) + self._record("train", "lr", step, epoch, lr) + self._record("train", "grad_norm", step, epoch, grad_norm) + for name, value in additional_metrics.items(): + self._record("train", name, step, epoch, value) # `**additional_metrics` is splatted first, matching report_validation, so a # backend metric cannot shadow the accumulated series or the step's own loss. # `step`/`epoch`/`lr`/`grad_norm` are named parameters and so already safe. @@ -126,15 +199,18 @@ def report_validation( ``val_loss`` is optional because not every algorithm produces one: GRPO validates on ``accuracy``/``avg_length`` and reports no loss at all. The - key is omitted rather than sent as null, which would chart as a real zero. + key is omitted rather than sent as null, which would chart as a real zero, + and the ``val_loss`` series simply stays empty for such runs. """ details: dict[str, object] = { "step": step, "epoch": epoch, **additional_metrics, } + for name, value in additional_metrics.items(): + self._record("val", name, step, epoch, value) if val_loss is not None: - self._val_metrics.append({"step": step, "epoch": epoch, "value": val_loss}) + self._record("val", "val_loss", step, epoch, val_loss) details["val_loss"] = val_loss details["metrics"] = self._build_metrics_summary() diff --git a/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py b/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py index e80c0d0252..e04bcab0e9 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py @@ -115,8 +115,15 @@ def update_task( logger.warning(f"Failed to update task progress: {e}") def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: + """Read back every stored metric series. + + Deliberately not restricted to a known set of names: backends decide what + they accumulate, and a resumed job that only seeded ``train_loss`` would + silently restart every other curve from empty. Non-list values are + dropped so a malformed blob cannot poison the accumulator. + """ if not self._enabled: - return {"train_loss": [], "val_loss": []} + return {} try: jobs = client_from_platform(self._sdk, JobsClient) @@ -127,16 +134,13 @@ def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: step=self._job_ctx.step, ).data() metrics = cast(dict[str, Any], (task.status_details or {}).get("metrics", {}) or {}) - return { - "train_loss": metrics.get("train_loss", []), - "val_loss": metrics.get("val_loss", []), - } + return {name: points for name, points in metrics.items() if isinstance(points, list)} except Exception as e: # Expected on a first run, where the task has no stored details yet. # Serves both resume seeding and update_task's metric preservation, # so the message stays neutral about which caller hit it. logger.info(f"No stored metrics available: {e}") - return {"train_loss": [], "val_loss": []} + return {} def report_running(self, phase: str, **details: Any) -> None: if "step" in details and "percentage_done" not in details and self._max_steps > 0: diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py index b0873381e4..dc0fd65fba 100644 --- a/packages/nmp_customization_common/tests/training/test_callbacks.py +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -117,22 +117,88 @@ def test_validation_with_loss_records_both_key_and_series(reporter: _RecordingRe # --------------------------------------------------------------------------- # -def test_additional_train_metrics_ride_along_without_entering_the_series( +def test_additional_train_metrics_become_series_and_ride_along( reporter: _RecordingReporter, ) -> None: - """The wide backend metric set is current-step only; series stay bounded.""" - _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, reward=0.62, kl_penalty=0.008) + """Each backend metric is both a curve and a current-step scalar.""" + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5, reward=0.62, kl_penalty=0.008) + callback.report_train_step(step=2, epoch=1, loss=0.4, reward=0.71, kl_penalty=0.009) report = reporter.reports[-1] - assert report["reward"] == 0.62 - assert report["kl_penalty"] == 0.008 - assert report["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] + assert report["reward"] == 0.71, "latest value still rides along at the top level" + assert report["metrics"]["train_reward"] == [ + {"step": 1, "epoch": 1, "value": 0.62}, + {"step": 2, "epoch": 1, "value": 0.71}, + ] + assert report["metrics"]["train_kl_penalty"] == [ + {"step": 1, "epoch": 1, "value": 0.008}, + {"step": 2, "epoch": 1, "value": 0.009}, + ] + + +def test_lr_and_grad_norm_accumulate(reporter: _RecordingReporter) -> None: + """Both are curves people read; neither is an `additional_metric`.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, lr=5e-06, grad_norm=1.9) + + metrics = reporter.reports[-1]["metrics"] + assert metrics["train_lr"] == [{"step": 1, "epoch": 1, "value": 5e-06}] + assert metrics["train_grad_norm"] == [{"step": 1, "epoch": 1, "value": 1.9}] -def test_additional_validation_metrics_ride_along(reporter: _RecordingReporter) -> None: +def test_absent_lr_and_grad_norm_create_no_series(reporter: _RecordingReporter) -> None: + """A backend that reports neither should not get two empty keys.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5) + + metrics = reporter.reports[-1]["metrics"] + assert "train_lr" not in metrics + assert "train_grad_norm" not in metrics + + +def test_additional_validation_metrics_become_series_and_ride_along(reporter: _RecordingReporter) -> None: _make_callback(reporter).report_validation(step=1, epoch=1, val_loss=0.25, accuracy=0.9) - assert reporter.reports[-1]["accuracy"] == 0.9 + report = reporter.reports[-1] + assert report["accuracy"] == 0.9 + assert report["metrics"]["val_accuracy"] == [{"step": 1, "epoch": 1, "value": 0.9}] + + +def test_train_and_validation_metrics_of_the_same_name_stay_separate( + reporter: _RecordingReporter, +) -> None: + """GRPO reports `truncation_rate` in both dicts; one series would interleave them.""" + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5, truncation_rate=0.18) + callback.report_validation(step=1, epoch=1, truncation_rate=0.04) + + metrics = reporter.reports[-1]["metrics"] + assert metrics["train_truncation_rate"] == [{"step": 1, "epoch": 1, "value": 0.18}] + assert metrics["val_truncation_rate"] == [{"step": 1, "epoch": 1, "value": 0.04}] + + +def test_non_numeric_metrics_are_dropped_from_the_series(reporter: _RecordingReporter) -> None: + """Histograms and tables ride in the same dict as the scalars upstream.""" + _make_callback(reporter).report_train_step( + step=1, epoch=1, loss=0.5, histogram=object(), nested={"a": 1}, flag=True, missing=float("nan") + ) + + metrics = reporter.reports[-1]["metrics"] + assert set(metrics) == {"train_loss", "val_loss"} + + +def test_series_survive_a_resume_beyond_the_loss_curves() -> None: + """A resumed job must continue every curve, not just train_loss.""" + prior = { + "train_loss": [{"step": 1, "epoch": 1, "value": 0.9}], + "train_reward": [{"step": 1, "epoch": 1, "value": 0.2}], + } + reporter = _RecordingReporter(prior) + _make_callback(reporter).report_train_step(step=2, epoch=1, loss=0.8, reward=0.3) + + assert reporter.reports[-1]["metrics"]["train_reward"] == [ + {"step": 1, "epoch": 1, "value": 0.2}, + {"step": 2, "epoch": 1, "value": 0.3}, + ] def test_additional_metrics_cannot_shadow_the_series(reporter: _RecordingReporter) -> None: diff --git a/services/automodel/tests/tasks/training/backends/test_callbacks.py b/services/automodel/tests/tasks/training/backends/test_callbacks.py index 46132c3903..03276934c2 100644 --- a/services/automodel/tests/tasks/training/backends/test_callbacks.py +++ b/services/automodel/tests/tasks/training/backends/test_callbacks.py @@ -104,8 +104,8 @@ def test_seeds_from_server_on_init(self): } callback, reporter = self._make_callback(prior_metrics=prior) - assert len(callback._train_metrics) == 2 - assert len(callback._val_metrics) == 1 + assert len(callback._series["train_loss"]) == 2 + assert len(callback._series["val_loss"]) == 1 reporter.fetch_current_metrics.assert_called_once() def test_seeded_metrics_included_in_first_report(self): diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py index 1ecec2fbc5..fe381d653e 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py @@ -9,12 +9,11 @@ # its affiliates is strictly prohibited. import logging -import math -import numbers from typing import Any, Mapping, Optional from nemo_rl.utils.logger import LoggerInterface from nmp.customization_common.service.context import NMPJobContext +from nmp.customization_common.training.callbacks import is_chartable from nmp.rl.tasks.training.backends.nemo_rl.callbacks import TrainingProgressCallback from nmp.rl.tasks.training.progress import JobsServiceProgressReporter @@ -71,12 +70,11 @@ def has_metric_value(metric: Any) -> bool: nested dict. ``math.isnan`` raises ``TypeError`` on all of those, so a bare None-check would turn a widened key list into a crash mid-training. - ``bool`` is rejected despite being an ``int`` subclass: no metric here is a - flag, and silently charting one as 0/1 is worse than dropping it. + Delegates to the shared predicate so the wire filter and the series filter + cannot drift apart -- a metric this forwards must be one the callback can + chart. """ - if isinstance(metric, bool) or not isinstance(metric, numbers.Real): - return False - return not math.isnan(float(metric)) + return is_chartable(metric) def resolve_log_interval(val_period: int | None) -> int: From b44e73a1b6c971fe3c76d4fba08ea3cb331c25f7 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 11:55:46 -0400 Subject: [PATCH 09/10] feat(rl): forward the GRPO diagnostics and reward spread ASTD-388 asks for AC6 of ASTD-388 ("Job monitoring shows reward, not loss") names six training- health metrics and a reward-distribution band. Two of the six were forwarded; the band had nothing behind it. Nothing was missing upstream. All four remaining health metrics -- gen_kl_error, policy_kl_error, js_divergence_error, sampling_importance_ratio -- are returned unconditionally by ClippedPGLossFn (algorithms/loss/loss_functions.py), land in all_mb_metrics, and are merged into the same `train` dict we already read. They are entries in that loss fn's `metric_normalizations` map, so the per-microbatch sum NeMo-RL applies is already the correct global value -- identical treatment to approx_entropy and token_mult_prob_error, which we forward today. The whitelist just never listed them. Same for the band: calculate_single_metric emits total_reward/{mean,median,min, max,stddev,histogram} and we took only /mean. The scalar four now come through; /histogram is a Histogram object and is dropped as non-scalar, as intended. These are what say a GRPO run is failing before the reward curve shows it: entropy collapse, and the logprob-error family that means the trainer and the generation engine have drifted apart. Signed-off-by: Albert Cui --- .../backends/nemo_rl/nemo_rl_logger.py | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py index fe381d653e..9aa5463288 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py @@ -39,16 +39,32 @@ "preference_loss", "rewards_chosen_mean", "rewards_rejected_mean", - # GRPO: reward and advantages — the signal that says whether RL is working + # GRPO: reward and advantages — the signal that says whether RL is working. + # The total_reward spread is what a reward-distribution band is drawn from; + # `calculate_single_metric` emits the whole set (a /histogram rides along too, + # but it is a Histogram object and is dropped as non-scalar). "reward", "total_reward/mean", + "total_reward/median", + "total_reward/min", + "total_reward/max", + "total_reward/stddev", "advantages/mean", "advantages/min", "advantages/max", - # GRPO: policy-optimization health + # GRPO: policy-optimization health. These are the diagnostics that predict a + # failing run -- entropy collapse, and the logprob-error family that says the + # trainer and the generation engine have drifted apart. All are returned + # unconditionally by ClippedPGLossFn and normalized via its + # `metric_normalizations` map, so the per-microbatch sum NeMo-RL applies is + # already the correct global value. "kl_penalty", "approx_entropy", "token_mult_prob_error", + "gen_kl_error", + "policy_kl_error", + "js_divergence_error", + "sampling_importance_ratio", # GRPO: rollout shape "truncation_rate", "natural_termination_rate", From ad427fb2ed86312db0659d3c6c62ab0837443b53 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 11:55:46 -0400 Subject: [PATCH 10/10] fix(customization): carry sticky status_details fields across updates Generalizes the metric preservation added earlier. `status_details` is REPLACED on every update, so a field survives only as long as the next report repeats it -- and three kinds of field were being lost to that, not one: metrics erased by the runner's checkpoint/completion/failure reports, which come from a different process max_steps, num_epochs stated once by report_training_start, gone from the first training step onward checkpoint_path published by one report, wiped by the next The last two are why ASTD-388 AC6's "Step / max steps" and "Latest checkpoint" could not be rendered: Studio reads both straight out of status_details (util/customizations.tsx), and both were absent for almost the whole run. The earlier fix special-cased `metrics` on the theory that everything else was a current-state snapshot that should expire. That split was wrong. There is a third category -- facts that stay true after the update that stated them -- and _CARRY_FORWARD now names it: cumulative (metrics), run constants (max_steps, num_epochs), monotonic progress (step, epoch), and sticky latest-values (checkpoint_path). Still excluded, deliberately: `phase`, which every report sets for itself, and the per-step observations (train_loss, lr, grad_norm, reward, ...). Those describe one instant and a stale copy would misrepresent "current" -- and nothing is lost, because each is now recoverable from its series. percentage_done is excluded too: it is derived from step and max_steps, both carried, so a consumer can recompute it rather than risk a copy that contradicts its inputs. Keeping the GET off the hot path is the whole design constraint. Values are remembered as they pass through, so a process that has already stated a field restates it for free; the blob is read back only when an update omits `metrics`, which is the tell that it did not come from TrainingProgressCallback. Per-step reports always carry `metrics` and never fetch. The runner's handful of reports always do. One subtlety: on resume, the driver's first report already carries `metrics`, so it would never read the blob back and would drop the previous run's checkpoint_path. _fetch_status_details therefore refreshes the cache as a side effect, which makes the resume-seeding fetch the callback already performs at construction double as the carry-forward seed -- no extra round-trip. Tests drive the SDK client seam rather than stubbing the fetch, so the real _fetch_status_details runs, cache side effect included. Signed-off-by: Albert Cui --- .../customization_common/training/progress.py | 144 ++++++--- .../tests/training/test_progress.py | 283 +++++++++++++----- 2 files changed, 318 insertions(+), 109 deletions(-) diff --git a/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py b/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py index e04bcab0e9..03522d2bb1 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py @@ -8,9 +8,9 @@ training runner; backends subclass it (or instantiate it directly) supplying their own ``service_name`` so the task SDK resolves the right credentials. -Every update REPLACES the task's ``status_details``. The accumulated metric -series is the one cumulative field in that blob, so ``update_task`` carries it -across updates that don't supply their own -- see :meth:`_preserve_metrics`. +Every update REPLACES the task's ``status_details``, so a field is only as +durable as the next report that omits it. ``update_task`` carries a defined set +of fields across updates that don't restate them -- see :data:`_CARRY_FORWARD`. For training-specific metrics (loss, validation, checkpoints) see the ``TrainingProgressCallback`` which composes this reporter. @@ -29,6 +29,43 @@ logger = logging.getLogger(__name__) +#: Fields restated on updates that don't supply their own. +#: +#: The rule is *what stays true after the update that stated it*: +#: +#: - ``metrics`` is cumulative -- the whole point is that it grows. +#: - ``max_steps``/``num_epochs`` are run constants, and are only ever stated +#: once, by ``report_training_start``. +#: - ``step``/``epoch`` are monotonic; a run does not un-reach step 30. +#: - ``checkpoint_path`` is a sticky latest-value, true until superseded. +#: +#: Deliberately excluded: ``phase`` (every report sets its own), and the +#: per-step observations (``train_loss``, ``lr``, ``grad_norm``, ``reward``, +#: ...). Those describe one instant, and a stale copy would misrepresent +#: "current" -- nothing is lost by letting them expire, because every one of +#: them is now recoverable from its series in ``metrics``. +#: +#: ``percentage_done`` is also excluded: it is derived from ``step`` and +#: ``max_steps``, both of which are carried, so a consumer can recompute it +#: rather than risk a copy that contradicts its own inputs. +_CARRY_FORWARD = frozenset({"metrics", "max_steps", "num_epochs", "step", "epoch", "checkpoint_path"}) + + +def _carries_information(value: Any) -> bool: + """Whether a stored value is worth restating on a later update. + + Empty containers are dropped so a task doesn't accumulate keys that say + nothing -- notably the all-empty ``metrics`` dict a job reports before its + first training step. + """ + if value is None: + return False + if isinstance(value, dict): + return any(_carries_information(item) for item in value.values()) + if isinstance(value, (list, str)): + return bool(value) + return True + class JobsServiceProgressReporter: """Reports high-level progress to the Jobs service.""" @@ -40,6 +77,10 @@ def __init__(self, job_ctx: NMPJobContext, service_name: str): self._max_steps = 0 self._num_epochs = 0 + #: Last-seen value of each :data:`_CARRY_FORWARD` field, populated as + #: updates pass through and from the stored blob when one is read back. + self._carried: dict[str, Any] = {} + # Gate on real job context, not bare truthiness: from_env() fills missing # identifiers with non-empty sentinel defaults, which would otherwise # enable reporting (and failing SDK calls) outside a real job run. @@ -57,31 +98,40 @@ def _calculate_percentage_done(self, step: int | None) -> int: # downstream progress consumers expect a bounded percentage. return min(100, int((step / self._max_steps) * 100)) - def _preserve_metrics(self, status_details: dict[str, Any] | None) -> dict[str, Any]: - """Re-attach the stored metric series to an update that doesn't carry one. - - ``status_details`` is REPLACED by the Jobs service, not merged, so any - update omitting ``metrics`` blanks the accumulated loss curve. - ``TrainingProgressCallback`` resends the series on every report it makes, - but the surrounding runner cannot: it reports ``processing_checkpoint``, - completion and failure from a *different process* than the training - driver that accumulated the series, so it holds nothing to resend. Every - job would otherwise end by erasing its own curve -- including, and most - expensively, on the failure path. - - Only ``metrics`` is carried over. The rest of the blob is deliberately a - current-state snapshot (``phase``, ``step``, ``lr``, ...); merging that - would leave a completed task advertising a mid-training step. - - Costs a GET only on updates that omit ``metrics``, which is the handful - the runner makes per job -- never the per-step training reports. + def _carry_forward(self, status_details: dict[str, Any] | None) -> dict[str, Any]: + """Restate the :data:`_CARRY_FORWARD` fields this update doesn't supply. + + ``status_details`` is REPLACED by the Jobs service, not merged, so a + field survives only as long as every subsequent report repeats it. Three + things were being lost to that: + + - the accumulated ``metrics``, on the runner's checkpoint/completion/ + failure reports -- so every job ended by erasing its own curves; + - ``max_steps``/``num_epochs``, stated once at training start and gone + from the first training step onward; + - ``checkpoint_path``, published by one report and wiped by the next. + + Values are remembered as they pass through (write-through), so a process + that has already stated a field can restate it for free. The stored blob + is read back only when the update omits ``metrics``, which is the tell + that it did not come from ``TrainingProgressCallback`` -- i.e. it is one + of the handful the runner makes, from a different process that holds no + state. Per-step training reports always carry ``metrics``, so the hot + path never pays for a round-trip. """ details = dict(status_details or {}) - if "metrics" in details: + self._remember(details) + + missing = [field for field in _CARRY_FORWARD if field not in details] + if not missing: return details - stored = self.fetch_current_metrics() - if any(stored.values()): - details["metrics"] = stored + + if "metrics" not in details: + self._fetch_status_details() + + for field in missing: + if field in self._carried: + details[field] = self._carried[field] return details def update_task( @@ -96,7 +146,7 @@ def update_task( if not self._is_main_rank: return - details = self._preserve_metrics(status_details) + details = self._carry_forward(status_details) try: jobs = client_from_platform(self._sdk, JobsClient) @@ -114,13 +164,14 @@ def update_task( except Exception as e: logger.warning(f"Failed to update task progress: {e}") - def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: - """Read back every stored metric series. + def _fetch_status_details(self) -> dict[str, Any]: + """Read back the task's stored ``status_details`` blob. - Deliberately not restricted to a known set of names: backends decide what - they accumulate, and a resumed job that only seeded ``train_loss`` would - silently restart every other curve from empty. Non-list values are - dropped so a malformed blob cannot poison the accumulator. + Refreshes the carry-forward cache as a side effect, so that the + resume-seeding fetch ``TrainingProgressCallback`` makes at construction + doubles as the seed for :meth:`_carry_forward`. Without that, a resumed + run would drop the previous run's ``checkpoint_path``: its first report + already carries ``metrics``, so it would never read the blob back. """ if not self._enabled: return {} @@ -133,15 +184,34 @@ def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: job=self._job_ctx.job_id, step=self._job_ctx.step, ).data() - metrics = cast(dict[str, Any], (task.status_details or {}).get("metrics", {}) or {}) - return {name: points for name, points in metrics.items() if isinstance(points, list)} + stored = cast(dict[str, Any], task.status_details or {}) except Exception as e: # Expected on a first run, where the task has no stored details yet. - # Serves both resume seeding and update_task's metric preservation, - # so the message stays neutral about which caller hit it. - logger.info(f"No stored metrics available: {e}") + # Serves both resume seeding and update_task's carry-forward, so the + # message stays neutral about which caller hit it. + logger.info(f"No stored status details available: {e}") return {} + self._remember(stored) + return stored + + def _remember(self, source: dict[str, Any]) -> None: + """Cache the carry-forward fields present in ``source``.""" + self._carried.update( + {field: value for field, value in source.items() if field in _CARRY_FORWARD and _carries_information(value)} + ) + + def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: + """Read back every stored metric series, for resume seeding. + + Deliberately not restricted to a known set of names: backends decide what + they accumulate, and a resumed job that only seeded ``train_loss`` would + silently restart every other curve from empty. Non-list values are + dropped so a malformed blob cannot poison the accumulator. + """ + metrics = cast(dict[str, Any], self._fetch_status_details().get("metrics", {}) or {}) + return {name: points for name, points in metrics.items() if isinstance(points, list)} + def report_running(self, phase: str, **details: Any) -> None: if "step" in details and "percentage_done" not in details and self._max_steps > 0: details["percentage_done"] = self._calculate_percentage_done(details["step"]) diff --git a/packages/nmp_customization_common/tests/training/test_progress.py b/packages/nmp_customization_common/tests/training/test_progress.py index 7dae541bf8..f34204794a 100644 --- a/packages/nmp_customization_common/tests/training/test_progress.py +++ b/packages/nmp_customization_common/tests/training/test_progress.py @@ -3,10 +3,11 @@ """Unit tests for JobsServiceProgressReporter's status_details handling. -Focused on metric preservation: the Jobs service REPLACES ``status_details``, and -the runner reports checkpoint processing, completion and failure from a different -process than the training driver that accumulated the loss curve. Without the -carry-over in ``update_task`` every job ends by erasing its own metrics. +Focused on carry-forward: the Jobs service REPLACES ``status_details``, so a +field lives only as long as the next report repeats it. The runner reports +checkpoint processing, completion and failure from a different process than the +training driver, and the driver states the schedule once and the checkpoint path +once. Without the carry-forward set, each of those is erased by the next update. """ from __future__ import annotations @@ -16,11 +17,22 @@ import pytest from nmp.customization_common.training.progress import JobsServiceProgressReporter -SERIES: dict[str, list[dict[str, float | int]]] = { +SERIES: dict[str, list[dict[str, Any]]] = { "train_loss": [{"step": 10, "epoch": 1, "value": 0.5}], - "val_loss": [{"step": 10, "epoch": 1, "value": 0.45}], + "train_reward": [{"step": 10, "epoch": 1, "value": 0.62}], +} +#: A blob a mid-run job would have stored: series plus the sticky facts. +STORED: dict[str, Any] = { + "phase": "training", + "step": 10, + "epoch": 1, + "max_steps": 30, + "num_epochs": 3, + "train_loss": 0.5, + "lr": 5e-06, + "checkpoint_path": "/ckpt/step-10", + "metrics": SERIES, } -EMPTY: dict[str, list[dict[str, float | int]]] = {"train_loss": [], "val_loss": []} class _JobCtx: @@ -38,114 +50,242 @@ class _Reporter(JobsServiceProgressReporter): Bypasses ``__init__`` rather than mocking the SDK factory: what is under test is the status_details logic, and the real constructor calls ``get_task_sdk``, which wants credentials. Every attribute ``update_task`` touches is set here. + + The SDK client itself is patched (see the ``jobs`` fixture) rather than + ``_fetch_status_details``, so the real fetch runs -- including its + carry-forward cache side effect. """ - def __init__(self, stored: dict[str, list[dict[str, float | int]]]) -> None: + def __init__(self) -> None: self._job_ctx = _JobCtx() # type: ignore[assignment] - duck-typed stand-in - self._sdk = object() # type: ignore[assignment] - never dereferenced; the client is patched + self._sdk = type("S", (), {"close": lambda self: None})() # type: ignore[assignment] self._is_main_rank = True self._enabled = True self._max_steps = 0 self._num_epochs = 0 - self._stored = stored - self.fetch_calls = 0 + self._carried = {} - def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: - self.fetch_calls += 1 - return self._stored +class _Task: + def __init__(self, status_details: dict[str, Any]) -> None: + self.status_details = status_details -class _StubJobsClient: - """Captures task updates instead of issuing them.""" + def data(self) -> "_Task": + return self - def __init__(self, sink: list[dict[str, Any]]) -> None: - self._sink = sink - def update_job_step_task(self, **kwargs: Any) -> None: - self._sink.append(kwargs) +class _Jobs: + """A mini Jobs service: replace-on-write, readable back, counting fetches.""" + def __init__(self) -> None: + self.sent: list[dict[str, Any]] = [] + self.stored: dict[str, Any] = {} + self.fetches = 0 + self.persist = False -@pytest.fixture -def sent(monkeypatch: pytest.MonkeyPatch) -> list[dict[str, Any]]: - """Capture what ``update_task`` would send, running its real body. + def client(self) -> Any: + harness = self - ``update_task`` swallows exceptions, so anything wrong with the stubs would - surface as an empty capture; the helpers below index into it eagerly so that - reads as a failure rather than a pass. - """ - sink: list[dict[str, Any]] = [] + class _Client: + def update_job_step_task(self, **kwargs: Any) -> None: + harness.sent.append(kwargs) + if harness.persist: + harness.stored = dict(kwargs["body"].status_details or {}) + + def get_job_step_task(self, **kwargs: Any) -> _Task: + harness.fetches += 1 + return _Task(harness.stored) + + return _Client() + + +@pytest.fixture +def jobs(monkeypatch: pytest.MonkeyPatch) -> _Jobs: + """Patch the SDK client seam so update_task and the fetch both run for real.""" + harness = _Jobs() monkeypatch.setattr( "nmp.customization_common.training.progress.client_from_platform", - lambda _sdk, _cls: _StubJobsClient(sink), + lambda _sdk, _cls: harness.client(), ) - return sink + return harness + +def _reporter(jobs: _Jobs, stored: dict[str, Any] | None = None) -> _Reporter: + """A reporter over the harness, with the server pre-seeded if given.""" + if stored is not None: + jobs.stored = stored + return _Reporter() -def _details(sent: list[dict[str, Any]]) -> dict[str, Any]: - assert len(sent) == 1, f"expected exactly one task update, got {len(sent)}" - return dict(sent[0]["body"].status_details or {}) + +def _details(jobs: _Jobs, index: int = -1) -> dict[str, Any]: + assert jobs.sent, "expected at least one task update" + return dict(jobs.sent[index]["body"].status_details or {}) # --------------------------------------------------------------------------- # -# Preservation +# What carries forward # --------------------------------------------------------------------------- # -def test_completion_preserves_the_accumulated_series(sent: list[dict[str, Any]]) -> None: - """The last write of a successful job must not blank the loss curve.""" - _Reporter(SERIES).report_completed("Training completed") +def test_completion_carries_series_schedule_and_checkpoint(jobs: _Jobs) -> None: + """The last write of a successful job must not blank what it took to get there.""" + _reporter(jobs, STORED).report_completed("Training completed") - details = _details(sent) + details = _details(jobs) assert details["metrics"] == SERIES - assert details["phase"] == "completed" + assert details["max_steps"] == 30 + assert details["num_epochs"] == 3 + assert details["step"] == 10 + assert details["checkpoint_path"] == "/ckpt/step-10" + assert details["phase"] == "completed", "the report's own phase still wins" -def test_failure_preserves_the_accumulated_series(sent: list[dict[str, Any]]) -> None: - """A failed run is exactly when the partial curve is most worth keeping.""" - _Reporter(SERIES).report_error("boom") +def test_failure_carries_the_same_set(jobs: _Jobs) -> None: + """A failed run is exactly when the partial curve and last checkpoint matter.""" + _reporter(jobs, STORED).report_error("boom") - assert _details(sent)["metrics"] == SERIES + details = _details(jobs) + assert details["metrics"] == SERIES + assert details["step"] == 10 + assert details["checkpoint_path"] == "/ckpt/step-10" -def test_intermediate_phase_preserves_the_accumulated_series(sent: list[dict[str, Any]]) -> None: +def test_intermediate_phase_carries_forward(jobs: _Jobs) -> None: """processing_checkpoint fires after the driver exits, before completion.""" - _Reporter(SERIES).report_running("processing_checkpoint") + _reporter(jobs, STORED).report_running("processing_checkpoint") - details = _details(sent) + details = _details(jobs) assert details["metrics"] == SERIES + assert details["max_steps"] == 30 assert details["phase"] == "processing_checkpoint" -def test_caller_supplied_metrics_win_and_skip_the_fetch(sent: list[dict[str, Any]]) -> None: - """Per-step training reports carry their own series; no round-trip for them.""" - fresher = {"train_loss": [{"step": 20, "epoch": 2, "value": 0.1}], "val_loss": []} - reporter = _Reporter(SERIES) +def test_per_step_observations_do_not_carry_forward(jobs: _Jobs) -> None: + """A completed task must not advertise a stale current loss or learning rate. - reporter.report_running("training", step=20, metrics=fresher) + Nothing is lost: each of these is recoverable from its series in `metrics`. + """ + _reporter(jobs, STORED).report_completed("Training completed") - assert _details(sent)["metrics"] == fresher - assert reporter.fetch_calls == 0 + details = _details(jobs) + assert "train_loss" not in details + assert "lr" not in details -def test_no_stored_metrics_adds_no_key(sent: list[dict[str, Any]]) -> None: - """Before training starts there is nothing to preserve; don't invent a key.""" - _Reporter(EMPTY).report_running("compiling_config") +def test_caller_supplied_values_win(jobs: _Jobs) -> None: + fresher = {"train_loss": [{"step": 20, "epoch": 2, "value": 0.1}]} + _reporter(jobs, STORED).report_running("training", step=20, metrics=fresher, max_steps=99) + + details = _details(jobs) + assert details["metrics"] == fresher + assert details["step"] == 20 + assert details["max_steps"] == 99 + + +def test_empty_stored_values_add_no_keys(jobs: _Jobs) -> None: + """Before training starts there is nothing to carry; don't invent keys.""" + stored = {"metrics": {"train_loss": [], "val_loss": []}, "checkpoint_path": ""} + _reporter(jobs, stored).report_running("compiling_config") + + details = _details(jobs) + assert "metrics" not in details + assert "checkpoint_path" not in details + + +# --------------------------------------------------------------------------- # +# Write-through cache: the per-step hot path must not pay for a round-trip +# --------------------------------------------------------------------------- # + + +def test_reports_carrying_metrics_never_fetch(jobs: _Jobs) -> None: + """`metrics` marks an update as coming from the accumulating callback. + + Those are the per-step reports. They omit max_steps and checkpoint_path, so a + naive implementation would read the blob back on every single training step. + """ + reporter = _reporter(jobs, STORED) + for step in range(1, 11): + reporter.report_running("training", step=step, metrics=SERIES) - assert "metrics" not in _details(sent) + assert jobs.fetches == 0 -def test_preservation_does_not_resurrect_stale_current_state(sent: list[dict[str, Any]]) -> None: - """Only the cumulative field carries over, not the step/lr snapshot.""" - _Reporter(SERIES).report_completed("Training completed") +def test_a_stated_value_is_restated_without_a_fetch(jobs: _Jobs) -> None: + """report_training_start states the schedule once; every later step needs it.""" + reporter = _reporter(jobs) + reporter.report_running("training", step=0, max_steps=30, num_epochs=3, metrics=SERIES) + reporter.report_running("training", step=1, metrics=SERIES) - assert set(_details(sent)) == {"message", "phase", "metrics"} + details = _details(jobs) + assert details["max_steps"] == 30 + assert details["num_epochs"] == 3 + assert jobs.fetches == 0 -def test_error_details_still_ride_along(sent: list[dict[str, Any]]) -> None: - """Preserving metrics must not displace the error payload.""" - _Reporter(SERIES).report_error({"message": "oom", "code": "OOM"}) +def test_checkpoint_path_survives_the_next_training_step(jobs: _Jobs) -> None: + """It was published by one report and wiped by the very next one.""" + reporter = _reporter(jobs) + reporter.report_running("checkpoint_saved", step=10, checkpoint_path="/ckpt/step-10", metrics=SERIES) + reporter.report_running("training", step=11, metrics=SERIES) + + assert _details(jobs)["checkpoint_path"] == "/ckpt/step-10" + + +def test_a_newer_checkpoint_supersedes_the_carried_one(jobs: _Jobs) -> None: + reporter = _reporter(jobs) + reporter.report_running("checkpoint_saved", step=10, checkpoint_path="/ckpt/step-10", metrics=SERIES) + reporter.report_running("checkpoint_saved", step=20, checkpoint_path="/ckpt/step-20", metrics=SERIES) + reporter.report_running("training", step=21, metrics=SERIES) + + assert _details(jobs)["checkpoint_path"] == "/ckpt/step-20" + + +def test_updates_without_metrics_read_the_blob_back(jobs: _Jobs) -> None: + """The runner's reports come from a process that holds no state at all.""" + reporter = _reporter(jobs, STORED) + reporter.report_running("processing_checkpoint") + + assert jobs.fetches == 1 + + +def test_error_details_still_ride_along(jobs: _Jobs) -> None: + """Carry-forward must not displace the error payload.""" + _reporter(jobs, STORED).report_error({"message": "oom", "code": "OOM"}) + + assert jobs.sent[0]["body"].error_details == {"message": "oom", "code": "OOM"} + + +# --------------------------------------------------------------------------- # +# Resume seeding +# --------------------------------------------------------------------------- # + + +def test_fetch_current_metrics_returns_every_series(jobs: _Jobs) -> None: + """A resumed job that only seeded train_loss would restart the other curves.""" + assert _reporter(jobs, STORED).fetch_current_metrics() == SERIES + + +def test_fetch_current_metrics_drops_non_list_values(jobs: _Jobs) -> None: + """A malformed blob must not poison the accumulator.""" + stored = {"metrics": {"train_loss": [{"step": 1, "epoch": 1, "value": 1.0}], "junk": 3}} + + assert set(_reporter(jobs, stored).fetch_current_metrics()) == {"train_loss"} + + +def test_resume_seeding_also_seeds_the_carry_forward_cache(jobs: _Jobs) -> None: + """The callback's construction-time fetch must double as the carry-forward seed. + + Otherwise a resumed run drops the previous run's checkpoint_path: its very + first report already carries `metrics`, so it never reads the blob back. + """ + reporter = _reporter(jobs, STORED) + reporter.fetch_current_metrics() # what TrainingProgressCallback.__init__ does + fetches_after_seeding = jobs.fetches + + reporter.report_running("training", step=1, metrics=SERIES) - assert sent[0]["body"].error_details == {"message": "oom", "code": "OOM"} + assert _details(jobs)["checkpoint_path"] == "/ckpt/step-10" + assert jobs.fetches == fetches_after_seeding, "no second round-trip" # --------------------------------------------------------------------------- # @@ -153,20 +293,19 @@ def test_error_details_still_ride_along(sent: list[dict[str, Any]]) -> None: # --------------------------------------------------------------------------- # -def test_disabled_reporter_sends_nothing_and_does_not_fetch(sent: list[dict[str, Any]]) -> None: - reporter = _Reporter(SERIES) +def test_disabled_reporter_sends_nothing(jobs: _Jobs) -> None: + reporter = _reporter(jobs, STORED) reporter._enabled = False reporter.report_completed("Training completed") - assert sent == [] - assert reporter.fetch_calls == 0 + assert jobs.sent == [] -def test_non_main_rank_sends_nothing(sent: list[dict[str, Any]]) -> None: - reporter = _Reporter(SERIES) +def test_non_main_rank_sends_nothing(jobs: _Jobs) -> None: + reporter = _reporter(jobs, STORED) reporter._is_main_rank = False reporter.report_completed("Training completed") - assert sent == [] + assert jobs.sent == []