---
orphan: true
---

# Orchestrator architecture

> Internals doc — for contributors. End users running
> `mhcflurry class1-train-pan-allele-models` don't need to read this.

## The one-line summary

The orchestrator owns the **run-wide resources** (parallel workers,
env tuning); workers consume those and own their per-fit device state.
Workers never build shared state and never set env knobs that affect
other workers.

If you find yourself adding orchestrator-shaped logic inside
`fit()`, `train_model()`, or any worker function, stop and think
whether it belongs in the orchestrator instead.

## Locus of control

```
orchestrator process               training worker process
─────────────────────              ───────────────────────

parse args                         (forked / spawned)
load training data                  receive WORKER_CONTEXT once
build WORKER_CONTEXT        ─┐
                             │
hoist env knobs              │
  TORCHINDUCTOR_COMPILE_THREADS
                             │
start worker pool            ─┘
                                    fit() builds device-resident
                                      AffinityDeviceTrainingData
                                      (peptide/allele/y/weights and
                                      random-negative slice live on
                                      the active torch device)
collect results  ◀──────────────────  return predictor
```

The orchestrator does the **expensive, single-threaded preparation**
ONCE. Workers do the **parallel, per-task** work N times.

## Tensor residency: device-resident affinity fit

Affinity `fit()` is **device-resident**. One container,
`mhcflurry.class1_affinity_training_data.AffinityDeviceTrainingData`, holds
the row space `[random negatives | real training rows]` as torch tensors on
the active device for the lifetime of one `fit()` call. The inner loop
forms batches via `index_select` against the combined buffer — no per-batch
host-to-device copies, no DataLoader workers feeding the affinity loop.

Random-negative peptides are refilled in-place once per epoch into the
top slice of the combined buffer, so the real-data block is never
recopied.

The pretrain-streaming path (`fit_streaming_batches`) is a separate code
path that *does* use a PyTorch DataLoader with optional prefetch workers.
That's controlled by `dataloader_num_workers`; 0 means in-process. It
feeds the pretrain generator only and does not affect the affinity row
layout above.

## Random negatives

Random negatives are worker-local. `RandomNegativesPool` amortizes
generation over `random_negative_pool_epochs`, but it does not try to
share encoded pools between model fits. This preserves the historical
behavior where independently trained workers sample independent negative
peptides.

For the default torch peptide-encoding path, the pool samples amino-acid
indices directly on the worker's active torch device and writes fixed-length
int8 rows. The model's embedding layer expands those indices to BLOSUM62 /
PMBEC / physchem features during the forward pass. Legacy host-vector
models fall back to the host encoder so their input shape remains unchanged.

## Parallelism backends

Two backends, one CLI surface:

- **Local** (`mhcflurry.parallelism`): `multiprocessing.Pool` of
  non-daemon workers. Workers can spawn DataLoader children for the
  pretrain streaming path. `resolve_local_parallelism_args` is the
  single pre-fork normalization point: it resolves
  `--max-workers-per-gpu=auto` and `--dataloader-num-workers=auto`
  without touching CUDA, caps local `--num-jobs` to GPU capacity when
  auto was requested, and hoists torch.compile's thread cap before the
  Pool starts. Fork workers inherit `WORKER_CONTEXT` through copy-on-write;
  spawn workers receive it once in their initializer, never once per queued
  task. Before a spawn pool starts, the planner deep-sizes that loaded context
  and rechecks current host RAM, so worker copies are capped even when the
  source CSV was compressed or torch compilation is disabled. Explicit numeric
  `--max-workers-per-gpu` keeps the
  historical CPU-overflow behavior.
- **Cluster** (`cluster_parallelism.py`): one job per work-item
  submitted via `bsub` / `sbatch` / `sh`. Workers serialize
  WORKER_CONTEXT to NFS and deserialize on the worker side.

Worker-side code (`train_model()`, etc.) is identical between
backends. Only the orchestrator branches on `args.cluster_parallelism`.

## Coverage matrix

