Skip to content

Allow multiple SLURM partitions for structure inference#50

Merged
DimaMolod merged 1 commit into
feature/batch-small-jobsfrom
feature/multi-partition-support
Jul 8, 2026
Merged

Allow multiple SLURM partitions for structure inference#50
DimaMolod merged 1 commit into
feature/batch-small-jobsfrom
feature/multi-partition-support

Conversation

@DimaMolod

Copy link
Copy Markdown
Collaborator

Problem

slurm_partition in config/config.yaml accepted only a single partition, so GPU structure_inference jobs were pinned to one queue and stuck waiting whenever that partition was busy — even if another GPU partition (e.g. gpu-training/transform) had idle nodes.

Change

SLURM's sbatch -p natively accepts a comma-separated partition list and starts the job on whichever partition frees up first. This PR lets slurm_partition be several partitions and forwards them to the plugin (which passes them straight to sbatch -p).

  • normalize_partitions() (new, in workflow/rules/common.smk, unit-testable): accepts any of
    • a single name — "gpu-el8" (unchanged behaviour),
    • a comma/space-separated string — "gpu-el8, gpu-training",
    • a YAML list —
      slurm_partition:
        - gpu-el8
        - gpu-training
    and emits a de-duplicated, order-preserving, comma-joined string (gpu-el8,gpu-training). No spaces, so it survives the SLURM plugin's shlex.quote unquoted and reaches sbatch verbatim.
  • Snakefile: DEFAULT_SLURM_PARTITION now runs the config value through normalize_partitions. Only structure_inference uses it; the CPU rules keep using the cluster default partition.
  • Docs: config.yaml + README document multi-partition usage; removed the now-obsolete "multi-partition routing is out of scope" note. Every listed partition must accept the job's GPUs/--mem/walltime (structure_inference_max_runtime ≤ each partition's MaxTime); a partition the job doesn't fit is simply skipped by SLURM.

Testing

  • Unit: test/test_normalize_partitions.py (8 cases: single/list/comma/whitespace inputs, dedup, None/empty, shlex.quote no-op). Full suite: 38 passed.
  • End-to-end on the EMBL cluster: with config slurm_partition: [gpu-el8, transform], Snakemake (slurm executor) built sbatch ... -p gpu-el8,transform ... and SLURM scheduled the job onto the free transform partition (H100 node hgx1), running on GPU to COMPLETED (exit 0). Verified the plugin path is version-agnostic (installed 1.9.2 and workspace 2.0.0 both pass the partition through shlex.quote unchanged).

🤖 Generated with Claude Code

`slurm_partition` in config.yaml accepted only a single partition, so GPU
inference jobs were pinned to one queue. SLURM's `sbatch -p` natively takes a
comma-separated partition list and starts the job on whichever partition frees
up first, so support that here.

Add `normalize_partitions()` in common.smk: it accepts a single name, a
comma/space-separated string, or a YAML list and emits a de-duplicated,
comma-joined string (no spaces, so it survives the plugin's shlex.quote and
reaches sbatch verbatim). The Snakefile runs the config value through it before
setting the `slurm_partition` resource; only structure_inference uses it.

Documented in config.yaml and README; removed the now-obsolete "multi-partition
routing is out of scope" note. Added unit tests (test_normalize_partitions.py).

Verified on the EMBL cluster: config `[gpu-el8, transform]` produced
`sbatch -p gpu-el8,transform`, and SLURM scheduled the job onto the free
`transform` partition (H100 node) to COMPLETED.

Co-Authored-By: Claude Opus 4.8 <[email protected]>
@DimaMolod
DimaMolod merged commit d0529d5 into feature/batch-small-jobs Jul 8, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant