| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from io import BytesIO |
| 4 | from typing import TYPE_CHECKING, Final |
| 5 | |
| 6 | import pybase64 |
| 7 | import torch |
| 8 | |
| 9 | from vllm.exceptions import VLLMValidationError |
| 10 | from vllm.utils.async_utils import make_async |
| 11 | from vllm.utils.sparse_utils import ( |
| 12 | check_sparse_tensor_invariants_threadsafe, |
| 13 | safe_to_dense, |
| 14 | ) |
| 15 | |
| 16 | if TYPE_CHECKING: |
| 17 | from vllm.config import ModelConfig |
| 18 | |
| 19 | |
| 20 | # How much of torch's own reason is repeated back to the caller. The reason is |
| 21 | # built from bytes the caller sent, so it is bounded rather than echoed whole. |
| 22 | _MAX_EMBED_ERROR_REASON_CHARS: Final = 200 |
| 23 | |
| 24 | |
| 25 | def _truncated_reason(exc: Exception) -> str: |
| 26 | """`exc`'s message, capped, saying how much was left out when it is capped. |
| 27 | |
| 28 | Without the count a capped reason reads like the whole one, so a caller |
| 29 | debugging a large payload cannot tell that torch said more than this. |
| 30 | """ |
| 31 | reason: Final = str(exc).strip() or type(exc).__name__ |
| 32 | omitted: Final = len(reason) - _MAX_EMBED_ERROR_REASON_CHARS |
| 33 | if omitted <= 0: |
| 34 | return reason |
| 35 | return ( |
| 36 | f"{reason[:_MAX_EMBED_ERROR_REASON_CHARS]}" |
| 37 | f"... ({omitted} more characters truncated)" |
| 38 | ) |
| 39 | |
| 40 | |
| 41 | def safe_load_prompt_embeds( |
| 42 | model_config: "ModelConfig", |
| 43 | embed: bytes, |
| 44 | ) -> torch.Tensor: |
| 45 | if not model_config.enable_prompt_embeds: |
| 46 | raise VLLMValidationError( |
| 47 | "You must set `--enable-prompt-embeds` to input `prompt_embeds`.", |
| 48 | parameter="prompt_embeds", |
| 49 | ) |
| 50 | |
| 51 | # Decoded outside the try below so that a body which is not base64 keeps |
| 52 | # raising binascii.Error, which is a ValueError and already answered 400. |
| 53 | payload: Final = pybase64.b64decode(embed, validate=True) |
| 54 | |
| 55 | with check_sparse_tensor_invariants_threadsafe(): |
| 56 | try: |
| 57 | tensor = torch.load( |
| 58 | BytesIO(payload), |
| 59 | weights_only=True, |
| 60 | map_location=torch.device("cpu"), |
| 61 | ) |
| 62 | except (torch.OutOfMemoryError, MemoryError): |
| 63 | # Running out of memory is the server's condition, not the caller's. |
| 64 | raise |
| 65 | except Exception as exc: |
| 66 | # Everything else here is about the bytes the caller sent, and torch |
| 67 | # has no single error type for "this is not a tensor file": the zip |
| 68 | # reader raises `RuntimeError`, the legacy pickle path raises |
| 69 | # `UnpicklingError`, `KeyError` or `IndexError`, an empty payload |
| 70 | # raises `EOFError`. None of them is a `ValueError`, so the |
| 71 | # entrypoints' fallback mapped them all to 500. Catching by |
| 72 | # behaviour rather than by type also survives torch changing which |
| 73 | # error it raises for a given malformed payload. |
| 74 | # |
| 75 | # torch's reason is worth keeping -- for a malformed sparse tensor |
| 76 | # it names the offending index -- but it is built from the caller's |
| 77 | # own bytes, so it is truncated rather than echoed whole. |
| 78 | raise VLLMValidationError( |
| 79 | "`prompt_embeds` could not be deserialized as a torch tensor: " |
| 80 | f"{_truncated_reason(exc)}", |
| 81 | parameter="prompt_embeds", |
| 82 | ) from exc |
| 83 | tensor = safe_to_dense(tensor, parameter="prompt_embeds") |
| 84 | |
| 85 | if tensor.dim() > 2: |
| 86 | tensor = tensor.squeeze(0) |
| 87 | if tensor.dim() != 2: |
| 88 | raise VLLMValidationError( |
| 89 | "`prompt_embeds` must be a 2D tensor of shape " |
| 90 | f"(num_tokens, hidden_size); got shape {tuple(tensor.shape)}.", |
| 91 | parameter="prompt_embeds", |
| 92 | ) |
| 93 | |
| 94 | # Pin each tensor to the model's hidden_size. Validating here |
| 95 | # also transitively guarantees cross-tensor consistency for requests that |
| 96 | # include multiple `prompt_embeds` parts, which is required by downstream |
| 97 | # concatenation in `_build_mixed_prompt_embeds`. |
| 98 | expected_hidden_size = model_config.get_hidden_size() |
| 99 | if tensor.shape[1] != expected_hidden_size: |
| 100 | raise VLLMValidationError( |
| 101 | f"`prompt_embeds` hidden_size {tensor.shape[1]} does not match " |
| 102 | f"the model's hidden_size {expected_hidden_size}.", |
| 103 | parameter="prompt_embeds", |
| 104 | ) |
| 105 | |
| 106 | # Cast to the model's dtype so API clients don't need to know the server's |
| 107 | # `--dtype` setting ahead of time. Only floating-point source dtypes are |
| 108 | # allowed. integer / bool / complex inputs almost certainly indicate caller |
| 109 | # error (e.g. quantized payloads, wrong tensor), and a silent `.to()` |
| 110 | # could hide a real mistake. |
| 111 | expected_dtype = model_config.dtype |
| 112 | if tensor.dtype != expected_dtype: |
| 113 | if not tensor.is_floating_point(): |
| 114 | raise VLLMValidationError( |
| 115 | f"`prompt_embeds` dtype {tensor.dtype} is not a floating-point " |
| 116 | f"type, cannot safely cast to the model's dtype {expected_dtype}.", |
| 117 | parameter="prompt_embeds", |
| 118 | ) |
| 119 | tensor = tensor.to(expected_dtype) |
| 120 | |
| 121 | return tensor |
| 122 | |
| 123 | |
| 124 | safe_load_prompt_embeds_async = make_async(safe_load_prompt_embeds) |
| 125 | """Async variant of `safe_load_prompt_embeds` that defers the decode to a |
| 126 | thread-pool executor, so the asyncio event loop is not blocked by the base64 |
| 127 | decode + `torch.load` work.""" |