From 09efe2a6c9a64493911f1ace6e4403923ce9d443 Mon Sep 17 00:00:00 2001
From: Tingyu Wang <tingyuw@nvidia.com>
Date: Tue, 2 Jan 2024 18:17:22 -0500
Subject: [PATCH] build wheels

---
 .github/workflows/pr.yaml     | 16 ++++++++++++++++
 ci/build_wheel.sh             |  6 ++++--
 ci/build_wheel_cugraph-dgl.sh |  6 ++++++
 ci/build_wheel_cugraph-pyg.sh |  6 ++++++
 4 files changed, 32 insertions(+), 2 deletions(-)
 create mode 100755 ci/build_wheel_cugraph-dgl.sh
 create mode 100755 ci/build_wheel_cugraph-pyg.sh

diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml
index 82c71efffdb..62d4b9439eb 100644
--- a/.github/workflows/pr.yaml
+++ b/.github/workflows/pr.yaml
@@ -25,6 +25,8 @@ jobs:
       - wheel-tests-cugraph
       - wheel-build-nx-cugraph
       - wheel-tests-nx-cugraph
+      - wheel-build-cugraph-dgl
+      - wheel-build-cugraph-pyg
       - devcontainer
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-24.02
@@ -127,6 +129,20 @@ jobs:
     with:
       build_type: pull-request
       script: ci/test_wheel_nx-cugraph.sh
+  wheel-build-cugraph-dgl:
+    needs: wheel-tests-cugraph
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-24.02
+    with:
+      build_type: pull-request
+      script: ci/build_wheel_cugraph-dgl.sh
+  wheel-build-cugraph-pyg:
+    needs: wheel-tests-cugraph
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-24.02
+    with:
+      build_type: pull-request
+      script: ci/build_wheel_cugraph-pyg.sh
   devcontainer:
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-24.02
diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index 163520ea1da..10e5f8e9709 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 
 set -euo pipefail
 
@@ -55,7 +55,9 @@ cd "${package_dir}"
 python -m pip wheel . -w dist -vvv --no-deps --disable-pip-version-check
 
 # pure-python packages should not have auditwheel run on them.
-if [[ ${package_name} == "nx-cugraph" ]]; then
+if [[ ${package_name} == "nx-cugraph" ]] || \
+   [[ ${package_name} == "cugraph-dgl" ]] || \
+   [[ ${package_name} == "cugraph-pyg" ]]; then
     RAPIDS_PY_WHEEL_NAME="${package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 dist
 else
     mkdir -p final_dist
diff --git a/ci/build_wheel_cugraph-dgl.sh b/ci/build_wheel_cugraph-dgl.sh
new file mode 100755
index 00000000000..d62f810cba4
--- /dev/null
+++ b/ci/build_wheel_cugraph-dgl.sh
@@ -0,0 +1,6 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+./ci/build_wheel.sh cugraph-dgl python/cugraph-dgl
diff --git a/ci/build_wheel_cugraph-pyg.sh b/ci/build_wheel_cugraph-pyg.sh
new file mode 100755
index 00000000000..97baa243f73
--- /dev/null
+++ b/ci/build_wheel_cugraph-pyg.sh
@@ -0,0 +1,6 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+./ci/build_wheel.sh cugraph-pyg python/cugraph-pyg