#!/bin/bash
set -e

cd /data/user1/yy
mkdir -p logs

VENV="/data/user1/yy/envs/vllm-dspark"
VLLM="$VENV/bin/vllm"
export PATH="$VENV/bin:$PATH"

MODEL="/data/models/DeepSeek-V4-Flash-DSpark"
PORT="${PORT:-30004}"
TP="${TP:-8}"
SPEC_TOKENS="${SPEC_TOKENS:-5}"
KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}"
ATTENTION_BACKEND="${ATTENTION_BACKEND:-}"
SPECULATIVE_CONFIG="${SPECULATIVE_CONFIG:-}"

LOG_TAG="tp${TP}_${KV_CACHE_DTYPE}_st${SPEC_TOKENS}"
if [[ -n "$ATTENTION_BACKEND" ]]; then
    LOG_TAG="${LOG_TAG}_${ATTENTION_BACKEND}"
fi
LOG="/data/user1/yy/logs/dsv4_flash_dspark_${LOG_TAG}_$(date +%Y%m%d_%H%M%S).log"
PID_FILE="/data/user1/yy/dsv4_flash_dspark.pid"

export TMPDIR=/data/user1/yy/tmp
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

echo "=== Starting DeepSeek-V4-Flash-DSpark (TP=$TP, spec-tokens=$SPEC_TOKENS, kv=$KV_CACHE_DTYPE) ==="
echo "Model: $MODEL"
echo "Port: $PORT"
echo "Log: $LOG"

rm -f "$PID_FILE"

cmd=(
    "$VLLM" serve "$MODEL"
    --trust-remote-code
    --tensor-parallel-size "$TP"
    --kv-cache-dtype "$KV_CACHE_DTYPE"
    --block-size 256
    --max-model-len auto
    --max-num-seqs 256
    --tokenizer-mode deepseek_v4
    --reasoning-parser deepseek_v4
    --spec-method dspark
    --spec-model "$MODEL"
    --spec-tokens "$SPEC_TOKENS"
    --no-disable-hybrid-kv-cache-manager
    --disable-uvicorn-access-log
    --port "$PORT"
)

if [[ -n "$ATTENTION_BACKEND" ]]; then
    cmd+=(--attention-backend "$ATTENTION_BACKEND")
fi

if [[ -n "$SPECULATIVE_CONFIG" ]]; then
    cmd+=(--speculative-config "$SPECULATIVE_CONFIG")
fi

nohup "${cmd[@]}" > "$LOG" 2>&1 &
PID=$!
echo $PID > "$PID_FILE"
echo "PID: $PID"
echo "Waiting for health..."

for i in $(seq 1 240); do
    if curl -s "http://127.0.0.1:$PORT/health" > /dev/null 2>&1; then
        echo "Server is ready at http://127.0.0.1:$PORT"
        echo "Log: $LOG"
        exit 0
    fi
    if ! kill -0 $PID 2>/dev/null; then
        echo "ERROR: Server exited early"
        tail -100 "$LOG"
        exit 1
    fi
    echo "Waiting... ($i/240)"
    sleep 5
done

echo "ERROR: Server not healthy after 240 retries"
tail -100 "$LOG"
exit 1