feat: Add support for custom ssm parameters in amiSelectorTerms

charts/karpenter-crd/templates/karpenter.k8s.aws_ec2nodeclasses.yaml

@@ -114,6 +114,9 @@ spec:
                           Owner is the owner for the ami.
                           You can specify a combination of AWS account IDs, "self", "amazon", and "aws-marketplace"
                         type: string
+                      ssmParameter:
+                        description: SSMParameter is the name (or ARN) of the SSM parameter containing the Image ID.
+                        type: string
                       tags:
                         additionalProperties:
                           type: string
@@ -130,8 +133,8 @@ spec:
                   minItems: 1
                   type: array
                   x-kubernetes-validations:
-                    - message: expected at least one, got none, ['tags', 'id', 'name', 'alias']
-                      rule: self.all(x, has(x.tags) || has(x.id) || has(x.name) || has(x.alias))
+                    - message: expected at least one, got none, ['tags', 'id', 'name', 'alias', 'ssmParameter']
+                      rule: self.all(x, has(x.tags) || has(x.id) || has(x.name) || has(x.alias) || has(x.ssmParameter))
                     - message: '''id'' is mutually exclusive, cannot be set with a combination of other fields in amiSelectorTerms'
                       rule: '!self.exists(x, has(x.id) && (has(x.alias) || has(x.tags) || has(x.name) || has(x.owner)))'
                     - message: '''alias'' is mutually exclusive, cannot be set with a combination of other fields in amiSelectorTerms'

pkg/apis/crds/karpenter.k8s.aws_ec2nodeclasses.yaml

@@ -111,6 +111,9 @@ spec:
                           Owner is the owner for the ami.
                           You can specify a combination of AWS account IDs, "self", "amazon", and "aws-marketplace"
                         type: string
+                      ssmParameter:
+                        description: SSMParameter is the name (or ARN) of the SSM parameter containing the Image ID.
+                        type: string
                       tags:
                         additionalProperties:
                           type: string
@@ -127,8 +130,8 @@ spec:
                   minItems: 1
                   type: array
                   x-kubernetes-validations:
-                    - message: expected at least one, got none, ['tags', 'id', 'name', 'alias']
-                      rule: self.all(x, has(x.tags) || has(x.id) || has(x.name) || has(x.alias))
+                    - message: expected at least one, got none, ['tags', 'id', 'name', 'alias', 'ssmParameter']
+                      rule: self.all(x, has(x.tags) || has(x.id) || has(x.name) || has(x.alias) || has(x.ssmParameter))
                     - message: '''id'' is mutually exclusive, cannot be set with a combination of other fields in amiSelectorTerms'
                       rule: '!self.exists(x, has(x.id) && (has(x.alias) || has(x.tags) || has(x.name) || has(x.owner)))'
                     - message: '''alias'' is mutually exclusive, cannot be set with a combination of other fields in amiSelectorTerms'

pkg/apis/v1/ec2nodeclass.go

