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,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,
|
||||
)
|
||||
Reference in New Issue
Block a user