fix

Author: hiworldwzj

lightllm/server/router/model_infer/mode_backend/pd_nixl/prefill_node_impl/prefill_trans_process.py

@@ -4,14 +4,12 @@
 import threading
 import setproctitle
 import torch.multiprocessing as mp
-import collections
 import queue
 import pickle
-from typing import List, Dict, Union, Deque, Optional
+from typing import List, Dict, Optional
 from lightllm.utils.log_utils import init_logger
 from lightllm.common.kv_cache_mem_manager import MemoryManager
-from lightllm.server.pd_io_struct import NIXLChunckedTransTask, NIXLChunckedTransTaskRet
-from lightllm.utils.device_utils import kv_trans_use_p2p
+from lightllm.server.pd_io_struct import NIXLChunckedTransTask
 from lightllm.utils.graceful_utils import graceful_registry
 from lightllm.server.core.objs import StartArgs
 from ..nixl_kv_transporter import NixlKVTransporter
@@ -45,9 +43,8 @@ def _init_env(
 
     import os
 
-    # prefill 节点不一定需要 mps 来协调，所以优先级设置为 1.
-    # 本身并不产生严重的阻塞。
-    os.environ["CUDA_MPS_CLIENT_PRIORITY"] = "1"
+    # prefill source-side page copy and UCX progress are on the request critical path.
+    os.environ["CUDA_MPS_CLIENT_PRIORITY"] = "0"
 
     torch.backends.cudnn.enabled = False
     setproctitle.setproctitle(f"lightllm::{get_unique_server_name()}::nixl_prefill_trans:Device{device_id}")
@@ -103,15 +100,18 @@ def __init__(
         kv_move_buffer = cur_mem_manager.alloc_paged_kv_move_buffer(
             page_num=self.args.nixl_pd_kv_page_num, page_size=self.args.nixl_pd_kv_page_size
         )
-        self.copy_cuda_stream = torch.cuda.Stream()
+        self.copy_cuda_stream = torch.cuda.Stream(priority=-1)
         self.transporter = NixlKVTransporter(
             node_id=self.args.pd_node_id, tp_idx=device_id, kv_move_buffer=kv_move_buffer
         )
         self.waiting_dict_lock = threading.Lock()
         self.waiting_dict: Dict[str, NIXLChunckedTransTask] = {}
+        self.waiting_nixl_write_task_lock = threading.Lock()
+        self.waiting_nixl_write_task_dict: Dict[str, NIXLChunckedTransTask] = {}
 
         self.local_copy_kv_queue = queue.Queue()
-        self.notify_peer_read_kv_queue = queue.Queue()
+        self.ready_transfer_queue = queue.Queue()
+        self.write_peer_kv_queue = queue.Queue()
         self.success_queue = queue.Queue()
         self.failed_queue = queue.Queue()
 
@@ -125,14 +125,30 @@ def __init__(
         for func in [
             self.recv_task_loop,
             self.local_copy_kv_loop,
-            self.notify_peer_to_read_kv_loop,
+            self.ready_transfer_loop,
+            self.accept_decode_write_task_loop,
+            self.write_peer_kv_loop,
             self.update_task_status_loop,
             self.success_loop,
             self.fail_loop,
         ]:
             threading.Thread(target=func, daemon=True).start()
         return
 
+    def _warmup(self):
+        for dp_index in range(self.args.dp // self.args.nnodes):
+            with torch.cuda.stream(stream=self.copy_cuda_stream):
+                cur_mem = self.mem_managers[self.device_id]
+                cur_mem.write_mem_to_page_kv_move_buffer(
+                    mem_indexes=[0],
+                    page_index=0,
+                    dp_index=dp_index,
+                    mem_managers=self.mem_managers,
+                    dp_world_size=self.dp_world_size,
+                )
+                torch.cuda.current_stream().synchronize()
+        return
+
     @log_exception
     def recv_task_loop(self):
         torch.cuda.set_device(self.device_id)
@@ -168,65 +184,26 @@ def local_copy_kv_loop(self):
                 sync_event = torch.cuda.Event()
                 sync_event.record()
 
-            self.notify_peer_read_kv_queue.put((sync_event, trans_task))
-        return
-
-    def _warmup(self):
-        for dp_index in range(self.args.dp // self.args.nnodes):
-            with torch.cuda.stream(stream=self.copy_cuda_stream):
-                cur_mem = self.mem_managers[self.device_id]
-                cur_mem.write_mem_to_page_kv_move_buffer(
-                    mem_indexes=[0],
-                    page_index=0,
-                    dp_index=dp_index,
-                    mem_managers=self.mem_managers,
-                    dp_world_size=self.dp_world_size,
-                )
-                torch.cuda.current_stream().synchronize()
+            self.ready_transfer_queue.put((sync_event, trans_task))
         return
 
     @log_exception
-    def notify_peer_to_read_kv_loop(self):
+    def ready_transfer_loop(self):
         torch.cuda.set_device(self.device_id)
         while True:
-            sync_event, trans_task = self.notify_peer_read_kv_queue.get()
+            sync_event, trans_task = self.ready_transfer_queue.get()
             trans_task: NIXLChunckedTransTask = trans_task
             sync_event: torch.cuda.Event = sync_event
-
             sync_event.synchronize()
-
-            trans_task.start_trans_time = time.time()
+            self.transporter.send_write_request_task_to_decode_node(trans_task)
+            logger.info(f"send WRITE request to decode: {trans_task.get_key()}")
             with self.waiting_dict_lock:
                 self.waiting_dict[trans_task.get_key()] = trans_task
-
-            try:
-                self.transporter.send_readtask_to_decode_node(trans_task=trans_task)
-            except BaseException as e:
-                logger.error(f"send readtask to decode node failed: {trans_task.to_str()}")
-                logger.exception(str(e))
-                self.transporter.remove_remote_agent(peer_name=trans_task.decode_agent_name)
-
-                with self.waiting_dict_lock:
-                    trans_task = self.waiting_dict.pop(trans_task.get_key(), None)
-
-                if trans_task is not None:
-                    trans_task.error_info = f"send readtask to decode node failed: {str(e)}"
-                    self.failed_queue.put(trans_task)
-                continue
-
-            logger.info(f"send readtask to decode: {trans_task.to_str()}")
         return
 
     @log_exception
-    def update_task_status_loop(
-        self,
-    ):
+    def accept_decode_write_task_loop(self):
         while True:
-            if len(self.waiting_dict) == 0:
-                time.sleep(0.001)
-                continue
-
-            # notify update
             try:
                 notifies_dict = self.transporter.get_new_notifs()
             except BaseException as e:
@@ -239,44 +216,133 @@ def update_task_status_loop(
                     for notify in _notify_list:
                         try:
                             notify_obj = pickle.loads(notify)
-                        except:
+                        except BaseException:
                             notify_obj = None
 
-                        if isinstance(notify_obj, NIXLChunckedTransTaskRet):
-                            key = notify_obj.get_key()
-                            with self.waiting_dict_lock:
-                                trans_task = self.waiting_dict.pop(key, None)
-
-                            if trans_task is not None:
-                                trans_task.error_info = notify_obj.error_info
-                                if trans_task.error_info is not None:
-                                    self.failed_queue.put(trans_task)
-                                else:
-                                    self.success_queue.put(trans_task)
-                            else:
-                                logger.warning(f"can not find trans task for ret: {notify_obj}")
-
-            # check time_out update
+                        if not isinstance(notify_obj, NIXLChunckedTransTask):
+                            continue
+
+                        if notify_obj.nixl_write_stage != "ready":
+                            logger.warning(f"ignore unknown prefill WRITE notify stage: {notify_obj.to_str()}")
+                            continue
+
+                        key = notify_obj.get_key()
+                        with self.waiting_dict_lock:
+                            trans_task = self.waiting_dict.pop(key, None)
+
+                        if trans_task is None:
+                            logger.warning(
+                                f"can not find pending WRITE request for ready notify: {notify_obj.to_str()}"
+                            )
+                            continue
+
+                        trans_task.nixl_dst_page_index = notify_obj.nixl_dst_page_index
+                        logger.info(
+                            f"recv WRITE ready from decode request_id={trans_task.request_id} "
+                            f"kv=[{trans_task.start_kv_index},{trans_task.end_kv_index}) "
+                            f"src_page={trans_task.nixl_src_page_index} dst_page={trans_task.nixl_dst_page_index}"
+                        )
+                        self.write_peer_kv_queue.put(trans_task)
+
             self._check_tasks_time_out()
 
+            if not notifies_dict:
+                time.sleep(0.001)
+        return
+
     def _check_tasks_time_out(self):
         with self.waiting_dict_lock:
-            keys = list(self.waiting_dict.keys())
+            timeout_tasks = []
+            for key, trans_task in list(self.waiting_dict.items()):
+                if trans_task.time_out():
+                    timeout_tasks.append(self.waiting_dict.pop(key))
+
+        for trans_task in timeout_tasks:
+            trans_task.error_info = "time out waiting decode WRITE ready"
+            self.failed_queue.put(trans_task)
+        return
 
-        for key in keys:
-            with self.waiting_dict_lock:
-                trans_task = self.waiting_dict.pop(key, None)
+    @log_exception
+    def write_peer_kv_loop(self):
+        torch.cuda.set_device(self.device_id)
+        while True:
+            trans_task = self.write_peer_kv_queue.get()
+            trans_task: NIXLChunckedTransTask = trans_task
 
-            if trans_task is not None and trans_task.time_out():
-                trans_task.error_info = "time out in update_task_status_loop"
+            try:
+                xfer_handle = self.transporter.write_blocks_paged(trans_task=trans_task)
+                trans_task.xfer_handle = xfer_handle
+            except BaseException as e:
+                logger.error(f"write_blocks_paged failed: {trans_task.to_str()}")
+                logger.exception(str(e))
+                self.transporter.remove_remote_agent(peer_name=trans_task.decode_agent_name)
+
+                trans_task.error_info = f"write_blocks_paged failed: {str(e)}"
                 self.failed_queue.put(trans_task)
                 continue
 
-            if trans_task is not None:
-                with self.waiting_dict_lock:
-                    self.waiting_dict[trans_task.get_key()] = trans_task
+            trans_task.start_trans_time = time.time()
+            with self.waiting_nixl_write_task_lock:
+                self.waiting_nixl_write_task_dict[trans_task.get_key()] = trans_task
+            logger.info(f"start WRITE to decode node: {trans_task.to_str()}")
         return
 
+    @log_exception
+    def update_task_status_loop(
+        self,
+    ):
+        while True:
+            if len(self.waiting_nixl_write_task_dict) == 0:
+                time.sleep(0.001)
+                continue
+
+            with self.waiting_nixl_write_task_lock:
+                tasks = list(self.waiting_nixl_write_task_dict.values())
+
+            for trans_task in tasks:
+                ret = self.transporter.check_task_status(trans_task=trans_task)
+                if ret == "DONE":
+                    with self.waiting_nixl_write_task_lock:
+                        trans_task = self.waiting_nixl_write_task_dict.pop(trans_task.get_key(), None)
+                    if trans_task is None:
+                        continue
+                    if self.transporter.capture_telemetry:
+                        telem = self.transporter.nixl_agent.get_xfer_telemetry(trans_task.xfer_handle)
+                        total_us = telem.xferDuration
+                        post_us = telem.postDuration
+                        backend_us = telem.xferDuration - telem.postDuration
+                        nixl_backend = self.transporter.nixl_agent.query_xfer_backend(trans_task.xfer_handle)
+                        logger.info(
+                            f"write trans task request_id={trans_task.request_id} "
+                            f"kv=[{trans_task.start_kv_index},{trans_task.end_kv_index}) "
+                            f"src_page={trans_task.nixl_src_page_index} dst_page={trans_task.nixl_dst_page_index} "
+                            f"xfer time: {total_us:.3f} us, "
+                            f"post time: {post_us:.3f} us, backend time: {backend_us:.3f} us, "
+                            f"nixl_backend: {nixl_backend}, total_bytes: {telem.totalBytes}"
+                        )
+                    self.transporter.send_write_done_task_to_decode_node(trans_task)
+                    logger.info(
+                        f"send WRITE done nixl notify "
+                        f"request_id={trans_task.request_id} "
+                        f"kv=[{trans_task.start_kv_index},{trans_task.end_kv_index}) "
+                        f"src_page={trans_task.nixl_src_page_index} dst_page={trans_task.nixl_dst_page_index}"
+                    )
+                    self.success_queue.put(trans_task)
+                elif ret == "ERR":
+                    with self.waiting_nixl_write_task_lock:
+                        trans_task = self.waiting_nixl_write_task_dict.pop(trans_task.get_key(), None)
+                    if trans_task is not None:
+                        trans_task.error_info = "xfer error"
+                        self.failed_queue.put(trans_task)
+                elif trans_task.time_out():
+                    with self.waiting_nixl_write_task_lock:
+                        trans_task = self.waiting_nixl_write_task_dict.pop(trans_task.get_key(), None)
+                    if trans_task is not None:
+                        trans_task.error_info = "time out in update_task_status_loop"
+                        self.failed_queue.put(trans_task)
+
+            time.sleep(0.001)
+
     @log_exception
     def success_loop(self):
         torch.cuda.set_device(self.device_id)
@@ -285,6 +351,8 @@ def success_loop(self):
             # 写回后，回收页面
             if trans_task.nixl_src_page_index is not None:
                 self.page_index_queue.put(trans_task.nixl_src_page_index)
+            if trans_task.xfer_handle is not None:
+                self.transporter.release_xfer_handle(trans_task.xfer_handle)
 
             ret = trans_task.createRetObj()
             ret.first_gen_token_id = None
@@ -301,6 +369,8 @@ def fail_loop(self):
             # 回收页面
             if trans_task.nixl_src_page_index is not None:
                 self.page_index_queue.put(trans_task.nixl_src_page_index)
+            if trans_task.xfer_handle is not None:
+                self.transporter.release_xfer_handle(trans_task.xfer_handle)
 
             ret = trans_task.createRetObj()
             self.task_out_queue.put(ret)