From 4752e17c46584760518e36ac116e0fa42d21b2b5 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 13 Oct 2026 09:41:07 +0200 Subject: [PATCH] fix: decouple inference memory settings and check budgets at startup Three inference memory settings each controlled a second thing: the base RAM was also the unified-memory margin, the host RAM coefficient also routed GPUs, and mem_safety_factor sized feature jobs too. - structure_inference_unified_memory_margin_mb sets the spill margin - structure_inference_vram_per_token_sq_mb sets GPU routing and demand - structure_inference_mem_safety_factor overrides the factor for inference only All default to today's values. With unified memory on and the fraction auto, startup now stops with a WorkflowError when any batch's estimated budget on its smallest routed GPU is below its estimated need, unless structure_inference_skip_memory_check is set. Routing and the check share routed_tier_min_vram_gb, so they cannot disagree. The default inference retry multiplier rises from 1.1 to 1.5. Closes #66. Co-Authored-By: Claude Opus 5.5 --- README.md | 51 ++++++++-- config/config.yaml | 16 ++- test/test_inference_command_line.py | 148 ++++++++++++++++++++++++++-- test/test_memory_resources.py | 56 +++++++++++ workflow/Snakefile | 72 +++++++++++++- workflow/rules/common.smk | 70 +++++++++++-- 6 files changed, 384 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 0815dd8..89b0989 100644 --- a/README.md +++ b/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. +
@@ -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 diff --git a/config/config.yaml b/config/config.yaml index 4f63247..9b02ee3 100644 --- a/config/config.yaml +++ b/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. diff --git a/test/test_inference_command_line.py b/test/test_inference_command_line.py index e3c5744..6626a11 100644 --- a/test/test_inference_command_line.py +++ b/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 ( 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,8 +68,9 @@ 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, @@ -74,14 +78,32 @@ def _dry_run(config_path): 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,8 +307,19 @@ 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) @@ -294,6 +327,109 @@ def test_auto_fraction_margin_follows_structure_inference_ram_bytes(tmp_path): _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) diff --git a/test/test_memory_resources.py b/test/test_memory_resources.py index 907567a..bae7d4d 100644 --- a/test/test_memory_resources.py +++ b/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 diff --git a/workflow/Snakefile b/workflow/Snakefile index 1644cb3..dcb9c4d 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -616,6 +616,9 @@ rule download_uniprot: # escalated on OOM retries via `scaling ** (attempt - 1)`. Defaults differ by # backend (AF2 is heavier than AF3); an explicit config value always wins. See README. MEM_SAFETY_FACTOR = float(config.get("mem_safety_factor", 1.25)) +STRUCTURE_INFERENCE_MEM_SAFETY = float(config.get( + "structure_inference_mem_safety_factor", MEM_SAFETY_FACTOR +)) MAX_MEM_MB = int(config.get("max_mem_mb", 0) or 0) # 0 = no cap # Backend defaults (FEATURE_BACKEND / INFERENCE_BACKEND are defined near the top). @@ -639,7 +642,13 @@ structure_inference_base_ram = config.get( structure_inference_ram_per_token_sq = config.get( "structure_inference_ram_per_token_sq_mb", _infer_ram_defaults["per_token_sq_mb"] ) -structure_inference_ram_scaling = config.get("structure_inference_ram_scaling", 1.1) +structure_inference_margin_mb = int(float(config.get( + "structure_inference_unified_memory_margin_mb", _infer_ram_defaults["base_mb"] +))) +structure_inference_vram_per_token_sq = config.get( + "structure_inference_vram_per_token_sq_mb", _infer_ram_defaults["per_token_sq_mb"] +) +structure_inference_ram_scaling = config.get("structure_inference_ram_scaling", 1.5) structure_inference_runtime_minutes = config.get( "structure_inference_runtime_minutes", _infer_ram_defaults["runtime_minutes"] ) @@ -654,6 +663,61 @@ structure_inference_runtime_minutes = config.get( GPU_TIERS = config.get("structure_inference_gpu_tiers") or [] GPU_VRAM_HEADROOM = float(config.get("structure_inference_gpu_vram_headroom", 1.0)) +if "structure_inference_ram_bytes" in config: + logger.info( + "[memory] structure_inference_ram_bytes controls host RAM only; " + f"structure_inference_unified_memory_margin_mb={structure_inference_margin_mb}." + ) +if "structure_inference_ram_per_token_sq_mb" in config: + logger.info( + "[memory] structure_inference_ram_per_token_sq_mb controls host RAM only; " + f"structure_inference_vram_per_token_sq_mb={structure_inference_vram_per_token_sq} " + "controls GPU routing and estimated total demand." + ) +if "mem_safety_factor" in config: + _inference_safety_source = ( + "inference override" if "structure_inference_mem_safety_factor" in config + else "fallback from mem_safety_factor" + ) + logger.info( + f"[memory] mem_safety_factor={MEM_SAFETY_FACTOR} controls feature RAM; " + f"structure_inference_mem_safety_factor={STRUCTURE_INFERENCE_MEM_SAFETY} " + f"controls inference RAM ({_inference_safety_source})." + ) + +if ( + not _as_bool(config.get("only_generate_features", False)) + and STRUCTURE_INFERENCE_UNIFIED_MEMORY + and STRUCTURE_INFERENCE_XLA_MEM_FRACTION == "auto" + and not _as_bool(config.get("structure_inference_skip_memory_check", False)) +): + _memory_shortfalls = unified_memory_budget_shortfalls( + {batch_id: batch_max_tokens_for(batch_id) for batch_id in BATCH_FOLDS}, + tiers=GPU_TIERS, + base_mb=structure_inference_base_ram, + per_token_sq_mb=structure_inference_ram_per_token_sq, + safety=STRUCTURE_INFERENCE_MEM_SAFETY, + margin_mb=structure_inference_margin_mb, + vram_per_token_sq_mb=structure_inference_vram_per_token_sq, + headroom=GPU_VRAM_HEADROOM, + cap_mb=MAX_MEM_MB, + min_gpu_vram_gb=config.get("structure_inference_min_gpu_vram_gb"), + ) + if _memory_shortfalls: + _worst = max(_memory_shortfalls, key=lambda item: item["need_mb"] - item["budget_mb"]) + raise WorkflowError( + f"[memory] {len(_memory_shortfalls)} inference batch(es) have an estimated " + "unified-memory budget below demand. Worst batch " + f"{_worst['batch']}: tokens={_worst['tokens']}, request_mb={_worst['request_mb']}, " + f"margin_mb={_worst['margin_mb']}, min_vram_gb={_worst['min_vram_gb']:g}, " + f"budget_mb={_worst['budget_mb']:.0f}, need_mb={_worst['need_mb']:.0f}. " + "Review structure_inference_mem_safety_factor (host RAM), " + "structure_inference_unified_memory_margin_mb (reserved host RAM), and " + "structure_inference_vram_per_token_sq_mb (routing and demand), plus " + "max_mem_mb if capped. To bypass this estimate, set " + "structure_inference_skip_memory_check: true." + ) + def _inference_slurm_extra(wildcards): """slurm_extra resource: --exclude small-GPU nodes for large complexes, merged @@ -662,7 +726,7 @@ def _inference_slurm_extra(wildcards): nodes = gpu_exclude_nodes( tokens, GPU_TIERS, - structure_inference_ram_per_token_sq, + structure_inference_vram_per_token_sq, GPU_VRAM_HEADROOM, DEFAULT_SLURM_EXCLUDE or "", ) @@ -964,7 +1028,7 @@ rule structure_inference: unified_memory="true" if STRUCTURE_INFERENCE_UNIFIED_MEMORY else "false", xla_mem_fraction=STRUCTURE_INFERENCE_XLA_MEM_FRACTION, # Host RAM that "auto" keeps for the process itself rather than for GPU spill. - host_ram_margin_mb=int(float(structure_inference_base_ram)), + host_ram_margin_mb=structure_inference_margin_mb, resident=lambda wc: "true" if len(BATCH_FOLDS[wc.batch]) > 1 else "false", resources: slurm_partition=(DEFAULT_SLURM_PARTITION or "gpu"), @@ -994,7 +1058,7 @@ rule structure_inference: base_mb=structure_inference_base_ram, per_token_sq_mb=structure_inference_ram_per_token_sq, scaling=structure_inference_ram_scaling, - safety=MEM_SAFETY_FACTOR, + safety=STRUCTURE_INFERENCE_MEM_SAFETY, attempt=attempt, cap_mb=MAX_MEM_MB, ), diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 2889670..aa16fbf 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -488,12 +488,27 @@ def required_gpu_vram_gb( ) -> float: """Estimated peak GPU VRAM (GB) for a complex of ``total_tokens``. - Uses the same O(N^2) coefficient as host-memory sizing as a proxy for on-device - peak demand, scaled by ``headroom`` (e.g. 0.8 tolerates ~20% spill to host). + Uses an independent O(N^2) coefficient for peak demand, scaled by ``headroom`` + (e.g. 0.8 tolerates ~20% spill to host). Tier routing uses 1000 MB per GB. """ return headroom * per_token_sq_mb * (max(int(total_tokens), 0) ** 2) / 1000.0 +def routed_tier_min_vram_gb( + total_tokens: int, tiers, per_token_sq_mb: float, headroom: float = 1.0 +): + """Minimum VRAM of the selected tier, or None when routing is not configured. + + Use the smallest tier that fits the estimate, or the largest tier with host + spill if no tier fits. Entries without nodes do not participate in routing. + """ + sizes = sorted(int(tier["min_vram_gb"]) for tier in tiers if tier and tier.get("nodes")) + if not sizes: + return None + required = required_gpu_vram_gb(total_tokens, per_token_sq_mb, headroom) + return next((size for size in sizes if size >= required), sizes[-1]) + + def gpu_exclude_nodes( total_tokens: int, tiers, @@ -518,13 +533,8 @@ def gpu_exclude_nodes( valid = [t for t in tiers if t and t.get("nodes")] if valid: ordered = sorted(valid, key=lambda t: int(t["min_vram_gb"])) - required = required_gpu_vram_gb(total_tokens, per_token_sq_mb, headroom) - chosen = len(ordered) - 1 - for index, tier in enumerate(ordered): - if int(tier["min_vram_gb"]) >= required: - chosen = index - break - parts.extend(str(tier["nodes"]) for tier in ordered[:chosen]) + minimum = routed_tier_min_vram_gb(total_tokens, valid, per_token_sq_mb, headroom) + parts.extend(str(tier["nodes"]) for tier in ordered if int(tier["min_vram_gb"]) < minimum) if extra_exclude: parts.append(str(extra_exclude)) return ",".join(part for part in parts if part) @@ -589,6 +599,48 @@ def estimate_inference_mem_mb( return _cap_mem(value, cap_mb) +def unified_memory_budget_shortfalls( + batches, + *, + tiers, + base_mb: float, + per_token_sq_mb: float, + safety: float, + margin_mb: int, + vram_per_token_sq_mb: float, + headroom: float = 1.0, + cap_mb: int = 0, + min_gpu_vram_gb: float | None = None, +): + """Batches whose first-attempt auto budget falls below estimated total demand. + + ``batches`` maps batch identities to the largest member's padded token count. + The caller uses the same tokens for scheduling. VRAM follows the routing + convention (1000 MB/GB); tier minima are estimates, not runtime measurements. + Without valid tiers, an explicit hardware lower bound is needed to check. + """ + shortfalls = [] + for batch, tokens in batches.items(): + minimum = routed_tier_min_vram_gb(tokens, tiers, vram_per_token_sq_mb, headroom) + if minimum is None: + minimum = min_gpu_vram_gb + if minimum is None or float(minimum) <= 0: + continue + request = estimate_inference_mem_mb( + tokens, base_mb=base_mb, per_token_sq_mb=per_token_sq_mb, + safety=safety, scaling=1.0, attempt=1, cap_mb=cap_mb, + ) + budget = float(minimum) * 1000 + max(0, request - margin_mb) + need = required_gpu_vram_gb(tokens, vram_per_token_sq_mb) * 1000 + if budget < need: + shortfalls.append({ + "batch": batch, "tokens": tokens, "request_mb": request, + "margin_mb": margin_mb, "min_vram_gb": float(minimum), + "budget_mb": budget, "need_mb": need, + }) + return shortfalls + + # Backend-specific memory defaults. AlphaFold-Multimer (AF2) inference carries a # substantially heavier host-RAM footprint than AlphaFold 3 at the same complex # size (measured host RSS ~4x higher around N~2300 in the benchmark campaign), and