Design a Shared Distributed ML Training Platform

Scenario

Autonomous vehicle development requires training diverse machine learning workloads—from large-scale perception transformers to multi-agent reinforcement learning trajectory models. Multiple engineering teams (Perception, Prediction, Controls, Simulation) share a central compute cluster containing thousands of accelerator nodes (GPUs/TPUs).

Design a multi-tenant distributed training platform that allows engineers to launch, track, and recover distributed training runs reliably and efficiently.

Requirements

Functional Requirements

• Job Submission & Gang Scheduling: Support distributed jobs (e.g., PyTorch DDP, DeepSpeed) requiring multi-node coordination. All workers for a job must start concurrently (gang scheduling) with topology-aware placement (e.g., intra-node NVLink, same leaf-switch InfiniBand).
• Fair Sharing & Quotas: Manage compute capacity across teams with hard quotas, soft borrowing, and job preemption based on priority.
• Dataset & Code Versioning: Ensure training runs bind deterministically to immutable dataset snapshots and code commits.
• Checkpointing & Fault Tolerance: Automatically detect worker/node failures, isolate unhealthy nodes, and resume jobs from the latest valid checkpoint with minimal lost progress.
• Experiment Tracking: Ingest and expose real-time metrics (loss, throughput, GPU utilization), artifacts, and lineage.

Non-Functional Requirements

• Scale: Support a pool of 4,000+ GPUs across 500+ physical hosts; handle hundreds of concurrent training runs and queue throughput of ~10,000 job launches/day.
• Reliability & MTTD/MTTR: Detect node failures within 10 seconds; resume affected distributed jobs within 90 seconds.
• Efficiency: Keep average GPU allocation idle overhead below 5% between epoch checkpoints or job transitions.

What Your Design Should Cover

1. High-Level Architecture: Control plane components (API gateway, resource manager/scheduler, node agents, experiment tracking store).
2. Gang Scheduling & Quota Engine: Algorithm and data model for fair-share allocation, preemption policy, and network topology-aware worker placement.
3. Fault Detection & Recovery Lifecycle: Workflow from worker heartbeat failure to node drain, health check, reschedule, and checkpoint reload.
4. Data Ingestion & Checkpointing Pipeline: Shared storage layer architecture (e.g., distributed object store vs. high-speed parallel file system cache) to prevent I/O bottlenecks during distributed reads and synchronous checkpoint writes.