Introduce workload.Finish helper function

mszadkow · mszadkow · commit 2d34fbdbd943 · 2025-11-07T16:13:29.000+01:00
diff --git a/pkg/controller/admissionchecks/multikueue/workload.go b/pkg/controller/admissionchecks/multikueue/workload.go
@@ -46,10 +46,8 @@ import (
 	config "sigs.k8s.io/kueue/apis/config/v1beta2"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta2"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
-	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	"sigs.k8s.io/kueue/pkg/util/api"
-	clientutil "sigs.k8s.io/kueue/pkg/util/client"
 	utilmaps "sigs.k8s.io/kueue/pkg/util/maps"
 	"sigs.k8s.io/kueue/pkg/workload"
 	"sigs.k8s.io/kueue/pkg/workloadslicing"
@@ -359,24 +357,8 @@ func (w *wlReconciler) reconcileGroup(ctx context.Context, group *wlGroup) (reco
 			log.V(3).Info("Group with no adapter, skip owner status copy", "workerCluster", remote)
 		}
 
-		// copy the status to the local one
-		finishCond := metav1.Condition{
-			Type:               kueue.WorkloadFinished,
-			Status:             metav1.ConditionTrue,
-			Reason:             remoteFinishedCond.Reason,
-			Message:            remoteFinishedCond.Message,
-			LastTransitionTime: metav1.NewTime(w.clock.Now()),
-		}
-		if features.Enabled(features.WorkloadRequestUseMergePatch) {
-			return reconcile.Result{}, clientutil.PatchStatus(ctx, w.client, group.local, func() (client.Object, bool, error) {
-				apimeta.SetStatusCondition(&group.local.Status.Conditions, finishCond)
-				return group.local, true, nil
-			})
-		}
-
-		wlPatch := workload.BaseSSAWorkload(group.local, false)
-		apimeta.SetStatusCondition(&wlPatch.Status.Conditions, finishCond)
-		return reconcile.Result{}, w.client.Status().Patch(ctx, wlPatch, client.Apply, client.FieldOwner(kueue.MultiKueueControllerName+"-finish"), client.ForceOwnership)
+		// finish workload and copy the status to the local one
+		return reconcile.Result{}, workload.Finish(ctx, w.client, group.local, remoteFinishedCond.Reason, remoteFinishedCond.Message, kueue.MultiKueueControllerName, w.clock, workload.WithForceOwnership())
 	}
 
 	// 2. delete all workloads that are out of sync (other than scaled-down elastic workloads)
