fix multinode abort

Author: shihaobai

lightllm/server/httpserver/manager.py

@@ -50,11 +50,12 @@ def __init__(
         context = zmq.asyncio.Context(2)
         self.send_to_router = context.socket(zmq.PUSH)
         self.send_to_router.connect(f"{args.zmq_mode}127.0.0.1:{router_port}")
-        
+
         self.multinode_req_manager = None
         self.child_node_events = {}
         self.waiting_objs = []
         self.child_node_lock = asyncio.Lock()
+        self.nnodes = args.nnodes
         if args.nnodes > 1:
             if args.node_rank == 0:
                 self.multinode_req_manager = []
@@ -64,12 +65,16 @@ def __init__(
                     context = zmq.asyncio.Context(2)
                     self.multinode_req_manager.append(context.socket(zmq.PUSH))
                     self.multinode_req_manager[-1].connect(f"tcp://{child_ip}:{args.multinode_httpmanager_port}")
-                    logger.info(f"HttpServerManager connected to child node at {child_ip}:{args.multinode_httpmanager_port}")
+                    logger.info(
+                        f"HttpServerManager connected to child node at {child_ip}:{args.multinode_httpmanager_port}"
+                    )
             else:
                 context = zmq.asyncio.Context(2)
                 self.multinode_req_manager = context.socket(zmq.PULL)
                 self.multinode_req_manager.bind(f"tcp://*:{args.multinode_httpmanager_port}")
-                logger.info(f"HttpServerManager listening for child node requests on *:{args.multinode_httpmanager_port}")
+                logger.info(
+                    f"HttpServerManager listening for child node requests on *:{args.multinode_httpmanager_port}"
+                )
 
         self.enable_multimodal = enable_multimodal
         if self.enable_multimodal:
@@ -145,16 +150,26 @@ async def _release_multimodal_resources(self, multimodal_params: MultimodalParam
     def tokens(self, prompt):
         prompt_ids = self.tokenizer.encode(prompt)
         return len(prompt_ids)
-    
+
     async def loop_for_request(self):
         assert self.args.node_rank > 0
         tasks = []
         while True:
-            request_id, prompt, sampling_params, multimodal_params, request_headers = await self.multinode_req_manager.recv_pyobj()
-            results_generator = self.generate(prompt, sampling_params, multimodal_params, None, request_headers, request_id)
+            (
+                request_id,
+                prompt,
+                sampling_params,
+                multimodal_params,
+                request_headers,
+            ) = await self.multinode_req_manager.recv_pyobj()
+            results_generator = self.generate(
+                prompt, sampling_params, multimodal_params, None, request_headers, request_id
+            )
+
             async def generate_wrapper(results_generator):
                 async for _, _, _, _ in results_generator:
                     pass
+
             tasks.append(asyncio.create_task(generate_wrapper(results_generator)))
             # cleanup
             while len(tasks) > 0 and tasks[0].done():
@@ -166,7 +181,7 @@ async def generate(
         sampling_params: SamplingParams,
         multimodal_params: MultimodalParams,
         request: Request,
-        request_headers = None,
+        request_headers=None,
         multinode_remote_request_id: Optional[int] = None,
     ) -> Tuple[int, str, dict, FinishStatus]:
         start_time = time.time()
@@ -178,7 +193,10 @@ async def generate(
             if multinode_remote_request_id is None:
                 group_request_id = self.id_gen.generate_id()
                 for sender in self.multinode_req_manager:
-                    sender.send_pyobj((group_request_id, prompt, sampling_params, multimodal_params, request_headers), protocol=pickle.HIGHEST_PROTOCOL)
+                    sender.send_pyobj(
+                        (group_request_id, prompt, sampling_params, multimodal_params, request_headers),
+                        protocol=pickle.HIGHEST_PROTOCOL,
+                    )
             else:
                 group_request_id = multinode_remote_request_id
             sampling_params.group_request_id = group_request_id
@@ -238,8 +256,12 @@ async def generate(
             await self.transfer_to_next_module(req_status.group_req_objs)
 
             results_generator = self._wait_to_token_package(
-                start_time, prompt_ids, group_request_id, sampling_params, req_status, # request,
-                request_headers,
+                start_time,
+                prompt_ids,
+                group_request_id,
+                sampling_params,
+                req_status,
+                request,
             )
             async for sub_req_id, request_output, metadata, finish_status in results_generator:
                 # p d 模式下，将 token 数据放入到转发队列中
@@ -368,8 +390,7 @@ async def _wait_to_token_package(
         group_request_id: int,
         sampling_params: SamplingParams,
         req_status: "ReqStatus",
-        request_headers,
-        # request: Request,
+        request: Request,
     ):
 
         event = req_status.event
@@ -385,10 +406,9 @@ async def _wait_to_token_package(
             except asyncio.TimeoutError:
                 pass
 
-            # TODO: abort() for multinode
-            # if request is not None and await request.is_disconnected():
-            #     await self.abort(group_request_id)
-            #     raise Exception(f"req_id {group_request_id} disconnected")
+            if request is not None and await request.is_disconnected() and self.nnodes == 1:
+                await self.abort(group_request_id)
+                raise Exception(f"req_id {group_request_id} disconnected")
 
             async with req_status.lock:
                 event.clear()
@@ -416,13 +436,12 @@ async def _wait_to_token_package(
                         unfinished_count -= 1
 
                     # 所有子请求完成后，就删除占用的资源
-                    if unfinished_count == 0:
+                    if unfinished_count == 0 and request is not None:
                         total_cost_time_ms = (time.time() - start_time) * 1000
                         mean_per_token_cost_time_ms = (total_cost_time_ms - first_token_cost_ms) / out_token_counter
                         self.per_token_costs.add(mean_per_token_cost_time_ms)
-                        x_request_id = request_headers.get("X-Request-Id", "") if request_headers is not None else ""
-                        x_session_id = request_headers.get("X-Session-Id", "") if request_headers is not None else ""
-
+                        x_request_id = request.headers.get("X-Request-Id", "") if request is not None else ""
+                        x_session_id = request.headers.get("X-Session-Id", "") if request is not None else ""
                         prompt_cache_ratio = prompt_cache_len / prompt_tokens
                         self.metric_client.histogram_observe("lightllm_cache_length", prompt_cache_len)
                         self.metric_client.histogram_observe("lightllm_cache_ratio", prompt_cache_ratio)
@@ -506,7 +525,7 @@ async def handle_loop(self):
         if self.pd_mode.is_P_or_D():
             self.forwarding_queue = AsyncQueue()
             asyncio.create_task(self.pd_handle_loop())
-        
+
         if self.args.node_rank > 0:
             asyncio.create_task(self.loop_for_request())