Ray Data Deep Dive: The AI-Era Distributed Engine Replacing Spark

Background & Positioning

Why Introduce Ray

In LLM, multimodal, and RAG workloads, the data side (preprocessing, embedding, batch inference, last-mile processing before training) is rapidly becoming the new bottleneck. Traditional solutions face three problems:

Spark/Flink: Excel at CPU batch processing but lack support for GPU heterogeneity, Python/Torch native objects, and model inference.

tf.data / PyTorch DataLoader: Single-node or embedded in training, difficult to scale independently, poor decoupling of preprocessing and training.

Custom pipelines: Ray Core can be hand-crafted but lacks a unified Dataset abstraction, shuffle, and fault tolerance.

Ray Data's goal: standardize, stream, and elastically scale the "last-mile" data pipeline between "CPU preprocessing" and "GPU training/inference".

This article covers:

- Ray Data principles and internal implementation
- Streaming Execution, Shuffle, Fault Tolerance
- Performance tuning, best practices

Introduction: Ray and Ray Data

2.1 Ray Ecosystem Overview

Ray consists of Ray Core (distributed primitives) and Ray AI Libraries (high-level libraries for AI workloads).

2.2 Ray Data Positioning: "Last-Mile Bridge"

Ray Data positions itself as the last-mile bridge from ETL to distributed training/inference.

2.3 Comparison with Other Solutions

2.4 Typical API Quick Look

import ray

# 1. Read
ds = ray.data.read_parquet(
    "s3://bucket/images/", override_num_blocks=256)

# 2. Vectorized transform (prefer map_batches over map)
def preprocess(batch):
    batch["image"] = [resize(x) for x in batch["image"]]
    return batch

ds = ds.map_batches(preprocess, batch_size=64, num_cpus=2)

# 3. GPU batch inference (Actor Pool holds model)
class CLIPEmbedder:
    def __init__(self):
        self.model = load_clip().cuda()
    def __call__(self, batch):
        batch["emb"] = self.model(batch["image"]).cpu().numpy()
        return batch

ds = ds.map_batches(
    CLIPEmbedder,
    batch_size=32,
    concurrency=4,  # 4 Actors
    num_gpus=1,     # each Actor uses 1 GPU
)

# 4. Write
ds.write_parquet("s3://bucket/embeddings/")

Principles

3.1 Ray Core Three Primitives

Ray Data is built entirely on Ray Core; understanding three primitives is essential:

Task: Stateless remote function decorated with @ray.remote. In Ray Data, stateless map, filter, read etc. are executed by TaskPool.

Actor: Stateful remote class instance, lifecycle managed by caller. In Ray Data, GPU batch inference and model residency scenarios are executed by ActorPool.

Object: Immutable object in Plasma shared memory, referenced by ObjectRef. In Ray Data, Blocks are passed as Objects between nodes with zero-copy.

3.2 Control Plane and Data Plane

Ray cluster logically splits into two planes, key to understanding scheduling and fault tolerance:

Control Plane: GCS + Scheduler embedded in each node's Raylet, responsible for metadata, task scheduling, resource accounting.

Data Plane: Plasma Object Store; large objects zero-copy via shared memory to same-node workers; cross-node via Object Store fetch.

3.3 Dataset = Blocks + Operation DAG

A Dataset is not actual data but a set of Block references + an operation DAG.

3.4 Lazy Execution and Materialize

Ray Data operations are lazy by default; only execution-triggering APIs actually launch the pipeline:

ds = ray.data.read_parquet("...")  # Lazy
ds = ds.map_batches(fn1)            # Lazy
ds = ds.map_batches(fn2)            # Lazy, DAG accumulates

# Execution-triggering APIs:
ds.show(5)                          # Streaming execution, fetch few rows
ds.take_batch(32)                   # Streaming execution, fetch one batch
ds.write_parquet("...")             # Streaming execution, write out
ds = ds.materialize()               # Fully materialize to Object Store (expensive)
ds.count()                          # Requires full scan, use cautiously

Note: count() , materialize() put intermediate results entirely into Object Store, easily triggering Spill or OOM; call on demand during training.

3.5 Streaming Execution

Unlike Spark's stage barriers, Ray Data uses streaming execution: downstream operators do not wait for upstream stage to fully complete; they produce and consume at Block granularity.

Advantage: I/O and compute, CPU and GPU can pipeline overlap, significantly improving overall throughput without intermediate disk writes.

Internal Implementation

4.1 Three-Layer Plan Model

Ray Data core source lives in python/ray/data/_internal/; execution path splits into three layers:

Logical Plan: User's read_parquet().map_batches(...) converted to logical operator tree.

Physical Plan: After optimization rules like Fusion and Pushdown, generates PhysicalOperator.

StreamingExecutor: The real "heart" that drives operators, manages Block queues, coordinates ResourceManager.

4.2 Block / RefBundle / Resource Manager

Block: Minimum data unit of Dataset, typically Arrow Table / Pandas DataFrame, default target size 1 MiB~128 MiB, controlled by target_min/max_block_size.

