Initialize AsyncLLMEngine bg loop correctly (#943)

Author: Yard1

vllm/engine/async_llm_engine.py

@@ -155,11 +155,15 @@ def __init__(self,
         self.finished_requests: Set[str] = set()
         self.background_loop = None
         if start_engine_loop:
-            self._start_background_loop()
+            self.start_background_loop()
 
-    def _start_background_loop(self) -> None:
+    @property
+    def is_running(self) -> bool:
+        return self.background_loop is not None
+
+    def start_background_loop(self) -> None:
         """Start the background loop."""
-        if self.background_loop is not None:
+        if self.is_running:
             raise RuntimeError("Background loop is already running.")
         self.background_loop = asyncio.get_event_loop().create_task(
             self.run_engine_loop())
@@ -323,7 +327,8 @@ async def get_model_config(self) -> ModelConfig:
 
     @classmethod
     def from_engine_args(cls,
-                         engine_args: AsyncEngineArgs) -> "AsyncLLMEngine":
+                         engine_args: AsyncEngineArgs,
+                         start_engine_loop: bool = False) -> "AsyncLLMEngine":
         """Creates an async LLM engine from the engine arguments."""
         # Create the engine configs.
         engine_configs = engine_args.create_engine_configs()
@@ -338,5 +343,6 @@ def from_engine_args(cls,
                      distributed_init_method,
                      placement_group,
                      log_requests=not engine_args.disable_log_requests,
-                     log_stats=not engine_args.disable_log_stats)
+                     log_stats=not engine_args.disable_log_stats,
+                     start_engine_loop=start_engine_loop)
         return engine

vllm/entrypoints/api_server.py

@@ -30,6 +30,10 @@ async def generate(request: Request) -> Response:
     stream = request_dict.pop("stream", False)
     sampling_params = SamplingParams(**request_dict)
     request_id = random_uuid()
+
+    if not engine.is_running:
+        engine.start_background_loop()
+
     results_generator = engine.generate(prompt, sampling_params, request_id)
 
     # Streaming case
@@ -75,7 +79,8 @@ async def abort_request() -> None:
     args = parser.parse_args()
 
     engine_args = AsyncEngineArgs.from_cli_args(args)
-    engine = AsyncLLMEngine.from_engine_args(engine_args)
+    engine = AsyncLLMEngine.from_engine_args(engine_args,
+                                             start_engine_loop=False)
 
     uvicorn.run(app,
                 host=args.host,

vllm/entrypoints/openai/api_server.py

@@ -191,6 +191,9 @@ async def create_chat_completion(request: ChatCompletionRequest,
     """
     logger.info(f"Received chat completion request: {request}")
 
+    if not engine.is_running:
+        engine.start_background_loop()
+
     error_check_ret = await check_model(request)
     if error_check_ret is not None:
         return error_check_ret
@@ -363,6 +366,9 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
     """
     logger.info(f"Received completion request: {request}")
 
+    if not engine.is_running:
+        engine.start_background_loop()
+
     error_check_ret = await check_model(request)
     if error_check_ret is not None:
         return error_check_ret
@@ -620,7 +626,8 @@ async def fake_stream_generator() -> AsyncGenerator[str, None]:
         served_model = args.model
 
     engine_args = AsyncEngineArgs.from_cli_args(args)
-    engine = AsyncLLMEngine.from_engine_args(engine_args)
+    engine = AsyncLLMEngine.from_engine_args(engine_args,
+                                             start_engine_loop=False)
     engine_model_config = asyncio.run(engine.get_model_config())
     max_model_len = engine_model_config.get_max_model_len()