Ray: A Distributed Framework for Emerging AI Applications

Abstract: The next generation of AI applications will continuously interact with the environment and learn from these interactions. These applications impose new and demanding systems requirements, both in terms of performance and flexibility. In this paper, we consider these requirements and present Ray---a distributed system to address them. Ray implements a unified interface that can express both task-parallel and actor-based computations, supported by a single dynamic execution engine. To meet the performance requirements, Ray employs a distributed scheduler and a distributed and fault-tolerant store to manage the system's control state. In our experiments, we demonstrate scaling beyond 1.8 million tasks per second and better performance than existing specialized systems for several challenging reinforcement learning applications.

• The paper presents Ray, a unified system integrating task-parallel and actor-based models to meet diverse AI workload requirements.
• It demonstrates near-linear scalability with over 1.8 million tasks per second and high-throughput object storage delivering up to 15GB/s.
• Ray’s robust design supports reinforcement learning, outperforming specialized systems in Evolution Strategies and Proximal Policy Optimization tasks.

An Overview of Ray: A Distributed Framework for Emerging AI Applications

The paper "Ray: A Distributed Framework for Emerging AI Applications" by Moritz et al. introduces Ray, a distributed system designed to meet the rigorous performance and flexibility demands of the next generation of AI applications. The system integrates task-parallel and actor-based computations using a single dynamic execution engine, which is supported by a distributed scheduler and a fault-tolerant store managing control state. The authors demonstrate Ray's scalability and performance through extensive evaluations, particularly highlighting its efficacy for reinforcement learning (RL) tasks.

Motivations and Requirements

Contemporary AI applications increasingly require systems capable of interacting dynamically with their environments, generating and evaluating actions in millisecond intervals, and handling heterogeneous workloads involving both CPU and GPU resources. Existing systems typically fall short in addressing these needs for fine-grained, dynamic computations and stateful interactions. This gap is particularly significant for reinforcement learning workloads, which demand tight coupling of training, simulation, and serving phases.

Design and Architecture

Ray is structured to meet these ambitious requirements through several key design principles:

1. Unified Interface: Ray offers a programming model that combines both task-parallel and actor-based abstractions. Tasks are used for stateless computations, enabling efficient load balancing and fault recovery, while actors cater to stateful computations, necessary for model training and maintaining shared mutable state.
2. Dynamic Execution: Ray's computation is modeled as a dynamic task graph where tasks and actor methods form nodes connected by data, control, and stateful edges. This approach allows for real-time task creation and dependency management, facilitating the development and execution of complex RL workflows.
3. Distributed Scheduler and Control Store: To handle the expected high throughput, Ray employs a bottom-up distributed scheduler coupled with a global control store (GCS). The GCS enables scalable storage of control state and metadata about tasks and objects, supporting lineage-based fault tolerance and efficient task scheduling.

Performance and Scalability

The paper presents strong empirical results to substantiate the system's design:

• Task Throughput: Ray achieves over 1.8 million tasks per second on a 100-node cluster, showing near-linear scalability. This performance is crucial for workloads that involve fine-grained tasks typical in AI applications.
• Object Store: The in-memory distributed object store in Ray provides high throughput (up to 15GB/s) and IOPS, ensuring efficient data exchange between tasks, particularly relevant for high-frequency, low-latency operations.
• Fault Tolerance: Ray's fault tolerance is validated through both task and actor reconstruction experiments. The system's ability to handle node failures seamlessly while maintaining task throughput underscores its robustness for production environments.

Application Insights

The utility of Ray is further highlighted through its performance in RL workloads:

• Evolution Strategies (ES): Ray's ES implementation scales effectively to 8192 cores, achieving more than twice the performance speed of specialized systems without requiring substantial custom optimizations.
• Proximal Policy Optimization (PPO): The PPO implementation in Ray outperforms specialized MPI-based implementations, achieving better performance with fewer resources. This signifies Ray's capability to handle complex, heterogeneous workloads efficiently.

Implications and Future Directions

Ray's architecture and performance make it a compelling choice for building and scaling the next generation of AI applications, particularly those involving dynamic interactions and stateful computations. The system's design principles could inspire future frameworks to adopt similar scalable and fault-tolerant approaches, particularly the use of a global control store and decoupled scheduling mechanisms.

Potential future developments may include enhancing the API to support higher-level abstractions, profiling-based scheduling optimizations, and more sophisticated garbage collection strategies for the GCS. Additionally, extending Ray to natively support more complex dataflows and distributed objects could further broaden its applicability.

In summary, Ray represents a significant step forward in distributed systems for AI, offering the flexibility, scalability, and performance necessary to meet the demands of modern applications. Its innovative integration of task and actor-based models provides a versatile platform poised to facilitate advancements in AI research and deployment.