feat: add resident inference batches - #636
Draft
DimaMolod wants to merge 7 commits into
Draft
Conversation
This was referenced Sep 3, 2026
AlphaFold2 inference is JAX-compiled, so a fresh process recompiles every model runner and a resident batch recompiles whenever the input shape changes. - accept --jax_compilation_cache_dir for the alphafold2 backend and apply it in setup(), so separate per-fold processes can share a persistent compile cache - allow --desired_num_res without --desired_num_msa, defaulting the MSA bound to the fold's own depth, so a batch can be padded to one shape from the residue count alone Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
pre_modelling_setup and AlphaPulldownPredictionAdapter._prepare_protein_object were near-identical retypes, and a `command -v` test in the workflow chose between them. They disagreed: the single-fold copy reassigned its metadata glob inside the feature-directory loop and tested it outside, so only the last feature directory could contribute, and it copied every match rather than the newest. The same fold produced different output depending on the container's age. - add alphapulldown/fold_preparation.py with prepare_fold(); both callers delegate - keep the batch behaviour: consider every feature directory, copy the newest match - name the warning after the description the search used, not the fragment name - retarget the entry-point tests at the shared module and drop its cached copy in the module-fabrication harness 426 passed; the one remaining failure predates this change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Which flags each backend accepts was stated once in run_structure_prediction.py and copied by hand into the workflow, where it drifted. The configuration dict handed to backend.setup was retyped verbatim in two places, and since every backend absorbs what it does not recognise in **kwargs, a mistyped key there was dropped in silence. - add alphapulldown/inference_flags.py: per-backend flag sets, one declarative key->attribute mapping for the model and postprocess configuration, and validate_model_configuration() - _validate_flags_for_backend, default_model_flags, default_postprocess_flags and the batch adapter's copies all delegate to it - both adapters validate the configuration before calling setup, so an unknown key is an error instead of a setting that never takes effect Co-Authored-By: Claude Opus 5 (1M context) <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.
Summary
Tests
Companion workflow PR: KosinskiLab/AlphaPulldownSnakemake#53.
Related MMseqs2-GPU core PR: #637.