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,
     });
@@ -232,22 +232,36 @@ export interface EcsOptimizedAmiProps {
   /**
    * What generation of Amazon Linux to use
    *
-   * @default AmazonLinux
+   * @deprecated Use amiType instead.
    */
   readonly generation?: ec2.AmazonLinuxGeneration;
+
+  /**
+   * What ECS Optimized AMI type to use
+   *
+   * @default is Amazon Linux
+   */
+  readonly amiType?: EcsOptimizedAmiType;
 }
 
 /**
  * Construct a Linux machine image from the latest ECS Optimized AMI published in SSM
  */
 export class EcsOptimizedAmi implements ec2.IMachineImageSource {
   private readonly generation: ec2.AmazonLinuxGeneration;
+  private readonly amiType: EcsOptimizedAmiType;
+
   private readonly amiParameterName: string;
 
   constructor(props?: EcsOptimizedAmiProps) {
     this.generation = (props && props.generation) || ec2.AmazonLinuxGeneration.AmazonLinux;
-    if (this.generation === ec2.AmazonLinuxGeneration.AmazonLinux2) {
+    this.amiType = (props && props.amiType) || EcsOptimizedAmiType.AmazonLinux;
+    if (this.generation === ec2.AmazonLinuxGeneration.AmazonLinux2 || this.amiType === EcsOptimizedAmiType.AmazonLinux2) {
       this.amiParameterName = "/aws/service/ecs/optimized-ami/amazon-linux-2/recommended";
+    } else if (this.amiType === EcsOptimizedAmiType.Gpu) {
+      this.amiParameterName = '/aws/service/ecs/optimized-ami/amazon-linux-2/gpu/recommended';
+    } else if (this.amiType === EcsOptimizedAmiType.Arm) {
+      this.amiParameterName = '/aws/service/ecs/optimized-ami/amazon-linux-2/arm64/recommended';
     } else {
       this.amiParameterName = "/aws/service/ecs/optimized-ami/amazon-linux/recommended";
     }
@@ -443,6 +457,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;
 }
 
 /**
@@ -490,3 +511,29 @@ export enum NamespaceType {
    */
   PublicDns = 'PublicDns',
 }
+
+/**
+ * The type of ECS OptimizedAMI to create
+ */
+export enum EcsOptimizedAmiType {
+
+  /**
+   * Create an Amazon Linux optimized AMI
+   */
+  AmazonLinux = 'AmazonLinux',
+
+  /**
+   * Create an Amazon Linux 2 optimized AMI
+   */
+  AmazonLinux2 = 'AmazonLinux2',
+
+  /**
+   * Create a GPU optimized AMI
+   */
+  Gpu = 'Gpu',
+
+  /**
+   * Create a ARM64 optimized AMI
+   */
+  Arm = 'Arm',
+}