From d84b14d0e8f539bd142e2931b44f808419e5d3e3 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Mon, 29 Jun 2026 09:51:10 +0200 Subject: [PATCH 1/9] Batch small inference jobs into one Slurm job (issue #48) Group folds into size-binned batches that share a single structure_inference Slurm job. The job runs run_structure_prediction.py once over the whole batch, so the model is loaded/compiled once and the queue is waited on once, instead of per fold. Composition is decided at parse time (sequence lengths are already resolved for length filtering / memory sizing), so no checkpoint is needed. - common.smk: add pure, unit-tested bin_folds() that groups (fold, tokens) pairs by size up to batch_size folds (and an optional batch_max_tokens cap). - Snakefile: compute BATCH_FOLDS / FOLD_BATCH at parse time. A batch's id is its first member, so batch_size=1 yields one singleton batch per fold whose id IS the fold -> identical prediction paths and DAG to before. structure_inference is re-keyed on {batch}: comma-joined --input/--output_directory (absl lists), memory sized from the largest fold, walltime scaled by fold count, GPU tier from the largest fold. analyze_structure stays per fold but depends on the shared batch sentinel via FOLD_BATCH. - config.yaml: add batch_size (default 1, unchanged behaviour) and batch_max_tokens. - Remove workflow/scripts/cluster_sequence_length.py (superseded by bin_folds; it relied on heavyweight AlphaPulldown imports and the injected snakemake object). - test/test_bin_folds.py: unit tests for the binning logic. Verified with snakemake dry-runs: batch_size=1 reproduces the baseline DAG (6/6/6), batch_size=3 yields size-grouped batched jobs with per-fold analysis, complex folds batch correctly (+ within a fold, , between folds), the analysis-off target path resolves to batch sentinels, and non-representative folds depend on their batch's representative sentinel. Co-Authored-By: Claude Opus 4.8 --- config/config.yaml | 13 ++ test/test_bin_folds.py | 65 +++++++++ workflow/Snakefile | 146 ++++++++++++++------ workflow/rules/common.smk | 49 +++++++ workflow/scripts/cluster_sequence_length.py | 119 ---------------- 5 files changed, 234 insertions(+), 158 deletions(-) create mode 100644 test/test_bin_folds.py delete mode 100755 workflow/scripts/cluster_sequence_length.py diff --git a/config/config.yaml b/config/config.yaml index ff20aa6..d3c1ed8 100644 --- a/config/config.yaml +++ b/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 diff --git a/test/test_bin_folds.py b/test/test_bin_folds.py new file mode 100644 index 0000000..8a12bb6 --- /dev/null +++ b/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() diff --git a/workflow/Snakefile b/workflow/Snakefile index 76e80a6..a90891f 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -252,9 +252,77 @@ for fold in kept_folds: _seen_features.add(_json) required_feature_files.append(_json) +# --- Small-job batching (issue #48) ------------------------------------------ +# Group folds into batches that share ONE structure_inference Slurm job. Each +# batch is run by a single run_structure_prediction.py invocation that loads the +# model once and predicts its folds back to back, so a batch pays the queue wait +# and model-load/compile cost once instead of per fold. Composition is decided +# here at parse time (sequence lengths are already resolved above for length +# filtering / memory sizing), so no checkpoint is needed. A batch's id is its +# first member: with batch_size <= 1 every batch is one fold whose id IS the fold, +# giving identical prediction paths and DAG to the unbatched pipeline. Folds are +# grouped by size; batch_max_tokens optionally caps the summed tokens per batch so +# a batch's total walltime stays within the partition limit. +BATCH_SIZE = int(config.get("batch_size", 1) or 1) +BATCH_MAX_TOKENS = int(config.get("batch_max_tokens", 0) or 0) + + +def _fold_tokens(fold): + return fold_total_tokens( + fold, + join(config["output_directory"], "data"), + protein_delimiter, + features_dir=join(config["output_directory"], "features"), + is_af3=(INFERENCE_BACKEND == "alphafold3"), + length_cache=sequence_length_cache, + ) + + +if BATCH_SIZE > 1 and kept_folds: + _fold_batches = bin_folds( + [(fold, _fold_tokens(fold)) for fold in kept_folds], + batch_size=BATCH_SIZE, + max_batch_tokens=BATCH_MAX_TOKENS, + ) +else: + _fold_batches = [[fold] for fold in kept_folds] + +# batch_id -> [fold, ...] and the inverse fold -> batch_id. The batch id is the +# first fold of the batch, so the sentinel predictions//completed_fold.txt +# lives in a real prediction directory (no synthetic dirs); singletons reduce to +# the original predictions//completed_fold.txt. +BATCH_FOLDS = {} +FOLD_BATCH = {} +for _batch_members in _fold_batches: + _batch_id = _batch_members[0] + BATCH_FOLDS[_batch_id] = _batch_members + for _member in _batch_members: + FOLD_BATCH[_member] = _batch_id + + +def batch_prediction_dirs(batch_id): + return [ + join(config["output_directory"], "predictions", fold) + for fold in BATCH_FOLDS[batch_id] + ] + + +def batch_max_tokens_for(batch_id): + """Largest fold in the batch sizes its memory/GPU tier (folds run one at a + time, so peak memory is the max, not the sum).""" + return max((_fold_tokens(fold) for fold in BATCH_FOLDS[batch_id]), default=0) + + +def completed_fold_target(fold): + """Sentinel marking a fold's batch as done (shared by all folds in the batch).""" + return join( + config["output_directory"], "predictions", FOLD_BATCH[fold], "completed_fold.txt" + ) + + required_folds = [ - join(config["output_directory"], "predictions", fold, "completed_fold.txt") - for fold in kept_folds + join(config["output_directory"], "predictions", batch_id, "completed_fold.txt") + for batch_id in BATCH_FOLDS ] ENABLE_STRUCTURE_ANALYSIS = config.get("enable_structure_analysis", True) GENERATE_RECURSIVE_REPORT = config.get("generate_recursive_report", False) @@ -380,15 +448,8 @@ GPU_VRAM_HEADROOM = float(config.get("structure_inference_gpu_vram_headroom", 1. def _inference_slurm_extra(wildcards): """slurm_extra resource: --exclude small-GPU nodes for large complexes, merged - with the static slurm_exclude_nodes.""" - tokens = fold_total_tokens( - wildcards.fold, - join(config["output_directory"], "data"), - protein_delimiter, - features_dir=join(config["output_directory"], "features"), - is_af3=(INFERENCE_BACKEND == "alphafold3"), - length_cache=sequence_length_cache, - ) + with the static slurm_exclude_nodes. Sized from the largest fold in the batch.""" + tokens = batch_max_tokens_for(wildcards.batch) nodes = gpu_exclude_nodes( tokens, GPU_TIERS, @@ -469,33 +530,39 @@ rule create_features: {params.cli_parameters} """ -def lookup_features(wildcards): +def lookup_batch_features(wildcards): # Inputs for inference: generated/precomputed protein features plus any direct # AF3 JSON inputs (e.g. ligands), which are required as the JSON file itself - # rather than as a generated _af3_input.json. - protein_bases, json_basenames = split_fold_inputs(wildcards.fold, protein_delimiter) - feature_files = [feature_name(base) for base in protein_bases] + list(json_basenames) - return [ - join(config["output_directory"], "features", feature_basename) - for feature_basename in feature_files - ] + # rather than as a generated _af3_input.json. Collected over every fold in + # the batch, de-duplicated (folds may share a protein) and order-preserving. + feature_files = [] + seen = set() + for fold in BATCH_FOLDS[wildcards.batch]: + protein_bases, json_basenames = split_fold_inputs(fold, protein_delimiter) + for feature_basename in [feature_name(b) for b in protein_bases] + list(json_basenames): + if feature_basename not in seen: + seen.add(feature_basename) + feature_files.append( + join(config["output_directory"], "features", feature_basename) + ) + return feature_files rule structure_inference: input: - features=lookup_features, + features=lookup_batch_features, output: - join(config["output_directory"], "predictions", "{fold}", "completed_fold.txt"), + join(config["output_directory"], "predictions", "{batch}", "completed_fold.txt"), params: data_directory=config["backend_weights_directory"], feature_directory=join(config["output_directory"], "features"), - output_directory=lambda wildcards: [ - join(config["output_directory"], "predictions", individual_fold) - for individual_fold in wildcards.fold.split(" ") - ], - requested_fold = ( - lambda wc: format_af3_requested_fold(wc.fold, protein_delimiter) - if IS_AF3 - else wc.fold + # Comma-joined because run_structure_prediction.py's --input/--output_directory + # are absl comma-separated lists; one output directory per fold in the batch. + output_directory=lambda wildcards: ",".join( + batch_prediction_dirs(wildcards.batch) + ), + requested_fold=lambda wc: ",".join( + (format_af3_requested_fold(fold, protein_delimiter) if IS_AF3 else fold) + for fold in BATCH_FOLDS[wc.batch] ), protein_delimiter=protein_delimiter, cli_parameters=" ".join( @@ -523,15 +590,10 @@ rule structure_inference: ), tasks_per_gpu=DEFAULT_STRUCTURE_INFERENCE_TASKS_PER_GPU, **linear_resources( + # Peak host RAM is the largest fold in the batch (folds run one at a + # time, memory released between). mem_fn=lambda wildcards, attempt: estimate_inference_mem_mb( - fold_total_tokens( - wildcards.fold, - join(config["output_directory"], "data"), - protein_delimiter, - features_dir=join(config["output_directory"], "features"), - is_af3=(INFERENCE_BACKEND == "alphafold3"), - length_cache=sequence_length_cache, - ), + batch_max_tokens_for(wildcards.batch), base_mb=structure_inference_base_ram, per_token_sq_mb=structure_inference_ram_per_token_sq, scaling=structure_inference_ram_scaling, @@ -539,8 +601,11 @@ rule structure_inference: attempt=attempt, cap_mb=MAX_MEM_MB, ), + # Walltime scales with the number of folds in the batch (they run + # sequentially in one process), capped at the partition MaxTime. runtime_fn=lambda wc, attempt: min( - structure_inference_runtime_minutes * attempt, STRUCTURE_INFERENCE_MAX_RUNTIME + structure_inference_runtime_minutes * len(BATCH_FOLDS[wc.batch]) * attempt, + STRUCTURE_INFERENCE_MAX_RUNTIME, ), ), threads: @@ -577,13 +642,16 @@ rule structure_inference: --features_directory={params.feature_directory} \ {params.cli_parameters} + mkdir -p "$(dirname "{output}")" echo "Completed" > "{output}" """ if ENABLE_STRUCTURE_ANALYSIS: rule analyze_structure: input: - rules.structure_inference.output, + # Depends on the (possibly shared) batch sentinel; analysis still runs + # per fold, reading this fold's own prediction directory. + lambda wildcards: completed_fold_target(wildcards.fold), output: join(config["output_directory"], "predictions", "{fold}", "interfaces.csv"), resources: diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index f5ba68d..410076c 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -563,3 +563,52 @@ def linear_resources( "runtime": _runtime, "attempt": _attempt, } + + +def bin_folds( + fold_tokens: Iterable[tuple[str, int]], + *, + batch_size: int = 1, + max_batch_tokens: int = 0, +) -> list[list[str]]: + """Group folds into batches that share a single inference job (issue #48). + + Each batch is run by one ``run_structure_prediction.py`` invocation, which + loads the model once and predicts the folds back to back. Batching trades + finer-grained retries for far less queue wait and a single model-load per + batch instead of one per fold. + + Folds are sorted by token count so a batch holds similarly sized folds: the + batch's memory is sized from its largest fold and its walltime from the sum, + so keeping sizes close stops a tiny fold from inheriting a huge fold's + allocation (and clusters the many small folds the issue is about). The number + of folds per batch is capped by ``batch_size``; the optional + ``max_batch_tokens`` additionally caps the summed tokens per batch so total + walltime stays within the partition limit. A single fold always forms a valid + batch even if it alone exceeds ``max_batch_tokens``. + + ``batch_size <= 1`` returns one fold per batch in the original input order, + i.e. the unbatched behaviour, so the default path is unchanged. + """ + items = [(str(fold), int(tokens or 0)) for fold, tokens in fold_tokens] + if batch_size <= 1: + return [[fold] for fold, _ in items] + + ordered = sorted(items, key=lambda ft: (ft[1], ft[0])) + cap = int(max_batch_tokens or 0) + + batches: list[list[str]] = [] + current: list[str] = [] + current_tokens = 0 + for fold, tokens in ordered: + too_many = len(current) >= batch_size + too_big = cap > 0 and bool(current) and (current_tokens + tokens) > cap + if current and (too_many or too_big): + batches.append(current) + current = [] + current_tokens = 0 + current.append(fold) + current_tokens += tokens + if current: + batches.append(current) + return batches diff --git a/workflow/scripts/cluster_sequence_length.py b/workflow/scripts/cluster_sequence_length.py deleted file mode 100755 index a0cd6e9..0000000 --- a/workflow/scripts/cluster_sequence_length.py +++ /dev/null @@ -1,119 +0,0 @@ -#!python3 -import argparse - -from alphapulldown.utils.modelling_setup import ( - parse_fold, - create_custom_info, - create_interactors, -) -from alphapulldown.objects import MultimericObject - - -def parse_args(): - parser = argparse.ArgumentParser( - description="Cluster folds by sequence length." - ) - parser.add_argument( - "--folds", - dest="folds", - type=str, - nargs="+", - default=None, - required=False, - help="All the jobs to fold", - ) - parser.add_argument( - "--protein_delimiter", - dest="protein_delimiter", - type=str, - default="+", - required=False, - help="protein list files", - ) - parser.add_argument( - "--features_directory", - dest="features_directory", - type=str, - nargs="+", - required=False, - help="Path to computed monomer features.", - ) - parser.add_argument( - "--bin_size", - dest="bin_size", - type=int, - required=False, - default=150, - help="Bin size used for clustering sequences.", - ) - parser.add_argument( - "--output_file", - dest="output_file", - type=str, - default="sequence_clusters.txt", - required=False, - help="Path to comma separated output file.", - ) - args = parser.parse_args() - - try: - args.folds = snakemake.params.folds - args.output_file = snakemake.output[0] - args.protein_delimiter = snakemake.params.protein_delimiter - args.bin_size = snakemake.params.cluster_bin_size - args.features_directory = [snakemake.params.feature_directory, ] - except Exception: - pass - - if args.folds is None: - raise ValueError("--folds needs to be specified.") - - return args - - -def main(): - args = parse_args() - - all_jobs = {"name": [], "msa_depth": [], "seq_length": []} - for idx, i in enumerate(args.folds): - parsed_input = parse_fold( - [i], args.features_directory, args.protein_delimiter - ) - - data = create_custom_info(parsed_input) - interactors = create_interactors(data, args.features_directory, 0) - multimer = MultimericObject(interactors[0]) - - msa_depth, seq_length = multimer.feature_dict["msa"].shape - all_jobs["name"].append(i) - all_jobs["msa_depth"].append(msa_depth) - all_jobs["seq_length"].append(seq_length) - - # Assign elements to bins - min_seq_length = max(min(all_jobs["seq_length"]), 1) - all_jobs["cluster"] = [ - int((value - min_seq_length) // args.bin_size) for value in all_jobs["seq_length"] - ] - label_stats = {} - for index, label in enumerate(all_jobs["cluster"]): - if label not in label_stats: - label_stats[label] = {"max_seq_length" : 0, "max_msa_depth" : 0} - label_stats[label]["max_seq_length"] = max( - label_stats[label]["max_seq_length"], all_jobs["seq_length"][index] - ) - label_stats[label]["max_msa_depth"] = max( - label_stats[label]["max_msa_depth"], all_jobs["msa_depth"][index] - ) - - all_jobs["max_msa_depth"], all_jobs["max_seq_length"] = [], [] - for label in all_jobs["cluster"]: - all_jobs["max_msa_depth"].append(label_stats[label]["max_msa_depth"]) - all_jobs["max_seq_length"].append(label_stats[label]["max_seq_length"]) - - with open(args.output_file, mode = "w", encoding = "utf-8") as ofile: - ofile.write(','.join([str(x) for x in list(all_jobs.keys())]) + "\n") - for fold in zip(*all_jobs.values()): - _ = ofile.write(','.join([str(x) for x in fold]) + "\n") - -if __name__ == "__main__": - main() From ef53d4a03cf7010967341b044bfc2787e04e8fcd Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Mon, 29 Jun 2026 10:06:00 +0200 Subject: [PATCH 2/9] Force --allow_resume for batched inference + document batching - Snakefile: when batch_size>1, ensure --allow_resume is on so a rerun of a crashed batch skips folds whose outputs already exist (both AF2 and AF3 backends honour it per fold). An explicit user value is left untouched. - README: add "Batching small jobs into one SLURM job" section. Validated end to end (no GPU): a real Snakemake run batched two AF3 complexes into one structure_inference job (comma-joined --input/--output_directory fanned out to both fold dirs), the non-representative fold's analyze_structure chained off the representative batch sentinel, and real AlphaJudge produced both interfaces.csv + report.pdf plus the aggregated summary. Co-Authored-By: Claude Opus 4.8 --- README.md | 26 ++++++++++++++++++++++++++ workflow/Snakefile | 16 +++++++++++++--- 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 5790979..bc371dc 100644 --- a/README.md +++ b/README.md @@ -412,6 +412,32 @@ length_filter_fetch_uniprot: true # set false for fully offline runs +### 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 — and the one-off model load/compile each job pays — +several folds can share a single `structure_inference` job that predicts them back to +back in one `run_structure_prediction.py` call (the model is loaded once): + +```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. `--allow_resume` is enabled automatically + for batches, so if a job is interrupted, a rerun skips folds whose outputs already + exist and only finishes the rest. +- 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: diff --git a/workflow/Snakefile b/workflow/Snakefile index a90891f..c549b20 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -266,6 +266,18 @@ for fold in kept_folds: BATCH_SIZE = int(config.get("batch_size", 1) or 1) BATCH_MAX_TOKENS = int(config.get("batch_max_tokens", 0) or 0) +# Inference CLI args. When batching, a crashed job re-runs the whole batch, so make +# sure the backend resumes folds whose outputs already exist instead of recomputing +# them (both AF2 and AF3 backends honour --allow_resume per fold). It defaults to +# True upstream; we set it explicitly so a batched run never silently recomputes +# finished folds. A user value in structure_inference_arguments is left untouched. +_structure_inference_args = dict(config.get("structure_inference_arguments", {})) +if BATCH_SIZE > 1: + _structure_inference_args.setdefault("--allow_resume", "true") +STRUCTURE_INFERENCE_CLI = " ".join( + f"{k}={v}" for k, v in _structure_inference_args.items() +) + def _fold_tokens(fold): return fold_total_tokens( @@ -565,9 +577,7 @@ rule structure_inference: for fold in BATCH_FOLDS[wc.batch] ), protein_delimiter=protein_delimiter, - cli_parameters=" ".join( - [f"{k}={v}" for k, v in config["structure_inference_arguments"].items()] - ), + cli_parameters=STRUCTURE_INFERENCE_CLI, unified_memory="true" if STRUCTURE_INFERENCE_UNIFIED_MEMORY else "false", xla_mem_fraction=STRUCTURE_INFERENCE_XLA_MEM_FRACTION, resources: From 59d55efd9ccbe144c0ca8668de798e14ac9afe1d Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Mon, 29 Jun 2026 14:56:15 +0200 Subject: [PATCH 3/9] Batch as a per-fold loop in one job; fix AF3 + Slurm-executor bugs Real GPU validation on SLURM (AF3, two complexes, batch_size=2) surfaced three issues the local/dry/no-GPU tests could not: 1. Parse-stable batch grouping. The Slurm executor re-parses the Snakefile on the compute node, where features are already staged, so size-grouping read from live token counts diverged from the login-node grouping that named the job -> the job's {batch} wildcard was an unknown batch id (KeyError). Group on the persisted sequence-length cache only (identical at every parse); empty cache -> group by name. Added _fold_grouping_tokens. 2. AF3 rejects --allow_resume ("not supported by backend 'alphafold3'"). Inject it for AlphaFold2 only. 3. AF3 merges folds passed together in one CLI call into a single combined complex (alphafold3_backend.predict: "merging ... into a single job"). The original comma-joined --input therefore produced one wrong N-chain complex for AF3. Redesign structure_inference to run run_structure_prediction.py ONCE PER FOLD in a shell loop sharing the one allocation -- correct for AF2 and AF3, still pays the queue wait once (issue #48's main goal). A shared --jax_compilation_cache_dir is set for batches so later folds reuse earlier compilations (esp. AF3 buckets). Verified end to end on GPU: one batched job predicted both folds SEPARATELY (two 2-chain complexes, not a 4-chain merge), AlphaJudge wrote per-fold interfaces.csv + report.pdf, and the recursive summary aggregated both. README updated. Co-Authored-By: Claude Opus 4.8 --- README.md | 17 +++++--- workflow/Snakefile | 104 +++++++++++++++++++++++++++++++-------------- 2 files changed, 83 insertions(+), 38 deletions(-) diff --git a/README.md b/README.md index bc371dc..a7c5e8e 100644 --- a/README.md +++ b/README.md @@ -415,9 +415,9 @@ length_filter_fetch_uniprot: true # set false for fully offline runs ### 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 — and the one-off model load/compile each job pays — -several folds can share a single `structure_inference` job that predicts them back to -back in one `run_structure_prediction.py` call (the model is loaded once): +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) @@ -427,9 +427,14 @@ 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. `--allow_resume` is enabled automatically - for batches, so if a job is interrupted, a rerun skips folds whose outputs already - exist and only finishes the rest. +- 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 diff --git a/workflow/Snakefile b/workflow/Snakefile index c549b20..0955030 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -253,27 +253,39 @@ for fold in kept_folds: required_feature_files.append(_json) # --- Small-job batching (issue #48) ------------------------------------------ -# Group folds into batches that share ONE structure_inference Slurm job. Each -# batch is run by a single run_structure_prediction.py invocation that loads the -# model once and predicts its folds back to back, so a batch pays the queue wait -# and model-load/compile cost once instead of per fold. Composition is decided -# here at parse time (sequence lengths are already resolved above for length -# filtering / memory sizing), so no checkpoint is needed. A batch's id is its -# first member: with batch_size <= 1 every batch is one fold whose id IS the fold, -# giving identical prediction paths and DAG to the unbatched pipeline. Folds are -# grouped by size; batch_max_tokens optionally caps the summed tokens per batch so -# a batch's total walltime stays within the partition limit. +# Group folds into batches that share ONE structure_inference Slurm job. The job +# runs run_structure_prediction.py once per fold (see the rule's shell loop), so the +# folds wait in the queue once between them instead of once each -- the main cost for +# many short inference jobs. A per-fold loop (rather than one call with all folds) is +# required because the AlphaFold3 backend merges folds passed together into a single +# combined complex; a shared --jax_compilation_cache_dir recovers most of the +# per-call compile cost. Composition is decided here at parse time (sequence lengths +# are already resolved above for length filtering / memory sizing), so no checkpoint +# is needed. A batch's id is its first member: with batch_size <= 1 every batch is one +# fold whose id IS the fold, giving identical prediction paths and DAG to the unbatched +# pipeline. Folds are grouped by size; batch_max_tokens optionally caps the summed +# tokens per batch so a batch's total walltime stays within the partition limit. BATCH_SIZE = int(config.get("batch_size", 1) or 1) BATCH_MAX_TOKENS = int(config.get("batch_max_tokens", 0) or 0) -# Inference CLI args. When batching, a crashed job re-runs the whole batch, so make -# sure the backend resumes folds whose outputs already exist instead of recomputing -# them (both AF2 and AF3 backends honour --allow_resume per fold). It defaults to -# True upstream; we set it explicitly so a batched run never silently recomputes -# finished folds. A user value in structure_inference_arguments is left untouched. +# Inference CLI args (shared by every per-fold call in a batch). Two batch tweaks, +# both leaving any user-set value untouched: +# * --allow_resume: a crashed batch re-runs all its folds, so resume the ones whose +# outputs already exist. Only AlphaFold2 accepts it; AlphaFold3 rejects unknown +# flags (ValueError "not supported by backend 'alphafold3'"), so AF2 only. +# * --jax_compilation_cache_dir: the batch runs one run_structure_prediction.py per +# fold (the AF3 backend MERGES folds passed together into one complex, so they must +# be predicted separately), which would otherwise recompile per fold. A shared +# on-disk JAX cache lets later folds reuse earlier compilations (esp. AF3 buckets), +# recovering most of the per-job compile cost the single-process approach would have. _structure_inference_args = dict(config.get("structure_inference_arguments", {})) if BATCH_SIZE > 1: - _structure_inference_args.setdefault("--allow_resume", "true") + if INFERENCE_BACKEND == "alphafold2": + _structure_inference_args.setdefault("--allow_resume", "true") + _structure_inference_args.setdefault( + "--jax_compilation_cache_dir", + join(config["output_directory"], ".jax_compilation_cache"), + ) STRUCTURE_INFERENCE_CLI = " ".join( f"{k}={v}" for k, v in _structure_inference_args.items() ) @@ -290,9 +302,25 @@ def _fold_tokens(fold): ) +def _fold_grouping_tokens(fold): + """Parse-stable size proxy for batching: read length ONLY from the persisted + sequence-length cache, never from staged features. The Slurm executor re-parses + the Snakefile on the compute node to run a single rule, where features are + already staged; reading them would change the size-based grouping (and thus the + batch id that named the job) versus the login-node parse that submitted it, + yielding a "no such batch" KeyError. The cache (`.sequence_lengths.tsv`, written + once at parse time when length resolution runs) is identical at every parse. When + it is empty (length filtering disabled) all folds score 0 and group by name -- + still deterministic, just not size-aware.""" + total = 0 + for name, copies in parse_fold_chains(fold, protein_delimiter): + total += int(sequence_length_cache.get(name, 0) or 0) * copies + return total + + if BATCH_SIZE > 1 and kept_folds: _fold_batches = bin_folds( - [(fold, _fold_tokens(fold)) for fold in kept_folds], + [(fold, _fold_grouping_tokens(fold)) for fold in kept_folds], batch_size=BATCH_SIZE, max_batch_tokens=BATCH_MAX_TOKENS, ) @@ -567,15 +595,18 @@ rule structure_inference: params: data_directory=config["backend_weights_directory"], feature_directory=join(config["output_directory"], "features"), - # Comma-joined because run_structure_prediction.py's --input/--output_directory - # are absl comma-separated lists; one output directory per fold in the batch. - output_directory=lambda wildcards: ",".join( - batch_prediction_dirs(wildcards.batch) - ), - requested_fold=lambda wc: ",".join( + # Parallel space-separated lists (one entry per fold in the batch) consumed by + # the shell loop below. The shell calls run_structure_prediction.py ONCE PER + # fold rather than passing them together, because the AlphaFold3 backend merges + # folds given in a single call into one combined complex. Fold specs and output + # paths never contain spaces, so space-splitting in bash is safe. + requested_folds=lambda wc: " ".join( (format_af3_requested_fold(fold, protein_delimiter) if IS_AF3 else fold) for fold in BATCH_FOLDS[wc.batch] ), + fold_output_dirs=lambda wildcards: " ".join( + batch_prediction_dirs(wildcards.batch) + ), protein_delimiter=protein_delimiter, cli_parameters=STRUCTURE_INFERENCE_CLI, unified_memory="true" if STRUCTURE_INFERENCE_UNIFIED_MEMORY else "false", @@ -601,7 +632,7 @@ rule structure_inference: tasks_per_gpu=DEFAULT_STRUCTURE_INFERENCE_TASKS_PER_GPU, **linear_resources( # Peak host RAM is the largest fold in the batch (folds run one at a - # time, memory released between). + # time in the loop, memory released between calls). mem_fn=lambda wildcards, attempt: estimate_inference_mem_mb( batch_max_tokens_for(wildcards.batch), base_mb=structure_inference_base_ram, @@ -612,7 +643,7 @@ rule structure_inference: cap_mb=MAX_MEM_MB, ), # Walltime scales with the number of folds in the batch (they run - # sequentially in one process), capped at the partition MaxTime. + # sequentially, one CLI call each), capped at the partition MaxTime. runtime_fn=lambda wc, attempt: min( structure_inference_runtime_minutes * len(BATCH_FOLDS[wc.batch]) * attempt, STRUCTURE_INFERENCE_MAX_RUNTIME, @@ -644,13 +675,22 @@ rule structure_inference: fi export XLA_CLIENT_MEM_FRACTION="$_aj_frac" fi - run_structure_prediction.py \ - --input {params.requested_fold} \ - --output_directory={params.output_directory} \ - --protein_delimiter={params.protein_delimiter} \ - --data_directory={params.data_directory} \ - --features_directory={params.feature_directory} \ - {params.cli_parameters} + # One run_structure_prediction.py call per fold in the batch (a single call + # with several folds would be merged into one complex by the AF3 backend). + # The folds share this one Slurm allocation -- the queue is waited on once -- + # and a shared --jax_compilation_cache_dir lets later folds reuse earlier + # compilations. Arrays split on spaces (fold specs / paths contain none). + read -ra _aj_inputs <<< "{params.requested_folds}" + read -ra _aj_outdirs <<< "{params.fold_output_dirs}" + for _aj_i in "${{!_aj_inputs[@]}}"; do + run_structure_prediction.py \ + --input "${{_aj_inputs[$_aj_i]}}" \ + --output_directory="${{_aj_outdirs[$_aj_i]}}" \ + --protein_delimiter={params.protein_delimiter} \ + --data_directory={params.data_directory} \ + --features_directory={params.feature_directory} \ + {params.cli_parameters} + done mkdir -p "$(dirname "{output}")" echo "Completed" > "{output}" From 0126ec3c877361c7d1606611308a083a9ffdf150 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 8 Jul 2026 10:01:46 +0200 Subject: [PATCH 4/9] Allow multiple SLURM partitions for structure inference `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 --- README.md | 14 ++++-- config/config.yaml | 13 ++++++ test/test_normalize_partitions.py | 77 +++++++++++++++++++++++++++++++ workflow/Snakefile | 6 ++- workflow/rules/common.smk | 35 ++++++++++++++ 5 files changed, 141 insertions(+), 4 deletions(-) create mode 100644 test/test_normalize_partitions.py diff --git a/README.md b/README.md index a7c5e8e..cb8ea7a 100644 --- a/README.md +++ b/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` / diff --git a/config/config.yaml b/config/config.yaml index d3c1ed8..80886a9 100644 --- a/config/config.yaml +++ b/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 diff --git a/test/test_normalize_partitions.py b/test/test_normalize_partitions.py new file mode 100644 index 0000000..caa237b --- /dev/null +++ b/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") diff --git a/workflow/Snakefile b/workflow/Snakefile index 0955030..8b58afb 100644 --- a/workflow/Snakefile +++ b/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) diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 410076c..80bf4aa 100644 --- a/workflow/rules/common.smk +++ b/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, From d726e9c6f4638b5f3b761d88d1a8b8923b89ffb2 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 8 Jul 2026 10:20:09 +0200 Subject: [PATCH 5/9] Docs: promote multi-partition to a visible subsection, fix wording The multiple-partitions feature was only described in a collapsed "avoiding unsuitable GPUs" details block, and both the README bullet and config.yaml comment were self-contradictory ("every listed partition must accept the job" vs "a partition the job doesn't fit is skipped"). Give it its own visible subsection under "SLURM defaults for structure inference" with a YAML example, and state the rule precisely: SLURM runs the job on the first listed partition that fits and skips the rest, so at least one must accommodate it. Co-Authored-By: Claude Opus 4.8 --- README.md | 28 ++++++++++++++++++++-------- config/config.yaml | 8 ++++---- 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index cb8ea7a..9ef40f3 100644 --- a/README.md +++ b/README.md @@ -233,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. @@ -262,14 +280,8 @@ you hit these. 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). 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. + 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` / diff --git a/config/config.yaml b/config/config.yaml index 80886a9..b7be0bf 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -149,10 +149,10 @@ alphafold_inference_threads: 8 # 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 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 From a2b941b92e26e38f996b7084578e9a88a46d8632 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 8 Jul 2026 15:53:53 +0200 Subject: [PATCH 6/9] Fix: don't pass --jax_compilation_cache_dir to the AlphaFold2 backend Batched inference (batch_size > 1) added --jax_compilation_cache_dir for every backend, but that flag is AlphaFold3-only (a JAX/XLA compile-cache path). run_structure_prediction.py validates flags per backend and hard-errors on any it doesn't accept, so every batched AlphaFold2 job died immediately with `ValueError: The following flags are not supported by backend 'alphafold2': ['jax_compilation_cache_dir']`. AF2 batching was completely broken. Gate --jax_compilation_cache_dir to AlphaFold3 only, mirroring how --allow_resume is already gated to AlphaFold2 only. Extract the gating into a testable common.smk helper (batch_inference_args) and add regression tests pinning each batch-only flag to the backend that accepts it. Found by an end-to-end AF2 run (batch_size 2, slurm_partition [gpu-el8,transform]) which now completes: both batched jobs COMPLETED and all four folds produced structures. Co-Authored-By: Claude Opus 4.8 --- test/test_batch_inference_args.py | 85 +++++++++++++++++++++++++++++++ workflow/Snakefile | 28 ++++------ workflow/rules/common.smk | 32 ++++++++++++ 3 files changed, 127 insertions(+), 18 deletions(-) create mode 100644 test/test_batch_inference_args.py diff --git a/test/test_batch_inference_args.py b/test/test_batch_inference_args.py new file mode 100644 index 0000000..7afad76 --- /dev/null +++ b/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 ''``). +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") diff --git a/workflow/Snakefile b/workflow/Snakefile index 8b58afb..43eb749 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -272,24 +272,16 @@ for fold in kept_folds: BATCH_SIZE = int(config.get("batch_size", 1) or 1) BATCH_MAX_TOKENS = int(config.get("batch_max_tokens", 0) or 0) -# Inference CLI args (shared by every per-fold call in a batch). Two batch tweaks, -# both leaving any user-set value untouched: -# * --allow_resume: a crashed batch re-runs all its folds, so resume the ones whose -# outputs already exist. Only AlphaFold2 accepts it; AlphaFold3 rejects unknown -# flags (ValueError "not supported by backend 'alphafold3'"), so AF2 only. -# * --jax_compilation_cache_dir: the batch runs one run_structure_prediction.py per -# fold (the AF3 backend MERGES folds passed together into one complex, so they must -# be predicted separately), which would otherwise recompile per fold. A shared -# on-disk JAX cache lets later folds reuse earlier compilations (esp. AF3 buckets), -# recovering most of the per-job compile cost the single-process approach would have. -_structure_inference_args = dict(config.get("structure_inference_arguments", {})) -if BATCH_SIZE > 1: - if INFERENCE_BACKEND == "alphafold2": - _structure_inference_args.setdefault("--allow_resume", "true") - _structure_inference_args.setdefault( - "--jax_compilation_cache_dir", - join(config["output_directory"], ".jax_compilation_cache"), - ) +# Inference CLI args shared by every per-fold call in a batch. The batch-only flags +# (--allow_resume for AF2, --jax_compilation_cache_dir for AF3) are BACKEND-SPECIFIC and +# each backend hard-errors on flags it does not accept, so the gating lives in +# batch_inference_args (see its docstring / test_batch_inference_args.py). +_structure_inference_args = batch_inference_args( + config.get("structure_inference_arguments", {}), + backend=INFERENCE_BACKEND, + batch_size=BATCH_SIZE, + jax_cache_dir=join(config["output_directory"], ".jax_compilation_cache"), +) STRUCTURE_INFERENCE_CLI = " ".join( f"{k}={v}" for k, v in _structure_inference_args.items() ) diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 80bf4aa..d59d843 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -506,6 +506,38 @@ def prepare_container_binds( os.environ.setdefault(var, "1") +def batch_inference_args( + base_args: dict, + *, + backend: str, + batch_size: int, + jax_cache_dir: str, +) -> dict: + """Return inference CLI args with the batch-only flags added, each gated to the + backend that accepts them. + + ``run_structure_prediction.py`` validates its flags per backend and hard-errors on + any it does not recognise (``ValueError: not supported by backend ''``), so a + batch-only flag must ONLY be added for the backend(s) that accept it: + + * ``--allow_resume`` (AlphaFold2 only): a crashed batch re-runs all its folds, so + resume the ones already done. AlphaFold3 rejects it. + * ``--jax_compilation_cache_dir`` (AlphaFold3 only): lets the per-fold calls in a + batch share one on-disk JAX compile cache. This is a JAX/XLA flag; AlphaFold2 + rejects it (its inference is not JAX-compiled). + + With ``batch_size <= 1`` nothing is added (the unbatched pipeline is untouched). Any + value the user already set is preserved (``setdefault``). + """ + args = dict(base_args) + if batch_size > 1: + if backend == "alphafold2": + args.setdefault("--allow_resume", "true") + if backend == "alphafold3": + args.setdefault("--jax_compilation_cache_dir", jax_cache_dir) + return args + + def normalize_partitions(value: Any) -> str | None: """Normalise a ``slurm_partition`` config value to a comma-separated string. From 65bb43ac093665a13ef91c08c4f6925b340f3cb0 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 8 Jul 2026 16:37:25 +0200 Subject: [PATCH 7/9] Docs: warn that inference flags are backend-exclusive; complete the lists The per-backend flag lists were marked "common options", incomplete, and gave no warning that run_structure_prediction.py hard-errors on a flag the selected backend doesn't accept (ValueError "not supported by backend ''") -- the same class of issue that broke batched AF2 (--jax_compilation_cache_dir is AF3-only). Add an IMPORTANT callout stating flags are backend-exclusive, note that batching auto-adds the correct one per backend, point to the authoritative per-image list (run_structure_prediction.py --help), and fill in the missing AF2 relaxation flags and the shared --use_ap_style AF3 flag. Co-Authored-By: Claude Opus 4.8 --- README.md | 30 +++++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 9ef40f3..f2b1ffd 100644 --- a/README.md +++ b/README.md @@ -516,7 +516,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 ''` 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 run_structure_prediction.py --help +> ``` +> (`alphalink` accepts the AlphaFold2 flags plus `--crosslinks`.)
AlphaFold2 flags @@ -528,7 +547,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 @@ -555,7 +578,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 @@ -565,6 +588,7 @@ structure_inference_arguments: --num_recycles: 10 --save_embeddings: False --save_distogram: False + --use_ap_style: False # shared with AlphaFold2 ```
From 1826bf21109b4c6e5989c01c8c67e481885270f9 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 8 Jul 2026 16:48:21 +0200 Subject: [PATCH 8/9] Warn at parse time on inference flags the backend won't accept run_structure_prediction.py aborts the inference job on the first structure_inference_arguments flag outside its per-backend allow set (e.g. --allow_resume on AF3, --jax_compilation_cache_dir on AF2), deep inside a Slurm job. Add a parse-time guard: unknown_inference_flags() mirrors the container's _validate_flags_for_backend allow-sets and lists any user flag the selected --fold_backend does not accept; the Snakefile logs a clear warning on the head node in seconds. Kept a warning (not a hard error) since the container is the source of truth, so list drift only ever produces a spurious note. Uses the raw --fold_backend so alphalink's --crosslinks is not mis-flagged. Adds unit tests. Co-Authored-By: Claude Opus 4.8 --- test/test_unknown_inference_flags.py | 78 ++++++++++++++++++++++++++++ workflow/Snakefile | 20 +++++++ workflow/rules/common.smk | 54 +++++++++++++++++++ 3 files changed, 152 insertions(+) create mode 100644 test/test_unknown_inference_flags.py diff --git a/test/test_unknown_inference_flags.py b/test/test_unknown_inference_flags.py new file mode 100644 index 0000000..21c73a5 --- /dev/null +++ b/test/test_unknown_inference_flags.py @@ -0,0 +1,78 @@ +"""Unit tests for ``unknown_inference_flags`` (parse-time backend-flag guard). + +Warns early when ``structure_inference_arguments`` contains a flag the selected backend +does not accept, instead of failing deep in a Slurm job. Mirrors the container's +``_validate_flags_for_backend`` allow-sets. + +Run with ``python test/test_unknown_inference_flags.py`` or pytest. +""" + +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) +unknown_inference_flags = common.unknown_inference_flags + + +def test_af2_flags_all_accepted(): + args = {"--fold_backend": "alphafold2", "--num_cycle": 3, "--allow_resume": "true"} + assert unknown_inference_flags(args, "alphafold2") == [] + + +def test_af3_only_flag_flagged_on_af2(): + args = {"--fold_backend": "alphafold2", "--jax_compilation_cache_dir": "/c"} + assert unknown_inference_flags(args, "alphafold2") == ["jax_compilation_cache_dir"] + + +def test_af2_only_flag_flagged_on_af3(): + args = {"--fold_backend": "alphafold3", "--allow_resume": "true"} + assert unknown_inference_flags(args, "alphafold3") == ["allow_resume"] + + +def test_af3_flags_all_accepted(): + args = {"--fold_backend": "alphafold3", "--num_diffusion_samples": 5, + "--jax_compilation_cache_dir": "/c", "--use_ap_style": False} + assert unknown_inference_flags(args, "alphafold3") == [] + + +def test_alphalink_accepts_crosslinks_and_af2_flags(): + args = {"--fold_backend": "alphalink", "--crosslinks": "x.pkl", "--num_cycle": 3} + assert unknown_inference_flags(args, "alphalink") == [] + # but crosslinks is NOT an AlphaFold2 flag + assert unknown_inference_flags({"--crosslinks": "x.pkl"}, "alphafold2") == ["crosslinks"] + + +def test_leading_dashes_and_equals_are_stripped(): + # keys may arrive with or without dashes, or with an inline value + assert unknown_inference_flags({"allow_resume": "true"}, "alphafold3") == ["allow_resume"] + assert unknown_inference_flags({"--bogus=1": "x"}, "alphafold3") == ["bogus"] + + +def test_unknown_backend_stays_silent(): + # can't judge an unrecognised backend -> no false warnings + assert unknown_inference_flags({"--anything": 1}, "some-future-backend") == [] + + +def test_backend_name_case_insensitive(): + assert unknown_inference_flags({"--allow_resume": "t"}, "AlphaFold3") == ["allow_resume"] + + +def test_order_preserved_and_deduped(): + args = {"--allow_resume": "t", "--models_to_relax": "best", "--allow_resume2": "t"} + # only genuinely-unknown AF3 flags, in order + out = unknown_inference_flags(args, "alphafold3") + assert out == ["allow_resume", "models_to_relax", "allow_resume2"] + + +if __name__ == "__main__": + for name, fn in sorted(globals().items()): + if name.startswith("test_") and callable(fn): + fn() + print(f"ok {name}") + print("all unknown_inference_flags tests passed") diff --git a/workflow/Snakefile b/workflow/Snakefile index 43eb749..29e5a7a 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -282,6 +282,26 @@ _structure_inference_args = batch_inference_args( batch_size=BATCH_SIZE, jax_cache_dir=join(config["output_directory"], ".jax_compilation_cache"), ) +# Parse-time guard: run_structure_prediction.py aborts the inference job on the first +# flag its backend does not accept (e.g. --allow_resume on AF3, --jax_compilation_cache_dir +# on AF2). Warn here -- on the head node, in seconds -- instead of failing deep in a Slurm +# job. Uses the RAW --fold_backend (so alphalink's --crosslinks is not mis-flagged); a +# warning only, since the container is the source of truth (see unknown_inference_flags). +_RAW_FOLD_BACKEND = str( + config.get("structure_inference_arguments", {}).get("--fold_backend", DATA_PIPELINE) +).strip().lower() +_unknown_inference_flags = unknown_inference_flags( + _structure_inference_args, _RAW_FOLD_BACKEND +) +if _unknown_inference_flags: + logger.warning( + f"[inference-flags] structure_inference_arguments has flag(s) the " + f"'{_RAW_FOLD_BACKEND}' backend does not recognise: {_unknown_inference_flags}. " + f"run_structure_prediction.py will abort the job with a \"not supported by " + f"backend\" error -- remove them or change --fold_backend. (If your prediction " + f"container is newer than this workflow this may be a false alarm; confirm with " + f"`run_structure_prediction.py --help`.)" + ) STRUCTURE_INFERENCE_CLI = " ".join( f"{k}={v}" for k, v in _structure_inference_args.items() ) diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index d59d843..cdf1c82 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -538,6 +538,60 @@ def batch_inference_args( return args +# Inference flags each backend accepts, mirroring run_structure_prediction.py's +# ``_validate_flags_for_backend``. Names are WITHOUT the leading ``--``. This is only +# used for a parse-time WARNING: the container is the source of truth and hard-errors, +# so if this list drifts (a newer image adds a flag) the worst case is a spurious +# warning, never a blocked run. Keep in sync with AlphaPulldown when convenient. +_COMMON_INFERENCE_FLAGS = { + "input", "output_directory", "data_directory", "features_directory", + "protein_delimiter", "fold_backend", "random_seed", "storage_mode", +} +_AF2_LIKE_INFERENCE_FLAGS = { + "compress_result_pickles", "remove_result_pickles", "models_to_relax", + "relax_best_score_threshold", "remove_keys_from_pickles", "convert_to_modelcif", + "allow_resume", "num_cycle", "num_predictions_per_model", "pair_msa", + "save_features_for_multimeric_object", "skip_templates", "msa_depth_scan", + "multimeric_template", "model_names", "msa_depth", "description_file", + "path_to_mmt", "threshold_clashes", "hb_allowance", "plddt_threshold", + "desired_num_res", "desired_num_msa", "benchmark", "model_preset", + "use_ap_style", "use_gpu_relax", "dropout", +} +_AF3_INFERENCE_FLAGS = { + "jax_compilation_cache_dir", "buckets", "flash_attention_implementation", + "num_diffusion_samples", "num_seeds", "debug_templates", "debug_msas", + "num_recycles", "save_embeddings", "save_distogram", "use_ap_style", +} +_ALPHALINK_EXTRA_FLAGS = {"crosslinks"} + +ALLOWED_INFERENCE_FLAGS = { + "alphafold2": _COMMON_INFERENCE_FLAGS | _AF2_LIKE_INFERENCE_FLAGS, + "alphalink": _COMMON_INFERENCE_FLAGS | _AF2_LIKE_INFERENCE_FLAGS | _ALPHALINK_EXTRA_FLAGS, + "alphafold3": _COMMON_INFERENCE_FLAGS | _AF3_INFERENCE_FLAGS, +} + + +def unknown_inference_flags(args, backend: str) -> list: + """Return the ``structure_inference_arguments`` keys the given backend does not + accept (leading ``--`` and any ``=value`` ignored), preserving input order. + + ``run_structure_prediction.py`` aborts the inference job on the first flag outside + its per-backend allow set (``ValueError: not supported by backend ''``), deep + inside a Slurm job. Calling this at parse time lets the workflow warn on the head + node in seconds instead. Returns ``[]`` when the backend name is unrecognised (we + cannot judge, so stay silent) or every flag is accepted. + """ + allowed = ALLOWED_INFERENCE_FLAGS.get(str(backend).strip().lower()) + if allowed is None: + return [] + unknown: list = [] + for key in (args or {}): + name = str(key).lstrip("-").split("=", 1)[0].strip() + if name and name not in allowed and name not in unknown: + unknown.append(name) + return unknown + + def normalize_partitions(value: Any) -> str | None: """Normalise a ``slurm_partition`` config value to a comma-separated string. From 3f59ba9588dff604aed9080f11e2bd0f71e3fa0c Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Thu, 9 Jul 2026 11:06:49 +0200 Subject: [PATCH 9/9] Docs: note AF3 batching depends on container + shared filesystem Batched AlphaFold3 shares one --jax_compilation_cache_dir under output_directory; with recent tokamax-based AF3 images this can fail on some clusters (XLA autotune cache "Device or resource busy" on BeeGFS; H100 tokamax autotuning cache aborts at load). Advise batch_size: 1 for AF3 when jobs crash during compilation. AF2 batching is unaffected. Co-Authored-By: Claude Opus 4.8 --- README.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/README.md b/README.md index f2b1ffd..9c66b13 100644 --- a/README.md +++ b/README.md @@ -461,6 +461,15 @@ batch_max_tokens: 0 # optional cap on summed residues per batch (0 = no cap) 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