DAT-1069: Resume/pause/stop functionality for training

datamint/client_cmd_tools/datamint_train.py

@@ -180,6 +180,11 @@ def _build_trainer_kwargs(args: argparse.Namespace, alias: str) -> dict[str, Any
         if args.max_epochs is not None:
             kwargs['max_epochs'] = args.max_epochs
         kwargs['configuration'] = '2d'
+        if args.resume is not None:
+            raise DatamintTrainCliError(
+                "--resume is not supported with --model nnunet (nnU-Net manages its own "
+                "checkpointing/resuming). Use the Python SDK's NNUNetTrainer(continue_training=True) instead."
+            )
     else:
         kwargs['max_epochs'] = args.max_epochs if args.max_epochs is not None else DEFAULT_MAX_EPOCHS
         if args.batch_size is not None:
@@ -189,6 +194,8 @@ def _build_trainer_kwargs(args: argparse.Namespace, alias: str) -> dict[str, Any
                 kwargs['image_size'] = args.image_size
             elif alias in DEFAULT_IMAGE_SIZE_FOR:
                 kwargs['image_size'] = DEFAULT_IMAGE_SIZE_FOR[alias]
+        if args.resume is not None:
+            kwargs['resume_from'] = args.resume
 
     if args.model_name is not None:
         kwargs['model_name'] = args.model_name
@@ -228,6 +235,17 @@ def _print_nnunet_notice(console: Console) -> None:
     )
 
 
+def _print_paused(console: Console, project: Project, model_alias: str, results: dict[str, Any]) -> None:
+    run_id = results['run_id']
+    console.print()
+    console.print("[warning]⏸ Training paused.[/warning]")
+    console.print(f"Checkpoint saved to: [key]{results['checkpoint_path']}[/key]")
+    console.print(
+        "Resume with: "
+        f"[key]datamint train --project {project.name!r} --model {model_alias} --resume {run_id}[/key]"
+    )
+
+
 def _print_results(console: Console, trainer, results: dict[str, Any]) -> None:
     console.print()
     console.print("[success]✅ Training finished![/success]")
@@ -346,6 +364,10 @@ def _execute(args: argparse.Namespace, api: Api, console: Console, *, show_plan:
     console.print(f"[bold]Starting training with {trainer_cls.__name__}...[/bold]")
     results = trainer.fit()
 
+    if results.get('paused'):
+        _print_paused(console, project, model_alias, results)
+        return 0
+
     _print_results(console, trainer, results)
 
     if args.show_in_web:
@@ -395,6 +417,10 @@ def _build_parser(subparsers: argparse._SubParsersAction | None = None) -> argpa
                         help='Target image size (square). Ignored for --model nnunet.')
     parser.add_argument('--model-name', type=str, default=None,
                         help='Name to register the trained model under in MLflow.')
+    parser.add_argument('--resume', type=str, default=None, metavar='RUN_ID',
+                        help='Resume a paused run by its MLflow run ID '
+                        '(printed when a previous run was interrupted with Ctrl+C). '
+                        'Not supported with --model nnunet.')
     parser.add_argument('--dry-run', action='store_true',
                         help='Show the detected training plan without training.')
     parser.add_argument('--show-in-web', action='store_true',
@@ -454,6 +480,7 @@ def main() -> None:
         _USER_LOGGER.error(f'❌ {e}')
         sys.exit(1)
     except KeyboardInterrupt:
+        # Only reached for Ctrl+C before trainer.fit() starts.
         CONSOLE.print("\nTraining cancelled by user.", style='warning')
         sys.exit(1)
 

datamint/lightning/trainers/base_trainer.py

@@ -9,12 +9,15 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Mapping
 from functools import cached_property
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar
 
 import lightning as L
 import mlflow
+from lightning.pytorch.trainer.states import TrainerStatus
 from torch import nn
 
+import datamint.configs
 from datamint._repr_utils import render_html_card, render_text_block
 from datamint.dataset.base import DatamintBaseDataset
 from datamint.lightning.datamodule import DatamintDataModule
@@ -71,6 +74,9 @@ class BaseTrainer(ABC):
             adapter after training.
         trainer_kwargs: Extra keyword arguments forwarded to
             :class:`lightning.Trainer`.
+        resume_from: Resume a previously paused run. Either the MLflow
+            ``run_id`` of the paused run or a full path to
+            a checkpoint file.
     """
 
     def __init__(
@@ -92,6 +98,7 @@ def __init__(
         model_name: str | None = None,
         auto_deploy_adapter: bool = True,
         trainer_kwargs: dict[str, Any] | None = None,
+        resume_from: str | None = None,
         **kwargs: Any,
     ) -> None:
         if dataset is None and project is None:
@@ -116,6 +123,7 @@ def __init__(
         self.auto_deploy_adapter = auto_deploy_adapter
         self.trainer_kwargs = trainer_kwargs or {}
         self.trainer_kwargs.update(kwargs)
+        self.resume_from = resume_from
 
         # Populated during fit()
         self._lightning_trainer: L.Trainer | None = None
@@ -185,15 +193,43 @@ def _prepare_pipeline(self) -> None:
         _ = self.datamodule
         _ = self.model
 
+    def _checkpoint_dir(self, run_id: str) -> Path:
+        """Local directory holding the resumable ``last.ckpt`` for *run_id*."""
+        if datamint.configs.DATAMINT_DATA_DIR is None:
+            raise RuntimeError("Could not determine a local data directory (home directory not found); "
+                              "pause/resume checkpointing is unavailable.")
+        return Path(datamint.configs.DATAMINT_DATA_DIR) / 'checkpoints' / run_id
+
+    def _resolve_resume_checkpoint(self) -> tuple[str, str] | None:
+        """Resolve ``resume_from`` into ``(run_id, checkpoint_path)``, or ``None`` if unset. """
+        if self.resume_from is None:
+            return None
+
+        candidate = Path(self.resume_from)
+        if candidate.is_file():
+            return candidate.parent.name, str(candidate)
+
+        run_id = self.resume_from
+        ckpt_path = self._checkpoint_dir(run_id) / 'last.ckpt'
+
+        if not ckpt_path.exists():
+            raise FileNotFoundError(
+                f"No checkpoint found for run '{run_id}' at {ckpt_path}. "
+                "Pass a full checkpoint path via 'resume_from' if it lives elsewhere."
+            )
+        return run_id, str(ckpt_path)
+
     def _create_lightning_trainer(
         self,
         run_id: str | None = None,
         override_params: dict[str, Any] | None = None,
         *,
         register_model: bool = True,
     ) -> L.Trainer:
+
         """Build the configured Lightning trainer instance for this wrapper."""
-        callbacks = self._build_default_callbacks(register_model=register_model) + list(self._build_callbacks())
+        callbacks = (self._build_default_callbacks(register_model=register_model, run_id=run_id)
+                     + list(self._build_callbacks()))
         logger = self._build_logger(run_id=run_id)
 
         trainer_params: dict[str, Any] = {
@@ -208,9 +244,15 @@ def _create_lightning_trainer(
 
         return L.Trainer(**trainer_params)
 
-    def _start_mlflow_run(self):
-        """Start an MLflow run for this trainer workflow."""
+    def _start_mlflow_run(self, resume_run_id: str | None = None):
+        """Start an MLflow run for this trainer workflow.
+
+        When *resume_run_id* is given, reattaches to that existing run.
+        """
+
         self._with_project()
+        if resume_run_id is not None:
+            return mlflow.start_run(run_id=resume_run_id)
         try:
             exp = mlflow.set_experiment(self.experiment_name)
         except mlflow.exceptions.MlflowException as e:
@@ -233,34 +275,64 @@ def _start_mlflow_run(self):
     def fit(self) -> dict[str, Any]:
         """Run the full training pipeline.
 
+        If ``resume_from`` was passed to the constructor, resumes training
+        from that checkpoint in the same MLflow run instead of starting a
+        fresh one.
+
+        If training is interrupted (e.g. Ctrl+C / SIGTERM) before
+        completion, Lightning saves a resumable checkpoint automatically and
+        this method returns early instead of raising. Resume later by
+        passing the returned ``'run_id'`` as ``resume_from`` to a new
+        trainer instance.
+
         Returns:
-            Dictionary with keys ``'trainer'``, ``'model'``,
-            ``'test_results'``, and ``'adapter'`` (when
+            On a completed run: dictionary with keys ``'trainer'``,
+            ``'model'``, ``'test_results'``, and ``'adapter'`` (when
             *auto_deploy_adapter* is enabled).
+            On a paused run: ``{'paused': True, 'run_id': ..., 'checkpoint_path': ...,
+            'trainer': ..., 'model': ...}`` -- ``test()`` and the deploy adapter are
+            not run.
         """
         self._reset_cached_pipeline_state()
 
-        with self._start_mlflow_run() as run:
-            self._prepare_pipeline()
-            self._lightning_trainer = self._create_lightning_trainer(run_id=run.info.run_id)
-
-            # 6. Train
-            _LOGGER.info("Starting training...")
-            self._lightning_trainer.fit(self.model, datamodule=self.datamodule)
+        resume = self._resolve_resume_checkpoint()
+        resume_run_id, ckpt_path = resume if resume is not None else (None, None)
 
-            # 7. Test
-            _LOGGER.info("Starting test...")
-            test_results = self._lightning_trainer.test(datamodule=self.datamodule)
-
-            # 8. Build deploy adapter (only needed when the model is not already a DatamintModel)
-            adapter = None
-            if self.auto_deploy_adapter and not isinstance(self.model, BaseDatamintModel):
-                _LOGGER.debug("Building deploy adapter...")
-                adapter = self._build_deploy_adapter()
-
-            # 9. Upload test predictions as annotations
-            predict_model = self.model if isinstance(self.model, BaseDatamintModel) else adapter
-            self._upload_test_predictions(predict_model)
+        try:
+            with self._start_mlflow_run(resume_run_id=resume_run_id) as run:
+                self._prepare_pipeline()
+                self._lightning_trainer = self._create_lightning_trainer(run_id=run.info.run_id)
+
+                # 6. Train
+                _LOGGER.info("Starting training...")
+                self._lightning_trainer.fit(self.model, datamodule=self.datamodule, ckpt_path=ckpt_path)
+
+                # 7. Test
+                _LOGGER.info("Starting test...")
+                test_results = self._lightning_trainer.test(datamodule=self.datamodule)
+
+                # 8. Build deploy adapter (only needed when the model is not already a DatamintModel)
+                adapter = None
+                if self.auto_deploy_adapter and not isinstance(self.model, BaseDatamintModel):
+                    _LOGGER.debug("Building deploy adapter...")
+                    adapter = self._build_deploy_adapter()
+
+                # 9. Upload test predictions as annotations
+                predict_model = self.model if isinstance(self.model, BaseDatamintModel) else adapter
+                self._upload_test_predictions(predict_model)
+        except SystemExit:
+            if self._lightning_trainer is None or self._lightning_trainer.state.status != TrainerStatus.INTERRUPTED:
+                raise
+            run_id = run.info.run_id
+            checkpoint_path = str(self._checkpoint_dir(run_id) / 'last.ckpt')
+            _LOGGER.info("Training paused (run_id=%s). Resume by passing resume_from=%r.", run_id, run_id)
+            return {
+                'paused': True,
+                'run_id': run_id,
+                'checkpoint_path': checkpoint_path,
+                'trainer': self._lightning_trainer,
+                'model': self.model,
+            }
 
         return {
             'trainer': self._lightning_trainer,
@@ -444,7 +516,7 @@ def _build_datamodule(
             pin_memory=False,
         )
 
-    def _build_default_callbacks(self, *, register_model: bool = True) -> list:
+    def _build_default_callbacks(self, *, register_model: bool = True, run_id: str | None = None) -> list:
         from mlflow.pyfunc.model import PythonModel
 
         from datamint.mlflow.lightning.callbacks import (
@@ -481,10 +553,25 @@ def _build_default_callbacks(self, *, register_model: bool = True) -> list:
 
         callbacks: list = [checkpoint_cls(**checkpoint_kwargs)]
 
+        if run_id is not None:
+            callbacks.append(self._build_resume_checkpoint_callback(run_id))
+
         callbacks.append(_LogDatasetSplitsCallback(self))
 
         return callbacks
 
+    def _build_resume_checkpoint_callback(self, run_id: str):
+        """Plain Lightning checkpoint that maintains a resumable ``last.ckpt``. """
+        from lightning.pytorch.callbacks import ModelCheckpoint
+
+        return ModelCheckpoint(
+            dirpath=str(self._checkpoint_dir(run_id)),
+            filename='last',
+            save_last=True,
+            save_top_k=0,
+            save_on_exception=True,
+        )
+
     def _build_callbacks(self) -> list:
         from lightning.pytorch.callbacks import EarlyStopping
 

docs/source/command_line_tools.rst

@@ -364,6 +364,16 @@ Advanced training options (custom losses, transforms, encoders, ``trainer_kwargs
 are intentionally not exposed here — use the Python SDK instead, see
 :doc:`Training your Model <trainer_api>`.
 
+If you interrupt a run with Ctrl+C, a resumable checkpoint is saved and the run's MLflow
+``run_id`` is printed. Resume it with ``--resume``:
+
+.. code-block:: bash
+
+    datamint train --project MyProject --model yolox --resume <run_id>
+
+``--resume`` is not supported with ``--model nnunet``, which manages its own
+checkpointing/resuming.
+
 See all available options by running ``datamint train --help``.
 
 Running local inference

docs/source/trainer_api.rst

@@ -166,6 +166,28 @@ If you only want evaluation, use ``test()`` instead:
 With ``register_model=True``, the trainer logs and registers the current model.
 
 
+Resuming a Paused Run
+----------------------
+
+If training is interrupted (e.g. Ctrl+C or SIGTERM) before it finishes, Lightning saves a
+resumable checkpoint automatically and the run's MLflow ``run_id`` is logged:
+
+.. code-block:: text
+
+   Training paused (run_id=abcd1234...). Resume by passing resume_from='abcd1234...'.
+
+To resume, pass that ``run_id`` (or a full checkpoint path) as ``resume_from`` when
+constructing a new trainer.
+
+.. code-block:: python
+
+   trainer = UNetPPTrainer(project="BUSI_Segmentation", resume_from="abcd1234...")
+   results = trainer.fit()
+
+Not supported by ``NNUNetTrainer``, which manages its own checkpointing/resuming — use
+``NNUNetTrainer(continue_training=True)`` instead.
+
+
 Passing Lightning Trainer Options
 ---------------------------------