api/api/services/docling/dsync.py

#!/usr/bin/env python3
"""
dsync.py — Redis-backed sync layer in front of docling-serve.

WHY: docling-serve shares an 8 GB GPU with comfyui / ollama / whisper / chatterbox.
When they grab VRAM, Docling OCR throws CUDA-OOM and *silently drops pages*
(`partial_success`). We can't evict the other apps and we are NOT pinning a GPU, so
instead we make extraction robust to OOM *by construction*:

  1. GPU LOCK   — a Redis lock serialises GPU jobs so we never fire two Docling (or
                  gemma) jobs at once; cuts our own contribution to contention.
  2. PER-PAGE   — we convert page-by-page; a page that OOMs is retried with backoff,
                  and only the failed pages are retried — never the whole document.
  3. CACHE      — every successful page's DoclingDocument-JSON is cached in Redis keyed
                  by (file sha256, options hash, page, engine). Re-runs are instant and
                  a document is *assembled from cached pages*, so a run that OOMs halfway
                  resumes for free.

Connection (env):
  DOCLING_REDIS_URL = redis://:[email protected]:30059/0
  (or DOCLING_REDIS_HOST/PORT/PASSWORD/DB). Falls back to no-cache if unset/unreachable.

Usage:
  from dsync import convert_document
  doc = convert_document("samples/AQA-Physics-Paper-1H-2022-with-qr.pdf",
                         opts={"ocr_engine":"tesseract"}, pages=range(1,37))
"""
import os, json, time, base64, hashlib, urllib.request, urllib.error

SERVE = os.environ.get("DOCLING_SERVE", "http://192.168.0.39:5001")
LOCK_KEY = "docling:gpulock"
LOCK_TTL = 900            # seconds; lock auto-expires so a crashed job can't deadlock us
CACHE_TTL = 7 * 24 * 3600
DEFAULT_OPTS = {"to_formats": ["json"], "image_export_mode": "placeholder", "do_ocr": True}


# ----------------------------------------------------------------- redis (optional)
def _redis():
    try:
        import redis
    except ImportError:
        return None
    url = os.environ.get("DOCLING_REDIS_URL")
    try:
        if url:
            c = redis.from_url(url, socket_timeout=4)
        else:
            host = os.environ.get("DOCLING_REDIS_HOST", "192.168.0.19")
            c = redis.Redis(host=host,
                            port=int(os.environ.get("DOCLING_REDIS_PORT", 30059)),
                            password=os.environ.get("DOCLING_REDIS_PASSWORD"),
                            db=int(os.environ.get("DOCLING_REDIS_DB", 0)),
                            socket_timeout=4)
        c.ping()
        return c
    except Exception as e:
        print(f"[dsync] redis unavailable ({e}); running without cache/lock")
        return None


class _GpuLock:
    """Best-effort distributed lock so only one GPU job runs at a time."""
    def __init__(self, r): self.r = r; self.tok = None
    def __enter__(self):
        if not self.r: return self
        self.tok = str(time.time())
        while not self.r.set(LOCK_KEY, self.tok, nx=True, ex=LOCK_TTL):
            time.sleep(1.5)
        return self
    def __exit__(self, *a):
        if self.r and self.tok and self.r.get(LOCK_KEY) == self.tok.encode():
            self.r.delete(LOCK_KEY)


# ----------------------------------------------------------------- keys
def _sha(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()[:16]


def _page_key(sha, opts, page):
    oh = hashlib.sha256(json.dumps(opts, sort_keys=True).encode()).hexdigest()[:12]
    return f"docling:page:{sha}:{oh}:{page}"


# ----------------------------------------------------------------- serve call
def _serve_convert(pdf_b64, fname, opts):
    body = {"options": opts,
            "sources": [{"kind": "file", "base64_string": pdf_b64, "filename": fname}],
            "target": {"kind": "inbody"}}
    req = urllib.request.Request(SERVE + "/v1/convert/source",
                                 data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    for _ in range(4):                              # tolerate the single-use 404 race
        try:
            return json.loads(urllib.request.urlopen(req, timeout=1200).read())
        except urllib.error.HTTPError as e:
            if e.code == 404:
                time.sleep(3); continue
            raise
    raise RuntimeError("serve: repeated 404")


def _is_oom(resp):
    return any("out of memory" in str(e).lower() for e in (resp.get("errors") or []))


# ----------------------------------------------------------------- public API
def convert_page(pdf, page, opts=None, *, r=None, retries=5):
    """Convert a single page, with cache + GPU-lock + OOM backoff. Returns the
    per-page DoclingDocument JSON (or None on hard failure)."""
    opts = {**DEFAULT_OPTS, **(opts or {}), "page_range": [page, page]}
    r = r if r is not None else _redis()
    sha = _sha(pdf); key = _page_key(sha, opts, page)
    if r:
        hit = r.get(key)
        if hit:
            print(f"[dsync] p{page} cache HIT")
            return json.loads(hit)
    b64 = base64.b64encode(open(pdf, "rb").read()).decode()
    fname = os.path.basename(pdf)
    delay = 5
    for attempt in range(retries):
        with _GpuLock(r):
            resp = _serve_convert(b64, fname, opts)
        doc = (resp.get("document") or {}).get("json_content")
        if doc and not _is_oom(resp):
            if r:
                r.set(key, json.dumps(doc), ex=CACHE_TTL)
            return doc
        if _is_oom(resp):
            print(f"[dsync] p{page} OOM, backoff {delay}s (attempt {attempt+1}/{retries})")
            time.sleep(delay); delay = min(delay * 2, 120)
            continue
        return doc                                  # non-OOM result (may be empty); don't loop
    print(f"[dsync] p{page} gave up after {retries} OOM retries")
    return None


def convert_document(pdf, opts=None, pages=None):
    """Convert all (or selected) pages page-by-page and merge into one structure.
    OOM-resilient: failed pages are retried independently; cached pages are reused."""
    r = _redis()
    if pages is None:
        import subprocess
        n = int(subprocess.check_output(["pdfinfo", pdf]).decode().split("Pages:")[1].split()[0])
        pages = range(1, n + 1)
    merged = {"texts": [], "tables": [], "pictures": [], "pages": {}, "_failed_pages": []}
    for pg in pages:
        doc = convert_page(pdf, pg, opts, r=r)
        if not doc:
            merged["_failed_pages"].append(pg); continue
        for k in ("texts", "tables", "pictures"):
            merged[k].extend(doc.get(k, []))
        merged["pages"].update(doc.get("pages", {}))
    return merged


if __name__ == "__main__":
    import sys
    pdf = sys.argv[1] if len(sys.argv) > 1 else "samples/AQA-Physics-Paper-1H-2022-with-qr.pdf"
    r = _redis()
    print("redis:", "connected" if r else "NOT connected (set DOCLING_REDIS_URL / _PASSWORD)")
    if r:
        d = convert_document(pdf, {"ocr_engine": "tesseract"}, pages=range(1, 5))
        print(f"merged texts={len(d['texts'])} failed_pages={d['_failed_pages']}")