Feature/batch small jobs#49
Merged
Merged
Conversation
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 <[email protected]>
- 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 <[email protected]>
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 <[email protected]>
`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 <[email protected]>
Allow multiple SLURM partitions for structure inference
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 <[email protected]>
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 <[email protected]>
…ists 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 '<name>'") -- 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 <[email protected]>
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 <[email protected]>
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 <[email protected]>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.