WG-Batch Updates: What’s New and What Is Next? - Yuki Iwai, CyberAgent, Inc. & Kevin Hannon, Red Hat

This talk provides an overview of the Kubernetes batch working group's progress and recent developments aimed at optimizing batch workloads, including HPC, AI, and ML tasks. The speakers, Yuki from CyberAgent and Kevin from Red Hat, discuss key enhancements such as the introduction of pod replacement policies and mutable pod resources specifically for suspended jobs. They emphasize the importance of queue management and scheduling for maximizing cluster utilization, highlighting new features like job sets designed to ease the management of complex distributed workloads. Additionally, they explore various scheduling algorithms tailored for AI training and the need for workload-aware scheduling APIs that simplify the handling of group-based scheduling. The session concludes with a demonstration of upcoming features that aim to seamlessly integrate gang scheduling into Kubernetes.

Thank you for joining me our batch working group update presentation. So, I'm Yuki, a software engineer working for CyberAgent and chair of Kubernetes batch working group. And thank you for training and yeah, some work working groups. And Hi everyone, I'm Kevin. Sorry for that. I work at Red Hat as a senior software engineer focusing on trying to improve HPC, AI, ML workloads on Kubernetes. I'm also a

chair for working group batch. Working group batch was the first group I joined the Kubernetes organization in and I haven't left. All right. Um Okay, let's get started. Um so, the first I before we talk about what's new in Kubernetes what's the new in Kubernetes batch working group, let me introduce what's batch working group. So, the Kubernetes batch working group is which gathers people to discuss enhancements

to better support batch By batch workloads, we mean HPC, AI, ML, data analytics, CI, or any other type of job which performs some computation and completes in contract to serving workloads which are meant to run indefinitely. The goal of the group is to reduce fragmentation in the Kubernetes ecosystem. We have people with expertise in different areas. The SIGs coming to the meeting include the first is SIG

SIG scheduling and SIG apps and SIG now and SIG auto scaling. so, the our scope is here. So, the first of all we work on the batch APIs, job API, and job set, and the many other ecosystem APIs. Second, we work on primitives to allow local variable queuing and scheduling. For example, queue. And we also work on queuing to maximize the cluster utilization and to improve batch

worker performance. Finally, we want to support our specialized hardware like GPUs and TPUs to collaborate with device management Okay. Uh let's dive to the job update. Okay. So, the actual So, we have four graduation update. The first is pod replacement policy GA from 1.34. Pod replacement policy is a knob to specify if a job should create replacement pods when the previous pods are terminating. It has two

strategies. First is terminating or failed. This is default strategy. And the uh second strategy is failed. And terminating or failed strategy is the same as previous behavior which allows the job controller to create replace pods even if the previous pods are still terminating. And the second strategy failed is a new strategy which prevents the job controller from creating replacement when the previous pods still uh exist regardless

of state active or And the second is mutable pod resources for suspended jobs. Uh mutable pod resources for suspended jobs allows us to change the job's pod spec resource requests limits parameters uh when a job is suspended. This doesn't support the aware resource claim and uh pod resource pod label resources. And uh third one Oh. And the third one is uh mutable scheduling directives for suspended Mutable

scheduling directives for suspended jobs gives us capabilities to automatically automatically reset the job's status start time when suspended. Previously, the external controller uh like queue had to manually reset the status start time when they want to resume the suspended job. This feature allows external controller to resume job by only one API call because suspend field uh exists in job spec and uh start time exists in status

uh field. In that case, uh external um controller had to call two type of two type of API calls. Uh which first is for spec, and the second is status. But, after we enable these functionalities, uh we can call only one API. And uh final uh feature is enable workload with Um actually, this enable workload with job is an enhancement uh collaboration with six scheduling and West

End Tas uh project. Um this allows us to automatically create workload, not Q worker, uh Kubernetes workload, and uh pod group resources for job. But, currently, uh we only uh support um non-elastic uh index jobs. And so Kevin uh will introduce more details after that. Yeah, thank you Yuki. So today I want to talk to you a little bit about what we've been up to in job

set to give an overview. I have a quick slide on what a job set is. Generally AI HPC ML workloads have the need for multiple templates. You can't really represent a workload with a single template when you're doing AI training. So for job set one of our goals was to streamline a lot of the implementation details that is required to run distributed workloads. So we have the

concept of basically it's a batch job with replicas, a replicated job. That creates pods under the hood cuz of the job controller and then we use a headless service. So all your pods are able to communicate together. One of the goals with these workloads is to always make sure we hide this complexity from the user as much as possible because at the end of day no one

really wants to care about a headless service or even the pods getting created. They want to interact with the job set and they'll assume everything will work. this project was introduced a few years ago I'm just going to talk more about what we've been up to in the last year. So the next the main thing that we are we work on is the concept of depends on.

Okay, that's good for you all. This is a concept of allowing an execution ordering for job set. This comes it's a very popular ask from the cube flow community to say I want a job pulling down a data set and then pull down a model and then run my training and we found that it was a kind of necessary to build this in in a controller to

kind of handle some of this complexity and then the depends on we had to add some status checks to make sure that we only start a certain workload when certain gates are satisfied. another feature that we've been really excited to introduce is the concept of it's kind of like bringing the stateful set to batch jobs where we have this idea of adding uh volume claim policies at

the job set level so that you can instantiate volume claims at the the individual jobs. And then one of the things why this is big in the Kubeflow community is they want the ability to have volumes that outlive the job. But some of them maybe be temporary like a lot of HPC style workloads have like a a temporary file system for running jobs, but once it's done

they don't care anymore. But in other cases for like checkpoints, you may want to write the checkpoints to disk so in case you want to recover from a job, you can use that. And one of the goals with job set, it's a critical component for the Kubeflow train job. And the idea there is that we want to make it as easy as possible for data scientists and

Python users to run Uh one feature we've worked on is uh failure policy. This is surprisingly complicated. I did not realize how hard it is to honestly run workloads on GPU nodes. There's always a lot of failures. There's also a lot of different kinds of failures and a lot of you like some are transient, some are not, and you kind of have to have this ability to

tune when to do retries versus when to not. So one of the ways we with failure policy, there's a feature in Kubernetes job called a pod failure policy. This is one of the first features the working group batch introduced like years ago. And the idea was that we can toggle how a job controller restarts jobs based off certain conditions or exit codes. We found that if job

set is using this, we still need to match on some of the regular expressions that come from the events or sorry from come from I think it's coming from the conditions like massing matching on the messages cuz they can be different. So, the idea was to kind of have some toggle-able way to restart job sets in certain cases. My favorite part about working group batch has always

been that our goal was always to contribute directly to core Kubernetes to improve the use case for all batch workloads. Uh I have I'm not going to go in much more details on the job set integration of this, cuz it's real a really quite involved uh feature. We have in a concept of an in-place restart for job set, because generally when you are a a large provider

like Google, uh you're running like you can imagine running like a 10,000 node workload, where you can't really afford all the workloads get recreated, cuz you have one failure. So, what we wanted to do was trigger jobs to either like triggering a job that's failing to reschedule itself, but not do all the jobs together. And we had to actually add a feature into the pod API called

restarts in place, that allows like a knit containers to kind of orchestrate the restarting of all the containers in a pod. So, that was that was how to and actually have this idea of an in-place job restart. I have all these QR codes I link to uh some features. If you are interested, please check that out. And with that, I want to hand it off to Okay,

thank you, Kenji. Um next uh I want to talk about Q. Okay. Uh so, let me introduce uh what is Q? Uh Q is a AI workable uh scheduler. Uh it is not custom Uh it is a second level scheduler in front of a Q scheduler. And uh don't try to replace uh core Q scheduler. actually Q provides many scheduling functionalities for AI training and uh inference

