Design an ML Training Platform

About 300 researchers and ML engineers share one cluster of roughly 4,000 GPUs. Their work ranges from single-GPU notebooks and fine-tunes to pre-training runs that use 1,000 GPUs for weeks. Today people reserve machines over chat, half the cluster sits idle at night, a single bad node kills a week-long run, and nobody can say which data and code produced the model now in production.

Design the internal training platform: how jobs are submitted and scheduled onto GPUs fairly between teams, how training data reaches the GPUs fast enough, how long jobs survive hardware failures, and how experiments, metrics and resulting models are tracked and promoted.

In 45 minutes, size the cluster workload, sketch the APIs and components, and go deep on gang scheduling with quotas and preemption, and on checkpointing and recovery.

Sketch on the HLD board, write your answer in the Design doc, and size it with the Estimator. Start the mock interview from the timer for a 45-minute run.

Design Path · AI & ML Systems · problem 60 of 70

Ask before you design. Try answering each yourself, then compare.

1. What does the cluster look like?
   Typical answer

   512 nodes with 8 GPUs each, fast interconnect inside each pod of 32 nodes and slower links between pods. One region; cloud burst is out of scope.

2. What is the job mix?
   Typical answer

   Most jobs are small (8 GPUs or fewer), but around 60% of GPU-hours go to large jobs of 64 to 1,024 GPUs.

3. How should capacity be shared?
   Typical answer

   Each team has a guaranteed quota and may borrow idle GPUs; borrowed GPUs can be taken back with a short warning.

4. How often does hardware fail?
   Typical answer

   Assume a large job sees an interruption (bad GPU, NIC, host) about every 8 hours.

5. Where does training data live?
   Typical answer

   In object storage, several petabytes, mostly pre-sharded into large files.

6. Do users bring their own code and frameworks?
   Typical answer

   Yes, as container images. The platform provides the launcher, rendezvous, storage mounts and an SDK for metrics and checkpoints.

7. What does the model registry need?
   Typical answer

   Versioned models with lineage (code commit, data snapshot, config, metrics), stages such as staging and production, and approvals.

Requirements

Functional

• Submit, monitor and cancel training jobs that request N GPUs of a given type.
• Schedule jobs fairly by team quota with priorities, borrowing and preemption.
• Stream training data from object storage to GPUs.
• Checkpoint jobs and resume them automatically after failures or preemption.
• Track experiments: parameters, metrics, artifacts, lineage.
• Register, version and promote models for deployment.

Non-functional

• Cluster GPU allocation above 85%, useful GPU time above 75%.
• Scheduling decision for a waiting job within 10 s once capacity exists.
• A large job loses at most ~30 minutes of work per failure and resumes within 10 minutes.
• Control plane 99.9% available; its outage must not kill running jobs.
• Metrics visible in dashboards within 30 s.

Out of scope

• Online model serving.
• Building the training frameworks themselves.
• Cross-region or cloud burst scheduling.

Back-of-the-envelope estimates

• Cost of idling: 4,096 GPUs at an assumed $2 per GPU-hour is ~$8 k per hour, so every 1% of utilisation is worth ~$700 k a year. Utilisation is the metric the design optimises.
• Jobs: 300 users × ~4 submissions/day ≈ 1,200 jobs/day, so scheduling throughput is trivial; placement quality is what matters.
• Checkpoints: a 70 B-parameter model with optimizer state in mixed precision is ~16 bytes/param ≈ 1.1 TB. Writing that in 30 s needs ~37 GB/s, only possible if all ranks write their shard in parallel, first to local disk.
• Data: an image job on 1,024 GPUs at 500 samples/s/GPU × 150 KB ≈ 77 GB/s of reads; text tokens are far smaller. Input bandwidth for vision and video drives the need for node-local caches.
• Metrics: 1,000 running jobs × 100 series × one point per 10 s ≈ 10 k points/s: a modest time-series workload.

API

POST /v1/jobs {image, command, resources:{gpus:256, gpu_type:"H-class"}, queue:"team-nlp",
               priority:"normal", preemptible:true, max_runtime:"72h",
               checkpoint:{dir:"s3://ckpt/run-81", every:"30m"}}      -> {job_id}
GET  /v1/jobs/{id}            -> {state, placement, attempts, last_checkpoint_step}
POST /v1/jobs/{id}:cancel
POST /v1/runs/{run_id}/metrics  [{step, name, value}]   # batched by the SDK
POST /v1/models/{name}/versions {run_id, artifact_uri, metrics} ; PATCH .../versions/{v} {stage}

