Fix workload slices when creation timestamps collide (#6203)

Author: mimowo

pkg/scheduler/scheduler.go

@@ -328,7 +328,7 @@ func (s *Scheduler) schedule(ctx context.Context) wait.SpeedSignal {
 		// Evict old workload-slice if any. Note: that oldWorkloadSlice is not nil only if
 		// this is a workload-slice enabled workload and there is an old slice to evict.
 		if features.Enabled(features.ElasticJobsViaWorkloadSlices) && oldWorkloadSlice != nil {
-			if err := s.replaceWorkloadSlice(ctx, oldWorkloadSlice.WorkloadInfo.ClusterQueue, e.Obj, oldWorkloadSlice.WorkloadInfo.Obj); err != nil {
+			if err := s.replaceWorkloadSlice(ctx, oldWorkloadSlice.WorkloadInfo.ClusterQueue, e.Obj, oldWorkloadSlice.WorkloadInfo.Obj.DeepCopy()); err != nil {
 				log.Error(err, "Failed to aggregate workload slice")
 				continue
 			}

pkg/workloadslicing/workloadslicing.go

@@ -28,6 +28,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/apimachinery/pkg/util/sets"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
@@ -124,9 +125,15 @@ func FindNotFinishedWorkloads(ctx context.Context, clnt client.Client, jobObject
 		return list.Items[i].CreationTimestamp.Before(&list.Items[j].CreationTimestamp)
 	})
 
+	replacedSlices := sets.New[workload.Reference]()
+	for _, w := range list.Items {
+		if replacedKey := ReplacementForKey(&w); replacedKey != nil {
+			replacedSlices.Insert(*replacedKey)
+		}
+	}
 	// Filter out workloads with activated "Finished" condition.
 	return slices.DeleteFunc(list.Items, func(w kueue.Workload) bool {
-		return apimeta.IsStatusConditionTrue(w.Status.Conditions, kueue.WorkloadFinished)
+		return apimeta.IsStatusConditionTrue(w.Status.Conditions, kueue.WorkloadFinished) || replacedSlices.Has(workload.Key(&w))
 	}), nil
 }
 

pkg/workloadslicing/workloadslicing_test.go

@@ -219,6 +219,23 @@ func TestFindNotFinishedWorkloads(t *testing.T) {
 				*testWorkload("test-2", testJobObject.Name, testJobObject.UID, now).Obj(),
 			},
 		},
+		"TwoActiveWorkloadsSelectNotReplaced": {
+			args: args{
+				ctx: t.Context(),
+				clnt: testWorkloadClientBuilder().WithLists(&kueue.WorkloadList{
+					Items: []kueue.Workload{
+						*testWorkload("test-2", testJobObject.Name, testJobObject.UID, now).
+							Annotation(WorkloadSliceReplacementFor, string(workload.NewReference("default", "test-2"))).Obj(),
+						*testWorkload("test-1", testJobObject.Name, testJobObject.UID, now).Obj(),
+					},
+				}).Build(),
+				jobObject:    testJobObject,
+				jobObjectGVK: testJobGVK,
+			},
+			want: []kueue.Workload{
+				*testWorkload("test-1", testJobObject.Name, testJobObject.UID, now).Obj(),
+			},
+		},
 	}
 	for name, tt := range tests {
 		t.Run(name, func(t *testing.T) {