From 3732b31e8895ac010a284054dcce254ddec14020 Mon Sep 17 00:00:00 2001
From: Sevag H <shanssian@nvidia.com>
Date: Fri, 6 Jan 2023 18:28:27 -0500
Subject: [PATCH] Build wheels alongside conda CI (#1182)

* Build wheels alongside conda CI

* Remove package-dir param from wheel test/publish

* wheel-publish needs wheel-build

* Adjust versioneer-override param name

* Add smoketest for abbreviated arm PR CI

* make wheel builds depend on `checks`

* Address PR feedback

* Removed some wheel parameters

Co-authored-by: AJ Schmidt <ajschmidt8@users.noreply.github.com>
Co-authored-by: AJ Schmidt <aschmidt@nvidia.com>
---
 .github/workflows/build.yaml | 23 ++++++++++++++++
 .github/workflows/pr.yaml    | 20 ++++++++++++++
 .github/workflows/test.yaml  | 11 ++++++++
 .github/workflows/wheels.yml | 52 ------------------------------------
 ci/wheel_smoke_test.py       |  5 ++++
 5 files changed, 59 insertions(+), 52 deletions(-)
 delete mode 100644 .github/workflows/wheels.yml
 create mode 100644 ci/wheel_smoke_test.py

diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index d7e81bcd3..732aae51a 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -52,3 +52,26 @@ jobs:
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
+  wheel-build:
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-build.yml@main
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      repo: rapidsai/rmm
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+      package-name: rmm
+      package-dir: python
+      skbuild-configure-options: "-DRMM_BUILD_WHEELS=ON"
+  wheel-publish:
+    needs: wheel-build
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-publish.yml@main
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      repo: rapidsai/rmm
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+      package-name: rmm
diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml
index be1c89255..07b234dda 100644
--- a/.github/workflows/pr.yaml
+++ b/.github/workflows/pr.yaml
@@ -17,6 +17,8 @@ jobs:
       - conda-cpp-tests
       - conda-python-build
       - conda-python-tests
+      - wheel-build
+      - wheel-tests
     secrets: inherit
     uses: rapidsai/shared-action-workflows/.github/workflows/pr-builder.yaml@cuda-118
   checks:
@@ -46,3 +48,21 @@ jobs:
     uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-tests.yaml@cuda-118
     with:
       build_type: pull-request
+  wheel-build:
+    needs: checks
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-build.yml@main
+    with:
+      build_type: pull-request
+      package-dir: python
+      package-name: rmm
+      skbuild-configure-options: "-DRMM_BUILD_WHEELS=ON"
+  wheel-tests:
+    needs: wheel-build
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-test.yml@main
+    with:
+      build_type: pull-request
+      package-name: rmm
+      test-unittest: "python -m pytest -v ./python/rmm/tests"
+      test-smoketest: "python ./ci/wheel_smoke_test.py"
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index c0961b587..49dd78071 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -30,3 +30,14 @@ jobs:
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
+  wheel-tests:
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-test.yml@main
+    with:
+      build_type: nightly
+      repo: rapidsai/rmm
+      branch: ${{ inputs.branch }}
+      date: ${{ inputs.date }}
+      sha: ${{ inputs.sha }}
+      package-name: rmm
+      test-unittest: "python -m pytest -v ./python/rmm/tests"
diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml
deleted file mode 100644
index abbe4f682..000000000
--- a/.github/workflows/wheels.yml
+++ /dev/null
@@ -1,52 +0,0 @@
-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/ci/wheel_smoke_test.py b/ci/wheel_smoke_test.py
new file mode 100644
index 000000000..737c0b605
--- /dev/null
+++ b/ci/wheel_smoke_test.py
@@ -0,0 +1,5 @@
+import rmm
+
+if __name__ == "__main__":
+    buf = rmm.DeviceBuffer(size=100)
+    assert buf.size == 100