Suya OCR API — vLLM-backed, OpenAI-compatible OCR service
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>
This commit is contained in:
@@ -0,0 +1,76 @@
|
||||
#!/usr/bin/env bash
|
||||
# MTP (multi-token-prediction) speculative-decode sweep for the latency report.
|
||||
# Serves the BF16 model three ways — no speculation (baseline), MTP with 1
|
||||
# speculative token, MTP with 2 — and captures the same page subset through the
|
||||
# OCR API for each. The model ships 1 nextn-predict layer, so MTP=1 is the
|
||||
# expected-valid setting and MTP=2 is exploratory.
|
||||
#
|
||||
# Run INSIDE the bench container: bash scripts/quant/tune_mtp.sh
|
||||
# Requires: api.py running on :5002 pointed at :8000; compat libs on path.
|
||||
set -uo pipefail
|
||||
|
||||
MODEL="${MODEL:-datalab-to/surya-ocr-2}"
|
||||
OCR_URL="http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/"
|
||||
N_PAGES="${N_PAGES:-5}"
|
||||
OUT_ROOT="results/quant/mtp"
|
||||
LOG=/tmp/tune_mtp.log
|
||||
export LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH:-}"
|
||||
|
||||
mapfile -t ALL < <(sed '/^#/d;/^$/d' eval_set/manifest.txt)
|
||||
IMAGES=("${ALL[@]:0:$N_PAGES}")
|
||||
|
||||
# config_name | extra vLLM args. The speculative-config JSON is written compact
|
||||
# (no spaces) so it survives word-splitting as a single argv token.
|
||||
CONFIGS=(
|
||||
"baseline|--enable-prefix-caching"
|
||||
"mtp1|--enable-prefix-caching --speculative-config {\"method\":\"mtp\",\"num_speculative_tokens\":1}"
|
||||
"mtp2|--enable-prefix-caching --speculative-config {\"method\":\"mtp\",\"num_speculative_tokens\":2}"
|
||||
)
|
||||
|
||||
base_args() {
|
||||
echo "--host 127.0.0.1 --port 8000 --model $MODEL --served-model-name $MODEL \
|
||||
--max-model-len 18000 --max-num-seqs 16 --gpu-memory-utilization 0.85 --no-enforce-eager \
|
||||
--mm-processor-kwargs {\"min_pixels\":3136,\"max_pixels\":6291456}"
|
||||
}
|
||||
|
||||
stop_server() {
|
||||
local pid
|
||||
pid=$(ss -ltnp 2>/dev/null | grep ":8000 " | grep -oP 'pid=\K[0-9]+' | head -1)
|
||||
[ -n "$pid" ] && kill -9 "$pid" 2>/dev/null
|
||||
pkill -9 -f "vllm.entrypoints" 2>/dev/null
|
||||
# EngineCore outlives the API server and holds the GPU; kill every remaining
|
||||
# CUDA compute process, then poll until the memory is actually freed.
|
||||
for p in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null); do
|
||||
kill -9 "$p" 2>/dev/null
|
||||
done
|
||||
for _ in $(seq 1 30); do
|
||||
used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | head -1)
|
||||
[ "${used:-99999}" -lt 2000 ] && break
|
||||
sleep 2
|
||||
done
|
||||
}
|
||||
|
||||
for entry in "${CONFIGS[@]}"; do
|
||||
name="${entry%%|*}"
|
||||
extra="${entry#*|}"
|
||||
echo "=== mtp config: $name ($extra) ===" | tee -a "$LOG"
|
||||
stop_server
|
||||
# shellcheck disable=SC2046
|
||||
nohup python3 -m vllm.entrypoints.openai.api_server $(base_args) $extra \
|
||||
> "/tmp/mtp_${name}.log" 2>&1 &
|
||||
ok=0
|
||||
for _ in $(seq 1 96); do
|
||||
if curl -fs http://127.0.0.1:8000/health >/dev/null 2>&1; then ok=1; break; fi
|
||||
if grep -qiE "Engine core initialization failed|ValueError|RuntimeError|Traceback" "/tmp/mtp_${name}.log" 2>/dev/null; then break; fi
|
||||
sleep 5
|
||||
done
|
||||
if [ "$ok" -ne 1 ]; then
|
||||
echo "$name: FAILED to start (see /tmp/mtp_${name}.log)" | tee -a "$LOG"
|
||||
continue
|
||||
fi
|
||||
python3 -m scripts.quant.capture --url "$OCR_URL" --out-dir "$OUT_ROOT/$name" "${IMAGES[@]}" \
|
||||
>> "$LOG" 2>&1
|
||||
echo "$name: captured $(ls "$OUT_ROOT/$name"/*.json 2>/dev/null | wc -l) pages" | tee -a "$LOG"
|
||||
done
|
||||
stop_server
|
||||
echo "MTP_DONE" | tee -a "$LOG"
|
||||
Reference in New Issue
Block a user