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

Author: Himani Anil Deshpande

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

@@ -20,3 +20,7 @@ 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/partial/_nvidia_imex_common.rb

@@ -24,23 +24,23 @@
     action :add
   end
 
-  directory "#{node['cluster']['shared_dir']}/nvidia-imex"
+  directory "#{node['cluster']['shared_dir']}/#{nvidia_imex_service}"
 
-  template "#{node['cluster']['shared_dir']}/nvidia-imex/config.cfg" do
+  template "#{node['cluster']['shared_dir']}/#{nvidia_imex_service}/config.cfg" do
     source 'nvidia-imex/nvidia-imex-config.erb'
     owner 'root'
     group 'root'
     mode '0755'
   end
 
-  template "#{node['cluster']['shared_dir']}/nvidia-imex/nodes_config.cfg" do
+  template "#{node['cluster']['shared_dir']}/#{nvidia_imex_service}/nodes_config.cfg" do
     source 'nvidia-imex/nvidia-imex-nodes.erb'
     owner 'root'
     group 'root'
     mode '0755'
   end
 
-  template "/etc/systemd/system/nvidia-imex.service" do
+  template "/etc/systemd/system/#{nvidia_imex_service}.service" do
     source 'nvidia-imex/nvidia-imex.service.erb'
     owner 'root'
     group 'root'
@@ -49,7 +49,7 @@
   end
 
   install_packages 'Install nvidia-imex' do
-    packages "nvidia-imex-#{_nvidia_imex_version}"
+    packages "#{nvidia_imex_service}-#{_nvidia_imex_version}"
     action :install
   end
   # Save Imex version in Node Attributes for InSpec Tests
@@ -61,15 +61,19 @@
   return unless imex_installed
   # Start nvidia-imex on p6e-gb200
   if get_nvswitch_count(get_device_ids['gb200']) > 1
-    service 'nvidia-imex' do
+    service nvidia_imex_service do
       action %i(start enable)
       supports status: true
-    end unless on_docker?
+    end
   end
 end
 
+def nvidia_imex_service
+  'nvidia-imex'
+end
+
 def imex_installed
-  ::File.exist?('/usr/bin/nvidia-imex') || ::File.exist?('/usr/bin/nvidia-imex-ctl')
+  ::File.exist?("/usr/bin/#{nvidia_imex_service}") || ::File.exist?("/usr/bin/#{nvidia_imex_service}-ctl")
 end
 
 def nvidia_enabled_or_installed?

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

@@ -14,6 +14,14 @@ def self.install(chef_run)
       end
     end
   end
+
+  def self.configure(chef_run)
+    chef_run.converge_dsl('aws-parallelcluster-platform') do
+      nvidia_imex 'configure' do
+        action :configure
+      end
+    end
+  end
 end
 
 describe 'nvidia_imex:nvidia_enabled_or_installed?' do
@@ -245,3 +253,72 @@ def self.install(chef_run)
     end
   end
 end
+
+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
+        cached(:chef_run) do
+          stubs_for_resource('nvidia_imex') do |res|
+            allow(res).to receive(:imex_installed).and_return(false)
+          end
+          runner = runner(platform: platform, version: version, step_into: ['nvidia_imex'])
+          ConvergeNvidiaImex.configure(runner)
+        end
+        cached(:node) { chef_run.node }
+
+        it 'does not configure nvidia-imex' do
+          is_expected.not_to configure_nvidia_imex('nvidia-imex')
+        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)
+          end
+          runner = runner(platform: platform, version: version, step_into: ['nvidia_imex'])
+          ConvergeNvidiaImex.configure(runner)
+        end
+        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 })
+          end
+        else
+          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
+
+      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(1)
+          end
+          runner = runner(platform: platform, version: version, step_into: ['nvidia_imex'])
+          ConvergeNvidiaImex.configure(runner)
+        end
+        cached(:node) { chef_run.node }
+
+        before do
+          chef_run.node.override['cluster']['region'] = 'aws_region'
+        end
+
+        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
+      end
+    end
+  end
+end

cookbooks/aws-parallelcluster-platform/test/controls/nvidia_imex_spec.rb

@@ -12,12 +12,14 @@
 control 'tag:install_expected_versions_of_nvidia_imex_installed' do
   only_if { ['yes', true, 'true'].include?(node['cluster']['nvidia']['enabled']) }
 
-  describe package('nvidia-imex') do
+  nvidia_imex_service = 'nvidia-imex'
+
+  describe package(nvidia_imex_service) do
     it { should be_installed }
     its('version') { should match /#{node['cluster']['nvidia']['imex']['version']}/ }
   end
 
-  %w(/usr/bin/nvidia-imex /usr/bin/nvidia-imex-ctl).each do |path|
+  ["/usr/bin/#{nvidia_imex_service}", "/usr/bin/#{nvidia_imex_service}-ctl"].each do |path|
     describe file(path) do
       it { should exist }
       its('owner') { should eq 'root' }
@@ -26,7 +28,7 @@
     end
   end
 
-  nvidia_imex_dir = "#{node['cluster']['shared_dir']}/nvidia-imex"
+  nvidia_imex_dir = "#{node['cluster']['shared_dir']}/#{nvidia_imex_service}"
 
   ["#{nvidia_imex_dir}/config.cfg", "#{nvidia_imex_dir}/nodes_config.cfg"].each do |conf_files|
     describe file(conf_files) do