Skip to content

fix: constrain NVIDIA device plugin to GPU nodes #8866

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
atharvaHJoshi wants to merge 1 commit into
base: main
Choose a base branch
from
+227 −0
Open

fix: constrain NVIDIA device plugin to GPU nodes #8866

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
102 changes: 102 additions & 0 deletions pkg/addons/device_plugin.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
_ "embed"
"errors"
"fmt"
"sort"
"time"

"github.com/kris-nova/logger"
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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{
Expand Down
125 changes: 125 additions & 0 deletions pkg/addons/device_plugin_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down