From 06012f9b4f0091cdcd4e949bd714c9342dfe8c0b Mon Sep 17 00:00:00 2001
From: Himani Anil Deshpande <himanidp@amazon.com>
Date: Tue, 20 May 2025 13:51:10 -0400
Subject: [PATCH] [Bug] Installing Pyxis if nvidia is installed or enabled

---
 .../recipes/install/install_pyxis.rb          |  2 +-
 .../spec/unit/recipes/install_pyxis_spec.rb   |  2 +
 .../test/controls/pyxis_spec.rb               | 43 +++++++++----------
 3 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/cookbooks/aws-parallelcluster-slurm/recipes/install/install_pyxis.rb b/cookbooks/aws-parallelcluster-slurm/recipes/install/install_pyxis.rb
index cfc6cad809..eabfc8afc3 100644
--- a/cookbooks/aws-parallelcluster-slurm/recipes/install/install_pyxis.rb
+++ b/cookbooks/aws-parallelcluster-slurm/recipes/install/install_pyxis.rb
@@ -15,7 +15,7 @@
 # OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
 # limitations under the License.
 
-return unless nvidia_installed?
+return unless nvidia_enabled? || nvidia_installed?
 return if pyxis_installed?
 
 pyxis_version = node['cluster']['pyxis']['version']
diff --git a/cookbooks/aws-parallelcluster-slurm/spec/unit/recipes/install_pyxis_spec.rb b/cookbooks/aws-parallelcluster-slurm/spec/unit/recipes/install_pyxis_spec.rb
index 04681fdd0e..0118e6edc3 100644
--- a/cookbooks/aws-parallelcluster-slurm/spec/unit/recipes/install_pyxis_spec.rb
+++ b/cookbooks/aws-parallelcluster-slurm/spec/unit/recipes/install_pyxis_spec.rb
@@ -33,6 +33,7 @@
           node.override['cluster']['pyxis']['version'] = pyxis_version
           node.override['cluster']['pyxis']['runtime_path'] = pyxis_runtime_dir
         end
+        allow_any_instance_of(Object).to receive(:nvidia_enabled?).and_return(true)
         allow_any_instance_of(Object).to receive(:nvidia_installed?).and_return(true)
         allow_any_instance_of(Object).to receive(:pyxis_installed?).and_return(false)
         runner.converge(described_recipe)
@@ -93,6 +94,7 @@
           runner = runner(platform: platform, version: version) do |_node|
             RSpec::Mocks.configuration.allow_message_expectations_on_nil = true
           end
+          allow_any_instance_of(Object).to receive(:nvidia_enabled?).and_return(true)
           allow_any_instance_of(Object).to receive(:nvidia_installed?).and_return(true)
           allow_any_instance_of(Object).to receive(:pyxis_installed?).and_return(true)
           runner.converge(described_recipe)
diff --git a/cookbooks/aws-parallelcluster-slurm/test/controls/pyxis_spec.rb b/cookbooks/aws-parallelcluster-slurm/test/controls/pyxis_spec.rb
index 77708aa770..43353f267c 100644
--- a/cookbooks/aws-parallelcluster-slurm/test/controls/pyxis_spec.rb
+++ b/cookbooks/aws-parallelcluster-slurm/test/controls/pyxis_spec.rb
@@ -9,25 +9,24 @@
 # 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.
 
-# FIXME: Re-enabled the following check and fix failures
-# control 'tag:install_pyxis_installed' do
-#   only_if { instance.nvidia_installed? }
-#
-#   title 'Checks Pyxis has been installed'
-#
-#   examples_dir = "/opt/parallelcluster/examples"
-#   dirs = [ examples_dir, "#{examples_dir}/spank", "#{examples_dir}/pyxis" ]
-#   dirs.each do |path|
-#     describe directory(path) do
-#       it { should exist }
-#     end
-#   end
-#
-#   describe file("#{examples_dir}/pyxis/pyxis.conf") do
-#     it { should exist }
-#   end
-#
-#   describe file("#{examples_dir}/spank/plugstack.conf") do
-#     it { should exist }
-#   end
-# end
+control 'tag:install_pyxis_installed' do
+  only_if { ['yes', true, 'true'].include?(node['cluster']['nvidia']['enabled']) || instance.nvidia_installed? }
+
+  title 'Checks Pyxis has been installed'
+
+  examples_dir = "/opt/parallelcluster/examples"
+  dirs = [ examples_dir, "#{examples_dir}/spank", "#{examples_dir}/pyxis" ]
+  dirs.each do |path|
+    describe directory(path) do
+      it { should exist }
+    end
+  end
+
+  describe file("#{examples_dir}/pyxis/pyxis.conf") do
+    it { should exist }
+  end
+
+  describe file("#{examples_dir}/spank/plugstack.conf") do
+    it { should exist }
+  end
+end