From 802e46cbb933368b04879f796cedaa512b6e6b83 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 21 Oct 2024 14:00:48 -0500
Subject: [PATCH] declare torch testing dependency for cugraph-dgl, combine pip
 installs again

---
 ci/test_wheel_cugraph-dgl.sh      | 10 +++-------
 ci/test_wheel_cugraph-pyg.sh      |  7 +------
 ci/test_wheel_pylibwholegraph.sh  |  9 ++-------
 dependencies.yaml                 |  3 ++-
 python/cugraph-pyg/pyproject.toml |  2 ++
 5 files changed, 10 insertions(+), 21 deletions(-)

diff --git a/ci/test_wheel_cugraph-dgl.sh b/ci/test_wheel_cugraph-dgl.sh
index 07db3b2..9c47a3a 100755
--- a/ci/test_wheel_cugraph-dgl.sh
+++ b/ci/test_wheel_cugraph-dgl.sh
@@ -26,15 +26,11 @@ DGL_URL="https://data.dgl.ai/wheels/torch-2.3/cu${PYTORCH_CUDA_VER}/repo.html"
 # echo to expand wildcard before adding `[extra]` requires for pip
 python -m pip install \
     -v \
+    --extra-index-url "${PYTORCH_URL}" \
     --find-links "${DGL_URL}" \
     "$(echo ./local-deps/pylibwholegraph_${RAPIDS_PY_CUDA_SUFFIX}*.whl)" \
     "$(echo ./dist/cugraph_dgl_${RAPIDS_PY_CUDA_SUFFIX}*.whl)[test]" \
-    'dgl==2.4.0'
-
-# install torch separately, to be sure we get a CUDA build
-python -m pip install \
-  --index-url "${PYTORCH_URL}" \
-  -v \
-  'torch>=2.0,<2.4.0a0'
+    'dgl==2.4.0' \
+    'torch>=2.0,<2.4.0a0'
 
 python -m pytest python/cugraph-dgl/cugraph_dgl/tests
diff --git a/ci/test_wheel_cugraph-pyg.sh b/ci/test_wheel_cugraph-pyg.sh
index dd9c9f3..25d25b7 100755
--- a/ci/test_wheel_cugraph-pyg.sh
+++ b/ci/test_wheel_cugraph-pyg.sh
@@ -24,6 +24,7 @@ fi
 # echo to expand wildcard before adding `[extra]` requires for pip
 python -m pip install \
     -v \
+    --index-url "${PYTORCH_URL}" \
     --find-links "${PYG_URL}" \
     "$(echo ./local-deps/pylibwholegraph_${RAPIDS_PY_CUDA_SUFFIX}*.whl)" \
     "$(echo ./dist/cugraph_pyg_${RAPIDS_PY_CUDA_SUFFIX}*.whl)[test]" \
@@ -33,12 +34,6 @@ python -m pip install \
     'torch_scatter' \
     'torch_sparse'
 
-# install torch separately, to be sure we get a CUDA build
-python -m pip install \
-  --index-url "${PYTORCH_URL}" \
-  -v \
-  'torch>=2.0,<2.4.0a0'
-
 # RAPIDS_DATASET_ROOT_DIR is used by test scripts
 export RAPIDS_DATASET_ROOT_DIR="$(realpath datasets)"
 
diff --git a/ci/test_wheel_pylibwholegraph.sh b/ci/test_wheel_pylibwholegraph.sh
index 823ae18..6d4b7d9 100755
--- a/ci/test_wheel_pylibwholegraph.sh
+++ b/ci/test_wheel_pylibwholegraph.sh
@@ -25,13 +25,8 @@ mkdir -p "${RAPIDS_TESTS_DIR}" "${RAPIDS_COVERAGE_DIR}"
 rapids-logger "Installing Packages"
 rapids-retry python -m pip install \
     --extra-index-url ${INDEX_URL} \
-    "$(echo ./dist/pylibwholegraph*.whl)[test]"
-
-# install torch separately, to be sure we get a CUDA build
-python -m pip install \
-  --index-url "${INDEX_URL}" \
-  -v \
-  'torch>=2.0,<2.4.0a0'
+    "$(echo ./dist/pylibwholegraph*.whl)[test]" \
+    'torch>=2.0,<2.4.0a0'
 
 rapids-logger "pytest pylibwholegraph"
 cd python/pylibwholegraph/pylibwholegraph/tests
diff --git a/dependencies.yaml b/dependencies.yaml
index a0f9e2d..0fe93b9 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -150,8 +150,9 @@ files:
       table: project.optional-dependencies
       key: test
     includes:
-      - test_python_common
       - depends_on_pylibwholegraph
+      - depends_on_pytorch
+      - test_python_common
 
   cugraph_dgl_dev:
     matrix:
diff --git a/python/cugraph-pyg/pyproject.toml b/python/cugraph-pyg/pyproject.toml
index 0ace928..d5c9786 100644
--- a/python/cugraph-pyg/pyproject.toml
+++ b/python/cugraph-pyg/pyproject.toml
@@ -47,6 +47,8 @@ test = [
     "pytest-benchmark",
     "pytest-cov",
     "pytest-xdist",
+    "tensordict>=0.1.2",
+    "torch>=2.0,<2.4.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [tool.setuptools]