refactor(nodeadm): PCIe detection for GPU instance types

Author: ndbaker1

nodeadm/internal/containerd/config.go

@@ -30,11 +30,15 @@ type containerdTemplateVars struct {
 	RuntimeBinaryName string
 }
 
-func writeContainerdConfig(cfg *api.NodeConfig) error {
-	if err := writeBaseRuntimeSpec(cfg); err != nil {
-		return err
-	}
+type modifier interface {
+	// Modify mutates a given containerd template configuration.
+	Modify(*containerdTemplateVars)
+
+	// Matches determines whether the configurator is relevant to the current instance.
+	Matches(*api.NodeConfig) bool
+}
 
+func writeContainerdConfig(cfg *api.NodeConfig) error {
 	containerdConfig, err := generateContainerdConfig(cfg)
 	if err != nil {
 		return err
@@ -59,12 +63,17 @@ func writeContainerdConfig(cfg *api.NodeConfig) error {
 }
 
 func generateContainerdConfig(cfg *api.NodeConfig) ([]byte, error) {
-	instanceOptions := applyInstanceTypeMixins(cfg.Status.Instance.Type)
-
 	configVars := containerdTemplateVars{
 		SandboxImage:      cfg.Status.Defaults.SandboxImage,
-		RuntimeBinaryName: instanceOptions.RuntimeBinaryName,
-		RuntimeName:       instanceOptions.RuntimeName,
+		RuntimeBinaryName: "/usr/sbin/runc",
+		RuntimeName:       "runc",
+	}
+	for _, configurator := range []modifier{
+		NewNvidiaModifier(),
+	} {
+		if configurator.Matches(cfg) {
+			configurator.Modify(&configVars)
+		}
 	}
 	var buf bytes.Buffer
 	if err := containerdConfigTemplate.Execute(&buf, configVars); err != nil {

nodeadm/internal/containerd/daemon.go

@@ -20,6 +20,9 @@ func NewContainerdDaemon(daemonManager daemon.DaemonManager) daemon.Daemon {
 }
 
 func (cd *containerd) Configure(c *api.NodeConfig) error {
+	if err := writeBaseRuntimeSpec(c); err != nil {
+		return err
+	}
 	return writeContainerdConfig(c)
 }
 

nodeadm/internal/containerd/nvidia.go

@@ -0,0 +1,65 @@
+package containerd
+
+import (
+	"os"
+	"slices"
+	"strings"
+
+	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/api"
+	"go.uber.org/zap"
+)
+
+type nvidiaModifier struct {
+	pcieDevicesPath string
+}
+
+func NewNvidiaModifier() *nvidiaModifier {
+	return &nvidiaModifier{
+		pcieDevicesPath: "/proc/bus/pci/devices",
+	}
+}
+
+func (m *nvidiaModifier) Matches(cfg *api.NodeConfig) bool {
+	return m.matchesInstanceType(cfg.Status.Instance.Type) || m.matchesPCIeVendor()
+}
+
+func (*nvidiaModifier) Modify(ctrdTemplate *containerdTemplateVars) {
+	zap.L().Info("Configuring NVIDIA runtime..")
+	ctrdTemplate.RuntimeName = "nvidia"
+	ctrdTemplate.RuntimeBinaryName = "/usr/bin/nvidia-container-runtime"
+}
+
+var nvidiaInstanceFamilies = []string{
+	"p3", "p3dn",
+	"p4d", "p4de",
+	"p5", "p5e", "p5en",
+	"g4", "g4dn",
+	"g5", "g5g",
+	"g6", "g6e",
+}
+
+// TODO: deprecate to avoid manual instance type tracking.
+func (*nvidiaModifier) matchesInstanceType(instanceType string) bool {
+	family := strings.Split(instanceType, ".")[0]
+	return slices.Contains(nvidiaInstanceFamilies, family)
+}
+
+func (m *nvidiaModifier) matchesPCIeVendor() bool {
+	devices, err := os.ReadFile(m.pcieDevicesPath)
+	if err != nil {
+		zap.L().Error("Failed to read PCIe devices", zap.Error(err))
+		return false
+	}
+	// The contents of '/proc/bus/pci/devices' looks like the following, where
+	// the last column contains the vendor name if present.
+	//
+	// something like the following:
+	//
+	// 0018 1d0f1111 0 c1000008         0        0         0 0 0 c0002  400000        0      0       0 0 0 20000
+	// 0020 1d0f8061 b c1508000         0        0         0 0 0     0    4000        0      0       0 0 0     0 nvme
+	// 0028 1d0fec20 0 c1504000         0 c1400008         0 0 0     0    4000        0 100000       0 0 0     0 ena
+	// 00f0 10de1eb8 a c0000000 44000000c        0 45000000c 0 0     0 1000000 10000000      0 2000000 0 0     0 nvidia
+	// 00f8 1d0fcd01 0 c1500000         0 c150c008         0 0 0     0    4000        0   2000       0 0 0     0 nvme
+	// 0030 1d0fec20 0 c1510000         0 c1600008         0 0 0     0    4000        0 100000       0 0 0     0 ena
+	return strings.Contains(string(devices), "nvidia")
+}

nodeadm/internal/containerd/nvidia_test.go

@@ -0,0 +1,47 @@
+package containerd
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/api"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNvidiaConfigurator(t *testing.T) {
+
+	t.Run("IsNvidiaUsingInstanceType", func(t *testing.T) {
+		configurator := nvidiaModifier{}
+		template := containerdTemplateVars{}
+		assert.True(t, configurator.Matches(nvidiaInstanceTypeNodeConfig("p5.48xlarge")))
+		configurator.Modify(&template)
+		assert.Equal(t, "nvidia", template.RuntimeName)
+		assert.Equal(t, "/usr/bin/nvidia-container-runtime", template.RuntimeBinaryName)
+	})
+
+	t.Run("IsNvidiaUsingPCIe", func(t *testing.T) {
+		configurator := nvidiaModifier{pcieDevicesPath: filepath.Join(t.TempDir(), "pcie-devices")}
+		os.WriteFile(configurator.pcieDevicesPath, []byte("nvidia"), 0777)
+		template := containerdTemplateVars{}
+		assert.True(t, configurator.Matches(nvidiaInstanceTypeNodeConfig("xxx.xxxxx")))
+		configurator.Modify(&template)
+		assert.Equal(t, "nvidia", template.RuntimeName)
+		assert.Equal(t, "/usr/bin/nvidia-container-runtime", template.RuntimeBinaryName)
+	})
+
+	t.Run("IsNotNvidia", func(t *testing.T) {
+		configurator := nvidiaModifier{}
+		assert.False(t, configurator.Matches(nvidiaInstanceTypeNodeConfig("m5.large")))
+	})
+}
+
+func nvidiaInstanceTypeNodeConfig(instanceType string) *api.NodeConfig {
+	return &api.NodeConfig{
+		Status: api.NodeConfigStatus{
+			Instance: api.InstanceDetails{
+				Type: instanceType,
+			},
+		},
+	}
+}