Add support for Slurm Accounting

Add support for Slurm Accounting at cookbook level
- Add Slurm configuration parameters to set up accounting (in slurm.conf and slurmdbd.conf)
- Define new user to run the slurmdbd daemon and assign it the slurmdbd.conf files
- Create slurmdbd.service unit file
- Create script to update database password from Secrets Manager
- Enable slurmdbd and activate accounting in the slurm configuration recipe
- Add action to wait for Slurm database to become responsive
- Bootstrap slurm database with cluster, default account and users in an idempotent way

Allow update-cluster to enable and disable Slurm Accounting
- Call the Slurm accounting configuration recipe if Slurm Accounting is enabled on an existing cluster
- Call a clean-up recipe to clear some of the Slurm Accounting configuration if it is disabled

Add unit tests for Slurm config file generation for Slurm Accounting
- Adapt existing tests to account for additional input needed by Slurm Accounting

Update Changelog

Signed-off-by: Jacopo De Amicis <[email protected]>

Author: jdeamicis

CHANGELOG.md

@@ -8,6 +8,7 @@ This file is used to list changes made in each version of the AWS ParallelCluste
 
 **ENHANCEMENTS**
 - Add support for AWS Trainium instances.
+- Add support for Slurm Accounting.
 - Upgrade third-party cookbook dependencies:
   - selinux-6.0.5 (from selinux-6.0.4)
   - nfs-5.0.0 (from nfs-2.6.4)

attributes/default.rb

@@ -134,6 +134,10 @@
 default['cluster']['slurm']['user_id'] = node['cluster']['reserved_base_uid'] + 1
 default['cluster']['slurm']['group'] = node['cluster']['slurm']['user']
 default['cluster']['slurm']['group_id'] = node['cluster']['slurm']['user_id']
+default['cluster']['slurm']['dbduser'] = 'slurmdbd'
+default['cluster']['slurm']['dbduser_id'] = node['cluster']['reserved_base_uid'] + 5
+default['cluster']['slurm']['dbdgroup'] = node['cluster']['slurm']['dbduser']
+default['cluster']['slurm']['dbdgroup_id'] = node['cluster']['slurm']['dbduser_id']
 default['cluster']['slurm']['install_dir'] = "/opt/slurm"
 default['cluster']['slurm']['fleet_config_path'] = "#{node['cluster']['slurm_plugin_dir']}/fleet-config.json"
 

cookbooks/aws-parallelcluster-slurm/files/default/head_node_slurm/slurm/pcluster_slurm_config_generator.py

@@ -18,6 +18,7 @@
 from os import makedirs, path
 from socket import gethostname
 from typing import Tuple
+from urllib.parse import ParseResult, urlparse
 
 import requests
 import yaml
