PodGroup / Workload API Integration with ComputeDomain

[johnbelamaric]

Today, the ComputeDomain CR is created by the user, and in turn it creates the necessary ResourceClaimTemplates. For workloads like need a ComputeDomain for a group of Pods (e.g., LWS, JobSet, TrainJob), this means that there is no way to match the lifecycle of the ComputeDomain to the lifecycle of those groups of Pods (e.g., an LWS replica).

SIG Scheduling has been working on a PodGroup concept, which models a group of related Pods. The PodGroup provides a few things:

• Gang scheduling policy for the pods in the group
• Topology constraints for the pods in the group
• Resource claims for resources that are shared across pods in a PodGroup
• A single integration point for workload controllers that operate on groups of pods

One idea that I think we should explore is that we can think of a ComputeDomain as a logical multi-node device that is allocatable via a ResourceClaim. This would allow us to lifecycle the ComputeDomain along with those PodGroups. This is a totally different flow from today. For example, with LWS it would work something like this:

• User creates:
  • ResourceClaimTemplate asking for a ComputeDomain (this would be a new DeviceClass). All it really needs is a name and the request for the DeviceClass.
  • Workload object, which contains a PodGroupTemplate that references the ResourceClaimTemplate
  • A LWS that references the Workload object / PodGroupTemplate
• LWS controller creates a replica, which means it creates a PodGroup using the PodGroupTemplate
• ResourceClaimController sees the PodGroup with the ResourceClaimTemplate, and uses it to create a ResourceClaim for the ComputeDomain
• ComputeDomainController sees that and either creates a ComputeDomain CR and just follows today's flow, or reads the RC directly and bypasses the CR as it may no longer be needed

Later, when the LWS is scaled down:

• LWS controller deletes the replica - meaning the PodGroup will be deleted, which is an explicit signal from the LWS controller that the group is no longer needed
• ResourceClaimController can de-allocate the RC and delete it
• ComputeDomainContrller sees that and deletes the CR and/or any RCTs it created

Overall the idea is to have these upstream abstractions (PodGroup, ResourceClaim) which workload controllers can create and use to represent relationships between them, and then controllers/drivers can actuate the real things following the lifecycle of those abstractions.

The RC has the "config" section directly in it as well as in the DeviceClass. That is a way to stuff implementation-specific details needed by those controllers into that object. So upstream things just understand the relationships/lifecycles of the abstractions and the implementation details are opaque to them.

This enables things like ComputeDomains and TPU slices to be lifecycled along with the groups of pods that use them.

See also:

• Gang Scheduling Support in Kubernetes kubernetes/enhancements#4671
• WAS: Decouple PodGroup API kubernetes/enhancements#5832
• DRA: ResourceClaim Support for Workloads kubernetes/enhancements#5729
• KEP-5732: Topology-aware workload scheduling kubernetes/enhancements#5732
• KEP-5732: Topology-aware workload scheduling kubernetes/enhancements#5732
• Support Multi-Node NVLink (MNNVL) for TrainJob with Topology Aware Scheduling kubeflow/trainer#3264
• https://docs.google.com/document/d/1GKhH-dDlMziun5p-fdTLXpNj-K9UK1Hi2nuk0SxztVg/edit?resourcekey=0-ZztAemssXwSJyK5bsn0tgQ&tab=t.0#heading=h.rsx94rglc48a
• https://docs.google.com/document/d/1Fg9ughIRMtt1HmDqiGWV-w9OKdrcKf_PsH4TjuP8Y40/edit?tab=t.0#heading=h.fd6kgam8o1zo

cc @mortent @klueska @nojnhuh @kannon92 @helayoty @wojtek-t @andreyvelich

[klueska]

One idea that I think we should explore is that we can think of a ComputeDomain as a logical multi-node device that is allocatable via a ResourceClaim. This would allow us to lifecycle the ComputeDomain along with those PodGroups.

This is how ComputeDomains were originally conceived, but we backed out of this design for one very specific reason. We wanted the consumption of the ComputeDomain "resource" to implicitly limit the number of pods that can be associated with that ComputeDomain on a given node to 1 (i.e. by advertising exactly 1 device for each possible ComputeDomain that can be set up on a node, which itself is also currently just 1).

Unless something has changed since I last looked at it, a global resource claim bound to a single multi-node device can't do this. In other words, nothing prevents multiple pods all pointing to this same resource claim from landing on the same node unless something else about the pods forces them onto different nodes.

I'm not saying there's not merit in having a ResourceClaim be the first object a user creates to trigger the creation of a ComputeDomain, but I don't immediately see how we can model membership in the ComputeDomain itself via this global ResourceClaim.

[johnbelamaric]

Unless something has changed since I last looked at it, a global resource claim bound to a single multi-node device can't do this.

Right, that has not changed. In practice, the individual pods would still have GPU claims for All gpus on the node, which will force the separation. But that's not the same as the per-node IMEX setup that ComputeDomain currently manages. I don't think that needs to change - we can just get the PodGroup-level integration with having an additional shim that creates the ComputeDomain CR when the ResourceClaim is allocated, and keep the existing flow.

A question - is the IMEX setup always full-node? I am wondering how (if?) ComputeDomain / IMEX is relevant to use cases like Nadia describes here, and if we need to consider those?

[klueska]

A question - is the IMEX setup always full-node?

You can only have one IMEX daemon running per node (or technically per GPU driver so you could have multiple per node if VMs are in the picture). Right now we enforce it so only one ComputeDomain can be associated with a given node at a time (and by extension only one pod on the node can be part of that one ComputeDomain). This model is not prescriptive though if you are willing to use IMEX channels for isolation between workloads instead of IMEX daemons.

Moving to channel isolation would, indeed, allow us to have multiple pods from the same ComputeDomain on the same node (using the same IMEX channel), or multiple pods from different ComputeDomains on the same node (using different IMEX channels). However, it's a rather large change to get there, and we would also lose quite a bit in moving to this model (as outlined in #920).

I think the ideal world would be where a user could (via a ComputeDomain CR or a top level ResourceClaim) say which isolation model they want, and the system adapts to it.

The basic idea being that you start long running IMEX daemons across all your nodes. They self discover all of their sibling nodes in the same clique and connect to them by default. At runtime (based on the user's request), some nodes get removed from the default "ensemble" and added to a ComputeDomain-specific IMEX domain just like we have today.

I do see a path towards getting here, but it would take a decent amount of design work to get right.

[jgehrcke]

About what @klueska said above:

We wanted the consumption of the ComputeDomain "resource" to implicitly limit the number of pods that can be associated with that ComputeDomain on a given node to 1

Adding to that -- in the ideal world, we want a clean (downside-free) way to associate more than one pod per ComputeDomain on the same node; this greatly simplifies building distributed applications. We track that in #309. That allows for each pod to grab just a single GPU; i.e. application code does not need to manage multiple local GPUs -- that removes one hierarchy level from the overall architecture (significant reduction in complexity).