Scaling AI Document Processing on Kubernetes with Ray & KubeRay

Every answer Needl.ai gives starts with reading. Before a document can be searched, quoted or cited, each of its pages has to become text, and for long PDFs that means running OCROCROptical character recognition: a model that recognises the letters in a picture of a page., layout and table models on GPUsGPUA graphics processor. It is the usual chip for running AI models and the most expensive part of a machine to rent.. On one GPU, a 463-page document took 38 minutes. Split into 50-page chunks and read by ten GPUs at the same time with Ray on Kubernetes, it took 10 minutes, GPU start-up included.

Part one: the plain-language version. Three short sections on the problem, the idea and where the work runs. Every technical term is underlined where it first appears; hover over it, or tap it on a phone, for a definition.

Why long PDFs are hard

Reading a page well takes more than recognising letters. OCR finds the words, layout analysisLayout analysisWorking out a page's structure: headings, columns, tables and figures, and the order to read them in. works out headings, columns and reading order, and table extractionTable extractionRebuilding a table's rows and columns from a page, so its numbers can be searched and compared. rebuilds rows and columns so the numbers can be searched. These models run on GPUs, the chips built for AI work, and each GPU has a fixed amount of memoryGPU memoryThe memory on the graphics card itself. A model and the pages it is working on must fit in it. of its own.

A short document fits easily. A 500-page annual report or prospectus doesn't. Fed to one machine in one go, it is slow, because the pages are read one after another, and its memory use spikes unpredictably as pages pass through the model. One spike past the GPU's memory and the run crashes. Engineers call that an OOMOOMOut of memory: a run that asks for more memory than the machine has, and crashes., for out of memory, and a crash sends the whole document back to the start.

The idea: split first, then read in parallel

The fix is to stop treating a document as one job. Each PDF is cut into chunksChunkA bounded slice of a document that one task handles. In our pipeline, 50 pages. of 50 pages, and each chunk is read on its own GPU, at the same time as the others. Inside a chunk, pages go to the model 20 at a time and the GPU's memory is cleaned after each batch, so memory use stays predictable however long the document is.

A failed chunk is retried on its own, without touching the other nine. The ten GPUs each read about a tenth of the pages, so the reading itself takes about as much GPU time as one GPU reading everything. Each GPU's few minutes of start-up come on top, and in exchange the document is ready in 10 minutes instead of 38.

Where the work runs

Three pieces of software share the work. KubernetesKubernetesSoftware that runs programs across a fleet of machines and restarts them when they fail. With a cluster autoscaler, the fleet grows and shrinks with the work. runs programs across a fleet of machines, and a cluster autoscaler adds or removes machines as the work needs them; each running program is a podPodThe smallest unit Kubernetes runs: one or more containers scheduled together on one machine.. RayRayAn open-source framework for running Python work in parallel across many machines. spreads Python work across many machines and decides which machine runs which piece. KubeRayKubeRayThe open-source operator that runs Ray clusters on Kubernetes as Kubernetes resources. lets Kubernetes run Ray, so the Ray cluster is described in a file and Kubernetes keeps it in that shape.

The cluster keeps two kinds of workerWorkerA process or machine that takes tasks, runs them and reports back.. CPU workers are always on and do the cheap work: splitting PDFs, converting files and coordinating jobs. GPU workers do the reading, and they exist only while OCR work is waiting. That is autoscalingAutoscalingAdding workers when work is waiting and removing them when it isn't.: when a job arrives, GPU workers start; when the work is done, they scale back down to zero, and an idle hour costs nothing in GPUs.

What this means for your team

• Long documents are ready sooner. A 500-page filing is read in parallel, so it becomes searchable in minutes rather than most of an hour.
• A bad page range costs one chunk. Each chunk is retried on its own, so one difficult section doesn't send the whole document back to the start.
• Capacity follows the work. GPU workers start when documents arrive and stop when they are done, so a quiet night costs nothing in GPUs.

Ten terms to know

Part two builds on these ten. The other terms the post uses are defined in the glossary at the end.

Part two: the engineering detail. From here on, the post follows the pipeline itself: the problem, the chunking model, the Ray building blocks, running it all on Kubernetes with KubeRay, and what it measured.

Running document processing pipelines at scale sounds straightforward, until PDFs grow beyond a few hundred pages.

When dealing with 500+ page documents, powered by GPU-based OCR, layout analysis, table extraction, and document understanding models, things get complicated very quickly:

• Jobs slow down
• GPUs run out of memory
• Costs grow unpredictably
• Failures become hard to debug

This post dives into a production-grade, distributed document processing approach using Ray and KubeRay that is:

• Fast
• Memory-safe
• Cost-efficient
• Kubernetes-native

The problem: large PDFs don't play nice with GPUs

A traditional document pipeline often looks like this:

• One job
• One machine
• One giant PDF
• One very unhappy GPU

This approach breaks down badly when documents reach hundreds of pages.

What goes wrong?

• Latency skyrocketsProcessing large PDFs sequentially wastes cluster capacity and increases turnaround time.
• GPU OOMs become commonOCR and layout models consume significant memory. Feeding large documents into a single inference pipeline causes unpredictable memory spikes.
• Costs stay highKeeping GPU nodes always running, even when there's no work, burns money fast, especially for bursty workloads.

Clearly, a different approach is needed.

What this approach solves

The goal is simple (but not easy):

• Parallelism within a single document
• Strict limits on GPU memory usage
• Automatic retries and failure isolation
• Ability to scale GPU nodes up and down to zero
• Minimal operational overhead

This is where Ray fits naturally.

So, why Ray?

Ray is a distributed execution framework that makes parallel computing feel surprisingly normal.

It provides:

• Simple primitives for parallel tasks and actors
• Native GPU-aware scheduling
• Fine-grained control over concurrency
• A production-ready Jobs API for submission

Instead of relying on a long-running driverDriverThe program that starts Ray work and waits for it. A long-lived driver is one more thing that can fail. process, workloads can be submitted as Ray JobsRay JobOne submitted unit of work with its own entry point, logs and lifetime, run inside the cluster and cleaned up after., executed inside the cluster, and cleaned up once finished.

However, Ray alone is not the full answer. The way the work itself is structured matters just as much.

The core idea: chunk first, then parallelize

The biggest breakthrough is simple:

Never run OCR on a full 300 to 500+ page document at once.

Instead, adopt a chunk-based processing model.

Step 1: Split large PDFs into bounded chunks (CPU)

Each PDF is split into fixed-size page ranges (for example, 50 pages per chunk).

Why this matters:

• Memory usage becomes predictable
• Failures are isolated to individual chunks
• Chunks can be retried independently

Each chunk is uploaded to S3S3Amazon's file storage. Documents, chunks and outputs are kept there. and represented only by metadataMetadataSmall facts about the data, such as where a chunk is stored and which pages it holds, rather than the data itself.:

{
  "bucket": "pdf-chunks",
  "key": "doc-123/chunk-0",
  "page_range": "0-49",
  "chunk_index": 0
}

Step 2: Parallel GPU OCR per chunk

Each chunk is processed independently on a GPU worker.

To avoid OOMs:

• Chunks are micro-batchedMicro-batchA smaller group of pages fed to the model in one go inside a task. In our pipeline, 20 pages. internally (for example, 20 pages at a time)
• GPU memory is explicitly cleaned after each batch
• Task concurrencyConcurrencyHow many tasks run at the same moment. is capped based on cluster size

This transforms one massive, fragile job into many small, safe, parallel tasks.

Architecture overview

A production-grade flow looks like this:

• A job request is pushed into a queueQueueA to-do list of jobs. A scheduler takes jobs from it one by one. (database-backed)
• A lightweight scheduler consumes jobs from the queue
• A Ray Job is submitted via the Ray Jobs API
• Inside the Ray cluster, the job:
  • Splits the PDF into chunks
  • Dispatches parallel GPU OCR tasks
  • Uploads intermediate and final outputs to S3

Processing metadata and job status are persisted to RDSRDSAmazon's managed relational database. Job status and processing metadata are kept there., while documents remain in S3. Once the final merged output is generated, intermediate chunk artifacts can be safely cleaned up to avoid unnecessary storage overhead. This separation keeps metadata lightweight and storage scalable.

