Fu DaiandClaude Opus 4.8
1a585693be
FastAPI service wrapping the Surya-OCR-2 model (datalab-to) served through vLLM: legacy /v1/api/ai/* endpoints, an OpenAI-compatible /v1/chat/completions endpoint, a coalescing request batcher, a local OCR CLI, Docker packaging, multilingual example outputs, and quantization/concurrency benchmarks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
78 lines
3.5 KiB
Bash
Executable File
78 lines
3.5 KiB
Bash
Executable File
#!/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
|