[codex] Update AlphaFold3 to Tokamax v3.0.2#623
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a076574c93
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| RUN wget -O /hmmer_build/jackhmmer_seq_limit.patch \ | ||
| https://raw.githubusercontent.com/google-deepmind/alphafold3/main/docker/jackhmmer_seq_limit.patch | ||
| COPY alphafold3/docker/jackhmmer_seq_limit.patch /hmmer_build/ |
There was a problem hiding this comment.
Check out submodules before copying AlphaFold3 files
In the build-alphafold3-container workflow I checked, the checkout step at .github/workflows/github_actions.yml:228 does not request submodules, while this Dockerfile now copies alphafold3/docker/jackhmmer_seq_limit.patch from the build context instead of downloading it/cloning with --recurse-submodules. On push/release builds that use that workflow, the alphafold3 gitlink is not populated, so BuildKit cannot find this source path and the container build stops before installing AF3; either the workflow needs submodules: recursive or the Dockerfile must not depend on checked-out submodule contents.
Useful? React with 👍 / 👎.
Drop the explicit chex>=0.1.91 dependency. chex>=0.1.91 requires Python >=3.11, which broke the smoke-tests (3.10) install step. chex is not imported anywhere in alphapulldown/, test/, or the bundled alphafold submodule and is not pulled into the base .[test] env, so removing the pin is safe and mirrors upstream alphafold3 (which does not pin chex either). Verified .[test] resolves under Python 3.10. Also add submodules: recursive to the build-alphafold3-container checkout. alphafold3.dockerfile copies alphafold3/docker/jackhmmer_ seq_limit.patch and the rest of the checked-out submodule from the build context, so without populating the gitlink the push/release container build fails before installing AF3. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Summary
alphafold3submodule from the old pre-Tokamax fork state toKosinskiLab/alphafold3@e3a1400ccc0bb2da2de19fd31cf50d088272175bonap-gapped-discontinuous-chains-v3.0.2.v3.0.2.Root Cause And Official History
The previous submodule SHA
6ad1a65predates the official AF3 compatibility work requested in google-deepmind/alphafold3#394. It was still on the old JAX/JAX-Triton/Triton stack (jax==0.4.34,jax-triton==0.2.0,triton==3.1.0) and could die immediately after Snakemake printedgpu_mem_mb=97887on RTX PRO 6000 / Blackwell cards.Official AF3 history checked for this update:
3890782: migrated AF3 totokamax.gated_linear_unitandtokamax.dot_product_attention.608edb6: updated to Tokamax0.0.11and JAX0.9.1.f6a5aec: officialv3.0.2, containing both commits.The new submodule SHA includes both official fixes:
git -C alphafold3 merge-base --is-ancestor 389078218c0604cf85caf4e2e9a830e456ff31d1 HEAD-> yes.git -C alphafold3 merge-base --is-ancestor 608edb684db9f6fd0e677fea01c4cefc60f8a8aa HEAD-> yes.Validation
pytest test/unit/test_alphafold3_backend_helpers.py -q->30 passed.pytest test/unit -q->319 passed, 2 skipped./scratch/dmolodenskiy/codex-af3-rtx/alphapulldown-af3-tokamax.sifafter local Docker/Podman/BuildKit paths were unavailable.alphafold3==3.0.2,jax==0.9.1,jaxlib==0.9.1,jax-cuda12-plugin==0.9.1, andtokamax==0.0.11.gpu50completed twice: jobs55520690and55520842, both onNVIDIA RTX PRO 6000 Blackwell Server Editionwithgpu_mem_mb=97887,jax_default_backend gpu, andrtx_smoke=ok.gpu50: job55521005,1 passed.gpu50: job55521729, produced the expectedP61626outputs includingcompleted_fold.txt.Coordinated Snakemake PR: KosinskiLab/AlphaPulldownSnakemake#47