polish code with new pre-commit rule (#2923)

fastdeploy/inter_communicator/engine_cache_queue.py

@@ -16,8 +16,13 @@
 
 import threading
 import time
-from multiprocessing.managers import (AcquirerProxy, BaseManager, ListProxy,
-                                      Value, ValueProxy)
+from multiprocessing.managers import (
+    AcquirerProxy,
+    BaseManager,
+    ListProxy,
+    Value,
+    ValueProxy,
+)
 from typing import Any, List, Tuple
 
 from fastdeploy.utils import get_logger
@@ -32,14 +37,14 @@ class EngineCacheQueue:
     """
 
     def __init__(
-            self,
-            address: Tuple[str, int] = ('127.0.0.1', 56666),
-            authkey: bytes = b'cache_queue_service',
-            is_server: bool = False,
-            num_client: int = 1,  # tensor parallel size
-            client_id: int = -1,  # tensor parallel id
-            local_data_parallel_size: int = 1,  # data parallel size
-            local_data_parallel_id: int = 0,  # local data parallel id
+        self,
+        address: Tuple[str, int] = ("127.0.0.1", 56666),
+        authkey: bytes = b"cache_queue_service",
+        is_server: bool = False,
+        num_client: int = 1,  # tensor parallel size
+        client_id: int = -1,  # tensor parallel id
+        local_data_parallel_size: int = 1,  # data parallel size
+        local_data_parallel_id: int = 0,  # local data parallel id
     ) -> None:
         """
         Initialize the cache communication queue.
@@ -64,19 +69,14 @@ class EngineCacheQueue:
             """
             Custom QueueManager for proxy object registration
             """
+
             pass
 
         if is_server:
             # Server-side initialization for shared resources
-            self.transfer_task_queue_init: List[List[Any]] = [
-                list() for _ in range(self.local_data_parallel_size)
-            ]
-            self.tansfer_done_queue_init: List[List[Any]] = [
-                list() for _ in range(self.local_data_parallel_size)
-            ]
-            self.cache_sync_value_init: List[Value] = [
-                Value("i", 0) for _ in range(self.local_data_parallel_size)
-            ]
+            self.transfer_task_queue_init: List[List[Any]] = [list() for _ in range(self.local_data_parallel_size)]
+            self.tansfer_done_queue_init: List[List[Any]] = [list() for _ in range(self.local_data_parallel_size)]
+            self.cache_sync_value_init: List[Value] = [Value("i", 0) for _ in range(self.local_data_parallel_size)]
             self.transfer_task_lock_init: List[threading.Lock] = [
                 threading.Lock() for _ in range(self.local_data_parallel_size)
             ]
@@ -85,84 +85,76 @@ class EngineCacheQueue:
             ]
 
             # Initialize barriers
-            self.barrier1_init = [
-                threading.Barrier(self.num_client)
-                for _ in range(self.local_data_parallel_size)
-            ]
-            self.barrier2_init = [
-                threading.Barrier(self.num_client)
-                for _ in range(self.local_data_parallel_size)
-            ]
-            self.barrier3_init = [
-                threading.Barrier(self.num_client)
-                for _ in range(self.local_data_parallel_size)
-            ]
+            self.barrier1_init = [threading.Barrier(self.num_client) for _ in range(self.local_data_parallel_size)]
+            self.barrier2_init = [threading.Barrier(self.num_client) for _ in range(self.local_data_parallel_size)]
+            self.barrier3_init = [threading.Barrier(self.num_client) for _ in range(self.local_data_parallel_size)]
             self.swap_to_cpu_barrier1_init = [
-                threading.Barrier(self.num_client)
-                for _ in range(self.local_data_parallel_size)
+                threading.Barrier(self.num_client) for _ in range(self.local_data_parallel_size)
             ]
             self.swap_to_cpu_barrier2_init = [
-                threading.Barrier(self.num_client)
-                for _ in range(self.local_data_parallel_size)
+                threading.Barrier(self.num_client) for _ in range(self.local_data_parallel_size)
             ]
             self.swap_to_gpu_barrier1_init = [
-                threading.Barrier(self.num_client)
-                for _ in range(self.local_data_parallel_size)
+                threading.Barrier(self.num_client) for _ in range(self.local_data_parallel_size)
             ]
             self.swap_to_gpu_barrier2_init = [
-                threading.Barrier(self.num_client)
-                for _ in range(self.local_data_parallel_size)
+                threading.Barrier(self.num_client) for _ in range(self.local_data_parallel_size)
             ]
 
             # Register shared objects with proxy types
             QueueManager.register(
                 "get_transfer_task_queue",
                 callable=lambda idx: self.transfer_task_queue_init[idx],
-                proxytype=ListProxy)
+                proxytype=ListProxy,
+            )
             QueueManager.register(
                 "get_tansfer_done_queue",
                 callable=lambda idx: self.tansfer_done_queue_init[idx],
-                proxytype=ListProxy)
+                proxytype=ListProxy,
+            )
             QueueManager.register(
                 "get_cache_sync_value",
                 callable=lambda idx: self.cache_sync_value_init[idx],
-                proxytype=ValueProxy)
+                proxytype=ValueProxy,
+            )
             QueueManager.register(
                 "get_transfer_task_lock",
                 callable=lambda idx: self.transfer_task_lock_init[idx],
-                proxytype=AcquirerProxy)
+                proxytype=AcquirerProxy,
+            )
             QueueManager.register(
                 "get_transfer_task_done_lock",
                 callable=lambda idx: self.transfer_task_done_lock_init[idx],
-                proxytype=AcquirerProxy)
-            QueueManager.register("get_barrier1",
-                                  callable=lambda idx: self.barrier1_init[idx])
-            QueueManager.register("get_barrier2",
-                                  callable=lambda idx: self.barrier2_init[idx])
-            QueueManager.register("get_barrier3",
-                                  callable=lambda idx: self.barrier3_init[idx])
+                proxytype=AcquirerProxy,
+            )
+            QueueManager.register("get_barrier1", callable=lambda idx: self.barrier1_init[idx])
+            QueueManager.register("get_barrier2", callable=lambda idx: self.barrier2_init[idx])
+            QueueManager.register("get_barrier3", callable=lambda idx: self.barrier3_init[idx])
             QueueManager.register(
                 "get_swap_to_cpu_barrier1",
-                callable=lambda idx: self.swap_to_cpu_barrier1_init[idx])
+                callable=lambda idx: self.swap_to_cpu_barrier1_init[idx],
+            )
             QueueManager.register(
                 "get_swap_to_cpu_barrier2",
-                callable=lambda idx: self.swap_to_cpu_barrier2_init[idx])
+                callable=lambda idx: self.swap_to_cpu_barrier2_init[idx],
+            )
             QueueManager.register(
                 "get_swap_to_gpu_barrier1",
-                callable=lambda idx: self.swap_to_gpu_barrier1_init[idx])
+                callable=lambda idx: self.swap_to_gpu_barrier1_init[idx],
+            )
             QueueManager.register(
                 "get_swap_to_gpu_barrier2",
-                callable=lambda idx: self.swap_to_gpu_barrier2_init[idx])
+                callable=lambda idx: self.swap_to_gpu_barrier2_init[idx],
+            )
 
-            self.manager: BaseManager = QueueManager(address=self.address,
-                                                     authkey=self.authkey)
+            self.manager: BaseManager = QueueManager(address=self.address, authkey=self.authkey)
             self.manager.start()
             logger.info(f"EngineCacheQueue server started at {self.address}")
         else:
             # Client-side connection setup
-            assert 0 <= self.client_id < self.num_client, (
-                f"client_id must be between 0 and {self.num_client-1}, got {self.client_id}"
-            )
+            assert (
+                0 <= self.client_id < self.num_client
+            ), f"client_id must be between 0 and {self.num_client-1}, got {self.client_id}"
             QueueManager.register("get_transfer_task_queue")
             QueueManager.register("get_tansfer_done_queue")
             QueueManager.register("get_cache_sync_value")
@@ -176,45 +168,32 @@ class EngineCacheQueue:
             QueueManager.register("get_swap_to_gpu_barrier1")
             QueueManager.register("get_swap_to_gpu_barrier2")
 
-            self.manager = QueueManager(address=self.address,
-                                        authkey=self.authkey)
+            self.manager = QueueManager(address=self.address, authkey=self.authkey)
             self._connect_with_retry()
 
         # Get proxy objects for shared resources
-        self.transfer_task_queue = self.manager.get_transfer_task_queue(
-            self.local_data_parallel_id)
-        self.tansfer_done_queue = self.manager.get_tansfer_done_queue(
-            self.local_data_parallel_id)
-        self.task_sync_value = self.manager.get_cache_sync_value(
-            self.local_data_parallel_id)
-        self.task_lock = self.manager.get_transfer_task_lock(
-            self.local_data_parallel_id)
-        self.task_done_lock = self.manager.get_transfer_task_done_lock(
-            self.local_data_parallel_id)
+        self.transfer_task_queue = self.manager.get_transfer_task_queue(self.local_data_parallel_id)
+        self.tansfer_done_queue = self.manager.get_tansfer_done_queue(self.local_data_parallel_id)
+        self.task_sync_value = self.manager.get_cache_sync_value(self.local_data_parallel_id)
+        self.task_lock = self.manager.get_transfer_task_lock(self.local_data_parallel_id)
+        self.task_done_lock = self.manager.get_transfer_task_done_lock(self.local_data_parallel_id)
 
         # Get barrier proxies
         self.barrier1 = self.manager.get_barrier1(self.local_data_parallel_id)
         self.barrier2 = self.manager.get_barrier2(self.local_data_parallel_id)
         self.barrier3 = self.manager.get_barrier3(self.local_data_parallel_id)
