Feature/batch small jobs

README.md

@@ -412,6 +412,37 @@ 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.
+
+`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:

config/config.yaml

@@ -123,6 +123,19 @@ 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
 

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()