Design a GPU scheduling platform

Reference
Start your answer

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

Share a large GPU cluster among research and production teams: queue training and inference jobs, enforce quotas and priorities, place jobs on topology-aware hardware, and preempt fairly.

Requirements

We are building the scheduler and control plane for one shared GPU cluster: research and production teams submit training and inference jobs, the platform queues them under per-team quotas and priorities, places each job on GPUs whose network topology suits it, and preempts lower-priority work fairly when capacity is short.

Clarifying questions worth asking first: How big is the cluster and how uniform is the hardware? Are training jobs gang-scheduled (all ranks or nothing)? Do jobs checkpoint, and how often? Is inference long-running serving or batch? Can production ever be preempted? Is this one cluster or several regions? Assume one cluster per region, mostly 8-GPU nodes, and that training jobs can checkpoint.

Functional requirements

  • Submit a job: image, command, GPU count and type, priority class, team, max runtime, topology hints (e.g. "all ranks in one network block").
  • Gang scheduling: a multi-node job starts only when all of its GPUs are allocated at once.
  • Quotas per team: a guaranteed share of each GPU type, plus the right to borrow idle capacity that can be taken back.
  • Priority classes: prod-serving > prod-batch > research > preemptible. Higher classes may preempt lower ones.
  • Preemption with a grace period so the victim can checkpoint, then automatic re-queue and resume from the latest checkpoint.
  • Query status, logs, events and placement; cancel a job; see the team's queue position and quota use.
  • Detect failed GPUs and nodes, drain them, and restart affected jobs elsewhere.

Non-functional requirements

  • Never double-allocate a GPU. Correctness of allocations beats scheduling latency.
  • Utilization: over 90% of healthy GPUs allocated, and high actual SM use (not just allocated).
  • Scheduling latency: small jobs (≤ 8 GPUs) placed in under 5 s when capacity exists; large gangs can wait for capacity but the decision loop itself stays under 1 s.
  • Control-plane availability 99.95%. If the scheduler is down, running jobs keep running; only new placements stop.
  • Fairness: no team starves; a borrowed GPU is reclaimed within minutes, not hours.
  • Auditable: every allocation, preemption and quota change has a record of who and why.

Back-of-the-envelope estimation

  • 24,000 GPUs = 3,000 nodes × 8 GPUs. Nodes grouped into network blocks of 32 nodes (256 GPUs) under one leaf switch; blocks joined by a spine with about 4:1 less bandwidth than inside a block.
  • At ~$2/GPU-hour the fleet costs ~$1.2M/day; each 1% of idle GPUs costs ~$12K/day. That is why utilization and fast reclaim matter.
  • Jobs: ~40,000 submissions/day, ~0.5/s average, bursts of 50/s when sweeps launch. 90% use ≤ 8 GPUs; a few dozen use 1,024–8,192 GPUs and hold most of the GPU-hours.
  • Node heartbeats: 3,000 nodes every 5 s = 600 writes/s, kept in memory, not in the durable store.
  • Metrics: 24,000 GPUs × 40 counters every 10 s ≈ 100K samples/s into a time-series store.
  • Durable state is small: ~1M job rows per month at ~2 KB = 2 GB/month. The allocation table is at most 24,000 live rows.
  • Checkpoints: a 70B-parameter job with optimizer state is ~1 TB. At 200 GB/s aggregate storage bandwidth that is about 5 s to write, which sets the grace period (60 s gives slack).
API design

Public REST (gRPC internally); all writes take an idempotency key.

POST /v1/jobs   Idempotency-Key: <uuid>
{ "team": "pretrain", "priority_class": "research", "image": "registry/llm:2026-10-01",
  "command": ["torchrun", "train.py"], "gpu_type": "h100", "replicas": 128, "gpus_per_replica": 8,
  "topology": { "pack": "block", "max_blocks": 4 }, "max_runtime": "72h",
  "checkpoint": { "preemptible": true, "grace_seconds": 60 } }
