Hongsheng Liu
d4096e01fe
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
908 lines
37 KiB
Python
908 lines
37 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 import Counter
|
||
from dataclasses import fields, is_dataclass
|
||
from pathlib import Path
|
||
from typing import Any, get_args, get_origin
|
||
|
||
from vllm.logger import init_logger
|
||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||
from vllm.transformers_utils.config import get_config, get_hf_file_to_dict
|
||
from vllm.transformers_utils.repo_utils import file_or_path_exists
|
||
|
||
from vllm_omni.config.config_factory import (
|
||
StageConfigFactory,
|
||
_materialize_object_storage_configs,
|
||
_name_match_candidate,
|
||
with_trust_remote_code_override,
|
||
)
|
||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
||
from vllm_omni.config.stage_config import _DEPLOY_DIR
|
||
from vllm_omni.config.yaml_util import create_config, load_yaml_config
|
||
from vllm_omni.diffusion.utils.hf_utils import (
|
||
_looks_like_dreamzero,
|
||
get_diffusion_model_index,
|
||
resolve_native_diffusion_model_class,
|
||
)
|
||
from vllm_omni.entrypoints.stage_utils import _to_dict
|
||
from vllm_omni.inputs.data import OmniSamplingParams
|
||
from vllm_omni.platforms import current_omni_platform
|
||
|
||
# Get the project root directory (2 levels up from this file)
|
||
PROJECT_ROOT = Path(__file__).parent.parent.parent
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
_DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = {
|
||
"GlmImagePipeline": "glm_image",
|
||
}
|
||
|
||
|
||
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 _try_get_class_name_from_diffusers_config(model: str) -> str | None:
|
||
"""Try to get class name from diffusers model configuration files.
|
||
|
||
Args:
|
||
model: Model name or path
|
||
|
||
Returns:
|
||
Model type string if found, None otherwise
|
||
"""
|
||
model_index = get_diffusion_model_index(model)
|
||
if model_index and isinstance(model_index, dict) and "_class_name" in model_index:
|
||
logger.debug(f"Found Diffusers model type '{model_index['_class_name']}'")
|
||
return model_index["_class_name"]
|
||
|
||
return None
|
||
|
||
|
||
def _filter_dict_like_object(obj: dict | Any) -> dict:
|
||
"""Filter dict-like object by removing callables and recursively converting values.
|
||
|
||
Converts dict-like objects to regular dicts while filtering out callable values
|
||
that are incompatible with OmegaConf. Recursively processes values through
|
||
_convert_dataclasses_to_dict for nested object conversion.
|
||
|
||
Args:
|
||
obj: Dict or dict-like object to filter
|
||
|
||
Returns:
|
||
Regular dict with callables filtered out and values recursively converted
|
||
|
||
Raises:
|
||
TypeError: If obj doesn't support .items() method
|
||
ValueError: If dict conversion fails unexpectedly
|
||
"""
|
||
|
||
def _is_callable_value(value: Any) -> bool:
|
||
if callable(value):
|
||
return True
|
||
return isinstance(
|
||
value,
|
||
types.FunctionType | types.MethodType | types.BuiltinFunctionType | types.BuiltinMethodType,
|
||
)
|
||
|
||
result = {}
|
||
filtered_keys = []
|
||
for k, v in obj.items():
|
||
# Preserve class objects by converting to a fully qualified name string
|
||
# so callers that resolve via import path (e.g. custom_pipeline_args.pipeline_class)
|
||
# still work after OmegaConf round-trip.
|
||
if isinstance(v, type):
|
||
module = getattr(v, "__module__", None)
|
||
qualname = getattr(v, "__qualname__", getattr(v, "__name__", None))
|
||
if module and qualname and module != "builtins":
|
||
result[k] = f"{module}.{qualname}"
|
||
else:
|
||
result[k] = qualname
|
||
elif _is_callable_value(v):
|
||
filtered_keys.append(str(k))
|
||
else:
|
||
result[k] = _convert_dataclasses_to_dict(v)
|
||
if filtered_keys:
|
||
logger.warning(
|
||
f"Filtered out {len(filtered_keys)} callable object(s) from base_engine_args "
|
||
f"that are not compatible with OmegaConf: {filtered_keys}. "
|
||
)
|
||
return result
|
||
|
||
|
||
def _convert_dataclasses_to_dict(obj: Any) -> Any:
|
||
"""Recursively convert non-serializable objects to OmegaConf-compatible types.
|
||
|
||
This is needed because OmegaConf cannot handle:
|
||
- Dataclass objects with Literal type annotations (e.g., StructuredOutputsConfig)
|
||
- Counter objects (from collections or vllm.utils)
|
||
- Set objects
|
||
- Callable objects (functions, methods, etc.)
|
||
- Other non-primitive types
|
||
"""
|
||
# IMPORTANT: Check Counter BEFORE dict, since Counter is a subclass of dict
|
||
# Handle Counter objects (convert to dict)
|
||
# Check by class name first to catch both collections.Counter and vllm.utils.Counter
|
||
if hasattr(obj, "__class__") and obj.__class__.__name__ == "Counter":
|
||
try:
|
||
return dict(obj)
|
||
except (TypeError, ValueError):
|
||
# If Counter can't be converted to dict, return empty dict
|
||
return {}
|
||
# Also check isinstance for collections.Counter (must be before dict check)
|
||
if isinstance(obj, Counter):
|
||
return dict(obj)
|
||
# Handle set objects (convert to list)
|
||
if isinstance(obj, set):
|
||
return list(obj)
|
||
# Handle dataclass objects
|
||
# Use field iteration instead of asdict() to:
|
||
# 1. Only include init fields (non-init fields cause "unexpected kwarg" errors)
|
||
# 2. Skip None values matching field defaults (avoids Pydantic validation
|
||
# failures when None is explicitly passed for non-Optional typed fields,
|
||
# e.g. CompilationConfig.cudagraph_capture_sizes: list[int] = None)
|
||
if is_dataclass(obj) and not isinstance(obj, type):
|
||
result = {}
|
||
for f in fields(obj):
|
||
if not f.init:
|
||
continue
|
||
value = getattr(obj, f.name)
|
||
if value is None and f.default is None:
|
||
continue
|
||
result[f.name] = _convert_dataclasses_to_dict(value)
|
||
return result
|
||
# Handle dictionaries (recurse into values) and filter out callables(cause error in OmegaConf.create)
|
||
# Note: This must come AFTER Counter check since Counter is a dict subclass
|
||
if isinstance(obj, dict):
|
||
return _filter_dict_like_object(obj)
|
||
# Preserve class objects by converting to a fully qualified name string.
|
||
if isinstance(obj, type):
|
||
module = getattr(obj, "__module__", None)
|
||
qualname = getattr(obj, "__qualname__", getattr(obj, "__name__", None))
|
||
if module and qualname and module != "builtins":
|
||
return f"{module}.{qualname}"
|
||
return qualname
|
||
# Handle callable objects (functions, methods, etc.) - skip them
|
||
# Note: This comes after dict/list checks to avoid misclassifying dict-like objects
|
||
if callable(obj):
|
||
return None
|
||
# Handle lists and tuples (recurse into items)
|
||
if isinstance(obj, list | tuple):
|
||
return type(obj)(_convert_dataclasses_to_dict(item) for item in obj if not callable(item))
|
||
# Try to convert any dict-like object (has keys/values methods) to dict
|
||
if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, str | bytes):
|
||
try:
|
||
return _filter_dict_like_object(obj)
|
||
except (TypeError, ValueError, AttributeError):
|
||
# If conversion fails, return as-is
|
||
return obj
|
||
# Primitive types and other objects that OmegaConf can handle
|
||
return obj
|
||
|
||
|
||
def _try_resolve_omni_model_type(model: str) -> str | None:
|
||
"""Try to resolve model_type for omni models with empty config.json.
|
||
|
||
Searches both the legacy ``stage_configs/*.yaml`` directory and the
|
||
migrated ``deploy/*.yaml`` directory for a stem that substring-matches
|
||
the model path basename (e.g. ``cosyvoice3`` in
|
||
``FunAudioLLM/Fun-CosyVoice3-0.5B-2512``). Only the basename is scanned
|
||
so URI segments such as the bucket name cannot select an unrelated
|
||
pipeline. The longest match wins so ``cosyvoice3`` beats ``cosyvoice``
|
||
and ``bagel_single_stage`` beats ``bagel``.
|
||
"""
|
||
model_lower = _name_match_candidate(model).lower().replace("-", "").replace("_", "")
|
||
best_match: str | None = None
|
||
best_len = 0
|
||
for subdir in ("model_executor/stage_configs", "deploy"):
|
||
config_dir = PROJECT_ROOT / "vllm_omni" / subdir
|
||
if not config_dir.exists():
|
||
continue
|
||
for config_file in sorted(config_dir.glob("*.yaml")):
|
||
candidate = config_file.stem.replace("-", "").replace("_", "")
|
||
if candidate and candidate in model_lower and len(candidate) > best_len:
|
||
best_match = config_file.stem
|
||
best_len = len(candidate)
|
||
return best_match
|
||
|
||
|
||
def _registry_default_deploy_path(model: str) -> str | None:
|
||
"""Return the registered pipeline's default deploy YAML path, if any.
|
||
|
||
A checkpoint's HF ``model_type`` is not always the registered Omni pipeline
|
||
key. MiniCPM-o 4.5, for example, reports ``minicpmo`` while its
|
||
predicate-selected pipeline is ``minicpmo_4_5``. Stage construction already
|
||
loads that pipeline's default deploy config, so connector discovery must
|
||
resolve the same path.
|
||
"""
|
||
pipeline_config = StageConfigFactory.get_pipeline_config(
|
||
model=model,
|
||
trust_remote_code=True,
|
||
)
|
||
if pipeline_config is None or pipeline_config.default_deploy_config_name is None:
|
||
return None
|
||
default_deploy_path = _DEPLOY_DIR / pipeline_config.default_deploy_config_name
|
||
if default_deploy_path.is_file():
|
||
return str(default_deploy_path)
|
||
return None
|
||
|
||
|
||
def resolve_model_config_path(model: str) -> str | None:
|
||
"""Resolve the stage/deploy config file path from the model name.
|
||
|
||
Resolves configuration path based on the model type and device type.
|
||
Order:
|
||
1. Device-specific stage config under ``stage_configs/{device_type}/``
|
||
2. If HF ``model_type`` is not an ``OMNI_PIPELINES`` key, the registered
|
||
pipeline's ``default_deploy_config_name`` (keeps connectors aligned with
|
||
stage construction when HF type and pipeline key differ)
|
||
3. ``deploy/{model_type}.yaml``
|
||
4. Legacy ``stage_configs/{model_type}.yaml``
|
||
5. Registered pipeline default deploy config (final fallthrough)
|
||
|
||
Args:
|
||
model: Model name or path (used to determine model_type)
|
||
|
||
Returns:
|
||
String path to the stage/deploy configuration file, or ``None`` if no
|
||
matching config file exists.
|
||
|
||
Raises:
|
||
ValueError: If model_type cannot be determined
|
||
"""
|
||
# Object-storage URIs are streamed by vLLM's Run:AI streamer only after
|
||
# each stage builds its ModelConfig, so config resolution here reads the
|
||
# local copy that vLLM-Omni materializes for such URIs. Name-based
|
||
# fallbacks keep the original string (the materialized path is a hash).
|
||
config_source = _materialize_object_storage_configs(model)
|
||
# Try to get config from standard transformers format first
|
||
try:
|
||
hf_config = get_config(config_source, trust_remote_code=True)
|
||
model_type = hf_config.model_type
|
||
except (ValueError, Exception):
|
||
# If standard transformers format fails, try diffusers format
|
||
native_model_class = resolve_native_diffusion_model_class(config_source)
|
||
if native_model_class is not None:
|
||
model_type = native_model_class
|
||
elif get_diffusion_model_index(config_source) is not None:
|
||
model_type = _try_get_class_name_from_diffusers_config(config_source)
|
||
if model_type is None:
|
||
raise ValueError(
|
||
f"Could not determine model_type for diffusers model: {model}. "
|
||
"Please ensure its Diffusers pipeline index contains '_class_name'"
|
||
)
|
||
elif file_or_path_exists(config_source, "config.json", revision=None):
|
||
# Try to read config.json manually for custom models like Bagel that fail get_config
|
||
# but have a valid config.json with model_type
|
||
try:
|
||
config_dict = get_hf_file_to_dict("config.json", config_source, revision=None)
|
||
if config_dict and "model_type" in config_dict:
|
||
model_type = config_dict["model_type"]
|
||
else:
|
||
# For models with empty config.json (e.g. CosyVoice3),
|
||
# try matching against registered omni stage configs.
|
||
model_type = _try_resolve_omni_model_type(model)
|
||
if model_type is None:
|
||
raise ValueError(f"config.json found but missing 'model_type' for model: {model}")
|
||
except Exception as e:
|
||
raise ValueError(f"Failed to read config.json for model: {model}. Error: {e}") from e
|
||
else:
|
||
# No config.json at repo root (e.g. GLM-TTS stores configs in
|
||
# subdirectories only). Try matching against registered deploy
|
||
# YAML filenames before giving up.
|
||
model_type = _try_resolve_omni_model_type(model)
|
||
if model_type is None:
|
||
raise ValueError(
|
||
f"Could not determine model_type for model: {model}. "
|
||
"Model is not in standard transformers or Diffusers format. "
|
||
f"Please ensure the model has proper configuration files with 'model_type' field"
|
||
)
|
||
|
||
default_config_path = current_omni_platform.get_default_stage_config_path()
|
||
if model_type == "vla" and _looks_like_dreamzero(config_source):
|
||
model_type = "dreamzero"
|
||
|
||
if model_type in _DIFFUSERS_CLASS_TO_CONFIG:
|
||
normalized_model_type = _DIFFUSERS_CLASS_TO_CONFIG[model_type]
|
||
else:
|
||
normalized_model_type = model_type.replace("-", "_")
|
||
model_type_str = f"{normalized_model_type}.yaml"
|
||
complete_config_path = PROJECT_ROOT / default_config_path / model_type_str
|
||
if os.path.exists(complete_config_path):
|
||
return str(complete_config_path)
|
||
|
||
# Prefer the registry default before deploy/<hf_model_type>.yaml when the
|
||
# HF type is not itself a pipeline key. Otherwise a future
|
||
# deploy/minicpmo.yaml would desync connectors from stages for MiniCPM-o 4.5.
|
||
if normalized_model_type not in OMNI_PIPELINES:
|
||
registry_path = _registry_default_deploy_path(model)
|
||
if registry_path is not None:
|
||
return registry_path
|
||
|
||
deploy_config_path = _DEPLOY_DIR / model_type_str
|
||
if os.path.exists(deploy_config_path):
|
||
return str(deploy_config_path)
|
||
|
||
stage_config_file = f"vllm_omni/model_executor/stage_configs/{normalized_model_type}.yaml"
|
||
stage_config_path = PROJECT_ROOT / stage_config_file
|
||
if os.path.exists(stage_config_path):
|
||
return str(stage_config_path)
|
||
|
||
return _registry_default_deploy_path(model)
|
||
|
||
|
||
def load_stage_configs_from_model(
|
||
model: str,
|
||
*,
|
||
trust_remote_code: bool | None,
|
||
base_engine_args: dict | None = None,
|
||
deploy_config_path: str | None = None,
|
||
stage_overrides: dict[str, dict[str, Any]] | None = None,
|
||
strategy_config_path: str | None = None,
|
||
) -> tuple[list, str | None]:
|
||
"""Load stage configurations from model's default config file.
|
||
|
||
For models registered in the pipeline registry, uses
|
||
``StageConfigFactory.create_legacy_stage_configs_from_model()`` which merges
|
||
PipelineConfig + DeployConfig + CLI overrides.
|
||
|
||
Models that cannot be resolved through the registry return no stages so the
|
||
caller can apply its default stage configuration, when available.
|
||
|
||
Args:
|
||
model: Model name or path (used to determine model_type)
|
||
trust_remote_code: Whether to trust remote code while resolving the model config.
|
||
base_engine_args: Base engine args to merge as CLI overrides.
|
||
deploy_config_path: Optional explicit deploy config path.
|
||
stage_overrides: Per-stage overrides from --stage-overrides.
|
||
strategy_config_path: Optional path to a composable-parallel
|
||
``strategy.yaml`` whose derived sizing is overlaid onto the
|
||
registry-merged stages.
|
||
|
||
Returns:
|
||
``(stage_configs, omni_lb_policy)``: the list of stage configuration
|
||
dictionaries plus the strategy-derived pipeline-wide ``omni_lb_policy``
|
||
(``None`` when no strategy set one). The policy is returned rather than
|
||
written into a caller-provided mutable dict.
|
||
"""
|
||
if base_engine_args is None:
|
||
base_engine_args = {}
|
||
|
||
cli_overrides = _convert_dataclasses_to_dict(dict(base_engine_args))
|
||
# A False inherited from the engine-args dump is the store_true flag's
|
||
# default, not an explicit choice — drop it so only the tri-state
|
||
# parameter below decides (see with_trust_remote_code_override).
|
||
if not cli_overrides.get("trust_remote_code"):
|
||
cli_overrides.pop("trust_remote_code", None)
|
||
cli_overrides = with_trust_remote_code_override(cli_overrides, trust_remote_code)
|
||
if stage_overrides:
|
||
for stage_id_str, overrides in stage_overrides.items():
|
||
for key, val in overrides.items():
|
||
cli_overrides[f"stage_{stage_id_str}_{key}"] = val
|
||
|
||
# Current runtime initialization still consumes legacy OmegaConf stage
|
||
# configs. ``StageConfigFactory.create_from_model`` now produces the
|
||
# structured ``VllmOmniConfig`` object; future RFC #4021 changes will
|
||
# migrate the engine/runtime consumers and replace this legacy resolver.
|
||
strategy_specs = None
|
||
if strategy_config_path is not None:
|
||
from vllm_omni.config.composable_parallel.strategy_loader import load_strategy_specs
|
||
|
||
strategy_specs = load_strategy_specs(strategy_config_path)
|
||
|
||
stages, omni_lb_policy = StageConfigFactory.create_legacy_stage_configs_from_model(
|
||
model,
|
||
trust_remote_code=trust_remote_code,
|
||
cli_overrides=cli_overrides,
|
||
deploy_config_path=deploy_config_path,
|
||
strategy_specs=strategy_specs,
|
||
)
|
||
if stages is not None:
|
||
# Convert StageConfig objects to OmegaConf for backward compat
|
||
return [stage.to_omegaconf() for stage in stages], omni_lb_policy
|
||
|
||
strategy_note = ""
|
||
if strategy_config_path is not None:
|
||
strategy_note = f" Strategy config {strategy_config_path!r} was not applied."
|
||
logger.warning(
|
||
"No registered PipelineConfig resolved for model %r. Legacy `stage_args` "
|
||
"YAMLs are no longer supported; register the pipeline and provide deployment "
|
||
"overrides through `deploy_config`.%s",
|
||
model,
|
||
strategy_note,
|
||
)
|
||
return [], None
|
||
|
||
|
||
def filter_stages(
|
||
config_path: str | None,
|
||
stage_configs: list,
|
||
kwargs: dict | None,
|
||
) -> list:
|
||
"""Filter stage configs by mode when YAML defines a `modes` section.
|
||
|
||
The YAML can define, e.g.:
|
||
|
||
modes:
|
||
- mode: text-to-image
|
||
stages: [1]
|
||
- mode: image-to-text
|
||
stages: [0]
|
||
|
||
When users pass `mode="image-to-text"` into Omni(**kwargs), only the stages
|
||
listed for that mode are returned. If no mode is provided, defaults to
|
||
"text-to-image". If no modes are defined or filtering fails, returns the
|
||
original stage_configs unchanged.
|
||
|
||
Args:
|
||
config_path: Path to the YAML config (used to read `modes`).
|
||
stage_configs: Loaded list of stage configs.
|
||
kwargs: Engine/caller kwargs; may contain "mode".
|
||
|
||
Returns:
|
||
Filtered list of stage configs (or original list if filtering not applied).
|
||
"""
|
||
if not stage_configs or config_path is None:
|
||
return stage_configs
|
||
|
||
try:
|
||
cfg = load_yaml_config(config_path)
|
||
yaml_modes = getattr(cfg, "modes", None)
|
||
if yaml_modes is None:
|
||
return stage_configs
|
||
|
||
mode_to_stage_ids: dict[str, list[int]] = {}
|
||
if yaml_modes is not None:
|
||
for entry in yaml_modes:
|
||
mode_name = None
|
||
stages = None
|
||
if hasattr(entry, "mode") or hasattr(entry, "stages"):
|
||
mode_name = getattr(entry, "mode", None)
|
||
stages = getattr(entry, "stages", None)
|
||
elif isinstance(entry, dict):
|
||
mode_name = entry.get("mode")
|
||
stages = entry.get("stages")
|
||
|
||
if mode_name is None or stages is None:
|
||
continue
|
||
|
||
if isinstance(stages, int):
|
||
stage_list = [stages]
|
||
else:
|
||
stage_list = list(stages)
|
||
|
||
mode_to_stage_ids[str(mode_name)] = [int(sid) for sid in stage_list]
|
||
|
||
# No modes section or empty mapping: use all stages and return early.
|
||
active_mode: str | None = None
|
||
if isinstance(kwargs, dict):
|
||
active_mode = kwargs.get("mode")
|
||
|
||
if active_mode is None:
|
||
active_mode = "text-to-image"
|
||
|
||
if active_mode not in mode_to_stage_ids:
|
||
logger.warning(
|
||
"Requested mode '%s' not found in config '%s'; available modes: %s. Using all stages.",
|
||
active_mode,
|
||
config_path,
|
||
sorted(mode_to_stage_ids.keys()),
|
||
)
|
||
return stage_configs
|
||
|
||
allowed_ids = set(mode_to_stage_ids[active_mode])
|
||
filtered_stage_configs = [sc for sc in stage_configs if getattr(sc, "stage_id", None) in allowed_ids]
|
||
if not filtered_stage_configs:
|
||
logger.warning(
|
||
"Mode '%s' in config '%s' resolved to stage ids %s, but none matched loaded stage_args. "
|
||
"Falling back to all stages.",
|
||
active_mode,
|
||
config_path,
|
||
sorted(allowed_ids),
|
||
)
|
||
return stage_configs
|
||
|
||
return filtered_stage_configs
|
||
except Exception as e:
|
||
logger.warning("Failed to apply mode-based stage filtering: %s", e)
|
||
return stage_configs
|
||
|
||
|
||
def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
|
||
"""Parse the ``--stage-overrides`` value into a per-stage override dict.
|
||
|
||
``value`` may be a raw JSON string (as supplied on the CLI) or an
|
||
already-parsed mapping. Returns ``None`` when no overrides are given.
|
||
|
||
Only the **shape** of the override mapping is validated here:
|
||
|
||
- top-level must be a dict
|
||
- keys must be non-negative ASCII integer strings (stage ids)
|
||
- values must be dicts
|
||
|
||
Field-name / type / range validation is intentionally not done at this
|
||
layer. The downstream resolver forwards every surviving key as
|
||
``stage_<id>_<key>`` into ``cli_overrides`` (see
|
||
``load_stage_configs_from_model``); field semantics live in
|
||
``StageConfigFactory`` for registered pipelines, and the
|
||
default-diffusion fallback in ``async_omni_engine.py`` reads
|
||
``stage_zero_overrides["extras"]`` for unregistered ones. Unknown
|
||
engine-arg keys that don't match ``OmniEngineArgs`` are dropped with
|
||
a warning at ``filter_dataclass_kwargs`` rather than failing here,
|
||
so a typo surfaces in the log instead of being silently lost.
|
||
|
||
Raises:
|
||
ValueError: when ``value`` is not a valid per-stage override
|
||
mapping (invalid JSON, non-dict top level, non-ASCII-digit
|
||
stage id, non-dict value).
|
||
"""
|
||
if not value:
|
||
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(
|
||
f"--stage-overrides must be a JSON object mapping stage_id -> overrides, "
|
||
f"got {type(parsed).__name__}: {parsed!r}"
|
||
)
|
||
if not parsed:
|
||
return None
|
||
|
||
for stage_id_str, overrides in parsed.items():
|
||
# ``str.isdigit()`` accepts Unicode digit classes (e.g. fullwidth
|
||
# ``"０"``); downstream's ``_STAGE_OVERRIDE_PATTERN`` is ASCII-only
|
||
# (``^stage_(\d+)_(.+)$``), so a non-ASCII-digit key would parse
|
||
# here and silently misroute as a global key. ``isascii()`` closes
|
||
# that hole.
|
||
if not isinstance(stage_id_str, str) or not stage_id_str.isdigit() or not stage_id_str.isascii():
|
||
raise ValueError(
|
||
f"--stage-overrides keys must be non-negative integer stage ids (as strings), got {stage_id_str!r}"
|
||
)
|
||
if not isinstance(overrides, dict):
|
||
raise ValueError(
|
||
f"--stage-overrides[{stage_id_str!r}] must be an object, got {type(overrides).__name__}: {overrides!r}"
|
||
)
|
||
|
||
return parsed
|
||
|
||
|
||
def load_and_resolve_stage_configs(
|
||
model: str,
|
||
kwargs: dict | None,
|
||
*,
|
||
trust_remote_code: bool | None,
|
||
default_stage_cfg_factory: Any = None,
|
||
deploy_config_path: str | None = None,
|
||
stage_overrides: dict[str, dict[str, Any]] | None = None,
|
||
strategy_config_path: str | None = None,
|
||
) -> tuple[str, list, str | None]:
|
||
"""Load stage configurations from a deploy YAML or model defaults.
|
||
|
||
Args:
|
||
model: Model name or path
|
||
kwargs: Engine arguments to merge with stage configs
|
||
trust_remote_code: Whether to trust remote code while resolving the model config.
|
||
default_stage_cfg_factory: Optional callable that takes no args and returns
|
||
default stage config list when no configs are found
|
||
deploy_config_path: Optional path to a deploy YAML.
|
||
stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON.
|
||
Keys are stage_id strings, values are dicts of overrides.
|
||
strategy_config_path: Optional path to a composable-parallel
|
||
``strategy.yaml`` overlaid onto the registry-merged stages.
|
||
|
||
Returns:
|
||
Tuple of ``(config_path, stage_configs, omni_lb_policy)`` — the last is
|
||
the strategy-derived pipeline-wide load-balancer policy (``None`` when no
|
||
strategy set one), returned for the engine to apply.
|
||
"""
|
||
config_path = deploy_config_path if deploy_config_path is not None else resolve_model_config_path(model)
|
||
stage_configs, omni_lb_policy = load_stage_configs_from_model(
|
||
model,
|
||
trust_remote_code=trust_remote_code,
|
||
base_engine_args=kwargs,
|
||
deploy_config_path=deploy_config_path,
|
||
stage_overrides=stage_overrides,
|
||
strategy_config_path=strategy_config_path,
|
||
)
|
||
if not stage_configs:
|
||
if default_stage_cfg_factory is not None:
|
||
default_stage_cfg = default_stage_cfg_factory()
|
||
stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
|
||
else:
|
||
stage_configs = []
|
||
|
||
stage_configs = filter_stages(config_path, stage_configs, kwargs)
|
||
|
||
# ``revision`` selects the model source for every stage. Keep the global
|
||
# CLI pin on each stage's engine args so diffusion config construction can
|
||
# pass it through to all component loaders. The registry-to-legacy-stage
|
||
# conversion currently drops this inherited vLLM EngineArgs field.
|
||
revision = (kwargs or {}).get("revision")
|
||
if revision is not None:
|
||
for stage_config in stage_configs:
|
||
engine_args = stage_config.get("engine_args") if hasattr(stage_config, "get") else None
|
||
if engine_args is not None and engine_args.get("revision") is None:
|
||
engine_args["revision"] = revision
|
||
|
||
logger.debug(f"stage_configs: {stage_configs}")
|
||
|
||
return config_path, stage_configs, omni_lb_policy
|
||
|
||
|
||
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()
|
||
|
||
|
||
### 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.warning("Request appears to be streaming, but got request type final only!")
|
||
elif not is_streaming and params.output_kind == RequestOutputKind.DELTA:
|
||
logger.warning("Request appears to not be streaming, but got request type delta!")
|
||
|
||
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)
|