41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
import importlib.metadata
|
|
import importlib.util
|
|
import os
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from tests.helpers import TmpCase
|
|
|
|
|
|
@unittest.skipIf(importlib.util.find_spec("vllm") is None, "needs vllm")
|
|
class Plugin(TmpCase):
|
|
def test_entry_point_is_installed(self):
|
|
plugins = importlib.metadata.entry_points(group="vllm.general_plugins")
|
|
self.assertIn("prain", [p.name for p in plugins])
|
|
|
|
def test_register_swaps_in_tiered_experts(self):
|
|
config = self.write("p.toml", f'pack_dir = "{self.dir}/pack"\n')
|
|
with mock.patch.dict(os.environ, {"PRAIN_CONFIG": config}):
|
|
from prain.plugin import register
|
|
|
|
register()
|
|
register() # vLLM may call it again; must not raise
|
|
from vllm.model_executor.custom_op import op_registry_oot
|
|
|
|
from prain.experts import TieredRoutedExperts
|
|
|
|
self.assertIs(op_registry_oot["RoutedExperts"], TieredRoutedExperts)
|
|
|
|
def test_draft_hook_wraps_the_v2_model_runner_once(self):
|
|
from vllm.v1.worker.gpu.input_batch import InputBatch
|
|
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
|
|
|
|
from prain.experts import _install_draft_hook
|
|
|
|
_install_draft_hook()
|
|
wrapped = GPUModelRunner.prepare_inputs
|
|
_install_draft_hook()
|
|
self.assertIs(GPUModelRunner.prepare_inputs, wrapped)
|
|
self.assertTrue(wrapped.prain)
|
|
self.assertLessEqual({"req_ids", "query_start_loc_np"}, set(InputBatch.__dataclass_fields__))
|