PyG Monorepo Refactor

build.sh

@@ -28,6 +28,7 @@ VALIDARGS="
    cugraph-service
    pylibcugraph
    cpp-mgtests
+   cugraph-pyg
    cugraph-dgl
    docs
    -v
@@ -50,6 +51,7 @@ HELP="$0 [<target> ...] [<flag> ...]
    libcugraph                 - build libcugraph.so and SG test binaries
    libcugraph_etl             - build libcugraph_etl.so and SG test binaries
    pylibcugraph               - build the pylibcugraph Python package
+   cugraph-pyg                - build the cugraph-pyg Python package
    cugraph                    - build the cugraph Python package
    cugraph-service            - build the cugraph-service_client and cugraph-service_server Python package
    cpp-mgtests                - build libcugraph and libcugraph_etl MG tests. Builds MPI communicator, adding MPI as a dependency.
@@ -67,7 +69,7 @@ HELP="$0 [<target> ...] [<flag> ...]
    --clean                    - clean an individual target (note: to do a complete rebuild, use the clean target described above)
    -h                         - print this text
 
- default action (no args) is to build and install 'libcugraph' then 'libcugraph_etl' then 'pylibcugraph' then 'cugraph' then 'cugraph-service' targets
+ default action (no args) is to build and install 'libcugraph' then 'libcugraph_etl' then 'pylibcugraph' then 'cugraph' targets
 
  libcugraph build dir is: ${LIBCUGRAPH_BUILD_DIR}
 
@@ -198,7 +200,7 @@ if hasArg uninstall; then
     # FIXME: if multiple versions of these packages are installed, this only
     # removes the latest one and leaves the others installed. build.sh uninstall
     # can be run multiple times to remove all of them, but that is not obvious.
-    pip uninstall -y pylibcugraph cugraph cugraph-service-client cugraph-service-server cugraph-dgl
+    pip uninstall -y pylibcugraph cugraph cugraph-service-client cugraph-service-server cugraph-dgl cugraph-pyg
 fi
 
 if hasArg clean; then
@@ -337,7 +339,7 @@ if buildAll || hasArg cugraph; then
 fi
 
 # Install the cugraph-service-client and cugraph-service-server Python packages
-if buildAll || hasArg cugraph-service; then
+if hasArg cugraph-service; then
     if hasArg --clean; then
         cleanPythonDir ${REPODIR}/python/cugraph-service
     else
@@ -350,8 +352,20 @@ if buildAll || hasArg cugraph-service; then
     fi
 fi
 
+# Build and install the cugraph-pyg Python package
+if hasArg cugraph-pyg; then
+    if hasArg --clean; then
+        cleanPythonDir ${REPODIR}/python/cugraph-pyg
+    else
+        if [[ ${INSTALL_TARGET} != "" ]]; then
+            cd ${REPODIR}/python/cugraph-pyg
+            python ${PYTHON_ARGS_FOR_INSTALL}
+        fi
+    fi
+fi
+
 # Install the cugraph-dgl extensions for DGL
-if buildAll || hasArg cugraph-dgl; then
+if hasArg cugraph-dgl; then
     if hasArg --clean; then
         cleanPythonDir ${REPODIR}/python/cugraph-dgl
     else

ci/checks/copyright.py

@@ -38,7 +38,7 @@
     re.compile(r"[.]flake8[.]cython$"),
     re.compile(r"meta[.]yaml$")
 ]
-ExemptFiles = []
+ExemptFiles = ['versioneer.py']
 
 # this will break starting at year 10000, which is probably OK :)
 CheckSimple = re.compile(

ci/cpu/build.sh

@@ -115,6 +115,8 @@ if [ "$BUILD_CUGRAPH" == "1" ]; then
     gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/cugraph --python=$PYTHON
     gpuci_logger "cugraph-service"
     gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/cugraph-service --python=$PYTHON
+    gpuci_logger "cugraph-pyg"
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/cugraph-pyg --python=$PYTHON
     gpuci_logger "cugraph-dgl"
     gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/cugraph-dgl --python=$PYTHON -c dglteam -c pytorch
   else
@@ -124,6 +126,8 @@ if [ "$BUILD_CUGRAPH" == "1" ]; then
     gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/cugraph -c ${CONDA_LOCAL_CHANNEL} --dirty --no-remove-work-dir --python=$PYTHON
     gpuci_logger "cugraph-service"
     gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/cugraph-service -c ${CONDA_LOCAL_CHANNEL} --dirty --no-remove-work-dir --python=$PYTHON
+    gpuci_logger "cugraph-pyg"
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/cugraph-pyg -c ${CONDA_LOCAL_CHANNEL} --dirty --no-remove-work-dir --python=$PYTHON
     gpuci_logger "cugraph-dgl"
     gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/cugraph-dgl -c ${CONDA_LOCAL_CHANNEL} --dirty --no-remove-work-dir --python=$PYTHON -c dglteam -c pytorch
 

ci/cpu/upload.sh

@@ -41,16 +41,21 @@ if [[ "$BUILD_CUGRAPH" == "1" ]]; then
   test -e ${PYLIBCUGRAPH_FILE}
   echo "Upload pylibcugraph file: ${PYLIBCUGRAPH_FILE}"
   gpuci_retry anaconda -t ${MY_UPLOAD_KEY} upload -u ${CONDA_USERNAME:-rapidsai} ${LABEL_OPTION} --skip-existing ${PYLIBCUGRAPH_FILE} --no-progress
-
+ 
   CUGRAPH_FILE=$(conda build --croot ${CONDA_BLD_DIR} conda/recipes/cugraph --python=$PYTHON --output)
   test -e ${CUGRAPH_FILE}
   echo "Upload cugraph file: ${CUGRAPH_FILE}"
   gpuci_retry anaconda -t ${MY_UPLOAD_KEY} upload -u ${CONDA_USERNAME:-rapidsai} ${LABEL_OPTION} --skip-existing ${CUGRAPH_FILE} --no-progress
-
+ 
   CUGRAPH_SERVICE_FILES=$(conda build --croot ${CONDA_BLD_DIR} conda/recipes/cugraph-service --python=$PYTHON --output)
   # no test -e since CUGRAPH_SERVICE_FILES has multiple files
   echo "Upload cugraph-server files: ${CUGRAPH_SERVICE_FILES}"
   gpuci_retry anaconda -t ${MY_UPLOAD_KEY} upload -u ${CONDA_USERNAME:-rapidsai} ${LABEL_OPTION} --skip-existing --no-progress ${CUGRAPH_SERVICE_FILES}
+
+  CUGRAPH_PYG_FILE=$(conda build --croot ${CONDA_BLD_DIR} conda/recipes/cugraph-pyg --python=$PYTHON --output)
+  test -e ${CUGRAPH_PYG_FILE}
+  echo "Upload cugraph file: ${CUGRAPH_PYG_FILE}"
+  gpuci_retry anaconda -t ${MY_UPLOAD_KEY} upload -u ${CONDA_USERNAME:-rapidsai} ${LABEL_OPTION} --skip-existing ${CUGRAPH_PYG_FILE} --no-progress
 
   CUGRAPH_DGL_FILE=$(conda build --croot ${CONDA_BLD_DIR} conda/recipes/cugraph-dgl --python=$PYTHON --output)
   test -e ${CUGRAPH_DGL_FILE}

ci/gpu/build.sh

@@ -114,7 +114,7 @@ conda list --show-channel-urls
 
 if [[ -z "$PROJECT_FLASH" || "$PROJECT_FLASH" == "0" ]]; then
     gpuci_logger "Build from source"
-    $WORKSPACE/build.sh -v clean libcugraph pylibcugraph cugraph cugraph-service-server cugraph-service-client
+    $WORKSPACE/build.sh -v clean libcugraph pylibcugraph cugraph cugraph-service-server cugraph-service-client cugraph-pyg
 else
     gpuci_logger "Installing libcugraph-tests"
     gpuci_mamba_retry install -c ${CONDA_ARTIFACT_PATH} libcugraph libcugraph_etl libcugraph-tests
@@ -132,8 +132,11 @@ else
     gpuci_logger "cugraph-service"
     gpuci_conda_retry mambabuild conda/recipes/cugraph-service --no-build-id --croot ${CONDA_BLD_DIR} -c ${CONDA_ARTIFACT_PATH} --python=${PYTHON}
 
-    gpuci_logger "Installing pylibcugraph, cugraph, and cugraph-service from build / artifact dirs"
-    gpuci_mamba_retry install -c ${CONDA_BLD_DIR} -c ${CONDA_ARTIFACT_PATH} --strict-channel-priority pylibcugraph cugraph cugraph-service-server cugraph-service-client
+    gpuci_logger "Building and installing cugraph-pyg..."
+    gpuci_conda_retry mambabuild conda/recipes/cugraph-pyg --no-build-id --croot ${CONDA_BLD_DIR} -c ${CONDA_ARTIFACT_PATH} --python=${PYTHON}
+
+    gpuci_logger "Installing pylibcugraph, cugraph, cugraph-pyg and cugraph-service from build / artifact dirs"
+    gpuci_mamba_retry install -c ${CONDA_BLD_DIR} -c ${CONDA_ARTIFACT_PATH} pylibcugraph cugraph cugraph-pyg cugraph-service-server cugraph-service-client
 fi
 
 ################################################################################

ci/test.sh

@@ -102,6 +102,13 @@ if hasArg "--run-python-tests"; then
     pytest -sv -m "managedmem_on and poolallocator_on and tiny" --benchmark-disable
     echo "Ran Python benchmarks for cuGraph (running as tests) : return code was: $?, test script exit code is now: $EXITCODE"
 
+    echo "Python pytest for cugraph_pyg (single-GPU only)..."
+    conda list
+    cd ${CUGRAPH_ROOT}/python/cugraph_pyg/cugraph_pyg
+    # rmat is not tested because of MG testing
+    pytest --cache-clear --junitxml=${CUGRAPH_ROOT}/junit-cugraph-pytests.xml -v --cov-config=.coveragerc --cov=cugraph_pyg --cov-report=xml:${WORKSPACE}/python/cugraph_pyg/cugraph-coverage.xml --cov-report term --ignore=raft --ignore=tests/mg --ignore=tests/int --ignore=tests/generators --benchmark-disable
+    echo "Ran Python pytest for cugraph_pyg : return code was: $?, test script exit code is now: $EXITCODE"
+
     echo "Python pytest for cugraph-service (single-GPU only)..."
     cd ${CUGRAPH_ROOT}/python/cugraph-service
     pytest -sv --cache-clear --junitxml=${CUGRAPH_ROOT}/junit-cugraph-service-pytests.xml --benchmark-disable -k "not mg" ./tests

conda/recipes/cugraph-pyg/build.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Copyright (c) 2022, NVIDIA CORPORATION.
+
+# This assumes the script is executed from the root of the repo directory
+./build.sh cugraph-pyg --allgpuarch

conda/recipes/cugraph-pyg/conda_build_config.yaml

@@ -0,0 +1,16 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.
+
+c_compiler_version:
+  - 9
+
+cxx_compiler_version:
+  - 9
+
+cuda_compiler:
+  - nvcc
+
+cmake_version:
+  - ">=3.23.1"
+
+sysroot_version:
+  - "2.17"

conda/recipes/cugraph-pyg/meta.yaml

@@ -0,0 +1,50 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.
+
+# Usage:
+#   conda build -c rapidsai -c conda-forge -c nvidia .
+{% set version = environ.get('GIT_DESCRIBE_TAG', '0.0.0.dev').lstrip('v') + environ.get('VERSION_SUFFIX', '') %}
+{% set minor_version =  version.split('.')[0] + '.' + version.split('.')[1] %}
+{% set cuda_version='.'.join(environ.get('CUDA', 'unknown').split('.')[:2]) %}
+{% set cuda_major=cuda_version.split('.')[0] %}
+{% set py_version=environ.get('CONDA_PY', 36) %}
+package:
+  name: cugraph-pyg
+  version: {{ version }}
+
+source:
+  git_url: ../../..
+
+build:
+  number: {{ GIT_DESCRIBE_NUMBER }}
+  string: cuda{{ cuda_major }}_py{{ py_version }}_{{ GIT_DESCRIBE_HASH }}_{{ GIT_DESCRIBE_NUMBER }}
+  script_env:
+    - PARALLEL_LEVEL
+  preserve_egg_dir: True
+
+requirements:
+  build:
+    - sysroot_{{ target_platform }} {{ sysroot_version }}
+  host:
+    - python x.x
+    - cython>=0.29,<0.30
+    - scikit-build>=0.13.1
+    # FIXME: this pin can be removed once we move to the GitHub Actions build process
+    - setuptools<=65.2.0
+  run:
+    - distributed>=2022.9.2
+    - numba>=0.56.2
+    - numpy
+    - pytorch
+    - cudatoolkit=11.6
+    # - pyg (uncomment once pyg 2.2 is released)
+
+tests:                                 # [linux64]
+  imports:                             # [linux64]
+    - cugraph_pyg                      # [linux64]
+    # - cugraph_service (uncomment once cugraph-service int. tests ready)
+
+about:
+  home: http://rapids.ai/
+  license: Apache-2.0
+  license_file: ../../../LICENSE
+  summary: cuGraph-pyg library

notebooks/gnn/pyg_hetero_mag.ipynb

@@ -173,7 +173,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from cugraph.experimental.pyg_extensions import to_pyg\n",
+    "from cugraph_pyg.data import to_pyg\n",
     "\n",
     "feature_store, graph_store = to_pyg(pG)"
    ]
@@ -184,7 +184,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from cugraph.experimental.pyg_extensions import CuGraphSampler\n",
+    "from cugraph_pyg.sampler import CuGraphSampler\n",
     "sampler = CuGraphSampler(\n",
     "    data=(feature_store, graph_store),\n",
     "    shuffle=True,\n",
@@ -224,6 +224,15 @@
     "### Create the Network"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%timeit next(iter(loader))"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

python/cugraph-pyg/MANIFEST.in

@@ -0,0 +1,2 @@
+include versioneer.py
+include cugraph-pyg/_version.py