feat: add warm queues

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"

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

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