1
2
3
4
5

@dataclass
class DataProto:
    batch: TensorDict = None              # 张量数据容器
    non_tensor_batch: dict = field(default_factory=dict)  # 非张量数据
    meta_info: dict = field(default_factory=dict)         # 元信息

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

def check_consistency(self):
    """检查 DataProto 的一致性"""
    if self.batch is not None:
        assert len(self.batch.batch_size) == 1, "只支持 num_batch_dims=1"
    
    if self.non_tensor_batch is not None:
        for key, val in self.non_tensor_batch.items():
            assert isinstance(val, np.ndarray)
            
    # 检查批次大小一致性
    if self.batch is not None and self.non_tensor_batch is not None:
        batch_size = self.batch.batch_size[0]
        for key, val in self.non_tensor_batch.items():
            assert val.shape[0] == batch_size

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

class RayPPOTrainer:
    def fit(self):
        # 控制流：训练循环
        for epoch in range(self.config.trainer.total_epochs):
            # 1. 数据准备
            batch = self._get_training_batch()
            
            # 2. 分发到计算流
            rollout_data = self.actor_rollout_wg.generate_sequences(batch)
            
            # 3. 收集结果
            advantages = self._compute_advantages(rollout_data)
            
            # 4. 策略更新
            self._update_policy(advantages)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

@register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
def generate_sequences(self, data: DataProto) -> DataProto:
    # 计算流：序列生成
    with torch.no_grad():
        # 1. 模型推理
        outputs = self.model.generate(
            input_ids=data.batch["input_ids"],
            attention_mask=data.batch["attention_mask"]
        )
        
        # 2. 返回结果
        return DataProto(
            batch=TensorDict({
                "generated_ids": outputs.sequences,
                "log_probs": outputs.log_probs
            }, batch_size=data.batch.batch_size),
            meta_info=data.meta_info
        )

1
2
3
4
5
6
7

class Dispatch(DynamicEnum):
    RANK_ZERO = "RANK_ZERO"                    # 只在rank 0执行
    ONE_TO_ALL = "ONE_TO_ALL"                  # 一对多广播
    ALL_TO_ALL = "ALL_TO_ALL"                  # 全对全通信
    DP_COMPUTE = "DP_COMPUTE"                   # 数据并行计算
    DP_COMPUTE_PROTO = "DP_COMPUTE_PROTO"       # DataProto数据并行
    DP_COMPUTE_PROTO_WITH_FUNC = "DP_COMPUTE_PROTO_WITH_FUNC"  # 带函数的DataProto并行

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs):
    """DataProto数据并行分发"""
    # 自动分割DataProto到worker数量
    splitted_args, splitted_kwargs = _split_args_kwargs_data_proto_with_auto_padding(
        worker_group.world_size, *args, **kwargs
    )
    return splitted_args, splitted_kwargs

def collect_dp_compute_data_proto(worker_group, output):
    """DataProto数据并行收集"""
    output = collect_dp_compute(worker_group, output)
    return _concat_data_proto_or_future(output)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

def _split_args_kwargs_data_proto_with_auto_padding(chunks, *args, **kwargs):
    """支持自动填充的数据分割"""
    data_proto_len = None
    padding_size = None
    
    def _padding_and_split_data(obj, chunks):
        nonlocal data_proto_len, padding_size
        if isinstance(obj, DataProto) and obj.is_padding_enabled():
            if data_proto_len is None:
                data_proto_len = len(obj)
                padding_size = (chunks - (data_proto_len % chunks)) if (data_proto_len % chunks > 0) else 0
            obj.padding(padding_size=padding_size)
        return obj.chunk(chunks=chunks)
    
    # 处理所有参数
    splitted_args = [_padding_and_split_data(arg, chunks) for arg in args]
    splitted_kwargs = {key: _padding_and_split_data(val, chunks) for key, val in kwargs.items()}
    
    return splitted_args, splitted_kwargs

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

def chunk(self, chunks: int) -> list["DataProto"]:
    """将DataProto分割成多个块"""
    if self.batch is not None:
        batch_lst = self.batch.chunk(chunks=chunks, dim=0)
    else:
        batch_lst = [None for _ in range(chunks)]
    
    # 处理非张量数据
    non_tensor_batch_lst = [{} for _ in range(chunks)]
    for key, val in self.non_tensor_batch.items():
        non_tensor_lst = np.array_split(val, chunks)
        for i in range(chunks):
            non_tensor_batch_lst[i][key] = non_tensor_lst[i]
    
    return [type(self)(batch=batch_lst[i], 
                      non_tensor_batch=non_tensor_batch_lst[i], 
                      meta_info=self.meta_info) for i in range(chunks)]

