#!/bin/bash
# 一次性演示：一个 Pod 占住 nvidia.com/gpu，另一个因此排队。
# 结束时删除这两个 Pod。vLLM 不在这次变更里。
set -u
NS=gpu-operator
HOLD=gpu-demo-hold
QUEUE=gpu-demo-queue
IMG=docker.m.daocloud.io/vllm/vllm-openai:latest

cleanup() {
  kubectl -n "$NS" delete pod "$HOLD" "$QUEUE" --ignore-not-found --wait=true >/dev/null 2>&1 || true
}
trap cleanup EXIT

section() {
  echo
  echo "========== $1 =========="
}

if kubectl get pods -A -o json | python3 -c '
import json,sys
data=json.load(sys.stdin)
bad=[]
for p in data.get("items",[]):
    name=p["metadata"]["name"]
    if name in ("gpu-demo-hold","gpu-demo-queue"):
        continue
    for c in p["spec"].get("containers",[]):
        lim=c.get("resources",{}).get("limits",{}) or {}
        req=c.get("resources",{}).get("requests",{}) or {}
        if "nvidia.com/gpu" in lim or "nvidia.com/gpu" in req:
            bad.append(p["metadata"]["namespace"]+"/"+name)
if bad:
    print("已有 Pod 在申请 nvidia.com/gpu: "+", ".join(bad))
    sys.exit(3)
'; then
  echo "当前没有别的 Pod 申请 nvidia.com/gpu，可以开始演示。"
else
  rc=$?
  if [ "$rc" = "3" ]; then
    exit 3
  fi
  echo "检查现有 GPU 申请失败"
  exit 1
fi

section "演示前：节点名额"
kubectl get node prod102 -o jsonpath='allocatable={.status.allocatable.nvidia\.com/gpu}{"\n"}'
echo -n "vLLM /health "
curl -s -o /dev/null -w '%{http_code}\n' --max-time 5 http://127.0.0.1:8000/health || echo fail

section "提交第一个任务 gpu-demo-hold"
cat <<EOF | kubectl apply -f -
apiVersion: v1
kind: Pod
metadata:
  name: ${HOLD}
  namespace: ${NS}
  labels:
    app: gpu-demo
spec:
  restartPolicy: Never
  containers:
    - name: smi
      image: ${IMG}
      imagePullPolicy: IfNotPresent
      command: ["bash","-lc","nvidia-smi -L; echo ENV_BEGIN; env | grep -E 'NVIDIA|CUDA' | sort; echo ENV_END; sleep 180"]
      resources:
        limits:
          nvidia.com/gpu: "1"
EOF

echo "等待 ${HOLD} 进入 Running ..."
ok=0
for i in $(seq 1 40); do
  phase=$(kubectl -n "$NS" get pod "$HOLD" -o jsonpath='{.status.phase}' 2>/dev/null || true)
  echo "t=${i} phase=${phase}"
  if [ "$phase" = "Running" ]; then
    ok=1
    break
  fi
  if [ "$phase" = "Failed" ] || [ "$phase" = "Succeeded" ]; then
    break
  fi
  sleep 3
done

section "第一个任务的事件"
kubectl -n "$NS" describe pod "$HOLD" | sed -n '/Events:/,$p'

section "节点上已分配的 nvidia.com/gpu"
kubectl describe node prod102 | sed -n '/Allocated resources:/,/Events:/p'

section "第一个任务的输出"
kubectl -n "$NS" logs "$HOLD" || true

if [ "$ok" != "1" ]; then
  echo "第一个任务没有进入 Running，停止演示并清理。"
  exit 1
fi

section "提交第二个任务 gpu-demo-queue"
cat <<EOF | kubectl apply -f -
apiVersion: v1
kind: Pod
metadata:
  name: ${QUEUE}
  namespace: ${NS}
  labels:
    app: gpu-demo
spec:
  restartPolicy: Never
  containers:
    - name: smi
      image: ${IMG}
      imagePullPolicy: IfNotPresent
      command: ["nvidia-smi","-L"]
      resources:
        limits:
          nvidia.com/gpu: "1"
EOF

sleep 8
section "两个任务并排"
kubectl -n "$NS" get pod "$HOLD" "$QUEUE" -o wide

section "第二个任务排队原因"
kubectl -n "$NS" describe pod "$QUEUE" | sed -n '/Events:/,$p'

section "删掉第一个任务，名额还给调度器"
# 先摘掉退出陷阱里的重复删除，后面按顺序删
trap - EXIT
kubectl -n "$NS" delete pod "$HOLD" --wait=true

echo "等待 ${QUEUE} 跑完 ..."
qok=0
for i in $(seq 1 40); do
  phase=$(kubectl -n "$NS" get pod "$QUEUE" -o jsonpath='{.status.phase}' 2>/dev/null || true)
  echo "t=${i} phase=${phase}"
  if [ "$phase" = "Succeeded" ] || [ "$phase" = "Failed" ]; then
    qok=1
    break
  fi
  sleep 3
done

section "第二个任务的事件"
kubectl -n "$NS" describe pod "$QUEUE" | sed -n '/Events:/,$p'

section "第二个任务的输出"
kubectl -n "$NS" logs "$QUEUE" || true

section "清理第二个任务"
kubectl -n "$NS" delete pod "$QUEUE" --wait=true

section "演示后"
echo -n "仍在申请 nvidia.com/gpu 的 Pod："
kubectl get pods -A -o json | python3 -c '
import json,sys
data=json.load(sys.stdin)
found=[]
for p in data.get("items",[]):
    for c in p["spec"].get("containers",[]):
        lim=c.get("resources",{}).get("limits",{}) or {}
        req=c.get("resources",{}).get("requests",{}) or {}
        if "nvidia.com/gpu" in lim or "nvidia.com/gpu" in req:
            found.append(p["metadata"]["namespace"]+"/"+p["metadata"]["name"])
print(", ".join(found) if found else "无")
'
kubectl get node prod102 -o jsonpath='allocatable={.status.allocatable.nvidia\.com/gpu}{"\n"}'
echo -n "vLLM /health "
curl -s -o /dev/null -w '%{http_code}\n' --max-time 5 http://127.0.0.1:8000/health || echo fail
kubectl -n gpu-operator get pods

if [ "$qok" != "1" ]; then
  echo "第二个任务没有在预期时间内结束"
  exit 1
fi
echo
echo "演示结束，两个任务都已删除。"