Skip to content

Commit 9a46540

Browse files
dgrove-ossdgrove-oss
authored
Support PreferNoSchedule labels for AutoPilot (#315)
1 parent 8e8fe11 commit 9a46540

File tree

5 files changed

+107
-51
lines changed

5 files changed

+107
-51
lines changed

‎internal/controller/appwrapper/appwrapper_controller_test.go

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ var _ = Describe("AppWrapper Controller", func() {
6666
awConfig.FaultTolerance.RetryLimit = 0
6767
awConfig.FaultTolerance.SuccessTTL = 0 * time.Second
6868
awConfig.Autopilot.ResourceTaints["nvidia.com/gpu"] = append(awConfig.Autopilot.ResourceTaints["nvidia.com/gpu"], v1.Taint{Key: "extra", Value: "test", Effect: v1.TaintEffectNoExecute})
69+
awConfig.Autopilot.ResourceTaints["nvidia.com/gpu"] = append(awConfig.Autopilot.ResourceTaints["nvidia.com/gpu"], v1.Taint{Key: "extra2", Value: "test2", Effect: v1.TaintEffectPreferNoSchedule})
6970

7071
awReconciler = &AppWrapperReconciler{
7172
Client: k8sClient,
@@ -187,11 +188,23 @@ var _ = Describe("AppWrapper Controller", func() {
187188
mes := p.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions
188189
for _, taint := range awReconciler.Config.Autopilot.ResourceTaints["nvidia.com/gpu"] {
189190
found := false
190-
for _, me := range mes {
191-
if me.Key == taint.Key {
192-
Expect(me.Operator).Should(Equal(v1.NodeSelectorOpNotIn))
193-
Expect(me.Values).Should(ContainElement(taint.Value))
194-
found = true
191+
if taint.Effect == v1.TaintEffectNoExecute || taint.Effect == v1.TaintEffectNoSchedule {
192+
for _, me := range mes {
193+
if me.Key == taint.Key {
194+
Expect(me.Operator).Should(Equal(v1.NodeSelectorOpNotIn))
195+
Expect(me.Values).Should(ContainElement(taint.Value))
196+
found = true
197+
}
198+
}
199+
} else {
200+
for _, st := range p.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution {
201+
for _, me := range st.Preference.MatchExpressions {
202+
if me.Key == taint.Key {
203+
Expect(me.Operator).Should(Equal(v1.NodeSelectorOpNotIn))
204+
Expect(me.Values).Should(ContainElement(taint.Value))
205+
found = true
206+
}
207+
}
195208
}
196209
}
197210
Expect(found).Should(BeTrue())

‎internal/controller/appwrapper/node_health_monitor.go

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -161,10 +161,12 @@ func (r *NodeHealthMonitor) updateNoScheduleNodes(ctx context.Context, node *v1.
161161
for key, value := range node.GetLabels() {
162162
for resourceName, taints := range r.Config.Autopilot.ResourceTaints {
163163
for _, taint := range taints {
164-
if key == taint.Key && value == taint.Value {
165-
quantity := node.Status.Capacity.Name(v1.ResourceName(resourceName), resource.DecimalSI)
166-
if !quantity.IsZero() {
167-
noScheduleResources[v1.ResourceName(resourceName)] = *quantity
164+
if taint.Effect == v1.TaintEffectNoExecute || taint.Effect == v1.TaintEffectNoSchedule {
165+
if key == taint.Key && value == taint.Value {
166+
quantity := node.Status.Capacity.Name(v1.ResourceName(resourceName), resource.DecimalSI)
167+
if !quantity.IsZero() {
168+
noScheduleResources[v1.ResourceName(resourceName)] = *quantity
169+
}
168170
}
169171
}
170172
}

‎internal/controller/appwrapper/node_health_monitor_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ var _ = Describe("NodeMonitor Controller", func() {
111111
Expect(noExecuteNodes[node1Name.Name]).Should(HaveKey("nvidia.com/gpu"))
112112

113113
By("Removing the EVICT label updates unhealthyNodes")
114-
node.Labels["autopilot.ibm.com/gpuhealth"] = "ERR"
114+
node.Labels["autopilot.ibm.com/gpuhealth"] = "WARN"
115115
Expect(k8sClient.Update(ctx, node)).Should(Succeed())
116116
_, err = nodeMonitor.Reconcile(ctx, reconcile.Request{NamespacedName: node1Name})
117117
Expect(err).NotTo(HaveOccurred())
@@ -151,7 +151,7 @@ var _ = Describe("NodeMonitor Controller", func() {
151151

152152
// remove another 4 gpus, lending limit should be 0 = max(0, 6-4-4)
153153
node2 := getNode(node2Name.Name)
154-
node2.Labels["autopilot.ibm.com/gpuhealth"] = "ERR"
154+
node2.Labels["autopilot.ibm.com/gpuhealth"] = "TESTING"
155155
Expect(k8sClient.Update(ctx, node2)).Should(Succeed())
156156
_, err = nodeMonitor.Reconcile(ctx, reconcile.Request{NamespacedName: node2Name})
157157
Expect(err).NotTo(HaveOccurred())

‎internal/controller/appwrapper/resource_management.go

Lines changed: 74 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import (
3030
kresource "k8s.io/apimachinery/pkg/api/resource"
3131
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
3232
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
33+
"k8s.io/utils/ptr"
3334
"sigs.k8s.io/controller-runtime/pkg/client"
3435
"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
3536
"sigs.k8s.io/controller-runtime/pkg/log"
@@ -108,7 +109,7 @@ func hasResourceRequest(spec map[string]interface{}, resource string) bool {
108109
return false
109110
}
110111

111-
func addNodeSelectorsToAffinity(spec map[string]interface{}, exprsToAdd []v1.NodeSelectorRequirement) error {
112+
func addNodeSelectorsToAffinity(spec map[string]interface{}, exprsToAdd []v1.NodeSelectorRequirement, required bool, weight int32) error {
112113
if _, ok := spec["affinity"]; !ok {
113114
spec["affinity"] = map[string]interface{}{}
114115
}
@@ -123,44 +124,64 @@ func addNodeSelectorsToAffinity(spec map[string]interface{}, exprsToAdd []v1.Nod
123124
if !ok {
124125
return fmt.Errorf("spec.affinity.nodeAffinity is not a map")
125126
}
126-
if _, ok := nodeAffinity["requiredDuringSchedulingIgnoredDuringExecution"]; !ok {
127-
nodeAffinity["requiredDuringSchedulingIgnoredDuringExecution"] = map[string]interface{}{}
128-
}
129-
nodeSelector, ok := nodeAffinity["requiredDuringSchedulingIgnoredDuringExecution"].(map[string]interface{})
130-
if !ok {
131-
return fmt.Errorf("spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution is not a map")
132-
}
133-
if _, ok := nodeSelector["nodeSelectorTerms"]; !ok {
134-
nodeSelector["nodeSelectorTerms"] = []interface{}{map[string]interface{}{}}
135-
}
136-
existingTerms, ok := nodeSelector["nodeSelectorTerms"].([]interface{})
137-
if !ok {
138-
return fmt.Errorf("spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms is not an array")
139-
}
140-
for idx, term := range existingTerms {
141-
selTerm, ok := term.(map[string]interface{})
127+
if required {
128+
if _, ok := nodeAffinity["requiredDuringSchedulingIgnoredDuringExecution"]; !ok {
129+
nodeAffinity["requiredDuringSchedulingIgnoredDuringExecution"] = map[string]interface{}{}
130+
}
131+
nodeSelector, ok := nodeAffinity["requiredDuringSchedulingIgnoredDuringExecution"].(map[string]interface{})
142132
if !ok {
143-
return fmt.Errorf("spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[%v] is not an map", idx)
133+
return fmt.Errorf("spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution is not a map")
144134
}
145-
if _, ok := selTerm["matchExpressions"]; !ok {
146-
selTerm["matchExpressions"] = []interface{}{}
135+
if _, ok := nodeSelector["nodeSelectorTerms"]; !ok {
136+
nodeSelector["nodeSelectorTerms"] = []interface{}{map[string]interface{}{}}
147137
}
148-
matchExpressions, ok := selTerm["matchExpressions"].([]interface{})
138+
existingTerms, ok := nodeSelector["nodeSelectorTerms"].([]interface{})
149139
if !ok {
150-
return fmt.Errorf("spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[%v].matchExpressions is not an map", idx)
140+
return fmt.Errorf("spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms is not an array")
151141
}
152-
for _, expr := range exprsToAdd {
153-
bytes, err := json.Marshal(expr)
154-
if err != nil {
155-
return fmt.Errorf("marshalling selectorTerm %v: %w", expr, err)
142+
for idx, term := range existingTerms {
143+
selTerm, ok := term.(map[string]interface{})
144+
if !ok {
145+
return fmt.Errorf("spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[%v] is not an map", idx)
156146
}
157-
var obj interface{}
158-
if err = json.Unmarshal(bytes, &obj); err != nil {
159-
return fmt.Errorf("unmarshalling selectorTerm %v: %w", expr, err)
147+
if _, ok := selTerm["matchExpressions"]; !ok {
148+
selTerm["matchExpressions"] = []interface{}{}
160149
}
161-
matchExpressions = append(matchExpressions, obj)
150+
matchExpressions, ok := selTerm["matchExpressions"].([]interface{})
151+
if !ok {
152+
return fmt.Errorf("spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[%v].matchExpressions is not an map", idx)
153+
}
154+
for _, expr := range exprsToAdd {
155+
bytes, err := json.Marshal(expr)
156+
if err != nil {
157+
return fmt.Errorf("marshalling selectorTerm %v: %w", expr, err)
158+
}
159+
var obj interface{}
160+
if err = json.Unmarshal(bytes, &obj); err != nil {
161+
return fmt.Errorf("unmarshalling selectorTerm %v: %w", expr, err)
162+
}
163+
matchExpressions = append(matchExpressions, obj)
164+
}
165+
selTerm["matchExpressions"] = matchExpressions
166+
}
167+
} else {
168+
if _, ok := nodeAffinity["preferredDuringSchedulingIgnoredDuringExecution"]; !ok {
169+
nodeAffinity["preferredDuringSchedulingIgnoredDuringExecution"] = []interface{}{}
162170
}
163-
selTerm["matchExpressions"] = matchExpressions
171+
terms, ok := nodeAffinity["preferredDuringSchedulingIgnoredDuringExecution"].([]interface{})
172+
if !ok {
173+
return fmt.Errorf("spec.affinity.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution is not an array")
174+
}
175+
bytes, err := json.Marshal(v1.PreferredSchedulingTerm{Weight: weight, Preference: v1.NodeSelectorTerm{MatchExpressions: exprsToAdd}})
176+
if err != nil {
177+
return fmt.Errorf("marshalling selectorTerms %v: %w", exprsToAdd, err)
178+
}
179+
var obj interface{}
180+
if err = json.Unmarshal(bytes, &obj); err != nil {
181+
return fmt.Errorf("unmarshalling selectorTerms %v: %w", exprsToAdd, err)
182+
}
183+
terms = append(terms, obj)
184+
nodeAffinity["preferredDuringSchedulingIgnoredDuringExecution"] = terms
164185
}
165186

166187
return nil
@@ -291,20 +312,35 @@ func (r *AppWrapperReconciler) createComponent(ctx context.Context, aw *workload
291312
}
292313

293314
if r.Config.Autopilot != nil && r.Config.Autopilot.InjectAntiAffinities {
294-
toAdd := map[string][]string{}
315+
toAddRequired := map[string][]string{}
316+
toAddPreferred := map[string][]string{}
295317
for resource, taints := range r.Config.Autopilot.ResourceTaints {
296318
if hasResourceRequest(spec, resource) {
297319
for _, taint := range taints {
298-
toAdd[taint.Key] = append(toAdd[taint.Key], taint.Value)
320+
if taint.Effect == v1.TaintEffectNoExecute || taint.Effect == v1.TaintEffectNoSchedule {
321+
toAddRequired[taint.Key] = append(toAddRequired[taint.Key], taint.Value)
322+
} else if taint.Effect == v1.TaintEffectPreferNoSchedule {
323+
toAddPreferred[taint.Key] = append(toAddPreferred[taint.Key], taint.Value)
324+
}
299325
}
300326
}
301327
}
302-
if len(toAdd) > 0 {
328+
if len(toAddRequired) > 0 {
329+
matchExpressions := []v1.NodeSelectorRequirement{}
330+
for k, v := range toAddRequired {
331+
matchExpressions = append(matchExpressions, v1.NodeSelectorRequirement{Operator: v1.NodeSelectorOpNotIn, Key: k, Values: v})
332+
}
333+
if err := addNodeSelectorsToAffinity(spec, matchExpressions, true, 0); err != nil {
334+
log.FromContext(ctx).Error(err, "failed to inject Autopilot affinities")
335+
}
336+
}
337+
if len(toAddPreferred) > 0 {
303338
matchExpressions := []v1.NodeSelectorRequirement{}
304-
for k, v := range toAdd {
339+
for k, v := range toAddPreferred {
305340
matchExpressions = append(matchExpressions, v1.NodeSelectorRequirement{Operator: v1.NodeSelectorOpNotIn, Key: k, Values: v})
306341
}
307-
if err := addNodeSelectorsToAffinity(spec, matchExpressions); err != nil {
342+
weight := ptr.Deref(r.Config.Autopilot.PreferNoScheduleWeight, 1)
343+
if err := addNodeSelectorsToAffinity(spec, matchExpressions, false, weight); err != nil {
308344
log.FromContext(ctx).Error(err, "failed to inject Autopilot affinities")
309345
}
310346
}
@@ -315,6 +351,8 @@ func (r *AppWrapperReconciler) createComponent(ctx context.Context, aw *workload
315351
return err, true
316352
}
317353

354+
log.FromContext(ctx).Info("After injection", "obj", obj)
355+
318356
orig := copyForStatusPatch(aw)
319357
if meta.FindStatusCondition(aw.Status.ComponentStatus[componentIdx].Conditions, string(workloadv1beta2.ResourcesDeployed)) == nil {
320358
aw.Status.ComponentStatus[componentIdx].Name = obj.GetName()

‎pkg/config/config.go

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222

2323
v1 "k8s.io/api/core/v1"
2424
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
25+
"k8s.io/utils/ptr"
2526
"sigs.k8s.io/kueue/apis/config/v1beta1"
2627
)
2728

@@ -51,9 +52,10 @@ type KueueJobReconcillerConfig struct {
5152
}
5253

5354
type AutopilotConfig struct {
54-
InjectAntiAffinities bool `json:"injectAntiAffinities,omitempty"`
55-
MonitorNodes bool `json:"monitorNodes,omitempty"`
56-
ResourceTaints map[string][]v1.Taint `json:"resourceTaints,omitempty"`
55+
InjectAntiAffinities bool `json:"injectAntiAffinities,omitempty"`
56+
MonitorNodes bool `json:"monitorNodes,omitempty"`
57+
ResourceTaints map[string][]v1.Taint `json:"resourceTaints,omitempty"`
58+
PreferNoScheduleWeight *int32 `json:"preferNoScheduleWeight,omitempty"`
5759
}
5860

5961
type FaultToleranceConfig struct {
@@ -116,10 +118,11 @@ func NewAppWrapperConfig() *AppWrapperConfig {
116118
MonitorNodes: true,
117119
ResourceTaints: map[string][]v1.Taint{
118120
"nvidia.com/gpu": {
119-
{Key: "autopilot.ibm.com/gpuhealth", Value: "ERR", Effect: v1.TaintEffectNoSchedule},
121+
{Key: "autopilot.ibm.com/gpuhealth", Value: "WARN", Effect: v1.TaintEffectPreferNoSchedule},
120122
{Key: "autopilot.ibm.com/gpuhealth", Value: "TESTING", Effect: v1.TaintEffectNoSchedule},
121123
{Key: "autopilot.ibm.com/gpuhealth", Value: "EVICT", Effect: v1.TaintEffectNoExecute}},
122124
},
125+
PreferNoScheduleWeight: ptr.To(int32(50)),
123126
},
124127
UserRBACAdmissionCheck: true,
125128
FaultTolerance: &FaultToleranceConfig{

0 commit comments

Comments
 (0)