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>
95 lines
2.3 KiB
Python
95 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import contextvars
|
|
import logging
|
|
import os
|
|
import threading
|
|
import time
|
|
from contextlib import contextmanager
|
|
from typing import Any, Iterator
|
|
|
|
|
|
_current_timing: contextvars.ContextVar["TimingCollector | None"] = contextvars.ContextVar(
|
|
"surya_current_timing",
|
|
default=None,
|
|
)
|
|
|
|
|
|
def timing_enabled() -> bool:
|
|
return os.getenv("SUYA_TIMING_ENABLED", "true").strip().lower() not in {
|
|
"0",
|
|
"false",
|
|
"no",
|
|
"off",
|
|
}
|
|
|
|
|
|
class TimingCollector:
|
|
def __init__(self, *, request_id: str = "-", batch_size: int = 0) -> None:
|
|
self.request_id = request_id
|
|
self.batch_size = batch_size
|
|
self._lock = threading.Lock()
|
|
self._events: list[dict[str, Any]] = []
|
|
|
|
def record(self, name: str, duration_ms: float, **metadata: Any) -> None:
|
|
if not timing_enabled():
|
|
return
|
|
event = {
|
|
"name": name,
|
|
"duration_ms": round(duration_ms, 2),
|
|
}
|
|
if metadata:
|
|
event["metadata"] = {
|
|
key: value
|
|
for key, value in metadata.items()
|
|
if value is not None
|
|
}
|
|
with self._lock:
|
|
self._events.append(event)
|
|
|
|
def summary(self) -> list[dict[str, Any]]:
|
|
with self._lock:
|
|
return list(self._events)
|
|
|
|
|
|
def get_current_timing() -> TimingCollector | None:
|
|
return _current_timing.get()
|
|
|
|
|
|
def set_current_timing(collector: TimingCollector | None):
|
|
return _current_timing.set(collector)
|
|
|
|
|
|
def reset_current_timing(token) -> None:
|
|
_current_timing.reset(token)
|
|
|
|
|
|
@contextmanager
|
|
def timing_span(name: str, **metadata: Any) -> Iterator[None]:
|
|
collector = get_current_timing()
|
|
if collector is None or not timing_enabled():
|
|
yield
|
|
return
|
|
start = time.perf_counter()
|
|
try:
|
|
yield
|
|
finally:
|
|
collector.record(name, (time.perf_counter() - start) * 1000, **metadata)
|
|
|
|
|
|
def log_timing_summary(
|
|
logger: logging.Logger,
|
|
collector: TimingCollector,
|
|
*,
|
|
message: str = "surya_timing_summary",
|
|
) -> None:
|
|
if not timing_enabled():
|
|
return
|
|
logger.info(
|
|
"%s request_id=%s batch_size=%s events=%s",
|
|
message,
|
|
collector.request_id,
|
|
collector.batch_size,
|
|
collector.summary(),
|
|
)
|