anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Refactor] let the AsyncOmni class inherit from Omni class (#511)

Browse Source 
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Signed-off-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Peiqi YinJiangyun ZhuHongsheng Liu
 authored and GitHub committed 2025-12-29 21:04:17 +08:00
1 parent 0aedbf4a18
commit ebe37e9f57
2 files changed
+252 -600
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/entrypoints/async_omni.py
+113 -479
View File
@@ -1,16 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import multiprocessing as mp
import os
import time
from collections.abc import AsyncGenerator, Iterable
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import asdict
from pprint import pformat
from typing import Any
from omegaconf import OmegaConf
from vllm.config import VllmConfig
from vllm.inputs.preprocess import InputPreprocessor
from vllm.logger import init_logger
@@ -23,50 +19,24 @@ from vllm.v1.engine.exceptions import EngineDeadError
# Internal imports (our code)
from vllm_omni.config import OmniModelConfig
from vllm_omni.diffusion.data import DiffusionParallelConfig
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.ray_utils.utils import (
create_placement_group,
get_ray_queue_class,
try_close_ray,
)
from vllm_omni.engine.input_processor import OmniInputProcessor
from vllm_omni.entrypoints.client_request_state import ClientRequestState
from vllm_omni.entrypoints.log_utils import (
OrchestratorMetrics,
)
from vllm_omni.entrypoints.omni import OmniBase
from vllm_omni.entrypoints.omni_stage import OmniStage
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,
load_stage_configs_from_model,
load_stage_configs_from_yaml,
resolve_model_config_path,
)
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
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 AsyncOmni:
class AsyncOmni(OmniBase):
"""Asynchronous unified entry point supporting multi-stage pipelines for LLM and Diffusion models.
Similar to the Omni class, but provides an asynchronous interface supporting
@@ -102,23 +72,6 @@ class AsyncOmni:
"""
def __init__(self, *args: Any, **kwargs: dict[str, Any]) -> None:
model = args[0] if args else kwargs.get("model", "")
assert model != "", "Null model id detected, please specify a model id."
model = omni_snapshot_download(model)
if args:
args[0] = model
elif kwargs.get("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
# Pause/resume control attributes
self._pause_cond: asyncio.Condition = asyncio.Condition()
self._paused: bool = False
@@ -127,352 +80,107 @@ class AsyncOmni:
self.request_states: dict[str, ClientRequestState] = {}
self.output_handler: asyncio.Task | None = None
# Initialize stages - each stage will create appropriate instances based on stage_type
# Stage workers will automatically create AsyncOmniLLM or AsyncOmniDiffusion instances
# based on stage_type in YAML config (handled in omni_stage.py)
logger.info(f"Initializing async stages for model: {model}")
# Use kwargs-based initialization logic to avoid conflicts with old _initialize_stages signature
self._initialize_stages_from_kwargs(model, kwargs)
super().__init__(*args, **kwargs)
def _initialize_stages_from_kwargs(self, model: str, kwargs: dict[str, Any]) -> None:
"""Initialize stage list management.
Each stage will create appropriate instances (AsyncOmniLLM or AsyncOmniDiffusion)
based on stage_type in YAML config.
"""
init_sleep_seconds = kwargs.get("init_sleep_seconds", 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)
# Load stage configs 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)
if not self.stage_configs:
default_stage_cfg = [
{
"stage_id": 0,
"stage_type": "diffusion",
"runtime": {
"process": True,
"devices": "0",
"max_batch_size": 1,
},
"engine_args": {
"parallel_config": DiffusionParallelConfig(
pipeline_parallel_size=1,
data_parallel_size=1,
tensor_parallel_size=1,
sequence_parallel_size=1,
ulysses_degree=1,
ring_degree=1,
cfg_parallel_size=1,
),
"vae_use_slicing": kwargs.get("vae_use_slicing", False),
"vae_use_tiling": kwargs.get("vae_use_tiling", False),
"cache_backend": kwargs.get("cache_backend", "none"),
"cache_config": kwargs.get("cache_config", None),
},
"final_output": True,
"final_output_type": "image",
}
]
default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
self.stage_configs = OmegaConf.create(default_stage_cfg)
def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]:
"""Create default diffusion stage configuration."""
# TODO: here is different from the Omni class. We should merge the two in the future.
devices = "0"
if "parallel_config" in kwargs:
parallel_config = kwargs["parallel_config"]
num_devices = kwargs["parallel_config"].world_size
for i in range(1, num_devices):
devices += f",{i}"
else:
self.config_path = stage_configs_path
self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
# 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
# Build OmniStage instances in parallel, preserving original order
def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
idx, cfg = idx_cfg
return idx, OmniStage(cfg)
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("[AsyncOrchestrator] Loaded %d stages", len(self.stage_list))
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._init_sleep_seconds = max(0, int(init_sleep_seconds))
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 _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"
parallel_config = DiffusionParallelConfig(
pipeline_parallel_size=1,
data_parallel_size=1,
tensor_parallel_size=1,
sequence_parallel_size=1,
ulysses_degree=1,
ring_degree=1,
cfg_parallel_size=1,
)
default_stage_cfg = [
{
"stage_id": 0,
"stage_type": "diffusion",
"runtime": {
"process": True,
"devices": devices,
"max_batch_size": 1,
},
"engine_args": {
"parallel_config": parallel_config,
"vae_use_slicing": kwargs.get("vae_use_slicing", False),
"vae_use_tiling": kwargs.get("vae_use_tiling", False),
"cache_backend": kwargs.get("cache_backend", "none"),
"cache_config": kwargs.get("cache_config", None),
},
"final_output": True,
"final_output_type": "image",
}
]
default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
return default_stage_cfg
for stage_id, stage in enumerate(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,
)
stage.init_stage_worker(
model,
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("[AsyncOrchestrator] Stage-%s process started", stage_id)
time.sleep(self._init_sleep_seconds)
def _process_stage_ready(self, stage: OmniStage, stage_id: int, result: dict[str, Any]) -> None:
# Store vllm_config received from worker process (may be None for diffusion stages)
vllm_config = result.get("vllm_config")
if vllm_config is not None:
stage.set_vllm_config(vllm_config)
tokenizer = result.get("tokenizer")
if tokenizer is not None:
stage.set_tokenizer(tokenizer)
is_tracing_enabled = result.get("is_tracing_enabled")
if is_tracing_enabled is not None:
stage.set_is_tracing_enabled(is_tracing_enabled)
super()._process_stage_ready(stage, stage_id, result)
def _wait_for_stages_ready(self, timeout: int = 120) -> None:
"""Wait for all stages to report readiness."""
deadline = time.time() + max(0, int(timeout))
num_stages = len(self.stage_list)
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
result = stage.try_collect()
if result is None:
continue
progressed = True
if result.get("type") == "stage_ready":
self._stages_ready.add(stage_id)
# Store vllm_config received from worker process (may be None for diffusion stages)
vllm_config = result.get("vllm_config")
if vllm_config is not None:
stage.set_vllm_config(vllm_config)
tokenizer = result.get("tokenizer")
if tokenizer is not None:
stage.set_tokenizer(tokenizer)
is_tracing_enabled = result.get("is_tracing_enabled")
if is_tracing_enabled is not None:
stage.set_is_tracing_enabled(is_tracing_enabled)
logger.info("[AsyncOrchestrator] Stage-%s reported ready", stage_id)
else:
# No user data should arrive before seeding; ignore other messages
pass
if not progressed:
time.sleep(0.01)
if len(self._stages_ready) < num_stages:
not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
super()._wait_for_stages_ready(timeout)
for stage in self.stage_list:
if stage.vllm_config is not None and stage.tokenizer is not None:
try:
vllm_config = stage.vllm_config
tokenizer = stage.tokenizer
# Initialize input_processor
self.input_processor = OmniInputProcessor(
vllm_config=vllm_config,
tokenizer=tokenizer,
)
# Initialize model_config
self.model_config = vllm_config.model_config
# Initialize io_processor
io_processor_plugin = self.model_config.io_processor_plugin
self.io_processor = get_io_processor(vllm_config, io_processor_plugin)
logger.info(
f"[{self._name}] Initialized input_processor, "
f"io_processor, and model_config from stage-{stage.stage_id}",
)
break
except Exception as e:
logger.warning(
f"[{self._name}] Failed to initialize processors from stage-{stage.stage_id}: {e}",
)
# If no LLM stage found, set processors to None
if not hasattr(self, "input_processor") or self.input_processor is None:
logger.warning(
"[AsyncOrchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s",
len(self._stages_ready),
num_stages,
not_ready,
f"[{self._name}] No LLM stage found, processors will not be available. "
"This may cause issues with OpenAIServingModels."
)
# Provide actionable suggestions before shutdown
try:
suggestions = [
"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 (init_sleep_seconds or call-site timeout).",
]
logger.error(
"[AsyncOrchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s",
"\n- ".join(suggestions),
)
except Exception:
# Best-effort logging of suggestions
logger.error(
"[AsyncOrchestrator] Stage initialization failed and an error occurred while logging suggestions",
)
elif len(self._stages_ready) == num_stages:
logger.info("[AsyncOrchestrator] All stages initialized successfully")
# Initialize input_processor, io_processor, and model_config for API server compatibility
# Find the first LLM stage (with vllm_config) to get vllm_config and tokenizer
for stage in self.stage_list:
if stage.vllm_config is not None and stage.tokenizer is not None:
try:
vllm_config = stage.vllm_config
tokenizer = stage.tokenizer
# Initialize input_processor
self.input_processor = OmniInputProcessor(
vllm_config=vllm_config,
tokenizer=tokenizer,
)
# Initialize model_config
self.model_config = vllm_config.model_config
# Initialize io_processor
io_processor_plugin = self.model_config.io_processor_plugin
self.io_processor = get_io_processor(vllm_config, io_processor_plugin)
logger.info(
"[AsyncOrchestrator] Initialized input_processor, "
"io_processor, and model_config from stage-%s",
stage.stage_id,
)
break
except Exception as e:
logger.warning(
"[AsyncOrchestrator] Failed to initialize processors from stage-%s: %s",
stage.stage_id,
e,
)
# If no LLM stage found, set processors to None
if not hasattr(self, "input_processor") or self.input_processor is None:
logger.warning(
"[AsyncOrchestrator] No LLM stage found, processors will not be available. "
"This may cause issues with OpenAIServingModels."
)
self.input_processor = None
self.io_processor = None
self.model_config = None
def _initialize_stages(
self,
model: str,
init_sleep_seconds: int,
shm_threshold_bytes: int,
init_timeout: int,
) -> None:
self.stage_list: list[OmniStage] = []
# 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)
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("[Orchestrator] Loaded %d stages", len(self.stage_list))
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_in_queues: list[mp.Queue] = []
self._stage_out_queues: list[mp.Queue] = []
self._init_sleep_seconds = max(0, int(init_sleep_seconds))
self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
self._start_stages(model)
# Wait for all stages to report readiness before seeding
self._stages_ready: set[int] = set()
self._wait_for_stages_ready(timeout=init_timeout)
def _start_stages(self, model: str) -> None:
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(self.stage_list):
# Use unbounded queues to avoid deadlock when seeding many requests
in_q = self._queue_cls()
out_q = self._queue_cls()
self._stage_in_queues.append(in_q)
self._stage_out_queues.append(out_q)
# Attach queues and start Stage-owned worker process
stage.attach_queues(in_q, out_q)
# Build connectors config for this stage
stage_connectors_config = get_stage_connector_config(
self.omni_transfer_config,
stage_id,
)
stage.init_stage_worker(
model,
is_async=True,
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("[Orchestrator] Stage-%s process started", stage_id)
time.sleep(self._init_sleep_seconds)
self.input_processor = None
self.io_processor = None
self.model_config = None
def close(self) -> None:
"""Close all stage processes and clean up resources.
Sends shutdown signals to all stage input queues and stops
all stage worker processes. This method should be called
when done using the AsyncOmni instance.
"""
for q in self._stage_in_queues:
try:
q.put_nowait(None)
except Exception as e:
logger.warning(
"[Orchestrator] Failed to send shutdown signal to \
stage input queue: %s",
e,
)
for stage in self.stage_list:
try:
stage.stop_stage_worker()
except Exception as e:
logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
"""Close all stage processes and clean up resources."""
if self.output_handler is not None:
self.output_handler.cancel()
self.output_handler = None
try_close_ray(self._ray_pg)
def __del__(self) -> None: # best-effort
try:
self.close()
except Exception as e:
logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
super().close()
def shutdown(self):
"""Shutdown, cleaning up the background proc and IPC.
@@ -523,7 +231,7 @@ class AsyncOmni:
async with self._pause_cond:
await self._pause_cond.wait_for(lambda: not self._paused)
logger.debug("[AsyncOrchestrator] generate() called")
logger.debug(f"[{self._name}] generate() called")
# Start output handler on the first call to generate()
self._run_output_handler()
@@ -577,7 +285,7 @@ class AsyncOmni:
_wall_start_ts,
)
# Seed stage-0 queue with all requests
logger.debug("[AsyncOrchestrator] Seeding request into stage-0")
logger.debug(f"[{self._name}] Seeding request into stage-0")
req_state = ClientRequestState(request_id)
self.request_states[request_id] = req_state
@@ -592,9 +300,9 @@ class AsyncOmni:
}
self.stage_list[0].submit(task)
_req_start_ts[request_id] = time.time()
logger.debug("[AsyncOrchestrator] Enqueued request %s to stage-0", request_id)
logger.debug(f"[{self._name}] Enqueued request {request_id} to stage-0")
logger.debug("[AsyncOrchestrator] Entering scheduling loop: stages=%d", num_stages)
logger.debug(f"[{self._name}] Entering scheduling loop: stages={num_stages}")
for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
result = await req_state.queue.get()
assert stage_id == req_state.stage_id
@@ -602,10 +310,7 @@ class AsyncOmni:
req_id = result.get("request_id")
if "error" in result:
logger.error(
"Stage %s error on request %s: %s",
stage_id,
req_id,
result["error"],
f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
)
raise RuntimeError(result) # Request Finished due to error
@@ -618,23 +323,16 @@ class AsyncOmni:
metrics.on_stage_metrics(stage_id, req_id, _m)
except Exception as e:
logger.exception(
"[AsyncOrchestrator] Failed to process metrics for stage %s, req %s: %s",
stage_id,
req_id,
e,
f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}",
)
logger.debug(
"[AsyncOrchestrator] Stage-%s completed request %s; forwarding or finalizing",
stage_id,
req_id,
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(
"[AsyncOrchestrator] Request %s finalized at stage-%s",
req_id,
stage_id,
f"[{self._name}] Request {req_id} finalized at stage-{stage_id}",
)
# End-to-end timing and time-per-token for final output
@@ -650,10 +348,7 @@ class AsyncOmni:
)
except Exception as e:
logger.exception(
"[AsyncOrchestrator] Finalize request handling error for req %s at stage %s: %s",
req_id,
stage_id,
e,
f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}",
)
if isinstance(engine_outputs, list):
@@ -709,93 +404,26 @@ class AsyncOmni:
# because continuing would cause the request to be silently dropped
# and the orchestrator to hang waiting for completion.
error_msg = (
f"[AsyncOrchestrator] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
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.error(error_msg)
raise RuntimeError(error_msg)
logger.debug(
"[AsyncOrchestrator] Forwarded request %s to stage-%s",
req_id,
next_stage_id,
)
logger.debug(f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}")
else:
logger.debug("[AsyncOrchestrator] Request %s fully completed", req_id)
logger.debug(f"[{self._name}] Request {req_id} fully completed")
logger.debug("[AsyncOrchestrator] All requests completed")
logger.debug(f"[{self._name}] All requests completed")
# Summarize and print stats
try:
summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
except Exception as e:
logger.exception("[AsyncOrchestrator] Failed to build/log summary: %s", e)
logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
finally:
self.request_states.pop(request_id, None)
def _wait_for_stages_ready(self, timeout: int = 120) -> None:
num_stages = len(self.stage_list)
while len(self._stages_ready) < num_stages:
progressed = False
for stage_id, stage in enumerate(self.stage_list):
if stage_id in self._stages_ready:
continue
result = stage.try_collect()
if result is None:
continue
progressed = True
if result.get("type") == "stage_ready":
self._stages_ready.add(stage_id)
# Store vllm_config received from worker process
vllm_config = result.get("vllm_config")
if vllm_config is not None:
stage.set_vllm_config(vllm_config)
tokenizer = result.get("tokenizer")
if tokenizer is not None:
stage.set_tokenizer(tokenizer)
# input_preprocessor = result.get("input_preprocessor")
# if input_preprocessor is not None:
# stage.set_input_preprocessor(input_preprocessor)
is_tracing_enabled = result.get("is_tracing_enabled")
if is_tracing_enabled is not None:
stage.set_is_tracing_enabled(is_tracing_enabled)
logger.debug("[Orchestrator] Stage-%s reported ready", stage_id)
else:
# No user data should arrive before seeding; ignore other messages
pass
if not progressed:
time.sleep(0.01)
if len(self._stages_ready) < num_stages:
not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
logger.warning(
"[Orchestrator] Initialization timeout: only %s/%s stages are \
ready; not ready: %s",
len(self._stages_ready),
num_stages,
not_ready,
)
# Provide actionable suggestions before shutdown
try:
suggestions = [
"Verify GPU/device assignment in config (runtime.devices) is \
correct.",
"Check GPU/host memory availability; reduce model or batch size if needed.", # noqa: E501
"Check model weights path and network reachability (if loading remotely).", # noqa: E501
"Increase initialization wait time (init_sleep_seconds or \
call-site timeout).",
]
logger.error(
"[Orchestrator] Stage initialization failed, shutting down. \
Suggestions:\n- %s",
"\n- ".join(suggestions),
)
except Exception:
# Best-effort logging of suggestions
logger.error(
"[Orchestrator] Stage initialization failed and an error \
occurred while logging suggestions",
)
def _run_output_handler(self) -> None:
if self.output_handler is not None:
return
@@ -821,10 +449,8 @@ class AsyncOmni:
req_state = request_states.get(req_id)
if req_state is None:
logger.debug(
"[Orchestrator] Request may have been aborted; \
dropping output for req %s at stage-%s ",
req_id,
stage_id,
f"[{self._name}] Request may have been aborted; \
dropping output for req {req_id} at stage-{stage_id}"
)
continue
await req_state.queue.put(result)
@@ -854,6 +480,14 @@ class AsyncOmni:
def errored(self) -> bool:
return not self.is_running
@property
def _name(self) -> str:
return "AsyncOrchestrator"
@property
def is_async(self) -> bool:
return True
@property
def dead_error(self) -> BaseException:
return EngineDeadError()
vllm_omni/entrypoints/omni.py
+139 -121
View File
@@ -54,8 +54,8 @@ def omni_snapshot_download(model_id) -> str:
return _dummy_snapshot_download(model_id)
class Omni:
"""Unified entrypoint for both LLM and Diffusion models for better usability.
class OmniBase:
"""Base class for serving Omni models.
Args:
*args: Variable length argument list.
@@ -75,11 +75,6 @@ class Omni:
- 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, *args: Any, **kwargs: dict[str, Any]) -> None:
@@ -106,12 +101,38 @@ class Omni:
logger.info(f"Initializing stages for model: {model}")
self._initialize_stages(model, kwargs)
def _initialize_stages(self, model: str, kwargs: dict[str, Any]) -> None:
"""Initialize stage list management.
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"])
# 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),
"final_output": True,
"final_output_type": "image",
}
]
default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
return default_stage_cfg
Each stage will create appropriate instance (OmniLLM or OmniDiffusion)
based on stage_type in YAML config (handled in omni_stage.py).
"""
def _initialize_stages(self, model: str, kwargs: dict[str, Any]) -> None:
"""Initialize stage list management."""
init_sleep_seconds = kwargs.get("init_sleep_seconds", 20)
shm_threshold_bytes = kwargs.get("shm_threshold_bytes", 65536)
init_timeout = kwargs.get("init_timeout", 300)
@@ -126,31 +147,7 @@ class Omni:
self.config_path = resolve_model_config_path(model)
self.stage_configs = load_stage_configs_from_model(model)
if not self.stage_configs:
# TODO: hack here, convert dtype to string to avoid non-premitive omegaconf create error.
if "dtype" in kwargs:
kwargs["dtype"] = str(kwargs["dtype"])
# 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}"
logger.info(f"model: {model}, kwargs: {kwargs}")
default_stage_cfg = [
{
"stage_id": 0,
"stage_type": "diffusion",
"runtime": {
"process": True,
"devices": devices,
"max_batch_size": 1,
},
"engine_args": OmegaConf.create(kwargs),
"final_output": True,
"final_output_type": "image",
}
]
default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
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
@@ -182,7 +179,7 @@ class Omni:
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("[Orchestrator] Loaded %d stages", len(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()
@@ -218,6 +215,7 @@ class Omni:
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,
@@ -226,9 +224,13 @@ class Omni:
ray_placement_group=self._ray_pg,
)
logger.debug("[Orchestrator] Stage-%s process started", stage_id)
logger.debug(f"[{self._name}] Stage-{stage_id} process started")
time.sleep(self._init_sleep_seconds)
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."""
deadline = time.time() + max(0, int(timeout))
@@ -243,40 +245,100 @@ class Omni:
continue
progressed = True
if result.get("type") == "stage_ready":
self._stages_ready.add(stage_id)
logger.info("[Orchestrator] Stage-%s reported ready", stage_id)
else:
# No user data should arrive before seeding; ignore other messages
pass
self._process_stage_ready(stage, stage_id, result)
if not progressed:
time.sleep(0.01)
if len(self._stages_ready) < num_stages:
not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
logger.warning(
"[Orchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s",
len(self._stages_ready),
num_stages,
not_ready,
f"[{self._name}] Initialization timeout: only {len(self._stages_ready)}/{num_stages}"
f" stages are ready; not ready: {not_ready}.",
)
# Provide actionable suggestions before shutdown
try:
suggestions = [
"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 (init_sleep_seconds or call-site timeout).",
]
suggestions = "".join(
[
"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 (init_sleep_seconds or call-site timeout).",
]
)
logger.error(
"[Orchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s",
"\n- ".join(suggestions),
f"[{self._name}] Stage initialization failed, shutting down. Suggestions:\n- {suggestions}",
)
except Exception:
# Best-effort logging of suggestions
logger.error(
"[Orchestrator] Stage initialization failed and an error occurred while logging suggestions",
f"[{self._name}] Stage initialization failed and an error occurred while logging suggestions",
)
elif len(self._stages_ready) == num_stages:
logger.info("[Orchestrator] All stages initialized successfully")
logger.info(f"[{self._name}] All stages initialized successfully")
def close(self) -> None:
"""Close all stage processes and clean up resources."""
# Close stages if they exist (for LLM models)
if self.stage_list:
for q in self._stage_in_queues:
try:
q.put_nowait(None)
except Exception as e:
logger.warning(
f"[{self._name}] Failed to send shutdown signal to stage input queue: {e}",
)
for stage in self.stage_list:
try:
stage.stop_stage_worker()
except Exception as e:
logger.warning(f"[{self._name}] Failed to stop stage worker: {e}")
try_close_ray(self._ray_pg)
def __del__(self): # pragma: no cover - best effort cleanup
try:
self.close()
except Exception:
logger.debug(f"[{self._name}] __del__ close() raised", exc_info=True)
@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:
*args: Variable length argument list.
- args[0]: Model name or path to load.
**kwargs: Arbitrary keyword arguments.
- model: Model name or path to load (if not in args).
- 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.
- init_sleep_seconds: Number of seconds to sleep between starting
each stage process during initialization
- 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, *args: Any, **kwargs: dict[str, Any]) -> None:
super().__init__(*args, **kwargs)
def generate(self, *args: Any, **kwargs: dict[str, Any]) -> list[OmniRequestOutput]:
"""Generate outputs for the given prompts.
@@ -336,7 +398,7 @@ class Omni:
sampling_params_list: Any | Sequence[Any] | None = None,
) -> list[OmniRequestOutput]:
"""Run generation through all stages in the pipeline."""
logger.debug("[Orchestrator] generate() called")
logger.debug(f"[{self._name}] generate() called")
if sampling_params_list is None:
raise ValueError("sampling_params_list is required for pipelined generation")
@@ -388,7 +450,7 @@ class Omni:
)
# Seed stage-0 queue with all requests
logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts))
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()
@@ -401,7 +463,7 @@ class Omni:
}
self.stage_list[0].submit(task)
_req_start_ts[req_id] = time.time()
logger.debug("[Orchestrator] Enqueued request %s to stage-0", req_id)
logger.debug(f"[{self._name}] Enqueued request {req_id} to stage-0")
# For each stage, forward results to next stage; collect finals at the end
# We pipeline by continually polling output queues in stage order
@@ -410,9 +472,7 @@ class Omni:
total_requests = len(request_prompts)
logger.debug(
"[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d",
total_requests,
num_stages,
f"[{self._name}] Entering scheduling loop: total_requests={total_requests}, stages={num_stages}",
)
while completed_requests < total_requests:
made_progress = False
@@ -425,10 +485,7 @@ class Omni:
req_id = result.get("request_id")
if "error" in result:
logger.error(
"Stage %s error on request %s: %s",
stage_id,
req_id,
result["error"],
f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
)
continue
@@ -447,15 +504,10 @@ class Omni:
metrics.on_stage_metrics(stage_id, req_id, _m)
except Exception as e:
logger.exception(
"[Orchestrator] Failed to process metrics for stage %s, req %s: %s",
stage_id,
req_id,
e,
f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}",
)
logger.debug(
"[Orchestrator] Stage-%s completed request %s; forwarding or finalizing",
stage_id,
req_id,
f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing",
)
stage.set_engine_outputs(engine_outputs)
@@ -468,9 +520,7 @@ class Omni:
)
)
logger.debug(
"[Orchestrator] Request %s finalized at stage-%s",
req_id,
stage_id,
f"[{self._name}] Request {req_id} finalized at stage-{stage_id}",
)
# End-to-end timing and time-per-token for final output
@@ -486,10 +536,7 @@ class Omni:
)
except Exception as e:
logger.exception(
"[Orchestrator] Finalize request handling error for req %s at stage %s: %s",
req_id,
stage_id,
e,
f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}",
)
next_stage_id = stage_id + 1
@@ -499,10 +546,8 @@ class Omni:
next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
except Exception as e:
logger.exception(
"[Orchestrator] Process engine inputs error for req %s at stage %s: %s",
req_id,
next_stage_id,
e,
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]
@@ -526,59 +571,32 @@ class Omni:
if not sent_via_connector:
raise RuntimeError(
f"[Orchestrator] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
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(
"[Orchestrator] Forwarded request %s to stage-%s",
req_id,
next_stage_id,
f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}",
)
remaining_by_stage[next_stage_id] += 1
else:
completed_requests += 1
logger.debug(
"[Orchestrator] Request %s fully completed (%d/%d)",
req_id,
completed_requests,
total_requests,
f"[{self._name}] Request {req_id} fully completed ({completed_requests}/{total_requests})",
)
if not made_progress:
time.sleep(0.005)
logger.debug("[Orchestrator] All requests completed")
logger.debug(f"[{self._name}] All requests completed")
# 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("[Orchestrator] Failed to build/log summary: %s", e)
logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
return final_outputs
def close(self) -> None:
"""Close all stage processes and clean up resources."""
# Close stages if they exist (for LLM models)
if self.stage_list:
for q in self._stage_in_queues:
try:
q.put_nowait(None)
except Exception as e:
logger.warning(
"[Orchestrator] Failed to send shutdown signal to stage input queue: %s",
e,
)
for stage in self.stage_list:
try:
stage.stop_stage_worker()
except Exception as e:
logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
try_close_ray(self._ray_pg)
def __del__(self): # pragma: no cover - best effort cleanup
try:
self.close()
except Exception:
logger.debug("[Orchestrator] __del__ close() raised", exc_info=True)
@property
def _name(self) -> str:
return "Orchestrator"