fix(contrib/pydantic): reuse TypeAdapters across payloads (#1703)

* fix(contrib/pydantic): reuse TypeAdapters across payloads

PydanticJSONPlainPayloadConverter.from_payload constructed a fresh
pydantic TypeAdapter for every payload, rebuilding the core schema each
time for non-class hints such as discriminated unions and generic
collections.

Cache adapters per converter instance, keyed on hashable type hints;
unhashable hints keep constructing fresh adapters. The cache is
unbounded by default and configurable via the new keyword-only
max_cached_type_adapters option on PydanticJSONPlainPayloadConverter
and PydanticPayloadConverter (positive bounds with LRU eviction, zero
disables caching, negative raises ValueError).

Fixes #1695

* fix(contrib/pydantic): default type adapter cache bound to 1024

Bound the per-converter type adapter cache to 1024 entries by default
with LRU eviction, capping worst-case memory even with runtime-generated
hints while never evicting for typical static hint sets. None remains
available for an unbounded cache and zero still disables caching.

* test(contrib/pydantic): cover re-imported class cache isolation

The workflow sandbox re-imports user modules, producing distinct class
objects with identical names. Verify each gets its own cache slot and
validates to its own world's class, even when one converter is shared.

* fix(contrib/pydantic): avoid double hash on cached decode path

Address review: try the cache directly instead of pre-hashing every
hint. On TypeError, hash the hint once only to distinguish an
unhashable hint (bypass the cache with a fresh adapter) from a
TypeError raised during adapter construction (re-raise), keeping
adapter errors unsuppressed.

Author: nathan-gage

CHANGELOG.md

@@ -22,6 +22,17 @@ 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)). Up
+  to 1024 type adapters are cached per converter instance by default, with
+  least-recently-used eviction. To change the bound, pass
+  ``max_cached_type_adapters`` to ``PydanticPayloadConverter`` (or
+  ``PydanticJSONPlainPayloadConverter``) from a nullary subclass used as the
+  ``DataConverter.payload_converter_class``; ``None`` makes the cache
+  unbounded and zero disables caching.
+
 ### 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,28 @@ 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 = 1024,
+    ) -> 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, with least-recently-used eviction. Defaults to 1024.
+                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 +110,26 @@ 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:
+            type_adapter = self._type_adapter(_type_hint)
+        except TypeError:
+            # Distinguish an unhashable hint (bypass the cache) from a
+            # TypeError raised while constructing the adapter (re-raise).
+            try:
+                hash(_type_hint)
+            except TypeError:
+                type_adapter = TypeAdapter(_type_hint)
+            else:
+                raise
+        return type_adapter.validate_json(payload.data)
 
 
 class PydanticPayloadConverter(CompositePayloadConverter):
@@ -106,9 +139,36 @@ 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 = 1024,
+    ) -> None:
+        """Initialize object.
+
+        Args:
+            to_json_options: Options for serializing values to JSON.
+            max_cached_type_adapters: Maximum number of type adapters to
+                cache, with least-recently-used eviction. Defaults to 1024.
+                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,157 @@
     clone_objects,
 )
 
+_MANY_TYPE_HINTS = tuple(
+    typing.cast(type, typing.cast(object, typing.Annotated[list[int], index]))
+    for index in range(1025)
+)
+_UNHASHABLE_TYPE_HINT = typing.cast(
+    type, typing.cast(object, typing.Annotated[list[int], []])
+)
+
+
+@pytest.mark.parametrize(
+    (
+        "converter_kwargs",
+        "type_hints",
+        "expected_type_adapter_constructions",
+    ),
+    [
+        # Default caches repeated hints
+        ({}, (list[int], list[int]), 1),
+        # Zero disables caching
+        ({"max_cached_type_adapters": 0}, (list[int], list[int]), 2),
+        # Unhashable hints bypass the cache
+        ({}, (_UNHASHABLE_TYPE_HINT, _UNHASHABLE_TYPE_HINT), 2),
+        # Default bound is 1024: 1025 distinct hints evict the first
+        ({}, _MANY_TYPE_HINTS + (_MANY_TYPE_HINTS[0],), 1026),
+        # None is unbounded: no eviction
+        (
+            {"max_cached_type_adapters": None},
+            _MANY_TYPE_HINTS + (_MANY_TYPE_HINTS[0],),
+            1025,
+        ),
+        # Explicit bound evicts least recently used
+        (
+            {"max_cached_type_adapters": 1},
+            (list[int], _MANY_TYPE_HINTS[0], list[int]),
+            3,
+        ),
+    ],
+)
+def test_type_adapter_reuse(
+    monkeypatch: pytest.MonkeyPatch,
+    converter_kwargs: dict[str, typing.Any],
+    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(**converter_kwargs)
+    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)
+
+
+def test_type_adapter_cache_distinguishes_reimported_classes():
+    # The workflow sandbox re-imports user modules, producing distinct class
+    # objects with identical names. Class hints hash by identity, so each
+    # world's class must get its own cache slot and validate to itself.
+    import types
+
+    source = "from pydantic import BaseModel\n\nclass Foo(BaseModel):\n    name: str\n"
+
+    def load_module() -> types.ModuleType:
+        module = types.ModuleType("test_reimported_models")
+        exec(compile(source, "test_reimported_models.py", "exec"), module.__dict__)
+        return module
+
+    foo_outside = load_module().Foo
+    foo_sandbox = load_module().Foo
+    assert foo_outside is not foo_sandbox
+    assert hash(foo_outside) != hash(foo_sandbox)
+
+    # Worst case: one converter shared across both worlds (the real sandbox
+    # creates a separate converter per workflow instance).
+    converter = PydanticJSONPlainPayloadConverter()
+    payload = converter.to_payload(foo_outside(name="x"))
+    assert payload is not None
+    decoded_outside = converter.from_payload(payload, foo_outside)
+    decoded_sandbox = converter.from_payload(payload, foo_sandbox)
+    assert type(decoded_outside) is foo_outside
+    assert type(decoded_sandbox) is foo_sandbox
+
+
+@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()