Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 28 additions & 11 deletions src/compute-plane-services/nvca/internal/util/k8sutil/affinity.go
Original file line number Diff line number Diff line change
Expand Up @@ -77,18 +77,35 @@ func SetCPUWorkloadNodeAffinity(pts *corev1.PodSpec, instanceTypeLabelKey string

RemoveInstanceTypeRequiredAffinity(pts.Affinity.NodeAffinity, instanceTypeLabelKey)

pts.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append(
pts.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution,
corev1.PreferredSchedulingTerm{
Weight: 100,
Preference: corev1.NodeSelectorTerm{
MatchExpressions: []corev1.NodeSelectorRequirement{{
Key: instanceTypeLabelKey,
Operator: corev1.NodeSelectorOpDoesNotExist,
}},
},
preferred := &pts.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution
for _, term := range *preferred {
if hasCPUWorkloadSoftAntiAffinity(term, instanceTypeLabelKey) {
return
}
}

*preferred = append(*preferred, corev1.PreferredSchedulingTerm{
Weight: 100,
Preference: corev1.NodeSelectorTerm{
MatchExpressions: []corev1.NodeSelectorRequirement{{
Key: instanceTypeLabelKey,
Operator: corev1.NodeSelectorOpDoesNotExist,
}},
},
)
})
}

// hasCPUWorkloadSoftAntiAffinity reports whether preferred already includes a
// preference that soft-avoids nodes with instanceTypeLabelKey, even
// when that preference also has other MatchExpressions/MatchFields.
func hasCPUWorkloadSoftAntiAffinity(term corev1.PreferredSchedulingTerm, instanceTypeLabelKey string) bool {
for _, expr := range term.Preference.MatchExpressions {
if expr.Key == instanceTypeLabelKey &&
expr.Operator == corev1.NodeSelectorOpDoesNotExist {
return true
}
}
return false
}

// RemoveInstanceTypeRequiredAffinity removes any RequiredDuringSchedulingIgnoredDuringExecution
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
/*
SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
SPDX-License-Identifier: Apache-2.0

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package k8sutil

import (
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
corev1 "k8s.io/api/core/v1"
)

func TestSetCPUWorkloadNodeAffinity_DedupesMultiExpressionPreference(t *testing.T) {
const instanceTypeKey = "nvca.nvcf.nvidia.io/instance-type"

podSpec := &corev1.PodSpec{
Affinity: &corev1.Affinity{
NodeAffinity: &corev1.NodeAffinity{
PreferredDuringSchedulingIgnoredDuringExecution: []corev1.PreferredSchedulingTerm{
{
Weight: 100,
Preference: corev1.NodeSelectorTerm{
MatchExpressions: []corev1.NodeSelectorRequirement{
{
Key: "topology.kubernetes.io/zone",
Operator: corev1.NodeSelectorOpIn,
Values: []string{"us-west"},
},
{
Key: instanceTypeKey,
Operator: corev1.NodeSelectorOpDoesNotExist,
},
{
Key: "kubernetes.io/arch",
Operator: corev1.NodeSelectorOpIn,
Values: []string{"amd64"},
},
},
MatchFields: []corev1.NodeSelectorRequirement{{
Key: "metadata.name",
Operator: corev1.NodeSelectorOpIn,
Values: []string{"node-a"},
}},
},
},
},
},
},
}

SetCPUWorkloadNodeAffinity(podSpec, instanceTypeKey)
SetCPUWorkloadNodeAffinity(podSpec, instanceTypeKey)

preferred := podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution
require.Len(t, preferred, 1, "should not append a duplicate soft anti-affinity term")
assert.Equal(t, int32(100), preferred[0].Weight)
assert.Len(t, preferred[0].Preference.MatchExpressions, 3)
assert.Len(t, preferred[0].Preference.MatchFields, 1)

found := false
for _, expr := range preferred[0].Preference.MatchExpressions {
if expr.Key == instanceTypeKey &&
expr.Operator == corev1.NodeSelectorOpDoesNotExist &&
len(expr.Values) == 0 {
found = true
}
}
assert.True(t, found)
}
Original file line number Diff line number Diff line change
Expand Up @@ -327,8 +327,6 @@ func (w *miniserviceMutatingWebhook) mutate(ctx context.Context, obj client.Obje

// Pod spec mutations must only be applied on creation events.
if isCreate {
w.mutatePodSpec(&t.Spec, meta)

// NVLink DRA mutations for claims/scheduling.
if w.fff.IsAttributeEnabled(featureflag.AttrNVLinkOptimized) {
w.mutateNVLinkDRA(obj.GetNamespace(), t)
Expand All @@ -337,6 +335,8 @@ func (w *miniserviceMutatingWebhook) mutate(ctx context.Context, obj client.Obje
if _, _, err := w.sharedStorageMutator.mutate(ctx, obj, meta); err != nil {
return fmt.Errorf("mutate shared storage: %w", err)
}

w.mutatePodSpec(&t.Spec, meta)
}
}

Expand Down Expand Up @@ -386,6 +386,8 @@ func (w *miniserviceMutatingWebhook) mutatePodSpec(ps *corev1.PodSpec, meta nvca
if meta.NodeAffinityKey != "" {
if w.fff.IsFeatureFlagEnabled(featureflag.HelmAllowCPUNodes) && !k8sutil.PodSpecRequestsGPU(ps) {
k8sutil.SetCPUWorkloadNodeAffinity(ps, meta.NodeAffinityKey)
} else {
k8sutil.SetInstanceTypeNodeAffinity(ps, meta.NodeAffinityKey, meta.NodeAffinityValue)
}
}

Expand Down
Loading
Loading