fix(retrieval): do not lose metric records on preemption or error

nemo_automodel/components/loggers/metric_logger.py

@@ -137,10 +137,23 @@ def _save(self, lines: List[str]) -> None:
             self._fp.flush()
             os.fsync(self._fp.fileno())
 
+    def _drain(self) -> None:
+        """Write buffered records out. Caller must hold ``self._lock``."""
+        self._save(self._move_to_cpu(self.buffer))
+        self.buffer = []
+
+    def flush(self) -> None:
+        """Write buffered records to disk without closing the file.
+
+        Lets a caller align durability with an external event -- a checkpoint boundary, say --
+        instead of relying on the record count happening to reach ``buffer_size`` there.
+        """
+        with self._lock:
+            self._drain()
+
     def close(self) -> None:
         with self._lock:
-            self._save(self._move_to_cpu(self.buffer))
-            self.buffer = []
+            self._drain()
             try:
                 self._fp.flush()
             except Exception:
@@ -173,6 +186,11 @@ def log(self, record: MetricsSample) -> None:
             return
         super().log(record)
 
+    def flush(self) -> None:
+        if self.rank != 0:
+            return
+        super().flush()
+
     def close(self) -> None:
         if self.rank != 0:
             return

nemo_automodel/recipes/retrieval/train_bi_encoder.py

@@ -31,7 +31,11 @@
 from nemo_automodel.components.distributed.init_utils import initialize_distributed
 from nemo_automodel.components.distributed.utils import FirstRankPerNode, get_sync_ctx
 from nemo_automodel.components.loggers.log_utils import setup_logging
-from nemo_automodel.components.loggers.metric_logger import MetricsSample, build_metric_logger
+from nemo_automodel.components.loggers.metric_logger import (
+    DEFAULT_BUFFER_SIZE,
+    MetricsSample,
+    build_metric_logger,
+)
 from nemo_automodel.components.loggers.wandb_utils import suppress_wandb_log_messages
 from nemo_automodel.components.optim.precision_warnings import warn_if_torch_adam_with_bf16_params
 from nemo_automodel.components.training.rng import ScopedRNG, StatefulRNG
@@ -384,11 +388,24 @@ def materialize_loader(config):
         )
         self._log_model_and_optimizer_details(self.model_parts, self.optimizer, self.lr_scheduler)
 
+        # buffer_size bounds how many records can be pending, so a hard kill that runs no
+        # cleanup loses at most one checkpoint interval of train metrics; it is NOT what keeps
+        # metrics in step with checkpoints, since it counts records rather than watching
+        # checkpoint boundaries. The explicit flush after each successful checkpoint is what
+        # provides that. flush=True because a buffered write only reaches the file object,
+        # whose own buffer dies with the process. Validation is rare enough to write every
+        # point.
+        train_logger_kwargs = {"flush": True}
+        if self.step_scheduler.ckpt_every_steps > 0:
+            train_logger_kwargs["buffer_size"] = min(DEFAULT_BUFFER_SIZE, self.step_scheduler.ckpt_every_steps)
         self.metric_logger_train = build_metric_logger(
-            pathlib.Path(self.checkpointer.config.checkpoint_dir) / "training.jsonl"
+            pathlib.Path(self.checkpointer.config.checkpoint_dir) / "training.jsonl",
+            **train_logger_kwargs,
         )
         self.metric_logger_valid = build_metric_logger(
-            pathlib.Path(self.checkpointer.config.checkpoint_dir) / "validation.jsonl"
+            pathlib.Path(self.checkpointer.config.checkpoint_dir) / "validation.jsonl",
+            buffer_size=1,
+            flush=True,
         )
         self.loss_average_window = deque(maxlen=self.step_scheduler.loss_average_window_steps)
 
@@ -431,13 +448,34 @@ def run_train_validation_loop(self):
                             train_loss=train_log_data.metrics["loss"],
                             val_loss=val_loss,
                         )
