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

Author: Himani Anil Deshpande

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/resources/nvidia_imex/partial/_nvidia_imex_common.rb

@@ -58,8 +58,8 @@
 end
 
 action :configure do
-  return unless imex_installed
-  # Start nvidia-imex on p6e-gb200
+  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
     service nvidia_imex_service do
       action %i(start enable)

cookbooks/aws-parallelcluster-platform/spec/unit/resources/nvidia_imex_spec.rb

@@ -257,7 +257,7 @@ def self.configure(chef_run)
 describe 'nvidia_imex:configure' do
   for_all_oses do |platform, version|
     context "on #{platform}#{version}" do
-      context 'when nvidia-imex binary is not installed' do
+      context "when nvidia-imex binary is not installed" do
         cached(:chef_run) do
           stubs_for_resource('nvidia_imex') do |res|
             allow(res).to receive(:imex_installed).and_return(false)
@@ -272,34 +272,37 @@ def self.configure(chef_run)
         end
       end
 
-      context 'when get_nvswitch_count > 1' do
-        cached(:chef_run) do
-          stubs_for_provider('nvidia_imex[configure]') do |pro|
-            allow(pro).to receive(:imex_installed).and_return(true)
-            allow(pro).to receive(:get_device_ids).and_return({ 'gb200' => 'test' })
-            allow(pro).to receive(:get_nvswitch_count).with('test').and_return(4)
+      %w(HeadNode LoginNode ComputeFleet).each do |node_type|
+        context "when get_nvswitch_count > 1 on #{node_type} node" do
+          cached(:chef_run) do
+            stubs_for_provider('nvidia_imex[configure]') do |pro|
+              allow(pro).to receive(:imex_installed).and_return(true)
+              allow(pro).to receive(:get_device_ids).and_return({ 'gb200' => 'test' })
+              allow(pro).to receive(:get_nvswitch_count).with('test').and_return(4)
+            end
+            runner(platform: platform, version: version, step_into: ['nvidia_imex'])
           end
-          runner = runner(platform: platform, version: version, step_into: ['nvidia_imex'])
-          ConvergeNvidiaImex.configure(runner)
-        end
-        cached(:node) { chef_run.node }
+          cached(:node) { chef_run.node }
 
-        before do
-          chef_run.node.override['cluster']['region'] = 'aws_region'
-        end
-
-        if platform == 'amazon' && version == '2'
-          it 'does not configure nvidia-imex' do
-            is_expected.not_to start_service('nvidia-imex').with_action(%i(start enable)).with_supports({ status: true })
+          before do
+            chef_run.node.override['cluster']['region'] = 'aws_region'
+            chef_run.node.override['cluster']['node_type'] = node_type
+            ConvergeNvidiaImex.configure(chef_run)
           end
-        else
-          it 'starts nvidia-imex service' do
-            is_expected.to start_service('nvidia-imex').with_action(%i(start enable)).with_supports({ status: true })
+
+          if (platform == 'amazon' && version == '2') || %w(HeadNode LoginNode).include?(node_type)
+            it 'does not configure nvidia-imex' do
+              is_expected.not_to start_service('nvidia-imex').with_action(%i(start enable)).with_supports({ status: true })
+            end
+          else
+            it 'it starts nvidia-imex service' do
+              is_expected.to start_service('nvidia-imex').with_action(%i(start enable)).with_supports({ status: true })
+            end
           end
         end
       end
 
-      context 'when get_nvswitch_count <= 1' do
+      context "when get_nvswitch_count <= 1" do
         cached(:chef_run) do
           stubs_for_provider('nvidia_imex[configure]') do |pro|
             allow(pro).to receive(:imex_installed).and_return(true)