Skip to content

Commit abf65b0

Browse files
mbobrovskyimbobrovskyi
committed
Add finalizer only if AdmissionCheck exists and Workload has Admission with TopologyAssignment.
1 parent 4abd683 commit abf65b0

File tree

3 files changed

+142
-86
lines changed

3 files changed

+142
-86
lines changed

‎slice/internal/controller/workload_controller.go

Lines changed: 31 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -114,17 +114,6 @@ func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
114114
return ctrl.Result{}, nil
115115
}
116116

117-
if controllerutil.AddFinalizer(wl, SliceControllerName) {
118-
if err = r.client.Update(ctx, wl); err != nil {
119-
if !apierrors.IsNotFound(err) {
120-
log.Error(err, "Failed to add finalizer")
121-
}
122-
return ctrl.Result{}, client.IgnoreNotFound(err)
123-
}
124-
log.V(5).Info("Added finalizer")
125-
return ctrl.Result{}, nil
126-
}
127-
128117
ac, err := r.sliceAC(ctx, wl)
129118
if err != nil {
130119
return reconcile.Result{}, err
@@ -137,6 +126,21 @@ func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
137126
log = log.WithValues("admissionCheck", ac.Name)
138127
ctrl.LoggerInto(ctx, log)
139128

129+
if !r.shouldAddFinalizer(wl) {
130+
return ctrl.Result{}, nil
131+
}
132+
133+
if controllerutil.AddFinalizer(wl, SliceControllerName) {
134+
if err = r.client.Update(ctx, wl); err != nil {
135+
if !apierrors.IsNotFound(err) {
136+
log.Error(err, "Failed to add finalizer")
137+
}
138+
return ctrl.Result{}, client.IgnoreNotFound(err)
139+
}
140+
log.V(5).Info("Added finalizer")
141+
return ctrl.Result{}, nil
142+
}
143+
140144
if ac.State == kueue.CheckStateReady {
141145
log.V(5).Info("Admission check is ready — nothing to do")
142146
return reconcile.Result{}, nil
@@ -160,6 +164,19 @@ func (r *WorkloadReconciler) shouldFinalize(wl *kueue.Workload) bool {
160164
return !wl.DeletionTimestamp.IsZero() || workload.IsFinished(wl) || workload.IsEvicted(wl) || !workload.IsActive(wl)
161165
}
162166

167+
func (r *WorkloadReconciler) shouldAddFinalizer(wl *kueue.Workload) bool {
168+
if !workload.HasQuotaReservation(wl) || wl.Status.Admission == nil {
169+
return false
170+
}
171+
for _, psa := range wl.Status.Admission.PodSetAssignments {
172+
// Only workloads with a TopologyAssignment should be processed.
173+
if psa.TopologyAssignment != nil {
174+
return true
175+
}
176+
}
177+
return false
178+
}
179+
163180
func (r *WorkloadReconciler) sliceAC(ctx context.Context, wl *kueue.Workload) (*kueue.AdmissionCheckState, error) {
164181
relevantChecks, err := admissioncheck.FilterForController(ctx, r.client, wl.Status.AdmissionChecks, SliceControllerName)
165182
if err != nil {
@@ -192,9 +209,6 @@ func (r *WorkloadReconciler) newSlice(wl *kueue.Workload) (*v1alpha1.Slice, erro
192209
if err := controllerutil.SetControllerReference(wl, slice, r.client.Scheme()); err != nil {
193210
return nil, err
194211
}
195-
if wl.Status.Admission == nil || wl.Status.Admission.PodSetAssignments == nil {
196-
return slice, nil
197-
}
198212

199213
nodeSelectors := sets.New[string]()
200214
for _, psa := range wl.Status.Admission.PodSetAssignments {
@@ -209,12 +223,14 @@ func (r *WorkloadReconciler) newSlice(wl *kueue.Workload) (*v1alpha1.Slice, erro
209223
}
210224

211225
func (r *WorkloadReconciler) createSlice(ctx context.Context, wl *kueue.Workload, ac *kueue.AdmissionCheckState) error {
226+
log := ctrl.LoggerFrom(ctx)
227+
212228
slice, err := r.newSlice(wl)
213229
if err != nil {
214230
return err
215231
}
216232

217-
log := ctrl.LoggerFrom(ctx).WithValues("slice", klog.KObj(slice))
233+
log = log.WithValues("slice", klog.KObj(slice))
218234

219235
err = r.client.Create(ctx, slice)
220236
if err != nil {

‎slice/internal/controller/workload_controller_test.go

Lines changed: 97 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,31 @@ func TestWorkloadReconciler(t *testing.T) {
6969
LastTransitionTime: metav1.NewTime(now),
7070
Message: "",
7171
})
72+
baseWorkloadWrapperWithAdmission := baseWorkloadWrapper.Clone().
73+
ReserveQuota(
74+
&kueue.Admission{
75+
PodSetAssignments: []kueue.PodSetAssignment{
76+
utiltesting.MakePodSetAssignment("psa1").
77+
TopologyAssignment(nil, []kueue.TopologyDomainAssignment{
78+
{
79+
Values: []string{"domain1", "domain2"},
80+
Count: 2,
81+
},
82+
}).Obj(),
83+
utiltesting.MakePodSetAssignment("psa2").
84+
TopologyAssignment(nil, []kueue.TopologyDomainAssignment{
85+
{
86+
Values: []string{"domain2", "domain3"},
87+
Count: 2,
88+
},
89+
}).
90+
Obj(),
91+
},
92+
}, now,
93+
)
7294
baseSliceWrapper := utiltesting.MakeSliceWrapper(baseWorkloadName, corev1.NamespaceDefault).
73-
ControllerReference(kueue.GroupVersion.WithKind("Workload"), baseWorkloadName, baseWorkloadName)
95+
ControllerReference(kueue.GroupVersion.WithKind("Workload"), baseWorkloadName, baseWorkloadName).
96+
NodeSelector(map[string][]string{TPUReservationSubblockLabel: {"domain1", "domain2", "domain3"}})
7497

7598
cases := map[string]struct {
7699
interceptorFuncsCreate func(ctx context.Context, client client.WithWatch, obj client.Object, opts ...client.CreateOption) error
@@ -82,9 +105,13 @@ func TestWorkloadReconciler(t *testing.T) {
82105
wantEvents []utiltesting.EventRecord
83106
}{
84107
"should skip reconciliation because the Workload was not found": {
85-
request: types.NamespacedName{Name: "other-workload", Namespace: corev1.NamespaceDefault},
86-
objs: []client.Object{baseWorkloadWrapper.DeepCopy()},
87-
wantWorkloads: []kueue.Workload{*baseWorkloadWrapper.DeepCopy()},
108+
request: types.NamespacedName{Name: "other-workload", Namespace: corev1.NamespaceDefault},
109+
objs: []client.Object{
110+
baseWorkloadWrapper.Clone().Finalizers(SliceControllerName).DeletionTimestamp(now).Obj(),
111+
},
112+
wantWorkloads: []kueue.Workload{
113+
*baseWorkloadWrapper.Clone().Finalizers(SliceControllerName).DeletionTimestamp(now).Obj(),
114+
},
88115
},
89116
"should delete the finalizer because the Workload has a DeletionTimestamp": {
90117
request: baseRequest,
@@ -117,24 +144,69 @@ func TestWorkloadReconciler(t *testing.T) {
117144
},
118145
wantWorkloads: []kueue.Workload{*baseWorkloadWrapper.Clone().Active(false).Obj()},
119146
},
120-
"should add finalizer": {
147+
"shouldn't add finalizer because there’s no Admission": {
121148
request: baseRequest,
122149
objs: []client.Object{
150+
baseAdmissionCheckWrapper.DeepCopy(),
123151
baseWorkloadWrapper.DeepCopy(),
124152
},
153+
wantWorkloads: []kueue.Workload{*baseWorkloadWrapper.DeepCopy()},
154+
},
155+
"shouldn't add finalizer because there’s no TopologyAssignment": {
156+
request: baseRequest,
157+
objs: []client.Object{
158+
baseAdmissionCheckWrapper.DeepCopy(),
159+
baseWorkloadWrapper.Clone().
160+
ReserveQuota(
161+
&kueue.Admission{
162+
PodSetAssignments: []kueue.PodSetAssignment{
163+
utiltesting.MakePodSetAssignment("psa1").Obj(),
164+
},
165+
}, now,
166+
).
167+
Obj(),
168+
},
125169
wantWorkloads: []kueue.Workload{
126-
*baseWorkloadWrapper.Clone().Finalizers(SliceControllerName).
170+
*baseWorkloadWrapper.Clone().
171+
ReserveQuota(
172+
&kueue.Admission{
173+
PodSetAssignments: []kueue.PodSetAssignment{
174+
utiltesting.MakePodSetAssignment("psa1").Obj(),
175+
},
176+
}, now,
177+
).
178+
Obj(),
179+
},
180+
},
181+
"shouldn't add finalizer because there’s no AdmissionCheck": {
182+
request: baseRequest,
183+
objs: []client.Object{
184+
baseWorkloadWrapperWithAdmission.DeepCopy(),
185+
},
186+
wantWorkloads: []kueue.Workload{
187+
*baseWorkloadWrapperWithAdmission.Clone().Obj(),
188+
},
189+
},
190+
"should add finalizer": {
191+
request: baseRequest,
192+
objs: []client.Object{
193+
baseAdmissionCheckWrapper.DeepCopy(),
194+
baseWorkloadWrapperWithAdmission.DeepCopy(),
195+
},
196+
wantWorkloads: []kueue.Workload{
197+
*baseWorkloadWrapperWithAdmission.Clone().
198+
Finalizers(SliceControllerName).
127199
Obj(),
128200
},
129201
},
130202
"should create a Slice": {
131203
request: baseRequest,
132204
objs: []client.Object{
133205
baseAdmissionCheckWrapper.DeepCopy(),
134-
baseWorkloadWrapper.Finalizers(SliceControllerName).DeepCopy(),
206+
baseWorkloadWrapperWithAdmission.Finalizers(SliceControllerName).DeepCopy(),
135207
},
136208
wantWorkloads: []kueue.Workload{
137-
*baseWorkloadWrapper.Clone().
209+
*baseWorkloadWrapperWithAdmission.Clone().
138210
Finalizers(SliceControllerName).
139211
AdmissionCheck(kueue.AdmissionCheckState{
140212
Name: kueue.AdmissionCheckReference(baseAdmissionCheckName),
@@ -158,56 +230,21 @@ func TestWorkloadReconciler(t *testing.T) {
158230
request: baseRequest,
159231
objs: []client.Object{
160232
baseAdmissionCheckWrapper.DeepCopy(),
161-
baseWorkloadWrapper.Clone().
162-
Finalizers(SliceControllerName).
163-
PodSetAssignments(utiltesting.MakePodSetAssignment("psa1").
164-
TopologyAssignment(nil, []kueue.TopologyDomainAssignment{
165-
{
166-
Values: []string{"domain1", "domain2"},
167-
Count: 2,
168-
},
169-
}).Obj(),
170-
utiltesting.MakePodSetAssignment("psa2").
171-
TopologyAssignment(nil, []kueue.TopologyDomainAssignment{
172-
{
173-
Values: []string{"domain2", "domain3"},
174-
Count: 2,
175-
},
176-
}).
177-
Obj(),
178-
).Obj(),
233+
baseWorkloadWrapperWithAdmission.Clone().Finalizers(SliceControllerName).Obj(),
179234
},
180235
wantWorkloads: []kueue.Workload{
181-
*baseWorkloadWrapper.Clone().
236+
*baseWorkloadWrapperWithAdmission.Clone().
182237
Finalizers(SliceControllerName).
183238
AdmissionCheck(kueue.AdmissionCheckState{
184239
Name: kueue.AdmissionCheckReference(baseAdmissionCheckName),
185240
State: kueue.CheckStatePending,
186241
LastTransitionTime: metav1.NewTime(now),
187242
Message: fmt.Sprintf(`The Slice "%s" has been created`, baseWorkloadName),
188243
}).
189-
PodSetAssignments(utiltesting.MakePodSetAssignment("psa1").
190-
TopologyAssignment(nil, []kueue.TopologyDomainAssignment{
191-
{
192-
Values: []string{"domain1", "domain2"},
193-
Count: 2,
194-
},
195-
}).Obj(),
196-
utiltesting.MakePodSetAssignment("psa2").
197-
TopologyAssignment(nil, []kueue.TopologyDomainAssignment{
198-
{
199-
Values: []string{"domain2", "domain3"},
200-
Count: 2,
201-
},
202-
}).
203-
Obj(),
204-
).
205244
Obj(),
206245
},
207246
wantSlices: []slice.Slice{
208-
*baseSliceWrapper.Clone().
209-
NodeSelector(map[string][]string{TPUReservationSubblockLabel: {"domain1", "domain2", "domain3"}}).
210-
Obj(),
247+
*baseSliceWrapper.DeepCopy(),
211248
},
212249
wantEvents: []utiltesting.EventRecord{
213250
{
@@ -228,10 +265,10 @@ func TestWorkloadReconciler(t *testing.T) {
228265
request: baseRequest,
229266
objs: []client.Object{
230267
baseAdmissionCheckWrapper.DeepCopy(),
231-
baseWorkloadWrapper.Finalizers(SliceControllerName).DeepCopy(),
268+
baseWorkloadWrapperWithAdmission.Finalizers(SliceControllerName).DeepCopy(),
232269
},
233270
wantWorkloads: []kueue.Workload{
234-
*baseWorkloadWrapper.Clone().
271+
*baseWorkloadWrapperWithAdmission.Clone().
235272
Finalizers(SliceControllerName).
236273
AdmissionCheck(kueue.AdmissionCheckState{
237274
Name: kueue.AdmissionCheckReference(baseAdmissionCheckName),
@@ -255,11 +292,11 @@ func TestWorkloadReconciler(t *testing.T) {
255292
request: baseRequest,
256293
objs: []client.Object{
257294
baseAdmissionCheckWrapper.DeepCopy(),
258-
baseWorkloadWrapper.Finalizers(SliceControllerName).DeepCopy(),
295+
baseWorkloadWrapperWithAdmission.Finalizers(SliceControllerName).DeepCopy(),
259296
baseSliceWrapper.Clone().Forming().Obj(),
260297
},
261298
wantWorkloads: []kueue.Workload{
262-
*baseWorkloadWrapper.Clone().
299+
*baseWorkloadWrapperWithAdmission.Clone().
263300
Finalizers(SliceControllerName).
264301
AdmissionCheck(kueue.AdmissionCheckState{
265302
Name: kueue.AdmissionCheckReference(baseAdmissionCheckName),
@@ -275,11 +312,11 @@ func TestWorkloadReconciler(t *testing.T) {
275312
request: baseRequest,
276313
objs: []client.Object{
277314
baseAdmissionCheckWrapper.DeepCopy(),
278-
baseWorkloadWrapper.Finalizers(SliceControllerName).DeepCopy(),
315+
baseWorkloadWrapperWithAdmission.Finalizers(SliceControllerName).DeepCopy(),
279316
baseSliceWrapper.Clone().Ready().Obj(),
280317
},
281318
wantWorkloads: []kueue.Workload{
282-
*baseWorkloadWrapper.Clone().
319+
*baseWorkloadWrapperWithAdmission.Clone().
283320
Finalizers(SliceControllerName).
284321
AdmissionCheck(kueue.AdmissionCheckState{
285322
Name: kueue.AdmissionCheckReference(baseAdmissionCheckName),
@@ -303,11 +340,11 @@ func TestWorkloadReconciler(t *testing.T) {
303340
request: baseRequest,
304341
objs: []client.Object{
305342
baseAdmissionCheckWrapper.DeepCopy(),
306-
baseWorkloadWrapper.Finalizers(SliceControllerName).DeepCopy(),
343+
baseWorkloadWrapperWithAdmission.Finalizers(SliceControllerName).DeepCopy(),
307344
baseSliceWrapper.Clone().Degraded().Obj(),
308345
},
309346
wantWorkloads: []kueue.Workload{
310-
*baseWorkloadWrapper.Clone().
347+
*baseWorkloadWrapperWithAdmission.Clone().
311348
Finalizers(SliceControllerName).
312349
AdmissionCheck(kueue.AdmissionCheckState{
313350
Name: kueue.AdmissionCheckReference(baseAdmissionCheckName),
@@ -331,11 +368,11 @@ func TestWorkloadReconciler(t *testing.T) {
331368
request: baseRequest,
332369
objs: []client.Object{
333370
baseAdmissionCheckWrapper.DeepCopy(),
334-
baseWorkloadWrapper.Finalizers(SliceControllerName).DeepCopy(),
371+
baseWorkloadWrapperWithAdmission.Finalizers(SliceControllerName).DeepCopy(),
335372
baseSliceWrapper.Clone().Deformed().Obj(),
336373
},
337374
wantWorkloads: []kueue.Workload{
338-
*baseWorkloadWrapper.Clone().
375+
*baseWorkloadWrapperWithAdmission.Clone().
339376
Finalizers(SliceControllerName).
340377
AdmissionCheck(kueue.AdmissionCheckState{
341378
Name: kueue.AdmissionCheckReference(baseAdmissionCheckName),
@@ -359,11 +396,11 @@ func TestWorkloadReconciler(t *testing.T) {
359396
request: baseRequest,
360397
objs: []client.Object{
361398
baseAdmissionCheckWrapper.DeepCopy(),
362-
baseWorkloadWrapper.Finalizers(SliceControllerName).DeepCopy(),
399+
baseWorkloadWrapperWithAdmission.Finalizers(SliceControllerName).DeepCopy(),
363400
baseSliceWrapper.Clone().Error().Obj(),
364401
},
365402
wantWorkloads: []kueue.Workload{
366-
*baseWorkloadWrapper.Clone().
403+
*baseWorkloadWrapperWithAdmission.Clone().
367404
Finalizers(SliceControllerName).
368405
AdmissionCheck(kueue.AdmissionCheckState{
369406
Name: kueue.AdmissionCheckReference(baseAdmissionCheckName),
@@ -388,11 +425,11 @@ func TestWorkloadReconciler(t *testing.T) {
388425
objs: []client.Object{
389426
baseAdmissionCheckWrapper.DeepCopy(),
390427
baseAdmissionCheckWrapper.Clone().Name(baseAdmissionCheckName + "2").Obj(),
391-
baseWorkloadWrapper.Finalizers(SliceControllerName).DeepCopy(),
428+
baseWorkloadWrapperWithAdmission.Finalizers(SliceControllerName).DeepCopy(),
392429
baseSliceWrapper.Clone().Ready().Obj(),
393430
},
394431
wantWorkloads: []kueue.Workload{
395-
*baseWorkloadWrapper.Clone().
432+
*baseWorkloadWrapperWithAdmission.Clone().
396433
Finalizers(SliceControllerName).
397434
AdmissionCheck(kueue.AdmissionCheckState{
398435
Name: kueue.AdmissionCheckReference(baseAdmissionCheckName),

0 commit comments

Comments
 (0)