867 lines
38 KiB
Python
867 lines
38 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 time
|
|
import uuid
|
|
import weakref
|
|
from collections.abc import Callable, Generator, Sequence
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from dataclasses import asdict
|
|
from pprint import pformat
|
|
from typing import Any, Literal, overload
|
|
|
|
from omegaconf import OmegaConf
|
|
from tqdm.auto import tqdm
|
|
from vllm import SamplingParams
|
|
from vllm.logger import init_logger
|
|
|
|
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.log_utils import OrchestratorMetrics
|
|
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 (
|
|
get_final_stage_id_for_e2e,
|
|
inject_omni_kv_config,
|
|
load_stage_configs_from_model,
|
|
load_stage_configs_from_yaml,
|
|
resolve_model_config_path,
|
|
)
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _weak_close_cleanup(stage_list, stage_in_queues, ray_pg):
|
|
"""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}")
|
|
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)
|
|
|
|
|
|
def _dummy_snapshot_download(model_id):
|
|
return model_id
|
|
|
|
|
|
def omni_snapshot_download(model_id) -> str:
|
|
# 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)
|
|
else:
|
|
return _dummy_snapshot_download(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[mp.Queue] = []
|
|
self._stage_out_queues: list[mp.Queue] = []
|
|
self._stages_ready: set[int] = set()
|
|
self._ray_pg = None
|
|
self._queue_cls = None
|
|
self._ctx = None
|
|
|
|
# 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]) -> dict[str, Any]:
|
|
"""Create default diffusion stage configuration."""
|
|
# We temporally create a default config for diffusion stage.
|
|
# In the future, we should merge the default config with the user-provided config.
|
|
# TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error.
|
|
if "dtype" in kwargs:
|
|
kwargs["dtype"] = str(kwargs["dtype"])
|
|
cache_backend = kwargs.get("cache_backend", "none")
|
|
cache_config = self._normalize_cache_config(cache_backend, kwargs.get("cache_config", None))
|
|
# TODO: hack, calculate devices based on parallel config.
|
|
devices = "0"
|
|
if "parallel_config" in kwargs:
|
|
num_devices = kwargs["parallel_config"].world_size
|
|
for i in range(1, num_devices):
|
|
devices += f",{i}"
|
|
default_stage_cfg = [
|
|
{
|
|
"stage_id": 0,
|
|
"stage_type": "diffusion",
|
|
"runtime": {
|
|
"process": True,
|
|
"devices": devices,
|
|
"max_batch_size": 1,
|
|
},
|
|
"engine_args": OmegaConf.create(
|
|
{
|
|
**kwargs,
|
|
"cache_backend": cache_backend,
|
|
"cache_config": cache_config,
|
|
}
|
|
),
|
|
"final_output": True,
|
|
"final_output_type": "image",
|
|
}
|
|
]
|
|
default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
|
|
return default_stage_cfg
|
|
|
|
def _initialize_stages(self, model: str, kwargs: dict[str, Any]) -> None:
|
|
"""Initialize stage list management."""
|
|
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)
|
|
stage_configs_path = kwargs.get("stage_configs_path", None)
|
|
log_stats = kwargs.get("log_stats", False)
|
|
|
|
### base engine args
|
|
tokenizer = kwargs.get("tokenizer", None)
|
|
|
|
base_engine_args = {"tokenizer": tokenizer} if tokenizer is not None else None
|
|
|
|
# Load stage configurations from YAML
|
|
if stage_configs_path is None:
|
|
self.config_path = resolve_model_config_path(model)
|
|
self.stage_configs = load_stage_configs_from_model(model, base_engine_args=base_engine_args)
|
|
if not self.stage_configs:
|
|
default_stage_cfg = self._create_default_diffusion_stage_cfg(kwargs)
|
|
self.stage_configs = OmegaConf.create(default_stage_cfg)
|
|
else:
|
|
self.config_path = stage_configs_path
|
|
self.stage_configs = load_stage_configs_from_yaml(stage_configs_path, base_engine_args=base_engine_args)
|
|
|
|
# Inject diffusion LoRA-related knobs from kwargs if not present in the stage config.
|
|
for cfg in self.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 = OmegaConf.create({})
|
|
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
|
|
except Exception as e:
|
|
logger.warning("Failed to inject LoRA config for stage: %s", e)
|
|
|
|
# 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._enable_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.debug(f"[{self._name}] Loaded {len(self.stage_list)} stages")
|
|
|
|
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)
|
|
|
|
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"
|
|
)
|
|
|
|
for stage_id, stage in enumerate[OmniStage](self.stage_list):
|
|
in_q = self._queue_cls()
|
|
out_q = self._queue_cls()
|
|
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)
|
|
|
|
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,
|
|
)
|
|
|
|
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."""
|
|
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 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):
|
|
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):
|
|
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()
|
|
|
|
@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._ray_pg,
|
|
)
|
|
|
|
@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()
|
|
|
|
# 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 = OrchestratorMetrics(
|
|
num_stages,
|
|
self._enable_stats,
|
|
_wall_start_ts,
|
|
)
|
|
|
|
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")
|
|
|
|
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)] + [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']}",
|
|
)
|
|
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
|
|
|
|
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 = result.get("metrics")
|
|
if _m is not None:
|
|
if not isinstance(_m, dict):
|
|
_m = asdict(_m)
|
|
metrics.on_stage_metrics(stage_id, req_id, _m)
|
|
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:
|
|
rid_key = str(req_id)
|
|
if stage_id == final_stage_id_to_prompt[req_id] and rid_key not in metrics.e2e_done:
|
|
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}",
|
|
)
|
|
yield OmniRequestOutput(
|
|
stage_id=stage_id,
|
|
final_output_type=stage.final_output_type, # type: ignore[attr-defined]
|
|
request_output=engine_outputs,
|
|
)
|
|
|
|
next_stage_id = stage_id + 1
|
|
if next_stage_id <= final_stage_id_to_prompt[req_id]:
|
|
next_stage: OmniStage = self.stage_list[next_stage_id]
|
|
try:
|
|
next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
|
|
except Exception as e:
|
|
logger.exception(
|
|
f"[{self._name}] Process engine inputs error for req {req_id}"
|
|
f" at stage {next_stage_id}: {e}",
|
|
)
|
|
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})",
|
|
)
|
|
|
|
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:
|
|
summary = metrics.build_and_log_summary(final_stage_id_to_prompt)
|
|
logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
|
|
except Exception as e:
|
|
logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
|
|
|
|
@property
|
|
def _name(self) -> str:
|
|
return "Orchestrator"
|