Merge branch 'main' into memory-overhead-controller

aws · Sep 29, 2024 · c63c9e6 · c63c9e6
2 parents 1667d85 + 8676ebf
commit c63c9e6
Showing 1 changed file with 36 additions and 15 deletions.
diff --git a/pkg/providers/instancetype/instancetype.go b/pkg/providers/instancetype/instancetype.go
@@ -24,6 +24,7 @@ import (
 	"sync"
 	"sync/atomic"
 
+	"github.com/awslabs/operatorpkg/option"
 	"github.com/mitchellh/hashstructure/v2"
 	"github.com/patrickmn/go-cache"
 	"github.com/prometheus/client_golang/prometheus"
@@ -58,7 +59,19 @@ type Provider interface {
 	UpdateInstanceTypeMemoryOverhead(ctx context.Context, kubeClient client.Client) error
 }
 
+type defaultProviderOptions struct {
+	instanceTypeFilterMap func(*cloudprovider.InstanceType, *v1.EC2NodeClass) (*cloudprovider.InstanceType, bool)
+}
+
+var WithInstanceTypeFilterMap = func(instanceTypeFilterMap func(*cloudprovider.InstanceType, *v1.EC2NodeClass) (*cloudprovider.InstanceType, bool)) func(*defaultProviderOptions) {
+	return func(opts *defaultProviderOptions) {
+		opts.instanceTypeFilterMap = instanceTypeFilterMap
+	}
+}
+
 type DefaultProvider struct {
+	defaultProviderOptions
+
 	region          string
 	ec2api          ec2iface.EC2API
 	subnetProvider  subnet.Provider
@@ -87,19 +100,27 @@ type DefaultProvider struct {
 }
 
 func NewDefaultProvider(region string, instanceTypesCache *cache.Cache, vmMemoryOverheadCache *cache.Cache, ec2api ec2iface.EC2API, subnetProvider subnet.Provider,
-	unavailableOfferingsCache *awscache.UnavailableOfferings, pricingProvider pricing.Provider) *DefaultProvider {
+	unavailableOfferingsCache *awscache.UnavailableOfferings, pricingProvider pricing.Provider, opts ...option.Function[defaultProviderOptions]) *DefaultProvider {
+	resolvedOpts := *option.Resolve(opts...)
+	if resolvedOpts.instanceTypeFilterMap == nil {
+		resolvedOpts.instanceTypeFilterMap = func(it *cloudprovider.InstanceType, _ *v1.EC2NodeClass) (*cloudprovider.InstanceType, bool) {
+			return it, true
+		}
+	}
 	return &DefaultProvider{
-		ec2api:                ec2api,
-		region:                region,
-		subnetProvider:        subnetProvider,
-		pricingProvider:       pricingProvider,
-		instanceTypesInfo:     []*ec2.InstanceTypeInfo{},
-		instanceTypeOfferings: map[string]sets.Set[string]{},
-		instanceTypesCache:    instanceTypesCache,
-		vmMemoryOverheadCache: vmMemoryOverheadCache,
-		unavailableOfferings:  unavailableOfferingsCache,
-		cm:                    pretty.NewChangeMonitor(),
-		instanceTypesSeqNum:   0,
+
+		defaultProviderOptions: resolvedOpts,
+		ec2api:                 ec2api,
+		region:                 region,
+		subnetProvider:         subnetProvider,
+		pricingProvider:        pricingProvider,
+		instanceTypesInfo:      []*ec2.InstanceTypeInfo{},
+		instanceTypeOfferings:  map[string]sets.Set[string]{},
+		instanceTypesCache:     instanceTypesCache,
+   	vmMemoryOverheadCache: vmMemoryOverheadCache,
+		unavailableOfferings:   unavailableOfferingsCache,
+		cm:                     pretty.NewChangeMonitor(),
+		instanceTypesSeqNum:    0,
 	}
 }
 
@@ -159,7 +180,7 @@ func (p *DefaultProvider) List(ctx context.Context, kc *v1.KubeletConfiguration,
 		log.FromContext(ctx).WithValues("zones", allZones.UnsortedList()).V(1).Info("discovered zones")
 	}
 	amiFamily := amifamily.GetAMIFamily(nodeClass.AMIFamily(), &amifamily.Options{})
-	result := lo.Map(p.instanceTypesInfo, func(i *ec2.InstanceTypeInfo, _ int) *cloudprovider.InstanceType {
+	result := lo.FilterMap(p.instanceTypesInfo, func(i *ec2.InstanceTypeInfo, _ int) (*cloudprovider.InstanceType, bool) {
 		instanceTypeVCPU.With(prometheus.Labels{
 			instanceTypeLabel: *i.InstanceType,
 		}).Set(float64(aws.Int64Value(i.VCpuInfo.DefaultVCpus)))
@@ -176,11 +197,11 @@ func (p *DefaultProvider) List(ctx context.Context, kc *v1.KubeletConfiguration,
 		// Any changes to the values passed into the NewInstanceType method will require making updates to the cache key
 		// so that Karpenter is able to cache the set of InstanceTypes based on values that alter the set of instance types
 		// !!! Important !!!
-		return NewInstanceType(ctx, i, p.region,
+		return p.instanceTypeFilterMap(NewInstanceType(ctx, i, p.region,
 			vmMemoryOverhead, nodeClass.Spec.BlockDeviceMappings, nodeClass.Spec.InstanceStorePolicy,
 			kc.MaxPods, kc.PodsPerCore, kc.KubeReserved, kc.SystemReserved, kc.EvictionHard, kc.EvictionSoft,
 			amiFamily, p.createOfferings(ctx, i, allZones, p.instanceTypeOfferings[aws.StringValue(i.InstanceType)], nodeClass.Status.Subnets),
-		)
+		), nodeClass)
 	})
 	p.instanceTypesCache.SetDefault(key, result)
 	return result, nil