The 2.3.0 modernization concentrates on the pan-allele affinity
training + percentile calibration paths. Other components inherit
shared backend selection and worker orchestration, but their data
paths remain intentionally smaller.

| | pretrain | finetune | select | calibrate |
|---|---|---|---|---|
| **affinity (pan-allele)** | streaming DataLoader + compact torch-index peptide batches | device-resident tensors + worker-local RN pool | filter ✓; pool/cache n/a (no fit) | filter ✓; pool/cache n/a |
| **affinity (allele-specific)** | n/a | device-resident tensors + worker-local RN pool | filter ✓ | shares calibrate command |
| **processing** | n/a | local+cluster worker pool | local+cluster worker pool | n/a (allele-independent) |
| **presentation** | n/a | parallel auto-sized feature generation + deterministic fit | n/a | filter ✓ (shares calibrate command) |

## Auto-tuned parallelism knobs

Three knobs auto-derive from the box's hardware so the orchestrator
keeps working when the recipe lands on a different tier. Every auto
resolver lives in `mhcflurry.parallelism` and is exercised by
the unit-test matrix in `test/test_orchestrator_helpers.py`. The
production recipes pass `auto` for each; pin a literal int only when
intentionally re-benchmarking.

### `--max-workers-per-gpu auto` → `auto_max_workers_per_gpu`

Picks per-GPU worker concurrency from the number of complete estimated worker
working sets that fit in detected free VRAM after one shared reserve (10%, at
least 1 GiB). There is no default worker hard cap. The working set is derived
from model artifacts/parameters, gradients and optimizer state for training,
resident encoded data, and the configured training minibatch where those facts
are available; otherwise the workload profile supplies a conservative fallback.
Free VRAM is read from `nvidia-smi` (no torch import — the parent process must
not initialize CUDA before forking). The final job count is also bounded by
available host RAM and vCPUs. In containers, host-memory sizing honors the
current cgroup limit and usage rather than the physical host's `/proc/meminfo`.
Spawn pools then refine the host cap from the loaded per-worker context and
current available RAM; fork pools retain copy-on-write sharing and skip that
copy multiplier. If no additional spawn copy fits safely, automatic local
execution falls back to serial work in the parent instead of forcing one child.

When torch compilation is enabled, the existing one-worker compile warmup also
records CUDA peak-reserved memory and process peak RSS. A measured peak can
tighten an automatic plan before the production pool starts, but can never
change explicit CLI values. Runtime validation and inference batches use the
same shared reserve against live free memory, so they do not independently
spend VRAM already assigned to workers.

For the generic 4 GiB fallback, 80 GiB free fits 18 workers and 40 GiB fits 9;
workload-specific estimates, host RAM, vCPUs, and the number of work items can
reduce those capacities. CPU-only execution resolves to one device worker.

### `--dataloader-num-workers auto` → `auto_dataloader_num_workers`

Picks the per-fit-worker DataLoader prefetch child count from box
capacity. This applies to the pretrain streaming path
(`fit_streaming_batches`); the affinity `fit()` path is device-resident
and ignores it. Inputs: total vCPUs, total RAM in GB, the post-cap
fit-worker count (`--num-jobs` resolved by
`resolve_local_parallelism_args`), and a hard cap (default 4 — the
empirical SM-scheduler-style wall, env-overridable via
`MHCFLURRY_AUTO_DATALOADER_HARD_CAP`).

Heuristic:

1. `cpu_per_fit = vcpus // num_fit_workers`
2. `cpu_cap = cpu_per_fit // 2` — each DL child needs ~2 effective
   cores (1 for the main loop, 1 for queue/collate/copy).
3. RAM cap (when `ram_gb` is provided): assume ~2 GB baseline per
   fit-worker plus ~0.5 GB per DL child;
   `ram_cap = max(0, (ram_gb / num_fit_workers - 2.0) / 0.5)`.
4. Result: `min(hard_cap, cpu_cap, ram_cap)` clamped to ≥1, except
   when any input cap is 0 (oversubscribed mains, RAM exhaustion, or
   user override `MHCFLURRY_AUTO_DATALOADER_HARD_CAP=0`) in which case
   the result is 0 — i.e. in-process batching, no children.