@@ -55,7 +55,7 @@ type EC2NodeClassSpec struct {
 	// +optional
 	AssociatePublicIPAddress *bool `json:"associatePublicIPAddress,omitempty"`
 	// AMISelectorTerms is a list of or ami selector terms. The terms are ORed.
-	// +kubebuilder:validation:XValidation:message="expected at least one, got none, ['tags', 'id', 'name', 'alias']",rule="self.all(x, has(x.tags) || has(x.id) || has(x.name) || has(x.alias))"
+	// +kubebuilder:validation:XValidation:message="expected at least one, got none, ['tags', 'id', 'name', 'alias', 'ssmParameter']",rule="self.all(x, has(x.tags) || has(x.id) || has(x.name) || has(x.alias) || has(x.ssmParameter))"
 	// +kubebuilder:validation:XValidation:message="'id' is mutually exclusive, cannot be set with a combination of other fields in amiSelectorTerms",rule="!self.exists(x, has(x.id) && (has(x.alias) || has(x.tags) || has(x.name) || has(x.owner)))"
 	// +kubebuilder:validation:XValidation:message="'alias' is mutually exclusive, cannot be set with a combination of other fields in amiSelectorTerms",rule="!self.exists(x, has(x.alias) && (has(x.id) || has(x.tags) || has(x.name) || has(x.owner)))"
 	// +kubebuilder:validation:XValidation:message="'alias' is mutually exclusive, cannot be set with a combination of other amiSelectorTerms",rule="!(self.exists(x, has(x.alias)) && self.size() != 1)"
@@ -227,6 +227,9 @@ type AMISelectorTerm struct {
 	// You can specify a combination of AWS account IDs, "self", "amazon", and "aws-marketplace"
 	// +optional
 	Owner string `json:"owner,omitempty"`
+	//SSMParameter is the name (or ARN) of the SSM parameter containing the Image ID.
+	// +optional
+	SSMParameter string `json:"ssmParameter,omitempty"`
 }
 
 // KubeletConfiguration defines args to be used when configuring kubelet on provisioned nodes.

pkg/errors/errors.go

@@ -40,6 +40,7 @@ var (
 		"InvalidLaunchTemplateId.NotFound",
 		"QueueDoesNotExist",
 		"NoSuchEntity",
+		"ParameterNotFound",
 	)
 	alreadyExistsErrorCodes = sets.New[string](
 		"EntityAlreadyExists",

pkg/providers/amifamily/ami.go

@@ -17,6 +17,7 @@ package amifamily
 import (
 	"context"
 	"fmt"
+	"strings"
 	"sync"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -27,10 +28,13 @@ import (
 	"github.com/samber/lo"
 	"k8s.io/utils/clock"
 
+	"github.com/aws/karpenter-provider-aws/pkg/errors"
+
 	v1 "github.com/aws/karpenter-provider-aws/pkg/apis/v1"
 	sdk "github.com/aws/karpenter-provider-aws/pkg/aws"
 	"github.com/aws/karpenter-provider-aws/pkg/providers/version"
 
+	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/karpenter/pkg/cloudprovider"
 	"sigs.k8s.io/karpenter/pkg/scheduling"
 
@@ -77,6 +81,7 @@ func (p *DefaultProvider) List(ctx context.Context, nodeClass *v1.EC2NodeClass)
 	return amis, nil
 }
 
+//nolint:gocyclo
 func (p *DefaultProvider) DescribeImageQueries(ctx context.Context, nodeClass *v1.EC2NodeClass) ([]DescribeImageQuery, error) {
 	// Aliases are mutually exclusive, both on the term level and field level within a term.
 	// This is enforced by a CEL validation, we will treat this as an invariant.
@@ -95,6 +100,23 @@ func (p *DefaultProvider) DescribeImageQueries(ctx context.Context, nodeClass *v
 		switch {
 		case term.ID != "":
 			idFilter.Values = append(idFilter.Values, term.ID)
+		case term.SSMParameter != "":
+			imageID, err := p.ssmProvider.Get(ctx, ssm.Parameter{
+				Name: term.SSMParameter,
+				Type: ssm.CustomParameterType,
+			})
+			if err != nil {
+				if !errors.IsNotFound(err) {
+					return []DescribeImageQuery{}, fmt.Errorf("resolving ssm parameter, %w", err)
+				}
+				log.FromContext(ctx).WithValues("ssmParameter", term.SSMParameter).V(1).Error(err, "parameter not found")
+				continue
+			}
+			if !strings.HasPrefix(imageID, "ami-") {
+				log.FromContext(ctx).WithValues("ssmParameter", term.SSMParameter, "id", imageID).V(1).Error(nil, "parameter value is an invalid AMI ID")
+				continue
+			}
+			idFilter.Values = append(idFilter.Values, imageID)
 		default:
 			query := DescribeImageQuery{
 				Owners: lo.Ternary(term.Owner != "", []string{term.Owner}, []string{}),

pkg/providers/amifamily/suite_test.go

@@ -59,10 +59,10 @@ func TestAWS(t *testing.T) {
 }
 
 const (
-	amd64AMI       = "amd64-ami-id"
-	arm64AMI       = "arm64-ami-id"
-	amd64NvidiaAMI = "amd64-nvidia-ami-id"
-	arm64NvidiaAMI = "arm64-nvidia-ami-id"
+	amd64AMI       = "ami-id-amd64"
+	arm64AMI       = "ami-id-arm64"
+	amd64NvidiaAMI = "ami-id-amd64-nvidia"
+	arm64NvidiaAMI = "ami-id-arm64-nvidia"
 )
 
 var _ = BeforeSuite(func() {
@@ -78,7 +78,7 @@ var _ = BeforeEach(func() {
 		Images: []ec2types.Image{
 			{
 				Name:         aws.String(amd64AMI),
-				ImageId:      aws.String("amd64-ami-id"),
+				ImageId:      aws.String("ami-id-amd64"),
 				CreationDate: aws.String(time.Time{}.Format(time.RFC3339)),
 				Architecture: "x86_64",
 				Tags: []ec2types.Tag{
@@ -89,7 +89,7 @@ var _ = BeforeEach(func() {
 			},
 			{
 				Name:         aws.String(arm64AMI),
-				ImageId:      aws.String("arm64-ami-id"),
+				ImageId:      aws.String("ami-id-arm64"),
 				CreationDate: aws.String(time.Time{}.Add(time.Minute).Format(time.RFC3339)),
 				Architecture: "arm64",
 				Tags: []ec2types.Tag{
@@ -100,7 +100,7 @@ var _ = BeforeEach(func() {
 			},
 			{
 				Name:         aws.String(amd64NvidiaAMI),
-				ImageId:      aws.String("amd64-nvidia-ami-id"),
+				ImageId:      aws.String("ami-id-amd64-nvidia"),
 				CreationDate: aws.String(time.Time{}.Add(2 * time.Minute).Format(time.RFC3339)),
 				Architecture: "x86_64",
 				Tags: []ec2types.Tag{
@@ -111,7 +111,7 @@ var _ = BeforeEach(func() {
 			},
 			{
 				Name:         aws.String(arm64NvidiaAMI),
-				ImageId:      aws.String("arm64-nvidia-ami-id"),
+				ImageId:      aws.String("ami-id-arm64-nvidia"),
 				CreationDate: aws.String(time.Time{}.Add(2 * time.Minute).Format(time.RFC3339)),
 				Architecture: "arm64",
 				Tags: []ec2types.Tag{
@@ -189,13 +189,14 @@ var _ = Describe("AMIProvider", func() {
 		Expect(err).ToNot(HaveOccurred())
 		Expect(amis).To(HaveLen(1))
 	})
+
 	It("should not cause data races when calling Get() simultaneously", func() {
 		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{
 			{
-				ID: "amd64-ami-id",
+				ID: "ami-id-amd64",
 			},
 			{
-				ID: "arm64-ami-id",
+				ID: "ami-id-arm64",
 			},
 		}
 		wg := sync.WaitGroup{}
@@ -213,15 +214,15 @@ var _ = Describe("AMIProvider", func() {
 				Expect(images).To(BeEquivalentTo([]amifamily.AMI{
 					{
 						Name:         arm64AMI,
-						AmiID:        "arm64-ami-id",
+						AmiID:        "ami-id-arm64",
 						CreationDate: time.Time{}.Add(time.Minute).Format(time.RFC3339),
 						Requirements: scheduling.NewLabelRequirements(map[string]string{
 							corev1.LabelArchStable: karpv1.ArchitectureArm64,
 						}),
 					},
 					{
 						Name:         amd64AMI,
-						AmiID:        "amd64-ami-id",
+						AmiID:        "ami-id-amd64",
 						CreationDate: time.Time{}.Format(time.RFC3339),
 						Requirements: scheduling.NewLabelRequirements(map[string]string{
 							corev1.LabelArchStable: karpv1.ArchitectureAmd64,
@@ -312,7 +313,7 @@ var _ = Describe("AMIProvider", func() {
 		BeforeEach(func() {
 			img = ec2types.Image{
 				Name:         aws.String(amd64AMI),
-				ImageId:      aws.String("amd64-ami-id"),
+				ImageId:      aws.String("ami-id-amd64"),
 				CreationDate: aws.String(time.Now().Format(time.RFC3339)),
 				Architecture: "x86_64",
 				Tags: []ec2types.Tag{
@@ -848,6 +849,52 @@ var _ = Describe("AMIProvider", func() {
 				},
 			))
 		})
+		It("should succeed to resolve AMIs that use an SSM parameter", func() {
+			nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{
+				SSMParameter: fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2/recommended/image_id", version),
+			}}
+			awsEnv.SSMAPI.Parameters = map[string]string{
+				fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2/recommended/image_id", version): amd64AMI,
+			}
+			amis, err := awsEnv.AMIProvider.List(ctx, nodeClass)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(amis).To(HaveLen(1))
+			Expect(amis[0].AmiID).To(Equal("ami-id-amd64"))
+			Expect(amis[0].Name).To(Equal(amd64AMI))
+		})
+		It("should succeed to resolve AMIs that use a custom SSM parameter", func() {
+			customParameter := "/my/custom/ami/parameter"
+			nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{
+				SSMParameter: customParameter,
+			}}
+			awsEnv.SSMAPI.Parameters = map[string]string{
+				customParameter: amd64AMI,
+			}
+			amis, err := awsEnv.AMIProvider.List(ctx, nodeClass)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(amis).To(HaveLen(1))
+			Expect(amis[0].AmiID).To(Equal("ami-id-amd64"))
+			Expect(amis[0].Name).To(Equal(amd64AMI))
+		})
+		It("should not throw an error if SSM parameter is not found", func() {
+			customParameter := "/my/custom/ami/parameter"
+			nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{
+				SSMParameter: customParameter,
+			}}
+			amis, err := awsEnv.AMIProvider.List(ctx, nodeClass)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(amis).To(HaveLen(0))
+		})
+		It("should throw an error if SSM parameter returns a different error", func() {
+			customParameter := "/my/custom/ami/parameter"
+			nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{
+				SSMParameter: customParameter,
+			}}
+			awsEnv.SSMAPI.WantErr = fmt.Errorf("some error")
+			amis, err := awsEnv.AMIProvider.List(ctx, nodeClass)
+			Expect(err).To(HaveOccurred())
+			Expect(amis).To(HaveLen(0))
+		})
 	})
 })
 

pkg/providers/ssm/provider.go

@@ -53,10 +53,10 @@ func (p *DefaultProvider) Get(ctx context.Context, parameter Parameter) (string,
 	if err != nil {
 		return "", fmt.Errorf("getting ssm parameter %q, %w", parameter.Name, err)
 	}
-	p.cache.SetDefault(parameter.CacheKey(), CacheEntry{
+	p.cache.Set(parameter.CacheKey(), CacheEntry{
 		Parameter: parameter,
 		Value:     lo.FromPtr(result.Parameter.Value),
-	})
+	}, parameter.GetCacheDuration())
 	log.FromContext(ctx).WithValues("parameter", parameter.Name, "value", result.Parameter.Value).Info("discovered ssm parameter")
 	return lo.FromPtr(result.Parameter.Value), nil
 }

pkg/providers/ssm/types.go

@@ -15,12 +15,19 @@ limitations under the License.
 package ssm
 
 import (
+	"time"
+
 	"github.com/aws/aws-sdk-go-v2/service/ssm"
 	"github.com/samber/lo"
 )
 
+const (
+	CustomParameterType = "custom"
+)
+
 type Parameter struct {
 	Name string
+	Type string
 	// IsMutable indicates if the value associated with an SSM parameter is expected to change. An example of a mutable
 	// parameter would be any of the "latest" or "recommended" AMI parameters which are updated each time a new AMI is
 	// released. On the otherhand, we would consider a parameter parameter for a specific AMI version to be immutable.
@@ -37,6 +44,14 @@ func (p *Parameter) CacheKey() string {
 	return p.Name
 }
 
+// GetCacheDuration returns the appropriate cache duration based on the parameter type
+func (p Parameter) GetCacheDuration() time.Duration {
+	if p.Type == CustomParameterType {
+		return 5 * time.Minute
+	}
+	return 24 * time.Hour
+}
+
 type CacheEntry struct {
 	Parameter Parameter
 	Value     string

test/suites/ami/suite_test.go

@@ -71,10 +71,13 @@ var _ = AfterEach(func() { env.Cleanup() })
 var _ = AfterEach(func() { env.AfterEach() })
 
 var _ = Describe("AMI", func() {
+	var ssmPath string
 	var customAMI string
 	var deprecatedAMI string
+
 	BeforeEach(func() {
-		customAMI = env.GetAMIBySSMPath(fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2023/x86_64/standard/recommended/image_id", env.K8sVersion()))
+		ssmPath = fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2023/x86_64/standard/recommended/image_id", env.K8sVersion())
+		customAMI = env.GetAMIBySSMPath(ssmPath)
 		deprecatedAMI = env.GetDeprecatedAMI(customAMI, "AL2023")
 	})
 
@@ -159,6 +162,21 @@ var _ = Describe("AMI", func() {
 
 		env.ExpectInstance(pod.Spec.NodeName).To(HaveField("ImageId", HaveValue(Equal(customAMI))))
 	})
+	It("should support ssm parameters by ARN", func() {
+		nodeClass.Spec.AMIFamily = lo.ToPtr(v1.AMIFamilyAL2023)
+		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{
+			{
+				SSMParameter: fmt.Sprintf("arn:aws:ssm:%s::parameter%s", env.Region, ssmPath),
+			},
+		}
+		pod := coretest.Pod()
+
+		env.ExpectCreated(pod, nodeClass, nodePool)
+		env.EventuallyExpectHealthy(pod)
+		env.ExpectCreatedNodeCount("==", 1)
+
+		env.ExpectInstance(pod.Spec.NodeName).To(HaveField("ImageId", HaveValue(Equal(customAMI))))
+	})
 	It("should support launching nodes with a deprecated ami", func() {
 		nodeClass.Spec.AMIFamily = lo.ToPtr(v1.AMIFamilyAL2023)
 		nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{
@@ -253,6 +271,16 @@ var _ = Describe("AMI", func() {
 			ExpectStatusConditions(env, env.Client, 1*time.Minute, nodeClass, status.Condition{Type: v1.ConditionTypeAMIsReady, Status: metav1.ConditionTrue})
 			ExpectStatusConditions(env, env.Client, 1*time.Minute, nodeClass, status.Condition{Type: status.ConditionReady, Status: metav1.ConditionTrue})
 		})
+		It("should have the EC2NodeClass status for AMIs using public ssm parameter ARN", func() {
+			nodeClass.Spec.AMIFamily = lo.ToPtr(v1.AMIFamilyAL2023)
+			nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{SSMParameter: fmt.Sprintf("arn:aws:ssm:%s::parameter%s", env.Region, ssmPath)}}
+			env.ExpectCreated(nodeClass)
+			nc := EventuallyExpectAMIsToExist(nodeClass)
+			Expect(len(nc.Status.AMIs)).To(BeNumerically("==", 1))
+			Expect(nc.Status.AMIs[0].ID).To(Equal(customAMI))
+			ExpectStatusConditions(env, env.Client, 1*time.Minute, nodeClass, status.Condition{Type: v1.ConditionTypeAMIsReady, Status: metav1.ConditionTrue})
+			ExpectStatusConditions(env, env.Client, 1*time.Minute, nodeClass, status.Condition{Type: status.ConditionReady, Status: metav1.ConditionTrue})
+		})
 		It("should have ec2nodeClass status as not ready since AMI was not resolved", func() {
 			nodeClass.Spec.AMIFamily = lo.ToPtr(v1.AMIFamilyAL2023)
 			nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{ID: "ami-123"}}