| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import dataclasses |
| 5 | import importlib |
| 6 | import pickle |
| 7 | from abc import ABC, abstractmethod |
| 8 | from collections.abc import Callable, Sequence |
| 9 | from functools import partial |
| 10 | from inspect import isclass |
| 11 | from types import FunctionType |
| 12 | from typing import Any, ClassVar, TypeAlias, cast, get_type_hints |
| 13 | |
| 14 | import cloudpickle |
| 15 | import msgspec |
| 16 | import numpy as np |
| 17 | import torch |
| 18 | import zmq |
| 19 | from msgspec import msgpack |
| 20 | from pydantic import GetCoreSchemaHandler |
| 21 | from pydantic_core import core_schema |
| 22 | |
| 23 | from vllm import envs |
| 24 | from vllm.logger import init_logger |
| 25 | from vllm.multimodal.inputs import ( |
| 26 | BaseMultiModalField, |
| 27 | MultiModalBatchedField, |
| 28 | MultiModalFieldConfig, |
| 29 | MultiModalFieldElem, |
| 30 | MultiModalFlatField, |
| 31 | MultiModalKwargsItem, |
| 32 | MultiModalKwargsItems, |
| 33 | MultiModalSharedField, |
| 34 | NestedTensors, |
| 35 | ) |
| 36 | from vllm.utils.torch_utils import PIN_MEMORY |
| 37 | from vllm.v1.utils import tensor_data |
| 38 | |
| 39 | logger = init_logger(__name__) |
| 40 | |
| 41 | CUSTOM_TYPE_PICKLE = 1 |
| 42 | CUSTOM_TYPE_CLOUDPICKLE = 2 |
| 43 | CUSTOM_TYPE_RAW_VIEW = 3 |
| 44 | |
| 45 | # MultiModalField class serialization type map. |
| 46 | # These need to list all possible field types and match them |
| 47 | # to factory methods in `MultiModalFieldConfig`. |
| 48 | MMF_CLASS_TO_FACTORY: dict[type[BaseMultiModalField], str] = { |
| 49 | MultiModalFlatField: "flat", |
| 50 | MultiModalSharedField: "shared", |
| 51 | MultiModalBatchedField: "batched", |
| 52 | } |
| 53 | |
| 54 | bytestr: TypeAlias = bytes | bytearray | memoryview | zmq.Frame |
| 55 | |
| 56 | |
| 57 | class OOBTensorConsumer(ABC): |
| 58 | @abstractmethod |
| 59 | def __call__(self, tensor: torch.Tensor) -> dict | None: |
| 60 | """Called with tensors for the current message. |
| 61 | Returns None to reject the tensor (falls back to regular serialization), |
| 62 | otherwise a dict with arbitrary placeholder data to be included |
| 63 | in the serialized message. |
| 64 | """ |
| 65 | return None |
| 66 | |
| 67 | @abstractmethod |
| 68 | def new_message(self) -> None: |
| 69 | """Called at the start of each new encoded message.""" |
| 70 | pass |
| 71 | |
| 72 | |
| 73 | # dtype, shape, metadata -> tensor |
| 74 | OOBTensorProvider = Callable[[str, tuple[int, ...], dict], torch.Tensor] |
| 75 | |
| 76 | |
| 77 | def _log_insecure_serialization_warning(): |
| 78 | logger.warning_once( |
| 79 | "Allowing insecure serialization using pickle due to " |
| 80 | "VLLM_ALLOW_INSECURE_SERIALIZATION=1" |
| 81 | ) |
| 82 | |
| 83 | |
| 84 | def _typestr(val: Any) -> tuple[str, str] | None: |
| 85 | if val is None: |
| 86 | return None |
| 87 | t = type(val) |
| 88 | return t.__module__, t.__qualname__ |
| 89 | |
| 90 | |
| 91 | def _encode_type_info_recursive(obj: Any) -> Any: |
| 92 | """Recursively encode type information for nested structures of |
| 93 | lists/dicts.""" |
| 94 | if obj is None: |
| 95 | return None |
| 96 | if type(obj) is list: |
| 97 | return [_encode_type_info_recursive(item) for item in obj] |
| 98 | if type(obj) is dict: |
| 99 | return {k: _encode_type_info_recursive(v) for k, v in obj.items()} |
| 100 | return _typestr(obj) |
| 101 | |
| 102 | |
| 103 | def _decode_type_info_recursive( |
| 104 | type_info: Any, data: Any, convert_fn: Callable[[Sequence[str], Any], Any] |
| 105 | ) -> Any: |
| 106 | """Recursively decode type information for nested structures of |
| 107 | lists/dicts.""" |
| 108 | if type_info is None: |
| 109 | return data |
| 110 | if isinstance(type_info, dict): |
| 111 | assert isinstance(data, dict) |
| 112 | return { |
| 113 | k: _decode_type_info_recursive(type_info[k], data[k], convert_fn) |
| 114 | for k in type_info |
| 115 | } |
| 116 | if isinstance(type_info, list) and ( |
| 117 | # Exclude serialized tensors/numpy arrays. |
| 118 | len(type_info) != 2 or not isinstance(type_info[0], str) |
| 119 | ): |
| 120 | assert isinstance(data, list) |
| 121 | return [ |
| 122 | _decode_type_info_recursive(ti, d, convert_fn) |
| 123 | for ti, d in zip(type_info, data) |
| 124 | ] |
| 125 | return convert_fn(type_info, data) |
| 126 | |
| 127 | |
| 128 | class UtilityResult: |
| 129 | """Wrapper for special handling when serializing/deserializing.""" |
| 130 | |
| 131 | def __init__(self, r: Any = None): |
| 132 | self.result = r |
| 133 | |
| 134 | |
| 135 | class MsgpackEncoder: |
| 136 | """Encoder with custom torch tensor and numpy array serialization. |
| 137 | |
| 138 | Note that unlike vanilla `msgspec` Encoders, this interface is generally |
| 139 | not thread-safe when encoding tensors / numpy arrays. |
| 140 | |
| 141 | By default, arrays below 256B are serialized inline Larger will get sent |
| 142 | via dedicated messages. Note that this is a per-tensor limit. |
| 143 | |
| 144 | When a ``oob_tensor_consumer`` is provided, tensors (CUDA and CPU) will be |
| 145 | offered to it for out-of-band handling. |
| 146 | """ |
| 147 | |
| 148 | def __init__( |
| 149 | self, |
| 150 | size_threshold: int | None = None, |
| 151 | oob_tensor_consumer: OOBTensorConsumer | None = None, |
| 152 | ): |
| 153 | if size_threshold is None: |
| 154 | size_threshold = envs.VLLM_MSGPACK_ZERO_COPY_THRESHOLD |
| 155 | self.encoder = msgpack.Encoder(enc_hook=self.enc_hook) |
| 156 | # This is used as a local stash of buffers that we can then access from |
| 157 | # our custom `msgspec` hook, `enc_hook`. We don't have a way to |
| 158 | # pass custom data to the hook otherwise. |
| 159 | self.aux_buffers: list[bytestr] | None = None |
| 160 | self.size_threshold = size_threshold |
| 161 | self.oob_tensor_consumer = oob_tensor_consumer |
| 162 | if envs.VLLM_ALLOW_INSECURE_SERIALIZATION: |
| 163 | _log_insecure_serialization_warning() |
| 164 | |
| 165 | def encode(self, obj: Any) -> Sequence[bytestr]: |
| 166 | try: |
| 167 | if self.oob_tensor_consumer is not None: |
| 168 | self.oob_tensor_consumer.new_message() |
| 169 | self.aux_buffers = bufs = [b""] |
| 170 | bufs[0] = self.encoder.encode(obj) |
| 171 | # This `bufs` list allows us to collect direct pointers to backing |
| 172 | # buffers of tensors and np arrays, and return them along with the |
| 173 | # top-level encoded buffer instead of copying their data into the |
| 174 | # new buffer. |
| 175 | return bufs |
| 176 | finally: |
| 177 | self.aux_buffers = None |
| 178 | |
| 179 | def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]: |
| 180 | try: |
| 181 | if self.oob_tensor_consumer is not None: |
| 182 | self.oob_tensor_consumer.new_message() |
| 183 | self.aux_buffers = [buf] |
| 184 | bufs = self.aux_buffers |
| 185 | self.encoder.encode_into(obj, buf) |
| 186 | return bufs |
| 187 | finally: |
| 188 | self.aux_buffers = None |
| 189 | |
| 190 | def enc_hook(self, obj: Any) -> Any: |
| 191 | if isinstance(obj, torch.Tensor): |
| 192 | return self._encode_tensor(obj) |
| 193 | |
| 194 | # Fall back to pickle for object or void kind ndarrays. |
| 195 | if isinstance(obj, np.ndarray) and obj.dtype.kind not in ("O", "V"): |
| 196 | return self._encode_ndarray(obj) |
| 197 | |
| 198 | if isinstance(obj, slice): |
| 199 | # We are assuming only int-based values will be used here. |
| 200 | return tuple( |
| 201 | int(v) if v is not None else None |
| 202 | for v in (obj.start, obj.stop, obj.step) |
| 203 | ) |
| 204 | |
| 205 | if isinstance(obj, MultiModalKwargsItem): |
| 206 | return self._encode_mm_item(obj) |
| 207 | |
| 208 | if isinstance(obj, MultiModalKwargsItems): |
| 209 | return self._encode_mm_items(obj) |
| 210 | |
| 211 | if isinstance(obj, UtilityResult): |
| 212 | result = obj.result |
| 213 | if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION: |
| 214 | return None, result |
| 215 | # Since utility results are not strongly typed, we recursively |
| 216 | # encode type information for nested structures of lists/dicts |
| 217 | # to help with correct msgspec deserialization. |
| 218 | return _encode_type_info_recursive(result), result |
| 219 | |
| 220 | if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION: |
| 221 | raise TypeError( |
| 222 | f"Object of type {type(obj)} is not serializable" |
| 223 | "Set VLLM_ALLOW_INSECURE_SERIALIZATION=1 to allow " |
| 224 | "fallback to pickle-based serialization." |
| 225 | ) |
| 226 | |
| 227 | if isinstance(obj, FunctionType): |
| 228 | # `pickle` is generally faster than cloudpickle, but can have |
| 229 | # problems serializing methods. |
| 230 | return msgpack.Ext(CUSTOM_TYPE_CLOUDPICKLE, cloudpickle.dumps(obj)) |
| 231 | |
| 232 | return msgpack.Ext( |
| 233 | CUSTOM_TYPE_PICKLE, pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) |
| 234 | ) |
| 235 | |
| 236 | def _encode_ndarray( |
| 237 | self, obj: np.ndarray |
| 238 | ) -> tuple[str, tuple[int, ...], int | memoryview]: |
| 239 | assert self.aux_buffers is not None |
| 240 | # If the array is non-contiguous, we need to copy it first |
| 241 | arr_data = obj.data if obj.flags.c_contiguous else obj.tobytes() |
| 242 | if not obj.shape or obj.nbytes < self.size_threshold: |
| 243 | # Encode small arrays and scalars inline. Using this extension type |
| 244 | # ensures we can avoid copying when decoding. |
| 245 | data = msgpack.Ext(CUSTOM_TYPE_RAW_VIEW, arr_data) |
| 246 | else: |
| 247 | # Otherwise encode index of backing buffer to avoid copy. |
| 248 | data = len(self.aux_buffers) |
| 249 | self.aux_buffers.append(arr_data) |
| 250 | |
| 251 | # We serialize the ndarray as a tuple of native types. |
| 252 | # The data is either inlined if small, or an index into a list of |
| 253 | # backing buffers that we've stashed in `aux_buffers`. |
| 254 | return obj.dtype.str, obj.shape, data |
| 255 | |
| 256 | def _encode_tensor( |
| 257 | self, obj: torch.Tensor |
| 258 | ) -> tuple[str, tuple[int, ...], int | dict | memoryview]: |
| 259 | oob_consumer = self.oob_tensor_consumer |
| 260 | # view the tensor as a contiguous 1D array of bytes |
| 261 | if obj.nbytes < self.size_threshold and obj.is_cpu: |
| 262 | # Smaller tensors are encoded inline, just like ndarrays. |
| 263 | data = msgpack.Ext(CUSTOM_TYPE_RAW_VIEW, tensor_data(obj)) |
| 264 | elif oob_consumer is not None and (data := oob_consumer(obj)) is not None: |
| 265 | assert isinstance(data, dict) |
| 266 | else: |
| 267 | # Otherwise encode index of backing buffer to avoid copy. |
| 268 | assert self.aux_buffers is not None |
| 269 | data = len(self.aux_buffers) |
| 270 | self.aux_buffers.append(tensor_data(obj)) |
| 271 | dtype = str(obj.dtype).removeprefix("torch.") |
| 272 | return dtype, obj.shape, data |
| 273 | |
| 274 | def _encode_mm_items(self, items: MultiModalKwargsItems) -> dict[str, Any]: |
| 275 | return { |
| 276 | modality: [self._encode_mm_item(item) for item in itemlist] |
| 277 | for modality, itemlist in items.items() |
| 278 | } |
| 279 | |
| 280 | def _encode_mm_item(self, item: MultiModalKwargsItem) -> dict[str, Any]: |
| 281 | return {key: self._encode_mm_field_elem(elem) for key, elem in item.items()} |
| 282 | |
| 283 | def _encode_mm_field_elem(self, elem: MultiModalFieldElem) -> dict[str, Any]: |
| 284 | return { |
| 285 | "data": ( |
| 286 | None if elem.data is None else self._encode_nested_tensors(elem.data) |
| 287 | ), |
| 288 | "field": self._encode_mm_field(elem.field), |
| 289 | } |
| 290 | |
| 291 | def _encode_nested_tensors(self, nt: NestedTensors) -> Any: |
| 292 | if isinstance(nt, torch.Tensor): |
| 293 | return self._encode_tensor(nt) |
| 294 | if isinstance(nt, (int, float)): |
| 295 | # Although it violates NestedTensors type, MultiModalKwargs |
| 296 | # values are sometimes floats. |
| 297 | return nt |
| 298 | return [self._encode_nested_tensors(x) for x in nt] |
| 299 | |
| 300 | def _encode_mm_field(self, field: BaseMultiModalField): |
| 301 | # Figure out the factory name for the field type. |
| 302 | name = MMF_CLASS_TO_FACTORY.get(field.__class__) |
| 303 | if not name: |
| 304 | raise TypeError(f"Unsupported field type: {field.__class__}") |
| 305 | |
| 306 | # We just need to copy all of the field values in order |
| 307 | # which will be then used to reconstruct the field. |
| 308 | factory_kw = {f.name: getattr(field, f.name) for f in dataclasses.fields(field)} |
| 309 | return name, factory_kw |
| 310 | |
| 311 | |
| 312 | class MsgpackDecoder: |
| 313 | """Decoder with custom torch tensor and numpy array serialization. |
| 314 | |
| 315 | Note that unlike vanilla `msgspec` Decoders, this interface is generally |
| 316 | not thread-safe when encoding tensors / numpy arrays. |
| 317 | |
| 318 | ``oob_tensor_provider`` must be used when an OOBTensorConsumer is used on the |
| 319 | encoder side. |
| 320 | """ |
| 321 | |
| 322 | def __init__( |
| 323 | self, |
| 324 | t: Any | None = None, |
| 325 | share_mem: bool = True, |
| 326 | oob_tensor_provider: OOBTensorProvider | None = None, |
| 327 | ): |
| 328 | self.share_mem = share_mem |
| 329 | self.pin_tensors = PIN_MEMORY |
| 330 | args = () if t is None else (t,) |
| 331 | self.decoder = msgpack.Decoder( |
| 332 | *args, ext_hook=self.ext_hook, dec_hook=self.dec_hook |
| 333 | ) |
| 334 | self.aux_buffers: Sequence[bytestr] = () |
| 335 | self.oob_tensor_provider = oob_tensor_provider |
| 336 | if envs.VLLM_ALLOW_INSECURE_SERIALIZATION: |
| 337 | _log_insecure_serialization_warning() |
| 338 | |
| 339 | def decode(self, bufs: bytestr | Sequence[bytestr]) -> Any: |
| 340 | if isinstance(bufs, bytestr): # type: ignore |
| 341 | return self.decoder.decode(bufs) |
| 342 | |
| 343 | self.aux_buffers = bufs |
| 344 | try: |
| 345 | return self.decoder.decode(bufs[0]) |
| 346 | finally: |
| 347 | self.aux_buffers = () |
| 348 | |
| 349 | def dec_hook(self, t: type, obj: Any) -> Any: |
| 350 | # Given native types in `obj`, convert to type `t`. |
| 351 | if isclass(t): |
| 352 | if issubclass(t, np.ndarray): |
| 353 | return self._decode_ndarray(obj) |
| 354 | if issubclass(t, torch.Tensor): |
| 355 | return self._decode_tensor(obj) |
| 356 | if t is slice: |
| 357 | return slice(*obj) |
| 358 | if issubclass(t, MultiModalKwargsItem): |
| 359 | return self._decode_mm_item(obj) |
| 360 | if issubclass(t, MultiModalKwargsItems): |
| 361 | return self._decode_mm_items(obj) |
| 362 | if t is UtilityResult: |
| 363 | return self._decode_utility_result(obj) |
| 364 | return obj |
| 365 | |
| 366 | def _decode_utility_result(self, obj: Any) -> UtilityResult: |
| 367 | result_type, result = obj |
| 368 | if result_type is not None: |
| 369 | if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION: |
| 370 | raise TypeError( |
| 371 | "VLLM_ALLOW_INSECURE_SERIALIZATION must " |
| 372 | "be set to use custom utility result types" |
| 373 | ) |
| 374 | # Use recursive decoding to handle nested structures |
| 375 | result = _decode_type_info_recursive( |
| 376 | result_type, result, self._convert_result |
| 377 | ) |
| 378 | return UtilityResult(result) |
| 379 | |
| 380 | def _convert_result(self, result_type: Sequence[str], result: Any) -> Any: |
| 381 | if result_type is None: |
| 382 | return result |
| 383 | mod_name, name = result_type |
| 384 | mod = importlib.import_module(mod_name) |
| 385 | result_type = getattr(mod, name) |
| 386 | return msgspec.convert(result, result_type, dec_hook=self.dec_hook) |
| 387 | |
| 388 | def _decode_ndarray(self, arr: Any) -> np.ndarray: |
| 389 | dtype, shape, data = arr |
| 390 | # zero-copy decode. We assume the ndarray will not be kept around, |
| 391 | # as it now locks the whole received message buffer in memory. |
| 392 | buffer = self.aux_buffers[data] if isinstance(data, int) else data |
| 393 | arr = np.frombuffer(buffer, dtype=dtype) |
| 394 | if not self.share_mem: |
| 395 | arr = arr.copy() |
| 396 | return arr.reshape(shape) |
| 397 | |
| 398 | def _decode_tensor(self, arr: Any) -> torch.Tensor: |
| 399 | dtype, shape, data = arr |
| 400 | if isinstance(data, dict): |
| 401 | assert self.oob_tensor_provider, ( |
| 402 | "Received OOB tensor but tensor provider is not set" |
| 403 | ) |
| 404 | return self.oob_tensor_provider(dtype, shape, data) |
| 405 | |
| 406 | is_aux = isinstance(data, int) |
| 407 | buffer = self.aux_buffers[data] if is_aux else data |
| 408 | buffer = buffer if isinstance(buffer, memoryview) else memoryview(buffer) |
| 409 | torch_dtype = getattr(torch, dtype) |
| 410 | assert isinstance(torch_dtype, torch.dtype) |
| 411 | if not buffer.nbytes: # torch.frombuffer doesn't like empty buffers |
| 412 | assert 0 in shape |
| 413 | return torch.empty(shape, dtype=torch_dtype) |
| 414 | # Create uint8 array |
| 415 | arr = torch.frombuffer(buffer, dtype=torch.uint8) |
| 416 | # Clone ensures tensor is backed by pytorch-owned memory for safe |
| 417 | # future async CPU->GPU transfer. |
| 418 | # Pin larger tensors for more efficient CPU->GPU transfer. |
| 419 | if not is_aux: |
| 420 | arr = arr.clone() |
| 421 | elif not self.share_mem: |
| 422 | arr = arr.pin_memory() if self.pin_tensors else arr.clone() |
| 423 | # Convert back to proper shape & type |
| 424 | return arr.view(torch_dtype).view(shape) |
| 425 | |
| 426 | def _decode_mm_items(self, obj: dict[str, Any]) -> MultiModalKwargsItems: |
| 427 | return MultiModalKwargsItems( |
| 428 | { |
| 429 | modality: [self._decode_mm_item(item) for item in itemlist] |
| 430 | for modality, itemlist in obj.items() |
| 431 | } |
| 432 | ) |
| 433 | |
| 434 | def _decode_mm_item(self, obj: dict[str, Any]) -> MultiModalKwargsItem: |
| 435 | return MultiModalKwargsItem( |
| 436 | {key: self._decode_mm_field_elem(elem) for key, elem in obj.items()} |
| 437 | ) |
| 438 | |
| 439 | def _decode_mm_field_elem(self, obj: dict[str, Any]) -> MultiModalFieldElem: |
| 440 | if obj["data"] is not None: |
| 441 | obj["data"] = self._decode_nested_tensors(obj["data"]) |
| 442 | |
| 443 | # Reconstruct the field processor using MultiModalFieldConfig |
| 444 | factory_meth_name, factory_kw = obj["field"] |
| 445 | factory_meth = getattr(MultiModalFieldConfig, factory_meth_name) |
| 446 | |
| 447 | # Special case: decode the union "slices" field of |
| 448 | # MultiModalFlatField |
| 449 | if factory_meth_name == "flat": |
| 450 | factory_kw["slices"] = self._decode_nested_slices(factory_kw["slices"]) |
| 451 | |
| 452 | obj["field"] = factory_meth("", **factory_kw).field |
| 453 | return MultiModalFieldElem(**obj) |
| 454 | |
| 455 | def _decode_nested_tensors(self, obj: Any) -> NestedTensors: |
| 456 | if isinstance(obj, (int, float)): |
| 457 | # Although it violates NestedTensors type, MultiModalKwargs |
| 458 | # values are sometimes floats. |
| 459 | return obj |
| 460 | if not isinstance(obj, list): |
| 461 | raise TypeError(f"Unexpected NestedTensors contents: {type(obj)}") |
| 462 | if obj and isinstance(obj[0], str): |
| 463 | return self._decode_tensor(obj) |
| 464 | return [self._decode_nested_tensors(x) for x in obj] |
| 465 | |
| 466 | def _decode_nested_slices(self, obj: Any) -> Any: |
| 467 | assert isinstance(obj, (list, tuple)) |
| 468 | if obj and not isinstance(obj[0], (list, tuple)): |
| 469 | return slice(*obj) |
| 470 | return [self._decode_nested_slices(x) for x in obj] |
| 471 | |
| 472 | def ext_hook(self, code: int, data: memoryview) -> Any: |
| 473 | if code == CUSTOM_TYPE_RAW_VIEW: |
| 474 | return data |
| 475 | |
| 476 | if envs.VLLM_ALLOW_INSECURE_SERIALIZATION: |
| 477 | if code == CUSTOM_TYPE_PICKLE: |
| 478 | return pickle.loads(data) |
| 479 | if code == CUSTOM_TYPE_CLOUDPICKLE: |
| 480 | return cloudpickle.loads(data) |
| 481 | |
| 482 | raise NotImplementedError(f"Extension type code {code} is not supported") |
| 483 | |
| 484 | |
| 485 | def run_method( |
| 486 | obj: Any, |
| 487 | method: str | bytes | Callable, |
| 488 | args: tuple[Any, ...], |
| 489 | kwargs: dict[str, Any], |
| 490 | ) -> Any: |
| 491 | """Run a method of an object with the given arguments and keyword arguments. |
| 492 | If the method is string, it will be converted to a method using getattr. |
| 493 | If the method is serialized bytes and will be deserialized using |
| 494 | cloudpickle. |
| 495 | If the method is a callable, it will be called directly. |
| 496 | """ |
| 497 | if isinstance(method, bytes): |
| 498 | func = partial(cloudpickle.loads(method), obj) |
| 499 | elif isinstance(method, str): |
| 500 | try: |
| 501 | func = getattr(obj, method) |
| 502 | except AttributeError: |
| 503 | raise NotImplementedError( |
| 504 | f"Method {method!r} is not implemented." |
| 505 | ) from None |
| 506 | else: |
| 507 | func = partial(method, obj) # type: ignore |
| 508 | return func(*args, **kwargs) |
| 509 | |
| 510 | |
| 511 | class PydanticMsgspecMixin: |
| 512 | """Make a ``msgspec.Struct`` compatible with Pydantic for both |
| 513 | **validation** (JSON/dict -> Struct) and **serialization** |
| 514 | (Struct -> JSON-safe dict). |
| 515 | |
| 516 | Subclasses may set ``__pydantic_msgspec_exclude__`` (a ``set[str]``) |
| 517 | to list non-underscore field names that should also be stripped from |
| 518 | serialized output. Fields whose names start with ``_`` are always |
| 519 | excluded automatically. |
| 520 | """ |
| 521 | |
| 522 | # Subclasses can override to exclude additional public-but-internal keys. |
| 523 | __pydantic_msgspec_exclude__: ClassVar[set[str]] = set() |
| 524 | |
| 525 | @classmethod |
| 526 | def __get_pydantic_core_schema__( |
| 527 | cls, source_type: Any, handler: GetCoreSchemaHandler |
| 528 | ) -> core_schema.CoreSchema: |
| 529 | """Make msgspec.Struct compatible with Pydantic, respecting defaults. |
| 530 | Handle JSON=>msgspec.Struct. Used when exposing msgspec.Struct to the |
| 531 | API as input or in `/docs`. Note this is cached by Pydantic and not |
| 532 | called on every validation. |
| 533 | """ |
| 534 | msgspec_fields = {f.name: f for f in msgspec.structs.fields(source_type)} |
| 535 | type_hints = get_type_hints(source_type) |
| 536 | |
| 537 | # Build the Pydantic typed_dict_field for each msgspec field |
| 538 | fields = {} |
| 539 | for name, hint in type_hints.items(): |
| 540 | if name not in msgspec_fields: |
| 541 | # Skip ClassVar and other non-struct annotations. |
| 542 | continue |
| 543 | # Skip private fields — they are excluded from serialization |
| 544 | # and should not appear in the generated JSON/OpenAPI schema. |
| 545 | if name.startswith("_"): |
| 546 | continue |
| 547 | msgspec_field = msgspec_fields[name] |
| 548 | |
| 549 | # typed_dict_field using the handler to get the schema |
| 550 | field_schema = handler(hint) |
| 551 | |
| 552 | # Add default value to the schema. |
| 553 | # Mark fields with defaults as not required so the generated |
| 554 | # JSON Schema stays consistent with ``omit_defaults=True`` |
| 555 | # serialization (fields at their default value may be absent). |
| 556 | if msgspec_field.default_factory is not msgspec.NODEFAULT: |
| 557 | wrapped_schema = core_schema.with_default_schema( |
| 558 | schema=field_schema, |
| 559 | default_factory=msgspec_field.default_factory, |
| 560 | ) |
| 561 | fields[name] = core_schema.typed_dict_field( |
| 562 | wrapped_schema, required=False |
| 563 | ) |
| 564 | elif msgspec_field.default is not msgspec.NODEFAULT: |
| 565 | wrapped_schema = core_schema.with_default_schema( |
| 566 | schema=field_schema, |
| 567 | default=msgspec_field.default, |
| 568 | ) |
| 569 | fields[name] = core_schema.typed_dict_field( |
| 570 | wrapped_schema, required=False |
| 571 | ) |
| 572 | else: |
| 573 | # No default, so Pydantic will treat it as required |
| 574 | fields[name] = core_schema.typed_dict_field(field_schema) |
| 575 | typed_dict_then_convert = core_schema.no_info_after_validator_function( |
| 576 | cls._validate_msgspec, |
| 577 | core_schema.typed_dict_schema(fields), |
| 578 | ) |
| 579 | |
| 580 | # Build a serializer that strips private / excluded fields. |
| 581 | serializer = core_schema.plain_serializer_function_ser_schema( |
| 582 | cls._serialize_msgspec, |
| 583 | info_arg=False, |
| 584 | ) |
| 585 | |
| 586 | # Accept either an already-constructed msgspec.Struct instance or a |
| 587 | # JSON/dict-like payload. |
| 588 | return core_schema.union_schema( |
| 589 | [ |
| 590 | core_schema.is_instance_schema(source_type), |
| 591 | typed_dict_then_convert, |
| 592 | ], |
| 593 | serialization=serializer, |
| 594 | ) |
| 595 | |
| 596 | @classmethod |
| 597 | def _validate_msgspec(cls, value: Any) -> Any: |
| 598 | """Validate and convert input to msgspec.Struct instance.""" |
| 599 | if isinstance(value, cls): |
| 600 | return value |
| 601 | if isinstance(value, dict): |
| 602 | return cls(**value) |
| 603 | return msgspec.convert(value, type=cls) |
| 604 | |
| 605 | @staticmethod |
| 606 | def _serialize_msgspec(value: Any) -> Any: |
| 607 | """Serialize a msgspec.Struct to a JSON-compatible dict, stripping |
| 608 | private (``_``-prefixed) and explicitly excluded fields. |
| 609 | |
| 610 | Uses ``msgspec.to_builtins`` which respects ``omit_defaults=True``, |
| 611 | so only fields that differ from their declared defaults are included. |
| 612 | """ |
| 613 | raw = msgspec.to_builtins(value) |
| 614 | if not isinstance(raw, dict): |
| 615 | return raw |
| 616 | |
| 617 | exclude: set[str] = cast( |
| 618 | set[str], |
| 619 | getattr(type(value), "__pydantic_msgspec_exclude__", set()), |
| 620 | ) |
| 621 | for key in list(raw): |
| 622 | if key.startswith("_") or key in exclude: |
| 623 | del raw[key] |
| 624 | |
| 625 | return raw |