72 lines
2.9 KiB
Python
72 lines
2.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from datetime import timedelta
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from vllm_omni.diffusion.distributed.autoencoders import distributed_vae_executor as vae_executor
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
class TileOperator:
|
|
def split(self, tensor):
|
|
return [vae_executor.TileTask(i, (i,), row[None]) for i, row in enumerate(tensor)], vae_executor.GridSpec(
|
|
split_dims=(0,), grid_shape=(tensor.shape[0],)
|
|
)
|
|
|
|
def exec(self, task):
|
|
return task.tensor * 2
|
|
|
|
def merge(self, tiles, grid):
|
|
return torch.cat([tiles[(i,)] for i in range(grid.grid_shape[0])])
|
|
|
|
|
|
def run_recovery_worker(rank, init_file):
|
|
torch.set_num_threads(1)
|
|
dist.init_process_group(
|
|
"gloo", init_method=f"file://{init_file}", rank=rank, world_size=2, timeout=timedelta(seconds=30)
|
|
)
|
|
try:
|
|
group = SimpleNamespace(device_group=dist.group.WORLD, cpu_group=dist.group.WORLD)
|
|
with pytest.MonkeyPatch.context() as patch:
|
|
patch.setattr(vae_executor, "get_world_group", lambda: group)
|
|
executor = vae_executor.DistributedVaeExecutor()
|
|
executor.set_parallel_size(2)
|
|
tensor = torch.arange(12, dtype=torch.float32).reshape(4, 3)
|
|
operator = TileOperator()
|
|
for stage, failing_rank in [("decode", 1), ("packing", 1), ("merge", 0), ("gather", 1)]:
|
|
with pytest.MonkeyPatch.context() as patch:
|
|
if rank == failing_rank:
|
|
|
|
def fail(*args, **kwargs):
|
|
raise torch.OutOfMemoryError("injected OOM")
|
|
|
|
target, attribute = {
|
|
"decode": (operator, "exec"),
|
|
"packing": (executor, "_pack_local_tiles"),
|
|
"merge": (operator, "merge"),
|
|
"gather": (torch, "empty_like"),
|
|
}[stage]
|
|
patch.setattr(target, attribute, fail)
|
|
with pytest.raises(torch.OutOfMemoryError, match="Distributed VAE"):
|
|
executor.execute(tensor, operator)
|
|
torch.testing.assert_close(executor.execute(tensor, operator), tensor * 2)
|
|
|
|
with pytest.raises(RuntimeError, match="Distributed VAE mixed failures failed"):
|
|
with executor._sync_errors("mixed failures"):
|
|
if rank == 0:
|
|
raise torch.OutOfMemoryError("retryable")
|
|
raise ValueError("not retryable")
|
|
torch.testing.assert_close(executor.execute(tensor, operator), tensor * 2)
|
|
finally:
|
|
dist.destroy_process_group()
|
|
|
|
|
|
def test_all_ranks_observe_failures_and_can_retry(tmp_path):
|
|
torch.multiprocessing.spawn(run_recovery_worker, args=(str(tmp_path / "rendezvous"),), nprocs=2, join=True)
|