Kubernetes for Multi-Host Training and Inference: Workload Aware Sc... Eric Tune & Dominik Marcinski

Pre-training Workloads [0:39]

• Pre-training involves creating new models and represents the largest AI workloads, often spanning multiple hosts, racks, and NVLink domains.
• These workloads run for days or weeks and face significant challenges in quickly detecting and restarting after software hangs or hardware failures, which occur frequently due to the sheer number of devices involved.
• Recovery time directly impacts cost reduction and the value derived from expensive hardware.
• Modern frameworks like Torch Elastic, Monarch, and Jax allow applications to continue running with minimal interruption, even if a portion of the hardware fails, by enabling the application to adapt and later rejoin.

With so many devices being used, there can be hardware failures like every couple of days and the time that it takes to recover and restart is is is significant reduction in cost.

Post-Training and Reinforcement Learning [03:07]

• Post-training activities, especially reinforcement learning, are becoming a primary focus for improving AI models.
• Reinforcement learning requires backpropagation through models, making it a multi-host and potentially multi-domain operation.
• A key challenge is the simultaneous execution and close interaction of accelerator workloads (for inference) and CPU workloads (for decision-making and backpropagation), often on different hardware types.
• Frameworks like Ray introduce dynamicity, with dynamic pod creation and resizing, posing questions about the scheduler's role when external frameworks manage many decisions.

The other challenge we see is that people are writing a lot of these reinforcement learning programs using frameworks especially ray that like want to be able to dynamically create and resize pods and have a lot of dynamicity.

Inference Serving [04:24]

• Once models are trained, they are deployed for inference to generate revenue, and these large models are typically multi-host even for minimal serving footprints.
• The concept of a "replica" is redefined to mean multiple pods in a specific configuration, with scaling up involving the addition of these multi-pod replicas, a pattern not native to the current scheduler.
• Online inference demands require handling variable daily traffic, avoiding idle machines, and responding to spikes, creating a complexity of managing lower-priority work alongside multi-pod replicas during scale-up and scale-down.
• Batch inference, while tolerant of longer latency, may not always be sufficient to fill capacity troughs.

So here replica used to mean like pod in deployment or stateful set but here replica mean like n pods in a certain configuration.

Research Workloads [05:47]

• Research workloads are focused on generating new ideas for future production cycles, distinct from the production-focused pre-training, post-training, and inference.
• Researchers typically run smaller, experimental workloads, necessitating fair sharing of machines and potentially preemption to maximize resource utilization and manage fragmentation.

So you generally have a bunch of researchers so they need to like fair share between the machines.

Scheduling Implications and Proposed Solutions [06:12]

• The challenges observed across these use cases can be grouped into four categories: higher throughput, orchestrating aware of workload groups of pods and topology, node groups, and extensibility for customer control over fair sharing.
• The core idea is to integrate these capabilities into the Kubernetes scheduler, starting with foundational concepts like "pod groups" and "node groups."
• The Workload API, introduced in alpha for version 1.35, aims to model workloads, enabling the scheduler to support workload-level scheduling constraints like topology and all-or-nothing requirements.
• Workload awareness means that every component affecting pod scheduling needs to know that a pod is part of a workload, understanding that any change to a pod affects the entire workload.
• Workload scheduling involves the scheduler assessing the impact of changes and attempting to reschedule workloads, considering both scheduling constraints and the current workload state, potentially rescheduling only a part or the entire workload.

So like how do you deal with all this these different things we need to change at once?

Workload API and Scheduling Framework Changes [09:24]

• The Workload API allows modeling workloads, enabling the Kubernetes scheduler to support workload-level scheduling constraints like topology and all-or-nothing requirements.
• Workload scheduling is intended to be fully customizable via a plug-in framework and can be optionally handled by external schedulers seamlessly integrated with the core scheduler.
• Key changes needed include extending the Kubernetes scheduler to accept workload placement proposals atomically, coordinate binding in an all-or-nothing fashion, and trigger rescheduling upon actuation failures, considering the conditions that led to the failure.
• Workload scheduling involves scheduling multiple pods together in a single cycle, which can be complex and may require finding efficient algorithms to solve the problem in a reasonably short time.
• The scheduling framework is evolving to separate logical scheduling from actuation, with a new phase for scheduling multiple pods simultaneously, producing placement proposals similar to those from external schedulers.

So what is workload uh awareness and why we need that?

Workload Mutability and Future Considerations [22:43]

• For version 1.35, the Workload API is initially immutable, but dynamic discovery of gangs is supported for specific use cases like leader-worker sets.
• Future iterations will prioritize allowing mutations like adding gangs to workloads and defining mutable aspects of pod templates relevant to workloads.
• The discussion touches on how Horizontal Pod Autoscaler (HPA) and Vertical Pod Autoscaler (VPA) might interact with workloads, with possibilities for VPA to predict future resource needs for better scheduling.
• Even deployments that don't require gang scheduling can be marked as workloads, allowing for better visibility and aggregate usage reporting, addressing a long-standing missing concept in Kubernetes.

So there's a design doc which explains like we call API rev one which is in 1.35 and then we talk about API rev N which is like here's what we think it could be like in like a year a year and a half from now.

Cross-Cluster and External Scheduler Integration [34:29]

• The Workload API is specific to a single Kubernetes cluster's control plane and does not directly address multi-cluster orchestration, although external multi-cluster orchestrators could potentially use the workload resource.
• The vision for external schedulers like Ray, Torch, and Jax is for them to integrate with the Kubernetes scheduler, which would handle lower-level resource tracking, validation, and atomic acceptance of placements to avoid deadlocks and races.
• This collaboration between different-scaled schedulers allows for a unified administration of the cluster.

So ideally the the cube scheduleuler would provide at least the the lowest level of keeping track of uh resources allocation uh validating whether the the placement is uh is basically correct.