Allow multiple SLURM partitions for structure inference

README.md

@@ -214,7 +214,8 @@ Then open `http://localhost:8501` in your browser.
 Override default values to match your cluster:
 
 ```yaml
-slurm_partition: "gpu"                      # which partition/queue to submit to
+slurm_partition: "gpu"                      # partition(s) to submit inference to; one name,
+                                            # "gpu-el8,gpu-training" or a YAML list for several
 slurm_qos: "normal"                         # optional QoS if your site uses it
 structure_inference_gpus_per_task: 1        # number of GPUs each inference job needs
 structure_inference_gpu_model: ""           # "" lets SLURM pick any GPU in the partition; set a model to pin
@@ -260,8 +261,15 @@ you hit these.
   When set this drives `--exclude` per job and **overrides** `structure_inference_gpu_model` (the two
   would conflict). It's the practical "fit to GPU" lever: requested host RAM is a separate pool and
   does not size GPU VRAM, but excluding too-small GPUs by length does. Use explicit comma node lists
-  (bracket ranges may be glob-expanded by the shell). Multi-partition routing (e.g. EMBL's bigger
-  `gpu-training` cards) is out of scope — keep one partition and let unified memory spill the tail.
+  (bracket ranges may be glob-expanded by the shell). VRAM-tier routing works *within* the listed
+  partition(s); it excludes nodes by name, so if you span partitions make sure the tier node lists
+  cover every partition you submit to.
+- **Span several partitions** by giving `slurm_partition` more than one name — a comma-separated
+  string (`"gpu-el8,gpu-training"`) or a YAML list. The plugin passes them straight to `sbatch -p`,
+  and SLURM starts each inference job on whichever partition frees up first, so jobs aren't stuck
+  behind one busy queue (e.g. spill onto EMBL's bigger `gpu-training` cards). Every listed partition
+  must accept the job's GPUs, `--mem` and walltime (`structure_inference_max_runtime` ≤ each
+  partition's `MaxTime`); a partition the job doesn't fit is simply skipped by SLURM.
 - **Exclude specific nodes** with `slurm_exclude_nodes` → passed verbatim to `sbatch --exclude`
   (e.g. `"gpu50,gpu51"`). Use it as a fallback for nodes whose GPU the container can't use — e.g.
   a CUDA compute capability newer than the container's bundled `ptxas` (fails `ptxas too old` /

config/config.yaml

@@ -140,6 +140,19 @@ batch_max_tokens: 0
 alphafold_inference_threads: 8
 
 # SLURM resources
+# Partition(s) for the GPU structure_inference jobs. Give ONE partition, or SEVERAL
+# to let SLURM start each job on whichever one frees up first (it schedules onto the
+# soonest-available partition). Accepts a single name, a comma/space-separated string,
+# or a YAML list -- all equivalent:
+#   slurm_partition: "gpu-el8"                     # single partition
+#   slurm_partition: "gpu-el8,gpu-training"        # comma-separated
+#   slurm_partition:                               # YAML list
+#     - gpu-el8
+#     - gpu-training
+# All listed partitions must accept the job's resources (GPUs, --mem, and a walltime
+# within each partition's MaxTime -- see structure_inference_max_runtime); a partition
+# the job cannot fit is simply skipped by SLURM. Only structure_inference uses this;
+# the other (CPU) rules run on the cluster's default partition.
 slurm_partition: "gpu-el8"
 slurm_qos: "normal"
 structure_inference_gpus_per_task: 1

test/test_normalize_partitions.py

@@ -0,0 +1,77 @@
+"""Unit tests for ``normalize_partitions`` (multi-partition support).
+
+``common.smk`` is plain Python (stdlib-only imports), so it is loaded by path and
+its functions tested directly. Run with ``python test/test_normalize_partitions.py``
+or ``pytest test/test_normalize_partitions.py``.
+"""
+
+import importlib.machinery
+import importlib.util
+import shlex
+from pathlib import Path
+
+_COMMON = Path(__file__).resolve().parents[1] / "workflow" / "rules" / "common.smk"
+_spec = importlib.util.spec_from_loader(
+    "aps_common", importlib.machinery.SourceFileLoader("aps_common", str(_COMMON))
+)
+common = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(common)
+normalize_partitions = common.normalize_partitions
+
+
+def test_single_partition_unchanged():
+    assert normalize_partitions("gpu-el8") == "gpu-el8"
+
+
+def test_none_and_empty_return_none():
+    assert normalize_partitions(None) is None
+    assert normalize_partitions("") is None
+    assert normalize_partitions("   ") is None
+    assert normalize_partitions([]) is None
+    assert normalize_partitions([" ", ""]) is None
+
+
+def test_comma_separated_string_is_normalized():
+    assert normalize_partitions("gpu-el8,transform") == "gpu-el8,transform"
+    # surrounding whitespace around commas is stripped
+    assert normalize_partitions("gpu-el8, transform ,training") == (
+        "gpu-el8,transform,training"
+    )
+
+
+def test_whitespace_separated_string():
+    assert normalize_partitions("gpu-el8 transform") == "gpu-el8,transform"
+
+
+def test_yaml_list_is_joined_with_commas():
+    assert normalize_partitions(["gpu-el8", "transform"]) == "gpu-el8,transform"
+    assert normalize_partitions(("gpu-el8", "gpu-training")) == "gpu-el8,gpu-training"
+
+
+def test_duplicates_removed_order_preserved():
+    assert normalize_partitions(["gpu-el8", "transform", "gpu-el8"]) == (
+        "gpu-el8,transform"
+    )
+    assert normalize_partitions("transform,gpu-el8,transform") == (
+        "transform,gpu-el8"
+    )
+
+
+def test_list_entries_are_stripped():
+    assert normalize_partitions([" gpu-el8 ", " transform"]) == "gpu-el8,transform"
+
+
+def test_result_survives_shlex_quote_unquoted():
+    # The SLURM plugin runs the partition through shlex.quote before passing it to
+    # `sbatch -p`. A comma-joined list must survive that untouched (comma is a
+    # shell-safe character) so sbatch receives the full partition list.
+    result = normalize_partitions("gpu-el8,transform")
+    assert shlex.quote(result) == result == "gpu-el8,transform"
+
+
+if __name__ == "__main__":
+    for name, fn in sorted(globals().items()):
+        if name.startswith("test_") and callable(fn):
+            fn()
+            print(f"ok  {name}")
+    print("all normalize_partitions tests passed")

workflow/Snakefile

@@ -63,7 +63,11 @@ input_files = config.get("input_files", [])
 if isinstance(input_files, (str, Path)):
     input_files = [input_files]
 
-DEFAULT_SLURM_PARTITION = config.get("slurm_partition")
+# slurm_partition may be a single name, a comma/space-separated string or a YAML
+# list; normalize to one comma-separated string (e.g. "gpu-el8,transform") that
+# sbatch -p understands natively (SLURM runs the job on whichever partition frees
+# up first). See config.yaml for usage.
+DEFAULT_SLURM_PARTITION = normalize_partitions(config.get("slurm_partition"))
 DEFAULT_SLURM_GRES = config.get("slurm_gres")
 DEFAULT_SLURM_QOS = config.get("slurm_qos")
 DEFAULT_STRUCTURE_INFERENCE_GPUS = config.get("structure_inference_gpus_per_task", 1)

workflow/rules/common.smk

@@ -506,6 +506,41 @@ def prepare_container_binds(
         os.environ.setdefault(var, "1")
 
 
+def normalize_partitions(value: Any) -> str | None:
+    """Normalise a ``slurm_partition`` config value to a comma-separated string.
+
+    SLURM's ``sbatch -p`` natively accepts several partitions as a comma list
+    (``-p gpu-el8,transform``) and schedules the job onto whichever one lets it
+    start soonest. This lets a user list every GPU partition they may run on so
+    inference jobs are not stuck behind one busy queue.
+
+    Accepts any of:
+
+    * a YAML list/tuple: ``[gpu-el8, transform]``
+    * a comma- and/or whitespace-separated string: ``"gpu-el8, transform"``
+    * a single partition string: ``"gpu-el8"`` (unchanged)
+    * ``None`` / empty -> ``None`` (caller supplies its own fallback)
+
+    Returns a de-duplicated, order-preserving comma-joined string (no spaces, so
+    it survives ``shlex.quote`` unquoted and reaches ``sbatch`` verbatim), or
+    ``None`` when no partition is given.
+    """
+    if value is None:
+        return None
+    if isinstance(value, (list, tuple, set)):
+        items = list(value)
+    else:
+        # A single scalar; split on commas and any surrounding whitespace so both
+        # "a,b", "a, b" and "a b" are accepted.
+        items = str(value).replace(",", " ").split()
+    names: list[str] = []
+    for item in items:
+        name = str(item).strip()
+        if name and name not in names:
+            names.append(name)
+    return ",".join(names) if names else None
+
+
 def linear_resources(
     *,
     mem: int = 800,