Feature/batch small jobs#49
Open
DimaMolod wants to merge 3 commits into
Open
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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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.