Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / transport.py: 84%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""CPU control, direct tensor broadcast, and configurable payload data planes.""" 

16# This distributed-data package is intentionally PyTorch-only. 

17 

18from __future__ import annotations 

19 

20import pickle 

21import struct 

22from dataclasses import dataclass 

23from typing import Any, Mapping, Sequence 

24 

25import torch # pylint: disable=forbidden-backend-import 

26import torch.distributed as dist # pylint: disable=forbidden-backend-import 

27 

28from hyper_parallel.distributed_data.model_batch import ModelBatchLayout 

29from hyper_parallel.distributed_data.schema import SampleKey 

30from hyper_parallel.distributed_data.topology import DataTopology 

31 

32@dataclass(frozen=True) 

33class DataGroups: 

34 """Control, payload, and model-consumer groups for distributed data.""" 

35 

36 data_plane_ranks: tuple[int, ...] 

37 control_group: Any 

38 payload_group: Any 

39 model_parallel_group: Any 

40 planner_rank: int 

41 distributed: bool 

42 model_parallel_tensor_group: Any = None 

43 

44 

45@dataclass(frozen=True) 

46class PreparedPayloadExchange: 

47 """Serialized host payload and collective tensor for one payload A2A.""" 

48 

49 input_splits: tuple[int, ...] 

50 send_storage: bytearray 

51 send_tensor: torch.Tensor 

52 local_segment: bytes 

53 

54 

55@dataclass 

56class ControlGatherWork: 

57 """Retain control buffers until the consumer waits and decodes on the launch stream.""" 

58 

59 encoded: torch.Tensor | None = None 

60 sizes: tuple[int, ...] = () 

61 gathered: tuple[torch.Tensor, ...] = () 

62 data_work: Any = None 

63 completed_values: tuple[Any, ...] | None = None 

64 

65 def wait(self) -> tuple[Any, ...]: 

66 """Wait for gathered bytes and decode them with the launch device and stream bound.""" 

67 if self.completed_values is None: 

68 if self.data_work is not None: 

69 self.data_work.wait() 

70 self.completed_values = tuple( 

71 _decode_control(payload, size) for payload, size in zip(self.gathered, self.sizes) 

72 ) 

73 return self.completed_values 

74 

75 

76@dataclass 

77class PayloadExchangeWork: 

78 """Asynchronous payload exchange owned by the launching device stream. 

79 

80 Waiting and decoding must use the launch stream and its rank-local device. 

81 Retaining the prepared send tensor keeps its backing storage valid until 

82 the accelerator has consumed it. 

83 """ 

84 

85 prepared: PreparedPayloadExchange 

86 output_splits: tuple[int, ...] 

87 received_tensor: torch.Tensor | None 

88 data_work: Any = None 

89 completed_payloads: dict[SampleKey, Any] | None = None 

90 

91 def wait(self) -> dict[SampleKey, Any]: 

92 """Wait and decode on the launch stream with its rank-local device bound.""" 

93 if self.completed_payloads is not None: 

94 return self.completed_payloads 

95 if self.data_work is not None: 

96 self.data_work.wait() 

97 if self.received_tensor is None: 

98 received_bytes = b"" 

99 else: 

100 received_bytes = self.received_tensor.cpu().numpy().tobytes() 

101 self.completed_payloads = _decode_received_payloads(received_bytes, self.output_splits) 

102 return self.completed_payloads 

103 

104 

105def _is_accelerator_backend(backend: str) -> bool: 

106 return "hccl" in backend.lower() or "nccl" in backend.lower() 

107 

108 

109def _control_backend(group: Any, backend: str | None = None) -> str: 

110 if backend is not None: 

111 return backend.lower() 

112 try: 

113 return str(dist.get_backend(group)).lower() 

114 except (RuntimeError, ValueError): 

115 # Lightweight unit-test groups and single-process fakes represent Gloo 

116 # groups with opaque sentinels rather than registered ProcessGroups. 

117 return "gloo" 

118 

119 

120def _require_control_device(device: Any, backend: str) -> torch.device: 

121 if device is None: 

122 raise ValueError(f"{backend} control collectives require an accelerator communication_device.") 

123 resolved = torch.device(device) 

124 expected = "npu" if "hccl" in backend else "cuda" 

125 if resolved.type != expected: 

126 raise ValueError(f"{backend} control collectives require a {expected} device, got {resolved}.") 

127 return resolved 

128 

129 

130def _encode_control(value: Any, device: Any) -> tuple[torch.Tensor, int]: 

131 encoded = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL) 

132 resolved = torch.device(device) 

133 payload = torch.zeros((max(1, len(encoded)),), dtype=torch.uint8, device=resolved) 

134 if encoded: 

135 payload[:len(encoded)] = torch.tensor(list(encoded), dtype=torch.uint8, device=resolved) 

136 return payload, len(encoded) 

137 

138 

139def _decode_control(payload: torch.Tensor, size: int) -> Any: 

140 encoded = payload[:size].cpu().numpy().tobytes() 

141 return pickle.loads(encoded) 

142 

143 

144def all_gather_control_object( 

145 value: Any, 

146 *, 

147 group: Any = None, 

148 device: Any = None, 

149 backend: str | None = None, 

150) -> tuple[Any, ...]: 

151 """Gather a small control object with Gloo or an accelerator backend. 

