Design a Job Queue for Multi-Stage Jobs on GPU and CPU Clusters

Design a system that accepts jobs and places them in a queue for execution. Each job has multiple stages, and the stages differ in what they need: some require mostly GPU compute and others mostly CPU. The interviewer expected a design built on a queue and a database, with a GPU cluster in addition to a CPU cluster to run the stages. ```hint Schedule stages, not jobs A job's stages want different hardware. Consider what unit of work goes into a queue, and who decides when the next stage of a job is ready to run. ``` ```hint Workers vanish mid-stage Decide how the system notices that a worker running a stage has disappeared, and what makes running that stage a second time safe. ``` ### Constraints and Clarifications - Given: jobs have multiple stages; some stages are GPU-heavy and others CPU-heavy; the design is expected to use a queue, a database, a GPU cluster and a CPU cluster. - Not given: job volume, stage durations, the size of data passed between stages, GPU types, priorities and deadlines. Ask, or state assumptions. ### Clarifying Questions - Are a job's stages a fixed sequence, or a graph in which some stages can run in parallel? - How many jobs arrive per day, how many stages does a job have, and does a stage run for seconds, minutes or hours? - How large are the inputs and outputs passed between stages, and where are they stored? - Is there more than one kind of GPU, and does a stage need a whole machine, one GPU or a share of one? - Do jobs have priorities, deadlines or per-team quotas? - When a stage fails, should it be retried, or should the whole job fail? Can a long stage save checkpoints and resume? - Do users need to watch progress and cancel jobs? ### What a Strong Answer Covers - Clarifies stage structure, volumes and durations, and states assumptions - A data model for jobs, stages and attempts, with the database as the source of truth - Scheduling at the level of stages, with separate queues per resource type and a component that advances each job - At-least-once execution with leases, heartbeats, retries and stage outputs that are safe to produce twice - Keeping expensive GPUs busy: right-sized requests, priorities, quotas, preemption and autoscaling - Failure handling, data passing between stages, and monitoring of queue depth and waiting time ### Follow-up Questions - GPU stages are queued for hours while the CPU cluster sits idle. What do you look at, and what do you change? - A high-priority job arrives while every GPU is busy with long low-priority stages. What happens? - A worker finishes a stage but crashes before reporting success. What does the system do? - How would users rerun a job from a failed stage without recomputing the stages before it?