From d8f8ab46b3ab0042cc9a2d42ebceba3488fc43a7 Mon Sep 17 00:00:00 2001 From: Bhavya Sharma Date: Sat, 28 Dec 2024 04:48:57 +0000 Subject: [PATCH 1/2] Fix generate image size report --- .github/workflows/check-image-size.yml | 14 ++++- src/main.py | 13 ++++- src/package_report.py | 72 ++++++++++++++------------ 3 files changed, 65 insertions(+), 34 deletions(-) diff --git a/.github/workflows/check-image-size.yml b/.github/workflows/check-image-size.yml index d8a7a7e5..ca2e9d88 100644 --- a/.github/workflows/check-image-size.yml +++ b/.github/workflows/check-image-size.yml @@ -6,12 +6,24 @@ on: image-version: required: true description: Image version= + staging-repo-name: + required: true + description: Staging repository name + staging-account: + required: true + description: Staging account ID # Call from other workflow workflow_call: inputs: image-version: type: string required: true + staging-repo-name: + type: string + required: true + staging-account: + type: string + required: true defaults: run: @@ -37,4 +49,4 @@ jobs: - name: Activate sagemaker-distribution run: micromamba activate sagemaker-distribution - name: Run size validation - run: python ./src/main.py generate-size-report --target-patch-version ${{ inputs.image-version }} --validate + run: python ./src/main.py generate-size-report --target-patch-version ${{ inputs.image-version }} --staging-repo-name ${{ inputs.staging-repo-name }} --staging-account ${{ inputs.staging-account }} --validate diff --git a/src/main.py b/src/main.py index 7d4429b2..55de8ec3 100644 --- a/src/main.py +++ b/src/main.py @@ -420,7 +420,8 @@ def get_arg_parser(): ) package_size_parser = subparsers.add_parser( "generate-size-report", - help="Generates toatl image size and package size report for each of the packages in the given " "image version.", + help="Generates toatl image size and package size report for each of the packages in the given " + "image version.", ) package_size_parser.set_defaults(func=generate_package_size_report) package_size_parser.add_argument( @@ -428,6 +429,16 @@ def get_arg_parser(): required=True, help="Specify the target patch version for which the package size report needs to be " "generated.", ) + package_size_parser.add_argument( + "--staging-repo-name", + required=True, + help="Specify the staging repository", + ) + package_size_parser.add_argument( + "--staging-account", + required=True, + help="Specify the staging account", + ) package_size_parser.add_argument( "--validate", action="store_true", diff --git a/src/package_report.py b/src/package_report.py index 17a2556b..2e06b692 100644 --- a/src/package_report.py +++ b/src/package_report.py @@ -269,10 +269,10 @@ def generate_package_size_report(args): source_patch_version = f.readline() base_version = get_semver(source_patch_version) base_version_dir = get_dir_for_version(base_version) if base_version else None - - #Generate the report for Total Image Size changed from Base Version - _generate_image_size_report(target_version, base_version) - + + # Generate the report for Total Image Size changed from Base Version + _generate_image_size_report(target_version, base_version, args.staging_repo_name, args.staging_account) + validate_results = [] for image_config in _image_generator_configs[target_version.major]: base_pkg_metadata = pull_conda_package_metadata(image_config, base_version_dir) if base_version else None @@ -312,61 +312,65 @@ def generate_package_dependency_report(args): print("## Image Type: " + "(" + image_config["image_type"].upper() + ")") _generate_python_package_dependency_report(image_config, base_version_dir, target_version_dir) -def get_image_size(image_tag): - try: - ecr_client = boto3.client('ecr', region_name="us-east-1") + +def get_image_size(image_tag, staging_repo_name, staging_account): + try: + ecr_client = boto3.client("ecr", region_name="us-east-1") response = ecr_client.describe_images( - registryId=STAGING_ACCOUNT, - repositoryName=STAGING_REPO_NAME, - imageIds=[{'imageTag': image_tag}] + registryId=staging_account, repositoryName=staging_repo_name, imageIds=[{"imageTag": image_tag}] ) - image_details = response.get('imageDetails', []) + image_details = response.get("imageDetails", []) if image_details: - image_size_bytes = image_details[0].get('imageSizeInBytes') + image_size_bytes = image_details[0].get("imageSizeInBytes") if image_size_bytes: - image_size_gb = image_size_bytes / (1024 ** 3) + image_size_gb = image_size_bytes / (1024**3) return image_size_gb else: - print(f"Image size not found for {STAGING_REPO_NAME}:{image_tag}") + print(f"Image size not found for {staging_repo_name}:{image_tag}") else: - print(f"No image details found for {STAGING_REPO_NAME}:{image_tag}") - + print(f"No image details found for {staging_repo_name}:{image_tag}") + except Exception as e: print(f"Error retrieving image size: {str(e)}") - + return None -def _generate_image_size_report(target_version, base_version): + +def _generate_image_size_report(target_version, base_version, staging_repo_name, staging_account): target_tag_gpu = f"{target_version.major}.{target_version.minor}.{target_version.patch}-gpu" base_tag_gpu = f"{base_version.major}.{base_version.minor}.{base_version.patch}-gpu" if base_version else None target_tag_cpu = f"{target_version.major}.{target_version.minor}.{target_version.patch}-cpu" base_tag_cpu = f"{base_version.major}.{base_version.minor}.{base_version.patch}-cpu" if base_version else None - - target_size_gpu = get_image_size(target_tag_gpu) - base_size_gpu = get_image_size(base_tag_gpu) if base_tag_gpu else None - target_size_cpu = get_image_size(target_tag_cpu) - base_size_cpu = get_image_size(base_tag_cpu) if base_tag_cpu else None - + target_size_gpu = get_image_size(target_tag_gpu, staging_repo_name, staging_account) + base_size_gpu = get_image_size(base_tag_gpu, staging_repo_name, staging_account) if base_tag_gpu else None + + target_size_cpu = get_image_size(target_tag_cpu, staging_repo_name, staging_account) + base_size_cpu = get_image_size(base_tag_cpu, staging_repo_name, staging_account) if base_tag_cpu else None + headers = ["Target Image Size", "Base Image Size", "Size Difference", "Size Change (%)"] print("\n# GPU Total Image Size Report\n") if target_size_gpu is not None and base_size_gpu is not None: size_diff_gpu = target_size_gpu - base_size_gpu size_diff_gpu_rel = size_diff_gpu / base_size_gpu if base_size_gpu != 0 else 0 - + size_change_gpu = f"{size_diff_gpu:.2f} GB {'Larger' if size_diff_gpu > 0 else ('Smaller' if size_diff_gpu < 0 else 'No Change')}" - - size_change_percentage_gpu = f"{round(abs(size_diff_gpu_rel * 100), 2)}% {'Larger' if size_diff_gpu > 0 else ('Smaller' if size_diff_gpu < 0 else 'No Change')}" if size_diff_gpu_rel else "-" + + size_change_percentage_gpu = ( + f"{round(abs(size_diff_gpu_rel * 100), 2)}% {'Larger' if size_diff_gpu > 0 else ('Smaller' if size_diff_gpu < 0 else 'No Change')}" + if size_diff_gpu_rel + else "-" + ) rows = [ { "Target Image Size": f"{target_tag_gpu}:{target_size_gpu:.2f} GB", "Base Image Size": f"{base_tag_gpu}:{base_size_gpu:.2f} GB", "Size Difference": size_change_gpu, - "Size Change (%)": size_change_percentage_gpu + "Size Change (%)": size_change_percentage_gpu, } ] print(create_markdown_table(headers, rows)) @@ -377,17 +381,21 @@ def _generate_image_size_report(target_version, base_version): if target_size_cpu is not None and base_size_cpu is not None: size_diff_cpu = target_size_cpu - base_size_cpu size_diff_cpu_rel = size_diff_cpu / base_size_cpu if base_size_cpu != 0 else 0 - + size_change_cpu = f"{size_diff_cpu:.2f} GB {'Larger' if size_diff_cpu > 0 else ('Smaller' if size_diff_cpu < 0 else 'No Change')}" - - size_change_percentage_cpu = f"{round(abs(size_diff_cpu_rel * 100), 2)}% {'Larger' if size_diff_cpu > 0 else ('Smaller' if size_diff_cpu < 0 else 'No Change')}" if size_diff_cpu_rel else "-" + + size_change_percentage_cpu = ( + f"{round(abs(size_diff_cpu_rel * 100), 2)}% {'Larger' if size_diff_cpu > 0 else ('Smaller' if size_diff_cpu < 0 else 'No Change')}" + if size_diff_cpu_rel + else "-" + ) rows = [ { "Target Image Size": f"{target_tag_cpu}:{target_size_cpu:.2f} GB", "Base Image Size": f"{base_tag_cpu}:{base_size_cpu:.2f} GB", "Size Difference": size_change_cpu, - "Size Change (%)": size_change_percentage_cpu + "Size Change (%)": size_change_percentage_cpu, } ] print(create_markdown_table(headers, rows)) From b7c93486521841bb1ac85dcd64489fe11387d043 Mon Sep 17 00:00:00 2001 From: Bhavya Sharma Date: Sat, 28 Dec 2024 04:48:57 +0000 Subject: [PATCH 2/2] Fix generate image size report --- src/package_report.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/package_report.py b/src/package_report.py index 2e06b692..bd4bccac 100644 --- a/src/package_report.py +++ b/src/package_report.py @@ -1,8 +1,8 @@ import json import os -import boto3 from itertools import islice +import boto3 import conda.cli.python_api from conda.models.match_spec import MatchSpec