add a parallelism manifest.yaml entry, export envvars at the bott…

…om so they can reference variables from the script (#326) * export the envvars twice so they can reference variables from the script * add a parallelism section to manifest.yaml, only export envvars once right before building * bump feature version
rapidsai · Jun 12, 2024 · 37e6231 · 37e6231
1 parent ad87b5d
commit 37e6231
Show file tree

Hide file tree

Showing 8 changed files with 68 additions and 40 deletions.
diff --git a/features/src/rapids-build-utils/devcontainer-feature.json b/features/src/rapids-build-utils/devcontainer-feature.json
@@ -1,7 +1,7 @@
 {
   "name": "NVIDIA RAPIDS devcontainer build utilities",
   "id": "rapids-build-utils",
-  "version": "24.8.5",
+  "version": "24.8.6",
   "description": "A feature to install the RAPIDS devcontainer build utilities",
   "containerEnv": {
     "BASH_ENV": "/etc/bash.bash_env"

diff --git a/features/src/rapids-build-utils/opt/rapids-build-utils/bin/generate-scripts.sh b/features/src/rapids-build-utils/opt/rapids-build-utils/bin/generate-scripts.sh
@@ -46,7 +46,7 @@ generate_script() {
     if test -n "${bin}"; then
         (
             cat - \
-          | envsubst '$HOME $NAME $SRC_PATH $PY_ENV $PY_SRC $PY_LIB $BIN_DIR $CPP_ENV $CPP_LIB $CPP_SRC $CPP_CMAKE_ARGS $CPP_CPACK_ARGS $CPP_DEPS $GIT_TAG $GIT_SSH_URL $GIT_HTTPS_URL $GIT_REPO $GIT_HOST $GIT_UPSTREAM $PIP_WHEEL_ARGS $PIP_INSTALL_ARGS' \
+          | envsubst '$HOME $NAME $SRC_PATH $PY_ENV $PY_SRC $PY_LIB $BIN_DIR $CPP_ENV $CPP_LIB $CPP_SRC $CPP_CMAKE_ARGS $CPP_CPACK_ARGS $CPP_DEPS $CPP_MAX_TOTAL_SYSTEM_MEMORY $CPP_MAX_DEVICE_OBJ_MEMORY_USAGE $CPP_MAX_DEVICE_OBJ_TO_COMPILE_IN_PARALLEL $GIT_TAG $GIT_SSH_URL $GIT_HTTPS_URL $GIT_REPO $GIT_HOST $GIT_UPSTREAM $PIP_WHEEL_ARGS $PIP_INSTALL_ARGS' \
           | tee "${TMP_SCRIPT_DIR}/${bin}" >/dev/null;
 
             chmod +x "${TMP_SCRIPT_DIR}/${bin}";
@@ -185,6 +185,9 @@ generate_scripts() {
     local cpp_cmake_args;
     local cpp_cpack_args;
     local cpp_depends_length;
+    local cpp_max_total_system_memory;
+    local cpp_max_device_obj_memory_usage;
+    local cpp_max_device_obj_to_compile_in_parallel;
 
     local py_env;
     local py_path;
@@ -249,6 +252,9 @@ generate_scripts() {
             cpp_cmake_args="${repo}_cpp_${j}_args_cmake";
             cpp_cpack_args="${repo}_cpp_${j}_args_cpack";
             cpp_depends_length="${repo}_cpp_${j}_depends_length";
+            cpp_max_total_system_memory="${repo}_cpp_${j}_parallelism_max_total_system_memory";
+            cpp_max_device_obj_memory_usage="${repo}_cpp_${j}_parallelism_max_device_obj_memory_usage";
+            cpp_max_device_obj_to_compile_in_parallel="${repo}_cpp_${j}_parallelism_max_device_obj_to_compile_in_parallel";
             cpp_path=~/"${!repo_path:-}${!cpp_sub_dir:+/${!cpp_sub_dir}}";
 
             cpp_dirs+=("${cpp_path}");
@@ -278,16 +284,19 @@ generate_scripts() {
             cpp_name_to_deps["${cpp_name}"]="${cpp_deps[*]}";
 
             if [[ -d ~/"${!repo_path:-}/.git" ]]; then
-                NAME="${repo_name:-}"                 \
-                SRC_PATH=~/"${!repo_path:-}"          \
-                BIN_DIR="${bin_dir}"                  \
-                CPP_ENV="${!cpp_env:-}"               \
-                CPP_LIB="${cpp_name:-}"               \
-                CPP_SRC="${!cpp_sub_dir:-}"           \
-                CPP_DEPS="${cpp_deps[*]}"             \
-                CPP_CMAKE_ARGS="${!cpp_cmake_args:-}" \
-                CPP_CPACK_ARGS="${!cpp_cpack_args:-}" \
-                generate_cpp_scripts                  ;
+                NAME="${repo_name:-}"                                                                       \
+                SRC_PATH=~/"${!repo_path:-}"                                                                \
+                BIN_DIR="${bin_dir}"                                                                        \
+                CPP_ENV="${!cpp_env:-}"                                                                     \
+                CPP_LIB="${cpp_name:-}"                                                                     \
+                CPP_SRC="${!cpp_sub_dir:-}"                                                                 \
+                CPP_DEPS="${cpp_deps[*]}"                                                                   \
+                CPP_CMAKE_ARGS="${!cpp_cmake_args:-}"                                                       \
+                CPP_CPACK_ARGS="${!cpp_cpack_args:-}"                                                       \
+                CPP_MAX_TOTAL_SYSTEM_MEMORY="${!cpp_max_total_system_memory:-}"                             \
+                CPP_MAX_DEVICE_OBJ_MEMORY_USAGE="${!cpp_max_device_obj_memory_usage:-}"                     \
+                CPP_MAX_DEVICE_OBJ_TO_COMPILE_IN_PARALLEL="${!cpp_max_device_obj_to_compile_in_parallel:-}" \
+                generate_cpp_scripts                                                                        ;
             fi
         done
 

diff --git a/features/src/rapids-build-utils/opt/rapids-build-utils/bin/get-num-archs-jobs-and-load.sh b/features/src/rapids-build-utils/opt/rapids-build-utils/bin/get-num-archs-jobs-and-load.sh
@@ -52,7 +52,16 @@ get_num_archs_jobs_and_load() {
 
     parallel="${j:-${JOBS:-${PARALLEL_LEVEL:-1}}}";
     max_archs="${max_archs:-${MAX_DEVICE_OBJ_TO_COMPILE_IN_PARALLEL:-${arch:-}}}";
-    max_device_obj_memory_usage="${max_device_obj_memory_usage:-${MAX_DEVICE_OBJ_MEMORY_USAGE:-1}}";
+    max_device_obj_memory_usage="${max_device_obj_memory_usage:-${MAX_DEVICE_OBJ_MEMORY_USAGE:-1Gi}}";
+
+    local num_re="^[0-9]+$";
+
+    # Assume un-suffixed inputs means gibibytes
+    if [[ "${max_device_obj_memory_usage}" =~ ${num_re} ]]; then
+        max_device_obj_memory_usage="${max_device_obj_memory_usage}Gi";
+    fi
+
+    max_device_obj_memory_usage="$(numfmt --from=auto "${max_device_obj_memory_usage}")";
 
     local n_arch="${archs:-1}";
 
@@ -94,10 +103,18 @@ get_num_archs_jobs_and_load() {
     fi
 
     local mem_for_device_objs="$((n_arch * max_device_obj_memory_usage))";
+    local mem_total="${max_total_system_memory:-${MAX_TOTAL_SYSTEM_MEMORY:-}}";
+
+    if test -z "${mem_total}"; then
+        local -r free_mem="$(free --bytes | grep -E '^Mem:' | tr -s '[:space:]' | cut -d' ' -f7 || echo '0')";
+        local -r freeswap="$(free --bytes | grep -E '^Swap:' | tr -s '[:space:]' | cut -d' ' -f4 || echo '0')";
+        mem_total="$((free_mem + freeswap))";
+    # Assume un-suffixed inputs means gibibytes
+    elif [[ "${mem_total}" =~ ${num_re} ]]; then
+        mem_total="${mem_total}Gi";
+    fi
+    mem_total="$(numfmt --from=auto "${mem_total}")";
 
-    local -r free_mem="$(free --gibi | grep -E '^Mem:' | tr -s '[:space:]' | cut -d' ' -f7 || echo '0')";
-    local -r freeswap="$(free --gibi | grep -E '^Swap:' | tr -s '[:space:]' | cut -d' ' -f4 || echo '0')";
-    local -r mem_total="${max_total_system_memory:-${MAX_TOTAL_SYSTEM_MEMORY:-$((free_mem + freeswap))}}";
     local n_load=$((parallel > n_cpus ? n_cpus : parallel));
     # shellcheck disable=SC2155
     local n_jobs="$(

diff --git a/features/src/rapids-build-utils/opt/rapids-build-utils/bin/tmpl/cpp.build.tmpl.sh b/features/src/rapids-build-utils/opt/rapids-build-utils/bin/tmpl/cpp.build.tmpl.sh
@@ -17,12 +17,13 @@ build_${CPP_LIB}_cpp() {
     local -;
     set -euo pipefail;
 
-    export ${CPP_ENV} PATH="$PATH";
-
-    eval "$(                                          \
-    PARALLEL_LEVEL=${PARALLEL_LEVEL:-$(nproc --all)}  \
-        rapids-get-num-archs-jobs-and-load "$@"       \
-        2>/dev/null                                   \
+    eval "$(\
+    PARALLEL_LEVEL=${PARALLEL_LEVEL:-$(nproc --all)} \
+    MAX_TOTAL_SYSTEM_MEMORY="${MAX_TOTAL_SYSTEM_MEMORY:-${CPP_MAX_TOTAL_SYSTEM_MEMORY}}" \
+    MAX_DEVICE_OBJ_MEMORY_USAGE="${MAX_DEVICE_OBJ_MEMORY_USAGE:-${CPP_MAX_DEVICE_OBJ_MEMORY_USAGE}}" \
+    MAX_DEVICE_OBJ_TO_COMPILE_IN_PARALLEL="${MAX_DEVICE_OBJ_TO_COMPILE_IN_PARALLEL:-${CPP_MAX_DEVICE_OBJ_TO_COMPILE_IN_PARALLEL}}" \
+        rapids-get-num-archs-jobs-and-load "$@" \
+        2>/dev/null \
     )";
 
     local -a cmake_args_="(${CMAKE_ARGS:-})";
@@ -44,6 +45,7 @@ build_${CPP_LIB}_cpp() {
 
     # Build C++ lib
     time (
+        export ${CPP_ENV} PATH="$PATH";
         local -a cmake_build_args="($(rapids-select-cmake-build-args ${n_jobs:+-j${n_jobs}} "${OPTS[@]}"))";
         cmake                               \
             --build "${CPP_SRC}/${BIN_DIR}" \

diff --git a/features/src/rapids-build-utils/opt/rapids-build-utils/bin/tmpl/cpp.configure.tmpl.sh b/features/src/rapids-build-utils/opt/rapids-build-utils/bin/tmpl/cpp.configure.tmpl.sh
@@ -19,12 +19,13 @@ configure_${CPP_LIB}_cpp() {
     local -;
     set -euo pipefail;
 
-    export ${CPP_ENV} PATH="$PATH";
-
-    eval "$(                                          \
-    PARALLEL_LEVEL=${PARALLEL_LEVEL:-$(nproc --all)}  \
-        rapids-get-num-archs-jobs-and-load "$@"       \
-        2>/dev/null                                   \
+    eval "$(\
+    PARALLEL_LEVEL=${PARALLEL_LEVEL:-$(nproc --all)} \
+    MAX_TOTAL_SYSTEM_MEMORY="${MAX_TOTAL_SYSTEM_MEMORY:-${CPP_MAX_TOTAL_SYSTEM_MEMORY}}" \
+    MAX_DEVICE_OBJ_MEMORY_USAGE="${MAX_DEVICE_OBJ_MEMORY_USAGE:-${CPP_MAX_DEVICE_OBJ_MEMORY_USAGE}}" \
+    MAX_DEVICE_OBJ_TO_COMPILE_IN_PARALLEL="${MAX_DEVICE_OBJ_TO_COMPILE_IN_PARALLEL:-${CPP_MAX_DEVICE_OBJ_TO_COMPILE_IN_PARALLEL}}" \
+        rapids-get-num-archs-jobs-and-load "$@" \
+        2>/dev/null \
     )";
 
     local -a cmake_args_="(${CMAKE_ARGS:-})";
@@ -55,6 +56,7 @@ configure_${CPP_LIB}_cpp() {
     cmake_args+=(-B "${bin_dir:-${CPP_SRC}/${BIN_DIR}}");
 
     time (
+        export ${CPP_ENV} PATH="$PATH";
         CUDAFLAGS="${CUDAFLAGS:+$CUDAFLAGS }-t=${n_arch}" \
             cmake "${cmake_args[@]}";
         { set +x; } 2>/dev/null; echo -n "lib${CPP_LIB} configure time:";

diff --git a/features/src/rapids-build-utils/opt/rapids-build-utils/bin/tmpl/python.build.wheel.tmpl.sh b/features/src/rapids-build-utils/opt/rapids-build-utils/bin/tmpl/python.build.wheel.tmpl.sh
@@ -21,8 +21,6 @@ build_${PY_LIB}_python_wheel() {
     local -;
     set -euo pipefail;
 
-    export ${PY_ENV} PATH="$PATH";
-
     eval "$(                                          \
     PARALLEL_LEVEL=${PARALLEL_LEVEL:-$(nproc --all)}  \
         rapids-get-num-archs-jobs-and-load "$@"       \
@@ -87,6 +85,7 @@ build_${PY_LIB}_python_wheel() {
     trap "rm -rf '${PY_SRC}/${py_lib//"-"/"_"}.egg-info'" EXIT;
 
     time (
+        export ${PY_ENV} PATH="$PATH";
         local cudaflags="${CUDAFLAGS:+$CUDAFLAGS }-t=${n_arch}";
         local build_type="$(rapids-select-cmake-build-type "${cmake_args_[@]}")";
         local nvcc_append_flags="${NVCC_APPEND_FLAGS:+$NVCC_APPEND_FLAGS }-t=${n_arch}";

diff --git a/features/src/rapids-build-utils/opt/rapids-build-utils/bin/tmpl/python.install.tmpl.sh b/features/src/rapids-build-utils/opt/rapids-build-utils/bin/tmpl/python.install.tmpl.sh
@@ -21,8 +21,6 @@ install_${PY_LIB}_python() {
     local -;
     set -euo pipefail;
 
-    export ${PY_ENV} PATH="$PATH";
-
     eval "$(                                          \
     PARALLEL_LEVEL=${PARALLEL_LEVEL:-$(nproc --all)}  \
         rapids-get-num-archs-jobs-and-load "$@"       \
@@ -105,6 +103,7 @@ install_${PY_LIB}_python() {
     trap "rm -rf '${PY_SRC}/${py_lib//"-"/"_"}.egg-info'" EXIT;
 
     time (
+        export ${PY_ENV} PATH="$PATH";
         local cudaflags="${CUDAFLAGS:+$CUDAFLAGS }-t=${n_arch}";
         local build_type="$(rapids-select-cmake-build-type "${cmake_args_[@]}")";
         local nvcc_append_flags="${NVCC_APPEND_FLAGS:+$NVCC_APPEND_FLAGS }-t=${n_arch}";

diff --git a/features/src/rapids-build-utils/opt/rapids-build-utils/manifest.yaml b/features/src/rapids-build-utils/opt/rapids-build-utils/manifest.yaml
@@ -84,8 +84,8 @@ repos:
     - name: raft
       sub_dir: cpp
       depends: [rmm]
-      env: |
-        MAX_DEVICE_OBJ_MEMORY_USAGE=3
+      parallelism:
+        max_device_obj_memory_usage: 3Gi
       args: {cmake: -DRAFT_COMPILE_LIBRARY=ON}
   python:
     - name: pylibraft
@@ -104,8 +104,8 @@ repos:
     - name: cuvs
       sub_dir: cpp
       depends: [raft]
-      env: |
-        MAX_DEVICE_OBJ_MEMORY_USAGE=3
+      parallelism:
+        max_device_obj_memory_usage: 3Gi
       args: {cmake: -DBUILD_C_LIBRARY=ON}
   python:
     - name: cuvs
@@ -128,8 +128,8 @@ repos:
     - name: cuml
       sub_dir: cpp
       depends: [cumlprims_mg]
-      env: |
-        MAX_DEVICE_OBJ_MEMORY_USAGE=3
+      parallelism:
+        max_device_obj_memory_usage: 3Gi
   python:
     - name: cuml
       sub_dir: python
@@ -188,8 +188,8 @@ repos:
     - name: cugraph
       sub_dir: cpp
       depends: [cugraph-ops]
-      env: |
-        MAX_DEVICE_OBJ_MEMORY_USAGE=6
+      parallelism:
+        max_device_obj_memory_usage: 6Gi
     - name: cugraph_etl
       sub_dir: cpp/libcugraph_etl
       depends: [cudf, cugraph]