#!/usr/bin/env bash
#!/usr/bin/env bash
# vLLM model serving: launch, health-check, and load-test an OpenAI-compatible endpoint.
#
# Usage: bash vllm_serve.sh [MODEL]
#   MODEL defaults to a small quantized checkpoint; override for real deployments.

set -euo pipefail

MODEL="${1:-Qwen/Qwen2.5-1.5B-Instruct}"
PORT="${PORT:-8000}"
GPU_MEM="${GPU_MEM:-0.9}"
MAX_SEQS="${MAX_SEQS:-128}"
SERVED_NAME="${SERVED_NAME:-my-llm}"
TP_SIZE="${TP_SIZE:-1}"

echo "=== vLLM Serving Configuration ==="
echo "Model:            $MODEL"
echo "Port:             $PORT"
echo "GPU mem util:     $GPU_MEM"
echo "Max sequences:    $MAX_SEQS"
echo "Served name:      $SERVED_NAME"
echo "Tensor parallel:  $TP_SIZE"

# 1. Start the OpenAI-compatible server (foreground; run in background in practice)
echo ""
echo "Starting vLLM server..."
vllm serve "$MODEL" \
  --tensor-parallel-size "$TP_SIZE" \
  --max-num-seqs "$MAX_SEQS" \
  --gpu-memory-utilization "$GPU_MEM" \
  --served-model-name "$SERVED_NAME" \
  --port "$PORT" \
  --host 0.0.0.0 &

VLLM_PID=$!
trap "kill $VLLM_PID 2>/dev/null || true" EXIT

# 2. Wait for readiness
echo "Waiting for server readiness..."
for i in $(seq 1 120); do
  if curl -sf "http://localhost:$PORT/health" >/dev/null 2>&1; then
    echo "Server ready after ${i}s."
    break
  fi
  if ! kill -0 "$VLLM_PID" 2>/dev/null; then
    echo "ERROR: vLLM process exited during startup."
    exit 1
  fi
  sleep 1
  if [ "$i" -eq 120 ]; then
    echo "ERROR: Timed out waiting for server."
    exit 1
  fi
done

# 3. Single request smoke test
echo ""
echo "=== Smoke test (single request) ==="
curl -s "http://localhost:$PORT/v1/chat/completions" \
  -H "Content-Type: application/json" \
  -d "{
    \"model\": \"$SERVED_NAME\",
    \"messages\": [{\"role\": \"user\", \"content\": \"Say hello in one sentence.\"}],
    \"max_tokens\": 32,
    \"temperature\": 0
  }" | python3 -c "import sys, json; d = json.load(sys.stdin); print(d['choices'][0]['message']['content'])"

# 4. Check model list
echo ""
echo "=== Models endpoint ==="
curl -s "http://localhost:$PORT/v1/models" | python3 -m json.tool

# 5. Load test (20 concurrent requests)
echo ""
echo "=== Load test (20 concurrent) ==="
python3 - <<'PY'
import concurrent.futures
import json
import time
import urllib.request

URL = f"http://localhost:{os.environ.get('PORT', '8000')}/v1/chat/completions"
SERVED = os.environ.get("SERVED_NAME", "my-llm")

def send_request(i):
    payload = json.dumps({
        "model": SERVED,
        "messages": [{"role": "user", "content": f"Write a short fact about number {i}."}],
        "max_tokens": 32,
        "temperature": 0,
    }).encode()
    start = time.perf_counter()
    req = urllib.request.Request(URL, data=payload, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=120) as resp:
        data = json.loads(resp.read())
    elapsed = time.perf_counter() - start
    return elapsed

start = time.perf_counter()
with concurrent.futures.ThreadPoolExecutor(max_workers=20) as pool:
    latencies = list(pool.map(send_request, range(20)))
total = time.perf_counter() - start

latencies.sort()
p50 = latencies[len(latencies)//2]
p95 = latencies[int(len(latencies)*0.95)]
print(f"  Total wall time: {total:.2f}s")
print(f"  p50 latency:     {p50:.3f}s")
print(f"  p95 latency:     {p95:.3f}s")
print(f"  Throughput:      {20/total:.1f} req/s")
PY

echo ""
echo "=== Done. Check nvidia-smi for memory usage ==="
nvidia-smi --query-gpu=memory.used,memory.total --format=csv,noheader 2>/dev/null || \
  echo "nvidia-smi not available (non-GPU environment)"