RefBundle: "Block reference bundle" passed between operators: a set of ObjectRef[Block] + metadata (row count, byte size, schema), avoiding direct data movement.

ResourceManager: Tracks each operator's current CPU/GPU/Object Store memory usage, feeds back to StreamingExecutor whether new Task/Actor can be scheduled.

4.3 Operator Fusion

Adjacent stateless operators are merged by operator_fusion rule, avoiding writing to Object Store at every step.

Effect: reduces Object Store writes, reduces serialization overhead; thus splitting two map_batches into multiple steps does not significantly slow down.

4.4 Actor Pool vs Task Pool

Comparison:

TaskPoolMapOperator : Suitable for stateless functions; concurrency: one Task per Block; model loading: reload every time; trigger: pass function to map_batches.

ActorPoolMapOperator : Suitable for stateful / re-initialization (model inference); concurrency: fixed / dynamic Actor pool; model loading: load once at Actor startup; trigger: pass Class to map_batches.

# ActorPool: use class + concurrency, model loads once
ds.map_batches(
    MyModelClass,
    concurrency=(2, 8),  # auto-scale between 2~8 Actors
    num_gpus=1,
    batch_size=64,
)

4.5 Shuffle Implementation Evolution

random_shuffle()

/ sort() / groupby() / repartition() all based on Shuffle.

Hash Shuffle v2 significantly reduces memory pressure for large dataset shuffles.

GPU scenarios can use rapidsmpf backend with NCCL / UCX, avoiding Host-Device copies.

4.6 Autoscaler Coordination

Ray Data's ResourceManager feeds each operator's "desired resources vs actual resources" to Cluster Autoscaler, driving cluster scale-up/down.

4.7 Fault Tolerance

Ray Data fault tolerance relies on two mechanisms:

Lineage Reconstruction: When Object lost (e.g., worker crash), Ray Core re-executes upstream tasks per task DAG to rebuild object.

Object Store Spill: When Plasma memory insufficient, Blocks automatically spill to disk ( /tmp/ray/session_*/spill).

Actor crash does not auto-restart by default; need explicit max_restarts , max_task_retries . GPU inference scenarios recommend non-zero values.

Performance Optimization

5.1 Optimizing Transforms

Vectorization + proper batch_size account for 90% of performance gains:

# Slow: per-row Python function
ds = ds.map(lambda row: {"x2": row["x"] * 2})

# Fast: vectorized (NumPy / Pandas / Polars)
def mul2(batch):
    batch["x2"] = batch["x"] * 2
    return batch
ds = ds.map_batches(mul2, batch_size=4096)

Enable Polars sort acceleration (large table sort / groupby):

import ray
ctx = ray.data.DataContext.get_current()
ctx.use_polars_sort = True

This switches internal sorting from PyArrow to Polars, faster for large tables.

5.2 Optimizing Reads

Ray Data automatically decides output Block count, heuristic order: start from default 200, constrained by min 1 MiB / max 128 MiB boundaries, then scale by CPU count (at least 2×CPU). read_op_min_num_blocks (default 200): initial Block count for read operator. target_min_block_size (default 1 MiB): too-small Blocks auto-merged. target_max_block_size (default 128 MiB): too-large Blocks auto-split. override_num_blocks: force specific output Block count.

# Scenario 1: many small files, merge
ds = ray.data.read_parquet(
    "s3://many_small_files/", override_num_blocks=64)

# Scenario 2: Parquet column pruning + predicate pushdown
ds = ray.data.read_parquet(
    "s3://table/",
    columns=["image_url", "label"],  # only needed columns
    filter=pa_ds.field("label") == 1,  # row-group level predicate
)

5.3 Memory and Spill

Block too large → Worker OOM; Block too small → task overhead high.

Object Store defaults to 30% of node memory; Spill to disk is costly.

Tuning: derive target_max_block_size from node memory and batch_size; avoid all-to-all full materialization ( random_shuffle, sort).

5.4 I/O and Network

Locality-aware scheduling: Map Tasks scheduled to nodes where Blocks reside, avoiding cross-node transfer.

Streaming Split for Trainer: ds.streaming_split(n) splits Dataset into n streams fed to n Trainers without full materialization.

5.5 Observability

# 1. Print stats after job
print(ds.stats())
# -> contains per-operator wall time, input/output block counts, CPU/GPU utilization

# 2. Real-time monitoring: Ray Dashboard
#    URL: http://<head>:8265 -> Jobs / Actors / Data

Key observation points: operator wall time proportion, Spill size, Actor utilization, per-Block P50/P99 processing latency.

5.6 Production Tuning Case Studies

5.6.1 Quantitative Relationship Between batch_size and OOM

Ray Data uses two memory types: Worker Heap Memory (process heap) and Object Store Memory (Plasma shared memory). Object Store spills to disk when full, but Heap Memory exceeding limit causes direct OOM crash.

Production-measured impact of batch_size on Heap memory (single Block 1000 rows, each row 125 KB):

