fix: support runner subnets in the same availability zone

README.md

@@ -228,7 +228,7 @@ Join our discord community via [this invite link](https://discord.gg/bxgXW8jJGh)
 | <a name="input_scale_up_reserved_concurrent_executions"></a> [scale\_up\_reserved\_concurrent\_executions](#input\_scale\_up\_reserved\_concurrent\_executions) | Amount of reserved concurrent executions for the scale-up lambda function. A value of 0 disables lambda from being triggered and -1 removes any concurrency limitations. | `number` | `1` | no |
 | <a name="input_ssm_paths"></a> [ssm\_paths](#input\_ssm\_paths) | The root path used in SSM to store configuration and secrets. | <pre>object({<br/>    root       = optional(string, "github-action-runners")<br/>    app        = optional(string, "app")<br/>    runners    = optional(string, "runners")<br/>    webhook    = optional(string, "webhook")<br/>    use_prefix = optional(bool, true)<br/>  })</pre> | `{}` | no |
 | <a name="input_state_event_rule_binaries_syncer"></a> [state\_event\_rule\_binaries\_syncer](#input\_state\_event\_rule\_binaries\_syncer) | Option to disable EventBridge Lambda trigger for the binary syncer, useful to stop automatic updates of binary distribution | `string` | `"ENABLED"` | no |
-| <a name="input_subnet_ids"></a> [subnet\_ids](#input\_subnet\_ids) | List of subnets in which the action runner instances will be launched. The subnets need to exist in the configured VPC (`vpc_id`), and must reside in different availability zones (see https://github.com/github-aws-runners/terraform-aws-github-runner/issues/2904) | `list(string)` | n/a | yes |
+| <a name="input_subnet_ids"></a> [subnet\_ids](#input\_subnet\_ids) | List of subnets in which the action runner instances will be launched. The subnets must exist in the configured VPC (`vpc_id`). | `list(string)` | n/a | yes |
 | <a name="input_syncer_lambda_s3_key"></a> [syncer\_lambda\_s3\_key](#input\_syncer\_lambda\_s3\_key) | S3 key for syncer lambda function. Required if using an S3 bucket to specify lambdas. | `string` | `null` | no |
 | <a name="input_syncer_lambda_s3_object_version"></a> [syncer\_lambda\_s3\_object\_version](#input\_syncer\_lambda\_s3\_object\_version) | S3 object version for syncer lambda function. Useful if S3 versioning is enabled on source bucket. | `string` | `null` | no |
 | <a name="input_tags"></a> [tags](#input\_tags) | Map of tags that will be added to created resources. By default resources will be tagged with name and environment. | `map(string)` | `{}` | no |

lambdas/libs/compute-providers/aws/ec2/src/control-plane/runners.test.ts

@@ -9,6 +9,8 @@ import {
   DeleteTagsCommand,
   DescribeInstancesCommand,
   type DescribeInstancesResult,
+  DescribeSubnetsCommand,
+  type DescribeSubnetsResult,
   EC2Client,
   FleetOnDemandAllocationStrategy,
   RunInstancesCommand,
@@ -77,6 +79,18 @@ const mockRunningInstancesJit: DescribeInstancesResult = {
     },
   ],
 };
+const mockDefaultSubnets: DescribeSubnetsResult = {
+  Subnets: [
+    { SubnetId: 'subnet-123', AvailabilityZoneId: 'euw1-az1' },
+    { SubnetId: 'subnet-456', AvailabilityZoneId: 'euw1-az2' },
+  ],
+};
+
+function getSubnetIdsFromFleetRequest(request: CreateFleetCommandInput): string[] {
+  const overrides = request.LaunchTemplateConfigs?.[0].Overrides ?? [];
+  const subnetIds = overrides.flatMap(({ SubnetId }) => (SubnetId ? [SubnetId] : []));
+  return [...new Set(subnetIds)];
+}
 
 describe('list instances', () => {
   beforeEach(() => {
@@ -338,13 +352,18 @@ describe('create runner', () => {
     mockEC2Client.reset();
     mockSSMClient.reset();
 
+    mockEC2Client.on(DescribeSubnetsCommand).resolves(mockDefaultSubnets);
     mockEC2Client.on(CreateFleetCommand).resolves({ Instances: [{ InstanceIds: ['i-1234'] }] });
     mockSSMClient.on(GetParameterCommand).resolves({});
   });
 
   it.each(RUNNER_TYPES)('calls create fleet of 1 instance with the default config for %p', async (type: RunnerType) => {
     await createRunner(createRunnerConfig({ ...defaultRunnerConfig, type: type }));
 
+    expect(mockEC2Client).toHaveReceivedCommandWith(DescribeSubnetsCommand, {
+      SubnetIds: ['subnet-123', 'subnet-456'],
+    });
+    expect(mockEC2Client).toHaveReceivedCommandTimes(CreateFleetCommand, 1);
     expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
       ...expectedCreateFleetRequest({
         ...defaultExpectedFleetRequestValues,
@@ -353,6 +372,71 @@ describe('create runner', () => {
     });
   });
 
+  it('partitions batch capacity across subnet sets that contain at most one subnet per Availability Zone', async () => {
+    mockEC2Client.on(DescribeSubnetsCommand).resolves({
+      Subnets: [
+        { SubnetId: 'subnet-123', AvailabilityZoneId: 'euw1-az1' },
+        { SubnetId: 'subnet-456', AvailabilityZoneId: 'euw1-az1' },
+        { SubnetId: 'subnet-789', AvailabilityZoneId: 'euw1-az2' },
+      ],
+    });
+    mockEC2Client
+      .on(CreateFleetCommand)
+      .resolvesOnce({ Instances: [{ InstanceIds: ['i-1234'] }] })
+      .resolvesOnce({ Instances: [{ InstanceIds: ['i-5678'] }] });
+
+    const result = await createRunner({
+      ...createRunnerConfig(defaultRunnerConfig),
+      numberOfRunners: 2,
+      subnets: ['subnet-123', 'subnet-456', 'subnet-789'],
+    });
+
+    expect(result).toEqual({
+      instances: ['i-1234', 'i-5678'],
+      retryableErrorCount: 0,
+      nonRetryableErrorCount: 0,
+    });
+    expect(mockEC2Client).toHaveReceivedCommandTimes(CreateFleetCommand, 2);
+    const fleetRequests = mockEC2Client.commandCalls(CreateFleetCommand).map(({ args: [command] }) => command.input);
+    const subnetSets = fleetRequests.map(getSubnetIdsFromFleetRequest);
+    expect(subnetSets.map((subnets) => subnets.sort()).sort()).toEqual(
+      [
+        ['subnet-123', 'subnet-789'],
+        ['subnet-456', 'subnet-789'],
+      ].sort(),
+    );
+  });
+
+  it('carries retryable subnet address failures to the next same-AZ subnet set', async () => {
+    mockEC2Client.on(DescribeSubnetsCommand).resolves({
+      Subnets: [
+        { SubnetId: 'subnet-123', AvailabilityZoneId: 'euw1-az1' },
+        { SubnetId: 'subnet-456', AvailabilityZoneId: 'euw1-az1' },
+      ],
+    });
+    mockEC2Client
+      .on(CreateFleetCommand)
+      .resolvesOnce({ Errors: [{ ErrorCode: 'InsufficientFreeAddressesInSubnet' }] })
+      .resolvesOnce({ Instances: [{ InstanceIds: ['i-1234'] }] });
+
+    const result = await createRunner(
+      createRunnerConfig({
+        ...defaultRunnerConfig,
+        scaleErrors: [...defaultRunnerConfig.scaleErrors, 'InsufficientFreeAddressesInSubnet'],
+      }),
+    );
+
+    expect(result).toEqual({
+      instances: ['i-1234'],
+      retryableErrorCount: 0,
+      nonRetryableErrorCount: 0,
+    });
+    expect(mockEC2Client).toHaveReceivedCommandTimes(CreateFleetCommand, 2);
+    const fleetRequests = mockEC2Client.commandCalls(CreateFleetCommand).map(({ args: [command] }) => command.input);
+    const attemptedSubnets = fleetRequests.flatMap(getSubnetIdsFromFleetRequest).sort();
+    expect(attemptedSubnets).toEqual(['subnet-123', 'subnet-456']);
+  });
+
   it('calls create fleet of 2 instances with the correct config for org ', async () => {
     const instances = [{ InstanceIds: ['i-1234', 'i-5678'] }];
 
@@ -545,6 +629,7 @@ describe('create runner', () => {
       },
     });
 
+    expect(mockEC2Client).not.toHaveReceivedCommand(DescribeSubnetsCommand);
     expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
       LaunchTemplateConfigs: [
         {
@@ -767,6 +852,7 @@ describe('create runner with errors', () => {
     mockEC2Client.reset();
     mockSSMClient.reset();
 
+    mockEC2Client.on(DescribeSubnetsCommand).resolves(mockDefaultSubnets);
     mockSSMClient.on(PutParameterCommand).resolves({});
     mockSSMClient.on(GetParameterCommand).resolves({});
     mockEC2Client.on(CreateFleetCommand).resolves({ Instances: [] });
@@ -1013,6 +1099,7 @@ describe('create runner with errors fail over to OnDemand', () => {
     mockEC2Client.reset();
     mockSSMClient.reset();
 
+    mockEC2Client.on(DescribeSubnetsCommand).resolves(mockDefaultSubnets);
     mockSSMClient.on(PutParameterCommand).resolves({});
     mockSSMClient.on(GetParameterCommand).resolves({});
     mockEC2Client.on(CreateFleetCommand).resolves({ Instances: [] });
@@ -1032,7 +1119,7 @@ describe('create runner with errors fail over to OnDemand', () => {
     expect(mockEC2Client).toHaveReceivedCommandTimes(CreateFleetCommand, 2);
 
     // first call with spot failure
-    expect(mockEC2Client).toHaveReceivedNthCommandWith(1, CreateFleetCommand, {
+    expect(mockEC2Client).toHaveReceivedNthSpecificCommandWith(1, CreateFleetCommand, {
       ...expectedCreateFleetRequest({
         ...defaultExpectedFleetRequestValues,
         totalTargetCapacity: 1,
@@ -1041,7 +1128,7 @@ describe('create runner with errors fail over to OnDemand', () => {
     });
 
     // second call with with OnDemand fallback, allocation strategy defaults to lowest-price
-    expect(mockEC2Client).toHaveReceivedNthCommandWith(2, CreateFleetCommand, {
+    expect(mockEC2Client).toHaveReceivedNthSpecificCommandWith(2, CreateFleetCommand, {
       ...expectedCreateFleetRequest({
         ...defaultExpectedFleetRequestValues,
         totalTargetCapacity: 1,
@@ -1079,7 +1166,7 @@ describe('create runner with errors fail over to OnDemand', () => {
     expect(mockEC2Client).toHaveReceivedCommandTimes(CreateFleetCommand, 2);
 
     // first call with spot failure
-    expect(mockEC2Client).toHaveReceivedNthCommandWith(1, CreateFleetCommand, {
+    expect(mockEC2Client).toHaveReceivedNthSpecificCommandWith(1, CreateFleetCommand, {
       ...expectedCreateFleetRequest({
         ...defaultExpectedFleetRequestValues,
         totalTargetCapacity: 2,
@@ -1088,7 +1175,7 @@ describe('create runner with errors fail over to OnDemand', () => {
     });
 
     // second call with with OnDemand failback, capacity is reduced by 1, allocation strategy defaults to lowest-price
-    expect(mockEC2Client).toHaveReceivedNthCommandWith(2, CreateFleetCommand, {
+    expect(mockEC2Client).toHaveReceivedNthSpecificCommandWith(2, CreateFleetCommand, {
       ...expectedCreateFleetRequest({
         ...defaultExpectedFleetRequestValues,
         totalTargetCapacity: 1,
@@ -1113,7 +1200,7 @@ describe('create runner with errors fail over to OnDemand', () => {
     expect(mockEC2Client).toHaveReceivedCommandTimes(CreateFleetCommand, 1);
 
     // first call with spot failure
-    expect(mockEC2Client).toHaveReceivedNthCommandWith(1, CreateFleetCommand, {
+    expect(mockEC2Client).toHaveReceivedNthSpecificCommandWith(1, CreateFleetCommand, {
       ...expectedCreateFleetRequest({
         ...defaultExpectedFleetRequestValues,
         totalTargetCapacity: 2,
@@ -1328,6 +1415,7 @@ describe('create runner with useDedicatedHost', () => {
     mockEC2Client.reset();
     mockSSMClient.reset();
 
+    mockEC2Client.on(DescribeSubnetsCommand).resolves(mockDefaultSubnets);
     mockEC2Client.on(RunInstancesCommand).resolves({
       Instances: [{ InstanceId: 'i-dedicated-1' }],
     });
@@ -1344,6 +1432,7 @@ describe('create runner with useDedicatedHost', () => {
     });
     expect(mockEC2Client).toHaveReceivedCommand(RunInstancesCommand);
     expect(mockEC2Client).not.toHaveReceivedCommand(CreateFleetCommand);
+    expect(mockEC2Client).not.toHaveReceivedCommand(DescribeSubnetsCommand);
   });
 
   it('uses CreateFleet when useDedicatedHost is false', async () => {

lambdas/libs/compute-providers/aws/ec2/src/control-plane/runners.ts

@@ -6,6 +6,7 @@ import {
   DeleteTagsCommand,
   DescribeInstancesCommand,
   DescribeInstancesResult,
+  DescribeSubnetsCommand,
   RunInstancesCommand,
   type RunInstancesCommandInput,
   RunInstancesCommandOutput,
@@ -32,6 +33,11 @@ interface Ec2Filter {
   Values: string[];
 }
 
+interface SubnetAllocation {
+  subnets: string[];
+  targetCapacity: number;
+}
+
 type FleetError = NonNullable<CreateFleetResult['Errors']>[number];
 
 export async function listEC2Runners(filters: Ec2ListRunnerFilters | undefined = undefined): Promise<RunnerInfo[]> {
@@ -196,6 +202,57 @@ function isRetryableAwsErrorName(errorName: string, configuredRetryableErrors: s
   return configuredRetryableErrors.includes(errorName) || RETRYABLE_AWS_ERROR_NAMES.has(errorName);
 }
 
+async function buildAzSafeSubnetSets(subnetIds: string[], ec2Client: EC2Client): Promise<string[][]> {
+  const uniqueSubnetIds = [...new Set(subnetIds)];
+  if (uniqueSubnetIds.length <= 1) {
+    return [uniqueSubnetIds];
+  }
+
+  const response = await ec2Client.send(new DescribeSubnetsCommand({ SubnetIds: uniqueSubnetIds }));
+  const availabilityZoneBySubnet = new Map<string, string>();
+  for (const subnet of response.Subnets || []) {
+    const availabilityZone = subnet.AvailabilityZoneId || subnet.AvailabilityZone;
+    if (subnet.SubnetId && availabilityZone) {
+      availabilityZoneBySubnet.set(subnet.SubnetId, availabilityZone);
+    }
+  }
+
+  const subnetsByAvailabilityZone = new Map<string, string[]>();
+  for (const subnetId of uniqueSubnetIds) {
+    const availabilityZone = availabilityZoneBySubnet.get(subnetId);
+    if (!availabilityZone) {
+      throw new Error(`Unable to resolve an Availability Zone for subnet '${subnetId}'.`);
+    }
+    const subnets = subnetsByAvailabilityZone.get(availabilityZone) || [];
+    subnets.push(subnetId);
+    subnetsByAvailabilityZone.set(availabilityZone, subnets);
+  }
+
+  const subnetSetCount = Math.max(...[...subnetsByAvailabilityZone.values()].map((subnets) => subnets.length));
+  const subnetSets = Array.from({ length: subnetSetCount }, (_, setIndex) =>
+    [...subnetsByAvailabilityZone.values()].map((subnets) => subnets[setIndex % subnets.length]),
+  );
+
+  logger.debug('Resolved AZ-safe subnet sets.', { subnetSets });
+  return subnetSets;
+}
+
+function buildSubnetAllocations(subnetSets: string[][], targetCapacity: number): SubnetAllocation[] {
+  if (subnetSets.length <= 1) {
+    return [{ subnets: subnetSets[0], targetCapacity }];
+  }
+
+  const startIndex = Math.floor(Math.random() * subnetSets.length);
+  const orderedSubnetSets = subnetSets.map((_, index) => subnetSets[(startIndex + index) % subnetSets.length]);
+  const baseTargetCapacity = Math.floor(targetCapacity / orderedSubnetSets.length);
+  const remainder = targetCapacity % orderedSubnetSets.length;
+
+  return orderedSubnetSets.map((subnets, index) => ({
+    subnets,
+    targetCapacity: baseTargetCapacity + (index < remainder ? 1 : 0),
+  }));
+}
+
 // The instance_allocation_strategy variable accepts the union of spot and on-demand strategies,
 // so a value valid for one capacity type can be invalid for the other. AWS rejects CreateFleet
 // when the strategy is not valid for the target capacity type, so fall back to 'lowest-price'
@@ -299,6 +356,57 @@ function buildRunInstancesOverrides(
 }
 
 export async function createRunner(runnerParameters: RunnerInputParameters): Promise<CreateRunnerResult> {
+  if (runnerParameters.useDedicatedHost || runnerParameters.ec2OverrideConfig?.SubnetId) {
+    return await createRunnerForSubnetSet(runnerParameters);
+  }
+
+  const ec2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION }));
+  let subnetSets: string[][];
+  try {
+    subnetSets = await buildAzSafeSubnetSets(runnerParameters.subnets, ec2Client);
+  } catch (error) {
+    const retryable = isRetryableAwsError(error, runnerParameters.scaleErrors);
+    logger.warn('Failed to resolve runner subnet Availability Zones.', {
+      error: error as Error,
+      retryable,
+    });
+    return failedCreateRunnerResult(runnerParameters.numberOfRunners, retryable);
+  }
+
+  if (subnetSets.length === 1) {
+    return await createRunnerForSubnetSet({ ...runnerParameters, subnets: subnetSets[0] });
+  }
+
+  const result: CreateRunnerResult = {
+    instances: [],
+    retryableErrorCount: 0,
+    nonRetryableErrorCount: 0,
+  };
+  let retryableCarry = 0;
+
+  const allocations = buildSubnetAllocations(subnetSets, runnerParameters.numberOfRunners);
+  for (const allocation of allocations) {
+    const targetCapacity = allocation.targetCapacity + retryableCarry;
+    retryableCarry = 0;
+    if (targetCapacity === 0) {
+      continue;
+    }
+
+    const allocationResult = await createRunnerForSubnetSet({
+      ...runnerParameters,
+      subnets: allocation.subnets,
+      numberOfRunners: targetCapacity,
+    });
+    result.instances.push(...allocationResult.instances);
+    result.nonRetryableErrorCount += allocationResult.nonRetryableErrorCount;
+    retryableCarry = allocationResult.retryableErrorCount;
+  }
+
+  result.retryableErrorCount = retryableCarry;
+  return result;
+}
+
+async function createRunnerForSubnetSet(runnerParameters: RunnerInputParameters): Promise<CreateRunnerResult> {
   logger.debug('Runner configuration.', {
     runner: {
       configuration: {

modules/multi-runner/README.md

@@ -202,7 +202,7 @@ module "multi-runner" {
 | <a name="input_scale_up_lambda_memory_size"></a> [scale\_up\_lambda\_memory\_size](#input\_scale\_up\_lambda\_memory\_size) | Memory size limit in MB for scale\_up lambda. | `number` | `512` | no |
 | <a name="input_ssm_paths"></a> [ssm\_paths](#input\_ssm\_paths) | The root path used in SSM to store configuration and secrets. | <pre>object({<br/>    root    = optional(string, "github-action-runners")<br/>    app     = optional(string, "app")<br/>    runners = optional(string, "runners")<br/>    webhook = optional(string, "webhook")<br/>  })</pre> | `{}` | no |
 | <a name="input_state_event_rule_binaries_syncer"></a> [state\_event\_rule\_binaries\_syncer](#input\_state\_event\_rule\_binaries\_syncer) | Option to disable EventBridge Lambda trigger for the binary syncer, useful to stop automatic updates of binary distribution | `string` | `"ENABLED"` | no |
-| <a name="input_subnet_ids"></a> [subnet\_ids](#input\_subnet\_ids) | List of subnets in which the action runners will be launched, the subnets needs to be subnets in the `vpc_id`. | `list(string)` | n/a | yes |
+| <a name="input_subnet_ids"></a> [subnet\_ids](#input\_subnet\_ids) | List of subnets in which the action runner instances will be launched. The subnets must exist in the configured VPC (`vpc_id`). | `list(string)` | n/a | yes |
 | <a name="input_syncer_lambda_s3_key"></a> [syncer\_lambda\_s3\_key](#input\_syncer\_lambda\_s3\_key) | S3 key for syncer lambda function. Required if using S3 bucket to specify lambdas. | `string` | `null` | no |
 | <a name="input_syncer_lambda_s3_object_version"></a> [syncer\_lambda\_s3\_object\_version](#input\_syncer\_lambda\_s3\_object\_version) | S3 object version for syncer lambda function. Useful if S3 versioning is enabled on source bucket. | `string` | `null` | no |
 | <a name="input_tags"></a> [tags](#input\_tags) | Map of tags that will be added to created resources. By default resources will be tagged with name and environment. | `map(string)` | `{}` | no |

modules/multi-runner/variables.tf

@@ -589,7 +589,7 @@ variable "vpc_id" {
 }
 
 variable "subnet_ids" {
-  description = "List of subnets in which the action runners will be launched, the subnets needs to be subnets in the `vpc_id`."
+  description = "List of subnets in which the action runner instances will be launched. The subnets must exist in the configured VPC (`vpc_id`)."
   type        = list(string)
 }