From dc9b21276072f14d93287aa7b4b97480e199dfba Mon Sep 17 00:00:00 2001 From: Sahel Sharifymoghaddam Date: Fri, 14 Aug 2026 15:50:45 +0000 Subject: [PATCH 1/3] fix(retrieval): do not lose metric records on preemption or error The bi-encoder recipe only closed its metric loggers on the success path, and records are buffered until close(), so a run that was preempted or raised lost every metric it had produced. - close() both loggers in the finally, so an exception cannot skip them - poll StepScheduler.sigterm_received per step; both of its iterators already stop on the flag, but nothing polled it inside the step loop, so it was only checked once per epoch - size the training buffer from ckpt_every_steps, capped at DEFAULT_BUFFER_SIZE - flush both loggers: draining the buffer only reaches the file object, whose own buffer dies with the process Signed-off-by: Sahel Sharifymoghaddam (cherry picked from commit d655b4f8008f45c47470f5c6ffc022227f8167f0) --- .../recipes/retrieval/train_bi_encoder.py | 33 ++++++++++++++++--- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/nemo_automodel/recipes/retrieval/train_bi_encoder.py b/nemo_automodel/recipes/retrieval/train_bi_encoder.py index 5eb9f3f33c..66dec4f8d1 100644 --- a/nemo_automodel/recipes/retrieval/train_bi_encoder.py +++ b/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,21 @@ def materialize_loader(config): ) self._log_model_and_optimizer_details(self.model_parts, self.optimizer, self.lr_scheduler) + # Train metrics reach disk at least as often as the checkpoint they accompany, + # capped at the default so a long checkpoint interval cannot buffer unboundedly. + # flush because a buffered write only reaches the file object, whose own buffer is + # lost 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) @@ -432,12 +446,21 @@ def run_train_validation_loop(self): val_loss=val_loss, ) 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): From 68bd47d73a94ece269066da32930ec6f43f7e77e Mon Sep 17 00:00:00 2001 From: Sahel Sharifymoghaddam Date: Wed, 26 Aug 2026 17:23:01 +0000 Subject: [PATCH 2/3] fix(retrieval): flush metrics at checkpoint boundaries, add regression coverage Addresses review feedback on the metric-durability change. Ruff formatting applied; `ruff format --check` and `ruff check` are clean on every file touched here. CHECKPOINT ALIGNMENT. The reviewer is right, and the previous comment overclaimed. buffer_size=min(DEFAULT_BUFFER_SIZE, ckpt_every_steps) counts RECORDS, and that count only coincides with checkpoint steps in the narrow case of a run starting at step 0 with purely periodic checkpoints. It drifts as soon as training resumes at a step that is not a multiple of ckpt_every_steps -- the buffer restarts empty while the checkpoint cadence does not -- and it never lines up for the checkpoints StepScheduler.is_ckpt_step also fires on: epoch boundaries, the final step, and the sigterm path. In those cases a checkpoint reaches disk while the metrics describing the steps it covers are still buffered, which is exactly what the change set out to prevent. Rather than narrow the claim, make it true: MetricLogger grows an explicit flush(), and the recipe calls it immediately after a successful save_checkpoint. buffer_size is retained but demoted to what it actually is -- a bound on how much can be pending if the process is killed without running any cleanup -- and its comment now says so instead of promising alignment. MetricLoggerDist.flush() carries the same rank-zero guard as close(). close() and flush() share a _drain() helper so their semantics cannot diverge. REGRESSION COVERAGE, CPU only, no distributed init: test_sigterm_on_checkpoint_step_saves_then_stops_the_loop -- a signal arriving on a scheduled checkpoint step still writes the checkpoint, and the post-checkpoint poll then stops the loop rather than running further steps. Asserts the loop ran exactly the steps up to the signal, the checkpoint was taken, and no metrics were lost. test_metrics_persisted_when_the_loop_raises -- the loop raises mid-step and the buffered records still reach disk, covering the move of logger.close() into the finally. test_checkpoint_flushes_metrics_when_not_buffer_aligned -- a checkpoint at a step no record-count boundary coincides with, which is the drift case above. Also pins the ordering: the flush happens after the save, not before. The broader StepScheduler explicit-poll/state-query redesign is left as the follow-up the reviewer described. Worth noting for that work: sigterm_received polls at most once per step (guarded by _sig_polled_step), so a signal that arrives DURING a long checkpoint save is not observed until the next step. The sticky flag means it is never lost, only deferred by one step, and the tests above cover the signal-before-the-step case rather than asserting the during-the-save case works. Co-Authored-By: Claude Opus 5 Signed-off-by: Sahel Sharifymoghaddam --- .../components/loggers/metric_logger.py | 22 ++- .../recipes/retrieval/train_bi_encoder.py | 27 ++- .../test_retrieval_bi_encoder_recipe.py | 157 ++++++++++++++++++ 3 files changed, 198 insertions(+), 8 deletions(-) diff --git a/nemo_automodel/components/loggers/metric_logger.py b/nemo_automodel/components/loggers/metric_logger.py index 549587794a..cfd1bb359b 100644 --- a/nemo_automodel/components/loggers/metric_logger.py +++ b/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 diff --git a/nemo_automodel/recipes/retrieval/train_bi_encoder.py b/nemo_automodel/recipes/retrieval/train_bi_encoder.py index 66dec4f8d1..bc8137ff55 100644 --- a/nemo_automodel/recipes/retrieval/train_bi_encoder.py +++ b/nemo_automodel/recipes/retrieval/train_bi_encoder.py @@ -388,10 +388,13 @@ def materialize_loader(config): ) self._log_model_and_optimizer_details(self.model_parts, self.optimizer, self.lr_scheduler) - # Train metrics reach disk at least as often as the checkpoint they accompany, - # capped at the default so a long checkpoint interval cannot buffer unboundedly. - # flush because a buffered write only reaches the file object, whose own buffer is - # lost with the process. Validation is rare enough to write every point. + # 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) @@ -445,12 +448,24 @@ 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) + logger.info( + "Preemption signal received at step %d; stopping cleanly.", self.step_scheduler.step + ) finally: if pbar is not None: pbar.close() diff --git a/tests/unit_tests/recipes/test_retrieval_bi_encoder_recipe.py b/tests/unit_tests/recipes/test_retrieval_bi_encoder_recipe.py index f7a2f599e5..aa58bd74a8 100644 --- a/tests/unit_tests/recipes/test_retrieval_bi_encoder_recipe.py +++ b/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" From f859ddfc3d692ddc34de8d469d75ef24467b45d5 Mon Sep 17 00:00:00 2001 From: Sahel Sharifymoghaddam Date: Fri, 28 Aug 2026 15:56:03 +0000 Subject: [PATCH 3/3] fix(retrieval): unshadow MetricLogger.flush and make it actually durable Addresses review feedback on #3583. `self.flush = flush` in __init__ replaced the flush() method with a bool on every instance, so `type(logger.flush)` was `bool` and the recipe's first checkpoint raised `TypeError: 'bool' object is not callable`. The flag is now stored as `_fsync_on_write`; the `flush=` keyword is unchanged. Writing a test with a real MetricLogger surfaced a second, deeper problem: flush() drained the record buffer into the file object but only reached the OS when the logger happened to be built with `flush=True`, which is not the default. Under the default the records sat in the process's stdio buffer -- still lost to a crash, and invisible to anything reading the file -- so the guarantee was empty exactly where it is relied on. flush() now flushes and fsyncs unconditionally; the constructor flag continues to govern only whether every buffer-size-triggered write pays for an fsync. test_checkpoint_flushes_metrics_when_not_buffer_aligned did not prove what it claimed: it sampled the count before the flush and then asserted a total that close() in the finally satisfies on its own, so deleting the recipe's flush calls left it green. Replaced with a test that drives a real MetricLogger and reads the real file at the top of each step, asserting the covered steps are durable mid-run, before close(). Each of the three defects above was mutation-checked against the new tests. Narrowed the recipe comment: flushing after the save is a post-save boundary, not an atomic one. A crash between the save completing and the flush leaves the checkpoint on disk with those records still buffered, so the comment no longer claims a checkpoint is never on disk without its metrics. Signed-off-by: Sahel Sharifymoghaddam Co-Authored-By: Claude Opus 5 --- .../components/loggers/metric_logger.py | 18 +++++- .../recipes/retrieval/train_bi_encoder.py | 7 ++- .../unit_tests/loggers/test_metric_logger.py | 52 ++++++++++++++++++ .../test_retrieval_bi_encoder_recipe.py | 55 +++++++++++++------ 4 files changed, 111 insertions(+), 21 deletions(-) diff --git a/nemo_automodel/components/loggers/metric_logger.py b/nemo_automodel/components/loggers/metric_logger.py index cfd1bb359b..0bc6ba273b 100644 --- a/nemo_automodel/components/loggers/metric_logger.py +++ b/nemo_automodel/components/loggers/metric_logger.py @@ -106,7 +106,9 @@ def __init__( if not isinstance(buffer_size, int) or buffer_size < 1: raise ValueError("buffer_size must be a positive integer") self.filepath = os.path.abspath(filepath) - self.flush = flush + # NOT `self.flush`: an instance attribute of that name shadows the flush() method + # below, so callers get a bool back and calling it raises TypeError. + self._fsync_on_write = flush self.buffer_size = buffer_size self.buffer: List[MetricsSample] = [] self._lock = threading.Lock() @@ -133,7 +135,7 @@ def _save(self, lines: List[str]) -> None: if len(lines) == 0: return self._fp.write("\n".join(lines) + "\n") - if self.flush: + if self._fsync_on_write: self._fp.flush() os.fsync(self._fp.fileno()) @@ -143,13 +145,23 @@ def _drain(self) -> None: self.buffer = [] def flush(self) -> None: - """Write buffered records to disk without closing the file. + """Write buffered records out to the file and fsync, without closing it. 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. + + The fsync is unconditional, NOT gated on the ``flush=`` constructor flag. That flag + controls whether every buffer-size-triggered write pays for an fsync; this method is + an explicit, caller-chosen durability point, so it has to be durable under the + default ``flush=False`` too. Draining into the file object alone would leave the + records in the process's stdio buffer -- still lost to a crash, and invisible to + anything reading the file -- which would make the guarantee empty exactly where it + is being relied on. """ with self._lock: self._drain() + self._fp.flush() + os.fsync(self._fp.fileno()) def close(self) -> None: with self._lock: diff --git a/nemo_automodel/recipes/retrieval/train_bi_encoder.py b/nemo_automodel/recipes/retrieval/train_bi_encoder.py index bc8137ff55..f6656941a1 100644 --- a/nemo_automodel/recipes/retrieval/train_bi_encoder.py +++ b/nemo_automodel/recipes/retrieval/train_bi_encoder.py @@ -448,8 +448,11 @@ 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 + # Flush once the checkpoint save returns, so the metrics describing the + # steps it covers are durable by the time the loop moves on rather than + # waiting for the buffer to fill. This is a post-save boundary, not an + # atomic one: a crash between the save completing and this flush leaves + # the checkpoint on disk with those records still buffered. 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 diff --git a/tests/unit_tests/loggers/test_metric_logger.py b/tests/unit_tests/loggers/test_metric_logger.py index 7ad88267f7..d641decc0b 100644 --- a/tests/unit_tests/loggers/test_metric_logger.py +++ b/tests/unit_tests/loggers/test_metric_logger.py @@ -166,3 +166,55 @@ def test_metric_logger_dist_nonzero_noop(tmp_path, monkeypatch): else: # If file was not created, that's also acceptable assert True + + +def test_flush_method_is_callable_and_not_shadowed_by_the_flag(tmp_path): + """The ``flush=`` constructor flag and the ``flush()`` method must coexist. + + Storing the flag as ``self.flush`` replaced the method with a bool on every instance, + so the first caller -- the recipe, at its first checkpoint -- got + ``TypeError: 'bool' object is not callable``. + """ + logger = MetricLogger(str(tmp_path / "metrics.jsonl"), flush=True, append=False) + try: + assert callable(logger.flush) + finally: + logger.close() + + +def test_flush_persists_buffered_records_without_closing(tmp_path): + """flush() is what lets a caller align durability with an external event. + + buffer_size is far larger than the number of records, so nothing reaches the file on + record count alone; only the explicit flush can put them there, and the file must stay + open for further writes afterwards. + """ + logfile = tmp_path / "metrics.jsonl" + logger = MetricLogger(str(logfile), append=False, buffer_size=1000) + try: + logger.log(metric_logger_mod.MetricsSample(step=1, epoch=0, metrics={"loss": 1.0})) + logger.log(metric_logger_mod.MetricsSample(step=2, epoch=0, metrics={"loss": 2.0})) + assert _read_jsonl(logfile) == [], "buffer_size not reached, so nothing is durable yet" + + logger.flush() + rows = _read_jsonl(logfile) + assert [r["step"] for r in rows] == [1, 2] + + # still usable: the flush drained the buffer, it did not close the file + logger.log(metric_logger_mod.MetricsSample(step=3, epoch=0, metrics={"loss": 3.0})) + logger.flush() + assert [r["step"] for r in _read_jsonl(logfile)] == [1, 2, 3] + finally: + logger.close() + + +def test_flush_on_an_empty_buffer_is_a_no_op(tmp_path): + """The recipe flushes at every checkpoint, including ones with nothing buffered.""" + logfile = tmp_path / "metrics.jsonl" + logger = MetricLogger(str(logfile), append=False, buffer_size=1000) + try: + logger.flush() + logger.flush() + assert logfile.read_text() == "" + finally: + logger.close() diff --git a/tests/unit_tests/recipes/test_retrieval_bi_encoder_recipe.py b/tests/unit_tests/recipes/test_retrieval_bi_encoder_recipe.py index aa58bd74a8..06063bd542 100644 --- a/tests/unit_tests/recipes/test_retrieval_bi_encoder_recipe.py +++ b/tests/unit_tests/recipes/test_retrieval_bi_encoder_recipe.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from collections import deque from contextlib import nullcontext from types import SimpleNamespace @@ -20,6 +21,7 @@ import torch from nemo_automodel.components.distributed.config import DDPConfig, FSDP2Config +from nemo_automodel.components.loggers.metric_logger import MetricLogger, MetricsSample from nemo_automodel.recipes.retrieval import train_bi_encoder from nemo_automodel.recipes.retrieval.train_bi_encoder import ( TrainBiEncoderRecipe, @@ -390,29 +392,50 @@ def test_metrics_persisted_when_the_loop_raises(): assert len(recipe.metric_logger_train.persisted) == 2 -def test_checkpoint_flushes_metrics_when_not_buffer_aligned(): +def _durable_steps(path): + """Steps whose records are readable from the file right now, by anything else.""" + if not path.exists(): + return [] + return [json.loads(line)["step"] for line in path.read_text().splitlines() if line.strip()] + + +def test_checkpoint_flushes_metrics_to_disk_mid_run(tmp_path): """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. + Uses a REAL MetricLogger and reads the real file, because the property under test is + whether records are on disk at a point in time -- which a stand-in that records calls + cannot show. It is also the only way this test can fail if the recipe's flush calls are + deleted: the earlier version sampled the count before the flush and then checked a total + that ``close()`` in the finally satisfies on its own, so it passed either way. + + The checkpoint 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. buffer_size is far above the length of the run, so nothing + reaches the file on count alone. """ + logfile = tmp_path / "train_metrics.jsonl" sched = _LoopStepScheduler(n_steps=6, ckpt_steps={3}) recipe = _make_loop_recipe(sched) - persisted_at_ckpt = {} - original_save = recipe.save_checkpoint + recipe.metric_logger_train = MetricLogger(str(logfile), append=False, buffer_size=1000) + recipe.log_train_metrics = lambda d: recipe.metric_logger_train.log( + MetricsSample(step=sched.step, epoch=0, metrics={"loss": 1.0}) + ) + + # Observe the file at the TOP of each step, so step N sees the state left by step N-1. + seen = {} + inner_step = recipe._run_train_optim_step - 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) + def _observing_step(batches, max_grad_norm): + seen[sched.step] = _durable_steps(logfile) + return inner_step(batches, max_grad_norm) - recipe.save_checkpoint = _save + recipe._run_train_optim_step = _observing_step 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" + assert seen[3] == [], "nothing durable before the checkpoint: the buffer is nowhere near full" + # THE ASSERTION THAT BITES: only the post-checkpoint flush can have put these on disk, + # and it is checked mid-run, before close() drains anything. + assert seen[4] == [1, 2, 3], "the checkpoint flush must have made the covered steps durable" + assert seen[6] == [1, 2, 3], "and no further flush happens until the run ends" + assert _durable_steps(logfile) == [1, 2, 3, 4, 5, 6], "close() drains the remainder"