From c37ee19e3d87c05d58ca494c69a948142462ae96 Mon Sep 17 00:00:00 2001
From: Joeegin <80816296+Joeegin@users.noreply.github.com>
Date: Wed, 28 Jan 2026 18:00:16 +0800
Subject: [PATCH] [CI] Add UT CI (#157)

Signed-off-by: Joeegin <3318329726@qq.com>
---
 .github/workflows/ut.yml |  57 ++++++++++++++
 tests/ut/test.py         | 156 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 213 insertions(+)
 create mode 100644 .github/workflows/ut.yml
 create mode 100644 tests/ut/test.py

diff --git a/.github/workflows/ut.yml b/.github/workflows/ut.yml
new file mode 100644
index 0000000..8ea5c23
--- /dev/null
+++ b/.github/workflows/ut.yml
@@ -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
\ No newline at end of file
diff --git a/tests/ut/test.py b/tests/ut/test.py
new file mode 100644
index 0000000..e5ea156
--- /dev/null
+++ b/tests/ut/test.py
@@ -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"])
\ No newline at end of file