checkpoint-engine 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

checkpoint_engine/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.1'
-__version_tuple__ = version_tuple = (0, 1, 1)
+__version__ = version = '0.1.3'
+__version_tuple__ = version_tuple = (0, 1, 3)
 
-__commit_id__ = commit_id = 'gf29b2e1c3'
+__commit_id__ = commit_id = None

checkpoint_engine/ps.py

@@ -1,30 +1,28 @@
-from __future__ import annotations
-
 import argparse
 import concurrent.futures
+import ctypes
 import os
 import pickle
 import random
 import socket
-import subprocess
 import threading
 import time
-import uuid
 from collections import defaultdict
+from collections.abc import Callable
 from datetime import timedelta
-from functools import cached_property, lru_cache
-from typing import Callable, NamedTuple
+from functools import lru_cache
+from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, NamedTuple
 
+import httpx
 import numpy as np
-import requests
 import torch
 import torch.distributed as dist
 import zmq
 from loguru import logger
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, PlainSerializer, PlainValidator, WithJsonSchema
 from safetensors.torch import safe_open
 from torch.multiprocessing.reductions import reduce_tensor
-from typing_extensions import TYPE_CHECKING
+
 
 if TYPE_CHECKING:
     from typing_extensions import TypedDict
@@ -37,16 +35,59 @@ if TYPE_CHECKING:
         tp_concat_dim: int
 
 
-class ParameterMeta(BaseModel):
-    # now all classes are changed to pydantic BaseModel
-    # it will directly report validation errors for unknown types
-    # like torch.dtype, torch.Size, so we need this configuration
-    # see https://docs.pydantic.dev/latest/api/config/#pydantic.config.ConfigDict.validate_assignment
-    model_config = ConfigDict(arbitrary_types_allowed=True)
+def _dt_validate(value: Any) -> torch.dtype:
+    if isinstance(value, str):
+        if not value.startswith("torch."):
+            raise ValueError(f"dtype {value} should start with torch.")
+        try:
+            value = getattr(torch, value.split(".")[1])
+        except AttributeError as e:
+            raise ValueError(f"unknown dtype: {value}") from e
+    if not isinstance(value, torch.dtype):
+        raise TypeError(f"dtype {value} should be torch.dtype, got {type(value)}")
+    return value
+
+
+_TorchDtype = Annotated[
+    torch.dtype,
+    PlainValidator(_dt_validate),
+    PlainSerializer(lambda x: str(x), return_type=str),
+    WithJsonSchema({"type": "string"}, mode="serialization"),
+]
+
+
+def _size_validate(value: Any) -> torch.Size:
+    if isinstance(value, list | tuple):
+        return torch.Size(value)
+    if not isinstance(value, torch.Size):
+        raise TypeError(f"size {value} should be torch.Size, got {type(value)}")
+    return value
+
+
+_TorchSize = Annotated[
+    torch.Size,
+    PlainValidator(_size_validate),
+    PlainSerializer(lambda x: tuple(x), return_type=tuple),
+    WithJsonSchema({"type": "array", "items": {"type": "integer"}}, mode="serialization"),
+]
+
 
+def _tensor_validate(value: Any) -> torch.Tensor:
+    if isinstance(value, torch.Tensor):
+        return value
+    raise TypeError(f"tensor {value} should be torch.Tensor, got {type(value)}")
+
+
+_TorchTensor = Annotated[
+    torch.Tensor,
+    PlainValidator(_tensor_validate),
+]
+
+
+class ParameterMeta(BaseModel):
     name: str
-    dtype: torch.dtype
-    shape: torch.Size
+    dtype: _TorchDtype
+    shape: _TorchSize
 
 
 class BucketRange(NamedTuple):
@@ -68,9 +109,7 @@ class MemoryBufferMetas(BaseModel):
 
 
 class MemoryBuffer(BaseModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
-    buffer: torch.Tensor
+    buffer: _TorchTensor
     size: int
     metas: list[ParameterMeta]
 
@@ -82,7 +121,7 @@ class MemoryBufferMetaList(BaseModel):
 
 class DataToGather(MemoryBufferMetaList):
     host_ip: str
-    zmq_socket_path: tuple[str, str]
+    device_uuid: str
 
 
 # 256 bytes alignment when flatten torch tensors to uint8 buffer
@@ -93,7 +132,7 @@ def _align_size(dtype: torch.dtype, shape: torch.Size) -> int:
     return (dtype.itemsize * shape.numel() + _ALIGN_SIZE - 1) // _ALIGN_SIZE * _ALIGN_SIZE
 
 
-def _to_named_tensor(metas: list[ParameterMeta], offset=0) -> list[dict]:
+def _to_named_tensor(metas: list[ParameterMeta], offset: int = 0) -> list[dict]:
     ret = []
     for meta in metas:
         size = _align_size(meta.dtype, meta.shape)
@@ -109,11 +148,11 @@ def _to_named_tensor(metas: list[ParameterMeta], offset=0) -> list[dict]:
     return ret
 
 
-def _load_checkpoint_file(file_path: str) -> tuple[int, dict[str, tuple[FileMeta, torch.Tensor]]]:
-    def _safetensors_load(fn) -> dict[str, tuple[FileMeta, torch.Tensor]]:
+def _load_checkpoint_file(file_path: str) -> tuple[int, dict[str, tuple["FileMeta", torch.Tensor]]]:
+    def _safetensors_load(fn: str) -> dict[str, tuple["FileMeta", torch.Tensor]]:
         ret = {}
         with safe_open(fn, framework="pt") as f:
-            for name in f.keys():
+            for name in f.keys():  # noqa: SIM118
                 weight = f.get_tensor(name)
                 meta = {
                     "key": name,
@@ -126,10 +165,10 @@ def _load_checkpoint_file(file_path: str) -> tuple[int, dict[str, tuple[FileMeta
         return ret
 
     # deprecated, will be removed in the future
-    def _fast_np_load(fn) -> dict[str, tuple[FileMeta, torch.Tensor]]:
+    def _fast_np_load(fn: str) -> dict[str, tuple["FileMeta", torch.Tensor]]:
         """load *.np file and return memmap and related tensor meta"""
 
-        def parse_npy_header(fin):
+        def parse_npy_header(fin: BinaryIO) -> dict[str, Any]:
             start = fin.tell()
             major, minor = np.lib.format.read_magic(fin)
             if major == 1 and minor == 0:
@@ -137,7 +176,9 @@ def _load_checkpoint_file(file_path: str) -> tuple[int, dict[str, tuple[FileMeta
             elif major == 2 and minor == 0:
                 read_header_fn = np.lib.format.read_array_header_2_0
             else:
-                raise ValueError(f"unknown version {major}.{minor} when parsing npy header from {fn}")
+                raise ValueError(
+                    f"unknown version {major}.{minor} when parsing npy header from {fn}"
+                )
             shape, is_fortran, dtype = read_header_fn(fin)
             return {
                 "shape": shape,
@@ -193,7 +234,9 @@ def _load_checkpoint_file(file_path: str) -> tuple[int, dict[str, tuple[FileMeta
     return tp_rank, ret
 
 
-def _concat_tp_weights(tp_weights: list[torch.Tensor], tp_concat_dim: int, tp_size: int) -> torch.Tensor:
+def _concat_tp_weights(
+    tp_weights: list[torch.Tensor], tp_concat_dim: int, tp_size: int
+) -> torch.Tensor:
     """Concat tp weights with meta info.
     If meta.concat_dim is -1, meas this is shared tp weights, just use the first weights.
     Else we will cat weights in concat_dim.
@@ -206,39 +249,54 @@ def _concat_tp_weights(tp_weights: list[torch.Tensor], tp_concat_dim: int, tp_si
     return torch.cat([w for w in tp_weights], dim=tp_concat_dim)
 
 
-def _get_physical_gpu_id(rank: int) -> str:
-    result = subprocess.run(["nvidia-smi", "-L"], capture_output=True, text=True)
-    if result.returncode != 0:
-        raise ValueError(result.stdout)
-    lines = result.stdout.strip().split("\n")
-    for line in lines:
-        if f"GPU {rank}" in line:
-            uuid = line.split("UUID: ")[1].strip(")")
-            return uuid
-    raise ValueError(f"not found gpu{rank} uuid")
+def _get_physical_gpu_id(device_index: int | None = None) -> str:
+    try:
+        return f"GPU-{torch.cuda.get_device_properties(device_index).uuid!s}"
+    except AssertionError as e:
+        raise ValueError(f"fail to get physical gpu id {device_index}") from e
 
 
 @lru_cache(maxsize=1)
-def _get_ip():
+def _get_ip() -> str:
     try:
         # try to get ip from network interface
         with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
             s.connect(("8.8.8.8", 80))
             return s.getsockname()[0]
-    except:
+    except Exception as e:  # noqa: BLE001
         # fallback to get ip from hostname
-        logger.warning("fail to get ip from network interface, fallback to get ip from hostname")
+        logger.warning(
+            f"fail to get ip from network interface, fallback to get ip from hostname: {e}"
+        )
         return socket.gethostbyname(socket.gethostname())
 
 
+def _ibv_get_device_list() -> list[str]:
+    lib = ctypes.CDLL("libibverbs.so.1")
+    lib.ibv_get_device_list.argtypes = [ctypes.POINTER(ctypes.c_int)]  # int *num_devices
+    lib.ibv_get_device_list.restype = ctypes.POINTER(ctypes.c_void_p)  # struct ibv_device **
+
+    lib.ibv_free_device_list.argtypes = [ctypes.POINTER(ctypes.c_void_p)]
+    lib.ibv_get_device_name.argtypes = [ctypes.c_void_p]  # struct ibv_device *
+    lib.ibv_get_device_name.restype = ctypes.c_char_p  # const char *
+
+    num = ctypes.c_int()
+    dev_array = lib.ibv_get_device_list(ctypes.byref(num))
+    if not dev_array or num.value <= 0:
+        return []
+
+    devices = []
+    for i in range(num.value):
+        dev_ptr = dev_array[i]  # struct ibv_device *
+        name = lib.ibv_get_device_name(dev_ptr)  # const char *
+        devices.append(name.decode())
+    lib.ibv_free_device_list(dev_array)
+    return devices
+
+
 def _get_rdma_devices() -> list[str]:
     """
-    use script like below to get RDMA devices, if NCCL_IB_HCA has multiple values, just return
-    ```bash
-    pushd /sys/class/infiniband/ > /dev/null;
-    for i in mlx5_*; do cat "$i"/ports/1/gid_attrs/types/* 2>/dev/null | grep v >/dev/null && echo "$i" ; done;
-    popd > /dev/null;
-    ```
+    use _ibv_get_device_list to get RDMA devices, if NCCL_IB_HCA has multiple values, just return
     """
     devices_str = os.getenv("PS_P2P_STORE_RDMA_DEVICES")
     if devices_str:
@@ -246,41 +304,27 @@ def _get_rdma_devices() -> list[str]:
     # if PS_P2P_STORE_RDMA_DEVICES is not set, try to use NCCL_IB_HCA to get RDMA devices
     hca = os.getenv("NCCL_IB_HCA", None)
     if hca:
-        l = hca.split(",")
-        if len(l) > 1:
+        hca_list = hca.split(",")
+        if len(hca_list) > 1:
             # if NCCL_IB_HCA has multiple values, just return
-            return l
+            return hca_list
         else:
-            hca = l[0]
-    basepath = "/sys/class/infiniband/"
-    port_path = "ports/1/gid_attrs/types"
-    devices = []
-    for device in sorted(os.listdir(basepath)):
-        if hca is not None and hca not in device:
-            continue
-        path = os.path.join(basepath, device, port_path)
-        if not os.path.exists(path) or not os.path.isdir(path):
-            continue
-        for port in os.listdir(path):
-            try:
-                content = open(os.path.join(path, port)).read()
-                if "v" in content:
-                    print(f"found rdma device {device} in port {port}: {content.strip()}")
-                    devices.append(device)
-                    break
-            except Exception:
-                pass
-    return devices
+            hca = hca_list[0]
+    return [device for device in sorted(_ibv_get_device_list()) if hca is None or hca in device]
 
 
-def _get_my_rdma_device(local_rank: int, gpu_count: int, devices: list[str]):
+def _get_my_rdma_device(local_rank: int, gpu_count: int, devices: list[str]) -> str:
     """
     implement network card device allocation, if network card is "mlx5_0,mlx5_1", then 0-3 will share mlx5_0, 4-7 will share mlx5_1, etc.
     """
     if not devices:
         raise RuntimeError("no rdma devices found")
-    assert len(devices) <= gpu_count, f"rdma devices count {len(devices)} should be less than or equal to gpu count {gpu_count}"
-    assert gpu_count % len(devices) == 0, f"gpu count {gpu_count} should be divisible by rdma devices count {len(devices)}"
+    assert len(devices) <= gpu_count, (
+        f"rdma devices count {len(devices)} should be less than or equal to gpu count {gpu_count}"
+    )
+    assert gpu_count % len(devices) == 0, (
+        f"gpu count {gpu_count} should be divisible by rdma devices count {len(devices)}"
+    )
     return devices[local_rank // (gpu_count // len(devices))]
 
 
@@ -305,8 +349,12 @@ def _load_checkpoint(files: list[str]) -> dict[str, torch.Tensor]:
                     size=1,
                 )
             if parameter_name not in parameter_metas:
-                assert isinstance(meta["dtype"], torch.dtype), f"meta {meta} dtype should be torch.dtype"
-                assert isinstance(meta["shape"], torch.Size), f"meta {meta} shape should be torch.Size"
+                assert isinstance(meta["dtype"], torch.dtype), (
+                    f"meta {meta} dtype should be torch.dtype"
+                )
+                assert isinstance(meta["shape"], torch.Size), (
+                    f"meta {meta} shape should be torch.Size"
+                )
                 parameter_metas[parameter_name] = ParameterMeta(
                     name=parameter_name,
                     shape=meta["shape"],
@@ -319,7 +367,9 @@ def _load_checkpoint(files: list[str]) -> dict[str, torch.Tensor]:
         if tp_meta.concat_dim != -1:
             shape = list(parameter_metas[name].shape)
             shape[tp_meta.concat_dim] = shape[tp_meta.concat_dim] * tp_meta.size
-            parameter_metas[name] = ParameterMeta(name=name, shape=torch.Size(shape), dtype=parameter_metas[name].dtype)
+            parameter_metas[name] = ParameterMeta(
+                name=name, shape=torch.Size(shape), dtype=parameter_metas[name].dtype
+            )
         weights_in_cpu = [parameters_with_tp[name][key] for key in sorted(parameters_with_tp[name])]
         # TODO: here concat is serial, which may be slow
         # but since tp storage is not used in the future
@@ -338,17 +388,19 @@ def _load_checkpoint(files: list[str]) -> dict[str, torch.Tensor]:
 
 def _register_checkpoint(
     *,
-    files: list[str] = [],
-    named_tensors: dict[str, torch.Tensor] = {},
+    files: list[str],
+    named_tensors: dict[str, torch.Tensor],
     rank: int | None = None,
 ) -> list[MemoryBuffer]:
-    logger.info(f"[rank{rank}] start to register checkpoint with {len(files)} files and {len(named_tensors)} named_tensors")
+    logger.info(
+        f"[rank{rank}] start to register checkpoint with {len(files)} files and {len(named_tensors)} named_tensors"
+    )
     if not files and not named_tensors:
         return []
     parameters = _load_checkpoint(files)
     if named_tensors:
         parameters.update(named_tensors)
-    bucket_size = max(4 << 30, max(map(lambda x: _align_size(x.dtype, x.shape), parameters.values())))
+    bucket_size = max(4 << 30, max(_align_size(x.dtype, x.shape) for x in parameters.values()))
 
     class MemoryBucket(BaseModel):
         size: int
@@ -363,7 +415,10 @@ def _register_checkpoint(
         buckets[-1].metas.append(ParameterMeta(name=name, shape=tensor.shape, dtype=tensor.dtype))
         buckets[-1].size += size
 
-    memory_buffers = [MemoryBuffer(buffer=torch.empty(0), size=bucket.size, metas=bucket.metas) for bucket in buckets]
+    memory_buffers = [
+        MemoryBuffer(buffer=torch.empty(0), size=bucket.size, metas=bucket.metas)
+        for bucket in buckets
+    ]
 
     def register_pin_memory(idx: int, size: int) -> tuple[int, torch.Tensor]:
         buffer = torch.empty(size, dtype=torch.uint8, pin_memory=True)
@@ -373,7 +428,10 @@ def _register_checkpoint(
         buffer[offset : offset + tensor.nbytes] = tensor.view(-1).view(dtype=torch.uint8)
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
-        futures = [executor.submit(register_pin_memory, idx, bucket.size) for idx, bucket in enumerate(buckets)]
+        futures = [
+            executor.submit(register_pin_memory, idx, bucket.size)
+            for idx, bucket in enumerate(buckets)
+        ]
         new_futures = []
         for future in concurrent.futures.as_completed(futures):
             idx, buffer = future.result()
@@ -400,8 +458,26 @@ def _register_checkpoint(
     return memory_buffers
 
 
-def request_inference_to_update(url: str, socket_paths: dict[str, str], timeout: float = 300.0):
-    resp = requests.post(
+def request_inference_to_update(
+    url: str,
+    socket_paths: dict[str, str],
+    timeout: float = 300.0,
+    uds: str | None = None,
+):
+    """Send an inference update request to inference server via HTTP or Unix socket.
+
+    Args:
+        url (str): The HTTP URL or request path (e.g., "http://localhost:19730/inference") to send the request to.
+        socket_paths (dict[str, str]): A dictionary containing device uuid and IPC socket paths for updating weights.
+        timeout (float, optional): Request timeout in seconds. Defaults to 300.0.
+        uds (str, optional): Path to a Unix domain socket. If provided, the request
+            will be sent via the Unix socket instead of HTTP. Defaults to None.
+
+    Raises:
+        httpx.HTTPStatusError: If the response contains an HTTP error status.
+        httpx.RequestError: If there was an issue while making the request.
+    """
+    resp = httpx.Client(transport=httpx.HTTPTransport(uds=uds)).post(
         url,
         json={
             "method": "update_weights_from_ipc",
@@ -413,7 +489,9 @@ def request_inference_to_update(url: str, socket_paths: dict[str, str], timeout:
     resp.raise_for_status()
 
 
-def _gen_h2d_buckets(global_metas: dict[int, MemoryBufferMetaList], bucket_size: int) -> list[tuple[int, H2DBucket]]:
+def _gen_h2d_buckets(
+    global_metas: dict[int, MemoryBufferMetaList], bucket_size: int
+) -> list[tuple[int, H2DBucket]]:
     buckets: list[tuple[int, H2DBucket]] = []
 
     for owner_rank, items in global_metas.items():
@@ -424,14 +502,18 @@ def _gen_h2d_buckets(global_metas: dict[int, MemoryBufferMetaList], bucket_size:
                 s = _align_size(meta.dtype, meta.shape)
                 if buckets[-1][1].size + s > bucket_size:
                     if offset - start_offset > 0:
-                        buckets[-1][1].ranges.append(BucketRange(idx, start_offset, offset - start_offset))
+                        buckets[-1][1].ranges.append(
+                            BucketRange(idx, start_offset, offset - start_offset)
+                        )
                     start_offset = offset
                     buckets.append((owner_rank, H2DBucket(size=0, ranges=[], items=[])))
                 offset += s
                 buckets[-1][1].size += s
                 buckets[-1][1].items.append(meta)
             buckets[-1][1].ranges.append(BucketRange(idx, start_offset, offset - start_offset))
-        assert buckets[-1][1].size > 0, f"buckets[-1][1].size {buckets[-1][1].size} should be greater than 0"
+        assert buckets[-1][1].size > 0, (
+            f"buckets[-1][1].size {buckets[-1][1].size} should be greater than 0"
+        )
     return buckets
 
 
@@ -470,7 +552,9 @@ class P2PStore:
             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 {device}")
+        logger.info(
+            f"[rank{self.rank}] p2p store initialized, addr is {self.addr}, rdma device is {device}"
+        )
 
     @property
     def addr(self) -> str:
@@ -492,37 +576,60 @@ class P2PStore:
         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])}")
+            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
+    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:
-    def __init__(self, *, auto_pg: bool = False):
+    def __init__(
+        self,
+        *,
+        rank: int | None = None,
+        world_size: int | None = None,
+        auto_pg: bool = False,
+        gpu_count: int | None = None,
+        mem_fraction: float | None = None,
+    ):
         """
         Initialize the parameter server. env RANK, WORLD_SIZE and MASTER_ADDR must be set.
 
         Args:
             auto_pg: Whether to automatically initialize the process group.
                 Notice that if auto_pg is True, will destroy the process group after update.
+            mem_fraction: The proportion (as a fraction) of the current free CUDA memory for allocation.
         """
-        self._rank = int(os.environ.get("RANK", None))
-        self._world_size = int(os.environ.get("WORLD_SIZE", None))
-        self._master_addr = os.getenv("MASTER_ADDR")
-        self._gpu_count = torch.cuda.device_count()
+        self._rank = rank or int(os.environ.get("RANK", None))
+        self._world_size = world_size or int(os.environ.get("WORLD_SIZE", None))
+        self._gpu_count = gpu_count or torch.cuda.device_count()
         self._local_rank = self._rank % self._gpu_count
         self._auto_pg = auto_pg
         self._all_hosts = []
-        self._global_socket_paths: list[tuple[str, str]] = []
+        self._global_device_uuids: list[str] = []
+        self._mem_fraction = mem_fraction or 0.9
 
         assert self._rank is not None and self._rank >= 0, self._rank
         assert self._world_size and self._world_size > 0, self._world_size
+        assert (
+            self._gpu_count is not None
+            and self._gpu_count > 0
+            and self._gpu_count <= torch.cuda.device_count()
+        ), self._gpu_count
+        assert (
+            self._mem_fraction is not None and self._mem_fraction > 0 and self._mem_fraction <= 1
+        ), self._mem_fraction
 
-        self._device_uuid = _get_physical_gpu_id(self._local_rank)
         self._zmq_ctx = zmq.Context()
+        self._zmq_addr_counter = 0
 
         self._memory_pool: dict[str, list[MemoryBuffer]] = {}
         # dict key is owner_rank, value is a bucket metas list in owner_rank
@@ -533,19 +640,27 @@ class ParameterServer:
             logger.warning(f"[rank{self._rank}] fail to initialize p2p store due to {e}")
             self._p2p_store = None
 
-        torch.cuda.set_device(self._local_rank)
+        device_index = self._local_rank
+        torch.cuda.set_device(device_index)
+        self._device_uuid = _get_physical_gpu_id(device_index)
 
-    def _logger_rank0(self, msg):
+    def _logger_rank0(self, msg: str):
         if self._local_rank == 0:
             logger.info(msg)
 
-    def get_metas(self):
+    def get_metas(self) -> dict[int, MemoryBufferMetaList]:
         return self._current_global_parameter_metas
 
     def load_metas(self, metas: dict[int, MemoryBufferMetaList]):
         self._current_global_parameter_metas = metas
 
-    def register_checkpoint(self, checkpoint_name: str, *, files: list[str] = [], named_tensors: dict[str, torch.Tensor] = {}):
+    def register_checkpoint(
+        self,
+        checkpoint_name: str,
+        *,
+        files: list[str] | None = None,
+        named_tensors: dict[str, torch.Tensor] | None = None,
+    ) -> None:
         """
         Register a checkpoint to the parameter server. Both files and named_tensors will be registered together.
 
@@ -555,12 +670,18 @@ class ParameterServer:
             named_tensors: The named tensors to register.
         """
         try:
-            assert checkpoint_name not in self._memory_pool, f"checkpoint {checkpoint_name} already registered"
-            self._memory_pool[checkpoint_name] = _register_checkpoint(files=files, named_tensors=named_tensors, rank=self._rank)
+            assert checkpoint_name not in self._memory_pool, (
+                f"checkpoint {checkpoint_name} already registered"
+            )
+            self._memory_pool[checkpoint_name] = _register_checkpoint(
+                files=files or [], named_tensors=named_tensors or {}, rank=self._rank
+            )
             if self._p2p_store is not None:
                 self._register_parameters_to_p2p_store(checkpoint_name)
         except Exception:
-            logger.exception(f"[rank{self._rank}] fail to register checkpoint {checkpoint_name} with files {files}")
+            logger.exception(
+                f"[rank{self._rank}] fail to register checkpoint {checkpoint_name} with files {files}"
+            )
             if self._p2p_store is not None:
                 self._unregister_parameters_from_p2p_store(checkpoint_name)
             self.unregister_checkpoint(checkpoint_name)
@@ -583,10 +704,6 @@ class ParameterServer:
         # this works by using torch>=2.5.0
         torch._C._host_emptyCache()
 
-    @cached_property
-    def _zmq_socket_path(self) -> str:
-        return f"ipc://@checkpoint-engine-{uuid.uuid4()}.sock"
-
     def gather_metas(self, checkpoint_name: str):
         """
         Gather the parameter metas from all ranks. This will gather memory_buffer, and other metadatas.
@@ -598,19 +715,17 @@ class ParameterServer:
         assert dist.is_initialized(), "process group is not initialized"
         metas_lst: list[DataToGather | None] = [None for _ in range(self._world_size)]  # type: ignore
         metas = DataToGather(
-            memory_buffer_metas_list=list(
-                map(
-                    lambda x: MemoryBufferMetas(
-                        metas=x.metas,
-                        ptr=x.buffer.data_ptr(),
-                        size=x.size,
-                    ),
-                    self._memory_pool.get(checkpoint_name, []),
-                ),
-            ),
+            memory_buffer_metas_list=[
+                MemoryBufferMetas(
+                    metas=x.metas,
+                    ptr=x.buffer.data_ptr(),
+                    size=x.size,
+                )
+                for x in self._memory_pool.get(checkpoint_name, [])
+            ],
             p2p_store_addr=None if self._p2p_store is None else self._p2p_store.addr,
             host_ip=_get_ip(),
-            zmq_socket_path=(self._device_uuid, self._zmq_socket_path),
+            device_uuid=self._device_uuid,
         )
 
         dist.all_gather_object(metas_lst, metas)
@@ -618,23 +733,31 @@ class ParameterServer:
         self._current_global_parameter_metas = {}
         num_parameters = 0
         all_hosts: list[str] = []
-        global_socket_paths: list[tuple[str, str]] = []
+        global_device_uuids: list[str] = []
         for i, metas_buckets in enumerate(metas_lst):
             assert metas_buckets is not None, f"metas_buckets {i} should not be None"
             if i % self._gpu_count == 0 and not self._all_hosts:
                 all_hosts.append(metas_buckets.host_ip)
-            if not self._global_socket_paths:
-                global_socket_paths.append(metas_buckets.zmq_socket_path)
+            if not self._global_device_uuids:
+                global_device_uuids.append(metas_buckets.device_uuid)
             if metas_buckets.memory_buffer_metas_list:
                 self._current_global_parameter_metas[i] = metas_buckets
-                num_parameters += sum(map(lambda x: len(x.metas), metas_buckets.memory_buffer_metas_list))
+                num_parameters += sum(len(x.metas) for x in metas_buckets.memory_buffer_metas_list)
         if not self._all_hosts:
             self._all_hosts = all_hosts
-        if not self._global_socket_paths:
-            self._global_socket_paths = global_socket_paths
-        logger.info(f"[rank{self._rank}] gather parameter metas finished, num_parameters: {num_parameters}")
+        if not self._global_device_uuids:
+            self._global_device_uuids = global_device_uuids
+        logger.info(
+            f"[rank{self._rank}] gather parameter metas finished, num_parameters: {num_parameters}"
+        )
 
-    def init_process_group(self, *, master_port: int | None = None, timeout: timedelta = timedelta(minutes=10)):
+    def init_process_group(
+        self,
+        *,
+        master_addr: str | None = None,
+        master_port: int | None = None,
+        timeout: timedelta = timedelta(minutes=10),
+    ):
         """
         Initialize the process group for the ranks. This global group can be easily destroyed by calling dist.destroy_process_group.
 
@@ -642,10 +765,22 @@ class ParameterServer:
             master_port: The specified port of the master node. If not set, will use _get_master_port to get the port.
             timeout: The timeout of the process group.
         """
+        master_addr = master_addr or os.getenv("MASTER_ADDR")
+        assert master_addr, "master_addr is required"
         store = dist.TCPStore(
-            self._master_addr, _get_master_port(master_port), self._world_size, timeout=timeout, is_master=self._rank == 0
+            master_addr,
+            _get_master_port(master_port),
+            self._world_size,
+            timeout=timeout,
+            is_master=self._rank == 0,
+        )
+        dist.init_process_group(
+            backend="nccl",
+            world_size=self._world_size,
+            rank=self._rank,
+            timeout=timeout,
+            store=store,
         )
-        dist.init_process_group(backend="nccl", world_size=self._world_size, rank=self._rank, timeout=timeout, store=store)
         logger.info(f"[rank{self._rank}] init process group successfully.")
 
     def update(
@@ -653,8 +788,8 @@ class ParameterServer:
         checkpoint_name: str,
         req_func: Callable[[list[tuple[str, str]]], None],
         *,
-        ranks: list[int] = [],
-    ):
+        ranks: list[int] | None = None,
+    ) -> None:
         """
         Update the checkpoint to inference engine. This function should be called after gather_metas.
 
@@ -667,18 +802,21 @@ class ParameterServer:
                 which is useful in disaggregated architecture.
         """
         try:
+            # if both ranks is None or [], it will use fully broadcast to update to all ranks
             if not ranks:
                 if self._auto_pg and not dist.is_initialized():
                     self.init_process_group()
                 self._update_per_bucket(checkpoint_name, req_func)
             else:
-                if self._rank not in ranks:
+                if not self._auto_pg and self._rank not in ranks:
                     return
                 if self._auto_pg:
                     if dist.is_initialized():
                         dist.destroy_process_group()
                         # HACK: wait 2s to ensure destroy is finished
                         time.sleep(2)
+                    if self._rank not in ranks:
+                        return
                     self.init_process_group_for_ranks(ranks)
                 self._update_per_bucket_p2p(checkpoint_name, req_func, ranks)
             if self._auto_pg:
@@ -692,21 +830,39 @@ class ParameterServer:
                 f"reserved {torch.cuda.memory_reserved() / 1024 / 1024} MB."
             )
         except Exception as e:
-            logger.exception(f"[rank{self._rank}] update checkpoint {checkpoint_name} with ranks {ranks} error {e}")
-            raise e
+            logger.exception(
+                f"[rank{self._rank}] update checkpoint {checkpoint_name} with ranks {ranks} error {e}"
+            )
+            raise
+
+    def _bind_zmq_socket(self) -> tuple[zmq.Socket, list[tuple[str, str]]]:
+        def zmq_handle(device_uuid: str) -> str:
+            return f"ipc://@checkpoint-engine-{device_uuid}-{self._zmq_addr_counter}.sock"
 
-    def _get_bucket_size(self, *, disable_h2d_buffer: bool = False) -> tuple[int, bool]:
-        GiB_bytes = 1 << 30
+        socket_paths = [(uid, zmq_handle(uid)) for uid in self._global_device_uuids]
+        socket = self._zmq_ctx.socket(zmq.REQ)
+        socket.bind(zmq_handle(self._device_uuid))
+        self._zmq_addr_counter += 1
+        return socket, socket_paths
+
+    def _detect_bucket_size(self, *, disable_h2d_buffer: bool = False) -> tuple[int, bool]:
+        GiB = 1 << 30  # noqa: N806
         # auto detect bucket size
-        free_bytes_tensor = torch.tensor(
-            int(float(torch.cuda.mem_get_info()[0]) * 0.9),
+        tensor = torch.tensor(
+            [
+                # proportion of current cuda free memory bytes
+                int(float(torch.cuda.mem_get_info()[0]) * self._mem_fraction),
+                # we use negative value to reuse allreduce min operation
+                # for getting the max value of zmq_addr_counter in all ranks
+                -self._zmq_addr_counter,
+            ],
             dtype=torch.int64,
             device="cuda",
         )
-        dist.all_reduce(free_bytes_tensor, op=dist.ReduceOp.MIN)
-        free_bytes = free_bytes_tensor.item()
+        dist.all_reduce(tensor, op=dist.ReduceOp.MIN)
+        tensor = tensor.cpu()
+        free_bytes, self._zmq_addr_counter = tensor[0].item(), -tensor[1].item()
         max_tensor_bytes = 0
-        max_bytes = int(os.getenv("PS_MAX_BUCKET_SIZE_GB", 8)) * GiB_bytes
         for items in self._current_global_parameter_metas.values():
             for metas_list in items.memory_buffer_metas_list:
                 for meta in metas_list.metas:
@@ -729,18 +885,27 @@ class ParameterServer:
                 f"max_tensor_bytes {max_tensor_bytes} should be less than free_bytes {free_bytes}"
             )
             disable_h2d_buffer = True
+        max_bytes = int(os.getenv("PS_MAX_BUCKET_SIZE_GB", 8)) * GiB
         bucket_size = min(max(max_bytes, max_tensor_bytes), free_bytes)
-        logger.info(f"[rank{self._rank}] auto detect bucket size {bucket_size / GiB_bytes:.2f} GiB")
+        logger.info(f"[rank{self._rank}] auto detect bucket size {bucket_size / GiB:.2f} GiB")
         return bucket_size, disable_h2d_buffer
 
-    def _copy_to_buffer(self, checkpoint_name: str, bucket: H2DBucket, buffer: torch.Tensor, owner_rank: int | None = None):
+    def _copy_to_buffer(
+        self,
+        checkpoint_name: str,
+        bucket: H2DBucket,
+        buffer: torch.Tensor,
+        owner_rank: int | None = None,
+    ):
         offset = 0
         if owner_rank is not None:
             buf_ptrs, remote_ptrs, lens = [], [], []
             ptr_base = buffer.data_ptr()
             target_addr, ptrs = self._get_addr_ptrs(owner_rank)
         for b in bucket.ranges:
-            assert offset + b.size <= bucket.size, f"offset {offset} + size {b.size} > bucket_size {bucket.size}"
+            assert offset + b.size <= bucket.size, (
+                f"offset {offset} + size {b.size} > bucket_size {bucket.size}"
+            )
             if owner_rank is not None:
                 buf_ptrs.append(ptr_base + offset)
                 remote_ptrs.append(ptrs[b.idx][0] + b.offset)
@@ -758,7 +923,11 @@ class ParameterServer:
         torch.cuda.synchronize()
 
     def init_process_group_for_ranks(
-        self, ranks: list[int], *, master_port: int | None = None, timeout: timedelta = timedelta(minutes=10)
+        self,
+        ranks: list[int],
+        *,
+        master_port: int | None = None,
+        timeout: timedelta = timedelta(minutes=10),
     ):
         """
         Initialize the process group for the ranks. This global group can be easily destroyed by calling dist.destroy_process_group.
@@ -787,8 +956,12 @@ class ParameterServer:
         # and will not participate in this update. Since they have registered memory addresses
         # to p2p_store at the beginning, update ranks can directly get the memory addresses
         # from other nodes and put the weights into the buffer.
-        store = dist.TCPStore(master_addr, master_port, len(ranks), is_master=rank == 0, timeout=timeout)
-        dist.init_process_group(backend="nccl", world_size=len(ranks), rank=rank, timeout=timeout, store=store)
+        store = dist.TCPStore(
+            master_addr, master_port, len(ranks), is_master=rank == 0, timeout=timeout
+        )
+        dist.init_process_group(
+            backend="nccl", world_size=len(ranks), rank=rank, timeout=timeout, store=store
+        )
 
     def _update_per_bucket_p2p(
         self,
@@ -800,7 +973,9 @@ class ParameterServer:
         assert ranks, "ranks should be set"
         if len(self._current_global_parameter_metas) == 0:
             raise ValueError("parameter metas is empty")
-        assert dist.is_initialized(), "process group is not initialized when update model per bucket p2p"
+        assert dist.is_initialized(), (
+            "process group is not initialized when update model per bucket p2p"
+        )
 
         need_update = self._rank in ranks
         logger.info(
@@ -814,26 +989,24 @@ class ParameterServer:
         # first execute a barrier to avoid subsequent cuda oom
         dist.barrier()
 
-        bucket_size, _ = self._get_bucket_size(disable_h2d_buffer=True)
+        bucket_size, _ = self._detect_bucket_size(disable_h2d_buffer=True)
         buffer = torch.empty(bucket_size * 2, dtype=torch.uint8, device="cuda")
-        IPC_BUFFER_NAME = "__ipc_buffer___"
-        self._p2p_store.register_named_tensors({IPC_BUFFER_NAME: buffer})
+        ipc_buffer_name = "__ipc_buffer___"
+        self._p2p_store.register_named_tensors({ipc_buffer_name: buffer})
         logger.info(
             f"[rank{self._rank}] register buffer, shape={buffer.shape}, dtype={buffer.dtype}, data_ptr={buffer.data_ptr()}, nbytes={buffer.nbytes}"
         )
         handle = reduce_tensor(buffer)
 
-        gidx = 0
         buckets = _gen_h2d_buckets(self._current_global_parameter_metas, bucket_size)
+        socket, socket_paths = self._bind_zmq_socket()
         req_thread = threading.Thread(
             target=req_func,
-            args=(self._global_socket_paths,),
+            args=(socket_paths,),
         )
         req_thread.start()
-        socket = self._zmq_ctx.socket(zmq.REQ)
-        socket.bind(self._zmq_socket_path)
         socket.send_pyobj(handle)
-        for owner_rank, bucket in buckets:
+        for gidx, (owner_rank, bucket) in enumerate(buckets):
             self._logger_rank0(
                 f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} owner_rank {owner_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. "
             )
@@ -845,7 +1018,6 @@ class ParameterServer:
             socket.recv()
             dist.barrier()
             socket.send_pyobj(_to_named_tensor(bucket.items, gidx % 2 * bucket_size))
-            gidx += 1
 
         socket.recv()
         socket.send_pyobj(None)
@@ -853,7 +1025,7 @@ class ParameterServer:
         req_thread.join()
         dist.barrier()
         socket.close()
-        self._p2p_store.unregister_named_tensors([IPC_BUFFER_NAME])
+        self._p2p_store.unregister_named_tensors([ipc_buffer_name])
         torch.cuda.empty_cache()
 
     def _get_addr_ptrs(self, owner_rank: int) -> tuple[str, list[tuple[int, int]]]:
@@ -877,7 +1049,9 @@ class ParameterServer:
         pool = self._memory_pool[checkpoint_name]
         if len(pool) == 0:
             return 0
-        return self._p2p_store.unregister_named_tensors([f"memory_pool_{checkpoint_name}_{idx}" for idx, _ in enumerate(pool)])
+        return self._p2p_store.unregister_named_tensors(
+            [f"memory_pool_{checkpoint_name}_{idx}" for idx, _ in enumerate(pool)]
+        )
 
     def _update_per_bucket(
         self,
@@ -891,11 +1065,13 @@ class ParameterServer:
 
         logger.info(f"[rank{self._rank}] update checkpoint {checkpoint_name}")
 
-        bucket_size, disable_h2d_buffer = self._get_bucket_size()
+        bucket_size, disable_h2d_buffer = self._detect_bucket_size()
         buckets = _gen_h2d_buckets(self._current_global_parameter_metas, bucket_size)
 
         h2d_buffer: torch.Tensor | None = (
-            None if disable_h2d_buffer else torch.empty(bucket_size, dtype=torch.uint8, device="cuda")
+            None
+            if disable_h2d_buffer
+            else torch.empty(bucket_size, dtype=torch.uint8, device="cuda")
         )
 
         owner_rank_buckets: list[H2DBucket] = []
@@ -914,13 +1090,12 @@ class ParameterServer:
             if len(buckets_by_owner_rank[owner_rank]) > max_len:
                 max_len = len(buckets_by_owner_rank[owner_rank])
 
+        socket, socket_paths = self._bind_zmq_socket()
         req_thread = threading.Thread(
             target=req_func,
-            args=(self._global_socket_paths,),
+            args=(socket_paths,),
         )
         req_thread.start()
-        socket = self._zmq_ctx.socket(zmq.REQ)
-        socket.bind(self._zmq_socket_path)
         socket.send_pyobj(handle)
 
         gidx = 0
@@ -932,7 +1107,10 @@ class ParameterServer:
                 if i >= len(_buckets):
                     continue
                 bucket = _buckets[i]
-                alloc, reserved = torch.cuda.memory_allocated() / 1024 / 1024, torch.cuda.memory_reserved() / 1024 / 1024
+                alloc, reserved = (
+                    torch.cuda.memory_allocated() / 1024 / 1024,
+                    torch.cuda.memory_reserved() / 1024 / 1024,
+                )
                 self._logger_rank0(
                     f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} owner_rank {owner_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. "
                     f"Current CUDA allocated {alloc:.2f} MB, "
@@ -960,7 +1138,7 @@ class ParameterServer:
         torch.cuda.empty_cache()
 
 
-def _init_api(ps: ParameterServer):
+def _init_api(ps: ParameterServer) -> Any:
     import fastapi
     from fastapi import Request
     from fastapi.responses import JSONResponse, Response
@@ -976,32 +1154,32 @@ def _init_api(ps: ParameterServer):
         inference_group_ranks: list[int] = []
         timeout: float = 300.0
 
-    def wrap_exception(func):
+    def wrap_exception(func: Callable[[], None]) -> Response:
         try:
             func()
-        except Exception as e:
+        except Exception as e:  # noqa: BLE001
             logger.exception(f"wrap exception {func} failed")
             return JSONResponse(content=str(e), status_code=500)
         return Response(status_code=200)
 
     @app.post("/v1/checkpoints/{checkpoint_name}/files")
-    async def register_files(checkpoint_name: str, req: RegisterRequest, raw: Request):
+    async def register_files(checkpoint_name: str, req: RegisterRequest, raw: Request) -> Response:
         return wrap_exception(lambda: ps.register_checkpoint(checkpoint_name, files=req.files))
 
     @app.delete("/v1/checkpoints/{checkpoint_name}")
-    async def unregister_checkpoint(checkpoint_name: str):
+    async def unregister_checkpoint(checkpoint_name: str) -> Response:
         return wrap_exception(lambda: ps.unregister_checkpoint(checkpoint_name))
 
     @app.get("/v1/healthz")
-    async def healthz():
+    async def healthz() -> Response:
         return Response(status_code=200)
 
     @app.post("/v1/checkpoints/{checkpoint_name}/gather-metas")
-    async def gather_metas(checkpoint_name: str):
+    async def gather_metas(checkpoint_name: str) -> Response:
         return wrap_exception(lambda: ps.gather_metas(checkpoint_name))
 
     @app.post("/v1/checkpoints/{checkpoint_name}/update")
-    async def update(checkpoint_name: str, req: UpdateRequest):
+    async def update(checkpoint_name: str, req: UpdateRequest) -> Response:
         def update_func(socket_paths: list[tuple[str, str]]):
             if req.update_url is None:
                 return
@@ -1018,11 +1196,13 @@ def _init_api(ps: ParameterServer):
 def run_from_cli():
     import uvicorn
 
-    parser = argparse.ArgumentParser(description="Paramter Server")
+    parser = argparse.ArgumentParser(description="Parameter Server")
     parser.add_argument("--uds", type=str)
 
     args = parser.parse_args()
-    logger.info(f"Parameter Server {args=}, master addr: {os.getenv('MASTER_ADDR')}, master port {os.getenv('MASTER_PORT')}")
+    logger.info(
+        f"Parameter Server {args=}, master addr: {os.getenv('MASTER_ADDR')}, master port {os.getenv('MASTER_PORT')}"
+    )
 
     assert args.uds and len(args.uds) > 0, args.uds
     ps = ParameterServer(auto_pg=True)

checkpoint_engine/worker.py

@@ -1,11 +1,12 @@
 import gc
-from typing import Callable, Optional, TypedDict
+from collections.abc import Callable
+from typing import TypedDict
 
 import torch
 import zmq
 
 
-def _rebuild_ipc(handle: tuple[Callable, tuple], device_id: Optional[int] = None) -> torch.Tensor:
+def _rebuild_ipc(handle: tuple[Callable, tuple], device_id: int | None = None) -> torch.Tensor:
     func, args = handle
     list_args = list(args)
     if device_id is not None:
@@ -24,12 +25,14 @@ class FlattenedTensorMetadata(TypedDict):
     offset: int
 
 
-def _extract_weights(payload: list[FlattenedTensorMetadata], buffer: torch.Tensor) -> list[tuple[str, torch.Tensor]]:
+def _extract_weights(
+    payload: list[FlattenedTensorMetadata], buffer: torch.Tensor
+) -> list[tuple[str, torch.Tensor]]:
     assert buffer is not None
     weights: list[tuple[str, torch.Tensor]] = []
     for item in payload:
         shape = item["shape"]
-        if isinstance(shape, (list, tuple)):
+        if isinstance(shape, list | tuple):
             shape = torch.Size(shape)
         assert isinstance(shape, torch.Size)
         dtype, offset = item["dtype"], item["offset"]
@@ -45,11 +48,11 @@ def update_weights_from_ipc(
     device_id: int,
     *,
     run: Callable[[list[tuple[str, torch.Tensor]]], None],
-    post_hook: Callable[[], None] = None,
+    post_hook: Callable[[], None] | None = None,
 ):
     socket = zmq_ctx.socket(zmq.REP)
     socket.connect(zmq_handle)
-    buffer: Optional[torch.Tensor] = None
+    buffer: torch.Tensor | None = None
     while True:
         payload: tuple[Callable, tuple] | list[FlattenedTensorMetadata] | None = socket.recv_pyobj()
         if payload is None:
@@ -100,5 +103,7 @@ class VllmColocateWorkerExtension:
             zmq_handles[device_uuid],
             device_id=self.device.index,
             run=self.model_runner.model.load_weights,
-            post_hook=lambda: process_weights_after_loading(self.model_runner.model, self.model_config, self.device),
+            post_hook=lambda: process_weights_after_loading(
+                self.model_runner.model, self.model_config, self.device
+            ),
         )

{checkpoint_engine-0.1.1.dist-info → checkpoint_engine-0.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpoint-engine
-Version: 0.1.1
+Version: 0.1.3
 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
@@ -11,14 +11,15 @@ Description-Content-Type: text/markdown
 License-File: LICENCE
 Requires-Dist: torch>=2.5.0
 Requires-Dist: fastapi
-Requires-Dist: pydantic
+Requires-Dist: pydantic>=2.0.0
 Requires-Dist: safetensors
 Requires-Dist: pyzmq
 Requires-Dist: uvicorn
 Requires-Dist: loguru
 Requires-Dist: numpy
+Requires-Dist: httpx
 Provides-Extra: p2p
-Requires-Dist: mooncake-transfer-engine; extra == "p2p"
+Requires-Dist: mooncake-transfer-engine>=0.3.5; extra == "p2p"
 Dynamic: license-file
 
 # Checkpoint Engine
@@ -41,7 +42,7 @@ The core weight update logic is in `ParameterServer` class, a service colocated
 - **P2P**: Used when new inference instances are dynamically added (due to restarts or dynamic availability) while the existing instances are already serving requests. Under this scenario, to avoid affecting the workloads on existing instances, we use the [`mooncake-transfer-engine`](https://github.com/kvcache-ai/Mooncake?tab=readme-ov-file#use-python-package) to P2P send weights from CPUs in existing instances to GPUs in new instances. See `_update_per_bucket_p2p`.
 
 ### Optimized Weight Broadcast
-In the *Broadcast* implementation, the checkpoint-engine holds references to sharded weights in CPU memory, and need to efficiently broadcast them to a cluster of inference instances, often under a different sharding pattern. 
+In the *Broadcast* implementation, the checkpoint-engine holds references to sharded weights in CPU memory, and need to efficiently broadcast them to a cluster of inference instances, often under a different sharding pattern.
 We arrange the data transfer into 3 stages:
 1. H2D: moving weights to GPU memory. These weights may come from disk or the training engine.
 2. broadcast: broadcast among checkpoint engine workers; the data results in a CUDA IPC buffer shared with inference engine.
@@ -73,9 +74,9 @@ Pipelining naturally requires more GPU memory. When memory is not enough, checkp
 All results above are tested by [`examples/update.py`](./examples/update.py) and use [vLLM v0.10.2rc1](https://github.com/vllm-project/vllm/tree/v0.10.2rc1) as inference engine. Some notes:
 
 * FP8 test needs additional vLLM patches, see [FP8 quantization](#fp8-quantization).
-* Device Info: we tested various combination of devices and paralleism setups. For exmaple, a 256-GPU TP16 setup means that we deploy 16 vLLM instances, each with 16-way tensor parallelism.
+* Device Info: we tested various combination of devices and parallelism setups. For example, a 256-GPU TP16 setup means that we deploy 16 vLLM instances, each with 16-way tensor parallelism.
 * Since update duration is related to IPC bucket size, we provide the bucket size in the table.
-* The P2P time were tested for updating no more than two nodes (16 GPUs) (`ParameterServer.update(ranks=range(0, 16))`) out of the entire cluster. 
+* The P2P time were tested for updating no more than two nodes (16 GPUs) (`ParameterServer.update(ranks=range(0, 16))`) out of the entire cluster.
 
 ## Installation
 
@@ -88,7 +89,7 @@ pip install checkpoint-engine
 Use the flexible P2P implementation, notice this will install `mooncake-transfer-engine` to support RDMA transfer between different ranks.
 
 ```Bash
-pip install checkpoint-engine[p2p]
+pip install 'checkpoint-engine[p2p]'
 ```
 
 If set `NCCL_IB_HCA` env, checkpoint-engine will use it to auto select net devices for different ranks. If not set, it will read all RDMA devices and try to divide them into each rank.
@@ -107,7 +108,7 @@ VLLM_USE_PRECOMPILED=1 uv pip install --editable .
 Install checkpoint-engine
 
 ```Bash
-uv pip install checkpoint-engine[p2p]
+uv pip install 'checkpoint-engine[p2p]'
 ```
 
 We use `Qwen/Qwen3-235B-A22B-Instruct-2507` (BF16) as the test model
@@ -133,7 +134,7 @@ torchrun --nproc-per-node 8 examples/update.py --update-method all --checkpoint-
 
 ### Reuse weights from existing instances
 
-New checkpoint-engine instances can join existing instances and reuse their weights. This is simple to achieve. 
+New checkpoint-engine instances can join existing instances and reuse their weights. This is simple to achieve.
 
 First, start the existing instances with `--save-metas-file global_metas.pkl` to save global metas to a file and use `--sleep-time 300` to make sure they stay alive.
 
@@ -150,7 +151,7 @@ torchrun --nproc-per-node 8 examples/update.py --load-metas-file global_metas.pk
 
 ### FP8 quantization
 
-FP8 quantization currently do not natively work in vLLM when updating weights. 
+FP8 quantization currently do not natively work in vLLM when updating weights.
 We provide a simple patch in [`patches/vllm_fp8.patch`](./patches/vllm_fp8.patch) to handle the correct weight update.
 Notice this patch is only tested in DeepSeek-V3.1 and Kimi-K2. Other models may meet some compatible issues.
 

checkpoint_engine-0.1.3.dist-info/RECORD

@@ -0,0 +1,9 @@
+checkpoint_engine/__init__.py,sha256=Zj4I008kn9R6fYr0lVBzcQSnvckLpX2s1ljCOOqV1c8,87
+checkpoint_engine/_version.py,sha256=q5nF98G8SoVeJqaknL0xdyxtv0egsqb0fK06_84Izu8,704
+checkpoint_engine/ps.py,sha256=9dXRXi0QDPoRYrgGKAYvdmDFBXejgusjR0ltbii5_B0,49134
+checkpoint_engine/worker.py,sha256=ZmJTHeNPbnE8sPInfrghj9jeRDkMUSQO906o1UoJv-E,3748
+checkpoint_engine-0.1.3.dist-info/licenses/LICENCE,sha256=D3gPmHKpGtF1yxYNhqjtBtZY_brZjDotJTzpnmClzlY,1067
+checkpoint_engine-0.1.3.dist-info/METADATA,sha256=y96dMjEOKWaO_PA0h5BX8G3Ku7Tt1jCU09uIf8iYgic,9322
+checkpoint_engine-0.1.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+checkpoint_engine-0.1.3.dist-info/top_level.txt,sha256=66sik_1eLakLYmcllOEJzFaNbSfjsueuP0tHYEzhMSs,18
+checkpoint_engine-0.1.3.dist-info/RECORD,,

{checkpoint_engine-0.1.1.dist-info → checkpoint_engine-0.1.3.dist-info}/licenses/LICENCE

@@ -18,4 +18,4 @@ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE
+SOFTWARE

checkpoint_engine-0.1.1.dist-info/RECORD

@@ -1,9 +0,0 @@
-checkpoint_engine/__init__.py,sha256=Zj4I008kn9R6fYr0lVBzcQSnvckLpX2s1ljCOOqV1c8,87
-checkpoint_engine/_version.py,sha256=yfx_VE-4lpqM4jnWOSq-8rihMWIwMaX9CQ7tNEpA4T0,712
-checkpoint_engine/ps.py,sha256=9u2rLOj-oQrXsnpYhdXFjv7ak2-f4BRUXB6KYlG3ah0,44422
-checkpoint_engine/worker.py,sha256=OrSeknjtECnO88I-YMdfkZj70TIRhjvTEeZkNyZk21M,3695
-checkpoint_engine-0.1.1.dist-info/licenses/LICENCE,sha256=0jqA0jrA_i9VUqd7FTVoI1KnN1ZRENwG_tlMRCjC63k,1066
-checkpoint_engine-0.1.1.dist-info/METADATA,sha256=WDc5tg3RQiCthzbEzI4D3t3I0AQhGUPcDpJn7TMhYbI,9286
-checkpoint_engine-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-checkpoint_engine-0.1.1.dist-info/top_level.txt,sha256=66sik_1eLakLYmcllOEJzFaNbSfjsueuP0tHYEzhMSs,18
-checkpoint_engine-0.1.1.dist-info/RECORD,,