# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import os import time import weakref from collections.abc import Mapping, Sequence from dataclasses import fields, is_dataclass from typing import TYPE_CHECKING, Any, Literal import huggingface_hub import vllm.envs as envs from omegaconf import OmegaConf from vllm.logger import init_logger from vllm.transformers_utils.repo_utils import file_or_path_exists from vllm.transformers_utils.runai_utils import is_runai_obj_uri from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.stage_config import merge_sampling_constraints from vllm_omni.engine.messages import ( EngineQueueMessage, ErrorMessage, OutputMessage, StageMetricsMessage, ) from vllm_omni.engine.omni_engine_base import OmniEngineBase 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.errors import raise_client_error_or from vllm_omni.metrics.modality import OmniModalityMetrics, observe_modality_at_finalize from vllm_omni.metrics.prometheus import OmniPrometheusMetrics from vllm_omni.metrics.stats import OrchestratorAggregator from vllm_omni.metrics.transfer import OmniTransferMetrics from vllm_omni.metrics.utils import ( extract_queue_wait_s, normalize_failure_reason, observe_stage_workload_metrics, ) from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.tracking_parser import TrackingNamespace if TYPE_CHECKING: from vllm_omni.engine.stage_pool import StagePool, StagePoolClient 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: OmniEngineBase) -> 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 # Object-storage models must remain URIs until each stage constructs its # ModelConfig. vLLM then materializes only config/tokenizer files locally # and keeps the URI in model_weights for Run:AI streaming. Treating the URI # as a Hugging Face repo here either fails validation or downloads through # the wrong backend before the stage processes are created. if is_runai_obj_uri(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` # Read through ``vllm.envs`` so the flag keeps vLLM's semantics (only # "true", case-insensitive, enables ModelScope). Reading the raw variable # here made every non-empty value truthy, so ``VLLM_USE_MODELSCOPE=0`` # took the ModelScope path while the rest of the stack stayed on HF. if envs.VLLM_USE_MODELSCOPE: from modelscope.hub.snapshot_download import snapshot_download return snapshot_download(model_id) # Modular Diffusers repositories describe independently loadable # components. Let the selected pipeline download only its component # sources instead of eagerly materializing the entire repository. try: if file_or_path_exists(model_id, "modular_model_index.json", revision=None): return model_id except (huggingface_hub.errors.GatedRepoError, huggingface_hub.errors.RepositoryNotFoundError): # Preserve the more helpful errors raised by the full-download path. pass 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: TrackingNamespace, model: str | None = None, ) -> OmniBase: """Build from a TrackingNamespace parsed by TrackingArgumentParser. Only args that are explicitly passed to parse_args are forwarded. """ if not isinstance(args, TrackingNamespace): raise TypeError( f"expected args to be of type TrackingNamespace, got {type(args)}. " "Hint: did you parse your args with TrackingArgumentParser?" ) explicit_kwargs = args.get_explicit_kwargs_dict() args_model = explicit_kwargs.pop("model", None) or args.model if model is not None and args_model is not None and model != args_model: raise ValueError( f"explicit model kwarg and args.model were both provided, but do not match [{model} != {args_model}]" ) if model is None and args_model is None: raise ValueError( "model must be explicitly passed as a parsed arg in the TrackingNamespace or directly provided." ) resolved_model = model or args_model return cls(model=resolved_model, **explicit_kwargs) def __init__( self, model: str, **kwargs: Any, ) -> None: if "engine_args" in kwargs: logger.warning( "engine_args were passed as a kwarg to an Omni instance; this is not supported. " "You should instead, pass the keyword arguments used to initialize the engine args " "directly to this object's initializer." ) 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_legacy_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) if "diffusion_batch_size" in kwargs: raise TypeError("`diffusion_batch_size` has been removed in Omni/AsyncOmni. Use `max_num_seqs`.") 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) # Construct transfer_metrics first so we can hand it to the engine # (which forwards it to the Orchestrator background thread for # TX-side emit; see OrchestratorBase._forward_to_next_stage). self.transfer_metrics = OmniTransferMetrics(model_name=model, log_stats=log_stats) self.prom_metrics = OmniPrometheusMetrics(model_name=model, log_stats=log_stats) st = time.time() self.engine = self._create_engine( model=model, init_timeout=init_timeout, stage_init_timeout=stage_init_timeout, transfer_emitter=self.transfer_metrics, prom_metrics=self.prom_metrics, log_stats=log_stats, **kwargs, ) self._shutdown_called = False self._weak_finalizer = weakref.finalize(self, _weak_shutdown_engine, self.engine) et = time.time() logger.info("[%s] %s initialized in %.2f seconds", self.__class__.__name__, type(self.engine).__name__, et - st) # Authoritative: the engine 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.mod_metrics = OmniModalityMetrics(model_name=model, log_stats=log_stats) self.default_sampling_params_list = self.engine.default_sampling_params_list self.sampling_constraints_list = self._get_sampling_constraints_list(self.engine.stage_configs) 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() def _create_engine(self, **engine_kwargs: Any) -> OmniEngineBase: """Construct this entrypoint's engine (``AsyncOmniEngine`` for turn-based use).""" raise NotImplementedError @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 @staticmethod def _replica_is_dead(client: StagePoolClient | None) -> bool: """Whether a replica slot is evicted (``None``) or its engine is dead. ``check_health()`` is the source of truth: it raises ``EngineDeadError`` when dead (inspecting the ``engine_dead`` flags internally), and the diffusion client additionally runs a synchronous ``proc.is_alive()`` probe that catches a silent SIGKILL/segfault the monitor thread has not flagged yet — so ``/health`` does not report 200 for a stage whose subprocess is already dead. A client exposing no ``check_health`` cannot confirm liveness, so it is treated as dead (fail closed). """ if client is None: return True check_health = getattr(client, "check_health", None) if not callable(check_health): return True try: check_health() except EngineDeadError: return True return False def _live_replica_count(self, pool: StagePool) -> int: """Number of replicas in ``pool`` that are neither evicted nor dead. A dead replica is evicted from its pool (slot set to ``None`` in ``StagePool.clients``); a replica that died but has not been evicted yet still carries an ``engine_dead`` flag. Both are excluded. """ return sum(1 for client in pool.clients if not self._replica_is_dead(client)) def _stage_has_no_live_replica(self, pool: StagePool) -> bool: """True when a non-empty stage pool has lost all of its replicas.""" return len(pool.clients) > 0 and self._live_replica_count(pool) == 0 @property def is_running(self) -> bool: return self.engine.is_alive() @property def errored(self) -> bool: """Whether the engine is in a process-fatal error state. True only when the orchestrator thread is dead. Per-stage liveness is deliberately excluded: the OpenAI serving paths precheck ``errored`` before request routing, so including it would reject requests that do not touch the dead stage (e.g. text-only chat when the talker stage is down). A fully dead stage instead fails only the requests routed through it (dispatch guards) and flips readiness to 503 via :meth:`check_health` (per-replica fault isolation, #4285). """ return not self.engine.is_alive() def check_health(self) -> None: if not self.engine.is_alive(): raise EngineDeadError("Orchestrator process is not alive") pools = getattr(self.engine, "stage_pools", None) if pools is None: return for pool in pools: if self._stage_has_no_live_replica(pool): raise EngineDeadError(f"Stage-{pool.stage_id} has no live replica") def resolve_sampling_params_list( self, sampling_params_list: Sequence[Any] | Any | None, allow_delta_coercion: bool = False, ) -> Sequence[Any]: """Resolve request parameters; pipeline sampling constraints override caller values.""" normalized: 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)}") # Streaming coercion may also change a constrained output kind in # the defaults (for example, an internal TTS stage's FINAL_ONLY). if sampling_params_list is not None or allow_delta_coercion: normalized = [ self._apply_sampling_constraints(params, constraints) for params, constraints in zip(normalized, self.sampling_constraints_list, strict=True) ] return normalized @staticmethod def _get_sampling_constraints_list(stage_configs: Sequence[Any]) -> list[dict[str, Any]]: """Extract each stage's required sampling settings from runtime configs.""" constraints_list: list[dict[str, Any]] = [] for stage_config in stage_configs: constraints = getattr(stage_config, "sampling_constraints", {}) if not isinstance(constraints, Mapping): constraints_list.append({}) continue if OmegaConf.is_config(constraints): constraints = OmegaConf.to_container(constraints, resolve=True) constraints_list.append(dict(constraints)) return constraints_list @staticmethod def _apply_sampling_constraints(params: Any, constraints: Mapping[str, Any]) -> Any: """Apply pipeline requirements, merging required stops with caller stops.""" if not constraints: return params if isinstance(params, Mapping): values = dict(params) elif is_dataclass(params): values = {field.name: getattr(params, field.name) for field in fields(params) if field.init} elif struct_fields := getattr(params, "__struct_fields__", None): values = { name: getattr(params, name) for name in struct_fields if not name.startswith("_") and name != "output_text_buffer_length" } else: raise TypeError(f"Expected a mapping, dataclass, or msgspec struct, got {type(params).__name__}") resolved = merge_sampling_constraints(values, constraints) if isinstance(params, Mapping): return resolved return type(params)(**resolved) def _record_request_failure_once(self, request_id: str, reason: str) -> None: req_state = self.request_states.get(request_id) prom = getattr(self, "prom_metrics", None) metrics = getattr(req_state, "metrics", None) if req_state is None or metrics is None or prom is None: return if str(request_id) in metrics.e2e_done or getattr(req_state, "failure_recorded", False): return req_state.failure_recorded = True prom.request_failed() prom.inc_requests_failed(normalize_failure_reason(reason)) def _publish_request_gauges(self, total: int) -> None: """Publish num_requests_running / num_requests_waiting for ``total`` in-flight requests. ``_running_counter`` counts requests from the moment the orchestrator dispatches them to a stage engine, which includes requests still queued inside a stage scheduler. The orchestrator reports those via ``_engines_waiting_counter``. """ prom = getattr(self, "prom_metrics", None) if prom is None: return engine = getattr(self, "engine", None) counter = getattr(engine, "_running_counter", None) waiting_counter = getattr(engine, "_engines_waiting_counter", None) dispatched = counter.value if counter is not None else total in_engine_waiting = waiting_counter.value if waiting_counter is not None else 0 in_engine_waiting = min(max(0, in_engine_waiting), dispatched) prom.set_running(max(0, dispatched - in_engine_waiting)) prom.set_waiting(max(0, total - dispatched) + in_engine_waiting) def _log_summary_and_cleanup(self, request_id: str, reason: str = "stage_error") -> None: req_state = self.request_states.get(request_id) try: if req_state is None or req_state.metrics is None: return if str(request_id) not in req_state.metrics.e2e_done: self._record_request_failure_once(request_id, reason) 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) # Republish gauges so any stale value left by the per-stage # publish in _process_single_result (which runs while the request # is still in self.request_states) is corrected after the pop. self._publish_request_gauges(len(self.request_states)) 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 _compute_final_output_stage_ids(self, output_modalities: list[str] | None) -> list[int]: requested_modalities = output_modalities or self.output_modalities requested_modalities = [m for m in requested_modalities if m in self.output_modalities] if not requested_modalities: requested_modalities = self.output_modalities return [ sid for sid, stage in enumerate(self._stage_meta_list) if getattr(stage, "final_output", False) and stage.final_output_type in requested_modalities ] 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 stage_meta = self.engine.get_stage_metadata(stage_id) req_state.metrics.on_stage_metrics(stage_id, req_id, _m, stage_meta.final_output_type) 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, ) self._raise_nonfatal_error_message(msg) 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 if msg.metrics is not None and not msg.finished and req_state.metrics is not None: stage_meta = self.engine.get_stage_metadata(stage_id) output_type = getattr(msg.engine_outputs, "final_output_type", stage_meta.final_output_type) msg_id = id(msg) consumed = req_state.consumed_metric_message_ids if msg_id not in consumed: req_state.metrics.on_stage_metrics(stage_id, req_id, msg.metrics, output_type) 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) consumed.add(msg_id) return False, req_id, stage_id, req_state def _raise_nonfatal_error_message(self, msg: ErrorMessage) -> None: """Raise the exception for a non-fatal, request-scoped error message.""" raise_client_error_or( msg.error, status_code=msg.status_code, error_type=msg.error_type, fallback=RuntimeError, ) 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 status_code = engine_outputs.error_status_code error_type = engine_outputs.error_type 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_client_error_or( error_text, status_code=status_code, error_type=error_type, fallback=EngineGenerateError, ) def _process_single_result( self, result: OutputMessage, stage_id: int, metrics: OrchestratorAggregator, req_start_ts: dict[str, float], wall_start_ts: float, final_stage_id_for_e2e: int, stage_event_cursor: int = 0, ) -> 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), [])[stage_event_cursor:] 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), [])[stage_event_cursor:]: 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 stage_meta = self.engine.get_stage_metadata(stage_id) output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type) if finished and _m is not None: msg_id = id(result) req_state = self.request_states.get(req_id) consumed = req_state.consumed_metric_message_ids if req_state is not None else None if consumed is not None and msg_id not in consumed: metrics.accumulate_diffusion_metrics(stage_meta.stage_type, req_id, engine_outputs) metrics.on_stage_metrics(stage_id, req_id, _m, output_type) consumed.add(msg_id) if peak_memory_mb > 0: self.prom_metrics.set_peak_memory(stage_id, peak_memory_mb) self.prom_metrics.observe_stage_gen_time( stage_id, stage_meta.stage_type, _m.stage_gen_time_ms / 1000.0, ) observe_stage_workload_metrics( self.prom_metrics, stage_type=stage_meta.stage_type, stage_metrics=_m, ) # Only meaningful for diffusion stages (text stages have no exec key). _diff_m = _m.diffusion_metrics or {} stage_in_queue_s = _diff_m.get("scheduler_queue_wait_s") if stage_in_queue_s is not None: self.prom_metrics.observe_stage_in_queue(stage_id, float(stage_in_queue_s)) if _m.output_unit_type == "image": if result.replica_id is not None: self.mod_metrics.observe_image_ttfp( str(stage_id), str(result.replica_id), _m.image_time_to_first_output_ms / 1000.0, ) if not stage_meta.final_output: return None output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type) 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), ) e2e_seconds = now - req_start_ts.get(req_id, wall_start_ts) # Extract finished_reason from upstream CompletionOutput so # the per-reason completion Counter is labelled correctly. completion_outputs = getattr(engine_outputs, "outputs", None) or [] fr = (getattr(completion_outputs[0], "finish_reason", None) if completion_outputs else None) or "stop" self.prom_metrics.request_succeeded( e2e_seconds, finished_reason=fr, ) # Token counters — aggregate across all stages for this request. _prompt_tok = 0 _gen_tok = 0 for evt in metrics.stage_events.get(rid_key, [])[stage_event_cursor:]: if evt.stage_id == 0: _prompt_tok += int(evt.num_tokens_in) _gen_tok += int(evt.num_tokens_out) self.prom_metrics.observe_tokens(_prompt_tok, _gen_tok) queue_wait_s = extract_queue_wait_s(_m.pipeline_timings if _m is not None else None) if queue_wait_s is not None: self.prom_metrics.observe_queue_wait(queue_wait_s) # Modality observe inside the same finalize guard so it fires # once per request and inherits the try/except isolation. observe_modality_at_finalize( self.mod_metrics, output_type=output_type, stage_id=stage_id, replica_id=result.replica_id, stage_metrics=_m, engine_outputs=engine_outputs, ) except Exception: logger.exception("[%s] Finalize request handling error", self.__class__.__name__) # When this result finalizes the request, the orchestrator has # already decremented _running_counter but _log_summary_and_cleanup # hasn't popped self.request_states yet — exclude the finalizing # request from `total` so waiting doesn't read 1 and stay stuck # there until the next request arrives. is_finalizing = finished and stage_id == final_stage_id_for_e2e self._publish_request_gauges(max(0, len(self.request_states) - (1 if is_finalizing else 0))) response_metrics: dict[str, Any] = {} stage_metrics: dict[str, dict[str, Any]] = {} rid_key = str(req_id) for evt in metrics.stage_events.get(rid_key, [])[stage_event_cursor:]: if evt.stage_id is None: continue sid = int(evt.stage_id) evt_stage_meta = self.engine.get_stage_metadata(sid) evt_output_type = evt.final_output_type if not evt_output_type: evt_output_type = evt_stage_meta.final_output_type stage_name = evt_stage_meta.model_stage sid_key = str(sid) stage_metrics[sid_key] = OrchestratorAggregator._merge_stage_metric_event(stage_metrics.get(sid_key), evt) stage_metrics[sid_key]["stage_name"] = stage_name or f"stage_{sid}" stage_metrics[sid_key]["final_output_type"] = evt_output_type if stage_metrics: response_metrics["stage_metrics"] = stage_metrics current_stage_metrics = stage_metrics.get(str(stage_id)) if current_stage_metrics is not None: response_metrics["stage_id"] = current_stage_metrics["stage_id"] 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.from_stage_output(). return OmniRequestOutput.from_stage_output( engine_outputs, request_id=req_id or "", finished=finished, stage_id=stage_id, replica_id=result.replica_id, final_output_type=output_type, metrics=response_metrics, stage_durations=stage_durations, peak_memory_mb=peak_memory_mb, ) def shutdown(self, timeout: float | None = None) -> 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()