Support hosts in multiple partitions when templating gres.conf

Author: jovial

defaults/main.yml

@@ -13,6 +13,7 @@ openhpc_retry_delay: 10
 openhpc_job_maxtime: '60-0' # quote this to avoid ansible converting some formats to seconds, which is interpreted as minutes by Slurm
 openhpc_config: "{{ openhpc_extra_config | default({}) }}"
 openhpc_gres_template: gres.conf.j2
+openhpc_gres_autodetect: 'disabled'
 openhpc_slurm_configless: "{{ 'enable_configless' in openhpc_config.get('SlurmctldParameters', []) }}"
 
 openhpc_state_save_location: /var/spool/slurm

filter_plugins/slurm_conf.py

@@ -15,6 +15,7 @@
 # NB: To test this from the repo root run:
 #   ansible-playbook -i tests/inventory -i tests/inventory-mock-groups tests/filter.yml
 
+from collections import defaultdict
 from ansible import errors
 import jinja2
 import re
@@ -65,6 +66,16 @@ def hostlist_expression(hosts):
             unmatchable.append(v)
     return ['{}[{}]'.format(k, _group_numbers(v)) for k, v in results.items()] + unmatchable
 
+@jinja2.pass_context
+def group_by_gres_autodetect(context, hosts):
+    """ Groups a list of hosts by the gres_autodetect variable
+    """
+    result = defaultdict(list)
+    for host in hosts:
+        gres_autodetect = _get_hostvar(context, 'gres_autodetect', inventory_hostname=host)
+        result[gres_autodetect].append(host)
+    return result
+
 def _group_numbers(numbers):
     units = []
     ints = [int(n) for n in numbers]
@@ -133,4 +144,5 @@ def filters(self):
             'error': error,
             'dict2parameters': dict2parameters,
             'config2dict': config2dict,
+            'group_by_gres_autodetect': group_by_gres_autodetect,
         }

templates/gres.conf.j2

@@ -1,20 +1,29 @@
 AutoDetect=off
 {% for part in openhpc_slurm_partitions %}
 {%  set nodelist = [] %}
+{%  set donehosts = set() %}
 {%  for group in part.get('groups', [part]) %}
-{%      set group_name = group.cluster_name|default(openhpc_cluster_name) ~ '_' ~ group.name %}
-{%      set inventory_group_hosts = groups.get(group_name, []) %}
-{%      if 'gres_autodetect' in group %}
-{%          for hostlist in (inventory_group_hosts | hostlist_expression) %}
-NodeName={{ hostlist }} AutoDetect={{ group.gres_autodetect }}
-{%          endfor %}
-{%      elif 'gres' in group %}
-{%          for gres in group.gres %}
-{%              set gres_name, gres_type, _ = gres.conf.split(':') %}
-{%              for hostlist in (inventory_group_hosts | hostlist_expression) %}
+{%      if 'gres' in group %}
+{%          set group_name = group.cluster_name|default(openhpc_cluster_name) ~ '_' ~ group.name %}
+{%          set inventory_group_hosts = groups.get(group_name, []) %}
+{%          set autodetect_mechanisms = inventory_group_hosts | group_by_gres_autodetect %}
+{%          for mechanism, _mechanism_hosts in autodetect_mechanisms.items() %}
+{%              set mechanism_hosts = _mechanism_hosts - donehosts %}
+{%              if mechanism != 'disabled' %}
+{%                  for hostlist in (mechanism_hosts | hostlist_expression) %}
+NodeName={{ hostlist }} AutoDetect={{ mechanism }}
+
+{%                  endfor %}
+{%              else %}
+{%                  for gres in group.gres %}
+{%                      set gres_name, gres_type, _ = gres.conf.split(':') %}
+{%                      for hostlist in (mechanism_hosts | hostlist_expression) %}
 NodeName={{ hostlist }} Name={{ gres_name }} Type={{ gres_type }} File={{ gres.file }}
 
-{%              endfor %}
+{%                      endfor %}
+{%                  endfor %}
+{%              endif %}
+                {% set _ = donehosts.update(mechanism_hosts) %}
 {%          endfor %}
 {%      endif %}
 {%  endfor %}