checkpoint-engine 0.3.0rc0__py3-none-any.whl → 0.3.1__py3-none-any.whl

checkpoint_engine/worker.py

@@ -1,6 +1,7 @@
 import gc
 import traceback
 from collections.abc import Callable
+from functools import cached_property
 from typing import TypedDict
 
 import torch
@@ -117,6 +118,21 @@ class VllmColocateWorkerExtension:
         `worker_extension_cls` argument when initializing the vLLM worker.
     """
 
+    @cached_property
+    def _device_uuid(self) -> str:
+        from vllm.platforms import current_platform
+
+        if current_platform.device_type == "cuda":
+            return current_platform.get_device_uuid(self.device.index)
+        elif current_platform.device_type == "npu":
+            return f"NPU-{npu_generate_uuid()}"
+        else:
+            raise ValueError(f"Unsupported device type: {current_platform.device_type}")
+
+    @cached_property
+    def _zmq_ctx(self) -> zmq.Context:
+        return zmq.Context()
+
     def update_weights_from_ipc(self, zmq_handles: dict[str, str]):
         """
         Update model weights from checkpoint-engine via IPC communication.
@@ -149,17 +165,10 @@ class VllmColocateWorkerExtension:
         if current_platform.device_type == "npu" and self.device is None:
             self.device = torch.device(f"npu:{self.local_rank}")
         assert self.device is not None
-        if not hasattr(self, "_zmq_ctx") or self._zmq_ctx is None:
-            self._zmq_ctx = zmq.Context()
-        if current_platform.device_type == "cuda":
-            device_uuid = current_platform.get_device_uuid(self.device.index)
-        elif current_platform.device_type == "npu":
-            device_uuid = f"NPU-{npu_generate_uuid()}"
-        else:
-            raise ValueError(f"Unsupported device type: {current_platform.device_type}")
+
         update_weights_from_ipc(
             self._zmq_ctx,
-            zmq_handles[device_uuid],
+            zmq_handles[self._device_uuid],
             device_id=self.device.index,
             run=self.model_runner.model.load_weights,
             post_hook=lambda: process_weights_after_loading(

{checkpoint_engine-0.3.0rc0.dist-info → checkpoint_engine-0.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpoint-engine
-Version: 0.3.0rc0
+Version: 0.3.1
 Summary: checkpoint-engine is a lightweight, decoupling and efficient weight update middleware
 Project-URL: Homepage, https://github.com/MoonshotAI/checkpoint-engine
 Project-URL: Repository, https://github.com/MoonshotAI/checkpoint-engine

checkpoint_engine-0.3.1.dist-info/RECORD

@@ -0,0 +1,15 @@
+checkpoint_engine/__init__.py,sha256=OeWxe9mxl2sZ6cW-blSTg6JbFlOMpGbBghLZtxGOqXk,942
+checkpoint_engine/__main__.py,sha256=yzQlApuYo6eIOqtqM018RosyxNzXzB5a-stxUvsh-dg,709
+checkpoint_engine/_version.py,sha256=gGLpQUQx-ty9SEy9PYw9OgJWWzJLBnCpfJOfzL7SjlI,704
+checkpoint_engine/api.py,sha256=JDiQ4i3Gb6GoaBhlp8lNuUPaVURoFFdeGJY9ZDDGvPc,3518
+checkpoint_engine/data_types.py,sha256=O9uAXjwB20iwrOHfEEQd8Y9CmaFspNJ9ks9noHqwQKk,2716
+checkpoint_engine/device_utils.py,sha256=iKrof60j3CY3fStRTq3DRTt_kE1vYoEWHhAeyh0lByA,3020
+checkpoint_engine/p2p_store.py,sha256=abiCDVmRISPt9QFfavHB9Jo7ZpBbSjUS1NevGuB-AVA,8721
+checkpoint_engine/pin_memory.py,sha256=gpoe_z5XxbWkCvFLaXXpyUUFetBXUjsOrxBSX-ksZTw,16141
+checkpoint_engine/ps.py,sha256=0d68Sqb_y3H6b5H37exMbghDJ294VKaGqoWkcKE-Ao8,40316
+checkpoint_engine/worker.py,sha256=ghj9d2u8hY_U2uiOZWIN2CqRNZH6PrzujT22fHUFBWI,6879
+checkpoint_engine-0.3.1.dist-info/licenses/LICENCE,sha256=D3gPmHKpGtF1yxYNhqjtBtZY_brZjDotJTzpnmClzlY,1067
+checkpoint_engine-0.3.1.dist-info/METADATA,sha256=RSkQaKNs4euXk162_9AQzuNWaJQT4gMTnQ-3QJeYY_E,11559
+checkpoint_engine-0.3.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+checkpoint_engine-0.3.1.dist-info/top_level.txt,sha256=66sik_1eLakLYmcllOEJzFaNbSfjsueuP0tHYEzhMSs,18
+checkpoint_engine-0.3.1.dist-info/RECORD,,

checkpoint_engine-0.3.0rc0.dist-info/RECORD

@@ -1,10 +0,0 @@
-checkpoint_engine/__init__.py,sha256=Zj4I008kn9R6fYr0lVBzcQSnvckLpX2s1ljCOOqV1c8,87
-checkpoint_engine/_version.py,sha256=v0iyeXv9HxMc4JmYu_bJTIGKXRQVfpijACyjq2P_sk0,714
-checkpoint_engine/device_utils.py,sha256=iKrof60j3CY3fStRTq3DRTt_kE1vYoEWHhAeyh0lByA,3020
-checkpoint_engine/ps.py,sha256=eIvg_eI7HMedacoQQer62NRnGDjANtxsHVxgM93ccXQ,66977
-checkpoint_engine/worker.py,sha256=f6kS1ushIXxkRCEHXM5wVofUer9OxRiVY03vmKYLzgo,6757
-checkpoint_engine-0.3.0rc0.dist-info/licenses/LICENCE,sha256=D3gPmHKpGtF1yxYNhqjtBtZY_brZjDotJTzpnmClzlY,1067
-checkpoint_engine-0.3.0rc0.dist-info/METADATA,sha256=iVd2qPdNyTPPX3XIEiuM0ASk8As72zSGfFIYicpZG3E,11562
-checkpoint_engine-0.3.0rc0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-checkpoint_engine-0.3.0rc0.dist-info/top_level.txt,sha256=66sik_1eLakLYmcllOEJzFaNbSfjsueuP0tHYEzhMSs,18
-checkpoint_engine-0.3.0rc0.dist-info/RECORD,,