From 634357f5ed559430b4a3ce5b14ca201652a864e2 Mon Sep 17 00:00:00 2001 From: johnbensnyder Date: Mon, 5 Apr 2021 00:54:20 +0000 Subject: [PATCH 1/3] added GPU environment support --- sagemaker_studio_image_build/builder.py | 8 ++++++-- sagemaker_studio_image_build/cli.py | 12 +++++++++--- sagemaker_studio_image_build/codebuild.py | 7 +++++-- 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/sagemaker_studio_image_build/builder.py b/sagemaker_studio_image_build/builder.py index 2733d92..9423022 100644 --- a/sagemaker_studio_image_build/builder.py +++ b/sagemaker_studio_image_build/builder.py @@ -64,12 +64,16 @@ 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, environment, 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, + environment=environment) as p: p.build(log) finally: delete_zip_file(bucket, key) diff --git a/sagemaker_studio_image_build/cli.py b/sagemaker_studio_image_build/cli.py index 9ef02b2..9c66536 100644 --- a/sagemaker_studio_image_build/cli.py +++ b/sagemaker_studio_image_build/cli.py @@ -50,7 +50,7 @@ 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, args.environment, extra_args, log=not args.no_logs ) @@ -72,11 +72,17 @@ def main(): ) build_parser.add_argument( "--compute-type", - help="The CodeBuild compute type (default: BUILD_GENERAL1_SMALL)", + help="The CodeBuild compute type (default: BUILD_GENERAL1_SMALL for CPU, BUILD_GENERAL1_LARGE for GPU)", 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( + "--environment", + help="The CodeBuild environment (default: LINUX_CONTAINER)", + choices=["LINUX_CONTAINER", "LINUX_GPU_CONTAINER"], + default="LINUX_CONTAINER" + ) build_parser.add_argument( "--role", help=f"The IAM role name for CodeBuild to use (default: the Studio execution role).", diff --git a/sagemaker_studio_image_build/codebuild.py b/sagemaker_studio_image_build/codebuild.py index 583859c..e9e867d 100644 --- a/sagemaker_studio_image_build/codebuild.py +++ b/sagemaker_studio_image_build/codebuild.py @@ -11,7 +11,7 @@ class TempCodeBuildProject: - def __init__(self, s3_location, role, repository=None, compute_type=None): + def __init__(self, s3_location, role, repository=None, compute_type=None, environment=None): self.s3_location = s3_location self.role = role @@ -19,6 +19,9 @@ def __init__(self, s3_location, role, repository=None, compute_type=None): 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.environment = environment or "LINUX_CONTAINER" + if self.environment=="LINUX_GPU_CONTAINER" and self.compute_type=="BUILD_GENERAL1_SMALL": + self.compute_type="BUILD_GENERAL1_LARGE" if repository: self.repo_name, self.tag = repository.split(":", maxsplit=1) @@ -61,7 +64,7 @@ def __enter__(self): "source": {"type": "S3", "location": self.s3_location}, "artifacts": {"type": "NO_ARTIFACTS"}, "environment": { - "type": "LINUX_CONTAINER", + "type": self.environment, "image": "aws/codebuild/standard:4.0", "computeType": self.compute_type, "environmentVariables": [ From 5ce3bf5b59aa629ad40f403faee6fa877169c400 Mon Sep 17 00:00:00 2001 From: johnbensnyder Date: Sun, 29 Aug 2021 23:36:19 +0000 Subject: [PATCH 2/3] removed automatic compute-type change for GPU --- sagemaker_studio_image_build/cli.py | 8 ++++---- sagemaker_studio_image_build/codebuild.py | 5 +++-- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/sagemaker_studio_image_build/cli.py b/sagemaker_studio_image_build/cli.py index 9c66536..b60ffc5 100644 --- a/sagemaker_studio_image_build/cli.py +++ b/sagemaker_studio_image_build/cli.py @@ -64,22 +64,22 @@ def main(): build_parser = subparsers.add_parser( "build", - help="Use AWS CodeBuild to build a Docker image and push to Amazon ECR", + help="Use AWS CodeBuild to build a Docker image and push to Amazon ECR.", ) build_parser.add_argument( "--repository", - help="The ECR repository:tag for the image (default: sagemaker-studio-${domain_id}:latest)", + help="The ECR repository:tag for the image (default: sagemaker-studio-${domain_id}:latest).", ) build_parser.add_argument( "--compute-type", - help="The CodeBuild compute type (default: BUILD_GENERAL1_SMALL for CPU, BUILD_GENERAL1_LARGE for GPU)", + help="The CodeBuild compute type (default: BUILD_GENERAL1_SMALL) set to BUILD_GENERAL1_LARGE for LINUX_GPU_CONTAINER environment.", choices=["BUILD_GENERAL1_SMALL", "BUILD_GENERAL1_MEDIUM", "BUILD_GENERAL1_LARGE", "BUILD_GENERAL1_2XLARGE"], default="BUILD_GENERAL1_SMALL" ) build_parser.add_argument( "--environment", - help="The CodeBuild environment (default: LINUX_CONTAINER)", + help="The CodeBuild environment (default: LINUX_CONTAINER).", choices=["LINUX_CONTAINER", "LINUX_GPU_CONTAINER"], default="LINUX_CONTAINER" ) diff --git a/sagemaker_studio_image_build/codebuild.py b/sagemaker_studio_image_build/codebuild.py index e9e867d..86ecf85 100644 --- a/sagemaker_studio_image_build/codebuild.py +++ b/sagemaker_studio_image_build/codebuild.py @@ -20,8 +20,9 @@ def __init__(self, s3_location, role, repository=None, compute_type=None, enviro self.repo_name = None self.compute_type = compute_type or "BUILD_GENERAL1_SMALL" self.environment = environment or "LINUX_CONTAINER" - if self.environment=="LINUX_GPU_CONTAINER" and self.compute_type=="BUILD_GENERAL1_SMALL": - self.compute_type="BUILD_GENERAL1_LARGE" + if self.environment=="LINUX_GPU_CONTAINER": + assert self.compute_type=="BUILD_GENERAL1_LARGE", \ + "LINUX_GPU_CONTAINER builds only available on BUILD_GENERAL1_LARGE. Please set `--compute-type BUILD_GENERAL1_LARGE`" if repository: self.repo_name, self.tag = repository.split(":", maxsplit=1) From c0b3c157cf55208f0ca8862e6b209ed965888275 Mon Sep 17 00:00:00 2001 From: johnbensnyder Date: Sun, 29 Aug 2021 23:41:36 +0000 Subject: [PATCH 3/3] removed unnecessary check for environment and compute type --- sagemaker_studio_image_build/codebuild.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sagemaker_studio_image_build/codebuild.py b/sagemaker_studio_image_build/codebuild.py index 86ecf85..e4d4da5 100644 --- a/sagemaker_studio_image_build/codebuild.py +++ b/sagemaker_studio_image_build/codebuild.py @@ -18,8 +18,8 @@ def __init__(self, s3_location, role, repository=None, compute_type=None, enviro 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.environment = environment or "LINUX_CONTAINER" + self.compute_type = compute_type + self.environment = environment if self.environment=="LINUX_GPU_CONTAINER": assert self.compute_type=="BUILD_GENERAL1_LARGE", \ "LINUX_GPU_CONTAINER builds only available on BUILD_GENERAL1_LARGE. Please set `--compute-type BUILD_GENERAL1_LARGE`"