Ray fundamentals

Once documents are broken into independent chunks, Ray offers a small set of building blocks that map well to large document processing pipelines. Rather than thinking in terms of distributed systems complexity, it helps to think in terms of clear units of work, resource boundaries, and controlled concurrency.

Remote tasks: stateless, parallel work

Remote tasksRemote taskA Ray function call that runs somewhere in the cluster, keeps no state and returns a result. are Ray's simplest primitive:

"Run this function somewhere in the cluster."

Remote tasks are a great fit for work that:

• Has clear inputs and outputs
• Can be retried independently
• Doesn't need shared state

Typical examples in a document pipeline include:

• PDF splitting
• PDF conversion
• OCR / text extraction per chunk
• Merging results

Example GPU OCR task:

@ray.remote(num_gpus=1)
def extract_text(chunk):
    # download chunk from S3
    # run OCR model
    # upload output back to S3
    return {"bucket": "...", "key": "..."}

Declaring num_gpus=1 ensures:

• The task runs only on GPU workers
• GPU workloads are not overscheduled

This becomes critical for memory-heavy OCR models.

Actors: stateful, long-lived workers

Some workloads benefit from maintaining state across multiple operations.

Ray ActorsActorA long-lived Ray worker that keeps state between calls. are long-lived workers that:

• Maintain state across multiple calls
• Encapsulate complex or multi-step workflows
• Enable caching of intermediate or deterministic results

Actors are well suited for workflows such as:

• Layout or structure post-processing
• Stateful document-level aggregation
• Reusing intermediate results across tasks

A simple way to think about actors:

"A long-running worker that holds state, instead of starting from scratch every time."

Actors are optional in many pipelines, but become useful when workflows require coordination, state retention, or reuse beyond a single task invocation.

Resource-aware scheduling (CPU vs GPU)

Ray makes it possible to be explicit about where work should run.

In this pattern:

• GPU-heavy OCR → GPU workers
• CPU-only tasks → CPU workers

Example CPU-only task:

@ray.remote(resources={"cpu_worker": 1})
def convert_pdf(task):
    ...

Why this matters:

• CPU jobs don't steal GPU capacity
• GPUs stay reserved for expensive OCR work
• The cluster remains predictable under load

As document volume grows, separating CPU and GPU responsibilities becomes essential.

Ray Jobs API: production-friendly submission

A key design choice in production pipelines is avoiding a long-lived driver.

Instead, each document (or job request) is submitted as a Ray Job using the Ray Jobs API.

This enables:

• Clear job boundaries
• Better failure isolation
• Easier retries and monitoring
• No single "forever-running" process to manage

Conceptually:

• A lightweight scheduler consumes work from a queue
• Submits a Ray Job
• Ray executes everything inside the cluster
• When the job finishes, resources are released

This model fits naturally with batch-style document processing workloads.

Keeping Ray's object store small

One subtle but critical lesson:

Avoid passing large OCR outputs through Ray's object storeObject storeRay's shared memory for passing results between tasks. It lives in the cluster's own memory..

Instead:

• Each task uploads results directly to S3
• Tasks return only small metadata (bucket, key, page range)

Benefits:

• Ray's object store stays small and fast
• Less memory pressure on the cluster
• Better overall stability under heavy workloads

This design choice prevents many hard-to-debug production issues.

Running Ray on Kubernetes with KubeRay

Ray is powerful, but managing Ray clusters manually is not fun.

KubeRay solves this by making Ray a first-class Kubernetes resource.

Why KubeRay?

Ray clusters can be set up in multiple ways.

One option is to run Ray directly on virtual machinesVirtual machineA rented computer in a cloud provider's data centre, which you set up and manage yourself. across cloud providers. While this works, it typically requires manually handling things like:

• Head and worker node provisioning
• Network configuration and service discovery
• Autoscaling logic
• Node lifecycle management
• GPU device setup and scheduling

For production environments, this quickly becomes operationally heavy.

