#!/bin/sh
# Memory watchdog for the first start of a new vLLM configuration on a machine
# with unified memory (e.g. NVIDIA DGX Spark). When the system runs out of
# memory there, it freezes instead of returning an error, so we stop the
# container before that happens.
#
# Checks every second. Kills the container (no grace period) when
#   - MemAvailable drops below MIN_AVAIL_MIB, or
#   - swap usage grows by more than MAX_SWAP_GROWTH_MIB over the value at start.
# All values in MiB, read from /proc/meminfo.
# Logs every sample, so the memory curve of the start can be read afterwards.
#
# Usage:  ./vllm-start-watchdog.sh <container-name> [minutes]
# Example: ./vllm-start-watchdog.sh vllm-mid 90 &
set -u

NAME="${1:?usage: $0 <container-name> [minutes]}"
MINUTES="${2:-90}"
MIN_AVAIL_MIB="${MIN_AVAIL_MIB:-12000}"
MAX_SWAP_GROWTH_MIB="${MAX_SWAP_GROWTH_MIB:-8000}"
LOG="${LOG:-./watchdog-$NAME.log}"

mem() { awk -v k="$1" '$1==k":"{print int($2/1024)}' /proc/meminfo; }  # kB -> MiB
swap_used() { echo $(( $(mem SwapTotal) - $(mem SwapFree) )); }

base_swap="$(swap_used)"
end=$(( $(date +%s) + MINUTES * 60 ))
echo "$(date '+%F %T') start: container=$NAME min_avail=${MIN_AVAIL_MIB}MiB swap_base=${base_swap}MiB max_growth=${MAX_SWAP_GROWTH_MIB}MiB" >> "$LOG"

while [ "$(date +%s)" -lt "$end" ]; do
  avail="$(mem MemAvailable)"
  swap="$(swap_used)"
  state="$(docker inspect -f '{{.State.Status}}' "$NAME" 2>/dev/null)" || state=missing
  echo "$(date '+%T') avail=${avail}MiB swap=${swap}MiB container=$state" >> "$LOG"
  if [ "$state" = running ]; then
    if [ "$avail" -lt "$MIN_AVAIL_MIB" ] || [ $(( swap - base_swap )) -gt "$MAX_SWAP_GROWTH_MIB" ]; then
      echo "$(date '+%T') ALARM avail=${avail}MiB swap=${swap}MiB -> docker kill $NAME" >> "$LOG"
      docker kill "$NAME" >> "$LOG" 2>&1
    fi
  fi
  sleep 1
done
echo "$(date '+%F %T') end" >> "$LOG"