| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import threading |
| 4 | from collections import defaultdict |
| 5 | from dataclasses import dataclass |
| 6 | from typing import TYPE_CHECKING, Generic, Protocol, TypeVar, cast |
| 7 | |
| 8 | from vllm.logger import init_logger |
| 9 | from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config |
| 10 | |
| 11 | from .processing import ( |
| 12 | BaseDummyInputsBuilder, |
| 13 | BaseMultiModalProcessor, |
| 14 | BaseProcessingInfo, |
| 15 | InputProcessingContext, |
| 16 | TimingContext, |
| 17 | ) |
| 18 | |
| 19 | if TYPE_CHECKING: |
| 20 | from vllm.config import ModelConfig, ObservabilityConfig |
| 21 | from vllm.model_executor.models.interfaces import SupportsMultiModal |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | N = TypeVar("N", bound=type["SupportsMultiModal"]) |
| 26 | _I = TypeVar("_I", bound=BaseProcessingInfo) |
| 27 | _I_co = TypeVar("_I_co", bound=BaseProcessingInfo, covariant=True) |
| 28 | |
| 29 | |
| 30 | class ProcessingInfoFactory(Protocol[_I_co]): |
| 31 | """Constructs a |
| 32 | [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor] |
| 33 | instance from the context. |
| 34 | """ |
| 35 | |
| 36 | def __call__( |
| 37 | self, |
| 38 | ctx: InputProcessingContext, |
| 39 | ) -> _I_co: ... |
| 40 | |
| 41 | |
| 42 | class DummyInputsBuilderFactory(Protocol[_I]): # type: ignore[misc] |
| 43 | """Constructs a |
| 44 | [`BaseDummyInputsBuilder`][vllm.multimodal.processing.BaseDummyInputsBuilder] |
| 45 | instance from the context. |
| 46 | """ |
| 47 | |
| 48 | def __call__(self, info: _I) -> BaseDummyInputsBuilder[_I]: ... |
| 49 | |
| 50 | |
| 51 | class MultiModalProcessorFactory(Protocol[_I]): # type: ignore[misc] |
| 52 | """Constructs a |
| 53 | [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor] |
| 54 | instance from the context. |
| 55 | """ |
| 56 | |
| 57 | def __call__( |
| 58 | self, |
| 59 | info: _I, |
| 60 | dummy_inputs: BaseDummyInputsBuilder[_I], |
| 61 | ) -> BaseMultiModalProcessor[_I]: ... |
| 62 | |
| 63 | |
| 64 | @dataclass(frozen=True) |
| 65 | class _ProcessorFactories(Generic[_I]): |
| 66 | info: ProcessingInfoFactory[_I] |
| 67 | processor: MultiModalProcessorFactory[_I] |
| 68 | dummy_inputs: DummyInputsBuilderFactory[_I] |
| 69 | |
| 70 | def build_processor( |
| 71 | self, |
| 72 | ctx: InputProcessingContext, |
| 73 | ): |
| 74 | info = self.info(ctx) |
| 75 | dummy_inputs_builder = self.dummy_inputs(info) |
| 76 | return self.processor(info, dummy_inputs_builder) |
| 77 | |
| 78 | |
| 79 | class MultiModalRegistry: |
| 80 | """A registry that dispatches data processing according to the model.""" |
| 81 | |
| 82 | def register_processor( |
| 83 | self, |
| 84 | processor: MultiModalProcessorFactory[_I], |
| 85 | *, |
| 86 | info: ProcessingInfoFactory[_I], |
| 87 | dummy_inputs: DummyInputsBuilderFactory[_I], |
| 88 | ): |
| 89 | """Register a multi-modal processor to a model class. The processor |
| 90 | is constructed lazily, hence a factory method should be passed. |
| 91 | |
| 92 | When the model receives multi-modal data, the provided function is |
| 93 | invoked to transform the data into a dictionary of model inputs. |
| 94 | """ |
| 95 | |
| 96 | def wrapper(model_cls: N) -> N: |
| 97 | if "_processor_factory" in model_cls.__dict__: |
| 98 | logger.warning( |
| 99 | "Model class %s already has a multi-modal processor " |
| 100 | "registered to %s. It is overwritten by the new one.", |
| 101 | model_cls, |
| 102 | self, |
| 103 | ) |
| 104 | |
| 105 | model_cls._processor_factory = _ProcessorFactories( |
| 106 | info=info, |
| 107 | dummy_inputs=dummy_inputs, |
| 108 | processor=processor, |
| 109 | ) |
| 110 | |
| 111 | return model_cls |
| 112 | |
| 113 | return wrapper |
| 114 | |
| 115 | def _get_model_cls(self, model_config: "ModelConfig") -> "SupportsMultiModal": |
| 116 | # Avoid circular import |
| 117 | from vllm.model_executor.model_loader import get_model_architecture |
| 118 | |
| 119 | model_cls, _ = get_model_architecture(model_config) |
| 120 | if not hasattr(model_cls, "_processor_factory"): |
| 121 | raise ValueError( |
| 122 | f"Model class {model_cls.__name__} has no registered " |
| 123 | "multimodal processor" |
| 124 | ) |
| 125 | return cast("SupportsMultiModal", model_cls) |
| 126 | |
| 127 | def _create_processing_ctx( |
| 128 | self, |
| 129 | model_config: "ModelConfig", |
| 130 | tokenizer: TokenizerLike | None = None, |
| 131 | ) -> InputProcessingContext: |
| 132 | if tokenizer is None: |
| 133 | tokenizer = cached_tokenizer_from_config(model_config) |
| 134 | |
| 135 | return InputProcessingContext(model_config, tokenizer) |
| 136 | |
| 137 | def _create_processing_info( |
| 138 | self, |
| 139 | model_config: "ModelConfig", |
| 140 | tokenizer: TokenizerLike | None = None, |
| 141 | ) -> BaseProcessingInfo: |
| 142 | model_cls = self._get_model_cls(model_config) |
| 143 | factories = model_cls._processor_factory |
| 144 | ctx = self._create_processing_ctx(model_config, tokenizer) |
| 145 | return factories.info(ctx) |
| 146 | |
| 147 | def get_processing_info(self, model_config: "ModelConfig") -> BaseProcessingInfo: |
| 148 | return self._create_processing_info(model_config, tokenizer=None) |
| 149 | |
| 150 | def create_processor( |
| 151 | self, |
| 152 | model_config: "ModelConfig", |
| 153 | *, |
| 154 | tokenizer: TokenizerLike | None = None, |
| 155 | ) -> BaseMultiModalProcessor[BaseProcessingInfo]: |
| 156 | """Create a multi-modal processor for a specific model and tokenizer.""" |
| 157 | if not model_config.is_multimodal_model: |
| 158 | model_name = model_config.served_model_name or model_config.model |
| 159 | raise ValueError(f"{model_name} is not a multimodal model") |
| 160 | |
| 161 | model_cls = self._get_model_cls(model_config) |
| 162 | factories = model_cls._processor_factory |
| 163 | |
| 164 | ctx = self._create_processing_ctx(model_config, tokenizer) |
| 165 | |
| 166 | return factories.build_processor(ctx) |
| 167 | |
| 168 | |
| 169 | class MultiModalTimingRegistry: |
| 170 | def __init__(self, observability_config: "ObservabilityConfig | None") -> None: |
| 171 | super().__init__() |
| 172 | |
| 173 | if observability_config and observability_config.enable_mm_processor_stats: |
| 174 | self._lock = threading.Lock() |
| 175 | self._ctx_by_request_id = defaultdict[str, TimingContext](TimingContext) |
| 176 | self._enabled = True |
| 177 | else: |
| 178 | self._enabled = False |
| 179 | |
| 180 | def get(self, request_id: str) -> TimingContext: |
| 181 | if not self._enabled: |
| 182 | return TimingContext(enabled=False) |
| 183 | |
| 184 | with self._lock: |
| 185 | return self._ctx_by_request_id[request_id] |
| 186 | |
| 187 | def stat(self) -> dict[str, dict[str, float]]: |
| 188 | if not self._enabled: |
| 189 | return {} |
| 190 | |
| 191 | with self._lock: |
| 192 | stats = { |
| 193 | req_id: ctx.get_stats_dict() |
| 194 | for req_id, ctx in self._ctx_by_request_id.items() |
| 195 | } |
| 196 | self._ctx_by_request_id.clear() |
| 197 | return stats |