GPU pod scheduler: find eligible nodes, then drain a node and reschedule its pods

A cluster has GPU nodes, each with a capacity of CPU cores, memory and GPUs, and some pods already running on it. Build the two scheduling steps a GPU cloud needs: placing a new pod, and draining a node so it can be taken out of service. Finish and test Part 1 before you move on.

Unless you agree on a different interface with the interviewer, assume:

• Node(name: str, cpu: int, memory_gb: int, gpu: int, pods: list[Pod]) and Pod(name: str, cpu: int, memory_gb: int, gpu: int); a node's free resources are its capacity minus its running pods. A pod is never split across nodes.

Part 1: nodes that can host a new pod

• eligible_nodes(nodes: list[Node], pod: Pod) -> list[str] returns the names of nodes whose free CPU, memory and GPUs all meet the pod's needs, sorted by name. A node with exactly enough free resources counts.

node  capacity (cpu, mem_gb, gpu)   running pods            free
a     (32, 256, 8)                  (8, 64, 4), (4, 32, 2)  (20, 160, 2)
b     (16, 128, 4)                  (8, 64, 2)              (8, 64, 2)
c     (64, 512, 8)                  none                    (64, 512, 8)
d     (16, 64, 2)                   none                    (16, 64, 2)

eligible_nodes(nodes, Pod("p", 8, 64, 2))   ->  ["a", "b", "c", "d"]   # b fits exactly
eligible_nodes(nodes, Pod("p", 4, 32, 3))   ->  ["c"]                  # only c has 3 free GPUs
eligible_nodes(nodes, Pod("p", 1, 1, 9))    ->  []                     # no node has 9 GPUs

Part 2: drain a node

• drain(nodes: list[Node], node_name: str) -> dict[str, str] | None removes the node and places every one of its pods on the remaining nodes, without moving any pod already running elsewhere. Returns pod name -> new node name with no node over its free CPU, memory or GPUs, or None if no such placement exists.

# n1 (16 GPUs) runs q1..q5 needing 2, 3, 3, 2 and 2 GPUs (each 2 CPUs, 16 GB);
# n2 and n3 each have 6 free GPUs and plenty of CPU and memory.
drain(nodes, "n1")  ->  {"q2": "n2", "q3": "n2",
                         "q1": "n3", "q4": "n3", "q5": "n3"}   # 3+3 and 2+2+2

Placing the pods in name order on the first node that fits fails: q1, q2 go to n2 (1 GPU left), q3, q4 to n3 (1 left), and q5 has nowhere to go. If n3 had only 5 free GPUs, the answer would be None (12 GPUs needed, 11 free).

Constraints

• Up to 10^4 nodes and 10^5 pods; all quantities are non-negative integers.
• A drained node runs at most about 20 pods.

Clarifying questions

• Should the eligible nodes be ranked (best fit first) or just filtered? Are there other filters, such as GPU type or nodes marked unschedulable?
• Is eligible_nodes called for every new pod, so free resources should be kept current?
• May existing pods move to make room during a drain? (No, here.)
• If not every pod fits, do we place as many as possible or refuse the whole drain?

What a strong answer covers

• Free resources computed once per node and kept current as pods start and stop, not recomputed per query; a filter on all three resources in O(nodes) with a deterministic order.
• Recognizing Part 2 as bin packing: greedy placement can fail when a valid placement exists, as the example shows.
• An exact search for the drain: backtracking with the largest pods first, pruning on total free resources and on nodes with identical free space, fine for a few dozen pods; first-fit or best-fit decreasing for large drains, and what it gives up.
• Applying the plan atomically: no node over capacity if the plan fails halfway, and pods created during the drain.
• Tests: exact fits, a pod bigger than any node, an empty cluster, infeasible drains, the example where greedy fails.

Follow-ups

• Schedule the pod: pick one node by best fit or by spreading, and say when each fits.
• Two schedulers place pods at the same time. How do you avoid putting both on the last free GPU?
• Some pods may move during a drain. How does the search change? Minimize the number of nodes that receive pods.