diff --git a/pkg/controller/jobframework/reconciler.go b/pkg/controller/jobframework/reconciler.go
@@ -452,12 +452,12 @@ func (r *JobReconciler) ReconcileGenericJob(ctx context.Context, req ctrl.Reques
 	// 2. handle job is finished.
 	if message, success, finished := job.Finished(ctx); finished {
 		log.V(3).Info("The workload is already finished")
-		if wl != nil && !apimeta.IsStatusConditionTrue(wl.Status.Conditions, kueue.WorkloadFinished) {
+		if wl != nil && !workload.IsFinished(wl) {
 			reason := kueue.WorkloadFinishedReasonSucceeded
 			if !success {
 				reason = kueue.WorkloadFinishedReasonFailed
 			}
-			err := workload.UpdateStatus(ctx, r.client, wl, kueue.WorkloadFinished, metav1.ConditionTrue, reason, message, constants.JobControllerName, r.clock)
+			err := workload.Finish(ctx, r.client, wl, reason, message, constants.JobControllerName, r.clock)
 			if err != nil && !apierrors.IsNotFound(err) {
 				return ctrl.Result{}, err
 			}
@@ -584,7 +584,7 @@ func (r *JobReconciler) ReconcileGenericJob(ctx context.Context, req ctrl.Reques
 				log.Error(err, "Unsuspending job")
 				if podset.IsPermanent(err) {
 					// Mark the workload as finished with failure since the is no point to retry.
-					errUpdateStatus := workload.UpdateStatus(ctx, r.client, wl, kueue.WorkloadFinished, metav1.ConditionTrue, FailedToStartFinishedReason, err.Error(), constants.JobControllerName, r.clock)
+					errUpdateStatus := workload.Finish(ctx, r.client, wl, FailedToStartFinishedReason, err.Error(), constants.JobControllerName, r.clock)
 					if errUpdateStatus != nil {
 						log.Error(errUpdateStatus, "Updating workload status, on start failure", "err", err)
 					}
@@ -1022,13 +1022,8 @@ func (r *JobReconciler) ensurePrebuiltWorkloadInSync(ctx context.Context, wl *ku
 			return false, err
 		}
 		// mark the workload as finished
-		err := workload.UpdateStatus(ctx, r.client, wl,
-			kueue.WorkloadFinished,
-			metav1.ConditionTrue,
-			kueue.WorkloadFinishedReasonOutOfSync,
-			"The prebuilt workload is out of sync with its user job",
-			constants.JobControllerName, r.clock)
-		return false, err
+		msg := "The prebuilt workload is out of sync with its user job"
+		return false, workload.Finish(ctx, r.client, wl, kueue.WorkloadFinishedReasonOutOfSync, msg, constants.JobControllerName, r.clock)
 	}
 	return true, nil
 }
diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go
@@ -6872,6 +6872,7 @@ func TestSchedule(t *testing.T) {
 						Status:             metav1.ConditionTrue,
 						Reason:             kueue.WorkloadSliceReplaced,
 						Message:            "Replaced to accommodate a workload (UID: , JobUID: ) due to workload slice aggregation",
+						ObservedGeneration: 1,
 						LastTransitionTime: metav1.NewTime(now),
 					}).
 					Obj(),
diff --git a/pkg/workload/workload.go b/pkg/workload/workload.go
@@ -819,6 +819,18 @@ func SetEvictedCondition(w *kueue.Workload, now time.Time, reason string, messag
 	return apimeta.SetStatusCondition(&w.Status.Conditions, condition)
 }
 
+func SetFinishedCondition(w *kueue.Workload, now time.Time, reason string, message string) bool {
+	condition := metav1.Condition{
+		Type:               kueue.WorkloadFinished,
+		Status:             metav1.ConditionTrue,
+		LastTransitionTime: metav1.NewTime(now),
+		Reason:             reason,
+		Message:            api.TruncateConditionMessage(message),
+		ObservedGeneration: w.Generation,
+	}
+	return apimeta.SetStatusCondition(&w.Status.Conditions, condition)
+}
+
 // PropagateResourceRequests synchronizes w.Status.ResourceRequests to
 // with info.TotalRequests if the feature gate is enabled and returns true if w was updated
 func PropagateResourceRequests(w *kueue.Workload, info *Info) bool {
@@ -1301,6 +1313,58 @@ func Evict(ctx context.Context, c client.Client, recorder record.EventRecorder,
 	return nil
 }
 
+type FinishOption func(*FinishOptions)
+
+type FinishOptions struct {
+	UsePatch       bool
+	ForceOwnership bool
+}
+
+func DefaultFinishOptions() *FinishOptions {
+	return &FinishOptions{
+		UsePatch:       false,
+		ForceOwnership: false,
+	}
+}
+
+func WithUsePatch() FinishOption {
+	return func(o *FinishOptions) {
+		o.UsePatch = true
+	}
+}
+
+func WithForceOwnership() FinishOption {
+	return func(o *FinishOptions) {
+		o.ForceOwnership = true
+	}
+}
+
+func Finish(ctx context.Context, c client.Client, wl *kueue.Workload, reason, msg, managerPrefix string, clock clock.Clock, options ...FinishOption) error {
+	optsFinish := DefaultFinishOptions()
+	for _, opt := range options {
+		opt(optsFinish)
+	}
+	usePatch := optsFinish.UsePatch
+	forceOwnership := optsFinish.ForceOwnership
+
+	if features.Enabled(features.WorkloadRequestUseMergePatch) || usePatch {
+		return clientutil.PatchStatus(ctx, c, wl, func() (client.Object, bool, error) {
+			update := SetFinishedCondition(wl, clock.Now(), reason, msg)
+			return wl, update, nil
+		})
+	} else {
+		newWl := BaseSSAWorkload(wl, false)
+		SetFinishedCondition(newWl, clock.Now(), reason, msg)
+
+		fieldOwner := client.FieldOwner(managerPrefix + "-" + kueue.WorkloadFinished)
+		applyOpts := []client.SubResourcePatchOption{fieldOwner}
+		if forceOwnership {
+			applyOpts = append(applyOpts, client.ForceOwnership)
+		}
+		return c.Status().Patch(ctx, newWl, client.Apply, applyOpts...)
+	}
+}
+
 func PriorityClassName(wl *kueue.Workload) string {
 	if wl.Spec.PriorityClassRef != nil {
 		return wl.Spec.PriorityClassRef.Name
diff --git a/pkg/workloadslicing/workloadslicing.go b/pkg/workloadslicing/workloadslicing.go
@@ -104,15 +104,7 @@ func Finish(ctx context.Context, clnt client.Client, clk clock.Clock, workloadSl
 	if apimeta.IsStatusConditionTrue(workloadSlice.Status.Conditions, kueue.WorkloadFinished) {
 		return nil
 	}
-	if err := clientutil.PatchStatus(ctx, clnt, workloadSlice, func() (client.Object, bool, error) {
-		return workloadSlice, apimeta.SetStatusCondition(&workloadSlice.Status.Conditions, metav1.Condition{
-			Type:               kueue.WorkloadFinished,
-			Status:             metav1.ConditionTrue,
-			Reason:             reason,
-			Message:            message,
-			LastTransitionTime: metav1.NewTime(clk.Now()),
-		}), nil
-	}); err != nil {
+	if err := workload.Finish(ctx, clnt, workloadSlice, reason, message, "", clk, workload.WithUsePatch()); err != nil {
 		return fmt.Errorf("failed to patch workload slice status: %w", err)
 	}
 	return nil
