A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Chore] Separate out system utilities from vllm.utils (#27201)
Signed-off-by: dongbo910220 <1275604947@qq.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
dongbo910220 committed
Oct 23, 2025 at 04:25 UTC
a0003b56b0b822c52bb0f3035c164370a802e6f5
25 files changed
+203
-159
tests/compile/test_async_tp.py
+1
-1
@@ -25,7 +25,7 @@ from vllm.distributed.parallel_state import (
25
initialize_model_parallel,
26
)
27
from vllm.platforms import current_platform
28
-from vllm.utils import update_environment_variables
28
+from vllm.utils.system_utils import update_environment_variables
29
30
from ..models.registry import HF_EXAMPLE_MODELS
31
from ..utils import (
tests/compile/test_fusion_all_reduce.py
+1
-1
@@ -31,7 +31,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
31
GroupShape,
32
)
33
from vllm.platforms import current_platform
34
-from vllm.utils import update_environment_variables
34
+from vllm.utils.system_utils import update_environment_variables
35
36
from ..utils import has_module_attribute, multi_gpu_test
37
from .backend import TestBackend
tests/compile/test_sequence_parallelism.py
+1
-1
@@ -29,7 +29,7 @@ from vllm.distributed.parallel_state import (
29
from vllm.model_executor.layers.layernorm import RMSNorm
30
from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp
31
from vllm.platforms import current_platform
32
-from vllm.utils import update_environment_variables
32
+from vllm.utils.system_utils import update_environment_variables
33
34
from ..utils import multi_gpu_test
35
from .backend import TestBackend
tests/distributed/test_eplb_execute.py
+1
-1
@@ -15,7 +15,7 @@ from vllm.distributed.parallel_state import (
15
get_tp_group,
16
init_distributed_environment,
17
)
18
-from vllm.utils import update_environment_variables
18
+from vllm.utils.system_utils import update_environment_variables
19
20
21
def distributed_run(fn, world_size):
tests/distributed/test_nccl_symm_mem_allreduce.py
+1
-1
@@ -23,7 +23,7 @@ from vllm.distributed.parallel_state import (
23
initialize_model_parallel,
24
)
25
from vllm.platforms import current_platform
26
-from vllm.utils import update_environment_variables
26
+from vllm.utils.system_utils import update_environment_variables
27
28
torch.manual_seed(42)
29
random.seed(44)
tests/distributed/test_pynccl.py
+1
-1
@@ -18,7 +18,7 @@ from vllm.distributed.parallel_state import (
18
graph_capture,
19
init_distributed_environment,
20
)
21
-from vllm.utils import update_environment_variables
21
+from vllm.utils.system_utils import update_environment_variables
22
23
24
def distributed_run(fn, world_size):
tests/distributed/test_shm_broadcast.py
+1
-1
@@ -10,8 +10,8 @@ import torch.distributed as dist
10
11
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
12
from vllm.distributed.utils import StatelessProcessGroup
13
-from vllm.utils import update_environment_variables
13
from vllm.utils.network_utils import get_open_port
14
+from vllm.utils.system_utils import update_environment_variables
15
16
17
def get_arrays(n: int, seed: int = 0) -> list[np.ndarray]:
tests/distributed/test_symm_mem_allreduce.py
+1
-1
@@ -23,7 +23,7 @@ from vllm.distributed.parallel_state import (
23
from vllm.engine.arg_utils import EngineArgs
24
from vllm.engine.llm_engine import LLMEngine
25
from vllm.platforms import current_platform
26
-from vllm.utils import update_environment_variables
26
+from vllm.utils.system_utils import update_environment_variables
27
28
torch.manual_seed(42)
29
random.seed(44)
tests/distributed/test_utils.py
+1
-1
@@ -10,8 +10,8 @@ import torch
10
import vllm.envs as envs
11
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
12
from vllm.distributed.utils import StatelessProcessGroup
13
-from vllm.utils import update_environment_variables
13
from vllm.utils.network_utils import get_open_port
14
+from vllm.utils.system_utils import update_environment_variables
15
from vllm.utils.torch_utils import cuda_device_count_stateless
16
17
from ..utils import multi_gpu_test
tests/kernels/mamba/test_mamba_mixer2.py
+1
-1
@@ -13,7 +13,7 @@ from vllm.distributed.parallel_state import (
13
)
14
from vllm.model_executor.layers.mamba.mamba_mixer2 import Mixer2RMSNormGated
15
from vllm.platforms import current_platform
16
-from vllm.utils import update_environment_variables
16
+from vllm.utils.system_utils import update_environment_variables
17
18
19
@multi_gpu_test(num_gpus=2)
tests/models/test_vision.py
+1
-1
@@ -19,8 +19,8 @@ from vllm.model_executor.models.vision import (
19
run_dp_sharded_vision_model,
20
)
21
from vllm.platforms import current_platform
22
-from vllm.utils import update_environment_variables
22
from vllm.utils.network_utils import get_open_port
23
+from vllm.utils.system_utils import update_environment_variables
24
25
pytestmark = pytest.mark.cpu_test
26
tests/utils_/test_system_utils.py
new
+19
@@ -0,0 +1,19 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import tempfile
5
+from pathlib import Path
6
+
7
+from vllm.utils.system_utils import unique_filepath
8
+
9
+
10
+def test_unique_filepath():
11
+ temp_dir = tempfile.mkdtemp()
12
+ path_fn = lambda i: Path(temp_dir) / f"file_{i}.txt"
13
+ paths = set()
14
+ for i in range(10):
15
+ path = unique_filepath(path_fn)
16
+ path.write_text("test")
17
+ paths.add(path)
18
+ assert len(paths) == 10
19
+ assert len(list(Path(temp_dir).glob("*.txt"))) == 10
tests/utils_/test_utils.py
-13
@@ -19,7 +19,6 @@ from vllm.transformers_utils.detokenizer_utils import convert_ids_list_to_tokens
19
from vllm.utils import (
20
FlexibleArgumentParser,
21
bind_kv_cache,
22
- unique_filepath,
22
)
23
from ..utils import create_new_process_for_each_test, flat_product
24
@@ -466,18 +465,6 @@ def test_load_config_file(tmp_path):
465
os.remove(str(config_file_path))
466
467
469
-def test_unique_filepath():
470
- temp_dir = tempfile.mkdtemp()
471
- path_fn = lambda i: Path(temp_dir) / f"file_{i}.txt"
472
- paths = set()
473
- for i in range(10):
474
- path = unique_filepath(path_fn)
475
- path.write_text("test")
476
- paths.add(path)
477
- assert len(paths) == 10
478
- assert len(list(Path(temp_dir).glob("*.txt"))) == 10
479
-
480
-
468
def test_flat_product():
469
# Check regular itertools.product behavior
470
result1 = list(flat_product([1, 2, 3], ["a", "b"]))
tests/v1/worker/test_gpu_model_runner.py
+1
-1
@@ -21,8 +21,8 @@ from vllm.distributed.parallel_state import (
21
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
22
from vllm.platforms import current_platform
23
from vllm.sampling_params import SamplingParams
24
-from vllm.utils import update_environment_variables
24
from vllm.utils.mem_constants import GiB_bytes
25
+from vllm.utils.system_utils import update_environment_variables
26
from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs
27
from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
28
from vllm.v1.kv_cache_interface import (
vllm/compilation/pass_manager.py
+1
-1
@@ -8,7 +8,7 @@ from vllm import envs
8
from vllm.config import VllmConfig, set_current_vllm_config
9
from vllm.logger import init_logger
10
from vllm.platforms import current_platform
11
-from vllm.utils import set_env_var
11
+from vllm.utils.system_utils import set_env_var
12
13
from .post_cleanup import PostCleanupPass
14
from .vllm_inductor_pass import VllmInductorPass
vllm/compilation/vllm_inductor_pass.py
+1
-1
@@ -114,7 +114,7 @@ class VllmPatternMatcherPass(VllmInductorPass):
114
115
debug_dump_path.mkdir(parents=True, exist_ok=True)
116
117
- from vllm.utils import unique_filepath
117
+ from vllm.utils.system_utils import unique_filepath
118
119
file_path = unique_filepath(
120
lambda i: debug_dump_path / f"patterns.{self.pass_name}.{i}.py"
vllm/distributed/device_communicators/all_reduce_utils.py
+1
-1
@@ -22,7 +22,7 @@ from vllm.logger import init_logger
22
from vllm.model_executor.layers.batch_invariant import (
23
vllm_is_batch_invariant,
24
)
25
-from vllm.utils import update_environment_variables
25
+from vllm.utils.system_utils import update_environment_variables
26
from vllm.utils.torch_utils import cuda_device_count_stateless
27
28
logger = init_logger(__name__)
vllm/entrypoints/cli/serve.py
+2
-5
@@ -18,12 +18,9 @@ from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_se
18
from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
19
from vllm.logger import init_logger
20
from vllm.usage.usage_lib import UsageContext
21
-from vllm.utils import (
22
- FlexibleArgumentParser,
23
- decorate_logs,
24
- set_process_title,
25
-)
21
+from vllm.utils import FlexibleArgumentParser
22
from vllm.utils.network_utils import get_tcp_uri
23
+from vllm.utils.system_utils import decorate_logs, set_process_title
24
from vllm.v1.engine.core import EngineCoreProc
25
from vllm.v1.engine.utils import CoreEngineProcManager, launch_core_engines
26
from vllm.v1.executor import Executor
vllm/entrypoints/openai/api_server.py
+2
-6
@@ -108,13 +108,9 @@ from vllm.entrypoints.utils import (
108
from vllm.logger import init_logger
109
from vllm.reasoning import ReasoningParserManager
110
from vllm.usage.usage_lib import UsageContext
111
-from vllm.utils import (
112
- Device,
113
- FlexibleArgumentParser,
114
- decorate_logs,
115
- set_ulimit,
116
-)
111
+from vllm.utils import Device, FlexibleArgumentParser, set_ulimit
112
from vllm.utils.network_utils import is_valid_ipv6_address
113
+from vllm.utils.system_utils import decorate_logs
114
from vllm.v1.engine.exceptions import EngineDeadError
115
from vllm.v1.metrics.prometheus import get_prometheus_registry
116
from vllm.version import __version__ as VLLM_VERSION
vllm/utils/__init__.py
+36
-108
@@ -5,6 +5,7 @@ import contextlib
5
import datetime
6
import enum
7
import getpass
8
+import importlib.util
9
import inspect
10
import json
11
import multiprocessing
@@ -33,13 +34,11 @@ from collections.abc import (
34
)
35
from concurrent.futures.process import ProcessPoolExecutor
36
from functools import cache, partial, wraps
36
-from pathlib import Path
37
-from typing import TYPE_CHECKING, Any, TextIO, TypeVar
37
+from typing import TYPE_CHECKING, Any, TypeVar
38
39
import cloudpickle
40
import psutil
41
import regex as re
42
-import setproctitle
42
import torch
43
import yaml
44
@@ -144,18 +143,6 @@ def random_uuid() -> str:
143
return str(uuid.uuid4().hex)
144
145
147
-def update_environment_variables(envs: dict[str, str]):
148
- for k, v in envs.items():
149
- if k in os.environ and os.environ[k] != v:
150
- logger.warning(
151
- "Overwriting environment variable %s from '%s' to '%s'",
152
- k,
153
- os.environ[k],
154
- v,
155
- )
156
- os.environ[k] = v
157
-
158
-
146
def cdiv(a: int, b: int) -> int:
147
"""Ceiling division."""
148
return -(a // -b)
@@ -1061,70 +1048,44 @@ def check_use_alibi(model_config: ModelConfig) -> bool:
1048
)
1049
1050
1064
-def set_process_title(
1065
- name: str, suffix: str = "", prefix: str = envs.VLLM_PROCESS_NAME_PREFIX
1066
-) -> None:
1067
- """
1068
- Set the current process title to a specific name with an
1069
- optional suffix.
1051
+@cache
1052
+def _has_module(module_name: str) -> bool:
1053
+ """Return True if *module_name* can be found in the current environment.
1054
1071
- Args:
1072
- name: The title to assign to the current process.
1073
- suffix: An optional suffix to append to the base name.
1074
- prefix: A prefix to prepend to the front separated by `::`.
1075
- """
1076
- if suffix:
1077
- name = f"{name}_{suffix}"
1078
- setproctitle.setproctitle(f"{prefix}::{name}")
1079
-
1080
-
1081
-def _add_prefix(file: TextIO, worker_name: str, pid: int) -> None:
1082
- """Prepend each output line with process-specific prefix"""
1083
-
1084
- prefix = f"{CYAN}({worker_name} pid={pid}){RESET} "
1085
- file_write = file.write
1086
-
1087
- def write_with_prefix(s: str):
1088
- if not s:
1089
- return
1090
- if file.start_new_line: # type: ignore[attr-defined]
1091
- file_write(prefix)
1092
- idx = 0
1093
- while (next_idx := s.find("\n", idx)) != -1:
1094
- next_idx += 1
1095
- file_write(s[idx:next_idx])
1096
- if next_idx == len(s):
1097
- file.start_new_line = True # type: ignore[attr-defined]
1098
- return
1099
- file_write(prefix)
1100
- idx = next_idx
1101
- file_write(s[idx:])
1102
- file.start_new_line = False # type: ignore[attr-defined]
1103
-
1104
- file.start_new_line = True # type: ignore[attr-defined]
1105
- file.write = write_with_prefix # type: ignore[method-assign]
1106
-
1107
-
1108
-def decorate_logs(process_name: str | None = None) -> None:
1055
+ The result is cached so that subsequent queries for the same module incur
1056
+ no additional overhead.
1057
"""
1110
- Adds a process-specific prefix to each line of output written to stdout and
1111
- stderr.
1058
+ return importlib.util.find_spec(module_name) is not None
1059
1113
- This function is intended to be called before initializing the api_server,
1114
- engine_core, or worker classes, so that all subsequent output from the
1115
- process is prefixed with the process name and PID. This helps distinguish
1116
- log output from different processes in multi-process environments.
1060
1118
- Args:
1119
- process_name: Optional; the name of the process to use in the prefix.
1120
- If not provided, the current process name from the multiprocessing
1121
- context is used.
1122
- """
1123
- if process_name is None:
1124
- process_name = get_mp_context().current_process().name
1125
- pid = os.getpid()
1126
- _add_prefix(sys.stdout, process_name, pid)
1127
- _add_prefix(sys.stderr, process_name, pid)
1061
+def has_pplx() -> bool:
1062
+ """Whether the optional `pplx_kernels` package is available."""
1063
+
1064
+ return _has_module("pplx_kernels")
1065
+
1066
+
1067
+def has_deep_ep() -> bool:
1068
+ """Whether the optional `deep_ep` package is available."""
1069
+
1070
+ return _has_module("deep_ep")
1071
+
1072
+
1073
+def has_deep_gemm() -> bool:
1074
+ """Whether the optional `deep_gemm` package is available."""
1075
+
1076
+ return _has_module("deep_gemm")
1077
+
1078
+
1079
+def has_triton_kernels() -> bool:
1080
+ """Whether the optional `triton_kernels` package is available."""
1081
+
1082
+ return _has_module("triton_kernels")
1083
+
1084
+
1085
+def has_tilelang() -> bool:
1086
+ """Whether the optional `tilelang` package is available."""
1087
+
1088
+ return _has_module("tilelang")
1089
1090
1091
def length_from_prompt_token_ids_or_embeds(
@@ -1149,36 +1110,3 @@ def length_from_prompt_token_ids_or_embeds(
1110
f" prompt_embeds={prompt_embeds_len}"
1111
)
1112
return prompt_token_len
1152
-
1153
-
1154
-@contextlib.contextmanager
1155
-def set_env_var(key, value):
1156
- old = os.environ.get(key)
1157
- os.environ[key] = value
1158
- try:
1159
- yield
1160
- finally:
1161
- if old is None:
1162
- del os.environ[key]
1163
- else:
1164
- os.environ[key] = old
1165
-
1166
-
1167
-def unique_filepath(fn: Callable[[int], Path]) -> Path:
1168
- """
1169
- unique_filepath returns a unique path by trying
1170
- to include an integer in increasing order.
1171
-
1172
- fn should be a callable that returns a path that
1173
- includes the passed int at a fixed location.
1174
-
1175
- Note: This function has a TOCTOU race condition.
1176
- Caller should use atomic operations (e.g., open with 'x' mode)
1177
- when creating the file to ensure thread safety.
1178
- """
1179
- i = 0
1180
- while True:
1181
- p = fn(i)
1182
- if not p.exists():
1183
- return p
1184
- i += 1
vllm/utils/system_utils.py
new
+123
@@ -0,0 +1,123 @@
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 os
8
+import sys
9
+from collections.abc import Callable, Iterator
10
+from pathlib import Path
11
+from typing import TextIO
12
+
13
+try:
14
+ import setproctitle
15
+except ImportError:
16
+ setproctitle = None # type: ignore[assignment]
17
+
18
+import vllm.envs as envs
19
+from vllm.logger import init_logger
20
+
21
+logger = init_logger(__name__)
22
+
23
+CYAN = "\033[1;36m"
24
+RESET = "\033[0;0m"
25
+
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
+# File path utilities
58
+
59
+
60
+def unique_filepath(fn: Callable[[int], Path]) -> Path:
61
+ """Generate a unique file path by trying incrementing integers.
62
+
63
+ Note: This function has a TOCTOU race condition.
64
+ Caller should use atomic operations (e.g., open with 'x' mode)
65
+ when creating the file to ensure thread safety.
66
+ """
67
+ i = 0
68
+ while True:
69
+ p = fn(i)
70
+ if not p.exists():
71
+ return p
72
+ i += 1
73
+
74
+
75
+# Process management utilities
76
+
77
+
78
+def set_process_title(
79
+ name: str, suffix: str = "", prefix: str = envs.VLLM_PROCESS_NAME_PREFIX
80
+) -> None:
81
+ """Set the current process title with optional suffix."""
82
+ if setproctitle is None:
83
+ return
84
+ if suffix:
85
+ name = f"{name}_{suffix}"
86
+ setproctitle.setproctitle(f"{prefix}::{name}")
87
+
88
+
89
+def _add_prefix(file: TextIO, worker_name: str, pid: int) -> None:
90
+ """Add colored prefix to file output for log decoration."""
91
+ prefix = f"{CYAN}({worker_name} pid={pid}){RESET} "
92
+ file_write = file.write
93
+
94
+ def write_with_prefix(s: str):
95
+ if not s:
96
+ return
97
+ if file.start_new_line: # type: ignore[attr-defined]
98
+ file_write(prefix)
99
+ idx = 0
100
+ while (next_idx := s.find("\n", idx)) != -1:
101
+ next_idx += 1
102
+ file_write(s[idx:next_idx])
103
+ if next_idx == len(s):
104
+ file.start_new_line = True # type: ignore[attr-defined]
105
+ return
106
+ file_write(prefix)
107
+ idx = next_idx
108
+ file_write(s[idx:])
109
+ file.start_new_line = False # type: ignore[attr-defined]
110
+
111
+ file.start_new_line = True # type: ignore[attr-defined]
112
+ file.write = write_with_prefix # type: ignore[method-assign]
113
+
114
+
115
+def decorate_logs(process_name: str | None = None) -> None:
116
+ """Decorate stdout/stderr with process name and PID prefix."""
117
+ from vllm.utils import get_mp_context
118
+
119
+ if process_name is None:
120
+ process_name = get_mp_context().current_process().name
121
+ pid = os.getpid()
122
+ _add_prefix(sys.stdout, process_name, pid)
123
+ _add_prefix(sys.stderr, process_name, pid)
vllm/v1/engine/coordinator.py
+2
-1
@@ -10,8 +10,9 @@ import zmq
10
11
from vllm.config import ParallelConfig
12
from vllm.logger import init_logger
13
-from vllm.utils import get_mp_context, set_process_title
13
+from vllm.utils import get_mp_context
14
from vllm.utils.network_utils import make_zmq_socket
15
+from vllm.utils.system_utils import set_process_title
16
from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequestType
17
from vllm.v1.serial_utils import MsgpackDecoder
18
from vllm.v1.utils import get_engine_client_zmq_addr, shutdown
vllm/v1/engine/core.py
+1
-4
@@ -28,14 +28,11 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
28
from vllm.multimodal.cache import engine_receiver_cache_from_config
29
from vllm.tasks import POOLING_TASKS, SupportedTask
30
from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
31
-from vllm.utils import (
32
- decorate_logs,
33
- set_process_title,
34
-)
31
from vllm.utils.gc_utils import maybe_attach_gc_debug_callback
32
from vllm.utils.hashing import get_hash_fn_by_name
33
from vllm.utils.import_utils import resolve_obj_by_qualname
34
from vllm.utils.network_utils import make_zmq_socket
35
+from vllm.utils.system_utils import decorate_logs, set_process_title
36
from vllm.v1.core.kv_cache_utils import (
37
BlockHash,
38
generate_scheduler_kv_cache_config,
vllm/v1/executor/multiproc_executor.py
+2
-6
@@ -35,17 +35,13 @@ from vllm.distributed.parallel_state import (
35
)
36
from vllm.envs import enable_envs_cache
37
from vllm.logger import init_logger
38
-from vllm.utils import (
39
- _maybe_force_spawn,
40
- decorate_logs,
41
- get_mp_context,
42
- set_process_title,
43
-)
38
+from vllm.utils import _maybe_force_spawn, get_mp_context
39
from vllm.utils.network_utils import (
40
get_distributed_init_method,
41
get_loopback_ip,
42
get_open_port,
43
)
44
+from vllm.utils.system_utils import decorate_logs, set_process_title
45
from vllm.v1.core.sched.output import SchedulerOutput
46
from vllm.v1.executor.abstract import Executor, FailureCallback
47
from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput
vllm/v1/worker/worker_base.py
+1
-1
@@ -16,10 +16,10 @@ from vllm.multimodal.cache import worker_receiver_cache_from_config
16
from vllm.utils import (
17
enable_trace_function_call_for_thread,
18
run_method,
19
- update_environment_variables,
19
warn_for_unimplemented_methods,
20
)
21
from vllm.utils.import_utils import resolve_obj_by_qualname
22
+from vllm.utils.system_utils import update_environment_variables
23
from vllm.v1.kv_cache_interface import KVCacheSpec
24
25
if TYPE_CHECKING: