#!/bin/bash
# Start vLLM with DSpark speculative decoding for the long-context matrix.
# Usage: start_dspark.sh <max_model_len> <max_num_seqs>
set -e

MAX_MODEL_LEN="${1:-80000}"
MAX_NUM_SEQS="${2:-64}"

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/config.env"

cd /data/user1/yy
mkdir -p logs tmp

VENV="${VENV_VLLM_DSPARK}"
export PATH="$VENV/bin:$PATH"
export PYTHONUNBUFFERED=1
export TMPDIR=/data/user1/yy/tmp
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}"

LOG="/data/user1/yy/logs/dsv4_h200_long_context_matrix_dspark_${MAX_MODEL_LEN}_$(date +%Y%m%d_%H%M%S).log"
PID_FILE="/data/user1/yy/dsv4_h200_long_context_matrix_dspark.pid"

rm -f "$PID_FILE"

echo "=== Starting DeepSeek-V4-Flash-DSpark server (TP=$TP, max_model_len=$MAX_MODEL_LEN, max_num_seqs=$MAX_NUM_SEQS) ==="
echo "Model: $DSPARK_MODEL_PATH"
echo "Port: $DSPARK_PORT"
echo "Log: $LOG"

nohup vllm serve "$DSPARK_MODEL_PATH" \
    --trust-remote-code \
    --tensor-parallel-size "$TP" \
    --kv-cache-dtype fp8 \
    --max-model-len "$MAX_MODEL_LEN" \
    --max-num-seqs "$MAX_NUM_SEQS" \
    --block-size 256 \
    --gpu-memory-utilization 0.90 \
    --tokenizer-mode deepseek_v4 \
    --reasoning-parser deepseek_v4 \
    --spec-method dspark \
    --spec-model "$DSPARK_MODEL_PATH" \
    --spec-tokens 5 \
    --no-disable-hybrid-kv-cache-manager \
    --disable-uvicorn-access-log \
    --port "$DSPARK_PORT" \
    > "$LOG" 2>&1 &

PID=$!
echo $PID > "$PID_FILE"
echo "PID: $PID"
echo "Waiting for health..."

for i in $(seq 1 240); do
    if curl --fail --silent --show-error --max-time 5 "http://127.0.0.1:${DSPARK_PORT}/health" >/dev/null 2>&1; then
        echo "DSpark server is ready at http://127.0.0.1:${DSPARK_PORT}"
        echo "Log: $LOG"
        exit 0
    fi
    if ! kill -0 $PID 2>/dev/null; then
        echo "ERROR: DSpark server exited early"
        tail -200 "$LOG"
        exit 1
    fi
    echo "Waiting... ($i/240)"
    sleep 5
done

echo "ERROR: DSpark server not healthy after 240 retries"
tail -200 "$LOG"
exit 1