Part II - Applications and libraries

Ray Data: Blocks, Streaming Execution, and ML Dataflow

Lesson 08 made Ray Tune a control plane that runs many trials and kills the losers early — but every trial still needs data, and so does every training run in lesson 10. The map-shuffle-reduce shape you built by hand in lesson 05 is the right primitive, yet writing it raw means you decide block sizes, materialization, and GPU feeding yourself. Ray Data packages that shape into a dataset abstraction that streams data larger than cluster RAM through a chain of transforms and hands batches to a trainer without ever stalling the GPU — if you tune it right. This lesson is about what a dataset actually is underneath, and the two knobs (compute strategy, block size) that decide whether your expensive GPUs sit idle.

1 · A dataset is a sequence of blocks

A Ray Dataset is not a single in-memory table. It is a logical sequence of blocks, where a block is a contiguous chunk of rows held as an Apache Arrow table (a columnar in-memory format) and stored as one immutable value in the object store (lesson 04). A 200 GB Parquet dataset read with default settings might become, say, 1,000 blocks of ~200 MB each, scattered across the object stores of every node in the cluster.

Why blocks, and not rows or one big table? Two reasons, and they are the whole design.

So the block is doing the same job the partition does in Spark or the chunk does in Dask: it is the atom of parallelism and the atom of placement. Everything in this lesson is a consequence of choosing the right number and size of those atoms.

2 · Streaming execution — bigger than RAM still flows