5. Serial / no-fit-worker case: returns 0.

Cross-checked configurations (see
`test_auto_dataloader_num_workers_hardware_tiers`):

| Box | vCPU | RAM | fit | resolved |
|---|---|---|---|---|
| 8×A100-80GB Verda | 176 | 400 G | 16 | **4** |
| 8×A100-40GB | 176 | 400 G | 8 | 4 |
| 8×L40S sweep box | 96 | 200 G | 16 | 3 |
| Single A100-80G Lambda | 30 | 200 G | 2 | 4 |
| Single A100-80G tight | 16 | 64 G | 2 | 4 |
| Single T4 / RTX | 8 | 16 G | 1 | 4 |
| Tight cluster node | 32 | 64 G | 16 | 1 |
| Very tight (8v / 8fit) | 8 | 16 G | 8 | **0** |
| RAM-starved (32G / 16fit) | 176 | 32 G | 16 | **0** |
| Serial / CPU | — | — | 0 | **0** |

**The heuristic is hardware-only by design** — `train_row_count` and
`random_negative_rate` do not enter the formula. Per-batch CPU work is
bounded by `minibatch_size`, not total dataset rows.

### `--num-jobs` (auto-derives from MWPG × GPUs)

`--num-jobs` defaults to `auto`, which resolves to `gpus × max_workers_per_gpu`
via `mhcflurry.parallelism.auto_num_jobs(num_gpus, max_workers_per_gpu)`
once `auto_max_workers_per_gpu` has resolved. Pass an integer to pin it.

### Cross-model coverage

| Model | `--max-workers-per-gpu auto` | `--dataloader-num-workers auto` | Notes |
|---|---|---|---|
| Pan-allele affinity | ✓ | ✓ (pretrain only) | Default in release recipe; affinity `fit()` is device-resident |
| Allele-specific affinity | ✓ | ✓ | Same `Class1NeuralNetwork` codebase; auto already wired |
| Processing | ✓ | n/a | Parallel model fits; the processing network has no DataLoader pretraining path. |
| Presentation | ✓ | n/a | Parallel feature generation followed by deterministic fitting. |

### Env overrides

| Env var | Default | Effect |
|---|---|---|
| `MHCFLURRY_AUTO_MAX_WORKERS_PER_GPU_PER_WORKER_GB` | 4.0 generic fallback | Expert per-worker VRAM override; workload-specific estimates normally replace the fallback. |
| `MHCFLURRY_AUTO_MAX_WORKERS_PER_GPU_HARD_CAP` | unset | Optional expert ceiling for workers per GPU. |
| `MHCFLURRY_AUTO_MAX_WORKERS_PER_GPU_FREE_VRAM_GB` | (auto-detect) | Pin free VRAM (CSV per GPU); for tests / hidden-`nvidia-smi` launchers |
| `MHCFLURRY_AUTO_DATALOADER_HARD_CAP` | 4 | DL child cap for `auto_dataloader_num_workers` |

## Env knobs vs CLI flags

A persistent question: when does a setting belong on `argparse` vs
the environment?

**Rule of thumb:**

- **CLI flag** when the orchestrator owns it and propagates it.
  (Examples: `--max-workers-per-gpu`, `--random-negative-pool-epochs`,
  `--num-jobs`.)
- **Env var with optional CLI relay** when the consumer is inside
  `fit()` or another worker-private code path, and the orchestrator
  only centralizes policy. (Examples: `MHCFLURRY_TORCH_COMPILE`,
  `MHCFLURRY_TORCH_COMPILE_LOSS`, `TORCHINDUCTOR_COMPILE_THREADS`.)

The orchestrator may **hoist** env vars: read its own args, compute a
sensible default, and put a concrete value in `os.environ` so workers
inherit it. `resolve_local_parallelism_args` calls
`hoist_torchinductor_compile_threads` for local runs — it sizes the
inductor compile pool against the resolved `--num-jobs` so N workers
don't each spawn `cpu_count()` compile threads. If
`TORCHINDUCTOR_COMPILE_THREADS=auto`, the orchestrator replaces it
with a numeric value before any worker sees it.

## Torch Compile Warmup

`torch.compile` is worker-local: the Python wrapper, graph guards, and
CUDA module handles cannot be shared across processes. What can be
shared on one machine is the Inductor/Triton on-disk cache. For local
Pool training, the orchestrator therefore runs one real work item
first in a one-worker warmup pool with a larger compile-thread budget,
saves that result, then restores production compile-thread sizing and
launches the full worker pool.

Cluster workers are different: they may land on different nodes, so
mhcflurry does not try to share Inductor cache across a cluster. Each
cluster worker process auto-sizes `TORCHINDUCTOR_COMPILE_THREADS`
locally when the env is unset or `auto`. If a scheduler packs multiple
mhcflurry work items onto one node, set
`MHCFLURRY_CLUSTER_WORKERS_PER_NODE` so each process uses a fair share
of cores.

Compiled losses are enabled by default when `MHCFLURRY_TORCH_COMPILE=1`
and can be disabled with `MHCFLURRY_TORCH_COMPILE_LOSS=0` or
`--torch-compile-loss 0`. CUDA workers run a one-op autograd warmup
before compiling losses to avoid the PyTorch 2.4 / Triton
`invalid device context` failure in the first compiled backward kernel.

## What is NOT the orchestrator's job

- Compiling models (`torch.compile`). Compilation is per-network and
  happens inside `fit()`. The orchestrator only sizes the compile
  worker pool via env.
- Calling `predict()` on individual alleles. The orchestrator
  builds work items; workers iterate.
- Validating data shape consistency between work items beyond what
  the worker-side fit/data constructors already check.

## Recipes

### Adding a new pre-fork resource

1. Add a `_initialize_<name>(args, all_work_items)` helper near the
   existing ones in `train_pan_allele_models_command.py`.
2. Stash the result in `WORKER_CONTEXT["<name>"]`.
3. Document the lookup key. Workers retrieve via
   `constant_data["<name>"]` (forked workers inherit; spawned workers receive
   one initializer payload per process; cluster workers receive the serialized
   context through the shared filesystem).
4. Add a `getattr(args, "<flag>", None)` gate so the helper is opt-in
   on the CLI side.

### Adding a new env knob

1. Decide: orchestrator-set or worker-private?
2. If orchestrator-set: add a `_hoist_<knob>(args)` helper that reads
   args + system info and `os.environ.setdefault(...)`s the value.
   Document the rule the orchestrator applies.
3. If worker-private: just read it from `os.environ` inside `fit()`
   or wherever it's consumed. Document it in
   `RELEASE_NOTES_<version>.md`.

### Adding a new worker-side filter

1. Drop into `mhcflurry/common.py` (e.g.,
   `filter_canonicalizable_alleles`).
2. Apply at every iteration site that could trip on the bad input.
   The 2.3.0 application sites are calibrate (affinity + presentation
   paths) and select (pan-allele + allele-specific paths).

## Fit-time data flow

Affinity training has two configuration knobs that are deliberately
configured separately. Do not collapse these into one.

### 1. Peptide amino-acid vector lookup

Controlled by the model hyperparameter
`peptide_amino_acid_encoding_torch`.

When enabled (the default), peptide strings are encoded as `(N, L)`
integer amino-acid indices. The model owns a frozen torch buffer for
the configured vector encoding (`BLOSUM62`, `PMBEC`, `simons1999_contact`,
or combinations), moves that buffer with `.to(device)`, and widens
indices to `(N, L, V)` inside `forward()` using
`torch.nn.functional.embedding`.

This path is device agnostic: CUDA, MPS, and CPU all execute the same
torch embedding operation on the active device. It is model semantics,
not DataLoader behavior.

Disable it only with `peptide_amino_acid_encoding_torch=False`, which
restores the old path where numpy expands peptide strings into
`(N, L, V)` vectors before they are moved to the device.

### 2. Pretrain DataLoader process parallelism

Controlled by the model hyperparameter `dataloader_num_workers`.

This applies to the pretrain streaming path
(`fit_streaming_batches`) only. It is pure process-count policy
for the pretrain data pipeline:

| value | meaning |
|---|---|
| `0` | build batches in the training worker process |
| `>0` | spawn that many DataLoader child processes per training worker |

Release recipes set only this knob. On a local 8-GPU run with
`NUM_JOBS=16`, `dataloader_num_workers=1` means up to 16 extra
fit-local DataLoader children while pretrain epochs are active; `2`
means up to 32. The thread-budget helper accounts for this when sizing
`OMP_NUM_THREADS`, `MKL_NUM_THREADS`, and `OPENBLAS_NUM_THREADS`.
Caller-provided values remain authoritative and are not replaced by the
planner's uniform runtime limit. For serial (`num_jobs=0`) execution, an
auto-owned budget is applied directly to the current native and PyTorch pools
because there is no worker initializer.

The affinity `fit()` path does not use a DataLoader and therefore
ignores this hyperparameter.

### Component consistency

The `Class1NeuralNetwork` affinity component model owns both
hyperparameters above, so allele-specific affinity models, pan-allele
affinity models, and affinity ensembles resolve the same internal rules
after config load. Missing keys in old component configs are filled
from defaults (`peptide_amino_acid_encoding_torch=True`,
`dataloader_num_workers=0`).

Processing models do not use the affinity `fit()` row layout and
therefore have no `AffinityDeviceTrainingData`.

## rsync hygiene (laptop ↔ remote training box)

`runplz` rsyncs the working tree up to the box on every launch, and
the run's `out/` directory back to the laptop on completion. Two
asymmetries to know about:

- **Up direction:** runplz's hardcoded exclude list (`.git`, `.venv`,
  `__pycache__`, etc.) doesn't anticipate per-project output dirs.
  In mhcflurry, `brev_runs/` accumulates 7-15 GB of past-run
  artifacts that ride along on every launch unless relocated. Run
  `bash scripts/dev/relocate_run_outputs.sh --apply` once to move
  `brev_runs/` and `results/` to `~/mhcflurry-brev-runs/` and
  `~/mhcflurry-results/`, with symlinks back into the repo. After
  that, rsync ships ~tiny symlinks instead of multi-GB directories.
- **Down direction:** the post-run rsync has NO excludes — everything
  under `out/` returns. Keep large throwaway run artifacts outside
  `out/` unless they are meant to ship back.

## Pointers to code

- Random-negative planning and pooling: `mhcflurry/random_negative_peptides.py`
- Affinity device-resident row space:
  `mhcflurry/class1_affinity_training_data.py` (`AffinityDeviceTrainingData`)
- Pseudogene/null filter: `mhcflurry/common.py`
  (`filter_canonicalizable_alleles`)
- Worker pool sizing: `mhcflurry/parallelism`
  (`auto_max_workers_per_gpu`, `resolve_max_workers_per_gpu`)
- Compile-thread hoist:
  `mhcflurry.parallelism.hoist_torchinductor_compile_threads`
- Cluster fork point:
  `mhcflurry/cluster_parallelism.cluster_results`

## Known asymmetries (deliberate)

These show up to readers as "why is this only done in pan-allele?"
The answer in each case is "the other components don't yet need it
and adding it would be feature work, not a fix":

- **`torch.compile`** is off by default everywhere; opt-in via
  `MHCFLURRY_TORCH_COMPILE=1` or `--torch-compile 1`. When enabled,
  the shared local-parallelism layer owns compile-thread sizing and
  local one-worker cache warmup for affinity and processing trainers.
  Presentation fitting is a separate model family and does not enter
  this central torch-training path.

## Future tightening (not in 2.3.0)

- **Presentation-training orchestration.** Today
  `train_presentation_models_command` runs single-process; mirror the
  pan-allele orchestration shape if presentation retraining becomes
  GPU-bound.