workloads such as owner nothing semantics and the right pod creation and a multi cluster dispatching. Additionally, Q is a resource quota manager which supports quota management across namespaces clusters and uh uh Q controls machine flavor preferences like spot or reserved instances. And also uh we support fair sharing. Um in this talk I want to slightly dive into topology of job scheduling. today we have a couple of

our topology of job scheduling algorithms. Uh best fit, multi layer topology, last least free capacity and balance placement. Uh each scheduling mode can be enabled by the combination of the preferred and constrained pod template level annotations and the Q feature gates. And um the first is a best of fit algorithm which tries to greedy schedule workloads to the largest free topology domains first. This scheduling algorithm supports

all topology preferences labels um required, preferred and constrained. And in the next let me quickly uh dive to the actual story. Okay. Um so let's imagine the situation where we have four layer topology domains in our cluster and uh and this three pods has already been assigned to the um left side um two racks and uh And then once we got a new training job with seven

seven replicas, how can we go how how can uh assign new nodes to all um all pods with topology aware scheduling best fit? Uh in that case, um could try to assign uh new pods um so into the largest free topology. So, the um almost six pods are scheduled into the entire free uh topology racks and the remaining one pods will be scheduled into the um so

already occupied uh topology uh domain rack. And the next is a uh these free capacity uh algorithm. Uh this is a um mostly uh similar to the uh cube scheduler bin packing uh algorithm, uh but let me re re-describes describe this uh algorithm. So, okay, again uh um we have a four-layer topology clusters and we already have three pods under the two topology racks. And uh after

we got new train training jobs with uh seven replicas, we will uh schedule all pods at into the same same topology uh as packed as possible. And the next is a a balance of balance of placement and a multi-layer topology. this uh balance of placement uh is supported of only for preferred topology preference rebel and so cluster administrator have to enable task balance placement feature gate. multi

multi multi layer topology is also supported both required and preferred topology preference rebel. this has not this has not been yet exposed to the public queue release, but we can leverage this functionality in the next 0.17 version. coming back to the our cluster and we have a three pods already. And after we got new training job with seven replicas. In case of a balance balance replacement and

preferred mode balance Oh, I guess this is a typo, so I mean both balance replacement and preferred mode plus multi layer topology algorithms. Um We we want to go go to a set of sub pod to the each topology rack um in case of a rack scale data center or server server solutions, for example, Nvidia NDL GB 200 or 300. interconnect so the number of interconnect are

exposed based on the a number of GPUs. In that case, so so that we can minimize the relationship between uh racks, uh we want to go we want to pack the minimum set of uh pods uh into a specific racks uh instead of uh being packing. Yep. Um the um actually we saw uh we have required mode and multi-area topology. And the same situation. And in this

uh we cannot make any assignment uh for this new training job uh because but required mode multi-area topology uh should be evenly divided uh by job replica job replica replica numbers. So, ac- I think uh this is a um primary different point between balance replacement and multi-area topology uh assignment. So, as I mentioned, so if um tra- workload uh the number of worker worker replicas can be

evenly divided by um remaining topology domain nodes, uh we can make uh topology assignment the same as balance replacement. So, this is a hard constraint um rather than um balance uh topology assignment. And uh so how we want to talk about uh two recent highlights uh from 0. Mm, I guess uh this is also type 0.15.0 to the B 0.16.4.4. So, first is uh, I think this

is important. Um, API graduation, um, from B1 beta 1 to B1 beta 2. And uh, before I go to the other uh, remaining uh, features, so let me quickly introduce our API graduation and migration, um, steps. in 0.15, we introduced the B1 beta 2 uh, Q API version with conversion mechanism. But, B1 beta 1 was still be served as a storage version. in the 0.16, uh, we

switched the storage version from B1 beta 2 to B1 beta 1 to And in the future 0.18, eventually we will stop B1 beta 1 uh, serving at all. So, please consider the API migration if you still use the B1 beta 1 API. And yeah, uh, coming back again. So, Yep. And uh, we have also uh, the array API upgrades. And uh, as I mentioned, we support uh,

some of the topology aware scheduling algorithms. And uh, so, additionally, so we extend the topology aware scheduling replace node node on node taints uh, to handle no schedule and no execute uh, node taints. Uh, which means so, Q try to automatically evict the no execute uh, um, I mean, uh, Q try to evict and re- assignment re- re-assign um, node to the workloads on top of no

no execute taint nodes. And we also uh uh extend the multi-queue some functionalities. Um okay, uh I guess time is coming. So, and so we we have a couple of uh um interesting uh priorities. And so uh finally, so uh we want to talk about the Q future plan. we plan to release 0.17 uh in the next week uh Tuesday or Monday, yeah, maybe. And so, we

will uh introduce uh the LA extended resource support and elastic job support. And uh So, yeah, as I mentioned, multi-area topology constraint support and uh uh supporting set of set or reader box set or subbing uh workloads in multi-queue. And uh so, uh support atomic preemption across across uh per workload admission uh to make uh preemption chance to to avoid preemption chance between uh multi-clusters. And we

have a couple of uh GA feature GA graduations. All right. And for the last part of our talk, I'm excited to talk about uh workload aware scheduling {slash} topology aware scheduling in going going core Kubernetes. generally, what is gang scheduling? Uh I think everyone always thinks it's all or nothing and it's simple, like, you know, you can just have all these pods scheduled together. What we find

is every workload has different groups. So, I have this cute little graphic showing different cats. You can and each each color cat represents a different group. And honestly every workload I've met or every person that's submitting workloads has a very different way that they want gang scheduling, and it's making our life pretty difficult. So, I have three workloads that I'm kind of showing. You have job, which

is just a set of pods. That's pretty That's somewhat easy. Unfortunately, there are even complications there, but for a mental model, a lot of people want all or nothing for jobs. Job set, you have a and this is even more It can even be more complicated, but it's you have a concept of a leader and a worker, where your leader like you need your leader to be

scheduled, and then your workers can start. In some cases, your workers don't need to all be complete, and but sometimes if your leader gets preempted, you need to kill the job. Like these things get really complicated, and then we have the beast of all workloads, which is disaggregated serving, that we're still kind of trying to figure out what the heck they actually want. Uh that was We

had a good meeting earlier this week on getting an understanding of that, but generally, it's represented by two leader worker sets that each represent a separate uh a separate way of doing inferencing. You have prefills and worker and decoders, and they communicate together through I think a KV cache, and then you have to have this like enforcing a certain ratio of prefill leaders to decoders, and I

don't fully understand it, but it's it's actually quite complicated, and it hurts my head a little bit, but at least we're getting an API for this to make this a little easier. now my mental model of AI workloads is they're honestly fickle beasts, and if there is one like basically, a lot of AI is doing a bunch of linear algebra or math, and if one pod gets

killed, you kind of lose a bunch of data in most frameworks. Some frameworks can handle the fault tolerance. A lot of A lot of ones can't. Basically, if a job gets preempted for some reason or a pod in the job gets preempted, you could potentially lose some data. So then people usually just destroy the job and recreate a new one. And there's a lot it's really hard

to do fault tolerance. But generally, this is why it's also it's not just gang scheduling, but also what what do you do in the process of when a gang is disrupted? And that semantics is not well solved. So the last thing So what are we going to do? Well, we introduced a project called workload aware scheduling. Basically, our goal is we want Kubernetes aware of how to

do group-based scheduling. We introduced two APIs, workload and pod group. Workload is a static definition. You can specify a min count. And then a pod group is like a runtime idea, and it also has min count, but it also has a status. Now, I like APIs. So this is kind of Well, you can see that. So a workload you can represent by it's just basically a these

are all namespace resources, and you can specify a min count. The idea is that a controller will read the workload, and then it will create the pod group, and then the way you link this together is you have a job, which links to a pod group template. And then this is pretty much what you need to do gang scheduling in 1.36. So yeah, workload aware disruption. Sorry,

