Playbook / AI system design / Design a distributed GPU training job scheduler with preemption and checkpointing

Design a distributed GPU training job scheduler with preemption and checkpointing

Expected question

"Design the scheduler for a shared cluster of ~10K GPUs serving thousands of researchers: jobs from 1 GPU to 4K-GPU multi-week pretraining runs, with priorities, fairness, preemption, and failure recovery."

Variant forms

Interviewers often ask the same design with different framing — recognize the archetype:

  • "Design a distributed job queue for 100K+ GPU training jobs with preemption and checkpointing."
  • "How do you gang-schedule a 512-GPU job without deadlocking the cluster on partial allocation?"
  • "Design topology-aware placement — why does NVLink vs. same-rack vs. cross-rack matter?"
  • "A 4K-GPU job has waited 6 hours while small jobs keep backfilling around it — fix it."
  • "Design checkpoint-interval policy for a training run with a known hardware failure rate."
  • "How do you keep GPU utilization high without starving low-priority teams, or letting them starve high-priority ones?"

Where this actually gets asked

Reported from an xAI backend/infra loop, 2025–26; adjacent to OpenAI/frontier-lab infra staples. Distinct from ../general-system-design/04 Distributed job scheduler / task queue: that entry is the generic CPU-oriented queue. GPUs add gang scheduling, topology-aware placement, preemption economics, and checkpoint-resume — where a generic-queue answer runs out of things to say.

Executive summary

30-second thesis

I'd treat GPU scheduling as fundamentally different from a generic queue on four axes — gang scheduling, topology, preemption cost, and failure rate — and make preemption a min-cost decision over priority/checkpoint-age/gang-size, never a boolean kill.

2-minute answer

A 512-GPU job needs all 512 simultaneously; partial allocation deadlocks the cluster, which kills naive work-stealing queues outright. So: reserve-and-drain for large gangs, with bounded backfill — small short jobs fill the accumulating hole only if their estimated runtime fits before the gang's start estimate. Unbounded backfill starves big jobs; no backfill craters utilization.

Placement is topology-aware: NVLink domain beats same-rack (InfiniBand) beats cross-rack, and placement quality alone can move training throughput 2–5x. Preemption isn't "kill and requeue" — you checkpoint, drain, then reschedule, and that cost has to enter the scheduling decision itself: victim selection minimizes total cost (lowest priority, most-recent checkpoint, smallest gang), on a grace period that checkpoints on signal before a hard kill deadline.

At 4K GPUs, hardware failures are the common case, not the exception — I'd quote the intuition behind the Young/Daly checkpoint-interval formula (checkpoint more often as failure rate rises, less often as checkpoints get more expensive) rather than treat checkpointing as a fixed cron. Fairness uses Dominant Resource Fairness with a preemptible opportunistic tier above quota, which is usually the single policy that does the most for utilization.

What I'd ask them: Elastic jobs (resume at a different world size) or must re-gang identically? Central scheduler acceptable at this scale, or does it need to be sharded? What's an acceptable lost-work bound after a failure?

Requirements

Functional

  • Accept jobs from 1 GPU to 4K+ GPUs with declared topology hints, priority, and checkpoint config.
  • Gang-schedule multi-GPU jobs; never partially allocate a gang.
  • Preempt lower-priority jobs to run higher-priority or larger-gang jobs; checkpoint before kill.
  • Detect node/job failure and auto-requeue from the last checkpoint.

Non-functional

  • High cluster utilization without starving large jobs or low-priority teams.
  • Placement-aware scheduling: prefer NVLink domain > same rack > cross-rack.
  • MTBF at 4K GPUs implies component failures every few hours (H) — recovery must be the common path, not an edge case.
  • Bounded lost work per failure via checkpoint-interval policy.

Core entities

  • Job: resource request (GPU count, topology hint), priority, checkpoint config, owning team.
  • Gang: the set of tasks that must be scheduled atomically for one job.
  • Reservation: an accumulating set of GPUs held toward a pending gang's start.
  • Checkpoint: a versioned, tiered snapshot (host RAM → NVMe → object store) tied to a job.
  • Node: GPU inventory, health state (ECC counters, NCCL probe results), cordon status.

