diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts b/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts index 1770dfa09a..2c3e72fe78 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts @@ -1,3 +1,4 @@ +import { DescribeLaunchTemplateVersionsCommand, EC2Client } from '@aws-sdk/client-ec2'; import { PutParameterCommand, SSMClient } from '@aws-sdk/client-ssm'; import { mockClient } from 'aws-sdk-client-mock'; import 'aws-sdk-client-mock-jest/vitest'; @@ -32,6 +33,7 @@ const mockOctokit = { const mockCreateRunner = vi.mocked(createRunner); const mockListRunners = vi.mocked(listEC2Runners); +const mockEC2Client = mockClient(EC2Client); const mockSSMClient = mockClient(SSMClient); const mockSSMgetParameter = vi.mocked(getParameter); const mockPublishRetryMessage = vi.mocked(publishRetryMessage); @@ -142,6 +144,22 @@ beforeEach(() => { vi.clearAllMocks(); setDefaults(); + mockEC2Client.reset(); + mockEC2Client.on(DescribeLaunchTemplateVersionsCommand).resolves({ + LaunchTemplateVersions: [ + { + LaunchTemplateData: { + BlockDeviceMappings: [ + { + DeviceName: '/dev/sda1', + Ebs: {}, + }, + ], + }, + }, + ], + }); + defaultSSMGetParameterMockImpl(); defaultOctokitMockImpl(); @@ -650,6 +668,86 @@ describe('scaleUp with GHES', () => { ); }); + it('loads the launch template block device name for dynamic EBS labels without DeviceName', async () => { + mockEC2Client.on(DescribeLaunchTemplateVersionsCommand).resolves({ + LaunchTemplateVersions: [ + { + LaunchTemplateData: { + BlockDeviceMappings: [ + { + DeviceName: '/dev/sdb', + VirtualName: 'ephemeral0', + }, + { + DeviceName: '/dev/sdf', + Ebs: {}, + }, + ], + }, + }, + ], + }); + + const testDataWithEbsLabels = [ + { + ...TEST_DATA_SINGLE, + labels: ['ghr-ec2-ebs-volume-size:100', 'ghr-ec2-ebs-volume-type:gp3'], + messageId: 'test-ebs-device-name', + }, + ]; + + await scaleUpModule.scaleUp(testDataWithEbsLabels); + + expect(mockEC2Client).toHaveReceivedCommandWith(DescribeLaunchTemplateVersionsCommand, { + LaunchTemplateName: 'lt-1', + Versions: ['$Default'], + }); + expect(createRunner).toBeCalledWith( + expect.objectContaining({ + ec2OverrideConfig: expect.objectContaining({ + BlockDeviceMappings: [ + { + DeviceName: '/dev/sdf', + Ebs: { + VolumeSize: 100, + VolumeType: 'gp3', + }, + }, + ], + }), + }), + ); + }); + + it('does not load launch template block device name when DeviceName is provided by labels', async () => { + const testDataWithEbsLabels = [ + { + ...TEST_DATA_SINGLE, + labels: ['ghr-ec2-block-device-name:/dev/sdg', 'ghr-ec2-ebs-volume-size:100', 'ghr-ec2-ebs-volume-type:gp3'], + messageId: 'test-explicit-ebs-device-name', + }, + ]; + + await scaleUpModule.scaleUp(testDataWithEbsLabels); + + expect(mockEC2Client).not.toHaveReceivedCommand(DescribeLaunchTemplateVersionsCommand); + expect(createRunner).toBeCalledWith( + expect.objectContaining({ + ec2OverrideConfig: expect.objectContaining({ + BlockDeviceMappings: [ + { + DeviceName: '/dev/sdg', + Ebs: { + VolumeSize: 100, + VolumeType: 'gp3', + }, + }, + ], + }), + }), + ); + }); + it('handles messages with no labels gracefully', async () => { const testDataWithNoLabels = [ { @@ -2459,11 +2557,16 @@ describe('parseEc2OverrideConfig', () => { }); describe('Placement', () => { - it('should parse placement-group label', () => { - const result = scaleUpModule.parseEc2OverrideConfig(['ghr-ec2-placement-group:my-placement-group']); + it('should parse placement-group-name label', () => { + const result = scaleUpModule.parseEc2OverrideConfig(['ghr-ec2-placement-group-name:my-placement-group']); expect(result?.Placement?.GroupName).toBe('my-placement-group'); }); + it('should parse placement-group-id label', () => { + const result = scaleUpModule.parseEc2OverrideConfig(['ghr-ec2-placement-group-id:pg-1234567890abcdef0']); + expect(result?.Placement?.GroupId).toBe('pg-1234567890abcdef0'); + }); + it('should parse placement-tenancy label', () => { const result = scaleUpModule.parseEc2OverrideConfig(['ghr-ec2-placement-tenancy:dedicated']); expect(result?.Placement?.Tenancy).toBe('dedicated'); @@ -2489,6 +2592,11 @@ describe('parseEc2OverrideConfig', () => { expect(result?.Placement?.AvailabilityZone).toBe('us-west-2b'); }); + it('should parse placement-availability-zone-id label', () => { + const result = scaleUpModule.parseEc2OverrideConfig(['ghr-ec2-placement-availability-zone-id:use1-az1']); + expect(result?.Placement?.AvailabilityZoneId).toBe('use1-az1'); + }); + it('should parse placement-spread-domain label', () => { const result = scaleUpModule.parseEc2OverrideConfig(['ghr-ec2-placement-spread-domain:my-spread-domain']); expect(result?.Placement?.SpreadDomain).toBe('my-spread-domain'); @@ -2505,7 +2613,7 @@ describe('parseEc2OverrideConfig', () => { it('should parse multiple placement labels', () => { const result = scaleUpModule.parseEc2OverrideConfig([ - 'ghr-ec2-placement-group:group-1', + 'ghr-ec2-placement-group-name:group-1', 'ghr-ec2-placement-tenancy:dedicated', 'ghr-ec2-placement-availability-zone:us-east-1b', ]); @@ -2516,6 +2624,17 @@ describe('parseEc2OverrideConfig', () => { }); describe('Block Device Mappings', () => { + it('should parse block-device-name label', () => { + const result = scaleUpModule.parseEc2OverrideConfig(['ghr-ec2-block-device-name:/dev/sdg']); + expect(result?.BlockDeviceMappings?.[0]?.DeviceName).toBe('/dev/sdg'); + }); + + it('should use default block device name when provided', () => { + const result = scaleUpModule.parseEc2OverrideConfig(['ghr-ec2-ebs-volume-size:100'], '/dev/sda1'); + expect(result?.BlockDeviceMappings?.[0]?.DeviceName).toBe('/dev/sda1'); + expect(result?.BlockDeviceMappings?.[0]?.Ebs?.VolumeSize).toBe(100); + }); + it('should parse ebs-volume-size label as number', () => { const result = scaleUpModule.parseEc2OverrideConfig(['ghr-ec2-ebs-volume-size:100']); expect(result?.BlockDeviceMappings?.[0]?.Ebs?.VolumeSize).toBe(100); @@ -2596,7 +2715,6 @@ describe('parseEc2OverrideConfig', () => { it('should initialize BlockDeviceMappings when not present', () => { const result = scaleUpModule.parseEc2OverrideConfig(['ghr-ec2-ebs-volume-size:50']); expect(result?.BlockDeviceMappings).toBeDefined(); - // expect(result?.BlockDeviceMappings?.[0]?.DeviceName).toBe('/dev/sda1'); }); }); @@ -3012,7 +3130,7 @@ describe('parseEc2OverrideConfig', () => { 'ghr-ec2-max-price:0.75', 'ghr-ec2-priority:1', // Placement - 'ghr-ec2-placement-group:my-group', + 'ghr-ec2-placement-group-name:my-group', 'ghr-ec2-placement-tenancy:dedicated', // Block Device 'ghr-ec2-ebs-volume-size:200', diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts index b742264842..8192a4071d 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts @@ -1,5 +1,9 @@ import { Octokit } from '@octokit/rest'; -import { addPersistentContextToChildLogger, createChildLogger } from '@aws-github-runner/aws-powertools-util'; +import { + addPersistentContextToChildLogger, + createChildLogger, + getTracedAWSV3Client, +} from '@aws-github-runner/aws-powertools-util'; import { getParameter, putParameter } from '@aws-github-runner/aws-ssm-util'; import yn from 'yn'; @@ -37,6 +41,8 @@ import { NetworkBandwidthGbpsRequest, TotalLocalStorageGBRequest, BaselineEbsBandwidthMbpsRequest, + DescribeLaunchTemplateVersionsCommand, + EC2Client, } from '@aws-sdk/client-ec2'; const logger = createChildLogger('scale-up'); @@ -469,7 +475,11 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise 0) { - ec2OverrideConfig = parseEc2OverrideConfig(dynamicEC2Labels); + const defaultBlockDeviceName = shouldLoadLaunchTemplateBlockDeviceName(dynamicEC2Labels) + ? await getDefaultBlockDeviceNameFromLaunchTemplate(launchTemplateName) + : undefined; + + ec2OverrideConfig = parseEc2OverrideConfig(dynamicEC2Labels, defaultBlockDeviceName); if (ec2OverrideConfig) { logger.debug('EC2 override config parsed from labels', { ec2OverrideConfig, @@ -838,16 +848,19 @@ async function createJitConfig( * - ghr-ec2-baseline-ebs-bandwidth-mbps-max: - Max baseline EBS bandwidth in Mbps * * Placement: - * - ghr-ec2-placement-group: - Placement group name + * - ghr-ec2-placement-group-name: - Placement group name + * - ghr-ec2-placement-group-id: - Placement group ID * - ghr-ec2-placement-tenancy: - Tenancy (default,dedicated,host) * - ghr-ec2-placement-host-id: - Dedicated host ID * - ghr-ec2-placement-affinity: - Affinity (default,host) * - ghr-ec2-placement-partition-number: - Partition number * - ghr-ec2-placement-availability-zone: - Placement availability zone + * - ghr-ec2-placement-availability-zone-id: - Placement availability zone ID * - ghr-ec2-placement-spread-domain: - Spread domain * - ghr-ec2-placement-host-resource-group-arn: - Host resource group ARN * * Block Device Mappings: + * - ghr-ec2-block-device-name: - Block device name * - ghr-ec2-ebs-volume-size: - EBS volume size in GB * - ghr-ec2-ebs-volume-type: - EBS volume type (gp2,gp3,io1,io2,st1,sc1) * - ghr-ec2-ebs-iops: - EBS IOPS @@ -871,9 +884,13 @@ async function createJitConfig( * runs-on: [self-hosted, linux, ghr-ec2-vcpu-count-min:4, ghr-ec2-memory-mib-min:16384, ghr-ec2-accelerator-types:gpu] * * @param labels - Array of GitHub workflow job labels + * @param defaultBlockDeviceName - Device name to use when dynamic block device labels create a mapping * @returns EC2 override configuration object or undefined if no valid config found */ -export function parseEc2OverrideConfig(labels: string[]): Ec2OverrideConfig | undefined { +export function parseEc2OverrideConfig( + labels: string[], + defaultBlockDeviceName?: string, +): Ec2OverrideConfig | undefined { const ec2Labels = labels.filter((l) => l.startsWith('ghr-ec2-')); const config: Ec2OverrideConfig = {}; @@ -906,61 +923,61 @@ export function parseEc2OverrideConfig(labels: string[]): Ec2OverrideConfig | un else if (key.startsWith('placement-')) { config.Placement = config.Placement || ({} as Placement); const placementKey = key.replace('placement-', ''); - if (placementKey === 'group') { - config.Placement.GroupName = value; - } else if (placementKey === 'tenancy') { - config.Placement.Tenancy = value as Tenancy; - } else if (placementKey === 'host-id') { - config.Placement.HostId = value; + if (placementKey === 'availability-zone-id') { + config.Placement.AvailabilityZoneId = value; } else if (placementKey === 'affinity') { config.Placement.Affinity = value; + } else if (placementKey === 'group-name') { + config.Placement.GroupName = value; } else if (placementKey === 'partition-number') { config.Placement.PartitionNumber = parseInt(value, 10); - } else if (placementKey === 'availability-zone') { - config.Placement.AvailabilityZone = value; + } else if (placementKey === 'host-id') { + config.Placement.HostId = value; + } else if (placementKey === 'tenancy') { + config.Placement.Tenancy = value as Tenancy; } else if (placementKey === 'spread-domain') { config.Placement.SpreadDomain = value; } else if (placementKey === 'host-resource-group-arn') { config.Placement.HostResourceGroupArn = value; + } else if (placementKey === 'group-id') { + config.Placement.GroupId = value; + } else if (placementKey === 'availability-zone') { + config.Placement.AvailabilityZone = value; } } - // Block Device Mappings (EBS) - else if (key.startsWith('ebs-')) { - config.BlockDeviceMappings = config.BlockDeviceMappings || ([{}] as FleetBlockDeviceMappingRequest[]); + // Block Device Mappings + else if (key === 'block-device-name') { + getOrCreateBlockDeviceMapping(config, defaultBlockDeviceName).DeviceName = value; + } else if (key === 'block-device-virtual-name') { + getOrCreateBlockDeviceMapping(config, defaultBlockDeviceName).VirtualName = value; + } else if (key.startsWith('ebs-')) { + const blockDeviceMapping = getOrCreateBlockDeviceMapping(config, defaultBlockDeviceName); const ebsKey = key.replace('ebs-', ''); - const ebs = - config.BlockDeviceMappings[0].Ebs || (config.BlockDeviceMappings[0].Ebs = {} as FleetEbsBlockDeviceRequest); + const ebs = blockDeviceMapping.Ebs || (blockDeviceMapping.Ebs = {} as FleetEbsBlockDeviceRequest); - if (ebsKey === 'volume-size') { - ebs.VolumeSize = parseInt(value, 10); - } else if (ebsKey === 'volume-type') { - ebs.VolumeType = value as VolumeType; + if (ebsKey === 'encrypted') { + ebs.Encrypted = value.toLowerCase() === 'true'; + } else if (ebsKey === 'delete-on-termination') { + ebs.DeleteOnTermination = value.toLowerCase() === 'true'; } else if (ebsKey === 'iops') { ebs.Iops = parseInt(value, 10); } else if (ebsKey === 'throughput') { ebs.Throughput = parseInt(value, 10); - } else if (ebsKey === 'encrypted') { - ebs.Encrypted = value.toLowerCase() === 'true'; } else if (ebsKey === 'kms-key-id') { ebs.KmsKeyId = value; - } else if (ebsKey === 'delete-on-termination') { - ebs.DeleteOnTermination = value.toLowerCase() === 'true'; } else if (ebsKey === 'snapshot-id') { ebs.SnapshotId = value; + } else if (ebsKey === 'volume-size') { + ebs.VolumeSize = parseInt(value, 10); + } else if (ebsKey === 'volume-type') { + ebs.VolumeType = value as VolumeType; } - } - - // Block Device Mappings (Non-EBS) - else if (key === 'block-device-virtual-name') { - config.BlockDeviceMappings = config.BlockDeviceMappings || ([{}] as FleetBlockDeviceMappingRequest[]); - config.BlockDeviceMappings[0].VirtualName = value; } else if (key === 'block-device-no-device') { - config.BlockDeviceMappings = config.BlockDeviceMappings || ([{}] as FleetBlockDeviceMappingRequest[]); - config.BlockDeviceMappings[0].NoDevice = value; + getOrCreateBlockDeviceMapping(config, defaultBlockDeviceName).NoDevice = value; } - // Instance Requirements - vCPU & Memory + // Instance Requirements else if (key.startsWith('vcpu-count-')) { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); config.InstanceRequirements.VCpuCount = config.InstanceRequirements.VCpuCount || ({} as VCpuCountRangeRequest); @@ -971,76 +988,43 @@ export function parseEc2OverrideConfig(labels: string[]): Ec2OverrideConfig | un config.InstanceRequirements.MemoryMiB = config.InstanceRequirements.MemoryMiB || ({} as MemoryMiBRequest); const subKey = key.replace('memory-mib-', ''); config.InstanceRequirements.MemoryMiB![subKey === 'min' ? 'Min' : 'Max'] = parseInt(value, 10); + } else if (key === 'cpu-manufacturers') { + config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); + config.InstanceRequirements.CpuManufacturers = value.split(',') as CpuManufacturer[]; } else if (key.startsWith('memory-gib-per-vcpu-')) { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); config.InstanceRequirements.MemoryGiBPerVCpu = config.InstanceRequirements.MemoryGiBPerVCpu || ({} as MemoryGiBPerVCpuRequest); const subKey = key.replace('memory-gib-per-vcpu-', ''); config.InstanceRequirements.MemoryGiBPerVCpu![subKey === 'min' ? 'Min' : 'Max'] = parseFloat(value); - } - - // Instance Requirements - CPU & Performance - else if (key === 'cpu-manufacturers') { + } else if (key === 'excluded-instance-types') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.CpuManufacturers = value.split(',') as CpuManufacturer[]; + config.InstanceRequirements.ExcludedInstanceTypes = value.split(','); } else if (key === 'instance-generations') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); config.InstanceRequirements.InstanceGenerations = value.split(',') as InstanceGeneration[]; - } else if (key === 'excluded-instance-types') { + } else if (key === 'spot-max-price-percentage-over-lowest-price') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.ExcludedInstanceTypes = value.split(','); - } else if (key === 'allowed-instance-types') { - config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.AllowedInstanceTypes = value.split(','); - } else if (key === 'burstable-performance') { + config.InstanceRequirements.SpotMaxPricePercentageOverLowestPrice = parseInt(value, 10); + } else if (key === 'on-demand-max-price-percentage-over-lowest-price') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.BurstablePerformance = value as BurstablePerformance; + config.InstanceRequirements.OnDemandMaxPricePercentageOverLowestPrice = parseInt(value, 10); } else if (key === 'bare-metal') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); config.InstanceRequirements.BareMetal = value as BareMetal; - } - - // Instance Requirements - Accelerators - else if (key.startsWith('accelerator-count-')) { - config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.AcceleratorCount = - config.InstanceRequirements.AcceleratorCount || ({} as AcceleratorCountRequest); - const subKey = key.replace('accelerator-count-', ''); - config.InstanceRequirements.AcceleratorCount![subKey === 'min' ? 'Min' : 'Max'] = parseInt(value, 10); - } else if (key === 'accelerator-types') { - config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.AcceleratorTypes = value.split(',') as AcceleratorType[]; - } else if (key === 'accelerator-manufacturers') { - config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.AcceleratorManufacturers = value.split(',') as AcceleratorManufacturer[]; - } else if (key === 'accelerator-names') { + } else if (key === 'burstable-performance') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.AcceleratorNames = value.split(',') as AcceleratorName[]; - } else if (key.startsWith('accelerator-total-memory-mib-')) { + config.InstanceRequirements.BurstablePerformance = value as BurstablePerformance; + } else if (key === 'require-hibernate-support') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.AcceleratorTotalMemoryMiB = - config.InstanceRequirements.AcceleratorTotalMemoryMiB || ({} as AcceleratorTotalMemoryMiBRequest); - const subKey = key.replace('accelerator-total-memory-mib-', ''); - config.InstanceRequirements.AcceleratorTotalMemoryMiB![subKey === 'min' ? 'Min' : 'Max'] = parseInt(value, 10); - } - - // Instance Requirements - Network - else if (key.startsWith('network-interface-count-')) { + config.InstanceRequirements.RequireHibernateSupport = value.toLowerCase() === 'true'; + } else if (key.startsWith('network-interface-count-')) { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); config.InstanceRequirements.NetworkInterfaceCount = config.InstanceRequirements.NetworkInterfaceCount || ({} as NetworkInterfaceCountRequest); const subKey = key.replace('network-interface-count-', ''); config.InstanceRequirements.NetworkInterfaceCount![subKey === 'min' ? 'Min' : 'Max'] = parseInt(value, 10); - } else if (key.startsWith('network-bandwidth-gbps-')) { - config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.NetworkBandwidthGbps = - config.InstanceRequirements.NetworkBandwidthGbps || ({} as NetworkBandwidthGbpsRequest); - const subKey = key.replace('network-bandwidth-gbps-', ''); - config.InstanceRequirements.NetworkBandwidthGbps![subKey === 'min' ? 'Min' : 'Max'] = parseFloat(value); - } - - // Instance Requirements - Storage - else if (key === 'local-storage') { + } else if (key === 'local-storage') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); config.InstanceRequirements.LocalStorage = value as LocalStorage; } else if (key === 'local-storage-types') { @@ -1058,24 +1042,39 @@ export function parseEc2OverrideConfig(labels: string[]): Ec2OverrideConfig | un config.InstanceRequirements.BaselineEbsBandwidthMbps || ({} as BaselineEbsBandwidthMbpsRequest); const subKey = key.replace('baseline-ebs-bandwidth-mbps-', ''); config.InstanceRequirements.BaselineEbsBandwidthMbps![subKey === 'min' ? 'Min' : 'Max'] = parseInt(value, 10); - } - - // Instance Requirements - Pricing & Other - else if (key === 'spot-max-price-percentage-over-lowest-price') { + } else if (key === 'accelerator-types') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.SpotMaxPricePercentageOverLowestPrice = parseInt(value, 10); - } else if (key === 'on-demand-max-price-percentage-over-lowest-price') { + config.InstanceRequirements.AcceleratorTypes = value.split(',') as AcceleratorType[]; + } else if (key.startsWith('accelerator-count-')) { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.OnDemandMaxPricePercentageOverLowestPrice = parseInt(value, 10); - } else if (key === 'max-spot-price-as-percentage-of-optimal-on-demand-price') { + config.InstanceRequirements.AcceleratorCount = + config.InstanceRequirements.AcceleratorCount || ({} as AcceleratorCountRequest); + const subKey = key.replace('accelerator-count-', ''); + config.InstanceRequirements.AcceleratorCount![subKey === 'min' ? 'Min' : 'Max'] = parseInt(value, 10); + } else if (key === 'accelerator-manufacturers') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.MaxSpotPriceAsPercentageOfOptimalOnDemandPrice = parseInt(value, 10); - } else if (key === 'require-hibernate-support') { + config.InstanceRequirements.AcceleratorManufacturers = value.split(',') as AcceleratorManufacturer[]; + } else if (key === 'accelerator-names') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.RequireHibernateSupport = value.toLowerCase() === 'true'; - } else if (key === 'require-encryption-in-transit') { + config.InstanceRequirements.AcceleratorNames = value.split(',') as AcceleratorName[]; + } else if (key.startsWith('accelerator-total-memory-mib-')) { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); - config.InstanceRequirements.RequireEncryptionInTransit = value.toLowerCase() === 'true'; + config.InstanceRequirements.AcceleratorTotalMemoryMiB = + config.InstanceRequirements.AcceleratorTotalMemoryMiB || ({} as AcceleratorTotalMemoryMiBRequest); + const subKey = key.replace('accelerator-total-memory-mib-', ''); + config.InstanceRequirements.AcceleratorTotalMemoryMiB![subKey === 'min' ? 'Min' : 'Max'] = parseInt(value, 10); + } else if (key.startsWith('network-bandwidth-gbps-')) { + config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); + config.InstanceRequirements.NetworkBandwidthGbps = + config.InstanceRequirements.NetworkBandwidthGbps || ({} as NetworkBandwidthGbpsRequest); + const subKey = key.replace('network-bandwidth-gbps-', ''); + config.InstanceRequirements.NetworkBandwidthGbps![subKey === 'min' ? 'Min' : 'Max'] = parseFloat(value); + } else if (key === 'allowed-instance-types') { + config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); + config.InstanceRequirements.AllowedInstanceTypes = value.split(','); + } else if (key === 'max-spot-price-as-percentage-of-optimal-on-demand-price') { + config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); + config.InstanceRequirements.MaxSpotPriceAsPercentageOfOptimalOnDemandPrice = parseInt(value, 10); } else if (key === 'baseline-performance-factors-cpu-reference-families') { config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); config.InstanceRequirements.BaselinePerformanceFactors = @@ -1085,12 +1084,62 @@ export function parseEc2OverrideConfig(labels: string[]): Ec2OverrideConfig | un config.InstanceRequirements.BaselinePerformanceFactors.Cpu.References = value .split(',') .map((family) => ({ InstanceFamily: family })) as PerformanceFactorReferenceRequest[]; + } else if (key === 'require-encryption-in-transit') { + config.InstanceRequirements = config.InstanceRequirements || ({} as InstanceRequirementsRequest); + config.InstanceRequirements.RequireEncryptionInTransit = value.toLowerCase() === 'true'; } } return Object.keys(config).length > 0 ? config : undefined; } +function getOrCreateBlockDeviceMapping( + config: Ec2OverrideConfig, + defaultBlockDeviceName?: string, +): FleetBlockDeviceMappingRequest { + config.BlockDeviceMappings = + config.BlockDeviceMappings || + ([defaultBlockDeviceName ? { DeviceName: defaultBlockDeviceName } : {}] as FleetBlockDeviceMappingRequest[]); + return config.BlockDeviceMappings[0]; +} + +function shouldLoadLaunchTemplateBlockDeviceName(labels: string[]): boolean { + const blockDeviceNameLabel = 'ghr-ec2-block-device-name:'; + let hasBlockDeviceOverride = false; + let hasBlockDeviceName = false; + + for (const label of labels) { + hasBlockDeviceOverride = + hasBlockDeviceOverride || label.startsWith('ghr-ec2-ebs-') || label.startsWith('ghr-ec2-block-device-'); + + hasBlockDeviceName = + hasBlockDeviceName || (label.startsWith(blockDeviceNameLabel) && label.slice(blockDeviceNameLabel.length) !== ''); + } + + return hasBlockDeviceOverride && !hasBlockDeviceName; +} + +async function getDefaultBlockDeviceNameFromLaunchTemplate(launchTemplateName: string): Promise { + const ec2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); + const launchTemplateVersions = await ec2Client.send( + new DescribeLaunchTemplateVersionsCommand({ + LaunchTemplateName: launchTemplateName, + Versions: ['$Default'], + }), + ); + const blockDeviceMappings = + launchTemplateVersions.LaunchTemplateVersions?.[0]?.LaunchTemplateData?.BlockDeviceMappings; + const blockDeviceName = + blockDeviceMappings?.find((blockDeviceMapping) => blockDeviceMapping.DeviceName && blockDeviceMapping.Ebs) + ?.DeviceName ?? blockDeviceMappings?.find((blockDeviceMapping) => blockDeviceMapping.DeviceName)?.DeviceName; + + if (!blockDeviceName) { + throw new Error(`Failed to determine block device name from launch template '${launchTemplateName}'.`); + } + + return blockDeviceName; +} + function labelsHash(labels: string[]): string { const prefix = 'ghr-'; diff --git a/modules/runners/policies/lambda-scale-up.json b/modules/runners/policies/lambda-scale-up.json index 86b7415637..4a38d23f6a 100644 --- a/modules/runners/policies/lambda-scale-up.json +++ b/modules/runners/policies/lambda-scale-up.json @@ -5,6 +5,7 @@ "Effect": "Allow", "Action": [ "ec2:DescribeInstances", + "ec2:DescribeLaunchTemplateVersions", "ec2:DescribeTags", "ec2:RunInstances", "ec2:CreateFleet",