From cc19d714792064c40258081bfca229c2e3fbbcc8 Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Thu, 4 Apr 2024 12:48:19 -0700 Subject: [PATCH 01/21] run on cuda 12, run examples for pyg --- ci/run_cugraph_pyg_pytests.sh | 6 +++ ci/test_python.sh | 95 +++++++++++++++++------------------ ci/test_wheel_cugraph-pyg.sh | 45 ++++++++++------- 3 files changed, 77 insertions(+), 69 deletions(-) diff --git a/ci/run_cugraph_pyg_pytests.sh b/ci/run_cugraph_pyg_pytests.sh index 47ed6ba000..1e5bca42d8 100755 --- a/ci/run_cugraph_pyg_pytests.sh +++ b/ci/run_cugraph_pyg_pytests.sh @@ -7,3 +7,9 @@ set -euo pipefail cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cugraph-pyg/cugraph_pyg pytest --cache-clear --ignore=tests/mg "$@" . + +# Test examples +for e in examples; do + rapids-logger "running example $e" + python $e +done diff --git a/ci/test_python.sh b/ci/test_python.sh index 3a47d7e149..b323c5922c 100755 --- a/ci/test_python.sh +++ b/ci/test_python.sh @@ -185,58 +185,53 @@ else rapids-logger "skipping cugraph_dgl pytest on CUDA!=11.8" fi -if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then - if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then - rapids-mamba-retry env create --force -f env.yaml -n test_cugraph_pyg - - # Temporarily allow unbound variables for conda activation. - set +u - conda activate test_cugraph_pyg - set -u - - # Will automatically install built dependencies of cuGraph-PyG - rapids-mamba-retry install \ - --channel "${CPP_CHANNEL}" \ - --channel "${PYTHON_CHANNEL}" \ - --channel pytorch \ - --channel nvidia \ - --channel pyg \ - --channel rapidsai-nightly \ - "cugraph-pyg" \ - "pytorch>=2.0,<2.1" \ - "pytorch-cuda=11.8" - - # Install pyg dependencies (which requires pip) - pip install \ - pyg_lib \ - torch_scatter \ - torch_sparse \ - torch_cluster \ - torch_spline_conv \ - -f https://data.pyg.org/whl/torch-2.0.0+cu118.html - - rapids-print-env - - rapids-logger "pytest cugraph_pyg (single GPU)" - # rmat is not tested because of multi-GPU testing - ./ci/run_cugraph_pyg_pytests.sh \ - --junitxml="${RAPIDS_TESTS_DIR}/junit-cugraph-pyg.xml" \ - --cov-config=../../.coveragerc \ - --cov=cugraph_pyg \ - --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/cugraph-pyg-coverage.xml" \ - --cov-report=term - - # Reactivate the test environment back - set +u - conda deactivate - conda activate test - set -u +if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then + rapids-mamba-retry env create --force -f env.yaml -n test_cugraph_pyg + + # Temporarily allow unbound variables for conda activation. + set +u + conda activate test_cugraph_pyg + set -u + + # Will automatically install built dependencies of cuGraph-PyG + rapids-mamba-retry install \ + --channel "${CPP_CHANNEL}" \ + --channel "${PYTHON_CHANNEL}" \ + --channel conda-forge \ + --channel nvidia \ + --channel pyg \ + --channel rapidsai-nightly \ + "cugraph-pyg" \ + "pytorch>=2.0,<2.1" + + # Install pyg dependencies (which requires pip) + pip install \ + pyg_lib \ + torch_scatter \ + torch_sparse \ + torch_cluster \ + torch_spline_conv \ + -f https://data.pyg.org/whl/torch-2.0.0+cu118.html + + rapids-print-env + + rapids-logger "pytest cugraph_pyg (single GPU)" + # rmat is not tested because of multi-GPU testing + ./ci/run_cugraph_pyg_pytests.sh \ + --junitxml="${RAPIDS_TESTS_DIR}/junit-cugraph-pyg.xml" \ + --cov-config=../../.coveragerc \ + --cov=cugraph_pyg \ + --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/cugraph-pyg-coverage.xml" \ + --cov-report=term + + # Reactivate the test environment back + set +u + conda deactivate + conda activate test + set -u - else - rapids-logger "skipping cugraph_pyg pytest on ARM64" - fi else - rapids-logger "skipping cugraph_pyg pytest on CUDA != 11.8" + rapids-logger "skipping cugraph_pyg pytest on ARM64" fi # test cugraph-equivariant diff --git a/ci/test_wheel_cugraph-pyg.sh b/ci/test_wheel_cugraph-pyg.sh index 50cbfb3e1f..d38eae4358 100755 --- a/ci/test_wheel_cugraph-pyg.sh +++ b/ci/test_wheel_cugraph-pyg.sh @@ -25,25 +25,32 @@ python -m pip install $(ls ./dist/${python_package_name}*.whl)[test] export RAPIDS_DATASET_ROOT_DIR="$(realpath datasets)" if [[ "${CUDA_VERSION}" == "11.8.0" ]]; then - rapids-logger "Installing PyTorch and PyG dependencies" PYTORCH_URL="https://download.pytorch.org/whl/cu118" - rapids-retry python -m pip install torch==2.1.0 --index-url ${PYTORCH_URL} - rapids-retry python -m pip install torch-geometric==2.4.0 - rapids-retry python -m pip install \ - pyg_lib \ - torch_scatter \ - torch_sparse \ - torch_cluster \ - torch_spline_conv \ - -f https://data.pyg.org/whl/torch-2.1.0+cu118.html - - rapids-logger "pytest cugraph-pyg (single GPU)" - pushd python/cugraph-pyg/cugraph_pyg - python -m pytest \ - --cache-clear \ - --ignore=tests/mg \ - tests - popd + PYG_URL="https://data.pyg.org/whl/torch-2.1.0+cu118.html" else - rapids-logger "skipping cugraph-pyg wheel test on CUDA!=11.8" + PYTORCH_URL="https://download.pytorch.org/whl/cu121" + PYG_URL="https://data.pyg.org/whl/torch-2.1.0+cu121.html" fi +rapids-logger "Installing PyTorch and PyG dependencies" +rapids-retry python -m pip install torch==2.1.0 --index-url ${PYTORCH_URL} +rapids-retry python -m pip install torch-geometric==2.4.0 +rapids-retry python -m pip install \ + pyg_lib \ + torch_scatter \ + torch_sparse \ + torch_cluster \ + torch_spline_conv \ + -f ${PYG_URL} + +rapids-logger "pytest cugraph-pyg (single GPU)" +pushd python/cugraph-pyg/cugraph_pyg +python -m pytest \ + --cache-clear \ + --ignore=tests/mg \ + tests +# Test examples +for e in examples; do + rapids-logger "running example $e" + python $e +done +popd From 77353f97ed4c03d3c162f3d6af878fe4f6fa5ba5 Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Fri, 5 Apr 2024 09:41:11 -0700 Subject: [PATCH 02/21] fix shell script --- ci/run_cugraph_pyg_pytests.sh | 6 +++--- ci/test_wheel_cugraph-pyg.sh | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/ci/run_cugraph_pyg_pytests.sh b/ci/run_cugraph_pyg_pytests.sh index 1e5bca42d8..818fc8f16c 100755 --- a/ci/run_cugraph_pyg_pytests.sh +++ b/ci/run_cugraph_pyg_pytests.sh @@ -9,7 +9,7 @@ cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cugraph-pyg/cugraph_ pytest --cache-clear --ignore=tests/mg "$@" . # Test examples -for e in examples; do - rapids-logger "running example $e" - python $e +for e in "$(pwd)"/examples/*; do + rapids-logger "running example $e" + python $e done diff --git a/ci/test_wheel_cugraph-pyg.sh b/ci/test_wheel_cugraph-pyg.sh index d38eae4358..d2c886ab9a 100755 --- a/ci/test_wheel_cugraph-pyg.sh +++ b/ci/test_wheel_cugraph-pyg.sh @@ -49,7 +49,7 @@ python -m pytest \ --ignore=tests/mg \ tests # Test examples -for e in examples; do +for e in "$(pwd)"/examples/*; do rapids-logger "running example $e" python $e done From 53dbfe8204e562f693eeb1211ff997fe88274f79 Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Fri, 5 Apr 2024 10:09:59 -0700 Subject: [PATCH 03/21] switch dgl pytorch to conda-forge version --- ci/test_python.sh | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/ci/test_python.sh b/ci/test_python.sh index 0bf48db157..972b8ea745 100755 --- a/ci/test_python.sh +++ b/ci/test_python.sh @@ -154,8 +154,7 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then rapids-mamba-retry install \ --channel "${CPP_CHANNEL}" \ --channel "${PYTHON_CHANNEL}" \ - --channel pytorch \ - --channel pytorch-nightly \ + --channel conda-forge \ --channel dglteam/label/cu118 \ --channel nvidia \ libcugraph \ @@ -164,8 +163,7 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then cugraph \ cugraph-dgl \ 'dgl>=1.1.0.cu*,<=2.0.0.cu*' \ - 'pytorch>=2.0' \ - 'pytorch-cuda>=11.8' + 'pytorch>=2.0' rapids-print-env From e83f848546341240cb70f951beffac83cbbb56c7 Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Tue, 9 Apr 2024 07:41:49 -0700 Subject: [PATCH 04/21] test only python files --- ci/run_cugraph_pyg_pytests.sh | 2 +- ci/test_wheel_cugraph-pyg.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ci/run_cugraph_pyg_pytests.sh b/ci/run_cugraph_pyg_pytests.sh index 818fc8f16c..0acc8aa462 100755 --- a/ci/run_cugraph_pyg_pytests.sh +++ b/ci/run_cugraph_pyg_pytests.sh @@ -9,7 +9,7 @@ cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cugraph-pyg/cugraph_ pytest --cache-clear --ignore=tests/mg "$@" . # Test examples -for e in "$(pwd)"/examples/*; do +for e in "$(pwd)"/examples/*.py; do rapids-logger "running example $e" python $e done diff --git a/ci/test_wheel_cugraph-pyg.sh b/ci/test_wheel_cugraph-pyg.sh index d2c886ab9a..e3c3b221ad 100755 --- a/ci/test_wheel_cugraph-pyg.sh +++ b/ci/test_wheel_cugraph-pyg.sh @@ -49,7 +49,7 @@ python -m pytest \ --ignore=tests/mg \ tests # Test examples -for e in "$(pwd)"/examples/*; do +for e in "$(pwd)"/examples/*.py; do rapids-logger "running example $e" python $e done From e67d637dba6ce6fab0d016ac500ec7f983ef6c80 Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Tue, 9 Apr 2024 07:46:00 -0700 Subject: [PATCH 05/21] update pr --- python/nx-cugraph/README.md | 3 +++ python/nx-cugraph/_nx_cugraph/__init__.py | 1 - 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/python/nx-cugraph/README.md b/python/nx-cugraph/README.md index 77066356a4..df80ae56a9 100644 --- a/python/nx-cugraph/README.md +++ b/python/nx-cugraph/README.md @@ -248,6 +248,9 @@ Below is the list of algorithms that are currently supported in nx-cugraph. ### Other 
+classes
+ └─ function
+     └─ is_negatively_weighted
 convert_matrix
  ├─ from_pandas_edgelist
  └─ from_scipy_sparse_array
diff --git a/python/nx-cugraph/_nx_cugraph/__init__.py b/python/nx-cugraph/_nx_cugraph/__init__.py
index 013f123672..7a073ffb29 100644
--- a/python/nx-cugraph/_nx_cugraph/__init__.py
+++ b/python/nx-cugraph/_nx_cugraph/__init__.py
@@ -207,7 +207,6 @@
         },
         "louvain_communities": {
             "dtype : dtype or None, optional": "The data type (np.float32, np.float64, or None) to use for the edge weights in the algorithm. If None, then dtype is determined by the edge values.",
-            "max_level : int, optional": "Upper limit of the number of macro-iterations (max: 500).",
         },
         "pagerank": {
             "dtype : dtype or None, optional": "The data type (np.float32, np.float64, or None) to use for the edge weights in the algorithm. If None, then dtype is determined by the edge values.",

From 531db1a2cf698314a10eb6be283faf87c779d784 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Tue, 9 Apr 2024 09:59:19 -0700
Subject: [PATCH 06/21] add pytorch to docs script

---
 ci/build_docs.sh | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/ci/build_docs.sh b/ci/build_docs.sh
index 9b72c98f1b..408fea5f79 100755
--- a/ci/build_docs.sh
+++ b/ci/build_docs.sh
@@ -23,6 +23,7 @@ PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
 rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \
   --channel "${PYTHON_CHANNEL}" \
+  --channel conda-forge \
   libcugraph \
   pylibcugraph \
   cugraph \
@@ -31,7 +32,8 @@ rapids-mamba-retry install \
   cugraph-service-client \
   libcugraph_etl \
   pylibcugraphops \
-  pylibwholegraph
+  pylibwholegraph \
+  "pytorch>=2.0"
 
 # This command installs `cugraph-dgl` without its dependencies
 # since this package can currently only run in `11.6` CTK environments

From cf2ce8fa9b812e18db2b8d4bfc6570a38559b047 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Tue, 9 Apr 2024 11:19:41 -0700
Subject: [PATCH 07/21] remove separate dgl install

---
 ci/build_docs.sh | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/ci/build_docs.sh b/ci/build_docs.sh
index 408fea5f79..c27a4ce455 100755
--- a/ci/build_docs.sh
+++ b/ci/build_docs.sh
@@ -28,19 +28,14 @@ rapids-mamba-retry install \
   pylibcugraph \
   cugraph \
   cugraph-pyg \
+  cugraph-dgl \
   cugraph-service-server \
   cugraph-service-client \
   libcugraph_etl \
   pylibcugraphops \
   pylibwholegraph \
   "pytorch>=2.0"
-
-# This command installs `cugraph-dgl` without its dependencies
-# since this package can currently only run in `11.6` CTK environments
-# due to the dependency version specifications in its conda recipe.
-rapids-logger "Install cugraph-dgl"
-rapids-mamba-retry install "${PYTHON_CHANNEL}/linux-64/cugraph-dgl-*.tar.bz2"
-
+  
 export RAPIDS_VERSION="$(rapids-version)"
 export RAPIDS_VERSION_MAJOR_MINOR="$(rapids-version-major-minor)"
 export RAPIDS_VERSION_NUMBER="$RAPIDS_VERSION_MAJOR_MINOR"

From 169b85b1b675d5d60c984bf5c5823b1f59eca672 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Tue, 9 Apr 2024 11:21:34 -0700
Subject: [PATCH 08/21] dgl

---
 ci/build_docs.sh            |  2 +-
 python/nx-cugraph/README.md | 12 ++++++------
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/ci/build_docs.sh b/ci/build_docs.sh
index c27a4ce455..5e4cc27e8e 100755
--- a/ci/build_docs.sh
+++ b/ci/build_docs.sh
@@ -35,7 +35,7 @@ rapids-mamba-retry install \
   pylibcugraphops \
   pylibwholegraph \
   "pytorch>=2.0"
-  
+
 export RAPIDS_VERSION="$(rapids-version)"
 export RAPIDS_VERSION_MAJOR_MINOR="$(rapids-version-major-minor)"
 export RAPIDS_VERSION_NUMBER="$RAPIDS_VERSION_MAJOR_MINOR"
diff --git a/python/nx-cugraph/README.md b/python/nx-cugraph/README.md
index 75b5c1c5aa..df80ae56a9 100644
--- a/python/nx-cugraph/README.md
+++ b/python/nx-cugraph/README.md
@@ -106,16 +106,16 @@ Below is the list of algorithms that are currently supported in nx-cugraph.
  └─ generators
      └─ complete_bipartite_graph
 centrality
- ├─ betweenness
+ ├─ betweenness
  │   ├─ betweenness_centrality
  │   └─ edge_betweenness_centrality
- ├─ degree_alg
+ ├─ degree_alg
  │   ├─ degree_centrality
  │   ├─ in_degree_centrality
  │   └─ out_degree_centrality
- ├─ eigenvector
+ ├─ eigenvector
  │   └─ eigenvector_centrality
- └─ katz
+ └─ katz
      └─ katz_centrality
 cluster
  ├─ average_clustering
@@ -126,12 +126,12 @@ Below is the list of algorithms that are currently supported in nx-cugraph.
  └─ louvain
      └─ louvain_communities
 components
- ├─ connected
+ ├─ connected
  │   ├─ connected_components
  │   ├─ is_connected
  │   ├─ node_connected_component
  │   └─ number_connected_components
- └─ weakly_connected
+ └─ weakly_connected
      ├─ is_weakly_connected
      ├─ number_weakly_connected_components
      └─ weakly_connected_components

From ecf3ea402c28e043be6de2ff793b6b039d513d06 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Tue, 9 Apr 2024 11:23:22 -0700
Subject: [PATCH 09/21] revert nx change

---
 python/nx-cugraph/README.md | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/python/nx-cugraph/README.md b/python/nx-cugraph/README.md
index df80ae56a9..75b5c1c5aa 100644
--- a/python/nx-cugraph/README.md
+++ b/python/nx-cugraph/README.md
@@ -106,16 +106,16 @@ Below is the list of algorithms that are currently supported in nx-cugraph.
  └─ generators
      └─ complete_bipartite_graph
 centrality
- ├─ betweenness
+ ├─ betweenness
  │   ├─ betweenness_centrality
  │   └─ edge_betweenness_centrality
- ├─ degree_alg
+ ├─ degree_alg
  │   ├─ degree_centrality
  │   ├─ in_degree_centrality
  │   └─ out_degree_centrality
- ├─ eigenvector
+ ├─ eigenvector
  │   └─ eigenvector_centrality
- └─ katz
+ └─ katz
      └─ katz_centrality
 cluster
  ├─ average_clustering
@@ -126,12 +126,12 @@ Below is the list of algorithms that are currently supported in nx-cugraph.
  └─ louvain
      └─ louvain_communities
 components
- ├─ connected
+ ├─ connected
  │   ├─ connected_components
  │   ├─ is_connected
  │   ├─ node_connected_component
  │   └─ number_connected_components
- └─ weakly_connected
+ └─ weakly_connected
      ├─ is_weakly_connected
      ├─ number_weakly_connected_components
      └─ weakly_connected_components

From cad86eec75a7b80b7ff565c08c68619ec483d898 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Tue, 9 Apr 2024 18:05:51 -0700
Subject: [PATCH 10/21] use pytorch channel instead of conda-forge

---
 ci/build_docs.sh             | 5 ++---
 ci/test_python.sh            | 4 ++--
 ci/test_wheel_cugraph-pyg.sh | 1 +
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/ci/build_docs.sh b/ci/build_docs.sh
index 5e4cc27e8e..aab8babe1f 100755
--- a/ci/build_docs.sh
+++ b/ci/build_docs.sh
@@ -23,7 +23,7 @@ PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
 rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \
   --channel "${PYTHON_CHANNEL}" \
-  --channel conda-forge \
+  --channel pytorch \
   libcugraph \
   pylibcugraph \
   cugraph \
@@ -33,8 +33,7 @@ rapids-mamba-retry install \
   cugraph-service-client \
   libcugraph_etl \
   pylibcugraphops \
-  pylibwholegraph \
-  "pytorch>=2.0"
+  pylibwholegraph
 
 export RAPIDS_VERSION="$(rapids-version)"
 export RAPIDS_VERSION_MAJOR_MINOR="$(rapids-version-major-minor)"
diff --git a/ci/test_python.sh b/ci/test_python.sh
index 972b8ea745..26cbe35c4f 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -154,7 +154,7 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
     rapids-mamba-retry install \
       --channel "${CPP_CHANNEL}" \
       --channel "${PYTHON_CHANNEL}" \
-      --channel conda-forge \
+      --channel pytorch \
       --channel dglteam/label/cu118 \
       --channel nvidia \
       libcugraph \
@@ -199,7 +199,7 @@ if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
   rapids-mamba-retry install \
     --channel "${CPP_CHANNEL}" \
     --channel "${PYTHON_CHANNEL}" \
-    --channel conda-forge \
+    --channel pytorch \
     --channel nvidia \
     --channel pyg \
     --channel rapidsai-nightly \
diff --git a/ci/test_wheel_cugraph-pyg.sh b/ci/test_wheel_cugraph-pyg.sh
index e3c3b221ad..ad615f0b3f 100755
--- a/ci/test_wheel_cugraph-pyg.sh
+++ b/ci/test_wheel_cugraph-pyg.sh
@@ -35,6 +35,7 @@ rapids-logger "Installing PyTorch and PyG dependencies"
 rapids-retry python -m pip install torch==2.1.0 --index-url ${PYTORCH_URL}
 rapids-retry python -m pip install torch-geometric==2.4.0
 rapids-retry python -m pip install \
+  ogb \
   pyg_lib \
   torch_scatter \
   torch_sparse \

From 6d35be0adc63f3c2e329ec6293d3ba5e10ea34bd Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Wed, 10 Apr 2024 09:20:49 -0700
Subject: [PATCH 11/21] switch back to conda-forge

---
 ci/build_docs.sh  | 12 ++++++++++--
 ci/test_python.sh | 16 ++++++++++++----
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/ci/build_docs.sh b/ci/build_docs.sh
index aab8babe1f..46ade982b5 100755
--- a/ci/build_docs.sh
+++ b/ci/build_docs.sh
@@ -20,10 +20,16 @@ rapids-logger "Downloading artifacts from previous jobs"
 CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
 PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
 
+if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
+  CONDA_CUDA_VERSION="11.8"
+else
+  CONDA_CUDA_VERSION="12.1"
+fi
+
 rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \
   --channel "${PYTHON_CHANNEL}" \
-  --channel pytorch \
+  --channel conda-forge \
   libcugraph \
   pylibcugraph \
   cugraph \
@@ -33,7 +39,9 @@ rapids-mamba-retry install \
   cugraph-service-client \
   libcugraph_etl \
   pylibcugraphops \
-  pylibwholegraph
+  pylibwholegraph \
+  pytorch \
+  "cuda-version=${CONDA_CUDA_VERSION}"
 
 export RAPIDS_VERSION="$(rapids-version)"
 export RAPIDS_VERSION_MAJOR_MINOR="$(rapids-version-major-minor)"
diff --git a/ci/test_python.sh b/ci/test_python.sh
index 26cbe35c4f..ff44fdf6f8 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -154,7 +154,7 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
     rapids-mamba-retry install \
       --channel "${CPP_CHANNEL}" \
       --channel "${PYTHON_CHANNEL}" \
-      --channel pytorch \
+      --channel conda-forge \
       --channel dglteam/label/cu118 \
       --channel nvidia \
       libcugraph \
@@ -163,7 +163,8 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
       cugraph \
       cugraph-dgl \
       'dgl>=1.1.0.cu*,<=2.0.0.cu*' \
-      'pytorch>=2.0'
+      'pytorch>=2.0' \
+      'cuda-version=11.8'
 
     rapids-print-env
 
@@ -195,16 +196,23 @@ if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
   conda activate test_cugraph_pyg
   set -u
 
+  if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
+    CONDA_CUDA_VERSION="11.8"
+  else
+    CONDA_CUDA_VERSION="12.1"
+  fi
+
   # Will automatically install built dependencies of cuGraph-PyG
   rapids-mamba-retry install \
     --channel "${CPP_CHANNEL}" \
     --channel "${PYTHON_CHANNEL}" \
-    --channel pytorch \
+    --channel conda-forge \
     --channel nvidia \
     --channel pyg \
     --channel rapidsai-nightly \
     "cugraph-pyg" \
-    "pytorch>=2.0,<2.1"
+    "pytorch>=2.0,<2.1" \
+    "cuda-version=${CONDA_CUDA_VERSION}"
 
   # Install pyg dependencies (which requires pip)
   pip install \

From 966b068100fcb70300882e265401b0249eceb1a7 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Wed, 10 Apr 2024 12:47:50 -0700
Subject: [PATCH 12/21] update dependencies, test

---
 ci/build_docs.sh  | 5 +++++
 ci/test_python.sh | 7 +++----
 dependencies.yaml | 1 -
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/ci/build_docs.sh b/ci/build_docs.sh
index 46ade982b5..089558bd45 100755
--- a/ci/build_docs.sh
+++ b/ci/build_docs.sh
@@ -22,14 +22,19 @@ PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
 
 if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
   CONDA_CUDA_VERSION="11.8"
+  DGL_CHANNEL="dglteam/label/cu118"
 else
   CONDA_CUDA_VERSION="12.1"
+  DGL_CHANNEL="dglteam/label/cu121"
 fi
 
 rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \
   --channel "${PYTHON_CHANNEL}" \
   --channel conda-forge \
+  --channel pyg \
+  --channel nvidia \
+  --channel "${DGL_CHANNEL}" \
   libcugraph \
   pylibcugraph \
   cugraph \
diff --git a/ci/test_python.sh b/ci/test_python.sh
index ff44fdf6f8..8e729fdf4c 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -207,11 +207,10 @@ if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
     --channel "${CPP_CHANNEL}" \
     --channel "${PYTHON_CHANNEL}" \
     --channel conda-forge \
-    --channel nvidia \
     --channel pyg \
-    --channel rapidsai-nightly \
+    --channel nvidia \
     "cugraph-pyg" \
-    "pytorch>=2.0,<2.1" \
+    "pytorch>=2.0" \
     "cuda-version=${CONDA_CUDA_VERSION}"
 
   # Install pyg dependencies (which requires pip)
@@ -254,7 +253,7 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
     rapids-mamba-retry install \
       --channel "${CPP_CHANNEL}" \
       --channel "${PYTHON_CHANNEL}" \
-      --channel pytorch \
+      --channel conda-forge \
       --channel nvidia \
       cugraph-equivariant
     pip install e3nn==0.5.1
diff --git a/dependencies.yaml b/dependencies.yaml
index 85fb1344ce..3fa19eb238 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -285,7 +285,6 @@ channels:
   - rapidsai
   - rapidsai-nightly
   - dask/label/dev
-  - pytorch
   - pyg
   - dglteam/label/cu118
   - conda-forge

From 22ebdd0a518e705f7f4c0a8329be7b716ec3e542 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Wed, 10 Apr 2024 14:03:59 -0700
Subject: [PATCH 13/21] style

---
 conda/environments/all_cuda-118_arch-x86_64.yaml     |  1 -
 conda/environments/all_cuda-122_arch-x86_64.yaml     |  1 -
 .../cugraph-dgl/conda/cugraph_dgl_dev_cuda-118.yaml  |  1 -
 .../cugraph-pyg/conda/cugraph_pyg_dev_cuda-118.yaml  |  1 -
 python/nx-cugraph/README.md                          | 12 ++++++------
 5 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml
index f656df9e32..f3aeedad64 100644
--- a/conda/environments/all_cuda-118_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-118_arch-x86_64.yaml
@@ -4,7 +4,6 @@ channels:
 - rapidsai
 - rapidsai-nightly
 - dask/label/dev
-- pytorch
 - pyg
 - dglteam/label/cu118
 - conda-forge
diff --git a/conda/environments/all_cuda-122_arch-x86_64.yaml b/conda/environments/all_cuda-122_arch-x86_64.yaml
index 2bd0746e49..4d64ce33a4 100644
--- a/conda/environments/all_cuda-122_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-122_arch-x86_64.yaml
@@ -4,7 +4,6 @@ channels:
 - rapidsai
 - rapidsai-nightly
 - dask/label/dev
-- pytorch
 - pyg
 - dglteam/label/cu118
 - conda-forge
diff --git a/python/cugraph-dgl/conda/cugraph_dgl_dev_cuda-118.yaml b/python/cugraph-dgl/conda/cugraph_dgl_dev_cuda-118.yaml
index feafe57246..bc8bf776a1 100644
--- a/python/cugraph-dgl/conda/cugraph_dgl_dev_cuda-118.yaml
+++ b/python/cugraph-dgl/conda/cugraph_dgl_dev_cuda-118.yaml
@@ -4,7 +4,6 @@ channels:
 - rapidsai
 - rapidsai-nightly
 - dask/label/dev
-- pytorch
 - pyg
 - dglteam/label/cu118
 - conda-forge
diff --git a/python/cugraph-pyg/conda/cugraph_pyg_dev_cuda-118.yaml b/python/cugraph-pyg/conda/cugraph_pyg_dev_cuda-118.yaml
index 81c53044a1..94e9f1decb 100644
--- a/python/cugraph-pyg/conda/cugraph_pyg_dev_cuda-118.yaml
+++ b/python/cugraph-pyg/conda/cugraph_pyg_dev_cuda-118.yaml
@@ -4,7 +4,6 @@ channels:
 - rapidsai
 - rapidsai-nightly
 - dask/label/dev
-- pytorch
 - pyg
 - dglteam/label/cu118
 - conda-forge
diff --git a/python/nx-cugraph/README.md b/python/nx-cugraph/README.md
index 75b5c1c5aa..df80ae56a9 100644
--- a/python/nx-cugraph/README.md
+++ b/python/nx-cugraph/README.md
@@ -106,16 +106,16 @@ Below is the list of algorithms that are currently supported in nx-cugraph.
  └─ generators
      └─ complete_bipartite_graph
 centrality
- ├─ betweenness
+ ├─ betweenness
  │   ├─ betweenness_centrality
  │   └─ edge_betweenness_centrality
- ├─ degree_alg
+ ├─ degree_alg
  │   ├─ degree_centrality
  │   ├─ in_degree_centrality
  │   └─ out_degree_centrality
- ├─ eigenvector
+ ├─ eigenvector
  │   └─ eigenvector_centrality
- └─ katz
+ └─ katz
      └─ katz_centrality
 cluster
  ├─ average_clustering
@@ -126,12 +126,12 @@ Below is the list of algorithms that are currently supported in nx-cugraph.
  └─ louvain
      └─ louvain_communities
 components
- ├─ connected
+ ├─ connected
  │   ├─ connected_components
  │   ├─ is_connected
  │   ├─ node_connected_component
  │   └─ number_connected_components
- └─ weakly_connected
+ └─ weakly_connected
      ├─ is_weakly_connected
      ├─ number_weakly_connected_components
      └─ weakly_connected_components

From 7690b66f754729c2ef0a58728b5e37bd564026f2 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Wed, 10 Apr 2024 14:04:41 -0700
Subject: [PATCH 14/21] revert nx change

---
 python/nx-cugraph/README.md | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/python/nx-cugraph/README.md b/python/nx-cugraph/README.md
index df80ae56a9..75b5c1c5aa 100644
--- a/python/nx-cugraph/README.md
+++ b/python/nx-cugraph/README.md
@@ -106,16 +106,16 @@ Below is the list of algorithms that are currently supported in nx-cugraph.
  └─ generators
      └─ complete_bipartite_graph
 centrality
- ├─ betweenness
+ ├─ betweenness
  │   ├─ betweenness_centrality
  │   └─ edge_betweenness_centrality
- ├─ degree_alg
+ ├─ degree_alg
  │   ├─ degree_centrality
  │   ├─ in_degree_centrality
  │   └─ out_degree_centrality
- ├─ eigenvector
+ ├─ eigenvector
  │   └─ eigenvector_centrality
- └─ katz
+ └─ katz
      └─ katz_centrality
 cluster
  ├─ average_clustering
@@ -126,12 +126,12 @@ Below is the list of algorithms that are currently supported in nx-cugraph.
  └─ louvain
      └─ louvain_communities
 components
- ├─ connected
+ ├─ connected
  │   ├─ connected_components
  │   ├─ is_connected
  │   ├─ node_connected_component
  │   └─ number_connected_components
- └─ weakly_connected
+ └─ weakly_connected
      ├─ is_weakly_connected
      ├─ number_weakly_connected_components
      └─ weakly_connected_components

From 2eac8eb7ed6e347ed2315a0ccced66a752f4b58a Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Wed, 10 Apr 2024 18:40:46 -0700
Subject: [PATCH 15/21] graph sage mg

---
 .../cugraph_pyg/examples/graph_sage_mg.py     | 31 ++++++++++++-------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_mg.py b/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_mg.py
index 4ca573504a..80d683e6c7 100644
--- a/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_mg.py
+++ b/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_mg.py
@@ -17,6 +17,7 @@
 import time
 import argparse
 import gc
+import warnings
 
 import torch
 import numpy as np
@@ -405,7 +406,8 @@ def parse_args():
         "--dask_scheduler_file",
         type=str,
         help="The path to the dask scheduler file",
-        required=True,
+        required=False,
+        default=None,
     )
 
     return parser.parse_args()
@@ -413,19 +415,24 @@ def parse_args():
 
 def main():
     args = parse_args()
+    if args.dask_scheduler_file is None:
+        warnings.warn(
+            "You must provide the dask scheduler file " "to run this example.  Exiting."
+        )
 
-    torch_devices = [int(d) for d in args.torch_devices.split(",")]
-
-    train_args = (
-        torch_devices,
-        args.torch_manager_ip,
-        args.torch_manager_port,
-        args.dask_scheduler_file,
-        args.num_epochs,
-        args.features_on_gpu,
-    )
+    else:
+        torch_devices = [int(d) for d in args.torch_devices.split(",")]
+
+        train_args = (
+            torch_devices,
+            args.torch_manager_ip,
+            args.torch_manager_port,
+            args.dask_scheduler_file,
+            args.num_epochs,
+            args.features_on_gpu,
+        )
 
-    tmp.spawn(train, args=train_args, nprocs=len(torch_devices))
+        tmp.spawn(train, args=train_args, nprocs=len(torch_devices))
 
 
 if __name__ == "__main__":

From 0ad4e0d08732159bad851b9683b02ae89bda61ce Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Wed, 10 Apr 2024 18:47:55 -0700
Subject: [PATCH 16/21] fix graph sage sg bug

---
 ci/test_python.sh                             | 13 ++-
 .../cugraph_pyg/examples/graph_sage_sg.py     | 79 +++++++++----------
 2 files changed, 45 insertions(+), 47 deletions(-)

