Skip to content

vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils

Shared constants, lazy imports and helpers for the NIXL connector.

Functions:

  • zmq_ctx –

    Context manager for a ZMQ socket

zmq_ctx(socket_type, addr)

Context manager for a ZMQ socket

Source code in vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
@contextlib.contextmanager
def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
    """Context manager for a ZMQ socket"""

    if socket_type not in (zmq.ROUTER, zmq.REQ):
        raise ValueError(f"Unexpected socket type: {socket_type}")

    ctx: zmq.Context | None = None
    try:
        ctx = zmq.Context()  # type: ignore[attr-defined]
        yield make_zmq_socket(
            ctx=ctx, path=addr, socket_type=socket_type, bind=socket_type == zmq.ROUTER
        )
    finally:
        if ctx is not None:
            ctx.destroy(linger=0)