change gres conf and partition conf with no gpu

Signed-off-by: chenwany <[email protected]>

Author: chenwany

CHANGELOG.md

@@ -3,6 +3,12 @@ aws-parallelcluster-cookbook CHANGELOG
 
 This file is used to list changes made in each version of the AWS ParallelCluster cookbook.
 
+3.x.x
+------
+
+**CHANGES**
+- Do not configure GPUs in Slurm when Nvidia driver is not installed.
+
 3.0.0
 ------
 

files/default/slurm/pcluster_slurm_config_generator.py

@@ -63,10 +63,8 @@ def generate_slurm_config_files(
     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
+                queue["Name"], queue, is_default_queue, file_type, env, pcluster_subdirectory, dryrun, no_gpu=no_gpu
             )
         is_default_queue = False
 
@@ -109,14 +107,23 @@ def _get_head_node_private_ip():
     return _get_metadata("local-ipv4")
 
 
-def _generate_queue_config(queue_name, queue_config, is_default_queue, file_type, jinja_env, output_dir, dryrun):
+def _generate_queue_config(
+    queue_name, queue_config, is_default_queue, file_type, jinja_env, output_dir, dryrun, no_gpu=False
+):
     log.info("Generating slurm_parallelcluster_%s_%s.conf", queue_name, file_type)
     rendered_template = jinja_env.get_template(f"slurm_parallelcluster_queue_{file_type}.conf").render(
-        queue_name=queue_name, queue_config=queue_config, is_default_queue=is_default_queue
+        queue_name=queue_name, queue_config=queue_config, is_default_queue=is_default_queue, no_gpu=no_gpu
     )
     if not dryrun:
         filename = path.join(output_dir, f"slurm_parallelcluster_{queue_name}_{file_type}.conf")
-        _write_rendered_template_to_file(rendered_template, filename)
+        if file_type == "gres" and no_gpu:
+            _write_rendered_template_to_file(
+                "# This file is automatically generated by pcluster\n"
+                "# Skipping GPUs configuration because Nvidia driver is not installed",
+                filename,
+            )
+        else:
+            _write_rendered_template_to_file(rendered_template, filename)
 
 
 def _generate_slurm_parallelcluster_configs(

files/default/slurm/templates/slurm_parallelcluster_queue_partition.conf

@@ -7,11 +7,11 @@
     {% set instance_type = compute_resource.InstanceType%}
     {% set sanitized_compute_name = compute_resource.Name | sanify_name %}
     {% if static_size > 0 %}
-NodeName={{ queue_name }}-st-{{ sanitized_compute_name }}-[1-{{ static_size }}] CPUs={{ compute_resource | vcpus }} State=CLOUD Feature=static,{{ instance_type }},{{ compute_resource.Name }}{% if compute_resource.Efa.Enabled %},efa{% endif %}{% if instance_type | gpus > 0 %},gpu Gres=gpu:{{ instance_type | gpu_type }}:{{ instance_type | gpus }}{% endif %}
+NodeName={{ queue_name }}-st-{{ sanitized_compute_name }}-[1-{{ static_size }}] CPUs={{ compute_resource | vcpus }} State=CLOUD Feature=static,{{ instance_type }},{{ compute_resource.Name }}{% if compute_resource.Efa.Enabled %},efa{% endif %}{% if instance_type | gpus > 0 %},gpu{% if not no_gpu %} Gres=gpu:{{ instance_type | gpu_type }}:{{ instance_type | gpus }}{% endif %}{% endif %}
 
     {% endif %}
     {% if dynamic_size > 0 %}
-NodeName={{ queue_name }}-dy-{{ sanitized_compute_name }}-[1-{{ dynamic_size }}] CPUs={{ compute_resource | vcpus }} State=CLOUD Feature=dynamic,{{ instance_type }},{{ compute_resource.Name }}{% if compute_resource.Efa.Enabled %},efa{% endif %}{% if instance_type | gpus > 0 %},gpu Gres=gpu:{{ instance_type | gpu_type }}:{{ instance_type | gpus }}{% endif %}
+NodeName={{ queue_name }}-dy-{{ sanitized_compute_name }}-[1-{{ dynamic_size }}] CPUs={{ compute_resource | vcpus }} State=CLOUD Feature=dynamic,{{ instance_type }},{{ compute_resource.Name }}{% if compute_resource.Efa.Enabled %},efa{% endif %}{% if instance_type | gpus > 0 %},gpu{% if not no_gpu %} Gres=gpu:{{ instance_type | gpu_type }}:{{ instance_type | gpus }}{% endif %}{% endif %}
 
     {% endif %}
     {% if static_size > 0 or dynamic_size > 0 %}

libraries/helpers.rb

@@ -132,15 +132,12 @@ def graphic_instance?
 end
 
 #
-# Check if the nvidia drive is installed
+# Check if Nvidia driver 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?
+  nvidia_installed = ::File.exist?('/usr/bin/nvidia-smi')
+  Chef::Log.warn("Nvidia driver is not installed") unless nvidia_installed
+  nvidia_installed
 end
 
 #

recipes/compute_slurm_config.rb

@@ -35,16 +35,8 @@
   retry_delay 6
 end
 
-# Check to see if there is GPU on the instance, only execute run_nvidiasmi if there is GPU and nvidia installed
-if graphic_instance?
-  if nvidia_installed?
-    execute "run_nvidiasmi" do
-      command 'nvidia-smi'
-    end
-  else
-    Chef::Log.warn("GPU instance but no Nvidia drivers found")
-  end
-end
+# Check to see if is GPU instance with Nvidia installed
+Chef::Log.warn("GPU instance but no Nvidia drivers found") if graphic_instance? && !nvidia_installed?
 
 cookbook_file '/etc/systemd/system/slurmd.service' do
   source 'slurmd.service'

recipes/head_node_slurm_config.rb

@@ -91,16 +91,11 @@
 end
 
 # Generate pcluster specific configs
+no_gpu = nvidia_installed? ? "" : "--no-gpu"
 execute "generate_pcluster_slurm_configs" do
-  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
+  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
 
 # all other OSs use /sys/fs/cgroup, which is the default

recipes/update_head_node_slurm.rb

@@ -34,16 +34,11 @@
     retry_delay 5
   end
   # Generate pcluster specific configs
+  no_gpu = nvidia_installed? ? "" : "--no-gpu"
   execute "generate_pcluster_slurm_configs" do
-    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
+    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
 
   execute 'stop clustermgtd' do

test/unit/slurm/test_slurm_config_generator.py

@@ -1,11 +1,16 @@
 import os
 
+import pytest
 import slurm
 from assertpy import assert_that
 from slurm.pcluster_slurm_config_generator import generate_slurm_config_files
 
 
-def test_generate_slurm_config_files(mocker, test_datadir, tmpdir):
+@pytest.mark.parametrize(
+    "no_gpu",
+    [False, True],
+)
+def test_generate_slurm_config_files(mocker, test_datadir, tmpdir, no_gpu):
     input_file = str(test_datadir / "sample_input.yaml")
     instance_types_data = str(test_datadir / "sample_instance_types_data.json")
 
@@ -14,12 +19,16 @@ 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, no_gpu=False)
+    generate_slurm_config_files(
+        tmpdir, template_directory, input_file, instance_types_data, dryrun=False, no_gpu=no_gpu
+    )
 
     for queue in ["efa", "gpu", "multiple_spot"]:
         for file_type in ["partition", "gres"]:
             file_name = f"pcluster/slurm_parallelcluster_{queue}_{file_type}.conf"
-            _assert_files_are_equal(tmpdir / file_name, test_datadir / "expected_outputs" / file_name)
+            no_nvidia = "_no_gpu" if (queue == "gpu" or file_type == "gres") and no_gpu else ""
+            output_file_name = f"pcluster/slurm_parallelcluster_{queue}_{file_type}{no_nvidia}.conf"
+            _assert_files_are_equal(tmpdir / file_name, test_datadir / "expected_outputs" / output_file_name)
 
     for file in ["slurm_parallelcluster.conf", "slurm_parallelcluster_gres.conf"]:
         _assert_files_are_equal(tmpdir / file, test_datadir / "expected_outputs" / file)

test/unit/slurm/test_slurm_config_generator/test_generate_slurm_config_files/expected_outputs/pcluster/slurm_parallelcluster_efa_gres_no_gpu.conf

@@ -0,0 +1,2 @@
+# This file is automatically generated by pcluster
+# Skipping GPUs configuration because Nvidia driver is not installed

test/unit/slurm/test_slurm_config_generator/test_generate_slurm_config_files/expected_outputs/pcluster/slurm_parallelcluster_gpu_gres_no_gpu.conf

@@ -0,0 +1,2 @@
+# This file is automatically generated by pcluster
+# Skipping GPUs configuration because Nvidia driver is not installed