feat: split files into multiple modules

Author: specture724

checkpoint_engine/__init__.py

@@ -2,3 +2,39 @@
     from ._version import __version__
 except ImportError:
     __version__ = "dev"
+
+from .api import request_inference_to_update
+from .data_types import (
+    BucketRange,
+    DataToGather,
+    H2DBucket,
+    MemoryBuffer,
+    MemoryBufferMetaList,
+    MemoryBufferMetas,
+    ParameterMeta,
+)
+from .device_utils import DeviceManager, get_ip, npu_generate_uuid
+from .p2p_store import P2PStore
+from .ps import ParameterServer
+from .worker import FlattenedTensorMetadata, VllmColocateWorkerExtension, update_weights_from_ipc
+
+
+__all__ = [
+    "BucketRange",
+    "DataToGather",
+    "DeviceManager",
+    "FlattenedTensorMetadata",
+    "H2DBucket",
+    "MemoryBuffer",
+    "MemoryBufferMetaList",
+    "MemoryBufferMetas",
+    "P2PStore",
+    "ParameterMeta",
+    "ParameterServer",
+    "VllmColocateWorkerExtension",
+    "__version__",
+    "get_ip",
+    "npu_generate_uuid",
+    "request_inference_to_update",
+    "update_weights_from_ipc",
+]

checkpoint_engine/ps.py

@@ -12,7 +12,6 @@
 from loguru import logger
 from torch.multiprocessing.reductions import reduce_tensor
 
-from checkpoint_engine.api import _init_api
 from checkpoint_engine.data_types import (
     BucketRange,
     DataToGather,
@@ -22,7 +21,6 @@
     MemoryBufferMetas,
     ParameterMeta,
 )
-
 from checkpoint_engine.device_utils import DeviceManager, get_ip, npu_generate_uuid
 from checkpoint_engine.p2p_store import P2PStore
 from checkpoint_engine.pin_memory import _ALIGN_SIZE, _register_checkpoint
@@ -31,6 +29,7 @@
 if TYPE_CHECKING:
     from checkpoint_engine.data_types import T
 
+
 def _to_named_tensor(metas: list[ParameterMeta], offset: int = 0) -> list[dict]:
     ret = []
     for meta in metas: