From cbbff02198121f0f006a697afc0a15a460204010 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Thu, 14 Nov 2024 18:15:22 -0600
Subject: [PATCH] enforce wheel size limits, README formatting in CI (#4754)

Contributes to https://github.com/rapidsai/build-planning/issues/110

Proposes adding 2 types of validation on wheels in CI, to ensure we continue to produce wheels that are suitable for PyPI.

* checks on wheel size (compressed),
  - *to be sure they're under PyPI limits*
  - *and to prompt discussion on PRs that significantly increase wheel sizes*
* checks on README formatting
  - *to ensure they'll render properly as the PyPI project homepages*
  - *e.g. like how https://github.com/scikit-learn/scikit-learn/blob/main/README.rst becomes https://pypi.org/project/scikit-learn/*

This intentionally omits the libraries that are about to be moved to other repos.

* `cugraph-dgl` / `cugraph-pyg` (https://github.com/rapidsai/cugraph-gnn/pull/67)
* `nx-cugraph` (https://github.com/rapidsai/nx-cugraph/pull/19)

Authors:
  - James Lamb (https://github.com/jameslamb)

Approvers:
  - Bradley Dice (https://github.com/bdice)

URL: https://github.com/rapidsai/cugraph/pull/4754
---
 ci/build_wheel_cugraph.sh          |  5 +++-
 ci/build_wheel_pylibcugraph.sh     |  5 +++-
 ci/validate_wheel.sh               | 39 ++++++++++++++++++++++++++++++
 python/cugraph/pyproject.toml      |  6 +++++
 python/pylibcugraph/pyproject.toml |  6 +++++
 5 files changed, 59 insertions(+), 2 deletions(-)
 create mode 100755 ci/validate_wheel.sh

diff --git a/ci/build_wheel_cugraph.sh b/ci/build_wheel_cugraph.sh
index 20d9bf47e3e..b937fd7f6d8 100755
--- a/ci/build_wheel_cugraph.sh
+++ b/ci/build_wheel_cugraph.sh
@@ -3,6 +3,8 @@
 
 set -euo pipefail
 
+package_dir="python/cugraph"
+
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
 
 # Download the pylibcugraph wheel built in the previous step and make it
@@ -31,4 +33,5 @@ esac
 export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_CUGRAPH_CPP=OFF;-DCPM_cugraph-ops_SOURCE=${GITHUB_WORKSPACE}/cugraph-ops/${EXTRA_CMAKE_ARGS}"
 export SKBUILD_BUILD_TOOL_ARGS="-j${PARALLEL_LEVEL};-l${PARALLEL_LEVEL}"
 
-./ci/build_wheel.sh cugraph python/cugraph
+./ci/build_wheel.sh cugraph ${package_dir}
+./ci/validate_wheel.sh ${package_dir} final_dist
diff --git a/ci/build_wheel_pylibcugraph.sh b/ci/build_wheel_pylibcugraph.sh
index fa967b0be29..928e12d2350 100755
--- a/ci/build_wheel_pylibcugraph.sh
+++ b/ci/build_wheel_pylibcugraph.sh
@@ -3,6 +3,8 @@
 
 set -euo pipefail
 
+package_dir="python/pylibcugraph"
+
 PARALLEL_LEVEL=$(python -c \
   "from math import ceil; from multiprocessing import cpu_count; print(ceil(cpu_count()/4))")
 
@@ -18,4 +20,5 @@ esac
 export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_CUGRAPH_CPP=OFF;-DCPM_cugraph-ops_SOURCE=${GITHUB_WORKSPACE}/cugraph-ops/${EXTRA_CMAKE_ARGS}"
 export SKBUILD_BUILD_TOOL_ARGS="-j${PARALLEL_LEVEL};-l${PARALLEL_LEVEL}"
 
-./ci/build_wheel.sh pylibcugraph python/pylibcugraph
+./ci/build_wheel.sh pylibcugraph ${package_dir}
+./ci/validate_wheel.sh ${package_dir} final_dist
diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh
new file mode 100755
index 00000000000..5ebac95a3f9
--- /dev/null
+++ b/ci/validate_wheel.sh
@@ -0,0 +1,39 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+package_dir=$1
+wheel_dir_relative_path=$2
+
+RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"
+
+# some packages are much larger on CUDA 11 than on CUDA 12
+if [[ "${package_dir}" == "python/cugraph" ]] || [[ "${package_dir}" == "python/pylibcugraph" ]]; then
+    if [[ "${RAPIDS_CUDA_MAJOR}" == "11" ]]; then
+        PYDISTCHECK_ARGS=(
+            --max-allowed-size-compressed '1.5G'
+        )
+    else
+        PYDISTCHECK_ARGS=(
+            --max-allowed-size-compressed '975M'
+        )
+    fi
+else
+    PYDISTCHECK_ARGS=()
+fi
+
+cd "${package_dir}"
+
+rapids-logger "validate packages with 'pydistcheck'"
+
+pydistcheck \
+    --inspect \
+    "${PYDISTCHECK_ARGS[@]}" \
+    "$(echo ${wheel_dir_relative_path}/*.whl)"
+
+rapids-logger "validate packages with 'twine'"
+
+twine check \
+    --strict \
+    "$(echo ${wheel_dir_relative_path}/*.whl)"
diff --git a/python/cugraph/pyproject.toml b/python/cugraph/pyproject.toml
index 8185a8d915d..d3960ab5d32 100644
--- a/python/cugraph/pyproject.toml
+++ b/python/cugraph/pyproject.toml
@@ -88,3 +88,9 @@ requires = [
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 dependencies-file = "../../dependencies.yaml"
 matrix-entry = "cuda_suffixed=true"
+
+[tool.pydistcheck]
+select = [
+    # NOTE: size threshold is managed via CLI args in CI scripts
+    "distro-too-large-compressed",
+]
diff --git a/python/pylibcugraph/pyproject.toml b/python/pylibcugraph/pyproject.toml
index c12280473b5..8d22f50cc93 100644
--- a/python/pylibcugraph/pyproject.toml
+++ b/python/pylibcugraph/pyproject.toml
@@ -78,3 +78,9 @@ requires = [
     "rmm==24.12.*,>=0.0.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true"
+
+[tool.pydistcheck]
+select = [
+    # NOTE: size threshold is managed via CLI args in CI scripts
+    "distro-too-large-compressed",
+]