Ray Datasets - The Last Mile Before Model Training

Ray 1.8+ introduced Datasets to address the last mile before model training. Since I have limited hands-on experience with Ray, this article mainly organizes the concepts. For an introduction to Ray, see Ray - A Distributed Framework for Emerging AI Applications.

Current Situation

From the Ray Docs:

1
Datasets is not intended as a replacement for more general data processing systems. Its utility is as the last-mile bridge from ETL pipeline outputs to distributed applications and libraries in Ray.

Datasets is intended as a connector between ETL pipelines and distributed training, rather than a replacement for existing ETL engines such as Spark, Dask, or Mars.

Training Pipeline

A typical training dataflow looks like this:

There are three stages:

• ETL/data processing: General processing for the model family, including feature computation, feature-value handling, and invalid-data filtering. Data is read from and written to a distributed filesystem (DFS).
• Preprocessing: Model-specific work, such as partitioning by primary key or shuffling data. Different models have different requirements. This stage reads ETL output, produces training samples, and persists them to DFS.
• Model training: Read samples from DFS to train the model.

A complete training run usually reads the samples multiple times. Each pass is an epoch; each epoch divides the data into batches containing multiple samples. The pseudocode below illustrates this (see epoch vs. iterations vs. batch for the terminology):

1
for epoch->epoches:
2
    for batch->epoch.batches:
3
        for sample->batch.samples:
4
            model.feed(sample)

Problems

Performance:

• Serialization overhead: The arrows above imply five serialization/deserialization steps. Some can be optimized away, but this describes the general case.
• Storage cost: Using DFS for intermediate data adds storage consumption, I/O, and bandwidth for reads and writes.

Missing functionality:

• No per-epoch shuffling: See why data should be shuffled for machine-learning tasks. Here I mainly mean global shuffling, which is rarely performed in practice because it consumes considerable resources and time.

Usability:

• Switching languages and frameworks: Model-oriented data processing requires development in both Spark and TensorFlow.
• Heterogeneous resources: Data processing and training may require different resource types, such as CPUs and GPUs.
• Scheduling: Airflow or Kubeflow coordinates tasks, with resources sitting idle during transitions.

Datasets Capabilities

Datasets is fundamentally an abstraction for distributed datasets in a Ray cluster. It can use Ray’s Plasma in-memory object store, scheduling, heterogeneous resource support, and fault tolerance. On this foundation, Ray Datasets provides data loading, preprocessing, and pipelining:

Data Loader：

• Frequent disk materialization causes serialization overhead. Based on Apache Arrow, Datasets converts inputs from various formats into Arrow and stores them in Plasma. Data can then move quickly between Ray tasks without unnecessary serialization and deserialization.

Pre-Processing：

• It provides map, batch map, and filter. As a newer framework, it builds on established ideas; for example, map_batch() directly supports vectorized reads.
• It provides global operations such as sorting, random shuffling, and groupBy. The principle is simple: zero-copy data fetching over distributed storage.

Pipelining CPU/GPU：

The following diagram comes from PyData Global 2021 - Unifying Large Scale Data Preprocessing and ML Pipelines with Ray Datasets. Training may repeat loading, preprocessing, and inference, with different resources needed at each stage. GPUs must remain reserved during CPU preprocessing; releasing them risks another task taking them and interrupting training. This leaves GPUs idle and wastes resources.

Ray’s scheduling and heterogeneous resource support allow Datasets to achieve the diagram’s ‘With Pipelining’ behavior. A convenient DataPipeline API expresses the pipeline.

Summary

In practice, the separation between training and feature-engineering teams and technology stacks causes problems such as:

• Training even a simple model requires familiarity with several systems.
• Frequent data exchange and switching between languages and systems can become the main training bottleneck.

Training is only the final stage of the dataflow, and Ray already integrates with TensorFlow and PyTorch. Previously, most Ray data-processing APIs were low-level and required extensive custom code for production use. Datasets extends Ray Train along the dataflow, giving potential users another reason to adopt Ray.

As a processing engine before training, Ray Datasets replaces Spark in some scenarios. Its API abstractions also resemble Spark’s, making Datasets-versus-Spark a common first question. The discussion above shows capabilities that Spark lacks, though this is understandable given the frameworks’ different original use cases.

This article is a brief collection of notes from reading about Ray Datasets. The links below provide more detail.

References

[1]. https://docs.ray.io/en/latest/data/dataset.html#datasets-distributed-data-loading-and-compute [2]. https://docs.ray.io/en/latest/train/train.html#train-docs [3]. https://www.anyscale.com/blog/why-third-generation-ml-platforms-are-more-performant [4]. https://www.youtube.com/watch?v=wl4tvru9_Cg [5]. https://www.anyscale.com/blog/the-third-generation-of-production-ml-architectures