Mutable JobSet spec with opportunistic updates

[GiuseppeTT]

cc @imreddy13 @kannon92 @andreyvelich

Introduction

Currently, the JobSet spec (specifically the Pod / Job template) is immutable. If a user needs to change a configuration in a running JobSet, they must delete and recreate the object. This interrupts running workloads and causes the loss of progress.

I would love to gather feedback from the community on making JobSets mutable:

• Have you heard similar requests from users where mutation of JobSets is required?
• What are your thoughts on allowing JobSets to become mutable in general?
• What do you think about an "opportunistic" update strategy (applying changes only during natural restarts)?

Real use case

Here is a real use case that I got from a big user:

1. Setup: The user has multiple JobSets running simultaneously and many more scheduled for when capacity becomes available
2. Trigger: The user develops a new Pod template (e.g., to fix a bug in the worker container or apply an optimization)
3. Current limitation: Currently, the user can only apply this new template to newly submitted JobSets. To fix existing ones, they have to recreate them, killing healthy running JobSets, which loses progress
4. Goal: The user wants to update the template for the currently running JobSets as well
5. Constraint: The user does not want the interrupt the running Jobs
6. Desired behavior: Instead, the user wants to use natural JobSet restarts as an opportunity to change the Pod template. If a Job fails and the JobSet controller recreates it, it should come back up with the new Pod template

Potential solution

One thing that I thought is the introduction of an updatePolicy for the Pod / Job template, similar to the idea of the existing failurePolicy, startupPolicy, and successPolicy. It could potentially include updateStrategies such as:

• Never: (Current behavior) The validating webhook blocks updates to the template
• Opportunistic: The webhook allows updates to the spec. The JobSet controller applies the updated Job template only when recreating the child Jobs during a restart. It does not force the deletion of running Jobs

[kannon92]

We do have a limitation on what we can update on job/pod templates.

https://github.com/kubernetes/enhancements/blob/master/keps/sig-apps/5440-mutable-job-pod-resource-updates/README.md

In this KEP I wanted to relax immutability while Jobs are suspended. Running jobs will reject any changes other than parallelism/completions.

The job immutability may limit updates to be delete and recreate.

On suspended Jobset and Job my hope would be that we can change requests/limits.

[kannon92]

Thinking about this, it sounds like what your customer wants is a rollout strategy for JobSet/Job.

I was thinking audit trail of this would be challenging. If we count restarts what pod template corresponds to what restart.

But thinking of this it sounds like Deployment has solutions for this ie rollout history so people can always revert back to old update.

[kannon92]

The only other usecase of mutability was elastic jobsets. I worked on relaxing updates to replicas but the usecase wasn't really there.

#463

[jmracek]

Just want to chime in here finally, since I've been lurking the comment section for a couple months and praying someone implements this, along with adding /scale for replicated jobs. I have a concrete use case for this that I'll share. I work on building inference services in the generative AI space. The models I work with are latent space diffusion models built from transformers, but more generally any autoregressive LLM will have similar use cases.

One algorithm for speeding up inference in long context length models is called Ulysses. Multi-head transformers can be thought of as ensembles of "heads". Each transformer layer has some large'ish number of heads, maybe something like 16 to 64. Ulysses allows for fast, long context models by partitioning the computation of the heads among a group of cooperating GPUs. Each GPU in the group computes the results from its assigned heads and then these are all-to-all gathered for subsequent computation.

Each node will have at most 8 GPUs, which means each pod in a deployment can have at most 8 GPUs. To allow scaling of model parallelism beyond 8 ranks, we need some abstraction to manage the deployment in which the atomic unit of deployment scaling is not a pod, but some fixed size collection of pods. E.g. for 64 heads, you could use up to 64 GPUs for each model replica, distributed over 8 nodes, with each node having a single pod. This scaling bottleneck meaningfully impacts the resolution and duration of the media that can be generated by these models, mainly because of GPU VRAM getting exhausted.

Kubernetes doesn't give a lot of good options for scaling replica sets where the atomic deployment unit involved multiple pods. In the past, I've hacked my own deployment controllers on the PytorchJob resource to create deployments of large models. This is not great, since you are scaling deployments manually and can't use things like Keda. When I learned about JobSet, I was happy, because it seems like this resource can address this use case.

If any of this is unclear please let me know. Happy to explain anything that doesn't make sense about my use case. Also willing to contribute and help out, although unclear if something like this is a great first issue.

[kannon92]

Just want to chime in here finally, since I've been lurking the comment section for a couple months and praying someone implements this, along with adding /scale for replicated jobs. I have a concrete use case for this that I'll share. I work on building inference services in the generative AI space. The models I work with are latent space diffusion models built from transformers, but more generally any autoregressive LLM will have similar use cases.

One algorithm for speeding up inference in long context length models is called Ulysses. Multi-head transformers can be thought of as ensembles of "heads". Each transformer layer has some large'ish number of heads, maybe something like 16 to 64. Ulysses allows for fast, long context models by partitioning the computation of the heads among a group of cooperating GPUs. Each GPU in the group computes the results from its assigned heads and then these are all-to-all gathered for subsequent computation.

Each node will have at most 8 GPUs, which means each pod in a deployment can have at most 8 GPUs. To allow scaling of model parallelism beyond 8 ranks, we need some abstraction to manage the deployment in which the atomic unit of deployment scaling is not a pod, but some fixed size collection of pods. E.g. for 64 heads, you could use up to 64 GPUs for each model replica, distributed over 8 nodes, with each node having a single pod. This scaling bottleneck meaningfully impacts the resolution and duration of the media that can be generated by these models, mainly because of GPU VRAM getting exhausted.

Kubernetes doesn't give a lot of good options for scaling replica sets where the atomic deployment unit involved multiple pods. In the past, I've hacked my own deployment controllers on the PytorchJob resource to create deployments of large models. This is not great, since you are scaling deployments manually and can't use things like Keda. When I learned about JobSet, I was happy, because it seems like this resource can address this use case.

If any of this is unclear please let me know. Happy to explain anything that doesn't make sense about my use case. Also willing to contribute and help out, although unclear if something like this is a great first issue.

For serving/inferencing one workload API that comes to mind is LeaderWorkerSet.

I believe scale is implemented on that object.

[jmracek]

I've always said my ideas are either good but not original, or original but not good. This one falls in the former category.

Thank you for the pointer! That's exactly what I need.

[andreyvelich]

Yes, we might run into limitations due to the immutability of Jobs and Pods, as @kannon92 pointed out above:
#1108 (comment)

I’d prefer that we start by adding support for Elastic JobSets, since this is a very common use case and broadly needed.

@GiuseppeTT, are there specific JobSet parameters you’re looking to modify?

We may need to evaluate this on a case-by-case basis, as some changes could introduce additional constraints or unintended side effects.

[GiuseppeTT]

The job immutability may limit updates to be delete and recreate.

Yes, we might run into limitations due to the immutability of Jobs and Pods, as @kannon92 pointed out above:
#1108 (comment)

I understand, but I don't think this will be a problem because JobSet restarts (either natural or forced) are a great opportunity to apply updates. This is because they naturally rely on recreating objects and naturally restart the workload (so that it can start from a clean state and incorporate changes)

Updating child Jobs / Pods while they are running can be good for other use cases but it is very difficult as you mentioned.

I was thinking audit trail of this would be challenging. If we count restarts what pod template corresponds to what restart.

Makes sense. One solution is the user annotating their Pod template (maybe this can be done automatically). At the very least the Pod objects already have the jobset.sigs.k8s.io/restart-attempt annotation.

@GiuseppeTT, are there specific JobSet parameters you’re looking to modify?

The use case I mentioned needs to change the Pod template inside of the Job template inside of the JobSet spec. But I suspect that this can be generalized to changing the Job template since it's the resource being created by JobSet.

I’d prefer that we start by adding support for Elastic JobSets, since this is a very common use case and broadly needed.

I like the idea of elastic JobSets and I would love to see it in JobSet. But if the feature development is up to me, I would prioritize the flexible Job / Pod template because of the explicit user request with a real use case.

[kannon92]

Most of those updates are forbidden via job controller.

I know that we relax immutability for scheduling parameters in the pod template while a job is suspended (node selector, taints, tolerations) and the KEP I linked allows you to change requests.

AFAIK every other mutation of pod template is forbidden and would require a KEP in Kubernetes to relax them.

[GiuseppeTT]

I don't follow.

Can you explain what is wrong with the following flow?

1. Create JobSet
2. All Jobs and Pods start running
3. Update the JobSet spec (let's say I change the image name, so it's a change in the Pod template, inside of the Job template)
4. JobSet controller notices the difference between the Job template and the running Job spec
5. JobSet controller forces a restart, which deletes all Jobs
6. When creating the new Jobs (as part of the restart), the JobSet controller uses the new Job template
7. Now all Jobs and Pods are running. The Pods use the new image

For an opportunistic upgrade, step 5 could be replaced by just waiting for one of the Jobs to fail, which triggers a restart.

The Pods and Jobs are never updated in-place, instead they are recreated to apply the update.

[kannon92]

Your usecase doesn't match.

Goal: The user wants to update the template for the currently running JobSets as well
Constraint: The user does not want the interrupt the running Jobs

I think recreating Jobs with updated templates would be fine. I think to support "Goal" you would have to add support for this in Job.

[sanposhiho]

It's not realistic, at least for a short term, to support the "real" mutation of arbitrary fields in template, that is, directly editing jobs or even running pods. In my understanding, the constraint to be more specific is that the user does not want to interrupt the running jobs, and hence want to apply the updates when the jobset controller recreates child jobs (opportunistic updates). That does not interrupt running jobs/pods because jobs have to be recreated anyways. That also does not require any new changes on k/k side (jobs or pods).

[GiuseppeTT]

Your usecase doesn't match.

Sorry. The comment from @sanposhiho clarifies it.

.

In summary, @kannon92 @andreyvelich, what do you think about supporting the following workflow?

1. Create a JobSet
2. JobSet controller creates the child Jobs, which creates the child Pods
3. Time passes ...
4. (This is where the new feature is used) Update the Pod template in the JobSet spec (e.g., add a Pod label, change a container image, etc)
5. Time passes ...
6. Eventually one of the child Jobs fails due to an unrelated reason. This triggers a JobSet restart
7. The JobSet controller recreates the child Jobs but uses the updated Job template (with the updated Pod template) for new child Jobs

Example of JobSet spec before and after the update.

Before:

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: js
spec:
  updatePolicy:
    updateStrategy: Opportunistic # New field
  replicatedJobs:
  - name: rj
    replicas: 2
    template:
      spec:
        completions: 2
        parallelism: 2
        template:
          spec:
            containers:
            - name: worker
              image: image-v1

After:

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: js
spec:
  updatePolicy:
    updateStrategy: Opportunistic # New field
  replicatedJobs:
  - name: rj
    replicas: 2
    template:
      spec:
        completions: 2
        parallelism: 2
        template:
          spec:
            containers:
            - name: worker
              image: image-v2 # Notice the changed image