EngineX/xc-llm-kunlun
3
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[CI] Add UT CI (#157)

Browse Source 
Signed-off-by: Joeegin <3318329726@qq.com>
This commit is contained in:
Joeegin
2026-01-28 18:00:16 +08:00
committed by GitHub
parent d18df18499
commit c37ee19e3d
2 changed files with 213 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
.github/workflows/ut.yml vendored Normal file
View File

@@ -0,0 +1,57 @@
name: Unit Test
on:
pull_request:
branches:
- main
jobs:
test-kunlun:
runs-on:
labels:
- self-hosted
- Linux
- X64
- test-1 # Actions Runner Label
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Install vLLM-Kunlun Dependencies
run: |
pip install -r requirements.txt
python setup.py build
python setup.py develop
# Install the KL3-customized build of PyTorch
wget -O xpytorch-cp310-torch251-ubuntu2004-x64.run https://baidu-kunlun-public.su.bcebos.com/v1/baidu-kunlun-share/1130/xpytorch-cp310-torch251-ubuntu2004-x64.run?authorization=bce-auth-v1%2FALTAKypXxBzU7gg4Mk4K4c6OYR%2F2025-12-02T05%3A01%3A27Z%2F-1%2Fhost%2Ff3cf499234f82303891aed2bcb0628918e379a21e841a3fac6bd94afef491ff7
bash xpytorch-cp310-torch251-ubuntu2004-x64.run
# Install custom ops
pip install "https://baidu-kunlun-public.su.bcebos.com/v1/baidu-kunlun-share/1130/xtorch_ops-0.1.2209%2B6752ad20-cp310-cp310-linux_x86_64.whl?authorization=bce-auth-v1%2FALTAKypXxBzU7gg4Mk4K4c6OYR%2F2025-12-05T06%3A18%3A00Z%2F-1%2Fhost%2F14936c2b7e7c557c1400e4c467c79f7a9217374a7aa4a046711ac4d948f460cd"
# Install the KLX3 custom Triton build
pip install "https://cce-ai-models.bj.bcebos.com/v1/vllm-kunlun-0.11.0/triton-3.0.0%2Bb2cde523-cp310-cp310-linux_x86_64.whl"
# Install the AIAK custom ops library
pip install "https://cce-ai-models.bj.bcebos.com/XSpeedGate-whl/release_merge/20251219_152418/xspeedgate_ops-0.0.0-cp310-cp310-linux_x86_64.whl"
- name: Install vLLM
run: |
pip install vllm==0.11.0 --no-build-isolation --no-deps --no-deps --index-url https://pip.baidu-int.com/simple/
- name: Install Test Dependencies
run: |
pip install pytest
- name: Run Unit Test
run: |
echo "Running full suite..."
export XPU_VISIBLE_DEVICES=1
pytest \
-vs \
--cov=vllm_kunlun \
--cov-report=term-missing \
-p no:warnings tests/ut

156
tests/ut/test.py Normal file
View File

@@ -0,0 +1,156 @@
from unittest.mock import MagicMock, Mock, patch
import pytest
import torch
def test_import():
"""Test that the module can be imported successfully."""
from vllm_kunlun.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
assert TorchCompileWrapperWithCustomDispatcher is not None
def test_basic_instantiation():
"""Test basic wrapper instantiation with mocked dependencies."""
from vllm_kunlun.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
# Create a concrete implementation
class TestWrapper(TorchCompileWrapperWithCustomDispatcher):
def forward(self, x):
return x * 2
# Mock all the dependencies
mock_config = MagicMock()
mock_config.compilation_config.init_backend.return_value = "eager"
mock_config.compilation_config.inductor_compile_config = None
with patch('vllm.config.get_current_vllm_config', return_value=mock_config):
with patch('vllm.config.CompilationLevel') as mock_level:
mock_level.DYNAMO_ONCE = 1
with patch('torch.compile', side_effect=lambda func, **kwargs: func):
with patch('torch._dynamo.convert_frame.register_bytecode_hook'):
wrapper = TestWrapper(compilation_level=0)
# Verify basic attributes exist
assert hasattr(wrapper, 'vllm_config')
assert hasattr(wrapper, 'compiled_callable')
assert hasattr(wrapper, 'original_code_object')
assert hasattr(wrapper, 'compiled_codes')
assert isinstance(wrapper.compiled_codes, list)
def test_forward_call():
"""Test that the forward method can be called."""
from vllm_kunlun.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
class TestWrapper(TorchCompileWrapperWithCustomDispatcher):
def forward(self, x):
return x * 2
mock_config = MagicMock()
mock_config.compilation_config.init_backend.return_value = "eager"
mock_config.compilation_config.inductor_compile_config = None
with patch('vllm.config.get_current_vllm_config', return_value=mock_config):
with patch('vllm.config.CompilationLevel') as mock_level:
mock_level.DYNAMO_ONCE = 1
with patch('torch.compile', side_effect=lambda func, **kwargs: func):
with patch('torch._dynamo.convert_frame.register_bytecode_hook'):
wrapper = TestWrapper(compilation_level=0)
# Test calling the wrapper
input_tensor = torch.tensor([1.0, 2.0, 3.0])
result = wrapper(input_tensor)
expected = input_tensor * 2
assert torch.allclose(result, expected)
def test_custom_callable():
"""Test wrapper with custom compiled callable."""
from vllm_kunlun.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
class TestWrapper(TorchCompileWrapperWithCustomDispatcher):
def forward(self, x):
return x * 2
custom_func = Mock(return_value=torch.tensor([5.0]))
mock_config = MagicMock()
mock_config.compilation_config.init_backend.return_value = "eager"
with patch('vllm.config.get_current_vllm_config', return_value=mock_config):
with patch('vllm.config.CompilationLevel') as mock_level:
mock_level.DYNAMO_ONCE = 1
with patch('torch._dynamo.convert_frame.register_bytecode_hook'):
wrapper = TestWrapper(
compiled_callable=custom_func,
compilation_level=0
)
# Verify custom callable is used
assert wrapper.compiled_callable is custom_func
# Call should use custom callable
result = wrapper(torch.tensor([1.0]))
assert custom_func.called
def test_bytecode_hook_basic():
"""Test that bytecode hook can be called without errors."""
from vllm_kunlun.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
from types import CodeType
class TestWrapper(TorchCompileWrapperWithCustomDispatcher):
def forward(self, x):
return x * 2
mock_config = MagicMock()
mock_config.compilation_config.init_backend.return_value = "eager"
mock_config.compilation_config.inductor_compile_config = None
mock_config.compilation_config.local_cache_dir = None
with patch('vllm.config.get_current_vllm_config', return_value=mock_config):
with patch('vllm.config.CompilationLevel') as mock_level:
mock_level.DYNAMO_ONCE = 1
with patch('torch.compile', side_effect=lambda func, **kwargs: func):
with patch('torch._dynamo.convert_frame.register_bytecode_hook'):
wrapper = TestWrapper(compilation_level=0)
# Test with wrong code object (should be ignored)
wrong_code = MagicMock(spec=CodeType)
new_code = MagicMock(spec=CodeType)
initial_count = len(wrapper.compiled_codes)
wrapper.bytecode_hook(wrong_code, new_code)
# Should not add anything
assert len(wrapper.compiled_codes) == initial_count
def test_use_custom_dispatcher_flag():
"""Test that use_custom_dispatcher flag is set based on compilation_level."""
from vllm_kunlun.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
class TestWrapper(TorchCompileWrapperWithCustomDispatcher):
def forward(self, x):
return x * 2
mock_config = MagicMock()
mock_config.compilation_config.init_backend.return_value = "eager"
mock_config.compilation_config.inductor_compile_config = None
with patch('vllm.config.get_current_vllm_config', return_value=mock_config):
with patch('vllm.config.CompilationLevel') as mock_level:
mock_level.DYNAMO_ONCE = 1
with patch('torch.compile', side_effect=lambda func, **kwargs: func):
with patch('torch._dynamo.convert_frame.register_bytecode_hook'):
# Test with low level
wrapper_low = TestWrapper(compilation_level=0)
assert wrapper_low.use_custom_dispatcher is False
# Test with high level
wrapper_high = TestWrapper(compilation_level=2)
assert wrapper_high.use_custom_dispatcher is True
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])