712 lines
29 KiB
Python
712 lines
29 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""AsyncOmniEngine: the turn-based engine (add_request / streaming updates / CFG companions / interaction)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import shutil
|
|
import time
|
|
from collections.abc import Mapping, Sequence
|
|
from typing import Any, Literal, cast
|
|
|
|
from vllm.inputs import PromptType
|
|
from vllm.logger import init_logger
|
|
from vllm.v1.engine import EngineCoreRequest
|
|
from vllm.v1.kv_hints import KvHintsEnvelope
|
|
|
|
from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY, TRANSFORM_OWNED_META_KEYS
|
|
from vllm_omni.engine import OmniEngineCoreRequest
|
|
from vllm_omni.engine.async_engine_utils import (
|
|
apply_omni_final_stage_metadata,
|
|
inject_global_id,
|
|
upgrade_to_omni_request,
|
|
)
|
|
from vllm_omni.engine.messages import (
|
|
AddCompanionRequestMessage,
|
|
InteractionMessage,
|
|
StageSubmissionMessage,
|
|
)
|
|
from vllm_omni.engine.omni_engine_base import OmniEngineBase, StageRuntimeInfo
|
|
from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorBase
|
|
from vllm_omni.engine.serialization import deserialize_additional_information
|
|
from vllm_omni.inputs.data import OmniInteractionPrompt, OmniSamplingParams
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class AsyncOmniEngine(OmniEngineBase):
|
|
"""Turn-based engine used by ``Omni`` / ``AsyncOmni``."""
|
|
|
|
def _create_orchestrator(self, **orchestrator_kwargs: Any) -> OrchestratorBase:
|
|
return Orchestrator(**orchestrator_kwargs)
|
|
|
|
# ---- request helpers ----
|
|
|
|
@staticmethod
|
|
def _iter_multimodal_items(value: Any) -> list[Any]:
|
|
if value is None:
|
|
return []
|
|
if isinstance(value, list):
|
|
return value
|
|
return [value]
|
|
|
|
_DEFAULT_MM_HASHER_ALGORITHM = "blake3"
|
|
|
|
def _resolve_mm_hasher_algorithm(self) -> str:
|
|
"""Return the MM hash algorithm configured for stage 0.
|
|
|
|
Upstream vLLM 0.29 removed ``MultiModalHasher``'s ``_get_mm_hasher_algorithm``
|
|
helper and the ``VLLM_MM_HASHER_ALGORITHM`` env var; the algorithm now lives
|
|
only on ``MultiModalConfig.mm_hasher_algorithm`` (``--mm-hasher-algorithm``,
|
|
default ``"blake3"``). Resolve it from stage-0's config so the pre-computed
|
|
replica-scoped uuid uses the same algorithm as the stage-0 renderer built by
|
|
``build_stage0_input_processor``. Falls back to the upstream default when the
|
|
engine has no stage config (unit tests built via ``object.__new__``) or the
|
|
model has no multimodal config.
|
|
"""
|
|
|
|
stage_vllm_configs = getattr(self, "stage_vllm_configs", None) or []
|
|
if not stage_vllm_configs:
|
|
return self._DEFAULT_MM_HASHER_ALGORITHM
|
|
|
|
model_config = getattr(stage_vllm_configs[0], "model_config", None)
|
|
if model_config is None:
|
|
return self._DEFAULT_MM_HASHER_ALGORITHM
|
|
|
|
mm_config = None
|
|
get_multimodal_config = getattr(model_config, "get_multimodal_config", None)
|
|
if callable(get_multimodal_config):
|
|
try:
|
|
mm_config = get_multimodal_config()
|
|
except ValueError:
|
|
# Model is not multimodal.
|
|
mm_config = None
|
|
if mm_config is None:
|
|
mm_config = getattr(model_config, "multimodal_config", None)
|
|
|
|
return getattr(mm_config, "mm_hasher_algorithm", None) or self._DEFAULT_MM_HASHER_ALGORITHM
|
|
|
|
def _ensure_stage_replica_mm_uuids(
|
|
self,
|
|
prompt: Any,
|
|
*,
|
|
stage_id: int,
|
|
replica_id: int,
|
|
) -> None:
|
|
"""Make multimodal processor-cache keys local to a stage replica.
|
|
|
|
vLLM's frontend multimodal sender cache is process-global, while each
|
|
vllm-omni stage replica owns a separate EngineCore receiver cache. If
|
|
two requests with the same image are routed to different stage-0
|
|
replicas, a plain content hash can make the sender omit the tensor for
|
|
a replica that has never received it. Prefixing user/content UUIDs with
|
|
the selected replica keeps cache reuse within the receiver that owns it.
|
|
"""
|
|
|
|
if not isinstance(prompt, dict):
|
|
return
|
|
|
|
mm_data = prompt.get("multi_modal_data")
|
|
if not isinstance(mm_data, dict) or not mm_data:
|
|
return
|
|
|
|
from vllm.multimodal.hasher import MultiModalHasher
|
|
|
|
mm_hasher_algorithm = self._resolve_mm_hasher_algorithm()
|
|
|
|
existing_uuids = prompt.get("multi_modal_uuids")
|
|
if not isinstance(existing_uuids, dict):
|
|
existing_uuids = {}
|
|
|
|
model_id = str(getattr(self, "model", ""))
|
|
scoped_uuids: dict[str, list[str | None]] = dict(existing_uuids)
|
|
for modality, raw_items in mm_data.items():
|
|
items = self._iter_multimodal_items(raw_items)
|
|
if not items:
|
|
continue
|
|
|
|
modality_existing = existing_uuids.get(modality)
|
|
if not isinstance(modality_existing, list):
|
|
modality_existing = [modality_existing] if modality_existing is not None else []
|
|
|
|
modality_uuids: list[str | None] = []
|
|
for idx, item in enumerate(items):
|
|
user_uuid = modality_existing[idx] if idx < len(modality_existing) else None
|
|
if user_uuid is not None:
|
|
base_uuid = str(user_uuid)
|
|
elif item is None:
|
|
base_uuid = None
|
|
else:
|
|
base_uuid = MultiModalHasher.hash_kwargs(
|
|
mm_hasher_algorithm,
|
|
model_id=model_id,
|
|
**{modality: item},
|
|
)
|
|
|
|
if base_uuid is None:
|
|
modality_uuids.append(None)
|
|
else:
|
|
modality_uuids.append(f"stage{stage_id}:rep{replica_id}:{base_uuid}")
|
|
|
|
scoped_uuids[modality] = modality_uuids
|
|
|
|
if scoped_uuids:
|
|
prompt["multi_modal_uuids"] = scoped_uuids
|
|
|
|
@staticmethod
|
|
def _stage_pool_replica_count(stage_pool: Any) -> int:
|
|
try:
|
|
live_num_replicas = getattr(stage_pool, "live_num_replicas", None)
|
|
if live_num_replicas is not None:
|
|
return int(live_num_replicas)
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
live_replica_ids = getattr(stage_pool, "live_replica_ids", None)
|
|
if callable(live_replica_ids):
|
|
return len(live_replica_ids())
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
clients = getattr(stage_pool, "clients", None)
|
|
if clients is not None:
|
|
return sum(1 for client in clients if client is not None)
|
|
except Exception:
|
|
pass
|
|
|
|
return int(getattr(stage_pool, "num_replicas", 1) or 1)
|
|
|
|
@staticmethod
|
|
def _stage_pool_is_distributed(stage_pool: Any) -> bool:
|
|
try:
|
|
is_distributed = getattr(stage_pool, "is_distributed", None)
|
|
if is_distributed is not None:
|
|
return bool(is_distributed() if callable(is_distributed) else is_distributed)
|
|
except Exception:
|
|
pass
|
|
|
|
return getattr(stage_pool, "_hub", None) is not None
|
|
|
|
def _scope_stage0_multimodal_cache_to_replica(
|
|
self,
|
|
request_id: str,
|
|
prompt: Any,
|
|
) -> int | None:
|
|
stage_pools = getattr(self, "stage_pools", None)
|
|
if isinstance(prompt, EngineCoreRequest) or not stage_pools:
|
|
return None
|
|
|
|
stage0_pool = stage_pools[0]
|
|
# TODO: Currently only supports the ar -> dit process.
|
|
# Future scenarios (e.g., dit -> ar) need to be added, which will require modifications here.
|
|
if stage0_pool.stage_type == "diffusion" or self._stage_pool_replica_count(stage0_pool) <= 1:
|
|
return None
|
|
|
|
prompts = prompt if isinstance(prompt, list) else [prompt]
|
|
if not any(isinstance(p, dict) and p.get("multi_modal_data") for p in prompts):
|
|
return None
|
|
|
|
if self._stage_pool_is_distributed(stage0_pool):
|
|
preselect_replica_id = getattr(stage0_pool, "preselect_replica_id", None)
|
|
if not callable(preselect_replica_id):
|
|
logger.debug(
|
|
"[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing "
|
|
"without preselect support req=%s",
|
|
request_id,
|
|
)
|
|
return None
|
|
replica_id = preselect_replica_id(request_id)
|
|
if replica_id is None:
|
|
logger.debug(
|
|
"[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing "
|
|
"because no serviceable replica is available yet req=%s",
|
|
request_id,
|
|
)
|
|
return None
|
|
else:
|
|
replica_id = stage0_pool.select_replica_id(request_id)
|
|
|
|
for p in prompts:
|
|
self._ensure_stage_replica_mm_uuids(
|
|
p,
|
|
stage_id=0,
|
|
replica_id=replica_id,
|
|
)
|
|
|
|
logger.debug(
|
|
"[AsyncOmniEngine] Scoped multimodal cache keys to stage-0 replica-%s for req=%s",
|
|
replica_id,
|
|
request_id,
|
|
)
|
|
return replica_id
|
|
|
|
def _build_add_request_message(
|
|
self,
|
|
request_id: str,
|
|
prompt: EngineCoreRequest | PromptType,
|
|
prompt_text: str | None = None,
|
|
sampling_params_list: Sequence[Any] | None = None,
|
|
final_stage_id: int = 0,
|
|
final_output_stage_ids: Sequence[int] | None = None,
|
|
arrival_time: float | None = None,
|
|
lora_request: Any = None,
|
|
tokenization_kwargs: dict[str, Any] | None = None,
|
|
trace_headers: Mapping[str, str] | None = None,
|
|
priority: int = 0,
|
|
data_parallel_rank: int | None = None,
|
|
reasoning_ended: bool | None = None,
|
|
*,
|
|
kv_hints: KvHintsEnvelope | None = None,
|
|
resumable: bool = False,
|
|
message_type: Literal["add_request", "streaming_update"] = "add_request",
|
|
) -> StageSubmissionMessage:
|
|
"""Build an add_request message after stage-0 preprocessing."""
|
|
request_timestamp = float(arrival_time) if arrival_time is not None else time.time()
|
|
effective_sampling_params_list: list[OmniSamplingParams] = (
|
|
list(cast(Sequence[OmniSamplingParams], sampling_params_list))
|
|
if sampling_params_list is not None
|
|
else list(self.default_sampling_params_list)
|
|
)
|
|
if not effective_sampling_params_list:
|
|
raise ValueError(
|
|
f"Missing sampling params for stage 0. Got {len(effective_sampling_params_list)} stage params."
|
|
)
|
|
params = effective_sampling_params_list[0]
|
|
|
|
# Keep the original prompt for downstream stages (they need the raw
|
|
# dict, e.g. for multi_modal_data).
|
|
if isinstance(prompt, dict):
|
|
raw_info = prompt.get("additional_information")
|
|
if isinstance(raw_info, dict):
|
|
raw_meta = raw_info.get("meta")
|
|
if isinstance(raw_meta, dict):
|
|
for key in TRANSFORM_OWNED_META_KEYS:
|
|
raw_meta.pop(key, None)
|
|
original_prompt = prompt
|
|
preselected_stage0_replica: int | None = None
|
|
request_artifact_dirs: list[str] = []
|
|
|
|
stage_type = self.stage_metadata[0].stage_type
|
|
if stage_type == "diffusion" and kv_hints is not None:
|
|
raise ValueError("kv_hints require an AR/LLM stage with a KV cache")
|
|
output_prompt_text: Any = None
|
|
_preprocess_ms = 0.0
|
|
if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest):
|
|
# Stage transforms and downstream stages must share the same
|
|
# request identity, including when the transform replaces the
|
|
# prompt object.
|
|
if isinstance(prompt, dict):
|
|
inject_global_id(prompt, request_id)
|
|
elif isinstance(prompt, list):
|
|
for item in prompt:
|
|
inject_global_id(item, request_id)
|
|
|
|
prompt_transform_func = getattr(self, "prompt_transform_func", None)
|
|
if prompt_transform_func is not None:
|
|
if isinstance(prompt, dict):
|
|
prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None)
|
|
prompt = prompt_transform_func(
|
|
copy.copy(prompt),
|
|
effective_sampling_params_list,
|
|
)
|
|
if isinstance(prompt, dict):
|
|
raw_dirs = prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None)
|
|
if isinstance(raw_dirs, list) and all(isinstance(path, str) for path in raw_dirs):
|
|
request_artifact_dirs = list(raw_dirs)
|
|
if isinstance(original_prompt, dict):
|
|
original_prompt[REQUEST_ARTIFACT_DIRS_KEY] = request_artifact_dirs
|
|
|
|
if isinstance(prompt, dict):
|
|
inject_global_id(prompt, request_id)
|
|
elif isinstance(prompt, list):
|
|
for item in prompt:
|
|
inject_global_id(item, request_id)
|
|
|
|
preselected_stage0_replica = self._scope_stage0_multimodal_cache_to_replica(
|
|
request_id,
|
|
prompt,
|
|
)
|
|
|
|
# Full input processing (tokenization, multimodal, etc.)
|
|
assert self.input_processor is not None
|
|
_t_preprocess = time.perf_counter()
|
|
try:
|
|
request = self.input_processor.process_inputs(
|
|
request_id=request_id,
|
|
prompt=prompt,
|
|
params=params,
|
|
supported_tasks=self.supported_tasks,
|
|
arrival_time=arrival_time,
|
|
lora_request=lora_request,
|
|
tokenization_kwargs=tokenization_kwargs,
|
|
trace_headers=trace_headers,
|
|
priority=priority,
|
|
data_parallel_rank=data_parallel_rank,
|
|
resumable=resumable,
|
|
**({"kv_hints": kv_hints} if kv_hints is not None else {}),
|
|
)
|
|
except Exception:
|
|
if preselected_stage0_replica is not None and self.stage_pools:
|
|
self.stage_pools[0].release_binding(request_id)
|
|
for artifact_dir in request_artifact_dirs:
|
|
shutil.rmtree(artifact_dir, ignore_errors=True)
|
|
raise
|
|
_preprocess_ms = (time.perf_counter() - _t_preprocess) * 1000.0
|
|
# TODO (Peiqi): add this for Qwen3-TTS only. Other models don't have
|
|
# additional_information field in the prompt.
|
|
request = upgrade_to_omni_request(request, prompt)
|
|
|
|
if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None:
|
|
processed_info = deserialize_additional_information(request.additional_information)
|
|
processed_meta = processed_info.get("meta")
|
|
if isinstance(processed_meta, dict):
|
|
if isinstance(original_prompt, dict):
|
|
original_info = dict(original_prompt.get("additional_information") or {})
|
|
original_meta = dict(original_info.get("meta") or {})
|
|
original_meta.update(processed_meta)
|
|
original_info["meta"] = original_meta
|
|
original_prompt["additional_information"] = original_info
|
|
|
|
if reasoning_ended is not None:
|
|
request.reasoning_ended = reasoning_ended
|
|
|
|
# Restore external_req_id to the original user-facing request_id.
|
|
# InputProcessor.process_inputs() renames request_id to an internal
|
|
# UUID (saving the original in external_req_id), but then overwrites
|
|
# external_req_id with the new internal ID. We need external_req_id
|
|
# to match the key used in Orchestrator.request_states so that
|
|
# output routing (output.request_id lookup) can find the req_state.
|
|
request.external_req_id = request_id
|
|
request = apply_omni_final_stage_metadata(request, final_stage_id)
|
|
|
|
# Registration with stage 0's output processor is deferred to the
|
|
# orchestrator thread (see Orchestrator._handle_add_request), which
|
|
# now routes admission through StagePool.submit_initial().
|
|
output_prompt_text = prompt_text
|
|
if output_prompt_text is None and isinstance(original_prompt, dict):
|
|
output_prompt_text = original_prompt.get("prompt")
|
|
prompt = request
|
|
else:
|
|
request_artifact_dirs = []
|
|
|
|
return StageSubmissionMessage(
|
|
type=message_type,
|
|
request_id=request_id,
|
|
prompt=prompt,
|
|
original_prompt=original_prompt,
|
|
output_prompt_text=output_prompt_text,
|
|
sampling_params_list=effective_sampling_params_list,
|
|
final_stage_id=final_stage_id,
|
|
final_output_stage_ids=list(final_output_stage_ids) if final_output_stage_ids is not None else None,
|
|
preprocess_ms=_preprocess_ms,
|
|
request_timestamp=request_timestamp,
|
|
enqueue_ts=time.perf_counter(),
|
|
request_artifact_dirs=request_artifact_dirs or None,
|
|
)
|
|
|
|
def _build_cfg_companions(
|
|
self,
|
|
parent_id: str,
|
|
original_prompt: Any,
|
|
stage0_params: Any,
|
|
sampling_params_list: list[Any],
|
|
) -> list[AddCompanionRequestMessage]:
|
|
"""Expand a prompt into its CFG companions, without enqueueing any.
|
|
|
|
Construction is separated from admission so a guided request is
|
|
all-or-nothing. A model whose guidance is mandatory cannot decode a
|
|
request whose companion never arrived: the pair never completes, the
|
|
request occupies scheduler and KV capacity for the scheduler's whole
|
|
hold budget, and then produces no audio. Raising here instead means the
|
|
caller learns immediately and nothing was admitted.
|
|
|
|
Raises:
|
|
Exception: Whatever prompt expansion or input processing raised.
|
|
The caller is expected to let it reach the client.
|
|
"""
|
|
assert self.prompt_expand_func is not None
|
|
expanded = self.prompt_expand_func(original_prompt, stage0_params)
|
|
if not expanded:
|
|
return []
|
|
|
|
companions: list[AddCompanionRequestMessage] = []
|
|
assert self.input_processor is not None
|
|
for ep in expanded:
|
|
cid = f"{parent_id}{ep.request_id_suffix}"
|
|
companion_prompt = ep.prompt
|
|
|
|
companion_params, companion_spl = ep.apply_overrides(stage0_params, sampling_params_list)
|
|
|
|
if isinstance(companion_prompt, dict):
|
|
inject_global_id(companion_prompt, cid)
|
|
|
|
request = self.input_processor.process_inputs(
|
|
request_id=cid,
|
|
prompt=companion_prompt,
|
|
params=companion_params,
|
|
supported_tasks=self.supported_tasks,
|
|
)
|
|
# Same restore the parent request gets: the upstream input
|
|
# processor drops omni-only prompt fields, so without this the
|
|
# companion reaches the worker with no additional_information at
|
|
# all. That is where ``global_request_id`` lives, and it is what
|
|
# was just injected above, so skipping it silently undoes the
|
|
# injection: the model sees the companion row with no id and
|
|
# cannot match it to its conditioned partner.
|
|
request = upgrade_to_omni_request(request, companion_prompt)
|
|
request.external_req_id = cid
|
|
# Companions are stage-0-final for ordinary downstream payloads,
|
|
# but diffusion still needs their CFG KV caches.
|
|
request = apply_omni_final_stage_metadata(request, 0, force_kv_transfer=True)
|
|
|
|
# Registration of this companion on stage-0's output processor is
|
|
# deferred to Orchestrator._handle_add_companion, which routes
|
|
# admission through StagePool.submit_initial(..., affinity_request_id=...).
|
|
companions.append(
|
|
AddCompanionRequestMessage(
|
|
companion_id=cid,
|
|
parent_id=parent_id,
|
|
role=ep.role,
|
|
prompt=request,
|
|
companion_prompt_text=companion_prompt,
|
|
sampling_params_list=companion_spl,
|
|
)
|
|
)
|
|
return companions
|
|
|
|
def _enqueue_cfg_companions(
|
|
self,
|
|
parent_id: str,
|
|
original_prompt: Any,
|
|
stage0_params: Any,
|
|
sampling_params_list: list[Any],
|
|
) -> None:
|
|
"""Build and enqueue CFG companions, tolerating a build failure.
|
|
|
|
Kept for callers that admit the parent first and cannot roll it back.
|
|
Prefer building with :meth:`_build_cfg_companions` before the parent is
|
|
admitted, so the pair is atomic.
|
|
"""
|
|
try:
|
|
companions = self._build_cfg_companions(parent_id, original_prompt, stage0_params, sampling_params_list)
|
|
except Exception:
|
|
logger.exception("[AsyncOmniEngine] CFG companion build failed for req %s", parent_id)
|
|
return
|
|
for companion in companions:
|
|
self.request_queue.sync_q.put(companion)
|
|
if not companions:
|
|
return
|
|
|
|
logger.info(
|
|
"[AsyncOmniEngine] CFG expansion for req %s: %d companions",
|
|
parent_id,
|
|
len(companions),
|
|
)
|
|
|
|
def add_request(
|
|
self,
|
|
request_id: str,
|
|
prompt: EngineCoreRequest | PromptType,
|
|
prompt_text: str | None = None,
|
|
sampling_params_list: Sequence[Any] | None = None,
|
|
final_stage_id: int = 0,
|
|
final_output_stage_ids: Sequence[int] | None = None,
|
|
arrival_time: float | None = None,
|
|
lora_request: Any = None,
|
|
tokenization_kwargs: dict[str, Any] | None = None,
|
|
trace_headers: Mapping[str, str] | None = None,
|
|
priority: int = 0,
|
|
data_parallel_rank: int | None = None,
|
|
reasoning_ended: bool | None = None,
|
|
*,
|
|
kv_hints: KvHintsEnvelope | None = None,
|
|
resumable: bool = False,
|
|
) -> None:
|
|
"""Process stage-0 input locally, then send to the Orchestrator.
|
|
|
|
Input processing and output
|
|
processor registration happen here in the caller's thread, avoiding
|
|
a queue + coroutine-switch round-trip. The Orchestrator receives a
|
|
ready-to-submit OmniEngineCoreRequest.
|
|
"""
|
|
try:
|
|
msg = self._build_add_request_message(
|
|
request_id=request_id,
|
|
prompt=prompt,
|
|
prompt_text=prompt_text,
|
|
sampling_params_list=sampling_params_list,
|
|
final_stage_id=final_stage_id,
|
|
final_output_stage_ids=final_output_stage_ids,
|
|
arrival_time=arrival_time,
|
|
lora_request=lora_request,
|
|
tokenization_kwargs=tokenization_kwargs,
|
|
trace_headers=trace_headers,
|
|
priority=priority,
|
|
data_parallel_rank=data_parallel_rank,
|
|
reasoning_ended=reasoning_ended,
|
|
resumable=resumable,
|
|
**({"kv_hints": kv_hints} if kv_hints is not None else {}),
|
|
)
|
|
except BaseException:
|
|
if isinstance(prompt, dict):
|
|
for artifact_dir in prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None) or ():
|
|
if isinstance(artifact_dir, str):
|
|
shutil.rmtree(artifact_dir, ignore_errors=True)
|
|
raise
|
|
# CFG companions are built before the parent is admitted, so the group
|
|
# is all-or-nothing: a build failure raises here, nothing is enqueued,
|
|
# and the caller sees the error. Admitting the parent first would leave
|
|
# an orphan holding scheduler and KV capacity that can never complete,
|
|
# because a model whose guidance is mandatory cannot decode a request
|
|
# whose companion never arrived.
|
|
companions: list[AddCompanionRequestMessage] = []
|
|
try:
|
|
if self.prompt_expand_func is not None and final_stage_id > 0:
|
|
effective_spl = msg.sampling_params_list
|
|
stage0_params = effective_spl[0] if effective_spl else None
|
|
if stage0_params is not None:
|
|
companions = self._build_cfg_companions(
|
|
request_id, msg.original_prompt, stage0_params, effective_spl
|
|
)
|
|
|
|
self.request_queue.sync_q.put(msg)
|
|
except BaseException:
|
|
for artifact_dir in msg.request_artifact_dirs or ():
|
|
shutil.rmtree(artifact_dir, ignore_errors=True)
|
|
raise
|
|
finally:
|
|
if isinstance(msg.original_prompt, dict):
|
|
msg.original_prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None)
|
|
for companion in companions:
|
|
self.request_queue.sync_q.put(companion)
|
|
if companions:
|
|
logger.info(
|
|
"[AsyncOmniEngine] CFG expansion for req %s: %d companions",
|
|
request_id,
|
|
len(companions),
|
|
)
|
|
|
|
async def add_request_async(
|
|
self,
|
|
request_id: str,
|
|
prompt: EngineCoreRequest | PromptType,
|
|
prompt_text: str | None = None,
|
|
sampling_params_list: Sequence[Any] | None = None,
|
|
final_stage_id: int = 0,
|
|
final_output_stage_ids: Sequence[int] | None = None,
|
|
arrival_time: float | None = None,
|
|
lora_request: Any = None,
|
|
tokenization_kwargs: dict[str, Any] | None = None,
|
|
trace_headers: Mapping[str, str] | None = None,
|
|
priority: int = 0,
|
|
data_parallel_rank: int | None = None,
|
|
reasoning_ended: bool | None = None,
|
|
*,
|
|
kv_hints: KvHintsEnvelope | None = None,
|
|
resumable: bool = False,
|
|
) -> None:
|
|
"""Async add_request API."""
|
|
self.add_request(
|
|
request_id=request_id,
|
|
prompt=prompt,
|
|
prompt_text=prompt_text,
|
|
sampling_params_list=sampling_params_list,
|
|
final_stage_id=final_stage_id,
|
|
final_output_stage_ids=final_output_stage_ids,
|
|
arrival_time=arrival_time,
|
|
lora_request=lora_request,
|
|
tokenization_kwargs=tokenization_kwargs,
|
|
trace_headers=trace_headers,
|
|
priority=priority,
|
|
data_parallel_rank=data_parallel_rank,
|
|
reasoning_ended=reasoning_ended,
|
|
resumable=resumable,
|
|
**({"kv_hints": kv_hints} if kv_hints is not None else {}),
|
|
)
|
|
|
|
def add_streaming_update(
|
|
self,
|
|
request_id: str,
|
|
prompt: EngineCoreRequest | PromptType,
|
|
prompt_text: str | None = None,
|
|
sampling_params_list: Sequence[Any] | None = None,
|
|
final_stage_id: int = 0,
|
|
final_output_stage_ids: Sequence[int] | None = None,
|
|
arrival_time: float | None = None,
|
|
lora_request: Any = None,
|
|
*,
|
|
resumable: bool = True,
|
|
) -> None:
|
|
"""Send an incremental streaming update for an existing request."""
|
|
msg = self._build_add_request_message(
|
|
request_id=request_id,
|
|
prompt=prompt,
|
|
prompt_text=prompt_text,
|
|
sampling_params_list=sampling_params_list,
|
|
final_stage_id=final_stage_id,
|
|
final_output_stage_ids=final_output_stage_ids,
|
|
arrival_time=arrival_time,
|
|
lora_request=lora_request,
|
|
resumable=resumable,
|
|
message_type="streaming_update",
|
|
)
|
|
self.request_queue.sync_q.put(msg)
|
|
|
|
async def add_streaming_update_async(
|
|
self,
|
|
request_id: str,
|
|
prompt: EngineCoreRequest | PromptType,
|
|
prompt_text: str | None = None,
|
|
sampling_params_list: Sequence[Any] | None = None,
|
|
final_stage_id: int = 0,
|
|
final_output_stage_ids: Sequence[int] | None = None,
|
|
arrival_time: float | None = None,
|
|
lora_request: Any = None,
|
|
*,
|
|
resumable: bool = True,
|
|
) -> None:
|
|
"""Async wrapper for add_streaming_update()."""
|
|
self.add_streaming_update(
|
|
request_id=request_id,
|
|
prompt=prompt,
|
|
prompt_text=prompt_text,
|
|
sampling_params_list=sampling_params_list,
|
|
final_stage_id=final_stage_id,
|
|
final_output_stage_ids=final_output_stage_ids,
|
|
arrival_time=arrival_time,
|
|
lora_request=lora_request,
|
|
resumable=resumable,
|
|
)
|
|
|
|
def submit_interaction(
|
|
self,
|
|
request_id: str,
|
|
interaction: OmniInteractionPrompt,
|
|
) -> None:
|
|
"""Send an interaction control message to the Orchestrator."""
|
|
if self.request_queue is None:
|
|
raise RuntimeError("request_queue is not initialized")
|
|
|
|
self.request_queue.sync_q.put_nowait(
|
|
InteractionMessage(
|
|
request_id=request_id,
|
|
interaction=interaction,
|
|
)
|
|
)
|
|
|
|
async def submit_interaction_async(
|
|
self,
|
|
request_id: str,
|
|
interaction: OmniInteractionPrompt,
|
|
) -> None:
|
|
"""Async interaction API."""
|
|
self.submit_interaction(request_id, interaction)
|
|
|
|
|
|
# ``StageRuntimeInfo`` moved to ``omni_engine_base``; keep the historical import path.
|
|
__all__ = ["AsyncOmniEngine", "StageRuntimeInfo"]
|