feat(adapter/nemo): add keep_mlf_checkpoint_on_train_end flag (#103)

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.

Author: kkkapu

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,
+                )

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"