KubeRay takes a different approach by making Ray a first-class Kubernetes resource.

Instead of managing Ray clusters at the VM level, KubeRay allows Ray to be defined and operated declarativelyDeclarativeDescribing the state you want, such as at most ten GPU workers, and letting the system make it so. inside Kubernetes.

Key benefits include:

• Declarative RayCluster CRDCustom resource (CRD)A new kind of object taught to Kubernetes, such as a RayCluster, which it then stores and manages like its own.Ray clusters are defined as Kubernetes resources and reconciled by an operatorOperatorA program inside Kubernetes that keeps a custom resource in the state it declares, creating and repairing what it describes..
• Stable head service for job submissionThe Ray head nodeHead nodeThe node that coordinates a Ray cluster: it accepts jobs and decides where tasks run. is exposed via a consistent Kubernetes service, simplifying job submission and discovery.
• Built-in autoscalingRay workers scale automatically based on pending resource demands, without custom VM logic.
• No manual VM managementNode provisioning, networking, and lifecycle are handled by Kubernetes and the cluster autoscalerCluster autoscalerThe Kubernetes component that adds machines when pods can't be placed and removes machines that sit empty..

This makes KubeRay a much more natural fit for running Ray in production Kubernetes environments, especially for GPU workloads and autoscaling, while remaining portable across clouds and environments through TerraformTerraformA tool that creates cloud infrastructure, such as networks and clusters, from files kept with the code. and HelmHelmA package manager for Kubernetes: one chart describes everything an application needs, with values per environment.-based deployments.

CPU and GPU worker separation

Two distinct worker pools help keep responsibilities clear and costs under control.

CPU workers

• Always on
• Handle splitting, conversion, orchestration
• Advertise a custom cpu_worker resource

GPU workers

• Handle OCR, layout analysis, table extraction
• Request nvidia.com/gpu
• Scale to zero when idle

This separation keeps GPU nodes focused only on expensive workloads.

GPU autoscaling and cost control

GPU workloads are spiky by nature.

With KubeRay:

• GPU workers have minReplicas = 0
• Ray autoscaling is enabled

What this means:

• No OCR jobs → zero GPU pods
• OCR jobs arrive → GPU pods spin up automatically
• Work finishes → GPU pods scale back down

GPU cost now matches real workload demand.

NVIDIA device plugin (GPU discovery)

To make GPUs schedulable in Kubernetes, the NVIDIA device pluginDevice pluginA Kubernetes add-on that tells the scheduler which nodes have special hardware, such as GPUs, and how many. is installed.

This exposes GPUs as:

nvidia.com/gpu

Without this, GPU workloads cannot be placed correctly.

Avoiding cold starts: preloading models

OCR models are large and slow to download.

To avoid delays:

• GPU workers preload models on startup
• Models are pulled from S3 into the correct cache paths
• Tasks start immediately when scheduled

This significantly reduces first-job latency and failure rates.

Performance results and constraints

Measured results

A 463-page document processed sequentially on a single GPU took 38 minutes to complete.

When the same document was split into chunks (50 pages per chunk) and processed in parallel across 10 GPU machines, it completed in 10 minutes, including GPU spin-up time.

This represents a ~3.8× reduction in overall processing time, while keeping GPU usage tightly scoped to the duration of actual work.

Around 4 to 5 minutes of the total runtime is attributable to GPU startup overhead, including instance spin-up, Kubernetes pod scheduling, and NVIDIA driver initialization. This overhead is expected when scaling GPU nodes on demand, and GPU nodes are scaled back down immediately after processing completes.

From a cost perspective, this approach is comparable to running a single GPU for a longer duration. Instead of keeping one GPU busy for 38 minutes, multiple GPUs are used briefly in parallel and then released, resulting in similar total GPU time consumption with significantly lower end-to-end latencyLatencyHow long one piece of work takes from start to finish..

Practical constraints

• GPU quotas are finiteCloud GPU limits (for example, A10 GPUs per region or account) naturally cap how much parallelism can be applied without requesting quotaGPU quotaThe number of GPUs a cloud account may run in a region at once. Raising it means asking the provider. increases.
• Concurrency is intentionally controlledAlthough Ray supports multiple tasks per node, chunk concurrency is tuned per machine to align with document processing library constraints and avoid GPU memory pressure.

