format

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
 
@@ -37,8 +38,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
 
@@ -216,29 +218,19 @@ 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)
 
-    # if args.served_model_name is not None:
-    #     served_model_names = args.served_model_name
-    # else:
-    #     served_model_names = [args.model]
-
-    served_model_names = "meta-llama/Meta-Llama-3-8B-Instruct"
-
-    from vllm.grpc.client import RPCClient
-    engine = RPCClient()
-
-    # 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()
+    if args.served_model_name is not None:
+        served_model_names = args.served_model_name
+    else:
+        served_model_names = [args.model]
+
+    print("HERE")
+    rpc_client = RPCClient()
+    model_config = await rpc_client.get_model_config()
+    print("HERE2")
 
     if args.disable_log_requests:
         request_logger = None
@@ -309,13 +301,17 @@ 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)
-
+
+    logger.info("Starting RPC Server")
+    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,
     )
 
@@ -332,10 +328,12 @@ 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__":
     # NOTE(simon):

vllm/entrypoints/openai/rpc/__init__.py

@@ -0,0 +1,25 @@
+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

@@ -1,62 +1,31 @@
-from vllm import AsyncLLMEngine
 from typing import AsyncIterator, Optional, Mapping
 
+from vllm.config import ModelConfig
 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 transformers import AutoTokenizer
-from dataclasses import dataclass
+from vllm.entrypoints.openai.rpc import (
+    VLLM_GENERATE_RPC_PATH, VLLM_GET_DATA_RPC_PATH, GenerateRequest, GetDataRequest)
 
 import zmq
 import zmq.asyncio
 import pickle
 
-MODEL = "meta-llama/Meta-Llama-3-8B-Instruct"
-ADDRESS = "ipc:///tmp/zmqtest"
 
-@dataclass
-class RCPRequest:
-    inputs: PromptInputs
-    sampling_params: SamplingParams
-    request_id: str
-
-
-class RPCClient(AsyncLLMEngine):
+class RPCClient:
     def __init__(self):
-        self.engine_use_ray = False
-        self.worker_use_ray = False
-        self.log_requests = False
-        self.engine = None
-
-        self.tokenizer = AutoTokenizer.from_pretrained(MODEL)
-
         self.context = zmq.asyncio.Context()
-
-
-    @property
-    def is_running(self) -> bool:
-        return True
+        self.is_ready_socket = self.context.socket(zmq.REP)
+        self.get_data_socket = self.context.socket(zmq.REQ)
+        self.get_data_socket.connect(VLLM_GET_DATA_RPC_PATH)
 
-    @property
-    def is_stopped(self) -> bool:
-        return False
 
