test.py

from types import SimpleNamespace
from unittest.mock import MagicMock, Mock, patch

import pytest
import torch


def _make_platform_config(hf_config):
    from vllm.config import CUDAGraphMode

    return SimpleNamespace(
        parallel_config=SimpleNamespace(worker_cls="manual", data_parallel_size=1),
        model_config=SimpleNamespace(
            use_mla=False,
            hf_config=hf_config,
            enforce_eager=False,
        ),
        cache_config=SimpleNamespace(block_size=16),
        speculative_config=None,
        compilation_config=SimpleNamespace(
            cudagraph_mode=CUDAGraphMode.NONE,
            pass_config=SimpleNamespace(enable_fusion=True),
            backend=None,
            custom_ops=[],
        ),
    )


def _check_platform_config(vllm_config):
    import vllm.envs as envs

    from vllm_kunlun.platforms.kunlun import KunlunPlatform

    with patch.object(envs, "VLLM_ALL2ALL_BACKEND", None, create=True):
        KunlunPlatform.check_and_update_config(vllm_config)


def test_qwen3_vl_text_config_inherits_top_level_tie_word_embeddings():
    text_config = SimpleNamespace()
    hf_config = SimpleNamespace(
        architectures=["Qwen3VLForConditionalGeneration"],
        text_config=text_config,
        tie_word_embeddings=False,
    )

    _check_platform_config(_make_platform_config(hf_config))

    assert text_config.tie_word_embeddings is False


def test_qwen3_vl_text_config_existing_tie_word_embeddings_is_preserved():
    text_config = SimpleNamespace(tie_word_embeddings=True)
    hf_config = SimpleNamespace(
        architectures=["Qwen3VLForConditionalGeneration"],
        text_config=text_config,
        tie_word_embeddings=False,
    )

    _check_platform_config(_make_platform_config(hf_config))

    assert text_config.tie_word_embeddings is True


def test_non_qwen3_vl_text_config_is_not_modified():
    text_config = SimpleNamespace()
    hf_config = SimpleNamespace(
        architectures=["Qwen3ForCausalLM"],
        text_config=text_config,
        tie_word_embeddings=False,
    )

    _check_platform_config(_make_platform_config(hf_config))

    assert not hasattr(text_config, "tie_word_embeddings")


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]))  # noqa
                assert custom_func.called


def test_bytecode_hook_basic():
    """Test that bytecode hook can be called without errors."""
    from types import CodeType

    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
    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"])