Final thoughts

Scaling document processing isn't about throwing bigger machines at the problem. It's about structuring the work correctly.

This architecture succeeds because it leans on a few proven ideas:

• Breaking work into small, bounded chunks instead of oversized batch jobs
• Letting parallelism follow resource constraints, not overwhelm them
• Treating GPUs as elastic resources, not always-on infrastructure
• Keeping infrastructure declarative and portable across environments

Ray handles execution and coordination.
KubeRay handles operations and scaling.

Together, they transform large-document OCR from a brittle pipeline into a resilient, production-ready system.

Want to go deeper?

If you'd like to explore the building blocks used in this setup in more detail, the following resources are helpful starting points:

• Ray on Kubernetes (Official Documentation)Overview of running Ray on Kubernetes, including KubeRay concepts and cluster setup
• NVIDIA Kubernetes Device PluginRequired for exposing GPUs (nvidia.com/gpu) to Kubernetes workloads
• AWS EC2 GPU Quotas and LimitsUnderstanding regional GPU limits and requesting quota increases

GPU autoscaling beyond Ray: event-driven scaling with KEDA

GPU workloads don't have to run all the time. Instead, GPU capacity can be brought up only when real work arrives and scaled back down when it's no longer needed.

One practical example is serving embedding models on Kubernetes, where KEDA is used to scale GPU-backed pods on demand based on incoming job requests or custom metrics. This approach avoids keeping expensive GPU nodes idle and helps align infrastructure costs directly with actual usage.

If you're interested in how event-driven autoscaling with KEDA can be used to manage GPU workloads efficiently, this post dives deeper into that pattern:

Got questions?

If you have thoughts, questions, or want to share how you're using Ray in your own stack, don't hesitate to reach out! Connect with me on LinkedIn.

Glossary

The ten core terms are defined in part one. These are the rest.

FAQs

What is Ray?

Ray is an open-source framework for running Python work in parallel across many machines. It offers remote tasks for stateless work, actors for work that keeps state, scheduling that knows which machines have GPUs, and a Jobs API for submitting work to a cluster. Needl.ai uses it to read the chunks of a long PDF on many GPUs at once.

What is KubeRay?

KubeRay is an open-source operator that runs Ray clusters on Kubernetes. A Ray cluster is declared as a RayCluster resource, and the operator creates and repairs the head and worker pods to match it. Workers scale on pending resource demand, and the head node gets a stable service for job submission.

Why split a PDF into chunks before OCR?

Bounded chunks make memory use predictable, keep a failure to one chunk and let many GPUs work on one document at the same time. In our pipeline each chunk is 50 pages, read 20 pages at a time, and a 463-page document went from 38 minutes on one GPU to 10 minutes on ten.

Why not use a bigger machine?

A bigger machine doesn't change the shape of the work: one long run still reads page after page, and its memory still spikes unpredictably as pages pass through the model. Splitting the work into bounded chunks lets many GPUs share one document and keeps each run within a predictable amount of memory.

How do GPU workers scale to zero?

The GPU worker group has a minimum of zero replicas and Ray autoscaling is on. With no OCR jobs there are no GPU pods; when jobs arrive, GPU pods start, and they scale back down when the work finishes. Starting a GPU worker takes around 4 to 5 minutes, covering instance spin-up, pod scheduling and NVIDIA driver initialization.

Does reading on ten GPUs cost more than on one?

Somewhat. The reading itself takes about the same GPU time either way, roughly 38 GPU-minutes in total, but each of the ten GPUs also spends 4 to 5 minutes starting up, so the total GPU time is higher. In exchange, the document is ready in 10 minutes rather than 38 (Figure 12).

How does this relate to Needl.ai's serverless extraction?

This pipeline reads pages with OCR models on GPUs that we run ourselves. Where a vision model reads the page through an API instead, the work is mostly waiting rather than computing, so that path has since moved to AWS Lambda. The third post in this series covers that move.