152 

153 Args: 

154 value: Rank-local serializable control value. 

155 group: Process group whose ranks participate in the gather. 

156 device: Rank-local accelerator used by HCCL/NCCL. 

157 backend: Explicit backend for startup calls that use WORLD. 

158 

159 Returns: 

160 Gathered values in process-group rank order. 

161 """ 

162 return _begin_all_gather_control_object(value, group=group, device=device, backend=backend).wait() 

163 

164 

165def _begin_all_gather_control_object( 

166 value: Any, 

167 *, 

168 group: Any = None, 

169 device: Any = None, 

170 backend: str | None = None, 

171) -> ControlGatherWork: 

172 """Exchange lengths on the caller, then launch accelerator bytes without decoding.""" 

173 effective_backend = _control_backend(group, backend) 

174 if not _is_accelerator_backend(effective_backend): 

175 gathered = [None] * dist.get_world_size(group) 

176 dist.all_gather_object(gathered, value, group=group) 

177 return ControlGatherWork(completed_values=tuple(gathered)) 

178 resolved = _require_control_device(device, effective_backend) 

179 encoded, encoded_size = _encode_control(value, resolved) 

180 local_size = torch.tensor([encoded_size], dtype=torch.int64, device=resolved) 

181 sizes = [torch.empty_like(local_size) for _ in range(dist.get_world_size(group))] 

182 dist.all_gather(sizes, local_size, group=group) 

183 host_sizes = tuple(int(item.item()) for item in sizes) 

184 max_size = max(host_sizes) 

185 if encoded.numel() < max_size: 

186 padded = torch.zeros((max_size,), dtype=torch.uint8, device=resolved) 

187 padded[:encoded.numel()] = encoded 

188 encoded = padded 

189 gathered = [torch.empty_like(encoded) for _ in sizes] 

190 work = dist.all_gather(gathered, encoded, group=group, async_op=True) 

191 return ControlGatherWork(encoded=encoded, sizes=host_sizes, gathered=tuple(gathered), data_work=work) 

192 

193 

194def broadcast_control_object( 

195 value: Any, 

196 *, 

197 src: int, 

198 group: Any = None, 

199 device: Any = None, 

200) -> Any: 

201 """Broadcast a small control object with Gloo or an accelerator backend. 

202 

203 Args: 

204 value: Source value or a placeholder on receivers. 

205 src: Global source rank. 

206 group: Process group whose ranks participate in the broadcast. 

207 device: Rank-local accelerator used by HCCL/NCCL. 

208 

209 Returns: 

210 The source value on every participating rank. 

211 """ 

212 effective_backend = _control_backend(group) 

213 if not _is_accelerator_backend(effective_backend): 

214 payload = [value] 

215 dist.broadcast_object_list(payload, src=src, group=group) 

216 return payload[0] 

217 resolved = _require_control_device(device, effective_backend) 

218 rank = dist.get_rank() 

219 encoded, encoded_size = _encode_control(value, resolved) if rank == src else ( 

220 torch.zeros((1,), dtype=torch.uint8, device=resolved), 0 

221 ) 

222 size = torch.tensor([encoded_size], dtype=torch.int64, device=resolved) 

223 dist.broadcast(size, src=src, group=group) 

224 target_size = int(size.item()) 

225 if encoded.numel() < max(1, target_size): 

226 encoded = torch.zeros((max(1, target_size),), dtype=torch.uint8, device=resolved) 

227 dist.broadcast(encoded, src=src, group=group) 

228 return _decode_control(encoded, target_size) 

229 

230 

231def create_data_groups( 

232 topology: DataTopology, 

233 dataset_reader_ranks: tuple[int, ...], 

234 planner_rank: int, 

235 *, 

236 cpu_backend: str, 

237 payload_backend: str | None, 

238 communication_device: Any, 

239 enable_payload_exchange: bool, 

240 model_device: Any = None, 

241) -> DataGroups: 

242 """Create control, optional payload, and model groups deterministically. 

243 

244 Args: 

245 topology: Constructor and model-consumer topology. 

246 dataset_reader_ranks: Global ranks that materialize raw samples. 

247 planner_rank: Global rank that creates loading plans. 

248 cpu_backend: Backend used for control and model-parallel batch broadcast. 

249 payload_backend: Optional backend for online payload A2A. When omitted 

250 with an accelerator communication device, the WORLD backend is used. 

251 communication_device: Optional rank-local device for payload tensors. 

252 enable_payload_exchange: Whether the online path needs a payload group. 

253 model_device: Training device for direct model-peer tensor broadcast, 

254 independent of the raw payload backend. 

255 

256 Returns: 

257 Groups relevant to the current rank. 

258 """ 

259 data_plane_ranks = tuple(sorted(set(dataset_reader_ranks) | set(topology.constructor_ranks))) 

260 if planner_rank not in data_plane_ranks: 

261 raise ValueError(f"planner_rank {planner_rank} must belong to data-plane ranks {data_plane_ranks}.") 

262 distributed = dist.is_available() and dist.is_initialized() 

263 if not distributed: 

264 return _single_rank_data_groups(topology, data_plane_ranks, planner_rank) 

265 _validate_distributed_topology(topology) 

266 _validate_group_backends(cpu_backend, payload_backend, enable_payload_exchange) 

267 effective_payload_backend, reuse_control_group = _resolve_payload_backend( 

268 data_plane_ranks=data_plane_ranks, 

269 cpu_backend=cpu_backend, 

270 payload_backend=payload_backend, 

271 communication_device=communication_device, 

272 enable_payload_exchange=enable_payload_exchange, 

273 ) 

274 control_group, payload_group = _create_data_plane_process_groups( 

275 topology=topology, 

276 data_plane_ranks=data_plane_ranks, 

277 cpu_backend=cpu_backend, 

278 payload_backend=effective_payload_backend, 

279 reuse_control_group=reuse_control_group, 

280 enable_payload_exchange=enable_payload_exchange, 

281 ) 

282 tensor_backend = _resolve_model_tensor_backend(communication_device, payload_backend) 

283 model_device = torch.device(model_device) if model_device is not None else None 

284 if model_device is not None and model_device.type in ("cuda", "npu"): 

285 tensor_backend = "nccl" if model_device.type == "cuda" else "hccl" 

286 model_parallel_group, model_parallel_tensor_group = _create_model_parallel_process_groups( 

287 topology, cpu_backend, tensor_backend, 

288 ) 

289 return DataGroups( 

290 data_plane_ranks, control_group, payload_group, model_parallel_group, planner_rank, True, 

291 model_parallel_tensor_group, 

292 ) 

293 

294 

295def _single_rank_data_groups( 

296 topology: DataTopology, 

297 data_plane_ranks: tuple[int, ...], 

298 planner_rank: int, 

299) -> DataGroups: 

300 if len(topology.rank_list) != 1: 

301 raise ValueError("torch.distributed must be initialized for a mesh containing more than one rank.") 

302 return DataGroups(data_plane_ranks, None, None, None, planner_rank, False, None) 

303 

304 

305def _validate_distributed_topology(topology: DataTopology) -> None: 

306 world_ranks = tuple(range(dist.get_world_size())) 

307 if set(topology.rank_list) != set(world_ranks): 

308 raise ValueError( 

309 "The distributed data root mesh must contain every torch.distributed world rank so all ranks create " 

310 "service groups in the same order." 

311 ) 

312 if dist.get_rank() != topology.global_rank: 

313 raise ValueError("Topology global_rank does not match torch.distributed rank.") 

314 

315 

316def _validate_group_backends( 

317 cpu_backend: str, 

318 payload_backend: str | None, 

319 enable_payload_exchange: bool, 

320) -> None: 

321 """Validate control and payload backend options before creating process groups.""" 

322 if not isinstance(cpu_backend, str) or not cpu_backend: 

323 raise ValueError("cpu_backend must be a non-empty backend name.") 

324 if cpu_backend.lower() not in ("gloo", "hccl"): 

325 raise ValueError("cpu_backend must be 'gloo' or 'hccl'.") 

326 if payload_backend is not None and (not isinstance(payload_backend, str) or not payload_backend): 

327 raise ValueError("payload_backend must be a non-empty backend name or None.") 

328 if not isinstance(enable_payload_exchange, bool): 

329 raise ValueError("enable_payload_exchange must be boolean.") 

330 

331 

332def _resolve_payload_backend( 

333 *, 

334 data_plane_ranks: tuple[int, ...], 

335 cpu_backend: str, 

336 payload_backend: str | None, 

337 communication_device: Any, 

338 enable_payload_exchange: bool, 

339) -> tuple[str | None, bool]: 

340 """Choose a compatible payload backend and whether it can reuse the control group.""" 

341 if not enable_payload_exchange or len(data_plane_ranks) == 1: 

342 return None, False 

343 device = torch.device(communication_device) if communication_device is not None else None 

344 effective_backend = payload_backend or _default_payload_backend(device, cpu_backend) 

345 normalized_backend = effective_backend.lower() 

346 _validate_payload_device(normalized_backend, effective_backend, device) 

347 reuse_control_group = normalized_backend == cpu_backend.lower() and (device is None or device.type == "cpu") 

348 return effective_backend, reuse_control_group 

349 

350 

351def _default_payload_backend(device: torch.device | None, cpu_backend: str) -> str: 

352 if device is not None and device.type != "cpu": 

353 return str(dist.get_backend()) 

354 return cpu_backend 

355 

356 

357def _validate_payload_device( 

358 normalized_backend: str, 

359 effective_backend: str, 

360 device: torch.device | None, 

361) -> None: 

362 """Require the communication device type expected by the payload backend.""" 

363 accelerator_backend = "hccl" in normalized_backend or "nccl" in normalized_backend 

364 if accelerator_backend and (device is None or device.type == "cpu"): 

365 raise ValueError( 

366 f"Payload backend {effective_backend!r} requires a rank-local accelerator communication_device." 

367 ) 

368 expected_device_type = None 

369 if "nccl" in normalized_backend: 

370 expected_device_type = "cuda" 

371 elif "hccl" in normalized_backend: 

372 expected_device_type = "npu" 

373 if expected_device_type is not None and device is not None and device.type != expected_device_type: 

374 raise ValueError( 

375 f"Payload backend {effective_backend!r} requires a {expected_device_type} device, but got {device}." 

376 ) 

377 if not accelerator_backend and device is not None and device.type != "cpu": 

378 raise ValueError(f"Payload backend {effective_backend!r} cannot exchange tensors on device {device}.") 

379 

380 

381def _create_data_plane_process_groups( 

382 *, 

383 topology: DataTopology, 

384 data_plane_ranks: tuple[int, ...], 

385 cpu_backend: str, 

386 payload_backend: str | None, 

387 reuse_control_group: bool, 

388 enable_payload_exchange: bool, 

389) -> tuple[Any, Any]: 

390 """Create control and payload groups in common rank order, returning owned handles.""" 

391 if len(data_plane_ranks) == 1: 

392 return None, None 

393 created_control_group = dist.new_group(ranks=list(data_plane_ranks), backend=cpu_backend) 

394 is_member = topology.global_rank in data_plane_ranks 

395 control_group = created_control_group if is_member else None 

396 if not enable_payload_exchange: 

397 return control_group, None 

398 if reuse_control_group: 

399 return control_group, control_group 

400 created_payload_group = dist.new_group(ranks=list(data_plane_ranks), backend=payload_backend) 

401 return control_group, created_payload_group if is_member else None 

402 

403 

404def _resolve_model_tensor_backend(communication_device: Any, payload_backend: str | None) -> str | None: 

405 """Select an accelerator backend for tensor model-group broadcasts when available.""" 

406 if communication_device is None: 

407 return None 

408 device = torch.device(communication_device) 

409 if device.type == "cpu": 

410 return None 

411 candidate = payload_backend or str(dist.get_backend()) 

412 normalized = candidate.lower() 

413 if "hccl" not in normalized and "nccl" not in normalized: 

414 return None 

415 return candidate 

416 

417 

418def _create_model_parallel_process_groups( 

419 topology: DataTopology, 

420 cpu_backend: str, 

421 tensor_backend: str | None, 

422) -> tuple[Any, Any]: 

423 """Create model-parallel control and tensor groups in common rank order.""" 

424 model_parallel_group = None 

425 model_parallel_tensor_group = None 

426 for rank_group in topology.model_parallel_rank_groups: 

427 if len(rank_group) == 1: 

428 continue 

429 is_member = topology.global_rank in rank_group 

430 created_group = dist.new_group(ranks=list(rank_group), backend=cpu_backend) 

431 if is_member: 

432 model_parallel_group = created_group 

433 if tensor_backend is not None: 

434 created_tensor_group = dist.new_group(ranks=list(rank_group), backend=tensor_backend) 

435 if is_member: 

436 model_parallel_tensor_group = created_tensor_group 

437 return model_parallel_group, model_parallel_tensor_group 

438 

439 

440def _encode_payload_segment(items: Sequence[tuple[SampleKey, Any]]) -> bytes: 

441 """Serialize an internal route with metadata and a contiguous binary buffer. 

