[Core] Add realtime AR-Diffusion tick sessions for LingBot World 2.0 (#5491)

Signed-off-by: zhangbin.327 <zhangbin.327@bytedance.com> Signed-off-by: Jack47 <scsvip@gmail.com> Co-authored-by: zhangbin.327 <zhangbin.327@bytedance.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-16 09:11:20 +08:00
1 parent 9798bc3a8d
commit 3875f8ded4
@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Sequence
+from typing import Any
+
+import pytest
+
+from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
+from vllm_omni.experimental.ar_diffusion.session import (
+    ARDiffusionSession,
+    ARDiffusionSessionEvent,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+    ARDiffusionChunkMetadata,
+    ARDiffusionControlInput,
+    ARDiffusionTickRequest,
+)
+from vllm_omni.inputs.data import (
+    OmniDiffusionSamplingParams,
+    OmniPromptType,
+    OmniSamplingParams,
+)
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeStagePool:
+    def __init__(self, num_replicas: int = 1) -> None:
+        self.num_replicas = num_replicas
+
+
+class FakeEngine:
+    def __init__(self, num_replicas: int = 1) -> None:
+        self.stage_pools = [FakeStagePool(num_replicas)]
+
+
+class FakeGenerateClient:
+    def __init__(
+        self,
+        *,
+        include_metadata: bool = True,
+        applied_event_ids: list[int] | None = None,
+        output_error: str | None = None,
+        num_replicas: int = 1,
+    ) -> None:
+        self.engine = FakeEngine(num_replicas)
+        self.include_metadata = include_metadata
+        self.applied_event_ids = applied_event_ids
+        self.output_error = output_error
+        self.calls: list[dict[str, Any]] = []
+
+    async def generate(
+        self,
+        prompt: OmniPromptType,
+        *,
+        request_id: str,
+        sampling_params_list: Sequence[OmniSamplingParams],
+        output_modalities: list[str] | None = None,
+    ) -> AsyncIterator[OmniRequestOutput]:
+        self.calls.append(
+            {
+                "prompt": prompt,
+                "request_id": request_id,
+                "sampling_params_list": sampling_params_list,
+                "output_modalities": output_modalities,
+            }
+        )
+        params = sampling_params_list[0]
+        assert isinstance(params, OmniDiffusionSamplingParams)
+        tick = ARDiffusionTickRequest.from_extra_args(params.extra_args)
+        assert tick is not None
+        chunk_metadata = ARDiffusionChunkMetadata.from_tick(tick).to_dict()
+        if self.applied_event_ids is not None:
+            chunk_metadata["applied_event_ids"] = self.applied_event_ids
+        metadata = {"metadata": {"ar_diffusion": chunk_metadata}} if self.include_metadata else {}
+        yield OmniRequestOutput(
+            request_id=request_id,
+            stage_id=0,
+            final_output_type="video",
+            images=["chunk"],  # type: ignore[list-item]
+            _multimodal_output=metadata,
+            finished=True,
+            error=self.output_error,
+        )
+
+
+class FakeLifecycle:
+    async def reset_session(self, session_id: str) -> None:
+        pass
+
+    async def close_session(self, session_id: str) -> None:
+        pass
+
+
+def make_tick() -> ARDiffusionTickRequest:
+    return ARDiffusionTickRequest(
+        session_id="world-1",
+        request_id="world-1-chunk-3",
+        chunk_index=3,
+        applied_event_ids=(7, 8),
+        prompt="turn left",
+        controls=(
+            ARDiffusionControlInput(
+                track="camera",
+                schema="example.camera.v1",
+                data={"samples": [{"x": 1.0}]},
+            ),
+        ),
+    )
+
+
+def test_consumer_rejects_replicated_ar_diffusion_stage() -> None:
+    client = FakeGenerateClient(num_replicas=2)
+
+    with pytest.raises(
+        ValueError,
+        match=r"require exactly one replica.*num_replicas=2",
+    ):
+        ARDiffusionOmniTickConsumer(
+            client,
+            prompt_provider=lambda tick: tick.prompt or "",
+            sampling_params_list=[OmniDiffusionSamplingParams()],
+            diffusion_stage_id=0,
+        )
+
+    assert client.calls == []
+
+
+@pytest.mark.asyncio
+async def test_consumer_submits_typed_tick_and_returns_standard_output() -> None:
+    client = FakeGenerateClient()
+    template = OmniDiffusionSamplingParams(
+        output_type="latent",
+        extra_args={"deployment_option": "kept"},
+    )
+    consumer = ARDiffusionOmniTickConsumer(
+        client,
+        prompt_provider=lambda tick: {
+            "prompt": tick.prompt,
+            "multi_modal_data": {"image": "initial-image"},
+        },
+        sampling_params_list=[template],
+        diffusion_stage_id=0,
+        output_modalities=["video"],
+    )
+    tick = make_tick()
+
+    output = await consumer.execute_tick(tick)
+
+    assert output.images == ["chunk"]
+    assert consumer.chunk_metadata(output) == ARDiffusionChunkMetadata.from_tick(tick)
+    assert len(client.calls) == 1
+    call = client.calls[0]
+    assert call["request_id"] == tick.request_id
+    assert call["output_modalities"] == ["video"]
+    assert call["prompt"] == {
+        "prompt": "turn left",
+        "multi_modal_data": {"image": "initial-image"},
+    }
+    submitted_params = call["sampling_params_list"][0]
+    assert submitted_params.extra_args["deployment_option"] == "kept"
+    assert ARDiffusionTickRequest.from_extra_args(submitted_params.extra_args) == tick
+    assert "ar_diffusion_tick" not in template.extra_args
+
+
+@pytest.mark.asyncio
+async def test_consumer_rejects_missing_chunk_metadata() -> None:
+    consumer = ARDiffusionOmniTickConsumer(
+        FakeGenerateClient(include_metadata=False),
+        prompt_provider=lambda tick: tick.prompt or "",
+        sampling_params_list=[OmniDiffusionSamplingParams()],
+        diffusion_stage_id=0,
+    )
+
+    with pytest.raises(ValueError, match="multimodal_output.metadata"):
+        await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_consumer_rejects_non_monotonic_applied_event_metadata() -> None:
+    consumer = ARDiffusionOmniTickConsumer(
+        FakeGenerateClient(applied_event_ids=[8, 7]),
+        prompt_provider=lambda tick: tick.prompt or "",
+        sampling_params_list=[OmniDiffusionSamplingParams()],
+        diffusion_stage_id=0,
+    )
+
+    with pytest.raises(ValueError, match="unique and strictly increasing"):
+        await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_consumer_preserves_engine_failure_before_metadata_validation() -> None:
+    consumer = ARDiffusionOmniTickConsumer(
+        FakeGenerateClient(include_metadata=False, output_error="model failed"),
+        prompt_provider=lambda tick: tick.prompt or "",
+        sampling_params_list=[OmniDiffusionSamplingParams()],
+        diffusion_stage_id=0,
+    )
+
+    with pytest.raises(RuntimeError, match="model failed"):
+        await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_session_commits_only_after_standard_engine_output_metadata() -> None:
+    client = FakeGenerateClient()
+    consumer = ARDiffusionOmniTickConsumer(
+        client,
+        prompt_provider=lambda tick: tick.prompt or "",
+        sampling_params_list=[OmniDiffusionSamplingParams()],
+        diffusion_stage_id=0,
+    )
+    session = ARDiffusionSession(
+        "world-1",
+        tick_consumer=consumer,
+        lifecycle=FakeLifecycle(),
+        max_pending_events=4,
+        request_id_factory=lambda session_id, chunk_index: f"{session_id}-chunk-{chunk_index}",
+    )
+    await session.accept_event(ARDiffusionSessionEvent(event_id=7, prompt="turn left"))
+
+    output = await session.next_chunk()
+
+    assert output.request_id == "world-1-chunk-0"
+    assert session.chunk_index == 1
+    assert session.pending_event_count == 0
+    assert client.calls[0]["request_id"] == "world-1-chunk-0"