[Diffusion] Add node-local runtime weight cache for DLO

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/model_loader/test_runtime_weight_cache.py

@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import fcntl
+import multiprocessing as mp
+import shutil
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
+    build_runtime_weight_cache_plan,
+    default_runtime_cache_root,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _LoaderMarker:
+    pass
+
+
+class _Transformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(3, 2)
+        self.register_buffer("persistent", torch.arange(2, dtype=torch.float32))
+        self.register_buffer("temporary", torch.ones(1), persistent=False)
+        with torch.no_grad():
+            self.proj.weight.copy_(torch.arange(6, dtype=torch.float32).reshape(2, 3))
+            self.proj.bias.copy_(torch.tensor([7.0, 8.0]))
+
+
+class _Pipeline(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.transformer = _Transformer()
+
+
+def test_default_root_is_not_split_by_replica_vllm_cache(monkeypatch):
+    monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/replica-specific-vllm-cache")
+
+    assert default_runtime_cache_root().endswith("/.cache/vllm-omni/dlo-runtime-weights")
+
+
+def _build(pipeline: nn.Module, cache_root: Path, **overrides):
+    kwargs = {
+        "dit_modules": (("transformer", pipeline.transformer),),
+        "loader_type": _LoaderMarker,
+        "cache_root": cache_root,
+        "lock_timeout_seconds": 2.0,
+        "max_shard_bytes": 32,
+        "model_identity": "test/model",
+        "revision": "revision",
+        "runtime_dtype": torch.float32,
+        "load_format": "default",
+        "loader_inputs": {"model_config": {"hidden_size": 3}},
+        "tensor_parallel_size": 1,
+        "tensor_parallel_rank": 0,
+        "sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"},
+        "use_hsdp": False,
+        "quantization_config": None,
+        "cfg_parallel_size": 1,
+        "pipeline_parallel_size": 1,
+    }
+    kwargs.update(overrides)
+    return build_runtime_weight_cache_plan(pipeline, **kwargs)
+
+
+def _multiprocess_builder(cache_root: str, queue) -> None:
+    result = _build(_Pipeline(), Path(cache_root))
+    queue.put(
+        (
+            result.fallback_code,
+            result.plan.runtime_layout_key if result.plan is not None else None,
+            sorted(binding.file_path for binding in result.plan.bindings.values()) if result.plan is not None else [],
+        )
+    )
+
+
+def test_equivalent_process_layouts_reuse_one_entry(tmp_path):
+    first = _build(_Pipeline(), tmp_path)
+    second = _build(_Pipeline(), tmp_path)
+
+    assert first.plan is not None
+    assert second.plan is not None
+    assert first.plan.runtime_layout_key == second.plan.runtime_layout_key
+    assert first.plan.post_load_complete
+    assert {binding.file_path for binding in first.plan.bindings.values()} == {
+        binding.file_path for binding in second.plan.bindings.values()
+    }
+    assert set(first.plan.bindings) == {
+        "transformer.persistent",
+        "transformer.proj.bias",
+        "transformer.proj.weight",
+    }
+
+
+def test_tp_coordinate_and_sp_implementation_guard_split_entries(tmp_path):
+    pipeline = _Pipeline()
+    tp0 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=0)
+    tp1 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=1)
+    sp2 = _build(
+        pipeline,
+        tmp_path,
+        sequence_parallel_guard={"sequence_parallel_size": 2, "backend": "ulysses"},
+    )
+
+    assert tp0.plan is not None and tp1.plan is not None and sp2.plan is not None
+    assert len({tp0.plan.runtime_layout_key, tp1.plan.runtime_layout_key, sp2.plan.runtime_layout_key}) == 3
+
+
+def test_final_runtime_content_is_authoritative(tmp_path):
+    first_pipeline = _Pipeline()
+    second_pipeline = _Pipeline()
+    with torch.no_grad():
+        second_pipeline.transformer.proj.weight[0, 0] += 1
+
+    first = _build(first_pipeline, tmp_path)
+    second = _build(second_pipeline, tmp_path)
+
+    assert first.plan is not None and second.plan is not None
+    assert first.plan.runtime_layout_key != second.plan.runtime_layout_key
+
+
+def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path):
+    first = _build(_Pipeline(), tmp_path)
+    assert first.plan is not None
+    shard = Path(next(iter(first.plan.bindings.values())).file_path)
+    old_inode = shard.stat().st_ino
+    with shard.open("r+b") as handle:
+        handle.seek(-1, 2)
+        byte = handle.read(1)
+        handle.seek(-1, 2)
+        handle.write(bytes([byte[0] ^ 0xFF]))
+
+    second = _build(_Pipeline(), tmp_path)
+
+    assert second.plan is not None
+    assert second.plan.runtime_layout_key == first.plan.runtime_layout_key
+    assert Path(next(iter(second.plan.bindings.values())).file_path).stat().st_ino != old_inode
+
+
+def test_lock_timeout_keeps_the_ordinary_tensors(tmp_path):
+    initial = _build(_Pipeline(), tmp_path)
+    assert initial.plan is not None
+    cache_key = initial.plan.runtime_layout_key
+    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
+    shutil.rmtree(entry_dir)
+    lock_path = tmp_path / ".locks" / f"{cache_key}.lock"
+
+    with lock_path.open("a+b") as lock_handle:
+        fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        result = _build(_Pipeline(), tmp_path, lock_timeout_seconds=0.01)
+
+    assert result.plan is None
+    assert result.fallback_code == "lock_timeout"
+
+
+def test_next_writer_removes_stale_temp_after_process_death(tmp_path):
+    initial = _build(_Pipeline(), tmp_path)
+    assert initial.plan is not None
+    cache_key = initial.plan.runtime_layout_key
+    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
+    shutil.rmtree(entry_dir)
+    stale_temp = tmp_path / ".tmp" / f"{cache_key}.stale-writer"
+    stale_temp.mkdir(parents=True)
+    (stale_temp / "partial").write_bytes(b"partial")
+
+    result = _build(_Pipeline(), tmp_path)
+
+    assert result.plan is not None
+    assert not stale_temp.exists()
+
+
+@pytest.mark.parametrize("case", ["alias", "external_alias", "noncontiguous"])
+def test_unsupported_tensor_layouts_fail_closed(tmp_path, case):
+    pipeline = _Pipeline()
+    if case == "alias":
+        pipeline.transformer.tied = pipeline.transformer.proj.weight
+        expected_code = "unsupported_alias"
+    elif case == "external_alias":
+        pipeline.encoder = nn.Module()
+        pipeline.encoder.shared = pipeline.transformer.proj.weight
+        expected_code = "unsupported_alias"
+    else:
+        pipeline.transformer.proj.weight = nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(3, 2).t())
+        expected_code = "unsupported_tensor"
+
+    result = _build(pipeline, tmp_path)
+
+    assert result.plan is None
+    assert result.fallback_code == expected_code
+    assert list(tmp_path.rglob("*.safetensors")) == []
+
+
+def test_concurrent_processes_publish_one_immutable_entry(tmp_path):
+    context = mp.get_context("spawn")
+    queue = context.Queue()
+    processes = [context.Process(target=_multiprocess_builder, args=(str(tmp_path), queue)) for _ in range(2)]
+    for process in processes:
+        process.start()
+    results = [queue.get(timeout=30) for _ in processes]
+    for process in processes:
+        process.join(timeout=30)
+
+    assert all(process.exitcode == 0 for process in processes)
+    assert all(code is None for code, _, _ in results)
+    assert len({key for _, key, _ in results}) == 1
+    assert len({tuple(paths) for _, _, paths in results}) == 1
+    assert len(list((tmp_path / "v1").iterdir())) == 1