Share conversion deadlines across dispatch and workers

backendN/app/common/app_factory.py

@@ -26,6 +26,7 @@
 
 from .config import Settings, deployment_revision
 from .errors import ConversionRejected, EngineUnavailable
+from .deadlines import DeadlineExceeded, DeadlineMiddleware, check_budget, compute_checkpoint
 from .files import read_limited, sanitize_filename, split_name, title_from_stem
 from .profiles import PROFILES, count_words, finalize_markdown
 from .security import require_internal_secret
@@ -75,6 +76,7 @@ def create_app(settings: Settings, converter: Converter, storage: StorageClient
     # No interactive docs or schema: these instances are internal compute,
     # not a public API surface.
     app = FastAPI(title=f"LLMiser backend ({settings.role})", docs_url=None, redoc_url=None, openapi_url=None)
+    app.add_middleware(DeadlineMiddleware)
     app.state.settings = settings
     app.state.converter = converter
     app.state.storage = storage if storage is not None else StorageClient.from_env()
@@ -95,9 +97,12 @@ def run(original_name: str, profile: str, job_id: str, load_bytes, *, raw: bool
             if profile not in PROFILES:
                 profile = "Standard"
             data = load_bytes()
+            compute_checkpoint()
             converter.validate(data, ext)
             title = title_from_stem(stem)
+            compute_checkpoint()
             markdown, warning = converter.convert(data, ext, filename=original_name, title=title)
+            compute_checkpoint()
             if raw:
                 content = markdown.strip() + "\n"
                 words = count_words(content)
@@ -112,6 +117,9 @@ def run(original_name: str, profile: str, job_id: str, load_bytes, *, raw: bool
                     profile=profile,
                 )
             outcome = "ok"
+        except DeadlineExceeded:
+            outcome = "deadline"
+            raise
         except ConversionRejected as err:
             outcome = err.code
             raise HTTPException(status_code=err.status_code, detail=err.detail) from None
@@ -214,9 +222,11 @@ def process(req: ProcessRequest) -> dict:
         )
 
         write_started = time.perf_counter()
+        check_budget()
         body = content.encode("utf-8")
         try:
             storage.upload(req.output_path, body, "text/markdown; charset=utf-8")
+            check_budget()
         except StorageError as err:
             # The conversion worked but the result isn't stored: never report
             # success. Transient → the dispatcher may retry on the peer.
@@ -248,6 +258,7 @@ def convert(
             job_id,
             lambda: read_limited(file.file, settings.max_upload_bytes),
         )
+        check_budget()
         return {**fields, "content": content}
 
     return app

backendN/app/common/deadlines.py

@@ -0,0 +1,91 @@
+"""Optional authenticated per-request budget; legacy callers keep their limits."""
+from contextvars import ContextVar
+from dataclasses import dataclass
+from functools import wraps
+import re
+import time
+
+from fastapi import HTTPException
+
+HEADER = b"x-llmiser-deadline-ms"
+FINALIZE_RESERVE_S = 10.0
+MAX_BUDGET_S = 300.0
+_current = ContextVar("processing_deadline", default=None)
+
+
+class DeadlineExceeded(HTTPException):
+    def __init__(self):
+        super().__init__(504, "Conversion timed out. Try a smaller file or retry later.",
+                         headers={"X-LLMiser-Deadline-Exceeded": "1"})
+
+
+@dataclass(frozen=True)
+class Budget:
+    end: float
+
+    def remaining(self):
+        return max(0.0, self.end - time.monotonic())
+
+
+class DeadlineMiddleware:
+    """Set context before body parsing/thread admission; auth validates it later."""
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        value = None
+        if scope["type"] == "http" and scope.get("path", "").startswith("/api/v1/internal/"):
+            headers = [v for k, v in scope.get("headers", []) if k.lower() == HEADER]
+            if headers:
+                raw = headers[0]
+                if len(headers) != 1 or not re.fullmatch(rb"[0-9]{1,16}", raw):
+                    value = "invalid"
+                else:
+                    seconds = min(MAX_BUDGET_S, max(0.0, int(raw) / 1000 - time.time()))
+                    value = Budget(time.monotonic() + seconds)
+        token = _current.set(value)
+        try:
+            await self.app(scope, receive, send)
+        finally:
+            _current.reset(token)
+
+
+def current_budget():
+    value = _current.get()
+    if value == "invalid":
+        raise HTTPException(400, "Invalid processing deadline")
+    return value
+
+
+def check_budget(reserve=0.0):
+    budget = current_budget()
+    if budget is not None and budget.remaining() <= reserve:
+        raise DeadlineExceeded()
+
+
+def compute_checkpoint():
+    check_budget(FINALIZE_RESERVE_S)
+
+
+def timeout_seconds(default, *, reserve=0.0):
+    budget = current_budget()
+    if budget is None:
+        return default
+    left = budget.remaining() - reserve
+    if left <= 0:
+        raise DeadlineExceeded()
+    return min(default, left)
+
+
+def budgeted_worker(fn):
+    """ThreadPoolExecutor does not inherit the request's ContextVar."""
+    value = _current.get()
+    @wraps(fn)
+    def run(*args, **kwargs):
+        token = _current.set(value)
+        try:
+            check_budget()
+            return fn(*args, **kwargs)
+        finally:
+            _current.reset(token)
+    return run

backendN/app/common/security.py

@@ -5,6 +5,7 @@
 import hmac
 
 from fastapi import Header, HTTPException, Request
+from .deadlines import check_budget
 
 
 def require_internal_secret(
@@ -23,3 +24,4 @@ def require_internal_secret(
         x_internal_secret.encode(), expected.encode()
     ):
         raise HTTPException(status_code=401, detail="Unauthorized")
+    check_budget()

backendN/app/common/storage.py

@@ -16,6 +16,7 @@
 from urllib.parse import quote
 
 import httpx
+from .deadlines import check_budget, timeout_seconds
 
 # Backends only touch the fixed object layout of a job (docs/BACKGROUND_JOBS.md).
 # Anything else is refused, so a bug or a compromised caller can't read or
@@ -54,6 +55,7 @@ class StorageClient:
     def __init__(self, base_url: str, service_key: str, bucket: str, *, timeout: float = 30.0, transport=None):
         self._base = base_url.rstrip("/") + "/storage/v1"
         self._bucket = bucket
+        self._timeout_s = timeout
         self._http = httpx.Client(
             timeout=timeout,
             transport=transport,
@@ -73,19 +75,21 @@ def _url(self, path: str) -> str:
 
     def download(self, path: str, max_bytes: int) -> bytes:
         try:
-            with self._http.stream("GET", self._url(path)) as res:
+            with self._http.stream("GET", self._url(path), timeout=timeout_seconds(self._timeout_s)) as res:
                 if res.status_code in (400, 404):
                     raise StorageError("Input object not found", transient=False, status=res.status_code)
                 if res.status_code >= 400:
                     raise StorageError(f"Storage read failed ({res.status_code})", transient=res.status_code >= 500, status=res.status_code)
                 chunks, size = [], 0
                 for chunk in res.iter_bytes():
+                    check_budget()
                     size += len(chunk)
                     if size > max_bytes:
                         raise StorageError("Input object exceeds the size limit", transient=False, status=413)
                     chunks.append(chunk)
                 return b"".join(chunks)
         except httpx.HTTPError as err:
+            check_budget()
             raise StorageError(f"Storage unreachable: {type(err).__name__}", transient=True) from None
 
     def upload(self, path: str, data: bytes, content_type: str) -> None:
@@ -94,9 +98,12 @@ def upload(self, path: str, data: bytes, content_type: str) -> None:
             res = self._http.post(
                 self._url(path),
                 content=data,
+                timeout=timeout_seconds(self._timeout_s),
                 headers={"Content-Type": content_type, "x-upsert": "true", "Cache-Control": "no-store"},
             )
         except httpx.HTTPError as err:
+            check_budget()
             raise StorageError(f"Storage unreachable: {type(err).__name__}", transient=True) from None
         if res.status_code >= 400:
             raise StorageError(f"Storage write failed ({res.status_code})", transient=res.status_code >= 500, status=res.status_code)
+        check_budget()

backendN/app/pdf_tasks.py

@@ -34,6 +34,7 @@
 from app.common.files import sanitize_filename, split_name, title_from_stem
 from app.common.profiles import PROFILES, finalize_markdown
 from app.common.security import require_internal_secret
+from app.common.deadlines import budgeted_worker, check_budget, compute_checkpoint
 from app.common.storage import StorageError
 
 logger = logging.getLogger("mdify.pdf")
@@ -64,6 +65,7 @@ def _pdfium_serialized(fn):
     @wraps(fn)
     def run(*args, **kwargs):
         with _PDFIUM_LOCK:
+            compute_checkpoint()
             return fn(*args, **kwargs)
     return run
 
@@ -89,6 +91,7 @@ def classify_pages(pdf) -> list[PageInfo]:
 
     pages = []
     for index in range(len(pdf)):
+        compute_checkpoint()
         page = pdf[index]
         try:
             width, height = page.get_size()
@@ -206,6 +209,7 @@ def analyze(req: ProcessRequest) -> dict:
         try:
             converter.validate(data, "pdf")
             with _PDFIUM_LOCK:
+                compute_checkpoint()
                 pdf = pdfium.PdfDocument(data)
         except ConversionRejected as err:
             raise HTTPException(status_code=err.status_code, detail=err.detail) from None
@@ -235,13 +239,15 @@ def analyze(req: ProcessRequest) -> dict:
             skipped = scans[MAX_OCR_PAGES:]
             segments, scan_units = [], []
             for kind, first, last in page_runs(pages):
+                compute_checkpoint()
                 node_id = str(uuid.uuid4())
                 if kind == "native" or first not in ocr_pages:
                     if kind == "scan":
                         continue  # over the OCR page limit: noted in the merge
                     markdown, _warning = converter.convert(
                         sub_pdf(pdf, first, last), "pdf", filename=name, title=name, note_low_text=False
                     )
+                    compute_checkpoint()
                     path = f"jobs/{job_id}/nodes/{node_id}/result.md"
                     store.upload(path, (markdown.strip() + "\n").encode("utf-8"), "text/markdown; charset=utf-8")
                     segments.append({"node_id": node_id, "page_from": first, "page_to": last, "output_path": path})
@@ -302,6 +308,7 @@ def merge(req: MergeRequest) -> dict:
         profile = req.profile if req.profile in PROFILES else "Standard"
 
         def load(part: MergePart) -> tuple[MergePart, str | None]:
+            compute_checkpoint()
             try:
                 return part, store.download(part.path, settings.max_upload_bytes).decode("utf-8", "replace")
             except StorageError as err:
@@ -311,12 +318,13 @@ def load(part: MergePart) -> tuple[MergePart, str | None]:
 
         try:
             with ThreadPoolExecutor(max_workers=UPLOAD_WORKERS) as pool:
-                loaded = list(pool.map(load, sorted(req.parts, key=lambda p: p.page_from)))
+                loaded = list(pool.map(budgeted_worker(load), sorted(req.parts, key=lambda p: p.page_from)))
         except StorageError:
             raise HTTPException(status_code=503, detail="Storage temporarily unavailable") from None
 
         blocks, failed_pages = [], []
         for part, text in loaded:
+            compute_checkpoint()
             if part.kind == "ocr":
                 label = f"<!-- page {part.page_from}: text recognition -->"
                 if text is None or not text.strip():
@@ -333,6 +341,7 @@ def load(part: MergePart) -> tuple[MergePart, str | None]:
                 f"(LLMiser reads up to {MAX_OCR_PAGES} scanned pages per PDF): {listed}."
             )
 
+        compute_checkpoint()
         content, words, tokens = finalize_markdown(
             "\n\n".join(blocks) or "_No text found in this document._",
             title=title_from_stem(stem),
@@ -342,6 +351,7 @@ def load(part: MergePart) -> tuple[MergePart, str | None]:
             profile=profile,
         )
         body = content.encode("utf-8")
+        check_budget()
         try:
             store.upload(req.output_path, body, "text/markdown; charset=utf-8")
         except StorageError as err: