# How the Beaker Cluster Integration Enables Distributed Processing in olmOCR

> Discover how Beaker cluster integration in olmOCR enables distributed PDF processing across hundreds of GPU nodes with a single command. Scale your OCR tasks efficiently.

- Repository: [Ai2/olmocr](https://github.com/allenai/olmocr)
- Tags: deep-dive
- Published: 2026-07-06

---

**The Beaker cluster integration in olmOCR allows you to scale PDF processing across hundreds of GPU nodes by submitting a single command that launches a distributed experiment on AI2's Beaker platform.**

The `allenai/olmocr` repository provides a seamless bridge between local development and large-scale distributed inference. When you activate the Beaker cluster integration, the pipeline submits itself as a containerized experiment that automatically shards work across multiple GPU replicas, eliminating the need for manual cluster management.

## How Beaker Cluster Integration Works

The integration operates as a thin wrapper around the existing `olmocr.pipeline` module. When triggered, it packages your current command-line arguments into a **BeakerExperimentSpec**, handles authentication automatically, and relies on the built-in **WorkQueue** abstraction to parallelize tasks without modifying core pipeline logic.

### CLI Flag Detection and Job Submission

The entry point for distributed processing begins in [`olmocr/pipeline.py`](https://github.com/allenai/olmocr/blob/main/olmocr/pipeline.py) at lines 1261–1270. When you append the `--beaker` flag to your command, the `main()` function detects this argument and routes execution to `submit_beaker_job(args)` instead of running locally.

The supported Beaker-specific flags include:
- `--beaker_workspace` – The target Beaker workspace (e.g., `ai2/olmocr`)
- `--beaker_cluster` – The cluster constraint (e.g., `gpu-a100-large`)
- `--beaker_gpus` – The number of GPU replicas to launch
- `--beaker_priority` – The job priority level (`low`, `normal`, or `high`)

### Secure Credential Management

Before submitting the experiment, `submit_beaker_job` (defined at lines 992–1065 in [`olmocr/pipeline.py`](https://github.com/allenai/olmocr/blob/main/olmocr/pipeline.py)) initializes the Beaker client using `Beaker.from_env()`. The function checks for required **Weka** and **AWS** credentials in your local environment. If these secrets are missing, the script interactively prompts you and writes them into the Beaker workspace, ensuring each GPU replica has access to the shared S3 storage backend without exposing credentials in the command history.

### Experiment Specification and Resource Allocation

The core of the integration constructs a `BeakerExperimentSpec` that defines how resources are allocated across the cluster. Key configuration fields include:

- **`replicas=args.beaker_gpus`** – Creates one container replica per requested GPU, so `--beaker_gpus 4` launches four identical GPU workers.
- **`constraints=BeakerConstraints(cluster=[...])`** – Pins the experiment to specific hardware clusters (e.g., `gpu-a100-xlarge`).
- **`priority=BeakerJobPriority[args.beaker_priority]`** – Controls scheduling precedence in the cluster queue.
- **`image=BeakerImageSource(beaker="jakep/olmocr-inference-<VERSION>")`** – References the pre-built container image containing the pipeline code and dependencies.
- **`command=["python", "-m", "olmocr.pipeline"] + args_list`** – Executes the same Python module on each worker, automatically stripping the `--beaker` flag to prevent recursive submission.

Each replica receives environment variables including `BEAKER_JOB_NAME`, `OWNER`, and `HF_HUB_OFFLINE=1` to ensure deterministic, air-gapped execution.

### Distributed Work Queue Execution

Once the experiment launches, each GPU replica executes the standard `worker()` coroutine defined in [`olmocr/pipeline.py`](https://github.com/allenai/olmocr/blob/main/olmocr/pipeline.py). The pipeline leverages the **`WorkQueue`** abstraction from [`olmocr/work_queue.py`](https://github.com/allenai/olmocr/blob/main/olmocr/work_queue.py), which supports both `LocalBackend` and `S3Backend`. 

Because all replicas connect to the same S3 workspace (specified by the positional argument in your original command), they automatically pull items from the shared queue. This architecture provides **true distributed processing** without requiring MPI or custom communication protocols—each worker independently processes PDFs or tarballs and writes results back to `args.workspace/results/output_<hash>.jsonl`.

## Running Your First Beaker Cluster Job

To process PDFs on a Beaker cluster with four A100 GPUs, use the following command structure:

```bash
python -m olmocr.pipeline s3://my-bucket/olmocr-workspace \
    --model allenai/olmOCR-2-7B-1025-FP8 \
    --workers 20 \
    --max_page_retries 8 \
    --beaker \
    --beaker_workspace ai2/olmocr \
    --beaker_cluster gpu-a100-large \
    --beaker_gpus 4 \
    --beaker_priority normal

```

Under the hood, the `submit_beaker_job` function converts this into the following Python specification:

```python
def submit_beaker_job(args):
    from beaker import (
        Beaker, BeakerConstraints, BeakerEnvVar, BeakerExperimentSpec,
        BeakerImageSource, BeakerJobPriority, BeakerResultSpec,
        BeakerRetrySpec, BeakerTaskContext, BeakerTaskResources,
        BeakerTaskSpec,
    )
    
    # ... credential handling omitted ...

    
    experiment_spec = BeakerExperimentSpec(
        budget="ai2/oe-base",
        description=task_name,
        tasks=[
            BeakerTaskSpec(
                name=task_name,
                replicas=args.beaker_gpus,  # One replica per GPU

                context=BeakerTaskContext(
                    priority=BeakerJobPriority[args.beaker_priority],
                    preemptible=True,
                ),
                image=BeakerImageSource(beaker=beaker_image),
                command=["python", "-m", "olmocr.pipeline"] + args_list,
                env_vars=[
                    BeakerEnvVar(name="BEAKER_JOB_NAME", value=task_name),
                    BeakerEnvVar(name="HF_HUB_OFFLINE", value="1"),
                    # ... additional secrets ...

                ],
                resources=BeakerTaskResources(gpu_count=1, memory="125GB"),
                constraints=BeakerConstraints(
                    cluster=[args.beaker_cluster]
                ),
                result=BeakerResultSpec(path="/noop-results"),
            )
        ],
        retry=BeakerRetrySpec(allowed_task_retries=10),
    )
    
    workload = b.experiment.create(spec=experiment_spec)
    print(f"Experiment URL: https://beaker.org/ex/{workload.experiment.id}")

```

After submission, the console prints the experiment URL where you can monitor progress and logs for each GPU replica.

## Summary

- The **Beaker cluster integration** is triggered by the `--beaker` flag in [`olmocr/pipeline.py`](https://github.com/allenai/olmocr/blob/main/olmocr/pipeline.py), which redirects execution to `submit_beaker_job()` at lines 992–1065.
- The function automatically manages **AWS and Weka credentials**, ensuring secure access to shared S3 workspaces across all replicas.
- Resource allocation is controlled through **BeakerExperimentSpec**, mapping `--beaker_gpus` to container replicas and `--beaker_cluster` to hardware constraints.
- **WorkQueue** ([`olmocr/work_queue.py`](https://github.com/allenai/olmocr/blob/main/olmocr/work_queue.py)) enables automatic work distribution, allowing each GPU replica to independently pull PDFs from the shared queue and write results without additional coordination code.
- Results aggregate to the shared S3 workspace and can be analyzed using the standard `print_stats` command, regardless of whether processing occurred locally or on a Beaker cluster.

## Frequently Asked Questions

### What is Beaker and why does olmOCR use it?

**Beaker** is AI2's internal experiment management platform for running containerized workloads on GPU clusters. According to the `allenai/olmocr` source code, the integration allows researchers to scale from single-GPU debugging to production-scale processing of thousands of PDFs without changing application code, leveraging Beaker's job scheduling, priority queuing, and automatic resource allocation.

### How do I configure AWS and Weka credentials for Beaker jobs?

The `submit_beaker_job` function in [`olmocr/pipeline.py`](https://github.com/allenai/olmocr/blob/main/olmocr/pipeline.py) (lines 992–1010) checks your local environment for existing Beaker secrets. If **AWS** or **Weka** credentials are missing, the script interactively prompts you for the access keys and writes them into the Beaker workspace namespace. This one-time setup ensures that every GPU replica launched by the experiment has authenticated access to your S3 workspace.

### How does work distribution work across multiple GPUs?

Work distribution relies on the **`WorkQueue`** abstraction defined in [`olmocr/work_queue.py`](https://github.com/allenai/olmocr/blob/main/olmocr/work_queue.py). When you specify `--beaker_gpus 4`, Beaker launches four identical containers, each running the standard pipeline worker. All replicas connect to the same S3-backed queue, atomically pulling PDF paths from the shared list and writing OCR results back to `results/output_<hash>.jsonl`. This design requires no master-worker communication; the queue itself handles load balancing.

### Can I run the pipeline locally and on Beaker with the same command?

Yes. The pipeline logic in [`olmocr/pipeline.py`](https://github.com/allenai/olmocr/blob/main/olmocr/pipeline.py) is identical between local and distributed modes. Simply remove the `--beaker` flag and associated Beaker arguments to run locally, or add them to submit to the cluster. The `--workers` flag controls thread-level parallelism within a single process, while `--beaker_gpus` controls the number of distributed replicas, allowing you to tune both levels of parallelism independently.