→ 202 { "job_id": "j_7Kq", "state": "QUEUED", "queue_position": 3 }

GET /v1/jobs/j_7Kq
→ { "state": "RUNNING", "attempt": 2, "allocation": { "nodes": ["n-0412", ...], "blocks": ["b17","b18"] },
    "started_at": "...", "last_checkpoint": "s3://ckpt/j_7Kq/step-41000", "preemptions": 1 }

DELETE /v1/jobs/j_7Kq                → 202 { "state": "CANCELLING" }
GET  /v1/jobs?team=pretrain&state=QUEUED&cursor=...
GET  /v1/jobs/j_7Kq/events           → scheduled, preempted, node_failed, restarted ...
GET  /v1/teams/pretrain/quota        → { "h100": { "guaranteed": 2048, "in_use": 2600, "borrowed": 552 } }
PUT  /v1/teams/pretrain/quota        (admin) { "h100": { "guaranteed": 3072 } }
POST /v1/nodes/n-0412:drain          (admin) { "reason": "xid 79" }

Scheduler to node agent (gRPC stream, opened by the agent):

Agent → Scheduler: Heartbeat { node_id, gpu_health[8], running: [{lease_id, job_id, epoch}], free_disk }
Scheduler → Agent: StartTask { lease_id, job_id, attempt, epoch, rank, image, env, gpu_indices }
                   StopTask  { lease_id, epoch, grace_seconds, reason }   // preempt or cancel
Agent → Scheduler: TaskStatus { lease_id, epoch, state, exit_code, checkpoint_uri }

The agent ignores any command whose epoch is lower than the newest one it has seen for that lease.

Data model

The source of truth is a strongly consistent, transactional store: Postgres with a hot standby at this size (or FoundationDB/Spanner at multi-cluster scale). The write rate is low (tens per second) and the cost of a wrong allocation is high, so we need transactions and unique constraints, not throughput.

Table Key columns Indexes / notes
jobs job_id PK, team, priority_class, state, spec JSON, attempt, submitted_at, last_checkpoint (state, priority_class, submitted_at) for the queue; (team, state)
allocations (node_id, gpu_index) PK, job_id, lease_id, epoch, granted_at the PK is the "no double allocation" guarantee
leases lease_id PK, job_id, node_id, epoch, expires_at one per (job attempt, node)
quotas (team, gpu_type) PK, guaranteed, max_borrow changed by admins, audited
nodes node_id PK, block_id, gpu_type, state (ready/cordoned/draining/down), labels (block_id, state)
job_events (job_id, ts) append-only audit trail

Access patterns: the scheduler reads all queued jobs plus free capacity every cycle (kept in memory, rebuilt from the store on startup); users read one job or a team's list; agents report status by lease. Core invariants: a GPU belongs to at most one live lease; a gang's allocations are all written in one transaction or none are; a team's guaranteed usage never exceeds its quota. Heartbeats and GPU metrics go to memory and a time-series store (Prometheus/VictoriaMetrics), never to Postgres.

High-level design
flowchart LR
  U[Users and CI] --> API[API server]
  API --> DB[(Postgres state store)]
  API --> Q[Admission and quota check]
  Q --> DB
  S[Scheduler leader] --> DB
  S <--> A[Node agents]
  A --> N[GPU nodes]
  A --> M[(Metrics store)]
  H[Health monitor] --> M
  H --> S
  N --> C[(Checkpoint storage)]
  R[Image cache] --> N
  1. Submit, schedule, run. The API checks the request against the team's quota and limits, writes the job as QUEUED, and returns 202. The scheduler leader runs a cycle every ~1 s (and on events): it sorts the queue by priority class, then by the team's fair-share deficit, then by age. For each job it finds a placement in its in-memory capacity map that meets the topology hint. It commits all allocation rows and leases for the gang in one transaction with a new epoch. It then sends StartTask to each agent; agents pull the image, bind the listed GPUs, start the rank, and report RUNNING. The job is RUNNING when all ranks report.
  2. Preempt, checkpoint, re-queue. When a higher-priority job can't fit, the scheduler picks victims (below), marks them PREEMPTING in a transaction, and sends StopTask with a 60 s grace. The job traps SIGTERM, writes a checkpoint, and exits. The agent reports checkpoint_uri; the scheduler frees the GPUs, bumps the victim's attempt, and puts it back in the queue keeping its original submit time. On restart it resumes from last_checkpoint. After the grace period the agent sends SIGKILL anyway.
  3. Failure. The health monitor sees a dead heartbeat or a fatal GPU error, cordons the node, and tells the scheduler. The scheduler fails the whole gang attempt, frees all of its GPUs, and re-queues the job at the front of its class.
