[Nvidia-Imex] Installing Nvidia-imex as part of ParallelCluster Build Image

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

cookbooks/aws-parallelcluster-platform/resources/nvidia_imex/nvidia_imex_ubuntu22+.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: 'ubuntu' do |node|
+  node['platform_version'].to_i >= 22
+end
+
+use 'partial/_nvidia_imex_common.rb'
+use 'partial/_nvidia_imex_debian.rb'
+
+def platform
+  "ubuntu#{node['platform_version'].delete('.')}"
+end

cookbooks/aws-parallelcluster-platform/resources/nvidia_imex/partial/_nvidia_imex_common.rb

@@ -0,0 +1,106 @@
+# 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.
+
+unified_mode true
+default_action :install
+
+action :install do
+  return unless nvidia_enabled_or_installed?
+  return if on_docker? || imex_installed? || aws_region.start_with?("us-iso")
+
+  directory node['cluster']['nvidia']['imex']['shared_dir']
+
+  action_install_imex
+  # Save Imex version in Node Attributes for InSpec Tests
+  node.default['cluster']['nvidia']['imex']['version'] = nvidia_imex_full_version
+  node.default['cluster']['nvidia']['imex']['package'] = nvidia_imex_package
+  node_attributes 'dump node attributes'
+end
+
+action :configure do
+  return unless imex_installed? && node['cluster']['node_type'] == "ComputeFleet"
+  # Start nvidia-imex on p6e-gb200 and only on ComputeFleet
+  if get_nvswitch_count(get_device_ids['gb200']) > 1
+    # For each Compute Resource, we generate a unique NVIDIA IMEX configuration file,
+    # if one doesn't already exist in a common, shared location.
+    template nvidia_imex_nodes_conf_file do
+      source 'nvidia-imex/nvidia-imex-nodes.erb'
+      owner 'root'
+      group 'root'
+      mode '0755'
+      action :create
+      not_if { file_exists_and_cluster_update?(nvidia_imex_nodes_conf_file) }
+    end
+
+    template nvidia_imex_main_conf_file do
+      source 'nvidia-imex/nvidia-imex-config.erb'
+      owner 'root'
+      group 'root'
+      mode '0755'
+      action :create
+      not_if { file_exists_and_cluster_update?(nvidia_imex_main_conf_file) }
+      variables(imex_nodes_config_file_path: nvidia_imex_nodes_conf_file)
+    end
+
+    template "/etc/systemd/system/#{nvidia_imex_service}.service" do
+      source 'nvidia-imex/nvidia-imex.service.erb'
+      owner 'root'
+      group 'root'
+      mode '0644'
+      action :create
+      variables(imex_main_config_file_path: nvidia_imex_main_conf_file)
+    end
+
+    service nvidia_imex_service do
+      action %i(enable start)
+      supports status: true
+    end
+  end
+end
+
+def nvidia_imex_package
+  "#{nvidia_imex_service}-#{nvidia_driver_major_version}"
+end
+
+def nvidia_driver_major_version
+  node['cluster']['nvidia']['driver_version'].split('.')[0]
+end
+
+def nvidia_imex_service
+  'nvidia-imex'
+end
+
+def nvidia_imex_full_version
+  "#{node['cluster']['nvidia']['driver_version']}-1"
+end
+
+def imex_installed?
+  ::File.exist?("/usr/bin/#{nvidia_imex_service}") || ::File.exist?("/usr/bin/#{nvidia_imex_service}-ctl")
+end
+
+def nvidia_enabled_or_installed?
+  nvidia_enabled? || nvidia_installed?
+end
+
+def file_exists_and_cluster_update?(file_path)
+  ::File.exist?(file_path) && !are_queues_updated?
+end
+
+def nvidia_imex_main_conf_file
+  "#{node['cluster']['nvidia']['imex']['shared_dir']}/config_#{node['cluster']['launch_template_id']}.cfg"
+end
+
+def nvidia_imex_nodes_conf_file
+  "#{node['cluster']['nvidia']['imex']['shared_dir']}/nodes_config_#{node['cluster']['launch_template_id']}.cfg"
+end

cookbooks/aws-parallelcluster-platform/resources/nvidia_imex/partial/_nvidia_imex_debian.rb

@@ -0,0 +1,42 @@
+# 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.
+
+action :install_imex do
+  remote_file "#{node['cluster']['sources_dir']}/#{nvidia_imex_package}-#{nvidia_imex_full_version}.deb" do
+    source "#{nvidia_imex_url}"
+    mode '0644'
+    retries 3
+    retry_delay 5
+    action :create_if_missing
+  end
+
+  bash "Install nvidia-imex" do
+    user 'root'
+    cwd node['cluster']['sources_dir']
+    code <<-NVIDIA_IMEX
+    set -e
+    dpkg -i #{nvidia_imex_package}-#{nvidia_imex_full_version}.deb && apt-mark hold #{nvidia_imex_package}
+    NVIDIA_IMEX
+    retries 3
+    retry_delay 5
+  end
+end
+
+def nvidia_imex_url
+  "#{node['cluster']['artifacts_s3_url']}/dependencies/nvidia_imex/#{platform}/#{nvidia_imex_package}_#{nvidia_imex_full_version}_#{arch_suffix}.deb"
+end
+
+def arch_suffix
+  arm_instance? ? 'arm64' : 'amd64'
+end