From c30376311bd15c1ec70933f2ad591ce5898da00e Mon Sep 17 00:00:00 2001
From: Ben Howe <141149032+bmhowe23@users.noreply.github.com>
Date: Mon, 27 Jan 2025 18:44:08 -0800
Subject: [PATCH] Update CI build_cudaq.sh to specify num parallel jobs (#63)

Related (and required) prerequisite:
https://github.com/NVIDIA/cuda-quantum/pull/2546

Signed-off-by: Ben Howe <bhowe@nvidia.com>
---
 .github/workflows/build_wheels.yaml      |  2 +-
 .github/workflows/scripts/build_cudaq.sh | 60 ++++++++++++++++++++++--
 2 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/build_wheels.yaml b/.github/workflows/build_wheels.yaml
index 190cf11..8a1e439 100644
--- a/.github/workflows/build_wheels.yaml
+++ b/.github/workflows/build_wheels.yaml
@@ -72,7 +72,7 @@ jobs:
 
       - name: Build CUDAQ toolchain
         run: |
-          .github/workflows/scripts/build_cudaq.sh ${{ matrix.python }}
+          .github/workflows/scripts/build_cudaq.sh --python-version ${{ matrix.python }}
 
       - name: Build CUDA-QX wheels
         run: |
diff --git a/.github/workflows/scripts/build_cudaq.sh b/.github/workflows/scripts/build_cudaq.sh
index 6fc49a2..c5f7884 100755
--- a/.github/workflows/scripts/build_cudaq.sh
+++ b/.github/workflows/scripts/build_cudaq.sh
@@ -8,6 +8,62 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 
+# ==============================================================================
+# Handling options
+# ==============================================================================
+
+show_help() {
+    echo "Usage: $0 [options]"
+    echo "Options:"
+    echo "  --python-version  Python version to build wheel for (e.g. 3.10)"
+    echo "  -j                Number of parallel jobs to build CUDA-Q with"
+    echo "                    (e.g. 8)"
+}
+
+parse_options() {
+    while (( $# > 0 )); do
+        case "$1" in
+            --python-version)
+                if [[ -n "$2" && "$2" != -* ]]; then
+                    python_version=("$2")
+                    shift 2
+                else
+                    echo "Error: Argument for $1 is missing" >&2
+                    exit 1
+                fi
+                ;;
+            -j)
+                if [[ -n "$2" && "$2" != -* ]]; then
+                    num_par_jobs=("$2")
+                    cudaq_ninja_jobs_arg="-j $num_par_jobs"
+                    shift 2
+                else
+                    echo "Error: Argument for $1 is missing" >&2
+                    exit 1
+                fi
+                ;;
+            -*)
+                echo "Error: Unknown option $1" >&2
+                show_help
+                exit 1
+                ;;
+            *)
+                echo "Error: Unknown argument $1" >&2
+                show_help
+                exit 1
+                ;;
+        esac
+    done
+}
+
+# Defaults
+python_version=3.10
+cudaq_ninja_jobs_arg=""
+
+# Parse options
+parse_options "$@"
+
+
 export CUDA_VERSION=12.0
 
 # We need to use a newer toolchain because CUDA-QX libraries rely on c++20
@@ -16,8 +72,6 @@ source /opt/rh/gcc-toolset-11/enable
 export CC=gcc
 export CXX=g++
 
-python_version=$1
-python_version=${python_version:-3.10}
 python=python${python_version}
 ${python} -m pip install --no-cache-dir numpy auditwheel
 
@@ -69,5 +123,5 @@ echo "$CUDAQ_PATCH" | git apply --verbose
 
 $python -m venv --system-site-packages .venv
 source .venv/bin/activate
-CUDAQ_BUILD_TESTS=FALSE bash scripts/build_cudaq.sh -v
+CUDAQ_BUILD_TESTS=FALSE bash scripts/build_cudaq.sh -v ${cudaq_ninja_jobs_arg}