feat: schedule local MMseqs2 features for AlphaFold 2 runs

[DimaMolod]

Local MMseqs2 features for AlphaFold 2 runs

mmseqs2_features refused every data pipeline except AlphaFold 3. AlphaPulldown (KosinskiLab/AlphaPulldown#642) now finalizes the same MSA bundles into AlphaFold 2 pickles, so this PR schedules the stage for --data_pipeline: alphafold2 too. It uses the same bounded GPU/CPU search shards and separate CPU finalization jobs, and writes .pkl output through the existing backend-aware FEATURE_OUTPUT_PATTERN. The finalize rule gets its backend from --data_pipeline rather than a hard-coded alphafold3. For AlphaFold 2 it forwards the template-stack arguments from create_feature_arguments (--use_hhsearch, the PDB seqres/PDB70/mmCIF paths, the search binaries), but never the native MSA arguments this stage replaces.

Running it needs a prediction image built from an AlphaPulldown that includes KosinskiLab/AlphaPulldown#642.

Fixes that affect AlphaFold 3 runs too

• Every local MMseqs2 shard failed with a KeyError on a run that starts from UniProt accessions, on every retry and before any search ran. The first real AlphaFold 2 end-to-end run found it. Snakemake re-parses the workflow in every SLURM job, and the shard plan depended on things that change between parses: lengths learned from FASTAs the run itself downloads, a shared feature store other runs write to, and (for repairs) the MSA cache directory's mtime, which every running shard changes. Planning is now one module, schedule_msa_shards, over an append-only shard registry beside the MSA cache. A protein once planned keeps its shard, only uncovered proteins are planned and appended, and every registered job id keeps resolving. A grown sample sheet re-plans nothing. A real-Snakemake regression test reproduces the original failure.
• Switching use_gpu silently reused the previous alignments. The workflow's MSA cache key decides whether the core stage runs at all, and it did not include search mode. It does now. num_iterations was unreachable and is now configurable and keyed. db_load_mode was documented but read by nothing; it is now passed through, and deliberately left out of the key because it is memory-only.
• AlphaFold 3 template overrides were dropped. AlphaFold 3's template search honours --pdb_seqres_database_path, --template_mmcif_dir and the hmmsearch/hmmbuild binaries, exactly as the native run does. The rule forwarded none of them, so an override silently searched the defaults under --data_dir. A run that overrides nothing keeps its feature cache namespace; a run that does override gets features built from the templates it asked for.

Template cache identity

For AF2 with --use_hhsearch: true, mmseqs2_features.template_database_ids.pdb70 is required and forwarded as --template_pdb70_database_id. Other searches require pdb_seqres; all require mmcif. The command and feature key share one database-selection helper. Update the selected immutable ID whenever its database is rebuilt, even at the same path: only finalized features are invalidated, and the MSA bundles stay reusable. Unused database identities do not move the cache. This needs the companion core update in KosinskiLab/AlphaPulldown#642.

Tests

137 tests pass at 96e92bc (2026-09-14, SLURM 61831784). Unit coverage reaches all eight new executable functions in the workflow module. PDB70 regressions verify command forwarding, missing-ID rejection, feature invalidation, and MSA reuse. They include the real-Snakemake regression test for the shard KeyError, an AlphaFold 2 dry-run DAG test, and schedule tests for inputs arriving between parses, growth, a protein leaving mid-run, repairs surviving other shards publishing, successive repairs, and an unusable registry.

End to end on the cluster: the historical 2026-09-11 run used workflow dd5b1e7 and core 71eef1c4, completing all 35 steps. It started from UniProt accessions and local FASTAs, put five proteins through one search shard (54 min on a MIG slice, first attempt), five finalizations (29–51 s each), nine AF2-multimer folds and their analysis. The folds cover a heterodimer, a homodimer, a chopped chain and a monomer, plus three edge cases on the real databases, alone and paired: thioredoxin (a deep MSA), a random 90-residue orphan with no homologs or templates, and a 26-residue peptide. A second invocation found nothing to do. This run used hmmsearch and predates the final resource-default commit and current PDB70 identity fix; the current unit and core integration/real-binary gates validate those revisions. The earlier run with the old shard planning is what surfaced the KeyError fixed here.

Benchmark

Numbers are in KosinskiLab/AlphaPulldown#642. For AlphaFold 2, top-ranked DockQ over 12 heterodimers averaged 0.56 with local MMseqs2 features against 0.59 with native reduced_dbs ones, and 9 of 12 interfaces were acceptable either way.

Also from the measurements

• AlphaFold 2 finalization has its own resource defaults. It had borrowed AlphaFold 3's (2 GB + 1 MB per residue, 30 min, ×1.1 per retry). AlphaFold 2's template featurization peaked at 18.8 GB and 87 min over 56 chains (median ~1 GB and 2 min), and it follows which structures the templates come from, not length. The defaults are now per backend: AlphaFold 2 gets a flat 16 GB base and 60 min, while AlphaFold 3 is unchanged and explicit settings still win.
• Cold GPU searches on network storage are documented. The search reads ~375 GB of padded databases in full, so from NFS a cold first attempt is I/O-bound: ~150 MB/s with the GPU idle, about an hour per shard, against minutes from page cache. The README says what to do when first attempts time out.

Follow-up

• Local MMseqs2: a database rebuilt in place under the same identifier is never re-searched #56: a database rebuilt in place under the same identifier is never re-searched. Keying it at parse time would break the every-parse-agrees invariant for databases the submit host cannot see, so it needs a design decision.

🤖 Generated with Claude Code