@@ -43,6 +44,7 @@ def generate_slurm_config_files(
     no_gpu,
     compute_node_bootstrap_timeout,
     realmemory_to_ec2memory_ratio,
+    slurmdbd_user,
 ):
     """
     Generate Slurm configuration files.
@@ -65,6 +67,7 @@ def generate_slurm_config_files(
     cluster_config = _load_cluster_config(input_file)
     head_node_config = _get_head_node_config()
     queues = cluster_config["Scheduling"]["SlurmQueues"]
+    cluster_name = next(tag["Value"] for tag in cluster_config["Tags"] if tag["Key"] == "parallelcluster:cluster-name")
 
     global instance_types_data
     with open(instance_types_data_path) as input_file:
@@ -86,16 +89,19 @@ def generate_slurm_config_files(
             )
         is_default_queue = False
 
-    # Generate slurm_parallelcluster.conf, slurm_parallelcluster_gres.conf and slurm_parallelcluster_cgroup.conf
+    # Generate include files for slurm configuration files
     for template_name in [
         "slurm_parallelcluster.conf",
         "slurm_parallelcluster_gres.conf",
         "slurm_parallelcluster_cgroup.conf",
+        "slurm_parallelcluster_slurmdbd.conf",
     ]:
         _generate_slurm_parallelcluster_configs(
             queues,
             head_node_config,
             cluster_config["Scheduling"]["SlurmSettings"],
+            cluster_name,
+            slurmdbd_user,
             template_name,
             compute_node_bootstrap_timeout,
             env,
@@ -161,6 +167,8 @@ def _generate_slurm_parallelcluster_configs(
     queues,
     head_node_config,
     scaling_config,
+    cluster_name,
+    slurmdbd_user,
     template_name,
     compute_node_bootstrap_timeout,
     jinja_env,
@@ -172,6 +180,8 @@ def _generate_slurm_parallelcluster_configs(
         queues=queues,
         head_node_config=head_node_config,
         scaling_config=scaling_config,
+        cluster_name=cluster_name,
+        slurmdbd_user=slurmdbd_user,
         compute_node_bootstrap_timeout=compute_node_bootstrap_timeout,
         output_dir=output_dir,
     )
@@ -195,6 +205,8 @@ def _get_jinja_env(template_directory, realmemory_to_ec2memory_ratio):
         _realmemory,
         realmemory_to_ec2memory_ratio=realmemory_to_ec2memory_ratio,
     )
+    env.filters["uri_host"] = functools.partial(_parse_uri, attr="host")
+    env.filters["uri_port"] = functools.partial(_parse_uri, attr="port")
 
     return env
 
@@ -307,6 +319,39 @@ def _realmemory(compute_resource, realmemory_to_ec2memory_ratio) -> int:
     return realmemory
 
 
+def _parse_netloc(uri: str, uri_parse: ParseResult, attr: str) -> str:
+    try:
+        netloc = uri_parse.netloc
+    except ValueError as e:
+        error_msg = f"Failure to parse uri with error '{str(e)}'. Please review the provided URI ('{uri}')"
+        log.critical(error_msg)
+        raise CriticalError(error_msg)
+    if not netloc:
+        error_msg = f"Invalid URI specified. Please review the provided URI ('{uri}')"
+        log.critical(error_msg)
+        raise CriticalError(error_msg)
+    if attr == "host":
+        ret = uri_parse.hostname
+    elif attr == "port":
+        ret = uri_parse.port
+        # Provide default MySQL port if port is not explicitely set
+        if not ret:
+            ret = "3306"
+    return ret
+
+
+def _parse_uri(uri, attr) -> str:
+    """Get a host from a URI/URL using urlparse."""
+    uri_parse = urlparse(uri)
+    if not uri_parse.netloc:
+        # This happens if users provide an URI without explicit scheme followed by ://
+        # (for example 'test.example.com:3306' instead of 'mysql://test.example.com:3306`).
+        uri_parse = urlparse("//" + uri)
+
+    # Parse netloc to get hostname or port
+    return _parse_netloc(uri, uri_parse, attr)
+
+
 def _write_rendered_template_to_file(rendered_template, filename):
     log.info("Writing contents of %s", filename)
     with open(filename, "w") as output_file:
@@ -399,6 +444,7 @@ def memory_ratio_float(arg):
             help="Configure ratio between RealMemory and memory advertised by EC2",
             required=True,
         )
+        parser.add_argument("--slurmdbd-user", help="User for the slurmdbd service.", required=True)
         args = parser.parse_args()
         generate_slurm_config_files(
             args.output_directory,
@@ -409,6 +455,7 @@ def memory_ratio_float(arg):
             args.no_gpu,
             args.compute_node_bootstrap_timeout,
             args.realmemory_to_ec2memory_ratio,
+            args.slurmdbd_user,
         )
     except Exception as e:
         log.exception("Failed to generate slurm configurations, exception: %s", e)

cookbooks/aws-parallelcluster-slurm/files/default/head_node_slurm/slurm/templates/slurm_parallelcluster.conf

@@ -11,6 +11,13 @@ SelectTypeParameters=CR_CPU_Memory
 {% else %}
 SelectTypeParameters=CR_CPU
 {% endif %}
+{% if scaling_config.Database.Uri is defined %}
+AccountingStorageType=accounting_storage/slurmdbd
+AccountingStorageHost={{ head_node_config.head_node_hostname }}
+AccountingStoragePort=6819
+AccountingStorageUser={{ slurmdbd_user }}
+JobAcctGatherType=jobacct_gather/cgroup
+{% endif %}
 
 {% for queue in queues %}
 include {{ output_dir }}/pcluster/slurm_parallelcluster_{{ queue.Name }}_partition.conf

cookbooks/aws-parallelcluster-slurm/files/default/head_node_slurm/slurm/templates/slurm_parallelcluster_slurmdbd.conf

@@ -0,0 +1,12 @@
+# slurm_parallelcluster_slurmdbd.conf is managed by the pcluster processes.
+# Do not modify.
+# Please add user-specific slurmdbd configuration options in slurmdbd.conf
+{% if scaling_config.Database.Uri is defined %}
+DbdHost={{ head_node_config.head_node_hostname }}
+StorageHost={{ scaling_config.Database.Uri | uri_host }}
+StoragePort={{ scaling_config.Database.Uri | uri_port }}
+{# Dashes in StorageLoc cause issues with the database creation #}
+StorageLoc={{ cluster_name | replace("-", "_") }}
+StorageUser={{ scaling_config.Database.UserName }}
+StoragePass=dummy
+{% endif %}

cookbooks/aws-parallelcluster-slurm/recipes/clear_slurm_accounting.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+#
+# Cookbook:: aws-parallelcluster-slurm
+# Recipe:: config_head_node
+#
+# Copyright:: 2013-2022 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.
+
+file "#{node['cluster']['scripts_dir']}/slurm/update_slurm_database_password.sh" do
+  action :delete
+end
+
+service "slurmdbd" do
+  supports restart: false
+  action %i(disable stop)
+end

cookbooks/aws-parallelcluster-slurm/recipes/config_head_node.rb

@@ -68,7 +68,8 @@
             " --output-directory #{node['cluster']['slurm']['install_dir']}/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']}"\
             " --compute-node-bootstrap-timeout #{node['cluster']['compute_node_bootstrap_timeout']} #{no_gpu}"\
-            " --realmemory-to-ec2memory-ratio #{node['cluster']['realmemory_to_ec2memory_ratio']}"
+            " --realmemory-to-ec2memory-ratio #{node['cluster']['realmemory_to_ec2memory_ratio']}"\
+            " --slurmdbd-user #{node['cluster']['slurm']['dbduser']}"
   end
 
   # Generate pcluster fleet config
@@ -192,6 +193,21 @@
   action :create
 end
 
+template '/etc/systemd/system/slurmdbd.service' do
+  source 'slurm/head_node/slurmdbd.service.erb'
+  owner 'root'
+  group 'root'
+  mode '0644'
+  action :create
+end
+
+ruby_block "Configure Slurm Accounting" do
+  block do
+    run_context.include_recipe "aws-parallelcluster-slurm::config_slurm_accounting"
+  end
+  not_if { node['cluster']['config'].dig(:Scheduling, :SlurmSettings, :Database).nil? }
+end
+
 service "slurmctld" do
   supports restart: false
   action %i(enable start)

cookbooks/aws-parallelcluster-slurm/recipes/config_slurm_accounting.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+#
+# Cookbook:: aws-parallelcluster-slurm
+# Recipe:: config_head_node
+#
+# Copyright:: 2013-2022 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.
+
+template "#{node['cluster']['slurm']['install_dir']}/etc/slurmdbd.conf" do
+  source 'slurm/slurmdbd.conf.erb'
+  owner "#{node['cluster']['slurm']['dbduser']}"
+  group "#{node['cluster']['slurm']['dbdgroup']}"
+  mode '0600'
+  # Do not overwrite possible user customization if the database credentials are updated
+  not_if { ::File.exist?("#{node['cluster']['slurm']['install_dir']}/etc/slurmdbd.conf") }
+end
+
+file "#{node['cluster']['slurm']['install_dir']}/etc/slurm_parallelcluster_slurmdbd.conf" do
+  owner "#{node['cluster']['slurm']['dbduser']}"
+  group "#{node['cluster']['slurm']['dbdgroup']}"
+  mode '0600'
+end
+
+template "#{node['cluster']['scripts_dir']}/slurm/update_slurm_database_password.sh" do
+  source 'slurm/head_node/update_slurm_database_password.sh.erb'
+  owner 'root'
+  group 'root'
+  mode '0700'
+  variables(
+    secret_arn: lazy { node['cluster']['config'].dig(:Scheduling, :SlurmSettings, :Database, :PasswordSecretArn) },
+    region: node['cluster']['region']
+  )
+  sensitive true
+end
+
+execute "update Slurm database password" do
+  user 'root'
+  group 'root'
+  command "#{node['cluster']['scripts_dir']}/slurm/update_slurm_database_password.sh"
+end
+
+service "slurmdbd" do
+  supports restart: false
+  action %i(enable start)
+end
+
+# After starting slurmdbd the database may not be fully responsive yet and
+# its bootstrapping may fail. We need to wait for sacctmgr to successfully
+# query the database before proceeding.
+execute "wait for slurm database" do
+  command "#{node['cluster']['slurm']['install_dir']}/bin/sacctmgr show clusters -Pn"
+  # Very large value to simulate infinite loop (we will hit some other timeout
+  # before this).
+  retries 100000
+  retry_delay 10
+end
+
+bash "bootstrap slurm database" do
+  user 'root'
+  group 'root'
+  code <<-BOOTSTRAP
+    SACCTMGR_CMD=#{node['cluster']['slurm']['install_dir']}/bin/sacctmgr
+    CLUSTER_NAME=#{node['cluster']['stack_name']}
+    DEF_ACCOUNT=pcdefault
+    SLURM_USER=#{node['cluster']['slurm']['user']}
+    DEF_USER=#{node['cluster']['cluster_user']}
+
+    # Add cluster to database if it is not present yet
+    [[ $($SACCTMGR_CMD show clusters -Pn cluster=$CLUSTER_NAME | grep $CLUSTER_NAME) ]] || \
+        $SACCTMGR_CMD -iQ add cluster $CLUSTER_NAME
+
+    # Add account-cluster association to database if it is not present yet
+    [[ $($SACCTMGR_CMD list associations -Pn cluster=$CLUSTER_NAME account=$DEF_ACCOUNT format=account | grep $DEF_ACCOUNT) ]] || \
+        $SACCTMGR_CMD -iQ add account $DEF_ACCOUNT Cluster=$CLUSTER_NAME \
+            Description="ParallelCluster default account" Organization="none"
+
+    # Add user-account associations to database if they are not present yet
+    [[ $($SACCTMGR_CMD list associations -Pn cluster=$CLUSTER_NAME account=$DEF_ACCOUNT user=$SLURM_USER format=user | grep $SLURM_USER) ]] || \
+        $SACCTMGR_CMD -iQ add user $SLURM_USER Account=$DEF_ACCOUNT AdminLevel=Admin
+    [[ $($SACCTMGR_CMD list associations -Pn cluster=$CLUSTER_NAME account=$DEF_ACCOUNT user=$DEF_USER format=user | grep $DEF_USER) ]] || \
+        $SACCTMGR_CMD -iQ add user $DEF_USER Account=$DEF_ACCOUNT AdminLevel=Admin
+
+    # sacctmgr might throw errors if the DEF_ACCOUNT is not associated to a cluster already defined on the database.
+    # This is not important for the scope of this script, so we return 0.
+    exit 0
+  BOOTSTRAP
+end

cookbooks/aws-parallelcluster-slurm/recipes/install_slurm.rb

@@ -50,6 +50,23 @@
   shell '/bin/bash'
 end
 
+# Setup slurmdbd group
+group node['cluster']['slurm']['dbdgroup'] do
+  comment 'slurm group'
+  gid node['cluster']['slurm']['dbdgroup_id']
+  system true
+end
+
+# Setup slurmdbd user
+user node['cluster']['slurm']['dbduser'] do
+  comment 'slurmdbd user'
+  uid node['cluster']['slurm']['dbduser_id']
+  gid node['cluster']['slurm']['dbdgroup_id']
+  manage_home false
+  system true
+  shell '/bin/bash'
+end
+
 include_recipe 'aws-parallelcluster-slurm::install_jwt'
 
 slurm_tarball = "#{node['cluster']['sources_dir']}/slurm-#{node['cluster']['slurm']['version']}.tar.gz"

cookbooks/aws-parallelcluster-slurm/recipes/update_head_node.rb

@@ -154,7 +154,8 @@ def update_nodes_in_queue(strategy, queues)
           " --instance-types-data #{node['cluster']['instance_types_data_path']}" \
           " --compute-node-bootstrap-timeout #{node['cluster']['compute_node_bootstrap_timeout']}" \
           " #{nvidia_installed? ? '' : '--no-gpu'}"\
-          " --realmemory-to-ec2memory-ratio #{node['cluster']['realmemory_to_ec2memory_ratio']}"
+          " --realmemory-to-ec2memory-ratio #{node['cluster']['realmemory_to_ec2memory_ratio']}"\
+          " --slurmdbd-user #{node['cluster']['slurm']['dbduser']}"
   not_if { ::File.exist?(node['cluster']['previous_cluster_config_path']) && !are_queues_updated? }
 end
 
@@ -171,6 +172,17 @@ def update_nodes_in_queue(strategy, queues)
   replace_only true
 end
 
+ruby_block "Update Slurm Accounting" do
+  block do
+    if node['cluster']['config'].dig(:Scheduling, :SlurmSettings, :Database).nil?
+      run_context.include_recipe "aws-parallelcluster-slurm::clear_slurm_accounting"
+    else
+      run_context.include_recipe "aws-parallelcluster-slurm::config_slurm_accounting"
+    end
+  end
+  only_if { ::File.exist?(node['cluster']['previous_cluster_config_path']) && is_slurm_database_updated? }
+end
+
 service 'slurmctld' do
   action :restart
   not_if { ::File.exist?(node['cluster']['previous_cluster_config_path']) && !are_queues_updated? }