Bending the Kubernetes scheduler

Introduction

In the previous post we showed how we built a system that lets us scale anywhere, using the Virtual Kubelet project. But we skipped a big chunk of that system: how we schedule around the VK node to maximize cluster usage and minimize cost.

Some approaches we have seen modify the Kubernetes scheduler, by changing its configuration or by adding plugins as steps. Others replace it entirely with schedulers like Volcano, Apache YuniKorn, or KAI-Scheduler. None of these capture our business needs out of the box, so we took a different route: make the default Kubernetes scheduler work for us.

We could have written plugins or modified the default scheduler ourselves. During the project, though, we noticed we did not need to. With a bit of cleverness and a few tools built on top, the default scheduler could still do the job.

As we mentioned in the previous post, the whole idea is for this system to act as spill capacity. We keep as much as possible in the cluster, and we spill only once the cluster is full: either research has taken it over, or inference demand has genuinely exhausted capacity. We can break this system into three parts.

1. Scheduling
2. Descheduling
3. Admission and priorities

Scheduling

Before we explain our design, it helps to know how the Kubernetes scheduler works. The scheduler runs a chain of plugins. Think of them as steps: each one filters and scores the available nodes, and the node with the highest final score wins. If you want more detail, these two references cover it well:

1. Scheduling Framework
2. Scheduler Configuration

One of these plugins is NodeResourcesFit:

NodeResourcesFit: For pod-by-pod scheduling checks if the node has all the resources that the Pod is requesting. The score can use one of three strategies: LeastAllocated (default), MostAllocated and RequestedToCapacityRatio. For PodGroup scheduling calculates the resource utilization in the entire evaluated placement. The score uses the MostAllocated strategy. Extension points: preFilter, filter, score, placementScore.

It defaults to LeastAllocated, which means the scheduler tries to distribute load across the nodes of the cluster. Now recall how we set up our VK node: it advertises an extremely high amount of resources. This step therefore weights the VK node disproportionately, and pods land on it even when we have capacity in the cluster.

The first idea that comes to mind is a soft node anti-affinity:

...
  affinity:
    nodeAffinity:
      preferredDuringSchedulingIgnoredDuringExecution:
      - weight: 100
        preference:
          matchExpressions:
          - key: kubernetes.io/hostname
            operator: NotIn
            values:
            - my-vk-node
...

Unfortunately, this does not work as expected. Kubernetes still schedules pods on the VK node even when in-cluster capacity is available, because the NodeResourcesFit score dominates.

Another idea was to change the configuration to MostAllocated. For us, at least, this is not desirable either. There are still cases where we want bin-packing, but in most cases we prefer to spread workloads around: it gives us resilience, and some workloads, like heavy data processing, consume a lot of CPU and RAM. We could and should set accurate CPU and RAM requests for these applications, but their usage is hard to predict.

One way to keep pods off a node is a taint. So we built our design around a dynamic mechanism: detect GPU usage in the cluster, and add or remove a taint on the VK node accordingly.

By default, not all workloads can run on the VK node. Most need manual migration first, such as downloading models on the External Provider side or fixing quirks. For that reason, the VK node always carries a taint.

...
  taints:
  - effect: NoSchedule
    key: vk.node.com
    value: "true"
...

Only migrated workloads carry the matching toleration. That alone is not enough, though. If those workloads always tolerate the taint, we are back to square one.

To close the loop, the system watches for periods of low GPU usage, meaning plenty of free GPUs in the cluster, and adds a second taint that no pod tolerates.

...
  taints:
  - effect: NoSchedule
    key: vk.node.com
    value: "true"
  - effect: NoSchedule
    key: vk.node.com/low-gpu-usage
    value: "true"
...

For us this also runs inside the VK process as a background task. It queries Prometheus for the number of available GPUs.

Note: We could have checked the state of Kubernetes directly, but for our case it was easy to encode the business logic in a PromQL query. For example, our lowest priority is data processing jobs. Even when they use GPUs, we exclude them from the count, since training jobs or inference should always evict them.