Deep dives

Scheduling atomicity and the lease protocol. One scheduler leader per cluster, elected through a lease in etcd or Postgres advisory locks; each new leader gets a higher fencing epoch. Placement is computed in memory, but it takes effect only when one transaction inserts every allocations row for the gang. The (node_id, gpu_index) primary key makes a double allocation impossible even with a bug or a stale leader. The transaction also checks leader_epoch = current, so a deposed leader's writes fail. Agents enforce the same rule locally: they refuse a StartTask with an epoch lower than one they've already seen, and they report what they're running on every heartbeat. The scheduler reconciles that report against the table: a task with no lease is killed, and a lease with no task after 2 minutes is released. A gang is not RUNNING until all ranks start; if any rank fails to start within 10 minutes, the whole attempt is released, so half-started gangs can't hold GPUs. Leases are renewed by heartbeats. If the scheduler is down, agents keep running their tasks; an agent kills a task only on an explicit stop or when its lease is lost for more than 15 minutes.

Topology-aware gang placement. Ring and tree all-reduce run at the speed of the slowest link, so a 1,024-GPU job split across five blocks can run 20–30% slower than the same job in four. Placement rules: never split a node across jobs once a job needs ≥ 8 GPUs; fill whole blocks first; choose the blocks with the fewest free GPUs that still fit (best fit) so other blocks stay empty for the next big job. Small jobs go to blocks that are already partly used, which keeps whole blocks free. A big gang that can't fit waits; to stop small jobs from taking freed GPUs forever, the scheduler reserves the blocks it plans to use for the head-of-line large job. While a block is reserved, only backfill jobs may run there, and only if their max_runtime ends before the reservation's expected start.

Strategy Utilization Large-job wait Notes
Greedy first fit High short term Can starve forever Fragments blocks
Strict FIFO Low Predictable Head-of-line job idles the cluster
Reservation + backfill High Bounded Pick this; needs honest runtime limits, so kill jobs that exceed max_runtime
Defragment by preemption Highest Lowest Costs checkpoint and restart time; use sparingly

Priority preemption and starvation prevention. Each team has a guaranteed share per GPU type; usage beyond it is "borrowed" and always preemptible. Victim selection for a job that needs G GPUs: (1) only lower priority classes, or borrowed capacity of any team in the same class; (2) prefer victims in the target blocks so the freed GPUs are contiguous; (3) prefer the fewest jobs, then the most recent checkpoint (least lost work), then the team furthest over its share. Production serving is never preempted; it only gets capacity from borrowed pools. To prevent thrashing and starvation: a job is not preempted again within 30 minutes of starting; each job has a preemption budget (e.g. 3 per day), after which it can only be displaced from capacity its team borrowed; queued jobs gain priority with age inside their class (+1 step per hour of waiting); and a team below its guarantee is always scheduled before a team above it. Between teams in the same class, order by dominant-resource fair share (GPU-hours used ÷ guarantee over a decaying 7-day window).

