| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | import ipaddress |
| 5 | import os |
| 6 | import socket |
| 7 | import sys |
| 8 | import tempfile |
| 9 | import warnings |
| 10 | from collections.abc import ( |
| 11 | Iterator, |
| 12 | Sequence, |
| 13 | ) |
| 14 | from typing import Any |
| 15 | from uuid import uuid4 |
| 16 | |
| 17 | import psutil |
| 18 | import zmq |
| 19 | import zmq.asyncio |
| 20 | from urllib3.util import parse_url |
| 21 | |
| 22 | import vllm.envs as envs |
| 23 | from vllm.logger import init_logger |
| 24 | |
| 25 | logger = init_logger(__name__) |
| 26 | |
| 27 | |
| 28 | def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]): |
| 29 | for sock in sockets: |
| 30 | if sock is not None: |
| 31 | sock.close(linger=0) |
| 32 | |
| 33 | |
| 34 | def get_ip() -> str: |
| 35 | host_ip = envs.VLLM_HOST_IP |
| 36 | if "HOST_IP" in os.environ and "VLLM_HOST_IP" not in os.environ: |
| 37 | logger.warning( |
| 38 | "The environment variable HOST_IP is deprecated and ignored, as" |
| 39 | " it is often used by Docker and other software to" |
| 40 | " interact with the container's network stack. Please " |
| 41 | "use VLLM_HOST_IP instead to set the IP address for vLLM processes" |
| 42 | " to communicate with each other." |
| 43 | ) |
| 44 | if host_ip: |
| 45 | return host_ip |
| 46 | |
| 47 | # IP is not set, try to get it from the network interface |
| 48 | |
| 49 | # try ipv4 |
| 50 | try: |
| 51 | with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s: |
| 52 | s.connect(("8.8.8.8", 80)) # Doesn't need to be reachable |
| 53 | return s.getsockname()[0] |
| 54 | except Exception: |
| 55 | pass |
| 56 | |
| 57 | # try ipv6 |
| 58 | try: |
| 59 | with socket.socket(socket.AF_INET6, socket.SOCK_DGRAM) as s: |
| 60 | # Google's public DNS server, see |
| 61 | # https://developers.google.com/speed/public-dns/docs/using#addresses |
| 62 | s.connect(("2001:4860:4860::8888", 80)) # Doesn't need to be reachable |
| 63 | return s.getsockname()[0] |
| 64 | except Exception: |
| 65 | pass |
| 66 | |
| 67 | warnings.warn( |
| 68 | "Failed to get the IP address, using 0.0.0.0 by default. " |
| 69 | "The value can be set by the environment variable" |
| 70 | " VLLM_HOST_IP or HOST_IP.", |
| 71 | stacklevel=2, |
| 72 | ) |
| 73 | return "0.0.0.0" |
| 74 | |
| 75 | |
| 76 | def test_loopback_bind(address: str, family: int) -> bool: |
| 77 | try: |
| 78 | s = socket.socket(family, socket.SOCK_DGRAM) |
| 79 | s.bind((address, 0)) # Port 0 = auto assign |
| 80 | s.close() |
| 81 | return True |
| 82 | except OSError: |
| 83 | return False |
| 84 | |
| 85 | |
| 86 | def get_loopback_ip() -> str: |
| 87 | loopback_ip = envs.VLLM_LOOPBACK_IP |
| 88 | if loopback_ip: |
| 89 | return loopback_ip |
| 90 | |
| 91 | # VLLM_LOOPBACK_IP is not set, try to get it based on network interface |
| 92 | |
| 93 | if test_loopback_bind("127.0.0.1", socket.AF_INET): |
| 94 | return "127.0.0.1" |
| 95 | elif test_loopback_bind("::1", socket.AF_INET6): |
| 96 | return "::1" |
| 97 | else: |
| 98 | raise RuntimeError( |
| 99 | "Neither 127.0.0.1 nor ::1 are bound to a local interface. " |
| 100 | "Set the VLLM_LOOPBACK_IP environment variable explicitly." |
| 101 | ) |
| 102 | |
| 103 | |
| 104 | def is_valid_ipv6_address(address: str) -> bool: |
| 105 | try: |
| 106 | ipaddress.IPv6Address(address) |
| 107 | return True |
| 108 | except ValueError: |
| 109 | return False |
| 110 | |
| 111 | |
| 112 | def split_host_port(host_port: str) -> tuple[str, int]: |
| 113 | # ipv6 |
| 114 | if host_port.startswith("["): |
| 115 | host, port = host_port.rsplit("]", 1) |
| 116 | host = host[1:] |
| 117 | port = port.split(":")[1] |
| 118 | return host, int(port) |
| 119 | else: |
| 120 | host, port = host_port.split(":") |
| 121 | return host, int(port) |
| 122 | |
| 123 | |
| 124 | def join_host_port(host: str, port: int) -> str: |
| 125 | if is_valid_ipv6_address(host): |
| 126 | return f"[{host}]:{port}" |
| 127 | else: |
| 128 | return f"{host}:{port}" |
| 129 | |
| 130 | |
| 131 | def get_distributed_init_method(ip: str, port: int) -> str: |
| 132 | return get_tcp_uri(ip, port) |
| 133 | |
| 134 | |
| 135 | def get_file_store_init_method() -> str: |
| 136 | return f"file://{tempfile.gettempdir()}/vllm_dist_{uuid4().hex}" |
| 137 | |
| 138 | |
| 139 | def aiter_requires_tcp_store() -> bool: |
| 140 | """AITER custom all-reduce requires a pure-TCP default store (its IPC |
| 141 | metadata exchange asserts on ``TCPStore``); the file:// rendezvous yields a |
| 142 | ``FileStore`` and trips that assertion. Prefer the TCP rendezvous |
| 143 | (pre-#50999) for ROCm + AITER custom AR until AITER accepts FileStore. |
| 144 | """ |
| 145 | from vllm._aiter_ops import rocm_aiter_ops |
| 146 | |
| 147 | return rocm_aiter_ops.is_custom_all_reduce_enabled() |
| 148 | |
| 149 | |
| 150 | def get_tcp_uri(ip: str, port: int) -> str: |
| 151 | if is_valid_ipv6_address(ip): |
| 152 | return f"tcp://[{ip}]:{port}" |
| 153 | else: |
| 154 | return f"tcp://{ip}:{port}" |
| 155 | |
| 156 | |
| 157 | def get_open_zmq_ipc_path() -> str: |
| 158 | base_rpc_path = envs.VLLM_RPC_BASE_PATH |
| 159 | return f"ipc://{base_rpc_path}/{uuid4()}" |
| 160 | |
| 161 | |
| 162 | def get_open_zmq_inproc_path() -> str: |
| 163 | return f"inproc://{uuid4()}" |
| 164 | |
| 165 | |
| 166 | def _get_reserved_port_range() -> range: |
| 167 | """Ports reserved for the data parallel master process (empty if unset).""" |
| 168 | if "VLLM_DP_MASTER_PORT" not in os.environ: |
| 169 | return range(0) |
| 170 | dp_master_port = envs.VLLM_DP_MASTER_PORT |
| 171 | return range(dp_master_port, dp_master_port + 10) |
| 172 | |
| 173 | |
| 174 | def get_open_port() -> int: |
| 175 | """Get an open port for the vLLM process to listen on. |
| 176 | An edge case to handle, is when we run data parallel, |
| 177 | we need to avoid ports that are potentially used by |
| 178 | the data parallel master process. |
| 179 | Right now we reserve 10 ports for the data parallel master |
| 180 | process. Currently it uses 2 ports. |
| 181 | """ |
| 182 | reserved_port_range = _get_reserved_port_range() |
| 183 | port = _get_open_port() |
| 184 | if port in reserved_port_range: |
| 185 | port = _get_open_port(start_port=reserved_port_range.stop, max_attempts=1000) |
| 186 | return port |
| 187 | |
| 188 | |
| 189 | def get_open_ports_list(count: int = 5) -> list[int]: |
| 190 | """Get a list of unique open ports. |
| 191 | |
| 192 | When VLLM_PORT is set, scans upward from that port, advancing |
| 193 | the start position after each find so every port is unique. |
| 194 | """ |
| 195 | ports_set = set[int]() |
| 196 | if envs.VLLM_PORT is not None: |
| 197 | reserved_port_range = _get_reserved_port_range() |
| 198 | next_port = envs.VLLM_PORT |
| 199 | for _ in range(count): |
| 200 | port = _get_open_port(start_port=next_port, max_attempts=1000) |
| 201 | if port in reserved_port_range: |
| 202 | port = _get_open_port( |
| 203 | start_port=reserved_port_range.stop, max_attempts=1000 |
| 204 | ) |
| 205 | ports_set.add(port) |
| 206 | next_port = port + 1 |
| 207 | return list(ports_set) |
| 208 | else: |
| 209 | while len(ports_set) < count: |
| 210 | ports_set.add(get_open_port()) |
| 211 | |
| 212 | return list(ports_set) |
| 213 | |
| 214 | |
| 215 | def _get_open_port( |
| 216 | start_port: int | None = None, |
| 217 | max_attempts: int | None = None, |
| 218 | ) -> int: |
| 219 | start_port = start_port if start_port is not None else envs.VLLM_PORT |
| 220 | port = start_port |
| 221 | if port is not None: |
| 222 | attempts = 0 |
| 223 | while True: |
| 224 | try: |
| 225 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
| 226 | s.bind(("", port)) |
| 227 | return port |
| 228 | except OSError: |
| 229 | port += 1 # Increment port number if already in use |
| 230 | logger.info("Port %d is already in use, trying port %d", port - 1, port) |
| 231 | attempts += 1 |
| 232 | if max_attempts is not None and attempts >= max_attempts: |
| 233 | raise RuntimeError( |
| 234 | f"Could not find open port after {max_attempts} " |
| 235 | f"attempts starting from port {start_port}" |
| 236 | ) |
| 237 | # try ipv4 |
| 238 | try: |
| 239 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
| 240 | s.bind(("", 0)) |
| 241 | return s.getsockname()[1] |
| 242 | except OSError: |
| 243 | # try ipv6 |
| 244 | with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s: |
| 245 | s.bind(("", 0)) |
| 246 | return s.getsockname()[1] |
| 247 | |
| 248 | |
| 249 | def find_process_using_port(port: int) -> psutil.Process | None: |
| 250 | # TODO: We can not check for running processes with network |
| 251 | # port on macOS. Therefore, we can not have a full graceful shutdown |
| 252 | # of vLLM. For now, let's not look for processes in this case. |
| 253 | # Ref: https://www.florianreinhard.de/accessdenied-in-psutil/ |
| 254 | if sys.platform.startswith("darwin"): |
| 255 | return None |
| 256 | |
| 257 | our_pid = os.getpid() |
| 258 | for conn in psutil.net_connections(): |
| 259 | if conn.laddr.port == port and (conn.pid is not None and conn.pid != our_pid): |
| 260 | try: |
| 261 | return psutil.Process(conn.pid) |
| 262 | except psutil.NoSuchProcess: |
| 263 | return None |
| 264 | return None |
| 265 | |
| 266 | |
| 267 | def split_zmq_path(path: str) -> tuple[str, str, str]: |
| 268 | """Split a zmq path into its parts.""" |
| 269 | parsed = parse_url(path) |
| 270 | if not parsed.scheme: |
| 271 | raise ValueError(f"Invalid zmq path: {path}") |
| 272 | |
| 273 | scheme = parsed.scheme |
| 274 | host = parsed.hostname or "" |
| 275 | port = "" if parsed.port is None else str(parsed.port) |
| 276 | if host.startswith("[") and host.endswith("]"): |
| 277 | host = host[1:-1] # Remove brackets for IPv6 address |
| 278 | |
| 279 | if scheme == "tcp" and not all((host, port)): |
| 280 | # The host and port fields are required for tcp |
| 281 | raise ValueError(f"Invalid zmq path: {path}") |
| 282 | |
| 283 | if scheme != "tcp" and port: |
| 284 | # port only makes sense with tcp |
| 285 | raise ValueError(f"Invalid zmq path: {path}") |
| 286 | |
| 287 | return scheme, host, port |
| 288 | |
| 289 | |
| 290 | def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str: |
| 291 | """Make a ZMQ path from its parts. |
| 292 | |
| 293 | Args: |
| 294 | scheme: The ZMQ transport scheme (e.g. tcp, ipc, inproc). |
| 295 | host: The host - can be an IPv4 address, IPv6 address, or hostname. |
| 296 | port: Optional port number, only used for TCP sockets. |
| 297 | |
| 298 | Returns: |
| 299 | A properly formatted ZMQ path string. |
| 300 | |
| 301 | """ |
| 302 | if port is None: |
| 303 | return f"{scheme}://{host}" |
| 304 | if is_valid_ipv6_address(host): |
| 305 | return f"{scheme}://[{host}]:{port}" |
| 306 | return f"{scheme}://{host}:{port}" |
| 307 | |
| 308 | |
| 309 | # Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501 |
| 310 | def make_zmq_socket( |
| 311 | ctx: zmq.asyncio.Context | zmq.Context, # type: ignore[name-defined] |
| 312 | path: str, |
| 313 | socket_type: Any, |
| 314 | bind: bool | None = None, |
| 315 | identity: bytes | None = None, |
| 316 | linger: int | None = None, |
| 317 | router_handover: bool = False, |
| 318 | ) -> zmq.Socket | zmq.asyncio.Socket: # type: ignore[name-defined] |
| 319 | """Make a ZMQ socket with the proper bind/connect semantics.""" |
| 320 | mem = psutil.virtual_memory() |
| 321 | socket = ctx.socket(socket_type) |
| 322 | |
| 323 | # Calculate buffer size based on system memory |
| 324 | total_mem = mem.total / 1024**3 |
| 325 | available_mem = mem.available / 1024**3 |
| 326 | # For systems with substantial memory (>32GB total, >16GB available): |
| 327 | # - Set a large 0.5GB buffer to improve throughput |
| 328 | # For systems with less memory: |
| 329 | # - Use system default (-1) to avoid excessive memory consumption |
| 330 | buf_size = int(0.5 * 1024**3) if total_mem > 32 and available_mem > 16 else -1 |
| 331 | |
| 332 | if bind is None: |
| 333 | bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB) |
| 334 | |
| 335 | if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER): |
| 336 | socket.setsockopt(zmq.RCVHWM, 0) |
| 337 | socket.setsockopt(zmq.RCVBUF, buf_size) |
| 338 | |
| 339 | if socket_type in (zmq.PUSH, zmq.DEALER, zmq.ROUTER): |
| 340 | socket.setsockopt(zmq.SNDHWM, 0) |
| 341 | socket.setsockopt(zmq.SNDBUF, buf_size) |
| 342 | |
| 343 | if socket_type == zmq.ROUTER and router_handover: |
| 344 | # Let a new connection take over an identity left behind by a dead one. |
| 345 | socket.setsockopt(zmq.ROUTER_HANDOVER, 1) |
| 346 | |
| 347 | if identity is not None: |
| 348 | socket.setsockopt(zmq.IDENTITY, identity) |
| 349 | |
| 350 | if linger is not None: |
| 351 | socket.setsockopt(zmq.LINGER, linger) |
| 352 | |
| 353 | if socket_type == zmq.XPUB: |
| 354 | socket.setsockopt(zmq.XPUB_VERBOSE, True) |
| 355 | |
| 356 | # Determine if the path is a TCP socket with an IPv6 address. |
| 357 | # Enable IPv6 on the zmq socket if so. |
| 358 | scheme, host, _ = split_zmq_path(path) |
| 359 | if scheme == "tcp" and is_valid_ipv6_address(host): |
| 360 | socket.setsockopt(zmq.IPV6, 1) |
| 361 | |
| 362 | if bind: |
| 363 | socket.bind(path) |
| 364 | else: |
| 365 | socket.connect(path) |
| 366 | |
| 367 | return socket |
| 368 | |
| 369 | |
| 370 | @contextlib.contextmanager |
| 371 | def zmq_socket_ctx( |
| 372 | path: str, |
| 373 | socket_type: Any, |
| 374 | bind: bool | None = None, |
| 375 | linger: int = 0, |
| 376 | identity: bytes | None = None, |
| 377 | router_handover: bool = False, |
| 378 | ) -> Iterator[zmq.Socket]: |
| 379 | """Context manager for a ZMQ socket.""" |
| 380 | ctx = zmq.Context() # type: ignore[attr-defined] |
| 381 | try: |
| 382 | yield make_zmq_socket( |
| 383 | ctx, |
| 384 | path, |
| 385 | socket_type, |
| 386 | bind=bind, |
| 387 | identity=identity, |
| 388 | router_handover=router_handover, |
| 389 | ) |
| 390 | except KeyboardInterrupt: |
| 391 | logger.debug("Got Keyboard Interrupt.") |
| 392 | |
| 393 | finally: |
| 394 | ctx.destroy(linger=linger) |