Support delayed registration for AWS KWOK

Author: jonathan-innis

kwok/ec2/ec2.go

@@ -38,11 +38,11 @@ import (
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/util/workqueue"
 	"k8s.io/utils/clock"
+	"k8s.io/utils/set"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/karpenter/kwok/apis/v1alpha1"
 	v1 "sigs.k8s.io/karpenter/pkg/apis/v1"
-	"sigs.k8s.io/karpenter/pkg/cloudprovider"
 
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 
@@ -64,7 +64,8 @@ type Client struct {
 	subnets       []ec2types.Subnet
 	strategy      strategy.Strategy
 
-	instances sync.Map
+	instances            sync.Map
+	instanceLaunchCancel sync.Map
 
 	launchTemplates        sync.Map
 	launchTemplateNameToID sync.Map
@@ -111,7 +112,12 @@ func NewClient(region, namespace string, ec2Client *ec2.Client, rateLimiterProvi
 
 func (c *Client) readBackup(ctx context.Context, cfg *rest.Config) {
 	configMaps := &corev1.ConfigMapList{}
-	lo.Must0(client.IgnoreNotFound(lo.Must(client.New(cfg, client.Options{})).List(ctx, configMaps, client.InNamespace(c.namespace))))
+	lo.Must0(lo.Must(client.New(cfg, client.Options{})).List(ctx, configMaps, client.InNamespace(c.namespace)))
+
+	nodeList := &corev1.NodeList{}
+	lo.Must0(lo.Must(client.New(cfg, client.Options{})).List(ctx, nodeList, client.MatchingLabels{v1alpha1.KwokLabelKey: v1alpha1.KwokLabelValue}))
+
+	instanceIDs := set.New[string](lo.Map(nodeList.Items, func(n corev1.Node, _ int) string { return lo.Must(utils.ParseInstanceID(n.Spec.ProviderID)) })...)
 
 	configMaps.Items = lo.Filter(configMaps.Items, func(c corev1.ConfigMap, _ int) bool {
 		return strings.Contains(c.Name, "kwok-aws-instances-")
@@ -123,6 +129,11 @@ func (c *Client) readBackup(ctx context.Context, cfg *rest.Config) {
 			lo.Must0(json.Unmarshal([]byte(cm.Data["instances"]), &instances))
 			for _, instance := range instances {
 				c.instances.Store(lo.FromPtr(instance.InstanceId), instance)
+				// Register nodes immediately if we killed the KWOK controller before actually registering the node
+				if !instanceIDs.Has(lo.FromPtr(instance.InstanceId)) {
+
+					lo.Must0(c.kubeClient.Create(ctx, c.toNode(ctx, instance)))
+				}
 			}
 			total += len(instances)
 		}
@@ -175,7 +186,7 @@ func (c *Client) backupInstances(ctx context.Context) error {
 	numConfigMaps := int(math.Ceil(float64(len(instances)) / float64(500)))
 	if numConfigMaps < len(configMaps.Items) {
 		errs := make([]error, numConfigMaps)
-		workqueue.ParallelizeUntil(ctx, 10, len(configMaps.Items)-numConfigMaps, func(i int) {
+		workqueue.ParallelizeUntil(ctx, len(configMaps.Items)-numConfigMaps, len(configMaps.Items)-numConfigMaps, func(i int) {
 			if err := c.kubeClient.Delete(ctx, &configMaps.Items[len(configMaps.Items)-i-1]); client.IgnoreNotFound(err) != nil {
 				errs[i] = fmt.Errorf("deleting configmap %q, %w", configMaps.Items[len(configMaps.Items)-i-1].Name, err)
 			}
@@ -186,7 +197,7 @@ func (c *Client) backupInstances(ctx context.Context) error {
 	}
 
 	errs := make([]error, numConfigMaps)
-	workqueue.ParallelizeUntil(ctx, 10, numConfigMaps, func(i int) {
+	workqueue.ParallelizeUntil(ctx, numConfigMaps, numConfigMaps, func(i int) {
 		cm := &corev1.ConfigMap{
 			ObjectMeta: metav1.ObjectMeta{
 				Name:      fmt.Sprintf("kwok-aws-instances-%d", i),
@@ -586,15 +597,21 @@ func (c *Client) CreateFleet(ctx context.Context, input *ec2.CreateFleetInput, _
 			VpcId:                    subnet.VpcId,
 		}
 		c.instances.Store(lo.FromPtr(instance.InstanceId), instance)
+		launchCtx, cancel := context.WithCancel(ctx)
+		c.instanceLaunchCancel.Store(lo.FromPtr(instance.InstanceId), cancel)
 
-		// Create the Node through the instance launch
-		// TODO: Eventually support delayed registration
-		nodePoolNameTag, _ := lo.Find(instance.Tags, func(t ec2types.Tag) bool {
-			return lo.FromPtr(t.Key) == v1.NodePoolLabelKey
-		})
-		if err := c.kubeClient.Create(ctx, toNode(lo.FromPtr(instance.InstanceId), lo.FromPtr(nodePoolNameTag.Value), it, lo.FromPtr(subnet.AvailabilityZone), v1.CapacityTypeOnDemand)); err != nil {
-			return nil, fmt.Errorf("creating node, %w", err)
-		}
+		go func() {
+			select {
+			case <-launchCtx.Done():
+				return
+			// This is meant to simulate instance startup time
+			case <-c.clock.After(30 * time.Second):
+			}
+			if err := c.kubeClient.Create(launchCtx, c.toNode(ctx, instance)); err != nil {
+				c.instances.Delete(lo.FromPtr(instance.InstanceId))
+				c.instanceLaunchCancel.Delete(lo.FromPtr(instance.InstanceId))
+			}
+		}()
 		fleetInstances = append(fleetInstances, ec2types.CreateFleetInstance{
 			InstanceIds:  []string{lo.FromPtr(instance.InstanceId)},
 			InstanceType: instance.InstanceType,
@@ -644,6 +661,9 @@ func (c *Client) TerminateInstances(_ context.Context, input *ec2.TerminateInsta
 
 	for _, id := range input.InstanceIds {
 		c.instances.Delete(id)
+		if cancel, ok := c.instanceLaunchCancel.LoadAndDelete(id); ok {
+			cancel.(context.CancelFunc)()
+		}
 	}
 	return &ec2.TerminateInstancesOutput{
 		TerminatingInstances: lo.Map(input.InstanceIds, func(id string, _ int) ec2types.InstanceStateChange {
@@ -862,7 +882,35 @@ func (c *Client) DeleteLaunchTemplate(_ context.Context, input *ec2.DeleteLaunch
 	}, nil
 }
 
-func toNode(instanceID, nodePoolName string, instanceType *cloudprovider.InstanceType, zone, capacityType string) *corev1.Node {
+func (c *Client) toNode(ctx context.Context, instance ec2types.Instance) *corev1.Node {
+	nodePoolNameTag, _ := lo.Find(instance.Tags, func(t ec2types.Tag) bool {
+		return lo.FromPtr(t.Key) == v1.NodePoolLabelKey
+	})
+	subnet := lo.Must(lo.Find(c.subnets, func(s ec2types.Subnet) bool {
+		return lo.FromPtr(s.SubnetId) == lo.FromPtr(instance.SubnetId)
+	}))
+	instanceTypeInfo := lo.Must(lo.Find(c.instanceTypes, func(i ec2types.InstanceTypeInfo) bool {
+		return i.InstanceType == instance.InstanceType
+	}))
+	// TODO: We need to get the capacity and allocatable information from the userData
+	it := instancetype.NewInstanceType(
+		ctx,
+		instanceTypeInfo,
+		c.region,
+		nil,
+		nil,
+		nil,
+		nil,
+		nil,
+		nil,
+		nil,
+		nil,
+		nil,
+		nil,
+		// TODO: Eventually support different AMIFamilies from userData
+		"al2023",
+		nil,
+	)
 	nodeName := fmt.Sprintf("%s-%d", strings.ReplaceAll(namesgenerator.GetRandomName(0), "_", "-"), rand.Uint32()) //nolint:gosec
 	return &corev1.Node{
 		ObjectMeta: metav1.ObjectMeta{
@@ -872,25 +920,25 @@ func toNode(instanceID, nodePoolName string, instanceType *cloudprovider.Instanc
 			},
 			// TODO: We can eventually add all the labels from the userData but for now we just add the NodePool labels
 			Labels: map[string]string{
-				corev1.LabelInstanceTypeStable: instanceType.Name,
+				corev1.LabelInstanceTypeStable: it.Name,
 				corev1.LabelHostname:           nodeName,
-				corev1.LabelTopologyRegion:     instanceType.Requirements.Get(corev1.LabelTopologyRegion).Any(),
-				corev1.LabelTopologyZone:       zone,
-				v1.CapacityTypeLabelKey:        capacityType,
-				corev1.LabelArchStable:         instanceType.Requirements.Get(corev1.LabelArchStable).Any(),
+				corev1.LabelTopologyRegion:     it.Requirements.Get(corev1.LabelTopologyRegion).Any(),
+				corev1.LabelTopologyZone:       lo.FromPtr(subnet.AvailabilityZone),
+				v1.CapacityTypeLabelKey:        v1.CapacityTypeOnDemand,
+				corev1.LabelArchStable:         it.Requirements.Get(corev1.LabelArchStable).Any(),
 				corev1.LabelOSStable:           string(corev1.Linux),
-				v1.NodePoolLabelKey:            nodePoolName,
+				v1.NodePoolLabelKey:            lo.FromPtr(nodePoolNameTag.Value),
 				v1alpha1.KwokLabelKey:          v1alpha1.KwokLabelValue,
 				v1alpha1.KwokPartitionLabelKey: "a",
 			},
 		},
 		Spec: corev1.NodeSpec{
-			ProviderID: fmt.Sprintf("kwok-aws:///%s/%s", zone, instanceID),
+			ProviderID: fmt.Sprintf("kwok-aws:///%s/%s", lo.FromPtr(subnet.AvailabilityZone), lo.FromPtr(instance.InstanceId)),
 			Taints:     []corev1.Taint{v1.UnregisteredNoExecuteTaint},
 		},
 		Status: corev1.NodeStatus{
-			Capacity:    instanceType.Capacity,
-			Allocatable: instanceType.Allocatable(),
+			Capacity:    it.Capacity,
+			Allocatable: it.Allocatable(),
 			Phase:       corev1.NodePending,
 		},
 	}