feat(ecs): add GPU support in container definition

packages/@aws-cdk/aws-ecs/lib/container-definition.ts

@@ -233,6 +233,13 @@ export interface ContainerDefinitionOptions {
    * @default - No Linux paramters.
    */
   readonly linuxParameters?: LinuxParameters;
+
+  /**
+   * The number of GPUs assigned to the container.
+   *
+   * @default - No GPUs assigned.
+   */
+  readonly gpuCount?: number;
 }
 
 /**
@@ -514,6 +521,7 @@ export class ContainerDefinition extends cdk.Construct {
       healthCheck: this.props.healthCheck && renderHealthCheck(this.props.healthCheck),
       links: cdk.Lazy.listValue({ produce: () => renderArray(this.links, l => l) }),
       linuxParameters: this.linuxParameters && this.linuxParameters.renderLinuxParameters(),
+      resourceRequirements: (this.props.gpuCount !== undefined) ? renderResourceRequirements(this.props.gpuCount) : undefined,
     };
   }
 }
@@ -614,6 +622,14 @@ function getHealthCheckCommand(hc: HealthCheck): string[] {
   return hcCommand.concat(cmd);
 }
 
+function renderResourceRequirements(gpuCount: number): CfnTaskDefinition.ResourceRequirementProperty[] | undefined {
+  if (gpuCount === 0) { return undefined; }
+  return [{
+    type: 'GPU',
+    value: gpuCount.toString(),
+  }];
+}
+
 /**
  * The ulimit settings to pass to the container.
  *

packages/@aws-cdk/aws-ecs/test/test.container-definition.ts

@@ -285,6 +285,38 @@ export = {
     test.done();
 
   },
+
+  'Given GPU count parameter': {
+    'will add resource requirements to container definition'(test: Test) {
+      // GIVEN
+      const stack = new cdk.Stack();
+      const taskDefinition = new ecs.Ec2TaskDefinition(stack, 'TaskDef');
+
+      // WHEN
+      taskDefinition.addContainer('cont', {
+        image: ecs.ContainerImage.fromRegistry('test'),
+        memoryLimitMiB: 1024,
+        gpuCount: 4,
+      });
+
+      // THEN
+      expect(stack).to(haveResourceLike('AWS::ECS::TaskDefinition', {
+        ContainerDefinitions: [
+          {
+            Image: 'test',
+            ResourceRequirements: [
+              {
+                Type: "GPU",
+                Value: "4"
+              }
+            ]
+          }
+        ]
+      }));
+
+      test.done();
+    },
+  },
 
   'can add secret environment variables to the container definition'(test: Test) {
     // GIVEN
@@ -706,6 +738,5 @@ export = {
       test.done();
     },
   },
-
   // render extra hosts test
 };