diff --git a/.github/workflows/github_actions.yml b/.github/workflows/github_actions.yml index 8c4ece6a..8b83ca8c 100644 --- a/.github/workflows/github_actions.yml +++ b/.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: diff --git a/.gitmodules b/.gitmodules index 08c43b92..47906c66 100644 --- a/.gitmodules +++ b/.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 diff --git a/README.md b/README.md index 97fc6908..bc6ff8f0 100644 --- a/README.md +++ b/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=:`. 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 diff --git a/alphafold3 b/alphafold3 index 6ad1a659..e3a1400c 160000 --- a/alphafold3 +++ b/alphafold3 @@ -1 +1 @@ -Subproject commit 6ad1a65994c2111d291a386cdc048d8c9bfae4af +Subproject commit e3a1400ccc0bb2da2de19fd31cf50d088272175b diff --git a/alphapulldown/folding_backend/alphafold3_backend.py b/alphapulldown/folding_backend/alphafold3_backend.py index ca5d375a..0ac260aa 100644 --- a/alphapulldown/folding_backend/alphafold3_backend.py +++ b/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, diff --git a/docker/alphafold3.dockerfile b/docker/alphafold3.dockerfile index 1d7f6631..7901b99b 100644 --- a/docker/alphafold3.dockerfile +++ b/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,6 +28,7 @@ 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) @@ -32,11 +36,10 @@ 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,16 +61,17 @@ 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 \ @@ -76,39 +79,23 @@ ENV PIP_INDEX_URL=${PIP_INDEX_URL} \ 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 - diff --git a/pyproject.toml b/pyproject.toml index 0f00a59e..6c32dece 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/test/unit/test_alphafold3_backend_helpers.py b/test/unit/test_alphafold3_backend_helpers.py index 6b6d9df2..75aa2cc9 100644 --- a/test/unit/test_alphafold3_backend_helpers.py +++ b/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")