From 7c3ccd9394c1c3abc0611ebac6cb9fb84d41cbb4 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 13 Oct 2026 14:48:11 +0200 Subject: [PATCH] feat: accept --allow_resume for AlphaFold 3 AlphaPulldown now resumes AlphaFold 3 folds as well (KosinskiLab/AlphaPulldown #655-#657): a rerun skips finished folds, and keeps the seeds an interrupted fold already wrote out. The flag defaults to True there, so batch reruns resume without the workflow passing it. It is still added for AlphaFold 2 batches only: AlphaFold 3 images up to 2.10.0 reject it. The AlphaFold 3 flag table accepts it, so --allow_resume: false (predict everything again) passes the parse-time check; config.yaml shows it. Docstrings, comments and two README statements that called the flag AlphaFold2-only are corrected in place, with no added README lines. Co-Authored-By: Claude Opus 5.5 --- README.md | 6 +++--- config/config.yaml | 3 +++ test/test_batch_inference_args.py | 6 ++++-- test/test_inference_command_line.py | 3 ++- test/test_inference_flag_tables.py | 8 +++++++- test/test_unknown_inference_flags.py | 12 ++++++------ workflow/Snakefile | 2 +- workflow/rules/common.smk | 7 +++++-- 8 files changed, 31 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 6eee826..fcbb14c 100644 --- a/README.md +++ b/README.md @@ -693,8 +693,8 @@ batch_max_tokens: 0 # optional cap on summed residues per batch (0 = no cap) share a shape, whereas for AlphaFold3 the point is queueing once instead of once per fold. - For AlphaFold2 batches, `--allow_resume` is enabled automatically, so if a job is - interrupted a rerun skips folds whose outputs already exist (AlphaFold3 does not accept - that flag, so its batches recompute the unfinished folds on rerun). + interrupted a rerun skips folds whose outputs already exist (AlphaFold3 images after + 2.10.0 resume by default as well, keeping finished folds and seeds). - Analysis and reports are unaffected — `alphajudge` still runs per fold (one `interfaces.csv` + `report.pdf` each) and the recursive summary still aggregates them. - **Trade-off:** a batch is one SLURM job, so a failure reruns the whole batch (minus the @@ -861,7 +861,7 @@ You can pass backend CLI switches through `structure_inference_arguments`. Commo > against the selected `--fold_backend` and aborts the job with > `ValueError: The following flags are not supported by backend ''` if you pass one the > backend does not accept. Only use flags from **your** backend's list below — e.g. -> `--allow_resume` is AlphaFold2-only. A single wrong flag fails the job immediately +> `--models_to_relax` is AlphaFold2-only. A single wrong flag fails the job immediately > (before any prediction runs). > > When **batching** (`batch_size > 1`) the workflow adds what each backend needs — diff --git a/config/config.yaml b/config/config.yaml index 3f04f13..2365a50 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -52,6 +52,9 @@ structure_inference_arguments: # 1.1-2x AlphaFold 3. Keep the quotes: YAML reads bare on/off as booleans. # README: "AlphaFold2 flags", "AlphaFold3 flags" # --fast_kernels: "auto" + # A rerun skips finished folds and resumes interrupted ones (AlphaFold3: images after + # 2.10.0). Set false to predict everything again. + # --allow_resume: false # SLURM only: replace "gpu" with your cluster's GPU partition. Several names can # be comma-separated. Desktop runs ignore these scheduler settings. diff --git a/test/test_batch_inference_args.py b/test/test_batch_inference_args.py index e489d59..44b7b7c 100644 --- a/test/test_batch_inference_args.py +++ b/test/test_batch_inference_args.py @@ -2,7 +2,8 @@ Regression guard: ``run_structure_prediction.py`` hard-errors on a flag its backend does not accept (``ValueError: not supported by backend ''``). Adding the wrong flag once -broke ALL batched AlphaFold2 inference. ``--allow_resume`` is AlphaFold2-only; +broke ALL batched AlphaFold2 inference. ``--allow_resume`` is added for AlphaFold2 only +(AlphaFold3 resumes by default where it accepts the flag, after AlphaPulldown 2.10.0); ``--jax_compilation_cache_dir`` is accepted by both JAX backends (AlphaFold2 since AlphaPulldown 2.8.0). These tests pin each flag to the backends that accept it. @@ -44,7 +45,8 @@ def test_af3_batch_gets_jax_cache_not_allow_resume(): jax_cache_dir=_CACHE, ) assert args["--jax_compilation_cache_dir"] == _CACHE - # the AF2-only flag must NOT be present for AF3 (it would ValueError at runtime) + # Not added for AF3: it resumes by default, and images up to AlphaPulldown 2.10.0 + # reject the flag (ValueError at runtime). assert "--allow_resume" not in args diff --git a/test/test_inference_command_line.py b/test/test_inference_command_line.py index 6626a11..da091e2 100644 --- a/test/test_inference_command_line.py +++ b/test/test_inference_command_line.py @@ -145,7 +145,8 @@ def test_fast_kernels_and_cache_reach_every_inference_call(tmp_path, backend): assert flags["fast_kernels"] == "True", script assert flags["jax_compilation_cache_dir"] == str(cache), script assert flags["fold_backend"] == backend, script - # --allow_resume is added to AlphaFold 2 batches only; AlphaFold 3 rejects it. + # --allow_resume is added to AlphaFold 2 batches only: AlphaFold 3 resumes by + # default, and its images up to AlphaPulldown 2.10.0 reject the flag. assert ("allow_resume" in flags) == (backend == "alphafold2"), script assert job.index(_XLA_CACHES_EXPORT) < job.index("run_structure_prediction") diff --git a/test/test_inference_flag_tables.py b/test/test_inference_flag_tables.py index 8a0ff27..cd3acfc 100644 --- a/test/test_inference_flag_tables.py +++ b/test/test_inference_flag_tables.py @@ -39,7 +39,7 @@ "jax_compilation_cache_dir", "buckets", "flash_attention_implementation", "num_diffusion_samples", "num_seeds", "debug_templates", "debug_msas", "num_recycles", "save_embeddings", "save_distogram", "use_ap_style", - "convert_to_modelcif", + "convert_to_modelcif", "allow_resume", } @@ -66,6 +66,12 @@ def test_fast_kernels_is_valid_on_both_backends(): assert _COMMON.unknown_inference_flags(args, "alphalink") == ["fast_kernels"] +def test_allow_resume_is_valid_on_both_backends(): + for backend in ("alphafold2", "alphafold3"): + args = {"--fold_backend": backend, "--allow_resume": False} + assert _COMMON.unknown_inference_flags(args, backend) == [] + + def test_jax_compile_cache_is_valid_on_both_backends(): for backend in ("alphafold2", "alphafold3"): args = {"--fold_backend": backend, "--jax_compilation_cache_dir": "/cache"} diff --git a/test/test_unknown_inference_flags.py b/test/test_unknown_inference_flags.py index 435976b..980588a 100644 --- a/test/test_unknown_inference_flags.py +++ b/test/test_unknown_inference_flags.py @@ -26,8 +26,8 @@ def test_af2_flags_all_accepted(): def test_af2_only_flag_flagged_on_af3(): - args = {"--fold_backend": "alphafold3", "--allow_resume": "true"} - assert unknown_inference_flags(args, "alphafold3") == ["allow_resume"] + args = {"--fold_backend": "alphafold3", "--models_to_relax": "best"} + assert unknown_inference_flags(args, "alphafold3") == ["models_to_relax"] def test_af3_flags_all_accepted(): @@ -45,7 +45,7 @@ def test_alphalink_accepts_crosslinks_and_af2_flags(): def test_leading_dashes_and_equals_are_stripped(): # keys may arrive with or without dashes, or with an inline value - assert unknown_inference_flags({"allow_resume": "true"}, "alphafold3") == ["allow_resume"] + assert unknown_inference_flags({"num_cycle": 3}, "alphafold3") == ["num_cycle"] assert unknown_inference_flags({"--bogus=1": "x"}, "alphafold3") == ["bogus"] @@ -55,14 +55,14 @@ def test_unknown_backend_stays_silent(): def test_backend_name_case_insensitive(): - assert unknown_inference_flags({"--allow_resume": "t"}, "AlphaFold3") == ["allow_resume"] + assert unknown_inference_flags({"--num_cycle": 3}, "AlphaFold3") == ["num_cycle"] def test_order_preserved_and_deduped(): - args = {"--allow_resume": "t", "--models_to_relax": "best", "--allow_resume2": "t"} + args = {"--num_cycle": 3, "--models_to_relax": "best", "--num_cycle2": "t"} # only genuinely-unknown AF3 flags, in order out = unknown_inference_flags(args, "alphafold3") - assert out == ["allow_resume", "models_to_relax", "allow_resume2"] + assert out == ["num_cycle", "models_to_relax", "num_cycle2"] if __name__ == "__main__": diff --git a/workflow/Snakefile b/workflow/Snakefile index 57fc20f..0a74a3b 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -353,7 +353,7 @@ _structure_inference_args = batch_inference_args( jax_cache_dir=JAX_CACHE_DIR, ) # Parse-time guard: run_structure_prediction.py aborts the inference job on the first -# flag its backend does not accept (e.g. --allow_resume on AF3). Warn here -- on the +# flag its backend does not accept (e.g. --models_to_relax on AF3). Warn here -- on the # head node, in seconds -- instead of failing deep in a Slurm job. Uses the RAW # --fold_backend (so alphalink's --crosslinks is not mis-flagged); a warning only, # since the container is the source of truth (see unknown_inference_flags). diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index ba13a2c..0bcc975 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -818,8 +818,9 @@ def batch_inference_args( any it does not recognise (``ValueError: not supported by backend ''``), so a flag must ONLY be added for the backend(s) that accept it: - * ``--allow_resume`` (AlphaFold2 only, ``batch_size > 1``): a crashed batch re-runs - all its folds, so resume the ones already done. AlphaFold3 rejects it. + * ``--allow_resume`` (AlphaFold2, ``batch_size > 1``): a crashed batch re-runs all + its folds, so resume the ones already done. Not added for AlphaFold3: its images + after AlphaPulldown 2.10.0 resume by default, and 2.10.0 and older reject the flag. * ``--jax_compilation_cache_dir`` (AlphaFold2 and AlphaFold3, every batch size): one on-disk JAX compile cache shared by every inference process. Without it every process compiles its models from scratch (AlphaFold3: once per token bucket, and @@ -937,6 +938,8 @@ _AF3_INFERENCE_FLAGS = { "num_diffusion_samples", "num_seeds", "debug_templates", "debug_msas", "num_recycles", "save_embeddings", "save_distogram", "use_ap_style", "convert_to_modelcif", + # AlphaPulldown releases after 2.10.0; 2.10.0 and older reject it for AlphaFold3. + "allow_resume", } _ALPHALINK_EXTRA_FLAGS = {"crosslinks"} # fast_kernels (off|on|auto) is accepted for alphafold2 and alphafold3, not AlphaLink: