Ray for Data Engineering: Distributed Python, Ray Data & Batch Inference at Scale

ray for data engineering is the Python-first distributed compute layer that senior engineers reach for when the workload is more Python than SQL, more model than aggregation, and more heterogeneous than a JVM cluster wants to be — the last-mile ML preprocessing, the terabyte-scale batch inference job, the embeddings pipeline that Spark can technically run but does grudgingly through a serialization boundary that taxes every UDF. Every distributed system forces a language decision, and for a decade that decision was "learn the JVM's data model or pay the PySpark UDF tax on every row"; Ray's bet is that a distributed runtime designed around Python objects, Python functions, and Python classes — not around a query planner that treats Python as a foreign function — is the right substrate for the half of modern data engineering that is really machine-learning plumbing.

This guide is the senior-DE walkthrough you wished existed the first time an interviewer asked "when would you pick Ray over Spark for a data pipeline," or "explain how Ray's object store gives you zero-copy sharing between tasks," or "walk me through a GPU batch-inference job in Ray Data and how you'd stop it from OOMing the cluster." It works through the five things every senior engineer must know: why Ray exists and where it fits versus Spark and Dask; the primitives of ray core — tasks, actors, futures, and the distributed object store; ray data as a streaming, lazy, block-based dataset engine built around map_batches; batch inference at scale with stateful GPU actors, autoscaling, and backpressure; and the operational reality of ray clusters — KubeRay, fault tolerance, and the workloads where Ray is the wrong tool. Each section pairs a teaching block with a Solution-Tail interview answer — runnable code, a step-by-step trace, an output table, then a concept-by-concept breakdown of why it works.

When you want hands-on reps immediately after reading, drill the data-processing practice library →, rehearse on the ETL practice library →, and sharpen the throughput axis with the optimization practice library →.

---

On this page

• Why Ray exists and where it fits for data engineers
• Ray Core — tasks, actors, the object store, and scheduling
• Ray Data — streaming datasets, lazy execution, and map_batches
• Batch inference at scale — GPU actors and autoscaling
• Ray clusters, ops, and interview signals
• Cheat sheet — Ray for data engineering recipes
• Frequently asked questions
• Practice on PipeCode

---

1. Why Ray exists and where it fits for data engineers

The Python-first distributed gap Spark leaves — and the workloads that fall into it

The one-sentence invariant: Ray is a distributed runtime whose unit of parallelism is a Python function (task) or a Python class instance (actor), sharing data through a shared-memory object store rather than a shuffle-oriented query engine — which makes it the right substrate for the Python-heavy, model-heavy, heterogeneous-hardware half of data engineering that Spark runs only through a costly JVM-to-Python serialization boundary. Spark is a magnificent SQL-and-DataFrame engine: for SELECT ... GROUP BY ... JOIN over columnar data at petabyte scale, nothing beats it. But the moment your pipeline stops being relational algebra and starts being "load a PyTorch model on each GPU, run inference over 400 million images, and write the embeddings back," Spark's model — Python code marshalled row-by-row across the JVM boundary, no first-class GPU scheduling, no stateful workers — starts fighting you. Ray was built at Berkeley's RISELab specifically for that class of work, and ray for data engineering is the practice of using it where it wins.

Where Ray fits — the workloads that fall into the Spark gap.

• Batch inference over large datasets. Run a model (LLM, vision, embeddings) over hundreds of millions of records. Needs GPU scheduling, a model loaded once per worker, and streaming so the dataset never has to fit in memory. This is Ray's flagship data-engineering use case.
• Last-mile ML preprocessing. Tokenization, image decoding, feature extraction — heavy Python UDFs that would pay the PySpark serialization tax on every batch. Ray runs them as native Python with zero-copy Arrow batches.
• Heterogeneous CPU + GPU pipelines. A pipeline where one stage is CPU-bound (decode) and the next is GPU-bound (infer). Ray schedules each stage on the right hardware in one job; Spark treats GPUs as an afterthought.
• Distributed Python that isn't a DataFrame. Hyperparameter sweeps, reinforcement learning, simulation, custom parallel algorithms. Anything that is "run this Python across a cluster" but not "run this SQL."

What Ray is not trying to be.

• A SQL warehouse. Ray Data is not Snowflake or BigQuery. If your workload is GROUP BY over a governed warehouse table, use the warehouse.
• A Spark replacement for relational ETL. Big shuffle-heavy joins with a mature Spark shop already running them? Ray does not obviously win. Ray wins on the Python/ML axis, not the relational axis.
• A streaming-first engine. Ray Data streams execution (it pipelines blocks), but it is a batch engine. For true event-time streaming with exactly-once sinks, that is Flink / Spark Structured Streaming / Kafka Streams territory.

The three engines senior engineers compare — Ray, Spark, Dask.

• Spark. JVM-native, SQL/DataFrame-first, best-in-class shuffle and query optimization, huge ecosystem. Python is a guest via Py4J; UDFs and ML serving pay a serialization boundary. Pick for relational ETL at scale.
• Dask. Python-native like Ray, DataFrame-and-array-first, familiar pandas/NumPy APIs. Excellent for scaling existing pandas code; weaker on GPU-first ML serving and stateful actors than Ray. Pick for "scale my pandas."
• Ray. Python-native, task/actor-first, first-class GPU scheduling, stateful actors, an object store for zero-copy sharing, and a data library (Ray Data) plus an ML stack (Ray Train / Serve / Tune). Pick for ML-adjacent distributed Python and batch inference.

What interviewers listen for.

• Do you frame Ray as "distributed Python built on tasks and actors," not as "a Spark competitor for everything"? — senior signal.
• Do you name the Spark UDF / JVM serialization boundary as the specific reason Python-heavy workloads underperform on Spark? — required answer.
• Do you name batch inference with GPU actors as Ray's flagship data-engineering use case? — senior signal.
• Do you name a workload where you would not pick Ray (relational ETL in a mature Spark shop, single-node pandas)? — senior signal, shows you know the envelope.

Worked example — the three-primitive mental model

Detailed explanation. Every Ray program is built from exactly three primitives, and being able to name them and say what each is for is the fastest way to sound fluent in an interview. Everything else — Ray Data, Ray Train, Ray Serve — is a library built on top of these three. Walk through the mental model before touching any higher-level API.

• Task — a stateless Python function you decorate with @ray.remote. Calling f.remote(x) schedules it on the cluster and returns a future immediately. Use tasks for embarrassingly parallel, stateless work.
• Actor — a Python class you decorate with @ray.remote. A.remote() creates one instance (a stateful worker process) on some node; a.method.remote() runs a method on it. Use actors when work needs to hold state — a loaded model, a running counter, a connection pool.
• Object (ObjectRef) — the result of a .remote() call or a ray.put(). It is a future and a handle into the distributed object store. ray.get(ref) blocks and materialises the value; the store gives zero-copy reads to any task on the same node.

Question. Using only the three primitives, sketch how a "decode images then run a model" pipeline maps onto Ray, naming which primitive each stage uses and why.

Input.

Stage	Nature	Primitive	Why
Read image bytes from S3	stateless, parallel	task	no state; fan out one task per shard
Decode + resize	stateless, CPU-bound	task	pure function over bytes
Load model + infer	stateful, GPU-bound	actor	model must load once, not per call
Share the decoded batch	data movement	object (ObjectRef)	zero-copy hand-off to the GPU actor

Code.

import ray

ray.init()  # connect to (or start) a local cluster

# --- Task: stateless, parallel ---
@ray.remote
def decode_and_resize(image_bytes: bytes, size: int = 224):
    # pure function: bytes in, tensor out; no state kept between calls
    import numpy as np
    # (illustrative decode; a real impl would use PIL/opencv)
    arr = np.frombuffer(image_bytes, dtype=np.uint8)
    return arr[:size * size].reshape(-1)  # toy "resized" vector

# --- Actor: stateful, holds the model ---
@ray.remote(num_gpus=1)
class Classifier:
    def __init__(self, weights_uri: str):
        # loads ONCE when the actor is created, not on every call
        self.model = load_model(weights_uri)  # your framework here

    def predict(self, batch):
        return self.model(batch)

def load_model(uri):        # stand-in for torch.load / from_pretrained
    return lambda b: [0.0] * len(b)

# Driver wires the primitives together
raw = [b"...", b"...", b"..."]                    # three image payloads
decoded_refs = [decode_and_resize.remote(b) for b in raw]   # 3 futures
clf = Classifier.remote("s3://models/resnet.pt")            # 1 GPU actor
preds = ray.get(clf.predict.remote(ray.get(decoded_refs)))  # gather + infer
print(preds)

Step-by-step explanation.

1. decode_and_resize is a task because decoding is stateless — each call is a pure function of its input bytes. decode_and_resize.remote(b) returns an ObjectRef immediately; the three calls run in parallel across the cluster's CPUs.
2. Classifier is an actor because the model must be loaded exactly once. If we made inference a task, every call would reload the weights from S3 — the single most common Ray anti-pattern. The num_gpus=1 reservation tells the scheduler to place this actor on a node with a free GPU.
3. The object store carries the decoded batches. decoded_refs are handles into the store; passing them to the actor lets Ray move only what is needed, and any task co-located on the same node reads them zero-copy.
4. ray.get is the single blocking point — the driver stays asynchronous until it actually needs the values. In production you would stream with ray.wait (section 2) rather than one big ray.get.
5. This three-primitive decomposition is the answer to "how do I think about a workload in Ray": classify each stage as stateless (task), stateful (actor), or data (object), and the design falls out.

Output.

Primitive	Created by	Blocking?	Data-engineering use
Task	@ray.remote on a function	no (returns ref)	parallel decode, parse, transform
Actor	@ray.remote on a class	no (returns handle)	model serving, stateful aggregation
Object	.remote() result / ray.put()	ray.get blocks	zero-copy batch hand-off

Rule of thumb. Classify every stage as stateless → task, stateful → actor, data → object before writing code. If you find yourself reloading a model inside a task, that stage wanted to be an actor.

Worked example — the Ray-vs-Spark fit decision

Detailed explanation. The most common Ray interview question is not "how does Ray work" but "when would you use it instead of Spark." The senior answer is a decision, not a preference — it walks a short set of axes and lands on a tool. Codify the axes so any scenario resolves quickly.

• Is the heavy work relational? Big joins, group-bys, SQL over columnar data → Spark. Python UDFs / model calls dominate → Ray.
• Is there GPU work? First-class GPU scheduling and model-once loading → Ray. CPU-only aggregation → Spark is fine.
• Does a stage need to hold state? A loaded model, a warm connection, an accumulator → actors → Ray. Stateless map/reduce → either.
• What does the team already run? A mature Spark platform doing relational ETL → don't rip it out for relational work. A Python/ML team with no JVM appetite → Ray lowers the tax.

Question. Given three workloads, apply the axes and pick the engine, with a one-line justification each.

