#!/usr/bin/env bash
# onstart для Vast: разворачивает агента LLM-ноды и держит его запущенным.
#
# Сам скрипт намеренно тонкий: вся логика в agent.py, который качается
# отдельно. Так починка ноды не требует пересоздания шаблона Vast —
# шаблоны там неизменяемы, а agent.py достаточно перезалить.
set -u

LOG=/var/log/vllm-node.log
STATE=/opt/vllm-node
AGENT="$STATE/agent.py"
AGENT_URL="${AGENT_URL:-https://huggingface.co/cloud19/vllm-vast-node/resolve/main/agent.py}"

mkdir -p "$STATE" /var/log/vllm-node
exec >>"$LOG" 2>&1
echo "=== bootstrap $(date -Is) ==="

# Скрипты из README прода ходят в /workspace; на образе vllm рабочий каталог
# другой, поэтому связываем, чтобы ручные команды работали как на проде.
[ -e /workspace ] || ln -s /root /workspace

# Переменные из `docker -e` достаются только процессу, которого Vast запустил
# при старте контейнера. Если поднять bootstrap руками из ssh, их не будет, и
# нода молча не зарегистрируется (менеджер ответит 403). Берём их у init.
if [ -z "${MANAGER_API_KEY:-}" ] && [ -r /proc/1/environ ]; then
    while IFS= read -r -d '' kv; do
        case "$kv" in
            MODEL_ID=*|SERVED_MODEL_NAME=*|WEIGHTS_GIB=*|VLLM_API_KEY=*|HF_TOKEN=*|\
            MANAGER_URL=*|MANAGER_API_KEY=*|NODE_POOLS=*|GPU_LAYOUT=*|PORTS=*|\
            TENSOR_PARALLEL_SIZE=*|MAX_MODEL_LEN=*|MAX_NUM_SEQS=*|REGISTER=*|\
            GPU_MEMORY_UTILIZATION=*|KV_CACHE_DTYPE=*|AGENT_URL=*|VLLM_EXTRA_ARGS=*)
                export "$kv" ;;
        esac
    done </proc/1/environ
    echo "переменные подхвачены у init-процесса контейнера"
fi

# Обновляем agent.py перед каждым запуском, а не один раз при загрузке:
# иначе `vllm-node restart` поднимал бы ту же версию, что и час назад,
# и выкатить исправление на живую ноду было бы нечем.
fetch_agent() {
    for attempt in 1 2 3 4 5 6 7 8 9 10; do
        if curl -fsSL --connect-timeout 10 --max-time 120 "$AGENT_URL" -o "$AGENT.new"; then
            mv "$AGENT.new" "$AGENT"
            echo "agent.py получен ($(wc -c <"$AGENT") байт)"
            return 0
        fi
        echo "agent.py не скачался (попытка $attempt из 10): $AGENT_URL"
        sleep 10
    done
    # Старая версия лучше, чем никакой: сеть могла моргнуть.
    [ -s "$AGENT" ] && { echo "остаюсь на прошлой версии agent.py"; return 0; }
    return 1
}

if ! fetch_agent; then
    echo "ФАТАЛЬНО: agent.py так и не скачался, нода не поднимется"
    exit 1
fi

# Небольшой помощник для рук: статус, логи, снятие с пулов.
cat >/usr/local/bin/vllm-node <<'CLI'
#!/usr/bin/env bash
set -u
STATE=/opt/vllm-node
case "${1:-status}" in
  status)
    if [ -f "$STATE/state.json" ]; then python3 -m json.tool "$STATE/state.json"; else echo "агент ещё не записал состояние"; fi ;;
  logs)
    if [ -n "${2:-}" ]; then tail -f "/var/log/vllm-node/gpu$2.log"; else tail -f /var/log/vllm-node.log; fi ;;
  vllm-logs)
    tail -f /var/log/vllm-node/gpu*.log ;;
  manual)
    # агент держит vLLM, но не трогает менеджер: удобно на отладке,
    # иначе снятая руками нода вернётся в пулы через минуту
    touch "$STATE/manual"; echo "ручной режим включён, нода снимется с пулов" ;;
  auto)
    rm -f "$STATE/manual"; echo "автоматический режим, нода вернётся в пулы" ;;
  stop)
    # мягко: агент сам снимет ноду с пулов менеджера
    pkill -TERM -f 'python3 .*agent\.py' && echo "агент получил TERM" ;;
  restart)
    touch "$STATE/restarting"; pkill -TERM -f 'python3 .*agent\.py'; echo "перезапуск, следи за vllm-node logs" ;;
  *)
    echo "использование: vllm-node {status|logs [N]|vllm-logs|manual|auto|stop|restart}" ;;
esac
CLI
chmod +x /usr/local/bin/vllm-node

# Перезапускаем агента, если он упал сам. Ручная остановка (`vllm-node stop`)
# выходит с нулём — тогда цикл заканчивается и нода остаётся снятой с пулов.
while true; do
    echo "--- запуск агента $(date -Is) ---"
    python3 "$AGENT"
    code=$?
    if [ -f "$STATE/restarting" ]; then
        rm -f "$STATE/restarting"
        echo "перезапуск по просьбе"
        sleep 5
        fetch_agent
        continue
    fi
    if [ "$code" = "0" ]; then
        echo "агент остановлен штатно, выхожу"
        break
    fi
    echo "агент упал с кодом $code, перезапуск через 30 с"
    sleep 30
    fetch_agent
done