From 73d02f093a00b1dcfe0fe1b960756694d929deb5 Mon Sep 17 00:00:00 2001 From: Atharva Joshi Date: Sat, 12 Sep 2026 17:46:28 +0530 Subject: [PATCH] fix: constrain NVIDIA device plugin to GPU nodes Signed-off-by: Atharva Joshi --- pkg/addons/device_plugin.go | 102 +++++++++++++++++++++++++ pkg/addons/device_plugin_test.go | 125 +++++++++++++++++++++++++++++++ 2 files changed, 227 insertions(+) diff --git a/pkg/addons/device_plugin.go b/pkg/addons/device_plugin.go index 5a72fe7799..977d06ae19 100644 --- a/pkg/addons/device_plugin.go +++ b/pkg/addons/device_plugin.go @@ -6,6 +6,7 @@ import ( _ "embed" "errors" "fmt" + "sort" "time" "github.com/kris-nova/logger" @@ -81,6 +82,7 @@ type DevicePlugin interface { Manifest() []byte SetImage(t *corev1.PodTemplateSpec) error SetTolerations(t *corev1.PodTemplateSpec) error + SetNodeAffinity(t *corev1.PodTemplateSpec) error Deploy() error } @@ -108,6 +110,9 @@ func applyDevicePlugin(dp DevicePlugin) error { if err := dp.SetTolerations(&daemonSet.Spec.Template); err != nil { return fmt.Errorf("adding tolerations to device plugin daemonset: %w", err) } + if err := dp.SetNodeAffinity(&daemonSet.Spec.Template); err != nil { + return fmt.Errorf("adding node affinity to device plugin daemonset: %w", err) + } msg, err := rawResource.CreateOrReplace(dp.PlanMode()) if err != nil { return fmt.Errorf("calling create or replace on raw device plugin daemonset: %w", err) @@ -163,6 +168,10 @@ func (n *NeuronDevicePlugin) SetTolerations(t *corev1.PodTemplateSpec) error { return nil } +func (n *NeuronDevicePlugin) SetNodeAffinity(t *corev1.PodTemplateSpec) error { + return nil +} + // Deploy deploys the Neuron device plugin to the specified cluster func (n *NeuronDevicePlugin) Deploy() error { return applyDevicePlugin(n) @@ -254,6 +263,95 @@ func (n *NvidiaDevicePlugin) SetTolerations(spec *corev1.PodTemplateSpec) error return nil } +// SetNodeAffinity sets a required node affinity on the DaemonSet pod template so +// that the device plugin is only scheduled on the NVIDIA GPU instance types +// defined in the cluster configuration. Without it the DaemonSet would schedule a +// replica on every node in the cluster, including CPU-only nodes where it has no +// GPU to manage and can end up crash-looping and draining node resources (see +// https://github.com/eksctl-io/eksctl/issues/8858). +func (n *NvidiaDevicePlugin) SetNodeAffinity(spec *corev1.PodTemplateSpec) error { + instanceTypes := n.nvidiaInstanceTypes() + if len(instanceTypes) == 0 { + return nil + } + + if spec.Spec.Affinity == nil { + spec.Spec.Affinity = &corev1.Affinity{} + } + if spec.Spec.Affinity.NodeAffinity == nil { + spec.Spec.Affinity.NodeAffinity = &corev1.NodeAffinity{} + } + nodeSelector := spec.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution + if nodeSelector == nil { + nodeSelector = &corev1.NodeSelector{} + spec.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = nodeSelector + } + + nodeSelector.NodeSelectorTerms = append(nodeSelector.NodeSelectorTerms, corev1.NodeSelectorTerm{ + MatchExpressions: []corev1.NodeSelectorRequirement{ + { + Key: corev1.LabelInstanceTypeStable, + Operator: corev1.NodeSelectorOpIn, + Values: instanceTypes, + }, + { + Key: "eks.amazonaws.com/compute-type", + Operator: corev1.NodeSelectorOpNotIn, + Values: []string{"fargate", "hybrid", "auto"}, + }, + }, + }) + return nil +} + +// nvidiaInstanceTypes returns the instance types in the cluster configuration that +// run NVIDIA GPUs on a supported AMI family (AmazonLinux2/AmazonLinux2023), i.e. +// the nodes the Nvidia device plugin needs to run on. +func (n *NvidiaDevicePlugin) nvidiaInstanceTypes() []string { + isSupportedFamily := func(family string) bool { + return family == api.NodeImageFamilyAmazonLinux2 || family == api.NodeImageFamilyAmazonLinux2023 + } + + instanceTypes := make(map[string]struct{}) + addType := func(instanceType string) { + // Only NVIDIA instance types actually expose GPUs; nodes of other types + // in a mixed nodegroup must not run the device plugin. + if instance.IsNvidiaInstanceType(instanceType) { + instanceTypes[instanceType] = struct{}{} + } + } + for _, ng := range n.spec.NodeGroups { + if !api.HasInstanceType(ng, instance.IsNvidiaInstanceType) || !isSupportedFamily(ng.GetAMIFamily()) { + continue + } + addType(ng.InstanceType) + if ng.InstancesDistribution != nil { + for _, it := range ng.InstancesDistribution.InstanceTypes { + addType(it) + } + } + } + for _, ng := range n.spec.ManagedNodeGroups { + if !api.HasInstanceTypeManaged(ng, instance.IsNvidiaInstanceType) || !isSupportedFamily(ng.GetAMIFamily()) { + continue + } + addType(ng.InstanceType) + for _, it := range ng.InstanceTypes { + addType(it) + } + } + + if len(instanceTypes) == 0 { + return nil + } + result := make([]string, 0, len(instanceTypes)) + for it := range instanceTypes { + result = append(result, it) + } + sort.Strings(result) + return result +} + // A EFADevicePlugin deploys the EFA Device Plugin to a cluster type EFADevicePlugin struct { rawClient kubernetes.RawClientInterface @@ -282,6 +380,10 @@ func (n *EFADevicePlugin) SetTolerations(spec *corev1.PodTemplateSpec) error { return nil } +func (n *EFADevicePlugin) SetNodeAffinity(spec *corev1.PodTemplateSpec) error { + return nil +} + // NewEFADevicePlugin creates a new EFADevicePlugin func NewEFADevicePlugin(rawClient kubernetes.RawClientInterface, region string, planMode bool, spec *api.ClusterConfig) DevicePlugin { return &EFADevicePlugin{ diff --git a/pkg/addons/device_plugin_test.go b/pkg/addons/device_plugin_test.go index 5f1bd67fc4..fb5a12c999 100644 --- a/pkg/addons/device_plugin_test.go +++ b/pkg/addons/device_plugin_test.go @@ -10,6 +10,131 @@ import ( ) var _ = Describe("NvidiaDevicePlugin", func() { + Describe("SetNodeAffinity", func() { + var ( + plugin *addons.NvidiaDevicePlugin + spec *corev1.PodTemplateSpec + config *api.ClusterConfig + ) + + BeforeEach(func() { + spec = &corev1.PodTemplateSpec{ + Spec: corev1.PodSpec{}, + } + config = &api.ClusterConfig{} + }) + + It("should add a node affinity for each NVIDIA instance type in the cluster", func() { + config.NodeGroups = []*api.NodeGroup{ + { + NodeGroupBase: &api.NodeGroupBase{ + Name: "nvidia-ng", + InstanceType: "g4dn.xlarge", + AMIFamily: api.NodeImageFamilyAmazonLinux2, + }, + }, + } + + plugin = addons.NewNvidiaDevicePlugin(nil, "us-west-2", false, config).(*addons.NvidiaDevicePlugin) + err := plugin.SetNodeAffinity(spec) + + Expect(err).NotTo(HaveOccurred()) + Expect(spec.Spec.Affinity).NotTo(BeNil()) + Expect(spec.Spec.Affinity.NodeAffinity).NotTo(BeNil()) + terms := spec.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms + Expect(terms).To(HaveLen(1)) + Expect(terms[0].MatchExpressions).To(ContainElement(corev1.NodeSelectorRequirement{ + Key: corev1.LabelInstanceTypeStable, + Operator: corev1.NodeSelectorOpIn, + Values: []string{"g4dn.xlarge"}, + })) + Expect(terms[0].MatchExpressions).To(ContainElement(corev1.NodeSelectorRequirement{ + Key: "eks.amazonaws.com/compute-type", + Operator: corev1.NodeSelectorOpNotIn, + Values: []string{"fargate", "hybrid", "auto"}, + })) + }) + + It("should collect NVIDIA instance types from managed nodegroups", func() { + config.NodeGroups = []*api.NodeGroup{ + { + NodeGroupBase: &api.NodeGroupBase{ + Name: "cpu-ng", + InstanceType: "m5.large", + AMIFamily: api.NodeImageFamilyAmazonLinux2023, + }, + }, + } + config.ManagedNodeGroups = []*api.ManagedNodeGroup{ + { + NodeGroupBase: &api.NodeGroupBase{ + Name: "managed-nvidia-ng", + InstanceType: "g4dn.2xlarge", + AMIFamily: api.NodeImageFamilyAmazonLinux2023, + }, + InstanceTypes: []string{"g4dn.4xlarge", "g5.xlarge"}, + }, + } + + plugin = addons.NewNvidiaDevicePlugin(nil, "us-west-2", false, config).(*addons.NvidiaDevicePlugin) + err := plugin.SetNodeAffinity(spec) + + Expect(err).NotTo(HaveOccurred()) + terms := spec.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms + Expect(terms).To(HaveLen(1)) + Expect(terms[0].MatchExpressions).To(ContainElement(corev1.NodeSelectorRequirement{ + Key: corev1.LabelInstanceTypeStable, + Operator: corev1.NodeSelectorOpIn, + Values: []string{"g4dn.2xlarge", "g4dn.4xlarge", "g5.xlarge"}, + })) + }) + + It("should only add NVIDIA instance types from unmanaged mixed instance nodegroups", func() { + config.NodeGroups = []*api.NodeGroup{ + { + NodeGroupBase: &api.NodeGroupBase{ + Name: "mixed-nvidia-ng", + InstanceType: "m5.large", + AMIFamily: api.NodeImageFamilyAmazonLinux2, + }, + InstancesDistribution: &api.NodeGroupInstancesDistribution{ + InstanceTypes: []string{"g4dn.xlarge", "m5.large", "g5.xlarge"}, + }, + }, + } + + plugin = addons.NewNvidiaDevicePlugin(nil, "us-west-2", false, config).(*addons.NvidiaDevicePlugin) + err := plugin.SetNodeAffinity(spec) + + Expect(err).NotTo(HaveOccurred()) + terms := spec.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms + Expect(terms).To(HaveLen(1)) + Expect(terms[0].MatchExpressions).To(ContainElement(corev1.NodeSelectorRequirement{ + Key: corev1.LabelInstanceTypeStable, + Operator: corev1.NodeSelectorOpIn, + Values: []string{"g4dn.xlarge", "g5.xlarge"}, + })) + }) + + It("should not add an affinity for non-NVIDIA nodegroups", func() { + config.NodeGroups = []*api.NodeGroup{ + { + NodeGroupBase: &api.NodeGroupBase{ + Name: "cpu-ng", + InstanceType: "m5.large", + AMIFamily: api.NodeImageFamilyAmazonLinux2, + }, + }, + } + + plugin = addons.NewNvidiaDevicePlugin(nil, "us-west-2", false, config).(*addons.NvidiaDevicePlugin) + err := plugin.SetNodeAffinity(spec) + + Expect(err).NotTo(HaveOccurred()) + Expect(spec.Spec.Affinity).To(BeNil()) + }) + }) + Describe("SetTolerations", func() { var ( plugin *addons.NvidiaDevicePlugin