Input.

Workload	Relational?	GPU?	Stateful stage?
Nightly fact-table join + aggregate (2 TB)	yes	no	no
Embed 300M product descriptions with an LLM	no	yes	yes (model)
Scale an existing pandas feature script 20×	partly	no	no

Code.

def pick_engine(relational: bool, needs_gpu: bool,
                stateful_stage: bool, has_spark_platform: bool) -> str:
    """Illustrative decision helper for engine selection."""
    if needs_gpu or stateful_stage:
        return "Ray"                      # GPU + model-once + actors = Ray's lane
    if relational and has_spark_platform:
        return "Spark"                    # relational ETL on an existing platform
    if relational:
        return "Spark (or Ray Data if Python-first)"
    return "Dask or Ray"                  # scale-my-pandas territory

print(pick_engine(True,  False, False, True))   # nightly join
# -> 'Spark'
print(pick_engine(False, True,  True,  True))   # LLM embeddings
# -> 'Ray'
print(pick_engine(True,  False, False, False))  # scale pandas, no Spark platform
# -> 'Spark (or Ray Data if Python-first)'

Step-by-step explanation.

1. The nightly join is pure relational algebra over 2 TB with an existing Spark platform — the axes short-circuit to Spark. Running it on Ray would mean rebuilding shuffle machinery Spark already has.
2. The LLM embedding job trips two Ray axes at once: it needs GPUs and it needs a model loaded once (stateful). Ray's GPU actors + Ray Data streaming are purpose-built for it; Spark would reload the model per task and marshal tensors across the JVM boundary.
3. The pandas-scaling job is the ambiguous one. If the team has no Spark platform and the logic is Python-heavy, Ray Data or Dask both fit; if there's a Spark platform and the logic is relational, Spark wins. "It depends on the existing platform" is the honest senior answer.
4. Naming the existing platform axis is what separates a senior answer from a benchmark-quoting one — the right engine is partly an organizational decision, not just a technical one.
5. The helper is illustrative, not gospel: real decisions weigh data size, latency SLOs, and team skills too. But walking axes out loud is exactly the structure interviewers reward.

Output.

Workload	Engine	One-line justification
Nightly fact-table join	Spark	relational shuffle at scale, platform exists
300M LLM embeddings	Ray	GPU scheduling + model-once actors + streaming
Scale pandas 20×	Dask or Ray	Python-native scaling; no JVM tax

Rule of thumb. Ray wins when the heavy work is Python or GPU or stateful; Spark wins when the heavy work is relational shuffle on an existing platform. Say "it depends on relational-vs-Python and on the existing platform" — never "Ray is faster."

Data engineering interview question on choosing Ray

A senior interviewer often opens with: "Your team runs a nightly Spark job that loads a PyTorch model in a mapPartitions UDF and scores 500 million rows. It takes 9 hours, the GPUs sit idle 60% of the time, and the model reloads on every partition. Walk me through why this is slow, whether Ray is the right move, and how you'd redesign it."

Solution Using a Ray task/actor/dataset fit map to redesign the scoring job

# Diagnosis + redesign sketch — why Spark struggles and how Ray fixes it.
#
# PROBLEM (Spark today):
#   - PyTorch model reloaded per partition (no stateful worker)
#   - tensors marshalled JVM <-> Python per batch (serialization tax)
#   - GPUs not first-class; scheduler over-provisions CPU executors
#   - result: 9h wall clock, 60% GPU idle
#
# REDESIGN (Ray Data + GPU actors): model loads ONCE per actor,
# Arrow batches stay in Python (zero JVM boundary), GPUs scheduled explicitly.

import ray
import numpy as np

ray.init()

class Scorer:
    def __init__(self, weights_uri: str):
        # Loads the model ONCE per actor process, not per batch.
        self.model = load_torch_model(weights_uri)

    def __call__(self, batch: dict) -> dict:
        # batch is a dict of numpy arrays (zero-copy from the object store)
        features = batch["features"]
        batch["score"] = self.model(features)
        return batch

def load_torch_model(uri):            # stand-in for torch.load(...).cuda().eval()
    return lambda x: np.asarray(x).sum(axis=1)

ds = ray.data.read_parquet("s3://warehouse/scoring_input/")   # 500M rows, lazy

scored = ds.map_batches(
    Scorer,
    fn_constructor_args=("s3://models/scorer.pt",),
    concurrency=8,          # 8 GPU actors in a pool
    num_gpus=1,             # each actor reserves one GPU
    batch_size=1024,        # rows per model call
    batch_format="numpy",
)

scored.write_parquet("s3://warehouse/scoring_output/")        # triggers execution

Step-by-step trace.

Step	Spark today	Ray redesign
Model load	per partition (reloaded ~2000×)	once per actor (8×)
Data movement	JVM ↔ Python per batch	Arrow batches stay in Python, zero-copy
GPU scheduling	best-effort, 60% idle	explicit num_gpus=1 per actor, kept fed
Parallelism unit	partition + UDF	streaming blocks → actor pool
Execution	materialise partitions	streaming executor, bounded memory
Wall clock	~9 h	~1–2 h (GPUs saturated)

The 500-million-row Parquet source is read lazily as Ray Data blocks; map_batches fans those blocks across a pool of eight GPU actors, each of which loaded the model exactly once at construction. Because batches are Arrow/NumPy that stay inside Python, there is no JVM boundary and no per-row serialization. The write triggers the streaming execution, so the cluster never has to hold all 500M scored rows in memory at once.

Output:

Metric	Spark (before)	Ray (after)
Model loads	~2000	8
GPU utilisation	~40%	~90%
Peak memory	full-partition	bounded (streaming)
Wall clock	9 h	1–2 h
Cost driver	idle GPUs + serialization	saturated GPUs

Why this works — concept by concept:

• Model-once actor — a Ray actor's __init__ runs a single time per worker process, so the model lives in GPU memory across thousands of batches. Replacing the per-partition reload is the single biggest win; it converts 2000 model loads into 8.
• No JVM boundary — Ray Data hands NumPy/Arrow batches to the actor inside one Python runtime. There is no Py4J marshalling, so the "serialization tax" that dominates PySpark UDFs disappears.
• Explicit GPU scheduling — num_gpus=1 makes the GPU a first-class schedulable resource; the actor pool is sized to keep every GPU busy rather than the scheduler guessing with CPU executors.
• Streaming execution — reading lazily and writing at the end lets Ray Data pipeline blocks through the actor pool, so peak memory is bounded by in-flight blocks, not by the 500M-row dataset size.
• Cost — O(rows) compute is unavoidable, but the constant factor collapses: 8 model loads instead of 2000, GPUs at ~90% instead of ~40%, and no per-batch serialization. Net effect is roughly a 5–8× wall-clock reduction at the same or lower GPU-hour cost.

Data Processing
Topic — data-processing
Distributed data-processing problems

Practice →

ETL Topic — etl ETL problems on large-scale scoring pipelines

Practice →

---

2. Ray Core — tasks, actors, the object store, and scheduling

Tasks are stateless futures, actors are stateful workers, and the object store is how they share data without copying

The mental model in one line: ray core gives you three composable primitives — ray tasks (@ray.remote functions that return ObjectRef futures), ray actors (@ray.remote classes that are long-lived stateful processes), and the shared-memory object store that holds every result and gives zero-copy reads to co-located workers — and a resource-aware scheduler that places each task or actor on a node with the CPUs, GPUs, or custom resources it declared it needs. Everything higher up the Ray stack, including Ray Data, is a library expressed in these primitives; understanding them is what lets you reason about performance, memory, and failure.

Tasks — stateless parallelism.

• Definition. Decorate a function with @ray.remote; call it with .remote(args). It runs on some worker in the cluster and returns an ObjectRef immediately — a future for the eventual result.
• Resolution. ray.get(ref) blocks until the value is ready and materialises it in the driver. ray.wait(refs, num_returns=k) returns as soon as k are done — the primitive for streaming/pipelined patterns.
• Resources. @ray.remote(num_cpus=2, num_gpus=0) declares what one invocation needs; the scheduler only places it where those resources are free.
• When to use. Embarrassingly parallel, stateless work: parse N files, transform N shards, run N simulations. No shared mutable state.

Actors — stateful parallelism.

• Definition. Decorate a class with @ray.remote; A.remote(args) creates one instance as a dedicated process on some node. a.method.remote(x) runs a method and returns an ObjectRef.
• State. The instance persists across calls — a loaded model, an in-memory counter, a DB connection pool. This is the primitive that batch inference and stateful aggregation are built on.
• Concurrency. By default an actor processes one method call at a time (its methods are serialized), which makes shared state safe without locks. Use max_concurrency or async methods for concurrent handling.
• When to use. Anything that must hold state between calls, or anything expensive to set up that you want to reuse.

The object store — zero-copy shared memory.

• What it is. Each node runs a shared-memory object store (historically "Plasma"). Every ObjectRef value lives there; tasks and actors on the same node read it without copying (a memory-mapped view).
• ray.put(obj). Explicitly place a large object into the store once and pass the ObjectRef to many tasks — the way to broadcast a big config/array without re-serializing it per call.
• Spilling. When the store fills, Ray spills objects to local disk (and back) automatically, so you are bounded by disk, not RAM, for intermediate results.
• Ownership + lineage. The store tracks which task produced each object so it can reconstruct lost objects on failure (section 5).

Scheduling — resource-aware placement.

• Logical resources. num_cpus, num_gpus, memory, and arbitrary resources={"tpu": 1} are logical accounting units the scheduler uses; they are reservations, not hard OS limits.
• Locality. The scheduler prefers to run a task on the node that already holds its input objects, avoiding network transfer — a big deal for large batches.
• Autoscaling hook. Pending tasks/actors that can't be placed signal the autoscaler to add nodes (section 5).

Common interview probes on Ray Core.

• "Task or actor?" — stateless → task, must-hold-state → actor.
• "How do two tasks share a big array without copying?" — ray.put once, pass the ObjectRef; co-located reads are zero-copy.
• "What does .remote() return?" — an ObjectRef (a future), not the value; ray.get materialises it.
• "How do you avoid blocking on the slowest task?" — ray.wait for streaming completion, not one big ray.get.

Worked example — parallel map with tasks and streaming with ray.wait

Detailed explanation. The canonical Ray Core pattern is "fan out N tasks, consume results as they finish." Beginners write ray.get([f.remote(x) for x in xs]), which blocks until all are done and buffers every result. The senior pattern uses ray.wait to process results as they complete, keeping memory flat and overlapping compute with consumption. Walk through both.

• Naive. ray.get(list_of_refs) — simple, but waits for the slowest task and holds all results.
• Streaming. Loop on ray.wait(refs, num_returns=1) — handle each result the instant it's ready; ideal when downstream can consume incrementally.

Question. Fan out a CPU-bound transform over 1000 shards and write each result as soon as it completes, never holding more than a handful in memory.

Input.

Parameter	Value
Shards	1000
Task	@ray.remote def process(shard)
Completion pattern	ray.wait(..., num_returns=1)
Memory goal	O(in-flight), not O(1000)

Code.

import ray
ray.init()

@ray.remote(num_cpus=1)
def process(shard_id: int) -> dict:
    # simulate CPU-bound work over one shard
    total = sum(i * i for i in range(shard_id % 100 + 1))
    return {"shard": shard_id, "checksum": total}

def write_result(r: dict) -> None:
    pass  # append to a sink: file, DB, queue

# 1. Fan out 1000 tasks — all return refs immediately (non-blocking)
pending = [process.remote(i) for i in range(1000)]

# 2. Consume as they finish — never hold all 1000 results at once
done_count = 0
while pending:
    ready, pending = ray.wait(pending, num_returns=1, timeout=None)
    for ref in ready:
        write_result(ray.get(ref))   # materialise ONE completed result
        done_count += 1

print(f"processed {done_count} shards, streaming")

Step-by-step explanation.

1. The list comprehension dispatches all 1000 tasks and returns 1000 ObjectRefs without blocking — the cluster starts working immediately while the driver moves on.
2. ray.wait(pending, num_returns=1) blocks only until one task finishes, returning (ready, still_pending). This is the streaming primitive: you react to the first result at ~the latency of the fastest task, not the slowest.
3. Inside the loop, ray.get(ref) materialises exactly one completed result, which is written and discarded. Peak memory is bounded by how many results are ready-but-unconsumed, not by the 1000-task fan-out.
4. Reassigning pending = still_pending shrinks the wait set each iteration, so the loop terminates when all tasks are done.
5. Contrast with ray.get(pending): correct, but it waits for the slowest of 1000 tasks before you can write anything, and it holds all 1000 dicts in driver memory. For large per-task outputs, that difference is the difference between flat and blown memory.

Output.

Pattern	First result at	Peak driver memory	Use when
ray.get(all_refs)	slowest task	O(N results)	small N, need all together
ray.wait loop	fastest task	O(in-flight)	large N, incremental sink

Rule of thumb. Default to ray.wait for large fan-outs with an incremental sink; reserve one-shot ray.get(list) for small result sets you genuinely need all at once. Streaming completion is how you keep driver memory flat.

Worked example — a stateful actor accumulator and ray.put broadcast

Detailed explanation. Two Ray Core patterns show up constantly in data pipelines: a stateful actor that accumulates a running result across many task outputs, and ray.put to broadcast a large read-only object (a lookup table, a config, an embedding matrix) to many tasks without re-serializing it per call. Walk through both together.

• Actor accumulator. One actor holds the running aggregate; tasks send their partial results to it. The actor's serialized method calls make the update race-free without locks.
• ray.put broadcast. Put the big lookup object into the store once; pass its single ObjectRef to every task. Co-located tasks read it zero-copy; remote tasks fetch it once per node.

Question. Compute a global word-frequency count over many shards using an actor accumulator, while broadcasting a large stop-word set to every mapper via ray.put.

Input.

Component	Mechanism
Stop-word set (large)	ray.put once, share the ObjectRef
Per-shard counting	@ray.remote task
Global merge	actor with a running dict
Final read	actor.result.remote() → ray.get

Code.

import ray
from collections import Counter
ray.init()

# 1. Broadcast a large read-only object ONCE via the object store
STOPWORDS = set(f"word{i}" for i in range(50_000))
stop_ref = ray.put(STOPWORDS)          # one copy in the store; share the ref

# 2. Stateless mapper task — receives the ObjectRef, not a fresh copy
@ray.remote
def count_shard(shard: list[str], stop_ref) -> dict:
    stop = stop_ref                    # Ray auto-resolves the ObjectRef arg
    return dict(Counter(w for w in shard if w not in stop))

# 3. Stateful actor that merges partial counts safely
@ray.remote
class Accumulator:
    def __init__(self):
        self.total = Counter()
    def merge(self, partial: dict) -> None:
        self.total.update(partial)     # serialized calls => no race
    def result(self) -> dict:
        return dict(self.total)

acc = Accumulator.remote()
shards = [["word1", "hello", "word2", "ray"], ["hello", "ray", "data"]]

# Fan out mappers, pipe each partial into the accumulator
merge_refs = [acc.merge.remote(count_shard.remote(s, stop_ref)) for s in shards]
ray.get(merge_refs)                    # wait for all merges to land
print(ray.get(acc.result.remote()))    # -> {'hello': 2, 'ray': 2, 'data': 1}

Step-by-step explanation.

1. ray.put(STOPWORDS) serializes the 50k-word set into the object store exactly once and returns stop_ref. Passing stop_ref (not STOPWORDS) to each task means Ray does not re-serialize the set per call — critical when the broadcast object is large.
2. When count_shard receives stop_ref as an argument, Ray transparently resolves it to the underlying value; on the node that already holds it, that read is zero-copy shared memory.
3. Accumulator is an actor because it holds mutable state (self.total). Its methods are executed one at a time, so concurrent merge calls from many mappers cannot corrupt the counter — no explicit lock needed.
4. count_shard.remote(s, stop_ref) returns a ref that is passed directly into acc.merge.remote(...). Ray sees the dependency and only runs the merge after the count finishes — automatic task chaining without manual ray.get in between.
5. The final ray.get(acc.result.remote()) is the one place the driver blocks, pulling the merged global counter back. This is the map-reduce shape expressed in pure Ray Core primitives.

Output.

Concept	Mechanism	Benefit
Broadcast	ray.put + shared ObjectRef	serialize once, not per task
Safe merge	actor with serialized methods	race-free, lock-free
Task chaining	pass ref into another .remote()	automatic dependency ordering
Final gather	single ray.get	one blocking point

Rule of thumb. ray.put any large read-only object you pass to more than a couple of tasks, and use an actor whenever many producers must merge into one shared result. Never pass a big object by value into a fan-out — you will pay the serialization cost N times.

Common beginner mistakes

• Reloading state inside a task. Loading a model or opening a connection at the top of a @ray.remote function reloads it on every call. That work belongs in an actor's __init__.
• ray.get inside a loop that dispatches. Calling ray.get right after each .remote() serializes the whole program — you lose all parallelism. Dispatch first, gather later (or use ray.wait).
• Passing big objects by value repeatedly. Sending a large array as a plain argument to many tasks re-serializes it each time. ray.put once and pass the ref.
• Assuming num_gpus limits the OS. Ray's resources are logical accounting; declaring num_gpus=0.5 does not physically fence GPU memory. Pack fractional GPUs only when you know the memory fits.
• Forgetting actors are single-threaded by default. A hot actor becomes a bottleneck because it processes one call at a time; scale out to an actor pool or raise max_concurrency deliberately.

Data engineering interview question on Ray Core

A senior interviewer might ask: "You have 10,000 CSV shards in S3. You need to parse each, look up every row against a 2 GB reference table, and produce a single merged summary. Design this in Ray Core: which primitives, how you broadcast the reference table, how you avoid blocking on the slowest shard, and how you keep driver memory flat."

Solution Using tasks + ray.put broadcast + an actor accumulator with ray.wait streaming

import ray
from collections import Counter
ray.init()

# 1. Broadcast the 2 GB reference table ONCE via the object store.
reference_table = load_reference_table()          # {key: enrichment}
ref_handle = ray.put(reference_table)             # single copy; share the ref

def load_reference_table() -> dict:
    return {f"k{i}": i for i in range(1_000_000)}  # stand-in

# 2. Stateless per-shard task: parse + enrich, return a compact partial summary.
@ray.remote(num_cpus=1)
def process_shard(s3_key: str, ref_handle) -> dict:
    table = ref_handle                            # resolved from the store
    summary = Counter()
    for row in read_csv_rows(s3_key):             # streaming reader, flat memory
        key = row["lookup_key"]
        if key in table:
            summary[row["category"]] += table[key]
    return dict(summary)                          # small object, not raw rows

def read_csv_rows(key):                           # stand-in streaming reader
    return [{"lookup_key": "k1", "category": "a"}]

# 3. Actor that merges partial summaries safely.
@ray.remote
class SummaryAccumulator:
    def __init__(self):
        self.total = Counter()
    def merge(self, partial: dict) -> int:
        self.total.update(partial)
        return len(self.total)
    def result(self) -> dict:
        return dict(self.total)

acc = SummaryAccumulator.remote()
shard_keys = [f"s3://bucket/shard_{i:05d}.csv" for i in range(10_000)]

# 4. Fan out all shards; stream completions into the accumulator via ray.wait.
pending = [process_shard.remote(k, ref_handle) for k in shard_keys]
while pending:
    ready, pending = ray.wait(pending, num_returns=1)
    for done in ready:
        acc.merge.remote(ray.get(done))           # merge one partial, discard it

final = ray.get(acc.result.remote())
print(f"merged summary over {len(shard_keys)} shards: {len(final)} categories")

Step-by-step trace.

Step	Mechanism	Why
Broadcast reference table	ray.put once → ref_handle	serialize 2 GB once, not 10,000×
Parse + enrich per shard	@ray.remote task	stateless, embarrassingly parallel
Keep per-task output small	return a Counter, not rows	driver never sees 10,000 row sets
Stream completions	ray.wait(num_returns=1)	react at fastest-task latency
Merge safely	actor, serialized merge	race-free, lock-free global merge
Final gather	one ray.get(acc.result)	single blocking point

The 2 GB table is placed in the object store once; every shard task reads it (zero-copy where co-located). Each task returns only a compact summary, so the driver never buffers raw rows. ray.wait feeds completed summaries into the accumulator actor as they finish, keeping both driver memory and the accumulator's update path bounded and race-free.

Output:

Metric	Value
Reference-table serializations	1 (via ray.put)
Parallelism	up to cluster-CPU count
Driver peak memory	O(in-flight partials)
Merge safety	serialized actor, no locks
First partial merged at	fastest-shard latency

Why this works — concept by concept:

• ray.put broadcast — placing the reference table in the object store once and sharing its ObjectRef avoids paying the 2 GB serialization cost per task. Co-located reads are zero-copy memory-mapped views, so lookups are fast and RAM-frugal.
• Compact task outputs — returning a per-shard Counter instead of raw rows keeps every object flowing through the driver tiny, which is what actually bounds driver memory in a 10,000-task fan-out.
• ray.wait streaming — consuming completions one at a time means the pipeline overlaps parsing, merging, and I/O, and never stalls on the single slowest shard.
• Actor accumulator — the actor's one-call-at-a-time execution makes the global merge race-free without any locking, the idiomatic Ray way to reduce many producers into one state.
• Cost — O(rows) parsing is unavoidable, but the broadcast is O(1) serialization instead of O(shards), and memory is O(in-flight) instead of O(shards). That is the difference between a job that scales to 10,000 shards and one that OOMs the driver at 500.

Data Processing
Topic — data-processing
Parallel map-reduce and fan-out problems

Practice →

Optimization Topic — optimization Optimization problems on broadcast and memory

Practice →

---

3. Ray Data — streaming datasets, lazy execution, and map_batches

A dataset is a stream of blocks, transforms are lazy operators, and map_batches is the workhorse

The mental model in one line: ray data is a distributed dataset library where a Dataset is a collection of **blocks (Arrow tables or pandas frames), transforms like map_batches and filter are lazy operators recorded into an execution plan, and a streaming executor pipelines blocks through those operators on execution triggers so the full dataset never has to materialise in memory — giving you distributed data processing with first-class GPU stages and Arrow-native zero-copy batches**. Ray Data is the layer most data engineers actually touch; it is the bridge between raw storage and the tasks/actors of Ray Core.

Blocks — the unit of parallelism.

• What a block is. A Dataset is partitioned into blocks, each an Arrow table (or pandas frame) of many rows. Operators run per-block in parallel; block count sets the parallelism ceiling.
• Reading. ray.data.read_parquet, read_csv, read_json, read_images, from_items, range create datasets. Readers auto-partition source files into blocks.
• Repartition. ds.repartition(n) changes block count to tune parallelism vs overhead; too few blocks under-utilise the cluster, too many add scheduling overhead.

Lazy execution — plan now, run on trigger.

• Lazy by default. map_batches, map, filter, add_column, select_columns build a logical plan; nothing runs yet.
• Triggers. Execution fires on consumption: .take(n), .iter_batches(), .write_parquet(), .count(), .show(), or an explicit .materialize().
• Optimization. Ray Data fuses adjacent operators (e.g. two map_batches into one pass) and plans a physical operator DAG before running.

The streaming executor — bounded memory.

• What it does. Instead of materialising each stage fully before the next, the streaming executor pipelines blocks through the operator DAG, so a 10 TB dataset flows through a fixed memory budget.
• Backpressure. It bounds the number of in-flight blocks per operator, so a slow GPU stage does not let the fast read stage flood memory (section 4).
• materialize(). Explicitly forces the whole dataset into the object store (useful to cache an expensive intermediate you will reuse) — the opt-out from streaming.

map_batches — the workhorse transform.

• Signature. ds.map_batches(fn, batch_size=..., batch_format="numpy"|"pandas"|"pyarrow", concurrency=..., num_gpus=...). fn receives a batch and returns a batch.
• Function vs class. Pass a function for stateless transforms; pass a class (a callable) for stateful ones (model loaded in __init__) — Ray Data runs the class as an actor pool.
• Batch format. numpy for tensor work, pandas for tabular logic, pyarrow for zero-copy columnar. Pick the one that avoids conversions.
• Row vs batch. map and filter are per-row (simpler, slower); map_batches is per-batch (vectorised, the default for performance).

Common interview probes on Ray Data.

• "Is Ray Data lazy?" — yes; transforms build a plan, execution fires on take/write/iter_batches/count.
• "Function or class in map_batches?" — function for stateless, class for stateful (model-once).
• "How does it avoid loading the whole dataset?" — streaming executor pipelines blocks with bounded in-flight memory.
• "What is a block?" — an Arrow/pandas partition; the unit of parallel execution.

Worked example — read_parquet → map_batches → write_parquet, and lazy vs eager

Detailed explanation. The bread-and-butter Ray Data pipeline reads Parquet, transforms batches, and writes Parquet — all streaming. The subtle point most beginners miss is that nothing runs until the write; the transforms only build a plan. Walk through the pipeline and prove the laziness.

• Read. read_parquet auto-partitions the files into blocks; lazy.
• Transform. map_batches records an operator; still lazy.
• Write. write_parquet triggers the streaming executor; blocks flow read → transform → write with bounded memory.

Question. Normalise a numeric column and drop invalid rows over a large Parquet dataset, and show which line actually triggers execution.

Input.

Parameter	Value
Source	s3://lake/events/ (Parquet, 200 GB)
Transform	scale amount, filter amount > 0
batch_format	pandas
Trigger	write_parquet

Code.

import ray
import pandas as pd
ray.init()

# 1. Read — lazy: builds a Read operator, reads nothing yet.
ds = ray.data.read_parquet("s3://lake/events/")

# 2. Transform batches — lazy: appends a MapBatches operator to the plan.
def normalize(batch: pd.DataFrame) -> pd.DataFrame:
    batch = batch[batch["amount"] > 0].copy()      # drop non-positive
    batch["amount_scaled"] = batch["amount"] / batch["amount"].max()
    return batch

transformed = ds.map_batches(normalize, batch_format="pandas", batch_size=10_000)

# 3. Filter — still lazy.
clean = transformed.filter(lambda row: row["status"] == "ok")

print(clean)          # prints the PLAN, not the data (no execution yet)

# 4. Write — TRIGGERS the streaming executor: read -> normalize -> filter -> write
clean.write_parquet("s3://lake/events_clean/")

Step-by-step explanation.

1. read_parquet returns a Dataset whose logical plan is just Read. It has looked at file metadata to plan blocks but read no row data — the operation is lazy.
2. map_batches(normalize, ...) appends a MapBatches operator to the plan. The normalize function is not called yet; Ray Data only records that it should be applied per batch when execution runs.
3. filter(...) appends a Filter operator. Printing clean shows the operator DAG (Read -> MapBatches -> Filter), confirming nothing has executed.
4. write_parquet is a trigger. Now the streaming executor runs the whole DAG: it reads a block, normalises it, filters it, writes it, and moves to the next — so peak memory is a handful of in-flight blocks, not 200 GB.
5. Because execution is deferred to the trigger, Ray Data can fuse MapBatches and Filter into a single pass over each block, avoiding an intermediate materialisation between them.

Output.

Line	Operation	Executes?
read_parquet(...)	Read operator	no (lazy)
map_batches(normalize)	MapBatches operator	no (lazy)
filter(...)	Filter operator	no (lazy)
print(clean)	show plan	no (plan only)
write_parquet(...)	sink	yes (triggers streaming run)

Rule of thumb. Assume every transform is lazy and only sinks/consumers (write_*, take, iter_batches, count, show, materialize) trigger work. If you want to cache an expensive intermediate, call .materialize() explicitly — otherwise it recomputes on each trigger.

Worked example — groupby aggregate and repartition for parallelism

Detailed explanation. Ray Data does relational-style aggregation too, though it is not its core strength. groupby(key).aggregate(...) performs a distributed group-by; repartition(n) tunes how many blocks (hence how much parallelism) the pipeline uses. Walk through a per-category sum with a deliberate repartition.

• groupby/aggregate. Distributed shuffle keyed by the group column, then per-group aggregation.
• repartition. Raise block count before a heavy map to increase parallelism; lower it before a write to control output-file count.
• Aggregations. Sum, Mean, Count, Max, Min, and custom AggregateFn.

Question. Compute total amount per category over a skewed dataset, repartitioning first so the heavy transform uses all cluster cores.

Input.

Parameter	Value
Source	50M rows, 20 categories
Pre-aggregation	map_batches cleanup
Parallelism control	repartition(200)
Aggregation	groupby("category").sum("amount")

Code.

import ray
from ray.data.aggregate import Sum
ray.init()

ds = ray.data.read_parquet("s3://lake/txns/")     # 50M rows, lazy

# Increase block count so the heavy transform saturates all cores.
ds = ds.repartition(200)

# Vectorised per-batch cleanup (lazy).
def clean(batch):
    batch["amount"] = batch["amount"].clip(lower=0)
    return batch

ds = ds.map_batches(clean, batch_format="pandas")

# Distributed group-by aggregation (triggers a shuffle + execution on consume).
totals = ds.groupby("category").aggregate(Sum("amount"))

for row in totals.take_all():                     # take_all triggers execution
    print(row)                                    # {'category': 'a', 'sum(amount)': 12345.0}

Step-by-step explanation.

1. read_parquet is lazy; repartition(200) records that the pipeline should run with 200 blocks. More blocks means more parallel map_batches tasks, which matters when the source produced only a few large blocks.
2. map_batches(clean, ...) is a vectorised per-batch cleanup — clip runs over the whole pandas batch at once, far cheaper than a per-row map.
3. groupby("category").aggregate(Sum("amount")) builds a distributed aggregation: Ray Data shuffles rows so all rows of a category land together, then sums per group. This is the one place Ray Data does a shuffle, so it is the expensive step.
4. take_all() is the trigger. Only now does the read → repartition → clean → group-by DAG execute, streaming where it can and shuffling for the group-by.
5. Repartitioning up before the heavy map improves parallelism; if the final output had too many small files you would repartition down before write_parquet. Block count is the main parallelism knob in Ray Data.

Output.

Stage	Blocks	Effect
after read	~few large	source-file driven
after repartition(200)	200	full-cluster parallelism for map
after groupby	20 (one per group)	shuffle-keyed
take_all()	—	triggers the whole DAG

Rule of thumb. Use repartition as the parallelism knob: raise block count before heavy maps to saturate cores, lower it before writes to control output-file count. Remember groupby is the one shuffle-heavy operator — keep pre-aggregation vectorised in map_batches.

Common beginner mistakes

• Expecting eager results. Building a pipeline and being surprised it "did nothing" — transforms are lazy; only sinks/consumers trigger. Call .materialize() or a consumer to run it.
• Using map where map_batches fits. Per-row map is much slower than vectorised map_batches; reserve map for genuinely per-row logic.
• Wrong batch_format. Returning a NumPy dict from a pandas transform forces conversions each batch. Match the format to your logic.
• Recomputing an expensive intermediate. Consuming the same lazy dataset twice recomputes it. .materialize() (or .write_parquet then re-read) to reuse.
• Ignoring block count. A dataset with 3 huge blocks can only run 3 tasks in parallel. Repartition to match cluster width.

Data engineering interview question on Ray Data

A senior interviewer might ask: "You need to read 2 TB of JSONL logs from S3, parse and enrich each record with a Python function, drop malformed rows, and write partitioned Parquet — on a cluster with 40 CPUs — without ever holding the whole dataset in memory. Walk me through the Ray Data pipeline, where execution triggers, and how you keep memory bounded."

Solution Using a streaming map_batches ETL with explicit repartition and a lazy plan

import ray
import pyarrow as pa
ray.init()

# 1. Read 2 TB JSONL — lazy; reader auto-partitions files into blocks.
ds = ray.data.read_json("s3://logs/raw/", file_extensions=["jsonl"])

# 2. Match block count to cluster width so all 40 CPUs stay busy.
ds = ds.repartition(400)          # ~10 blocks per CPU for good pipelining

# 3. Vectorised parse + enrich per batch (lazy). Return only valid rows.
def parse_enrich(batch: dict) -> dict:
    import numpy as np
    ts = batch["ts"]
    user = batch["user_id"]
    # enrich: derive an hour bucket; mark malformed (missing user) rows
    valid = np.array([u is not None for u in user])
    out = {
        "user_id":  np.array(user)[valid],
        "ts":       np.array(ts)[valid],
        "hour":     np.array([str(t)[:13] for t in np.array(ts)[valid]]),
        "enriched": np.array([lookup_region(u) for u in np.array(user)[valid]]),
    }
    return out

def lookup_region(u):             # stand-in enrichment
    return "us" if hash(u) % 2 else "eu"

parsed = ds.map_batches(parse_enrich, batch_format="numpy", batch_size=20_000)

# 4. Drop any remaining nulls (lazy filter).
clean = parsed.filter(lambda r: r["ts"] is not None)

# 5. Write partitioned Parquet — TRIGGERS the streaming run.
clean.write_parquet(
    "s3://logs/clean/",
    partition_cols=["hour"],      # Hive-style partitioning by hour bucket
)
print("streaming ETL complete")

Step-by-step trace.

Step	Operator	Lazy?	Memory behaviour
read_json	Read	yes	reads metadata only
repartition(400)	Repartition	yes	plans 400 blocks (~10/CPU)
map_batches(parse_enrich)	MapBatches	yes	vectorised per batch
filter(...)	Filter	yes	fused with map where possible
write_parquet(partition_cols)	Write sink	no	triggers streaming; bounded in-flight

The 2 TB read is lazy; repartitioning to 400 blocks gives roughly ten blocks per CPU so the streaming executor always has work to pipeline. map_batches parses and enriches each batch as vectorised NumPy, and the filter is fused into the same block pass. The write_parquet sink is the single trigger — from there the executor streams read → transform → filter → write one block at a time, partitioning output by hour, so the cluster's peak memory is a bounded set of in-flight blocks rather than 2 TB.

Output:

Metric	Value
Peak memory	O(in-flight blocks), not 2 TB
Parallelism	40 CPUs, ~10 blocks each
Execution trigger	write_parquet only
Output layout	Hive-partitioned by hour
Malformed rows	dropped in parse_enrich + filter

Why this works — concept by concept:

• Lazy plan — every transform records an operator instead of running, so Ray Data can fuse the map and filter into one block pass and defer all I/O until the write trigger. Nothing touches 2 TB until it must.
• Streaming executor — the write triggers a pipelined run where blocks flow read → transform → write with a bounded number in flight, which is exactly why a 2 TB dataset fits in a fixed memory budget.
• Repartition for parallelism — sizing to ~10 blocks per CPU keeps all 40 cores fed and lets the streaming executor overlap I/O with compute; too few blocks would idle cores, too many would add scheduling overhead.
• Vectorised map_batches — parsing and enriching a whole NumPy batch at once, rather than per row, is the difference between a compute-bound and an interpreter-bound transform.
• Cost — O(rows) parse work is unavoidable, but memory is O(in-flight blocks) instead of O(dataset), and the partitioned write is O(rows) I/O with no shuffle. That combination is what lets one 40-CPU cluster stream 2 TB without OOM.

Data Processing
Topic — data-processing
Streaming transform and block-parallel problems

Practice →

ETL Topic — etl ETL problems on lazy plans and partitioned writes

Practice →

---

4. Batch inference at scale — GPU actors and autoscaling

Load the model once in a stateful actor, keep the GPUs fed with backpressure, and let the autoscaler right-size the cluster

The mental model in one line: batch inference in Ray Data is map_batches with a **class instead of a function — the class loads the model once in __init__ so it stays resident in GPU memory, concurrency sets how many such GPU actors run as a pool, num_gpus reserves a device per actor, and the streaming executor applies backpressure so the fast read stage never floods memory ahead of the slower GPU stage — while the autoscaler adds GPU nodes when the actor pool has pending demand**. This is Ray's flagship data-engineering workload and the single most-probed Ray interview scenario.

The stateful callable class — model loaded once.

• Why a class. A function passed to map_batches is reconstructed per task; a class is instantiated once per actor and reused across batches. Put the expensive load_model in __init__.
• __call__. The batch transform. Receives a batch (NumPy/pandas/Arrow), runs the model, returns predictions merged into the batch.
• concurrency. An int (fixed pool size) or a (min, max) tuple (autoscaling pool). This is how many model replicas run in parallel.
• num_gpus. Reserve one (or a fraction of a) GPU per actor. Fractional GPUs let multiple small models share a device when memory permits.

Backpressure — keep the GPU fed without OOM.

• The problem. Reading and decoding are fast; GPU inference is slow. Without limits, the read stage races ahead and buffers gigabytes of decoded batches waiting for the GPU — OOM.
• The mechanism. The streaming executor bounds in-flight blocks per operator and the object-store memory budget, so the read stage throttles to match GPU throughput.
• Tuning. batch_size trades GPU utilisation (bigger batches = better throughput) against memory (bigger batches = more GPU RAM). DataContext knobs cap object-store usage.

Autoscaling — right-size the cluster.

• How. The autoscaler watches pending actors/tasks that cannot be placed and adds worker nodes (up to max); it removes idle nodes after a timeout.
• GPU pools. A concurrency=(4, 16) actor pool tells the autoscaler it may grow the GPU replica count from 4 to 16 as demand warrants.
• Cost control. min_replicas low + max_replicas bounded keeps the GPU bill proportional to actual work; nodes scale to zero when idle.

The end-to-end shape.

• Read (CPU, streaming) → decode/preprocess (map_batches function, CPU) → infer (map_batches class, GPU actor pool) → write (CPU, streaming). Each stage runs on the hardware it declared.

Common interview probes on batch inference.

• "Function or class for the model stage?" — class, so the model loads once.
• "How do you keep the GPU busy?" — right-size batch_size and pool concurrency; let backpressure throttle the reader.
• "How do you avoid OOM?" — streaming execution + bounded in-flight blocks + a batch_size the GPU memory can hold.
• "How does the cluster grow?" — autoscaler adds GPU nodes for pending pool replicas; scales down when idle.

Worked example — a GPU batch-inference pipeline with an actor pool

Detailed explanation. The canonical Ray Data inference job: read images, decode on CPU, run a vision model on a GPU actor pool, write embeddings. The load-once class is the crux. Walk through the pipeline stage by stage, noting which hardware each stage runs on.

• Read + decode. CPU stages via map_batches functions; cheap and parallel.
• Infer. A class in map_batches with concurrency and num_gpus=1; model loaded once in __init__.
• Write. Streaming sink; triggers the run.

Question. Score 100M images with a vision model on 8 GPUs, loading the model once per GPU, and write the embeddings to Parquet.

Input.

Parameter	Value
Source	s3://images/ (100M images)
Decode stage	CPU map_batches function
Infer stage	GPU actor class, concurrency=8, num_gpus=1
batch_size	256 (fits GPU memory)
Sink	write_parquet

Code.

import ray
import numpy as np
ray.init()

# --- CPU stage: decode + preprocess (stateless function) ---
def preprocess(batch: dict) -> dict:
    # batch["image"] is raw bytes; decode + resize to a fixed tensor
    imgs = [decode_resize(b) for b in batch["image"]]
    batch["pixels"] = np.stack(imgs).astype("float32")
    return batch

def decode_resize(b):                 # stand-in for PIL/opencv decode+resize
    return np.zeros((3, 224, 224), dtype="float32")

# --- GPU stage: stateful class; model loads ONCE per actor ---
class Embedder:
    def __init__(self, weights_uri: str):
        self.model = load_vision_model(weights_uri)   # once, into GPU memory

    def __call__(self, batch: dict) -> dict:
        pixels = batch["pixels"]                       # (N, 3, 224, 224)
        batch["embedding"] = self.model(pixels)        # (N, 512)
        del batch["pixels"]                            # drop heavy intermediate
        return batch

def load_vision_model(uri):           # stand-in for a torch model .cuda().eval()
    return lambda x: np.asarray(x).reshape(len(x), -1)[:, :512]

# --- Pipeline: read -> decode (CPU) -> embed (GPU pool) -> write ---
ds = ray.data.read_images("s3://images/", size=(224, 224))

embeddings = (
    ds
    .map_batches(preprocess, batch_format="numpy", batch_size=256)
    .map_batches(
        Embedder,
        fn_constructor_args=("s3://models/clip.pt",),
        concurrency=8,        # 8 GPU actors
        num_gpus=1,           # one GPU each
        batch_size=256,       # rows per model call
        batch_format="numpy",
    )
    .drop_columns(["image"])  # don't write raw bytes
)

embeddings.write_parquet("s3://embeddings/")   # triggers the streaming run

Step-by-step explanation.

1. read_images builds a lazy dataset of image blocks. preprocess is a stateless function — decoding has no state to keep, so a function (not a class) is correct and cheapest.
2. Embedder is a class passed to map_batches. Ray Data instantiates it once per actor; load_vision_model runs a single time and the weights stay resident in GPU memory across every batch that actor handles.
3. concurrency=8 + num_gpus=1 creates a pool of eight GPU actors, one per device. The scheduler places each on a node with a free GPU; eight model replicas run in parallel.
4. batch_size=256 is chosen so a batch of 256 preprocessed tensors fits in one GPU's memory alongside the model. Dropping pixels after inference frees the heavy intermediate before the batch flows onward.
5. write_parquet triggers the streaming executor: read → decode (CPU) → embed (GPU pool) → write, all pipelined. Backpressure keeps the CPU decode stage from racing ahead of the eight GPUs, so decoded batches don't pile up in memory.

Output.

Stage	Hardware	Function or class	Runs model?
read_images	CPU	reader	no
preprocess	CPU	function	no
Embedder	GPU (×8)	class (actor pool)	yes (loaded once each)
write_parquet	CPU	sink	no (triggers run)

Rule of thumb. The model stage is always a class so the model loads once; decode/parse stages are functions. Size batch_size to fill GPU memory, concurrency to the GPU count, and let backpressure protect memory. That is the entire batch-inference recipe.

Worked example — the OOM-from-no-backpressure failure and its fix

Detailed explanation. A common production incident: an inference job runs fine on a sample but OOMs on the full dataset. The cause is almost always a fast producer flooding a slow GPU consumer, combined with a batch_size too big for GPU memory. Walk through the diagnosis and the two-part fix.

• Symptom. Object-store memory climbs, spilling to disk thrashes, then the GPU worker OOMs mid-run.
• Root cause A. batch_size too large — a single batch's tensors exceed GPU memory.
• Root cause B. No effective throttle — the read/decode stage outruns the GPU and buffers decoded batches.

Question. Diagnose and fix an inference job that OOMs at scale; show the config that keeps memory bounded.

Input.

Symptom	Cause	Fix
GPU OOM per batch	batch_size too big	lower batch_size to fit device memory
Object store fills / spills	fast producer, slow GPU	rely on streaming backpressure; cap object-store fraction
Cluster idles then bursts	wrong pool size	set concurrency to real GPU count

Code.

import ray
from ray.data import DataContext
ray.init()

# 1. Bound the object-store memory the streaming executor may use.
ctx = DataContext.get_current()
ctx.execution_options.resource_limits.object_store_memory = 8 * 1024**3  # 8 GB

ds = ray.data.read_images("s3://images/", size=(224, 224))

class Embedder:
    def __init__(self, uri):
        self.model = load_vision_model(uri)
    def __call__(self, batch):
        batch["embedding"] = self.model(batch["pixels"])
        del batch["pixels"]                 # free the heavy intermediate ASAP
        return batch

def load_vision_model(uri):
    return lambda x: x

out = (
    ds
    .map_batches(preprocess, batch_format="numpy", batch_size=64)   # was 256
    .map_batches(
        Embedder,
        fn_constructor_args=("s3://models/clip.pt",),
        concurrency=8,
        num_gpus=1,
        batch_size=64,       # smaller batch fits GPU memory with headroom
        batch_format="numpy",
    )
)
out.write_parquet("s3://embeddings/")

Step-by-step explanation.

1. The OOM had two roots. First, batch_size=256 produced tensors too large for one GPU's memory once the model's own activations were added. Lowering to 64 leaves headroom — the first and most direct fix.
2. Second, the streaming executor already applies backpressure, but its object-store budget was effectively unbounded, so decoded batches accumulated. Setting resource_limits.object_store_memory = 8 GB caps how much the executor buffers, forcing the read/decode stage to throttle to GPU speed.
3. Deleting batch["pixels"] right after inference frees the largest intermediate before the batch moves to the write stage, cutting steady-state memory further.
4. concurrency=8 matched to the eight real GPUs stops the scheduler from over- or under-provisioning replicas — over-provisioning would multiply GPU-memory pressure, under-provisioning would idle devices.
5. After the fix, object-store memory holds flat under the cap, the GPUs stay near-saturated because backpressure keeps just enough decoded batches in flight, and the job completes without spilling-thrash.

Output.

Metric	Before (OOM)	After (fixed)
batch_size	256	64
Object-store cap	unbounded	8 GB
Heavy intermediate	kept	freed after infer
GPU utilisation	crash	~90% steady
Outcome	OOM mid-run	completes, memory flat

Rule of thumb. When an inference job OOMs at scale, first shrink batch_size to fit GPU memory, then cap the streaming executor's object-store budget so the fast producer throttles to the slow GPU, and free heavy intermediates inside __call__. Backpressure only helps if you give it a memory ceiling to enforce.

Common beginner mistakes

• Model in __call__ instead of __init__. Reloads the model every batch — the classic Ray inference blunder. Load in __init__.
• Passing a function for the model stage. A function can't hold the loaded model; use a class so Ray runs it as a reusable actor.
• batch_size sized for throughput only. Too large and a single batch OOMs the GPU. Size for GPU memory first, then push it up.
• concurrency unrelated to GPU count. More actors than GPUs oversubscribe devices; fewer waste them. Match the pool to real GPUs (or use fractional num_gpus).
• No object-store ceiling. Without a memory cap, a fast reader floods the store ahead of a slow GPU and the job spills or OOMs.

Data engineering interview question on batch inference

A senior interviewer might ask: "Design a Ray batch-inference job to embed 400 million text documents with a transformer model on a cluster that can autoscale from 2 to 16 GPUs. Cover the actor pool, model-once loading, batch sizing, backpressure, autoscaling config, and how you'd write the results idempotently."

Solution Using an actor-pool map_batches inference pipeline with autoscaling and backpressure

import ray
from ray.data import DataContext
ray.init()

# 1. Cap the streaming executor's object-store budget (backpressure ceiling).
ctx = DataContext.get_current()
ctx.execution_options.resource_limits.object_store_memory = 16 * 1024**3

# 2. Read 400M documents lazily; reader partitions into blocks.
ds = ray.data.read_parquet("s3://corpus/docs/")

# 3. CPU stage: tokenize (stateless function).
def tokenize(batch: dict) -> dict:
    batch["input_ids"] = [encode(t) for t in batch["text"]]
    return batch

def encode(t):                        # stand-in tokenizer
    return [ord(c) % 128 for c in t[:512]]

# 4. GPU stage: stateful class; transformer loaded ONCE per actor.
class Embedder:
    def __init__(self, model_uri: str):
        self.model = load_transformer(model_uri)     # once, resident on GPU
    def __call__(self, batch: dict) -> dict:
        batch["embedding"] = self.model(batch["input_ids"])
        del batch["input_ids"]                        # free heavy intermediate
        return batch

def load_transformer(uri):
    return lambda ids: [[0.0] * 768 for _ in ids]

# 5. Pipeline with an AUTOSCALING actor pool: 2..16 GPU replicas.
embedded = (
    ds
    .map_batches(tokenize, batch_format="pandas", batch_size=512)
    .map_batches(
        Embedder,
        fn_constructor_args=("s3://models/bge.pt",),
        concurrency=(2, 16),   # autoscale GPU replicas with demand
        num_gpus=1,            # one GPU per replica
        batch_size=128,        # fits GPU memory with headroom
        batch_format="pandas",
    )
)

# 6. Idempotent write: partition by a stable id bucket so re-runs overwrite cleanly.
embedded.write_parquet(
    "s3://corpus/embeddings/",
    partition_cols=["shard_id"],
    mode="overwrite",          # re-running a shard replaces it, not appends
)
print("400M-doc embedding job submitted")

Step-by-step trace.

Layer	Component	Purpose
Backpressure ceiling	object_store_memory = 16 GB	fast tokenizer throttles to GPU speed
Read	read_parquet (lazy)	partition 400M docs into blocks
Tokenize	map_batches function (CPU)	stateless, vectorised per batch
Embed	map_batches class, concurrency=(2,16), num_gpus=1	model-once GPU actor pool
Autoscale	pool grows 2→16 on pending demand	GPU count tracks backlog
Write	write_parquet(partition_cols, mode="overwrite")	idempotent, re-run-safe

The 400M-document read is lazy. Tokenization runs as a CPU function; embedding runs on an autoscaling pool of GPU actors that each load the transformer exactly once. concurrency=(2, 16) lets the autoscaler add GPU nodes as the streaming executor reports pending inference work, and drop them when the backlog clears. The 16 GB object-store ceiling makes the fast tokenizer throttle to GPU throughput, and the partitioned overwrite write means re-running any shard replaces its output rather than duplicating it.

Output:

Metric	Value
Model loads	once per active GPU replica (≤16)
GPU replicas	autoscales 2→16 with demand
Peak object-store memory	≤16 GB (capped)
Write semantics	idempotent (partitioned overwrite)
Steady-state GPU utilisation	~90% (backpressure-fed)

Why this works — concept by concept:

• Model-once GPU actor — the embedder class loads the transformer in __init__, so each of the up-to-16 GPU replicas pays the load cost once and reuses the resident model across millions of documents.
• Autoscaling actor pool — concurrency=(2, 16) hands the autoscaler a range; pending inference work grows the pool (and the cluster) toward 16 GPUs, and idle capacity shrinks it back, so the GPU bill tracks real backlog.
• Backpressure ceiling — capping object-store memory forces the fast CPU tokenizer to run only as far ahead as the GPUs can consume, which is what prevents the OOM/spill-thrash failure at 400M scale.
• Idempotent partitioned write — writing Parquet partitioned by shard_id with mode="overwrite" makes retries safe: re-running a failed shard replaces exactly that partition instead of appending duplicates.
• Cost — inference is O(docs) GPU work no matter what, but the design collapses the constant factor (model loaded ≤16 times, GPUs ~90% busy) and bounds memory to a fixed ceiling, so the job scales from a sample to 400M documents without a rewrite.

Optimization
Topic — optimization
Throughput and backpressure optimization problems

Practice →

Data Processing Topic — data-processing GPU batch-inference pipeline problems

Practice →

---

5. Ray clusters, ops, and interview signals

A Ray cluster is a head node plus autoscaling workers; KubeRay runs it on Kubernetes; fault tolerance is retries plus lineage — and knowing when NOT to use Ray is the senior signal

The mental model in one line: a ray clusters deployment is one **head node (running the Global Control Store and the autoscaler) plus a pool of worker nodes, most often provisioned on Kubernetes by the KubeRay operator through RayCluster / RayJob / RayService custom resources, with fault tolerance layered as task retries, actor restarts, and object lineage reconstruction — and the mark of a senior engineer is naming the workloads where Ray is the wrong tool**. Ops is where Ray knowledge stops being trivia and starts being production judgement.

Cluster anatomy — head and workers.

• Head node. Runs the Global Control Store (GCS) — cluster metadata, actor registry, resource accounting — plus the autoscaler and the dashboard. Losing the head loses the cluster unless GCS fault tolerance is configured.
• Worker nodes. Run tasks and actors; each has its own object store (shared memory). Workers are cattle — added and removed by the autoscaler.
• Autoscaler. Watches pending tasks/actors that cannot be placed and adds nodes up to a max; removes idle nodes after a timeout. Node types can be heterogeneous (CPU pools, GPU pools).

KubeRay — Ray on Kubernetes.

• The operator. KubeRay is a Kubernetes operator that manages Ray clusters via CRDs. It is the standard production deployment path in 2026.
• RayCluster. A long-lived cluster (head + worker groups) you submit jobs to. Worker groups define replicas, resources, and autoscaling bounds.
• RayJob. A run-to-completion batch job — creates a cluster (or uses an existing one), runs an entrypoint, tears down. The idiomatic way to run a batch-inference or ETL job.
• RayService. A long-running Ray Serve deployment with rolling upgrades and health checks — for online model serving, not batch.

Fault tolerance — retries, restarts, lineage.

• Task retries. @ray.remote(max_retries=3) re-executes a failed task (e.g. on a lost worker). Idempotent tasks make this safe.
• Actor restarts. @ray.remote(max_restarts=3) recreates a crashed actor; combine with max_task_retries to re-send in-flight calls. State in __init__ is rebuilt (reload the model).
• Object reconstruction. If an object in the store is lost with its worker, Ray re-runs the task that produced it using the recorded lineage — automatic recovery of intermediate data.
• GCS fault tolerance. Back the head's GCS with external Redis so a head restart does not lose the cluster; essential for long-running RayService.

When NOT to use Ray — the senior signal.

• Single-node pandas fits. If the data fits on one machine and pandas/DuckDB handles it, a cluster is pure overhead and operational risk. Use the single-node tool.
• Pure relational ETL on a mature Spark shop. Big shuffle-heavy joins with a working Spark platform? Don't migrate for its own sake — Ray wins on Python/ML, not relational shuffle.
• A governed SQL warehouse workload. GROUP BY over Snowflake/BigQuery belongs in the warehouse, not in Ray Data.
• Tiny or latency-critical online serving with no ML. The cluster's coordination overhead isn't worth it for small jobs; and hard real-time request/response has other tools.

