@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 651cf0fec1

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 16:00 UTC

[V1] DP scale-out (1/N): Use zmq ROUTER/DEALER sockets for input queue (#15906)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Apr 4, 2025 at 12:56 UTC 651cf0fec19ed1da44ad266066f86b74e5246006

4 files changed +113 -69

vllm/utils.py

+25 -10

@@ -2189,6 +2189,8 @@ def make_zmq_socket(

2189 ctx: Union[zmq.asyncio.Context, zmq.Context], # type: ignore[name-defined]

2190 path: str,

2191 socket_type: Any,

2192 + bind: Optional[bool] = None,

2193 + identity: Optional[bytes] = None,

2194 ) -> Union[zmq.Socket, zmq.asyncio.Socket]: # type: ignore[name-defined]

2195 """Make a ZMQ socket with the proper bind/connect semantics."""

2196

@@ -2207,16 +2209,24 @@ def make_zmq_socket(

2209 else:

2210 buf_size = -1 # Use system default buffer size

2211

2210 - if socket_type == zmq.constants.PULL:

2211 - socket.setsockopt(zmq.constants.RCVHWM, 0)

2212 - socket.setsockopt(zmq.constants.RCVBUF, buf_size)

2212 + if bind is None:

2213 + bind = socket_type != zmq.PUSH

2214 +

2215 + if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER):

2216 + socket.setsockopt(zmq.RCVHWM, 0)

2217 + socket.setsockopt(zmq.RCVBUF, buf_size)

2218 +

2219 + if socket_type in (zmq.PUSH, zmq.DEALER, zmq.ROUTER):

2220 + socket.setsockopt(zmq.SNDHWM, 0)

2221 + socket.setsockopt(zmq.SNDBUF, buf_size)

2222 +

2223 + if identity is not None:

2224 + socket.setsockopt(zmq.IDENTITY, identity)

2225 +

2226 + if bind:

2227 socket.bind(path)

2214 - elif socket_type == zmq.constants.PUSH:

2215 - socket.setsockopt(zmq.constants.SNDHWM, 0)

2216 - socket.setsockopt(zmq.constants.SNDBUF, buf_size)

2217 - socket.connect(path)

2228 else:

2219 - raise ValueError(f"Unknown Socket Type: {socket_type}")

2229 + socket.connect(path)

2230

2231 return socket

2232

@@ -2225,14 +2235,19 @@ def make_zmq_socket(

2235 def zmq_socket_ctx(

2236 path: str,

2237 socket_type: Any,

2238 + bind: Optional[bool] = None,

2239 linger: int = 0,

2240 + identity: Optional[bytes] = None,

2241 ) -> Iterator[zmq.Socket]:

2242 """Context manager for a ZMQ socket"""

2243

2244 ctx = zmq.Context() # type: ignore[attr-defined]

2245 try:

2234 - yield make_zmq_socket(ctx, path, socket_type)

2235 -

2246 + yield make_zmq_socket(ctx,

2247 + path,

2248 + socket_type,

2249 + bind=bind,

2250 + identity=identity)

2251 except KeyboardInterrupt:

2252 logger.debug("Got Keyboard Interrupt.")

2253

vllm/v1/engine/core.py

+16 -12

@@ -318,6 +318,11 @@ class EngineCoreProc(EngineCore):

318 ):

319 super().__init__(vllm_config, executor_class, log_stats)

320

321 + self.step_fn = (self.step if self.batch_queue is None else

322 + self.step_with_batch_queue)

323 +

324 + self.global_unfinished_reqs = False

325 +

326 # Background Threads and Queues for IO. These enable us to

327 # overlap ZMQ socket IO with GPU since they release the GIL,

328 # and to overlap some serialization/deserialization with the

@@ -327,22 +332,16 @@ class EngineCoreProc(EngineCore):

332 Any]] = queue.Queue()

333 self.output_queue: queue.Queue[EngineCoreOutputs] = queue.Queue()

334 threading.Thread(target=self.process_input_socket,

330 - args=(input_path, ),

335 + args=(input_path, engine_index),

336 daemon=True).start()

337 threading.Thread(target=self.process_output_socket,

338 args=(output_path, engine_index),

339 daemon=True).start()

340

336 - self.global_unfinished_reqs = False

