Main Architecture of allenai/olmocr: Async PDF Processing with vLLM

The main architecture of allenai/olmocr is a highly-parallel, asynchronous pipeline that converts PDF documents into Dolma-style JSONL files by distributing page-rendering and LLM inference tasks across multiple workers backed by a local vLLM server.

The olmocr repository implements a modular, production-ready system designed to process large-scale PDF corpora. Its main architecture centers on three distinct layers—orchestration, work distribution, and worker processing—that together handle everything from CLI argument parsing to final JSONL output generation. The system supports both local filesystem and S3 backends, implements robust retry logic for rotation errors, and optionally exports clean Markdown alongside structured Dolma documents.

Three Core Architectural Layers

The codebase organizes functionality into three primary layers that communicate through async boundaries and shared queue abstractions.

Orchestration Layer (pipeline.py)

The entry point resides in olmocr/pipeline.py, which coordinates the entire lifecycle of the processing job. This layer handles four critical responsibilities:

1. CLI Parsing – The if __name__ == "__main__" block (lines 1-124) and argument parser (lines 124-210) ingest parameters for workspace paths, model location, worker count, and PDF glob patterns.
2. Model Acquisition – The download_model() function fetches models from HuggingFace, S3, or local paths, while vllm_server_host() launches the vLLM inference server as an async subprocess.
3. Queue Initialization – Based on the --pdfs argument, the system builds a work index and invokes WorkQueue.populate_queue() to group files (defaulting to ~200 pages per group) into a compressed CSV index.
4. Worker Spawning – asyncio.create_task(worker(...)) launches the configured number of concurrent workers, while a background metrics task tracks token usage and queue depth.

Work Distribution Layer (work_queue.py)

Located in olmocr/work_queue.py, this layer provides a pluggable backend abstraction that ensures disjoint work distribution across distributed or local environments.

The Backend base class defines async methods for index persistence and lock management, with two concrete implementations:

• LocalBackend – Uses the local filesystem for indices and lock files.
• S3Backend – Uses S3 storage with zstd-compressed CSVs for the work index.

The WorkQueue class manages the index format—a gz-compressed CSV where each row contains <hash>,path1,path2,…. The hash is a deterministic SHA-1 of the sorted path group (_compute_workgroup_hash). Before processing, workers check is_worker_lock_taken() to detect stale locks, then create a lock file via create_worker_lock(). Upon completion, the worker removes the lock and creates a done flag using create_done_flag(). The initialize_queue() method removes already-completed items and shuffles the remaining work for even distribution.

Worker Processing Layer (pipeline.py)

Worker logic resides primarily in olmocr/pipeline.py (lines 155-440) and executes the core PDF-to-text transformation:

1. PDF Acquisition – Workers download PDFs or extract them from tarballs using S3 utilities.
2. Page Rendering – Each page is rendered to a base-64 PNG using render_pdf_to_base64png from olmocr/data/renderpdf.py.
3. LLM Query Construction – The build_page_query() function (lines 106-145) constructs a JSON payload containing the image and a static YAML prompt from build_no_anchoring_v4_yaml_prompt, sending it to the vLLM server via apost with exponential backoff.
4. Response Parsing – The FrontMatterParser extracts YAML frontmatter into a PageResponse dataclass containing fields like natural_text, primary_language, and is_rotation_valid.
5. Retry Logic – Rotation errors trigger sequential retries with corrected rotation; other failures retry in parallel once the vLLM queue empties.
6. Document Aggregation – build_dolma_document() concatenates page text, computes token totals, builds page spans, and creates the final Dolma JSON object.
7. Output Writing – Results are written as JSONL files (uploaded to S3 or copied locally), with optional Markdown output mapped via get_markdown_path() to the workspace's markdown/ tree.

Supporting Services

Beyond the core layers, several specialized modules provide cross-cutting functionality:

• S3 Utilities (olmocr/s3_utils.py) – Handles download_zstd_csv, upload_zstd_csv, download_directory, and parse_s3_path for both backends and worker file operations.
• Image Utilities (olmocr/image_utils.py) – Detects JPEG/PNG formats, converts images to PDF bytes, and handles rotation/resize operations before LLM ingestion.
• PDF Filtering (olmocr/filter/filter.py) – Pre-screens PDFs for English language content and spam detection to avoid expensive inference on low-quality documents.
• Metrics Tracking (olmocr/metrics.py) – The MetricsKeeper class records token usage, while WorkerTracker maintains per-worker status tables for real-time monitoring.
• Prompt Templates (olmocr/prompts/*.py) – Contains YAML prompt templates that guide the LLM's extraction behavior, including anchoring instructions.

End-to-End Execution Flow

The system follows a deterministic async flow from CLI invocation to final output:

flowchart TD
    A[CLI (pipeline.py)] --> B[download_model / launch vLLM]
    B --> C[WorkQueue.populate_queue]
    C --> D[WorkQueue.initialize_queue]
    D --> E[Worker tasks]
    E --> F[process_pdf / process_tarball]
    F --> G[render each page → PNG]
    G --> H[build LLM request (JSON)]
    H --> I[vLLM server (HTTP POST)]
    I --> J[parse model response → PageResponse]
    J --> K[retry logic (rotation / parallel)]
    K --> L[build_dolma_document]
    L --> M[Write JSONL (S3/local)]
    M --> N[Optional Markdown output]

Practical Code Examples

Running the Pipeline Locally

To process PDFs with 8 workers and generate Markdown output:

python -m olmocr.pipeline \
    s3://my-bucket/workspace \
    --pdfs s3://my-bucket/pdfs/*.pdf \
    --model allenai/olmOCR-2-7B-1025-FP8 \
    --workers 8 \
    --markdown

This command creates a work index, starts the vLLM server, spawns eight asynchronous workers, and writes both JSONL and Markdown results to the specified S3 workspace.

Adding New PDFs to an Existing Queue

Using the WorkQueue API directly to populate work items:

from olmocr.work_queue import WorkQueue, LocalBackend

backend = LocalBackend("/tmp/olmocr_workspace")
queue = WorkQueue(backend)

# Add new PDF paths (group size = 200 pages)

await queue.populate_queue(["/data/papers/paper1.pdf", "/data/papers/paper2.pdf"], items_per_group=200)

Using the S3 Backend Directly

For custom integrations with S3 storage:

import boto3
from olmocr.work_queue import WorkQueue, S3Backend

s3 = boto3.client("s3")
backend = S3Backend(s3, "s3://my-bucket/olmocr_workspace")
queue = WorkQueue(backend)

await queue.initialize_queue()
work_item = await queue.get_work()
print(work_item.hash, work_item.work_paths)

Summary

• Main architecture consists of three layers: orchestration (pipeline.py), work distribution (work_queue.py), and worker processing (pipeline.py worker functions).
• WorkQueue abstraction supports both local filesystem and S3 backends with deterministic hashing, file locking, and completion tracking.
• Worker processing renders PDF pages to images, queries a local vLLM server, parses YAML frontmatter, and aggregates results into Dolma-style JSONL documents.
• Retry logic handles rotation errors sequentially and other failures in parallel when the inference queue empties.
• Key files include olmocr/pipeline.py for orchestration, olmocr/work_queue.py for queue management, and olmocr/s3_utils.py for cloud storage operations.

Frequently Asked Questions

How does olmocr handle work distribution across multiple workers?

The WorkQueue class in olmocr/work_queue.py manages distribution through a compressed CSV index where PDF paths are grouped into work items (default ~200 pages per group). Each work item receives a deterministic SHA-1 hash, and workers acquire exclusive locks via create_worker_lock() before processing. The system supports both LocalBackend for filesystem-based coordination and S3Backend for distributed cloud processing, using done-flag files to track completion and prevent duplicate work.

What retry mechanisms exist for failed page processing?

According to the worker implementation in olmocr/pipeline.py, the system distinguishes between rotation errors and generic failures. When is_rotation_valid is false in the PageResponse, the worker performs sequential retries with corrected rotation angles. For other failures, the worker may retry in parallel once the vLLM server queue empties, using exponential backoff via the apost function for HTTP POST requests to the inference server.

Can olmocr process PDFs stored locally instead of S3?

Yes. While the repository emphasizes S3 support through S3Backend, you can use LocalBackend from olmocr/work_queue.py to process local PDF files. The LocalBackend class implements the same Backend interface using local filesystem operations for index storage, lock files, and completion flags. Simply pass local paths to the --pdfs argument and specify a local workspace directory when running pipeline.py.

What is the output format of the processed documents?

The pipeline produces Dolma-style JSONL documents where each line represents a processed PDF containing aggregated page text, token counts, and page spans. If the --markdown flag is enabled, the system also generates plain .md files containing the extracted natural text. The JSON structure includes fields such as natural_text, primary_language, and is_rotation_valid parsed from the LLM's YAML frontmatter response.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}