vllm-omni/tests/diffusion/test_diffusion_worker.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""
Unit tests for DiffusionWorker class.

This module tests the DiffusionWorker implementation:
- sleep: Putting worker into sleep mode (levels 1 and 2)
- wake_up: Waking worker from sleep mode
"""

import pytest
import torch
from pytest_mock import MockerFixture

from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker

pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]


def patch_cumem_allocator(mocker: MockerFixture):
    mock_allocator_class = mocker.Mock()
    mocker.patch(
        "vllm_omni.diffusion.worker.diffusion_worker._get_cumem_allocator_class",
        return_value=mock_allocator_class,
    )
    return mock_allocator_class


@pytest.fixture
def mock_od_config(mocker: MockerFixture):
    """Create a mock OmniDiffusionConfig."""
    config = mocker.Mock()
    config.num_gpus = 1
    config.master_port = 12345
    config.enable_sleep_mode = False
    config.cache_backend = None
    config.cache_config = None
    config.model = "test-model"
    return config


@pytest.fixture
def mock_gpu_worker(mocker: MockerFixture, mock_od_config):
    """Create a DiffusionWorker with mocked initialization."""
    mocker.patch.object(DiffusionWorker, "init_device")
    mocker.patch.object(DiffusionWorker, "load_model")
    worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
    # Mock the model_runner with pipeline
    worker.model_runner = mocker.Mock()
    worker.model_runner.pipeline = mocker.Mock()
    worker.device = torch.device("cuda", 0)
    worker._sleep_saved_buffers = {}
    return worker


class TestDiffusionWorkerSleep:
    """Test DiffusionWorker.sleep method."""

    @pytest.fixture(autouse=True)
    def setup_allocator(self, mocker: MockerFixture):
        """
        Unified interception of Allocators, and provision of default security values.
        """
        self.mock_allocator_class = patch_cumem_allocator(mocker)
        self.mock_allocator = mocker.Mock()
        self.mock_allocator_class.get_instance.return_value = self.mock_allocator
        self.mock_allocator.get_current_usage.return_value = 4 * 1024**3
        self.mock_allocator.sleep = mocker.Mock()

    def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
        """Test sleep mode level 1 (offload weights only)."""
        mock_allocator_class = patch_cumem_allocator(mocker)
        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
        mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
        mock_platform.get_device_total_memory.return_value = 80 * 1024**3
        mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")

        # Setup process-scoped memory mocks
        # Before sleep: 3GB used
        # After sleep: 1GB used (freed 2GB)
        initial_usage = 3 * 1024**3
        mock_get_process_memory.side_effect = [
            initial_usage,
            1 * 1024**3,
        ]

        # Setup allocator mock
        mock_allocator = mocker.Mock()
        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
        mock_allocator.sleep = mocker.Mock()
        mock_allocator.get_current_usage.return_value = initial_usage

        # Call sleep with level 1
        result = mock_gpu_worker.sleep(level=1)

        # Verify sleep was called with correct tags
        mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
        assert bool(result) is True
        # Verify buffers were NOT saved (level 1 doesn't save buffers)
        assert len(mock_gpu_worker._sleep_saved_buffers) == 0

    def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker):
        """Test sleep mode level 2 (offload all, save buffers)."""
        mock_allocator_class = patch_cumem_allocator(mocker)
        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
        mock_platform.get_free_memory.side_effect = [5 * 1024**3, 10 * 1024**3]
        mock_platform.get_device_total_memory.return_value = 80 * 1024**3
        mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")

        # Setup process-scoped memory mocks
        initial_usage = 5 * 1024**3
        mock_get_process_memory.side_effect = [
            initial_usage,  # Before sleep
            1 * 1024**3,  # After sleep (freed 4GB)
        ]

        # Setup allocator mock
        mock_allocator = mocker.Mock()
        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
        mock_allocator.sleep = mocker.Mock()
        mock_allocator.get_current_usage.return_value = initial_usage

        # Mock pipeline buffers
        mock_buffer1 = torch.randn(10, 10)
        mock_buffer2 = torch.randn(20, 20)
        mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
            return_value=[
                ("buffer1", mock_buffer1),
                ("buffer2", mock_buffer2),
            ]
        )

        # Call sleep with level 2
        result = mock_gpu_worker.sleep(level=2)

        # Verify sleep was called with empty tags (offload all)
        mock_allocator.sleep.assert_called_once_with(offload_tags=tuple())
        assert bool(result) is True

        # Verify buffers were saved
        assert len(mock_gpu_worker._sleep_saved_buffers) == 2
        assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
        assert "buffer2" in mock_gpu_worker._sleep_saved_buffers

    def test_sleep_level_2_releases_captures_through_the_model_runner(self, mocker: MockerFixture, mock_gpu_worker):
        """Level 2 discards the memory captures were recorded against, and the
        runner owns the release. The worker asks it once and does not reach past
        it into the pipeline."""
        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
        mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
        mock_platform.get_device_total_memory.return_value = 80 * 1024**3
        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory", side_effect=[0, 0])
        mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(return_value=[])

        mock_gpu_worker.sleep(level=2)

        mock_gpu_worker.model_runner.release_captured_graphs.assert_called_once_with()
        mock_gpu_worker.model_runner.pipeline.release_captured_graphs.assert_not_called()

    def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker):
        """Test that sleep validates memory was actually freed."""
        mock_allocator_class = patch_cumem_allocator(mocker)
        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
        mock_platform.get_free_memory.return_value = 10 * 1024**3
        mock_platform.get_device_total_memory.return_value = 80 * 1024**3
        mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")

        # Simulate process memory increase (should trigger assertion error)
        initial_usage = 1 * 1024**3
        mock_get_process_memory.side_effect = [
            initial_usage,  # Before sleep: 1GB used
            3 * 1024**3,  # After sleep: 3GB used (negative freed)
        ]

        mock_allocator = mocker.Mock()
        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
        mock_allocator.sleep = mocker.Mock()
        mock_allocator.get_current_usage.return_value = initial_usage

        # This should raise an assertion error
        result = mock_gpu_worker.sleep(level=1)
        assert result == initial_usage

    def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker):
        """Test sleep uses device-scoped fallback when NVML is unavailable."""

        mock_allocator_class = patch_cumem_allocator(mocker)
        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
        mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
        mock_get_process_memory.side_effect = [None, None]
        mock_platform.get_free_memory.side_effect = [
            1 * 1024**3,  # Before sleep
            3 * 1024**3,  # After sleep
        ]
        mock_platform.get_device_total_memory.return_value = 8 * 1024**3

        mock_allocator = mocker.Mock()
        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
        mock_allocator.sleep = mocker.Mock()
        mock_allocator.get_current_usage.return_value = 2 * 1024**3

        result = mock_gpu_worker.sleep(level=1)

        mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
        assert bool(result) is True


class TestDiffusionWorkerWakeUp:
    """Test DiffusionWorker.wake_up method."""

    def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker):
        """Test wake_up without saved buffers (level 1 sleep)."""
        mock_allocator_class = patch_cumem_allocator(mocker)

        # Setup allocator mock
        mock_allocator = mocker.Mock()
        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
        mock_allocator.wake_up = mocker.Mock()
        mock_allocator.get_current_usage.return_value = 10 * 1024**3

        # Ensure no saved buffers
        mock_gpu_worker._sleep_saved_buffers = {}

        # Call wake_up
        result = mock_gpu_worker.wake_up(tags=["weights"])

        # Verify allocator.wake_up was called
        mock_allocator.wake_up.assert_called_once_with(["weights"])
        assert bool(result) is True

    def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker):
        """Test wake_up with saved buffers (level 2 sleep)."""
        mock_allocator_class = patch_cumem_allocator(mocker)

        # Setup allocator mock
        mock_allocator = mocker.Mock()
        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
        mock_allocator.wake_up = mocker.Mock()
        mock_allocator.get_current_usage.return_value = 10 * 1024**3

        # Create saved buffers
        saved_buffer1 = torch.randn(10, 10)
        saved_buffer2 = torch.randn(20, 20)
        mock_gpu_worker._sleep_saved_buffers = {
            "buffer1": saved_buffer1,
            "buffer2": saved_buffer2,
        }

        # Mock pipeline buffers (these will be restored)
        mock_buffer1 = mocker.Mock()
        mock_buffer1.data = mocker.Mock()
        mock_buffer2 = mocker.Mock()
        mock_buffer2.data = mocker.Mock()

        mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
            return_value=[
                ("buffer1", mock_buffer1),
                ("buffer2", mock_buffer2),
            ]
        )

        # Call wake_up
        result = mock_gpu_worker.wake_up(tags=None)

        # Verify allocator.wake_up was called
        mock_allocator.wake_up.assert_called_once_with(None)

        # Verify buffers were restored
        mock_buffer1.data.copy_.assert_called_once()
        mock_buffer2.data.copy_.assert_called_once()

        # Verify saved buffers were cleared
        assert len(mock_gpu_worker._sleep_saved_buffers) == 0
        assert bool(result) is True

    def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker):
        """Test wake_up only restores buffers that were saved."""
        mock_allocator_class = patch_cumem_allocator(mocker)

        # Setup allocator mock
        mock_allocator = mocker.Mock()
        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
        mock_allocator.wake_up = mocker.Mock()
        mock_allocator.get_current_usage.return_value = 10 * 1024**3

        # Only save buffer1, not buffer2
        saved_buffer1 = torch.randn(10, 10)
        mock_gpu_worker._sleep_saved_buffers = {
            "buffer1": saved_buffer1,
        }

        # Mock pipeline has both buffers
        mock_buffer1 = mocker.Mock()
        mock_buffer1.data = mocker.Mock()
        mock_buffer2 = mocker.Mock()
        mock_buffer2.data = mocker.Mock()

        mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
            return_value=[
                ("buffer1", mock_buffer1),
                ("buffer2", mock_buffer2),
            ]
        )

        # Call wake_up
        result = mock_gpu_worker.wake_up()

        # Verify only buffer1 was restored
        mock_buffer1.data.copy_.assert_called_once()
        # buffer2 should NOT be restored since it wasn't saved
        mock_buffer2.data.copy_.assert_not_called()

        assert result is True


class TestDiffusionWorkerInitLoraManager:
    """Test DiffusionWorker.init_lora_manager method."""

    def test_skips_when_lora_already_fused(self, mocker: MockerFixture, mock_gpu_worker):
        """When pipeline.lora_is_fused is True, init_lora_manager returns immediately."""
        from vllm_omni.diffusion.lora.manager import LoRABackend

        mock_gpu_worker.od_config.lora_backend = LoRABackend.DISTILL
        mock_gpu_worker.od_config.lora_path = "/path/to/lora.safetensors"
        mock_gpu_worker.od_config.lora_scale = 1.0

        pipeline = mock_gpu_worker.model_runner.pipeline
        pipeline.lora_is_fused = True
        pipeline.load_lora_weights = mocker.Mock()

        mock_gpu_worker.init_lora_manager()

        pipeline.load_lora_weights.assert_not_called()
        assert mock_gpu_worker.lora_manager is None

    def test_distill_backend_fuses_and_marks_flag_when_not_pre_fused(self, mocker: MockerFixture, mock_gpu_worker):
        """When lora_is_fused is False, distill backend calls load_lora_weights and sets flag."""
        from vllm_omni.diffusion.lora.manager import LoRABackend

        mock_gpu_worker.od_config.lora_backend = LoRABackend.DISTILL
        mock_gpu_worker.od_config.lora_path = "/path/to/lora.safetensors"
        mock_gpu_worker.od_config.lora_scale = 1.0

        pipeline = mock_gpu_worker.model_runner.pipeline
        pipeline.lora_is_fused = False
        pipeline.load_lora_weights = mocker.Mock()

        mock_gpu_worker.init_lora_manager()

        pipeline.load_lora_weights.assert_called_once_with("/path/to/lora.safetensors")
        assert pipeline.lora_is_fused is True