fix(operator): fix TrainJob suspend/resume webhook error (#3008)

Author: JEETDESAI25

pkg/controller/trainjob_controller.go

@@ -70,8 +70,10 @@ func WithWatchers(watchers ...TrainJobWatcher) TrainJobReconcilerOption {
 	}
 }
 
-var _ reconcile.Reconciler = (*TrainJobReconciler)(nil)
-var _ predicate.TypedPredicate[*trainer.TrainJob] = (*TrainJobReconciler)(nil)
+var (
+	_ reconcile.Reconciler                        = (*TrainJobReconciler)(nil)
+	_ predicate.TypedPredicate[*trainer.TrainJob] = (*TrainJobReconciler)(nil)
+)
 
 func NewTrainJobReconciler(client client.Client, recorder record.EventRecorder, runtimes map[string]jobruntimes.Runtime, opts ...TrainJobReconcilerOption) *TrainJobReconciler {
 	options := &TrainJobReconcilerOptions{}
@@ -145,6 +147,12 @@ func (r *TrainJobReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 }
 
 func (r *TrainJobReconciler) reconcileObjects(ctx context.Context, runtime jobruntimes.Runtime, trainJob *trainer.TrainJob) error {
+	if ptr.Deref(trainJob.Spec.Suspend, false) {
+		if err := runtime.SyncSuspend(ctx, trainJob); err != nil {
+			return err
+		}
+	}
+
 	objects, err := runtime.NewObjects(ctx, trainJob)
 	if err != nil {
 		return err
@@ -154,6 +162,11 @@ func (r *TrainJobReconciler) reconcileObjects(ctx context.Context, runtime jobru
 			return err
 		}
 	}
+	if !ptr.Deref(trainJob.Spec.Suspend, false) {
+		if err := runtime.SyncSuspend(ctx, trainJob); err != nil {
+			return err
+		}
+	}
 	return nil
 }
 

pkg/runtime/core/clustertrainingruntime.go

@@ -77,6 +77,10 @@ func (r *ClusterTrainingRuntime) TrainJobStatus(ctx context.Context, trainJob *t
 	return r.TrainingRuntime.TrainJobStatus(ctx, trainJob)
 }
 
+func (r *ClusterTrainingRuntime) SyncSuspend(ctx context.Context, trainJob *trainer.TrainJob) error {
+	return r.TrainingRuntime.SyncSuspend(ctx, trainJob)
+}
+
 func (r *ClusterTrainingRuntime) EventHandlerRegistrars() []runtime.ReconcilerBuilder {
 	return nil
 }

pkg/runtime/core/trainingruntime.go

@@ -263,3 +263,7 @@ func (r *TrainingRuntime) ValidateObjects(ctx context.Context, old, new *trainer
 	info, _ := r.newRuntimeInfo(new, trainingRuntime.Spec.Template, trainingRuntime.Spec.MLPolicy, trainingRuntime.Spec.PodGroupPolicy) // ignoring the error here as the runtime configured should be valid
 	return r.framework.RunCustomValidationPlugins(ctx, info, old, new)
 }
+
+func (r *TrainingRuntime) SyncSuspend(ctx context.Context, trainJob *trainer.TrainJob) error {
+	return r.framework.RunSuspendSyncPlugins(ctx, trainJob)
+}

pkg/runtime/framework/core/framework.go

@@ -44,6 +44,7 @@ type Framework struct {
 	podNetworkPlugins            []framework.PodNetworkPlugin
 	componentBuilderPlugins      []framework.ComponentBuilderPlugin
 	trainJobStatusPlugin         framework.TrainJobStatusPlugin
+	suspendSyncPlugins           []framework.SuspendSyncPlugin
 }
 
 func New(ctx context.Context, c client.Client, r fwkplugins.Registry, indexer client.FieldIndexer) (*Framework, error) {
@@ -85,6 +86,9 @@ func New(ctx context.Context, c client.Client, r fwkplugins.Registry, indexer cl
 			}
 			f.trainJobStatusPlugin = p
 		}
+		if p, ok := plugin.(framework.SuspendSyncPlugin); ok {
+			f.suspendSyncPlugins = append(f.suspendSyncPlugins, p)
+		}
 	}
 	f.plugins = plugins
 	return f, nil
@@ -151,6 +155,15 @@ func (f *Framework) RunTrainJobStatusPlugin(ctx context.Context, trainJob *train
 	return nil, nil
 }
 
