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/alphafold3 b/alphafold3 index 7809148b..98785987 160000 --- a/alphafold3 +++ b/alphafold3 @@ -1 +1 @@ -Subproject commit 7809148bfefdb42daa3bf56c884e345b2b7a30c6 +Subproject commit 98785987139da9781ba271b159c93efaaf23a06f 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/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/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/af3_gpu_checks.py b/test/cluster/af3_gpu_checks.py new file mode 100644 index 00000000..9451f2b1 --- /dev/null +++ b/test/cluster/af3_gpu_checks.py @@ -0,0 +1,225 @@ +"""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 math +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}") + + +# -------------------------------------------------------------------------------------- +# 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_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.""" diff --git a/test/cluster/check_alphafold3_predictions.py b/test/cluster/check_alphafold3_predictions.py index 2231a758..10bb1a21 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 # @@ -4112,6 +4117,250 @@ def test_af3_rejects_alphafold2_flag(self): ) +class TestAlphaFold3FastKernels(_TestBase): + """--fast_kernels for AF3 on a GPU, and one model compile per process. + + 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" + # 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.""" + 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, *flags: str, fold: str = FOLD) -> subprocess.CompletedProcess: + self._require_af3_functional_environment() + return subprocess.run( + [ + sys.executable, + str(self.script_single), + f"--input={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}", + *flags, + ], + 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, 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()} + 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: + problems = af3_gpu_checks.confidence_problems(path.name, json.loads(path.read_text())) + 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)) + 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.NOT_SUPPORTED}" in log: + reason = next(line for line in log.splitlines() if self.NOT_SUPPORTED in line) + 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")) + + 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.NOT_SUPPORTED}"), + (key, record), + ) + 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() + 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.""" 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/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") diff --git a/test/unit/test_af3_fused_triangles.py b/test/unit/test_af3_fused_triangles.py index c424f0a2..2a5bb925 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", @@ -86,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): @@ -284,3 +338,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]) diff --git a/test/unit/test_af3_gpu_checks.py b/test/unit/test_af3_gpu_checks.py new file mode 100644 index 00000000..544014d7 --- /dev/null +++ b/test/unit/test_af3_gpu_checks.py @@ -0,0 +1,256 @@ +"""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 + +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 + + +# -------------------------------------------------------------------------------------- +# 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"] 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 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, 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 ): 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" diff --git a/test/unit/test_mmseqs2_gpu_dockerfiles.py b/test/unit/test_mmseqs2_gpu_dockerfiles.py index a3564b8e..043efd20 100644 --- a/test/unit/test_mmseqs2_gpu_dockerfiles.py +++ b/test/unit/test_mmseqs2_gpu_dockerfiles.py @@ -101,3 +101,64 @@ 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_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) + 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"] + + +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:]