Skip to content

Commit 6973065

Browse files
committed
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.
1 parent e4e0250 commit 6973065

File tree

14 files changed

+60
-1
lines changed

14 files changed

+60
-1
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ export interface RunnerInputParameters {
4242
launchTemplateName: string;
4343
ec2instanceCriteria: {
4444
instanceTypes: string[];
45+
instanceTypePriorities?: Record<string, number>;
4546
targetCapacityType: DefaultTargetCapacityType;
4647
maxSpotPrice?: string;
4748
instanceAllocationStrategy: SpotAllocationStrategy | FleetOnDemandAllocationStrategy;

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -419,6 +419,26 @@ describe('create runner', () => {
419419
});
420420
});
421421

422+
it('calls create fleet with custom instance type priorities', async () => {
423+
const priorities = { 'm5.large': 10, 'c5.large': 5 };
424+
await createRunner(
425+
createRunnerConfig({
426+
...defaultRunnerConfig,
427+
capacityType: 'on-demand',
428+
allocationStrategy: FleetOnDemandAllocationStrategy.PRIORITIZED,
429+
instanceTypePriorities: priorities,
430+
}),
431+
);
432+
expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
433+
...expectedCreateFleetRequest({
434+
...defaultExpectedFleetRequestValues,
435+
capacityType: 'on-demand',
436+
allocationStrategy: FleetOnDemandAllocationStrategy.PRIORITIZED,
437+
instanceTypePriorities: priorities,
438+
}),
439+
});
440+
});
441+
422442
it('calls run instances with the on-demand capacity', async () => {
423443
await createRunner(createRunnerConfig({ ...defaultRunnerConfig, maxSpotPrice: '0.1' }));
424444
expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
@@ -772,6 +792,7 @@ interface RunnerConfig {
772792
type: RunnerType;
773793
capacityType: DefaultTargetCapacityType;
774794
allocationStrategy: SpotAllocationStrategy | FleetOnDemandAllocationStrategy;
795+
instanceTypePriorities?: Record<string, number>;
775796
maxSpotPrice?: string;
776797
amiIdSsmParameterName?: string;
777798
tracingEnabled?: boolean;
@@ -789,6 +810,7 @@ function createRunnerConfig(runnerConfig: RunnerConfig): RunnerInputParameters {
789810
launchTemplateName: LAUNCH_TEMPLATE,
790811
ec2instanceCriteria: {
791812
instanceTypes: ['m5.large', 'c5.large'],
813+
instanceTypePriorities: runnerConfig.instanceTypePriorities,
792814
targetCapacityType: runnerConfig.capacityType,
793815
maxSpotPrice: runnerConfig.maxSpotPrice,
794816
instanceAllocationStrategy: runnerConfig.allocationStrategy,
@@ -806,6 +828,7 @@ interface ExpectedFleetRequestValues {
806828
type: 'Repo' | 'Org';
807829
capacityType: DefaultTargetCapacityType;
808830
allocationStrategy: SpotAllocationStrategy | FleetOnDemandAllocationStrategy;
831+
instanceTypePriorities?: Record<string, number>;
809832
maxSpotPrice?: string;
810833
totalTargetCapacity: number;
811834
imageId?: string;
@@ -838,18 +861,22 @@ function expectedCreateFleetRequest(expectedValues: ExpectedFleetRequestValues):
838861
{
839862
InstanceType: 'm5.large',
840863
SubnetId: 'subnet-123',
864+
Priority: expectedValues.instanceTypePriorities?.['m5.large'] ?? 0,
841865
},
842866
{
843867
InstanceType: 'c5.large',
844868
SubnetId: 'subnet-123',
869+
Priority: expectedValues.instanceTypePriorities?.['c5.large'] ?? 1,
845870
},
846871
{
847872
InstanceType: 'm5.large',
848873
SubnetId: 'subnet-456',
874+
Priority: expectedValues.instanceTypePriorities?.['m5.large'] ?? 0,
849875
},
850876
{
851877
InstanceType: 'c5.large',
852878
SubnetId: 'subnet-456',
879+
Priority: expectedValues.instanceTypePriorities?.['c5.large'] ?? 1,
853880
},
854881
],
855882
},

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,14 +127,16 @@ function generateFleetOverrides(
127127
subnetIds: string[],
128128
instancesTypes: string[],
129129
amiId?: string,
130+
instanceTypePriorities?: Record<string, number>,
130131
): FleetLaunchTemplateOverridesRequest[] {
131132
const result: FleetLaunchTemplateOverridesRequest[] = [];
132133
subnetIds.forEach((s) => {
133-
instancesTypes.forEach((i) => {
134+
instancesTypes.forEach((i, index) => {
134135
const item: FleetLaunchTemplateOverridesRequest = {
135136
SubnetId: s,
136137
InstanceType: i as _InstanceType,
137138
ImageId: amiId,
139+
Priority: instanceTypePriorities?.[i] ?? index,
138140
};
139141
result.push(item);
140142
});
@@ -277,6 +279,7 @@ async function createInstances(
277279
runnerParameters.subnets,
278280
runnerParameters.ec2instanceCriteria.instanceTypes,
279281
amiIdOverride,
282+
runnerParameters.ec2instanceCriteria.instanceTypePriorities,
280283
),
281284
},
282285
],

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,9 @@ export async function adjust(event: PoolEvent): Promise<void> {
3636
const launchTemplateName = process.env.LAUNCH_TEMPLATE_NAME;
3737
const instanceMaxSpotPrice = process.env.INSTANCE_MAX_SPOT_PRICE;
3838
const instanceAllocationStrategy = process.env.INSTANCE_ALLOCATION_STRATEGY || 'lowest-price'; // same as AWS default
39+
const instanceTypePriorities = process.env.INSTANCE_TYPE_PRIORITIES
40+
? (JSON.parse(process.env.INSTANCE_TYPE_PRIORITIES) as Record<string, number>)
41+
: undefined;
3942
const runnerOwner = process.env.RUNNER_OWNER;
4043
const amiIdSsmParameterName = process.env.AMI_ID_SSM_PARAMETER_NAME;
4144
const tracingEnabled = yn(process.env.POWERTOOLS_TRACE_ENABLED, { default: false });
@@ -92,6 +95,7 @@ export async function adjust(event: PoolEvent): Promise<void> {
9295
{
9396
ec2instanceCriteria: {
9497
instanceTypes,
98+
instanceTypePriorities,
9599
targetCapacityType: instanceTargetCapacityType,
96100
maxSpotPrice: instanceMaxSpotPrice,
97101
instanceAllocationStrategy: instanceAllocationStrategy,

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,6 +308,9 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
308308
const launchTemplateName = process.env.LAUNCH_TEMPLATE_NAME;
309309
const instanceMaxSpotPrice = process.env.INSTANCE_MAX_SPOT_PRICE;
310310
const instanceAllocationStrategy = process.env.INSTANCE_ALLOCATION_STRATEGY || 'lowest-price'; // same as AWS default
311+
const instanceTypePriorities = process.env.INSTANCE_TYPE_PRIORITIES
312+
? (JSON.parse(process.env.INSTANCE_TYPE_PRIORITIES) as Record<string, number>)
313+
: undefined;
311314
const enableJobQueuedCheck = yn(process.env.ENABLE_JOB_QUEUED_CHECK, { default: true });
312315
const amiIdSsmParameterName = process.env.AMI_ID_SSM_PARAMETER_NAME;
313316
const runnerNamePrefix = process.env.RUNNER_NAME_PREFIX || '';
@@ -497,6 +500,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
497500
{
498501
ec2instanceCriteria: {
499502
instanceTypes,
503+
instanceTypePriorities,
500504
targetCapacityType: instanceTargetCapacityType,
501505
maxSpotPrice: instanceMaxSpotPrice,
502506
instanceAllocationStrategy: instanceAllocationStrategy,

main.tf

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,7 @@ module "runners" {
175175
instance_types = var.instance_types
176176
instance_target_capacity_type = var.instance_target_capacity_type
177177
instance_allocation_strategy = var.instance_allocation_strategy
178+
instance_type_priorities = var.instance_type_priorities
178179
instance_max_spot_price = var.instance_max_spot_price
179180
block_device_mappings = var.block_device_mappings
180181

modules/multi-runner/runners.tf

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ module "runners" {
2222
instance_types = each.value.runner_config.instance_types
2323
instance_target_capacity_type = each.value.runner_config.instance_target_capacity_type
2424
instance_allocation_strategy = each.value.runner_config.instance_allocation_strategy
25+
instance_type_priorities = each.value.runner_config.instance_type_priorities
2526
instance_max_spot_price = each.value.runner_config.instance_max_spot_price
2627
block_device_mappings = each.value.runner_config.block_device_mappings
2728

modules/multi-runner/variables.tf

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ variable "multi_runner_config" {
9595
enable_ssm_on_runners = optional(bool, false)
9696
enable_userdata = optional(bool, true)
9797
instance_allocation_strategy = optional(string, "lowest-price")
98+
instance_type_priorities = optional(map(number), null)
9899
instance_max_spot_price = optional(string, null)
99100
instance_target_capacity_type = optional(string, "spot")
100101
instance_types = list(string)
@@ -215,6 +216,7 @@ variable "multi_runner_config" {
215216
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."
216217
enable_userdata: "Should the userdata script be enabled for the runner. Set this to false if you are using your own prebuilt AMI."
217218
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`."
219+
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`."
218220
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."
219221
instance_target_capacity_type: "Default lifecycle used for runner instances, can be either `spot` or `on-demand`."
220222
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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ module "pool" {
1212
user_agent = var.user_agent
1313
github_app_parameters = var.github_app_parameters
1414
instance_allocation_strategy = var.instance_allocation_strategy
15+
instance_type_priorities = var.instance_type_priorities
1516
instance_max_spot_price = var.instance_max_spot_price
1617
instance_target_capacity_type = var.instance_target_capacity_type
1718
instance_types = var.instance_types

modules/runners/pool/main.tf

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ resource "aws_lambda_function" "pool" {
2727
INSTANCE_ALLOCATION_STRATEGY = var.config.instance_allocation_strategy
2828
INSTANCE_MAX_SPOT_PRICE = var.config.instance_max_spot_price
2929
INSTANCE_TARGET_CAPACITY_TYPE = var.config.instance_target_capacity_type
30+
INSTANCE_TYPE_PRIORITIES = var.config.instance_type_priorities != null ? jsonencode(var.config.instance_type_priorities) : ""
3031
INSTANCE_TYPES = join(",", var.config.instance_types)
3132
LAUNCH_TEMPLATE_NAME = var.config.runner.launch_template.name
3233
LOG_LEVEL = var.config.lambda.log_level

0 commit comments

Comments
 (0)