Added support for running the AWS CodeBuild Project within a VPC and expanded the list of supported ECR repos containing pre-built Docker images maintained by AWS

Author: Kopas

sagemaker_studio_image_build/builder.py

@@ -64,12 +64,13 @@ def delete_zip_file(bucket, key):
     s3.delete_object(Bucket=bucket, Key=key)
 
 
-def build_image(repository, role, bucket, compute_type, extra_args, log=True):
+def build_image(repository, role, bucket, compute_type, vpc_config, extra_args, log=True):
     bucket, key = upload_zip_file(repository, bucket, " ".join(extra_args))
     try:
         from sagemaker_studio_image_build.codebuild import TempCodeBuildProject
 
-        with TempCodeBuildProject(f"{bucket}/{key}", role, repository=repository, compute_type=compute_type) as p:
+        with TempCodeBuildProject(f"{bucket}/{key}", role, repository=repository, 
+                                    compute_type=compute_type, vpc_config=vpc_config) as p:
             p.build(log)
     finally:
         delete_zip_file(bucket, key)

sagemaker_studio_image_build/cli.py

@@ -15,6 +15,15 @@ def validate_args(args, extra_args):
                 f'Error parsing reference: "{args.repository}" is not a valid repository/tag'
             )
 
+    vpc_config = [args.vpc_id, args.subnets, args.security_groups]
+    none_arg_count = sum(arg is None for arg in [args.vpc_id, args.subnets, args.security_groups])
+
+    if none_arg_count > 0 and none_arg_count < 3:
+        raise ValueError(
+            'Invalid input of the VPC configuration. Please either provide all of the VPC arguments or none of them,'\
+            'in which case the CodeBuild Project, by default, will not run within a VPC.'
+        )
+
     # Validate extra_args
     for idx, extra_arg in enumerate(extra_args):
         # Validate that the path to the Dockerfile is within the PWD.
@@ -46,11 +55,24 @@ def get_role(args):
         )
 
 
+def construct_vpc_config(args):
+    if args.vpc_id is None:
+        return None
+    else:
+        vpc_config = {
+                        'vpcId': args.vpc_id,
+                        'subnets': args.subnets.split(','),
+                        'securityGroupIds': args.security_groups.split(',')
+                        }
+        return vpc_config
+
+
 def build_image(args, extra_args):
     validate_args(args, extra_args)
 
     builder.build_image(
-        args.repository, get_role(args), args.bucket, args.compute_type, extra_args, log=not args.no_logs
+        args.repository, get_role(args), args.bucket, args.compute_type, 
+        construct_vpc_config(args), extra_args, log=not args.no_logs
     )
 
 
@@ -74,7 +96,7 @@ def main():
         "--compute-type",
         help="The CodeBuild compute type (default: BUILD_GENERAL1_SMALL)",
         choices=["BUILD_GENERAL1_SMALL", "BUILD_GENERAL1_MEDIUM",
-                 "BUILD_GENERAL1_LARGE", "BUILD_GENERAL1_2XLARGE"]
+                 "BUILD_GENERAL1_LARGE", "BUILD_GENERAL1_2XLARGE"],
         default="BUILD_GENERAL1_SMALL"
     )
     build_parser.add_argument(
@@ -85,6 +107,18 @@ def main():
         "--bucket",
         help="The S3 bucket to use for sending data to CodeBuild (if None, use the SageMaker SDK default bucket).",
     )
+    build_parser.add_argument(
+        "--vpc-id",
+        help="The Id of the VPC that will host the CodeBuild Project (such as vpc-05c09f91d48831c8c).",
+    )
+    build_parser.add_argument(
+        "--subnets",
+        help="The comma-separated list of subnet ids for the CodeBuild Project (such as subnet-0b31f1863e9d31a67)",
+    )
+    build_parser.add_argument(
+        "--security-groups",
+        help="The comma-separated list of security group ids for the CodeBuild Project (such as sg-0ce4ec0d0414d2ddc).",
+    )
     build_parser.add_argument(
         "--no-logs",
         action="store_true",

sagemaker_studio_image_build/codebuild.py

@@ -11,14 +11,15 @@
 
 
 class TempCodeBuildProject:
-    def __init__(self, s3_location, role, repository=None, compute_type=None):
+    def __init__(self, s3_location, role, repository=None, compute_type=None, vpc_config=None):
         self.s3_location = s3_location
         self.role = role
 
         self.session = boto3.session.Session()
         self.domain_id, self.user_profile_name = self._get_studio_metadata()
         self.repo_name = None
         self.compute_type = compute_type or "BUILD_GENERAL1_SMALL"
+        self.vpc_config = vpc_config
 
         if repository:
             self.repo_name, self.tag = repository.split(":", maxsplit=1)
@@ -75,6 +76,9 @@ def __enter__(self):
             "serviceRole": f"arn:{partition}:iam::{account}:role/{self.role}",
         }
 
+        if self.vpc_config is not None:
+            args["vpcConfig"] = self.vpc_config
+
         client.create_project(**args)
         return self
 

sagemaker_studio_image_build/data/buildspec.template.yml

@@ -9,6 +9,9 @@ phases:
       - $(aws ecr get-login --no-include-email --region $AWS_DEFAULT_REGION --registry-ids 217643126080)
       - $(aws ecr get-login --no-include-email --region $AWS_DEFAULT_REGION --registry-ids 727897471807)
       - $(aws ecr get-login --no-include-email --region $AWS_DEFAULT_REGION --registry-ids 626614931356)
+      - $(aws ecr get-login --no-include-email --region $AWS_DEFAULT_REGION --registry-ids 683313688378)
+      - $(aws ecr get-login --no-include-email --region $AWS_DEFAULT_REGION --registry-ids 520713654638)
+      - $(aws ecr get-login --no-include-email --region $AWS_DEFAULT_REGION --registry-ids 462105765813)
   build:
     commands:
       - echo Build started on `date`