feat: influence AZ selection with env hint (#699)

Author: ndbaker1

internal/deployers/eksapi/common.go

@@ -0,0 +1,28 @@
+package eksapi
+
+import (
+	"os"
+	"slices"
+	"strings"
+)
+
+const AvailabilityZonePriorityEnv = "EKSAPI_AZ_PRIORITY"
+
+func availabilityZoneHintedOrder(availabilityZones []string) []string {
+	var priorityAZs []string
+	if priorityAZsString, ok := os.LookupEnv(AvailabilityZonePriorityEnv); ok {
+		priorityAZs = strings.Split(priorityAZsString, ",")
+	}
+	if len(priorityAZs) == 0 {
+		return availabilityZones
+	}
+	return slices.SortedStableFunc(slices.Values(availabilityZones), func(az1, az2 string) int {
+		if slices.Contains(priorityAZs, az1) {
+			if slices.Contains(priorityAZs, az2) {
+				return 0
+			}
+			return -1
+		}
+		return 0
+	})
+}

internal/deployers/eksapi/common_test.go

@@ -0,0 +1,15 @@
+package eksapi
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_AZ_PRIORITY(t *testing.T) {
+	t.Setenv(AvailabilityZonePriorityEnv, "us-west-2d")
+	assert.Equal(t,
+		[]string{"us-west-2d", "us-west-2b", "us-west-2c"},
+		availabilityZoneHintedOrder([]string{"us-west-2b", "us-west-2c", "us-west-2d"}),
+	)
+}

internal/deployers/eksapi/infra.go

@@ -416,6 +416,8 @@ func (m *InfrastructureManager) normalizeAZs(opts *deployerOptions, subnetAZs []
 		}
 	}
 
+	// enforce users' preferred ordering over AZs
+	filteredAZs = availabilityZoneHintedOrder(filteredAZs)
 	// truncate the list if we went over the max
 	filteredAZs = filteredAZs[:min(len(filteredAZs), expectedCount)]
 

internal/deployers/eksapi/node.go

@@ -699,8 +699,12 @@ func (m *nodeManager) getValidAvailabilityZonesFilter(opts *deployerOptions, inf
 	if describeResponse == nil || len(describeResponse.InstanceTypeOfferings) == 0 {
 		return nil, fmt.Errorf("no instance type offerings in current region with filters %v", describeFilters)
 	}
+	var candidateAZs []string
+	for _, offering := range describeResponse.InstanceTypeOfferings {
+		candidateAZs = append(candidateAZs, aws.ToString(offering.Location))
+	}
 	// EFA traffic cannot cross an AZ https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa.html#efa-limits
-	targetAZ := aws.ToString(describeResponse.InstanceTypeOfferings[0].Location)
+	targetAZ := availabilityZoneHintedOrder(candidateAZs)[0]
 	klog.Infof("Found availability zone %q with offering for instance types %v", targetAZ, opts.InstanceTypes)
 	return []string{targetAZ}, nil
 }
@@ -718,8 +722,7 @@ func (m *nodeManager) getValidSubnets(opts *deployerOptions, infra *Infrastructu
 		describeFilters = append(describeFilters, ec2types.Filter{
 			Name:   aws.String("availability-zone"),
 			Values: availabilityZoneFilter,
-		},
-		)
+		})
 	}
 	describeResponse, err := m.clients.EC2().DescribeSubnets(context.TODO(), &ec2.DescribeSubnetsInput{
 		Filters:   describeFilters,