Ray - A Distributed Framework for Emerging AI Applications

Liao Jiayi Liao Jiayi #Ray#Technology

Notes on Ray, a distributed framework designed for reinforcement learning and emerging AI workloads.

Translated from Chinese with AI · Read the original

Some notes and thoughts on Ray.

Background

Open-sourced in 2018, Ray was hailed by technology media as a future replacement for Spark. Its paper, however, explicitly states that it was designed for reinforcement learning. The introduction observes that supervised learning trains on offline data before deployment online, with a lengthy processing pipeline that often takes hours or days to turn data into business value. This is clearly insufficient for emerging applications such as autonomous driving.

xx

For this RL system, policy optimization can be abstracted as follows:

def rollout(policy, environment):
state <- environment.initial_state()
// 不断迭代直到模拟失败
while (not environment.has_terminated):
action <- policy.compute(state)
state,reward <- environment.step(action)
trajectory.append(state, reward)
return trajectory
def train_policy(environment):
policy <- initial_policy()
while (policy not converged):
// 对于每个 policy 进行 k 次模拟实验
for i from 1 to k:
trajectory.append(rollout(policy, environment))
policy := policy.update(trajectories)
return policy

Reinforcement learning differs from traditional supervised learning in several ways:

  • Results depend heavily on extensive trial and error and feedback in simulated environments, requiring many tasks and substantial computation.
  • The computation DAG can evolve continuously: each simulation result may determine the next computation topology.
  • Uncertain environmental changes demand extremely low latency, with millions of tasks computed on a millisecond timescale.

Popular frameworks such as MapReduce, Spark, and TensorFlow cannot fully satisfy all these requirements.

Design Goals

Ray’s design goals include:

  • Flexibility
    • Heterogeneous tasks
      • Parallel tasks may perform different computations.
      • Tasks may have different durations. Parallel experiments can run for different lengths of time, requiring flexible control, such as promptly releasing resources.
      • Tasks may use different resources.
    • Dynamic topology: Unlike Spark and Flink jobs compiled into graphs before execution, Ray’s topology forms dynamically at runtime.
  • High performance
  • Developer friendliness
    • Data replay
    • Fault tolerance
    • Parallelizing algorithms

Architecture

Application Layer

  • Driver: The process running the user’s program.
  • Worker: A stateless process that executes stateless tasks, analogous to a Flink TaskManager or Spark executor.
  • Actor: A stateful process started by a worker that executes stateful tasks.

System Layer

Bottom-Up Distributed Scheduler

Each node has a local scheduler, with a global scheduler coordinating the cluster. Scheduling is bottom-up: a task first goes to the local scheduler, which asks the global scheduler for help if necessary. The global scheduler scales horizontally. The scheduling flow is shown below:

xx

In-Memory Distributed Object Store(plasma)

Each node has an in-memory object store for state. Shared memory enables zero-copy sharing between processes on the same node. If a task’s data is absent locally, it is copied from another node. Lost objects or failed nodes can be recovered by reconstructing data from graph lineage.

Global Control Store (GCS)

The metadata management server scales horizontally.

  • Object Table: Maps objects in object stores to their node locations.
  • Task Table: Stores task execution information.
  • Function Table: Stores user-defined functions and computation-graph information.

Combining these three components gives the following diagram:

xx

Example Program

A simple a+b example illustrates how the architecture works:

xx

(a)

  1. Define add and broadcast the function to another node, N2.
  2. N1 invokes id = add.remote(a,b), submitting it first to the local scheduler.
  3. The local scheduler finds a on N1 but not b, and asks the global scheduler for help.
  4. The global scheduler queries the GCS Object Table and locates b on N2.
  5. The global scheduler schedules through N2’s local scheduler.
  6. N2’s local scheduler discovers that a is missing.
  7. N2 queries the GCS Object Table and locates a on N1.
  8. Copy a from N1’s object store into the local object store.
  9. Execute a+b on N2.

(b)

Diagram b is straightforward and needs no further explanation.

Programming & Computation Model

The paper repeatedly discusses the programming model, whose main features are:

  • Instead of dividing a graph into stages as traditional engines do, Ray can use ray.wait() to decide what to do next based on the results already received, potentially only a subset.
  • Heterogeneous resources, as illustrated below.
  • Nested remote functions: A remote function can call other remote functions.
  • Actors: Stateful operators.

Nested remote functions are particularly interesting: a task in a distributed job can start another distributed job. In Spark terms, this is like calling an RDD from a task within an RDD. I have encountered use cases for this, but Spark’s static topology does not allow such flexibility.

Conclusion

Ray’s design does not emphasize batch versus streaming. It provides relatively low-level distributed capabilities and leaves higher-level mechanisms unimplemented. For example, users may need to hand-code shuffle logic with Ray APIs. Further adoption will require abstractions for these higher-level mechanisms.

Several design choices stand out:

  1. All centralized components scale horizontally. The global scheduler can scale out, and GCS data can be sharded. Considering scalability from multiple angles at the beginning is a valuable lesson for software designers.
  2. The bottom-up distributed scheduler is an inventive approach that substantially accelerates task scheduling and increases cluster scheduling throughput.
  3. The in-memory object store resembles Flink’s StateBackend in some ways. Because StateBackend wraps RocksDB or other existing stores, Flink has limited control over state behavior. Production storage-performance problems often resist repeated tuning. Ray’s decision to build its own object store early recognized the component’s central role and seems forward-looking.
  4. Nested remote functions and FaaS-style programming: Dynamic graphs make flexible invocation possible. Flink and Spark programming can feel cumbersome and rigid; I believe this more flexible style will become mainstream.

References