feat: Adding Custom AMIFamily and validations

pkg/apis/awsnodetemplate/v1alpha1/awsnodetemplate_validation.go

@@ -16,13 +16,23 @@ package v1alpha1
 
 import (
 	"context"
+	"fmt"
+	"regexp"
 
 	"knative.dev/pkg/apis"
+
+	"github.com/aws/karpenter/pkg/cloudprovider/aws/apis/v1alpha1"
+	"github.com/aws/karpenter/pkg/utils/functional"
 )
 
 const (
 	launchTemplatePath = "launchTemplate"
 	userDataPath       = "userData"
+	amiSelectorPath    = "amiSelector"
+)
+
+var (
+	amiRegex = regexp.MustCompile("ami-[0-9a-z]+")
 )
 
 func (a *AWSNodeTemplate) Validate(ctx context.Context) (errs *apis.FieldError) {
@@ -36,6 +46,8 @@ func (a *AWSNodeTemplateSpec) validate(ctx context.Context) (errs *apis.FieldErr
 	return errs.Also(
 		a.AWS.Validate(),
 		a.validateUserData(),
+		a.validateAMISelector(),
+		a.validateAMIFamily(),
 	)
 }
 
@@ -48,3 +60,37 @@ func (a *AWSNodeTemplateSpec) validateUserData() (errs *apis.FieldError) {
 	}
 	return errs
 }
+
+func (a *AWSNodeTemplateSpec) validateAMIFamily() (errs *apis.FieldError) {
+	if a.AMIFamily == nil {
+		return nil
+	}
+	if *a.AMIFamily == v1alpha1.AMIFamilyCustom && a.AMISelector == nil {
+		errs = errs.Also(apis.ErrMissingField(amiSelectorPath))
+	}
+	return errs
+}
+
+func (a *AWSNodeTemplateSpec) validateAMISelector() (errs *apis.FieldError) {
+	if a.AMISelector == nil {
+		return nil
+	}
+	if a.LaunchTemplateName != nil {
+		errs = errs.Also(apis.ErrMultipleOneOf(amiSelectorPath, launchTemplatePath))
+	}
+	for key, value := range a.AMISelector {
+		if key == "" || value == "" {
+			errs = errs.Also(apis.ErrInvalidValue("\"\"", fmt.Sprintf("%s['%s']", amiSelectorPath, key)))
+		}
+		if key == "aws-ids" {
+			for _, amiID := range functional.SplitCommaSeparatedString(value) {
+				if !amiRegex.MatchString(amiID) {
+					fieldValue := fmt.Sprintf("\"%s\"", amiID)
+					message := fmt.Sprintf("%s['%s'] must be a valid ami-id (regex: %s)", amiSelectorPath, key, amiRegex.String())
+					errs = errs.Also(apis.ErrInvalidValue(fieldValue, message))
+				}
+			}
+		}
+	}
+	return errs
+}

pkg/cloudprovider/aws/amifamily/bootstrap/custom.go

@@ -0,0 +1,29 @@
+/*
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package bootstrap
+
+import (
+	"encoding/base64"
+
+	"github.com/aws/aws-sdk-go/aws"
+)
+
+type Custom struct {
+	Options
+}
+
+func (e Custom) Script() (string, error) {
+	return base64.StdEncoding.EncodeToString([]byte(aws.StringValue(e.Options.CustomUserData))), nil
+}

pkg/cloudprovider/aws/amifamily/custom.go

@@ -0,0 +1,62 @@
+/*
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package amifamily
+
+import (
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+
+	"github.com/aws/karpenter/pkg/apis/provisioning/v1alpha5"
+	"github.com/aws/karpenter/pkg/cloudprovider"
+	"github.com/aws/karpenter/pkg/cloudprovider/aws/amifamily/bootstrap"
+	"github.com/aws/karpenter/pkg/cloudprovider/aws/apis/v1alpha1"
+)
+
+type Custom struct {
+	*Options
+}
+
+// UserData returns the default userdata script for the AMI Family
+func (c Custom) UserData(kubeletConfig *v1alpha5.KubeletConfiguration, taints []v1.Taint, labels map[string]string, caBundle *string, _ []cloudprovider.InstanceType, customUserData *string) bootstrap.Bootstrapper {
+	return bootstrap.Custom{
+		Options: bootstrap.Options{
+			CustomUserData: customUserData,
+		},
+	}
+}
+
+func (c Custom) SSMAlias(version string, instanceType cloudprovider.InstanceType) string {
+	return "/unknown"
+}
+
+func (c Custom) DefaultBlockDeviceMappings() []*v1alpha1.BlockDeviceMapping {
+	// By returning nil, we ensure that EC2 will automatically choose the volumes defined by the AMI
+	// and we don't need to describe the AMI ourselves.
+	return nil
+}
+
+// EphemeralBlockDevice is the block device that the pods on the node will use. For an AMI of a custom family, this is unknown
+// to us.
+func (c Custom) EphemeralBlockDevice() *string {
+	return nil
+}
+
+func (c Custom) EphemeralBlockDeviceOverhead() resource.Quantity {
+	return resource.MustParse("5Gi")
+}
+
+func (c Custom) ENILimitedMemoryOverhead() bool {
+	return true
+}

pkg/cloudprovider/aws/amifamily/resolver.go

@@ -132,6 +132,8 @@ func GetAMIFamily(amiFamily *string, options *Options) AMIFamily {
 		return &Bottlerocket{Options: options}
 	case v1alpha1.AMIFamilyUbuntu:
 		return &Ubuntu{Options: options}
+	case v1alpha1.AMIFamilyCustom:
+		return &Custom{Options: options}
 	default:
 		return &AL2{Options: options}
 	}

pkg/cloudprovider/aws/apis/v1alpha1/register.go

@@ -42,10 +42,12 @@ var (
 	AMIFamilyBottlerocket = "Bottlerocket"
 	AMIFamilyAL2          = "AL2"
 	AMIFamilyUbuntu       = "Ubuntu"
+	AMIFamilyCustom       = "Custom"
 	SupportedAMIFamilies  = []string{
 		AMIFamilyBottlerocket,
 		AMIFamilyAL2,
 		AMIFamilyUbuntu,
+		AMIFamilyCustom,
 	}
 	SupportedContainerRuntimesByAMIFamily = map[string]sets.String{
 		AMIFamilyBottlerocket: sets.NewString("containerd"),

pkg/cloudprovider/aws/instancetype.go

@@ -167,8 +167,13 @@ func (i *InstanceType) memory() *resource.Quantity {
 
 // Setting ephemeral-storage to be either the default value or what is defined in blockDeviceMappings
 func (i *InstanceType) ephemeralStorage() *resource.Quantity {
-	ephemeralBlockDevice := amifamily.GetAMIFamily(i.provider.AMIFamily, &amifamily.Options{}).EphemeralBlockDevice()
-	if i.provider.BlockDeviceMappings != nil {
+	if len(i.provider.BlockDeviceMappings) != 0 {
+		if aws.StringValue(i.provider.AMIFamily) == v1alpha1.AMIFamilyCustom {
+			// For Custom AMIFamily, use the volume size of the last defined block device mapping.
+			// TODO: Consider giving better control to define which block device will be used for pods.
+			return i.provider.BlockDeviceMappings[len(i.provider.BlockDeviceMappings)-1].EBS.VolumeSize
+		}
+		ephemeralBlockDevice := amifamily.GetAMIFamily(i.provider.AMIFamily, &amifamily.Options{}).EphemeralBlockDevice()
 		for _, blockDevice := range i.provider.BlockDeviceMappings {
 			// If a block device mapping exists in the provider for the root volume, set the volume size specified in the provider
 			if *blockDevice.DeviceName == *ephemeralBlockDevice {

pkg/cloudprovider/aws/launchtemplate.go

@@ -204,6 +204,10 @@ func (p *LaunchTemplateProvider) createLaunchTemplate(ctx context.Context, optio
 }
 
 func (p *LaunchTemplateProvider) blockDeviceMappings(blockDeviceMappings []*v1alpha1.BlockDeviceMapping) []*ec2.LaunchTemplateBlockDeviceMappingRequest {
+	if len(blockDeviceMappings) == 0 {
+		// The EC2 API fails with empty slices and expects nil.
+		return nil
+	}
 	blockDeviceMappingsRequest := []*ec2.LaunchTemplateBlockDeviceMappingRequest{}
 	for _, blockDeviceMapping := range blockDeviceMappings {
 		blockDeviceMappingsRequest = append(blockDeviceMappingsRequest, &ec2.LaunchTemplateBlockDeviceMappingRequest{

pkg/cloudprovider/aws/suite_test.go

@@ -1262,6 +1262,32 @@ var _ = Describe("Allocation", func() {
 					input := fakeEC2API.CalledWithCreateLaunchTemplateInput.Pop()
 					Expect("ami-123").To(Equal(*input.LaunchTemplateData.ImageId))
 				})
+				It("should copy over userData untouched when AMIFamily is Custom", func() {
+					opts.AWSENILimitedPodDensity = false
+					provider, _ := v1alpha1.Deserialize(provisioner.Spec.Provider)
+					provider.AMIFamily = &v1alpha1.AMIFamilyCustom
+					providerRefName := strings.ToLower(randomdata.SillyName())
+					providerRef := &v1alpha5.ProviderRef{
+						Name: providerRefName,
+					}
+					nodeTemplate := test.AWSNodeTemplate(test.AWSNodeTemplateOptions{
+						UserData:    aws.String("special user data"),
+						AMISelector: map[string]string{"karpenter.sh/discovery": "my-cluster"},
+						AWS:         *provider,
+						ObjectMeta:  metav1.ObjectMeta{Name: providerRefName}})
+					fakeEC2API.DescribeImagesOutput.Set(&ec2.DescribeImagesOutput{Images: []*ec2.Image{
+						{ImageId: aws.String("ami-123"), Architecture: aws.String("x86_64")},
+					}})
+					ExpectApplied(ctx, env.Client, nodeTemplate)
+					newProvisioner := test.Provisioner(test.ProvisionerOptions{ProviderRef: providerRef})
+					ExpectApplied(ctx, env.Client, newProvisioner)
+					pod := ExpectProvisioned(ctx, env.Client, controller, test.UnschedulablePod())[0]
+					ExpectScheduled(ctx, env.Client, pod)
+					Expect(fakeEC2API.CalledWithCreateLaunchTemplateInput.Len()).To(Equal(1))
+					input := fakeEC2API.CalledWithCreateLaunchTemplateInput.Pop()
+					userData, _ := base64.StdEncoding.DecodeString(*input.LaunchTemplateData.UserData)
+					Expect("special user data").To(Equal(string(userData)))
+				})
 				It("should correctly use ami selector with specific IDs in AWSNodeTemplate", func() {
 					opts.AWSENILimitedPodDensity = false
 					provider, _ := v1alpha1.Deserialize(provisioner.Spec.Provider)
@@ -1527,6 +1553,45 @@ var _ = Describe("Allocation", func() {
 				Expect(*input.LaunchTemplateData.BlockDeviceMappings[1].Ebs.VolumeType).To(Equal("gp3"))
 				Expect(input.LaunchTemplateData.BlockDeviceMappings[1].Ebs.Iops).To(BeNil())
 			})
+			It("should not default block device mappings for custom AMIFamilies", func() {
+				provider, _ := v1alpha1.Deserialize(provisioner.Spec.Provider)
+				provider.AMIFamily = &v1alpha1.AMIFamilyCustom
+				ExpectApplied(ctx, env.Client, test.Provisioner(test.ProvisionerOptions{Provider: provider}))
+				pod := ExpectProvisioned(ctx, env.Client, controller, test.UnschedulablePod())[0]
+				ExpectScheduled(ctx, env.Client, pod)
+				Expect(fakeEC2API.CalledWithCreateLaunchTemplateInput.Len()).To(Equal(1))
+				input := fakeEC2API.CalledWithCreateLaunchTemplateInput.Pop()
+				Expect(len(input.LaunchTemplateData.BlockDeviceMappings)).To(Equal(0))
+			})
+			It("should use custom block device mapping for custom AMIFamilies", func() {
+				provider, _ := v1alpha1.Deserialize(provisioner.Spec.Provider)
+				provider.AMIFamily = &v1alpha1.AMIFamilyCustom
+				provider.BlockDeviceMappings = []*v1alpha1.BlockDeviceMapping{
+					{
+						DeviceName: aws.String("/dev/xvda"),
+						EBS: &v1alpha1.BlockDevice{
+							DeleteOnTermination: aws.Bool(true),
+							Encrypted:           aws.Bool(true),
+							VolumeType:          aws.String("io2"),
+							VolumeSize:          resource.NewScaledQuantity(40, resource.Giga),
+							IOPS:                aws.Int64(10_000),
+							KMSKeyID:            aws.String("arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab"),
+						},
+					},
+				}
+				ExpectApplied(ctx, env.Client, test.Provisioner(test.ProvisionerOptions{Provider: provider}))
+				pod := ExpectProvisioned(ctx, env.Client, controller, test.UnschedulablePod())[0]
+				ExpectScheduled(ctx, env.Client, pod)
+				Expect(fakeEC2API.CalledWithCreateLaunchTemplateInput.Len()).To(Equal(1))
+				input := fakeEC2API.CalledWithCreateLaunchTemplateInput.Pop()
+				Expect(len(input.LaunchTemplateData.BlockDeviceMappings)).To(Equal(1))
+				Expect(*input.LaunchTemplateData.BlockDeviceMappings[0].Ebs.VolumeSize).To(Equal(int64(40)))
+				Expect(*input.LaunchTemplateData.BlockDeviceMappings[0].Ebs.VolumeType).To(Equal("io2"))
+				Expect(*input.LaunchTemplateData.BlockDeviceMappings[0].Ebs.Iops).To(Equal(int64(10_000)))
+				Expect(*input.LaunchTemplateData.BlockDeviceMappings[0].Ebs.DeleteOnTermination).To(BeTrue())
+				Expect(*input.LaunchTemplateData.BlockDeviceMappings[0].Ebs.Encrypted).To(BeTrue())
+				Expect(*input.LaunchTemplateData.BlockDeviceMappings[0].Ebs.KmsKeyId).To(Equal("arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab"))
+			})
 		})
 		Context("Ephemeral Storage", func() {
 			It("should pack pods when a daemonset has an ephemeral-storage request", func() {
@@ -1640,6 +1705,36 @@ var _ = Describe("Allocation", func() {
 				// capacity isn't recorded on the node any longer, but we know the pod should schedule
 				ExpectScheduled(ctx, env.Client, pod)
 			})
+			It("should pack pods using blockdevicemappings for Custom AMIFamily", func() {
+				provider, _ = v1alpha1.Deserialize(provisioner.Spec.Provider)
+				provider.AMIFamily = &v1alpha1.AMIFamilyCustom
+				provider.BlockDeviceMappings = []*v1alpha1.BlockDeviceMapping{
+					{
+						DeviceName: aws.String("/dev/xvda"),
+						EBS: &v1alpha1.BlockDevice{
+							VolumeSize: resource.NewScaledQuantity(20, resource.Giga),
+						},
+					},
+					{
+						DeviceName: aws.String("/dev/xvdb"),
+						EBS: &v1alpha1.BlockDevice{
+							VolumeSize: resource.NewScaledQuantity(40, resource.Giga),
+						},
+					},
+				}
+				ExpectApplied(ctx, env.Client, test.Provisioner(test.ProvisionerOptions{Provider: provider}))
+				pod := ExpectProvisioned(ctx, env.Client, controller,
+					test.UnschedulablePod(test.PodOptions{ResourceRequirements: v1.ResourceRequirements{
+						Requests: map[v1.ResourceName]resource.Quantity{
+							// this pod can only be satisifed if `/dev/xvdb` will house all the pods.
+							v1.ResourceEphemeralStorage: resource.MustParse("25Gi"),
+						},
+					},
+					}))[0]
+
+				// capacity isn't recorded on the node any longer, but we know the pod should schedule
+				ExpectScheduled(ctx, env.Client, pod)
+			})
 		})
 	})
 	Context("Defaulting", func() {

test/suites/integration/suite_test.go

@@ -23,13 +23,15 @@ import (
 )
 
 var env *environment.Environment
+var customAMI string
 
 func TestIntegration(t *testing.T) {
 	RegisterFailHandler(Fail)
 	BeforeSuite(func() {
 		var err error
 		env, err = environment.NewEnvironment(t)
 		Expect(err).ToNot(HaveOccurred())
+		customAMI = selectCustomAMI("/aws/service/eks/optimized-ami/%s/amazon-linux-2/recommended/image_id")
 	})
 	RunSpecs(t, "Integration")
 }
@@ -152,14 +154,12 @@ var _ = Describe("Sanity Checks", func() {
 	})
 	Context("LaunchTemplates", func() {
 		It("should use the AMI defined by the AMI Selector", func() {
-			// TODO - Cache the AMI so it can be re-used across multiple tests.
-			amiUnderTest := selectCustomAMI("/aws/service/eks/optimized-ami/%s/amazon-linux-2/recommended/image_id")
 			provider := test.AWSNodeTemplate(test.AWSNodeTemplateOptions{AWS: v1alpha1.AWS{
 				SecurityGroupSelector: map[string]string{"karpenter.sh/discovery": env.Options.ClusterName},
 				SubnetSelector:        map[string]string{"karpenter.sh/discovery": env.Options.ClusterName},
 				AMIFamily:             &v1alpha1.AMIFamilyAL2,
 			},
-				AMISelector: map[string]string{"aws-ids": amiUnderTest}, // TODO - Retrieve recommended EKS AMI and use here.
+				AMISelector: map[string]string{"aws-ids": customAMI},
 			})
 			provisioner := test.Provisioner(test.ProvisionerOptions{ProviderRef: &v1alpha5.ProviderRef{Name: provider.Name}})
 			pod := test.Pod()
@@ -177,7 +177,37 @@ var _ = Describe("Sanity Checks", func() {
 			instance, _ := env.Ec2API.DescribeInstances(&ec2.DescribeInstancesInput{
 				InstanceIds: aws.StringSlice([]string{instanceID}),
 			})
-			Expect(*instance.Reservations[0].Instances[0].ImageId).To(Equal(amiUnderTest))
+			Expect(*instance.Reservations[0].Instances[0].ImageId).To(Equal(customAMI))
+			env.ExpectDeleted(pod)
+			env.EventuallyExpectScaleDown()
+			env.ExpectNoCrashes()
+		})
+		It("should support Custom AMIFamily with AMI Selectors", func() {
+			provider := test.AWSNodeTemplate(test.AWSNodeTemplateOptions{AWS: v1alpha1.AWS{
+				SecurityGroupSelector: map[string]string{"karpenter.sh/discovery": env.Options.ClusterName},
+				SubnetSelector:        map[string]string{"karpenter.sh/discovery": env.Options.ClusterName},
+				AMIFamily:             &v1alpha1.AMIFamilyCustom,
+			},
+				AMISelector: map[string]string{"aws-ids": customAMI},
+				UserData:    aws.String(fmt.Sprintf("#!/bin/bash\n/etc/eks/bootstrap.sh '%s'", env.Options.ClusterName)),
+			})
+			provisioner := test.Provisioner(test.ProvisionerOptions{ProviderRef: &v1alpha5.ProviderRef{Name: provider.Name}})
+			pod := test.Pod()
+
+			env.ExpectCreatedNodeCount("==", 0)
+			env.ExpectCreated(pod, provider, provisioner)
+
+			env.EventuallyExpectHealthy(pod)
+			env.ExpectCreatedNodeCount("==", 1)
+
+			var node v1.Node
+			env.Client.Get(env.Context, types.NamespacedName{Name: pod.Spec.NodeName}, &node)
+			providerIDSplit := strings.Split(node.Spec.ProviderID, "/")
+			instanceID := providerIDSplit[len(providerIDSplit)-1]
+			instance, _ := env.Ec2API.DescribeInstances(&ec2.DescribeInstancesInput{
+				InstanceIds: aws.StringSlice([]string{instanceID}),
+			})
+			Expect(*instance.Reservations[0].Instances[0].ImageId).To(Equal(customAMI))
 			env.ExpectDeleted(pod)
 			env.EventuallyExpectScaleDown()
 			env.ExpectNoCrashes()