[codex] Update AlphaFold3 to Tokamax v3.0.2

.github/workflows/github_actions.yml

@@ -226,6 +226,8 @@ jobs:
         with:
           ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
       - uses: actions/checkout@v4
+        with:
+          submodules: recursive
       - run: rm -rf /opt/hostedtoolcache
       - uses: docker/login-action@v3
         with:

.gitmodules

@@ -17,4 +17,4 @@
 [submodule "alphafold3"]
 	path = alphafold3
 	url = https://github.com/KosinskiLab/alphafold3.git
-	branch = main
+	branch = ap-gapped-discontinuous-chains-v3.0.2

README.md

@@ -247,8 +247,11 @@ you hit these.
 - **Restrict to one model** with `structure_inference_gpu_model` (e.g. `"A100"`) → the plugin emits
   `--gpus=<model>:<count>`. Accepts a single model name; leave `""` for any.
 - **Exclude specific nodes** with `slurm_exclude_nodes` → passed verbatim to `sbatch --exclude`
-  (e.g. `"gpu50,gpu51"`). Use it for nodes whose GPU the container can't use — e.g. a CUDA compute
-  capability newer than the container's bundled `ptxas` (fails `ptxas too old` / `UNIMPLEMENTED`).
+  (e.g. `"gpu50,gpu51"`). Use it as a fallback for nodes whose GPU the container can't use — e.g.
+  a CUDA compute capability newer than the container's bundled `ptxas` (fails `ptxas too old` /
+  `UNIMPLEMENTED`). RTX PRO 6000 / Blackwell failures with old AlphaFold 3 images are fixed by
+  updating to AF3 v3.0.2/Tokamax; excluding those nodes is only an old-image workaround, not the
+  compatibility proof.
   `--exclude` is allowed in `slurm_extra` whereas `--constraint`/`--gres`/`--gpus` are not, so it is
   the supported way to drop a few nodes while keeping the rest of the partition.
 - **`structure_inference_max_runtime`** caps per-job wall time (minutes). Wall time scales as

alphapulldown/folding_backend/alphafold3_backend.py

@@ -26,12 +26,12 @@
 import haiku as hk
 import jax
 import numpy as np
+import tokamax
 from alphafold3.common import base_config
 from alphafold3.common import folding_input
 from alphafold3.constants import chemical_components
 from alphafold3.data import featurisation
 from alphafold3.data import parsers as af3_parsers
-from alphafold3.jax.attention import attention
 from alphafold3.model import features, params, post_processing
 from alphafold3.model import model
 from alphafold3.model.components import utils
@@ -145,8 +145,8 @@ def run_inference(
             self.device,
         )
         result = self._model(rng_key, featurised_example)