Key conclusion: reducing batch_size not only lowers OOM risk, but in large-object scenarios can even be faster — because it reduces single deserialization and memory copy overhead. GPU inference must explicitly specify integer batch_size ; default batch_size=None treats entire Block as one batch, easily exploding GPU memory.

5.6.2 Object Store Memory Tuning

Object Store defaults to 30% of node memory, often insufficient in large-scale production. Adjust via environment variable:

# Increase Object Store memory ratio from default 30% to 50%
# e.g., node 32GB RAM: Object Store max from 9.6GB to 16GB
export ARNOLD_BYTEDRAY_ray_io_start_object_memory_ratio=0.5

Note: adjusting this parameter does not affect function.remote(memory=xxx) scheduling logic, but squeezes Worker process available memory. Balance based on total node memory and Worker actual needs; bigger is not always better.

5.6.3 Multimodal Op Fusion Optimization

In multimodal data pipelines, operator fusion is easily overlooked but yields huge gains. Example: "image download → Embedding inference":

Without fusion: download operator outputs image data → write to Object Store → inference operator reads from Object Store → deserialize → inference. Large objects (multi-MB images) incur massive cross-node transfer and serialization overhead.

With fusion: download and Embedding inference run inside same Actor; images passed directly in memory to model, completely avoiding network transfer and serialization, achieving zero-copy flow.

Implementation: merge multiple steps into a single Actor's call method, instead of splitting into multiple map_batches operators.

5.6.4 Avoiding Meaningless Shuffle Overhead

Multimodal business scenarios often lack key-based grouping/aggregation needs, but underlying Flink/DPS frameworks force multiple keyby operations. Ray Data natively avoids such meaningless shuffle overhead, a major benefit when migrating from Flink.

Best Practices & Troubleshooting

7.1 Actor vs Task Trade-offs

- Model inference, heavy init (>1s): use Actor (class + concurrency)
- Lightweight CPU transform (<100ms init): use Task (function)
- Need cross-Batch shared state (stats, cache): use Actor

7.2 Common Anti-Patterns

Multiple count(): each triggers full scan → cache result or call only at end.

Unnecessary materialize(): full data into Object Store causes Spill → use streaming iter.

All-to-all operation abuse: random_shuffle / sort full materialization → avoid if possible.

Single-node concentration: read/write limited to single machine → add override_num_blocks to distribute.

Frequent Spill: Block too large or Object Store too small → reduce batch or expand memory.

Output file fragmentation: too many small files → repartition before write.

7.3 Common Failure Quick Reference

Worker OOM: Block too large / batch too large → reduce target_max_block_size, batch_size.

Task Hang: Actor deadlock / downstream backpressure → check Dashboard Actor state, py-spy stack.

Low Throughput: no fusion / no vectorization → use map_batches; check stats() for operator latency.

HDFS Read Failure: Kerberos / permission / path → verify PyArrow HDFS config.

Repeated Spill: object_store_memory insufficient → expand memory or reduce Block size.

7.5 Production Exception Handling & Fault Tolerance Strategies

7.5.1 User Code Exceptions vs System Exceptions

Ray Data handles exceptions fundamentally differently; understanding this is key to production stability:

User code exception: Fail Fast — task fails immediately, throws UserCodeException, no retry. Typical: data format error, missing field, business logic bug.

System exception: Auto-retry — worker crash, node loss, etc., automatically retried on other nodes. Typical: node crash, preemptible resource revocation, bad GPU.

Best practice: when data cannot be processed, raise Exception to fail fast and expose issue, rather than swallowing errors. Ray Data wraps user code exceptions into UserCodeException with stack trace pointing to your business code, aiding debugging.

7.5.2 Multimodal Data Quality Gates

In OCR, ASR, Embedding preprocessing, extraction quality directly impacts downstream training. Production recommendations:

OCR/ASR outputs carry confidence; low-confidence samples filtered or sent for human review.

ASR handles filler words, segmentation errors, proper nouns; configure business vocabulary correction.

Multimodal retain original image/audio alignment with text; text serves only as searchable/trainable view.

Regular human sampling evaluation; never assume extraction results are inherently correct.

7.5.3 Spark vs Ray Data Selection Decision

Not all scenarios suit Ray Data; decision tree for quick selection:

Core operation: SQL, Join, Shuffle, dedup, aggregation → Spark; Model inference, Embedding, OCR/ASR preprocessing → Ray Data.

Compute resource: CPU-intensive → Spark; GPU-intensive / CPU-GPU hybrid → Ray Data.

Data format: Structured tabular primary → Spark; Multimodal objects (image/video/audio) primary → Ray Data.

Ecosystem dependency: Existing Spark/lakehouse stack → Spark; Existing Ray/Torch training stack → Ray Data.

Team familiarity: SQL/Spark tuning expertise → Spark; Python/deep learning expertise → Ray Data.

Best combo: structured heavy lifting (cleaning, dedup, Join, sampling) stays in Spark; model-inference multimodal operators (OCR, ASR, Embedding, rerank) go to Ray Data. Decide boundary by measured throughput and cost, not technology for technology's sake.