diff --git a/temporalio/activity.py b/temporalio/activity.py index 4e632701e..bec4d1d76 100644 --- a/temporalio/activity.py +++ b/temporalio/activity.py @@ -30,7 +30,7 @@ import temporalio.common import temporalio.converter -from .types import CallableType +from .types import AnyType, CallableType if TYPE_CHECKING: from temporalio.client import Client @@ -206,6 +206,7 @@ class _Context: runtime_metric_meter: temporalio.common.MetricMeter | None client: Client | None cancellation_details: _ActivityCancellationDetailsHolder + data_converter: temporalio.converter.DataConverter | None = None _logger_details: Mapping[str, Any] | None = None _payload_converter: temporalio.converter.PayloadConverter | None = None _metric_meter: temporalio.common.MetricMeter | None = None @@ -458,6 +459,54 @@ def payload_converter() -> temporalio.converter.PayloadConverter: return _Context.current().payload_converter +async def resolve_value_handle( + handle: temporalio.common.ValueHandle[AnyType], +) -> AnyType: + """Acquire the value a :py:class:`temporalio.common.ValueHandle` refers to. + + Uses this activity's data converter to run the handle's deferred inbound + pipeline (external-storage retrieval if offloaded, codec decode, then + deserialization) under the activity's serialization context. Call it from + activity code, where acquisition I/O is permitted; a workflow forwards + handles but does not acquire their values. + """ + context = _Context.current() + if context.data_converter is None: + raise RuntimeError( + "No data converter is available in this activity context; " + "cannot acquire a ValueHandle value." + ) + return await context.data_converter.resolve_value_handle(handle) + + +async def create_value_handle( + value: AnyType, + *, + metadata: Mapping[str, str] | None = None, +) -> temporalio.common.ValueHandle[AnyType]: + """Produce a :py:class:`~temporalio.common.ValueHandle` from a value, deferring the store. + + Converts the value immediately (the convert does no I/O) and returns a + *pending* handle. The stored form -- codec encode plus external-storage upload + if configured, with the ``metadata`` attached -- is produced only when the + activity's result is committed, using the activity's serialization context. If + the activity faults before returning the handle, nothing is uploaded and no + external-storage blob is orphaned. A consumer can read the metadata without + acquiring the value. A workflow forwards handles but does not create or acquire + their values. + + Async for consistency with :py:func:`resolve_value_handle` and the SDK's other + boundary operations, even though the convert itself does no I/O. + """ + context = _Context.current() + if context.data_converter is None: + raise RuntimeError( + "No data converter is available in this activity context; " + "cannot create a ValueHandle." + ) + return await context.data_converter.create_value_handle(value, metadata=metadata) + + def metric_meter() -> temporalio.common.MetricMeter: """Get the metric meter for the current activity. diff --git a/temporalio/bridge/worker.py b/temporalio/bridge/worker.py index 6554c508c..66e63bda9 100644 --- a/temporalio/bridge/worker.py +++ b/temporalio/bridge/worker.py @@ -303,31 +303,72 @@ async def visit_payloads(self, payloads: PayloadSequence) -> None: payloads.extend(new_payloads) +def _skip_deferred( + f: Callable[[Sequence[Payload]], Awaitable[list[Payload]]], + defer: Callable[[], bool], +) -> Callable[[Sequence[Payload]], Awaitable[list[Payload]]]: + """Wrap a transform so payloads at a deferred position pass through. + + ``defer`` is a content-neutral predicate that reads the visitor's current + command / run-argument context (set during traversal) and returns True when + the payload(s) being visited should keep their opaque form -- i.e. skip + external-storage retrieval and codec decode so they surface as forward-only + handles. The decision never inspects the payload proto itself. + """ + + async def wrapped(payloads: Sequence[Payload]) -> list[Payload]: + if defer(): + return list(payloads) + return await f(payloads) + + return wrapped + + async def decode_activation( activation: temporalio.bridge.proto.workflow_activation.WorkflowActivation, data_converter: temporalio.converter.DataConverter, decode_headers: bool, storage_concurrency_limit: int, + defer: Callable[[], bool] | None = None, + index_run_args: bool = False, ) -> temporalio.converter._extstore.StorageOperationMetrics: """Decode all payloads in the activation. + Args: + defer: content-neutral predicate (reads the visitor's current command / + run-argument context) returning True when the payload(s) being + visited should be left un-retrieved and un-codec-decoded so they + surface as forward-only ValueHandles. + index_run_args: visit run arguments one at a time so ``defer`` can decide + per argument position (only needed when a run arg is a ValueHandle). + Returns: Metrics from any external storage retrieval operations that occurred. """ + retrieve: Callable[[Sequence[Payload]], Awaitable[list[Payload]]] = ( + data_converter._external_retrieve_payload_sequence + ) + decode: Callable[[Sequence[Payload]], Awaitable[list[Payload]]] = ( + data_converter._decode_payload_sequence + ) + if defer is not None: + retrieve = _skip_deferred(retrieve, defer) + decode = _skip_deferred(decode, defer) + metrics = temporalio.converter._extstore.StorageOperationMetrics() with metrics.track(): await CommandAwarePayloadVisitor( skip_search_attributes=True, skip_headers=not decode_headers, concurrency_limit=storage_concurrency_limit, - ).visit( - _Visitor(data_converter._external_retrieve_payload_sequence), activation - ) + index_run_args=index_run_args, + ).visit(_Visitor(retrieve), activation) await CommandAwarePayloadVisitor( skip_search_attributes=True, skip_headers=not decode_headers, - ).visit(_Visitor(data_converter._decode_payload_sequence), activation) + index_run_args=index_run_args, + ).visit(_Visitor(decode), activation) return metrics diff --git a/temporalio/common.py b/temporalio/common.py index ad75b56b9..972ef6e0f 100644 --- a/temporalio/common.py +++ b/temporalio/common.py @@ -9,7 +9,7 @@ import warnings from abc import ABC, abstractmethod from collections.abc import Callable, Collection, Iterator, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime, timedelta from enum import IntEnum from typing import ( @@ -378,6 +378,90 @@ def __setstate__(self, state: object) -> None: ) +# User metadata attached to a value handle is stored under payload-metadata keys +# with this prefix. It is server-opaque and travels with the reference, so +# consumers can read it without acquiring the value. +_HANDLE_METADATA_PREFIX = "_temporal-handle-meta." + + +@dataclass(frozen=True) +class ValueHandle(Generic[temporalio.types.AnyType]): + """A lazy, immutable, pass-by-reference handle to a value. + + The lazy sibling of :py:class:`RawValue`. Annotate a workflow/activity/signal + parameter or return value as ``ValueHandle[T]`` to receive one of these + instead of the materialized value. Forward it onward without cost, and + acquire its value at a boundary (an activity or client) where fetching data + is permitted, via :py:func:`temporalio.activity.resolve_value_handle` or + :py:meth:`temporalio.converter.DataConverter.resolve_value_handle`. A handle + does not acquire its own value. + """ + + # The handle's payload. For a realized handle this is the stored reference + # the boundary converter acquires from. For a *pending* handle (see create at + # the activity boundary) it is the value already converted at create time, + # still awaiting the deferred codec encode and external-storage upload. + _payload: temporalio.api.common.v1.Payload | None = None + # Inner type ``T`` captured from the annotation, used as the decode hint + # when acquiring the value. May be None for a bare ``ValueHandle``. + _type: type | None = field(default=None, compare=False) + # Pending-handle state: True while the store (codec encode + external-storage + # upload) is still deferred to commit, plus the metadata to attach when it is + # stored. This state is transient and producer-side only; a handle only ever + # crosses the wire realized. + _pending: bool = field(default=False, compare=False) + _pending_metadata: Mapping[str, str] | None = field( + default=None, compare=False, repr=False + ) + + @property + def metadata(self) -> Mapping[str, str]: + """User metadata attached at creation, readable without acquiring the value. + + Lets a consumer route or filter on small descriptive data (size, format, + etc.) without downloading the underlying payload. Empty if none was + attached. + """ + if self._pending_metadata is not None: + return dict(self._pending_metadata) + if self._payload is None: + return {} + prefix = _HANDLE_METADATA_PREFIX + return { + key[len(prefix) :]: value.decode() + for key, value in self._payload.metadata.items() + if key.startswith(prefix) + } + + def __getstate__(self) -> object: + """Pickle support (workflow sandbox caching).""" + return {"payload": self._payload.SerializeToString(), "type": self._type} + + def __setstate__(self, state: object) -> None: + """Pickle support.""" + if not isinstance(state, dict): + raise TypeError(f"Expected dict state, got {type(state)}") + object.__setattr__( + self, + "_payload", + temporalio.api.common.v1.Payload.FromString(state["payload"]), + ) + object.__setattr__(self, "_type", state.get("type")) + + +class AsHandle: + """Marker for ``Annotated[T, AsHandle]``: consume ``T`` as a forward-only + :py:class:`ValueHandle` without ``T`` leaving the shared contract. + + ``Annotated[T, AsHandle]`` is transparent to type checkers (both a caller and + the callee still see ``T``), so a caller passes a plain ``T`` -- no coupling. + The SDK recovers the marker via ``get_type_hints(..., include_extras=True)`` + and delivers a handle instead of materializing. This is the *forward-only* + (type-erased) option: the callee's variable is still statically ``T``, so it + can forward the value but cannot call handle methods on it statically. + """ + + # We choose to make this a list instead of an sequence so we can catch if people # are not sending lists each time but maybe accidentally sending a string (which # is a sequence) diff --git a/temporalio/converter/_data_converter.py b/temporalio/converter/_data_converter.py index 823d1cc13..45707e784 100644 --- a/temporalio/converter/_data_converter.py +++ b/temporalio/converter/_data_converter.py @@ -29,10 +29,16 @@ from temporalio.converter._payload_converter import ( PayloadConverter, ) +from temporalio.converter._payload_handle import ( + ValueHandle, + _attach_metadata, + _is_payload_handle_hint, +) from temporalio.converter._serialization_context import ( SerializationContext, WithSerializationContext, ) +from temporalio.types import AnyType _REFERENCE_MESSAGE_TYPE = ExternalStorageReference.DESCRIPTOR.full_name.encode() @@ -108,6 +114,17 @@ async def encode( same number as values given, but must be at least one and cannot be more than was given. """ + # Realize any pending (deferred-store) value handles here, at commit. + # This is where a deferred create_value_handle actually stores. If a + # producing activity faults before returning its result, encode is never + # called for that result, so a pending handle is simply dropped: nothing + # is uploaded and no external-storage blob is orphaned. + values = [ + await self._realize_value_handle(value) + if isinstance(value, ValueHandle) and value._pending + else value + for value in values + ] payloads = self.payload_converter.to_payloads(values) payloads = await self._encode_payload_sequence(payloads) payloads = await self._external_store_payload_sequence(payloads) @@ -128,10 +145,117 @@ async def decode( Returns: Decoded and converted values. """ + # Positions annotated as ValueHandle defer acquisition: keep their + # opaque payload and skip eager external-storage retrieval + codec + # decode, producing a data-only handle. Acquisition is a boundary + # operation (resolve_value_handle), not something the handle does itself, so + # no converter is captured here. + if type_hints is not None and any( + _is_payload_handle_hint(h) for h in type_hints + ): + payloads = list(payloads) + transform_indexes = [ + i + for i in range(len(payloads)) + if i >= len(type_hints) or not _is_payload_handle_hint(type_hints[i]) + ] + to_transform = [payloads[i] for i in transform_indexes] + if to_transform: + transformed = await self._external_retrieve_payload_sequence( + to_transform + ) + transformed = await self._decode_payload_sequence(transformed) + for i, payload in zip(transform_indexes, transformed): + payloads[i] = payload + return self.payload_converter.from_payloads(payloads, type_hints) + payloads = await self._external_retrieve_payload_sequence(payloads) payloads = await self._decode_payload_sequence(payloads) return self.payload_converter.from_payloads(payloads, type_hints) + async def resolve_value_handle(self, handle: ValueHandle[AnyType]) -> AnyType: + """Acquire the value a :py:class:`ValueHandle` refers to. + + This is the boundary operation for handles: the handle is a plain value + carrying the opaque payload and its inner type, and this converter + supplies the machinery. It runs the deferred inbound pipeline + (external-storage retrieval if offloaded, then codec decode) under this + converter's serialization context, and deserializes into the handle's + captured type ``T``. + + Call it where acquisition I/O is permitted (an activity or client + boundary), never inside the workflow sandbox. In activity code, prefer + :py:func:`temporalio.activity.resolve_value_handle`, which uses the + activity's converter. + + Raises: + RuntimeError: if the handle carries no concrete type (a bare + ``ValueHandle`` annotation), since conversion needs a type. + """ + # A pending handle (created but not yet committed) holds the value as an + # already-converted payload; deserialize it directly, with no I/O. + if handle._pending: + return self.payload_converter.from_payloads( + [handle._payload], [handle._type] + )[0] + inner_type = handle._type + if inner_type is None: + raise RuntimeError( + "[TMPRL1106] ValueHandle has no type to acquire into. " + "Annotate the value as ValueHandle[T] with a concrete type T." + ) + payload = await self._transform_inbound_payload(handle._payload) + values = self.payload_converter.from_payloads([payload], [inner_type]) + return values[0] + + async def create_value_handle( + self, + value: Any, + *, + metadata: Mapping[str, str] | None = None, + ) -> ValueHandle[Any]: + """Produce a :py:class:`ValueHandle` from a value, deferring the store. + + The producer-side counterpart to :py:meth:`resolve_value_handle`. The + value is *converted* immediately (as a workflow method invocation does), + so it is snapshotted and any serialization error surfaces here at the call + site. But the I/O-bearing tail -- codec encode and external-storage upload + -- is deferred: it runs only when the handle is encoded at result/input + commit (see :py:meth:`encode` and :py:meth:`_realize_value_handle`). So if + a producing activity faults before returning the handle, nothing is + uploaded and no external-storage blob is orphaned. + + This is async for consistency with :py:meth:`resolve_value_handle` and the + SDK's other boundary operations, and to leave room to perform I/O (such as + an eager store) without a breaking signature change; the convert itself + does none. Call it where the deferred upload will be permitted at commit + (an activity or client boundary), never inside the workflow sandbox. In + activity code, prefer :py:func:`temporalio.activity.create_value_handle`. + """ + [payload] = self.payload_converter.to_payloads([value]) + return ValueHandle( + _payload=payload, + _type=type(value), + _pending=True, + _pending_metadata=metadata, + ) + + async def _realize_value_handle( + self, handle: ValueHandle[Any] + ) -> ValueHandle[Any]: + """Store a pending handle's converted value now, returning a realized handle. + + The value was already converted at create time; this runs the deferred, + I/O-bearing tail of the outbound pipeline -- codec encode, then + external-storage offload if configured -- and attaches the pending + metadata to the resulting reference. Called from :py:meth:`encode` at commit. + """ + [payload] = await self._encode_payload_sequence([handle._payload]) + [payload] = await self._external_store_payload_sequence([payload]) + if handle._pending_metadata: + _attach_metadata(payload, handle._pending_metadata) + return ValueHandle(payload, handle._type) + async def encode_wrapper( self, values: Sequence[Any] ) -> temporalio.api.common.v1.Payloads: diff --git a/temporalio/converter/_payload_converter.py b/temporalio/converter/_payload_converter.py index 8ee85ef72..7fc3a9845 100644 --- a/temporalio/converter/_payload_converter.py +++ b/temporalio/converter/_payload_converter.py @@ -45,6 +45,12 @@ if sys.version_info >= (3, 11): from enum import StrEnum # type: ignore[reportUnreachable] +from temporalio.converter._payload_handle import ( + ValueHandle, + _create_handle, + _is_payload_handle_hint, + _payload_handle_inner_type, +) from temporalio.converter._serialization_context import ( SerializationContext, WithSerializationContext, @@ -261,6 +267,11 @@ def to_payloads( # RawValue should just pass through if isinstance(value, temporalio.common.RawValue): payload = value.payload + # A ValueHandle re-emits its opaque payload unchanged, so + # forwarding it (e.g. workflow -> activity) neither downloads nor + # re-stores the underlying data. + elif isinstance(value, ValueHandle): + payload = value._payload else: for converter in self.converters.values(): payload = converter.to_payload(value) @@ -286,13 +297,22 @@ def from_payloads( KeyError: Unknown payload encoding RuntimeError: Error during decode """ - values = [] + values: list[Any] = [] type_hints = type_hints or [] for index, (payload, type_hint) in enumerate(zip_longest(payloads, type_hints)): # Raw value should just wrap if type_hint == temporalio.common.RawValue: values.append(temporalio.common.RawValue(payload)) continue + # A ValueHandle[T] hint defers acquisition: wrap the opaque payload + # regardless of its encoding rather than materializing it now. The + # handle binds to the current boundary converter (via contextvar) if + # one is set, else it is forward-only. + if type_hint is not None and _is_payload_handle_hint(type_hint): + values.append( + _create_handle(payload, _payload_handle_inner_type(type_hint)) + ) + continue encoding = payload.metadata.get("encoding", b"") converter = self.converters.get(encoding) if converter is None: diff --git a/temporalio/converter/_payload_handle.py b/temporalio/converter/_payload_handle.py new file mode 100644 index 000000000..17b7cd458 --- /dev/null +++ b/temporalio/converter/_payload_handle.py @@ -0,0 +1,66 @@ +"""Converter-internal helpers for value handles. + +The user-facing :py:class:`temporalio.common.ValueHandle` type and the +``AsHandle`` marker live in :py:mod:`temporalio.common`, next to +:py:class:`temporalio.common.RawValue`. This module holds only the +converter-internal helpers that recognize handle type hints and build handles +during payload conversion. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, Optional, get_args, get_origin + +import temporalio.api.common.v1 +from temporalio.common import ( + _HANDLE_METADATA_PREFIX, + AsHandle, + ValueHandle, +) + + +def _is_payload_handle_hint(hint: Any) -> bool: + """Return True for ``ValueHandle``, ``ValueHandle[T]``, or ``Annotated[T, AsHandle]``.""" + return ( + hint is ValueHandle + or get_origin(hint) is ValueHandle + or AsHandle in getattr(hint, "__metadata__", ()) + ) + + +def _payload_handle_inner_type(hint: Any) -> Optional[type]: + """Return ``T`` from ``ValueHandle[T]`` or ``Annotated[T, AsHandle]``, else None.""" + # Annotated[T, AsHandle]: the base type T is the inner (materialized) type. + if AsHandle in getattr(hint, "__metadata__", ()): + return hint.__origin__ + args = get_args(hint) + return args[0] if args else None + + +def _payload_handle_hint(inner_type: Optional[type]) -> Any: + """Build a ``ValueHandle[inner_type]`` hint (bare if ``inner_type`` is None). + + Used to upgrade a call's result type so an unchanged activity/child result + is consumed as a handle: the declared return type becomes the handle's ``T``. + """ + return ( + ValueHandle[inner_type] # type: ignore[valid-type] + if inner_type is not None + else ValueHandle + ) + + +def _create_handle( + payload: temporalio.api.common.v1.Payload, inner_type: Optional[type] +) -> ValueHandle[Any]: + """Build a data-only handle (no captured converter).""" + return ValueHandle(payload, inner_type) + + +def _attach_metadata( + payload: temporalio.api.common.v1.Payload, metadata: Mapping[str, str] +) -> None: + """Attach user metadata to a payload as server-opaque, prefixed keys.""" + for key, value in metadata.items(): + payload.metadata[_HANDLE_METADATA_PREFIX + key] = value.encode() diff --git a/temporalio/worker/_activity.py b/temporalio/worker/_activity.py index 0304b3b75..89ae296b1 100644 --- a/temporalio/worker/_activity.py +++ b/temporalio/worker/_activity.py @@ -666,6 +666,7 @@ async def _execute_activity( runtime_metric_meter=None if sync_non_threaded else self._metric_meter, client=self._client if not running_activity.sync else None, cancellation_details=running_activity.cancellation_details, + data_converter=data_converter, ) ) temporalio.activity.logger.debug("Starting activity") diff --git a/temporalio/worker/_command_aware_visitor.py b/temporalio/worker/_command_aware_visitor.py index 500fc4db5..8b08b22a6 100644 --- a/temporalio/worker/_command_aware_visitor.py +++ b/temporalio/worker/_command_aware_visitor.py @@ -9,6 +9,7 @@ from temporalio.bridge._visitor import PayloadVisitor from temporalio.bridge._visitor_functions import VisitorFunctions from temporalio.bridge.proto.workflow_activation.workflow_activation_pb2 import ( + InitializeWorkflow, ResolveActivity, ResolveChildWorkflowExecution, ResolveChildWorkflowExecutionStart, @@ -39,6 +40,13 @@ class CommandInfo: contextvars.ContextVar("current_command_info", default=None) ) +# Set to the positional index of a workflow run argument while that argument is +# being visited. Lets deferral be decided per argument position (content-neutral) +# rather than by inspecting the payload. +current_run_arg_index: contextvars.ContextVar[int | None] = contextvars.ContextVar( + "current_run_arg_index", default=None +) + class CommandAwarePayloadVisitor(PayloadVisitor): """Payload visitor that sets command context during traversal. @@ -53,6 +61,7 @@ def __init__( skip_search_attributes: bool = False, skip_headers: bool = False, concurrency_limit: int = 1, + index_run_args: bool = False, ) -> None: """Creates a new command-aware payload visitor. @@ -61,12 +70,48 @@ def __init__( skip_headers: If True, headers are not visited. concurrency_limit: Maximum number of payload visits that may run concurrently during a single call to visit(). Defaults to 1. + index_run_args: If True, workflow run arguments are visited one at a + time with :py:data:`current_run_arg_index` set, so retrieval can + be deferred per argument position. Left False (batched) unless a + run argument is consumed as a ValueHandle. """ super().__init__( skip_search_attributes=skip_search_attributes, skip_headers=skip_headers, concurrency_limit=concurrency_limit, ) + self._index_run_args = index_run_args + + async def _visit_coresdk_workflow_activation_InitializeWorkflow( + self, fs: VisitorFunctions, o: InitializeWorkflow + ) -> None: + if not self._index_run_args: + await super()._visit_coresdk_workflow_activation_InitializeWorkflow(fs, o) + return + # Visit each run argument individually so per-argument-position deferral + # can be decided, then visit the remaining fields normally. Keep the + # "remaining fields" list in sync with the generated base method. + for index, argument in enumerate(o.arguments): + token = current_run_arg_index.set(index) + try: + await self._visit_temporal_api_common_v1_Payload(fs, argument) + finally: + current_run_arg_index.reset(token) + if not self.skip_headers: + for header in o.headers.values(): + await self._visit_temporal_api_common_v1_Payload(fs, header) + if o.HasField("continued_failure"): + await self._visit_temporal_api_failure_v1_Failure(fs, o.continued_failure) + if o.HasField("last_completion_result"): + await self._visit_temporal_api_common_v1_Payloads( + fs, o.last_completion_result + ) + if o.HasField("memo"): + await self._visit_temporal_api_common_v1_Memo(fs, o.memo) + if o.HasField("search_attributes"): + await self._visit_temporal_api_common_v1_SearchAttributes( + fs, o.search_attributes + ) # Workflow commands with payloads async def _visit_coresdk_workflow_commands_CompleteWorkflowExecution( diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index b9513068d..7e9252a71 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -16,6 +16,7 @@ from types import TracebackType import temporalio.api.common.v1 +import temporalio.api.enums.v1.command_type_pb2 import temporalio.bridge.proto.workflow_activation import temporalio.bridge.proto.workflow_completion import temporalio.bridge.runtime @@ -23,6 +24,7 @@ import temporalio.common import temporalio.converter import temporalio.converter._extstore +import temporalio.converter._payload_handle import temporalio.exceptions import temporalio.workflow from temporalio.bridge.worker import PollShutdownError @@ -288,6 +290,53 @@ def run_inline() -> None: loop.call_soon(run_inline) return await future + def _defer_retrieval( + self, + init_job: temporalio.bridge.proto.workflow_activation.InitializeWorkflow | None, + workflow: _RunningWorkflow | None, + ) -> tuple[Callable[[], bool] | None, bool]: + """Build a content-neutral deferral predicate for ``decode_activation``. + + The predicate reads the payload visitor's current context (run-argument + index, or command seq for a resolved activity/child result) and returns + True when that position's payload should stay a forward-only ValueHandle + -- so retrieval and codec decode are skipped. It decides purely from + API-declared types by position, never by inspecting the payload proto. + + Returns the predicate (or None if nothing is deferrable) and whether run + arguments must be visited per-index (only when a run arg is a handle). + """ + is_handle = temporalio.converter._payload_handle._is_payload_handle_hint + run_arg_indices: set[int] = set() + if init_job: + defn = self._workflows.get(init_job.workflow_type) + if defn and defn.arg_types: + run_arg_indices = { + i for i, t in enumerate(defn.arg_types) if is_handle(t) + } + + if not run_arg_indices and workflow is None: + return None, False + + instance = workflow.instance if workflow is not None else None + command_type = temporalio.api.enums.v1.command_type_pb2.CommandType + result_command_types = { + command_type.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK, + command_type.COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION, + } + + def defer() -> bool: + arg_index = _command_aware_visitor.current_run_arg_index.get() + if arg_index is not None: + return arg_index in run_arg_indices + if instance is not None: + info = _command_aware_visitor.current_command_info.get() + if info is not None and info.command_type in result_command_types: + return instance.is_result_deferred(info) + return False + + return defer, bool(run_arg_indices) + async def _handle_activation( self, act: temporalio.bridge.proto.workflow_activation.WorkflowActivation ) -> None: @@ -364,11 +413,14 @@ async def _handle_activation( workflow_context_dc=data_converter, workflow_context=workflow_context, ) + defer_predicate, index_run_args = self._defer_retrieval(init_job, workflow) download_metrics = await temporalio.bridge.worker.decode_activation( act, data_converter, decode_headers=self._encode_headers, storage_concurrency_limit=self._max_workflow_task_external_storage_concurrency, + defer=defer_predicate, + index_run_args=index_run_args, ) if not workflow: assert init_job diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 726ff85e0..f7b59616a 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -56,6 +56,7 @@ import temporalio.bridge.proto.workflow_completion import temporalio.common import temporalio.converter +import temporalio.converter._payload_handle import temporalio.exceptions import temporalio.nexus.system import temporalio.workflow @@ -82,6 +83,7 @@ logger = logging.getLogger(__name__) + # Set to true to log all cases where we're ignoring things during delete LOG_IGNORE_DURING_DELETE = False @@ -222,6 +224,19 @@ def get_external_store_context( """ raise NotImplementedError + def is_result_deferred( + self, + command_info: _command_aware_visitor.CommandInfo | None, + ) -> bool: + """Whether the resolved result for this command should stay a handle. + + Not abstract: defaults to False (eager retrieval, the prior behavior). + Overridden to return True when the pending activity/child-workflow + requested its result as a ``ValueHandle``, so the worker can skip + downloading an offloaded result the workflow only forwards. + """ + return False + @abstractmethod def get_info(self) -> temporalio.workflow.Info: """Return the workflow info for this instance.""" @@ -2266,6 +2281,33 @@ def _failure_converter_with_context( failure_converter = failure_converter.with_context(context) return failure_converter + def is_result_deferred( + self, + command_info: _command_aware_visitor.CommandInfo | None, + ) -> bool: + if command_info is None: + return False + command_type = temporalio.api.enums.v1.command_type_pb2.CommandType + seq = command_info.command_seq + handle: Any = None + if ( + command_info.command_type + == command_type.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK + ): + handle = self._pending_activities.get(seq) + elif ( + command_info.command_type + == command_type.COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION + ): + handle = self._pending_child_workflows.get(seq) + if handle is None: + return False + ret_type = handle._input.ret_type + return ( + ret_type is not None + and temporalio.converter._payload_handle._is_payload_handle_hint(ret_type) + ) + def get_serialization_context( self, command_info: _command_aware_visitor.CommandInfo | None, @@ -3179,6 +3221,12 @@ def __init__( ) ) + def as_value_handle(self) -> temporalio.workflow.ActivityHandle[Any]: + ph = temporalio.converter._payload_handle + if not ph._is_payload_handle_hint(self._input.ret_type): + self._input.ret_type = ph._payload_handle_hint(self._input.ret_type) + return self + def cancel(self, msg: Any | None = None) -> bool: # Allow the cancel to go through for the task even if we're deleting, # just don't do any commands @@ -3332,6 +3380,12 @@ def __init__( workflow_context ) + def as_value_handle(self) -> temporalio.workflow.ChildWorkflowHandle[Any, Any]: + ph = temporalio.converter._payload_handle + if not ph._is_payload_handle_hint(self._input.ret_type): + self._input.ret_type = ph._payload_handle_hint(self._input.ret_type) + return self + @property def id(self) -> str: return self._input.id diff --git a/temporalio/worker/workflow_sandbox/_in_sandbox.py b/temporalio/worker/workflow_sandbox/_in_sandbox.py index d18374899..4bd2b8230 100644 --- a/temporalio/worker/workflow_sandbox/_in_sandbox.py +++ b/temporalio/worker/workflow_sandbox/_in_sandbox.py @@ -96,3 +96,10 @@ def get_external_store_context( ) -> StorageDriverStoreContext: """Get store context for external storage.""" return self.instance.get_external_store_context(command_info) + + def is_result_deferred( + self, + command_info: _command_aware_visitor.CommandInfo | None, + ) -> bool: + """Whether the command's resolved result should stay a handle.""" + return self.instance.is_result_deferred(command_info) diff --git a/temporalio/worker/workflow_sandbox/_runner.py b/temporalio/worker/workflow_sandbox/_runner.py index 17f473d64..24d946ab8 100644 --- a/temporalio/worker/workflow_sandbox/_runner.py +++ b/temporalio/worker/workflow_sandbox/_runner.py @@ -229,3 +229,20 @@ def get_external_store_context( ) # type: ignore finally: self.importer.restriction_context.is_runtime = False + + def is_result_deferred( + self, + command_info: _command_aware_visitor.CommandInfo | None, + ) -> bool: + # Forward call to the sandboxed instance + self.importer.restriction_context.is_runtime = True + try: + self._run_code( + "with __temporal_importer.applied():\n" + " __temporal_deferred = __temporal_in_sandbox.is_result_deferred(__temporal_command_info)\n", + __temporal_importer=self.importer, + __temporal_command_info=command_info, + ) + return bool(self.globals_and_locals.pop("__temporal_deferred", False)) + finally: + self.importer.restriction_context.is_runtime = False diff --git a/temporalio/workflow/_activities.py b/temporalio/workflow/_activities.py index ef883c016..2c8192353 100644 --- a/temporalio/workflow/_activities.py +++ b/temporalio/workflow/_activities.py @@ -8,6 +8,7 @@ import temporalio.bridge.proto.workflow_commands import temporalio.common +import temporalio.converter from ..types import ( AnyType, @@ -65,7 +66,19 @@ class ActivityHandle(_AsyncioTask[ReturnType]): # type: ignore[type-var] This extends :py:class:`asyncio.Task` and supports all task features. """ - pass + def as_value_handle( + self, + ) -> ActivityHandle[temporalio.common.ValueHandle[ReturnType]]: + """Consume this activity's result as a lazy ValueHandle. + + The activity is unchanged and still returns its declared type; awaiting + this handle yields a :py:class:`temporalio.common.ValueHandle` of + that type instead of the materialized value. If the result was offloaded + to external storage it is not downloaded into the workflow -- forward the + handle to an activity (or materialize it there) to avoid paying for data + the workflow only routes. + """ + raise NotImplementedError class ActivityCancellationType(IntEnum): diff --git a/temporalio/workflow/_workflow_ops.py b/temporalio/workflow/_workflow_ops.py index f80ca1bdb..9cbffc272 100644 --- a/temporalio/workflow/_workflow_ops.py +++ b/temporalio/workflow/_workflow_ops.py @@ -7,6 +7,7 @@ import temporalio.bridge.proto.child_workflow import temporalio.common +import temporalio.converter from ..types import ( MethodAsyncNoParam, @@ -68,6 +69,19 @@ def first_execution_run_id(self) -> str | None: """Run ID for the workflow.""" raise NotImplementedError + def as_value_handle( + self, + ) -> ChildWorkflowHandle[SelfType, temporalio.common.ValueHandle[ReturnType]]: + """Consume this child workflow's result as a lazy ValueHandle. + + The child workflow is unchanged and still returns its declared type; + awaiting this handle yields a + :py:class:`temporalio.common.ValueHandle` of that type instead of + the materialized value. If the result was offloaded to external storage + it is not downloaded into the parent workflow. + """ + raise NotImplementedError + @overload async def signal( self, diff --git a/tests/test_payload_handle.py b/tests/test_payload_handle.py new file mode 100644 index 000000000..2e046ee32 --- /dev/null +++ b/tests/test_payload_handle.py @@ -0,0 +1,190 @@ +"""Unit tests for ValueHandle (Phase 1 prototype), server-free. + +These exercise the converter-level behavior: a ValueHandle[T] annotation +defers acquisition (external-storage retrieval, codec decode, deserialization) +until the value is acquired at a boundary via DataConverter.resolve_value_handle, +and forwarding a handle re-emits its opaque payload without downloading. The +proof is the driver's retrieve-call count. +""" + +from __future__ import annotations + +import pickle +from collections.abc import Sequence + +import pytest + +from temporalio.api.common.v1 import Payload +from temporalio.common import ValueHandle +from temporalio.converter import ( + DataConverter, + ExternalStorage, + PayloadCodec, +) +from tests.test_extstore import InMemoryTestDriver + +# A value large enough to be worth offloading; threshold=0 offloads everything. +_BIG = "x" * 1000 + + +def _storage_converter( + driver: InMemoryTestDriver, codec: PayloadCodec | None = None +) -> DataConverter: + return DataConverter( + payload_codec=codec, + external_storage=ExternalStorage(drivers=[driver], payload_size_threshold=0), + ) + + +async def test_toplevel_reference_becomes_handle() -> None: + driver = InMemoryTestDriver() + dc = _storage_converter(driver) + + payloads = await dc.encode([_BIG]) + assert driver._store_calls == 1 + + [handle] = await dc.decode(payloads, [ValueHandle[str]]) + assert isinstance(handle, ValueHandle) + # No download happened just by receiving the handle. + assert driver._retrieve_calls == 0 + + # The value is acquired at the boundary, through the converter. + assert await dc.resolve_value_handle(handle) == _BIG + assert driver._retrieve_calls == 1 + + +async def test_handle_is_data_only_and_forwards_without_download() -> None: + driver = InMemoryTestDriver() + dc = _storage_converter(driver) + [reference] = await dc.encode([_BIG]) + + # A handle is a data-only value: holding it downloads nothing, and it owns + # no acquire behavior (acquisition is a boundary operation, not a method on + # the handle). Forward-only-ness lives on the workflow surface, which simply + # does not expose resolve_value_handle, not in the handle's state. + [handle] = dc.payload_converter.from_payloads([reference], [ValueHandle[str]]) + assert isinstance(handle, ValueHandle) + assert not hasattr(handle, "materialize") + assert not hasattr(handle, "resolve_value_handle") + assert driver._retrieve_calls == 0 + + # Forwarding re-emits a byte-identical reference payload, still no download. + [out] = dc.payload_converter.to_payloads([handle]) + assert out.SerializeToString() == reference.SerializeToString() + assert driver._retrieve_calls == 0 + + # The value is acquired only through a boundary converter. + assert await dc.resolve_value_handle(handle) == _BIG + assert driver._retrieve_calls == 1 + + +async def test_non_handle_annotation_is_eager() -> None: + driver = InMemoryTestDriver() + dc = _storage_converter(driver) + payloads = await dc.encode([_BIG]) + + # Default behavior is unchanged: a real-type hint materializes eagerly. + [value] = await dc.decode(payloads, [str]) + assert value == _BIG + assert not isinstance(value, ValueHandle) + assert driver._retrieve_calls == 1 + + +async def test_tmprl1105_preserved_for_non_handle() -> None: + driver = InMemoryTestDriver() + [reference] = await _storage_converter(driver).encode([_BIG]) + + # A reference decoded as a real type without storage still raises, exactly + # as before this feature. + with pytest.raises(RuntimeError, match="TMPRL1105"): + await DataConverter().decode([reference], [str]) + + +class _MarkerCodec(PayloadCodec): + """Reversible codec that prefixes data and counts decode calls.""" + + def __init__(self) -> None: + self.decode_calls = 0 + + async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: + return [ + Payload(metadata=dict(p.metadata), data=b"C" + p.data) for p in payloads + ] + + async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: + self.decode_calls += 1 + out = [] + for p in payloads: + data = p.data[1:] if p.data.startswith(b"C") else p.data + out.append(Payload(metadata=dict(p.metadata), data=data)) + return out + + +async def test_codec_deferred_until_acquired() -> None: + driver = InMemoryTestDriver() + codec = _MarkerCodec() + dc = _storage_converter(driver, codec=codec) + payloads = await dc.encode([_BIG]) + + [handle] = await dc.decode(payloads, [ValueHandle[str]]) + # The reference is not codec-decoded when the handle is produced. + assert codec.decode_calls == 0 + + assert await dc.resolve_value_handle(handle) == _BIG + assert codec.decode_calls == 1 + + +async def test_pickled_handle_survives_and_forwards() -> None: + driver = InMemoryTestDriver() + dc = _storage_converter(driver) + payloads = await dc.encode([_BIG]) + [handle] = await dc.decode(payloads, [ValueHandle[str]]) + + restored = pickle.loads(pickle.dumps(handle)) + assert isinstance(restored, ValueHandle) + # The opaque payload survives, so a rehydrated handle can still be forwarded + # and acquired at a boundary. Compare with proto equality: re-parsing may + # reorder the metadata map, so serialized bytes are not a reliable check. + [out] = dc.payload_converter.to_payloads([restored]) + assert out == payloads[0] + assert await dc.resolve_value_handle(restored) == _BIG + + +async def test_create_value_handle_defers_store_until_commit() -> None: + driver = InMemoryTestDriver() + # A realistic threshold: the value offloads, its small reference does not. + dc = DataConverter( + external_storage=ExternalStorage(drivers=[driver], payload_size_threshold=1024) + ) + value = "x" * 4096 + + # Creating a handle stores nothing: the store is deferred to commit (encode). + # create_value_handle does no I/O at call time (the convert is synchronous). + handle = await dc.create_value_handle(value, metadata={"pages": "42"}) + assert isinstance(handle, ValueHandle) + assert driver._store_calls == 0 + # Metadata is known without any store. + assert handle.metadata == {"pages": "42"} + + # Encoding the handle, as at result/input commit, is where the store happens. + [payload] = await dc.encode([handle]) + assert driver._store_calls == 1 + + # The committed payload is a realized reference carrying the metadata: a + # consumer probes the metadata without downloading, then resolves to the value. + [realized] = await dc.decode([payload], [ValueHandle[str]]) + assert realized.metadata == {"pages": "42"} + assert driver._retrieve_calls == 0 + assert await dc.resolve_value_handle(realized) == value + assert driver._retrieve_calls == 1 + + +async def test_pending_handle_dropped_without_commit_is_never_stored() -> None: + driver = InMemoryTestDriver() + dc = _storage_converter(driver) + + # A producer that creates a handle but never commits it (e.g. an activity that + # faults before returning) uploads nothing, so no external-storage blob is + # orphaned. + _ = await dc.create_value_handle(_BIG, metadata={"pages": "42"}) + assert driver._store_calls == 0 diff --git a/tests/test_payload_handle_annotated.py b/tests/test_payload_handle_annotated.py new file mode 100644 index 000000000..7df58e3e1 --- /dev/null +++ b/tests/test_payload_handle_annotated.py @@ -0,0 +1,81 @@ +"""Demonstration of the ``Annotated[T, AsHandle]`` forward-only consumption option. + +This is the "type-erased" option from the consumption-decoupling design: the +shared contract type stays ``T``, and a *marker* on the annotation (not the type +itself) drives handle consumption. A caller passes a plain ``T`` because +``Annotated[T, AsHandle]`` is transparent to type checkers, so there is no +coupling; the SDK recovers the marker and delivers a forward-only +:py:class:`ValueHandle` instead of materializing. + +Scope note: this exercises the converter layer directly, which is where the +mechanism lives. It does NOT yet work end to end through a running workflow, +because the SDK resolves argument/return types with ``get_type_hints()`` WITHOUT +``include_extras`` (``temporalio/common.py`` around line 1398), which strips the +``Annotated`` metadata before the converter ever sees it. Turning that on (and +threading it through arg-type resolution) is the remaining wiring; this test +isolates and demonstrates the underlying mechanism. +""" + +from __future__ import annotations + +from typing import Annotated, Any, get_type_hints + +from temporalio.common import AsHandle, ValueHandle +from temporalio.converter import DataConverter +from temporalio.converter._payload_handle import ( + _is_payload_handle_hint, + _payload_handle_inner_type, +) + + +def test_annotated_is_transparent_but_marker_recoverable() -> None: + def handler(data: Annotated[str, AsHandle]) -> None: ... + + # A caller and a type checker see a plain `str` -- the contract is unchanged, + # so there is no coupling: the caller passes a `str`, not a ValueHandle. + assert get_type_hints(handler)["data"] is str + # The SDK can still recover the marker when it asks for extras. + assert ( + get_type_hints(handler, include_extras=True)["data"] == Annotated[str, AsHandle] + ) + + +def test_annotated_marker_is_recognized_as_a_handle_hint() -> None: + hint = Annotated[str, AsHandle] + assert _is_payload_handle_hint(hint) + # The inner (materialized) type is the base type of the annotation. + assert _payload_handle_inner_type(hint) is str + + +async def test_same_payload_consumed_as_value_or_handle_by_marker() -> None: + dc = DataConverter() + # An ordinary inline payload (no external storage): the mechanism is neutral + # to the payload's shape -- it does not require an offloaded reference. + [payload] = await dc.encode(["big-value"]) + + # Consumed as the contract type -> the materialized value. + [value] = await dc.decode([payload], [str]) + assert value == "big-value" + + # Consumed via the marker -> a ValueHandle, with no change to the `str` + # contract. Its value is acquired at the boundary through the converter. + handle_hint: Any = Annotated[str, AsHandle] + [handle] = await dc.decode([payload], [handle_hint]) + assert isinstance(handle, ValueHandle) + assert await dc.resolve_value_handle(handle) == "big-value" + + +async def test_annotated_handle_is_a_data_only_value() -> None: + dc = DataConverter() + [payload] = dc.payload_converter.to_payloads(["big-value"]) + + # Sync conversion, as inside the workflow sandbox, yields a data-only handle: + # it carries the payload but owns no acquire behavior. Forward-only-ness is a + # property of the workflow surface (which does not expose resolve_value_handle), + # not of the handle; acquisition is a boundary operation. + handle_hint: Any = Annotated[str, AsHandle] + [handle] = dc.payload_converter.from_payloads([payload], [handle_hint]) + assert isinstance(handle, ValueHandle) + assert not hasattr(handle, "materialize") + # Through a boundary converter the value is acquirable. + assert await dc.resolve_value_handle(handle) == "big-value" diff --git a/tests/worker/test_payload_handle.py b/tests/worker/test_payload_handle.py new file mode 100644 index 000000000..ef65df3e6 --- /dev/null +++ b/tests/worker/test_payload_handle.py @@ -0,0 +1,286 @@ +"""End-to-end tests for ValueHandle (Phase 1 prototype). + +Demonstrates the headline behavior: a workflow whose run argument is annotated +``ValueHandle[T]`` receives a forward-only handle instead of an eagerly +downloaded value, forwards it to an activity without downloading, and the +activity acquires it on demand. The proof is the driver's retrieve count: +0 for pass-through (and on replay), exactly 1 when an activity materializes. +""" + +from __future__ import annotations + +import dataclasses +import uuid +from datetime import timedelta + +import temporalio.converter +from temporalio import activity, workflow +from temporalio.client import Client +from temporalio.common import ValueHandle +from temporalio.converter import ExternalStorage +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Replayer +from tests.helpers import new_worker +from tests.test_extstore import InMemoryTestDriver + +# Larger than the offload threshold so the run argument is externalized. +_BIG = "x" * 2000 +_THRESHOLD = 1024 + + +@activity.defn +async def consume_handle(data: ValueHandle[str]) -> int: + # The activity needs the bytes, so it acquires them on demand at the boundary. + value = await activity.resolve_value_handle(data) + return len(value) + + +@activity.defn +async def ignore_handle(data: ValueHandle[str]) -> str: + # Never materializes; the handle is just passed through. + return "ignored" + + +@workflow.defn +class ForwardToConsumeWorkflow: + @workflow.run + async def run(self, data: ValueHandle[str]) -> int: + # Forward the handle to an activity without materializing it here. + return await workflow.execute_activity( + consume_handle, data, start_to_close_timeout=timedelta(seconds=30) + ) + + +@workflow.defn +class ForwardToIgnoreWorkflow: + @workflow.run + async def run(self, data: ValueHandle[str]) -> str: + return await workflow.execute_activity( + ignore_handle, data, start_to_close_timeout=timedelta(seconds=30) + ) + + +@activity.defn +async def produce_big() -> str: + # An ordinary activity returning a large value; it is unchanged / unaware of + # handles. The value is offloaded to external storage on completion. + return _BIG + + +@workflow.defn +class ResultAsHandleConsumeWorkflow: + @workflow.run + async def run(self) -> int: + # Upgrade an unchanged activity's result to a handle, then forward it to + # an activity that materializes it. + handle = await workflow.start_activity( + produce_big, start_to_close_timeout=timedelta(seconds=30) + ).as_value_handle() + return await workflow.execute_activity( + consume_handle, handle, start_to_close_timeout=timedelta(seconds=30) + ) + + +@workflow.defn +class ResultAsHandlePassThroughWorkflow: + @workflow.run + async def run(self) -> str: + handle = await workflow.start_activity( + produce_big, start_to_close_timeout=timedelta(seconds=30) + ).as_value_handle() + return await workflow.execute_activity( + ignore_handle, handle, start_to_close_timeout=timedelta(seconds=30) + ) + + +@workflow.defn +class ChildProducerWorkflow: + @workflow.run + async def run(self) -> str: + return _BIG + + +@workflow.defn +class ParentChildResultAsHandleWorkflow: + @workflow.run + async def run(self) -> str: + # Upgrade an unchanged child workflow's result to a handle and forward it + # without materializing it in this (parent) workflow. + child = await workflow.start_child_workflow(ChildProducerWorkflow.run) + handle = await child.as_value_handle() + return await workflow.execute_activity( + ignore_handle, handle, start_to_close_timeout=timedelta(seconds=30) + ) + + +@activity.defn +async def produce_handle_with_metadata() -> ValueHandle[str]: + # The activity creates a handle from its (large) result and attaches metadata + # that a consumer can probe without downloading the value. + return await activity.create_value_handle(_BIG, metadata={"length": str(len(_BIG))}) + + +@workflow.defn +class ProbeMetadataThenForwardWorkflow: + @workflow.run + async def run(self) -> int: + handle = await workflow.execute_activity( + produce_handle_with_metadata, + start_to_close_timeout=timedelta(seconds=30), + ) + # Probe metadata in the workflow to decide, without downloading the value. + assert handle.metadata["length"] == str(len(_BIG)) + # Forward the handle to an activity that materializes it. + return await workflow.execute_activity( + consume_handle, handle, start_to_close_timeout=timedelta(seconds=30) + ) + + +def _data_converter(driver: InMemoryTestDriver) -> temporalio.converter.DataConverter: + return dataclasses.replace( + temporalio.converter.default(), + external_storage=ExternalStorage( + drivers=[driver], payload_size_threshold=_THRESHOLD + ), + ) + + +async def _client(env: WorkflowEnvironment, driver: InMemoryTestDriver) -> Client: + return await Client.connect( + env.client.service_client.config.target_host, + namespace=env.client.namespace, + data_converter=_data_converter(driver), + ) + + +async def test_activity_materializes_handle_once(env: WorkflowEnvironment) -> None: + driver = InMemoryTestDriver() + client = await _client(env, driver) + async with new_worker( + client, ForwardToConsumeWorkflow, activities=[consume_handle] + ) as worker: + # The caller sends the real value; the workflow opts to receive it as a + # handle, so pass via the loosely-typed args form. + result = await client.execute_workflow( + ForwardToConsumeWorkflow.run, + args=[_BIG], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == len(_BIG) + # Downloaded exactly once, at the point the activity materialized it. + assert driver._retrieve_calls == 1 + + +async def test_workflow_pass_through_no_download(env: WorkflowEnvironment) -> None: + driver = InMemoryTestDriver() + client = await _client(env, driver) + async with new_worker( + client, ForwardToIgnoreWorkflow, activities=[ignore_handle] + ) as worker: + handle = await client.start_workflow( + ForwardToIgnoreWorkflow.run, + args=[_BIG], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == "ignored" + + # The workflow forwarded the handle without ever downloading it. + assert driver._retrieve_calls == 0 + + # Replaying the same history must also avoid any download. + history = await handle.fetch_history() + replay_result = await Replayer( + workflows=[ForwardToIgnoreWorkflow], + data_converter=_data_converter(driver), + ).replay_workflow(history, raise_on_replay_failure=True) + assert replay_result is not None + assert driver._retrieve_calls == 0 + + +async def test_activity_result_as_handle_materializes_once( + env: WorkflowEnvironment, +) -> None: + driver = InMemoryTestDriver() + client = await _client(env, driver) + async with new_worker( + client, + ResultAsHandleConsumeWorkflow, + activities=[produce_big, consume_handle], + ) as worker: + result = await client.execute_workflow( + ResultAsHandleConsumeWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == len(_BIG) + # The producing activity's result was offloaded, upgraded to a handle in the + # workflow (not downloaded there), and materialized exactly once downstream. + assert driver._retrieve_calls == 1 + + +async def test_activity_result_as_handle_pass_through_no_download( + env: WorkflowEnvironment, +) -> None: + driver = InMemoryTestDriver() + client = await _client(env, driver) + async with new_worker( + client, + ResultAsHandlePassThroughWorkflow, + activities=[produce_big, ignore_handle], + ) as worker: + result = await client.execute_workflow( + ResultAsHandlePassThroughWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == "ignored" + # The unchanged activity's offloaded result was never downloaded: the + # workflow received it as a handle and forwarded it without materializing. + assert driver._retrieve_calls == 0 + + +async def test_child_workflow_result_as_handle_pass_through( + env: WorkflowEnvironment, +) -> None: + driver = InMemoryTestDriver() + client = await _client(env, driver) + async with new_worker( + client, + ParentChildResultAsHandleWorkflow, + ChildProducerWorkflow, + activities=[ignore_handle], + ) as worker: + result = await client.execute_workflow( + ParentChildResultAsHandleWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == "ignored" + # The child workflow's offloaded result was upgraded to a handle in the + # parent and forwarded without ever being downloaded. + assert driver._retrieve_calls == 0 + + +async def test_activity_creates_handle_with_probeable_metadata( + env: WorkflowEnvironment, +) -> None: + driver = InMemoryTestDriver() + client = await _client(env, driver) + async with new_worker( + client, + ProbeMetadataThenForwardWorkflow, + activities=[produce_handle_with_metadata, consume_handle], + ) as worker: + result = await client.execute_workflow( + ProbeMetadataThenForwardWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == len(_BIG) + # Stored once (when the activity created the handle) and downloaded once (when + # the consuming activity materialized it). The workflow probed metadata and + # forwarded the handle without any download. + assert driver._store_calls == 1 + assert driver._retrieve_calls == 1