From 5acaeeacb23efb09a089dfa80da28d873581509a Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 07:11:44 +0200 Subject: [PATCH 01/14] ci: run the AlphaFold fork's fused-kernel hook tests pytest.ini skips the alphafold/ submodule, so the fork's alphafold/model/fused_kernels_test.py (plain-JAX stand-ins for colabfold-kernels, CPU) never ran in CI. A step in smoke-tests now names it. It runs on Python 3.11 only: the fork's base_config needs typing.dataclass_transform, which AlphaPulldown backfills only at runtime on 3.10. A workflow contract test keeps the step in place. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/github_actions.yml | 8 ++++++++ test/unit/test_mmseqs2_gpu_dockerfiles.py | 19 +++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/.github/workflows/github_actions.yml b/.github/workflows/github_actions.yml index 1bc49788..ce5d01e5 100644 --- a/.github/workflows/github_actions.yml +++ b/.github/workflows/github_actions.yml @@ -52,6 +52,14 @@ jobs: #pytest -s test/test_pdb_analyser.py #pytest -s test/test_get_good_inter_pae.py + - name: AlphaFold fork fused-kernel hooks (fake ops, CPU) + # pytest.ini does not collect the alphafold/ submodule, so the file is named + # here. The fork's base_config needs typing.dataclass_transform (Python 3.11+, + # as in the AlphaFold 2 image); AlphaPulldown backfills it only at runtime. + if: matrix.python-version != '3.10' + run: | + pytest -q alphafold/alphafold/model/fused_kernels_test.py + coverage: runs-on: ubuntu-latest defaults: diff --git a/test/unit/test_mmseqs2_gpu_dockerfiles.py b/test/unit/test_mmseqs2_gpu_dockerfiles.py index a3564b8e..8d9e9414 100644 --- a/test/unit/test_mmseqs2_gpu_dockerfiles.py +++ b/test/unit/test_mmseqs2_gpu_dockerfiles.py @@ -101,3 +101,22 @@ def test_alphafold3_pr_build_runs_compiled_feature_batch_contracts(): assert "MMSEQS_INTEGRATION_BINARY=/opt/mmseqs/bin/mmseqs" in dockerfile assert "test/integration/test_mmseqs2_command_contract.py" in dockerfile assert 'addopts="-ra --strict-markers"' in dockerfile + + +def test_smoke_tests_run_the_alphafold_fork_fused_kernel_hooks(): + # pytest.ini never collects the alphafold/ submodule; CI must name the file. + workflow = yaml.load(WORKFLOW.read_text(encoding="utf-8"), Loader=yaml.BaseLoader) + steps = workflow["jobs"]["smoke-tests"]["steps"] + hook_tests = [ + step + for step in steps + if "alphafold/alphafold/model/fused_kernels_test.py" in step.get("run", "") + ] + + assert len(hook_tests) == 1 + assert ( + REPOSITORY / "alphafold" / "alphafold" / "model" / "fused_kernels_test.py" + ).is_file() + # The fork needs Python 3.11's typing; the matrix's other entries must run it. + assert hook_tests[0]["if"] == "matrix.python-version != '3.10'" + assert "3.11" in workflow["jobs"]["smoke-tests"]["strategy"]["matrix"]["python-version"] From 8127a26d084e7f00171ac0fe54093a5012fa5655 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 07:16:10 +0200 Subject: [PATCH 02/14] test: one AF3 compile per process, with the real tokamax AlphaPulldown calls tokamax's private get_autotuning_cache_overlay_state before the first trace, because tokamax otherwise creates a JAX user context during it and the next prediction recompiles the model (1afe779e). The unit test with a stub only checks that the hook is called by name. These run the real tokamax and JAX on CPU (AF3 image; skipped without tokamax): the hook exists and is what a kernel's autotuning-cache lookup goes through; a jitted function doing that lookup and calling a tokamax op compiles twice in three calls without the hook (the regression) and once after alphafold3_backend._initialise_tokamax_trace_context. Each case runs in a fresh process. Passes on jax 0.9.1/tokamax 0.0.11 and jax 0.10.2/tokamax 0.0.12. Co-Authored-By: Claude Opus 5.5 --- test/unit/test_af3_tokamax_trace_context.py | 96 +++++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 test/unit/test_af3_tokamax_trace_context.py diff --git a/test/unit/test_af3_tokamax_trace_context.py b/test/unit/test_af3_tokamax_trace_context.py new file mode 100644 index 00000000..08c3142d --- /dev/null +++ b/test/unit/test_af3_tokamax_trace_context.py @@ -0,0 +1,96 @@ +"""AF3 must compile its model once per process, not again on the second prediction. + +tokamax creates a JAX user context the first time an op looks up its autotuning cache, +which is while the model is first traced. That changes JAX's trace context, part of +every jit cache key, so the next call re-traced and recompiled the whole model +(~50 s per prediction). AlphaPulldown calls tokamax's private hook before the first +trace (``alphafold3_backend._initialise_tokamax_trace_context``). + +These tests use the real tokamax and JAX, on CPU, so they run in the AF3 image and +skip elsewhere. On CPU tokamax's XLA implementations need no config and never look up +the cache, so the jitted function does the lookup a GPU kernel does while it is traced. +Each case runs in a fresh process: the context is created once per process (thread). +""" + +from __future__ import annotations + +import importlib +import inspect +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("tokamax") + +REPO_ROOT = Path(__file__).resolve().parents[2] + +_PROBE = r""" +import json, sys +import jax +import jax.numpy as jnp +import tokamax +from tokamax._src.ops.gated_linear_unit import base + +if sys.argv[1] == "alphapulldown": + from alphapulldown.folding_backend.alphafold3_backend import ( + _initialise_tokamax_trace_context, + ) + _initialise_tokamax_trace_context() + +op = base.GatedLinearUnit() + +@jax.jit +def forward(x, weights): + # What a tokamax GPU kernel does while it is traced: resolve its config. + op.bind(x, weights, activation=jax.nn.swish).cached_autotuning_data + return tokamax.gated_linear_unit( + x=x, weights=weights, activation=jax.nn.swish, implementation="xla" + ) + +args = (jnp.ones((8, 16)), jnp.ones((16, 2, 32))) +for _ in range(3): + jax.block_until_ready(forward(*args)) +print("PROBE", json.dumps({"compiled": forward._cache_size()})) +""" + + +def _compiled_versions(variant: str) -> int: + env = os.environ.copy() + env["JAX_PLATFORMS"] = "cpu" + env["PYTHONPATH"] = os.pathsep.join( + [str(REPO_ROOT), *filter(None, [env.get("PYTHONPATH")])] + ) + result = subprocess.run( + [sys.executable, "-c", _PROBE, variant], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + timeout=300, + ) + assert result.returncode == 0, result.stdout + result.stderr + line = next(l for l in result.stdout.splitlines() if l.startswith("PROBE ")) + return json.loads(line[len("PROBE "):])["compiled"] + + +def test_tokamax_still_has_the_hook_alphapulldown_calls(): + op = importlib.import_module("tokamax._src.ops.op") + + assert callable(getattr(op, "get_autotuning_cache_overlay_state", None)) + # The lookup a kernel makes while traced goes through the same hook. + lookup = inspect.getsource(op.BoundArguments.cached_autotuning_data.fget) + assert "get_autotuning_cache_overlay_state" in lookup + + +def test_a_jitted_tokamax_op_recompiles_on_its_second_call_without_the_hook(): + # The regression itself: if this stops recompiling, tokamax no longer creates its + # context lazily and the AlphaPulldown hook call may be unnecessary. + assert _compiled_versions("stock") == 2 + + +def test_alphapulldown_hook_keeps_a_jitted_tokamax_op_to_one_compile(): + assert _compiled_versions("alphapulldown") == 1 From f3584e4a55531ba40f1b0c9ec6a0a34421ce041f Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 07:16:10 +0200 Subject: [PATCH 03/14] ci: image builds check the fused kernels they ship The AF3 image's af3-compatibility-tests stage now also runs test_af3_fused_triangles.py, test_alphafold3_backend_helpers.py, the new tokamax one-compile test, and the fork's fused_triangle_test.py (Pallas interpret mode, as a script so it imports the installed AF3 and its CCD data rather than src/). All pass on CPU in the 2.9.0 image with the fork's v3.0.4 model/jax trees, on jax 0.10.2 and 0.9.1. The AF2 image has no test stage; after its last dependency step it now imports colabfold-kernels through alphafold.model.fused_kernels.fused_ops() (no GPU needed), so a package/fork/jax mismatch fails the build rather than a --fast_kernels prediction. Contract tests pin both. Co-Authored-By: Claude Opus 5.5 --- docker/alphafold2.dockerfile | 10 ++++++ docker/alphafold3.dockerfile | 9 ++++- test/unit/test_mmseqs2_gpu_dockerfiles.py | 42 +++++++++++++++++++++++ 3 files changed, 60 insertions(+), 1 deletion(-) diff --git a/docker/alphafold2.dockerfile b/docker/alphafold2.dockerfile index 92953994..0c8542b8 100644 --- a/docker/alphafold2.dockerfile +++ b/docker/alphafold2.dockerfile @@ -150,6 +150,16 @@ RUN pip3 install --upgrade pip --no-cache-dir \ # numpy<2 as the LAST dependency step. jax 0.5.3 runs fine with numpy 1.26.x. RUN pip install --no-cache-dir "numpy<2" +# --fast_kernels reaches colabfold-kernels through the fork's hooks, which add the +# Pallas names jax 0.5.3 lacks. Check that this package, fork and jax import together +# (no GPU needed), so an incompatible release fails the build, not a prediction. +RUN python -c "from alphafold.model import fused_kernels; \ +ops = fused_kernels.fused_ops(); \ +assert all(callable(getattr(ops, name)) for name in ('attention', 'layer_norm', 'gated_dual_proj')); \ +from alphapulldown.prediction import fast_kernels; \ +assert fast_kernels._package_version(), 'colabfold-kernels version unknown'; \ +print('colabfold-kernels', fast_kernels._package_version(), 'imports through alphafold.model.fused_kernels')" + # AlphaFold's template code formats a hit's `sum_probs` with %.2f in the error # path of `_process_single_hit`, but sum_probs is legitimately None for some hits # - the same function's *warning* path already prints it with %s, and diff --git a/docker/alphafold3.dockerfile b/docker/alphafold3.dockerfile index c3604967..253ec4c1 100644 --- a/docker/alphafold3.dockerfile +++ b/docker/alphafold3.dockerfile @@ -143,6 +143,9 @@ EOF # --------------------------------------------------------------------- FROM base AS af3-compatibility-tests +# The fork's fused_triangle_test.py checks the triangle kernels and their dispatch in +# Pallas interpret mode, on CPU. It runs as a script so that it imports the installed +# AF3 (with its built CCD data) rather than the submodule's src/ tree. WORKDIR /app/AlphaPulldown RUN uv pip install --no-cache "modelcif>=1.6" && \ python -c "from alphafold3.common import folding_input; import alphafold3.structure.mmcif; import ihm; import modelcif.reader" && \ @@ -151,7 +154,11 @@ RUN uv pip install --no-cache "modelcif>=1.6" && \ test/unit/test_create_batch_features.py \ test/unit/test_feature_metadata.py::test_embedded_json_is_accepted_and_round_tripped_by_vanilla_alphafold3 \ test/unit/test_af3_modelcif.py::test_augment_real_af3_modelcif_preserves_comments_and_is_modelcif_readable \ - test/unit/test_convert_to_modelcif_helpers.py && \ + test/unit/test_convert_to_modelcif_helpers.py \ + test/unit/test_af3_fused_triangles.py \ + test/unit/test_alphafold3_backend_helpers.py \ + test/unit/test_af3_tokamax_trace_context.py && \ + JAX_PLATFORMS=cpu python alphafold3/src/alphafold3/model/network/fused_triangle_test.py && \ MMSEQS_INTEGRATION_BINARY=/opt/mmseqs/bin/mmseqs \ python -m pytest -q -o addopts="-ra --strict-markers" \ test/integration/test_mmseqs2_command_contract.py && \ diff --git a/test/unit/test_mmseqs2_gpu_dockerfiles.py b/test/unit/test_mmseqs2_gpu_dockerfiles.py index 8d9e9414..043efd20 100644 --- a/test/unit/test_mmseqs2_gpu_dockerfiles.py +++ b/test/unit/test_mmseqs2_gpu_dockerfiles.py @@ -103,6 +103,34 @@ def test_alphafold3_pr_build_runs_compiled_feature_batch_contracts(): assert 'addopts="-ra --strict-markers"' in dockerfile +def test_alphafold3_image_tests_the_fused_triangles_and_one_compile_per_process(): + dockerfile = (REPOSITORY / "docker" / "alphafold3.dockerfile").read_text( + encoding="utf-8" + ) + stage = dockerfile.split("FROM base AS af3-compatibility-tests", 1)[1].split( + "FROM base AS runtime", 1 + )[0] + + for test_path in ( + "test/unit/test_af3_fused_triangles.py", + "test/unit/test_alphafold3_backend_helpers.py", + "test/unit/test_af3_tokamax_trace_context.py", + "alphafold3/src/alphafold3/model/network/fused_triangle_test.py", + ): + assert test_path in stage, test_path + assert (REPOSITORY / test_path).is_file(), test_path + # As a script, so it imports the installed AF3 and its built CCD data. + assert ( + "JAX_PLATFORMS=cpu python alphafold3/src/alphafold3/model/network/" + "fused_triangle_test.py" + ) in stage + # The runtime image cannot be built without the test stage passing. + assert "FROM base AS runtime" in dockerfile + assert "COPY --from=af3-compatibility-tests" in dockerfile.split( + "FROM base AS runtime", 1 + )[1] + + def test_smoke_tests_run_the_alphafold_fork_fused_kernel_hooks(): # pytest.ini never collects the alphafold/ submodule; CI must name the file. workflow = yaml.load(WORKFLOW.read_text(encoding="utf-8"), Loader=yaml.BaseLoader) @@ -120,3 +148,17 @@ def test_smoke_tests_run_the_alphafold_fork_fused_kernel_hooks(): # The fork needs Python 3.11's typing; the matrix's other entries must run it. assert hook_tests[0]["if"] == "matrix.python-version != '3.10'" assert "3.11" in workflow["jobs"]["smoke-tests"]["strategy"]["matrix"]["python-version"] + + +def test_alphafold2_image_checks_the_fused_kernels_import_with_its_final_jax(): + dockerfile = (REPOSITORY / "docker" / "alphafold2.dockerfile").read_text( + encoding="utf-8" + ) + kernels = dockerfile.index('"colabfold-kernels>=0.4.0,<0.5"') + jax = dockerfile.index('"jax[cuda12]"==0.5.3') + numpy = dockerfile.index('RUN pip install --no-cache-dir "numpy<2"') + check = dockerfile.index("fused_kernels.fused_ops()") + + # After the last dependency change, so it checks what the image ships. + assert kernels < jax < numpy < check + assert "pip install" not in dockerfile[check:] From 112a86ff0d8557e764576b367dcededc4edf616a Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 07:17:45 +0200 Subject: [PATCH 04/14] test: padded AF2 prediction reports scores over the real residues, end to end 4519a56a's tests call recalculate_confidence directly. This one runs --desired_num_res through predict_individual_job and postprocess: the real pad_input_features, two fake multimer RunModels that see and score the padded length (pLDDT, PAE with aligned_confidence_probs, pTM and ipTM including the padding), and AlphaFold's real confidence module for the JSON. It checks that the models ran padded while the object kept its unpadded features, and that ranking_debug.json (order, iptm+ptm, iptm), pae_*.json, confidence_*.json, the PAE plots and the rewritten result pickles cover only the 7 real residues and match the unpadded scores. With trimming disabled, the two models rank in the opposite order and it fails. Co-Authored-By: Claude Opus 5.5 --- test/unit/test_alphafold2_backend_helpers.py | 177 ++++++++++++++++++- 1 file changed, 173 insertions(+), 4 deletions(-) diff --git a/test/unit/test_alphafold2_backend_helpers.py b/test/unit/test_alphafold2_backend_helpers.py index 17888ec8..e35ae8c7 100644 --- a/test/unit/test_alphafold2_backend_helpers.py +++ b/test/unit/test_alphafold2_backend_helpers.py @@ -9,6 +9,9 @@ import numpy as np import pytest +# The real padding, imported before af2_backend_module replaces its module with a stub. +from alphapulldown.utils.modelling_setup import pad_input_features as real_pad_input_features + MODULE_PATH = ( Path(__file__).resolve().parents[2] @@ -917,12 +920,17 @@ def _padded_run_model_result(confidence_module, n_real, n_pad, seed=0): return result, truth, asym_id -def test_padded_scores_are_recomputed_over_the_real_residues(af2_backend_module, monkeypatch): - # AlphaFold's own confidence module (numpy and scipy only), not the test stub. +def _real_af2_confidence(): + """AlphaFold's own confidence module (numpy and scipy only), not the test stub.""" path = Path(__file__).resolve().parents[2] / "alphafold/alphafold/common/confidence.py" spec = importlib.util.spec_from_file_location("real_af2_confidence", path) - real_confidence = importlib.util.module_from_spec(spec) - spec.loader.exec_module(real_confidence) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_padded_scores_are_recomputed_over_the_real_residues(af2_backend_module, monkeypatch): + real_confidence = _real_af2_confidence() monkeypatch.setattr(af2_backend_module, "confidence", real_confidence) result, truth, asym_id = _padded_run_model_result(real_confidence, n_real=7, n_pad=5) @@ -969,6 +977,167 @@ def predicted_tm_score(logits, breaks, asym_id=None, interface=False): assert "iptm" not in monomer +class _PaddedRunModel: + """AlphaFold-Multimer's RunModel as far as scoring goes: outputs at the padded length. + + It keeps fixed PAE-head logits for a padded length and returns what RunModel + returns for them (pLDDT, the PAE array with its bin probabilities, and pTM/ipTM + over every residue it was given, padding included). + """ + + multimer_mode = True + + def __init__(self, confidence_module, n_padded, seed): + rng = np.random.default_rng(seed) + self.confidence = confidence_module + self.logits = rng.normal(size=(n_padded, n_padded, 64)) * 3 + self.breaks = np.linspace(0.0, 31.0, 63) + self.plddt = rng.uniform(30, 95, size=n_padded) + self.seen_num_res = [] + + def process_features(self, feature_dict, random_seed): + return dict(feature_dict) # multimer models take the features as they are + + def predict(self, processed_feature_dict, random_seed): + n = processed_feature_dict["aatype"].shape[0] + self.seen_num_res.append(n) + logits = self.logits[:n, :n] + result = self.confidence.compute_predicted_aligned_error(logits=logits, breaks=self.breaks) + result["plddt"] = self.plddt[:n] + result["ptm"] = self.confidence.predicted_tm_score(logits=logits, breaks=self.breaks) + result["iptm"] = self.confidence.predicted_tm_score( + logits=logits, breaks=self.breaks, + asym_id=processed_feature_dict["asym_id"], interface=True, + ) + result["ranking_confidence"] = 0.8 * result["iptm"] + 0.2 * result["ptm"] + return result + + def real_scores(self, asym_id): + n = len(asym_id) + logits = self.logits[:n, :n] + ptm = self.confidence.predicted_tm_score(logits=logits, breaks=self.breaks) + iptm = self.confidence.predicted_tm_score( + logits=logits, breaks=self.breaks, asym_id=asym_id, interface=True + ) + pae = self.confidence.compute_predicted_aligned_error(logits=logits, breaks=self.breaks) + return { + "ptm": ptm, + "iptm": iptm, + "ranking_confidence": 0.8 * iptm + 0.2 * ptm, + "plddt": self.plddt[:n], + "predicted_aligned_error": pae["predicted_aligned_error"], + } + + +def _dimer_features(chain_lengths, num_msa): + """Unpadded AlphaFold-Multimer features of a heterodimer, as pad_input_features sees them.""" + num_res = sum(chain_lengths) + asym_id = np.concatenate( + [np.full(length, chain, dtype=np.int32) for chain, length in enumerate(chain_lengths, 1)] + ) + return { + "aatype": np.arange(num_res, dtype=np.int32) % 20, + "residue_index": np.concatenate([np.arange(length) for length in chain_lengths]).astype(np.int32), + "asym_id": asym_id, + "entity_id": asym_id.copy(), + "sym_id": np.ones(num_res, dtype=np.int32), + "seq_mask": np.ones(num_res, dtype=np.float32), + "msa": np.zeros((num_msa, num_res), dtype=np.int32), + "deletion_matrix": np.zeros((num_msa, num_res), dtype=np.float32), + "msa_mask": np.ones((num_msa, num_res), dtype=np.float32), + "template_aatype": np.zeros((1, num_res), dtype=np.int32), + "template_all_atom_positions": np.ones((1, num_res, 37, 3), dtype=np.float32), + "template_all_atom_mask": np.ones((1, num_res, 37), dtype=np.float32), + "assembly_num_chains": np.array(len(chain_lengths)), + "num_templates": np.array(1), + "seq_length": np.array(num_res), + "num_alignments": np.array(num_msa), + } + + +def test_padded_prediction_reports_scores_and_pae_over_the_real_residues( + af2_backend_module, monkeypatch, tmp_path +): + # --desired_num_res end to end: the real padding, a model that sees and scores the + # padded length, AlphaFold's real confidence code, then the files postprocess writes. + confidence = _real_af2_confidence() + monkeypatch.setattr(af2_backend_module, "confidence", confidence) + monkeypatch.setattr(af2_backend_module, "pad_input_features", real_pad_input_features) + plotted = [] + monkeypatch.setattr( + af2_backend_module, + "plot_pae_from_matrix", + lambda **kwargs: plotted.append(np.shape(kwargs["pae_matrix"])), + ) + monkeypatch.setattr(af2_backend_module, "post_prediction_process", lambda *a, **k: None) + monkeypatch.setattr(af2_backend_module, "_resolve_gpu_relax", lambda use_gpu: False) + + chain_lengths, n_real, n_padded = (3, 4), 7, 12 + features = _dimer_features(chain_lengths, num_msa=5) + dimer = af2_backend_module.MultimericObject( + description="dimer", + input_seqs=["MKV", "LLAG"], + feature_dict=features, + multimeric_mode=False, + ) + runners = { + "model_1_multimer_v3_pred_0": _PaddedRunModel(confidence, n_padded, seed=1), + "model_2_multimer_v3_pred_0": _PaddedRunModel(confidence, n_padded, seed=2), + } + + results = af2_backend_module.AlphaFold2Backend.predict_individual_job( + model_runners=runners, + multimeric_object=dimer, + allow_resume=False, + skip_templates=False, + output_dir=tmp_path, + random_seed=0, + desired_num_res=n_padded, + ) + # The models ran padded; the object keeps its unpadded features (and chain IDs). + assert [runner.seen_num_res for runner in runners.values()] == [[n_padded], [n_padded]] + assert dimer.feature_dict["asym_id"].shape == (n_real,) + truth = {name: runner.real_scores(features["asym_id"]) for name, runner in runners.items()} + for name in runners: + assert results[name]["iptm"] != pytest.approx(truth[name]["iptm"], abs=1e-3) + + af2_backend_module.AlphaFold2Backend.postprocess( + prediction_results=results, + multimeric_object=dimer, + output_dir=str(tmp_path), + features_directory=str(tmp_path), + models_to_relax=af2_backend_module.ModelsToRelax.NONE, + convert_to_modelcif=False, + ) + + ranking = json.loads((tmp_path / "ranking_debug.json").read_text()) + assert ranking["order"] == sorted( + truth, key=lambda name: truth[name]["ranking_confidence"], reverse=True + ) + assert plotted == [(n_real, n_real)] * 2 + for name, expected in truth.items(): + assert ranking["iptm+ptm"][name] == pytest.approx(expected["ranking_confidence"], abs=1e-6) + assert ranking["iptm"][name] == pytest.approx(expected["iptm"], abs=1e-6) + + (pae_record,) = json.loads((tmp_path / f"pae_{name}.json").read_text()) + pae = np.asarray(pae_record["predicted_aligned_error"]) + assert pae.shape == (n_real, n_real) + np.testing.assert_allclose(pae, expected["predicted_aligned_error"], atol=0.05) + assert pae_record["max_predicted_aligned_error"] == pytest.approx(31.75) + + plddt = json.loads((tmp_path / f"confidence_{name}.json").read_text()) + assert plddt["residueNumber"] == list(range(1, n_real + 1)) + np.testing.assert_allclose(plddt["confidenceScore"], expected["plddt"], atol=0.005) + + with open(tmp_path / f"result_{name}.pkl", "rb") as handle: + saved = pickle.load(handle) + assert saved["iptm"] == pytest.approx(expected["iptm"], abs=1e-6) + assert saved["ptm"] == pytest.approx(expected["ptm"], abs=1e-6) + assert saved["predicted_aligned_error"].shape == (n_real, n_real) + assert saved["aligned_confidence_probs"].shape == (n_real, n_real, 64) + assert saved["plddt"].shape == (n_real,) + + def test_postprocess_ranks_models_relaxes_best_and_runs_cleanup( af2_backend_module, monkeypatch, From c1b12f02b1c2838cdfe56d6c278b4bf84e5c1ba2 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 07:19:07 +0200 Subject: [PATCH 05/14] test: real JAX compile cache shared between processes test_jax_compilation_cache.py checks the jax.config calls with a fake. This integration test runs real JAX on CPU in fresh interpreters through enable_persistent_compilation_cache: a second process with the same cache directory gets cache hits, no misses and writes nothing new. With JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES unset (AlphaPulldown's default) or "none", JAX hands XLA no autotune or kernel cache paths and the directory holds only JAX's entries; "all" in the environment is respected, which shows the probe can see them. ~10 s; passes on jax 0.5.3, 0.9.1, 0.10.2. Co-Authored-By: Claude Opus 5.5 --- .../test_jax_compilation_cache_real.py | 125 ++++++++++++++++++ 1 file changed, 125 insertions(+) create mode 100644 test/integration/test_jax_compilation_cache_real.py diff --git a/test/integration/test_jax_compilation_cache_real.py b/test/integration/test_jax_compilation_cache_real.py new file mode 100644 index 00000000..07bdcf99 --- /dev/null +++ b/test/integration/test_jax_compilation_cache_real.py @@ -0,0 +1,125 @@ +"""AlphaPulldown's persistent JAX compile cache with real JAX on CPU. + +Each case runs a fresh interpreter that turns the cache on through +``enable_persistent_compilation_cache`` and calls one tiny jitted function, so the +second process of a pair must load it from disk instead of compiling it. The unit tests +in test/unit/test_jax_compilation_cache.py only check the configuration calls. +""" + +from __future__ import annotations + +import importlib.util +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +if importlib.util.find_spec("jaxlib") is None: # conftest may stub jax itself + pytest.skip("needs a real JAX installation", allow_module_level=True) + +REPO_ROOT = Path(__file__).resolve().parents[2] +XLA_CACHES_ENV = "JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES" + +_PROBE = r""" +import json, sys +import jax +import jax.numpy as jnp +from jax._src import compiler + +events = [] +jax.monitoring.register_event_listener(lambda event, **kwargs: events.append(event)) + +from alphapulldown.prediction.jax_compilation_cache import ( + enable_persistent_compilation_cache, +) + +used = enable_persistent_compilation_cache(sys.argv[1]) +forward = jax.jit(lambda x: jnp.tanh(x @ x.T).sum()) +forward(jnp.ones((8, 8))).block_until_ready() + +# The XLA cache paths JAX hands the compiler when the persistent cache is on. Keep +# every object alive: the protobuf accessors return views into their parent. +options = compiler.get_compile_options(num_replicas=1, num_partitions=1) +build = options.executable_build_options +debug = build.debug_options +print("PROBE " + json.dumps({ + "cache_dir": used, + "hits": events.count("/jax/compilation_cache/cache_hits"), + "misses": events.count("/jax/compilation_cache/cache_misses"), + "xla_caches": jax.config.jax_persistent_cache_enable_xla_caches, + "autotune_cache_dir": debug.xla_gpu_per_fusion_autotune_cache_dir, + "kernel_cache_file": debug.xla_gpu_kernel_cache_file, +})) +""" + + +def _run_probe(cache_dir: Path, xla_caches: str | None = None) -> dict: + env = os.environ.copy() + for name in (XLA_CACHES_ENV, "JAX_COMPILATION_CACHE_DIR", "PYTEST_ADDOPTS"): + env.pop(name, None) + if xla_caches is not None: + env[XLA_CACHES_ENV] = xla_caches + env["JAX_PLATFORMS"] = "cpu" + env["PYTHONPATH"] = os.pathsep.join( + [str(REPO_ROOT), *filter(None, [os.environ.get("PYTHONPATH")])] + ) + result = subprocess.run( + [sys.executable, "-c", _PROBE, str(cache_dir)], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + timeout=180, + ) + assert result.returncode == 0, result.stdout + result.stderr + line = next(l for l in result.stdout.splitlines() if l.startswith("PROBE ")) + return json.loads(line[len("PROBE "):]) + + +def _entries(cache_dir: Path) -> dict[str, bool]: + """Name -> is a regular file, for everything in the cache directory.""" + return {path.name: path.is_file() for path in cache_dir.iterdir()} + + +def test_a_second_process_loads_the_compiled_function_from_the_cache(tmp_path): + cache_dir = tmp_path / "jax-cache" + + first = _run_probe(cache_dir) + written = _entries(cache_dir) + second = _run_probe(cache_dir) + + assert first["cache_dir"] == str(cache_dir) + assert first["hits"] == 0 and first["misses"] >= 1 + assert written and all(written.values()) + assert second["hits"] >= 1 and second["misses"] == 0 + assert _entries(cache_dir) == written # nothing compiled, nothing new written + + +@pytest.mark.parametrize("xla_caches", [None, "none"], ids=["default", "env-none"]) +def test_only_jax_entries_are_written_without_xla_caches(tmp_path, xla_caches): + # XLA's autotune cache segfaults the compiler on shared filesystems, so AlphaPulldown + # turns XLA's caches off unless the environment chooses otherwise. + cache_dir = tmp_path / "jax-cache" + + probe = _run_probe(cache_dir, xla_caches) + + assert probe["xla_caches"] == "none" + assert probe["autotune_cache_dir"] == "" + assert probe["kernel_cache_file"] == "" + entries = _entries(cache_dir) + assert entries and all(entries.values()) + assert not [name for name in entries if name.startswith("xla_gpu_")] + + +def test_the_environment_can_still_turn_xla_caches_on(tmp_path): + # The probe above would see them: with "all", JAX points XLA into the cache dir. + cache_dir = tmp_path / "jax-cache" + + probe = _run_probe(cache_dir, "all") + + assert probe["xla_caches"] == "all" + assert probe["autotune_cache_dir"] == str(cache_dir / "xla_gpu_per_fusion_autotune_cache_dir") + assert probe["kernel_cache_file"] == str(cache_dir / "xla_gpu_kernel_cache_file") From b5cef3fbe4769fdf44e41947fb410d58732a1a95 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 07:20:32 +0200 Subject: [PATCH 06/14] test: AF3 fused-triangle settings are fields the fork's GlobalConfig declares AlphaFold3Backend.setup copies FusedTriangleChoice.global_config_update() onto config.global_config with setattr, so a renamed or mistyped field, or a value outside the fork's Literal, would silently do nothing. This parses the submodule's model_config.py with ast (no AF3 install needed) and checks, for every compute capability the fork's dispatch has measured, that each key device_settings returns is a GlobalConfig field with a value its annotation allows; that the triangle implementations default to "default" so off changes nothing; and that the fields metadata() reads are declared. Co-Authored-By: Claude Opus 5.5 --- test/unit/test_af3_fused_triangles.py | 91 +++++++++++++++++++++++++++ 1 file changed, 91 insertions(+) diff --git a/test/unit/test_af3_fused_triangles.py b/test/unit/test_af3_fused_triangles.py index c424f0a2..e71fd8ce 100644 --- a/test/unit/test_af3_fused_triangles.py +++ b/test/unit/test_af3_fused_triangles.py @@ -2,6 +2,7 @@ from __future__ import annotations +import ast import contextlib import importlib.util import sys @@ -17,6 +18,9 @@ Path(__file__).resolve().parents[2] / "alphafold3/src/alphafold3/jax/fused_triangle/dispatch.py" ) +MODEL_CONFIG_PATH = ( + Path(__file__).resolve().parents[2] / "alphafold3/src/alphafold3/model/model_config.py" +) SETTINGS = { "triangle_multiplication_implementation": "pallas", "triangle_attention_implementation": "auto", @@ -284,3 +288,90 @@ def test_metadata_float32_layers_run_the_original_body(dispatch): assert {operation["reason"] for operation in record["operations"].values()} == { "dtype_not_bfloat16" } + + +def _global_config_fields() -> dict[str, tuple[ast.expr, ast.expr | None]]: + """GlobalConfig's annotated fields in the submodule's model_config.py, without AF3.""" + tree = ast.parse(MODEL_CONFIG_PATH.read_text(encoding="utf-8")) + (global_config,) = [ + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == "GlobalConfig" + ] + return { + node.target.id: (node.annotation, node.value) + for node in global_config.body + if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) + } + + +def _literal_values(annotation: ast.expr) -> set | None: + """The allowed values of a ``Literal[...]`` annotation, or None for any other type.""" + if not ( + isinstance(annotation, ast.Subscript) + and isinstance(annotation.value, (ast.Name, ast.Attribute)) + and getattr(annotation.value, "id", getattr(annotation.value, "attr", None)) + == "Literal" + ): + return None + values = annotation.slice + elements = values.elts if isinstance(values, ast.Tuple) else [values] + return {ast.literal_eval(element) for element in elements} + + +def _assert_declared(fields, key, value): + assert key in fields, f"GlobalConfig declares no {key!r}" + annotation, _ = fields[key] + allowed = _literal_values(annotation) + if allowed is not None: + assert value in allowed, f"{key}={value!r} is not one of {sorted(allowed)}" + elif isinstance(annotation, ast.Name) and annotation.id in ("str", "float", "int", "bool"): + types = {"str": str, "float": (int, float), "int": int, "bool": bool}[annotation.id] + assert isinstance(value, types), f"{key}={value!r} is not a {annotation.id}" + + +def _measured_devices(): + spec = importlib.util.spec_from_file_location("af3_dispatch_contract", DISPATCH_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module._MEASURED_COMPUTE_CAPABILITIES + + +@pytest.mark.parametrize("compute_capability", _measured_devices()) +def test_kernel_settings_are_declared_by_the_forks_global_config( + dispatch, model_config, compute_capability +): + # AlphaFold3Backend.setup sets every key of the choice on config.global_config, and + # an undeclared one would be a silent no-op. Parsed, so it needs no AF3 install. + fields = _global_config_fields() + device = types.SimpleNamespace( + compute_capability=compute_capability, + device_kind="GPU", + memory_stats=lambda: {"bytes_limit": 40 * 2**30}, + ) + settings = af3_fused_triangles.device_settings(device) + choice = af3_fused_triangles.FusedTriangleChoice(True, "--fast_kernels=on", "on", settings) + + assert choice.global_config_update() == settings + for key, value in settings.items(): + _assert_declared(fields, key, value) + # Off changes nothing, so the fork's defaults must be the original layers. + assert af3_fused_triangles.OFF.global_config_update() == {} + for operation in ("triangle_multiplication", "triangle_attention"): + _, default = fields[f"{operation}_implementation"] + assert ast.literal_eval(default) == "default" + + +def test_kernel_record_reads_fields_the_forks_global_config_declares(): + # metadata() reads these from the runner's config to record what ran. + fields = _global_config_fields() + + for key in ( + "triangle_multiplication_implementation", + "triangle_attention_implementation", + "fused_triangle_compute_capability", + "fused_triangle_memory_gib", + "bfloat16", + ): + assert key in fields, key + assert "none" in _literal_values(fields["bfloat16"][0]) From 9229093e1bbcd345f6efe564c2d871dfd1d90be1 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 07:29:28 +0200 Subject: [PATCH 07/14] test: AF2 GPU checks for --fast_kernels and padded batches TestFastKernelsAndPadding in the AF2 cluster suite (run manually with RUN_GPU_FUNCTIONAL_TESTS=1; skipped on CI or without an NVIDIA GPU) folds the TEST homodimer with model_1_multimer_v3: - --fast_kernels=off and auto: inference_kernels.json records the stock code and then the kernels with the GPU's compute capability, so auto cannot fall back silently; ipTM and pTM differ by less than 0.02. - a monomer with --fast_kernels=auto keeps the stock code and says so. - the batch command with --desired_num_res=160 against the unpadded fold: pickle, PAE JSON and ranking_debug.json cover the 128 real residues and ipTM, pTM and ranking confidence agree within 0.02. The kernel cases skip when colabfold-kernels is missing or the GPU is older than compute capability 8.0. Co-Authored-By: Claude Opus 5.5 --- test/cluster/check_alphafold2_predictions.py | 187 +++++++++++++++++++ 1 file changed, 187 insertions(+) diff --git a/test/cluster/check_alphafold2_predictions.py b/test/cluster/check_alphafold2_predictions.py index a9a7715f..305c211e 100755 --- a/test/cluster/check_alphafold2_predictions.py +++ b/test/cluster/check_alphafold2_predictions.py @@ -6,7 +6,9 @@ from __future__ import annotations import os +import importlib.util import json +import math import pickle import shutil import subprocess @@ -74,6 +76,43 @@ def _gpu_functional_test_skip_reason() -> str | None: return None +def _gpu_compute_capability() -> float | None: + """The first visible NVIDIA GPU's compute capability, from nvidia-smi.""" + nvidia_smi = shutil.which("nvidia-smi") + if not nvidia_smi: + return None + try: + result = subprocess.run( + [nvidia_smi, "--query-gpu=compute_cap", "--format=csv,noheader"], + capture_output=True, + text=True, + check=False, + ) + return float(result.stdout.split()[0]) + except (OSError, IndexError, ValueError): + return None + + +def _fast_kernels_skip_reason() -> str | None: + """Why --fast_kernels cannot turn the kernels on here, or None if it should.""" + if importlib.util.find_spec("colabfold_kernels") is None: + return "colabfold-kernels is not installed (pip install 'alphapulldown[fast-kernels]')" + capability = _gpu_compute_capability() + if capability is None or capability < 8.0: + return f"the fused kernels need compute capability 8.0 or newer, not {capability}" + return None + + +def _af2_result_dir(root: Path) -> Path: + """The one directory under ``root`` holding an AF2 prediction.""" + if (root / "ranking_debug.json").exists(): + return root + candidates = sorted(path.parent for path in root.rglob("ranking_debug.json")) + if len(candidates) != 1: + raise AssertionError(f"Expected one AF2 result directory under {root}, found {candidates}") + return candidates[0] + + def _mmseqs_functional_test_skip_reason() -> str | None: if os.getenv("RUN_MMSEQS_FUNCTIONAL_TESTS", "").lower() in ("1", "true", "yes"): return None @@ -520,6 +559,154 @@ def test_dropout_increases_diversity(self): # The test passes if calculations succeed - the diversity check is informational +class TestFastKernelsAndPadding(_TestBase): + """--fast_kernels and --desired_num_res on a GPU, against the stock unpadded fold. + + Folds the TEST homodimer (2 x 64 residues) with one multimer model, so each case + runs a few short predictions. Select with -k TestFastKernelsAndPadding. + """ + + DIMER = "TEST+TEST" + NUM_REAL_RESIDUES = 128 + PADDED_NUM_RES = 160 + MULTIMER_MODEL = "model_1_multimer_v3" + # Fused and stock kernels, or padded and unpadded shapes, differ in their numerics. + SCORE_TOLERANCE = 0.02 + + def _predict(self, name: str, fold: str, *extra: str) -> Path: + output_dir = self.output_dir / name + output_dir.mkdir(parents=True, exist_ok=True) + res = self._run_prediction_subprocess( + [ + sys.executable, + str(self.script_single), + f"--input={fold}", + f"--output_directory={output_dir}", + "--num_cycle=3", + "--num_predictions_per_model=1", + "--random_seed=42", + f"--data_directory={DATA_DIR}", + f"--features_directory={self.test_features_dir}", + *extra, + ] + ) + self.assertEqual( + res.returncode, 0, f"{name} failed.\nSTDOUT:\n{res.stdout}\nSTDERR:\n{res.stderr}" + ) + self.last_log = res.stdout + res.stderr + return _af2_result_dir(output_dir) + + def _predict_padded_batch(self, name: str, fold: str, *extra: str) -> Path: + """Fold through the batch command with --desired_num_res, as the workflow does.""" + output_dir = self.output_dir / name + output_dir.mkdir(parents=True, exist_ok=True) + manifest = self.output_dir / f"{name}.jsonl" + manifest.write_text( + json.dumps({"job_id": name, "input": fold, "output_directory": str(output_dir)}) + + "\n", + encoding="utf-8", + ) + res = self._run_prediction_subprocess( + [ + sys.executable, + str(self.script_single.parent / "run_structure_prediction_batch.py"), + f"--manifest={manifest}", + "--num_cycle=3", + "--num_predictions_per_model=1", + "--random_seed=42", + f"--data_directory={DATA_DIR}", + f"--features_directory={self.test_features_dir}", + *extra, + ] + ) + self.assertEqual( + res.returncode, 0, f"{name} failed.\nSTDOUT:\n{res.stdout}\nSTDERR:\n{res.stderr}" + ) + return _af2_result_dir(output_dir) + + @staticmethod + def _kernels(result_dir: Path) -> dict: + return json.loads((result_dir / "inference_kernels.json").read_text(encoding="utf-8")) + + @staticmethod + def _scores(result_dir: Path) -> dict: + (pickle_path,) = sorted(result_dir.glob("result_*.pkl")) + with pickle_path.open("rb") as handle: + result = pickle.load(handle) + return { + "iptm": float(result["iptm"]), + "ptm": float(result["ptm"]), + "ranking_confidence": float(result["ranking_confidence"]), + "mean_plddt": float(sum(result["plddt"]) / len(result["plddt"])), + "num_residues": len(result["plddt"]), + "pae_shape": tuple(result["predicted_aligned_error"].shape), + } + + def test_dimer_fast_kernels_match_the_stock_code(self): + skip_reason = _fast_kernels_skip_reason() + if skip_reason: + self.skipTest(skip_reason) + model = f"--model_names={self.MULTIMER_MODEL}" + + stock = self._predict("stock", self.DIMER, model, "--fast_kernels=off") + fused = self._predict("fused", self.DIMER, model, "--fast_kernels=auto") + + model_name = f"{self.MULTIMER_MODEL}_pred_0" + self.assertEqual( + self._kernels(stock), + {model_name: {"compute_capability": None, "fused_kernels": False}}, + ) + fused_record = self._kernels(fused) + self.assertEqual(set(fused_record), {model_name}) + # auto would fall back silently; the requirements are met, so it must not. + self.assertTrue(fused_record[model_name]["fused_kernels"], fused_record) + self.assertGreaterEqual(fused_record[model_name]["compute_capability"], 80) + + stock_scores, fused_scores = self._scores(stock), self._scores(fused) + logger.info("stock %s, fused %s", stock_scores, fused_scores) + self.assertLess(abs(fused_scores["iptm"] - stock_scores["iptm"]), self.SCORE_TOLERANCE) + self.assertLess(abs(fused_scores["ptm"] - stock_scores["ptm"]), self.SCORE_TOLERANCE) + + def test_monomer_models_keep_the_stock_code(self): + skip_reason = _fast_kernels_skip_reason() + if skip_reason: + self.skipTest(skip_reason) + + result_dir = self._predict( + "monomer", "TEST", "--model_names=model_1_ptm", "--fast_kernels=auto" + ) + + record = self._kernels(result_dir) + self.assertTrue(record) + self.assertFalse(any(entry["fused_kernels"] for entry in record.values()), record) + self.assertIn("monomer models run in fp32 and keep the standard code", self.last_log) + + def test_padded_batch_scores_match_the_unpadded_fold(self): + model = f"--model_names={self.MULTIMER_MODEL}" + + unpadded = self._predict("unpadded", self.DIMER, model) + padded = self._predict_padded_batch( + "padded", self.DIMER, model, f"--desired_num_res={self.PADDED_NUM_RES}" + ) + + n = self.NUM_REAL_RESIDUES + expected, got = self._scores(unpadded), self._scores(padded) + logger.info("unpadded %s, padded %s", expected, got) + self.assertEqual(got["num_residues"], n) + self.assertEqual(got["pae_shape"], (n, n)) + for key in ("iptm", "ptm", "ranking_confidence"): + self.assertTrue(math.isfinite(got[key]), key) + self.assertLess(abs(got[key] - expected[key]), self.SCORE_TOLERANCE, key) + self.assertLess(abs(got["mean_plddt"] - expected["mean_plddt"]), 1.0) + # The JSON written for viewers and ranking covers the real residues only. + for pae_path in padded.glob("pae_*.json"): + (pae,) = json.loads(pae_path.read_text(encoding="utf-8")) + self.assertEqual(len(pae["predicted_aligned_error"]), n, pae_path) + ranking = json.loads((padded / "ranking_debug.json").read_text(encoding="utf-8")) + for name, score in ranking["iptm+ptm"].items(): + self.assertLess(abs(score - expected["ranking_confidence"]), self.SCORE_TOLERANCE, name) + + class TestMmseqsIssue588Inference(_TestBase): """Opt-in end-to-end regression for freshly generated mmseq AF2 features.""" From 2205c7e01bdbec025079fdf52cfd94354373df2b Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 07:29:28 +0200 Subject: [PATCH 08/14] test: AF3 GPU checks for --fast_kernels and one compile per process TestAlphaFold3FastKernels in the AF3 cluster suite (RUN_GPU_FUNCTIONAL_TESTS=1; skipped on CI or without an NVIDIA GPU): - --fast_kernels=off, on and auto each fold A0A075B6L2 with two seeds and check inference_kernels.json has one seed-N record per seed with the requested mode, every confidence JSON is finite, off records the original layers, and on records the fused layers per operation. on skips with the backend's reason on a GPU the fork's policy does not enable; auto must then record the fallback, and otherwise the fused layers. - three identical folds through the batch command in one process, with JAX_LOG_COMPILES=1 and the persistent cache off, compile and trace jit(apply_fn) exactly once (the 1afe779e regression; matches the manual retrace_fixed.sbatch check). Co-Authored-By: Claude Opus 5.5 --- test/cluster/check_alphafold3_predictions.py | 167 +++++++++++++++++++ 1 file changed, 167 insertions(+) diff --git a/test/cluster/check_alphafold3_predictions.py b/test/cluster/check_alphafold3_predictions.py index 2231a758..1f7f5823 100755 --- a/test/cluster/check_alphafold3_predictions.py +++ b/test/cluster/check_alphafold3_predictions.py @@ -4112,6 +4112,173 @@ def test_af3_rejects_alphafold2_flag(self): ) +def _non_finite_values(payload: Any, where: str = "") -> list[str]: + """Paths of NaN or infinite numbers in a parsed JSON payload.""" + if isinstance(payload, dict): + items = ((f"{where}.{key}", value) for key, value in payload.items()) + elif isinstance(payload, list): + items = ((f"{where}[{index}]", value) for index, value in enumerate(payload)) + else: + items = () + nested = [bad for path, value in items for bad in _non_finite_values(value, path)] + if nested: + return nested + if isinstance(payload, float) and not np.isfinite(payload): + return [where] + return [] + + +class TestAlphaFold3FastKernels(_TestBase): + """--fast_kernels for AF3 on a GPU, and one model compile per process. + + Select with -k TestAlphaFold3FastKernels. ``on`` skips with the backend's reason on + a GPU the fork's policy does not enable; ``auto`` must then fall back. + """ + + FOLD = "A0A075B6L2:1" + UNAVAILABLE = "AF3 fused triangle kernels are unavailable" + + def _fold(self, mode: str) -> subprocess.CompletedProcess: + self._require_af3_functional_environment() + return subprocess.run( + [ + sys.executable, + str(self.script_single), + f"--input={self.FOLD}", + f"--output_directory={self.output_dir}", + f"--data_directory={DATA_DIR}", + f"--features_directory={self.test_features_dir}", + "--fold_backend=alphafold3", + f"--flash_attention_implementation={self._af3_flash_attention_impl()}", + "--num_diffusion_samples=1", + "--num_seeds=2", + f"--fast_kernels={mode}", + ], + capture_output=True, + text=True, + env=self._make_af3_test_env(), + ) + + def _kernel_records(self, mode: str) -> dict[str, dict[str, Any]]: + """inference_kernels.json of a finished fold: one record per seed, outputs finite.""" + result_dir = self._resolve_single_af3_result_dir() + with (result_dir / "ranking_scores.csv").open() as handle: + seeds = {line.split(",")[0] for line in handle.readlines()[1:] if line.strip()} + records = json.loads((result_dir / "inference_kernels.json").read_text()) + self.assertEqual(set(records), {f"seed-{seed}" for seed in seeds}) + self.assertLen(records, 2) + for key, record in records.items(): + self.assertEqual(record["backend"], "alphafold3", key) + self.assertEqual(record["requested_mode"], mode, key) + self.assertIsInstance(record["padded_tokens"], int, key) + confidences = sorted(result_dir.rglob("*confidences.json")) + self.assertTrue(confidences, f"no confidences in {result_dir}") + for path in confidences: + self.assertEqual(_non_finite_values(json.loads(path.read_text())), [], path) + return records + + def _assert_fused(self, records: dict[str, dict[str, Any]]) -> None: + for key, record in records.items(): + self.assertTrue(record["fused_kernels"], (key, record)) + self.assertIn("policy_version", record, key) + self.assertEqual( + set(record["operations"]), + { + f"{operation}_c{channels}" + for operation in ("triangle_multiplication", "triangle_attention") + for channels in (128, 64) + }, + key, + ) + + def test_af3_fast_kernels_off_keeps_the_original_layers(self): + res = self._fold("off") + self._runCommonTests(res) + + for key, record in self._kernel_records("off").items(): + self.assertFalse(record["fused_kernels"], key) + self.assertEqual(record["reason"], "--fast_kernels=off", key) + self.assertNotIn("operations", record, key) + + def test_af3_fast_kernels_on_runs_the_fused_layers(self): + res = self._fold("on") + log = res.stdout + res.stderr + if res.returncode != 0 and f"--fast_kernels=on, but {self.UNAVAILABLE}" in log: + reason = next(line for line in log.splitlines() if self.UNAVAILABLE in line) + self.skipTest(f"this GPU cannot run the fused triangles: {reason.strip()}") + self._runCommonTests(res) + + self._assert_fused(self._kernel_records("on")) + + def test_af3_fast_kernels_auto_uses_them_where_the_gpu_can(self): + res = self._fold("auto") + self._runCommonTests(res) + + records = self._kernel_records("auto") + if any(record["fused_kernels"] for record in records.values()): + self._assert_fused(records) + else: + for key, record in records.items(): + self.assertTrue( + record["reason"].startswith(f"--fast_kernels=auto: {self.UNAVAILABLE}"), + (key, record), + ) + + def test_af3_identical_folds_compile_the_model_once(self): + """Three identical folds in one process: one jit(apply_fn) compile (1afe779e).""" + self._require_af3_functional_environment() + manifest = self.output_dir / "manifest.jsonl" + jobs = [f"repeat_{index}" for index in range(1, 4)] + manifest.write_text( + "".join( + json.dumps( + { + "job_id": job, + "input": self.FOLD, + "output_directory": str(self.output_dir / job), + } + ) + + "\n" + for job in jobs + ) + ) + env = self._make_af3_test_env() + env["JAX_LOG_COMPILES"] = "1" + + res = subprocess.run( + [ + sys.executable, + str(self.script_single.parent / "run_structure_prediction_batch.py"), + f"--manifest={manifest}", + f"--data_directory={DATA_DIR}", + f"--features_directory={self.test_features_dir}", + "--fold_backend=alphafold3", + f"--flash_attention_implementation={self._af3_flash_attention_impl()}", + "--num_diffusion_samples=1", + # A cache load would hide a recompile. + "--jax_compilation_cache_dir=none", + ], + capture_output=True, + text=True, + env=env, + ) + self.assertEqual( + res.returncode, 0, f"batch failed.\nSTDOUT:\n{res.stdout}\nSTDERR:\n{res.stderr}" + ) + + log = res.stdout + res.stderr + for job in jobs: + self.assertTrue( + list((self.output_dir / job).rglob("ranking_scores.csv")), + f"{job} wrote no predictions", + ) + self.assertEqual(log.count("Finished XLA compilation of jit(apply_fn)"), 1) + # "tracing + transforming apply_fn for pjit" on jax 0.9, "tracing apply_fn for jit" + # on jax 0.10. + traces = re.findall(r"Finished tracing (?:\+ transforming )?apply_fn for p?jit", log) + self.assertLen(traces, 1) + + # --------------------------------------------------------------------------- # def _parse_test_args(): """Parse test-specific arguments that work with both absltest and pytest.""" From b46a0e72d31b29604531923ffc5f954be7f66308 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 08:57:59 +0200 Subject: [PATCH 09/14] test: GPU test runners judge a job by pytest's result line Passing AF3 runs log a library traceback (tokamax's autotuning cache), and the runners marked any log containing "Traceback" as FAILED. pytest's last result line now decides when it is there; the old heuristics remain for jobs that died before pytest reported. Co-Authored-By: Claude Opus 5.5 --- test/cluster/run_alphafold2_predictions.py | 23 ++++++++ test/cluster/run_alphafold3_predictions.py | 23 ++++++++ .../test_cluster_runner_classification.py | 55 +++++++++++++++++++ 3 files changed, 101 insertions(+) create mode 100644 test/unit/test_cluster_runner_classification.py diff --git a/test/cluster/run_alphafold2_predictions.py b/test/cluster/run_alphafold2_predictions.py index a921ec11..f76c5192 100644 --- a/test/cluster/run_alphafold2_predictions.py +++ b/test/cluster/run_alphafold2_predictions.py @@ -411,6 +411,21 @@ def _extract_reason_from_log(text: str) -> str: return "\n".join(non_empty_lines[-20:]) +def _pytest_summary(text: str) -> str | None: + """FAILED, PASSED or SKIPPED from pytest's last result line ("== 1 passed, 2 skipped in 3.4s =="), or None.""" + lines = re.findall(r"(?m)^=+ (.+?) in [0-9.]+s(?: \([0-9:]+\))? =+\s*$", text) + if not lines: + return None + last = lines[-1] + if re.search(r"\b\d+ (failed|errors?)\b", last): + return "FAILED" + if re.search(r"\b\d+ passed\b", last): + return "PASSED" + if re.search(r"\b\d+ skipped\b", last): + return "SKIPPED" + return None + + def classify_job(job: JobSpec) -> None: job.slurm_state, job.exit_code = query_sacct(job.job_id or "") text = _combined_log_text(job) @@ -421,6 +436,14 @@ def classify_job(job: JobSpec) -> None: job.reason = f"Slurm state: {state}\n{_extract_reason_from_log(text)}" return + # pytest's own result line decides when it is there: logs of passing AF3 runs contain library tracebacks that are + # only logged (e.g. tokamax's autotuning cache), so "Traceback" alone does not mean the test failed. + pytest_summary = _pytest_summary(text) + if pytest_summary is not None: + job.outcome = pytest_summary + job.reason = _extract_reason_from_log(text) + return + if re.search(r"(?im)\bkilled\b", text): job.outcome = "FAILED" job.reason = _extract_reason_from_log(text) diff --git a/test/cluster/run_alphafold3_predictions.py b/test/cluster/run_alphafold3_predictions.py index bd0fc1c9..f7634036 100644 --- a/test/cluster/run_alphafold3_predictions.py +++ b/test/cluster/run_alphafold3_predictions.py @@ -405,6 +405,21 @@ def _extract_reason_from_log(text: str) -> str: return "\n".join(non_empty_lines[-20:]) +def _pytest_summary(text: str) -> str | None: + """FAILED, PASSED or SKIPPED from pytest's last result line ("== 1 passed, 2 skipped in 3.4s =="), or None.""" + lines = re.findall(r"(?m)^=+ (.+?) in [0-9.]+s(?: \([0-9:]+\))? =+\s*$", text) + if not lines: + return None + last = lines[-1] + if re.search(r"\b\d+ (failed|errors?)\b", last): + return "FAILED" + if re.search(r"\b\d+ passed\b", last): + return "PASSED" + if re.search(r"\b\d+ skipped\b", last): + return "SKIPPED" + return None + + def classify_job(job: JobSpec) -> None: job.slurm_state, job.exit_code = query_sacct(job.job_id or "") text = _combined_log_text(job) @@ -415,6 +430,14 @@ def classify_job(job: JobSpec) -> None: job.reason = f"Slurm state: {state}\n{_extract_reason_from_log(text)}" return + # pytest's own result line decides when it is there: logs of passing AF3 runs contain library tracebacks that are + # only logged (e.g. tokamax's autotuning cache), so "Traceback" alone does not mean the test failed. + pytest_summary = _pytest_summary(text) + if pytest_summary is not None: + job.outcome = pytest_summary + job.reason = _extract_reason_from_log(text) + return + if re.search(r"(?im)\bkilled\b", text): job.outcome = "FAILED" job.reason = _extract_reason_from_log(text) diff --git a/test/unit/test_cluster_runner_classification.py b/test/unit/test_cluster_runner_classification.py new file mode 100644 index 00000000..9d03532c --- /dev/null +++ b/test/unit/test_cluster_runner_classification.py @@ -0,0 +1,55 @@ +"""The GPU test runners judge a job by pytest's own result line, not by tracebacks a passing run happens to log.""" +import importlib.util +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[2] + + +def _load(name): + spec = importlib.util.spec_from_file_location(name, ROOT / "test" / "cluster" / f"{name}.py") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module # dataclasses look their module up while the class is created + spec.loader.exec_module(module) + return module + + +@pytest.fixture(params=["run_alphafold2_predictions", "run_alphafold3_predictions"]) +def runner(request): + return _load(request.param) + + +def _classify(runner, monkeypatch, text, state="COMPLETED"): + monkeypatch.setattr(runner, "query_sacct", lambda job_id: (state, "0:0")) + monkeypatch.setattr(runner, "_combined_log_text", lambda job: text) + job = runner.JobSpec(index=1, nodeid="t::x", slug="x", stdout_path=Path("o"), stderr_path=Path("e"), + script_path=Path("s"), rerun_command="", job_id="1") + runner.classify_job(job) + return job.outcome + + +LOGGED_TRACEBACK = "Traceback (most recent call last):\npydantic_core.ValidationError: ...\n" + + +def test_logged_traceback_does_not_fail_a_passing_test(runner, monkeypatch): + text = LOGGED_TRACEBACK + "PASSED\n\n======== 1 passed in 95.36s (0:01:35) ========\n" + assert _classify(runner, monkeypatch, text) == "PASSED" + + +@pytest.mark.parametrize( + "line, outcome", + [("== 1 failed in 3.1s ==", "FAILED"), ("== 1 passed, 1 error in 3.1s ==", "FAILED"), + ("== 1 skipped in 0.5s ==", "SKIPPED"), ("== 2 passed, 1 skipped in 9.0s ==", "PASSED")], +) +def test_pytest_result_line_decides(runner, monkeypatch, line, outcome): + assert _classify(runner, monkeypatch, LOGGED_TRACEBACK + line + "\n") == outcome + + +def test_without_a_result_line_a_traceback_still_fails(runner, monkeypatch): + assert _classify(runner, monkeypatch, LOGGED_TRACEBACK) == "FAILED" + + +def test_a_failed_slurm_job_fails_whatever_the_log_says(runner, monkeypatch): + assert _classify(runner, monkeypatch, "== 1 passed in 1.0s ==\n", state="FAILED") == "FAILED" From 853050d5724582c3f04d9df1b2d9e4137d347648 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 09:25:49 +0200 Subject: [PATCH 10/14] fix: tell an unsupported GPU from AF3 fused kernels that fail their self-check resolve() reported every failure of device_settings() and smoke_test() as "AF3 fused triangle kernels are unavailable", so kernels that do not compile, crash or return non-finite values on a supported GPU looked like unsupported hardware: the GPU tests skipped (on) or accepted the fallback (auto). Now a missing fork or a device the fork's policy refuses is "not supported here" (auto: warning), and a failed compile/run/finite check on a supported device is "failed their self-check" (auto: error with traceback). on still raises ValueError in both cases, naming which. Co-Authored-By: Claude Opus 5.5 --- .../prediction/af3_fused_triangles.py | 47 +++++++++++--- docs/af3_fused_triangles.md | 26 +++++--- test/cluster/check_alphafold3_predictions.py | 8 +-- test/unit/test_af3_fused_triangles.py | 62 +++++++++++++++++-- test/unit/test_alphafold3_backend_helpers.py | 37 +++++++++++ 5 files changed, 154 insertions(+), 26 deletions(-) diff --git a/alphapulldown/prediction/af3_fused_triangles.py b/alphapulldown/prediction/af3_fused_triangles.py index 3d4b4314..fde7cb58 100644 --- a/alphapulldown/prediction/af3_fused_triangles.py +++ b/alphapulldown/prediction/af3_fused_triangles.py @@ -16,6 +16,12 @@ * ``auto``: as ``on`` where the device can run them, the original layers (with a log line) elsewhere. +Two reasons are told apart. The kernels are *not supported here* when the installed AF3 +has no fused-triangle hooks or the fork's policy does not enable this device; ``auto`` +then logs a warning. They *failed their self-check* when the device is supported but +compiling or running them fails, or they return non-finite values: that is a kernel, +driver or JAX bug, so ``auto`` logs it as an error before falling back. + JAX and the fork are imported only when a mode other than ``off`` is resolved, or when the record of a model that requested the kernels is written. """ @@ -133,27 +139,54 @@ def resolve(mode: Any, device: jax.Device) -> FusedTriangleChoice: """Turn a ``--fast_kernels`` value into a decision for ``device``, checking it first. Raises: - ValueError: for an unknown mode, or for ``on`` when ``device`` cannot run the - kernels. + ValueError: for an unknown mode, or for ``on`` when the kernels are not supported + on ``device`` or fail their self-check there; the message says which. """ mode = normalise_mode(mode) if mode == "off": return OFF try: settings = device_settings(device) + except Exception as exc: # no hooks, or a device the fork's policy does not enable + return _fall_back( + mode, f"AF3 fused triangle kernels are not supported here: {exc}", exc + ) + try: smoke_test(settings, device) except Exception as exc: # any failure here would recur inside the model - problem = f"AF3 fused triangle kernels are unavailable: {exc}" - if mode == "on": - raise ValueError(f"--fast_kernels=on, but {problem}.") from exc - logging.warning("Fused kernels off (--fast_kernels=auto): %s.", problem) - return FusedTriangleChoice(False, f"--fast_kernels=auto: {problem}", mode) + return _fall_back( + mode, + f"AF3 fused triangle kernels failed their self-check: " + f"{type(exc).__name__}: {exc}", + exc, + self_check=True, + ) logging.info( "AF3 fused triangle kernels on; per-operation size limits apply: %s", settings ) return FusedTriangleChoice(True, f"--fast_kernels={mode}", mode, settings) +def _fall_back( + mode: str, problem: str, exc: Exception, *, self_check: bool = False +) -> FusedTriangleChoice: + """Raise for ``on``; for ``auto``, log ``problem`` and keep the original layers. + + A device without support is expected, so it is a warning. Kernels that fail their + self-check on a supported device are a bug that must not pass unnoticed, so that is + an error, with the traceback. + """ + if mode == "on": + raise ValueError(f"--fast_kernels=on, but {problem}.") from exc + if self_check: + logging.error( + "Fused kernels off (--fast_kernels=auto): %s.", problem, exc_info=exc + ) + else: + logging.warning("Fused kernels off (--fast_kernels=auto): %s.", problem) + return FusedTriangleChoice(False, f"--fast_kernels=auto: {problem}", mode) + + def metadata( global_config: Any, choice: FusedTriangleChoice, *, num_tokens: int ) -> dict[str, Any]: diff --git a/docs/af3_fused_triangles.md b/docs/af3_fused_triangles.md index 0b0d07c7..277e6718 100644 --- a/docs/af3_fused_triangles.md +++ b/docs/af3_fused_triangles.md @@ -138,15 +138,23 @@ logged. ## Troubleshooting `--fast_kernels=on` stops before loading weights and gives the reason; `auto` logs the -same reason as `Fused kernels off (--fast_kernels=auto): ...`: - -| reason contains | meaning | -| --- | --- | -| `no fused-triangle hooks` | the installed `alphafold3` predates the kernels; rebuild it from the submodule | -| `unvalidated_compute_capability` | the GPU's architecture was not measured (table above) | -| `unknown_memory_budget` | JAX reports no allocator limit for the GPU | -| `memory_budget_below_12_gib` | JAX may allocate less than 12 GiB | -| anything else | compiling or running the small check on this GPU failed | +same reason as `Fused kernels off (--fast_kernels=auto): ...` and keeps the original +layers. The reason starts with one of two phrases: + +- `AF3 fused triangle kernels are not supported here: ...` (a warning under `auto`): + this installation or GPU is not one the kernels are enabled on. + + | reason contains | meaning | + | --- | --- | + | `no fused-triangle hooks` | the installed `alphafold3` predates the kernels; rebuild it from the submodule | + | `unvalidated_compute_capability` | the GPU's architecture was not measured (table above) | + | `unknown_memory_budget` | JAX reports no allocator limit for the GPU | + | `memory_budget_below_12_gib` | JAX may allocate less than 12 GiB | + +- `AF3 fused triangle kernels failed their self-check: ...` (an error, with the + traceback, under `auto`): the GPU is supported, but compiling or running the small + check on it failed, or the kernels returned non-finite values. That points to a + kernel, driver or JAX problem; please report it with the log. During a prediction, a warning such as `triangle_attention_implementation='auto' was requested, but a layer ... runs the default implementation (size_limit)` is expected diff --git a/test/cluster/check_alphafold3_predictions.py b/test/cluster/check_alphafold3_predictions.py index 1f7f5823..297222f3 100755 --- a/test/cluster/check_alphafold3_predictions.py +++ b/test/cluster/check_alphafold3_predictions.py @@ -4136,7 +4136,7 @@ class TestAlphaFold3FastKernels(_TestBase): """ FOLD = "A0A075B6L2:1" - UNAVAILABLE = "AF3 fused triangle kernels are unavailable" + NOT_SUPPORTED = "AF3 fused triangle kernels are not supported here" def _fold(self, mode: str) -> subprocess.CompletedProcess: self._require_af3_functional_environment() @@ -4203,8 +4203,8 @@ def test_af3_fast_kernels_off_keeps_the_original_layers(self): def test_af3_fast_kernels_on_runs_the_fused_layers(self): res = self._fold("on") log = res.stdout + res.stderr - if res.returncode != 0 and f"--fast_kernels=on, but {self.UNAVAILABLE}" in log: - reason = next(line for line in log.splitlines() if self.UNAVAILABLE in line) + if res.returncode != 0 and f"--fast_kernels=on, but {self.NOT_SUPPORTED}" in log: + reason = next(line for line in log.splitlines() if self.NOT_SUPPORTED in line) self.skipTest(f"this GPU cannot run the fused triangles: {reason.strip()}") self._runCommonTests(res) @@ -4220,7 +4220,7 @@ def test_af3_fast_kernels_auto_uses_them_where_the_gpu_can(self): else: for key, record in records.items(): self.assertTrue( - record["reason"].startswith(f"--fast_kernels=auto: {self.UNAVAILABLE}"), + record["reason"].startswith(f"--fast_kernels=auto: {self.NOT_SUPPORTED}"), (key, record), ) diff --git a/test/unit/test_af3_fused_triangles.py b/test/unit/test_af3_fused_triangles.py index e71fd8ce..2a5bb925 100644 --- a/test/unit/test_af3_fused_triangles.py +++ b/test/unit/test_af3_fused_triangles.py @@ -90,21 +90,71 @@ def test_unknown_mode_is_rejected(): af3_fused_triangles.resolve("fast", object()) -@pytest.mark.parametrize("stage", ["device_settings", "smoke_test"]) -def test_on_fails_loudly_and_auto_falls_back(monkeypatch, stage): +NOT_SUPPORTED = "AF3 fused triangle kernels are not supported here: " +SELF_CHECK_FAILED = "AF3 fused triangle kernels failed their self-check: RuntimeError: " + + +@pytest.mark.parametrize( + "stage, problem, level", + [ + # No hooks, or a device the fork's policy does not enable: expected, a warning. + ("device_settings", NOT_SUPPORTED, "WARNING"), + # A supported device whose kernels do not compile, crash or return NaN: a bug. + ("smoke_test", SELF_CHECK_FAILED, "ERROR"), + ], +) +def test_on_fails_loudly_and_auto_falls_back(monkeypatch, caplog, stage, problem, level): monkeypatch.setattr(af3_fused_triangles, "device_settings", lambda device: SETTINGS) monkeypatch.setattr(af3_fused_triangles, "smoke_test", lambda *args: None) monkeypatch.setattr(af3_fused_triangles, stage, _unavailable) - with pytest.raises(ValueError, match="--fast_kernels=on, but .*unavailable test device"): + with pytest.raises(ValueError) as raised: af3_fused_triangles.resolve("on", object()) + assert str(raised.value) == f"--fast_kernels=on, but {problem}unavailable test device." - choice = af3_fused_triangles.resolve("auto", object()) + with caplog.at_level("INFO"): + choice = af3_fused_triangles.resolve("auto", object()) assert not choice.enabled assert choice.mode == "auto" - assert choice.reason.startswith("--fast_kernels=auto: ") - assert "unavailable test device" in choice.reason + assert choice.reason == f"--fast_kernels=auto: {problem}unavailable test device" assert choice.global_config_update() == {} + (record,) = [r for r in caplog.records if "Fused kernels off" in r.getMessage()] + assert record.levelname == level + assert record.getMessage() == ( + f"Fused kernels off (--fast_kernels=auto): {problem}unavailable test device." + ) + # An error carries the traceback of the failed self-check; a warning needs none. + assert bool(record.exc_info) == (level == "ERROR") + + +def test_a_device_without_support_is_not_given_a_self_check(monkeypatch): + monkeypatch.setattr(af3_fused_triangles, "device_settings", _unavailable) + monkeypatch.setattr(af3_fused_triangles, "smoke_test", _fail) + + choice = af3_fused_triangles.resolve("auto", object()) + + assert choice.reason.startswith(f"--fast_kernels=auto: {NOT_SUPPORTED}") + + +@pytest.mark.parametrize( + "compute_capability, memory_stats, reason", + [ + ("10.0", lambda: {"bytes_limit": 80 * 2**30}, "unvalidated_compute_capability"), + ("8.0", lambda: {"bytes_limit": 8 * 2**30}, "memory_budget_below_12_gib"), + ], +) +def test_a_device_the_policy_refuses_is_not_supported_here( + dispatch, model_config, monkeypatch, compute_capability, memory_stats, reason +): + monkeypatch.setattr(af3_fused_triangles, "smoke_test", _fail) + device = types.SimpleNamespace( + compute_capability=compute_capability, device_kind="GPU", memory_stats=memory_stats + ) + + with pytest.raises(ValueError) as raised: + af3_fused_triangles.resolve("on", device) + + assert str(raised.value) == f"--fast_kernels=on, but {NOT_SUPPORTED}GPU: {reason}." def test_checks_the_given_device_before_enabling(monkeypatch): diff --git a/test/unit/test_alphafold3_backend_helpers.py b/test/unit/test_alphafold3_backend_helpers.py index 0433cb3f..b3206685 100644 --- a/test/unit/test_alphafold3_backend_helpers.py +++ b/test/unit/test_alphafold3_backend_helpers.py @@ -1117,6 +1117,43 @@ def test_af3_setup_applies_the_fused_triangle_choice_for_its_gpu( assert getattr(runner.config.global_config, key) == value +def _raise(message): + def fail(*args): + raise RuntimeError(message) + + return fail + + +@pytest.mark.parametrize( + "stage, problem", + [ + ("device_settings", "AF3 fused triangle kernels are not supported here"), + ("smoke_test", "AF3 fused triangle kernels failed their self-check"), + ], +) +def test_af3_setup_says_whether_the_gpu_lacks_support_or_the_kernels_failed( + af3_backend_module, monkeypatch, tmp_path, stage, problem +): + fused_triangles = af3_backend_module.af3_fused_triangles + _stub_af3_setup(af3_backend_module, monkeypatch, SimpleNamespace(compute_capability=9.0)) + monkeypatch.setattr( + fused_triangles, + "device_settings", + lambda device: {"triangle_multiplication_implementation": "pallas"}, + ) + monkeypatch.setattr(fused_triangles, "smoke_test", lambda *args: None) + monkeypatch.setattr(fused_triangles, stage, _raise("test device")) + + with pytest.raises(ValueError, match=f"^--fast_kernels=on, but {problem}: .*test device"): + _af3_setup(af3_backend_module, tmp_path, fast_kernels="on") + + runner = _af3_setup(af3_backend_module, tmp_path, fast_kernels="auto") + assert not runner.fused_triangles.enabled + assert runner.fused_triangles.reason.startswith(f"--fast_kernels=auto: {problem}: ") + # The original layers: no kernel settings reach AF3's GlobalConfig. + assert vars(runner.config.global_config) == {"flash_attention_implementation": "triton"} + + def test_af3_tokamax_trace_context_is_created_by_its_own_hook( af3_backend_module, monkeypatch ): From d94114361cc05fa44cf74cbd716e629b64f51ead Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 09:28:52 +0200 Subject: [PATCH 11/14] test: AF3 fast-kernel GPU tests skip only when nvidia-smi agrees the GPU is unsupported The on test skipped, and the auto test accepted the fallback, on the backend's word alone. Both now also ask the fork's device_policy about the GPU as nvidia-smi describes it (compute capability, and memory times the fold's JAX memory fraction; a MIG slice's memory from its profile), without importing JAX in the test process. A "not supported here" verdict on a GPU that check says is supported, or cannot describe, fails the test. The checks live in test/cluster/af3_gpu_checks.py so CPU unit tests cover them against the submodule's dispatch.py. Co-Authored-By: Claude Opus 5.5 --- test/cluster/af3_gpu_checks.py | 133 ++++++++++++++++++ test/cluster/check_alphafold3_predictions.py | 34 ++++- test/unit/test_af3_gpu_checks.py | 135 +++++++++++++++++++ 3 files changed, 299 insertions(+), 3 deletions(-) create mode 100644 test/cluster/af3_gpu_checks.py create mode 100644 test/unit/test_af3_gpu_checks.py diff --git a/test/cluster/af3_gpu_checks.py b/test/cluster/af3_gpu_checks.py new file mode 100644 index 00000000..7be46df9 --- /dev/null +++ b/test/cluster/af3_gpu_checks.py @@ -0,0 +1,133 @@ +"""Checks for the AlphaFold 3 GPU tests that need neither JAX nor a GPU. + +check_alphafold3_predictions.py uses them, and test/unit/test_af3_gpu_checks.py tests +them on a CPU. Nothing here imports JAX: the test process must leave the GPU's memory to +the fold it starts. +""" + +from __future__ import annotations + +import dataclasses +import re +import shutil +import subprocess +from typing import Any, Callable, Mapping, NamedTuple + +# JAX's allocator takes this share of GPU memory unless the environment says otherwise. +_JAX_DEFAULT_MEMORY_FRACTION = 0.75 + + +@dataclasses.dataclass(frozen=True) +class Gpu: + """What nvidia-smi says about the visible GPU.""" + + compute_capability: str + memory_gib: float # the card's memory, or its MIG slice's + + +class Support(NamedTuple): + """Whether the fork's fused triangle kernels are enabled on a GPU. + + ``supported`` is None when that cannot be told; ``detail`` says what was compared. + """ + + supported: bool | None + detail: str + + +def parse_nvidia_smi(query: str, listing: str = "") -> Gpu: + """The first GPU in nvidia-smi's output. + + Args: + query: the output of ``nvidia-smi --query-gpu=compute_cap,memory.total, + mig.mode.current --format=csv,noheader``, e.g. ``9.0, 81559 MiB, Disabled``. + listing: the output of ``nvidia-smi -L``. With MIG enabled, memory.total is the + whole card's, so the memory is read from the first MIG profile listed there + (``MIG 2g.20gb``). + + Raises: + ValueError: if the output does not have that form. + """ + lines = query.strip().splitlines() + fields = [field.strip() for field in lines[0].split(",")] if lines else [] + if len(fields) < 2: + raise ValueError(f"unexpected nvidia-smi output: {query!r}") + memory = fields[1].split() + if len(memory) != 2 or memory[1] != "MiB": + raise ValueError(f"unexpected nvidia-smi memory.total: {fields[1]!r}") + memory_gib = float(memory[0]) / 1024 + if len(fields) > 2 and fields[2] == "Enabled": + profile = re.search(r"\bMIG\s+\S*?(\d+)gb\b", listing) + if profile is None: + raise ValueError(f"MIG is enabled, but nvidia-smi -L names no MIG device: {listing!r}") + memory_gib = int(profile.group(1)) * 1e9 / 2**30 + return Gpu(fields[0], memory_gib) + + +def visible_gpu() -> Gpu | None: + """The first GPU that nvidia-smi reports, or None if it reports none we can read.""" + nvidia_smi = shutil.which("nvidia-smi") + if not nvidia_smi: + return None + try: + query, listing = ( + subprocess.run( + [nvidia_smi, *args], capture_output=True, text=True, check=True + ).stdout + for args in ( + [ + "--query-gpu=compute_cap,memory.total,mig.mode.current", + "--format=csv,noheader", + ], + ["-L"], + ) + ) + return parse_nvidia_smi(query, listing) + except (OSError, subprocess.CalledProcessError, ValueError): + return None + + +def jax_memory_fraction(env: Mapping[str, str]) -> float: + """The share of GPU memory JAX's allocator may take in a process with ``env``.""" + for name in ("XLA_CLIENT_MEM_FRACTION", "XLA_PYTHON_CLIENT_MEM_FRACTION"): + if env.get(name): + return float(env[name]) + return _JAX_DEFAULT_MEMORY_FRACTION + + +def installed_device_policy() -> Callable[[str, float], Any] | None: + """The installed fork's ``device_policy``, or None for an AF3 without fused kernels. + + The fork's dispatch module decides without importing JAX, so this import is cheap. + """ + try: + from alphafold3.jax.fused_triangle import dispatch + except ImportError: + return None + return dispatch.device_policy + + +def fused_triangle_support( + gpu: Gpu | None, + memory_fraction: float, + device_policy: Callable[[str, float], Any] | None, +) -> Support: + """Whether ``device_policy`` enables the fused kernels on ``gpu``, judged without JAX. + + The backend asks the same policy with the compute capability and allocator budget + that JAX reports. Here the budget is estimated as ``memory_fraction`` of the memory + nvidia-smi reports, so a test can check the backend's verdict against the hardware. + """ + if device_policy is None: + return Support(False, "the installed alphafold3 has no alphafold3.jax.fused_triangle") + if gpu is None: + return Support(None, "nvidia-smi reports no GPU it can describe") + budget_gib = gpu.memory_gib * memory_fraction + policy = device_policy(gpu.compute_capability, budget_gib) + detail = ( + f"compute capability {gpu.compute_capability}, about {budget_gib:.1f} GiB for JAX " + f"({memory_fraction:g} of {gpu.memory_gib:.1f} GiB)" + ) + if policy.enabled: + return Support(True, detail) + return Support(False, f"{detail}: {policy.reason}") diff --git a/test/cluster/check_alphafold3_predictions.py b/test/cluster/check_alphafold3_predictions.py index 297222f3..6bc84bff 100755 --- a/test/cluster/check_alphafold3_predictions.py +++ b/test/cluster/check_alphafold3_predictions.py @@ -42,6 +42,11 @@ ) from alphapulldown_input_parser import generate_fold_specifications +# The sibling helper module; pytest puts this directory on sys.path, a direct load may not. +if str(Path(__file__).resolve().parent) not in sys.path: + sys.path.insert(0, str(Path(__file__).resolve().parent)) +import af3_gpu_checks # noqa: E402 + # --------------------------------------------------------------------------- # # configuration / environment guards # @@ -4131,13 +4136,34 @@ def _non_finite_values(payload: Any, where: str = "") -> list[str]: class TestAlphaFold3FastKernels(_TestBase): """--fast_kernels for AF3 on a GPU, and one model compile per process. - Select with -k TestAlphaFold3FastKernels. ``on`` skips with the backend's reason on - a GPU the fork's policy does not enable; ``auto`` must then fall back. + Select with -k TestAlphaFold3FastKernels. ``on`` skips, and ``auto`` may fall back, + only when the backend says the kernels are not supported here and nvidia-smi, read + against the fork's device policy, agrees. Kernels that fail on a GPU the policy + enables are a failure, never a skip. """ FOLD = "A0A075B6L2:1" NOT_SUPPORTED = "AF3 fused triangle kernels are not supported here" + def _fused_triangle_support(self) -> af3_gpu_checks.Support: + """The fork's policy for this GPU, from nvidia-smi rather than the code under test.""" + return af3_gpu_checks.fused_triangle_support( + af3_gpu_checks.visible_gpu(), + af3_gpu_checks.jax_memory_fraction(self._make_af3_test_env()), + af3_gpu_checks.installed_device_policy(), + ) + + def _assert_gpu_unsupported(self, backend_reason: str) -> str: + """The independent check's detail, once it agrees the kernels are not enabled here.""" + support = self._fused_triangle_support() + self.assertIs( + support.supported, + False, + f"the backend says {backend_reason!r}, but nvidia-smi and the fork's policy " + f"say otherwise ({support.detail})", + ) + return support.detail + def _fold(self, mode: str) -> subprocess.CompletedProcess: self._require_af3_functional_environment() return subprocess.run( @@ -4205,7 +4231,8 @@ def test_af3_fast_kernels_on_runs_the_fused_layers(self): log = res.stdout + res.stderr if res.returncode != 0 and f"--fast_kernels=on, but {self.NOT_SUPPORTED}" in log: reason = next(line for line in log.splitlines() if self.NOT_SUPPORTED in line) - self.skipTest(f"this GPU cannot run the fused triangles: {reason.strip()}") + detail = self._assert_gpu_unsupported(reason.strip()) + self.skipTest(f"this GPU cannot run the fused triangles: {reason.strip()} ({detail})") self._runCommonTests(res) self._assert_fused(self._kernel_records("on")) @@ -4223,6 +4250,7 @@ def test_af3_fast_kernels_auto_uses_them_where_the_gpu_can(self): record["reason"].startswith(f"--fast_kernels=auto: {self.NOT_SUPPORTED}"), (key, record), ) + self._assert_gpu_unsupported(next(iter(records.values()))["reason"]) def test_af3_identical_folds_compile_the_model_once(self): """Three identical folds in one process: one jit(apply_fn) compile (1afe779e).""" diff --git a/test/unit/test_af3_gpu_checks.py b/test/unit/test_af3_gpu_checks.py new file mode 100644 index 00000000..d0097c2d --- /dev/null +++ b/test/unit/test_af3_gpu_checks.py @@ -0,0 +1,135 @@ +"""The AF3 GPU tests' own checks, which must judge a GPU and its outputs without JAX.""" + +from __future__ import annotations + +import importlib.util +import subprocess +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[2] +DISPATCH_PATH = ROOT / "alphafold3/src/alphafold3/jax/fused_triangle/dispatch.py" + + +def _load(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module # dataclasses look their module up while the class is created + spec.loader.exec_module(module) + return module + + +checks = _load("af3_gpu_checks", ROOT / "test/cluster/af3_gpu_checks.py") +dispatch = _load("af3_dispatch_gpu_checks", DISPATCH_PATH) + + +# -------------------------------------------------------------------------------------- +# Is the GPU one the fork's fused triangle kernels are enabled on? + + +def test_nvidia_smi_output_gives_the_first_gpu(): + gpu = checks.parse_nvidia_smi( + "9.0, 81559 MiB, Disabled\n9.0, 81559 MiB, Disabled\n", "GPU 0: NVIDIA H100" + ) + + assert gpu.compute_capability == "9.0" + assert gpu.memory_gib == pytest.approx(79.65, abs=0.01) + + +def test_a_mig_slice_has_the_memory_of_its_profile(): + listing = ( + "GPU 0: NVIDIA RTX PRO 4500 Blackwell (UUID: GPU-1)\n" + " MIG 1g.8gb Device 0: (UUID: MIG-2)\n" + ) + + gpu = checks.parse_nvidia_smi("12.0, 32623 MiB, Enabled", listing) + + assert gpu.compute_capability == "12.0" + assert gpu.memory_gib == pytest.approx(8e9 / 2**30) + + +@pytest.mark.parametrize( + "query, listing", + [ + ("", ""), + ("9.0", ""), + ("9.0, [N/A], Disabled", ""), + ("12.0, 32623 MiB, Enabled", "GPU 0: NVIDIA RTX PRO 4500 Blackwell"), + ], +) +def test_unreadable_nvidia_smi_output_is_an_error(query, listing): + with pytest.raises(ValueError): + checks.parse_nvidia_smi(query, listing) + + +def test_the_jax_memory_fraction_comes_from_the_fold_environment(): + assert checks.jax_memory_fraction({"XLA_CLIENT_MEM_FRACTION": "0.95"}) == 0.95 + assert checks.jax_memory_fraction({"XLA_PYTHON_CLIENT_MEM_FRACTION": "0.5"}) == 0.5 + assert checks.jax_memory_fraction({}) == 0.75 + + +@pytest.mark.parametrize( + "compute_capability, memory_gib", + [ + ("8.0", 39.5), # A100 40 GB + ("8.6", 23.6), # RTX 3090, A40 below 48 + ("8.9", 44.4), # L40S + ("9.0", 79.6), # H100 + ("12.0", 95.0), # RTX Pro 6000 Blackwell + ], +) +def test_the_measured_gpus_are_supported(compute_capability, memory_gib): + gpu = checks.Gpu(compute_capability, memory_gib) + + support = checks.fused_triangle_support(gpu, 0.95, dispatch.device_policy) + + assert support.supported is True, support.detail + + +@pytest.mark.parametrize( + "gpu, fraction, reason", + [ + (("10.0", 178.0), 0.95, "unvalidated_compute_capability"), # B200 + (("7.0", 31.7), 0.95, "unvalidated_compute_capability"), # V100 + (("12.0", 8e9 / 2**30), 0.95, "memory_budget_below_12_gib"), # MIG 1g.8gb + (("9.0", 15.0), 0.75, "memory_budget_below_12_gib"), # 11.25 GiB for JAX + ], +) +def test_other_gpus_are_not_supported_and_the_detail_says_why(gpu, fraction, reason): + support = checks.fused_triangle_support(checks.Gpu(*gpu), fraction, dispatch.device_policy) + + assert support.supported is False + assert support.detail.endswith(f": {reason}") + assert f"compute capability {gpu[0]}" in support.detail + + +def test_an_alphafold3_without_the_fork_supports_no_gpu(): + support = checks.fused_triangle_support(checks.Gpu("9.0", 79.6), 0.95, None) + + assert support.supported is False + assert "alphafold3.jax.fused_triangle" in support.detail + + +def test_without_a_described_gpu_support_cannot_be_told(): + support = checks.fused_triangle_support(None, 0.95, dispatch.device_policy) + + assert support.supported is None + + +def test_the_check_reads_the_forks_policy_without_importing_jax(): + # The test process must leave the GPU to the fold it starts. + code = ( + "import sys\n" + f"sys.path[:0] = [{str(ROOT / 'test/cluster')!r}, {str(ROOT / 'alphafold3/src')!r}]\n" + "import af3_gpu_checks\n" + "policy = af3_gpu_checks.installed_device_policy()\n" + "assert policy is not None\n" + "assert policy('9.0', 75.7).enabled\n" + "assert 'jax' not in sys.modules, sorted(m for m in sys.modules if 'jax' in m)\n" + ) + + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + + assert result.returncode == 0, result.stderr From d0142dae3a087cbc615f196a3cce8b8e3d3368cc Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 09:32:07 +0200 Subject: [PATCH 12/14] test: AF3 GPU tests reject null scores in confidences files AF3 writes NaN scores as JSON null, which the non-finite check let through. Each summary_confidences.json now needs ptm, ranking_score, fraction_disordered, has_clash, chain_ptm, chain_pair_iptm and chain_pair_pae_min present, numeric and finite (arrays element by element), and iptm and chain_iptm too unless the prediction has one chain, where AF3 defines ipTM as NaN. Each confidences.json needs finite atom_plddts, pae and contact_probs. Co-Authored-By: Claude Opus 5.5 --- test/cluster/af3_gpu_checks.py | 92 ++++++++++++++ test/cluster/check_alphafold3_predictions.py | 21 +--- test/unit/test_af3_gpu_checks.py | 123 ++++++++++++++++++- 3 files changed, 217 insertions(+), 19 deletions(-) diff --git a/test/cluster/af3_gpu_checks.py b/test/cluster/af3_gpu_checks.py index 7be46df9..9451f2b1 100644 --- a/test/cluster/af3_gpu_checks.py +++ b/test/cluster/af3_gpu_checks.py @@ -8,6 +8,7 @@ from __future__ import annotations import dataclasses +import math import re import shutil import subprocess @@ -131,3 +132,94 @@ def fused_triangle_support( if policy.enabled: return Support(True, detail) return Support(False, f"{detail}: {policy.reason}") + + +# -------------------------------------------------------------------------------------- +# Are the scores AF3 wrote real numbers? + +# summary_confidences.json, as confidence_types.StructureConfidenceSummary writes it: +# field -> whether it is an array (chain_pair_* are [num_chains, num_chains]). +_SUMMARY_FIELDS = { + "ptm": False, + "iptm": False, + "ranking_score": False, + "fraction_disordered": False, + "has_clash": False, + "chain_ptm": True, + "chain_iptm": True, + "chain_pair_iptm": True, + "chain_pair_pae_min": True, +} +# AF3 defines ipTM only between chains: with one chain, iptm and the cross-chain +# chain_iptm are NaN (model.py, confidences.get_ranking_score). Nothing else may be. +_SINGLE_CHAIN_NULLABLE = frozenset({"iptm", "chain_iptm"}) +# AF3 writes has_clash as 0.0 or 1.0; a boolean means the same. +_BOOLEAN_FIELDS = frozenset({"has_clash"}) +# confidences.json (StructureConfidenceFull): per-atom pLDDT, per-token-pair PAE and +# contact probabilities. +_FULL_ARRAYS = ("atom_plddts", "pae", "contact_probs") +# Problems listed per field, so that an all-NaN PAE does not flood the report. +_MAX_PROBLEMS_PER_FIELD = 5 + + +def confidence_problems(name: str, payload: Any) -> list[str]: + """What is wrong with the scores in one AF3 confidences file; empty when nothing. + + ``name`` (the file name) picks the format: ``*summary_confidences.json`` or + ``*confidences.json``. Every score must be present, a number and finite, and arrays + element by element. AF3 writes NaN as JSON null, so null counts as NaN, except + where AF3 defines the score as NaN: ipTM of a single-chain prediction. + """ + if name.endswith("summary_confidences.json"): + fields = _SUMMARY_FIELDS + elif name.endswith("confidences.json"): + fields = dict.fromkeys(_FULL_ARRAYS, True) + else: + raise ValueError(f"not an AF3 confidences file: {name}") + if not isinstance(payload, dict): + return [f"not a JSON object: {type(payload).__name__}"] + chain_ptm = payload.get("chain_ptm") + single_chain = isinstance(chain_ptm, list) and len(chain_ptm) == 1 + problems = [] + for field, is_array in fields.items(): + if field not in payload: + problems.append(f"{field} is missing") + continue + value = payload[field] + if is_array and not (isinstance(value, list) and value): + problems.append(f"{field} is not a non-empty array: {value!r:.60}") + continue + if not is_array and isinstance(value, list): + problems.append(f"{field} is an array, not a number") + continue + found = _number_problems( + field, + value, + nullable=single_chain and field in _SINGLE_CHAIN_NULLABLE, + boolean=field in _BOOLEAN_FIELDS, + ) + if len(found) > _MAX_PROBLEMS_PER_FIELD: + more = len(found) - _MAX_PROBLEMS_PER_FIELD + found = [*found[:_MAX_PROBLEMS_PER_FIELD], f"{field}: {more} more"] + problems.extend(found) + return problems + + +def _number_problems(where: str, value: Any, *, nullable: bool, boolean: bool) -> list[str]: + if isinstance(value, list): + return [ + problem + for index, element in enumerate(value) + for problem in _number_problems( + f"{where}[{index}]", element, nullable=nullable, boolean=boolean + ) + ] + if value is None: + return [] if nullable else [f"{where} is null (NaN)"] + if isinstance(value, bool): + return [] if boolean else [f"{where} is a boolean, not a number"] + if not isinstance(value, (int, float)): + return [f"{where} is not a number: {value!r:.60}"] + if not math.isfinite(value): + return [f"{where} is {value}"] + return [] diff --git a/test/cluster/check_alphafold3_predictions.py b/test/cluster/check_alphafold3_predictions.py index 6bc84bff..57f9e39b 100755 --- a/test/cluster/check_alphafold3_predictions.py +++ b/test/cluster/check_alphafold3_predictions.py @@ -4117,22 +4117,6 @@ def test_af3_rejects_alphafold2_flag(self): ) -def _non_finite_values(payload: Any, where: str = "") -> list[str]: - """Paths of NaN or infinite numbers in a parsed JSON payload.""" - if isinstance(payload, dict): - items = ((f"{where}.{key}", value) for key, value in payload.items()) - elif isinstance(payload, list): - items = ((f"{where}[{index}]", value) for index, value in enumerate(payload)) - else: - items = () - nested = [bad for path, value in items for bad in _non_finite_values(value, path)] - if nested: - return nested - if isinstance(payload, float) and not np.isfinite(payload): - return [where] - return [] - - class TestAlphaFold3FastKernels(_TestBase): """--fast_kernels for AF3 on a GPU, and one model compile per process. @@ -4186,7 +4170,7 @@ def _fold(self, mode: str) -> subprocess.CompletedProcess: ) def _kernel_records(self, mode: str) -> dict[str, dict[str, Any]]: - """inference_kernels.json of a finished fold: one record per seed, outputs finite.""" + """inference_kernels.json of a finished fold: one record per seed, scores finite.""" result_dir = self._resolve_single_af3_result_dir() with (result_dir / "ranking_scores.csv").open() as handle: seeds = {line.split(",")[0] for line in handle.readlines()[1:] if line.strip()} @@ -4200,7 +4184,8 @@ def _kernel_records(self, mode: str) -> dict[str, dict[str, Any]]: confidences = sorted(result_dir.rglob("*confidences.json")) self.assertTrue(confidences, f"no confidences in {result_dir}") for path in confidences: - self.assertEqual(_non_finite_values(json.loads(path.read_text())), [], path) + problems = af3_gpu_checks.confidence_problems(path.name, json.loads(path.read_text())) + self.assertEqual(problems, [], path) return records def _assert_fused(self, records: dict[str, dict[str, Any]]) -> None: diff --git a/test/unit/test_af3_gpu_checks.py b/test/unit/test_af3_gpu_checks.py index d0097c2d..544014d7 100644 --- a/test/unit/test_af3_gpu_checks.py +++ b/test/unit/test_af3_gpu_checks.py @@ -1,8 +1,9 @@ -"""The AF3 GPU tests' own checks, which must judge a GPU and its outputs without JAX.""" +"""The AF3 GPU tests' own checks, which judge a GPU and its outputs without JAX.""" from __future__ import annotations import importlib.util +import json import subprocess import sys from pathlib import Path @@ -133,3 +134,123 @@ def test_the_check_reads_the_forks_policy_without_importing_jax(): result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) assert result.returncode == 0, result.stderr + + +# -------------------------------------------------------------------------------------- +# Are the scores AF3 wrote real numbers? AF3 writes NaN as null. + +MONOMER_SUMMARY = """{ + "chain_iptm": [null], "chain_pair_iptm": [[0.84]], "chain_pair_pae_min": [[0.76]], + "chain_ptm": [0.84], "fraction_disordered": 0.0, "has_clash": 0.0, "iptm": null, + "ptm": 0.84, "ranking_score": 0.84, "chain_ids": ["A"] +}""" +DIMER_SUMMARY = """{ + "chain_iptm": [0.61, 0.61], "chain_pair_iptm": [[0.8, 0.61], [0.61, 0.79]], + "chain_pair_pae_min": [[0.8, 4.2], [4.3, 0.8]], "chain_ptm": [0.8, 0.79], + "fraction_disordered": 0.05, "has_clash": 0.0, "iptm": 0.61, "ptm": 0.7, + "ranking_score": 0.66 +}""" +FULL = """{ + "atom_chain_ids": ["A", "A"], "atom_plddts": [91.2, 88.4], + "contact_probs": [[1.0, 0.9], [0.9, 1.0]], "pae": [[0.8, 1.4], [1.3, 0.8]], + "token_chain_ids": ["A", "A"], "token_res_ids": [1, 2] +}""" + + +def _summary(text: str, **changes) -> dict: + payload = json.loads(text) + payload.update(changes) + return payload + + +def test_a_monomer_may_have_no_iptm(): + assert checks.confidence_problems("summary_confidences.json", _summary(MONOMER_SUMMARY)) == [] + + +def test_complete_dimer_and_full_confidences_pass(): + assert checks.confidence_problems("x_summary_confidences.json", _summary(DIMER_SUMMARY)) == [] + assert checks.confidence_problems("confidences.json", json.loads(FULL)) == [] + + +@pytest.mark.parametrize( + "changes, problem", + [ + ({"ptm": None}, "ptm is null (NaN)"), + ({"ranking_score": None}, "ranking_score is null (NaN)"), + ({"chain_ptm": [None]}, "chain_ptm[0] is null (NaN)"), + ({"chain_pair_iptm": [[None]]}, "chain_pair_iptm[0][0] is null (NaN)"), + ({"ptm": float("nan")}, "ptm is nan"), + ({"fraction_disordered": float("inf")}, "fraction_disordered is inf"), + ({"ptm": True}, "ptm is a boolean, not a number"), + ({"ptm": "0.84"}, "ptm is not a number: '0.84'"), + ({"ptm": [0.84]}, "ptm is an array, not a number"), + ({"chain_ptm": []}, "chain_ptm is not a non-empty array: []"), + ], +) +def test_missing_or_non_finite_monomer_scores_fail(changes, problem): + problems = checks.confidence_problems( + "summary_confidences.json", _summary(MONOMER_SUMMARY, **changes) + ) + + assert problem in problems, problems + + +def test_nan_written_by_python_json_fails(): + payload = json.loads(MONOMER_SUMMARY.replace('"ptm": 0.84', '"ptm": NaN')) + + assert checks.confidence_problems("summary_confidences.json", payload) == ["ptm is nan"] + + +def test_a_missing_score_fails(): + payload = _summary(MONOMER_SUMMARY) + del payload["ranking_score"] + + assert checks.confidence_problems("summary_confidences.json", payload) == [ + "ranking_score is missing" + ] + + +def test_has_clash_may_be_a_boolean(): + payload = _summary(MONOMER_SUMMARY, has_clash=False) + + assert checks.confidence_problems("summary_confidences.json", payload) == [] + + +@pytest.mark.parametrize( + "changes, problem", + [ + ({"iptm": None}, "iptm is null (NaN)"), + ({"chain_iptm": [0.61, None]}, "chain_iptm[1] is null (NaN)"), + ], +) +def test_a_complex_must_have_an_iptm(changes, problem): + problems = checks.confidence_problems( + "summary_confidences.json", _summary(DIMER_SUMMARY, **changes) + ) + + assert problems == [problem] + + +@pytest.mark.parametrize( + "field, value, problem", + [ + ("atom_plddts", [91.2, None], "atom_plddts[1] is null (NaN)"), + ("pae", [[0.8, float("nan")], [1.3, 0.8]], "pae[0][1] is nan"), + ("contact_probs", None, "contact_probs is not a non-empty array: None"), + ], +) +def test_full_confidences_need_finite_plddt_pae_and_contacts(field, value, problem): + payload = json.loads(FULL) + payload[field] = value + + assert checks.confidence_problems("seed-1_sample-0/confidences.json", payload) == [problem] + + +def test_an_all_nan_array_is_reported_briefly(): + payload = json.loads(FULL) + payload["pae"] = [[None] * 64 for _ in range(64)] + + problems = checks.confidence_problems("confidences.json", payload) + + assert problems[:5] == [f"pae[0][{index}] is null (NaN)" for index in range(5)] + assert problems[5:] == [f"pae: {64 * 64 - 5} more"] From 22045684aba95b4e6fbe14d423c215a74cb39ab0 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 09:36:53 +0200 Subject: [PATCH 13/14] test: AF3 fused triangle kernels score like the original layers Like AF2's test_dimer_fast_kernels_match_the_stock_code: on a GPU the fork's policy enables (per nvidia-smi), fold a homodimer with the same seeds with --fast_kernels off and on and require, per seed, ranking score, pTM and ipTM within 0.02 and mean pLDDT within 1, with the on run's records showing the fused layers. Elsewhere the test skips with the reason. Co-Authored-By: Claude Opus 5.5 --- test/cluster/check_alphafold3_predictions.py | 73 +++++++++++++++++++- 1 file changed, 71 insertions(+), 2 deletions(-) diff --git a/test/cluster/check_alphafold3_predictions.py b/test/cluster/check_alphafold3_predictions.py index 57f9e39b..10bb1a21 100755 --- a/test/cluster/check_alphafold3_predictions.py +++ b/test/cluster/check_alphafold3_predictions.py @@ -4127,7 +4127,14 @@ class TestAlphaFold3FastKernels(_TestBase): """ FOLD = "A0A075B6L2:1" + # A homodimer, so that the comparison also covers ipTM. + DIMER = "A0A075B6L2:2" NOT_SUPPORTED = "AF3 fused triangle kernels are not supported here" + RANDOM_SEED = 42 + # Fused bf16 kernels change rounding, not the model: per seed, ranking score, pTM + # and ipTM within 0.02 of the original layers, mean pLDDT (0-100) within 1. + SCORE_TOLERANCE = 0.02 + PLDDT_TOLERANCE = 1.0 def _fused_triangle_support(self) -> af3_gpu_checks.Support: """The fork's policy for this GPU, from nvidia-smi rather than the code under test.""" @@ -4148,13 +4155,13 @@ def _assert_gpu_unsupported(self, backend_reason: str) -> str: ) return support.detail - def _fold(self, mode: str) -> subprocess.CompletedProcess: + def _fold(self, mode: str, *flags: str, fold: str = FOLD) -> subprocess.CompletedProcess: self._require_af3_functional_environment() return subprocess.run( [ sys.executable, str(self.script_single), - f"--input={self.FOLD}", + f"--input={fold}", f"--output_directory={self.output_dir}", f"--data_directory={DATA_DIR}", f"--features_directory={self.test_features_dir}", @@ -4163,6 +4170,7 @@ def _fold(self, mode: str) -> subprocess.CompletedProcess: "--num_diffusion_samples=1", "--num_seeds=2", f"--fast_kernels={mode}", + *flags, ], capture_output=True, text=True, @@ -4188,6 +4196,25 @@ def _kernel_records(self, mode: str) -> dict[str, dict[str, Any]]: self.assertEqual(problems, [], path) return records + def _sample_scores(self) -> dict[str, dict[str, float | None]]: + """Ranking score, pTM, ipTM (None if undefined) and mean pLDDT of each sample.""" + scores = {} + for sample_dir in sorted(self._resolve_single_af3_result_dir().glob("seed-*_sample-*")): + (summary_path,) = sample_dir.glob("*summary_confidences.json") + (full_path,) = [ + path + for path in sample_dir.glob("*confidences.json") + if not path.name.endswith("summary_confidences.json") + ] + summary = json.loads(summary_path.read_text()) + scores[sample_dir.name] = { + "ranking_score": summary["ranking_score"], + "ptm": summary["ptm"], + "iptm": summary["iptm"], + "mean_plddt": float(np.mean(json.loads(full_path.read_text())["atom_plddts"])), + } + return scores + def _assert_fused(self, records: dict[str, dict[str, Any]]) -> None: for key, record in records.items(): self.assertTrue(record["fused_kernels"], (key, record)) @@ -4237,6 +4264,48 @@ def test_af3_fast_kernels_auto_uses_them_where_the_gpu_can(self): ) self._assert_gpu_unsupported(next(iter(records.values()))["reason"]) + def test_af3_fast_kernels_match_the_original_layers(self): + """The same dimer and seeds with --fast_kernels off and on: the scores agree.""" + support = self._fused_triangle_support() + if not support.supported: + self.skipTest( + f"the fork's policy does not enable the fused triangles here ({support.detail})" + ) + + root = self.output_dir + runs = {} + for mode in ("off", "on"): + # _fold and the helpers that read its outputs work in self.output_dir. + self.output_dir = root / mode + self.output_dir.mkdir(parents=True, exist_ok=True) + self._runCommonTests( + self._fold(mode, f"--random_seed={self.RANDOM_SEED}", fold=self.DIMER) + ) + runs[mode] = self._kernel_records(mode), self._sample_scores() + + (stock_records, stock), (fused_records, fused) = runs["off"], runs["on"] + self.assertFalse( + any(record["fused_kernels"] for record in stock_records.values()), stock_records + ) + self._assert_fused(fused_records) + print(f"original layers: {stock}\nfused kernels: {fused}") + self.assertEqual(set(fused), set(stock)) + self.assertLen(stock, 2) # two seeds, one sample each + for sample, expected in stock.items(): + got = fused[sample] + for key in ("ranking_score", "ptm", "iptm"): + context = (sample, key, expected, got) + self.assertEqual(got[key] is None, expected[key] is None, context) + if expected[key] is not None: + self.assertLessEqual( + abs(got[key] - expected[key]), self.SCORE_TOLERANCE, context + ) + self.assertLessEqual( + abs(got["mean_plddt"] - expected["mean_plddt"]), + self.PLDDT_TOLERANCE, + (sample, expected, got), + ) + def test_af3_identical_folds_compile_the_model_once(self): """Three identical folds in one process: one jit(apply_fn) compile (1afe779e).""" self._require_af3_functional_environment() From ab4a8ea0f0fa01515ab6180a5261eb1871480523 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 09:49:17 +0200 Subject: [PATCH 14/14] chore: alphafold3 submodule at the merged fork (KosinskiLab/alphafold3#5) Summary confidences write one chain ID per chain (v3.0.4 wrote one per token). Co-Authored-By: Claude Opus 5.5 --- alphafold3 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/alphafold3 b/alphafold3 index 7809148b..98785987 160000 --- a/alphafold3 +++ b/alphafold3 @@ -1 +1 @@ -Subproject commit 7809148bfefdb42daa3bf56c884e345b2b7a30c6 +Subproject commit 98785987139da9781ba271b159c93efaaf23a06f