veRL代码阅读-2.Ray

看VeRL代码之前发现代码里主要使用了ray框架来进行调度和通信. 所以先对ray进行初步学习, 后续有空闲时间再细看下Ray的代码.

框架原理

构成

架构图如下, ray里主要分为系统层面的layer和应用层的layer.

系统层layer:

GCS(Global Control Store): 中心数据存储，是 Worker 之间传递消息的纽带, 储存了代码, 输入参数, 返回值.
Scheduler: 分成Global和Local两种, 严重怀疑VeRL论文里说的中心节点其实就是这个Global. 待确认. Local是每个单机上的调度器(又名Raylet), worker通过Local和Global进行通信. 下图是一个交互示意图:
Object Store: 主要作用是通过RPC传递worker间的数据,

应用层Layer:

Driver: 执行用户程序的进程
Worker: 无状态的执行remote task的进程, worker是框架自动启动的. 当用户声明了一个remote方法时, 这个remote方法会被发布到所有的worker上
Actor: 有状态进程(这里的actor概念和rl里的actor不一样注意一下)，在被调用时只执行其暴露的函数。与worker不同的地方在于，actor需要worker或driver显式实例化, 特殊点在于每个方法的执行依赖于前一个方法产生的状态

调用流程

以一个remote请求发送和获取结果为例. 发送的步骤依次为:

driver把remote函数提交给localScheduler -> localScheduler把任务告知global -> globalScheduler查询GCS拿到对应的函数参数存的位置 -> global把这个任务调度有参数b的Node2 -> Node2的LocalScheduler检查所有的参数是不是本地都有 -> 把没有的参数查GCS,从而知道需要的参数存在哪个节点上 -> 从对应的节点拉取本地没有的参数 -> 所有入参数据ready后, 执行计算步骤 -> 把计算的结果存到本地的ObjectStore里

接收计算结果(ray.get)的步骤:

通过localScheduler在ObjectStore里检查返回的future对象\(id_c\)是否在本地 -> 向GCS查询\(id_c\)的存储位置, 如果这时候\(id_c\)还没产出, N1就会往GCS注册一个回调, 当产出时通知 -> N2完成计算后把结果存储到本地ObjectStore, 把\(id_c\)的meta信息添加到GCS, 这时候就会触发之前注册的回调 -> GCS通知N1数据可用,并且发送位置信息 -> N1向N2发送RPC请求拉取\(id_c\), 整个过程完成.

使用方法

task模式, 也就是上面说的worker, 适合无状态的逻辑执行

@ray.remote  #定义无状态分布式任务

def add(x, y):

    return x + y

# 异步提交任务，立即返回 future 对象（对象引用）

future = add.remote(1, 2)

results = ray.get([future]) #走上一章节说的get模式从remote拉取结果

Actor模式, 能够维护状态并封装方法。适合需要有状态的场景，如参数服务器、计数器. 特点是同一个Actor的方法调用按顺序执行，保证状态一致性, 而不同Actor实例之间可以并行执行

@ray.remote  # 使用 @ray.remote 将类转换为分布式 Actor, 实例在其生命周期内可以维持状态

class Counter:

    def __init__(self):

        self.value = 0

    def increment(self):

        self.value += 1

        return self.value

    def get_value(self):

        return self.value

counter = Counter.remote()  # 创建 Actor 实例, 存储在GCS里

future1 = counter.increment.remote()  # 第一次增加

future2 = counter.increment.remote()  # 第二次增加

future3 = counter.get_value.remote()  # 获取当前值

print(ray.get([future1, future2]))  # 输出: [1, 2]

print(ray.get(future3))  # 输出: 2

主要接口:

接口	功能
ray.put()	通过对象存储实现跨节点数据访问, put到remote的同时在GCS注册, 返回对象引用
ray.get()	同上, 功能变为数据拉取
ray.exceptions & @ray.remote(max_retries=3)	异常处理与重试, 搭配使用提升容错
@ray.remote(num_gpus=1) def gpu_task():	静态资源配置
future = task.options(num_cpus=2, num_gpus=1).remote()	动态资源配置
ray.kill(Actor)	强制终止 Actor 实例
refs = [task.remote(i) for i in range(4)] ready_refs, remaining_refs = ray.wait(refs, num_returns=len(refs))	等待异步提交的任务完成

参考

ray论文: https://arxiv.org/abs/1712.05889