Batch Inference at Scale with Ray

Batch inference is the cost and throughput play: run thousands of LLM calls reliably, cheaply, and observably — instead of paying for per-interactive-request spikes.

When you need batch inference

Interactive chat needs per-request latency. But a lot of LLM work is offline: scoring a corpus, summarizing last week’s logs, classifying a million tickets, regenerating evals, or backfilling a vector store. For those, latency per call isn’t the concern — total time, cost, and reliability are. A crash ten minutes into a marathon should not restart from scratch.

Ray: the distributed compute substrate

Ray is an open-source framework for distributing Python. Its core abstractions are tasks (an arbitrary Python function executed on a cluster) and actors (long-lived objects that keep state between calls). It is the right tool when you want to map over a large number of LLM calls with retries, checkpointing, and a rich graph.

Install: pip install "ray[default]" then ray start --head for a local cluster.

The batch worker: an actor

A @ray.remote actor wraps the model client. Each remote call runs on the cluster and keeps a single client, avoiding re-initialization overhead per request.

import ray
from google import genai

@ray.remote
class Generator:
    def __init__(self, model="gemini-2.5-pro"):
        self.client = genai.Client()
        self.model = model

    def generate(self, prompt: str) -> str:
        resp = self.client.models.generate_content(model=self.model, contents=prompt)
        return resp.text

Submitting a big batch

Fan out a list of tasks to the actor pool and gather them:

actors = [Generator.remote() for _ in range(NUM_WORKERS)]
futures = [actors[i % NUM_WORKERS].generate.remote(p) for i, p in enumerate(prompts)]
results = ray.get(futures)   # blocks until all done

for doc, text in zip(documents, results):
    doc.generated = text

ray.get waits for all, and because each task is separate, Ray resubmits a failed task without rerunning the whole batch.

Reliability: retries and rate limits

LLM APIs apply rate limits (HTTP 429) and transient errors. Wrap the call in retry logic with exponential backoff:

from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type

class RateLimitError(Exception):
    pass

@retry(reraise=True, stop=stop_after_attempt(5),
       wait=wait_exponential(max=60),
       retry=retry_if_exception_type(RateLimitError))
def generate_with_retry(actor, prompt):
    return actor.generate.remote(prompt)

This is the ref’s “make cost visible and operations professional”: a named job handled in backoff on rate-limit and resumable on failure.

Observability and tracking

Wrap each generation with a trace_id and log a structured row (model, tokens, prompt ID, latency). Ray ships a dashboard (default :8265). The ref’s theme — “You cannot manage what you cannot measure” — collect throughput, failure rate, and a token-to-cost tally per run and store them for post-hoc audits.

Putting It All Together

import ray
from my_executor import Generator, build_prompts

ray.init()
gen = Generator.remote()

prompts = build_prompts()                     # e.g. 5,000 tagged sentences
futures = [gen.generate.remote(p) for p in prompts]
outputs = ray.get(futures)

for rec, text in zip(records, outputs):
    log_row({"model": model, "tokens": count(text), "cost_usd": estimate_cost(text)})
    store(rec, text)

Conclusion & Next Steps

Batch inference with Ray turns heavy offline LLM work into a resumable, retried, observable pipeline. Next: add checkpoint progress so a failed run continues where it left off, add back-pressure for very large sets, and attach the trace/cost pipeline from the observability draft. — make cost visible and operations professional.

References / Sources

• Ray docs. https://docs.ray.io
• tenacity. https://github.com/jd/tenacity
• Gemini text generation. https://ai.google.dev/gemini-api/docs/text-generation