[DimaMolod]

Three inference memory keys each controlled a second thing their names don't mention, so tuning host RAM also changed the spill margin, GPU routing or feature memory. In the October campaign, structure_inference_ram_bytes: 60000 with mem_safety_factor: 0.2 set the margin above the request. That clamped the auto fraction to 1.000, and about 100 AF3 jobs failed with RESOURCE_EXHAUSTED.

Changes

• New keys. Each defaults to today's value, so an unmodified config submits the same first attempts:

  Key	Default	Takes over from
  structure_inference_unified_memory_margin_mb	AF2 16000, AF3 8000	structure_inference_ram_bytes (spill margin)
  structure_inference_vram_per_token_sq_mb	AF2 0.0055, AF3 0.0045	structure_inference_ram_per_token_sq_mb (GPU routing)
  structure_inference_mem_safety_factor	mem_safety_factor	mem_safety_factor (inference only)

• Startup feasibility check. It runs when unified memory is on and the fraction is auto:

  • For every batch it compares budget = smallest routed tier's VRAM + max(0, request − margin) with need = VRAM coefficient × padded tokens².
  • The request includes the inference safety factor and max_mem_mb.
  • Any shortfall stops the parse with a WorkflowError. The error names how many batches fall short, the worst batch (tokens, request, margin, VRAM, budget, need) and the knobs to adjust.
  • Without GPU tiers it is skipped unless structure_inference_min_gpu_vram_gb is set, and structure_inference_skip_memory_check: true bypasses it.
  • Defaults cannot trip it: VRAM + 1.25·(base + kN²) − base ≥ kN².

• routed_tier_min_vram_gb() is factored out of gpu_exclude_nodes() and used by the check, so routing and the check cannot disagree.

• [memory] startup lines appear when structure_inference_ram_bytes, structure_inference_ram_per_token_sq_mb or mem_safety_factor is set explicitly. They say what each key now controls.

• The default structure_inference_ram_scaling goes from 1.1 to 1.5, in the Snakefile, config/config.yaml and the README.

• README "Unified memory" now covers the budget formula, the three keys, why peak RSS undercounts need, and the check.

Testing

• Every coupling reproduced on main, with tests that failed there:
  • Margin. structure_inference_ram_bytes: 60000 logged margin_mb=60000 and fraction 1.000 for both backends. The margin now stays at the backend default, giving fractions 1.328 (AF2) and 1.502 (AF3).
  • Routing. Lowering structure_inference_ram_per_token_sq_mb to 0.001 dropped the --exclude for a 4,608-token fold. Routing now follows only structure_inference_vram_per_token_sq_mb.
  • Safety factor. structure_inference_mem_safety_factor: 0.5 was ignored: inference mem_mb was 20068 where 8027 was expected. Feature mem_mb stays at 95000 either way.
  • Retries. Attempt 2 requested 142383 MB (×1.1); it now requests 194159 MB (×1.5).
  • Startup check. The incident config (AF3, 4,500 residues → 4,608 padded tokens, 48 GB tier) passed the dry run on main. It now stops at startup with need_mb=95551.
    • Passes: the default config, the skip flag, no tiers, unified memory off, a pinned fraction, feature-only runs.
    • Trips: a minimum-VRAM floor without tiers, the inference-only safety override, a max_mem_mb cap.
• Unit tests: routed_tier_min_vram_gb matches gpu_exclude_nodes across tiers and headroom; unified_memory_budget_shortfalls is tested across caps, tiers, margin and headroom.
• test_auto_fraction_margin_follows_structure_inference_ram_bytes is replaced, as the issue proposes, because it pinned the coupling being removed.
• Full suite: 196 passed, 1 skipped; with Container rules inherit a host TMPDIR that isn't bound inside the container, so XLA's ptxas fails #65 and Local MMseqs2: a database rebuilt in place under the same identifier is never re-searched #56 merged on top, 213 passed, 1 skipped.

Compatibility

• Configs that set the old keys explicitly now behave differently. The [memory] lines say so:
  • structure_inference_ram_bytes no longer moves the spill margin;
  • structure_inference_ram_per_token_sq_mb no longer moves GPU routing;
  • mem_safety_factor still applies to both stages unless structure_inference_mem_safety_factor is set.
• Retries ask for more memory by default. In fix: let the auto XLA fraction spill into the job's host RAM #63's throughput review, 75% of folds killed by a SLURM memory limit finished after the ×1.1 escalation, and 1.1 was kept. ×1.5 spends more host RAM on retries to stop undersized folds failing every attempt at almost the same size: attempt 4 now gets 3.375× rather than 1.33×. Configs that set structure_inference_ram_scaling: 1.1 explicitly keep the old behaviour.
• A config that runs today can now stop at startup, for example one whose max_mem_mb caps a large fold below its need. The check is an estimate from the tier minima, not from nvidia-smi. Set structure_inference_skip_memory_check: true to proceed anyway.

Not in this PR (follow-ups listed in the issue): sizing host RAM from need minus the routed VRAM, and recording each fold's device peak memory.

Closes #66.

🤖 Generated with Claude Code