-        result = jax.tree_map(np.asarray, result)
-        result = jax.tree_map(
+        result = jax.tree_util.tree_map(np.asarray, result)
+        result = jax.tree_util.tree_map(
             lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x,
             result,
         )
@@ -742,7 +742,7 @@ def setup(
         def make_model_config(
             *,
             model_class: type[ModelT] = MyNewModel,
-            flash_attention_implementation: attention.Implementation,
+            flash_attention_implementation: tokamax.DotProductAttentionImplementation,
             num_diffusion_samples: int = 5,
             num_recycles: int = 10,
             return_embeddings: bool = False,
@@ -787,7 +787,8 @@ def make_model_config(
             model_class=MyNewModel,
             config=make_model_config(
                 flash_attention_implementation=typing.cast(
-                    attention.Implementation, flash_attention_implementation
+                    tokamax.DotProductAttentionImplementation,
+                    flash_attention_implementation,
                 ),
                 num_diffusion_samples=num_diffusion_samples,
                 num_recycles=num_recycles,

docker/alphafold3.dockerfile

@@ -2,6 +2,9 @@ FROM nvidia/cuda:12.6.3-base-ubuntu24.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV VIRTUAL_ENV=/opt/venv
+ENV UV_PROJECT_ENVIRONMENT=${VIRTUAL_ENV}
+ENV UV_COMPILE_BYTECODE=1
+ENV UV_LINK_MODE=copy
 ENV PATH="/hmmer/bin:${VIRTUAL_ENV}/bin:${PATH}"
 ENV PIP_DISABLE_PIP_VERSION_CHECK=1
 
@@ -10,9 +13,9 @@ ENV PIP_DISABLE_PIP_VERSION_CHECK=1
 # ---------------------------------------------------------------------
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
-        python3 \
-        python3-venv \
-        python3-dev \
+        python3.12 \
+        python3.12-venv \
+        python3.12-dev \
         gcc-12 g++-12 \
         build-essential \
         libc6-dev \
@@ -25,18 +28,18 @@ RUN apt-get update && \
         make \
         pkg-config \
         zlib1g-dev \
+        zstd \
     && rm -rf /var/lib/apt/lists/*
 
 # Force gcc-12 (avoid gcc-13 ICE on 24.04)
 ENV CC=gcc-12
 ENV CXX=g++-12
 
 # ---------------------------------------------------------------------
-# Python venv  (PIN pip so pip-tools works)
+# Python venv and package installer
 # ---------------------------------------------------------------------
-RUN python3 -m venv ${VIRTUAL_ENV} && \
-    pip install --no-cache-dir --upgrade "pip<25.3" setuptools wheel && \
-    pip install --no-cache-dir "pip-tools==7.5.2"
+COPY --from=ghcr.io/astral-sh/uv:0.9.24 /uv /uvx /bin/
+RUN uv venv ${VIRTUAL_ENV}
 
 # ---------------------------------------------------------------------
 # HMMER (with seq_limit patch)
@@ -46,8 +49,7 @@ RUN mkdir /hmmer_build /hmmer && \
     echo "ca70d94fd0cf271bd7063423aabb116d42de533117343a9b27a65c17ff06fbf3  /hmmer_build/hmmer-3.4.tar.gz" | sha256sum -c - && \
     tar -xzf /hmmer_build/hmmer-3.4.tar.gz -C /hmmer_build
 
-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/
 
 RUN cd /hmmer_build && \
     patch -p0 < jackhmmer_seq_limit.patch && \
@@ -59,56 +61,41 @@ RUN cd /hmmer_build && \
     rm -rf /hmmer_build
 
 # ---------------------------------------------------------------------
-# Clone AlphaPulldown with submodules
+# Copy AlphaPulldown with its checked-out submodules. Building from the local
+# checkout is important for PR/SIF validation: it exercises this branch's
+# alphafold3 submodule pointer instead of cloning the repository default branch.
 # ---------------------------------------------------------------------
-RUN git clone --recurse-submodules https://github.com/KosinskiLab/AlphaPulldown.git /app/AlphaPulldown
+COPY . /app/AlphaPulldown
 
 # ---------------------------------------------------------------------
-# Install AlphaFold3 (regenerate dev-requirements inside image)
+# Install AlphaFold3 from its locked v3.0.2/Tokamax environment
 # ---------------------------------------------------------------------
 WORKDIR /app/AlphaPulldown/alphafold3
 
-# force real PyPI, not a mirror
 ARG PIP_INDEX_URL=https://pypi.org/simple
 ENV PIP_INDEX_URL=${PIP_INDEX_URL} \
     PIP_DEFAULT_TIMEOUT=600 \
     PIP_RETRIES=10 \
     PIP_NO_CACHE_DIR=1 \
     PIP_DISABLE_PIP_VERSION_CHECK=1
 
-ENV PIP_CACHE_DIR=/tmp/pip-cache
 ENV CMAKE_BUILD_PARALLEL_LEVEL=1
 ENV CFLAGS="-O2 -pipe"
 ENV CXXFLAGS="-O2 -pipe"
-RUN apt-get update && apt-get install -y clang lld
-ENV CC=clang
-ENV CXX=clang++
-ENV CMAKE_BUILD_PARALLEL_LEVEL=2
-
-
-RUN rm -rf "$PIP_CACHE_DIR" && mkdir -p "$PIP_CACHE_DIR" && \
-    pip-compile \
-        --no-reuse-hashes \
-        --extra=dev \
-        --generate-hashes \
-        --resolver=backtracking \
-        --output-file=dev-requirements.txt \
-        pyproject.toml && \
-    pip install --require-hashes -r dev-requirements.txt && \
-    pip install git+https://github.com/openmm/pdbfixer.git && \
-    pip install --no-deps . && \
-    rm -rf "$PIP_CACHE_DIR"
+
+RUN uv sync --frozen --all-groups --no-editable && \
+    uv pip install --no-cache git+https://github.com/openmm/pdbfixer.git
 
 # ---------------------------------------------------------------------
 # Build CCD database
 # ---------------------------------------------------------------------
-RUN build_data
+RUN uv run --no-sync build_data
 
 # ---------------------------------------------------------------------
 # Install AlphaPulldown
 # ---------------------------------------------------------------------
 WORKDIR /app/AlphaPulldown
-RUN pip install --no-cache-dir .
+RUN uv pip install --no-cache .
 
 # ---------------------------------------------------------------------
 # Runtime env
@@ -128,4 +115,3 @@ from alphafold3.constants import chemical_components
 from alphapulldown.folding_backend.folding_backend import FoldingBackend
 print("AF3 + AlphaPulldown import OK, CCD present")
 EOF
-

pyproject.toml

@@ -38,9 +38,8 @@ dependencies = [
   "jupyterlab>=3.0",
   "ipywidgets",
   "ml-dtypes",
-  "chex>=0.1.86",
   "immutabledict>=2.0.0",
-  "typing-extensions==4.14.0",
+  "typing-extensions>=4.15.0",
 ]
 
 [project.optional-dependencies]
@@ -53,17 +52,15 @@ alphafold2 = [
   "pdbfixer>=1.10",
 ]
 alphafold3 = [
-  "jaxtyping==0.2.34",
-  "jax==0.5.3",
-  "jax[cuda12]==0.5.3; platform_system == 'Linux' and platform_machine == 'x86_64'",
-  "jax-triton==0.2.0; platform_system == 'Linux' and platform_machine == 'x86_64'",
+  "dm-haiku==0.0.16",
+  "jax==0.9.1",
+  "jax[cuda12]==0.9.1; platform_system == 'Linux' and platform_machine == 'x86_64'",
   "modelcif>=1.6",
-  "numpy<2",
+  "numpy",
   "openmm>=8.2",
   "pdbfixer>=1.10",
-  "rdkit==2024.3.5",
-  "triton==3.1.0; platform_system == 'Linux' and platform_machine == 'x86_64'",
-  "typeguard==2.13.3",
+  "rdkit==2025.9.4",
+  "tokamax==0.0.11",
   "zstandard",
 ]
 test = [

test/unit/test_alphafold3_backend_helpers.py

@@ -57,6 +57,11 @@ def _transform(fn):
         haiku.transform = _transform
         sys.modules["haiku"] = haiku
 
+    if "tokamax" not in sys.modules:
+        tokamax = types.ModuleType("tokamax")
+        tokamax.DotProductAttentionImplementation = str
+        sys.modules["tokamax"] = tokamax
+
     terms_dir = tmp_path / "af3cpp"
     terms_dir.mkdir(parents=True, exist_ok=True)
     (terms_dir / "OUTPUT_TERMS_OF_USE.md").write_text("stub terms", encoding="utf-8")