def upload_zip_file()

in sagemaker_studio_image_build/builder.py [0:0]

• 35 lines of code
• 6 McCabe index (conditional complexity)

def upload_zip_file(repo_name, bucket, extra_args, dir="."):
    """
    1. Zip up the PWD.
    2. Replace placeholders in buildspec.yml and add to the Zip.
    3. Upload to S3.
    """
    if not bucket:
        # The SDK logs a warning for not having pandas and this is the only way to suppress it.
        # https://github.com/aws/sagemaker-python-sdk/issues/1696
        logging.basicConfig(level=logging.CRITICAL)
        import sagemaker.session as session

        logging.basicConfig(level=logging.INFO)

        bucket = session.Session().default_bucket()

    random_suffix = "".join(random.choices(string.ascii_letters, k=16))
    key = f"codebuild-sagemaker-container-{random_suffix}.zip"
    origdir = os.getcwd()
    os.chdir(dir)
    try:
        with tempfile.TemporaryFile() as tmp:
            with zipfile.ZipFile(tmp, "w") as zip:
                # Zip all files in "dir"
                for dirname, _, filelist in os.walk("."):
                    for file in filelist:
                        zip.write(f"{dirname}/{file}")
                # Add buildspec.yml
                data_dir = os.path.join(
                    os.path.dirname(os.path.abspath(__file__)), "data"
                )
                with tempfile.NamedTemporaryFile() as buildspec:
                    with open(
                        os.path.join(data_dir, "buildspec.template.yml")
                    ) as buildspec_template:
                        buildspec_replaced = buildspec_template.read().replace(
                            "REPLACE_ME_BUILD_ARGS", extra_args
                        )
                        buildspec.write(buildspec_replaced.encode())
                    buildspec.seek(0)
                    zip.write(buildspec.name, "buildspec.yml")
            tmp.seek(0)
            s3 = boto3.session.Session().client("s3")
            s3.upload_fileobj(tmp, bucket, key)
        return (bucket, key)
    finally:
        os.chdir(origdir)