[2/N] Elastic EP Milestone 2: Integrating NIXL-EP (#35627)

Signed-off-by: Itay Alroy <ialroy@nvidia.com>
Co-authored-by: Yongji Wu <wuyongji317@gmail.com>
Co-authored-by: Ron Tourgeman <rtourgeman@nvidia.com>

vllm/utils/import_utils.py

@@ -412,6 +412,11 @@ def has_deep_gemm() -> bool:
     return _has_module("deep_gemm")
 
 
+def has_nixl_ep() -> bool:
+    """Whether the optional `nixl_ep` package is available."""
+    return _has_module("nixl_ep")
+
+
 def has_triton_kernels() -> bool:
     """Whether the optional `triton_kernels` package is available."""
     is_available = _has_module("triton_kernels") or _has_module(

vllm/utils/network_utils.py

@@ -288,6 +288,7 @@ def make_zmq_socket(
     bind: bool | None = None,
     identity: bytes | None = None,
     linger: int | None = None,
+    router_handover: bool = False,
 ) -> zmq.Socket | zmq.asyncio.Socket:  # type: ignore[name-defined]
     """Make a ZMQ socket with the proper bind/connect semantics."""
 
@@ -314,6 +315,10 @@ def make_zmq_socket(
         socket.setsockopt(zmq.SNDHWM, 0)
         socket.setsockopt(zmq.SNDBUF, buf_size)
 
+    if socket_type == zmq.ROUTER and router_handover:
+        # Let a new connection take over an identity left behind by a dead one.
+        socket.setsockopt(zmq.ROUTER_HANDOVER, 1)
+
     if identity is not None:
         socket.setsockopt(zmq.IDENTITY, identity)
 
@@ -344,12 +349,20 @@ def zmq_socket_ctx(
     bind: bool | None = None,
     linger: int = 0,
     identity: bytes | None = None,
+    router_handover: bool = False,
 ) -> Iterator[zmq.Socket]:
     """Context manager for a ZMQ socket"""
 
     ctx = zmq.Context()  # type: ignore[attr-defined]
     try:
-        yield make_zmq_socket(ctx, path, socket_type, bind=bind, identity=identity)
+        yield make_zmq_socket(
+            ctx,
+            path,
+            socket_type,
+            bind=bind,
+            identity=identity,
+            router_handover=router_handover,
+        )
     except KeyboardInterrupt:
         logger.debug("Got Keyboard Interrupt.")