lightning config for requeing jobs on a preempted slurm partition - #478
Open
edyoshikun wants to merge 4 commits into
Open
lightning config for requeing jobs on a preempted slurm partition#478edyoshikun wants to merge 4 commits into
edyoshikun wants to merge 4 commits into
Conversation
edyoshikun
commented
Jul 1, 2026
edyoshikun
commented
Member
- PR adds adds optional flag to support the requining of slurm jobs in a preemtable SLURM cluster (e.g Reef)
Contributor
|
@edyoshikun Any suggestions on how to test this? |
Member
Author
|
Just tested this on Reef and it worked. I sent SIGUSR1 mid-training -> Lightning caught it, checkpointed, called scontrol requeue, and then the same job was requeued again, and on restart, train.sh found checkpoints/last.ckpt and resumed with the same run ID |
srivarra
approved these changes
Jul 9, 2026
srivarra
left a comment
srivarra
left a comment
Contributor
There was a problem hiding this comment.
Glad it worked out, I think the code looks good as well.
| @@ -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 | |||
Contributor
There was a problem hiding this comment.
do you still need that signal? #SBATCH --signal=B:USR1@300?
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces an opt-in mechanism to auto-requeue Lightning training jobs when running under SLURM on preemptible partitions (e.g., Reef), aiming to checkpoint on SIGUSR1 and resume after SLURM requeues the job.
Changes:
- Added a
--slurm_auto_requeueCLI flag and wiring to attachSLURMEnvironment(auto_requeue=True, requeue_signal=SIGUSR1)when running under SLURM. - Updated the DynaCLR SLURM training launcher to auto-resume from
checkpoints/last.ckptand persist a W&B run id across requeues. - Updated one example SLURM sbatch script to enable requeueing and to emit
SIGUSR1ahead of termination.
Reviewed changes
Copilot reviewed 3 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
packages/viscy-utils/src/viscy_utils/cli.py |
Adds an opt-in CLI flag and config rewrite to attach Lightning’s SLURM auto-requeue environment plugin. |
applications/dynaclr/configs/training/slurm/train.sh |
Adds auto-resume logic and stable W&B run id persistence to support continuity across SLURM requeues. |
applications/dynaclr/configs/training/DynaCLR-2D/DynaCLR-2D-MIP-BagOfChannels.sh |
Enables SLURM requeueing and sends SIGUSR1 before termination for checkpoint/requeue behavior. |
uv.lock |
Updates locked dependencies/extras (includes changes not described by the PR’s stated SLURM-focused scope). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Pin ModelCheckpoint.dirpath to default_root_dir and scope the run's checkpoints/wandb-id/hpc_ckpt by SLURM_JOB_ID so a fresh submission starts clean while a genuine preemption+requeue (same SLURM_JOB_ID) resumes its own state. Append an instantiated SLURMEnvironment rather than a lazy_instance/ class_path spec, which fails Union-slot validation in before_instantiate_classes. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…luster Add docs/clusters/reef.md (CoreWeave SLURM adaptation guide) and split the DynaCLR-2D training configs into per-cluster bruno/ and reef/ subdirectories. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This was referenced Jul 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.