From 9b29aec66107c2bff45d6aa7bc71963122ae063a Mon Sep 17 00:00:00 2001
From: dgriffi2 <dgriffi2@andrew.cmu.edu>
Date: Wed, 29 Jan 2025 20:06:07 -0500
Subject: [PATCH] initial dvosler qos role

---
 slurm_qos/defaults/main/accounts.yml |  2 +
 slurm_qos/defaults/main/main.yml     |  3 +
 slurm_qos/defaults/main/qos.yml      |  2 +
 slurm_qos/defaults/main/users.yml    |  2 +
 slurm_qos/tasks/main.yml             | 92 ++++++++++++++++++++++++++++
 5 files changed, 101 insertions(+)
 create mode 100644 slurm_qos/defaults/main/accounts.yml
 create mode 100644 slurm_qos/defaults/main/main.yml
 create mode 100644 slurm_qos/defaults/main/qos.yml
 create mode 100644 slurm_qos/defaults/main/users.yml
 create mode 100644 slurm_qos/tasks/main.yml

diff --git a/slurm_qos/defaults/main/accounts.yml b/slurm_qos/defaults/main/accounts.yml
new file mode 100644
index 0000000..e95ecc5
--- /dev/null
+++ b/slurm_qos/defaults/main/accounts.yml
@@ -0,0 +1,2 @@
+---
+slurm_qos_accounts: []
diff --git a/slurm_qos/defaults/main/main.yml b/slurm_qos/defaults/main/main.yml
new file mode 100644
index 0000000..3017fe3
--- /dev/null
+++ b/slurm_qos/defaults/main/main.yml
@@ -0,0 +1,3 @@
+---
+slurm_qos_assert_vars: true
+slurm_qos_enabled: true
diff --git a/slurm_qos/defaults/main/qos.yml b/slurm_qos/defaults/main/qos.yml
new file mode 100644
index 0000000..b1be15d
--- /dev/null
+++ b/slurm_qos/defaults/main/qos.yml
@@ -0,0 +1,2 @@
+---
+slurm_qos_settings: []
diff --git a/slurm_qos/defaults/main/users.yml b/slurm_qos/defaults/main/users.yml
new file mode 100644
index 0000000..c9d6d0b
--- /dev/null
+++ b/slurm_qos/defaults/main/users.yml
@@ -0,0 +1,2 @@
+---
+slurm_qos_users: []
diff --git a/slurm_qos/tasks/main.yml b/slurm_qos/tasks/main.yml
new file mode 100644
index 0000000..abcb446
--- /dev/null
+++ b/slurm_qos/tasks/main.yml
@@ -0,0 +1,92 @@
+---
+- name: Ensure Slurm QoS exist
+  when: slurm_qos_settings | default([]) | length > 0
+  block:
+    - name: Check Slurm QoS list
+      ansible.builtin.command: "sacctmgr -Pn show qos format=name"
+      register: qos_list
+      changed_when: false
+      failed_when: false
+
+    - name: Add Slurm QoS (if missing)
+      ansible.builtin.command: >
+        sacctmgr -i add qos {{ item.name }}
+      loop: "{{ slurm_qos_settings }}"
+      loop_control:
+        label: "{{ item.name }}"
+      when: qos_list.stdout is not search(item.name)
+      changed_when: true
+
+    - name: Modify Slurm QoS (if already exists)
+      ansible.builtin.command: >
+        sacctmgr -i modify qos {{ item.name }} set
+        {% if item.MaxNodes is defined %}MaxNodes={{ item.MaxNodes }}{% endif %}
+        {% if item.MaxSubmitJobs is defined %}MaxSubmitJobs={{ item.MaxSubmitJobs }}{% endif %}
+        {% if item.MaxJobsPerUser is defined %}MaxJobsPerUser={{ item.MaxJobsPerUser }}{% endif %}
+        {% if item.GraceTime is defined %}GraceTime={{ item.GraceTime }}{% endif %}
+        {% if item.Preempt is defined %}Preempt={{ item.Preempt }}{% endif %}
+      loop: "{{ slurm_qos_settings }}"
+      loop_control:
+        label: "{{ item.name }}"
+      when: qos_list.stdout is search(item.name)
+      changed_when: true
+
+- name: Ensure Slurm accounts exist
+  when: slurm_qos_accounts | default([]) | length > 0
+  block:
+    - name: Get Slurm cluster list
+      ansible.builtin.command: "sacctmgr -Pn show cluster format=cluster"
+      register: cluster_list
+      changed_when: false
+      failed_when: false
+
+    - name: Set cluster fact if only one cluster is found
+      set_fact:
+        slurm_cluster_name: "{{ cluster_list.stdout_lines[0] }}"
+      when: cluster_list.stdout_lines | length == 1
+
+    - name: Check if Slurm account exists
+      ansible.builtin.command: "sacctmgr -Pn list account name={{ item.name }} format=account"
+      register: account_check
+      changed_when: false
+      failed_when: false
+      loop: "{{ slurm_qos_accounts }}"
+      loop_control:
+        label: "{{ item.name }}"
+
+    - name: Add Slurm account
+      ansible.builtin.command: >
+        sacctmgr -i add account name={{ item.item.name }}
+        cluster={{ slurm_cluster_name }}
+        {% 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
+  when: slurm_qos_users | default([]) | length > 0
+  block:
+    - name: Check Slurm users
+      ansible.builtin.command: "sacctmgr -Pn list users format=user"
+      register: user_list
+      changed_when: false
+      failed_when: false
+
+    - name: Add Slurm user
+      ansible.builtin.command: >
+        sacctmgr -i add user name={{ item.username }}
+        {% if item.sponsor is defined %}Account={{ item.sponsor }}{% 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