feat: added option to specify additional codebuild env variables

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.
 
 

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)

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",

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,
             },