wheels: build with CUDA 13.0, test against mix of CTK versions, make 'torch-geometric' fully optional for 'cugraph-pyg'

ci/test_wheel_cugraph-pyg.sh

@@ -38,7 +38,19 @@ if [ -z "$(ls -A ${TORCH_WHEEL_DIR} 2>/dev/null)" ]; then
   rapids-echo-stderr "No 'torch' wheels downloaded."
   torch_downloaded=false
 else
-  PIP_INSTALL_ARGS+=("${TORCH_WHEEL_DIR}"/torch-*.whl)
+  # if we were able to install 'torch', also install other dependencies that need 'torch',
+  # like 'torch-geometric' and 'sentence-transformers'
+  TORCH_DEPS_REQS_FILE=$(mktemp)
+  rapids-dependency-file-generator \
+    --file-key deps_that_require_torch \
+    --output requirements \
+    --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION};dependencies=${RAPIDS_DEPENDENCIES};require_gpu=true" \
+  | tee "${TORCH_DEPS_REQS_FILE}"
+
+  PIP_INSTALL_ARGS+=(
+    "${TORCH_WHEEL_DIR}"/torch-*.whl
+    -r "${TORCH_DEPS_REQS_FILE}"
+  )
 fi
 
 # notes:
@@ -61,33 +73,19 @@ popd
 export TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1
 
 if [[ "${torch_downloaded}" == "true" ]]; then
-  # TODO: remove this when RAPIDS wheels and 'torch' CUDA wheels have compatible package requirements
-  #
-  #    * https://github.com/rapidsai/cugraph/issues/5443
-  #    * https://github.com/rapidsai/build-planning/issues/257
-  #    * https://github.com/rapidsai/build-planning/issues/255
-  #
-  CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"
-  CUDA_MINOR=$(echo "${RAPIDS_CUDA_VERSION}" | cut -d'.' -f2)
-  if [[ "${CUDA_MAJOR}" == "13" ]]; then
-    pip install \
-      --upgrade \
-      "nvidia-nvjitlink>=${CUDA_MAJOR}.${CUDA_MINOR}"
-  fi
-
   # 'torch' is an optional dependency of 'cugraph-pyg'... confirm that it's actually
   # installed here and that we've installed a package with CUDA support.
   rapids-logger "Confirming that PyTorch is installed"
   python -c "import torch; assert torch.cuda.is_available()"
 
-  rapids-logger "pytest cugraph-pyg (single GPU, with 'torch')"
+  rapids-logger "pytest cugraph-pyg (single GPU, with 'torch' and 'torch-geometric')"
   ./ci/run_cugraph_pyg_pytests.sh
 fi
 
-rapids-logger "import cugraph-pyg (no 'torch')"
+rapids-logger "import cugraph-pyg (no 'torch' or 'torch-geometric')"
 ./ci/uninstall-torch-wheels.sh
 
 python -c "import cugraph_pyg; print(f'cugraph-pyg version: {cugraph_pyg.__version__}')"
 
-rapids-logger "pytest cugraph-pyg (no 'torch')"
+rapids-logger "pytest cugraph-pyg (no 'torch' or 'torch-geometric')"
 ./ci/run_cugraph_pyg_pytests.sh

ci/test_wheel_pylibwholegraph.sh

@@ -51,20 +51,6 @@ rapids-pip-retry install \
 
 
 if [[ "${torch_downloaded}" == "true" ]]; then
-  # TODO: remove this when RAPIDS wheels and 'torch' CUDA wheels have compatible package requirements
-  #
-  #    * https://github.com/rapidsai/cugraph/issues/5443
-  #    * https://github.com/rapidsai/build-planning/issues/257
-  #    * https://github.com/rapidsai/build-planning/issues/255
-  #
-  CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"
-  CUDA_MINOR=$(echo "${RAPIDS_CUDA_VERSION}" | cut -d'.' -f2)
-  if [[ "${CUDA_MAJOR}" == "13" ]]; then
-    pip install \
-      --upgrade \
-      "nvidia-nvjitlink>=${CUDA_MAJOR}.${CUDA_MINOR}"
-  fi
-
   # 'torch' is an optional dependency of 'pylibwholegraph'... confirm that it's actually
   # installed here and that we've installed a package with CUDA support.
   rapids-logger "Confirming that PyTorch is installed"

ci/uninstall-torch-wheels.sh

@@ -4,7 +4,7 @@
 
 set -euo pipefail
 
-pip uninstall --yes 'torch'
+pip uninstall --yes 'torch' 'torch-geometric'
 
 # 'pytest' leaves behind some pycache files in site-packages/torch that make 'import torch'
 # seem to "work" even though there's not really a package there, leading to errors like

dependencies.yaml

@@ -151,7 +151,6 @@ files:
       table: project
     includes:
       - depends_on_cupy
-      - depends_on_pyg
       - depends_on_pylibwholegraph
       - depends_on_pylibcugraph
       - python_run_cugraph_pyg
@@ -164,8 +163,6 @@ files:
     includes:
       - depends_on_cuml
       - depends_on_cugraph
-      - depends_on_ogb
-      - depends_on_sentence_transformers
       - test_python_common
   cugraph_pyg_dev:
     matrix:
@@ -179,6 +176,11 @@ files:
       - depends_on_pyg
       - depends_on_pytorch
       - test_python_common
+  deps_that_require_torch:
+    output: none
+    includes:
+      - depends_on_pyg
+      - depends_on_sentence_transformers
   torch_only:
     output: none
     includes:
@@ -193,6 +195,8 @@ dependencies:
       - output_types: [conda, requirements]
         packages:
           - pre-commit
+  # 'cuda_version' intentionally does not contain fallback entries... we want
+  # a loud error if an unsupported 'cuda' value is passed
   cuda_version:
     specific:
       - output_types: conda
@@ -201,18 +205,10 @@ dependencies:
               cuda: "12.2"
             packages:
               - cuda-version=12.2
-          - matrix:
-              cuda: "12.4"
-            packages:
-              - cuda-version=12.4
           - matrix:
               cuda: "12.5"
             packages:
               - cuda-version=12.5
-          - matrix:
-              cuda: "12.6"
-            packages:
-              - cuda-version=12.6
           - matrix:
               cuda: "12.8"
             packages:
@@ -229,6 +225,51 @@ dependencies:
               cuda: "13.1"
             packages:
               - cuda-version=13.1
+      - output_types: requirements
+        matrices:
+          # if use_cuda_wheels=false is provided, do not add dependencies on any CUDA wheels
+          # (e.g. for DLFW and pip devcontainers)
+          - matrix:
+              use_cuda_wheels: "false"
+            packages:
+          - matrix:
+              arch: aarch64
+              cuda: "12.2"
+              use_cuda_wheels: "true"
+            packages:
+              # some components (like nvidia-cublas-cu12 and nvidia-cuda-nvcc-cu12) didn't have
+              # aarch64 wheels until CTK 12.3, so allow a slightly looser bound here
+              - cuda-toolkit>=12.2,<12.4
+          - matrix:
+              cuda: "12.2"
+              use_cuda_wheels: "true"
+            packages:
+              - cuda-toolkit==12.2.*
+          - matrix:
+              cuda: "12.5"
+              use_cuda_wheels: "true"
+            packages:
+              - cuda-toolkit==12.5.*
+          - matrix:
+              cuda: "12.8"
+              use_cuda_wheels: "true"
+            packages:
+              - cuda-toolkit==12.8.*
+          - matrix:
+              cuda: "12.9"
+              use_cuda_wheels: "true"
+            packages:
+              - cuda-toolkit==12.9.*
+          - matrix:
+              cuda: "13.0"
+              use_cuda_wheels: "true"
+            packages:
+              - cuda-toolkit==13.0.*
+          - matrix:
+              cuda: "13.1"
+              use_cuda_wheels: "true"
+            packages:
+              - cuda-toolkit==13.1.*
   cuda:
     common:
       - output_types: [conda]
@@ -458,37 +499,26 @@ dependencies:
       - output_types: [conda]
         packages:
           - ogb
-    specific:
-      - output_types: [requirements, pyproject]
-        matrices:
-          - matrix:
-              no_pytorch: "true"
-            packages:
-          - matrix:
-            packages:
-              - ogb
   # for MovieLens example
   depends_on_sentence_transformers:
-    common:
-      - output_types: [conda]
-        packages:
-          - sentence-transformers
     specific:
-      - output_types: [requirements, pyproject]
+      - output_types: [conda, requirements, pyproject]
         matrices:
           - matrix:
               no_pytorch: "true"
             packages:
           - matrix:
             packages:
-              - sentence-transformers
+              - sentence-transformers>=3.0.1
   depends_on_pyg:
     common:
       - output_types: [conda]
         packages:
           - pytorch_geometric>=2.5,<2.8
     specific:
-      - output_types: [requirements, pyproject]
+      # no 'pyproject'... we intentionally keep 'torch-geometric' out of wheel
+      # metadata so it's possible to install 'cugraph-pyg' without 'torch'
+      - output_types: requirements
         matrices:
           - matrix:
               no_pytorch: "true"

python/cugraph-pyg/pyproject.toml

@@ -38,7 +38,6 @@ dependencies = [
     "pandas",
     "pylibcugraph==26.4.*,>=0.0.0a0",
     "pylibwholegraph==26.4.*,>=0.0.0a0",
-    "torch-geometric>=2.5,<2.8",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [project.entry-points."rapids_doctor_check"]
@@ -52,12 +51,10 @@ Documentation = "https://docs.rapids.ai/api/cugraph/stable/"
 test = [
     "cugraph==26.4.*,>=0.0.0a0",
     "cuml==26.4.*,>=0.0.0a0",
-    "ogb",
     "pytest",
     "pytest-benchmark",
     "pytest-cov",
     "pytest-xdist",
-    "sentence-transformers",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [tool.setuptools.dynamic]