diff --git a/CHANGELOG.md b/CHANGELOG.md index 483c73d1d..12cfe91fe 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/temporalio/contrib/pydantic.py b/temporalio/contrib/pydantic.py index c5f2deb41..6a5704b53 100644 --- a/temporalio/contrib/pydantic.py +++ b/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,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 +134,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 diff --git a/tests/contrib/pydantic/test_pydantic.py b/tests/contrib/pydantic/test_pydantic.py index 69a723a56..7d70dfd19 100644 --- a/tests/contrib/pydantic/test_pydantic.py +++ b/tests/contrib/pydantic/test_pydantic.py @@ -2,6 +2,7 @@ import datetime import os import pathlib +import typing import uuid import pydantic @@ -9,7 +10,11 @@ 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()