Repository: kubernetes-sigs/kueue

Issue:
[flaky workload slicing unit test] sigs.k8s.io/kueue/pkg/workloadslicing: TestFinish/AlreadyFinishedWorkload
**What happened**:

failure on unrelated branch: https://prow.k8s.io/view/gs/kubernetes-ci-logs/pr-logs/pull/kubernetes-sigs_kueue/7075/pull-kueue-test-unit-main/1973017516351426560

**What you expected to happen**:
no failure
**How to reproduce it (as minimally and precisely as possible)**:
ci
**Anything else we need to know?**:
```
sigs.k8s.io/kueue/pkg/workloadslicing: TestFinish/AlreadyFinishedWorkload expand_less
Run #0: Failed expand_less	0s
{Failed  === RUN   TestFinish/AlreadyFinishedWorkload
    workloadslicing_test.go:377: Deactivated() (-want,+got):
          &v1beta1.Workload{
          	TypeMeta:   {},
          	ObjectMeta: {Name: "test", Namespace: "default", ResourceVersion: "1", CreationTimestamp: {Time: s"2025-09-30 13:48:12 +0000 UTC"}, ...},
          	Spec:       {PodSets: {{Name: "main", Template: {Spec: {Containers: {{Name: "c", Resources: {Requests: {s"cpu": {i: {...}, s: "100m", Format: "DecimalSI"}}}}}, RestartPolicy: "Never"}}, Count: 1}}},
          	Status: v1beta1.WorkloadStatus{
          		Conditions: []v1.Condition{
          			{
          				Type:               "Finished",
          				Status:             "True",
          				ObservedGeneration: 0,
        - 				LastTransitionTime: v1.Time{Time: s"2025-09-30 13:48:13.013682336 +0000 UTC m=+1.677315711"},
        + 				LastTransitionTime: v1.Time{Time: s"2025-09-30 13:48:12 +0000 UTC"},
          				Reason:             "ByTest",
          				Message:            "Finished by test",
          			},
          		},
          		Admission:    nil,
          		RequeueState: nil,
          		... // 8 identical fields
          	},
          }
--- FAIL: TestFinish/AlreadyFinishedWorkload (0.02s)
}
```

Hints from issue comments:
@mimowo: /kind flake 
looks very similar to https://github.com/kubernetes-sigs/kueue/pull/7076
cc @mszadkow @ichekrygin ptal@mszadkow: /assign

