auto gres - v1

Author: sjpb

defaults/main.yml

@@ -12,6 +12,7 @@ openhpc_packages:
 openhpc_resume_timeout: 300
 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_gres_autodetect: 'off'
 openhpc_default_config:
   # This only defines values which are not Slurm defaults
   SlurmctldHost: "{{ openhpc_slurm_control_host }}{% if openhpc_slurm_control_host_address is defined %}({{ openhpc_slurm_control_host_address }}){% endif %}"
@@ -50,6 +51,9 @@ openhpc_cgroup_default_config:
 openhpc_config: {}
 openhpc_cgroup_config: {}
 ohpc_nodegroup_gres_types: >-
+  {% if openhpc_gres_autodetect == 'nvml' %}
+  gpu
+  {% else %}
   {{
     openhpc_nodegroups |
     community.general.json_query('[].gres[].conf') |
@@ -58,6 +62,7 @@ ohpc_nodegroup_gres_types: >-
     unique |
     join(',')
   }}
+  {% endif %}
 openhpc_gres_template: gres.conf.j2
 openhpc_cgroup_template: cgroup.conf.j2
 

library/gpu_info.py

@@ -0,0 +1,88 @@
+#!/usr/bin/python
+
+# Copyright: (c) 2025, StackHPC
+# Apache 2 License
+
+from ansible.module_utils.basic import AnsibleModule
+
+ANSIBLE_METADATA = {
+    "metadata_version": "0.1",
+    "status": ["preview"],
+    "supported_by": "community",
+}
+
+DOCUMENTATION = """
+---
+module: sacct_cluster
+short_description: Manages clusters in the accounting database
+version_added: "2.9"
+description:
+    - "Adds/removes a cluster from the accounting database"
+options:
+    name:
+        description:
+            - Name of the cluster
+        required: true
+        type: str
+    state:
+        description:
+        - If C(present), cluster will be added if it does't already exist
+        - If C(absent), cluster will be removed if it exists
+        type: str
+        required: true
+        choices: [ absent, present]
+
+requirements:
+    - "python >= 3.6"
+author:
+    - Will Szumski, StackHPC
+"""
+
+EXAMPLES = """
+"""
+
+import collections
+
+def run_module():
+    module_args = dict({})
+
+    module = AnsibleModule(argument_spec=module_args, supports_check_mode=True)
+
+    try:
+        rc ,stdout, stderr = module.run_command("nvidia-smi --query-gpu=name --format=noheader", check_rc=False, handle_exceptions=False)
+    except FileNotFoundError: # nvidia-smi not installed
+        rc = None
+    
+    # nvidia-smi return codes: https://docs.nvidia.com/deploy/nvidia-smi/index.html
+    gpus = {}
+    result = {'changed': False, 'gpus': gpus, 'gres':''}
+    if rc == 0:
+        # stdout line e.g. 'NVIDIA H200' for each GPU
+        lines = [line for line in stdout.splitlines() if line != ''] # defensive: currently no blank lines
+        models = [line.split()[1] for line in lines]
+        gpus.update(collections.Counter(models))
+    elif rc == 9:
+        # nvidia-smi installed but driver not running
+        pass
+    elif rc == None:
+        # nvidia-smi not installed
+        pass
+    else:
+        result.update({'stdout': stdout, 'rc': rc, 'stderr':stderr})
+        module.fail_json(**result)
+    
+    if len(gpus) > 0:
+        gres_parts = []
+        for model, count in gpus.items():
+            gres_parts.append(f"gpu:{model}:{count}")
+        result.update({'gres': ','.join(gres_parts)})
+
+    module.exit_json(**result)
+
+
+def main():
+    run_module()
+
+
+if __name__ == "__main__":
+    main()

tasks/runtime.yml

@@ -63,6 +63,16 @@
   notify: Restart slurmdbd service
   when: openhpc_enable.database | default(false) | bool
 
+- name: Query GPU info
+  gpu_info:
+  register: _gpu_info
+  when: openhpc_enable.batch | default(false)
+
+- name: Set fact for node GPU GRES
+  set_fact:
+    ohpc_node_gpu_gres: "{{ _gpu_info.gres }}"
+  when: openhpc_enable.batch | default(false)
+
 - name: Template slurm.conf
   template:
     src: "{{ openhpc_slurm_conf_template }}"

templates/gres.conf.j2

@@ -1,16 +1,16 @@
-AutoDetect=off
+AutoDetect={{ openhpc_gres_autodetect }}
 {% for nodegroup in openhpc_nodegroups %}
 {%      set gres_list = nodegroup.gres | default([]) %}
-{%      set gres_autodetect = nodegroup.gres_autodetect | default('off') %}
+{%      set nodegroup_gres_autodetect = nodegroup.gres_autodetect | default('off') %}
 {%      set inventory_group_name = openhpc_cluster_name ~ '_' ~ nodegroup.name %}
 {%      set inventory_group_hosts = groups.get(inventory_group_name, []) %}
 {%      set hostlist_string  = inventory_group_hosts | hostlist_expression | join(',') %}
-{%      if gres_autodetect != 'off'  %}
-NodeName={{ hostlist_string }} AutoDetect={{ gres_autodetect }}
+{%      if nodegroup_gres_autodetect != 'off'  %}
+NodeName={{ hostlist_string }} AutoDetect={{ nodegroup_gres_autodetect }}
 {%      else %}
 {%          for gres in gres_list %}
 {%              set gres_name, gres_type, _ = gres.conf.split(':') %}
-NodeName={{ hostlist_string }} Name={{ gres_name }} Type={{ gres_type }} File={{ gres.file | mandatory('The gres configuration dictionary: ' ~  gres ~ ' is missing the file key, but gres_autodetect is set to off. The error occured on node group: ' ~  nodegroup.name ~ '. Please add the file key or set gres_autodetect.') }}
+NodeName={{ hostlist_string }} Name={{ gres_name }} Type={{ gres_type }} File={{ gres.file | mandatory('The gres configuration dictionary: ' ~  gres ~ ' is missing the file key, but gres_autodetect is not specified. The error occured on node group: ' ~  nodegroup.name ~ '. Please add the file key or set gres_autodetect.') }}
 {%          endfor %}{# gres #}
 {%      endif %}{# autodetect #}
 {% endfor %}{# nodegroup #}

templates/slurm.conf.j2

@@ -38,7 +38,11 @@ NodeName={{ hostlists | join(',') }} {{ '' -}}
     CoresPerSocket={{ first_host_hv['ansible_processor_cores'] }} {{ '' -}}
     ThreadsPerCore={{ first_host_hv['ansible_processor_threads_per_core'] }} {{ '' -}}
     {{ nodegroup.node_params | default({}) | dict2parameters }} {{ '' -}}
-    {% if 'gres' in nodegroup %}Gres={{ ','.join(nodegroup.gres | map(attribute='conf')) }}{% endif %}
+    {% if 'gres' in nodegroup -%}
+    Gres={{ ','.join(nodegroup.gres | map(attribute='conf')) -}}
+    {% elif openhpc_gres_autodetect == 'nvml' and first_host_hv['ohpc_node_gpu_gres'] != '' -%}
+    Gres={{ first_host_hv['ohpc_node_gpu_gres'] -}}
+    {% endif %}
 
 {%      endif %}{# 1 or more hosts in inventory #}
 NodeSet=nodegroup_{{ nodegroup.name }} Feature=nodegroup_{{ nodegroup.name }}