fix: Custom capacity provider for ECS was broken (#336)

Fixes #333
CloudSnorkel · May 11, 2023 · 5bb58e4 · 5bb58e4
1 parent c7313ac
commit 5bb58e4
Show file tree

Hide file tree

Showing 3 changed files with 113 additions and 48 deletions.
diff --git a/API.md b/API.md
diff --git a/src/providers/ecs.ts b/src/providers/ecs.ts
@@ -83,6 +83,8 @@ export interface EcsRunnerProviderProps extends RunnerProviderProps {
   /**
    * Existing capacity provider to use.
    *
+   * Make sure the AMI used by the capacity provider is compatible with ECS.
+   *
    * @default new capacity provider
    */
   readonly capacityProvider?: ecs.AsgCapacityProvider;
@@ -325,54 +327,59 @@ export class EcsRunnerProvider extends BaseProvider implements IRunnerProvider {
     const imageBuilder = props?.imageBuilder ?? EcsRunnerProvider.imageBuilder(this, 'Image Builder');
     const image = this.image = imageBuilder.bindDockerImage();
 
-    if (props?.capacityProvider && (props?.minInstances || props?.maxInstances || props?.instanceType || props?.storageSize)) {
-      cdk.Annotations.of(this).addWarning('When using a custom capacity provider, minInstances, maxInstances, instanceType and storageSize will be ignored.');
-    }
+    if (props?.capacityProvider) {
+      if (props?.minInstances || props?.maxInstances || props?.instanceType || props?.storageSize || props?.spot || props?.spotMaxPrice) {
+        cdk.Annotations.of(this).addWarning('When using a custom capacity provider, minInstances, maxInstances, instanceType, storageSize, spot, and spotMaxPrice will be ignored.');
+      }
 
-    const spot = props?.spot ?? props?.spotMaxPrice !== undefined;
-
-    const launchTemplate = new ec2.LaunchTemplate(this, 'Launch Template', {
-      machineImage: this.defaultClusterInstanceAmi(),
-      instanceType: props?.instanceType ?? this.defaultClusterInstanceType(),
-      blockDevices: props?.storageSize ? [
-        {
-          deviceName: '/dev/sda1',
-          volume: {
-            ebsDevice: {
-              volumeSize: props?.storageSize?.toGibibytes(),
-              deleteOnTermination: true,
+      this.capacityProvider = props.capacityProvider;
+    } else {
+      const spot = props?.spot ?? props?.spotMaxPrice !== undefined;
+
+      const launchTemplate = new ec2.LaunchTemplate(this, 'Launch Template', {
+        machineImage: this.defaultClusterInstanceAmi(),
+        instanceType: props?.instanceType ?? this.defaultClusterInstanceType(),
+        blockDevices: props?.storageSize ? [
+          {
+            deviceName: '/dev/sda1',
+            volume: {
+              ebsDevice: {
+                volumeSize: props?.storageSize?.toGibibytes(),
+                deleteOnTermination: true,
+              },
             },
           },
-        },
-      ] : undefined,
-      spotOptions: spot ? {
-        requestType: ec2.SpotRequestType.ONE_TIME,
-        maxPrice: props?.spotMaxPrice ? parseFloat(props?.spotMaxPrice) : undefined,
-      } : undefined,
-      requireImdsv2: true,
-      securityGroup: this.securityGroups[0],
-      role: new iam.Role(this, 'Launch Template Role', {
-        assumedBy: new iam.ServicePrincipal('ec2.amazonaws.com'),
-      }),
-      userData: ec2.UserData.forOperatingSystem(image.os.is(Os.WINDOWS) ? ec2.OperatingSystemType.WINDOWS : ec2.OperatingSystemType.LINUX),
-    });
-    this.securityGroups.slice(1).map(sg => launchTemplate.connections.addSecurityGroup(sg));
-
-    const autoScalingGroup = new autoscaling.AutoScalingGroup(this, 'Auto Scaling Group', {
-      vpc: this.vpc,
-      launchTemplate,
-      vpcSubnets: this.subnetSelection,
-      minCapacity: props?.minInstances ?? 0,
-      maxCapacity: props?.maxInstances ?? 5,
-    });
-    autoScalingGroup.addUserData(this.loginCommand(), this.pullCommand());
-    autoScalingGroup.role.addManagedPolicy(iam.ManagedPolicy.fromAwsManagedPolicyName('AmazonSSMManagedInstanceCore'));
-    image.imageRepository.grantPull(autoScalingGroup);
+        ] : undefined,
+        spotOptions: spot ? {
+          requestType: ec2.SpotRequestType.ONE_TIME,
+          maxPrice: props?.spotMaxPrice ? parseFloat(props?.spotMaxPrice) : undefined,
+        } : undefined,
+        requireImdsv2: true,
+        securityGroup: this.securityGroups[0],
+        role: new iam.Role(this, 'Launch Template Role', {
+          assumedBy: new iam.ServicePrincipal('ec2.amazonaws.com'),
+        }),
+        userData: ec2.UserData.forOperatingSystem(image.os.is(Os.WINDOWS) ? ec2.OperatingSystemType.WINDOWS : ec2.OperatingSystemType.LINUX),
+      });
+      this.securityGroups.slice(1).map(sg => launchTemplate.connections.addSecurityGroup(sg));
 
-    this.capacityProvider = props?.capacityProvider ?? new ecs.AsgCapacityProvider(this, 'Capacity Provider', {
-      autoScalingGroup,
-      spotInstanceDraining: false, // waste of money to restart jobs as the restarted job won't have a token
-    });
+      const autoScalingGroup = new autoscaling.AutoScalingGroup(this, 'Auto Scaling Group', {
+        vpc: this.vpc,
+        launchTemplate,
+        vpcSubnets: this.subnetSelection,
+        minCapacity: props?.minInstances ?? 0,
+        maxCapacity: props?.maxInstances ?? 5,
+      });
+
+      this.capacityProvider = props?.capacityProvider ?? new ecs.AsgCapacityProvider(this, 'Capacity Provider', {
+        autoScalingGroup,
+        spotInstanceDraining: false, // waste of money to restart jobs as the restarted job won't have a token
+      });
+    }
+
+    this.capacityProvider.autoScalingGroup.addUserData(this.loginCommand(), this.pullCommand());
+    this.capacityProvider.autoScalingGroup.role.addManagedPolicy(iam.ManagedPolicy.fromAwsManagedPolicyName('AmazonSSMManagedInstanceCore'));
+    image.imageRepository.grantPull(this.capacityProvider.autoScalingGroup);
 
     this.cluster.addAsgCapacityProvider(
       this.capacityProvider,

diff --git a/test/providers.test.ts b/test/providers.test.ts
@@ -1,9 +1,8 @@
 import * as cdk from 'aws-cdk-lib';
-import {
-  aws_ec2 as ec2,
-} from 'aws-cdk-lib';
+import { aws_ec2 as ec2, aws_ecs as ecs } from 'aws-cdk-lib';
 import { Match, Template } from 'aws-cdk-lib/assertions';
-import { CodeBuildRunnerProvider, FargateRunnerProvider, LambdaRunnerProvider } from '../src';
+import * as autoscaling from 'aws-cdk-lib/aws-autoscaling';
+import { CodeBuildRunnerProvider, EcsRunnerProvider, FargateRunnerProvider, LambdaRunnerProvider } from '../src';
 
 test('CodeBuild provider', () => {
   const app = new cdk.App();
@@ -90,3 +89,60 @@ test('Fargate provider', () => {
     ],
   }));
 });
+
+describe('ECS provider', () => {
+  test('Basic', () => {
+    const app = new cdk.App();
+    const stack = new cdk.Stack(app, 'test');
+
+    const vpc = new ec2.Vpc(stack, 'vpc');
+    const sg = new ec2.SecurityGroup(stack, 'sg', { vpc });
+
+    new EcsRunnerProvider(stack, 'provider', {
+      vpc: vpc,
+      securityGroups: [sg],
+    });
+
+    const template = Template.fromStack(stack);
+
+    template.resourceCountIs('AWS::ECS::Cluster', 1);
+    template.resourceCountIs('AWS::AutoScaling::AutoScalingGroup', 1);
+
+    template.hasResourceProperties('AWS::ECS::TaskDefinition', Match.objectLike({
+      NetworkMode: 'bridge',
+      RequiresCompatibilities: ['EC2'],
+      ContainerDefinitions: [
+        {
+          Name: 'runner',
+        },
+      ],
+    }));
+  });
+
+  test('Custom capacity provider', () => {
+    const app = new cdk.App();
+    const stack = new cdk.Stack(app, 'test');
+
+    const vpc = new ec2.Vpc(stack, 'vpc');
+    const sg = new ec2.SecurityGroup(stack, 'sg', { vpc });
+
+    new EcsRunnerProvider(stack, 'provider', {
+      vpc: vpc,
+      securityGroups: [sg],
+      capacityProvider: new ecs.AsgCapacityProvider(stack, 'Capacity Provider', {
+        autoScalingGroup: new autoscaling.AutoScalingGroup(stack, 'Auto Scaling Group', {
+          vpc: vpc,
+          instanceType: ec2.InstanceType.of(ec2.InstanceClass.T3, ec2.InstanceSize.MICRO),
+          machineImage: ecs.EcsOptimizedImage.amazonLinux2(),
+          minCapacity: 1,
+          maxCapacity: 3,
+        }),
+      }),
+    });
+
+    const template = Template.fromStack(stack);
+
+    // don't create our own autoscaling group when capacity provider is specified
+    template.resourceCountIs('AWS::AutoScaling::AutoScalingGroup', 1);
+  });
+});