diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml new file mode 100644 index 0000000..a0f72b2 --- /dev/null +++ b/.github/workflows/test.yaml @@ -0,0 +1,182 @@ +name: Tests + +on: + push: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + unit-tests: + name: Unit Tests + runs-on: blacksmith-4vcpu-ubuntu-2404 + timeout-minutes: 20 + steps: + - uses: actions/checkout@v6 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + + - name: Install system dependencies + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + ffmpeg libimage-exiftool-perl libmagic1 + + - name: Install Python dependencies + run: | + pip install -r requirements.txt + pip install pytest reportlab + + - name: Run unit tests + env: + PYTHONPATH: ${{ github.workspace }} + run: pytest tests/ -v + + k8s-benchmark: + name: PDF Benchmark (Kubernetes) + runs-on: ubuntu-latest + timeout-minutes: 45 + permissions: + contents: read + pull-requests: write + env: + IMAGE_TAG: ci-${{ github.sha }} + NAMESPACE: markitdown-server + steps: + - uses: actions/checkout@v6 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Build image + uses: docker/build-push-action@v6 + with: + context: . + platforms: linux/amd64 + push: false + load: true + tags: markitdown-server:${{ env.IMAGE_TAG }} + cache-from: type=gha + cache-to: type=gha,mode=max + + - name: Start minikube + uses: medyagh/setup-minikube@v0.0.21 + with: + driver: docker + cpus: 4 + memory: 6000m + addons: metrics-server + + - name: Load image into minikube + run: minikube image load "markitdown-server:${IMAGE_TAG}" + + - 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. + kubectl create secret generic markitdown-server-secret \ + -n "$NAMESPACE" \ + --from-literal=OPENAI_API_KEY=ci-placeholder \ + --from-literal=ADMIN_API_KEY=ci-test-key \ + --dry-run=client -o yaml | kubectl apply -f - + + - name: Deploy + run: | + cd k8s + 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 + + - 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 + + # No port-forward here on purpose: the benchmark restarts the pod between + # books, which tears a forward down, so the script manages its own. + - name: Benchmark Foundation books + env: + ADMIN_API_KEY: ci-test-key + LOCAL_PORT: "8080" + run: | + set -o pipefail + bash scripts/k8s-benchmark.sh | tee /tmp/benchmark.txt + + - name: Build benchmark report + if: always() + run: | + # On failure the script appends per-book pod logs, so cap the size: + # a GitHub comment is rejected outright past ~65k characters. + if [ -f /tmp/benchmark.txt ]; then + head -c 55000 /tmp/benchmark.txt > /tmp/benchmark-trimmed.txt + if [ "$(wc -c < /tmp/benchmark.txt)" -gt 55000 ]; then + echo '... output truncated, see the k8s-logs artifact ...' \ + >> /tmp/benchmark-trimmed.txt + fi + else + echo 'benchmark did not produce output' > /tmp/benchmark-trimmed.txt + fi + { + echo '' + echo '## PDF conversion benchmark' + echo + echo '```' + cat /tmp/benchmark-trimmed.txt + echo '```' + echo + echo "Commit \`${GITHUB_SHA:0:7}\` · [workflow run](${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID})" + } > /tmp/benchmark-report.md + cat /tmp/benchmark-report.md >> "$GITHUB_STEP_SUMMARY" + + # Posted even when the benchmark fails -- a budget violation is exactly + # the result worth surfacing on the PR. + - name: Comment benchmark on PR + if: always() + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + # This workflow runs on push, not pull_request, so the PR number has + # to be looked up from the branch. + PR=$(gh pr list --head "$GITHUB_REF_NAME" --state open \ + --json number --jq '.[0].number // empty') + if [ -z "$PR" ]; then + echo "No open PR for branch $GITHUB_REF_NAME; skipping comment." + exit 0 + fi + # Edit the previous report rather than stacking a new comment on + # every push; fall back to creating one on the first run. + gh pr comment "$PR" --body-file /tmp/benchmark-report.md --edit-last \ + || gh pr comment "$PR" --body-file /tmp/benchmark-report.md + + - name: Collect pod diagnostics + if: always() + run: | + mkdir -p /tmp/k8s-logs + # 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 get pods -n "$NAMESPACE" -o wide > /tmp/k8s-logs/pods.txt 2>&1 || true + kubectl describe pods -n "$NAMESPACE" -l app=markitdown-server \ + > /tmp/k8s-logs/describe.txt 2>&1 || true + + - name: Upload diagnostics + uses: actions/upload-artifact@v7 + if: always() + with: + name: k8s-logs + path: /tmp/k8s-logs/ + retention-days: 7 diff --git a/converter.py b/converter.py index f289e42..5e553e0 100644 --- a/converter.py +++ b/converter.py @@ -9,6 +9,8 @@ import requests import tempfile +from pdf_chunk import convert_pdf_chunked, should_chunk + logger = logging.getLogger(__name__) client = OpenAI( @@ -93,9 +95,17 @@ def convert(url: str) -> DocumentConverterResult: downloaded_at = time.monotonic() logger.info("downloaded %s in %.3fs", url, downloaded_at - started) - md = MarkItDown(llm_client=client, llm_model=model) try: - converted = md.convert(temp_file) + # 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) diff --git a/k8s/deployment.yaml b/k8s/deployment.yaml index 2490675..47f2e8a 100644 --- a/k8s/deployment.yaml +++ b/k8s/deployment.yaml @@ -24,18 +24,29 @@ spec: image: ghcr.io/sirily11/markitdown-server:latest resources: limits: - memory: "1024Mi" - cpu: "500m" + # 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: "384Mi" - cpu: "250m" + 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 diff --git a/k8s/ingress.yaml b/k8s/ingress.yaml index a43c48b..abdd920 100644 --- a/k8s/ingress.yaml +++ b/k8s/ingress.yaml @@ -6,6 +6,11 @@ metadata: 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. + nginx.ingress.kubernetes.io/proxy-connect-timeout: "60" + 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" diff --git a/pdf_chunk.py b/pdf_chunk.py new file mode 100644 index 0000000..5f2b31c --- /dev/null +++ b/pdf_chunk.py @@ -0,0 +1,134 @@ +""" +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 3f5c8d5..2f76154 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,14 @@ 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)" + "redis (>=5.2.1,<6.0.0)", + "pypdf (>=6.1.1,<7.0.0)" +] + +[project.optional-dependencies] +test = [ + "pytest (>=8.0.0,<9.0.0)", + "reportlab (>=4.0.0,<5.0.0)", ] diff --git a/requirements.txt b/requirements.txt index 7b9785d..7bd5712 100644 --- a/requirements.txt +++ b/requirements.txt @@ -43,6 +43,7 @@ pdfplumber==0.11.10 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 diff --git a/scripts/k8s-benchmark.sh b/scripts/k8s-benchmark.sh new file mode 100644 index 0000000..a7266db --- /dev/null +++ b/scripts/k8s-benchmark.sh @@ -0,0 +1,335 @@ +#!/usr/bin/env bash +# +# Convert real Asimov Foundation PDFs through the service running in the +# cluster, enforcing a per-book time and memory budget. +# +# The budgets are regression guards, not performance targets: they are sized to +# catch the failure modes that actually matter (chunking silently falling back +# to whole-document conversion, or peak memory creeping back toward the pod +# limit) without failing on normal runner variance. +# +# Memory comes from the pod's cgroup high-water mark rather than `kubectl top`, +# which samples on an interval and routinely misses the spike that matters. +# That mark is cumulative and the cgroup is mounted read-only, so it cannot be +# 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. +# +# 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 + +NAMESPACE="${NAMESPACE:-markitdown-server}" +LOCAL_PORT="${LOCAL_PORT:-8080}" +BASE_URL="http://localhost:${LOCAL_PORT}" +# 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. +MEM_BUDGET_PCT="${MEM_BUDGET_PCT:-90}" +# Per-book pod logs land here; the workflow uploads the directory as an artifact. +LOG_DIR="${LOG_DIR:-/tmp/benchmark-pod-logs}" +RAW="https://raw.githubusercontent.com/GSimas/Asimov/master/Books" + +# name|url|budget_seconds +# +# Budgets are ~2.5x the time measured on a native CI runner at 2 CPU / 2 +# workers (I Robot 12.9s, Prelude 21.0s, Foundation and Earth 71.4s, Trilogy +# 32.5s). That leaves room for runner variance while still catching a gross +# regression such as chunking silently falling back to whole-document +# conversion. Tighten further once several runs have established the spread. +# +# Ordered smallest -> largest so a failure surfaces on a cheap case first. +BOOKS=( + "I Robot.pdf|${RAW}/I_Robot/I%20Robot.pdf|35" + "Prelude to Foundation.pdf|${RAW}/Prelude_Foundation/Prelude%20to%20Foundation.pdf|55" + # Slower than the Trilogy despite having half the pages: its pages take + # markitdown's expensive form-detection path rather than the fast pdfminer + # one, costing ~4x per page. + "Foundation and Earth.pdf|${RAW}/Foundation_Earth/Foundation%20and%20Earth.pdf|180" + "Foundation Trilogy.pdf|${RAW}/Foundation/Foundation%20Trilogy.pdf|85" +) + +PF_PID="" + +stop_port_forward() { + [ -z "$PF_PID" ] && return 0 + kill "$PF_PID" 2>/dev/null + # Reap it, so bash job control does not print "Terminated" over the report. + wait "$PF_PID" 2>/dev/null + PF_PID="" +} +trap stop_port_forward EXIT + +# Convert a Kubernetes memory quantity to MiB. The API server normalises what +# the manifest says, so a limit written as "2048Mi" reads back as "2Gi" -- any +# parser that only handles one suffix silently produces nonsense percentages. +to_mib() { + awk -v q="$1" 'BEGIN{ + if (q ~ /Ki$/) { sub(/Ki$/,"",q); printf "%.0f", q/1024 } + else if (q ~ /Mi$/) { sub(/Mi$/,"",q); printf "%.0f", q } + else if (q ~ /Gi$/) { sub(/Gi$/,"",q); printf "%.0f", q*1024 } + else if (q ~ /Ti$/) { sub(/Ti$/,"",q); printf "%.0f", q*1048576 } + else if (q ~ /k$/) { sub(/k$/,"",q); printf "%.0f", q*1000/1048576 } + else if (q ~ /M$/) { sub(/M$/,"",q); printf "%.0f", q*1000000/1048576 } + else if (q ~ /G$/) { sub(/G$/,"",q); printf "%.0f", q*1000000000/1048576 } + else if (q ~ /^[0-9.]+$/) { printf "%.0f", q/1048576 } + else { printf "0" } + }' +} + +# Serving pods only. +# +# `--field-selector=status.phase=Running` is NOT enough: a pod being deleted +# stays in phase Running for the whole termination grace period (120s here), so +# 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 \ + | jq -r '.items[] + | select(.metadata.deletionTimestamp == null) + | select(.status.phase == "Running") + | select(any(.status.containerStatuses[]?; .ready)) + | .metadata.name' +} + +# Wait for exactly one serving pod, so before/after samples cannot straddle two +# different pods and report a phantom restart. Returns non-zero (and prints +# nothing) if the set never settles -- callers must treat that as an error, not +# as a zero reading. +pod_name() { + local names count + for _ in $(seq 1 90); do + names="$(running_pod_names)" + count="$(printf '%s\n' "$names" | grep -c .)" + if [ "$count" = "1" ]; then + printf '%s' "$names" + return 0 + fi + sleep 2 + done + return 1 +} + +# An OOM kill restarts the container in place, keeping the pod name, so the +# restart counter is the signal that actually catches it. +restart_count() { + kubectl get pod -n "$NAMESPACE" "$1" \ + -o jsonpath='{.status.containerStatuses[0].restartCount}' 2>/dev/null +} + +peak_mb() { + kubectl exec -n "$NAMESPACE" "$1" -- sh -c \ + 'cat /sys/fs/cgroup/memory.peak 2>/dev/null \ + || cat /sys/fs/cgroup/memory/memory.max_usage_in_bytes 2>/dev/null \ + || echo 0' 2>/dev/null | awk '{printf "%.0f", $1/1048576}' +} + +start_port_forward() { + stop_port_forward + kubectl port-forward -n "$NAMESPACE" svc/markitdown-server \ + "${LOCAL_PORT}:8080" >/dev/null 2>&1 & + PF_PID=$! + for _ in $(seq 1 30); do + curl -sf "${BASE_URL}/" >/dev/null 2>&1 && return 0 + sleep 2 + done + echo "ERROR: service did not become reachable on ${BASE_URL}" + return 1 +} + +slugify() { + printf '%s' "$1" | tr -c 'A-Za-z0-9' '-' +} + +# Capture logs for a book now: the next book restarts the pod, which discards +# them. Grabbing logs only at the end would report the last book's pod no +# matter which one actually failed. +# +# 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. +capture_logs() { + local slug="$1" pod i=0 + while IFS= read -r pod; do + [ -z "$pod" ] && continue + i=$((i + 1)) + kubectl logs -n "$NAMESPACE" "$pod" --tail=200 \ + > "${LOG_DIR}/${slug}.${i}-${pod}.log" 2>&1 || true + # A restarted container's interesting output is in the dead instance. + kubectl logs -n "$NAMESPACE" "$pod" --previous --tail=200 \ + > "${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 \ + -o jsonpath='{range .items[*]}{.metadata.name}{"\n"}{end}' 2>/dev/null)" +} + +restart_pod() { + kubectl rollout restart deployment/markitdown-server -n "$NAMESPACE" >/dev/null + kubectl rollout status deployment/markitdown-server -n "$NAMESPACE" \ + --timeout=300s >/dev/null || return 1 + start_port_forward +} + +MEM_LIMIT_RAW=$(kubectl get deployment markitdown-server -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" \ + -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 " 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}" + 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 +printf '%-30s %7s %8s %8s %8s %9s %s\n' \ + BOOK SECONDS BUDGET PEAK_MB LIMIT_PCT CHARS VERDICT +printf '%s\n' "-----------------------------------------------------------------------------------------" + +failures=0 +FAILED_BOOKS=() +for entry in "${BOOKS[@]}"; do + IFS='|' read -r name url budget <<< "$entry" + budget=$(awk -v b="$budget" -v s="$BUDGET_SCALE" 'BEGIN{printf "%.0f", b*s}') + slug="$(slugify "$name")" + + if ! restart_pod; then + printf '%-30s %7s %8s %8s %8s %9s %s\n' \ + "$name" "-" "$budget" "-" "-" "-" "FAIL pod did not come up" + capture_logs "$slug" + FAILED_BOOKS+=("$slug|$name") + failures=$((failures + 1)) + continue + fi + + # Never default an unreadable value to zero: that turns "could not measure" + # into a confident wrong number, which is how the phantom OOM verdicts got + # reported. Bail out loudly instead. + 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 + capture_logs "$slug" + FAILED_BOOKS+=("$slug|$name") + failures=$((failures + 1)) + continue + fi + restarts_before="$(restart_count "$pod")" + [ -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. + 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)" + 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="" + restarts_after="$(restart_count "$pod_after")" + [ -z "$restarts_after" ] && restarts_after=0 + peak="$(peak_mb "$pod_after")" + # A failed read must stay distinct from a genuine low number. The pod always + # uses some memory, so 0 here means the cgroup read did not work. + if [ -z "$peak" ] || [ "$peak" = "0" ]; then + peak="?" + fi + + verdict="ok" + if [ -z "$pod_after" ]; then + verdict="FAIL lost track of the serving pod after the request" + failures=$((failures + 1)) + elif [ "$pod_after" != "$pod" ] || [ "$restarts_after" != "$restarts_before" ]; then + # The peak reading is meaningless here: a container restart zeroes the + # cgroup high-water mark, so a low number does not mean low usage. + verdict="FAIL container restarted mid-request (likely OOM kill)" + failures=$((failures + 1)) + 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" + failures=$((failures + 1)) + elif [ "$(awk -v s="$seconds" -v b="$budget" 'BEGIN{print (s>b)?1:0}')" = "1" ]; then + verdict="FAIL over time budget" + failures=$((failures + 1)) + fi + + pct="-" + if [ "$MEM_LIMIT_MB" != "0" ] && [ "$peak" != "?" ]; then + raw_pct="$(awk -v p="$peak" -v l="$MEM_LIMIT_MB" 'BEGIN{printf "%.0f", 100*p/l}')" + if [ "$verdict" = "ok" ] && \ + [ "$(awk -v p="$raw_pct" -v m="$MEM_BUDGET_PCT" 'BEGIN{print (p>m)?1:0}')" = "1" ]; then + verdict="FAIL over memory budget (>${MEM_BUDGET_PCT}%)" + failures=$((failures + 1)) + fi + 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") + fi +done + +echo +if [ "$failures" -gt 0 ]; then + echo "FAILED: $failures budget violation(s)" + # Guarded: expanding an empty array under `set -u` errors on bash < 4.4, and + # a failure before any pod resolved leaves this list empty. + if [ "${#FAILED_BOOKS[@]}" -gt 0 ]; then + for entry in "${FAILED_BOOKS[@]}"; do + slug="${entry%%|*}" + label="${entry##*|}" + found=0 + for logfile in "${LOG_DIR}/${slug}."*.log; do + [ -f "$logfile" ] || continue + found=1 + echo + echo "--- ${label}: $(basename "$logfile") (last 60 lines) ---" + tail -60 "$logfile" + done + if [ "$found" = "0" ]; then + echo + echo "--- ${label}: no pod logs captured ---" + fi + done + fi + echo + echo "--- pod status ---" + kubectl get pods -n "$NAMESPACE" -l app=markitdown-server 2>&1 | head -10 + exit 1 +fi +echo "All books converted within time and memory budgets." +echo "(per-book pod logs in ${LOG_DIR})" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..431f3d0 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,45 @@ +import pytest +from reportlab.lib.pagesizes import letter +from reportlab.pdfgen import canvas + + +def write_pdf(path: str, pages: int, lines_per_page: int = 20) -> str: + """ + Write a prose-style PDF where every page carries a unique marker. + + The markers let tests assert on both content preservation and page ordering + after a document has been split and reassembled. + """ + c = canvas.Canvas(path, pagesize=letter) + for page in range(1, pages + 1): + text = c.beginText(72, 720) + text.textLine(f"PAGEMARKER{page:05d}") + for line in range(lines_per_page): + text.textLine( + f"Lorem ipsum dolor sit amet consectetur adipiscing elit, " + f"line {line} of page {page}." + ) + c.drawText(text) + c.showPage() + c.save() + return path + + +@pytest.fixture +def make_pdf(tmp_path): + """Factory fixture: make_pdf(pages) -> path to a generated PDF.""" + def _make(pages: int, name: str = "doc.pdf", lines_per_page: int = 20) -> str: + return write_pdf(str(tmp_path / name), pages, lines_per_page) + + return _make + + +@pytest.fixture +def small_pdf(make_pdf): + return make_pdf(3) + + +@pytest.fixture +def multi_chunk_pdf(make_pdf): + """25 pages — splits into multiple chunks at the default 20/chunk.""" + return make_pdf(25) diff --git a/tests/test_pdf_chunk.py b/tests/test_pdf_chunk.py new file mode 100644 index 0000000..b39dc14 --- /dev/null +++ b/tests/test_pdf_chunk.py @@ -0,0 +1,158 @@ +""" +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)