36b8f80266
Signed-off-by: wuhang <wuhang6@huawei.com> Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: wuhang <whlbx@hotmail.com> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""ZMQ-based queue utilities for Omni IPC."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import queue
|
|
from typing import Any
|
|
|
|
import zmq
|
|
from vllm.utils.network_utils import make_zmq_socket
|
|
|
|
|
|
class ZmqQueue:
|
|
"""Queue-like wrapper on a ZMQ socket."""
|
|
|
|
def __init__(
|
|
self,
|
|
ctx: zmq.Context,
|
|
socket_type: int,
|
|
*,
|
|
bind: str | None = None,
|
|
connect: str | None = None,
|
|
recv_timeout_ms: int | None = None,
|
|
send_timeout_ms: int | None = None,
|
|
) -> None:
|
|
# Determine path and bind mode
|
|
path = bind if bind is not None else connect
|
|
if path is None:
|
|
raise ValueError("Either bind or connect must be specified")
|
|
bind_mode = bind is not None
|
|
|
|
self._socket = make_zmq_socket(ctx, path, socket_type, bind=bind_mode, linger=5000)
|
|
|
|
# Reusable poller for efficient polling operations
|
|
self._poller = zmq.Poller()
|
|
self._poller.register(self._socket, zmq.POLLIN)
|
|
|
|
# Store default timeout settings
|
|
self._default_recv_timeout = recv_timeout_ms
|
|
self._default_send_timeout = send_timeout_ms
|
|
|
|
# Apply timeout settings if specified
|
|
if recv_timeout_ms is not None:
|
|
self._socket.rcvtimeo = recv_timeout_ms
|
|
if send_timeout_ms is not None:
|
|
self._socket.sndtimeo = send_timeout_ms
|
|
|
|
self.endpoint = path
|
|
|
|
def put(self, obj: Any) -> None:
|
|
"""Send an object to the queue. Blocks until sent or timeout."""
|
|
try:
|
|
self._socket.send_pyobj(obj)
|
|
except zmq.Again as e:
|
|
raise queue.Full() from e
|
|
|
|
def put_nowait(self, obj: Any) -> None:
|
|
"""Send an object to the queue without blocking."""
|
|
try:
|
|
self._socket.send_pyobj(obj, flags=zmq.NOBLOCK)
|
|
except zmq.Again as e:
|
|
raise queue.Full() from e
|
|
|
|
def get(self, timeout: float | None = None) -> Any:
|
|
"""Receive an object from the queue with optional timeout in seconds."""
|
|
if timeout is None:
|
|
return self._socket.recv_pyobj()
|
|
|
|
# Use the reusable poller for timeout handling
|
|
events = dict(self._poller.poll(int(timeout * 1000)))
|
|
if events.get(self._socket) == zmq.POLLIN:
|
|
return self._socket.recv_pyobj()
|
|
raise queue.Empty()
|
|
|
|
def get_nowait(self) -> Any:
|
|
"""Receive an object from the queue without blocking."""
|
|
try:
|
|
return self._socket.recv_pyobj(flags=zmq.NOBLOCK)
|
|
except zmq.Again as e:
|
|
raise queue.Empty() from e
|
|
|
|
def empty(self) -> bool:
|
|
"""Check if the queue is empty without blocking."""
|
|
events = dict(self._poller.poll(0))
|
|
return events.get(self._socket) != zmq.POLLIN
|
|
|
|
def close(self) -> None:
|
|
self._socket.close(0)
|
|
|
|
|
|
def create_zmq_queue(ctx: zmq.Context, endpoint: str, socket_type: int) -> ZmqQueue:
|
|
"""Create a ZmqQueue from an endpoint string and socket type."""
|
|
return ZmqQueue(ctx, socket_type, connect=endpoint)
|