Fix generate image size report

Author: bhavya2109sharma

.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

src/main.py

@@ -428,6 +428,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",

src/package_report.py

@@ -271,7 +271,7 @@ def generate_package_size_report(args):
     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_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]:
@@ -312,13 +312,13 @@ 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):
+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,
+            registryId=staging_account,
+            repositoryName=staging_repo_name,
             imageIds=[{'imageTag': image_tag}]
         )
         image_details = response.get('imageDetails', [])
@@ -328,27 +328,27 @@ def get_image_size(image_tag):
                 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_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)
-    base_size_cpu = get_image_size(base_tag_cpu) if base_tag_cpu 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 (%)"]