Common interview probes on Ray ops.

• "What runs on the head node?" — GCS, autoscaler, dashboard; back GCS with Redis for HA.
• "RayJob vs RayService?" — batch run-to-completion vs long-running online serving.
• "How does Ray recover a lost object?" — lineage reconstruction re-runs the producing task.
• "When would you not use Ray?" — single-node data, pure relational ETL, warehouse SQL, tiny jobs.

Worked example — a KubeRay RayJob for a batch pipeline

Detailed explanation. In production you rarely start a cluster by hand; you declare a RayJob and let KubeRay create an ephemeral cluster, run the entrypoint, and tear down. Walk through the manifest and what each field controls.

• entrypoint. The command that runs your Ray script on the head.
• rayClusterSpec. Head and worker group definitions, including autoscaling bounds and GPU resources.
• shutdownAfterJobFinishes. Tears the cluster down on completion — the cost-control default for batch.

Question. Write a KubeRay RayJob that runs a batch-inference script on an autoscaling GPU worker group and tears down when done.

Input.

Field	Value
Kind	RayJob
Entrypoint	python batch_infer.py
Worker group	GPU, autoscale 1→8
Teardown	shutdownAfterJobFinishes: true

Code.

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: batch-infer-embeddings
spec:
  entrypoint: python /home/ray/batch_infer.py
  shutdownAfterJobFinishes: true        # ephemeral cluster; tears down on finish
  rayClusterSpec:
    rayVersion: "2.9.0"
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
            - name: ray-head
              image: rayproject/ray:2.9.0
              resources:
                requests: {cpu: "2", memory: "8Gi"}
    workerGroupSpecs:
      - groupName: gpu-workers
        replicas: 1
        minReplicas: 1
        maxReplicas: 8                   # autoscaler grows the GPU pool to 8
        rayStartParams: {}
        template:
          spec:
            containers:
              - name: ray-worker
                image: rayproject/ray-ml:2.9.0-gpu
                resources:
                  requests: {cpu: "8", memory: "32Gi", nvidia.com/gpu: "1"}
                  limits:   {nvidia.com/gpu: "1"}

Step-by-step explanation.

1. kind: RayJob tells KubeRay this is a run-to-completion batch job, not a long-lived cluster or a serving deployment. The operator provisions a cluster, runs the entrypoint, and (because of the shutdown flag) tears it down.
2. entrypoint is the command executed on the head once the cluster is ready — here the batch-inference script from section 4. Its concurrency=(1,8) actor pool will drive the worker autoscaling.
3. headGroupSpec sizes the head, which runs GCS and the autoscaler. It needs modest CPU/memory; it is not where the GPU work happens.
4. workerGroupSpecs defines the GPU pool: minReplicas: 1, maxReplicas: 8, and nvidia.com/gpu: "1" per worker. The autoscaler adds workers (up to 8) as the Ray job reports pending GPU actors, matching cluster size to demand.
5. shutdownAfterJobFinishes: true is the cost-control default for batch — the GPU nodes exist only for the job's duration, so you pay for GPUs only while inference runs.

Output.

CRD	Lifecycle	Use
RayCluster	long-lived	shared cluster for many jobs
RayJob	run-to-completion	batch ETL / inference (this example)
RayService	long-running	online Ray Serve with rolling upgrades

Rule of thumb. Use RayJob with shutdownAfterJobFinishes: true for batch — an ephemeral, autoscaling, GPU cluster that exists only while the job runs. Reserve RayCluster for shared interactive use and RayService for online serving.

Worked example — fault-tolerance configuration and recovery

Detailed explanation. Long-running jobs must survive worker loss, actor crashes, and object loss. Ray gives three knobs; a senior answer configures all three deliberately. Walk through each with the idempotency assumption it requires.

• Task retries. Re-run failed tasks; safe only if the task is idempotent (or its output is overwritten).
• Actor restarts + task retries. Recreate a crashed actor and re-send its in-flight calls; the actor rebuilds state in __init__.
• Lineage reconstruction. Automatic — Ray re-runs the producing task to rebuild a lost object.

Question. Configure a batch job to tolerate a lost GPU worker mid-run without failing the whole job, and explain what happens on failure.

Input.

Failure	Mechanism	Requirement
Task fails (worker lost)	max_retries	idempotent task
Actor crashes (OOM, node loss)	max_restarts + max_task_retries	state rebuildable in __init__
Object lost with worker	lineage reconstruction	deterministic producer

Code.

import ray
ray.init()

# Idempotent task: retried safely on worker loss.
@ray.remote(max_retries=3)
def transform_shard(shard_uri: str) -> str:
    out_uri = shard_uri.replace("/raw/", "/clean/")
    write_overwrite(out_uri, process(shard_uri))   # overwrite => retry-safe
    return out_uri

def process(u): return b"..."
def write_overwrite(u, data): pass

# Restartable GPU actor: recreated on crash; model reloads in __init__.
@ray.remote(num_gpus=1, max_restarts=3, max_task_retries=2)
class Embedder:
    def __init__(self, uri: str):
        self.model = load_model(uri)      # rebuilt automatically on restart
    def embed(self, batch):
        return self.model(batch)

def load_model(uri): return lambda b: b

# On a lost worker: pending transform_shard tasks retry (up to 3);
# the Embedder actor is recreated (up to 3) and in-flight embed calls
# retry (up to 2); any lost intermediate object is reconstructed from lineage.
refs = [transform_shard.remote(f"s3://raw/shard_{i}.parquet") for i in range(100)]
print(ray.get(refs)[:3])

Step-by-step explanation.

1. max_retries=3 on transform_shard means a task killed by a lost worker is re-scheduled up to three times. This is only safe because the task overwrites its output — a re-run produces the same file, so retries can't duplicate data.
2. max_restarts=3 on the Embedder actor tells Ray to recreate the actor process if it crashes (GPU OOM, node loss). Because the model is loaded in __init__, the restarted actor rebuilds its full state automatically — no external checkpoint needed.
3. max_task_retries=2 re-sends method calls that were in flight when the actor died to the restarted actor, so a crash mid-embed doesn't lose that batch.
4. Lineage reconstruction is automatic and needs no flag: if a worker holding an intermediate object dies, Ray re-runs the deterministic task that produced it to rebuild the object, then continues. This is why Ray Data pipelines survive worker churn.
5. The combination — idempotent retried tasks, restartable actors, and lineage — lets a 100-shard job survive losing a GPU worker mid-run: pending work retries elsewhere, the model reloads on a new node, and no completed work is redone unnecessarily.

Output.

Failure event	Ray response	Data effect
GPU worker lost	tasks retry, actor restarts	none (overwrite + rebuild)
Actor OOM	actor recreated, calls retried	batch reprocessed
Object lost	lineage re-runs producer	intermediate rebuilt
Head lost (no GCS FT)	cluster down	job fails — configure Redis GCS

Rule of thumb. Set max_retries on idempotent tasks, max_restarts + max_task_retries on stateful actors, and rely on automatic lineage for intermediates — but remember the head node is the single point of failure unless you back GCS with external Redis. Retries are only safe if the work is idempotent.

Common beginner mistakes

• Assuming the head node is HA by default. Without GCS-Redis fault tolerance, a head restart kills the cluster. Configure it for anything long-running.
• Retrying non-idempotent tasks. max_retries on a task that appends to a sink duplicates data on retry. Make the task overwrite or dedupe first.
• Running batch on a RayService. Serving CRDs are for online Ray Serve; batch belongs in a RayJob that tears down after.
• Over-provisioning the cluster manually. Fixing worker count defeats the autoscaler; set min/max and let it size to demand.
• Using Ray where single-node fits. A cluster for data that fits in one machine's RAM is cost and operational overhead for no benefit.

Data engineering interview question on Ray ops

A senior interviewer might ask: "You're moving a nightly Ray batch-inference job to production on Kubernetes. It runs for ~3 hours on an autoscaling GPU pool and must survive worker preemption without failing or duplicating output. Walk me through the KubeRay resource you'd use, the fault-tolerance config, the idempotent-write strategy, and where you'd still be exposed."

Solution Using a KubeRay RayJob with autoscaling GPUs, task/actor fault tolerance, and idempotent writes

# 1. KubeRay RayJob — ephemeral, autoscaling GPU cluster, torn down on finish.
apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: nightly-batch-infer
spec:
  entrypoint: python /home/ray/nightly_infer.py
  shutdownAfterJobFinishes: true
  rayClusterSpec:
    rayVersion: "2.9.0"
    headGroupSpec:
      template:
        spec:
          containers:
            - name: ray-head
              image: rayproject/ray:2.9.0
              resources: {requests: {cpu: "2", memory: "8Gi"}}
    workerGroupSpecs:
      - groupName: gpu-workers
        replicas: 2
        minReplicas: 2
        maxReplicas: 12            # autoscale GPUs with pending demand
        template:
          spec:
            containers:
              - name: ray-worker
                image: rayproject/ray-ml:2.9.0-gpu
                resources:
                  requests: {cpu: "8", memory: "32Gi", nvidia.com/gpu: "1"}
                  limits:   {nvidia.com/gpu: "1"}

# 2. nightly_infer.py — fault-tolerant, idempotent inference.
import ray
from ray.data import DataContext
ray.init()

DataContext.get_current().execution_options.resource_limits.object_store_memory = 12 * 1024**3

class Embedder:                      # restartable GPU actor; model-once
    def __init__(self, uri):
        self.model = load_model(uri)
    def __call__(self, batch):
        batch["embedding"] = self.model(batch["input_ids"])
        del batch["input_ids"]
        return batch

def load_model(uri): return lambda x: x

ds = ray.data.read_parquet("s3://corpus/docs/")
out = ds.map_batches(
    Embedder,
    fn_constructor_args=("s3://models/bge.pt",),
    concurrency=(2, 12),            # matches worker autoscaling range
    num_gpus=1,
    batch_size=128,
    max_restarts=3,                 # actor recreated on preemption
    max_task_retries=2,             # in-flight batches re-sent
)

# Idempotent write: partitioned overwrite so preemption re-runs replace, not append.
out.write_parquet(
    "s3://corpus/embeddings/dt=2026-08-03/",
    partition_cols=["shard_id"],
    mode="overwrite",
)
print("nightly batch inference complete")

Step-by-step trace.

Layer	Config	Failure it covers
Ephemeral cluster	RayJob + shutdownAfterJobFinishes	pays for GPUs only during the run
Autoscaling GPUs	minReplicas 2 → maxReplicas 12	scales to backlog, down when idle
Actor restart	max_restarts=3	GPU worker preemption / OOM
Call retry	max_task_retries=2	batch in flight when actor died
Lineage	automatic	lost intermediate objects
Idempotent write	partitioned mode="overwrite"	re-run duplication

