Design a GPU-aware pod scheduler
Company: Together AI
Role: Software Engineer
Category: System Design
Difficulty: hard
Interview Round: Technical Screen
Design an object-oriented, GPU-aware pod scheduler and cluster manager. Each Node has the shape {name: string, total_gpu: int, running_pods: Pod[]}. Each Pod has the shape {name: string, gpu_required: int}. Implement APIs: add_node(name, total_gpu), remove_node(name), add_pod(name, gpu_required), schedule_pod(pod_name) that assigns the pod to a node with enough free GPUs, remove_pod(pod_name), get_node_utilization(name), and list_nodes()/list_pods(). Specify data structures to support efficient lookups of nodes by available GPUs and pods by name. Describe and justify a placement strategy (e.g., best-fit or first-fit) and how you'd update indexes on every add/remove/schedule/evict operation. Discuss concurrency control (simultaneous adds/schedules), idempotency, and failure handling (e.g., removing a node that still has running pods, pod rescheduling on node removal). Provide time and space complexity for each API and write pseudocode for schedule_pod using your chosen strategy. Include edge cases like gpu_required > total_gpu on any node and fragmentation when multiple small pods occupy a large node.
Overview: This interview question evaluates requirements, scale assumptions, API/data design, architecture, trade-offs, failure modes, and rollout in a realistic interview setting. A strong answer for Design a GPU-aware pod scheduler states assumptions, handles edge cases, explains trade-offs, and shows how to validate the result clearly.
Community answers
Answer by dongreanay
I’d model this as a ClusterManager with nodes_by_name, pods_by_name, and an ordered free GPU index like free_gpu -> set(nodes). Each node tracks total_gpu, used_gpu, and running pod names. Each pod tracks gpu_required, status, and assigned node.
For placement, I’d use best-fit. Pick the node with the smallest available GPU count that still fits the pod. That avoids wasting large nodes on small pods and preserves capacity for bigger requests.
The main invariant is simple. Whenever a node’s free GPU count changes, remove it from its old bucket in the index, update the node and pod state, then reinsert it with the new free count. schedule_pod needs to be atomic because two schedulers should not be able to assign the same GPUs. For a simplified design, I’d use a cluster-level lock. In a real system, I’d use transactional updates or CAS/versioning on node state.
For remove_node, I’d either reject if it still has running pods or support drain mode. Mark the node unschedulable, evict its pods to pending, remove it from the index, and try to reschedule them. If rescheduling fails, those pods stay pending.
The main edge case is fragmentation. The cluster may have enough total free GPUs, but no single node has enough because pods cannot be split. Best-fit helps reduce that, but it does not eliminate it.