feat: Add feature to enable dynamic ec2 config via workflow labels

README.md

@@ -121,6 +121,7 @@ Join our discord community via [this invite link](https://discord.gg/bxgXW8jJGh)
 | <a name="input_disable_runner_autoupdate"></a> [disable\_runner\_autoupdate](#input\_disable\_runner\_autoupdate) | Disable the auto update of the github runner agent. Be aware there is a grace period of 30 days, see also the [GitHub article](https://github.blog/changelog/2022-02-01-github-actions-self-hosted-runners-can-now-disable-automatic-updates/) | `bool` | `false` | no |
 | <a name="input_enable_ami_housekeeper"></a> [enable\_ami\_housekeeper](#input\_enable\_ami\_housekeeper) | Option to disable the lambda to clean up old AMIs. | `bool` | `false` | no |
 | <a name="input_enable_cloudwatch_agent"></a> [enable\_cloudwatch\_agent](#input\_enable\_cloudwatch\_agent) | Enables the cloudwatch agent on the ec2 runner instances. The runner uses a default config that can be overridden via `cloudwatch_config`. | `bool` | `true` | no |
+| <a name="input_enable_dynamic_labels"></a> [enable\_dynamic\_labels](#input\_enable\_dynamic\_labels) | Enable dynamic EC2 configs based on workflow job labels. When enabled, jobs can request specific configs via the 'gh-ec2-<config type key>:<config type value>' label (e.g., 'gh-ec2-instance-type:t3.large'). | `bool` | `false` | no |
 | <a name="input_enable_ephemeral_runners"></a> [enable\_ephemeral\_runners](#input\_enable\_ephemeral\_runners) | Enable ephemeral runners, runners will only be used once. | `bool` | `false` | no |
 | <a name="input_enable_jit_config"></a> [enable\_jit\_config](#input\_enable\_jit\_config) | Overwrite the default behavior for JIT configuration. By default JIT configuration is enabled for ephemeral runners and disabled for non-ephemeral runners. In case of GHES check first if the JIT config API is available. In case you are upgrading from 3.x to 4.x you can set `enable_jit_config` to `false` to avoid a breaking change when having your own AMI. | `bool` | `null` | no |
 | <a name="input_enable_job_queued_check"></a> [enable\_job\_queued\_check](#input\_enable\_job\_queued\_check) | Only scale if the job event received by the scale up lambda is in the queued state. By default enabled for non ephemeral runners and disabled for ephemeral. Set this variable to overwrite the default behavior. | `bool` | `null` | no |

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

@@ -1,4 +1,11 @@
-import { DefaultTargetCapacityType, SpotAllocationStrategy } from '@aws-sdk/client-ec2';
+import {
+  DefaultTargetCapacityType,
+  InstanceRequirementsRequest,
+  SpotAllocationStrategy,
+  _InstanceType,
+  Placement,
+  FleetBlockDeviceMappingRequest,
+} from '@aws-sdk/client-ec2';
 
 export type RunnerType = 'Org' | 'Repo';
 
@@ -29,6 +36,20 @@ export interface ListRunnerFilters {
   statuses?: string[];
 }
 
+export interface Ec2OverrideConfig {
+  InstanceType?: _InstanceType;
+  MaxPrice?: string;
+  SubnetId?: string;
+  AvailabilityZone?: string;
+  WeightedCapacity?: number;
+  Priority?: number;
+  Placement?: Placement;
+  BlockDeviceMappings?: FleetBlockDeviceMappingRequest[];
+  InstanceRequirements?: InstanceRequirementsRequest;
+  ImageId?: string;
+  AvailabilityZoneId?: string;
+}
+
 export interface RunnerInputParameters {
   environment: string;
   runnerType: RunnerType;
@@ -41,6 +62,7 @@ export interface RunnerInputParameters {
     maxSpotPrice?: string;
     instanceAllocationStrategy: SpotAllocationStrategy;
   };
+  ec2OverrideConfig?: Ec2OverrideConfig;
   numberOfRunners: number;
   amiIdSsmParameterName?: string;
   tracingEnabled?: boolean;

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

@@ -318,6 +318,7 @@ describe('create runner', () => {
     allocationStrategy: SpotAllocationStrategy.CAPACITY_OPTIMIZED,
     capacityType: 'spot',
     type: 'Org',
+    scaleErrors: ['UnfulfillableCapacity', 'MaxSpotInstanceCountExceeded'],
   };
 
   const defaultExpectedFleetRequestValues: ExpectedFleetRequestValues = {
@@ -425,6 +426,215 @@ describe('create runner', () => {
       }),
     });
   });
+
+  it('overrides SubnetId when specified in ec2OverrideConfig', async () => {
+    await createRunner({
+      ...createRunnerConfig(defaultRunnerConfig),
+      ec2OverrideConfig: {
+        SubnetId: 'subnet-override',
+      },
+    });
+
+    expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
+      LaunchTemplateConfigs: [
+        {
+          LaunchTemplateSpecification: {
+            LaunchTemplateName: 'lt-1',
+            Version: '$Default',
+          },
+          Overrides: [
+            {
+              InstanceType: 'm5.large',
+              SubnetId: 'subnet-override',
+            },
+            {
+              InstanceType: 'c5.large',
+              SubnetId: 'subnet-override',
+            },
+          ],
+        },
+      ],
+      SpotOptions: {
+        AllocationStrategy: SpotAllocationStrategy.CAPACITY_OPTIMIZED,
+      },
+      TagSpecifications: expect.any(Array),
+      TargetCapacitySpecification: {
+        DefaultTargetCapacityType: 'spot',
+        TotalTargetCapacity: 1,
+      },
+      Type: 'instant',
+    });
+  });
+
+  it('overrides InstanceType when specified in ec2OverrideConfig', async () => {
+    await createRunner({
+      ...createRunnerConfig(defaultRunnerConfig),
+      ec2OverrideConfig: {
+        InstanceType: 't3.xlarge',
+      },
+    });
+
+    expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
+      LaunchTemplateConfigs: [
+        {
+          LaunchTemplateSpecification: {
+            LaunchTemplateName: 'lt-1',
+            Version: '$Default',
+          },
+          Overrides: [
+            {
+              InstanceType: 't3.xlarge',
+              SubnetId: 'subnet-123',
+            },
+            {
+              InstanceType: 't3.xlarge',
+              SubnetId: 'subnet-456',
+            },
+          ],
+        },
+      ],
+      SpotOptions: {
+        AllocationStrategy: SpotAllocationStrategy.CAPACITY_OPTIMIZED,
+      },
+      TagSpecifications: expect.any(Array),
+      TargetCapacitySpecification: {
+        DefaultTargetCapacityType: 'spot',
+        TotalTargetCapacity: 1,
+      },
+      Type: 'instant',
+    });
+  });
+
+  it('overrides ImageId when specified in ec2OverrideConfig', async () => {
+    await createRunner({
+      ...createRunnerConfig(defaultRunnerConfig),
+      ec2OverrideConfig: {
+        ImageId: 'ami-override-123',
+      },
+    });
+
+    expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
+      LaunchTemplateConfigs: [
+        {
+          LaunchTemplateSpecification: {
+            LaunchTemplateName: 'lt-1',
+            Version: '$Default',
+          },
+          Overrides: [
+            {
+              InstanceType: 'm5.large',
+              SubnetId: 'subnet-123',
+              ImageId: 'ami-override-123',
+            },
+            {
+              InstanceType: 'c5.large',
+              SubnetId: 'subnet-123',
+              ImageId: 'ami-override-123',
+            },
+            {
+              InstanceType: 'm5.large',
+              SubnetId: 'subnet-456',
+              ImageId: 'ami-override-123',
+            },
+            {
+              InstanceType: 'c5.large',
+              SubnetId: 'subnet-456',
+              ImageId: 'ami-override-123',
+            },
+          ],
+        },
+      ],
+      SpotOptions: {
+        AllocationStrategy: SpotAllocationStrategy.CAPACITY_OPTIMIZED,
+      },
+      TagSpecifications: expect.any(Array),
+      TargetCapacitySpecification: {
+        DefaultTargetCapacityType: 'spot',
+        TotalTargetCapacity: 1,
+      },
+      Type: 'instant',
+    });
+  });
+
+  it('overrides all three fields (SubnetId, InstanceType, ImageId) when specified in ec2OverrideConfig', async () => {
+    await createRunner({
+      ...createRunnerConfig(defaultRunnerConfig),
+      ec2OverrideConfig: {
+        SubnetId: 'subnet-custom',
+        InstanceType: 'c5.2xlarge',
+        ImageId: 'ami-custom-456',
+      },
+    });
+
+    expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
+      LaunchTemplateConfigs: [
+        {
+          LaunchTemplateSpecification: {
+            LaunchTemplateName: 'lt-1',
+            Version: '$Default',
+          },
+          Overrides: [
+            {
+              InstanceType: 'c5.2xlarge',
+              SubnetId: 'subnet-custom',
+              ImageId: 'ami-custom-456',
+            },
+          ],
+        },
+      ],
+      SpotOptions: {
+        AllocationStrategy: SpotAllocationStrategy.CAPACITY_OPTIMIZED,
+      },
+      TagSpecifications: expect.any(Array),
+      TargetCapacitySpecification: {
+        DefaultTargetCapacityType: 'spot',
+        TotalTargetCapacity: 1,
+      },
+      Type: 'instant',
+    });
+  });
+
+  it('spreads additional ec2OverrideConfig properties to Overrides', async () => {
+    await createRunner({
+      ...createRunnerConfig(defaultRunnerConfig),
+      ec2OverrideConfig: {
+        SubnetId: 'subnet-override',
+        InstanceType: 't3.medium',
+        MaxPrice: '0.05',
+        Priority: 1.5,
+        WeightedCapacity: 2.0,
+      },
+    });
+
+    expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
+      LaunchTemplateConfigs: [
+        {
+          LaunchTemplateSpecification: {
+            LaunchTemplateName: 'lt-1',
+            Version: '$Default',
+          },
+          Overrides: [
+            {
+              InstanceType: 't3.medium',
+              SubnetId: 'subnet-override',
+              MaxPrice: '0.05',
+              Priority: 1.5,
+              WeightedCapacity: 2.0,
+            },
+          ],
+        },
+      ],
+      SpotOptions: {
+        AllocationStrategy: SpotAllocationStrategy.CAPACITY_OPTIMIZED,
+      },
+      TagSpecifications: expect.any(Array),
+      TargetCapacitySpecification: {
+        DefaultTargetCapacityType: 'spot',
+        TotalTargetCapacity: 1,
+      },
+      Type: 'instant',
+    });
+  });
 });
 
 describe('create runner with errors', () => {
@@ -546,6 +756,7 @@ describe('create runner with errors fail over to OnDemand', () => {
     capacityType: 'spot',
     type: 'Repo',
     onDemandFailoverOnError: ['InsufficientInstanceCapacity'],
+    scaleErrors: ['UnfulfillableCapacity', 'MaxSpotInstanceCountExceeded'],
   };
   const defaultExpectedFleetRequestValues: ExpectedFleetRequestValues = {
     type: 'Repo',

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

@@ -125,14 +125,22 @@ function generateFleetOverrides(
   subnetIds: string[],
   instancesTypes: string[],
   amiId?: string,
+  ec2OverrideConfig?: Runners.Ec2OverrideConfig,
 ): FleetLaunchTemplateOverridesRequest[] {
   const result: FleetLaunchTemplateOverridesRequest[] = [];
-  subnetIds.forEach((s) => {
-    instancesTypes.forEach((i) => {
+
+  // Use override values if available, otherwise use parameter arrays
+  const subnetsToUse = ec2OverrideConfig?.SubnetId ? [ec2OverrideConfig.SubnetId] : subnetIds;
+  const instanceTypesToUse = ec2OverrideConfig?.InstanceType ? [ec2OverrideConfig.InstanceType] : instancesTypes;
+  const amiIdToUse = ec2OverrideConfig?.ImageId ?? amiId;
+
+  subnetsToUse.forEach((s) => {
+    instanceTypesToUse.forEach((i) => {
       const item: FleetLaunchTemplateOverridesRequest = {
         SubnetId: s,
         InstanceType: i as _InstanceType,
-        ImageId: amiId,
+        ImageId: amiIdToUse,
+        ...ec2OverrideConfig,
       };
       result.push(item);
     });
@@ -265,6 +273,7 @@ async function createInstances(
             runnerParameters.subnets,
             runnerParameters.ec2instanceCriteria.instanceTypes,
             amiIdOverride,
+            runnerParameters.ec2OverrideConfig,
           ),
         },
       ],

lambdas/functions/control-plane/src/scale-runners/ScaleError.test.ts

@@ -23,10 +23,42 @@ describe('ScaleError', () => {
 
   describe('toBatchItemFailures', () => {
     const mockMessages: ActionRequestMessageSQS[] = [
-      { messageId: 'msg-1', id: 1, eventType: 'workflow_job' },
-      { messageId: 'msg-2', id: 2, eventType: 'workflow_job' },
-      { messageId: 'msg-3', id: 3, eventType: 'workflow_job' },
-      { messageId: 'msg-4', id: 4, eventType: 'workflow_job' },
+      {
+        messageId: 'msg-1',
+        id: 1,
+        eventType: 'workflow_job',
+        repositoryName: 'repo',
+        repositoryOwner: 'owner',
+        installationId: 123,
+        repoOwnerType: 'Organization',
+      },
+      {
+        messageId: 'msg-2',
+        id: 2,
+        eventType: 'workflow_job',
+        repositoryName: 'repo',
+        repositoryOwner: 'owner',
+        installationId: 123,
+        repoOwnerType: 'Organization',
+      },
+      {
+        messageId: 'msg-3',
+        id: 3,
+        eventType: 'workflow_job',
+        repositoryName: 'repo',
+        repositoryOwner: 'owner',
+        installationId: 123,
+        repoOwnerType: 'Organization',
+      },
+      {
+        messageId: 'msg-4',
+        id: 4,
+        eventType: 'workflow_job',
+        repositoryName: 'repo',
+        repositoryOwner: 'owner',
+        installationId: 123,
+        repoOwnerType: 'Organization',
+      },
     ];
 
     it.each([