attempt to commonize Kubeflow jobs Multikueue support methods

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

@@ -45,7 +45,7 @@ func init() {
 		JobType:                &kftraining.PaddleJob{},
 		AddToScheme:            kftraining.AddToScheme,
 		IsManagingObjectsOwner: isPaddleJob,
-		MultiKueueAdapter:      &multikueueAdapter{},
+		MultiKueueAdapter:      kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk),
 	}))
 }
 

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

@@ -17,101 +17,27 @@ 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/controller/jobs/kubeflow/kubeflowjob"
 	"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
-	}
+var _ jobframework.MultiKueueAdapter = kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
 
-	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 copyJobStatus(dst, src *kftraining.PaddleJob) {
+	dst.Status = src.Status
 }
 
-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)
+func copyJobSpec(dst, src *kftraining.PaddleJob) {
+	*dst = kftraining.PaddleJob{
+		ObjectMeta: api.CloneObjectMetaForCreation(&src.ObjectMeta),
+		Spec:       *src.Spec.DeepCopy(),
 	}
-	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 {
+func 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

@@ -31,6 +31,8 @@ import (
 
 	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/controller/jobs/kubeflow/kubeflowjob"
 	"sigs.k8s.io/kueue/pkg/util/slices"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	kfutiltesting "sigs.k8s.io/kueue/pkg/util/testingjobs/paddlejob"
@@ -52,7 +54,7 @@ func TestMultikueueAdapter(t *testing.T) {
 		managersPaddleJobs []kftraining.PaddleJob
 		workerPaddleJobs   []kftraining.PaddleJob
 
-		operation func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error
+		operation func(ctx context.Context, adapter jobframework.MultiKueueAdapter, managerClient, workerClient client.Client) error
 
 		wantError              error
 		wantManagersPaddleJobs []kftraining.PaddleJob
@@ -62,7 +64,8 @@ func TestMultikueueAdapter(t *testing.T) {
 			managersPaddleJobs: []kftraining.PaddleJob{
 				*paddleJobBuilder.Clone().Obj(),
 			},
-			operation: func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error {
+
+			operation: func(ctx context.Context, adapter jobframework.MultiKueueAdapter, managerClient, workerClient client.Client) error {
 				return adapter.SyncJob(ctx, managerClient, workerClient, types.NamespacedName{Name: "paddlejob1", Namespace: TestNamespace}, "wl1", "origin1")
 			},
 
@@ -87,7 +90,7 @@ func TestMultikueueAdapter(t *testing.T) {
 					StatusConditions(kftraining.JobCondition{Type: kftraining.JobSucceeded, Status: corev1.ConditionTrue}).
 					Obj(),
 			},
-			operation: func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error {
+			operation: func(ctx context.Context, adapter jobframework.MultiKueueAdapter, managerClient, workerClient client.Client) error {
 				return adapter.SyncJob(ctx, managerClient, workerClient, types.NamespacedName{Name: "paddlejob1", Namespace: TestNamespace}, "wl1", "origin1")
 			},
 
@@ -111,7 +114,7 @@ func TestMultikueueAdapter(t *testing.T) {
 					Label(kueuealpha.MultiKueueOriginLabel, "origin1").
 					Obj(),
 			},
-			operation: func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error {
+			operation: func(ctx context.Context, adapter jobframework.MultiKueueAdapter, managerClient, workerClient client.Client) error {
 				return adapter.DeleteRemoteObject(ctx, workerClient, types.NamespacedName{Name: "paddlejob1", Namespace: TestNamespace})
 			},
 		},
@@ -129,7 +132,7 @@ func TestMultikueueAdapter(t *testing.T) {
 
 			ctx, _ := utiltesting.ContextWithLog(t)
 
-			adapter := &multikueueAdapter{}
+			adapter := kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
 
 			gotErr := tc.operation(ctx, adapter, managerClient, workerClient)
 

pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorch_multikueue_adapter.go

@@ -17,101 +17,27 @@ limitations under the License.
 package pytorchjob
 
 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/controller/jobs/kubeflow/kubeflowjob"
 	"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.PyTorchJob{}
-	err := localClient.Get(ctx, key, &localJob)
-	if err != nil {
-		return err
-	}
+var _ jobframework.MultiKueueAdapter = kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
 
-	remoteJob := &kftraining.PyTorchJob{}
-	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.PyTorchJob{
-		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 copyJobStatus(dst, src *kftraining.PyTorchJob) {
+	dst.Status = src.Status
 }
 
-func (b *multikueueAdapter) DeleteRemoteObject(ctx context.Context, remoteClient client.Client, key types.NamespacedName) error {
-	job := kftraining.PyTorchJob{}
-	err := remoteClient.Get(ctx, key, &job)
-	if err != nil {
-		return client.IgnoreNotFound(err)
+func copyJobSpec(dst, src *kftraining.PyTorchJob) {
+	*dst = kftraining.PyTorchJob{
+		ObjectMeta: api.CloneObjectMetaForCreation(&src.ObjectMeta),
+		Spec:       *src.Spec.DeepCopy(),
 	}
-	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 {
+func getEmptyList() client.ObjectList {
 	return &kftraining.PyTorchJobList{}
 }
-
-func (*multikueueAdapter) WorkloadKeyFor(o runtime.Object) (types.NamespacedName, error) {
-	pyTorchJob, isPyTorchJob := o.(*kftraining.PyTorchJob)
-	if !isPyTorchJob {
-		return types.NamespacedName{}, errors.New("not a PyTorchJob")
-	}
-
-	prebuiltWl, hasPrebuiltWorkload := pyTorchJob.Labels[constants.PrebuiltWorkloadLabel]
-	if !hasPrebuiltWorkload {
-		return types.NamespacedName{}, fmt.Errorf("no prebuilt workload found for PyTorchJob: %s", klog.KObj(pyTorchJob))
-	}
-
-	return types.NamespacedName{Name: prebuiltWl, Namespace: pyTorchJob.Namespace}, nil
-}

pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorch_multikueue_adapter_test.go

@@ -32,6 +32,8 @@ import (
 
 	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/controller/jobs/kubeflow/kubeflowjob"
 	"sigs.k8s.io/kueue/pkg/util/slices"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	kfutiltesting "sigs.k8s.io/kueue/pkg/util/testingjobs/pytorchjob"
@@ -53,7 +55,7 @@ func TestMultikueueAdapter(t *testing.T) {
 		managersPyTorchJobs []kftraining.PyTorchJob
 		workerPyTorchJobs   []kftraining.PyTorchJob
 
-		operation func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error
+		operation func(ctx context.Context, adapter jobframework.MultiKueueAdapter, managerClient, workerClient client.Client) error
 
 		wantError               error
 		wantManagersPyTorchJobs []kftraining.PyTorchJob
@@ -63,7 +65,7 @@ func TestMultikueueAdapter(t *testing.T) {
 			managersPyTorchJobs: []kftraining.PyTorchJob{
 				*pyTorchJobBuilder.Clone().Obj(),
 			},
-			operation: func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error {
+			operation: func(ctx context.Context, adapter jobframework.MultiKueueAdapter, managerClient, workerClient client.Client) error {
 				return adapter.SyncJob(ctx, managerClient, workerClient, types.NamespacedName{Name: "pytorchjob1", Namespace: TestNamespace}, "wl1", "origin1")
 			},
 
@@ -88,7 +90,7 @@ func TestMultikueueAdapter(t *testing.T) {
 					StatusConditions(kftraining.JobCondition{Type: kftraining.JobSucceeded, Status: corev1.ConditionTrue}).
 					Obj(),
 			},
-			operation: func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error {
+			operation: func(ctx context.Context, adapter jobframework.MultiKueueAdapter, managerClient, workerClient client.Client) error {
 				return adapter.SyncJob(ctx, managerClient, workerClient, types.NamespacedName{Name: "pytorchjob1", Namespace: TestNamespace}, "wl1", "origin1")
 			},
 
@@ -112,7 +114,7 @@ func TestMultikueueAdapter(t *testing.T) {
 					Label(kueuealpha.MultiKueueOriginLabel, "origin1").
 					Obj(),
 			},
-			operation: func(ctx context.Context, adapter *multikueueAdapter, managerClient, workerClient client.Client) error {
+			operation: func(ctx context.Context, adapter jobframework.MultiKueueAdapter, managerClient, workerClient client.Client) error {
 				return adapter.DeleteRemoteObject(ctx, workerClient, types.NamespacedName{Name: "pytorchjob1", Namespace: TestNamespace})
 			},
 		},
@@ -130,7 +132,7 @@ func TestMultikueueAdapter(t *testing.T) {
 
 			ctx, _ := utiltesting.ContextWithLog(t)
 
-			adapter := &multikueueAdapter{}
+			adapter := kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
 
 			gotErr := tc.operation(ctx, adapter, managerClient, workerClient)
 

pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorchjob_controller.go

@@ -45,7 +45,7 @@ func init() {
 		JobType:                &kftraining.PyTorchJob{},
 		AddToScheme:            kftraining.AddToScheme,
 		IsManagingObjectsOwner: isPyTorchJob,
-		MultiKueueAdapter:      &multikueueAdapter{},
+		MultiKueueAdapter:      kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk),
 	}))
 }
 

pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller.go

@@ -45,7 +45,7 @@ func init() {
 		JobType:                &kftraining.TFJob{},
 		AddToScheme:            kftraining.AddToScheme,
 		IsManagingObjectsOwner: isTFJob,
-		MultiKueueAdapter:      &multikueueAdapter{},
+		MultiKueueAdapter:      kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk),
 	}))
 }