897 lines
39 KiB
Python
897 lines
39 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Tests for midway prompt update across runner, batch, pipeline, and entrypoint layers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import queue
|
|
import time
|
|
from contextlib import contextmanager
|
|
from types import SimpleNamespace
|
|
from typing import Any, cast
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tests.engine.test_orchestrator import (
|
|
OrchestratorFixture,
|
|
_build_harness,
|
|
_wait_for,
|
|
)
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
|
|
from vllm_omni.diffusion.interaction.coordinator import InteractionCoordinator
|
|
from vllm_omni.diffusion.interaction.modality_handlers.prompt import (
|
|
PromptInteractionHandler,
|
|
PromptSession,
|
|
)
|
|
from vllm_omni.diffusion.models.helios.pipeline_helios import HeliosPipeline
|
|
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
|
from vllm_omni.diffusion.worker.input_batch import InputBatch
|
|
from vllm_omni.diffusion.worker.utils import StepRequestState
|
|
from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
|
|
from vllm_omni.engine.messages import (
|
|
ErrorMessage,
|
|
InteractionMessage,
|
|
ShutdownRequestMessage,
|
|
StageSubmissionMessage,
|
|
)
|
|
from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
|
|
from vllm_omni.engine.stage_client import StagePoolClient
|
|
from vllm_omni.engine.stage_init_utils import StageMetadata
|
|
from vllm_omni.engine.stage_pool import StagePool
|
|
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
|
from vllm_omni.entrypoints.async_omni_base import AsyncEventResolver
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.fixture
|
|
def pipeline() -> HeliosPipeline:
|
|
pipeline = object.__new__(HeliosPipeline)
|
|
pipeline.device = torch.device("cpu")
|
|
pipeline.transformer = SimpleNamespace(dtype=torch.float32) # pyright: ignore[reportAttributeAccessIssue]
|
|
pipeline.encode_prompt = MagicMock(
|
|
return_value=(
|
|
torch.full((1, 4, 2), 2.0),
|
|
None,
|
|
)
|
|
)
|
|
pipeline.prepare_next_chunk = MagicMock()
|
|
od_config = SimpleNamespace(model_class_name="HeliosPipeline")
|
|
pipeline._interaction_coordinator = InteractionCoordinator.build(pipeline, od_config)
|
|
return pipeline
|
|
|
|
|
|
def _make_diffusion_request_state(*, request_id: str = "req-1", fps: int | None = None) -> StepRequestState:
|
|
state = StepRequestState(
|
|
request_id=request_id,
|
|
sampling=SimpleNamespace( # pyright: ignore[reportArgumentType]
|
|
num_outputs_per_prompt=1,
|
|
max_sequence_length=226,
|
|
fps=fps,
|
|
),
|
|
prompt="hello",
|
|
)
|
|
state.prompt_embeds = torch.zeros(1, 4, 2)
|
|
state.extra = {"window_num_frames": 8, "num_latent_frames_per_chunk": 3}
|
|
return state
|
|
|
|
|
|
def _make_diffusion_model_runner(
|
|
*,
|
|
pipeline,
|
|
streaming_output: bool = True,
|
|
model_class_name: str = "HeliosPipeline",
|
|
) -> DiffusionModelRunner:
|
|
runner = object.__new__(DiffusionModelRunner)
|
|
runner.pipeline = pipeline
|
|
runner.state_cache = {}
|
|
runner.od_config = SimpleNamespace( # pyright: ignore[reportAttributeAccessIssue]
|
|
model_class_name=model_class_name,
|
|
streaming_output=streaming_output,
|
|
cache_backend=None,
|
|
parallel_config=SimpleNamespace(use_hsdp=False),
|
|
)
|
|
runner.vllm_config = object()
|
|
runner.device = torch.device("cpu")
|
|
runner.cache_backend = None
|
|
runner.offload_backend = None
|
|
runner.prompt_embed_cache = None
|
|
runner.kv_transfer_manager = SimpleNamespace(
|
|
receive_multi_kv_cache_distributed=lambda *a, **k: None,
|
|
)
|
|
# HeliosPipeline skeletons from object.__new__ are not full SupportsStepExecution
|
|
# instances; unit tests of submit_interaction stub this. Stepwise path tests
|
|
# rebind to the real method.
|
|
runner._supports_step_mode = lambda: True
|
|
# Mirror load_model: coordinator is always present once the runner is "loaded".
|
|
# Build now so handlers capture the pipeline's current encode_prompt (tests may
|
|
# rebind the mock before calling this helper).
|
|
runner._interaction_coordinator = InteractionCoordinator.build(pipeline, runner.od_config)
|
|
if hasattr(pipeline, "_interaction_coordinator"):
|
|
pipeline._interaction_coordinator = runner._interaction_coordinator
|
|
return runner
|
|
|
|
|
|
def _prompt_interaction(
|
|
prompt: str = "updated",
|
|
*,
|
|
event_id: str = "ui-update-1",
|
|
transition_chunks: int | None = None,
|
|
) -> OmniInteractionPrompt:
|
|
interaction: dict[str, Any] = {
|
|
"event_id": event_id,
|
|
"event": {"prompt": prompt},
|
|
}
|
|
if transition_chunks is not None:
|
|
interaction["transition_chunks"] = transition_chunks
|
|
return cast(OmniInteractionPrompt, interaction)
|
|
|
|
|
|
class TestPromptUpdateExecution:
|
|
"""Runner, InputBatch cache, and pipeline chunk-boundary behavior."""
|
|
|
|
def test_runner_interaction_delegates_to_helios_pipeline(self, pipeline: HeliosPipeline) -> None:
|
|
"""Runner encodes the new prompt and queues pending embeds on the request state."""
|
|
runner = _make_diffusion_model_runner(pipeline=pipeline)
|
|
state = _make_diffusion_request_state()
|
|
runner.state_cache["req-1"] = state
|
|
|
|
runner.submit_interaction("req-1", _prompt_interaction(transition_chunks=2))
|
|
|
|
pipeline.encode_prompt.assert_called_once() # pyright: ignore[reportAttributeAccessIssue]
|
|
session = state.interaction_sessions["prompt"]
|
|
assert isinstance(session, PromptSession)
|
|
pending = session.pending_event
|
|
assert pending is not None
|
|
assert pending.event_id == "ui-update-1"
|
|
assert pending.transition_chunks == 2
|
|
assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0))
|
|
|
|
def test_runner_prompt_update_rejects_unsupported_pipeline(self) -> None:
|
|
"""Runner rejects prompt updates when the pipeline lacks prompt-update support."""
|
|
|
|
class _UnsupportedPipeline:
|
|
supports_step_execution = True
|
|
|
|
runner = _make_diffusion_model_runner(pipeline=_UnsupportedPipeline())
|
|
runner.state_cache["req-1"] = _make_diffusion_request_state()
|
|
with pytest.raises(ValueError, match="not supported"):
|
|
runner.submit_interaction("req-1", _prompt_interaction())
|
|
|
|
def test_runner_prompt_update_rejects_missing_request(self, pipeline) -> None:
|
|
"""Runner rejects prompt updates when no active request state exists."""
|
|
runner = _make_diffusion_model_runner(pipeline=pipeline)
|
|
with pytest.raises(ValueError, match="No active request state"):
|
|
runner.submit_interaction("missing", _prompt_interaction())
|
|
|
|
@pytest.mark.parametrize(
|
|
"interaction",
|
|
[
|
|
{},
|
|
{"event_id": "bad", "event": {"multi_modal_data": {}}},
|
|
{"event_id": "bad", "event": {"multi_modal_data": None}},
|
|
{"event_id": "bad", "event": {"prompt": "updated", "multi_modal_data": "bad"}},
|
|
{
|
|
"event_id": "cam-only",
|
|
"event": {
|
|
"multi_modal_data": {
|
|
"camera": {
|
|
"mode": "velocity",
|
|
"data": {"translation": [0.0, 0.0, 0.05], "rotation": [0.0, 0.0, 0.0, 1.0]},
|
|
},
|
|
}
|
|
},
|
|
},
|
|
{
|
|
"event_id": "cam-and-prompt",
|
|
"event": {
|
|
"prompt": "updated",
|
|
"multi_modal_data": {
|
|
"camera": {
|
|
"mode": "velocity",
|
|
"data": {"translation": [0.0, 0.0, 0.05], "rotation": [0.0, 0.0, 0.0, 1.0]},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
],
|
|
)
|
|
def test_runner_rejects_unsupported_interactions(
|
|
self,
|
|
pipeline: HeliosPipeline,
|
|
interaction: dict[str, Any],
|
|
) -> None:
|
|
"""Helios has no camera handler; malformed / unsupported payloads are rejected."""
|
|
runner = _make_diffusion_model_runner(pipeline=pipeline)
|
|
state = _make_diffusion_request_state()
|
|
runner.state_cache["req-1"] = state
|
|
|
|
with pytest.raises((ValueError, KeyError, AttributeError, TypeError)):
|
|
runner.submit_interaction("req-1", cast(Any, interaction))
|
|
|
|
# Rejected composite/malformed events must leave interaction queues unchanged.
|
|
pipeline.encode_prompt.assert_not_called() # pyright: ignore[reportAttributeAccessIssue]
|
|
assert "prompt" not in state.interaction_sessions
|
|
assert "camera" not in state.interaction_sessions
|
|
|
|
def test_rejected_composite_leaves_existing_prompt_queue_unchanged(
|
|
self,
|
|
pipeline: HeliosPipeline,
|
|
) -> None:
|
|
"""Helios must not apply a composite prompt when camera modality is unsupported."""
|
|
runner = _make_diffusion_model_runner(pipeline=pipeline)
|
|
state = _make_diffusion_request_state()
|
|
runner.state_cache["req-1"] = state
|
|
runner.submit_interaction("req-1", _prompt_interaction("already-queued", event_id="prior"))
|
|
session = state.interaction_sessions["prompt"]
|
|
assert isinstance(session, PromptSession)
|
|
prior = session.pending_event
|
|
assert prior is not None
|
|
assert prior.event_id == "prior"
|
|
pipeline.encode_prompt.reset_mock() # pyright: ignore[reportAttributeAccessIssue]
|
|
|
|
with pytest.raises(ValueError, match="camera"):
|
|
runner.submit_interaction(
|
|
"req-1",
|
|
cast(
|
|
Any,
|
|
{
|
|
"event_id": "cam-and-prompt",
|
|
"event": {
|
|
"prompt": "should-not-replace",
|
|
"multi_modal_data": {
|
|
"camera": {
|
|
"mode": "velocity",
|
|
"data": {
|
|
"translation": [0.0, 0.0, 0.05],
|
|
"rotation": [0.0, 0.0, 0.0, 1.0],
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
),
|
|
)
|
|
|
|
pipeline.encode_prompt.assert_not_called() # pyright: ignore[reportAttributeAccessIssue]
|
|
pending = session.pending_event
|
|
assert pending is not None
|
|
assert pending.event_id == "prior"
|
|
assert pending.prompt == "already-queued"
|
|
assert "camera" not in state.interaction_sessions
|
|
|
|
@pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HeliosPyramidPipeline"])
|
|
def test_coordinator_registers_prompt_for_helios_aliases(
|
|
self,
|
|
pipeline: HeliosPipeline,
|
|
model_class_name: str,
|
|
) -> None:
|
|
"""Both Helios architecture names resolve the prompt handler registry."""
|
|
od_config = SimpleNamespace(model_class_name=model_class_name)
|
|
coordinator = InteractionCoordinator.build(pipeline, od_config)
|
|
assert coordinator.has_modality("prompt")
|
|
|
|
runner = _make_diffusion_model_runner(pipeline=pipeline, model_class_name=model_class_name)
|
|
runner.state_cache["req-1"] = _make_diffusion_request_state()
|
|
runner.submit_interaction("req-1", _prompt_interaction())
|
|
session = runner.state_cache["req-1"].interaction_sessions["prompt"]
|
|
assert isinstance(session, PromptSession)
|
|
assert session.pending_event is not None
|
|
|
|
def test_helios_peek_chunk_media_requires_positive_fps(self, pipeline: HeliosPipeline) -> None:
|
|
"""Media-timeline peek must raise on missing/invalid fps instead of coercing to 0.0."""
|
|
state = _make_diffusion_request_state(fps=None)
|
|
with pytest.raises(ValueError, match="sampling.fps is required and must be > 0"):
|
|
pipeline.peek_chunk_media(state)
|
|
|
|
state.sampling.fps = 0 # pyright: ignore[reportAttributeAccessIssue]
|
|
with pytest.raises(ValueError, match="sampling.fps is required and must be > 0"):
|
|
pipeline.peek_chunk_media(state)
|
|
|
|
def test_input_batch_refreshes_prompt_embeds_on_version_change(self) -> None:
|
|
"""InputBatch rebuilds prompt_embeds when prompt_update_version changes."""
|
|
state = StepRequestState(
|
|
request_id="req-1",
|
|
sampling=SimpleNamespace(), # pyright: ignore[reportArgumentType]
|
|
prompt="hello",
|
|
)
|
|
state.prompt_embeds = torch.zeros(1, 2, 3)
|
|
state.latents = torch.zeros(1, 2)
|
|
state.timesteps = torch.tensor([1.0])
|
|
state.interaction_sessions["prompt"] = PromptSession(version=0)
|
|
|
|
batch = InputBatch.make_batch([state])
|
|
assert torch.equal(batch.prompt_embeds, torch.zeros(1, 2, 3)) # pyright: ignore[reportArgumentType]
|
|
|
|
state.prompt_embeds = torch.ones(1, 2, 3)
|
|
state.interaction_sessions["prompt"] = PromptSession(version=1)
|
|
refreshed = InputBatch.make_batch([state], cached_batch=batch)
|
|
assert torch.equal(refreshed.prompt_embeds, torch.ones(1, 2, 3)) # pyright: ignore[reportArgumentType]
|
|
|
|
def test_helios_prompt_handler_enqueue_queues_pending_target(self, pipeline: HeliosPipeline) -> None:
|
|
"""Prompt handler queues target embeds without mutating current prompt_embeds."""
|
|
state = _make_diffusion_request_state()
|
|
|
|
PromptInteractionHandler.from_pipeline(pipeline).enqueue(
|
|
state,
|
|
event_id="ui-update-1",
|
|
received_at=0.0,
|
|
payload={"prompt": "new scene"},
|
|
transition_chunks=2,
|
|
)
|
|
|
|
session = state.interaction_sessions["prompt"]
|
|
assert isinstance(session, PromptSession)
|
|
pending = session.pending_event
|
|
assert pending is not None
|
|
assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0))
|
|
assert pending.transition_chunks == 2
|
|
assert torch.equal(state.prompt_embeds, torch.zeros(1, 4, 2)) # pyright: ignore[reportArgumentType]
|
|
|
|
def test_helios_prompt_handler_enqueue_rejects_before_initial_generation(self, pipeline: HeliosPipeline) -> None:
|
|
"""Reject prompt updates submitted before initial prompt embeds exist."""
|
|
state = _make_diffusion_request_state()
|
|
state.prompt_embeds = None
|
|
|
|
with pytest.raises(
|
|
ValueError,
|
|
match="prompt_update is not allowed before initial generation has started",
|
|
):
|
|
PromptInteractionHandler.from_pipeline(pipeline).enqueue(
|
|
state,
|
|
event_id="ui-update-1",
|
|
received_at=0.0,
|
|
payload={"prompt": "new scene"},
|
|
transition_chunks=2,
|
|
)
|
|
|
|
session = state.interaction_sessions.get("prompt")
|
|
assert session is None or (isinstance(session, PromptSession) and session.pending_event is None)
|
|
|
|
def test_helios_apply_interaction_at_chunk_boundary_starts_transition(self, pipeline: HeliosPipeline) -> None:
|
|
"""At chunk boundary, starts transition state and bumps prompt_update_version."""
|
|
state = _make_diffusion_request_state(fps=None)
|
|
PromptInteractionHandler.from_pipeline(pipeline).enqueue(
|
|
state,
|
|
event_id="ui-update-1",
|
|
received_at=0.0,
|
|
payload={"prompt": "new scene"},
|
|
transition_chunks=2,
|
|
)
|
|
|
|
pipeline.apply_interaction_at_chunk_boundary(state)
|
|
|
|
session = state.interaction_sessions["prompt"]
|
|
assert isinstance(session, PromptSession)
|
|
assert session.active_event is not None
|
|
assert torch.equal(state.prompt_embeds, torch.ones(1, 4, 2)) # pyright: ignore[reportArgumentType]
|
|
assert session.version == 1
|
|
assert state.interaction_chunk_metadata is not None
|
|
assert state.interaction_chunk_metadata.as_dict() == {
|
|
"started_event_ids": ["ui-update-1"],
|
|
"active_event_ids": ["ui-update-1"],
|
|
"completed_event_ids": [],
|
|
}
|
|
|
|
def test_helios_apply_interaction_advances_transition_over_chunks(self, pipeline: HeliosPipeline) -> None:
|
|
"""At chunk boundary, interpolates embeds until the target prompt is reached."""
|
|
state = _make_diffusion_request_state(fps=None)
|
|
PromptInteractionHandler.from_pipeline(pipeline).enqueue(
|
|
state,
|
|
event_id="ui-update-1",
|
|
received_at=0.0,
|
|
payload={"prompt": "new scene"},
|
|
transition_chunks=3,
|
|
)
|
|
pipeline.apply_interaction_at_chunk_boundary(state)
|
|
|
|
assert torch.allclose(state.prompt_embeds, torch.full((1, 4, 2), 2.0 / 3.0)) # pyright: ignore[reportArgumentType]
|
|
|
|
for _ in range(2):
|
|
pipeline.apply_interaction_at_chunk_boundary(state)
|
|
|
|
assert torch.allclose(state.prompt_embeds, torch.full((1, 4, 2), 2.0)) # pyright: ignore[reportArgumentType]
|
|
session = state.interaction_sessions["prompt"]
|
|
assert isinstance(session, PromptSession)
|
|
assert session.version == 3
|
|
|
|
# Further chunk boundaries after completion must not bump the version.
|
|
pipeline.apply_interaction_at_chunk_boundary(state)
|
|
assert session.version == 3
|
|
|
|
def test_stepwise_runner_applies_prompt_update_and_acks(
|
|
self,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Drive the production runner loop: submit → apply(no media args) → ACK on next chunk.
|
|
|
|
Also regresses ``sampling.fps is None`` on the prompt-only path (no peek).
|
|
ACK ids land on the *next* chunk's ``DiffusionOutput`` via
|
|
``_attach_stepwise_metadata`` consuming the prior boundary metadata.
|
|
"""
|
|
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
|
|
from vllm_omni.diffusion.interaction.mixin import InteractionMixin
|
|
from vllm_omni.diffusion.interaction.types import ChunkMediaSpec
|
|
|
|
chunks = [
|
|
DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2), # pyright: ignore[reportArgumentType]
|
|
DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2), # pyright: ignore[reportArgumentType]
|
|
]
|
|
|
|
class _InteractiveChunkPipeline(InteractionMixin):
|
|
device = torch.device("cpu")
|
|
supports_step_execution = True
|
|
|
|
def __init__(self) -> None:
|
|
self._outputs = chunks
|
|
self.decode_calls = 0
|
|
self.transformer = SimpleNamespace(dtype=torch.float32)
|
|
self.encode_prompt = MagicMock(return_value=(torch.full((1, 4, 2), 2.0), None))
|
|
self.prepare_next_chunk = MagicMock()
|
|
od_config = SimpleNamespace(model_class_name="HeliosPipeline")
|
|
self._interaction_coordinator = InteractionCoordinator.build(self, od_config) # pyright: ignore[reportArgumentType]
|
|
|
|
def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec:
|
|
# Same contract as Helios: never coerce missing fps to 0.0.
|
|
fps = state.sampling.fps
|
|
if fps is None or float(fps) <= 0:
|
|
raise ValueError(
|
|
"sampling.fps is required and must be > 0 for interaction modalities "
|
|
f"that use the chunk media timeline, got {fps!r} "
|
|
f"(request_id={state.request_id!r})"
|
|
)
|
|
return ChunkMediaSpec(
|
|
num_media_frames=int(state.extra["window_num_frames"]),
|
|
fps=float(fps),
|
|
num_latent_frames=int(state.extra["window_num_frames"]),
|
|
)
|
|
|
|
def prepare_encode(self, state: StepRequestState) -> StepRequestState:
|
|
state.prompt_embeds = torch.zeros(1, 4, 2)
|
|
state.latents = torch.zeros(1, 1)
|
|
state.timesteps = torch.tensor([1.0, 0.0, 1.0, 0.0])
|
|
state.step_index = 0
|
|
state.step_in_chunk = 0
|
|
state.chunk_num_steps = 2
|
|
state.total_chunks = 2
|
|
state.extra = {"window_num_frames": 8}
|
|
return state
|
|
|
|
def denoise_step(self, input_batch: Any, states: Any) -> torch.Tensor:
|
|
del states
|
|
return torch.ones_like(input_batch.latents)
|
|
|
|
def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor | None) -> None:
|
|
state.latents = noise_pred
|
|
state.step_index += 1
|
|
state.step_in_chunk += 1
|
|
|
|
def post_decode(self, state: StepRequestState) -> DiffusionOutput:
|
|
output = self._outputs[self.decode_calls]
|
|
self.decode_calls += 1
|
|
state.chunk_index += 1
|
|
state.step_index = 0
|
|
state.step_in_chunk = 0
|
|
if not state.request_denoise_completed:
|
|
state.latents = torch.zeros(1, 1)
|
|
return output
|
|
|
|
pipeline = _InteractiveChunkPipeline()
|
|
runner = _make_diffusion_model_runner(pipeline=pipeline)
|
|
runner._supports_step_mode = lambda: DiffusionModelRunner._supports_step_mode(runner)
|
|
runner.od_config.streaming_output = True
|
|
runner.od_config.step_execution = True
|
|
runner.diffusion_kv_backend = SimpleNamespace(remove_diffusion_kv_requests=lambda *_: 0) # pyright: ignore[reportAttributeAccessIssue]
|
|
|
|
sampling = SimpleNamespace(
|
|
generator=None,
|
|
seed=None,
|
|
generator_device=None,
|
|
num_inference_steps=4,
|
|
num_outputs_per_prompt=1,
|
|
max_sequence_length=226,
|
|
fps=None,
|
|
)
|
|
req = SimpleNamespace(
|
|
request_id="req",
|
|
prompt="a prompt",
|
|
sampling_params=sampling,
|
|
kv_sender_info=None,
|
|
)
|
|
|
|
@contextmanager
|
|
def _noop_forward_context(*args: Any, **kwargs: Any):
|
|
del args, kwargs
|
|
yield
|
|
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
|
|
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
|
|
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
|
|
monkeypatch.setattr(model_runner_module.current_omni_platform, "is_available", lambda: False)
|
|
|
|
scheduler_output = SimpleNamespace(
|
|
finished_req_ids=set(),
|
|
scheduled_new_reqs=[SimpleNamespace(request_id="req", req=req, diffusion_kv_metadata=None)],
|
|
scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
|
|
)
|
|
# Admit request and run first denoise step (no chunk boundary yet).
|
|
DiffusionModelRunner.execute_stepwise(runner, scheduler_output) # pyright: ignore[reportArgumentType]
|
|
assert "req" in runner.state_cache
|
|
assert runner.state_cache["req"].sampling.fps is None
|
|
|
|
runner.submit_interaction("req", _prompt_interaction("new scene", transition_chunks=2))
|
|
|
|
cached = SimpleNamespace(
|
|
finished_req_ids=set(),
|
|
scheduled_new_reqs=[],
|
|
scheduled_cached_reqs=SimpleNamespace(request_ids=["req"]),
|
|
)
|
|
# Second step completes chunk 0: apply runs (no peek), metadata staged.
|
|
chunk0 = DiffusionModelRunner.execute_stepwise(runner, cached).get_request_output("req") # pyright: ignore[reportArgumentType]
|
|
assert chunk0 is not None and chunk0.result is not None
|
|
assert chunk0.result.started_event_ids == []
|
|
state = runner.state_cache["req"]
|
|
assert torch.equal(state.prompt_embeds, torch.ones(1, 4, 2)) # pyright: ignore[reportArgumentType]
|
|
assert state.interaction_chunk_metadata is not None
|
|
assert state.interaction_chunk_metadata.started_event_ids == ["ui-update-1"]
|
|
pipeline.prepare_next_chunk.assert_called() # pyright: ignore[reportAttributeAccessIssue]
|
|
|
|
# Next chunk decode attaches prior-boundary ACK ids onto DiffusionOutput.
|
|
DiffusionModelRunner.execute_stepwise(runner, cached) # pyright: ignore[reportArgumentType]
|
|
chunk1 = DiffusionModelRunner.execute_stepwise(runner, cached).get_request_output("req") # pyright: ignore[reportArgumentType]
|
|
assert chunk1 is not None and chunk1.result is not None
|
|
assert chunk1.result.started_event_ids == ["ui-update-1"]
|
|
assert chunk1.result.active_event_ids == ["ui-update-1"]
|
|
assert chunk1.finished is True
|
|
|
|
|
|
class TestPromptUpdateIntegration:
|
|
"""AsyncOmni prompt update through orchestrator, inline client, and runner."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_orchestrator_prompt_update_failure_surfaces_non_fatal_error(
|
|
self,
|
|
) -> None:
|
|
"""Worker-side prompt_update rejection is reported without crashing the orchestrator."""
|
|
|
|
class _RejectingStagePool:
|
|
async def submit_interaction(self, *args: Any, **kwargs: Any) -> None:
|
|
del args, kwargs
|
|
raise ValueError("prompt embeds are not ready")
|
|
|
|
output_queue: asyncio.Queue[ErrorMessage] = asyncio.Queue()
|
|
orchestrator = object.__new__(Orchestrator)
|
|
orchestrator.stage_pools = [_RejectingStagePool()] # pyright: ignore[reportAttributeAccessIssue]
|
|
orchestrator.output_async_queue = output_queue # pyright: ignore[reportAttributeAccessIssue]
|
|
orchestrator.request_states = {"req-1": OrchestratorRequestState(request_id="req-1")}
|
|
|
|
await orchestrator._handle_interaction(
|
|
InteractionMessage(
|
|
request_id="req-1",
|
|
interaction=_prompt_interaction("new scene", event_id="ui-update-1", transition_chunks=2),
|
|
)
|
|
)
|
|
|
|
msg = output_queue.get_nowait()
|
|
assert msg == ErrorMessage(
|
|
error="Failed interaction for request req-1: prompt embeds are not ready",
|
|
fatal=False,
|
|
request_id="req-1",
|
|
event_id="ui-update-1",
|
|
stage_id=0,
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runner_prompt_update_failure_surfaces_non_fatal_error(self, pipeline: HeliosPipeline) -> None:
|
|
"""Runner-side prompt encode rejection is reported through the orchestrator."""
|
|
pipeline.encode_prompt = MagicMock( # pyright: ignore[reportAttributeAccessIssue]
|
|
side_effect=ValueError("prompt embeds are not ready")
|
|
)
|
|
runner = _make_diffusion_model_runner(pipeline=pipeline)
|
|
runner.state_cache["req-1"] = _make_diffusion_request_state(request_id="req-1")
|
|
prompt_update_engine = self._PromptUpdateEngine(self._IncrementalStreamingPipeline(), runner)
|
|
inline_client = self._make_inline_pipeline_client(prompt_update_engine)
|
|
|
|
output_queue: asyncio.Queue[ErrorMessage] = asyncio.Queue()
|
|
orchestrator = object.__new__(Orchestrator)
|
|
orchestrator.stage_pools = [ # pyright: ignore[reportAttributeAccessIssue]
|
|
StagePool(0, cast(StagePoolClient, inline_client))
|
|
]
|
|
orchestrator.output_async_queue = output_queue # pyright: ignore[reportAttributeAccessIssue]
|
|
orchestrator.request_states = {"req-1": OrchestratorRequestState(request_id="req-1")}
|
|
|
|
try:
|
|
await orchestrator._handle_interaction(
|
|
InteractionMessage(
|
|
request_id="req-1",
|
|
interaction=_prompt_interaction("new scene", event_id="ui-update-1", transition_chunks=2),
|
|
)
|
|
)
|
|
finally:
|
|
inline_client.shutdown()
|
|
|
|
pipeline.encode_prompt.assert_called_once() # pyright: ignore[reportAttributeAccessIssue]
|
|
msg = output_queue.get_nowait()
|
|
assert msg == ErrorMessage(
|
|
error="Failed interaction for request req-1: prompt embeds are not ready",
|
|
fatal=False,
|
|
request_id="req-1",
|
|
event_id="ui-update-1",
|
|
stage_id=0,
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prompt_update_reaches_runner_from_async_omni(self, pipeline: HeliosPipeline) -> None:
|
|
"""Midway prompt update submitted via AsyncOmni reaches the diffusion runner."""
|
|
streaming_pipeline = self._IncrementalStreamingPipeline()
|
|
runner = _make_diffusion_model_runner(pipeline=pipeline)
|
|
prompt_update_engine = self._PromptUpdateEngine(streaming_pipeline, runner)
|
|
inline_client = self._make_inline_pipeline_client(prompt_update_engine)
|
|
fixture = _build_harness([inline_client])
|
|
omni = self._make_async_omni(self._OrchestratorBridgeEngine(fixture))
|
|
|
|
generate_task: asyncio.Task[list[OmniRequestOutput]] | None = None
|
|
try:
|
|
generate_task = asyncio.create_task(self._collect_generate_outputs(omni))
|
|
|
|
await _wait_for(
|
|
lambda: (
|
|
len(runner.state_cache) > 0
|
|
and any(state.external_request_id == "req-omni" for state in omni.request_states.values())
|
|
)
|
|
)
|
|
internal_request_id = next(iter(runner.state_cache))
|
|
|
|
await omni.submit_interaction_async(
|
|
"req-omni",
|
|
interaction=_prompt_interaction("new scene", event_id="ui-update-1", transition_chunks=2),
|
|
)
|
|
|
|
outputs = await generate_task
|
|
|
|
pending = runner.state_cache[internal_request_id].interaction_sessions["prompt"].pending_event
|
|
assert pending is not None
|
|
assert pending.transition_chunks == 2
|
|
assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0))
|
|
pipeline.encode_prompt.assert_called_once() # pyright: ignore[reportAttributeAccessIssue]
|
|
assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
|
|
finally:
|
|
if generate_task is not None and not generate_task.done():
|
|
generate_task.cancel()
|
|
await asyncio.gather(generate_task, return_exceptions=True)
|
|
await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
|
|
|
|
@staticmethod
|
|
async def _collect_generate_outputs(omni: AsyncOmni) -> list[OmniRequestOutput]:
|
|
outputs: list[OmniRequestOutput] = []
|
|
async for output in omni.generate(
|
|
prompt={"prompt": "a cat"},
|
|
request_id="req-omni",
|
|
sampling_params_list=[OmniDiffusionSamplingParams()],
|
|
output_modalities=["image"],
|
|
):
|
|
outputs.append(output)
|
|
return outputs
|
|
|
|
@classmethod
|
|
def _make_inline_pipeline_client(cls, engine: _PromptUpdateEngine) -> InlineStageDiffusionClient:
|
|
metadata = StageMetadata(
|
|
stage_id=0,
|
|
stage_type="diffusion",
|
|
engine_output_type="image",
|
|
is_comprehension=False,
|
|
requires_multimodal_data=False,
|
|
engine_input_source=[],
|
|
final_output=True,
|
|
final_output_type="image",
|
|
default_sampling_params=OmniDiffusionSamplingParams(),
|
|
custom_process_input_func=None,
|
|
model_stage=None,
|
|
runtime_cfg=None,
|
|
)
|
|
with patch.object(InlineStageDiffusionClient, "_enrich_config"):
|
|
with patch(
|
|
"vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine.make_engine",
|
|
return_value=engine,
|
|
):
|
|
od_config = MagicMock(spec=OmniDiffusionConfig)
|
|
od_config.streaming_output = True
|
|
return InlineStageDiffusionClient(
|
|
model="test_model",
|
|
od_config=od_config,
|
|
metadata=metadata,
|
|
)
|
|
|
|
@staticmethod
|
|
def _make_async_omni(engine: _OrchestratorBridgeEngine) -> AsyncOmni:
|
|
omni = object.__new__(AsyncOmni)
|
|
omni.engine = engine # pyright: ignore[reportAttributeAccessIssue]
|
|
omni.log_stats = False
|
|
omni._pause_cond = asyncio.Condition()
|
|
omni._paused = False
|
|
omni.request_states = {}
|
|
omni.final_output_task = None
|
|
omni.event_resolver = AsyncEventResolver()
|
|
omni._enable_ar_profiler = False
|
|
omni.prom_metrics = MagicMock()
|
|
omni.mod_metrics = MagicMock()
|
|
omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params # pyright: ignore[reportAttributeAccessIssue]
|
|
omni._compute_final_stage_id = lambda output_modalities: 0
|
|
omni._compute_final_output_stage_ids = lambda output_modalities: [0]
|
|
omni.default_sampling_params_list = engine.default_sampling_params_list # pyright: ignore[reportAttributeAccessIssue]
|
|
omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None) # pyright: ignore[reportAttributeAccessIssue]
|
|
return omni
|
|
|
|
@staticmethod
|
|
async def _shutdown_pipeline_omni_harness(
|
|
omni: AsyncOmni,
|
|
fixture: OrchestratorFixture,
|
|
inline_client: InlineStageDiffusionClient,
|
|
) -> None:
|
|
if omni.final_output_task is not None:
|
|
omni.final_output_task.cancel()
|
|
await asyncio.gather(omni.final_output_task, return_exceptions=True)
|
|
inline_client.shutdown()
|
|
fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
|
|
await asyncio.to_thread(fixture.thread.join, 5)
|
|
|
|
class _IncrementalStreamingPipeline:
|
|
"""Returns one streaming chunk per ``step_outputs`` call."""
|
|
|
|
supports_step_execution = True
|
|
|
|
def __init__(self) -> None:
|
|
self.requests: list[Any] = []
|
|
self._step_index = 0
|
|
|
|
def step_outputs(self, request):
|
|
self.requests.append(request)
|
|
if self._step_index == 0:
|
|
self._step_index += 1
|
|
return [DiffusionOutput(output={"chunk": 0}, finished=False)]
|
|
return [DiffusionOutput(output={"chunk": 1}, finished=True)]
|
|
|
|
class _OrchestratorBridgeEngine:
|
|
"""Minimal AsyncOmni engine facade backed by a live Orchestrator harness."""
|
|
|
|
def __init__(self, fixture: OrchestratorFixture) -> None:
|
|
self._fixture = fixture
|
|
self.stage_metadata = [
|
|
StageRuntimeInfo(
|
|
stage_type="diffusion",
|
|
final_output=True,
|
|
final_output_type="image",
|
|
)
|
|
]
|
|
self.stage_configs: list[Any] = [SimpleNamespace(stage_type="diffusion")]
|
|
self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
|
|
self.num_stages = 1
|
|
self.supported_tasks = ("generate",)
|
|
self._alive = True
|
|
|
|
async def add_request_async(
|
|
self,
|
|
*,
|
|
request_id: str,
|
|
prompt: Any,
|
|
sampling_params_list: list[Any],
|
|
final_stage_id: int,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self._fixture.request_sync_q.put_nowait(
|
|
StageSubmissionMessage(
|
|
type="add_request",
|
|
request_id=request_id,
|
|
prompt=prompt,
|
|
original_prompt=prompt,
|
|
output_prompt_text=None,
|
|
sampling_params_list=sampling_params_list,
|
|
final_stage_id=final_stage_id,
|
|
final_output_stage_ids=kwargs.get("final_output_stage_ids"),
|
|
preprocess_ms=0.0,
|
|
request_timestamp=kwargs.get("arrival_time", time.time()),
|
|
enqueue_ts=time.perf_counter(),
|
|
)
|
|
)
|
|
|
|
async def submit_interaction_async(
|
|
self,
|
|
request_id: str,
|
|
*,
|
|
interaction: dict[str, Any],
|
|
) -> None:
|
|
self._fixture.request_sync_q.put_nowait(
|
|
InteractionMessage(
|
|
request_id=request_id,
|
|
interaction=cast(Any, interaction),
|
|
)
|
|
)
|
|
|
|
async def try_get_output_async(self) -> Any | None:
|
|
try:
|
|
return self._fixture.output_sync_q.get_nowait()
|
|
except queue.Empty:
|
|
return None
|
|
|
|
def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo:
|
|
return self.stage_metadata[stage_id]
|
|
|
|
def is_alive(self) -> bool:
|
|
return self._fixture.thread.is_alive()
|
|
|
|
async def abort_async(self, request_ids: list[str]) -> None:
|
|
del request_ids
|
|
|
|
class _PromptUpdateEngine:
|
|
"""DiffusionEngine stand-in that streams chunks and routes prompt_update to a runner."""
|
|
|
|
def __init__(
|
|
self,
|
|
streaming_pipeline: TestPromptUpdateIntegration._IncrementalStreamingPipeline,
|
|
runner: DiffusionModelRunner,
|
|
) -> None:
|
|
self.streaming_pipeline = streaming_pipeline
|
|
self._runner = runner
|
|
self.executor = SimpleNamespace(
|
|
register_failure_callback=MagicMock(),
|
|
check_health=MagicMock(),
|
|
)
|
|
|
|
async def step_streaming(self, request):
|
|
if request.request_id not in self._runner.state_cache:
|
|
state = _make_diffusion_request_state(request_id=request.request_id)
|
|
self._runner.state_cache[request.request_id] = state
|
|
|
|
state = self._runner.state_cache[request.request_id]
|
|
|
|
outputs = self.streaming_pipeline.step_outputs(request)
|
|
output = outputs[0]
|
|
custom_output = cast(dict[str, Any], output.output) or {}
|
|
yield [
|
|
OmniRequestOutput.from_diffusion(
|
|
request_id=request.request_id,
|
|
images=list(custom_output.get("images") or []),
|
|
custom_output=custom_output,
|
|
finished=output.finished,
|
|
)
|
|
]
|
|
|
|
deadline = time.monotonic() + 5.0
|
|
while (
|
|
not isinstance(state.interaction_sessions.get("prompt"), PromptSession)
|
|
or state.interaction_sessions["prompt"].pending_event is None
|
|
):
|
|
if time.monotonic() >= deadline:
|
|
raise TimeoutError("timed out waiting for prompt update during streaming")
|
|
await asyncio.sleep(0.01)
|
|
|
|
outputs = self.streaming_pipeline.step_outputs(request)
|
|
output = outputs[0]
|
|
custom_output = cast(dict[str, Any], output.output) or {}
|
|
yield [
|
|
OmniRequestOutput.from_diffusion(
|
|
request_id=request.request_id,
|
|
images=list(custom_output.get("images") or []),
|
|
custom_output=custom_output,
|
|
finished=output.finished,
|
|
)
|
|
]
|
|
|
|
def collective_rpc(self, method, timeout, args, kwargs, unique_reply_rank):
|
|
del timeout, kwargs, unique_reply_rank
|
|
if method == "submit_interaction":
|
|
request_id, interaction = args
|
|
self._runner.submit_interaction(request_id, interaction)
|
|
return None
|
|
raise NotImplementedError(f"collective_rpc not mocked for {method!r}")
|
|
|
|
def abort(self, request_id: str) -> None:
|
|
del request_id
|