Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 7 additions & 66 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -340,26 +340,6 @@ prediction_container: "/path/to/images/alphafold3-2.5.0.sif"

</details>

### Container temporary directories

Errors such as `ptxas fatal: Could not open output file` for a `tmpxft_*` file can
mean the job's inherited `TMPDIR` points to scratch that is not mounted inside
the container. Python may silently choose another temporary directory, but
native tools such as XLA's `ptxas` use `TMPDIR` directly.

Every container rule checks `TMPDIR` before starting work. A usable directory is
kept without logging; an unset or inaccessible path falls back to `/tmp`, with
one `[container-tmpdir]` message in the job log. To choose another fallback:

```yaml
container_tmpdir: /path/visible/inside/container
```

The fallback must be writable and searchable inside the container. The guard
creates it when possible and stops with a clear error if it is unusable. This
setting does not add a host bind mount; configure site mounts separately if the
fallback is outside the container's existing mounts.

### GPU compatibility

The containers carry their own CUDA runtime (pip `nvidia-*` wheels), so GPU support depends on the
Expand Down Expand Up @@ -542,14 +522,13 @@ again of host RAM. With the default `auto`, the fraction is computed **per job a
time** as

```text
1 + max(0, allocated host RAM - structure_inference_unified_memory_margin_mb) / GPU VRAM
1 + (allocated host RAM - structure_inference_unified_memory_margin_mb) / (physical 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 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
spill into all of its host RAM except a 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
Expand All @@ -566,45 +545,9 @@ 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
`structure_inference_mem_safety_factor`) rather than
that request (`structure_inference_ram_per_token_sq_mb` or `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>
Expand Down Expand Up @@ -636,13 +579,11 @@ structure_inference mem = safety * (structure_inference_ram_bytes + per_token_s
| `alphafold3` | 40000 MB | 25 MB | 8000 MB | 0.0045 |

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
memory, the `auto` limit (GPU VRAM plus the host RAM request minus the margin) covers large
complexes instead of OOM-ing.
- 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.
- The first attempt already includes `mem_safety_factor` (default `1.25`) of head-room.
**OOM retries still escalate** on top, multiplying by `..._ram_scaling ** (attempt - 1)`,
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`.
so a bad estimate self-heals.
- 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
Expand Down
Loading