[Nvidia-Imex] Installing Nvidia-imex as part of ParallelCluster Build Image (#2996)

* [NVIDIA_IMEX] Add resource to install Nvidia-imex

* [NVIDIA_IMEX] Adding Unit test for IMEX installation

* [NVIDIA_IMEX] Not Install Nvidia-imex for Isolated regions

* [NVIDIA_IMEX] Install Nvdia-Imex as part of Build Image

* [Nvidia-imex] Never Install NVIDIA Imex for AL2

* [Nvidia-imex] Add unit tests for NVidia Imex

* [Nvidia-imex] Cookstyle changes

* [NVIDIA_IMEX] Adding Kitchen test for Installation and Configuration

* [FABRIC MANAGER] Using common library for getting NVSwitch count

* [NVIDIA-IMEX] Configure Nvidia-imex only if we use Gb200 instance

* [NVIDIA-IMEX] USe specific Version naming for nvidia-imex installation

* [NVIDIA-IMEX] Install Nvidia-imex and flush cache before it

* [NVIDIA-IMEX] Redirect nvidia-imex to system logs which are pushed in CW

* [NVIDIA-IMEX] Install With specific version in name

* [NVIDIA-IMEX] Removing flush cache as it does not exist for package resource

* [NVIDIA-IMEX] Installing NVIDIA IMEx using install_packages resource

* [NVIDIA-IMEX] Adding Unit tests for Configuration of nvidia-imex

* [NVIDIA-IMEX] Configuring nvidia-imex only for gb200 and ComputeFleet node

* [NVIDIA-IMEX] Not check installation of nvidia-imex for Alinux2

* [NVIDIA-IMEX] Inspec Test

* [NVIDIA-IMEX] Setting Nvidia-imex node attributes which should show the version and package name as seen on the repositories

* [NVIDIA-IMEX] Test epoch version

* [NVIDIA-IMEX] Add Version and package name for debian installation

* [NVIDIA-IMEX] Add changelog

* Add unit test for checking configuration of nvidia-imex

* [Nvidia-Imex] Use nvidia-imex shared directory for Inspec and configuration files

* [Nvidia-Imex] Update copyright year

* [Nvidia-Imex] Adding correct comments

* [Nvidia-Imex] Updating function names

* [Nvidia-Imex] Remove _nvidia_imex_version as it is not needed

* [Nvidia-Imex] Update action sequence for service

* [NVIDIA-IMEX] Comment the official docs for nvidia-imex service file

* [NVIDIA-IMEX] Using common naming convention for package name

* [NVIDIA-IMEX] Correcting kitchen test

* [NVIDIA_IMEX] Install nvidia-imex from s3

* [NVIDIA_IMEX] Install nvidia-imex from s3

* [NVIDIA_IMEX] Update unit tests

* [NVIDIA Driver] Upgrade NVIDIA driver to 570.172.08 for all except AL2

Removing test URL

---------

Co-authored-by: Himani Anil Deshpande <[email protected]>

Author: himani2411

CHANGELOG.md

@@ -15,6 +15,7 @@ This file is used to list changes made in each version of the AWS ParallelCluste
 - Addressed cluster id mismatch known issue by deleting the file `/var/spool/slurm.state/clustername` before configuring Slurm accounting.
 - Upgrade DCV to version 2024.0-19030.
 - Remove `berkshelf`. All cookbooks are local and do not need `berkshelf` dependency management.
+- Install nvidia-imex for all OSs except AL2.
 
 **BUG FIXES**
 - Fix a race condition in CloudWatch Agent startup that could cause nodes bootstrap failures.
@@ -38,6 +39,7 @@ This file is used to list changes made in each version of the AWS ParallelCluste
   - Libfabric-aws: libfabric-aws-2.1.0-1
   - Rdma-core: rdma-core-57.0-1
   - Open MPI: openmpi40-aws-4.1.7-2 and openmpi50-aws-5.0.6
+- Upgrade NVIDIA driver to version 570.172.08 (from 570.86.15) for all OSs except AL2.
 
 **BUG FIXES**
 - Fix a bug in the installation of ARM Performance Library that was causing the build image fail in isolated environments.

cookbooks/aws-parallelcluster-platform/attributes/platform.rb

@@ -16,12 +16,15 @@
 
 # NVidia
 default['cluster']['nvidia']['enabled'] = 'no'
-default['cluster']['nvidia']['driver_version'] = '570.86.15'
+default['cluster']['nvidia']['driver_version'] = '570.172.08'
 default['cluster']['nvidia']['dcgm_version'] = '3.3.6'
 if platform?('amazon') && node['platform_version'] == "2"
   default['cluster']['nvidia']['driver_version'] = '550.127.08'
 end
 
+# nvidia-imex
+default['cluster']['nvidia']['imex']['shared_dir'] = "#{node['cluster']['shared_dir']}/nvidia-imex"
+
 # DCV
 default['cluster']['dcv']['authenticator']['user'] = "dcvextauth"
 default['cluster']['dcv']['authenticator']['user_id'] = node['cluster']['reserved_base_uid'] + 3

cookbooks/aws-parallelcluster-platform/libraries/nvidia.rb

@@ -18,3 +18,19 @@ def is_process_running(process_name)
 
   !ps.stdout.strip.empty?
 end
+
+#
+# Get Count of GPUs in instance
+#
+def get_nvswitch_count(device_id)
+  shell_out("lspci -d #{device_id} | wc -l").stdout.strip.to_i
+end
+
+def get_device_ids
+  #  A100 (P4), H100(P5), B200(P6) and GB200()p6e) systems have NVSwitches
+  # NVSwitch device id is 10de:1af1 for P4 instance
+  # NVSwitch device id is 10de:22a3 for P5 instance
+  # NVSwitch device id is 10de:2901 for P6 instance
+  # NVSwitch device id is 10de:2941 for P6e instance
+  { 'a100' => '10de:1af1', 'h100' => '10de:22a3', 'b200' => '10de:2901', 'gb200' => '10de:2941' }
+end

cookbooks/aws-parallelcluster-platform/recipes/config/nvidia_config.rb

@@ -24,3 +24,7 @@
 end
 
 include_recipe "aws-parallelcluster-platform::nvidia_uvm"
+
+nvidia_imex 'Configure nvidia-imex' do
+  action :configure
+end

cookbooks/aws-parallelcluster-platform/recipes/install/nvidia_install.rb

@@ -24,3 +24,5 @@
 fabric_manager 'Install Nvidia Fabric Manager'
 
 nvidia_dcgm 'install Nvidia datacenter-gpu-manager'
+
+nvidia_imex 'Install nvidia-imex'

cookbooks/aws-parallelcluster-platform/resources/fabric_manager/partial/_fabric_manager_common.rb

@@ -54,12 +54,8 @@ def _nvidia_driver_version
 
 # Get number of nv switches
 def get_nvswitches
-  #  A100 (P4), H100(P5) and B200(P6) systems have NVSwitches
-  # NVSwitch device id is 10de:1af1 for P4 instance
-  # NVSwitch device id is 10de:22a3 for P5 instance
-  # NVSwitch device id is 10de:2901 for P6 instance
   # We sum the count for all these deviceIds as output of lscpi command will be >0
   # for only one device ID based on the instance type
-  nvswitch_device_ids = ['10de:1af1', '10de:22a3', '10de:2901']
-  nvswitch_device_ids.sum { |id| shell_out("lspci -d #{id} | wc -l").stdout.strip.to_i }
+  nvswitch_device_ids = get_device_ids.values
+  nvswitch_device_ids.sum { |id| get_nvswitch_count(id) }
 end

cookbooks/aws-parallelcluster-platform/resources/nvidia_imex/nvidia_imex_alinux2023.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+# Copyright:: 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file.
+# This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied.
+# See the License for the specific language governing permissions and limitations under the License.
+
+provides :nvidia_imex, platform: 'amazon' do |node|
+  node['platform_version'].to_i == 2023
+end
+
+use 'partial/_nvidia_imex_common.rb'
+use 'partial/_nvidia_imex_rhel.rb'
+
+def platform
+  "amzn#{node['platform_version'].to_i}"
+end

cookbooks/aws-parallelcluster-platform/resources/nvidia_imex/nvidia_imex_amazon2.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+# Copyright:: 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file.
+# This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied.
+# See the License for the specific language governing permissions and limitations under the License.
+
+provides :nvidia_imex, platform: 'amazon', platform_version: '2'
+
+use 'partial/_nvidia_imex_common.rb'
+use 'partial/_nvidia_imex_rhel.rb'
+
+def imex_installed?
+  # We do not install NVIDIA-Imex for Alinux2 due to restriction on NVIDIA driver
+  true
+end
+
+action :configure do
+  # Do nothing
+end

cookbooks/aws-parallelcluster-platform/resources/nvidia_imex/nvidia_imex_redhat8.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+# Copyright:: 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file.
+# This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied.
+# See the License for the specific language governing permissions and limitations under the License.
+
+provides :nvidia_imex, platform: 'redhat' do |node|
+  node['platform_version'].to_i >= 8
+end
+
+use 'partial/_nvidia_imex_common.rb'
+use 'partial/_nvidia_imex_rhel.rb'
+
+def platform
+  "rhel#{node['platform_version'].to_i}"
+end

cookbooks/aws-parallelcluster-platform/resources/nvidia_imex/nvidia_imex_rocky8.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+# Copyright:: 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file.
+# This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied.
+# See the License for the specific language governing permissions and limitations under the License.
+
+provides :nvidia_imex, platform: 'rocky' do |node|
+  node['platform_version'].to_i >= 8
+end
+
+use 'partial/_nvidia_imex_common.rb'
+use 'partial/_nvidia_imex_rhel.rb'
+
+def platform
+  "rhel#{node['platform_version'].to_i}"
+end