1
2
3
4
5

def to(self, device) -> "DataProto":
    """设备间数据移动"""
    if self.batch is not None:
        self.batch = self.batch.to(device)
    return self

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

@dataclass
class DataProtoFuture:
    """异步DataProto，避免阻塞控制流"""
    collect_fn: Callable
    futures: list[ray.ObjectRef]
    dispatch_fn: Callable = None
    
    def get(self):
        output = ray.get(self.futures)
        output = self.collect_fn(output)
        if self.dispatch_fn is not None:
            output = self.dispatch_fn(output)
        return output

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83

# 原始训练数据
raw_data = {
    "prompts": [
        "请计算 15 + 27 = ?",
        "求解方程 2x + 5 = 13",
        "一个圆的半径是3，求面积"
    ],
    "responses": [
        "15 + 27 = 42",
        "2x + 5 = 13\nx = 4", 
        "面积 = πr² = 9π"
    ],
    "rewards": [0.8, 0.9, 0.7]
}

# 转换为DataProto
def create_training_dataproto(raw_data):
    # 1. 文本编码
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
    
    prompt_ids = []
    response_ids = []
    attention_masks = []
    
    for prompt, response in zip(raw_data["prompts"], raw_data["responses"]):
        # 编码prompt
        prompt_tokens = tokenizer.encode(prompt, add_special_tokens=True)
        prompt_ids.append(prompt_tokens)
        
        # 编码response
        response_tokens = tokenizer.encode(response, add_special_tokens=False)
        response_ids.append(response_tokens)
        
        # 创建attention mask
        total_length = len(prompt_tokens) + len(response_tokens)
        attention_masks.append([1] * total_length)
    
    # 2. 填充到相同长度
    max_length = max(len(p) + len(r) for p, r in zip(prompt_ids, response_ids))
    
    padded_prompt_ids = []
    padded_response_ids = []
    padded_attention_masks = []
    
    for p_ids, r_ids, mask in zip(prompt_ids, response_ids, attention_masks):
        # 填充prompt
        padded_p = p_ids + [tokenizer.pad_token_id] * (max_length - len(p_ids) - len(r_ids))
        padded_prompt_ids.append(padded_p)
        
        # 填充response
        padded_r = r_ids + [tokenizer.pad_token_id] * (max_length - len(p_ids) - len(r_ids))
        padded_response_ids.append(padded_r)
        
        # 更新attention mask
        padded_mask = mask + [0] * (max_length - len(mask))
        padded_attention_masks.append(padded_mask)
    
    # 3. 创建DataProto
    training_dataproto = DataProto.from_dict(
        tensors={
            "prompt_ids": torch.tensor(padded_prompt_ids, dtype=torch.long),
            "response_ids": torch.tensor(padded_response_ids, dtype=torch.long),
            "attention_mask": torch.tensor(padded_attention_masks, dtype=torch.long),
        },
        non_tensors={
            "raw_prompts": np.array(raw_data["prompts"], dtype=object),
            "raw_responses": np.array(raw_data["responses"], dtype=object),
            "rewards": np.array(raw_data["rewards"], dtype=np.float32),
        },
        meta_info={
            "dataset_name": "math_training",
            "batch_size": len(raw_data["prompts"]),
            "max_length": max_length,
            "tokenizer_name": "Qwen/Qwen2.5-7B-Instruct"
        }
    )
    
    return training_dataproto

# 创建训练数据
training_data = create_training_dataproto(raw_data)
print("DataProto结构:")
print(training_data.get_data_info())

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

# 假设有4个GPU worker
world_size = 4
batch_size = 3

# 原始DataProto
original_dataproto = training_data  # batch_size=3

# 自动填充到能被4整除的大小
padded_dataproto, pad_size = pad_dataproto_to_divisor(original_dataproto, world_size)
# 现在padded_dataproto的batch_size=4 (填充了1个样本)

# 分割成4个chunk
chunks = padded_dataproto.chunk(chunks=world_size)

print(f"原始batch_size: {len(original_dataproto)}")
print(f"填充后batch_size: {len(padded_dataproto)}")
print(f"分割后chunk数量: {len(chunks)}")
print(f"每个chunk的batch_size: {len(chunks[0])}")

