P6 round 1 / P6-26: Use the transfer selector for P2P checksum verification

Consolidate the raw fixes without changing their final tree.

Raw commits:
- bc2e74eefbcca7ab18b5950efcfac89df8211014

miles/backends/training_utils/weight_update/protocols/p2p.py

@@ -161,7 +161,7 @@ class UpdateWeightP2P(WeightTransferProtocol):
         self.disconnect()
         self.rollout_engines = rollout_engines
         self.cell_updaters_of_cell_id = {
-            cell_id: _P2PRolloutCellUpdater(args=self.args, cell_id=cell_id, api_client=api_client)
+            cell_id: _P2PRolloutCellUpdater(args=self.args, cell_id=cell_id, api_client=api_client, selector=selector)
             for api_client, cell_id in zip(rollout_engines, engine_cell_ids, strict=True)
         }
 

miles/backends/training_utils/weight_update/protocols/p2p_rollout_cell_updater.py

@@ -21,8 +21,10 @@ class _P2PRolloutCellUpdater(_RolloutCellUpdater):
         args: Namespace,
         cell_id: str,
         api_client: SGLangApiClient,
+        selector: str,
     ) -> None:
         super().__init__(args=args, cell_id=cell_id, api_client=api_client)
+        self.selector = selector
         self._executor = ThreadPoolExecutor(max_workers=1)
         self.targets_by_rollout_engine_rank: dict[int, RemoteWeightInfo] = {}
         self._pending_writes: list[Future[None]] = []
@@ -137,7 +139,9 @@ def _verify_transfer_checksums(
     names: list[str],
     sent_checksums: dict[str, str],
 ) -> None:
-    engine_body = cell_updater.submit_client_call("check_weights", action="raw_checksum", names=names).result()
+    engine_body = cell_updater.submit_client_call(
+        "check_weights", action="raw_checksum", names=names, selector=cell_updater.selector
+    ).result()
     if engine_body is None:
         return
     checksum_utils.verify_transfer_checksums(