Design a scheduler for ML training and batch jobs on a shared cluster

Questions first: how big is the cluster and how many job submissions a day? What is the mix of long training runs and short batch jobs? Do jobs need several machines at the same time (gang scheduling)? Are there per-team quotas, and may idle quota be borrowed? Can a running job be preempted, and do jobs checkpoint? Are there job dependencies, or is that a separate workflow layer? What queue-wait target matters?

Assumed scope: submit, queue, place, run, retry and cancel jobs; fair share with priorities; gang placement; preemption; failure recovery. Out of scope: the workflow or DAG engine, the training framework, and data storage.

Planning figures for the exercise, not facts about any real cluster.

• Fleet: 500 GPU machines x 8 GPUs = 4,000 GPUs, plus 2,000 CPU machines: 2,500 machines.
• Submissions: 20,000 jobs a day / 86,400 = 0.23 a second; at a 5x peak, about 1.2 a second. At 20 tasks a job that is 400,000 tasks a day, about 4.6 a second.
• State changes: about 5 persisted transitions per task (queued, placed, running, finished, plus a retry now and then) = 2 million a day, about 23 a second. Heartbeats are not persisted one by one.
• Heartbeats: 2,500 machines every 5 seconds = 500 a second, held in the scheduler's memory.
• Capacity: 4,000 GPUs x 24 hours = 96,000 GPU-hours a day. If demand is 60,000 GPU-hours, average utilization is 62.5%, so queues build at peaks and in fragmented moments, not all day.
• Placement work: a naive pass of 5,000 pending tasks against 2,500 machines is 12.5 million checks, a second or less on one core, and filtering by resource class makes it smaller.

What the numbers force: this is not a throughput problem. 23 durable writes a second sit far below the 5,000 a second one primary can commit, and one active scheduler can do the placement. The hard parts are correctness (never run a task twice by mistake), fairness, and recovering from failure.

1. API and job store. Submit takes a job spec (resources, priority class, team, image, retry policy, checkpoint flag) and a client-supplied job id, which makes a repeated submit harmless. Jobs and tasks live in a strongly consistent store: a small replicated cluster using a consensus protocol (three or five nodes), or a relational primary with a standby. Each task has a state machine: PENDING, ASSIGNED, RUNNING, SUCCEEDED, FAILED, PREEMPTED, CANCELLED, and every transition is a conditional update on the previous state.
2. Active scheduler. One leader is elected through a lease held in the consensus store, and each leadership term carries an epoch number. The leader keeps the cluster picture in memory (machine capacity, running tasks, last heartbeat) and rebuilds it from the store and agent reports after a failover. A warm standby makes failover a matter of seconds.
3. Queueing and fairness. One queue per team within a priority class. Each team has a guaranteed share and may borrow idle capacity from others; borrowed capacity can be reclaimed. Order pending work by priority class, then by how far a team is below its share, then by age so nothing starves. For jobs using several resource types, compare shares by the dominant resource.
4. Placement. Filter machines by resource class, label and locality, then pick by a score (best fit to limit fragmentation, spread across racks for fault tolerance). Gang jobs are all or nothing: reserve the full set of machines before starting any task. While a big gang waits, backfill smaller, shorter jobs into the gaps if they cannot delay the reservation.
5. Preemption. A high-priority task that cannot fit may evict lower-priority tasks. Choose victims that checkpoint, send a grace signal, wait a configured period, then kill. Preempted tasks go back to the queue and do not use up their retry budget.
6. Node agents. Each machine runs an agent that heartbeats every 5 seconds, starts containers, and reports status. Assignments carry the scheduler epoch and a task attempt number; an agent rejects an assignment from an older epoch.
7. Failure detection. A machine is suspect after three missed heartbeats (15 seconds) and its tasks are rescheduled only after the agent's lease has certainly expired. An agent that cannot renew its lease kills its own tasks first. That ordering stops a partitioned machine from running a zombie copy while the replacement runs too.
8. Retries. Classify failures: infrastructure (machine lost, preempted) retries freely with a cap; application errors (non-zero exit) count against the user's retry budget with backoff. Retries happen only in the scheduler, never also in the client and the agent, to avoid multiplied attempts.
9. Observability. Queue wait by team and priority, utilization per resource class, fragmentation (free GPUs that are not in a usable block), preemption count, scheduler loop time, and job outcome by failure class. Per-team cost reports come from the same data.

Hard problem one: a single writer that really is single. Leader lease plus epoch fencing means a deposed scheduler that wakes up late can issue commands, but agents and the store reject them.

Hard problem two: gangs without deadlock or starvation. Reservation with backfill and a bounded wait: if a gang has waited past a threshold, the scheduler stops placing small jobs on machines it needs.

• Centralized scheduler vs distributed (sharded) schedulers. One scheduler has a global view, which makes gang placement and fair share straightforward, and the numbers say it is fast enough. Pick it. Flip to sharded schedulers with optimistic placement when the cluster is tens of thousands of machines or the loop time no longer fits the latency target.
• Preemption vs waiting. Preemption keeps urgent work fast and utilization high but wastes progress when jobs do not checkpoint. Waiting is predictable but leaves urgent jobs stuck. Pick preemption for checkpointing jobs only, with a grace period. Flip to waiting with reserved capacity if most jobs cannot checkpoint.
• Strict quotas vs borrowing. Strict quotas are simple and predictable but leave GPUs idle. Borrowing lifts utilization but needs reclaim logic. Pick borrowing with reclaim. Flip to strict quotas for teams that need guaranteed capacity for a launch.
• Consensus store vs a plain relational primary. A consensus cluster survives a node loss with no failover step; a primary with a standby is easier to operate. At 23 writes a second either works. Pick the one your team already runs well, and keep the state machine and epoch rules the same.

• Designing for huge throughput. The submission rate is about one a second, and the real difficulty is correctness and fairness.
• Persisting every heartbeat to the database.
• Having one scheduler with no answer for what happens when it is wrong about being the leader.
• Forgetting gang scheduling, or doing it as "place tasks as machines free up", which deadlocks two big jobs that each hold half of what they need.
• Rescheduling a task the moment heartbeats stop, so a slow network produces two copies of a training run.
• Treating every failure alike, so infrastructure faults burn the user's retry budget.
• Ignoring fragmentation: 100 free GPUs scattered across machines cannot run an 8-GPU-per-machine job.