-        self.swap_to_cpu_barrier1 = self.manager.get_swap_to_cpu_barrier1(
-            self.local_data_parallel_id)
-        self.swap_to_cpu_barrier2 = self.manager.get_swap_to_cpu_barrier2(
-            self.local_data_parallel_id)
-        self.swap_to_gpu_barrier1 = self.manager.get_swap_to_gpu_barrier1(
-            self.local_data_parallel_id)
-        self.swap_to_gpu_barrier2 = self.manager.get_swap_to_gpu_barrier2(
-            self.local_data_parallel_id)
+        self.swap_to_cpu_barrier1 = self.manager.get_swap_to_cpu_barrier1(self.local_data_parallel_id)
+        self.swap_to_cpu_barrier2 = self.manager.get_swap_to_cpu_barrier2(self.local_data_parallel_id)
+        self.swap_to_gpu_barrier1 = self.manager.get_swap_to_gpu_barrier1(self.local_data_parallel_id)
+        self.swap_to_gpu_barrier2 = self.manager.get_swap_to_gpu_barrier2(self.local_data_parallel_id)
         self.total_num: int = (1 << self.num_client) - 1
 
         if not is_server:
             # Setup position and total_num for sync operations
             self.position: int = 1 << self.client_id
-            logger.info(
-                f"Connected EngineCacheQueue client_id: {self.client_id}")
+            logger.info(f"Connected EngineCacheQueue client_id: {self.client_id}")
 
-    def _connect_with_retry(self,
-                            max_retries: int = 5,
-                            interval: int = 3) -> None:
+    def _connect_with_retry(self, max_retries: int = 5, interval: int = 3) -> None:
         """
         Connect to the server with retry mechanism.
 
@@ -231,8 +210,7 @@ class EngineCacheQueue:
                 return
             except ConnectionRefusedError:
                 time.sleep(interval)
-        raise ConnectionError(
-            f"EngineCacheQueue cannot connect to {self.address}")
+        raise ConnectionError(f"EngineCacheQueue cannot connect to {self.address}")
 
     def put_transfer_task(self, item):
         """
@@ -246,8 +224,7 @@ class EngineCacheQueue:
             self.task_lock.acquire()
         self.task_sync_value.set(0)
         self.transfer_task_queue.append(item)
-        logger.info(
-            f"put_transfer_task: put swap task {item[-1]} to queue successful")
+        logger.info(f"put_transfer_task: put swap task {item[-1]} to queue successful")
         self.task_lock.release()
 
     def get_transfer_task(self):
@@ -257,15 +234,11 @@ class EngineCacheQueue:
         data = None
         read_finish = False
         self.task_lock.acquire()
-        if (self.task_sync_value.get() & self.position == 0
-                and len(self.transfer_task_queue) > 0):
+        if self.task_sync_value.get() & self.position == 0 and len(self.transfer_task_queue) > 0:
             data = self.transfer_task_queue[0]
-            logger.debug(
-                f"get_transfer_task: Get {data} by {self.client_id} from queue successful"
-            )
+            logger.debug(f"get_transfer_task: Get {data} by {self.client_id} from queue successful")
             set_value = self.task_sync_value.get() | self.position
-            logger.info("get_transfer_task: rank: {0} set_value: {1}".format(
-                self.client_id, set_value))
+            logger.info(f"get_transfer_task: rank: {self.client_id} set_value: {set_value}")
             if set_value >= self.total_num:
                 self.transfer_task_queue.pop(0)
                 set_value = 0
@@ -281,9 +254,7 @@ class EngineCacheQueue:
         self.task_done_lock.acquire()
         self.tansfer_done_queue.append(item)
         self.task_done_lock.release()
-        logger.info(
-            f"put_transfer_done_signal: put swap task {item[-1]} finished signal to queue successful"
-        )
+        logger.info(f"put_transfer_done_signal: put swap task {item[-1]} finished signal to queue successful")
 
     def get_transfer_done_signal(self):
         """
@@ -293,9 +264,7 @@ class EngineCacheQueue:
         self.task_done_lock.acquire()
         if len(self.tansfer_done_queue) > 0:
             data = self.tansfer_done_queue.pop(0)
-            logger.info(
-                f"get_transfer_done_signal: Get swap task {data[-1]} finished signal from queue successful"
-            )
+            logger.info(f"get_transfer_done_signal: Get swap task {data[-1]} finished signal from queue successful")
         self.task_done_lock.release()
         return data