Merge pull request #17 from ikopas/feature/add-vpc-config

Feature/add vpc config

Author: jaipreet-s

.gitignore

@@ -11,3 +11,4 @@ manual_dist/**
 /build
 .idea/
 .vscode/
+.DS_Store

README.md

@@ -20,12 +20,12 @@ Any additional arguments supported with `docker build` are supported
 sm-docker build . --file /path/to/Dockerfile --build-arg foo=bar
 ```
 
-By default, the image will be pushed to a repository `sagemakerstudio` with the tag `latest`, and use the Studio App's execution role and the default SageMaker Python SDK S3 bucket
+By default, the CodeBuild project will not run within a VPC, the image will be pushed to a repository `sagemakerstudio` with the tag `latest`, and use the Studio App's execution role and the default SageMaker Python SDK S3 bucket
 
 These can be overridden with the relevant CLI options.
 
 ```bash
-sm-docker build . --repository mynewrepo:1.0 --role MyRoleName 
+sm-docker build . --repository mynewrepo:1.0 --role SampleDockerBuildRole --bucket sagemaker-us-east-1-326543455535 --vpc-id vpc-0c70e76ef1c603b94 --subnet-ids subnet-0d984f080338960bb,subnet-0ac3e96808c8092f2 --security-group-ids sg-0d31b4042f2902cd0
 ``` 
 
 The CLI will take care of packaging the current directory and uploading to S3, creating a CodeBuild project, starting a build with the S3 artifacts, tailing the build logs, and uploading the built image to ECR.
@@ -101,16 +101,38 @@ The following permissions are required in the execution role to execute a build
                 "ecr:DescribeRepositories",
                 "ecr:UploadLayerPart",
                 "ecr:ListImages",
-                "ecr:InitiateLayerUpload",
+                "ecr:InitiateLayerUpload", 
                 "ecr:BatchCheckLayerAvailability",
                 "ecr:PutImage"
             ],
             "Resource": "arn:aws:ecr:*:*:repository/sagemaker-studio*"
         },
         {
+            "Sid": "ReadAccessToPrebuiltAwsImages",
             "Effect": "Allow",
-            "Action": "ecr:GetAuthorizationToken",
-            "Resource": "*"
+            "Action": [
+                "ecr:BatchGetImage",
+                "ecr:GetDownloadUrlForLayer"
+            ],
+            "Resource": [
+                "arn:aws:ecr:*:763104351884:repository/*",
+                "arn:aws:ecr:*:217643126080:repository/*",
+                "arn:aws:ecr:*:727897471807:repository/*",
+                "arn:aws:ecr:*:626614931356:repository/*",
+                "arn:aws:ecr:*:683313688378:repository/*",
+                "arn:aws:ecr:*:520713654638:repository/*",
+                "arn:aws:ecr:*:462105765813:repository/*"
+            ]
+        },
+        {
+            "Sid": "EcrAuthorizationTokenRetrieval",
+            "Effect": "Allow",
+            "Action": [
+                "ecr:GetAuthorizationToken"
+            ],
+            "Resource": [
+                "*"
+            ]
         },
         {
             "Effect": "Allow",
@@ -151,6 +173,26 @@ The following permissions are required in the execution role to execute a build
 
 ```
 
+If you need to run your CodeBuild project within a VPC, please add the following actions to your execution role that the CodeBuild Project will assume:
+
+```json
+        {
+            "Sid": "VpcAccessActions",
+            "Effect": "Allow",
+            "Action": [
+                "ec2:CreateNetworkInterface",
+                "ec2:CreateNetworkInterfacePermission",
+                "ec2:DescribeDhcpOptions",
+                "ec2:DescribeNetworkInterfaces",
+                "ec2:DeleteNetworkInterface",
+                "ec2:DescribeSubnets",
+                "ec2:DescribeSecurityGroups",
+                "ec2:DescribeVpcs"
+            ],
+            "Resource": "*"
+        }
+```
+
 ### Development
 
 Checkout the repository.

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.subnet_ids, args.security_group_ids]
+    none_arg_count = sum(arg is None for arg in [args.vpc_id, args.subnet_ids, args.security_group_ids])
+
+    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.subnet_ids.split(','),
+                        'securityGroupIds': args.security_group_ids.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(
+        "--subnet-ids",
+        help="The comma-separated list of subnet ids for the CodeBuild Project (such as subnet-0b31f1863e9d31a67)",
+    )
+    build_parser.add_argument(
+        "--security-group-ids",
+        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`