[GLUTEN-4122][VL] Add num_threads option to limit build threads (#4123)

apache · Jan 16, 2024 · 4a385c3 · 4a385c3
1 parent 46527ca
commit 4a385c3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/dev/builddeps-veloxbe.sh b/dev/builddeps-veloxbe.sh
@@ -32,6 +32,7 @@ VELOX_BRANCH=""
 VELOX_HOME=""
 VELOX_PARAMETER=""
 COMPILE_ARROW_JAVA=OFF
+NUM_THREADS=${NUM_THREADS:-$(nproc --ignore=2)}
 
 for arg in "$@"
 do
@@ -125,6 +126,10 @@ do
         --compile_arrow_java=*)
         COMPILE_ARROW_JAVA=("${arg#*=}")
         shift # Remove argument name from processing
+        ;;
+        --num_threads=*)
+        NUM_THREADS=("${arg#*=}")
+        shift # Remove argument name from processing
         ;;
 	      *)
         OTHER_ARGUMENTS+=("$1")
@@ -163,7 +168,7 @@ cd $GLUTEN_DIR/ep/build-velox/src
 # When BUILD_TESTS is on for gluten cpp, we need turn on VELOX_BUILD_TEST_UTILS via build_test_utils.
 ./build_velox.sh --run_setup_script=$RUN_SETUP_SCRIPT --enable_s3=$ENABLE_S3 --enable_gcs=$ENABLE_GCS --build_type=$BUILD_TYPE --enable_hdfs=$ENABLE_HDFS \
                  --enable_abfs=$ENABLE_ABFS --enable_ep_cache=$ENABLE_EP_CACHE --build_test_utils=$BUILD_TESTS --build_tests=$BUILD_VELOX_TESTS --build_benchmarks=$BUILD_VELOX_BENCHMARKS \
-                 --compile_arrow_java=$COMPILE_ARROW_JAVA
+                 --compile_arrow_java=$COMPILE_ARROW_JAVA  --num_threads=$NUM_THREADS
 
 ## compile gluten cpp
 cd $GLUTEN_DIR/cpp
@@ -173,4 +178,4 @@ cd build
 cmake -DBUILD_VELOX_BACKEND=ON -DCMAKE_BUILD_TYPE=$BUILD_TYPE \
       -DBUILD_TESTS=$BUILD_TESTS -DBUILD_EXAMPLES=$BUILD_EXAMPLES -DBUILD_BENCHMARKS=$BUILD_BENCHMARKS -DBUILD_JEMALLOC=$BUILD_JEMALLOC \
       -DENABLE_HBM=$ENABLE_HBM -DENABLE_QAT=$ENABLE_QAT -DENABLE_IAA=$ENABLE_IAA -DENABLE_GCS=$ENABLE_GCS -DENABLE_S3=$ENABLE_S3 -DENABLE_HDFS=$ENABLE_HDFS -DENABLE_ABFS=$ENABLE_ABFS ..
-make -j
+make -j $NUM_THREADS
diff --git a/ep/build-velox/src/build_velox.sh b/ep/build-velox/src/build_velox.sh
@@ -35,6 +35,7 @@ ENABLE_TESTS=OFF
 BUILD_TEST_UTILS=OFF
 RUN_SETUP_SCRIPT=ON
 COMPILE_ARROW_JAVA=OFF
+NUM_THREADS=""
 OTHER_ARGUMENTS=""
 
 OS=`uname -s`
@@ -90,6 +91,10 @@ for arg in "$@"; do
     COMPILE_ARROW_JAVA=("${arg#*=}")
     shift # Remove argument name from processing
     ;;
+  --num_threads=*)
+    NUM_THREADS=("${arg#*=}")
+    shift # Remove argument name from processing
+    ;;
   *)
     OTHER_ARGUMENTS+=("$1")
     shift # Remove generic argument from processing
@@ -139,11 +144,17 @@ function compile {
   COMPILE_TYPE=$(if [[ "$BUILD_TYPE" == "debug" ]] || [[ "$BUILD_TYPE" == "Debug" ]]; then echo 'debug'; else echo 'release'; fi)
   echo "COMPILE_OPTION: "$COMPILE_OPTION
 
+  NUM_THREADS_OPTS=""
+  if [ -n "${NUM_THREADS:-}" ]; then
+    NUM_THREADS_OPTS="NUM_THREADS=$NUM_THREADS MAX_HIGH_MEM_JOBS=$NUM_THREADS MAX_LINK_JOBS=$NUM_THREADS"
+  fi
+  echo "NUM_THREADS_OPTS: $NUM_THREADS_OPTS"
+
   export simdjson_SOURCE=BUNDLED
   if [ $ARCH == 'x86_64' ]; then
-    make $COMPILE_TYPE EXTRA_CMAKE_FLAGS="${COMPILE_OPTION}"
+    make $COMPILE_TYPE $NUM_THREADS_OPTS EXTRA_CMAKE_FLAGS="${COMPILE_OPTION}"
   elif [[ "$ARCH" == 'arm64' || "$ARCH" == 'aarch64' ]]; then
-    CPU_TARGET=$ARCH make $COMPILE_TYPE EXTRA_CMAKE_FLAGS="${COMPILE_OPTION}"
+    CPU_TARGET=$ARCH make $COMPILE_TYPE $NUM_THREADS_OPTS EXTRA_CMAKE_FLAGS="${COMPILE_OPTION}"
   else
     echo "Unsupported arch: $ARCH"
     exit 1