Provide Kubeflow PaddleJob support for MultiKueue

charts/kueue/templates/rbac/role.yaml

@@ -253,6 +253,7 @@ rules:
       - paddlejobs/status
     verbs:
       - get
+      - patch
       - update
   - apiGroups:
       - kubeflow.org

config/components/rbac/role.yaml

@@ -252,6 +252,7 @@ rules:
   - paddlejobs/status
   verbs:
   - get
+  - patch
   - update
 - apiGroups:
   - kubeflow.org

pkg/controller/jobframework/validation.go

@@ -39,7 +39,8 @@ var (
 	supportedPrebuiltWlJobGVKs    = sets.New(
 		batchv1.SchemeGroupVersion.WithKind("Job").String(),
 		jobset.SchemeGroupVersion.WithKind("JobSet").String(),
-		kftraining.SchemeGroupVersion.WithKind(kftraining.TFJobKind).String())
+		kftraining.SchemeGroupVersion.WithKind(kftraining.TFJobKind).String(),
+		kftraining.SchemeGroupVersion.WithKind(kftraining.PaddleJobKind).String())
 )
 
 // ValidateJobOnCreate encapsulates all GenericJob validations that must be performed on a Create operation

pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_controller.go

@@ -45,13 +45,14 @@ func init() {
 		JobType:                &kftraining.PaddleJob{},
 		AddToScheme:            kftraining.AddToScheme,
 		IsManagingObjectsOwner: isPaddleJob,
+		MultiKueueAdapter:      &multikueueAdapter{},
 	}))
 }
 
 // +kubebuilder:rbac:groups=scheduling.k8s.io,resources=priorityclasses,verbs=list;get;watch
 // +kubebuilder:rbac:groups="",resources=events,verbs=create;watch;update;patch
 // +kubebuilder:rbac:groups=kubeflow.org,resources=paddlejobs,verbs=get;list;watch;update;patch
-// +kubebuilder:rbac:groups=kubeflow.org,resources=paddlejobs/status,verbs=get;update
+// +kubebuilder:rbac:groups=kubeflow.org,resources=paddlejobs/status,verbs=get;update;patch
 // +kubebuilder:rbac:groups=kubeflow.org,resources=paddlejobs/finalizers,verbs=get;update
 // +kubebuilder:rbac:groups=kueue.x-k8s.io,resources=workloads,verbs=get;list;watch;create;update;patch;delete
 // +kubebuilder:rbac:groups=kueue.x-k8s.io,resources=workloads/status,verbs=get;update;patch

pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_controller_test.go