I have to go to skip this. Basically, on the pod group you can say whether or not you allow pod or pod group preemption. Great. Workload aware resource claims. This has been the bane of my existence as a job set maintainer because Google and Nvidia keep coming up with new ways to use their hardware. And one of the the problems we run into is that people want

a group of pods to share the same resource claim for compute domains and and TPUs have reasons that I don't even understand, but they also have a similar use case where they want a bunch of pods to share the same claim and this makes it really hard to represent in the Kubernetes model. That's why DRA is the future, but and so we know that we want to

force like the sharing of resource claims, but we also need to have a resource claims that can share multiple pods and that's why we have the pod group have the concept of resource So, basically my goal as a working group batch organizer is I want to make this as simple as possible for our users. I don't think anyone should ever worry about workloads or pod groups. If

you come from the HPC world, they don't think about gang scheduling. It's just done for you. That's kind of one of our our hope in the long run is that gang scheduling is a critical component for Kubernetes controllers. So, we are working actively with some of the maintainers of job and job set and queue array to kind of bring gang scheduling into these constructs as soon as

the APIs are available. Uh it was a really busy release. Uh I will give a demo next slide, but I had to build the main as of 2 days ago cuz some of the stuff just merged 2 days ago to get to show all this. But yeah, we did six features uh for workload aware scheduling uh and the area that we're still in flight on is how

to do like queue awareness of workload aware scheduling is one that we are still figuring out and auto scaling. And now I like to switch to a demo to kind of show it pull it all together. So, uh Can you all see that? Oh, sorry. Thank you. Can you all see that? Okay, cool. Uh thank you. Uh so, I'm just going to we have kind of so

I'm okay. I got it. Uh we I'm just going to show cuz I think we're running out of time. I'm just going to show the the job Actually, I'll I'll show both. It's fine. Last conference. Uh so, I have the concept of a bring your own pod group. So, I'm going to go ahead and create a workload. And I'm going to create a pod group. And so,

uh you can see pod group. And this kind of shows you have 1,000 pods, points to a template. And then now all I want to do is I I'm going to create a job. Yeah. template. And then I link to the pod group. And if you're familiar with Kubernetes gang scheduling, you'll know that I by by convenience, this will not be placed. So, if I go ahead

and create the pods, you'll see all 1,000 uh pods are stuck in pending. This is kind of how we get like the scheduler is actually knowing about the whole group. So, this is how like if controllers are lagging behind, you could still bring gang scheduling to them even if it's not handled in the controller itself. The next thing I'm going to switch to is uh gang job.

This is Uh sorry, I'm reading the names of my files, but So, the idea here is I have a batch job that already handles gang scheduling for you. And this actually is going to go ahead And this go ahead and creates the pod group and the uh workload. And you can see that uh in this case, it the status means it's still pending, but it's actually going

to be scheduled. And then uh the last thing I just wanted to like I'm pretty excited about all this. It's been a long time coming. and I wanted to switch to the presentation just to finally close it So, basically, if any of this interests you, we I managed to even some people. We need a Slack channel. I'm glad that I did that, so it's easier for people

to come to. It was spread all over the place cuz we had six scheduling queue and a bunch of people all involved but we kind of now have a great place to hang out as workload aware scheduling in the Kubernetes slack. if you are a controller maintainer, try to think about how to integrate with this because it's coming. I'm hoping it's not going to be like DRA

where it took a lot longer. Like our goal is to try and make this like available for people as soon as possible because I'm feeling the the fact that we don't have a true gang scheduling solution in Kubernetes and it makes supporting AI workloads really challenging. Unfortunately, the conference is over so you can't meet us at the solutions showcase but we were always at the solutions showcase

every time now. Uh we were we want to be good friends with Q. Uh also, this is working group batch. Please join us if this interests you and thank you. That's uh that's all our talk.