[Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951)

Signed-off-by: wangyu <410167048@qq.com>

tests/diffusion/models/dreamzero/test_session_state_accounting.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-device byte accounting and the bounded VAE encoder ring (RFC #4480).
+
+Two things are pinned here. First, that a session reports the bytes it actually
+holds: both buckets walked, storages rather than logical element counts, and
+aliases counted once. Second, that bounding the adapter's VAE encoder history
+changes no numbers -- ``quant_conv`` is pointwise in time, so the frames the
+reader keeps do not depend on the ones the ring drops.
+
+All CPU, tiny tensors, no model.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.experimental.world_models.adapters.state_dreamzero_adapter import (
+    DreamZeroStateAdapter,
+)
+from vllm_omni.experimental.world_models.session_state import (
+    LatentBuffer,
+    SessionStateManager,
+    device_bytes,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _bytes(value: object) -> dict[str, int]:
+    return device_bytes(value, set())
+
+
+# -- device_bytes ------------------------------------------------------------
+
+
+def test_counts_tensor_storage_under_its_device() -> None:
+    tensor = torch.zeros(64, dtype=torch.float32)
+    assert _bytes(tensor) == {"cpu": 256}
+
+
+def test_views_on_one_storage_counted_once() -> None:
+    tensor = torch.zeros(64, dtype=torch.float32)
+    seen: set[tuple[str, int]] = set()
+    first = device_bytes(tensor, seen)
+    second = device_bytes(tensor[:8], seen)
+    assert first == {"cpu": 256}
+    assert second == {}
+
+
+def test_slice_reports_the_storage_it_pins() -> None:
+    # The whole point of counting storages: this slice keeps all 256 bytes
+    # alive, and ``numel() * element_size()`` would claim 32.
+    sliced = torch.zeros(64, dtype=torch.float32)[:8]
+    assert _bytes(sliced) == {"cpu": 256}
+
+
+def test_containers_recurse_and_scalars_contribute_nothing() -> None:
+    payload = {
+        "tensors": [torch.zeros(16, dtype=torch.float32), None],
+        "count": 7,
+        "flag": True,
+        "name": "not-a-buffer",
+    }
+    assert _bytes(payload) == {"cpu": 64}
+
+
+def test_numpy_arrays_land_on_the_host() -> None:
+    assert _bytes(np.zeros(32, dtype=np.float32)) == {"cpu": 128}
+
+
+# -- SessionState / manager --------------------------------------------------
+
+
+def test_session_counts_both_buckets_and_dedupes_across_them() -> None:
+    manager = SessionStateManager()
+    session = manager.get_or_create_session("s")
+
+    buffer: LatentBuffer[torch.Tensor] = LatentBuffer()
+    buffer.allocate(maxlen=None)
+    shared = torch.zeros(64, dtype=torch.float32)  # 256 bytes
+    buffer.append(shared)
+    session.put("latents", buffer)
+    session.attrs["conditioning"] = torch.zeros(16, dtype=torch.float32)  # 64 bytes
+
+    assert session.nbytes_by_device() == {"cpu": 320}
+
+    # The same storage reachable from both buckets is one allocation, not two.
+    session.attrs["alias"] = shared
+    assert session.nbytes_by_device() == {"cpu": 320}
+
+
+def test_stats_reports_per_device_totals() -> None:
+    manager = SessionStateManager()
+    manager.get_or_create_session("s").attrs["x"] = torch.zeros(64, dtype=torch.float32)
+
+    stats = manager.stats()
+    assert stats["nbytes:cpu"] == 256
+    assert stats["total_nbytes"] == 256
+
+
+def test_attribute_bytes_were_previously_invisible() -> None:
+    """The defect behind tzhouam's review comment on PR #4487.
+
+    Every byte here lives in ``attrs``, and the old ``nbytes`` walked only the
+    typed objects -- so a session holding real memory reported zero.
+    """
+    manager = SessionStateManager()
+    session = manager.get_or_create_session("s")
+    session.attrs["conv_cache"] = [torch.zeros(1024, dtype=torch.float32)]
+
+    assert session.nbytes_by_device() == {"cpu": 4096}
+
+
+# -- the largest value DreamZero holds ---------------------------------------
+
+
+def test_feat_map_bytes_are_counted() -> None:
+    """``vae_enc_feat_map`` is the largest thing a DreamZero session holds --
+    603 MiB at the shipped config -- and it stays an attribute, so this is
+    exactly the memory the old accounting missed."""
+    manager = SessionStateManager()
+    adapter = DreamZeroStateAdapter("s", manager)
+    adapter.vae_enc_feat_map = [torch.zeros(1024, dtype=torch.float32), None]
+
+    assert manager.stats()["nbytes:cpu"] == 4096
+
+    # Freeing it piecemeal is not something the manager can do: the cache has no
+    # recompute source and the stream breaks without it. Ending the stream (or
+    # the session) is the only release, and that reclaims attributes too.
+    adapter.reset_vae_encoder_stream()
+    assert manager.stats().get("nbytes:cpu", 0) == 0
+
+
+# -- bounded VAE encoder ring ------------------------------------------------
+
+
+def _quantize(encoder_out: torch.Tensor) -> torch.Tensor:
+    """Stand-in for ``_vae_quantize_encoder_out``.
+
+    The real one is ``WanCausalConv3d(z_dim*2, z_dim*2, 1)`` followed by a
+    per-channel affine rescale. Every operation is pointwise in time, which is
+    the property the ring relies on; a 1x1x1 convolution reproduces it.
+    """
+    weight = torch.arange(1, encoder_out.shape[1] + 1, dtype=encoder_out.dtype).view(1, -1, 1, 1, 1)
+    return encoder_out * weight + 0.5
+
+
+def test_ring_is_bounded() -> None:
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=2)
+    adapter.vae_encoder_out = torch.randn(1, 4, 1, 2, 2)
+    for _ in range(10):
+        adapter.append_vae_encoder_chunk(torch.randn(1, 4, 1, 2, 2))
+
+    assert adapter.vae_encoder_out is not None
+    assert adapter.vae_encoder_out.shape[2] == 2
+
+
+def test_ring_matches_unbounded_history_for_the_frames_that_are_read() -> None:
+    window = 2
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=window)
+    seed = torch.randn(1, 4, 1, 2, 2)
+    chunks = [torch.randn(1, 4, 1, 2, 2) for _ in range(9)]
+
+    adapter.vae_encoder_out = seed
+    unbounded = seed
+    for chunk in chunks:
+        adapter.append_vae_encoder_chunk(chunk)
+        unbounded = torch.cat([unbounded, chunk], dim=2)
+
+    # What the reader actually returns: quantise, then keep the last `window`.
+    from_ring = _quantize(adapter.vae_encoder_out)[:, :, -window:]
+    from_history = _quantize(unbounded)[:, :, -window:]
+    torch.testing.assert_close(from_ring, from_history, rtol=0, atol=0)
+
+
+def test_ring_holds_a_single_frame_before_the_first_body_chunk() -> None:
+    # After `_vae_stream_seed` there is exactly one latent frame, which is what
+    # sends the reader down its padding branch.
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=2)
+    seed = torch.randn(1, 4, 1, 2, 2)
+    adapter.vae_encoder_out = seed
+
+    assert adapter.vae_encoder_out is not None
+    assert adapter.vae_encoder_out.shape[2] == 1
+    torch.testing.assert_close(adapter.vae_encoder_out, seed, rtol=0, atol=0)
+
+
+def test_resetting_the_stream_clears_the_ring() -> None:
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=2)
+    adapter.vae_encoder_out = torch.randn(1, 4, 1, 2, 2)
+    adapter.reset_vae_encoder_stream()
+
+    assert adapter.vae_encoder_out is None
+
+
+def test_window_must_be_positive() -> None:
+    with pytest.raises(ValueError, match="vae_encoder_window"):
+        DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=0)