Skip to content

Commit 0ead2c1

Browse files
mbobrovskyimbobrovskyi
committed
Use WorkloadRequestUseMergePatch in UpdateReclaimablePods().
1 parent 81cae06 commit 0ead2c1

File tree

4 files changed

+77
-37
lines changed

4 files changed

+77
-37
lines changed

‎pkg/workload/workload.go‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1098,9 +1098,28 @@ func HasQuotaReservation(w *kueue.Workload) bool {
10981098

10991099
// UpdateReclaimablePods updates the ReclaimablePods list for the workload with SSA.
11001100
func UpdateReclaimablePods(ctx context.Context, c client.Client, w *kueue.Workload, reclaimablePods []kueue.ReclaimablePod) error {
1101-
patch := BaseSSAWorkload(w, false)
1102-
patch.Status.ReclaimablePods = reclaimablePods
1103-
return c.Status().Patch(ctx, patch, client.Apply, client.FieldOwner(constants.ReclaimablePodsMgr))
1101+
var (
1102+
wlCopy *kueue.Workload
1103+
err error
1104+
)
1105+
1106+
if features.Enabled(features.WorkloadRequestUseMergePatch) {
1107+
wlCopy = w.DeepCopy()
1108+
err = clientutil.PatchStatus(ctx, c, wlCopy, func() (bool, error) {
1109+
wlCopy.Status.ReclaimablePods = reclaimablePods
1110+
return true, nil
1111+
})
1112+
} else {
1113+
wlCopy = BaseSSAWorkload(w, true)
1114+
wlCopy.Status.ReclaimablePods = reclaimablePods
1115+
err = c.Status().Patch(ctx, wlCopy, client.Apply, client.FieldOwner(constants.ReclaimablePodsMgr))
1116+
}
1117+
if err != nil {
1118+
return err
1119+
}
1120+
1121+
wlCopy.DeepCopyInto(w)
1122+
return nil
11041123
}
11051124

11061125
// ReclaimablePodsAreEqual checks if two Reclaimable pods are semantically equal

‎test/integration/singlecluster/controller/core/clusterqueue_controller_test.go‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -513,8 +513,9 @@ var _ = ginkgo.Describe("ClusterQueue controller", ginkgo.Ordered, ginkgo.Contin
513513
}, util.Timeout, util.Interval).Should(gomega.Succeed())
514514

515515
ginkgo.By("Mark two workers as reclaimable", func() {
516-
gomega.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{{Name: "workers", Count: 2}})).To(gomega.Succeed())
517-
516+
gomega.Eventually(func(g gomega.Gomega) {
517+
g.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{{Name: "workers", Count: 2}})).To(gomega.Succeed())
518+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
518519
util.ExpectReservingActiveWorkloadsMetric(clusterQueue, 1)
519520
util.ExpectLQReservingActiveWorkloadsMetric(localQueue, 1)
520521
gomega.Eventually(func(g gomega.Gomega) {
@@ -552,8 +553,9 @@ var _ = ginkgo.Describe("ClusterQueue controller", ginkgo.Ordered, ginkgo.Contin
552553
})
553554

554555
ginkgo.By("Mark all workers and a driver as reclaimable", func() {
555-
gomega.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{{Name: "workers", Count: 5}, {Name: "driver", Count: 1}})).To(gomega.Succeed())
556-
556+
gomega.Eventually(func(g gomega.Gomega) {
557+
g.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{{Name: "workers", Count: 5}, {Name: "driver", Count: 1}})).To(gomega.Succeed())
558+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
557559
util.ExpectReservingActiveWorkloadsMetric(clusterQueue, 1)
558560
util.ExpectLQReservingActiveWorkloadsMetric(localQueue, 1)
559561
gomega.Eventually(func(g gomega.Gomega) {
@@ -940,7 +942,9 @@ var _ = ginkgo.Describe("ClusterQueue controller", ginkgo.Ordered, ginkgo.Contin
940942
})
941943

942944
ginkgo.By("Marking two workers as reclaimable", func() {
943-
gomega.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{{Name: "workers", Count: 2}})).To(gomega.Succeed())
945+
gomega.Eventually(func(g gomega.Gomega) {
946+
g.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{{Name: "workers", Count: 2}})).To(gomega.Succeed())
947+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
944948
})
945949

