4. Ray对象存储与共享内存

任务目标

本案例通过一个可运行的最小示例,演示Ray对象存储(Object Store)的核心用法:将大对象放入共享内存、跨任务传递对象引用(ObjectRef),以及使用ray.get检索对象。你将直观理解Ray如何通过共享内存机制避免大对象在分布式任务间的重复序列化拷贝。

环境依赖

需要已安装Ray(pip install ray),Python版本3.8及以上。本案例无需额外依赖。

核心实现

以下代码创建了一个约80MB的整数列表(模拟大对象),将其放入Ray对象存储,并通过两个远程任务验证对象引用的传递与共享内存读取:

import ray
import numpy as np
import time

# 初始化Ray,使用本地模式(无需集群)
ray.init(ignore_reinit_error=True)

# 模拟一个较大的对象:1000万个float64元素,约80MB
large_data = np.random.rand(10_000_000)

# 将大对象放入Ray Object Store,返回对象引用(ObjectRef)
data_ref = ray.put(large_data)
print(f"已放入对象存储,ObjectRef: {data_ref}")

# 远程任务1:接收对象引用,读取并计算数据总和
@ray.remote
def compute_sum(ref):
    data = ray.get(ref)  # 从共享内存读取对象
    return float(np.sum(data))

# 远程任务2:接收对象引用,读取并计算数据均值
@ray.remote
def compute_mean(ref):
    data = ray.get(ref)  # 同一对象,无需再次序列化传输
    return float(np.mean(data))

# 并行提交两个任务,传入同一个ObjectRef
sum_ref = compute_sum.remote(data_ref)
mean_ref = compute_mean.remote(data_ref)

# 获取结果
total = ray.get(sum_ref)
mean = ray.get(mean_ref)

# 验证结果与本地计算一致
local_sum = float(np.sum(large_data))
local_mean = float(np.mean(large_data))

print(f"远程求和: {total:.2f}, 本地求和: {local_sum:.2f}")
print(f"远程均值: {mean:.6f}, 本地均值: {local_mean:.6f}")

# 验证通过条件:误差小于1e-6
assert abs(total - local_sum) < 1e-6, "求和结果不一致"
assert abs(mean - local_mean) < 1e-6, "均值结果不一致"
print("验证通过:远程计算结果与本地一致")

# 查看对象存储统计信息
print(f"对象存储中的对象数量: {len(ray.objects())}")
ray.shutdown()

运行与验证

将上述代码保存为ray_object_store.py,在终端执行:

python ray_object_store.py

程序正常结束且输出包含验证通过即表示成功。关键观察点:

  • ray.put返回的ObjectRef是轻量级句柄,可跨任务传递;
  • 两个远程任务通过ray.get(ref)读取的是同一份共享内存数据,无需重复拷贝;
  • 最终断言确保分布式计算结果与本地直接计算一致,误差在1e-6以内。

内存管理机制说明

Ray的Object Store采用共享内存架构:ray.put将数据序列化后写入本机的共享内存文件系统(/dev/shm),返回的ObjectRef仅包含元数据指针。当远程任务需要该对象时,ray.get直接从共享内存映射读取,避免了跨进程的序列化传输开销。对于大对象(如本例的80MB数组),这种机制显著降低了分布式训练中的数据拷贝成本。

Ray对象存储共享内存数据流 主进程 large_data (80MB) ray.put() 序列化写入 Ray Object Store 共享内存 /dev/shm ObjectRef: 轻量级句柄 数据仅存一份 返回ObjectRef 传递ObjectRef 传递ObjectRef 远程任务 compute_sum ray.get(ref) → 直接读共享内存 远程任务 compute_mean ray.get(ref) → 同一份数据 结果返回 主进程 共享内存对象存储 远程任务 ObjectRef传递 结果返回

常见问题

  • 内存不足:若/dev/shm空间不足(默认约物理内存一半),可减小np.random.rand的参数(如改为1_000_000)或通过ray.init(object_store_memory=2_000_000_000)显式指定内存上限。
  • 对象生命周期:ObjectRef超出作用域后,Ray的引用计数机制会自动释放对应共享内存。本例中ray.shutdown()会清空所有对象。
  • 嵌套对象:ray.put支持任意可序列化Python对象(包括嵌套字典、自定义类实例),但要求对象可被pickle序列化。