worker ut

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

Author: ForBetterCodeNine

tests/ut/torchair/test_torchair_model_runner.py

@@ -1,12 +1,13 @@
-import os
-import unittest
-import pytest
-from unittest.mock import Mock, MagicMock, patch
+from unittest.mock import MagicMock, Mock, patch
 
+import pytest
 import torch
-from vllm_ascend.torchair_model_runner import NPUTorchairModelRunner
+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):
 
@@ -20,7 +21,6 @@ def setup_npu_torchair_model_runner(self, mocker: MockerFixture):
         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()
@@ -29,30 +29,37 @@ def setup_npu_torchair_model_runner(self, mocker: MockerFixture):
         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.worker.model_runner_v1.NPUModelRunner.__init__",
+            return_value=None)
 
-        mocker.patch("vllm_ascend.get_ascend_config", return_value=ascend_config)
+        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)
+        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()):
+        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):
+    def test_init(self, mocker: MockerFixture,
+                  setup_npu_torchair_model_runner):
         runner = setup_npu_torchair_model_runner
-        assert isinstance(runner, NPUTorchairModelRunner)
+        assert isinstance(runner, NPUTorchairModelRunner)

tests/ut/torchair/test_torchair_mtp_proposer.py

@@ -1,15 +1,12 @@
-import pytest
-from unittest.mock import Mock, MagicMock, patch
-from vllm_ascend import torchair_mtp_proposer
-from vllm.config import VllmConfig
+from unittest.mock import MagicMock, Mock
 
+import pytest
 import torch
+from vllm.config import VllmConfig
 
-def vllm_version_is(version):
-    return version == "0.11.0"
-
-import sys
-sys.modules[__name__].vllm_version_is = vllm_version_is
+from tests.ut.base import PytestBase
+from vllm_ascend.torchair.torchair_mtp_proposer import TorchairMtpProposer
+from vllm_ascend.utils import vllm_version_is
 
 
 class TestTorchairMtpProposer(PytestBase):
@@ -29,90 +26,91 @@ def setup_torchair_mtp_proposer(self, mocker: pytest.MockerFixture):
 
         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'
-            )
+                '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()}
+        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")
+        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):
+        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
-        )
+        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 = mocker.patch(
+            "vllm.config.get_layers_from_vllm_config")
         mock_get_layers.assert_called_with(
             proposer.vllm_config,
-            mocker.patch("vllm.model_executor.layers.attention_layer_base.AttentionLayerBase")
-        )
+            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
-        )
+        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
-        )
+            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)
         mock_torchair_deepseek_mtp.load_weights.assert_called_once_with(
-            mock_model_loader.get_all_weights.return_value
-        )
+            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 = 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
-        )
+            proposer.vllm_config.device_config.device)

tests/ut/torchair/test_torchair_worker.py

@@ -1,5 +1,3 @@
-import os
-import unittest
 from unittest.mock import MagicMock, patch
 
 import torch
@@ -69,24 +67,24 @@ def test_init_device(self, mock_platform, mock_init_dist_env):
 
             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(
-            )
+            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):
+    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):
+        with patch.object(NPUTorchairWorker, "__init__",
+                          lambda x, **kwargs: None):
             worker = NPUTorchairWorker
             worker.local_rank = 1
             worker.model_config = MagicMock()
@@ -100,13 +98,8 @@ def test_init_device_torchair_worker(self, mock_platform, mock_init_dist_env):
 
             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(
-            )
+            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)
-
-
-