+                        # Flush AFTER the checkpoint lands, so a checkpoint is never on disk
+                        # without the metrics describing the steps it covers. Buffer size alone
+                        # does not give this: it counts records, and the count only lines up
+                        # with checkpoint steps while training runs from step 0 with periodic
+                        # checkpoints. Resuming at a step that is not a multiple of
+                        # ckpt_every_steps, or a checkpoint taken at an epoch boundary or the
+                        # final step, leaves the two out of phase.
+                        if self.metric_logger_train is not None:
+                            self.metric_logger_train.flush()
+                        if self.metric_logger_valid is not None:
+                            self.metric_logger_valid.flush()
                     self._maybe_collect_garbage()
+
+                    # Poll per step; the StepScheduler iterators stop on the flag.
+                    if self.step_scheduler.sigterm_received:
+                        logger.info(
+                            "Preemption signal received at step %d; stopping cleanly.", self.step_scheduler.step
+                        )
         finally:
             if pbar is not None:
                 pbar.close()
+            # In the finally, not after it: an exception (OOM, NCCL timeout) would
+            # otherwise propagate past these and discard every buffered record.
+            if self.metric_logger_train is not None:
+                self.metric_logger_train.close()
+            if self.metric_logger_valid is not None:
+                self.metric_logger_valid.close()
 
-        self.metric_logger_train.close()
-        self.metric_logger_valid.close()
         self._finalize_and_close_checkpointer()
 
     def _forward_backward_step(self, idx, batch, *, loss_buffer, num_batches, is_train: bool = True):

tests/unit_tests/recipes/test_retrieval_bi_encoder_recipe.py

@@ -16,6 +16,7 @@
 from contextlib import nullcontext
 from types import SimpleNamespace
 
+import pytest
 import torch
 
 from nemo_automodel.components.distributed.config import DDPConfig, FSDP2Config
@@ -259,3 +260,159 @@ def _fake_scale_grads_and_clip_grad_norm(*args, **kwargs):
     recipe._run_train_optim_step([{}], max_grad_norm=1.0)
 
     assert captured["use_torch_clip_grad_norm"] is False
