Repository navigation
[WIP] feat: Workload-Aware Scheduling for TrainJob – Phase 1 (KEP-3015) - #4254
andreyvelich wants to merge 16 commits into
Conversation
|
@andreyvelich: GitHub didn't allow me to assign the following users: VassilisVassiliadis. Note that only kubeflow members with read permissions, repo collaborators and people who have commented on this issue/PR can be assigned. Additionally, issues/PRs can only have 10 assignees at the same time. DetailsIn response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository. |
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
There was a problem hiding this comment.
🟡 Changes recommended
Shared-device claim propagation and InOrder gang scheduling have correctness gaps, and the WAS workflow does not execute its suite.
3 open issues
6 files not reviewed
pkg/apis/trainer/v1alpha1/zz_generated.deepcopy.go: Generated filepkg/client/applyconfiguration/trainer/v1alpha1/jobscheduling.go: Generated filepkg/client/applyconfiguration/trainer/v1alpha1/replicatedjobscheduling.go: Generated filepkg/client/applyconfiguration/trainer/v1alpha1/trainingruntimescheduling.go: Generated filepkg/client/applyconfiguration/trainer/v1alpha1/trainingruntimespec.go: Generated filepkg/client/applyconfiguration/utils.go: Generated file
What changed in this PR
Introduces Phase 1 of KEP-3015, integrating Kubernetes Workload-Aware Scheduling into Trainer’s runtime APIs and plugin framework.
Changes:
- Adds feature-gated scheduling configuration, validation, and generated API support.
- Connects scheduling groups and resource claims to the JobSet pipeline.
- Adds documentation, admission tests, and E2E scaffolding.
| File | Description |
|---|---|
| test/integration/webhooks/trainjob_test.go | Tests resource-claim conflicts. |
| test/integration/webhooks/clustertrainingruntime_webhook_test.go | Tests scheduling admission. |
| test/integration/framework/framework.go | Enables scheduling in integration tests. |
| test/e2e-was/suite_test.go | Initializes the WAS suite. |
| test/e2e-was/e2e_was_test.go | Adds namespace lifecycle scaffolding. |
| proposals/3015-workload-aware-scheduling/README.md | Updates design and limitations. |
| pkg/webhooks/trainjob_webhook.go | Reuses claim-path helpers. |
| pkg/webhooks/trainingruntime_webhook.go | Validates runtime scheduling. |
| pkg/webhooks/trainingruntime_webhook_test.go | Tests scheduling validation. |
| pkg/webhooks/clustertrainingruntime_webhook.go | Applies shared scheduling validation. |
| pkg/webhooks/clustertrainingruntime_webhook_test.go | Tests cluster-runtime validation. |
| pkg/util/trainjob/trainjob.go | Adds claim-path traversal helpers. |
| pkg/util/trainjob/trainjob_test.go | Tests claim-path helpers. |
| pkg/util/testing/wrapper.go | Adds scheduling fixture support. |
| pkg/runtime/runtime.go | Carries scheduling and claim data. |
| pkg/runtime/runtime_test.go | Tests runtime data propagation. |
| pkg/runtime/interface.go | Extends RuntimeInfo parameters. |
| pkg/runtime/framework/plugins/workload/builder.go | Builds scheduling templates. |
| pkg/runtime/framework/plugins/registry.go | Registers the feature-gated plugin. |
| pkg/runtime/framework/plugins/registry_test.go | Tests plugin registration. |
| pkg/runtime/framework/plugins/jobset/jobset.go | Propagates scheduling and claims. |
| pkg/runtime/framework/plugins/jobset/jobset_test.go | Tests JobSet propagation. |
| pkg/runtime/framework/interface.go | Defines scheduling enforcement interface. |
| pkg/runtime/framework/core/framework.go | Runs scheduling enforcement plugins. |
| pkg/runtime/framework/core/framework_test.go | Tests scheduling enforcement. |
| pkg/runtime/core/trainingruntime.go | Connects scheduling to runtime processing. |
| pkg/runtime/core/trainingruntime_test.go | Tests scheduling and claim validation. |
| pkg/runtime/core/clustertrainingruntime.go | Passes cluster-runtime scheduling. |
| pkg/runtime/core/clustertrainingruntime_test.go | Tests shared-claim conflicts. |
| pkg/features/features.go | Defines the alpha feature gate. |
| pkg/client/applyconfiguration/utils.go | Registers scheduling apply types. |
| pkg/client/applyconfiguration/trainer/v1alpha1/trainingruntimespec.go | Adds scheduling apply field. |
| pkg/client/applyconfiguration/trainer/v1alpha1/trainingruntimescheduling.go | Adds runtime scheduling apply type. |
| pkg/client/applyconfiguration/trainer/v1alpha1/replicatedjobscheduling.go | Adds ReplicatedJob scheduling apply type. |
| pkg/client/applyconfiguration/trainer/v1alpha1/jobscheduling.go | Adds Job scheduling apply type. |
| pkg/apply/apply.go | Adds claim upsert helpers. |
| pkg/apply/apply_test.go | Tests claim upserts. |
| pkg/apis/trainer/v1alpha1/zz_generated.deepcopy.go | Generates scheduling deep copies. |
| pkg/apis/trainer/v1alpha1/trainingruntime_types.go | Defines scheduling APIs and validation. |
| manifests/base/rbac/role.yaml | Grants scheduling resource permissions. |
| manifests/base/crds/trainer.kubeflow.org_trainingruntimes.yaml | Adds runtime scheduling schema. |
| manifests/base/crds/trainer.kubeflow.org_clustertrainingruntimes.yaml | Adds cluster-runtime scheduling schema. |
| hack/kind-config-was.yaml | Configures Kubernetes scheduling features. |
| hack/e2e-setup-cluster-was.sh | Provisions the WAS test cluster. |
| docs/operator-guides/workload-aware-scheduling.md | Documents configuration and prerequisites. |
| docs/operator-guides/index.md | Links the scheduling guide. |
| charts/kubeflow-trainer/tests/rbac/clusterrole_test.yaml | Tests scheduling permissions. |
| charts/kubeflow-trainer/templates/rbac/clusterrole.yaml | Adds Helm scheduling permissions. |
| charts/kubeflow-trainer/templates/crd/trainer.kubeflow.org_trainingruntimes.yaml | Adds Helm runtime scheduling schema. |
| api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_training_runtime_spec.py | Exposes scheduling in Python. |
| api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_training_runtime_scheduling.py | Adds runtime scheduling model. |
| api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_replicated_job_scheduling.py | Adds ReplicatedJob scheduling model. |
| api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_job_scheduling.py | Adds Job scheduling model. |
| api/python_api/kubeflow_trainer_api/models/__init__.py | Exports scheduling models. |
| api/openapi-spec/swagger.json | Publishes scheduling API definitions. |
| .github/workflows/test-e2e-was.yaml | Adds WAS cluster-setup workflow. |
Review effort: Balanced
| Policy: workloadbuilder.PolicyInput{PodGroupData: entry.Job.SchedulingPolicy, PathElements: []string{"schedulingPolicy"}}, | ||
| Constraints: workloadbuilder.ConstraintsInput{PodGroupData: entry.Job.SchedulingConstraints, PathElements: []string{"schedulingConstraints"}}, | ||
| DisruptionMode: workloadbuilder.DisruptionModeInput{PodGroupData: entry.Job.DisruptionMode, PathElements: []string{"disruptionMode"}}, | ||
| ResourceClaims: workloadbuilder.ResourceClaimsInput{PodGroupData: entry.Job.ResourceClaims, PathElements: []string{"resourceClaims"}}, |
There was a problem hiding this comment.
Good point, for now I will validate the the Pod template has correct the resource claims. In the future, we can automatically set RCs if needed.
| replicas[entry.TargetReplicatedJobs[0]], schedulingJobReplicasErrorMsg)) | ||
| } | ||
| // The targets of a per-ReplicatedJob (level 2) entry share a PodGroup, which has a single priority. | ||
| if entry.Job == nil { |
There was a problem hiding this comment.
StartupPolicy is deprecated, and we don't need to support it in validation.
|
|
||
| - name: Setup Kind Cluster with Workload-Aware Scheduling | ||
| run: | | ||
| KIND=./bin/kind K8S_VERSION=${{ matrix.kubernetes-version }} ./hack/e2e-setup-cluster-was.sh |
|
Should we consider how to leverage the JobSet integration as much as possible here? It would be ideal to use the JobSet API as is for now and delegate as much of the lifecycle handling to JobSet. Otherwise you have to essentially replicate everything we did in JobSet. It seems there is a discussion in wg-was slack about the right approach here. |
|
|
||
| // template for the JobSet which will be used by TrainJob. | ||
| // +optional | ||
| Template JobSetTemplateSpec `json:"template,omitzero"` |
There was a problem hiding this comment.
Can we rely on the JobSet API itself? It would be less work to set the right API fields on JobSet for TrainJOb if possible.
There was a problem hiding this comment.
Do you propose that inside our plugin we translate the runtime.spec.scheduling API into JobSet's .spec.template.spec.scheduling API @kannon92? We listed delegating Workload creation to JobSet as a non-goal in the KEP
I see pros and cons for your approach.
Pros:
- We get creation of the
WorkloadandPodGroupsfor free (andCompositePodGroupsonce JobSet supports them), including per-JobPodGroupsfor ReplicatedJobs with multiple replicas, cleanup on suspend, andminCountupdates. - One API for JobSet and TrainJob users.
Cons:
- It goes against the principle that @wojtek-t @mm4tt added in KEP-6089:
https://github.com/kubernetes/enhancements/tree/master/keps/sig-scheduling/6089-was-controller-apis#core-principles--assumptions. TheWorkloadshould always be created and managed by the
root controller, which is TrainJob here. - Validation and defaulting become tricky, since JobSet doesn't know anything about the TrainJob or ClusterTrainingRuntime APIs. For example, we compute the gang
minCountfrom
numNodes, excluding the MPI launcher anddependsOnReplicatedJobs. For DRA (cc @Sridhar1030), we rejectresourceClaimsPerNodeand Pod claims inruntimePatcheswhen the runtime setsjob.resourceClaims, and we require a matching PodresourceClaimsentry. Moving this to JobSet would force us to add validation for the JobSet.spec.template.spec.schedulingfields, which can conflict with JobSet's own validation webhook.
As we discussed in this thread, we will delegate creation of PodGroups to JobSet, since it creates the Jobs for each ReplicatedJob: https://kubernetes.slack.com/archives/C0AHLJ0EAEL/p1792420422824019?thread_ts=1792117492.596419&cid=C0AHLJ0EAEL
I'm still unsure about the Workload itself.
@tenzen-y @mm4tt @robert-bell any thoughts on this?
There was a problem hiding this comment.
We could think of a more user friendly API for trainer but still set the API on JobSet side.
There was a problem hiding this comment.
I am not even sure if we need to expose any .scheduling API to the TrainJob yet.
Even if users want to override the default scheduling behaviour, we can use TrainJob's runtimePatches field for that. For the majority of use-cases the scheduling setup per Runtime might be sufficient.
But we need more discussions around this, since I might be wrong.
For the Phase 1, we need to agree which controller (TrainJob or JobSet) will create the Workload object.
There was a problem hiding this comment.
Yea for now I wonder if the simplest fix for TrainJob is best.
It seems the work with CompositePodGroups may require workoads at TrainJob level but I think jobSet should still manage the lifecycle of podgroups.
For non CPG, maybe it just makes sense to have CTR create the JobSet with scheduling fields. And let Jobset handle lifecycle.
There was a problem hiding this comment.
Yes, the PodGroups will be managed by JobSet since it creates the ReplicatedJobs.
For non CPG, maybe it just makes sense to have CTR create the JobSet with scheduling fields. And let Jobset handle lifecycle.
IIRC, once we support CPG in Phase 2, we will always create CPG -> CPG -> PG for every TrainJob workload (despite how users configure .scheduling API).
We should do the same for JobSet, let's discuss it here: kubernetes-sigs/jobset#1367
cc @yindia
|
Thanks for the ping @andreyvelich! This works well with KEP-4064 (#4145), where each MPI node group is its own replicatedJob. Listing the launcher and all groups in one entry puts them in a single gang: scheduling:
replicatedJobs:
- targetReplicatedJobs: ["launcher", "cpu", "l40s", "w7900"]Two small changes would help:
|
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
64d60b3 to
dcee4d2
Compare
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
| // Delete returns the Workload and PodGroups of a suspended, complete, or failed TrainJob. | ||
| func (w *Workload) Delete(ctx context.Context, trainJob *trainer.TrainJob) ([]client.Object, error) { | ||
| if trainJob == nil || (!trainjob.IsTrainJobFinished(trainJob) && !ptr.Deref(trainJob.Spec.Suspend, false)) { | ||
| return nil, nil | ||
| } |
There was a problem hiding this comment.
Added Delete() step to cleanup Workload and PodGroups.
Thanks for adding this capability @daehyun-kim-94!
|
Thanks @matbun!
Let's discuss it in: #4145
I think one entry can have as much ReplicatedJobs as we want, since it translates to a single PodGroup. On the other hand, we have a bug in JobSet, since number of items in ReplicatedJobScheduling should not exceed 8. |
Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
robert-bell
left a comment
robert-bell
left a comment
There was a problem hiding this comment.
Thanks @andreyvelich! I've left some initial thoughts.
| // Requires the DRAWorkloadResourceClaims Kubernetes feature gate. | ||
| // +listType=atomic | ||
| // +kubebuilder:validation:MaxItems=4 | ||
| // +kubebuilder:validation:XValidation:rule="self.all(c, has(c.resourceClaimName) != has(c.resourceClaimTemplateName))", message="exactly one of resourceClaimName or resourceClaimTemplateName must be set" |
There was a problem hiding this comment.
[P2] Reject duplicate claim names here; runtime admission currently accepts them, but the generated Workload is rejected and its Pods cannot be scheduled.
| // +kubebuilder:validation:XValidation:rule="self.all(c, has(c.resourceClaimName) != has(c.resourceClaimTemplateName))", message="exactly one of resourceClaimName or resourceClaimTemplateName must be set" | |
| // +kubebuilder:validation:XValidation:rule="self.all(c, has(c.resourceClaimName) != has(c.resourceClaimTemplateName))", message="exactly one of resourceClaimName or resourceClaimTemplateName must be set" | |
| // +kubebuilder:validation:XValidation:rule="self.all(c, self.filter(o, o.name == c.name).size() == 1)", message="resourceClaims names must be unique" |
There was a problem hiding this comment.
Do we need a separate action? We're duplicating maintenance.
Can we use the existing action, maybe with a flag to enable WAS (or even just always enable it in 1.37 cluster)?
|
|
||
| The TrainJob level can't be combined with `replicatedJobs`. Within `replicatedJobs`, each entry | ||
| can use either level 2 or level 3, and every ReplicatedJob of the runtime must be targeted by | ||
| exactly one entry. |
There was a problem hiding this comment.
I've only given the docs a light pass for now, but this bit could be made clearer - maybe an extra column summarising when you'd use each level?
| trainingRuntime := testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Obj() | ||
| trainJob := testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").Obj() | ||
|
|
||
| info, err := (&TrainingRuntime{}).newRuntimeInfo(trainJob, trainingRuntime.Spec.Template, nil, nil, scheduling) |
There was a problem hiding this comment.
nit: worth adding a clone here?
| info, err := (&TrainingRuntime{}).newRuntimeInfo(trainJob, trainingRuntime.Spec.Template, nil, nil, scheduling) | |
| info, err := (&TrainingRuntime{}).newRuntimeInfo(trainJob, trainingRuntime.Spec.Template, nil, nil, scheduling.DeepCopy()) |
|
|
||
| func TestRunEnforceSchedulingPlugins(t *testing.T) { | ||
| // Do not use t.Parallel() — SetFeatureGateDuringTest mutates global state. | ||
| featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.TrainJobWorkloadAwareScheduling, true) |
There was a problem hiding this comment.
is it worth adding a test if the feature gate is disabled?
There was a problem hiding this comment.
ah I see there's on in registry_test.go. I don't mind not having it, but I think it would be a good addition.
| schedulingFeatureGateErrorMsg = "TrainJobWorkloadAwareScheduling feature gate is disabled" | ||
| schedulingTargetNotFoundErrorMsg = "must reference a replicatedJob in spec.template.spec.replicatedJobs" | ||
| schedulingUntargetedErrorMsg = "must be targeted by a spec.scheduling.replicatedJobs entry when spec.scheduling.replicatedJobs is set" | ||
| schedulingJobReplicasErrorMsg = "must target a ReplicatedJob with at most 1 replica; a PodGroup per replica is not supported yet" |
There was a problem hiding this comment.
| schedulingJobReplicasErrorMsg = "must target a ReplicatedJob with at most 1 replica; a PodGroup per replica is not supported yet" | |
| schedulingJobReplicasErrorMsg = "must target a ReplicatedJob with at most 1 replicas; a PodGroup per replica is not supported yet" |
| if _, err := w.restMapper.RESTMapping( | ||
| schema.GroupKind{Group: schedulingv1beta1.SchemeGroupVersion.Group, Kind: kind}, | ||
| schedulingv1beta1.SchemeGroupVersion.Version, | ||
| ); err != nil { | ||
| w.logger.Error(err, "The scheduling.k8s.io/v1beta1 API must be enabled in advance") | ||
| return nil | ||
| } |
There was a problem hiding this comment.
What should happen on clusters without WAS apis once the feature is enabled by default? Should we block runtime/trainjob creation?
| } | ||
|
|
||
| var objects []apiruntime.ApplyConfiguration | ||
| if err := w.client.Get(ctx, client.ObjectKeyFromObject(workload), &schedulingv1beta1.Workload{}); err != nil { |
There was a problem hiding this comment.
are we able to use the existing ssa pattern here instead of getting the object? Looks like we're just checking for existence. Is Workload immutable?
Fixes: #3015
This adds an initial support for Kubernetes Workload-Aware Scheduling in the Trainer APIs 🚀
In the Phase 1, we only support PodGroups and Workloads via
.schedulingAPI in the TrainingRuntime.Here is the working example on my Kind cluster with 8 CPUs:
Example WAS with TrainJob
After TrainJob is created:
TODOs:
While I am working on the TODO items, I would appreciate your review of the implementation, so we can finish it before Trainer 2.4 release: #4188.
/assign @robert-bell @tenzen-y @akshaychitneni @Sridhar1030 @abhijeet-dhumal @tariq-hasan @VassilisVassiliadis
cc @kannon92 @helayoty @wojtek-t @mm4tt @vsoch @astefanutti @RobuRishabh @kaisoz @aniket2405 @matbun @killianmuldoon