feat: add support for all fields in FleetLaunchTemplateOverridesRequest

Author: edersonbrilhante

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

@@ -1,4 +1,4 @@
-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 +29,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 +55,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

@@ -425,6 +425,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', () => {

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,
           ),
         },
       ],