| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import contextlib |
| 7 | import logging |
| 8 | import multiprocessing |
| 9 | import os |
| 10 | import signal |
| 11 | import sys |
| 12 | from collections.abc import Callable, Iterator |
| 13 | from pathlib import Path |
| 14 | |
| 15 | import psutil |
| 16 | |
| 17 | import vllm.envs as envs |
| 18 | from vllm.logger import init_logger, set_vllm_process_name |
| 19 | from vllm.platforms import current_platform |
| 20 | from vllm.platforms.interface import in_wsl |
| 21 | from vllm.ray.lazy_utils import is_in_ray_actor |
| 22 | |
| 23 | from .platform_utils import cuda_is_initialized, xpu_is_initialized |
| 24 | |
| 25 | logger = init_logger(__name__) |
| 26 | |
| 27 | # Environment variable utilities |
| 28 | |
| 29 | |
| 30 | def update_environment_variables(envs_dict: dict[str, str]): |
| 31 | """Update multiple environment variables with logging.""" |
| 32 | for k, v in envs_dict.items(): |
| 33 | if k in os.environ and os.environ[k] != v: |
| 34 | logger.warning( |
| 35 | "Overwriting environment variable %s from '%s' to '%s'", |
| 36 | k, |
| 37 | os.environ[k], |
| 38 | v, |
| 39 | ) |
| 40 | os.environ[k] = v |
| 41 | |
| 42 | |
| 43 | @contextlib.contextmanager |
| 44 | def set_env_var(key: str, value: str) -> Iterator[None]: |
| 45 | """Temporarily set an environment variable.""" |
| 46 | old = os.environ.get(key) |
| 47 | os.environ[key] = value |
| 48 | try: |
| 49 | yield |
| 50 | finally: |
| 51 | if old is None: |
| 52 | os.environ.pop(key, None) |
| 53 | else: |
| 54 | os.environ[key] = old |
| 55 | |
| 56 | |
| 57 | @contextlib.contextmanager |
| 58 | def suppress_stdout(): |
| 59 | """Suppress stdout from C libraries at the file descriptor level. |
| 60 | |
| 61 | Only suppresses stdout, not stderr, to preserve error messages. |
| 62 | Suppression is disabled when vLLM debug logging is enabled. |
| 63 | |
| 64 | Example: |
| 65 | with suppress_stdout(): |
| 66 | # C library calls that would normally print to stdout |
| 67 | torch.distributed.new_group(ranks, backend="gloo") |
| 68 | |
| 69 | """ |
| 70 | # Don't suppress if logging level is DEBUG |
| 71 | if logger.isEnabledFor(logging.DEBUG): |
| 72 | yield |
| 73 | return |
| 74 | |
| 75 | stdout_fd = sys.stdout.fileno() |
| 76 | stdout_dup = os.dup(stdout_fd) |
| 77 | devnull_fd = os.open(os.devnull, os.O_WRONLY) |
| 78 | |
| 79 | try: |
| 80 | sys.stdout.flush() |
| 81 | os.dup2(devnull_fd, stdout_fd) |
| 82 | yield |
| 83 | finally: |
| 84 | sys.stdout.flush() |
| 85 | os.dup2(stdout_dup, stdout_fd) |
| 86 | os.close(stdout_dup) |
| 87 | os.close(devnull_fd) |
| 88 | |
| 89 | |
| 90 | # File path utilities |
| 91 | |
| 92 | |
| 93 | def unique_filepath(fn: Callable[[int], Path]) -> Path: |
| 94 | """Generate a unique file path by trying incrementing integers. |
| 95 | |
| 96 | Note: This function has a TOCTOU race condition. |
| 97 | Caller should use atomic operations (e.g., open with 'x' mode) |
| 98 | when creating the file to ensure thread safety. |
| 99 | """ |
| 100 | i = 0 |
| 101 | while True: |
| 102 | p = fn(i) |
| 103 | if not p.exists(): |
| 104 | return p |
| 105 | i += 1 |
| 106 | |
| 107 | |
| 108 | # Process management utilities |
| 109 | |
| 110 | |
| 111 | def _sync_visible_devices_env_vars(): |
| 112 | """Sync HIP/CUDA visibility env vars before spawning (ROCm only).""" |
| 113 | if not current_platform.is_rocm(): |
| 114 | return |
| 115 | |
| 116 | from vllm.platforms.rocm import _sync_hip_cuda_env_vars |
| 117 | |
| 118 | _sync_hip_cuda_env_vars() |
| 119 | |
| 120 | |
| 121 | def _maybe_force_spawn(): |
| 122 | """Check if we need to force the use of the `spawn` multiprocessing start |
| 123 | method. |
| 124 | """ |
| 125 | if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") == "spawn": |
| 126 | return |
| 127 | |
| 128 | reasons = [] |
| 129 | if is_in_ray_actor(): |
| 130 | # even if we choose to spawn, we need to pass the ray address |
| 131 | # to the subprocess so that it knows how to connect to the ray cluster. |
| 132 | # env vars are inherited by subprocesses, even if we use spawn. |
| 133 | import ray |
| 134 | |
| 135 | os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address |
| 136 | reasons.append("In a Ray actor and can only be spawned") |
| 137 | |
| 138 | # Force spawn if NUMA binding is enabled via --numa-bind. |
| 139 | # NUMA binding uses executable hijacking which requires spawn |
| 140 | if "--numa-bind" in sys.argv: |
| 141 | reasons.append("NUMA binding requires spawn method") |
| 142 | |
| 143 | if cuda_is_initialized(): |
| 144 | reasons.append("CUDA is initialized") |
| 145 | elif xpu_is_initialized(): |
| 146 | reasons.append("XPU is initialized") |
| 147 | |
| 148 | if in_wsl(): |
| 149 | reasons.append("WSL is detected and NVML is not compatible with fork") |
| 150 | |
| 151 | if reasons: |
| 152 | logger.warning( |
| 153 | "We must use the `spawn` multiprocessing start method. " |
| 154 | "Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. " |
| 155 | "See https://docs.vllm.ai/en/latest/usage/" |
| 156 | "troubleshooting.html#python-multiprocessing " |
| 157 | "for more information. Reasons: %s", |
| 158 | "; ".join(reasons), |
| 159 | ) |
| 160 | os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" |
| 161 | |
| 162 | |
| 163 | def get_mp_context(): |
| 164 | """Get a multiprocessing context with a particular method (spawn or fork). |
| 165 | By default we follow the value of the VLLM_WORKER_MULTIPROC_METHOD to |
| 166 | determine the multiprocessing method (default is fork). However, under |
| 167 | certain conditions, we may enforce spawn and override the value of |
| 168 | VLLM_WORKER_MULTIPROC_METHOD. |
| 169 | """ |
| 170 | _maybe_force_spawn() |
| 171 | # (ROCm): Sync GPU visibility env vars so spawned children inherit |
| 172 | # consistent values. Must run after _maybe_force_spawn and regardless |
| 173 | # of whether spawn was already set. |
| 174 | _sync_visible_devices_env_vars() |
| 175 | mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD |
| 176 | return multiprocessing.get_context(mp_method) |
| 177 | |
| 178 | |
| 179 | def set_process_title( |
| 180 | name: str, |
| 181 | suffix: str = "", |
| 182 | prefix: str = envs.VLLM_PROCESS_NAME_PREFIX, |
| 183 | ) -> None: |
| 184 | """Set the current process title with optional suffix.""" |
| 185 | try: |
| 186 | import setproctitle |
| 187 | except ImportError: |
| 188 | return |
| 189 | |
| 190 | if suffix: |
| 191 | name = f"{name}_{suffix}" |
| 192 | |
| 193 | setproctitle.setproctitle(f"{prefix}::{name}") |
| 194 | |
| 195 | |
| 196 | def decorate_logs( |
| 197 | process_name: str | None = None, *, skip_if_decorated: bool = False |
| 198 | ) -> None: |
| 199 | """Add vLLM process metadata to subsequent log records.""" |
| 200 | if process_name is None: |
| 201 | process_name = get_mp_context().current_process().name |
| 202 | |
| 203 | set_vllm_process_name(process_name, skip_if_set=skip_if_decorated) |
| 204 | |
| 205 | |
| 206 | def kill_process_tree(pid: int): |
| 207 | """Kills all descendant processes of the given pid by sending SIGKILL. |
| 208 | |
| 209 | Args: |
| 210 | pid (int): Process ID of the parent process |
| 211 | |
| 212 | """ |
| 213 | try: |
| 214 | parent = psutil.Process(pid) |
| 215 | except psutil.NoSuchProcess: |
| 216 | return |
| 217 | |
| 218 | # Get all children recursively |
| 219 | children = parent.children(recursive=True) |
| 220 | |
| 221 | # Send SIGKILL to all children first |
| 222 | for child in children: |
| 223 | with contextlib.suppress(ProcessLookupError): |
| 224 | logger.warning( |
| 225 | "[shutdown] force killing sub process %s pid %d", |
| 226 | child.name(), |
| 227 | child.pid, |
| 228 | ) |
| 229 | os.kill(child.pid, signal.SIGKILL) |
| 230 | |
| 231 | # Finally kill the parent |
| 232 | with contextlib.suppress(ProcessLookupError): |
| 233 | os.kill(pid, signal.SIGKILL) |
| 234 | |
| 235 | |
| 236 | # Resource utilities |
| 237 | |
| 238 | |
| 239 | # Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L630 |
| 240 | def set_ulimit(target_soft_limit: int = 65535): |
| 241 | if sys.platform.startswith("win"): |
| 242 | logger.debug("Windows detected, skipping ulimit adjustment.") |
| 243 | return |
| 244 | |
| 245 | import resource |
| 246 | |
| 247 | resource_type = resource.RLIMIT_NOFILE |
| 248 | current_soft, current_hard = resource.getrlimit(resource_type) |
| 249 | |
| 250 | if current_soft < target_soft_limit: |
| 251 | try: |
| 252 | resource.setrlimit(resource_type, (target_soft_limit, current_hard)) |
| 253 | except ValueError as e: |
| 254 | logger.warning( |
| 255 | "Found ulimit of %s and failed to automatically increase " |
| 256 | "with error %s. This can cause fd limit errors like " |
| 257 | "`OSError: [Errno 24] Too many open files`. Consider " |
| 258 | "increasing with ulimit -n", |
| 259 | current_soft, |
| 260 | e, |
| 261 | ) |
| 262 | |
| 263 | |
| 264 | def find_loaded_library(lib_name: str) -> str | None: |
| 265 | """ |
| 266 | According to https://man7.org/linux/man-pages/man5/proc_pid_maps.5.html, |
| 267 | the file `/proc/self/maps` contains the memory maps of the process, which includes the |
| 268 | shared libraries loaded by the process. We can use this file to find the path of the |
| 269 | loaded library. |
| 270 | """ # noqa |
| 271 | # Match the mapped file's name, not the whole line: an unrelated library |
| 272 | # whose name merely contains lib_name (e.g. TileLang's libcudart_stub.so |
| 273 | # when looking for libcudart) or a directory component containing it must |
| 274 | # not win. Legitimate filenames are {lib_name}.so[.*] or a name-mangled |
| 275 | # {lib_name}-<hash>.so[.*], and /proc/self/maps is ordered by mapping |
| 276 | # address rather than load order, so a substring hit is a |
| 277 | # nondeterministic hijack. |
| 278 | with open("/proc/self/maps") as f: |
| 279 | for line in f: |
| 280 | start = line.find("/") |
| 281 | if start == -1: |
| 282 | continue |
| 283 | path = line[start:].strip() |
| 284 | filename = path.rsplit("/", maxsplit=1)[-1] |
| 285 | if filename.startswith((f"{lib_name}.", f"{lib_name}-")): |
| 286 | return path |
| 287 | # the library is not loaded in the current process |
| 288 | return None |