Skip to content

Commit

Permalink
[ BugFix ] Move zmq frontend to IPC instead of TCP (vllm-project#7222)
Browse files Browse the repository at this point in the history
  • Loading branch information
robertgshaw2-neuralmagic authored and sfc-gh-mkeralapura committed Aug 12, 2024
1 parent 21559ad commit ab69234
Show file tree
Hide file tree
Showing 5 changed files with 29 additions and 22 deletions.
12 changes: 8 additions & 4 deletions vllm/entrypoints/openai/api_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@
OpenAIServingTokenization)
from vllm.logger import init_logger
from vllm.usage.usage_lib import UsageContext
from vllm.utils import FlexibleArgumentParser, get_open_port
from vllm.utils import FlexibleArgumentParser, get_open_zmq_ipc_path
from vllm.version import __version__ as VLLM_VERSION

TIMEOUT_KEEP_ALIVE = 5 # seconds
Expand Down Expand Up @@ -106,16 +106,20 @@ async def build_async_engine_client(args) -> AsyncIterator[AsyncEngineClient]:

# Otherwise, use the multiprocessing AsyncLLMEngine.
else:
# Select random path for IPC.
rpc_path = get_open_zmq_ipc_path()
logger.info("Multiprocessing frontend to use %s for RPC Path.",
rpc_path)

# Start RPCServer in separate process (holds the AsyncLLMEngine).
port = get_open_port(envs.VLLM_RPC_PORT)
rpc_server_process = Process(target=run_rpc_server,
args=(engine_args,
UsageContext.OPENAI_API_SERVER,
port))
rpc_path))
rpc_server_process.start()

# Build RPCClient, which conforms to AsyncEngineClient Protocol.
async_engine_client = AsyncEngineRPCClient(port)
async_engine_client = AsyncEngineRPCClient(rpc_path)
await async_engine_client.setup()

try:
Expand Down
6 changes: 3 additions & 3 deletions vllm/entrypoints/openai/rpc/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,9 @@

class AsyncEngineRPCClient:

def __init__(self, port: int):
def __init__(self, rpc_path: str):
self.context = zmq.asyncio.Context()
self.path = f"tcp://localhost:{port}"
self.rpc_path = rpc_path

async def setup(self):
"""Setup the client before it starts sending server requests."""
Expand Down Expand Up @@ -58,7 +58,7 @@ def socket(self):
# to enable streaming.
socket = self.context.socket(zmq.constants.DEALER)
try:
socket.connect(self.path)
socket.connect(self.rpc_path)
yield socket
finally:
socket.close()
Expand Down
10 changes: 4 additions & 6 deletions vllm/entrypoints/openai/rpc/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
class AsyncEngineRPCServer:

def __init__(self, async_engine_args: AsyncEngineArgs,
usage_context: UsageContext, port: int):
usage_context: UsageContext, rpc_path: str):
# Initialize engine first.
self.engine = AsyncLLMEngine.from_engine_args(async_engine_args,
usage_context)
Expand All @@ -30,9 +30,7 @@ def __init__(self, async_engine_args: AsyncEngineArgs,

# Init socket for readiness state.
self.socket = self.context.socket(zmq.constants.ROUTER)
# Note numeric form of localhost should be used for zmq bind(),
# see https://stackoverflow.com/a/8958414
self.socket.bind(f"tcp://127.0.0.1:{port}")
self.socket.bind(rpc_path)

def cleanup(self):
"""Cleanup all resources."""
Expand Down Expand Up @@ -213,6 +211,6 @@ def signal_handler() -> None:


def run_rpc_server(async_engine_args: AsyncEngineArgs,
usage_context: UsageContext, port: int):
server = AsyncEngineRPCServer(async_engine_args, usage_context, port)
usage_context: UsageContext, rpc_path: str):
server = AsyncEngineRPCServer(async_engine_args, usage_context, rpc_path)
asyncio.run(run_server(server))
11 changes: 6 additions & 5 deletions vllm/envs.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import os
import tempfile
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional

if TYPE_CHECKING:
VLLM_HOST_IP: str = ""
VLLM_PORT: Optional[int] = None
VLLM_RPC_PORT: int = 5570
VLLM_RPC_BASE_PATH: str = tempfile.gettempdir()
VLLM_USE_MODELSCOPE: bool = False
VLLM_RINGBUFFER_WARNING_INTERVAL: int = 60
VLLM_INSTANCE_ID: Optional[str] = None
Expand Down Expand Up @@ -142,10 +143,10 @@ def get_default_config_root():
lambda: int(os.getenv('VLLM_PORT', '0'))
if 'VLLM_PORT' in os.environ else None,

# used when the frontend api server is running in multi-processing mode,
# to communicate with the backend engine process over ZMQ.
'VLLM_RPC_PORT':
lambda: int(os.getenv('VLLM_RPC_PORT', '5570')),
# path used for ipc when the frontend api server is running in
# multi-processing mode to communicate with the backend engine process.
'VLLM_RPC_BASE_PATH':
lambda: os.getenv('VLLM_RPC_BASE_PATH', tempfile.gettempdir()),

# If true, will load models from ModelScope instead of Hugging Face Hub.
# note that the value is true or false, not numbers
Expand Down
12 changes: 8 additions & 4 deletions vllm/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from typing import (Any, AsyncGenerator, Awaitable, Callable, Dict, Generic,
Hashable, List, Optional, OrderedDict, Set, Tuple, TypeVar,
Union, overload)
from uuid import uuid4

import numpy as np
import numpy.typing as npt
Expand Down Expand Up @@ -484,10 +485,13 @@ def get_distributed_init_method(ip: str, port: int) -> str:
return f"tcp://[{ip}]:{port}" if ":" in ip else f"tcp://{ip}:{port}"


def get_open_port(port: Optional[int] = None) -> int:
if port is None:
# Default behavior here is to return a port for multi-gpu communication
port = envs.VLLM_PORT
def get_open_zmq_ipc_path() -> str:
base_rpc_path = envs.VLLM_RPC_BASE_PATH
return f"ipc://{base_rpc_path}/{uuid4()}"


def get_open_port() -> int:
port = envs.VLLM_PORT
if port is not None:
while True:
try:
Expand Down

0 comments on commit ab69234

Please sign in to comment.