Tuning this query took some work, because capturing the real state of the cluster is hard. Once set up for the intended use case, though, it represents the cluster state well and is inexpensive to evaluate, especially if you use a recording rule.

(
    sum(
        (1-kube_node_spec_unschedulable{node=~"<nodes in cluster prefix>-.*"}) 
        * on(node)
        kube_node_status_condition{node=~"<nodes in cluster prefix>-.*", condition="Ready", status="true"}
        * on(node)
        kube_node_status_allocatable{resource="nvidia_com_gpu", node=~"<nodes in cluster prefix>-.*"}
    )
    -
    sum(
        kube_pod_container_resource_requests{resource="nvidia_com_gpu", pod!~"<regex to exclude undesirable pods>", node=~"<nodes in cluster prefix>-.*"} 
        * on (pod, namespace, uid) 
        kube_pod_status_phase{phase="Running"}
    )
) >= 0 or vector(0)

This gives us the signal that tells us the state of the cluster. To avoid thrashing, we average it over time and stabilize it on top of that: we add or remove the taint only after the value stays above or below the configured threshold for a set duration.

Descheduling

Another important part of the system: how do we migrate workloads back into the cluster once GPUs free up? A naive idea would be to swap NoSchedule for NoExecute. But that taint evicts every pod on the node the moment it lands. If a big chunk of inference is running on the VK node, all of it gets disrupted at the same time.

Thankfully, there is a project built exactly for this purpose: Descheduler. It ships with several policies, and one fits our case well: RemovePodsViolatingNodeTaints.

This strategy makes sure that pods violating NoSchedule taints on nodes are removed. For example there is a pod “podA” with a toleration to tolerate a taint key=value:NoSchedule scheduled and running on the tainted node. If the node’s taint is subsequently updated/removed, taint is no longer satisfied by its pods’ tolerations and will be evicted.

- name: taintedNodes
  pluginConfig:
    - name: DefaultEvictor
      args:
        podProtections:
          defaultDisabled:
            - PodsWithLocalStorage
        minPodAge: 15m
    - name: RemovePodsViolatingNodeTaints
      args:
        includedTaints:
          - vk.node.com/low-gpu-usage
  plugins:
    balance:
      enabled: []
    deschedule:
      enabled:
        - RemovePodsViolatingNodeTaints

It lets us configure a minimum pod age to avoid thrashing, along with the descheduling frequency and the maximum number of pods to evict per cycle. Each removal counts as a regular eviction, so Pod Disruption Budgets are respected. No workload goes to zero replicas, or below whatever floor you configured.

Admission and priorities

For admission we use Kueue. It gives us gang scheduling, and our use of VK enables a nice trick. We run four queues:

• training: holds almost all the GPUs in the cluster. It can lend to and borrow from the other queues.
• inference: a small pool of GPUs for inference processes that cannot be offloaded to the VK node.
• low-priority: holds zero GPUs. It only borrows from other queues, and gets reclaimed when needed.
• inference-vk: a huge pool of GPUs. It cannot lend to other queues, and it cannot borrow from them.

Kueue already gives us workload prioritization, so we built our pod priority system around a single question: how many GPUs does the workload use?

• 1-gpu
• 2-gpu
• 4-gpu
• 8-gpu

1-gpu is the lowest priority and 8-gpu is the highest. This fixes GPU fragmentation. Say the cluster has 8 GPUs available in total, but split between 2 nodes. Without this priority system, an 8-GPU pod that wants to schedule would be stuck in Pending. With it, the 8-GPU pod evicts the pods on one of the nodes and schedules there, and the evicted pods migrate to the other node.

The way we like to think of it: Kueue does admission, and the scheduler does the reshuffling of pods and GPUs.

Limitations

The system works well, but it is not perfect. Sometimes we need to update the PromQL query to capture new logic. When the query and the real cluster state drift out of sync, the node can stay tainted while the cluster is actually full (a state the query does not capture), and pods sit in Pending. The stabilization delay cuts both ways too: it prevents thrashing, but it also means the system sometimes does not respond as quickly as we would like. Both have happened a few times. Neither has been a big problem.