| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """EC connector helper utilities.""" |
| 4 | |
| 5 | from typing import TYPE_CHECKING, Any |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.utils.collection_utils import is_list_of |
| 11 | from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput |
| 12 | |
| 13 | if TYPE_CHECKING: |
| 14 | from vllm.config import ModelConfig |
| 15 | from vllm.multimodal.inputs import MultiModalFeatureSpec |
| 16 | |
| 17 | logger = init_logger(__name__) |
| 18 | |
| 19 | |
| 20 | class PlaceholderMetadataResolver: |
| 21 | """Resolves which processed keys a model needs published per modality. |
| 22 | |
| 23 | Reads `MultiModalDataParser.embedding_fields`, the same declaration the |
| 24 | consumer's parser requires, so the two cannot drift. An empty set means |
| 25 | the modality cannot be delivered out of band, and the consumer will |
| 26 | process the media itself. |
| 27 | """ |
| 28 | |
| 29 | def __init__(self, model_config: "ModelConfig") -> None: |
| 30 | self._model_config = model_config |
| 31 | self._cache: dict[str, set[str]] = {} |
| 32 | |
| 33 | def fields_for(self, modality: str) -> set[str]: |
| 34 | if modality in self._cache: |
| 35 | return self._cache[modality] |
| 36 | |
| 37 | fields: set[str] = set() |
| 38 | try: |
| 39 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 40 | |
| 41 | info = MULTIMODAL_REGISTRY.create_processor(self._model_config).info |
| 42 | fields = info.data_parser.placeholder_metadata_fields(modality) |
| 43 | except Exception: |
| 44 | logger.warning( |
| 45 | "Could not determine the placeholder metadata fields for " |
| 46 | "modality %s; the consumer will preprocess the media itself.", |
| 47 | modality, |
| 48 | exc_info=True, |
| 49 | ) |
| 50 | |
| 51 | self._cache[modality] = fields |
| 52 | return fields |
| 53 | |
| 54 | |
| 55 | def collect_ec_item_metadata( |
| 56 | mm_features: "list[MultiModalFeatureSpec]", |
| 57 | resolver: PlaceholderMetadataResolver, |
| 58 | ) -> dict[str, dict[str, Any]]: |
| 59 | """Build one `ec_transfer_params` entry per feature for `request_finished()`. |
| 60 | |
| 61 | Keyed by mm_hash, each entry carries a `metadata` dict with whatever |
| 62 | placeholder fields `resolver` says this model needs published for its |
| 63 | modality, so a consumer can skip media preprocessing. `item_indices` |
| 64 | identifies every occurrence in `mm_features`, including repeated hashes. |
| 65 | Audio token counts come from the placeholder; other metadata requires |
| 66 | `data` and falls back to raw media when it is unavailable on cache hits. |
| 67 | A connector that also has transfer coordinates to report |
| 68 | (e.g. NIXL peer_host/peer_port/size_bytes) merges those in alongside |
| 69 | `metadata`, not into it. |
| 70 | """ |
| 71 | items: dict[str, dict[str, Any]] = {} |
| 72 | for index, feature in enumerate(mm_features): |
| 73 | metadata: dict[str, Any] = {} |
| 74 | wanted = resolver.fields_for(feature.modality) |
| 75 | if "audio_num_tokens" in wanted: |
| 76 | metadata["audio_num_tokens"] = [feature.mm_position.get_num_embeds()] |
| 77 | if feature.data is not None: |
| 78 | for key, value in feature.data.get_data().items(): |
| 79 | if key not in wanted: |
| 80 | continue |
| 81 | if isinstance(value, torch.Tensor): |
| 82 | metadata[key] = torch.atleast_1d(value).tolist() |
| 83 | elif is_list_of(value, (int, float)): |
| 84 | # Some metadata (e.g. Qwen3-VL video timestamps) is |
| 85 | # produced as a plain list rather than a tensor. |
| 86 | metadata[key] = value |
| 87 | entry = items.setdefault(feature.identifier, {"item_indices": []}) |
| 88 | entry["metadata"] = metadata |
| 89 | entry["item_indices"].append(index) |
| 90 | return items |
| 91 | |
| 92 | |
| 93 | class ECOutputAggregator: |
| 94 | """Merge every worker's EC connector output onto the single |
| 95 | ModelRunnerOutput that reaches the scheduler. |
| 96 | |
| 97 | Mirrors KVOutputAggregator: only `output_rank`'s output is returned to the |
| 98 | scheduler, but the EC connector may have run on any rank. |
| 99 | """ |
| 100 | |
| 101 | def aggregate( |
| 102 | self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0 |
| 103 | ) -> ModelRunnerOutput | None: |
| 104 | output = outputs[output_rank] |
| 105 | if not output: |
| 106 | return None |
| 107 | |
| 108 | finished_sending = set[str]() |
| 109 | finished_recving = set[str]() |
| 110 | worker_meta = None |
| 111 | connector_stats = None |
| 112 | for model_runner_output in outputs: |
| 113 | assert model_runner_output is not None |
| 114 | ec_output = model_runner_output.ec_connector_output |
| 115 | if not ec_output: |
| 116 | continue |
| 117 | |
| 118 | finished_sending |= ec_output.finished_sending or set() |
| 119 | finished_recving |= ec_output.finished_recving or set() |
| 120 | |
| 121 | if meta := ec_output.ec_connector_worker_meta: |
| 122 | worker_meta = ( |
| 123 | meta if worker_meta is None else worker_meta.aggregate(meta) |
| 124 | ) |
| 125 | |
| 126 | if stats := ec_output.ec_connector_stats: |
| 127 | connector_stats = ( |
| 128 | stats |
| 129 | if connector_stats is None |
| 130 | else connector_stats.aggregate(stats) |
| 131 | ) |
| 132 | |
| 133 | aggregated = ECConnectorOutput( |
| 134 | finished_sending=finished_sending or None, |
| 135 | finished_recving=finished_recving or None, |
| 136 | ec_connector_stats=connector_stats, |
| 137 | ec_connector_worker_meta=worker_meta, |
| 138 | ) |
| 139 | if aggregated.is_empty(): |
| 140 | output.ec_connector_output = None |
| 141 | return output |
| 142 | |
| 143 | # `output` is the shared empty output whenever `output_rank` had no work, |
| 144 | # so attach through the copy-on-write helper. |
| 145 | return ModelRunnerOutput.with_ec_conn_output(output, aggregated) |