Skip to content

Commit

Permalink
attempt to commonize Kubeflow jobs Multikueue support methods (#2795)
Browse files Browse the repository at this point in the history 
* attempt to commonize Kubeflow jobs Multikueue support methods

* make generic adapter to be specialized by job type

* Move common Multikueue adapter to Kubeflowjob package

* Changes after code review
  • Loading branch information
@mszadkow
mszadkow authored Aug 19, 2024
1 parent d2b257a commit 801734b
Show file tree
Hide file tree
Showing 13 changed files with 214 additions and 356 deletions.
2 changes: 1 addition & 1 deletion pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_controller.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ func init() {
JobType: &kftraining.PaddleJob{},
AddToScheme: kftraining.AddToScheme,
IsManagingObjectsOwner: isPaddleJob,
MultiKueueAdapter: &multikueueAdapter{},
MultiKueueAdapter: kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk),
}))
}

Expand Down
92 changes: 9 additions & 83 deletions pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_multikueue_adapter.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
13 changes: 8 additions & 5 deletions pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_multikueue_adapter_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand All @@ -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")
},

Expand All @@ -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")
},

Expand All @@ -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})
},
},
Expand All @@ -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)

Expand Down
92 changes: 9 additions & 83 deletions pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorch_multikueue_adapter.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
12 changes: 7 additions & 5 deletions pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorch_multikueue_adapter_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand All @@ -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")
},

Expand All @@ -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")
},

Expand All @@ -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})
},
},
Expand All @@ -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)

Expand Down
2 changes: 1 addition & 1 deletion pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorchjob_controller.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ func init() {
JobType: &kftraining.PyTorchJob{},
AddToScheme: kftraining.AddToScheme,
IsManagingObjectsOwner: isPyTorchJob,
MultiKueueAdapter: &multikueueAdapter{},
MultiKueueAdapter: kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk),
}))
}

Expand Down
2 changes: 1 addition & 1 deletion pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ func init() {
JobType: &kftraining.TFJob{},
AddToScheme: kftraining.AddToScheme,
IsManagingObjectsOwner: isTFJob,
MultiKueueAdapter: &multikueueAdapter{},
MultiKueueAdapter: kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk),
}))
}

Expand Down
Loading

0 comments on commit 801734b

Please sign in to comment.