From 175ed7cb067b860cee6212b37f13e9041fc4970d Mon Sep 17 00:00:00 2001 From: sirily11 <32106111+sirily11@users.noreply.github.com> Date: Sun, 19 Jul 2026 21:30:58 +0800 Subject: [PATCH] feat: move to mq for markdown parsing --- .env.example | 82 +++++++ .github/workflows/release.yaml | 11 +- .github/workflows/test.yaml | 59 +++-- .gitignore | 12 +- Dockerfile | 3 +- Makefile | 36 ++- README.md | 133 ++++++++++- api/index.py | 415 ++++++++++++++++++++++++++++++--- cache.py | 149 ------------ chunker.py | 80 +++++++ config.py | 107 +++++++++ converter.py | 118 +++++----- db.py | 285 ++++++++++++++++++++++ docker-compose.yml | 35 +++ jobstore.py | 273 ++++++++++++++++++++++ k8s/api-deployment.yaml | 100 ++++++++ k8s/deployment.yaml | 84 ------- k8s/hpa.yaml | 60 ++++- k8s/ingress.yaml | 16 +- k8s/kustomization.yaml | 10 +- k8s/libsql.yaml | 100 ++++++++ k8s/minio.yaml | 143 ++++++++++++ k8s/rabbitmq.yaml | 108 +++++++++ k8s/redis.yaml | 62 ----- k8s/secrets.example.yaml | 45 ++++ k8s/worker-deployment.yaml | 95 ++++++++ migrations/001_init.sql | 61 +++++ pagination.py | 62 +++++ pdf_chunk.py | 134 ----------- pyproject.toml | 20 +- requirements.txt | 14 +- scripts/k8s-benchmark.sh | 149 +++++++++--- scripts/local-smoke.sh | 114 +++++++++ storage.py | 240 +++++++++++++++++++ taskqueue.py | 269 +++++++++++++++++++++ test_main.http | 39 +++- tests/conftest.py | 128 ++++++++++ tests/test_api.py | 324 +++++++++++++++++++++++++ tests/test_chunker.py | 108 +++++++++ tests/test_converter.py | 35 +++ tests/test_db.py | 175 ++++++++++++++ tests/test_integration.py | 113 +++++++++ tests/test_jobstore.py | 168 +++++++++++++ tests/test_pdf_chunk.py | 158 ------------- tests/test_storage.py | 108 +++++++++ tests/test_taskqueue.py | 145 ++++++++++++ tests/test_worker.py | 290 +++++++++++++++++++++++ worker.py | 281 ++++++++++++++++++++++ 48 files changed, 4981 insertions(+), 775 deletions(-) create mode 100644 .env.example delete mode 100644 cache.py create mode 100644 chunker.py create mode 100644 config.py create mode 100644 db.py create mode 100644 docker-compose.yml create mode 100644 jobstore.py create mode 100644 k8s/api-deployment.yaml delete mode 100644 k8s/deployment.yaml create mode 100644 k8s/libsql.yaml create mode 100644 k8s/minio.yaml create mode 100644 k8s/rabbitmq.yaml delete mode 100644 k8s/redis.yaml create mode 100644 k8s/secrets.example.yaml create mode 100644 k8s/worker-deployment.yaml create mode 100644 migrations/001_init.sql create mode 100644 pagination.py delete mode 100644 pdf_chunk.py create mode 100755 scripts/local-smoke.sh create mode 100644 storage.py create mode 100644 taskqueue.py create mode 100644 tests/test_api.py create mode 100644 tests/test_chunker.py create mode 100644 tests/test_converter.py create mode 100644 tests/test_db.py create mode 100644 tests/test_integration.py create mode 100644 tests/test_jobstore.py delete mode 100644 tests/test_pdf_chunk.py create mode 100644 tests/test_storage.py create mode 100644 tests/test_taskqueue.py create mode 100644 tests/test_worker.py create mode 100644 worker.py diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..84f3907 --- /dev/null +++ b/.env.example @@ -0,0 +1,82 @@ +# Copy to .env and fill in. Defaults for everything here live in config.py; +# .env only needs the values you actually want to override. +# +# The Makefile loads .env and supplies sensible local defaults for anything it +# does not find, so a minimal .env is just OPENAI_API_KEY. + +# ---- Conversion ---- +OPENAI_API_KEY= +OPENAI_MODEL=google/gemini-3.1-flash-lite-preview +# LLM-assisted conversion (image descriptions). Slow and costs tokens per page. +# Folded into the storage cache key, so changing it re-converts rather than +# serving the previous setting's output. +CONVERT_USE_LLM=false + +# ---- API ---- +ADMIN_API_KEY=dev-key + +# ---- Database (libSQL / Turso) ---- +# db.py picks the driver from this value: +# a path or file: URL -> stdlib sqlite3 (local single-process, and CI) +# libsql:// or http:// -> libsql driver (hosted Turso, or self-hosted sqld) +# +# The API and workers are separate processes and must share one database, so a +# local SQLite file only works if you run a single process. +TURSO_DATABASE_URL=http://localhost:8080 +# Hosted Turso only. Setting this forces the libsql driver regardless of URL. +TURSO_AUTH_TOKEN= +LIBSQL_PORT=8080 + +# ---- Object storage ---- +# Any S3-compatible endpoint: minio locally, R2 or AWS in production. +# Ports default off 9000/9001 since those commonly collide with other projects. +S3_ENDPOINT=http://localhost:9100 +S3_BUCKET=markitdown +S3_ACCESS_KEY_ID=minioadmin +S3_SECRET_ACCESS_KEY=minioadmin +S3_REGION=us-east-1 +MINIO_PORT=9100 +MINIO_CONSOLE_PORT=9101 +# Recorded on the document row. Actual expiry is a bucket lifecycle rule, so +# every object for a document ages out together. +DOC_TTL_DAYS=7 + +# ---- Chunking ---- +# Smaller chunks lower peak memory per worker but add per-chunk parser startup. +PDF_PAGES_PER_CHUNK=20 +# Below this page count a PDF is converted whole rather than split. +PDF_MIN_PAGES_FOR_PARALLEL=40 +PAGE_SIZE=5000 + +# ---- Queue (RabbitMQ) ---- +# Job distribution. The producer publishes one message per chunk; workers +# consume with prefetch=1, so the broker hands each chunk to whoever is free. +# Adding a worker pod adds throughput with no coordination. +# +# Port 5673 locally to avoid colliding with an existing 5672 broker. +RABBITMQ_URL=amqp://guest:guest@localhost:5673/%2F +RABBITMQ_QUEUE=markitdown.chunks +RABBITMQ_EXCHANGE=markitdown +RABBITMQ_PORT=5673 +RABBITMQ_MANAGEMENT_PORT=15673 +RABBITMQ_USER=guest +RABBITMQ_PASSWORD=guest +# Unacked messages a worker holds at once. 1 is what bounds memory to a chunk +# rather than a document, and keeps dispatch fair. +RABBITMQ_PREFETCH=1 + +# Retry backoff tiers in seconds, one dedicated delay queue each; attempt N +# waits tier N, then dead-letters back onto the work queue. +# +# Per-tier queues rather than a per-message TTL on one queue: RabbitMQ only +# expires the message at the *head*, so a 300s message parked in front would +# hold back a 5s message queued behind it. +RETRY_DELAYS=5,30,300 +# Attempts per chunk. Exhausting them parks the message on +# markitdown.chunks.failed and fails the whole job — a document silently +# missing pages is worse than an error. +MAX_ATTEMPTS=3 + +# ---- Worker ---- +SOURCE_CACHE_DIR=/tmp/markitdown-src +SOURCE_CACHE_MAX_BYTES=2147483648 diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 87d7bb4..72e29cf 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -100,12 +100,17 @@ jobs: kubectl kustomize . | kubectl apply -f - # Update the deployment with the specific release tag - kubectl set image deployment/markitdown-server markitdown-server=ghcr.io/${{ github.repository_owner }}/markitdown-server:${VERSION} -n markitdown-server + # One image, two deployments: the producer (API) and the consumer + # (worker) must move together or a new worker could read job rows + # written by an old producer. + kubectl set image deployment/markitdown-api markitdown-api=ghcr.io/${{ github.repository_owner }}/markitdown-server:${VERSION} -n markitdown-server + kubectl set image deployment/markitdown-worker markitdown-worker=ghcr.io/${{ github.repository_owner }}/markitdown-server:${VERSION} -n markitdown-server # Wait for rollout to complete - kubectl rollout status deployment/markitdown-server -n markitdown-server --timeout=300s + kubectl rollout status deployment/markitdown-api -n markitdown-server --timeout=300s + kubectl rollout status deployment/markitdown-worker -n markitdown-server --timeout=300s - name: Verify deployment run: | - kubectl get pods -n markitdown-server -l app=markitdown-server + kubectl get pods -n markitdown-server kubectl get service -n markitdown-server markitdown-server diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index a0f72b2..feee1e4 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -29,12 +29,14 @@ jobs: - name: Install Python dependencies run: | pip install -r requirements.txt - pip install pytest reportlab + pip install pytest reportlab httpx "moto[s3]" + # No database or object-store service container needed: the tests run + # against stdlib sqlite3 and an in-process moto S3. db.py picks the + # sqlite driver for any non-server DATABASE_URL, so the same code paths + # are exercised here as in the cluster. - name: Run unit tests - env: - PYTHONPATH: ${{ github.workspace }} - run: pytest tests/ -v + run: pytest -v k8s-benchmark: name: PDF Benchmark (Kubernetes) @@ -77,12 +79,21 @@ jobs: - name: Create secret run: | kubectl create namespace "$NAMESPACE" --dry-run=client -o yaml | kubectl apply -f - - # PDF conversion never calls the LLM (that path is images only), so a - # placeholder key is enough to satisfy the import-time lookup. + # Conversion only calls the LLM when use_llm is requested, which the + # benchmark does not do, so a placeholder key is enough. + # RABBITMQ_URL points at the in-cluster broker from rabbitmq.yaml: + # the runner has no route to the shared linda-namespace broker the + # real secret uses. RABBITMQ_USER/PASSWORD are what that broker boots + # with, so the two must agree. kubectl create secret generic markitdown-server-secret \ -n "$NAMESPACE" \ --from-literal=OPENAI_API_KEY=ci-placeholder \ --from-literal=ADMIN_API_KEY=ci-test-key \ + --from-literal=S3_ACCESS_KEY_ID=ci-minio-user \ + --from-literal=S3_SECRET_ACCESS_KEY=ci-minio-password \ + --from-literal=RABBITMQ_USER=ci-rabbit-user \ + --from-literal=RABBITMQ_PASSWORD=ci-rabbit-password \ + --from-literal=RABBITMQ_URL='amqp://ci-rabbit-user:ci-rabbit-password@markitdown-rabbitmq:5672/%2F' \ --dry-run=client -o yaml | kubectl apply -f - - name: Deploy @@ -91,18 +102,24 @@ jobs: kustomize edit set image \ "ghcr.io/sirily11/markitdown-server=markitdown-server:${IMAGE_TAG}" kubectl apply -k . - # Drop the HPA for the benchmark. Conversion pegs the CPU, so the HPA - # scales to maxReplicas mid-run and then holds there for its 300s - # scale-down window. Two replicas behind the service make per-book - # timing and peak-memory attribution meaningless, since requests and - # the cgroup reads can land on different pods. - kubectl delete hpa markitdown-server-hpa -n "$NAMESPACE" --ignore-not-found - kubectl scale deployment/markitdown-server -n "$NAMESPACE" --replicas=1 + # Drop the worker HPA and pin the worker to one replica. Conversion + # pegs the CPU, so the HPA would scale up mid-run and hold there for + # its scale-down window; with a worker fleet, peak memory is spread + # across pods and no single cgroup reading means anything. + kubectl delete hpa markitdown-worker-hpa -n "$NAMESPACE" --ignore-not-found + kubectl delete hpa markitdown-api-hpa -n "$NAMESPACE" --ignore-not-found + kubectl scale deployment/markitdown-worker -n "$NAMESPACE" --replicas=1 - name: Wait for rollout run: | - kubectl rollout status deployment/markitdown-redis -n "$NAMESPACE" --timeout=300s || true - kubectl rollout status deployment/markitdown-server -n "$NAMESPACE" --timeout=300s + # State first: the API and worker both fail readiness until these are + # up, so rolling them out in order keeps the failure legible. + kubectl rollout status deployment/markitdown-libsql -n "$NAMESPACE" --timeout=300s + kubectl rollout status deployment/markitdown-minio -n "$NAMESPACE" --timeout=300s + kubectl wait --for=condition=complete job/markitdown-minio-setup \ + -n "$NAMESPACE" --timeout=300s + kubectl rollout status deployment/markitdown-api -n "$NAMESPACE" --timeout=300s + kubectl rollout status deployment/markitdown-worker -n "$NAMESPACE" --timeout=300s # No port-forward here on purpose: the benchmark restarts the pod between # books, which tears a forward down, so the script manages its own. @@ -167,10 +184,16 @@ jobs: # Per-book logs captured during the run. The pod is restarted between # books, so these are the only copy of an earlier book's output. cp -r /tmp/benchmark-pod-logs /tmp/k8s-logs/per-book 2>/dev/null || true - kubectl logs -n "$NAMESPACE" -l app=markitdown-server --tail=-1 \ - > /tmp/k8s-logs/server-final.log 2>&1 || true + kubectl logs -n "$NAMESPACE" -l app=markitdown-api --tail=-1 \ + > /tmp/k8s-logs/api-final.log 2>&1 || true + kubectl logs -n "$NAMESPACE" -l app=markitdown-worker --tail=-1 \ + > /tmp/k8s-logs/worker-final.log 2>&1 || true + kubectl logs -n "$NAMESPACE" -l app=markitdown-libsql --tail=200 \ + > /tmp/k8s-logs/libsql.log 2>&1 || true + kubectl logs -n "$NAMESPACE" -l app=markitdown-minio --tail=200 \ + > /tmp/k8s-logs/minio.log 2>&1 || true kubectl get pods -n "$NAMESPACE" -o wide > /tmp/k8s-logs/pods.txt 2>&1 || true - kubectl describe pods -n "$NAMESPACE" -l app=markitdown-server \ + kubectl describe pods -n "$NAMESPACE" \ > /tmp/k8s-logs/describe.txt 2>&1 || true - name: Upload diagnostics diff --git a/.gitignore b/.gitignore index 87746c2..20d78c4 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,14 @@ __pycache__ .DS_Store .vercel -secrets.yaml \ No newline at end of file +# Real credentials. The trailing * also covers editor/backup spill +# (secrets.yaml.bak, .orig, .save); secrets.example.yaml is not matched and +# stays tracked as the template. +secrets.yaml* +.venv +venv +.pytest_cache +# Local sqlite database, used when TURSO_DATABASE_URL points at a file. +*.db +*.db-wal +*.db-shm diff --git a/Dockerfile b/Dockerfile index a67b8da..7af507f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -28,5 +28,6 @@ EXPOSE 8000 ENV PYTHONPATH=/app ENV PYTHONUNBUFFERED=1 -# Run the server +# One image, two roles. The default is the API (producer); worker pods override +# the command with ["python", "-m", "worker"] to run the consumer instead. CMD ["uvicorn", "api.index:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/Makefile b/Makefile index 7b6f3f9..f12a17e 100644 --- a/Makefile +++ b/Makefile @@ -1,3 +1,37 @@ +# Local development. +# +# make up rabbitmq + minio +# make api the producer (uvicorn, with reload) +# make worker a consumer — run several to add capacity +# make test the test suite +# +# The database is remote (Turso); only the broker and object storage run here. + +ifneq (,$(wildcard .env)) +-include .env +export +endif + +PY := $(if $(wildcard .venv/bin/python),.venv/bin/python,python3) +PORT ?= 8000 + +.PHONY: up down api worker test docker + +up: + docker compose up -d + +down: + docker compose down + +api: + $(PY) -m uvicorn api.index:app --reload --host 0.0.0.0 --port $(PORT) + +worker: + $(PY) -m worker + +test: + $(PY) -m pytest -q + docker: docker buildx create --use --name multi-arch-builder || true docker buildx build \ @@ -8,4 +42,4 @@ docker: --build-arg BUILD_TIME=$$(date -u +'%Y-%m-%d_%H:%M:%S') \ --build-arg COMMIT_HASH=$$(git rev-parse --short HEAD) \ --push \ - . \ No newline at end of file + . diff --git a/README.md b/README.md index 297c6a3..e0e9f02 100644 --- a/README.md +++ b/README.md @@ -1 +1,132 @@ -# Markitdown Server \ No newline at end of file +# markitdown-server + +Converts documents at a URL to markdown, using a producer/consumer queue so +large PDFs can be split across worker pods. The original blocking API remains +available, with a separate endpoint for asynchronous callers. + +## How it works + +``` +POST /convert or /async/convert + │ download → sha256(bytes) → doc_key = file_hash:config_fingerprint + │ + ├─ document already converted? ──► return it. No work queued. + │ + └─ upload source to S3, plan page ranges, insert job + N chunk rows + │ + ┌───────────────┼───────────────┐ + worker A worker B worker C + claim a chunk (atomic UPDATE), extract its page range, + convert it, write the markdown to S3 + │ + last chunk to finish assembles the document, + paginates it, and marks the job done +``` + +Two processes share one image: the **API** (producer) and the **worker** +(consumer, `python -m worker`). Scale by adding worker pods — they coordinate +only through the database, so there is no leader and nothing to reconfigure. + +### The queue is a table + +`job_chunks` *is* the queue. It already had to track chunk state for progress +and retries, so making it the queue too means one store to keep consistent +rather than two that can disagree. Claiming is a single atomic statement, and +because SQLite serialises writers, two workers cannot claim the same row. + +Everything a dedicated queue would provide falls out of three columns: + +| Column | Provides | +|---|---| +| `attempts` | retry limit; incremented at *claim* time so an OOM-killed worker still burns an attempt | +| `available_at` | exponential backoff | +| `lease_expires_at` | crash recovery — a lease that stops being renewed is a worker that died | + +Workers heartbeat to extend their lease while converting, so the lease means +"the owner is gone" rather than "this is taking a while" — decoupled from how +long a chunk legitimately takes. + +### Caching + +Storage keys are content-addressed: `sha256(file bytes)` plus a fingerprint of +everything that changes the output (markitdown version, LLM flag, model, page +size, chunk size). Re-submitting a file returns instantly with no conversion +and no LLM spend, and upgrading markitdown invalidates the cache automatically +rather than serving stale markdown forever. + +## API + +| Endpoint | Purpose | +|---|---| +| `POST /convert` | Convert a URL synchronously. Returns the original `{id, content, pagination}` payload after all queued chunks finish. | +| `POST /async/convert` | Submit a URL asynchronously. Returns `202 {job_id, doc_key}`, or `200` with the first page on a cache hit. | +| `GET /convert/jobs/{job_id}` | Job status and chunk progress. | +| `GET /convert/{doc_key}/pages/{page}` | A page of the converted document. | +| `GET /` | Liveness. Checks nothing external. | +| `GET /readyz` | Readiness. Checks the database and object store. | + +All except `/` and `/readyz` require an `x-api-key` header matching +`ADMIN_API_KEY`. + +## Configuration + +See `config.py` — every environment variable is defined there with its default. +The ones that matter most: + +| Variable | Default | Notes | +|---|---|---| +| `TURSO_DATABASE_URL` | `file:markitdown.db` | A path uses stdlib sqlite3; `libsql://` or `http://` uses a libSQL server | +| `TURSO_AUTH_TOKEN` | — | Set for hosted Turso | +| `S3_ENDPOINT` / `S3_BUCKET` | — / `markitdown` | Any S3-compatible store | +| `CONVERT_USE_LLM` | `false` | LLM-assisted conversion (image descriptions); slow and costs tokens per page | +| `PDF_PAGES_PER_CHUNK` | `20` | Smaller chunks lower peak memory per worker | +| `MAX_ATTEMPTS` | `3` | Per chunk; exhausting it fails the whole job | +| `LEASE_TTL` / `HEARTBEAT_INTERVAL` | `120` / `30` | Lease must comfortably exceed the heartbeat | + +## Development + +```bash +python -m venv .venv && .venv/bin/pip install -r requirements.txt +.venv/bin/pip install pytest reportlab httpx "moto[s3]" +.venv/bin/python -m pytest +``` + +Tests run against stdlib sqlite3 and an in-process moto S3 — no service +containers, no network. `tests/test_integration.py` is the slow one: it does a +real PDF over real HTTP through real markitdown conversion. + +### Running locally + +`docker-compose.yml` brings up the two dependencies — libSQL (shared database) +and minio (object storage). There is no broker: the queue is a table. The API +and workers run on the host so you keep reload and a debugger. + +```bash +make venv # once +make up # libsql + minio, bucket created with the 7-day expiry rule +make api # producer, in another shell +make worker # consumer — run this several times to add workers +make smoke # submit a real PDF, poll it, read every page back +``` + +`make smoke` is the end-to-end check: it converts a real Asimov PDF, walks all +pages back out, and asserts a resubmission hits the content-addressed cache. +Adding worker shells is how you watch chunks distribute — the Foundation +Trilogy splits into 25 chunks and finishes in ~12s across three workers. + +Ports default to 8080 (libsql) and 9100/9101 (minio) to avoid the usual +collisions; override with `LIBSQL_PORT` / `MINIO_PORT` in `.env`. `make up` +reads `.env`, and defaults everything it does not find, so `.env` only needs +your real secrets. `make reset` wipes the volumes. + +## Deployment + +`k8s/` is a kustomize stack: an API deployment, a worker deployment (HPA'd on +CPU), plus self-hosted libSQL and minio for state. To run fully managed, drop +`libsql.yaml` and `minio.yaml` from `kustomization.yaml` and point the secret at +hosted Turso and R2/S3 — note that the inline `env` entries in the deployments +take precedence over `envFrom`, so those must be removed too. + +`scripts/k8s-benchmark.sh` converts four real Asimov PDFs against a cluster and +enforces per-book time and memory budgets; it is the performance regression +guard and runs in CI on every push. diff --git a/api/index.py b/api/index.py index 74326ec..d8bfdbf 100644 --- a/api/index.py +++ b/api/index.py @@ -1,11 +1,24 @@ -from fastapi import FastAPI, HTTPException, Depends, Header -from fastapi.concurrency import run_in_threadpool -from converter import convert -from cache import store_document, get_page, PAGE_SIZE -from pydantic import BaseModel -from typing import Optional +import asyncio import logging import os +import time +from contextlib import asynccontextmanager +from typing import Optional + +import yaml +from fastapi import Depends, FastAPI, HTTPException, Response +from fastapi.concurrency import run_in_threadpool +from fastapi.security import APIKeyHeader +from pydantic import BaseModel, Field + +import chunker +import config +import converter +import db +import jobstore +import pagination +import storage +import taskqueue logging.basicConfig( level=os.environ.get("LOG_LEVEL", "INFO"), @@ -13,15 +26,114 @@ ) logger = logging.getLogger(__name__) -app = FastAPI() + +@asynccontextmanager +async def lifespan(_app: FastAPI): + # Migrations are idempotent (IF NOT EXISTS), so every API pod can run them + # on boot. Object-storage provisioning belongs to deployment setup, not the + # request-serving process; readiness reports when that dependency is down. + await run_in_threadpool(db.migrate) + yield + + +app = FastAPI( + lifespan=lifespan, + title="markitdown-server", + version="1.0.0", + description=( + "Convert documents at a URL to paginated markdown.\n\n" + "`POST /convert` preserves the synchronous API and returns only after " + "conversion finishes. `POST /async/convert` returns a `job_id` to poll " + "at `GET /convert/jobs/{job_id}`. Both routes use the same distributed " + "worker queue, and finished pages are read from " + "`GET /convert/{doc_key}/pages/{page}`. Conversion is content-addressed " + "on the source bytes plus `use_llm`, so re-submitting an identical file " + "skips the queue and returns its first page directly." + ), +) class ConvertRequest(BaseModel): - file: Optional[str] = None + file: Optional[str] = Field( + None, description="URL of the document to convert. Fetched by the server.", + examples=["https://example.com/book.pdf"], + ) + # Whether to run markitdown's LLM-assisted conversion. Folded into the + # cache key, so LLM and non-LLM results are stored separately. + use_llm: Optional[bool] = Field( + None, + description=( + "Run markitdown's LLM-assisted conversion (image descriptions). " + "Part of the cache key, so results do not collide with non-LLM " + "ones. Defaults to the server's CONVERT_USE_LLM." + ), + ) + + +class Pagination(BaseModel): + id: str + page: int + page_size: int + total_pages: int + total_length: int + has_next: bool + has_prev: bool + next_page: Optional[int] = Field(None, description="null on the last page.") + prev_page: Optional[int] = Field(None, description="null on the first page.") + + +class PageResponse(BaseModel): + """One page of converted markdown. Also the 200 body of `POST /convert`.""" + id: str = Field(description="The document's `doc_key`.") + content: str + pagination: Pagination + + +class JobAccepted(BaseModel): + """202 body of `POST /async/convert`: work was queued.""" + job_id: str + status: str = Field(examples=["queued", "running"]) + doc_key: str + total_chunks: int + + +class JobProgress(BaseModel): + completed: int + total: int + percent: int + + +class JobStatus(BaseModel): + job_id: str + status: str = Field( + description="queued -> running -> done | failed.", + examples=["queued", "running", "done", "failed"], + ) + doc_key: str + progress: JobProgress + error: Optional[str] = Field(None, description="Set only when status is failed.") + created_at: int = Field(description="Unix epoch seconds.") + updated_at: int = Field(description="Unix epoch seconds.") + + +class Health(BaseModel): + status: str = Field(examples=["ok"]) -async def verify_api_key(x_api_key: str = Header(None)): - admin_api_key = os.environ.get("ADMIN_API_KEY") +class Readiness(BaseModel): + status: str = Field(examples=["ok", "degraded"]) + checks: dict[str, bool] = Field( + description="Per-dependency health: database, storage, queue." + ) + + +# auto_error=False so a missing header falls through to the check below and +# yields 401 rather than the 403 the security scheme would raise on its own. +api_key_header = APIKeyHeader(name="x-api-key", auto_error=False) + + +async def verify_api_key(x_api_key: str = Depends(api_key_header)): + admin_api_key = config.ADMIN_API_KEY if not admin_api_key: raise HTTPException(status_code=500, detail="API key not configured on server") @@ -31,47 +143,288 @@ async def verify_api_key(x_api_key: str = Header(None)): return x_api_key -@app.post("/convert", dependencies=[Depends(verify_api_key)]) -async def convert_endpoint( - request: Optional[ConvertRequest] = None, -): +def _first_page_payload(doc_key: str, manifest: dict) -> Optional[dict]: + content = storage.get_page(doc_key, 1) + if content is None: + return None + return { + "id": doc_key, + "content": content, + "pagination": pagination.build_pagination( + doc_key, 1, manifest["total_pages"], + manifest["total_length"], manifest["page_size"], + ), + } + + +def submit(url: str, use_llm: bool) -> dict: """ - Convert a document at the given URL to markdown. + Producer: hash the source, then either serve it from cache or enqueue work. - The markdown is paginated and stored in Redis (5 min TTL). The response - contains the first page plus pagination metadata; subsequent pages can be - fetched from ``GET /convert/{id}/pages/{page}`` using the returned ``id``. + Blocking (network + hashing); callers on the event loop must offload this. """ + started = time.monotonic() + path = converter.download(url) + try: + file_hash = storage.hash_file(path) + doc_key = storage.make_doc_key(file_hash, use_llm) + + # Content-addressed, so an existing document for this key was produced + # from identical bytes under identical settings. Nothing to do. + cached = jobstore.find_cached_document(doc_key) + if cached is not None: + manifest = storage.get_manifest(doc_key) + if manifest is not None: + payload = _first_page_payload(doc_key, manifest) + if payload is not None: + logger.info("cache hit for %s (%.2fs)", url, time.monotonic() - started) + return {"cached": True, "payload": payload} + # Row survived its objects (lifecycle expiry). Reconvert. + logger.info("stale document row for %s, reconverting", doc_key) + + pages = chunker.page_count(path) + ranges = chunker.plan_chunks( + pages, config.PAGES_PER_CHUNK, config.MIN_PAGES_FOR_PARALLEL + ) + storage.put_source(file_hash, path) + finally: + if os.path.exists(path): + os.remove(path) + + job_id = jobstore.new_job_id() + # State first, then publish. A message whose job row does not exist yet + # would be delivered to a worker that cannot find it; the reverse — rows + # with no message — is merely a job that never starts, which is visible in + # the status endpoint rather than a confusing worker-side error. + jobstore.create_job(job_id, doc_key, file_hash, url, ranges, use_llm) + taskqueue.publish([ + taskqueue.ChunkTask( + job_id=job_id, + chunk_index=index, + doc_key=doc_key, + file_hash=file_hash, + start_page=start, + end_page=end, + total_chunks=len(ranges), + use_llm=use_llm, + ) + for index, (start, end) in enumerate(ranges) + ]) + logger.info( + "queued job %s for %s: %s pages -> %d chunks (%.2fs)", + job_id, url, pages, len(ranges), time.monotonic() - started, + ) + return {"cached": False, "job_id": job_id, "doc_key": doc_key, + "total_chunks": len(ranges)} + + +def _job_payload(job: dict) -> dict: + return { + "job_id": job["job_id"], + "status": job["status"], + "doc_key": job["doc_key"], + "progress": jobstore.job_progress(job["job_id"]), + "error": job["error"], + "created_at": job["created_at"], + "updated_at": job["updated_at"], + } + + +async def _wait_for_job(job_id: str) -> dict: + """Poll a queued job until the worker fleet finishes or fails it.""" + while True: + job = await run_in_threadpool(jobstore.get_job, job_id) + if job is not None and job["status"] in ("done", "failed"): + return job + await asyncio.sleep(config.SYNC_POLL_INTERVAL) + + +async def _submit_request(request: Optional[ConvertRequest]) -> dict: + """Validate and enqueue a request, translating producer errors to HTTP.""" if request is None or not request.file: raise HTTPException(status_code=422, detail="`file` (a URL) is required") + use_llm = config.CONVERT_USE_LLM if request.use_llm is None else request.use_llm + try: - # `convert` blocks on network IO and CPU-bound parsing; running it on the - # event loop would stall the health probes and get the pod killed. - result = await run_in_threadpool(convert, request.file) - return await run_in_threadpool(store_document, result.text_content) - except Exception as e: - logger.exception("conversion failed for %s", request.file) - raise HTTPException(status_code=400, detail=f"Error processing URL: {str(e)}") + return await run_in_threadpool(submit, request.file, use_llm) + except Exception as exc: + # Download and size-cap failures are the caller's problem and are + # reported synchronously, rather than as a job that instantly fails. + logger.exception("submission failed for %s", request.file) + raise HTTPException( + status_code=400, detail=f"Error processing URL: {str(exc)}" + ) from exc + + +async def _completed_job_payload(job: dict) -> dict: + """Turn a successful terminal job into the original first-page response.""" + if job["status"] == "failed": + raise HTTPException( + status_code=400, detail=f"Conversion failed: {job['error']}" + ) + + manifest = await run_in_threadpool(storage.get_manifest, job["doc_key"]) + if manifest is None: + raise HTTPException(status_code=500, detail="Converted document is missing") + payload = await run_in_threadpool(_first_page_payload, job["doc_key"], manifest) + if payload is None: + raise HTTPException(status_code=500, detail="Converted document is missing") + return payload + + +@app.post( + "/convert", + dependencies=[Depends(verify_api_key)], + summary="Convert a document synchronously", + response_model=PageResponse, + responses={ + 200: {"description": "Conversion finished: first page."}, + 400: {"description": "The URL could not be downloaded, or conversion failed."}, + 422: {"description": "`file` is missing."}, + }, +) +async def convert_endpoint( + request: Optional[ConvertRequest] = None, +): + """ + Preserve the original blocking API while distributing work to worker pods. + + This request returns only after every chunk has finished and the document + has been assembled. Callers that want a job id immediately should use + ``POST /async/convert`` instead. + """ + result = await _submit_request(request) + + if result["cached"]: + return result["payload"] + + job = await _wait_for_job(result["job_id"]) + return await _completed_job_payload(job) + + +@app.post( + "/async/convert", + dependencies=[Depends(verify_api_key)], + summary="Submit a document for asynchronous conversion", + responses={ + 200: {"model": PageResponse, "description": "Cache hit: first page."}, + 202: {"model": JobAccepted, "description": "Queued; poll the job id."}, + 400: {"description": "The URL could not be downloaded."}, + 422: {"description": "`file` is missing."}, + }, +) +async def async_convert_endpoint( + response: Response, + request: Optional[ConvertRequest] = None, +): + """Queue conversion and return immediately unless the document is cached.""" + result = await _submit_request(request) + if result["cached"]: + return result["payload"] -@app.get("/convert/{doc_id}/pages/{page}", dependencies=[Depends(verify_api_key)]) + response.status_code = 202 + return { + "job_id": result["job_id"], + "status": "queued", + "doc_key": result["doc_key"], + "total_chunks": result["total_chunks"], + } + + +@app.get( + "/convert/jobs/{job_id}", + dependencies=[Depends(verify_api_key)], + summary="Poll a conversion job", + response_model=JobStatus, + responses={404: {"description": "No such job."}}, +) +async def get_job_endpoint(job_id: str): + """Status and progress for a submitted conversion job.""" + job = await run_in_threadpool(jobstore.get_job, job_id) + if job is None: + raise HTTPException(status_code=404, detail="Job not found") + return await run_in_threadpool(_job_payload, job) + + +@app.get( + "/convert/{doc_id}/pages/{page}", + dependencies=[Depends(verify_api_key)], + summary="Read one page of a converted document", + response_model=PageResponse, + responses={404: {"description": "Unknown document, expired, or page out of range."}}, +) async def get_page_endpoint(doc_id: str, page: int): """ Fetch a single page of a previously converted document. - Returns 404 if the document has expired from the cache or the page is out + Returns 404 if the document has expired from storage or the page is out of range. """ - result = await run_in_threadpool(get_page, doc_id, page) - if result is None: + manifest = await run_in_threadpool(storage.get_manifest, doc_id) + if manifest is None or page < 1 or page > manifest["total_pages"]: + raise HTTPException( + status_code=404, + detail="Document or page not found (it may have expired).", + ) + + content = await run_in_threadpool(storage.get_page, doc_id, page) + if content is None: raise HTTPException( status_code=404, detail="Document or page not found (it may have expired).", ) - return result + return { + "id": doc_id, + "content": content, + "pagination": pagination.build_pagination( + doc_id, page, manifest["total_pages"], + manifest["total_length"], manifest["page_size"], + ), + } -@app.get("/") + +_openapi_yaml: Optional[str] = None + + +@app.get("/openapi.yaml", include_in_schema=False) +async def openapi_yaml(): + """ + The generated OpenAPI document, as YAML. + + Serialised from `app.openapi()` rather than kept as a checked-in file, so + it cannot drift from the routes. FastAPI already serves the JSON form at + `/openapi.json`; this is the same document for tooling that wants YAML. + """ + global _openapi_yaml + if _openapi_yaml is None: + # app.openapi() memoises its dict, so this only pays the dump cost once. + _openapi_yaml = yaml.safe_dump(app.openapi(), sort_keys=False) + return Response(_openapi_yaml, media_type="application/yaml") + + +@app.get("/", response_model=Health, summary="Liveness") async def root(): + """Liveness. Deliberately checks nothing external — a database or storage + blip should not get the pod restarted.""" return {"status": "ok"} + + +@app.get( + "/readyz", + response_model=Readiness, + summary="Readiness", + responses={503: {"model": Readiness, "description": "A dependency is down."}}, +) +async def readyz(response: Response): + """Readiness. Fails the pod out of the Service while dependencies are down.""" + checks = { + "database": await run_in_threadpool(db.health), + "storage": await run_in_threadpool(storage.health), + "queue": await run_in_threadpool(taskqueue.health), + } + if not all(checks.values()): + response.status_code = 503 + return {"status": "ok" if all(checks.values()) else "degraded", "checks": checks} diff --git a/cache.py b/cache.py deleted file mode 100644 index b318ed1..0000000 --- a/cache.py +++ /dev/null @@ -1,149 +0,0 @@ -import json -import os -import uuid -from typing import Optional - -import redis - -# ---- Configuration (overridable via env) ---- -REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379/0") -# Number of characters per page. -PAGE_SIZE = int(os.environ.get("PAGE_SIZE", "5000")) -# How long (seconds) paginated content lives in Redis. -CACHE_TTL = int(os.environ.get("CACHE_TTL", "300")) # 5 minutes - -_KEY_PREFIX = "markitdown:doc" - -_client: Optional[redis.Redis] = None - - -def get_client() -> redis.Redis: - """Return a lazily-initialised, shared Redis client.""" - global _client - if _client is None: - _client = redis.Redis.from_url(REDIS_URL, decode_responses=True) - return _client - - -def _meta_key(doc_id: str) -> str: - return f"{_KEY_PREFIX}:{doc_id}:meta" - - -def _page_key(doc_id: str, page: int) -> str: - return f"{_KEY_PREFIX}:{doc_id}:page:{page}" - - -def paginate(content: str, page_size: int = PAGE_SIZE) -> list[str]: - """ - Split markdown content into pages of roughly ``page_size`` characters. - - Pages are broken on paragraph boundaries (``\\n\\n``) where possible so a - paragraph is not split across pages. A single paragraph larger than - ``page_size`` is hard-split. - """ - if not content: - return [""] - - pages: list[str] = [] - current = "" - - for block in content.split("\n\n"): - # Hard-split a single block that is itself larger than a page. - while len(block) > page_size: - if current: - pages.append(current) - current = "" - pages.append(block[:page_size]) - block = block[page_size:] - - candidate = block if not current else f"{current}\n\n{block}" - if len(candidate) > page_size: - pages.append(current) - current = block - else: - current = candidate - - if current or not pages: - pages.append(current) - - return pages - - -def build_pagination(doc_id: str, page: int, total_pages: int, total_length: int, - page_size: int = PAGE_SIZE) -> dict: - has_next = page < total_pages - has_prev = page > 1 - return { - "id": doc_id, - "page": page, - "page_size": page_size, - "total_pages": total_pages, - "total_length": total_length, - "has_next": has_next, - "has_prev": has_prev, - "next_page": page + 1 if has_next else None, - "prev_page": page - 1 if has_prev else None, - } - - -def store_document(content: str, page_size: int = PAGE_SIZE, - ttl: int = CACHE_TTL) -> dict: - """ - Paginate ``content`` and store every page + metadata in Redis with a TTL. - - Returns a payload for the first page:: - - {"id", "content", "pagination": {...}} - """ - pages = paginate(content, page_size) - doc_id = uuid.uuid4().hex - total_pages = len(pages) - total_length = len(content) - - meta = { - "total_pages": total_pages, - "total_length": total_length, - "page_size": page_size, - } - - client = get_client() - pipe = client.pipeline() - pipe.set(_meta_key(doc_id), json.dumps(meta), ex=ttl) - for idx, page_content in enumerate(pages, start=1): - pipe.set(_page_key(doc_id, idx), page_content, ex=ttl) - pipe.execute() - - return { - "id": doc_id, - "content": pages[0], - "pagination": build_pagination(doc_id, 1, total_pages, total_length, page_size), - } - - -def get_page(doc_id: str, page: int) -> Optional[dict]: - """ - Fetch a single page for a previously stored document. - - Returns ``None`` if the document/page is not found (expired or invalid). - """ - client = get_client() - raw_meta = client.get(_meta_key(doc_id)) - if raw_meta is None: - return None - - meta = json.loads(raw_meta) - total_pages = meta["total_pages"] - if page < 1 or page > total_pages: - return None - - content = client.get(_page_key(doc_id, page)) - if content is None: - return None - - return { - "id": doc_id, - "content": content, - "pagination": build_pagination( - doc_id, page, total_pages, meta["total_length"], meta["page_size"] - ), - } diff --git a/chunker.py b/chunker.py new file mode 100644 index 0000000..286122d --- /dev/null +++ b/chunker.py @@ -0,0 +1,80 @@ +""" +Splitting a document into independently-convertible page ranges. + +The producer only *plans* ranges — it never writes chunk files. Each worker +materialises its own range from the shared source, so no chunk bytes have to +travel between pods. +""" +import logging +from typing import Optional + +logger = logging.getLogger(__name__) + +# A chunk covering the whole file, used for non-PDFs and small PDFs. -1 rather +# than a real page number because the worker must not try to page-slice a file +# format that has no pages. +WHOLE_FILE = (0, -1) + + +def page_count(path: str) -> Optional[int]: + """ + Return the number of pages in ``path``, or None if it is not a readable PDF. + + Doubles as the PDF sniff, so a corrupt or non-PDF file transparently falls + back to whole-file conversion. + """ + try: + from pypdf import PdfReader + + return len(PdfReader(path).pages) + except Exception: + return None + + +def plan_chunks( + pages: Optional[int], + pages_per_chunk: int, + min_pages: int, +) -> list[tuple[int, int]]: + """ + Plan the page ranges a document should be split into. + + Returns a list of ``(start, end)`` pairs with ``end`` exclusive, or a single + ``WHOLE_FILE`` range when splitting is not applicable: a non-PDF (``pages`` + is None), or a PDF small enough that per-chunk parser startup would cost + more than the parallelism saves. + + Pure and side-effect free so the partitioning can be tested without + touching a real document. + """ + if pages_per_chunk < 1: + raise ValueError("pages_per_chunk must be >= 1") + + if pages is None or pages < min_pages: + return [WHOLE_FILE] + + return [ + (start, min(start + pages_per_chunk, pages)) + for start in range(0, pages, pages_per_chunk) + ] + + +def extract_pages(src_path: str, start: int, end: int, out_path: str) -> str: + """ + Write pages ``[start, end)`` of ``src_path`` to ``out_path``. + + ``end == -1`` means the whole file, in which case the source is used + directly rather than copied — the caller must not mutate or delete it. + """ + if end == -1: + return src_path + + from pypdf import PdfReader, PdfWriter + + reader = PdfReader(src_path) + writer = PdfWriter() + for page in reader.pages[start:end]: + writer.add_page(page) + with open(out_path, "wb") as fh: + writer.write(fh) + return out_path diff --git a/config.py b/config.py new file mode 100644 index 0000000..f1ff5c4 --- /dev/null +++ b/config.py @@ -0,0 +1,107 @@ +""" +Every environment-driven setting in one place. + +Two entrypoints (the API pod and the worker pod) run off the same image and +read overlapping subsets of this, so scattering ``os.environ`` reads across +modules made it impossible to see what a given pod actually needs configured. +""" +import os + +# Load .env before anything reads the environment, so `uvicorn api.index:app` +# and `python -m worker` work on their own rather than only through `make`. +# Real environment variables always win, so a container's env is not shadowed +# by a stray .env baked into an image. +try: + from dotenv import load_dotenv + + load_dotenv(override=False) +except ImportError: # pragma: no cover - python-dotenv is a declared dependency + pass + + +def _int(name: str, default: int) -> int: + return int(os.environ.get(name, str(default))) + + +def _float(name: str, default: float) -> float: + return float(os.environ.get(name, str(default))) + + +def _bool(name: str, default: bool) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() in ("1", "true", "yes", "on") + + +# ---- Auth ---- +ADMIN_API_KEY = os.environ.get("ADMIN_API_KEY") + +# ---- Database (libSQL / Turso) ---- +# A local file path works as-is; point at libsql://.turso.io with a token +# for the hosted service. Tests use ":memory:". +DATABASE_URL = os.environ.get("TURSO_DATABASE_URL", "file:markitdown.db") +DATABASE_AUTH_TOKEN = os.environ.get("TURSO_AUTH_TOKEN") + +# ---- Object storage ---- +S3_ENDPOINT = os.environ.get("S3_ENDPOINT") or None +S3_BUCKET = os.environ.get("S3_BUCKET", "markitdown") +S3_REGION = os.environ.get("S3_REGION", "auto") +S3_ACCESS_KEY_ID = os.environ.get("S3_ACCESS_KEY_ID") +S3_SECRET_ACCESS_KEY = os.environ.get("S3_SECRET_ACCESS_KEY") +# Informational only: expiry is enforced by an S3 lifecycle rule on the bucket, +# not by us. Stored on the document row so callers can see when it will vanish. +DOC_TTL_DAYS = _int("DOC_TTL_DAYS", 7) + +# ---- Conversion ---- +OPENAI_MODEL = os.environ.get("OPENAI_MODEL", "google/gemini-3.1-flash-lite-preview") +OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://ai-gateway.vercel.sh/v1") +# Whether to run markitdown's LLM-assisted conversion (image descriptions) by +# default. Off by default: it is dramatically slower and costs tokens per page, +# which does not pay for itself on book-sized PDFs. +CONVERT_USE_LLM = _bool("CONVERT_USE_LLM", False) + +# ---- Chunking ---- +PAGES_PER_CHUNK = _int("PDF_PAGES_PER_CHUNK", 20) +# Below this page count the per-chunk parser startup cost outweighs any +# parallelism, so the document is converted as a single chunk. +MIN_PAGES_FOR_PARALLEL = _int("PDF_MIN_PAGES_FOR_PARALLEL", 40) + +# ---- Pagination ---- +PAGE_SIZE = _int("PAGE_SIZE", 5000) + +# ---- Download ---- +MAX_DOWNLOAD_BYTES = _int("MAX_DOWNLOAD_BYTES", 100 * 1024 * 1024) + +# ---- Queue (RabbitMQ) ---- +RABBITMQ_URL = os.environ.get("RABBITMQ_URL", "amqp://guest:guest@localhost:5672/%2F") +RABBITMQ_EXCHANGE = os.environ.get("RABBITMQ_EXCHANGE", "markitdown") +RABBITMQ_QUEUE = os.environ.get("RABBITMQ_QUEUE", "markitdown.chunks") +# Unacked messages a worker may hold. 1 means a worker takes exactly one chunk +# at a time, which is what bounds its memory to a chunk rather than a document +# and keeps dispatch fair — without it RabbitMQ round-robins eagerly and a slow +# worker sits on a backlog while others idle. +RABBITMQ_PREFETCH = _int("RABBITMQ_PREFETCH", 1) +# How long a blocked consumer waits before yielding to heartbeats/shutdown. +RABBITMQ_CONSUME_TIMEOUT = _float("RABBITMQ_CONSUME_TIMEOUT", 5.0) +RABBITMQ_HEARTBEAT = _int("RABBITMQ_HEARTBEAT", 60) +# Delay tiers, in seconds, one retry queue each. Attempt N waits tier N. +# +# Per-tier queues rather than a per-message TTL on one queue: RabbitMQ only +# expires the message at the *head*, so a 300s message parked in front would +# hold back a 5s message queued behind it. +RETRY_DELAYS = [ + int(d) for d in os.environ.get("RETRY_DELAYS", "5,30,300").split(",") if d.strip() +] +MAX_ATTEMPTS = _int("MAX_ATTEMPTS", 3) +# Worker-local cache of downloaded source files, so a worker converting several +# chunks of the same document only fetches it from S3 once. +SOURCE_CACHE_DIR = os.environ.get("SOURCE_CACHE_DIR", "/tmp/markitdown-src") +SOURCE_CACHE_MAX_BYTES = _int("SOURCE_CACHE_MAX_BYTES", 2 * 1024 * 1024 * 1024) +# Touched by the worker loop each iteration; the k8s liveness probe checks its +# mtime to catch a wedged loop that is still technically running. +WORKER_LIVENESS_FILE = os.environ.get("WORKER_LIVENESS_FILE", "/tmp/worker-alive") + +# ---- API ---- +# Polling cadence used while the synchronous endpoint awaits its queued job. +SYNC_POLL_INTERVAL = _float("SYNC_POLL_INTERVAL", 0.5) diff --git a/converter.py b/converter.py index 5e553e0..ffbd65f 100644 --- a/converter.py +++ b/converter.py @@ -1,36 +1,51 @@ +""" +Downloading source documents and converting them to markdown. + +The LLM client is built lazily. It used to be constructed at import time from a +required ``OPENAI_API_KEY``, which meant any process that merely imported this +module — including workers that never take the LLM path — crashed at startup +without a key it did not need. + +markitdown is likewise imported lazily. It costs ~150-190 MiB of RSS, and the +API pod imports this module only to download and hash files. +""" +import functools import logging import os import os.path -import time +import tempfile from urllib.parse import urlparse -from markitdown import MarkItDown, DocumentConverterResult -from openai import OpenAI import requests -import tempfile -from pdf_chunk import convert_pdf_chunked, should_chunk +import config logger = logging.getLogger(__name__) -client = OpenAI( - base_url="https://ai-gateway.vercel.sh/v1", - api_key=os.environ["OPENAI_API_KEY"], -) -model = os.environ.get("OPENAI_MODEL", "google/gemini-3.1-flash-lite-preview") +_CHUNK_SIZE = 64 * 1024 -# Seconds to wait for the connection and for each chunk of the response body. -DOWNLOAD_CONNECT_TIMEOUT = float(os.environ.get("DOWNLOAD_CONNECT_TIMEOUT", "10")) -DOWNLOAD_READ_TIMEOUT = float(os.environ.get("DOWNLOAD_READ_TIMEOUT", "30")) -# Hard cap on the downloaded file size, in bytes (default 100 MiB). -MAX_DOWNLOAD_BYTES = int(os.environ.get("MAX_DOWNLOAD_BYTES", str(100 * 1024 * 1024))) -_CHUNK_SIZE = 64 * 1024 +@functools.lru_cache(maxsize=1) +def get_llm_client(): + """ + Build the OpenAI client on first use. + + Raises only when the LLM path is actually taken, so a missing key is a + per-request error rather than a startup crash. + """ + from openai import OpenAI + + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + raise RuntimeError( + "OPENAI_API_KEY is not set, but LLM-assisted conversion was requested" + ) + return OpenAI(base_url=config.OPENAI_BASE_URL, api_key=api_key) def download(url: str) -> str: """ - Download a URL and return a file path. + Download a URL and return a temporary file path. Args: url (str): The URL to download. @@ -38,13 +53,12 @@ def download(url: str) -> str: Returns: str: The temporary file path where the downloaded content is stored. """ - with requests.get( - url, - stream=True, - timeout=(DOWNLOAD_CONNECT_TIMEOUT, DOWNLOAD_READ_TIMEOUT), - ) as response: + with requests.get(url, stream=True) as response: if response.status_code != 200: - raise Exception(f"Failed to download file from {url}, status code: {response.status_code}") + raise Exception( + f"Failed to download file from {url}, " + f"status code: {response.status_code}" + ) # Extract filename from URL or use a default name parsed_url = urlparse(url) @@ -52,7 +66,6 @@ def download(url: str) -> str: if not filename: filename = "downloaded_file" - # Create a temporary file with a name derived from the URL fd, temp_path = tempfile.mkstemp(suffix=f"_{filename}") os.close(fd) @@ -64,9 +77,10 @@ def download(url: str) -> str: if not chunk: continue written += len(chunk) - if written > MAX_DOWNLOAD_BYTES: + if written > config.MAX_DOWNLOAD_BYTES: raise Exception( - f"File from {url} exceeds the maximum size of {MAX_DOWNLOAD_BYTES} bytes" + f"File from {url} exceeds the maximum size of " + f"{config.MAX_DOWNLOAD_BYTES} bytes" ) file.write(chunk) except Exception: @@ -77,44 +91,20 @@ def download(url: str) -> str: return temp_path -def convert(url: str) -> DocumentConverterResult: +def convert_file(path: str, use_llm: bool = False) -> str: """ - Convert a URL to a MarkItDown object. - - This is fully synchronous and CPU/IO blocking; callers running inside an - event loop must offload it to a worker thread. - - Args: - url (str): The URL to convert. + Convert a local file to markdown. - Returns: - MarkItDown: The converted MarkItDown object. + The single conversion primitive. Previously the LLM was used or skipped + depending on whether a document happened to cross the 40-page chunking + threshold, so a 39-page PDF got image descriptions and a 40-page one did + not — a behaviour flip no caller could predict. ``use_llm`` now says so + explicitly and applies uniformly to every chunk of a document. """ - started = time.monotonic() - temp_file = download(url) - downloaded_at = time.monotonic() - logger.info("downloaded %s in %.3fs", url, downloaded_at - started) - - try: - # Large PDFs are converted chunk-by-chunk across worker processes. - # Done whole-document, markitdown holds the parsed object graph for - # every page at once, which pins the container memory limit and - # serialises all the parsing onto one core. - if should_chunk(temp_file): - markdown = convert_pdf_chunked(temp_file) - converted = DocumentConverterResult(markdown=markdown) - else: - md = MarkItDown(llm_client=client, llm_model=model) - converted = md.convert(temp_file) - finally: - if os.path.exists(temp_file): - os.remove(temp_file) - - logger.info( - "converted %s in %.3fs (download %.3fs, convert %.3fs)", - url, - time.monotonic() - started, - downloaded_at - started, - time.monotonic() - downloaded_at, - ) - return converted + from markitdown import MarkItDown + + if use_llm: + md = MarkItDown(llm_client=get_llm_client(), llm_model=config.OPENAI_MODEL) + else: + md = MarkItDown() + return md.convert(path).markdown diff --git a/db.py b/db.py new file mode 100644 index 0000000..db54d5c --- /dev/null +++ b/db.py @@ -0,0 +1,285 @@ +""" +Database access: connections, migrations, and small query helpers. + +Two interchangeable drivers sit behind the same DB-API 2.0 surface: + +* **libsql** for hosted Turso (``libsql://…`` or any URL with an auth token). +* **stdlib sqlite3** for everything else — local development and CI. + +Both speak the same SQL, so nothing above this module knows which is in use. +Keeping sqlite3 as the local/CI path means tests need no native libsql build +and no network, and the queue's claim semantics are exercised against a real +SQLite engine either way. + +Connections are thread-local. Both entrypoints are multi-threaded — the API +runs blocking work in FastAPI's threadpool, and each worker runs a heartbeat +thread alongside its conversion — and neither driver's connection is safe to +share across threads. +""" +import logging +import os +import threading +from contextlib import contextmanager +from typing import Any, Iterable, Optional + +import config + +logger = logging.getLogger(__name__) + +_local = threading.local() +_migrated = False +_migrate_lock = threading.Lock() + +MIGRATIONS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "migrations") + +# How long a writer waits for a competing write lock before giving up. Workers +# claim chunks concurrently, so brief contention is normal and must block +# rather than raise "database is locked". +_BUSY_TIMEOUT_SECONDS = 30 + + +def uses_libsql() -> bool: + """ + True when the database is a libSQL *server* rather than a local SQLite file. + + Covers hosted Turso (libsql://) and a self-hosted sqld over plain HTTP, + which is what the in-cluster deployment uses. Anything else — a path, or + ":memory:" — is stdlib sqlite3. + """ + url = config.DATABASE_URL + return ( + url.startswith(("libsql://", "http://", "https://")) + or bool(config.DATABASE_AUTH_TOKEN) + ) + + +def _connect_libsql(): + import libsql + + url = config.DATABASE_URL + if url.startswith("file:"): + url = url[len("file:"):] + if config.DATABASE_AUTH_TOKEN: + return libsql.connect(url, auth_token=config.DATABASE_AUTH_TOKEN) + return libsql.connect(url) + + +def _connect_sqlite(): + import sqlite3 + + url = config.DATABASE_URL + if url.startswith("file:"): + url = url[len("file:"):] + + if url == ":memory:": + # A plain ":memory:" database is private to the connection that opened + # it, so with thread-local connections every thread would silently get + # its own empty database. Shared cache keeps them pointed at one. + conn = sqlite3.connect( + "file:markitdown-mem?mode=memory&cache=shared", + uri=True, + timeout=_BUSY_TIMEOUT_SECONDS, + ) + else: + conn = sqlite3.connect(url, timeout=_BUSY_TIMEOUT_SECONDS) + # WAL lets readers proceed during a write, which matters because the + # API polls job status while workers are claiming and completing. + conn.execute("PRAGMA journal_mode=WAL") + + conn.execute(f"PRAGMA busy_timeout={_BUSY_TIMEOUT_SECONDS * 1000}") + return conn + + +def get_connection(): + """Return this thread's connection, opening it on first use.""" + conn = getattr(_local, "conn", None) + if conn is None: + conn = _connect_libsql() if uses_libsql() else _connect_sqlite() + _local.conn = conn + return conn + + +def close_connection() -> None: + """Close and forget this thread's connection.""" + conn = getattr(_local, "conn", None) + if conn is not None: + try: + conn.close() + except Exception: + logger.debug("error closing connection", exc_info=True) + _local.conn = None + + +def reset() -> None: + """Drop the cached connection and re-run migrations on next use.""" + global _migrated + close_connection() + with _migrate_lock: + _migrated = False + + +def migrate() -> None: + """ + Apply every .sql file in migrations/ in name order. + + Migrations are written to be idempotent (IF NOT EXISTS), so this is safe to + run from every pod on every start rather than needing a separate job. + """ + global _migrated + with _migrate_lock: + if _migrated: + return + conn = get_connection() + for name in sorted(os.listdir(MIGRATIONS_DIR)): + if not name.endswith(".sql"): + continue + with open(os.path.join(MIGRATIONS_DIR, name), "r") as fh: + sql = fh.read() + for statement in _split_statements(sql): + conn.execute(statement) + conn.commit() + _migrated = True + logger.info("migrations applied") + + +def _split_statements(sql: str) -> list[str]: + """ + Split a migration file into statements. + + Both drivers' execute() take one statement at a time. Line comments are + stripped before splitting because a ';' inside a comment would otherwise + cut a statement in half — which is exactly what "-- exclusive; -1 means + the whole file" did. + + The migrations contain no string literals, so this does not need to be a + real parser; if one is ever added, this must become quote-aware. + """ + without_comments = "\n".join( + line.split("--")[0] for line in sql.splitlines() + ) + return [s.strip() for s in without_comments.split(";") if s.strip()] + + +def _rows_to_dicts(cursor) -> list[dict]: + if cursor.description is None: + return [] + # libsql returns None from fetchall() when a statement matched no rows, + # where sqlite3 returns []. A RETURNING clause that matched nothing — the + # normal "queue is empty" case — hits this on every poll. + rows = cursor.fetchall() + if not rows: + return [] + columns = [d[0] for d in cursor.description] + return [dict(zip(columns, row)) for row in rows] + + +# libsql speaks hrana over HTTP, where a connection holds a server-side stream +# handle that expires when idle ("Stream handle N is expired"). A pooled +# thread-local connection therefore goes stale on its own, and every subsequent +# statement on it fails until it is reopened. sqlite3 has no equivalent. +_STALE_CONNECTION_MARKERS = ( + "stream handle", + "stream is closed", + "stream not found", + "connection reset", + "connection closed", + "broken pipe", +) + + +def _is_stale_connection(exc: Exception) -> bool: + message = str(exc).lower() + return any(marker in message for marker in _STALE_CONNECTION_MARKERS) + + +def _run(sql: str, params: tuple, commit: bool) -> list[dict]: + """ + Execute one statement, reopening the connection once if it went stale. + + The retry is deliberately narrow: it fires only on errors that mean the + server never received the statement, so a write cannot be applied twice. + Any other failure propagates untouched. + """ + ensure_ready() + for attempt in (1, 2): + conn = get_connection() + try: + cursor = conn.execute(sql, params) + rows = _rows_to_dicts(cursor) + if commit: + conn.commit() + return rows + except Exception as exc: + if attempt == 1 and _is_stale_connection(exc): + logger.info("reopening stale database connection: %s", exc) + close_connection() + continue + raise + return [] + + +def query(sql: str, params: Iterable[Any] = ()) -> list[dict]: + """Run a read and return rows as dicts.""" + return _run(sql, tuple(params), commit=False) + + +def query_one(sql: str, params: Iterable[Any] = ()) -> Optional[dict]: + rows = query(sql, params) + return rows[0] if rows else None + + +def execute(sql: str, params: Iterable[Any] = ()) -> list[dict]: + """ + Run a write and commit it. + + Returns any RETURNING rows, which is how the queue's claim and completion + statements report what they actually changed. + """ + return _run(sql, tuple(params), commit=True) + + +@contextmanager +def transaction(): + """ + Run a block of statements as one transaction. + + Yields the raw connection. Commits on clean exit, rolls back on exception. + + The connection is pinged first so an expired libsql stream handle is + replaced *before* the block starts. A transaction cannot be retried + transparently once the caller's statements have begun, so the cheap check + up front is the only safe place to recover. + """ + ensure_ready() + query("SELECT 1 AS ok") + conn = get_connection() + try: + yield conn + conn.commit() + except Exception: + try: + conn.rollback() + except Exception: + logger.warning("rollback failed", exc_info=True) + raise + + +def changes(conn) -> int: + """Rows modified by the most recent statement on ``conn``.""" + return conn.execute("SELECT changes() AS n").fetchall()[0][0] + + +def ensure_ready() -> None: + if not _migrated: + migrate() + + +def health() -> bool: + try: + # Goes through _run, so a connection that expired while the pod was + # idle is reopened here rather than reported as a dead dependency. + query("SELECT 1 AS ok") + return True + except Exception: + logger.warning("database health check failed", exc_info=True) + return False diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..a84b9ca --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,35 @@ +# Local dependencies only. +# +# The database is remote (Turso) — see TURSO_DATABASE_URL in .env. What has to +# run locally is the broker that distributes chunks, and object storage for the +# source files and converted pages. + +services: + rabbitmq: + image: rabbitmq:4-management-alpine + container_name: markitdown-rabbitmq + ports: + - "${RABBITMQ_PORT:-5673}:5672" + - "${RABBITMQ_MANAGEMENT_PORT:-15673}:15672" + environment: + RABBITMQ_DEFAULT_USER: ${RABBITMQ_USER:-guest} + RABBITMQ_DEFAULT_PASS: ${RABBITMQ_PASSWORD:-guest} + volumes: + - rabbitmq-data:/var/lib/rabbitmq + + minio: + image: quay.io/minio/minio:latest + container_name: markitdown-minio + command: server /data --console-address ":9001" + ports: + - "${MINIO_PORT:-9100}:9000" + - "${MINIO_CONSOLE_PORT:-9101}:9001" + environment: + MINIO_ROOT_USER: ${S3_ACCESS_KEY_ID:-minioadmin} + MINIO_ROOT_PASSWORD: ${S3_SECRET_ACCESS_KEY:-minioadmin} + volumes: + - minio-data:/data + +volumes: + rabbitmq-data: + minio-data: diff --git a/jobstore.py b/jobstore.py new file mode 100644 index 0000000..a4b56fd --- /dev/null +++ b/jobstore.py @@ -0,0 +1,273 @@ +""" +Job and document state. + +RabbitMQ distributes the work (see taskqueue.py); this owns everything the +broker cannot answer: + +* the per-job remaining-chunk counter, which decides *who assembles* the + finished document — a question that has to be answered atomically and + exactly once across the whole worker fleet; +* progress for ``GET /convert/jobs/{id}``; +* the content-addressed document cache, so a re-submitted file skips + conversion entirely. + +``job_chunks`` mirrors each chunk's status for visibility and progress. It is +not the queue — RabbitMQ is — so it carries no leases, no claim, and no polling. +""" +import logging +import time +import uuid +from typing import Optional + +import config +import db + +logger = logging.getLogger(__name__) + +# complete_chunk() sentinels. +JOB_ALREADY_TERMINAL = -1 +CHUNK_ALREADY_DONE = -2 + + +def now() -> int: + return int(time.time()) + + +def new_job_id() -> str: + return uuid.uuid4().hex + + +# ---- Document cache ---- + + +def find_cached_document(doc_key: str) -> Optional[dict]: + """Return a previously converted document by cache key, if any.""" + return db.query_one("SELECT * FROM documents WHERE doc_key = ?", (doc_key,)) + + +def record_document( + doc_key: str, + file_hash: str, + total_pages: int, + total_chunks: int, + total_length: int, + page_size: int, +) -> None: + """ + Record a finished document. + + INSERT OR REPLACE rather than INSERT: two jobs for the same content can race + (both submitted before either finished), and since the key is content + addressed they produce identical output, so last writer wins harmlessly. + """ + created = now() + db.execute( + """ + INSERT OR REPLACE INTO documents + (doc_key, file_hash, config_fp, total_pages, total_chunks, + total_length, page_size, created_at, expires_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + doc_key, + file_hash, + doc_key.split(":")[-1], + total_pages, + total_chunks, + total_length, + page_size, + created, + created + config.DOC_TTL_DAYS * 86400, + ), + ) + + +# ---- Job lifecycle ---- + + +def create_job( + job_id: str, + doc_key: str, + file_hash: str, + url: str, + ranges: list[tuple[int, int]], + use_llm: bool, +) -> None: + """ + Create a job and its chunk rows. + + Inserting the chunk rows *is* the enqueue — there is no separate publish + step, so a job can never exist with its work missing from the queue. + """ + timestamp = now() + with db.transaction() as conn: + conn.execute( + """ + INSERT INTO jobs + (job_id, doc_key, file_hash, url, status, total_chunks, + remaining_chunks, use_llm, created_at, updated_at) + VALUES (?, ?, ?, ?, 'queued', ?, ?, ?, ?, ?) + """, + (job_id, doc_key, file_hash, url, len(ranges), len(ranges), + 1 if use_llm else 0, timestamp, timestamp), + ) + for index, (start, end) in enumerate(ranges): + conn.execute( + """ + INSERT INTO job_chunks + (job_id, chunk_index, start_page, end_page, status, + attempts, available_at, updated_at) + VALUES (?, ?, ?, ?, 'pending', 0, ?, ?) + """, + (job_id, index, start, end, timestamp, timestamp), + ) + logger.info("created job %s with %d chunks", job_id, len(ranges)) + + +def get_job(job_id: str) -> Optional[dict]: + return db.query_one("SELECT * FROM jobs WHERE job_id = ?", (job_id,)) + + +def job_progress(job_id: str) -> dict: + job = get_job(job_id) + if job is None: + return {"completed": 0, "total": 0, "percent": 0} + total = job["total_chunks"] + completed = total - job["remaining_chunks"] + return { + "completed": completed, + "total": total, + "percent": round(100 * completed / total) if total else 0, + } + + +def finish_job(job_id: str) -> None: + db.execute( + "UPDATE jobs SET status = 'done', updated_at = ? WHERE job_id = ?", + (now(), job_id), + ) + + +def fail_job(job_id: str, error: str, chunk_index: Optional[int] = None) -> None: + """ + Mark a job failed. + + One poison chunk fails the whole job: silently returning a document that is + missing twenty pages is worse than returning an error. Sibling chunks still + in flight discover this via complete_chunk() and drop their work. + """ + db.execute( + """ + UPDATE jobs SET status = 'failed', error = ?, error_chunk = ?, updated_at = ? + WHERE job_id = ? AND status NOT IN ('done', 'failed') + """, + (error[:2000], chunk_index, now(), job_id), + ) + logger.warning("job %s failed: %s", job_id, error) + + +# ---- Chunk state ---- + + +def start_chunk(job_id: str, chunk_index: int, owner: str, attempt: int) -> None: + """ + Record that a worker has picked a chunk up. + + Bookkeeping for progress and debugging only — the broker, not this row, + decides who holds the message. + """ + timestamp = now() + db.execute( + """ + UPDATE job_chunks + SET status = 'running', lease_owner = ?, attempts = ?, updated_at = ? + WHERE job_id = ? AND chunk_index = ? AND status <> 'done' + """, + (owner, attempt, timestamp, job_id, chunk_index), + ) + db.execute( + """ + UPDATE jobs SET status = 'running', updated_at = ? + WHERE job_id = ? AND status = 'queued' + """, + (timestamp, job_id), + ) + + +def fail_chunk(job_id: str, chunk_index: int, error: str) -> None: + """Mark a chunk as permanently failed after exhausting its attempts.""" + db.execute( + """ + UPDATE job_chunks + SET status = 'failed', lease_owner = NULL, last_error = ?, updated_at = ? + WHERE job_id = ? AND chunk_index = ? + """, + (error[:2000], now(), job_id, chunk_index), + ) + + +def reschedule_chunk(job_id: str, chunk_index: int, error: str) -> None: + """Mark a chunk as awaiting retry; RabbitMQ owns the actual redelivery.""" + db.execute( + """ + UPDATE job_chunks + SET status = 'pending', lease_owner = NULL, last_error = ?, updated_at = ? + WHERE job_id = ? AND chunk_index = ? AND status <> 'done' + """, + (error[:2000], now(), job_id, chunk_index), + ) + + +def complete_chunk(job_id: str, chunk_index: int, s3_key: str) -> int: + """ + Record a finished chunk and return the job's remaining chunk count. + + Returns 0 when this caller completed the *last* chunk and is therefore + responsible for assembling the document, ``CHUNK_ALREADY_DONE`` if this + chunk was already recorded, or ``JOB_ALREADY_TERMINAL`` if the job has + since finished or failed. + + The guarded first UPDATE is what makes this idempotent, and that is load + bearing: RabbitMQ delivers at least once, so a worker that converts a chunk + and dies before acking will have it redelivered and converted again. + Without the guard, ``remaining_chunks`` would be decremented twice, drive + past zero, and no worker would ever see 0 to assemble the document. + """ + timestamp = now() + with db.transaction() as conn: + conn.execute( + """ + UPDATE job_chunks + SET status = 'done', s3_key = ?, lease_owner = NULL, + last_error = NULL, updated_at = ? + WHERE job_id = ? AND chunk_index = ? AND status <> 'done' + """, + (s3_key, timestamp, job_id, chunk_index), + ) + if db.changes(conn) == 0: + return CHUNK_ALREADY_DONE + + cursor = conn.execute( + """ + UPDATE jobs + SET remaining_chunks = remaining_chunks - 1, + status = 'running', + updated_at = ? + WHERE job_id = ? AND status NOT IN ('done', 'failed') + RETURNING remaining_chunks + """, + (timestamp, job_id), + ) + rows = cursor.fetchall() + if not rows: + return JOB_ALREADY_TERMINAL + return rows[0][0] + + +def chunk_indices(job_id: str) -> list[int]: + """Every chunk index for a job, in page order.""" + rows = db.query( + "SELECT chunk_index FROM job_chunks WHERE job_id = ? ORDER BY chunk_index", + (job_id,), + ) + return [row["chunk_index"] for row in rows] diff --git a/k8s/api-deployment.yaml b/k8s/api-deployment.yaml new file mode 100644 index 0000000..32ab1a8 --- /dev/null +++ b/k8s/api-deployment.yaml @@ -0,0 +1,100 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: markitdown-api + namespace: markitdown-server + labels: + app: markitdown-api +spec: + replicas: 2 + selector: + matchLabels: + app: markitdown-api + template: + metadata: + labels: + app: markitdown-api + spec: + # The API no longer parses documents — it downloads, hashes and enqueues. + # /convert still holds the HTTP request open while workers finish, so let + # in-flight synchronous callers drain during a rollout. + terminationGracePeriodSeconds: 120 + containers: + - name: markitdown-api + image: ghcr.io/sirily11/markitdown-server:latest + resources: + limits: + # Sized for download + SHA-256 + enqueue. markitdown is imported + # lazily and never on this path, so its ~150-190Mi is not paid + # for here. + memory: "512Mi" + cpu: "500m" + requests: + memory: "256Mi" + cpu: "100m" + env: + # NOTE: `env` takes precedence over `envFrom`, so these in-cluster + # defaults win over anything of the same name in the secret. To run + # against hosted Turso or R2, delete the relevant entries and supply + # them via the secret instead — changing only the secret will not + # take effect while they are here. + # RABBITMQ_URL is intentionally absent here: it arrives whole from + # the secret via envFrom, so the broker (in-cluster rabbitmq.yaml + # vs a shared one in another namespace) is a secret-only change. + # It used to be pinned here only because it was assembled from + # $(RABBITMQ_USER) — and $(VAR) expands solely from earlier entries + # in this same list, never from envFrom. + - name: TURSO_DATABASE_URL + value: "http://markitdown-libsql:8080" + - name: S3_ENDPOINT + value: "http://markitdown-minio:9000" + - name: S3_BUCKET + value: "markitdown" + - name: OPENAI_MODEL + value: "google/gemini-3.1-flash-lite-preview" + - name: PDF_PAGES_PER_CHUNK + value: "20" + - name: PDF_MIN_PAGES_FOR_PARALLEL + value: "40" + envFrom: + # ADMIN_API_KEY, OPENAI_API_KEY, S3_ACCESS_KEY_ID, + # S3_SECRET_ACCESS_KEY (and TURSO_AUTH_TOKEN when managed). + - secretRef: + name: markitdown-server-secret + livenessProbe: + # Deliberately the dependency-free endpoint: a database or storage + # blip should take the pod out of the Service (readiness), not + # restart it. + httpGet: + path: / + port: 8000 + initialDelaySeconds: 5 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + readinessProbe: + httpGet: + path: /readyz + port: 8000 + initialDelaySeconds: 5 + periodSeconds: 5 + timeoutSeconds: 5 + failureThreshold: 3 + +--- +apiVersion: v1 +kind: Service +metadata: + # Name kept from the pre-split deployment so the ingress and the benchmark's + # port-forward do not need to change. + name: markitdown-server + namespace: markitdown-server + labels: + app: markitdown-api +spec: + selector: + app: markitdown-api + ports: + - port: 8080 + targetPort: 8000 + type: ClusterIP diff --git a/k8s/deployment.yaml b/k8s/deployment.yaml deleted file mode 100644 index 47f2e8a..0000000 --- a/k8s/deployment.yaml +++ /dev/null @@ -1,84 +0,0 @@ -apiVersion: apps/v1 -kind: Deployment -metadata: - name: markitdown-server - namespace: markitdown-server - labels: - app: markitdown-server -spec: - replicas: 1 - selector: - matchLabels: - app: markitdown-server - template: - metadata: - labels: - app: markitdown-server - spec: - # Conversions can run for ~60s; let in-flight requests drain on rollout - # instead of being SIGKILLed after the default 30s. - terminationGracePeriodSeconds: 120 - volumes: - containers: - - name: markitdown-server - image: ghcr.io/sirily11/markitdown-server:latest - resources: - limits: - # PDF conversion is CPU-bound and now fans chunks out across - # worker processes; the limit has to exceed PDF_MAX_WORKERS - # cores or the workers just contend for the same slice. - memory: "2048Mi" - cpu: "2000m" - requests: - # Baseline RSS is ~120-140Mi; a 128Mi request pinned memory - # utilization at ~100% and kept the HPA at max replicas. - memory: "512Mi" - cpu: "500m" - env: - - name: REDIS_URL - value: "redis://markitdown-redis:6379/0" - - name: OPENAI_MODEL - value: "google/gemini-2.5-flash-preview-05-20" - # Keep in step with the CPU limit above: each worker saturates - # roughly one core, and each costs ~150-200Mi of resident memory - # for its own markitdown import, so this trades against the - # memory limit too. - - name: PDF_MAX_WORKERS - value: "2" - - name: PDF_PAGES_PER_CHUNK - value: "20" - envFrom: - - secretRef: - name: markitdown-server-secret - livenessProbe: - httpGet: - path: / - port: 8000 - initialDelaySeconds: 10 - periodSeconds: 10 - timeoutSeconds: 5 - failureThreshold: 3 - readinessProbe: - httpGet: - path: / - port: 8000 - initialDelaySeconds: 10 - periodSeconds: 10 - timeoutSeconds: 5 - failureThreshold: 3 - ---- -apiVersion: v1 -kind: Service -metadata: - name: markitdown-server - namespace: markitdown-server - labels: - app: markitdown-server -spec: - selector: - app: markitdown-server - ports: - - port: 8080 - targetPort: 8000 - type: ClusterIP diff --git a/k8s/hpa.yaml b/k8s/hpa.yaml index 7328d12..8d3d56a 100644 --- a/k8s/hpa.yaml +++ b/k8s/hpa.yaml @@ -1,15 +1,15 @@ apiVersion: autoscaling/v2 kind: HorizontalPodAutoscaler metadata: - name: markitdown-server-hpa + name: markitdown-api-hpa namespace: markitdown-server spec: scaleTargetRef: apiVersion: apps/v1 kind: Deployment - name: markitdown-server - minReplicas: 1 - maxReplicas: 2 + name: markitdown-api + minReplicas: 2 + maxReplicas: 4 metrics: - type: Resource resource: @@ -17,19 +17,57 @@ spec: target: type: Utilization averageUtilization: 70 + behavior: + scaleDown: + stabilizationWindowSeconds: 300 + policies: + - type: Percent + value: 50 + periodSeconds: 60 + scaleUp: + stabilizationWindowSeconds: 0 + policies: + - type: Percent + value: 100 + periodSeconds: 60 + +--- +apiVersion: autoscaling/v2 +kind: HorizontalPodAutoscaler +metadata: + name: markitdown-worker-hpa + namespace: markitdown-server +spec: + scaleTargetRef: + apiVersion: apps/v1 + kind: Deployment + name: markitdown-worker + minReplicas: 1 + maxReplicas: 8 + metrics: + # CPU only, deliberately. A worker mid-chunk legitimately sits near its + # memory limit, and that memory does not fall when the queue drains, so a + # memory target would scale up on load and then refuse to scale back down. + # + # The honest metric is queue depth — jobstore.pending_depth() — which needs + # KEDA or a custom metrics adapter. CPU is a lagging proxy: an idle worker + # blocks at ~0% CPU, so a cold burst scales up only once work is already + # running. Acceptable for throughput, and worth revisiting. - type: Resource resource: - name: memory + name: cpu target: type: Utilization - averageUtilization: 80 + averageUtilization: 70 behavior: scaleDown: - stabilizationWindowSeconds: 300 + # Long window: chunks take minutes, and killing a worker mid-chunk means + # waiting out its lease before the work is retried. + stabilizationWindowSeconds: 600 policies: - - type: Percent - value: 10 - periodSeconds: 60 + - type: Pods + value: 1 + periodSeconds: 120 scaleUp: stabilizationWindowSeconds: 0 policies: @@ -37,6 +75,6 @@ spec: value: 100 periodSeconds: 60 - type: Pods - value: 2 + value: 4 periodSeconds: 60 selectPolicy: Max diff --git a/k8s/ingress.yaml b/k8s/ingress.yaml index abdd920..026eb29 100644 --- a/k8s/ingress.yaml +++ b/k8s/ingress.yaml @@ -5,17 +5,17 @@ metadata: namespace: markitdown-server annotations: cert-manager.io/cluster-issuer: letsencrypt-prod - nginx.ingress.kubernetes.io/proxy-body-size: "1000m" - # Large PDFs take several minutes to parse; nginx's 60s default cut the - # connection mid-conversion and returned a 504 to the caller. + # Request bodies are a JSON object holding a URL; the source document is + # fetched server-side, so the old 1000m cap was vestigial. + nginx.ingress.kubernetes.io/proxy-body-size: "1m" nginx.ingress.kubernetes.io/proxy-connect-timeout: "60" + # /convert preserves the original blocking contract, and large PDFs can + # take several minutes. /async/convert is available to avoid a long request. nginx.ingress.kubernetes.io/proxy-send-timeout: "600" nginx.ingress.kubernetes.io/proxy-read-timeout: "600" - nginx.ingress.kubernetes.io/affinity: "cookie" - nginx.ingress.kubernetes.io/affinity-mode: "persistent" - nginx.ingress.kubernetes.io/session-cookie-name: "mcp-router-affinity" - nginx.ingress.kubernetes.io/session-cookie-max-age: "86400" - nginx.ingress.kubernetes.io/session-cookie-path: "/" + # No session affinity: every pod reads the same database and object store, + # so any pod can answer any request. Affinity is what made replicas > 1 + # useless before. spec: ingressClassName: nginx tls: diff --git a/k8s/kustomization.yaml b/k8s/kustomization.yaml index 9bd1039..093d7bf 100644 --- a/k8s/kustomization.yaml +++ b/k8s/kustomization.yaml @@ -5,7 +5,13 @@ namespace: markitdown-server resources: - namespace.yaml - - redis.yaml - - deployment.yaml + # Self-hosted infrastructure. Drop libsql/minio and repoint the secret at + # hosted Turso and R2/S3 to run those managed; the application is identical + # either way. + - rabbitmq.yaml + - libsql.yaml + - minio.yaml + - api-deployment.yaml + - worker-deployment.yaml - hpa.yaml - ingress.yaml diff --git a/k8s/libsql.yaml b/k8s/libsql.yaml new file mode 100644 index 0000000..a857c61 --- /dev/null +++ b/k8s/libsql.yaml @@ -0,0 +1,100 @@ +# Self-hosted libSQL server (sqld). +# +# The API and worker pods must share one database — a SQLite file on a pod's +# local disk cannot be, which is why this exists rather than a PVC. For a +# managed setup, delete this file from kustomization.yaml and point +# TURSO_DATABASE_URL / TURSO_AUTH_TOKEN at hosted Turso instead; nothing in the +# application changes. +# +# Unlike the Redis it replaces, this is a durable store, not a cache: it holds +# the queue and every job's state. Hence a PVC and Recreate rather than an +# ephemeral pod with an eviction policy. +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: markitdown-libsql-data + namespace: markitdown-server +spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 5Gi + +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: markitdown-libsql + namespace: markitdown-server + labels: + app: markitdown-libsql +spec: + replicas: 1 + strategy: + # SQLite takes a single writer lock on the data directory, so the new pod + # must not start until the old one has released it. + type: Recreate + selector: + matchLabels: + app: markitdown-libsql + template: + metadata: + labels: + app: markitdown-libsql + spec: + containers: + - name: libsql + image: ghcr.io/tursodatabase/libsql-server:latest + ports: + - containerPort: 8080 + name: http + env: + - name: SQLD_NODE + value: "primary" + - name: SQLD_DB_PATH + value: "/var/lib/sqld/markitdown.db" + - name: SQLD_HTTP_LISTEN_ADDR + value: "0.0.0.0:8080" + resources: + limits: + memory: "512Mi" + cpu: "500m" + requests: + memory: "128Mi" + cpu: "100m" + volumeMounts: + - name: data + mountPath: /var/lib/sqld + readinessProbe: + httpGet: + path: /health + port: 8080 + initialDelaySeconds: 5 + periodSeconds: 5 + livenessProbe: + httpGet: + path: /health + port: 8080 + initialDelaySeconds: 15 + periodSeconds: 20 + volumes: + - name: data + persistentVolumeClaim: + claimName: markitdown-libsql-data + +--- +apiVersion: v1 +kind: Service +metadata: + name: markitdown-libsql + namespace: markitdown-server + labels: + app: markitdown-libsql +spec: + selector: + app: markitdown-libsql + ports: + - port: 8080 + targetPort: 8080 + type: ClusterIP diff --git a/k8s/minio.yaml b/k8s/minio.yaml new file mode 100644 index 0000000..2a1a150 --- /dev/null +++ b/k8s/minio.yaml @@ -0,0 +1,143 @@ +# In-cluster S3-compatible object storage. +# +# Holds source files, per-chunk markdown and the final paginated pages. For a +# managed setup, delete this from kustomization.yaml and point the S3_* secret +# values at R2 or AWS instead — the application only speaks the S3 API. +# +# Object expiry is a bucket lifecycle rule (see the bucket-setup Job below), +# not something the application enforces. +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: markitdown-minio-data + namespace: markitdown-server +spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 20Gi + +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: markitdown-minio + namespace: markitdown-server + labels: + app: markitdown-minio +spec: + replicas: 1 + strategy: + type: Recreate + selector: + matchLabels: + app: markitdown-minio + template: + metadata: + labels: + app: markitdown-minio + spec: + containers: + - name: minio + image: quay.io/minio/minio:latest + args: ["server", "/data", "--console-address", ":9001"] + ports: + - containerPort: 9000 + name: s3 + - containerPort: 9001 + name: console + env: + - name: MINIO_ROOT_USER + valueFrom: + secretKeyRef: + name: markitdown-server-secret + key: S3_ACCESS_KEY_ID + - name: MINIO_ROOT_PASSWORD + valueFrom: + secretKeyRef: + name: markitdown-server-secret + key: S3_SECRET_ACCESS_KEY + resources: + limits: + memory: "1024Mi" + cpu: "1000m" + requests: + memory: "256Mi" + cpu: "100m" + volumeMounts: + - name: data + mountPath: /data + readinessProbe: + httpGet: + path: /minio/health/ready + port: 9000 + initialDelaySeconds: 5 + periodSeconds: 5 + livenessProbe: + httpGet: + path: /minio/health/live + port: 9000 + initialDelaySeconds: 15 + periodSeconds: 20 + volumes: + - name: data + persistentVolumeClaim: + claimName: markitdown-minio-data + +--- +apiVersion: v1 +kind: Service +metadata: + name: markitdown-minio + namespace: markitdown-server + labels: + app: markitdown-minio +spec: + selector: + app: markitdown-minio + ports: + - port: 9000 + targetPort: 9000 + name: s3 + type: ClusterIP + +--- +# Creates the bucket and sets the 7-day expiry. Idempotent, so re-running a +# deploy is safe; `backoffLimit` retries while minio is still starting. +apiVersion: batch/v1 +kind: Job +metadata: + name: markitdown-minio-setup + namespace: markitdown-server +spec: + backoffLimit: 10 + template: + spec: + restartPolicy: OnFailure + containers: + - name: mc + image: quay.io/minio/mc:latest + env: + - name: S3_ACCESS_KEY_ID + valueFrom: + secretKeyRef: + name: markitdown-server-secret + key: S3_ACCESS_KEY_ID + - name: S3_SECRET_ACCESS_KEY + valueFrom: + secretKeyRef: + name: markitdown-server-secret + key: S3_SECRET_ACCESS_KEY + command: + - sh + - -c + - | + set -e + mc alias set local http://markitdown-minio:9000 \ + "$S3_ACCESS_KEY_ID" "$S3_SECRET_ACCESS_KEY" + mc mb --ignore-existing local/markitdown + # Expiry is enforced here rather than in application code, so + # every object for a document ages out together. + mc ilm rule add local/markitdown --expire-days 7 || true + echo "bucket ready" diff --git a/k8s/rabbitmq.yaml b/k8s/rabbitmq.yaml new file mode 100644 index 0000000..87e4ebd --- /dev/null +++ b/k8s/rabbitmq.yaml @@ -0,0 +1,108 @@ +# RabbitMQ — job distribution. +# +# The API publishes one message per chunk; workers consume with prefetch=1, so +# the broker hands each chunk to whichever worker is free. Adding a worker pod +# adds throughput with no coordination. +# +# Durable queues and persistent messages, backed by a PVC: a broker restart +# must not lose queued work, since a dropped message means a job that never +# finishes. +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: markitdown-rabbitmq-data + namespace: markitdown-server +spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 5Gi + +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: markitdown-rabbitmq + namespace: markitdown-server + labels: + app: markitdown-rabbitmq +spec: + replicas: 1 + strategy: + # One node owns the message store on the PVC; a rolling update would have + # two pods contend for it. + type: Recreate + selector: + matchLabels: + app: markitdown-rabbitmq + template: + metadata: + labels: + app: markitdown-rabbitmq + spec: + containers: + - name: rabbitmq + image: rabbitmq:4-management-alpine + ports: + - containerPort: 5672 + name: amqp + - containerPort: 15672 + name: management + env: + - name: RABBITMQ_DEFAULT_USER + valueFrom: + secretKeyRef: + name: markitdown-server-secret + key: RABBITMQ_USER + - name: RABBITMQ_DEFAULT_PASS + valueFrom: + secretKeyRef: + name: markitdown-server-secret + key: RABBITMQ_PASSWORD + resources: + limits: + memory: "1024Mi" + cpu: "1000m" + requests: + memory: "256Mi" + cpu: "100m" + volumeMounts: + - name: data + mountPath: /var/lib/rabbitmq + readinessProbe: + exec: + command: ["rabbitmq-diagnostics", "-q", "check_port_connectivity"] + initialDelaySeconds: 20 + periodSeconds: 10 + timeoutSeconds: 10 + livenessProbe: + exec: + command: ["rabbitmq-diagnostics", "-q", "ping"] + initialDelaySeconds: 60 + periodSeconds: 30 + timeoutSeconds: 15 + volumes: + - name: data + persistentVolumeClaim: + claimName: markitdown-rabbitmq-data + +--- +apiVersion: v1 +kind: Service +metadata: + name: markitdown-rabbitmq + namespace: markitdown-server + labels: + app: markitdown-rabbitmq +spec: + selector: + app: markitdown-rabbitmq + ports: + - port: 5672 + targetPort: 5672 + name: amqp + - port: 15672 + targetPort: 15672 + name: management + type: ClusterIP diff --git a/k8s/redis.yaml b/k8s/redis.yaml deleted file mode 100644 index bb7a8c0..0000000 --- a/k8s/redis.yaml +++ /dev/null @@ -1,62 +0,0 @@ -apiVersion: apps/v1 -kind: Deployment -metadata: - name: markitdown-redis - namespace: markitdown-server - labels: - app: markitdown-redis -spec: - replicas: 1 - selector: - matchLabels: - app: markitdown-redis - template: - metadata: - labels: - app: markitdown-redis - spec: - containers: - - name: redis - image: redis:7-alpine - # Ephemeral cache only (5 min TTL); cap memory and evict on pressure. - args: - - "--maxmemory" - - "256mb" - - "--maxmemory-policy" - - "allkeys-lru" - - "--save" - - "" - ports: - - containerPort: 6379 - resources: - limits: - memory: "320Mi" - cpu: "250m" - requests: - memory: "64Mi" - cpu: "50m" - livenessProbe: - tcpSocket: - port: 6379 - initialDelaySeconds: 10 - periodSeconds: 10 - readinessProbe: - exec: - command: ["redis-cli", "ping"] - initialDelaySeconds: 5 - periodSeconds: 10 ---- -apiVersion: v1 -kind: Service -metadata: - name: markitdown-redis - namespace: markitdown-server - labels: - app: markitdown-redis -spec: - selector: - app: markitdown-redis - ports: - - port: 6379 - targetPort: 6379 - type: ClusterIP diff --git a/k8s/secrets.example.yaml b/k8s/secrets.example.yaml new file mode 100644 index 0000000..15f647b --- /dev/null +++ b/k8s/secrets.example.yaml @@ -0,0 +1,45 @@ +# Copy to k8s/secrets.yaml (gitignored) and fill in, then: +# +# kubectl apply -f k8s/secrets.yaml +# +# Applied on its own rather than through kustomization.yaml, so that `kubectl +# apply -k k8s` never carries credentials. +# +# Every manifest reads from this one secret. The keys under "required" are +# referenced by secretKeyRef without `optional: true`, so a missing one leaves +# the pod in CreateContainerConfigError instead of falling back to a default. +apiVersion: v1 +kind: Secret +metadata: + name: markitdown-server-secret + namespace: markitdown-server +type: Opaque +stringData: + # ---- Required ---- + # Conversion only calls the LLM when use_llm is requested; a placeholder is + # enough for a cluster that never sets it. + OPENAI_API_KEY: CHANGEME + # The x-api-key callers present to the API. + ADMIN_API_KEY: CHANGEME + # Consumed by rabbitmq.yaml (as RABBITMQ_DEFAULT_USER/PASS) and by the api + # and worker deployments, which interpolate them into RABBITMQ_URL. + RABBITMQ_USER: CHANGEME + RABBITMQ_PASSWORD: CHANGEME + # Both the credentials minio.yaml boots with and the ones the app + # authenticates using — they are the same pair on purpose. + S3_ACCESS_KEY_ID: CHANGEME + S3_SECRET_ACCESS_KEY: CHANGEME + + # ---- Managed Turso / R2 / S3 (optional) ---- + # The default stack runs libsql and minio in-cluster, addressed by + # TURSO_DATABASE_URL and S3_ENDPOINT under `env:` in api-deployment.yaml and + # worker-deployment.yaml. + # + # `env:` takes precedence over `envFrom:`, so uncommenting these alone does + # NOTHING — delete the matching entries from both deployments first, and drop + # libsql.yaml / minio.yaml from kustomization.yaml. + # TURSO_DATABASE_URL: libsql://-.turso.io + # TURSO_AUTH_TOKEN: CHANGEME + # S3_ENDPOINT: https://.r2.cloudflarestorage.com + # S3_BUCKET: markitdown + # S3_REGION: auto diff --git a/k8s/worker-deployment.yaml b/k8s/worker-deployment.yaml new file mode 100644 index 0000000..5388321 --- /dev/null +++ b/k8s/worker-deployment.yaml @@ -0,0 +1,95 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: markitdown-worker + namespace: markitdown-server + labels: + app: markitdown-worker +spec: + replicas: 2 + selector: + matchLabels: + app: markitdown-worker + template: + metadata: + labels: + app: markitdown-worker + spec: + # A worker stops claiming on SIGTERM and finishes the chunk it holds. + # This is a courtesy, not a guarantee: if the pod is killed first, the + # chunk's lease expires and another worker reclaims it. + terminationGracePeriodSeconds: 180 + containers: + - name: markitdown-worker + image: ghcr.io/sirily11/markitdown-server:latest + command: ["python", "-m", "worker"] + resources: + limits: + # One chunk at a time, so peak memory is bounded by chunk size + # rather than document size. This is what lets a fleet of small + # pods replace one 2Gi pod. + memory: "1024Mi" + cpu: "1000m" + requests: + memory: "512Mi" + cpu: "500m" + env: + # NOTE: `env` takes precedence over `envFrom`, so these in-cluster + # defaults win over anything of the same name in the secret. To run + # against hosted Turso or R2, delete the relevant entries and supply + # them via the secret instead — changing only the secret will not + # take effect while they are here. + # RABBITMQ_URL is intentionally absent here: it arrives whole from + # the secret via envFrom, so the broker (in-cluster rabbitmq.yaml + # vs a shared one in another namespace) is a secret-only change. + # It used to be pinned here only because it was assembled from + # $(RABBITMQ_USER) — and $(VAR) expands solely from earlier entries + # in this same list, never from envFrom. + - name: TURSO_DATABASE_URL + value: "http://markitdown-libsql:8080" + - name: S3_ENDPOINT + value: "http://markitdown-minio:9000" + - name: S3_BUCKET + value: "markitdown" + - name: OPENAI_MODEL + value: "google/gemini-3.1-flash-lite-preview" + - name: PDF_PAGES_PER_CHUNK + value: "20" + - name: MAX_ATTEMPTS + value: "3" + # One unacked message at a time: bounds memory to a chunk and keeps + # dispatch fair, so a slow worker stops receiving instead of + # hoarding a backlog while its peers idle. + - name: RABBITMQ_PREFETCH + value: "1" + - name: LEASE_TTL + value: "120" + - name: HEARTBEAT_INTERVAL + value: "30" + - name: SOURCE_CACHE_DIR + value: "/tmp/markitdown-src" + envFrom: + - secretRef: + name: markitdown-server-secret + volumeMounts: + - name: source-cache + mountPath: /tmp + livenessProbe: + # The worker serves no HTTP, so liveness is the freshness of a file + # the main loop touches each iteration. A process that is running + # but no longer claiming work would otherwise look healthy forever. + exec: + command: + - sh + - -c + - test $(( $(date +%s) - $(stat -c %Y /tmp/worker-alive) )) -lt 300 + initialDelaySeconds: 30 + periodSeconds: 30 + timeoutSeconds: 5 + failureThreshold: 3 + volumes: + # Holds the source-file cache. emptyDir so a rescheduled pod starts + # cold rather than inheriting another node's files. + - name: source-cache + emptyDir: + sizeLimit: 4Gi diff --git a/migrations/001_init.sql b/migrations/001_init.sql new file mode 100644 index 0000000..1711dc2 --- /dev/null +++ b/migrations/001_init.sql @@ -0,0 +1,61 @@ +-- Documents that have been fully converted, keyed by content hash + config +-- fingerprint. A hit here means a re-submitted file needs no work at all. +CREATE TABLE IF NOT EXISTS documents ( + doc_key TEXT PRIMARY KEY, + file_hash TEXT NOT NULL, + config_fp TEXT NOT NULL, + total_pages INTEGER NOT NULL, + total_chunks INTEGER NOT NULL, + total_length INTEGER NOT NULL, + page_size INTEGER NOT NULL, + created_at INTEGER NOT NULL, + expires_at INTEGER NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_documents_expires ON documents(expires_at); + +CREATE TABLE IF NOT EXISTS jobs ( + job_id TEXT PRIMARY KEY, + doc_key TEXT NOT NULL, + file_hash TEXT NOT NULL, + url TEXT NOT NULL, + status TEXT NOT NULL, -- queued|running|assembling|done|failed + total_chunks INTEGER NOT NULL, + remaining_chunks INTEGER NOT NULL, + use_llm INTEGER NOT NULL, + error TEXT, + error_chunk INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_jobs_doc_key ON jobs(doc_key); + +-- This table is the work queue. Claiming is a single atomic UPDATE, and +-- because SQLite serialises writers two workers can never claim the same row. +-- +-- Retry, crash recovery, backoff and dead-lettering all fall out of three +-- columns: attempts (incremented at claim time so an OOM-killed worker still +-- burns an attempt), available_at (backoff gate) and lease_expires_at (a lease +-- that stops being renewed is a worker that died). +CREATE TABLE IF NOT EXISTS job_chunks ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + job_id TEXT NOT NULL, + chunk_index INTEGER NOT NULL, + start_page INTEGER NOT NULL, + end_page INTEGER NOT NULL, -- exclusive; -1 means "the whole file" + status TEXT NOT NULL, -- pending|running|done|failed + attempts INTEGER NOT NULL DEFAULT 0, + available_at INTEGER NOT NULL, + lease_owner TEXT, + lease_expires_at INTEGER, + s3_key TEXT, + last_error TEXT, + updated_at INTEGER NOT NULL, + UNIQUE (job_id, chunk_index) +); + +CREATE INDEX IF NOT EXISTS idx_chunks_claim + ON job_chunks(status, available_at, lease_expires_at); + +CREATE INDEX IF NOT EXISTS idx_chunks_job ON job_chunks(job_id, chunk_index); diff --git a/pagination.py b/pagination.py new file mode 100644 index 0000000..ca6f41d --- /dev/null +++ b/pagination.py @@ -0,0 +1,62 @@ +""" +Splitting converted markdown into fixed-size pages. + +Previously part of cache.py, which also owned Redis storage. Storage now lives +in storage.py; this is purely the page-splitting arithmetic. +""" +import config + +PAGE_SIZE = config.PAGE_SIZE + + +def paginate(content: str, page_size: int = PAGE_SIZE) -> list[str]: + """ + Split markdown content into pages of roughly ``page_size`` characters. + + Pages are broken on paragraph boundaries (``\\n\\n``) where possible so a + paragraph is not split across pages. A single paragraph larger than + ``page_size`` is hard-split. + """ + if not content: + return [""] + + pages: list[str] = [] + current = "" + + for block in content.split("\n\n"): + # Hard-split a single block that is itself larger than a page. + while len(block) > page_size: + if current: + pages.append(current) + current = "" + pages.append(block[:page_size]) + block = block[page_size:] + + candidate = block if not current else f"{current}\n\n{block}" + if len(candidate) > page_size: + pages.append(current) + current = block + else: + current = candidate + + if current or not pages: + pages.append(current) + + return pages + + +def build_pagination(doc_id: str, page: int, total_pages: int, total_length: int, + page_size: int = PAGE_SIZE) -> dict: + has_next = page < total_pages + has_prev = page > 1 + return { + "id": doc_id, + "page": page, + "page_size": page_size, + "total_pages": total_pages, + "total_length": total_length, + "has_next": has_next, + "has_prev": has_prev, + "next_page": page + 1 if has_next else None, + "prev_page": page - 1 if has_prev else None, + } diff --git a/pdf_chunk.py b/pdf_chunk.py deleted file mode 100644 index 5f2b31c..0000000 --- a/pdf_chunk.py +++ /dev/null @@ -1,134 +0,0 @@ -""" -Page-chunked, multi-process PDF conversion. - -markitdown's ``PdfConverter`` holds the parsed pdfplumber object graph for the -*whole* document at once (and, for prose PDFs, then re-parses everything with -pdfminer), so both peak memory and wall-clock grow with document size in a -single process. - -Splitting the PDF into page ranges and converting each range in its own process -bounds both: a worker only ever holds N pages of parsed objects, and it returns -that memory to the OS when it exits. - -Both pdfplumber and pdfminer.six are pure Python, so this has to use processes — -threads would serialise on the GIL and buy nothing. -""" -import logging -import os -import shutil -import tempfile -from concurrent.futures import ProcessPoolExecutor -from typing import Optional - -logger = logging.getLogger(__name__) - -# Pages per chunk. Smaller chunks lower peak memory per worker but add -# per-chunk parser startup overhead. -PAGES_PER_CHUNK = int(os.environ.get("PDF_PAGES_PER_CHUNK", "20")) -# Worker processes. Each holds ~140-190 MiB just for the markitdown import, so -# this must be sized against the container memory limit, not just its CPU. -MAX_WORKERS = int(os.environ.get("PDF_MAX_WORKERS", "2")) -# Below this page count the process startup cost outweighs the parallelism. -MIN_PAGES_FOR_PARALLEL = int(os.environ.get("PDF_MIN_PAGES_FOR_PARALLEL", "40")) - - -def page_count(path: str) -> Optional[int]: - """ - Return the number of pages in ``path``, or None if it is not a readable PDF. - - Used both to decide whether chunking is worthwhile and as the PDF sniff, so - a corrupt or non-PDF file transparently falls back to the normal path. - """ - try: - from pypdf import PdfReader - - return len(PdfReader(path).pages) - except Exception: - return None - - -def split_pdf(path: str, pages_per_chunk: int, out_dir: str) -> list[str]: - """ - Split ``path`` into single-file PDFs of ``pages_per_chunk`` pages each. - - Chunks are written to ``out_dir`` and passed to workers by path rather than - by value — pickling the bytes through the process pool would duplicate the - document in memory, which is the thing this module exists to avoid. - """ - from pypdf import PdfReader, PdfWriter - - if pages_per_chunk < 1: - raise ValueError("pages_per_chunk must be >= 1") - - reader = PdfReader(path) - total = len(reader.pages) - chunk_paths: list[str] = [] - - for index, start in enumerate(range(0, total, pages_per_chunk)): - writer = PdfWriter() - for page in reader.pages[start:start + pages_per_chunk]: - writer.add_page(page) - chunk_path = os.path.join(out_dir, f"chunk_{index:05d}.pdf") - with open(chunk_path, "wb") as fh: - writer.write(fh) - chunk_paths.append(chunk_path) - - return chunk_paths - - -def convert_chunk(chunk_path: str) -> str: - """ - Convert one chunk file to markdown. Runs in a worker process. - - Must stay importable at module level so the pool can pickle it. markitdown - is imported lazily here so the parent does not pay for it when the chunked - path is never taken. - """ - from markitdown._stream_info import StreamInfo - from markitdown.converters._pdf_converter import PdfConverter - - stream_info = StreamInfo(extension=".pdf", mimetype="application/pdf") - with open(chunk_path, "rb") as fh: - return PdfConverter().convert(fh, stream_info).markdown - - -def convert_pdf_chunked( - path: str, - pages_per_chunk: int = PAGES_PER_CHUNK, - max_workers: int = MAX_WORKERS, -) -> str: - """ - Convert a PDF by splitting it into page ranges converted in parallel. - - Chunk results are concatenated in page order, so output ordering matches a - whole-document conversion. - """ - work_dir = tempfile.mkdtemp(prefix="pdf_chunks_") - try: - chunk_paths = split_pdf(path, pages_per_chunk, work_dir) - logger.info("split %s into %d chunks of <=%d pages", - path, len(chunk_paths), pages_per_chunk) - - if len(chunk_paths) == 1: - return convert_chunk(chunk_paths[0]) - - # "spawn" rather than the Linux default "fork": this runs inside a - # FastAPI threadpool worker, and forking a process that holds threads - # can deadlock if a lock is held at fork time. - import multiprocessing - - context = multiprocessing.get_context("spawn") - workers = max(1, min(max_workers, len(chunk_paths))) - with ProcessPoolExecutor(max_workers=workers, mp_context=context) as pool: - # map() preserves input order, so pages stay sequential. - parts = list(pool.map(convert_chunk, chunk_paths)) - - return "\n\n".join(part for part in parts if part and part.strip()) - finally: - shutil.rmtree(work_dir, ignore_errors=True) - - -def should_chunk(path: str, min_pages: int = MIN_PAGES_FOR_PARALLEL) -> bool: - """True if ``path`` is a PDF large enough to be worth chunking.""" - pages = page_count(path) - return pages is not None and pages >= min_pages diff --git a/pyproject.toml b/pyproject.toml index 2f76154..11aefa1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,14 +13,30 @@ dependencies = [ "markitdown[all] (>=0.1.5,<0.2.0)", "openai (>=1.82.1,<2.0.0)", "python-multipart (>=0.0.20,<0.0.21)", - "redis (>=5.2.1,<6.0.0)", - "pypdf (>=6.1.1,<7.0.0)" + "pypdf (>=6.1.1,<7.0.0)", + "boto3 (>=1.40.0,<2.0.0)", + # Only needed to talk to hosted Turso; local/CI runs use stdlib sqlite3. + "libsql (>=0.1.11,<0.2.0)", + "pika (>=1.3.0,<2.0.0)", + # Serialising the generated OpenAPI document at /openapi.yaml. Arrives + # transitively via markitdown too, but this is a direct import now. + "PyYAML (>=6.0,<7.0)" ] [project.optional-dependencies] test = [ "pytest (>=8.0.0,<9.0.0)", "reportlab (>=4.0.0,<5.0.0)", + "httpx (>=0.28.0,<0.29.0)", + "moto[s3] (>=5.0.0,<6.0.0)", +] + +[tool.pytest.ini_options] +# Source modules live at the repo root and are imported flat. +pythonpath = ["."] +testpaths = ["tests"] +markers = [ + "slow: exercises real markitdown conversion rather than a stub", ] diff --git a/requirements.txt b/requirements.txt index 7bd5712..2d71fb0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,8 @@ azure-ai-documentintelligence==1.0.2 azure-core==1.34.0 azure-identity==1.23.0 beautifulsoup4==4.13.4 +boto3==1.43.51 +botocore==1.43.51 certifi==2025.4.26 cffi==1.17.1 charset-normalizer==3.4.2 @@ -23,14 +25,16 @@ humanfriendly==10.0 idna==3.10 isodate==0.7.2 jiter==0.10.0 +jmespath==1.1.0 +libsql==0.1.11 lxml==5.4.0 magika==0.6.2 mammoth==1.9.1 markdownify==1.1.0 markitdown==0.1.5 mpmath==1.3.0 -msal==1.32.3 msal-extensions==1.3.1 +msal==1.32.3 numpy==2.2.6 olefile==0.47 onnxruntime==1.22.0 @@ -40,22 +44,24 @@ packaging==25.0 pandas==2.2.3 pdfminer.six==20260107 pdfplumber==0.11.10 +pika==1.4.1 pillow==12.3.0 protobuf==6.31.1 pycparser==2.22 -pypdf==6.1.1 -pypdfium2==5.12.1 pydantic==2.11.5 pydantic_core==2.33.2 pydub==0.25.1 PyJWT==2.10.1 +pypdf==6.1.1 +pypdfium2==5.12.1 python-dateutil==2.9.0.post0 python-dotenv==1.1.0 python-multipart==0.0.20 python-pptx==1.0.2 pytz==2025.2 -redis==5.2.1 +PyYAML==6.0.3 requests==2.32.3 +s3transfer==0.19.1 six==1.17.0 sniffio==1.3.1 soupsieve==2.7 diff --git a/scripts/k8s-benchmark.sh b/scripts/k8s-benchmark.sh index a7266db..078d8df 100644 --- a/scripts/k8s-benchmark.sh +++ b/scripts/k8s-benchmark.sh @@ -14,6 +14,13 @@ # reset in place -- the deployment is restarted before each book instead, which # both gives a true per-book peak and makes every book start from a cold pod. # +# Conversion is asynchronous now, so a book is timed from submit to the job +# reporting `done`, not from a single blocking POST. The memory that matters +# moved with it: the API pod only downloads and hashes, so the tracked pod is +# the *worker*. Worker replicas must be pinned to 1 for the same reason a single +# replica was needed before -- with a fleet, peak memory is spread across pods +# and no single cgroup reading means anything. +# # Expects: kubectl context pointing at the cluster and $ADMIN_API_KEY set. The # port-forward is managed here, since restarting the pod tears it down. set -uo pipefail @@ -21,6 +28,11 @@ set -uo pipefail NAMESPACE="${NAMESPACE:-markitdown-server}" LOCAL_PORT="${LOCAL_PORT:-8080}" BASE_URL="http://localhost:${LOCAL_PORT}" +# The deployment whose memory is measured and which is restarted between books. +WORKER_LABEL="${WORKER_LABEL:-markitdown-worker}" +API_LABEL="${API_LABEL:-markitdown-api}" +# How long to poll a job before declaring it hung. Set per-book from the budget. +POLL_INTERVAL="${POLL_INTERVAL:-2}" # Scales every time budget at once, for slower runners or local use. BUDGET_SCALE="${BUDGET_SCALE:-1.0}" # Fail if peak memory exceeds this share of the pod's limit. @@ -83,7 +95,7 @@ to_mib() { # after a rollout restart the outgoing pod keeps being counted long after the # new one is serving. Filter on deletionTimestamp and readiness instead. running_pod_names() { - kubectl get pods -n "$NAMESPACE" -l app=markitdown-server -o json 2>/dev/null \ + kubectl get pods -n "$NAMESPACE" -l "app=${WORKER_LABEL}" -o json 2>/dev/null \ | jq -r '.items[] | select(.metadata.deletionTimestamp == null) | select(.status.phase == "Running") @@ -147,6 +159,8 @@ slugify() { # Logs every matching pod rather than just the one we tracked -- when pod # resolution is what failed, the logs are exactly what explains why, and # reporting "(no logs captured)" in that case is the least useful outcome. +# Both deployments are logged: a failure can now be in the producer (submit) or +# the consumer (conversion), and guessing wrong wastes a whole CI run. capture_logs() { local slug="$1" pod i=0 while IFS= read -r pod; do @@ -159,41 +173,119 @@ capture_logs() { > "${LOG_DIR}/${slug}.${i}-${pod}.previous.log" 2>/dev/null || true [ -s "${LOG_DIR}/${slug}.${i}-${pod}.previous.log" ] \ || rm -f "${LOG_DIR}/${slug}.${i}-${pod}.previous.log" - done <<< "$(kubectl get pods -n "$NAMESPACE" -l app=markitdown-server \ + done <<< "$(kubectl get pods -n "$NAMESPACE" \ + -l "app in (${WORKER_LABEL},${API_LABEL})" \ -o jsonpath='{range .items[*]}{.metadata.name}{"\n"}{end}' 2>/dev/null)" } +# Only the worker is restarted: it owns the cgroup high-water mark being +# measured, and bouncing the API would just drop the port-forward for nothing. restart_pod() { - kubectl rollout restart deployment/markitdown-server -n "$NAMESPACE" >/dev/null - kubectl rollout status deployment/markitdown-server -n "$NAMESPACE" \ + kubectl rollout restart "deployment/${WORKER_LABEL}" -n "$NAMESPACE" >/dev/null + kubectl rollout status "deployment/${WORKER_LABEL}" -n "$NAMESPACE" \ --timeout=300s >/dev/null || return 1 start_port_forward } -MEM_LIMIT_RAW=$(kubectl get deployment markitdown-server -n "$NAMESPACE" \ +# Submit a book and poll until the job reaches a terminal state. +# +# Echoes " "; status is `done`, `failed`, `timeout`, or +# `http:` when the submission itself was rejected. +convert_book() { + local url="$1" deadline="$2" submit job_id code status doc_key + + submit="$(curl -sS --max-time 120 -w '\n%{http_code}' \ + -X POST "${BASE_URL}/async/convert" \ + -H "x-api-key: ${ADMIN_API_KEY}" \ + -H 'Content-Type: application/json' \ + -d "{\"file\": \"${url}\"}" 2>&1)" + code="$(printf '%s' "$submit" | tail -n1)" + payload="$(printf '%s' "$submit" | sed '$d')" + + # 200 means a cache hit served it outright, which is a valid (and instant) + # outcome; 202 is the normal queued path. + if [ "$code" = "200" ]; then + printf 'done %s' "$(printf '%s' "$payload" | python3 -c \ + 'import json,sys; print(json.load(sys.stdin)["pagination"]["total_length"])' \ + 2>/dev/null || echo '-')" + return 0 + fi + if [ "$code" != "202" ]; then + printf 'http:%s -' "$code" + return 0 + fi + + job_id="$(printf '%s' "$payload" | python3 -c \ + 'import json,sys; print(json.load(sys.stdin)["job_id"])' 2>/dev/null)" + doc_key="$(printf '%s' "$payload" | python3 -c \ + 'import json,sys; print(json.load(sys.stdin)["doc_key"])' 2>/dev/null)" + if [ -z "$job_id" ]; then + printf 'http:bad-submit-payload -' + return 0 + fi + + while [ "$(date +%s)" -lt "$deadline" ]; do + status="$(curl -sS --max-time 30 \ + "${BASE_URL}/convert/jobs/${job_id}" \ + -H "x-api-key: ${ADMIN_API_KEY}" 2>/dev/null \ + | python3 -c 'import json,sys; print(json.load(sys.stdin)["status"])' \ + 2>/dev/null)" + case "$status" in + done) + printf 'done %s' "$(curl -sS --max-time 30 \ + "${BASE_URL}/convert/${doc_key}/pages/1" \ + -H "x-api-key: ${ADMIN_API_KEY}" 2>/dev/null \ + | python3 -c \ + 'import json,sys; print(json.load(sys.stdin)["pagination"]["total_length"])' \ + 2>/dev/null || echo '-')" + return 0 + ;; + failed) + printf 'failed -' + return 0 + ;; + esac + sleep "$POLL_INTERVAL" + done + printf 'timeout -' +} + +MEM_LIMIT_RAW=$(kubectl get deployment "${WORKER_LABEL}" -n "$NAMESPACE" \ -o jsonpath='{.spec.template.spec.containers[0].resources.limits.memory}' 2>/dev/null) MEM_LIMIT_MB=$(to_mib "$MEM_LIMIT_RAW") [ -z "$MEM_LIMIT_MB" ] && MEM_LIMIT_MB=0 mkdir -p "$LOG_DIR" -# The benchmark assumes a single replica: with more than one pod behind the -# service, the request and the cgroup read can land on different pods and every -# per-book number becomes meaningless. An active HPA breaks that assumption the -# moment conversion saturates the CPU, so fail up front with the actual reason -# rather than timing out later waiting for the pod set to settle. -HPA_MAX="$(kubectl get hpa markitdown-server-hpa -n "$NAMESPACE" \ +# The benchmark assumes a single *worker* replica: peak memory is read from one +# pod's cgroup, and with a fleet the work is spread across pods so that number +# describes nothing. An active HPA breaks the assumption the moment conversion +# saturates the CPU, so fail up front with the actual reason rather than timing +# out later waiting for the pod set to settle. +# +# The API HPA is irrelevant here and deliberately not checked -- the API is +# stateless and any pod can answer, which is the whole point of the split. +HPA_MAX="$(kubectl get hpa "${WORKER_LABEL}-hpa" -n "$NAMESPACE" \ -o jsonpath='{.spec.maxReplicas}' 2>/dev/null)" if [ -n "$HPA_MAX" ] && [ "$HPA_MAX" -gt 1 ]; then - echo "ERROR: HorizontalPodAutoscaler allows up to ${HPA_MAX} replicas." + echo "ERROR: HorizontalPodAutoscaler allows up to ${HPA_MAX} worker replicas." echo " It will scale up under conversion load and invalidate the" echo " measurements. Delete it before benchmarking:" - echo " kubectl delete hpa markitdown-server-hpa -n ${NAMESPACE}" + echo " kubectl delete hpa ${WORKER_LABEL}-hpa -n ${NAMESPACE}" + exit 1 +fi + +WORKER_REPLICAS="$(kubectl get deployment "${WORKER_LABEL}" -n "$NAMESPACE" \ + -o jsonpath='{.spec.replicas}' 2>/dev/null)" +if [ -n "$WORKER_REPLICAS" ] && [ "$WORKER_REPLICAS" -gt 1 ]; then + echo "ERROR: ${WORKER_LABEL} has ${WORKER_REPLICAS} replicas; peak memory" + echo " would be split across pods. Pin it to one:" + echo " kubectl scale deployment/${WORKER_LABEL} -n ${NAMESPACE} --replicas=1" exit 1 fi echo "namespace=$NAMESPACE memory_limit=${MEM_LIMIT_MB}Mi (${MEM_LIMIT_RAW}) budget_scale=${BUDGET_SCALE}" -echo "(pod restarted before each book, so PEAK_MB is that book's own high-water mark)" +echo "(worker restarted before each book, so PEAK_MB is that book's own high-water mark)" echo printf '%-30s %7s %8s %8s %8s %9s %s\n' \ BOOK SECONDS BUDGET PEAK_MB LIMIT_PCT CHARS VERDICT @@ -221,7 +313,7 @@ for entry in "${BOOKS[@]}"; do if ! pod="$(pod_name)" || [ -z "$pod" ]; then printf '%-30s %7s %8s %8s %8s %9s %s\n' \ "$name" "-" "$budget" "-" "-" "-" "FAIL could not resolve a single serving pod" - kubectl get pods -n "$NAMESPACE" -l app=markitdown-server 2>&1 | head -10 + kubectl get pods -n "$NAMESPACE" -l "app in (${WORKER_LABEL},${API_LABEL})" 2>&1 | head -10 capture_logs "$slug" FAILED_BOOKS+=("$slug|$name") failures=$((failures + 1)) @@ -231,18 +323,13 @@ for entry in "${BOOKS[@]}"; do [ -z "$restarts_before" ] && restarts_before=0 start=$(date +%s.%N) - # Cap the request just past the budget: a hang should fail on the budget, not - # sit until the job-level timeout kills the whole run. + # Cap the poll just past the budget: a hung job should fail on the budget, + # not sit until the job-level timeout kills the whole run. max_time=$(awk -v b="$budget" 'BEGIN{printf "%.0f", b*1.5+60}') - body="$(curl -sS --max-time "$max_time" -w '\n%{http_code}' \ - -X POST "${BASE_URL}/convert" \ - -H "x-api-key: ${ADMIN_API_KEY}" \ - -H 'Content-Type: application/json' \ - -d "{\"file\": \"${url}\"}" 2>&1)" + deadline=$(( $(date +%s) + max_time )) + read -r status chars <<< "$(convert_book "$url" "$deadline")" end=$(date +%s.%N) - status="$(printf '%s' "$body" | tail -n1)" - payload="$(printf '%s' "$body" | sed '$d')" seconds="$(awk -v a="$start" -v b="$end" 'BEGIN{printf "%.1f", b-a}')" pod_after="$(pod_name)" || pod_after="" @@ -267,8 +354,8 @@ for entry in "${BOOKS[@]}"; do elif [ "$peak" = "?" ]; then verdict="FAIL could not read peak memory from the pod cgroup" failures=$((failures + 1)) - elif [ "$status" != "200" ]; then - verdict="FAIL http=$status" + elif [ "$status" != "done" ]; then + verdict="FAIL job status=$status" failures=$((failures + 1)) elif [ "$(awk -v s="$seconds" -v b="$budget" 'BEGIN{print (s>b)?1:0}')" = "1" ]; then verdict="FAIL over time budget" @@ -286,17 +373,9 @@ for entry in "${BOOKS[@]}"; do pct="${raw_pct}%" fi - chars="$(printf '%s' "$payload" | python3 -c \ - 'import json,sys; print(json.load(sys.stdin)["pagination"]["total_length"])' \ - 2>/dev/null || echo "-")" - printf '%-30s %7s %8s %8s %8s %9s %s\n' \ "$name" "$seconds" "$budget" "$peak" "$pct" "$chars" "$verdict" - if [ "$status" != "200" ]; then - echo " $(printf '%s' "$payload" | head -c 300)" - fi - capture_logs "$slug" if [ "$verdict" != "ok" ]; then FAILED_BOOKS+=("$slug|$name") @@ -328,7 +407,7 @@ if [ "$failures" -gt 0 ]; then fi echo echo "--- pod status ---" - kubectl get pods -n "$NAMESPACE" -l app=markitdown-server 2>&1 | head -10 + kubectl get pods -n "$NAMESPACE" -l "app in (${WORKER_LABEL},${API_LABEL})" 2>&1 | head -10 exit 1 fi echo "All books converted within time and memory budgets." diff --git a/scripts/local-smoke.sh b/scripts/local-smoke.sh new file mode 100755 index 0000000..8be45d6 --- /dev/null +++ b/scripts/local-smoke.sh @@ -0,0 +1,114 @@ +#!/usr/bin/env bash +# +# End-to-end smoke test against a locally running api + worker. +# +# Submits a real PDF, polls the job to completion, then walks every page back +# out and checks the content actually survived. This exercises the whole path +# the k8s benchmark does — download, hash, chunk, claim, convert, assemble, +# paginate — without needing a cluster. +# +# Usage: +# make up +# make api # in another shell +# make worker # in a third shell +# make smoke +set -uo pipefail + +BASE_URL="${BASE_URL:-http://localhost:8000}" +API_KEY="${ADMIN_API_KEY:-dev-key}" +# Small enough to finish quickly, large enough to split into several chunks at +# the default 20 pages/chunk. +PDF_URL="${PDF_URL:-https://raw.githubusercontent.com/GSimas/Asimov/master/Books/I_Robot/I%20Robot.pdf}" +TIMEOUT="${TIMEOUT:-300}" + +fail() { echo "FAIL: $*" >&2; exit 1; } + +curl -sf "${BASE_URL}/" >/dev/null 2>&1 \ + || fail "api is not reachable at ${BASE_URL} (run 'make api')" + +ready="$(curl -sS "${BASE_URL}/readyz" 2>/dev/null)" +echo "$ready" | grep -q '"status": *"ok"' \ + || fail "dependencies are not ready (run 'make up'): $ready" + +echo "submitting ${PDF_URL}" +submit="$(curl -sS -w '\n%{http_code}' \ + -X POST "${BASE_URL}/async/convert" \ + -H "x-api-key: ${API_KEY}" \ + -H 'Content-Type: application/json' \ + -d "{\"file\": \"${PDF_URL}\"}")" + +code="$(printf '%s' "$submit" | tail -n1)" +payload="$(printf '%s' "$submit" | sed '$d')" + +case "$code" in + 200) + echo "cache hit — already converted. 'make reset' to start clean." + printf '%s' "$payload" | python3 -m json.tool | head -20 + exit 0 + ;; + 202) ;; + *) fail "submit returned HTTP ${code}: $(printf '%s' "$payload" | head -c 400)" ;; +esac + +job_id="$(printf '%s' "$payload" | python3 -c 'import json,sys; print(json.load(sys.stdin)["job_id"])')" +doc_key="$(printf '%s' "$payload" | python3 -c 'import json,sys; print(json.load(sys.stdin)["doc_key"])')" +chunks="$(printf '%s' "$payload" | python3 -c 'import json,sys; print(json.load(sys.stdin)["total_chunks"])')" +echo "job ${job_id} queued as ${chunks} chunk(s)" + +start=$(date +%s) +deadline=$(( start + TIMEOUT )) +last="" +while [ "$(date +%s)" -lt "$deadline" ]; do + job="$(curl -sS "${BASE_URL}/convert/jobs/${job_id}" -H "x-api-key: ${API_KEY}")" + status="$(printf '%s' "$job" | python3 -c 'import json,sys; print(json.load(sys.stdin)["status"])' 2>/dev/null)" + progress="$(printf '%s' "$job" | python3 -c \ + 'import json,sys; p=json.load(sys.stdin)["progress"]; print(f"{p[\"completed\"]}/{p[\"total\"]} ({p[\"percent\"]}%)")' 2>/dev/null)" + + if [ "$progress" != "$last" ]; then + echo " ${status}: ${progress}" + last="$progress" + fi + + case "$status" in + done) break ;; + failed) + fail "job failed: $(printf '%s' "$job" | python3 -c 'import json,sys; print(json.load(sys.stdin)["error"])')" + ;; + esac + sleep 2 +done + +[ "$status" = "done" ] || fail "job did not finish within ${TIMEOUT}s (last status: ${status:-unknown}). Is 'make worker' running?" + +elapsed=$(( $(date +%s) - start )) +echo "converted in ${elapsed}s" + +# Walk every page back out, so a job that reports done but stored nothing +# readable still fails the smoke test. +page1="$(curl -sS "${BASE_URL}/convert/${doc_key}/pages/1" -H "x-api-key: ${API_KEY}")" +total_pages="$(printf '%s' "$page1" | python3 -c 'import json,sys; print(json.load(sys.stdin)["pagination"]["total_pages"])' 2>/dev/null)" \ + || fail "could not read page 1: $(printf '%s' "$page1" | head -c 300)" +total_length="$(printf '%s' "$page1" | python3 -c 'import json,sys; print(json.load(sys.stdin)["pagination"]["total_length"])')" + +chars=0 +for page in $(seq 1 "$total_pages"); do + body="$(curl -sS "${BASE_URL}/convert/${doc_key}/pages/${page}" -H "x-api-key: ${API_KEY}")" + n="$(printf '%s' "$body" | python3 -c 'import json,sys; print(len(json.load(sys.stdin)["content"]))' 2>/dev/null)" \ + || fail "page ${page} of ${total_pages} was not readable" + chars=$(( chars + n )) +done + +echo "read ${total_pages} page(s), ${chars} chars (manifest says ${total_length})" +[ "$chars" -gt 0 ] || fail "document is empty" + +# Re-submitting identical bytes must hit the cache and queue nothing. +echo "checking the content-addressed cache" +again_code="$(curl -sS -o /dev/null -w '%{http_code}' \ + -X POST "${BASE_URL}/async/convert" \ + -H "x-api-key: ${API_KEY}" \ + -H 'Content-Type: application/json' \ + -d "{\"file\": \"${PDF_URL}\"}")" +[ "$again_code" = "200" ] || fail "expected a 200 cache hit on resubmit, got ${again_code}" + +echo +echo "PASS: converted in ${elapsed}s, ${total_pages} pages, cache hit on resubmit" diff --git a/storage.py b/storage.py new file mode 100644 index 0000000..aa8f5ec --- /dev/null +++ b/storage.py @@ -0,0 +1,240 @@ +""" +Content-addressed object storage for sources, chunk output and final pages. + +Keys are derived from the file's content hash rather than its URL, so the same +document submitted from two different URLs converts once. The key also folds in +a fingerprint of every setting that changes the output, so bumping markitdown +or flipping the LLM flag invalidates cached results automatically instead of +serving stale markdown forever. + +Because keys are content-addressed, writes are idempotent: a chunk that gets +converted twice (a lease expiring while its owner is still alive) writes +identical bytes to an identical key, so duplicates are harmless rather than +corrupting. +""" +import hashlib +import json +import logging +import threading +from typing import Optional + +import boto3 +from botocore.config import Config as BotoConfig +from botocore.exceptions import ClientError + +import config + +logger = logging.getLogger(__name__) + +_local = threading.local() +_CHUNK_SIZE = 1024 * 1024 + + +def get_client(): + """Return this thread's S3 client (boto3 clients are not thread-safe).""" + client = getattr(_local, "client", None) + if client is None: + client = boto3.client( + "s3", + endpoint_url=config.S3_ENDPOINT, + region_name=config.S3_REGION, + aws_access_key_id=config.S3_ACCESS_KEY_ID, + aws_secret_access_key=config.S3_SECRET_ACCESS_KEY, + config=BotoConfig(retries={"max_attempts": 3, "mode": "standard"}), + ) + _local.client = client + return client + + +def reset() -> None: + """Drop the cached client (used by tests).""" + _local.client = None + + +# ---- Key derivation ---- + + +def hash_file(path: str) -> str: + """SHA-256 of a file's contents, streamed so memory stays bounded.""" + digest = hashlib.sha256() + with open(path, "rb") as fh: + for block in iter(lambda: fh.read(_CHUNK_SIZE), b""): + digest.update(block) + return digest.hexdigest() + + +def config_fingerprint(use_llm: bool) -> str: + """ + Fingerprint everything that changes conversion output. + + markitdown's version is included because an upgrade can legitimately + produce different markdown for the same input; without it, cached results + from the old version would be served indefinitely. + """ + try: + from importlib.metadata import version + + markitdown_version = version("markitdown") + except Exception: + markitdown_version = "unknown" + + parts = "|".join([ + markitdown_version, + "llm" if use_llm else "nollm", + config.OPENAI_MODEL if use_llm else "-", + str(config.PAGE_SIZE), + str(config.PAGES_PER_CHUNK), + str(config.MIN_PAGES_FOR_PARALLEL), + ]) + return hashlib.sha256(parts.encode()).hexdigest()[:16] + + +def make_doc_key(file_hash: str, use_llm: bool) -> str: + return f"{file_hash}:{config_fingerprint(use_llm)}" + + +def source_key(file_hash: str) -> str: + return f"src/{file_hash}" + + +def part_key(doc_key: str, chunk_index: int) -> str: + return f"parts/{doc_key}/{chunk_index:05d}.md" + + +def manifest_key(doc_key: str) -> str: + return f"docs/{doc_key}/manifest.json" + + +def page_key(doc_key: str, page: int) -> str: + return f"docs/{doc_key}/pages/{page}.md" + + +# ---- Operations ---- + + +def _get_bytes(key: str) -> Optional[bytes]: + try: + response = get_client().get_object(Bucket=config.S3_BUCKET, Key=key) + return response["Body"].read() + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") in ("NoSuchKey", "404"): + return None + raise + + +def _put_bytes(key: str, data: bytes, content_type: str) -> None: + get_client().put_object( + Bucket=config.S3_BUCKET, Key=key, Body=data, ContentType=content_type + ) + + +def ensure_bucket() -> None: + """ + Create the bucket if it does not exist. + + Idempotent and safe to call from every process on boot, which is what + replaces the separate bucket-setup job. A pre-existing bucket (R2, S3, or a + minio volume that survived a restart) short-circuits on the head_bucket. + """ + client = get_client() + try: + client.head_bucket(Bucket=config.S3_BUCKET) + return + except ClientError: + pass + + try: + client.create_bucket(Bucket=config.S3_BUCKET) + logger.info("created bucket %s", config.S3_BUCKET) + except ClientError as exc: + # Another process won the race, or the credentials are read-only + # against a bucket that already exists. Both are fine. + code = exc.response.get("Error", {}).get("Code") + if code not in ("BucketAlreadyOwnedByYou", "BucketAlreadyExists"): + raise + + +def exists(key: str) -> bool: + try: + get_client().head_object(Bucket=config.S3_BUCKET, Key=key) + return True + except ClientError: + return False + + +def put_source(file_hash: str, path: str) -> str: + """ + Upload a source file, skipping the transfer if it is already stored. + + Content-addressed, so an existing object with this key is by definition the + same bytes. + """ + key = source_key(file_hash) + if exists(key): + logger.info("source %s already stored", file_hash[:12]) + return key + with open(path, "rb") as fh: + get_client().put_object(Bucket=config.S3_BUCKET, Key=key, Body=fh) + return key + + +def get_source(file_hash: str, dest_path: str) -> str: + """Download a source file to ``dest_path``.""" + get_client().download_file(config.S3_BUCKET, source_key(file_hash), dest_path) + return dest_path + + +def put_part(doc_key: str, chunk_index: int, markdown: str) -> str: + key = part_key(doc_key, chunk_index) + _put_bytes(key, markdown.encode("utf-8"), "text/markdown; charset=utf-8") + return key + + +def get_part(doc_key: str, chunk_index: int) -> Optional[str]: + data = _get_bytes(part_key(doc_key, chunk_index)) + return data.decode("utf-8") if data is not None else None + + +def put_document(doc_key: str, pages: list[str], total_length: int) -> dict: + """ + Store every page plus a manifest, and return the manifest. + + Pages are written before the manifest so a manifest is never visible + pointing at pages that do not exist yet. + """ + for index, content in enumerate(pages, start=1): + _put_bytes( + page_key(doc_key, index), + content.encode("utf-8"), + "text/markdown; charset=utf-8", + ) + + manifest = { + "doc_key": doc_key, + "total_pages": len(pages), + "total_length": total_length, + "page_size": config.PAGE_SIZE, + } + _put_bytes( + manifest_key(doc_key), json.dumps(manifest).encode("utf-8"), "application/json" + ) + return manifest + + +def get_manifest(doc_key: str) -> Optional[dict]: + data = _get_bytes(manifest_key(doc_key)) + return json.loads(data) if data is not None else None + + +def get_page(doc_key: str, page: int) -> Optional[str]: + data = _get_bytes(page_key(doc_key, page)) + return data.decode("utf-8") if data is not None else None + + +def health() -> bool: + try: + get_client().head_bucket(Bucket=config.S3_BUCKET) + return True + except Exception: + logger.warning("storage health check failed", exc_info=True) + return False diff --git a/taskqueue.py b/taskqueue.py new file mode 100644 index 0000000..4630bc3 --- /dev/null +++ b/taskqueue.py @@ -0,0 +1,269 @@ +""" +RabbitMQ job distribution. + +The producer publishes one message per chunk; workers consume them with a +prefetch of 1, so the broker hands each chunk to whichever worker is free. +Adding a worker pod adds capacity with no coordination — that is what RabbitMQ +buys over a database-polled queue. + +Division of labour with the database: + +* **RabbitMQ owns delivery** — dispatch, redelivery when a worker dies without + acking, retry scheduling, and dead-lettering. +* **Turso owns state** — the per-job remaining-chunk counter, progress, and the + document manifest. A broker cannot answer "did I just finish the last of 25 + chunks?" atomically, and that question decides who assembles the document. + +Delivery is at-least-once: a worker that converts a chunk and then dies before +acking will have that chunk redelivered. jobstore.complete_chunk is written to +be idempotent for exactly this reason, and storage keys are content-addressed +so a duplicate conversion rewrites identical bytes. + +Topology:: + + markitdown (direct) -> markitdown.chunks + markitdown.retry (direct) -> markitdown.retry.{1,2,3} + each with a TTL, dead-lettering back + into markitdown.chunks + markitdown.dlx (direct) -> markitdown.chunks.failed +""" +import json +import logging +import threading +from dataclasses import asdict, dataclass +from typing import Callable, Optional + +import pika + +import config + +logger = logging.getLogger(__name__) + +RETRY_EXCHANGE = f"{config.RABBITMQ_EXCHANGE}.retry" +DLX_EXCHANGE = f"{config.RABBITMQ_EXCHANGE}.dlx" +FAILED_QUEUE = f"{config.RABBITMQ_QUEUE}.failed" + +_local = threading.local() + + +@dataclass(frozen=True) +class ChunkTask: + """One unit of work: a page range of one document.""" + + job_id: str + chunk_index: int + doc_key: str + file_hash: str + start_page: int + end_page: int # exclusive; -1 means the whole file + total_chunks: int + use_llm: bool + attempt: int = 1 + + def to_bytes(self) -> bytes: + return json.dumps(asdict(self)).encode() + + @classmethod + def from_bytes(cls, raw: bytes) -> "ChunkTask": + return cls(**json.loads(raw)) + + def next_attempt(self) -> "ChunkTask": + return ChunkTask(**{**asdict(self), "attempt": self.attempt + 1}) + + +def retry_queue_name(attempt: int) -> str: + """Retry queue for a given attempt, clamped to the last tier.""" + tier = min(attempt, len(config.RETRY_DELAYS)) + return f"{config.RABBITMQ_QUEUE}.retry.{tier}" + + +def connect() -> "pika.BlockingConnection": + """Open a connection for this thread, declaring the topology once.""" + connection = getattr(_local, "connection", None) + if connection is not None and connection.is_open: + return connection + + params = pika.URLParameters(config.RABBITMQ_URL) + params.heartbeat = config.RABBITMQ_HEARTBEAT + # Without this, a broker that is up but refusing connections (still booting) + # fails the whole process instead of settling. + params.blocked_connection_timeout = 30 + connection = pika.BlockingConnection(params) + _local.connection = connection + _local.channel = connection.channel() + declare(_local.channel) + return connection + + +def channel() -> "pika.adapters.blocking_connection.BlockingChannel": + connect() + return _local.channel + + +def close() -> None: + connection = getattr(_local, "connection", None) + if connection is not None: + try: + if connection.is_open: + connection.close() + except Exception: + logger.debug("error closing rabbitmq connection", exc_info=True) + _local.connection = None + _local.channel = None + + +def declare(ch) -> None: + """ + Declare exchanges and queues. Idempotent, so every process does it on boot. + + Everything is durable and messages are published persistent, so a broker + restart does not lose queued work. + """ + ch.exchange_declare(config.RABBITMQ_EXCHANGE, "direct", durable=True) + ch.exchange_declare(RETRY_EXCHANGE, "direct", durable=True) + ch.exchange_declare(DLX_EXCHANGE, "direct", durable=True) + + # Main work queue. Rejected messages land in the DLX. + ch.queue_declare( + config.RABBITMQ_QUEUE, + durable=True, + arguments={"x-dead-letter-exchange": DLX_EXCHANGE}, + ) + ch.queue_bind(config.RABBITMQ_QUEUE, config.RABBITMQ_EXCHANGE, + routing_key=config.RABBITMQ_QUEUE) + + # One retry queue per delay tier. A message sits here for the tier's TTL, + # then dead-letters back into the main queue for another attempt. + for tier, delay in enumerate(config.RETRY_DELAYS, start=1): + name = f"{config.RABBITMQ_QUEUE}.retry.{tier}" + ch.queue_declare( + name, + durable=True, + arguments={ + "x-message-ttl": delay * 1000, + "x-dead-letter-exchange": config.RABBITMQ_EXCHANGE, + "x-dead-letter-routing-key": config.RABBITMQ_QUEUE, + }, + ) + ch.queue_bind(name, RETRY_EXCHANGE, routing_key=name) + + # Poison messages. Retained rather than dropped so a failure can be read + # back after the fact. + ch.queue_declare(FAILED_QUEUE, durable=True) + ch.queue_bind(FAILED_QUEUE, DLX_EXCHANGE, routing_key=config.RABBITMQ_QUEUE) + + +_PERSISTENT = pika.BasicProperties( + delivery_mode=pika.DeliveryMode.Persistent, + content_type="application/json", +) + + +def publish(tasks: list[ChunkTask]) -> None: + """Publish chunk tasks onto the work queue.""" + ch = channel() + for task in tasks: + ch.basic_publish( + exchange=config.RABBITMQ_EXCHANGE, + routing_key=config.RABBITMQ_QUEUE, + body=task.to_bytes(), + properties=_PERSISTENT, + ) + + +def publish_retry(task: ChunkTask) -> None: + """Schedule a failed chunk for another attempt after its backoff delay.""" + ch = channel() + queue = retry_queue_name(task.attempt) + ch.basic_publish( + exchange=RETRY_EXCHANGE, + routing_key=queue, + body=task.next_attempt().to_bytes(), + properties=_PERSISTENT, + ) + logger.info("chunk %d of job %s retrying via %s", + task.chunk_index, task.job_id, queue) + + +def publish_failed(task: ChunkTask, error: str) -> None: + """Park a chunk that exhausted its attempts on the failed queue.""" + ch = channel() + ch.basic_publish( + exchange=DLX_EXCHANGE, + routing_key=config.RABBITMQ_QUEUE, + body=json.dumps({**asdict(task), "error": error[:2000]}).encode(), + properties=_PERSISTENT, + ) + + +def consume(handler: Callable[[ChunkTask], None], should_stop: Callable[[], bool]) -> None: + """ + Consume chunk tasks until ``should_stop()`` returns True. + + Uses ``consume`` with an inactivity timeout rather than ``start_consuming`` + so the loop regains control periodically — otherwise a SIGTERM could not be + honoured until the next message arrived. + + Messages are acked only after the handler returns. A worker that dies + mid-chunk therefore has its message redelivered to another worker, which is + what replaces lease expiry in a database-polled queue. + """ + ch = channel() + ch.basic_qos(prefetch_count=config.RABBITMQ_PREFETCH) + + for method, _properties, body in ch.consume( + config.RABBITMQ_QUEUE, + inactivity_timeout=config.RABBITMQ_CONSUME_TIMEOUT, + ): + if should_stop(): + break + if method is None: # inactivity timeout, no message + continue + + try: + task = ChunkTask.from_bytes(body) + except Exception: + # Unparseable message: ack it away rather than let it redeliver + # forever, but keep a copy on the failed queue. + logger.exception("dropping malformed message") + ch.basic_publish( + exchange=DLX_EXCHANGE, + routing_key=config.RABBITMQ_QUEUE, + body=body, + properties=_PERSISTENT, + ) + ch.basic_ack(method.delivery_tag) + continue + + try: + handler(task) + finally: + # Ack regardless: the handler owns its own retry decision and has + # already republished if it wants another attempt. Nacking here + # too would duplicate the message. + ch.basic_ack(method.delivery_tag) + + try: + ch.cancel() + except Exception: + logger.debug("error cancelling consumer", exc_info=True) + + +def queue_depth() -> Optional[int]: + """Messages waiting on the work queue — the metric to autoscale workers on.""" + try: + result = channel().queue_declare( + config.RABBITMQ_QUEUE, durable=True, passive=True + ) + return result.method.message_count + except Exception: + logger.warning("could not read queue depth", exc_info=True) + return None + + +def health() -> bool: + try: + return connect().is_open + except Exception: + logger.warning("rabbitmq health check failed", exc_info=True) + return False diff --git a/test_main.http b/test_main.http index c9ac2f6..971ecdd 100644 --- a/test_main.http +++ b/test_main.http @@ -5,9 +5,16 @@ Accept: application/json ### -# Convert a (possibly long) document to markdown. -# Returns the first page + pagination metadata. The full paginated content -# is cached in Redis for 5 minutes under the returned `id`. +# Readiness: checks the database and object store. +GET http://127.0.0.1:8000/readyz +Accept: application/json + +### + +# Convert a document synchronously. +# +# Workers still convert chunks in parallel, but this request waits for all of +# them and returns the original {id, content, pagination} response. POST http://127.0.0.1:8000/convert Content-Type: application/json x-api-key: NxQUCbbezCZ8KUhMtg5s0MQTQU4gDQ7 @@ -18,9 +25,29 @@ x-api-key: NxQUCbbezCZ8KUhMtg5s0MQTQU4gDQ7 ### -# Fetch a subsequent page of a previously converted document. -# Use the `id` and `pagination.next_page` returned by /convert. -GET http://127.0.0.1:8000/convert/REPLACE_WITH_DOC_ID/pages/2 +# Submit a document asynchronously. Returns 202 with a `job_id` and `doc_key`; +# if this exact file was converted before, returns the first page immediately. +POST http://127.0.0.1:8000/async/convert +Content-Type: application/json +x-api-key: NxQUCbbezCZ8KUhMtg5s0MQTQU4gDQ7 + +{ + "file": "https://example.com/some-long-document.pdf" +} + +### + +# Poll a job. `status` goes queued -> running -> done (or failed), and +# `progress` reports how many chunks have finished. +GET http://127.0.0.1:8000/convert/jobs/REPLACE_WITH_JOB_ID +Accept: application/json +x-api-key: NxQUCbbezCZ8KUhMtg5s0MQTQU4gDQ7 + +### + +# Fetch a page once the job is done. Use the `doc_key` from the submission +# response and `pagination.next_page` to walk the document. +GET http://127.0.0.1:8000/convert/REPLACE_WITH_DOC_KEY/pages/1 Accept: application/json x-api-key: NxQUCbbezCZ8KUhMtg5s0MQTQU4gDQ7 diff --git a/tests/conftest.py b/tests/conftest.py index 431f3d0..87d83ab 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,6 @@ +import os +import threading + import pytest from reportlab.lib.pagesizes import letter from reportlab.pdfgen import canvas @@ -43,3 +46,128 @@ def small_pdf(make_pdf): def multi_chunk_pdf(make_pdf): """25 pages — splits into multiple chunks at the default 20/chunk.""" return make_pdf(25) + + +@pytest.fixture +def database(tmp_path, monkeypatch): + """ + A fresh SQLite database per test. + + A real file rather than ":memory:" so the concurrency tests, which open a + connection per thread, all see the same database and exercise real write + contention. + """ + import config + import db + + monkeypatch.setattr(config, "DATABASE_URL", str(tmp_path / "test.db")) + monkeypatch.setattr(config, "DATABASE_AUTH_TOKEN", None) + db.reset() + db.migrate() + yield db + db.reset() + + +@pytest.fixture +def s3(monkeypatch): + """An in-memory S3 via moto, with the bucket already created.""" + import boto3 + from moto import mock_aws + + import config + import storage + + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test") + monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1") + monkeypatch.setattr(config, "S3_ENDPOINT", None) + monkeypatch.setattr(config, "S3_REGION", "us-east-1") + monkeypatch.setattr(config, "S3_BUCKET", "test-bucket") + monkeypatch.setattr(config, "S3_ACCESS_KEY_ID", "test") + monkeypatch.setattr(config, "S3_SECRET_ACCESS_KEY", "test") + + with mock_aws(): + boto3.client("s3", region_name="us-east-1").create_bucket(Bucket="test-bucket") + storage.reset() + yield storage + storage.reset() + + +class FakeBroker: + """ + In-memory stand-in for RabbitMQ. + + Models only what the code depends on: FIFO delivery, one delivery at a + time, and retry/failed queues that messages are republished onto. Delays + are not simulated — retried messages become immediately available, since + the delay is RabbitMQ's TTL and not our logic. + """ + + def __init__(self): + self._lock = threading.Lock() + self.ready: list = [] + self.retried: list = [] + self.failed: list = [] + self.published: list = [] + + def publish(self, tasks): + for task in tasks: + self.ready.append(task) + self.published.append(task) + + def publish_retry(self, task): + nxt = task.next_attempt() + self.retried.append(nxt) + self.ready.append(nxt) + + def publish_failed(self, task, error): + self.failed.append((task, error)) + + def _take(self): + """Pop one message, or None. Locked: the integration test drains from + several threads, and two of them must never get the same message.""" + with self._lock: + return self.ready.pop(0) if self.ready else None + + def consume(self, handler, should_stop): + while not should_stop(): + task = self._take() + if task is None: + return + handler(task) + + def drain(self, handler, limit: int = 200): + """Deliver every message, including ones handlers republish.""" + for _ in range(limit): + task = self._take() + if task is None: + return + handler(task) + raise AssertionError("drain did not converge — retry loop?") + + +@pytest.fixture +def broker(monkeypatch): + """Replace RabbitMQ with an in-memory broker.""" + import taskqueue + + fake = FakeBroker() + monkeypatch.setattr(taskqueue, "publish", fake.publish) + monkeypatch.setattr(taskqueue, "publish_retry", fake.publish_retry) + monkeypatch.setattr(taskqueue, "publish_failed", fake.publish_failed) + monkeypatch.setattr(taskqueue, "consume", fake.consume) + monkeypatch.setattr(taskqueue, "health", lambda: True) + return fake + + +@pytest.fixture +def fast_queue(monkeypatch): + """Collapse queue timings so tests do not sleep.""" + import config + + monkeypatch.setattr(config, "RABBITMQ_CONSUME_TIMEOUT", 0.01) + monkeypatch.setattr( + config, "WORKER_LIVENESS_FILE", + os.path.join(os.environ.get("TMPDIR", "/tmp"), "worker-alive-test"), + ) + return config diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 0000000..7b7d8ff --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,324 @@ +""" +The producer HTTP surface: submission, cache hits, job polling, page reads. +""" +import pytest +from fastapi.testclient import TestClient + +import config +import jobstore +import storage +import worker + +API_KEY = "test-key" +HEADERS = {"x-api-key": API_KEY} + + +def test_startup_does_not_provision_storage(monkeypatch): + import api.index as index + + def fail_if_called(): + raise AssertionError("API startup must not provision object storage") + + monkeypatch.setattr(index.db, "migrate", lambda: None) + monkeypatch.setattr(index.storage, "ensure_bucket", fail_if_called) + + with TestClient(index.app) as test_client: + assert test_client.get("/").status_code == 200 + + +@pytest.fixture +def client(database, s3, broker, fast_queue, monkeypatch, tmp_path): + import api.index as index + + monkeypatch.setattr(config, "ADMIN_API_KEY", API_KEY) + + # Stand in for the network: "downloading" a URL yields a local file whose + # bytes are derived from the URL, so distinct URLs hash distinctly. + def fake_download(url: str) -> str: + path = tmp_path / f"dl-{abs(hash(url))}.bin" + path.write_bytes(f"content of {url}".encode()) + return str(path) + + monkeypatch.setattr(index.converter, "download", fake_download) + monkeypatch.setattr(index.chunker, "page_count", lambda path: None) + monkeypatch.setattr( + worker, "convert_range", + lambda file_hash, start, end, use_llm: "# converted markdown", + ) + + with TestClient(index.app) as test_client: + yield test_client + + +def drain(broker, limit=200): + """Deliver every queued chunk, as a worker fleet would.""" + broker.drain(worker.handle, limit=limit) + + +class TestAuth: + def test_missing_key_is_rejected(self, client): + assert client.post("/convert", json={"file": "http://x/a.pdf"}).status_code == 401 + + def test_wrong_key_is_rejected(self, client): + response = client.post( + "/convert", json={"file": "http://x/a.pdf"}, + headers={"x-api-key": "nope"}, + ) + assert response.status_code == 401 + + def test_async_route_also_requires_a_key(self, client): + response = client.post("/async/convert", json={"file": "http://x/a.pdf"}) + assert response.status_code == 401 + + def test_health_needs_no_key(self, client): + assert client.get("/").status_code == 200 + + +class TestAsyncSubmit: + def test_returns_202_with_a_job_id(self, client): + response = client.post( + "/async/convert", json={"file": "http://x/a.pdf"}, headers=HEADERS + ) + assert response.status_code == 202 + body = response.json() + assert body["status"] == "queued" + assert body["job_id"] + assert body["total_chunks"] == 1 + + def test_missing_file_is_422(self, client): + response = client.post("/async/convert", json={}, headers=HEADERS) + assert response.status_code == 422 + + def test_download_failure_is_reported_synchronously(self, client, monkeypatch): + """A bad URL is the caller's problem — better a 400 now than a job that + instantly fails.""" + import api.index as index + + monkeypatch.setattr( + index.converter, "download", + lambda url: (_ for _ in ()).throw(Exception("404 from origin")), + ) + response = client.post( + "/async/convert", + json={"file": "http://x/missing.pdf"}, + headers=HEADERS, + ) + assert response.status_code == 400 + assert "404 from origin" in response.json()["detail"] + + +class TestJobStatus: + def test_reports_progress_then_completion(self, broker, client): + job_id = client.post( + "/async/convert", json={"file": "http://x/a.pdf"}, headers=HEADERS + ).json()["job_id"] + + queued = client.get(f"/convert/jobs/{job_id}", headers=HEADERS).json() + assert queued["status"] == "queued" + assert queued["progress"] == {"completed": 0, "total": 1, "percent": 0} + + drain(broker) + + done = client.get(f"/convert/jobs/{job_id}", headers=HEADERS).json() + assert done["status"] == "done" + assert done["progress"]["percent"] == 100 + + def test_unknown_job_is_404(self, client): + assert client.get("/convert/jobs/nope", headers=HEADERS).status_code == 404 + + def test_failed_job_reports_its_error(self, client): + job_id = client.post( + "/async/convert", json={"file": "http://x/a.pdf"}, headers=HEADERS + ).json()["job_id"] + jobstore.fail_job(job_id, "parser exploded", 0) + + body = client.get(f"/convert/jobs/{job_id}", headers=HEADERS).json() + assert body["status"] == "failed" + assert body["error"] == "parser exploded" + + +class TestCacheHit: + def test_resubmitting_the_same_file_skips_the_queue(self, broker, client): + first = client.post( + "/async/convert", json={"file": "http://x/a.pdf"}, headers=HEADERS + ) + assert first.status_code == 202 + drain(broker) + + second = client.post( + "/async/convert", json={"file": "http://x/a.pdf"}, headers=HEADERS + ) + assert second.status_code == 200 + body = second.json() + assert body["content"] == "# converted markdown" + assert body["pagination"]["total_pages"] == 1 + + def test_cache_hit_creates_no_new_job(self, broker, client, database): + client.post( + "/async/convert", json={"file": "http://x/a.pdf"}, headers=HEADERS + ) + drain(broker) + client.post( + "/async/convert", json={"file": "http://x/a.pdf"}, headers=HEADERS + ) + + rows = database.query("SELECT COUNT(*) AS n FROM jobs") + assert rows[0]["n"] == 1 + + def test_different_files_do_not_share_a_cache_entry(self, broker, client): + client.post( + "/async/convert", json={"file": "http://x/a.pdf"}, headers=HEADERS + ) + drain(broker) + response = client.post( + "/async/convert", json={"file": "http://x/b.pdf"}, headers=HEADERS + ) + assert response.status_code == 202 + + def test_llm_flag_does_not_share_a_cache_entry(self, broker, client): + client.post( + "/async/convert", + json={"file": "http://x/a.pdf", "use_llm": False}, + headers=HEADERS, + ) + drain(broker) + response = client.post( + "/async/convert", + json={"file": "http://x/a.pdf", "use_llm": True}, + headers=HEADERS, + ) + assert response.status_code == 202 + + +class TestSyncConvert: + def test_default_route_waits_and_returns_the_legacy_payload( + self, broker, client, + ): + """Existing callers get the old response without changing their body.""" + import threading + import time + + def convert_after_a_moment(): + time.sleep(0.2) + drain(broker) + + threading.Thread(target=convert_after_a_moment, daemon=True).start() + + response = client.post( + "/convert", json={"file": "http://x/a.pdf"}, + headers=HEADERS, + ) + assert response.status_code == 200 + body = response.json() + assert set(body) == {"id", "content", "pagination"} + assert body["content"] == "# converted markdown" + + def test_surfaces_a_conversion_failure(self, broker, client, monkeypatch): + import threading + import time + + monkeypatch.setattr(config, "MAX_ATTEMPTS", 1) + monkeypatch.setattr( + worker, "convert_range", + lambda *a: (_ for _ in ()).throw(RuntimeError("bad pdf")), + ) + + def fail_it(): + time.sleep(0.2) + drain(broker) + + threading.Thread(target=fail_it, daemon=True).start() + + response = client.post( + "/convert", json={"file": "http://x/a.pdf"}, + headers=HEADERS, + ) + assert response.status_code == 400 + assert "bad pdf" in response.json()["detail"] + + +class TestPages: + def test_reads_a_page(self, broker, client): + doc_key = client.post( + "/async/convert", json={"file": "http://x/a.pdf"}, headers=HEADERS + ).json()["doc_key"] + drain(broker) + + response = client.get(f"/convert/{doc_key}/pages/1", headers=HEADERS) + assert response.status_code == 200 + assert response.json()["content"] == "# converted markdown" + + def test_out_of_range_page_is_404(self, broker, client): + body = client.post( + "/async/convert", json={"file": "http://x/a.pdf"}, headers=HEADERS + ).json() + drain(broker) + assert client.get( + f"/convert/{body['doc_key']}/pages/99", headers=HEADERS + ).status_code == 404 + + def test_unknown_document_is_404(self, client): + assert client.get( + "/convert/nosuchdoc/pages/1", headers=HEADERS + ).status_code == 404 + + +class TestReadiness: + def test_reports_healthy_dependencies(self, client): + body = client.get("/readyz").json() + assert body["status"] == "ok" + assert body["checks"] == {"database": True, "storage": True, "queue": True} + + def test_reports_degraded_when_storage_is_down(self, client, monkeypatch): + import api.index as index + + monkeypatch.setattr(index.storage, "health", lambda: False) + response = client.get("/readyz") + assert response.status_code == 503 + assert response.json()["status"] == "degraded" + + +class TestOpenAPI: + def test_yaml_needs_no_key(self, client): + assert client.get("/openapi.yaml").status_code == 200 + + def test_yaml_matches_the_json_document(self, client): + import yaml + + spec = yaml.safe_load(client.get("/openapi.yaml").text) + assert spec == client.get("/openapi.json").json() + + def test_documents_every_route(self, client): + import yaml + + paths = yaml.safe_load(client.get("/openapi.yaml").text)["paths"] + assert set(paths) == { + "/", "/readyz", "/convert", "/async/convert", + "/convert/jobs/{job_id}", "/convert/{doc_id}/pages/{page}", + } + + def test_documents_sync_and_async_response_contracts(self, client): + import yaml + + paths = yaml.safe_load(client.get("/openapi.yaml").text)["paths"] + sync_responses = paths["/convert"]["post"]["responses"] + assert "202" not in sync_responses + assert sync_responses["200"]["content"]["application/json"]["schema"] == { + "$ref": "#/components/schemas/PageResponse" + } + + async_responses = paths["/async/convert"]["post"]["responses"] + assert async_responses["200"]["content"]["application/json"]["schema"] == { + "$ref": "#/components/schemas/PageResponse" + } + assert async_responses["202"]["content"]["application/json"]["schema"] == { + "$ref": "#/components/schemas/JobAccepted" + } + + def test_declares_the_api_key_security_scheme(self, client): + import yaml + + spec = yaml.safe_load(client.get("/openapi.yaml").text) + assert spec["components"]["securitySchemes"]["APIKeyHeader"] == { + "type": "apiKey", "in": "header", "name": "x-api-key", + } diff --git a/tests/test_chunker.py b/tests/test_chunker.py new file mode 100644 index 0000000..a109c90 --- /dev/null +++ b/tests/test_chunker.py @@ -0,0 +1,108 @@ +""" +Chunk planning and page-range extraction. + +plan_chunks() is pure, so the partitioning rules — which used to be entangled +with writing files to a temp dir — can be asserted directly. +""" +import pytest + +from chunker import WHOLE_FILE, extract_pages, page_count, plan_chunks + + +class TestPlanChunks: + def test_non_pdf_is_a_single_whole_file_chunk(self): + assert plan_chunks(None, 20, 40) == [WHOLE_FILE] + + def test_small_pdf_is_not_split(self): + assert plan_chunks(39, 20, 40) == [WHOLE_FILE] + + def test_at_threshold_is_split(self): + assert plan_chunks(40, 20, 40) == [(0, 20), (20, 40)] + + def test_exact_multiple_has_no_empty_trailing_chunk(self): + assert plan_chunks(60, 20, 40) == [(0, 20), (20, 40), (40, 60)] + + def test_remainder_becomes_a_short_final_chunk(self): + assert plan_chunks(45, 20, 40) == [(0, 20), (20, 40), (40, 45)] + + def test_ranges_are_contiguous_and_cover_every_page(self): + ranges = plan_chunks(137, 20, 40) + assert ranges[0][0] == 0 + assert ranges[-1][1] == 137 + for (_, end), (start, _) in zip(ranges, ranges[1:]): + assert end == start + assert sum(end - start for start, end in ranges) == 137 + + def test_chunk_size_of_one(self): + assert plan_chunks(3, 1, 1) == [(0, 1), (1, 2), (2, 3)] + + def test_zero_pages(self): + # A zero-page PDF is below any threshold, so it stays whole-file and + # the converter deals with whatever it is. + assert plan_chunks(0, 20, 40) == [WHOLE_FILE] + + def test_invalid_chunk_size_rejected(self): + with pytest.raises(ValueError): + plan_chunks(100, 0, 40) + + +class TestPageCount: + def test_counts_pages(self, make_pdf): + assert page_count(make_pdf(7)) == 7 + + def test_non_pdf_returns_none(self, tmp_path): + path = tmp_path / "notes.txt" + path.write_text("this is not a pdf") + assert page_count(str(path)) is None + + def test_missing_file_returns_none(self): + assert page_count("/nonexistent/file.pdf") is None + + +class TestExtractPages: + def test_extracts_the_requested_range(self, make_pdf, tmp_path): + src = make_pdf(25) + out = str(tmp_path / "range.pdf") + extract_pages(src, 5, 15, out) + assert page_count(out) == 10 + + def test_final_short_range(self, make_pdf, tmp_path): + src = make_pdf(25) + out = str(tmp_path / "range.pdf") + extract_pages(src, 20, 25, out) + assert page_count(out) == 5 + + def test_whole_file_sentinel_returns_the_source_untouched(self, make_pdf): + src = make_pdf(3) + assert extract_pages(src, 0, -1, "/tmp/unused.pdf") == src + + def test_extracted_range_holds_exactly_its_own_pages(self, make_pdf, tmp_path): + """The marker on each page proves ranges neither drop nor overlap.""" + from markitdown import MarkItDown + + src = make_pdf(10) + out = str(tmp_path / "range.pdf") + extract_pages(src, 3, 6, out) + + markdown = MarkItDown().convert(out).markdown + # Pages are 0-indexed in the range, 1-indexed in the marker. + assert "PAGEMARKER00004" in markdown + assert "PAGEMARKER00005" in markdown + assert "PAGEMARKER00006" in markdown + assert "PAGEMARKER00003" not in markdown + assert "PAGEMARKER00007" not in markdown + + def test_full_split_reassembles_in_order(self, make_pdf, tmp_path): + """Every page survives a plan/extract round trip, in the right order.""" + from markitdown import MarkItDown + + src = make_pdf(25) + seen = [] + for index, (start, end) in enumerate(plan_chunks(25, 10, 1)): + out = str(tmp_path / f"chunk_{index}.pdf") + extract_pages(src, start, end, out) + seen.append(MarkItDown().convert(out).markdown) + + combined = "\n\n".join(seen) + positions = [combined.index(f"PAGEMARKER{n:05d}") for n in range(1, 26)] + assert positions == sorted(positions) diff --git a/tests/test_converter.py b/tests/test_converter.py new file mode 100644 index 0000000..b033aca --- /dev/null +++ b/tests/test_converter.py @@ -0,0 +1,35 @@ +import os + +import converter + + +def test_download_has_no_fixed_request_timeout(monkeypatch): + request_options = None + + class Response: + status_code = 200 + + def __enter__(self): + return self + + def __exit__(self, *_args): + return None + + def iter_content(self, chunk_size): + assert chunk_size > 0 + yield b"document" + + def get(_url, **options): + nonlocal request_options + request_options = options + return Response() + + monkeypatch.setattr(converter.requests, "get", get) + + path = converter.download("https://example.com/document.pdf") + try: + assert request_options == {"stream": True} + with open(path, "rb") as downloaded: + assert downloaded.read() == b"document" + finally: + os.remove(path) diff --git a/tests/test_db.py b/tests/test_db.py new file mode 100644 index 0000000..21bbe63 --- /dev/null +++ b/tests/test_db.py @@ -0,0 +1,175 @@ +""" +Driver-level behaviour that differs between sqlite3 and libsql. + +Tests run against sqlite3, so anything libsql does differently has to be +asserted explicitly here or it only shows up in production. +""" +import db + + +class FakeCursor: + """Stands in for a driver cursor with a chosen fetchall() return value.""" + + def __init__(self, description, rows): + self.description = description + self._rows = rows + + def fetchall(self): + return self._rows + + +class TestRowsToDicts: + def test_none_from_fetchall_is_an_empty_result(self): + """ + libsql returns None from fetchall() when a statement matched no rows, + where sqlite3 returns []. The queue hits this on every idle poll: the + claim is an UPDATE ... RETURNING that matches nothing when there is no + work. + + Getting this wrong is not merely a crash. The UPDATE commits before the + result is read, so the raise happened *after* `attempts` had already + been incremented — every poll burned a retry against a chunk no worker + ever received, and the whole job became unclaimable within seconds. + """ + cursor = FakeCursor(description=[("id",), ("job_id",)], rows=None) + assert db._rows_to_dicts(cursor) == [] + + def test_empty_list_from_fetchall_is_an_empty_result(self): + cursor = FakeCursor(description=[("id",)], rows=[]) + assert db._rows_to_dicts(cursor) == [] + + def test_no_description_is_an_empty_result(self): + """A plain UPDATE with no RETURNING clause has no description.""" + cursor = FakeCursor(description=None, rows=None) + assert db._rows_to_dicts(cursor) == [] + + def test_rows_are_zipped_to_column_names(self): + cursor = FakeCursor( + description=[("id",), ("name",)], rows=[(1, "a"), (2, "b")] + ) + assert db._rows_to_dicts(cursor) == [ + {"id": 1, "name": "a"}, + {"id": 2, "name": "b"}, + ] + + +class TestStaleConnection: + """ + libsql holds a server-side stream handle that expires when idle, so a + pooled connection dies on its own and every later statement fails with + "Stream handle N is expired". sqlite3 has no equivalent, so this is only + reachable in production unless asserted here. + """ + + def test_expired_stream_handle_is_treated_as_stale(self): + exc = Exception("hrana server: Stream handle for 926860691 is expired") + assert db._is_stale_connection(exc) is True + + def test_dropped_socket_is_treated_as_stale(self): + assert db._is_stale_connection(Exception("Connection reset by peer")) is True + assert db._is_stale_connection(Exception("broken pipe")) is True + + def test_ordinary_sql_errors_are_not_stale(self): + """A retry must not paper over a real error — especially not a write, + which could then apply twice.""" + assert db._is_stale_connection(Exception("no such table: jobs")) is False + assert db._is_stale_connection(Exception("UNIQUE constraint failed")) is False + + def test_stale_connection_is_reopened_and_the_statement_retried( + self, database, monkeypatch + ): + calls = {"n": 0} + real_get = db.get_connection + + class Flaky: + """Fails the first execute with a stale-handle error, as an + expired connection would.""" + + def __init__(self, conn): + self._conn = conn + + def execute(self, sql, params=()): + calls["n"] += 1 + if calls["n"] == 1: + raise Exception("Stream handle for 42 is expired") + return self._conn.execute(sql, params) + + def __getattr__(self, name): + return getattr(self._conn, name) + + def flaky_get(): + conn = real_get() + return Flaky(conn) if calls["n"] == 0 else conn + + monkeypatch.setattr(db, "get_connection", flaky_get) + assert db.query("SELECT 1 AS ok") == [{"ok": 1}] + assert calls["n"] >= 1 + + def test_a_real_error_still_propagates(self, database): + import pytest + + with pytest.raises(Exception): + db.query("SELECT * FROM table_that_does_not_exist") + + +class TestDriverSelection: + def test_local_path_uses_sqlite(self, monkeypatch): + import config + + monkeypatch.setattr(config, "DATABASE_URL", "/tmp/x.db") + monkeypatch.setattr(config, "DATABASE_AUTH_TOKEN", None) + assert db.uses_libsql() is False + + def test_file_scheme_uses_sqlite(self, monkeypatch): + import config + + monkeypatch.setattr(config, "DATABASE_URL", "file:markitdown.db") + monkeypatch.setattr(config, "DATABASE_AUTH_TOKEN", None) + assert db.uses_libsql() is False + + def test_libsql_scheme_uses_libsql(self, monkeypatch): + import config + + monkeypatch.setattr(config, "DATABASE_URL", "libsql://db.turso.io") + monkeypatch.setattr(config, "DATABASE_AUTH_TOKEN", None) + assert db.uses_libsql() is True + + def test_self_hosted_http_uses_libsql(self, monkeypatch): + """The in-cluster sqld is plain HTTP, not the libsql:// scheme.""" + import config + + monkeypatch.setattr(config, "DATABASE_URL", "http://markitdown-libsql:8080") + monkeypatch.setattr(config, "DATABASE_AUTH_TOKEN", None) + assert db.uses_libsql() is True + + def test_auth_token_forces_libsql(self, monkeypatch): + import config + + monkeypatch.setattr(config, "DATABASE_URL", "somewhere.db") + monkeypatch.setattr(config, "DATABASE_AUTH_TOKEN", "token") + assert db.uses_libsql() is True + + +class TestMigrations: + def test_split_statements_ignores_semicolons_in_comments(self): + """ + A ';' inside a line comment used to cut a statement in half, which is + how "-- exclusive; -1 means the whole file" broke every migration. + """ + sql = """ + CREATE TABLE t ( + a INTEGER, -- exclusive; -1 means the whole file + b TEXT + ); + CREATE INDEX i ON t(a); + """ + statements = db._split_statements(sql) + assert len(statements) == 2 + assert statements[0].startswith("CREATE TABLE") + assert statements[1].startswith("CREATE INDEX") + + def test_migrations_are_idempotent(self, database): + """Every pod runs migrations on boot, so re-running must be a no-op.""" + db._migrated = False + db.migrate() + assert db.query_one("SELECT COUNT(*) AS n FROM jobs")["n"] == 0 diff --git a/tests/test_integration.py b/tests/test_integration.py new file mode 100644 index 0000000..e06a9f8 --- /dev/null +++ b/tests/test_integration.py @@ -0,0 +1,113 @@ +""" +Full-stack pass: real PDF, real HTTP download, real markitdown conversion. + +Everything else stubs conversion to stay fast. This one does not — it is the +test that proves the pieces actually fit together, so it is deliberately the +slow one. +""" +import functools +import http.server +import threading + +import pytest + +import config +import jobstore +import storage +import worker + +pytestmark = pytest.mark.slow + + +@pytest.fixture +def pdf_server(tmp_path, make_pdf): + """Serve a generated PDF over real HTTP on localhost.""" + make_pdf(25, name="book.pdf") + + handler = functools.partial( + http.server.SimpleHTTPRequestHandler, directory=str(tmp_path) + ) + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + yield f"http://127.0.0.1:{server.server_port}/book.pdf" + finally: + server.shutdown() + + +@pytest.fixture +def client(database, s3, monkeypatch, tmp_path): + from fastapi.testclient import TestClient + + import api.index as index + + monkeypatch.setattr(config, "ADMIN_API_KEY", "k") + monkeypatch.setattr(config, "SOURCE_CACHE_DIR", str(tmp_path / "srccache")) + # Force the 25-page document to actually split, so assembly across chunk + # boundaries is exercised rather than short-circuited. + monkeypatch.setattr(config, "PAGES_PER_CHUNK", 10) + monkeypatch.setattr(config, "MIN_PAGES_FOR_PARALLEL", 5) + + with TestClient(index.app) as test_client: + yield test_client + + +def test_converts_a_real_pdf_end_to_end(client, pdf_server, database, s3, broker): + headers = {"x-api-key": "k"} + + submitted = client.post( + "/async/convert", json={"file": pdf_server}, headers=headers + ) + assert submitted.status_code == 202 + body = submitted.json() + assert body["total_chunks"] == 3 # 25 pages at 10/chunk + job_id, doc_key = body["job_id"], body["doc_key"] + + # Two workers race for the three chunks, as they would across pods. + threads = [ + threading.Thread(target=_drain, args=(broker,)) for _ in range(2) + ] + for t in threads: + t.start() + for t in threads: + t.join(timeout=300) + + status = client.get(f"/convert/jobs/{job_id}", headers=headers).json() + assert status["status"] == "done", status + assert status["progress"] == {"completed": 3, "total": 3, "percent": 100} + + # Walk every page through the HTTP API and reassemble it. + manifest = storage.get_manifest(doc_key) + text = "" + for page in range(1, manifest["total_pages"] + 1): + response = client.get(f"/convert/{doc_key}/pages/{page}", headers=headers) + assert response.status_code == 200 + text += response.json()["content"] + + # Every page of the source survived... + missing = [n for n in range(1, 26) if f"PAGEMARKER{n:05d}" not in text] + assert not missing, f"lost pages: {missing}" + + # ...and they are still in order across chunk boundaries. + positions = [text.index(f"PAGEMARKER{n:05d}") for n in range(1, 26)] + assert positions == sorted(positions) + + +def test_resubmitting_hits_the_cache(broker, client, pdf_server, database, s3): + headers = {"x-api-key": "k"} + + client.post("/async/convert", json={"file": pdf_server}, headers=headers) + _drain(broker) + + # The original synchronous endpoint shares the same content-addressed + # result and still returns the legacy first-page payload. + again = client.post("/convert", json={"file": pdf_server}, headers=headers) + assert again.status_code == 200 # served from cache, no new job + assert "PAGEMARKER00001" in again.json()["content"] + + rows = database.query("SELECT COUNT(*) AS n FROM jobs") + assert rows[0]["n"] == 1 + + +def _drain(broker) -> None: + broker.drain(worker.handle) diff --git a/tests/test_jobstore.py b/tests/test_jobstore.py new file mode 100644 index 0000000..2bc4867 --- /dev/null +++ b/tests/test_jobstore.py @@ -0,0 +1,168 @@ +""" +Job and document state. + +The completion counter is the piece that matters most: RabbitMQ distributes the +work, but it cannot answer "did I just finish the last of 25 chunks?" — and that +answer decides who assembles the document, exactly once. +""" +import threading + +import jobstore +from jobstore import CHUNK_ALREADY_DONE, JOB_ALREADY_TERMINAL + + +def make_job(chunks: int = 3, use_llm: bool = False) -> str: + job_id = jobstore.new_job_id() + ranges = [(i * 10, (i + 1) * 10) for i in range(chunks)] + jobstore.create_job(job_id, f"key-{job_id}", "hash", "http://x/f.pdf", + ranges, use_llm) + return job_id + + +class TestCreateJob: + def test_creates_job_and_chunk_rows(self, database): + job_id = make_job(4) + job = jobstore.get_job(job_id) + assert job["status"] == "queued" + assert job["total_chunks"] == 4 + assert job["remaining_chunks"] == 4 + assert jobstore.chunk_indices(job_id) == [0, 1, 2, 3] + + def test_progress_starts_at_zero(self, database): + job_id = make_job(4) + assert jobstore.job_progress(job_id) == { + "completed": 0, "total": 4, "percent": 0 + } + + +class TestCompleteChunk: + def test_counts_down_to_zero_for_the_assembler(self, database): + job_id = make_job(3) + results = [jobstore.complete_chunk(job_id, i, "k") for i in range(3)] + assert results == [2, 1, 0] + + def test_duplicate_completion_does_not_double_decrement(self, database): + """ + RabbitMQ delivers at least once, so a chunk whose worker died before + acking gets redelivered and converted twice. The second recording must + be a no-op, or the counter goes negative and nobody ever assembles. + """ + job_id = make_job(3) + assert jobstore.complete_chunk(job_id, 0, "k") == 2 + assert jobstore.complete_chunk(job_id, 0, "k") == CHUNK_ALREADY_DONE + assert jobstore.get_job(job_id)["remaining_chunks"] == 2 + + def test_completion_on_a_failed_job_is_rejected(self, database): + job_id = make_job(2) + jobstore.fail_job(job_id, "poisoned") + assert jobstore.complete_chunk(job_id, 0, "k") == JOB_ALREADY_TERMINAL + + def test_progress_tracks_completions(self, database): + job_id = make_job(4) + jobstore.complete_chunk(job_id, 0, "k") + assert jobstore.job_progress(job_id) == { + "completed": 1, "total": 4, "percent": 25 + } + + def test_concurrent_completion_yields_exactly_one_assembler(self, database): + """Exactly one worker must see remaining == 0, or a document is + assembled twice or never.""" + job_id = make_job(16) + results: list[int] = [] + lock = threading.Lock() + barrier = threading.Barrier(16) + + def complete(index): + barrier.wait() + value = jobstore.complete_chunk(job_id, index, "k") + with lock: + results.append(value) + + threads = [threading.Thread(target=complete, args=(i,)) for i in range(16)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=60) + + assert results.count(0) == 1 + assert sorted(results) == list(range(16)) + + +class TestFailure: + def test_fail_job_is_terminal(self, database): + job_id = make_job(2) + jobstore.fail_job(job_id, "first", 0) + jobstore.fail_job(job_id, "second", 1) + assert jobstore.get_job(job_id)["error"] == "first" + + def test_finish_job_blocks_later_failure(self, database): + job_id = make_job(1) + jobstore.finish_job(job_id) + jobstore.fail_job(job_id, "too late") + assert jobstore.get_job(job_id)["status"] == "done" + + def test_fail_chunk_records_the_error(self, database): + job_id = make_job(2) + jobstore.fail_chunk(job_id, 1, "corrupt page range") + row = database.query_one( + "SELECT status, last_error FROM job_chunks WHERE job_id = ? AND chunk_index = 1", + (job_id,), + ) + assert row["status"] == "failed" + assert row["last_error"] == "corrupt page range" + + def test_reschedule_leaves_a_done_chunk_alone(self, database): + """A late retry must not undo a chunk that already succeeded.""" + job_id = make_job(2) + jobstore.complete_chunk(job_id, 0, "k") + jobstore.reschedule_chunk(job_id, 0, "stale failure") + row = database.query_one( + "SELECT status FROM job_chunks WHERE job_id = ? AND chunk_index = 0", + (job_id,), + ) + assert row["status"] == "done" + + +class TestChunkState: + def test_start_chunk_marks_running_and_flips_the_job(self, database): + job_id = make_job(2) + jobstore.start_chunk(job_id, 0, "worker-1", attempt=1) + assert jobstore.get_job(job_id)["status"] == "running" + row = database.query_one( + "SELECT status, lease_owner, attempts FROM job_chunks " + "WHERE job_id = ? AND chunk_index = 0", + (job_id,), + ) + assert row["status"] == "running" + assert row["lease_owner"] == "worker-1" + assert row["attempts"] == 1 + + def test_start_chunk_does_not_reopen_a_done_chunk(self, database): + """Redelivery of an already-converted chunk must not un-finish it.""" + job_id = make_job(2) + jobstore.complete_chunk(job_id, 0, "k") + jobstore.start_chunk(job_id, 0, "worker-2", attempt=2) + row = database.query_one( + "SELECT status FROM job_chunks WHERE job_id = ? AND chunk_index = 0", + (job_id,), + ) + assert row["status"] == "done" + + +class TestDocuments: + def test_round_trips_a_document(self, database): + jobstore.record_document("k1", "hash", 5, 2, 9000, 5000) + doc = jobstore.find_cached_document("k1") + assert doc["total_pages"] == 5 + assert doc["total_chunks"] == 2 + assert doc["expires_at"] > doc["created_at"] + + def test_missing_document_is_none(self, database): + assert jobstore.find_cached_document("nope") is None + + def test_recording_twice_is_harmless(self, database): + """Two jobs for identical content can race; both produce the same + output, so last writer wins.""" + jobstore.record_document("k1", "hash", 5, 2, 9000, 5000) + jobstore.record_document("k1", "hash", 5, 2, 9000, 5000) + assert jobstore.find_cached_document("k1")["total_pages"] == 5 diff --git a/tests/test_pdf_chunk.py b/tests/test_pdf_chunk.py deleted file mode 100644 index b39dc14..0000000 --- a/tests/test_pdf_chunk.py +++ /dev/null @@ -1,158 +0,0 @@ -""" -Unit tests for page-chunked PDF conversion. - -These run against real generated PDFs rather than mocks — the whole point of -the module is how pypdf/pdfplumber/pdfminer behave on actual documents, and a -mocked parser would test nothing. -""" -import os - -import pytest - -import pdf_chunk - - -class TestPageCount: - def test_counts_pages(self, make_pdf): - assert pdf_chunk.page_count(make_pdf(7)) == 7 - - def test_returns_none_for_non_pdf(self, tmp_path): - path = tmp_path / "not.pdf" - path.write_text("this is plainly not a pdf") - assert pdf_chunk.page_count(str(path)) is None - - def test_returns_none_for_missing_file(self, tmp_path): - assert pdf_chunk.page_count(str(tmp_path / "nope.pdf")) is None - - -class TestShouldChunk: - def test_true_at_or_above_threshold(self, make_pdf): - assert pdf_chunk.should_chunk(make_pdf(10), min_pages=10) is True - - def test_false_below_threshold(self, make_pdf): - assert pdf_chunk.should_chunk(make_pdf(9), min_pages=10) is False - - def test_false_for_non_pdf(self, tmp_path): - path = tmp_path / "not.pdf" - path.write_text("nope") - assert pdf_chunk.should_chunk(str(path), min_pages=1) is False - - -class TestSplitPdf: - def test_splits_evenly(self, make_pdf, tmp_path): - out = tmp_path / "chunks" - out.mkdir() - chunks = pdf_chunk.split_pdf(make_pdf(20), 5, str(out)) - assert len(chunks) == 4 - assert all(pdf_chunk.page_count(c) == 5 for c in chunks) - - def test_last_chunk_holds_remainder(self, make_pdf, tmp_path): - out = tmp_path / "chunks" - out.mkdir() - chunks = pdf_chunk.split_pdf(make_pdf(22), 5, str(out)) - assert len(chunks) == 5 - assert [pdf_chunk.page_count(c) for c in chunks] == [5, 5, 5, 5, 2] - - def test_single_chunk_when_smaller_than_chunk_size(self, make_pdf, tmp_path): - out = tmp_path / "chunks" - out.mkdir() - chunks = pdf_chunk.split_pdf(make_pdf(3), 20, str(out)) - assert len(chunks) == 1 - assert pdf_chunk.page_count(chunks[0]) == 3 - - def test_preserves_total_page_count(self, make_pdf, tmp_path): - out = tmp_path / "chunks" - out.mkdir() - chunks = pdf_chunk.split_pdf(make_pdf(37), 8, str(out)) - assert sum(pdf_chunk.page_count(c) for c in chunks) == 37 - - def test_chunks_are_ordered_by_name(self, make_pdf, tmp_path): - out = tmp_path / "chunks" - out.mkdir() - chunks = pdf_chunk.split_pdf(make_pdf(30), 10, str(out)) - assert chunks == sorted(chunks) - - def test_rejects_zero_pages_per_chunk(self, make_pdf, tmp_path): - out = tmp_path / "chunks" - out.mkdir() - with pytest.raises(ValueError): - pdf_chunk.split_pdf(make_pdf(5), 0, str(out)) - - -class TestConvertChunk: - def test_extracts_text_from_chunk(self, small_pdf): - markdown = pdf_chunk.convert_chunk(small_pdf) - assert "PAGEMARKER00001" in markdown - assert "PAGEMARKER00003" in markdown - - -class TestConvertPdfChunked: - def test_single_chunk_document(self, make_pdf): - markdown = pdf_chunk.convert_pdf_chunked( - make_pdf(5), pages_per_chunk=20, max_workers=2 - ) - assert "PAGEMARKER00001" in markdown - assert "PAGEMARKER00005" in markdown - - def test_retains_every_page_across_chunks(self, multi_chunk_pdf): - markdown = pdf_chunk.convert_pdf_chunked( - multi_chunk_pdf, pages_per_chunk=10, max_workers=2 - ) - for page in range(1, 26): - assert f"PAGEMARKER{page:05d}" in markdown, f"page {page} missing" - - def test_preserves_page_order(self, multi_chunk_pdf): - markdown = pdf_chunk.convert_pdf_chunked( - multi_chunk_pdf, pages_per_chunk=10, max_workers=2 - ) - positions = [ - markdown.index(f"PAGEMARKER{page:05d}") for page in range(1, 26) - ] - assert positions == sorted(positions), "pages came back out of order" - - def test_chunk_size_does_not_change_content(self, make_pdf): - """Page count per chunk is a tuning knob, not a semantic one.""" - path = make_pdf(24) - a = pdf_chunk.convert_pdf_chunked(path, pages_per_chunk=6, max_workers=2) - b = pdf_chunk.convert_pdf_chunked(path, pages_per_chunk=12, max_workers=2) - for page in range(1, 25): - marker = f"PAGEMARKER{page:05d}" - assert marker in a and marker in b - - def test_matches_sequential_conversion(self, make_pdf): - """ - Chunked output must carry the same page content as the whole-document - path. Exact string equality is not asserted: markitdown picks its - extraction strategy per document, so deciding it per chunk can shift - whitespace at chunk boundaries. - """ - from markitdown import MarkItDown - - path = make_pdf(24) - sequential = MarkItDown().convert(path).text_content - chunked = pdf_chunk.convert_pdf_chunked( - path, pages_per_chunk=8, max_workers=2 - ) - - for page in range(1, 25): - marker = f"PAGEMARKER{page:05d}" - assert marker in sequential and marker in chunked - - # Allow small boundary differences but catch wholesale content loss. - ratio = len(chunked) / len(sequential) - assert 0.95 < ratio < 1.05, f"length drifted: {ratio:.3f}" - - def test_cleans_up_temp_chunks(self, multi_chunk_pdf, tmp_path): - before = set(os.listdir("/tmp")) if os.path.isdir("/tmp") else set() - pdf_chunk.convert_pdf_chunked( - multi_chunk_pdf, pages_per_chunk=10, max_workers=2 - ) - after = set(os.listdir("/tmp")) if os.path.isdir("/tmp") else set() - leaked = {d for d in after - before if d.startswith("pdf_chunks_")} - assert not leaked, f"left temp dirs behind: {leaked}" - - def test_propagates_worker_failure(self, tmp_path): - bad = tmp_path / "bad.pdf" - bad.write_text("not a pdf at all") - with pytest.raises(Exception): - pdf_chunk.convert_pdf_chunked(str(bad), pages_per_chunk=5) diff --git a/tests/test_storage.py b/tests/test_storage.py new file mode 100644 index 0000000..2717560 --- /dev/null +++ b/tests/test_storage.py @@ -0,0 +1,108 @@ +""" +Content-addressed storage: key derivation and object round trips. + +The key derivation tests are the interesting ones — the cache is only safe if +the key changes whenever anything that affects the output changes. +""" +import config +import storage + + +class TestKeyDerivation: + def test_same_bytes_hash_the_same(self, tmp_path): + a, b = tmp_path / "a.bin", tmp_path / "b.bin" + a.write_bytes(b"identical content") + b.write_bytes(b"identical content") + assert storage.hash_file(str(a)) == storage.hash_file(str(b)) + + def test_different_bytes_hash_differently(self, tmp_path): + a, b = tmp_path / "a.bin", tmp_path / "b.bin" + a.write_bytes(b"one") + b.write_bytes(b"two") + assert storage.hash_file(str(a)) != storage.hash_file(str(b)) + + def test_llm_flag_changes_the_doc_key(self): + """LLM and non-LLM output differ, so they must not share a cache + entry.""" + assert storage.make_doc_key("h", True) != storage.make_doc_key("h", False) + + def test_page_size_changes_the_doc_key(self, monkeypatch): + before = storage.make_doc_key("h", False) + monkeypatch.setattr(config, "PAGE_SIZE", config.PAGE_SIZE + 1) + assert storage.make_doc_key("h", False) != before + + def test_chunk_size_changes_the_doc_key(self, monkeypatch): + """Chunk boundaries affect the assembled output, so they belong in the + fingerprint.""" + before = storage.make_doc_key("h", False) + monkeypatch.setattr(config, "PAGES_PER_CHUNK", config.PAGES_PER_CHUNK + 1) + assert storage.make_doc_key("h", False) != before + + def test_doc_key_is_stable_across_calls(self): + assert storage.make_doc_key("h", False) == storage.make_doc_key("h", False) + + def test_part_keys_sort_in_page_order(self): + """Zero-padded so lexicographic listing matches page order.""" + keys = [storage.part_key("d", i) for i in (2, 10, 1)] + assert sorted(keys) == [ + storage.part_key("d", 1), + storage.part_key("d", 2), + storage.part_key("d", 10), + ] + + +class TestObjects: + def test_source_round_trip(self, s3, tmp_path): + src = tmp_path / "in.pdf" + src.write_bytes(b"pdf bytes here") + storage.put_source("hash1", str(src)) + + dest = tmp_path / "out.pdf" + storage.get_source("hash1", str(dest)) + assert dest.read_bytes() == b"pdf bytes here" + + def test_put_source_skips_an_existing_object(self, s3, tmp_path): + src = tmp_path / "in.pdf" + src.write_bytes(b"original") + storage.put_source("hash1", str(src)) + + # Same key, different bytes: content addressing means this cannot + # legitimately happen, so the upload should be skipped entirely. + src.write_bytes(b"different") + storage.put_source("hash1", str(src)) + + dest = tmp_path / "out.pdf" + storage.get_source("hash1", str(dest)) + assert dest.read_bytes() == b"original" + + def test_part_round_trip(self, s3): + storage.put_part("doc1", 3, "# chunk three") + assert storage.get_part("doc1", 3) == "# chunk three" + + def test_missing_part_is_none(self, s3): + assert storage.get_part("doc1", 99) is None + + def test_document_round_trip(self, s3): + pages = ["page one", "page two", "page three"] + manifest = storage.put_document("doc1", pages, 42) + + assert manifest["total_pages"] == 3 + assert manifest["total_length"] == 42 + assert storage.get_manifest("doc1") == manifest + assert storage.get_page("doc1", 1) == "page one" + assert storage.get_page("doc1", 3) == "page three" + + def test_pages_are_one_indexed(self, s3): + storage.put_document("doc1", ["only page"], 9) + assert storage.get_page("doc1", 0) is None + assert storage.get_page("doc1", 1) == "only page" + + def test_missing_manifest_is_none(self, s3): + assert storage.get_manifest("nope") is None + + def test_unicode_survives_the_round_trip(self, s3): + storage.put_document("doc1", ["héllo — wörld 🎉"], 15) + assert storage.get_page("doc1", 1) == "héllo — wörld 🎉" + + def test_health_reports_true_for_a_live_bucket(self, s3): + assert storage.health() is True diff --git a/tests/test_taskqueue.py b/tests/test_taskqueue.py new file mode 100644 index 0000000..0d057a1 --- /dev/null +++ b/tests/test_taskqueue.py @@ -0,0 +1,145 @@ +""" +RabbitMQ message contract and topology. + +The broker itself is not under test; what is, is the message round trip and the +declaration arguments, since a wrong dead-letter setting silently sends retries +to the wrong place and only shows up under load. +""" +import config +import taskqueue +from taskqueue import ChunkTask + + +def make_task(**overrides) -> ChunkTask: + base = dict( + job_id="job-1", chunk_index=3, doc_key="doc-1", file_hash="hash-1", + start_page=60, end_page=80, total_chunks=5, use_llm=False, + ) + return ChunkTask(**{**base, **overrides}) + + +class TestSerialisation: + def test_round_trips_every_field(self): + task = make_task() + assert ChunkTask.from_bytes(task.to_bytes()) == task + + def test_whole_file_sentinel_survives(self): + """end_page == -1 means "not a page range"; losing it would make the + worker try to page-slice a non-PDF.""" + task = make_task(end_page=-1) + assert ChunkTask.from_bytes(task.to_bytes()).end_page == -1 + + def test_use_llm_stays_boolean(self): + task = make_task(use_llm=True) + assert ChunkTask.from_bytes(task.to_bytes()).use_llm is True + + def test_attempt_defaults_to_one(self): + assert make_task().attempt == 1 + + def test_next_attempt_increments_and_preserves_the_rest(self): + task = make_task(attempt=2) + nxt = task.next_attempt() + assert nxt.attempt == 3 + assert nxt.job_id == task.job_id + assert nxt.chunk_index == task.chunk_index + assert nxt.start_page == task.start_page + + def test_tasks_are_immutable(self): + """Frozen so a handler cannot mutate a task and republish something + subtly different from what it received.""" + import dataclasses + + import pytest + + with pytest.raises(dataclasses.FrozenInstanceError): + make_task().chunk_index = 99 + + +class TestRetryTiers: + def test_each_attempt_maps_to_its_own_delay_queue(self, monkeypatch): + monkeypatch.setattr(config, "RETRY_DELAYS", [5, 30, 300]) + assert taskqueue.retry_queue_name(1).endswith(".retry.1") + assert taskqueue.retry_queue_name(2).endswith(".retry.2") + assert taskqueue.retry_queue_name(3).endswith(".retry.3") + + def test_attempts_beyond_the_last_tier_clamp(self, monkeypatch): + """MAX_ATTEMPTS could be raised past the configured tiers; that must + reuse the longest delay, not address a queue that was never declared.""" + monkeypatch.setattr(config, "RETRY_DELAYS", [5, 30]) + assert taskqueue.retry_queue_name(7).endswith(".retry.2") + + +class RecordingChannel: + """Captures declarations instead of talking to a broker.""" + + def __init__(self): + self.exchanges = {} + self.queues = {} + self.bindings = [] + + def exchange_declare(self, exchange, kind, durable=False, **kw): + self.exchanges[exchange] = {"type": kind, "durable": durable} + + def queue_declare(self, queue, durable=False, arguments=None, **kw): + self.queues[queue] = {"durable": durable, "arguments": arguments or {}} + + def queue_bind(self, queue, exchange, routing_key=None, **kw): + self.bindings.append((queue, exchange, routing_key)) + + +class TestTopology: + def declared(self, monkeypatch): + monkeypatch.setattr(config, "RETRY_DELAYS", [5, 30, 300]) + ch = RecordingChannel() + taskqueue.declare(ch) + return ch + + def test_work_queue_is_durable(self, monkeypatch): + """Messages outlive a broker restart, so a queued job is not lost.""" + ch = self.declared(monkeypatch) + assert ch.queues[config.RABBITMQ_QUEUE]["durable"] is True + + def test_work_queue_dead_letters_to_the_dlx(self, monkeypatch): + ch = self.declared(monkeypatch) + args = ch.queues[config.RABBITMQ_QUEUE]["arguments"] + assert args["x-dead-letter-exchange"] == taskqueue.DLX_EXCHANGE + + def test_each_retry_tier_has_its_own_ttl(self, monkeypatch): + """Per-tier queues, not per-message TTL: RabbitMQ only expires the + message at the head, so a 300s message would block a 5s one behind it.""" + ch = self.declared(monkeypatch) + ttls = [ + ch.queues[f"{config.RABBITMQ_QUEUE}.retry.{tier}"]["arguments"]["x-message-ttl"] + for tier in (1, 2, 3) + ] + assert ttls == [5000, 30000, 300000] + + def test_retry_queues_dead_letter_back_into_the_work_queue(self, monkeypatch): + """This is what makes a delayed retry actually get redelivered.""" + ch = self.declared(monkeypatch) + for tier in (1, 2, 3): + args = ch.queues[f"{config.RABBITMQ_QUEUE}.retry.{tier}"]["arguments"] + assert args["x-dead-letter-exchange"] == config.RABBITMQ_EXCHANGE + assert args["x-dead-letter-routing-key"] == config.RABBITMQ_QUEUE + + def test_failed_queue_exists_and_is_bound_to_the_dlx(self, monkeypatch): + """Poison messages are retained for inspection rather than dropped.""" + ch = self.declared(monkeypatch) + assert taskqueue.FAILED_QUEUE in ch.queues + assert (taskqueue.FAILED_QUEUE, taskqueue.DLX_EXCHANGE, + config.RABBITMQ_QUEUE) in ch.bindings + + def test_retry_queues_do_not_dead_letter_to_the_failed_queue(self, monkeypatch): + """A retry must go back to work, not straight to the graveyard — the + easiest way to get this wrong is to reuse the DLX everywhere.""" + ch = self.declared(monkeypatch) + for tier in (1, 2, 3): + args = ch.queues[f"{config.RABBITMQ_QUEUE}.retry.{tier}"]["arguments"] + assert args["x-dead-letter-exchange"] != taskqueue.DLX_EXCHANGE + + def test_declaration_is_idempotent(self, monkeypatch): + """Every pod declares on boot.""" + ch = self.declared(monkeypatch) + before = (len(ch.queues), len(ch.exchanges)) + taskqueue.declare(ch) + assert (len(ch.queues), len(ch.exchanges)) == before diff --git a/tests/test_worker.py b/tests/test_worker.py new file mode 100644 index 0000000..72dea9b --- /dev/null +++ b/tests/test_worker.py @@ -0,0 +1,290 @@ +""" +The consumer: receive a chunk -> convert -> record -> assemble. + +Conversion is stubbed (markitdown is slow and covered by test_chunker); what is +under test is the orchestration around it — retries, at-least-once duplicate +delivery, and who assembles the finished document. +""" +import pytest + +import config +import jobstore +import storage +import taskqueue +import worker +from taskqueue import ChunkTask + + +@pytest.fixture +def stub_convert(monkeypatch): + """Replace conversion with a deterministic, instant stand-in.""" + calls = [] + + def _convert(file_hash, start, end, use_llm): + calls.append((file_hash, start, end, use_llm)) + return f"CHUNK[{start}:{end}]" + + monkeypatch.setattr(worker, "convert_range", _convert) + return calls + + +def make_job(broker, chunks=3, use_llm=False): + """Create a job and publish its chunks, as the producer does.""" + job_id = jobstore.new_job_id() + doc_key = f"doc-{job_id}" + ranges = [(i * 10, (i + 1) * 10) for i in range(chunks)] + jobstore.create_job(job_id, doc_key, "filehash", "http://x/f.pdf", + ranges, use_llm) + taskqueue.publish([ + ChunkTask(job_id=job_id, chunk_index=i, doc_key=doc_key, + file_hash="filehash", start_page=s, end_page=e, + total_chunks=chunks, use_llm=use_llm) + for i, (s, e) in enumerate(ranges) + ]) + return job_id + + +class TestHappyPath: + def test_job_completes_and_publishes_a_readable_document( + self, database, s3, broker, fast_queue, stub_convert + ): + job_id = make_job(broker, 3) + broker.drain(worker.handle) + + job = jobstore.get_job(job_id) + assert job["status"] == "done" + assert job["remaining_chunks"] == 0 + assert storage.get_page(job["doc_key"], 1) is not None + + def test_chunks_are_assembled_in_page_order( + self, database, s3, broker, fast_queue, stub_convert + ): + job_id = make_job(broker, 5) + broker.drain(worker.handle) + + doc_key = jobstore.get_job(job_id)["doc_key"] + manifest = storage.get_manifest(doc_key) + text = "".join( + storage.get_page(doc_key, n) + for n in range(1, manifest["total_pages"] + 1) + ) + positions = [text.index(f"CHUNK[{i * 10}:{(i + 1) * 10}]") for i in range(5)] + assert positions == sorted(positions) + + def test_out_of_order_delivery_still_assembles_in_page_order( + self, database, s3, broker, fast_queue, stub_convert + ): + """ + RabbitMQ gives no cross-consumer ordering guarantee, and a retry pushes + its chunk arbitrarily late. Assembly must therefore sort by chunk + index, not by completion order. + """ + job_id = make_job(broker, 5) + broker.ready.reverse() # deliver chunk 4 first, chunk 0 last + broker.drain(worker.handle) + + doc_key = jobstore.get_job(job_id)["doc_key"] + manifest = storage.get_manifest(doc_key) + text = "".join( + storage.get_page(doc_key, n) + for n in range(1, manifest["total_pages"] + 1) + ) + positions = [text.index(f"CHUNK[{i * 10}:{(i + 1) * 10}]") for i in range(5)] + assert positions == sorted(positions) + + def test_use_llm_flag_reaches_the_converter( + self, database, s3, broker, fast_queue, stub_convert + ): + make_job(broker, 1, use_llm=True) + broker.drain(worker.handle) + assert stub_convert[0][3] is True + + def test_document_is_recorded_for_the_cache( + self, database, s3, broker, fast_queue, stub_convert + ): + job_id = make_job(broker, 2) + broker.drain(worker.handle) + + cached = jobstore.find_cached_document(jobstore.get_job(job_id)["doc_key"]) + assert cached is not None + assert cached["total_chunks"] == 2 + + +class TestRetries: + def test_transient_failure_is_republished_and_succeeds( + self, database, s3, broker, fast_queue, monkeypatch + ): + monkeypatch.setattr(config, "MAX_ATTEMPTS", 3) + attempts = [] + + def flaky(file_hash, start, end, use_llm): + attempts.append(start) + if len(attempts) == 1: + raise RuntimeError("transient parser hiccup") + return "recovered" + + monkeypatch.setattr(worker, "convert_range", flaky) + + job_id = make_job(broker, 1) + broker.drain(worker.handle) + + assert len(attempts) == 2 + assert len(broker.retried) == 1 + assert broker.retried[0].attempt == 2 + assert jobstore.get_job(job_id)["status"] == "done" + + def test_retry_carries_an_incrementing_attempt( + self, database, s3, broker, fast_queue, monkeypatch + ): + monkeypatch.setattr(config, "MAX_ATTEMPTS", 3) + monkeypatch.setattr( + worker, "convert_range", + lambda *a: (_ for _ in ()).throw(RuntimeError("always fails")), + ) + make_job(broker, 1) + broker.drain(worker.handle) + + assert [t.attempt for t in broker.retried] == [2, 3] + + def test_poison_chunk_goes_to_the_failed_queue_and_fails_the_job( + self, database, s3, broker, fast_queue, monkeypatch + ): + """A document silently missing pages is worse than an error.""" + monkeypatch.setattr(config, "MAX_ATTEMPTS", 2) + + def selective(file_hash, start, end, use_llm): + if start == 10: + raise RuntimeError("corrupt page range") + return "fine" + + monkeypatch.setattr(worker, "convert_range", selective) + + job_id = make_job(broker, 3) + broker.drain(worker.handle) + + job = jobstore.get_job(job_id) + assert job["status"] == "failed" + assert job["error_chunk"] == 1 + assert len(broker.failed) == 1 + assert "corrupt page range" in broker.failed[0][1] + + def test_no_document_is_published_for_a_failed_job( + self, database, s3, broker, fast_queue, monkeypatch + ): + monkeypatch.setattr(config, "MAX_ATTEMPTS", 1) + monkeypatch.setattr( + worker, "convert_range", + lambda *a: (_ for _ in ()).throw(RuntimeError("nope")), + ) + job_id = make_job(broker, 2) + broker.drain(worker.handle) + + doc_key = jobstore.get_job(job_id)["doc_key"] + assert jobstore.get_job(job_id)["status"] == "failed" + assert storage.get_manifest(doc_key) is None + assert jobstore.find_cached_document(doc_key) is None + + def test_chunks_of_an_already_failed_job_are_dropped( + self, database, s3, broker, fast_queue, stub_convert + ): + job_id = make_job(broker, 3) + jobstore.fail_job(job_id, "failed elsewhere") + broker.drain(worker.handle) + + assert jobstore.get_job(job_id)["status"] == "failed" + assert stub_convert == [] # no work done on a dead job + + def test_assembly_failure_fails_the_job( + self, database, s3, broker, fast_queue, stub_convert, monkeypatch + ): + job_id = make_job(broker, 1) + monkeypatch.setattr(storage, "get_part", lambda *a: None) + broker.drain(worker.handle) + assert jobstore.get_job(job_id)["status"] == "failed" + + +class TestAtLeastOnceDelivery: + def test_duplicate_delivery_does_not_corrupt_the_counter( + self, database, s3, broker, fast_queue, stub_convert + ): + """ + RabbitMQ redelivers a message whose consumer died before acking, so a + chunk can genuinely be converted twice. The second completion must be a + no-op — otherwise remaining_chunks goes negative and nobody ever sees 0 + to assemble the document. + """ + job_id = make_job(broker, 2) + first = broker.ready[0] + + worker.handle(first) + worker.handle(first) # redelivery of the same message + + job = jobstore.get_job(job_id) + assert job["remaining_chunks"] == 1 + assert job["status"] != "failed" + + def test_redelivery_after_completion_never_double_assembles( + self, database, s3, broker, fast_queue, stub_convert + ): + job_id = make_job(broker, 2) + tasks = list(broker.ready) + for task in tasks: + worker.handle(task) + assert jobstore.get_job(job_id)["status"] == "done" + + # Every message redelivered after the job finished. + for task in tasks: + worker.handle(task) + assert jobstore.get_job(job_id)["status"] == "done" + assert jobstore.get_job(job_id)["remaining_chunks"] == 0 + + +class TestSourceCache: + def test_source_is_fetched_once_for_repeated_chunks( + self, database, s3, fast_queue, tmp_path, monkeypatch + ): + """A worker receiving several chunks of one document downloads once.""" + monkeypatch.setattr(config, "SOURCE_CACHE_DIR", str(tmp_path / "cache")) + src = tmp_path / "src.bin" + src.write_bytes(b"the source document") + storage.put_source("filehash", str(src)) + + downloads = [] + original = storage.get_source + + def counting(file_hash, dest): + downloads.append(file_hash) + return original(file_hash, dest) + + monkeypatch.setattr(storage, "get_source", counting) + for _ in range(4): + worker.fetch_source("filehash") + + assert len(downloads) == 1 + + def test_prune_evicts_until_within_budget(self, fast_queue, tmp_path, monkeypatch): + cache = tmp_path / "cache" + cache.mkdir() + monkeypatch.setattr(config, "SOURCE_CACHE_DIR", str(cache)) + for name in ("a", "b", "c"): + (cache / name).write_bytes(b"x" * 1000) + + worker.prune_source_cache(max_bytes=1500) + assert len(list(cache.iterdir())) <= 2 + + def test_failed_download_leaves_no_partial_cache_entry( + self, database, s3, fast_queue, tmp_path, monkeypatch + ): + """A truncated download must not later look like a valid cache hit.""" + monkeypatch.setattr(config, "SOURCE_CACHE_DIR", str(tmp_path / "cache")) + monkeypatch.setattr( + storage, "get_source", + lambda *a: (_ for _ in ()).throw(RuntimeError("network died")), + ) + + with pytest.raises(RuntimeError): + worker.fetch_source("filehash") + + cache_dir = tmp_path / "cache" + assert not (cache_dir / "filehash").exists() + assert list(cache_dir.iterdir()) == [] diff --git a/worker.py b/worker.py new file mode 100644 index 0000000..bd34c10 --- /dev/null +++ b/worker.py @@ -0,0 +1,281 @@ +""" +The consumer. Converts chunks delivered by RabbitMQ and assembles finished +documents. + +Run with ``python -m worker``. Scale by adding pods: RabbitMQ dispatches each +chunk to whichever worker is free, so there is no leader, no partition +assignment, and nothing to reconfigure when the fleet size changes. + +Each worker holds one unacked message at a time (prefetch=1). That bounds +memory to the size of a *chunk* rather than a document, and keeps dispatch fair +— a worker stuck on a slow chunk stops receiving new ones instead of hoarding a +backlog while its peers idle. +""" +import logging +import os +import signal +import socket +import tempfile +import threading +import time + +import chunker +import config +import converter +import jobstore +import pagination +import storage +import taskqueue +from taskqueue import ChunkTask + +logger = logging.getLogger(__name__) + +_stop = threading.Event() + + +def consumer_name() -> str: + return f"{socket.gethostname()}-{os.getpid()}" + + +def _touch_liveness() -> None: + """ + Mark the loop as alive for the k8s liveness probe. + + The worker serves no HTTP, so the probe checks this file's mtime instead. + Touched on every consume tick, including idle ones. + """ + try: + with open(config.WORKER_LIVENESS_FILE, "w") as fh: + fh.write(str(int(time.time()))) + except OSError: + logger.debug("could not touch liveness file", exc_info=True) + + +def fetch_source(file_hash: str) -> str: + """ + Get the source file locally, caching it under SOURCE_CACHE_DIR. + + A worker often receives several chunks of the same document, so caching + turns N downloads into one. + """ + os.makedirs(config.SOURCE_CACHE_DIR, exist_ok=True) + path = os.path.join(config.SOURCE_CACHE_DIR, file_hash) + if os.path.exists(path): + os.utime(path, None) # mark as recently used for the LRU prune + return path + + # Download to a temp name and rename, so a crashed download cannot leave a + # truncated file that later looks like a valid cache hit. + fd, temp_path = tempfile.mkstemp(dir=config.SOURCE_CACHE_DIR) + os.close(fd) + try: + storage.get_source(file_hash, temp_path) + os.replace(temp_path, path) + except Exception: + if os.path.exists(temp_path): + os.remove(temp_path) + raise + + prune_source_cache() + return path + + +def prune_source_cache(max_bytes: int = None) -> None: + """Evict least-recently-used sources until the cache fits in its budget.""" + max_bytes = config.SOURCE_CACHE_MAX_BYTES if max_bytes is None else max_bytes + try: + entries = [] + total = 0 + for name in os.listdir(config.SOURCE_CACHE_DIR): + full = os.path.join(config.SOURCE_CACHE_DIR, name) + if not os.path.isfile(full): + continue + stat = os.stat(full) + entries.append((stat.st_atime, stat.st_size, full)) + total += stat.st_size + + for _, size, full in sorted(entries): + if total <= max_bytes: + break + os.remove(full) + total -= size + except OSError: + logger.debug("source cache prune failed", exc_info=True) + + +def convert_range(file_hash: str, start: int, end: int, use_llm: bool) -> str: + """Materialise one page range from the shared source and convert it.""" + source_path = fetch_source(file_hash) + + if end == -1: + # Whole-file chunk (non-PDF, or a PDF too small to split). + return converter.convert_file(source_path, use_llm) + + fd, chunk_path = tempfile.mkstemp(suffix=".pdf") + os.close(fd) + try: + chunker.extract_pages(source_path, start, end, chunk_path) + return converter.convert_file(chunk_path, use_llm) + finally: + if os.path.exists(chunk_path): + os.remove(chunk_path) + + +def assemble(job: dict) -> None: + """ + Join every chunk's markdown into the final document and publish it. + + Called by whichever worker completed the last chunk. Parts are fetched in + chunk-index order, which is page order. + """ + job_id = job["job_id"] + doc_key = job["doc_key"] + logger.info("assembling job %s (%d chunks)", job_id, job["total_chunks"]) + + parts = [] + for index in jobstore.chunk_indices(job_id): + part = storage.get_part(doc_key, index) + if part is None: + raise RuntimeError(f"chunk {index} missing from storage") + parts.append(part) + + markdown = "\n\n".join(part for part in parts if part and part.strip()) + pages = pagination.paginate(markdown, config.PAGE_SIZE) + + storage.put_document(doc_key, pages, len(markdown)) + jobstore.record_document( + doc_key=doc_key, + file_hash=job["file_hash"], + total_pages=len(pages), + total_chunks=job["total_chunks"], + total_length=len(markdown), + page_size=config.PAGE_SIZE, + ) + jobstore.finish_job(job_id) + logger.info("job %s done: %d pages, %d chars", job_id, len(pages), len(markdown)) + + +def handle(task: ChunkTask) -> None: + """ + Convert one delivered chunk. + + Never raises: the message is acked by the consumer loop once this returns, + so any retry decision has to be made here and republished explicitly. + """ + owner = consumer_name() + job = jobstore.get_job(task.job_id) + if job is None or job["status"] in ("done", "failed"): + logger.info("dropping chunk %d: job %s is no longer active", + task.chunk_index, task.job_id) + return + + jobstore.start_chunk(task.job_id, task.chunk_index, owner, task.attempt) + started = time.monotonic() + + try: + markdown = convert_range( + task.file_hash, task.start_page, task.end_page, task.use_llm + ) + key = storage.put_part(task.doc_key, task.chunk_index, markdown) + except Exception as exc: + _handle_failure(task, exc) + return + + remaining = jobstore.complete_chunk(task.job_id, task.chunk_index, key) + logger.info( + "converted chunk %d/%d of job %s in %.1fs (%d remaining)", + task.chunk_index + 1, task.total_chunks, task.job_id, + time.monotonic() - started, max(remaining, 0), + ) + + if remaining == 0: + try: + assemble(jobstore.get_job(task.job_id)) + except Exception as exc: + logger.exception("assembly failed for job %s", task.job_id) + jobstore.fail_job(task.job_id, f"assembly failed: {exc}") + + +def _handle_failure(task: ChunkTask, exc: Exception) -> None: + """Retry a failed chunk, or fail the whole job once attempts run out.""" + error = f"{type(exc).__name__}: {exc}" + logger.exception("chunk %d of job %s failed (attempt %d/%d)", + task.chunk_index, task.job_id, task.attempt, + config.MAX_ATTEMPTS) + + if task.attempt >= config.MAX_ATTEMPTS: + # One poison chunk fails the job: silently returning a document that is + # missing pages is worse than returning an error. + taskqueue.publish_failed(task, error) + jobstore.fail_chunk(task.job_id, task.chunk_index, error) + jobstore.fail_job( + task.job_id, + f"chunk {task.chunk_index} failed after {task.attempt} attempts: {exc}", + task.chunk_index, + ) + return + + jobstore.reschedule_chunk(task.job_id, task.chunk_index, error) + taskqueue.publish_retry(task) + + +def run_forever() -> None: + owner = consumer_name() + logger.info("worker %s consuming from %s", owner, config.RABBITMQ_QUEUE) + + def on_task(task: ChunkTask) -> None: + _touch_liveness() + try: + handle(task) + except Exception: + # handle() is meant to absorb its own errors; anything escaping + # here would otherwise kill the consumer loop. + logger.exception("unhandled error on chunk %d of job %s", + task.chunk_index, task.job_id) + + while not _stop.is_set(): + try: + _touch_liveness() + taskqueue.consume(on_task, _stop.is_set) + except Exception: + if _stop.is_set(): + break + # A dropped broker connection is expected during a RabbitMQ + # restart; reconnect rather than crashlooping the pod. + logger.exception("consumer connection lost, reconnecting") + taskqueue.close() + _stop.wait(5) + + taskqueue.close() + logger.info("worker %s stopped", owner) + + +def _handle_signal(signum, _frame) -> None: + """ + Stop after the current chunk. + + The in-flight message is deliberately not nacked: if the process is killed + before it finishes, the unacked message is redelivered by the broker + automatically, which is crash-safe in a way a signal handler is not. + """ + logger.info("received signal %s, finishing current chunk", signum) + _stop.set() + + +def main() -> None: + logging.basicConfig( + level=os.environ.get("LOG_LEVEL", "INFO"), + format="%(asctime)s %(levelname)s %(name)s %(message)s", + ) + signal.signal(signal.SIGTERM, _handle_signal) + signal.signal(signal.SIGINT, _handle_signal) + + import db + + db.migrate() + storage.ensure_bucket() + run_forever() + + +if __name__ == "__main__": + main()