From 2d74a4695992b21d1adc2ccbe6874e1128e996db Mon Sep 17 00:00:00 2001
From: Luca Pizzini <lpizzini7@gmail.com>
Date: Mon, 12 Jun 2023 15:52:23 +0200
Subject: [PATCH] fix(codebuild): add possibility to specify
 `BUILD_GENERAL1_SMALL` compute type with Linux GPU build image (#25880)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This fix allows specifying the `BUILD_GENERAL1_SMALL` compute type when using a Linux GPU to build an image as defined in the [docs](https://docs.aws.amazon.com/codebuild/latest/userguide/build-env-ref-compute-types.html).

Closes #25857.

----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
---
 .../lib/linux-gpu-build-image.ts              |  5 +--
 .../test/linux-gpu-build-image.test.ts        | 33 +++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/packages/aws-cdk-lib/aws-codebuild/lib/linux-gpu-build-image.ts b/packages/aws-cdk-lib/aws-codebuild/lib/linux-gpu-build-image.ts
index a31678a5a6567..218703089b643 100644
--- a/packages/aws-cdk-lib/aws-codebuild/lib/linux-gpu-build-image.ts
+++ b/packages/aws-cdk-lib/aws-codebuild/lib/linux-gpu-build-image.ts
@@ -142,8 +142,9 @@ export class LinuxGpuBuildImage implements IBindableBuildImage {
   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 !== ComputeType.LARGE &&
+        buildEnvironment.computeType !== ComputeType.SMALL) {
+      ret.push(`GPU images only support ComputeType '${ComputeType.LARGE}' and '${ComputeType.SMALL}' - ` +
         `'${buildEnvironment.computeType}' was given`);
     }
     return ret;
diff --git a/packages/aws-cdk-lib/aws-codebuild/test/linux-gpu-build-image.test.ts b/packages/aws-cdk-lib/aws-codebuild/test/linux-gpu-build-image.test.ts
index f1fdada6c68e5..e090c778f9bae 100644
--- a/packages/aws-cdk-lib/aws-codebuild/test/linux-gpu-build-image.test.ts
+++ b/packages/aws-cdk-lib/aws-codebuild/test/linux-gpu-build-image.test.ts
@@ -57,6 +57,39 @@ describe('Linux GPU build image', () => {
         },
       });
     });
+
+    test('allows to specify a BUILD_GENERAL1_SMALL build environment', () => {
+      const stack = new cdk.Stack();
+
+      new codebuild.Project(stack, 'Project', {
+        buildSpec: codebuild.BuildSpec.fromObject({
+          version: '0.2',
+          phases: {
+            build: { commands: ['ls'] },
+          },
+        }),
+        environment: {
+          buildImage: codebuild.LinuxGpuBuildImage.awsDeepLearningContainersImage(
+            'my-repo', 'my-tag', '123456789012'),
+          computeType: codebuild.ComputeType.SMALL,
+        },
+      });
+
+      Template.fromStack(stack).hasResourceProperties('AWS::CodeBuild::Project', {
+        Environment: {
+          ComputeType: 'BUILD_GENERAL1_SMALL',
+          Image: {
+            'Fn::Join': ['', [
+              '123456789012.dkr.ecr.',
+              { Ref: 'AWS::Region' },
+              '.',
+              { Ref: 'AWS::URLSuffix' },
+              '/my-repo:my-tag',
+            ]],
+          },
+        },
+      });
+    });
   });
 
   describe('ECR Repository', () => {