#!/bin/bash
# GPU Process → Container → Pod 역추적 스크립트
# 환경: cgroup v1, containerd, systemd slice 형식
# gpu-4 노드에서 root로 실행

echo "============================================"
echo "GPU Process → Pod 매핑 (cgroup v1)"
echo "시각: $(date)"
echo "노드: $(hostname)"
echo "============================================"

# --- 함수 정의 ---

extract_container_id() {
  local pid=$1
  # cgroup v1 systemd slice: cri-containerd-<64hex>.scope
  grep ':memory:' /proc/$pid/cgroup 2>/dev/null | grep -oP 'cri-containerd-\K[a-f0-9]{64}'
}

extract_pod_uid() {
  local pid=$1
  # kubepods-<qos>-pod<UUID_underscores>.slice → 하이픈 변환
  local raw=$(grep ':memory:' /proc/$pid/cgroup 2>/dev/null | grep -oP 'pod\K[a-f0-9_]+(?=\.slice)')
  if [ -n "$raw" ]; then
    echo "$raw" | tr '_' '-'
  fi
}

extract_qos() {
  local pid=$1
  local qos=$(grep ':memory:' /proc/$pid/cgroup 2>/dev/null | grep -oP 'kubepods-\K(besteffort|burstable)(?=[\.-])' | head -1)
  if [ -n "$qos" ]; then
    echo "$qos"
  else
    grep ':memory:' /proc/$pid/cgroup 2>/dev/null | grep -q 'kubepods' && echo "guaranteed" || echo "host"
  fi
}

get_pod_info() {
  local container_id=$1
  local info=$(crictl inspect "$container_id" 2>/dev/null)
  [ -z "$info" ] && return

  local pod=$(echo "$info" | grep -oP '"io.kubernetes.pod.name":\s*"\K[^"]+' | head -1)
  local ns=$(echo "$info" | grep -oP '"io.kubernetes.pod.namespace":\s*"\K[^"]+' | head -1)
  local cont=$(echo "$info" | grep -oP '"io.kubernetes.container.name":\s*"\K[^"]+' | head -1)
  echo "${pod}|${ns}|${cont}"
}

# --- 1단계: GPU 점유 PID 수집 ---

echo ""
echo "[1] GPU 점유 프로세스 수집 (nvidia-smi + fuser)"

NVIDIA_PIDS=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits 2>/dev/null | sort -u)

FUSER_PIDS=""
for dev in /dev/nvidia[0-9]* /dev/nvidia-uvm /dev/nvidiactl; do
  [ -e "$dev" ] || continue
  PIDS=$(fuser "$dev" 2>/dev/null)
  FUSER_PIDS="$FUSER_PIDS $PIDS"
done
FUSER_PIDS=$(echo $FUSER_PIDS | tr ' ' '\n' | sort -u | grep -v '^$')

ALL_PIDS=$(echo -e "${NVIDIA_PIDS}\n${FUSER_PIDS}" | sort -u | grep -v '^$')

if [ -z "$ALL_PIDS" ]; then
  echo "  GPU 점유 프로세스 없음"
  echo ""
  echo "[!] Ghost Context 체크:"
  nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits 2>/dev/null | while IFS=', ' read -r idx mem; do
    if [ "$mem" -gt 100 ]; then
      echo "  ⚠️ GPU $idx: 프로세스 없는데 ${mem}MiB 점유 → GHOST CONTEXT"
      echo "     추적 불가. nvidia-smi --gpu-reset -i $idx 또는 노드 재부팅 필요"
    else
      echo "  ✅ GPU $idx: ${mem}MiB (정상)"
    fi
  done
  exit 0
fi

echo "  nvidia-smi PIDs: $(echo $NVIDIA_PIDS | tr '\n' ' ')"
echo "  fuser PIDs:      $(echo $FUSER_PIDS | tr '\n' ' ')"

# --- 2단계: PID → Container → Pod 매핑 ---

echo ""
echo "[2] PID → Container → Pod 매핑"
echo "--------------------------------------------------------------------------------------------------------------"
printf "%-8s %-6s %-12s %-8s %-12s %-30s %-20s %-15s\n" \
  "PID" "STATE" "GPU_MEM" "NSMI" "QOS" "POD" "NAMESPACE" "CONTAINER"
echo "--------------------------------------------------------------------------------------------------------------"

