feat(ecs): add GPU support in container definition (#3044)

* feat(ecs) add GPU support in container definition * feat(ecs): add GPU support in container definition * fix(ecs): removed exta space in test file
aws · Aug 22, 2019 · 2590327 · 2590327
1 parent 6dcae29
commit 2590327
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 1 deletion.
diff --git a/packages/@aws-cdk/aws-ecs/lib/container-definition.ts b/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;
 }
 
 /**
@@ -519,6 +526,7 @@ export class ContainerDefinition extends cdk.Construct {
       healthCheck: this.props.healthCheck && renderHealthCheck(this.props.healthCheck),
       links: cdk.Lazy.listValue({ produce: () => this.links }, { omitEmpty: true }),
       linuxParameters: this.linuxParameters && this.linuxParameters.renderLinuxParameters(),
+      resourceRequirements: (this.props.gpuCount !== undefined) ? renderResourceRequirements(this.props.gpuCount) : undefined,
     };
   }
 }
@@ -612,6 +620,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.
  *

diff --git a/packages/@aws-cdk/aws-ecs/test/test.container-definition.ts b/packages/@aws-cdk/aws-ecs/test/test.container-definition.ts
@@ -486,6 +486,38 @@ export = {
 
   },
 
+  '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
     const stack = new cdk.Stack();
@@ -1090,6 +1122,5 @@ export = {
       test.done();
     }
   },
-
   // render extra hosts test
 };