lightning config for requeing jobs on a preempted slurm partition

applications/dynaclr/configs/training/DynaCLR-2D/DynaCLR-2D-MIP-BagOfChannels.sh

@@ -16,6 +16,8 @@
 #SBATCH --cpus-per-task=15
 #SBATCH --mem-per-cpu=8G
 #SBATCH --time=3-00:00:00
+#SBATCH --requeue
+#SBATCH --signal=B:USR1@300
 
 # ── Run identity ──────────────────────────────────────────────────────
 # Warm-started from prior mixed-markers run (s1f8kgtp/last.ckpt, Apr 22)

applications/dynaclr/configs/training/slurm/train.sh

@@ -61,18 +61,41 @@ for cfg in $CONFIGS; do
   CONFIG_FLAGS="${CONFIG_FLAGS} --config ${WORKSPACE_DIR}/${cfg}"
 done
 
+# Auto-resume after SLURM preemption/requeue: if no explicit CKPT_PATH was
+# given but a checkpoint exists in the (stable) run directory, resume from it.
+# A fresh run has no last.ckpt, so the first launch trains from scratch.
+if [ -z "${CKPT_PATH:-}" ] && [ -f "${RUN_DIR}/checkpoints/last.ckpt" ]; then
+  CKPT_PATH="${RUN_DIR}/checkpoints/last.ckpt"
+  echo "Resuming from ${CKPT_PATH}"
+fi
+
 CKPT_FLAG=""
 if [ -n "${CKPT_PATH:-}" ]; then
   CKPT_FLAG="--ckpt_path ${CKPT_PATH}"
 fi
 
-WANDB_ID_FLAG=""
-if [ -n "${WANDB_RUN_ID:-}" ]; then
-  WANDB_ID_FLAG="--trainer.logger.init_args.id=${WANDB_RUN_ID} --trainer.logger.init_args.resume=must"
+# Persist the W&B run id so a requeued job continues the same run (continuous
+# metrics across preemptions). The id is generated once on first launch and
+# reused on every resume. Generating it shell-side (rather than reading it back
+# from wandb) avoids touching logger.experiment in a callback, which can
+# deadlock DDP.
+WANDB_ID_FILE="${RUN_DIR}/.wandb_run_id"
+if [ -z "${WANDB_RUN_ID:-}" ]; then
+  if [ -f "${WANDB_ID_FILE}" ]; then
+    WANDB_RUN_ID="$(cat "${WANDB_ID_FILE}")"
+  else
+    WANDB_RUN_ID="$(python -c 'import secrets; print(secrets.token_hex(4))')"
+    echo "${WANDB_RUN_ID}" > "${WANDB_ID_FILE}"
+  fi
 fi
 
+# resume=allow (not must) so the first launch can create the run; subsequent
+# requeues find the existing id and continue it.
+WANDB_ID_FLAG="--trainer.logger.init_args.id=${WANDB_RUN_ID} --trainer.logger.init_args.resume=allow"
+
 srun uv run --project "$WORKSPACE_DIR" dynaclr fit \
   ${CONFIG_FLAGS} \
+  --slurm_auto_requeue \
   --trainer.default_root_dir="${RUN_DIR}" \
   --trainer.logger.init_args.project="${PROJECT}" \
   --trainer.logger.init_args.name="${RUN_NAME}" \

packages/viscy-utils/src/viscy_utils/cli.py

@@ -4,6 +4,7 @@
 import logging
 import os
 import re
+import signal
 import sys
 import tempfile
 from collections.abc import Callable
@@ -16,6 +17,7 @@
 from lightning.pytorch import LightningDataModule, LightningModule
 from lightning.pytorch.cli import LightningCLI
 from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch.plugins.environments import SLURMEnvironment
 
 from viscy_utils.compose import load_composed_config
 from viscy_utils.trainer import VisCyTrainer
@@ -69,6 +71,40 @@ def _configure_wandb_logger(
         init_args["group"] = base_name
 
 
+def _configure_slurm_requeue(config: Namespace, subcommand: str | None) -> None:
+    """Attach :class:`SLURMEnvironment` with auto-requeue under SLURM batch jobs.
+
+    On a preemptible cluster, SLURM sends ``SIGUSR1`` before killing the job.
+    With :class:`SLURMEnvironment` attached, Lightning catches that signal,
+    writes a checkpoint, and calls ``scontrol requeue`` so the job resumes
+    from the checkpoint when resources free up.
+
+    Opt in with ``--slurm_auto_requeue``; when the flag is absent, normal
+    Lightning behavior is kept. Even when set, the plugin is attached only if
+    ``SLURMEnvironment.detect()`` is true, so passing it off SLURM (or in an
+    interactive allocation without it) is a no-op.
+    """
+    root = config[subcommand] if subcommand is not None else config
+    if not isinstance(root, Namespace):
+        return
+    if not root.get("slurm_auto_requeue", False):
+        return
+    if not SLURMEnvironment.detect():
+        return
+    trainer = root.get("trainer")
+    if not isinstance(trainer, Namespace):
+        return
+    plugins = trainer.get("plugins")
+    if plugins is None:
+        plugins = []
+    elif not isinstance(plugins, list):
+        plugins = [plugins]
+    if any(isinstance(p, SLURMEnvironment) for p in plugins):
+        return
+    plugins.append(SLURMEnvironment(auto_requeue=True, requeue_signal=signal.SIGUSR1))
+    trainer["plugins"] = plugins
+
+
 class VisCyCLI(LightningCLI):
     """Extending lightning CLI arguments and defaults."""
 
@@ -84,12 +120,22 @@ def subcommands() -> dict[str, set[str]]:
         return subcommands
 
     def add_arguments_to_parser(self, parser) -> None:
-        """Set default logger."""
+        """Set default logger and SLURM auto-requeue toggle."""
         parser.set_defaults(
             {
                 "trainer.logger": lazy_instance(WandbLogger),
             }
         )
+        parser.add_argument(
+            "--slurm_auto_requeue",
+            action="store_true",
+            help=(
+                "Opt in to SLURMEnvironment(auto_requeue=True): when running "
+                "under SLURM, preempted jobs checkpoint and requeue "
+                "automatically. Absent (default) keeps normal Lightning "
+                "behavior. No effect off SLURM."
+            ),
+        )
 
     def _parse_ckpt_path(self) -> None:
         # For predict/test/validate: snapshot model init_args before checkpoint
@@ -127,6 +173,7 @@ def _parse_ckpt_path(self) -> None:
     def before_instantiate_classes(self) -> None:
         """Apply shared config rewrites before Lightning object creation."""
         _configure_wandb_logger(self.config, self.subcommand)
+        _configure_slurm_requeue(self.config, self.subcommand)
 
 
 def _setup_environment() -> None: