Fix CI build-time parallelism + support the same env var in cuda.core (#743)

* fix CI parallelism

* make cuda.core also support CUDA_PYTHON_PARALLEL_LEVEL

* [pre-commit.ci] auto code formatting

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: leofang

.github/workflows/build-wheel.yml

@@ -78,6 +78,8 @@ jobs:
           CIBW_BUILD_VERBOSITY: 1
           CIBW_BEFORE_BUILD_WINDOWS: "pip install delvewheel"
           CIBW_REPAIR_WHEEL_COMMAND_WINDOWS: "delvewheel repair --namespace-pkg cuda -w {dest_dir} {wheel}"
+          CIBW_ENVIRONMENT: >
+            CUDA_PYTHON_PARALLEL_LEVEL=${{ env.CUDA_PYTHON_PARALLEL_LEVEL }}
         with:
           package-dir: ./cuda_core/
           output-dir: ${{ env.CUDA_CORE_ARTIFACTS_DIR }}
@@ -121,11 +123,11 @@ jobs:
           CIBW_ENVIRONMENT_LINUX: >
             CUDA_PATH=/host/${{ env.CUDA_PATH }}
             LIBRARY_PATH=/host/${{ env.CUDA_PATH }}/lib
-            CUDA_BINDINGS_PARALLEL_LEVEL=${{ env.CUDA_BINDINGS_PARALLEL_LEVEL }}
+            CUDA_PYTHON_PARALLEL_LEVEL=${{ env.CUDA_PYTHON_PARALLEL_LEVEL }}
           CIBW_ENVIRONMENT_WINDOWS: >
             CUDA_HOME="$(cygpath -w ${{ env.CUDA_PATH }})"
             LIB="${CUDA_HOME}\\lib\\x64;${LIB}"
-            CUDA_BINDINGS_PARALLEL_LEVEL=${{ env.CUDA_BINDINGS_PARALLEL_LEVEL }}
+            CUDA_PYTHON_PARALLEL_LEVEL=${{ env.CUDA_PYTHON_PARALLEL_LEVEL }}
           CIBW_BEFORE_BUILD_WINDOWS: "pip install delvewheel"
           CIBW_REPAIR_WHEEL_COMMAND_WINDOWS: "delvewheel repair --namespace-pkg cuda -w {dest_dir} {wheel}"
         with:

ci/tools/env-vars

@@ -28,7 +28,7 @@ elif [[ "${HOST_PLATFORM}" == win* ]]; then
 fi
 
 echo "${TOOLS_PATH}" >> $GITHUB_PATH
-echo "CUDA_BINDINGS_PARALLEL_LEVEL=$(nproc)" >> $GITHUB_ENV
+echo "CUDA_PYTHON_PARALLEL_LEVEL=$(nproc)" >> $GITHUB_ENV
 CUDA_CORE_ARTIFACT_BASENAME="cuda-core-python${PYTHON_VERSION_FORMATTED}-${HOST_PLATFORM}"
 echo "CUDA_CORE_ARTIFACT_BASENAME=${CUDA_CORE_ARTIFACT_BASENAME}" >> $GITHUB_ENV
 echo "CUDA_CORE_ARTIFACT_NAME=${CUDA_CORE_ARTIFACT_BASENAME}-${SHA}" >> $GITHUB_ENV

cuda_core/setup.py

@@ -9,6 +9,9 @@
 from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext as _build_ext
 
+nthreads = int(os.environ.get("CUDA_PYTHON_PARALLEL_LEVEL", os.cpu_count() // 2))
+
+
 # It seems setuptools' wildcard support has problems for namespace packages,
 # so we explicitly spell out all Extension instances.
 root_module = "cuda.core.experimental"
@@ -33,7 +36,7 @@ def strip_prefix_suffix(filename):
 
 class build_ext(_build_ext):
     def build_extensions(self):
-        self.parallel = os.cpu_count() // 2
+        self.parallel = nthreads
         super().build_extensions()