Design a distributed AI model downloader
Remix this answer and make it yours, or start blank. Either way you work beside the reference, one chapter at a time, and it stays private until you publish it.
Problem statement
Get a new multi-hundred-gigabyte model checkpoint onto thousands of GPU hosts as fast as possible without overloading the storage it comes from.
Requirements
We are building the system that takes a newly written model checkpoint (hundreds of GB in object storage) and lands a verified copy on the local disk of every GPU host in a fleet of thousands, as fast as the network allows, while reading each byte from the origin store roughly once per data center instead of once per host.
Functional requirements
- An operator or the training pipeline starts a distribution job: checkpoint (model, version, manifest), target hosts (a selector such as
cluster=inference-a, pool=sonnet), priority, and an origin bandwidth cap. - Every target host ends up with a byte-exact copy, verified against a signed manifest of per-chunk hashes, placed atomically at
/models/<model>/<version>/. - Hosts that join late, crash, or reboot mid-job resume from what they already have; they never restart from zero.
- Progress is visible per job: hosts done, bytes moved, origin bytes read, stragglers, ETA.
- When enough hosts are ready, the system hands off to the deployer ("version V ready on hosts H"), which does the rolling switch.
- Old versions are kept for fast rollback (N-1 at least) and garbage-collected by policy.
Non-functional requirements
- Speed: 800 GB to 4,000 hosts in under 3 minutes p50, under 5 minutes to 99.5% of hosts. The bound is host NIC bandwidth, not the origin.
- Origin protection: origin reads stay near 1x the checkpoint size per data center, and never above the configured cap (e.g. 100 Gbps), regardless of fleet size.
- Correctness: a host never serves or loads a corrupt or partial checkpoint. Integrity is end-to-end (hash at write time, verify at every receiver).
- Fault tolerance: 1-2% of hosts failing or going slow during a job must not stall the rest. The coordinator going down must not stop bytes that are already flowing.
- No interference: distribution uses the frontend Ethernet NICs and must not starve serving traffic; it is rate-limited per host when the host is serving.
Back-of-the-envelope estimation
- Checkpoint S = 800 GB (bf16, a few hundred billion parameters). Chunk = 64 MB (unit of hashing and bookkeeping) → 12,500 chunks. Block = 4 MB (unit of forwarding) → 200,000 blocks.
- Fleet N = 4,096 hosts, 128 racks of 32. Frontend NIC 100 Gbps → ~10 GB/s usable. Local NVMe RAID writes ~15-20 GB/s, so the NIC is the bottleneck.
- Naive (everyone pulls from object storage): 4,096 × 800 GB = 3.3 PB of origin reads. At a 1 Tbps (125 GB/s) origin cap that is ~26,000 s, about 7 hours, and it melts the storage tier for everyone else.
- Lower bound with peer forwarding: every host must receive S bytes, so time ≥ S / B = 800 GB / 10 GB/s = 80 s. A good design gets within ~1.5x of this: ~2 minutes.
- Origin load with peer forwarding: 800 GB per data center, read at ~10 GB/s (80 Gbps), under the cap.
- Coordinator state: one bit per (host, chunk) = 4,096 × 12,500 = 51 Mbit ≈ 6.4 MB per job. Progress reports every 2 s from 4,096 hosts ≈ 2,000 RPC/s. One process handles this easily.
API design
Control plane is gRPC between agents (one per host) and the coordinator. Data plane is a direct host-to-host streaming RPC.
// Operator / pipeline
CreateJob({model: "claude-x", version: "2026-10-12.3", manifest_uri: "s3://ckpt/.../manifest.json",
targets: {selector: "cluster=inf-a,pool=x"}, origin_cap_gbps: 100,
priority: "high", ready_threshold: 0.995})
→ {job_id: "dj_81f2", state: "PLANNING"}
GetJob({job_id}) → {state, hosts_total: 4096, hosts_ready: 3911, bytes_moved, origin_bytes, eta_s, stragglers: [...]}
CancelJob({job_id}) → {}
// Agent ↔ coordinator
Register({job_id, host_id, rack, have: <chunk bitmap>, free_disk_gb}) → {manifest_ref, topology_version}
GetAssignment({job_id, host_id}) →
{topology_version: 7, stripes: [{stripe: 3, upstream: "h-0412" | "ORIGIN", downstreams: ["h-0413"], resume_block: 18240}]}
ReportProgress({job_id, host_id, topology_version, stripe_offsets: [..8], verified_chunks_delta: [...],
rx_gbps, tx_gbps, errors: [{kind: "HASH_MISMATCH", chunk: 9012, from: "h-0410"}]})
→ {topology_version: 8, actions: ["REFETCH_ASSIGNMENT"]} // piggybacked changes
ReportSuspect({job_id, reporter, suspect_host, reason: "STALLED" | "CONN_RESET" | "BAD_DATA"}) → {new_upstream}
// Peer data plane (server-streaming)
StreamStripe({job_id, stripe, from_block}) → stream {block_idx, chunk_idx, offset, bytes[4 MB], crc32c}
FetchChunk({job_id, chunk_idx}) → bytes[64 MB] // repair path for holes
Assignments are versioned. An agent acts only on the newest topology_version it has seen, so a stale coordinator reply cannot rewire a pipeline backwards.
Data model
Small, strongly consistent control state lives in Postgres (or etcd for a lighter footprint). High-churn progress lives in coordinator memory and on hosts.
| Table / key | Key columns | Notes |
|---|---|---|
manifests |
(model, version) PK, total_bytes, chunk_size, signature, files JSON |
Written by the checkpoint writer. Immutable. |
manifest_chunks |
(model, version, chunk_idx) PK, file, offset, len, sha256
|
12,500 rows per version; loaded once per job. |
jobs |
job_id PK, model, version, state, origin_cap, created_at; index (state)
|
PLANNING → RUNNING → READY → DONE / FAILED / CANCELLED. |
job_hosts |
(job_id, host_id) PK, rack, state, ready_at; index (job_id, state)
|
PENDING → RECEIVING → VERIFIED → READY, or QUARANTINED. |
topology |
(job_id, version) PK, edges blob |
Each rewire writes a new version; the coordinator keeps the latest in memory. |
host_versions |
(host_id, model, version) PK, state, bytes
|
What each disk holds; drives skip-if-present and GC. |
On each host, a local journal /models/.staging/<version>/journal records verified chunk indexes and per-stripe block offsets with fsync every few seconds. This is the source of truth for resume; the coordinator's bitmap is a cache rebuilt from Register and ReportProgress.
Invariants: a chunk counts as present only after its SHA-256 matches the manifest; a host is READY only after every chunk is verified and the directory has been renamed from staging; manifests and chunks are content-addressed and never mutated.
High-level design
flowchart LR
W[Training job] --> OS[(Object storage)]
W --> M[Manifest service]
OP[Operator or pipeline] --> CO[Coordinator]
CO --> DB[(Postgres)]
M --> CO
CO -. assignments .-> A1
CO -. assignments .-> A2
OS -->|"stripe 0 to 7"| R1[Stripe roots]
R1 --> A1[Rack leaders]
A1 --> A2[Hosts in rack]
A2 --> D[Local NVMe]
CO --> DEP[Deployer]
-
Start a job. The training pipeline uploads shards and a manifest (per-chunk SHA-256, signed).
CreateJobresolves the target selector to 4,096 hosts, skips hosts whosehost_versionsrow already shows the version, and builds a topology. Agents learn about the job through their long-poll and callRegisterwith whatever chunks they already hold. -
Topology. The checkpoint is split into k = 8 stripes (chunk index mod 8). For each stripe the coordinator builds a two-level pipeline: one chain through a leader in every rack (128 hops), and from each leader a chain through the other 31 hosts in its rack. Leaders differ per stripe (host
sin each rack leads stripes), so no host uploads much more than it downloads. - External fetch. For each stripe, the first leader in the chain is the stripe root. It reads its stripe from object storage with range GETs at ≤ cap/8 and forwards each 4 MB block downstream as soon as it lands (cut-through), before the whole chunk arrives. Origin reads total 1x the checkpoint.
-
Pipeline forwarding. Every host receives 8 stripes at once (~1.25 GB/s each), writes blocks to staging files with
pwriteat the manifest offset, and streams them to its downstreams. When a 64 MB chunk completes, it is hashed and journaled. -
Completion. When all 12,500 chunks verify, the agent fsyncs, renames staging to
/models/<model>/<version>/, and reportsREADY. When ≥ 99.5% of hosts are ready (or a deadline hits), the coordinator emitsDistributionComplete{job, version, ready_hosts}to the deployer and keeps repairing stragglers in the background.
Deep dives
Topology selection. Total time ≈ S / B (bandwidth term) + depth × block / per-stripe rate (pipeline fill term) + repair tail.
| Topology | Origin reads | Per-host upload | Depth / fill time | Failure blast radius | Verdict |
|---|---|---|---|---|---|
| Everyone from origin | N × S | 0 | 1 | none | 7 h; kills the origin |
| One chain of 4,096 | 1 × S | 1 × S | 4,096 hops; one slow host gates all below it | everything downstream | Bandwidth-optimal, too fragile |
| k-ary tree (fanout 4) | 1 × S | 4 × S for interior nodes | log₄ N ≈ 6 | a subtree | Interior hosts are upload-bound: ~4x slower |
| BitTorrent-style swarm | ~1-2 × S | ~1 × S | dynamic | small | Robust, but chatty metadata and poor rack locality |
| Striped two-level chains | 1 × S | ≤ 1.125 × S | 160 hops × 3.2 ms ≈ 0.5 s | one rack segment per stripe | Pick this, with swarm-style repair for holes |
With 4 MB blocks and ~1.25 GB/s per stripe, each hop adds ~3.2 ms, so the 160-hop fill is negligible next to the 80 s bandwidth term. Rack-local chains keep 31/32 of the hops inside the top-of-rack switch; only the 128-hop leader chain crosses the spine, 8 stripes × 1.25 GB/s per rack.
Failure recovery and pipeline repair. Detection: each downstream expects a block at least every 2 s; a stall, a connection reset, or a report from the host's own health agent triggers ReportSuspect. Repair is a splice: the coordinator removes the dead host from each of its 8 chains, connects its upstream directly to its downstream, bumps topology_version, and pushes it. The downstream reconnects with StreamStripe(from_block = last contiguous block it holds). The upstream holds every block it already forwarded on disk, so it replays from disk without going back to the origin. One splice costs ~2-4 s for one rack segment of one stripe; other stripes keep flowing. If a leader dies, the next host in that rack takes over for that stripe. If a stripe root dies, its successor resumes the range GET from the last offset. A rejoining host re-registers its bitmap and is appended at the tail of each rack chain, so it never blocks anyone ahead of it.
Slow hosts and bandwidth saturation. In a chain, one host at 30% NIC speed throttles everything behind it, so slowness must count as a failure. The coordinator compares each host's rx_gbps to the stripe median; under 50% for 10 s moves the host to the tail of its rack chain (it becomes a leaf with no downstreams). Flow control is credit-based: a downstream grants 16 blocks (64 MB) of credit per stripe, which bounds memory to 8 × 64 MB per host and stops a fast upstream from overrunning a slow disk. For holes left by repairs or quarantined peers, an agent near the end switches to swarm mode: it asks the coordinator for 3 peers that hold the missing chunks (rarest-first, same rack preferred) and pulls with FetchChunk. This removes the last-1% tail that pure pipelines suffer from. Origin pacing: stripe roots share a token bucket held by the coordinator (cap / 8 each, rebalanced when a stripe falls behind), and S3 range GETs use 16 parallel 64 MB ranges per root to reach per-connection throughput.
Integrity. The manifest is produced by the writer at checkpoint time (the only place that saw the original bytes) and signed. Every receiver verifies each 64 MB chunk's SHA-256 before journaling it; CRC32C per 4 MB block catches wire errors early, before forwarding. On a mismatch the host re-fetches that chunk from a different peer. If two different downstreams report bad data from the same upstream, that host is quarantined (excluded from topologies, disk scrubbed) because it is likely a bad disk or bad memory. A host never loads from staging: the rename to the final path happens only after every chunk passes.
Coordinator availability. The data plane does not need the coordinator after assignment: chains keep streaming if it disappears, and only repairs pause. Run 3 coordinator replicas with leader election on an etcd lease (10 s TTL). Durable state (job, manifest reference, latest topology version) is in Postgres; the rest is rebuilt. A new leader loads the job and the latest topology, then waits one report interval (2 s) to rebuild the 6.4 MB bitmap from ReportProgress, which agents send to whichever replica holds the lease. Agents buffer reports and retry with backoff. Topology versions are monotonic and stored transactionally, so a split-brain old leader's assignments are rejected by agents as stale. Shard coordinators by job if many jobs run at once; one job's 2,000 RPC/s never needs more than one leader.
Trade-offs
- Striped chains over a swarm: near-optimal bandwidth and predictable origin load, at the cost of central topology management and a repair protocol. A swarm needs no coordinator for routing but wastes upload on duplicate chunks and is harder to keep rack-local.
- Cut-through 4 MB blocks with 64 MB hash units: low pipeline latency, but a host forwards bytes before the chunk hash is checked. A corrupt block can propagate one chunk deep; it is caught at every receiver and re-fetched, so correctness holds but some bandwidth is wasted.
- Bounded coordinator role: keeping it out of the data path makes it easy to make highly available, but repairs wait for leader failover (~10-15 s).
- Ready at 99.5%, not 100%: handoff is fast, but the deployer must route only to hosts listed as ready, and stragglers join later.
- No compression by default: bf16 weights compress only ~10-20% with byte-splitting plus zstd, which does not pay off on 100 Gbps in-DC links where CPU would become the bottleneck. It does pay off on cross-DC WAN links, so compress only there.
- Delta downloads: chunk-level content addressing means unchanged shards (frozen embeddings, a shared tokenizer, LoRA-only updates) are skipped for free. A full fine-tune changes almost every weight, so true binary diffs rarely help and are not worth their complexity.
- Multi-DC: one stripe-root set per data center reads from a regional replica, filled by a single cross-region copy over the WAN with compression. This gives up some freshness (the regional copy has to finish first, or at least run ahead) in exchange for keeping WAN cost at 1x per region.
- Keeping N-1 on disk: rollback is instant, but it costs 800 GB of NVMe per host. GC runs only after the deployer confirms the new version has served cleanly.
Interviewer expectations
Mid-level (L4)
- Sees that N hosts pulling from object storage is the problem, and proposes peer-to-peer or tree distribution.
- Does the basic math: per-host bandwidth gives the time floor, and origin reads should not scale with N.
- Includes per-chunk hashing and atomic placement so a partial checkpoint is never loaded.
Senior (L5)
- Compares chain, tree and swarm on origin load, per-host upload and depth, and picks one with numbers (fill time, ~80 s floor).
- Designs pipeline repair from local disk, resumable journals, and treating slow hosts as failures.
- Keeps the coordinator out of the data path and explains how it fails over without losing progress.
Staff (L6+)
- Plans for rack and spine topology, origin rate caps, interference with serving traffic, and credit-based flow control.
- Covers the tail: swarm-style hole repair, quarantining bad hosts, a ready threshold, and a clean handoff contract with the deployer.
- Extends to multiple data centers, version retention and GC, and where compression and dedup are and are not worth it.