worker ut test

Signed-off-by: CodeNine-CJ <[email protected]>

Author: ForBetterCodeNine

tests/ut/torchair/test_torchair_model_runner.py

@@ -0,0 +1,58 @@
+import os
+import unittest
+import pytest
+from unittest.mock import Mock, MagicMock, patch
+
+import torch
+from vllm_ascend.torchair_model_runner import NPUTorchairModelRunner
+from vllm.config import VllmConfig
+
+
+class TestNPUTorchairModelRunner(PytestBase):
+
+    @pytest.fixture
+    def setup_npu_torchair_model_runner(self, mocker: MockerFixture):
+        vllm_config = MagicMock(spec=VllmConfig)
+        vllm_config.model_config = MagicMock()
+        vllm_config.model_config.hf_config = MagicMock()
+        vllm_config.model_config.hf_config.index_topk = 2
+
+        device = torch.device("npu:0")
+
+        ascend_config = MagicMock()
+        ascend_config = enable_shared_expert_dp = False
+        ascend_config.max_num_batched_tokens = 2048
+        ascend_config.max_model_len = 1024
+        ascend_config.torchair_graph_config = MagicMock()
+        ascend_config.torchair_graph_config.use_cached_graph = True
+        ascend_config.torchair_graph_config.use_cached_kv_cache_bytes = False
+        ascend_config.torchair_graph_config.graph_batch_sizes = [1, 2, 4]
+        ascend_config.torchair_graph_config.graph_batch_sizes_init = True
+
+        mocker.patch("vllm_ascend.worker.model_runner_v1.NPUModelRunner.__init__",
+                    return_value=None)
+
+        mocker.patch("vllm_ascend.get_ascend_config", return_value=ascend_config)
+        mocker.patch("vllm_ascend.torchair.utils.register_torchair_model")
+        mocker.patch("vllm_ascend.torchair.utils.torchair_ops_patch")
+        mocker.patch("vllm_ascend.torchair.utils.torchair_quant_method_register")
+        mocker.patch("vllm_ascend.envs.VLLM_ASCEND_TRACE_RECOMPILES", return_value=False)
+
+        mock_attn_builder = Mock()
+        mock_attn_backend = Mock()
+        mock_attn_backend.get_builder_cls.return_value = lambda *args, **kwargs: mock_attn_builder
+        with patch.object(NPUTorchairModelRunner, 'attn_backend', mock_attn_backend):
+            with patch.object(NPUTorchairModelRunner, 'speculative_config', MagicMock()):
+                NPUTorchairModelRunner.decode_token_per_req = 1
+                NPUTorchairModelRunner.max_num_tokens = 10
+
+                runner = NPUTorchairModelRunner(vllm_config, device)
+                runner.vllm_config = vllm_config
+                runner.device = device
+                runner.attn_backend = mock_attn_backend
+        
+        return runner
+
+    def test_init(self, mocker: MockerFixture, setup_npu_torchair_model_runner):
+        runner = setup_npu_torchair_model_runner
+        assert isinstance(runner, NPUTorchairModelRunner)

tests/ut/torchair/test_torchair_mtp_proposer.py

@@ -0,0 +1,118 @@
+import pytest
+from unittest.mock import Mock, MagicMock, patch
+from vllm_ascend import torchair_mtp_proposer
+from vllm.config import VllmConfig
+
+import torch
+
+def vllm_version_is(version):
+    return version == "0.11.0"
+
+import sys
+sys.modules[__name__].vllm_version_is = vllm_version_is
+
+
+class TestTorchairMtpProposer(PytestBase):
+
+    @pytest.fixture
+    def setup_torchair_mtp_proposer(self, mocker: pytest.MockerFixture):
+        vllm_config = MagicMock(spec=VllmConfig)
+        vllm_config.device_config = MagicMock()
+        vllm_config.device_config.device = torch.device("npu:0")
+        vllm_config.speculative_config = MagicMock()
+        vllm_config.speculative_config.draft_model_config = MagicMock()
+        vllm_config.speculative_config.draft_model_config.dtype = torch.float16
+        vllm_config.load_config = MagicMock()
+
+        device = torch.device("npu:0")
+        runner = MagicMock()
+
+        mocker.patch("vllm_ascend.torchair_mtp_proposer.__init__",
+                     return_value=None)
+        
+        if vllm_version_is("0.11.0"):
+            mock_set_default_dtype = mocker.patch(
+                'vllm.model_executor.model_loader.utils.set_default_torch_dtype'
+            )
+        else:
+            mock_set_default_dtype = mocker.patch(
+                'vllm.utls.torch_utils.set_default_torch_dtype'
+            )
+        mock_set_default_dtype.return_value.__enter__.return_value = None
+
+        mock_model_loader = MagicMock()
+        mocker.patch("vllm.model_executor.model_loader.get_model_loader",
+                     return_value=mock_model_loader)
+        mock_layers = {"target_attn_layer_1": Mock(), "draft_attn_layer_2": Mock()}
+        mocker.patch("vllm.config.get_layers_from_vllm_config",
+                     return_value=mock_layers)
+        mock_set_current = mocker.patch("vllm.config.set_current_vllm_config")
+        mock_set_current.return_value.__enter__.return_value = None
+        mock_torchair_deepseek_mtp = MagicMock()
+        mock_torchair_deepseek_mtp.to.return_value = mock_torchair_deepseek_mtp
+        mocker.patch("vllm_ascend.torchair.models.torchair_deepseek_mtp.TorchairDeepSeekMTP",
+                     return_value=mock_torchair_deepseek_mtp)
+        mocker.patch("vllm.model_executor.model_loader.utils.process_weights_after_loading")
+
+        proposer = TorchairMtpProposer(vllm_config, device, runner)
+        proposer.vllm_config = vllm_config
+        proposer.device = device
+        proposer.runner = runner
+
+        return proposer, mock_model_loader, mock_torchair_deepseek_mtp
+    
+    def test_init(self, setup_torchair_mtp_proposer):
+        proposer, _, _, = setup_torchair_mtp_proposer
+
+        assert isinstance(proposer, setup_torchair_mtp_proposer)
+        assert proposer.torchair_compiled_model is None
+        assert proposer.torchair_compiled_models = {}
+        Mock.assert_called_once_with(
+            proposer.__class__.__bases__[0],
+            proposer.vllm_config,
+            proposer.device,
+            proposer.runner
+        )
+    
+    def test_load_model(self, setup_torchair_mtp_proposer, mocker: pytest.MockerFixture):
+        proposer, mock_model_loader, mock_torchair_deepseek_mtp = setup_torchair_mtp_proposer
+        dummpy_model = Mock()
+
+        proposer.load_model(dummpy_model)
+
+        mocker.patch("vllm.model_executor.model_loader.get_model_loader").assert_called_once_with(
+            proposer.vllm_config.load_config
+        )
+
+        mock_get_layers = mocker.patch("vllm.config.get_layers_from_vllm_config")
+        assert mock_get_layers.call_count = 2
+        mock_get_layers.assert_called_with(
+            proposer.vllm_config,
+            mocker.patch("vllm.model_executor.layers.attention_layer_base.AttentionLayerBase")
+        )
+
+        mocker.patch("vllm_ascend.torchair.models.torchair_deepseek_mtp.TorchairDeepSeekMTP").assert_called_once_with(
+            vllm_config=proposer.vllm_config
+        )
+        mock_torchair_deepseek_mtp.to.assert_called_once(
+            proposer.vllm_config.device_config.device
+        )
+
+        assert len(proposer.attn_layer_name) == 1
+        mocker_layers_keys = mock_get_layers.return_value.keys()
+        assert proposer.attn_layer_name[0] in mocker_layers_keys
+
+        mock_model_loader.get_all_weights.assert_called_once_with(
+            proposer.vllm_config.speculative_config.draft_model_config,
+            mock_torchair_deepseek_mtp
+        )
+        mock_torchair_deepseek_mtp.load_weights.assert_called_once_with(
+            mock_model_loader.get_all_weights.return_value
+        )
+
+        mock_process_weights = mocker.patch("vllm.model_executor.model_loader.utils.process_weights_after_loading")
+        mock_process_weights.assert_called_once_with(
+            mock_torchair_deepseek_mtp,
+            proposer.vllm_config.speculative_config.draft_model_config,
+            proposer.vllm_config.device_config.device
+        )

tests/ut/torchair/test_torchair_worker.py

@@ -0,0 +1,112 @@
+import os
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+from vllm.config import CacheConfig, ModelConfig, ParallelConfig, VllmConfig
+
+from tests.ut.base import TestBase
+from vllm_ascend.utils import vllm_version_is
+
+init_cache_hf_modules_path = "vllm.utils.init_cached_hf_modules" if vllm_version_is(
+    "0.11.0") else "vllm.utils.import_utils.init_cached_hf_modules"
+
+
+class TestNPUTorchairWorker(TestBase):
+
+    def setUp(self):
+        self.cache_config_mock = MagicMock(spec=CacheConfig)
+        self.cache_config_mock.cache_type = "auto"
+
+        self.model_config_mock = MagicMock(spec=ModelConfig)
+        self.model_config_mock.dtype = torch.float16
+        self.model_config_mock.trust_remote_code = False
+
+        self.hf_config_mock = MagicMock()
+        self.hf_config_mock.model_type = "test_model"
+        if hasattr(self.hf_config_mock, 'index_topk'):
+            delattr(self.hf_config_mock, 'index_topk')
+
+        self.model_config_mock.hf_config = self.hf_config_mock
+
+        self.parallel_config_mock = MagicMock(spec=ParallelConfig)
+
+        self.vllm_config_mock = MagicMock(spec=VllmConfig)
+        self.vllm_config_mock.cache_config = self.cache_config_mock
+        self.vllm_config_mock.model_config = self.model_config_mock
+        self.vllm_config_mock.parallel_config = self.parallel_config_mock
+        self.vllm_config_mock.additional_config = None
+        self.vllm_config_mock.load_config = None
+        self.vllm_config_mock.scheduler_config = None
+        self.vllm_config_mock.device_config = None
+        self.vllm_config_mock.compilation_config = None
+
+        self.local_rank = 0
+        self.rank = 0
+        self.distributed_init_method = "tcp://localhost:12345"
+        self.is_driver_worker = False
+
+    @patch(
+        "vllm_ascend.worker.worker_v1.NPUWorker._init_worker_distributed_environment"
+    )
+    @patch("vllm_ascend.worker.worker_v1.NPUPlatform")
+    def test_init_device(self, mock_platform, mock_init_dist_env):
+        from vllm_ascend.worker.worker_v1 import NPUWorker
+
+        mock_platform.mem_get_info.return_value = (1000, 2000)
+
+        with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
+            worker = NPUWorker
+            worker.local_rank = 1
+            worker.model_config = MagicMock()
+            worker.model_config.seed = 42
+
+            result = worker._init_device()
+
+            mock_platform.set_device.assert_called_once()
+            call_args = mock_platform.set_device.call_args[0][0]
+            self.assertEqual(str(call_args), "npu:1")
+
+            mock_platform.empty_cache.assert_called_once()
+            mock_platform.seed_everything.assert_called_once_with(42)
+            mock_platform.mem_get_info.assert_called_once(
+            )
+            mock_init_dist_env.assert_called_once(
+            )
+
+            self.assertEqual(str(result), "npu:1")
+            self.assertEqual(worker.init_npu_memory, 1000)
+    
+    @patch(
+        "vllm_ascend.worker.worker_v1.NPUWorker._init_worker_distributed_environment"
+    )
+    @patch("vllm_ascend.worker.worker_v1.NPUPlatform")
+    def test_init_device_torchair_worker(self, mock_platform, mock_init_dist_env):
+        from vllm_ascend.torchair.torchair_worker import NPUTorchairWorker
+
+        mock_platform.mem_get_info.return_value = (1000, 2000)
+
+        with patch.object(NPUTorchairWorker, "__init__", lambda x, **kwargs: None):
+            worker = NPUTorchairWorker
+            worker.local_rank = 1
+            worker.model_config = MagicMock()
+            worker.model_config.seed = 42
+
+            result = worker._init_device()
+
+            mock_platform.set_device.assert_called_once()
+            call_args = mock_platform.set_device.call_args[0][0]
+            self.assertEqual(str(call_args), "npu:1")
+
+            mock_platform.empty_cache.assert_called_once()
+            mock_platform.seed_everything.assert_called_once_with(42)
+            mock_platform.mem_get_info.assert_called_once(
+            )
+            mock_init_dist_env.assert_called_once(
+            )
+
+            self.assertEqual(str(result), "npu:1")
+            self.assertEqual(worker.init_npu_memory, 1000)
+
+
+