@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 158 lines 5.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import builtins
4 import logging
5 from unittest.mock import mock_open, patch
6
7 import pytest
8
9 from vllm.platforms import (
10 _is_amd_zen_cpu,
11 cpu_platform_plugin,
12 resolve_current_platform_cls_qualname,
13 )
14
15
16 @pytest.fixture
17 def caplog_vllm(caplog):
18 """`caplog`, but also captures vLLM's loggers.
19
20 vllm/logger.py configures the "vllm" logger with propagate=False, so
21 records from child loggers like "vllm.platforms" reach vLLM's own
22 handler but never bubble up to the root logger where plain `caplog`
23 listens. Attach caplog's handler directly to "vllm" to see them too.
24 """
25 logger = logging.getLogger("vllm")
26 logger.addHandler(caplog.handler)
27 try:
28 yield caplog
29 finally:
30 logger.removeHandler(caplog.handler)
31
32
33 def test_is_amd_zen_cpu_detects_amd_with_avx512():
34 cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2 avx512f avx512bw"
35 with (
36 patch("os.path.exists", return_value=True),
37 patch("builtins.open", mock_open(read_data=cpuinfo)),
38 ):
39 assert _is_amd_zen_cpu()
40
41
42 def test_is_amd_zen_cpu_returns_false_for_amd_without_avx512():
43 cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2"
44 with (
45 patch("os.path.exists", return_value=True),
46 patch("builtins.open", mock_open(read_data=cpuinfo)),
47 ):
48 assert not _is_amd_zen_cpu()
49
50
51 def test_is_amd_zen_cpu_returns_false_for_intel_with_avx512():
52 cpuinfo = "vendor_id: GenuineIntel\nflags: avx avx2 avx512f"
53 with (
54 patch("os.path.exists", return_value=True),
55 patch("builtins.open", mock_open(read_data=cpuinfo)),
56 ):
57 assert not _is_amd_zen_cpu()
58
59
60 def test_is_amd_zen_cpu_returns_false_when_cpuinfo_missing():
61 with patch("os.path.exists", return_value=False):
62 assert not _is_amd_zen_cpu()
63
64
65 def test_cpu_target_selects_cpu_platform_from_non_cpu_wheel(
66 monkeypatch: pytest.MonkeyPatch,
67 ):
68 monkeypatch.setenv("VLLM_TARGET_DEVICE", "cpu")
69
70 with (
71 patch("vllm.platforms.vllm_version_matches_substr") as version_matches,
72 patch("vllm.platforms._is_amd_zen_cpu", return_value=False),
73 patch("vllm.platforms.rocm_platform_plugin") as rocm_plugin,
74 ):
75 assert (
76 resolve_current_platform_cls_qualname() == "vllm.platforms.cpu.CpuPlatform"
77 )
78
79 # An explicit target does not depend on the installed wheel's version
80 # suffix or host accelerators (a native CI job can reuse a ROCm wheel).
81 version_matches.assert_not_called()
82 rocm_plugin.assert_not_called()
83
84
85 def test_broken_zentorch_falls_back_to_cpu_platform(caplog_vllm):
86 original_import = builtins.__import__
87
88 def import_with_broken_zentorch(name, *args, **kwargs):
89 if name == "zentorch":
90 raise OSError("incompatible shared library")
91 return original_import(name, *args, **kwargs)
92
93 with (
94 patch("vllm.platforms.envs.VLLM_TARGET_DEVICE", "cuda"),
95 patch.dict(
96 "vllm.platforms.builtin_platform_plugins",
97 {"cpu": cpu_platform_plugin},
98 clear=True,
99 ),
100 patch("vllm.platforms.load_plugins_by_group", return_value={}),
101 patch("vllm.platforms.vllm_version_matches_substr", return_value=True),
102 patch("vllm.platforms._is_amd_zen_cpu", return_value=True),
103 patch.object(builtins, "__import__", side_effect=import_with_broken_zentorch),
104 caplog_vllm.at_level(logging.DEBUG, logger="vllm.platforms"),
105 ):
106 platform = resolve_current_platform_cls_qualname()
107
108 assert platform == "vllm.platforms.cpu.CpuPlatform"
109 assert "zentorch failed to import" in caplog_vllm.text
110 assert "OSError: incompatible shared library" in caplog_vllm.text
111
112
113 def test_zentorch_failure_other_than_oserror_is_not_recovered(caplog_vllm):
114 original_import = builtins.__import__
115
116 def import_with_broken_zentorch(name, *args, **kwargs):
117 if name == "zentorch":
118 raise RuntimeError("unknown zentorch failure")
119 return original_import(name, *args, **kwargs)
120
121 with (
122 patch("vllm.platforms.envs.VLLM_TARGET_DEVICE", "cuda"),
123 patch.dict(
124 "vllm.platforms.builtin_platform_plugins",
125 {"cpu": cpu_platform_plugin},
126 clear=True,
127 ),
128 patch("vllm.platforms.load_plugins_by_group", return_value={}),
129 patch("vllm.platforms.vllm_version_matches_substr", return_value=True),
130 patch("vllm.platforms._is_amd_zen_cpu", return_value=True),
131 patch.object(builtins, "__import__", side_effect=import_with_broken_zentorch),
132 caplog_vllm.at_level(logging.DEBUG, logger="vllm.platforms"),
133 ):
134 platform = resolve_current_platform_cls_qualname()
135
136 assert platform == "vllm.platforms.interface.UnspecifiedPlatform"
137 assert "RuntimeError: unknown zentorch failure" in caplog_vllm.text
138
139
140 def test_platform_plugin_failure_is_logged(caplog_vllm):
141 def failing_plugin():
142 raise RuntimeError("plugin exploded")
143
144 with (
145 patch("vllm.platforms.envs.VLLM_TARGET_DEVICE", "cuda"),
146 patch.dict(
147 "vllm.platforms.builtin_platform_plugins",
148 {"cpu": failing_plugin},
149 clear=True,
150 ),
151 patch("vllm.platforms.load_plugins_by_group", return_value={}),
152 caplog_vllm.at_level(logging.DEBUG, logger="vllm.platforms"),
153 ):
154 platform = resolve_current_platform_cls_qualname()
155
156 assert platform == "vllm.platforms.interface.UnspecifiedPlatform"
157 assert "Platform plugin cpu failed during detection" in caplog_vllm.text
158 assert "RuntimeError: plugin exploded" in caplog_vllm.text