Merge pull request #93 from amy-why-3459/bug_fix_v0.9.1

【bugfix】Support individual health check calls

Author: zengchuang-hw

examples/offline_inference/epd/chat_with_image.py

@@ -8,6 +8,7 @@
 from PIL import Image
 
 from vllm import SamplingParams
+from vllm.disaggregated.protocol import ServerType
 from vllm.disaggregated.proxy import Proxy
 
 parser = argparse.ArgumentParser()
@@ -56,11 +57,40 @@ async def main():
         pd_addr_list=args.pd_addr_list,
         model_name=args.model_name,
     )
+    retry_times = 20
+    sleep_times = 10
+    timeout_times = 3
+    for i in range(retry_times):
+        tasks_0 = [
+            asyncio.create_task(
+                asyncio.wait_for(
+                    p.check_health(ServerType.E_INSTANCE, iid), timeout=timeout_times
+                )
+            )
+            for iid in range(len(args.encode_addr_list))
+        ]
+        tasks_1 = [
+            asyncio.create_task(
+                asyncio.wait_for(
+                    p.check_health(ServerType.PD_INSTANCE, iid), timeout=timeout_times
+                )
+            )
+            for iid in range(len(args.pd_addr_list))
+        ]
+        tasks = tasks_0 + tasks_1
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        if all([isinstance(result, bool) and result for result in results]):
+            print("All instances are ready")
+            break
+        else:
+            print(f"retry_times:{i}, results:{results}")
+            await asyncio.sleep(sleep_times)
+
     prompt = (
         "<|im_start|> system\n"
         "You are a helpful assistant.<|im_end|> \n"
         "<|im_start|> user\n"
-        "<image> \n"
+        "<|image_pad|> \n"
         "What is the text in the illustrate?<|im_end|> \n"
         "<|im_start|> assistant\n"
     )
@@ -73,6 +103,7 @@ async def main():
         for i in range(10)
     ]
     await asyncio.gather(*tasks)
+    p.shutdown()
 
 
 if __name__ == "__main__":

vllm/disaggregated/proxy.py

@@ -49,7 +49,7 @@ def __init__(
         pd_addr_list: list[str],
         model_name: str,
         router: type[RoutingInterface] = RandomRouter,
-        enable_health_monitor=True,
+        enable_health_monitor=False,
         health_check_interval=10,
         health_threshold=3,
     ):
@@ -369,9 +369,10 @@ async def _run_output_handler(self) -> None:
                         f"Unknown response type from worker: {resp_type}")
 
                 if resp.request_id not in self.queues:
-                    logger.warning(
-                        "Request %s may have been aborted, ignore response.",
-                        resp.request_id)
+                    if resp_type != ResponseType.HEARTBEAT:
+                        logger.warning(
+                            "Request %s may have been aborted, "
+                            "ignore response.", resp.request_id)
                 elif isinstance(resp, FailureResponse):
                     self.queues[resp.request_id].put_nowait(
                         RuntimeError(f"Request error: {resp.error_message}"))
@@ -429,6 +430,10 @@ async def do_log_stats(
         pass
 
     async def check_health(self, server_type: ServerType, id: int):
+        # lazy initialization
+        if self.output_handler is None:
+            self.output_handler = asyncio.create_task(
+                self._run_output_handler())
         request_id = str(uuid.uuid4())
         request = HeartbeatRequest(request_id=request_id)
         q = asyncio.Queue()