lightning config for requeing jobs on a preempted slurm partition

applications/dynaclr/configs/training/DynaCLR-2D/DynaCLR-2D-MIP-BagOfChannels.sh → applications/dynaclr/configs/training/DynaCLR-2D/bruno/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/DynaCLR-2D/reef/DynaCLR-2D-MIP-pretrain-classical.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+# DynaCLR-2D-MIP-pretrain CLASSICAL (SimCLR-style) variant — Reef/Kelp smoke test.
+# Adapted from
+# applications/dynaclr/configs/training/DynaCLR-2D/bruno/DynaCLR-2D-MIP-BagOfChannels-classical.sh
+# per docs/clusters/reef.md: named GPU partition instead of --constraint,
+# explicit --qos (Reef requires one). Uses the requeing_slurm branch's
+# checkpoint+requeue support (train.sh always passes --slurm_auto_requeue;
+# it only attaches SLURMEnvironment(auto_requeue=True) when actually running
+# under SLURM). --requeue/--signal are harmless at --qos dev (non-preemptible)
+# and keep this script ready to resubmit at --qos mid/low unchanged.
+#
+#   sbatch applications/dynaclr/configs/training/DynaCLR-2D/reef/DynaCLR-2D-MIP-pretrain-classical.sh
+
+#SBATCH --job-name=dynaclr_2d_pretrain_smoke
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=2
+#SBATCH --gpus=2
+#SBATCH --partition=h100-reserved
+#SBATCH --qos=dev
+#SBATCH --cpus-per-task=15
+#SBATCH --mem-per-cpu=8G
+#SBATCH --time=1-00:00:00
+#SBATCH --requeue
+#SBATCH --signal=B:USR1@300
+
+export WORKSPACE_DIR="/mnt/main0/home/eduardo.hirata/repos/VisCy"
+export MODEL_ROOT="/bio/projects/compimaging/models"
+export PROJECT="DynaCLR-2D-MIP-pretrain"
+export RUN_NAME="2d-mip-classical-ntxent-t0p2-lr2e5-bs256-192to160-zext11-single-marker-reef-smoke"
+export CONFIGS="applications/dynaclr/configs/training/DynaCLR-2D/reef/DynaCLR-2D-MIP-pretrain.yml applications/dynaclr/configs/training/DynaCLR-2D/reef/DynaCLR-2D-MIP-pretrain-classical.yml"
+
+# Smoke test: cap epochs/batches so we quickly get a checkpoint to validate
+# against (single dirpath, no duplication) and to test SLURM preemption
+# (auto-requeue + resume). Drop this override once the pipeline is validated.
+export EXTRA_ARGS="--trainer.max_epochs=30 --trainer.limit_train_batches=5 --trainer.limit_val_batches=3"
+
+source "${WORKSPACE_DIR}/applications/dynaclr/configs/training/slurm/train.sh"

applications/dynaclr/configs/training/DynaCLR-2D/reef/DynaCLR-2D-MIP-pretrain-classical.yml

@@ -0,0 +1,36 @@
+# Override: SimCLR-style "classical" DynaCLR-2D-MIP-pretrain.
+# Anchor and positive are the same crop; augmentation creates two views.
+# Same mock parquet as the base leaf; only the positive-sampling strategy
+# and batching differ from the (not-yet-written) temporal-positive variant.
+
+data:
+  init_args:
+    cell_index_path: /bio/projects/compimaging/models/collections/DynaCLR-2D-MIP-pretrain-mock.parquet
+    # Self-augmented positives (SimCLR-style). Fully self-supervised:
+    # no lineage/track lookup, no tau scheduling. Augmentation pipeline
+    # provides view diversity. Single-marker batches still supply
+    # same-marker different-cell negatives, preserving discrimination
+    # signal at the marker/perturbation level.
+    positive_cell_source: self
+    positive_match_columns: null
+    positive_channel_source: same
+    # Single-marker batches (OPS strategy) — every batch is one marker,
+    # forcing the model to learn cellular features instead of channel
+    # shortcuts.
+    batch_group_by: marker
+    # Within a marker's draw, balance across the experiments containing
+    # that marker.
+    stratify_by: experiment
+    # Marker-uniform weights, restricted to the 5 markers actually present
+    # in DynaCLR-2D-MIP-pretrain-mock.parquet (G3BP1/TOMM20/SEC61B/
+    # viral_sensor/Phase3D from the 5 mock 2026_* datasets). A marker key
+    # not present in the data is silently ignored; a marker present but
+    # missing from this dict gets sampling weight 0 (see
+    # viscy_data.sampler.FlexibleBatchSampler._precompute_groups) — so
+    # this list must be kept in sync with the collection's marker set.
+    group_weights:
+      Phase3D: 1.0
+      G3BP1: 1.0
+      SEC61B: 1.0
+      TOMM20: 1.0
+      viral_sensor: 1.0

