attempt to commonize Kubeflow jobs Multikueue support methods

pkg/controller/jobs/kubeflow/common/common.go

@@ -0,0 +1,134 @@
+package common
+
+import (
+	"context"
+	"fmt"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/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"
+	clientutil "sigs.k8s.io/kueue/pkg/util/client"
+)
+
+type objAsPtr[T any] interface {
+	metav1.Object
+	client.Object
+	*T
+}
+
+type adapter[PtrT objAsPtr[T], T any] struct {
+	copySpec   func(dst, src PtrT)
+	copyStatus func(dst, src PtrT)
+	emptyList  func() client.ObjectList
+	gwk        schema.GroupVersionKind
+	typeName   string
+}
+
+type fullInterface interface {
+	jobframework.MultiKueueAdapter
+	jobframework.MultiKueueWatcher
+}
+
+func NewAdapter[PtrT objAsPtr[T], T any](
+	copySpec func(dst, src PtrT),
+	copyStatus func(dst, src PtrT),
+	emptyList func() client.ObjectList,
+	gwk schema.GroupVersionKind,
+	typeName string,
+) fullInterface {
+	return &adapter[PtrT, T]{
+		copySpec:   copySpec,
+		copyStatus: copyStatus,
+		emptyList:  emptyList,
+		gwk:        gwk,
+		typeName:   typeName,
+	}
+}
+
+func (a adapter[PtrT, T]) GVK() schema.GroupVersionKind {
+	return a.gwk
+}
+
+func (a adapter[PtrT, T]) KeepAdmissionCheckPending() bool {
+	return false
+}
+
+func (a adapter[PtrT, T]) IsJobManagedByKueue(_ context.Context, _ client.Client, _ types.NamespacedName) (bool, string, error) {
+	return true, "", nil
+}
+
+func (a adapter[PtrT, T]) SyncJob(
+	ctx context.Context,
+	localClient client.Client,
+	remoteClient client.Client,
+	key types.NamespacedName,
+	workloadName, origin string) error {
+	localJob := PtrT(new(T))
+	err := localClient.Get(ctx, key, localJob)
+	if err != nil {
+		return err
+	}
+
+	remoteJob := PtrT(new(T))
+	err = remoteClient.Get(ctx, key, remoteJob)
+	if client.IgnoreNotFound(err) != nil {
+		return err
+	}
+
+	if err == nil {
+		return clientutil.PatchStatus(ctx, localClient, localJob, func() (bool, error) {
+			// if the remote exists, just copy the status
+			a.copyStatus(localJob, remoteJob)
+			return true, nil
+		})
+	}
+
+	remoteJob = PtrT(new(T))
+	a.copySpec(remoteJob, localJob)
+
+	// add the prebuilt workload
+	labels := remoteJob.GetLabels()
+	if remoteJob.GetLabels() == nil {
+		labels = make(map[string]string, 2)
+	}
+	labels[constants.PrebuiltWorkloadLabel] = workloadName
+	labels[kueuealpha.MultiKueueOriginLabel] = origin
+	remoteJob.SetLabels(labels)
+
+	return remoteClient.Create(ctx, remoteJob)
+}
+
+func (a adapter[PtrT, T]) DeleteRemoteObject(ctx context.Context, remoteClient client.Client, key types.NamespacedName) error {
+	job := PtrT(new(T))
+	err := remoteClient.Get(ctx, key, job)
+	if err != nil {
+		return client.IgnoreNotFound(err)
+	}
+	return client.IgnoreNotFound(remoteClient.Delete(ctx, job))
+}
+
+func (a adapter[PtrT, T]) GetEmptyList() client.ObjectList {
+	return a.emptyList()
+}
+
+func (a adapter[PtrT, T]) WorkloadKeyFor(o runtime.Object) (types.NamespacedName, error) {
+	job, isTheJob := o.(PtrT)
+	if !isTheJob {
+		return types.NamespacedName{}, fmt.Errorf("not a %s", a.typeName)
+	}
+
+	prebuiltWl, hasPrebuiltWorkload := job.GetLabels()[constants.PrebuiltWorkloadLabel]
+	if !hasPrebuiltWorkload {
+		return types.NamespacedName{}, fmt.Errorf("no prebuilt workload found for %s: %s", a.typeName, klog.KObj(job))
+	}
+
+	return types.NamespacedName{Name: prebuiltWl, Namespace: job.GetNamespace()}, nil
+}

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

@@ -29,6 +29,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
+	kfcommon "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/common"
 	"sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/kubeflowjob"
 )
 
@@ -45,7 +46,7 @@ func init() {
 		JobType:                &kftraining.PaddleJob{},
 		AddToScheme:            kftraining.AddToScheme,
 		IsManagingObjectsOwner: isPaddleJob,
-		MultiKueueAdapter:      &multikueueAdapter{},
+		MultiKueueAdapter:      kfcommon.NewAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, kftraining.PaddleJobKind),
 	}))
 }
 

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"
+	kfcommon "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/common"
 	"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 = kfcommon.NewAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, kftraining.PaddleJobKind)
 
-	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"
+	kfcommon "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/common"
 	"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 := kfcommon.NewAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, kftraining.PaddleJobKind)
 
 			gotErr := tc.operation(ctx, adapter, managerClient, workerClient)