@@ -255,8 +255,8 @@ func TestReconciler(t *testing.T) {
 			reconcilerOptions: []jobframework.Option{
 				jobframework.WithManageJobsWithoutQueueName(true),
 			},
-			job:     testingpaddlejob.MakePaddleJob("paddlejob", "ns").Parallelism(2).Obj(),
-			wantJob: testingpaddlejob.MakePaddleJob("paddlejob", "ns").Parallelism(2).Obj(),
+			job:     testingpaddlejob.MakePaddleJob("paddlejob", "ns").PaddleReplicaSpecsDefault().Parallelism(2).Obj(),
+			wantJob: testingpaddlejob.MakePaddleJob("paddlejob", "ns").PaddleReplicaSpecsDefault().Parallelism(2).Obj(),
 			wantWorkloads: []kueue.Workload{
 				*utiltesting.MakeWorkload("paddlejob", "ns").
 					PodSets(
@@ -270,12 +270,13 @@ func TestReconciler(t *testing.T) {
 			reconcilerOptions: []jobframework.Option{
 				jobframework.WithManageJobsWithoutQueueName(false),
 			},
-			job:           testingpaddlejob.MakePaddleJob("paddlejob", "ns").Parallelism(2).Obj(),
-			wantJob:       testingpaddlejob.MakePaddleJob("paddlejob", "ns").Parallelism(2).Obj(),
+			job:           testingpaddlejob.MakePaddleJob("paddlejob", "ns").PaddleReplicaSpecsDefault().Parallelism(2).Obj(),
+			wantJob:       testingpaddlejob.MakePaddleJob("paddlejob", "ns").PaddleReplicaSpecsDefault().Parallelism(2).Obj(),
 			wantWorkloads: []kueue.Workload{},
 		},
 		"when workload is evicted, suspended is reset, restore node affinity": {
 			job: testingpaddlejob.MakePaddleJob("paddlejob", "ns").
+				PaddleReplicaSpecsDefault().
 				Image("").
 				Args(nil).
 				Queue("foo").
@@ -319,6 +320,7 @@ func TestReconciler(t *testing.T) {
 					Obj(),
 			},
 			wantJob: testingpaddlejob.MakePaddleJob("paddlejob", "ns").
+				PaddleReplicaSpecsDefault().
 				Image("").
 				Args(nil).
 				Queue("foo").

pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_multikueue_adapter.go

@@ -0,0 +1,117 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package paddlejob
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	kftraining "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/klog/v2"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
+	"sigs.k8s.io/kueue/pkg/controller/constants"
+	"sigs.k8s.io/kueue/pkg/controller/jobframework"
+	"sigs.k8s.io/kueue/pkg/util/api"
+	clientutil "sigs.k8s.io/kueue/pkg/util/client"
+)
+
+type multikueueAdapter struct{}
+
+var _ jobframework.MultiKueueAdapter = (*multikueueAdapter)(nil)
+
+func (b *multikueueAdapter) SyncJob(ctx context.Context, localClient client.Client, remoteClient client.Client, key types.NamespacedName, workloadName, origin string) error {
+	localJob := kftraining.PaddleJob{}
+	err := localClient.Get(ctx, key, &localJob)
+	if err != nil {
+		return err
+	}
+
+	remoteJob := &kftraining.PaddleJob{}
+	err = remoteClient.Get(ctx, key, remoteJob)
+	if client.IgnoreNotFound(err) != nil {
+		return err
+	}
+
+	// if the remote exists, just copy the status
+	if err == nil {
+		return clientutil.PatchStatus(ctx, localClient, &localJob, func() (bool, error) {
+			localJob.Status = remoteJob.Status
+			return true, nil
+		})
+	}
+
+	remoteJob = &kftraining.PaddleJob{
+		ObjectMeta: api.CloneObjectMetaForCreation(&localJob.ObjectMeta),
+		Spec:       *localJob.Spec.DeepCopy(),
+	}
+
+	// add the prebuilt workload
+	if remoteJob.Labels == nil {
+		remoteJob.Labels = make(map[string]string, 2)
+	}
+	remoteJob.Labels[constants.PrebuiltWorkloadLabel] = workloadName
+	remoteJob.Labels[kueuealpha.MultiKueueOriginLabel] = origin
+
+	return remoteClient.Create(ctx, remoteJob)
+}
+
+func (b *multikueueAdapter) DeleteRemoteObject(ctx context.Context, remoteClient client.Client, key types.NamespacedName) error {
+	job := kftraining.PaddleJob{}
+	err := remoteClient.Get(ctx, key, &job)
+	if err != nil {
+		return client.IgnoreNotFound(err)
+	}
+	return client.IgnoreNotFound(remoteClient.Delete(ctx, &job))
+}
+
+func (b *multikueueAdapter) KeepAdmissionCheckPending() bool {
+	return false
+}
+
+func (b *multikueueAdapter) IsJobManagedByKueue(context.Context, client.Client, types.NamespacedName) (bool, string, error) {
+	return true, "", nil
+}
+
+func (b *multikueueAdapter) GVK() schema.GroupVersionKind {
+	return gvk
+}
+
+var _ jobframework.MultiKueueWatcher = (*multikueueAdapter)(nil)
+
+func (*multikueueAdapter) GetEmptyList() client.ObjectList {
+	return &kftraining.PaddleJobList{}
+}
+
+func (*multikueueAdapter) WorkloadKeyFor(o runtime.Object) (types.NamespacedName, error) {
+	paddleJob, isPaddleJob := o.(*kftraining.PaddleJob)
+	if !isPaddleJob {
+		return types.NamespacedName{}, errors.New("not a PaddleJob")
+	}
+
+	prebuiltWl, hasPrebuiltWorkload := paddleJob.Labels[constants.PrebuiltWorkloadLabel]
+	if !hasPrebuiltWorkload {
+		return types.NamespacedName{}, fmt.Errorf("no prebuilt workload found for PaddleJob: %s", klog.KObj(paddleJob))
+	}
+
+	return types.NamespacedName{Name: prebuiltWl, Namespace: paddleJob.Namespace}, nil
+}

pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_multikueue_adapter_test.go

@@ -0,0 +1,159 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package paddlejob
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	kftraining "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
+
+	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
+	"sigs.k8s.io/kueue/pkg/controller/constants"
+	"sigs.k8s.io/kueue/pkg/util/slices"
+	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
+	kfutiltesting "sigs.k8s.io/kueue/pkg/util/testingjobs/paddlejob"
+)
+
+const (
+	TestNamespace = "ns"
+)
+
+func TestMultikueueAdapter(t *testing.T) {
+	objCheckOpts := []cmp.Option{
+		cmpopts.IgnoreFields(metav1.ObjectMeta{}, "ResourceVersion"),
+		cmpopts.EquateEmpty(),
+	}
+
+	paddleJobBuilder := kfutiltesting.MakePaddleJob("paddlejob1", TestNamespace).Queue("queue").Suspend(false)
+
+	cases := map[string]struct {
+		managersPaddleJobs []kftraining.PaddleJob
+		workerPaddleJobs   []kftraining.PaddleJob
+
+		operation func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error
+
+		wantError              error
+		wantManagersPaddleJobs []kftraining.PaddleJob
+		wantWorkerPaddleJobs   []kftraining.PaddleJob
+	}{
+		"sync creates missing remote PaddleJob": {
+			managersPaddleJobs: []kftraining.PaddleJob{
+				*paddleJobBuilder.Clone().Obj(),
+			},
+			operation: func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error {
+				return adapter.SyncJob(ctx, managerClient, workerClient, types.NamespacedName{Name: "paddlejob1", Namespace: TestNamespace}, "wl1", "origin1")
+			},
+
+			wantManagersPaddleJobs: []kftraining.PaddleJob{
+				*paddleJobBuilder.Clone().Obj(),
+			},
+			wantWorkerPaddleJobs: []kftraining.PaddleJob{
+				*paddleJobBuilder.Clone().
+					Label(constants.PrebuiltWorkloadLabel, "wl1").
+					Label(kueuealpha.MultiKueueOriginLabel, "origin1").
+					Obj(),
+			},
+		},
+		"sync status from remote PaddleJob": {
+			managersPaddleJobs: []kftraining.PaddleJob{
+				*paddleJobBuilder.Clone().Obj(),
+			},
+			workerPaddleJobs: []kftraining.PaddleJob{
+				*paddleJobBuilder.Clone().
+					Label(constants.PrebuiltWorkloadLabel, "wl1").
+					Label(kueuealpha.MultiKueueOriginLabel, "origin1").
+					StatusConditions(kftraining.JobCondition{Type: kftraining.JobSucceeded, Status: corev1.ConditionTrue}).
+					Obj(),
+			},
+			operation: func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error {
+				return adapter.SyncJob(ctx, managerClient, workerClient, types.NamespacedName{Name: "paddlejob1", Namespace: TestNamespace}, "wl1", "origin1")
+			},
+
+			wantManagersPaddleJobs: []kftraining.PaddleJob{
+				*paddleJobBuilder.Clone().
+					StatusConditions(kftraining.JobCondition{Type: kftraining.JobSucceeded, Status: corev1.ConditionTrue}).
+					Obj(),
+			},
+			wantWorkerPaddleJobs: []kftraining.PaddleJob{
+				*paddleJobBuilder.Clone().
+					Label(constants.PrebuiltWorkloadLabel, "wl1").
+					Label(kueuealpha.MultiKueueOriginLabel, "origin1").
+					StatusConditions(kftraining.JobCondition{Type: kftraining.JobSucceeded, Status: corev1.ConditionTrue}).
+					Obj(),
+			},
+		},
+		"remote PaddleJob is deleted": {
+			workerPaddleJobs: []kftraining.PaddleJob{
+				*paddleJobBuilder.Clone().
+					Label(constants.PrebuiltWorkloadLabel, "wl1").
+					Label(kueuealpha.MultiKueueOriginLabel, "origin1").
+					Obj(),
+			},
+			operation: func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error {
+				return adapter.DeleteRemoteObject(ctx, workerClient, types.NamespacedName{Name: "paddlejob1", Namespace: TestNamespace})
+			},
+		},
+	}
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			managerBuilder := utiltesting.NewClientBuilder(kftraining.AddToScheme).WithInterceptorFuncs(interceptor.Funcs{SubResourcePatch: utiltesting.TreatSSAAsStrategicMerge})
+			managerBuilder = managerBuilder.WithLists(&kftraining.PaddleJobList{Items: tc.managersPaddleJobs})
+			managerBuilder = managerBuilder.WithStatusSubresource(slices.Map(tc.managersPaddleJobs, func(w *kftraining.PaddleJob) client.Object { return w })...)
+			managerClient := managerBuilder.Build()
+
+			workerBuilder := utiltesting.NewClientBuilder(kftraining.AddToScheme).WithInterceptorFuncs(interceptor.Funcs{SubResourcePatch: utiltesting.TreatSSAAsStrategicMerge})
+			workerBuilder = workerBuilder.WithLists(&kftraining.PaddleJobList{Items: tc.workerPaddleJobs})
+			workerClient := workerBuilder.Build()
+
+			ctx, _ := utiltesting.ContextWithLog(t)
+
+			adapter := &multikueueAdapter{}
+
+			gotErr := tc.operation(ctx, adapter, managerClient, workerClient)
+
+			if diff := cmp.Diff(tc.wantError, gotErr, cmpopts.EquateErrors()); diff != "" {
+				t.Errorf("unexpected error (-want/+got):\n%s", diff)
+			}
+
+			gotManagersPaddleJob := &kftraining.PaddleJobList{}
+			if err := managerClient.List(ctx, gotManagersPaddleJob); err != nil {
+				t.Errorf("unexpected list manager's PaddleJobs error %s", err)
+			} else {
+				if diff := cmp.Diff(tc.wantManagersPaddleJobs, gotManagersPaddleJob.Items, objCheckOpts...); diff != "" {
+					t.Errorf("unexpected manager's PaddleJobs (-want/+got):\n%s", diff)
+				}
+			}
+
+			gotWorkerPaddleJobs := &kftraining.PaddleJobList{}
+			if err := workerClient.List(ctx, gotWorkerPaddleJobs); err != nil {
+				t.Errorf("unexpected list worker's PaddleJobs error %s", err)
+			} else {
+				if diff := cmp.Diff(tc.wantWorkerPaddleJobs, gotWorkerPaddleJobs.Items, objCheckOpts...); diff != "" {
+					t.Errorf("unexpected worker's PaddleJobs (-want/+got):\n%s", diff)
+				}
+			}
+		})
+	}
+}