for pid in $ALL_PIDS; do

  IN_NSMI="N"
  echo "$NVIDIA_PIDS" | grep -q "^${pid}$" && IN_NSMI="Y"

  # 프로세스 생존
  if [ ! -d "/proc/$pid" ]; then
    printf "%-8s %-6s %-12s %-8s %-12s %-30s\n" \
      "$pid" "DEAD" "orphaned" "$IN_NSMI" "-" "❌ /proc 없음 — Pod 추적 불가"
    continue
  fi

  STATE=$(cat /proc/$pid/status 2>/dev/null | grep "^State:" | awk '{print $2}')

  GPU_MEM=$(nvidia-smi --query-compute-apps=pid,used_gpu_memory \
    --format=csv,noheader,nounits 2>/dev/null | grep "^${pid}," | awk -F', ' '{print $2}' | xargs)
  [ -z "$GPU_MEM" ] && GPU_MEM="fd_only"

  CONTAINER_ID=$(extract_container_id $pid)
  QOS=$(extract_qos $pid)

  if [ -z "$CONTAINER_ID" ]; then
    CMD=$(cat /proc/$pid/cmdline 2>/dev/null | tr '\0' ' ' | head -c 40)
    printf "%-8s %-6s %-12s %-8s %-12s %-30s\n" \
      "$pid" "$STATE" "${GPU_MEM}" "$IN_NSMI" "$QOS" "호스트: $CMD"
    continue
  fi

  POD_INFO=$(get_pod_info "$CONTAINER_ID")

  if [ -n "$POD_INFO" ]; then
    POD_NAME=$(echo "$POD_INFO" | cut -d'|' -f1)
    NAMESPACE=$(echo "$POD_INFO" | cut -d'|' -f2)
    CONT_NAME=$(echo "$POD_INFO" | cut -d'|' -f3)
    printf "%-8s %-6s %-12s %-8s %-12s %-30s %-20s %-15s\n" \
      "$pid" "$STATE" "${GPU_MEM}" "$IN_NSMI" "$QOS" "$POD_NAME" "$NAMESPACE" "$CONT_NAME"
  else
    SHORT_ID=$(echo "$CONTAINER_ID" | head -c 12)
    POD_UID=$(extract_pod_uid $pid)
    printf "%-8s %-6s %-12s %-8s %-12s %-30s\n" \
      "$pid" "$STATE" "${GPU_MEM}" "$IN_NSMI" "$QOS" "⚠️ 삭제됨 (ctr:${SHORT_ID} pod:${POD_UID})"
  fi
done

echo "--------------------------------------------------------------------------------------------------------------"

# --- 3단계: 숨겨진 프로세스 ---

HIDDEN_PIDS=$(comm -13 <(echo "$NVIDIA_PIDS" | sort) <(echo "$FUSER_PIDS" | sort) 2>/dev/null | grep -v '^$')
if [ -n "$HIDDEN_PIDS" ]; then
  echo ""
  echo "[3] ⚠️ nvidia-smi 미등록 but GPU fd 보유 프로세스:"
  for pid in $HIDDEN_PIDS; do
    if [ -d "/proc/$pid" ]; then
      CMD=$(cat /proc/$pid/cmdline 2>/dev/null | tr '\0' ' ' | head -c 60)
      FDS=$(ls -la /proc/$pid/fd 2>/dev/null | grep nvidia | wc -l)
      CONTAINER_ID=$(extract_container_id $pid)
      if [ -n "$CONTAINER_ID" ]; then
        POD_INFO=$(get_pod_info "$CONTAINER_ID")
        POD_NAME=$(echo "$POD_INFO" | cut -d'|' -f1)
        echo "  PID $pid: nvidia fd ${FDS}개, Pod: ${POD_NAME:-unknown}, CMD: $CMD"
      else
        echo "  PID $pid: nvidia fd ${FDS}개, 호스트 프로세스, CMD: $CMD"
      fi
    fi
  done
fi

# --- 4단계: 요약 ---

echo ""
echo "[4] 요약"
TOTAL=$(echo "$ALL_PIDS" | grep -v '^$' | wc -l)
DEAD=0; ZOMBIE=0; HIDDEN_COUNT=0
for pid in $ALL_PIDS; do
  [ ! -d "/proc/$pid" ] && DEAD=$((DEAD+1))
  [ -d "/proc/$pid" ] && grep -q "^State:.*Z" /proc/$pid/status 2>/dev/null && ZOMBIE=$((ZOMBIE+1))
done
[ -n "$HIDDEN_PIDS" ] && HIDDEN_COUNT=$(echo "$HIDDEN_PIDS" | grep -v '^$' | wc -l)

echo "  전체 GPU 관련 프로세스: $TOTAL"
echo "  정상:                   $((TOTAL - DEAD - ZOMBIE))"
echo "  죽은 프로세스 (orphaned ctx): $DEAD"
echo "  좀비 프로세스:          $ZOMBIE"
echo "  nvidia-smi 미표시:      $HIDDEN_COUNT"

if [ "$DEAD" -gt 0 ] || [ "$ZOMBIE" -gt 0 ]; then
  echo ""
  echo "  ⚠️ 조치 필요:"
  [ "$ZOMBIE" -gt 0 ] && echo "     좀비 → 부모 프로세스 kill 또는 Pod 재시작"
  [ "$DEAD" -gt 0 ] && echo "     orphaned ctx → Pod 재시작, 안 되면 nvidia-smi --gpu-reset 또는 노드 재부팅"
fi

echo ""
echo "============================================"