use qos partitions

slurm_qos/tasks/main.yml

@@ -62,32 +62,72 @@
         {% if item.item.organization is defined %}Organization={{ item.item.organization }}{% endif %}
         {% if item.item.description is defined %}Description="{{ item.item.description }}"{% endif %}
         {% if item.item.parent is defined %}Parent={{ item.item.parent }}{% endif %}
-        {% if item.item.default_qos is defined %}DefaultQOS={{ item.item.default_qos }}{% endif %}
-        {% if item.item.qos is defined %}QOS={{ item.item.qos | join(',') }}{% endif %}
       loop: "{{ account_check.results }}"
       loop_control:
         label: "{{ item.item.name }}"
       when: item.stdout is not search(item.item.name)
       changed_when: true
 
-- name: Ensure Slurm users exist
+- name: Ensure Slurm users exist with partition limits
   when: slurm_qos_users | default([]) | length > 0
   block:
-    - name: Check Slurm users
+    - name: Check existing Slurm users
       ansible.builtin.command: "sacctmgr -Pn list users format=user"
       register: user_list
       changed_when: false
       failed_when: false
 
-    - name: Add Slurm user
+    - name: Add Slurm user if not already present
       ansible.builtin.command: >
         sacctmgr -i add user name={{ item.username }}
-        {% if item.sponsor is defined %}Account={{ item.sponsor }}{% endif %}
+        cluster={{ slurm_cluster_name }}
+        account={{ item.sponsor | default('orcd') }}
+        {% if item.partitions is defined and item.partitions | length > 0 %}
+          partition={{ item.partitions | map(attribute='name') | join(',') }}
+        {% else %}
+          partition=debug,preempt
+        {% endif %}
+        {% if item.qos_list is defined and item.qos_list | length > 0 %}
+          QOS={{ item.qos_list | join(',') }}
+        {% else %}
+          QOS=debug_qos,preempt_qos
+        {% endif %}
+        {% if item.qos_default is defined %}
+          DefaultQOS={{ item.qos_default }}
+        {% else %}
+          DefaultQOS=preempt_qos
+        {% endif %}
         {% if item.comment is defined %}Comment="{{ item.comment }}"{% endif %}
-        {% if item.qos_list is defined and item.qos_list | length > 0 %}QOS={{ item.qos_list | join(',') }}{% endif %}
-        {% if item.qos_default is defined %}DefaultQOS={{ item.qos_default }}{% endif %}
       loop: "{{ slurm_qos_users }}"
       loop_control:
         label: "{{ item.username }}"
       when: user_list.stdout is not search(item.username)
       changed_when: true
+      register: user_creation_result
+      failed_when: user_creation_result.rc != 0
+
+    - name: Set partition-specific node limits
+      ansible.builtin.command: >
+        sacctmgr -i modify user where
+        name={{ item.0.username }}
+        cluster={{ slurm_cluster_name }}
+        account={{ item.0.sponsor | default('orcd') }}
+        partition={{ item.1.name }}
+        set
+        GrpTRES={{ item.1.grptres | default(
+          slurm_qos_partition_defaults[item.1.name].grptres if item.1.name in slurm_qos_partition_defaults else
+          'node=' + (
+            item.1.max_nodes | default(
+              slurm_qos_partition_defaults[item.1.name].max_nodes if item.1.name in slurm_qos_partition_defaults else
+              '4' if item.1.name == 'preempt' else
+              '2' if item.1.name == 'debug' else
+              ''
+            )
+          ) | string
+        ) }}
+      loop: "{{ slurm_qos_users | subelements('partitions') }}"
+      loop_control:
+        label: "{{ item.0.username }} - {{ item.1.name }}"
+      changed_when: true
+
+      #GrpTRES={{ item.1.grptres | default(slurm_qos_partition_defaults[item.1.name].grptres if item.1.name in slurm_qos_partition_defaults else 'node=' + (item.1.max_nodes | default(slurm_qos_partition_defaults[item.1.name].max_nodes if item.1.name in slurm_qos_partition_defaults else '4' if item.1.name == 'preempt' else '2' if item.1.name == 'debug' else '')) | string) }}