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,116 @@
|
||||
"""Surya inference manager.
|
||||
|
||||
One process owns one SuryaInferenceManager. The manager wraps a single backend
|
||||
(vllm | llamacpp) which speaks OpenAI-compatible chat completions.
|
||||
|
||||
Predictors take the manager via explicit injection at construction time.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from typing import List, Optional
|
||||
|
||||
from surya.inference.backends.base import Backend
|
||||
from surya.inference.schema import BatchInputItem, BatchOutputItem
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
from surya.timing import timing_span
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def _has_nvidia_gpu() -> bool:
|
||||
"""True if an NVIDIA GPU is present on this host.
|
||||
|
||||
We deliberately do *not* rely solely on ``torch.cuda.is_available()``:
|
||||
the installed torch wheel's CUDA build can be newer than the host driver
|
||||
(PyPI's default wheel tracks the latest CUDA), in which case torch reports
|
||||
no CUDA even on a perfectly good GPU box. That would silently route us to
|
||||
the CPU llama.cpp backend on a machine that should be running vllm. So we
|
||||
take torch's word when it *does* see CUDA, and otherwise fall back to
|
||||
probing for the GPU directly via ``nvidia-smi``.
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Instant, load-independent check: the NVIDIA device node only exists when
|
||||
# a GPU + driver are present. Preferred over nvidia-smi because nvidia-smi
|
||||
# can block for several seconds on a GPU under heavy load, which would race
|
||||
# a timeout and falsely report "no GPU".
|
||||
if os.path.exists("/dev/nvidia0"):
|
||||
return True
|
||||
|
||||
nvidia_smi = shutil.which("nvidia-smi")
|
||||
if not nvidia_smi:
|
||||
return False
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[nvidia_smi, "-L"], capture_output=True, text=True, timeout=15
|
||||
)
|
||||
return result.returncode == 0 and "GPU" in result.stdout
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _autodetect_backend() -> str:
|
||||
if settings.SURYA_INFERENCE_BACKEND:
|
||||
return settings.SURYA_INFERENCE_BACKEND
|
||||
# NVIDIA GPU → vllm, mps/cpu → llamacpp
|
||||
if _has_nvidia_gpu():
|
||||
return "vllm"
|
||||
return "llamacpp"
|
||||
|
||||
|
||||
def _build_backend(method: str) -> Backend:
|
||||
method = method.lower()
|
||||
if method == "vllm":
|
||||
from surya.inference.backends.vllm import VllmBackend
|
||||
|
||||
return VllmBackend()
|
||||
if method == "llamacpp":
|
||||
from surya.inference.backends.llamacpp import LlamaCppBackend
|
||||
return LlamaCppBackend()
|
||||
raise ValueError(
|
||||
f"Unknown inference backend {method!r}. Supported: 'vllm', 'llamacpp'."
|
||||
)
|
||||
|
||||
|
||||
class SuryaInferenceManager:
|
||||
"""Single entry point for VLM inference. Construct once per process."""
|
||||
|
||||
def __init__(self, method: Optional[str] = None, lazy: bool = True):
|
||||
self.method = method or _autodetect_backend()
|
||||
self.backend: Backend = _build_backend(self.method)
|
||||
if not lazy:
|
||||
self.backend.start()
|
||||
|
||||
def start(self) -> None:
|
||||
self.backend.start()
|
||||
|
||||
def stop(self) -> None:
|
||||
self.backend.stop()
|
||||
|
||||
def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
|
||||
with timing_span("surya_manager_generate", backend=self.method, item_count=len(batch)):
|
||||
return self.backend.generate(batch)
|
||||
|
||||
|
||||
# Module-level lazy singleton for callers that don't want explicit construction
|
||||
# (notebooks, ad-hoc scripts). Surya's own models.py and marker should use
|
||||
# explicit construction.
|
||||
_default_manager: Optional[SuryaInferenceManager] = None
|
||||
|
||||
|
||||
def get_default_manager() -> SuryaInferenceManager:
|
||||
global _default_manager
|
||||
if _default_manager is None:
|
||||
_default_manager = SuryaInferenceManager()
|
||||
return _default_manager
|
||||
@@ -0,0 +1,2 @@
|
||||
from surya.inference.backends.base import Backend as Backend
|
||||
from surya.inference.backends.base import ServerHandle as ServerHandle
|
||||
@@ -0,0 +1,31 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import List
|
||||
|
||||
|
||||
from surya.inference.schema import BatchInputItem, BatchOutputItem
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServerHandle:
|
||||
base_url: str # e.g. "http://127.0.0.1:8765/v1"
|
||||
model_name: str # what gets passed in OpenAI `model` field
|
||||
spawned_by_us: bool # if True, we manage atexit cleanup
|
||||
|
||||
|
||||
class Backend:
|
||||
"""Abstract backend. Concrete backends own server lifecycle + generation."""
|
||||
|
||||
name: str # "vllm" | "llamacpp"
|
||||
|
||||
def start(self) -> ServerHandle:
|
||||
"""Idempotent: probe → attach if alive, else spawn. Returns handle."""
|
||||
raise NotImplementedError
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop the server if we spawned it."""
|
||||
raise NotImplementedError
|
||||
|
||||
def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,207 @@
|
||||
"""llama.cpp backend: spawns the upstream `llama-server` binary natively.
|
||||
|
||||
Install:
|
||||
- macOS: brew install llama.cpp (Metal build, MPS)
|
||||
- Linux: brew install llama.cpp OR github.com/ggml-org/llama.cpp/releases
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
from openai import OpenAI
|
||||
|
||||
from surya.inference.backends.base import Backend, ServerHandle
|
||||
from surya.inference.backends.openai_client import chat_completions_batch
|
||||
from surya.inference.backends.spawn import (
|
||||
SpawnHandle,
|
||||
SpawnError,
|
||||
attach_or_spawn,
|
||||
)
|
||||
from surya.inference.schema import BatchInputItem, BatchOutputItem
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def _resolve_llama_server_binary() -> str:
|
||||
binary = settings.LLAMA_CPP_BINARY
|
||||
if binary and os.path.isfile(binary):
|
||||
return binary
|
||||
found = shutil.which(binary or "llama-server")
|
||||
if found:
|
||||
return found
|
||||
raise SpawnError(
|
||||
"llama-server binary not found. Install with:\n"
|
||||
" macOS: brew install llama.cpp\n"
|
||||
" Linux: brew install llama.cpp OR download from\n"
|
||||
" https://github.com/ggml-org/llama.cpp/releases\n"
|
||||
"Or set LLAMA_CPP_BINARY in your env to the binary path."
|
||||
)
|
||||
|
||||
|
||||
def _download_gguf_files() -> tuple[str, str]:
|
||||
"""Download model + mmproj GGUFs from HF Hub. Returns local paths."""
|
||||
repo = settings.SURYA_GGUF_REPO
|
||||
model_file = settings.SURYA_GGUF_MODEL_FILE
|
||||
mmproj_file = settings.SURYA_GGUF_MMPROJ_FILE
|
||||
logger.info(f"Downloading {model_file} and {mmproj_file} from {repo}")
|
||||
model_path = hf_hub_download(repo_id=repo, filename=model_file)
|
||||
mmproj_path = hf_hub_download(repo_id=repo, filename=mmproj_file)
|
||||
return model_path, mmproj_path
|
||||
|
||||
|
||||
def _health_url(port: int) -> str:
|
||||
return f"http://{settings.SURYA_INFERENCE_HOST}:{port}"
|
||||
|
||||
|
||||
def _openai_url(port: int) -> str:
|
||||
return f"http://{settings.SURYA_INFERENCE_HOST}:{port}/v1"
|
||||
|
||||
|
||||
class LlamaCppBackend(Backend):
|
||||
name = "llamacpp"
|
||||
|
||||
def __init__(self):
|
||||
self.handle: Optional[ServerHandle] = None
|
||||
self._client: Optional[OpenAI] = None
|
||||
|
||||
def start(self) -> ServerHandle:
|
||||
if self.handle is not None:
|
||||
return self.handle
|
||||
|
||||
# If user pinned an external server, attach without spawning.
|
||||
# No binary or GGUF download needed in that case.
|
||||
if settings.SURYA_INFERENCE_URL:
|
||||
spawned = attach_or_spawn(
|
||||
backend=self.name,
|
||||
expected_model_name=settings.SURYA_MODEL_CHECKPOINT,
|
||||
spawn_fn=lambda port: SpawnHandle(
|
||||
pid=None, cleanup_id="", cleanup_kind="process"
|
||||
), # never called
|
||||
health_url_for=_health_url,
|
||||
openai_url_for=_openai_url,
|
||||
startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT,
|
||||
)
|
||||
self.handle = ServerHandle(
|
||||
base_url=spawned.base_url,
|
||||
model_name=spawned.model_name,
|
||||
spawned_by_us=spawned.spawned_by_us,
|
||||
)
|
||||
self._client = OpenAI(api_key="EMPTY", base_url=self.handle.base_url)
|
||||
return self.handle
|
||||
|
||||
binary = _resolve_llama_server_binary()
|
||||
|
||||
# Pre-download GGUFs so the spawn doesn't race the download
|
||||
if (
|
||||
settings.SURYA_GGUF_LOCAL_MODEL_PATH
|
||||
and settings.SURYA_GGUF_LOCAL_MMPROJ_PATH
|
||||
):
|
||||
model_path = settings.SURYA_GGUF_LOCAL_MODEL_PATH
|
||||
mmproj_path = settings.SURYA_GGUF_LOCAL_MMPROJ_PATH
|
||||
else:
|
||||
model_path, mmproj_path = _download_gguf_files()
|
||||
|
||||
# Total KV-cache budget. llama-server divides --ctx-size across
|
||||
# --parallel slots, so a too-small total silently truncates outputs
|
||||
# once each slot's share fills. Scale with parallel by default;
|
||||
# SURYA_INFERENCE_CTX_SIZE overrides to a fixed value if set.
|
||||
parallel = settings.SURYA_INFERENCE_PARALLEL
|
||||
per_slot = settings.SURYA_INFERENCE_CTX_PER_SLOT
|
||||
ctx_size = settings.SURYA_INFERENCE_CTX_SIZE
|
||||
if ctx_size is None:
|
||||
ctx_size = max(16384, parallel * per_slot)
|
||||
effective_per_slot = ctx_size // max(parallel, 1)
|
||||
logger.info(
|
||||
f"llama-server ctx-size={ctx_size} "
|
||||
f"(~{effective_per_slot}/slot × {parallel} parallel slots)"
|
||||
)
|
||||
if effective_per_slot < per_slot:
|
||||
logger.warning(
|
||||
f"per-slot ctx ({effective_per_slot}) is below recommended "
|
||||
f"{per_slot}; outputs may truncate. Raise "
|
||||
f"SURYA_INFERENCE_CTX_SIZE or SURYA_INFERENCE_CTX_PER_SLOT, "
|
||||
f"or lower SURYA_INFERENCE_PARALLEL."
|
||||
)
|
||||
|
||||
def spawn_fn(port: int) -> SpawnHandle:
|
||||
cmd = [
|
||||
binary,
|
||||
"-m",
|
||||
model_path,
|
||||
"--mmproj",
|
||||
mmproj_path,
|
||||
"-ngl",
|
||||
str(settings.LLAMA_CPP_NGL),
|
||||
"--host",
|
||||
settings.SURYA_INFERENCE_HOST,
|
||||
"--port",
|
||||
str(port),
|
||||
"--parallel",
|
||||
str(parallel),
|
||||
"--ctx-size",
|
||||
str(ctx_size),
|
||||
"--no-mmproj-offload" if settings.LLAMA_CPP_NO_MMPROJ_OFFLOAD else "",
|
||||
"--alias",
|
||||
settings.SURYA_MODEL_CHECKPOINT,
|
||||
"--jinja",
|
||||
]
|
||||
cmd = [c for c in cmd if c]
|
||||
for extra in (settings.LLAMA_CPP_EXTRA_ARGS or "").split():
|
||||
cmd.append(extra)
|
||||
logger.info(f"Spawning: {' '.join(cmd)}")
|
||||
log_path = Path("~/.cache/datalab/surya/llamacpp_server.log").expanduser()
|
||||
log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
log_fp = open(log_path, "ab")
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=log_fp,
|
||||
stderr=subprocess.STDOUT,
|
||||
start_new_session=True,
|
||||
)
|
||||
return SpawnHandle(
|
||||
pid=proc.pid, cleanup_id=str(proc.pid), cleanup_kind="process"
|
||||
)
|
||||
|
||||
spawned = attach_or_spawn(
|
||||
backend=self.name,
|
||||
expected_model_name=settings.SURYA_MODEL_CHECKPOINT,
|
||||
spawn_fn=spawn_fn,
|
||||
health_url_for=_health_url,
|
||||
openai_url_for=_openai_url,
|
||||
startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT,
|
||||
)
|
||||
self.handle = ServerHandle(
|
||||
base_url=spawned.base_url,
|
||||
model_name=spawned.model_name,
|
||||
spawned_by_us=spawned.spawned_by_us,
|
||||
)
|
||||
self._client = OpenAI(
|
||||
api_key="EMPTY",
|
||||
base_url=self.handle.base_url,
|
||||
)
|
||||
return self.handle
|
||||
|
||||
def stop(self) -> None:
|
||||
# atexit handler in spawn.py owns cleanup; nothing to do here.
|
||||
self.handle = None
|
||||
self._client = None
|
||||
|
||||
def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
|
||||
if self.handle is None or self._client is None:
|
||||
self.start()
|
||||
return chat_completions_batch(
|
||||
batch,
|
||||
client=self._client,
|
||||
model_name=self.handle.model_name,
|
||||
timeout=settings.SURYA_INFERENCE_TIMEOUT_SECONDS,
|
||||
max_workers=settings.SURYA_INFERENCE_PARALLEL,
|
||||
request_logprobs_default=settings.SURYA_INFERENCE_LOGPROBS,
|
||||
)
|
||||
@@ -0,0 +1,257 @@
|
||||
"""Shared OpenAI-compatible chat completions client. Used by vllm + llama.cpp.
|
||||
|
||||
Both servers expose `/v1/chat/completions` with the same request/response shape,
|
||||
so this module is the single point of HTTP contact for both backends.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List, Optional
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from surya.inference.prompts import PROMPT_MAPPING
|
||||
from surya.inference.schema import (
|
||||
BatchInputItem,
|
||||
BatchOutputItem,
|
||||
GenerationResult,
|
||||
)
|
||||
from surya.inference.util import detect_repeat_token, scale_to_fit
|
||||
from surya.logging import get_logger
|
||||
from surya.timing import get_current_timing, set_current_timing, timing_span
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def resolve_max_workers(batch_len: int, max_inflight: int) -> int:
|
||||
"""Concurrent HTTP workers for a batch: as many as the batch needs, capped
|
||||
by max_inflight so we keep vLLM's sequence slots full without over-queueing."""
|
||||
return max(1, min(batch_len, max_inflight))
|
||||
|
||||
|
||||
def encode_image_b64(image: Image.Image) -> tuple[str, str]:
|
||||
image_format = os.getenv("SUYA_VLLM_IMAGE_FORMAT", "JPEG").upper()
|
||||
if image_format not in {"JPEG", "PNG"}:
|
||||
raise ValueError("SUYA_VLLM_IMAGE_FORMAT must be JPEG or PNG")
|
||||
|
||||
buf = io.BytesIO()
|
||||
if image_format == "JPEG":
|
||||
quality = int(os.getenv("SUYA_VLLM_JPEG_QUALITY", "92"))
|
||||
image.save(buf, format="JPEG", quality=quality, subsampling=0)
|
||||
mime_type = "image/jpeg"
|
||||
else:
|
||||
image.save(buf, format="PNG")
|
||||
mime_type = "image/png"
|
||||
view = buf.getbuffer()
|
||||
try:
|
||||
return base64.b64encode(view).decode("ascii"), mime_type
|
||||
finally:
|
||||
view.release()
|
||||
|
||||
|
||||
def _build_messages(image: Image.Image, prompt: str):
|
||||
with timing_span("openai_encode_image_b64", image_size=image.size):
|
||||
image_b64, mime_type = encode_image_b64(image)
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:{mime_type};base64,{image_b64}"},
|
||||
},
|
||||
{"type": "text", "text": prompt},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _mean_token_prob(logprobs_content) -> Optional[float]:
|
||||
if not logprobs_content:
|
||||
return None
|
||||
probs = []
|
||||
for tok in logprobs_content:
|
||||
lp = (
|
||||
tok.get("logprob")
|
||||
if isinstance(tok, dict)
|
||||
else getattr(tok, "logprob", None)
|
||||
)
|
||||
if lp is None:
|
||||
continue
|
||||
probs.append(math.exp(lp))
|
||||
if not probs:
|
||||
return None
|
||||
return sum(probs) / len(probs)
|
||||
|
||||
|
||||
def _generate_one(
|
||||
item: BatchInputItem,
|
||||
client,
|
||||
model_name: str,
|
||||
max_tokens_default: int,
|
||||
temperature: float,
|
||||
top_p: float,
|
||||
timeout: float,
|
||||
request_logprobs_default: bool,
|
||||
) -> GenerationResult:
|
||||
prompt = item.prompt or PROMPT_MAPPING[item.prompt_type]
|
||||
with timing_span("openai_scale_image", prompt_type=item.prompt_type, image_size=item.image.size):
|
||||
image = scale_to_fit(item.image)
|
||||
with timing_span("openai_build_messages", prompt_type=item.prompt_type):
|
||||
messages = _build_messages(image, prompt)
|
||||
|
||||
max_tokens = item.max_tokens or max_tokens_default
|
||||
request_logprobs = item.request_logprobs or request_logprobs_default
|
||||
|
||||
kwargs = dict(
|
||||
model=model_name,
|
||||
messages=messages,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
timeout=timeout,
|
||||
)
|
||||
if request_logprobs:
|
||||
kwargs["logprobs"] = True
|
||||
|
||||
# Structured output: prefer OpenAI-standard response_format (works on both
|
||||
# vllm and llama.cpp). Fall back to vllm's extra_body for guided_regex.
|
||||
if item.guided_json is not None:
|
||||
kwargs["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "structured_output",
|
||||
"schema": item.guided_json,
|
||||
"strict": True,
|
||||
},
|
||||
}
|
||||
if item.guided_regex is not None:
|
||||
kwargs.setdefault("extra_body", {})["guided_regex"] = item.guided_regex
|
||||
|
||||
try:
|
||||
with timing_span(
|
||||
"openai_chat_completion",
|
||||
prompt_type=item.prompt_type,
|
||||
max_tokens=max_tokens,
|
||||
):
|
||||
completion = client.chat.completions.create(**kwargs)
|
||||
raw = completion.choices[0].message.content or ""
|
||||
token_count = completion.usage.completion_tokens if completion.usage else 0
|
||||
with timing_span(
|
||||
"openai_parse_completion",
|
||||
prompt_type=item.prompt_type,
|
||||
token_count=token_count,
|
||||
):
|
||||
mean_p = None
|
||||
logprobs_content = None
|
||||
if request_logprobs:
|
||||
choice = completion.choices[0]
|
||||
lp = getattr(choice, "logprobs", None)
|
||||
if lp is not None:
|
||||
content = getattr(lp, "content", None)
|
||||
if content is not None:
|
||||
logprobs_content = [
|
||||
c.model_dump() if hasattr(c, "model_dump") else c
|
||||
for c in content
|
||||
]
|
||||
mean_p = _mean_token_prob(content)
|
||||
return GenerationResult(
|
||||
raw=raw,
|
||||
token_count=token_count,
|
||||
error=False,
|
||||
mean_token_prob=mean_p,
|
||||
logprobs=logprobs_content,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Inference error: {e}")
|
||||
return GenerationResult(raw="", token_count=0, error=True)
|
||||
|
||||
|
||||
def _should_retry(
|
||||
result: GenerationResult,
|
||||
retries: int,
|
||||
max_retries: int,
|
||||
) -> bool:
|
||||
if retries >= max_retries:
|
||||
return False
|
||||
if result.error:
|
||||
return True
|
||||
has_repeat = detect_repeat_token(result.raw) or (
|
||||
len(result.raw) > 50 and detect_repeat_token(result.raw, cut_from_end=50)
|
||||
)
|
||||
return has_repeat
|
||||
|
||||
|
||||
def chat_completions_batch(
|
||||
batch: List[BatchInputItem],
|
||||
client,
|
||||
model_name: str,
|
||||
max_tokens_default: int = 2048,
|
||||
temperature: float = 0.0,
|
||||
top_p: float = 0.1,
|
||||
timeout: float = 600.0,
|
||||
max_workers: Optional[int] = None,
|
||||
max_retries: int = 3,
|
||||
request_logprobs_default: bool = True,
|
||||
) -> List[BatchOutputItem]:
|
||||
"""Run a batch of items through the chat completions endpoint with concurrent workers."""
|
||||
if not batch:
|
||||
return []
|
||||
if max_workers is None:
|
||||
max_workers = min(64, len(batch))
|
||||
collector = get_current_timing()
|
||||
|
||||
def _process(item: BatchInputItem) -> BatchOutputItem:
|
||||
if collector is not None:
|
||||
set_current_timing(collector)
|
||||
result = _generate_one(
|
||||
item,
|
||||
client=client,
|
||||
model_name=model_name,
|
||||
max_tokens_default=max_tokens_default,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
timeout=timeout,
|
||||
request_logprobs_default=request_logprobs_default,
|
||||
)
|
||||
retries = 0
|
||||
while _should_retry(result, retries, max_retries):
|
||||
backoff = 1.5 * (retries + 1) if result.error else 0
|
||||
if backoff:
|
||||
time.sleep(backoff)
|
||||
retry_temp = min(temperature + 0.2 * (retries + 1), 0.8)
|
||||
retry_top_p = 0.95 if not result.error else top_p
|
||||
result = _generate_one(
|
||||
item,
|
||||
client=client,
|
||||
model_name=model_name,
|
||||
max_tokens_default=max_tokens_default,
|
||||
temperature=retry_temp,
|
||||
top_p=retry_top_p,
|
||||
timeout=timeout,
|
||||
request_logprobs_default=request_logprobs_default,
|
||||
)
|
||||
retries += 1
|
||||
return BatchOutputItem(
|
||||
raw=result.raw,
|
||||
token_count=result.token_count,
|
||||
error=result.error,
|
||||
mean_token_prob=result.mean_token_prob,
|
||||
logprobs=result.logprobs,
|
||||
metadata=item.metadata,
|
||||
)
|
||||
|
||||
with timing_span(
|
||||
"openai_batch_threadpool",
|
||||
item_count=len(batch),
|
||||
max_workers=max_workers,
|
||||
max_retries=max_retries,
|
||||
):
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
return list(executor.map(_process, batch))
|
||||
@@ -0,0 +1,351 @@
|
||||
"""Server lifecycle: probe, filelock, sentinel, atexit cleanup.
|
||||
|
||||
Pattern: probe `/health` → if alive return handle → else acquire lock, re-probe,
|
||||
spawn detached, write sentinel, register atexit kill (only the spawner cleans up).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def _cache_dir() -> Path:
|
||||
base = Path(os.path.expanduser("~/.cache/datalab/surya"))
|
||||
base.mkdir(parents=True, exist_ok=True)
|
||||
return base
|
||||
|
||||
|
||||
def _sentinel_path(backend: str) -> Path:
|
||||
return _cache_dir() / f"{backend}_server.json"
|
||||
|
||||
|
||||
def _lock_path(backend: str) -> Path:
|
||||
return _cache_dir() / f"{backend}_server.lock"
|
||||
|
||||
|
||||
def find_free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def probe_health(base_url: str, timeout: float = 1.0) -> bool:
|
||||
"""Returns True if the server reports healthy at /health."""
|
||||
try:
|
||||
# llama.cpp returns 200 on /health when ready; vllm returns 200 on /health too.
|
||||
with httpx.Client(timeout=timeout) as client:
|
||||
r = client.get(f"{base_url}/health")
|
||||
return r.status_code == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def wait_for_health(
|
||||
base_url: str, total_timeout: float = 300.0, interval: float = 1.0
|
||||
) -> bool:
|
||||
deadline = time.time() + total_timeout
|
||||
while time.time() < deadline:
|
||||
if probe_health(base_url):
|
||||
return True
|
||||
time.sleep(interval)
|
||||
return False
|
||||
|
||||
|
||||
def probe_model_id(openai_base: str, timeout: float = 5.0) -> Optional[str]:
|
||||
"""Returns the model id reported by the running server, or None on failure."""
|
||||
try:
|
||||
with httpx.Client(timeout=timeout) as client:
|
||||
r = client.get(f"{openai_base}/models")
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
models = data.get("data") or []
|
||||
if models:
|
||||
return models[0].get("id")
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class SpawnedServer:
|
||||
base_url: str # full openai base, e.g. "http://127.0.0.1:8765/v1"
|
||||
health_url: str # base for /health, e.g. "http://127.0.0.1:8765"
|
||||
model_name: str # what to pass as `model`
|
||||
pid: Optional[int]
|
||||
backend: str
|
||||
spawned_by_us: bool
|
||||
|
||||
|
||||
class SpawnError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def _read_sentinel(backend: str) -> Optional[dict]:
|
||||
p = _sentinel_path(backend)
|
||||
if not p.exists():
|
||||
return None
|
||||
try:
|
||||
return json.loads(p.read_text())
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _write_sentinel(backend: str, data: dict) -> None:
|
||||
_sentinel_path(backend).write_text(json.dumps(data))
|
||||
|
||||
|
||||
def _delete_sentinel(backend: str) -> None:
|
||||
p = _sentinel_path(backend)
|
||||
if p.exists():
|
||||
try:
|
||||
p.unlink()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _stop_process(pid: int, name: str) -> None:
|
||||
try:
|
||||
# Graceful first
|
||||
os.kill(pid, 15) # SIGTERM
|
||||
for _ in range(20):
|
||||
try:
|
||||
os.kill(pid, 0) # still alive?
|
||||
except ProcessLookupError:
|
||||
logger.info(f"Stopped {name} (pid {pid})")
|
||||
return
|
||||
time.sleep(0.5)
|
||||
# Hard
|
||||
os.kill(pid, 9)
|
||||
logger.warning(f"Force-killed {name} (pid {pid})")
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to stop {name} (pid {pid}): {e}")
|
||||
|
||||
|
||||
def _capture_server_logs(handle: "SpawnHandle", tail: int = 100) -> str:
|
||||
"""Best-effort tail of a server's logs, for surfacing startup failures."""
|
||||
try:
|
||||
if handle.cleanup_kind == "docker":
|
||||
r = subprocess.run(
|
||||
["docker", "logs", "--tail", str(tail), handle.cleanup_id],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=15,
|
||||
)
|
||||
return (r.stdout or "") + (r.stderr or "") or "(no docker logs)"
|
||||
# llama.cpp process backend logs to this file (see llamacpp.py)
|
||||
log_path = Path("~/.cache/datalab/surya/llamacpp_server.log").expanduser()
|
||||
if log_path.exists():
|
||||
lines = log_path.read_text(errors="replace").splitlines()
|
||||
return "\n".join(lines[-tail:]) or "(empty log)"
|
||||
except Exception as e:
|
||||
return f"(could not capture logs: {e})"
|
||||
return "(no logs available)"
|
||||
|
||||
|
||||
def _stop_docker_container(name: str) -> None:
|
||||
try:
|
||||
subprocess.run(
|
||||
["docker", "stop", name], check=False, capture_output=True, timeout=30
|
||||
)
|
||||
logger.info(f"Stopped docker container {name}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to stop docker container {name}: {e}")
|
||||
|
||||
|
||||
def attach_or_spawn(
|
||||
backend: str,
|
||||
expected_model_name: str,
|
||||
spawn_fn: Callable[[int], "SpawnHandle"],
|
||||
health_url_for: Callable[[int], str],
|
||||
openai_url_for: Callable[[int], str],
|
||||
startup_timeout: float = 600.0,
|
||||
) -> SpawnedServer:
|
||||
"""Generic attach-or-spawn with file lock and sentinel.
|
||||
|
||||
`spawn_fn(port)` must launch the server detached and return a SpawnHandle
|
||||
with `pid` (int or None for docker) and a `cleanup_id` (e.g. container name).
|
||||
"""
|
||||
# 0. If user pinned an external URL, attach without lock
|
||||
if settings.SURYA_INFERENCE_URL:
|
||||
base_url = settings.SURYA_INFERENCE_URL.rstrip("/")
|
||||
health_url = base_url[: -len("/v1")] if base_url.endswith("/v1") else base_url
|
||||
if not probe_health(health_url):
|
||||
raise SpawnError(
|
||||
f"SURYA_INFERENCE_URL={base_url} is not reachable at /health. "
|
||||
"Start the server or unset the variable."
|
||||
)
|
||||
model_name = probe_model_id(base_url) or expected_model_name
|
||||
if model_name != expected_model_name:
|
||||
raise SpawnError(
|
||||
f"Model mismatch at {base_url}: expected {expected_model_name!r}, got {model_name!r}. "
|
||||
"Stop the running server or unset SURYA_INFERENCE_URL."
|
||||
)
|
||||
return SpawnedServer(
|
||||
base_url=base_url,
|
||||
health_url=health_url,
|
||||
model_name=model_name,
|
||||
pid=None,
|
||||
backend=backend,
|
||||
spawned_by_us=False,
|
||||
)
|
||||
|
||||
# 1. Probe sentinel without lock
|
||||
existing = _read_sentinel(backend)
|
||||
if existing:
|
||||
port = existing.get("port")
|
||||
pid = existing.get("pid")
|
||||
if port and probe_health(health_url_for(port)):
|
||||
running_model = probe_model_id(openai_url_for(port)) or expected_model_name
|
||||
if running_model != expected_model_name:
|
||||
raise SpawnError(
|
||||
f"Existing {backend} server on port {port} serves {running_model!r}, "
|
||||
f"expected {expected_model_name!r}. Stop it before continuing."
|
||||
)
|
||||
logger.info(f"Attaching to existing {backend} server on port {port}")
|
||||
return SpawnedServer(
|
||||
base_url=openai_url_for(port),
|
||||
health_url=health_url_for(port),
|
||||
model_name=running_model,
|
||||
pid=pid,
|
||||
backend=backend,
|
||||
spawned_by_us=False,
|
||||
)
|
||||
else:
|
||||
_delete_sentinel(backend)
|
||||
|
||||
if not settings.SURYA_INFERENCE_AUTOSTART:
|
||||
raise SpawnError(
|
||||
f"No running {backend} server and SURYA_INFERENCE_AUTOSTART is False. "
|
||||
"Set the variable to True or start the server manually."
|
||||
)
|
||||
|
||||
# 2. Acquire filelock to prevent races
|
||||
try:
|
||||
from filelock import FileLock
|
||||
except ImportError as e:
|
||||
raise SpawnError(
|
||||
"filelock is required for server spawn. pip install filelock"
|
||||
) from e
|
||||
|
||||
lock = FileLock(str(_lock_path(backend)), timeout=120)
|
||||
with lock:
|
||||
# Re-check sentinel inside the lock
|
||||
existing = _read_sentinel(backend)
|
||||
if existing:
|
||||
port = existing.get("port")
|
||||
if port and probe_health(health_url_for(port)):
|
||||
running_model = (
|
||||
probe_model_id(openai_url_for(port)) or expected_model_name
|
||||
)
|
||||
if running_model != expected_model_name:
|
||||
raise SpawnError(
|
||||
f"Existing {backend} server on port {port} serves {running_model!r}, "
|
||||
f"expected {expected_model_name!r}."
|
||||
)
|
||||
return SpawnedServer(
|
||||
base_url=openai_url_for(port),
|
||||
health_url=health_url_for(port),
|
||||
model_name=running_model,
|
||||
pid=existing.get("pid"),
|
||||
backend=backend,
|
||||
spawned_by_us=False,
|
||||
)
|
||||
|
||||
# 3. Spawn fresh
|
||||
port = settings.SURYA_INFERENCE_PORT or find_free_port()
|
||||
logger.info(f"Spawning {backend} server on port {port}")
|
||||
spawn_handle = spawn_fn(port)
|
||||
|
||||
# 4. Write sentinel
|
||||
_write_sentinel(
|
||||
backend,
|
||||
{
|
||||
"port": port,
|
||||
"pid": spawn_handle.pid,
|
||||
"model": expected_model_name,
|
||||
"backend": backend,
|
||||
"cleanup_id": spawn_handle.cleanup_id,
|
||||
"cleanup_kind": spawn_handle.cleanup_kind,
|
||||
},
|
||||
)
|
||||
|
||||
# 5. Register atexit cleanup (only spawner). Skipped when keep-alive is
|
||||
# set so the server outlives this process and later commands attach to
|
||||
# it via the sentinel. (_cleanup is still callable below on startup
|
||||
# failure, where we always tear a half-started server down.)
|
||||
def _cleanup():
|
||||
try:
|
||||
if spawn_handle.cleanup_kind == "docker":
|
||||
_stop_docker_container(spawn_handle.cleanup_id)
|
||||
elif spawn_handle.cleanup_kind == "process":
|
||||
if spawn_handle.pid:
|
||||
_stop_process(spawn_handle.pid, backend)
|
||||
finally:
|
||||
_delete_sentinel(backend)
|
||||
|
||||
if settings.SURYA_INFERENCE_KEEP_ALIVE:
|
||||
logger.info(
|
||||
f"keep-alive: {backend} server on port {port} will stay up "
|
||||
f"after exit (cleanup_id={spawn_handle.cleanup_id!r})"
|
||||
)
|
||||
else:
|
||||
atexit.register(_cleanup)
|
||||
|
||||
# 6. Wait for health
|
||||
health_url = health_url_for(port)
|
||||
if not wait_for_health(health_url, total_timeout=startup_timeout):
|
||||
# Grab the server's own logs *before* cleanup tears the (--rm)
|
||||
# container down, otherwise the actual failure reason is lost and
|
||||
# all the caller sees is this timeout.
|
||||
logs = _capture_server_logs(spawn_handle)
|
||||
_cleanup()
|
||||
raise SpawnError(
|
||||
f"{backend} server failed to become healthy at {health_url} "
|
||||
f"within {startup_timeout}s.\n"
|
||||
f"--- last {backend} server logs ---\n{logs}"
|
||||
)
|
||||
|
||||
# 7. Verify model name
|
||||
running_model = probe_model_id(openai_url_for(port))
|
||||
if running_model and running_model != expected_model_name:
|
||||
logger.warning(
|
||||
f"{backend} server reports model={running_model!r} "
|
||||
f"but expected {expected_model_name!r}; using reported name."
|
||||
)
|
||||
expected_model_name = running_model
|
||||
|
||||
logger.info(
|
||||
f"{backend} server ready on port {port} (model={expected_model_name})"
|
||||
)
|
||||
return SpawnedServer(
|
||||
base_url=openai_url_for(port),
|
||||
health_url=health_url,
|
||||
model_name=expected_model_name,
|
||||
pid=spawn_handle.pid,
|
||||
backend=backend,
|
||||
spawned_by_us=True,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SpawnHandle:
|
||||
pid: Optional[int]
|
||||
cleanup_id: str # container name for docker, str(pid) for process
|
||||
cleanup_kind: str # "docker" | "process"
|
||||
@@ -0,0 +1,224 @@
|
||||
"""vllm backend: spawns the vllm/vllm-openai docker image with MTP=2."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from typing import List, Optional
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
from surya.inference.backends.base import Backend, ServerHandle
|
||||
from surya.inference.backends.openai_client import chat_completions_batch, resolve_max_workers
|
||||
from surya.inference.backends.spawn import (
|
||||
SpawnHandle,
|
||||
SpawnError,
|
||||
attach_or_spawn,
|
||||
)
|
||||
from surya.inference.schema import BatchInputItem, BatchOutputItem
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
from surya.timing import timing_span
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
# 24GB baseline (re-tune for surya-2 once benchmarks land)
|
||||
BASELINE_VRAM_GB = 24
|
||||
BASELINE_MAX_BATCHED_TOKENS = 8192
|
||||
BASELINE_MAX_NUM_SEQS = 32
|
||||
|
||||
GPU_VRAM_GB = {
|
||||
"b300": 270,
|
||||
"b200": 180,
|
||||
"h200": 141,
|
||||
"h100": 80,
|
||||
"a100-80": 80,
|
||||
"a100": 40,
|
||||
"a100-40": 40,
|
||||
"l40s": 48,
|
||||
"a10": 24,
|
||||
"l4": 24,
|
||||
"5090": 32,
|
||||
"4090": 24,
|
||||
"3090": 24,
|
||||
"t4": 16,
|
||||
}
|
||||
|
||||
|
||||
def _gpu_settings(gpu: str) -> tuple[int, int]:
|
||||
vram = GPU_VRAM_GB.get(gpu)
|
||||
if vram is None:
|
||||
available = ", ".join(sorted(GPU_VRAM_GB.keys()))
|
||||
raise SpawnError(f"Unknown VLLM_GPU_TYPE {gpu!r}. Available: {available}")
|
||||
ratio = vram / BASELINE_VRAM_GB
|
||||
raw_tokens = BASELINE_MAX_BATCHED_TOKENS * ratio
|
||||
max_batched_tokens = max(1024, 2 ** math.floor(math.log2(raw_tokens)))
|
||||
max_num_seqs = max(8, (int(BASELINE_MAX_NUM_SEQS * ratio) // 8) * 8)
|
||||
return max_batched_tokens, max_num_seqs
|
||||
|
||||
|
||||
def _resolve_docker_binary() -> str:
|
||||
found = shutil.which("docker")
|
||||
if found:
|
||||
return found
|
||||
raise SpawnError(
|
||||
"docker binary not found. Install Docker (https://docs.docker.com/get-docker/) "
|
||||
"and ensure the daemon is running."
|
||||
)
|
||||
|
||||
|
||||
def _health_url(port: int) -> str:
|
||||
return f"http://{settings.SURYA_INFERENCE_HOST}:{port}"
|
||||
|
||||
|
||||
def _openai_url(port: int) -> str:
|
||||
return f"http://{settings.SURYA_INFERENCE_HOST}:{port}/v1"
|
||||
|
||||
|
||||
class VllmBackend(Backend):
|
||||
name = "vllm"
|
||||
|
||||
def __init__(self):
|
||||
self.handle: Optional[ServerHandle] = None
|
||||
self._client: Optional[OpenAI] = None
|
||||
|
||||
def start(self) -> ServerHandle:
|
||||
if self.handle is not None:
|
||||
return self.handle
|
||||
|
||||
# If user pinned an external server, attach without spawning docker.
|
||||
if settings.SURYA_INFERENCE_URL:
|
||||
spawned = attach_or_spawn(
|
||||
backend=self.name,
|
||||
expected_model_name=settings.SURYA_MODEL_CHECKPOINT,
|
||||
spawn_fn=lambda port: SpawnHandle(
|
||||
pid=None, cleanup_id="", cleanup_kind="docker"
|
||||
),
|
||||
health_url_for=_health_url,
|
||||
openai_url_for=_openai_url,
|
||||
startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT,
|
||||
)
|
||||
self.handle = ServerHandle(
|
||||
base_url=spawned.base_url,
|
||||
model_name=spawned.model_name,
|
||||
spawned_by_us=spawned.spawned_by_us,
|
||||
)
|
||||
self._client = OpenAI(
|
||||
api_key=settings.VLLM_API_KEY, base_url=self.handle.base_url
|
||||
)
|
||||
return self.handle
|
||||
|
||||
if os.getenv("SUYA_ALLOW_NESTED_DOCKER", "false").lower() not in {"1", "true", "yes"}:
|
||||
raise SpawnError(
|
||||
"Nested Docker vLLM startup is disabled. Start vLLM in this "
|
||||
"container and set SURYA_INFERENCE_URL, for example "
|
||||
"http://127.0.0.1:8000/v1."
|
||||
)
|
||||
|
||||
docker = _resolve_docker_binary()
|
||||
max_batched_tokens, max_num_seqs = _gpu_settings(settings.VLLM_GPU_TYPE)
|
||||
|
||||
def spawn_fn(port: int) -> SpawnHandle:
|
||||
container_name = f"surya-vllm-{port}"
|
||||
hf_cache = os.path.expanduser(settings.DOCKER_HF_CACHE_PATH)
|
||||
cmd = [
|
||||
docker,
|
||||
"run",
|
||||
"--rm",
|
||||
"-d",
|
||||
"--name",
|
||||
container_name,
|
||||
"--runtime",
|
||||
"nvidia",
|
||||
"--gpus",
|
||||
f"device={settings.VLLM_GPUS}",
|
||||
"-v",
|
||||
f"{hf_cache}:/root/.cache/huggingface",
|
||||
"-p",
|
||||
f"{port}:8000",
|
||||
"--ipc=host",
|
||||
settings.VLLM_DOCKER_IMAGE,
|
||||
"--model",
|
||||
settings.SURYA_MODEL_CHECKPOINT,
|
||||
"--no-enforce-eager",
|
||||
"--max-num-seqs",
|
||||
str(max_num_seqs),
|
||||
"--dtype",
|
||||
settings.VLLM_DTYPE,
|
||||
"--max-model-len",
|
||||
str(settings.VLLM_MAX_MODEL_LEN),
|
||||
"--max-num-batched-tokens",
|
||||
str(max_batched_tokens),
|
||||
"--gpu-memory-utilization",
|
||||
str(settings.VLLM_GPU_MEMORY_UTILIZATION),
|
||||
"--enable-prefix-caching",
|
||||
"--mm-processor-kwargs",
|
||||
json.dumps({"min_pixels": 3136, "max_pixels": 6291456}),
|
||||
"--served-model-name",
|
||||
settings.SURYA_MODEL_CHECKPOINT,
|
||||
]
|
||||
if settings.VLLM_ENABLE_MTP:
|
||||
spec_config = json.dumps(
|
||||
{
|
||||
"method": "mtp",
|
||||
"num_speculative_tokens": settings.VLLM_MTP_TOKENS,
|
||||
}
|
||||
)
|
||||
cmd.extend(["--speculative-config", spec_config])
|
||||
for extra in (settings.VLLM_EXTRA_ARGS or "").split():
|
||||
cmd.append(extra)
|
||||
logger.info(f"Spawning: {' '.join(cmd)}")
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, check=False)
|
||||
if result.returncode != 0:
|
||||
raise SpawnError(f"docker run failed: {result.stderr or result.stdout}")
|
||||
return SpawnHandle(
|
||||
pid=None, cleanup_id=container_name, cleanup_kind="docker"
|
||||
)
|
||||
|
||||
spawned = attach_or_spawn(
|
||||
backend=self.name,
|
||||
expected_model_name=settings.SURYA_MODEL_CHECKPOINT,
|
||||
spawn_fn=spawn_fn,
|
||||
health_url_for=_health_url,
|
||||
openai_url_for=_openai_url,
|
||||
startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT,
|
||||
)
|
||||
self.handle = ServerHandle(
|
||||
base_url=spawned.base_url,
|
||||
model_name=spawned.model_name,
|
||||
spawned_by_us=spawned.spawned_by_us,
|
||||
)
|
||||
self._client = OpenAI(
|
||||
api_key=settings.VLLM_API_KEY,
|
||||
base_url=self.handle.base_url,
|
||||
)
|
||||
return self.handle
|
||||
|
||||
def stop(self) -> None:
|
||||
self.handle = None
|
||||
self._client = None
|
||||
|
||||
def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
|
||||
if self.handle is None or self._client is None:
|
||||
with timing_span("vllm_backend_start"):
|
||||
self.start()
|
||||
with timing_span(
|
||||
"vllm_backend_generate",
|
||||
item_count=len(batch),
|
||||
parallel=settings.SURYA_INFERENCE_PARALLEL,
|
||||
):
|
||||
return chat_completions_batch(
|
||||
batch,
|
||||
client=self._client,
|
||||
model_name=self.handle.model_name,
|
||||
timeout=settings.SURYA_INFERENCE_TIMEOUT_SECONDS,
|
||||
max_workers=resolve_max_workers(
|
||||
len(batch), settings.SURYA_INFERENCE_MAX_INFLIGHT
|
||||
),
|
||||
max_retries=settings.SURYA_INFERENCE_MAX_RETRIES,
|
||||
request_logprobs_default=settings.SURYA_INFERENCE_LOGPROBS,
|
||||
)
|
||||
@@ -0,0 +1,182 @@
|
||||
"""Parsers for the three task outputs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Tuple
|
||||
|
||||
|
||||
from surya.logging import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
# ---- Layout (LAYOUT_PROMPT) -------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParsedLayoutBlock:
|
||||
label: str
|
||||
bbox: Tuple[float, float, float, float] # 0-1000 normalized
|
||||
count: int # multiple of 50, model's token estimate
|
||||
|
||||
|
||||
_JSON_ARRAY_RE = re.compile(r"\[.*\]", re.DOTALL)
|
||||
|
||||
|
||||
def _strip_fences(text: str) -> str:
|
||||
cleaned = text.strip()
|
||||
if cleaned.startswith("```"):
|
||||
cleaned = re.sub(r"^```[a-zA-Z]*\n", "", cleaned)
|
||||
cleaned = re.sub(r"\n```\s*$", "", cleaned)
|
||||
return cleaned
|
||||
|
||||
|
||||
def _coerce_bbox(bbox) -> Tuple[float, float, float, float]:
|
||||
if isinstance(bbox, str):
|
||||
parts = [float(x) for x in bbox.replace(",", " ").split()]
|
||||
else:
|
||||
parts = [float(x) for x in bbox]
|
||||
if len(parts) != 4:
|
||||
raise ValueError(f"Bad bbox: {bbox!r}")
|
||||
return (parts[0], parts[1], parts[2], parts[3])
|
||||
|
||||
|
||||
def _coerce_count(value) -> int:
|
||||
if value is None:
|
||||
return 0
|
||||
try:
|
||||
return max(0, int(value))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def parse_layout(text: str) -> List[ParsedLayoutBlock]:
|
||||
"""Pull the JSON array out of LAYOUT_PROMPT output and convert to typed blocks.
|
||||
|
||||
Tolerates code fences, missing fields, and stringified bboxes.
|
||||
"""
|
||||
cleaned = _strip_fences(text)
|
||||
m = _JSON_ARRAY_RE.search(cleaned)
|
||||
if not m:
|
||||
raise ValueError(f"No JSON array found in layout output: {text[:500]!r}")
|
||||
raw = json.loads(m.group(0))
|
||||
out: List[ParsedLayoutBlock] = []
|
||||
for item in raw:
|
||||
try:
|
||||
bbox = _coerce_bbox(item["bbox"])
|
||||
except (KeyError, ValueError) as e:
|
||||
logger.warning(f"Skipping layout block with bad bbox: {e}")
|
||||
continue
|
||||
label = str(item.get("label", "block"))
|
||||
count = _coerce_count(item.get("count"))
|
||||
out.append(ParsedLayoutBlock(label=label, bbox=bbox, count=count))
|
||||
return out
|
||||
|
||||
|
||||
# ---- Table rec (TABLE_REC_PROMPT) ------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParsedTableElement:
|
||||
label: str # "Row" or "Col"
|
||||
bbox: Tuple[float, float, float, float]
|
||||
|
||||
|
||||
def parse_table_rec(text: str) -> List[ParsedTableElement]:
|
||||
"""Parse JSON array of {label: "Row"|"Col", bbox: "x0 y0 x1 y1"} from
|
||||
TABLE_REC_PROMPT output. Returns a flat list of Row + Col elements;
|
||||
cell derivation is the caller's job."""
|
||||
cleaned = _strip_fences(text)
|
||||
m = _JSON_ARRAY_RE.search(cleaned)
|
||||
if not m:
|
||||
raise ValueError(f"No JSON array found in table_rec output: {text[:500]!r}")
|
||||
raw = json.loads(m.group(0))
|
||||
out: List[ParsedTableElement] = []
|
||||
for item in raw:
|
||||
label = str(item.get("label", "")).strip()
|
||||
if label not in ("Row", "Col"):
|
||||
continue
|
||||
try:
|
||||
bbox = _coerce_bbox(item["bbox"])
|
||||
except (KeyError, ValueError):
|
||||
continue
|
||||
out.append(ParsedTableElement(label=label, bbox=bbox))
|
||||
return out
|
||||
|
||||
|
||||
# ---- Block HTML (BLOCK_PROMPT for full table path / general block path) ---
|
||||
|
||||
|
||||
def clean_block_html(html: str) -> str:
|
||||
"""Light cleanup of model-emitted HTML for a single block.
|
||||
|
||||
Strips code fences, leading/trailing whitespace. Does NOT validate against
|
||||
ALLOWED_TAGS — the model is expected to comply, and downstream consumers
|
||||
can sanitize further if needed.
|
||||
"""
|
||||
cleaned = _strip_fences(html).strip()
|
||||
return cleaned
|
||||
|
||||
|
||||
# ---- Full-page fallback (HIGH_ACCURACY_BBOX_PROMPT) -----------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParsedFullPageBlock:
|
||||
label: str
|
||||
bbox: Tuple[float, float, float, float] # 0-1000 normalized
|
||||
html: str # inner HTML of the wrapping div
|
||||
|
||||
|
||||
def parse_full_page_html(text: str) -> List[ParsedFullPageBlock]:
|
||||
"""Parse output of HIGH_ACCURACY_BBOX_PROMPT — top-level <div data-bbox=...
|
||||
data-label=...>inner HTML</div> blocks. Returns one entry per top-level div."""
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
cleaned = _strip_fences(text).strip()
|
||||
if not cleaned:
|
||||
return []
|
||||
# The model outputs a sequence of top-level divs (no surrounding root).
|
||||
# BeautifulSoup parses fine without one.
|
||||
soup = BeautifulSoup(cleaned, "html.parser")
|
||||
divs = soup.find_all("div", recursive=False)
|
||||
out: List[ParsedFullPageBlock] = []
|
||||
for div in divs:
|
||||
label = div.get("data-label")
|
||||
bbox_str = div.get("data-bbox")
|
||||
if not label or not bbox_str:
|
||||
continue
|
||||
try:
|
||||
parts = [float(x) for x in bbox_str.split()]
|
||||
except ValueError:
|
||||
continue
|
||||
if len(parts) != 4:
|
||||
continue
|
||||
# Strip nested data-bbox attrs from the inner HTML so downstream
|
||||
# consumers don't see model debug info on every child element.
|
||||
for tag in div.find_all(attrs={"data-bbox": True}):
|
||||
del tag["data-bbox"]
|
||||
for tag in div.find_all(attrs={"data-label": True}):
|
||||
del tag["data-label"]
|
||||
inner = "".join(str(c) for c in div.contents).strip()
|
||||
out.append(
|
||||
ParsedFullPageBlock(
|
||||
label=str(label),
|
||||
bbox=(parts[0], parts[1], parts[2], parts[3]),
|
||||
html=inner,
|
||||
)
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def denorm_bbox(bbox, img_w: int, img_h: int, scale: int = 1000):
|
||||
x0, y0, x1, y1 = bbox
|
||||
return (
|
||||
x0 / scale * img_w,
|
||||
y0 / scale * img_h,
|
||||
x1 / scale * img_w,
|
||||
y1 / scale * img_h,
|
||||
)
|
||||
@@ -0,0 +1,158 @@
|
||||
"""Prompt strings for surya2. The exact wording is the model's training-time
|
||||
contract — do not paraphrase without retraining."""
|
||||
|
||||
from surya.inference.schema import PROMPT_TYPE_BLOCK as PROMPT_TYPE_BLOCK
|
||||
from surya.inference.schema import (
|
||||
PROMPT_TYPE_HIGH_ACCURACY_BBOX as PROMPT_TYPE_HIGH_ACCURACY_BBOX,
|
||||
)
|
||||
from surya.inference.schema import PROMPT_TYPE_LAYOUT as PROMPT_TYPE_LAYOUT
|
||||
from surya.inference.schema import PROMPT_TYPE_TABLE_REC as PROMPT_TYPE_TABLE_REC
|
||||
|
||||
ALLOWED_TAGS = [
|
||||
"math",
|
||||
"br",
|
||||
"i",
|
||||
"b",
|
||||
"u",
|
||||
"del",
|
||||
"sup",
|
||||
"sub",
|
||||
"table",
|
||||
"tr",
|
||||
"td",
|
||||
"p",
|
||||
"th",
|
||||
"div",
|
||||
"pre",
|
||||
"h1",
|
||||
"h2",
|
||||
"h3",
|
||||
"h4",
|
||||
"h5",
|
||||
"ul",
|
||||
"ol",
|
||||
"li",
|
||||
"input",
|
||||
"a",
|
||||
"span",
|
||||
"img",
|
||||
"hr",
|
||||
"tbody",
|
||||
"small",
|
||||
"caption",
|
||||
"strong",
|
||||
"thead",
|
||||
"big",
|
||||
"code",
|
||||
"chem",
|
||||
]
|
||||
|
||||
ALLOWED_ATTRIBUTES = [
|
||||
"class",
|
||||
"colspan",
|
||||
"rowspan",
|
||||
"display",
|
||||
"checked",
|
||||
"type",
|
||||
"border",
|
||||
"value",
|
||||
"style",
|
||||
"href",
|
||||
"alt",
|
||||
"align",
|
||||
"data-bbox",
|
||||
"data-label",
|
||||
]
|
||||
|
||||
# Block labels we don't run OCR on.
|
||||
SKIP_OCR_LABELS = {"Figure", "Image", "Diagram", "Blank-Page"}
|
||||
|
||||
LAYOUT_PROMPT = (
|
||||
"Output the layout of this image as JSON. Each entry is a dict with "
|
||||
'"label", "bbox", and "count" fields. Bbox is x0 y0 x1 y1, normalized 0-1000.'
|
||||
)
|
||||
|
||||
BLOCK_PROMPT = "OCR this block image to HTML."
|
||||
|
||||
TABLE_REC_PROMPT = (
|
||||
"Output the table rows then columns as JSON. Each entry is a dict with "
|
||||
'"label" ("Row" or "Col") and "bbox" (x0 y0 x1 y1, normalized 0-1000).'
|
||||
)
|
||||
|
||||
HIGH_ACCURACY_BBOX_PROMPT = (
|
||||
"OCR this image to HTML. Each block is a div with data-label and data-bbox "
|
||||
"(x0 y0 x1 y1, normalized 0-1000)."
|
||||
)
|
||||
|
||||
|
||||
PROMPT_MAPPING = {
|
||||
"layout": LAYOUT_PROMPT,
|
||||
"block": BLOCK_PROMPT,
|
||||
"table_rec": TABLE_REC_PROMPT,
|
||||
"high_accuracy_bbox": HIGH_ACCURACY_BBOX_PROMPT,
|
||||
}
|
||||
|
||||
|
||||
# JSON schema for LAYOUT_PROMPT — enforced via vllm guided decoding so the
|
||||
# model can't emit malformed JSON. bbox is a "x0 y0 x1 y1" string (model's
|
||||
# training-time format); count is a non-negative integer.
|
||||
LAYOUT_LABEL_SET = [
|
||||
"Caption",
|
||||
"Footnote",
|
||||
"Equation-Block",
|
||||
"List-Group",
|
||||
"Page-Header",
|
||||
"Page-Footer",
|
||||
"Image",
|
||||
"Section-Header",
|
||||
"Table",
|
||||
"Text",
|
||||
"Complex-Block",
|
||||
"Code-Block",
|
||||
"Form",
|
||||
"Table-Of-Contents",
|
||||
"Figure",
|
||||
"Chemical-Block",
|
||||
"Diagram",
|
||||
"Bibliography",
|
||||
"Blank-Page",
|
||||
]
|
||||
|
||||
LAYOUT_JSON_SCHEMA = {
|
||||
"type": "array",
|
||||
"maxItems": 200,
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"label": {"type": "string", "enum": LAYOUT_LABEL_SET},
|
||||
"bbox": {
|
||||
"type": "string",
|
||||
"pattern": r"^\d{1,4} \d{1,4} \d{1,4} \d{1,4}$",
|
||||
},
|
||||
"count": {"type": "integer", "minimum": 0, "maximum": 10000},
|
||||
},
|
||||
"required": ["label", "bbox", "count"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# JSON schema for TABLE_REC_PROMPT — array of {label: Row|Col, bbox: "x0 y0 x1 y1"}.
|
||||
TABLE_REC_LABEL_SET = ["Row", "Col"]
|
||||
|
||||
TABLE_REC_JSON_SCHEMA = {
|
||||
"type": "array",
|
||||
"maxItems": 200,
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"label": {"type": "string", "enum": TABLE_REC_LABEL_SET},
|
||||
"bbox": {
|
||||
"type": "string",
|
||||
"pattern": r"^\d{1,4} \d{1,4} \d{1,4} \d{1,4}$",
|
||||
},
|
||||
},
|
||||
"required": ["label", "bbox"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
PROMPT_TYPE_LAYOUT = "layout"
|
||||
PROMPT_TYPE_BLOCK = "block"
|
||||
PROMPT_TYPE_TABLE_REC = "table_rec"
|
||||
PROMPT_TYPE_HIGH_ACCURACY_BBOX = "high_accuracy_bbox"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchInputItem:
|
||||
image: Image.Image
|
||||
prompt_type: str
|
||||
prompt: Optional[str] = None # If set, overrides the default prompt for prompt_type
|
||||
max_tokens: Optional[int] = None
|
||||
request_logprobs: bool = False
|
||||
# vllm-native guided decoding — JSON schema, regex, or grammar string.
|
||||
# When set, the server constrains the decode tokens to match the schema.
|
||||
guided_json: Optional[dict] = None
|
||||
guided_regex: Optional[str] = None
|
||||
metadata: dict = field(default_factory=dict) # Free-form, passes through to output
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationResult:
|
||||
raw: str
|
||||
token_count: int
|
||||
error: bool = False
|
||||
# Mean of exp(logprob) across response tokens, if logprobs requested
|
||||
mean_token_prob: Optional[float] = None
|
||||
# Per-token logprobs (raw OpenAI-style content list), if requested - phase 2 use
|
||||
logprobs: Optional[List[Any]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchOutputItem:
|
||||
raw: str
|
||||
token_count: int
|
||||
error: bool
|
||||
mean_token_prob: Optional[float] = None
|
||||
logprobs: Optional[List[Any]] = None
|
||||
metadata: dict = field(default_factory=dict)
|
||||
@@ -0,0 +1,96 @@
|
||||
from typing import Tuple
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def scale_to_fit(
|
||||
img: Image.Image,
|
||||
max_size: Tuple[int, int] = (3072, 2048),
|
||||
min_size: Tuple[int, int] = (1792, 28),
|
||||
grid_size: int = 28,
|
||||
) -> Image.Image:
|
||||
resample_method = Image.Resampling.LANCZOS
|
||||
|
||||
width, height = img.size
|
||||
|
||||
if width <= 0 or height <= 0:
|
||||
return img
|
||||
|
||||
original_ar = width / height
|
||||
current_pixels = width * height
|
||||
max_pixels = max_size[0] * max_size[1]
|
||||
min_pixels = min_size[0] * min_size[1]
|
||||
|
||||
scale = 1.0
|
||||
if current_pixels > max_pixels:
|
||||
scale = (max_pixels / current_pixels) ** 0.5
|
||||
elif current_pixels < min_pixels:
|
||||
scale = (min_pixels / current_pixels) ** 0.5
|
||||
|
||||
w_blocks = max(1, round((width * scale) / grid_size))
|
||||
h_blocks = max(1, round((height * scale) / grid_size))
|
||||
|
||||
while (w_blocks * h_blocks * grid_size * grid_size) > max_pixels:
|
||||
if w_blocks == 1 and h_blocks == 1:
|
||||
break
|
||||
|
||||
if w_blocks == 1:
|
||||
h_blocks -= 1
|
||||
continue
|
||||
if h_blocks == 1:
|
||||
w_blocks -= 1
|
||||
continue
|
||||
|
||||
ar_w_loss = abs(((w_blocks - 1) / h_blocks) - original_ar)
|
||||
ar_h_loss = abs((w_blocks / (h_blocks - 1)) - original_ar)
|
||||
|
||||
if ar_w_loss < ar_h_loss:
|
||||
w_blocks -= 1
|
||||
else:
|
||||
h_blocks -= 1
|
||||
|
||||
new_width = w_blocks * grid_size
|
||||
new_height = h_blocks * grid_size
|
||||
|
||||
if (new_width, new_height) == (width, height):
|
||||
return img
|
||||
|
||||
return img.resize((new_width, new_height), resample=resample_method)
|
||||
|
||||
|
||||
def detect_repeat_token(
|
||||
predicted_tokens: str,
|
||||
base_max_repeats: int = 4,
|
||||
window_size: int = 500,
|
||||
cut_from_end: int = 0,
|
||||
scaling_factor: float = 3.0,
|
||||
) -> bool:
|
||||
if cut_from_end > 0:
|
||||
predicted_tokens = predicted_tokens[:-cut_from_end]
|
||||
|
||||
for seq_len in range(1, window_size // 2 + 1):
|
||||
candidate_seq = predicted_tokens[-seq_len:]
|
||||
|
||||
max_repeats = int(base_max_repeats * (1 + scaling_factor / seq_len))
|
||||
|
||||
repeat_count = 0
|
||||
pos = len(predicted_tokens) - seq_len
|
||||
if pos < 0:
|
||||
continue
|
||||
|
||||
while pos >= 0:
|
||||
if predicted_tokens[pos : pos + seq_len] == candidate_seq:
|
||||
repeat_count += 1
|
||||
pos -= seq_len
|
||||
else:
|
||||
break
|
||||
|
||||
if repeat_count > max_repeats:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def image_token_budget(block_count: int, ceiling: int = 4096, floor: int = 64) -> int:
|
||||
"""Per-block max_tokens: count + 100, clamped to [floor, ceiling]."""
|
||||
return min(max(block_count + 100, floor), ceiling)
|
||||
Reference in New Issue
Block a user