Fault detection and job recovery. At this scale something fails every few hours: GPU Xid errors, ECC errors, NVLink faults, IB link flaps, hung NCCL collectives. Signals: a missed heartbeat for 30 s marks the node suspect, and at 90 s it is down; DCGM fatal Xids cordon the node at once; a job whose step counter stops for 10 minutes while its GPUs sit at 0% use is treated as a hang. One bad GPU in an 8,192-GPU gang fails the whole gang, so the recovery steps are: kill all ranks, release the allocation, cordon the bad node, and re-queue the job ahead of others in its class. Keep a hot spare pool of ~1–2% of nodes per large reservation so a restart doesn't wait for capacity. Cordoned nodes run burn-in tests (NCCL all-reduce, a matmul stress test) before returning to service. With a 2-hour checkpoint interval, average lost work per failure is about 1 hour × job size, so big jobs should checkpoint every 15–30 minutes (async checkpoint to local NVMe, then upload in the background).

Admission control and quotas. Admission rejects at submit what can never run: a GPU count larger than the cluster or than the team's guarantee plus borrow limit, an image that isn't in the allowed registry, a missing max runtime. Quotas are checked again at schedule time, not only at submit time, because usage changes. Per-team limits on queued jobs (e.g. 5,000) stop a sweep from flooding the queue and slowing every cycle. The scheduling loop cost grows with queue length times candidate blocks, so cap how far it looks each cycle (the top 2,000 jobs by priority) and index free capacity by block.

Other considerations. Image cold starts: LLM images are 10–30 GB; 1,024 nodes pulling at once overwhelms the registry. Use a per-block pull-through cache, lazy-loading image formats (stargz/SOCI-style) so containers start before the full pull, and pre-pull the images of reserved jobs onto their blocks while they wait. Kubernetes path: in practice this is Kubernetes with a batch scheduler (Kueue/Volcano-style queues and gang admission), the NVIDIA device plugin, and topology labels per node. The concepts above map onto it: quota = ClusterQueue, lease = pod binding, agent = kubelet. Its limits at this scale are etcd write load and scheduler throughput for big gangs, so many labs run a custom gang scheduler on top. Observability: track allocated GPU % and actual use (SM active %, tensor-core use, from DCGM) per team, queue wait p50/p99 by class, GPU-hours lost to preemption and failures, and time from failure to restart. SLOs: 90% allocation, p95 wait under 10 minutes for ≤ 8 GPU jobs, under 2 minutes of scheduler-caused idle per freed node.

Trade-offs
  • A single scheduler leader gives a global view and simple correctness, but it limits decision throughput; sharding by GPU type or pool scales it and gives up placing one job across shards.
  • Gang scheduling with reservations stops large jobs from starving but leaves reserved GPUs idle when no backfill job fits; that is the price of predictable start times.
  • Borrowing idle quota keeps utilization high but makes borrowed work preemptible, so research jobs see restarts and must checkpoint.
  • Strict topology packing improves training speed but increases wait time; letting users set max_blocks lets them choose between starting sooner and running faster.
  • Preemption with a grace period saves work but delays the high-priority job by up to the grace time; a hard kill is faster and wastes everything since the last checkpoint.
  • Postgres as the source of truth is simple and transactional but needs care with failover (fencing epochs); etcd or Spanner remove some of that work at higher cost.
  • Killing jobs that exceed max_runtime makes backfill reliable but annoys users who guess low; give a short extension request flow.
Interviewer expectations

Mid-level (L4)

  • Clear job lifecycle (queued, running, preempting, done) and a working submit, status and cancel API.
  • A scheduler that matches jobs to free GPUs and a node agent that runs them.
  • Notices that multi-node training jobs need all GPUs at once.

Senior (L5)

  • Gang scheduling done atomically, with a constraint that makes double allocation impossible.
  • Quotas with guaranteed vs borrowed capacity, priority classes, and preemption with checkpoint and re-queue.
  • Topology-aware placement and a plan for failed nodes that restarts the whole gang.

Staff (L6+)

  • Leader fencing and reconciliation between the store and agents; running jobs survive a control-plane outage.
  • Reservation plus backfill to stop large-job starvation, anti-thrashing rules, and fair-share across teams.
  • Ties design choices to cost (idle GPU-hours, lost work per failure, checkpoint interval) and knows how it maps onto Kubernetes and where that breaks.