[CI/UT] Fix UTs on register customop and warm up model (#2862)
### What this PR does / why we need it?
Fix UTs on register customop and warm up model
### How was this patch tested?
CI passed with existing test.
Co-authored-by: Icey <1790571317@qq.com>
- vLLM version: main
- vLLM main:
cc99baf14d
---------
Signed-off-by: MengqingCao <cmq0113@163.com>
This commit is contained in:
@@ -24,6 +24,7 @@ from vllm.config import (CompilationConfig, ModelConfig, ParallelConfig,
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend import utils
|
||||
from vllm_ascend.utils import REGISTERED_ASCEND_OPS
|
||||
|
||||
|
||||
class TestUtils(TestBase):
|
||||
@@ -302,14 +303,14 @@ class TestUtils(TestBase):
|
||||
|
||||
# ascend custom op is not registered
|
||||
utils.register_ascend_customop()
|
||||
# should call register_oot three
|
||||
self.assertEqual(mock_customop.register_oot.call_count, 13)
|
||||
self.assertEqual(mock_customop.register_oot.call_count,
|
||||
len(REGISTERED_ASCEND_OPS))
|
||||
self.assertTrue(utils._ASCEND_CUSTOMOP_IS_REIGISTERED)
|
||||
|
||||
# ascend custom op is already registered
|
||||
utils.register_ascend_customop()
|
||||
# should not register_oot again, thus only called three in this ut
|
||||
self.assertEqual(mock_customop.register_oot.call_count, 13)
|
||||
self.assertEqual(mock_customop.register_oot.call_count,
|
||||
len(REGISTERED_ASCEND_OPS))
|
||||
|
||||
|
||||
class TestProfileExecuteDuration(TestBase):
|
||||
|
||||
@@ -1009,9 +1009,8 @@ class TestNPUWorker(TestBase):
|
||||
|
||||
@patch("vllm_ascend.worker.worker_v1.NPUPlatform.seed_everything")
|
||||
@patch("vllm_ascend.worker.worker_v1.logger")
|
||||
@patch("torch_npu._npu_matmul_add_fp32")
|
||||
def test_compile_or_warm_up_model_with_eager_mode(self,
|
||||
mock_npu_matmul_add,
|
||||
@patch("vllm_ascend.worker.worker_v1.NPUWorker._warm_up_atb")
|
||||
def test_compile_or_warm_up_model_with_eager_mode(self, mock_warm_up_atb,
|
||||
mock_logger,
|
||||
mock_seed_everything):
|
||||
"""Test compile_or_warm_up_model method - eager mode"""
|
||||
@@ -1054,14 +1053,14 @@ class TestNPUWorker(TestBase):
|
||||
# Verify seed setting
|
||||
mock_seed_everything.assert_called_once_with(12345)
|
||||
|
||||
# Verify calls
|
||||
mock_npu_matmul_add.assert_called_once()
|
||||
# Verify atb warm up
|
||||
mock_warm_up_atb.assert_called_once()
|
||||
|
||||
@patch("vllm_ascend.worker.worker_v1.NPUPlatform.seed_everything")
|
||||
@patch("vllm_ascend.worker.worker_v1.logger")
|
||||
@patch("torch_npu._npu_matmul_add_fp32")
|
||||
@patch("vllm_ascend.worker.worker_v1.NPUWorker._warm_up_atb")
|
||||
def test_compile_or_warm_up_model_with_graph_capture(
|
||||
self, mock_npu_matmul_add, mock_logger, mock_seed_everything):
|
||||
self, mock_warm_up_atb, mock_logger, mock_seed_everything):
|
||||
"""Test compile_or_warm_up_model method - with graph capture enabled"""
|
||||
from vllm_ascend.worker.worker_v1 import NPUWorker
|
||||
|
||||
@@ -1094,8 +1093,8 @@ class TestNPUWorker(TestBase):
|
||||
# Verify seed setting
|
||||
mock_seed_everything.assert_called_once_with(67890)
|
||||
|
||||
# Verify calls
|
||||
mock_npu_matmul_add.assert_called_once()
|
||||
# Verify atb warm up
|
||||
mock_warm_up_atb.assert_called_once()
|
||||
|
||||
@patch("vllm_ascend.worker.worker_v1.CaMemAllocator")
|
||||
def test_initialize_from_config_with_sleep_mode(self,
|
||||
|
||||
Reference in New Issue
Block a user