+func (f *Framework) RunSuspendSyncPlugins(ctx context.Context, trainJob *trainer.TrainJob) error {
+	for _, plugin := range f.suspendSyncPlugins {
+		if err := plugin.SyncSuspend(ctx, trainJob); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (f *Framework) WatchExtensionPlugins() []framework.WatchExtensionPlugin {
 	return f.watchExtensionPlugins
 }

pkg/runtime/framework/core/framework_test.go

@@ -114,6 +114,9 @@ func TestNew(t *testing.T) {
 					&mpi.MPI{},
 				},
 				trainJobStatusPlugin: &jobset.JobSet{},
+				suspendSyncPlugins: []framework.SuspendSyncPlugin{
+					&jobset.JobSet{},
+				},
 			},
 		},
 		"indexer key for trainingRuntime and runtimeClass is an empty": {

pkg/runtime/framework/interface.go

@@ -65,3 +65,8 @@ type TrainJobStatusPlugin interface {
 	Plugin
 	Status(ctx context.Context, trainJob *trainer.TrainJob) (*trainer.TrainJobStatus, error)
 }
+
+type SuspendSyncPlugin interface {
+	Plugin
+	SyncSuspend(ctx context.Context, trainJob *trainer.TrainJob) error
+}

pkg/runtime/framework/plugins/jobset/jobset.go

@@ -60,11 +60,14 @@ type JobSet struct {
 	logger     logr.Logger
 }
 
-var _ framework.WatchExtensionPlugin = (*JobSet)(nil)
-var _ framework.PodNetworkPlugin = (*JobSet)(nil)
-var _ framework.ComponentBuilderPlugin = (*JobSet)(nil)
-var _ framework.TrainJobStatusPlugin = (*JobSet)(nil)
-var _ framework.CustomValidationPlugin = (*JobSet)(nil)
+var (
+	_ framework.WatchExtensionPlugin   = (*JobSet)(nil)
+	_ framework.PodNetworkPlugin       = (*JobSet)(nil)
+	_ framework.ComponentBuilderPlugin = (*JobSet)(nil)
+	_ framework.TrainJobStatusPlugin   = (*JobSet)(nil)
+	_ framework.CustomValidationPlugin = (*JobSet)(nil)
+	_ framework.SuspendSyncPlugin      = (*JobSet)(nil)
+)
 
 const Name = constants.JobSetKind
 
@@ -145,7 +148,6 @@ func (j *JobSet) Validate(ctx context.Context, info *runtime.Info, oldObj, newOb
 				}
 			}
 		}
-
 	}
 
 	return nil, allErrs
@@ -239,14 +241,15 @@ func (j *JobSet) Build(ctx context.Context, info *runtime.Info, trainJob *traine
 		return nil, fmt.Errorf("runtime info or object is missing")
 	}
 
-	// Do not update the JobSet if it already exists and is not suspended
+	// Check if JobSet already exists
 	oldJobSet := &jobsetv1alpha2.JobSet{}
 	if err := j.client.Get(ctx, client.ObjectKeyFromObject(trainJob), oldJobSet); err != nil {
 		if !apierrors.IsNotFound(err) {
 			return nil, err
 		}
 		oldJobSet = nil
 	}
+
 	if oldJobSet != nil &&
 		!ptr.Deref(trainJob.Spec.Suspend, false) &&
 		!ptr.Deref(oldJobSet.Spec.Suspend, false) {
@@ -288,12 +291,20 @@ func (j *JobSet) Build(ctx context.Context, info *runtime.Info, trainJob *traine
 
 	// TODO (andreyvelich): Refactor the builder with wrappers for PodSpec.
 	// TODO: Once we remove deprecated runtime.Info.Trainer, we should remove JobSet Builder with DeprecatedTrainer().
-	jobSet := jobSetBuilder.
+	jobSetBuilder = jobSetBuilder.
 		Initializer(trainJob).
 		Trainer(info, trainJob).
 		PodLabels(info.Scheduler.PodLabels).
-		PodAnnotations(info.Scheduler.PodAnnotations).
-		Suspend(trainJob.Spec.Suspend).
+		PodAnnotations(info.Scheduler.PodAnnotations)
+
+	// Only set suspend on initial creation. For existing JobSets, suspend is
+	// managed separately via SyncSuspend using merge patch. SSA without this
+	// field leaves the existing suspend value unchanged (releases ownership only).
+	if oldJobSet == nil {
+		jobSetBuilder = jobSetBuilder.Suspend(trainJob.Spec.Suspend)
+	}
+
+	jobSet := jobSetBuilder.
 		Build().
 		WithOwnerReferences(metav1ac.OwnerReference().
 			WithAPIVersion(trainer.GroupVersion.String()).
@@ -337,3 +348,23 @@ func (j *JobSet) Status(ctx context.Context, trainJob *trainer.TrainJob) (*train
 
 	return status, nil
 }
+
+func (j *JobSet) SyncSuspend(ctx context.Context, trainJob *trainer.TrainJob) error {
+	jobSet := &jobsetv1alpha2.JobSet{}
+	if err := j.client.Get(ctx, client.ObjectKeyFromObject(trainJob), jobSet); err != nil {
+		return client.IgnoreNotFound(err)
+	}
+
+	trainJobSuspend := ptr.Deref(trainJob.Spec.Suspend, false)
+	jobSetSuspend := ptr.Deref(jobSet.Spec.Suspend, false)
+
+	if trainJobSuspend != jobSetSuspend {
+		patch := client.MergeFrom(jobSet.DeepCopy())
+		jobSet.Spec.Suspend = ptr.To(trainJobSuspend)
+		if err := j.client.Patch(ctx, jobSet, patch); err != nil {
+			return fmt.Errorf("failed to patch JobSet suspend field: %w", err)
+		}
+	}
+
+	return nil
+}

pkg/runtime/interface.go

@@ -37,4 +37,5 @@ type Runtime interface {
 	TrainJobStatus(ctx context.Context, trainJob *trainer.TrainJob) (*trainer.TrainJobStatus, error)
 	EventHandlerRegistrars() []ReconcilerBuilder
 	ValidateObjects(ctx context.Context, old, new *trainer.TrainJob) (admission.Warnings, field.ErrorList)
+	SyncSuspend(ctx context.Context, trainJob *trainer.TrainJob) error
 }