chore: revert to using ssm:GetParameter

pkg/cache/cache.go

@@ -37,7 +37,7 @@ const (
 	AssociatePublicIPAddressTTL = 5 * time.Minute
 	// SSMGetParametersByPathTTL is the time to drop SSM Parameters by path data. This only queries EKS Optimized AMI
 	// releases, so we should expect this to be updated relatively infrequently.
-	SSMGetParametersByPathTTL = 30 * time.Minute
+	SSMGetParametersByPathTTL = 24 * time.Hour
 )
 
 const (

pkg/fake/ssmapi.go

@@ -17,120 +17,68 @@ package fake
 import (
 	"context"
 	"fmt"
-	"log"
-	"regexp"
-	"strings"
 
 	"github.com/Pallinder/go-randomdata"
 	"github.com/samber/lo"
 
-	"github.com/aws/karpenter-provider-aws/pkg/providers/version"
-
 	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/ssm"
 	"github.com/aws/aws-sdk-go/service/ssm/ssmiface"
 )
 
 type SSMAPI struct {
 	ssmiface.SSMAPI
-	Parameters                map[string]string
-	GetParametersByPathOutput *ssm.GetParametersByPathOutput
-	WantErr                   error
+	Parameters         map[string]string
+	GetParameterOutput *ssm.GetParameterOutput
+	WantErr            error
 
-	defaultParametersForPath map[string][]*ssm.Parameter
+	defaultParameters map[string]string
 }
 
 func NewSSMAPI() *SSMAPI {
 	return &SSMAPI{
-		defaultParametersForPath: map[string][]*ssm.Parameter{},
+		defaultParameters: map[string]string{},
 	}
 }
 
-func (a SSMAPI) GetParametersByPathPagesWithContext(_ context.Context, input *ssm.GetParametersByPathInput, f func(*ssm.GetParametersByPathOutput, bool) bool, _ ...request.Option) error {
-	if !lo.FromPtr(input.Recursive) {
-		log.Fatalf("fake SSM API currently only supports GetParametersByPathPagesWithContext when recursive is true")
-	}
+func (a SSMAPI) GetParameterWithContext(_ context.Context, input *ssm.GetParameterInput, _ ...request.Option) (*ssm.GetParameterOutput, error) {
+	parameter := lo.FromPtr(input.Name)
 	if a.WantErr != nil {
-		return a.WantErr
+		return &ssm.GetParameterOutput{}, a.WantErr
 	}
-	if a.GetParametersByPathOutput != nil {
-		f(a.GetParametersByPathOutput, true)
-		return nil
+	if a.GetParameterOutput != nil {
+		return a.GetParameterOutput, nil
 	}
 	if len(a.Parameters) != 0 {
-		f(&ssm.GetParametersByPathOutput{
-			Parameters: lo.FilterMap(lo.Entries(a.Parameters), func(p lo.Entry[string, string], _ int) (*ssm.Parameter, bool) {
-				// The parameter does not start with the path
-				if !strings.HasPrefix(p.Key, lo.FromPtr(input.Path)) {
-					return nil, false
-				}
-				// The parameter starts with the input path, but the last segment of the input path is only a subset of the matching segment of the parameters path.
-				// Ex: "/aws/service/eks-optimized-ami/amazon-linux-2" is a prefix for "/aws/service/eks-optimized-ami/amazon-linux-2-gpu/..." but we shouldn't match
-				if strings.TrimPrefix(p.Key, lo.FromPtr(input.Path))[0] != '/' {
-					return nil, false
-				}
-				return &ssm.Parameter{
-					Name:  lo.ToPtr(p.Key),
-					Value: lo.ToPtr(p.Value),
-				}, true
-			}),
-		}, true)
-		return nil
-	}
-	if params := a.getDefaultParametersForPath(lo.FromPtr(input.Path)); params != nil {
-		f(&ssm.GetParametersByPathOutput{Parameters: params}, true)
-		return nil
+		value, ok := a.Parameters[parameter]
+		if !ok {
+			return &ssm.GetParameterOutput{}, fmt.Errorf("parameter %q not found", lo.FromPtr(input.Name))
+		}
+		return &ssm.GetParameterOutput{
+			Parameter: &ssm.Parameter{
+				Name:  lo.ToPtr(parameter),
+				Value: lo.ToPtr(value),
+			},
+		}, nil
 	}
-	return fmt.Errorf("path %q does not exist", lo.FromPtr(input.Path))
-}
 
-func (a SSMAPI) getDefaultParametersForPath(path string) []*ssm.Parameter {
-	// If we've already generated default parameters, return the same parameters across calls. This ensures we don't
-	// drift due to different results from one call to the next.
-	if params, ok := a.defaultParametersForPath[path]; ok {
-		return params
+	// Cache default parameters that was successive calls for the same parameter return the same result
+	value, ok := a.defaultParameters[parameter]
+	if !ok {
+		value = fmt.Sprintf("ami-%s", randomdata.Alphanumeric(16))
+		a.defaultParameters[parameter] = value
 	}
-	suffixes := map[string][]string{
-		`^\/aws\/service\/eks/optimized-ami\/.*\/amazon-linux-2$`:       []string{"recommended/image_id"},
-		`^\/aws\/service\/eks/optimized-ami\/.*\/amazon-linux-2-arm64$`: []string{"recommended/image_id"},
-		`^\/aws\/service\/eks/optimized-ami\/.*\/amazon-linux-2-gpu$`:   []string{"recommended/image_id"},
-		`^\/aws\/service\/eks/optimized-ami\/.*\/amazon-linux-2023$`: []string{
-			"x86_64/standard/recommended/image_id",
-			"arm64/standard/recommended/image_id",
-			"x86_64/nvidia/recommended/image_id",
-			"arm64/nvidia/recommended/image_id",
-			"x86_64/neuron/recommended/image_id",
-			"arm64/neuron/recommended/image_id",
+	return &ssm.GetParameterOutput{
+		Parameter: &ssm.Parameter{
+			Name:  lo.ToPtr(parameter),
+			Value: lo.ToPtr(value),
 		},
-		`\/aws\/service\/bottlerocket\/aws-k8s-.*`: []string{
-			"x86_64/latest/image_id",
-			"arm64/latest/image_id",
-		},
-		`\/aws\/service\/ami-windows-latest`: lo.FlatMap(version.SupportedK8sVersions(), func(version string, _ int) []string {
-			return []string{
-				fmt.Sprintf("Windows_Server-2019-English-Core-EKS_Optimized-%s/image_id", version),
-				fmt.Sprintf("Windows_Server-2022-English-Core-EKS_Optimized-%s/image_id", version),
-			}
-		}),
-	}
-	for matchStr, suffixes := range suffixes {
-		if !regexp.MustCompile(matchStr).MatchString(path) {
-			continue
-		}
-		params := lo.Map(suffixes, func(suffix string, _ int) *ssm.Parameter {
-			return &ssm.Parameter{
-				Name:  lo.ToPtr(fmt.Sprintf("%s/%s", path, suffix)),
-				Value: lo.ToPtr(fmt.Sprintf("ami-%s", randomdata.Alphanumeric(16))),
-			}
-		})
-		a.defaultParametersForPath[path] = params
-		return params
-	}
-	return nil
+	}, nil
 }
 
 func (a *SSMAPI) Reset() {
-	a.GetParametersByPathOutput = nil
 	a.Parameters = nil
+	a.GetParameterOutput = nil
 	a.WantErr = nil
+	a.defaultParameters = map[string]string{}
 }

pkg/providers/amifamily/al2.go

@@ -17,16 +17,13 @@ package amifamily
 import (
 	"context"
 	"fmt"
-	"regexp"
-	"strings"
 
 	"github.com/aws/aws-sdk-go/aws"
 	corev1 "k8s.io/api/core/v1"
 
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/samber/lo"
 
-	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/karpenter/pkg/scheduling"
 
 	v1 "github.com/aws/karpenter-provider-aws/pkg/apis/v1"
@@ -43,59 +40,46 @@ type AL2 struct {
 }
 
 func (a AL2) DescribeImageQuery(ctx context.Context, ssmProvider ssm.Provider, k8sVersion string, amiVersion string) (DescribeImageQuery, error) {
-	imageIDs := make([]*string, 0, 5)
-	requirements := make(map[string][]scheduling.Requirements)
-	// Example Paths:
-	// - Latest EKS 1.30 Standard Image: /aws/service/eks/optimized-ami/1.30/amazon-linux-2/recommended/image_id
-	// - Specific EKS 1.30 GPU Image: /aws/service/eks/optimized-ami/1.30/amazon-linux-2-gpu/amazon-eks-node-1.30-v20240625/image_id
-	for rootPath, variants := range map[string][]Variant{
-		fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2", k8sVersion):       {VariantStandard},
-		fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2-arm64", k8sVersion): {VariantStandard},
-		fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2-gpu", k8sVersion):   {VariantNeuron, VariantNvidia},
+	ids := map[string][]Variant{}
+	for path, variants := range map[string][]Variant{
+		fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2/%s/image_id", k8sVersion, lo.Ternary(
+			amiVersion == AMIVersionLatest,
+			"recommended",
+			fmt.Sprintf("amazon-eks-node-%s-%s", k8sVersion, amiVersion),
+		)): {VariantStandard},
+		fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2-arm64/%s/image_id", k8sVersion, lo.Ternary(
+			amiVersion == AMIVersionLatest,
+			"recommended",
+			fmt.Sprintf("amazon-eks-arm64-node-%s-%s", k8sVersion, amiVersion),
+		)): {VariantStandard},
+		fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2-gpu/%s/image_id", k8sVersion, lo.Ternary(
+			amiVersion == AMIVersionLatest,
+			"recommended",
+			fmt.Sprintf("amazon-eks-gpu-node-%s-%s", k8sVersion, amiVersion),
+		)): {VariantNeuron, VariantNvidia},
 	} {
-		results, err := ssmProvider.List(ctx, rootPath)
+		imageID, err := ssmProvider.Get(ctx, path)
 		if err != nil {
-			log.FromContext(ctx).WithValues("path", rootPath, "family", "al2").Error(err, "discovering AMIs from ssm")
 			continue
 		}
-		for path, value := range results {
-			pathComponents := strings.Split(path, "/")
-			// Only select image_id paths which match the desired AMI version
-			if len(pathComponents) != 9 || pathComponents[8] != "image_id" {
-				continue
-			}
-			if av, err := a.extractAMIVersion(pathComponents[7]); err != nil || av != amiVersion {
-				continue
-			}
-			imageIDs = append(imageIDs, lo.ToPtr(value))
-			requirements[value] = lo.Map(variants, func(v Variant, _ int) scheduling.Requirements { return v.Requirements() })
-		}
+		ids[imageID] = variants
 	}
 	// Failed to discover any AMIs, we should short circuit AMI discovery
-	if len(imageIDs) == 0 {
+	if len(ids) == 0 {
 		return DescribeImageQuery{}, fmt.Errorf(`failed to discover any AMIs for alias "al2@%s"`, amiVersion)
 	}
+
 	return DescribeImageQuery{
 		Filters: []*ec2.Filter{{
 			Name:   lo.ToPtr("image-id"),
-			Values: imageIDs,
+			Values: lo.ToSlicePtr(lo.Keys(ids)),
 		}},
-		KnownRequirements: requirements,
+		KnownRequirements: lo.MapValues(ids, func(variants []Variant, _ string) []scheduling.Requirements {
+			return lo.Map(variants, func(v Variant, _ int) scheduling.Requirements { return v.Requirements() })
+		}),
 	}, nil
 }
 
-func (a AL2) extractAMIVersion(versionStr string) (string, error) {
-	if versionStr == "recommended" {
-		return AMIVersionLatest, nil
-	}
-	rgx := regexp.MustCompile(`^.*(v\d{8})$`)
-	matches := rgx.FindStringSubmatch(versionStr)
-	if len(matches) != 2 {
-		return "", fmt.Errorf("failed to extract AMI version")
-	}
-	return matches[1], nil
-}
-
 // UserData returns the exact same string for equivalent input,
 // even if elements of those inputs are in differing orders,
 // guaranteeing it won't cause spurious hash differences.

pkg/providers/amifamily/al2023.go

@@ -17,12 +17,9 @@ package amifamily
 import (
 	"context"
 	"fmt"
-	"regexp"
-	"strings"
 
 	"github.com/samber/lo"
 	corev1 "k8s.io/api/core/v1"
-	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/karpenter/pkg/cloudprovider"
 	"sigs.k8s.io/karpenter/pkg/scheduling"
 
@@ -39,71 +36,43 @@ type AL2023 struct {
 }
 
 func (a AL2023) DescribeImageQuery(ctx context.Context, ssmProvider ssm.Provider, k8sVersion string, amiVersion string) (DescribeImageQuery, error) {
-	requirements := make(map[string][]scheduling.Requirements)
-	imageIDs := make([]*string, 0, 5)
-	// Example Paths:
-	// - Latest EKS 1.30 arm64 Standard Image: /aws/service/eks/optimized-ami/1.30/amazon-linux-2023/arm64/standard/recommended/image_id
-	// - Specific EKS 1.30 amd64 Nvidia Image: /aws/service/eks/optimized-ami/1.30/amazon-linux-2023/x86_64/nvidia/amazon-eks-node-al2023-x86_64-nvidia-1.30-v20240625/image_id
-	rootPath := fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2023", k8sVersion)
-	results, err := ssmProvider.List(ctx, rootPath)
-	if err != nil {
-		log.FromContext(ctx).WithValues("path", rootPath, "family", "al2023").Error(err, "discovering AMIs from ssm")
-		return DescribeImageQuery{}, fmt.Errorf(`failed to discover any AMIs for alias "al2023@%s"`, amiVersion)
-	}
-
 	ids := map[string]Variant{}
-	for path, value := range results {
-		pathComponents := strings.Split(path, "/")
-		if len(pathComponents) != 11 || pathComponents[10] != "image_id" {
-			continue
-		}
-		if av, err := a.extractAMIVersion(pathComponents[9]); err != nil || av != amiVersion {
-			continue
-		}
-		variant, err := NewVariant(pathComponents[8])
-		if err != nil {
-			continue
-		}
-		ids[value] = variant
-	}
-
-	// EKS doesn't currently vend any accelerated AL2023 AMIs. We should schedule all workloads to
-	// these standard AMIs until accelerated AMIs are available. This approach ensures Karpenter is
-	// forwards compatible with acclerated AMIs once they become available.
-	hasAcceleratedAMIs := lo.ContainsBy(lo.Values(ids), func(v Variant) bool {
-		return v != VariantStandard
-	})
-	for id, variant := range ids {
-		imageIDs = append(imageIDs, lo.ToPtr(id))
-		if hasAcceleratedAMIs {
-			requirements[id] = []scheduling.Requirements{variant.Requirements()}
+	for arch, variants := range map[string][]Variant{
+		"x86_64": []Variant{VariantStandard, VariantNvidia, VariantNeuron},
+		"arm64":  []Variant{VariantStandard},
+	} {
+		for _, variant := range variants {
+			path := a.resolvePath(arch, string(variant), k8sVersion, amiVersion)
+			imageID, err := ssmProvider.Get(ctx, path)
+			if err != nil {
+				continue
+			}
+			ids[imageID] = variant
 		}
 	}
-
 	// Failed to discover any AMIs, we should short circuit AMI discovery
-	if len(imageIDs) == 0 {
+	if len(ids) == 0 {
 		return DescribeImageQuery{}, fmt.Errorf(`failed to discover AMIs for alias "al2023@%s"`, amiVersion)
 	}
 
 	return DescribeImageQuery{
 		Filters: []*ec2.Filter{{
 			Name:   lo.ToPtr("image-id"),
-			Values: imageIDs,
+			Values: lo.ToSlicePtr(lo.Keys(ids)),
 		}},
-		KnownRequirements: requirements,
+		KnownRequirements: lo.MapValues(ids, func(v Variant, _ string) []scheduling.Requirements {
+			return []scheduling.Requirements{v.Requirements()}
+		}),
 	}, nil
 }
 
-func (a AL2023) extractAMIVersion(versionStr string) (string, error) {
-	if versionStr == "recommended" {
-		return AMIVersionLatest, nil
-	}
-	rgx := regexp.MustCompile(`^.*(v\d{8})$`)
-	matches := rgx.FindStringSubmatch(versionStr)
-	if len(matches) != 2 {
-		return "", fmt.Errorf("failed to extract AMI version")
-	}
-	return matches[1], nil
+func (a AL2023) resolvePath(architecture, variant, k8sVersion, amiVersion string) string {
+	name := lo.Ternary(
+		amiVersion == AMIVersionLatest,
+		"recommended",
+		fmt.Sprintf("amazon-eks-node-al2023-%s-%s-%s-%s", architecture, variant, k8sVersion, amiVersion),
+	)
+	return fmt.Sprintf("/aws/service/eks/optimized-ami/%s/amazon-linux-2023/%s/%s/%s/image_id", k8sVersion, architecture, variant, name)
 }
 
 func (a AL2023) UserData(kubeletConfig *v1.KubeletConfiguration, taints []corev1.Taint, labels map[string]string, caBundle *string, _ []*cloudprovider.InstanceType, customUserData *string, instanceStorePolicy *v1.InstanceStorePolicy) bootstrap.Bootstrapper {