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/_streams/_activity.py b/temporalio/common/_streams/_activity.py index 36779e3cf..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.""" @@ -50,8 +53,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 +68,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 new file mode 100644 index 000000000..2405cc1cd --- /dev/null +++ b/temporalio/common/_streams/_activity_stream.py @@ -0,0 +1,227 @@ +"""Streams in Activity code, which write and read as the attempt.""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +from collections.abc import Coroutine, Generator, Mapping +from datetime import timedelta +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, 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") + +# 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 _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." + ) + + +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 Stream in a SyncStream." + ) from None + + +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 + 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. + """ + + 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, + 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 self._access().publish(values) + + def subscribe( + self, + read: ReadSpec, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> Subscription[Any]: + """Start a read.""" + _check_loop() + access = self._access() + return subscription( + lambda position: access.items(read, position), + read, + _backend.options_of(rpc_metadata, rpc_timeout), + ) + + async def get_offset( + 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 self._access().get_offset() + + async def close( + self, + 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(self._handle) + + async def enter(self) -> None: + """Raise before the block, since the close at its end would raise.""" + raise _not_owner(self._handle) + + async def exit(self, exc: BaseException | None) -> None: + """Raise, since an Activity can't close a Stream.""" + raise _not_owner(self._handle) + + 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 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 + + +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. + + Raises: + StreamConfigError: No Stream backend runs the Activity. + """ + 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: + 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 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: + return await items.__anext__() diff --git a/temporalio/common/_streams/_extension.py b/temporalio/common/_streams/_extension.py index 37f03d922..f460b9363 100644 --- a/temporalio/common/_streams/_extension.py +++ b/temporalio/common/_streams/_extension.py @@ -174,6 +174,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 @@ -252,6 +253,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 @@ -605,4 +607,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 new file mode 100644 index 000000000..3fdb9bb2e --- /dev/null +++ b/tests/contrib/streams/test_activity_handles.py @@ -0,0 +1,995 @@ +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 ( + BufferedPublisher, + Client, + Stream, + SyncStream, + WorkflowHandle, +) +from temporalio.common import RetryPolicy +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 + + +@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.StreamHandle[str] + + +@dataclass +class SyncInput: + prompt: str + stream: common.StreamHandle[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]: + 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 stream.get_offset()) + return offsets + + +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(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) + 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, stream.handle), + start_to_close_timeout=timedelta(seconds=30), + ) + 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.handle, + 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, + SyncActivities(stream_client).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 activity.client().to_stream(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, stream.handle), + 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.StreamHandle[str] + expired: common.StreamHandle[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 = activity.client().to_stream(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.StreamHandle[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, + stream.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, expired.handle], + 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 = activity.client().to_stream(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, stream.handle), + 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(handle: common.StreamHandle[str]) -> str: + try: + await activity.client().to_stream(handle).publish("x") + except common.StreamConfigError: + return "config" + return "wrote" + + +@activity.defn +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") + except common.StreamNotOwnerError as err: + return StreamErrorReason.Name(StreamErrorReason.ValueType(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.StreamHandle[str](id=stream.handle.id, backend="elsewhere"), + start_to_close_timeout=timedelta(seconds=30), + ) + owner = await workflow.execute_activity( + close_it, + 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.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.handle.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" + + +@activity.defn +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__) + try: + stream.producer("ingest", 1) + except common.StreamUnsupportedError as err: + refused.append(type(err).__name__) + return refused + + +@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.handle, start_to_close_timeout=timedelta(seconds=30) + ) + refused.append( + await workflow.execute_activity_method( + SyncActivities.publishes_without_the_sync_stream, + stream.handle, + start_to_close_timeout=timedelta(seconds=30), + ) + ) + try: + # 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() + 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, + SyncActivities(stream_client).publishes_without_the_sync_stream, + ], + 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"] * 3, + ] + items = await read(handle) + assert records(items) == [] + + +@activity.defn +async def leaves_a_failed_write(handle: common.StreamHandle[str]) -> str: + # Nobody created the Stream, and no caller waits for the write. + activity.client().to_stream(handle).publish("lost") + return "returned" + + +@activity.defn +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") + 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.handle], + 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(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__(common.StreamHandle(id="batches"), None, cast(Any, None)) + self.batches: list[tuple[int, ...]] = [] + self.fail = fail + self.in_flight = 0 + self.most_in_flight = 0 + + 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: + 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 = 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 = 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 BufferedPublisher(stream) as publisher: + publisher.publish(1) + with pytest.raises(TypeError): + BufferedPublisher(stream).publish() + + +@activity.defn +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) + 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, + stream.handle, + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.publish("after") + await stream.close() + 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( + 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 = 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: + stream = activity.client().to_stream(input.stream) + publisher = BufferedPublisher(stream, max_delay=timedelta(minutes=10)) + publisher.publish(*input.prompt.split()) + + +@dataclass +class ManyStreams: + key: str + targets: list[common.StreamHandle[str]] + maybe: common.StreamHandle[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 handle in [*input.targets, input.maybe]: + await activity.client().to_stream(handle).publish("attempt-1") + raise RuntimeError("first attempt fails") + + +@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", stream.handle), + start_to_close_timeout=timedelta(seconds=30), + ) + await workflow.execute_activity( + flaky_many, + ManyStreams( + key, + [stream.handle], + other.handle, + ), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + return await workflow.execute_activity_method( + SyncActivities.sync_read, + stream.handle, + 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, SyncActivities(stream_client).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: + token = _activity.current.set(cast(Any, SimpleNamespace(loop=loop))) + + def run() -> None: + try: + _activity_stream.run_blocking(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_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) + 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, handle: common.StreamHandle[str]) -> str: + # The Activity gets no Stream, but the run received one. + return await workflow.execute_activity( + plain, handle.id, start_to_close_timeout=timedelta(seconds=10) + ) + + +@workflow.defn +class RetriesAfterContinueAsNew: + @workflow.run + 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.handle]) + await workflow.execute_activity( + flaky, + CallModelInput(key, handle), + 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, + workflow_stream(creates, Tokens).handle, + 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..c10f27be5 --- /dev/null +++ b/tests/contrib/streams/test_pydantic.py @@ -0,0 +1,68 @@ +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.StreamHandle[str] + + +class TurnResult(pydantic.BaseModel): + text: str + stream: common.StreamHandle[str] + + +@activity.defn +async def call_model(input: CallModelInput) -> int: + stream = activity.client().to_stream(input.stream) + return await 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=stream.handle), + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.close(result=prompt.upper()) + return TurnResult(text=prompt.upper(), stream=stream.handle) + + +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" + 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.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 71d438cc0..b23c5e165 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, Stream from temporalio.worker import ( UnsandboxedWorkflowRunner, @@ -15,6 +16,12 @@ class Tokens(common.StreamDefinition[str]): ... +@dataclass +class CallModelInput: + prompt: str + stream: common.StreamHandle[str] + + async def read_all(stream: Stream[Any]) -> list[Any]: return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] @@ -64,3 +71,95 @@ 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 activity.client().to_stream(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", stream.handle), + 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", stream.handle), + 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 20815c438..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 @@ -358,6 +358,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 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]): ...