3.0 Make slurm work when using GPU instances with no NVIDIA drivers installed

* if nvidia drive is not installed on AMI of headnode, don't configure gres slurm gpu configuration

Signed-off-by: Yulei Wang <[email protected]>

Author: yuleiwan

files/default/slurm/pcluster_slurm_config_generator.py

@@ -30,7 +30,9 @@ class CriticalError(Exception):
     pass
 
 
-def generate_slurm_config_files(output_directory, template_directory, input_file, instance_types_data_path, dryrun):
+def generate_slurm_config_files(
+    output_directory, template_directory, input_file, instance_types_data_path, dryrun, no_gpu
+):
     """
     Generate Slurm configuration files.
 
@@ -61,6 +63,8 @@ def generate_slurm_config_files(output_directory, template_directory, input_file
     is_default_queue = True  # The first queue in the queues list is the default queue
     for queue in queues:
         for file_type in ["partition", "gres"]:
+            if file_type == "gres" and no_gpu:
+                continue
             _generate_queue_config(
                 queue["Name"], queue, is_default_queue, file_type, env, pcluster_subdirectory, dryrun
             )
@@ -272,9 +276,21 @@ def main():
             required=False,
             default=False,
         )
+        parser.add_argument(
+            "--no-gpu",
+            action="store_true",
+            help="no gpu configuration",
+            required=False,
+            default=False,
+        )
         args = parser.parse_args()
         generate_slurm_config_files(
-            args.output_directory, args.template_directory, args.input_file, args.instance_types_data, args.dryrun
+            args.output_directory,
+            args.template_directory,
+            args.input_file,
+            args.instance_types_data,
+            args.dryrun,
+            args.no_gpu,
         )
     except Exception as e:
         log.exception("Failed to generate slurm configurations, exception: %s", e)

libraries/helpers.rb

@@ -131,6 +131,18 @@ def graphic_instance?
   !has_gpu.stdout.strip.empty?
 end
 
+#
+# Check if the nvidia drive is installed
+#
+def nvidia_installed?
+  nvidia_installed = Mixlib::ShellOut.new("which nvidia-smi")
+  nvidia_installed.run_command
+
+  Chef::Log.info("Nvidia drive is not installed") if nvidia_installed.stdout.strip.empty?
+
+  !nvidia_installed.stdout.strip.empty?
+end
+
 #
 # Check if the AMI is bootstrapped
 #

recipes/compute_slurm_config.rb

@@ -35,10 +35,14 @@
   retry_delay 6
 end
 
-# Check to see if there is GPU on the instance, only execute run_nvidiasmi if there is GPU
+# Check to see if there is GPU on the instance, only execute run_nvidiasmi if there is GPU and nvidia installed
 if graphic_instance?
-  execute "run_nvidiasmi" do
-    command 'nvidia-smi'
+  if nvidia_installed?
+    execute "run_nvidiasmi" do
+      command 'nvidia-smi'
+    end
+  else
+    Chef::Log.warn("GPU instance but no Nvidia drivers found")
   end
 end
 

recipes/head_node_slurm_config.rb

@@ -92,9 +92,15 @@
 
 # Generate pcluster specific configs
 execute "generate_pcluster_slurm_configs" do
-  command "#{node['cluster']['cookbook_virtualenv_path']}/bin/python #{node['cluster']['scripts_dir']}/slurm/pcluster_slurm_config_generator.py"\
-          " --output-directory /opt/slurm/etc/ --template-directory #{node['cluster']['scripts_dir']}/slurm/templates/"\
-          " --input-file #{node['cluster']['cluster_config_path']}  --instance-types-data #{node['cluster']['instance_types_data_path']}"
+  if nvidia_installed?
+    command "#{node['cluster']['cookbook_virtualenv_path']}/bin/python #{node['cluster']['scripts_dir']}/slurm/pcluster_slurm_config_generator.py"\
+            " --output-directory /opt/slurm/etc/ --template-directory #{node['cluster']['scripts_dir']}/slurm/templates/"\
+            " --input-file #{node['cluster']['cluster_config_path']}  --instance-types-data #{node['cluster']['instance_types_data_path']}"
+  else
+    command "#{node['cluster']['cookbook_virtualenv_path']}/bin/python #{node['cluster']['scripts_dir']}/slurm/pcluster_slurm_config_generator.py"\
+            " --output-directory /opt/slurm/etc/ --template-directory #{node['cluster']['scripts_dir']}/slurm/templates/"\
+            " --input-file #{node['cluster']['cluster_config_path']}  --instance-types-data #{node['cluster']['instance_types_data_path']} --no-gpu"
+  end
 end
 
 # all other OSs use /sys/fs/cgroup, which is the default

recipes/update_head_node_slurm.rb

@@ -35,9 +35,15 @@
   end
   # Generate pcluster specific configs
   execute "generate_pcluster_slurm_configs" do
-    command "#{node['cluster']['cookbook_virtualenv_path']}/bin/python #{node['cluster']['scripts_dir']}/slurm/pcluster_slurm_config_generator.py" \
-            " --output-directory /opt/slurm/etc/ --template-directory #{node['cluster']['scripts_dir']}/slurm/templates/"\
-            " --input-file #{updated_cluster_config_path} --instance-types-data #{node['cluster']['instance_types_data_path']}"
+    if nvidia_installed?
+      command "#{node['cluster']['cookbook_virtualenv_path']}/bin/python #{node['cluster']['scripts_dir']}/slurm/pcluster_slurm_config_generator.py" \
+              " --output-directory /opt/slurm/etc/ --template-directory #{node['cluster']['scripts_dir']}/slurm/templates/"\
+              " --input-file #{updated_cluster_config_path} --instance-types-data #{node['cluster']['instance_types_data_path']}"
+    else
+      command "#{node['cluster']['cookbook_virtualenv_path']}/bin/python #{node['cluster']['scripts_dir']}/slurm/pcluster_slurm_config_generator.py" \
+              " --output-directory /opt/slurm/etc/ --template-directory #{node['cluster']['scripts_dir']}/slurm/templates/"\
+              " --input-file #{updated_cluster_config_path} --instance-types-data #{node['cluster']['instance_types_data_path']} --no-gpu"
+    end
   end
 
   execute 'stop clustermgtd' do

test/unit/slurm/test_slurm_config_generator.py

@@ -14,7 +14,7 @@ def test_generate_slurm_config_files(mocker, test_datadir, tmpdir):
         "slurm.pcluster_slurm_config_generator._get_head_node_private_ip", return_value="ip.1.0.0.0", autospec=True
     )
     template_directory = os.path.dirname(slurm.__file__) + "/templates"
-    generate_slurm_config_files(tmpdir, template_directory, input_file, instance_types_data, dryrun=False)
+    generate_slurm_config_files(tmpdir, template_directory, input_file, instance_types_data, dryrun=False, no_gpu=False)
 
     for queue in ["efa", "gpu", "multiple_spot"]:
         for file_type in ["partition", "gres"]: