[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:
Mengqing Cao
2025-09-11 11:30:16 +08:00
committed by GitHub
parent b7df04de9b
commit c2fdd4b8bc
3 changed files with 36 additions and 43 deletions

View File

@@ -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):

View File

@@ -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,