feat: Allow instance store policy to be configured on BR family (#7044)

hack/tools/launchtemplate_counter/main.go

@@ -0,0 +1,145 @@
+/*
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"os"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/patrickmn/go-cache"
+	"github.com/samber/lo"
+	corev1 "k8s.io/api/core/v1"
+	karpv1 "sigs.k8s.io/karpenter/pkg/apis/v1"
+	"sigs.k8s.io/karpenter/pkg/cloudprovider"
+	coreoptions "sigs.k8s.io/karpenter/pkg/operator/options"
+	"sigs.k8s.io/karpenter/pkg/scheduling"
+	coretest "sigs.k8s.io/karpenter/pkg/test"
+
+	v1 "github.com/aws/karpenter-provider-aws/pkg/apis/v1"
+	awscache "github.com/aws/karpenter-provider-aws/pkg/cache"
+	"github.com/aws/karpenter-provider-aws/pkg/operator/options"
+	"github.com/aws/karpenter-provider-aws/pkg/providers/amifamily"
+	"github.com/aws/karpenter-provider-aws/pkg/providers/instancetype"
+	"github.com/aws/karpenter-provider-aws/pkg/providers/pricing"
+	"github.com/aws/karpenter-provider-aws/pkg/providers/subnet"
+	"github.com/aws/karpenter-provider-aws/pkg/test"
+)
+
+func main() {
+	lo.Must0(os.Setenv("AWS_SDK_LOAD_CONFIG", "true"))
+
+	ctx := coreoptions.ToContext(context.Background(), coretest.Options())
+	ctx = options.ToContext(ctx, test.Options(test.OptionsFields{
+		ClusterName:     lo.ToPtr("docs-gen"),
+		ClusterEndpoint: lo.ToPtr("https://docs-gen.aws"),
+		IsolatedVPC:     lo.ToPtr(true), // disable pricing lookup
+	}))
+
+	region := "us-west-2"
+	sess := session.Must(session.NewSession(&aws.Config{Region: lo.ToPtr(region)}))
+	ec2api := ec2.New(sess)
+	subnetProvider := subnet.NewDefaultProvider(ec2api, cache.New(awscache.DefaultTTL, awscache.DefaultCleanupInterval), cache.New(awscache.AvailableIPAddressTTL, awscache.DefaultCleanupInterval), cache.New(awscache.AssociatePublicIPAddressTTL, awscache.DefaultCleanupInterval))
+	instanceTypeProvider := instancetype.NewDefaultProvider(
+		cache.New(awscache.InstanceTypesAndZonesTTL, awscache.DefaultCleanupInterval),
+		ec2api,
+		subnetProvider,
+		instancetype.NewDefaultResolver(
+			region,
+			pricing.NewDefaultProvider(
+				ctx,
+				pricing.NewAPI(sess, *sess.Config.Region),
+				ec2api,
+				*sess.Config.Region,
+			),
+			awscache.NewUnavailableOfferings(),
+		),
+	)
+	if err := instanceTypeProvider.UpdateInstanceTypes(ctx); err != nil {
+		log.Fatalf("updating instance types, %s", err)
+	}
+	if err := instanceTypeProvider.UpdateInstanceTypeOfferings(ctx); err != nil {
+		log.Fatalf("updating instance types offerings, %s", err)
+	}
+	// Fake a NodeClass, so we can use it to get InstanceTypes
+	nodeClass := &v1.EC2NodeClass{
+		Spec: v1.EC2NodeClassSpec{
+			AMISelectorTerms: []v1.AMISelectorTerm{{
+				Alias: "al2023@latest",
+			}},
+			SubnetSelectorTerms: []v1.SubnetSelectorTerm{
+				{
+					Tags: map[string]string{
+						"*": "*",
+					},
+				},
+			},
+		},
+	}
+	subnets, err := subnetProvider.List(ctx, nodeClass)
+	if err != nil {
+		log.Fatalf("listing subnets, %s", err)
+	}
+	nodeClass.Status.Subnets = lo.Map(subnets, func(ec2subnet *ec2.Subnet, _ int) v1.Subnet {
+		return v1.Subnet{
+			ID:   *ec2subnet.SubnetId,
+			Zone: *ec2subnet.AvailabilityZone,
+		}
+	})
+	nodeClass.Status.AMIs = []v1.AMI{
+		{
+			ID:   coretest.RandomName(),
+			Name: coretest.RandomName(),
+			Requirements: []corev1.NodeSelectorRequirement{
+				{
+					Key:      corev1.LabelArchStable,
+					Operator: corev1.NodeSelectorOpIn,
+					Values:   []string{karpv1.ArchitectureAmd64},
+				},
+			},
+		},
+		{
+			ID:   coretest.RandomName(),
+			Name: coretest.RandomName(),
+			Requirements: []corev1.NodeSelectorRequirement{
+				{
+					Key:      corev1.LabelArchStable,
+					Operator: corev1.NodeSelectorOpIn,
+					Values:   []string{karpv1.ArchitectureArm64},
+				},
+			},
+		},
+	}
+	instanceTypes, err := instanceTypeProvider.List(ctx, nodeClass)
+
+	// See how many launch templates we get by constraining our instance types to just be "c", "m", and "r"
+	reqs := scheduling.NewRequirements(scheduling.NewRequirement(v1.LabelInstanceCategory, corev1.NodeSelectorOpIn, "c", "m", "r"))
+	instanceTypes = lo.Filter(instanceTypes, func(it *cloudprovider.InstanceType, _ int) bool {
+		return it.Requirements.Compatible(reqs) == nil
+	})
+	fmt.Printf("Got %d instance types after filtering\n", len(instanceTypes))
+
+	resolver := amifamily.NewDefaultResolver()
+	launchTemplates, err := resolver.Resolve(nodeClass, &karpv1.NodeClaim{}, lo.Slice(instanceTypes, 0, 60), karpv1.CapacityTypeOnDemand, &amifamily.Options{InstanceStorePolicy: lo.ToPtr(v1.InstanceStorePolicyRAID0)})
+
+	if err != nil {
+		log.Fatalf("resolving launchTemplates, %s", err)
+	}
+	fmt.Printf("Got %d launch templates back from the resolver\n", len(launchTemplates))
+}

pkg/providers/amifamily/bootstrap/bottlerocket.go

@@ -23,6 +23,8 @@ import (
 	"github.com/samber/lo"
 
 	"github.com/aws/aws-sdk-go/aws"
+
+	v1 "github.com/aws/karpenter-provider-aws/pkg/apis/v1"
 )
 
 type Bottlerocket struct {
@@ -76,6 +78,17 @@ func (b Bottlerocket) Script() (string, error) {
 	for _, taint := range b.Taints {
 		s.Settings.Kubernetes.NodeTaints[taint.Key] = append(s.Settings.Kubernetes.NodeTaints[taint.Key], fmt.Sprintf("%s:%s", taint.Value, taint.Effect))
 	}
+
+	if lo.FromPtr(b.InstanceStorePolicy) == v1.InstanceStorePolicyRAID0 {
+		if s.Settings.BootstrapCommands == nil {
+			s.Settings.BootstrapCommands = map[string]BootstrapCommand{}
+		}
+		s.Settings.BootstrapCommands["000-mount-instance-storage"] = BootstrapCommand{
+			Commands:  [][]string{{"apiclient", "ephemeral-storage", "init"}, {"apiclient", "ephemeral-storage", "bind", "--dirs", "/var/lib/containerd", "/var/lib/kubelet", "/var/log/pods"}},
+			Essential: true,
+			Mode:      BootstrapCommandModeAlways,
+		}
+	}
 	script, err := s.MarshalTOML()
 	if err != nil {
 		return "", fmt.Errorf("constructing toml UserData %w", err)

pkg/providers/amifamily/bootstrap/bottlerocketsettings.go

@@ -38,7 +38,8 @@ type BottlerocketConfig struct {
 // BottlerocketSettings is a subset of all configuration in https://github.com/bottlerocket-os/bottlerocket/blob/d427c40931cba6e6bedc5b75e9c084a6e1818db9/sources/models/src/lib.rs#L260
 // These settings apply across all K8s versions that karpenter supports.
 type BottlerocketSettings struct {
-	Kubernetes BottlerocketKubernetes `toml:"kubernetes"`
+	Kubernetes        BottlerocketKubernetes      `toml:"kubernetes"`
+	BootstrapCommands map[string]BootstrapCommand `toml:"bootstrap-commands,omitempty"`
 }
 
 // BottlerocketKubernetes is k8s specific configuration for bottlerocket api
@@ -94,6 +95,22 @@ type BottlerocketCredentialProvider struct {
 	Environment   map[string]string `toml:"environment,omitempty"`
 }
 
+type BootstrapCommandMode string
+
+const (
+	BootstrapCommandModeAlways BootstrapCommandMode = "always"
+	BootstrapCommandModeOnce   BootstrapCommandMode = "once"
+	BootstrapCommandModeOff    BootstrapCommandMode = "off"
+)
+
+// BootstrapCommand model defined in the Bottlerocket Core Kit in
+// https://github.com/bottlerocket-os/bottlerocket-core-kit/blob/fdf32c291ad18370de3a5fdc4c20a9588bc14177/sources/bootstrap-commands/src/main.rs#L57
+type BootstrapCommand struct {
+	Commands  [][]string           `toml:"commands"`
+	Mode      BootstrapCommandMode `toml:"mode"`
+	Essential bool                 `toml:"essential"`
+}
+
 func (c *BottlerocketConfig) UnmarshalTOML(data []byte) error {
 	// unmarshal known settings
 	s := struct {
@@ -115,5 +132,8 @@ func (c *BottlerocketConfig) MarshalTOML() ([]byte, error) {
 		c.SettingsRaw = map[string]interface{}{}
 	}
 	c.SettingsRaw["kubernetes"] = c.Settings.Kubernetes
+	if c.Settings.BootstrapCommands != nil {
+		c.SettingsRaw["bootstrap-commands"] = c.Settings.BootstrapCommands
+	}
 	return toml.Marshal(c)
 }

pkg/providers/amifamily/bottlerocket.go

@@ -72,16 +72,17 @@ func (b Bottlerocket) DescribeImageQuery(ctx context.Context, ssmProvider ssm.Pr
 }
 
 // UserData returns the default userdata script for the AMI Family
-func (b Bottlerocket) UserData(kubeletConfig *v1.KubeletConfiguration, taints []corev1.Taint, labels map[string]string, caBundle *string, _ []*cloudprovider.InstanceType, customUserData *string, _ *v1.InstanceStorePolicy) bootstrap.Bootstrapper {
+func (b Bottlerocket) UserData(kubeletConfig *v1.KubeletConfiguration, taints []corev1.Taint, labels map[string]string, caBundle *string, _ []*cloudprovider.InstanceType, customUserData *string, instanceStorePolicy *v1.InstanceStorePolicy) bootstrap.Bootstrapper {
 	return bootstrap.Bottlerocket{
 		Options: bootstrap.Options{
-			ClusterName:     b.Options.ClusterName,
-			ClusterEndpoint: b.Options.ClusterEndpoint,
-			KubeletConfig:   kubeletConfig,
-			Taints:          taints,
-			Labels:          labels,
-			CABundle:        caBundle,
-			CustomUserData:  customUserData,
+			ClusterName:         b.Options.ClusterName,
+			ClusterEndpoint:     b.Options.ClusterEndpoint,
+			KubeletConfig:       kubeletConfig,
+			Taints:              taints,
+			Labels:              labels,
+			CABundle:            caBundle,
+			CustomUserData:      customUserData,
+			InstanceStorePolicy: instanceStorePolicy,
 		},
 	}
 }

pkg/providers/instancetype/types.go

@@ -235,7 +235,7 @@ func computeRequirements(info *ec2.InstanceTypeInfo, offerings cloudprovider.Off
 		requirements.Get(v1.LabelInstanceFamily).Insert(instanceTypeParts[0])
 		requirements.Get(v1.LabelInstanceSize).Insert(instanceTypeParts[1])
 	}
-	if info.InstanceStorageInfo != nil && aws.StringValue(info.InstanceStorageInfo.NvmeSupport) != ec2.EphemeralNvmeSupportUnsupported {
+	if info.InstanceStorageInfo != nil && aws.StringValue(info.InstanceStorageInfo.NvmeSupport) != ec2.EphemeralNvmeSupportUnsupported && info.InstanceStorageInfo.TotalSizeInGB != nil {
 		requirements[v1.LabelInstanceLocalNVME].Insert(fmt.Sprint(aws.Int64Value(info.InstanceStorageInfo.TotalSizeInGB)))
 	}
 	// Network bandwidth

pkg/providers/launchtemplate/suite_test.go

@@ -1374,8 +1374,54 @@ var _ = Describe("LaunchTemplate Provider", func() {
 			pod := coretest.UnschedulablePod()
 			ExpectProvisioned(ctx, env.Client, cluster, cloudProvider, prov, pod)
 			ExpectScheduled(ctx, env.Client, pod)
+
+			Expect(awsEnv.EC2API.CalledWithCreateLaunchTemplateInput.Len()).To(BeNumerically("==", 5))
 			ExpectLaunchTemplatesCreatedWithUserDataContaining("--local-disks raid0")
 		})
+		It("should specify RAID0 bootstrap-command when instance-store policy is set on Bottlerocket", func() {
+			nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{Alias: "bottlerocket@latest"}}
+			nodeClass.Spec.InstanceStorePolicy = lo.ToPtr(v1.InstanceStorePolicyRAID0)
+			ExpectApplied(ctx, env.Client, nodePool, nodeClass)
+			pod := coretest.UnschedulablePod()
+			ExpectProvisioned(ctx, env.Client, cluster, cloudProvider, prov, pod)
+			ExpectScheduled(ctx, env.Client, pod)
+
+			Expect(awsEnv.EC2API.CalledWithCreateLaunchTemplateInput.Len()).To(BeNumerically("==", 5))
+			ExpectLaunchTemplatesCreatedWithUserDataContaining(`
+[settings.bootstrap-commands.000-mount-instance-storage]
+commands = [['apiclient', 'ephemeral-storage', 'init'], ['apiclient', 'ephemeral-storage', 'bind', '--dirs', '/var/lib/containerd', '/var/lib/kubelet', '/var/log/pods']]
+mode = 'always'
+essential = true
+`)
+		})
+		It("should merge bootstrap-commands when instance-store policy is set on Bottlerocket", func() {
+			nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{Alias: "bottlerocket@latest"}}
+			nodeClass.Spec.InstanceStorePolicy = lo.ToPtr(v1.InstanceStorePolicyRAID0)
+			nodeClass.Spec.UserData = lo.ToPtr(`
+[settings.bootstrap-commands.111-say-hello]
+commands = [['echo', 'hello']]
+mode = 'always'
+essential = true
+`)
+			ExpectApplied(ctx, env.Client, nodePool, nodeClass)
+			pod := coretest.UnschedulablePod()
+			ExpectProvisioned(ctx, env.Client, cluster, cloudProvider, prov, pod)
+			ExpectScheduled(ctx, env.Client, pod)
+
+			Expect(awsEnv.EC2API.CalledWithCreateLaunchTemplateInput.Len()).To(BeNumerically("==", 5))
+			ExpectLaunchTemplatesCreatedWithUserDataContaining(`
+[settings.bootstrap-commands]
+[settings.bootstrap-commands.000-mount-instance-storage]
+commands = [['apiclient', 'ephemeral-storage', 'init'], ['apiclient', 'ephemeral-storage', 'bind', '--dirs', '/var/lib/containerd', '/var/lib/kubelet', '/var/log/pods']]
+mode = 'always'
+essential = true
+
+[settings.bootstrap-commands.111-say-hello]
+commands = [['echo', 'hello']]
+mode = 'always'
+essential = true
+`)
+		})
 		Context("Bottlerocket", func() {
 			BeforeEach(func() {
 				nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{Alias: "bottlerocket@latest"}}
@@ -2240,7 +2286,7 @@ func ExpectUserDataCreatedWithNodeConfigs(userData string) []admv1alpha1.NodeCon
 	GinkgoHelper()
 	archive, err := mime.NewArchive(userData)
 	Expect(err).To(BeNil())
-	nodeConfigs := lo.FilterMap([]mime.Entry(archive), func(entry mime.Entry, _ int) (admv1alpha1.NodeConfig, bool) {
+	nodeConfigs := lo.FilterMap(archive, func(entry mime.Entry, _ int) (admv1alpha1.NodeConfig, bool) {
 		config := admv1alpha1.NodeConfig{}
 		if entry.ContentType != mime.ContentTypeNodeConfig {
 			return config, false

test/suites/ami/suite_test.go

@@ -159,7 +159,7 @@ var _ = Describe("AMI", func() {
 
 	Context("AMIFamily", func() {
 		DescribeTable(
-			"should providion a node using an alias",
+			"should provision a node using an alias",
 			func(alias string) {
 				pod := coretest.Pod()
 				nodeClass.Spec.AMISelectorTerms = []v1.AMISelectorTerm{{Alias: alias}}