Ray Data: A dynamic and scalable system for heterogeneous data processing

Frank Sifei Luan; Ron Yifeng Wang; Yile Gu; Ziming Mao; Charlotte Lin; Amog Kamsetty; Hao Chen; Cheng Su

cs.DCJan 16, 2025

Ray Data: A dynamic and scalable system for heterogeneous data processing

Authors: Frank Sifei Luan, Ron Yifeng Wang, Yile Gu, Ziming Mao, Charlotte Lin, Amog Kamsetty, Hao Chen, Cheng Su,

Organizations: UC Berkeley · University of Washington · Anyscale

Abstract

While ML model training and inference are both GPU-intensive, CPU-based data processing is often the bottleneck. Existing distributed data processing systems assume homogeneous resource requirements. They excel at CPU-based computation but either under-utilize heterogeneous resources or impose high overheads on failure and reconfiguration. We introduce Ray Data, a scalable system for heterogeneous workloads spanning batch inference and ML training. The key idea is to combine dynamic task creation for memory efficiency with dynamic task placement for operator elasticity. Making both decisions dynamically allows the system to adapt to heterogeneous pipelines where the memory and compute requirements per operator are unknown and variable over time. Ray Data addresses two challenges: (1) leveraging dynamic task creation for memory-efficient pipelining, and (2) designing a task placement policy that jointly allocates heterogeneous compute and memory resources to maximize throughput while limiting memory usage. Ray Data improves throughput on heterogeneous batch inference pipelines by 2.6-12 $\times$ compared to traditional data processing systems. By leveraging heterogeneous clusters, Ray Data improves end-to-end training time for multimodal models such as Stable Diffusion by 31% compared to popular ML data loaders.

Figures & tables

Panel of Figure 1 — Figure 1: Logical dataflow graphs representing heterogeneous ML applications. Nodes are operators. (a) A pipeline for video or image generation. (b) A multimodal training pipeline. The UNet model is replicated for distributed data-parallel training.

Panel of Figure 2 — Figure 2: Execution plans for video generation ( Figure 1(a) ). A0, A1, etc. are data records. Consecutive operators that require the same resources are fused, e.g., load and decode are fused into A. A1 is a large video that decodes to 3 video segments, B $1_{1}$ , B $1_{2}$ , B $1_{3}$ . (a) Before execution, the system creates tasks by dividing the dataset into evenly sized partitions, e.g., the box with A0 and A1. At run time, tasks are assigned resources dynamically. (b) Before execution, the system assigns operators to executors and a range partition per operator instance (e.g., even rows for CPU0). At run time, executors dynamically buffer and batch intermediate records. (c) Before execution, Ray Data divides only the initial dataset into partitions. At run time, each task is dynamically assigned resources and can dynamically produce output partitions. Lighter partitions indicate ones that have not yet been materialized.

Figure 3: Execution timelines for Figure 2 . GPU1 is added to the cluster during execution to show how each system handles reconfiguration. (a) Tasks are dynamically assigned to resources, so GPU1 can be used immediately. However, tasks execute one stage at a time, and the number of partitions per task must be set before run time. (b) Executors locally decide how many records to batch together into a task for better pipelining. Task assignment to executors is static, so cluster reconfiguration can require expensive synchronization. (c) Ray Data dynamically creates and assigns tasks, and tasks dynamically produce partitions, producing even-sized partitions.

System	Dyn task creation	Dyn task placement	Recovery	Rollback	Heterogeneous clusters
Batch systems [ 7 , 2 , 44 , 15 ]	$\times$	$\checkmark$	Logical log	Task	$\checkmark$
Microbatch systems [ 45 , 40 ]	$\times$	$\checkmark$	Checkpt+Logical log	Task	$\checkmark$
Naiad [ 25 ] , Flink [ 6 ]	$\checkmark$	$\times$	Checkpt	Epoch	$\checkmark$
MillWheel [ 1 ] , Kafka [ 19 ]	$\checkmark$	$\checkmark$	Physical log	Record	$\checkmark$
PyTorch DL [ 31 ]	$\checkmark$	$\times$	None	Job	$\times$
tf.data [ 26 ]	$\checkmark$	$\checkmark$	Checkpt	Epoch	$\times$

Table 1: Systems for heterogeneous data processing. Dynamic task creation : the system decides how to batch records between operators. Dynamic task assignment : the system dynamically assigns tasks to resources. Recovery : the recovery approach. Rollback : the minimum amount to re-execute upon failure or cluster reconfiguration. Heterogeneous clusters : whether the system can scale using heterogeneous nodes.

Panel of Figure 4 — Figure 4: Arrows represent data dependencies. (a) Batch systems allow repartitioning, but it must be specified by the user before execution. Execution across operators cannot be pipelined. (b) A system with dynamic task creation and placement automatically repartitions outputs to reduce peak memory usage. This also allows the next operator to start in parallel with A1.

Panel of Figure 5 — Figure 5: Scheduling under memory pressure. Green represents CPU executors’ local memory capacity (1 partition per CPU, 2 total). Pink represents the system’s shared memory capacity for intermediate data (1 partition total). Arrows represent data dependencies. (1) indicates a period where executors must stall and buffer outputs locally until there is sufficient space in shared memory. (a) further stalls CPU0 during (2) until there is enough shared memory to schedule A3 , while (b) schedules A3 as soon as possible.

Method	Description
read	Read items from files.
map	Transform each item.
map_batches	Transform a batch of items. Useful for controlling GPU batch size.
flat_map	Transform each item and flatten the results.
filter	Return items that match a predicate.
limit	Truncate to the first N items.

Table 2: A subset of the Ray Data Dataset API. The bottom four are consumption APIs that trigger execution, while the others are lazy.

Figure 6: Ray Data architecture overview. Ray Data executes as a Ray library. The Ray Data scheduler maintains the partition metadata ( Figure 2(c) ) and dispatches tasks to Ray workers (dashed arrows).

Panel of Figure 7 — Figure 7: (a) RAG job completion time. Ray Data-staged runs stages sequentially in one process; its Retrieval run time is the optimal JCT if CPU is the bottleneck. (b) Video classification throughput comparison of batch (Spark and Ray Data-staged ), streaming (Flink and Ray Data-static ), microbatch ( Ray Data-microbatch ), and dynamic task creation and assignment ( Ray Data-dynamic ) systems. (c) Fault tolerance comparison for video classification on a heterogeneous cluster with 1 GPU node and 1 CPU-only node. Executor ( blue ) or node ( red ) failure is injected at $t$ =15min. For node failure curves, the node is added back to the cluster at $t$ =30min. The dashed line is a variant of Ray Data that emulates checkpoint-and-restore. This strategy leads to job downtime after every cluster change and longer completion time overall. Unmodified Ray Data (solid) smoothly handles executor failure and node failure and addition.

Panel of Figure 8 — Figure 8: (a) Training ResNet-50, loading data from local disk vs. cloud storage (S3). All experiments run on 1 g5.2xlarge node (1 NVIDIA A10G GPU, 8 vCPUs), except for Ray Data (S3), which uses an additional m7i.2xlarge (8 vCPUs) node to scale out S3 loading. (b) Run time and cost for one epoch of Stable Diffusion pre-training.

Figure 9: Run times for systems under different memory limits on a synthetic benchmark. Grey means the system is unable to finish due to OOM. Ray Data(-Part.) means Ray Data without dynamic repartition (§ 4.2.1 ). Ray Data(-Adapt.) means Ray Data without adaptive memory-aware scheduling (§ 4.3.1 ).

Panel of Figure 10 — Figure 10: (a) Effect of partition sizes on throughput in Ray Data. (b) Scalability experiment showing throughput over number of nodes compared to using raw Ray tasks or generator tasks ( Ray Generator ).

Heterogeneous Computing

Heterogeneous computing studies how to coordinate different processors, accelerators, and memory systems to improve performance, efficiency, and cost.

See all 80 papers

BatchGen: An Architecture for Scalable and Efficient Batch Inference
Jun 19, 2026Tairan Xu, Leyang Xue, Zhan Lu +7Efficient Neural Network Inference Efficient Inference
University of Edinburgh · Tencent
Understand and Accelerate Memory Processing Pipeline for Large Language Model Inference
Mar 30, 2026Zifan He, Rui Ma, Yizhou Sun +1LLM Inference LLM Inference Acceleration
Department of Computer Science, University of California, Los Angeles · Microsoft Research Asia
ETH: Data-Balanced Heterogeneous Pipeline Parallelism for Training Large Models
Oct 15, 2026Mianjie Yu, Zhongmin Zhao, Zizhao Mo +9Distributed Training Pipeline Parallelism
University of Macau · Great Bay University · Sun Yat-sen University +1

Jun 19, 2026J/K move · Enter open · S save

BatchGen: An Architecture for Scalable and Efficient Batch Inference

Tairan Xu, Leyang Xue, Zhan Lu, Jinfu Deng+6

University of Edinburgh · Tencent

Efficient Neural Network Inference Efficient Inference Distributed Inference

Open paper AskPDF arXiv

Batch inference has become a central mode of AI computation, yet existing inference engines still rely on execution models designed for interactive serving. When scaled to millions of sequences, batch workloads reveal two fundamental requirements: the ability to handle extreme inter- and intra-sequence load variation that emerges only at runtime, and the ability to sustain high utilization across large fleets of GPUs. Existing systems fail to meet these requirements, losing substantial fractions of achievable throughput. We introduce a new architectural foundation for batch inference: the sequence coroutine compute model, which represents each sequence as a fine-grained, event-driven coroutine. This model exposes expressive primitives that allow the runtime to reorganize work dynamically, enabling larger expert-level batches, mitigating stragglers, reallocating work across devices, and maintaining utilization even on cost-effective or memory-constrained GPUs. Building on this abstraction, we implement BatchGen, a production-ready system that uses the coroutine model at cluster scale. On a 128-GPU cluster, BatchGen reduces batch completion time by up to

2.3\times

, and on memory-constrained accelerators it outperforms the strongest offloading baseline by up to

9.6\times

. We will open-source BatchGen at https://github.com/batchgen-project/batchgen

Abstract

Heterogeneous Computing

Explore similar work

BatchGen: An Architecture for Scalable and Efficient Batch Inference

Understand and Accelerate Memory Processing Pipeline for Large Language Model Inference

ETH: Data-Balanced Heterogeneous Pipeline Parallelism for Training Large Models