From 76b401d3b158e04a13be1687fd5109e37bb48168 Mon Sep 17 00:00:00 2001 From: Kang Pu Date: Fri, 17 Apr 2026 14:32:46 +0000 Subject: [PATCH 1/2] feat(adapter/nemo): add keep_mlf_checkpoint_on_train_end flag This change will add a `keep_mlf_checkpoint_on_train_end` flag to MLFlashpointCheckpointCallback. It defaults to False, but when set to True, it skips the deletion of the final checkpoint at the end of training. This ensures the last checkpoint can be preserved for use in E2E testing. Change-Id: Ic238a9dc94fbd9d3c556f0ad1484e277a21fc71e --- .../adapter/nemo/checkpoint_callback.py | 13 ++++- .../adapter/nemo/test_checkpoint_callback.py | 54 +++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/src/ml_flashpoint/adapter/nemo/checkpoint_callback.py b/src/ml_flashpoint/adapter/nemo/checkpoint_callback.py index 6435c7b..cd7cc14 100644 --- a/src/ml_flashpoint/adapter/nemo/checkpoint_callback.py +++ b/src/ml_flashpoint/adapter/nemo/checkpoint_callback.py @@ -55,6 +55,7 @@ def __init__( every_n_steps: int, skip_every_n_steps: Optional[int] = None, enabled: bool = True, + keep_mlf_checkpoint_on_train_end: bool = False, ): """ Initializes and validates the callback. @@ -68,12 +69,15 @@ def __init__( skip_every_n_steps (int, optional): The step frequency to skip checkpointing. This is suggested to be set to the interval used for long-term checkpointing by the alternative strategy. Defaults to 0 (no skipping). enabled (bool): Whether this callback should be enabled. Defaults to True. + keep_mlf_checkpoint_on_train_end (bool): Whether to keep the ML Flashpoint checkpoint after training ends. + Defaults to False. """ self.base_container = CheckpointContainerId(checkpoint_base_container) self.every_n_steps = every_n_steps self.skip_every_n_steps = skip_every_n_steps if skip_every_n_steps is not None else 0 self._enabled = enabled self._replication_manager = None + self.keep_mlf_checkpoint_on_train_end = keep_mlf_checkpoint_on_train_end self._validate() @property @@ -184,5 +188,10 @@ def on_train_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") - self.replication_manager.shutdown() if trainer.local_rank == 0: - _LOGGER.info("Local rank 0: Performing final checkpoint cleanup...") - trainer.strategy.checkpoint_io.remove_checkpoint(self.base_container.data) + if not self.keep_mlf_checkpoint_on_train_end: + _LOGGER.info("Local rank 0: Performing final checkpoint cleanup...") + trainer.strategy.checkpoint_io.remove_checkpoint(self.base_container.data) + else: + _LOGGER.info( + "Local rank 0: Skipping final checkpoint cleanup due to keep_mlf_checkpoint_on_train_end=True." + ) diff --git a/tests/adapter/nemo/test_checkpoint_callback.py b/tests/adapter/nemo/test_checkpoint_callback.py index 114aa41..7704bfa 100644 --- a/tests/adapter/nemo/test_checkpoint_callback.py +++ b/tests/adapter/nemo/test_checkpoint_callback.py @@ -497,3 +497,57 @@ def test_on_train_end_is_idempotent(mocker, tmp_path): # Verify file deletion assert not base_container_path.exists(), "Base container directory should have been deleted" + + +def test_on_train_end_skips_cleanup_when_flag_is_true(mocker, tmp_path): + """ + Tests that the final checkpoint cleanup is skipped when + keep_mlf_checkpoint_on_train_end is set to True. + + This ensures that for E2E tests or specific debugging scenarios, + the last ML Flashpoint checkpoint remains on disk after training ends. + """ + + # Given + trainer = mocker.MagicMock(spec=pl.Trainer) + trainer.local_rank = 0 + chkpt_obj_manager = CheckpointObjectManager() + + checkpoint_io = MLFlashpointCheckpointIO( + flashpoint_base_path=str(tmp_path / "ckpt_base"), + alt_checkpoint_io=mocker.MagicMock(), + chkpt_obj_manager=chkpt_obj_manager, + save_strategy=mocker.MagicMock(), + load_strategy=mocker.MagicMock(), + trainer=trainer, + ) + checkpoint_io.maybe_finalize_save_checkpoint = mocker.MagicMock() + mocker.spy(checkpoint_io, "remove_checkpoint") + trainer.strategy.checkpoint_io = checkpoint_io + + pl_module = mocker.MagicMock(spec=pl.LightningModule) + + # Create a base container directory and a dummy file inside it + base_container_path = tmp_path / "ckpt_base" + base_container_path.mkdir() + dummy_file = base_container_path / "dummy.txt" + dummy_file.write_text("dummy") + + base_container = CheckpointContainerId(str(base_container_path)) + callback = MLFlashpointCheckpointCallback( + checkpoint_base_container=base_container, every_n_steps=1, keep_mlf_checkpoint_on_train_end=True + ) + callback.replication_manager = mocker.MagicMock() + + # When + callback.on_train_end(trainer, pl_module) + + # Then + checkpoint_io.maybe_finalize_save_checkpoint.assert_called_once_with(blocking=True) + trainer.strategy.barrier.assert_called_once_with("mlf_cleanup_barrier") + callback.replication_manager.shutdown.assert_called_once() + + checkpoint_io.remove_checkpoint.assert_not_called() + + assert base_container_path.exists(), "Base container directory should NOT have been deleted" + assert dummy_file.exists(), "Dummy file should NOT have been deleted" From c51b934568b02ccc0303b1698a341273628ba97e Mon Sep 17 00:00:00 2001 From: Kang Pu Date: Fri, 17 Apr 2026 18:02:59 +0000 Subject: [PATCH 2/2] address review comments Change-Id: I5ca3c29cf84ede858c39b280ec263210039d83ea --- src/ml_flashpoint/adapter/nemo/checkpoint_callback.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/ml_flashpoint/adapter/nemo/checkpoint_callback.py b/src/ml_flashpoint/adapter/nemo/checkpoint_callback.py index cd7cc14..385ae69 100644 --- a/src/ml_flashpoint/adapter/nemo/checkpoint_callback.py +++ b/src/ml_flashpoint/adapter/nemo/checkpoint_callback.py @@ -77,7 +77,7 @@ def __init__( self.skip_every_n_steps = skip_every_n_steps if skip_every_n_steps is not None else 0 self._enabled = enabled self._replication_manager = None - self.keep_mlf_checkpoint_on_train_end = keep_mlf_checkpoint_on_train_end + self._keep_mlf_checkpoint_on_train_end = keep_mlf_checkpoint_on_train_end self._validate() @property @@ -188,10 +188,11 @@ def on_train_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") - self.replication_manager.shutdown() if trainer.local_rank == 0: - if not self.keep_mlf_checkpoint_on_train_end: + if not self._keep_mlf_checkpoint_on_train_end: _LOGGER.info("Local rank 0: Performing final checkpoint cleanup...") trainer.strategy.checkpoint_io.remove_checkpoint(self.base_container.data) else: _LOGGER.info( - "Local rank 0: Skipping final checkpoint cleanup due to keep_mlf_checkpoint_on_train_end=True." + "Local rank 0: Skipping final checkpoint cleanup because keep_mlf_checkpoint_on_train_end=%s.", + self._keep_mlf_checkpoint_on_train_end, )