fix(gpu): Stabilize NodeManager startup and enhance script robustness

This commit incorporates several fixes to address NodeManager startup issues, particularly on Rocky Linux, and improves the overall robustness and idempotency of the GPU initialization script.

Key changes:

1.  **Corrected Systemd Service Order:**
    -   The `dataproc-gpu-config.service` is now configured to run *Before* `hadoop-yarn-nodemanager.service` and `hadoop-yarn-resourcemanager.service`.
    -   YARN services are now restarted using `ExecStartPost` directives within the `dataproc-gpu-config.service` unit, ensuring configurations are applied prior to service startup.
    -   The internal restart loop in the generated `apply-dataproc-gpu-config.sh` is no longer needed and has been removed.

2.  **Rocky Linux NodeManager Fix:**
    -   Added `export HADOOP_NICENESS=0` to `/etc/hadoop/conf/hadoop-env.sh` on Rocky Linux systems to prevent NodeManager startup failures due to `renice` permission issues.

3.  **Improved Idempotency:**
    -   `configure_gpu_isolation` now checks for existing lines in `container-executor.cfg` and `yarn-env.sh` before adding configurations, preventing duplicates.

4.  **Hardened MIG Script Fetching:**
    -   The script will now exit with an error if MIG is enabled but the required MIG scripts cannot be fetched.

5.  **Deferred Script Logging:**
    -   The generated `apply-dataproc-gpu-config.sh` script now defines `install_log` to ensure `execute_with_retries` can log properly.

6.  **Debugging:**
    -   Added more verbose output around the `systemctl daemon-reload` call in the generated script to help diagnose any potential issues.
    -   Redirected stderr to stdout when executing the generated script in the main script for better log capture.

Author: cjac

gpu/install_gpu_driver.sh