Jobs declare their gang size and whether they are preemptible up front, so the scheduler can place them all-or-nothing and knows exactly which work it may evict.

Data model

Jobs and queues (relational): jobs(job_id, queue, spec, state, priority, attempts, placement), queues(name, parent, guaranteed_gpus, max_gpus). Low volume, needs transactions for state transitions.

Node inventory: held in the scheduler’s memory, rebuilt from node-agent heartbeats; health history persisted for flaky-node detection.

Runs (relational): runs(run_id, job_id, commit, image_digest, dataset_snapshot, params). Metrics in a time-series or columnar store keyed by (run_id, metric, step).

Checkpoints and artifacts in object storage under /runs/{run}/step-{n}/shard-{rank} with a manifest written last, so a checkpoint only counts once its manifest exists.

Registry (relational): model_versions(model, version, run_id, artifact_uri, sha256, stage, approved_by), linking every version back to its run.

High-level design

Submit. A researcher uses the CLI or SDK to call the job API, which validates the spec against the team’s quota, stores it in the job database, and notifies the gang scheduler.

Schedule. The scheduler keeps a live view of every node from the node health service. It orders pending jobs by queue fairness and priority, finds a topology-aware placement for the whole gang, preempting borrowed capacity if a team is below its guarantee, and binds the job to nodes in the GPU pools. Node agents start the containers and run rendezvous so all ranks find each other.

Train. Each rank streams its deterministic slice of data shards from the training data store through a node-local dataset cache. Every ~30 minutes ranks write sharded checkpoints to local disk and upload them asynchronously to the checkpoint store. The SDK sends metrics and parameters to the experiment tracker, which writes them to the metrics store.

Recover. When a node fails, health reports it, the scheduler marks the job interrupted, swaps in a spare node and restarts all ranks from the newest complete manifest.

Promote. A good run is registered in the model registry, pointing at its weights; after approval the deploy pipeline picks up the version.

Open the HLD board and press Reference to see the architecture diagram.

Deep dives

Gang scheduling without fragmentation

Distributed jobs need every rank to start together, so placement is all-or-nothing. Large jobs should land inside one pod with whole nodes, because cross-pod links slow collective operations. Small jobs are best-fit packed onto nodes that are already partly used, which keeps whole nodes free for big gangs. A waiting 512-GPU job can starve forever while small jobs keep grabbing freed GPUs, so the scheduler makes a reservation: it picks the nodes that will free up soonest, stops placing new long jobs there, and backfills them only with jobs whose declared max runtime ends before the reservation starts.

Quotas, borrowing and preemption

Queues form a tree (org, team, project) with a guaranteed GPU count and a ceiling. Idle guaranteed capacity is lent to other queues, but borrowed jobs must be marked preemptible. When an owner reclaims capacity, the scheduler picks victims that lose the least: preemptible jobs on borrowed GPUs, ranked by GPUs × minutes since their last checkpoint. Victims get a signal and a grace period of about two minutes to save, then are requeued with their attempt count. Fair share between queues uses recent usage with decay, so a team that just ran a huge job yields to others for a while.

Checkpointing and fast recovery

Each rank copies its shard of model and optimizer state to host memory, which pauses training only for seconds, then writes to local NVMe and uploads in the background. A manifest naming all shards is written last, making the checkpoint atomic. On failure the job restarts from the newest manifest, not the newest files. The interval balances save cost against lost work: with a 30-second save and a failure every 8 hours, the optimum is about the square root of 2 × 30 s × 8 h, roughly 22 minutes, which loses about 5% of GPU time to saving and rework before counting restart time. A pool of warm spare nodes cuts restart time to a few minutes.

Feeding GPUs and reproducibility

Millions of small files kill throughput, so datasets are packed into shards of a few hundred megabytes. The sampler is deterministic: shard order comes from (seed, epoch), and each rank takes a fixed slice, so a restart resumes at exactly the same sample when the loader position is saved in the checkpoint. A node-local cache keeps shards for repeated epochs and lets jobs that share a dataset avoid rereading object storage. For lineage, the tracker records the image digest, code commit, dataset snapshot ID and full config for each run, and the registry links every model version to that run.

Failure modes

• If a GPU node fails mid-job, the job is restarted from the last complete manifest on a spare node; the node is drained and sent to diagnostics after repeated faults.
• If the scheduler or API is down, running jobs continue untouched because node agents own the processes; only new placements pause.
• If a checkpoint upload is interrupted, its manifest is never written, so recovery falls back to the previous complete checkpoint.
• If the object store throttles, node-local caches keep current epochs running and the scheduler stops starting new data-heavy jobs.
• If the metrics store is unavailable, the SDK buffers to local disk and replays later; training never blocks on tracking.

