fix pre-commit

Author: kip-cxj

checkpoint_engine/distributed/__init__.py

@@ -1,23 +1,23 @@
 from .base import (
     Distributed,
-    init_process_group,
-    destroy_process_group,
-    is_initialized,
     all_gather_object,
     all_reduce,
-    broadcast,
     barrier,
+    broadcast,
+    destroy_process_group,
+    init_process_group,
+    is_initialized,
     new_group,
 )
 
 __all__ = [
     "Distributed",
-    "init_process_group",
-    "destroy_process_group",
-    "is_initialized",
     "all_gather_object",
     "all_reduce",
-    "broadcast",
     "barrier",
+    "broadcast",
+    "destroy_process_group",
+    "init_process_group",
+    "is_initialized",
     "new_group",
 ]

checkpoint_engine/distributed/base.py

@@ -1,9 +1,9 @@
-from abc import ABC, abstractmethod
+import importlib
 import io
 import pickle
+from abc import ABC, abstractmethod
 from datetime import timedelta
 from typing import Any, List
-import importlib
 
 import torch
 import torch.distributed as torch_dist
@@ -169,6 +169,7 @@ def is_initialized() -> bool:
         return torch_dist.is_initialized()
     return _BACKEND_INSTANCE.is_initialized()
 
+
 def all_gather_object(
     object_list: list[Any],
     obj: Any,

checkpoint_engine/distributed/hccl.py

@@ -14,12 +14,10 @@
     hcclComm_t,
     hcclDataType_t,
     hcclDataTypeEnum,
-    hcclRedOp_t,
-    hcclRedOpTypeEnum,
     hcclResult_t,
-    hcclUniqueId,
 )
 from vllm_ascend.utils import current_stream
+
 from checkpoint_engine.distributed.base import Distributed, _common_all_gather_object
 
 
@@ -214,7 +212,6 @@ def init_process_group(
         self.comm = self.pyhccl.comm
         self.initialized = True
 
-
     def destroy_process_group(
         self,
         group=None,
@@ -236,13 +233,7 @@ def destroy_process_group(
     def is_initialized(self) -> bool:
         return self.initialized
 
-
-    def all_gather_object(
-        self,
-        object_list: list[Any],
-        obj: Any,
-        group=None
-    ):
+    def all_gather_object(self, object_list: list[Any], obj: Any, group=None):
         assert self.initialized, "not initialized"
 
         if group:
@@ -257,12 +248,7 @@ def all_gather_object(
             self.pyhccl.comm = self.comm
 
 
-    def all_reduce(
-        self,
-        tensor: torch.Tensor,
-        op=ReduceOp.SUM,
-        group=None
-    ):
+    def all_reduce(self, tensor: torch.Tensor, op=ReduceOp.SUM, group=None):
         assert self.initialized, "not initialized"
 
         if group:
@@ -278,12 +264,7 @@ def all_reduce(
             self.pyhccl.comm = self.comm
 
 
-    def broadcast(
-        self,
-        tensor: torch.Tensor,
-        src=None,
-        group=None
-    ):
+    def broadcast(self, tensor: torch.Tensor, src=None, group=None):
         assert self.initialized, "not initialized"
 
         if group:
@@ -303,10 +284,7 @@ def broadcast(
             self.pyhccl.rank = self.rank
 
 
-    def barrier(
-        self,
-        group=None
-    ):
+    def barrier(self, group=None):
         assert self.initialized, "not initialized"
 
         if group:
@@ -322,10 +300,7 @@ def barrier(
             self.pyhccl.comm = self.comm
 
 
-    def new_group(
-        self,
-        ranks
-    ):
+    def new_group(self, ranks):
         assert self.initialized, "not initialized"
 
         # if ranks is None or [], using the world instead

checkpoint_engine/distributed/nccl.py

@@ -1,19 +1,19 @@
 import ctypes
 from datetime import timedelta
-from typing import Any, List, Optional
+from typing import Any
 
 import torch
 from torch.distributed import ReduceOp
 from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
 from vllm.distributed.device_communicators.pynccl_wrapper import (
     Function,
     NCCLLibrary,
-    buffer_type,
     ncclComm_t,
     ncclResult_t,
 )
 from vllm.distributed.utils import StatelessProcessGroup
 from vllm.utils import current_stream
+
 from checkpoint_engine.distributed.base import Distributed, _common_all_gather_object
 
 
@@ -132,7 +132,6 @@ def init_process_group(
         self.comm = self.pynccl.comm
         self.initialized = True
 
-
     def destroy_process_group(
         self,
         group=None,
@@ -155,12 +154,7 @@ def is_initialized(self) -> bool:
         return self.initialized
 
 
-    def all_gather_object(
-        self,
-        object_list: list[Any],
-        obj: Any,
-        group=None
-    ):
+    def all_gather_object(self, object_list: list[Any], obj: Any, group=None):
         assert self.initialized, "not initialized"
 
         if group:
@@ -175,12 +169,7 @@ def all_gather_object(
             self.pynccl.comm = self.comm
 
 
-    def all_reduce(
-        self,
-        tensor: torch.Tensor,
-        op=ReduceOp.SUM,
-        group=None
-    ):
+    def all_reduce(self, tensor: torch.Tensor, op=ReduceOp.SUM, group=None):
         assert self.initialized, "not initialized"
 
         if group:
@@ -196,12 +185,7 @@ def all_reduce(
             self.pynccl.comm = self.comm
 
 
-    def broadcast(
-        self,
-        tensor: torch.Tensor,
-        src=None,
-        group=None
-    ):
+    def broadcast(self, tensor: torch.Tensor, src=None, group=None):
         assert self.initialized, "not initialized"
 
         if group:
@@ -221,10 +205,7 @@ def broadcast(
             self.pynccl.rank = self.rank
 
 
-    def barrier(
-        self,
-        group=None
-    ):
+    def barrier(self, group=None):
         assert self.initialized, "not initialized"
 
         if group:
@@ -240,10 +221,7 @@ def barrier(
             self.pynccl.comm = self.comm
 
 
-    def new_group(
-        self,
-        ranks
-    ):
+    def new_group(self, ranks):
         assert self.initialized, "not initialized"
 
         # ranks is None or []

checkpoint_engine/ps.py

@@ -519,9 +519,7 @@ def init_process_group(
             )
         logger.info(f"[rank{self._rank}] init process group successfully.")
 
-    def store_based_barrier(
-        self, store, timeout: timedelta = timedelta(minutes=5)
-    ) -> None:
+    def store_based_barrier(self, store, timeout: timedelta = timedelta(minutes=5)) -> None:
         """
         Perform a store-based barrier synchronization across all ranks.
 
@@ -571,11 +569,10 @@ def update(
         try:
             master_addr = os.getenv("MASTER_ADDR") or master_addr
             assert master_addr, "master_addr is required"
-            if self._auto_pg:
-                if not dist.is_initialized():
-                    self.init_process_group(
-                        timeout=timeout, master_addr=master_addr, master_port=master_port
-                    )
+            if self._auto_pg and not dist.is_initialized():
+                self.init_process_group(
+                    timeout=timeout, master_addr=master_addr, master_port=master_port
+                )
             # if ranks is None or [], it will use fully broadcast to update to all ranks
             ranks_group = dist.new_group(ranks) if ranks else None
             self._update_per_bucket(checkpoint_name, req_func, ranks_group, ranks)

examples/update.py

@@ -13,9 +13,9 @@
 from loguru import logger
 from safetensors import safe_open
 
+import checkpoint_engine.distributed as dist
 from checkpoint_engine import request_inference_to_update
 from checkpoint_engine.ps import ParameterServer
-import checkpoint_engine.distributed as dist
 
 
 @contextmanager