try using librmm wheels

rapidsai · Sep 16, 2024 · b7c151e · b7c151e
1 parent 2a6aeea
commit b7c151e
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 1 deletion.
diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
@@ -39,8 +39,16 @@ case "${RAPIDS_CUDA_VERSION}" in
   ;;
 esac
 
+LIBRMM_CHANNEL=$(rapids-get-pr-wheel-artifact rmm 1678 cpp)
+RMM_CHANNEL=$(rapids-get-pr-wheel-artifact rmm 1678 python)
+
+echo "rmm-${RAPIDS_PY_CUDA_SUFFIX} @ $(echo ${RMM_CHANNEL}/rmm*.whl)" > /tmp/constraints.txt
+echo "librmm-${RAPIDS_PY_CUDA_SUFFIX} @ $(echo ${LIBRMM_CHANNEL}/librmm*.whl)" >> /tmp/constraints.txt
+echo "" >> /tmp/constraints.txt
+export PIP_CONSTRAINT=/tmp/constraints.txt
+
 # Hardcode the output dir
-python -m pip wheel . -w dist -vvv --no-deps --disable-pip-version-check
+python -m pip wheel . -w dist -v --no-deps --disable-pip-version-check
 
 mkdir -p final_dist
 python -m auditwheel repair -w final_dist "${EXCLUDE_ARGS[@]}" dist/*

diff --git a/ci/test_wheel_pylibraft.sh b/ci/test_wheel_pylibraft.sh
@@ -7,6 +7,14 @@ 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
 
+LIBRMM_CHANNEL=$(rapids-get-pr-wheel-artifact rmm 1678 cpp)
+RMM_CHANNEL=$(rapids-get-pr-wheel-artifact rmm 1678 python)
+
+echo "rmm-${RAPIDS_PY_CUDA_SUFFIX} @ $(echo ${RMM_CHANNEL}/rmm*.whl)" > /tmp/constraints.txt
+echo "librmm-${RAPIDS_PY_CUDA_SUFFIX} @ $(echo ${LIBRMM_CHANNEL}/librmm*.whl)" >> /tmp/constraints.txt
+echo "" >> /tmp/constraints.txt
+export PIP_CONSTRAINT=/tmp/constraints.txt
+
 # echo to expand wildcard before adding `[extra]` requires for pip
 python -m pip install $(echo ./dist/pylibraft*.whl)[test]
 

diff --git a/ci/test_wheel_raft_dask.sh b/ci/test_wheel_raft_dask.sh
@@ -7,6 +7,14 @@ 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
 
+LIBRMM_CHANNEL=$(rapids-get-pr-wheel-artifact rmm 1678 cpp)
+RMM_CHANNEL=$(rapids-get-pr-wheel-artifact rmm 1678 python)
+
+echo "rmm-${RAPIDS_PY_CUDA_SUFFIX} @ $(echo ${RMM_CHANNEL}/rmm*.whl)" > /tmp/constraints.txt
+echo "librmm-${RAPIDS_PY_CUDA_SUFFIX} @ $(echo ${LIBRMM_CHANNEL}/librmm*.whl)" >> /tmp/constraints.txt
+echo "" >> /tmp/constraints.txt
+export PIP_CONSTRAINT=/tmp/constraints.txt
+
 # 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
 python -m pip install --no-deps ./local-pylibraft-dep/pylibraft*.whl