Ray Data: A dynamic and scalable system for heterogeneous data processing
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.