[ Frontend ] Multiprocessing for OpenAI Server with zeromq

examples/openai_completion_client.py

@@ -14,14 +14,13 @@
 model = models.data[0].id
 
 # Completion API
-stream = False
+stream = True
 completion = client.completions.create(
     model=model,
     prompt="A robot may not injure a human being",
     echo=False,
-    n=2,
-    stream=stream,
-    logprobs=3)
+    n=1,
+    stream=stream)
 
 print("Completion results:")
 if stream:

vllm/entrypoints/openai/api_server.py

@@ -4,6 +4,7 @@
 import re
 import signal
 from contextlib import asynccontextmanager
+from multiprocessing import Process
 from http import HTTPStatus
 from typing import Optional, Set
 
@@ -16,6 +17,8 @@
 from prometheus_client import make_asgi_app
 from starlette.routing import Mount
 
+from transformers import AutoTokenizer
+
 import vllm.envs as envs
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.engine.async_llm_engine import AsyncLLMEngine
@@ -37,8 +40,9 @@
 from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
 from vllm.entrypoints.openai.serving_tokenization import (
     OpenAIServingTokenization)
+from vllm.entrypoints.openai.rpc.client import RPCClient
+from vllm.entrypoints.openai.rpc.server import run_rpc_server
 from vllm.logger import init_logger
-from vllm.usage.usage_lib import UsageContext
 from vllm.utils import FlexibleArgumentParser
 from vllm.version import __version__ as VLLM_VERSION
 
@@ -64,10 +68,10 @@ async def _force_log():
             await asyncio.sleep(10)
             await engine.do_log_stats()
 
-    if not engine_args.disable_log_stats:
-        task = asyncio.create_task(_force_log())
-        _running_tasks.add(task)
-        task.add_done_callback(_running_tasks.remove)
+    # if not engine_args.disable_log_stats:
+    #     task = asyncio.create_task(_force_log())
+    #     _running_tasks.add(task)
+    #     task.add_done_callback(_running_tasks.remove)
 
     yield
 
@@ -216,7 +220,6 @@ async def authentication(request: Request, call_next):
 
 async def build_server(
     args,
-    llm_engine: Optional[AsyncLLMEngine] = None,
     **uvicorn_kwargs,
 ) -> uvicorn.Server:
     app = build_app(args)
