feat(aws-ecs): include GPU & ARM based ECS optimized AMI options

packages/@aws-cdk/aws-ecs/lib/cluster.ts

@@ -119,7 +119,7 @@ export class Cluster extends Resource implements ICluster {
     const autoScalingGroup = new autoscaling.AutoScalingGroup(this, id, {
       ...options,
       vpc: this.vpc,
-      machineImage: new EcsOptimizedAmi(),
+      machineImage: options.machineImage || new EcsOptimizedAmi(),
       updateType: options.updateType || autoscaling.UpdateType.ReplacingUpdate,
       instanceType: options.instanceType,
     });
@@ -265,6 +265,48 @@ export class EcsOptimizedAmi implements ec2.IMachineImageSource {
   }
 }
 
+/**
+ * Construct a GPU-based Linux machine image from the latest ECS Optimized AMI published in SSM
+ */
+export class EcsGpuOptimizedAmi implements ec2.IMachineImageSource {
+  private readonly amiParameterName: string = '/aws/service/ecs/optimized-ami/amazon-linux-2/gpu/recommended';
+
+  /**
+   * Return the correct image
+   */
+  public getImage(scope: Construct): ec2.MachineImage {
+    const ssmProvider = new SSMParameterProvider(scope, {
+      parameterName: this.amiParameterName
+    });
+
+    const json = ssmProvider.parameterValue("{\"image_id\": \"\"}");
+    const ami = JSON.parse(json).image_id;
+
+    return new ec2.MachineImage(ami, new ec2.LinuxOS());
+  }
+}
+
+/**
+ * Construct a ARM-based Linux machine image from the latest ECS Optimized AMI published in SSM
+ */
+export class EcsArmOptimizedAmi implements ec2.IMachineImageSource {
+  private readonly amiParameterName: string = '/aws/service/ecs/optimized-ami/amazon-linux-2/arm64/recommended';
+
+  /**
+   * Return the correct image
+   */
+  public getImage(scope: Construct): ec2.MachineImage {
+    const ssmProvider = new SSMParameterProvider(scope, {
+      parameterName: this.amiParameterName
+    });
+
+    const json = ssmProvider.parameterValue("{\"image_id\": \"\"}");
+    const ami = JSON.parse(json).image_id;
+
+    return new ec2.MachineImage(ami, new ec2.LinuxOS());
+  }
+}
+
 /**
  * An ECS cluster
  */
@@ -432,6 +474,13 @@ export interface AddAutoScalingGroupCapacityOptions {
    * @default 300
    */
   readonly taskDrainTimeSeconds?: number;
+
+  /**
+   * The machine image for the ECS instances
+   *
+   * @default - No default machine image
+   */
+  readonly machineImage?: ec2.IMachineImageSource;
 }
 
 /**