OmniRequestOutput inherits RequestOutput

Signed-off-by: Boao Shi <aoibosh@connect.hku.hk>

vllm_omni/distributed/omni_connectors/utils/serialization.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from collections.abc import Mapping
-from dataclasses import asdict, is_dataclass
+from dataclasses import asdict, fields, is_dataclass
 from typing import Any
 
 import msgspec
@@ -59,7 +59,10 @@ class OmniMsgpackEncoder:
         if isinstance(obj, Image.Image):
             return self._encode_pil_image(obj)
 
-        # RequestOutput (not a dataclass, needs special handling)
+        # RequestOutput (not a dataclass, needs special handling).
+        # Note: OmniRequestOutput subclasses RequestOutput but is a dataclass,
+        # so msgspec encodes it natively (field map) and it never reaches
+        # this hook.
         if isinstance(obj, RequestOutput):
             return self._encode_request_output(obj)
 
@@ -252,23 +255,30 @@ class OmniMsgpackDecoder:
     def _decode_omni_request_output(self, obj: dict[str, Any]) -> Any:
         """Decode dict to OmniRequestOutput.
 
-        OmniRequestOutput is a dataclass, so we can use msgspec.convert
-        or construct it directly.
+        By the time this runs, ``_post_process`` has already reconstructed
+        nested values bottom-up (CompletionOutput, tensors, images, ...), so
+        the object is built directly from the known dataclass fields.
         """
         from vllm_omni.outputs import OmniRequestOutput
 
+        data = dict(obj)
+        # Legacy wire format (pre-inheritance): the stage output was nested
+        # under "request_output". Merge its content via the compat setter.
+        legacy_inner = data.pop("request_output", None)
+        multi_modal_placeholders = data.pop("multi_modal_placeholders", None)
+        field_names = {f.name for f in fields(OmniRequestOutput)}
+        known = {k: v for k, v in data.items() if k in field_names}
         try:
-            # Use msgspec.convert for dataclass reconstruction
-            return msgspec.convert(obj, OmniRequestOutput)
+            out = OmniRequestOutput(**known)
         except Exception:
-            try:
-                # Fallback: construct directly if msgspec.convert fails
-                # (e.g., if some fields are missing or have wrong types)
-                return OmniRequestOutput(**obj)
-            except Exception:
-                # If both attempts fail, return dict as-is (defensive fallback)
-                # This should rarely happen if _is_omni_request_output is correct
+            # Defensive fallback; should rarely happen if
+            # _is_omni_request_output is correct.
             return obj
+        if legacy_inner is not None and not isinstance(legacy_inner, dict):
+            out.request_output = legacy_inner
+        if multi_modal_placeholders is not None:
+            setattr(out, "multi_modal_placeholders", multi_modal_placeholders)
+        return out
 
     def _decode_tensor(self, obj: dict[str, Any]) -> torch.Tensor:
         """Decode dict to torch.Tensor."""

vllm_omni/entrypoints/omni_base.py

@@ -577,7 +577,6 @@ class OmniBase(PDDisaggregationMixin):
         self.prom_metrics.set_running(running)
         self.prom_metrics.set_waiting(max(0, total - running))
 
-        images = getattr(engine_outputs, "images", []) if output_type == "image" else []
         response_metrics: dict[str, Any] = {}
         stage_metrics: dict[str, dict[str, Any]] = {}
         rid_key = str(req_id)
@@ -602,17 +601,14 @@ class OmniBase(PDDisaggregationMixin):
                 response_metrics["final_output_type"] = current_stage_metrics["final_output_type"]
                 response_metrics["num_tokens_in"] = current_stage_metrics["num_tokens_in"]
                 response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"]
+        # Generation content (outputs, prompt, images, trajectory_*, ...) is
+        # copied from engine_outputs onto the returned object by
+        # OmniRequestOutput.__post_init__.
         return OmniRequestOutput(
             request_id=req_id or "",
             stage_id=stage_id,
             final_output_type=output_type,
             request_output=engine_outputs,
-            images=images,
-            trajectory_latents=getattr(engine_outputs, "trajectory_latents", None),
-            trajectory_timesteps=getattr(engine_outputs, "trajectory_timesteps", None),
-            trajectory_log_probs=getattr(engine_outputs, "trajectory_log_probs", None),
-            trajectory_decoded=getattr(engine_outputs, "trajectory_decoded", None),
-            _custom_output=getattr(engine_outputs, "_custom_output", {}),
             metrics=response_metrics,
             stage_durations=stage_durations,
             peak_memory_mb=peak_memory_mb,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -3113,7 +3113,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 request_id=request_id,
             )
 
