Build wheels alongside conda CI (#1116)

This PR adds pip wheel CI to the Conda CI, instead of having them work separately.

Authors:
  - Sevag H (https://github.com/sevagh)
  - AJ Schmidt (https://github.com/ajschmidt8)
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - AJ Schmidt (https://github.com/ajschmidt8)
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: #1116

.github/workflows/build.yaml

@@ -52,3 +52,46 @@ jobs:
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
+  wheel-build-pylibraft:
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-build.yml@main
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+      package-name: pylibraft
+      package-dir: python/pylibraft
+      skbuild-configure-options: "-DRAFT_BUILD_WHEELS=ON -DDETECT_CONDA_ENV=OFF -DFIND_RAFT_CPP=OFF"
+  wheel-publish-pylibraft:
+    needs: wheel-build-pylibraft
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-publish.yml@main
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+      package-name: pylibraft
+  wheel-build-raft-dask:
+    needs: wheel-publish-pylibraft
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-build.yml@main
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+      package-name: raft_dask
+      package-dir: python/raft-dask
+      skbuild-configure-options: "-DRAFT_BUILD_WHEELS=ON -DDETECT_CONDA_ENV=OFF -DFIND_RAFT_CPP=OFF"
+  wheel-publish-raft-dask:
+    needs: wheel-build-raft-dask
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-publish.yml@main
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+      package-name: raft_dask

.github/workflows/pr.yaml

@@ -17,6 +17,10 @@ jobs:
       - conda-cpp-tests
       - conda-python-build
       - conda-python-tests
+      - wheel-build-pylibraft
+      - wheel-tests-pylibraft
+      - wheel-build-raft-dask
+      - wheel-tests-raft-dask
     secrets: inherit
     uses: rapidsai/shared-action-workflows/.github/workflows/pr-builder.yaml@main
   checks:
@@ -47,3 +51,45 @@ jobs:
     uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-tests.yaml@main
     with:
       build_type: pull-request
+  wheel-build-pylibraft:
+    needs: checks
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-build.yml@main
+    with:
+      build_type: pull-request
+      package-name: pylibraft
+      package-dir: python/pylibraft
+      skbuild-configure-options: "-DRAFT_BUILD_WHEELS=ON -DDETECT_CONDA_ENV=OFF -DFIND_RAFT_CPP=OFF"
+  wheel-tests-pylibraft:
+    needs: wheel-build-pylibraft
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-test.yml@main
+    with:
+      build_type: pull-request
+      package-name: pylibraft
+      test-before-amd64: "pip install cupy-cuda11x"
+      test-before-arm64: "pip install cupy-cuda11x -f https://pip.cupy.dev/aarch64"
+      test-unittest: "python -m pytest -v ./python/pylibraft/pylibraft/test"
+      test-smoketest: "python ./ci/wheel_smoke_test_pylibraft.py"
+  wheel-build-raft-dask:
+    needs: wheel-tests-pylibraft
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-build.yml@main
+    with:
+      build_type: pull-request
+      package-name: raft_dask
+      package-dir: python/raft-dask
+      before-wheel: "RAPIDS_PY_WHEEL_NAME=pylibraft_cu11 rapids-download-wheels-from-s3 ./local-pylibraft-dep"
+      cibw-before-build: "pip install --no-deps ./local-pylibraft-dep/pylibraft*.whl"
+      skbuild-configure-options: "-DRAFT_BUILD_WHEELS=ON -DDETECT_CONDA_ENV=OFF -DFIND_RAFT_CPP=OFF"
+  wheel-tests-raft-dask:
+    needs: wheel-build-raft-dask
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-test.yml@main
+    with:
+      build_type: pull-request
+      package-name: raft_dask
+      test-before-amd64: "RAPIDS_PY_WHEEL_NAME=pylibraft_cu11 rapids-download-wheels-from-s3 ./local-pylibraft-dep && pip install --no-deps ./local-pylibraft-dep/pylibraft*.whl && pip install git+https://github.com/dask/dask.git@main git+https://github.com/dask/distributed.git@main git+https://github.com/rapidsai/[email protected]"
+      test-before-arm64: "RAPIDS_PY_WHEEL_NAME=pylibraft_cu11 rapids-download-wheels-from-s3 ./local-pylibraft-dep && pip install --no-deps ./local-pylibraft-dep/pylibraft*.whl && pip install git+https://github.com/dask/dask.git@main git+https://github.com/dask/distributed.git@main git+https://github.com/rapidsai/[email protected]"
+      test-unittest: "python -m pytest -v ./python/raft-dask/raft_dask/test"
+      test-smoketest: "python ./ci/wheel_smoke_test_raft_dask.py"

.github/workflows/test.yaml

@@ -30,3 +30,27 @@ jobs:
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
+  wheel-tests-pylibraft:
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-test.yml@main
+    with:
+      build_type: nightly
+      branch: ${{ inputs.branch }}
+      date: ${{ inputs.date }}
+      sha: ${{ inputs.sha }}
+      package-name: pylibraft
+      test-before-amd64: "pip install cupy-cuda11x"
+      test-before-arm64: "pip install cupy-cuda11x -f https://pip.cupy.dev/aarch64"
+      test-unittest: "python -m pytest -v ./python/pylibraft/pylibraft/test"
+  wheel-tests-raft-dask:
+    secrets: inherit
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-test.yml@main
+    with:
+      build_type: nightly
+      branch: ${{ inputs.branch }}
+      date: ${{ inputs.date }}
+      sha: ${{ inputs.sha }}
+      package-name: raft_dask
+      test-before-amd64: "pip install git+https://github.com/dask/dask.git@main git+https://github.com/dask/distributed.git@main git+https://github.com/rapidsai/[email protected]"
+      test-before-arm64: "pip install git+https://github.com/dask/dask.git@main git+https://github.com/dask/distributed.git@main git+https://github.com/rapidsai/[email protected]"
+      test-unittest: "python -m pytest -v ./python/raft-dask/raft_dask/test"

ci/release/update-version.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2020-2022, NVIDIA CORPORATION.
+# Copyright (c) 2020-2023, NVIDIA CORPORATION.
 ########################
 # RAFT Version Updater #
 ########################
@@ -53,3 +53,6 @@ done
 sed_runner "s/export UCX_PY_VERSION=.*/export UCX_PY_VERSION='${NEXT_UCX_PY_VERSION}'/g" ci/gpu/build.sh
 sed_runner "s/export UCX_PY_VERSION=.*/export UCX_PY_VERSION='${NEXT_UCX_PY_VERSION}'/g" ci/cpu/build.sh
 sed_runner "/^ucx_py_version:$/ {n;s/.*/  - \"${NEXT_UCX_PY_VERSION}\"/}" conda/recipes/raft-dask/conda_build_config.yaml
+
+# Wheel builds install dask-cuda from source, update its branch
+sed_runner "s/dask-cuda.git@branch-[^\"\s]\+/dask-cuda.git@branch-${NEXT_SHORT_TAG}/g" .github/workflows/*.yaml

ci/wheel_smoke_test_pylibraft.py

@@ -0,0 +1,38 @@
+import numpy as np
+from scipy.spatial.distance import cdist
+
+from pylibraft.common import Handle, Stream, device_ndarray
+from pylibraft.distance import pairwise_distance
+
+
+if __name__ == "__main__":
+    metric = "euclidean"
+    n_rows = 1337
+    n_cols = 1337
+
+    input1 = np.random.random_sample((n_rows, n_cols))
+    input1 = np.asarray(input1, order="C").astype(np.float64)
+
+    output = np.zeros((n_rows, n_rows), dtype=np.float64)
+
+    expected = cdist(input1, input1, metric)
+
+    expected[expected <= 1e-5] = 0.0
+
+    input1_device = device_ndarray(input1)
+    output_device = None
+
+    s2 = Stream()
+    handle = Handle(stream=s2)
+    ret_output = pairwise_distance(
+        input1_device, input1_device, output_device, metric, handle=handle
+    )
+    handle.sync()
+
+    output_device = ret_output
+
+    actual = output_device.copy_to_host()
+
+    actual[actual <= 1e-5] = 0.0
+
+    assert np.allclose(expected, actual, rtol=1e-4)

ci/wheel_smoke_test_raft_dask.py

@@ -0,0 +1,92 @@
+from dask.distributed import Client, wait
+from dask_cuda import LocalCUDACluster, initialize
+
+from raft_dask.common import (
+    Comms,
+    local_handle,
+    perform_test_comm_split,
+    perform_test_comms_allgather,
+    perform_test_comms_allreduce,
+    perform_test_comms_bcast,
+    perform_test_comms_device_multicast_sendrecv,
+    perform_test_comms_device_send_or_recv,
+    perform_test_comms_device_sendrecv,
+    perform_test_comms_gather,
+    perform_test_comms_gatherv,
+    perform_test_comms_reduce,
+    perform_test_comms_reducescatter,
+    perform_test_comms_send_recv,
+)
+
+import os
+os.environ["UCX_LOG_LEVEL"] = "error"
+
+
+def func_test_send_recv(sessionId, n_trials):
+    handle = local_handle(sessionId)
+    return perform_test_comms_send_recv(handle, n_trials)
+
+
+def func_test_collective(func, sessionId, root):
+    handle = local_handle(sessionId)
+    return func(handle, root)
+
+
+if __name__ == "__main__":
+    # initial setup
+    cluster = LocalCUDACluster(protocol="tcp", scheduler_port=0)
+    client = Client(cluster)
+
+    n_trials = 5
+    root_location = "client"
+
+    # p2p test for ucx
+    cb = Comms(comms_p2p=True, verbose=True)
+    cb.init()
+
+    dfs = [
+        client.submit(
+            func_test_send_recv,
+            cb.sessionId,
+            n_trials,
+            pure=False,
+            workers=[w],
+        )
+        for w in cb.worker_addresses
+    ]
+
+    wait(dfs, timeout=5)
+
+    assert list(map(lambda x: x.result(), dfs))
+
+    cb.destroy()
+
+    # collectives test for nccl
+
+    cb = Comms(
+        verbose=True, client=client, nccl_root_location=root_location
+    )
+    cb.init()
+
+    for k, v in cb.worker_info(cb.worker_addresses).items():
+
+        dfs = [
+            client.submit(
+                func_test_collective,
+                perform_test_comms_allgather,
+                cb.sessionId,
+                v["rank"],
+                pure=False,
+                workers=[w],
+            )
+            for w in cb.worker_addresses
+        ]
+        wait(dfs, timeout=5)
+
+        assert all([x.result() for x in dfs])
+
+    cb.destroy()
+
+    # final client and cluster teardown
+    client.close()
+    cluster.close()

python/raft-dask/CMakeLists.txt

@@ -1,5 +1,5 @@
 # =============================================================================
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 # in compliance with the License. You may obtain a copy of the License at
@@ -69,6 +69,8 @@ if(NOT raft_FOUND)
   endif()
 
   add_subdirectory(../../cpp raft-cpp ${_exclude_from_all})
+  list(APPEND CMAKE_MODULE_PATH ${CMAKE_BINARY_DIR}/cmake/find_modules)
+  find_package(NCCL REQUIRED)
 endif()
 
 include(rapids-cython)