diff --git a/config/config.yaml b/config/config.yaml index 92fa550..10dd62c 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -242,7 +242,11 @@ analyze_structure_arguments: --pae_filter: 100.0 # interface PAE cutoff; lower values are stricter --models_to_analyse: "all" +# The report scores with analyze_structure_arguments, so it reuses each fold's +# interfaces.csv; flags here are added to them or override them. # Optional combined validation PDF; use an absolute output path. recursive_report_arguments: {} # Add this entry inside recursive_report_arguments (replace {} with the block): # --aggregate_report: /path/to/reports/aggregate_report.pdf +# AlphaJudge processes for the report (--cores); memory scales with them. +# recursive_report_threads: 8 diff --git a/test/test_recursive_report.py b/test/test_recursive_report.py new file mode 100644 index 0000000..d02b700 --- /dev/null +++ b/test/test_recursive_report.py @@ -0,0 +1,118 @@ +"""The recursive report reuses each fold's AlphaJudge table and is sized for its workers. + +The helpers come from ``workflow/rules/common.smk``, loaded directly. The dry run +renders the real rule with the resident-batch fixtures; no container, GPU, +database, or cluster allocation is needed. +""" + +from __future__ import annotations + +import importlib.machinery +import importlib.util +import os +import re +import shutil +import subprocess +from pathlib import Path + +import pytest +import yaml + +_REPOSITORY = Path(__file__).resolve().parents[1] +_FIXTURES = _REPOSITORY / "test" / "fixtures" / "resident_batch" +_COMMON = _REPOSITORY / "workflow" / "rules" / "common.smk" +_loader = importlib.machinery.SourceFileLoader("aps_common", str(_COMMON)) +_spec = importlib.util.spec_from_loader("aps_common", _loader) +common = importlib.util.module_from_spec(_spec) +_loader.exec_module(common) + +ANALYSIS = {"--contact_thresh": 8.0, "--pae_filter": 100.0, "--models_to_analyse": "all"} + + +def test_report_scores_with_the_analysis_settings(): + assert common.recursive_report_arguments(ANALYSIS, {}, 4) == {**ANALYSIS, "--cores": 4} + + +def test_report_flags_extend_or_override_the_analysis_settings(): + report = {"--aggregate_report": "/reports/all.pdf", "--pae_filter": 50.0, "--cores": 2} + assert common.recursive_report_arguments(ANALYSIS, report, 4) == { + "--contact_thresh": 8.0, + "--pae_filter": 50.0, + "--models_to_analyse": "all", + "--aggregate_report": "/reports/all.pdf", + "--cores": 2, + } + + +@pytest.mark.parametrize("report", [["--models_to_analyse=best"], "--models_to_analyse=best"]) +def test_list_or_string_report_arguments_are_used_as_written(report): + assert common.recursive_report_arguments(ANALYSIS, report, 4) == report + + +@pytest.mark.parametrize("tokens,peak_mb", [(6836, 5866), (6910, 5912), (7164, 6264)]) +def test_run_memory_exceeds_measured_alphajudge_peaks(tokens, peak_mb): + # The AF3 runs closest to the bound in a 1,900-run campaign. + assert common.alphajudge_run_mem_mb(tokens) >= 1.25 * peak_mb + + +def test_report_memory_holds_the_largest_run_and_median_runs(): + tokens = [9000, 1000, 3000, 2000] + largest, median = common.alphajudge_run_mem_mb(9000), common.alphajudge_run_mem_mb(3000) + assert common.recursive_report_mem_mb(tokens, workers=4, attempt=1) == int(largest + 3 * median) + assert common.recursive_report_mem_mb(tokens, workers=4, attempt=2) == int(1.5 * (largest + 3 * median)) + assert common.recursive_report_mem_mb(tokens, workers=4, attempt=1, cap_mb=10000) == 10000 + assert common.recursive_report_mem_mb(tokens, workers=1, attempt=1) == int(largest) + + +def test_report_memory_without_lengths_is_the_base_per_worker(): + assert common.recursive_report_mem_mb([], workers=4, attempt=1) == 4 * common.ALPHAJUDGE_BASE_MB + assert common.recursive_report_mem_mb([0, 0], workers=4, attempt=1) == 4 * common.ALPHAJUDGE_BASE_MB + + +def test_report_command_uses_the_analysis_settings_and_its_threads(tmp_path): + snakemake = shutil.which("snakemake") + if snakemake is None: + pytest.skip("Snakemake executable is not available") + config = yaml.safe_load((_FIXTURES / "af3.yaml").read_text(encoding="utf-8")) + config.update( + output_directory=str(tmp_path / "output"), + structure_inference_unified_memory=False, + enable_structure_analysis=True, + generate_recursive_report=True, + analyze_structure_arguments=ANALYSIS, + recursive_report_arguments={"--aggregate_report": "/reports/all.pdf"}, + recursive_report_threads=3, + ) + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + completed = subprocess.run( + [ + snakemake, + "--snakefile", str(_REPOSITORY / "workflow" / "Snakefile"), + "--configfile", str(config_path), + "--cores", "4", + "--rerun-triggers", "mtime", + "--dry-run", + "--printshellcmds", + ], + cwd=_REPOSITORY, + capture_output=True, + text=True, + timeout=120, + env={**os.environ, "APPTAINER_BINDPATH": "", "SINGULARITY_BINDPATH": ""}, + ) + assert completed.returncode == 0, completed.stdout[-2000:] + completed.stderr[-2000:] + job, = [ + job for job in re.split(r"^(?=(?:local)?rule \w+:$)", completed.stdout, flags=re.MULTILINE) + if job.startswith("rule recursive_report:") + ] + assert re.search(r"^\s*threads: 3$", job, re.MULTILINE), job + shell = job.split("Shell command:", 1)[1] + for flag in ( + "--contact_thresh=8.0", + "--pae_filter=100.0", + "--models_to_analyse=all", + "--aggregate_report=/reports/all.pdf", + "--cores=3", + ): + assert flag in shell, (flag, shell) diff --git a/workflow/Snakefile b/workflow/Snakefile index 43e8930..1078784 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -494,6 +494,7 @@ required_folds = [ ENABLE_STRUCTURE_ANALYSIS = config.get("enable_structure_analysis", True) GENERATE_RECURSIVE_REPORT = config.get("generate_recursive_report", False) RECURSIVE_REPORT_ARGUMENTS = config.get("recursive_report_arguments", {}) +RECURSIVE_REPORT_THREADS = max(int(config.get("recursive_report_threads", 8)), 1) RECURSIVE_REPORT = ( join(config["output_directory"], "reports", "all_interfaces.csv") if ENABLE_STRUCTURE_ANALYSIS and GENERATE_RECURSIVE_REPORT @@ -1236,12 +1237,27 @@ if ENABLE_STRUCTURE_ANALYSIS and GENERATE_RECURSIVE_REPORT: tmpdir_guard=lambda wildcards: container_tmpdir_guard(CONTAINER_TMPDIR), prediction_dir=lambda wildcards: join(config["output_directory"], "predictions"), report_dir=join(config["output_directory"], "reports"), - extra_args=lambda wildcards: " ".join( - _format_recursive_arguments(RECURSIVE_REPORT_ARGUMENTS) + extra_args=lambda wildcards, threads: " ".join( + _format_recursive_arguments( + recursive_report_arguments( + config.get("analyze_structure_arguments", {}), + RECURSIVE_REPORT_ARGUMENTS, + threads, + ) + ) ), + threads: RECURSIVE_REPORT_THREADS resources: qos=DEFAULT_SLURM_QOS, - **linear_resources(mem=2000, runtime=1440), + **linear_resources( + mem_fn=lambda wildcards, attempt: recursive_report_mem_mb( + [_fold_grouping_tokens(fold) for fold in kept_folds], + workers=RECURSIVE_REPORT_THREADS, + attempt=attempt, + cap_mb=MAX_MEM_MB, + ), + runtime=1440, + ), container: config["analysis_container"], shell: diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index d20821e..270c9ea 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -600,6 +600,64 @@ def estimate_inference_mem_mb( return _cap_mem(value, cap_mb) +ALPHAJUDGE_BASE_MB = 2000 +ALPHAJUDGE_PER_TOKEN_SQ_MB = 1.2e-4 + + +def alphajudge_run_mem_mb( + tokens: int, + *, + base_mb: float = ALPHAJUDGE_BASE_MB, + per_token_sq_mb: float = ALPHAJUDGE_PER_TOKEN_SQ_MB, +) -> float: + """Peak RAM (MB) of AlphaJudge scoring one run of ``tokens`` tokens. + + Its PAE and contact matrices are O(N^2). The coefficients exceed the measured + peaks of ~1,900 AF2 and AF3 runs of 600-8,800 tokens by at least 1.3x + (largest: 6.3 GB at 7,164 tokens). + """ + return base_mb + per_token_sq_mb * (max(int(tokens), 0) ** 2) + + +def recursive_report_mem_mb( + fold_tokens, + *, + workers: int, + attempt: int, + cap_mb: int = 0, + scaling: float = 1.5, +) -> int: + """Host RAM (MB) for the recursive AlphaJudge report. + + Its ``workers`` AlphaJudge processes score one run each at a time: one may + hold the largest run while the others hold typical (median) runs. Retries + escalate via ``scaling ** (attempt - 1)``. + """ + tokens = sorted(max(int(value), 0) for value in fold_tokens) or [0] + estimate = alphajudge_run_mem_mb(tokens[-1]) + ( + (max(int(workers), 1) - 1) * alphajudge_run_mem_mb(tokens[len(tokens) // 2]) + ) + value = estimate * (scaling ** max(int(attempt) - 1, 0)) + return _cap_mem(value, cap_mb) + + +def recursive_report_arguments(analysis_arguments, report_arguments, threads: int): + """AlphaJudge flags for the recursive report. + + The report reuses a fold's interfaces.csv only when it scores with the same + settings as analyze_structure; any other setting rescores every fold and + replaces its table. So analyze_structure_arguments are the base, and + recursive_report_arguments add report-only flags (``--aggregate_report``) or + override them. ``--cores`` defaults to the job's threads. A list or string is + passed through as written. + """ + if not isinstance(report_arguments, dict): + return report_arguments + arguments = {**(analysis_arguments or {}), **report_arguments} + arguments.setdefault("--cores", int(threads)) + return arguments + + def unified_memory_budget_shortfalls( batches, *,