[xpu] use cpu barrier (#4181)

This commit is contained in:
zhupengyang
2025-09-23 12:19:03 +08:00
committed by GitHub
parent 813befadfa
commit 9082f625ba
2 changed files with 17 additions and 2 deletions

View File

@@ -101,6 +101,9 @@ class EngineWorkerQueue:
self.finish_request_barrier = [
threading.Barrier(self.num_client) for _ in range(self.local_data_parallel_size)
]
self.worker_process_tp_barrier = [
threading.Barrier(self.num_client) for _ in range(self.local_data_parallel_size)
]
self.finish_add_cache_task_barrier = [
threading.Barrier(self.num_client) for _ in range(self.local_data_parallel_size)
@@ -193,6 +196,10 @@ class EngineWorkerQueue:
"get_finish_add_cache_task_barrier",
callable=lambda idx: self.finish_add_cache_task_barrier[idx],
)
QueueManager.register(
"get_worker_process_tp_barrier",
callable=lambda idx: self.worker_process_tp_barrier[idx],
)
self.manager: BaseManager = QueueManager(address=self.address, authkey=self.authkey)
self.manager.start()
else:
@@ -217,6 +224,7 @@ class EngineWorkerQueue:
QueueManager.register("get_connect_rdma_tasks")
QueueManager.register("get_connect_rdma_tasks_responses")
QueueManager.register("get_connect_task_lock")
QueueManager.register("get_worker_process_tp_barrier")
self.manager = QueueManager(address=self.address, authkey=self.authkey)
self._connect_with_retry()
@@ -239,6 +247,7 @@ class EngineWorkerQueue:
self.finish_add_cache_task_barrier = self.manager.get_finish_add_cache_task_barrier(
self.local_data_parallel_id
)
self.worker_process_tp_barrier = self.manager.get_worker_process_tp_barrier(self.local_data_parallel_id)
self.finished_req_queue = self.manager.get_finish_request_queue(self.local_data_parallel_id)
self.finished_add_cache_task_queue = self.manager.get_finish_add_cache_task_queue(
self.local_data_parallel_id