The Streaming Batch Model for Efficient and Fault-Tolerant Heterogeneous Execution

TL;DR

The paper tackles the bottleneck of CPU-based data processing in GPU-centric ML workflows by introducing the streaming batch model, a partition-based, memory-aware approach that enables elastic, pipelined execution across heterogeneous resources. It presents Ray Data, a centralized, online scheduler that dynamically partitions data, streams partitions between heterogeneous operators, and adapts to memory pressure and cluster changes, while maintaining scalability and fault tolerance. Empirical results show Ray Data substantially outperforms traditional batch and streaming systems (up to $12\times$ throughput improvements) and delivers significant training speedups (e.g., $31\%$ on Stable Diffusion) by leveraging disaggregated CPU/GPU clusters. The work advances practical heterogeneous data processing for ML by combining memory-efficient partitioning, dynamic resource allocation, and robust recovery, with broad implications for RAG, video classification, and large-scale multimodal training.

Abstract

While ML model training and inference are both GPU-intensive, CPU-based data processing is often the bottleneck. Distributed data processing systems based on the batch or stream processing models 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 the streaming batch model, a hybrid of batch and streaming that enables efficient and fault-tolerant heterogeneous execution. The key idea is to use partitions as the unit of execution to achieve elasticity, but to allow partitions to be dynamically created and streamed between heterogeneous operators for memory-efficient pipelining. We present Ray Data, a streaming batch system that improves throughput on heterogeneous batch inference pipelines by 2.5-12$\times$ compared to traditional batch and stream processing systems. By leveraging heterogeneous clusters, Ray Data improves training throughput for multimodal models such as Stable Diffusion by 31% compared to single-node ML data loaders.

Figures (11)

• 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 model is replicated for distributed data-parallel training.
• Figure 2: Execution plans for video generation (\ref{['fig:model:dag']}). A0, A1, etc. are data records. Consecutive operators that require the same resources are fused. A1 is a large video that decodes to 3 video segments, B$1_1$, B$1_2$, B$1_3$. (a) Before execution, batch systems divide the dataset into evenly sized partitions, e.g., the box with A0 and A1. At run time, partitions are assigned resources dynamically. (b) Before execution, streaming systems assign each executor resources and a range partition for an operator (e.g., even keys for CPU0). At run time, executors stream records in dynamically sized batches. (c) Before execution, streaming batch systems divide the initial dataset into partitions. At run time, partitions can be dynamically split and dynamically assigned resources. Lighter partitions indicate ones that have not yet been materialized.
• Figure 3: Execution timelines for \ref{['fig:system']}. (a) Batch systems dynamically assign partitions to resources but execute one stage at a time, materializing all intermediate partitions. (b) Streaming systems pipeline records between executors in dynamically sized batches but statically assign resources and records to executors, which can reduce utilization. (c) Streaming batch systems dynamically size partitions and dynamically assign partitions to resources for efficient pipelining and load-balancing.
• 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) Streaming batch systems dynamically and automatically repartition outputs to reduce peak memory usage. This also allows the next operator to start in parallel with A1.
• 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 , while (b) schedules as soon as possible.