applications/dynaclr/configs/training/DynaCLR-2D/reef/DynaCLR-2D-MIP-pretrain.yml

@@ -0,0 +1,145 @@
+# DynaCLR-2D-MIP-pretrain (Reef/Kelp smoke test)
+# ==============================
+# First DynaCLR-2D-MIP training run on the Reef/Kelp (CoreWeave) cluster.
+# Same 2D bag-of-channels contrastive recipe as
+# applications/dynaclr/configs/training/DynaCLR-2D/bruno/DynaCLR-2D-MIP-BagOfChannels.yml,
+# pointed at a mock collection (5 2026_* A549 infectomics experiments,
+# Globus-copied to /bio) to validate the Reef training path end to end
+# before scaling up to the full production collection.
+#
+# Launch:
+#   sbatch applications/dynaclr/configs/training/DynaCLR-2D/reef/DynaCLR-2D-MIP-pretrain-classical.sh
+#
+# Resume:
+#   CKPT_PATH=.../last.ckpt sbatch .../DynaCLR-2D-MIP-pretrain-classical.sh
+
+base:
+  - ../../recipes/trainer/fit.yml
+  - ../../recipes/topology/ddp_2gpu.yml
+  - ../../recipes/model/contrastive_encoder_convnext_tiny.yml
+
+trainer:
+  precision: bf16-mixed
+  max_epochs: 150
+  limit_train_batches: 800
+  limit_val_batches: 200
+  logger:
+    init_args:
+      project: DynaCLR-2D-MIP-pretrain
+      name: null
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+      init_args:
+        logging_interval: step
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        monitor: loss/val
+        every_n_epochs: 1
+        save_top_k: 5
+        save_last: true
+    - class_path: viscy_utils.callbacks.OnlineEvalCallback
+      init_args:
+        every_n_epochs: 5
+        label_key: perturbation
+        k: 20
+        track_id_key: global_track_id
+        timepoint_key: t
+
+model:
+  init_args:
+    encoder:
+      init_args:
+        in_stack_depth: 1
+        stem_kernel_size: [1, 4, 4]
+        stem_stride: [1, 4, 4]
+        projection_dim: 32
+        drop_path_rate: 0.1
+    loss_function:
+      init_args:
+        temperature: 0.2
+    lr: 0.00002
+    pca_color_keys: "[perturbation,hours_post_perturbation,experiment,marker]"
+    log_negative_metrics_every_n_epochs: 2
+    example_input_array_shape: [1, 1, 1, 160, 160]
+
+data:
+  class_path: dynaclr.data.datamodule.MultiExperimentDataModule
+  init_args:
+    cell_index_path: /bio/projects/compimaging/models/collections/DynaCLR-2D-MIP-pretrain-mock.parquet
+    focus_channel: Phase3D
+    reference_pixel_size_xy_um: 0.1494
+    z_window: 1
+    z_extraction_window: 16
+    z_focus_offset: 0.3
+    yx_patch_size: [256, 256]
+    final_yx_patch_size: [160, 160]
+    channels_per_sample: 1
+    positive_cell_source: lookup
+    positive_match_columns: [lineage_id]
+    positive_channel_source: same
+    tau_range: [0.5, 2.0]
+    tau_decay_rate: 2.0
+    stratify_by: [perturbation, marker]
+    split_ratio: 0.8
+    batch_size: 256
+    num_workers: 4
+    prefetch_factor: 1
+    buffer_size: 1
+    cache_pool_bytes: 0
+    file_io_concurrency: 32
+    seed: 42
+    normalizations:
+      - class_path: viscy_transforms.NormalizeSampled
+        init_args:
+          keys: [channel_0]
+          level: timepoint_statistics
+          subtrahend: mean
+          divisor: std
+    augmentations:
+      - class_path: viscy_transforms.BatchedRandAffined
+        init_args:
+          keys: [channel_0]
+          prob: 0.8
+          scale_range: [[0.8, 1.3], [0.8, 1.3], [0.8, 1.3]]
+          rotate_range: [3.14, 0.0, 0.0]
+          shear_range: [0.05, 0.05, 0.0, 0.05, 0.0, 0.05]
+      - class_path: viscy_transforms.BatchedRandFlipd
+        init_args:
+          keys: [channel_0]
+          spatial_axes: [1, 2]
+          prob: 0.5
+      - class_path: viscy_transforms.BatchedRandAdjustContrastd
+        init_args:
+          keys: [channel_0]
+          prob: 0.5
+          gamma: [0.6, 1.6]
+      - class_path: viscy_transforms.BatchedRandScaleIntensityd
+        init_args:
+          keys: [channel_0]
+          prob: 0.5
+          factors: 0.5
+      - class_path: viscy_transforms.BatchedRandGaussianSmoothd
+        init_args:
+          keys: [channel_0]
+          prob: 0.5
+          sigma_x: [0.25, 0.50]
+          sigma_y: [0.25, 0.50]
+          sigma_z: [0.0, 0.0]
+      - class_path: viscy_transforms.BatchedRandGaussianNoised
+        init_args:
+          keys: [channel_0]
+          prob: 0.5
+          mean: 0.0
+          std: 0.1
+      # Random Z crop: select 10 of 20 extracted slices for Z-invariance.
+      # Must come before ZReduction so MIP sees a variable sub-stack.
+      - class_path: viscy_transforms.BatchedRandSpatialCropd
+        init_args:
+          keys: [channel_0]
+          roi_size: [10, 192, 192]
+      # Z-reduction: MIP for fluorescence, center-slice for label-free.
+      # Must be LAST augmentation (before implicit final spatial crop).
+      - class_path: viscy_transforms.BatchedChannelWiseZReductiond
+        init_args:
+          keys: [channel_0]
+          allow_missing_keys: true

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

