378 lines
14 KiB
Python
378 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import json
|
|
import os
|
|
import types
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass, fields, is_dataclass
|
|
from typing import Any, get_args, get_origin
|
|
|
|
from vllm.logger import init_logger
|
|
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
|
|
|
from vllm_omni.config.config_factory import with_trust_remote_code_override
|
|
from vllm_omni.entrypoints.stage_utils import _to_dict
|
|
from vllm_omni.inputs.data import OmniSamplingParams
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: str, omni_to: str) -> None:
|
|
"""Inject connector configuration into stage engine arguments."""
|
|
# Prepare omni_kv_config dict
|
|
omni_conf_dict = {}
|
|
try:
|
|
# Access engine_args safely (might be OmegaConf or dict)
|
|
existing_args = stage.engine_args
|
|
if hasattr(existing_args, "get"):
|
|
_oc = existing_args.get("omni_kv_config", None)
|
|
if _oc:
|
|
if hasattr(_oc, "items"): # dict-like
|
|
omni_conf_dict = dict(_oc)
|
|
else: # object?
|
|
omni_conf_dict = _to_dict(_oc)
|
|
except Exception:
|
|
omni_conf_dict = {}
|
|
|
|
# Inject connector info
|
|
omni_conf_dict["connector_config"] = omni_conn_cfg
|
|
omni_conf_dict["omni_from_stage"] = omni_from
|
|
omni_conf_dict["omni_to_stage"] = omni_to
|
|
|
|
# Write back to engine_args
|
|
try:
|
|
if hasattr(stage.engine_args, "__setitem__"):
|
|
stage.engine_args["omni_kv_config"] = omni_conf_dict
|
|
else:
|
|
setattr(stage.engine_args, "omni_kv_config", omni_conf_dict)
|
|
except Exception as e:
|
|
# Fallback for OmegaConf or similar if direct set fails?
|
|
logger.error(f"Failed to inject omni connector config into stage: {e}")
|
|
|
|
|
|
def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
|
|
"""Parse and validate the shape of per-stage JSON overrides."""
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, str):
|
|
try:
|
|
parsed = json.loads(value)
|
|
except json.JSONDecodeError as exc:
|
|
raise ValueError(f"--stage-overrides is not valid JSON: {exc}. Got: {value!r}") from exc
|
|
else:
|
|
parsed = value
|
|
|
|
if not isinstance(parsed, dict):
|
|
raise ValueError(
|
|
"--stage-overrides must be a JSON object mapping stage_id -> overrides, "
|
|
f"got {type(parsed).__name__}: {parsed!r}"
|
|
)
|
|
for stage_id, overrides in parsed.items():
|
|
if not isinstance(stage_id, str) or not stage_id.isascii() or not stage_id.isdigit():
|
|
raise ValueError(
|
|
f"--stage-overrides keys must be non-negative integer stage ids (as strings), got {stage_id!r}"
|
|
)
|
|
if not isinstance(overrides, dict):
|
|
raise ValueError(
|
|
f"--stage-overrides[{stage_id!r}] must be an object, got {type(overrides).__name__}: {overrides!r}"
|
|
)
|
|
|
|
return parsed
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StageConfigInputs:
|
|
"""Normalized inputs shared by all stage-config loading paths."""
|
|
|
|
model: str
|
|
kwargs: dict[str, Any]
|
|
trust_remote_code: bool | None
|
|
deploy_config_path: str | None
|
|
stage_overrides: dict[str, dict[str, Any]] | None
|
|
strategy_config_path: str | None
|
|
|
|
|
|
def prepare_stage_config_inputs(
|
|
model: str,
|
|
kwargs: Mapping[str, Any],
|
|
*,
|
|
trust_remote_code: bool | None,
|
|
snapshot_model: bool = False,
|
|
) -> StageConfigInputs:
|
|
"""Normalize model/config arguments before resolving stage configs.
|
|
|
|
The standard API worker, multi-API parent, and headless entrypoint must
|
|
apply the same legacy-key filtering, trust-remote-code precedence, and
|
|
config-file extraction before calling :func:`resolve_omni_config`.
|
|
"""
|
|
resolved_model = model
|
|
if snapshot_model:
|
|
# Keep model materialization at the CLI boundary for the parent-owned
|
|
# launch path. The API worker's OmniBase has already done this work.
|
|
from vllm_omni.entrypoints.omni_base import omni_snapshot_download
|
|
|
|
resolved_model = omni_snapshot_download(model)
|
|
|
|
resolved_kwargs = dict(kwargs)
|
|
for legacy_arg in ("stage_configs_path", "stage_configs"):
|
|
if legacy_arg in resolved_kwargs:
|
|
raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.")
|
|
# Frontend process count is not a stage engine override.
|
|
resolved_kwargs.pop("api_server_count", None)
|
|
resolved_kwargs.pop("disable_log_stats", None)
|
|
if resolved_kwargs.get("diffusion_streaming_output") and resolved_kwargs.get("streaming_output") is None:
|
|
resolved_kwargs["streaming_output"] = True
|
|
resolved_kwargs.pop("model_tag", None)
|
|
|
|
resolved_kwargs = with_trust_remote_code_override(resolved_kwargs, trust_remote_code)
|
|
deploy_config_path = resolved_kwargs.pop("deploy_config", None)
|
|
strategy_config_path = resolved_kwargs.pop("strategy_config", None)
|
|
stage_overrides = parse_stage_overrides(resolved_kwargs.pop("stage_overrides", None))
|
|
return StageConfigInputs(
|
|
model=resolved_model,
|
|
kwargs=resolved_kwargs,
|
|
trust_remote_code=trust_remote_code,
|
|
deploy_config_path=deploy_config_path,
|
|
stage_overrides=stage_overrides,
|
|
strategy_config_path=strategy_config_path,
|
|
)
|
|
|
|
|
|
def get_final_stage_id_for_e2e(
|
|
output_modalities: list[str] | None, default_modalities: list[str], stage_list: list
|
|
) -> int:
|
|
"""Get the final stage id for e2e.
|
|
|
|
Args:
|
|
stage_list: List of stage configurations
|
|
|
|
Returns:
|
|
Final stage id for e2e
|
|
"""
|
|
last_stage_id = len(stage_list) - 1
|
|
if output_modalities is not None:
|
|
prompt_modalities = []
|
|
for modality in output_modalities:
|
|
if modality not in default_modalities:
|
|
logger.warning(f"Invalid output modality: {modality}, ignoring it")
|
|
# TODO: if user specifies unsupported modalities, invalid it and raise an error
|
|
continue
|
|
prompt_modalities.append(modality)
|
|
output_modalities = prompt_modalities
|
|
else:
|
|
output_modalities = default_modalities
|
|
|
|
try:
|
|
final_stage_id_for_e2e = last_stage_id
|
|
for _sid in range(last_stage_id, -1, -1):
|
|
if (
|
|
getattr(stage_list[_sid], "final_output", False)
|
|
and stage_list[_sid].final_output_type in output_modalities
|
|
):
|
|
final_stage_id_for_e2e = _sid
|
|
break
|
|
except Exception as e:
|
|
logger.debug(
|
|
"[Orchestrator] Failed to determine final stage for E2E; \
|
|
falling back to last: %s",
|
|
e,
|
|
exc_info=True,
|
|
)
|
|
final_stage_id_for_e2e = last_stage_id
|
|
|
|
return final_stage_id_for_e2e
|
|
|
|
|
|
def filter_dataclass_kwargs(cls: Any, kwargs: dict) -> dict:
|
|
"""Filter kwargs to only include fields defined in the dataclass.
|
|
|
|
Args:
|
|
cls: Dataclass type
|
|
kwargs: Keyword arguments to filter
|
|
|
|
Returns:
|
|
Filtered keyword arguments containing only valid dataclass fields
|
|
"""
|
|
if not is_dataclass(cls):
|
|
raise ValueError(f"{cls} is not a dataclass")
|
|
if not isinstance(kwargs, dict):
|
|
raise ValueError("kwargs must be a dictionary")
|
|
|
|
def _filter_value(value: Any, annotation: Any) -> Any:
|
|
"""Recursively filter nested dict/list values based on dataclass annotations."""
|
|
if annotation is None:
|
|
return value
|
|
|
|
origin = get_origin(annotation)
|
|
if origin is None:
|
|
if isinstance(annotation, type) and is_dataclass(annotation) and isinstance(value, dict):
|
|
return filter_dataclass_kwargs(annotation, value)
|
|
return value
|
|
|
|
if origin in (list, tuple, set):
|
|
args = get_args(annotation)
|
|
inner = args[0] if args else None
|
|
if isinstance(value, list | tuple | set):
|
|
return type(value)(_filter_value(v, inner) for v in value)
|
|
return value
|
|
|
|
if origin is dict:
|
|
args = get_args(annotation)
|
|
val_type = args[1] if len(args) > 1 else None
|
|
if isinstance(value, dict):
|
|
return {k: _filter_value(v, val_type) for k, v in value.items()}
|
|
return value
|
|
|
|
if origin is types.UnionType or origin is getattr(types, "UnionType", None):
|
|
for arg in get_args(annotation):
|
|
if isinstance(arg, type) and is_dataclass(arg) and isinstance(value, dict):
|
|
return filter_dataclass_kwargs(arg, value)
|
|
# Try container-style filtering for union members
|
|
filtered = _filter_value(value, arg)
|
|
if filtered is not value:
|
|
return filtered
|
|
return value
|
|
|
|
return value
|
|
|
|
valid_fields = {f.name: f for f in fields(cls) if f.init}
|
|
filtered_kwargs = {}
|
|
for k, v in kwargs.items():
|
|
if k not in valid_fields:
|
|
logger.warning(
|
|
"Dropping unknown %s field %r (not declared on the dataclass)",
|
|
cls.__name__,
|
|
k,
|
|
)
|
|
continue
|
|
field = valid_fields[k]
|
|
filtered_kwargs[k] = _filter_value(v, field.type)
|
|
|
|
return filtered_kwargs
|
|
|
|
|
|
# The following code detects if the process is running in a container and if
|
|
# PID host is available. If so, we can use process-scoped memory tracking;
|
|
# otherwise we need sequential init locks.
|
|
|
|
|
|
def _read_text(path: str) -> str | None:
|
|
try:
|
|
with open(path, encoding="utf-8", errors="replace") as f:
|
|
return f.read()
|
|
except (FileNotFoundError, PermissionError, OSError):
|
|
return None
|
|
|
|
|
|
def in_container() -> bool:
|
|
# Common Docker signal
|
|
if os.path.exists("/.dockerenv"):
|
|
return True
|
|
|
|
# cgroup markers (works for Docker/containerd/K8s/Podman in many setups)
|
|
cg = _read_text("/proc/1/cgroup") or ""
|
|
markers = ("docker", "containerd", "kubepods", "libpod", "podman")
|
|
return any(m in cg for m in markers)
|
|
|
|
|
|
def has_pid_host() -> bool | None:
|
|
"""
|
|
Returns:
|
|
True -> very likely running with --pid=host (host PID namespace)
|
|
False -> very likely isolated PID namespace (default)
|
|
None -> cannot determine
|
|
"""
|
|
# Strong signal: in host pid namespace, PID 2 is usually kthreadd
|
|
comm2 = _read_text("/proc/2/comm")
|
|
if comm2 is not None:
|
|
comm2 = comm2.strip()
|
|
if comm2 == "kthreadd":
|
|
return True
|
|
# If PID 2 exists and is NOT kthreadd, we're almost certainly not in host pid ns
|
|
return False
|
|
|
|
# Fallback: check for other low-numbered kernel threads (best-effort)
|
|
for pid, name in [(3, "rcu_gp"), (4, "rcu_par_gp"), (10, "ksoftirqd/0")]:
|
|
comm = _read_text(f"/proc/{pid}/comm")
|
|
if comm is not None:
|
|
if comm.strip() == name:
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
return False
|
|
|
|
|
|
def detect_pid_host() -> bool:
|
|
ic = in_container()
|
|
if not ic:
|
|
return True
|
|
|
|
return has_pid_host() is True
|
|
|
|
|
|
### Helpers for handling delta messages
|
|
def coerce_param_message_types(params: list[OmniSamplingParams], is_streaming: bool):
|
|
"""Iterate over the sampling params and convert to the message types
|
|
to DELTA messages, if streaming is enabled, or FINAL_ONLY if
|
|
it's disabled, while respecting `.skip_clone` on the params.
|
|
|
|
This is needed to avoid emitting redundant multimodal data.
|
|
"""
|
|
# Coerce vLLM's default output kinds as needed to handle streaming
|
|
# (i.e., DELTA output kind). Note that this is only applied to non
|
|
# Diffusion sampling params.
|
|
#
|
|
# NOTE: Hidden states will still be passed between stages.
|
|
for idx, sp in enumerate(params):
|
|
# For OmniDiffusionParams don't set output kind
|
|
if isinstance(sp, SamplingParams):
|
|
params[idx] = maybe_coerce_to_message_type(sp, is_streaming)
|
|
return params
|
|
|
|
|
|
def maybe_coerce_to_message_type(params: SamplingParams, is_streaming: bool):
|
|
"""If this is a CUMULATIVE message, coerce it to DELTA if streaming, otherwise FINAL_ONLY."""
|
|
target_type = RequestOutputKind.DELTA if is_streaming else RequestOutputKind.FINAL_ONLY
|
|
if params.output_kind == target_type:
|
|
return params
|
|
elif is_streaming and params.output_kind == RequestOutputKind.FINAL_ONLY:
|
|
logger.debug("Coercing FINAL_ONLY output to DELTA for streaming")
|
|
elif not is_streaming and params.output_kind == RequestOutputKind.DELTA:
|
|
logger.debug("Coercing DELTA output to FINAL_ONLY for non-streaming")
|
|
|
|
if not params.skip_clone:
|
|
params = params.clone()
|
|
params.skip_clone = True
|
|
params.output_kind = target_type
|
|
return params
|
|
|
|
|
|
class PureDiffusionLauncherAdapter:
|
|
"""vLLM launcher compatibility shim for pure-diffusion mode.
|
|
|
|
The upstream launcher's shutdown path reads
|
|
``app.state.engine_client.vllm_config.shutdown_timeout``
|
|
(vllm/entrypoints/launcher.py), but ``AsyncOmni.vllm_config`` returns
|
|
``None`` when the pipeline has no comprehension stage (pure diffusion),
|
|
which crashes ``handle_shutdown`` with AttributeError and hangs server
|
|
teardown (workers force-killed, spurious resource_tracker noise).
|
|
|
|
This adapter only overrides the ``vllm_config`` property with a minimal
|
|
fallback carrying ``shutdown_timeout`` and forwards every other attribute
|
|
to the wrapped engine client, so the pure-diffusion detection
|
|
(``get_vllm_config()`` still returns ``None``) is unaffected.
|
|
"""
|
|
|
|
def __init__(self, engine_client: Any, shutdown_timeout: float) -> None:
|
|
object.__setattr__(self, "_wrapped", engine_client)
|
|
object.__setattr__(self, "_shutdown_timeout", float(shutdown_timeout))
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
return getattr(self._wrapped, name)
|
|
|
|
@property
|
|
def vllm_config(self) -> Any:
|
|
return types.SimpleNamespace(shutdown_timeout=self._shutdown_timeout)
|