test_platform_plugins.py 1.58 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import pytest
5
6
import torch

7
from vllm.plugins import load_general_plugins
8
9


10
11
12
def test_platform_plugins():
    # simulate workload by running an example
    import runpy
13

14
15
    current_file = __file__
    import os
16

17
18
    example_file = os.path.join(
        os.path.dirname(os.path.dirname(os.path.dirname(current_file))),
19
20
21
        "examples",
        "offline_inference/basic/basic.py",
    )
22
23
24
25
    runpy.run_path(example_file)

    # check if the plugin is loaded correctly
    from vllm.platforms import _init_trace, current_platform
26

27
28
29
    assert current_platform.device_name == "DummyDevice", (
        f"Expected DummyDevice, got {current_platform.device_name}, "
        "possibly because current_platform is imported before the plugin"
30
31
        f" is loaded. The first import:\n{_init_trace}"
    )
32
33


34
# def test_oot_custom_op(default_vllm_config, monkeypatch: pytest.MonkeyPatch):
35
36
37
#     # simulate workload by running an example
#     load_general_plugins()
#     from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
38

39
40
41
#     layer = RotaryEmbedding(16, 16, 16, 16, True, torch.float16)
#     assert layer.__class__.__name__ == "DummyRotaryEmbedding", (
#         f"Expected DummyRotaryEmbedding, got {layer.__class__.__name__}, "
42
43
#         "possibly because the custom op is not registered correctly."
#     )
44
45
#     assert hasattr(layer, "addition_config"), (
#         "Expected DummyRotaryEmbedding to have an 'addition_config' attribute, "
46
47
#         "which is set by the custom op."
#     )