logic for admin partition adn qos

slurm_qos/tasks/main.yml

@@ -42,7 +42,7 @@
       failed_when: false
 
     - name: Set cluster fact if only one cluster is found
-      set_fact:
+      ansible.builtin.set_fact:
         slurm_cluster_name: "{{ cluster_list.stdout_lines[0] }}"
       when: cluster_list.stdout_lines | length == 1
 
@@ -106,6 +106,51 @@
       register: user_creation_result
       failed_when: user_creation_result.rc != 0
 
+    - name: Check existing partition associations
+      ansible.builtin.command: >
+        sacctmgr -Pn show assoc where user={{ item.username }} cluster={{ slurm_cluster_name }} account={{ item.sponsor | default('orcd') }} format=Partition,QOS,DefaultQOS
+      register: assoc_check
+      loop: "{{ slurm_qos_users }}"
+      loop_control:
+        label: "{{ item.username }}"
+      when: user_list.stdout is search(item.username)
+      changed_when: false
+      failed_when: false
+
+    - name: Ensure partition associations exist for existing users
+      ansible.builtin.command: >
+        sacctmgr -i add user name={{ item.0.username }}
+        cluster={{ slurm_cluster_name }}
+        account={{ item.0.sponsor | default('orcd') }}
+        partition={{ item.1.name }}
+        {% if item.1.name == 'admin' %}
+          QOS=admin_qos DefaultQOS=admin_qos
+        {% endif %}
+      loop: "{{ slurm_qos_users | subelements('partitions') }}"
+      loop_control:
+        label: "{{ item.0.username }} - {{ item.1.name }}"
+      when:
+        - user_list.stdout is search(item.0.username)
+        - assoc_check.results | selectattr('item.username', 'equalto', item.0.username) | map(attribute='stdout') | first | default('') is not search(item.1.name)
+      changed_when: true
+
+    - name: Update QoS for existing admin partition associations
+      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 QOS=admin_qos DefaultQOS=admin_qos
+      loop: "{{ slurm_qos_users | subelements('partitions') }}"
+      loop_control:
+        label: "{{ item.0.username }} - {{ item.1.name }}"
+      when:
+        - user_list.stdout is search(item.0.username)
+        - item.1.name == 'admin'
+        - assoc_check.results | selectattr('item.username', 'equalto', item.0.username) | map(attribute='stdout') | first | default('') is search(item.1.name)
+      changed_when: true
+
     - name: Set partition-specific node limits
       ansible.builtin.command: >
         sacctmgr -i modify user where
@@ -114,6 +159,11 @@
         account={{ item.0.sponsor | default('orcd') }}
         partition={{ item.1.name }}
         set
+        {% if item.1.max_nodes is defined and item.1.max_nodes | int == -1 %}
+          GrpTRES=
+        {% elif item.1.max_nodes is defined and item.1.max_nodes | int >= 0 %}
+          GrpTRES=node={{ item.1.max_nodes }}
+        {% else %}
           GrpTRES={{ item.1.grptres | default(
             slurm_qos_partition_defaults[item.1.name].grptres if item.1.name in slurm_qos_partition_defaults else
             'node=' + (
@@ -121,13 +171,13 @@
                 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
+                '' if item.1.name == 'admin' else
                 ''
               )
             ) | string
           ) }}
+        {% endif %}
       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) }}