#!/usr/bin/env bash
# Start vLLM server in Docker for a given TPxDP configuration.
# Usage: start_vllm_docker.sh <TP> <DP>
#
# Uses the vllm/vllm-openai image and the same argument set as the
# bare-metal start script. The container is removed automatically on stop.
set -e

TP="${1}"
DP="${2}"

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

RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}"
mkdir -p "${RUNTIME_BASE}/logs" "${RUNTIME_BASE}/tmp"

# Persistent kernel caches on the host so TileLang/DeepGEMM JIT compilation
# survives container restarts (each TP×DP config starts a fresh container).
CACHE_DIR="${RUNTIME_BASE}/cache"
mkdir -p "${CACHE_DIR}/tilelang" "${CACHE_DIR}/deep_gemm"

IMAGE="${DOCKER_IMAGE:-vllm/vllm-openai:latest}"
PORT="${VLLM_PORT:-30032}"
NAME="${EXPERIMENT}_vllm_tp${TP}_dp${DP}"
PID_FILE="${RUNTIME_BASE}/${EXPERIMENT}_vllm_tp${TP}_dp${DP}.pid"

LOG="${RUNTIME_BASE}/logs/${EXPERIMENT}_vllm_docker_tp${TP}_dp${DP}_$(date +%Y%m%d_%H%M%S).log"
rm -f "$PID_FILE"

# Clean up any stale container with the same name.
docker rm -f "$NAME" >/dev/null 2>&1 || true

SERVER_ARGS=(
  "$MODEL_PATH"
  --trust-remote-code
  --kv-cache-dtype "$KV_CACHE_DTYPE"
  --block-size "$BLOCK_SIZE"
  --tensor-parallel-size "$TP"
  --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION"
  --max-model-len "$MAX_MODEL_LEN"
  --max-num-seqs "$MAX_NUM_SEQS"
  --no-enable-flashinfer-autotune
  --tokenizer-mode "$TOKENIZER_MODE"
  --reasoning-parser "$REASONING_PARSER"
  --spec-method "$SPEC_METHOD"
  --spec-model "$MODEL_PATH"
  --spec-tokens "$SPEC_TOKENS"
  --no-disable-hybrid-kv-cache-manager
  --disable-uvicorn-access-log
  --host 0.0.0.0
  --port "$PORT"
)

if [[ "$DP" -gt 1 ]]; then
  SERVER_ARGS+=(
    --data-parallel-size "$DP"
  )
fi

SERVER_ARGS_STR="${SERVER_ARGS[*]}"

echo "=== Starting vLLM server in Docker (TP=${TP}, DP=${DP}) ==="
echo "Image: $IMAGE"
echo "Model: $MODEL_PATH"
echo "Container name: $NAME"
echo "Host port: $PORT"
echo "Command: vllm ${SERVER_ARGS_STR}"
echo "Log: $LOG"

# Run docker in the foreground so that killing the host process stops the
# container (the --rm flag ensures cleanup). nohup lets us background it and
# capture the host PID in the same way as the bare-metal start script.
nohup docker run --rm \
  --name "$NAME" \
  --gpus all \
  --ipc host \
  --shm-size 16g \
  --ulimit memlock=-1 \
  -p "${PORT}:${PORT}" \
  -v "${MODEL_PATH}:${MODEL_PATH}:ro" \
  -v "${RUNTIME_BASE}/tmp:/tmp" \
  -v "${CACHE_DIR}/tilelang:/root/.tilelang" \
  -v "${CACHE_DIR}/deep_gemm:/root/.cache/deep_gemm" \
  -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \
  -e PYTHONUNBUFFERED=1 \
  -e HF_HUB_OFFLINE=1 \
  -e TRANSFORMERS_OFFLINE=1 \
  -e VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-3600}" \
  "$IMAGE" \
  "${SERVER_ARGS[@]}" \
  > "$LOG" 2>&1 &

PID=$!
echo $PID > "$PID_FILE"
echo "PID: $PID"
echo "Waiting for health on port ${PORT}..."

for i in $(seq 1 "${HEALTH_CHECK_RETRIES:-720}"); do
  if curl --fail --silent --show-error --max-time 5 "http://127.0.0.1:${PORT}/health" >/dev/null 2>&1; then
    echo "vLLM 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: Docker vLLM server exited early"
    tail -200 "$LOG"
    exit 1
  fi
  echo "Waiting... ($i/${HEALTH_CHECK_RETRIES:-720})"
  sleep 5
done

echo "ERROR: Docker vLLM server not healthy after ${HEALTH_CHECK_RETRIES:-720} retries"
tail -200 "$LOG"
exit 1