The naive mental model of a data pipeline is "load the dataset into memory, then map over it, then map again." That model breaks the moment the dataset is larger than cluster RAM: you cannot hold 2 TB of features in 512 GB of object-store memory, so a bulk-materialize-then-transform design either spills catastrophically (lesson 04's spill cliff) or crashes.

Ray Data instead uses streaming execution (pipelined execution): it does not materialize the whole dataset between operators. Operators in the chain — read → map_batches → map_batches → consumer — run concurrently, and blocks flow through them one at a time. At any instant only a working set of blocks is resident; a block is read, transformed, consumed, and its memory freed while the next block is still being read. The dataset's total size can exceed cluster RAM because the cluster never holds all of it at once.

The mechanism that keeps a streaming pipeline from blowing up is backpressure: a control signal that throttles a fast upstream operator when a slow downstream operator cannot keep up. If the reader is producing blocks faster than featurize can consume them, the resident blocks would pile up in the object store until it spills. Backpressure pauses the reader until the buffer drains, capping the working set. The consequence to internalize: in a streaming pipeline, the slowest operator sets the throughput of the whole chain, exactly like a factory line — which is why GPU starvation (§4) is a backpressure problem in disguise.

3 · map_batches and the compute-strategy choice

The workhorse transform is map_batches(fn): it hands your function a batch — a slice of one block, as a dict of NumPy arrays or a pandas frame — and you return the transformed batch. Batching is what makes it fast: a vectorized NumPy or framework call over 1,024 rows is far cheaper than 1,024 per-row Python calls (lesson 02's granularity argument, applied to data). Prefer map_batches over a per-row map for anything that vectorizes.

The decision that actually matters is the compute strategy: should each batch run as a stateless task or on a pool of stateful actors?

You select actors by passing a class (not a function) plus compute=ray.data.ActorPoolStrategy(...). Ray Data builds the actors, calls __init__ once each, then streams batches through their __call__. This is the batch-inference pattern: load the model once per actor, score many blocks.

import ray
from ray.data import ActorPoolStrategy

ds = ray.data.read_parquet("s3://bucket/images/")   # → a Dataset of blocks

# cheap, stateless: run as tasks (the default)
def to_float(batch):                                 # batch = dict of np arrays
    batch["image"] = batch["image"].astype("float32") / 255.0
    return batch

# expensive init: load the model ONCE per actor, reuse across batches
class Classifier:
    def __init__(self):
        self.model = load_model().cuda()             # ~20 s — paid once per actor

    def __call__(self, batch):
        import torch
        with torch.no_grad():
            batch["pred"] = self.model(
                torch.as_tensor(batch["image"]).cuda()).argmax(1).cpu().numpy()
        return batch

scored = (
    ds.map_batches(to_float, batch_format="numpy")              # tasks
      .map_batches(Classifier,                                   # actors:
                   compute=ActorPoolStrategy(size=4),            # 4 long-lived workers
                   num_gpus=1,                                    # one GPU each
                   batch_size=1024)
)

4 · The producer/consumer link to Ray Train — don't starve the GPU

Ray Data exists largely to feed Ray Train (lesson 10). The relationship is producer/consumer: the data pipeline (CPU work — read, decode, augment) produces batches; the training loop (GPU work — forward/backward) consumes them. The hard rule of any such pipeline: the GPU can only run as fast as the slowest upstream operator can feed it. If preprocessing cannot deliver a batch in the time the GPU takes to consume one, the GPU idles — and an idle A100 is the most expensive idle resource in your system.

The fixes follow directly from the math, and all of them are knobs Ray Data gives you:

The handoff itself: in training you call iter_batches (or iter_torch_batches) on the dataset, and Ray streams shards to each training worker. Lesson 10 shows the sharded side — train.get_dataset_shard gives each DDP rank its slice; here the point is only that streaming + prefetch + enough producers is what keeps that shard's GPU fed.

# consumer side — stream batches to the trainer, prefetching to hide latency
for batch in scored.iter_torch_batches(batch_size=256, prefetch_batches=4):
    loss = train_step(batch)          # 50 ms GPU; prefetch keeps the next batch ready

5 · Block size, the CPU/GPU balance, and the integration boundary

Block size is the single tuning knob with the most leverage, and it is a genuine trade-off with a wrong answer on both ends.

The other balance is CPU-preprocess vs GPU-train: spend too little on preprocessing parallelism and you starve the GPU (§4); spend too much and CPU nodes sit idle while the GPU is the bottleneck. The target is to provision just enough producers that the GPU is the bottleneck — because the GPU is what you are paying the most for, you want it to be the constraint, not the data path.

Finally, the integration boundary. Ray Data reads and writes the formats your data already lives in — read_parquet, read_csv, from_pandas, from_spark, from_huggingface — and is the glue that moves data from those systems into Python ML transforms. It is explicitly not a Spark-SQL replacement: it has no mature query optimizer, no SQL engine, no decades of enterprise ETL hardening. The honest division of labor — do heavy SQL ETL in Spark/Snowflake, then hand the result to Ray Data for the last-mile ML preprocessing and GPU feeding — is the subject of lesson 15's ecosystem positioning. Treat each integration as a contract: schema, partitioning, and ordering must stay explicit across the boundary.

6 · Failure modes & checklist

Checkpoint exercise

Where this points next

We can now keep a GPU fed; lesson 10 puts something on the other end of that pipe. Ray Train takes the streamed shards and runs N worker actors — one DDP rank per GPU — each holding a model replica and exchanging gradients via allreduce (the collective shape from lesson 05, and from the system_ml track). The open question this lesson leaves: once data feeding is solved, what bounds the speedup of adding more training workers, and what happens when the model itself no longer fits on one GPU? That is where data-parallel training meets the memory wall.

Interview prompts

• What is a block, and why is it the right unit? (§1 — a chunk of rows as an Arrow table stored as one object-store value; it is both the atom of parallelism (one task per block) and the atom of placement (locality, spilling, lineage).)
• How does a dataset bigger than cluster RAM run at all? (§2 — streaming execution pipelines blocks through operators concurrently; only a bounded working set is resident at once, so total size can exceed RAM; backpressure caps the buffer.)
• When do you use ActorPoolStrategy instead of the default tasks? (§3 — when the transform has expensive one-time init, e.g. loading a model on a GPU; actors construct the callable once and reuse it, amortizing init across many batches (~155× on the worked job) instead of paying it per batch.)
• Your GPU sits at 60% utilization during training — why and how do you fix it? (§4 — the data pipeline can't supply a batch within the GPU's step time, so the GPU idles between steps; add CPU preprocessing producers until producer-time ÷ workers ≤ GPU-step time, and prefetch to hide latency.)
• What goes wrong with blocks that are too big? Too small? (§5 — too big: stragglers and memory pressure that triggers spilling; too small: per-block scheduling overhead dominates the actual transform. Target ~128–256 MB.)
• Should you do your heavy SQL joins in Ray Data? (§5 — no; Ray Data has no query optimizer or SQL engine. Do heavy ETL in Spark/Snowflake and hand the result to Ray Data for last-mile ML preprocessing and GPU feeding, per lesson 15.)
• Why is calling materialize() between every operator an anti-pattern? (§5 — it defeats streaming, forces the whole dataset into the object store, and triggers spilling; stream unless you must inspect or reuse a stage.)