[Test] Add ut test for torchair (#4287)
### What this PR does / why we need it?
The current community lacks unit tests (UT) for files such as
torchair_worker, mtp_proposer, and model_runner. Therefore, UT coverage
for these files needs to be added.
### Does this PR introduce _any_ user-facing change?
NO
### How was this patch tested?
- vLLM version: v0.11.0
- vLLM main:
2918c1b49c
---------
Signed-off-by: CodeNine-CJ <chenjian343@huawei.com>
This commit is contained in:
45
tests/ut/torchair/test_torchair_model_runner.py
Normal file
45
tests/ut/torchair/test_torchair_model_runner.py
Normal file
@@ -0,0 +1,45 @@
|
||||
from unittest.mock import MagicMock, Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm.config import VllmConfig
|
||||
|
||||
from tests.ut.base import PytestBase
|
||||
from vllm_ascend.torchair.torchair_model_runner import NPUTorchairModelRunner
|
||||
|
||||
|
||||
class TestNPUTorchairModelRunner(PytestBase):
|
||||
|
||||
@pytest.fixture
|
||||
def setup_npu_torchair_model_runner(self, mocker: MockerFixture):
|
||||
mocker.patch.object(NPUTorchairModelRunner, "__init__",
|
||||
lambda self, *args, **kwargs: None)
|
||||
runner = NPUTorchairModelRunner(Mock(), Mock())
|
||||
|
||||
runner.device = torch.device("cpu")
|
||||
runner.vllm_config = MagicMock(spec=VllmConfig)
|
||||
|
||||
runner.speculative_config = MagicMock(
|
||||
method="deepseek_mtp",
|
||||
num_speculative_tokens=4,
|
||||
disable_padded_drafter_batch=False)
|
||||
|
||||
runner.ascend_config = MagicMock(enable_shared_expert_dp=False,
|
||||
torchair_graph_config=MagicMock(
|
||||
use_cached_graph=True,
|
||||
graph_batch_sizes=[1, 2, 4]))
|
||||
|
||||
runner.decode_token_per_req = 2
|
||||
runner.is_kv_consumer = True
|
||||
runner.max_num_reqs = 100
|
||||
|
||||
runner.model_config = MagicMock(hf_config=MagicMock(index_topk=2))
|
||||
runner.attn_backend = MagicMock(get_builder_cls=lambda: Mock())
|
||||
|
||||
return runner
|
||||
|
||||
def test_init(self, mocker: MockerFixture,
|
||||
setup_npu_torchair_model_runner):
|
||||
runner = setup_npu_torchair_model_runner
|
||||
assert isinstance(runner, NPUTorchairModelRunner)
|
||||
Reference in New Issue
Block a user