-    @property
-    def errored(self) -> bool:
-        return False
-
-    async def get_tokenizer(
-        self,
-        lora_request: Optional[LoRARequest] = None,
-    ) -> "PreTrainedTokenizer":
-        # TODO: what to return :/
-        return self.tokenizer
-
-    def start_background_loop(self):
-        # TODO something lol
-        pass
+    async def get_model_config(self) -> ModelConfig:
+        self.get_data_socket.send(pickle.dumps(GetDataRequest.MODEL_CONFIG))
+        return pickle.loads(await self.get_data_socket.recv())
+
 
     async def generate(
         self,
@@ -67,19 +36,28 @@ async def generate(
         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(ADDRESS)
+        socket.connect(VLLM_GENERATE_RPC_PATH)
 
+        # Send GenerateRequest to the RPC Server.
         await socket.send_multipart([
             pickle.dumps(
-                RCPRequest(
+                GenerateRequest(
                     inputs=inputs,
                     sampling_params=sampling_params,
-                    request_id=request_id
+                    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)

vllm/entrypoints/openai/rpc/server.py

@@ -1,21 +1,52 @@
-from vllm import AsyncEngineArgs, AsyncLLMEngine
 import asyncio
 import pickle
 import zmq
 import zmq.asyncio
 
-from .client import MODEL, ADDRESS
+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)
 
 class RPCServer:
-    def __init__(self):
+    def __init__(self, async_engine_args):        
+        # Initialize engine first.
+        self.engine = AsyncLLMEngine.from_engine_args(
+            async_engine_args, UsageContext.OPENAI_API_SERVER)
+
+        # Initialize context.
         self.context = zmq.asyncio.Context()
-        self.socket = self.context.socket(zmq.ROUTER)
-        self.socket.bind(ADDRESS)
+
+        # Init socket for readiness state.
+        self.is_ready_socket = self.context.socket(zmq.REP)
+        self.is_ready_socket.bind(VLLM_IS_READY_RPC_PATH)
+
+        # Init socket for generation.
+        self.generate_socket = self.context.socket(zmq.ROUTER)
+        self.generate_socket.bind(VLLM_GENERATE_RPC_PATH)
+
+        # TODO (robertgshaw2-neuralmagic): 
+        # add socket for generation without streaming
+
+        # Init socket for simple data requests.
+        self.get_data_socket = self.context.socket(zmq.REP)
+        self.get_data_socket.bind(VLLM_GET_DATA_RPC_PATH)
+
+        # Setup polling so we can listen on both sockets.
+        self.poller = zmq.asyncio.Poller()
+        self.poller.register(self.generate_socket, zmq.POLLIN)
+        self.poller.register(self.get_data_socket, zmq.POLLIN)
+
+
+    async def get_data(self, message):
+        request_type = pickle.loads(message)
+        if request_type == GetDataRequest.MODEL_CONFIG:
+            return await self.engine.get_model_config()
+        else:
+            raise ValueError(f"Unknown request type: {request_type}")
 
-        self.running_tasks = set()
-        self.engine = AsyncLLMEngine.from_engine_args(
-            AsyncEngineArgs(model=MODEL,
-                            enable_chunked_prefill=True))
 
     async def generate(self, identity, message):
         request = pickle.loads(message)
@@ -29,23 +60,40 @@ async def generate(self, identity, message):
                 identity, 
                 pickle.dumps(request_output, pickle.HIGHEST_PROTOCOL)
             ])
-        
+
     async def run_loop(self):
+        # Notify the RPC client that we are ready to recieve requests.
+        await self.is_ready_socket.send_string("Ready!")
+        self.is_ready_socket.close()
+
+        # Avoid GC of running tasks.
+        running_tasks = set()
         while True:
-            identity, message = await self.socket.recv_multipart()
+            try:
+                socks = dict(await self.poller.poll())
+            except KeyboardInterrupt:
+                # TODO: should there be some other exception here?
+                break
 
-            # Process the request in the background.
-            task = asyncio.create_task(self.generate(identity=identity,
-                                                     message=message))
+            task = None
+            if self.generate_socket in socks:
+                identity, message = await self.generate_socket.recv_multipart()
+                task = asyncio.create_task(self.generate(identity, message)) 
+
+            elif self.get_data_socket in socks:
+                message = await self.get_data_socket.recv()
+                task = asyncio.create_task(self.get_data(message))
 
             # We need to keep around a strong reference to the task, 
             # to avoid the task disappearing mid-execution as running tasks
             # can be GC'ed. Below is a common "fire-and-forget" tasks
             # https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
-            self.running_tasks.add(task)
-            task.add_done_callback(self.running_tasks.discard)
+            if task is not None:
+                running_tasks.add(task)
+                task.add_done_callback(running_tasks.discard)
 
+        # TODO: Do I need to close the generate / get_data sockets?
 
-if __name__ == "__main__":
-    server = RPCServer()
+def run_rpc_server(async_engine_args):
+    server = RPCServer(async_engine_args=async_engine_args)
     asyncio.run(server.run_loop())

vllm/entrypoints/openai/serving_completion.py

@@ -45,7 +45,7 @@ class OpenAIServingCompletion(OpenAIServing):
     def __init__(
         self,
         engine: AsyncLLMEngine,
-        # model_config: ModelConfig,
+        model_config: ModelConfig,
         served_model_names: List[str],
         *,
         lora_modules: Optional[List[LoRAModulePath]],
@@ -54,7 +54,7 @@ def __init__(
         return_tokens_as_token_ids: bool = False,
     ):
         super().__init__(engine=engine,
-                        #  model_config=model_config,
+                         model_config=model_config,
                          served_model_names=served_model_names,
                          lora_modules=lora_modules,
                          prompt_adapters=prompt_adapters,