Skip to content

feat: schedule batched local MMseqs2-GPU features - #54

Draft
DimaMolod wants to merge 8 commits into
mainfrom
feat/mmseqs2-gpu
Draft

feat: schedule batched local MMseqs2-GPU features#54
DimaMolod wants to merge 8 commits into
mainfrom
feat/mmseqs2-gpu

Conversation

@DimaMolod

@DimaMolod DimaMolod commented Sep 3, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Summary

  • split opt-in AF3 feature generation into bounded create_batch_msas.py GPU shards and independent per-protein finalize_batch_features.py CPU jobs, so GPUs are released before native AF3 template processing
  • shard deterministically by query count and residues; one shard maps to one core MSA batch and loads each MMseqs database once per GPU allocation
  • size GPU host RAM from the padded database footprint plus one query chunk, and scale walltime by shard work
  • preserve successful per-protein MSA bundles across failed shard jobs; only the atomic completion manifest is a Snakemake output
  • validate manifest schema and every expected bundle using size/mtime as the fast path and SHA-256 after metadata changes; missing, corrupt, or stale bundles select a fresh repair target
  • keep final MSA semantic validation in the core finalizer, which removes only invalid bundles so the next DAG repairs them automatically
  • namespace MSA and final-feature caches by output-affecting MMseqs, database, and AF3 template provenance
  • use only the pinned MMseqs binary bundled in both prediction images, preserve existing container binds, and document padded-database creation, RAM/VRAM/storage, gpuserver guidance, and the continued need for native AF3 databases
  • run the complete workflow test suite in CI rather than an exact single test module

Native AF3 behavior is preserved: template search receives the merged UniRef90 + small-BFD + MGnify unpaired MSA.

Tests

  • full workflow suite with Snakemake 9: 81 passed
  • real DAG dry-runs cover bounded GPU shards, parallel CPU finalizers, resources, missing-bundle repair, and disabled-mode behavior
  • an executed failure/retry test verifies that Snakemake removes only the failed summary, keeps completed MSA bundles, and reuses them on retry
  • manifest tests cover non-object/corrupt manifests, stat fast path, digest fallback, same-size corruption, and repaired-cache stabilization

Dependency

Requires draft core PR KosinskiLab/AlphaPulldown#637, specifically its separate create_batch_msas.py / finalize_batch_features.py interfaces and schemaVersion: 2 artifact manifest (name, file, sizeBytes, mtimeNs, sha256).

Related resident inference workflow PR: #53.

This was referenced Sep 3, 2026
Draft
Draft
DimaMolod and others added 6 commits September 3, 2026 12:34
@DimaMolod
test: keep bind checks standalone
81a61e8
@DimaMolod
fix: harden mmseqs gpu workflow scheduling
7410200
@DimaMolod @claude
docs: point the MMseqs database example at the shared EMBL build
7083800 
The padded databases now exist at /g/alphafold/AlphaFold_DBs/mmseqs, so the block
can be uncommented as-is on the cluster instead of being retyped.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@DimaMolod @claude
feat: support CPU search and size both stages from measurement
0f6f2d8 
The search mode was fixed to GPU with no way to express it, and both stages borrowed
resource models that did not describe them.

- rename the config block to mmseqs2_features and add use_gpu, stated up front: it
  selects the search and decides whether the shard rule asks for a GPU at all, so a
  CPU run lands on slurm_partition_cpu with no GPU held
- replace gpu_database_ram_mb/gpu_chunk_ram_per_residue_mb with one search_ram_mb.
  Search memory measured FLAT in shard size - 1, 8, 32 and 128 queries against the
  full AF3 database set all peaked at 149 GB - so a per-residue term was describing
  something that does not vary. Default raised 64 GB -> 160 GB accordingly.
- runtime base raised 15 -> 90 min: 1 query took 65 min and 128 took 66, because the
  database scan dominates. CPU search gets cpu_runtime_multiplier (measured 2.4x).
- give finalization its own model instead of the jackhmmer one. It does not build
  MSAs; measured 0.24 GB / 13 s for one protein against the 24 GB and 24 h requested.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@DimaMolod @claude
fix: tell MMseqs2 the Slurm allocation, and document MSA recall
4412438 
- pass --mmseqs_split_memory_limit derived from the rule's own mem_mb (90% of it).
  MMseqs2 otherwise sizes its database splits from 90% of the physical node memory
  via sysconf(_SC_PHYS_PAGES), ignoring the cgroup: on a 755 GB node with --mem=64G
  it sizes for ~680 GB, declines to split, and is OOM-killed instead.
- record measured MSA recall against the native pipeline on the same four databases:
  90.2% unpaired and 98.0% paired overall, but 54-68% unpaired on the shallowest
  families, with identical template counts. Documented as something to spot-check,
  not as a verdict; matched inference and DockQ would be needed for that.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@DimaMolod @claude
fix: derive search memory from the largest configured database
dfe0bd8 
Measured against the real padded databases, one at a time:

  mgnify     170 GB on disk -> 128.7 GB peak RSS
  uniprot    117 GB         ->  81.3 GB
  uniref90    79 GB         ->  54.4 GB
  small_bfd   22 GB         ->  18.8 GB

All four chained in one stage peaked at 149.4 GB - 1.16x the largest database, 0.53x
the sum of the individual peaks. They are searched one after another, so a fixed
constant was both wrong in size and wrong in shape.

- size the estimate from the largest configured database (0.85 of its padded size,
  times the 1.16 chaining overhead), falling back to search_ram_mb when the paths
  cannot be read at parse time
- memory is flat in shard size (1/8/32/128 queries) and in query length (1144 vs 8268
  residues) in both CPU and GPU mode, matching MMseqs2's own estimator taking no
  query argument, so no per-query or per-residue term is added
- document db_load_mode: mgnify measured 128.7 -> 92.5 GB and 21% faster under mode 2,
  but on a warm page cache, so it is documented rather than defaulted

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@DimaMolod