#!/usr/bin/env bash
# Start a long-running Docker container for SGLang benchmarking.
# The container stays alive (sleep infinity) so we can docker exec into it
# to run different TP×DP configurations without losing DeepGEMM JIT cache.
#
# Usage: start_sglang_container.sh
set -e

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"

IMAGE="${DOCKER_IMAGE:-lmsysorg/sglang:latest}"
PORT="${SGLANG_PORT:-30033}"
NAME="${EXPERIMENT}_sglang_container"
PID_FILE="${RUNTIME_BASE}/${EXPERIMENT}_container.pid"

# Persistent cache directory on host for DeepGEMM JIT kernels.
CACHE_DIR="${RUNTIME_BASE}/cache"
mkdir -p "${CACHE_DIR}/deep_gemm" "${CACHE_DIR}/tvm-ffi"

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

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

echo "=== Starting SGLang benchmark container ==="
echo "Image: $IMAGE"
echo "Container name: $NAME"
echo "Host port: $PORT"
echo "Cache dir: $CACHE_DIR"
echo "Log: $LOG"

# Start a long-running container.
# We override the entrypoint to sleep infinity so the container stays alive.
docker run -d \
  --name "$NAME" \
  --gpus all \
  --ipc host \
  --shm-size 16g \
  --entrypoint /bin/bash \
  -p "${PORT}:${PORT}" \
  -v "${MODEL_PATH}:${MODEL_PATH}:ro" \
  -v "${CACHE_DIR}/deep_gemm:/root/.cache/deep_gemm" \
  -v "${CACHE_DIR}/tvm-ffi:/root/.cache/tvm-ffi" \
  -v "${RUNTIME_BASE}/tmp:/tmp" \
  -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \
  -e PYTHONUNBUFFERED=1 \
  "$IMAGE" \
  -c "sleep infinity" \
  > "$LOG" 2>&1

# Wait a moment for container to be ready.
sleep 2

if ! docker ps --filter "name=$NAME" --format '{{.Names}}' | grep -q "$NAME"; then
  echo "ERROR: Container failed to start"
  tail -50 "$LOG"
  exit 1
fi

# Record the container ID for later use.
docker inspect -f '{{.Id}}' "$NAME" > "$PID_FILE"

echo "Container $NAME is running"
echo "Log: $LOG"