# 分发到各个worker
for i, chunk in enumerate(chunks):
    print(f"Worker {i} 接收数据:")
    print(f"  - prompt_ids shape: {chunk.batch['prompt_ids'].shape}")
    print(f"  - 包含样本数: {len(chunk)}")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

class RayPPOTrainer:
    def __init__(self, config):
        self.config = config
        self.actor_rollout_wg = None  # Actor和Rollout的WorkerGroup
        self.critic_wg = None         # Critic的WorkerGroup
        self.ref_policy_wg = None     # Reference Policy的WorkerGroup
        
    def fit(self):
        """PPO训练的主控制循环"""
        for epoch in range(self.config.trainer.total_epochs):
            print(f"开始第 {epoch} 轮训练")
            
            # 1. 获取训练批次
            batch_dataproto = self._get_training_batch()
            print(f"获取训练批次，batch_size: {len(batch_dataproto)}")
            
            # 2. 分发到Actor进行序列生成
            print("开始序列生成...")
            rollout_dataproto = self.actor_rollout_wg.generate_sequences(batch_dataproto)
            print(f"序列生成完成，生成 {len(rollout_dataproto)} 个序列")
            
            # 3. 分发到Critic计算价值
            print("开始价值计算...")
            value_dataproto = self.critic_wg.compute_values(rollout_dataproto)
            print(f"价值计算完成")
            
            # 4. 分发到Reference Policy计算log概率
            print("开始参考策略计算...")
            ref_log_prob_dataproto = self.ref_policy_wg.compute_log_probs(rollout_dataproto)
            print(f"参考策略计算完成")
            
            # 5. 在控制流中计算优势函数
            print("计算优势函数...")
            advantages = self._compute_advantages(
                rollout_dataproto, 
                value_dataproto, 
                ref_log_prob_dataproto
            )
            
            # 6. 更新Actor策略
            print("更新Actor策略...")
            self.actor_rollout_wg.update_actor(advantages)
            
            # 7. 更新Critic价值网络
            print("更新Critic价值网络...")
            self.critic_wg.update_critic(advantages)
            
            print(f"第 {epoch} 轮训练完成\n")
    
    def _get_training_batch(self):
        """获取训练批次"""
        # 从数据集中采样
        batch_data = {
            "prompt_ids": torch.randint(0, 1000, (self.config.data.train_batch_size, 512)),
            "attention_mask": torch.ones(self.config.data.train_batch_size, 512),
        }
        
        return DataProto.from_dict(
            tensors=batch_data,
            meta_info={"epoch": self.current_epoch}
        )
    
    def _compute_advantages(self, rollout_data, value_data, ref_log_prob_data):
        """计算优势函数 - 在控制流中执行"""
        # 从各个DataProto中提取数据
        rewards = rollout_data.batch["rewards"]  # [batch_size, seq_len]
        values = value_data.batch["values"]      # [batch_size, seq_len]
        log_probs = rollout_data.batch["log_probs"]  # [batch_size, seq_len]
        ref_log_probs = ref_log_prob_data.batch["log_probs"]  # [batch_size, seq_len]
        
        # 计算优势函数
        advantages = rewards - values
        advantages = advantages * rollout_data.batch["attention_mask"]
        
        # 计算策略比率
        log_ratio = log_probs - ref_log_probs
        ratio = torch.exp(log_ratio)
        
        # 创建优势DataProto
        advantages_dataproto = DataProto.from_dict(
            tensors={
                "advantages": advantages,
                "ratio": ratio,
                "rewards": rewards,
                "values": values
            },
            meta_info=rollout_data.meta_info
        )
        
        return advantages_dataproto

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

