Repository navigation
fix + test: AF3 kernel self-check, GPU test runners, chain_ids; cover the changes since 2.9.1 - #647
Merged
Merged
Conversation
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 <noreply@anthropic.com>
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 (1afe779). 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
…d to end 4519a56'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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
…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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 1afe779 regression; matches the manual retrace_fixed.sbatch check). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
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 <noreply@anthropic.com>
…elf-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 <noreply@anthropic.com>
…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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
DimaMolod
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Two fixes and tests. Fixes, in separate commits:
--fast_kernels:resolve()now tells an unsupported GPU ("not supported here") from kernels that fail their self-check ("failed their self-check"). Withauto, a failed self-check still falls back to the stock layers, but is logged as an ERROR, so a kernel bug can't hide as a quiet fallback.chain_idsone entry per chain (v3.0.4 wrote one per token).Tests cover everything new since 2.9.1 that wasn't covered yet: the compile cache, AF2
--fast_kernels, the padded AF2 scores, AF3--fast_kernelsand the AF3 one-compile fix.alphafold/alphafold/model/fused_kernels_test.py), on Python 3.11, which the fork needs.pytest.iniskipsalphafold/, so they never ran before.test_af3_fused_triangles.py, the wholetest_alphafold3_backend_helpers.py, the new tokamax test and the fork'sfused_triangle_test.py.fused_kernels.fused_ops()imports.test_mmseqs2_gpu_dockerfiles.pykeeps these steps in place.test_af3_tokamax_trace_context.py(real tokamax): the hook AlphaPulldown uses exists and is the path a kernel's cache lookup takes. Three calls compile twice without the hook and once with it.test_alphafold2_backend_helpers.py::test_padded_prediction_reports_scores_and_pae_over_the_real_residues: runs predict and postprocess with the real padding and confidence code. Ranking, PAE/pLDDT JSON, plots and pickles cover only the real residues.test/integration/test_jax_compilation_cache_real.py: with real jax, a second process hits the cache and writes nothing new; no XLA cache paths are set unless asked for.test_af3_fused_triangles.py: every setting AlphaPulldown passes is aGlobalConfigfield the fork declares, with an allowed value (parsed withast).test/cluster, run withRUN_GPU_FUNCTIONAL_TESTS=1; they skip otherwise):TestFastKernelsAndPadding(AF2): off vs auto on a dimer (kernel record, ipTM/pTM within 0.02); a monomer keeps the stock code; a padded batch matches the unpadded fold.TestAlphaFold3FastKernels(AF3):seed-Nrecords, and every required score in the confidences files is a finite number.null, which is how AF3 writes NaN, fails; only a monomer's ipTM may be null.onskips, andautoaccepts a fallback, only when the backend says "not supported here" and an independentnvidia-smicheck against the fork's device policy agrees.jit(apply_fn)once.Validation
fused_ops()check passes with colabfold-kernels 0.4.0.main.🤖 Generated with Claude Code