+
+
+class _RecordingMetricLogger:
+    """Buffers like MetricLogger but records what actually reached "disk" and when."""
+
+    def __init__(self):
+        self.buffer = []
+        self.persisted = []
+        self.closed = False
+
+    def log(self, record):
+        self.buffer.append(record)
+
+    def flush(self):
+        self.persisted.extend(self.buffer)
+        self.buffer = []
+
+    def close(self):
+        self.flush()
+        self.closed = True
+
+
+class _LoopStepScheduler:
+    """Minimal StepScheduler stand-in driving a fixed number of steps.
+
+    ``sigterm_at`` makes the signal appear at that step, mirroring the real scheduler's
+    sticky flag: once raised it stays raised, and the epoch/step iterators stop on it.
+    """
+
+    def __init__(self, n_steps, ckpt_steps=(), sigterm_at=None, val_steps=()):
+        self.step = 0
+        self.epoch = 0
+        self._n = n_steps
+        self._ckpt = set(ckpt_steps)
+        self._val = set(val_steps)
+        self._sigterm_at = sigterm_at
+        self.sigterm_flag = False
+        self.steps_run = []
+
+    @property
+    def epochs(self):
+        for e in range(1):
+            if self.sigterm_received:
+                return
+            yield e
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+    def __iter__(self):
+        while self.step < self._n:
+            self.step += 1
+            if self.sigterm_flag:
+                return
+            self.steps_run.append(self.step)
+            yield [{}]
+
+    @property
+    def sigterm_received(self):
+        if self.sigterm_flag:
+            return True
+        if self._sigterm_at is not None and self.step >= self._sigterm_at:
+            self.sigterm_flag = True
+        return self.sigterm_flag
+
+    @property
+    def is_ckpt_step(self):
+        return self.step in self._ckpt or self.sigterm_received
+
+    @property
+    def is_val_step(self):
+        return self.step in self._val
+
+
+def _make_loop_recipe(step_scheduler, *, raise_at=None):
+    recipe = TrainBiEncoderRecipe.__new__(TrainBiEncoderRecipe)
+    recipe.model_parts = [torch.nn.Linear(1, 1)]
+    recipe.step_scheduler = step_scheduler
+    recipe.dataloader = SimpleNamespace(dataset=SimpleNamespace())
+    recipe.val_dataloader = None
+    recipe.max_grad_norm = 1.0
+    recipe.timestamp = 0.0
+    recipe.metric_logger_train = _RecordingMetricLogger()
+    recipe.metric_logger_valid = _RecordingMetricLogger()
+    recipe.saved_at = []
+
+    def _optim_step(batches, max_grad_norm):
+        if raise_at is not None and step_scheduler.step == raise_at:
+            raise RuntimeError("simulated CUDA OOM")
+        return SimpleNamespace(metrics={"loss": 1.0})
+
+    recipe._run_train_optim_step = _optim_step
+    recipe.log_train_metrics = lambda d: recipe.metric_logger_train.log(d)
+    recipe._make_progress_bar = lambda: None
+    recipe._update_progress_bar = lambda pbar, metrics: None
+    recipe._maybe_collect_garbage = lambda: None
+    recipe._finalize_and_close_checkpointer = lambda: None
+    recipe.save_checkpoint = lambda *a, **k: recipe.saved_at.append(step_scheduler.step)
+    return recipe
+
+
+def test_sigterm_on_checkpoint_step_saves_then_stops_the_loop():
+    """A signal on a scheduled checkpoint step: the checkpoint is written, then the
+    post-checkpoint poll stops the loop instead of running further steps."""
+    sched = _LoopStepScheduler(n_steps=10, ckpt_steps={4}, sigterm_at=4)
+    recipe = _make_loop_recipe(sched)
+
+    recipe.run_train_validation_loop()
+
+    assert sched.steps_run == [1, 2, 3, 4], "loop must stop after the signalled step"
+    assert recipe.saved_at == [4], "checkpoint must still be taken"
+    assert recipe.metric_logger_train.closed
+    assert len(recipe.metric_logger_train.persisted) == 4, "every step's metrics must survive"
+
+
+def test_metrics_persisted_when_the_loop_raises():
+    """An exception mid-loop must still persist buffered metrics, because the loggers are
+    closed in the finally rather than after it."""
+    sched = _LoopStepScheduler(n_steps=10)
+    recipe = _make_loop_recipe(sched, raise_at=3)
+
+    with pytest.raises(RuntimeError, match="simulated CUDA OOM"):
+        recipe.run_train_validation_loop()
+
+    assert recipe.metric_logger_train.closed, "logger must be closed from the finally"
+    assert recipe.metric_logger_valid.closed
+    # steps 1 and 2 logged before the raise; both must have reached disk
+    assert len(recipe.metric_logger_train.persisted) == 2
+
+
+def test_checkpoint_flushes_metrics_when_not_buffer_aligned():
+    """The durability guarantee comes from the explicit flush, not from record counting.
+
+    The checkpoint here lands at step 3, which no record-count boundary coincides with --
+    the situation after resuming at a step that is not a multiple of ckpt_every_steps, or at
+    an epoch-boundary checkpoint. The metrics for the covered steps must be on disk by the
+    time the checkpoint is taken.
+    """
+    sched = _LoopStepScheduler(n_steps=6, ckpt_steps={3})
+    recipe = _make_loop_recipe(sched)
+    persisted_at_ckpt = {}
+    original_save = recipe.save_checkpoint
+
+    def _save(*a, **k):
+        original_save(*a, **k)
+        # sample AFTER the recipe's flush by deferring the read to the next step
+        persisted_at_ckpt["at_save"] = len(recipe.metric_logger_train.persisted)
+
+    recipe.save_checkpoint = _save
+    recipe.run_train_validation_loop()
+
+    # the flush runs immediately after save_checkpoint returns, so by step 4 the first three
+    # steps are durable rather than sitting in the buffer
+    assert recipe.saved_at == [3]
+    assert len(recipe.metric_logger_train.persisted) == 6
+    assert persisted_at_ckpt["at_save"] == 0, "flush is after the save, not before"