| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from vllm.plugins import load_general_plugins |
| 8 | |
| 9 | |
| 10 | def test_platform_plugins(): |
| 11 | # simulate workload by running an example |
| 12 | import runpy |
| 13 | |
| 14 | current_file = __file__ |
| 15 | import os |
| 16 | |
| 17 | example_file = os.path.join( |
| 18 | os.path.dirname(os.path.dirname(os.path.dirname(current_file))), |
| 19 | "examples", |
| 20 | "basic/offline_inference/basic.py", |
| 21 | ) |
| 22 | runpy.run_path(example_file) |
| 23 | |
| 24 | # check if the plugin is loaded correctly |
| 25 | from vllm.platforms import _init_trace, current_platform |
| 26 | |
| 27 | assert current_platform.device_name == "DummyDevice", ( |
| 28 | f"Expected DummyDevice, got {current_platform.device_name}, " |
| 29 | "possibly because current_platform is imported before the plugin" |
| 30 | f" is loaded. The first import:\n{_init_trace}" |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | def test_oot_custom_op(default_vllm_config, monkeypatch: pytest.MonkeyPatch): |
| 35 | # simulate workload by running an example |
| 36 | load_general_plugins() |
| 37 | from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding |
| 38 | |
| 39 | layer = RotaryEmbedding(16, 16, 16, 16, True, torch.float16) |
| 40 | assert layer.__class__.__name__ == "DummyRotaryEmbedding", ( |
| 41 | f"Expected DummyRotaryEmbedding, got {layer.__class__.__name__}, " |
| 42 | "possibly because the custom op is not registered correctly." |
| 43 | ) |
| 44 | assert hasattr(layer, "addition_config"), ( |
| 45 | "Expected DummyRotaryEmbedding to have an 'addition_config' attribute, " |
| 46 | "which is set by the custom op." |
| 47 | ) |