337 -

338 - self.step_fn = (self.step if self.batch_queue is None else

339 - self.step_with_batch_queue)

340 -

341 @staticmethod

342 def run_engine_core(*args,

343 dp_rank: int = 0,

344 local_dp_rank: int = 0,

345 - ready_pipe,

345 **kwargs):

346 """Launch EngineCore busy loop in background process."""

347

@@ -377,9 +376,6 @@ class EngineCoreProc(EngineCore):

376 else:

377 engine_core = EngineCoreProc(*args, **kwargs)

378

380 - # Send Readiness signal to EngineClient.

381 - ready_pipe.send({"status": "READY"})

382 -

379 engine_core.run_busy_loop()

380

381 except SystemExit:

@@ -476,14 +472,22 @@ class EngineCoreProc(EngineCore):

472 and not isinstance(v, p.annotation) else v

473 for v, p in zip(args, arg_types))

474

479 - def process_input_socket(self, input_path: str):

475 + def process_input_socket(self, input_path: str, engine_index: int):

476 """Input socket IO thread."""

477

478 # Msgpack serialization decoding.

479 add_request_decoder = MsgpackDecoder(EngineCoreRequest)

480 generic_decoder = MsgpackDecoder()

481 + identity = engine_index.to_bytes(length=2, byteorder="little")

482 +

483 + with zmq_socket_ctx(input_path,

484 + zmq.DEALER,

485 + identity=identity,

486 + bind=False) as socket:

487 +

488 + # Send ready message to front-end once input socket is connected.

489 + socket.send(b'READY')

490

486 - with zmq_socket_ctx(input_path, zmq.constants.PULL) as socket:

491 while True:

492 # (RequestType, RequestData)

493 type_frame, data_frame = socket.recv_multipart(copy=False)

vllm/v1/engine/core_client.py

+71 -37

@@ -8,7 +8,7 @@ import threading

8 import uuid

9 import weakref

10 from abc import ABC, abstractmethod

11 -from collections.abc import Awaitable, Sequence

11 +from collections.abc import Awaitable

12 from concurrent.futures import Future

13 from dataclasses import dataclass, field

14 from threading import Thread

@@ -35,6 +35,8 @@ AnyFuture = Union[asyncio.Future[Any], Future[Any]]

36 _R = TypeVar('_R') # Return type for collective_rpc

38 +STARTUP_POLL_PERIOD_MS = 10000

39 +

41 class EngineCoreClient(ABC):

