Skip to content

Commit 46b74af

Browse files
olekzablolekzabl
committed
Addressing comments
1 parent 75179cd commit 46b74af

File tree

5 files changed

+31
-53
lines changed

5 files changed

+31
-53
lines changed

‎pkg/controller/tas/node_failure_controller.go‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -191,10 +191,9 @@ func (r *nodeFailureReconciler) getWorkloadsOnNode(ctx context.Context, nodeName
191191
if !utiltas.IsLowestLevelHostname(topologyAssignment.Levels) {
192192
continue
193193
}
194-
for value := range utiltas.ValuesAtLevel(topologyAssignment, len(topologyAssignment.Levels)-1) {
194+
for value := range utiltas.LowestLevelValues(topologyAssignment) {
195195
if value == nodeName {
196196
tasWorkloadsOnNode.Insert(types.NamespacedName{Name: wl.Name, Namespace: wl.Namespace})
197-
break
198197
}
199198
}
200199
}

‎pkg/util/tas/tas_assignment.go‎

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -52,10 +52,7 @@ func countAtIndex(slice kueue.TopologyAssignmentSlice, idx int) int32 {
5252
return slice.PodCounts.Individual[idx]
5353
}
5454

55-
func ValuesAtLevel(ta *kueue.TopologyAssignment, levelIdx int) iter.Seq[string] {
56-
if ta == nil {
57-
return nil
58-
}
55+
func valuesAtLevel(ta *kueue.TopologyAssignment, levelIdx int) iter.Seq[string] {
5956
return func(yield func(string) bool) {
6057
for _, slice := range ta.Slices {
6158
values := slice.ValuesPerLevel[levelIdx]
@@ -68,6 +65,13 @@ func ValuesAtLevel(ta *kueue.TopologyAssignment, levelIdx int) iter.Seq[string]
6865
}
6966
}
7067

68+
func LowestLevelValues(ta *kueue.TopologyAssignment) iter.Seq[string] {
69+
if ta == nil {
70+
return nil
71+
}
72+
return valuesAtLevel(ta, len(ta.Levels)-1)
73+
}
74+
7175
func PodCounts(ta *kueue.TopologyAssignment) iter.Seq[int32] {
7276
if ta == nil {
7377
return nil
@@ -132,23 +136,17 @@ func fillSingleCompactSliceValues(
132136
) {
133137
var prefix, suffix string
134138
var maxLen, minLen, count int
135-
start := true
136139
for s := range inputProvider() {
137140
count++
138-
if start {
141+
if count == 1 {
139142
prefix = s
140143
suffix = s
141144
maxLen = len(s)
142145
minLen = len(s)
143-
start = false
144146
} else {
145147
n := len(s)
146-
if n < minLen {
147-
minLen = n
148-
}
149-
if n > maxLen {
150-
maxLen = n
151-
}
148+
minLen = min(minLen, n)
149+
maxLen = max(maxLen, n)
152150
if n < len(prefix) {
153151
prefix = prefix[:n]
154152
}

‎pkg/util/tas/tas_assignment_test.go‎

Lines changed: 14 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ import (
2727
kueue "sigs.k8s.io/kueue/apis/kueue/v1beta2"
2828
)
2929

30-
type TestCase struct {
30+
type testCase struct {
3131
name string
3232
internal *TopologyAssignment
3333
v1beta2 *kueue.TopologyAssignment
@@ -36,7 +36,7 @@ type TestCase struct {
3636
}
3737

3838
// bothWaysTestCases expect that internal <-> v1beta2 maps in both ways.
39-
var bothWaysTestCases = []TestCase{
39+
var bothWaysTestCases = []testCase{
4040
{
4141
name: "empty",
4242
internal: &TopologyAssignment{
@@ -421,7 +421,7 @@ var bothWaysTestCases = []TestCase{
421421

422422
// oneWayTestCases expect that v1beta2 -> internal (via InternalFrom).
423423
// (these v1beta2 values cannot be produced by V1Beta2From as of now - though this may change in the future)
424-
var oneWayTestCases = []TestCase{
424+
var oneWayTestCases = []testCase{
425425
{
426426
name: "multiple slices, no prefixes/suffixes or universal values",
427427
internal: &TopologyAssignment{
@@ -655,7 +655,7 @@ func TestInternalSeqFromV1Beta2_forNil(t *testing.T) {
655655
}
656656
}
657657

658-
func TestInternalSeqFrmV1Beta2_iteratorStops(t *testing.T) {
658+
func TestInternalSeqFromV1Beta2_iteratorStops(t *testing.T) {
659659
for range InternalSeqFrom(twoDomains) {
660660
// Break the loop prematurely.
661661
// If the iterator isn't smart enough to stop, this will panic.
@@ -709,7 +709,7 @@ func TestTotalDomainCount_forNil(t *testing.T) {
709709
}
710710
}
711711

712-
func TestValuesAtLevel(t *testing.T) {
712+
func TestLowestLevel(t *testing.T) {
713713
v1beta2 := &kueue.TopologyAssignment{
714714
Levels: []string{"a", "b"},
715715
Slices: []kueue.TopologyAssignmentSlice{
@@ -767,42 +767,23 @@ func TestValuesAtLevel(t *testing.T) {
767767
},
768768
},
769769
}
770-
testCases := []struct {
771-
name string
772-
levelIdx int
773-
want []string
774-
}{
775-
{
776-
name: "at level 0",
777-
levelIdx: 0,
778-
want: []string{"a1", "a1", "a2", "a3", "a4", "a10", "a10"},
779-
},
780-
{
781-
name: "at level 1",
782-
levelIdx: 1,
783-
want: []string{"b1-s", "b2-s", "x-t", "y-t", "z-t", "b10", "b10"},
784-
},
785-
}
786-
for _, tc := range testCases {
787-
t.Run(tc.name, func(t *testing.T) {
788-
seq := ValuesAtLevel(v1beta2, tc.levelIdx)
789-
got := slices.Collect(seq)
790-
if diff := cmp.Diff(tc.want, got, cmpopts.EquateEmpty()); diff != "" {
791-
t.Errorf("unexpected result (-want,+got):\n%s", diff)
792-
}
793-
})
770+
want := []string{"b1-s", "b2-s", "x-t", "y-t", "z-t", "b10", "b10"}
771+
seq := LowestLevelValues(v1beta2)
772+
got := slices.Collect(seq)
773+
if diff := cmp.Diff(want, got, cmpopts.EquateEmpty()); diff != "" {
774+
t.Errorf("unexpected result (-want,+got):\n%s", diff)
794775
}
795776
}
796777

797-
func TestValuesAtLevel_forNil(t *testing.T) {
798-
got := ValuesAtLevel(nil, 0)
778+
func TestLowestLevelValues_forNil(t *testing.T) {
779+
got := LowestLevelValues(nil)
799780
if got != nil {
800781
t.Errorf("unexpected result for nil: %+v", got)
801782
}
802783
}
803784

804-
func TestValuesAtLevel_iteratorStops(t *testing.T) {
805-
for range ValuesAtLevel(twoDomains, 0) {
785+
func TestLowestLevelValues_iteratorStops(t *testing.T) {
786+
for range LowestLevelValues(twoDomains) {
806787
// Break the loop prematurely.
807788
// If the iterator isn't smart enough to stop, this will panic.
808789
break

‎pkg/workload/workload.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1183,7 +1183,7 @@ func HasTopologyAssignmentWithUnhealthyNode(w *kueue.Workload) bool {
11831183
if psa.TopologyAssignment == nil {
11841184
continue
11851185
}
1186-
for value := range tas.ValuesAtLevel(psa.TopologyAssignment, len(psa.TopologyAssignment.Levels)-1) {
1186+
for value := range tas.LowestLevelValues(psa.TopologyAssignment) {
11871187
if HasUnhealthyNode(w, value) {
11881188
return true
11891189
}

‎test/integration/singlecluster/controller/jobs/pod/pod_controller_test.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2488,7 +2488,7 @@ var _ = ginkgo.Describe("Pod controller with TASReplaceNodeOnPodTermination", gi
24882488
ginkgo.By("verify the workload is admitted", func() {
24892489
util.ExpectWorkloadsToBeAdmitted(ctx, k8sClient, wl)
24902490
gomega.Expect(k8sClient.Get(ctx, wlKey, wl)).To(gomega.Succeed())
2491-
nodeNames := slices.Collect(tas.ValuesAtLevel(wl.Status.Admission.PodSetAssignments[0].TopologyAssignment, 0))
2491+
nodeNames := slices.Collect(tas.LowestLevelValues(wl.Status.Admission.PodSetAssignments[0].TopologyAssignment))
24922492
gomega.Expect(nodeNames).To(gomega.HaveLen(1))
24932493
nodeName = nodeNames[0]
24942494
gomega.Expect(nodeName).To(gomega.Or(gomega.Equal("x1"), gomega.Equal("x3")))
@@ -2534,7 +2534,7 @@ var _ = ginkgo.Describe("Pod controller with TASReplaceNodeOnPodTermination", gi
25342534
ginkgo.By("verify the workload is assigned a new node", func() {
25352535
gomega.Eventually(func(g gomega.Gomega) string {
25362536
gomega.Expect(k8sClient.Get(ctx, wlKey, wl)).To(gomega.Succeed())
2537-
nodeNames := slices.Collect(tas.ValuesAtLevel(wl.Status.Admission.PodSetAssignments[0].TopologyAssignment, 0))
2537+
nodeNames := slices.Collect(tas.LowestLevelValues(wl.Status.Admission.PodSetAssignments[0].TopologyAssignment))
25382538
gomega.Expect(nodeNames).To(gomega.HaveLen(1))
25392539
return nodeNames[0]
25402540
}, util.Timeout, util.Interval).ShouldNot(gomega.Equal(nodeName))
@@ -2572,7 +2572,7 @@ var _ = ginkgo.Describe("Pod controller with TASReplaceNodeOnPodTermination", gi
25722572
ginkgo.By("verify the workload is admitted", func() {
25732573
util.ExpectWorkloadsToBeAdmitted(ctx, k8sClient, wl)
25742574
gomega.Expect(k8sClient.Get(ctx, wlKey, wl)).To(gomega.Succeed())
2575-
nodeNames := slices.Collect(tas.ValuesAtLevel(wl.Status.Admission.PodSetAssignments[0].TopologyAssignment, 0))
2575+
nodeNames := slices.Collect(tas.LowestLevelValues(wl.Status.Admission.PodSetAssignments[0].TopologyAssignment))
25762576
gomega.Expect(nodeNames).To(gomega.HaveLen(1))
25772577
nodeName = nodeNames[0]
25782578
gomega.Expect(nodeName).To(gomega.Or(gomega.Equal("x1"), gomega.Equal("x3")))
@@ -2618,7 +2618,7 @@ var _ = ginkgo.Describe("Pod controller with TASReplaceNodeOnPodTermination", gi
26182618
ginkgo.By("verify the workload is assigned a new node", func() {
26192619
gomega.Eventually(func(g gomega.Gomega) string {
26202620
gomega.Expect(k8sClient.Get(ctx, wlKey, wl)).To(gomega.Succeed())
2621-
nodeNames := slices.Collect(tas.ValuesAtLevel(wl.Status.Admission.PodSetAssignments[0].TopologyAssignment, 0))
2621+
nodeNames := slices.Collect(tas.LowestLevelValues(wl.Status.Admission.PodSetAssignments[0].TopologyAssignment))
26222622
gomega.Expect(nodeNames).To(gomega.HaveLen(1))
26232623
return nodeNames[0]
26242624
}, util.Timeout, util.Interval).ShouldNot(gomega.Equal(nodeName))

0 commit comments

Comments
 (0)