From ad13d08045a4309f489ad2dc540885c33073ef30 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Tue, 17 Dec 2024 12:55:03 -0800
Subject: [PATCH] download libraft in wheel test jobs, change file-size limits

---
 ci/test_wheel_pylibraft.sh | 8 ++++++--
 ci/test_wheel_raft_dask.sh | 8 ++++----
 ci/validate_wheel.sh       | 2 +-
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/ci/test_wheel_pylibraft.sh b/ci/test_wheel_pylibraft.sh
index b38f5a690b..1e0b34d609 100755
--- a/ci/test_wheel_pylibraft.sh
+++ b/ci/test_wheel_pylibraft.sh
@@ -5,9 +5,13 @@ set -euo pipefail
 
 mkdir -p ./dist
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
-RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist
+RAPIDS_PY_WHEEL_NAME="libraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 cpp ./local-libraft-dep
+RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 python ./dist
+
 
 # echo to expand wildcard before adding `[extra]` requires for pip
-python -m pip install $(echo ./dist/pylibraft*.whl)[test]
+python -m pip install \
+    ./local-libraft-dep/libraft*.whl \
+    "$(echo ./dist/pylibraft*.whl)[test]"
 
 python -m pytest ./python/pylibraft/pylibraft/test
diff --git a/ci/test_wheel_raft_dask.sh b/ci/test_wheel_raft_dask.sh
index a778a3ec51..011de4d409 100755
--- a/ci/test_wheel_raft_dask.sh
+++ b/ci/test_wheel_raft_dask.sh
@@ -5,13 +5,13 @@ set -euo pipefail
 
 mkdir -p ./dist
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
-RAPIDS_PY_WHEEL_NAME="raft_dask_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist
-
-# Download the pylibraft built in the previous step
-RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./local-pylibraft-dep
+RAPIDS_PY_WHEEL_NAME="libraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 cpp ./local-libraft-dep
+RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 python ./local-pylibraft-dep
+RAPIDS_PY_WHEEL_NAME="raft_dask_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 python ./dist
 
 # echo to expand wildcard before adding `[extra]` requires for pip
 python -m pip install -v \
+    ./local-libraft-dep/libraft*.whl \
     ./local-pylibraft-dep/pylibraft*.whl \
     "$(echo ./dist/raft_dask_${RAPIDS_PY_CUDA_SUFFIX}*.whl)[test]"
 
diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh
index 018abcba23..708aeb029c 100755
--- a/ci/validate_wheel.sh
+++ b/ci/validate_wheel.sh
@@ -28,7 +28,7 @@ elif [[ "${package_name}" == "pylibraft" ]] || [[ "${package_name}" == "raft-das
         )
     else
         PYDISTCHECK_ARGS=(
-            --max-allowed-size-compressed '100M'
+            --max-allowed-size-compressed '400M'
         )
     fi
 else