fix: decouple inference memory settings and check budgets at startup

README.md

@@ -522,13 +522,14 @@ again of host RAM. With the default `auto`, the fraction is computed **per job a
 time** as
 
 ```text
-1 + (allocated host RAM - structure_inference_ram_bytes) / (physical GPU VRAM)
+1 + max(0, allocated host RAM - structure_inference_unified_memory_margin_mb) / GPU VRAM
 ```
 
 The GPU VRAM is read with `nvidia-smi` once the job lands on a node, and the host RAM is
 the job's SLURM `--mem` allocation (which scales with retry attempts). The job can thus
-spill into all of its host RAM except the base inference RAM (`structure_inference_ram_bytes`,
-by default 8000 MB for AlphaFold 3 and 16000 MB for AlphaFold 2), which stays free for the
+spill into all of its host RAM except the independent process margin
+(`structure_inference_unified_memory_margin_mb`, by default 8000 MB for AlphaFold 3 and
+16000 MB for AlphaFold 2), which stays free for the
 process itself. For example, 128 GB of host RAM on an 80 GB GPU gives 1 + (128 - 8) / 80 = 2.5,
 a 200 GB limit with 120 GB of host spill. The fraction never goes below 1, and the chosen value
 is logged as a `[unified-memory]` line at the top of the job log. Pin a number instead if
@@ -545,9 +546,45 @@ from everything else on the node.
 
 Under `auto` the spill allowance grows with the job's host RAM request (see the
 length-aware memory requests below), so if a large complex still runs out of memory, raise
-that request (`structure_inference_ram_per_token_sq_mb` or `mem_safety_factor`) rather than
+that request (`structure_inference_ram_per_token_sq_mb` or
+`structure_inference_mem_safety_factor`) rather than
 pinning the fraction.
 
+Peak host RSS measures only the part of the XLA pool that spilled into host memory;
+it does not measure the model's total demand. Sizing the request from RSS alone can
+therefore make the total pool too small. Host RAM, GPU routing and the spill margin
+can be tuned independently:
+
+| Setting | Default | Controls |
+|---|---|---|
+| `structure_inference_unified_memory_margin_mb` | AF2: 16000; AF3: 8000 | Host RAM reserved for the process before allowing spill |
+| `structure_inference_vram_per_token_sq_mb` | AF2: 0.0055; AF3: 0.0045 | GPU-tier routing and estimated total inference demand |
+| `structure_inference_mem_safety_factor` | `mem_safety_factor` (normally 1.25) | Inference host RAM only |
+
+With unified memory enabled and the fraction set to `auto`, startup checks every
+inference batch before submission. It compares `min VRAM + max(0, request - margin)`
+with `VRAM coefficient * padded tokens²`. The request includes the inference safety
+factor and `max_mem_mb` cap; routing uses its own headroom and the same padded token
+count. If the estimated budget is too small, the error reports the number of affected
+batches and the worst batch's tokens, request, margin, budget and need.
+
+This is an estimate using the configured tiers, which must cover the schedulable
+GPU pool. It preserves the routing convention of 1000 MB per GB; use conservative
+`min_vram_gb` values for each tier. The runtime fraction still uses the actual
+`nvidia-smi` reading, including MIG slice sizes, so a successful startup check cannot
+guarantee that inference fits. Without tiers the check is skipped unless
+`structure_inference_min_gpu_vram_gb` supplies a hardware lower bound.
+`structure_inference_skip_memory_check: true` bypasses it. Feature-only workflows,
+disabled unified memory and pinned numeric fractions also bypass the check.
+
+**Compatibility:** `structure_inference_ram_bytes` now controls only the host-RAM
+base, and `structure_inference_ram_per_token_sq_mb` only its quadratic coefficient;
+they no longer change the spill margin or GPU routing. Startup `[memory]` notes show
+the effective settings. `mem_safety_factor` still controls both stages unless the
+inference override is set. The default first-attempt estimates are unchanged, but
+the default inference retry multiplier increases from 1.1 to **1.5**. Explicit retry
+multipliers continue to take precedence.
+
 </details>
 
 <details>