@ray.remote
class ActorRolloutWorker:
    def __init__(self, model_config):
        self.model = None
        self.config = model_config
        
    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
    def generate_sequences(self, data: DataProto) -> DataProto:
        """生成序列 - 在计算流中执行"""
        print(f"Worker {self.rank} 开始生成序列，batch_size: {len(data)}")
        
        # 1. 模型推理
        with torch.no_grad():
            input_ids = data.batch["prompt_ids"]
            attention_mask = data.batch["attention_mask"]
            
            # 生成序列
            outputs = self.model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=256,
                do_sample=True,
                temperature=0.7,
                return_dict_in_generate=True,
                output_scores=True
            )
            
            # 提取生成的token ids
            generated_ids = outputs.sequences
            log_probs = torch.stack(outputs.scores, dim=1).log_softmax(dim=-1)
            
            # 计算每个token的log概率
            token_log_probs = []
            for i, seq in enumerate(generated_ids):
                seq_log_probs = []
                for j, token_id in enumerate(seq):
                    if j < len(log_probs[i]):
                        seq_log_probs.append(log_probs[i][j][token_id].item())
                token_log_probs.append(seq_log_probs)
            
            # 计算奖励（这里使用简单的长度奖励作为示例）
            rewards = torch.tensor([[len(seq) * 0.1] * len(seq) for seq in generated_ids])
        
        # 2. 创建结果DataProto
        result_dataproto = DataProto.from_dict(
            tensors={
                "generated_ids": generated_ids,
                "log_probs": torch.tensor(token_log_probs),
                "rewards": rewards,
                "attention_mask": torch.ones_like(generated_ids)
            },
            non_tensors={
                "raw_generated_texts": np.array([
                    self.tokenizer.decode(seq, skip_special_tokens=True) 
                    for seq in generated_ids
                ], dtype=object)
            },
            meta_info=data.meta_info
        )
        
        print(f"Worker {self.rank} 序列生成完成")
        return result_dataproto
    
    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
    def update_actor(self, advantages: DataProto) -> None:
        """更新Actor策略 - 在计算流中执行"""
        print(f"Worker {self.rank} 开始更新Actor策略")
        
        # 1. 提取数据
        advantages_tensor = advantages.batch["advantages"]
        ratio = advantages.batch["ratio"]
        rewards = advantages.batch["rewards"]
        
        # 2. 计算PPO损失
        clip_ratio = 0.2
        policy_loss_1 = -advantages_tensor * ratio
        policy_loss_2 = -advantages_tensor * torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio)
        policy_loss = torch.maximum(policy_loss_1, policy_loss_2).mean()
        
        # 3. 反向传播
        self.optimizer.zero_grad()
        policy_loss.backward()
        self.optimizer.step()
        
        print(f"Worker {self.rank} Actor策略更新完成，损失: {policy_loss.item():.4f}")

@ray.remote
class CriticWorker:
    def __init__(self, model_config):
        self.model = None
        self.config = model_config
        
    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
    def compute_values(self, data: DataProto) -> DataProto:
        """计算价值函数 - 在计算流中执行"""
        print(f"Worker {self.rank} 开始计算价值函数")
        
        with torch.no_grad():
            input_ids = data.batch["generated_ids"]
            attention_mask = data.batch["attention_mask"]
            
            # 前向传播计算价值
            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
            values = outputs.value  # [batch_size, seq_len]
        
        # 创建价值DataProto
        value_dataproto = DataProto.from_dict(
            tensors={"values": values},
            meta_info=data.meta_info
        )
        
        print(f"Worker {self.rank} 价值函数计算完成")
        return value_dataproto

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# 示例：batch_size=250, world_size=4
original_batch_size = 250
world_size = 4

# 计算需要填充的数量
padding_needed = (world_size - (original_batch_size % world_size)) % world_size
# padding_needed = 2

# 填充后的batch_size = 252，可以被4整除
final_batch_size = original_batch_size + padding_needed  # 252

# 每个worker获得 252 // 4 = 63 个样本
samples_per_worker = final_batch_size // world_size  # 63

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

# 控制流中的异步调用
def async_training_step(self):
    # 1. 异步生成序列
    rollout_future = self.actor_rollout_wg.generate_sequences.remote(training_data)
    
    # 2. 控制流可以继续其他工作
    print("序列生成正在进行中...")
    
    # 3. 当需要结果时再等待
    rollout_data = ray.get(rollout_future)
    print("序列生成完成")
    
    # 4. 继续后续步骤
    value_data = self.critic_wg.compute_values(rollout_data)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

# 1. 启动Head节点
ray start --head --port=6379 --dashboard-host=0.0.0.0 --dashboard-port=8265

# 2. 启动Worker节点
ray start --address='head_node_ip:6379'

# 3. 在Driver中初始化集群
import ray
ray.init(address='ray://head_node_ip:10001')

# 4. 创建WorkerGroup
resource_pool = RayResourcePool(
    process_on_nodes=[4, 4, 4],  # 每个节点4个进程
    use_gpu=True,
    max_colocate_count=1
)

