From 3fd56dc4ab5cd33b9b72d37a3f4eb8dbc4885bac Mon Sep 17 00:00:00 2001
From: Tingyu Wang <tingyuw@nvidia.com>
Date: Fri, 12 Jan 2024 20:51:33 -0500
Subject: [PATCH] absorb conda build into build_python script

---
 .github/workflows/pr.yaml              |  7 ------
 ci/build_python.sh                     |  5 ++++
 ci/build_python_cugraph-equivariant.sh | 28 ---------------------
 ci/test_python.sh                      | 34 ++++++++++++++++++++++++++
 4 files changed, 39 insertions(+), 35 deletions(-)
 delete mode 100755 ci/build_python_cugraph-equivariant.sh

diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml
index 080b7653580..1bb2e0ab0a7 100644
--- a/.github/workflows/pr.yaml
+++ b/.github/workflows/pr.yaml
@@ -29,7 +29,6 @@ jobs:
       - wheel-tests-cugraph-dgl
       - wheel-build-cugraph-pyg
       - wheel-tests-cugraph-pyg
-      - conda-python-build-cugraph-equivariant
       - wheel-build-cugraph-equivariant
       - wheel-tests-cugraph-equivariant
       - devcontainer
@@ -59,12 +58,6 @@ jobs:
     uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-24.02
     with:
       build_type: pull-request
-  conda-python-build-cugraph-equivariant:
-    secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-24.02
-    with:
-      build_type: pull-request
-      build_script: ci/build_python_cugraph-equivariant.sh
   conda-python-tests:
     needs: conda-python-build
     secrets: inherit
diff --git a/ci/build_python.sh b/ci/build_python.sh
index a99e5ce63e8..07a4f59396b 100755
--- a/ci/build_python.sh
+++ b/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
diff --git a/ci/build_python_cugraph-equivariant.sh b/ci/build_python_cugraph-equivariant.sh
deleted file mode 100755
index b48d203a210..00000000000
--- a/ci/build_python_cugraph-equivariant.sh
+++ /dev/null
@@ -1,28 +0,0 @@
-#!/bin/bash
-# Copyright (c) 2024, NVIDIA CORPORATION.
-
-set -euo pipefail
-
-source rapids-env-update
-
-rapids-print-env
-
-version=$(rapids-generate-version)
-git_commit=$(git rev-parse HEAD)
-export RAPIDS_PACKAGE_VERSION=${version}
-echo "${version}" > VERSION
-
-rapids-logger "Begin py build"
-
-package_dir="python"
-for package_name in cugraph-equivariant; do
-  underscore_package_name=$(echo "${package_name}" | tr "-" "_")
-  sed -i "/^__git_commit__/ s/= .*/= \"${git_commit}\"/g" "${package_dir}/${package_name}/${underscore_package_name}/_version.py"
-done
-
-rapids-conda-retry mambabuild \
-  --no-test \
-  --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
-  conda/recipes/cugraph-equivariant
-
-rapids-upload-conda-to-s3 python
diff --git a/ci/test_python.sh b/ci/test_python.sh
index 7eb5a08edc8..f5a554e6414 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -247,5 +247,39 @@ else
   rapids-logger "skipping cugraph_pyg pytest on CUDA != 11.8"
 fi
 
+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
+    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-dgl.xml" \
+      --cov-config=../../.coveragerc \
+      --cov=cugraph_dgl \
+      --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/cugraph-dgl-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}