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/wheels-build.yaml@branch-24.02
+ 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/wheels-publish.yaml@branch-24.02
+ 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/pr-builder.yaml@branch-24.02
@@ -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/wheels-build.yaml@branch-24.02
+ 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/wheels-test.yaml@branch-24.02
+ 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/build-in-devcontainer.yaml@branch-24.02

.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/wheels-test.yaml@branch-24.02
+ 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",
+]