From df814d043c291ccdb7cdf5c46d72579faeb7796d Mon Sep 17 00:00:00 2001 From: Marcel Gwerder Date: Fri, 17 Sep 2021 14:35:26 +0200 Subject: [PATCH] feat: added option to specify codebuild env variables --- README.md | 5 +++++ sagemaker_studio_image_build/builder.py | 4 ++-- sagemaker_studio_image_build/cli.py | 15 ++++++++++++++- sagemaker_studio_image_build/codebuild.py | 4 +++- 4 files changed, 24 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 1868de7..9c9e45f 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,11 @@ These can be overridden with the relevant CLI options. 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 ``` +Additional environment variables can be specified for the CodeBuild project with the `--env` argument. Use `--env` multiple times to specify more than one variable. +```bash +sm-docker build . --file /path/to/Dockerfile --env DOCKER_BUILDKIT=1 +``` + 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. diff --git a/sagemaker_studio_image_build/builder.py b/sagemaker_studio_image_build/builder.py index 18e8364..aa9c07e 100644 --- a/sagemaker_studio_image_build/builder.py +++ b/sagemaker_studio_image_build/builder.py @@ -64,13 +64,13 @@ def delete_zip_file(bucket, key): s3.delete_object(Bucket=bucket, Key=key) -def build_image(repository, role, bucket, compute_type, vpc_config, extra_args, log=True): +def build_image(repository, role, bucket, compute_type, vpc_config, extra_args, env, 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, vpc_config=vpc_config) as p: + compute_type=compute_type, vpc_config=vpc_config, env=env) 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 146ccd6..e20b44a 100644 --- a/sagemaker_studio_image_build/cli.py +++ b/sagemaker_studio_image_build/cli.py @@ -72,10 +72,18 @@ def build_image(args, extra_args): builder.build_image( args.repository, get_role(args), args.bucket, args.compute_type, - construct_vpc_config(args), extra_args, log=not args.no_logs + construct_vpc_config(args), extra_args, args.env, log=not args.no_logs ) +class AppendEnvAction(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + items = getattr(namespace, self.dest, None) + items = [] if items is None else items[:] + items.append(dict(zip(["name", "value"], values.split("=", 1)))) + setattr(namespace, self.dest, items) + + def main(): parser = argparse.ArgumentParser( @@ -119,6 +127,11 @@ def main(): "--security-group-ids", help="The comma-separated list of security group ids for the CodeBuild Project (such as sg-0ce4ec0d0414d2ddc).", ) + build_parser.add_argument( + "--env", + action=AppendEnvAction, + help="An environment variable to set for the CodeBuild Project (such as DOCKER_BUILDKIT=1). Can be specified multiple times.", + ) build_parser.add_argument( "--no-logs", action="store_true", diff --git a/sagemaker_studio_image_build/codebuild.py b/sagemaker_studio_image_build/codebuild.py index a9c99b7..8c1b0a3 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, vpc_config=None): + def __init__(self, s3_location, role, repository=None, compute_type=None, vpc_config=None, env=[]): self.s3_location = s3_location self.role = role @@ -20,6 +20,7 @@ def __init__(self, s3_location, role, repository=None, compute_type=None, vpc_co self.repo_name = None self.compute_type = compute_type or "BUILD_GENERAL1_SMALL" self.vpc_config = vpc_config + self.env = env if repository: self.repo_name, self.tag = repository.split(":", maxsplit=1) @@ -70,6 +71,7 @@ def __enter__(self): {"name": "AWS_ACCOUNT_ID", "value": account}, {"name": "IMAGE_REPO_NAME", "value": self.repo_name}, {"name": "IMAGE_TAG", "value": self.tag}, + *self.env ], "privilegedMode": True, },