feat(nodeadm): pass nvidia gpu startup labels to kubelet

Author: ndbaker1

nodeadm/internal/kubelet/config.go

@@ -213,6 +213,26 @@ func (ksc *kubeletConfig) withOutpostSetup(cfg *api.NodeConfig) error {
 	return nil
 }
 
+func (ksc *kubeletConfig) withNodeLabels(flags map[string]string, nodeLabelFuncs map[string]LabelValueFunc) {
+	var nodeLabels []string
+	for nodeLabelKey, nodeLabelFunc := range nodeLabelFuncs {
+		nodeLabelValue, ok, err := nodeLabelFunc()
+		if err != nil {
+			zap.L().Error("Failed to get node label value", zap.String("key", nodeLabelKey), zap.Error(err))
+			continue
+		}
+		if !ok {
+			continue
+		}
+		nodeLabel := fmt.Sprintf("%s=%s", nodeLabelKey, nodeLabelValue)
+		zap.L().Info("Adding node label", zap.String("label", nodeLabel))
+		nodeLabels = append(nodeLabels, nodeLabel)
+	}
+	if len(nodeLabels) > 0 {
+		flags["node-labels"] = strings.Join(nodeLabels, ",")
+	}
+}
+
 func (ksc *kubeletConfig) withNodeIp(cfg *api.NodeConfig, flags map[string]string) error {
 	nodeIp, err := getNodeIp(context.TODO(), cfg, imds.DefaultClient())
 	if err != nil {
@@ -284,14 +304,13 @@ func (ksc *kubeletConfig) withDefaultReservedResources(cfg *api.NodeConfig, reso
 //
 // TODO: revisit once the minimum supportted version catches up or the container
 // runtime is moved to containerd 2.0
-func (ksc *kubeletConfig) withPodInfraContainerImage(cfg *api.NodeConfig, flags map[string]string) error {
+func (ksc *kubeletConfig) withPodInfraContainerImage(cfg *api.NodeConfig, flags map[string]string) {
 	// the flag is a noop on 1.29+, since the behavior was changed to use the
 	// CRI image pinning behavior and no longer considers the flag value.
 	// see: https://github.com/kubernetes/kubernetes/pull/118544
 	if semver.Compare(cfg.Status.KubeletVersion, "v1.29.0") < 0 {
 		flags["pod-infra-container-image"] = cfg.Status.Defaults.SandboxImage
 	}
-	return nil
 }
 
 func (ksc *kubeletConfig) withImageServiceEndpoint(cfg *api.NodeConfig, resources system.Resources) {
@@ -312,14 +331,15 @@ func (k *kubelet) GenerateKubeletConfig(cfg *api.NodeConfig) (*kubeletConfig, er
 	if err := kubeletConfig.withNodeIp(cfg, k.flags); err != nil {
 		return nil, err
 	}
-	if err := kubeletConfig.withPodInfraContainerImage(cfg, k.flags); err != nil {
-		return nil, err
-	}
 
 	kubeletConfig.withVersionToggles(cfg)
 	kubeletConfig.withCloudProvider(cfg, k.flags)
 	kubeletConfig.withDefaultReservedResources(cfg, k.resources)
 	kubeletConfig.withImageServiceEndpoint(cfg, k.resources)
+	kubeletConfig.withNodeLabels(k.flags, map[string]LabelValueFunc{
+		// see: https://github.com/NVIDIA/gpu-operator/commit/e25291b86cf4542ac62d8635cda4bd653c4face3
+		"nvidia.com/gpu.present": getNvidiaGPULabel,
+	})
 
 	return &kubeletConfig, nil
 }

nodeadm/internal/kubelet/labels.go

@@ -0,0 +1,18 @@
+package kubelet
+
+import (
+	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/system"
+)
+
+type LabelValueFunc func() (string, bool, error)
+
+func getNvidiaGPULabel() (string, bool, error) {
+	ok, err := system.IsPCIVendorAttached(system.NVIDIA_VENDOR_ID)
+	if err != nil {
+		return "", false, err
+	}
+	if !ok {
+		return "", false, nil
+	}
+	return "true", true, nil
+}

nodeadm/internal/system/devices.go

@@ -0,0 +1,28 @@
+package system
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+const NVIDIA_VENDOR_ID = "0x10de"
+
+// IsPCIVendorAttached returns whether any pcie devices with a given vendor id
+// are attached to the instance.
+func IsPCIVendorAttached(vendorId string) (bool, error) {
+	vendorPaths, err := filepath.Glob("/sys/bus/pci/devices/*/vendor")
+	if err != nil {
+		return false, err
+	}
+	for _, vendorPath := range vendorPaths {
+		vendorIdBytes, err := os.ReadFile(filepath.Clean(vendorPath))
+		if err != nil {
+			continue
+		}
+		if strings.TrimSpace(string(vendorIdBytes)) == vendorId {
+			return true, nil
+		}
+	}
+	return false, nil
+}

nodeadm/test/e2e/cases/nvidia-gpu-labels/config.yaml

@@ -0,0 +1,12 @@
+---
+apiVersion: node.eks.aws/v1alpha1
+kind: NodeConfig
+spec:
+  cluster:
+    name: my-cluster
+    apiServerEndpoint: https://example.com
+    certificateAuthority: Y2VydGlmaWNhdGVBdXRob3JpdHk=
+    cidr: 10.100.0.0/16
+  kubelet:
+    flags:
+      - --node-labels=foo=bar

nodeadm/test/e2e/cases/nvidia-gpu-labels/run.sh

@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+
+set -o errexit
+set -o nounset
+set -o pipefail
+
+source /helpers.sh
+
+mock::aws
+mock::kubelet 1.32.0
+wait::dbus-ready
+
+nodeadm init --skip run --config-source file://config.yaml
+
+assert::file-not-contains /etc/eks/kubelet/environment '--node-labels=nvidia.com/gpu.present'
+assert::file-contains /etc/eks/kubelet/environment '--node-labels=foo=bar'
+
+# mock a pci device with the nvidia vendor-id.
+mock::pci-device "0x10de"
+
+nodeadm init --skip run --config-source file://config.yaml
+
+assert::file-contains /etc/eks/kubelet/environment '--node-labels=nvidia.com/gpu.present=true'
+assert::file-contains /etc/eks/kubelet/environment '--node-labels=foo=bar'

nodeadm/test/e2e/helpers.sh

@@ -99,6 +99,19 @@ function mock::setup-local-disks() {
   chmod +x /usr/bin/setup-local-disks
 }
 
+function mock::pci-device() {
+  # since we cannot modify the kernel sysfs paths, we bind mount our mock
+  # directory on top of an existing pci device. the only requirement for the check
+  # is that the vendor ID file is correct.
+  #
+  # NOTE: this currently only supports mocking a single device at a time.
+  local pcie_vendor=$1
+  pci_mock_dst=/sys/bus/pci/devices/$(ls /sys/bus/pci/devices/ | head -n 1)
+  pci_mock_src=$(mktemp -d)
+  echo "$pcie_vendor" > $pci_mock_src/vendor
+  mount --bind $pci_mock_src $pci_mock_dst
+}
+
 function wait::path-exists() {
   if [ "$#" -ne 1 ]; then
     echo "Usage: wait::path-exists TARGET_PATH"