Katkat
Back

Design a Fair GPU Training Job Platform

FluidstackSoftware Engineer · Onsite
EasySoftware EngineerOnsiteSystem Design

Last updated

Design a centralized scheduling and execution platform for one-off ML training jobs across a shared, heterogeneous GPU cluster, balancing fair resource sharing, multi-GPU network topology, and fault tolerance.

Scenario

Modern machine learning teams frequently submit ad-hoc, one-off training and fine-tuning jobs that demand varying GPU hardware profiles (e.g., NVIDIA A100 40GB/80GB, H100, L40S). You need to design an internal, multi-tenant GPU training platform that accepts user job specifications, schedules workloads fairly without starvation, enforces organizational quotas, and reliably executes distributed training runs across a finite physical fleet.

Requirements

Functional Requirements

  • Job Contracts: Allow users to submit job specifications including container image, entrypoint command, requested GPU type/count (single-GPU and multi-GPU), maximum runtime, and checkpoint storage paths.
  • Two-Tier Scheduling (Admission vs. Placement): Verify user/team quotas upon submission, queue admitted jobs, and bind them to concrete physical nodes according to topology constraints.
  • Multi-GPU Placement: Guarantee that multi-GPU jobs satisfy locality and interconnect requirements (e.g., co-locating 8-GPU jobs within a single node on NVLink, or multi-node jobs within the same InfiniBand rail/leaf switch).
  • Fair Sharing & Anti-Starvation: Implement scheduling logic (such as Dominant Resource Fairness or Weighted Fair Queuing) across multiple teams so that small jobs are not starved by large, long-running jobs.
  • Preemption & Recovery: Support graceful preemption of lower-priority jobs (signaling to save a checkpoint before SIGKILL) and automatic rescheduling on worker node failures.
  • Observability: Expose real-time fleet utilization, queue wait times, and per-tenant quota consumption.

Non-Functional Requirements

  • High Availability: The control plane should not lose queued or active job states during scheduler restarts.
  • Low Overhead: Scheduling decisions should be made within 1–2 seconds of capacity becoming free.
  • Idempotency: Prevent duplicate allocations if a worker re-registers or heartbeats are delayed.

Scale Assumptions

  • Fleet Size: 5,000 total GPUs across 750 physical nodes.
  • Tenants: 50 distinct teams, submitting a combined peak of 500 job submissions per hour.
  • Job Length: Ranges from 15 minutes (hyperparameter sweeps) to 48 hours (large fine-tuning runs).

What Your Design Should Cover

  1. High-Level Architecture: Control plane components (API service, State Store, Admission Controller, Scheduler, Node Agents).
  2. Data Model: Core entities for Jobs, Allocations, Nodes, and Tenant Quotas.
  3. Scheduling Pipeline: Separation between queuing/fair-share priority calculation and topological gang placement.
  4. Failure & Preemption Flow: Node heartbeat failures, GPU health checking (e.g., Xid events), and graceful shutdown protocols.

Hints

Hint 1

Decouple admission control (verifying team quotas and setting queue priority) from the topological placement engine (matching jobs to physical GPU nodes).

Hint 2

Distributed multi-GPU training requires gang scheduling; ensure that either all required GPUs with appropriate interconnects (e.g., NVLink) are allocated simultaneously or none are.

Hint 3

To prevent starvation caused by large reservations, look into priority aging, dynamic share deficits, or time-sliced preemption using checkpoint intervals.

Follow-up questions

  • How do you handle node health degradation (such as uncorrectable ECC memory errors or thermal throttling) during an ongoing multi-node training run?
  • How would you extend this scheduler to support elastic training workloads that can dynamically scale the number of GPUs up or down without failing the job?
  • What queue metrics would you expose to end users to give them accurate estimates of queue wait times before their job starts?