diff --git a/ci/test_python.sh b/ci/test_python.sh
index 8e729fdf4c..c7847d5615 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -214,13 +214,12 @@ if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
     "cuda-version=${CONDA_CUDA_VERSION}"
 
   # Install pyg dependencies (which requires pip)
-  pip install \
-      pyg_lib \
-      torch_scatter \
-      torch_sparse \
-      torch_cluster \
-      torch_spline_conv \
-    -f https://data.pyg.org/whl/torch-2.0.0+cu118.html
+  # Install from source due to conda-forge pytorch incompatibility
+  pip install --verbose git+https://github.com/pyg-team/pyg-lib.git
+  pip install --verbose torch_scatter
+  pip install --verbose torch_sparse
+  pip install --verbose torch_cluster
+  pip install --verbose torch_spline_conv
 
   rapids-print-env
 
diff --git a/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_sg.py b/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_sg.py
index 9c96a707e4..ff4f101ca1 100644
--- a/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_sg.py
+++ b/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_sg.py
@@ -128,47 +128,46 @@ def train(device: int, features_device: Union[str, int] = "cpu", num_epochs=2) -
         # barrier() cannot do this since the number of ops per rank is
         # different.  It essentially acts like barrier would if the
         # number of ops per rank was the same.
-        for epoch in range(num_epochs):
-            for iter_i, hetero_data in enumerate(cugraph_bulk_loader):
-                num_batches += 1
-                if iter_i % 20 == 0:
-                    print(f"iteration {iter_i}")
-
-                # train
-                train_mask = hetero_data.train_dict["paper"]
-                y_true = hetero_data.y_dict["paper"]
-
-                y_pred = model(
-                    hetero_data.x_dict["paper"].to(device).to(torch.float32),
-                    hetero_data.edge_index_dict[("paper", "cites", "paper")].to(device),
-                    (len(y_true), len(y_true)),
-                )
-
-                y_true = F.one_hot(
-                    y_true[train_mask].to(torch.int64), num_classes=349
-                ).to(torch.float32)
-
-                y_pred = y_pred[train_mask]
-
-                loss = F.cross_entropy(y_pred, y_true)
-
-                optimizer.zero_grad()
-                loss.backward()
-                optimizer.step()
-                total_loss += loss.item()
-
-                del y_true
-                del y_pred
-                del loss
-                del hetero_data
-                gc.collect()
-
-            end_time_train = time.perf_counter_ns()
-            print(
-                f"epoch {epoch} time: "
-                f"{(end_time_train - start_time_train) / 1e9:3.4f} s"
+        for iter_i, hetero_data in enumerate(cugraph_bulk_loader):
+            num_batches += 1
+            if iter_i % 20 == 0:
+                print(f"iteration {iter_i}")
+
+            # train
+            train_mask = hetero_data.train_dict["paper"]
+            y_true = hetero_data.y_dict["paper"]
+
+            y_pred = model(
+                hetero_data.x_dict["paper"].to(device).to(torch.float32),
+                hetero_data.edge_index_dict[("paper", "cites", "paper")].to(device),
+                (len(y_true), len(y_true)),
             )
-            print(f"loss after epoch {epoch}: {total_loss / num_batches}")
+
+            y_true = F.one_hot(
+                y_true[train_mask].to(torch.int64), num_classes=349
+            ).to(torch.float32)
+
+            y_pred = y_pred[train_mask]
+
+            loss = F.cross_entropy(y_pred, y_true)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()
+
+            del y_true
+            del y_pred
+            del loss
+            del hetero_data
+            gc.collect()
+
+        end_time_train = time.perf_counter_ns()
+        print(
+            f"epoch {epoch} time: "
+            f"{(end_time_train - start_time_train) / 1e9:3.4f} s"
+        )
+        print(f"loss after epoch {epoch}: {total_loss / num_batches}")
 
 
 def parse_args():

From 2d22fdf08be1dfdddb5cb21c9f3441f864f40e79 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Wed, 10 Apr 2024 18:50:53 -0700
Subject: [PATCH 17/21] reformat

---
 python/cugraph-pyg/cugraph_pyg/examples/graph_sage_sg.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_sg.py b/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_sg.py
index ff4f101ca1..58a403084d 100644
--- a/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_sg.py
+++ b/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_sg.py
@@ -97,10 +97,13 @@ def train(device: int, features_device: Union[str, int] = "cpu", num_epochs=2) -
 
     num_papers = data[0]["num_nodes_dict"]["paper"]
     train_perc = 0.1
+
     train_nodes = torch.randperm(num_papers)
     train_nodes = train_nodes[: int(train_perc * num_papers)]
+
     train_mask = torch.full((num_papers,), -1, device=device)
     train_mask[train_nodes] = 1
+
     fs.add_data(train_mask, "paper", "train")
 
     cugraph_store = CuGraphStore(fs, G, N)
@@ -143,9 +146,9 @@ def train(device: int, features_device: Union[str, int] = "cpu", num_epochs=2) -
                 (len(y_true), len(y_true)),
             )
 
-            y_true = F.one_hot(
-                y_true[train_mask].to(torch.int64), num_classes=349
-            ).to(torch.float32)
+            y_true = F.one_hot(y_true[train_mask].to(torch.int64), num_classes=349).to(
+                torch.float32
+            )
 
             y_pred = y_pred[train_mask]
 

From a0a8a5063e6728e747c477b2c1349baabf0b2362 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Wed, 10 Apr 2024 19:25:49 -0700
Subject: [PATCH 18/21] switch back to pytorch-cuda for PyG

---
 ci/test_python.sh | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/ci/test_python.sh b/ci/test_python.sh
index c7847d5615..0586efee88 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -198,28 +198,32 @@ if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
 
   if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
     CONDA_CUDA_VERSION="11.8"
+    PYG_URL="https://data.pyg.org/whl/torch-2.1.0+cu118.html"
   else
     CONDA_CUDA_VERSION="12.1"
+    PYG_URL="https://data.pyg.org/whl/torch-2.1.0+cu121.html"
   fi
 
   # Will automatically install built dependencies of cuGraph-PyG
   rapids-mamba-retry install \
     --channel "${CPP_CHANNEL}" \
     --channel "${PYTHON_CHANNEL}" \
-    --channel conda-forge \
+    --channel pytorch \
     --channel pyg \
     --channel nvidia \
     "cugraph-pyg" \
-    "pytorch>=2.0" \
-    "cuda-version=${CONDA_CUDA_VERSION}"
+    "pytorch=2.1.0" \
+    "pytorch-cuda=${CONDA_CUDA_VERSION}"
 
   # Install pyg dependencies (which requires pip)
-  # Install from source due to conda-forge pytorch incompatibility
-  pip install --verbose git+https://github.com/pyg-team/pyg-lib.git
-  pip install --verbose torch_scatter
-  pip install --verbose torch_sparse
-  pip install --verbose torch_cluster
-  pip install --verbose torch_spline_conv
+
+  pip install \
+      pyg_lib \
+      torch_scatter \
+      torch_sparse \
+      torch_cluster \
+      torch_spline_conv \
+    -f PYG_URL
 
   rapids-print-env
 

From 5a96d8340d15b62047ef61eb8fc78ed67077d0b6 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Thu, 11 Apr 2024 13:52:14 -0700
Subject: [PATCH 19/21] disable cuda 12

---
 ci/test_python.sh | 101 ++++++++++++++++++++++++----------------------
 1 file changed, 52 insertions(+), 49 deletions(-)

diff --git a/ci/test_python.sh b/ci/test_python.sh
index 0586efee88..c2a0212dc1 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -188,62 +188,65 @@ else
   rapids-logger "skipping cugraph_dgl pytest on CUDA!=11.8"
 fi
 
-if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
-  rapids-mamba-retry env create --force -f env.yaml -n test_cugraph_pyg
+if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
+  if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
+    rapids-mamba-retry env create --force -f env.yaml -n test_cugraph_pyg
 
-  # Temporarily allow unbound variables for conda activation.
-  set +u
-  conda activate test_cugraph_pyg
-  set -u
+    # Temporarily allow unbound variables for conda activation.
+    set +u
+    conda activate test_cugraph_pyg
+    set -u
 
-  if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
+    # TODO re-enable logic once CUDA 12 is testable
+    #if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
     CONDA_CUDA_VERSION="11.8"
     PYG_URL="https://data.pyg.org/whl/torch-2.1.0+cu118.html"
-  else
-    CONDA_CUDA_VERSION="12.1"
-    PYG_URL="https://data.pyg.org/whl/torch-2.1.0+cu121.html"
-  fi
+    #else
+    #  CONDA_CUDA_VERSION="12.1"
+    #  PYG_URL="https://data.pyg.org/whl/torch-2.1.0+cu121.html"
+    #fi
+
+    # Will automatically install built dependencies of cuGraph-PyG
+    rapids-mamba-retry install \
+      --channel "${CPP_CHANNEL}" \
+      --channel "${PYTHON_CHANNEL}" \
+      --channel pytorch \
+      --channel pyg \
+      --channel nvidia \
+      "cugraph-pyg" \
+      "pytorch=2.1.0" \
+      "pytorch-cuda=${CONDA_CUDA_VERSION}"
+
+    # Install pyg dependencies (which requires pip)
 
-  # Will automatically install built dependencies of cuGraph-PyG
-  rapids-mamba-retry install \
-    --channel "${CPP_CHANNEL}" \
-    --channel "${PYTHON_CHANNEL}" \
-    --channel pytorch \
-    --channel pyg \
-    --channel nvidia \
-    "cugraph-pyg" \
-    "pytorch=2.1.0" \
-    "pytorch-cuda=${CONDA_CUDA_VERSION}"
-
-  # Install pyg dependencies (which requires pip)
-
-  pip install \
-      pyg_lib \
-      torch_scatter \
-      torch_sparse \
-      torch_cluster \
-      torch_spline_conv \
-    -f PYG_URL
-
-  rapids-print-env
-
-  rapids-logger "pytest cugraph_pyg (single GPU)"
-  # rmat is not tested because of multi-GPU testing
-  ./ci/run_cugraph_pyg_pytests.sh \
-    --junitxml="${RAPIDS_TESTS_DIR}/junit-cugraph-pyg.xml" \
-    --cov-config=../../.coveragerc \
-    --cov=cugraph_pyg \
-    --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/cugraph-pyg-coverage.xml" \
-    --cov-report=term
-
-  # Reactivate the test environment back
-  set +u
-  conda deactivate
-  conda activate test
-  set -u
+    pip install \
+        pyg_lib \
+        torch_scatter \
+        torch_sparse \
+        torch_cluster \
+        torch_spline_conv \
+      -f PYG_URL
 
+    rapids-print-env
+
+    rapids-logger "pytest cugraph_pyg (single GPU)"
+    # rmat is not tested because of multi-GPU testing
+    ./ci/run_cugraph_pyg_pytests.sh \
+      --junitxml="${RAPIDS_TESTS_DIR}/junit-cugraph-pyg.xml" \
+      --cov-config=../../.coveragerc \
+      --cov=cugraph_pyg \
+      --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/cugraph-pyg-coverage.xml" \
+      --cov-report=term
+
+    # Reactivate the test environment back
+    set +u
+    conda deactivate
+    conda activate test
+    set -u
+  else
+    rapids-logger "skipping cugraph_pyg pytest on ARM64"
 else
-  rapids-logger "skipping cugraph_pyg pytest on ARM64"
+  rapids-logger "skipping cugraph_pyg pytest on CUDA!=11.8"
 fi
 
 # test cugraph-equivariant

From 4b4be901c7bd27b5a51113cf3158d19514b0fa72 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Thu, 11 Apr 2024 16:29:12 -0700
Subject: [PATCH 20/21] fix syntax error in python script

---
 ci/test_python.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/ci/test_python.sh b/ci/test_python.sh
index c2a0212dc1..908e097550 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -245,6 +245,7 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
     set -u
   else
     rapids-logger "skipping cugraph_pyg pytest on ARM64"
+  fi
 else
   rapids-logger "skipping cugraph_pyg pytest on CUDA!=11.8"
 fi

From 5d58d52233630c32a09058dbe0e4d69b1f516a91 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi 
Date: Thu, 11 Apr 2024 19:27:14 -0700
Subject: [PATCH 21/21] install from pyg url

---
 ci/test_python.sh | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/ci/test_python.sh b/ci/test_python.sh
index 908e097550..dfcf6816de 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -219,13 +219,14 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
 
     # Install pyg dependencies (which requires pip)
 
+    pip install ogb
     pip install \
         pyg_lib \
         torch_scatter \
         torch_sparse \
         torch_cluster \
         torch_spline_conv \
-      -f PYG_URL
+      -f ${PYG_URL}
 
     rapids-print-env