diff --git a/README.md b/README.md index 6eee826..a2c81a9 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 from + 2.11.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..3f7780e 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 from + # 2.11.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..ad178b9 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, from AlphaPulldown 2.11.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 fa22288..e1aac59 100644 --- a/test/test_inference_flag_tables.py +++ b/test/test_inference_flag_tables.py @@ -24,8 +24,8 @@ _COMMON = SourceFileLoader("common_flag_tables", str(COMMON_SMK)).load_module() -# The AlphaPulldown the shipped default containers run (2.10.0). -SHIPPED_ALPHAPULLDOWN_COMMIT = "6e7f98695f576cf7ec087dbaef39b1f39652054e" +# The AlphaPulldown the tables come from: main after 2.10.0, where AF3 accepts --allow_resume. +SHIPPED_ALPHAPULLDOWN_COMMIT = "228c5eed6c7066467e46b87fb3054cc2c2e73ae1" # Shaped like AlphaPulldown's inference_flags.py: named sets, unions, an annotated dict. _SOURCE = ''' @@ -197,6 +197,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 9d5a5e2..1da795f 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 + from AlphaPulldown 2.11.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 @@ -913,7 +914,7 @@ def prediction_batch_id(folds: Iterable[str]) -> str: # to update for a new release; CI regenerates from the recorded commit and fails on any # difference. # BEGIN GENERATED: inference flags (tools/generate_inference_flag_tables.py; do not edit by hand) -# Source: KosinskiLab/AlphaPulldown@6e7f98695f576cf7ec087dbaef39b1f39652054e +# Source: KosinskiLab/AlphaPulldown@228c5eed6c7066467e46b87fb3054cc2c2e73ae1 # alphapulldown/prediction/inference_flags.py, FLAGS_BY_BACKEND ALLOWED_INFERENCE_FLAGS = { "alphafold2": frozenset({ @@ -957,6 +958,7 @@ ALLOWED_INFERENCE_FLAGS = { "use_gpu_relax", }), "alphafold3": frozenset({ + "allow_resume", "buckets", "convert_to_modelcif", "data_directory",