Feature/batch small jobs

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
@@ -232,6 +233,24 @@ fields keeps the job submission consistent across clusters.
 the default `0` prevents that flag, which avoids conflicting with the Tres-per-task request on many
 systems. Set it to a positive integer only if your site explicitly requires `--ntasks-per-gpu`.
 
+**Multiple partitions.** `slurm_partition` may name more than one partition — as a comma-separated
+string (`"gpu-el8,gpu-training"`) or a YAML list:
+
+```yaml
+slurm_partition:          # inference runs on whichever of these frees up first
+  - gpu-el8
+  - gpu-training
+```
+
+The value is passed straight to `sbatch -p`, and SLURM starts each inference job on whichever listed
+partition can run it soonest, so jobs aren't stuck behind one busy queue (e.g. they spill onto a
+site's larger `gpu-training` cards when the default GPU partition is full). SLURM runs the job on the
+first listed partition that fits its GPUs, `--mem` and walltime and skips the ones that don't (e.g. a
+partition whose `MaxTime` is below `structure_inference_max_runtime`, or with no matching GPU) — so
+make sure **at least one** listed partition can accommodate the job. Only `structure_inference` uses
+this; the other (CPU) rules run on the cluster's default partition. A single name (the default) is
+unchanged.
+
 The remaining optional fields help with two common cluster issues: keeping inference off GPUs it
 can't use, and large complexes running out of GPU memory. Defaults are sensible; expand below only if
 you hit these.
@@ -260,8 +279,9 @@ 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 **multiple partitions** (see above) make
+  sure the tier node lists cover every partition you submit to.
 - **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` /
@@ -412,6 +432,46 @@ length_filter_fetch_uniprot: true     # set false for fully offline runs
 
 </details>
 
+### Batching small jobs into one SLURM job
+
+Many short, inference-only predictions can spend more time waiting in the SLURM queue
+than running. To amortise that wait, several folds can share a single
+`structure_inference` job: the job runs `run_structure_prediction.py` once per fold in a
+loop, so the folds queue **once** between them instead of once each.
+
+```yaml
+batch_size: 4          # max folds per inference job (1 = one job per fold, the default)
+batch_max_tokens: 0    # optional cap on summed residues per batch (0 = no cap)
+```
+
+- Folds are grouped **by size**, so a batch's memory tracks its largest fold and its
+  walltime scales with the number of folds. `batch_max_tokens` keeps a batch's total
+  work within the partition's `MaxTime`; a single oversized fold always runs alone.
+- Works with both AlphaFold2 and AlphaFold3. Each fold is predicted by its own CLI call
+  (a single call with several folds would be **merged into one complex** by the AF3
+  backend), so the per-fold model load is paid each time; a shared
+  `--jax_compilation_cache_dir` is set automatically so later folds reuse earlier
+  compilations (especially AF3 buckets), recovering most of the compile cost.
+- For AlphaFold2 batches, `--allow_resume` is enabled automatically, so if a job is
+  interrupted a rerun skips folds whose outputs already exist (AlphaFold3 does not accept
+  that flag, so its batches recompute the unfinished folds on rerun).
+- Analysis and reports are unaffected — `alphajudge` still runs per fold (one
+  `interfaces.csv` + `report.pdf` each) and the recursive summary still aggregates them.
+- **Trade-off:** a batch is one SLURM job, so a failure reruns the whole batch (minus the
+  folds resume can skip) and the allocation is sized for the batch's largest fold. Keep
+  `batch_size` modest and pair it with `batch_max_tokens` for heterogeneous fold sizes.
+
+> [!NOTE]
+> **AlphaFold3 batching depends on your container + shared filesystem.** A batched AF3 job
+> shares one `--jax_compilation_cache_dir` under `output_directory`. With recent
+> tokamax-based AF3 images this can fail on some cluster setups: the XLA autotune cache
+> write may return `Device or resource busy` on network filesystems (e.g. BeeGFS), and on
+> **H100** the image's bundled tokamax autotuning cache can abort at load. If AF3 jobs
+> crash during compilation, run AF3 with `batch_size: 1` (AlphaFold2 batching is
+> unaffected, and single-fold AF3 avoids the shared cache entirely).
+
+`batch_size: 1` (the default) is exactly the original one-job-per-fold behaviour.
+
 ### Using precomputed features
 
 If you have precomputed protein features, specify the directory:
@@ -465,7 +525,26 @@ structure_inference_arguments:
 
 ### Backend-specific flags
 
-You can pass any backend CLI switches through `structure_inference_arguments`. Common options are listed below; keep or remove lines based on your needs.
+You can pass backend CLI switches through `structure_inference_arguments`. Common options are listed below; keep or remove lines based on your needs.
+
+> [!IMPORTANT]
+> **These flags are backend-exclusive.** `run_structure_prediction.py` validates every flag
+> against the selected `--fold_backend` and aborts the job with
+> `ValueError: The following flags are not supported by backend '<name>'` if you pass one the
+> backend does not accept. Only use flags from **your** backend's list below — e.g.
+> `--allow_resume` is AlphaFold2-only and `--jax_compilation_cache_dir` is AlphaFold3-only.
+> A single wrong flag fails the job immediately (before any prediction runs).
+>
+> When **batching** (`batch_size > 1`) the workflow adds the correct one for you —
+> `--allow_resume` for AlphaFold2, `--jax_compilation_cache_dir` for AlphaFold3 — so you don't
+> set them yourself.
+>
+> The authoritative, always-current list for your image is the backend validation inside the
+> container. Print it with:
+> ```bash
+> singularity exec <prediction_container> run_structure_prediction.py --help
+> ```
+> (`alphalink` accepts the AlphaFold2 flags plus `--crosslinks`.)
 
 <details>
 <summary>AlphaFold2 flags</summary>
@@ -477,7 +556,11 @@ structure_inference_arguments:
   --models_to_relax: None                 # all | best | none
   --remove_keys_from_pickles: True        # strip large tensors from pickle outputs
   --convert_to_modelcif: True             # additionally write ModelCIF files
-  --allow_resume: True                    # resume from partial runs
+  --allow_resume: True                    # resume from partial runs (auto-added when batching)
+  --relax_best_score_threshold: null      # only relax models above this score
+  --threshold_clashes: null               # clash threshold for relaxation
+  --hb_allowance: null                    # H-bond allowance for relaxation
+  --plddt_threshold: null                 # pLDDT cutoff for relaxation
   --num_cycle: 3
   --num_predictions_per_model: 1
   --pair_msa: True
@@ -504,7 +587,7 @@ structure_inference_arguments:
 
 ```yaml
 structure_inference_arguments:
