#!/bin/sh
# /usr/local/bin/prestop.sh
# Pod preStop hook용 GPU cleanup 스크립트
# 시나리오 3·4 공통 사용
#   - 시나리오 3 (:tini, worker가 SIGTERM 무시): fallback 경로 트리거
#   - 시나리오 4 (:fullfix, worker graceful): fallback 경로 미트리거

echo "[preStop] GPU cleanup 시작: $(date)"

# 1. nvidia-smi + fuser로 GPU 점유 PID 수집
GPU_PIDS=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits 2>/dev/null)
FUSER_PIDS=""
for dev in /dev/nvidia[0-9]* /dev/nvidia-uvm; do
  [ -e "$dev" ] || continue
  PIDS=$(fuser "$dev" 2>/dev/null)
  FUSER_PIDS="$FUSER_PIDS $PIDS"
done
ALL_PIDS=$(echo "$GPU_PIDS $FUSER_PIDS" | tr ' ' '\n' | sort -u | grep -v '^$')
echo "[preStop] 전체 GPU 관련 PIDs: $ALL_PIDS"

# 2. PID 1(tini) 제외하고 SIGTERM 전송
for pid in $ALL_PIDS; do
  if [ "$pid" != "1" ]; then
    echo "[preStop] SIGTERM -> PID $pid"
    kill -TERM $pid 2>/dev/null
  fi
done

# 3. 최대 20초 대기 — graceful cleanup
for i in $(seq 1 20); do
  REMAINING=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits 2>/dev/null | wc -l)
  if [ "$REMAINING" -eq 0 ]; then
    echo "[preStop] GPU 프로세스 정리 완료 ✅"
    break
  fi
  echo "[preStop] 대기 ($i/20) 남은 프로세스: $REMAINING"
  sleep 1
done

# 4. SIGKILL fallback (시나리오 3에서 트리거됨)
STILL=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits 2>/dev/null | wc -l)
if [ "$STILL" -gt 0 ]; then
  echo "[preStop] ⚠️ SIGTERM 무시됨 → SIGKILL fallback"
  for pid in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits 2>/dev/null); do
    kill -9 $pid 2>/dev/null
  done
fi

echo "[preStop] 완료: $(date)"