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:
Fu Dai
@@ -0,0 +1,185 @@
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import statistics
|
||||
import subprocess
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
DEFAULT_OLD_ENDPOINT = "http://127.0.0.1:5002/v1/api/ai/suya_ocr/"
|
||||
DEFAULT_NEW_ENDPOINT = "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/"
|
||||
|
||||
|
||||
def _encode_image(path: Path) -> str:
|
||||
return base64.b64encode(path.read_bytes()).decode("utf-8")
|
||||
|
||||
|
||||
def _preflight() -> Dict[str, Any]:
|
||||
checks: Dict[str, Any] = {}
|
||||
commands = {
|
||||
"docker_runtimes": ["docker", "info", "--format", "{{json .Runtimes}}"],
|
||||
"nvidia_smi": ["nvidia-smi", "-L"],
|
||||
"nvidia_container_runtime": ["which", "nvidia-container-runtime"],
|
||||
"nvidia_ctk": ["which", "nvidia-ctk"],
|
||||
"docker_gpus": [
|
||||
"docker",
|
||||
"run",
|
||||
"--rm",
|
||||
"--gpus",
|
||||
"all",
|
||||
"--entrypoint",
|
||||
"nvidia-smi",
|
||||
"vllm/vllm-openai:v0.20.1",
|
||||
"-L",
|
||||
],
|
||||
}
|
||||
for name, cmd in commands.items():
|
||||
try:
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
|
||||
checks[name] = {
|
||||
"ok": result.returncode == 0,
|
||||
"stdout": result.stdout.strip(),
|
||||
"stderr": result.stderr.strip(),
|
||||
}
|
||||
except Exception as exc:
|
||||
checks[name] = {"ok": False, "stdout": "", "stderr": str(exc)}
|
||||
runtimes = checks.get("docker_runtimes", {}).get("stdout", "")
|
||||
checks["docker_has_nvidia_runtime"] = '"nvidia"' in runtimes
|
||||
if not checks["docker_has_nvidia_runtime"] and checks.get("docker_gpus", {}).get("ok"):
|
||||
checks["recommended_fix"] = (
|
||||
"Docker GPU passthrough works with --gpus. Use the local "
|
||||
"scripts/docker_nvidia_runtime_compat.sh wrapper to strip the legacy "
|
||||
"--runtime nvidia flag from Surya's vLLM spawn command."
|
||||
)
|
||||
else:
|
||||
checks["recommended_fix"] = (
|
||||
"Run: sudo nvidia-ctk runtime configure --runtime=docker "
|
||||
"--config=/etc/docker/daemon.json && sudo systemctl restart docker"
|
||||
)
|
||||
return checks
|
||||
|
||||
|
||||
def _percentile(values: List[float], percentile: float) -> float:
|
||||
if not values:
|
||||
return 0.0
|
||||
ordered = sorted(values)
|
||||
index = min(len(ordered) - 1, int(round((percentile / 100) * (len(ordered) - 1))))
|
||||
return ordered[index]
|
||||
|
||||
|
||||
def _post_once(url: str, payload: Dict[str, Any], timeout: float) -> Dict[str, Any]:
|
||||
start = time.perf_counter()
|
||||
try:
|
||||
response = requests.post(url, json=payload, timeout=timeout)
|
||||
elapsed = time.perf_counter() - start
|
||||
body = response.json() if response.content else {}
|
||||
return {
|
||||
"ok": response.status_code == 200 and body.get("code") == 200,
|
||||
"status_code": response.status_code,
|
||||
"api_code": body.get("code"),
|
||||
"elapsed_seconds": elapsed,
|
||||
"error": body.get("message") if body.get("code") != 200 else None,
|
||||
}
|
||||
except Exception as exc:
|
||||
return {
|
||||
"ok": False,
|
||||
"status_code": None,
|
||||
"api_code": None,
|
||||
"elapsed_seconds": time.perf_counter() - start,
|
||||
"error": str(exc),
|
||||
}
|
||||
|
||||
|
||||
def run_endpoint(
|
||||
name: str,
|
||||
url: str,
|
||||
payload: Dict[str, Any],
|
||||
requests_count: int,
|
||||
concurrency: int,
|
||||
timeout: float,
|
||||
) -> Dict[str, Any]:
|
||||
start = time.perf_counter()
|
||||
samples: List[Dict[str, Any]] = []
|
||||
with ThreadPoolExecutor(max_workers=concurrency) as pool:
|
||||
futures = [
|
||||
pool.submit(_post_once, url, payload, timeout)
|
||||
for _ in range(requests_count)
|
||||
]
|
||||
for future in as_completed(futures):
|
||||
samples.append(future.result())
|
||||
wall_seconds = time.perf_counter() - start
|
||||
latencies = [sample["elapsed_seconds"] for sample in samples]
|
||||
success_count = sum(1 for sample in samples if sample["ok"])
|
||||
return {
|
||||
"name": name,
|
||||
"url": url,
|
||||
"requests": requests_count,
|
||||
"concurrency": concurrency,
|
||||
"success": success_count,
|
||||
"failed": requests_count - success_count,
|
||||
"wall_seconds": wall_seconds,
|
||||
"throughput_rps": requests_count / wall_seconds if wall_seconds else 0.0,
|
||||
"latency_seconds": {
|
||||
"min": min(latencies) if latencies else 0.0,
|
||||
"mean": statistics.mean(latencies) if latencies else 0.0,
|
||||
"p50": _percentile(latencies, 50),
|
||||
"p95": _percentile(latencies, 95),
|
||||
"max": max(latencies) if latencies else 0.0,
|
||||
},
|
||||
"errors": [sample for sample in samples if not sample["ok"]][:10],
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Compare legacy and vLLM OCR endpoint concurrency.")
|
||||
parser.add_argument("--image", required=True, help="Path to an input png/jpg image.")
|
||||
parser.add_argument("--old-url", default=DEFAULT_OLD_ENDPOINT)
|
||||
parser.add_argument("--new-url", default=DEFAULT_NEW_ENDPOINT)
|
||||
parser.add_argument("--requests", type=int, default=20)
|
||||
parser.add_argument("--concurrency", type=int, default=8)
|
||||
parser.add_argument("--timeout", type=float, default=900)
|
||||
parser.add_argument("--output", default="concurrency_test_results.json")
|
||||
parser.add_argument("--skip-old", action="store_true")
|
||||
parser.add_argument("--skip-new", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
image_path = Path(args.image)
|
||||
image_type = image_path.suffix.lstrip(".").lower() or "png"
|
||||
payload = {
|
||||
"file": _encode_image(image_path),
|
||||
"type": "jpg" if image_type == "jpeg" else image_type,
|
||||
"skip_text_detection": False,
|
||||
"skip_table_detection": False,
|
||||
"recognize_math": False,
|
||||
"ocr_with_boxes": True,
|
||||
}
|
||||
|
||||
results = {
|
||||
"started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
|
||||
"image": str(image_path),
|
||||
"requests": args.requests,
|
||||
"concurrency": args.concurrency,
|
||||
"preflight": _preflight(),
|
||||
"endpoints": [],
|
||||
}
|
||||
if not args.skip_old:
|
||||
results["endpoints"].append(
|
||||
run_endpoint("legacy", args.old_url, payload, args.requests, args.concurrency, args.timeout)
|
||||
)
|
||||
if not args.skip_new:
|
||||
results["endpoints"].append(
|
||||
run_endpoint("vllm", args.new_url, payload, args.requests, args.concurrency, args.timeout)
|
||||
)
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.write_text(json.dumps(results, indent=2), encoding="utf-8")
|
||||
print(json.dumps(results, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user