templates/spark-rapids/mig.sh.in:

* insert additional templates, common/yarn_functions and gpu/mig_functions
* MIG code is inline in the main function for mig
* call additional exit handlers: pip and yarn
* call additional prepare handler: pip

templates/spark-rapids/spark-rapids.sh.in:
* remove redundant template disclaimer
* insert additional template, common/install_functions, gpu/install_functions, gpu/yarn_functions, gpu/spark_functions
* remove redundant call to configure_gpu_script
* call additional exit handlers: gpu_install_exit_handler, yarn_exit_handler
* call additional prepare handler: prepare_gpu_install_env

Author: cjac

templates/spark-rapids/mig.sh.in

@@ -14,13 +14,55 @@
 #
 [% PROCESS common/template_disclaimer %]
 
-set -euxo pipefail
-
 [% INSERT common/util_functions %]
 
+[% INSERT common/yarn_functions %]
+
+[% INSERT gpu/mig_functions %]
+
 [% INSERT gpu/util_functions %]
 
+set -euxo pipefail
+
 function main() {
+  if [[ "${nvsmi_works}" == "1" ]] ; then
+    # if this is called without the MIG script then the drivers are not installed
+    query_nvsmi
+    local xpath='//nvidia_smi_log/*/mig_mode/current_mig/text()'
+    set +e
+    migquery_result="$("${xmllint}" --xpath "${xpath}" "${nvsmi_query_xml}" | grep -v 'N/A')"
+    set -e
+    NUM_MIG_GPUS="$(echo ${migquery_result} | uniq | wc -l)"
+
+    if [[ "${NUM_MIG_GPUS}" -gt "0" ]] ; then
+      if [[ "${NUM_MIG_GPUS}" -eq "1" ]]; then
+        if (echo "${migquery_result}" | grep Enabled); then
+          IS_MIG_ENABLED=1
+          NVIDIA_SMI_PATH='/usr/local/yarn-mig-scripts/'
+          MIG_MAJOR_CAPS=`grep nvidia-caps /proc/devices | cut -d ' ' -f 1`
+          fetch_mig_scripts
+        fi
+      fi
+    fi
+  fi
+
+  # if mig is enabled drivers would have already been installed
+  if [[ $IS_MIG_ENABLED -eq 0 ]]; then
+    install_nvidia_gpu_driver
+    install_cuda
+    load_kernel_module
+
+    #Install GPU metrics collection in Stackdriver if needed
+    if [[ "${INSTALL_GPU_AGENT}" == "true" ]]; then
+      install_gpu_agent
+#      install_gpu_monitoring_agent
+      echo 'GPU metrics agent successfully deployed.'
+    else
+      echo 'GPU metrics agent has not been installed.'
+    fi
+    configure_gpu_exclusive_mode
+  fi
+
   setup_gpu_yarn
 
   echo "yarn setup complete"
@@ -33,12 +75,15 @@ function main() {
 
 function exit_handler() {
   gpu_exit_handler
+  pip_exit_handler
+  yarn_exit_handler
   common_exit_handler
   return 0
 }
 
 function prepare_to_install(){
   prepare_common_env
+  prepare_pip_env
   prepare_gpu_env
   trap exit_handler EXIT
 }

templates/spark-rapids/spark-rapids.sh.in

@@ -22,37 +22,54 @@
 # For details see
 # github.com/GoogleCloudDataproc/custom-images/tree/main/examples/secure-boot
 #
-[% PROCESS common/template_disclaimer %]
 
 set -euxo pipefail
 
 [% INSERT common/util_functions %]
-
+[% INSERT common/install_functions %]
 [% INSERT gpu/util_functions %]
+[% INSERT gpu/install_functions %]
+[% INCLUDE gpu/yarn_functions %]
+[% INSERT gpu/spark_functions %]
 
 function main() {
+  install_gpu_driver_and_cuda
+
+  #Install GPU metrics collection in Stackdriver if needed
+  if [[ "${INSTALL_GPU_AGENT}" == "true" ]]; then
+#    install_gpu_agent
+    install_gpu_monitoring_agent
+    echo 'GPU metrics agent successfully deployed.'
+  else
+    echo 'GPU metrics agent has not been installed.'
+  fi
+  configure_gpu_exclusive_mode
+
   setup_gpu_yarn
 
   echo "yarn setup complete"
 
   if [[ "${RAPIDS_RUNTIME}" == "SPARK" ]]; then
     install_spark_rapids
-    configure_gpu_script
     echo "RAPIDS initialized with Spark runtime"
   elif [[ "${RAPIDS_RUNTIME}" == "DASK" ]]; then
-    # we are not currently tooled for installing dask in this action.
-    echo "RAPIDS recognizes DASK runtime - currently supported using dask/dask.sh or rapids/rapids.sh"
+    echo "This action only installs spark-rapids"
+    exit 1
   else
     echo "Unrecognized RAPIDS Runtime: ${RAPIDS_RUNTIME}"
+    exit 1
   fi
 
   echo "main complete"
   return 0
 }
 
 function exit_handler() {
+  set +e
+  gpu_install_exit_handler
   gpu_exit_handler
   pip_exit_handler
+  yarn_exit_handler
   common_exit_handler
   return 0
 }
@@ -61,6 +78,7 @@ function prepare_to_install(){
   prepare_common_env
   prepare_pip_env
   prepare_gpu_env
+  prepare_gpu_install_env
   trap exit_handler EXIT
 }