🥅 handle shutdown and request errors

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>

vllm/entrypoints/openai/api_server.py

@@ -54,6 +54,7 @@
 openai_serving_completion: OpenAIServingCompletion
 openai_serving_embedding: OpenAIServingEmbedding
 openai_serving_tokenization: OpenAIServingTokenization
+rpc_client: RPCClient
 
 logger = init_logger('vllm.entrypoints.openai.api_server')
 
@@ -230,6 +231,7 @@ async def build_server(
         served_model_names = [args.model]
 
     # TODO: figure out a way around passing the token
+    global rpc_client
     rpc_client = RPCClient(tokenizer=AutoTokenizer.from_pretrained(args.model))
     await rpc_client.wait_for_server()
     logger.info("RPC Client connected to RPC server.")
@@ -334,6 +336,8 @@ def signal_handler() -> None:
     except asyncio.CancelledError:
         print("Gracefully stopping http server")
         await server.shutdown()
+        print("Cleaning up ZMQ client context")
+        rpc_client.close()
         rpc_server_process.join()
 
 

vllm/entrypoints/openai/rpc/client.py

@@ -37,6 +37,10 @@ def __init__(self, tokenizer):
     async def wait_for_server(self):
         await self.is_ready_socket.recv()
 
+    def close(self):
+        """Destroy the zmq context and close all sockets"""
+        self.context.destroy()
+
     async def get_model_config(self) -> ModelConfig:
         self.get_data_socket.send(pickle.dumps(GetDataRequest.MODEL_CONFIG))
         model_config = await self.get_data_socket.recv()
@@ -90,9 +94,13 @@ async def generate(
             message = await socket.recv()
             request_output = pickle.loads(message)
 
+            if isinstance(request_output, Exception):
+                socket.close()
+                raise request_output
+
             if request_output.finished:
                 break
             yield request_output
 
-        socket.close()
         yield request_output
+        socket.close()

vllm/entrypoints/openai/rpc/server.py

@@ -2,13 +2,17 @@
 import pickle
 import zmq
 import zmq.asyncio
+import signal
 
 from vllm import AsyncLLMEngine
 from vllm.usage.usage_lib import UsageContext
 from vllm.entrypoints.openai.rpc import (VLLM_GENERATE_RPC_PATH,
                                          VLLM_GET_DATA_RPC_PATH,
                                          VLLM_IS_READY_RPC_PATH,
                                          GetDataRequest)
+from vllm.logger import init_logger
+
+logger = init_logger('vllm.entrypoints.openai.rpc.server')
 
 
 class RPCServer:
@@ -40,6 +44,13 @@ def __init__(self, async_engine_args):
         self.poller.register(self.generate_socket, zmq.POLLIN)
         self.poller.register(self.get_data_socket, zmq.POLLIN)
 
+    def cleanup(self):
+        """Shuts down the zmq context and closes all sockets"""
+        self.context.destroy()
+        del self.get_data_socket
+        del self.generate_socket
+        del self.is_ready_socket
+
     async def get_data(self, message):
         request_type = pickle.loads(message)
 
@@ -52,18 +63,26 @@ async def get_data(self, message):
             [pickle.dumps(data, pickle.HIGHEST_PROTOCOL)])
 
     async def generate(self, identity, message):
-        request = pickle.loads(message)
-
-        results_generator = self.engine.generate(
-            request.inputs,
-            sampling_params=request.sampling_params,
-            request_id=request.request_id)
-
-        async for request_output in results_generator:
+        try:
+            request = pickle.loads(message)
+
+            results_generator = self.engine.generate(
+                request.inputs,
+                sampling_params=request.sampling_params,
+                request_id=request.request_id)
+
+            async for request_output in results_generator:
+                self.generate_socket.send_multipart([
+                    identity,
+                    pickle.dumps(request_output, pickle.HIGHEST_PROTOCOL)
+                ])
+        except Exception as e:
+            ### Notify client of all failures
             self.generate_socket.send_multipart([
-                identity,
-                pickle.dumps(request_output, pickle.HIGHEST_PROTOCOL)
-            ])
+                    identity,
+                    pickle.dumps(e, pickle.HIGHEST_PROTOCOL)
+                ])
+
 
     async def run_loop(self):
         # Notify the RPC client that we are ready to recieve requests.
@@ -73,11 +92,8 @@ async def run_loop(self):
         # Avoid GC of running tasks.
         running_tasks = set()
         while True:
-            try:
-                socks = dict(await self.poller.poll())
-            except KeyboardInterrupt:
-                # TODO: should there be some other exception here?
-                break
+            self.poll_future = self.poller.poll()
+            socks = dict(await self.poll_future)
 
             task = None
             if self.generate_socket in socks:
@@ -99,6 +115,30 @@ async def run_loop(self):
         # TODO: Do I need to close the generate / get_data sockets?
 
 
+async def run_server(server: RPCServer):
+    # Run with proper interrupt handling
+    logger.info("Booting up vLLM zmq backend")
+
+    loop = asyncio.get_running_loop()
+
+    server_task = loop.create_task(server.run_loop())
+    def signal_handler() -> None:
+        # Kill the server on interrupt / terminate
+        server_task.cancel()
+
+    loop.add_signal_handler(signal.SIGINT, signal_handler)
+    loop.add_signal_handler(signal.SIGTERM, signal_handler)
+
+    try:
+        await server_task
+    except asyncio.CancelledError:
+        logger.info("ZMQ Backend was interrupted")
+    finally:
+        # Clean up all the zmq resources before exiting
+        server.cleanup()
+    logger.info("vLLM ZMQ Backend shut down")
+
+
 def run_rpc_server(async_engine_args):
     server = RPCServer(async_engine_args=async_engine_args)
-    asyncio.run(server.run_loop())
+    asyncio.run(run_server(server))