fix(contrib/pydantic): reuse TypeAdapters across payloads

CHANGELOG.md

@@ -22,6 +22,15 @@ to include examples, links to docs, or any other relevant information.
 
 ### Changed
 
+- `temporalio.contrib.pydantic` converters now reuse Pydantic type adapters
+  for repeated type hints instead of rebuilding their schemas for every
+  payload, greatly speeding up decode of non-model hints such as discriminated
+  unions ([#1695](https://github.com/temporalio/sdk-python/issues/1695)). The
+  per-converter-instance cache is unbounded by default; to bound or disable
+  it, pass ``max_cached_type_adapters`` to ``PydanticPayloadConverter`` (or
+  ``PydanticJSONPlainPayloadConverter``) from a nullary subclass used as the
+  ``DataConverter.payload_converter_class``.
+
 ### Deprecated
 
 ### :boom: Breaking Changes

temporalio/contrib/pydantic.py

@@ -13,6 +13,7 @@
 Pydantic v1 is not supported.
 """
 
+import functools
 from dataclasses import dataclass
 from typing import Any
 
@@ -53,10 +54,26 @@ class PydanticJSONPlainPayloadConverter(EncodingPayloadConverter):
     See https://docs.pydantic.dev/latest/api/standard_library_types/
     """
 
-    def __init__(self, to_json_options: ToJsonOptions | None = None):
-        """Create a new payload converter."""
+    def __init__(
+        self,
+        to_json_options: ToJsonOptions | None = None,
+        *,
+        max_cached_type_adapters: int | None = None,
+    ) -> None:
+        """Create a new payload converter.
+
+        Args:
+            to_json_options: Options for serializing values to JSON.
+            max_cached_type_adapters: Maximum number of type adapters to cache.
+                If ``None``, the cache is unbounded. If zero, caching is disabled.
+        """
+        if max_cached_type_adapters is not None and max_cached_type_adapters < 0:
+            raise ValueError("max_cached_type_adapters cannot be negative")
         self._schema_serializer = SchemaSerializer(any_schema())
         self._to_json_options = to_json_options
+        self._type_adapter = functools.lru_cache(maxsize=max_cached_type_adapters)(
+            TypeAdapter
+        )
 
     @property
     def encoding(self) -> str:
@@ -91,12 +108,21 @@ def from_payload(
 
         Uses ``pydantic.TypeAdapter.validate_json`` to construct an
         instance of the type specified by ``type_hint`` from the JSON payload.
+        Type adapters are cached per hashable type hint; see
+        ``max_cached_type_adapters`` on the constructor.
 
         See
         https://docs.pydantic.dev/latest/api/type_adapter/#pydantic.type_adapter.TypeAdapter.validate_json.
         """
         _type_hint = type_hint if type_hint is not None else Any
-        return TypeAdapter(_type_hint).validate_json(payload.data)
+        type_adapter: TypeAdapter[Any]
+        try:
+            hash(_type_hint)
+        except TypeError:
+            type_adapter = TypeAdapter(_type_hint)
+        else:
+            type_adapter = self._type_adapter(_type_hint)
+        return type_adapter.validate_json(payload.data)
 
 
 class PydanticPayloadConverter(CompositePayloadConverter):
@@ -106,9 +132,34 @@ class PydanticPayloadConverter(CompositePayloadConverter):
     :py:class:`PydanticJSONPlainPayloadConverter`.
     """
 
-    def __init__(self, to_json_options: ToJsonOptions | None = None) -> None:
-        """Initialize object"""
-        json_payload_converter = PydanticJSONPlainPayloadConverter(to_json_options)
+    def __init__(
+        self,
+        to_json_options: ToJsonOptions | None = None,
+        *,
+        max_cached_type_adapters: int | None = None,
+    ) -> None:
+        """Initialize object.
+
+        Args:
+            to_json_options: Options for serializing values to JSON.
+            max_cached_type_adapters: Maximum number of type adapters to cache.
+                If ``None``, the cache is unbounded. If zero, caching is disabled.
+
+                To configure this through a :py:class:`DataConverter`, use a
+                nullary subclass as the payload converter class::
+
+                    class MyPayloadConverter(PydanticPayloadConverter):
+                        def __init__(self) -> None:
+                            super().__init__(max_cached_type_adapters=128)
+
+                    my_data_converter = DataConverter(
+                        payload_converter_class=MyPayloadConverter
+                    )
+        """
+        json_payload_converter = PydanticJSONPlainPayloadConverter(
+            to_json_options,
+            max_cached_type_adapters=max_cached_type_adapters,
+        )
         super().__init__(
             *(
                 c

tests/contrib/pydantic/test_pydantic.py

@@ -2,14 +2,19 @@
 import datetime
 import os
 import pathlib
+import typing
 import uuid
 
 import pydantic
 import pytest
 from pydantic import BaseModel
 
 from temporalio.client import Client
-from temporalio.contrib.pydantic import pydantic_data_converter
+from temporalio.contrib.pydantic import (
+    PydanticJSONPlainPayloadConverter,
+    PydanticPayloadConverter,
+    pydantic_data_converter,
+)
 from temporalio.worker import Worker
 from temporalio.worker.workflow_sandbox._restrictions import (
     RestrictionContext,
@@ -41,6 +46,115 @@
     clone_objects,
 )
 
+_MANY_TYPE_HINTS = tuple(
+    typing.cast(type, typing.cast(object, typing.Annotated[list[int], index]))
+    for index in range(129)
+)
+_UNHASHABLE_TYPE_HINT = typing.cast(
+    type, typing.cast(object, typing.Annotated[list[int], []])
+)
+
+
+@pytest.mark.parametrize(
+    (
+        "max_cached_type_adapters",
+        "type_hints",
+        "expected_type_adapter_constructions",
+    ),
+    [
+        (None, (list[int], list[int]), 1),
+        (0, (list[int], list[int]), 2),
+        (None, (_UNHASHABLE_TYPE_HINT, _UNHASHABLE_TYPE_HINT), 2),
+        (None, _MANY_TYPE_HINTS + (_MANY_TYPE_HINTS[0],), 129),
+        (128, _MANY_TYPE_HINTS + (_MANY_TYPE_HINTS[0],), 130),
+    ],
+)
+def test_type_adapter_reuse(
+    monkeypatch: pytest.MonkeyPatch,
+    max_cached_type_adapters: int | None,
+    type_hints: tuple[type, ...],
+    expected_type_adapter_constructions: int,
+):
+    actual_type_adapter = pydantic.TypeAdapter
+    type_adapter_constructions = 0
+
+    def counting_type_adapter(
+        type_hint: typing.Any,
+    ) -> pydantic.TypeAdapter[typing.Any]:
+        nonlocal type_adapter_constructions
+        type_adapter_constructions += 1
+        return actual_type_adapter(type_hint)
+
+    monkeypatch.setattr(
+        "temporalio.contrib.pydantic.TypeAdapter", counting_type_adapter
+    )
+    converter = PydanticJSONPlainPayloadConverter(
+        max_cached_type_adapters=max_cached_type_adapters
+    )
+    payload = converter.to_payload([1])
+    assert payload is not None
+    for type_hint in type_hints:
+        assert converter.from_payload(payload, type_hint) == [1]
+    assert type_adapter_constructions == expected_type_adapter_constructions
+
+
+@pytest.mark.parametrize(
+    ("max_cached_type_adapters", "expected_type_adapter_constructions"),
+    [(None, 1), (0, 2)],
+)
+def test_composite_converter_forwards_type_adapter_cache_size(
+    monkeypatch: pytest.MonkeyPatch,
+    max_cached_type_adapters: int | None,
+    expected_type_adapter_constructions: int,
+):
+    actual_type_adapter = pydantic.TypeAdapter
+    type_adapter_constructions = 0
+
+    def counting_type_adapter(
+        type_hint: typing.Any,
+    ) -> pydantic.TypeAdapter[typing.Any]:
+        nonlocal type_adapter_constructions
+        type_adapter_constructions += 1
+        return actual_type_adapter(type_hint)
+
+    monkeypatch.setattr(
+        "temporalio.contrib.pydantic.TypeAdapter", counting_type_adapter
+    )
+    converter = PydanticPayloadConverter(
+        max_cached_type_adapters=max_cached_type_adapters
+    )
+    payloads = converter.to_payloads([[1], [2]])
+    assert converter.from_payloads(payloads, [list[int], list[int]]) == [[1], [2]]
+    assert type_adapter_constructions == expected_type_adapter_constructions
+
+
+def test_type_adapter_reuse_across_threads_with_deferred_build():
+    import concurrent.futures
+
+    class DeferredModel(BaseModel):
+        model_config = pydantic.ConfigDict(defer_build=True)
+        value: int
+
+    converter = PydanticJSONPlainPayloadConverter()
+    payload = converter.to_payload(DeferredModel(value=1))
+    assert payload is not None
+
+    def decode() -> DeferredModel:
+        return converter.from_payload(payload, DeferredModel)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
+        results = list(executor.map(lambda _: decode(), range(64)))
+    assert all(result == DeferredModel(value=1) for result in results)
+
+
+@pytest.mark.parametrize(
+    "converter_type",
+    [PydanticJSONPlainPayloadConverter, PydanticPayloadConverter],
+)
+def test_type_adapter_cache_rejects_negative_size(converter_type: type):
+    with pytest.raises(ValueError, match="max_cached_type_adapters cannot be negative"):
+        converter_type(max_cached_type_adapters=-1)
+
 
 async def test_instantiation_outside_sandbox():
     make_list_of_pydantic_objects()