@@ -581,9 +618,11 @@ structure_inference  mem = safety * (structure_inference_ram_bytes + per_token_s
   The AF3 inference quadratic is sized to the observed GPU-VRAM demand so that, with unified
   memory, the `auto` limit (GPU VRAM plus the host RAM request minus its base) covers large
   complexes instead of OOM-ing.
-- The first attempt already includes `mem_safety_factor` (default `1.25`) of head-room.
+- The first attempt already includes `mem_safety_factor` (default `1.25`) of head-room;
+  `structure_inference_mem_safety_factor` can override it for inference alone.
   **OOM retries still escalate** on top, multiplying by `..._ram_scaling ** (attempt - 1)`,
-  so a bad estimate self-heals.
+  so a bad estimate self-heals. Inference defaults to `1.5` (+50% per retry, 3.375×
+  on attempt 4 before the memory cap); feature generation retains `1.1`.
 - Override any backend default by setting the matching key in `config/config.yaml`
   (`feature_create_ram_bytes`, `feature_create_ram_per_residue_mb`,
   `structure_inference_ram_bytes`, `structure_inference_ram_per_token_sq_mb`); an explicit

config/config.yaml

@@ -103,15 +103,16 @@ batch_max_tokens: 0
 # All runtime values are minutes. README: "SLURM defaults for structure inference"
 alphafold_inference_threads: 8
 # feature_threads: 8                    # CPUs per feature-generation job
-mem_safety_factor: 1.25                 # margin added to the memory estimate
+mem_safety_factor: 1.25                 # feature margin; inference fallback too
+# structure_inference_mem_safety_factor: 1.25  # inference-only override of the above
 max_mem_mb: 0                           # cap per job; 0 means no cap
 feature_create_ram_scaling: 1.1         # multiply RAM by this on each retry
-structure_inference_ram_scaling: 1.1
+structure_inference_ram_scaling: 1.5     # 50% more host RAM on each OOM retry
 # Overrides replace backend defaults; leave commented unless you need to tune.
 # feature_create_ram_bytes: 64000
 # feature_create_ram_per_residue_mb: 120
 # structure_inference_ram_bytes: 16000
-# structure_inference_ram_per_token_sq_mb: 0.0045
+# structure_inference_ram_per_token_sq_mb: 0.0055  # AF2; AF3 defaults to 0.0045
 # structure_inference_runtime_minutes: 1440
 structure_inference_max_runtime: 10080  # runtime ceiling; fit it to your partition
 
@@ -132,6 +133,8 @@ length_filter_fetch_uniprot: true       # set false to avoid online length looku
 structure_inference_gpus_per_task: 1
 structure_inference_gpu_model: ""
 # structure_inference_gpu_vram_headroom: 1.0
+# Routing and demand use their own coefficient, independent of host RAM tuning.
+# structure_inference_vram_per_token_sq_mb: 0.0055  # AF2; AF3 defaults to 0.0045
 # structure_inference_gpu_tiers:
 #   - {min_vram_gb: 24, nodes: "gpu-24gb-01,gpu-24gb-02"}
 #   - {min_vram_gb: 80, nodes: "gpu-80gb-01,gpu-80gb-02"}
@@ -140,9 +143,14 @@ structure_inference_gpu_model: ""
 # Unified memory lets large predictions use host RAM when GPU memory fills.
 # It can be slower. The fraction is XLA's total limit (VRAM plus host spill) as a
 # multiple of GPU VRAM: 2.0 = all VRAM plus the same amount again of host RAM.
-# "auto" spills into the job's host RAM, minus structure_inference_ram_bytes.
+# "auto" spills into the job's host RAM, minus an independent process margin.
+# structure_inference_unified_memory_margin_mb: 16000  # AF2; AF3 defaults to 8000
 structure_inference_unified_memory: true
 structure_inference_xla_mem_fraction: auto
+# Check estimated first-attempt budgets against the smallest routed GPU at startup.
+# Without tiers the check is skipped, unless a minimum VRAM is supplied (GB).
+# structure_inference_min_gpu_vram_gb: 24
+# structure_inference_skip_memory_check: false
 
 ## -- Optional local MMseqs2 features -----------------------------------------
 # Alternative alignment search for AlphaFold 2/3 proteins and AlphaFold 3 RNA.

test/test_inference_command_line.py

@@ -9,6 +9,7 @@
 
 from __future__ import annotations
 
+import json
 import os
 import re
 import shutil
@@ -54,7 +55,9 @@ def _write_config(tmp_path, backend, *, batch_size, inference_arguments, **extra
     return path, yaml.safe_load(text)
 
 
-def _dry_run(config_path):
+def _dry_run(config_path, *, expect_success=True, extra_args=(), touch=False):
+    # Touch mode schedules jobs to resolve input-dependent resources (<TBD> in
+    # Snakemake 9 dry runs), but never runs the feature/inference shell commands.
     snakemake = shutil.which("snakemake")
     if snakemake is None:
         pytest.skip("Snakemake executable is not available")
@@ -65,23 +68,42 @@ def _dry_run(config_path):
             "--configfile", str(config_path),
             "--cores", "1",
             "--rerun-triggers", "mtime",
-            "--dry-run",
+            "--touch" if touch else "--dry-run",
             "--printshellcmds",
+            *extra_args,
         ],
         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:]
+    if expect_success:
+        assert completed.returncode == 0, completed.stdout[-2000:] + completed.stderr[-2000:]
     return completed
 
 
+def _write_large_fold_config(tmp_path, *, tokens=4608, **extra):
+    """One precomputed AF3 fold: offline length resolution and concrete job resources."""
+    sample_sheet = tmp_path / "sample_sheet.csv"
+    sample_sheet.write_text("Large\n", encoding="utf-8")
+    features = tmp_path / "output" / "features"
+    features.mkdir(parents=True)
+    (features / "Large_af3_input.json").write_text(
+        json.dumps({"sequences": [{"protein": {"id": "A", "sequence": "A" * tokens}}]}),
+        encoding="utf-8",
+    )
+    config_path, _config = _write_config(
+        tmp_path, "alphafold3", batch_size=1, inference_arguments="",
+        input_files=[str(sample_sheet)], feature_directory=[str(features)], **extra,
+    )
+    return config_path
+
+
 def _structure_inference_jobs(stdout):
     """The printed structure_inference jobs, each from its header to the next job's."""
     jobs = re.split(r"^(?=(?:local)?rule \w+:$)", stdout, flags=re.MULTILINE)
-    return [job for job in jobs if job.startswith("rule structure_inference:")]
+    return [job for job in jobs if job.startswith(("rule structure_inference:", "localrule structure_inference:"))]
 
 
 def _invocations(job):
@@ -285,15 +307,129 @@ def test_auto_fraction_lets_the_job_spill_into_its_host_ram(tmp_path, backend):
         _assert_unified_memory_exported(variables, "3.2")
 
 
-def test_auto_fraction_margin_follows_structure_inference_ram_bytes(tmp_path):
-    block = _unified_memory_block(tmp_path, structure_inference_ram_bytes=20000)
+@pytest.mark.parametrize("backend", sorted(_BACKENDS))
+def test_auto_fraction_margin_is_independent_of_structure_inference_ram_bytes(tmp_path, backend):
+    block = _unified_memory_block(tmp_path, backend, structure_inference_ram_bytes=60000)
+
+    logged, variables = _run_unified_memory_block(tmp_path, block, "echo 46068\n", 31110)
+
+    expected = {"alphafold3": "1.502", "alphafold2": "1.328"}[backend]
+    assert logged == {"margin_mb": _BASE_RAM_MB[backend], "gpu_mb": 46068, "fraction": expected}
+    _assert_unified_memory_exported(variables, expected)
+
+
+def test_auto_fraction_uses_explicit_unified_memory_margin(tmp_path):
+    block = _unified_memory_block(tmp_path, structure_inference_unified_memory_margin_mb=20000)
 
     logged, variables = _run_unified_memory_block(tmp_path, block, "echo 81559\n", 16000)
 
     assert logged == {"margin_mb": 20000, "gpu_mb": 81559, "fraction": "1.000"}
     _assert_unified_memory_exported(variables, "1.000")
 
 
