omerpaz95andClaude Opus 5
8a0a7ee40a
Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
132 lines
4.9 KiB
Python
132 lines
4.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""EC connector helper utilities."""
|
|
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import torch
|
|
|
|
from vllm.logger import init_logger
|
|
from vllm.utils.collection_utils import is_list_of
|
|
from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.config import ModelConfig
|
|
from vllm.multimodal.inputs import MultiModalFeatureSpec
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class PlaceholderMetadataResolver:
|
|
"""Resolves which processed keys a model needs published per modality.
|
|
|
|
Reads `MultiModalDataParser.embedding_fields`, the same declaration the
|
|
consumer's parser requires, so the two cannot drift. An empty set means
|
|
the modality cannot be delivered out of band, and the consumer will
|
|
process the media itself.
|
|
"""
|
|
|
|
def __init__(self, model_config: "ModelConfig") -> None:
|
|
self._model_config = model_config
|
|
self._cache: dict[str, set[str]] = {}
|
|
|
|
def fields_for(self, modality: str) -> set[str]:
|
|
if modality in self._cache:
|
|
return self._cache[modality]
|
|
|
|
fields: set[str] = set()
|
|
try:
|
|
from vllm.multimodal import MULTIMODAL_REGISTRY
|
|
|
|
info = MULTIMODAL_REGISTRY.create_processor(self._model_config).info
|
|
fields = info.data_parser.placeholder_metadata_fields(modality)
|
|
except Exception:
|
|
logger.warning(
|
|
"Could not determine the placeholder metadata fields for "
|
|
"modality %s; the consumer will preprocess the media itself.",
|
|
modality,
|
|
exc_info=True,
|
|
)
|
|
|
|
self._cache[modality] = fields
|
|
return fields
|
|
|
|
|
|
def collect_ec_item_metadata(
|
|
mm_features: "list[MultiModalFeatureSpec]",
|
|
resolver: PlaceholderMetadataResolver,
|
|
) -> dict[str, dict[str, Any]]:
|
|
"""Build one `ec_transfer_params` entry per feature for `request_finished()`.
|
|
|
|
Keyed by mm_hash, each entry carries a `metadata` dict with whatever
|
|
placeholder fields `resolver` says this model needs published for its
|
|
modality, so a consumer can skip the image transform. `data` is None for
|
|
items served from the processor cache, in which case the metadata is
|
|
unavailable here and the consumer has to fall back to processing the
|
|
media itself. A connector that also has transfer coordinates to report
|
|
(e.g. NIXL peer_host/peer_port/size_bytes) merges those in alongside
|
|
`metadata`, not into it.
|
|
"""
|
|
items: dict[str, dict[str, Any]] = {}
|
|
for feature in mm_features:
|
|
metadata: dict[str, Any] = {}
|
|
if feature.data is not None:
|
|
wanted = resolver.fields_for(feature.modality)
|
|
for key, value in feature.data.get_data().items():
|
|
if key not in wanted:
|
|
continue
|
|
if isinstance(value, torch.Tensor):
|
|
metadata[key] = value.tolist()
|
|
elif is_list_of(value, (int, float)):
|
|
# Some metadata (e.g. Qwen3-VL video timestamps) is
|
|
# produced as a plain list rather than a tensor.
|
|
metadata[key] = value
|
|
items[feature.identifier] = {"metadata": metadata}
|
|
return items
|
|
|
|
|
|
class ECOutputAggregator:
|
|
"""Merge every worker's EC connector output onto the single
|
|
ModelRunnerOutput that reaches the scheduler.
|
|
|
|
Mirrors KVOutputAggregator: only `output_rank`'s output is returned to the
|
|
scheduler, but the EC connector may have run on any rank.
|
|
"""
|
|
|
|
def aggregate(
|
|
self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0
|
|
) -> ModelRunnerOutput | None:
|
|
output = outputs[output_rank]
|
|
if not output:
|
|
return None
|
|
|
|
finished_sending = set[str]()
|
|
finished_recving = set[str]()
|
|
worker_meta = None
|
|
for model_runner_output in outputs:
|
|
assert model_runner_output is not None
|
|
ec_output = model_runner_output.ec_connector_output
|
|
if not ec_output:
|
|
continue
|
|
|
|
finished_sending |= ec_output.finished_sending or set()
|
|
finished_recving |= ec_output.finished_recving or set()
|
|
|
|
if meta := ec_output.ec_connector_worker_meta:
|
|
worker_meta = (
|
|
meta if worker_meta is None else worker_meta.aggregate(meta)
|
|
)
|
|
|
|
aggregated = ECConnectorOutput(
|
|
finished_sending=finished_sending or None,
|
|
finished_recving=finished_recving or None,
|
|
ec_connector_worker_meta=worker_meta,
|
|
)
|
|
if aggregated.is_empty():
|
|
output.ec_connector_output = None
|
|
return output
|
|
|
|
# `output` is the shared empty output whenever `output_rank` had no work,
|
|
# so attach through the copy-on-write helper.
|
|
return ModelRunnerOutput.with_ec_conn_output(output, aggregated)
|