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>
126 lines
4.0 KiB
Python
126 lines
4.0 KiB
Python
import argparse
|
|
import csv
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from concurrency_test import DEFAULT_NEW_ENDPOINT, _encode_image, run_endpoint
|
|
|
|
|
|
def parse_levels(value: str) -> list[int]:
|
|
return [int(item.strip()) for item in value.split(",") if item.strip()]
|
|
|
|
|
|
def write_csv(path: Path, rows: list[dict]) -> None:
|
|
columns = [
|
|
"concurrency",
|
|
"requests",
|
|
"success",
|
|
"failed",
|
|
"wall_seconds",
|
|
"throughput_rps",
|
|
"latency_min",
|
|
"latency_mean",
|
|
"latency_p50",
|
|
"latency_p95",
|
|
"latency_max",
|
|
]
|
|
with path.open("w", newline="", encoding="utf-8") as f:
|
|
writer = csv.DictWriter(f, fieldnames=columns)
|
|
writer.writeheader()
|
|
for row in rows:
|
|
lat = row["latency_seconds"]
|
|
writer.writerow(
|
|
{
|
|
"concurrency": row["concurrency"],
|
|
"requests": row["requests"],
|
|
"success": row["success"],
|
|
"failed": row["failed"],
|
|
"wall_seconds": row["wall_seconds"],
|
|
"throughput_rps": row["throughput_rps"],
|
|
"latency_min": lat["min"],
|
|
"latency_mean": lat["mean"],
|
|
"latency_p50": lat["p50"],
|
|
"latency_p95": lat["p95"],
|
|
"latency_max": lat["max"],
|
|
}
|
|
)
|
|
|
|
|
|
def write_plot(path: Path, rows: list[dict]) -> None:
|
|
import matplotlib
|
|
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
|
|
x = [row["concurrency"] for row in rows]
|
|
mean = [row["latency_seconds"]["mean"] for row in rows]
|
|
p50 = [row["latency_seconds"]["p50"] for row in rows]
|
|
p95 = [row["latency_seconds"]["p95"] for row in rows]
|
|
|
|
plt.figure(figsize=(10, 6))
|
|
plt.plot(x, mean, marker="o", label="mean")
|
|
plt.plot(x, p50, marker="o", label="p50")
|
|
plt.plot(x, p95, marker="o", label="p95")
|
|
plt.xlabel("Concurrency")
|
|
plt.ylabel("Latency per request (seconds)")
|
|
plt.title("vLLM OCR latency vs concurrency")
|
|
plt.grid(True, alpha=0.3)
|
|
plt.legend()
|
|
plt.tight_layout()
|
|
plt.savefig(path, dpi=160)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Sweep vLLM OCR endpoint concurrency levels.")
|
|
parser.add_argument("--image", required=True)
|
|
parser.add_argument("--url", default=DEFAULT_NEW_ENDPOINT)
|
|
parser.add_argument("--levels", default="1,2,4,6,8,10,12,20,30,40,50")
|
|
parser.add_argument("--timeout", type=float, default=900)
|
|
parser.add_argument("--output-json", default="concurrency_sweep_results.json")
|
|
parser.add_argument("--output-csv", default="concurrency_sweep_results.csv")
|
|
parser.add_argument("--output-plot", default="concurrency_sweep_latency.png")
|
|
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,
|
|
}
|
|
|
|
levels = parse_levels(args.levels)
|
|
rows = []
|
|
for concurrency in levels:
|
|
print(f"running concurrency={concurrency}", flush=True)
|
|
rows.append(
|
|
run_endpoint(
|
|
"vllm",
|
|
args.url,
|
|
payload,
|
|
requests_count=concurrency,
|
|
concurrency=concurrency,
|
|
timeout=args.timeout,
|
|
)
|
|
)
|
|
|
|
result = {
|
|
"started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
|
|
"image": str(image_path),
|
|
"url": args.url,
|
|
"levels": levels,
|
|
"results": rows,
|
|
}
|
|
Path(args.output_json).write_text(json.dumps(result, indent=2), encoding="utf-8")
|
|
write_csv(Path(args.output_csv), rows)
|
|
write_plot(Path(args.output_plot), rows)
|
|
print(json.dumps(result, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|