Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import json | |
| import os | |
| import inspect | |
| import time | |
| from concurrent.futures import TimeoutError | |
| from pathlib import Path | |
| from typing import Any | |
| from gradio_client import Client, handle_file | |
| class WorkerError(RuntimeError): | |
| pass | |
| def normalize_endpoint(endpoint: str) -> str: | |
| value = (endpoint or "").strip().rstrip("/") | |
| if not value: | |
| return "" | |
| if not value.startswith(("http://", "https://")): | |
| value = "https://" + value | |
| return value | |
| def _client(endpoint: str) -> Client: | |
| hf_token = os.getenv("HF_TOKEN", "").strip() or None | |
| if not hf_token: | |
| return Client(endpoint) | |
| signature = inspect.signature(Client.__init__) | |
| if "hf_token" in signature.parameters: | |
| return Client(endpoint, hf_token=hf_token) | |
| if "token" in signature.parameters: | |
| return Client(endpoint, token=hf_token) | |
| if "headers" in signature.parameters: | |
| return Client(endpoint, headers={"Authorization": f"Bearer {hf_token}"}) | |
| return Client(endpoint) | |
| def check_health(endpoint: str, timeout_seconds: float = 20.0) -> dict[str, Any]: | |
| endpoint = normalize_endpoint(endpoint) | |
| if not endpoint: | |
| raise WorkerError("Worker URL is not configured.") | |
| try: | |
| job = _client(endpoint).submit(api_name="/health") | |
| payload = job.result(timeout=timeout_seconds) | |
| if not isinstance(payload, dict): | |
| raise WorkerError("Worker health response is not a JSON object.") | |
| return payload | |
| except TimeoutError as exc: | |
| raise WorkerError(f"Worker health check timed out after {timeout_seconds:.0f}s.") from exc | |
| except Exception as exc: | |
| raise WorkerError(f"Worker health check failed: {exc}") from exc | |
| def run_page( | |
| *, | |
| endpoint: str, | |
| model_id: str, | |
| page_path: Path, | |
| prompt: str, | |
| options: dict[str, Any], | |
| timeout_seconds: float, | |
| ) -> dict[str, Any]: | |
| endpoint = normalize_endpoint(endpoint) | |
| if not endpoint: | |
| raise WorkerError("Worker URL is not configured.") | |
| retry_delays = (0.0, 2.0, 6.0) | |
| last_error: Exception | None = None | |
| for attempt, delay in enumerate(retry_delays, start=1): | |
| if delay: | |
| time.sleep(delay) | |
| job = None | |
| try: | |
| job = _client(endpoint).submit( | |
| handle_file(str(page_path)), | |
| model_id, | |
| prompt or "", | |
| json.dumps(options, ensure_ascii=False), | |
| os.getenv("WORKER_API_TOKEN", "").strip(), | |
| api_name="/ocr", | |
| ) | |
| payload = job.result(timeout=timeout_seconds) | |
| if not isinstance(payload, dict): | |
| raise WorkerError("Worker response is not a JSON object.") | |
| required = {"model", "text", "markdown", "metrics"} | |
| missing = sorted(required - payload.keys()) | |
| if missing: | |
| raise WorkerError(f"Worker response is missing fields: {missing}") | |
| return payload | |
| except TimeoutError as exc: | |
| last_error = WorkerError(f"Worker request timed out after {timeout_seconds:.0f}s.") | |
| if job is not None: | |
| try: | |
| job.cancel() | |
| except Exception: | |
| pass | |
| if attempt >= len(retry_delays): | |
| break | |
| except (ValueError, OSError, WorkerError, Exception) as exc: | |
| last_error = exc | |
| if attempt >= len(retry_delays): | |
| break | |
| raise WorkerError(f"OCR request failed after retries: {last_error}") | |