Prevent updating Job status on suspended jobs (#4070)

kubernetes-sigs · Jan 29, 2025 · 83e680d · 83e680d
1 parent 3142908
commit 83e680d
Show file tree

Hide file tree

Showing 17 changed files with 44 additions and 14 deletions.
diff --git a/pkg/controller/jobs/jobset/jobset_multikueue_adapter.go b/pkg/controller/jobs/jobset/jobset_multikueue_adapter.go
@@ -26,6 +26,7 @@ import (
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/klog/v2"
 	"k8s.io/utils/ptr"
+	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	jobset "sigs.k8s.io/jobset/api/jobset/v1alpha2"
 
@@ -41,6 +42,8 @@ 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 {
+	log := ctrl.LoggerFrom(ctx)
+
 	localJob := jobset.JobSet{}
 	err := localClient.Get(ctx, key, &localJob)
 	if err != nil {
@@ -55,6 +58,12 @@ func (b *multikueueAdapter) SyncJob(ctx context.Context, localClient client.Clie
 
 	// if the remote exists, just copy the status
 	if err == nil {
+		if fromObject(&localJob).IsSuspended() {
+			// Ensure the job is unsuspended before updating its status; otherwise, it will fail when patching the spec.
+			log.V(2).Info("Skipping the sync since the local job is still suspended")
+			return nil
+		}
+
 		return clientutil.PatchStatus(ctx, localClient, &localJob, func() (bool, error) {
 			localJob.Status = remoteJob.Status
 			return true, nil

diff --git a/pkg/controller/jobs/jobset/jobset_multikueue_adapter_test.go b/pkg/controller/jobs/jobset/jobset_multikueue_adapter_test.go
@@ -46,7 +46,7 @@ func TestMultikueueAdapter(t *testing.T) {
 		cmpopts.EquateEmpty(),
 	}
 
-	baseJobSetBuilder := utiltestingjobset.MakeJobSet("jobset1", TestNamespace)
+	baseJobSetBuilder := utiltestingjobset.MakeJobSet("jobset1", TestNamespace).Suspend(false)
 	baseJobSetManagedByKueueBuilder := baseJobSetBuilder.DeepCopy().ManagedBy(kueue.MultiKueueControllerName)
 
 	cases := map[string]struct {

diff --git a/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_controller.go b/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_controller.go
@@ -57,7 +57,7 @@ func init() {
 		JobType:                &kftraining.PaddleJob{},
 		AddToScheme:            kftraining.AddToScheme,
 		IsManagingObjectsOwner: isPaddleJob,
-		MultiKueueAdapter:      kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk),
+		MultiKueueAdapter:      kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject),
 	}))
 }
 

diff --git a/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_multikueue_adapter.go b/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_multikueue_adapter.go
@@ -25,7 +25,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/util/api"
 )
 
-var _ jobframework.MultiKueueAdapter = kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
+var _ jobframework.MultiKueueAdapter = kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject)
 
 func copyJobStatus(dst, src *kftraining.PaddleJob) {
 	dst.Status = src.Status

diff --git a/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_multikueue_adapter_test.go b/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_multikueue_adapter_test.go
@@ -132,7 +132,7 @@ func TestMultikueueAdapter(t *testing.T) {
 
 			ctx, _ := utiltesting.ContextWithLog(t)
 
-			adapter := kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
+			adapter := kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject)
 
 			gotErr := tc.operation(ctx, adapter, managerClient, workerClient)
 

diff --git a/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorch_multikueue_adapter.go b/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorch_multikueue_adapter.go
@@ -25,7 +25,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/util/api"
 )
 
-var _ jobframework.MultiKueueAdapter = kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
+var _ jobframework.MultiKueueAdapter = kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject)
 
 func copyJobStatus(dst, src *kftraining.PyTorchJob) {
 	dst.Status = src.Status

diff --git a/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorch_multikueue_adapter_test.go b/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorch_multikueue_adapter_test.go
@@ -131,7 +131,7 @@ func TestMultikueueAdapter(t *testing.T) {
 
 			ctx, _ := utiltesting.ContextWithLog(t)
 
-			adapter := kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
+			adapter := kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject)
 
 			gotErr := tc.operation(ctx, adapter, managerClient, workerClient)
 

diff --git a/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorchjob_controller.go b/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorchjob_controller.go
@@ -57,7 +57,7 @@ func init() {
 		JobType:                &kftraining.PyTorchJob{},
 		AddToScheme:            kftraining.AddToScheme,
 		IsManagingObjectsOwner: isPyTorchJob,
-		MultiKueueAdapter:      kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk),
+		MultiKueueAdapter:      kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject),
 	}))
 }
 

diff --git a/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller.go b/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller.go
@@ -57,7 +57,7 @@ func init() {
 		JobType:                &kftraining.TFJob{},
 		AddToScheme:            kftraining.AddToScheme,
 		IsManagingObjectsOwner: isTFJob,
-		MultiKueueAdapter:      kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk),
+		MultiKueueAdapter:      kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject),
 	}))
 }
 

diff --git a/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_multikueue_adapter.go b/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_multikueue_adapter.go
@@ -25,7 +25,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/util/api"
 )
 
-var _ jobframework.MultiKueueAdapter = kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
+var _ jobframework.MultiKueueAdapter = kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject)
 
 func copyJobStatus(dst, src *kftraining.TFJob) {
 	dst.Status = src.Status

diff --git a/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_multikueue_adapter_test.go b/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_multikueue_adapter_test.go
@@ -131,7 +131,7 @@ func TestMultikueueAdapter(t *testing.T) {
 
 			ctx, _ := utiltesting.ContextWithLog(t)
 
-			adapter := kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
+			adapter := kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject)
 
 			gotErr := tc.operation(ctx, adapter, managerClient, workerClient)
 

diff --git a/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_controller.go b/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_controller.go
@@ -57,7 +57,7 @@ func init() {
 		JobType:                &kftraining.XGBoostJob{},
 		AddToScheme:            kftraining.AddToScheme,
 		IsManagingObjectsOwner: isXGBoostJob,
-		MultiKueueAdapter:      kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk),
+		MultiKueueAdapter:      kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject),
 	}))
 }
 

diff --git a/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_multikueue_adapter.go b/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_multikueue_adapter.go
@@ -25,7 +25,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/util/api"
 )
 
-var _ jobframework.MultiKueueAdapter = kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
+var _ jobframework.MultiKueueAdapter = kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject)
 
 func copyJobStatus(dst, src *kftraining.XGBoostJob) {
 	dst.Status = src.Status

diff --git a/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_multikueue_adapter_test.go b/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_multikueue_adapter_test.go
@@ -131,7 +131,7 @@ func TestMultikueueAdapter(t *testing.T) {
 
 			ctx, _ := utiltesting.ContextWithLog(t)
 
-			adapter := kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk)
+			adapter := kubeflowjob.NewMKAdapter(copyJobSpec, copyJobStatus, getEmptyList, gvk, fromObject)
 
 			gotErr := tc.operation(ctx, adapter, managerClient, workerClient)
 

diff --git a/pkg/controller/jobs/kubeflow/kubeflowjob/kubeflowjob_multikueue_adapter.go b/pkg/controller/jobs/kubeflow/kubeflowjob/kubeflowjob_multikueue_adapter.go
@@ -25,6 +25,7 @@ import (
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/klog/v2"
+	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
@@ -44,6 +45,7 @@ type adapter[PtrT objAsPtr[T], T any] struct {
 	copyStatus func(dst, src PtrT)
 	emptyList  func() client.ObjectList
 	gvk        schema.GroupVersionKind
+	fromObject func(runtime.Object) *KubeflowJob
 }
 
 type fullInterface interface {
@@ -56,12 +58,14 @@ func NewMKAdapter[PtrT objAsPtr[T], T any](
 	copyStatus func(dst, src PtrT),
 	emptyList func() client.ObjectList,
 	gvk schema.GroupVersionKind,
+	fromObject func(runtime.Object) *KubeflowJob,
 ) fullInterface {
 	return &adapter[PtrT, T]{
 		copySpec:   copySpec,
 		copyStatus: copyStatus,
 		emptyList:  emptyList,
 		gvk:        gvk,
+		fromObject: fromObject,
 	}
 }
 
@@ -83,6 +87,8 @@ func (a adapter[PtrT, T]) SyncJob(
 	remoteClient client.Client,
 	key types.NamespacedName,
 	workloadName, origin string) error {
+	log := ctrl.LoggerFrom(ctx)
+
 	localJob := PtrT(new(T))
 	err := localClient.Get(ctx, key, localJob)
 	if err != nil {
@@ -96,6 +102,12 @@ func (a adapter[PtrT, T]) SyncJob(
 	}
 
 	if err == nil {
+		if a.fromObject(localJob).IsSuspended() {
+			// Ensure the job is unsuspended before updating its status; otherwise, it will fail when patching the spec.
+			log.V(2).Info("Skipping the sync since the local job is still suspended")
+			return nil
+		}
+
 		return clientutil.PatchStatus(ctx, localClient, localJob, func() (bool, error) {
 			// if the remote exists, just copy the status
 			a.copyStatus(localJob, remoteJob)

diff --git a/pkg/controller/jobs/mpijob/mpijob_multikueue_adapter.go b/pkg/controller/jobs/mpijob/mpijob_multikueue_adapter.go
@@ -27,6 +27,7 @@ import (
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/klog/v2"
 	"k8s.io/utils/ptr"
+	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
@@ -41,6 +42,8 @@ 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 {
+	log := ctrl.LoggerFrom(ctx)
+
 	localJob := kfmpi.MPIJob{}
 	err := localClient.Get(ctx, key, &localJob)
 	if err != nil {
@@ -55,6 +58,12 @@ func (b *multikueueAdapter) SyncJob(ctx context.Context, localClient client.Clie
 
 	// if the remote exists, just copy the status
 	if err == nil {
+		if fromObject(&localJob).IsSuspended() {
+			// Ensure the job is unsuspended before updating its status; otherwise, it will fail when patching the spec.
+			log.V(2).Info("Skipping the sync since the local job is still suspended")
+			return nil
+		}
+
 		return clientutil.PatchStatus(ctx, localClient, &localJob, func() (bool, error) {
 			localJob.Status = remoteJob.Status
 			return true, nil

diff --git a/pkg/controller/jobs/mpijob/mpijob_multikueue_adapter_test.go b/pkg/controller/jobs/mpijob/mpijob_multikueue_adapter_test.go
@@ -47,7 +47,7 @@ func TestMultikueueAdapter(t *testing.T) {
 		cmpopts.EquateEmpty(),
 	}
 
-	mpiJobBuilder := utiltestingmpijob.MakeMPIJob("mpijob1", TestNamespace)
+	mpiJobBuilder := utiltestingmpijob.MakeMPIJob("mpijob1", TestNamespace).Suspend(false)
 
 	cases := map[string]struct {
 		managersMpiJobs []kfmpi.MPIJob