-        images = getattr(result.request_output, "images", [])
+        images = getattr(result, "images", [])
         stage_durations = result.stage_durations
         peak_memory_mb = result.peak_memory_mb
         cot_output = None
@@ -3184,7 +3184,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         )
                         yield f"data: {chunk.model_dump_json()}\n\n"
                 elif final_output_type == "image":
-                    images = self._flatten_diffusion_images(getattr(output.request_output, "images", []))
+                    images = self._flatten_diffusion_images(getattr(output, "images", []))
                     if not images:
                         raise RuntimeError("Streaming image edit produced an empty final image output.")
                     image_data = [
@@ -3488,8 +3488,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
             # Image output path (text2img / img2img)
             final_output_type = getattr(result, "final_output_type", "image")
-            # Handle nested OmniRequestOutput structure where images might be in request_output
-            images = getattr(result.request_output, "images", [])
+            images = getattr(result, "images", [])
             multimodal_output = getattr(result, "multimodal_output", {}) or {}
             stage_durations = result.stage_durations
             peak_memory_mb = result.peak_memory_mb

vllm_omni/outputs/__init__.py

@@ -3,7 +3,7 @@ from typing import Any
 
 import torch
 from PIL import Image
-from vllm.outputs import RequestOutput
+from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.v1.outputs import ModelRunnerOutput
 
 from vllm_omni.inputs.data import OmniPromptType
@@ -59,12 +59,58 @@ class OmniModelRunnerOutput(ModelRunnerOutput):
     omni_connector_output: OmniConnectorOutput | None = None
 
 
+# RequestOutput attributes that hold generation content.  When wrapping a
+# stage output, these are copied onto the OmniRequestOutput itself so that
+# the object IS the RequestOutput instead of holding one.
+_REQUEST_OUTPUT_CONTENT_ATTRS = (
+    "prompt",
+    "prompt_token_ids",
+    "prompt_logprobs",
+    "outputs",
+    "finished",
+    "lora_request",
+    "encoder_prompt",
+    "encoder_prompt_token_ids",
+    "num_cached_tokens",
+    "kv_transfer_params",
+)
+
+# Omni-specific content copied when the wrapped stage output is itself an
+# OmniRequestOutput (e.g. a diffusion stage inside a multi-stage pipeline).
+_OMNI_CONTENT_ATTRS = (
+    "images",
+    "latents",
+    "trajectory_latents",
+    "trajectory_timesteps",
+    "trajectory_log_probs",
+    "trajectory_decoded",
+    "_multimodal_output",
+    "_custom_output",
+)
+
+# Dynamic attributes that must not be copied verbatim onto the wrapper:
+# they collide with properties defined on OmniRequestOutput.
+_COPY_SKIP_ATTRS = frozenset({"request_output", "multimodal_output", "custom_output"})
+
+
 @dataclass
-class OmniRequestOutput:
+class OmniRequestOutput(RequestOutput):
     """Unified request output for both pipeline stages and diffusion models.
 
+    Extends vLLM's ``RequestOutput`` so that omni outputs can flow directly
+    through vLLM serving codepaths (which expect ``prompt_token_ids``,
+    ``outputs``, etc. as real attributes).  The inherited fields store the
+    LLM generation content; omni-specific fields store pipeline/diffusion
+    extras.
+
+    Note: ``RequestOutput`` is a plain class (not a dataclass), so all of its
+    attributes are redeclared below as dataclass fields with defaults — the
+    dataclass-generated ``__init__`` replaces ``RequestOutput.__init__`` and
+    must set them itself.
+
     This class handles outputs from:
-    1. Multi-stage LLM pipelines (with stage_id, final_output_type, request_output)
+    1. Multi-stage LLM pipelines (with stage_id, final_output_type, and the
+       inherited RequestOutput fields carrying the stage's generation content)
     2. Diffusion models (with images, prompt, metrics)
 
     Attributes:
@@ -72,30 +118,38 @@ class OmniRequestOutput:
         finished: Whether generation is complete
         stage_id: Identifier of the stage that produced this output (pipeline mode)
         final_output_type: Type of output ("text", "image", "audio", "latents")
-        request_output: The underlying RequestOutput from the stage (pipeline mode)
         images: List of generated PIL images (diffusion mode)
-        prompt: The prompt used for generation (diffusion mode)
+        prompt: The prompt used for generation
         latents: Optional tensor of latent representations (diffusion mode)
-        metrics: Optional dictionary of generation metrics
+        metrics: Generation metrics. A plain dict for omni outputs; may carry
+            vLLM's request stats object when copied from a raw RequestOutput.
     """
 
+    # --- Inherited RequestOutput attributes (redeclared as fields) ---
     request_id: str = ""
+    prompt: OmniPromptType | None = None
+    prompt_token_ids: list[int] | None = None
+    prompt_logprobs: Any = None
+    outputs: list[CompletionOutput] = field(default_factory=list)
     finished: bool = True
+    metrics: Any = field(default_factory=dict)
+    lora_request: Any = None
+    encoder_prompt: str | None = None
+    encoder_prompt_token_ids: list[int] | None = None
+    num_cached_tokens: int | None = None
+    kv_transfer_params: dict[str, Any] | None = None
 
-    # Pipeline stage fields
+    # --- Pipeline stage fields ---
     stage_id: int | None = None
     final_output_type: str = "text"
-    request_output: RequestOutput | None = None
 
-    # Diffusion model fields
+    # --- Diffusion model fields ---
     images: list[Image.Image] = field(default_factory=list)
-    prompt: OmniPromptType | None = None
     latents: torch.Tensor | None = None
     trajectory_latents: torch.Tensor | None = None
     trajectory_timesteps: torch.Tensor | None = None
     trajectory_log_probs: torch.Tensor | None = None
     trajectory_decoded: list | None = None
-    metrics: dict[str, Any] = field(default_factory=dict)
     _multimodal_output: dict[str, Any] = field(default_factory=dict)
     _custom_output: dict[str, Any] = field(default_factory=dict)
 
@@ -110,6 +164,44 @@ class OmniRequestOutput:
     error_status_code: int | None = None
     error_type: str | None = None
 
+    def _copy_content_from(self, source: Any) -> None:
+        """Copy generation content from a stage output into this object.
+
+        ``source`` may be a vLLM ``RequestOutput``, another
+        ``OmniRequestOutput``, or a duck-typed stand-in (tests use mocks and
+        SimpleNamespace). Attribute presence is checked against the source's
+        instance ``__dict__`` so that auto-created Mock attributes are not
+        copied. Control fields set by the wrapper (stage_id,
+        final_output_type, metrics, stage_durations, ...) are intentionally
+        not copied.
+        """
+        source_vars = getattr(source, "__dict__", None)
+
+        def _has(name: str) -> bool:
+            if source_vars is not None:
+                return name in source_vars
+            return hasattr(source, name)
+
+        for name in _REQUEST_OUTPUT_CONTENT_ATTRS + _OMNI_CONTENT_ATTRS:
+            if _has(name):
+                setattr(self, name, getattr(source, name))
+        if not self.request_id:
+            self.request_id = getattr(source, "request_id", "") or ""
+        # Dynamic attribute that vLLM code attaches to RequestOutput; it is a
+        # read-only property on this class, so store it in _multimodal_output.
+        if not isinstance(source, OmniRequestOutput) and _has("multimodal_output"):
+            mm = source.multimodal_output
+            if mm is not None:
+                self._multimodal_output = mm
+        # Preserve remaining dynamic attributes (e.g. dynamically annotated
+        # stage outputs) so they stay reachable through `request_output`.
+        if source_vars:
+            known = type(self).__dataclass_fields__.keys()
+            for name, value in source_vars.items():
+                if name.startswith("_") or name in known or name in _COPY_SKIP_ATTRS:
+                    continue
+                setattr(self, name, value)
+
     @classmethod
     def from_error(
         cls,
@@ -143,7 +235,10 @@ class OmniRequestOutput:
         final_output_type: str,
         request_output: RequestOutput,
     ) -> "OmniRequestOutput":
-        """Create output from pipeline stage.
+        """Create output from a pipeline stage.
+
+        The stage output's content is copied into the new object's own
+        (inherited) fields.
 
         Args:
             stage_id: Stage identifier
@@ -153,13 +248,13 @@ class OmniRequestOutput:
         Returns:
             OmniRequestOutput configured for pipeline mode
         """
-        return cls(
-            request_id=getattr(request_output, "request_id", ""),
+        out = cls(
             stage_id=stage_id,
             final_output_type=final_output_type,
             request_output=request_output,
-            finished=True,
         )
+        out.finished = True
+        return out
 
     @classmethod
     def from_diffusion(
@@ -220,35 +315,21 @@ class OmniRequestOutput:
 
     @property
     def multimodal_output(self) -> Any:
-        """Return multimodal output from the underlying request output or local field.
+        """Return the multimodal output payload.
 
-        For pipeline outputs, this checks completion outputs first, then request_output.
-        For diffusion outputs, this returns the local _multimodal_output field.
+        Checks completion outputs first (where multimodal_output is attached
+        by AR stages), then the local _multimodal_output field.
 
         Returns either a MultimodalPayload (Phase 3+) or a plain dict (legacy).
         """
-        if self.request_output is None:
-            return self._multimodal_output
-
-        # Check completion outputs first (where multimodal_output is attached)
-        for output in getattr(self.request_output, "outputs", []):
+        for output in self.outputs:
             if mm := getattr(output, "multimodal_output", None):
                 return mm
-        if mm := getattr(self.request_output, "multimodal_output", None):
-            return mm
         return self._multimodal_output
 
     @property
     def custom_output(self) -> dict[str, Any]:
-        """Return custom output data from diffusion pipelines.
-
-        For diffusion outputs, returns the local _custom_output field.
-        For pipeline outputs with an inner OmniRequestOutput, forwards
-        the custom_output from the inner request output.
-        """
-        if self.request_output is not None:
-            if isinstance(self.request_output, OmniRequestOutput):
-                return self.request_output._custom_output
+        """Return custom output data from diffusion pipelines."""
         return self._custom_output
 
     @custom_output.setter
@@ -260,58 +341,6 @@ class OmniRequestOutput:
         """Return the number of generated images."""
         return len(self.images)
 
-    # Pass-through properties keep vLLM serving codepaths compatible with
-    # OmniRequestOutput for pipeline outputs (Issue #345).
-    @property
-    def prompt_token_ids(self) -> list[int] | None:
-        """Return prompt token IDs from the underlying request output.
-
-        This property is required for compatibility with vLLM's streaming
-        chat completion generator which checks res.prompt_token_ids.
-        """
-        if self.request_output is not None:
-            return getattr(self.request_output, "prompt_token_ids", None)
-        return None
-
-    @property
-    def outputs(self) -> list[Any]:
-        """Return outputs from the underlying request output.
-
-        This property is required for compatibility with vLLM's streaming
-        and non-streaming chat completion generators.
-        """
-        if self.request_output is not None:
-            return getattr(self.request_output, "outputs", [])
-        return []
-
-    @property
-    def encoder_prompt_token_ids(self) -> list[int] | None:
-        """Return encoder prompt token IDs from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "encoder_prompt_token_ids", None)
-        return None
-
-    @property
-    def prompt_logprobs(self) -> Any:
-        """Return prompt logprobs from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "prompt_logprobs", None)
-        return None
-
-    @property
-    def num_cached_tokens(self) -> int | None:
-        """Return number of cached tokens from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "num_cached_tokens", None)
-        return None
-
-    @property
-    def kv_transfer_params(self) -> Any:
-        """Return KV transfer params from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "kv_transfer_params", None)
-        return None
-
     @property
     def is_diffusion_output(self) -> bool:
         """Check if this is a diffusion model output."""
