1512 lines
64 KiB
Python
1512 lines
64 KiB
Python
|
import unittest
|
||
|
|
from unittest.mock import MagicMock, patch
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from vllm.config import CacheConfig, ModelConfig, ParallelConfig, ProfilerConfig, VllmConfig
|
||
|
|
|
||
|
|
from tests.ut.base import TestBase
|
||
|
|
|
||
|
|
init_cached_hf_modules_path = "vllm.utils.import_utils.init_cached_hf_modules"
|
||
|
|
|
||
|
|
|
||
|
|
class TestNPUWorker(TestBase):
|
||
|
|
def setUp(self):
|
||
|
|
"""Setup test environment"""
|
||
|
|
# Create configuration mocks
|
||
|
|
self.cache_config_mock = MagicMock(spec=CacheConfig)
|
||
|
|
self.cache_config_mock.cache_dtype = "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 = MagicMock()
|
||
|
|
self.vllm_config_mock.compilation_config.ir_enable_torch_wrap = False
|
||
|
|
self.vllm_config_mock.kernel_config = MagicMock()
|
||
|
|
self.vllm_config_mock.kernel_config.ir_op_priority = MagicMock()
|
||
|
|
self.vllm_config_mock.kernel_config.ir_op_priority.set_default = MagicMock()
|
||
|
|
self.vllm_config_mock.profiler_config = MagicMock()
|
||
|
|
self.vllm_config_mock.quant_config = MagicMock()
|
||
|
|
self.vllm_config_mock.speculative_config = None
|
||
|
|
self.vllm_config_mock.observability_config = None
|
||
|
|
self.vllm_config_mock.weight_transfer_config = None
|
||
|
|
|
||
|
|
self.local_rank = 0
|
||
|
|
self.rank = 0
|
||
|
|
self.distributed_init_method = "tcp://localhost:12345"
|
||
|
|
self.is_driver_worker = False
|
||
|
|
|
||
|
|
@patch("vllm_ascend.utils.adapt_patch")
|
||
|
|
@patch("vllm_ascend.ops")
|
||
|
|
@patch("vllm_ascend.worker.worker._register_atb_extensions")
|
||
|
|
@patch("vllm_ascend.worker.worker.register_ascend_customop")
|
||
|
|
@patch("vllm_ascend.worker.worker.get_ascend_config")
|
||
|
|
@patch("vllm_ascend.worker.worker.init_ascend_config")
|
||
|
|
@patch("vllm_ascend.worker.worker.check_ascend_device_type")
|
||
|
|
@patch(init_cached_hf_modules_path, create=True)
|
||
|
|
@patch("vllm_ascend.worker.worker.TorchNPUProfilerWrapper")
|
||
|
|
def test_init_npu_worker_normal_case(
|
||
|
|
self,
|
||
|
|
mock_profiler_wrapper,
|
||
|
|
mock_init_cached_hf_modules,
|
||
|
|
mock_check_ascend_device_type,
|
||
|
|
mock_init_ascend_config,
|
||
|
|
mock_get_ascend_config,
|
||
|
|
mock_register_ascend_customop,
|
||
|
|
mock_register_atb_extensions,
|
||
|
|
mock_ops,
|
||
|
|
mock_adapt_patch,
|
||
|
|
):
|
||
|
|
"""Test NPUWorker normal initialization"""
|
||
|
|
# Setup mock behavior
|
||
|
|
mock_ops.register_dummy_fusion_op.return_value = None
|
||
|
|
mock_ascend_config = MagicMock()
|
||
|
|
mock_ascend_config.enable_cpu_binding = True
|
||
|
|
mock_get_ascend_config.return_value = mock_ascend_config
|
||
|
|
|
||
|
|
# Import and create NPUWorker instance
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
worker = NPUWorker(
|
||
|
|
vllm_config=self.vllm_config_mock,
|
||
|
|
local_rank=self.local_rank,
|
||
|
|
rank=self.rank,
|
||
|
|
distributed_init_method=self.distributed_init_method,
|
||
|
|
is_driver_worker=self.is_driver_worker,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Verify initialization call order
|
||
|
|
mock_adapt_patch.assert_called_once()
|
||
|
|
mock_ops.register_dummy_fusion_op.assert_called_once()
|
||
|
|
mock_register_atb_extensions.assert_called_once()
|
||
|
|
mock_register_ascend_customop.assert_called_once()
|
||
|
|
mock_init_ascend_config.assert_called_once_with(self.vllm_config_mock)
|
||
|
|
mock_check_ascend_device_type.assert_called_once()
|
||
|
|
|
||
|
|
# Verify cache_dtype setting
|
||
|
|
self.assertEqual(worker.cache_dtype, torch.float16)
|
||
|
|
# Profiler is lazily initialized - not created during __init__ (RFC #6954)
|
||
|
|
mock_profiler_wrapper.assert_not_called()
|
||
|
|
|
||
|
|
# Verify init_cached_hf_modules is not called (trust_remote_code=False)
|
||
|
|
mock_init_cached_hf_modules.assert_not_called()
|
||
|
|
|
||
|
|
@patch("vllm_ascend.utils.adapt_patch")
|
||
|
|
@patch("vllm_ascend.ops")
|
||
|
|
@patch("vllm_ascend.worker.worker._register_atb_extensions")
|
||
|
|
@patch("vllm_ascend.worker.worker.register_ascend_customop")
|
||
|
|
@patch("vllm_ascend.worker.worker.get_ascend_config")
|
||
|
|
@patch("vllm_ascend.worker.worker.init_ascend_config")
|
||
|
|
@patch("vllm_ascend.worker.worker.check_ascend_device_type")
|
||
|
|
@patch(init_cached_hf_modules_path, create=True)
|
||
|
|
@patch("vllm_ascend.worker.worker.TorchNPUProfilerWrapper")
|
||
|
|
def test_init_npu_worker_with_trust_remote_code(
|
||
|
|
self,
|
||
|
|
mock_profiler_wrapper,
|
||
|
|
mock_init_cached_hf_modules,
|
||
|
|
mock_check_ascend_device_type,
|
||
|
|
mock_init_ascend_config,
|
||
|
|
mock_get_ascend_config,
|
||
|
|
mock_register_ascend_customop,
|
||
|
|
mock_register_atb_extensions,
|
||
|
|
mock_ops,
|
||
|
|
mock_adapt_patch,
|
||
|
|
):
|
||
|
|
"""Test NPUWorker initialization with trust_remote_code=True"""
|
||
|
|
# Set trust_remote_code=True
|
||
|
|
self.model_config_mock.trust_remote_code = True
|
||
|
|
mock_ops.register_dummy_fusion_op.return_value = None
|
||
|
|
mock_ascend_config = MagicMock()
|
||
|
|
mock_ascend_config.enable_cpu_binding = True
|
||
|
|
mock_get_ascend_config.return_value = mock_ascend_config
|
||
|
|
|
||
|
|
# Create NPUWorker instance
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
_ = NPUWorker(
|
||
|
|
vllm_config=self.vllm_config_mock,
|
||
|
|
local_rank=self.local_rank,
|
||
|
|
rank=self.rank,
|
||
|
|
distributed_init_method=self.distributed_init_method,
|
||
|
|
is_driver_worker=self.is_driver_worker,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Verify init_cached_hf_modules is called (trust_remote_code=True)
|
||
|
|
mock_init_cached_hf_modules.assert_not_called()
|
||
|
|
|
||
|
|
@patch("vllm_ascend.utils.adapt_patch")
|
||
|
|
@patch("vllm_ascend.ops")
|
||
|
|
@patch("vllm_ascend.worker.worker._register_atb_extensions")
|
||
|
|
@patch("vllm_ascend.worker.worker.register_ascend_customop")
|
||
|
|
@patch("vllm_ascend.worker.worker.get_ascend_config")
|
||
|
|
@patch("vllm_ascend.worker.worker.init_ascend_config")
|
||
|
|
@patch("vllm_ascend.worker.worker.check_ascend_device_type")
|
||
|
|
@patch(init_cached_hf_modules_path, create=True)
|
||
|
|
@patch("vllm_ascend.worker.worker.TorchNPUProfilerWrapper")
|
||
|
|
def test_init_npu_worker_with_custom_cache_dtype(
|
||
|
|
self,
|
||
|
|
mock_profiler_wrapper,
|
||
|
|
mock_init_cached_hf_modules,
|
||
|
|
mock_check_ascend_device_type,
|
||
|
|
mock_init_ascend_config,
|
||
|
|
mock_get_ascend_config,
|
||
|
|
mock_register_ascend_customop,
|
||
|
|
mock_register_atb_extensions,
|
||
|
|
mock_ops,
|
||
|
|
mock_adapt_patch,
|
||
|
|
):
|
||
|
|
"""Test NPUWorker initialization with custom cache_dtype"""
|
||
|
|
# Set custom cache_dtype
|
||
|
|
self.cache_config_mock.cache_dtype = "float32"
|
||
|
|
mock_ops.register_dummy_fusion_op.return_value = None
|
||
|
|
mock_ascend_config = MagicMock()
|
||
|
|
mock_ascend_config.enable_cpu_binding = True
|
||
|
|
mock_get_ascend_config.return_value = mock_ascend_config
|
||
|
|
|
||
|
|
# Create NPUWorker instance
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
with patch("vllm.utils.torch_utils.STR_DTYPE_TO_TORCH_DTYPE", {"float32": torch.float32}):
|
||
|
|
worker = NPUWorker(
|
||
|
|
vllm_config=self.vllm_config_mock,
|
||
|
|
local_rank=self.local_rank,
|
||
|
|
rank=self.rank,
|
||
|
|
distributed_init_method=self.distributed_init_method,
|
||
|
|
is_driver_worker=self.is_driver_worker,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Verify cache_dtype is set to custom value
|
||
|
|
self.assertEqual(worker.cache_dtype, torch.float32)
|
||
|
|
|
||
|
|
def test_initialize_cache(self):
|
||
|
|
"""Test initialize_cache method"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create a simple worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.cache_config = MagicMock()
|
||
|
|
|
||
|
|
# Test initialize_cache
|
||
|
|
worker.initialize_cache(100, 50)
|
||
|
|
|
||
|
|
# Verify parameter setting
|
||
|
|
self.assertEqual(worker.cache_config.num_gpu_blocks, 100)
|
||
|
|
self.assertEqual(worker.cache_config.num_cpu_blocks, 50)
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.CaMemAllocator")
|
||
|
|
@patch("vllm_ascend.worker.worker.get_ascend_config")
|
||
|
|
def test_wake_up_mode_enabled(self, mock_get_config, mock_allocator_class):
|
||
|
|
mock_config = MagicMock()
|
||
|
|
mock_config.weight_nz_mode = 0
|
||
|
|
mock_config.enable_sleep_mode_extra_cleanup = True
|
||
|
|
mock_get_config.return_value = mock_config
|
||
|
|
"""Test wake_up method when sleep mode is enabled"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Setup mock
|
||
|
|
mock_allocator = MagicMock()
|
||
|
|
mock_allocator_class.get_instance.return_value = mock_allocator
|
||
|
|
|
||
|
|
mock_hidden_size = MagicMock()
|
||
|
|
mock_hf_config = MagicMock()
|
||
|
|
mock_hf_config.hidden_size = mock_hidden_size
|
||
|
|
mock_model_config = MagicMock()
|
||
|
|
mock_model_config.hf_config = mock_hf_config
|
||
|
|
mock_vllm_config = MagicMock()
|
||
|
|
mock_vllm_config.model_config = mock_model_config
|
||
|
|
|
||
|
|
mock_model_runner = MagicMock()
|
||
|
|
mock_model_runner.model = MagicMock()
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = mock_model_runner
|
||
|
|
worker.vllm_config = mock_vllm_config
|
||
|
|
worker._sleep_saved_buffers = {}
|
||
|
|
worker.sleep_wakeup_manager = MagicMock()
|
||
|
|
# Test wake_up method
|
||
|
|
worker.wake_up(tags=["test_tag"])
|
||
|
|
|
||
|
|
mock_allocator.wake_up.assert_called_once_with(tags=["test_tag"])
|
||
|
|
worker.sleep_wakeup_manager.wakeup.assert_called_once_with(["test_tag"])
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.current_platform")
|
||
|
|
@patch("vllm_ascend.worker.worker.MemorySnapshot")
|
||
|
|
@patch("vllm_ascend.worker.worker.NPUWorker._init_worker_distributed_environment")
|
||
|
|
@patch("vllm_ascend.worker.worker.init_device_properties_triton")
|
||
|
|
@patch("vllm_ascend.worker.worker.get_ascend_device_type")
|
||
|
|
@patch("torch.npu.set_device")
|
||
|
|
@patch("torch.npu.empty_cache")
|
||
|
|
@patch("torch.npu.mem_get_info")
|
||
|
|
def test_init_device(
|
||
|
|
self,
|
||
|
|
mock_mem_get_info,
|
||
|
|
mock_empty_cache,
|
||
|
|
mock_set_device,
|
||
|
|
mock_get_device_type,
|
||
|
|
mock_init_triton,
|
||
|
|
mock_init_dist_env,
|
||
|
|
mock_snapshot_cls,
|
||
|
|
mock_current_platform,
|
||
|
|
):
|
||
|
|
"""Test _init_device method"""
|
||
|
|
from vllm_ascend.worker.worker import AscendDeviceType, NPUWorker
|
||
|
|
|
||
|
|
# Setup mock
|
||
|
|
mock_mem_get_info.return_value = (1000, 2000)
|
||
|
|
mock_get_device_type.return_value = AscendDeviceType.A2
|
||
|
|
|
||
|
|
# Mock MemorySnapshot
|
||
|
|
mock_snapshot = MagicMock()
|
||
|
|
mock_snapshot.free_memory = 1000
|
||
|
|
mock_snapshot.total_memory = 2000
|
||
|
|
mock_snapshot_cls.return_value = mock_snapshot
|
||
|
|
|
||
|
|
# Mock current_platform for v0.24.0 init_device path
|
||
|
|
mock_current_platform.logical_device_id_to_visible_device_id.return_value = 0
|
||
|
|
mock_current_platform.device_type = "npu"
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.local_rank = 0
|
||
|
|
worker.model_config = MagicMock()
|
||
|
|
worker.model_config.seed = 42
|
||
|
|
worker.parallel_config = MagicMock()
|
||
|
|
worker.parallel_config.local_world_size = 0
|
||
|
|
worker.parallel_config.data_parallel_size = 1
|
||
|
|
worker.parallel_config.assigned_physical_gpu_ids = None
|
||
|
|
worker.parallel_config.distributed_executor_backend = "ray"
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.vllm_config.kv_transfer_config = None
|
||
|
|
worker.cache_config = MagicMock()
|
||
|
|
worker.cache_config.gpu_memory_utilization = 0.5
|
||
|
|
|
||
|
|
# Test _init_device
|
||
|
|
result = worker._init_device()
|
||
|
|
|
||
|
|
mock_init_dist_env.assert_called_once()
|
||
|
|
self.assertEqual(str(result), "npu:0")
|
||
|
|
self.assertEqual(worker.init_snapshot, mock_snapshot)
|
||
|
|
self.assertEqual(worker.requested_memory, 2000 * 0.5)
|
||
|
|
|
||
|
|
def test_profile_start_stop(self):
|
||
|
|
"""Test profile method start and stop"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
profiler_config = ProfilerConfig(
|
||
|
|
profiler="torch",
|
||
|
|
torch_profiler_dir="/path/to/traces",
|
||
|
|
)
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.profiler_config = profiler_config
|
||
|
|
worker.rank = 0
|
||
|
|
mock_profiler = MagicMock()
|
||
|
|
worker.profiler = mock_profiler
|
||
|
|
|
||
|
|
with patch("vllm.distributed.utils.get_worker_rank_suffix", return_value="dp0_pp0_tp0_dcp0_ep0_rank0"):
|
||
|
|
worker.profile(is_start=True)
|
||
|
|
mock_profiler.start.assert_called_once()
|
||
|
|
|
||
|
|
worker.profile(is_start=False)
|
||
|
|
mock_profiler.stop.assert_called_once()
|
||
|
|
|
||
|
|
def test_profile_no_profiler_raises_error(self):
|
||
|
|
"""Test profile method raises exception when profiler is not available"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock - profiler_config indicates profiling disabled
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.profiler = None
|
||
|
|
worker.profiler_config = ProfilerConfig(profiler=None, torch_profiler_dir="")
|
||
|
|
|
||
|
|
# Test should raise exception
|
||
|
|
with self.assertRaises(RuntimeError) as cm:
|
||
|
|
worker.profile()
|
||
|
|
|
||
|
|
self.assertIn("Profiling is not enabled", str(cm.exception))
|
||
|
|
|
||
|
|
def test_profile_with_prefix_uses_trace_name(self):
|
||
|
|
"""[RFC #6954] profile() accepts profile_prefix and passes trace_name to TorchNPUProfilerWrapper"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
profiler_config = ProfilerConfig(
|
||
|
|
profiler="torch",
|
||
|
|
torch_profiler_dir="/path/to/traces",
|
||
|
|
)
|
||
|
|
vllm_config_mock = MagicMock()
|
||
|
|
vllm_config_mock.profiler_config = profiler_config
|
||
|
|
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.profiler_config = profiler_config
|
||
|
|
worker.profiler = None
|
||
|
|
worker.rank = 0
|
||
|
|
|
||
|
|
with (
|
||
|
|
patch("vllm.distributed.utils.get_worker_rank_suffix", return_value="dp0_pp0_tp0_dcp0_ep0_rank0"),
|
||
|
|
patch("vllm_ascend.worker.worker.TorchNPUProfilerWrapper") as mock_profiler_wrapper,
|
||
|
|
):
|
||
|
|
worker.profile(is_start=True, profile_prefix="warmup")
|
||
|
|
|
||
|
|
mock_profiler_wrapper.assert_called_once_with(
|
||
|
|
profiler_config,
|
||
|
|
"warmup_dp0_pp0_tp0_dcp0_ep0_rank0",
|
||
|
|
)
|
||
|
|
mock_profiler_wrapper.return_value.start.assert_called_once()
|
||
|
|
|
||
|
|
def test_profile_lazy_init(self):
|
||
|
|
"""[RFC #6954] Profiler is lazily created on first profile(is_start=True) call"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
profiler_config = ProfilerConfig(
|
||
|
|
profiler="torch",
|
||
|
|
torch_profiler_dir="/path/to/traces",
|
||
|
|
)
|
||
|
|
vllm_config_mock = MagicMock()
|
||
|
|
vllm_config_mock.profiler_config = profiler_config
|
||
|
|
|
||
|
|
with patch("vllm_ascend.worker.worker.TorchNPUProfilerWrapper") as mock_profiler_wrapper:
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.profiler_config = profiler_config
|
||
|
|
worker.profiler = None
|
||
|
|
worker.rank = 0
|
||
|
|
|
||
|
|
self.assertIsNone(worker.profiler)
|
||
|
|
mock_profiler_wrapper.assert_not_called()
|
||
|
|
|
||
|
|
with patch("vllm.distributed.utils.get_worker_rank_suffix", return_value="dp0_pp0_tp0_dcp0_ep0_rank0"):
|
||
|
|
worker.profile(is_start=True)
|
||
|
|
|
||
|
|
mock_profiler_wrapper.assert_called_once_with(
|
||
|
|
profiler_config,
|
||
|
|
"dp0_pp0_tp0_dcp0_ep0_rank0",
|
||
|
|
)
|
||
|
|
self.assertIs(worker.profiler, mock_profiler_wrapper.return_value)
|
||
|
|
mock_profiler_wrapper.return_value.start.assert_called_once()
|
||
|
|
|
||
|
|
def test_profile_restart_reuses_existing_profiler(self):
|
||
|
|
"""[RFC #6954] Restarting profile reuses existing profiler."""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
profiler_config = ProfilerConfig(
|
||
|
|
profiler="torch",
|
||
|
|
torch_profiler_dir="/path/to/traces",
|
||
|
|
)
|
||
|
|
mock_profiler = MagicMock()
|
||
|
|
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.profiler_config = profiler_config
|
||
|
|
worker.profiler = None
|
||
|
|
worker.rank = 0
|
||
|
|
|
||
|
|
with (
|
||
|
|
patch("vllm.distributed.utils.get_worker_rank_suffix", return_value="dp0_pp0_tp0_dcp0_ep0_rank0"),
|
||
|
|
patch("vllm_ascend.worker.worker.TorchNPUProfilerWrapper", return_value=mock_profiler) as mock_wrapper,
|
||
|
|
):
|
||
|
|
worker.profile(is_start=True, profile_prefix="session1")
|
||
|
|
mock_wrapper.assert_called_once_with(
|
||
|
|
profiler_config,
|
||
|
|
"session1_dp0_pp0_tp0_dcp0_ep0_rank0",
|
||
|
|
)
|
||
|
|
|
||
|
|
worker.profile(is_start=False)
|
||
|
|
worker.profile(is_start=True) # Restart without new prefix
|
||
|
|
# Should NOT create new profiler, just restart existing
|
||
|
|
mock_wrapper.assert_called_once()
|
||
|
|
self.assertEqual(mock_profiler.start.call_count, 2)
|
||
|
|
mock_profiler.stop.assert_called_once()
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.logger")
|
||
|
|
def test_profile_stop_without_start_logs_warning(self, mock_logger):
|
||
|
|
"""Test stopping profiling before start logs a warning and returns."""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
profiler_config = ProfilerConfig(
|
||
|
|
profiler="torch",
|
||
|
|
torch_profiler_dir="/path/to/traces",
|
||
|
|
)
|
||
|
|
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.profiler_config = profiler_config
|
||
|
|
worker.profiler = None
|
||
|
|
|
||
|
|
worker.profile(is_start=False)
|
||
|
|
|
||
|
|
mock_logger.warning.assert_called_once_with("Profiler was not started, nothing to stop.")
|
||
|
|
|
||
|
|
def test_lora_methods(self):
|
||
|
|
"""Test LoRA related methods"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
mock_model_runner = MagicMock()
|
||
|
|
worker.model_runner = mock_model_runner
|
||
|
|
|
||
|
|
# Set return values
|
||
|
|
mock_model_runner.add_lora.return_value = True
|
||
|
|
mock_model_runner.remove_lora.return_value = True
|
||
|
|
mock_model_runner.list_loras.return_value = {1, 2, 3}
|
||
|
|
mock_model_runner.pin_lora.return_value = True
|
||
|
|
|
||
|
|
# Test each method
|
||
|
|
mock_request = MagicMock()
|
||
|
|
self.assertTrue(worker.add_lora(mock_request))
|
||
|
|
mock_model_runner.add_lora.assert_called_once_with(mock_request)
|
||
|
|
|
||
|
|
self.assertTrue(worker.remove_lora(1))
|
||
|
|
mock_model_runner.remove_lora.assert_called_once_with(1)
|
||
|
|
|
||
|
|
self.assertEqual(worker.list_loras(), {1, 2, 3})
|
||
|
|
mock_model_runner.list_loras.assert_called_once()
|
||
|
|
|
||
|
|
self.assertTrue(worker.pin_lora(2))
|
||
|
|
mock_model_runner.pin_lora.assert_called_once_with(2)
|
||
|
|
|
||
|
|
def test_get_methods(self):
|
||
|
|
"""Test various get methods"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
mock_model_runner = MagicMock()
|
||
|
|
worker.model_runner = mock_model_runner
|
||
|
|
|
||
|
|
# Set return values
|
||
|
|
mock_model = MagicMock()
|
||
|
|
mock_kv_cache_spec = {"test": "spec"}
|
||
|
|
mock_pooling_tasks = ["task1", "task2"]
|
||
|
|
mock_supported_tasks = ("task1", "task2")
|
||
|
|
|
||
|
|
mock_model_runner.get_model.return_value = mock_model
|
||
|
|
mock_model_runner.get_kv_cache_spec.return_value = mock_kv_cache_spec
|
||
|
|
mock_model_runner.get_supported_pooling_tasks.return_value = mock_pooling_tasks
|
||
|
|
mock_model_runner.get_supported_tasks.return_value = mock_supported_tasks
|
||
|
|
|
||
|
|
# Test each get method
|
||
|
|
self.assertEqual(worker.get_model(), mock_model)
|
||
|
|
self.assertEqual(worker.get_kv_cache_spec(), mock_kv_cache_spec)
|
||
|
|
self.assertEqual(worker.get_supported_pooling_tasks(), mock_pooling_tasks)
|
||
|
|
self.assertEqual(worker.get_supported_tasks(), mock_supported_tasks)
|
||
|
|
|
||
|
|
def test_execute_dummy_batch(self):
|
||
|
|
"""Test execute_dummy_batch method"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.compilation_config = MagicMock()
|
||
|
|
worker.compilation_config.cudagraph_mode = MagicMock()
|
||
|
|
mock_model_runner = MagicMock()
|
||
|
|
mock_decode_token_per_req = mock_model_runner.decode_token_per_req
|
||
|
|
worker.model_runner = mock_model_runner
|
||
|
|
|
||
|
|
# Test execute_dummy_batch
|
||
|
|
worker.execute_dummy_batch()
|
||
|
|
|
||
|
|
# Verify call
|
||
|
|
mock_model_runner._dummy_run.assert_called_once_with(
|
||
|
|
num_tokens=mock_decode_token_per_req, uniform_decode=True
|
||
|
|
)
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.memory_profiling")
|
||
|
|
@patch("torch.npu.reset_peak_memory_stats")
|
||
|
|
@patch("torch.npu.empty_cache")
|
||
|
|
@patch("torch_npu.npu.memory_stats")
|
||
|
|
@patch("torch_npu.npu.mem_get_info")
|
||
|
|
@patch("vllm_ascend.worker.worker.logger")
|
||
|
|
def test_determine_available_memory_normal_case(
|
||
|
|
self,
|
||
|
|
mock_logger,
|
||
|
|
mock_torch_mem_get_info,
|
||
|
|
mock_torch_memory_stats,
|
||
|
|
mock_torch_empty_cache,
|
||
|
|
mock_torch_reset_peak_memory_stats,
|
||
|
|
mock_memory_profiling,
|
||
|
|
):
|
||
|
|
"""Test determine_available_memory normal case (no non-torch memory allocation)"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Mock memory_profiling context manager
|
||
|
|
mock_profile_result = MagicMock()
|
||
|
|
mock_profile_result.non_torch_increase = 1000
|
||
|
|
mock_profile_result.torch_peak_increase = 2000
|
||
|
|
mock_profile_result.weights_memory = 500
|
||
|
|
mock_profile_result.before_profile = MagicMock()
|
||
|
|
mock_profile_result.before_profile.torch_peak = 0
|
||
|
|
mock_profile_result.after_profile = MagicMock()
|
||
|
|
mock_profile_result.after_profile.free_memory = 6500
|
||
|
|
|
||
|
|
mock_context = MagicMock()
|
||
|
|
mock_context.__enter__ = MagicMock(return_value=mock_profile_result)
|
||
|
|
mock_context.__exit__ = MagicMock(return_value=False)
|
||
|
|
mock_memory_profiling.return_value = mock_context
|
||
|
|
|
||
|
|
# Mock init_snapshot
|
||
|
|
mock_init_snapshot = MagicMock()
|
||
|
|
mock_init_snapshot.free_memory = 8000
|
||
|
|
mock_init_snapshot.total_memory = 10000
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.init_snapshot = mock_init_snapshot
|
||
|
|
worker.requested_memory = 10000 * 0.8
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.model_runner.model_memory_usage = 500
|
||
|
|
worker.cache_config = MagicMock()
|
||
|
|
worker.cache_config.gpu_memory_utilization = 0.8
|
||
|
|
worker.cache_config.kv_cache_memory_bytes = None
|
||
|
|
worker.device = torch.device("npu:0")
|
||
|
|
|
||
|
|
# Mock torch.npu.memory_stats for profile_torch_peak
|
||
|
|
# profile_torch_peak = memory_stats()["allocated_bytes.all.peak"] = 2000
|
||
|
|
mock_torch_memory_stats.return_value = {"allocated_bytes.all.peak": 2000}
|
||
|
|
|
||
|
|
result = worker.determine_available_memory()
|
||
|
|
|
||
|
|
worker.model_runner.profile_run.assert_called_once()
|
||
|
|
|
||
|
|
# non_kv_cache_memory = non_torch_increase(1000) + torch_peak_increase(2000-0) + weights_memory(500) = 3500
|
||
|
|
# result = requested_memory(8000) - non_kv_cache_memory(3500) = 4500
|
||
|
|
expected_result = int(10000 * 0.8 - 3500)
|
||
|
|
self.assertEqual(result, expected_result)
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.memory_profiling")
|
||
|
|
@patch("torch.npu.reset_peak_memory_stats")
|
||
|
|
@patch("torch.npu.empty_cache")
|
||
|
|
@patch("torch_npu.npu.memory_stats")
|
||
|
|
@patch("torch_npu.npu.mem_get_info")
|
||
|
|
def test_determine_available_memory_with_non_torch_allocations(
|
||
|
|
self,
|
||
|
|
mock_torch_mem_get_info,
|
||
|
|
mock_torch_memory_stats,
|
||
|
|
mock_torch_empty_cache,
|
||
|
|
mock_torch_reset_peak_memory_stats,
|
||
|
|
mock_memory_profiling,
|
||
|
|
):
|
||
|
|
"""Test determine_available_memory with significant non-torch memory allocation"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Mock memory_profiling context manager with large non-torch allocation
|
||
|
|
mock_profile_result = MagicMock()
|
||
|
|
mock_profile_result.non_torch_increase = 4000
|
||
|
|
mock_profile_result.torch_peak_increase = 1500
|
||
|
|
mock_profile_result.weights_memory = 500
|
||
|
|
mock_profile_result.before_profile = MagicMock()
|
||
|
|
mock_profile_result.before_profile.torch_peak = 0
|
||
|
|
mock_profile_result.after_profile = MagicMock()
|
||
|
|
mock_profile_result.after_profile.free_memory = 4000
|
||
|
|
|
||
|
|
mock_context = MagicMock()
|
||
|
|
mock_context.__enter__ = MagicMock(return_value=mock_profile_result)
|
||
|
|
mock_context.__exit__ = MagicMock(return_value=False)
|
||
|
|
mock_memory_profiling.return_value = mock_context
|
||
|
|
|
||
|
|
# Mock init_snapshot
|
||
|
|
mock_init_snapshot = MagicMock()
|
||
|
|
mock_init_snapshot.free_memory = 8500
|
||
|
|
mock_init_snapshot.total_memory = 10000
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.init_snapshot = mock_init_snapshot
|
||
|
|
worker.requested_memory = 10000 * 0.9
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.model_runner.model_memory_usage = 500
|
||
|
|
worker.cache_config = MagicMock()
|
||
|
|
worker.cache_config.gpu_memory_utilization = 0.9
|
||
|
|
worker.cache_config.kv_cache_memory_bytes = None
|
||
|
|
worker.device = torch.device("npu:0")
|
||
|
|
|
||
|
|
mock_torch_memory_stats.return_value = {"allocated_bytes.all.peak": 1500}
|
||
|
|
|
||
|
|
result = worker.determine_available_memory()
|
||
|
|
|
||
|
|
# non_kv_cache_memory = non_torch_increase(4000) + torch_peak_increase(1500-0) + weights_memory(500) = 6000
|
||
|
|
# result = requested_memory(9000) - non_kv_cache_memory(6000) = 3000
|
||
|
|
expected_result = int(10000 * 0.9 - 6000)
|
||
|
|
self.assertEqual(result, expected_result)
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.memory_profiling")
|
||
|
|
@patch("torch.npu.mem_get_info")
|
||
|
|
@patch("torch.npu.reset_peak_memory_stats")
|
||
|
|
@patch("torch.npu.empty_cache")
|
||
|
|
def test_determine_available_memory_memory_profiling_error(
|
||
|
|
self, mock_torch_empty_cache, mock_torch_reset_peak_memory_stats, mock_torch_mem_get_info, mock_memory_profiling
|
||
|
|
):
|
||
|
|
"""Test determine_available_memory throws exception on memory profiling error"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Mock memory_profiling where free memory after profile > init free memory (error case)
|
||
|
|
mock_profile_result = MagicMock()
|
||
|
|
mock_profile_result.non_kv_cache_memory = 2000
|
||
|
|
mock_profile_result.after_profile = MagicMock()
|
||
|
|
mock_profile_result.after_profile.free_memory = 9000 # More free than init!
|
||
|
|
mock_profile_result.non_torch_increase = 0
|
||
|
|
mock_profile_result.torch_peak_increase = 0
|
||
|
|
mock_profile_result.weights_memory = 0
|
||
|
|
|
||
|
|
mock_context = MagicMock()
|
||
|
|
mock_context.__enter__ = MagicMock(return_value=mock_profile_result)
|
||
|
|
mock_context.__exit__ = MagicMock(return_value=False)
|
||
|
|
mock_memory_profiling.return_value = mock_context
|
||
|
|
|
||
|
|
mock_init_snapshot = MagicMock()
|
||
|
|
mock_init_snapshot.free_memory = 8500 # Less than after_profile free (9000)
|
||
|
|
mock_init_snapshot.total_memory = 10000
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.init_snapshot = mock_init_snapshot
|
||
|
|
worker.requested_memory = 10000 * 0.8
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.cache_config = MagicMock()
|
||
|
|
worker.cache_config.gpu_memory_utilization = 0.8
|
||
|
|
worker.cache_config.kv_cache_memory_bytes = None
|
||
|
|
worker.device = torch.device("npu:0")
|
||
|
|
|
||
|
|
# Test should throw assertion error
|
||
|
|
with self.assertRaises(AssertionError) as cm:
|
||
|
|
worker.determine_available_memory()
|
||
|
|
|
||
|
|
self.assertIn("Error in memory profiling", str(cm.exception))
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.memory_profiling")
|
||
|
|
@patch("torch.npu.reset_peak_memory_stats")
|
||
|
|
@patch("torch.npu.empty_cache")
|
||
|
|
@patch("torch_npu.npu.memory_stats")
|
||
|
|
@patch("torch_npu.npu.mem_get_info")
|
||
|
|
def test_determine_available_memory_negative_result(
|
||
|
|
self,
|
||
|
|
mock_torch_mem_get_info,
|
||
|
|
mock_torch_memory_stats,
|
||
|
|
mock_torch_empty_cache,
|
||
|
|
mock_torch_reset_peak_memory_stats,
|
||
|
|
mock_memory_profiling,
|
||
|
|
):
|
||
|
|
"""Test determine_available_memory returns 0 when result is negative"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Mock memory_profiling where non_kv_cache_memory > requested_memory
|
||
|
|
mock_profile_result = MagicMock()
|
||
|
|
mock_profile_result.non_torch_increase = 1000
|
||
|
|
mock_profile_result.torch_peak_increase = 9000
|
||
|
|
mock_profile_result.weights_memory = 500
|
||
|
|
mock_profile_result.before_profile = MagicMock()
|
||
|
|
mock_profile_result.before_profile.torch_peak = 0
|
||
|
|
mock_profile_result.after_profile = MagicMock()
|
||
|
|
mock_profile_result.after_profile.free_memory = 2000
|
||
|
|
|
||
|
|
mock_context = MagicMock()
|
||
|
|
mock_context.__enter__ = MagicMock(return_value=mock_profile_result)
|
||
|
|
mock_context.__exit__ = MagicMock(return_value=False)
|
||
|
|
mock_memory_profiling.return_value = mock_context
|
||
|
|
|
||
|
|
# Mock init_snapshot
|
||
|
|
mock_init_snapshot = MagicMock()
|
||
|
|
mock_init_snapshot.free_memory = 8500
|
||
|
|
mock_init_snapshot.total_memory = 10000
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.init_snapshot = mock_init_snapshot
|
||
|
|
worker.requested_memory = 10000 * 0.8
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.model_runner.model_memory_usage = 500
|
||
|
|
worker.cache_config = MagicMock()
|
||
|
|
worker.cache_config.gpu_memory_utilization = 0.8
|
||
|
|
worker.cache_config.kv_cache_memory_bytes = None
|
||
|
|
worker.device = torch.device("npu:0")
|
||
|
|
|
||
|
|
mock_torch_memory_stats.return_value = {"allocated_bytes.all.peak": 9000}
|
||
|
|
|
||
|
|
result = worker.determine_available_memory()
|
||
|
|
|
||
|
|
# non_kv_cache_memory = 1000 + 9000 + 500 = 10500
|
||
|
|
# available = requested(8000) - non_kv_cache(10500) = -2500
|
||
|
|
# upstream no longer clamps to 0, returns int(negative)
|
||
|
|
self.assertEqual(result, int(8000 - 10500))
|
||
|
|
|
||
|
|
def test_execute_model_first_rank(self):
|
||
|
|
"""Test execute_model method - first rank case"""
|
||
|
|
from vllm.v1.outputs import ModelRunnerOutput
|
||
|
|
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with (
|
||
|
|
patch.object(NPUWorker, "__init__", lambda x, **kwargs: None),
|
||
|
|
patch("vllm_ascend.worker.worker.get_pp_group") as mock_get_pp_group,
|
||
|
|
patch("vllm_ascend.worker.worker.get_ascend_config") as mock_get_ascend_config,
|
||
|
|
):
|
||
|
|
mock_ascend_config = MagicMock()
|
||
|
|
mock_ascend_config.msmonitor_use_daemon = False
|
||
|
|
mock_get_ascend_config.return_value = mock_ascend_config
|
||
|
|
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.vllm_config.parallel_config = MagicMock()
|
||
|
|
worker.vllm_config.parallel_config.distributed_executor_backend = "ray"
|
||
|
|
worker.profiler = None
|
||
|
|
worker._pp_send_work = []
|
||
|
|
|
||
|
|
# Set as first rank
|
||
|
|
mock_pp_group = MagicMock()
|
||
|
|
mock_pp_group.is_first_rank = True
|
||
|
|
mock_pp_group.is_last_rank = True
|
||
|
|
mock_get_pp_group.return_value = mock_pp_group
|
||
|
|
|
||
|
|
# Mock scheduler_output and return result
|
||
|
|
mock_scheduler_output = MagicMock()
|
||
|
|
mock_scheduler_output.total_num_scheduled_tokens = 1
|
||
|
|
mock_model_output = MagicMock(spec=ModelRunnerOutput)
|
||
|
|
worker.model_runner.execute_model.return_value = mock_model_output
|
||
|
|
|
||
|
|
# Test execute_model
|
||
|
|
result = worker.execute_model(mock_scheduler_output)
|
||
|
|
|
||
|
|
# Verify call
|
||
|
|
worker.model_runner.execute_model.assert_called_once_with(mock_scheduler_output, None)
|
||
|
|
self.assertEqual(result, mock_model_output)
|
||
|
|
|
||
|
|
def test_execute_model_calls_profiler_step_when_enabled(self):
|
||
|
|
"""Test execute_model steps the profiler before model execution."""
|
||
|
|
from vllm.v1.outputs import ModelRunnerOutput
|
||
|
|
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
call_order = []
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with (
|
||
|
|
patch.object(NPUWorker, "__init__", lambda x, **kwargs: None),
|
||
|
|
patch("vllm_ascend.worker.worker.get_pp_group") as mock_get_pp_group,
|
||
|
|
patch("vllm_ascend.worker.worker.get_ascend_config") as mock_get_ascend_config,
|
||
|
|
):
|
||
|
|
mock_ascend_config = MagicMock()
|
||
|
|
mock_ascend_config.msmonitor_use_daemon = False
|
||
|
|
mock_get_ascend_config.return_value = mock_ascend_config
|
||
|
|
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.vllm_config.parallel_config = MagicMock()
|
||
|
|
worker.vllm_config.parallel_config.distributed_executor_backend = "ray"
|
||
|
|
worker.profiler = MagicMock()
|
||
|
|
worker.profiler.step.side_effect = lambda: call_order.append("step")
|
||
|
|
worker._pp_send_work = []
|
||
|
|
|
||
|
|
mock_pp_group = MagicMock()
|
||
|
|
mock_pp_group.is_first_rank = True
|
||
|
|
mock_pp_group.is_last_rank = True
|
||
|
|
mock_get_pp_group.return_value = mock_pp_group
|
||
|
|
|
||
|
|
mock_scheduler_output = MagicMock()
|
||
|
|
mock_scheduler_output.total_num_scheduled_tokens = 1
|
||
|
|
mock_model_output = MagicMock(spec=ModelRunnerOutput)
|
||
|
|
|
||
|
|
def execute_model(*args):
|
||
|
|
call_order.append("execute")
|
||
|
|
return mock_model_output
|
||
|
|
|
||
|
|
worker.model_runner.execute_model.side_effect = execute_model
|
||
|
|
|
||
|
|
result = worker.execute_model(mock_scheduler_output)
|
||
|
|
|
||
|
|
worker.profiler.step.assert_called_once()
|
||
|
|
worker.model_runner.execute_model.assert_called_once_with(mock_scheduler_output, None)
|
||
|
|
self.assertEqual(call_order, ["step", "execute"])
|
||
|
|
self.assertEqual(result, mock_model_output)
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.get_ascend_config")
|
||
|
|
@patch("vllm_ascend.worker.worker.enable_sp", return_value=False)
|
||
|
|
@patch("vllm_ascend.worker.worker.get_pp_group")
|
||
|
|
@patch("vllm_ascend.worker.worker.get_tp_group")
|
||
|
|
def test_execute_model_middle_rank(
|
||
|
|
self, mock_get_tp_group, mock_get_pp_group, mock_enable_sp, mock_get_ascend_config
|
||
|
|
):
|
||
|
|
"""Test execute_model method - middle rank case"""
|
||
|
|
from vllm.sequence import IntermediateTensors
|
||
|
|
|
||
|
|
mock_ascend_config = MagicMock()
|
||
|
|
mock_ascend_config.msmonitor_use_daemon = False
|
||
|
|
mock_get_ascend_config.return_value = mock_ascend_config
|
||
|
|
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.vllm_config.parallel_config = MagicMock()
|
||
|
|
worker.vllm_config.parallel_config.distributed_executor_backend = "ray"
|
||
|
|
worker.profiler = None
|
||
|
|
worker._pp_send_work = []
|
||
|
|
|
||
|
|
# Set as middle rank (not first, not last)
|
||
|
|
mock_pp_group = MagicMock()
|
||
|
|
mock_pp_group.is_first_rank = False
|
||
|
|
mock_pp_group.is_last_rank = False
|
||
|
|
mock_get_pp_group.return_value = mock_pp_group
|
||
|
|
|
||
|
|
# Setup tensor reception data
|
||
|
|
mock_pp_group.irecv_tensor_dict.return_value = ({"tensor": "data"}, None, None)
|
||
|
|
mock_pp_group.isend_tensor_dict.return_value = []
|
||
|
|
|
||
|
|
# Mock return IntermediateTensors - use real type
|
||
|
|
mock_intermediate_output = MagicMock(spec=IntermediateTensors)
|
||
|
|
mock_intermediate_output.tensors = {"output_tensor": "data"}
|
||
|
|
mock_intermediate_output.kv_connector_output = None # Set to None to trigger return None
|
||
|
|
worker.model_runner.execute_model.return_value = mock_intermediate_output
|
||
|
|
|
||
|
|
mock_scheduler_output = MagicMock()
|
||
|
|
mock_scheduler_output.total_num_scheduled_tokens = 1
|
||
|
|
|
||
|
|
# Test execute_model
|
||
|
|
result = worker.execute_model(mock_scheduler_output)
|
||
|
|
|
||
|
|
# Verify tensor reception
|
||
|
|
mock_pp_group.irecv_tensor_dict.assert_called_once()
|
||
|
|
|
||
|
|
# Verify model execution with intermediate_tensors
|
||
|
|
# Second parameter should be AsyncIntermediateTensors instance
|
||
|
|
worker.model_runner.execute_model.assert_called_once()
|
||
|
|
args, kwargs = worker.model_runner.execute_model.call_args
|
||
|
|
self.assertEqual(args[0], mock_scheduler_output)
|
||
|
|
|
||
|
|
# Verify tensor sending
|
||
|
|
mock_pp_group.isend_tensor_dict.assert_called_once()
|
||
|
|
|
||
|
|
# Middle rank without kv_transfer_group should return None
|
||
|
|
self.assertIsNone(result)
|
||
|
|
|
||
|
|
def test_execute_model_external_launcher(self):
|
||
|
|
"""Test execute_model method - external_launcher mode"""
|
||
|
|
from vllm.v1.outputs import ModelRunnerOutput
|
||
|
|
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with (
|
||
|
|
patch.object(NPUWorker, "__init__", lambda x, **kwargs: None),
|
||
|
|
patch("vllm_ascend.worker.worker.get_pp_group") as mock_get_pp_group,
|
||
|
|
patch("vllm_ascend.worker.worker.get_ascend_config") as mock_get_ascend_config,
|
||
|
|
):
|
||
|
|
mock_ascend_config = MagicMock()
|
||
|
|
mock_ascend_config.msmonitor_use_daemon = False
|
||
|
|
mock_get_ascend_config.return_value = mock_ascend_config
|
||
|
|
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.vllm_config.parallel_config = MagicMock()
|
||
|
|
worker.vllm_config.parallel_config.distributed_executor_backend = "external_launcher"
|
||
|
|
worker.profiler = None
|
||
|
|
worker._pp_send_work = []
|
||
|
|
|
||
|
|
# Set as non-last rank
|
||
|
|
mock_pp_group = MagicMock()
|
||
|
|
mock_pp_group.is_first_rank = True
|
||
|
|
mock_pp_group.is_last_rank = False
|
||
|
|
mock_get_pp_group.return_value = mock_pp_group
|
||
|
|
|
||
|
|
# Mock return result
|
||
|
|
mock_scheduler_output = MagicMock()
|
||
|
|
mock_scheduler_output.total_num_scheduled_tokens = 1
|
||
|
|
mock_model_output = MagicMock(spec=ModelRunnerOutput)
|
||
|
|
worker.model_runner.execute_model.return_value = mock_model_output
|
||
|
|
|
||
|
|
# Test execute_model
|
||
|
|
result = worker.execute_model(mock_scheduler_output)
|
||
|
|
|
||
|
|
# In external_launcher mode, it doesn't enter middle processing logic, returns result directly
|
||
|
|
self.assertEqual(result, mock_model_output)
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.CaMemAllocator")
|
||
|
|
def test_load_model_with_sleep_mode(self, mock_allocator_class):
|
||
|
|
"""Test load_model method - with sleep mode enabled"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.vllm_config.model_config = MagicMock()
|
||
|
|
worker.vllm_config.model_config.enable_sleep_mode = True
|
||
|
|
worker.vllm_config.weight_transfer_config = None
|
||
|
|
worker.vllm_config.kv_transfer_config = None
|
||
|
|
|
||
|
|
# Setup allocator mock
|
||
|
|
mock_allocator = MagicMock()
|
||
|
|
mock_allocator.get_current_usage.return_value = 0
|
||
|
|
mock_context = MagicMock()
|
||
|
|
mock_allocator.use_memory_pool.return_value = mock_context
|
||
|
|
mock_allocator_class.get_instance.return_value = mock_allocator
|
||
|
|
|
||
|
|
# Test load_model
|
||
|
|
worker.load_model()
|
||
|
|
|
||
|
|
# Verify calls
|
||
|
|
mock_allocator_class.get_instance.assert_called_once()
|
||
|
|
mock_allocator.get_current_usage.assert_called_once()
|
||
|
|
mock_allocator.use_memory_pool.assert_called_once_with(tag="weights")
|
||
|
|
worker.model_runner.load_model.assert_called_once()
|
||
|
|
|
||
|
|
def test_load_model_without_sleep_mode(self):
|
||
|
|
"""Test load_model method - without sleep mode enabled"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.vllm_config.model_config = MagicMock()
|
||
|
|
worker.vllm_config.model_config.enable_sleep_mode = False
|
||
|
|
worker.vllm_config.weight_transfer_config = None
|
||
|
|
|
||
|
|
# Test load_model
|
||
|
|
worker.load_model()
|
||
|
|
|
||
|
|
# Verify calls
|
||
|
|
worker.model_runner.load_model.assert_called_once()
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.CaMemAllocator")
|
||
|
|
def test_load_model_sleep_mode_assertion_error(self, mock_allocator_class):
|
||
|
|
"""Test load_model method - assertion error in sleep mode"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.vllm_config.model_config = MagicMock()
|
||
|
|
worker.vllm_config.model_config.enable_sleep_mode = True
|
||
|
|
|
||
|
|
# Setup allocator mock - current usage is not 0
|
||
|
|
mock_allocator = MagicMock()
|
||
|
|
mock_allocator.get_current_usage.return_value = 100 # Non-zero value
|
||
|
|
mock_allocator_class.get_instance.return_value = mock_allocator
|
||
|
|
|
||
|
|
# Test should throw assertion error
|
||
|
|
with self.assertRaises(AssertionError) as cm:
|
||
|
|
worker.load_model()
|
||
|
|
|
||
|
|
self.assertIn("Sleep mode can only be", str(cm.exception))
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.set_random_seed")
|
||
|
|
@patch("vllm_ascend.worker.worker.get_ascend_device_type")
|
||
|
|
@patch("vllm_ascend.worker.worker.AscendDeviceType")
|
||
|
|
@patch("vllm_ascend.worker.worker.get_ascend_config")
|
||
|
|
@patch("vllm_ascend.worker.worker.logger")
|
||
|
|
@patch("vllm_ascend.worker.worker.NPUWorker._warm_up_atb")
|
||
|
|
def test_compile_or_warm_up_model_with_eager_mode(
|
||
|
|
self,
|
||
|
|
mock_warm_up_atb,
|
||
|
|
mock_logger,
|
||
|
|
mock_get_ascend_config,
|
||
|
|
mock_ascend_device_type,
|
||
|
|
mock_get_ascend_device_type,
|
||
|
|
mock_set_random_seed,
|
||
|
|
):
|
||
|
|
"""Test compile_or_warm_up_model method - eager mode"""
|
||
|
|
mock_ascend_config = MagicMock()
|
||
|
|
mock_ascend_config.ascend_compilation_config = MagicMock()
|
||
|
|
mock_ascend_config.ascend_compilation_config.enable_npugraph_ex = False
|
||
|
|
mock_ascend_config.enable_cpu_binding = False
|
||
|
|
mock_get_ascend_config.return_value = mock_ascend_config
|
||
|
|
mock_get_ascend_device_type.return_value = mock_ascend_device_type.A9B
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.model_config = MagicMock()
|
||
|
|
worker.model_config.enforce_eager = True
|
||
|
|
worker.model_config.seed = 12345
|
||
|
|
worker.cache_config = MagicMock()
|
||
|
|
worker.cache_config.kv_cache_memory_bytes = 1024
|
||
|
|
|
||
|
|
# Setup compilation config
|
||
|
|
worker.vllm_config.compilation_config = MagicMock()
|
||
|
|
worker.vllm_config.compilation_config.compile_sizes = [1, 4, 8, 16]
|
||
|
|
worker.vllm_config.compilation_config.cudagraph_capture_sizes = [4, 8]
|
||
|
|
|
||
|
|
# Test compile_or_warm_up_model
|
||
|
|
worker.compile_or_warm_up_model()
|
||
|
|
|
||
|
|
# Verify _dummy_run call count and order (by size descending)
|
||
|
|
expected_calls = [
|
||
|
|
unittest.mock.call(16),
|
||
|
|
unittest.mock.call(8),
|
||
|
|
unittest.mock.call(4),
|
||
|
|
unittest.mock.call(1),
|
||
|
|
]
|
||
|
|
worker.model_runner._dummy_run.assert_has_calls(expected_calls)
|
||
|
|
|
||
|
|
# Should not call capture_model in eager mode
|
||
|
|
worker.model_runner.capture_model.assert_not_called()
|
||
|
|
|
||
|
|
# Verify log output
|
||
|
|
self.assertEqual(mock_logger.info.call_count, 4)
|
||
|
|
|
||
|
|
# Verify atb warm up
|
||
|
|
mock_warm_up_atb.assert_called_once()
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.set_random_seed")
|
||
|
|
@patch("vllm_ascend.worker.worker.get_ascend_device_type")
|
||
|
|
@patch("vllm_ascend.worker.worker.AscendDeviceType")
|
||
|
|
@patch("vllm_ascend.worker.worker.CUDAGraphMode")
|
||
|
|
@patch("vllm_ascend.worker.worker.get_ascend_config")
|
||
|
|
@patch("vllm_ascend.worker.worker.logger")
|
||
|
|
@patch("vllm_ascend.worker.worker.NPUWorker._warm_up_atb")
|
||
|
|
def test_compile_or_warm_up_model_with_graph_capture(
|
||
|
|
self,
|
||
|
|
mock_warm_up_atb,
|
||
|
|
mock_logger,
|
||
|
|
mock_get_ascend_config,
|
||
|
|
mock_cudagraph_mode,
|
||
|
|
mock_ascend_device_type,
|
||
|
|
mock_get_ascend_device_type,
|
||
|
|
mock_set_random_seed,
|
||
|
|
):
|
||
|
|
"""Test compile_or_warm_up_model method - with graph capture enabled"""
|
||
|
|
mock_ascend_config = MagicMock()
|
||
|
|
mock_ascend_config.ascend_compilation_config = MagicMock()
|
||
|
|
mock_ascend_config.ascend_compilation_config.enable_npugraph_ex = False
|
||
|
|
mock_ascend_config.enable_cpu_binding = False
|
||
|
|
mock_get_ascend_config.return_value = mock_ascend_config
|
||
|
|
mock_get_ascend_device_type.return_value = mock_ascend_device_type.A9B
|
||
|
|
mock_cudagraph_mode.NONE = mock_cudagraph_mode.NONE
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.model_config = MagicMock()
|
||
|
|
worker.model_config.enforce_eager = False # Enable graph capture
|
||
|
|
worker.model_config.seed = 67890
|
||
|
|
worker.cache_config = MagicMock()
|
||
|
|
worker.cache_config.kv_cache_memory_bytes = 1024
|
||
|
|
|
||
|
|
# Setup compilation config
|
||
|
|
worker.vllm_config.compilation_config = MagicMock()
|
||
|
|
worker.vllm_config.compilation_config.compile_sizes = [1, 4, 8, 16]
|
||
|
|
worker.vllm_config.compilation_config.cudagraph_capture_sizes = [4, 8]
|
||
|
|
worker.vllm_config.compilation_config.cudagraph_mode = mock_cudagraph_mode.FULL
|
||
|
|
worker.vllm_config.compilation_config.get_compile_ranges.return_value = []
|
||
|
|
|
||
|
|
# Test compile_or_warm_up_model
|
||
|
|
worker.compile_or_warm_up_model()
|
||
|
|
|
||
|
|
# Verify only call _dummy_run for sizes not in cudagraph_capture_sizes
|
||
|
|
expected_calls = [unittest.mock.call(16), unittest.mock.call(1)]
|
||
|
|
worker.model_runner._dummy_run.assert_has_calls(expected_calls)
|
||
|
|
|
||
|
|
# Should call capture_model in non-eager mode
|
||
|
|
worker.model_runner.capture_model.assert_called_once()
|
||
|
|
|
||
|
|
# Verify atb warm up
|
||
|
|
mock_warm_up_atb.assert_called_once()
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.ensure_kv_transfer_initialized")
|
||
|
|
@patch("vllm_ascend.worker.worker.CaMemAllocator")
|
||
|
|
def test_initialize_from_config_with_sleep_mode(self, mock_allocator_class, mock_ensure_kv_transfer):
|
||
|
|
"""Test initialize_from_config method - with sleep mode enabled"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with (
|
||
|
|
patch.object(NPUWorker, "__init__", lambda x, **kwargs: None),
|
||
|
|
patch("vllm_ascend.worker.worker.ensure_kv_transfer_initialized"),
|
||
|
|
):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.vllm_config.speculative_config = None
|
||
|
|
worker.vllm_config.model_config = MagicMock()
|
||
|
|
worker.vllm_config.model_config.enable_sleep_mode = True
|
||
|
|
worker.vllm_config.kv_transfer_config = None
|
||
|
|
|
||
|
|
# Setup allocator mock
|
||
|
|
mock_allocator = MagicMock()
|
||
|
|
mock_context = MagicMock()
|
||
|
|
mock_allocator.use_memory_pool.return_value = mock_context
|
||
|
|
mock_allocator_class.get_instance.return_value = mock_allocator
|
||
|
|
|
||
|
|
# Create mock kv_cache_config
|
||
|
|
mock_kv_cache_config = MagicMock()
|
||
|
|
|
||
|
|
# Test initialize_from_config
|
||
|
|
worker.initialize_from_config(mock_kv_cache_config)
|
||
|
|
|
||
|
|
# Verify calls
|
||
|
|
mock_allocator_class.get_instance.assert_called_once()
|
||
|
|
mock_allocator.use_memory_pool.assert_called_once_with(tag="kv_cache")
|
||
|
|
worker.model_runner.initialize_kv_cache.assert_called_once_with(mock_kv_cache_config)
|
||
|
|
|
||
|
|
def test_acl_graph_sleep_wakeup_manager_sleep_resets_acl_graph_state(self):
|
||
|
|
from vllm_ascend.device_allocator.sleep_mem_optimized import AclGraphSleepWakeupManager
|
||
|
|
|
||
|
|
model_runner = MagicMock()
|
||
|
|
model_runner.use_aclgraph = True
|
||
|
|
graph_manager = MagicMock()
|
||
|
|
graph_manager.graphs = MagicMock()
|
||
|
|
graph_manager.pool = None
|
||
|
|
model_runner.cudagraph_manager = graph_manager
|
||
|
|
saver = AclGraphSleepWakeupManager(MagicMock(), lambda: model_runner)
|
||
|
|
with (
|
||
|
|
patch(
|
||
|
|
"vllm_ascend.device_allocator.sleep_mem_optimized.AclGraphSleepWakeupManager"
|
||
|
|
".clear_all_attention_workspaces"
|
||
|
|
) as mock_clear,
|
||
|
|
patch(
|
||
|
|
"vllm_ascend.device_allocator.sleep_mem_optimized.AclGraphSleepWakeupManager.reset_all_graph_params"
|
||
|
|
) as mock_reset,
|
||
|
|
):
|
||
|
|
saver.sleep()
|
||
|
|
mock_clear.assert_called_once()
|
||
|
|
mock_reset.assert_called_once()
|
||
|
|
graph_manager.graphs.clear.assert_called_once()
|
||
|
|
|
||
|
|
def test_hccl_sleep_wakeup_manager_sleep_waits_and_destroys(self):
|
||
|
|
from vllm_ascend.device_allocator.sleep_mem_optimized import HcclSleepWakeupManager
|
||
|
|
|
||
|
|
worker = MagicMock()
|
||
|
|
handle = MagicMock()
|
||
|
|
worker._pp_send_work = [handle]
|
||
|
|
saver = HcclSleepWakeupManager(MagicMock(), worker)
|
||
|
|
saver._destroyed = False
|
||
|
|
|
||
|
|
with (
|
||
|
|
patch("vllm_ascend.device_allocator.sleep_mem_optimized.torch.distributed.is_available", return_value=True),
|
||
|
|
patch(
|
||
|
|
"vllm_ascend.device_allocator.sleep_mem_optimized.torch.distributed.is_initialized",
|
||
|
|
return_value=True,
|
||
|
|
),
|
||
|
|
patch("vllm_ascend.device_allocator.sleep_mem_optimized.torch.npu.synchronize") as mock_synchronize,
|
||
|
|
patch(
|
||
|
|
"vllm_ascend.device_allocator.sleep_mem_optimized.HcclSleepWakeupManager.destroy_hccl",
|
||
|
|
return_value=2,
|
||
|
|
) as mock_destroy,
|
||
|
|
):
|
||
|
|
saver.sleep()
|
||
|
|
|
||
|
|
handle.wait.assert_called_once()
|
||
|
|
self.assertEqual(worker._pp_send_work, [])
|
||
|
|
mock_synchronize.assert_called_once()
|
||
|
|
mock_destroy.assert_called_once()
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.ensure_kv_transfer_initialized")
|
||
|
|
def test_initialize_from_config_without_sleep_mode(self, mock_ensure_kv_transfer):
|
||
|
|
"""Test initialize_from_config method - without sleep mode enabled"""
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with (
|
||
|
|
patch.object(NPUWorker, "__init__", lambda x, **kwargs: None),
|
||
|
|
patch("vllm_ascend.worker.worker.ensure_kv_transfer_initialized"),
|
||
|
|
):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.vllm_config.speculative_config = None
|
||
|
|
worker.vllm_config.model_config = MagicMock()
|
||
|
|
worker.vllm_config.model_config.enable_sleep_mode = False
|
||
|
|
worker.vllm_config.kv_transfer_config = None
|
||
|
|
|
||
|
|
# Create mock kv_cache_config
|
||
|
|
mock_kv_cache_config = MagicMock()
|
||
|
|
|
||
|
|
# Test initialize_from_config
|
||
|
|
worker.initialize_from_config(mock_kv_cache_config)
|
||
|
|
|
||
|
|
# Verify calls
|
||
|
|
worker.model_runner.initialize_kv_cache.assert_called_once_with(mock_kv_cache_config)
|
||
|
|
|
||
|
|
@patch("vllm_ascend.worker.worker.get_ascend_config")
|
||
|
|
@patch("vllm_ascend.worker.worker.enable_sp", return_value=False)
|
||
|
|
@patch("vllm_ascend.worker.worker.get_pp_group")
|
||
|
|
@patch("vllm_ascend.worker.worker.get_tp_group")
|
||
|
|
@patch("vllm_ascend.worker.worker.EMPTY_MODEL_RUNNER_OUTPUT")
|
||
|
|
def test_execute_model_kv_connector_not_finished(
|
||
|
|
self, mock_empty_output, mock_get_tp_group, mock_get_pp_group, mock_enable_sp, mock_get_ascend_config
|
||
|
|
):
|
||
|
|
"""Test execute_model method - kv_connector_output not finished sending/recving case"""
|
||
|
|
from vllm.sequence import IntermediateTensors
|
||
|
|
|
||
|
|
mock_ascend_config = MagicMock()
|
||
|
|
mock_ascend_config.msmonitor_use_daemon = False
|
||
|
|
mock_get_ascend_config.return_value = mock_ascend_config
|
||
|
|
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
# Create worker mock
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.vllm_config = MagicMock()
|
||
|
|
worker.vllm_config.parallel_config = MagicMock()
|
||
|
|
worker.vllm_config.parallel_config.distributed_executor_backend = "ray"
|
||
|
|
worker.profiler = None
|
||
|
|
worker._pp_send_work = []
|
||
|
|
|
||
|
|
# Set as middle rank (not first, not last)
|
||
|
|
mock_pp_group = MagicMock()
|
||
|
|
mock_pp_group.is_first_rank = False
|
||
|
|
mock_pp_group.is_last_rank = False
|
||
|
|
mock_get_pp_group.return_value = mock_pp_group
|
||
|
|
|
||
|
|
# Setup tensor reception data
|
||
|
|
mock_pp_group.irecv_tensor_dict.return_value = ({"tensor": "data"}, None, None)
|
||
|
|
mock_pp_group.isend_tensor_dict.return_value = []
|
||
|
|
|
||
|
|
# Create mock kv_connector_output - both finished_sending and finished_recving are False
|
||
|
|
mock_kv_connector_output = MagicMock()
|
||
|
|
mock_kv_connector_output.finished_sending = False
|
||
|
|
mock_kv_connector_output.finished_recving = False
|
||
|
|
|
||
|
|
# Mock return IntermediateTensors with kv_connector_output
|
||
|
|
mock_intermediate_output = MagicMock(spec=IntermediateTensors)
|
||
|
|
mock_intermediate_output.tensors = {"output_tensor": "data"}
|
||
|
|
mock_intermediate_output.kv_connector_output = mock_kv_connector_output
|
||
|
|
worker.model_runner.execute_model.return_value = mock_intermediate_output
|
||
|
|
|
||
|
|
mock_scheduler_output = MagicMock()
|
||
|
|
mock_scheduler_output.total_num_scheduled_tokens = 1
|
||
|
|
|
||
|
|
# Test execute_model
|
||
|
|
result = worker.execute_model(mock_scheduler_output)
|
||
|
|
|
||
|
|
# Verify tensor reception and sending
|
||
|
|
mock_pp_group.irecv_tensor_dict.assert_called_once()
|
||
|
|
mock_pp_group.isend_tensor_dict.assert_called_once()
|
||
|
|
|
||
|
|
# When both flags are False, return EMPTY_MODEL_RUNNER_OUTPUT directly.
|
||
|
|
self.assertEqual(result, mock_empty_output)
|
||
|
|
|
||
|
|
|
||
|
|
class TestNPUWorkerWeightUpdate(TestBase):
|
||
|
|
def _make_worker(self, engine=None):
|
||
|
|
from vllm_ascend.worker.worker import NPUWorker
|
||
|
|
|
||
|
|
with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None):
|
||
|
|
worker = NPUWorker()
|
||
|
|
worker.weight_transfer_engine = engine
|
||
|
|
worker._weight_update_active = False
|
||
|
|
worker._is_checkpoint_format = True
|
||
|
|
worker.device = torch.device("cpu")
|
||
|
|
worker.model_runner = MagicMock()
|
||
|
|
worker.model_runner.model = MagicMock()
|
||
|
|
worker.model_config = MagicMock()
|
||
|
|
return worker
|
||
|
|
|
||
|
|
def test_check_engine_raises_when_unconfigured(self):
|
||
|
|
worker = self._make_worker(engine=None)
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
worker.init_weight_transfer_engine({})
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
worker.start_weight_update()
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
worker.update_weights({})
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
worker.finish_weight_update()
|
||
|
|
|
||
|
|
def test_init_weight_transfer_engine_dispatches_to_engine(self):
|
||
|
|
engine = MagicMock()
|
||
|
|
engine.parse_init_info.return_value = "typed_init"
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
|
||
|
|
init_info = {"master_address": "127.0.0.1", "master_port": 12345}
|
||
|
|
worker.init_weight_transfer_engine(init_info)
|
||
|
|
|
||
|
|
engine.parse_init_info.assert_called_once_with(init_info)
|
||
|
|
engine.init_transfer_engine.assert_called_once_with("typed_init")
|
||
|
|
|
||
|
|
@patch("vllm.model_executor.model_loader.reload.initialize_layerwise_reload")
|
||
|
|
@patch.dict("os.environ", {"VLLM_ASCEND_ENABLE_NZ": "0"})
|
||
|
|
def test_start_weight_update_checkpoint_format(self, mock_init_reload):
|
||
|
|
engine = MagicMock()
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
|
||
|
|
worker.start_weight_update(is_checkpoint_format=True)
|
||
|
|
|
||
|
|
mock_init_reload.assert_called_once_with(worker.model_runner.model)
|
||
|
|
self.assertTrue(worker._weight_update_active)
|
||
|
|
self.assertTrue(worker._is_checkpoint_format)
|
||
|
|
|
||
|
|
@patch("vllm.model_executor.model_loader.reload.initialize_layerwise_reload")
|
||
|
|
@patch.dict("os.environ", {"VLLM_ASCEND_ENABLE_NZ": "0"})
|
||
|
|
def test_start_weight_update_kernel_format(self, mock_init_reload):
|
||
|
|
engine = MagicMock()
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
|
||
|
|
worker.start_weight_update(is_checkpoint_format=False)
|
||
|
|
|
||
|
|
mock_init_reload.assert_not_called()
|
||
|
|
self.assertTrue(worker._weight_update_active)
|
||
|
|
self.assertFalse(worker._is_checkpoint_format)
|
||
|
|
|
||
|
|
@patch.dict("os.environ", {"VLLM_ASCEND_ENABLE_NZ": "0"})
|
||
|
|
def test_start_weight_update_rejects_reentry(self):
|
||
|
|
engine = MagicMock()
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
worker._weight_update_active = True
|
||
|
|
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
worker.start_weight_update()
|
||
|
|
|
||
|
|
@patch.dict("os.environ", {"VLLM_ASCEND_ENABLE_NZ": "1"})
|
||
|
|
def test_start_weight_update_rejects_nz(self):
|
||
|
|
engine = MagicMock()
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
worker.start_weight_update()
|
||
|
|
|
||
|
|
def test_update_weights_requires_start(self):
|
||
|
|
engine = MagicMock()
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
worker.update_weights({"names": [], "dtype_names": [], "shapes": []})
|
||
|
|
|
||
|
|
@patch("torch.npu.synchronize", create=True)
|
||
|
|
@patch("vllm.model_executor.model_loader.reload.finalize_layerwise_reload")
|
||
|
|
@patch("vllm.model_executor.model_loader.reload.initialize_layerwise_reload")
|
||
|
|
@patch.dict("os.environ", {"VLLM_ASCEND_ENABLE_NZ": "0"})
|
||
|
|
def test_update_weights_checkpoint_format(self, mock_init_reload, mock_finalize_reload, mock_sync):
|
||
|
|
engine = MagicMock()
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
|
||
|
|
engine.parse_update_info.return_value = "typed_update"
|
||
|
|
worker._weight_update_active = True
|
||
|
|
worker._is_checkpoint_format = True
|
||
|
|
|
||
|
|
worker.update_weights({"foo": "bar"})
|
||
|
|
|
||
|
|
engine.parse_update_info.assert_called_once_with({"foo": "bar"})
|
||
|
|
engine.receive_weights.assert_called_once()
|
||
|
|
_, kwargs = engine.receive_weights.call_args
|
||
|
|
self.assertIs(kwargs["load_weights"], worker.model_runner.model.load_weights)
|
||
|
|
mock_sync.assert_called_once()
|
||
|
|
|
||
|
|
# reload lifecycle is split across start_weight_update / finish_weight_update
|
||
|
|
mock_init_reload.assert_not_called()
|
||
|
|
mock_finalize_reload.assert_not_called()
|
||
|
|
|
||
|
|
@patch("torch.npu.synchronize", create=True)
|
||
|
|
@patch.dict("os.environ", {"VLLM_ASCEND_ENABLE_NZ": "0"})
|
||
|
|
def test_update_weights_kernel_format(self, mock_sync):
|
||
|
|
engine = MagicMock()
|
||
|
|
|
||
|
|
def fake_receive(update_info, load_weights):
|
||
|
|
load_weights([("layer.weight", torch.zeros(2))])
|
||
|
|
|
||
|
|
engine.receive_weights.side_effect = fake_receive
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
param = torch.nn.Parameter(torch.ones(2), requires_grad=True)
|
||
|
|
worker.model_runner.model.get_parameter.return_value = param
|
||
|
|
|
||
|
|
engine.parse_update_info.return_value = "typed_update"
|
||
|
|
worker._weight_update_active = True
|
||
|
|
worker._is_checkpoint_format = False
|
||
|
|
|
||
|
|
worker.update_weights({"foo": "bar"})
|
||
|
|
|
||
|
|
worker.model_runner.model.get_parameter.assert_called_once_with("layer.weight")
|
||
|
|
torch.testing.assert_close(param.detach(), torch.zeros(2))
|
||
|
|
self.assertTrue(param.requires_grad)
|
||
|
|
|
||
|
|
@patch("vllm.model_executor.model_loader.reload.finalize_layerwise_reload")
|
||
|
|
def test_finish_weight_update_resets_state(self, mock_finalize_reload):
|
||
|
|
engine = MagicMock()
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
worker._weight_update_active = True
|
||
|
|
worker._is_checkpoint_format = True
|
||
|
|
|
||
|
|
worker.finish_weight_update()
|
||
|
|
|
||
|
|
mock_finalize_reload.assert_called_once_with(worker.model_runner.model, worker.model_config)
|
||
|
|
self.assertFalse(worker._weight_update_active)
|
||
|
|
self.assertTrue(worker._is_checkpoint_format)
|
||
|
|
|
||
|
|
def test_finish_without_start_raises(self):
|
||
|
|
engine = MagicMock()
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
worker.finish_weight_update()
|
||
|
|
|
||
|
|
def test_double_finish_raises(self):
|
||
|
|
engine = MagicMock()
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
worker._weight_update_active = True
|
||
|
|
worker._is_checkpoint_format = False
|
||
|
|
|
||
|
|
worker.finish_weight_update()
|
||
|
|
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
worker.finish_weight_update()
|
||
|
|
|
||
|
|
@patch("torch.npu.synchronize", create=True)
|
||
|
|
def test_update_after_finish_requires_restart(self, _mock_sync):
|
||
|
|
engine = MagicMock()
|
||
|
|
engine.parse_update_info.return_value = "typed"
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
worker._weight_update_active = True
|
||
|
|
worker._is_checkpoint_format = False
|
||
|
|
worker.finish_weight_update()
|
||
|
|
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
worker.update_weights({"names": [], "dtype_names": [], "shapes": []})
|
||
|
|
|
||
|
|
@patch("vllm.distributed.kv_transfer.ensure_kv_transfer_shutdown", create=True)
|
||
|
|
def test_shutdown_releases_engine(self, _mock_kv_shutdown):
|
||
|
|
engine = MagicMock()
|
||
|
|
worker = self._make_worker(engine=engine)
|
||
|
|
worker.profiler = None
|
||
|
|
|
||
|
|
worker.shutdown()
|
||
|
|
|
||
|
|
engine.shutdown.assert_called_once()
|