@@ -1708,15 +1708,38 @@ function configure_gpu_isolation() {
      echo "Hadoop conf dir ${HADOOP_CONF_DIR} not found. Skipping GPU isolation config."
      return
   fi
+
   # enable GPU isolation
   sed -i "s/yarn\.nodemanager\.linux\-container\-executor\.group\=.*$/yarn\.nodemanager\.linux\-container\-executor\.group\=yarn/g" "${HADOOP_CONF_DIR}/container-executor.cfg"
+
+  local cfg_file="${HADOOP_CONF_DIR}/container-executor.cfg"
+  # Ensure [gpu] section exists and module.enabled=true
+  if ! grep -q "^\[gpu\]" "${cfg_file}"; then
+    printf '\n[gpu]\nmodule.enabled=true\n' >> "${cfg_file}"
+  elif ! grep -q "^module.enabled=true" "${cfg_file}"; then
+    # Assuming [gpu] exists, add module.enabled if missing
+    sed -i "/\[gpu\]/a module.enabled=true" "${cfg_file}"
+  fi
+
   if [[ $IS_MIG_ENABLED -ne 0 ]]; then
-    # configure the container-executor.cfg to have major caps
-    printf '\n[gpu]\nmodule.enabled=true\ngpu.major-device-number=%s\n\n[cgroups]\nroot=/sys/fs/cgroup\nyarn-hierarchy=yarn\n' $MIG_MAJOR_CAPS >> "${HADOOP_CONF_DIR}/container-executor.cfg"
-    printf 'export MIG_AS_GPU_ENABLED=1\n' >> "${HADOOP_CONF_DIR}/yarn-env.sh"
-    printf 'export ENABLE_MIG_GPUS_FOR_CGROUPS=1\n' >> "${HADOOP_CONF_DIR}/yarn-env.sh"
-  else
-    printf '\n[gpu]\nmodule.enabled=true\n[cgroups]\nroot=/sys/fs/cgroup\nyarn-hierarchy=yarn\n' >> "${HADOOP_CONF_DIR}/container-executor.cfg"
+    if ! grep -q "^gpu.major-device-number=" "${cfg_file}"; then
+      sed -i "/\[gpu\]/a gpu.major-device-number=${MIG_MAJOR_CAPS}" "${cfg_file}"
+    fi
+  fi
+
+  # Ensure [cgroups] section and its contents
+  if ! grep -q "^\[cgroups\]" "${cfg_file}"; then
+    printf '\n[cgroups]\nroot=/sys/fs/cgroup\nyarn-hierarchy=yarn\n' >> "${cfg_file}"
+  fi
+
+  local yarn_env_file="${HADOOP_CONF_DIR}/yarn-env.sh"
+  if [[ $IS_MIG_ENABLED -ne 0 ]]; then
+    if ! grep -q "^export MIG_AS_GPU_ENABLED=" "${yarn_env_file}"; then
+      printf 'export MIG_AS_GPU_ENABLED=1\n' >> "${yarn_env_file}"
+    fi
+    if ! grep -q "^export ENABLE_MIG_GPUS_FOR_CGROUPS=" "${yarn_env_file}"; then
+      printf 'export ENABLE_MIG_GPUS_FOR_CGROUPS=1\n' >> "${yarn_env_file}"
+    fi
   fi
 
   # Configure a systemd unit to ensure that permissions are set on restart
@@ -1974,14 +1997,30 @@ function run_hadoop_spark_config() {
           NUM_MIG_GPUS="$(echo ${migquery_result} | uniq | wc -l)"
           if [[ "${NUM_MIG_GPUS}" -eq "1" ]] && (echo "${migquery_result}" | grep -q Enabled); then
             IS_MIG_ENABLED=1
+            if [[ ! -d "/usr/local/yarn-mig-scripts" ]]; then
+              if ! fetch_mig_scripts; then
+                echo "ERROR: Failed to fetch MIG scripts, which are required because MIG is enabled." >&2
+                return 1
+              fi
+            fi
             NVIDIA_SMI_PATH='/usr/local/yarn-mig-scripts/' # Set MIG path
             MIG_MAJOR_CAPS=$(grep nvidia-caps /proc/devices | cut -d ' ' -f 1 || echo 0)
-            if [[ ! -d "/usr/local/yarn-mig-scripts" ]]; then fetch_mig_scripts || echo "WARN: Failed to fetch MIG scripts." >&2; fi
           fi
       fi
   fi
 
+  # Prevent renice errors on systems like Rocky Linux
+  if is_rocky; then
+    if [[ -f "${HADOOP_CONF_DIR}/hadoop-env.sh" ]]; then
+      if ! grep -q "^export HADOOP_NICENESS=" "${HADOOP_CONF_DIR}/hadoop-env.sh"; then
+        echo "export HADOOP_NICENESS=0" >> "${HADOOP_CONF_DIR}/hadoop-env.sh"
+        echo "Added HADOOP_NICENESS=0 to hadoop-env.sh"
+      fi
+    fi
+  fi
+
   # Ensure config directories exist
+
   if [[ ! -d "${HADOOP_CONF_DIR}" || ! -d "${SPARK_CONF_DIR}" ]]; then
      echo "ERROR: Config directories (${HADOOP_CONF_DIR}, ${SPARK_CONF_DIR}) not found. Cannot apply configuration."
      return 1 # Use return instead of exit in a function
@@ -2025,6 +2064,9 @@ function create_deferred_config_files() {
 # Deferred configuration script generated by install_gpu_driver.sh
 set -xeuo pipefail
 
+# Define install_log for execute_with_retries
+readonly install_log="/tmp/apply-dataproc-gpu-config.log"
+
 # --- Minimal necessary functions and variables ---
 # Define constants
 readonly HADOOP_CONF_DIR='/etc/hadoop/conf'
@@ -2157,25 +2199,13 @@ curl_retry_args="-fsSL --retry-connrefused --retry 10 --retry-max-time 30"
 $(declare -f run_hadoop_spark_config)
 
 # --- Execute the config logic ---
-if run_hadoop_spark_config; then
-  # Configuration successful, disable the service
-  systemctl disable ${service_name}.service
-  rm -f "${config_script_path}" "${service_file}"
-  systemctl daemon-reload
-else
-  echo "ERROR: Deferred configuration script (${config_script_path}) failed." >&2
-  # Keep the service enabled to allow for manual inspection/retry
-  exit 1
-fi
-
-# Restart services after applying config
-for svc in resourcemanager nodemanager; do
-  if (systemctl is-active --quiet hadoop-yarn-\${svc}.service); then
-    systemctl stop  hadoop-yarn-\${svc}.service || echo "WARN: Failed to stop \${svc}"
-    systemctl start hadoop-yarn-\${svc}.service || echo "WARN: Failed to start \${svc}"
-  fi
-done
-
+run_hadoop_spark_config
+# Configuration successful, disable the service
+systemctl disable dataproc-gpu-config.service
+rm -f "${config_script_path}" "${service_file}"
+echo "Reloading systemd manager configuration..."
+systemctl daemon-reload
+echo "systemd daemon-reload complete."
 exit 0
 EOF
 
@@ -2184,13 +2214,15 @@ EOF
   cat <<EOF > "${service_file}"
 [Unit]
 Description=Apply Dataproc GPU configuration on first boot
-# Ensure it runs after Dataproc agent and YARN services are likely up
-After=google-dataproc-agent.service network-online.target hadoop-yarn-resourcemanager.service hadoop-yarn-nodemanager.service
+After=network-online.target google-dataproc-agent.service
 Wants=network-online.target google-dataproc-agent.service
+Before=hadoop-yarn-resourcemanager.service hadoop-yarn-nodemanager.service
 
 [Service]
 Type=oneshot
 ExecStart=${config_script_path} # Execute the generated config script
+ExecStartPost=/bin/systemctl try-restart hadoop-yarn-resourcemanager.service
+ExecStartPost=/bin/systemctl try-restart hadoop-yarn-nodemanager.service
 RemainAfterExit=no # Service is done after exec
 StandardOutput=journal+console
 StandardError=journal+console
@@ -2301,7 +2333,7 @@ function main() {
     # Running as a standard init action: execute the generated script immediately
     local -r config_script_path="/usr/local/sbin/apply-dataproc-gpu-config.sh"
     if [[ -x "${config_script_path}" ]]; then
-        bash -x "${config_script_path}"
+        bash -x "${config_script_path}" 2>&1
     else
         echo "ERROR: Generated config script ${config_script_path} not found or not executable."
         exit 1
@@ -2550,8 +2582,8 @@ function set_proxy(){
              storage  datafusion    dataproc certificatemanager   networksecurity
              dataflow privateca     logging )
 
-  for svc in "${services[@]}"; do
-    no_proxy_list+=("${svc}.googleapis.com")
+  for gcpsvc in "${services[@]}"; do
+    no_proxy_list+=("${gcpsvc}.googleapis.com")
   done
 
   no_proxy="$( IFS=',' ; echo "${no_proxy_list[*]}" )"
@@ -2749,9 +2781,9 @@ function harden_sshd_config() {
     # TODO: test whether sshd will reload with this change before mv
     mv -f /tmp/sshd_config_new /etc/ssh/sshd_config
   done
-  local svc=ssh
-  if is_rocky ; then svc="sshd" ; fi
-  systemctl reload "${svc}"
+  local sshsvc=ssh
+  if is_rocky ; then sshsvc="sshd" ; fi
+  systemctl reload "${sshsvc}"
 }
 
 function prepare_to_install(){