hao-ai-lab-FastVideo/fastvideo/tests/worker/test_gpu_worker.py

import os
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
import torch

import fastvideo.envs as envs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.pipelines import ForwardBatch
from fastvideo.tests.api.config_snapshot import isolated_environment
from fastvideo.worker.gpu_worker import Worker, _log_cuda_device_uuid


def test_cuda_device_uuid_receipt_is_disabled_without_nvtx_profiling(monkeypatch, env_overrides) -> None:
    """Avoid NVIDIA property access during ordinary worker initialization."""
    get_device_properties = Mock()
    env_overrides.enter_context(envs.FASTVIDEO_NVTX_PROFILE.override(False))
    monkeypatch.setattr(torch.cuda, "get_device_properties", get_device_properties)

    _log_cuda_device_uuid(0, torch.device("cuda:0"))

    get_device_properties.assert_not_called()


def test_cuda_device_uuid_receipt_identifies_profiled_worker(monkeypatch, env_overrides) -> None:
    """Bind one profiled worker rank to its NVIDIA device UUID in logs."""
    log_info = Mock()
    env_overrides.enter_context(envs.FASTVIDEO_NVTX_PROFILE.override(True))
    monkeypatch.setattr(torch.cuda, "get_device_properties", lambda device: SimpleNamespace(uuid="device-uuid"))
    monkeypatch.setattr("fastvideo.worker.gpu_worker.logger.info", log_info)

    _log_cuda_device_uuid(2, torch.device("cuda:0"))

    log_info.assert_called_once_with(
        "Worker %d CUDA device UUID: GPU-%s",
        2,
        "device-uuid",
        local_main_process_only=False,
    )


@pytest.mark.parametrize("executor_backend", ["mp", "ray"])
def test_init_device_binds_the_worker_device_and_keeps_the_resolved_config(monkeypatch, env_overrides,
                                                                          executor_backend: str) -> None:
    """The worker binds its own device and builds the pipeline on the config it received, without a device query."""
    events = []
    with isolated_environment():
        args = resolve_inference_config({
            "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
            "engine": {
                "num_gpus": 1,
                "execution_backend": executor_backend
            },
        })
    probe = Mock(side_effect=AssertionError("device probe ran in the worker"))
    monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", probe)
    worker = Worker(args, local_rank=3, rank=3, distributed_init_method="env://")

    env_overrides.enter_context(envs.override_external("LOCAL_RANK", "0"))
    monkeypatch.setattr("fastvideo.platforms.current_platform.is_cuda_alike", lambda: True)
    monkeypatch.setattr("fastvideo.platforms.current_platform.is_cuda", lambda: False)
    monkeypatch.setattr(torch.cuda, "set_device", lambda device: events.append(("set_device", device.index)))
    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (123, 456))
    monkeypatch.setattr(
        "fastvideo.worker.gpu_worker.maybe_init_distributed_environment_and_model_parallel",
        lambda *args: events.append(("distributed", None)),
    )
    monkeypatch.setattr("fastvideo.worker.gpu_worker.build_pipeline", lambda args: events.append(("pipeline", None)))

    worker.init_device()

    assert events == [
        ("set_device", 3),
        ("distributed", None),
        ("pipeline", None),
    ]
    assert os.environ["LOCAL_RANK"] == "3"
    assert worker.device == torch.device("cuda:3")
    assert worker.init_gpu_memory == 123
    assert worker.resolved_config is args
    probe.assert_not_called()


def _worker_returning(output_batch: ForwardBatch) -> Worker:
    worker = Worker.__new__(Worker)
    worker.resolved_config = SimpleNamespace()
    worker.pipeline = SimpleNamespace(forward=lambda batch, args: output_batch)
    return worker


def test_execute_forward_runs_the_pipeline_on_the_worker_config():
    worker = Worker.__new__(Worker)
    worker.resolved_config = SimpleNamespace()
    seen = []
    worker.pipeline = SimpleNamespace(forward=lambda batch, args: seen.append(args) or batch)

    worker.execute_forward(ForwardBatch(data_type="video", return_frames=True), _request_config())

    assert seen == [worker.resolved_config]


def _request_config(output_type: str = "pil") -> SimpleNamespace:
    """Stand-in for the request's resolved config with the typed path that ``Worker.execute_forward`` reads."""
    return SimpleNamespace(pipeline=SimpleNamespace(output_type=output_type))


def test_execute_forward_drops_metadata_only_output_before_transport():
    output = torch.ones((1, 3, 2, 4, 4))
    output_batch = ForwardBatch(data_type="video", output=output)
    worker = _worker_returning(output_batch)
    request_batch = ForwardBatch(data_type="video", save_video=False, return_frames=False)

    result = worker.execute_forward(request_batch, _request_config())

    assert result.output is not None
    assert result.output.device.type == "cpu"
    assert result.output.numel() == 0


def test_execute_forward_preserves_missing_metadata_only_output():
    output_batch = ForwardBatch(data_type="video", output=None)
    worker = _worker_returning(output_batch)
    request_batch = ForwardBatch(data_type="video", save_video=False, return_frames=False)

    result = worker.execute_forward(request_batch, _request_config())

    assert result.output is None


def test_execute_forward_drops_save_only_latent_output():
    output = torch.ones((1, 16, 1, 2, 2))
    output_batch = ForwardBatch(data_type="video", output=output)
    worker = _worker_returning(output_batch)
    request_batch = ForwardBatch(data_type="video", save_video=True, return_frames=False)

    result = worker.execute_forward(request_batch, _request_config(output_type="latent"))

    assert result.output is not None
    assert result.output.device.type == "cpu"
    assert result.output.numel() == 0


def test_execute_forward_drops_save_only_audio_placeholder():
    output = torch.ones((1, 3, 1, 8, 8))
    output_batch = ForwardBatch(data_type="audio", output=output, extra={"audio_only": True})
    worker = _worker_returning(output_batch)
    request_batch = ForwardBatch(data_type="audio", save_video=True, return_frames=False)

    result = worker.execute_forward(request_batch, _request_config())

    assert result.output is not None
    assert result.output.device.type == "cpu"
    assert result.output.numel() == 0


@pytest.mark.parametrize(
    ("save_video", "return_frames"),
    [
        (True, False),
        (False, True),
        (True, True),
    ],
)
def test_execute_forward_preserves_requested_output(save_video, return_frames):
    output = torch.ones((1, 3, 2, 4, 4))
    output_batch = ForwardBatch(data_type="video", output=output)
    worker = _worker_returning(output_batch)
    request_batch = ForwardBatch(data_type="video", save_video=save_video, return_frames=return_frames)

    result = worker.execute_forward(request_batch, _request_config())

    assert result.output is output