From 86d50a10c445c1c3b2f11dd90280ecbcd8989174 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 25 Oct 2026 15:05:00 -0700 Subject: [PATCH 1/4] Added Activity Stream handles and `BufferedPublisher`. An Activity writes as its attempt, so a retry withdraws the Records of the earlier attempts. An Activity's `BufferedPublisher` flushes at the end of the attempt. --- temporalio/common/__init__.py | 14 + temporalio/common/_streams/__init__.py | 10 + .../common/_streams/_activity_stream.py | 295 +++++++ temporalio/common/_streams/_buffered.py | 176 ++++ temporalio/common/_streams/_stream_client.py | 7 +- .../contrib/streams/test_activity_handles.py | 801 ++++++++++++++++++ tests/contrib/streams/test_pydantic.py | 71 ++ tests/contrib/streams/test_redis.py | 103 ++- tests/contrib/streams/test_stream_client.py | 14 + 9 files changed, 1487 insertions(+), 4 deletions(-) create mode 100644 temporalio/common/_streams/_activity_stream.py create mode 100644 temporalio/common/_streams/_buffered.py create mode 100644 tests/contrib/streams/test_activity_handles.py create mode 100644 tests/contrib/streams/test_pydantic.py diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index ba622ecb5..dd7853807 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -34,6 +34,12 @@ import temporalio.types if TYPE_CHECKING: + from temporalio.common._streams._activity_stream import ( + ActivityStream, + SyncActivityStream, + SyncSubscription, + ) + from temporalio.common._streams._buffered import BufferedPublisher from temporalio.common._streams._defn import StreamDefinition, defn from temporalio.common._streams._errors import ( StreamClosedError, @@ -70,7 +76,9 @@ __all__ = [ "ActivityIDConflictPolicy", "ActivityIDReusePolicy", + "ActivityStream", "AutoUpgradeVersioningOverride", + "BufferedPublisher", "ClientStream", "CloseStatus", "Closed", @@ -125,6 +133,8 @@ "StreamUnsupportedError", "Subscription", "Superseded", + "SyncActivityStream", + "SyncSubscription", "TypedSearchAttributes", "VersioningBehavior", "VersioningOverride", @@ -1591,6 +1601,8 @@ def wait_sync(self, timeout: float | None = None) -> None: # so this module imports the Stream names on first use. _STREAM_NAMES = frozenset( [ + "ActivityStream", + "BufferedPublisher", "ClientStream", "CloseStatus", "Closed", @@ -1616,6 +1628,8 @@ def wait_sync(self, timeout: float | None = None) -> None: "StreamUnsupportedError", "Subscription", "Superseded", + "SyncActivityStream", + "SyncSubscription", "WorkflowStream", "defn", ] diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index 4c828426f..ace0c5ba7 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -1,5 +1,11 @@ """The Streams API, which :py:mod:`temporalio.common` exports.""" +from temporalio.common._streams._activity_stream import ( + ActivityStream, + SyncActivityStream, + SyncSubscription, +) +from temporalio.common._streams._buffered import BufferedPublisher from temporalio.common._streams._defn import StreamDefinition, defn from temporalio.common._streams._errors import ( StreamClosedError, @@ -31,6 +37,8 @@ __all__ = [ "END", + "ActivityStream", + "BufferedPublisher", "ClientStream", "CloseStatus", "Closed", @@ -55,6 +63,8 @@ "StreamUnsupportedError", "Subscription", "Superseded", + "SyncActivityStream", + "SyncSubscription", "WorkflowStream", "defn", ] diff --git a/temporalio/common/_streams/_activity_stream.py b/temporalio/common/_streams/_activity_stream.py new file mode 100644 index 000000000..0013f3a50 --- /dev/null +++ b/temporalio/common/_streams/_activity_stream.py @@ -0,0 +1,295 @@ +"""Stream handles in Activity code.""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +from collections.abc import Coroutine, Generator +from datetime import timedelta +from typing import Any, Generic, TypeVar + +from temporalio.api.enums.v1 import StreamErrorReason +from temporalio.common._streams import _activity, _errors +from temporalio.common._streams._access import Access, read_spec, subscription +from temporalio.common._streams._handle import StreamHandle +from temporalio.common._streams._items import Closed, End, Record, Superseded +from temporalio.common._streams._reader import Subscription + +T = TypeVar("T") +_R = TypeVar("_R") + +# How long a blocked sync call waits at a time. +_WAIT_STEP = 0.1 + + +def _start( + loop: asyncio.AbstractEventLoop, + call: Coroutine[Any, Any, _R], + future: concurrent.futures.Future[_R], +) -> None: + # Runs in the Worker's loop. The Activity thread can cancel the future before this runs. + if future.cancelled(): + call.close() + return + task = loop.create_task(call) + + def task_done(task: asyncio.Task[_R]) -> None: + if task.cancelled(): + future.cancel() + elif future.set_running_or_notify_cancel(): + error = task.exception() + if error is None: + future.set_result(task.result()) + else: + future.set_exception(error) + + def future_done(future: concurrent.futures.Future[_R]) -> None: + if future.cancelled() and not loop.is_closed(): + loop.call_soon_threadsafe(task.cancel) + + task.add_done_callback(task_done) + future.add_done_callback(future_done) + + +def _bound(handle: StreamHandle[Any]) -> _activity.ActivityStreams: + streams = _activity.current.get() + if streams is None: + raise _errors.StreamConfigError( + "No Stream backend runs this Activity. Register one on the Worker's client " + "with stream_backends=." + ) + if handle.backend and handle.backend != streams.backend_name: + raise _errors.StreamConfigError( + f"The Stream {handle.id} is of the backend {handle.backend}. The Worker has only " + f"{streams.backend_name}." + ) + if not handle.backend and not streams.default: + raise _errors.no_default_backend( + "Worker", [streams.backend_name], "a handle that has its backend name" + ) + return streams + + +def _not_owner(handle: StreamHandle[Any]) -> _errors.StreamNotOwnerError: + return _errors.StreamNotOwnerError( + f"Can't close the Stream {handle.id} from an Activity. A Workflow's Stream closes " + "only from Workflow code, and a client's only from a client.", + reason=StreamErrorReason.STREAM_ERROR_REASON_NOT_OWNER, + ) + + +def _access(handle: StreamHandle[Any]) -> Access: + streams = _bound(handle) + stream_id = handle.id + return Access( + backend=streams.backend, + converter=streams.converter, + namespace=streams.namespace, + stream_id=stream_id, + writer=lambda: streams.writer(stream_id), + ) + + +class ActivityStream(StreamHandle[T]): + """A Stream in an async Activity. The Activity can write and read it. + + An Activity gets a Stream only as an argument, such as a field of its input. Each publish + writes before it returns. Use :py:class:`BufferedPublisher` to batch. + + .. warning:: + This API is experimental. + """ + + def _access(self) -> Access: + return _access(self) + + async def publish(self, *values: T) -> int: + """Write one Record for each value, in order, and return the last Record's offset. + + A list is one value. + + Raises: + StreamNotFoundError: Nobody created the Stream. + """ + return await self._access().publish(values) + + def subscribe( + self, + *, + offset: int | End | None = None, + last: int | None = None, + wait_new_records: bool = True, + create_wait: timedelta | None = None, + ) -> Subscription[T]: + """Read the Stream's items in order. See :py:meth:`ClientStream.subscribe`.""" + read = read_spec( + self._payload_type, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + access = self._access() + return subscription(lambda position: access.items(read, position), read) + + async def get_offset(self) -> int: + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. + + Offsets are sparse and above zero, so ``offset=0`` reads from the start. + """ + return await self._access().get_offset() + + async def close( + self, *, result: Any | None = None, failure: BaseException | None = None + ) -> None: + """Raise :py:class:`StreamNotOwnerError`, since an Activity can't close a Stream. + + A Workflow's Stream closes only from Workflow code, and a client's Stream only from a + client. The rule includes a Local Activity of the run chain that created the Stream. A + Local Activity runs before its Workflow Task commits. If the task then failed, the close + would stay, and the task's retry couldn't create the Stream again. + """ + del result, failure + raise _not_owner(self) + + +class SyncSubscription(Generic[T]): + """The items of one read of a Stream in a sync Activity, in order. Iterate it with ``for``. + + Each iteration blocks the Activity's thread. See :py:class:`Subscription`. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, stream: SyncActivityStream[Any], subscription: Subscription[T] + ) -> None: + """Create a read. Use a handle's ``subscribe`` instead.""" + self._subscription = subscription + self._steps = stream._items(subscription) + + @property + def next_offset(self) -> int | None: + """The offset to resume from. See :py:attr:`Subscription.next_offset`.""" + return self._subscription.next_offset + + def __iter__(self) -> SyncSubscription[T]: + """The read is its own iterator.""" + return self + + def __next__(self) -> Record[T] | Superseded | Closed: + """The next item.""" + return next(self._steps) + + def close(self) -> None: + """End the read early. It makes no more calls.""" + self._steps.close() + + +class SyncActivityStream(StreamHandle[T]): + """A Stream in a sync Activity. Its calls block the Activity's thread. + + The calls run on the Worker's event loop, so the Activity thread only waits for them. + + .. warning:: + This API is experimental. + """ + + def _access(self) -> Access: + return _access(self) + + def _run(self, call: Coroutine[Any, Any, _R]) -> _R: + loop = _bound(self).loop + try: + on_loop = asyncio.get_running_loop() is loop + except RuntimeError: + on_loop = False + if on_loop: + call.close() + raise RuntimeError( + "A SyncActivityStream blocks, so use ActivityStream in an async Activity" + ) + # The future exists before the call is scheduled, so a cancel raised in this thread at + # any point after the schedule reaches the except below, which cancels the call. + future: concurrent.futures.Future[_R] = concurrent.futures.Future() + try: + try: + loop.call_soon_threadsafe(_start, loop, call, future) + except RuntimeError: + # The Worker's loop is closed, such as at shutdown, so the call never starts. + call.close() + raise + while True: + # Short waits let the Worker raise an exception in this thread, such as a cancel. + try: + return future.result(timeout=_WAIT_STEP) + except concurrent.futures.TimeoutError: + continue + except BaseException: + # Cancel the call too, so a long poll doesn't keep running in the Worker's loop after + # the Activity thread is canceled, such as from a heartbeat. + future.cancel() + raise + + def publish(self, *values: T) -> int: + """Write one Record for each value, in order, and return the last Record's offset. + + A list is one value. + + Raises: + StreamNotFoundError: Nobody created the Stream. + """ + return self._run(self._access().publish(values)) + + def subscribe( + self, + *, + offset: int | End | None = None, + last: int | None = None, + wait_new_records: bool = True, + create_wait: timedelta | None = None, + ) -> SyncSubscription[T]: + """Read the Stream's items in order. See :py:meth:`ClientStream.subscribe`.""" + read = read_spec( + self._payload_type, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + access = self._access() + return SyncSubscription( + self, subscription(lambda position: access.items(read, position), read) + ) + + def _items( + self, items: Subscription[T] + ) -> Generator[Record[T] | Superseded | Closed, None, None]: + try: + while True: + try: + item = self._run(_next(items)) + except StopAsyncIteration: + return + yield item + finally: + self._run(items.aclose()) + + def get_offset(self) -> int: + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. + + Offsets are sparse and above zero, so ``offset=0`` reads from the start. + """ + return self._run(self._access().get_offset()) + + def close( + self, *, result: Any | None = None, failure: BaseException | None = None + ) -> None: + """Raise :py:class:`StreamNotOwnerError`. See :py:meth:`ActivityStream.close`.""" + del result, failure + raise _not_owner(self) + + +async def _next(items: Subscription[T]) -> Record[T] | Superseded | Closed: + return await items.__anext__() diff --git a/temporalio/common/_streams/_buffered.py b/temporalio/common/_streams/_buffered.py new file mode 100644 index 000000000..fdf65c637 --- /dev/null +++ b/temporalio/common/_streams/_buffered.py @@ -0,0 +1,176 @@ +"""A publisher that batches the Records of one Stream.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta +from types import TracebackType +from typing import Generic, TypeVar + +from temporalio.common._streams import _activity +from temporalio.common._streams._activity_stream import ActivityStream +from temporalio.common._streams._stream_client import ClientStream + +T = TypeVar("T") + +# Records in one batch of the crate's writer. A full batch is written at once. +_BATCH = 1000 + + +class BufferedPublisher(Generic[T]): + """Batches the publishes of one Stream. + + Use it in an Activity or a client that writes often. :py:meth:`publish` returns at once. + The publisher writes the buffered values as one batch at most ``max_delay`` after the first + of them, or at once when they fill a batch. One write is in flight at a time, and values + published during it wait for the next one. Use it in ``async with``, which writes the + remaining values at the end of the block. In an Activity, the end of the attempt also + writes the remaining values. + + ``await flush()`` is the backpressure point. A caller that publishes faster than the Store + accepts Records must await it from time to time, so the buffer doesn't grow without bound. + + After a batch fails, the publisher writes nothing more, so a reader sees no gap. Each later + call raises the error. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + target: ActivityStream[T] | ClientStream[T], + *, + max_delay: timedelta = timedelta(milliseconds=100), + ) -> None: + """Create a publisher. + + Args: + target: An ``ActivityStream`` or a ``ClientStream``. + max_delay: The longest time that a value waits in the buffer. + """ + self._stream = target + self._max_delay = max_delay + self._buffer: list[T] = [] + # When the oldest buffered value arrived, so the next batch is at most max_delay later. + self._first_at = 0.0 + self._timer: asyncio.TimerHandle | None = None + self._writing: asyncio.Task[None] | None = None + self._published = 0 + # Count of the values that a write took, written or not. + self._settled = 0 + self._error: BaseException | None = None + self._last_offset = 0 + self._unwritten = 0 + self._reported = False + streams = _activity.current.get() + if streams is not None: + streams.publishers.append(self) + + def publish(self, *values: T) -> None: + """Add one Record for each value to the buffer, and return at once. + + A list is one value. + + Raises: + StreamError: A batch failed to write, so the publisher takes no more values. + """ + self._raise_error() + if not values: + raise TypeError("publish takes one or more values") + loop = asyncio.get_running_loop() + if not self._buffer: + self._first_at = loop.time() + self._buffer.extend(values) + self._published += len(values) + self._schedule() + + async def flush(self) -> int: + """Write every value published before the call, and return the last written offset. + + Returns 0 if the publisher hasn't written anything. + """ + target = self._published + while self._settled < target: + if self._writing is None: + self._start() + assert self._writing is not None + await asyncio.shield(self._writing) + self._raise_error() + return self._last_offset + + def _raise_error(self) -> None: + if self._error is not None: + self._reported = True + raise self._error + + async def _flush_at_attempt_end(self) -> None: + # Fail the attempt for unwritten values, unless the Activity already got the error and + # continued. + reported = self._reported + try: + await self.flush() + except Exception: + if not reported: + raise + + def _schedule(self) -> None: + if self._writing is not None or not self._buffer: + return + if len(self._buffer) >= _BATCH: + self._start() + elif self._timer is None: + self._timer = asyncio.get_running_loop().call_at( + self._first_at + self._max_delay.total_seconds(), self._start + ) + + def _start(self) -> None: + if self._timer is not None: + self._timer.cancel() + self._timer = None + if self._writing is not None or not self._buffer: + return + batch, self._buffer = self._buffer, [] + self._writing = asyncio.get_running_loop().create_task(self._write(batch)) + + async def _write(self, batch: list[T]) -> None: + try: + if self._error is None: + self._last_offset = await self._send(batch) + else: + # A failed batch fails the later ones too, so a reader sees no gap. + self._unwritten += len(batch) + except Exception as err: + self._unwritten += len(batch) + self._error = err + finally: + self._settled += len(batch) + self._writing = None + # Schedule the values published during the write, for when their delay ends or they + # fill a batch. + self._schedule() + + async def _send(self, batch: list[T]) -> int: + return await self._stream._access().publish(batch) + + @property + def unwritten(self) -> int: + """How many published values a failed batch left unwritten.""" + return self._unwritten + + async def __aenter__(self) -> BufferedPublisher[T]: + """Use the publisher in ``async with``.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Write the buffered values. A write error does not hide an error of the block.""" + try: + await self.flush() + except Exception: + if exc is None: + raise diff --git a/temporalio/common/_streams/_stream_client.py b/temporalio/common/_streams/_stream_client.py index 210f9bf3e..eeb533e9b 100644 --- a/temporalio/common/_streams/_stream_client.py +++ b/temporalio/common/_streams/_stream_client.py @@ -90,9 +90,10 @@ class ClientStream(StreamHandle[T]): :py:meth:`temporalio.client.Client.stream`. A ``ClientStream`` that a query of a client returns, such as a Workflow's ``stream`` query, uses that client. - Each publish writes before it returns. All handles of one client and Stream write as one - Producer. :py:meth:`producer` returns a handle that writes as a named Producer. Use a handle - in ``async with`` to close the Stream at the end of the block. + Each publish writes before it returns. Use :py:class:`BufferedPublisher` to batch. All + handles of one client and Stream write as one Producer. :py:meth:`producer` returns a + handle that writes as a named Producer. Use a handle in ``async with`` to close the Stream + at the end of the block. .. warning:: This API is experimental. diff --git a/tests/contrib/streams/test_activity_handles.py b/tests/contrib/streams/test_activity_handles.py new file mode 100644 index 000000000..55fff5d36 --- /dev/null +++ b/tests/contrib/streams/test_activity_handles.py @@ -0,0 +1,801 @@ +import asyncio +import contextlib +import contextvars +import ctypes +import gc +import inspect +import sys +import threading +import time +import uuid +from collections.abc import Coroutine, Iterator, Sequence +from dataclasses import dataclass +from datetime import timedelta +from types import SimpleNamespace +from typing import Any, cast + +import pytest + +from temporalio import activity, common, workflow +from temporalio.api.enums.v1 import StreamErrorReason +from temporalio.client import Client, WorkflowHandle +from temporalio.common import RetryPolicy +from temporalio.common._streams import _activity +from temporalio.exceptions import CancelledError +from tests.contrib.streams._support import workflow_stream +from tests.helpers import new_worker + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +@common.defn(name="other") +class Other(common.StreamDefinition[str]): ... + + +@dataclass +class CallModelInput: + prompt: str + stream: common.ActivityStream[str] + + +@dataclass +class SyncInput: + prompt: str + stream: common.SyncActivityStream[str] + + +async def read( + handle: WorkflowHandle[Any, Any], + definition: type[common.StreamDefinition[Any]] = Tokens, +) -> list[object]: + stream = workflow_stream(handle, definition) + return [item async for item in stream.subscribe(wait_new_records=False)] + + +def records(items: Sequence[object]) -> list[common.Record[Any]]: + return [item for item in items if isinstance(item, common.Record)] + + +@activity.defn +async def call_model(input: CallModelInput) -> list[int]: + assert input.stream._payload_type is str + offsets = [await input.stream.publish(word) for word in input.prompt.split()] + async with common.BufferedPublisher( + input.stream, max_delay=timedelta(seconds=10) + ) as buffered: + buffered.publish("model-a") + buffered.publish("model-b") + offsets.append(await input.stream.get_offset()) + return offsets + + +@activity.defn +def call_model_sync(input: SyncInput) -> int: + last = input.stream.publish(*input.prompt.split()) + assert input.stream.get_offset() == last + 1 + return last + + +@activity.defn +async def read_back(stream: common.ActivityStream[str]) -> list[Any]: + return [ + item.data + async for item in stream.subscribe(wait_new_records=False) + if isinstance(item, common.Record) + ] + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, prompt: str) -> list[Any]: + stream = await workflow.create_stream(Tokens) + offsets = await workflow.execute_activity( + call_model, + CallModelInput(prompt, common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + ) + last = await workflow.execute_activity( + call_model_sync, + SyncInput("sync words", common.SyncActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + ) + read = await workflow.execute_activity( + read_back, + common.ActivityStream.from_handle(stream), + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.close() + return [offsets, last, read] + + +async def test_activity_publishes_and_returns_offsets(stream_client: Client): + async with new_worker( + stream_client, + Agent, + activities=[call_model, call_model_sync, read_back], + activity_executor=_executor(), + ) as worker: + handle = await stream_client.start_workflow( + Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + offsets, last, read_tokens = await handle.result() + items = await read(handle) + data = records(items) + written = ["hi", "there", "model-a", "model-b", "sync", "words"] + assert [record.data for record in data] == written + assert offsets[:2] == [data[0].offset, data[1].offset] + assert offsets[2] == data[3].offset + 1 + assert last == data[5].offset + assert read_tokens == written + assert isinstance(items[-1], common.Closed) + + +def _executor() -> Any: + import concurrent.futures + + return concurrent.futures.ThreadPoolExecutor(max_workers=2) + + +attempts: dict[str, int] = {} + + +@activity.defn +async def flaky(input: CallModelInput) -> None: + attempts[input.prompt] = attempts.get(input.prompt, 0) + 1 + if attempts[input.prompt] == 1: + await input.stream.publish("attempt-1") + raise RuntimeError("first attempt fails") + + +@workflow.defn +class Retrying: + @workflow.run + async def run(self, key: str) -> None: + stream = await workflow.create_stream(Tokens) + await stream.publish("start") + await workflow.execute_activity( + flaky, + CallModelInput(key, common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + + +async def test_retry_withdraws_a_received_stream(stream_client: Client): + key = str(uuid.uuid4()) + async with new_worker(stream_client, Retrying, activities=[flaky]) as worker: + handle = await stream_client.start_workflow( + Retrying.run, key, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + items = await read(handle) + superseded = [item for item in items if isinstance(item, common.Superseded)] + # The second attempt writes nothing, and its start still withdraws the first one's Record. + assert len(superseded) == 1 + withdrawn = [r.data for r in records(items) if superseded[0].withdraws(r)] + assert withdrawn == ["attempt-1"] + + +@dataclass +class FinishedStreams: + key: str + closed: common.ActivityStream[str] + expired: common.ActivityStream[str] + + +@activity.defn +async def reads_after_a_failed_attempt(input: FinishedStreams) -> list[Any]: + attempts[input.key] = attempts.get(input.key, 0) + 1 + if attempts[input.key] == 1: + raise RuntimeError("first attempt fails") + items = input.closed.subscribe(wait_new_records=False) + return [item.data async for item in items if isinstance(item, common.Record)] + + +@workflow.defn +class ReadsFinishedStreams: + @workflow.run + async def run(self, key: str, expired: common.WorkflowStream[str]) -> list[Any]: + stream = await workflow.create_stream(Tokens) + await stream.publish("turn") + await stream.close() + return await workflow.execute_activity( + reads_after_a_failed_attempt, + FinishedStreams( + key, + common.ActivityStream.from_handle(stream), + common.ActivityStream.from_handle(expired), + ), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + + +async def test_retry_goes_on_with_a_closed_or_expired_stream(stream_client: Client): + # On a Stream that accepts no write, the second attempt's open withdraws nothing, and the + # attempt only reads. + expired = await stream_client.create_stream( + Tokens, id=f"expired-{uuid.uuid4()}", retention=timedelta(milliseconds=200) + ) + await expired.publish("old") + await asyncio.sleep(0.5) + key = str(uuid.uuid4()) + async with new_worker( + stream_client, ReadsFinishedStreams, activities=[reads_after_a_failed_attempt] + ) as worker: + result = await stream_client.execute_workflow( + ReadsFinishedStreams.run, + args=[key, common.WorkflowStream.from_handle(expired)], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == ["turn"] and attempts[key] == 2 + + +# Set by a test after its readers showed the first attempt's Record. +shown: dict[str, asyncio.Event] = {} + + +@activity.defn +async def fails_after_it_is_shown(input: CallModelInput) -> None: + attempts[input.prompt] = attempts.get(input.prompt, 0) + 1 + tokens = input.stream + if attempts[input.prompt] == 1: + await tokens.publish("attempt-1") + await shown[input.prompt].wait() + raise RuntimeError("first attempt fails") + await tokens.publish("attempt-2") + + +@workflow.defn +class ShownThenWithdrawn: + @workflow.run + async def run(self, key: str) -> None: + stream = await workflow.create_stream(Tokens) + await workflow.execute_activity( + fails_after_it_is_shown, + CallModelInput(key, common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + await stream.close() + + +async def test_reader_withdraws_a_shown_record_after_a_reconnect( + stream_client: Client, +): + key = str(uuid.uuid4()) + shown[key] = asyncio.Event() + async with new_worker( + stream_client, ShownThenWithdrawn, activities=[fails_after_it_is_shown] + ) as worker: + handle = await stream_client.start_workflow( + ShownThenWithdrawn.run, + key, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + stream = workflow_stream(handle, Tokens) + live = stream.subscribe(create_wait=timedelta(seconds=20)) + reconnects = stream.subscribe(create_wait=timedelta(seconds=20)) + first = await anext(live) + before = await anext(reconnects) + # The reader closes its connection and keeps only the Records that it showed and its + # resume offset, as a UI does across a page reload. + resume_at = reconnects.next_offset + await reconnects.aclose() + shown[key].set() + live_items = [first, *[item async for item in live]] + resumed = [item async for item in stream.subscribe(offset=resume_at)] + assert isinstance(first, common.Record) and first.data == "attempt-1" + assert isinstance(before, common.Record) and before.offset == first.offset + [live_superseded] = [i for i in live_items if isinstance(i, common.Superseded)] + assert live_superseded.withdrawn_offsets == (first.offset,) + # A new read got no Record of the failed attempt, so the item lists no offset. The kept + # Record is still withdrawn. + [superseded] = [i for i in resumed if isinstance(i, common.Superseded)] + assert superseded.withdrawn_offsets == () + assert superseded.withdraws(before) + assert [r.data for r in records(resumed)] == ["attempt-2"] + assert isinstance(resumed[-1], common.Closed) + + +@activity.defn +async def wrong_backend(stream: common.ActivityStream[str]) -> str: + try: + await stream.publish("x") + except common.StreamConfigError: + return "config" + return "wrote" + + +@activity.defn +async def close_it(stream: common.ActivityStream[str]) -> str: + await stream.publish("last") + try: + await stream.close(result="closed by activity") + except common.StreamNotOwnerError as err: + return StreamErrorReason.Name(err.reason) + return "closed" + + +@workflow.defn +class Closing: + @workflow.run + async def run(self) -> list[str]: + stream = await workflow.create_stream(Tokens) + backend = await workflow.execute_activity( + wrong_backend, + common.ActivityStream[str](id=stream.id, backend="elsewhere"), + start_to_close_timeout=timedelta(seconds=30), + ) + owner = await workflow.execute_activity( + close_it, + common.ActivityStream.from_handle(stream), + start_to_close_timeout=timedelta(seconds=30), + ) + # A Local Activity runs before its Workflow Task commits, so it can't close it either. + local = await workflow.execute_local_activity( + close_it, + common.ActivityStream.from_handle(stream), + start_to_close_timeout=timedelta(seconds=30), + ) + # A handle with only the id names no Workflow, and no Activity can close a Stream. + by_id = await workflow.execute_local_activity( + close_it, + common.ActivityStream[str](id=stream.id), + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.close(result="closed by the run") + return [backend, owner, local, by_id] + + +async def test_activity_checks_the_backend_and_the_owner(stream_client: Client): + async with new_worker( + stream_client, Closing, activities=[wrong_backend, close_it] + ) as worker: + handle = await stream_client.start_workflow( + Closing.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + # An Activity writes to the Workflow's Stream but can't close it. + assert await handle.result() == [ + "config", + *["STREAM_ERROR_REASON_NOT_OWNER"] * 3, + ] + items = await read(handle) + assert [r.data for r in records(items)] == ["last", "last", "last"] + assert isinstance(items[-1], common.Closed) + assert items[-1].result == "closed by the run" + + +class _Batches(common.ActivityStream[int]): + # A Stream whose writes go to a list, so the batches show with no backend. + def __init__(self, fail: bool = False) -> None: + super().__init__(id="batches") + self.batches: list[tuple[int, ...]] = [] + self.fail = fail + self.in_flight = 0 + self.most_in_flight = 0 + + def _access(self) -> Any: + return SimpleNamespace(publish=self._write) + + async def _write(self, values: Sequence[int]) -> int: + if self.fail: + raise common.StreamStorageError("store down") + self.in_flight += 1 + self.most_in_flight = max(self.most_in_flight, self.in_flight) + await asyncio.sleep(0.01) + self.in_flight -= 1 + self.batches.append(tuple(values)) + return sum(len(batch) for batch in self.batches) + + +async def test_buffered_publisher_batches_by_delay(): + stream = _Batches() + publisher = common.BufferedPublisher(stream, max_delay=timedelta(milliseconds=20)) + publisher.publish(1) + publisher.publish(2, 3) + await asyncio.sleep(0.1) + assert stream.batches == [(1, 2, 3)] + publisher.publish(4) + assert await publisher.flush() == 4 + assert stream.batches == [(1, 2, 3), (4,)] + assert await publisher.flush() == 4 + + +async def test_buffered_publisher_writes_a_full_batch_with_one_write_in_flight(): + stream = _Batches() + publisher = common.BufferedPublisher(stream, max_delay=timedelta(minutes=10)) + publisher.publish(*range(1000)) + publisher.publish(*range(1000, 2500)) + await asyncio.sleep(0.2) + # A full batch is written without waiting for the delay. The values published during its + # write are in the next batch. + assert [len(batch) for batch in stream.batches] == [1000, 1500] + assert stream.most_in_flight == 1 + assert await publisher.flush() == 2500 + + +async def test_buffered_publisher_raises_a_failed_batch(): + stream = _Batches(fail=True) + with pytest.raises(common.StreamStorageError): + async with common.BufferedPublisher(stream) as publisher: + publisher.publish(1) + with pytest.raises(TypeError): + common.BufferedPublisher(stream).publish() + + +@activity.defn +async def local_step(stream: common.ActivityStream[str]) -> list[Any]: + # Each Workflow write is an awaited Local Activity, so the run's create and Records are in + # the Store before a later Local Activity starts. + seen = [ + item.data + async for item in stream.subscribe(wait_new_records=False) + if isinstance(item, common.Record) + ] + await stream.publish("local") + return seen + + +@workflow.defn +class LocalOrder: + @workflow.run + async def run(self) -> list[Any]: + stream = await workflow.create_stream(Tokens) + await stream.publish("before") + seen = await workflow.execute_local_activity( + local_step, + common.ActivityStream.from_handle(stream), + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.publish("after") + await stream.close() + return seen + + +async def test_workflow_writes_go_out_before_a_local_activity(stream_client: Client): + async with new_worker(stream_client, LocalOrder, activities=[local_step]) as worker: + handle = await stream_client.start_workflow( + LocalOrder.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert await handle.result() == ["before"] + items = await read(handle) + written = records(items) + assert [record.data for record in written] == ["before", "local", "after"] + producer = written[1].producer + assert producer is not None and producer.id.startswith("la:") + assert isinstance(items[-1], common.Closed) + + +async def test_buffered_publisher_stays_failed(): + stream = _Batches(fail=True) + publisher = common.BufferedPublisher(stream) + publisher.publish(1, 2) + with pytest.raises(common.StreamStorageError): + await publisher.flush() + stream.fail = False + # A later batch would leave a gap, so the publisher refuses it and keeps the error. + with pytest.raises(common.StreamStorageError): + publisher.publish(3) + with pytest.raises(common.StreamStorageError): + await publisher.flush() + assert stream.batches == [] and publisher.unwritten == 2 + + +@activity.defn +async def never_flushes(input: CallModelInput) -> None: + publisher = common.BufferedPublisher(input.stream, max_delay=timedelta(minutes=10)) + publisher.publish(*input.prompt.split()) + + +@dataclass +class ManyStreams: + key: str + targets: list[common.ActivityStream[str]] + maybe: common.ActivityStream[str] | None = None + + +@activity.defn +async def flaky_many(input: ManyStreams) -> None: + attempts[input.key] = attempts.get(input.key, 0) + 1 + assert input.maybe is not None + if attempts[input.key] == 1: + for stream in [*input.targets, input.maybe]: + await stream.publish("attempt-1") + raise RuntimeError("first attempt fails") + + +@activity.defn +def sync_read(stream: common.SyncActivityStream[str]) -> list[Any]: + return [ + item.data + for item in stream.subscribe(wait_new_records=False) + if isinstance(item, common.Record) + ] + + +@workflow.defn +class Leftovers: + @workflow.run + async def run(self, key: str) -> list[Any]: + stream = await workflow.create_stream(Tokens) + other = await workflow.create_stream(Other) + await workflow.execute_activity( + never_flushes, + CallModelInput("a b", common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + ) + await workflow.execute_activity( + flaky_many, + ManyStreams( + key, + [common.ActivityStream.from_handle(stream)], + common.ActivityStream.from_handle(other), + ), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + return await workflow.execute_activity( + sync_read, + common.SyncActivityStream.from_handle(stream), + start_to_close_timeout=timedelta(seconds=30), + ) + + +async def test_attempt_end_flushes_and_retry_finds_nested_streams( + stream_client: Client, +): + key = str(uuid.uuid4()) + async with new_worker( + stream_client, + Leftovers, + activities=[never_flushes, flaky_many, sync_read], + activity_executor=_executor(), + ) as worker: + handle = await stream_client.start_workflow( + Leftovers.run, key, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + # The publisher was never flushed, so the end of the attempt wrote its values. + assert await handle.result() == ["a", "b", "attempt-1"] + items = await read(handle) + other = await read(handle, Other) + for read_items in (items, other): + superseded = [i for i in read_items if isinstance(i, common.Superseded)] + withdrawn = [r.data for r in records(read_items) if superseded[0].withdraws(r)] + assert withdrawn == ["attempt-1"] + + +@contextlib.contextmanager +def _loop_in_thread() -> Iterator[asyncio.AbstractEventLoop]: + # Each test closes its loop. The finalizer of a loop left open can run in a later caller + # thread and swallow the cancel that a test raises there. + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + try: + yield loop + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(5) + loop.close() + + +def _long_poll( + started: threading.Event, stopped: threading.Event +) -> Coroutine[Any, Any, None]: + async def long_poll() -> None: + started.set() + try: + await asyncio.sleep(60) + finally: + stopped.set() + + return long_poll() + + +def _caller( + loop: asyncio.AbstractEventLoop, + call: Coroutine[Any, Any, None], + errors: list[BaseException], +) -> threading.Thread: + stream: common.SyncActivityStream[Any] = common.SyncActivityStream(id="s") + token = _activity.current.set( + cast(Any, SimpleNamespace(loop=loop, backend_name="", default=True)) + ) + + def run() -> None: + try: + stream._run(call) + except BaseException as err: + errors.append(err) + + caller = threading.Thread(target=contextvars.copy_context().run, args=(run,)) + _activity.current.reset(token) + return caller + + +def _waits_for_the_call(caller: threading.Thread) -> bool: + frame = sys._current_frames().get(caller.ident or 0) + while frame is not None: + if ( + frame.f_code.co_name == "result" + and "concurrent" in frame.f_code.co_filename + ): + return True + frame = frame.f_back + return False + + +def test_sync_call_stops_when_the_thread_is_cancelled(): + started, stopped = threading.Event(), threading.Event() + errors: list[BaseException] = [] + with _loop_in_thread() as loop: + caller = _caller(loop, _long_poll(started, stopped), errors) + # A garbage collection in the caller thread runs finalizers there, and a finalizer + # swallows an exception that is raised in its thread. + gc.collect() + gc.disable() + try: + caller.start() + assert started.wait(5) + # Raise only while the caller waits on the call's future, so the test checks _run and + # not where the scheduler happens to stop the caller thread. + deadline = time.monotonic() + 5 + while not _waits_for_the_call(caller) and time.monotonic() < deadline: + time.sleep(0.001) + assert _waits_for_the_call(caller) + # The Worker cancels a sync Activity by raising in its thread. + assert caller.ident is not None + ctypes.pythonapi.PyThreadState_SetAsyncExc( + ctypes.c_ulong(caller.ident), ctypes.py_object(CancelledError) + ) + caller.join(5) + finally: + gc.enable() + assert stopped.wait(5) + assert errors and isinstance(errors[0], CancelledError) + + +def test_sync_call_on_a_closed_loop_closes_its_coroutine(): + errors: list[BaseException] = [] + loop = asyncio.new_event_loop() + loop.close() + call = _long_poll(threading.Event(), threading.Event()) + caller = _caller(loop, call, errors) + caller.start() + caller.join(5) + assert errors and isinstance(errors[0], RuntimeError) + # A closed coroutine can't warn at shutdown that nobody awaited it. + assert inspect.getcoroutinestate(call) == inspect.CORO_CLOSED + + +def test_sync_call_stops_when_the_cancel_lands_right_after_the_schedule(): + started, stopped = threading.Event(), threading.Event() + errors: list[BaseException] = [] + with _loop_in_thread() as loop: + schedule = loop.call_soon_threadsafe + cancelled: list[bool] = [] + + def schedule_then_cancel(*args: Any) -> Any: + handle = schedule(*args) + # The Worker's cancel lands in the Activity thread as soon as the call is scheduled. + if threading.current_thread() is caller and not cancelled: + cancelled.append(True) + assert started.wait(5) + raise CancelledError() + return handle + + setattr(loop, "call_soon_threadsafe", schedule_then_cancel) + caller = _caller(loop, _long_poll(started, stopped), errors) + caller.start() + caller.join(5) + assert stopped.wait(5) + assert errors and isinstance(errors[0], CancelledError) + + +@activity.defn +async def plain(name: str) -> str: + return f"hello {name}" + + +@workflow.defn +class NeverTouchesStreams: + @workflow.run + async def run(self) -> str: + return await workflow.execute_activity( + plain, "ada", start_to_close_timeout=timedelta(seconds=10) + ) + + +@workflow.defn +class ReceivesAStream: + @workflow.run + async def run(self, stream: common.WorkflowStream[str]) -> str: + # The Activity gets no Stream, but the run received one. + return await workflow.execute_activity( + plain, stream.id, start_to_close_timeout=timedelta(seconds=10) + ) + + +@workflow.defn +class RetriesAfterContinueAsNew: + @workflow.run + async def run(self, key: str, stream: common.WorkflowStream[str] | None) -> None: + if stream is None: + created = await workflow.create_stream(Tokens) + workflow.continue_as_new(args=[key, created]) + await workflow.execute_activity( + flaky, + CallModelInput(key, common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + + +async def keyed(handle: WorkflowHandle[Any, Any], activity_type: str) -> list[bool]: + # Whether each schedule of the Activity has the Producer key. + history = await handle.fetch_history() + return [ + "__temporal_stream_key" in attributes.header.fields + for event in history.events + if ( + attributes := event.activity_task_scheduled_event_attributes + ).activity_type.name + == activity_type + ] + + +async def test_only_a_run_that_uses_streams_keys_its_activities(stream_client: Client): + # The key is in History, so a run with no Stream adds none. A run that created, bound or + # received a Stream adds it to each Activity, so an Activity that writes keeps its Producer id. + key = str(uuid.uuid4()) + async with new_worker( + stream_client, + NeverTouchesStreams, + ReceivesAStream, + Retrying, + RetriesAfterContinueAsNew, + activities=[plain, flaky], + ) as worker: + none = await stream_client.start_workflow( + NeverTouchesStreams.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await none.result() == "hello ada" + assert await keyed(none, "plain") == [False] + creates = await stream_client.start_workflow( + Retrying.run, key, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await creates.result() + assert await keyed(creates, "flaky") == [True] + receives = await stream_client.start_workflow( + ReceivesAStream.run, + common.WorkflowStream.from_handle(workflow_stream(creates, Tokens)), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await receives.result() + assert await keyed(receives, "plain") == [True] + # The next run gets the Stream as its argument. Its retried Activity withdraws the first + # attempt's Records, as in the run that created the Stream. + continued_key = str(uuid.uuid4()) + continued = await stream_client.start_workflow( + RetriesAfterContinueAsNew.run, + args=[continued_key, None], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await continued.result() + assert await keyed(continued, "flaky") == [True] + items = await read(continued) + superseded = [item for item in items if isinstance(item, common.Superseded)] + assert len(superseded) == 1 + assert [r.data for r in records(items) if superseded[0].withdraws(r)] == [ + "attempt-1" + ] diff --git a/tests/contrib/streams/test_pydantic.py b/tests/contrib/streams/test_pydantic.py new file mode 100644 index 000000000..d7fadc069 --- /dev/null +++ b/tests/contrib/streams/test_pydantic.py @@ -0,0 +1,71 @@ +import uuid +from datetime import timedelta + +import pydantic + +from temporalio import activity, common, workflow +from temporalio.client import Client +from temporalio.contrib.pydantic import pydantic_data_converter +from temporalio.contrib.stream.testing import MemoryStreamBackend +from tests.helpers import new_worker + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +class CallModelInput(pydantic.BaseModel): + prompt: str + stream: common.ActivityStream[str] + + +class TurnResult(pydantic.BaseModel): + text: str + stream: common.ClientStream[str] + + +@activity.defn +async def call_model(input: CallModelInput) -> int: + assert input.stream._payload_type is str + return await input.stream.publish(*input.prompt.split()) + + +@workflow.defn +class TurnWorkflow: + @workflow.run + async def run(self, prompt: str) -> TurnResult: + stream = await workflow.create_stream(Tokens) + await workflow.execute_activity( + call_model, + CallModelInput( + prompt=prompt, stream=common.ActivityStream.from_handle(stream) + ), + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.close(result=prompt.upper()) + return TurnResult( + text=prompt.upper(), stream=common.ClientStream.from_handle(stream) + ) + + +async def test_handles_in_models_round_trip_on_the_memory_store(client: Client): + config = client.config() + config["stream_backends"] = [MemoryStreamBackend(default=True)] + config["data_converter"] = pydantic_data_converter + client = Client(**config) + async with new_worker(client, TurnWorkflow, activities=[call_model]) as worker: + handle = await client.start_workflow( + TurnWorkflow.run, + "hi there", + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + result = await handle.result() + assert result.text == "HI THERE" + assert result.stream.workflow_id == handle.id + assert result.stream.first_execution_run_id == handle.first_execution_run_id + assert result.stream.backend == "memory" + stream = client.stream(Tokens, id=result.stream.id) + items = [item async for item in stream.subscribe()] + assert [i.data for i in items if isinstance(i, common.Record)] == ["hi", "there"] + assert isinstance(items[-1], common.Closed) and items[-1].result is not None diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index b1ec38ae2..02d249ad8 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -1,8 +1,9 @@ import uuid +from dataclasses import dataclass from datetime import timedelta from typing import Any -from temporalio import common, workflow +from temporalio import activity, common, workflow from temporalio.client import Client from temporalio.worker import ( UnsandboxedWorkflowRunner, @@ -15,6 +16,12 @@ class Tokens(common.StreamDefinition[str]): ... +@dataclass +class CallModelInput: + prompt: str + stream: common.ActivityStream[str] + + async def read_all(stream: common.ClientStream[Any]) -> list[Any]: return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] @@ -64,3 +71,97 @@ async def test_failed_workflow_task_is_withdrawn(stream_client: Client): withdrawn = [r for r in written if any(i.withdraws(r) for i in superseded)] assert withdrawn == [written[1]] assert isinstance(items[-1], common.Closed) + + +@activity.defn +async def local_tokens(input: CallModelInput) -> None: + await input.stream.publish(*input.prompt.split()) + + +@workflow.defn +class LocalThenFailTask: + @workflow.run + async def run(self, key: str) -> None: + stream = await workflow.create_stream(Tokens) + first = attempts.setdefault(key, 0) == 0 + await workflow.execute_local_activity( + local_tokens, + CallModelInput( + "0 1" if first else "2 3", common.ActivityStream.from_handle(stream) + ), + start_to_close_timeout=timedelta(seconds=30), + ) + if first: + attempts[key] = 1 + raise RuntimeError("fail this Workflow Task once") + await stream.close() + + +async def test_rerun_local_activity_withdraws_its_failed_task(stream_client: Client): + async with new_worker( + stream_client, + LocalThenFailTask, + activities=[local_tokens], + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: + handle = await stream_client.start_workflow( + LocalThenFailTask.run, + str(uuid.uuid4()), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + items = await read_all(workflow_stream(handle, Tokens)) + written = records(items) + superseded = [item for item in items if isinstance(item, common.Superseded)] + visible = [r for r in written if not any(i.withdraws(r) for i in superseded)] + assert [r.data for r in visible] == ["2", "3"] + # The rerun keeps the Producer id and uses the retry task's epoch. + first, second = written[0].producer, visible[0].producer + assert first is not None and second is not None and first.id == second.id + assert first.epoch is not None and second.epoch is not None + assert (first.epoch.task_attempt, first.epoch.task_started_event_id) < ( + second.epoch.task_attempt, + second.epoch.task_started_event_id, + ) + + +@workflow.defn +class LocalOnlyInFailedTask: + @workflow.run + async def run(self, key: str) -> None: + stream = await workflow.create_stream(Tokens) + if attempts.setdefault(key, 0) == 0: + attempts[key] = 1 + await workflow.execute_local_activity( + local_tokens, + CallModelInput("0 1", common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + ) + raise RuntimeError("fail this Workflow Task once") + await stream.publish("other") + await stream.close() + + +async def test_later_write_withdraws_a_skipped_local_activity(stream_client: Client): + # The retry task skips the Local Activity, so only its own write can withdraw "0 1". The + # write also uses the skipped call's key. + async with new_worker( + stream_client, + LocalOnlyInFailedTask, + activities=[local_tokens], + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: + handle = await stream_client.start_workflow( + LocalOnlyInFailedTask.run, + str(uuid.uuid4()), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + items = await read_all(workflow_stream(handle, Tokens)) + written = records(items) + superseded = [item for item in items if isinstance(item, common.Superseded)] + visible = [r for r in written if not any(s.withdraws(r) for s in superseded)] + assert [r.data for r in visible] == ["other"] + assert isinstance(items[-1], common.Closed) diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index 7d3125174..aaf818a22 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -308,6 +308,20 @@ async def read(definition: type[common.StreamDefinition[Any]]) -> list[object]: assert await read(TurnDicts) == [{"kind": "model", "text": ""}] +async def test_buffered_publisher_keeps_the_publish_order(stream_client: Client): + stream = await stream_client.create_stream(Turns, id=f"buffered-{uuid.uuid4()}") + async with common.BufferedPublisher(stream) as publisher: + publisher.publish(Turn("user")) + publisher.publish(Turn("a"), Turn("b")) + publisher.publish(Turn("model")) + written = [ + item.data + async for item in stream.subscribe(wait_new_records=False) + if isinstance(item, common.Record) + ] + assert written == [Turn("user"), Turn("a"), Turn("b"), Turn("model")] + + @common.defn(name="turns") class LooseTurns(common.StreamDefinition[str]): ... From 355a4edcec075fb7b72d60a3cd298f0731dd97ed Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sat, 31 Oct 2026 02:11:14 -0700 Subject: [PATCH 2/4] Named the handle parameter of `SyncSubscription` `handle`. --- temporalio/common/_streams/_activity_stream.py | 8 +++----- tests/contrib/streams/test_activity_handles.py | 2 +- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/temporalio/common/_streams/_activity_stream.py b/temporalio/common/_streams/_activity_stream.py index 0013f3a50..2515a43b5 100644 --- a/temporalio/common/_streams/_activity_stream.py +++ b/temporalio/common/_streams/_activity_stream.py @@ -8,7 +8,6 @@ from datetime import timedelta from typing import Any, Generic, TypeVar -from temporalio.api.enums.v1 import StreamErrorReason from temporalio.common._streams import _activity, _errors from temporalio.common._streams._access import Access, read_spec, subscription from temporalio.common._streams._handle import StreamHandle @@ -73,8 +72,7 @@ def _bound(handle: StreamHandle[Any]) -> _activity.ActivityStreams: def _not_owner(handle: StreamHandle[Any]) -> _errors.StreamNotOwnerError: return _errors.StreamNotOwnerError( f"Can't close the Stream {handle.id} from an Activity. A Workflow's Stream closes " - "only from Workflow code, and a client's only from a client.", - reason=StreamErrorReason.STREAM_ERROR_REASON_NOT_OWNER, + "only from Workflow code, and a client's only from a client." ) @@ -163,11 +161,11 @@ class SyncSubscription(Generic[T]): """ def __init__( - self, stream: SyncActivityStream[Any], subscription: Subscription[T] + self, handle: SyncActivityStream[Any], subscription: Subscription[T] ) -> None: """Create a read. Use a handle's ``subscribe`` instead.""" self._subscription = subscription - self._steps = stream._items(subscription) + self._steps = handle._items(subscription) @property def next_offset(self) -> int | None: diff --git a/tests/contrib/streams/test_activity_handles.py b/tests/contrib/streams/test_activity_handles.py index 55fff5d36..9e81e27de 100644 --- a/tests/contrib/streams/test_activity_handles.py +++ b/tests/contrib/streams/test_activity_handles.py @@ -318,7 +318,7 @@ async def close_it(stream: common.ActivityStream[str]) -> str: try: await stream.close(result="closed by activity") except common.StreamNotOwnerError as err: - return StreamErrorReason.Name(err.reason) + return StreamErrorReason.Name(StreamErrorReason.ValueType(err.reason)) return "closed" From dee7a4a5408631395be7675e90c53ae34315eeea Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 1 Nov 2026 02:04:45 -0700 Subject: [PATCH 3/4] Gave `StreamHandle` the Activity calls, and added `SyncStreamHandle`. In an Activity a handle passes each call to the Activity calls, which write as the attempt's Producer. A write error that no caller got fails the attempt. Only the creator's side closes a Stream, so a close and `async with` raise before any call. `SyncStreamHandle` wraps a handle for a sync Activity, and `BufferedPublisher` takes a handle in client and Activity code. --- temporalio/common/__init__.py | 9 +- temporalio/common/_streams/__init__.py | 6 +- temporalio/common/_streams/_activity.py | 13 +- .../common/_streams/_activity_stream.py | 177 +++++++++------ temporalio/common/_streams/_buffered.py | 23 +- temporalio/common/_streams/_handle.py | 24 +- .../contrib/streams/test_activity_handles.py | 214 ++++++++++++++---- tests/contrib/streams/test_pydantic.py | 14 +- tests/contrib/streams/test_redis.py | 8 +- 9 files changed, 337 insertions(+), 151 deletions(-) diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index cda35b1d1..5896cb26f 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -35,8 +35,7 @@ if TYPE_CHECKING: from temporalio.common._streams._activity_stream import ( - ActivityStream, - SyncActivityStream, + SyncStreamHandle, SyncSubscription, ) from temporalio.common._streams._buffered import BufferedPublisher @@ -74,7 +73,6 @@ __all__ = [ "ActivityIDConflictPolicy", "ActivityIDReusePolicy", - "ActivityStream", "AutoUpgradeVersioningOverride", "BufferedPublisher", "CloseStatus", @@ -130,7 +128,7 @@ "StreamUnsupportedError", "Subscription", "Superseded", - "SyncActivityStream", + "SyncStreamHandle", "SyncSubscription", "TypedSearchAttributes", "VersioningBehavior", @@ -1597,7 +1595,6 @@ def wait_sync(self, timeout: float | None = None) -> None: # so this module imports the Stream names on first use. _STREAM_NAMES = frozenset( [ - "ActivityStream", "BufferedPublisher", "CloseStatus", "Closed", @@ -1623,7 +1620,7 @@ def wait_sync(self, timeout: float | None = None) -> None: "StreamUnsupportedError", "Subscription", "Superseded", - "SyncActivityStream", + "SyncStreamHandle", "SyncSubscription", "defn", ] diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index 2c889e3e8..d0156304a 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -1,8 +1,7 @@ """The Streams API, which :py:mod:`temporalio.common` exports.""" from temporalio.common._streams._activity_stream import ( - ActivityStream, - SyncActivityStream, + SyncStreamHandle, SyncSubscription, ) from temporalio.common._streams._buffered import BufferedPublisher @@ -35,7 +34,6 @@ __all__ = [ "END", - "ActivityStream", "BufferedPublisher", "CloseStatus", "Closed", @@ -60,7 +58,7 @@ "StreamUnsupportedError", "Subscription", "Superseded", - "SyncActivityStream", + "SyncStreamHandle", "SyncSubscription", "defn", ] diff --git a/temporalio/common/_streams/_activity.py b/temporalio/common/_streams/_activity.py index d6e87b288..34d9c3491 100644 --- a/temporalio/common/_streams/_activity.py +++ b/temporalio/common/_streams/_activity.py @@ -50,8 +50,11 @@ def writer(self, stream_id: str) -> StreamWriter: return writer async def close(self) -> None: - """Flush the attempt's publishers, then close every writer of the attempt.""" - errors = [] + """Flush the attempt's publishers, then close every writer of the attempt. + + A write error that no caller got fails the attempt. + """ + errors: list[BaseException] = [] for publisher in self.publishers: try: await publisher._flush_at_attempt_end() @@ -62,6 +65,12 @@ async def close(self) -> None: await writer.close() except Exception as err: errors.append(err) + # As the end of a Workflow run does, the end of the attempt raises a write error that + # no caller got. A retry then writes the Records again, and the attempt doesn't end + # with a gap. + lost = writer.unretrieved_error() + if lost is not None: + errors.append(lost) if errors: raise errors[0] diff --git a/temporalio/common/_streams/_activity_stream.py b/temporalio/common/_streams/_activity_stream.py index 2515a43b5..9e367f801 100644 --- a/temporalio/common/_streams/_activity_stream.py +++ b/temporalio/common/_streams/_activity_stream.py @@ -1,16 +1,16 @@ -"""Stream handles in Activity code.""" +"""Streams in Activity code.""" from __future__ import annotations import asyncio import concurrent.futures -from collections.abc import Coroutine, Generator +from collections.abc import Coroutine, Generator, Mapping from datetime import timedelta from typing import Any, Generic, TypeVar -from temporalio.common._streams import _activity, _errors +from temporalio.common._streams import _activity, _backend, _errors from temporalio.common._streams._access import Access, read_spec, subscription -from temporalio.common._streams._handle import StreamHandle +from temporalio.common._streams._handle import ContextCalls, StreamHandle from temporalio.common._streams._items import Closed, End, Record, Superseded from temporalio.common._streams._reader import Subscription @@ -70,6 +70,8 @@ def _bound(handle: StreamHandle[Any]) -> _activity.ActivityStreams: def _not_owner(handle: StreamHandle[Any]) -> _errors.StreamNotOwnerError: + # Only the creator's side closes a Stream: the run chain a Workflow's Stream, and a client a + # client's Stream. Code in an Activity only writes, so the refusal comes before any call. return _errors.StreamNotOwnerError( f"Can't close the Stream {handle.id} from an Activity. A Workflow's Stream closes " "only from Workflow code, and a client's only from a client." @@ -88,67 +90,93 @@ def _access(handle: StreamHandle[Any]) -> Access: ) -class ActivityStream(StreamHandle[T]): - """A Stream in an async Activity. The Activity can write and read it. +def _check_loop() -> None: + try: + asyncio.get_running_loop() + except RuntimeError: + raise _errors.StreamUnsupportedError( + "A sync Activity can't wait for a Stream call. Wrap the handle in a " + "SyncStreamHandle." + ) from None - An Activity gets a Stream only as an argument, such as a field of its input. Each publish - writes before it returns. Use :py:class:`BufferedPublisher` to batch. - .. warning:: - This API is experimental. - """ +class ActivityCalls(ContextCalls): + """The Stream calls in an Activity. The Activity can write and read a Stream. - def _access(self) -> Access: - return _access(self) - - async def publish(self, *values: T) -> int: - """Write one Record for each value, in order, and return the last Record's offset. - - A list is one value. + Each publish starts its write at the call, as the Producer of the attempt. An Activity + can't close a Stream. A Workflow's Stream closes only from Workflow code, and a client's + Stream only from a client. The rule includes a Local Activity of the run chain that created + the Stream. A Local Activity runs before its Workflow Task commits. If the task then failed, + the close would stay, and the task's retry couldn't create the Stream again. + """ - Raises: - StreamNotFoundError: Nobody created the Stream. - """ - return await self._access().publish(values) + def publish( + self, + handle: StreamHandle[Any], + values: tuple[Any, ...], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> asyncio.Future[int]: + """Start a write. The write's task takes the call options.""" + _check_loop() + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + return _access(handle).publish(values) def subscribe( self, + handle: StreamHandle[Any], *, - offset: int | End | None = None, - last: int | None = None, - wait_new_records: bool = True, - create_wait: timedelta | None = None, - ) -> Subscription[T]: - """Read the Stream's items in order. See :py:meth:`ClientStream.subscribe`.""" + offset: int | End | None, + last: int | None, + wait_new_records: bool, + create_wait: timedelta | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> Subscription[Any]: + """Start a read.""" + _check_loop() read = read_spec( - self._payload_type, + handle._payload_type, offset=offset, last=last, wait_new_records=wait_new_records, create_wait=create_wait, ) - access = self._access() - return subscription(lambda position: access.items(read, position), read) - - async def get_offset(self) -> int: - """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. + access = _access(handle) + return subscription( + lambda position: access.items(read, position), + read, + _backend.options_of(rpc_metadata, rpc_timeout), + ) - Offsets are sparse and above zero, so ``offset=0`` reads from the start. - """ - return await self._access().get_offset() + async def get_offset( + self, + handle: StreamHandle[Any], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> int: + """Get the next offset.""" + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + return await _access(handle).get_offset() async def close( - self, *, result: Any | None = None, failure: BaseException | None = None + self, + handle: StreamHandle[Any], + result: Any, + failure: BaseException | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, ) -> None: - """Raise :py:class:`StreamNotOwnerError`, since an Activity can't close a Stream. + """Raise :py:class:`StreamNotOwnerError`, since an Activity can't close a Stream.""" + raise _not_owner(handle) - A Workflow's Stream closes only from Workflow code, and a client's Stream only from a - client. The rule includes a Local Activity of the run chain that created the Stream. A - Local Activity runs before its Workflow Task commits. If the task then failed, the close - would stay, and the task's retry couldn't create the Stream again. - """ - del result, failure - raise _not_owner(self) + async def enter(self, handle: StreamHandle[Any]) -> None: + """Raise before the block, since the close at its end would raise.""" + raise _not_owner(handle) + + async def exit(self, handle: StreamHandle[Any], exc: BaseException | None) -> None: + """Raise, since an Activity can't close a Stream.""" + raise _not_owner(handle) class SyncSubscription(Generic[T]): @@ -161,7 +189,7 @@ class SyncSubscription(Generic[T]): """ def __init__( - self, handle: SyncActivityStream[Any], subscription: Subscription[T] + self, handle: SyncStreamHandle[Any], subscription: Subscription[T] ) -> None: """Create a read. Use a handle's ``subscribe`` instead.""" self._subscription = subscription @@ -185,20 +213,23 @@ def close(self) -> None: self._steps.close() -class SyncActivityStream(StreamHandle[T]): - """A Stream in a sync Activity. Its calls block the Activity's thread. +class SyncStreamHandle(Generic[T]): + """A Stream handle in a sync Activity. Its calls block the Activity's thread. - The calls run on the Worker's event loop, so the Activity thread only waits for them. + It wraps a :py:class:`StreamHandle`, such as a field of the Activity's input. Each call runs + the handle's call of the same name on the Worker's event loop, and blocks until it ends. A + publish returns the offset, not a future. .. warning:: This API is experimental. """ - def _access(self) -> Access: - return _access(self) + def __init__(self, handle: StreamHandle[T]) -> None: + """Wrap ``handle`` for a sync Activity.""" + self._handle = handle def _run(self, call: Coroutine[Any, Any, _R]) -> _R: - loop = _bound(self).loop + loop = _bound(self._handle).loop try: on_loop = asyncio.get_running_loop() is loop except RuntimeError: @@ -206,7 +237,7 @@ def _run(self, call: Coroutine[Any, Any, _R]) -> _R: if on_loop: call.close() raise RuntimeError( - "A SyncActivityStream blocks, so use ActivityStream in an async Activity" + "A SyncStreamHandle blocks, so use the StreamHandle in an async Activity" ) # The future exists before the call is scheduled, so a cancel raised in this thread at # any point after the schedule reaches the except below, which cancels the call. @@ -238,7 +269,11 @@ def publish(self, *values: T) -> int: Raises: StreamNotFoundError: Nobody created the Stream. """ - return self._run(self._access().publish(values)) + + async def publish() -> int: + return await self._handle.publish(*values) + + return self._run(publish()) def subscribe( self, @@ -248,18 +283,17 @@ def subscribe( wait_new_records: bool = True, create_wait: timedelta | None = None, ) -> SyncSubscription[T]: - """Read the Stream's items in order. See :py:meth:`ClientStream.subscribe`.""" - read = read_spec( - self._payload_type, - offset=offset, - last=last, - wait_new_records=wait_new_records, - create_wait=create_wait, - ) - access = self._access() - return SyncSubscription( - self, subscription(lambda position: access.items(read, position), read) - ) + """Read the Stream's items in order. See :py:meth:`StreamHandle.subscribe`.""" + + async def subscribe() -> Subscription[T]: + return self._handle.subscribe( + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + + return SyncSubscription(self, self._run(subscribe())) def _items( self, items: Subscription[T] @@ -279,14 +313,13 @@ def get_offset(self) -> int: Offsets are sparse and above zero, so ``offset=0`` reads from the start. """ - return self._run(self._access().get_offset()) + return self._run(self._handle.get_offset()) def close( - self, *, result: Any | None = None, failure: BaseException | None = None + self, *, result: Any = None, failure: BaseException | None = None ) -> None: - """Raise :py:class:`StreamNotOwnerError`. See :py:meth:`ActivityStream.close`.""" - del result, failure - raise _not_owner(self) + """Raise :py:class:`StreamNotOwnerError`, since an Activity can't close a Stream.""" + self._run(self._handle.close(result=result, failure=failure)) async def _next(items: Subscription[T]) -> Record[T] | Superseded | Closed: diff --git a/temporalio/common/_streams/_buffered.py b/temporalio/common/_streams/_buffered.py index fdf65c637..845c26caa 100644 --- a/temporalio/common/_streams/_buffered.py +++ b/temporalio/common/_streams/_buffered.py @@ -7,9 +7,9 @@ from types import TracebackType from typing import Generic, TypeVar -from temporalio.common._streams import _activity -from temporalio.common._streams._activity_stream import ActivityStream -from temporalio.common._streams._stream_client import ClientStream +import temporalio.workflow +from temporalio.common._streams import _activity, _errors +from temporalio.common._streams._handle import StreamHandle T = TypeVar("T") @@ -39,17 +39,26 @@ class BufferedPublisher(Generic[T]): def __init__( self, - target: ActivityStream[T] | ClientStream[T], + handle: StreamHandle[T], *, max_delay: timedelta = timedelta(milliseconds=100), ) -> None: """Create a publisher. Args: - target: An ``ActivityStream`` or a ``ClientStream``. + handle: The handle of the Stream. max_delay: The longest time that a value waits in the buffer. + + Raises: + StreamUnsupportedError: The code is Workflow code. There, the publishes that wait + behind a write to the same Stream go as one Local Activity. """ - self._stream = target + if temporalio.workflow.in_workflow(): + raise _errors.StreamUnsupportedError( + "Workflow code can't use a BufferedPublisher. Its publishes that wait behind " + "a write already go as one Local Activity." + ) + self._handle = handle self._max_delay = max_delay self._buffer: list[T] = [] # When the oldest buffered value arrived, so the next batch is at most max_delay later. @@ -151,7 +160,7 @@ async def _write(self, batch: list[T]) -> None: self._schedule() async def _send(self, batch: list[T]) -> int: - return await self._stream._access().publish(batch) + return await self._handle.publish(*batch) @property def unwritten(self) -> int: diff --git a/temporalio/common/_streams/_handle.py b/temporalio/common/_streams/_handle.py index 839b64b34..4cbc6ffb8 100644 --- a/temporalio/common/_streams/_handle.py +++ b/temporalio/common/_streams/_handle.py @@ -288,6 +288,11 @@ class StreamHandle(Generic[T]): Workflow passes a handle to an Activity, a child Workflow or the next run of continue-as-new as an argument, or hands it out through a query. + An Activity, including a Local Activity, writes to and reads a Stream as the Producer of + its attempt, and can't close it. The end of the attempt waits for the attempt's writes, and + a write error that no caller got fails the attempt. A sync Activity wraps the handle in + :py:class:`SyncStreamHandle`. + In client code, a handle uses the client that made it, such as with :py:meth:`temporalio.client.Client.get_stream_handle`, or the client that decoded it, such as from a query, an Update or a Workflow result. A handle that no client decoded, such as @@ -349,6 +354,12 @@ def _calls(self) -> ContextCalls: from temporalio.common._streams import _workflow return _workflow.WorkflowCalls() + from temporalio.common._streams import _activity + + if _activity.current.get() is not None: + from temporalio.common._streams import _activity_stream + + return _activity_stream.ActivityCalls() from temporalio.common._streams import _stream_client return _stream_client.calls(self) @@ -395,7 +406,8 @@ def publish( Stream's backend, which fails the Workflow Task and not the run. temporalio.workflow.ReadOnlyContextError: The Workflow code must not change state, such as a query handler. - StreamUnsupportedError: Workflow code gave ``rpc_metadata`` or ``rpc_timeout``. + StreamUnsupportedError: Workflow code gave ``rpc_metadata`` or ``rpc_timeout``, or + the code is a sync Activity, which uses :py:class:`SyncStreamHandle`. """ return self._calls().publish(self, values, rpc_metadata, rpc_timeout) @@ -430,7 +442,8 @@ def subscribe( that waits for new Records returns before the deadline. Raises: - StreamUnsupportedError: The code is Workflow code, which can't read a Stream. + StreamUnsupportedError: The code is Workflow code, which can't read a Stream, or a + sync Activity, which uses :py:class:`SyncStreamHandle`. """ return self._calls().subscribe( self, @@ -483,9 +496,10 @@ async def close( Later writes in the run raise :py:class:`StreamClosedError`. Raises: - StreamNotOwnerError: The code doesn't own the Stream, such as a client with a - Workflow's Stream. Only the run chain closes a Workflow's Stream, and only a - client closes a client's Stream. + StreamNotOwnerError: The code doesn't own the Stream. Only the run chain closes a + Workflow's Stream, and only a client closes a client's Stream, so an Activity + closes none. A client with a Workflow's handle, Workflow code with another + owner's handle, and an Activity raise with no call. """ await self._calls().close(self, result, failure, rpc_metadata, rpc_timeout) diff --git a/tests/contrib/streams/test_activity_handles.py b/tests/contrib/streams/test_activity_handles.py index 9e81e27de..495988c7e 100644 --- a/tests/contrib/streams/test_activity_handles.py +++ b/tests/contrib/streams/test_activity_handles.py @@ -21,7 +21,7 @@ from temporalio.client import Client, WorkflowHandle from temporalio.common import RetryPolicy from temporalio.common._streams import _activity -from temporalio.exceptions import CancelledError +from temporalio.exceptions import ActivityError, ApplicationError, CancelledError from tests.contrib.streams._support import workflow_stream from tests.helpers import new_worker @@ -37,13 +37,13 @@ class Other(common.StreamDefinition[str]): ... @dataclass class CallModelInput: prompt: str - stream: common.ActivityStream[str] + stream: common.StreamHandle[str] @dataclass class SyncInput: prompt: str - stream: common.SyncActivityStream[str] + stream: common.StreamHandle[str] async def read( @@ -73,13 +73,14 @@ async def call_model(input: CallModelInput) -> list[int]: @activity.defn def call_model_sync(input: SyncInput) -> int: - last = input.stream.publish(*input.prompt.split()) - assert input.stream.get_offset() == last + 1 + stream = common.SyncStreamHandle(input.stream) + last = stream.publish(*input.prompt.split()) + assert stream.get_offset() == last + 1 return last @activity.defn -async def read_back(stream: common.ActivityStream[str]) -> list[Any]: +async def read_back(stream: common.StreamHandle[str]) -> list[Any]: return [ item.data async for item in stream.subscribe(wait_new_records=False) @@ -94,17 +95,17 @@ async def run(self, prompt: str) -> list[Any]: stream = await workflow.create_stream(Tokens) offsets = await workflow.execute_activity( call_model, - CallModelInput(prompt, common.ActivityStream.from_handle(stream)), + CallModelInput(prompt, stream), start_to_close_timeout=timedelta(seconds=30), ) last = await workflow.execute_activity( call_model_sync, - SyncInput("sync words", common.SyncActivityStream.from_handle(stream)), + SyncInput("sync words", stream), start_to_close_timeout=timedelta(seconds=30), ) read = await workflow.execute_activity( read_back, - common.ActivityStream.from_handle(stream), + stream, start_to_close_timeout=timedelta(seconds=30), ) await stream.close() @@ -158,7 +159,7 @@ async def run(self, key: str) -> None: await stream.publish("start") await workflow.execute_activity( flaky, - CallModelInput(key, common.ActivityStream.from_handle(stream)), + CallModelInput(key, stream), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), ) @@ -182,8 +183,8 @@ async def test_retry_withdraws_a_received_stream(stream_client: Client): @dataclass class FinishedStreams: key: str - closed: common.ActivityStream[str] - expired: common.ActivityStream[str] + closed: common.StreamHandle[str] + expired: common.StreamHandle[str] @activity.defn @@ -198,7 +199,7 @@ async def reads_after_a_failed_attempt(input: FinishedStreams) -> list[Any]: @workflow.defn class ReadsFinishedStreams: @workflow.run - async def run(self, key: str, expired: common.WorkflowStream[str]) -> list[Any]: + async def run(self, key: str, expired: common.StreamHandle[str]) -> list[Any]: stream = await workflow.create_stream(Tokens) await stream.publish("turn") await stream.close() @@ -206,8 +207,8 @@ async def run(self, key: str, expired: common.WorkflowStream[str]) -> list[Any]: reads_after_a_failed_attempt, FinishedStreams( key, - common.ActivityStream.from_handle(stream), - common.ActivityStream.from_handle(expired), + stream, + expired, ), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), @@ -228,7 +229,7 @@ async def test_retry_goes_on_with_a_closed_or_expired_stream(stream_client: Clie ) as worker: result = await stream_client.execute_workflow( ReadsFinishedStreams.run, - args=[key, common.WorkflowStream.from_handle(expired)], + args=[key, expired], id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) @@ -257,7 +258,7 @@ async def run(self, key: str) -> None: stream = await workflow.create_stream(Tokens) await workflow.execute_activity( fails_after_it_is_shown, - CallModelInput(key, common.ActivityStream.from_handle(stream)), + CallModelInput(key, stream), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), ) @@ -304,7 +305,7 @@ async def test_reader_withdraws_a_shown_record_after_a_reconnect( @activity.defn -async def wrong_backend(stream: common.ActivityStream[str]) -> str: +async def wrong_backend(stream: common.StreamHandle[str]) -> str: try: await stream.publish("x") except common.StreamConfigError: @@ -313,7 +314,7 @@ async def wrong_backend(stream: common.ActivityStream[str]) -> str: @activity.defn -async def close_it(stream: common.ActivityStream[str]) -> str: +async def close_it(stream: common.StreamHandle[str]) -> str: await stream.publish("last") try: await stream.close(result="closed by activity") @@ -329,24 +330,24 @@ async def run(self) -> list[str]: stream = await workflow.create_stream(Tokens) backend = await workflow.execute_activity( wrong_backend, - common.ActivityStream[str](id=stream.id, backend="elsewhere"), + common.StreamHandle[str](id=stream.id, backend="elsewhere"), start_to_close_timeout=timedelta(seconds=30), ) owner = await workflow.execute_activity( close_it, - common.ActivityStream.from_handle(stream), + stream, start_to_close_timeout=timedelta(seconds=30), ) # A Local Activity runs before its Workflow Task commits, so it can't close it either. local = await workflow.execute_local_activity( close_it, - common.ActivityStream.from_handle(stream), + stream, start_to_close_timeout=timedelta(seconds=30), ) # A handle with only the id names no Workflow, and no Activity can close a Stream. by_id = await workflow.execute_local_activity( close_it, - common.ActivityStream[str](id=stream.id), + common.StreamHandle[str](id=stream.id), start_to_close_timeout=timedelta(seconds=30), ) await stream.close(result="closed by the run") @@ -371,7 +372,138 @@ async def test_activity_checks_the_backend_and_the_owner(stream_client: Client): assert items[-1].result == "closed by the run" -class _Batches(common.ActivityStream[int]): +@activity.defn +async def refused_calls(stream: common.StreamHandle[str]) -> list[str]: + refused = [] + try: + async with stream: + refused.append("entered") + except common.StreamNotOwnerError as err: + refused.append(type(err).__name__) + for call in ( + lambda: stream.producer("ingest", 1), + lambda: stream.publish_nowait("x"), + ): + try: + call() + except common.StreamUnsupportedError as err: + refused.append(type(err).__name__) + return refused + + +@activity.defn +def publishes_without_the_sync_handle(stream: common.StreamHandle[str]) -> str: + try: + stream.publish("x") + except common.StreamUnsupportedError as err: + return type(err).__name__ + return "wrote" + + +@workflow.defn +class RefusedCalls: + @workflow.run + async def run(self) -> list[str]: + stream = await workflow.create_stream(Tokens) + refused = await workflow.execute_activity( + refused_calls, stream, start_to_close_timeout=timedelta(seconds=30) + ) + refused.append( + await workflow.execute_activity( + publishes_without_the_sync_handle, + stream, + start_to_close_timeout=timedelta(seconds=30), + ) + ) + try: + common.BufferedPublisher(stream) + except common.StreamUnsupportedError as err: + refused.append(type(err).__name__) + await stream.close() + return refused + + +async def test_calls_that_an_activity_cant_make_raise(stream_client: Client): + async with new_worker( + stream_client, + RefusedCalls, + activities=[refused_calls, publishes_without_the_sync_handle], + activity_executor=_executor(), + ) as worker: + handle = await stream_client.start_workflow( + RefusedCalls.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert await handle.result() == [ + "StreamNotOwnerError", + *["StreamUnsupportedError"] * 4, + ] + items = await read(handle) + assert records(items) == [] + + +@activity.defn +async def leaves_a_failed_write(stream: common.StreamHandle[str]) -> str: + # Nobody created the Stream, and no caller waits for the write. + stream.publish("lost") + return "returned" + + +@activity.defn +async def enters_a_client_stream(stream: common.StreamHandle[str]) -> str: + try: + async with stream: + await stream.publish("in the block") + except common.StreamNotOwnerError as err: + return StreamErrorReason.Name(StreamErrorReason.ValueType(err.reason)) + return "closed" + + +@workflow.defn +class LostAndClientWrites: + @workflow.run + async def run( + self, missing_id: str, client_stream: common.StreamHandle[str] + ) -> list[str]: + try: + lost = await workflow.execute_activity( + leaves_a_failed_write, + common.StreamHandle[str](id=missing_id), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(maximum_attempts=1), + ) + except ActivityError as err: + assert isinstance(err.cause, ApplicationError) + lost = err.cause.type or "" + closed = await workflow.execute_activity( + enters_a_client_stream, + client_stream, + start_to_close_timeout=timedelta(seconds=30), + ) + return [lost, closed] + + +async def test_attempt_end_raises_a_lost_write_and_a_client_stream_refuses_the_block( + stream_client: Client, +): + stream = await stream_client.create_stream(Tokens, id=f"client-{uuid.uuid4()}") + async with new_worker( + stream_client, + LostAndClientWrites, + activities=[leaves_a_failed_write, enters_a_client_stream], + ) as worker: + result = await stream_client.execute_workflow( + LostAndClientWrites.run, + args=[f"missing-{uuid.uuid4()}", stream], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + # Only a client closes a client's Stream, so the block fails before it runs, with no call. + assert result == ["StreamNotFoundError", "STREAM_ERROR_REASON_NOT_OWNER"] + items = [item async for item in stream.subscribe(wait_new_records=False)] + assert records(items) == [] + + +class _Batches(common.StreamHandle[int]): # A Stream whose writes go to a list, so the batches show with no backend. def __init__(self, fail: bool = False) -> None: super().__init__(id="batches") @@ -380,8 +512,8 @@ def __init__(self, fail: bool = False) -> None: self.in_flight = 0 self.most_in_flight = 0 - def _access(self) -> Any: - return SimpleNamespace(publish=self._write) + def publish(self, *values: int, **options: Any) -> asyncio.Future[int]: + return asyncio.ensure_future(self._write(values)) async def _write(self, values: Sequence[int]) -> int: if self.fail: @@ -430,7 +562,7 @@ async def test_buffered_publisher_raises_a_failed_batch(): @activity.defn -async def local_step(stream: common.ActivityStream[str]) -> list[Any]: +async def local_step(stream: common.StreamHandle[str]) -> list[Any]: # Each Workflow write is an awaited Local Activity, so the run's create and Records are in # the Store before a later Local Activity starts. seen = [ @@ -450,7 +582,7 @@ async def run(self) -> list[Any]: await stream.publish("before") seen = await workflow.execute_local_activity( local_step, - common.ActivityStream.from_handle(stream), + stream, start_to_close_timeout=timedelta(seconds=30), ) await stream.publish("after") @@ -496,8 +628,8 @@ async def never_flushes(input: CallModelInput) -> None: @dataclass class ManyStreams: key: str - targets: list[common.ActivityStream[str]] - maybe: common.ActivityStream[str] | None = None + targets: list[common.StreamHandle[str]] + maybe: common.StreamHandle[str] | None = None @activity.defn @@ -511,10 +643,10 @@ async def flaky_many(input: ManyStreams) -> None: @activity.defn -def sync_read(stream: common.SyncActivityStream[str]) -> list[Any]: +def sync_read(stream: common.StreamHandle[str]) -> list[Any]: return [ item.data - for item in stream.subscribe(wait_new_records=False) + for item in common.SyncStreamHandle(stream).subscribe(wait_new_records=False) if isinstance(item, common.Record) ] @@ -527,22 +659,22 @@ async def run(self, key: str) -> list[Any]: other = await workflow.create_stream(Other) await workflow.execute_activity( never_flushes, - CallModelInput("a b", common.ActivityStream.from_handle(stream)), + CallModelInput("a b", stream), start_to_close_timeout=timedelta(seconds=30), ) await workflow.execute_activity( flaky_many, ManyStreams( key, - [common.ActivityStream.from_handle(stream)], - common.ActivityStream.from_handle(other), + [stream], + other, ), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), ) return await workflow.execute_activity( sync_read, - common.SyncActivityStream.from_handle(stream), + stream, start_to_close_timeout=timedelta(seconds=30), ) @@ -603,7 +735,7 @@ def _caller( call: Coroutine[Any, Any, None], errors: list[BaseException], ) -> threading.Thread: - stream: common.SyncActivityStream[Any] = common.SyncActivityStream(id="s") + stream = common.SyncStreamHandle(common.StreamHandle[Any](id="s")) token = _activity.current.set( cast(Any, SimpleNamespace(loop=loop, backend_name="", default=True)) ) @@ -715,7 +847,7 @@ async def run(self) -> str: @workflow.defn class ReceivesAStream: @workflow.run - async def run(self, stream: common.WorkflowStream[str]) -> str: + async def run(self, stream: common.StreamHandle[str]) -> str: # The Activity gets no Stream, but the run received one. return await workflow.execute_activity( plain, stream.id, start_to_close_timeout=timedelta(seconds=10) @@ -725,13 +857,13 @@ async def run(self, stream: common.WorkflowStream[str]) -> str: @workflow.defn class RetriesAfterContinueAsNew: @workflow.run - async def run(self, key: str, stream: common.WorkflowStream[str] | None) -> None: + async def run(self, key: str, stream: common.StreamHandle[str] | None) -> None: if stream is None: created = await workflow.create_stream(Tokens) workflow.continue_as_new(args=[key, created]) await workflow.execute_activity( flaky, - CallModelInput(key, common.ActivityStream.from_handle(stream)), + CallModelInput(key, stream), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), ) @@ -776,7 +908,7 @@ async def test_only_a_run_that_uses_streams_keys_its_activities(stream_client: C assert await keyed(creates, "flaky") == [True] receives = await stream_client.start_workflow( ReceivesAStream.run, - common.WorkflowStream.from_handle(workflow_stream(creates, Tokens)), + workflow_stream(creates, Tokens), id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) diff --git a/tests/contrib/streams/test_pydantic.py b/tests/contrib/streams/test_pydantic.py index d7fadc069..cbcdd2677 100644 --- a/tests/contrib/streams/test_pydantic.py +++ b/tests/contrib/streams/test_pydantic.py @@ -16,12 +16,12 @@ class Tokens(common.StreamDefinition[str]): ... class CallModelInput(pydantic.BaseModel): prompt: str - stream: common.ActivityStream[str] + stream: common.StreamHandle[str] class TurnResult(pydantic.BaseModel): text: str - stream: common.ClientStream[str] + stream: common.StreamHandle[str] @activity.defn @@ -37,15 +37,11 @@ async def run(self, prompt: str) -> TurnResult: stream = await workflow.create_stream(Tokens) await workflow.execute_activity( call_model, - CallModelInput( - prompt=prompt, stream=common.ActivityStream.from_handle(stream) - ), + CallModelInput(prompt=prompt, stream=stream), start_to_close_timeout=timedelta(seconds=30), ) await stream.close(result=prompt.upper()) - return TurnResult( - text=prompt.upper(), stream=common.ClientStream.from_handle(stream) - ) + return TurnResult(text=prompt.upper(), stream=stream) async def test_handles_in_models_round_trip_on_the_memory_store(client: Client): @@ -65,7 +61,7 @@ async def test_handles_in_models_round_trip_on_the_memory_store(client: Client): assert result.stream.workflow_id == handle.id assert result.stream.first_execution_run_id == handle.first_execution_run_id assert result.stream.backend == "memory" - stream = client.stream(Tokens, id=result.stream.id) + stream = client.get_stream_handle(Tokens, id=result.stream.id) items = [item async for item in stream.subscribe()] assert [i.data for i in items if isinstance(i, common.Record)] == ["hi", "there"] assert isinstance(items[-1], common.Closed) and items[-1].result is not None diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index ec1ce8ee8..5f432aa1f 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -19,7 +19,7 @@ class Tokens(common.StreamDefinition[str]): ... @dataclass class CallModelInput: prompt: str - stream: common.ActivityStream[str] + stream: common.StreamHandle[str] async def read_all(stream: common.StreamHandle[Any]) -> list[Any]: @@ -86,9 +86,7 @@ async def run(self, key: str) -> None: first = attempts.setdefault(key, 0) == 0 await workflow.execute_local_activity( local_tokens, - CallModelInput( - "0 1" if first else "2 3", common.ActivityStream.from_handle(stream) - ), + CallModelInput("0 1" if first else "2 3", stream), start_to_close_timeout=timedelta(seconds=30), ) if first: @@ -135,7 +133,7 @@ async def run(self, key: str) -> None: attempts[key] = 1 await workflow.execute_local_activity( local_tokens, - CallModelInput("0 1", common.ActivityStream.from_handle(stream)), + CallModelInput("0 1", stream), start_to_close_timeout=timedelta(seconds=30), ) raise RuntimeError("fail this Workflow Task once") From d344a3b1ee9c0252598bf97658f1cb992f22cb04 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 3 Nov 2026 01:02:11 -0700 Subject: [PATCH 4/4] Moved the Activity calls to `Stream`, and added `SyncStream`. --- temporalio/client/__init__.py | 5 +- temporalio/client/_streams.py | 322 +++++++++++++++++- temporalio/common/__init__.py | 11 - temporalio/common/_streams/__init__.py | 8 - temporalio/common/_streams/_activity.py | 3 + .../common/_streams/_activity_stream.py | 299 ++++++---------- temporalio/common/_streams/_buffered.py | 185 ---------- temporalio/common/_streams/_extension.py | 3 + temporalio/common/_streams/_stream_client.py | 10 +- .../contrib/streams/test_activity_handles.py | 272 +++++++++------ tests/contrib/streams/test_pydantic.py | 15 +- tests/contrib/streams/test_redis.py | 6 +- tests/contrib/streams/test_stream_client.py | 4 +- 13 files changed, 615 insertions(+), 528 deletions(-) delete mode 100644 temporalio/common/_streams/_buffered.py diff --git a/temporalio/client/__init__.py b/temporalio/client/__init__.py index dafe658fc..70ff482ad 100644 --- a/temporalio/client/__init__.py +++ b/temporalio/client/__init__.py @@ -185,7 +185,7 @@ ScheduleUpdate, ScheduleUpdateInput, ) -from ._streams import Stream +from ._streams import BufferedPublisher, Stream, SyncStream, SyncSubscription from ._worker_versioning import ( BuildIdOp, BuildIdOpAddNewCompatible, @@ -398,5 +398,8 @@ "ParamType", "ReturnType", "SelfType", + "BufferedPublisher", "Stream", + "SyncStream", + "SyncSubscription", ] diff --git a/temporalio/client/_streams.py b/temporalio/client/_streams.py index e06c5d39c..0ac827ea0 100644 --- a/temporalio/client/_streams.py +++ b/temporalio/client/_streams.py @@ -7,15 +7,24 @@ from __future__ import annotations import asyncio -from collections.abc import Mapping +from collections.abc import Generator, Mapping from datetime import timedelta from types import TracebackType from typing import TYPE_CHECKING, Any, Generic, TypeVar from typing_extensions import Self +import temporalio.workflow + if TYPE_CHECKING: - from temporalio.common import End, StreamHandle, Subscription + from temporalio.common import ( + Closed, + End, + Record, + StreamHandle, + Subscription, + Superseded, + ) from temporalio.common._streams._stream_client import StreamCalls T = TypeVar("T") @@ -33,6 +42,12 @@ class Stream(Generic[T]): returns a Stream that writes as a named Producer. A write error that no caller got raises at the next call of the Stream. + In an Activity, including a Local Activity, a Stream from the Worker's client, such as from + ``activity.client()``, writes and reads as the Producer of the attempt. So a retry withdraws + the Records of the failed attempt. It can't close the Stream. The end of the attempt waits + for the attempt's writes, and a write error that no caller got fails the attempt. A sync + Activity wraps the Stream in :py:class:`SyncStream`. + .. warning:: This API is experimental. """ @@ -78,6 +93,8 @@ def publish( Raises: StreamClosedError: This object closed the Stream. + StreamUnsupportedError: The code is a sync Activity, which uses + :py:class:`SyncStream`. """ return self._calls.publish(values, rpc_metadata, rpc_timeout) @@ -110,6 +127,10 @@ def subscribe( rpc_metadata: Headers of each call of the read. See :py:meth:`publish`. rpc_timeout: The deadline of each call of the read, not of the whole read. A call that waits for new Records returns before the deadline. + + Raises: + StreamUnsupportedError: The code is a sync Activity, which uses + :py:class:`SyncStream`. """ from temporalio.common._streams._access import read_spec @@ -150,10 +171,10 @@ async def close( canceled. ``rpc_metadata`` and ``rpc_timeout`` are as for :py:meth:`publish`. Raises: - StreamNotOwnerError: The code doesn't own the Stream, such as a client with a - Workflow's Stream. Only the run chain closes a Workflow's Stream, and only a - client closes a client's Stream. A client with a Workflow's Stream raises with - no call. + StreamNotOwnerError: The code doesn't own the Stream. Only the run chain closes a + Workflow's Stream, and only a client closes a client's Stream, so an Activity + closes none. A client with a Workflow's Stream and an Activity raise with no + call. """ await self._calls.close(result, failure, rpc_metadata, rpc_timeout) @@ -192,6 +213,9 @@ def producer(self, name: str, epoch: int) -> Stream[T]: Args: name: The Producer name, the same across restarts. epoch: The writer's epoch, 1 or more. + + Raises: + StreamUnsupportedError: The Stream writes as an Activity attempt. """ return Stream( self._handle, self._payload_type, self._calls.producer(name, epoch) @@ -200,3 +224,289 @@ def producer(self, name: str, epoch: int) -> Stream[T]: def __repr__(self) -> str: """Show the handle.""" return f"{type(self).__name__}({self._handle!r})" + + +class SyncSubscription(Generic[T]): + """The items of one read of a Stream in a sync Activity, in order. Iterate it with ``for``. + + Each iteration blocks the Activity's thread. See :py:class:`temporalio.common.Subscription`. + + .. warning:: + This API is experimental. + """ + + def __init__(self, subscription: Subscription[T]) -> None: + """Create a read. Use a ``SyncStream``'s ``subscribe`` instead.""" + from temporalio.common._streams import _activity_stream + + self._subscription = subscription + self._steps: Generator[Record[T] | Superseded | Closed, None, None] = ( + _activity_stream.sync_items(subscription) + ) + + @property + def next_offset(self) -> int | None: + """The offset to resume from. See :py:attr:`temporalio.common.Subscription.next_offset`.""" + return self._subscription.next_offset + + def __iter__(self) -> SyncSubscription[T]: + """The read is its own iterator.""" + return self + + def __next__(self) -> Record[T] | Superseded | Closed: + """The next item.""" + return next(self._steps) + + def close(self) -> None: + """End the read early. It makes no more calls.""" + self._steps.close() + + +class SyncStream(Generic[T]): + """A Stream in a sync Activity. Its calls block the Activity's thread. + + It wraps a :py:class:`Stream` of the Worker's client, such as for a handle of the Activity's + input. A sync Activity has no ``activity.client()``, so it keeps the Worker's client, such as + in an Activity class. Each call runs the Stream's call of the same name on the Worker's event + loop, and blocks until it ends. A publish returns the offset, not a future. + + .. warning:: + This API is experimental. + """ + + def __init__(self, stream: Stream[T]) -> None: + """Wrap ``stream`` for a sync Activity.""" + self._stream = stream + + def publish(self, *values: T) -> int: + """Write one Record for each value, in order, and return the last Record's offset. + + A list is one value. + + Raises: + StreamNotFoundError: Nobody created the Stream. + """ + from temporalio.common._streams import _activity_stream + + async def publish() -> int: + return await self._stream.publish(*values) + + return _activity_stream.run_blocking(publish()) + + def subscribe( + self, + *, + offset: int | End | None = None, + last: int | None = None, + wait_new_records: bool = True, + create_wait: timedelta | None = None, + ) -> SyncSubscription[T]: + """Read the Stream's items in order. See :py:meth:`Stream.subscribe`.""" + from temporalio.common._streams import _activity_stream + + async def subscribe() -> Subscription[T]: + return self._stream.subscribe( + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + + return SyncSubscription(_activity_stream.run_blocking(subscribe())) + + def get_offset(self) -> int: + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. + + Offsets are sparse and above zero, so ``offset=0`` reads from the start. + """ + from temporalio.common._streams import _activity_stream + + return _activity_stream.run_blocking(self._stream.get_offset()) + + def close( + self, *, result: Any = None, failure: BaseException | None = None + ) -> None: + """Close the Stream. See :py:meth:`Stream.close`. + + Raises: + StreamNotOwnerError: The Stream writes as the Activity attempt, which can't close it. + """ + from temporalio.common._streams import _activity_stream + + _activity_stream.run_blocking( + self._stream.close(result=result, failure=failure) + ) + + +# Records in one batch of the crate's writer. A full batch is written at once. +_BATCH = 1000 + + +class BufferedPublisher(Generic[T]): + """Batches the publishes of one Stream. + + Use it in an Activity or a client that writes often. :py:meth:`publish` returns at once. + The publisher writes the buffered values as one batch at most ``max_delay`` after the first + of them, or at once when they fill a batch. One write is in flight at a time, and values + published during it wait for the next one. Use it in ``async with``, which writes the + remaining values at the end of the block. For a Stream that writes as an Activity attempt, + the end of the attempt also writes the remaining values. + + ``await flush()`` is the backpressure point. A caller that publishes faster than the Store + accepts Records must await it from time to time, so the buffer doesn't grow without bound. + + After a batch fails, the publisher writes nothing more, so a reader sees no gap. Each later + call raises the error. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + stream: Stream[T], + *, + max_delay: timedelta = timedelta(milliseconds=100), + ) -> None: + """Create a publisher. + + Args: + stream: The Stream to publish to. + max_delay: The longest time that a value waits in the buffer. + + Raises: + StreamUnsupportedError: The code is Workflow code. There, the publishes that wait + behind a write to the same Stream go as one Local Activity. + """ + from temporalio.common._streams import _activity_stream, _errors + + if temporalio.workflow.in_workflow(): + raise _errors.StreamUnsupportedError( + "Workflow code can't use a BufferedPublisher. Its publishes that wait behind " + "a write already go as one Local Activity." + ) + self._stream = stream + self._max_delay = max_delay + self._buffer: list[T] = [] + # When the oldest buffered value arrived, so the next batch is at most max_delay later. + self._first_at = 0.0 + self._timer: asyncio.TimerHandle | None = None + self._writing: asyncio.Task[None] | None = None + self._published = 0 + # Count of the values that a write took, written or not. + self._settled = 0 + self._error: BaseException | None = None + self._last_offset = 0 + self._unwritten = 0 + self._reported = False + attempt = _activity_stream.attempt_of(stream) + if attempt is not None: + attempt.publishers.append(self) + + def publish(self, *values: T) -> None: + """Add one Record for each value to the buffer, and return at once. + + A list is one value. + + Raises: + StreamError: A batch failed to write, so the publisher takes no more values. + """ + self._raise_error() + if not values: + raise TypeError("publish takes one or more values") + loop = asyncio.get_running_loop() + if not self._buffer: + self._first_at = loop.time() + self._buffer.extend(values) + self._published += len(values) + self._schedule() + + async def flush(self) -> int: + """Write every value published before the call, and return the last written offset. + + Returns 0 if the publisher hasn't written anything. + """ + target = self._published + while self._settled < target: + if self._writing is None: + self._start() + assert self._writing is not None + await asyncio.shield(self._writing) + self._raise_error() + return self._last_offset + + def _raise_error(self) -> None: + if self._error is not None: + self._reported = True + raise self._error + + async def _flush_at_attempt_end(self) -> None: + # Fail the attempt for unwritten values, unless the Activity already got the error and + # continued. + reported = self._reported + try: + await self.flush() + except Exception: + if not reported: + raise + + def _schedule(self) -> None: + if self._writing is not None or not self._buffer: + return + if len(self._buffer) >= _BATCH: + self._start() + elif self._timer is None: + self._timer = asyncio.get_running_loop().call_at( + self._first_at + self._max_delay.total_seconds(), self._start + ) + + def _start(self) -> None: + if self._timer is not None: + self._timer.cancel() + self._timer = None + if self._writing is not None or not self._buffer: + return + batch, self._buffer = self._buffer, [] + self._writing = asyncio.get_running_loop().create_task(self._write(batch)) + + async def _write(self, batch: list[T]) -> None: + try: + if self._error is None: + self._last_offset = await self._send(batch) + else: + # A failed batch fails the later ones too, so a reader sees no gap. + self._unwritten += len(batch) + except Exception as err: + self._unwritten += len(batch) + self._error = err + finally: + self._settled += len(batch) + self._writing = None + # Schedule the values published during the write, for when their delay ends or they + # fill a batch. + self._schedule() + + async def _send(self, batch: list[T]) -> int: + return await self._stream.publish(*batch) + + @property + def unwritten(self) -> int: + """How many published values a failed batch left unwritten.""" + return self._unwritten + + async def __aenter__(self) -> BufferedPublisher[T]: + """Use the publisher in ``async with``.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Write the buffered values. A write error does not hide an error of the block.""" + try: + await self.flush() + except Exception: + if exc is None: + raise diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index 12bdb4fec..adc62d05e 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -34,11 +34,6 @@ import temporalio.types if TYPE_CHECKING: - from temporalio.common._streams._activity_stream import ( - SyncStreamHandle, - SyncSubscription, - ) - from temporalio.common._streams._buffered import BufferedPublisher from temporalio.common._streams._defn import StreamDefinition, defn from temporalio.common._streams._errors import ( StreamClosedError, @@ -82,7 +77,6 @@ "ActivityIDConflictPolicy", "ActivityIDReusePolicy", "AutoUpgradeVersioningOverride", - "BufferedPublisher", "CloseStatus", "Closed", "Creating", @@ -141,8 +135,6 @@ "StreamUnsupportedError", "Subscription", "Superseded", - "SyncStreamHandle", - "SyncSubscription", "TypedSearchAttributes", "VersioningBehavior", "VersioningOverride", @@ -1609,7 +1601,6 @@ def wait_sync(self, timeout: float | None = None) -> None: # so this module imports the Stream names on first use. _STREAM_NAMES = frozenset( [ - "BufferedPublisher", "CloseStatus", "Closed", "Creating", @@ -1639,8 +1630,6 @@ def wait_sync(self, timeout: float | None = None) -> None: "StreamUnsupportedError", "Subscription", "Superseded", - "SyncStreamHandle", - "SyncSubscription", "WorkflowTaskEpoch", "defn", ] diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index 9f6027ab1..4fa4e2bf0 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -1,10 +1,5 @@ """The Streams API, which :py:mod:`temporalio.common` exports.""" -from temporalio.common._streams._activity_stream import ( - SyncStreamHandle, - SyncSubscription, -) -from temporalio.common._streams._buffered import BufferedPublisher from temporalio.common._streams._defn import StreamDefinition, defn from temporalio.common._streams._errors import ( StreamClosedError, @@ -44,7 +39,6 @@ __all__ = [ "END", - "BufferedPublisher", "CloseStatus", "Closed", "Creating", @@ -73,8 +67,6 @@ "StreamUnsupportedError", "Subscription", "Superseded", - "SyncStreamHandle", - "SyncSubscription", "WorkflowTaskEpoch", "defn", ] diff --git a/temporalio/common/_streams/_activity.py b/temporalio/common/_streams/_activity.py index 1c289fe7f..69fa6d68d 100644 --- a/temporalio/common/_streams/_activity.py +++ b/temporalio/common/_streams/_activity.py @@ -33,6 +33,9 @@ class ActivityStreams: writers: dict[str, StreamWriter] = field(default_factory=dict) # The attempt's BufferedPublishers, which the end of the attempt flushes. publishers: list[Any] = field(default_factory=list) + # The Worker's client. Its Streams in the attempt write as the attempt, and other clients' + # Streams write as those clients. + client: Any = None def writer(self, stream_id: str) -> StreamWriter: """The attempt's writer for a Stream, one per Stream.""" diff --git a/temporalio/common/_streams/_activity_stream.py b/temporalio/common/_streams/_activity_stream.py index 9e367f801..2405cc1cd 100644 --- a/temporalio/common/_streams/_activity_stream.py +++ b/temporalio/common/_streams/_activity_stream.py @@ -1,4 +1,4 @@ -"""Streams in Activity code.""" +"""Streams in Activity code, which write and read as the attempt.""" from __future__ import annotations @@ -6,13 +6,15 @@ import concurrent.futures from collections.abc import Coroutine, Generator, Mapping from datetime import timedelta -from typing import Any, Generic, TypeVar +from typing import Any, TypeVar +from temporalio.client import Stream from temporalio.common._streams import _activity, _backend, _errors -from temporalio.common._streams._access import Access, read_spec, subscription -from temporalio.common._streams._handle import ContextCalls, StreamHandle -from temporalio.common._streams._items import Closed, End, Record, Superseded +from temporalio.common._streams._access import Access, ReadSpec, subscription +from temporalio.common._streams._handle import StreamHandle +from temporalio.common._streams._items import Closed, Record, Superseded from temporalio.common._streams._reader import Subscription +from temporalio.common._streams._stream_client import StreamCalls T = TypeVar("T") _R = TypeVar("_R") @@ -50,25 +52,6 @@ def future_done(future: concurrent.futures.Future[_R]) -> None: future.add_done_callback(future_done) -def _bound(handle: StreamHandle[Any]) -> _activity.ActivityStreams: - streams = _activity.current.get() - if streams is None: - raise _errors.StreamConfigError( - "No Stream backend runs this Activity. Register one on the Worker's client " - "with stream_backends=." - ) - if handle.backend and handle.backend != streams.backend_name: - raise _errors.StreamConfigError( - f"The Stream {handle.id} is of the backend {handle.backend}. The Worker has only " - f"{streams.backend_name}." - ) - if not handle.backend and not streams.default: - raise _errors.no_default_backend( - "Worker", [streams.backend_name], "a handle that has its backend name" - ) - return streams - - def _not_owner(handle: StreamHandle[Any]) -> _errors.StreamNotOwnerError: # Only the creator's side closes a Stream: the run chain a Workflow's Stream, and a client a # client's Stream. Code in an Activity only writes, so the refusal comes before any call. @@ -78,30 +61,17 @@ def _not_owner(handle: StreamHandle[Any]) -> _errors.StreamNotOwnerError: ) -def _access(handle: StreamHandle[Any]) -> Access: - streams = _bound(handle) - stream_id = handle.id - return Access( - backend=streams.backend, - converter=streams.converter, - namespace=streams.namespace, - stream_id=stream_id, - writer=lambda: streams.writer(stream_id), - ) - - def _check_loop() -> None: try: asyncio.get_running_loop() except RuntimeError: raise _errors.StreamUnsupportedError( - "A sync Activity can't wait for a Stream call. Wrap the handle in a " - "SyncStreamHandle." + "A sync Activity can't wait for a Stream call. Wrap the Stream in a SyncStream." ) from None -class ActivityCalls(ContextCalls): - """The Stream calls in an Activity. The Activity can write and read a Stream. +class ActivityCalls(StreamCalls): + """The calls of one Stream in an Activity attempt. They write and read as the attempt. Each publish starts its write at the call, as the Producer of the attempt. An Activity can't close a Stream. A Workflow's Stream closes only from Workflow code, and a client's @@ -110,9 +80,26 @@ class ActivityCalls(ContextCalls): the close would stay, and the task's retry couldn't create the Stream again. """ + def __init__( + self, streams: _activity.ActivityStreams, handle: StreamHandle[Any] + ) -> None: + """Create the calls of ``handle`` in the attempt of ``streams``.""" + self.streams = streams + self._handle = handle + + def _access(self) -> Access: + streams = self.streams + stream_id = self._handle.id + return Access( + backend=streams.backend, + converter=streams.converter, + namespace=streams.namespace, + stream_id=stream_id, + writer=lambda: streams.writer(stream_id), + ) + def publish( self, - handle: StreamHandle[Any], values: tuple[Any, ...], rpc_metadata: Mapping[str, str | bytes], rpc_timeout: timedelta | None, @@ -120,29 +107,17 @@ def publish( """Start a write. The write's task takes the call options.""" _check_loop() with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): - return _access(handle).publish(values) + return self._access().publish(values) def subscribe( self, - handle: StreamHandle[Any], - *, - offset: int | End | None, - last: int | None, - wait_new_records: bool, - create_wait: timedelta | None, + read: ReadSpec, rpc_metadata: Mapping[str, str | bytes], rpc_timeout: timedelta | None, ) -> Subscription[Any]: """Start a read.""" _check_loop() - read = read_spec( - handle._payload_type, - offset=offset, - last=last, - wait_new_records=wait_new_records, - create_wait=create_wait, - ) - access = _access(handle) + access = self._access() return subscription( lambda position: access.items(read, position), read, @@ -150,176 +125,102 @@ def subscribe( ) async def get_offset( - self, - handle: StreamHandle[Any], - rpc_metadata: Mapping[str, str | bytes], - rpc_timeout: timedelta | None, + self, rpc_metadata: Mapping[str, str | bytes], rpc_timeout: timedelta | None ) -> int: """Get the next offset.""" with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): - return await _access(handle).get_offset() + return await self._access().get_offset() async def close( self, - handle: StreamHandle[Any], result: Any, failure: BaseException | None, rpc_metadata: Mapping[str, str | bytes], rpc_timeout: timedelta | None, ) -> None: """Raise :py:class:`StreamNotOwnerError`, since an Activity can't close a Stream.""" - raise _not_owner(handle) + raise _not_owner(self._handle) - async def enter(self, handle: StreamHandle[Any]) -> None: + async def enter(self) -> None: """Raise before the block, since the close at its end would raise.""" - raise _not_owner(handle) + raise _not_owner(self._handle) - async def exit(self, handle: StreamHandle[Any], exc: BaseException | None) -> None: + async def exit(self, exc: BaseException | None) -> None: """Raise, since an Activity can't close a Stream.""" - raise _not_owner(handle) - - -class SyncSubscription(Generic[T]): - """The items of one read of a Stream in a sync Activity, in order. Iterate it with ``for``. - - Each iteration blocks the Activity's thread. See :py:class:`Subscription`. - - .. warning:: - This API is experimental. - """ - - def __init__( - self, handle: SyncStreamHandle[Any], subscription: Subscription[T] - ) -> None: - """Create a read. Use a handle's ``subscribe`` instead.""" - self._subscription = subscription - self._steps = handle._items(subscription) - - @property - def next_offset(self) -> int | None: - """The offset to resume from. See :py:attr:`Subscription.next_offset`.""" - return self._subscription.next_offset - - def __iter__(self) -> SyncSubscription[T]: - """The read is its own iterator.""" - return self + raise _not_owner(self._handle) - def __next__(self) -> Record[T] | Superseded | Closed: - """The next item.""" - return next(self._steps) + def producer(self, name: str, epoch: int) -> StreamCalls: + """Raise, since an Activity writes as its attempt.""" + raise _errors.StreamUnsupportedError( + "A Stream in an Activity writes as the attempt, not as a named Producer" + ) - def close(self) -> None: - """End the read early. It makes no more calls.""" - self._steps.close() +def attempt_of(stream: Stream[Any]) -> _activity.ActivityStreams | None: + """The Activity attempt that ``stream`` writes as, or None for a client's Stream.""" + calls = stream._calls + return calls.streams if isinstance(calls, ActivityCalls) else None -class SyncStreamHandle(Generic[T]): - """A Stream handle in a sync Activity. Its calls block the Activity's thread. - It wraps a :py:class:`StreamHandle`, such as a field of the Activity's input. Each call runs - the handle's call of the same name on the Worker's event loop, and blocks until it ends. A - publish returns the offset, not a future. +def run_blocking(call: Coroutine[Any, Any, _R]) -> _R: + """Run ``call`` on the Worker's event loop, and block the Activity's thread until it ends. - .. warning:: - This API is experimental. + Raises: + StreamConfigError: No Stream backend runs the Activity. """ - - def __init__(self, handle: StreamHandle[T]) -> None: - """Wrap ``handle`` for a sync Activity.""" - self._handle = handle - - def _run(self, call: Coroutine[Any, Any, _R]) -> _R: - loop = _bound(self._handle).loop + streams = _activity.current.get() + if streams is None: + call.close() + raise _errors.StreamConfigError( + "No Stream backend runs this Activity. Register one on the Worker's client " + "with stream_backends=." + ) + loop = streams.loop + try: + on_loop = asyncio.get_running_loop() is loop + except RuntimeError: + on_loop = False + if on_loop: + call.close() + raise RuntimeError( + "A SyncStream blocks, so use the Stream in an async Activity" + ) + # The future exists before the call is scheduled, so a cancel raised in this thread at + # any point after the schedule reaches the except below, which cancels the call. + future: concurrent.futures.Future[_R] = concurrent.futures.Future() + try: try: - on_loop = asyncio.get_running_loop() is loop + loop.call_soon_threadsafe(_start, loop, call, future) except RuntimeError: - on_loop = False - if on_loop: + # The Worker's loop is closed, such as at shutdown, so the call never starts. call.close() - raise RuntimeError( - "A SyncStreamHandle blocks, so use the StreamHandle in an async Activity" - ) - # The future exists before the call is scheduled, so a cancel raised in this thread at - # any point after the schedule reaches the except below, which cancels the call. - future: concurrent.futures.Future[_R] = concurrent.futures.Future() - try: - try: - loop.call_soon_threadsafe(_start, loop, call, future) - except RuntimeError: - # The Worker's loop is closed, such as at shutdown, so the call never starts. - call.close() - raise - while True: - # Short waits let the Worker raise an exception in this thread, such as a cancel. - try: - return future.result(timeout=_WAIT_STEP) - except concurrent.futures.TimeoutError: - continue - except BaseException: - # Cancel the call too, so a long poll doesn't keep running in the Worker's loop after - # the Activity thread is canceled, such as from a heartbeat. - future.cancel() raise - - def publish(self, *values: T) -> int: - """Write one Record for each value, in order, and return the last Record's offset. - - A list is one value. - - Raises: - StreamNotFoundError: Nobody created the Stream. - """ - - async def publish() -> int: - return await self._handle.publish(*values) - - return self._run(publish()) - - def subscribe( - self, - *, - offset: int | End | None = None, - last: int | None = None, - wait_new_records: bool = True, - create_wait: timedelta | None = None, - ) -> SyncSubscription[T]: - """Read the Stream's items in order. See :py:meth:`StreamHandle.subscribe`.""" - - async def subscribe() -> Subscription[T]: - return self._handle.subscribe( - offset=offset, - last=last, - wait_new_records=wait_new_records, - create_wait=create_wait, - ) - - return SyncSubscription(self, self._run(subscribe())) - - def _items( - self, items: Subscription[T] - ) -> Generator[Record[T] | Superseded | Closed, None, None]: - try: - while True: - try: - item = self._run(_next(items)) - except StopAsyncIteration: - return - yield item - finally: - self._run(items.aclose()) - - def get_offset(self) -> int: - """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. - - Offsets are sparse and above zero, so ``offset=0`` reads from the start. - """ - return self._run(self._handle.get_offset()) - - def close( - self, *, result: Any = None, failure: BaseException | None = None - ) -> None: - """Raise :py:class:`StreamNotOwnerError`, since an Activity can't close a Stream.""" - self._run(self._handle.close(result=result, failure=failure)) + while True: + # Short waits let the Worker raise an exception in this thread, such as a cancel. + try: + return future.result(timeout=_WAIT_STEP) + except concurrent.futures.TimeoutError: + continue + except BaseException: + # Cancel the call too, so a long poll doesn't keep running in the Worker's loop after + # the Activity thread is canceled, such as from a heartbeat. + future.cancel() + raise + + +def sync_items( + items: Subscription[T], +) -> Generator[Record[T] | Superseded | Closed, None, None]: + """The items of a read, each of which blocks the Activity's thread.""" + try: + while True: + try: + item = run_blocking(_next(items)) + except StopAsyncIteration: + return + yield item + finally: + run_blocking(items.aclose()) async def _next(items: Subscription[T]) -> Record[T] | Superseded | Closed: diff --git a/temporalio/common/_streams/_buffered.py b/temporalio/common/_streams/_buffered.py deleted file mode 100644 index 845c26caa..000000000 --- a/temporalio/common/_streams/_buffered.py +++ /dev/null @@ -1,185 +0,0 @@ -"""A publisher that batches the Records of one Stream.""" - -from __future__ import annotations - -import asyncio -from datetime import timedelta -from types import TracebackType -from typing import Generic, TypeVar - -import temporalio.workflow -from temporalio.common._streams import _activity, _errors -from temporalio.common._streams._handle import StreamHandle - -T = TypeVar("T") - -# Records in one batch of the crate's writer. A full batch is written at once. -_BATCH = 1000 - - -class BufferedPublisher(Generic[T]): - """Batches the publishes of one Stream. - - Use it in an Activity or a client that writes often. :py:meth:`publish` returns at once. - The publisher writes the buffered values as one batch at most ``max_delay`` after the first - of them, or at once when they fill a batch. One write is in flight at a time, and values - published during it wait for the next one. Use it in ``async with``, which writes the - remaining values at the end of the block. In an Activity, the end of the attempt also - writes the remaining values. - - ``await flush()`` is the backpressure point. A caller that publishes faster than the Store - accepts Records must await it from time to time, so the buffer doesn't grow without bound. - - After a batch fails, the publisher writes nothing more, so a reader sees no gap. Each later - call raises the error. - - .. warning:: - This API is experimental. - """ - - def __init__( - self, - handle: StreamHandle[T], - *, - max_delay: timedelta = timedelta(milliseconds=100), - ) -> None: - """Create a publisher. - - Args: - handle: The handle of the Stream. - max_delay: The longest time that a value waits in the buffer. - - Raises: - StreamUnsupportedError: The code is Workflow code. There, the publishes that wait - behind a write to the same Stream go as one Local Activity. - """ - if temporalio.workflow.in_workflow(): - raise _errors.StreamUnsupportedError( - "Workflow code can't use a BufferedPublisher. Its publishes that wait behind " - "a write already go as one Local Activity." - ) - self._handle = handle - self._max_delay = max_delay - self._buffer: list[T] = [] - # When the oldest buffered value arrived, so the next batch is at most max_delay later. - self._first_at = 0.0 - self._timer: asyncio.TimerHandle | None = None - self._writing: asyncio.Task[None] | None = None - self._published = 0 - # Count of the values that a write took, written or not. - self._settled = 0 - self._error: BaseException | None = None - self._last_offset = 0 - self._unwritten = 0 - self._reported = False - streams = _activity.current.get() - if streams is not None: - streams.publishers.append(self) - - def publish(self, *values: T) -> None: - """Add one Record for each value to the buffer, and return at once. - - A list is one value. - - Raises: - StreamError: A batch failed to write, so the publisher takes no more values. - """ - self._raise_error() - if not values: - raise TypeError("publish takes one or more values") - loop = asyncio.get_running_loop() - if not self._buffer: - self._first_at = loop.time() - self._buffer.extend(values) - self._published += len(values) - self._schedule() - - async def flush(self) -> int: - """Write every value published before the call, and return the last written offset. - - Returns 0 if the publisher hasn't written anything. - """ - target = self._published - while self._settled < target: - if self._writing is None: - self._start() - assert self._writing is not None - await asyncio.shield(self._writing) - self._raise_error() - return self._last_offset - - def _raise_error(self) -> None: - if self._error is not None: - self._reported = True - raise self._error - - async def _flush_at_attempt_end(self) -> None: - # Fail the attempt for unwritten values, unless the Activity already got the error and - # continued. - reported = self._reported - try: - await self.flush() - except Exception: - if not reported: - raise - - def _schedule(self) -> None: - if self._writing is not None or not self._buffer: - return - if len(self._buffer) >= _BATCH: - self._start() - elif self._timer is None: - self._timer = asyncio.get_running_loop().call_at( - self._first_at + self._max_delay.total_seconds(), self._start - ) - - def _start(self) -> None: - if self._timer is not None: - self._timer.cancel() - self._timer = None - if self._writing is not None or not self._buffer: - return - batch, self._buffer = self._buffer, [] - self._writing = asyncio.get_running_loop().create_task(self._write(batch)) - - async def _write(self, batch: list[T]) -> None: - try: - if self._error is None: - self._last_offset = await self._send(batch) - else: - # A failed batch fails the later ones too, so a reader sees no gap. - self._unwritten += len(batch) - except Exception as err: - self._unwritten += len(batch) - self._error = err - finally: - self._settled += len(batch) - self._writing = None - # Schedule the values published during the write, for when their delay ends or they - # fill a batch. - self._schedule() - - async def _send(self, batch: list[T]) -> int: - return await self._handle.publish(*batch) - - @property - def unwritten(self) -> int: - """How many published values a failed batch left unwritten.""" - return self._unwritten - - async def __aenter__(self) -> BufferedPublisher[T]: - """Use the publisher in ``async with``.""" - return self - - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: - """Write the buffered values. A write error does not hide an error of the block.""" - try: - await self.flush() - except Exception: - if exc is None: - raise diff --git a/temporalio/common/_streams/_extension.py b/temporalio/common/_streams/_extension.py index 0fb2fe734..3898f0e88 100644 --- a/temporalio/common/_streams/_extension.py +++ b/temporalio/common/_streams/_extension.py @@ -172,6 +172,7 @@ def _configure_worker( # Records get no serialization context, so Workflows, Activities and clients encode # them alike. mine.converter = client.data_converter + mine.client = client runner = config.get("workflow_runner") if isinstance(runner, SandboxedWorkflowRunner): # The interceptor and temporalio.workflow must share the run's Stream state, even @@ -250,6 +251,7 @@ def __init__( self.plugin = plugin self.backend = backend self.converter = converter or temporalio.converter.DataConverter.default + self.client: temporalio.client.Client | None = None def workflow_interceptor_class( self, input: temporalio.worker.WorkflowInterceptorClassInput @@ -603,4 +605,5 @@ def _activity_streams( workflow_id=info.workflow_id, previous_completed_task=previous, failed_tasks=failed, + client=interceptor.client, ) diff --git a/temporalio/common/_streams/_stream_client.py b/temporalio/common/_streams/_stream_client.py index 44c3a473c..a8e02cdc5 100644 --- a/temporalio/common/_streams/_stream_client.py +++ b/temporalio/common/_streams/_stream_client.py @@ -16,7 +16,7 @@ import temporalio.client from temporalio.api.enums.v1 import StreamProducerKind from temporalio.client._streams import Stream -from temporalio.common._streams import _backend, _errors, _ids +from temporalio.common._streams import _activity, _backend, _errors, _ids from temporalio.common._streams._access import Access, ReadSpec, subscription from temporalio.common._streams._backend import BackendCalls from temporalio.common._streams._defn import StreamDefinition, definition_of @@ -334,4 +334,12 @@ def _stream( handle: StreamHandle[T], payload_type: Any, ) -> Stream[T]: + streams = _activity.current.get() + if streams is not None and streams.client is client: + # The Activity module imports this one. + from temporalio.common._streams._activity_stream import ActivityCalls + + # A Stream of the Worker's client writes as the attempt, so a retry withdraws the + # Records of the failed attempt. + return Stream(handle, payload_type, ActivityCalls(streams, handle)) return Stream(handle, payload_type, ClientCalls(client, backend, handle)) diff --git a/tests/contrib/streams/test_activity_handles.py b/tests/contrib/streams/test_activity_handles.py index 495988c7e..3fdb9bb2e 100644 --- a/tests/contrib/streams/test_activity_handles.py +++ b/tests/contrib/streams/test_activity_handles.py @@ -18,9 +18,15 @@ from temporalio import activity, common, workflow from temporalio.api.enums.v1 import StreamErrorReason -from temporalio.client import Client, WorkflowHandle +from temporalio.client import ( + BufferedPublisher, + Client, + Stream, + SyncStream, + WorkflowHandle, +) from temporalio.common import RetryPolicy -from temporalio.common._streams import _activity +from temporalio.common._streams import _activity, _activity_stream from temporalio.exceptions import ActivityError, ApplicationError, CancelledError from tests.contrib.streams._support import workflow_stream from tests.helpers import new_worker @@ -60,27 +66,50 @@ def records(items: Sequence[object]) -> list[common.Record[Any]]: @activity.defn async def call_model(input: CallModelInput) -> list[int]: - assert input.stream._payload_type is str - offsets = [await input.stream.publish(word) for word in input.prompt.split()] - async with common.BufferedPublisher( - input.stream, max_delay=timedelta(seconds=10) - ) as buffered: + stream = activity.client().to_stream(input.stream) + offsets = [await stream.publish(word) for word in input.prompt.split()] + async with BufferedPublisher(stream, max_delay=timedelta(seconds=10)) as buffered: buffered.publish("model-a") buffered.publish("model-b") - offsets.append(await input.stream.get_offset()) + offsets.append(await stream.get_offset()) return offsets -@activity.defn -def call_model_sync(input: SyncInput) -> int: - stream = common.SyncStreamHandle(input.stream) - last = stream.publish(*input.prompt.split()) - assert stream.get_offset() == last + 1 - return last +class SyncActivities: + # A sync Activity has no activity.client(), so it keeps the Worker's client. + def __init__(self, client: Client) -> None: + self._client = client + + @activity.defn + def call_model_sync(self, input: SyncInput) -> int: + stream = SyncStream(self._client.to_stream(input.stream)) + last = stream.publish(*input.prompt.split()) + assert stream.get_offset() == last + 1 + return last + + @activity.defn + def publishes_without_the_sync_stream( + self, handle: common.StreamHandle[str] + ) -> str: + try: + self._client.to_stream(handle).publish("x") + except common.StreamUnsupportedError as err: + return type(err).__name__ + return "wrote" + + @activity.defn + def sync_read(self, handle: common.StreamHandle[str]) -> list[Any]: + stream = SyncStream(self._client.to_stream(handle)) + return [ + item.data + for item in stream.subscribe(wait_new_records=False) + if isinstance(item, common.Record) + ] @activity.defn -async def read_back(stream: common.StreamHandle[str]) -> list[Any]: +async def read_back(handle: common.StreamHandle[str]) -> list[Any]: + stream = activity.client().to_stream(handle) return [ item.data async for item in stream.subscribe(wait_new_records=False) @@ -95,17 +124,17 @@ async def run(self, prompt: str) -> list[Any]: stream = await workflow.create_stream(Tokens) offsets = await workflow.execute_activity( call_model, - CallModelInput(prompt, stream), + CallModelInput(prompt, stream.handle), start_to_close_timeout=timedelta(seconds=30), ) - last = await workflow.execute_activity( - call_model_sync, - SyncInput("sync words", stream), + last = await workflow.execute_activity_method( + SyncActivities.call_model_sync, + SyncInput("sync words", stream.handle), start_to_close_timeout=timedelta(seconds=30), ) read = await workflow.execute_activity( read_back, - stream, + stream.handle, start_to_close_timeout=timedelta(seconds=30), ) await stream.close() @@ -116,7 +145,11 @@ async def test_activity_publishes_and_returns_offsets(stream_client: Client): async with new_worker( stream_client, Agent, - activities=[call_model, call_model_sync, read_back], + activities=[ + call_model, + SyncActivities(stream_client).call_model_sync, + read_back, + ], activity_executor=_executor(), ) as worker: handle = await stream_client.start_workflow( @@ -147,7 +180,7 @@ def _executor() -> Any: async def flaky(input: CallModelInput) -> None: attempts[input.prompt] = attempts.get(input.prompt, 0) + 1 if attempts[input.prompt] == 1: - await input.stream.publish("attempt-1") + await activity.client().to_stream(input.stream).publish("attempt-1") raise RuntimeError("first attempt fails") @@ -159,7 +192,7 @@ async def run(self, key: str) -> None: await stream.publish("start") await workflow.execute_activity( flaky, - CallModelInput(key, stream), + CallModelInput(key, stream.handle), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), ) @@ -192,7 +225,7 @@ async def reads_after_a_failed_attempt(input: FinishedStreams) -> list[Any]: attempts[input.key] = attempts.get(input.key, 0) + 1 if attempts[input.key] == 1: raise RuntimeError("first attempt fails") - items = input.closed.subscribe(wait_new_records=False) + items = activity.client().to_stream(input.closed).subscribe(wait_new_records=False) return [item.data async for item in items if isinstance(item, common.Record)] @@ -207,7 +240,7 @@ async def run(self, key: str, expired: common.StreamHandle[str]) -> list[Any]: reads_after_a_failed_attempt, FinishedStreams( key, - stream, + stream.handle, expired, ), start_to_close_timeout=timedelta(seconds=30), @@ -229,7 +262,7 @@ async def test_retry_goes_on_with_a_closed_or_expired_stream(stream_client: Clie ) as worker: result = await stream_client.execute_workflow( ReadsFinishedStreams.run, - args=[key, expired], + args=[key, expired.handle], id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) @@ -243,7 +276,7 @@ async def test_retry_goes_on_with_a_closed_or_expired_stream(stream_client: Clie @activity.defn async def fails_after_it_is_shown(input: CallModelInput) -> None: attempts[input.prompt] = attempts.get(input.prompt, 0) + 1 - tokens = input.stream + tokens = activity.client().to_stream(input.stream) if attempts[input.prompt] == 1: await tokens.publish("attempt-1") await shown[input.prompt].wait() @@ -258,7 +291,7 @@ async def run(self, key: str) -> None: stream = await workflow.create_stream(Tokens) await workflow.execute_activity( fails_after_it_is_shown, - CallModelInput(key, stream), + CallModelInput(key, stream.handle), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), ) @@ -305,16 +338,17 @@ async def test_reader_withdraws_a_shown_record_after_a_reconnect( @activity.defn -async def wrong_backend(stream: common.StreamHandle[str]) -> str: +async def wrong_backend(handle: common.StreamHandle[str]) -> str: try: - await stream.publish("x") + await activity.client().to_stream(handle).publish("x") except common.StreamConfigError: return "config" return "wrote" @activity.defn -async def close_it(stream: common.StreamHandle[str]) -> str: +async def close_it(handle: common.StreamHandle[str]) -> str: + stream = activity.client().to_stream(handle) await stream.publish("last") try: await stream.close(result="closed by activity") @@ -330,24 +364,24 @@ async def run(self) -> list[str]: stream = await workflow.create_stream(Tokens) backend = await workflow.execute_activity( wrong_backend, - common.StreamHandle[str](id=stream.id, backend="elsewhere"), + common.StreamHandle[str](id=stream.handle.id, backend="elsewhere"), start_to_close_timeout=timedelta(seconds=30), ) owner = await workflow.execute_activity( close_it, - stream, + stream.handle, start_to_close_timeout=timedelta(seconds=30), ) # A Local Activity runs before its Workflow Task commits, so it can't close it either. local = await workflow.execute_local_activity( close_it, - stream, + stream.handle, start_to_close_timeout=timedelta(seconds=30), ) # A handle with only the id names no Workflow, and no Activity can close a Stream. by_id = await workflow.execute_local_activity( close_it, - common.StreamHandle[str](id=stream.id), + common.StreamHandle[str](id=stream.handle.id), start_to_close_timeout=timedelta(seconds=30), ) await stream.close(result="closed by the run") @@ -373,31 +407,19 @@ async def test_activity_checks_the_backend_and_the_owner(stream_client: Client): @activity.defn -async def refused_calls(stream: common.StreamHandle[str]) -> list[str]: +async def refused_calls(handle: common.StreamHandle[str]) -> list[str]: + stream = activity.client().to_stream(handle) refused = [] try: async with stream: refused.append("entered") except common.StreamNotOwnerError as err: refused.append(type(err).__name__) - for call in ( - lambda: stream.producer("ingest", 1), - lambda: stream.publish_nowait("x"), - ): - try: - call() - except common.StreamUnsupportedError as err: - refused.append(type(err).__name__) - return refused - - -@activity.defn -def publishes_without_the_sync_handle(stream: common.StreamHandle[str]) -> str: try: - stream.publish("x") + stream.producer("ingest", 1) except common.StreamUnsupportedError as err: - return type(err).__name__ - return "wrote" + refused.append(type(err).__name__) + return refused @workflow.defn @@ -406,17 +428,18 @@ class RefusedCalls: async def run(self) -> list[str]: stream = await workflow.create_stream(Tokens) refused = await workflow.execute_activity( - refused_calls, stream, start_to_close_timeout=timedelta(seconds=30) + refused_calls, stream.handle, start_to_close_timeout=timedelta(seconds=30) ) refused.append( - await workflow.execute_activity( - publishes_without_the_sync_handle, - stream, + await workflow.execute_activity_method( + SyncActivities.publishes_without_the_sync_stream, + stream.handle, start_to_close_timeout=timedelta(seconds=30), ) ) try: - common.BufferedPublisher(stream) + # Workflow code has no Stream, so this is for a caller that has no type checks. + BufferedPublisher(cast(Any, stream)) except common.StreamUnsupportedError as err: refused.append(type(err).__name__) await stream.close() @@ -427,7 +450,10 @@ async def test_calls_that_an_activity_cant_make_raise(stream_client: Client): async with new_worker( stream_client, RefusedCalls, - activities=[refused_calls, publishes_without_the_sync_handle], + activities=[ + refused_calls, + SyncActivities(stream_client).publishes_without_the_sync_stream, + ], activity_executor=_executor(), ) as worker: handle = await stream_client.start_workflow( @@ -435,21 +461,22 @@ async def test_calls_that_an_activity_cant_make_raise(stream_client: Client): ) assert await handle.result() == [ "StreamNotOwnerError", - *["StreamUnsupportedError"] * 4, + *["StreamUnsupportedError"] * 3, ] items = await read(handle) assert records(items) == [] @activity.defn -async def leaves_a_failed_write(stream: common.StreamHandle[str]) -> str: +async def leaves_a_failed_write(handle: common.StreamHandle[str]) -> str: # Nobody created the Stream, and no caller waits for the write. - stream.publish("lost") + activity.client().to_stream(handle).publish("lost") return "returned" @activity.defn -async def enters_a_client_stream(stream: common.StreamHandle[str]) -> str: +async def enters_a_client_stream(handle: common.StreamHandle[str]) -> str: + stream = activity.client().to_stream(handle) try: async with stream: await stream.publish("in the block") @@ -493,7 +520,7 @@ async def test_attempt_end_raises_a_lost_write_and_a_client_stream_refuses_the_b ) as worker: result = await stream_client.execute_workflow( LostAndClientWrites.run, - args=[f"missing-{uuid.uuid4()}", stream], + args=[f"missing-{uuid.uuid4()}", stream.handle], id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) @@ -503,10 +530,10 @@ async def test_attempt_end_raises_a_lost_write_and_a_client_stream_refuses_the_b assert records(items) == [] -class _Batches(common.StreamHandle[int]): +class _Batches(Stream[int]): # A Stream whose writes go to a list, so the batches show with no backend. def __init__(self, fail: bool = False) -> None: - super().__init__(id="batches") + super().__init__(common.StreamHandle(id="batches"), None, cast(Any, None)) self.batches: list[tuple[int, ...]] = [] self.fail = fail self.in_flight = 0 @@ -528,7 +555,7 @@ async def _write(self, values: Sequence[int]) -> int: async def test_buffered_publisher_batches_by_delay(): stream = _Batches() - publisher = common.BufferedPublisher(stream, max_delay=timedelta(milliseconds=20)) + publisher = BufferedPublisher(stream, max_delay=timedelta(milliseconds=20)) publisher.publish(1) publisher.publish(2, 3) await asyncio.sleep(0.1) @@ -541,7 +568,7 @@ async def test_buffered_publisher_batches_by_delay(): async def test_buffered_publisher_writes_a_full_batch_with_one_write_in_flight(): stream = _Batches() - publisher = common.BufferedPublisher(stream, max_delay=timedelta(minutes=10)) + publisher = BufferedPublisher(stream, max_delay=timedelta(minutes=10)) publisher.publish(*range(1000)) publisher.publish(*range(1000, 2500)) await asyncio.sleep(0.2) @@ -555,16 +582,17 @@ async def test_buffered_publisher_writes_a_full_batch_with_one_write_in_flight() async def test_buffered_publisher_raises_a_failed_batch(): stream = _Batches(fail=True) with pytest.raises(common.StreamStorageError): - async with common.BufferedPublisher(stream) as publisher: + async with BufferedPublisher(stream) as publisher: publisher.publish(1) with pytest.raises(TypeError): - common.BufferedPublisher(stream).publish() + BufferedPublisher(stream).publish() @activity.defn -async def local_step(stream: common.StreamHandle[str]) -> list[Any]: +async def local_step(handle: common.StreamHandle[str]) -> list[Any]: # Each Workflow write is an awaited Local Activity, so the run's create and Records are in # the Store before a later Local Activity starts. + stream = activity.client().to_stream(handle) seen = [ item.data async for item in stream.subscribe(wait_new_records=False) @@ -582,7 +610,7 @@ async def run(self) -> list[Any]: await stream.publish("before") seen = await workflow.execute_local_activity( local_step, - stream, + stream.handle, start_to_close_timeout=timedelta(seconds=30), ) await stream.publish("after") @@ -590,6 +618,51 @@ async def run(self) -> list[Any]: return seen +class TwoClients: + def __init__(self, other: Client) -> None: + self._other = other + + @activity.defn + async def writes_with_both(self, handle: common.StreamHandle[str]) -> None: + await activity.client().to_stream(handle).publish("worker client") + await self._other.to_stream(handle).publish("other client") + + +@workflow.defn +class WritesWithTwoClients: + @workflow.run + async def run(self) -> None: + stream = await workflow.create_stream(Tokens) + await workflow.execute_activity_method( + TwoClients.writes_with_both, + stream.handle, + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.close() + + +async def test_only_the_workers_client_writes_as_the_attempt(stream_client: Client): + # Another client in the Activity writes as that client, so a retry doesn't withdraw it. + other = Client(**stream_client.config()) + async with new_worker( + stream_client, + WritesWithTwoClients, + activities=[TwoClients(other).writes_with_both], + ) as worker: + handle = await stream_client.start_workflow( + WritesWithTwoClients.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + items = await read(handle) + kinds = {r.data: r.producer.kind for r in records(items) if r.producer} + assert kinds == { + "worker client": common.StreamProducerKind.StreamProducerKindActivity, + "other client": common.StreamProducerKind.StreamProducerKindClient, + } + + async def test_workflow_writes_go_out_before_a_local_activity(stream_client: Client): async with new_worker(stream_client, LocalOrder, activities=[local_step]) as worker: handle = await stream_client.start_workflow( @@ -606,7 +679,7 @@ async def test_workflow_writes_go_out_before_a_local_activity(stream_client: Cli async def test_buffered_publisher_stays_failed(): stream = _Batches(fail=True) - publisher = common.BufferedPublisher(stream) + publisher = BufferedPublisher(stream) publisher.publish(1, 2) with pytest.raises(common.StreamStorageError): await publisher.flush() @@ -621,7 +694,8 @@ async def test_buffered_publisher_stays_failed(): @activity.defn async def never_flushes(input: CallModelInput) -> None: - publisher = common.BufferedPublisher(input.stream, max_delay=timedelta(minutes=10)) + stream = activity.client().to_stream(input.stream) + publisher = BufferedPublisher(stream, max_delay=timedelta(minutes=10)) publisher.publish(*input.prompt.split()) @@ -637,20 +711,11 @@ async def flaky_many(input: ManyStreams) -> None: attempts[input.key] = attempts.get(input.key, 0) + 1 assert input.maybe is not None if attempts[input.key] == 1: - for stream in [*input.targets, input.maybe]: - await stream.publish("attempt-1") + for handle in [*input.targets, input.maybe]: + await activity.client().to_stream(handle).publish("attempt-1") raise RuntimeError("first attempt fails") -@activity.defn -def sync_read(stream: common.StreamHandle[str]) -> list[Any]: - return [ - item.data - for item in common.SyncStreamHandle(stream).subscribe(wait_new_records=False) - if isinstance(item, common.Record) - ] - - @workflow.defn class Leftovers: @workflow.run @@ -659,22 +724,22 @@ async def run(self, key: str) -> list[Any]: other = await workflow.create_stream(Other) await workflow.execute_activity( never_flushes, - CallModelInput("a b", stream), + CallModelInput("a b", stream.handle), start_to_close_timeout=timedelta(seconds=30), ) await workflow.execute_activity( flaky_many, ManyStreams( key, - [stream], - other, + [stream.handle], + other.handle, ), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), ) - return await workflow.execute_activity( - sync_read, - stream, + return await workflow.execute_activity_method( + SyncActivities.sync_read, + stream.handle, start_to_close_timeout=timedelta(seconds=30), ) @@ -686,7 +751,7 @@ async def test_attempt_end_flushes_and_retry_finds_nested_streams( async with new_worker( stream_client, Leftovers, - activities=[never_flushes, flaky_many, sync_read], + activities=[never_flushes, flaky_many, SyncActivities(stream_client).sync_read], activity_executor=_executor(), ) as worker: handle = await stream_client.start_workflow( @@ -735,14 +800,11 @@ def _caller( call: Coroutine[Any, Any, None], errors: list[BaseException], ) -> threading.Thread: - stream = common.SyncStreamHandle(common.StreamHandle[Any](id="s")) - token = _activity.current.set( - cast(Any, SimpleNamespace(loop=loop, backend_name="", default=True)) - ) + token = _activity.current.set(cast(Any, SimpleNamespace(loop=loop))) def run() -> None: try: - stream._run(call) + _activity_stream.run_blocking(call) except BaseException as err: errors.append(err) @@ -775,8 +837,8 @@ def test_sync_call_stops_when_the_thread_is_cancelled(): try: caller.start() assert started.wait(5) - # Raise only while the caller waits on the call's future, so the test checks _run and - # not where the scheduler happens to stop the caller thread. + # Raise only while the caller waits on the call's future, so the test checks + # run_blocking and not where the scheduler happens to stop the caller thread. deadline = time.monotonic() + 5 while not _waits_for_the_call(caller) and time.monotonic() < deadline: time.sleep(0.001) @@ -847,23 +909,23 @@ async def run(self) -> str: @workflow.defn class ReceivesAStream: @workflow.run - async def run(self, stream: common.StreamHandle[str]) -> str: + async def run(self, handle: common.StreamHandle[str]) -> str: # The Activity gets no Stream, but the run received one. return await workflow.execute_activity( - plain, stream.id, start_to_close_timeout=timedelta(seconds=10) + plain, handle.id, start_to_close_timeout=timedelta(seconds=10) ) @workflow.defn class RetriesAfterContinueAsNew: @workflow.run - async def run(self, key: str, stream: common.StreamHandle[str] | None) -> None: - if stream is None: + async def run(self, key: str, handle: common.StreamHandle[str] | None) -> None: + if handle is None: created = await workflow.create_stream(Tokens) - workflow.continue_as_new(args=[key, created]) + workflow.continue_as_new(args=[key, created.handle]) await workflow.execute_activity( flaky, - CallModelInput(key, stream), + CallModelInput(key, handle), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), ) @@ -908,7 +970,7 @@ async def test_only_a_run_that_uses_streams_keys_its_activities(stream_client: C assert await keyed(creates, "flaky") == [True] receives = await stream_client.start_workflow( ReceivesAStream.run, - workflow_stream(creates, Tokens), + workflow_stream(creates, Tokens).handle, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) diff --git a/tests/contrib/streams/test_pydantic.py b/tests/contrib/streams/test_pydantic.py index cbcdd2677..c10f27be5 100644 --- a/tests/contrib/streams/test_pydantic.py +++ b/tests/contrib/streams/test_pydantic.py @@ -26,8 +26,8 @@ class TurnResult(pydantic.BaseModel): @activity.defn async def call_model(input: CallModelInput) -> int: - assert input.stream._payload_type is str - return await input.stream.publish(*input.prompt.split()) + stream = activity.client().to_stream(input.stream) + return await stream.publish(*input.prompt.split()) @workflow.defn @@ -37,11 +37,11 @@ async def run(self, prompt: str) -> TurnResult: stream = await workflow.create_stream(Tokens) await workflow.execute_activity( call_model, - CallModelInput(prompt=prompt, stream=stream), + CallModelInput(prompt=prompt, stream=stream.handle), start_to_close_timeout=timedelta(seconds=30), ) await stream.close(result=prompt.upper()) - return TurnResult(text=prompt.upper(), stream=stream) + return TurnResult(text=prompt.upper(), stream=stream.handle) async def test_handles_in_models_round_trip_on_the_memory_store(client: Client): @@ -58,10 +58,11 @@ async def test_handles_in_models_round_trip_on_the_memory_store(client: Client): ) result = await handle.result() assert result.text == "HI THERE" - assert result.stream.workflow_id == handle.id - assert result.stream.first_execution_run_id == handle.first_execution_run_id + owner = result.stream.owner + assert owner is not None and owner.business_id == handle.id + assert owner.run_id == handle.first_execution_run_id assert result.stream.backend == "memory" - stream = client.get_stream_handle(Tokens, id=result.stream.id) + stream = client.to_stream(result.stream) items = [item async for item in stream.subscribe()] assert [i.data for i in items if isinstance(i, common.Record)] == ["hi", "there"] assert isinstance(items[-1], common.Closed) and items[-1].result is not None diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index 4ffe7e8d8..b23c5e165 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -75,7 +75,7 @@ async def test_failed_workflow_task_is_withdrawn(stream_client: Client): @activity.defn async def local_tokens(input: CallModelInput) -> None: - await input.stream.publish(*input.prompt.split()) + await activity.client().to_stream(input.stream).publish(*input.prompt.split()) @workflow.defn @@ -86,7 +86,7 @@ async def run(self, key: str) -> None: first = attempts.setdefault(key, 0) == 0 await workflow.execute_local_activity( local_tokens, - CallModelInput("0 1" if first else "2 3", stream), + CallModelInput("0 1" if first else "2 3", stream.handle), start_to_close_timeout=timedelta(seconds=30), ) if first: @@ -133,7 +133,7 @@ async def run(self, key: str) -> None: attempts[key] = 1 await workflow.execute_local_activity( local_tokens, - CallModelInput("0 1", stream), + CallModelInput("0 1", stream.handle), start_to_close_timeout=timedelta(seconds=30), ) raise RuntimeError("fail this Workflow Task once") diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index e5c9bcf64..abab2bb3f 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -9,7 +9,7 @@ from temporalio import common, workflow from temporalio.api.enums.v1 import StreamErrorReason -from temporalio.client import Client +from temporalio.client import BufferedPublisher, Client from temporalio.common._streams._handle import workflow_owner from temporalio.contrib.stream.testing import MemoryStreamBackend from temporalio.converter import DataConverter, DefaultFailureConverter @@ -360,7 +360,7 @@ async def read(definition: type[common.StreamDefinition[Any]]) -> list[object]: async def test_buffered_publisher_keeps_the_publish_order(stream_client: Client): stream = await stream_client.create_stream(Turns, id=f"buffered-{uuid.uuid4()}") - async with common.BufferedPublisher(stream) as publisher: + async with BufferedPublisher(stream) as publisher: publisher.publish(Turn("user")) publisher.publish(Turn("a"), Turn("b")) publisher.publish(Turn("model"))