-  --jax_compilation_cache_dir: null
+  --jax_compilation_cache_dir: null       # AF3-only; auto-added when batching
   --buckets: ['64','128','256','512','768','1024','1280','1536','2048','2560','3072','3584','4096','4608','5120']
   --flash_attention_implementation: triton
   --num_diffusion_samples: 5
@@ -514,6 +597,7 @@ structure_inference_arguments:
   --num_recycles: 10
   --save_embeddings: False
   --save_distogram: False
+  --use_ap_style: False                   # shared with AlphaFold2
 ```
 </details>
 

config/config.yaml

@@ -123,10 +123,36 @@ max_protein_length: 0
 # have no local FASTA / cached length yet. Set false for fully offline runs.
 length_filter_fetch_uniprot: true
 
+# Small-job batching (issue #48): group several folds into ONE structure_inference
+# Slurm job to amortise queue wait and per-job model loading. The job runs
+# run_structure_prediction.py once over the whole batch (model loaded once, folds
+# predicted back to back). Folds are grouped by size so a batch's memory tracks its
+# largest fold and its walltime the number of folds. batch_size is the max folds per
+# job (1 = today's one-job-per-fold behaviour, fully unchanged). batch_max_tokens
+# optionally caps the summed residues per batch so total walltime stays within the
+# partition MaxTime; 0 disables that cap (a single oversized fold always runs alone).
+# Trade-off: a failed batch reruns all its folds, so keep batches modest; folds whose
+# outputs already exist are skipped on rerun when the backend supports resuming.
+batch_size: 1
+batch_max_tokens: 0
+
 # Number of threads for AlphaFold inference
 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
+# SLURM runs the job on the first listed partition that fits its GPUs, --mem and
+# walltime and skips those that don't (e.g. a MaxTime below structure_inference_max_runtime,
+# or no matching GPU), so make sure at least one listed partition can accommodate it.
+# Only structure_inference uses this; the other (CPU) rules run on the default partition.
 slurm_partition: "gpu-el8"
 slurm_qos: "normal"
 structure_inference_gpus_per_task: 1

test/test_batch_inference_args.py

@@ -0,0 +1,85 @@
+"""Unit tests for ``batch_inference_args`` (issue #48 batching, backend flag gating).
+
+Regression guard: ``--jax_compilation_cache_dir`` is an AlphaFold3-only (JAX) flag and
+``--allow_resume`` is an AlphaFold2-only flag; ``run_structure_prediction.py`` hard-errors
+on a flag its backend does not accept (``ValueError: not supported by backend '<name>'``).
+Adding the wrong flag broke ALL batched AlphaFold2 inference. These tests pin each flag to
+the backend that accepts it.
+
+``common.smk`` is plain Python, loaded by path. Run with
+``python test/test_batch_inference_args.py`` or ``pytest test/test_batch_inference_args.py``.
+"""
+
+import importlib.machinery
+import importlib.util
+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)
+batch_inference_args = common.batch_inference_args
+
+_CACHE = "/out/.jax_compilation_cache"
+
+
+def test_af2_batch_gets_allow_resume_not_jax_cache():
+    args = batch_inference_args(
+        {"--fold_backend": "alphafold2"},
+        backend="alphafold2",
+        batch_size=2,
+        jax_cache_dir=_CACHE,
+    )
+    assert args["--allow_resume"] == "true"
+    # the AF3-only flag must NOT be present for AF2 (it would ValueError at runtime)
+    assert "--jax_compilation_cache_dir" not in args
+
+
+def test_af3_batch_gets_jax_cache_not_allow_resume():
+    args = batch_inference_args(
+        {"--fold_backend": "alphafold3"},
+        backend="alphafold3",
+        batch_size=2,
+        jax_cache_dir=_CACHE,
+    )
+    assert args["--jax_compilation_cache_dir"] == _CACHE
+    # the AF2-only flag must NOT be present for AF3
+    assert "--allow_resume" not in args
+
+
+def test_batch_size_one_adds_nothing():
+    base = {"--fold_backend": "alphafold3"}
+    for backend in ("alphafold2", "alphafold3"):
+        args = batch_inference_args(
+            base, backend=backend, batch_size=1, jax_cache_dir=_CACHE
+        )
+        assert args == base
+        assert "--allow_resume" not in args
+        assert "--jax_compilation_cache_dir" not in args
+
+
+def test_user_values_are_preserved():
+    # explicit user settings win over the batch defaults (setdefault semantics)
+    args = batch_inference_args(
+        {"--fold_backend": "alphafold3", "--jax_compilation_cache_dir": "/custom"},
+        backend="alphafold3",
+        batch_size=4,
+        jax_cache_dir=_CACHE,
+    )
+    assert args["--jax_compilation_cache_dir"] == "/custom"
+
+
+def test_does_not_mutate_input():
+    base = {"--fold_backend": "alphafold2"}
+    batch_inference_args(base, backend="alphafold2", batch_size=2, jax_cache_dir=_CACHE)
+    assert base == {"--fold_backend": "alphafold2"}
+
+
+if __name__ == "__main__":
+    for name, fn in sorted(globals().items()):
+        if name.startswith("test_") and callable(fn):
+            fn()
+            print(f"ok  {name}")
+    print("all batch_inference_args tests passed")

test/test_bin_folds.py

@@ -0,0 +1,65 @@
+"""Unit tests for ``bin_folds`` (issue #48 small-job batching).
+
+``common.smk`` is plain Python (stdlib-only imports), so it is loaded by path and
+its functions tested directly. Run with ``python test/test_bin_folds.py`` or
+``pytest test/test_bin_folds.py``.
+"""
+
+import importlib.machinery
+import importlib.util
+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)
+bin_folds = common.bin_folds
+
+
+def test_batch_size_one_is_unbatched_and_preserves_order():
+    folds = [("B+C", 800), ("A", 100), ("D", 400)]
+    assert bin_folds(folds, batch_size=1) == [["B+C"], ["A"], ["D"]]
+    # batch_size <= 0 behaves the same (disabled)
+    assert bin_folds(folds, batch_size=0) == [["B+C"], ["A"], ["D"]]
+
+
+def test_groups_by_size_up_to_batch_size():
+    folds = [("big", 900), ("t1", 100), ("t2", 120), ("t3", 110), ("t4", 130)]
+    batches = bin_folds(folds, batch_size=2)
+    # sorted ascending, packed two-per-batch -> small folds cluster together
+    assert batches == [["t1", "t3"], ["t2", "t4"], ["big"]]
+    # every fold appears exactly once
+    assert sorted(f for b in batches for f in b) == ["big", "t1", "t2", "t3", "t4"]
+
+
+def test_max_batch_tokens_caps_summed_work():
+    folds = [("a", 300), ("b", 300), ("c", 300), ("d", 300)]
+    batches = bin_folds(folds, batch_size=10, max_batch_tokens=700)
+    # 300+300 <= 700 but +300 > 700 -> two per batch
+    assert batches == [["a", "b"], ["c", "d"]]
+
+
+def test_single_oversized_fold_still_forms_a_batch():
+    folds = [("huge", 5000), ("small", 50)]
+    batches = bin_folds(folds, batch_size=5, max_batch_tokens=1000)
+    # 'huge' alone exceeds the cap but must not be dropped
+    assert ["huge"] in batches
+    assert sorted(f for b in batches for f in b) == ["huge", "small"]
+
+
+def test_empty_input():
+    assert bin_folds([], batch_size=4) == []
+
+
+def _run():
+    fns = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
+    for fn in fns:
+        fn()
+        print(f"ok  {fn.__name__}")
+    print(f"\n{len(fns)} passed")
+
+
+if __name__ == "__main__":
+    _run()