From cc39cabedc94930c5abf6c979930f4d27d4ec9b9 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Thu, 14 Nov 2024 17:06:29 -0600
Subject: [PATCH] use different wheel-size thresholds based on CUDA version

---
 ci/validate_wheel.sh       | 14 ++++++++++++++
 python/cuvs/pyproject.toml |  4 +---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh
index 5910a5c59f..f2b235765c 100755
--- a/ci/validate_wheel.sh
+++ b/ci/validate_wheel.sh
@@ -6,12 +6,26 @@ 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 [[ "${RAPIDS_CUDA_MAJOR}" == "11" ]]; then
+    PYDISTCHECK_ARGS=(
+        --max-allowed-size-compressed '1.4G'
+    )
+else
+    PYDISTCHECK_ARGS=(
+        --max-allowed-size-compressed '950M'
+    )
+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'"
diff --git a/python/cuvs/pyproject.toml b/python/cuvs/pyproject.toml
index d400267760..92e4993c76 100644
--- a/python/cuvs/pyproject.toml
+++ b/python/cuvs/pyproject.toml
@@ -135,12 +135,10 @@ 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",
 ]
 
-# detect when package size grows significantly
-max_allowed_size_compressed = '1.4G'
-
 [tool.pytest.ini_options]
 filterwarnings = [
     "error",