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
[Misc] Clean up more utils (#27567)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Oct 27, 2025 at 23:30 UTC
6ebffafbb609963d696d27c5e334fd0b9fe0add7
24 files changed
+282
-315
requirements/docs.txt
+2
@@ -13,6 +13,8 @@ ruff
13
# Required for argparse hook only
14
-f https://download.pytorch.org/whl/cpu
15
cachetools
16
+cloudpickle
17
+py-cpuinfo
18
msgspec
19
pydantic
20
torch
tools/pre_commit/check_pickle_imports.py
-1
@@ -39,7 +39,6 @@ ALLOWED_FILES = {
39
"vllm/v1/executor/multiproc_executor.py",
40
"vllm/v1/executor/ray_executor.py",
41
"vllm/entrypoints/llm.py",
42
- "vllm/utils/__init__.py",
42
"tests/utils.py",
43
# pickle and cloudpickle
44
"vllm/v1/serial_utils.py",
vllm/config/model.py
+23
@@ -1618,6 +1618,29 @@ class ModelConfig:
1618
"""Extract the HF encoder/decoder model flag."""
1619
return is_encoder_decoder(self.hf_config)
1620
1621
+ @property
1622
+ def uses_alibi(self) -> bool:
1623
+ cfg = self.hf_text_config
1624
+
1625
+ return (
1626
+ getattr(cfg, "alibi", False) # Falcon
1627
+ or "BloomForCausalLM" in self.architectures # Bloom
1628
+ or getattr(cfg, "position_encoding_type", "") == "alibi" # codellm_1b_alibi
1629
+ or (
1630
+ hasattr(cfg, "attn_config") # MPT
1631
+ and (
1632
+ (
1633
+ isinstance(cfg.attn_config, dict)
1634
+ and cfg.attn_config.get("alibi", False)
1635
+ )
1636
+ or (
1637
+ not isinstance(cfg.attn_config, dict)
1638
+ and getattr(cfg.attn_config, "alibi", False)
1639
+ )
1640
+ )
1641
+ )
1642
+ )
1643
+
1644
@property
1645
def uses_mrope(self) -> bool:
1646
return uses_mrope(self.hf_config)
vllm/config/vllm.py
+27
-1
@@ -2,12 +2,16 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import copy
5
+import getpass
6
import hashlib
7
import json
8
import os
9
+import tempfile
10
+import threading
11
import time
12
from contextlib import contextmanager
13
from dataclasses import replace
14
+from datetime import datetime
15
from functools import lru_cache
16
from pathlib import Path
17
from typing import TYPE_CHECKING, Any, TypeVar
@@ -17,7 +21,7 @@ from pydantic import ConfigDict, Field
21
from pydantic.dataclasses import dataclass
22
23
import vllm.envs as envs
20
-from vllm.logger import init_logger
24
+from vllm.logger import enable_trace_function_call, init_logger
25
from vllm.transformers_utils.runai_utils import is_runai_obj_uri
26
from vllm.utils import random_uuid
27
@@ -206,6 +210,28 @@ class VllmConfig:
210
# i.e., batch_size <= self.compilation_config.max_cudagraph_capture_size
211
return self.compilation_config.bs_to_padded_graph_size[batch_size]
212
213
+ def enable_trace_function_call_for_thread(self) -> None:
214
+ """
215
+ Set up function tracing for the current thread,
216
+ if enabled via the `VLLM_TRACE_FUNCTION` environment variable.
217
+ """
218
+ if envs.VLLM_TRACE_FUNCTION:
219
+ tmp_dir = tempfile.gettempdir()
220
+ # add username to tmp_dir to avoid permission issues
221
+ tmp_dir = os.path.join(tmp_dir, getpass.getuser())
222
+ filename = (
223
+ f"VLLM_TRACE_FUNCTION_for_process_{os.getpid()}"
224
+ f"_thread_{threading.get_ident()}_at_{datetime.now()}.log"
225
+ ).replace(" ", "_")
226
+ log_path = os.path.join(
227
+ tmp_dir,
228
+ "vllm",
229
+ f"vllm-instance-{self.instance_id}",
230
+ filename,
231
+ )
232
+ os.makedirs(os.path.dirname(log_path), exist_ok=True)
233
+ enable_trace_function_call(log_path)
234
+
235
@staticmethod
236
def _get_quantization_config(
237
model_config: ModelConfig, load_config: LoadConfig
vllm/engine/arg_utils.py
+1
-2
@@ -73,7 +73,7 @@ from vllm.config.utils import get_field
73
from vllm.logger import init_logger
74
from vllm.platforms import CpuArchEnum, current_platform
75
from vllm.plugins import load_general_plugins
76
-from vllm.ray.lazy_utils import is_ray_initialized
76
+from vllm.ray.lazy_utils import is_in_ray_actor, is_ray_initialized
77
from vllm.reasoning import ReasoningParserManager
78
from vllm.test_utils import MODEL_WEIGHTS_S3_BUCKET, MODELS_ON_S3
79
from vllm.transformers_utils.config import (
@@ -82,7 +82,6 @@ from vllm.transformers_utils.config import (
82
maybe_override_with_speculators,
83
)
84
from vllm.transformers_utils.utils import check_gguf_file
85
-from vllm.utils import is_in_ray_actor
85
from vllm.utils.argparse_utils import FlexibleArgumentParser
86
from vllm.utils.mem_constants import GiB_bytes
87
from vllm.utils.network_utils import get_ip
vllm/entrypoints/anthropic/api_server.py
+1
-1
@@ -51,9 +51,9 @@ from vllm.entrypoints.utils import (
51
with_cancellation,
52
)
53
from vllm.logger import init_logger
54
-from vllm.utils import set_ulimit
54
from vllm.utils.argparse_utils import FlexibleArgumentParser
55
from vllm.utils.network_utils import is_valid_ipv6_address
56
+from vllm.utils.system_utils import set_ulimit
57
from vllm.version import __version__ as VLLM_VERSION
58
59
prometheus_multiproc_dir: tempfile.TemporaryDirectory
vllm/entrypoints/api_server.py
+2
-1
@@ -26,8 +26,9 @@ from vllm.entrypoints.utils import with_cancellation
26
from vllm.logger import init_logger
27
from vllm.sampling_params import SamplingParams
28
from vllm.usage.usage_lib import UsageContext
29
-from vllm.utils import random_uuid, set_ulimit
29
+from vllm.utils import random_uuid
30
from vllm.utils.argparse_utils import FlexibleArgumentParser
31
+from vllm.utils.system_utils import set_ulimit
32
from vllm.version import __version__ as VLLM_VERSION
33
34
logger = init_logger("vllm.entrypoints.api_server")
vllm/entrypoints/openai/api_server.py
+2
-2
@@ -108,10 +108,10 @@ 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 Device, set_ulimit
111
+from vllm.utils import Device
112
from vllm.utils.argparse_utils import FlexibleArgumentParser
113
from vllm.utils.network_utils import is_valid_ipv6_address
114
-from vllm.utils.system_utils import decorate_logs
114
+from vllm.utils.system_utils import decorate_logs, set_ulimit
115
from vllm.v1.engine.exceptions import EngineDeadError
116
from vllm.v1.metrics.prometheus import get_prometheus_registry
117
from vllm.version import __version__ as VLLM_VERSION
vllm/platforms/__init__.py
+1
-1
@@ -60,7 +60,7 @@ def cuda_platform_plugin() -> str | None:
60
is_cuda = False
61
logger.debug("Checking if CUDA platform is available.")
62
try:
63
- from vllm.utils import import_pynvml
63
+ from vllm.utils.import_utils import import_pynvml
64
65
pynvml = import_pynvml()
66
pynvml.nvmlInit()
vllm/platforms/cuda.py
+1
-1
@@ -16,7 +16,7 @@ from typing_extensions import ParamSpec
16
import vllm._C # noqa
17
import vllm.envs as envs
18
from vllm.logger import init_logger
19
-from vllm.utils import import_pynvml
19
+from vllm.utils.import_utils import import_pynvml
20
from vllm.utils.torch_utils import cuda_device_count_stateless
21
22
from .interface import DeviceCapability, Platform, PlatformEnum
vllm/utils/__init__.py
+4
-242
@@ -1,34 +1,22 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import contextlib
5
-import datetime
4
import enum
7
-import getpass
5
import inspect
9
-import multiprocessing
10
-import os
11
-import signal
12
-import sys
13
-import tempfile
6
import threading
7
import uuid
8
import warnings
17
-from collections.abc import Callable
18
-from functools import partial, wraps
19
-from typing import TYPE_CHECKING, Any, TypeVar
9
+from functools import wraps
10
+from typing import Any, TypeVar
11
21
-import cloudpickle
22
-import psutil
12
import torch
13
25
-import vllm.envs as envs
26
-from vllm.logger import enable_trace_function_call, init_logger
27
-from vllm.ray.lazy_utils import is_in_ray_actor
14
+from vllm.logger import init_logger
15
16
_DEPRECATED_MAPPINGS = {
17
"cprofile": "profiling",
18
"cprofile_context": "profiling",
19
+ # Used by lm-eval
20
"get_open_port": "network_utils",
21
}
22
@@ -53,12 +41,6 @@ def __dir__() -> list[str]:
41
return sorted(list(globals().keys()) + list(_DEPRECATED_MAPPINGS.keys()))
42
43
56
-if TYPE_CHECKING:
57
- from vllm.config import ModelConfig, VllmConfig
58
-else:
59
- ModelConfig = object
60
- VllmConfig = object
61
-
44
logger = init_logger(__name__)
45
46
# This value is chosen to have a balance between ITL and TTFT. Note it is
@@ -83,13 +65,7 @@ STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"
65
STR_INVALID_VAL: str = "INVALID"
66
67
86
-# ANSI color codes
87
-CYAN = "\033[1;36m"
88
-RESET = "\033[0;0m"
89
-
90
-
68
T = TypeVar("T")
92
-U = TypeVar("U")
69
70
71
class Device(enum.Enum):
@@ -144,195 +120,6 @@ def random_uuid() -> str:
120
return str(uuid.uuid4().hex)
121
122
147
-# TODO: This function can be removed if transformer_modules classes are
148
-# serialized by value when communicating between processes
149
-def init_cached_hf_modules() -> None:
150
- """
151
- Lazy initialization of the Hugging Face modules.
152
- """
153
- from transformers.dynamic_module_utils import init_hf_modules
154
-
155
- init_hf_modules()
156
-
157
-
158
-def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:
159
- """Set up function tracing for the current thread,
160
- if enabled via the VLLM_TRACE_FUNCTION environment variable
161
- """
162
-
163
- if envs.VLLM_TRACE_FUNCTION:
164
- tmp_dir = tempfile.gettempdir()
165
- # add username to tmp_dir to avoid permission issues
166
- tmp_dir = os.path.join(tmp_dir, getpass.getuser())
167
- filename = (
168
- f"VLLM_TRACE_FUNCTION_for_process_{os.getpid()}"
169
- f"_thread_{threading.get_ident()}_"
170
- f"at_{datetime.datetime.now()}.log"
171
- ).replace(" ", "_")
172
- log_path = os.path.join(
173
- tmp_dir, "vllm", f"vllm-instance-{vllm_config.instance_id}", filename
174
- )
175
- os.makedirs(os.path.dirname(log_path), exist_ok=True)
176
- enable_trace_function_call(log_path)
177
-
178
-
179
-def kill_process_tree(pid: int):
180
- """
181
- Kills all descendant processes of the given pid by sending SIGKILL.
182
-
183
- Args:
184
- pid (int): Process ID of the parent process
185
- """
186
- try:
187
- parent = psutil.Process(pid)
188
- except psutil.NoSuchProcess:
189
- return
190
-
191
- # Get all children recursively
192
- children = parent.children(recursive=True)
193
-
194
- # Send SIGKILL to all children first
195
- for child in children:
196
- with contextlib.suppress(ProcessLookupError):
197
- os.kill(child.pid, signal.SIGKILL)
198
-
199
- # Finally kill the parent
200
- with contextlib.suppress(ProcessLookupError):
201
- os.kill(pid, signal.SIGKILL)
202
-
203
-
204
-# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L630 # noqa: E501
205
-def set_ulimit(target_soft_limit=65535):
206
- if sys.platform.startswith("win"):
207
- logger.info("Windows detected, skipping ulimit adjustment.")
208
- return
209
-
210
- import resource
211
-
212
- resource_type = resource.RLIMIT_NOFILE
213
- current_soft, current_hard = resource.getrlimit(resource_type)
214
-
215
- if current_soft < target_soft_limit:
216
- try:
217
- resource.setrlimit(resource_type, (target_soft_limit, current_hard))
218
- except ValueError as e:
219
- logger.warning(
220
- "Found ulimit of %s and failed to automatically increase "
221
- "with error %s. This can cause fd limit errors like "
222
- "`OSError: [Errno 24] Too many open files`. Consider "
223
- "increasing with ulimit -n",
224
- current_soft,
225
- e,
226
- )
227
-
228
-
229
-def _maybe_force_spawn():
230
- """Check if we need to force the use of the `spawn` multiprocessing start
231
- method.
232
- """
233
- if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") == "spawn":
234
- return
235
-
236
- reasons = []
237
- if is_in_ray_actor():
238
- # even if we choose to spawn, we need to pass the ray address
239
- # to the subprocess so that it knows how to connect to the ray cluster.
240
- # env vars are inherited by subprocesses, even if we use spawn.
241
- import ray
242
-
243
- os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
244
- reasons.append("In a Ray actor and can only be spawned")
245
-
246
- from .platform_utils import cuda_is_initialized, xpu_is_initialized
247
-
248
- if cuda_is_initialized():
249
- reasons.append("CUDA is initialized")
250
- elif xpu_is_initialized():
251
- reasons.append("XPU is initialized")
252
-
253
- if reasons:
254
- logger.warning(
255
- "We must use the `spawn` multiprocessing start method. "
256
- "Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "
257
- "See https://docs.vllm.ai/en/latest/usage/"
258
- "troubleshooting.html#python-multiprocessing "
259
- "for more information. Reasons: %s",
260
- "; ".join(reasons),
261
- )
262
- os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
263
-
264
-
265
-def get_mp_context():
266
- """Get a multiprocessing context with a particular method (spawn or fork).
267
- By default we follow the value of the VLLM_WORKER_MULTIPROC_METHOD to
268
- determine the multiprocessing method (default is fork). However, under
269
- certain conditions, we may enforce spawn and override the value of
270
- VLLM_WORKER_MULTIPROC_METHOD.
271
- """
272
- _maybe_force_spawn()
273
- mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD
274
- return multiprocessing.get_context(mp_method)
275
-
276
-
277
-def run_method(
278
- obj: Any,
279
- method: str | bytes | Callable,
280
- args: tuple[Any],
281
- kwargs: dict[str, Any],
282
-) -> Any:
283
- """
284
- Run a method of an object with the given arguments and keyword arguments.
285
- If the method is string, it will be converted to a method using getattr.
286
- If the method is serialized bytes and will be deserialized using
287
- cloudpickle.
288
- If the method is a callable, it will be called directly.
289
- """
290
- if isinstance(method, bytes):
291
- func = partial(cloudpickle.loads(method), obj)
292
- elif isinstance(method, str):
293
- try:
294
- func = getattr(obj, method)
295
- except AttributeError:
296
- raise NotImplementedError(
297
- f"Method {method!r} is not implemented."
298
- ) from None
299
- else:
300
- func = partial(method, obj) # type: ignore
301
- return func(*args, **kwargs)
302
-
303
-
304
-def import_pynvml():
305
- """
306
- Historical comments:
307
-
308
- libnvml.so is the library behind nvidia-smi, and
309
- pynvml is a Python wrapper around it. We use it to get GPU
310
- status without initializing CUDA context in the current process.
311
- Historically, there are two packages that provide pynvml:
312
- - `nvidia-ml-py` (https://pypi.org/project/nvidia-ml-py/): The official
313
- wrapper. It is a dependency of vLLM, and is installed when users
314
- install vLLM. It provides a Python module named `pynvml`.
315
- - `pynvml` (https://pypi.org/project/pynvml/): An unofficial wrapper.
316
- Prior to version 12.0, it also provides a Python module `pynvml`,
317
- and therefore conflicts with the official one. What's worse,
318
- the module is a Python package, and has higher priority than
319
- the official one which is a standalone Python file.
320
- This causes errors when both of them are installed.
321
- Starting from version 12.0, it migrates to a new module
322
- named `pynvml_utils` to avoid the conflict.
323
- It is so confusing that many packages in the community use the
324
- unofficial one by mistake, and we have to handle this case.
325
- For example, `nvcr.io/nvidia/pytorch:24.12-py3` uses the unofficial
326
- one, and it will cause errors, see the issue
327
- https://github.com/vllm-project/vllm/issues/12847 for example.
328
- After all the troubles, we decide to copy the official `pynvml`
329
- module to our codebase, and use it directly.
330
- """
331
- import vllm.third_party.pynvml as pynvml
332
-
333
- return pynvml
334
-
335
-
123
def warn_for_unimplemented_methods(cls: type[T]) -> type[T]:
124
"""
125
A replacement for `abc.ABC`.
@@ -376,31 +163,6 @@ def warn_for_unimplemented_methods(cls: type[T]) -> type[T]:
163
return cls
164
165
379
-# Only relevant for models using ALiBi (e.g, MPT)
380
-def check_use_alibi(model_config: ModelConfig) -> bool:
381
- cfg = model_config.hf_text_config
382
- return (
383
- getattr(cfg, "alibi", False) # Falcon
384
- or (
385
- "BloomForCausalLM" in getattr(model_config.hf_config, "architectures", [])
386
- ) # Bloom
387
- or getattr(cfg, "position_encoding_type", "") == "alibi" # codellm_1b_alibi
388
- or (
389
- hasattr(cfg, "attn_config") # MPT
390
- and (
391
- (
392
- isinstance(cfg.attn_config, dict)
393
- and cfg.attn_config.get("alibi", False)
394
- )
395
- or (
396
- not isinstance(cfg.attn_config, dict)
397
- and getattr(cfg.attn_config, "alibi", False)
398
- )
399
- )
400
- )
401
- )
402
-
403
-
166
def length_from_prompt_token_ids_or_embeds(
167
prompt_token_ids: list[int] | None,
168
prompt_embeds: torch.Tensor | None,
vllm/utils/argparse_utils.py
+3
-23
@@ -10,37 +10,21 @@ from argparse import (
10
ArgumentDefaultsHelpFormatter,
11
ArgumentParser,
12
ArgumentTypeError,
13
+ Namespace,
14
RawDescriptionHelpFormatter,
15
_ArgumentGroup,
16
)
17
from collections import defaultdict
17
-from typing import TYPE_CHECKING, Any
18
+from typing import Any
19
20
import regex as re
21
import yaml
22
23
from vllm.logger import init_logger
24
24
-if TYPE_CHECKING:
25
- from argparse import Namespace
26
-else:
27
- Namespace = object
28
-
25
logger = init_logger(__name__)
26
27
32
-class StoreBoolean(Action):
33
- def __call__(self, parser, namespace, values, option_string=None):
34
- if values.lower() == "true":
35
- setattr(namespace, self.dest, True)
36
- elif values.lower() == "false":
37
- setattr(namespace, self.dest, False)
38
- else:
39
- raise ValueError(
40
- f"Invalid boolean value: {values}. Expected 'true' or 'false'."
41
- )
42
-
43
-
28
class SortedHelpFormatter(ArgumentDefaultsHelpFormatter, RawDescriptionHelpFormatter):
29
"""SortedHelpFormatter that sorts arguments by their option strings."""
30
@@ -487,12 +471,8 @@ class FlexibleArgumentParser(ArgumentParser):
471
)
472
raise ex
473
490
- store_boolean_arguments = [
491
- action.dest for action in self._actions if isinstance(action, StoreBoolean)
492
- ]
493
-
474
for key, value in config.items():
495
- if isinstance(value, bool) and key not in store_boolean_arguments:
475
+ if isinstance(value, bool):
476
if value:
477
processed_args.append("--" + key)
478
elif isinstance(value, list):
vllm/utils/import_utils.py
+43
@@ -19,6 +19,49 @@ import regex as re
19
from typing_extensions import Never
20
21
22
+# TODO: This function can be removed if transformer_modules classes are
23
+# serialized by value when communicating between processes
24
+def init_cached_hf_modules() -> None:
25
+ """
26
+ Lazy initialization of the Hugging Face modules.
27
+ """
28
+ from transformers.dynamic_module_utils import init_hf_modules
29
+
30
+ init_hf_modules()
31
+
32
+
33
+def import_pynvml():
34
+ """
35
+ Historical comments:
36
+
37
+ libnvml.so is the library behind nvidia-smi, and
38
+ pynvml is a Python wrapper around it. We use it to get GPU
39
+ status without initializing CUDA context in the current process.
40
+ Historically, there are two packages that provide pynvml:
41
+ - `nvidia-ml-py` (https://pypi.org/project/nvidia-ml-py/): The official
42
+ wrapper. It is a dependency of vLLM, and is installed when users
43
+ install vLLM. It provides a Python module named `pynvml`.
44
+ - `pynvml` (https://pypi.org/project/pynvml/): An unofficial wrapper.
45
+ Prior to version 12.0, it also provides a Python module `pynvml`,
46
+ and therefore conflicts with the official one. What's worse,
47
+ the module is a Python package, and has higher priority than
48
+ the official one which is a standalone Python file.
49
+ This causes errors when both of them are installed.
50
+ Starting from version 12.0, it migrates to a new module
51
+ named `pynvml_utils` to avoid the conflict.
52
+ It is so confusing that many packages in the community use the
53
+ unofficial one by mistake, and we have to handle this case.
54
+ For example, `nvcr.io/nvidia/pytorch:24.12-py3` uses the unofficial
55
+ one, and it will cause errors, see the issue
56
+ https://github.com/vllm-project/vllm/issues/12847 for example.
57
+ After all the troubles, we decide to copy the official `pynvml`
58
+ module to our codebase, and use it directly.
59
+ """
60
+ import vllm.third_party.pynvml as pynvml
61
+
62
+ return pynvml
63
+
64
+
65
def import_from_path(module_name: str, file_path: str | os.PathLike):
66
"""
67
Import a Python file according to its file path.
vllm/utils/system_utils.py
+114
-8
@@ -4,19 +4,21 @@
4
from __future__ import annotations
5
6
import contextlib
7
+import multiprocessing
8
import os
9
+import signal
10
import sys
11
from collections.abc import Callable, Iterator
12
from pathlib import Path
13
from typing import TextIO
14
13
-try:
14
- import setproctitle
15
-except ImportError:
16
- setproctitle = None # type: ignore[assignment]
15
+import psutil
16
17
import vllm.envs as envs
18
from vllm.logger import init_logger
19
+from vllm.ray.lazy_utils import is_in_ray_actor
20
+
21
+from .platform_utils import cuda_is_initialized, xpu_is_initialized
22
23
logger = init_logger(__name__)
24
@@ -75,14 +77,66 @@ def unique_filepath(fn: Callable[[int], Path]) -> Path:
77
# Process management utilities
78
79
80
+def _maybe_force_spawn():
81
+ """Check if we need to force the use of the `spawn` multiprocessing start
82
+ method.
83
+ """
84
+ if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") == "spawn":
85
+ return
86
+
87
+ reasons = []
88
+ if is_in_ray_actor():
89
+ # even if we choose to spawn, we need to pass the ray address
90
+ # to the subprocess so that it knows how to connect to the ray cluster.
91
+ # env vars are inherited by subprocesses, even if we use spawn.
92
+ import ray
93
+
94
+ os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
95
+ reasons.append("In a Ray actor and can only be spawned")
96
+
97
+ if cuda_is_initialized():
98
+ reasons.append("CUDA is initialized")
99
+ elif xpu_is_initialized():
100
+ reasons.append("XPU is initialized")
101
+
102
+ if reasons:
103
+ logger.warning(
104
+ "We must use the `spawn` multiprocessing start method. "
105
+ "Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "
106
+ "See https://docs.vllm.ai/en/latest/usage/"
107
+ "troubleshooting.html#python-multiprocessing "
108
+ "for more information. Reasons: %s",
109
+ "; ".join(reasons),
110
+ )
111
+ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
112
+
113
+
114
+def get_mp_context():
115
+ """Get a multiprocessing context with a particular method (spawn or fork).
116
+ By default we follow the value of the VLLM_WORKER_MULTIPROC_METHOD to
117
+ determine the multiprocessing method (default is fork). However, under
118
+ certain conditions, we may enforce spawn and override the value of
119
+ VLLM_WORKER_MULTIPROC_METHOD.
120
+ """
121
+ _maybe_force_spawn()
122
+ mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD
123
+ return multiprocessing.get_context(mp_method)
124
+
125
+
126
def set_process_title(
79
- name: str, suffix: str = "", prefix: str = envs.VLLM_PROCESS_NAME_PREFIX
127
+ name: str,
128
+ suffix: str = "",
129
+ prefix: str = envs.VLLM_PROCESS_NAME_PREFIX,
130
) -> None:
131
"""Set the current process title with optional suffix."""
82
- if setproctitle is None:
132
+ try:
133
+ import setproctitle
134
+ except ImportError:
135
return
136
+
137
if suffix:
138
name = f"{name}_{suffix}"
139
+
140
setproctitle.setproctitle(f"{prefix}::{name}")
141
142
@@ -114,10 +168,62 @@ def _add_prefix(file: TextIO, worker_name: str, pid: int) -> None:
168
169
def decorate_logs(process_name: str | None = None) -> None:
170
"""Decorate stdout/stderr with process name and PID prefix."""
117
- from vllm.utils import get_mp_context
118
-
171
if process_name is None:
172
process_name = get_mp_context().current_process().name
173
+
174
pid = os.getpid()
175
_add_prefix(sys.stdout, process_name, pid)
176
_add_prefix(sys.stderr, process_name, pid)
177
+
178
+
179
+def kill_process_tree(pid: int):
180
+ """
181
+ Kills all descendant processes of the given pid by sending SIGKILL.
182
+
183
+ Args:
184
+ pid (int): Process ID of the parent process
185
+ """
186
+ try:
187
+ parent = psutil.Process(pid)
188
+ except psutil.NoSuchProcess:
189
+ return
190
+
191
+ # Get all children recursively
192
+ children = parent.children(recursive=True)
193
+
194
+ # Send SIGKILL to all children first
195
+ for child in children:
196
+ with contextlib.suppress(ProcessLookupError):
197
+ os.kill(child.pid, signal.SIGKILL)
198
+
199
+ # Finally kill the parent
200
+ with contextlib.suppress(ProcessLookupError):
201
+ os.kill(pid, signal.SIGKILL)
202
+
203
+
204
+# Resource utilities
205
+
206
+
207
+# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L630
208
+def set_ulimit(target_soft_limit: int = 65535):
209
+ if sys.platform.startswith("win"):
210
+ logger.info("Windows detected, skipping ulimit adjustment.")
211
+ return
212
+
213
+ import resource
214
+
215
+ resource_type = resource.RLIMIT_NOFILE
216
+ current_soft, current_hard = resource.getrlimit(resource_type)
217
+
218
+ if current_soft < target_soft_limit:
219
+ try:
220
+ resource.setrlimit(resource_type, (target_soft_limit, current_hard))
221
+ except ValueError as e:
222
+ logger.warning(
223
+ "Found ulimit of %s and failed to automatically increase "
224
+ "with error %s. This can cause fd limit errors like "
225
+ "`OSError: [Errno 24] Too many open files`. Consider "
226
+ "increasing with ulimit -n",
227
+ current_soft,
228
+ e,
229
+ )
vllm/v1/engine/coordinator.py
+1
-2
@@ -10,9 +10,8 @@ import zmq
10
11
from vllm.config import ParallelConfig
12
from vllm.logger import init_logger
13
-from vllm.utils import get_mp_context
13
from vllm.utils.network_utils import make_zmq_socket
15
-from vllm.utils.system_utils import set_process_title
14
+from vllm.utils.system_utils import get_mp_context, set_process_title
15
from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequestType
16
from vllm.v1.serial_utils import MsgpackDecoder
17
from vllm.v1.utils import get_engine_client_zmq_addr, shutdown
vllm/v1/engine/utils.py
+1
-1
@@ -20,8 +20,8 @@ from vllm.config import CacheConfig, ParallelConfig, VllmConfig
20
from vllm.logger import init_logger
21
from vllm.platforms import current_platform
22
from vllm.ray.ray_env import get_env_vars_to_copy
23
-from vllm.utils import get_mp_context
23
from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
24
+from vllm.utils.system_utils import get_mp_context
25
from vllm.v1.engine.coordinator import DPCoordinator
26
from vllm.v1.executor import Executor
27
from vllm.v1.utils import get_engine_client_zmq_addr, shutdown
vllm/v1/executor/multiproc_executor.py
+6
-2
@@ -35,13 +35,17 @@ 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 _maybe_force_spawn, get_mp_context
38
from vllm.utils.network_utils import (
39
get_distributed_init_method,
40
get_loopback_ip,
41
get_open_port,
42
)
44
-from vllm.utils.system_utils import decorate_logs, set_process_title
43
+from vllm.utils.system_utils import (
44
+ _maybe_force_spawn,
45
+ decorate_logs,
46
+ get_mp_context,
47
+ set_process_title,
48
+)
49
from vllm.v1.core.sched.output import SchedulerOutput
50
from vllm.v1.executor.abstract import Executor, FailureCallback
51
from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput
vllm/v1/executor/uniproc_executor.py
+1
-1
@@ -12,11 +12,11 @@ import torch.distributed as dist
12
13
import vllm.envs as envs
14
from vllm.logger import init_logger
15
-from vllm.utils import run_method
15
from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
16
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
17
from vllm.v1.executor.abstract import Executor
18
from vllm.v1.outputs import AsyncModelRunnerOutput
19
+from vllm.v1.serial_utils import run_method
20
from vllm.v1.worker.worker_base import WorkerWrapperBase
21
22
logger = init_logger(__name__)
vllm/v1/serial_utils.py
+28
@@ -5,6 +5,7 @@ import dataclasses
5
import importlib
6
import pickle
7
from collections.abc import Callable, Sequence
8
+from functools import partial
9
from inspect import isclass
10
from types import FunctionType
11
from typing import Any, TypeAlias
@@ -429,3 +430,30 @@ class MsgpackDecoder:
430
return cloudpickle.loads(data)
431
432
raise NotImplementedError(f"Extension type code {code} is not supported")
433
+
434
+
435
+def run_method(
436
+ obj: Any,
437
+ method: str | bytes | Callable,
438
+ args: tuple[Any, ...],
439
+ kwargs: dict[str, Any],
440
+) -> Any:
441
+ """
442
+ Run a method of an object with the given arguments and keyword arguments.
443
+ If the method is string, it will be converted to a method using getattr.
444
+ If the method is serialized bytes and will be deserialized using
445
+ cloudpickle.
446
+ If the method is a callable, it will be called directly.
447
+ """
448
+ if isinstance(method, bytes):
449
+ func = partial(cloudpickle.loads(method), obj)
450
+ elif isinstance(method, str):
451
+ try:
452
+ func = getattr(obj, method)
453
+ except AttributeError:
454
+ raise NotImplementedError(
455
+ f"Method {method!r} is not implemented."
456
+ ) from None
457
+ else:
458
+ func = partial(method, obj) # type: ignore
459
+ return func(*args, **kwargs)
vllm/v1/utils.py
+1
-1
@@ -25,8 +25,8 @@ from torch.autograd.profiler import record_function
25
import vllm.envs as envs
26
from vllm.logger import init_logger
27
from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message
28
-from vllm.utils import kill_process_tree
28
from vllm.utils.network_utils import get_open_port, get_open_zmq_ipc_path, get_tcp_uri
29
+from vllm.utils.system_utils import kill_process_tree
30
31
if TYPE_CHECKING:
32
import numpy as np
vllm/v1/worker/gpu_model_runner.py
+2
-5
@@ -69,10 +69,7 @@ from vllm.pooling_params import PoolingParams
69
from vllm.sampling_params import SamplingType
70
from vllm.sequence import IntermediateTensors
71
from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
72
-from vllm.utils import (
73
- check_use_alibi,
74
- length_from_prompt_token_ids_or_embeds,
75
-)
72
+from vllm.utils import length_from_prompt_token_ids_or_embeds
73
from vllm.utils.jsontree import json_map_leaves
74
from vllm.utils.math_utils import cdiv, round_up
75
from vllm.utils.mem_constants import GiB_bytes
@@ -266,7 +263,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
263
self.hidden_size = model_config.get_hidden_size()
264
self.attention_chunk_size = model_config.attention_chunk_size
265
# Only relevant for models using ALiBi (e.g, MPT)
269
- self.use_alibi = check_use_alibi(model_config)
266
+ self.use_alibi = model_config.uses_alibi
267
268
self.cascade_attn_enabled = not self.model_config.disable_cascade_attn
269
vllm/v1/worker/gpu_worker.py
+1
-1
@@ -72,7 +72,7 @@ class Worker(WorkerBase):
72
73
if self.model_config.trust_remote_code:
74
# note: lazy import to avoid importing torch before initializing
75
- from vllm.utils import init_cached_hf_modules
75
+ from vllm.utils.import_utils import init_cached_hf_modules
76
77
init_cached_hf_modules()
78
vllm/v1/worker/tpu_worker.py
+1
-1
@@ -89,7 +89,7 @@ class TPUWorker:
89
90
if self.model_config.trust_remote_code:
91
# note: lazy import to avoid importing torch before initializing
92
- from vllm.utils import init_cached_hf_modules
92
+ from vllm.utils.import_utils import init_cached_hf_modules
93
94
init_cached_hf_modules()
95
vllm/v1/worker/worker_base.py
+16
-18
@@ -13,14 +13,11 @@ from vllm.logger import init_logger
13
from vllm.lora.request import LoRARequest
14
from vllm.multimodal import MULTIMODAL_REGISTRY
15
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
- warn_for_unimplemented_methods,
20
-)
16
+from vllm.utils import warn_for_unimplemented_methods
17
from vllm.utils.import_utils import resolve_obj_by_qualname
18
from vllm.utils.system_utils import update_environment_variables
19
from vllm.v1.kv_cache_interface import KVCacheSpec
20
+from vllm.v1.serial_utils import run_method
21
22
if TYPE_CHECKING:
23
from vllm.v1.core.sched.output import SchedulerOutput
@@ -182,19 +179,20 @@ class WorkerWrapperBase:
179
"""
180
self.rpc_rank = rpc_rank
181
self.worker: WorkerBase | None = None
185
- self.vllm_config: VllmConfig | None = None
182
+
183
# do not store this `vllm_config`, `init_worker` will set the final
187
- # one. TODO: investigate if we can remove this field in
188
- # `WorkerWrapperBase`, `init_cached_hf_modules` should be
189
- # unnecessary now.
190
- if vllm_config.model_config is not None:
191
- # it can be None in tests
192
- trust_remote_code = vllm_config.model_config.trust_remote_code
193
- if trust_remote_code:
194
- # note: lazy import to avoid importing torch before initializing
195
- from vllm.utils import init_cached_hf_modules
196
-
197
- init_cached_hf_modules()
184
+ # one.
185
+ # TODO: investigate if we can remove this field in `WorkerWrapperBase`,
186
+ # `init_cached_hf_modules` should be unnecessary now.
187
+ self.vllm_config: VllmConfig | None = None
188
+
189
+ # `model_config` can be None in tests
190
+ model_config = vllm_config.model_config
191
+ if model_config and model_config.trust_remote_code:
192
+ # note: lazy import to avoid importing torch before initializing
193
+ from vllm.utils.import_utils import init_cached_hf_modules
194
+
195
+ init_cached_hf_modules()
196
197
def shutdown(self) -> None:
198
if self.worker is not None:
@@ -231,7 +229,7 @@ class WorkerWrapperBase:
229
assert self.vllm_config is not None, (
230
"vllm_config is required to initialize the worker"
231
)
234
- enable_trace_function_call_for_thread(self.vllm_config)
232
+ self.vllm_config.enable_trace_function_call_for_thread()
233
234
from vllm.plugins import load_general_plugins
235