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 4, 01:15 UTC
[Bugfix] Fix LOGITPROC_SOURCE_ENTRYPOINT test to use spawn-compatible dist-info registration for XPU/ROCm (#42040)
Signed-off-by: dqzhengAP <dqzheng1996@gmail.com> Signed-off-by: David Zheng <153074367+dzhengAP@users.noreply.github.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
David Zheng committed
May 8, 2026 at 21:32 UTC
df2636a9d8139dd1d53ec1f1e666abe66fb76a3c
4 files changed
+129
-48
tests/utils.py
+56
-11
@@ -20,9 +20,9 @@ import time
20
import warnings
21
from collections.abc import Callable, Iterable, Sequence
22
from contextlib import ExitStack, contextmanager
23
-from multiprocessing import Process
23
+from multiprocessing import Process, get_context
24
from pathlib import Path
25
-from typing import Any, Literal
25
+from typing import Any, Literal, cast
26
from unittest.mock import patch
27
28
import anthropic
@@ -127,6 +127,21 @@ ROCM_ENGINE_KWARGS: dict = (
127
)
128
129
130
+def requires_spawn_multiprocessing() -> bool:
131
+ """Whether this platform requires spawn instead of fork for test processes."""
132
+ return current_platform.is_rocm() or current_platform.is_xpu()
133
+
134
+
135
+def _run_in_new_process_group(
136
+ child_process_fxn: Callable[[dict[str, str] | None, str, list[str]], None],
137
+ env_dict: dict[str, str] | None,
138
+ model: str,
139
+ vllm_serve_args: list[str],
140
+) -> None:
141
+ os.setsid()
142
+ child_process_fxn(env_dict, model, vllm_serve_args)
143
+
144
+
145
class RemoteVLLMServer:
146
"""Base class for launching vLLM server subprocesses for testing.
147
@@ -738,8 +753,11 @@ class RemoteOpenAIServerCustom(RemoteOpenAIServer):
753
def _start_server(
754
self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
755
) -> None:
741
- self.proc: Process = Process(
742
- target=self.child_process_fxn, args=(env_dict, model, vllm_serve_args)
756
+ method = "spawn" if requires_spawn_multiprocessing() else "fork"
757
+ ctx = get_context(method)
758
+ self.proc: Process = cast(Any, ctx).Process(
759
+ target=_run_in_new_process_group,
760
+ args=(self.child_process_fxn, env_dict, model, vllm_serve_args),
761
) # type: ignore[assignment]
762
self.proc.start()
763
@@ -769,12 +787,40 @@ class RemoteOpenAIServerCustom(RemoteOpenAIServer):
787
def _poll(self) -> int | None:
788
return self.proc.exitcode
789
772
- def __exit__(self, exc_type, exc_value, traceback):
773
- self.proc.terminate()
774
- self.proc.join(8)
790
+ def _terminate_process_tree(self) -> None:
791
+ pid = self.proc.pid
792
+ if pid is None:
793
+ return
794
+
795
+ pgid: int | None
796
+ try:
797
+ pgid = os.getpgid(pid)
798
+ # _run_in_new_process_group should make the child the group
799
+ # leader. Avoid signaling pytest's process group if startup failed
800
+ # before os.setsid() ran.
801
+ if pgid != pid:
802
+ pgid = None
803
+ except (ProcessLookupError, OSError):
804
+ pgid = None
805
+
806
+ with contextlib.suppress(ProcessLookupError, OSError):
807
+ self.proc.terminate()
808
+ print(f"[RemoteOpenAIServerCustom] Sent SIGTERM to process {pid}")
809
+
810
+ self.proc.join(15)
811
if self.proc.is_alive():
776
- # force kill if needed
777
- self.proc.kill()
812
+ print(
813
+ f"[RemoteOpenAIServerCustom] Server {pid} did not respond "
814
+ "to SIGTERM, sending SIGKILL to process group"
815
+ )
816
+ if pgid is not None:
817
+ with contextlib.suppress(ProcessLookupError, OSError):
818
+ os.killpg(pgid, signal.SIGKILL)
819
+ else:
820
+ self.proc.kill()
821
+ self.proc.join(10)
822
+
823
+ self._kill_process_group_survivors(pgid)
824
825
826
def _test_completion(
@@ -1633,8 +1679,7 @@ def create_new_process_for_each_test(
1679
A decorator to run test functions in separate processes.
1680
"""
1681
if method is None:
1636
- use_spawn = current_platform.is_rocm() or current_platform.is_xpu()
1637
- method = "spawn" if use_spawn else "fork"
1682
+ method = "spawn" if requires_spawn_multiprocessing() else "fork"
1683
1684
assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
1685
tests/v1/logits_processors/test_custom_offline.py
+7
-16
@@ -16,8 +16,8 @@ from tests.v1.logits_processors.utils import (
16
DummyLogitsProcessor,
17
WrappedPerReqLogitsProcessor,
18
prompts,
19
+ setup_fake_entrypoint,
20
)
20
-from tests.v1.logits_processors.utils import entry_points as fake_entry_points
21
from vllm import LLM, SamplingParams
22
from vllm.v1.sample.logits_processor import (
23
STR_POOLING_REJECTS_LOGITSPROCS,
@@ -145,13 +145,9 @@ def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource
145
146
if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_ENTRYPOINT:
147
# Scenario: vLLM loads a logitproc from a preconfigured entrypoint
148
- # To that end, mock a dummy logitproc entrypoint
149
- import importlib.metadata
150
-
151
- importlib.metadata.entry_points = fake_entry_points # type: ignore
152
-
153
- # fork is required for workers to see entrypoint patch
154
- monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork")
148
+ # To that end, register a real dist-info package so spawned
149
+ # workers can discover the entrypoint via PYTHONPATH
150
+ setup_fake_entrypoint(monkeypatch)
151
_run_test({}, logitproc_loaded=True)
152
return
153
@@ -266,14 +262,9 @@ def test_rejects_custom_logitsprocs(
262
# Scenario: vLLM loads a model and ignores a logitproc that is
263
# available at a preconfigured entrypoint
264
269
- # Patch in dummy logitproc entrypoint
270
- import importlib.metadata
271
-
272
- importlib.metadata.entry_points = fake_entry_points # type: ignore
273
-
274
- # fork is required for entrypoint patch to be visible to workers,
275
- # although they should ignore the entrypoint patch anyway
276
- monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork")
265
+ # Register real dist-info package so spawned workers can
266
+ # discover the entrypoint via PYTHONPATH (spawn-compatible)
267
+ setup_fake_entrypoint(monkeypatch)
268
269
llm = LLM(**llm_kwargs)
270
# Require that no custom logitsprocs have been loaded
tests/v1/logits_processors/test_custom_online.py
+6
-13
@@ -18,8 +18,8 @@ from tests.v1.logits_processors.utils import (
18
MODEL_NAME,
19
TEMP_GREEDY,
20
prompts,
21
+ setup_fake_entrypoint,
22
)
22
-from tests.v1.logits_processors.utils import entry_points as fake_entry_points
23
24
25
def _server_with_logitproc_entrypoint(
@@ -27,16 +27,9 @@ def _server_with_logitproc_entrypoint(
27
model: str,
28
vllm_serve_args: list[str],
29
) -> None:
30
- """Start vLLM server, inject dummy logitproc entrypoint"""
31
-
32
- # Patch `entry_points` to inject logitproc entrypoint
33
- import importlib.metadata
34
-
35
- importlib.metadata.entry_points = fake_entry_points # type: ignore
30
+ """Start vLLM server with dummy logitproc entrypoint."""
31
from vllm.entrypoints.cli import main
32
38
- # fork is required for workers to see entrypoint patch
39
- os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "fork"
33
if env_dict is not None:
34
os.environ.update(env_dict)
35
@@ -50,7 +43,7 @@ def _server_with_logitproc_fqcn(
43
model: str,
44
vllm_serve_args: list[str],
45
) -> None:
53
- """Start vLLM server, inject module with dummy logitproc"""
46
+ """Start vLLM server with dummy logitproc specified by FQCN."""
47
from vllm.entrypoints.cli import main
48
49
if env_dict is not None:
@@ -80,8 +73,8 @@ def default_server_args():
73
def server(default_server_args, request, monkeypatch):
74
"""Consider two server configurations:
75
(1) --logits-processors cli arg specifies dummy logits processor via fully-
83
- qualified class name (FQCN); patch in a dummy logits processor module
84
- (2) No --logits-processors cli arg; patch in a dummy logits processor
76
+ qualified class name (FQCN)
77
+ (2) No --logits-processors cli arg; inject a dummy logits processor
78
entrypoint
79
"""
80
@@ -94,6 +87,7 @@ def server(default_server_args, request, monkeypatch):
87
_server_fxn = _server_with_logitproc_fqcn
88
else:
89
# Launch server, inject dummy logitproc entrypoint
90
+ setup_fake_entrypoint(monkeypatch)
91
args = default_server_args
92
_server_fxn = _server_with_logitproc_entrypoint
93
@@ -119,7 +113,6 @@ api_keyword_args = {
113
}
114
115
122
-@pytest.mark.asyncio
116
@pytest.mark.parametrize(
117
"model_name",
118
[MODEL_NAME],
tests/v1/logits_processors/utils.py
+60
-8
@@ -1,12 +1,15 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import types
4
+import os
5
+import tempfile
6
from enum import Enum, auto
7
+from pathlib import Path
8
from typing import Any
9
10
import torch
11
12
+from tests.utils import requires_spawn_multiprocessing
13
from vllm.config import VllmConfig
14
from vllm.logger import init_logger
15
from vllm.sampling_params import SamplingParams
@@ -102,11 +105,6 @@ class DummyLogitsProcessor(LogitsProcessor):
105
return logits
106
107
105
-"""Dummy module with dummy logitproc class"""
106
-dummy_module = types.ModuleType(DUMMY_LOGITPROC_MODULE)
107
-dummy_module.DummyLogitsProcessor = DummyLogitsProcessor # type: ignore
108
-
109
-
108
class EntryPoint:
109
"""Dummy entrypoint class for logitsprocs testing"""
110
@@ -187,5 +185,59 @@ class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
185
return DummyPerReqLogitsProcessor(target_token)
186
187
190
-"""Fake version of importlib.metadata.entry_points"""
191
-entry_points = lambda group: EntryPoints(group)
188
+def register_fake_entrypoint(monkeypatch) -> str:
189
+ """Register the dummy logitsproc entrypoint in a way that is visible
190
+ to spawned subprocesses by creating a real dist-info directory on disk.
191
+
192
+ Unlike monkey-patching importlib.metadata.entry_points (which only works
193
+ with fork), this approach writes a real dist-info package that
194
+ importlib.metadata can discover in any subprocess via PYTHONPATH.
195
+
196
+ Returns the temp directory path.
197
+ """
198
+ tmpdir = Path(tempfile.mkdtemp(prefix="dummy-logitproc-"))
199
+ dist_info = tmpdir / "dummy_logitproc-0.1.dist-info"
200
+ dist_info.mkdir()
201
+
202
+ # Write METADATA file (required by importlib.metadata)
203
+ (dist_info / "METADATA").write_text(
204
+ "Metadata-Version: 2.1\nName: dummy-logitproc\nVersion: 0.1\n",
205
+ encoding="utf-8",
206
+ )
207
+
208
+ # Write entry_points.txt
209
+ (dist_info / "entry_points.txt").write_text(
210
+ f"[{LOGITSPROCS_GROUP}]\n"
211
+ f"{DUMMY_LOGITPROC_ENTRYPOINT} = {DUMMY_LOGITPROC_FQCN}\n",
212
+ encoding="utf-8",
213
+ )
214
+
215
+ # Add to PYTHONPATH so spawned subprocesses can discover it
216
+ existing = os.environ.get("PYTHONPATH", "")
217
+ monkeypatch.setenv(
218
+ "PYTHONPATH", str(tmpdir) + (os.pathsep + existing if existing else "")
219
+ )
220
+
221
+ # Also update sys.path for the current process so the driver can
222
+ # discover the entrypoint.
223
+ monkeypatch.syspath_prepend(str(tmpdir))
224
+
225
+ return str(tmpdir)
226
+
227
+
228
+def fake_entry_points(group: str) -> EntryPoints:
229
+ """Fake version of importlib.metadata.entry_points."""
230
+ return EntryPoints(group)
231
+
232
+
233
+def setup_fake_entrypoint(monkeypatch) -> None:
234
+ """Expose the dummy logitproc entrypoint for the current platform."""
235
+ if requires_spawn_multiprocessing():
236
+ register_fake_entrypoint(monkeypatch)
237
+ monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
238
+ return
239
+
240
+ import importlib.metadata
241
+
242
+ monkeypatch.setattr(importlib.metadata, "entry_points", fake_entry_points)
243
+ monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork")