Schedule GPU Pods and Drain a Node

Problem Statement

You are tasked with building a deterministic GPU scheduler for a cluster manager. The cluster consists of several worker nodes, each with a fixed total number of GPUs. Workload requests are submitted as "pods", where each pod requests a specific integer number of GPUs that must be allocated entirely on a single node.

Implement the GPUClusterManager class with the following methods:

• __init__(nodes: dict[str, int]): Initializes the cluster with a mapping of node_id to its total GPU capacity.
• schedule_pod(pod_id: str, gpus_needed: int) -> str | None: Schedules a pod onto an eligible node using a Best-Fit policy. An eligible node must have at least gpus_needed available GPUs. Among all eligible nodes:
  1. Choose the node with the minimum remaining GPUs that can still accommodate the pod (exact fits preferred).
  2. If there is a tie, select the node with the lexicographically smallest node_id. If scheduled, deduct gpus_needed from that node's available GPUs, record the placement, and return the node_id. If no node can satisfy the request, return None (the cluster state remains unchanged).
• can_drain(node_id: str) -> bool: Performs a dry-run check to determine if node_id can be safely drained. Draining requires that all pods currently running on node_id can be reassigned to the other nodes in the cluster without evicting or moving workloads already running on those other nodes. To simulate the reassignment deterministically, pods on node_id are rescheduled one by one in descending order of GPU requirement (ties broken by pod_id lexicographically). Each pod is placed using the same Best-Fit policy described above on the remaining nodes. Return True if every pod on node_id can be placed, or False otherwise. This check must not mutate the cluster's actual state.

Example

manager = GPUClusterManager({"node-1": 8, "node-2": 8, "node-3": 4})

# Schedule pods
manager.schedule_pod("p1", 4)  # Returns "node-3" (exact fit: 4 remaining)
manager.schedule_pod("p2", 4)  # Returns "node-1" (tie between node-1 and node-2; node-1 is smaller)
manager.schedule_pod("p3", 6)  # Returns "node-2" (remaining 8 >= 6; node-1 only has 4)

# Node state:
# node-1: 4 used, 4 free (holds ["p2"])
# node-2: 6 used, 2 free (holds ["p3"])
# node-3: 4 used, 0 free (holds ["p1"])

manager.can_drain("node-1")
# "p2" needs 4 GPUs. Other nodes free: node-2 has 2, node-3 has 0.
# Cannot fit "p2". Returns False.

manager.schedule_pod("p4", 2)  # Returns "node-2" (exact fit: 2 remaining)
manager.can_drain("node-3")
# "p1" needs 4 GPUs. node-1 has 4 free, node-2 has 0 free.
# "p1" fits on node-1. Returns True.

Constraints

• 1 <= len(nodes) <= 500
• Total GPU capacity per node: 1 <= capacity <= 64
• 1 <= gpus_needed <= 64
• pod_id and node_id are non-empty alphanumeric strings.
• All pod_id values passed to schedule_pod are unique.