Source-linked AI summary
Ray: A Distributed Framework for Emerging AI Applications
Philipp Moritz, Robert Nishihara, Stephanie Wang, Alexey Tumanov, Richard Liaw, Eric Liang, Melih Elibol, Zongheng Yang, William Paul, Michael I. Jordan, Ion Stoica
TL;DR
Emerging AI applications need systems that continuously interact with environments while handling demanding performance and flexibility requirements. Ray unifies task-parallel and actor-based computation in a dynamic execution engine with distributed control-state management. The system scales to 1.8 million tasks per second and outperforms specialized systems on representative reinforcement-learning applications.
Problem
Emerging AI applications require continuous environment interaction, exploration, and learning, but existing systems do not jointly support their simulation, training, and serving workloads.
Method
Ray unifies task-parallel and actor-based computations in a dynamic task graph, using a sharded metadata store and bottom-up distributed scheduling.
Results
1.8 million tasks per second: Ray demonstrates linear scalability and matches or outperforms custom systems on representative reinforcement-learning applications.
Takeaways & Limitations
Ray’s programming flexibility supports application-level optimizations across tightly coupled training, simulation, and serving workloads.
Takeaways & Limitations
Ray’s generality makes specialized scheduling optimizations difficult and requires garbage-collection policies to bound lineage storage costs.
Abstract
from arXiv · showhide
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.
1 Introduction
Emerging AI applications require systems that handle dynamic, heterogeneous reinforcement-learning workloads spanning simulation, training, and serving. Ray addresses this need with unified task and actor abstractions over a scalable, fault-tolerant execution engine.
- Emerging AI applications must react to changing environments, learn from delayed feedback, and explore action sequences over long-term goals.
- RL systems require fine-grained, heterogeneous, and dynamic execution, including millions of tasks per second with millisecond-level latencies.Workloads can vary from milliseconds to hours and may use CPUs, GPUs, or TPUs.
- Existing Big Data, task-parallel, streaming, and distributed deep-learning frameworks do not jointly support RL simulation, training, and serving requirements.
- Stitching together specialized systems is impractical because tightly coupled components create a substantial systems-engineering burden.
- Ray provides a general-purpose framework for RL simulation, training, and serving through unified task-parallel and actor-based computations.Tasks support dynamic load balancing and failure recovery, while actors support stateful computations such as model training.
- Ray is not intended to replace general-purpose serving or data-parallel frameworks with broader deployment features and APIs.
- Ray stores control state in a sharded metadata store and uses bottom-up distributed scheduling to achieve scalability and fault tolerance.The metadata store maintains computation lineage and object locations, while the distributed scheduler supports high task throughput.
2 Motivation and Requirements
Reinforcement learning alternates policy evaluation through environment interaction with policy improvement from collected trajectories. Its systems must support diverse computations, hardware, statefulness, and dynamic execution across training, serving, and simulation.
- An RL agent repeatedly interacts with an environment to learn a policy that maps environment states to actions and maximizes reward.
- Policy learning alternates evaluating the current policy through trajectories with improving it using those trajectories.Evaluation generates sequences of state-reward tuples; improvement updates the policy toward higher reward.
- RL frameworks must efficiently support training, serving, and simulation as interconnected workloads.Training updates policies, serving renders actions, and simulation evaluates policies.
- Requirements: RL computations range from milliseconds to hours and may require heterogeneous CPUs, GPUs, or TPUs.
- Requirements: RL applications need both stateless computations for simulation and data processing and stateful computations for parameter servers and repeated training.
- Requirements: Dynamic execution is required because computation completion order and simulation results can determine which computations occur next.
- Requirements: A framework must handle millions of tasks per second; a 200-node example reaches 1.28M tasks/sec for 5ms single-core tasks.
3 Programming and Computation Model
Ray unifies task-parallel and actor-based programming in a dynamic task graph, combining stateless remote functions with stateful actors. Its API and graph semantics support heterogeneous execution, dependency tracking, parallelism, and lineage-based recovery.
- Programming Model: Ray provides both task-parallel and actor abstractions over a dynamic task graph computation model.The graph evolves during execution, and remote functions or actor methods trigger when inputs become available.
- Tasks and Actors: Tasks execute stateless remote functions, return futures immediately, and use futures to express parallelism while capturing data dependencies.Remote functions operate on immutable objects and are expected to be stateless and side-effect free, simplifying fault tolerance through re-execution.
- Tasks and Actors: Actors execute remotely invoked methods serially on stateful workers, allowing internal state updates and actor handles to be passed to other computations.Actor methods return futures like tasks, but operate on persistent state shared across method invocations.
- Programming API: Ray adds ray.wait(), resource requirements, and nested remote functions to support heterogeneous durations, heterogeneous resources, and flexible distributed execution.ray.wait() returns the first k available results instead of waiting for all results, while nested calls support distributed invocation and scalability.
- Computation Model: The computation graph represents data objects and task invocations with data and control edges, while actor methods additionally use stateful edges to preserve invocation order and lineage.Control edges represent nested calls; stateful edges connect successive methods on the same actor and support reconstruction of lost outputs.
4 Architecture
Ray separates its application and system layers, combining task and actor abstractions with scalable scheduling, data management, and fault tolerance. Its global control store centralizes durable control state while distributed components remain scalable and largely stateless.
- Application and system layers: Ray’s application layer exposes drivers, stateless workers, and stateful actors, while the system layer provides scalable scheduling and data management.Workers execute remote functions without persistent local state; actors execute state-dependent methods serially.
- Global control store: The global control store maintains Ray’s control state using sharding for scale and per-shard chain replication for fault tolerance.The GCS stores lineage and object-directory metadata needed for dynamically spawning millions of tasks per second.
- Fault tolerance and implementation: The GCS enables stateless system components that can restart from shared lineage while the scheduler and object store scale independently.The implementation uses sharded, chain-replicated Redis-backed GCS tables and event-driven local and global schedulers.
- Distributed scheduler: Ray’s bottom-up scheduler tries local scheduling first and forwards tasks to the global scheduler only when local capacity or resource constraints require it.The global scheduler selects feasible nodes using estimated queueing and remote-input transfer time, with replicas available if it becomes a bottleneck.
- Distributed object store: Ray’s in-memory object store uses shared memory for zero-copy local data sharing and replicates remote inputs before execution.Objects are immutable, may be evicted to disk with LRU, and required objects can be recovered through lineage re-execution after node failure.
- Fault tolerance and implementation: Ray’s object store does not support distributed objects, so each object must fit on a single node unless applications decompose it into collections of futures.This limitation applies to large matrices, trees, and other objects that would otherwise span nodes.
5 Evaluation
The evaluation examines whether Ray meets its latency, scalability, and fault-tolerance targets, and measures overheads of distributed primitives. It also compares Ray with specialized systems and studies its advantages over custom RL systems.
- Evaluation questions: The evaluation tests Ray’s latency, scalability, and fault tolerance against the requirements introduced earlier.These measurements are reported in Section 5.1.
- Evaluation questions: The evaluation measures overheads for distributed primitives such as allreduce implemented through Ray’s API.This asks whether Ray’s programming interface can support communication-intensive primitives efficiently.
- Evaluation questions: The RL evaluation compares Ray with specialized systems for training, serving, and simulation and examines advantages over custom systems.The study addresses both comparative performance and the benefits of using Ray for RL applications.
5.1 Microbenchmarks
Microbenchmarks evaluate Ray’s locality, scalability, object storage, control-store fault tolerance, task and actor recovery, and allreduce performance. The results show strong scaling, low-latency recovery, and competitive distributed communication performance.
- Locality and scalability: Ray’s locality-aware task placement avoids the 1–2 orders of magnitude latency increase observed for non-local placement with 10–100MB inputs.The shared object store lets developers use tasks when locality matters, while actors cannot move computation to large remote objects.
- Object store: Single-client object-store write throughput exceeds 15GB/s for large objects, while small-object performance is measured using IOPS.Large-object creation is dominated by memcpy, whereas small-object overheads arise mainly from serialization and client–store IPC.
- Locality and scalability: Ray exceeds 1 million tasks per second at 60 nodes and scales beyond 1.8 million tasks per second at 100 nodes.The system processes 100 million tasks in 54 seconds, with near-perfect linear throughput scaling on the evaluated workload.
- GCS fault tolerance: GCS reconfiguration keeps maximum client-observed latency under 30ms during chain-member failure and recovery.The experiment starts with two replicas, kills one member, adds a replacement, and restores two-way replication.
- Task and actor recovery: Lineage recovery reconstructs lost task dependencies after node removal, while actor checkpoints reduce recovery to 500 method re-executions instead of 10k.The task workload returns to original throughput after nodes are added back, and checkpointing bounds actor reconstruction time.
- Allreduce: Ray completes 100MB allreduce across 16 nodes in ∼200ms and 1GB in ∼1200ms, outperforming OpenMPI by 1.5× and 2× respectively.The evaluation attributes the performance to multiple network-transfer threads and shows that extra scheduling latency substantially reduces allreduce performance.
5.2 Building blocks
Ray provides a unified framework for training, embedded serving, and simulation, matching specialized systems in training and improving serving throughput while handling heterogeneous simulations.
- Distributed Training: Ray uses actors for data-parallel synchronous SGD, with model weights synchronized through allreduce or a parameter server implemented on Ray’s API.
- Distributed Training: Ray matches Horovod and remains within 10% of distributed TensorFlow by expressing application-level optimizations such as pipelined gradient computation and transfer.
- Model Serving: Ray targets embedded serving to simulators in the same dynamic task graph, unlike Clipper’s external-client serving model.
- Model Serving: Ray achieves an order of magnitude higher throughput than Clipper for colocated embedded serving because of low-overhead serialization and shared memory abstractions.
- Simulation: Ray and MPI both scale for heterogeneous simulations, but Ray improves utilization by gathering variable-length simulation results concurrently rather than using BSP-style rounds and barriers.
5.3 RL Applications
Ray’s flexible programming model tightly couples reinforcement-learning training, simulation, and serving, enabling performance that matches or exceeds specialized implementations with less engineering effort.
- RL Applications: Ray’s dynamic task graph supports application-level optimizations across tightly coupled training, simulation, and serving that are difficult to incorporate into one-off systems.
- Evolution Strategies: Ray’s evolution-strategies implementation scales to 8192 cores and reaches a median completion time of 3.7 minutes, more than twice as fast as the best published result.
- Evolution Strategies: Ray required modifying only 7 lines to parallelize a serial evolution-strategies implementation, while the reference system used several hundred lines for worker communication.
- PPO: Ray’s PPO implementation outperforms the optimized MPI implementation in all experiments while using a fraction of the GPUs.
- PPO: Ray’s resource heterogeneity decreases PPO’s cost by a factor of 4.5 by scheduling CPU-only tasks on cheaper high-CPU instances.
6 Related Work
Ray combines dynamic task graphs, actors, distributed control state, and scalable scheduling to address limitations of dataflow, machine-learning, and actor systems.
- Dynamic Task Graphs: Ray extends dynamic task graphs with actors, supporting stateful computation needed for distributed training and serving.
- Dataflow Systems: Dataflow systems such as MapReduce and Spark are restrictive for fine-grained dynamic simulation, while Dryad lacks dynamic task graphs and actor abstractions.
- Machine Learning Frameworks: TensorFlow and MXNet efficiently train static linear-algebra DAGs but provide limited support for tightly coupling training, simulation, and embedded serving.
- Machine Learning Frameworks: OpenMPI can achieve high performance but requires explicit coordination for heterogeneous dynamic task graphs and explicit fault-tolerance handling.
- Actor Systems: Ray differs from Orleans and Akka through stronger support for recovery from data loss, including lineage-based recovery mechanisms.
- Global Control Store and Scheduling: Ray’s distributed bottom-up scheduler horizontally scales and handles dynamically constructed task graphs, unlike most centralized cluster schedulers.
7 Discussion and Experiences
Ray’s development experience emphasizes a minimal but extensible API, while its generality introduces optimization and lineage-storage challenges; centralized control state supports scalability and debugging.
- API: Ray’s API evolved from tasks to include wait() for heterogeneous rollouts and actors for stateful simulators and expensive initialization.
- Limitations: Ray’s generality makes specialized optimizations harder because scheduling decisions occur without full knowledge of the computation graph.
- Limitations: Lineage storage requires garbage-collection policies to bound Global Control Store costs, and this capability was still under development.
- Fault Tolerance: Fault tolerance simplifies application reasoning, helps debug stochastic algorithms, and permits use of cheaper resources such as AWS spot instances, with minimal overhead for target workloads.
- Global Control Store: The Global Control Store supports whole-system queries and timeline visualization for debugging Ray and applications.
- Global Control Store: Adding GCS shards and scheduler replicas enabled horizontal scaling as control-state bottlenecks emerged.
8 Conclusion
Ray combines programming flexibility, scalability, low latency, and fault tolerance for emerging AI applications. Its evaluation shows linear scaling to 1.8 million tasks per second and substantial improvements on contemporary reinforcement learning workloads.
- Ray unifies task-parallel and actor programming models within a single dynamic task graph.
- Its architecture combines a global control store with a bottom-up distributed scheduler.
- Ray achieves programming flexibility, high throughput, low latency, and transparent fault tolerance for heterogeneous AI workloads.
- 1.8 million tasks per second is Ray’s demonstrated linear scalability limit in the evaluation.
- Ray substantially improves performance on several contemporary reinforcement learning workloads.