From 742ef4e3ac2954ae506e2e96d4352e5328db8d65 Mon Sep 17 00:00:00 2001
From: Leopold Talirz <leopold.talirz@microsoft.com>
Date: Tue, 12 Sep 2023 17:42:53 +0200
Subject: [PATCH] feat: add warm queues

Warm queues will always hold at least 1 idling node so that new jobs can start instantly.
The cronjob that checks for this is run every 5 minutes, on workdays only.
---
 config.schema.json                  |  4 ++++
 playbooks/scheduler.yml             | 20 +++++++++++++++++++
 playbooks/templates/queue-warmup.sh | 31 +++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+)
 create mode 100644 playbooks/templates/queue-warmup.sh

diff --git a/config.schema.json b/config.schema.json
index c7ab5963b..405e07e92 100644
--- a/config.schema.json
+++ b/config.schema.json
@@ -1077,6 +1077,10 @@
                     "description": "Whether to create placement groups (SLURM only). Defaults to true",
                     "type": "boolean"
                 },
+                "warmup": {
+                    "description": "Always keep 1 node running. Defaults to false",
+                    "type": "boolean"
+                },
                 "idle_timeout": {
                     "description": "Idle time (seconds) before shutting down VMs. Choose lower than autoscale.idle_timeout",
                     "type": "integer"
diff --git a/playbooks/scheduler.yml b/playbooks/scheduler.yml
index 9db1820e0..08c029177 100644
--- a/playbooks/scheduler.yml
+++ b/playbooks/scheduler.yml
@@ -54,3 +54,23 @@
       name: '*'
       state: latest
 #      exclude: kernel*,kmod*,amlfs*
+
+  - name: copy queue warmup script
+    template:
+      src: 'queue-warmup.sh'
+      dest: '/usr/local/sbin/queue-warmup.sh'
+      mode: '0733'
+    when: ( queue_manager is defined and queue_manager == "slurm" )
+    tags: scheduler-warmup
+  - name: set up cronjob for queue warmup
+    cron:
+      name: "queue-warmup"
+      job: "/usr/local/sbin/queue-warmup.sh {{ warmup_queues | map(attribute='name') | join(' ') }}"
+      minute: "*/5"
+      weekday: 1-5
+      user: "root"
+      state: "present"
+    vars:
+      warmup_queues: "{{ queues | selectattr('warmup', 'defined') | selectattr('warmup', 'equalto', true) }}"
+    when: ( queue_manager is defined and queue_manager == "slurm" )
+    tags: scheduler-warmup
diff --git a/playbooks/templates/queue-warmup.sh b/playbooks/templates/queue-warmup.sh
new file mode 100644
index 000000000..f2d3861c4
--- /dev/null
+++ b/playbooks/templates/queue-warmup.sh
@@ -0,0 +1,31 @@
+#!/bin/bash
+# Usage: ./warmup-queues.sh viz hb2la
+set -e
+
+# SLURM node states & state flags on AZ-HOP
+
+# idle   VM allocated and idling
+# idle~  VM not allocated from Azure
+# idle#  VM being allocated from Azure
+# idle%  VM being powered down
+# mix    Some CPUs allocated but not all
+
+for queue in "$@"; do
+  available=`sinfo -p $queue --states=mix,idle --noheader | grep -v idle~ | grep -v idle# | grep -v idle% | wc -l`
+  allocating=`sinfo -p $queue --states=idle --noheader | grep idle# | wc -l`
+
+  if [[ $available == 0 && $allocating == 0 ]]; then
+    echo "Allocating 1 node on queue $queue"
+    srun --partition $queue bash > /dev/null 2>&1 &
+    PID=$!
+    sleep 2
+    set +e
+    kill $PID
+    set -e
+  elif [[ $available -gt 0 ]]; then
+    # "touch" one available node so that it won't be deallocated by slurm after timeout
+    set +e
+    srun --partition $queue "exit" > /dev/null 2>&1 &
+    set -e
+  fi
+done