@@ -226,14 +229,11 @@ async def build_server(
     else:
         served_model_names = [args.model]
 
-    global engine, engine_args
-
-    engine_args = AsyncEngineArgs.from_cli_args(args)
-    engine = (llm_engine
-              if llm_engine is not None else AsyncLLMEngine.from_engine_args(
-                  engine_args, usage_context=UsageContext.OPENAI_API_SERVER))
-
-    model_config = await engine.get_model_config()
+    # TODO: figure out a way around passing the token
+    rpc_client = RPCClient(tokenizer=AutoTokenizer.from_pretrained(args.model))
+    await rpc_client.wait_for_server()
+    logger.info("RPC Client connected to RPC server.")
+    model_config = await rpc_client.get_model_config()
 
     if args.disable_log_requests:
         request_logger = None
@@ -245,40 +245,40 @@ async def build_server(
     global openai_serving_embedding
     global openai_serving_tokenization
 
-    openai_serving_chat = OpenAIServingChat(
-        engine,
-        model_config,
-        served_model_names,
-        args.response_role,
-        lora_modules=args.lora_modules,
-        prompt_adapters=args.prompt_adapters,
-        request_logger=request_logger,
-        chat_template=args.chat_template,
-        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
-    )
+    # openai_serving_chat = OpenAIServingChat(
+    #     engine,
+    #     model_config,
+    #     served_model_names,
+    #     args.response_role,
+    #     lora_modules=args.lora_modules,
+    #     prompt_adapters=args.prompt_adapters,
+    #     request_logger=request_logger,
+    #     chat_template=args.chat_template,
+    #     return_tokens_as_token_ids=args.return_tokens_as_token_ids,
+    # )
     openai_serving_completion = OpenAIServingCompletion(
-        engine,
+        rpc_client,
         model_config,
         served_model_names,
         lora_modules=args.lora_modules,
         prompt_adapters=args.prompt_adapters,
         request_logger=request_logger,
         return_tokens_as_token_ids=args.return_tokens_as_token_ids,
     )
-    openai_serving_embedding = OpenAIServingEmbedding(
-        engine,
-        model_config,
-        served_model_names,
-        request_logger=request_logger,
-    )
-    openai_serving_tokenization = OpenAIServingTokenization(
-        engine,
-        model_config,
-        served_model_names,
-        lora_modules=args.lora_modules,
-        request_logger=request_logger,
-        chat_template=args.chat_template,
-    )
+    # openai_serving_embedding = OpenAIServingEmbedding(
+    #     engine,
+    #     model_config,
+    #     served_model_names,
+    #     request_logger=request_logger,
+    # )
+    # openai_serving_tokenization = OpenAIServingTokenization(
+    #     engine,
+    #     model_config,
+    #     served_model_names,
+    #     lora_modules=args.lora_modules,
+    #     request_logger=request_logger,
+    #     chat_template=args.chat_template,
+    # )
     app.root_path = args.root_path
 
     logger.info("Available routes are:")
@@ -304,13 +304,16 @@ async def build_server(
     return uvicorn.Server(config)
 
 
-async def run_server(args, llm_engine=None, **uvicorn_kwargs) -> None:
+async def run_server(args, **uvicorn_kwargs) -> None:
     logger.info("vLLM API server version %s", VLLM_VERSION)
     logger.info("args: %s", args)
 
+    rpc_server_process = Process(target=run_rpc_server,
+                                 args=(AsyncEngineArgs.from_cli_args(args), ))
+    rpc_server_process.start()
+
     server = await build_server(
         args,
-        llm_engine,
         **uvicorn_kwargs,
     )
 
@@ -327,9 +330,11 @@ def signal_handler() -> None:
 
     try:
         await server_task
+        rpc_server_process.join()
     except asyncio.CancelledError:
         print("Gracefully stopping http server")
         await server.shutdown()
+        rpc_server_process.join()
 
 
 if __name__ == "__main__":

vllm/entrypoints/openai/rpc/__init__.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+from typing import Optional, Mapping
+from enum import Enum
+
+from vllm.inputs import PromptInputs
+from vllm.lora.request import LoRARequest
+from vllm.prompt_adapter.request import PromptAdapterRequest
+from vllm.sampling_params import SamplingParams
+
+VLLM_GENERATE_RPC_PATH = "tcp://localhost:5570"
+VLLM_GET_DATA_RPC_PATH = "tcp://localhost:5571"
+VLLM_IS_READY_RPC_PATH = "tcp://localhost:5572"
+
+
+@dataclass
+class GenerateRequest:
+    inputs: PromptInputs
+    sampling_params: SamplingParams
+    request_id: str
+    lora_request: Optional[LoRARequest] = None
+    trace_headers: Optional[Mapping[str, str]] = None
+    prompt_adapter_request: Optional[PromptAdapterRequest] = None
+
+
+class GetDataRequest(Enum):
+    MODEL_CONFIG = 1

vllm/entrypoints/openai/rpc/client.py

@@ -0,0 +1,98 @@
+from typing import AsyncIterator, Optional, Mapping
+
+from vllm.config import ModelConfig, DecodingConfig
+from vllm.inputs import PromptInputs
+from vllm.lora.request import LoRARequest
+from vllm.outputs import RequestOutput
+from vllm.prompt_adapter.request import PromptAdapterRequest
+from vllm.sampling_params import SamplingParams
+from vllm.entrypoints.openai.rpc import (VLLM_GENERATE_RPC_PATH,
+                                         VLLM_GET_DATA_RPC_PATH,
+                                         VLLM_IS_READY_RPC_PATH,
+                                         GenerateRequest, GetDataRequest)
+
+import zmq
+import zmq.asyncio
+import pickle
+
+
+class RPCClient:
+
+    # TODO: check if opening all these sockets is an antipattern?
+    def __init__(self, tokenizer):
+        self.context = zmq.asyncio.Context()
+
+        # TODO: do the tokenizer properly.
+        self.tokenizer = tokenizer
+        self.decoding_config = DecodingConfig()
+
+        # Socket to check if the RPC server is ready.
+        self.is_ready_socket = self.context.socket(zmq.REP)
+        self.is_ready_socket.connect(VLLM_IS_READY_RPC_PATH)
+
+        # Socket to query data (e.g. get_model_config)
+        self.get_data_socket = self.context.socket(zmq.REQ)
+        self.get_data_socket.connect(VLLM_GET_DATA_RPC_PATH)
+
+    async def wait_for_server(self):
+        await self.is_ready_socket.recv()
+
+    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()
+        return pickle.loads(model_config)
+
+    async def get_tokenizer(self, lora_request: LoRARequest):
+        # TODO: handle this via get data? - or avoid doing via RPC
+        return self.tokenizer
+
+    async def get_decoding_config(self):
+        # TODO: handle this via get data? -  or avoid doing via RPC
+        return self.decoding_config
+
+    async def abort(self, request_id: str):
+        # TODO: actually handle this with a new socket.
+        pass
+
+    async def is_tracing_enabled(self):
+        return False
+
+    async def generate(
+        self,
+        inputs: PromptInputs,
+        sampling_params: SamplingParams,
+        request_id: str,
+        lora_request: Optional[LoRARequest] = None,
+        trace_headers: Optional[Mapping[str, str]] = None,
+        prompt_adapter_request: Optional[PromptAdapterRequest] = None
+    ) -> AsyncIterator[RequestOutput]:
+
+        # Connect to RPC socket for Request-Reply pattern,
+        # Note that we use DEALER to enable asynchronous communication
+        # to enable streaming.
+        socket = self.context.socket(zmq.DEALER)
+        socket.connect(VLLM_GENERATE_RPC_PATH)
+
+        # Send GenerateRequest to the RPC Server.
+        await socket.send_multipart([
+            pickle.dumps(
+                GenerateRequest(inputs=inputs,
+                                sampling_params=sampling_params,
+                                request_id=request_id,
+                                lora_request=lora_request,
+                                trace_headers=trace_headers,
+                                prompt_adapter_request=prompt_adapter_request),
+                pickle.HIGHEST_PROTOCOL)
+        ])
+
+        # Stream back the results from the RPC Server.
+        while True:
+            message = await socket.recv()
+            request_output = pickle.loads(message)
+
+            if request_output.finished:
+                break
+            yield request_output
+
+        socket.close()
+        yield request_output