From 80e7a645795be90f161b231328054acd0164515d Mon Sep 17 00:00:00 2001
From: dvosler <dvosler@cs.cmu.edu>
Date: Wed, 26 Feb 2025 13:51:17 -0500
Subject: [PATCH] logic for admin partition adn qos

---
 slurm_qos/tasks/main.yml | 78 ++++++++++++++++++++++++++++++++--------
 1 file changed, 64 insertions(+), 14 deletions(-)

diff --git a/slurm_qos/tasks/main.yml b/slurm_qos/tasks/main.yml
index b6739ff..4b2b262 100644
--- a/slurm_qos/tasks/main.yml
+++ b/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,20 +159,25 @@
         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
-        ) }}
+        {% 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=' + (
+              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
+                '' 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) }}