Use nvcomp wheel instead of bundling nvcomp (#478)

Contributes to https://github.com/rapdisai/build-planning/issues/74

Authors:
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)

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

URL: #478

ci/build_wheel_python.sh

@@ -24,9 +24,13 @@ cd "${package_dir}"
 echo "libkvikio-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo ${CPP_WHEELHOUSE}/libkvikio_*.whl)" > ./constraints.txt
 
 PIP_CONSTRAINT="${PWD}/constraints.txt" \
+SKBUILD_CMAKE_ARGS="-DUSE_NVCOMP_RUNTIME_WHEEL=ON" \
     python -m pip wheel . -w dist -vvv --no-deps --disable-pip-version-check
 
 mkdir -p final_dist
-python -m auditwheel repair -w final_dist dist/*
+python -m auditwheel repair \
+    --exclude libnvcomp.so.4 \
+    -w final_dist \
+    dist/*
 
 RAPIDS_PY_WHEEL_NAME="${package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 final_dist

dependencies.yaml

@@ -13,6 +13,7 @@ files:
       - cuda
       - cuda_version
       - depends_on_cupy
+      - depends_on_nvcomp
       - docs
       - py_version
       - rapids_build_skbuild
@@ -41,37 +42,38 @@ files:
       - cuda_version
       - docs
       - py_version
-  py_build:
+  py_build_kvikio:
     output: pyproject
     pyproject_dir: python/kvikio
     extras:
       table: build-system
     includes:
       - rapids_build_skbuild
-  py_build_cpp_wheel:
+  py_build_libkvikio:
     output: pyproject
     pyproject_dir: python/libkvikio
     extras:
       table: build-system
     includes:
       - rapids_build_skbuild
-  py_run:
+  py_run_kvikio:
     output: pyproject
     pyproject_dir: python/kvikio
     extras:
       table: project
     includes:
       - depends_on_cupy
+      - depends_on_nvcomp
       - run
-  py_wheel_cpp:
+  py_rapids_build_libkvikio:
     output: pyproject
     pyproject_dir: python/libkvikio
     extras:
       table: tool.rapids-build-backend
       key: requires
     includes:
       - build-universal
-  py_wheel_python:
+  py_rapids_build_kvikio:
     output: pyproject
     pyproject_dir: python/kvikio
     extras:
@@ -199,10 +201,6 @@ dependencies:
             packages:
               - cuda-version=12.5
   cuda:
-    common:
-      - output_types: conda
-        packages:
-          - nvcomp==4.0.1
     specific:
       - output_types: conda
         matrices:
@@ -264,6 +262,25 @@ dependencies:
             packages: &cupy_packages_cu11
               - cupy-cuda11x>=12.0.0
           - {matrix: null, packages: *cupy_packages_cu11}
+  depends_on_nvcomp:
+    common:
+      - output_types: conda
+        packages:
+          - nvcomp==4.0.1
+    specific:
+      - output_types: [requirements, pyproject]
+        matrices:
+          - matrix:
+              cuda: "12.*"
+            packages:
+              - nvidia-nvcomp-cu12==4.0.1
+          - matrix:
+              cuda: "11.*"
+            packages:
+              - nvidia-nvcomp-cu11==4.0.1
+          - matrix:
+            packages:
+              - nvidia-nvcomp==4.0.1
   docs:
     common:
       - output_types: [conda, requirements]

python/kvikio/CMakeLists.txt

@@ -26,6 +26,8 @@ project(
   LANGUAGES CXX CUDA
 )
 
+option(USE_NVCOMP_RUNTIME_WHEEL "Use the nvcomp wheel at runtime instead of the system library" OFF)
+
 # TODO: Should we symlink FindcuFile.cmake into python/cmake? find cuFile
 include(../../cpp/cmake/Modules/FindcuFile.cmake)
 
@@ -40,11 +42,4 @@ add_subdirectory(cmake)
 
 set(cython_lib_dir kvikio)
 
-# It would be better to factor nvcomp out into its own wheel. Until that is available, we vendor it
-# here.
-install_aliased_imported_targets(
-  TARGETS nvcomp::nvcomp nvcomp::nvcomp_gdeflate nvcomp::nvcomp_bitcomp DESTINATION
-  ${cython_lib_dir}/_lib
-)
-
 add_subdirectory(kvikio/_lib)

python/kvikio/kvikio/_lib/CMakeLists.txt

@@ -22,3 +22,13 @@ rapids_cython_create_modules(
   SOURCE_FILES "${cython_modules}"
   LINKED_LIBRARIES kvikio::kvikio nvcomp::nvcomp
 )
+if(USE_NVCOMP_RUNTIME_WHEEL)
+  set(rpaths "$ORIGIN/../../nvidia/nvcomp")
+  foreach(tgt IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
+    set_property(
+      TARGET ${tgt}
+      PROPERTY INSTALL_RPATH ${rpaths}
+      APPEND
+    )
+  endforeach()
+endif()

python/kvikio/pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "cupy-cuda11x>=12.0.0",
     "numcodecs !=0.12.0",
     "numpy>=1.23,<3.0a0",
+    "nvidia-nvcomp==4.0.1",
     "packaging",
     "zarr",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.