diff --git a/src/ml_flashpoint/adapter/nemo/checkpoint_callback.py b/src/ml_flashpoint/adapter/nemo/checkpoint_callback.py index 6435c7b..385ae69 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,11 @@ 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 because keep_mlf_checkpoint_on_train_end=%s.", + self._keep_mlf_checkpoint_on_train_end, + ) 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"