Core algorithm in code

import math

GPUS_PER_NODE = 8


def place(job_gpus, pods):
    """Gang placement. pods: {pod: [free GPUs per node]}. Returns [(pod, node, gpus)] or None (job waits).
    Small jobs best-fit onto partly used nodes; large jobs take whole free nodes, ideally in one pod."""
    if job_gpus < GPUS_PER_NODE:
        fits = [(free, pod, i) for pod, nodes in pods.items() for i, free in enumerate(nodes) if free >= job_gpus]
        if not fits:
            return None
        free, pod, i = min(fits)                       # tightest fit keeps whole nodes free
        pods[pod][i] -= job_gpus
        return [(pod, i, job_gpus)]
    need = math.ceil(job_gpus / GPUS_PER_NODE)
    whole = {pod: [i for i, f in enumerate(n) if f == GPUS_PER_NODE] for pod, n in pods.items()}
    single = [(len(ids), pod) for pod, ids in whole.items() if len(ids) >= need]
    if single:
        _, pod = min(single)                            # best-fit pod: least leftover
        chosen = [(pod, i) for i in whole[pod][:need]]
    else:
        chosen = []
        for pod in sorted(whole, key=lambda p: -len(whole[p])):   # span as few pods as possible
            chosen += [(pod, i) for i in whole[pod]]
        if len(chosen) < need:
            return None                                 # all-or-nothing: never start a partial gang
        chosen = chosen[:need]
    for pod, i in chosen:
        pods[pod][i] = 0
    return [(pod, i, GPUS_PER_NODE) for pod, i in chosen]


def checkpoint_interval(save_s, mtbf_h):
    """Young-style optimum: sqrt(2 * cost of one save * mean time between failures)."""
    return math.sqrt(2 * save_s * mtbf_h * 3600)


def wasted_fraction(interval_s, save_s, mtbf_h, restart_s=600):
    mtbf = mtbf_h * 3600
    return save_s / interval_s + (interval_s / 2 + restart_s) / mtbf


pods = {"podA": [8, 8, 8, 5, 2], "podB": [8, 8, 8, 8, 8, 8], "podC": [8, 8, 0, 0]}
for name, gpus in [("notebook", 2), ("finetune", 4), ("pretrain", 40), ("big", 64)]:
    print(f"{name:9} {gpus:3} GPUs ->", place(gpus, pods))
print("free after:", pods)

save_s, mtbf_h = 30, 8      # 30 s blocking save, job-level failure every ~8 h
best = checkpoint_interval(save_s, mtbf_h)
for iv in (300, best, 7200):
    print(f"interval {iv/60:5.1f} min -> {wasted_fraction(iv, save_s, mtbf_h):.1%} of GPU time lost")

Trade-offs

• Strict all-or-nothing gang placement avoids deadlocked half-started jobs but costs idle GPUs while a reservation drains.
• Lending idle quota raises utilisation but costs borrowed jobs occasional preemption and lost minutes of work.
• Frequent checkpoints cut rework after failures but cost save time and storage bandwidth.
• Topology-aware placement speeds large jobs but makes them wait longer for a pod with enough free nodes.
• Keeping spare nodes idle speeds recovery but costs a few percent of cluster capacity.

What a strong answer shows

• Recognises that distributed jobs need gang scheduling and addresses fragmentation and starvation.
• Designs quotas with borrowing and a concrete preemption policy.
• Sizes checkpoints with numbers and chooses an interval with a reason.
• Makes checkpoints atomic and recovery automatic.
• Keeps the data path fast and resumable.
• Links runs, data, code and model versions for lineage.

Mistakes to avoid

• Scheduling GPUs one at a time like generic containers, so large jobs never start.
• Treating checkpointing as the user’s problem and ignoring the time a 1 TB save takes.
• Static per-team partitions that leave the cluster half idle.
• Forgetting that preemption without a save signal throws away hours of work.
• Putting experiment tracking on the critical path of training steps.

Follow-ups to expect

• How would you add interactive notebooks without wasting GPUs? (Fractional or time-sliced GPUs with idle timeouts.)
• How would you support elastic jobs? (Let jobs declare min and max size; grow into idle GPUs and shrink before preemption.)
• How do you detect a slow GPU that drags a whole job? (Compare per-rank step times; drain stragglers.)
• How would you charge teams? (Meter GPU-hours by queue with lower rates for preemptible use.)

Practice next