feat(runners): add instance_type_priorities for prioritized allocation

Add optional `instance_type_priorities` variable (map of instance type
to priority number) to control EC2 Fleet override priorities. When not
set, priorities default to the index position in `instance_types`,
preserving the user's ordering. This makes the `prioritized` allocation
strategy work correctly for both spot and on-demand fleets.

Author: piotrj

lambdas/functions/control-plane/src/aws/runners.d.ts

@@ -37,6 +37,7 @@ export interface RunnerInputParameters {
   launchTemplateName: string;
   ec2instanceCriteria: {
     instanceTypes: string[];
+    instanceTypePriorities?: Record<string, number>;
     targetCapacityType: DefaultTargetCapacityType;
     maxSpotPrice?: string;
     instanceAllocationStrategy: SpotAllocationStrategy | FleetOnDemandAllocationStrategy;

lambdas/functions/control-plane/src/aws/runners.test.ts

@@ -396,6 +396,26 @@ describe('create runner', () => {
     });
   });
 
+  it('calls create fleet with custom instance type priorities', async () => {
+    const priorities = { 'm5.large': 10, 'c5.large': 5 };
+    await createRunner(
+      createRunnerConfig({
+        ...defaultRunnerConfig,
+        capacityType: 'on-demand',
+        allocationStrategy: FleetOnDemandAllocationStrategy.PRIORITIZED,
+        instanceTypePriorities: priorities,
+      }),
+    );
+    expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
+      ...expectedCreateFleetRequest({
+        ...defaultExpectedFleetRequestValues,
+        capacityType: 'on-demand',
+        allocationStrategy: FleetOnDemandAllocationStrategy.PRIORITIZED,
+        instanceTypePriorities: priorities,
+      }),
+    });
+  });
+
   it('calls run instances with the on-demand capacity', async () => {
     await createRunner(createRunnerConfig({ ...defaultRunnerConfig, maxSpotPrice: '0.1' }));
     expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
@@ -722,6 +742,7 @@ interface RunnerConfig {
   type: RunnerType;
   capacityType: DefaultTargetCapacityType;
   allocationStrategy: SpotAllocationStrategy | FleetOnDemandAllocationStrategy;
+  instanceTypePriorities?: Record<string, number>;
   maxSpotPrice?: string;
   amiIdSsmParameterName?: string;
   tracingEnabled?: boolean;
@@ -738,6 +759,7 @@ function createRunnerConfig(runnerConfig: RunnerConfig): RunnerInputParameters {
     launchTemplateName: LAUNCH_TEMPLATE,
     ec2instanceCriteria: {
       instanceTypes: ['m5.large', 'c5.large'],
+      instanceTypePriorities: runnerConfig.instanceTypePriorities,
       targetCapacityType: runnerConfig.capacityType,
       maxSpotPrice: runnerConfig.maxSpotPrice,
       instanceAllocationStrategy: runnerConfig.allocationStrategy,
@@ -754,6 +776,7 @@ interface ExpectedFleetRequestValues {
   type: 'Repo' | 'Org';
   capacityType: DefaultTargetCapacityType;
   allocationStrategy: SpotAllocationStrategy | FleetOnDemandAllocationStrategy;
+  instanceTypePriorities?: Record<string, number>;
   maxSpotPrice?: string;
   totalTargetCapacity: number;
   imageId?: string;
@@ -785,18 +808,22 @@ function expectedCreateFleetRequest(expectedValues: ExpectedFleetRequestValues):
           {
             InstanceType: 'm5.large',
             SubnetId: 'subnet-123',
+            Priority: expectedValues.instanceTypePriorities?.['m5.large'] ?? 0,
           },
           {
             InstanceType: 'c5.large',
             SubnetId: 'subnet-123',
+            Priority: expectedValues.instanceTypePriorities?.['c5.large'] ?? 1,
           },
           {
             InstanceType: 'm5.large',
             SubnetId: 'subnet-456',
+            Priority: expectedValues.instanceTypePriorities?.['m5.large'] ?? 0,
           },
           {
             InstanceType: 'c5.large',
             SubnetId: 'subnet-456',
+            Priority: expectedValues.instanceTypePriorities?.['c5.large'] ?? 1,
           },
         ],
       },

lambdas/functions/control-plane/src/aws/runners.ts

@@ -127,14 +127,16 @@ function generateFleetOverrides(
   subnetIds: string[],
   instancesTypes: string[],
   amiId?: string,
+  instanceTypePriorities?: Record<string, number>,
 ): FleetLaunchTemplateOverridesRequest[] {
   const result: FleetLaunchTemplateOverridesRequest[] = [];
   subnetIds.forEach((s) => {
-    instancesTypes.forEach((i) => {
+    instancesTypes.forEach((i, index) => {
       const item: FleetLaunchTemplateOverridesRequest = {
         SubnetId: s,
         InstanceType: i as _InstanceType,
         ImageId: amiId,
+        Priority: instanceTypePriorities?.[i] ?? index,
       };
       result.push(item);
     });
@@ -277,6 +279,7 @@ async function createInstances(
             runnerParameters.subnets,
             runnerParameters.ec2instanceCriteria.instanceTypes,
             amiIdOverride,
+            runnerParameters.ec2instanceCriteria.instanceTypePriorities,
           ),
         },
       ],

lambdas/functions/control-plane/src/pool/pool.ts

@@ -36,6 +36,9 @@ export async function adjust(event: PoolEvent): Promise<void> {
   const launchTemplateName = process.env.LAUNCH_TEMPLATE_NAME;
   const instanceMaxSpotPrice = process.env.INSTANCE_MAX_SPOT_PRICE;
   const instanceAllocationStrategy = process.env.INSTANCE_ALLOCATION_STRATEGY || 'lowest-price'; // same as AWS default
+  const instanceTypePriorities = process.env.INSTANCE_TYPE_PRIORITIES
+    ? (JSON.parse(process.env.INSTANCE_TYPE_PRIORITIES) as Record<string, number>)
+    : undefined;
   const runnerOwner = process.env.RUNNER_OWNER;
   const amiIdSsmParameterName = process.env.AMI_ID_SSM_PARAMETER_NAME;
   const tracingEnabled = yn(process.env.POWERTOOLS_TRACE_ENABLED, { default: false });
@@ -92,6 +95,7 @@ export async function adjust(event: PoolEvent): Promise<void> {
       {
         ec2instanceCriteria: {
           instanceTypes,
+          instanceTypePriorities,
           targetCapacityType: instanceTargetCapacityType,
           maxSpotPrice: instanceMaxSpotPrice,
           instanceAllocationStrategy: instanceAllocationStrategy,

lambdas/functions/control-plane/src/scale-runners/scale-up.ts

@@ -304,6 +304,9 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
   const launchTemplateName = process.env.LAUNCH_TEMPLATE_NAME;
   const instanceMaxSpotPrice = process.env.INSTANCE_MAX_SPOT_PRICE;
   const instanceAllocationStrategy = process.env.INSTANCE_ALLOCATION_STRATEGY || 'lowest-price'; // same as AWS default
+  const instanceTypePriorities = process.env.INSTANCE_TYPE_PRIORITIES
+    ? (JSON.parse(process.env.INSTANCE_TYPE_PRIORITIES) as Record<string, number>)
+    : undefined;
   const enableJobQueuedCheck = yn(process.env.ENABLE_JOB_QUEUED_CHECK, { default: true });
   const amiIdSsmParameterName = process.env.AMI_ID_SSM_PARAMETER_NAME;
   const runnerNamePrefix = process.env.RUNNER_NAME_PREFIX || '';
@@ -493,6 +496,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
       {
         ec2instanceCriteria: {
           instanceTypes,
+          instanceTypePriorities,
           targetCapacityType: instanceTargetCapacityType,
           maxSpotPrice: instanceMaxSpotPrice,
           instanceAllocationStrategy: instanceAllocationStrategy,

main.tf

@@ -175,6 +175,7 @@ module "runners" {
   instance_types                = var.instance_types
   instance_target_capacity_type = var.instance_target_capacity_type
   instance_allocation_strategy  = var.instance_allocation_strategy
+  instance_type_priorities      = var.instance_type_priorities
   instance_max_spot_price       = var.instance_max_spot_price
   block_device_mappings         = var.block_device_mappings
 

modules/multi-runner/runners.tf

@@ -22,6 +22,7 @@ module "runners" {
   instance_types                = each.value.runner_config.instance_types
   instance_target_capacity_type = each.value.runner_config.instance_target_capacity_type
   instance_allocation_strategy  = each.value.runner_config.instance_allocation_strategy
+  instance_type_priorities      = each.value.runner_config.instance_type_priorities
   instance_max_spot_price       = each.value.runner_config.instance_max_spot_price
   block_device_mappings         = each.value.runner_config.block_device_mappings
 

modules/multi-runner/variables.tf

@@ -95,6 +95,7 @@ variable "multi_runner_config" {
       enable_ssm_on_runners                   = optional(bool, false)
       enable_userdata                         = optional(bool, true)
       instance_allocation_strategy            = optional(string, "lowest-price")
+      instance_type_priorities                = optional(map(number), null)
       instance_max_spot_price                 = optional(string, null)
       instance_target_capacity_type           = optional(string, "spot")
       instance_types                          = list(string)
@@ -215,6 +216,7 @@ variable "multi_runner_config" {
         enable_ssm_on_runners: "Enable to allow access the runner instances for debugging purposes via SSM. Note that this adds additional permissions to the runner instances."
         enable_userdata: "Should the userdata script be enabled for the runner. Set this to false if you are using your own prebuilt AMI."
         instance_allocation_strategy: "The allocation strategy for creating instances. For spot, AWS recommends `capacity-optimized`; for on-demand, use `lowest-price` or `prioritized`. The AWS default is `lowest-price`."
+        instance_type_priorities: "A map of instance type to priority for the `prioritized` allocation strategy. Lower numbers mean higher priority. If not provided, priorities are assigned based on the order of `instance_types`."
         instance_max_spot_price: "Max price price for spot instances per hour. This variable will be passed to the create fleet as max spot price for the fleet."
         instance_target_capacity_type: "Default lifecycle used for runner instances, can be either `spot` or `on-demand`."
         instance_types: "List of instance types for the action runner. Defaults are based on runner_os (al2023 for linux and Windows Server Core for win)."

modules/runners/pool.tf

@@ -12,6 +12,7 @@ module "pool" {
     user_agent                    = var.user_agent
     github_app_parameters         = var.github_app_parameters
     instance_allocation_strategy  = var.instance_allocation_strategy
+    instance_type_priorities      = var.instance_type_priorities
     instance_max_spot_price       = var.instance_max_spot_price
     instance_target_capacity_type = var.instance_target_capacity_type
     instance_types                = var.instance_types

modules/runners/pool/main.tf

@@ -27,6 +27,7 @@ resource "aws_lambda_function" "pool" {
       INSTANCE_ALLOCATION_STRATEGY             = var.config.instance_allocation_strategy
       INSTANCE_MAX_SPOT_PRICE                  = var.config.instance_max_spot_price
       INSTANCE_TARGET_CAPACITY_TYPE            = var.config.instance_target_capacity_type
+      INSTANCE_TYPE_PRIORITIES                 = var.config.instance_type_priorities != null ? jsonencode(var.config.instance_type_priorities) : ""
       INSTANCE_TYPES                           = join(",", var.config.instance_types)
       LAUNCH_TEMPLATE_NAME                     = var.config.runner.launch_template.name
       LOG_LEVEL                                = var.config.lambda.log_level