+@pytest.mark.parametrize(
+    "extra, excluded",
+    [
+        ({}, "small,medium"),
+        ({"structure_inference_ram_per_token_sq_mb": 0.001}, "small,medium"),
+        ({"structure_inference_vram_per_token_sq_mb": 0.002}, "small"),
+    ],
+)
+def test_gpu_routing_uses_its_own_coefficient(tmp_path, extra, excluded):
+    config_path = _write_large_fold_config(
+        tmp_path,
+        structure_inference_gpu_tiers=[
+            {"min_vram_gb": 24, "nodes": "small"},
+            {"min_vram_gb": 48, "nodes": "medium"},
+            {"min_vram_gb": 80, "nodes": "large"},
+        ],
+        **extra,
+    )
+
+    jobs = _structure_inference_jobs(_dry_run(config_path).stdout)
+
+    assert len(jobs) == 1
+    match = re.search(r"slurm_extra=--exclude=(\S+)", jobs[0])
+    assert match, jobs[0]
+    assert match.group(1).rstrip(",") == excluded
+
+
+@pytest.mark.parametrize("safety, inference_mb", [(None, 20068), (0.5, 8027)])
+def test_inference_safety_override_preserves_feature_memory(tmp_path, safety, inference_mb):
+    data = tmp_path / "output" / "data"
+    data.mkdir(parents=True)
+    for name in ("A", "B"):
+        (data / f"{name}.fasta").write_text(f">{name}\n" + "A" * 100 + "\n", encoding="utf-8")
+    extra = {} if safety is None else {"structure_inference_mem_safety_factor": safety}
+    config_path, _config = _write_config(
+        tmp_path, "alphafold2", batch_size=1, inference_arguments="",
+        feature_directory=[], **extra,
+    )
+
+    completed = _dry_run(config_path, touch=True)
+    stdout = completed.stdout + completed.stderr
+
+    feature_jobs = [job for job in re.split(r"^(?=(?:local)?rule \w+:$)", stdout, flags=re.MULTILINE)
+                    if job.startswith(("rule create_features:", "localrule create_features:"))]
+    assert len(feature_jobs) == 2
+    for job in feature_jobs:
+        assert re.search(r"\bmem_mb=(\d+)", job).group(1) == "95000"
+    for job in _structure_inference_jobs(stdout):
+        assert re.search(r"\bmem_mb=(\d+)", job).group(1) == str(inference_mb)
+
+
+def test_default_inference_retry_requests_fifty_percent_more_memory(tmp_path):
+    config_path = _write_large_fold_config(tmp_path)
+
+    completed = _dry_run(config_path, extra_args=("--attempt", "2"), touch=True)
+    stdout = completed.stdout + completed.stderr
+
+    jobs = _structure_inference_jobs(stdout)
+    assert len(jobs) == 1
+    assert re.search(r"\bmem_mb=(\d+)", jobs[0]).group(1) == "194159"
+
+
+@pytest.mark.parametrize(
+    "extra, fails",
+    [
+        ({}, True),
+        ({"structure_inference_ram_bytes": 8000, "mem_safety_factor": 1.25}, False),
+        ({"structure_inference_skip_memory_check": True}, False),
+        ({"structure_inference_gpu_tiers": []}, False),
+        ({"structure_inference_gpu_tiers": [], "structure_inference_min_gpu_vram_gb": 48}, True),
+        ({"structure_inference_unified_memory": False}, False),
+        ({"structure_inference_xla_mem_fraction": 2.4}, False),
+        ({"only_generate_features": True}, False),
+        ({"mem_safety_factor": 1.25, "structure_inference_mem_safety_factor": 0.2}, True),
+        ({"structure_inference_ram_bytes": 8000, "mem_safety_factor": 1.25, "max_mem_mb": 16000}, True),
+    ],
+)
+def test_startup_checks_unified_memory_budget_before_submission(tmp_path, extra, fails):
+    options = {
+        "structure_inference_ram_bytes": 60000,
+        "mem_safety_factor": 0.2,
+        "structure_inference_gpu_tiers": [{"min_vram_gb": 48, "nodes": "n48"}],
+        **extra,
+    }
+    # 4500 residues run at 4608 padded tokens; using the unpadded length understates demand.
+    config_path = _write_large_fold_config(tmp_path, tokens=4500, **options)
+
+    completed = _dry_run(config_path, expect_success=not fails)
+
+    log = completed.stdout + completed.stderr
+    if fails:
+        assert completed.returncode != 0
+        assert "1 inference batch(es)" in log
+        assert "tokens=4608" in log
+        assert "margin_mb=8000" in log
+        assert "need_mb=95551" in log
+        assert "structure_inference_unified_memory_margin_mb" in log
+        assert "structure_inference_vram_per_token_sq_mb" in log
+        assert "structure_inference_mem_safety_factor" in log
+    else:
+        assert "inference batch(es)" not in log
+
+
 def test_pinned_fraction_is_exported_as_written(tmp_path):
     block = _unified_memory_block(tmp_path, structure_inference_xla_mem_fraction=2.4)
 

