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

Author: Himani Anil Deshpande

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

@@ -17,6 +17,7 @@
 end
 
 use 'partial/_nvidia_imex_common.rb'
+use 'partial/_nvidia_imex_rhel.rb'
 
 def _nvidia_imex_version
   "#{nvidia_driver_major_version}-#{nvidia_imex_full_version}"

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

@@ -15,6 +15,7 @@
 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

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

@@ -17,6 +17,7 @@
 end
 
 use 'partial/_nvidia_imex_common.rb'
+use 'partial/_nvidia_imex_rhel.rb'
 
 def _nvidia_imex_version
   "#{nvidia_driver_major_version}-#{nvidia_imex_full_version}"

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

@@ -17,6 +17,7 @@
 end
 
 use 'partial/_nvidia_imex_common.rb'
+use 'partial/_nvidia_imex_rhel.rb'
 
 def _nvidia_imex_version
   "#{nvidia_driver_major_version}-#{nvidia_imex_full_version}"

cookbooks/aws-parallelcluster-platform/resources/nvidia_imex/nvidia_imex_ubuntu22+.rb

@@ -17,10 +17,8 @@
 end
 
 use 'partial/_nvidia_imex_common.rb'
+use 'partial/_nvidia_imex_debian.rb'
 
 def _nvidia_imex_version
-  "1:#{nvidia_driver_major_version}_#{nvidia_imex_full_version}"
-  # The single digit "1" is epoch version. Without the "1", package install fails because version does not exist.
-  # See details here:
-  # https://askubuntu.com/questions/441879/why-do-some-packages-have-extra-numbers-before-a-colon-on-the-front-of-their-ver
+  "#{nvidia_driver_major_version}_#{nvidia_imex_full_version}"
 end

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

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 #
-# Copyright:: 2013-2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# 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.
@@ -48,14 +48,15 @@
     action :create
   end
 
-  install_packages 'Install nvidia-imex' do
-    packages "#{nvidia_imex_service}-#{_nvidia_imex_version}"
-    action :install
-  end
+  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_service}-#{nvidia_driver_major_version}"
   node_attributes 'dump node attributes'
+
+  nvidia_repo 'remove nvidia repository' do
+    action :remove
+  end
 end
 
 action :configure do

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

@@ -0,0 +1,22 @@
+# 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
+  apt_package "Install nvidia-imex" do
+    package_name "#{nvidia_imex_service}-#{nvidia_driver_major_version}"
+    version nvidia_imex_full_version
+    retries 10
+    retry_delay 5
+  end
+end

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

@@ -0,0 +1,20 @@
+# 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
+  install_packages 'Install nvidia-imex' do
+    packages "#{nvidia_imex_service}-#{_nvidia_imex_version}"
+    action :install
+  end
+end

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

@@ -178,11 +178,12 @@ def self.configure(chef_run)
           runner(platform: platform, version: version, step_into: ['nvidia_imex'])
         end
         cached(:nvidia_imex_version) { "1.2.3-1" }
+        cached(:nvidia_imex_package) { "nvidia-imex-1" }
         cached(:nvidia_imex_name) do
           if %(redhat rocky).include?(platform) || platform == 'amazon' && version == '2023'
-            "nvidia-imex-1-#{nvidia_imex_version}"
+            "#{nvidia_imex_package}-#{nvidia_imex_version}"
           else
-            "nvidia-imex-1_#{nvidia_imex_version}"
+            "#{nvidia_imex_package}_#{nvidia_imex_version}"
           end
         end
         cached(:node) { chef_run.node }
@@ -219,7 +220,9 @@ def self.configure(chef_run)
           end
           it 'does not set nvidia-imex version' do
             expect(node.default['cluster']['nvidia']['imex']['version']).not_to eq(nvidia_imex_version)
+            expect(node.default['cluster']['nvidia']['imex']['package']).not_to eq(nvidia_imex_package)
             is_expected.not_to write_node_attributes('dump node attributes')
+            is_expected.not_to remove_nvidia_repo('remove nvidia repository')
           end
         else
           it 'installs nvidia-imex' do
@@ -241,13 +244,23 @@ def self.configure(chef_run)
               .with(user: 'root')
               .with(group: 'root')
               .with(mode: '0644')
-            is_expected.to install_install_packages('Install nvidia-imex')
-              .with(packages: "#{nvidia_imex_name}")
-              .with(action: %i(install))
+            if platform == 'ubuntu'
+              is_expected.to install_apt_package('Install nvidia-imex')
+                .with(package_name: nvidia_imex_package)
+                .with(version: nvidia_imex_version)
+                .with(retries: 10)
+                .with(retry_delay: 5)
+            else
+              is_expected.to install_install_packages('Install nvidia-imex')
+                .with(packages: nvidia_imex_name)
+                .with(action: %i(install))
+            end
           end
           it 'sets nvidia-imex version' do
             expect(node.default['cluster']['nvidia']['imex']['version']).to eq(nvidia_imex_version)
+            expect(node.default['cluster']['nvidia']['imex']['package']).to eq(nvidia_imex_package)
             is_expected.to write_node_attributes('dump node attributes')
+            is_expected.to remove_nvidia_repo('remove nvidia repository')
           end
         end
       end