Create cugraph-equivariant package (#4036)

Bring up `cugraph-equivariant` package and add TensorProduct conv layers.

Authors:
  - Tingyu Wang (https://github.com/tingyu66)

Approvers:
  - Jake Awe (https://github.com/AyodeAwe)
  - https://github.com/DejunL
  - Maximilian Stadler (https://github.com/stadlmax)
  - Mario Geiger (https://github.com/mariogeiger)
  - Rick Ratzel (https://github.com/rlratzel)

URL: #4036

.github/workflows/build.yaml

@@ -173,3 +173,22 @@ jobs:
       sha: ${{ inputs.sha }}
       date: ${{ inputs.date }}
       package-name: cugraph-pyg
+  wheel-build-cugraph-equivariant:
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/[email protected]
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+      script: ci/build_wheel_cugraph-equivariant.sh
+  wheel-publish-cugraph-equivariant:
+    needs: wheel-build-cugraph-equivariant
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/[email protected]
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+      package-name: cugraph-equivariant

.github/workflows/pr.yaml

@@ -29,6 +29,8 @@ jobs:
       - wheel-tests-cugraph-dgl
       - wheel-build-cugraph-pyg
       - wheel-tests-cugraph-pyg
+      - wheel-build-cugraph-equivariant
+      - wheel-tests-cugraph-equivariant
       - devcontainer
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]
@@ -161,6 +163,20 @@ jobs:
       build_type: pull-request
       script: ci/test_wheel_cugraph-pyg.sh
       matrix_filter: map(select(.ARCH == "amd64" and .CUDA_VER == "11.8.0"))
+  wheel-build-cugraph-equivariant:
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/[email protected]
+    with:
+      build_type: pull-request
+      script: ci/build_wheel_cugraph-equivariant.sh
+  wheel-tests-cugraph-equivariant:
+    needs: wheel-build-cugraph-equivariant
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/[email protected]
+    with:
+      build_type: pull-request
+      script: ci/test_wheel_cugraph-equivariant.sh
+      matrix_filter: map(select(.ARCH == "amd64"))
   devcontainer:
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]

.github/workflows/test.yaml

@@ -75,3 +75,12 @@ jobs:
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
       script: ci/test_wheel_cugraph-pyg.sh
+  wheel-tests-cugraph-equivariant:
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/[email protected]
+    with:
+      build_type: nightly
+      branch: ${{ inputs.branch }}
+      date: ${{ inputs.date }}
+      sha: ${{ inputs.sha }}
+      script: ci/test_wheel_cugraph-equivariant.sh

build.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-# Copyright (c) 2019-2023, NVIDIA CORPORATION.
+# Copyright (c) 2019-2024, NVIDIA CORPORATION.
 
 # cugraph build script
 
@@ -31,6 +31,7 @@ VALIDARGS="
    cugraph-service
    cugraph-pyg
    cugraph-dgl
+   cugraph-equivariant
    nx-cugraph
    cpp-mgtests
    cpp-mtmgtests
@@ -60,6 +61,7 @@ HELP="$0 [<target> ...] [<flag> ...]
    cugraph-service            - build the cugraph-service_client and cugraph-service_server Python package
    cugraph-pyg                - build the cugraph-pyg Python package
    cugraph-dgl                - build the cugraph-dgl extensions for DGL
+   cugraph-equivariant        - build the cugraph-equivariant Python package
    nx-cugraph                 - build the nx-cugraph Python package
    cpp-mgtests                - build libcugraph and libcugraph_etl MG tests. Builds MPI communicator, adding MPI as a dependency.
    cpp-mtmgtests              - build libcugraph MTMG tests. Adds UCX as a dependency (temporary).
@@ -222,7 +224,7 @@ if hasArg uninstall; then
     # 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 cugraph-pyg nx-cugraph
+        cugraph-dgl cugraph-pyg cugraph-equivariant nx-cugraph
 fi
 
 if hasArg clean; then
@@ -359,6 +361,15 @@ if hasArg cugraph-dgl || hasArg all; then
     fi
 fi
 
+# Build and install the cugraph-equivariant Python package
+if hasArg cugraph-equivariant || hasArg all; then
+    if hasArg --clean; then
+        cleanPythonDir ${REPODIR}/python/cugraph-equivariant
+    else
+        python ${PYTHON_ARGS_FOR_INSTALL} ${REPODIR}/python/cugraph-equivariant
+    fi
+fi
+
 # Build and install the nx-cugraph Python package
 if hasArg nx-cugraph || hasArg all; then
     if hasArg --clean; then

ci/build_python.sh

@@ -89,4 +89,9 @@ if [[ ${RAPIDS_CUDA_MAJOR} == "11" ]]; then
     conda/recipes/cugraph-dgl
 fi
 
+rapids-conda-retry mambabuild \
+  --no-test \
+  --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
+  conda/recipes/cugraph-equivariant
+
 rapids-upload-conda-to-s3 python

ci/build_wheel.sh

@@ -57,7 +57,8 @@ python -m pip wheel . -w dist -vvv --no-deps --disable-pip-version-check
 # pure-python packages should not have auditwheel run on them.
 if [[ ${package_name} == "nx-cugraph" ]] || \
    [[ ${package_name} == "cugraph-dgl" ]] || \
-   [[ ${package_name} == "cugraph-pyg" ]]; then
+   [[ ${package_name} == "cugraph-pyg" ]] || \
+   [[ ${package_name} == "cugraph-equivariant" ]]; then
     RAPIDS_PY_WHEEL_NAME="${package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 dist
 else
     mkdir -p final_dist

ci/build_wheel_cugraph-equivariant.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+./ci/build_wheel.sh cugraph-equivariant python/cugraph-equivariant

ci/test_python.sh

@@ -247,5 +247,46 @@ else
   rapids-logger "skipping cugraph_pyg pytest on CUDA != 11.8"
 fi
 
+# test cugraph-equivariant
+if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
+  if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
+    # Reuse cugraph-dgl's test env for cugraph-equivariant
+    set +u
+    conda activate test_cugraph_dgl
+    set -u
+    rapids-mamba-retry install \
+      --channel "${CPP_CHANNEL}" \
+      --channel "${PYTHON_CHANNEL}" \
+      --channel pytorch \
+      --channel nvidia \
+      cugraph-equivariant
+    pip install e3nn==0.5.1
+
+    rapids-print-env
+
+    rapids-logger "pytest cugraph-equivariant"
+    pushd python/cugraph-equivariant/cugraph_equivariant
+    pytest \
+      --cache-clear \
+      --junitxml="${RAPIDS_TESTS_DIR}/junit-cugraph-equivariant.xml" \
+      --cov-config=../../.coveragerc \
+      --cov=cugraph_equivariant \
+      --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/cugraph-equivariant-coverage.xml" \
+      --cov-report=term \
+      .
+    popd
+
+    # Reactivate the test environment back
+    set +u
+    conda deactivate
+    conda activate test
+    set -u
+  else
+    rapids-logger "skipping cugraph-equivariant pytest on ARM64"
+  fi
+else
+  rapids-logger "skipping cugraph-equivariant pytest on CUDA!=11.8"
+fi
+
 rapids-logger "Test script exiting with value: $EXITCODE"
 exit ${EXITCODE}

ci/test_wheel_cugraph-equivariant.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -eoxu pipefail
+
+package_name="cugraph-equivariant"
+package_dir="python/cugraph-equivariant"
+
+python_package_name=$(echo ${package_name}|sed 's/-/_/g')
+
+mkdir -p ./dist
+RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
+
+# use 'ls' to expand wildcard before adding `[extra]` requires for pip
+RAPIDS_PY_WHEEL_NAME="${package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist
+# pip creates wheels using python package names
+python -m pip install $(ls ./dist/${python_package_name}*.whl)[test]
+
+
+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
+  PYTORCH_CUDA_VER="121"
+else
+  PYTORCH_CUDA_VER=$PKG_CUDA_VER
+fi
+PYTORCH_URL="https://download.pytorch.org/whl/cu${PYTORCH_CUDA_VER}"
+
+rapids-logger "Installing PyTorch and e3nn"
+rapids-retry python -m pip install torch --index-url ${PYTORCH_URL}
+rapids-retry python -m pip install e3nn
+
+python -m pytest python/cugraph-equivariant/cugraph_equivariant/tests

conda/recipes/cugraph-equivariant/build.sh

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

conda/recipes/cugraph-equivariant/meta.yaml

@@ -0,0 +1,37 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+{% set version = environ['RAPIDS_PACKAGE_VERSION'].lstrip('v') + environ.get('VERSION_SUFFIX', '') %}
+{% set minor_version = version.split('.')[0] + '.' + version.split('.')[1] %}
+{% set py_version = environ['CONDA_PY'] %}
+{% set date_string = environ['RAPIDS_DATE_STRING'] %}
+
+package:
+  name: cugraph-equivariant
+  version: {{ version }}
+
+source:
+  path: ../../..
+
+build:
+  number: {{ GIT_DESCRIBE_NUMBER }}
+  build:
+      number: {{ GIT_DESCRIBE_NUMBER }}
+      string: py{{ py_version }}_{{ date_string }}_{{ GIT_DESCRIBE_HASH }}_{{ GIT_DESCRIBE_NUMBER }}
+
+requirements:
+  host:
+    - python
+  run:
+    - pylibcugraphops ={{ minor_version }}
+    - python
+
+tests:
+  imports:
+    - cugraph_equivariant
+
+about:
+  home: https://rapids.ai/
+  dev_url: https://github.com/rapidsai/cugraph
+  license: Apache-2.0
+  license_file: ../../../LICENSE
+  summary: GPU-accelerated equivariant convolutional layers.

dependencies.yaml

@@ -198,6 +198,28 @@ files:
       key: test
     includes:
       - test_python_common
+  py_build_cugraph_equivariant:
+    output: pyproject
+    pyproject_dir: python/cugraph-equivariant
+    extras:
+      table: build-system
+    includes:
+      - python_build_wheel
+  py_run_cugraph_equivariant:
+    output: pyproject
+    pyproject_dir: python/cugraph-equivariant
+    extras:
+      table: project
+    includes:
+      - depends_on_pylibcugraphops
+  py_test_cugraph_equivariant:
+    output: pyproject
+    pyproject_dir: python/cugraph-equivariant
+    extras:
+      table: project.optional-dependencies
+      key: test
+    includes:
+      - test_python_common
   py_build_cugraph_service_client:
     output: pyproject
     pyproject_dir: python/cugraph-service/client

python/cugraph-equivariant/LICENSE

@@ -0,0 +1 @@
+../../LICENSE

python/cugraph-equivariant/README.md

@@ -0,0 +1,5 @@
+# cugraph-equivariant
+
+## Description
+
+cugraph-equivariant library provides fast symmetry-preserving (equivariant) operations and convolutional layers, to accelerate the equivariant neural networks in drug discovery and other domains.

python/cugraph-equivariant/cugraph_equivariant/VERSION

@@ -0,0 +1 @@
+../../../VERSION

python/cugraph-equivariant/cugraph_equivariant/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from cugraph_equivariant._version import __git_commit__, __version__

python/cugraph-equivariant/cugraph_equivariant/_version.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.resources
+
+# Read VERSION file from the module that is symlinked to VERSION file
+# in the root of the repo at build time or copied to the module at
+# installation. VERSION is a separate file that allows CI build-time scripts
+# to update version info (including commit hashes) without modifying
+# source files.
+__version__ = (
+    importlib.resources.files("cugraph_equivariant")
+    .joinpath("VERSION")
+    .read_text()
+    .strip()
+)
+__git_commit__ = ""

python/cugraph-equivariant/cugraph_equivariant/nn/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .tensor_product_conv import FullyConnectedTensorProductConv
+
+DiffDockTensorProductConv = FullyConnectedTensorProductConv
+
+__all__ = [
+    "FullyConnectedTensorProductConv",
+    "DiffDockTensorProductConv",
+]