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>
60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
"""Launch a vLLM OpenAI server for one quant variant, health-check it, and tear
|
|
it down by port (never by command-line match — that bit us in Approach A)."""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import time
|
|
import urllib.request
|
|
from typing import List, Optional
|
|
|
|
# Default seconds to wait for a server to become healthy. Overridable via env so a
|
|
# method that fails to boot is recorded quickly instead of blocking the sweep.
|
|
DEFAULT_HEALTH_TIMEOUT = float(os.environ.get("SURYA_QUANT_HEALTH_TIMEOUT", "900"))
|
|
|
|
|
|
def health_url(port: int) -> str:
|
|
return f"http://127.0.0.1:{port}/health"
|
|
|
|
|
|
def parse_listening_pid(ss_output: str, port: int) -> Optional[int]:
|
|
for line in ss_output.splitlines():
|
|
if f":{port} " in line or line.rstrip().endswith(f":{port}"):
|
|
m = re.search(r"pid=(\d+)", line)
|
|
if m:
|
|
return int(m.group(1))
|
|
return None
|
|
|
|
|
|
def wait_healthy(port: int, timeout: float = DEFAULT_HEALTH_TIMEOUT) -> bool:
|
|
deadline = time.time() + timeout
|
|
url = health_url(port)
|
|
while time.time() < deadline:
|
|
try:
|
|
with urllib.request.urlopen(url, timeout=2) as resp:
|
|
if resp.status == 200:
|
|
return True
|
|
except Exception:
|
|
time.sleep(3)
|
|
return False
|
|
|
|
|
|
def start_server(serve_args: List[str], log_path: str) -> subprocess.Popen:
|
|
cmd = ["python", "-m", "vllm.entrypoints.openai.api_server", *serve_args]
|
|
log = open(log_path, "w")
|
|
return subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT)
|
|
|
|
|
|
def stop_server(port: int, proc: Optional[subprocess.Popen] = None) -> None:
|
|
if proc is not None:
|
|
proc.terminate()
|
|
try:
|
|
proc.wait(timeout=30)
|
|
except Exception:
|
|
proc.kill()
|
|
out = subprocess.run(["ss", "-ltnp"], capture_output=True, text=True).stdout
|
|
pid = parse_listening_pid(out, port)
|
|
if pid:
|
|
subprocess.run(["kill", str(pid)])
|