vllm-omni/vllm_omni/entrypoints/omni_base.py

from __future__ import annotations

import argparse
import os
import sys
import time
import warnings
import weakref
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Literal

import huggingface_hub
from vllm.logger import init_logger
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError

from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.messages import (
    EngineQueueMessage,
    ErrorMessage,
    OutputMessage,
    StageMetricsMessage,
)
from vllm_omni.entrypoints.client_request_state import ClientRequestState
from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e
from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
from vllm_omni.outputs import OmniRequestOutput

if TYPE_CHECKING:
    from vllm_omni.engine.arg_utils import OmniEngineArgs

logger = init_logger(__name__)


class OmniEngineDeadError(EngineDeadError):
    _DEFAULT_MESSAGE = EngineDeadError().args[0]
    error_stage_id: int | None

    def __init__(
        self,
        message: str | None = None,
        *,
        error_stage_id: int | None = None,
        suppress_context: bool = False,
    ) -> None:
        resolved_message = message or self._DEFAULT_MESSAGE
        Exception.__init__(self, resolved_message)
        self.__suppress_context__ = suppress_context
        self.error_stage_id = error_stage_id


def _weak_shutdown_engine(engine: AsyncOmniEngine) -> None:
    """Best-effort engine cleanup for GC finalization."""
    try:
        engine.shutdown()
    except Exception:
        pass


def omni_snapshot_download(model_id: str) -> str:
    if os.path.exists(model_id):
        return model_id

    # TODO: this is just a workaround for quickly use modelscope, we should support
    # modelscope in weight loading feature instead of using `snapshot_download`
    if os.environ.get("VLLM_USE_MODELSCOPE", False):
        from modelscope.hub.snapshot_download import snapshot_download

        return snapshot_download(model_id)

    try:
        download_weights_from_hf_specific(
            model_name_or_path=model_id,
            cache_dir=None,
            allow_patterns=["*"],
            require_all=True,
        )
    except huggingface_hub.errors.GatedRepoError:
        raise ValueError(
            f"Access to model '{model_id}' is restricted. "
            f"Visit https://huggingface.co/{model_id} to accept "
            f"the license and request access."
        )
    except huggingface_hub.errors.RepositoryNotFoundError:
        raise ValueError(f"Repository not found for '{model_id}'. Please check the model name or path.")
    except PermissionError:
        logger.warning(
            "Permission denied when downloading '%s'. Assuming the model is already cached locally.",
            model_id,
        )

    return model_id


OutputMessageHandleResult = tuple[Literal[True], None, None, None] | tuple[Literal[False], str, int, ClientRequestState]


class OmniBase(PDDisaggregationMixin):
    """Shared runtime foundation for AsyncOmni and Omni."""

    @classmethod
    def from_cli_args(
        cls,
        args: argparse.Namespace,
        *,
        parser: argparse.ArgumentParser | None = None,
        **overrides: Any,
    ) -> OmniBase:
        """Deprecated argparse builder.

        Build from argparse. If ``parser`` is passed and not yet nullified,
        un-typed engine fields are reset to ``None``. New callers should
        nullify deploy-overriding parser defaults with
        ``nullify_stage_engine_defaults(parser)`` and construct Omni/AsyncOmni
        directly.
        """
        warnings.warn(
            "`from_cli_args()` is deprecated. Nullify deploy-overriding parser defaults "
            "with `nullify_stage_engine_defaults(parser)` and construct Omni/AsyncOmni "
            "directly from `vars(args)`.",
            DeprecationWarning,
            stacklevel=2,
        )
        kwargs: dict[str, Any] = {k: v for k, v in vars(args).items() if not k.startswith("_")}

        if parser is not None and not getattr(parser, "_omni_nullified", False):
            from vllm_omni.config.stage_config import deploy_override_field_names
            from vllm_omni.entrypoints.utils import detect_explicit_cli_keys

            explicit = detect_explicit_cli_keys(sys.argv[1:], parser) or set()
            override_dests = deploy_override_field_names()
            for key in list(kwargs):
                if key in override_dests and key not in explicit:
                    kwargs[key] = None

        kwargs.update(overrides)
        return cls(**kwargs)

    def __init__(
        self,
        model: str,
        **kwargs: Any,
    ) -> None:
        engine_args: OmniEngineArgs | None = kwargs.pop("engine_args", None)

        stage_init_timeout = kwargs.pop("stage_init_timeout", 300)
        init_timeout = kwargs.pop("init_timeout", 600)
        log_stats = kwargs.pop("log_stats", False)
        self._enable_ar_profiler = kwargs.pop("enable_ar_profiler", False)
        # NOTE: read-only lookup — must NOT pop. Popping here drops the key
        # before it reaches ``StageConfigFactory._create_from_registry``, so
        # ``--no-async-chunk`` (``async_chunk=False``) silently fails to
        # override the deploy YAML's ``async_chunk: true`` default.
        async_chunk = kwargs.get("async_chunk")
        output_modalities = kwargs.pop("output_modalities", None)
        diffusion_batch_size: int = kwargs.pop("diffusion_batch_size", 1)

        if "log_requests" in kwargs:
            raise TypeError("`log_requests` has been removed in Omni/AsyncOmni. Use `log_stats`.")
        model = omni_snapshot_download(model)
        self.__dict__["_name"] = self.__class__.__name__
        self.model = model
        self.log_stats = log_stats
        # Provisional value (mirrors the CLI/caller kwarg); the engine resolves
        # pipeline + deploy YAML + CLI precedence below and the final value is
        # re-assigned from ``self.engine.async_chunk`` after init.
        self.async_chunk = bool(async_chunk) if async_chunk is not None else False
        self.output_modalities = output_modalities or []
        self.tts_batch_max_items: int = kwargs.pop("tts_batch_max_items", 32)

        logger.info("[%s] Initializing with model %s", self.__class__.__name__, model)
        st = time.time()
        self.engine = AsyncOmniEngine(
            model=model,
            engine_args=engine_args,
            init_timeout=init_timeout,
            stage_init_timeout=stage_init_timeout,
            diffusion_batch_size=diffusion_batch_size,
            **kwargs,
        )
        self._shutdown_called = False
        self._weak_finalizer = weakref.finalize(self, _weak_shutdown_engine, self.engine)
        et = time.time()
        logger.info("[%s] AsyncOmniEngine initialized in %.2f seconds", self.__class__.__name__, et - st)
        # Authoritative: ``AsyncOmniEngine`` resolves (pipeline + deploy YAML +
        # CLI overrides) through ``StageConfigFactory`` and stores the final
        # value on ``engine.async_chunk``; mirror it here so ``--no-async-chunk``
        # (explicit ``False``) is not fallen-back-through by ``or``.
        self.async_chunk = bool(getattr(self.engine, "async_chunk", False))

        self.request_states: dict[str, ClientRequestState] = {}

        self.default_sampling_params_list = self.engine.default_sampling_params_list
        if not self.output_modalities:
            self.output_modalities = [
                self.engine.get_stage_metadata(i).final_output_type for i in range(self.engine.num_stages)
            ]

        self._stage_meta_list = [self.engine.get_stage_metadata(i) for i in range(self.engine.num_stages)]

        logger.info(
            "[%s] Initialized with %s stages for model %s",
            self.__class__.__name__,
            self.engine.num_stages,
            model,
        )

        # PD disaggregation state (detects if a prefill/decode stage pair is configured)
        self._init_pd_state()

    @property
    def num_stages(self) -> int:
        return self.engine.num_stages

    @property
    def stage_configs(self) -> list:
        """Expose engine stage configs for PD disaggregation detection and validation."""
        return self.engine.stage_configs

    def _has_dead_stage(self) -> bool:
        for stage_client in self.engine.stage_clients:
            if getattr(stage_client, "_engine_dead", False):
                return True
            resources = getattr(stage_client, "resources", None)
            if resources is not None and getattr(resources, "engine_dead", False):
                return True
        return False

    @property
    def is_running(self) -> bool:
        return self.engine.is_alive() and not self._has_dead_stage()

    @property
    def errored(self) -> bool:
        """Whether the engine is in a non-recoverable error state.

        True when the orchestrator thread is dead **or** any stage client
        has been marked dead (e.g. diffusion worker OOM / process death).

        Checks both ``_engine_dead`` (StageDiffusionClient) and
        ``resources.engine_dead`` (StageEngineCoreClient / AsyncMPClient)
        since the two client types store the flag differently.
        """
        return not self.engine.is_alive() or self._has_dead_stage()

    def check_health(self) -> None:
        if not self.engine.is_alive():
            raise EngineDeadError("Orchestrator process is not alive")
        for stage_client in self.engine.stage_clients:
            if hasattr(stage_client, "check_health"):
                stage_client.check_health()

    def resolve_sampling_params_list(
        self,
        sampling_params_list: Sequence[Any] | Any | None,
        allow_delta_coercion: bool = False,
    ) -> Sequence[Any]:
        if sampling_params_list is None:
            normalized = self.default_sampling_params_list
            # Set the output kind to delta since no params were specified
            if allow_delta_coercion:
                normalized = coerce_param_message_types(list(normalized), is_streaming=True)

        elif isinstance(sampling_params_list, Sequence) and not isinstance(sampling_params_list, (str, bytes)):
            normalized = sampling_params_list
        elif self.num_stages == 1:
            normalized = [sampling_params_list]
        else:
            raise ValueError(f"Expected {self.num_stages} sampling params, got a single sampling params object")
        if len(normalized) != self.num_stages:
            raise ValueError(f"Expected {self.num_stages} sampling params, got {len(normalized)}")
        return normalized

    def _log_summary_and_cleanup(self, request_id: str) -> None:
        req_state = self.request_states.get(request_id)
        try:
            if req_state is None or req_state.metrics is None:
                return
            if self.log_stats:
                # Emit per-request orchestrator timing (including e2e_total_ms)
                # before dropping request state.
                req_state.metrics.build_and_log_summary()
        except Exception:
            logger.exception(
                "[%s] Failed to build/log summary for req=%s",
                self.__class__.__name__,
                request_id,
            )
        finally:
            self.request_states.pop(request_id, None)

    def _compute_final_stage_id(self, output_modalities: list[str] | None) -> int:
        return get_final_stage_id_for_e2e(
            output_modalities,
            self.output_modalities,
            self._stage_meta_list,
        )

    def _process_stage_metrics_message(self, msg: StageMetricsMessage) -> None:
        req_id = msg.request_id
        req_state = self.request_states.get(req_id)
        if req_state is None or req_state.metrics is None:
            return
        _m = msg.metrics
        stage_id = msg.stage_id
        req_state.metrics.on_stage_metrics(stage_id, req_id, _m)
        submit_ts = msg.stage_submit_ts
        now = time.time()
        if req_state.metrics.stage_first_ts[stage_id] is None:
            req_state.metrics.stage_first_ts[stage_id] = submit_ts if submit_ts is not None else now
        req_state.metrics.stage_last_ts[stage_id] = max(req_state.metrics.stage_last_ts[stage_id] or 0.0, now)

    def _handle_output_message(
        self,
        msg: EngineQueueMessage | None,
    ) -> OutputMessageHandleResult:
        """Handle one Orchestrator output-queue message."""
        if msg is None:
            return True, None, None, None

        if isinstance(msg, StageMetricsMessage):
            self._process_stage_metrics_message(msg)
            return True, None, None, None

        if isinstance(msg, ErrorMessage):
            if msg.fatal:
                raise OmniEngineDeadError(
                    msg.error,
                    error_stage_id=msg.stage_id,
                )
            raise RuntimeError(msg.error)

        if not isinstance(msg, OutputMessage):
            logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg.type)
            return True, None, None, None

        req_id = msg.request_id
        stage_id = msg.stage_id

        req_state = self.request_states.get(req_id)
        if req_state is None:
            logger.debug(
                "[%s] dropping output for unknown req %s",
                self.__class__.__name__,
                req_id,
            )
            return True, None, None, None

        req_state.stage_id = stage_id

        return False, req_id, stage_id, req_state

    def _check_engine_output_error(
        self,
        result: OutputMessage,
        request_id: str,
        stage_id: int,
    ) -> None:
        """Raise if ``engine_outputs`` carries an error field.

        Raises :class:`EngineDeadError` when ``self.errored`` indicates the
        engine is unrecoverable, otherwise raises :class:`EngineGenerateError`
        (recoverable, single-request failure).
        """
        engine_outputs = result.engine_outputs
        error_text = getattr(engine_outputs, "error", None)
        if error_text is None:
            return
        logger.error(
            "[%s] Stage error for req=%s stage-%s: %s",
            self.__class__.__name__,
            request_id,
            stage_id,
            error_text,
        )
        # NOTE: O(n_stages) check for every error.
        if self.errored:
            raise OmniEngineDeadError(
                error_text,
                error_stage_id=stage_id,
            )
        raise EngineGenerateError(error_text)

    def _process_single_result(
        self,
        result: OutputMessage,
        stage_id: int,
        metrics: OrchestratorMetrics,
        req_start_ts: dict[str, float],
        wall_start_ts: float,
        final_stage_id_for_e2e: int,
    ) -> OmniRequestOutput | None:
        req_id = result.request_id
        engine_outputs = result.engine_outputs
        stage_durations = getattr(engine_outputs, "stage_durations", {})
        peak_memory_mb = getattr(engine_outputs, "peak_memory_mb", 0.0)

        # Merge AR stage timing from OrchestratorAggregator.stage_events
        if self._enable_ar_profiler:
            ar_events = metrics.stage_events.get(str(req_id), [])
            for evt in ar_events:
                if evt.stage_id != stage_id:
                    stage_durations[f"ar_stage_{evt.stage_id}"] = evt.stage_gen_time_ms / 1000.0

        # Merge pipeline timings from Orchestrator into stage_durations
        _m = result.metrics
        if _m is not None and hasattr(_m, "pipeline_timings") and _m.pipeline_timings:
            for key, value in _m.pipeline_timings.items():
                if key not in stage_durations:
                    stage_durations[key] = value

        # Merge per-stage gen times into stage_durations
        for evt in metrics.stage_events.get(str(req_id), []):
            key = f"stage_{evt.stage_id}_gen_ms"
            if key not in stage_durations:
                stage_durations[key] = evt.stage_gen_time_ms
        # Current stage gen time (not yet in stage_events at this point)
        if _m is not None:
            stage_durations.setdefault(f"stage_{stage_id}_gen_ms", _m.stage_gen_time_ms)

        finished = engine_outputs.finished

        submit_ts = result.stage_submit_ts
        now = time.time()
        if metrics.stage_first_ts[stage_id] is None:
            metrics.stage_first_ts[stage_id] = submit_ts if submit_ts is not None else now
        metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, now)

        _m = result.metrics
        if finished and _m is not None:
            metrics.on_stage_metrics(stage_id, req_id, _m)

        stage_meta = self.engine.get_stage_metadata(stage_id)
        if not stage_meta.final_output:
            return None

        try:
            rid_key = str(req_id)
            if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done and finished:
                metrics.on_finalize_request(
                    stage_id,
                    req_id,
                    req_start_ts.get(req_id, wall_start_ts),
                )
        except Exception:
            logger.exception("[%s] Finalize request handling error", self.__class__.__name__)

        output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type)
        images = getattr(engine_outputs, "images", []) if output_type == "image" else []
        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", {}),
            stage_durations=stage_durations,
            peak_memory_mb=peak_memory_mb,
        )

    def shutdown(self) -> None:
        logger.info("[%s] Shutting down", self.__class__.__name__)
        self._shutdown_base()

    def close(self) -> None:
        self.shutdown()

    def start_profile(
        self,
        profile_prefix: str | None = None,
        stages: list[int] | None = None,
    ) -> list[Any]:
        """Start profiling specified stages.

        Uses vLLM-compatible profile(is_start=True, profile_prefix) interface.

        Args:
            profile_prefix: Optional prefix for the trace file names.
            stages: List of stage IDs to profile. If None, profiles all stages.

        Returns:
            List of results from each stage.
        """
        return self.engine.collective_rpc(method="profile", args=(True, profile_prefix), stage_ids=stages)

    def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
        """Stop profiling specified stages.

        Uses vLLM-compatible profile(is_start=False) interface.

        Args:
            stages: List of stage IDs to profile. If None, stops all stages.

        Returns:
            List of results from each stage.
        """
        return self.engine.collective_rpc(method="profile", args=(False, None), stage_ids=stages)

    def _shutdown_base(self) -> None:
        if getattr(self, "_shutdown_called", False):
            return
        self._shutdown_called = True
        finalizer = getattr(self, "_weak_finalizer", None)
        if finalizer is not None and finalizer.alive:
            finalizer.detach()
        self.engine.shutdown()