support checkpoint engine[P2P] recovery in colocate mode - #2027
Conversation
Checkpoint-Engine P2P 推理引擎故障恢复在共卡 RL 训练中，rollout worker 负责启动推理引擎并处理生成请求。若某个推理引擎进程异常退出，原有流程会由 WorkerLifecycleState新增状态PENDING_WEIGHTS恢复后的 rollout worker group 不会立刻进入
checkpoint-engine P2P 更新 pending 目标
1.. 从 后台监控线程在 checkpoint-engine 模式下， 该线程会在满足以下条件时触发恢复：
权重同步语义
|
36a86ff to
8a53757
Compare
4f909e0 to
5b31599
Compare
|
@claude review |
1 similar comment
|
@claude review |
|
Claude: Summary本 PR 为 colocate 模式引入 checkpoint-engine P2P 故障恢复：新增 ProduceBatchResult impact: RoutedExperts impact: 新增的 Ray concurrency impact: 有效分组仍为 IssuesCritical
Warning
Main Flowchart after this PRflowchart TD
A[RLColocateTrainer._fit] --> B[_rollout_resources_available.set]
B --> C[_start_check_pending_rollout_worker_thread]
C --> D[produce_batch]
D --> E{rollout worker 崩溃?}
E -- 是 --> F[HealthManager.run_once: mark INACTIVE]
F --> G[run_once 内同步 _restart_inactive_workers]
G --> H[set_groups_state -> PENDING_WEIGHTS]
H --> I[后台线程 _update_pending_rollout_weights_from_checkpoint_engine]
I --> J[bind pending targets + weight_update need_update]
J --> K[mark_worker_groups_lifecycle_state PENDING->ACTIVE]
E -- 否 --> L[_sync_weights_and_save]
L --> M[restart_inactive_workers + bind_train_rollout]
M --> N[weight_update: update_targets 含 INACTIVE]
N --> O[mark pending -> ACTIVE]
D --> P[controller.generate except Exception -> FAILED]
classDef changed fill:#dff0d8,stroke:#3c763d
classDef problem fill:#f2dede,stroke:#a94442,stroke-width:2px
class B,C,H,I,J,L changed
class G,K,N,O,P problem
核心原理实现与单测
抽象与信息隐藏评估
公开 Interface 的线性业务流程评估
单测建议
其他 Issues
VerdictREQUEST_CHANGES |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| try: | ||
| self._ps.register_checkpoint(name, files=[], named_tensors=shard, use_shared_memory_pool=True) | ||
| if self._sync_after_register: | ||
| DEVICE_MODULE.synchronize() | ||
| except Exception: | ||
| self.logger.error("[checkpoint_engine] register_checkpoint failed rank={self.rank} name={name}") | ||
| self._checkpoint_name = name |
There was a problem hiding this comment.
Claude: [正确性] register_checkpoint 抛异常后仍执行 self._checkpoint_name = name，导致 has_registered_checkpoint() 误报成功，后台恢复线程会用未注册的 checkpoint 调 gather_metas/update。另外这条 error 日志缺 f 前缀，会原样打印 {self.rank}。建议异常时不设置 _checkpoint_name 并向上抛出。
| except Exception as e: | ||
| self.logger.exception(f"RolloutController.generate failed: session_id={session_id}") | ||
| rollout_state.status = Status.FAILED | ||
| rollout_state.error_msg = f"Rollout request failed: {type(e).__name__}: {str(e)[:1024]}" | ||
| return rollout_state |
There was a problem hiding this comment.
Claude: [正确性] 兜底 except Exception 把所有 worker 异常统一降级为终态 Status.FAILED，原本可区分/可重试的失败会被掩盖。ProduceBatchResult impact: 直接改变 leftover_failed 与 produced_samples/produced_tokens 计数。建议只捕获明确可归类为 FAILED 的异常类型。
There was a problem hiding this comment.
这里增加这个except的出发点是啥？如果是rollout worker出现了问题，会在rollout worker内部的完成故障处理，将返回的rollout state的status标为FAILED。那么从rollout worker到rollout controller的返回就不应该会有其他的报错了
There was a problem hiding this comment.
这儿加except是因为，在训练时，如果有挂掉的推理引擎，开partial rollout的情况，这里并不会立即把请求转给其他worker，反而是一直在重试，rollout_state.error_msg 这个会很大，然后一直返回这个很大的信息，然后机器就会崩掉。
| def get_target_state_worker_groups(self, target_state: WorkerLifecycleState) -> tuple[WorkerGroup, ...]: | ||
| """Return lifecycle groups containing workers in the requested | ||
| state.""" | ||
| with self._lock: | ||
| worker_groups = self._build_worker_groups() | ||
| matched_groups = [ | ||
| group | ||
| for group in worker_groups.values() | ||
| if any(worker.lifecycle_state is target_state for worker in group.workers) | ||
| ] | ||
| return tuple(sorted(matched_groups, key=lambda group: group.ranks)) |
There was a problem hiding this comment.
Claude: [正确性] group 用 any() 匹配、set_groups_state 按 rank 过滤 source_state：只要 group 内有一个 rank 是 PENDING_WEIGHTS 就整组被选中，而组内仍为 INACTIVE 的 rank 保持不变，随后整组被 notify_worker_group_recovered 重新注册进路由，entrypoint 会收到实际不可用的流量。建议整组状态一致才允许迁移。
| except Exception as e: | ||
| self.logger.exception(f"RolloutController.generate failed: session_id={session_id}") | ||
| rollout_state.status = Status.FAILED | ||
| rollout_state.error_msg = f"Rollout request failed: {type(e).__name__}: {str(e)[:1024]}" | ||
| return rollout_state |
There was a problem hiding this comment.
这里增加这个except的出发点是啥？如果是rollout worker出现了问题，会在rollout worker内部的完成故障处理，将返回的rollout state的status标为FAILED。那么从rollout worker到rollout controller的返回就不应该会有其他的报错了
|
|
||
| def _broadcast_to_active_workers(self, method_name: str, **kwargs): | ||
| workers = self.registry.active_workers() | ||
| def _broadcast_to_workers(self, method_name: str, target_state: WorkerLifecycleState, **kwargs): |
There was a problem hiding this comment.
这个函数为啥要修改？看上去调用点还是 target_state=WorkerLifecycleState.ACTIVE
There was a problem hiding this comment.
这里是因为增加了pending状态，所有给active worker的特性函数，比如onload_weights、onload_kvcache等都需要再给pending状态再写一遍，这里加了一个状态参数，这样就可以使用同一个函数，传递不同的参数控制对哪些状态的worker进行操作了
|
@claude review |
2 participants
No description provided.