Boyang ZHONG
dc72a35f4e
For n > 1, the image endpoints returned one preview grid instead of n images. The endpoints now run one model batch, save each decoded sample as its own file, and return one URL or base64 item for each image. This applies to generations and edits. A keyword-only return_samples field (default False) is added to OutputConfig, SamplingParam and ForwardBatch. Batched image requests set it and do not build preview frames. The worker keeps the sample tensor for these requests, and non-output ranks clear the field. Defaults and positional arguments do not change.
193 lines
8.4 KiB
Python
193 lines
8.4 KiB
Python
import os
|
|
from datetime import timedelta
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import fastvideo.envs as envs
|
|
import fastvideo.platforms as platforms
|
|
from fastvideo.fastvideo_args import FastVideoArgs
|
|
from fastvideo.pipelines import ForwardBatch
|
|
from fastvideo.worker.gpu_worker import Worker, _log_cuda_device_uuid
|
|
|
|
|
|
@pytest.mark.parametrize("dist_timeout", [0, -1])
|
|
def test_fastvideo_args_rejects_non_positive_dist_timeout(dist_timeout) -> None:
|
|
with pytest.raises(ValueError, match="greater than zero seconds"):
|
|
FastVideoArgs(model_path="test", dist_timeout=dist_timeout)
|
|
|
|
|
|
def test_worker_threads_dist_timeout_to_distributed_groups(monkeypatch, env_overrides) -> None:
|
|
init_distributed = Mock()
|
|
# init_device() writes LOCAL_RANK; restore it after the test.
|
|
env_overrides.enter_context(envs.override_external("LOCAL_RANK", "0"))
|
|
monkeypatch.setattr("fastvideo.worker.gpu_worker.get_local_torch_device", lambda: torch.device("cpu"))
|
|
# Fake the platform behind fastvideo.platforms.current_platform, the seam the
|
|
# attention-selector tests use. Patching current_platform itself would fetch
|
|
# it through the module's __getattr__, and monkeypatch's undo would then pin
|
|
# the real platform as a module attribute that hides every later
|
|
# _current_platform fake for the rest of the session.
|
|
monkeypatch.setattr(platforms, "_current_platform",
|
|
SimpleNamespace(is_mps=lambda: False,
|
|
is_cuda_alike=lambda: False,
|
|
is_cuda=lambda: False,
|
|
has_unified_memory=lambda device_id=0: False))
|
|
monkeypatch.setattr("fastvideo.worker.gpu_worker.maybe_init_distributed_environment_and_model_parallel",
|
|
init_distributed)
|
|
monkeypatch.setattr("fastvideo.worker.gpu_worker.build_pipeline", lambda args: SimpleNamespace())
|
|
args = FastVideoArgs(model_path="test",
|
|
num_gpus=2,
|
|
sp_size=2,
|
|
dist_timeout=7,
|
|
distributed_executor_backend="external_launcher")
|
|
worker = Worker(args, local_rank=1, rank=1, distributed_init_method="env://")
|
|
|
|
worker.init_device()
|
|
|
|
init_distributed.assert_called_once_with(1, 2, "env://", timeout=timedelta(seconds=7))
|
|
|
|
|
|
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_applies_offload_policy_after_binding_worker_device(monkeypatch, env_overrides,
|
|
executor_backend: str) -> None:
|
|
"""The runtime probe must see this worker's device, never driver device 0."""
|
|
events = []
|
|
args = FastVideoArgs(model_path="test", num_gpus=1, distributed_executor_backend=executor_backend)
|
|
args.finalize_device_offload_policy = Mock(side_effect=lambda device_id: events.append(("policy", device_id)))
|
|
worker = Worker(args, local_rank=3, rank=3, distributed_init_method="env://")
|
|
|
|
env_overrides.enter_context(envs.override_external("LOCAL_RANK", "0"))
|
|
# init_device() also writes RANK and WORLD_SIZE for these backends. A leaked
|
|
# RANK reaches every later child process (the profiler names its per-rank
|
|
# summary after it).
|
|
env_overrides.enter_context(envs.override_external("RANK", None))
|
|
env_overrides.enter_context(envs.override_external("WORLD_SIZE", None))
|
|
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, **kwargs: 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),
|
|
("policy", 3),
|
|
("distributed", None),
|
|
("pipeline", None),
|
|
]
|
|
assert os.environ["LOCAL_RANK"] == "3"
|
|
assert worker.device == torch.device("cuda:3")
|
|
assert worker.init_gpu_memory == 123
|
|
|
|
|
|
def _worker_returning(output_batch: ForwardBatch) -> Worker:
|
|
worker = Worker.__new__(Worker)
|
|
worker.fastvideo_args = SimpleNamespace(is_output_rank=True)
|
|
worker.pipeline = SimpleNamespace(forward=lambda batch, args: output_batch)
|
|
return worker
|
|
|
|
|
|
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, FastVideoArgs(model_path="test"))
|
|
|
|
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, FastVideoArgs(model_path="test"))
|
|
|
|
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, FastVideoArgs(model_path="test", 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, FastVideoArgs(model_path="test"))
|
|
|
|
assert result.output is not None
|
|
assert result.output.device.type == "cpu"
|
|
assert result.output.numel() == 0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("save_video", "return_frames", "return_samples"),
|
|
[
|
|
(True, False, False),
|
|
(False, True, False),
|
|
(True, True, False),
|
|
(False, False, True),
|
|
],
|
|
)
|
|
def test_execute_forward_preserves_requested_output(save_video, return_frames, return_samples):
|
|
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,
|
|
return_samples=return_samples)
|
|
|
|
result = worker.execute_forward(request_batch, FastVideoArgs(model_path="test"))
|
|
|
|
assert result.output is output
|