diff --git a/README.md b/README.md index 5790979..a7c5e8e 100644 --- a/README.md +++ b/README.md @@ -412,6 +412,37 @@ 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, 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: 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..0955030 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -252,9 +252,117 @@ 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. 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 (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"), + ) +STRUCTURE_INFERENCE_CLI = " ".join( + f"{k}={v}" for k, v in _structure_inference_args.items() +) + + +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, + ) + + +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_grouping_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 +488,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,38 +570,45 @@ 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 + # 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] ), - protein_delimiter=protein_delimiter, - cli_parameters=" ".join( - [f"{k}={v}" for k, v in config["structure_inference_arguments"].items()] + 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", xla_mem_fraction=STRUCTURE_INFERENCE_XLA_MEM_FRACTION, resources: @@ -523,15 +631,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 in the loop, memory released between calls). 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 +642,11 @@ rule structure_inference: attempt=attempt, cap_mb=MAX_MEM_MB, ), + # Walltime scales with the number of folds in the batch (they run + # sequentially, one CLI call each), 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: @@ -569,21 +675,33 @@ 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}" """ 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()