| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Loading of custom logits processor classes for the V2 model runner. |
| 4 | |
| 5 | This module is deliberately import-light: the frontend process uses it to |
| 6 | validate per-request params, so importing it must not pull in model-runner |
| 7 | side modules (torch, triton, worker state). Imports needed to instantiate |
| 8 | processors are deferred to ``build_custom_logits_processors``. |
| 9 | """ |
| 10 | |
| 11 | import importlib |
| 12 | from collections.abc import Callable, Sequence |
| 13 | from functools import lru_cache |
| 14 | from importlib.metadata import entry_points |
| 15 | from typing import TYPE_CHECKING |
| 16 | |
| 17 | from vllm.exceptions import VLLMValidationError |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.v1.worker.gpu.sample.logits_processor.interface import ( |
| 20 | LogitsProcessor, |
| 21 | LogitsProcRequestState, |
| 22 | ) |
| 23 | |
| 24 | if TYPE_CHECKING: |
| 25 | from vllm.config import VllmConfig |
| 26 | from vllm.sampling_params import SamplingParams |
| 27 | from vllm.v1.worker.gpu.states import RequestState |
| 28 | |
| 29 | logger = init_logger(__name__) |
| 30 | |
| 31 | # Same entry-points group as |
| 32 | # vllm.v1.sample.logits_processor.LOGITSPROCS_GROUP; duplicated here so this |
| 33 | # module stays importable without the V1 sampler machinery. |
| 34 | LOGITSPROCS_GROUP = "vllm.logits_processors" |
| 35 | |
| 36 | |
| 37 | def _load_v2_logitsprocs_plugins() -> list[type[LogitsProcessor]]: |
| 38 | """Load installed logits processor plugins. |
| 39 | |
| 40 | Plugins registered under ``LOGITSPROCS_GROUP`` ship V1-interface |
| 41 | processors today; those crash under the V2 sampler, so reject them at |
| 42 | load time. |
| 43 | """ |
| 44 | from vllm.utils.torch_utils import guard_cuda_initialization |
| 45 | |
| 46 | installed_logitsprocs_plugins = entry_points(group=LOGITSPROCS_GROUP) |
| 47 | if len(installed_logitsprocs_plugins) == 0: |
| 48 | logger.debug("No logitsprocs plugins installed (group %s).", LOGITSPROCS_GROUP) |
| 49 | return [] |
| 50 | |
| 51 | logger.debug("Loading installed logitsprocs plugins (group %s):", LOGITSPROCS_GROUP) |
| 52 | classes: list[type[LogitsProcessor]] = [] |
| 53 | for entrypoint in installed_logitsprocs_plugins: |
| 54 | try: |
| 55 | logger.debug( |
| 56 | "- Loading logitproc plugin entrypoint=%s target=%s", |
| 57 | entrypoint.name, |
| 58 | entrypoint.value, |
| 59 | ) |
| 60 | with guard_cuda_initialization(): |
| 61 | cls = entrypoint.load() |
| 62 | except Exception as e: |
| 63 | logger.error("Failed to load LogitsProcessor plugin %s: %s", entrypoint, e) |
| 64 | raise RuntimeError( |
| 65 | f"Failed to load LogitsProcessor plugin {entrypoint}" |
| 66 | ) from e |
| 67 | if not (isinstance(cls, type) and issubclass(cls, LogitsProcessor)): |
| 68 | raise ValueError( |
| 69 | f"LogitsProcessor plugin {entrypoint.name} is not a subclass of " |
| 70 | "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor; " |
| 71 | "V1-interface plugins are not supported by the V2 model runner." |
| 72 | ) |
| 73 | classes.append(cls) |
| 74 | return classes |
| 75 | |
| 76 | |
| 77 | def _load_v2_logitsprocs_by_fqcns( |
| 78 | logits_processors: Sequence[str | type], |
| 79 | ) -> list[type[LogitsProcessor]]: |
| 80 | """Resolve a mixed list of processor types and FQCN strings into types. |
| 81 | |
| 82 | FQCN syntax is <module>:<type> i.e. x.y.z:CustomLogitProc. |
| 83 | """ |
| 84 | from vllm.utils.torch_utils import guard_cuda_initialization |
| 85 | |
| 86 | classes: list[type[LogitsProcessor]] = [] |
| 87 | for ldx, logitproc in enumerate(logits_processors): |
| 88 | if isinstance(logitproc, type): |
| 89 | logger.debug(" - Already-loaded logit processor: %s", logitproc.__name__) |
| 90 | if not issubclass(logitproc, LogitsProcessor): |
| 91 | raise ValueError( |
| 92 | f"{logitproc.__name__} is not a subclass of " |
| 93 | "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor" |
| 94 | ) |
| 95 | classes.append(logitproc) |
| 96 | continue |
| 97 | |
| 98 | logger.debug("- Loading logits processor %s", logitproc) |
| 99 | parts = logitproc.split(":") |
| 100 | if len(parts) != 2: |
| 101 | raise ValueError( |
| 102 | f"Invalid logits processor FQCN {logitproc!r}. " |
| 103 | "Expected format: '<module>:<type>'" |
| 104 | ) |
| 105 | module_path, qualname = parts |
| 106 | |
| 107 | try: |
| 108 | with guard_cuda_initialization(): |
| 109 | module = importlib.import_module(module_path) |
| 110 | except Exception as e: |
| 111 | logger.error( |
| 112 | "Failed to load %sth LogitsProcessor plugin %s: %s", |
| 113 | ldx, |
| 114 | logitproc, |
| 115 | e, |
| 116 | ) |
| 117 | raise RuntimeError( |
| 118 | f"Failed to load {ldx}th LogitsProcessor plugin {logitproc}" |
| 119 | ) from e |
| 120 | |
| 121 | obj = module |
| 122 | for attr in qualname.split("."): |
| 123 | obj = getattr(obj, attr) |
| 124 | if not isinstance(obj, type): |
| 125 | raise ValueError("Loaded logit processor must be a type.") |
| 126 | if not issubclass(obj, LogitsProcessor): |
| 127 | raise ValueError( |
| 128 | f"{obj.__name__} is not a subclass of " |
| 129 | "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor" |
| 130 | ) |
| 131 | classes.append(obj) |
| 132 | |
| 133 | return classes |
| 134 | |
| 135 | |
| 136 | def _load_v2_logitsprocs( |
| 137 | logits_processors: Sequence[str | type] | None, |
| 138 | ) -> list[type[LogitsProcessor]]: |
| 139 | """Load all custom logits processors for the V2 model runner. |
| 140 | |
| 141 | Combines installed plugins (validated as V2 processors) with the |
| 142 | user-specified list of processor types and FQCN strings. |
| 143 | """ |
| 144 | return _load_v2_logitsprocs_plugins() + _load_v2_logitsprocs_by_fqcns( |
| 145 | logits_processors or () |
| 146 | ) |
| 147 | |
| 148 | |
| 149 | @lru_cache |
| 150 | def _cached_load_v2_logitsprocs( |
| 151 | custom_logitsprocs: tuple[str | type, ...], |
| 152 | ) -> list[type[LogitsProcessor]]: |
| 153 | return _load_v2_logitsprocs(list(custom_logitsprocs)) |
| 154 | |
| 155 | |
| 156 | def build_custom_logits_processors( |
| 157 | vllm_config: "VllmConfig", |
| 158 | req_states: "RequestState", |
| 159 | is_pooling_model: bool, |
| 160 | custom_logitsprocs: Sequence[str | type] = (), |
| 161 | ) -> list[LogitsProcessor]: |
| 162 | """Load and instantiate custom logits processors, entrypoint plugins first. |
| 163 | |
| 164 | Raises: |
| 165 | ValueError: if a pooling model specifies custom processors, or a |
| 166 | loaded class does not implement the V2 interface. |
| 167 | RuntimeError: if an FQCN fails to import. |
| 168 | |
| 169 | """ |
| 170 | from vllm.v1.sample.logits_processor import STR_POOLING_REJECTS_LOGITSPROCS |
| 171 | |
| 172 | if is_pooling_model: |
| 173 | if custom_logitsprocs: |
| 174 | raise ValueError(STR_POOLING_REJECTS_LOGITSPROCS) |
| 175 | logger.debug( |
| 176 | "Skipping logits processor loading because pooling models" |
| 177 | " do not support logits processors." |
| 178 | ) |
| 179 | return [] |
| 180 | |
| 181 | custom_logitsprocs_classes = _load_v2_logitsprocs(custom_logitsprocs) |
| 182 | lp_req_state = LogitsProcRequestState.from_request_state(req_states) |
| 183 | return [ctor(vllm_config, lp_req_state) for ctor in custom_logitsprocs_classes] |
| 184 | |
| 185 | |
| 186 | def build_custom_logits_processors_params_validator( |
| 187 | custom_logitsprocs: Sequence[str | type] | None, |
| 188 | ) -> "Callable[[SamplingParams], None]": |
| 189 | """Load custom processor classes once and return a params validator. |
| 190 | |
| 191 | Called from the frontend at startup. The returned callable runs each |
| 192 | processor's ``validate_params`` at request admission. |
| 193 | |
| 194 | Raises: |
| 195 | ValueError: if a loaded class does not implement the V2 interface. |
| 196 | RuntimeError: if an FQCN fails to import. |
| 197 | |
| 198 | """ |
| 199 | classes = _cached_load_v2_logitsprocs(tuple(custom_logitsprocs or ())) |
| 200 | |
| 201 | if not classes: |
| 202 | return lambda _: None |
| 203 | |
| 204 | def validate_params(sampling_params: "SamplingParams") -> None: |
| 205 | for cls in classes: |
| 206 | try: |
| 207 | cls.validate_params(sampling_params) |
| 208 | except ValueError as e: |
| 209 | raise VLLMValidationError(str(e)) from e |
| 210 | |
| 211 | return validate_params |