From a4d76675710dc21ae76c55dfea9eb9cc4600a94e Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Thu, 17 Oct 2024 14:05:45 -0500
Subject: [PATCH] manage more dependencies in dependencies.yaml, add 'numpy'
 run dependency to conda packages

---
 ci/test_wheel.sh                              | 13 ++++++----
 conda/recipes/pylibwholegraph/meta.yaml       |  1 +
 dependencies.yaml                             | 24 +++++++++++++++----
 .../binding/wholememory_binding.pyx           |  1 -
 python/pylibwholegraph/pyproject.toml         | 10 ++++++++
 5 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/ci/test_wheel.sh b/ci/test_wheel.sh
index 9fdeec250..b61a1a745 100755
--- a/ci/test_wheel.sh
+++ b/ci/test_wheel.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 
 set -e          # abort the script on error
 set -o pipefail # piped commands propagate their error
@@ -9,9 +9,7 @@ mkdir -p ./dist
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
 RAPIDS_PY_WHEEL_NAME="pylibwholegraph_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist
 
-# echo to expand wildcard before adding `[extra]` requires for pip
-python -m pip install $(echo ./dist/pylibwholegraph*.whl)
-
+# determine PyTorch source
 PKG_CUDA_VER="$(echo ${CUDA_VERSION} | cut -d '.' -f1,2 | tr -d '.')"
 PKG_CUDA_VER_MAJOR=${PKG_CUDA_VER:0:2}
 if [[ "${PKG_CUDA_VER_MAJOR}" == "12" ]]; then
@@ -19,13 +17,18 @@ if [[ "${PKG_CUDA_VER_MAJOR}" == "12" ]]; then
 else
   INDEX_URL="https://download.pytorch.org/whl/cu${PKG_CUDA_VER}"
 fi
+
+# echo to expand wildcard before adding `[extra]` requires for pip
+python -m pip install \
+  "$(echo ./dist/pylibwholegraph*.whl)[test]"
+
+
 RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}
 RAPIDS_COVERAGE_DIR=${RAPIDS_COVERAGE_DIR:-"${PWD}/coverage-results"}
 mkdir -p "${RAPIDS_TESTS_DIR}" "${RAPIDS_COVERAGE_DIR}"
 
 rapids-logger "Installing PyTorch"
 rapids-retry python -m pip install --pre torch --index-url ${INDEX_URL}
-rapids-retry python -m pip install pytest pytest-forked numpy
 rapids-logger "pytest pylibwholegraph"
 cd python/pylibwholegraph/pylibwholegraph/tests
 python -m pytest \
diff --git a/conda/recipes/pylibwholegraph/meta.yaml b/conda/recipes/pylibwholegraph/meta.yaml
index 8149fdc95..b78d2e6a4 100644
--- a/conda/recipes/pylibwholegraph/meta.yaml
+++ b/conda/recipes/pylibwholegraph/meta.yaml
@@ -71,6 +71,7 @@ requirements:
     - cudatoolkit
     {% endif %}
     - libwholegraph ={{ version }}
+    - numpy>=1.23,<3.0a0
     - python
 
 about:
diff --git a/dependencies.yaml b/dependencies.yaml
index 4d26f8a0c..8bcccd70b 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -64,6 +64,21 @@ files:
       key: requires
     includes:
       - python_build_wheel
+  py_run_pylibwholegraph:
+    output: pyproject
+    pyproject_dir: python/pylibwholegraph
+    extras:
+      table: project
+    includes:
+      - run
+  py_test_pylibwholegraph:
+    output: pyproject
+    pyproject_dir: python/pylibwholegraph
+    extras:
+      table: project.optional-dependencies
+      key: test
+    includes:
+      - test_python
 channels:
   - rapidsai
   - rapidsai-nightly
@@ -211,7 +226,8 @@ dependencies:
   run:
     common:
       - output_types: [conda, requirements]
-        packages: []
+        packages:
+          - &numpy numpy>=1.23,<3.0a0
   test_cpp:
     common:
       - output_types: [conda]
@@ -224,10 +240,10 @@ dependencies:
           - c-compiler
           - cxx-compiler
           - *nccl
-      - output_types: [conda, requirements]
-        packages:
           - ninja
-          - numpy>=1.23,<3.0a0
+      - output_types: [conda, pyproject, requirements]
+        packages:
+          - *numpy
           - pytest
           - pytest-forked
           - pytest-xdist
diff --git a/python/pylibwholegraph/pylibwholegraph/binding/wholememory_binding.pyx b/python/pylibwholegraph/pylibwholegraph/binding/wholememory_binding.pyx
index 5d8de9ce7..1e1298ff7 100644
--- a/python/pylibwholegraph/pylibwholegraph/binding/wholememory_binding.pyx
+++ b/python/pylibwholegraph/pylibwholegraph/binding/wholememory_binding.pyx
@@ -27,7 +27,6 @@ from libcpp cimport bool
 from cpython cimport Py_buffer
 from cpython cimport array
 import array
-import numpy as np
 from cpython.ref cimport PyObject, Py_INCREF, Py_DECREF
 from cpython.object cimport Py_TYPE, PyObject_CallObject
 from cpython.tuple cimport *
diff --git a/python/pylibwholegraph/pyproject.toml b/python/pylibwholegraph/pyproject.toml
index 0c4233529..62c9111a9 100644
--- a/python/pylibwholegraph/pyproject.toml
+++ b/python/pylibwholegraph/pyproject.toml
@@ -36,7 +36,17 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
 ]
+dependencies = [
+] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
+
 
+[project.optional-dependencies]
+test = [
+    "numpy>=1.23,<3.0a0",
+    "pytest",
+    "pytest-forked",
+    "pytest-xdist",
+] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 [tool.rapids-build-backend]
 build-backend = "scikit_build_core.build"
 dependencies-file = "../../dependencies.yaml"