Update the checkpoint to inference engine. This function should be called after gather_metas. Warning: if _auto_pg is False when initializing ParameterServer, please make sure ALL ranks in the WORLD_SIZE call `update` function, otherwise, it will hang. Args:
(
self,
checkpoint_name: str,
req_func: Callable[[list[tuple[str, str]]], None],
*,
timeout: timedelta = timedelta(minutes=10),
ranks: list[int] | None = None,
)
| 542 | ) |
| 543 | |
| 544 | def update( |
| 545 | self, |
| 546 | checkpoint_name: str, |
| 547 | req_func: Callable[[list[tuple[str, str]]], None], |
| 548 | *, |
| 549 | timeout: timedelta = timedelta(minutes=10), |
| 550 | ranks: list[int] | None = None, |
| 551 | ) -> None: |
| 552 | """ |
| 553 | Update the checkpoint to inference engine. This function should be called after gather_metas. |
| 554 | Warning: if _auto_pg is False when initializing ParameterServer, please make sure ALL ranks in the WORLD_SIZE call `update` function, |
| 555 | otherwise, it will hang. |
| 556 | |
| 557 | Args: |
| 558 | checkpoint_name: The name of the checkpoint. |
| 559 | req_func: The function to request the inference of inference engine. |
| 560 | ranks: The ranks to update. If not set, will use fully broadcast to update to all ranks, |
| 561 | which is the fastest way to update weights, especially in colocated architecture. |
| 562 | If set, will use p2p to update to the ranks, this is flexible to update to a group of ranks, |
| 563 | which is useful in disaggregated architecture. |
| 564 | master_addr: The master address for process group initialization. If not set, will use env MASTER_ADDR. |
| 565 | master_port: The master port for process group initialization. If not set, will use _get_master_port to get the port, which will use MASTER_PORT+1. |
| 566 | timeout: The timeout of the barrier operation. |
| 567 | """ |
| 568 | assert req_func is not None, "req_func is required" |
| 569 | ranks_group = None |
| 570 | try: |
| 571 | if self._auto_pg and not dist.is_initialized(): |
| 572 | self.init_process_group(timeout=timeout) |
| 573 | # if ranks is None or [], it will use fully broadcast to update to all ranks |
| 574 | ranks_group = dist.new_group(ranks) if ranks else None |
| 575 | self._update_per_bucket(checkpoint_name, req_func, ranks_group, ranks) |
| 576 | self.store_based_barrier() |
| 577 | except Exception as e: |
| 578 | logger.exception( |
| 579 | f"[rank{self._rank}] update checkpoint {checkpoint_name} with ranks {ranks} error {e}" |
| 580 | ) |
| 581 | raise |
| 582 | finally: |
| 583 | if ranks_group: |
| 584 | dist.destroy_process_group(ranks_group) |
| 585 | if self._auto_pg and dist.is_initialized(): |
| 586 | dist.destroy_process_group() |
| 587 | self.device_manager.device_module.empty_cache() |
| 588 | logger.info( |
| 589 | f"[rank{self._rank}] update checkpoint {checkpoint_name} with ranks {ranks} done. " |
| 590 | f"Current device allocated {self.device_manager.device_module.memory_allocated() / 1024 / 1024} MB, " |
| 591 | f"reserved {self.device_manager.device_module.memory_reserved() / 1024 / 1024} MB." |
| 592 | ) |
| 593 | |
| 594 | def _bind_zmq_socket(self) -> tuple[zmq.Socket, list[tuple[str, str]]]: |
| 595 | def zmq_handle(device_uuid: str) -> str: |