address comments

Author: joshuakguo

cmd/controller/main.go

@@ -62,7 +62,7 @@ func main() {
 			overlayUndecoratedCloudProvider,
 			clusterState,
 			op.InstanceTypeStore,
-			corecontrollers.WithRegistrationHook(registrationhooks.NewPlacementGroupRegistrationHook(op.GetClient(), op.InstanceProvider, op.PlacementGroupProvider)),
+			corecontrollers.WithRegistrationHook(registrationhooks.NewPlacementGroupRegistrationHook(op.InstanceProvider)),
 		)...).
 		WithControllers(ctx, controllers.NewControllers(
 			ctx,

kwok/operator/operator.go

@@ -149,7 +149,7 @@ func NewOperator(ctx context.Context, operator *operator.Operator) (context.Cont
 	amiProvider := amifamily.NewDefaultProvider(operator.Clock, versionProvider, ssmProvider, ec2api, cache.New(awscache.DefaultTTL, awscache.DefaultCleanupInterval))
 	placementGroupProvider := placementgroup.NewProvider(
 		ec2api,
-		cache.New(awscache.PlacementGroupTTL, awscache.DefaultCleanupInterval),
+		cache.New(awscache.DefaultTTL, awscache.DefaultCleanupInterval),
 	)
 	amiResolver := amifamily.NewDefaultResolver(cfg.Region)
 	launchTemplateProvider := launchtemplate.NewDefaultProvider(

pkg/apis/v1/ec2nodeclass.go

@@ -57,7 +57,7 @@ type EC2NodeClassSpec struct {
 	// +kubebuilder:validation:XValidation:message="expected at least one, got none, ['name', 'id']",rule="has(self.name) || has(self.id)"
 	// +kubebuilder:validation:XValidation:message="'name' and 'id' are mutually exclusive",rule="!(has(self.name) && has(self.id))"
 	// +optional
-	PlacementGroupSelector *PlacementGroupSelectorTerm `json:"placementGroupSelector,omitempty"`
+	PlacementGroupSelector *PlacementGroupSelector `json:"placementGroupSelector,omitempty"`
 	// AssociatePublicIPAddress controls if public IP addresses are assigned to instances that are launched with the nodeclass.
 	// +optional
 	AssociatePublicIPAddress *bool `json:"associatePublicIPAddress,omitempty"`
@@ -211,7 +211,7 @@ type CapacityReservationSelectorTerm struct {
 	InstanceMatchCriteria string `json:"instanceMatchCriteria,omitempty"`
 }
 
-type PlacementGroupSelectorTerm struct {
+type PlacementGroupSelector struct {
 	// Name is the placement group name in EC2
 	// +kubebuilder:validation:MinLength:=1
 	// +optional

pkg/apis/v1/ec2nodeclass_status.go

@@ -173,13 +173,11 @@ func (in *EC2NodeClass) StatusConditions() status.ConditionSet {
 		ConditionTypeSecurityGroupsReady,
 		ConditionTypeInstanceProfileReady,
 		ConditionTypeValidationSucceeded,
+		ConditionTypePlacementGroupReady,
 	}
 	if CapacityReservationsEnabled {
 		conds = append(conds, ConditionTypeCapacityReservationsReady)
 	}
-	if in.Spec.PlacementGroupSelector != nil {
-		conds = append(conds, ConditionTypePlacementGroupReady)
-	}
 	return status.NewReadyConditions(conds...).For(in)
 }
 

pkg/apis/v1/zz_generated.deepcopy.go

@@ -31,9 +31,6 @@ const (
 	// updated every minute, but we want to persist the data longer in the event of an EC2 API outage. 24 hours was the
 	// compormise made for API outage reseliency and gargage collecting entries for orphaned reservations.
 	CapacityReservationAvailabilityTTL = 24 * time.Hour
-	// PlacementGroupTTL is the time we will persist cached placement group data. Like capacity reservations, we persist
-	// longer than the reconciliation interval to maintain availability during EC2 API outages.
-	PlacementGroupTTL = 6 * time.Hour
 	// InstanceTypesZonesAndOfferingsTTL is the time before we refresh instance types, zones, and offerings at EC2
 	InstanceTypesZonesAndOfferingsTTL = 5 * time.Minute
 	// InstanceProfileTTL is the time before we refresh checking instance profile existence at IAM

pkg/cache/cache.go

@@ -22,19 +22,28 @@ import (
 	"sync/atomic"
 
 	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
+	"github.com/awslabs/operatorpkg/option"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
 	"github.com/patrickmn/go-cache"
 )
 
-// PlacementGroupScope optionally scopes ICE cache entries to a specific placement group and partition.
-// When set, ICE entries are isolated so that placement-group-specific failures don't block non-PG launches.
-type PlacementGroupScope struct {
-	// ID is the placement group ID (e.g., "pg-0123456789abcdef0")
-	ID string
-	// Partition is the partition number as a string (e.g., "3"). Empty when no specific partition is targeted
-	// or for non-partition placement groups.
-	Partition string
+// UnavailableOfferingsOption is a functional option for scoping ICE cache entries.
+type UnavailableOfferingsOption = option.Function[unavailableOfferingsOptions]
+
+type unavailableOfferingsOptions struct {
+	placementGroupID        string
+	placementGroupPartition string
+}
+
+// WithPlacementGroup scopes an ICE cache entry to a specific placement group ID.
+func WithPlacementGroup(id string) UnavailableOfferingsOption {
+	return func(o *unavailableOfferingsOptions) { o.placementGroupID = id }
+}
+
+// WithPlacementGroupPartition further scopes an ICE cache entry to a specific partition.
+func WithPlacementGroupPartition(partition string) UnavailableOfferingsOption {
+	return func(o *unavailableOfferingsOptions) { o.placementGroupPartition = partition }
 }
 
 // UnavailableOfferings stores any offerings that return ICE (insufficient capacity errors) when
@@ -93,8 +102,8 @@ func (u *UnavailableOfferings) SeqNum(instanceType ec2types.InstanceType) uint64
 // The pgScope parameter scopes the lookup so that an ICE from a placement group launch
 // does not incorrectly prevent launches of the same instance type + zone without that placement group.
 // When a partition is specified in the scope, the lookup is further scoped to that partition.
-func (u *UnavailableOfferings) IsUnavailable(instanceType ec2types.InstanceType, zone, capacityType string, pgScope ...PlacementGroupScope) bool {
-	_, offeringFound := u.offeringCache.Get(u.key(instanceType, zone, capacityType, pgScope...))
+func (u *UnavailableOfferings) IsUnavailable(instanceType ec2types.InstanceType, zone, capacityType string, opts ...UnavailableOfferingsOption) bool {
+	_, offeringFound := u.offeringCache.Get(u.key(instanceType, zone, capacityType, opts...))
 	_, capacityTypeFound := u.capacityTypeCache.Get(capacityType)
 	_, azFound := u.azCache.Get(zone)
 	return offeringFound || capacityTypeFound || azFound
@@ -104,7 +113,8 @@ func (u *UnavailableOfferings) IsUnavailable(instanceType ec2types.InstanceType,
 // The pgScope parameter scopes the cache entry so that placement-group-specific ICEs don't
 // block non-PG launches of the same instance type + zone. When a partition is specified, the cache
 // entry is further scoped so only that partition is marked unavailable.
-func (u *UnavailableOfferings) MarkUnavailable(ctx context.Context, instanceType ec2types.InstanceType, zone, capacityType string, unavailableReason map[string]string, pgScope ...PlacementGroupScope) {
+func (u *UnavailableOfferings) MarkUnavailable(ctx context.Context, instanceType ec2types.InstanceType, zone, capacityType string, unavailableReason map[string]string, opts ...UnavailableOfferingsOption) {
+	resolved := option.Resolve(opts...)
 	// even if the key is already in the cache, we still need to call Set to extend the cached entry's TTL
 	logValues := []any{
 		"reason", unavailableReason["reason"],
@@ -113,10 +123,10 @@ func (u *UnavailableOfferings) MarkUnavailable(ctx context.Context, instanceType
 		"capacity-type", capacityType,
 		"ttl", UnavailableOfferingsTTL,
 	}
-	if len(pgScope) > 0 && pgScope[0].ID != "" {
-		logValues = append(logValues, "placement-group-id", pgScope[0].ID)
-		if pgScope[0].Partition != "" {
-			logValues = append(logValues, "placement-group-partition", pgScope[0].Partition)
+	if resolved.placementGroupID != "" {
+		logValues = append(logValues, "placement-group-id", resolved.placementGroupID)
+		if resolved.placementGroupPartition != "" {
+			logValues = append(logValues, "placement-group-partition", resolved.placementGroupPartition)
 		}
 	}
 	// Add fleetID if provided
@@ -126,7 +136,7 @@ func (u *UnavailableOfferings) MarkUnavailable(ctx context.Context, instanceType
 		logValues = append(logValues, key, unavailableReason[key])
 	}
 	log.FromContext(ctx).WithValues(logValues...).V(1).Info("removing offering from offerings")
-	u.offeringCache.SetDefault(u.key(instanceType, zone, capacityType, pgScope...), struct{}{})
+	u.offeringCache.SetDefault(u.key(instanceType, zone, capacityType, opts...), struct{}{})
 	u.offeringCacheSeqNumMu.Lock()
 	u.offeringCacheSeqNum[instanceType]++
 	u.offeringCacheSeqNumMu.Unlock()
@@ -156,12 +166,13 @@ func (u *UnavailableOfferings) Flush() {
 // When a placement group scope is provided, the PG ID (and optionally partition) is included in the key
 // to scope ICE entries per placement group and partition.
 // Format: <capacityType>:<instanceType>:<zone>[:<pgID>[:<partition>]]
-func (u *UnavailableOfferings) key(instanceType ec2types.InstanceType, zone string, capacityType string, pgScope ...PlacementGroupScope) string {
-	if len(pgScope) > 0 && pgScope[0].ID != "" {
-		if pgScope[0].Partition != "" {
-			return fmt.Sprintf("%s:%s:%s:%s:%s", capacityType, instanceType, zone, pgScope[0].ID, pgScope[0].Partition)
+func (u *UnavailableOfferings) key(instanceType ec2types.InstanceType, zone string, capacityType string, opts ...UnavailableOfferingsOption) string {
+	resolved := option.Resolve(opts...)
+	if resolved.placementGroupID != "" {
+		if resolved.placementGroupPartition != "" {
+			return fmt.Sprintf("%s:%s:%s:%s:%s", capacityType, instanceType, zone, resolved.placementGroupID, resolved.placementGroupPartition)
 		}
-		return fmt.Sprintf("%s:%s:%s:%s", capacityType, instanceType, zone, pgScope[0].ID)
+		return fmt.Sprintf("%s:%s:%s:%s", capacityType, instanceType, zone, resolved.placementGroupID)
 	}
 	return fmt.Sprintf("%s:%s:%s", capacityType, instanceType, zone)
 }

pkg/cache/unavailableofferings.go

@@ -149,7 +149,7 @@ func (c *CloudProvider) Create(ctx context.Context, nodeClaim *karpv1.NodeClaim)
 	instanceType, _ := lo.Find(instanceTypes, func(i *cloudprovider.InstanceType) bool {
 		return i.Name == string(instance.Type)
 	})
-	nc := c.instanceToNodeClaim(instance, instanceType, nodeClass)
+	nc := c.instanceToNodeClaim(ctx, instance, instanceType, nodeClass)
 	nc.Annotations = lo.Assign(nc.Annotations, map[string]string{
 		v1.AnnotationEC2NodeClassHash:        nodeClass.Hash(),
 		v1.AnnotationEC2NodeClassHashVersion: v1.EC2NodeClassHashVersion,
@@ -173,7 +173,7 @@ func (c *CloudProvider) List(ctx context.Context) ([]*karpv1.NodeClaim, error) {
 		if client.IgnoreNotFound(err) != nil {
 			return nil, fmt.Errorf("resolving nodeclass, %w", err)
 		}
-		nodeClaims = append(nodeClaims, c.instanceToNodeClaim(it, instanceType, nc))
+		nodeClaims = append(nodeClaims, c.instanceToNodeClaim(ctx, it, instanceType, nc))
 	}
 	return nodeClaims, nil
 }
@@ -196,7 +196,7 @@ func (c *CloudProvider) Get(ctx context.Context, providerID string) (*karpv1.Nod
 	if client.IgnoreNotFound(err) != nil {
 		return nil, fmt.Errorf("resolving nodeclass, %w", err)
 	}
-	return c.instanceToNodeClaim(instance, instanceType, nc), nil
+	return c.instanceToNodeClaim(ctx, instance, instanceType, nc), nil
 }
 
 // GetInstanceTypes returns all available InstanceTypes
@@ -411,7 +411,7 @@ func (c *CloudProvider) resolveNodePoolFromInstance(ctx context.Context, instanc
 }
 
 //nolint:gocyclo
-func (c *CloudProvider) instanceToNodeClaim(i *instance.Instance, instanceType *cloudprovider.InstanceType, nodeClass *v1.EC2NodeClass) *karpv1.NodeClaim {
+func (c *CloudProvider) instanceToNodeClaim(ctx context.Context, i *instance.Instance, instanceType *cloudprovider.InstanceType, nodeClass *v1.EC2NodeClass) *karpv1.NodeClaim {
 	nodeClaim := &karpv1.NodeClaim{}
 	labels := map[string]string{}
 	annotations := map[string]string{}
@@ -466,7 +466,7 @@ func (c *CloudProvider) instanceToNodeClaim(i *instance.Instance, instanceType *
 	}
 	// Placement group labels
 	if nodeClass != nil {
-		if pg := c.placementGroupProvider.GetForNodeClass(nodeClass); pg != nil {
+		if pg, _ := c.placementGroupProvider.Get(ctx, nodeClass); pg != nil {
 			labels[v1.LabelPlacementGroupID] = pg.ID
 		}
 	}

pkg/cloudprovider/cloudprovider.go

@@ -66,7 +66,7 @@ func (c *CloudProvider) isNodeClassDrifted(ctx context.Context, nodeClaim *karpv
 		return "", fmt.Errorf("calculating subnet drift, %w", err)
 	}
 	capacityReservationsDrifted := c.isCapacityReservationDrifted(instance, nodeClass)
-	placementGroupDrifted := c.isPlacementGroupDrifted(nodeClaim, nodeClass)
+	placementGroupDrifted := c.isPlacementGroupDrifted(ctx, nodeClaim, nodeClass)
 	drifted := lo.FindOrElse([]cloudprovider.DriftReason{
 		securitygroupDrifted,
 		subnetDrifted,
@@ -151,10 +151,10 @@ func (c *CloudProvider) isCapacityReservationDrifted(instance *instance.Instance
 
 // isPlacementGroupDrifted checks if the node's placement group ID label no longer matches the EC2NodeClass's
 // resolved placement group. This covers scenarios where placementGroupSelector was added, removed, or changed.
-func (c *CloudProvider) isPlacementGroupDrifted(nodeClaim *karpv1.NodeClaim, nodeClass *v1.EC2NodeClass) cloudprovider.DriftReason {
+func (c *CloudProvider) isPlacementGroupDrifted(ctx context.Context, nodeClaim *karpv1.NodeClaim, nodeClass *v1.EC2NodeClass) cloudprovider.DriftReason {
 	nodeClaimPGID := nodeClaim.Labels[v1.LabelPlacementGroupID]
 	var nodeClassPGID string
-	if pg := c.placementGroupProvider.GetForNodeClass(nodeClass); pg != nil {
+	if pg, _ := c.placementGroupProvider.Get(ctx, nodeClass); pg != nil {
 		nodeClassPGID = pg.ID
 	}
 	if nodeClaimPGID != nodeClassPGID {

pkg/cloudprovider/drift.go

@@ -19,15 +19,11 @@ import (
 	"fmt"
 	"strconv"
 
-	"k8s.io/apimachinery/pkg/types"
-	"sigs.k8s.io/controller-runtime/pkg/client"
-
 	karpv1 "sigs.k8s.io/karpenter/pkg/apis/v1"
 	"sigs.k8s.io/karpenter/pkg/cloudprovider"
 
 	v1 "github.com/aws/karpenter-provider-aws/pkg/apis/v1"
 	"github.com/aws/karpenter-provider-aws/pkg/providers/instance"
-	"github.com/aws/karpenter-provider-aws/pkg/providers/placementgroup"
 	"github.com/aws/karpenter-provider-aws/pkg/utils"
 )
 
@@ -38,16 +34,12 @@ import (
 // taint is removed, so that TopologySpreadConstraints using the partition topology key always
 // see accurate partition data.
 type PlacementGroupRegistrationHook struct {
-	kubeClient             client.Client
-	instanceProvider       instance.Provider
-	placementGroupProvider placementgroup.Provider
+	instanceProvider instance.Provider
 }
 
-func NewPlacementGroupRegistrationHook(kubeClient client.Client, instanceProvider instance.Provider, placementGroupProvider placementgroup.Provider) *PlacementGroupRegistrationHook {
+func NewPlacementGroupRegistrationHook(instanceProvider instance.Provider) *PlacementGroupRegistrationHook {
 	return &PlacementGroupRegistrationHook{
-		kubeClient:             kubeClient,
-		instanceProvider:       instanceProvider,
-		placementGroupProvider: placementGroupProvider,
+		instanceProvider: instanceProvider,
 	}
 }
 
@@ -56,46 +48,34 @@ func (h *PlacementGroupRegistrationHook) Name() string {
 }
 
 func (h *PlacementGroupRegistrationHook) Registered(ctx context.Context, nodeClaim *karpv1.NodeClaim) (cloudprovider.NodeLifecycleHookResult, error) {
-	// Resolve the EC2NodeClass from the NodeClaim's nodeClassRef
-	nodeClass := &v1.EC2NodeClass{}
-	if err := h.kubeClient.Get(ctx, types.NamespacedName{Name: nodeClaim.Spec.NodeClassRef.Name}, nodeClass); err != nil {
-		return cloudprovider.NodeLifecycleHookResult{}, fmt.Errorf("resolving ec2nodeclass for placement group hook, %w", err)
-	}
-
-	// Check if the EC2NodeClass has a partition placement group
-	pg := h.placementGroupProvider.GetForNodeClass(nodeClass)
-	if pg == nil || pg.Strategy != placementgroup.StrategyPartition {
+	if _, ok := nodeClaim.Labels[v1.LabelPlacementGroupID]; !ok {
 		return cloudprovider.NodeLifecycleHookResult{}, nil
 	}
 
-	// Check if the partition label is already populated on the NodeClaim
 	if _, ok := nodeClaim.Labels[v1.LabelPlacementGroupPartition]; ok {
 		return cloudprovider.NodeLifecycleHookResult{}, nil
 	}
 
-	// We need the providerID to look up the instance
 	if nodeClaim.Status.ProviderID == "" {
 		return cloudprovider.NodeLifecycleHookResult{Requeue: true}, nil
 	}
 
-	// Parse the instance ID from the provider ID
 	instanceID, err := utils.ParseInstanceID(nodeClaim.Status.ProviderID)
 	if err != nil {
-		return cloudprovider.NodeLifecycleHookResult{Requeue: true}, fmt.Errorf("parsing instance ID from provider ID, %w", err)
+		return cloudprovider.NodeLifecycleHookResult{}, fmt.Errorf("parsing instance ID from provider ID, %w", err)
 	}
 
-	// Get the instance details, skipping cache to get fresh partition data
 	inst, err := h.instanceProvider.Get(ctx, instanceID, instance.SkipCache)
 	if err != nil {
-		return cloudprovider.NodeLifecycleHookResult{Requeue: true}, fmt.Errorf("describing instance for partition number, %w", err)
+		return cloudprovider.NodeLifecycleHookResult{}, fmt.Errorf("describing instance for partition number, %w", err)
 	}
 
-	// Check if the partition number has been assigned
+	// If the instance doesn't have a partition number, it's not in a partition placement group
+	// (e.g., cluster or spread PGs don't have partitions). Nothing to gate.
 	if inst.PartitionNumber == nil {
-		return cloudprovider.NodeLifecycleHookResult{Requeue: true}, nil
+		return cloudprovider.NodeLifecycleHookResult{}, nil
 	}
 
-	// Set the partition label on the NodeClaim so it gets synced to the Node during registration
 	nodeClaim.Labels[v1.LabelPlacementGroupPartition] = strconv.FormatInt(int64(*inst.PartitionNumber), 10)
 	return cloudprovider.NodeLifecycleHookResult{}, nil
 }