test/test_memory_resources.py

@@ -404,6 +404,62 @@ def test_gpu_exclude_nodes_vram_routing():
     assert common.gpu_exclude_nodes(2400, [], c) == ""
 
 
+def test_routed_tier_min_vram_gb_matches_gpu_exclusions():
+    tiers = [
+        {"min_vram_gb": 80, "nodes": "n80"},
+        {"min_vram_gb": 24, "nodes": "n24"},
+        {"min_vram_gb": 48, "nodes": "n48"},
+        {"min_vram_gb": 16, "nodes": ""},
+    ]
+    for tokens, headroom, minimum, excluded in (
+        (800, 1.0, 24, ""),
+        (2400, 1.0, 48, "n24"),
+        (3500, 1.0, 80, "n24,n48"),
+        (20000, 1.0, 80, "n24,n48"),
+        (3500, 0.5, 48, "n24"),
+    ):
+        assert common.routed_tier_min_vram_gb(tokens, tiers, 0.0045, headroom) == minimum
+        assert common.gpu_exclude_nodes(tokens, tiers, 0.0045, headroom) == excluded
+    assert common.routed_tier_min_vram_gb(4608, [], 0.0045) is None
+
+
+def test_unified_memory_budget_shortfalls_match_requests_and_routing():
+    options = dict(
+        tiers=[{"min_vram_gb": 48, "nodes": "n48"}],
+        base_mb=60000, per_token_sq_mb=0.0045, safety=0.2,
+        margin_mb=8000, vram_per_token_sq_mb=0.0045,
+    )
+    shortfalls = common.unified_memory_budget_shortfalls({"large": 4608, "small": 512}, **options)
+    assert len(shortfalls) == 1
+    worst = shortfalls[0]
+    assert worst["batch"] == "large"
+    assert worst["tokens"] == 4608
+    assert worst["request_mb"] == 31110
+    assert worst["margin_mb"] == 8000
+    assert worst["budget_mb"] == 71110
+    assert round(worst["need_mb"]) == 95551
+
+    for changes, expected in (
+        ({"base_mb": 8000, "safety": 1.25}, []),
+        # The cap applies before checking, just as it does to the submitted request.
+        ({"base_mb": 8000, "safety": 1.25, "cap_mb": 16000}, ["large"]),
+        ({"tiers": []}, []),  # no hardware lower bound: unknown, not infeasible
+        ({"tiers": [], "min_gpu_vram_gb": 48}, ["large"]),
+        ({"tiers": [], "min_gpu_vram_gb": 96}, []),
+        # A margin above the request leaves the whole GPU usable for a small fold.
+        ({"margin_mb": 200000}, ["large"]),
+        # Headroom affects routing only; the demand estimate stays at headroom=1.
+        ({"tiers": [{"min_vram_gb": 48, "nodes": "n48"},
+                    {"min_vram_gb": 96, "nodes": "n96"}], "headroom": 0.5}, ["large"]),
+        ({"tiers": [{"min_vram_gb": 48, "nodes": "n48"},
+                    {"min_vram_gb": 96, "nodes": "n96"}]}, []),
+    ):
+        result = common.unified_memory_budget_shortfalls(
+            {"large": 4608, "small": 512}, **{**options, **changes}
+        )
+        assert [item["batch"] for item in result] == expected, changes
+
+
 def test_af3_padded_tokens_follow_buckets():
     assert common.af3_padded_tokens(1537) == 2048
     assert common.af3_padded_tokens(2048) == 2048