From 58806fd475da1a5311afd03cf6c60bd23088e355 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Fri, 11 Oct 2024 15:53:39 -0500
Subject: [PATCH] pytorch pins, some CI script cleanup

---
 ci/test_notebooks.sh                |  6 +++++-
 ci/test_python.sh                   | 20 +++++++++++---------
 ci/test_wheel.sh                    |  1 -
 ci/test_wheel_cugraph-dgl.sh        |  1 -
 ci/test_wheel_cugraph-pyg.sh        |  1 -
 conda/recipes/cugraph-dgl/meta.yaml |  4 ++--
 conda/recipes/cugraph-pyg/meta.yaml |  4 ++--
 7 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/ci/test_notebooks.sh b/ci/test_notebooks.sh
index 31ec560..59b60e4 100755
--- a/ci/test_notebooks.sh
+++ b/ci/test_notebooks.sh
@@ -5,6 +5,8 @@ set -Eeuo pipefail
 
 . /opt/conda/etc/profile.d/conda.sh
 
+RAPIDS_VERSION="$(rapids-version)"
+
 rapids-logger "Generate notebook testing dependencies"
 rapids-dependency-file-generator \
   --output conda \
@@ -27,7 +29,9 @@ PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
 rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \
   --channel "${PYTHON_CHANNEL}" \
-  libcugraph pylibcugraph cugraph
+  "libcugraph=${RAPIDS_VERSION}" \
+  "pylibcugraph=${RAPIDS_VERSION}" \
+  "cugraph=${RAPIDS_VERSION}"
 
 NBTEST="$(realpath "$(dirname "$0")/utils/nbtest.sh")"
 NOTEBOOK_LIST="$(realpath "$(dirname "$0")/notebook_list.py")"
diff --git a/ci/test_python.sh b/ci/test_python.sh
index 7e1be3a..9fec454 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -8,6 +8,8 @@ cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../
 
 . /opt/conda/etc/profile.d/conda.sh
 
+RAPIDS_VERSION="$(rapids-version)"
+
 rapids-logger "Generate Python testing dependencies"
 rapids-dependency-file-generator \
   --output conda \
@@ -59,10 +61,10 @@ if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
     --channel conda-forge \
     --channel dglteam/label/th23_cu118 \
     --channel nvidia \
-    "pylibwholegraph=${RAPIDS_VERSION_MAJOR_MINOR}.*" \
-    "pylibcugraphops=${RAPIDS_VERSION_MAJOR_MINOR}.*" \
-    "cugraph=${RAPIDS_VERSION_MAJOR_MINOR}.*" \
-    "cugraph-dgl=${RAPIDS_VERSION_MAJOR_MINOR}.*" \
+    "pylibwholegraph=${RAPIDS_VERSION}.*" \
+    "pylibcugraphops=${RAPIDS_VERSION}.*" \
+    "cugraph=${RAPIDS_VERSION}.*" \
+    "cugraph-dgl=${RAPIDS_VERSION}.*" \
     'pytorch::pytorch>=2.3,<2.4' \
     'cuda-version=11.8' \
     "ogb"
@@ -111,10 +113,10 @@ if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
     --channel "${CPP_CHANNEL}" \
     --channel "${PYTHON_CHANNEL}" \
     --channel pytorch \
-    "pylibwholegraph=${RAPIDS_VERSION_MAJOR_MINOR}.*" \
-    "pylibcugraphops=${RAPIDS_VERSION_MAJOR_MINOR}.*" \
-    "cugraph=${RAPIDS_VERSION_MAJOR_MINOR}.*" \
-    "cugraph-pyg=${RAPIDS_VERSION_MAJOR_MINOR}.*" \
+    "pylibwholegraph=${RAPIDS_VERSION}.*" \
+    "pylibcugraphops=${RAPIDS_VERSION}.*" \
+    "cugraph=${RAPIDS_VERSION}.*" \
+    "cugraph-pyg=${RAPIDS_VERSION}.*" \
     "pytorch::pytorch>=2.3,<2.4" \
     "ogb"
 
@@ -160,7 +162,7 @@ if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
     --channel "${CPP_CHANNEL}" \
     --channel "${PYTHON_CHANNEL}" \
     --channel pytorch \
-    "pylibwholegraph=${RAPIDS_VERSION_MAJOR_MINOR}.*" \
+    "pylibwholegraph=${RAPIDS_VERSION}.*" \
     "pytorch::pytorch>=2.3,<2.4" \
     "ogb"
 
diff --git a/ci/test_wheel.sh b/ci/test_wheel.sh
index 39d2c1f..25a59fd 100755
--- a/ci/test_wheel.sh
+++ b/ci/test_wheel.sh
@@ -4,7 +4,6 @@
 set -eoxu pipefail
 
 package_name=$1
-package_dir=$2
 
 python_package_name=$(echo ${package_name}|sed 's/-/_/g')
 
diff --git a/ci/test_wheel_cugraph-dgl.sh b/ci/test_wheel_cugraph-dgl.sh
index 564b46c..1dbf44f 100755
--- a/ci/test_wheel_cugraph-dgl.sh
+++ b/ci/test_wheel_cugraph-dgl.sh
@@ -4,7 +4,6 @@
 set -eoxu pipefail
 
 package_name="cugraph-dgl"
-package_dir="python/cugraph-dgl"
 
 python_package_name=$(echo ${package_name}|sed 's/-/_/g')
 
diff --git a/ci/test_wheel_cugraph-pyg.sh b/ci/test_wheel_cugraph-pyg.sh
index 1004063..1dc94bb 100755
--- a/ci/test_wheel_cugraph-pyg.sh
+++ b/ci/test_wheel_cugraph-pyg.sh
@@ -4,7 +4,6 @@
 set -eoxu pipefail
 
 package_name="cugraph-pyg"
-package_dir="python/cugraph-pyg"
 
 python_package_name=$(echo ${package_name}|sed 's/-/_/g')
 
diff --git a/conda/recipes/cugraph-dgl/meta.yaml b/conda/recipes/cugraph-dgl/meta.yaml
index 3c90075..86bbd4a 100644
--- a/conda/recipes/cugraph-dgl/meta.yaml
+++ b/conda/recipes/cugraph-dgl/meta.yaml
@@ -23,13 +23,13 @@ requirements:
     - python
   run:
     - cugraph ={{ version }}
-    - dgl >=1.1.0.cu*
+    - dgl >=2.4.0.th23.cu*
     - numba >=0.57
     - numpy >=1.23,<3.0a0
     - pylibcugraphops ={{ minor_version }}
     - tensordict >=0.1.2
     - python
-    - pytorch >=2.0
+    - pytorch >=2.3,<2.4.0a0
     - cupy >=12.0.0
 
 tests:
diff --git a/conda/recipes/cugraph-pyg/meta.yaml b/conda/recipes/cugraph-pyg/meta.yaml
index 5bce780..27ccaa4 100644
--- a/conda/recipes/cugraph-pyg/meta.yaml
+++ b/conda/recipes/cugraph-pyg/meta.yaml
@@ -30,12 +30,12 @@ requirements:
     - numba >=0.57
     - numpy >=1.23,<3.0a0
     - python
-    - pytorch >=2.0
+    - pytorch >=2.3,<2.4.0a0
     - cupy >=12.0.0
     - cugraph ={{ version }}
     - pylibcugraphops ={{ minor_version }}
     - tensordict >=0.1.2
-    - pyg >=2.5,<2.6
+    - pytorch_geometric >=2.5,<2.6
 
 tests:
   imports: