misc

Author: specture724

checkpoint_engine/ps.py

@@ -1,6 +1,5 @@
 import os
 import threading
-import time
 from collections import defaultdict
 from collections.abc import Callable
 from datetime import timedelta
@@ -162,80 +161,6 @@ def _get_master_port(master_port: int | None = None) -> int:
     return master_port
 
 
-class P2PStore:
-    def __init__(self, device_manager: DeviceManager):
-        from mooncake.engine import TransferEngine
-
-        self.rank = int(os.getenv("RANK"))
-        gpu_count = device_manager.device_module.device_count()
-        local_rank = self.rank % gpu_count
-        device_type = device_manager.device_type
-        if device_type == "npu" and os.getenv("PS_P2P_STORE_RDMA_DEVICES") is None:
-            self.device = ""
-        else:
-            self.device = _get_my_rdma_device(local_rank, gpu_count, _get_rdma_devices())
-        self.ip = get_ip()
-
-        # we will start at most 8 ps processes, so we use 8 retries to avoid port conflicts in extreme cases
-        retry_count = 8
-        for i in range(retry_count):
-            self.engine = TransferEngine()
-            ret = self.engine.initialize(
-                self.ip,
-                "P2PHANDSHAKE",
-                "ascend_direct" if device_type == "npu" else "rdma",
-                self.device,
-            )
-            if ret == 0:
-                break
-            # sleep 0.5 ~ 2.0s, to avoid port conflicts when two processes retry at the same time
-            sleep_ms = random.randint(500, 2000)
-            logger.warning(
-                f"[rank{self.rank}] fail to initialize transfer engine, ret {ret}, retry {i + 1}/{retry_count} in {sleep_ms}ms"
-            )
-            time.sleep(sleep_ms / 1000)
-        else:
-            raise RuntimeError(f"[rank{self.rank}] fail to initialize transfer engine")
-        self.port = self.engine.get_rpc_port()
-        self.named_tensors: dict[str, torch.Tensor] = {}
-        logger.info(
-            f"[rank{self.rank}] p2p store initialized, addr is {self.addr}, rdma device is {self.device}"
-        )
-
-    @property
-    def addr(self) -> str:
-        return f"{self.ip}:{self.port}"
-
-    def register_named_tensors(self, named_tensors: dict[str, torch.Tensor]):
-        buffer_addresses = [tensor.data_ptr() for tensor in named_tensors.values()]
-        capacities = [tensor.nbytes for tensor in named_tensors.values()]
-        self.named_tensors.update(named_tensors)
-        for i, name in enumerate(named_tensors.keys()):
-            logger.info(
-                f"[rank{self.rank}] p2p store register tensor {name} with addr {hex(buffer_addresses[i])} and capacity {capacities[i]}"
-            )
-        assert self.engine.batch_register_memory(buffer_addresses, capacities) == 0
-
-    def unregister_named_tensors(self, names: list[str]) -> int:
-        buffer_addresses = [self.named_tensors[name].data_ptr() for name in names]
-        assert self.engine.batch_unregister_memory(buffer_addresses) == 0
-        num_unregistered = 0
-        for i, name in enumerate(names):
-            del self.named_tensors[name]
-            logger.info(
-                f"[rank{self.rank}] p2p store unregister tensor {name} with addr {hex(buffer_addresses[i])}"
-            )
-            num_unregistered += 1
-        return num_unregistered
-
-    def batch_transfer_sync_read(
-        self, target_hostname: str, buf_ptrs: list[int], remote_ptrs: list[int], lens: list[int]
-    ):
-        assert (
-            self.engine.batch_transfer_sync_read(target_hostname, buf_ptrs, remote_ptrs, lens) == 0
-        )
-
-
 class ParameterServer:
     shared_memory_pool_name = "__shared_memory_pool__"