@@ -320,60 +349,33 @@ class OmniRequestOutput:
     @property
     def is_pipeline_output(self) -> bool:
         """Check if this is a pipeline stage output."""
-        return self.stage_id is not None and self.request_output is not None
+        return self.stage_id is not None
 
     def unwrap(self) -> "OmniRequestOutput":
-        """Unwrap nested OmniRequestOutput to get the innermost result.
+        """Return the output itself.
 
-        This helper handles the common pattern where pipeline outputs may wrap
-        other OmniRequestOutput instances. It recursively unwraps until it reaches
-        the final output with actual content (images, text, etc.).
-
-        Returns:
-            The innermost OmniRequestOutput containing the actual generation results.
-
-        Example:
-            ```python
-            result = omni.generate(...)
-            output = OmniRequestOutput.unwrap_result(result)
-            if output.images:
-                # Access images directly
-                video_frames = output.images
-            ```
+        Kept for backward compatibility: outputs used to nest a stage's
+        OmniRequestOutput inside ``request_output``; content is now flattened
+        onto the object itself, so there is nothing to unwrap.
         """
-        current = self
-        # Unwrap nested pipeline outputs
-        while current.is_pipeline_output and current.request_output is not None:
-            if isinstance(current.request_output, OmniRequestOutput):
-                current = current.request_output
-            else:
-                break
-        return current
+        return self
 
     @staticmethod
     def unwrap_result(result: Any) -> "OmniRequestOutput":
-        """Unwrap result from omni.generate() to get the final OmniRequestOutput.
+        """Extract the final OmniRequestOutput from omni.generate() results.
 
-        This static helper handles the full unwrapping pattern including:
+        Handles:
         1. Extracting from list if needed
         2. Type validation
-        3. Recursive unwrapping of nested pipeline outputs
 
         Args:
             result: The result from omni.generate() - may be a list or OmniRequestOutput
 
         Returns:
-            The innermost OmniRequestOutput with actual content
+            The OmniRequestOutput with actual content
 
         Raises:
             ValueError: If result is not an OmniRequestOutput or list containing one
-
-        Example:
-            ```python
-            result = omni.generate(...)
-            output = OmniRequestOutput.unwrap_result(result)
-            # output is guaranteed to be the final OmniRequestOutput
-            ```
         """
         # Handle list wrapper
         if isinstance(result, list):
