Merge remote-tracking branch 'origin/main' into release-3.0.2

Author: claytonparnell

src/main.py

@@ -424,6 +424,11 @@ def get_arg_parser():
         required=True,
         help="Specify the base patch version for which the package staleness report needs to be " "generated.",
     )
+    package_staleness_parser.add_argument(
+        "--download-stats",
+        action="store_true",
+        help="Specify if download count stats should be included in the package staleness report.",
+    )
     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 "

src/package_report.py

@@ -64,19 +64,20 @@ def _get_package_versions_in_upstream(target_packages_match_spec_out, target_ver
 
 
 def _generate_staleness_report_per_image(
-    package_versions_in_upstream, target_packages_match_spec_out, image_config, version
+    package_versions_in_upstream, target_packages_match_spec_out, image_config, version, download_stats
 ):
     print("\n# Staleness Report: " + str(version) + "(" + image_config["image_type"] + ")\n")
     staleness_report_rows = []
 
-    # Get conda download statistics for all installed packages
-    # Use the month before last to get full month of data
-    previous_month = (datetime.now() - relativedelta(months=2)).strftime("%Y-%m")
-    pkg_list = list(package_versions_in_upstream.keys())
-    # Suppress FutureWarning from pandas so it doesn't show in report
-    with warnings.catch_warnings():
-        warnings.filterwarnings("ignore", category=FutureWarning)
-        conda_download_stats = overall(pkg_list, month=previous_month)
+    if download_stats:
+        # Get conda download statistics for all installed packages
+        # Use the month before last to get full month of data
+        previous_month = (datetime.now() - relativedelta(months=2)).strftime("%Y-%m")
+        pkg_list = list(package_versions_in_upstream.keys())
+        # Suppress FutureWarning from pandas so it doesn't show in report
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning)
+            conda_download_stats = overall(pkg_list, month=previous_month)
 
     for package in package_versions_in_upstream:
         version_in_sagemaker_distribution = str(target_packages_match_spec_out[package].get("version")).removeprefix(
@@ -87,32 +88,45 @@ def _generate_staleness_report_per_image(
             if version_in_sagemaker_distribution == package_versions_in_upstream[package]
             else "${\color{red}" + package + "}$"
         )
-        staleness_report_rows.append(
-            {
-                "package": package_string,
-                "version_in_sagemaker_distribution": version_in_sagemaker_distribution,
-                "latest_relavant_version": package_versions_in_upstream[package],
-                "downloads": conda_download_stats[package],
-            }
-        )
 
-    staleness_report_rows.sort(
-        key=lambda x: (
-            not x["package"].startswith("${\\color"),  # Stale packages at top of list
-            -x["downloads"],  # Sorted by downloads
-        )
-    )
-    print(
-        create_markdown_table(
-            [
-                "Package",
-                "Current Version in the Distribution image",
-                "Latest Relevant Version in " "Upstream",
-                "Downloads (Conda, previous month)",
-            ],
-            staleness_report_rows,
+        if download_stats:
+            # Get download count with error handling
+            try:
+                download_count = conda_download_stats[package]
+            except (KeyError, TypeError):
+                download_count = 0
+
+            staleness_report_rows.append(
+                {
+                    "package": package_string,
+                    "version_in_sagemaker_distribution": version_in_sagemaker_distribution,
+                    "latest_relavant_version": package_versions_in_upstream[package],
+                    "downloads": download_count,
+                }
+            )
+        else:
+            staleness_report_rows.append(
+                {
+                    "package": package_string,
+                    "version_in_sagemaker_distribution": version_in_sagemaker_distribution,
+                    "latest_relavant_version": package_versions_in_upstream[package],
+                }
+            )
+
+    markdown_table_columns = [
+        "Package",
+        "Current Version in the Distribution image",
+        "Latest Relevant Version in " "Upstream",
+    ]
+    if download_stats:
+        markdown_table_columns.append("Downloads (Conda, previous month)")
+        staleness_report_rows.sort(
+            key=lambda x: (
+                not x["package"].startswith("${\\color"),  # Stale packages at top of list
+                -x["downloads"],  # Sorted by downloads
+            )
         )
-    )
+    print(create_markdown_table(markdown_table_columns, staleness_report_rows))
 
 
 def _get_installed_package_versions_and_conda_versions(
@@ -281,7 +295,11 @@ def generate_package_staleness_report(args):
             latest_package_versions_in_upstream,
         ) = _get_installed_package_versions_and_conda_versions(image_config, target_version_dir, target_version)
         _generate_staleness_report_per_image(
-            latest_package_versions_in_upstream, target_packages_match_spec_out, image_config, target_version
+            latest_package_versions_in_upstream,
+            target_packages_match_spec_out,
+            image_config,
+            target_version,
+            args.download_stats,
         )