@@ -38,7 +38,27 @@ function cleanup() {
 }
 trap cleanup EXIT
 
-mkdir -p "${RUN_DIR}/checkpoints"
+# Scope checkpoints (our own last.ckpt/epoch=*.ckpt, the W&B run id below, AND
+# Lightning's own SLURM auto-requeue hpc_ckpt_*.ckpt -- see below) by
+# SLURM_JOB_ID so each distinct submission gets its own state instead of
+# silently resuming a previous, unrelated run's checkpoint. SLURM preserves
+# the same SLURM_JOB_ID across a genuine preemption+requeue (Restarts
+# increments, the JobID doesn't), so real resumes still find their own
+# checkpoint here; a fresh `sbatch` submission always gets a new SLURM_JOB_ID
+# and starts clean.
+CKPT_DIR="${RUN_DIR}/checkpoints"
+if [ -n "${SLURM_JOB_ID:-}" ]; then
+  CKPT_DIR="${CKPT_DIR}/${SLURM_JOB_ID}"
+fi
+mkdir -p "${CKPT_DIR}"
+
+# #SBATCH --output can't reference $RUN_DIR (SLURM parses #SBATCH directives
+# before this script's env vars exist), so redirect stdout/stderr here once
+# RUN_DIR is known. Everything from this point on (scontrol, srun, dynaclr
+# fit) lands in the run directory instead of the sbatch submission dir.
+if [ -n "${SLURM_JOB_ID:-}" ]; then
+  exec > "${RUN_DIR}/slurm-${SLURM_JOB_ID}.out" 2>&1
+fi
 
 # Rotate existing config.yaml before Lightning overwrites it
 if [ -f "${RUN_DIR}/config.yaml" ]; then
@@ -61,19 +81,51 @@ 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 this job's checkpoint dir, resume from it.
+# A fresh run (new SLURM_JOB_ID) has no last.ckpt, so it trains from scratch.
+if [ -z "${CKPT_PATH:-}" ] && [ -f "${CKPT_DIR}/last.ckpt" ]; then
+  CKPT_PATH="${CKPT_DIR}/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="${CKPT_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="$(uv run --project "$WORKSPACE_DIR" 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"
+
+# default_root_dir=CKPT_DIR (not RUN_DIR): Lightning's own SLURMEnvironment
+# writes/reads its auto-requeue checkpoint (hpc_ckpt_N.ckpt) directly under
+# default_root_dir with no separate config knob -- if this stayed RUN_DIR, a
+# brand new SLURM_JOB_ID would silently auto-resume from a DIFFERENT, earlier
+# job's leftover hpc_ckpt file. Pointing it at the per-job CKPT_DIR keeps that
+# mechanism scoped consistently with our own ModelCheckpoint dirpath (see
+# _configure_checkpoint_dirpath in viscy_utils/cli.py, which now just uses
+# default_root_dir directly). WandbLogger's save_dir stays RUN_DIR (unscoped)
+# since the W&B run itself is meant to persist across preemption+resume.
 srun uv run --project "$WORKSPACE_DIR" dynaclr fit \
   ${CONFIG_FLAGS} \
-  --trainer.default_root_dir="${RUN_DIR}" \
+  --slurm_auto_requeue \
+  --trainer.default_root_dir="${CKPT_DIR}" \
   --trainer.logger.init_args.project="${PROJECT}" \
   --trainer.logger.init_args.name="${RUN_NAME}" \
   --trainer.logger.init_args.save_dir="${RUN_DIR}" \