base.py — verl/single_controller/ray/base.py¶
文件路径¶
verl/single_controller/ray/base.py
文件概述¶
这是 single_controller 模块中最大、最重要的实现文件（约 1100 行）。它是 base/ 中抽象类的 Ray 后端具体实现，包含：
- RayResourcePool：基于 Ray PlacementGroup 的资源池
- SubRayResourcePool：资源池的子集，用于拆分
- ResourcePoolManager：多资源池管理器
- RayClassWithInitArgs：Ray Actor 的延迟创建包装器
- RayWorkerGroup：核心类，管理一组 Ray Actor Worker
- FusedWorker 机制：多个 Worker 共存在同一个 Ray Actor 中
- 多种辅助函数：资源池合并、拆分、PlacementGroup 排序等
关键代码讲解¶
1. func_generator：代理函数生成器¶
def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking):
class Functor:
def __call__(this, *args, **kwargs):
args, kwargs = dispatch_fn(self, *args, **kwargs)
padding_count = kwargs.pop(_padding_size_key, 0)
output = execute_fn(method_name, *args, **kwargs)
if blocking:
output = ray.get(output)
output = collect_fn(self, output)
if padding_count > 0:
if isinstance(output, DataProto):
indices = [i for i in range(len(output))][:-padding_count]
output = output.select_idxs(indices)
elif isinstance(output, list):
output = output[:-padding_count]
return output
return type(method_name, (Functor,), {})()
这是整个分布式调用的执行管线。当用户调用 worker_group.some_method(data) 时，实际执行的就是这个 Functor：
- dispatch_fn：数据分发（如切分 DataProto）
- 弹出 padding 信息：如果分发时做了 padding，记录 padding 大小
- execute_fn：通过 Ray 远程调用各 Worker 的方法
- ray.get：如果是阻塞模式，等待所有 Worker 返回
- collect_fn：收集并合并结果
- 去除 padding：如果之前做了 padding，把多余的部分去掉
注意 type(method_name, (Functor,), {})() 这个技巧：动态创建一个以方法名命名的类（继承 Functor），这样在调试时看到的类名就是方法名，便于排查问题。
2. sort_placement_group_by_node_ip：PlacementGroup 排序¶
def sort_placement_group_by_node_ip(pgs: list[PlacementGroup]) -> list[PlacementGroup]:
node_ip = {node["NodeID"]: node["NodeManagerAddress"] for node in ray.nodes()}
pg_ip = {}
for pg in pgs:
specs = ray._private.state.state.placement_group_table(pg.id)
...
return sorted(pgs, key=lambda pg: pg_ip.get(pg, ""))
按节点 IP 排序 PlacementGroup。这对于分布式检查点恢复很重要：确保相同物理节点上的 Worker 获得相同的 RANK，即使 Ray 集群重启后 PlacementGroup 的分配顺序可能不同。
3. RayResourcePool：Ray 资源池¶
class RayResourcePool(ResourcePool):
def __init__(self, process_on_nodes=None, use_gpu=True, name_prefix=None,
max_colocate_count=10, detached=False, accelerator_type=None):
super().__init__(process_on_nodes, max_colocate_count)
self.use_gpu = use_gpu
self.name_prefix = get_random_string(length=6) if name_prefix is None else name_prefix
self.pgs = None
self.detached = detached
self.accelerator_type = accelerator_type
继承 ResourcePool，增加了 Ray 特有的属性：
- use_gpu：是否使用 GPU
- name_prefix：PlacementGroup 名称前缀
- pgs：缓存已创建的 PlacementGroup
- detached：是否为 detached 模式（PlacementGroup 在创建者退出后仍然存在）
get_placement_groups：创建 PlacementGroup¶
def get_placement_groups(self, strategy="STRICT_PACK", name=None, device_name="cuda"):
if self.pgs is not None:
return self.pgs
bundle = {"CPU": self.max_colocate_count}
if self.use_gpu:
bundle[device_name] = 1
pg_scheme = [[bundle.copy() for _ in range(process_count)] for process_count in self._store]
pgs = [
placement_group(bundles=bundles, strategy=strategy, name=pg_name_prefix + str(idx), lifetime=lifetime)
for idx, bundles in enumerate(pg_scheme)
]
ray.get([pg.ready() for pg in pgs])
self.pgs = sort_placement_group_by_node_ip(pgs)
return pgs
PlacementGroup 是 Ray 的资源预留机制。 这段代码的逻辑：
- 如果已经创建过，直接返回缓存
- 定义每个 bundle 的资源需求：
max_colocate_count个 CPU + 1 个 GPU - 按照
_store（每个节点的进程数）创建 PlacementGroup - 使用
STRICT_PACK策略确保同一个 PlacementGroup 的所有 bundle 在同一个节点上 - 等待所有 PlacementGroup 就绪
- 按节点 IP 排序
为什么 CPU 数量是 max_colocate_count？ 因为一张 GPU 上可能共存多个 Worker（如 Actor 和 Critic），每个 Worker 需要 1 个 CPU，所以一个 bundle 需要预留 max_colocate_count 个 CPU。
4. SubRayResourcePool：子资源池¶
class SubRayResourcePool(RayResourcePool):
def __init__(self, placement_groups, start_bundle_index, subgroup_world_size, **kwargs):
super().__init__(**kwargs)
self.pgs = placement_groups
self.start_bundle_index = start_bundle_index
self.subgroup_world_size = subgroup_world_size
@property
def world_size(self):
return self.subgroup_world_size
当需要把一个大的资源池拆分给不同的模型使用时，SubRayResourcePool 表示资源池的一个子集。它共享父池的 PlacementGroup，但只使用从 start_bundle_index 开始的 subgroup_world_size 个 bundle。
5. ResourcePoolManager：资源池管理器¶
@dataclass
class ResourcePoolManager:
resource_pool_spec: dict[str, list[int]]
mapping: dict[int, str]
resource_pool_dict: dict[str, RayResourcePool] = field(default_factory=dict)
def create_resource_pool(self):
for resource_pool_name, process_on_nodes in self.resource_pool_spec.items():
resource_pool = RayResourcePool(
process_on_nodes=process_on_nodes, use_gpu=True, max_colocate_count=3,
name_prefix=resource_pool_name
)
self.resource_pool_dict[resource_pool_name] = resource_pool
self._check_resource_available()
高级接口，用于管理多个资源池：
- resource_pool_spec：定义各资源池的规格（如 {"pool1": [4, 4], "pool2": [2]}）
- mapping：将角色映射到资源池（如 {actor_role: "pool1", reward_role: "pool2"}）
- _check_resource_available()：检查集群的 GPU 数量是否满足需求
6. split_resource_pool 和 merge_resource_pool¶
def split_resource_pool(resource_pool, split_size):
# 将一个资源池拆成多个子资源池
...
def merge_resource_pool(rp1, rp2):
# 合并两个资源池
new_store = rp1.store + rp2.store
merged = type(rp1)(new_store, rp1.use_gpu, ...)
merged.pgs = rp1.get_placement_groups() + rp2.get_placement_groups()
return merged
资源池的拆分和合并操作： - 拆分：把 8 GPU 的资源池拆成两个 4 GPU 的子池 - 合并：把两个资源池合并成一个更大的池
7. RayClassWithInitArgs：Ray Actor 创建器¶
class RayClassWithInitArgs(ClassWithInitArgs):
def __call__(self, placement_group, placement_group_bundle_idx,
use_gpu=True, num_gpus=1, sharing_with=None, device_name="cuda"):
if sharing_with is not None:
target_node_id = ray.get(sharing_with.get_node_id.remote())
visible_devices = ray.get(sharing_with.get_cuda_visible_devices.remote())
options = {"scheduling_strategy": NodeAffinitySchedulingStrategy(
node_id=target_node_id, soft=False)}
return self.cls.options(**options).remote(*self.args,
cuda_visible_devices=visible_devices, **self.kwargs)
options = {
"scheduling_strategy": PlacementGroupSchedulingStrategy(
placement_group=placement_group,
placement_group_bundle_index=placement_group_bundle_idx)
}
if use_gpu and device_name == "cuda":
options["num_gpus"] = num_gpus
return self.cls.options(**options).remote(*self.args, **self.kwargs)
继承 ClassWithInitArgs，增加了 Ray Actor 的创建逻辑：
- 使用 PlacementGroup 调度：确保 Worker 被放到指定的节点和 GPU 上
- sharing_with 模式：如果需要和现有 Actor 共享同一节点/GPU，使用
NodeAffinitySchedulingStrategy - GPU 资源分配：通过
num_gpus参数告诉 Ray 需要多少 GPU 资源
8. RayWorkerGroup：Ray Worker 组（核心类）¶
class RayWorkerGroup(WorkerGroup):
def __init__(self, resource_pool=None, ray_cls_with_init=None,
bin_pack=True, name_prefix=None, detached=False,
worker_names=None, worker_handles=None, **kwargs):
...
if self._is_init_with_detached_workers:
self._init_with_detached_workers(worker_names=worker_names, worker_handles=worker_handles)
elif isinstance(resource_pool, SubRayResourcePool):
self._init_with_subresource_pool(...)
else:
self._init_with_resource_pool(...)
if ray_cls_with_init is not None:
self._bind_worker_method(self.ray_cls_with_init.cls, func_generator)
RayWorkerGroup 有三种初始化方式： 1. 从已有的 detached workers 初始化：附加到已存在的 Ray Actor 上 2. 从 SubRayResourcePool 初始化：使用子资源池创建 Worker 3. 从 RayResourcePool 初始化：使用完整资源池创建 Worker
最后调用 _bind_worker_method，将 Worker 类上的 @register 方法绑定为 WorkerGroup 的代理方法。
_init_with_resource_pool：创建新 Worker¶
def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, detached, worker_env=None):
pgs = resource_pool.get_placement_groups(strategy=strategy, device_name=self.device_name)
world_size = resource_pool.world_size
rank = -1
local_world_size = resource_pool.store[0]
for pg_idx, pg in enumerate(sort_placement_group_by_node_ip(pgs)):
if pg_idx == 0:
self._get_master_addr_port(pg, bundle_index=0, ...)
for local_rank in range(local_world_size):
rank += 1
self._create_worker(rank=rank, pg_idx=pg_idx, pg=pg, local_rank=local_rank, ...)
遍历所有节点的所有 GPU，逐个创建 Worker。第一个 PlacementGroup 的第一个 bundle 用于获取 master_addr 和 master_port。
_create_worker：创建单个 Worker¶
def _create_worker(self, rank, pg_idx, pg, local_rank, resource_pool, ray_cls_with_init, worker_env, detached):
num_gpus = 1 / resource_pool.max_colocate_count
env_vars = {
"WORLD_SIZE": str(world_size),
"RANK": str(rank),
"WG_PREFIX": self.name_prefix,
"WG_BACKEND": "ray",
"MASTER_ADDR": self._master_addr,
"MASTER_PORT": self._master_port,
}
ray_cls_with_init.update_options({"runtime_env": {"env_vars": env_vars}, "name": name})
worker = ray_cls_with_init(
placement_group=pg, placement_group_bundle_idx=local_rank,
use_gpu=self.use_gpu, num_gpus=num_gpus, device_name=self.device_name,
)
self._workers.append(worker)
self._worker_names.append(name)
创建单个 Worker 的流程：
1. 计算每个 Worker 需要的 GPU 分数（1 / max_colocate_count）
2. 设置环境变量（RANK、WORLD_SIZE 等）
3. 通过 ray_cls_with_init 创建 Ray Actor
4. 将 Worker 添加到管理列表
为什么 num_gpus 是分数？ 如果 max_colocate_count=3（三个 Worker 共享一张 GPU），每个 Worker 只需要 1/3 个 GPU 的 Ray 资源配额。
execute_all_async：异步执行¶
def execute_all_async(self, method_name, *args, **kwargs):
length = len(self._workers)
if all(isinstance(arg, list) for arg in args) and all(isinstance(kwarg, list) for kwarg in kwargs.values()):
if all(len(arg) == length for arg in args) and all(len(kwarg) == length for kwarg in kwargs.values()):
result = []
for i in range(length):
sliced_args = tuple(arg[i] for arg in args)
sliced_kwargs = {k: v[i] for k, v in kwargs.items()}
result.append(
self._execute_remote_single_worker(self._workers[i], method_name, *sliced_args, **sliced_kwargs)
)
return result
return [self._execute_remote_single_worker(worker, method_name, *args, **kwargs) for worker in self._workers]
异步执行策略： - 如果所有参数都是列表且长度等于 worker 数量，则把每个参数的第 i 个元素发给第 i 个 Worker（分片模式） - 否则，把完整参数发给每个 Worker（广播模式）
spawn 和 spawn_fused：Worker 组拆分¶
def spawn(self, prefix_set):
if self.fused_worker_used:
return self.spawn_fused(prefix_set)
new_worker_group_dict = {}
for prefix in prefix_set:
new_worker_group = self.from_detached(
name_prefix=self.name_prefix,
worker_names=self._worker_names,
worker_handles=self._workers,
ray_cls_with_init=self.ray_cls_with_init,
)
_rebind_actor_methods(new_worker_group, prefix)
new_worker_group_dict[prefix] = new_worker_group
return new_worker_group_dict
spawn 用于将一个共存的 WorkerGroup 拆分成多个独立的 WorkerGroup。比如一个共存 Actor+Critic 的 WorkerGroup，spawn 后变成两个独立的 WorkerGroup：一个只有 Actor 方法，一个只有 Critic 方法。
9. FusedWorker 机制¶
def create_colocated_worker_raw_cls(class_dict: dict[str, RayClassWithInitArgs]):
raw_cls_dict = {cls_name: _unwrap_ray_remote(cia.cls) for cls_name, cia in class_dict.items()}
class FusedWorker(Worker):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
for cls_name, udc, ud_args, ud_kwargs in zip(...):
with temp_env_var("DISABLE_WORKER_INIT", "1"):
self.fused_worker_dict[cls_name] = udc(*ud_args, **ud_kwargs)
setattr(self, cls_name, self.fused_worker_dict[cls_name])
# 让每个子 Worker 都能访问其他子 Worker
for _, worker in self.fused_worker_dict.items():
setattr(worker, Worker.fused_worker_attr_name, self.fused_worker_dict)
def _fuw_execute(self, method_name: str, *args, **kwargs):
names = method_name.split("_fwmn_")
cls_name = names[0]
method_name = names[1]
udc_method = getattr(self.fused_worker_dict[cls_name], method_name)
return udc_method(*args, **kwargs)
return type(class_name_renamed, (FusedWorker,), {})
FusedWorker 是 verl 的一个重要优化。 在 RLHF 中，Actor 和 Critic 经常共享同一张 GPU。与其为每个模型创建独立的 Ray Actor（浪费资源），不如把多个 Worker 融合到一个 Ray Actor 中：
create_colocated_worker_raw_cls接收一个字典（如{"actor": ActorCIA, "critic": CriticCIA}）- 创建一个
FusedWorker类，它的__init__会在内部实例化 Actor 和 Critic _fuw_execute方法根据方法名前缀路由到对应的子 Worker- 每个子 Worker 通过
fused_worker_dict可以访问其他子 Worker
def create_colocated_worker_cls_fused(class_dict: dict[str, RayClassWithInitArgs]):
raw_colocated_worker_cls = create_colocated_worker_raw_cls(class_dict)
remote_cls = ray.remote(raw_colocated_worker_cls)
cia = RayClassWithInitArgs(cls=remote_cls)
cia.fused_worker_used = True
return cia
最终用户调用 create_colocated_worker_cls_fused 来创建融合的 Worker 类。
FusedWorker 结构图：
单个 Ray Actor (FusedWorker)
+--------------------------------------------------+
| |
| fused_worker_dict = { |
| "actor": ActorWorker 实例, |
| "critic": CriticWorker 实例, |
| } |
| |
| _fuw_execute("actor_fwmn_train_step", data) |
| -> self.fused_worker_dict["actor"] |
| .train_step(data) |
| |
+--------------------------------------------------+
10. create_colocated_worker_cls（旧版，已弃用）¶
def create_colocated_worker_cls(class_dict: dict[str, RayClassWithInitArgs]):
class WorkerDict(worker_cls):
def __init__(self):
super().__init__()
self.worker_dict = {}
for key, user_defined_cls in cls_dict.items():
user_defined_cls = _unwrap_ray_remote(user_defined_cls)
with temp_env_var("DISABLE_WORKER_INIT", "1"):
self.worker_dict[key] = user_defined_cls(...)
...
旧版的共存 Worker 实现，已被 create_colocated_worker_cls_fused 取代。核心差异在于旧版直接在 WorkerDict 上绑定带前缀的方法（如 actor_train_step），而新版通过 _fuw_execute 动态路由。
核心类/函数列表¶
| 名称 | 类型 | 说明 |
|---|---|---|
func_generator() |
函数 | 生成分发-执行-收集的代理函数 |
sort_placement_group_by_node_ip() |
函数 | 按节点 IP 排序 PlacementGroup |
get_master_addr_port() |
Ray 远程函数 | 获取 master 地址和端口 |
RayResourcePool |
类 | 基于 Ray PlacementGroup 的资源池 |
SubRayResourcePool |
类 | 资源池的子集 |
ResourcePoolManager |
dataclass | 管理多个资源池 |
split_resource_pool() |
函数 | 拆分资源池 |
merge_resource_pool() |
函数 | 合并资源池 |
RayClassWithInitArgs |
类 | Ray Actor 的延迟创建包装器 |
RayWorkerGroup |
类 | 基于 Ray 的 WorkerGroup 实现 |
RayWorkerGroup.execute_all_async() |
方法 | 异步在所有 Worker 上执行 |
RayWorkerGroup.execute_rank_zero() |
方法 | 在 rank 0 Worker 上执行 |
RayWorkerGroup.spawn() |
方法 | 将共存 WorkerGroup 拆分 |
RayWorkerGroup.fuse() |
方法 | 融合多个角色到一个 WorkerGroup |
create_colocated_worker_cls() |
函数 | 创建共存 Worker（旧版，弃用） |
create_colocated_worker_raw_cls() |
函数 | 创建 FusedWorker 原始类 |
create_colocated_worker_cls_fused() |
函数 | 创建融合 Worker 类（推荐使用） |
与其他模块的关系¶
- 继承
base/中的基类：RayResourcePool继承ResourcePool，RayWorkerGroup继承WorkerGroup，RayClassWithInitArgs继承ClassWithInitArgs - 使用
base/decorator.py：导入MAGIC_ATTR、Dispatch - 使用
base/worker.py：FusedWorker继承Worker - 依赖 Ray 库：
ray.remote、PlacementGroup、ray.get等 - 依赖
verl.protocol：使用DataProto进行数据操作 - 被上层 Trainer 使用：Trainer 创建
RayWorkerGroup来管理训练中的各个角色
Worker 创建完整流程图¶
用户代码:
resource_pool = RayResourcePool(process_on_nodes=[4, 4])
cls_with_args = RayClassWithInitArgs(cls=MyWorker)
wg = RayWorkerGroup(resource_pool, ray_cls_with_init=cls_with_args)
内部流程:
1. resource_pool.get_placement_groups()
-> 创建 2 个 PlacementGroup，每个 4 个 bundle
-> 每个 bundle: {CPU: max_colocate_count, GPU: 1}
-> ray.get([pg.ready()...]) 等待资源就绪
2. _get_master_addr_port(pg[0])
-> 在第一个节点上获取 master IP 和空闲端口
3. for each (pg, local_rank):
_create_worker(rank, pg, local_rank, ...)
-> 设置环境变量: RANK, WORLD_SIZE, MASTER_ADDR, ...
-> ray_cls_with_init(pg, bundle_idx, num_gpus=1/3)
-> MyWorker.options(scheduling_strategy=..., num_gpus=...).remote()
-> 创建出 Ray Actor
4. _bind_worker_method(MyWorker, func_generator)
-> 遍历 MyWorker 的 @register 方法
-> 为每个方法生成代理函数
-> setattr(wg, method_name, proxy_fn)
结果:
wg.train_step(data) 可以直接调用
内部自动: 切分数据 -> 分发到8个Worker -> 各自执行 -> 收集合并 -> 返回
小结¶
ray/base.py 是 single_controller 模块的实现核心。它把抽象的 Worker、WorkerGroup、ResourcePool 概念落地到 Ray 分布式框架上。关键要点：
- RayResourcePool 通过 Ray PlacementGroup 管理 GPU 资源
- RayWorkerGroup 负责创建 Ray Actor 并绑定代理方法
- func_generator 实现了 dispatch -> execute -> collect 的完整管线
- FusedWorker 让多个角色共享同一个 Ray Actor，节省资源
- 所有复杂的分布式通信细节都被封装在这一层，上层用户只需关心业务逻辑