nvidia-driver version configurable

common/library/module_utils/input_validation/schema/omnia_config.json

@@ -72,6 +72,14 @@
             },
             "additionalProperties": false
           },
+          "nvidia_driver_version_x86_64": {
+            "type": "string",
+            "description": "NVIDIA driver version to install on x86_64 GPU nodes"
+          },
+          "nvidia_driver_version_aarch64": {
+            "type": "string",
+            "description": "NVIDIA driver version to install on aarch64 GPU nodes"
+          },
           "config_sources": {
             "type": "object",
             "description": "Config can be a file path or inline mapping",

input/omnia_config.yml

@@ -137,11 +137,24 @@
 #   cgroup: /opt/omnia/input/project_default/cgroup.conf
 #   slurmdbd: /opt/omnia/input/project_default/slurmdbd.conf
 
+# nvidia_driver_version
+# NVIDIA CUDA driver version to install on compute nodes.
+# These driver versions are used by install_cuda_driver.sh on Slurm nodes.
+# Two versions can be specified per cluster:
+#   - nvidia_driver_version_x86_64: driver version for x86_64 compute nodes
+#   - nvidia_driver_version_aarch64: driver version for aarch64 compute nodes
+# Default: "580.159.04" if not specified.
+# 580.159.04 is the NVIDIA CUDA driver version.
+# This version was validated on x86_64 and aarch64 compute nodes
+# running RHEL 10.0.
+
 slurm_cluster:
   - cluster_name: slurm_cluster
+    nvidia_driver_version_x86_64: "580.159.04"
+    nvidia_driver_version_aarch64: "580.159.04"
     nfs_storage_name: nfs_slurm
     vast_storage_name: vast_storage
-  
+
 # ----------------------------SERVICE K8S------------------------------------------------------
 # For service k8s cluster below parameters are required,(List)
 # - cluster_name is required field

provision/roles/configure_ochami/templates/hpc_tools/install_cuda_driver.sh.j2

@@ -3,7 +3,30 @@
 # Never touches NFS lock artifacts. Never touches /hpc_tools/cuda contents.
 set -euo pipefail
 
-DRIVER_VERSION="580.159.04"
+# NVIDIA driver version defaults.
+#
+# 580.159.04 is NOT the latest NVIDIA driver release.
+# It was chosen as the default because it is a validated and packaged
+# version in the Omnia local repository for Omnia 2.2, for both x86_64
+# and aarch64 compute nodes running RHEL 10.0.
+#
+# We do not automatically use the latest NVIDIA release because the
+# matching nvidia-driver-<version> and nvidia-driver-cuda-<version>
+# packages must exist in the local repository and be tested on the
+# target OS/arch before they can be installed safely. If you change
+# this to a newer NVIDIA release, ensure the matching packages are
+# available; otherwise the install will fail.
+
+DRIVER_VERSION_x86_64="{{ hostvars['localhost']['nvidia_driver_version_x86_64'] | default('580.159.04') }}"
+DRIVER_VERSION_aarch64="{{ hostvars['localhost']['nvidia_driver_version_aarch64'] | default('580.159.04') }}"
+
+ARCH=$(uname -m)
+if [ "$ARCH" == "aarch64" ]; then
+    DRIVER_VERSION="$DRIVER_VERSION_aarch64"
+else
+    DRIVER_VERSION="$DRIVER_VERSION_x86_64"
+fi
+
 LOGFILE="/var/log/nvidia_install.log"
 
 # ── Ensure full PATH (non-interactive shells may miss /usr/sbin) ──

provision/roles/provision_validations/tasks/include_software_config.yml

@@ -90,3 +90,8 @@
 - name: Check if openmpi support is true
   ansible.builtin.set_fact:
     openmpi_support: "{{ software_config.softwares | selectattr('name', 'equalto', 'openmpi') | list | length > 0 }}"
+
+- name: Set NVIDIA driver version from omnia_config
+  ansible.builtin.set_fact:
+    nvidia_driver_version_x86_64: "{{ slurm_cluster[0].nvidia_driver_version_x86_64 | default('580.159.04') }}"
+    nvidia_driver_version_aarch64: "{{ slurm_cluster[0].nvidia_driver_version_aarch64 | default('580.159.04') }}"

upgrade/roles/import_input_parameters/templates/omnia_config.j2

@@ -137,11 +137,24 @@
 #   cgroup: /opt/omnia/input/project_default/cgroup.conf
 #   slurmdbd: /opt/omnia/input/project_default/slurmdbd.conf
 
+# nvidia_driver_version
+# NVIDIA CUDA driver version to install on compute nodes.
+# These driver versions are used by install_cuda_driver.sh on Slurm nodes.
+# Two versions can be specified per cluster:
+#   - nvidia_driver_version_x86_64: driver version for x86_64 compute nodes
+#   - nvidia_driver_version_aarch64: driver version for aarch64 compute nodes
+# Default: "580.159.04" if not specified.
+# 580.159.04 is the NVIDIA CUDA driver version.
+# This version was validated on x86_64 and aarch64 compute nodes
+# running RHEL 10.0.
+
 slurm_cluster:
 {% set _slurm_cluster = omnia_slurm_cluster | default([], true) %}
 {% if (_slurm_cluster | length) > 0 %}
 {% for _cluster in _slurm_cluster %}
   - cluster_name: {{ _cluster.cluster_name | default('') }}
+    nvidia_driver_version_x86_64: {{ _cluster.nvidia_driver_version_x86_64 | default('580.159.04') }}
+    nvidia_driver_version_aarch64: {{ _cluster.nvidia_driver_version_aarch64 | default('580.159.04') }}
     nfs_storage_name: {{ _cluster.nfs_storage_name | default('') }}
     # vast_storage_name: "vast_storage"