API / interface

POST /v1/jobs
{ "gpus": 512, "topology_hint": "nvlink-domain", "priority": "P1",
  "checkpoint": {"interval_s": 900, "elastic": false} }
→ 202 { "job_id": "job_...", "state": "queued" }

GET /v1/jobs/{job_id}
→ 200 { "state": "running", "gang_size": 512, "last_checkpoint": "..." }

POST /v1/jobs/{job_id}/preempt
{ "reason": "higher_priority_gang" }
→ 202 { "checkpoint_triggered": true, "grace_period_s": 120 }

Staff+ callout: preemption is a first-class API call with a grace period, not a kill -9 the scheduler happens to also do.

Data Flow

Rendering architecture diagram…

High-level design

Rendering architecture diagram…

Central scheduler with optimistic concurrency (Omega-style) keeps placement globally coherent at this scale; shard by cluster cell only if a single leader becomes the bottleneck.

Deep dive 1: gang scheduling and backfill

Reserve-and-drain accumulates freed GPUs toward a large gang's reservation while it waits. Unbounded backfill of small jobs into the accumulating hole starves the big gang indefinitely; bounded backfill — only admit small jobs whose estimated runtime fits before the gang's start estimate — keeps utilization high without indefinite starvation. Add a hard start deadline with escalating preemption as the deadline nears, so a 4K-GPU job can't wait forever behind an endless stream of small jobs.

Deep dive 2: preemption policy

Victim selection minimizes total cost: prefer lowest priority, most-recent checkpoint (least lost work), smallest gang. Grace period is checkpoint-on-signal (SIGTERM → checkpoint → exit) with a hard kill deadline behind it. Preempted jobs get a requeue priority boost to prevent starvation/thrash, plus a preemption rate limit per job so one job isn't preempted repeatedly before it can make progress.

Deep dive 3: checkpoint interval and topology

Tiered/async checkpointing — snapshot to host memory in seconds so training resumes immediately, then trickle to NVMe and object store. Interval is an optimization problem: the Young/Daly intuition (checkpoint more often as failure rate rises, less often as checkpoints get more expensive) is worth quoting even without deriving the formula live. Elastic jobs can resume at a different world size; non-elastic jobs must re-gang identically. Topology placement (NVLink > rack > cross-rack) is scored the same way at schedule time, since a bad placement silently taxes throughput 2–5x with no error anywhere to point at.

Deep dive 4: fairness and failure handling

Dominant Resource Fairness across teams with weighted shares; quota is not the same as reservation — allow opportunistic burst above quota at preemptible priority, which is usually the single policy that lifts cluster utilization the most. On failure: health checks at the node agent (ECC error counters, idle-time NCCL probes), fail-fast the whole gang on mid-run failure, requeue from checkpoint, cordon the node. Distinguish job-fault from infra-fault for retry-budget accounting so a bad checkpoint doesn't get infinite infra-fault retries.

What's expected at each level

  • Mid-level: priority queue + basic preemption, no gang-scheduling awareness.
  • Senior: gang scheduling, checkpoint/resume, aware that topology matters.
  • Staff+: bounded backfill with a starvation bound, min-cost preemption victim selection, checkpoint-interval reasoning (even informal Young/Daly intuition).
  • Principal: opportunistic preemptible tier as a utilization lever, job-fault vs. infra-fault retry-budget distinction, and a story for testing the scheduler itself (simulation, shadow, chaos).

Follow-up questions to expect

  • "A 4K-GPU job has waited 6 hours while backfill keeps running." Reservation with a hard start deadline; backfill only jobs whose estimate fits; escalating preemption as the deadline nears.
  • "Checkpoints for a 400B model are 5TB — object store write takes 20 minutes." Tiered async (host-RAM snapshot unblocks training immediately), sharded parallel writes per rank, incremental/differential checkpoints for optimizer state.
  • "How do you test the scheduler itself?" Discrete-event simulation against production traces, shadow scheduling, chaos testing (kill nodes mid-gang).
← Design an enterprise LLM gateway (unified multi-provider platform)Design a unified query engine across dispersed data sources (email, calendar, docs, chat) →