Recursive report: reuse per-fold tables, size memory and CPUs

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

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)

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:

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,
     *,