Design a GPU Cluster Scheduler for ML Training

A GPU cluster scheduler for ML training looks superficially like "Kubernetes but for GPUs," and that surface reading is why most first attempts fail. The real problem is dispatching distributed training jobs that need N GPUs simultaneously, placed on the right network topology, preempting correctly under priority pressure, enforcing per-team quotas, and recovering from node failures mid-run without losing weeks of training progress. The generic Kubernetes default scheduler cannot do this, and that is the reason systems like Volcano, YuniKorn, Slurm, Kueue, and the internal schedulers at Meta, Google, Microsoft, and OpenAI exist as distinct pieces of infrastructure.

The canonical use cases pin down what must hold:

• Distributed training job — a PyTorch or JAX job asks for 64 GPUs. All 64 must start at the same time. If only 63 are ready, the job must not start; 63 running ranks idle while waiting for rank 0 wastes GPUs at thousands of dollars per hour.
• Priority preemption — a P0 oncall incident job arrives needing 32 GPUs. The cluster is full of P2 research jobs. The scheduler must pick a victim, signal a clean checkpoint, and start the P0 within minutes.
• Topology-aware placement — a 64-GPU job on eight hosts in one rack trains 3x faster than the same 64 GPUs spread across four racks, because AllReduce bandwidth is dominated by the slowest link in the ring.
• Spot/preemptible handling — cloud providers yank spot GPUs with short notice (AWS Spot 2 minutes; GCP preemptible and Azure Spot 30 seconds). The scheduler must checkpoint, release the node, allocate a replacement, and resume training with no human involvement.
• Fair share across teams — Research gets 40% of the cluster, Production gets 60%, and within Research each team gets a guaranteed floor but can burst into idle capacity.

Let us size for a realistic large-scale GPU cluster — think of a Tier-2 model training organization with a few thousand GPUs, roughly what mid-size AI labs or large enterprise ML teams operate.

Fleet size

• Total GPUs: 10,000. At 8 GPUs per host, that is 1,250 hosts.
• Network topology: 32 racks, roughly 40 hosts per rack. Within-rack InfiniBand at 200 Gbps. Cross-rack Ethernet at 40 Gbps. Within-host NVLink at 600 GB/s.
• Availability zones: 2 AZs for redundancy, but training jobs never span AZs (cross-AZ latency kills AllReduce).
• GPU types: mixed fleet — roughly 70% H100, 20% A100, 10% L40S. Node inventory tracks which GPUs are on which hosts.

Three APIs carry the weight: the Job Submission API (how users create and manage training jobs), the Cluster Query API (how users see queue state and node inventory), and the Internal Node Agent API (how hosts report status and receive commands). The submission API is the most important because it is the surface every user touches.

Job Submission API (REST, user-facing)

1POST /api/v1/jobs
2Content-Type: application/json
3Authorization: Bearer <user-token>
4
5{
6  "name": "llama-70b-pretrain-run-47",
7  "team": "research-foundation-models",
8  "owner": "[email protected]",
9  "image": "registry.example.com/training/llama-70b:2026-04-19",
10  "command": ["python", "-m", "train", "--config", "llama-70b.yaml"],
11  "resources": {
12    "gpu_count": 64,
13    "gpu_type": "H100",
14    "cpu_per_rank": 16,
15    "memory_per_rank_gb": 256,
16    "scratch_disk_gb": 2000
17  },
18  "topology_hint": "prefer_same_rack",
19  "priority": "P1",
20  "max_runtime_hours": 168,
21  "checkpoint": {
22    "path": "s3://my-checkpoints/llama-70b/run-47/",
23    "interval_minutes": 30,
24    "retain_last_n": 3
25  },
26  "restart_policy": {
27    "on_preemption": "resume_from_checkpoint",
28    "on_node_failure": "resume_from_checkpoint",
29    "max_restarts": 20
30  },
31  "env": {
32    "WANDB_PROJECT": "llama-70b",
33    "NCCL_DEBUG": "INFO"
34  }
35}

Three very different stores, each optimized for a different access pattern: etcd (cluster scheduling state — small, strongly consistent, low-latency reads), PostgreSQL (job history and quota ledger — medium scale, transactional, queryable), and Object Storage (checkpoints and logs — huge, append-heavy, cheap).

Cluster State Store — etcd

Technology choice: etcd (or FoundationDB for larger scale).

The system decomposes into three planes with very different performance and availability characteristics:

1. Control plane — API server, admission controller, queue manager, core scheduler, placement engine, preemption engine, quota enforcer. Runs on a small fleet of machines that have nothing to do with the GPUs they schedule. Owns the scheduling decision.
2. Data plane — the 1,250 GPU hosts themselves, each running a node agent that reports telemetry and executes launch/drain commands. The data plane never makes scheduling decisions; it just reports state and obeys.
3. State plane — etcd for scheduling state, PostgreSQL for history, object storage for checkpoints. The control plane reads/writes here; the data plane never touches these directly (it only talks to the control plane via gRPC).

Four flows carry the system's weight: gang-scheduled job submission and dispatch (the hot path), priority preemption (the adversarial case), coordinated checkpointing (the durability path), and queue backfill (the efficiency path).

Flow 1: Normal job submission and gang dispatch

The weight lives in four components: the Placement Engine (topology-aware scoring and bin-packing), the Preemption Engine (victim selection and grace orchestration), the Hierarchical Fair-Share / DRF module (quota and ordering), and the Fragmentation Detector (cluster-level utilization recovery).

1. Placement Engine — topology-aware scoring

The scheduler lives in a space shaped by four fundamental tradeoffs: linearizable state store vs throughput (etcd vs PostgreSQL), Kubernetes extension vs standalone scheduler, gang scheduling safety vs GPU fragmentation, and preemption vs queue wait.

etcd versus PostgreSQL for scheduler state

• etcd: linearizable, transactional multi-key writes, watch API, sub-15 ms commit. Ideal for scheduling hot path. Modest scale (GBs of state, 10K-50K writes/sec ceiling).
• PostgreSQL: rich query language, unlimited scale with sharding, mature ecosystem. Strong consistency via single primary but no native watch and higher commit latency (20-100 ms).
• Our choice: etcd for scheduling state (live jobs, allocations, queue, nodes), PostgreSQL for historical data (completed jobs, audit trail, quota ledger). This is the Kubernetes pattern and it wins because the two stores have complementary strengths.
• Why not both in PostgreSQL: LISTEN/NOTIFY is best-effort and session-bound; the watch API is essential for controllers to react to state changes in real-time. Also, PostgreSQL commit latency under multi-table transactions with index maintenance is 3-5x etcd's, which matters at 100+ decisions/sec.
• Why not both in etcd: etcd's query surface is "get key, get prefix, range scan" — fine for scheduling but terrible for "sum GPU-hours by team over last quarter." Forcing analytics into etcd would kill it.

Scheduler leader crashes

Scenario: the active scheduler leader process OOMs and restarts.

Response:

ML-predicted queue wait times

Current ETAs are computed from simple heuristics (average job runtime, current queue depth, recent dispatch rate). A learned model using historical job completion data, hour-of-day seasonality, team-specific patterns, and priority mix would produce much better estimates. Benefits: users make better submission timing decisions, on-call knows when capacity will actually free up. Challenge: training the model requires good labeled data (actual queue times vs predicted), and the model must retrain frequently as cluster composition shifts.

Topology-aware preemption