From d0d4634a21260b8ba6d6bd71150f92b5ced751fd Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 13 Nov 2024 10:21:37 -0600
Subject: [PATCH] enforce wheel size limits, README formatting in CI

---
 ci/build_wheel_distributed_ucxx.sh     |  1 +
 ci/build_wheel_libucxx.sh              |  2 ++
 ci/build_wheel_ucxx.sh                 |  2 ++
 ci/validate_wheel.sh                   | 21 +++++++++++++++++++++
 python/distributed-ucxx/pyproject.toml |  8 ++++++++
 python/libucxx/pyproject.toml          |  8 ++++++++
 python/ucxx/pyproject.toml             |  8 ++++++++
 7 files changed, 50 insertions(+)
 create mode 100755 ci/validate_wheel.sh

diff --git a/ci/build_wheel_distributed_ucxx.sh b/ci/build_wheel_distributed_ucxx.sh
index f6ee95a7..512e36bd 100755
--- a/ci/build_wheel_distributed_ucxx.sh
+++ b/ci/build_wheel_distributed_ucxx.sh
@@ -8,5 +8,6 @@ package_dir="python/distributed-ucxx"
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
 
 ./ci/build_wheel.sh distributed-ucxx "${package_dir}"
+./ci/validate_wheel.sh "${package_dir}" dist
 
 RAPIDS_PY_WHEEL_NAME="distributed_ucxx_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 python "${package_dir}/dist"
diff --git a/ci/build_wheel_libucxx.sh b/ci/build_wheel_libucxx.sh
index e9262077..4ae818b5 100755
--- a/ci/build_wheel_libucxx.sh
+++ b/ci/build_wheel_libucxx.sh
@@ -37,4 +37,6 @@ python -m auditwheel repair \
     -w "${package_dir}/final_dist" \
     ${package_dir}/dist/*
 
+./ci/validate_wheel.sh "${package_dir}" final_dist
+
 RAPIDS_PY_WHEEL_NAME="${package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 cpp "${package_dir}/final_dist"
diff --git a/ci/build_wheel_ucxx.sh b/ci/build_wheel_ucxx.sh
index 78602085..b8e60c4c 100755
--- a/ci/build_wheel_ucxx.sh
+++ b/ci/build_wheel_ucxx.sh
@@ -28,4 +28,6 @@ python -m auditwheel repair \
     -w "${package_dir}/final_dist" \
     ${package_dir}/dist/*
 
+./ci/validate_wheel.sh "${package_dir}" final_dist
+
 RAPIDS_PY_WHEEL_NAME="${package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 python "${package_dir}/final_dist"
diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh
new file mode 100755
index 00000000..5910a5c5
--- /dev/null
+++ b/ci/validate_wheel.sh
@@ -0,0 +1,21 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+package_dir=$1
+wheel_dir_relative_path=$2
+
+cd "${package_dir}"
+
+rapids-logger "validate packages with 'pydistcheck'"
+
+pydistcheck \
+    --inspect \
+    "$(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/distributed-ucxx/pyproject.toml b/python/distributed-ucxx/pyproject.toml
index 1005df8d..da7291cf 100644
--- a/python/distributed-ucxx/pyproject.toml
+++ b/python/distributed-ucxx/pyproject.toml
@@ -125,6 +125,14 @@ matrix-entry = "cuda_suffixed=true"
 [tool.setuptools.dynamic]
 version = {file = "distributed_ucxx/VERSION"}
 
+[tool.pydistcheck]
+select = [
+    "distro-too-large-compressed",
+]
+
+# PyPI limit is 100 MiB, fail CI before we get too close to that
+max_allowed_size_compressed = '75M'
+
 [tool.pytest.ini_options]
 markers = [
     "ignore_alive_references",
diff --git a/python/libucxx/pyproject.toml b/python/libucxx/pyproject.toml
index e035effb..8309badf 100644
--- a/python/libucxx/pyproject.toml
+++ b/python/libucxx/pyproject.toml
@@ -65,3 +65,11 @@ requires = [
     "libucx==1.15.0",
     "ninja",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
+
+[tool.pydistcheck]
+select = [
+    "distro-too-large-compressed",
+]
+
+# PyPI limit is 100 MiB, fail CI before we get too close to that
+max_allowed_size_compressed = '75M'
diff --git a/python/ucxx/pyproject.toml b/python/ucxx/pyproject.toml
index c548e85e..730fc0cb 100644
--- a/python/ucxx/pyproject.toml
+++ b/python/ucxx/pyproject.toml
@@ -79,3 +79,11 @@ wheel.exclude = ["*.pyx", "CMakeLists.txt"]
 provider = "scikit_build_core.metadata.regex"
 input = "ucxx/VERSION"
 regex = "(?P<value>.*)"
+
+[tool.pydistcheck]
+select = [
+    "distro-too-large-compressed",
+]
+
+# PyPI limit is 100 MiB, fail CI before we get too close to that
+max_allowed_size_compressed = '75M'