@@ -385,7 +387,6 @@ class OmniRequestOutput:
         if not isinstance(result, OmniRequestOutput):
             raise ValueError(f"Expected OmniRequestOutput, got {type(result)}")
 
-        # Unwrap nested outputs
         return result.unwrap()
 
     def to_dict(self) -> dict[str, Any]:
@@ -425,7 +426,8 @@ class OmniRequestOutput:
             f"finished={self.finished}",
             f"stage_id={self.stage_id}",
             f"final_output_type={self.final_output_type!r}",
-            f"request_output={self.request_output}",
+            f"prompt_token_ids={self.prompt_token_ids}",
+            f"outputs={self.outputs}",
             f"images={images_repr}",
             f"prompt={self.prompt!r}",
             f"latents={self.latents}",
@@ -437,3 +439,45 @@ class OmniRequestOutput:
         ]
 
         return f"OmniRequestOutput({', '.join(parts)})"
+
+
+def _request_output_get(self: OmniRequestOutput) -> RequestOutput | None:
+    """Backward-compat view of the wrapped stage output.
+
+    The omni output now IS the RequestOutput, so pipeline/LLM outputs return
+    ``self``; pure diffusion or error outputs return None (matching the old
+    composition semantics where ``request_output`` was unset for them).
+    """
+    if self.stage_id is not None or self.outputs:
+        return self
+    return None
+
+
+def _request_output_set(self: OmniRequestOutput, value: Any) -> None:
+    if value is not None:
+        self._copy_content_from(value)
+
+
+# `request_output` is attached after class creation, NOT declared as a field
+# or InitVar: msgspec serializes dataclasses natively via __dataclass_fields__
+# (which includes InitVar pseudo-fields), so a request_output entry would be
+# read through the property during encoding — returning `self` and recursing
+# forever. As a plain property it is invisible to both dataclasses and
+# msgspec.
+OmniRequestOutput.request_output = property(  # type: ignore[assignment]
+    _request_output_get, _request_output_set
+)
+
+# Keep `request_output=` working as a constructor keyword for
+# backward compatibility: the wrapped stage output's content is copied onto
+# the new object (see _copy_content_from).
+_dataclass_init = OmniRequestOutput.__init__
+
+
+def _omni_request_output_init(self: OmniRequestOutput, *args: Any, request_output: Any = None, **kwargs: Any) -> None:
+    _dataclass_init(self, *args, **kwargs)
+    if request_output is not None:
+        self._copy_content_from(request_output)
+
+
+OmniRequestOutput.__init__ = _omni_request_output_init  # type: ignore[method-assign]