442 

443 Tensor and byte leaves are kept out of the metadata pickle. The metadata 

444 contains only the nested sample structure and descriptors pointing into a 

445 contiguous byte region. A legacy pickle route remains available for 

446 payloads without binary leaves so old scalar-only unit/test routes retain 

447 their wire compatibility. 

448 """ 

449 if not items: 

450 return b"" 

451 if not any(_payload_contains_binary(payload) for _, payload in items): 

452 return pickle.dumps(tuple(items), protocol=pickle.HIGHEST_PROTOCOL) 

453 encoded_items = [] 

454 binary_chunks: list[bytes] = [] 

455 binary_offset = 0 

456 for key, payload in items: 

457 schema, binary_offset = _encode_payload_value(payload, binary_chunks, binary_offset) 

458 encoded_items.append((key, schema)) 

459 metadata = pickle.dumps(tuple(encoded_items), protocol=pickle.HIGHEST_PROTOCOL) 

460 binary = b"".join(binary_chunks) 

461 return _PAYLOAD_HEADER.pack(_PAYLOAD_MAGIC, len(metadata), len(binary)) + metadata + binary 

462 

463 

464def _decode_payload_segment(frame: bytes) -> tuple[tuple[SampleKey, Any], ...]: 

465 """Deserialize a route produced by peers running the same codec.""" 

466 if not frame: 

467 return () 

468 if not frame.startswith(_PAYLOAD_MAGIC): 

469 # Preserve compatibility with routes produced before the tensor-buffer 

470 # codec was introduced. This path is used only for scalar-only routes. 

471 return pickle.loads(frame) 

472 if len(frame) < _PAYLOAD_HEADER.size: 

473 raise ValueError("Payload frame is truncated before its header.") 

474 _, metadata_size, binary_size = _PAYLOAD_HEADER.unpack_from(frame) 

475 payload_offset = _PAYLOAD_HEADER.size 

476 payload_size = metadata_size + binary_size 

477 if payload_offset + payload_size != len(frame): 

478 raise ValueError( 

479 "Payload frame size does not match its metadata and binary descriptors." 

480 ) 

481 metadata_end = payload_offset + metadata_size 

482 encoded_items = pickle.loads(frame[payload_offset:metadata_end]) 

483 binary = bytearray(frame[metadata_end:]) 

484 decoded_items = [] 

485 for key, schema in encoded_items: 

486 decoded_items.append((key, _decode_payload_value(schema, binary))) 

487 return tuple(decoded_items) 

488 

489 

490_PAYLOAD_MAGIC = b"HPB1" 

491_PAYLOAD_HEADER = struct.Struct("!4sQQ") 

492_TENSOR_TAG = "__hp_tensor__" 

493_BYTES_TAG = "__hp_bytes__" 

494_BYTEARRAY_TAG = "__hp_bytearray__" 

495 

496 

497def _payload_contains_binary(value: Any) -> bool: 

498 """Return whether a payload contains a tensor or raw byte leaf.""" 

499 if torch.is_tensor(value) or isinstance(value, (bytes, bytearray)): 

500 return True 

501 if isinstance(value, Mapping): 

502 return any(_payload_contains_binary(child) for child in value.values()) 

503 if isinstance(value, (tuple, list)): 

504 return any(_payload_contains_binary(child) for child in value) 

505 return False 

506 

507 

508def _encode_payload_value( 

509 value: Any, 

510 binary_chunks: list[bytes], 

511 binary_offset: int, 

512) -> tuple[Any, int]: 

513 """Replace binary leaves with descriptors and append their bytes.""" 

514 if torch.is_tensor(value): 

515 tensor = value.detach() 

516 if tensor.device.type != "cpu": 

517 tensor = tensor.cpu() 

518 tensor = tensor.contiguous() 

519 # Flatten first: PyTorch does not permit a scalar tensor to change 

520 # dtype via ``view`` directly because it has no dimension to resize. 

521 raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes() 

522 descriptor = ( 

523 _TENSOR_TAG, 

524 str(tensor.dtype), 

525 tuple(tensor.shape), 

526 binary_offset, 

527 len(raw), 

528 ) 

529 binary_chunks.append(raw) 

530 return descriptor, binary_offset + len(raw) 

531 if isinstance(value, (bytes, bytearray)): 

532 raw = bytes(value) 

533 tag = _BYTES_TAG if isinstance(value, bytes) else _BYTEARRAY_TAG 

534 descriptor = (tag, binary_offset, len(raw)) 

535 binary_chunks.append(raw) 

536 return descriptor, binary_offset + len(raw) 

537 if isinstance(value, Mapping): 

538 encoded_items = [] 

539 for key, child in value.items(): 

540 encoded_child, binary_offset = _encode_payload_value(child, binary_chunks, binary_offset) 

541 encoded_items.append((key, encoded_child)) 

542 return ("__hp_mapping__", tuple(encoded_items)), binary_offset 

543 if isinstance(value, (tuple, list)): 

544 encoded_items = [] 

545 for child in value: 

546 encoded_child, binary_offset = _encode_payload_value(child, binary_chunks, binary_offset) 

547 encoded_items.append(encoded_child) 

548 tag = "__hp_tuple__" if isinstance(value, tuple) else "__hp_list__" 

549 return (tag, tuple(encoded_items)), binary_offset 

550 if value is None or isinstance(value, (bool, int, float, str)): 

551 return ("__hp_value__", value), binary_offset 

552 raise TypeError( 

553 "Payload tensor-buffer codec does not support a non-metadata value of " 

554 f"type {type(value).__qualname__}." 

555 ) 

556 

557 

558def _decode_payload_value(schema: Any, binary: bytearray) -> Any: 

559 """Reconstruct one sample from metadata descriptors and binary storage.""" 

560 if not isinstance(schema, tuple) or not schema: 

561 raise ValueError("Payload metadata contains an invalid schema node.") 

562 tag = schema[0] 

563 if tag == _TENSOR_TAG: 

564 return _decode_tensor_payload(schema, binary) 

565 if tag in (_BYTES_TAG, _BYTEARRAY_TAG): 

566 return _decode_bytes_payload(schema, binary) 

567 if tag == "__hp_mapping__": 

568 return { 

569 key: _decode_payload_value(child, binary) 

570 for key, child in schema[1] 

571 } 

572 if tag == "__hp_tuple__": 

573 return tuple(_decode_payload_value(child, binary) for child in schema[1]) 

574 if tag == "__hp_list__": 

575 return [_decode_payload_value(child, binary) for child in schema[1]] 

576 if tag == "__hp_value__": 

577 return schema[1] 

578 raise ValueError(f"Unknown payload metadata schema tag: {tag!r}.") 

579 

580 

581def _decode_tensor_payload(schema: tuple, binary: bytearray) -> torch.Tensor: 

582 """Reconstruct a tensor view after validating its descriptor.""" 

583 if len(schema) != 5: 

584 raise ValueError("Tensor payload descriptor has an invalid field count.") 

585 _, dtype_name, shape, offset, size = schema 

586 dtype = _payload_dtype(dtype_name) 

587 shape = tuple(shape) 

588 if any(not isinstance(dim, int) or dim < 0 for dim in shape): 

589 raise ValueError(f"Tensor payload descriptor has an invalid shape: {shape!r}.") 

590 if not isinstance(offset, int) or not isinstance(size, int) or offset < 0 or size < 0: 

591 raise ValueError("Tensor payload descriptor has an invalid byte range.") 

592 end = offset + size 

593 if end > len(binary): 

594 raise ValueError("Tensor payload descriptor points outside the binary buffer.") 

595 element_count = 1 

596 for dim in shape: 

597 element_count *= dim 

598 expected_size = element_count * torch.empty((), dtype=dtype).element_size() 

599 if expected_size != size: 

600 raise ValueError( 

601 f"Tensor payload byte size mismatch: expected {expected_size}, got {size}." 

602 ) 

603 if size == 0: 

604 return torch.empty(shape, dtype=dtype) 

605 tensor = torch.frombuffer(memoryview(binary)[offset:end], dtype=dtype, count=element_count) 

606 return tensor.reshape(shape) 

607 

608 

609def _decode_bytes_payload(schema: tuple, binary: bytearray) -> bytes | bytearray: 

610 """Reconstruct a byte sequence after validating its descriptor.""" 

611 if len(schema) != 3: 

612 raise ValueError("Byte payload descriptor has an invalid field count.") 

613 tag, offset, size = schema 

614 if not isinstance(offset, int) or not isinstance(size, int) or offset < 0 or size < 0: 

615 raise ValueError("Byte payload descriptor has an invalid byte range.") 

616 end = offset + size 

617 if end > len(binary): 

618 raise ValueError("Byte payload descriptor points outside the binary buffer.") 

619 raw = bytes(binary[offset:end]) 

620 return bytearray(raw) if tag == _BYTEARRAY_TAG else raw 

621 

622 

623def _payload_dtype(dtype_name: Any) -> torch.dtype: 

624 """Resolve the stable string form emitted by ``str(torch.dtype)``.""" 

625 if not isinstance(dtype_name, str) or not dtype_name.startswith("torch."): 

626 raise ValueError(f"Payload tensor descriptor has an invalid dtype: {dtype_name!r}.") 

627 dtype = getattr(torch, dtype_name[6:], None) 

628 if not isinstance(dtype, torch.dtype): 

629 raise ValueError(f"Unsupported payload tensor dtype: {dtype_name!r}.") 

630 return dtype 

631 

632 

633def _exchange_payload_sizes( 

634 input_splits: Sequence[int], 

635 control_group: Any, 

636 device: Any = None, 

637) -> list[int]: 

638 """Exchange per-peer byte counts and wait before reading receive sizes on Host.""" 

639 backend = _control_backend(control_group) 

640 size_device = _require_control_device(device, backend) if _is_accelerator_backend(backend) else torch.device("cpu") 

641 size_input = torch.tensor(tuple(input_splits), dtype=torch.int64, device=size_device) 

642 size_output = torch.empty_like(size_input) 

643 size_work = dist.all_to_all_single(size_output, size_input, group=control_group, async_op=True) 

644 if size_work is not None: 

645 size_work.wait() 

646 return [int(size) for size in size_output.tolist()] 

647 

648 

649def _decode_received_payloads( 

650 received_bytes: bytes, 

651 output_splits: Sequence[int], 

652) -> dict[SampleKey, Any]: 

653 """Merge routes without silently overwriting duplicate sample occurrences.""" 

654 payloads: dict[SampleKey, Any] = {} 

655 cursor = 0 

656 for segment_size in output_splits: 

657 segment_items = _decode_payload_segment(received_bytes[cursor:cursor + segment_size]) 

658 for key, payload in segment_items: 

659 if key in payloads: 

660 raise ValueError(f"Data Constructor received duplicate payload for {key}.") 

661 payloads[key] = payload 

662 cursor += segment_size 

663 return payloads 

664 

665 

666class DataPlaneTransport: 

667 """Metadata/control collectives and variable-byte sample all-to-all. 

668 

669 The coordinator calls these methods only on data-plane members. Group 

670 topology is fixed at build time; it is not revalidated for each collective. 

671 """ 

672 

673 def __init__(self, groups: DataGroups, global_rank: int, communication_device: Any = None) -> None: 

674 """Store service-group membership.""" 

675 self._ranks = groups.data_plane_ranks 

676 self._control_group = groups.control_group 

677 self._payload_group = groups.payload_group 

678 self._planner_rank = groups.planner_rank 

679 self._global_rank = global_rank 

680 self._is_member = global_rank in self._ranks 

681 self._local_index = self._ranks.index(global_rank) if self._is_member else None 

682 if self._is_member and len(self._ranks) > 1 and (not groups.distributed or self._control_group is None): 

683 raise ValueError("Multi-rank data-plane communication requires an initialized process group.") 

684 self._communication_device = ( 

685 torch.device(communication_device) if communication_device is not None else None 

686 ) 

687 

688 @property 

689 def is_member(self) -> bool: 

690 """Return whether this rank participates in Dataset Reader/Planner communication.""" 

691 return self._is_member 

692 

693 @property 

694 def planner_rank(self) -> int: 

695 """Return the global Planner rank.""" 

696 return self._planner_rank 

697 

698 @property 

699 def ranks(self) -> tuple[int, ...]: 

700 """Return data-plane ranks in collective order.""" 

701 return self._ranks 

702 

703 @property 

704 def communication_device(self) -> torch.device | None: 

705 """Return the rank-local accelerator used by payload collectives.""" 

706 return self._communication_device 

707 

708 @property 

709 def communication_backend(self) -> str: 

710 """Return the backend used by this data-plane control group.""" 

711 return _control_backend(self._control_group) 

712 

713 def all_gather_object(self, value: Any) -> tuple[Any, ...]: 

714 """Gather small control objects on every data-plane rank. 

