1459 lines
64 KiB
Python
1459 lines
64 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
import json
|
||
import multiprocessing as mp
|
||
import os
|
||
import threading
|
||
import time
|
||
import uuid
|
||
import weakref
|
||
from collections.abc import Callable, Generator, Sequence
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from typing import Any, Literal, TypeVar, overload
|
||
|
||
import huggingface_hub
|
||
import msgspec.msgpack
|
||
import torch
|
||
import zmq
|
||
from tqdm.auto import tqdm
|
||
from vllm import SamplingParams
|
||
from vllm.logger import init_logger
|
||
from vllm.utils.network_utils import make_zmq_socket
|
||
from vllm.v1.utils import get_engine_client_zmq_addr
|
||
|
||
from vllm_omni.config.stage_config import StageConfigFactory
|
||
from vllm_omni.config.yaml_util import create_config
|
||
from vllm_omni.distributed.omni_connectors import (
|
||
get_stage_connector_config,
|
||
initialize_orchestrator_connectors,
|
||
)
|
||
from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector
|
||
from vllm_omni.distributed.omni_connectors.utils.initialization import (
|
||
resolve_omni_kv_config_for_stage,
|
||
)
|
||
from vllm_omni.distributed.ray_utils.utils import (
|
||
create_placement_group,
|
||
get_ray_queue_class,
|
||
try_close_ray,
|
||
)
|
||
from vllm_omni.entrypoints.cfg_companion_tracker import CfgCompanionTracker
|
||
from vllm_omni.entrypoints.omni_stage import OmniStage
|
||
from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
|
||
from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
|
||
from vllm_omni.entrypoints.utils import (
|
||
filter_dataclass_kwargs,
|
||
get_final_stage_id_for_e2e,
|
||
inject_omni_kv_config,
|
||
load_and_resolve_stage_configs,
|
||
)
|
||
from vllm_omni.entrypoints.zmq_utils import ZmqQueue
|
||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams
|
||
from vllm_omni.lora.request import LoRARequest
|
||
from vllm_omni.metrics import OrchestratorAggregator, StageRequestStats
|
||
from vllm_omni.model_executor.model_loader.weight_utils import (
|
||
download_weights_from_hf_specific,
|
||
)
|
||
from vllm_omni.outputs import OmniRequestOutput
|
||
|
||
_R = TypeVar("_R")
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
def _weak_close_cleanup(
|
||
stage_list,
|
||
stage_in_queues,
|
||
stage_out_queues,
|
||
ray_pg,
|
||
zmq_ctx=None,
|
||
handshake_stop: threading.Event | None = None,
|
||
zmq_handshake_socket: zmq.Socket | None = None,
|
||
handshake_thread: threading.Thread | None = None,
|
||
):
|
||
"""Weak reference cleanup function for OmniBase instances."""
|
||
if stage_list:
|
||
for q in stage_in_queues:
|
||
try:
|
||
q.put_nowait(SHUTDOWN_TASK)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to send shutdown signal to stage input queue: {e}")
|
||
close_fn = getattr(q, "close", None)
|
||
if callable(close_fn):
|
||
close_fn()
|
||
for q in stage_out_queues:
|
||
close_fn = getattr(q, "close", None)
|
||
if callable(close_fn):
|
||
close_fn()
|
||
for stage in stage_list:
|
||
try:
|
||
stage.stop_stage_worker()
|
||
except Exception as e:
|
||
logger.warning(f"Failed to stop stage worker: {e}")
|
||
try_close_ray(ray_pg)
|
||
|
||
# Gracefully shutdown handshake server thread
|
||
if handshake_stop is not None:
|
||
handshake_stop.set()
|
||
if handshake_thread is not None:
|
||
handshake_thread.join(timeout=2.0)
|
||
if handshake_thread.is_alive():
|
||
logger.warning("Handshake server thread did not terminate gracefully within timeout")
|
||
|
||
# Close ZMQ resources after thread has exited
|
||
if zmq_handshake_socket is not None:
|
||
zmq_handshake_socket.close(0)
|
||
if zmq_ctx is not None:
|
||
zmq_ctx.term()
|
||
|
||
|
||
def _dummy_snapshot_download(model_id):
|
||
return model_id
|
||
|
||
|
||
def omni_snapshot_download(model_id) -> str:
|
||
# If it's already a local path, just return it
|
||
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)
|
||
# For other cases (Hugging Face), perform a real download to ensure all
|
||
# necessary files (including *.pt for audio/diffusion) are available locally
|
||
# before stage workers are spawned. This prevents initialization timeouts.
|
||
# Return the original model_id so that model_config.model preserves
|
||
# HuggingFace semantics (e.g. "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice")
|
||
# instead of the resolved cache path.
|
||
try:
|
||
download_weights_from_hf_specific(
|
||
model_name_or_path=model_id,
|
||
cache_dir=None,
|
||
allow_patterns=["*"],
|
||
require_all=True,
|
||
)
|
||
except huggingface_hub.errors.RepositoryNotFoundError:
|
||
logger.warning(f"Repository not found for '{model_id}'.")
|
||
return model_id
|
||
|
||
|
||
class OmniBase:
|
||
"""Base class for serving Omni models.
|
||
|
||
Args:
|
||
model: Model name or path to load.
|
||
**kwargs: Arbitrary keyword arguments.
|
||
- stage_configs_path: Optional path to YAML file containing stage
|
||
configurations. If None, configurations are loaded from the model.
|
||
- log_stats: Whether to enable statistics logging
|
||
be written to files with stage-specific suffixes.
|
||
- stage_init_timeout: Per-stage init watchdog (seconds). Measured from
|
||
when the previous stage finished (possibly a prior Omni run with GPU
|
||
reuse/overlap) to when the current stage starts to initialize.
|
||
- shm_threshold_bytes: Threshold in bytes for using shared memory
|
||
for IPC. Objects larger than this threshold will use shared memory.
|
||
- worker_backend: Backend for worker processes. Default is "multi_process".
|
||
- ray_address: Address of Ray cluster for Ray backend, if using Ray backend.
|
||
- batch_timeout: Timeout in seconds for batching requests within a stage
|
||
- init_timeout: Timeout in seconds for waiting for all stages to initialize
|
||
- Additional keyword arguments passed to stage engines.
|
||
"""
|
||
|
||
def __init__(self, model: str, **kwargs: Any) -> None:
|
||
model = omni_snapshot_download(model)
|
||
kwargs["model"] = model
|
||
|
||
# Stage management attributes
|
||
self.stage_list: list[OmniStage] = []
|
||
self._stage_in_queues: list[Any] = []
|
||
self._stage_out_queues: list[Any] = []
|
||
self._stages_ready: set[int] = set()
|
||
self._ray_pg = None
|
||
self._queue_cls = None
|
||
self._ctx = None
|
||
self._zmq_ctx: zmq.Context | None = None
|
||
self._zmq_master_address: str | None = None
|
||
self._zmq_master_port: int | None = None
|
||
self._zmq_handshake_socket: zmq.Socket | None = None
|
||
self._handshake_thread: threading.Thread | None = None
|
||
self._handshake_stop: threading.Event | None = None
|
||
self._handshake_endpoints: dict[int, tuple[str, str]] = {}
|
||
self._handshake_seen: set[int] = set() # Track which stage IDs have completed ZMQ handshake
|
||
self._single_stage_id: int | None = None # Optional: deploy only a specific stage ID
|
||
|
||
# Sleep mode tracking
|
||
self._is_sleeping: bool = False
|
||
|
||
# RPC results storage: {stage_id: {rpc_id: result}}
|
||
# Used by collective_rpc to retrieve results collected from the output queue
|
||
self._rpc_results: dict[int, dict[str, dict[str, Any]]] = {}
|
||
|
||
# Initialize stages - each stage will create appropriate instance based on stage_type
|
||
# Stage workers will automatically create OmniLLM or OmniDiffusion instances
|
||
# based on stage_type in YAML config (handled in omni_stage.py)
|
||
logger.info(f"Initializing stages for model: {model}")
|
||
self._initialize_stages(model, kwargs)
|
||
|
||
def _get_default_cache_config(self, cache_backend: str | None) -> dict[str, Any] | None:
|
||
if cache_backend == "cache_dit":
|
||
return {
|
||
"Fn_compute_blocks": 1,
|
||
"Bn_compute_blocks": 0,
|
||
"max_warmup_steps": 4,
|
||
"residual_diff_threshold": 0.24,
|
||
"max_continuous_cached_steps": 3,
|
||
"enable_taylorseer": False,
|
||
"taylorseer_order": 1,
|
||
"scm_steps_mask_policy": None,
|
||
"scm_steps_policy": "dynamic",
|
||
}
|
||
if cache_backend == "tea_cache":
|
||
return {
|
||
"rel_l1_thresh": 0.2,
|
||
}
|
||
return None
|
||
|
||
def _normalize_cache_config(self, cache_backend: str | None, cache_config: Any | None) -> Any | None:
|
||
if isinstance(cache_config, str):
|
||
try:
|
||
cache_config = json.loads(cache_config)
|
||
except json.JSONDecodeError:
|
||
logger.warning("Invalid cache_config JSON, using defaults.")
|
||
cache_config = None
|
||
if cache_config is None and cache_backend not in (None, "", "none"):
|
||
cache_config = self._get_default_cache_config(cache_backend)
|
||
return cache_config
|
||
|
||
def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
|
||
"""Create default diffusion stage configuration.
|
||
|
||
Uses StageConfigFactory for typed configuration creation while
|
||
maintaining backward compatibility with the legacy format.
|
||
|
||
Args:
|
||
kwargs: Engine arguments from CLI/API.
|
||
|
||
Returns:
|
||
List containing a single OmegaConf config for the diffusion stage.
|
||
"""
|
||
# Normalize dtype
|
||
if "dtype" in kwargs and not isinstance(kwargs["dtype"], str):
|
||
if not isinstance(kwargs["dtype"], torch.dtype):
|
||
raise TypeError(f"Provided dtype must be a string or torch.dtype, got {type(kwargs['dtype']).__name__}")
|
||
kwargs["dtype"] = str(kwargs["dtype"]).removeprefix("torch.")
|
||
|
||
# Normalize cache config before passing to factory
|
||
cache_backend = kwargs.get("cache_backend", "none")
|
||
cache_config = self._normalize_cache_config(cache_backend, kwargs.get("cache_config", None))
|
||
|
||
# Update kwargs with normalized values
|
||
kwargs_copy = dict(kwargs)
|
||
kwargs_copy["cache_backend"] = cache_backend
|
||
kwargs_copy["cache_config"] = cache_config
|
||
|
||
# Use the factory to create default diffusion config
|
||
return StageConfigFactory.create_default_diffusion(kwargs_copy)
|
||
|
||
def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[str, list[Any]]:
|
||
"""Resolve stage configs and inject defaults shared by orchestrator/headless."""
|
||
# TODO(wuhang):
|
||
# Remove kwargs as parameters in the future.
|
||
# Use dataclass directly for engine args.
|
||
|
||
stage_configs_path = kwargs.get("stage_configs_path", None)
|
||
|
||
# TTS-specific CLI overrides
|
||
self.tts_max_instructions_length: int | None = kwargs.get("tts_max_instructions_length", None)
|
||
|
||
# Load stage configurations from YAML
|
||
config_path, stage_configs = load_and_resolve_stage_configs(
|
||
model,
|
||
stage_configs_path,
|
||
kwargs,
|
||
default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs),
|
||
)
|
||
|
||
# Inject diffusion LoRA-related knobs from kwargs if not present in the stage config.
|
||
for cfg in stage_configs:
|
||
try:
|
||
if getattr(cfg, "stage_type", None) != "diffusion":
|
||
continue
|
||
if not hasattr(cfg, "engine_args") or cfg.engine_args is None:
|
||
cfg.engine_args = create_config({})
|
||
if kwargs.get("lora_path") is not None:
|
||
if not hasattr(cfg.engine_args, "lora_path") or cfg.engine_args.lora_path is None:
|
||
cfg.engine_args.lora_path = kwargs["lora_path"]
|
||
lora_scale = kwargs.get("lora_scale")
|
||
if lora_scale is None:
|
||
# Backwards compatibility for older callers.
|
||
lora_scale = kwargs.get("static_lora_scale")
|
||
if lora_scale is not None:
|
||
if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None:
|
||
cfg.engine_args.lora_scale = lora_scale
|
||
quantization_config = kwargs.get("quantization_config")
|
||
if quantization_config is not None:
|
||
if (
|
||
not hasattr(cfg.engine_args, "quantization_config")
|
||
or cfg.engine_args.quantization_config is None
|
||
):
|
||
cfg.engine_args.quantization_config = quantization_config
|
||
except Exception as e:
|
||
logger.warning("Failed to inject LoRA config for stage: %s", e)
|
||
|
||
return config_path, stage_configs
|
||
|
||
def _initialize_stages(self, model: str, kwargs: dict[str, Any]) -> None:
|
||
"""Initialize stage list management."""
|
||
self._inline_diffusion = False
|
||
self._inline_engine = None
|
||
|
||
stage_init_timeout = kwargs.get("stage_init_timeout", 20)
|
||
shm_threshold_bytes = kwargs.get("shm_threshold_bytes", 65536)
|
||
init_timeout = kwargs.get("init_timeout", 300)
|
||
worker_backend = kwargs.get("worker_backend", "multi_process")
|
||
ray_address = kwargs.get("ray_address", None)
|
||
batch_timeout = kwargs.get("batch_timeout", 10)
|
||
log_stats = kwargs.get("log_stats", False)
|
||
self._single_stage_id = kwargs.get("stage_id", None)
|
||
self._zmq_master_address = kwargs.get("omni_master_address", None)
|
||
if self._zmq_master_address is None:
|
||
self._zmq_master_address = "127.0.0.1"
|
||
logger.info("No omni_master_address provided, defaulting to localhost (127.0.0.1)")
|
||
self._zmq_master_port = kwargs.get("omni_master_port", None)
|
||
|
||
# Resolve stage configs shared by orchestrator/headless paths.
|
||
self.config_path, self.stage_configs = self._resolve_stage_configs(model, kwargs)
|
||
|
||
# Initialize connectors
|
||
self.omni_transfer_config, self.connectors = initialize_orchestrator_connectors(
|
||
self.config_path, worker_backend=worker_backend, shm_threshold_bytes=shm_threshold_bytes
|
||
)
|
||
|
||
# Initialize stats paths
|
||
self.log_stats: bool = bool(log_stats)
|
||
|
||
self.worker_backend = worker_backend
|
||
self.ray_address = ray_address
|
||
self.batch_timeout = batch_timeout
|
||
# async chunk remains the same for each stage
|
||
self.async_chunk = self._is_async_chunk_enable(self.stage_configs)
|
||
|
||
# Build OmniStage instances in parallel, preserve original order
|
||
def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
|
||
idx, cfg = idx_cfg
|
||
return idx, OmniStage(cfg, stage_init_timeout=stage_init_timeout)
|
||
|
||
with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor:
|
||
futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)]
|
||
results: list[tuple[int, OmniStage]] = []
|
||
for fut in as_completed(futures):
|
||
results.append(fut.result())
|
||
results.sort(key=lambda x: x[0])
|
||
self.stage_list = [st for _, st in results]
|
||
self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list]
|
||
self.output_modalities = [st.final_output_type for st in self.stage_list]
|
||
logger.info(f"[{self._name}] Loaded {len(self.stage_list)} stages")
|
||
|
||
# Phase 1 optimization: for a single diffusion stage in async mode,
|
||
# run the engine directly in the orchestrator process to eliminate
|
||
# the stage worker subprocess and its IPC serialization overhead.
|
||
if len(self.stage_list) == 1 and self.stage_list[0].stage_type == "diffusion" and self.is_async:
|
||
self._init_inline_diffusion_engine(model, self.stage_configs[0], kwargs)
|
||
return
|
||
|
||
if self.worker_backend == "ray":
|
||
self._queue_cls = get_ray_queue_class()
|
||
else:
|
||
self._ctx = mp.get_context("spawn")
|
||
self._queue_cls = lambda: self._ctx.Queue(maxsize=0)
|
||
|
||
self._stage_init_timeout = max(0, int(stage_init_timeout))
|
||
self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
|
||
self._start_stages(model)
|
||
# Wait for all stages to report readiness before seeding
|
||
self._wait_for_stages_ready(timeout=init_timeout)
|
||
# Set up RPC result checkers so that collective_rpc works
|
||
self._setup_rpc_result_checkers()
|
||
|
||
def _init_inline_diffusion_engine(
|
||
self,
|
||
model: str,
|
||
stage_config: Any,
|
||
kwargs: dict[str, Any],
|
||
) -> None:
|
||
"""Initialize diffusion engine directly in the orchestrator process.
|
||
|
||
For single-stage diffusion pipelines, this eliminates the stage worker
|
||
subprocess and the associated Hop3 IPC serialization overhead.
|
||
GPU workers for tensor parallelism are still spawned by the
|
||
DiffusionExecutor as separate processes.
|
||
"""
|
||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||
from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
|
||
from vllm_omni.entrypoints.stage_utils import (
|
||
_to_dict,
|
||
load_func_from_config,
|
||
set_stage_devices,
|
||
)
|
||
|
||
stage_id = stage_config.stage_id
|
||
engine_args = _to_dict(stage_config.engine_args)
|
||
runtime_cfg = _to_dict(getattr(stage_config, "runtime", {}))
|
||
|
||
if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn":
|
||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||
|
||
try:
|
||
from vllm_omni.platforms import current_omni_platform
|
||
|
||
device_type = current_omni_platform.device_type
|
||
set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
|
||
except Exception as e:
|
||
logger.warning("Device setup for inline diffusion failed: %s", e)
|
||
|
||
engine_args = filter_dataclass_kwargs(OmniDiffusionConfig, engine_args)
|
||
engine_args.pop("model_stage", None)
|
||
engine_args.pop("model", None)
|
||
|
||
cfg_kv_collect_func = load_func_from_config(getattr(stage_config, "cfg_kv_collect_func", None))
|
||
|
||
self._inline_engine = OmniDiffusion(
|
||
model=model,
|
||
stage_id=stage_id,
|
||
engine_input_source=getattr(stage_config, "engine_input_source", []),
|
||
cfg_kv_collect_func=cfg_kv_collect_func,
|
||
**engine_args,
|
||
)
|
||
self._inline_diffusion = True
|
||
|
||
# These attributes are normally set by AsyncOmni._wait_for_stages_ready
|
||
# but we skip that for inline mode. Set them to None since there is no
|
||
# LLM stage to provide them.
|
||
self.input_processor = None
|
||
self.io_processor = None
|
||
self.model_config = None
|
||
|
||
logger.info(
|
||
"[%s] Inline diffusion mode active – stage worker subprocess bypassed",
|
||
self._name,
|
||
)
|
||
|
||
def _is_async_chunk_enable(self, stage_args: list) -> bool:
|
||
"""get async chunk flag"""
|
||
engine_args = getattr(stage_args[0], "engine_args", None)
|
||
return bool(getattr(engine_args, "async_chunk", False))
|
||
|
||
def _start_stages(self, model: str) -> None:
|
||
"""Start all stage processes."""
|
||
if self.worker_backend == "ray":
|
||
# Initialize Ray Cluster
|
||
self._ray_pg = create_placement_group(
|
||
number_of_stages=len(self.stage_list), address=self.ray_address, strategy="PACK"
|
||
)
|
||
else:
|
||
# Initialize ZMQ context
|
||
if self._zmq_ctx is None:
|
||
self._zmq_ctx = zmq.Context()
|
||
|
||
# Allocate endpoints for each stage
|
||
total_stages = len(self.stage_configs)
|
||
self._handshake_endpoints = {}
|
||
|
||
# If --stage-id is not set, use local_only mode
|
||
local_only = self._single_stage_id is None
|
||
|
||
for sid in range(total_stages):
|
||
in_endpoint = get_engine_client_zmq_addr(local_only=local_only, host=self._zmq_master_address)
|
||
out_endpoint = get_engine_client_zmq_addr(local_only=local_only, host=self._zmq_master_address)
|
||
self._handshake_endpoints[sid] = (in_endpoint, out_endpoint)
|
||
logger.debug(
|
||
f"[{self._name}] Allocated endpoints for stage-{sid}: in={in_endpoint}, out={out_endpoint}"
|
||
)
|
||
|
||
# Start handshake server
|
||
self.start_handshake_server()
|
||
|
||
for stage_id, stage in enumerate[OmniStage](self.stage_list):
|
||
if self.worker_backend == "ray":
|
||
in_q = self._queue_cls()
|
||
out_q = self._queue_cls()
|
||
else:
|
||
in_endpoint, out_endpoint = self._handshake_endpoints[stage_id]
|
||
in_q = ZmqQueue(self._zmq_ctx, zmq.PUSH, bind=in_endpoint)
|
||
out_q = ZmqQueue(self._zmq_ctx, zmq.PULL, bind=out_endpoint)
|
||
|
||
self._stage_in_queues.append(in_q)
|
||
self._stage_out_queues.append(out_q)
|
||
stage.attach_queues(in_q, out_q)
|
||
|
||
stage_connectors_config = get_stage_connector_config(
|
||
self.omni_transfer_config,
|
||
stage_id,
|
||
)
|
||
|
||
# Inject YAML-resolved connector config into omni_kv_config for
|
||
# in-engine usage (GPU model runner reads model_config.omni_kv_config).
|
||
try:
|
||
omni_conn_cfg, omni_from, omni_to = resolve_omni_kv_config_for_stage(
|
||
self.omni_transfer_config, stage_id
|
||
)
|
||
if omni_conn_cfg:
|
||
inject_omni_kv_config(stage, omni_conn_cfg, omni_from, omni_to) # type: ignore
|
||
|
||
except Exception as e:
|
||
logger.debug("[Omni] Failed to inject omni connector config into stage-%s: %s", stage_id, e)
|
||
|
||
if self._single_stage_id is not None and stage_id != int(self._single_stage_id):
|
||
logger.info(
|
||
f"[{self._name}] Skipping initialization of stage-{stage_id} worker due to single_stage_id setting"
|
||
)
|
||
continue
|
||
|
||
stage.init_stage_worker(
|
||
model,
|
||
is_async=self.is_async,
|
||
shm_threshold_bytes=self._shm_threshold_bytes,
|
||
ctx=self._ctx if self.worker_backend != "ray" else None,
|
||
batch_timeout=self.batch_timeout,
|
||
connectors_config=stage_connectors_config,
|
||
worker_backend=self.worker_backend,
|
||
ray_placement_group=self._ray_pg,
|
||
ignore_runtime_config=True if self._single_stage_id is not None else False,
|
||
)
|
||
|
||
logger.debug(f"[{self._name}] Stage-{stage_id} process started")
|
||
|
||
def _process_stage_ready(self, stage: OmniStage, stage_id: int, result: dict[str, Any]) -> None:
|
||
self._stages_ready.add(stage_id)
|
||
logger.info(f"[{self._name}] Stage-{stage_id} reported ready")
|
||
|
||
def _wait_for_stages_ready(self, timeout: int = 120) -> None:
|
||
"""Wait for all stages to report readiness with optimized polling."""
|
||
if self._single_stage_id is not None and self.worker_backend != "ray":
|
||
timeout = self._wait_for_handshakes(timeout)
|
||
|
||
num_stages = len(self.stage_list)
|
||
deadline = time.time() + max(0, int(timeout))
|
||
|
||
logger.info(f"[{self._name}] Waiting for {num_stages} stages to initialize (timeout: {timeout}s)")
|
||
|
||
while len(self._stages_ready) < num_stages and time.time() < deadline:
|
||
progressed = False
|
||
for stage_id, stage in enumerate(self.stage_list):
|
||
if stage_id in self._stages_ready:
|
||
continue
|
||
|
||
# Check if the stage has reported status
|
||
if result := stage.try_collect():
|
||
progressed = True
|
||
if result.get("type") == "stage_ready":
|
||
self._process_stage_ready(stage, stage_id, result)
|
||
|
||
if not progressed:
|
||
time.sleep(0.05)
|
||
|
||
# Handle Final State
|
||
if len(self._stages_ready) == num_stages:
|
||
logger.info(f"[{self._name}] All stages initialized successfully")
|
||
return
|
||
|
||
# Handle Timeout/Failure
|
||
not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
|
||
logger.warning(
|
||
f"[{self._name}] Initialization timeout: {len(self._stages_ready)}/{num_stages} "
|
||
f"stages ready. Missing stages: {not_ready}"
|
||
)
|
||
|
||
suggestions = [
|
||
f"Ignore this warning if the model weight download / load from disk time is longer than {timeout}s.",
|
||
"Verify GPU/device assignment in config (runtime.devices) is correct.",
|
||
"Check GPU/host memory availability; reduce model or batch size if needed.",
|
||
"Check model weights path and network reachability (if loading remotely).",
|
||
"Increase initialization wait time (stage_init_timeout or call-site timeout).",
|
||
]
|
||
|
||
formatted_suggestions = "\n".join(f" {i + 1}) {msg}" for i, msg in enumerate(suggestions))
|
||
|
||
logger.warning(f"[{self._name}] Stage initialization timeout. Troubleshooting Steps:\n{formatted_suggestions}")
|
||
|
||
def _is_profiler_enabled(self, stage_id: int) -> bool:
|
||
"""Check if profiler config is set for a given stage."""
|
||
stage = self.stage_list[stage_id]
|
||
# For diffusion stages, profiling is controlled by VLLM_TORCH_PROFILER_DIR env var
|
||
if stage.stage_type == "diffusion":
|
||
return True
|
||
# For LLM stages, check if profiler_config is set in engine_args
|
||
engine_args = getattr(stage.stage_config, "engine_args", None)
|
||
if engine_args is None:
|
||
return False
|
||
profiler_config = getattr(engine_args, "profiler_config", None)
|
||
if profiler_config is None:
|
||
return False
|
||
profiler = getattr(profiler_config, "profiler", None)
|
||
return profiler is not None
|
||
|
||
def _setup_rpc_result_checkers(self) -> None:
|
||
"""Set up RPC result checkers for all stages.
|
||
|
||
Each checker reads from the shared ``_rpc_results`` dict so that
|
||
``OmniStage.collective_rpc`` can retrieve results that were
|
||
collected by the orchestrator (or, in the sync path, by the
|
||
stage-level checker that drains the output queue).
|
||
|
||
Uses a weak reference to ``self`` to avoid a circular reference
|
||
(OmniBase → stage_list → OmniStage → closure → OmniBase) that
|
||
would prevent the instance from being freed by reference counting.
|
||
"""
|
||
weak_self = weakref.ref(self)
|
||
for stage in self.stage_list:
|
||
sid = stage.stage_id
|
||
|
||
def make_rpc_checker(stage_id: int):
|
||
def rpc_checker(rpc_id: str) -> dict[str, Any] | None:
|
||
_self = weak_self()
|
||
if _self is None:
|
||
return None
|
||
# First check the shared dict
|
||
if stage_id in _self._rpc_results and rpc_id in _self._rpc_results[stage_id]:
|
||
return _self._rpc_results[stage_id].pop(rpc_id)
|
||
# In the sync path there is no background output handler,
|
||
# so drain the output queue ourselves and stash any
|
||
# non-RPC results back.
|
||
out_q = _self._stage_out_queues[stage_id] if stage_id < len(_self._stage_out_queues) else None
|
||
if out_q is not None:
|
||
import queue as _queue
|
||
|
||
try:
|
||
while True:
|
||
item = out_q.get_nowait()
|
||
if isinstance(item, dict) and item.get("type") == "collective_rpc_result":
|
||
item_rpc_id = item.get("rpc_id")
|
||
if item_rpc_id == rpc_id:
|
||
return item
|
||
# Stash for another caller
|
||
if stage_id not in _self._rpc_results:
|
||
_self._rpc_results[stage_id] = {}
|
||
_self._rpc_results[stage_id][item_rpc_id] = item
|
||
else:
|
||
# Non-RPC item — put it back
|
||
out_q.put(item)
|
||
break
|
||
except _queue.Empty:
|
||
pass
|
||
return None
|
||
|
||
return rpc_checker
|
||
|
||
stage._rpc_result_checker = make_rpc_checker(sid)
|
||
|
||
def collective_rpc(
|
||
self,
|
||
method: str | Callable[..., _R],
|
||
timeout: float | None = None,
|
||
args: tuple = (),
|
||
kwargs: dict[str, Any] | None = None,
|
||
) -> list[_R]:
|
||
"""Execute an RPC call on all stage workers.
|
||
|
||
Args:
|
||
method: Name of the worker method to execute, or a callable.
|
||
timeout: Maximum time in seconds to wait for execution.
|
||
args: Positional arguments to pass to the worker method.
|
||
kwargs: Keyword arguments to pass to the worker method.
|
||
|
||
Returns:
|
||
A list containing the results from each stage.
|
||
"""
|
||
results: list[_R] = []
|
||
for stage in self.stage_list:
|
||
results.append(
|
||
stage.collective_rpc(
|
||
method=method,
|
||
timeout=timeout,
|
||
args=args,
|
||
kwargs=kwargs,
|
||
)
|
||
)
|
||
return results
|
||
|
||
def sleep(self, level: int = 1) -> None:
|
||
"""Put the engine to sleep to free up resources.
|
||
|
||
Args:
|
||
level: Sleep level (1 = light sleep, higher = deeper sleep).
|
||
"""
|
||
self._is_sleeping = True
|
||
self.collective_rpc(method="sleep", args=(level,))
|
||
|
||
def wake_up(self, tags: list[str] | None = None) -> None:
|
||
"""Wake the engine up from sleep.
|
||
|
||
Args:
|
||
tags: Optional list of tags to selectively wake components.
|
||
"""
|
||
self._is_sleeping = False
|
||
self.collective_rpc(method="wake_up", args=(tags,))
|
||
|
||
def is_sleeping(self) -> bool:
|
||
"""Check whether the engine is sleeping."""
|
||
return getattr(self, "_is_sleeping", False)
|
||
|
||
def add_lora(self, lora_request: LoRARequest) -> bool:
|
||
"""Load a new LoRA adapter into the engine for future requests."""
|
||
result = self.collective_rpc(method="add_lora", args=(lora_request,))
|
||
return result[0][0]
|
||
|
||
def start_profile(self, stages: list[int] | None = None) -> None:
|
||
"""Start profiling for specified stages.
|
||
|
||
Sends start_profile command to stage workers. Profiling must be enabled
|
||
via VLLM_TORCH_PROFILER_DIR environment variable.
|
||
|
||
Args:
|
||
stages: List of stage IDs to start profiling. If None, starts
|
||
profiling for all stages that have profiling enabled.
|
||
|
||
Example:
|
||
>>> # Profile all stages
|
||
>>> omni.start_profile()
|
||
>>> outputs = omni.generate(prompts, sampling_params)
|
||
>>> omni.stop_profile()
|
||
|
||
>>> # Profile only stage 0 and 2
|
||
>>> omni.start_profile(stages=[0, 2])
|
||
"""
|
||
if stages is None:
|
||
stages = list(range(len(self.stage_list)))
|
||
|
||
for stage_id in stages:
|
||
if stage_id < len(self.stage_list):
|
||
if not self._is_profiler_enabled(stage_id):
|
||
logger.info(
|
||
"[%s] Skipping start_profile for stage-%s: profiler config not set",
|
||
self._name,
|
||
stage_id,
|
||
)
|
||
continue
|
||
try:
|
||
self.stage_list[stage_id].submit({"type": OmniStageTaskType.PROFILER_START})
|
||
logger.info("[%s] Sent start_profile to stage-%s", self._name, stage_id)
|
||
except Exception as e:
|
||
logger.warning(
|
||
"[%s] Failed to send start_profile to stage-%s: %s",
|
||
self._name,
|
||
stage_id,
|
||
e,
|
||
)
|
||
|
||
def stop_profile(self, stages: list[int] | None = None) -> dict:
|
||
"""
|
||
Synchronously stop profiling for specified stages and collect
|
||
the file paths for traces and tables.
|
||
"""
|
||
if stages is None:
|
||
stages = list(range(len(self.stage_list)))
|
||
|
||
all_results = {"traces": [], "tables": []}
|
||
|
||
for stage_id in stages:
|
||
if stage_id < len(self.stage_list):
|
||
if not self._is_profiler_enabled(stage_id):
|
||
logger.info(
|
||
"[%s] Skipping stop_profile for stage-%s: profiler config not set",
|
||
self._name,
|
||
stage_id,
|
||
)
|
||
continue
|
||
stage = self.stage_list[stage_id]
|
||
|
||
# Check if the stage object has our new bridge method
|
||
if hasattr(stage, "stop_profile"):
|
||
logger.info("[%s] Requesting profile data collection from stage-%s", self._name, stage_id)
|
||
|
||
# This is the blocking call that triggers the RPC chain
|
||
stage_data = stage.stop_profile()
|
||
|
||
if isinstance(stage_data, dict):
|
||
# FIX: Handle both single key and list key formats
|
||
traces = stage_data.get("trace") or stage_data.get("traces")
|
||
tables = stage_data.get("table") or stage_data.get("tables")
|
||
|
||
# Debug logging
|
||
logger.debug(f"[{self._name}] Stage-{stage_id} returned: {stage_data.keys()}")
|
||
if traces:
|
||
logger.debug(f"[{self._name}] Stage-{stage_id} traces type: {type(traces)}")
|
||
if tables:
|
||
logger.debug(f"[{self._name}] Stage-{stage_id} tables type: {type(tables)}")
|
||
|
||
# Handle single strings
|
||
if traces:
|
||
if isinstance(traces, str):
|
||
all_results["traces"].append(traces)
|
||
elif isinstance(traces, list):
|
||
all_results["traces"].extend(traces)
|
||
|
||
# Handle single strings
|
||
if tables:
|
||
if isinstance(tables, str):
|
||
all_results["tables"].append(tables)
|
||
elif isinstance(tables, list):
|
||
all_results["tables"].extend(tables)
|
||
else:
|
||
logger.warning(f"[{self._name}] Stage-{stage_id} returned no table data")
|
||
else:
|
||
logger.warning(f"[{self._name}] Stage-{stage_id} returned non-dict data: {type(stage_data)}")
|
||
else:
|
||
# Fallback for non-diffusion stages
|
||
logger.warning(
|
||
"[%s] Stage-%s does not support synchronous stop_profile. Falling back to async.",
|
||
self._name,
|
||
stage_id,
|
||
)
|
||
stage.submit({"type": OmniStageTaskType.PROFILER_STOP})
|
||
|
||
# Final debug output
|
||
logger.info(
|
||
f"[{self._name}] Collected {len(all_results['traces'])} trace(s) and {len(all_results['tables'])} table(s)"
|
||
)
|
||
|
||
return all_results
|
||
|
||
def close(self) -> None:
|
||
"""Close all stage processes and clean up resources."""
|
||
if hasattr(self, "_weak_finalizer"):
|
||
self._weak_finalizer()
|
||
|
||
def _process_handshake_message(self, msg: Any) -> dict[str, Any]:
|
||
"""Process incoming handshake message and generate response.
|
||
|
||
Args:
|
||
msg: Decoded message from client
|
||
|
||
Returns:
|
||
Response dictionary with ok status and either endpoints or error
|
||
"""
|
||
if not isinstance(msg, dict) or msg.get("type") != "handshake":
|
||
return {"ok": False, "error": "invalid handshake payload"}
|
||
|
||
try:
|
||
stage_id = int(msg.get("stage_id"))
|
||
except (TypeError, ValueError) as e:
|
||
return {"ok": False, "error": f"invalid stage_id: {e}"}
|
||
|
||
endpoints = self._handshake_endpoints.get(stage_id)
|
||
if endpoints is None:
|
||
return {"ok": False, "error": f"unknown stage_id: {stage_id}"}
|
||
|
||
# Mark stage as seen and prepare success response
|
||
self._handshake_seen.add(stage_id)
|
||
in_endpoint, out_endpoint = endpoints
|
||
|
||
logger.info(
|
||
"[%s] Handshake received from stage-%s",
|
||
self._name,
|
||
stage_id,
|
||
)
|
||
|
||
return {
|
||
"ok": True,
|
||
"in_endpoint": in_endpoint,
|
||
"out_endpoint": out_endpoint,
|
||
}
|
||
|
||
def _run_handshake_server_loop(self) -> None:
|
||
"""Main loop for handshake server - polls for messages and responds."""
|
||
poller = zmq.Poller()
|
||
poller.register(self._zmq_handshake_socket, zmq.POLLIN)
|
||
|
||
try:
|
||
while not self._handshake_stop.is_set():
|
||
events = poller.poll(1000)
|
||
has_message = any(sock == self._zmq_handshake_socket and event == zmq.POLLIN for sock, event in events)
|
||
if not has_message:
|
||
continue
|
||
|
||
msg = msgspec.msgpack.decode(self._zmq_handshake_socket.recv())
|
||
response = msgspec.msgpack.encode(self._process_handshake_message(msg))
|
||
self._zmq_handshake_socket.send(response)
|
||
finally:
|
||
poller.unregister(self._zmq_handshake_socket)
|
||
|
||
def start_handshake_server(self) -> None:
|
||
"""Start the ZMQ handshake server.
|
||
|
||
The handshake server allows distributed stages to discover their
|
||
queue endpoints by querying the orchestrator with their stage_id.
|
||
Skips starting if the server is already running or ZMQ is not initialized.
|
||
"""
|
||
# Skip if already running or ZMQ not initialized
|
||
if self._handshake_thread is not None or self._zmq_ctx is None:
|
||
return
|
||
|
||
# Skip if master address/port not configured
|
||
if not self._zmq_master_address or self._zmq_master_port is None:
|
||
return
|
||
|
||
# Create server endpoint and socket
|
||
endpoint = get_engine_client_zmq_addr(
|
||
local_only=False, host=self._zmq_master_address, port=int(self._zmq_master_port)
|
||
)
|
||
|
||
self._handshake_stop = threading.Event()
|
||
self._zmq_handshake_socket = make_zmq_socket(self._zmq_ctx, endpoint, zmq.REP, bind=True, linger=5000)
|
||
|
||
# Start server thread
|
||
self._handshake_thread = threading.Thread(
|
||
target=self._run_handshake_server_loop, daemon=True, name="zmq-handshake-server"
|
||
)
|
||
self._handshake_thread.start()
|
||
|
||
def _wait_for_handshakes(self, timeout: int = 120) -> int:
|
||
"""Wait for handshakes from all expected stages.
|
||
|
||
Args:
|
||
timeout: Timeout in seconds for waiting for handshakes. Default is 120s.
|
||
|
||
Returns:
|
||
Remaining timeout in seconds after waiting for handshakes.
|
||
"""
|
||
total_stages = len(self.stage_configs)
|
||
expected = set(range(total_stages)) - {int(self._single_stage_id)}
|
||
if not expected:
|
||
return timeout
|
||
|
||
deadline = time.time() + max(0, int(timeout))
|
||
logger.info(f"[{self._name}] Waiting for handshakes from stages: {expected} (timeout: {timeout}s)")
|
||
|
||
# NOTE: _handshake_seen may be updated from the handshake server thread.
|
||
# It is intentionally used here without additional locking because:
|
||
# - _handshake_seen only ever grows (stages are added but never removed), and
|
||
# - we only check membership and set inclusion relative to `expected`.
|
||
# Under these monotonic semantics and the CPython GIL, concurrent reads/writes
|
||
# are safe for this usage and cannot violate correctness: we may observe a
|
||
# slightly stale view, but the loop condition remains valid and eventually
|
||
# becomes true once all expected stages have handshaked or the timeout elapses.
|
||
while not expected.issubset(self._handshake_seen) and time.time() < deadline:
|
||
time.sleep(1.0)
|
||
|
||
remaining_timeout = max(0, int(deadline - time.time()))
|
||
|
||
if not expected.issubset(self._handshake_seen):
|
||
missing = sorted(expected - self._handshake_seen)
|
||
logger.warning(
|
||
f"[{self._name}] Handshake timeout: {len(self._handshake_seen)}/{len(expected)} "
|
||
f"stages completed handshake. Missing stages: {missing}"
|
||
)
|
||
|
||
return remaining_timeout
|
||
|
||
@property
|
||
def _name(self) -> str:
|
||
return "OmniBase"
|
||
|
||
@property
|
||
def is_async(self) -> bool:
|
||
return False
|
||
|
||
|
||
class Omni(OmniBase):
|
||
"""Unified entrypoint for both LLM and Diffusion models for better usability.
|
||
|
||
Args:
|
||
model: Model name or path to load.
|
||
**kwargs: Arbitrary keyword arguments.
|
||
- stage_configs_path: Optional path to YAML file containing stage
|
||
configurations. If None, configurations are loaded from the model.
|
||
- log_stats: Whether to enable statistics logging
|
||
be written to files with stage-specific suffixes.
|
||
- stage_init_timeout: Per-stage init watchdog (seconds). Measured from
|
||
when the previous stage finished (possibly a prior Omni run with GPU
|
||
reuse/overlap) to when the current stage starts to initialize.
|
||
- shm_threshold_bytes: Threshold in bytes for using shared memory
|
||
for IPC. Objects larger than this threshold will use shared memory.
|
||
- worker_backend: Backend for worker processes. Default is "multi_process".
|
||
- ray_address: Address of Ray cluster for Ray backend, if using Ray backend.
|
||
- batch_timeout: Timeout in seconds for batching requests within a stage
|
||
- init_timeout: Timeout in seconds for waiting for all stages to initialize
|
||
- Additional keyword arguments passed to stage engines.
|
||
|
||
Example:
|
||
>>> omni = Omni(model="Qwen/Qwen2.5-Omni-7B")
|
||
>>> outputs = omni.generate(prompts="Hello, world!", sampling_params_list=[SamplingParams()])
|
||
>>> print(outputs)
|
||
"""
|
||
|
||
def __init__(self, model: str, **kwargs: Any) -> None:
|
||
super().__init__(model, **kwargs)
|
||
|
||
# Register weak reference cleanup (called on garbage collection)
|
||
self._weak_finalizer = weakref.finalize(
|
||
self,
|
||
_weak_close_cleanup,
|
||
self.stage_list,
|
||
self._stage_in_queues,
|
||
self._stage_out_queues,
|
||
self._ray_pg,
|
||
self._zmq_ctx,
|
||
self._handshake_stop,
|
||
self._zmq_handshake_socket,
|
||
self._handshake_thread,
|
||
)
|
||
|
||
@overload
|
||
def generate(
|
||
self,
|
||
prompts: OmniPromptType | Sequence[OmniPromptType],
|
||
sampling_params_list: OmniSamplingParams | Sequence[OmniSamplingParams] | None = None,
|
||
*,
|
||
py_generator: Literal[True],
|
||
) -> Generator[OmniRequestOutput, None, None]: ...
|
||
|
||
@overload
|
||
def generate(
|
||
self,
|
||
prompts: OmniPromptType | Sequence[OmniPromptType],
|
||
sampling_params_list: OmniSamplingParams | Sequence[OmniSamplingParams] | None = None,
|
||
*,
|
||
py_generator: Literal[False] = False,
|
||
) -> list[OmniRequestOutput]: ...
|
||
|
||
def generate(
|
||
self,
|
||
prompts: OmniPromptType | Sequence[OmniPromptType],
|
||
sampling_params_list: OmniSamplingParams | Sequence[OmniSamplingParams] | None = None,
|
||
*,
|
||
py_generator: bool = False,
|
||
use_tqdm: bool | Callable[..., tqdm] = True,
|
||
) -> Generator[OmniRequestOutput, None, None] | list[OmniRequestOutput]:
|
||
"""Generate outputs for the given prompts.
|
||
|
||
Orchestrates the multi-stage pipeline based on YAML configuration.
|
||
Each stage will use OmniLLM or OmniDiffusion based on stage_type.
|
||
|
||
Args:
|
||
prompts: Input prompt(s) for generation.
|
||
sampling_params_list: Optional list of per-stage parameters.
|
||
py_generator: Whether the returned result(s) are wrapped in a generator instead of a list.
|
||
use_tqdm: Whether to use tqdm progress bar
|
||
|
||
Returns:
|
||
List of OmniRequestOutput objects, one for each input prompt.
|
||
Each output contains the stage_id, final_output_type, and
|
||
the request_output from the final stage.
|
||
|
||
Raises:
|
||
ValueError: If sampling_params_list is None or has incorrect length.
|
||
"""
|
||
if sampling_params_list is None:
|
||
sampling_params_list = self.default_sampling_params_list
|
||
elif not isinstance(sampling_params_list, Sequence):
|
||
# TODO: After the recent introduction of BAGEL model (one LLM and one Diffusion),
|
||
# expect the text_to_image example code to run when only passing one OmniDiffusionSamplingParams
|
||
# This behavior may be confusing, and future PR can improve it.
|
||
per_stage_params: list[OmniSamplingParams] = []
|
||
for default_stage_sp in self.default_sampling_params_list:
|
||
default_sp_type = default_stage_sp.__class__
|
||
if default_sp_type == sampling_params_list.__class__:
|
||
per_stage_params.append(sampling_params_list)
|
||
else:
|
||
per_stage_params.append(default_stage_sp)
|
||
sampling_params_list = per_stage_params
|
||
|
||
try:
|
||
if py_generator:
|
||
return self._run_generation_with_generator(prompts, sampling_params_list)
|
||
else:
|
||
outputs = list(self._run_generation(prompts, sampling_params_list, use_tqdm))
|
||
return outputs
|
||
except Exception as e:
|
||
logger.exception("[Orchestrator] Failed to run generation: %s", e)
|
||
# Always close on exception to ensure cleanup
|
||
self.close()
|
||
raise e
|
||
|
||
def _run_generation_with_generator(
|
||
self,
|
||
prompts: OmniPromptType | Sequence[OmniPromptType],
|
||
sampling_params_list: Sequence[OmniSamplingParams],
|
||
) -> Generator[OmniRequestOutput, None, None]:
|
||
"""Run generation through all stages in the pipeline and return a generator."""
|
||
gen = self._run_generation(prompts, sampling_params_list)
|
||
try:
|
||
yield from gen
|
||
except Exception as e:
|
||
logger.exception("[Orchestrator] Failed to run generation: %s", e)
|
||
raise e
|
||
finally:
|
||
# Cleanup when generator is exhausted or closed
|
||
self.close()
|
||
|
||
def _run_generation(
|
||
self,
|
||
prompts: OmniPromptType | Sequence[OmniPromptType],
|
||
sampling_params_list: Sequence[OmniSamplingParams],
|
||
use_tqdm: bool | Callable[..., tqdm] = True,
|
||
) -> Generator[OmniRequestOutput, None, None]:
|
||
"""Run generation through all stages in the pipeline."""
|
||
logger.debug(f"[{self._name}] generate() called")
|
||
if sampling_params_list is None:
|
||
raise ValueError("sampling_params_list is required for pipelined generation")
|
||
|
||
if len(sampling_params_list) != len(self.stage_list):
|
||
raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
|
||
|
||
for i, (stage, sp) in enumerate(zip(self.stage_list, sampling_params_list)):
|
||
ExpectedSPType = OmniDiffusionSamplingParams if stage.stage_type == "diffusion" else SamplingParams
|
||
if not isinstance(sp, ExpectedSPType):
|
||
raise ValueError(
|
||
f"Expected sampling parameters with type {ExpectedSPType} in stage {i}, got {sp.__class__}"
|
||
)
|
||
|
||
# Normalize prompts to a list for per-request iteration
|
||
# str is also Sequence but only test list-like containers here
|
||
if isinstance(prompts, str) or not isinstance(prompts, Sequence):
|
||
request_prompts: list[OmniPromptType] = [prompts]
|
||
else:
|
||
request_prompts = list(prompts)
|
||
|
||
# Orchestrator keeps stage objects for input derivation
|
||
num_stages = len(self.stage_list)
|
||
|
||
# Generate globally unique request IDs and map them to original prompts
|
||
request_ids = [f"{i}_{uuid.uuid4()}" for i in range(len(request_prompts))]
|
||
request_id_to_prompt = {rid: p for rid, p in zip(request_ids, request_prompts)}
|
||
|
||
# Track per-request start time for end-to-end timing
|
||
_req_start_ts: dict[str, float] = {}
|
||
_wall_start_ts: float = time.time()
|
||
|
||
# CFG companion tracking (prompt expansion + lifecycle management)
|
||
cfg = CfgCompanionTracker(
|
||
prompt_expand_func=getattr(self.stage_list[0], "prompt_expand_func", None),
|
||
stage0_sampling_params=sampling_params_list[0],
|
||
)
|
||
expanded_companions = cfg.expand_prompts(request_id_to_prompt)
|
||
|
||
# Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
|
||
final_stage_id_to_prompt: dict[str, int] = {}
|
||
for rid, prompt in request_id_to_prompt.items():
|
||
if isinstance(prompt, dict):
|
||
prompt_modalities = prompt.get("modalities", None)
|
||
else:
|
||
prompt_modalities = None
|
||
final_stage_id_for_e2e = get_final_stage_id_for_e2e(
|
||
prompt_modalities, self.output_modalities, self.stage_list
|
||
)
|
||
final_stage_id_to_prompt[rid] = final_stage_id_for_e2e
|
||
|
||
# Metrics/aggregation helper
|
||
metrics = OrchestratorAggregator(
|
||
num_stages,
|
||
self.log_stats,
|
||
_wall_start_ts,
|
||
final_stage_id_to_prompt,
|
||
)
|
||
|
||
it = request_id_to_prompt.items()
|
||
if use_tqdm:
|
||
tqdm_func = use_tqdm if callable(use_tqdm) else tqdm
|
||
it = tqdm_func(it, desc="Adding requests")
|
||
|
||
# Seed stage-0 queue with all requests
|
||
logger.debug(f"[{self._name}] Seeding {len(request_prompts)} requests into stage-0")
|
||
# Mark first input time for stage-0
|
||
metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
|
||
|
||
for req_id, prompt in request_id_to_prompt.items():
|
||
sp0 = sampling_params_list[0] # type: ignore[index]
|
||
task = {
|
||
"request_id": req_id,
|
||
"engine_inputs": prompt,
|
||
"sampling_params": sp0,
|
||
}
|
||
self.stage_list[0].submit(task)
|
||
_req_start_ts[req_id] = time.time()
|
||
logger.debug(f"[{self._name}] Enqueued request {req_id} to stage-0")
|
||
|
||
# Submit CFG companion requests to stage-0
|
||
if cfg.is_active:
|
||
for companion_id, companion_prompt in expanded_companions:
|
||
task = {
|
||
"request_id": companion_id,
|
||
"engine_inputs": companion_prompt,
|
||
"sampling_params": cfg.stage0_sampling_params,
|
||
}
|
||
self.stage_list[0].submit(task)
|
||
_req_start_ts[companion_id] = time.time()
|
||
logger.debug(f"[{self._name}] Enqueued CFG companion {companion_id} to stage-0")
|
||
|
||
pbar = None
|
||
if use_tqdm:
|
||
tqdm_func = use_tqdm if callable(use_tqdm) else tqdm
|
||
pbar = tqdm_func(
|
||
total=len(request_prompts),
|
||
desc="Processed prompts",
|
||
dynamic_ncols=True,
|
||
postfix=(f"est. speed input: {0:.2f} unit/s, output: {0:.2f} unit/s"),
|
||
)
|
||
# For each stage, forward results to next stage; collect finals at the end
|
||
# We pipeline by continually polling output queues in stage order
|
||
remaining_by_stage: list[int] = [len(request_prompts) + cfg.num_companions] + [0] * (num_stages - 1)
|
||
completed_requests = 0
|
||
total_requests = len(request_prompts)
|
||
|
||
logger.debug(
|
||
f"[{self._name}] Entering scheduling loop: total_requests={total_requests}, stages={num_stages}",
|
||
)
|
||
while completed_requests < total_requests:
|
||
made_progress = False
|
||
for stage_id, stage in enumerate(self.stage_list):
|
||
result = stage.try_collect()
|
||
if result is None:
|
||
continue
|
||
|
||
made_progress = True
|
||
req_id = result.get("request_id")
|
||
if "error" in result:
|
||
logger.error(
|
||
f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
|
||
)
|
||
if cfg.is_companion(req_id) and stage_id == 0:
|
||
parent_id, parent_aborted = cfg.on_companion_error(req_id)
|
||
if parent_aborted:
|
||
completed_requests += 1
|
||
logger.error(
|
||
f"[{self._name}] Parent {parent_id} aborted due to "
|
||
f"companion failure ({completed_requests}/{total_requests})",
|
||
)
|
||
continue
|
||
|
||
if result.get("type") == "stage_ready":
|
||
# Only happens when stage is initialized slower than expected,
|
||
# so we wait for a short time and try again
|
||
time.sleep(0.05)
|
||
continue
|
||
|
||
# CFG: companion requests only run through Stage-0
|
||
if cfg.is_companion(req_id) and stage_id == 0:
|
||
ready_parent = cfg.on_companion_completed(req_id)
|
||
if ready_parent is not None:
|
||
success = cfg.forward_parent_with_cfg(
|
||
ready_parent,
|
||
cfg.pop_pending_parent(ready_parent),
|
||
self.stage_list,
|
||
self.connectors,
|
||
sampling_params_list,
|
||
request_id_to_prompt,
|
||
final_stage_id_to_prompt,
|
||
metrics,
|
||
remaining_by_stage,
|
||
)
|
||
if not success:
|
||
cfg.consume_parent_failure(ready_parent)
|
||
completed_requests += 1
|
||
logger.error(
|
||
f"[{self._name}] Parent {ready_parent} dropped due to CFG forwarding failure "
|
||
f"({completed_requests}/{total_requests})",
|
||
)
|
||
continue
|
||
|
||
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
|
||
# Mark last output time for this stage whenever we receive outputs
|
||
metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
|
||
try:
|
||
_m: StageRequestStats = result.get("metrics")
|
||
if _m is not None:
|
||
# Accumulate generation time
|
||
metrics.accumulated_gen_time_ms[req_id][stage_id] += _m.stage_gen_time_ms
|
||
|
||
# For diffusion stages, we also accumulate diffusion time
|
||
metrics.accumulate_diffusion_metrics(stage.stage_type, req_id, engine_outputs)
|
||
|
||
metrics.on_stage_metrics(stage_id, req_id, _m, stage.final_output_type)
|
||
if pbar:
|
||
elapsed = pbar.format_dict["elapsed"] or 1e-6
|
||
# Aggregate total tokens/images across all stages
|
||
total_out = sum(metrics.stage_total_tokens)
|
||
out_spd = total_out / elapsed
|
||
|
||
modality = self.output_modalities[stage_id]
|
||
unit = "img" if modality == "image" else "tok"
|
||
|
||
# Pre-calculate for cleaner string formatting
|
||
if metrics.e2e_count > 0:
|
||
avg_lat = metrics.e2e_total_ms / metrics.e2e_count
|
||
else:
|
||
avg_lat = 0
|
||
|
||
# Align with vLLM's wording "est. speed" using multi-line parentheses
|
||
pbar.postfix = (
|
||
f"est. speed stage-{stage_id} {unit}/s: {out_spd:.2f}, avg e2e_lat: {avg_lat:.1f}ms"
|
||
)
|
||
except Exception as e:
|
||
logger.exception(
|
||
f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}",
|
||
)
|
||
logger.debug(
|
||
f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing",
|
||
)
|
||
stage.set_engine_outputs(engine_outputs)
|
||
|
||
if getattr(stage, "final_output", False):
|
||
logger.debug(
|
||
f"[{self._name}] Request {req_id} finalized at stage-{stage_id}",
|
||
)
|
||
|
||
# End-to-end timing and time-per-token for final output
|
||
# (only once per request at the designated final stage)
|
||
try:
|
||
if stage_id == final_stage_id_to_prompt[req_id]:
|
||
metrics.on_finalize_request(
|
||
stage_id,
|
||
req_id,
|
||
_req_start_ts.get(req_id, _wall_start_ts),
|
||
)
|
||
except Exception as e:
|
||
logger.exception(
|
||
f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}",
|
||
)
|
||
output_to_yield = OmniRequestOutput(
|
||
stage_id=stage_id,
|
||
final_output_type=stage.final_output_type, # type: ignore[attr-defined]
|
||
request_output=engine_outputs,
|
||
)
|
||
|
||
# Record audio generated frames (only when finished)
|
||
try:
|
||
finished = (
|
||
engine_outputs.finished
|
||
if hasattr(engine_outputs, "finished")
|
||
else (
|
||
engine_outputs[0].finished
|
||
if isinstance(engine_outputs, list)
|
||
and engine_outputs
|
||
and hasattr(engine_outputs[0], "finished")
|
||
else False
|
||
)
|
||
)
|
||
if finished:
|
||
metrics.record_audio_generated_frames(output_to_yield, stage_id, req_id)
|
||
except Exception as e:
|
||
logger.exception(
|
||
f"[{self._name}] Failed to record audio metrics for req {req_id} at stage {stage_id}: {e}",
|
||
)
|
||
|
||
yield output_to_yield
|
||
|
||
next_stage_id = stage_id + 1
|
||
if next_stage_id <= final_stage_id_to_prompt[req_id]:
|
||
# CFG: if this parent has companions, defer forwarding
|
||
if cfg.has_companions(req_id) and stage_id == 0:
|
||
if cfg.is_parent_failed(req_id):
|
||
cfg.consume_parent_failure(req_id)
|
||
completed_requests += 1
|
||
logger.error(
|
||
f"[{self._name}] Parent {req_id} skipped CFG forwarding due to "
|
||
f"companion failure ({completed_requests}/{total_requests})",
|
||
)
|
||
continue
|
||
|
||
if cfg.all_companions_done(req_id):
|
||
success = cfg.forward_parent_with_cfg(
|
||
req_id,
|
||
{"engine_outputs": engine_outputs, "stage_id": stage_id},
|
||
self.stage_list,
|
||
self.connectors,
|
||
sampling_params_list,
|
||
request_id_to_prompt,
|
||
final_stage_id_to_prompt,
|
||
metrics,
|
||
remaining_by_stage,
|
||
)
|
||
if not success:
|
||
cfg.consume_parent_failure(req_id)
|
||
completed_requests += 1
|
||
logger.error(
|
||
f"[{self._name}] Parent {req_id} dropped due to CFG forwarding failure "
|
||
f"({completed_requests}/{total_requests})",
|
||
)
|
||
else:
|
||
cfg.defer_parent(req_id, engine_outputs, stage_id)
|
||
continue
|
||
|
||
next_stage: OmniStage = self.stage_list[next_stage_id]
|
||
try:
|
||
# Derive inputs for the next stage, record preprocess time
|
||
with metrics.stage_postprocess_timer(stage_id, req_id):
|
||
next_inputs = next_stage.process_engine_inputs(
|
||
self.stage_list, [request_id_to_prompt[req_id]]
|
||
)
|
||
except Exception as e:
|
||
completed_requests += 1
|
||
logger.exception(
|
||
f"[{self._name}] Process engine inputs error for req {req_id}"
|
||
f" at stage {next_stage_id}: {e} ({completed_requests}/{total_requests})",
|
||
)
|
||
continue
|
||
sp_next = sampling_params_list[next_stage_id] # type: ignore[index]
|
||
|
||
# Check if we have a connector for this edge
|
||
connector_key = (str(stage_id), str(next_stage_id))
|
||
connector = self.connectors.get(connector_key)
|
||
sent_via_connector = False
|
||
if connector:
|
||
sent_via_connector = try_send_via_connector(
|
||
connector=connector,
|
||
stage_id=stage_id,
|
||
next_stage_id=next_stage_id,
|
||
req_id=req_id,
|
||
next_inputs=next_inputs,
|
||
sampling_params=sp_next,
|
||
original_prompt=request_id_to_prompt[req_id],
|
||
next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
|
||
metrics=metrics,
|
||
)
|
||
|
||
if not sent_via_connector:
|
||
raise RuntimeError(
|
||
f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
|
||
"Configure a connector for this edge or inspect connector logs for details."
|
||
)
|
||
|
||
logger.debug(
|
||
f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}",
|
||
)
|
||
remaining_by_stage[next_stage_id] += 1
|
||
else:
|
||
completed_requests += 1
|
||
if pbar:
|
||
final_mod = self.output_modalities[final_stage_id_to_prompt[req_id]]
|
||
pbar.unit = "img" if final_mod == "image" else "req"
|
||
pbar.update(1)
|
||
logger.debug(
|
||
f"[{self._name}] Request {req_id} fully completed ({completed_requests}/{total_requests})",
|
||
)
|
||
for timed_out_id in cfg.check_timeouts():
|
||
completed_requests += 1
|
||
logger.error(
|
||
f"[{self._name}] Parent {timed_out_id} timed out; counting as failed "
|
||
f"({completed_requests}/{total_requests})",
|
||
)
|
||
|
||
if not made_progress:
|
||
time.sleep(0.005)
|
||
logger.debug(f"[{self._name}] All requests completed")
|
||
|
||
if pbar:
|
||
pbar.close()
|
||
|
||
# Summarize and print stats
|
||
try:
|
||
metrics.build_and_log_summary()
|
||
except Exception as e:
|
||
logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
|
||
|
||
@property
|
||
def _name(self) -> str:
|
||
return "Orchestrator"
|