diff --git a/backendN/app/common/app_factory.py b/backendN/app/common/app_factory.py index 08f92f8..a0c5bda 100644 --- a/backendN/app/common/app_factory.py +++ b/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 diff --git a/backendN/app/common/deadlines.py b/backendN/app/common/deadlines.py new file mode 100644 index 0000000..116ae31 --- /dev/null +++ b/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 diff --git a/backendN/app/common/security.py b/backendN/app/common/security.py index f1213db..bc36944 100644 --- a/backendN/app/common/security.py +++ b/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() diff --git a/backendN/app/common/storage.py b/backendN/app/common/storage.py index da6b48c..584e4d4 100644 --- a/backendN/app/common/storage.py +++ b/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() diff --git a/backendN/app/pdf_tasks.py b/backendN/app/pdf_tasks.py index 33c15df..badb284 100644 --- a/backendN/app/pdf_tasks.py +++ b/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,6 +239,7 @@ 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": @@ -242,6 +247,7 @@ def analyze(req: ProcessRequest) -> dict: 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"" 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: diff --git a/backendN/tests/test_deadlines.py b/backendN/tests/test_deadlines.py new file mode 100644 index 0000000..9c51eed --- /dev/null +++ b/backendN/tests/test_deadlines.py @@ -0,0 +1,136 @@ +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +import time +import threading +from types import SimpleNamespace +import uuid + +from fastapi.testclient import TestClient +import httpx +import pytest + +from app.common import deadlines as budget +from app.common.app_factory import create_app +from app.common.config import Settings +from app.common.storage import StorageClient +from conftest import SECRET + +HEADER = 'X-LLMiser-Deadline-Ms' +def headers(seconds=60): + return {'X-Internal-Secret':SECRET, HEADER:str(int((time.time()+seconds)*1000))} + +@contextmanager +def active(seconds=60): + value = budget.Budget(time.monotonic()+seconds) + token = budget._current.set(value) + try: yield value + finally: budget._current.reset(token) + +def expire(value): + object.__setattr__(value,'end',time.monotonic()-1) + +@pytest.mark.parametrize('raw',['','abc','-1','1.5','1e13','9'*17]) +def test_invalid_deadline_is_authenticated_before_validation(client,raw): + form = {'files':{'file':('a.txt',b'complete text')}} + response = client.post('/api/v1/internal/convert',headers={'X-Internal-Secret':'wrong',HEADER:raw},**form) + assert response.status_code == 401 + response = client.post('/api/v1/internal/convert',headers={'X-Internal-Secret':SECRET,HEADER:raw},**form) + assert response.status_code == 400 + +def test_expired_deadline_does_not_call_converter(client,monkeypatch): + monkeypatch.setattr(client.app.state.converter,'convert',lambda *a,**k:pytest.fail('expired parser started')) + result = client.post('/api/v1/internal/convert',headers=headers(-1),files={'file':('a.txt',b'text')}) + assert result.status_code == 504 + assert result.headers['X-LLMiser-Deadline-Exceeded'] == '1' + +def test_missing_deadline_and_valid_deadline_produce_identical_complete_results(client): + results = [client.post('/api/v1/internal/convert',headers=h,files={'file':('a.txt','complete café 日本語 🙂'.encode())}).json() for h in [{'X-Internal-Secret':SECRET},headers()]] + for key in ['content','filename','word_count','tokens_est']: + assert results[0][key] == results[1][key] + assert budget.current_budget() is None + +def test_parse_consuming_budget_cannot_write_or_report_success(): + writes=[] + def convert(*args,**kwargs): + assert budget.current_budget() is not None # ASGI → sync worker propagation + expire(budget.current_budget()) + return 'complete output',None + engine=SimpleNamespace(engine_name='test',extensions={'txt'},validate=lambda *a:None,convert=convert) + store=SimpleNamespace(download=lambda *a:b'text',upload=lambda *a:writes.append(a)) + settings=Settings(role='normal',instance='test',shared_secret=SECRET,max_upload_bytes=100,version='test') + client=TestClient(create_app(settings,engine,storage=store)) + job=str(uuid.uuid4()) + response=client.post('/api/v1/internal/process',headers=headers(),json={'job_id':job,'input_path':f'jobs/{job}/input/source.txt','output_path':f'jobs/{job}/output/result.md','original_filename':'a.txt'}) + assert response.status_code == 504 + assert writes == [] + +def test_upload_workers_inherit_budget_and_do_not_leak_it(): + with active() as value, ThreadPoolExecutor(max_workers=2) as pool: + assert list(pool.map(budget.budgeted_worker(lambda n:budget.current_budget()),range(6))) == [value]*6 + expire(value) + with pytest.raises(budget.DeadlineExceeded): + list(pool.map(budget.budgeted_worker(lambda n:pytest.fail('expired upload started')),range(6))) + assert budget.current_budget() is None + + +def test_concurrent_requests_have_independent_budgets(): + barrier=threading.Barrier(2) + seen=[] + def convert(*args,**kwargs): + value=budget.current_budget() + barrier.wait(timeout=5) + seen.append(value.remaining()) + return 'complete',None + engine=SimpleNamespace(engine_name='test',extensions={'txt'},validate=lambda *a:None,convert=convert) + settings=Settings(role='normal',instance='test',shared_secret=SECRET,max_upload_bytes=100,version='test') + client=TestClient(create_app(settings,engine)) + def post(seconds): + return client.post('/api/v1/internal/convert',headers=headers(seconds),files={'file':('a.txt',b'text')}).status_code + with ThreadPoolExecutor(max_workers=2) as pool: + assert list(pool.map(post,[60,120])) == [200,200] + assert min(seen) < 90 < max(seen) + assert budget.current_budget() is None + +def test_storage_timeout_shrinks_and_drip_read_checks_total_budget(): + with active(2) as value: + class Stream(httpx.SyncByteStream): + def __iter__(self): + yield b'first' + expire(value) + yield b'last' + def handler(request): + assert all(0 < seconds <= 2 for seconds in request.extensions['timeout'].values()) + return httpx.Response(200,stream=Stream()) + store=StorageClient('https://storage.test','test','test',transport=httpx.MockTransport(handler)) + with pytest.raises(budget.DeadlineExceeded): + store.download(f'jobs/{uuid.uuid4()}/input/source.txt',100) + + +def test_slow_upload_deadline_is_not_reported_as_storage_success(): + with active(60) as value: + def handler(request): + expire(value) + return httpx.Response(200,json={}) + store=StorageClient('https://storage.test','test','test',transport=httpx.MockTransport(handler)) + with pytest.raises(budget.DeadlineExceeded): + store.upload(f'jobs/{uuid.uuid4()}/output/result.md',b'complete','text/markdown') + + +def test_deadline_io_timeout_retains_deadline_marker_instead_of_peer_failover(): + with active(60) as value: + def handler(request): + expire(value) + raise httpx.ReadTimeout('generated deadline timeout',request=request) + store=StorageClient('https://storage.test','test','test',transport=httpx.MockTransport(handler)) + with pytest.raises(budget.DeadlineExceeded): + store.download(f'jobs/{uuid.uuid4()}/input/source.txt',100) + +def test_far_future_budget_is_clamped_and_duplicate_headers_rejected(client,monkeypatch): + def convert(*a,**k): + assert 0 < budget.current_budget().remaining() <= 300 + return 'complete',None + monkeypatch.setattr(client.app.state.converter,'convert',convert) + result=client.post('/api/v1/internal/convert',headers={'X-Internal-Secret':SECRET,HEADER:'9'*16},files={'file':('a.txt',b'text')}) + assert result.status_code == 200 + result=client.post('/api/v1/internal/convert',headers=[('X-Internal-Secret',SECRET),(HEADER,'1'),(HEADER,'2')],files={'file':('a.txt',b'text')}) + assert result.status_code == 400 diff --git a/backendN/tests/test_pdf_tasks.py b/backendN/tests/test_pdf_tasks.py index e15982a..7f01e72 100644 --- a/backendN/tests/test_pdf_tasks.py +++ b/backendN/tests/test_pdf_tasks.py @@ -65,6 +65,22 @@ def job_paths(): return job, f"jobs/{job}/input/source.pdf", f"jobs/{job}/output/result.md" +def test_deadline_after_pdf_classification_does_not_write_partial_success(client,storage,monkeypatch): + import time + from app.common.deadlines import current_budget + original = pdf_tasks.classify_pages + def classify(pdf): + pages = original(pdf) + object.__setattr__(current_budget(),'end',time.monotonic()-1) + return pages + monkeypatch.setattr(pdf_tasks,'classify_pages',classify) + job,src,out = job_paths() + storage.objects[src] = MIXED + response = client.post('/api/v1/internal/pdf/analyze',headers={**HEADERS,'X-LLMiser-Deadline-Ms':str(int((time.time()+60)*1000))},json={'job_id':job,'input_path':src,'output_path':out,'original_filename':'mixed.pdf'}) + assert response.status_code == 504 + assert set(storage.objects) == {src} + + def test_pages_are_classified_and_grouped(): pdf = pdfium.PdfDocument(MIXED) try: diff --git a/backendO/app/common/app_factory.py b/backendO/app/common/app_factory.py index 08f92f8..a0c5bda 100644 --- a/backendO/app/common/app_factory.py +++ b/backendO/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 diff --git a/backendO/app/common/deadlines.py b/backendO/app/common/deadlines.py new file mode 100644 index 0000000..116ae31 --- /dev/null +++ b/backendO/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 diff --git a/backendO/app/common/security.py b/backendO/app/common/security.py index f1213db..bc36944 100644 --- a/backendO/app/common/security.py +++ b/backendO/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() diff --git a/backendO/app/common/storage.py b/backendO/app/common/storage.py index da6b48c..584e4d4 100644 --- a/backendO/app/common/storage.py +++ b/backendO/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() diff --git a/backendO/app/converter.py b/backendO/app/converter.py index 318e972..d7f828f 100644 --- a/backendO/app/converter.py +++ b/backendO/app/converter.py @@ -22,6 +22,7 @@ from PIL import Image, ImageOps, JpegImagePlugin, UnidentifiedImageError from app.common.errors import ConversionRejected, EngineUnavailable +from app.common.deadlines import FINALIZE_RESERVE_S, compute_checkpoint, timeout_seconds PIL_FORMAT_BY_EXT = { "jpg": "JPEG", @@ -269,7 +270,7 @@ def _ocr(self, img: Image.Image, deadline: float) -> tuple[str, float, int]: img, lang=OCR_LANG, config=f"--oem 1 --psm {OCR_PSM}", - timeout=max(1, int(remaining)), + timeout=max(0.001, min(remaining, timeout_seconds(OCR_TIMEOUT_S, reserve=FINALIZE_RESERVE_S))), output_type=pytesseract.Output.DICT, ) return _text_from_data(data) @@ -287,7 +288,7 @@ def _osd_rotation(self, img: Image.Image) -> int | None: img, config="--psm 0 -c min_characters_to_try=5", output_type=pytesseract.Output.DICT, - timeout=10, + timeout=timeout_seconds(10, reserve=FINALIZE_RESERVE_S), ) except (pytesseract.TesseractError, RuntimeError): return None # too little text, or OSD timed out @@ -301,7 +302,7 @@ def _recognize(self, img: Image.Image) -> tuple[str, int]: decide, the other three orientations) is kept only if it reads with clearly higher confidence. """ - deadline = time.monotonic() + OCR_TIMEOUT_S + deadline = time.monotonic() + timeout_seconds(OCR_TIMEOUT_S, reserve=FINALIZE_RESERVE_S) best_text, best_conf, best_words = self._ocr(img, deadline) best_rotation = 0 proposed = None @@ -315,6 +316,7 @@ def _recognize(self, img: Image.Image) -> tuple[str, int]: candidates = [90, 180, 270] for rotation in candidates: + compute_checkpoint() if time.monotonic() >= deadline: break text, conf, words = self._ocr(_rotate_upright(img, rotation), deadline) @@ -338,6 +340,7 @@ def convert(self, data: bytes, ext: str, *, filename: str, title: str) -> tuple[ # Hold capacity before decoding: otherwise waiting requests retain # their own full-size prepared images while another job runs OCR. with _OCR_LOCK: + compute_checkpoint() img = _open_checked(data, ext) try: frames = getattr(img, "n_frames", 1) @@ -349,16 +352,19 @@ def convert(self, data: bytes, ext: str, *, filename: str, title: str) -> tuple[ img.close() try: + compute_checkpoint() raw, rotation = self._recognize(prepared) except pytesseract.TesseractNotFoundError as err: raise EngineUnavailable(str(err)) from None except RuntimeError as err: + compute_checkpoint() if "timeout" in str(err).lower(): raise ConversionRejected(422, "OCR took too long for this image", "ocr_timeout") from None raise finally: prepared.close() + compute_checkpoint() text = _clean_text(raw) if not text: return "_No text detected in this image._", "no-text-detected" diff --git a/backendO/tests/test_deadlines.py b/backendO/tests/test_deadlines.py new file mode 100644 index 0000000..a96c878 --- /dev/null +++ b/backendO/tests/test_deadlines.py @@ -0,0 +1,47 @@ +from contextlib import contextmanager +import time +import pytest +from app import converter +from app.common import deadlines +from conftest import image_bytes + +@contextmanager +def active(seconds): + value=deadlines.Budget(time.monotonic()+seconds) + token=deadlines._current.set(value) + try: yield value + finally: deadlines._current.reset(token) + +def test_waiting_expired_ocr_does_not_decode(monkeypatch): + monkeypatch.setattr(converter,'_open_checked',lambda *a:pytest.fail('expired OCR decoded image')) + with active(-1),pytest.raises(deadlines.DeadlineExceeded): + converter.OcrConverter().convert(image_bytes('PNG'),'png',filename='a.png',title='a') + +def test_tesseract_timeout_preserves_result_save_reserve(monkeypatch): + import pytesseract + calls=[] + def image_to_data(image,**kwargs): + calls.append(kwargs['timeout']) + return {'text':[],'conf':[],'block_num':[],'par_num':[],'line_num':[]} + monkeypatch.setattr(pytesseract,'image_to_data',image_to_data) + with active(15): + converter.OcrConverter()._ocr(None,time.monotonic()+85) + assert 0 < calls[0] <= 5 + +def test_ocr_budget_timeout_is_504_and_prepared_pixels_close(monkeypatch): + prepared=[] + original=converter._prepare + def prepare(source): + image=original(source) + prepared.append(image) + return image + def recognize(*args): + object.__setattr__(deadlines.current_budget(),'end',time.monotonic()-1) + raise RuntimeError('Tesseract process timeout') + engine=converter.OcrConverter() + monkeypatch.setattr(converter,'_prepare',prepare) + monkeypatch.setattr(engine,'_recognize',recognize) + with active(60),pytest.raises(deadlines.DeadlineExceeded): + engine.convert(image_bytes('PNG'),'png',filename='a.png',title='a') + with pytest.raises(ValueError,match='closed'): + prepared[0].getpixel((0,0)) diff --git a/backendZ/app/archive.py b/backendZ/app/archive.py index 32fac8e..3892156 100644 --- a/backendZ/app/archive.py +++ b/backendZ/app/archive.py @@ -23,6 +23,7 @@ from dataclasses import asdict, dataclass, field from app.common.errors import ConversionRejected +from app.common.deadlines import DeadlineExceeded, compute_checkpoint def _env_int(name: str, default: int) -> int: @@ -221,6 +222,7 @@ def process_archive(data: bytes, convert_document, *, ocr_available: bool) -> li def _walk(archive, infos, prefix, depth, state, out, convert_document, ocr_available): for info in sorted(infos, key=lambda i: i.filename.casefold()): + compute_checkpoint() state.entries_seen += 1 rel = normalize(info.filename) if rel is None: @@ -312,6 +314,7 @@ def _walk(archive, infos, prefix, depth, state, out, convert_document, ocr_avail continue text = _convert(entry, raw, convert_document) + compute_checkpoint() if text is None: continue size = len(text.encode("utf-8")) @@ -327,6 +330,8 @@ def _convert(entry: Entry, raw: bytes, convert_document) -> str | None: if entry.node_type == "DOCUMENT": try: markdown = (convert_document(raw, ext, entry.path.rsplit("/", 1)[-1]) or "").strip() + except DeadlineExceeded: + raise except Exception as err: # noqa: BLE001 - one bad file must not fail the archive entry.status, entry.skip_reason = "FAILED", f"could not be converted ({type(err).__name__})" return None diff --git a/backendZ/app/archive_tasks.py b/backendZ/app/archive_tasks.py index 73453ba..3014988 100644 --- a/backendZ/app/archive_tasks.py +++ b/backendZ/app/archive_tasks.py @@ -32,6 +32,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 DeadlineExceeded, budgeted_worker, check_budget, compute_checkpoint from app.common.storage import StorageError logger = logging.getLogger("mdify.archive") @@ -115,6 +116,7 @@ def partial_blobs(entries: list[archive.Entry]) -> list[bytes]: def write_outputs(store, *, job_id: str, name: str, profile: str, source_bytes: int, entries: list[archive.Entry], ocr_texts: dict[str, str | None], settings) -> dict: + compute_checkpoint() if len(entries) > archive.MAX_ENTRIES: raise resource_limit() stem, _ext = split_name(name) @@ -154,10 +156,12 @@ def write_outputs(store, *, job_id: str, name: str, profile: str, source_bytes: raise resource_limit() try: + check_budget() with ThreadPoolExecutor(max_workers=WORKERS) as pool: # Encoding inside the worker bounds live application byte buffers # to WORKERS objects instead of retaining every encoded part. - list(pool.map(lambda f: store.upload(f[0], f[1].encode("utf-8"), f[2]), files)) + list(pool.map(budgeted_worker(lambda f: store.upload(f[0], f[1].encode("utf-8"), f[2])), files)) + check_budget() except StorageError as err: raise HTTPException(status_code=503 if err.transient else 502, detail="Could not store the project result") from None @@ -208,6 +212,8 @@ def process(req: ProcessRequest) -> dict: raise HTTPException(status_code=err.status_code, detail=err.detail) from None except EngineUnavailable: raise HTTPException(status_code=503, detail="Conversion engine unavailable") from None + except DeadlineExceeded: + raise except Exception: logger.exception("archive processing failed job=%s", job_id) raise HTTPException(status_code=422, detail="This ZIP file couldn't be processed.") from None @@ -244,6 +250,7 @@ def process(req: ProcessRequest) -> dict: # Converted entries wait in Storage for the merge, in parts <= 4 MB. parts = partial_blobs(entries) + check_budget() partial_paths = [f"jobs/{job_id}/nodes/archive-{i}/partial.json" for i in range(1, len(parts) + 1)] uploads += [ (path, chunk, "application/json; charset=utf-8") @@ -251,7 +258,8 @@ def process(req: ProcessRequest) -> dict: ] try: with ThreadPoolExecutor(max_workers=WORKERS) as pool: - list(pool.map(lambda u: store.upload(*u), uploads)) + list(pool.map(budgeted_worker(lambda u: store.upload(*u)), uploads)) + check_budget() except StorageError as err: raise HTTPException(status_code=503 if err.transient else 502, detail="Could not store the archive contents") from None @@ -281,6 +289,7 @@ def merge(req: MergeRequest) -> dict: def fetch(path: str) -> str | None: nonlocal read_bytes + compute_checkpoint() try: data = store.download(path, PARTIAL_BYTES) except StorageError as err: diff --git a/backendZ/app/common/app_factory.py b/backendZ/app/common/app_factory.py index 08f92f8..a0c5bda 100644 --- a/backendZ/app/common/app_factory.py +++ b/backendZ/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 diff --git a/backendZ/app/common/deadlines.py b/backendZ/app/common/deadlines.py new file mode 100644 index 0000000..116ae31 --- /dev/null +++ b/backendZ/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 diff --git a/backendZ/app/common/security.py b/backendZ/app/common/security.py index f1213db..bc36944 100644 --- a/backendZ/app/common/security.py +++ b/backendZ/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() diff --git a/backendZ/app/common/storage.py b/backendZ/app/common/storage.py index da6b48c..584e4d4 100644 --- a/backendZ/app/common/storage.py +++ b/backendZ/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() diff --git a/backendZ/tests/test_deadlines.py b/backendZ/tests/test_deadlines.py new file mode 100644 index 0000000..890e1c8 --- /dev/null +++ b/backendZ/tests/test_deadlines.py @@ -0,0 +1,31 @@ +import time +import pytest +from app import archive +from app.common.deadlines import current_budget +from test_archive_tasks import HEADERS,PROCESS,process_body,upload + +def test_archive_expiry_is_not_hidden_as_one_bad_document(client,storage,converter,monkeypatch): + original = converter.documents.__class__.__call__ + calls=[] + def convert(self,*args): + calls.append(True) + text=original(self,*args) + object.__setattr__(current_budget(),'end',time.monotonic()-1) + return text + monkeypatch.setattr(converter.documents.__class__,'__call__',convert) + job=upload(storage,{'first.html':'