# 5. 启动Worker进程
actor_rollout_wg = RayWorkerGroup(
    resource_pool=resource_pool,
    ray_cls_with_init=RayClassWithInitArgs(ActorRolloutWorker, config)
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

# Object Store基本操作示例
import ray

# 1. 存储DataProto到Object Store
data_proto = DataProto.from_dict(tensors={"input_ids": torch.randn(256, 512)})
object_ref = ray.put(data_proto)  # 存储到本地Object Store

# 2. 发送ObjectRef到远程Worker
@ray.remote
def worker_function(ref):
    data = ray.get(ref)  # 从Object Store获取数据
    return process_data(data)

# 3. 执行远程调用
future = worker_function.remote(object_ref)
result = ray.get(future)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

# DataProto的序列化过程
def __getstate__(self):
    """序列化DataProto"""
    import io
    buffer = io.BytesIO()
    
    # 1. 序列化TensorDict
    if self.batch is not None:
        batch_to_save = self.batch.contiguous().consolidate()
        torch.save(batch_to_save, buffer)
    
    # 2. 序列化numpy数组和meta_info
    buffer_bytes = buffer.getvalue()
    return buffer_bytes, self.non_tensor_batch, self.meta_info

def __setstate__(self, data):
    """反序列化DataProto"""
    import io
    batch_deserialized_bytes, non_tensor_batch, meta_info = data
    
    # 1. 反序列化TensorDict
    batch_deserialized = io.BytesIO(initial_bytes=batch_deserialized_bytes)
    batch = torch.load(batch_deserialized, map_location="cpu")
    
    # 2. 重建DataProto
    self.batch = batch
    self.non_tensor_batch = non_tensor_batch
    self.meta_info = meta_info

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

# 实际的数据分发过程
def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs):
    """DataProto数据并行分发"""
    world_size = worker_group.world_size  # 例如：8个worker
    
    # 1. 自动填充
    data_proto = args[0]  # 原始DataProto，batch_size=250
    padded_data, pad_size = pad_dataproto_to_divisor(data_proto, world_size)
    # 现在padded_data的batch_size=256 (填充了6个样本)
    
    # 2. 分割数据
    chunks = padded_data.chunk(chunks=world_size)
    # 每个chunk的batch_size=32
    
    # 3. 分发到各个worker
    futures = []
    for i, worker in enumerate(worker_group.workers):
        chunk = chunks[i]
        # 序列化chunk并通过Ray发送
        future = worker.generate_sequences.remote(chunk)
        futures.append(future)
    
    return futures

# 在物理层面的实际传输
def physical_data_transfer_example():
    """展示物理层面的数据传输"""
    
    # 原始数据大小
    batch_size = 250
    seq_len = 512
    vocab_size = 32000
    
    # 计算数据大小
    prompt_ids_size = batch_size * seq_len * 4  # int32, bytes
    attention_mask_size = batch_size * seq_len * 4  # int32, bytes
    tensor_data_size = prompt_ids_size + attention_mask_size  # ~2MB
    
    # 序列化开销
    serialization_overhead = tensor_data_size * 0.1  # 约10%开销
    
    # 网络传输
    network_bandwidth = 10 * 1024 * 1024  # 10Gbps
    transfer_time = (tensor_data_size + serialization_overhead) / network_bandwidth
    
    print(f"数据大小: {tensor_data_size / 1024 / 1024:.2f} MB")
    print(f"序列化开销: {serialization_overhead / 1024 / 1024:.2f} MB")
    print(f"传输时间: {transfer_time * 1000:.2f} ms")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

# 内存优化示例
class MemoryOptimizedDataProto:
    def __init__(self):
        self.object_refs = {}  # 缓存ObjectRef
        
    def send_to_worker(self, data_proto, worker_id):
        """优化的数据传输"""
        
        # 1. 检查是否已有缓存
        cache_key = f"{worker_id}_{hash(data_proto)}"
        if cache_key in self.object_refs:
            return self.object_refs[cache_key]
        
        # 2. 序列化并存储
        object_ref = ray.put(data_proto)
        self.object_refs[cache_key] = object_ref
        
        # 3. 发送到worker
        return object_ref
    
    def cleanup_cache(self):
        """清理缓存"""
        for ref in self.object_refs.values():
            ray.delete(ref)
        self.object_refs.clear()