Add build scripts for building Nvidia and Neuron AMIs based on AL2023

Makefile

@@ -14,6 +14,7 @@ K8S_VERSION_MINOR := $(word 1,${K8S_VERSION_PARTS}).$(word 2,${K8S_VERSION_PARTS
 
 AMI_VARIANT ?= amazon-eks
 AMI_VERSION ?= v$(shell date '+%Y%m%d')
+NVIDIA_MAJOR_DRIVER_VERSION_DEFAULT=555
 os_distro ?= al2
 arch ?= x86_64
 
@@ -30,6 +31,16 @@ ifeq ($(enable_fips), true)
 	AMI_VARIANT := $(AMI_VARIANT)-fips
 endif
 
+ifeq ($(os_distro), al2023)
+	ifdef accelerator_vendor
+		AMI_VARIANT := $(AMI_VARIANT)-$(accelerator_vendor)
+		ifeq ($(accelerator_vendor), nvidia)
+			nvidia_major_driver_version := $(NVIDIA_MAJOR_DRIVER_VERSION_DEFAULT)
+			AMI_VARIANT := $(AMI_VARIANT)-$(nvidia_major_driver_version)
+		endif
+	endif
+endif
+
 ami_name ?= $(AMI_VARIANT)-node-$(K8S_VERSION_MINOR)-$(AMI_VERSION)
 
 # ami owner overrides for cn/gov-cloud
@@ -91,7 +102,7 @@ validate: ## Validate packer config
 
 .PHONY: k8s
 k8s: validate ## Build default K8s version of EKS Optimized AMI
-	@echo "Building AMI [os_distro=$(os_distro) kubernetes_version=$(kubernetes_version) arch=$(arch)]"
+	@echo "Building AMI [os_distro=$(os_distro) kubernetes_version=$(kubernetes_version) arch=$(arch) $(if $(accelerator_vendor),accelerator_vendor=$(accelerator_vendor))]"
 	$(PACKER_BINARY) build -timestamp-ui -color=false $(PACKER_ARGS) $(PACKER_TEMPLATE_FILE)
 
 # DEPRECATION NOTICE: `make` targets for each Kubernetes minor version will not be added after 1.28

README.md

@@ -42,10 +42,22 @@ make k8s=1.29
 # build an AMI with a specific Kubernetes version and a specific OS distro
 make k8s=1.29 os_distro=al2023
 
+# build Neuron AMI with a specific Kubernetes version
+make k8s=1.29 os_distro=al2023 accelerator_vendor=neuron
+
+# build Nvidia AMI with a specific Kubernetes version and default driver major version
+make k8s=1.29 os_distro=al2023 accelerator_vendor=nvidia
+
+# build Nvidia AMI with a specific Kubernetes version and a specific driver major version
+make k8s=1.29 os_distro=al2023 accelerator_vendor=nvidia nvidia_major_driver_version=555
+
 # check default value and options in help doc
 make help
 ```
 
+> **Note**
+> Accelerated AMIs (Neuron and Nvidia) are currently only supported for AL2023 operating system.
+
 The Makefile chooses a particular kubelet binary to use per Kubernetes version which you can [view here](Makefile).
 
 > **Note**
@@ -58,4 +70,18 @@ For security issues or concerns, please do not open an issue or pull request on
 
 ## ⚖️ License Summary
 
-This sample code is made available under a modified MIT license. See the LICENSE file.
+This sample code is made available under a MIT-0 license. See the LICENSE file.
+
+Although this repository is released under the MIT license, when using Nvidia accelerated AMIs you agree to the NVIDIA Cloud End User License Agreement: https://s3.amazonaws.com/EULA/NVidiaEULAforAWS.pdf.
+
+Although this repository is released under the MIT license, Nvidia accelerated AMIs  
+use the third party [open-gpu-kernel-modules](https://github.com/NVIDIA/open-gpu-kernel-modules). The open-gpu-kernel-modules project's licensing includes the dual MIT/GPLv2 license.
+
+Although this repository is released under the MIT license, Nvidia accelerated AMIs  
+use the third party [nvidia-container-toolkit](https://github.com/NVIDIA/nvidia-container-toolkit). The nvidia-container-toolkit project's licensing includes the Apache-2.0 license.
+
+Although this repository is released under the MIT license, Neuron accelerated AMIs  
+use the third party [Neuron Driver](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/runtime/aws-neuronx-dkms/index.html). The Neuron Driver project's licensing includes the GPLv2 license.
+
+Although this repository is released under the MIT license, accelerated AMIs  
+use the third party [Elastic Fabric Adapter Driver](https://github.com/amzn/amzn-drivers/tree/master/kernel/linux/efa). The Elastic Fabric Adapter Driver project's licensing includes the GPLv2 license.

doc/usage/al2023.md

@@ -5,6 +5,7 @@
 <!-- template-variable-table-boundary -->
 | Variable | Description |
 | - | - |
+| `accelerator_vendor` | Vendor that provides the GPU or accelerator hardware. Currently we support Neuron and Nvidia. |
 | `ami_component_description` |  |
 | `ami_description` |  |
 | `ami_name` |  |
@@ -28,6 +29,7 @@
 | `kubernetes_build_date` |  |
 | `kubernetes_version` |  |
 | `launch_block_device_mappings_volume_size` |  |
+| `nvidia_major_driver_version` | Driver version to install, depends on what is available in Nvidia upstream yum repository. |
 | `remote_folder` | Directory path for shell provisioner scripts on the builder instance |
 | `runc_version` |  |
 | `security_group_id` |  |

nodeadm/internal/containerd/config.go

@@ -37,6 +37,12 @@ func writeContainerdConfig(cfg *api.NodeConfig) error {
 	if err != nil {
 		return err
 	}
+
+	err = applyInstanceTypeMixins(cfg, &containerdConfig)
+	if err != nil {
+		return err
+	}
+
 	// because the logic in containerd's import merge decides to completely
 	// overwrite entire sections, we want to implement this merging ourselves.
 	// see: https://github.com/containerd/containerd/blob/a91b05d99ceac46329be06eb43f7ae10b89aad45/cmd/containerd/server/config/config.go#L407-L431

nodeadm/internal/containerd/oci_config.go

@@ -0,0 +1,118 @@
+package containerd
+
+import (
+	"errors"
+	"io/fs"
+	"os/exec"
+	"reflect"
+	"slices"
+	"strings"
+
+	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/api"
+	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/util"
+	"github.com/pelletier/go-toml/v2"
+	"go.uber.org/zap"
+)
+
+type instanceTypeMixin struct {
+	instanceFamilies []string
+	apply            func(*[]byte) error
+}
+
+func (m *instanceTypeMixin) matches(cfg *api.NodeConfig) bool {
+	instanceFamily := strings.Split(cfg.Status.Instance.Type, ".")[0]
+	return slices.Contains(m.instanceFamilies, instanceFamily)
+}
+
+var (
+	// TODO: fetch this list dynamically
+	nvidiaInstances         = []string{"p3", "p3dn", "p4d", "p4de", "p5", "g4", "g4dn", "g5", "g6"}
+	NvidiaInstanceTypeMixin = instanceTypeMixin{
+		instanceFamilies: nvidiaInstances,
+		apply:            applyNvidia,
+	}
+
+	mixins = []instanceTypeMixin{
+		NvidiaInstanceTypeMixin,
+	}
+)
+
+// applyInstanceTypeMixins adds the needed OCI hook options to containerd config.toml
+// based on the instance family
+func applyInstanceTypeMixins(cfg *api.NodeConfig, containerdConfig *[]byte) error {
+	for _, mixin := range mixins {
+		if mixin.matches(cfg) {
+			if err := mixin.apply(containerdConfig); err != nil {
+				return err
+			}
+			return nil
+		}
+	}
+	zap.L().Info("No containerd OCI configuration needed..", zap.String("instanceType", cfg.Status.Instance.Type))
+	return nil
+}
+
+type CommandExecutor interface {
+	CombinedOutput(name string, arg ...string) ([]byte, error)
+}
+
+type RealCommandExecutor struct{}
+
+func (r RealCommandExecutor) CombinedOutput(name string, arg ...string) ([]byte, error) {
+	if isAllowedCommand(name, arg...) {
+		cmd := exec.Command(name, arg...)
+		return cmd.CombinedOutput()
+	}
+	return nil, errors.New("unrecognised command")
+
+}
+
+type FileWriter interface {
+	WriteFileWithDir(filePath string, data []byte, perm fs.FileMode) error
+}
+
+type RealFileWriter struct{}
+
+func (r RealFileWriter) WriteFileWithDir(filePath string, data []byte, perm fs.FileMode) error {
+	return util.WriteFileWithDir(filePath, data, perm)
+}
+
+var execCommand CommandExecutor = RealCommandExecutor{}
+var fileWriter FileWriter = RealFileWriter{}
+var nvidiaCtkCommand = "/usr/bin/nvidia-ctk"
+var nvidiaCtkParams = []string{"--quiet", "runtime", "configure", "--runtime=containerd", "--set-as-default", "--dry-run"}
+
+func isAllowedCommand(command string, args ...string) bool {
+	return command == "/usr/bin/nvidia-ctk" && reflect.DeepEqual(args, nvidiaCtkParams)
+}
+
+// applyNvidia adds the needed Nvidia containerd options using nvidia container toolkit:
+// https://github.com/NVIDIA/nvidia-container-toolkit
+func applyNvidia(containerdConfig *[]byte) error {
+	zap.L().Info("Configuring Nvidia OCI hook..")
+	// before calling nvidia-ctk, we'll write the generated containerd config first
+	// nvidia-ctk will use it as a base
+	err := fileWriter.WriteFileWithDir(containerdConfigFile, *containerdConfig, containerdConfigPerm)
+
+	if err != nil {
+		return err
+	}
+
+	output, err := execCommand.CombinedOutput(nvidiaCtkCommand, nvidiaCtkParams...)
+	if err != nil {
+		return err
+	}
+
+	if containerdConfig != nil {
+		containerdConfigMap, err := util.Merge(*containerdConfig, output, toml.Marshal, toml.Unmarshal)
+		if err != nil {
+			return err
+		}
+		*containerdConfig, err = toml.Marshal(containerdConfigMap)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

nodeadm/internal/containerd/oci_config_test.go

@@ -0,0 +1,118 @@
+package containerd
+
+import (
+	"io/fs"
+	"reflect"
+	"testing"
+
+	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/api"
+)
+
+type MockCommandExecutor struct {
+	Output []byte
+	Err    error
+}
+
+func (m MockCommandExecutor) CombinedOutput(name string, arg ...string) ([]byte, error) {
+	return m.Output, m.Err
+}
+
+type MockFileWriter struct {
+	Err error
+}
+
+func (m MockFileWriter) WriteFileWithDir(filePath string, data []byte, perm fs.FileMode) error {
+	return m.Err
+}
+
+func TestApplyInstanceTypeMixins(t *testing.T) {
+
+	var nvidiaExpectedOutput = []byte(`version = 2
+
+[grpc]
+address = '/run/foo/foo.sock'
+
+[plugins]
+[plugins.'io.containerd.grpc.v1.cri']
+[plugins.'io.containerd.grpc.v1.cri'.containerd]
+default_runtime_name = 'nvidia'
+
+[plugins.'io.containerd.grpc.v1.cri'.containerd.runtimes]
+[plugins.'io.containerd.grpc.v1.cri'.containerd.runtimes.nvidia]
+privileged_without_host_devices = false
+runtime_engine = ''
+runtime_root = ''
+runtime_type = 'io.containerd.runc.v2'
+
+[plugins.'io.containerd.grpc.v1.cri'.containerd.runtimes.nvidia.options]
+BinaryName = '/usr/bin/nvidia-container-runtime'
+`)
+
+	var nonAcceleratedExpectedOutput = []byte(`
+version = 2
+[grpc]
+address = '/run/foo/foo.sock'
+`)
+	var tests = []struct {
+		name           string
+		instanceType   string
+		expectedOutput []byte
+		expectedError  error
+	}{
+		{instanceType: "p5.xlarge", expectedOutput: nvidiaExpectedOutput, expectedError: nil},
+		// non accelerated instance
+		{instanceType: "m5.xlarge", expectedOutput: nonAcceleratedExpectedOutput, expectedError: nil},
+	}
+	for _, test := range tests {
+		var mockConfig = []byte(`
+version = 2
+[grpc]
+address = '/run/foo/foo.sock'
+`)
+		mockOutput := []byte(`
+version=2
+
+[plugins]
+
+[plugins."io.containerd.grpc.v1.cri"]
+
+[plugins."io.containerd.grpc.v1.cri".containerd]
+default_runtime_name="nvidia"
+
+[plugins."io.containerd.grpc.v1.cri".containerd.runtimes]
+
+[plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia]
+privileged_without_host_devices=false
+runtime_engine=""
+runtime_root=""
+runtime_type="io.containerd.runc.v2"
+
+[plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia.options]
+BinaryName="/usr/bin/nvidia-container-runtime"
+`)
+
+		execCommand = MockCommandExecutor{
+			Output: mockOutput,
+			Err:    nil,
+		}
+		fileWriter = MockFileWriter{
+			Err: nil,
+		}
+
+		err := applyInstanceTypeMixins(&api.NodeConfig{
+			Status: api.NodeConfigStatus{
+				Instance: api.InstanceDetails{
+					Type: test.instanceType,
+				},
+			},
+		}, &mockConfig)
+
+		if err != test.expectedError {
+			t.Fatalf("unexpected error: %v", err)
+		}
+
+		if !reflect.DeepEqual(mockConfig, test.expectedOutput) {
+			t.Fatalf("unexpected output: %s, expecting: %s", mockConfig, test.expectedOutput)
+		}
+	}
+}

nodeadm/test/e2e/cases/containerd-oci-config-neuron/config.yaml

@@ -0,0 +1,18 @@
+---
+apiVersion: node.eks.aws/v1alpha1
+kind: NodeConfig
+spec:
+  cluster:
+    name: my-cluster
+    apiServerEndpoint: https://example.com
+    certificateAuthority: Y2VydGlmaWNhdGVBdXRob3JpdHk=
+    cidr: 10.100.0.0/16
+  containerd:
+    config: |
+      version = 2
+
+      [grpc]
+      address = "/run/foo/foo.sock"
+
+      [plugins."io.containerd.grpc.v1.cri".containerd]
+      discard_unpacked_layers = false

nodeadm/test/e2e/cases/containerd-oci-config-neuron/expected-containerd-config.toml

@@ -0,0 +1,29 @@
+root = '/var/lib/containerd'
+state = '/run/containerd'
+version = 2
+
+[grpc]
+address = '/run/foo/foo.sock'
+
+[plugins]
+[plugins.'io.containerd.grpc.v1.cri']
+sandbox_image = '602401143452.dkr.ecr.us-west-2.amazonaws.com/eks/pause:3.5'
+
+[plugins.'io.containerd.grpc.v1.cri'.cni]
+bin_dir = '/opt/cni/bin'
+conf_dir = '/etc/cni/net.d'
+
+[plugins.'io.containerd.grpc.v1.cri'.containerd]
+default_runtime_name = 'runc'
+discard_unpacked_layers = false
+
+[plugins.'io.containerd.grpc.v1.cri'.containerd.runtimes]
+[plugins.'io.containerd.grpc.v1.cri'.containerd.runtimes.runc]
+base_runtime_spec = '/etc/containerd/base-runtime-spec.json'
+runtime_type = 'io.containerd.runc.v2'
+
+[plugins.'io.containerd.grpc.v1.cri'.containerd.runtimes.runc.options]
+SystemdCgroup = true
+
+[plugins.'io.containerd.grpc.v1.cri'.registry]
+config_path = '/etc/containerd/certs.d:/etc/docker/certs.d'