#!/usr/bin/env bash
set -euo pipefail

export SURYA_MODEL_CHECKPOINT="${SURYA_MODEL_CHECKPOINT:-datalab-to/surya-ocr-2}"
export SURYA_INFERENCE_BACKEND="${SURYA_INFERENCE_BACKEND:-vllm}"
export SURYA_INFERENCE_URL="${SURYA_INFERENCE_URL:-http://127.0.0.1:8000/v1}"
export SURYA_INFERENCE_AUTOSTART="${SURYA_INFERENCE_AUTOSTART:-false}"
export SURYA_INFERENCE_PARALLEL="${SURYA_INFERENCE_PARALLEL:-8}"
export SURYA_INFERENCE_LOGPROBS="${SURYA_INFERENCE_LOGPROBS:-false}"
export SURYA_INFERENCE_MAX_RETRIES="${SURYA_INFERENCE_MAX_RETRIES:-1}"
# Keep in sync with VLLM_MAX_NUM_SEQS so the batcher fills every vLLM sequence
# slot (see docs/diagnosis_baseline.md). Overcommitting past it only queues.
export SURYA_INFERENCE_MAX_INFLIGHT="${SURYA_INFERENCE_MAX_INFLIGHT:-16}"
export SURYA_MAX_TOKENS_FULL_PAGE="${SURYA_MAX_TOKENS_FULL_PAGE:-6144}"
export SURYA_MAX_BLOCKS_PER_PAGE="${SURYA_MAX_BLOCKS_PER_PAGE:-80}"
export SUYA_OCR_MODE="${SUYA_OCR_MODE:-block}"
export SUYA_MAX_BATCH_SIZE="${SUYA_MAX_BATCH_SIZE:-8}"
export SUYA_BATCH_WAIT_MS="${SUYA_BATCH_WAIT_MS:-25}"
export SUYA_VLLM_IMAGE_FORMAT="${SUYA_VLLM_IMAGE_FORMAT:-JPEG}"
export SUYA_VLLM_JPEG_QUALITY="${SUYA_VLLM_JPEG_QUALITY:-92}"
export VLLM_DTYPE="${VLLM_DTYPE:-float16}"
export VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.85}"
export VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-18000}"
export VLLM_MAX_NUM_SEQS="${VLLM_MAX_NUM_SEQS:-16}"
export VLLM_MAX_BATCHED_TOKENS="${VLLM_MAX_BATCHED_TOKENS:-4096}"

# vLLM serving parameters are tuned for latency — see
# docs/quantization_benchmark_results.md §5 (A100, BF16). Three flags carry the
# win and MUST stay on; do not pass their negations via VLLM_EXTRA_ARGS:
#   --no-enforce-eager      CUDA-graph capture. The single biggest lever: eager
#                           mode (--enforce-eager) measured ~9.5x slower
#                           (48.97s vs 5.15s/page). Set explicitly so a future
#                           vLLM default flip can't silently disable graphs.
#   --enable-prefix-caching ~15% win; OCR prompts share a long fixed prefix.
#   chunked prefill (on by default) is REQUIRED — the qwen3_5 mamba/SSM cache
#                           fails engine init with --no-enable-chunked-prefill.
python3 -m vllm.entrypoints.openai.api_server \
  --host 127.0.0.1 \
  --port 8000 \
  --model "${SURYA_MODEL_CHECKPOINT}" \
  --served-model-name "${SURYA_MODEL_CHECKPOINT}" \
  --dtype "${VLLM_DTYPE}" \
  --max-model-len "${VLLM_MAX_MODEL_LEN}" \
  --max-num-seqs "${VLLM_MAX_NUM_SEQS}" \
  --max-num-batched-tokens "${VLLM_MAX_BATCHED_TOKENS}" \
  --gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION}" \
  --no-enforce-eager \
  --enable-prefix-caching \
  --mm-processor-kwargs '{"min_pixels":3136,"max_pixels":6291456}' \
  ${VLLM_EXTRA_ARGS:-} &
VLLM_PID=$!

cleanup() {
  kill "${VLLM_PID}" 2>/dev/null || true
}
trap cleanup EXIT INT TERM

python3 - <<'PY'
import os
import time
import urllib.request

base = os.environ.get("SURYA_INFERENCE_URL", "http://127.0.0.1:8000/v1")
health = base[:-3] + "/health" if base.endswith("/v1") else base.rstrip("/") + "/health"
deadline = time.time() + float(os.environ.get("SURYA_INFERENCE_STARTUP_TIMEOUT", "900"))
while time.time() < deadline:
    try:
        with urllib.request.urlopen(health, timeout=2) as response:
            if response.status == 200:
                print(f"vLLM health check passed: {health}", flush=True)
                raise SystemExit(0)
    except Exception:
        time.sleep(2)
raise SystemExit(f"vLLM did not become healthy: {health}")
PY

python3 api.py