feat(codebuild): add support for GPU build images

packages/@aws-cdk/aws-codebuild/README.md

@@ -238,6 +238,40 @@ The following example shows how to define an image from a private docker registr
 
 [Docker Registry example](./test/integ.docker-registry.lit.ts)
 
+### GPU images
+
+The class `LinuxGpuBuildImage` contains constants for working with
+[AWS Deep Learning Container images](https://aws.amazon.com/releasenotes/available-deep-learning-containers-images):
+
+
+```typescript
+new codebuild.Project(this, 'Project', {
+  environment: {
+    buildImage: codebuild.LinuxGpuBuildImage.DLC_TENSORFLOW_2_1_0_INFERENCE,
+  },
+  ...
+})
+```
+
+One complication is that the repositories for the DLC images are in
+different accounts in different AWS regions.
+In most cases, the CDK will handle providing the correct account for you;
+in rare cases (for example, deploying to new regions)
+where our information might be out of date,
+you can always specify the account
+(along with the repository name and tag)
+explicitly using the `awsDeepLearningContainersImage` method:
+
+```typescript
+new codebuild.Project(this, 'Project', {
+  environment: {
+    buildImage: codebuild.LinuxGpuBuildImage.awsDeepLearningContainersImage(
+      'tensorflow-inference', '2.1.0-gpu-py36-cu101-ubuntu18.04', '123456789012'),
+  },
+  ...
+})
+```
+
 ## Credentials
 
 CodeBuild allows you to store credentials used when communicating with various sources,

packages/@aws-cdk/aws-codebuild/lib/index.ts

@@ -8,6 +8,7 @@ export * from './artifacts';
 export * from './cache';
 export * from './build-spec';
 export * from './file-location';
+export * from './linux-gpu-build-image';
 
 // AWS::CodeBuild CloudFormation Resources:
 export * from './codebuild.generated';

packages/@aws-cdk/aws-codebuild/lib/linux-gpu-build-image.ts

@@ -0,0 +1,141 @@
+import * as ecr from '@aws-cdk/aws-ecr';
+import * as core from '@aws-cdk/core';
+import { FactName, RegionInfo } from '@aws-cdk/region-info';
+import { BuildSpec } from './build-spec';
+import { runScriptLinuxBuildSpec } from './private/run-script-linux-build-spec';
+import {
+  BuildEnvironment, BuildImageBindOptions, BuildImageConfig, ComputeType, IBindableBuildImage, IBuildImage,
+  ImagePullPrincipalType, IProject,
+} from './project';
+
+const mappingName = 'AwsDeepLearningContainersRepositoriesAccounts';
+
+/**
+ * A CodeBuild GPU image running Linux.
+ *
+ * This class has public constants that represent the most popular GPU images from AWS Deep Learning Containers.
+ *
+ * @see https://aws.amazon.com/releasenotes/available-deep-learning-containers-images
+ */
+export class LinuxGpuBuildImage implements IBindableBuildImage {
+  /** Tensorflow 1.14.0 GPU image from AWS Deep Learning Containers. */
+  public static readonly DLC_TENSORFLOW_1_14_0 = LinuxGpuBuildImage.awsDeepLearningContainersImage('tensorflow-training',
+    '1.14.0-gpu-py36-cu100-ubuntu16.04');
+  /** Tensorflow 1.15.0 GPU image from AWS Deep Learning Containers. */
+  public static readonly DLC_TENSORFLOW_1_15_0 = LinuxGpuBuildImage.awsDeepLearningContainersImage('tensorflow-training',
+    '1.15.0-gpu-py36-cu100-ubuntu18.04');
+  /** Tensorflow 1.15.2 GPU training image from AWS Deep Learning Containers. */
+  public static readonly DLC_TENSORFLOW_1_15_2_TRAINING = LinuxGpuBuildImage.awsDeepLearningContainersImage('tensorflow-training',
+    '1.15.2-gpu-py37-cu100-ubuntu18.04');
+  /** Tensorflow 1.15.2 GPU inference image from AWS Deep Learning Containers. */
+  public static readonly DLC_TENSORFLOW_1_15_2_INFERENCE = LinuxGpuBuildImage.awsDeepLearningContainersImage('tensorflow-inference',
+    '1.15.2-gpu-py36-cu100-ubuntu18.04');
+  /** Tensorflow 2.0.0 GPU image from AWS Deep Learning Containers. */
+  public static readonly DLC_TENSORFLOW_2_0_0 = LinuxGpuBuildImage.awsDeepLearningContainersImage('tensorflow-training',
+    '2.0.0-gpu-py36-cu100-ubuntu18.04');
+  /** Tensorflow 2.0.1 GPU image from AWS Deep Learning Containers. */
+  public static readonly DLC_TENSORFLOW_2_0_1 = LinuxGpuBuildImage.awsDeepLearningContainersImage('tensorflow-training',
+    '2.0.1-gpu-py36-cu100-ubuntu18.04');
+  /** Tensorflow 2.1.0 GPU training image from AWS Deep Learning Containers. */
+  public static readonly DLC_TENSORFLOW_2_1_0_TRAINING = LinuxGpuBuildImage.awsDeepLearningContainersImage('tensorflow-training',
+    '2.1.0-gpu-py36-cu101-ubuntu18.04');
+  /** Tensorflow 2.1.0 GPU inference image from AWS Deep Learning Containers. */
+  public static readonly DLC_TENSORFLOW_2_1_0_INFERENCE = LinuxGpuBuildImage.awsDeepLearningContainersImage('tensorflow-inference',
+    '2.1.0-gpu-py36-cu101-ubuntu18.04');
+  /** Tensorflow 2.2.0 GPU training image from AWS Deep Learning Containers. */
+  public static readonly DLC_TENSORFLOW_2_2_0_TRAINING = LinuxGpuBuildImage.awsDeepLearningContainersImage('tensorflow-training',
+    '2.2.0-gpu-py37-cu101-ubuntu18.04');
+
+  /** PyTorch 1.2.0 GPU image from AWS Deep Learning Containers. */
+  public static readonly DLC_PYTORCH_1_2_0 = LinuxGpuBuildImage.awsDeepLearningContainersImage('pytorch-training',
+    '1.2.0-gpu-py36-cu100-ubuntu16.04');
+  /** PyTorch 1.3.1 GPU image from AWS Deep Learning Containers. */
+  public static readonly DLC_PYTORCH_1_3_1 = LinuxGpuBuildImage.awsDeepLearningContainersImage('pytorch-training',
+    '1.3.1-gpu-py36-cu101-ubuntu16.04');
+  /** PyTorch 1.4.0 GPU training image from AWS Deep Learning Containers. */
+  public static readonly DLC_PYTORCH_1_4_0_TRAINING = LinuxGpuBuildImage.awsDeepLearningContainersImage('pytorch-training',
+    '1.4.0-gpu-py36-cu101-ubuntu16.04');
+  /** PyTorch 1.4.0 GPU inference image from AWS Deep Learning Containers. */
+  public static readonly DLC_PYTORCH_1_4_0_INFERENCE = LinuxGpuBuildImage.awsDeepLearningContainersImage('pytorch-inference',
+    '1.4.0-gpu-py36-cu101-ubuntu16.04');
+  /** PyTorch 1.5.0 GPU training image from AWS Deep Learning Containers. */
+  public static readonly DLC_PYTORCH_1_5_0_TRAINING = LinuxGpuBuildImage.awsDeepLearningContainersImage('pytorch-training',
+    '1.5.0-gpu-py36-cu101-ubuntu16.04');
+  /** PyTorch 1.5.0 GPU inference image from AWS Deep Learning Containers. */
+  public static readonly DLC_PYTORCH_1_5_0_INFERENCE = LinuxGpuBuildImage.awsDeepLearningContainersImage('pytorch-inference',
+    '1.5.0-gpu-py36-cu101-ubuntu16.04');
+
+  /** MXNet 1.4.1 GPU image from AWS Deep Learning Containers. */
+  public static readonly DLC_MXNET_1_4_1 = LinuxGpuBuildImage.awsDeepLearningContainersImage('mxnet-training',
+    '1.4.1-gpu-py36-cu100-ubuntu16.04');
+  /** MXNet 1.6.0 GPU image from AWS Deep Learning Containers. */
+  public static readonly DLC_MXNET_1_6_0 = LinuxGpuBuildImage.awsDeepLearningContainersImage('mxnet-training',
+    '1.6.0-gpu-py36-cu101-ubuntu16.04');
+
+  /**
+   * Returns a Linux GPU build image from AWS Deep Learning Containers.
+   *
+   * @param repositoryName the name of the repository,
+   *   for example "pytorch-inference"
+   * @param tag the tag of the image, for example "1.5.0-gpu-py36-cu101-ubuntu16.04"
+   * @param account the AWS account ID where the DLC repository for this region is hosted in.
+   *   In many cases, the CDK can infer that for you, but for some newer region our information
+   *   might be out of date; in that case, you can specify the region explicitly using this optional parameter
+   * @see https://aws.amazon.com/releasenotes/available-deep-learning-containers-images
+   */
+  public static awsDeepLearningContainersImage(repositoryName: string, tag: string, account?: string): IBuildImage {
+    return new LinuxGpuBuildImage(repositoryName, tag, account);
+  }
+
+  public readonly type = 'LINUX_GPU_CONTAINER';
+  public readonly defaultComputeType = ComputeType.LARGE;
+  public readonly imageId: string;
+  public readonly imagePullPrincipalType?: ImagePullPrincipalType = ImagePullPrincipalType.SERVICE_ROLE;
+
+  private readonly accountExpression: string;
+
+  private constructor(private readonly repositoryName: string, tag: string, private readonly account: string | undefined) {
+    this.accountExpression = account ?? core.Fn.findInMap(mappingName, core.Aws.REGION, 'account');
+    this.imageId = `${this.accountExpression}.dkr.ecr.${core.Aws.REGION}.${core.Aws.URL_SUFFIX}/${repositoryName}:${tag}`;
+  }
+
+  public bind(scope: core.Construct, project: IProject, _options: BuildImageBindOptions): BuildImageConfig {
+    if (!this.account) {
+      const scopeStack = core.Stack.of(scope);
+      // Unfortunately, the account IDs of the DLC repositories are not the same in all regions.
+      // Because of that, use a (singleton) Mapping to find the correct account
+      if (!scopeStack.node.tryFindChild(mappingName)) {
+        const mapping: { [k1: string]: { [k2: string]: any } } = {};
+        // get the accounts from the region-info module
+        const region2Accounts = RegionInfo.regionMap(FactName.DLC_REPOSITORY_ACCOUNT);
+        for (const [region, account] of Object.entries(region2Accounts)) {
+          mapping[region] = { account };
+        }
+        new core.CfnMapping(scopeStack, mappingName, { mapping });
+      }
+    }
+
+    const repository = ecr.Repository.fromRepositoryAttributes(scope, 'AwsDlcRepositoryCodeBuild', {
+      repositoryName: this.repositoryName,
+      repositoryArn: ecr.Repository.arnForLocalRepository(this.repositoryName, scope, this.accountExpression),
+    });
+    repository.grantPull(project);
+
+    return {
+    };
+  }
+
+  public validate(buildEnvironment: BuildEnvironment): string[] {
+    const ret = [];
+    if (buildEnvironment.computeType &&
+        buildEnvironment.computeType !== ComputeType.LARGE) {
+      ret.push(`GPU images only support ComputeType '${ComputeType.LARGE}' - ` +
+        `'${buildEnvironment.computeType}' was given`);
+    }
+    return ret;
+  }
+
+  public runScriptBuildspec(entrypoint: string): BuildSpec {
+    return runScriptLinuxBuildSpec(entrypoint);
+  }
+}

packages/@aws-cdk/aws-codebuild/lib/private/run-script-linux-build-spec.ts

@@ -0,0 +1,32 @@
+import { BuildSpec } from '../build-spec';
+
+export const S3_BUCKET_ENV = 'SCRIPT_S3_BUCKET';
+export const S3_KEY_ENV = 'SCRIPT_S3_KEY';
+
+export function runScriptLinuxBuildSpec(entrypoint: string) {
+  return BuildSpec.fromObject({
+    version: '0.2',
+    phases: {
+      pre_build: {
+        commands: [
+          // Better echo the location here; if this fails, the error message only contains
+          // the unexpanded variables by default. It might fail if you're running an old
+          // definition of the CodeBuild project--the permissions will have been changed
+          // to only allow downloading the very latest version.
+          `echo "Downloading scripts from s3://\${${S3_BUCKET_ENV}}/\${${S3_KEY_ENV}}"`,
+          `aws s3 cp s3://\${${S3_BUCKET_ENV}}/\${${S3_KEY_ENV}} /tmp`,
+          'mkdir -p /tmp/scriptdir',
+          `unzip /tmp/$(basename \$${S3_KEY_ENV}) -d /tmp/scriptdir`,
+        ],
+      },
+      build: {
+        commands: [
+          'export SCRIPT_DIR=/tmp/scriptdir',
+          `echo "Running ${entrypoint}"`,
+          `chmod +x /tmp/scriptdir/${entrypoint}`,
+          `/tmp/scriptdir/${entrypoint}`,
+        ],
+      },
+    },
+  });
+}

packages/@aws-cdk/aws-codebuild/lib/project.ts

@@ -16,13 +16,11 @@ import { CodePipelineArtifacts } from './codepipeline-artifacts';
 import { IFileSystemLocation } from './file-location';
 import { NoArtifacts } from './no-artifacts';
 import { NoSource } from './no-source';
+import { runScriptLinuxBuildSpec, S3_BUCKET_ENV, S3_KEY_ENV } from './private/run-script-linux-build-spec';
 import { renderReportGroupArn } from './report-group-utils';
 import { ISource } from './source';
 import { CODEPIPELINE_SOURCE_ARTIFACTS_TYPE, NO_SOURCE_TYPE } from './source-types';
 
-const S3_BUCKET_ENV = 'SCRIPT_S3_BUCKET';
-const S3_KEY_ENV = 'SCRIPT_S3_KEY';
-
 export interface IProject extends IResource, iam.IGrantable, ec2.IConnectable {
   /**
    * The ARN of this Project.
@@ -796,6 +794,12 @@ export class Project extends ProjectBase {
     if (props.encryptionKey) {
       this.encryptionKey = props.encryptionKey;
     }
+
+    // bind
+    const bindFunction = (this.buildImage as any).bind;
+    if (bindFunction) {
+      bindFunction.call(this.buildImage, this, this, {});
+    }
   }
 
   /**
@@ -1196,6 +1200,21 @@ export interface IBuildImage {
   runScriptBuildspec(entrypoint: string): BuildSpec;
 }
 
+/** Optional arguments to {@link IBuildImage.binder} - currently empty. */
+export interface BuildImageBindOptions {}
+
+/** The return type from {@link IBuildImage.binder} - currently empty. */
+export interface BuildImageConfig {}
+
+// @deprecated(not in tsdoc on purpose): add bind() to IBuildImage
+// and get rid of IBindableBuildImage
+
+/** A variant of {@link IBuildImage} that allows binding to the project. */
+export interface IBindableBuildImage extends IBuildImage {
+  /** Function that allows the build image access to the construct tree. */
+  bind(scope: Construct, project: IProject, options: BuildImageBindOptions): BuildImageConfig;
+}
+
 class ArmBuildImage implements IBuildImage {
   public readonly type = 'ARM_CONTAINER';
   public readonly defaultComputeType = ComputeType.LARGE;
@@ -1423,34 +1442,6 @@ export class LinuxBuildImage implements IBuildImage {
   }
 }
 
-function runScriptLinuxBuildSpec(entrypoint: string) {
-  return BuildSpec.fromObject({
-    version: '0.2',
-    phases: {
-      pre_build: {
-        commands: [
-          // Better echo the location here; if this fails, the error message only contains
-          // the unexpanded variables by default. It might fail if you're running an old
-          // definition of the CodeBuild project--the permissions will have been changed
-          // to only allow downloading the very latest version.
-          `echo "Downloading scripts from s3://\${${S3_BUCKET_ENV}}/\${${S3_KEY_ENV}}"`,
-          `aws s3 cp s3://\${${S3_BUCKET_ENV}}/\${${S3_KEY_ENV}} /tmp`,
-          'mkdir -p /tmp/scriptdir',
-          `unzip /tmp/$(basename \$${S3_KEY_ENV}) -d /tmp/scriptdir`,
-        ],
-      },
-      build: {
-        commands: [
-          'export SCRIPT_DIR=/tmp/scriptdir',
-          `echo "Running ${entrypoint}"`,
-          `chmod +x /tmp/scriptdir/${entrypoint}`,
-          `/tmp/scriptdir/${entrypoint}`,
-        ],
-      },
-    },
-  });
-}
-
 /**
  * Construction properties of {@link WindowsBuildImage}.
  * Module-private, as the constructor of {@link WindowsBuildImage} is private.

packages/@aws-cdk/aws-codebuild/package.json

@@ -91,6 +91,7 @@
     "@aws-cdk/aws-s3-assets": "0.0.0",
     "@aws-cdk/aws-secretsmanager": "0.0.0",
     "@aws-cdk/core": "0.0.0",
+    "@aws-cdk/region-info": "0.0.0",
     "constructs": "^3.0.2"
   },
   "homepage": "https://github.com/aws/aws-cdk",
@@ -108,6 +109,7 @@
     "@aws-cdk/aws-s3-assets": "0.0.0",
     "@aws-cdk/aws-secretsmanager": "0.0.0",
     "@aws-cdk/core": "0.0.0",
+    "@aws-cdk/region-info": "0.0.0",
     "constructs": "^3.0.2"
   },
   "engines": {