The RayJob provisions an autoscaling GPU cluster that grows from 2 to 12 workers as the streaming executor reports pending inference, then tears down on completion. The Embedder actor loads the model once and is configured to restart on preemption, with in-flight batches re-sent; lost intermediates are rebuilt from lineage automatically. The partitioned-overwrite write makes retries safe — a re-run of any shard replaces its partition rather than appending duplicates. The residual exposure is the head node: without GCS-Redis fault tolerance, a head-pod eviction fails the whole job.

Output:

Metric	Value
GPU workers	autoscale 2→12, ephemeral
Worker preemption	survived (restart + retry + lineage)
Output duplication on retry	none (partitioned overwrite)
Cost profile	GPUs billed only during the run
Residual risk	head node (needs GCS-Redis HA)

Why this works — concept by concept:

• RayJob ephemerality — a run-to-completion CRD that provisions and tears down an autoscaling cluster means the expensive GPU pool exists only while the nightly job runs, which is the correct cost posture for batch.
• Actor restart + task retry — max_restarts recreates a preempted GPU actor (reloading the model in __init__), and max_task_retries re-sends the batch it was mid-processing, so a lost worker costs a few reprocessed batches, not the job.
• Lineage reconstruction — Ray rebuilds any intermediate object lost with a worker by re-running its deterministic producer, so worker churn during a 3-hour run is transparent to the pipeline.
• Idempotent partitioned write — writing partitioned by shard_id with mode="overwrite" makes every retry safe: replacing a partition can never duplicate rows, which is the invariant that lets you enable retries at all.
• Cost — inference is O(docs) GPU work; the design keeps the constant factor low (model-once, GPUs backpressure-fed) and bills GPUs only for the run, while the one uncovered failure mode — the head node — is a known, fixable gap (GCS-Redis) rather than a silent one.

Optimization
Topic — optimization
Cluster sizing and cost-optimization problems

Practice →

Streaming
Topic — streaming
Fault-tolerance and recovery problems

Practice →

---

Cheat sheet — Ray for data engineering recipes

• Three primitives. ray tasks = @ray.remote functions (stateless, return ObjectRef futures); ray actors = @ray.remote classes (stateful workers, one call at a time); object store = shared-memory holding every result with zero-copy co-located reads. Classify each pipeline stage as stateless → task, stateful → actor, data → object before writing code.
• Fan-out + streaming completion. Dispatch all tasks first (refs = [f.remote(x) for x in xs]), then consume with while pending: ready, pending = ray.wait(pending, num_returns=1). Never ray.get inside the dispatch loop — it serializes the program and kills parallelism.
• Broadcast large read-only objects. ref = ray.put(big_obj) once, pass ref to every task; Ray serializes it a single time and co-located tasks read it zero-copy. Passing the object by value into a fan-out re-serializes it per task.
• Ray Data is lazy. read_*, map_batches, map, filter, add_column, select_columns build a plan; execution triggers on write_*, take, take_all, iter_batches, count, show, or explicit materialize(). Print a dataset to see the operator DAG, not the rows.
• map_batches function vs class. Function for stateless transforms (decode, parse, tokenize); class for stateful ones (model-once) — Ray runs the class as an actor pool. batch_format = numpy/pandas/pyarrow; pick the one that avoids conversions.
• Batch inference recipe. Model loads in __init__, inference in __call__, concurrency = GPU replica count (or (min,max) to autoscale), num_gpus=1 per actor, batch_size sized to GPU memory. Free heavy intermediates (del batch["pixels"]) after inference.
• Backpressure + OOM defense. The streaming executor bounds in-flight blocks; cap it with DataContext.get_current().execution_options.resource_limits.object_store_memory. When an inference job OOMs: shrink batch_size to fit GPU memory first, then cap object-store memory so the fast reader throttles to the GPU.
• Repartition for parallelism. ds.repartition(n) is the main parallelism knob — raise block count (~10× CPU) before heavy maps to saturate cores; lower it before writes to control output-file count. A dataset with 3 huge blocks can only run 3 tasks in parallel.
• groupby is the one shuffle. ds.groupby(key).aggregate(Sum("col")) shuffles; keep pre-aggregation cleanup vectorised in map_batches. Ray Data is a batch engine, not a SQL warehouse — for governed GROUP BY at scale, use the warehouse.
• KubeRay CRDs. RayCluster = long-lived shared cluster; RayJob (+ shutdownAfterJobFinishes: true) = ephemeral run-to-completion batch (ETL/inference); RayService = long-running online Ray Serve with rolling upgrades. Batch jobs are RayJob, not RayService.
• Fault tolerance. @ray.remote(max_retries=3) on idempotent tasks; @ray.remote(max_restarts=3, max_task_retries=2) on stateful actors (state rebuilds in __init__); object loss is auto-recovered via lineage. The head node's GCS is the single point of failure — back it with external Redis for long-running clusters.
• When NOT to use Ray. Data fits on one node (use pandas/DuckDB); pure relational shuffle-ETL on a mature Spark platform (keep Spark); governed warehouse SQL (use the warehouse); tiny or hard-real-time jobs (coordination overhead not worth it). Naming the anti-pattern unprompted is the senior signal.

Frequently asked questions

What is Ray and how does it relate to data engineering?

Ray is an open-source distributed runtime whose unit of parallelism is a Python function (a task) or a Python class instance (an actor), with a shared-memory object store that lets those units exchange data without copying. For ray for data engineering, the relevant layer is usually Ray Data — a streaming, lazy, block-based dataset library — plus Ray Core underneath it. Ray fits the Python-heavy, model-heavy, heterogeneous-hardware half of data engineering: batch inference over hundreds of millions of records, last-mile ML preprocessing, embeddings pipelines, and any "run this Python across a cluster" workload that isn't naturally a SQL query. It is not a SQL warehouse and not a replacement for Spark on relational ETL; it wins specifically where Python and GPUs dominate the work.

When should I use Ray instead of Spark?

Use Ray when the heavy work is Python, GPU, or stateful; use Spark when the heavy work is relational shuffle on an existing platform. Spark runs Python through a JVM boundary, so mapPartitions UDFs and model serving pay a per-batch serialization tax and reload state per partition; Ray runs the same logic as native Python with Arrow-native zero-copy batches and lets a stateful actor load a model exactly once. The clearest Ray wins are batch inference with GPU actors, tokenization/decoding-heavy preprocessing, and heterogeneous CPU→GPU pipelines. The clearest Spark wins are big JOIN/GROUP BY shuffles over columnar data where a mature Spark platform already exists. The senior framing is "it depends on relational-vs-Python and on the existing platform," never "Ray is faster."

What is Ray Data and how is it different from a Spark DataFrame?

ray data is a distributed dataset library where a Dataset is a collection of blocks (Arrow tables or pandas frames), transforms are lazy operators recorded into a plan, and a streaming executor pipelines blocks through those operators so the full dataset never has to fit in memory. The workhorse transform is map_batches, which applies a function (stateless) or a class (stateful, model-once) to each batch, with first-class num_gpus scheduling. Compared to a Spark DataFrame, Ray Data is Python-native (no JVM boundary), designed around ML/GPU stages rather than SQL, and streams execution by default. It is weaker than Spark at large relational shuffles and is not a query-optimized SQL engine — it is a data-loading-and-transform layer optimized for feeding models.

How does batch inference work in Ray, and how do I keep the GPU busy without OOM?

Batch inference is map_batches with a class: the model loads once in __init__ (so it stays resident on the GPU across every batch), inference runs in __call__, concurrency sets the number of GPU-actor replicas, and num_gpus=1 reserves a device per actor. To keep the GPU busy, size batch_size to fill GPU memory and set concurrency to the real GPU count (or a (min,max) range for autoscaling). To avoid OOM, rely on the streaming executor's backpressure — it bounds in-flight blocks so the fast read/decode stage throttles to GPU speed — and cap the object-store budget via DataContext. The two most common OOM fixes are shrinking batch_size to fit device memory and freeing heavy intermediates (del batch["pixels"]) right after inference.

What is the Ray object store and why does it matter?

The object store is a shared-memory store on each node that holds the value behind every ObjectRef — every .remote() result and every ray.put(). It matters for two reasons. First, zero-copy sharing: tasks and actors co-located on a node read a stored object as a memory-mapped view without copying it, so passing a large batch between stages is cheap. Second, broadcast efficiency: ray.put(big_obj) once and sharing the ObjectRef avoids re-serializing a large read-only object (a lookup table, an embedding matrix) per task. The store also underpins fault tolerance — it tracks which task produced each object so lost objects can be reconstructed from lineage — and it spills to disk automatically when full, so intermediates are bounded by disk rather than RAM.

When should I NOT use Ray?

Do not use Ray when the data fits comfortably on a single machine — pandas, Polars, or DuckDB will be simpler and cheaper than standing up a cluster. Do not migrate pure relational, shuffle-heavy ETL off a mature Spark platform just to use Ray; Ray wins on the Python/ML axis, not on relational shuffle, and rewriting working Spark for its own sake is pure risk. Do not use Ray Data as a substitute for governed warehouse SQL — a GROUP BY over Snowflake or BigQuery belongs in the warehouse. And do not reach for a cluster for tiny jobs or hard real-time request/response serving where coordination overhead outweighs the benefit. Naming this anti-pattern unprompted — "every technology has a workload envelope, and here's Ray's" — is the strongest senior signal in a Ray interview.

Practice on PipeCode

• Drill the data-processing practice library → for the map-reduce, fan-out, block-parallel, and batch-inference-pipeline problems that Ray Core and Ray Data workloads are built from.
• Rehearse on the ETL practice library → for the streaming map_batches, lazy-plan, partitioned-write, and large-scale scoring patterns senior interviewers probe when Ray comes up.
• Sharpen the throughput axis with the optimization practice library → for the ray.put broadcast, backpressure, batch-sizing, and cluster-cost trade-offs that separate a fluent Ray answer from a stumbling one.
• Add the streaming practice library → for the fault-tolerance, retry, and recovery scenarios that Ray's lineage-and-restart model is designed to survive.
• Stack the prerequisites against PipeCode's broader 450+ data-engineering catalogue to anchor the tasks-vs-actors-vs-datasets decision map against real graded inputs.

Lock in Ray decision muscle memory

Docs explain the API. PipeCode drills explain the decision — when Ray beats Spark, when a stage wants an actor instead of a task, when a batch-inference job needs backpressure before it OOMs, when the workload belongs in a warehouse instead of a cluster. Pipecode.ai is Leetcode for Data Engineering — pattern-first practice tuned for the production trade-offs senior data engineers actually face.

Practice data-processing problems →
Practice optimization problems →