From 4da70d53986a66fa11588cab7ab20fe26d172abe Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Tue, 15 Nov 2022 11:50:09 -0500
Subject: [PATCH] Add wheel builds to rmm (#1148)

This PR enables building wheels in RMM.

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Sevag H (https://github.com/sevagh)
  - Paul Taylor (https://github.com/trxcllnt)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Sevag H (https://github.com/sevagh)

URL: https://github.com/rapidsai/rmm/pull/1148
---
 .github/workflows/wheels.yml | 52 ++++++++++++++++++++++++++++++++++++
 .gitignore                   |  3 +++
 python/CMakeLists.txt        | 15 ++++++++---
 python/LICENSE               |  1 +
 python/setup.cfg             |  9 -------
 python/setup.py              | 14 +++++++---
 6 files changed, 77 insertions(+), 17 deletions(-)
 create mode 100644 .github/workflows/wheels.yml
 create mode 120000 python/LICENSE

diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml
new file mode 100644
index 000000000..abbe4f682
--- /dev/null
+++ b/.github/workflows/wheels.yml
@@ -0,0 +1,52 @@
+name: RMM wheels
+
+on:
+  workflow_call:
+    inputs:
+      versioneer-override:
+        type: string
+        default: ''
+      build-tag:
+        type: string
+        default: ''
+      branch:
+        required: true
+        type: string
+      date:
+        required: true
+        type: string
+      sha:
+        required: true
+        type: string
+      build-type:
+        type: string
+        default: nightly
+
+concurrency:
+  # Hardcode rmm rather than using github.repository to support calling this
+  # from other workflows that build multiple wheels.
+  group: "rmm-${{ github.workflow }}-${{ github.ref }}"
+  cancel-in-progress: true
+
+jobs:
+  rmm-wheels:
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux.yml@main
+    with:
+      repo: rapidsai/rmm
+
+      build-type: ${{ inputs.build-type }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+
+      package-dir: python
+      package-name: rmm
+
+      python-package-versioneer-override: ${{ inputs.versioneer-override }}
+      python-package-build-tag: ${{ inputs.build-tag }}
+
+      skbuild-configure-options: "-DRMM_BUILD_WHEELS=ON"
+
+      test-extras: test
+      test-unittest: "python -m pytest -v ./python/rmm/tests"
+    secrets: inherit
diff --git a/.gitignore b/.gitignore
index 2c603209d..1ab57e4d4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -150,3 +150,6 @@ ENV/
 
 # RMM log files
 rmm_log.txt
+
+# cibuildwheel
+/wheelhouse
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
index 3803d718a..352ca5b75 100644
--- a/python/CMakeLists.txt
+++ b/python/CMakeLists.txt
@@ -31,6 +31,7 @@ project(
 
 option(FIND_RMM_CPP "Search for existing RMM C++ installations before defaulting to local files"
        OFF)
+option(RMM_BUILD_WHEELS "Whether this build is generating a Python wheel." OFF)
 
 # If the user requested it we attempt to find RMM.
 if(FIND_RMM_CPP)
@@ -40,12 +41,18 @@ else()
 endif()
 
 if(NOT rmm_FOUND)
-  # TODO: This will not be necessary once we upgrade to CMake 3.22, which will pull in the required
-  # languages for the C++ project even if this project does not require those languges.
-  enable_language(CUDA)
   set(BUILD_TESTS OFF)
   set(BUILD_BENCHMARKS OFF)
-  add_subdirectory(../ rmm-cpp)
+
+  set(_exclude_from_all "")
+  if(RMM_BUILD_WHEELS)
+    # Statically link dependencies if building wheels
+    set(CUDA_STATIC_RUNTIME ON)
+    # Don't install the rmm C++ targets into wheels
+    set(_exclude_from_all EXCLUDE_FROM_ALL)
+  endif()
+
+  add_subdirectory(../ rmm-cpp ${_exclude_from_all})
 endif()
 
 include(rapids-cython)
diff --git a/python/LICENSE b/python/LICENSE
new file mode 120000
index 000000000..ea5b60640
--- /dev/null
+++ b/python/LICENSE
@@ -0,0 +1 @@
+../LICENSE
\ No newline at end of file
diff --git a/python/setup.cfg b/python/setup.cfg
index 06b44e59f..9e0227b94 100644
--- a/python/setup.cfg
+++ b/python/setup.cfg
@@ -47,12 +47,3 @@ skip=
     build
     dist
     __init__.py
-
-
-[options]
-packages = find:
-install_requires =
-    cuda-python>=11.5,<11.7.1
-    numpy>=1.19
-    numba>=0.49
-python_requires = >=3.8
diff --git a/python/setup.py b/python/setup.py
index 2f91d200b..7b8f2b4b0 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -1,13 +1,17 @@
 # Copyright (c) 2019-2022, NVIDIA CORPORATION.
 
+import os
+
 from setuptools import find_packages
 from skbuild import setup
 
 import versioneer
 
 setup(
-    name="rmm",
-    version=versioneer.get_version(),
+    name="rmm" + os.getenv("RAPIDS_PY_WHEEL_CUDA_SUFFIX", default=""),
+    version=os.getenv(
+        "RAPIDS_PY_WHEEL_VERSIONEER_OVERRIDE", default=versioneer.get_version()
+    ),
     description="rmm - RAPIDS Memory Manager",
     url="https://github.com/rapidsai/rmm",
     author="NVIDIA Corporation",
@@ -22,8 +26,10 @@
         "Programming Language :: Python :: 3.9",
     ],
     # Include the separately-compiled shared library
-    extras_require={"test": ["pytest", "pytest-xdist"]},
+    extras_require={"test": ["pytest"]},
     packages=find_packages(include=["rmm", "rmm.*"]),
+    include_package_data=True,
+    python_requires=">=3.8",
     package_data={
         # Note: A dict comprehension with an explicit copy is necessary (rather
         # than something simpler like a dict.fromkeys) because otherwise every
@@ -35,7 +41,7 @@
     },
     cmdclass=versioneer.get_cmdclass(),
     install_requires=[
-        "cuda-python>=11.5,<11.7.1",
+        "cuda-python>=11.7.1,<12.0",
         "numpy>=1.19",
         "numba>=0.49",
     ],