946950
ginkgo.By("Validating CQ status hasn't changed", func() {

‎test/integration/singlecluster/scheduler/scheduler_test.go‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -470,8 +470,9 @@ var _ = ginkgo.Describe("Scheduler", func() {
470470
})
471471

472472
ginkgo.By("Reclaim one pod from the first workload", func() {
473-
gomega.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, firstWl, []kueue.ReclaimablePod{{Name: "third", Count: 1}})).To(gomega.Succeed())
474-
473+
gomega.Eventually(func(g gomega.Gomega) {
474+
g.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, firstWl, []kueue.ReclaimablePod{{Name: "third", Count: 1}})).To(gomega.Succeed())
475+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
475476
util.ExpectPendingWorkloadsMetric(preemptionClusterQ, 0, 0)
476477
util.ExpectAdmittedWorkloadsTotalMetric(preemptionClusterQ, "", 1)
477478
})
@@ -494,8 +495,9 @@ var _ = ginkgo.Describe("Scheduler", func() {
494495
})
495496

496497
ginkgo.By("Reclaim two pods from the second workload so that the first workload is resumed", func() {
497-
gomega.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, secondWl, []kueue.ReclaimablePod{{Name: "first", Count: 1}, {Name: "second", Count: 1}})).To(gomega.Succeed())
498-
498+
gomega.Eventually(func(g gomega.Gomega) {
499+
g.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, secondWl, []kueue.ReclaimablePod{{Name: "first", Count: 1}, {Name: "second", Count: 1}})).To(gomega.Succeed())
500+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
499501
util.ExpectWorkloadsToBeAdmitted(ctx, k8sClient, firstWl, secondWl)
500502
util.ExpectPendingWorkloadsMetric(preemptionClusterQ, 0, 0)
501503
util.ExpectReservingActiveWorkloadsMetric(preemptionClusterQ, 2)
@@ -531,7 +533,9 @@ var _ = ginkgo.Describe("Scheduler", func() {
531533
})
532534

533535
ginkgo.By("Reclaim one pod from the first workload and admitting the second one", func() {
534-
gomega.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, firstWl, []kueue.ReclaimablePod{{Name: "third", Count: 1}})).To(gomega.Succeed())
536+
gomega.Eventually(func(g gomega.Gomega) {
537+
g.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, firstWl, []kueue.ReclaimablePod{{Name: "third", Count: 1}})).To(gomega.Succeed())
538+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
535539
util.ExpectPendingWorkloadsMetric(preemptionClusterQ, 0, 0)
536540
util.ExpectAdmittedWorkloadsTotalMetric(preemptionClusterQ, "", 2)
537541
})
@@ -567,7 +571,9 @@ var _ = ginkgo.Describe("Scheduler", func() {
567571
})
568572

569573
ginkgo.By("Reclaim one pod from the first workload and admitting the second one", func() {
570-
gomega.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, firstWl, []kueue.ReclaimablePod{{Name: "third", Count: 1}})).To(gomega.Succeed())
574+
gomega.Eventually(func(g gomega.Gomega) {
575+
g.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, firstWl, []kueue.ReclaimablePod{{Name: "third", Count: 1}})).To(gomega.Succeed())
576+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
571577
util.ExpectWorkloadsToBeAdmitted(ctx, k8sClient, firstWl)
572578
util.ExpectPendingWorkloadsMetric(preemptionClusterQ, 0, 1)
573579
util.ExpectAdmittedWorkloadsTotalMetric(preemptionClusterQ, "", 1)
@@ -734,8 +740,9 @@ var _ = ginkgo.Describe("Scheduler", func() {
734740
util.ExpectAdmittedWorkloadsTotalMetric(prodClusterQ, "", 0)
735741
})
736742
ginkgo.By("Mark one pod as reclaimable", func() {
737-
gomega.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{{Name: kueue.DefaultPodSetName, Count: 1}})).To(gomega.Succeed())
738-
743+
gomega.Eventually(func(g gomega.Gomega) {
744+
g.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{{Name: kueue.DefaultPodSetName, Count: 1}})).To(gomega.Succeed())
745+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
739746
util.ExpectWorkloadsToHaveQuotaReservation(ctx, k8sClient, prodClusterQ.Name, wl)
740747
util.ExpectPendingWorkloadsMetric(prodClusterQ, 0, 0)
741748
util.ExpectAdmittedWorkloadsTotalMetric(prodClusterQ, "", 1)

‎test/integration/singlecluster/webhook/core/workload_test.go‎

Lines changed: 31 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -528,11 +528,13 @@ var _ = ginkgo.Describe("Workload validating webhook", ginkgo.Ordered, func() {
528528
Obj()
529529
util.MustCreate(ctx, k8sClient, wl)
530530

531-
err := workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{
532-
{Name: "ps1", Count: 4},
533-
{Name: "ps2", Count: 1},
534-
})
535-
gomega.Expect(err).Should(utiltesting.BeForbiddenError())
531+
gomega.Eventually(func(g gomega.Gomega) {
532+
err := workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{
533+
{Name: "ps1", Count: 4},
534+
{Name: "ps2", Count: 1},
535+
})
536+
g.Expect(err).Should(utiltesting.BeForbiddenError())
537+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
536538
})
537539
})
538540

@@ -1139,9 +1141,11 @@ var _ = ginkgo.Describe("Workload validating webhook", ginkgo.Ordered, func() {
11391141
).
11401142
Obj()
11411143
util.MustCreate(ctx, k8sClient, wl)
1142-
gomega.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{
1143-
{Name: "ps1", Count: 1},
1144-
})).Should(gomega.Succeed())
1144+
gomega.Eventually(func(g gomega.Gomega) {
1145+
g.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{
1146+
{Name: "ps1", Count: 1},
1147+
})).Should(gomega.Succeed())
1148+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
11451149

11461150
util.SetQuotaReservation(ctx, k8sClient, client.ObjectKeyFromObject(wl),
11471151
utiltestingapi.MakeAdmission("cluster-queue").
@@ -1151,11 +1155,13 @@ var _ = ginkgo.Describe("Workload validating webhook", ginkgo.Ordered, func() {
11511155
Obj())
11521156

11531157
ginkgo.By("Updating reclaimable pods")
1154-
err := workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{
1155-
{Name: "ps1", Count: 2},
1156-
{Name: "ps2", Count: 1},
1157-
})
1158-
gomega.Expect(err).Should(gomega.Succeed())
1158+
gomega.Eventually(func(g gomega.Gomega) {
1159+
err := workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{
1160+
{Name: "ps1", Count: 2},
1161+
{Name: "ps2", Count: 1},
1162+
})
1163+
g.Expect(err).Should(gomega.Succeed())
1164+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
11591165
})
11601166

11611167
ginkgo.It("reclaimable pod count cannot change down", func() {
@@ -1167,10 +1173,12 @@ var _ = ginkgo.Describe("Workload validating webhook", ginkgo.Ordered, func() {
11671173
).
11681174
Obj()
11691175
util.MustCreate(ctx, k8sClient, wl)
1170-
gomega.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{
1171-
{Name: "ps1", Count: 2},
1172-
{Name: "ps2", Count: 1},
1173-
})).Should(gomega.Succeed())
1176+
gomega.Eventually(func(g gomega.Gomega) {
1177+
g.Expect(workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{
1178+
{Name: "ps1", Count: 2},
1179+
{Name: "ps2", Count: 1},
1180+
})).Should(gomega.Succeed())
1181+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
11741182

11751183
util.SetQuotaReservation(ctx, k8sClient, client.ObjectKeyFromObject(wl),
11761184
utiltestingapi.MakeAdmission("cluster-queue").
@@ -1180,10 +1188,12 @@ var _ = ginkgo.Describe("Workload validating webhook", ginkgo.Ordered, func() {
11801188
Obj())
11811189

11821190
ginkgo.By("Updating reclaimable pods")
1183-
err := workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{
1184-
{Name: "ps1", Count: 1},
1185-
})
1186-
gomega.Expect(err).Should(utiltesting.BeForbiddenError())
1191+
gomega.Eventually(func(g gomega.Gomega) {
1192+
err := workload.UpdateReclaimablePods(ctx, k8sClient, wl, []kueue.ReclaimablePod{
1193+
{Name: "ps1", Count: 1},
1194+
})
1195+
g.Expect(err).Should(utiltesting.BeForbiddenError())
1196+
}, util.Timeout, util.Interval).Should(gomega.Succeed())
11871197
})
11881198

11891199
ginkgo.It("podSetUpdates should be immutable when state is ready", func() {

0 commit comments

Comments
 (0)