108 lines
3.4 KiB
Python
108 lines
3.4 KiB
Python
"""Multi-replica eval through the public ``Evaluator`` API.
|
|
|
|
Skipped automatically when fewer than 2 CUDA devices are visible.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from fastvideo.eval import create_evaluator
|
|
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
not torch.cuda.is_available() or torch.cuda.device_count() < 2,
|
|
reason="multi-GPU evaluator tests require at least 2 visible CUDA devices",
|
|
)
|
|
|
|
|
|
_T, _C, _H, _W = 6, 3, 32, 32
|
|
|
|
|
|
def _make_samples(n: int) -> list[dict]:
|
|
torch.manual_seed(7)
|
|
samples: list[dict] = []
|
|
for i in range(n):
|
|
gen = torch.rand(_T, _C, _H, _W)
|
|
# Slightly perturb each row's reference so scores vary by index.
|
|
ref = gen + 0.01 * (i + 1) * torch.rand_like(gen)
|
|
samples.append({"video": gen, "reference": ref})
|
|
return samples
|
|
|
|
|
|
@pytest.fixture
|
|
def baseline_scores():
|
|
"""Reference scores computed on a single-GPU evaluator. The multi-GPU
|
|
runs must reproduce these exactly when handed the same input list —
|
|
that's the only way to verify round-robin dispatch isn't dropping or
|
|
reordering samples."""
|
|
samples = _make_samples(8)
|
|
ev = create_evaluator(
|
|
metrics=["common.psnr", "common.ssim"],
|
|
device="cuda:0",
|
|
num_gpus=1,
|
|
)
|
|
try:
|
|
out = ev.evaluate(samples=samples)
|
|
finally:
|
|
ev.shutdown()
|
|
return samples, out
|
|
|
|
|
|
def test_multi_gpu_evaluator_reports_two_workers():
|
|
ev = create_evaluator(metrics=["common.psnr"], num_gpus=2)
|
|
try:
|
|
assert ev.num_gpus == 2
|
|
finally:
|
|
ev.shutdown()
|
|
|
|
|
|
def test_multi_gpu_dispatch_preserves_order_and_scores(baseline_scores):
|
|
"""Same samples, multi-GPU dispatch — results must match the single-GPU
|
|
baseline element-for-element. This verifies (a) the round-robin doesn't
|
|
reorder, (b) every sample is scored exactly once, (c) the workers
|
|
don't share mutable state."""
|
|
samples, expected = baseline_scores
|
|
|
|
ev = create_evaluator(
|
|
metrics=["common.psnr", "common.ssim"],
|
|
num_gpus=2,
|
|
)
|
|
try:
|
|
got = ev.evaluate(samples=samples)
|
|
finally:
|
|
ev.shutdown()
|
|
|
|
assert len(got) == len(expected)
|
|
for i, (g, e) in enumerate(zip(got, expected)):
|
|
assert set(g.keys()) == {"common.psnr", "common.ssim"}, f"row {i}"
|
|
assert g["common.psnr"].score == pytest.approx(e["common.psnr"].score), \
|
|
f"row {i} psnr drift"
|
|
assert g["common.ssim"].score == pytest.approx(e["common.ssim"].score), \
|
|
f"row {i} ssim drift"
|
|
|
|
|
|
def test_multi_gpu_evaluator_kwargs_form_runs_on_one_replica():
|
|
"""The kwargs form (single sample) is documented to always hit worker
|
|
0; this test pins the contract so future refactors don't accidentally
|
|
fan out a single call."""
|
|
ev = create_evaluator(metrics=["common.psnr"], num_gpus=2)
|
|
try:
|
|
torch.manual_seed(0)
|
|
gen = torch.rand(_T, _C, _H, _W)
|
|
out = ev.evaluate(video=gen, reference=gen)
|
|
assert out["common.psnr"].score > 50.0 # PSNR(x, x) is huge
|
|
finally:
|
|
ev.shutdown()
|
|
|
|
|
|
def test_multi_gpu_release_cuda_memory_runs_clean():
|
|
"""``release_cuda_memory`` must hit every replica without crashing."""
|
|
ev = create_evaluator(metrics=["common.psnr"], num_gpus=2)
|
|
try:
|
|
samples = _make_samples(2)
|
|
_ = ev.evaluate(samples=samples)
|
|
ev.release_cuda_memory() # should not raise
|
|
finally:
|
|
ev.shutdown()
|