42 """

@@ -261,15 +263,13 @@ class CoreEngine:

263 vllm_config: VllmConfig,

264 executor_class: type[Executor],

265 log_stats: bool,

264 - ctx: Union[zmq.Context, zmq.asyncio.Context],

266 + input_path: str,

267 output_path: str,

268 index: int = 0,

269 local_dp_rank: int = 0,

270 ):

269 - # Paths and sockets for IPC.

270 - input_path = get_open_zmq_ipc_path()

271 - self.input_socket = make_zmq_socket(ctx, input_path,

272 - zmq.constants.PUSH)

271 + self.index = index

272 + self.identity = index.to_bytes(length=2, byteorder="little")

273 try:

274 # Start EngineCore in background process.

275 self.proc_handle = BackgroundProcHandle(

@@ -291,14 +291,9 @@ class CoreEngine:

291 # Ensure socket is closed if process fails to start.

292 self.close()

293

294 - def send_multipart(self, msg_parts: Sequence):

295 - return self.input_socket.send_multipart(msg_parts, copy=False)

296 -

294 def close(self):

295 if proc_handle := getattr(self, "proc_handle", None):

296 proc_handle.shutdown()

300 - if socket := getattr(self, "input_socket", None):

301 - socket.close(linger=0)

297

298

299 @dataclass

@@ -309,6 +304,7 @@ class BackgroundResources:

304 ctx: Union[zmq.Context]

305 core_engines: list[CoreEngine] = field(default_factory=list)

306 output_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None

307 + input_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None

308 shutdown_path: Optional[str] = None

309

310 def __call__(self):

@@ -321,6 +317,8 @@ class BackgroundResources:

317 # aren't explicitly closed first.

318 if self.output_socket is not None:

319 self.output_socket.close(linger=0)

320 + if self.input_socket is not None:

321 + self.input_socket.close(linger=0)

322 if self.shutdown_path is not None:

323 # We must ensure that the sync output socket is

324 # closed cleanly in its own thread.

@@ -387,21 +385,51 @@ class MPClient(EngineCoreClient):

385

386 # Paths and sockets for IPC.

387 self.output_path = get_open_zmq_ipc_path()

388 + input_path = get_open_zmq_ipc_path()

389 + self.input_socket = make_zmq_socket(self.ctx,

390 + input_path,

391 + zmq.ROUTER,

392 + bind=True)

393 + self.resources.input_socket = self.input_socket

394

395 new_core_engine = lambda index, local_dp_rank=None: CoreEngine(

392 - vllm_config, executor_class, log_stats, self.ctx, self.output_path,

393 - index, local_dp_rank)

396 + vllm_config, executor_class, log_stats, input_path, self.

397 + output_path, index, local_dp_rank)

398

399 # Start engine core process(es).

400 self._init_core_engines(vllm_config, new_core_engine,

401 self.resources.core_engines)

402

403 # Wait for engine core process(es) to start.

400 - for engine in self.resources.core_engines:

401 - engine.proc_handle.wait_for_startup()

404 + self._wait_for_engine_startup()

405

406 self.utility_results: dict[int, AnyFuture] = {}

407

408 + def _wait_for_engine_startup(self):

409 + # Get a sync handle to the socket which can be sync or async.

410 + sync_input_socket = zmq.Socket.shadow(self.input_socket)

411 +

412 + # Wait for engine core process(es) to send ready messages.

413 + identities = set(eng.index for eng in self.resources.core_engines)

414 + while identities:

415 + while not sync_input_socket.poll(timeout=STARTUP_POLL_PERIOD_MS):

416 + logger.info("Waiting for %d core engine proc(s) to start: %s",

417 + len(identities), identities)

418 + eng_id_bytes, msg = sync_input_socket.recv_multipart()

419 + eng_id = int.from_bytes(eng_id_bytes, byteorder="little")

420 + if eng_id not in identities:

421 + raise RuntimeError(f"Unexpected or duplicate engine: {eng_id}")

422 + if msg != b'READY':

423 + raise RuntimeError(f"Engine {eng_id} failed: {msg.decode()}")

424 + logger.info("Core engine process %d ready.", eng_id)

425 + identities.discard(eng_id)

426 +

427 + # Double check that the process are running.

428 + for engine in self.resources.core_engines:

429 + proc = engine.proc_handle.proc

430 + if proc.exitcode is not None:

431 + raise RuntimeError(f"Engine proc {proc.name} not running")

432 +

433 def _init_core_engines(

434 self,

435 vllm_config: VllmConfig,

@@ -494,9 +522,10 @@ class SyncMPClient(MPClient):

522 return self.outputs_queue.get()

523

524 def _send_input(self, request_type: EngineCoreRequestType, request: Any):

497 - # (RequestType, SerializedRequest)

498 - msg = (request_type.value, self.encoder.encode(request))

499 - self.core_engine.send_multipart(msg)

525 + # (Identity, RequestType, SerializedRequest)

526 + msg = (self.core_engine.identity, request_type.value,

527 + self.encoder.encode(request))

528 + self.input_socket.send_multipart(msg, copy=False)

529

530 def call_utility(self, method: str, *args) -> Any:

531 call_id = uuid.uuid1().int >> 64

@@ -625,30 +654,34 @@ class AsyncMPClient(MPClient):

654 assert self.outputs_queue is not None

655 return await self.outputs_queue.get()

656

628 - async def _send_input(self, request_type: EngineCoreRequestType,

629 - request: Any) -> None:

630 - await self.core_engine.send_multipart(

631 - (request_type.value, self.encoder.encode(request)))

657 + def _send_input(self,

658 + request_type: EngineCoreRequestType,

659 + request: Any,

660 + engine: Optional[CoreEngine] = None) -> Awaitable[None]:

661 + if engine is None:

662 + engine = self.core_engine

663

633 - self._ensure_output_queue_task()

664 + message = (request_type.value, self.encoder.encode(request))

665 + return self._send_input_message(message, engine)

666 +

667 + def _send_input_message(self, message: tuple[bytes, bytes],

668 + engine: CoreEngine) -> Awaitable[None]:

669 + message = (engine.identity, ) + message # type: ignore[assignment]

670 + return self.input_socket.send_multipart(message, copy=False)

671

672 async def call_utility_async(self, method: str, *args) -> Any:

673 return await self._call_utility_async(method,

674 *args,

675 engine=self.core_engine)

676

640 - async def _call_utility_async(

641 - self,

642 - method: str,

643 - *args,

644 - engine: CoreEngine,

645 - ) -> Any:

677 + async def _call_utility_async(self, method: str, *args,

678 + engine: CoreEngine) -> Any:

679 call_id = uuid.uuid1().int >> 64

680 future = asyncio.get_running_loop().create_future()

681 self.utility_results[call_id] = future

682 message = (EngineCoreRequestType.UTILITY.value,

683 self.encoder.encode((call_id, method, args)))

651 - await engine.send_multipart(message)

684 + await self._send_input_message(message, engine)

685 self._ensure_output_queue_task()

686 return await future

687

@@ -657,6 +690,7 @@ class AsyncMPClient(MPClient):

690 # tokenized.

691 request.prompt = None

692 await self._send_input(EngineCoreRequestType.ADD, request)

693 + self._ensure_output_queue_task()

694

695 async def abort_requests_async(self, request_ids: list[str]) -> None:

696 if len(request_ids) > 0:

@@ -761,15 +795,15 @@ class DPAsyncMPClient(AsyncMPClient):

795 self.reqs_in_flight[request.request_id] = chosen_engine

796 chosen_engine.num_reqs_in_flight += 1

797 if self.num_engines_running >= len(self.core_engines):

764 - await chosen_engine.send_multipart(msg)

798 + await self._send_input_message(msg, chosen_engine)

799 else:

800 # Send request to chosen engine and dp start loop

801 # control message to all other engines.

802 self.num_engines_running += len(self.core_engines)

803 await asyncio.gather(*[

770 - engine.send_multipart(msg if engine is

771 - chosen_engine else self.start_dp_msg)

772 - for engine in self.core_engines

804 + self._send_input_message(

805 + msg if engine is chosen_engine else self.start_dp_msg,

806 + engine) for engine in self.core_engines

807 ])

808

809 self._ensure_output_queue_task()

@@ -794,7 +828,7 @@ class DPAsyncMPClient(AsyncMPClient):

828 # sure to start the other engines:

829 self.num_engines_running = len(self.core_engines)

830 coros = [

797 - engine.send_multipart(self.start_dp_msg)

831 + self._send_input_message(self.start_dp_msg, engine)

832 for engine in self.core_engines

833 if not engine.num_reqs_in_flight

834 ]

@@ -820,5 +854,5 @@ class DPAsyncMPClient(AsyncMPClient):

854

855 async def _abort_requests(self, request_ids: list[str],

856 engine: CoreEngine) -> None:

823 - await engine.send_multipart((EngineCoreRequestType.ABORT.value,

824 - self.encoder.encode(request_ids)))

857 + await self._send_input(EngineCoreRequestType.ABORT, request_ids,

858 + engine)

vllm/v1/utils.py

+1 -10

@@ -105,12 +105,9 @@ class BackgroundProcHandle:

105 process_kwargs: dict[Any, Any],

106 ):

107 context = get_mp_context()

108 - self.reader, writer = context.Pipe(duplex=False)

108

110 - assert ("ready_pipe" not in process_kwargs

111 - and "input_path" not in process_kwargs

109 + assert ("input_path" not in process_kwargs

110 and "output_path" not in process_kwargs)

113 - process_kwargs["ready_pipe"] = writer

111 process_kwargs["input_path"] = input_path

112 process_kwargs["output_path"] = output_path

113

@@ -122,12 +119,6 @@ class BackgroundProcHandle:

119 input_path, output_path)

120 self.proc.start()

121

125 - def wait_for_startup(self):

126 - # Wait for startup.

127 - if self.reader.recv()["status"] != "READY":

128 - raise RuntimeError(f"{self.proc.name} initialization failed. "

129 - "See root cause above.")

130 -

122 def shutdown(self):

123 self._finalizer()

124