715 

716 Args: 

717 value: Rank-local control value. 

718 

719 Returns: 

720 Values in data-plane rank order. 

721 """ 

722 return self.begin_all_gather_object(value).wait() 

723 

724 def begin_all_gather_object(self, value: Any) -> ControlGatherWork: 

725 """Launch control gathering and defer accelerator result decoding to the consumer. 

726 

727 Args: 

728 value: Rank-local control value. 

729 

730 Returns: 

731 Work retaining communication buffers. Wait on the launch device and 

732 stream before using the result. Accelerator length exchange remains 

733 synchronous; Gloo and single-rank results are already complete. 

734 """ 

735 if len(self._ranks) == 1: 

736 return ControlGatherWork(completed_values=(value,)) 

737 return _begin_all_gather_control_object( 

738 value, 

739 group=self._control_group, 

740 device=self._communication_device, 

741 ) 

742 

743 def all_ranks_true(self, value: bool) -> bool: 

744 """Return whether every data-plane rank supplied ``True``. 

745 

746 Args: 

747 value: Rank-local boolean value. 

748 

749 Returns: 

750 Whether every participating rank supplied ``True``. 

751 """ 

752 if len(self._ranks) == 1: 

753 return value 

754 backend = _control_backend(self._control_group) 

755 flag_device = ( 

756 _require_control_device(self._communication_device, backend) 

757 if _is_accelerator_backend(backend)