Solution patch:
diff --git a/pkg/controller/jobframework/reconciler.go b/pkg/controller/jobframework/reconciler.go
index c68d43079..f5cbcd774 100644
--- a/pkg/controller/jobframework/reconciler.go
+++ b/pkg/controller/jobframework/reconciler.go
@@ -834,7 +834,7 @@ func (r *JobReconciler) ensureOneWorkload(ctx context.Context, job GenericJob, o
 		// Workload slices allow modifications only to PodSet.Count.
 		// Any other changes will result in the slice being marked as incompatible,
 		// and the workload will fall back to being processed by the original ensureOneWorkload function.
-		wl, compatible, err := workloadslicing.EnsureWorkloadSlices(ctx, r.client, podSets, object, job.GVK())
+		wl, compatible, err := workloadslicing.EnsureWorkloadSlices(ctx, r.client, r.clock, podSets, object, job.GVK())
 		if err != nil {
 			return nil, err
 		}
diff --git a/pkg/controller/jobs/trainjob/trainjob_controller.go b/pkg/controller/jobs/trainjob/trainjob_controller.go
index 0d36ec091..3b7efea00 100644
--- a/pkg/controller/jobs/trainjob/trainjob_controller.go
+++ b/pkg/controller/jobs/trainjob/trainjob_controller.go
@@ -27,6 +27,7 @@ import (
 	kftrainerruntime "github.com/kubeflow/trainer/v2/pkg/runtime"
 	kftrainerruntimecore "github.com/kubeflow/trainer/v2/pkg/runtime/core"
 	kftrainerjobset "github.com/kubeflow/trainer/v2/pkg/runtime/framework/plugins/jobset"
+	trainjobutil "github.com/kubeflow/trainer/v2/pkg/util/trainjob"
 	apimeta "k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -213,7 +214,7 @@ func jobsetApplyToJobset(jobsetApply *jobsetapplyapi.JobSetApplyConfiguration) (
 }
 
 func getRuntimeSpec(trainJob *kftrainer.TrainJob) (*kftrainer.TrainingRuntimeSpec, error) {
-	if *trainJob.Spec.RuntimeRef.Kind == kftrainer.ClusterTrainingRuntimeKind {
+	if trainjobutil.RuntimeRefIsClusterTrainingRuntime(trainJob.Spec.RuntimeRef) {
 		var ctr kftrainer.ClusterTrainingRuntime
 		err := reconciler.client.Get(reconciler.ctx, client.ObjectKey{Name: trainJob.Spec.RuntimeRef.Name}, &ctr)
 		if err != nil {
diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
index a9909d4da..31bcd2da6 100644
--- a/pkg/scheduler/scheduler.go
+++ b/pkg/scheduler/scheduler.go
@@ -865,7 +865,7 @@ func (s *Scheduler) replaceWorkloadSlice(ctx context.Context, oldQueue kueue.Clu
 	}
 	reason := kueue.WorkloadSliceReplaced
 	message := fmt.Sprintf("Replaced to accommodate a workload (UID: %s, JobUID: %s) due to workload slice aggregation", newSlice.UID, newSlice.Labels[controllerconstants.JobUIDLabel])
-	if err := workloadslicing.Finish(ctx, s.client, oldSlice, reason, message); err != nil {
+	if err := workloadslicing.Finish(ctx, s.client, s.clock, oldSlice, reason, message); err != nil {
 		return fmt.Errorf("failed to replace workload slice: %w", err)
 	}
 
diff --git a/pkg/workloadslicing/workloadslicing.go b/pkg/workloadslicing/workloadslicing.go
index 4891ca4aa..56d2e7e1e 100644
--- a/pkg/workloadslicing/workloadslicing.go
+++ b/pkg/workloadslicing/workloadslicing.go
@@ -27,6 +27,7 @@ import (
 	apimeta "k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/utils/clock"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
@@ -98,17 +99,18 @@ func ReplacementForKey(wl *kueue.Workload) *workload.Reference {
 // 1. It checks if the "Finished" condition is already applied. If true, it returns immediately, doing nothing.
 // 2. If the "Finished" condition is not set, it patches the workload slice's status to add the "Finished" condition.
 // 3. If the patch fails, it returns an error.
-func Finish(ctx context.Context, clnt client.Client, workloadSlice *kueue.Workload, reason, message string) error {
+func Finish(ctx context.Context, clnt client.Client, clk clock.Clock, workloadSlice *kueue.Workload, reason, message string) error {
 	// NOOP if the workload already has "Finished" condition (irrespective of reason and message values).
 	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,
+			Type:               kueue.WorkloadFinished,
+			Status:             metav1.ConditionTrue,
+			Reason:             reason,
+			Message:            message,
+			LastTransitionTime: metav1.NewTime(clk.Now()),
 		}), nil
 	}); err != nil {
 		return fmt.Errorf("failed to patch workload slice status: %w", err)
@@ -171,7 +173,7 @@ func ScaledUp(workload *kueue.Workload) bool {
 // - *Workload, true, nil: when a compatible workload exists or a new slice is needed.
 // - nil, false, nil: when an incompatible workload exists and no update is performed.
 // - error: on failure to fetch, update, or deactivate a workload slice.
-func EnsureWorkloadSlices(ctx context.Context, clnt client.Client, jobPodSets []kueue.PodSet, jobObject client.Object, jobObjectGVK schema.GroupVersionKind) (*kueue.Workload, bool, error) {
+func EnsureWorkloadSlices(ctx context.Context, clnt client.Client, clk clock.Clock, jobPodSets []kueue.PodSet, jobObject client.Object, jobObjectGVK schema.GroupVersionKind) (*kueue.Workload, bool, error) {
 	jobPodSetsCounts := workload.ExtractPodSetCounts(jobPodSets)
 
 	workloads, err := FindNotFinishedWorkloads(ctx, clnt, jobObject, jobObjectGVK)
@@ -225,7 +227,7 @@ func EnsureWorkloadSlices(ctx context.Context, clnt client.Client, jobPodSets []
 		// explicitly evicted.
 		if evictedCondition := apimeta.FindStatusCondition(oldWorkload.Status.Conditions, kueue.WorkloadEvicted); !workload.HasQuotaReservation(&oldWorkload) || evictedCondition != nil {
 			// Finish the old workload slice as out of sync.
-			if err := Finish(ctx, clnt, &oldWorkload, kueue.WorkloadFinishedReasonOutOfSync, "The workload slice is out of sync with its parent job"); err != nil {
+			if err := Finish(ctx, clnt, clk, &oldWorkload, kueue.WorkloadFinishedReasonOutOfSync, "The workload slice is out of sync with its parent job"); err != nil {
 				return nil, true, err
 			}
 		}



Test patch:
diff --git a/pkg/util/testing/wrappers.go b/pkg/util/testing/wrappers.go
index a0b7b4114..6f8190ba9 100644
--- a/pkg/util/testing/wrappers.go
+++ b/pkg/util/testing/wrappers.go
@@ -207,10 +207,14 @@ func (w *WorkloadWrapper) AdmittedAt(a bool, t time.Time) *WorkloadWrapper {
 }
 
 func (w *WorkloadWrapper) Finished() *WorkloadWrapper {
+	return w.FinishedAt(time.Now())
+}
+
+func (w *WorkloadWrapper) FinishedAt(t time.Time) *WorkloadWrapper {
 	cond := metav1.Condition{
 		Type:               kueue.WorkloadFinished,
 		Status:             metav1.ConditionTrue,
-		LastTransitionTime: metav1.Now(),
+		LastTransitionTime: metav1.NewTime(t),
 		Reason:             "ByTest",
 		Message:            "Finished by test",
 	}
diff --git a/pkg/workloadslicing/workloadslicing_test.go b/pkg/workloadslicing/workloadslicing_test.go
index 7163b33d7..1131a7992 100644
--- a/pkg/workloadslicing/workloadslicing_test.go
+++ b/pkg/workloadslicing/workloadslicing_test.go
@@ -34,6 +34,7 @@ import (
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/uuid"
 	"k8s.io/client-go/kubernetes/scheme"
+	testingclock "k8s.io/utils/clock/testing"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 	"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
@@ -309,6 +310,7 @@ func TestFinish(t *testing.T) {
 		workload *kueue.Workload
 	}
 	now := time.Now()
+	fakeClock := testingclock.NewFakeClock(now)
 	tests := map[string]struct {
 		args args
 		want want
@@ -331,13 +333,13 @@ func TestFinish(t *testing.T) {
 		"AlreadyFinishedWorkload": {
 			args: args{
 				clnt: testWorkloadClientBuilder().
-					WithObjects(testWorkload("test", "test-job", "job-uid", now).Finished().Obj()).Build(),
-				workloadSlice: testWorkload("test", "test-job", "job-uid", now).Finished().Obj(),
+					WithObjects(testWorkload("test", "test-job", "job-uid", now).FinishedAt(now).Obj()).Build(),
+				workloadSlice: testWorkload("test", "test-job", "job-uid", now).FinishedAt(now).Obj(),
 				reason:        "TestReason",
 				message:       "Test Message.",
 			},
 			want: want{
-				workload: testWorkload("test", "test-job", "job-uid", now).Finished().Obj(),
+				workload: testWorkload("test", "test-job", "job-uid", now).FinishedAt(now).Obj(),
 			},
 		},
 		"NotFinished": {
@@ -353,10 +355,11 @@ func TestFinish(t *testing.T) {
 				workload: testWorkload("test", "test-job", "job-uid", now).
 					ResourceVersion("2").
 					Condition(metav1.Condition{
-						Type:    kueue.WorkloadFinished,
-						Status:  metav1.ConditionTrue,
-						Reason:  "TestReason",
-						Message: "Test Message.",
+						Type:               kueue.WorkloadFinished,
+						Status:             metav1.ConditionTrue,
+						Reason:             "TestReason",
+						Message:            "Test Message.",
+						LastTransitionTime: metav1.NewTime(now),
 					}).Obj(),
 			},
 		},
@@ -364,7 +367,7 @@ func TestFinish(t *testing.T) {
 	for name, tt := range tests {
 		t.Run(name, func(t *testing.T) {
 			ctx, _ := utiltesting.ContextWithLog(t)
-			if err := Finish(ctx, tt.args.clnt, tt.args.workloadSlice, tt.args.reason, tt.args.message); (err != nil) != tt.want.err {
+			if err := Finish(ctx, tt.args.clnt, fakeClock, tt.args.workloadSlice, tt.args.reason, tt.args.message); (err != nil) != tt.want.err {
 				t.Errorf("Finish() error = %v, wantErr %v", err, tt.want.err)
 			}
 			if tt.want.workload != nil {
@@ -394,6 +397,7 @@ func TestEnsureWorkloadSlices(t *testing.T) {
 		error      bool
 	}
 	now := time.Now()
+	fakeClock := testingclock.NewFakeClock(now)
 	fiveMinutesAgo := now.Add(-5 * time.Minute)
 	testPodSets := func(count int32) []kueue.PodSet {
 		return []kueue.PodSet{
@@ -984,7 +988,7 @@ func TestEnsureWorkloadSlices(t *testing.T) {
 	for name, tt := range tests {
 		t.Run(name, func(t *testing.T) {
 			ctx, _ := utiltesting.ContextWithLog(t)
-			gotWorkload, gotCompatible, gotError := EnsureWorkloadSlices(ctx, tt.args.clnt, tt.args.jobPodSets, tt.args.jobObject, tt.args.jobObjectGVK)
+			gotWorkload, gotCompatible, gotError := EnsureWorkloadSlices(ctx, tt.args.clnt, fakeClock, tt.args.jobPodSets, tt.args.jobObject, tt.args.jobObjectGVK)
 			if (gotError != nil) != tt.want.error {
 				t.Errorf("EnsureWorkloadSlices() error = %v, wantErr %v", gotError, tt.want.error)
 				return
diff --git a/test/integration/multikueue/jobs_test.go b/test/integration/multikueue/jobs_test.go
index a6e0e76c7..52ef34c90 100644
--- a/test/integration/multikueue/jobs_test.go
+++ b/test/integration/multikueue/jobs_test.go
@@ -35,6 +35,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/sets"
+	"k8s.io/utils/clock"
 	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/manager"
@@ -1628,6 +1629,7 @@ var _ = ginkgo.Describe("MultiKueue", ginkgo.Ordered, ginkgo.ContinueOnFailure,
 		manager := managerTestCluster
 		worker1 := worker1TestCluster
 		worker2 := worker2TestCluster
+		realClock := clock.RealClock{}
 
 		features.SetFeatureGateDuringTest(ginkgo.GinkgoTB(), features.ElasticJobsViaWorkloadSlices, true)
 
@@ -1758,7 +1760,7 @@ var _ = ginkgo.Describe("MultiKueue", ginkgo.Ordered, ginkgo.ContinueOnFailure,
 		ginkgo.By("admit the new workload and finish the old workload in the manager cluster", func() {
 			gomega.Eventually(func(g gomega.Gomega) {
 				oldWorkload := getWorkload(g, manager.ctx, manager.client, workloadKey)
-				g.Expect(workloadslicing.Finish(manager.ctx, manager.client, oldWorkload, kueue.WorkloadSliceReplaced, "Replaced to accommodate a new slice")).To(gomega.Succeed())
+				g.Expect(workloadslicing.Finish(manager.ctx, manager.client, realClock, oldWorkload, kueue.WorkloadSliceReplaced, "Replaced to accommodate a new slice")).To(gomega.Succeed())
 			}, util.Timeout, util.Interval).Should(gomega.Succeed())
 			util.SetQuotaReservation(manager.ctx, manager.client, newWorkloadKey, utiltesting.MakeAdmission(managerCq.Name).Obj())
 		})
@@ -1795,7 +1797,7 @@ var _ = ginkgo.Describe("MultiKueue", ginkgo.Ordered, ginkgo.ContinueOnFailure,
 			util.SetQuotaReservation(worker1.ctx, worker1.client, newWorkloadKey, utiltesting.MakeAdmission(managerCq.Name).Obj())
 			gomega.Eventually(func(g gomega.Gomega) {
 				workload := getWorkload(g, worker1.ctx, worker1.client, workloadKey)
-				g.Expect(workloadslicing.Finish(worker1.ctx, worker1.client, workload, kueue.WorkloadSliceReplaced, "Replaced to accommodate a new slice")).To(gomega.Succeed())
+				g.Expect(workloadslicing.Finish(worker1.ctx, worker1.client, realClock, workload, kueue.WorkloadSliceReplaced, "Replaced to accommodate a new slice")).To(gomega.Succeed())
 			}, util.Timeout, util.Interval).Should(gomega.Succeed())
 		})