diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index 0721be43d..7237ec3cd 100644 --- a/temporalio/contrib/streams/__init__.py +++ b/temporalio/contrib/streams/__init__.py @@ -1,16 +1,24 @@ """Streams of Records that Workflows, Activities and clients write and read. -Workflow code writes with :py:func:`workflow_stream` and :py:func:`create_workflow_stream`. -Activities use :py:func:`activity_stream`. Clients read with :py:func:`for_handle` or :py:func:`for_ref` and create -standalone Streams with :py:func:`create_stream`. A Stream plugin, such as -:py:class:`temporalio.contrib.streams.redis.RedisStreamPlugin`, carries the calls. +A Stream definition, a class with ``@streams.defn``, gives the Stream name and the type of each +Topic. Workflow code creates a Stream with +:py:func:`temporalio.contrib.streams.workflow.create_stream`. A client creates and gets Streams +with :py:class:`StreamClient`. A Stream is data: a Workflow passes it to an Activity, a child +Workflow or the next run as an argument, and the parameter's type gives the handle class of +that context, such as :py:class:`ActivityStream`. + +A Stream plugin carries the calls and registers a backend name with ``name=``, such as +:py:class:`temporalio.contrib.streams.redis.RedisStreamPlugin` (``redis``) and +:py:class:`temporalio.contrib.streams.testing.MemoryStreamPlugin` (``memory``). A handle +carries the name of its backend. A client may have two plugins, and ``default=True`` marks the +one that a create with no ``backend=`` uses. For now a Worker has one Stream plugin, so it +serves one backend. .. warning:: This package is experimental. """ from temporalio.contrib.streams import workflow -from temporalio.contrib.streams._activity import activity_stream from temporalio.contrib.streams._activity_stream import ( ActivityStream, ActivityTopic, @@ -19,11 +27,6 @@ ) from temporalio.contrib.streams._backend import StreamBackend from temporalio.contrib.streams._buffered import BufferedPublisher -from temporalio.contrib.streams._client import ( - create_stream, - for_handle, - for_ref, -) from temporalio.contrib.streams._defn import ( StreamDefinition, TopicDefinition, @@ -45,18 +48,12 @@ ) from temporalio.contrib.streams._handle import StreamHandle from temporalio.contrib.streams._items import Closed, Record, Superseded -from temporalio.contrib.streams._outside import OutsideStream -from temporalio.contrib.streams._stream import StreamRef, Topic from temporalio.contrib.streams._stream_client import ( ClientStream, ClientTopic, StreamClient, ) -from temporalio.contrib.streams._workflow import ( - WorkflowStream, - create_workflow_stream, - workflow_stream, -) +from temporalio.contrib.streams.workflow import WorkflowStream, WorkflowTopic __all__ = [ "ActivityStream", @@ -65,37 +62,29 @@ "ClientStream", "ClientTopic", "Closed", - "OutsideStream", "ReadOnlyContextError", "Record", "StreamBackend", "StreamBufferFullError", "StreamClient", - "StreamDefinition", "StreamClosedError", + "StreamDefinition", "StreamError", "StreamExpiredError", "StreamFencedError", "StreamHandle", "StreamNotFoundError", "StreamOffsetError", - "StreamRef", "StreamStorageError", "StreamTransportMissingError", "StreamUnsupportedError", "Superseded", "SyncActivityStream", "SyncActivityTopic", - "Topic", "TopicDefinition", "WorkflowStream", - "activity_stream", - "create_stream", - "create_workflow_stream", + "WorkflowTopic", "defn", - "for_handle", - "for_ref", "topic", "workflow", - "workflow_stream", ] diff --git a/temporalio/contrib/streams/_access.py b/temporalio/contrib/streams/_access.py index d0c4b1ca6..80d1eda29 100644 --- a/temporalio/contrib/streams/_access.py +++ b/temporalio/contrib/streams/_access.py @@ -74,7 +74,7 @@ async def _read( types=types, from_offset=from_offset, wait_new_records=wait_new_records, - wait_for_create=create_wait, + create_wait=create_wait, ) async for item in items: yield item diff --git a/temporalio/contrib/streams/_activity.py b/temporalio/contrib/streams/_activity.py index b3ce24a4f..4c171d937 100644 --- a/temporalio/contrib/streams/_activity.py +++ b/temporalio/contrib/streams/_activity.py @@ -1,59 +1,46 @@ -"""Streams in Activity code.""" +"""The Stream state of one Activity attempt.""" from __future__ import annotations import asyncio import contextvars -from collections.abc import Awaitable, Callable from dataclasses import dataclass, field -from datetime import timedelta from typing import Any import temporalio.api.streamservice.v1 as stream_pb import temporalio.converter -from temporalio.contrib.streams import _errors from temporalio.contrib.streams._backend import StreamBackend -from temporalio.contrib.streams._outside import OutsideStream, batch_options -from temporalio.contrib.streams._stream import StreamRef -from temporalio.contrib.streams._writer import BatchOptions, StreamWriter +from temporalio.contrib.streams._writer import StreamWriter @dataclass class ActivityStreams: - """What a Stream plugin gives the facade for one Activity attempt. - - .. warning:: - This API is experimental. - """ + """What a Stream plugin gives the Activity handles for one Activity attempt.""" backend: StreamBackend converter: temporalio.converter.DataConverter namespace: str producer: stream_pb.Producer - default_ref: StreamRef | None loop: asyncio.AbstractEventLoop backend_name: str = "" - batch: BatchOptions = BatchOptions() 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) - def writer(self, stream_id: str, batch: BatchOptions | None = None) -> StreamWriter: + def writer(self, stream_id: str) -> StreamWriter: """The attempt's writer for a Stream, one per Stream.""" writer = self.writers.get(stream_id) if writer is None: + + async def known_id() -> str: + return stream_id + writer = StreamWriter( backend=self.backend, converter=self.converter, namespace=self.namespace, - stream_id=_known(stream_id), + stream_id=known_id, producer=self.producer, - batch=batch or self.batch, - # The default Stream has no key until its first write, so its create goes - # first. A received Stream must exist already. - create_first=self.default_ref is not None - and stream_id == self.default_ref.id, - loop=self.loop, ) self.writers[stream_id] = writer return writer @@ -75,56 +62,6 @@ async def close(self) -> None: raise errors[0] -def _known(stream_id: str) -> Callable[[], Awaitable[str]]: - async def find() -> str: - return stream_id - - return find - - current: contextvars.ContextVar[ActivityStreams | None] = contextvars.ContextVar( "__temporal_activity_streams", default=None ) - - -def activity_stream( - ref: StreamRef | None = None, - *, - batch_interval: timedelta | None = None, - max_batch_size: int | None = None, - max_retry_duration: timedelta | None = None, -) -> OutsideStream: - """A Stream for this Activity to write to and read. - - Writes are buffered and flushed every ``batch_interval``. The plugin flushes before the - Activity completes, and a failed flush fails the attempt. - - Args: - ref: A Stream that the Activity received. With none, the default Stream of the - Workflow that scheduled the Activity. - batch_interval: How often the writer flushes. Defaults to 100 ms. - max_batch_size: At most this many Records in one batch. - max_retry_duration: How long a flush retries a Store failure. - """ - streams = current.get() - if streams is None: - raise _errors.StreamTransportMissingError( - "No Stream plugin runs this Activity, such as RedisStreamPlugin" - ) - if ref is None: - if streams.default_ref is None: - raise _errors.StreamError( - "A Standalone Activity has no default Stream, so pass a StreamRef" - ) - ref = streams.default_ref - stream_id = ref.id - batch = None - if batch_interval or max_batch_size or max_retry_duration: - batch = batch_options(batch_interval, max_batch_size, max_retry_duration) - return OutsideStream( - backend=streams.backend, - converter=streams.converter, - namespace=streams.namespace, - stream_id=stream_id, - writer=lambda: streams.writer(stream_id, batch), - ) diff --git a/temporalio/contrib/streams/_backend.py b/temporalio/contrib/streams/_backend.py index 577d45e7b..a2a6ac994 100644 --- a/temporalio/contrib/streams/_backend.py +++ b/temporalio/contrib/streams/_backend.py @@ -2,9 +2,8 @@ from __future__ import annotations -import weakref from datetime import timedelta -from typing import Any, TypeVar +from typing import TypeVar import google.protobuf.message @@ -48,8 +47,6 @@ class StreamBackend: def __init__(self, extension: temporalio.worker.WorkerExtension) -> None: """Create a backend over an extension.""" self._extension = extension - # Writers whose open or encode may still run on the loop, for close_writers. - self.writers: weakref.WeakSet[Any] = weakref.WeakSet() @property def extension(self) -> temporalio.worker.WorkerExtension: @@ -84,23 +81,10 @@ async def open_writer(self, request: ext_pb.OpenWriterRequest) -> int: ) return response.writer_id - def publish( - self, writer_id: int, records: list[stream_pb.AppendRecord], force_flush: bool - ) -> None: - """Add encoded Records to a writer's buffer.""" - self.call( - "Publish", - ext_pb.PublishRequest( - writer_id=writer_id, records=records, force_flush=force_flush - ), - ) - async def write_records( self, writer_id: int, records: list[stream_pb.AppendRecord] ) -> int: - """Write Records at once for a writer, after its buffered ones. Returns the last - Record's offset. - """ + """Write Records at once for a writer. Returns the last Record's offset.""" response = await self.call_async( "WriteRecords", ext_pb.WriteRecordsRequest(writer_id=writer_id, records=records), @@ -124,15 +108,6 @@ async def flush_run_buffer( ) return response.records - async def flush_writer(self, writer_id: int) -> int: - """Wait until the writer's Records are in the Store.""" - response = await self.call_async( - "FlushWriter", - ext_pb.FlushWriterRequest(writer_id=writer_id), - ext_pb.FlushWriterResponse, - ) - return response.last_offset - async def close_writer(self, writer_id: int) -> int: """Flush and drop a writer. The Stream stays open.""" response = await self.call_async( @@ -143,17 +118,11 @@ async def close_writer(self, writer_id: int) -> int: return response.last_offset async def close_writers(self) -> int: - """Flush and drop every open writer of the extension, and return how many failed.""" - failed = 0 - for writer in list(self.writers): - try: - await writer.settle() - except Exception: - failed += 1 + """Drop every open writer of the extension, and return how many failed.""" response = await self.call_async( "CloseWriters", ext_pb.CloseWritersRequest(), ext_pb.CloseWritersResponse ) - return failed + response.failed + return response.failed async def drain_closes(self, timeout: timedelta) -> int: """Wait for the Workflow closes in flight, and return how many are left.""" @@ -194,17 +163,6 @@ async def call( ) -def backend_of_client(client: temporalio.client.Client) -> StreamBackend: - """The backend that a client's Stream plugin gives.""" - for plugin in client.config()["plugins"]: - backend = getattr(plugin, "_temporal_stream_backend", None) - if isinstance(backend, StreamBackend): - return backend - raise _errors.StreamTransportMissingError( - "The client has no Stream plugin, such as RedisStreamPlugin" - ) - - def client_backend( client: temporalio.client.Client, name: str | None ) -> tuple[str, StreamBackend]: diff --git a/temporalio/contrib/streams/_client.py b/temporalio/contrib/streams/_client.py deleted file mode 100644 index 94e89748a..000000000 --- a/temporalio/contrib/streams/_client.py +++ /dev/null @@ -1,251 +0,0 @@ -"""Streams in client code.""" - -from __future__ import annotations - -import asyncio -import itertools -import uuid -from collections.abc import Awaitable, Callable -from datetime import timedelta -from typing import Any - -import temporalio.api.streamservice.v1 as stream_pb -import temporalio.client -from temporalio.contrib.streams import _convert, _errors, _ids -from temporalio.contrib.streams._backend import StreamBackend, backend_of_client -from temporalio.contrib.streams._outside import OutsideStream, batch_options -from temporalio.contrib.streams._stream import StreamRef -from temporalio.contrib.streams._writer import BatchOptions, StreamWriter - -# A client's Producer id has a key per process and writer, so two writers never share a -# fence or a sequence. The crate takes two parts after `client:`, so the writer number goes into -# the key. -_PROCESS_KEY = str(uuid.uuid4()) -_WRITERS = itertools.count(1) - - -def _client_producer(client: temporalio.client.Client) -> stream_pb.Producer: - key = f"{_PROCESS_KEY}-{next(_WRITERS)}" - return stream_pb.Producer( - kind=stream_pb.ProducerKind.PRODUCER_KIND_CLIENT, - id=_ids.client_producer_id(client.identity, key), - epoch=stream_pb.ProducerEpoch(attempt=1), - ) - - -class ClientStream(OutsideStream): - """A Stream that a client reads, writes and closes. - - .. warning:: - This API is experimental. - """ - - def __init__( - self, - client: temporalio.client.Client, - backend: StreamBackend, - stream_id: str | Callable[[], Awaitable[str]], - batch: BatchOptions, - *, - wait_for_create: timedelta | None = None, - handle: temporalio.client.WorkflowHandle[Any, Any] | None = None, - ) -> None: - """Create a Stream handle. Use :py:func:`for_handle` or :py:func:`create_stream`.""" - self._client = client - self._handle = handle - converter = client.data_converter - super().__init__( - backend=backend, - converter=converter, - namespace=client.namespace, - stream_id=stream_id, - writer=lambda: StreamWriter( - backend=backend, - converter=converter, - namespace=client.namespace, - stream_id=self._stream_id, - producer=_client_producer(client), - batch=batch, - create_first=False, - loop=asyncio.get_running_loop(), - ), - wait_for_create=wait_for_create, - run_ended=self._handle_closed if handle is not None else None, - ) - - async def _handle_closed(self) -> bool: - assert self._handle is not None - description = await self._handle.describe() - return ( - description.status is not None - and description.status != temporalio.client.WorkflowExecutionStatus.RUNNING - ) - - async def close( - self, *, result: Any | None = None, failure: BaseException | None = None - ) -> None: - """Flush this client's Records and close the Stream. Only its creator closes it.""" - producer = self._get_writer().producer - await self._get_writer().close() - request = stream_pb.CloseStreamRequest( - namespace=self._namespace, - stream=stream_pb.StreamRef(id=await self._stream_id()), - producer=producer, - ) - if result is not None: - request.result.CopyFrom( - (await _convert.encode(self._converter, result)).payload - ) - if failure is not None: - await self._converter.encode_failure(failure, request.failure) - await self._backend.call_async( - "CloseStream", request, stream_pb.CloseStreamResponse - ) - - -def for_handle( - handle: temporalio.client.WorkflowHandle[Any, Any], - *, - name: str | None = None, - wait_for_create: timedelta | None = None, - end_with_run: bool = False, - batch_interval: timedelta | None = None, - max_batch_size: int | None = None, - max_retry_duration: timedelta | None = None, -) -> ClientStream: - """The default Stream of a Workflow, or the Stream it created with ``name``. - - It builds the id from the Workflow id and the first run id, so it makes no call. A handle - with no first run id makes one Describe call on the Stream's first use. - - Args: - handle: A handle of the Workflow, best one that ``start_workflow`` returned. - name: The name that the Workflow gave ``create_workflow_stream``. - wait_for_create: How long a read waits for a Stream that doesn't exist yet. Defaults - to no wait. - end_with_run: End a read when the run is no longer running, also with no close. It - also ends a wait for the create. - batch_interval: How often this client's writes flush. Defaults to 100 ms. - max_batch_size: At most this many Records in one batch. - max_retry_duration: How long a flush retries a Store failure. - """ - - async def find_id() -> str: - # The design allows one Describe call when the handle has no first run id. - description = await handle.describe() - first = description.raw_description.workflow_execution_info.first_run_id - if not first: - raise ValueError( - "The server gave no first run id, so the Stream id is unknown" - ) - stream._owner = (handle.id, first) - return stream_id_of(first) - - def stream_id_of(first_run_id: str) -> str: - if name is not None: - return _ids.scoped_stream_id(handle.id, first_run_id, name) - return _ids.default_stream_id(handle.id, first_run_id) - - stream_id = ( - stream_id_of(handle.first_execution_run_id) - if handle.first_execution_run_id - else find_id - ) - stream = ClientStream( - handle._client, - backend_of_client(handle._client), - stream_id, - batch_options(batch_interval, max_batch_size, max_retry_duration), - wait_for_create=wait_for_create, - handle=handle if end_with_run else None, - ) - if handle.first_execution_run_id: - stream._owner = (handle.id, handle.first_execution_run_id) - return stream - - -def for_ref( - client: temporalio.client.Client, - ref: StreamRef, - *, - wait_for_create: timedelta | None = None, - end_with_run: bool = False, - batch_interval: timedelta | None = None, - max_batch_size: int | None = None, - max_retry_duration: timedelta | None = None, -) -> ClientStream: - """A Stream from a :py:class:`StreamRef`, such as one that a Workflow published. - - It makes no call. - - Args: - client: The client with the Stream plugin. - ref: The reference of the Stream. - wait_for_create: How long a read waits for a Stream that doesn't exist yet. Defaults - to no wait. - end_with_run: End a read when the owning Workflow's run is no longer running, as - :py:func:`for_handle` does. Only a Workflow's Stream has an owner. - batch_interval: How often this client's writes flush. Defaults to 100 ms. - max_batch_size: At most this many Records in one batch. - max_retry_duration: How long a flush retries a Store failure. - - Raises: - ValueError: ``end_with_run`` is set, and the ref has no owning Workflow. - """ - handle = None - if end_with_run: - if not ref.workflow_id or not ref.first_execution_run_id: - raise ValueError("end_with_run needs a ref of a Workflow's Stream") - handle = client.get_workflow_handle( - ref.workflow_id, first_execution_run_id=ref.first_execution_run_id - ) - stream = ClientStream( - client, - backend_of_client(client), - ref.id, - batch_options(batch_interval, max_batch_size, max_retry_duration), - wait_for_create=wait_for_create, - handle=handle, - ) - if ref.workflow_id and ref.first_execution_run_id: - stream._owner = (ref.workflow_id, ref.first_execution_run_id) - return stream - - -async def create_stream( - client: temporalio.client.Client, - stream_id: str, - *, - retention: timedelta | None = None, - batch_interval: timedelta | None = None, - max_batch_size: int | None = None, - max_retry_duration: timedelta | None = None, -) -> ClientStream: - """Create a standalone Stream. The client owns it, so only the client closes it. - - It is idempotent for one id and one Retention. - - Raises: - StreamError: The Stream exists with another Retention, with the reason - ``stream_already_exists``. - """ - backend = backend_of_client(client) - request = stream_pb.CreateStreamRequest( - namespace=client.namespace, stream_id=stream_id - ) - if retention is not None: - request.retention.FromTimedelta(retention) - response = await backend.call_async( - "CreateStream", request, stream_pb.CreateStreamResponse - ) - if response.status == stream_pb.StreamStatus.STREAM_STATUS_CLOSED: - raise _errors.StreamClosedError( - f"The Stream {stream_id} was created and closed before", - reason="stream_closed", - ) - return ClientStream( - client, - backend, - stream_id, - batch_options(batch_interval, max_batch_size, max_retry_duration), - ) diff --git a/temporalio/contrib/streams/_convert.py b/temporalio/contrib/streams/_convert.py index dd70d651c..e0980c2e4 100644 --- a/temporalio/contrib/streams/_convert.py +++ b/temporalio/contrib/streams/_convert.py @@ -55,15 +55,6 @@ async def encode( return temporalio.common.RawValue(payload) -def encode_now( - converter: temporalio.converter.DataConverter, value: Any -) -> temporalio.common.RawValue | None: - """Convert one value at once, or None when a codec makes it async.""" - if converter.payload_codec is not None: - return None - return temporalio.common.RawValue(converter.payload_converter.to_payload(value)) - - async def decode( converter: temporalio.converter.DataConverter, body: Any, type_hint: Any ) -> Any: diff --git a/temporalio/contrib/streams/_ids.py b/temporalio/contrib/streams/_ids.py index 732426aba..d688a5768 100644 --- a/temporalio/contrib/streams/_ids.py +++ b/temporalio/contrib/streams/_ids.py @@ -14,11 +14,6 @@ def _join(*parts: str) -> str: return ":".join(encode_part(part) for part in parts) -def default_stream_id(workflow_id: str, first_run_id: str) -> str: - """Id of a Workflow's default Stream.""" - return _join(workflow_id, first_run_id) - - def scoped_stream_id(workflow_id: str, first_run_id: str, name: str) -> str: """Id of a Stream that a Workflow creates with a name.""" return _join(workflow_id, first_run_id, name) diff --git a/temporalio/contrib/streams/_outside.py b/temporalio/contrib/streams/_outside.py deleted file mode 100644 index 9ebf0f86c..000000000 --- a/temporalio/contrib/streams/_outside.py +++ /dev/null @@ -1,192 +0,0 @@ -"""Streams outside Workflow code, for Activities and clients.""" - -from __future__ import annotations - -from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence -from datetime import timedelta -from typing import Any, TypeVar, overload - -import temporalio.api.streamservice.v1 as stream_pb -import temporalio.converter -from temporalio.contrib.streams import _reader -from temporalio.contrib.streams._backend import StreamBackend -from temporalio.contrib.streams._reader import Item -from temporalio.contrib.streams._stream import StreamRef, Topic, _Readable -from temporalio.contrib.streams._writer import BatchOptions, StreamWriter - -T = TypeVar("T") - - -class OutsideStream(_Readable): - """A Stream that an Activity or a client reads and writes. - - .. warning:: - This API is experimental. - """ - - def __init__( - self, - *, - backend: StreamBackend, - converter: temporalio.converter.DataConverter, - namespace: str, - stream_id: str | Callable[[], Awaitable[str]], - writer: Callable[[], StreamWriter], - wait_for_create: timedelta | None = None, - run_ended: Callable[[], Awaitable[bool]] | None = None, - ) -> None: - """Create a Stream handle. Use the facade functions instead.""" - self._backend = backend - self._converter = converter - self._namespace = namespace - # A client that has only a Workflow id learns the first run id on its first call. - self._id = stream_id if isinstance(stream_id, str) else "" - # The owning Workflow of a Workflow's Stream, which its ref carries. - self._owner: tuple[str, str] | None = None - self._find_id = None if isinstance(stream_id, str) else stream_id - self._writer_factory = writer - self._writer: StreamWriter | None = None - self._wait_for_create = wait_for_create - self._run_ended = run_ended - # A read of every Topic still decodes each known Topic with its type. - self._types: dict[str, Any] = {} - - @property - def ref(self) -> StreamRef: - """Reference to pass to Workflows, Activities and other clients. - - Raises: - RuntimeError: The Stream id needs the first run id, and no call has found it yet. - Call :py:meth:`get_offset` first. - """ - if not self._id: - raise RuntimeError( - "The Stream id is not known yet. Call get_offset() first." - ) - if self._owner is None: - return StreamRef(self._id) - return StreamRef( - self._id, - workflow_id=self._owner[0], - first_execution_run_id=self._owner[1], - ) - - async def _stream_id(self) -> str: - if not self._id: - assert self._find_id is not None - self._id = await self._find_id() - return self._id - - @overload - def topic(self, name: str) -> Topic[Any]: ... - - @overload - def topic(self, name: str, *, type: type[T]) -> Topic[T]: ... - - def topic(self, name: str, *, type: Any = None) -> Topic[Any]: - """A Topic of this Stream. ``type`` decodes its values, and ``RawValue`` keeps the - Payload. With no type, values decode as the converter gives them. - """ - self._types[name] = type - return Topic(self, name, type) - - def _get_writer(self) -> StreamWriter: - if self._writer is None: - self._writer = self._writer_factory() - return self._writer - - def _publish(self, topic: str, value: Any, force_flush: bool) -> None: - self._get_writer().publish(topic, value, force_flush) - - async def flush(self) -> None: - """Wait until the Records published so far are in the Store.""" - if self._writer is not None: - await self._writer.flush() - - def subscribe( - self, - *, - topics: Sequence[Topic[Any]] | None = None, - topic_prefixes: Sequence[str] = (), - from_offset: int | None = None, - wait_new_records: bool = True, - type: Any = None, - ) -> AsyncGenerator[Item, None]: - """Read the Stream's items in order. - - The read ends after a ``Closed`` item, or at the newest Record when - ``wait_new_records`` is False. A ``Superseded`` item withdraws earlier Records of - a failed attempt, and it passes every Topic filter. - - Args: - topics: Read these Topics, by exact match. - topic_prefixes: Also read the Topics that start with one of these. With no - ``topics`` and no prefixes, the read gives every Topic. - from_offset: Start at this offset, inclusive. Defaults to the first Record. - wait_new_records: Wait for new Records at the end of the Stream. - type: Decode the values of Topics with no type as this type. A Topic from - :py:meth:`topic` with a type keeps its own. - """ - return self._read( - topics=[topic.name for topic in topics] if topics else None, - topic_prefixes=topic_prefixes, - types={ - **self._types, - **{topic.name: topic.type for topic in topics or []}, - }, - default_type=type, - from_offset=from_offset, - wait_new_records=wait_new_records, - ) - - async def _read( - self, - *, - topics: list[str] | None, - topic_prefixes: Sequence[str], - types: dict[str, Any], - default_type: Any, - from_offset: int | None, - wait_new_records: bool, - ) -> AsyncGenerator[Item, None]: - items = _reader.read( - self._backend.service(self._namespace), - self._converter, - await self._stream_id(), - topics=topics, - topic_prefixes=topic_prefixes, - types=types, - default_type=default_type, - from_offset=from_offset, - wait_new_records=wait_new_records, - wait_for_create=self._wait_for_create, - run_ended=self._run_ended, - ) - async for item in items: - yield item - - async def get_offset(self) -> int: - """The next offset of the Stream, the newest offset plus one.""" - response = await self._backend.call_async( - "DescribeStream", - stream_pb.DescribeStreamRequest( - namespace=self._namespace, - stream=stream_pb.StreamRef(id=await self._stream_id()), - ), - stream_pb.DescribeStreamResponse, - ) - # Offsets start above zero, so zero means the Stream has no Record. - return response.last_offset + 1 if response.last_offset else 0 - - -def batch_options( - batch_interval: timedelta | None, - max_batch_size: int | None, - max_retry_duration: timedelta | None, -) -> BatchOptions: - """Batch options from the facade keywords.""" - return BatchOptions( - batch_interval=batch_interval, - max_batch_size=max_batch_size, - max_retry_duration=max_retry_duration, - ) diff --git a/temporalio/contrib/streams/_reader.py b/temporalio/contrib/streams/_reader.py index 52d8e70f6..defd0ffd5 100644 --- a/temporalio/contrib/streams/_reader.py +++ b/temporalio/contrib/streams/_reader.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Mapping, Sequence from datetime import timedelta from typing import Any @@ -18,10 +18,7 @@ # The service waits at most 60 s per read, so a longer wait is a loop of reads. _MAX_WAIT = timedelta(seconds=60) -# How often a reader that ends with the run checks the run, while it waits for the create -# or for new Records. -_RUN_CHECK = timedelta(seconds=2) -# How long one loop of the generated long poll may run before the facade looks again. +# How long one loop of the generated long poll may run before the reader looks again. _POLL_DEADLINE = timedelta(minutes=5) @@ -33,24 +30,17 @@ async def read( topics: Sequence[str] | None, topic_prefixes: Sequence[str] = (), types: Mapping[str, Any], - default_type: Any = None, from_offset: int | None, wait_new_records: bool, - wait_for_create: timedelta | None, - run_ended: Callable[[], Awaitable[bool]] | None = None, + create_wait: timedelta | None, ) -> AsyncGenerator[Item, None]: """Give the items of a Stream in order, until a close, or the newest Record when ``wait_new_records`` is False. """ loop = asyncio.get_running_loop() - create_until = ( - loop.time() + wait_for_create.total_seconds() if wait_for_create else None - ) + create_until = loop.time() + create_wait.total_seconds() if create_wait else None offset = from_offset names = list(topics or []) - # A reader that ends with the run waits for the create in short steps, so it sees the - # run end soon. - create_step = _RUN_CHECK if run_ended is not None else _MAX_WAIT while True: create_left = ( timedelta(seconds=max(create_until - loop.time(), 0)) @@ -63,28 +53,22 @@ async def read( start=None if offset is not None else ReadStart.READ_START_BEGINNING, topics=names, topic_prefixes=list(topic_prefixes), - create_wait=min(create_left, create_step) if create_left else None, + create_wait=min(create_left, _MAX_WAIT) if create_left else None, ) try: with _convert.raw_bodies(converter.failure_converter): if wait_new_records: - # A reader that ends with the run looks at the run between short polls. answer = await service.read_records_until_records( - request, - deadline=_RUN_CHECK - if run_ended is not None - else _POLL_DEADLINE, + request, deadline=_POLL_DEADLINE ) else: answer = await service.read_records(request) except _errors.StreamNotFoundError: if create_until is not None and loop.time() < create_until: - if run_ended is not None and await run_ended(): - return continue raise for record in answer.records or []: - item = await _item(converter, record, types, default_type) + item = await _item(converter, record, types) if item is None: continue yield item @@ -96,20 +80,17 @@ async def read( continue if answer.done or not wait_new_records: return - if run_ended is not None and await run_ended(): - return async def _item( converter: temporalio.converter.DataConverter, record: models.StreamRecord, types: Mapping[str, Any], - default_type: Any, ) -> Item | None: offset = record.offset or 0 if record.kind == StreamRecordKind.STREAM_RECORD_KIND_DATA: topic = record.topic or "" - type_hint = types.get(topic) or default_type + type_hint = types.get(topic) data = await _convert.decode(converter, record.body, type_hint) return Record(topic=topic, data=data, offset=offset, producer=record.producer) if record.kind == StreamRecordKind.STREAM_RECORD_KIND_SUPERSEDED: diff --git a/temporalio/contrib/streams/_stream.py b/temporalio/contrib/streams/_stream.py deleted file mode 100644 index 3c73ce516..000000000 --- a/temporalio/contrib/streams/_stream.py +++ /dev/null @@ -1,166 +0,0 @@ -"""Stream references and Topics.""" - -from __future__ import annotations - -import abc -from collections.abc import AsyncGenerator, Sequence -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Generic, TypeVar - -import temporalio.converter - -if TYPE_CHECKING: - from temporalio.contrib.streams._reader import Item - -T = TypeVar("T") - - -class _StreamRefJson( - temporalio.converter.TransferTypeConverter["StreamRef", dict[str, str]] -): - # The JSON form is the proto JSON of temporal.api.streamservice.v1.StreamRef, so every SDK - # reads a ref that another wrote. Snake case keys from older refs still decode. - transfer_type: type[dict[str, str]] | None = dict - - def to_transfer_type(self, value: StreamRef) -> dict[str, str]: - json = {"id": value.id} - if value.workflow_id: - json["workflowId"] = value.workflow_id - if value.first_execution_run_id: - json["firstExecutionRunId"] = value.first_execution_run_id - return json - - def from_transfer_type( - self, value: dict[str, str], type_hint: type[StreamRef] - ) -> StreamRef: - return StreamRef( - value["id"], - workflow_id=value.get("workflowId") or value.get("workflow_id") or None, - first_execution_run_id=value.get("firstExecutionRunId") - or value.get("first_execution_run_id") - or None, - ) - - -def _ref_from_json(value: Any) -> StreamRef: - if isinstance(value, StreamRef): - return value - if not isinstance(value, dict): - raise ValueError(f"Expected a StreamRef or its JSON object, got {type(value)}") - return _StreamRefJson().from_transfer_type(value, StreamRef) - - -@temporalio.converter.transfer_type_convertible(_StreamRefJson) -@dataclass(frozen=True) -class StreamRef: - """Names one Stream. Pass it to Workflows and Activities as an argument. - - .. warning:: - This API is experimental. - """ - - id: str - """The Stream id. It is opaque to callers.""" - - workflow_id: str | None = None - """For a Workflow's Stream, the id of the Workflow that owns it. None for a standalone - Stream.""" - - first_execution_run_id: str | None = None - """For a Workflow's Stream, the first run id of the owning run chain. None for a - standalone Stream.""" - - @classmethod - def __get_pydantic_core_schema__(cls, source: Any, handler: Any) -> Any: - """Give pydantic models the proto JSON form of the ref, as the payload converter does. - - Pydantic is optional, so it is imported only when a model uses the ref. - """ - from pydantic_core import core_schema - - return core_schema.no_info_plain_validator_function( - _ref_from_json, - json_schema_input_schema=core_schema.typed_dict_schema( - { - "id": core_schema.typed_dict_field(core_schema.str_schema()), - "workflowId": core_schema.typed_dict_field( - core_schema.str_schema(), required=False - ), - "firstExecutionRunId": core_schema.typed_dict_field( - core_schema.str_schema(), required=False - ), - } - ), - serialization=core_schema.plain_serializer_function_ser_schema( - _StreamRefJson().to_transfer_type - ), - ) - - -class _Writable(abc.ABC): - # A Stream kind that Topics publish to. - @abc.abstractmethod - def _publish(self, topic: str, value: Any, force_flush: bool) -> None: ... - - -class _Readable(_Writable): - # A Stream kind that Topics read from: Activities and clients. - @abc.abstractmethod - def subscribe( - self, - *, - topics: Sequence[Topic[Any]] | None = None, - topic_prefixes: Sequence[str] = (), - from_offset: int | None = None, - wait_new_records: bool = True, - type: Any = None, - ) -> AsyncGenerator[Item, None]: ... - - -class Topic(Generic[T]): - """One Topic of a Stream. All Topics of a Stream share one order. - - .. warning:: - This API is experimental. - """ - - def __init__(self, stream: _Writable, name: str, type: Any) -> None: - """Create a Topic. Use ``topic()`` of a Stream instead.""" - self._stream = stream - self._name = name - self._type = type - - @property - def name(self) -> str: - """Name of the Topic.""" - return self._name - - @property - def type(self) -> Any: - """Type that a read decodes the Topic's values with.""" - return self._type - - def publish(self, value: T, *, force_flush: bool = False) -> None: - """Add a Record to the Topic. It returns at once and the Record is buffered. - - Args: - value: The value. It goes through the payload converter and codec. - force_flush: Outside Workflow code, start a flush at once. Workflow code writes - its Records when the Workflow Task completes, so it has no effect there. - """ - self._stream._publish(self._name, value, force_flush) - - def subscribe( - self, - *, - from_offset: int | None = None, - wait_new_records: bool = True, - ) -> AsyncGenerator[Item, None]: - """Read this Topic only. See the Stream's ``subscribe``.""" - if not isinstance(self._stream, _Readable): - raise TypeError("Workflow code can't read a Stream") - return self._stream.subscribe( - topics=[self], - from_offset=from_offset, - wait_new_records=wait_new_records, - ) diff --git a/temporalio/contrib/streams/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index 3fffa2a0d..f318e3016 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -3,6 +3,8 @@ from __future__ import annotations import asyncio +import itertools +import uuid from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence from datetime import timedelta from typing import Any, TypeVar, overload @@ -12,15 +14,29 @@ from temporalio.contrib.streams import _errors, _ids from temporalio.contrib.streams._access import Access, _AsyncTopic, read_spec from temporalio.contrib.streams._backend import StreamBackend, client_backend -from temporalio.contrib.streams._client import _client_producer from temporalio.contrib.streams._defn import TopicDefinition, definition_of from temporalio.contrib.streams._handle import StreamHandle from temporalio.contrib.streams._reader import Item -from temporalio.contrib.streams._writer import BatchOptions, StreamWriter +from temporalio.contrib.streams._writer import StreamWriter E = TypeVar("E") T = TypeVar("T") +# A client's Producer id has a key per process and handle, so two handles never share a fence +# or a sequence. The crate takes two parts after `client:`, so the handle number goes into the +# key. +_PROCESS_KEY = str(uuid.uuid4()) +_WRITERS = itertools.count(1) + + +def _client_producer(client: temporalio.client.Client) -> stream_pb.Producer: + key = f"{_PROCESS_KEY}-{next(_WRITERS)}" + return stream_pb.Producer( + kind=stream_pb.ProducerKind.PRODUCER_KIND_CLIENT, + id=_ids.client_producer_id(client.identity, key), + epoch=stream_pb.ProducerEpoch(attempt=1), + ) + class ClientTopic(_AsyncTopic[T]): """One Topic of a :py:class:`ClientStream`. All Topics of a Stream share one order. @@ -128,9 +144,6 @@ def writer() -> StreamWriter: namespace=client.namespace, stream_id=self._stream_id, producer=_client_producer(client), - batch=BatchOptions(), - create_first=False, - loop=asyncio.get_running_loop(), ) return self._writer diff --git a/temporalio/contrib/streams/_workflow.py b/temporalio/contrib/streams/_workflow.py deleted file mode 100644 index bce3e8d0d..000000000 --- a/temporalio/contrib/streams/_workflow.py +++ /dev/null @@ -1,257 +0,0 @@ -"""Streams in Workflow code.""" - -from __future__ import annotations - -import asyncio -import weakref -from collections.abc import Coroutine -from datetime import timedelta -from types import TracebackType -from typing import Any, TypeVar, overload - -import temporalio.exceptions -import temporalio.workflow -from temporalio.contrib.streams import _errors, _ids -from temporalio.contrib.streams._stream import StreamRef, Topic, _Writable -from temporalio.contrib.streams.system import models -from temporalio.contrib.streams.system.operations.append_records import ( - append_records, -) -from temporalio.contrib.streams.system.operations.close_stream import close_stream -from temporalio.contrib.streams.system.operations.create_stream import ( - create_stream, -) - -T = TypeVar("T") -_R = TypeVar("_R") - - -class _RunState: - def __init__(self) -> None: - self.streams: dict[str, WorkflowStream] = {} - - -# One state per Workflow instance. The temporalio modules pass through the sandbox, so the key -# is the instance, and the state goes when the run leaves the cache. -_states: weakref.WeakKeyDictionary[Any, _RunState] = weakref.WeakKeyDictionary() - - -def _state() -> _RunState: - runtime = temporalio.workflow._Runtime.current() - state = _states.get(runtime) - if state is None: - state = _RunState() - _states[runtime] = state - return state - - -def _has_transport() -> bool: - # A Stream plugin marks its outbound interceptor. With none, the call would go to the - # server as a System Nexus command, which only native Streams serve. - outbound = getattr(temporalio.workflow._Runtime.current(), "_outbound", None) - while outbound is not None: - if getattr(outbound, "_temporal_stream_transport", False): - return True - outbound = getattr(outbound, "next", None) - return False - - -def _now(call: Coroutine[Any, Any, _R]) -> _R: - # The plugin's interceptor answers at once, so the call ends in its first step. - try: - call.send(None) - except StopIteration as done: - return done.value - call.close() - raise _errors.StreamTransportMissingError( - "No Stream plugin took the call, such as RedisStreamPlugin" - ) - - -class WorkflowStream(_Writable): - """A Stream that Workflow code writes to. - - Writes are buffered, and the Records go to the Store when the Workflow Task completes. - Workflow code can't read a Stream. - - .. warning:: - This API is experimental. - """ - - def __init__( - self, - ref: StreamRef, - *, - create: bool, - retention: timedelta | None, - auto_close: bool, - ) -> None: - """Create a Stream handle. Use :py:func:`workflow_stream` or - :py:func:`create_workflow_stream` instead. - """ - self._ref = ref - self._id = ref.id - self._create = create - self._retention = retention - self._auto_close = auto_close - self._created = False - self._used = False - self._closed = False - - @property - def ref(self) -> StreamRef: - """Reference to pass to Activities, child Workflows and clients.""" - return self._ref - - def _model_ref(self) -> models.StreamRef: - return models.StreamRef( - id=self._id, - workflow_id=self._ref.workflow_id, - first_execution_run_id=self._ref.first_execution_run_id, - ) - - @property - def auto_close(self) -> bool: - """Whether the run closes the Stream when it closes.""" - return self._auto_close - - @overload - def topic(self, name: str) -> Topic[Any]: ... - - @overload - def topic(self, name: str, *, type: type[T]) -> Topic[T]: ... - - def topic(self, name: str, *, type: Any = None) -> Topic[Any]: - """A Topic of this Stream. ``type`` decodes its values, and ``RawValue`` keeps the - Payload. With no type, values decode as the converter gives them. - """ - return Topic(self, name, type) - - def _check(self, action: str) -> None: - if temporalio.workflow.unsafe.is_read_only(): - raise _errors.ReadOnlyContextError( - f"Can't {action} a Stream in read-only code" - ) - if self._closed: - raise _errors.StreamClosedError( - f"Can't {action} the closed Stream {self._id}", reason="stream_closed" - ) - if not _has_transport(): - raise _errors.StreamTransportMissingError( - "No Stream plugin is on the Worker, such as RedisStreamPlugin" - ) - - def _ensure_created(self) -> None: - if self._create and not self._created: - self._created = True - _now(create_stream(stream_id=self._id, retention=self._retention)) - - def _publish(self, topic: str, value: Any, force_flush: bool) -> None: - self._check("write to") - self._ensure_created() - self._used = True - _now( - append_records( - stream=self._model_ref(), - records=[models.AppendRecord(topic=topic, body=value)], - ) - ) - - def close( - self, *, result: Any | None = None, failure: BaseException | None = None - ) -> None: - """Close the Stream. It goes out after the server accepts this Workflow Task. - - Readers get the Records before it, then a ``Closed`` item with ``result`` or - ``failure``. Later writes raise :py:class:`StreamClosedError`. - """ - self._check("close") - self._ensure_created() - self._closed = True - _now(close_stream(stream=self._model_ref(), result=result, failure=failure)) - - def __enter__(self) -> WorkflowStream: - """Use the Stream in a ``with`` block, which closes it at the end.""" - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: - """Close the Stream, with the exception as its failure if there is one. - - A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as - an eviction, leaves it open, since the block did not end in Workflow code. - """ - if self._closed: - return - if isinstance(exc, asyncio.CancelledError): - self.close( - failure=temporalio.exceptions.CancelledError(str(exc) or "Cancelled") - ) - elif exc is None or isinstance(exc, Exception): - self.close(failure=exc) - - -def workflow_stream( - ref: StreamRef | None = None, *, auto_close: bool | None = None -) -> WorkflowStream: - """A Stream for this Workflow to write to. - - Args: - ref: A Stream that the Workflow received. With none, the run chain's default Stream, - whose id is ``:``. - auto_close: Whether the run closes the Stream when the run closes, with the run's - result or failure. Defaults to True for the default Stream and False for a - received one. A run that continues as new does not close it. - """ - state = _state() - if ref is None: - info = temporalio.workflow.info() - ref = StreamRef( - _ids.default_stream_id(info.workflow_id, info.first_execution_run_id), - workflow_id=info.workflow_id, - first_execution_run_id=info.first_execution_run_id, - ) - create, close = True, True if auto_close is None else auto_close - else: - create, close = False, bool(auto_close) - stream_id = ref.id - stream = state.streams.get(stream_id) - if stream is None: - stream = WorkflowStream(ref, create=create, retention=None, auto_close=close) - state.streams[stream_id] = stream - elif auto_close is not None: - stream._auto_close = auto_close - return stream - - -def create_workflow_stream( - name: str, *, retention: timedelta | None = None -) -> WorkflowStream: - """Create a Stream with a name, scoped to this run chain. - - Its id is ``::``, so a client finds it with - :py:func:`for_handle`. The name must be unique in the run chain. Use it in a ``with`` - block, which closes it at the end. - - Args: - name: The name of the Stream. - retention: The age after which a Record is dropped. Defaults to the Store's. - """ - info = temporalio.workflow.info() - stream_id = _ids.scoped_stream_id( - info.workflow_id, info.first_execution_run_id, name - ) - ref = StreamRef( - stream_id, - workflow_id=info.workflow_id, - first_execution_run_id=info.first_execution_run_id, - ) - stream = WorkflowStream(ref, create=True, retention=retention, auto_close=False) - stream._check("create") - stream._ensure_created() - _state().streams[stream_id] = stream - return stream diff --git a/temporalio/contrib/streams/_writer.py b/temporalio/contrib/streams/_writer.py index baefdf736..11cffb501 100644 --- a/temporalio/contrib/streams/_writer.py +++ b/temporalio/contrib/streams/_writer.py @@ -1,33 +1,20 @@ -"""The buffered writer of an Activity or a client Producer.""" +"""The writer of an Activity or a client Producer.""" from __future__ import annotations import asyncio from collections.abc import Awaitable, Callable, Sequence -from dataclasses import dataclass -from datetime import timedelta from typing import Any import temporalio.api.streamservice.v1 as stream_pb -import temporalio.common import temporalio.converter -from temporalio.contrib.streams import _convert, _errors +from temporalio.contrib.streams import _errors from temporalio.contrib.streams._backend import StreamBackend from temporalio.contrib.streams.proto import v1 as ext_pb -@dataclass(frozen=True) -class BatchOptions: - """How a writer batches Records. An unset value takes the crate's default.""" - - batch_interval: timedelta | None = None - max_batch_size: int | None = None - max_retry_duration: timedelta | None = None - max_buffered_records: int | None = None - - class StreamWriter: - """Buffers the Records of one Producer on one Stream, through the crate's writer.""" + """Writes the Records of one Producer on one Stream, through the crate's writer.""" def __init__( self, @@ -37,29 +24,16 @@ def __init__( namespace: str, stream_id: Callable[[], Awaitable[str]], producer: stream_pb.Producer, - batch: BatchOptions, - create_first: bool, - loop: asyncio.AbstractEventLoop, ) -> None: - """Create a writer. It opens on its first publish. - - ``loop`` runs the writer's async work, so a sync Activity on another thread can - publish too. - """ + """Create a writer. It opens on its first write.""" self._backend = backend self._converter = converter self._namespace = namespace self._resolve_id = stream_id - self._stream_id = "" self._producer = producer - self._batch = batch - self._create_first = create_first - self._loop = loop self._writer_id: int | None = None self._open_lock = asyncio.Lock() - self._pending: asyncio.Task[None] | None = None self._closed = False - backend.writers.add(self) @property def producer(self) -> stream_pb.Producer: @@ -69,111 +43,34 @@ def producer(self) -> stream_pb.Producer: async def withdraw_earlier_attempts(self) -> None: """Open the writer of a new attempt, which withdraws the Records of the earlier ones. - It creates no Stream. A Stream that does not exist has nothing to withdraw, and the - first publish creates it. + A Stream that does not exist has nothing to withdraw. """ try: - await self.open(create=False) + await self.open() except _errors.StreamNotFoundError: pass - async def open(self, *, create: bool = True) -> int: + async def open(self) -> int: """Open the crate's writer once. The open of a new attempt withdraws the earlier one's - Records, even when it publishes nothing. + Records, even when it writes nothing. """ async with self._open_lock: if self._writer_id is None: - if not self._stream_id: - self._stream_id = await self._resolve_id() - if self._create_first and create: - await self._backend.call_async( - "CreateStream", - stream_pb.CreateStreamRequest( - namespace=self._namespace, stream_id=self._stream_id - ), - stream_pb.CreateStreamResponse, + self._writer_id = await self._backend.open_writer( + ext_pb.OpenWriterRequest( + namespace=self._namespace, + stream=stream_pb.StreamRef(id=await self._resolve_id()), + producer=self._producer, ) - request = ext_pb.OpenWriterRequest( - namespace=self._namespace, - stream=stream_pb.StreamRef(id=self._stream_id), - producer=self._producer, - max_batch_size=self._batch.max_batch_size or 0, - max_buffered_records=self._batch.max_buffered_records or 0, ) - if self._batch.batch_interval is not None: - request.batch_interval.FromTimedelta(self._batch.batch_interval) - if self._batch.max_retry_duration is not None: - request.max_retry_duration.FromTimedelta( - self._batch.max_retry_duration - ) - self._writer_id = await self._backend.open_writer(request) return self._writer_id - def publish(self, topic: str, value: Any, force_flush: bool) -> None: - """Add one Record to the buffer and return at once.""" - try: - on_loop = asyncio.get_running_loop() is self._loop - except RuntimeError: - on_loop = False - if not on_loop: - # A sync Activity runs on another thread, so the loop does the publish, and the - # thread waits only for the hand-off. - asyncio.run_coroutine_threadsafe( - self._publish_on_loop(topic, value, force_flush), self._loop - ).result() - return - self._publish(topic, value, force_flush) - - async def _publish_on_loop(self, topic: str, value: Any, force_flush: bool) -> None: - self._publish(topic, value, force_flush) - - def _publish(self, topic: str, value: Any, force_flush: bool) -> None: - if self._closed: - raise _errors.StreamClosedError( - "The writer is closed", reason="stream_closed" - ) - body = _convert.encode_now(self._converter, value) - if body is not None and self._writer_id is not None and self._pending is None: - self._send(topic, body, force_flush) - return - # The open and the codec are async, so they run in order in the background, and a - # flush raises what they raised. - previous = self._pending - self._pending = self._loop.create_task( - self._publish_later(previous, topic, value, force_flush) - ) - - async def _publish_later( - self, - previous: asyncio.Task[None] | None, - topic: str, - value: Any, - force_flush: bool, - ) -> None: - if previous is not None: - await previous - await self.open() - self._send(topic, await _convert.encode(self._converter, value), force_flush) - - def _send( - self, topic: str, body: temporalio.common.RawValue, force_flush: bool - ) -> None: - assert self._writer_id is not None - self._backend.publish( - self._writer_id, - [stream_pb.AppendRecord(topic=topic, body=body.payload)], - force_flush, - ) - async def write(self, topic: str, values: Sequence[Any]) -> int: - """Write the values as Records now, after the ones published so far. Returns the last - Record's offset. - """ + """Write the values as Records now. Returns the last Record's offset.""" if self._closed: raise _errors.StreamClosedError( "The writer is closed", reason="stream_closed" ) - await self._settle() writer_id = await self.open() bodies = await self._converter.encode(list(values)) return await self._backend.write_records( @@ -181,38 +78,10 @@ async def write(self, topic: str, values: Sequence[Any]) -> int: [stream_pb.AppendRecord(topic=topic, body=body) for body in bodies], ) - async def settle(self) -> None: - """Wait for the publishes that still open the writer or run the codec.""" - await self._settle() - - async def _settle(self) -> None: - pending, self._pending = self._pending, None - if pending is not None: - await pending - - async def flush(self) -> int: - """Wait until every Record published so far is in the Store. Returns the newest - offset that the writer wrote, or 0. - """ - await self._settle() - if self._writer_id is None: - return 0 - return await self._backend.flush_writer(self._writer_id) - - async def close(self) -> int: - """Flush and drop the writer. The Stream stays open.""" + async def close(self) -> None: + """Drop the writer. The Stream stays open.""" if self._closed: - return 0 + return self._closed = True - error: BaseException | None = None - try: - await self._settle() - except BaseException as err: - error = err - # The crate's writer must go even when a publish failed, or its flusher stays. - last = 0 if self._writer_id is not None: - last = await self._backend.close_writer(self._writer_id) - if error is not None: - raise error - return last + await self._backend.close_writer(self._writer_id) diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 5ac06e31f..0b3e698cc 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -24,11 +24,10 @@ import temporalio.runtime import temporalio.worker import temporalio.workflow -from temporalio.contrib.streams import _activity, _errors, _ids, _workflow +from temporalio.contrib.streams import _activity, _errors, _ids from temporalio.contrib.streams import workflow as streams_workflow from temporalio.contrib.streams._backend import StreamBackend from temporalio.contrib.streams._handle import StreamHandle -from temporalio.contrib.streams._stream import StreamRef from temporalio.contrib.streams.proto import v1 as ext_pb from temporalio.contrib.streams.system import models from temporalio.plugin import SimplePlugin @@ -41,9 +40,8 @@ _SERVICE = "temporal.api.streamservice.v1.StreamService" # How long a Worker's shutdown waits for the Stream closes that are still going out. _DRAIN_TIMEOUT = timedelta(seconds=10) -# The plugin puts these into the headers of each Activity and Local Activity that a Workflow -# schedules. The key keeps the Producer id the same across retries. -_REF_HEADER = "__temporal_stream_ref" +# The plugin puts this key into the headers of each Activity and Local Activity that a Workflow +# schedules. It keeps the Producer id the same across retries. _KEY_HEADER = "__temporal_stream_key" # Continue-as-new passes the ids of the open Streams that close with the run chain to the next # run, in this header. @@ -179,7 +177,9 @@ class RedisStreamPlugin(StreamPlugin): """Carries Stream calls to Redis, through the temporalio-streams runtime in the bridge. Put it on the client. A Worker takes its client's plugins, so one plugin serves Workflow - code, Activities and the client. + code, Activities and the client. It registers the backend name ``redis``, which the + Stream handles it creates carry. For now a Worker takes one Stream plugin, so it serves one + backend. A client may have more than one, and ``backend=`` picks one. .. warning:: This API is experimental. @@ -422,9 +422,9 @@ def chain_continues( ) -> bool: """Whether the server starts another run of the chain after this run ends. - Then the next run goes on with the default Stream, so the runtime skips this run's closes. - A cron run always continues. A failed run continues when its retry policy has attempts and - time left for the failure. + Then the next run goes on with the run chain's Streams, so the runtime skips this run's + closes. A cron run always continues. A failed run continues when its retry policy has + attempts and time left for the failure. """ if info.cron_schedule: return True @@ -462,45 +462,16 @@ def _close_with_run( failure: BaseException | None = None, chain_continues: bool = False, ) -> None: - info = temporalio.workflow.info() - default_id = _ids.default_stream_id(info.workflow_id, info.first_execution_run_id) - state = _workflow._state() - # The default Stream closes with the run even when this run wrote nothing, since an - # earlier run of the chain may have written it. A close of a Stream with no key does - # nothing. - ids = ( - [] - if (default := state.streams.get(default_id)) and default._closed - else [default_id] - ) - for stream in state.streams.values(): - if stream.ref.id == default_id or stream._closed: - continue - if stream.auto_close and stream._used: - ids.append(stream.ref.id) - for stream in state.streams.values(): - if stream.ref.id in ids: - stream._closed = True - # A received ref keeps its own owner. - refs = { - stream_id: state.streams[stream_id]._model_ref() - if stream_id in state.streams - else _owned_ref(stream_id) - for stream_id in ids - } for run_stream in streams_workflow._state().streams.values(): - if run_stream.close_with_run and not run_stream.closed: - run_stream.closed = True - refs[run_stream.ref.id] = run_stream.ref - for ref in refs.values(): + if not run_stream.close_with_run or run_stream.closed: + continue + run_stream.closed = True temporalio.workflow.send_to_plugins( _WorkflowCall( "CloseStream", _convert_in_sandbox( models.CloseStreamRequest( - stream=ref, - result=result, - failure=failure, + stream=run_stream.ref, result=result, failure=failure ) ), when_run_closes=True, @@ -536,7 +507,7 @@ def _convert_in_sandbox(request: Any) -> Any: def _owned_ref(stream_id: str) -> models.StreamRef: - # A Stream that the run creates, or its default Stream, is the run chain's own. + # A Stream that the run chain creates is the chain's own. info = temporalio.workflow.info() return models.StreamRef( id=stream_id, @@ -572,15 +543,10 @@ def __init__( def _activity_headers(self) -> dict[str, temporalio.api.common.v1.Payload]: self._activity_calls += 1 - info = temporalio.workflow.info() # Both ends use the default converter, so a custom user converter can't break the - # headers. + # header. converter = temporalio.converter.PayloadConverter.default - ref = _ids.default_stream_id(info.workflow_id, info.first_execution_run_id) - return { - _REF_HEADER: converter.to_payload(ref), - _KEY_HEADER: converter.to_payload(str(self._activity_calls)), - } + return {_KEY_HEADER: converter.to_payload(str(self._activity_calls))} async def start_system_nexus_operation( self, input: temporalio.worker.StartNexusOperationInput[Any, Any] @@ -653,8 +619,6 @@ async def execute_activity( # The open of a new attempt withdraws the Records of the earlier ones, also # when this attempt publishes nothing. It creates no Stream. stream_ids = [handle.id for handle in _stream_arguments(input.args)] - if streams.default_ref is not None: - stream_ids.append(streams.default_ref.id) for stream_id in dict.fromkeys(stream_ids): await streams.writer(stream_id).withdraw_earlier_attempts() result = await super().execute_activity(input) @@ -740,7 +704,6 @@ def _activity_streams( # A sync Activity has no client, so the plugin gives the backend and the converter. info = temporalio.activity.info() converter = temporalio.converter.PayloadConverter.default - ref = input.headers.get(_REF_HEADER) key = input.headers.get(_KEY_HEADER) key_text = converter.from_payload(key, str) if key is not None else "" epoch = stream_pb.ProducerEpoch(attempt=info.attempt) @@ -779,9 +742,6 @@ def _activity_streams( converter=plugin._converter, namespace=info.namespace, producer=producer, - default_ref=StreamRef(converter.from_payload(ref, str)) - if ref is not None - else None, loop=asyncio.get_running_loop(), backend_name=plugin.backend_name, ) diff --git a/temporalio/contrib/streams/testing.py b/temporalio/contrib/streams/testing.py index 8b98eee4c..c933bc16a 100644 --- a/temporalio/contrib/streams/testing.py +++ b/temporalio/contrib/streams/testing.py @@ -19,7 +19,8 @@ class MemoryStreamPlugin(StreamPlugin): """Keeps Streams in this process, with the same rules as on Redis, for tests. Put it on the client, as with ``RedisStreamPlugin``. Every Worker and client that shares one - plugin shares its Streams, and they go when the process ends. + plugin shares its Streams, and they go when the process ends. It registers the backend + name ``memory``, so tests swap the Store with no change to Workflow code. .. warning:: This API is experimental. diff --git a/tests/contrib/streams/fake_extension.py b/tests/contrib/streams/fake_extension.py index 31eaeeaf4..9704944a5 100644 --- a/tests/contrib/streams/fake_extension.py +++ b/tests/contrib/streams/fake_extension.py @@ -76,6 +76,25 @@ async def call_async(self, method: str, input: bytes) -> bytes: writer_id = len(self.writers) + 1 self.writers[writer_id] = (open_writer.stream.id, open_writer.producer) return ext_pb.OpenWriterResponse(writer_id=writer_id).SerializeToString() + if method == "WriteRecords": + write = ext_pb.WriteRecordsRequest.FromString(input) + stream_id, producer = self.writers[write.writer_id] + offsets = [ + self._append( + stream_id, + stream_pb.StreamRecord( + kind=stream_pb.StreamRecordKind.STREAM_RECORD_KIND_DATA, + topic=record.topic, + body=record.body, + producer=producer, + ), + ) + for record in write.records + ] + return ext_pb.WriteRecordsResponse( + first_offset=offsets[0] if offsets else 0, + last_offset=offsets[-1] if offsets else 0, + ).SerializeToString() if method in ("FlushWriter", "CloseWriter"): return ext_pb.FlushWriterResponse( last_offset=self.next_offset - 1 diff --git a/tests/contrib/streams/test_activity_handles.py b/tests/contrib/streams/test_activity_handles.py index 450f04c75..8328dac03 100644 --- a/tests/contrib/streams/test_activity_handles.py +++ b/tests/contrib/streams/test_activity_handles.py @@ -44,13 +44,11 @@ class SyncInput: stream: streams.SyncActivityStream[Events] -async def read(handle: WorkflowHandle[Any, Any], name: str = "events") -> list[object]: - stream = streams.for_handle( - handle, name=name, wait_for_create=timedelta(seconds=20) - ) - stream.topic("turns", type=Turn) - stream.topic("tokens", type=str) - return [item async for item in stream.subscribe()] +async def read( + client: Client, handle: WorkflowHandle[Any, Any], name: str = "events" +) -> list[object]: + stream = streams.StreamClient(client).workflow_stream(handle, Events, name=name) + return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] def records(items: list[object]) -> list[streams.Record[Any]]: @@ -121,7 +119,7 @@ async def test_activity_publishes_and_returns_offsets(stream_client: Client): 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) + items = await read(stream_client, handle) data = records(items) assert [record.data for record in data] == [ "hi", @@ -176,7 +174,7 @@ async def test_retry_withdraws_a_received_stream(stream_client: Client): Retrying.run, key, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) await handle.result() - items = await read(handle) + items = await read(stream_client, handle) superseded = [item for item in items if isinstance(item, streams.Superseded)] # The second attempt writes nothing, and its start still withdraws the first one's Record. assert len(superseded) == 1 @@ -225,7 +223,7 @@ async def test_activity_closes_and_checks_the_backend(stream_client: Client): Closing.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) assert await handle.result() == "stream_unsupported" - items = await read(handle) + items = await read(stream_client, handle) assert [r.data for r in records(items)] == ["last"] assert isinstance(items[-1], streams.Closed) assert items[-1].result == "closed by activity" @@ -302,7 +300,7 @@ async def test_run_buffer_goes_out_before_a_local_activity(stream_client: Client LocalOrder.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) assert await handle.result() == ["before"] - items = await read(handle) + items = await read(stream_client, handle) written = records(items) assert [record.data for record in written] == ["before", "local", "after"] producer = written[1].producer @@ -402,8 +400,8 @@ async def test_attempt_end_flushes_and_retry_finds_nested_streams( ) # 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") + items = await read(stream_client, handle) + other = await read(stream_client, handle, "other") for read_items in (items, other): superseded = [i for i in read_items if isinstance(i, streams.Superseded)] withdrawn = [r.data for r in records(read_items) if superseded[0].withdraws(r)] diff --git a/tests/contrib/streams/test_facade.py b/tests/contrib/streams/test_facade.py index 43b38d571..b3c178e1c 100644 --- a/tests/contrib/streams/test_facade.py +++ b/tests/contrib/streams/test_facade.py @@ -1,23 +1,14 @@ import asyncio -import json -import uuid from dataclasses import dataclass from datetime import timedelta -from typing import Any -import google.protobuf.json_format import pytest -import temporalio.api.common.v1 import temporalio.api.streamservice.v1 as stream_pb -import temporalio.converter -from temporalio import workflow from temporalio.client import Client from temporalio.common import RawValue from temporalio.contrib import streams from temporalio.contrib.streams import _errors, _ids -from temporalio.contrib.streams._stream import _Writable -from temporalio.contrib.streams._writer import BatchOptions, StreamWriter from temporalio.contrib.streams.proto import v1 as ext_pb from temporalio.contrib.streams.system import models from temporalio.exceptions import ApplicationError @@ -25,7 +16,6 @@ from temporalio.service import RPCStatusCode from temporalio.worker import ExtensionError from tests.contrib.streams.fake_extension import FakeExtension -from tests.helpers import new_worker @dataclass @@ -44,7 +34,6 @@ def fake_client(client: Client) -> tuple[Client, FakeExtension]: def test_ids_match_the_crate(): - assert _ids.default_stream_id("agent-1", "run-1") == "agent-1:run-1" assert ( _ids.scoped_stream_id("a:b/c", "r%1", "é name") == "a%3Ab%2Fc:r%251:%C3%A9%20name" @@ -90,14 +79,19 @@ def producer(attempt: int) -> models.Producer: assert not item.withdraws(new) +@streams.defn(name="events") +class Events: + turns = streams.topic(Turn) + tokens = streams.topic(str) + + async def test_client_writes_reads_and_closes(client: Client): client, extension = fake_client(client) - stream = await streams.create_stream(client, "chat-42") - turns = stream.topic("turns", type=Turn) - turns.publish(Turn("model")) - stream.topic("tokens", type=str).publish("hello", force_flush=True) - await stream.flush() - assert await stream.get_offset() == 3 + stream = await streams.StreamClient(client).create(Events, id="chat-42") + first = await stream.topic(Events.turns).publish(Turn("model")) + second = await stream.topic("tokens").publish("hello") + assert second == first + 1 + assert await stream.get_offset() == second + 1 await stream.close(result="bye") items = [item async for item in stream.subscribe()] assert [type(item) for item in items] == [ @@ -105,72 +99,41 @@ async def test_client_writes_reads_and_closes(client: Client): streams.Record, streams.Closed, ] - first, second, closed = items - assert isinstance(first, streams.Record) and first.data == Turn("model") - assert isinstance(second, streams.Record) and second.data == "hello" - assert first.producer is not None - assert first.producer.kind == stream_pb.ProducerKind.PRODUCER_KIND_CLIENT + turn, token, closed = items + assert isinstance(turn, streams.Record) and turn.data == Turn("model") + assert isinstance(token, streams.Record) and token.data == "hello" + assert turn.producer is not None + assert turn.producer.kind == stream_pb.ProducerKind.PRODUCER_KIND_CLIENT assert isinstance(closed, streams.Closed) and closed.result == "bye" assert extension.namespaces == {client.namespace} - only_tokens = [ - item async for item in stream.topic("tokens", type=RawValue).subscribe() - ] - assert isinstance(only_tokens[0], streams.Record) - assert isinstance(only_tokens[0].data, RawValue) - resumed = [item async for item in stream.subscribe(from_offset=second.offset)] - assert [item.offset for item in resumed] == [second.offset, closed.offset] + raw = stream.topic("raw", type=RawValue) + assert raw.type is RawValue + resumed = [item async for item in stream.subscribe(from_offset=token.offset)] + assert [item.offset for item in resumed] == [token.offset, closed.offset] async def test_read_waits_for_create(client: Client): client, _ = fake_client(client) + stream_client = streams.StreamClient(client) handle = client.get_workflow_handle("wf-1", first_execution_run_id="run-1") + stream = stream_client.workflow_stream(handle, Events) with pytest.raises(streams.StreamNotFoundError): - async for _ in streams.for_handle(handle).subscribe(): + async for _ in stream.subscribe(): pass - waiting = streams.for_handle(handle, wait_for_create=timedelta(seconds=5)) async def create_later() -> None: - created = await streams.create_stream( - client, _ids.default_stream_id("wf-1", "run-1") - ) - created.topic("t", type=int).publish(5) + created = await stream_client.create(Events, id=stream.id) + await created.topic(Events.tokens).publish("five") await created.close() task = asyncio.create_task(create_later()) - items = [item async for item in waiting.subscribe()] + items = [item async for item in stream.subscribe(create_wait=timedelta(seconds=5))] await task - assert isinstance(items[0], streams.Record) and items[0].data == 5 + assert isinstance(items[0], streams.Record) and items[0].data == "five" assert isinstance(items[-1], streams.Closed) -def test_client_without_plugin_has_no_transport(client: Client): - handle = client.get_workflow_handle("wf-1", first_execution_run_id="run-1") - with pytest.raises(streams.StreamTransportMissingError): - streams.for_handle(handle) - - -@workflow.defn -class PublishWithoutPlugin: - @workflow.run - async def run(self) -> str: - try: - streams.workflow_stream().topic("t", type=str).publish("x") - except streams.StreamTransportMissingError as err: - return type(err).__name__ - return "published" - - -async def test_workflow_without_plugin_has_no_transport(client: Client): - async with new_worker(client, PublishWithoutPlugin) as worker: - result = await client.execute_workflow( - PublishWithoutPlugin.run, - id=f"wf-{uuid.uuid4()}", - task_queue=worker.task_queue, - ) - assert result == "StreamTransportMissingError" - - def test_stream_errors_fail_the_workflow_not_the_task(): err = streams.StreamClosedError("closed", reason="stream_closed") assert isinstance(err, ApplicationError) @@ -183,17 +146,13 @@ def test_stream_errors_fail_the_workflow_not_the_task(): assert mapped.code == RPCStatusCode.FAILED_PRECONDITION -async def test_each_client_writer_is_its_own_producer(client: Client): +async def test_each_client_handle_is_its_own_producer(client: Client): client, _ = fake_client(client) - stream = await streams.create_stream(client, "two-writers") - other = streams.for_handle( - client.get_workflow_handle("unused", first_execution_run_id="run") - ) - other._id = stream.ref.id - stream.topic("t", type=int).publish(1) - other.topic("t", type=int).publish(2) - await stream.flush() - await other.flush() + stream_client = streams.StreamClient(client) + stream = await stream_client.create(Events, id="two-writers") + other = stream_client.get_handle(Events, id="two-writers") + await stream.topic(Events.tokens).publish("1") + await other.topic(Events.tokens).publish("2") items = [item async for item in stream.subscribe(wait_new_records=False)] producers = { item.producer.id @@ -203,179 +162,9 @@ async def test_each_client_writer_is_its_own_producer(client: Client): assert len(producers) == 2 -async def test_for_handle_finds_the_first_run_with_one_describe(client: Client): - async with new_worker(client, PublishWithoutPlugin) as worker: - handle = await client.start_workflow( - PublishWithoutPlugin.run, - id=f"wf-{uuid.uuid4()}", - task_queue=worker.task_queue, - ) - await handle.result() - fake, _ = fake_client(client) - first_run = handle.first_execution_run_id - assert first_run - await streams.create_stream(fake, _ids.default_stream_id(handle.id, first_run)) - by_id = streams.for_handle(fake.get_workflow_handle(handle.id)) - with pytest.raises(RuntimeError): - by_id.ref - assert await by_id.get_offset() == 0 - assert by_id.ref.id == _ids.default_stream_id(handle.id, first_run) - - -async def test_writer_publishes_from_another_thread(client: Client): - extension = FakeExtension() - backend = streams.StreamBackend(extension) - await backend.call_async( - "CreateStream", - stream_pb.CreateStreamRequest(stream_id="threads"), - stream_pb.CreateStreamResponse, - ) - - async def stream_id() -> str: - return "threads" - - writer = StreamWriter( - backend=backend, - converter=client.data_converter, - namespace="default", - stream_id=stream_id, - producer=stream_pb.Producer( - kind=stream_pb.ProducerKind.PRODUCER_KIND_CLIENT, id="client:a:b" - ), - batch=BatchOptions(), - create_first=False, - loop=asyncio.get_running_loop(), - ) - await asyncio.to_thread(writer.publish, "t", "from a thread", False) - await writer.flush() - assert len(extension.streams["threads"].records) == 1 - - -async def test_for_ref_opens_a_client_stream(client: Client): - client, _ = fake_client(client) - created = await streams.create_stream(client, "by-ref") - created.topic("t", type=int).publish(7) - await created.flush() - opened = streams.for_ref(client, created.ref) - reader = opened.topic("t", type=int).subscribe(wait_new_records=False) - first = await reader.__anext__() - await reader.aclose() - assert isinstance(first, streams.Record) and first.data == 7 - - -def test_topic_types_check(): - untyped: streams.Topic[Any] = streams.Topic(_NoWrites(), "t", None) - typed: streams.Topic[int] = streams.Topic(_NoWrites(), "t", int) - assert untyped.type is None and typed.type is int - - -class _NoWrites(_Writable): - def _publish(self, topic: str, value: Any, force_flush: bool) -> None: - raise AssertionError("no writes") - - -async def test_create_wait_ends_with_the_run(client: Client): - async with new_worker(client, PublishWithoutPlugin) as worker: - handle = await client.start_workflow( - PublishWithoutPlugin.run, - id=f"wf-{uuid.uuid4()}", - task_queue=worker.task_queue, - ) - await handle.result() - fake, _ = fake_client(client) - stream = streams.for_handle( - fake.get_workflow_handle( - handle.id, first_execution_run_id=handle.first_execution_run_id - ), - wait_for_create=timedelta(minutes=5), - end_with_run=True, - ) - start = asyncio.get_running_loop().time() - assert [item async for item in stream.subscribe()] == [] - assert asyncio.get_running_loop().time() - start < 30 - - -async def test_read_default_type_for_untyped_topics(client: Client): - client, _ = fake_client(client) - stream = await streams.create_stream(client, "default-type") - stream.topic("turns").publish(Turn("model")) - stream.topic("tokens").publish("hi") - await stream.flush() - reader = streams.for_ref(client, stream.ref) - reader.topic("tokens", type=RawValue) - items = [item async for item in reader.subscribe(type=Turn, wait_new_records=False)] - first, second = items - assert isinstance(first, streams.Record) and first.data == Turn("model") - # A typed Topic keeps its own type. - assert isinstance(second, streams.Record) and isinstance(second.data, RawValue) - - -async def test_read_by_topic_prefix(client: Client): - client, _ = fake_client(client) - stream = await streams.create_stream(client, "prefixes") - for topic in ("agent/a", "agent/b", "other", "exact"): - stream.topic(topic).publish(topic) - await stream.flush() - read = [ - item.topic - async for item in stream.subscribe( - topics=[stream.topic("exact")], - topic_prefixes=["agent/"], - wait_new_records=False, - ) - if isinstance(item, streams.Record) - ] - assert read == ["agent/a", "agent/b", "exact"] - - -def test_ref_without_owner_cannot_end_with_run(client: Client): - client, _ = fake_client(client) - with pytest.raises(ValueError): - streams.for_ref(client, streams.StreamRef("standalone"), end_with_run=True) - - -def test_old_refs_still_decode(): - converter = temporalio.converter.DataConverter.default.payload_converter - payload = converter.to_payload({"id": "old"}) - ref = converter.from_payload(payload, streams.StreamRef) - assert ref == streams.StreamRef("old") - assert ref.workflow_id is None and ref.first_execution_run_id is None - - -@pytest.mark.parametrize( - "ref", - [ - streams.StreamRef("only-id"), - streams.StreamRef("wf:run", workflow_id="wf", first_execution_run_id="run"), - ], -) -def test_ref_json_is_the_proto_json(ref: streams.StreamRef): - converter = temporalio.converter.DataConverter.default.payload_converter - payload = converter.to_payload(ref) - proto = stream_pb.StreamRef( - id=ref.id, - workflow_id=ref.workflow_id or "", - first_execution_run_id=ref.first_execution_run_id or "", - ) - assert json.loads(payload.data) == json.loads( - google.protobuf.json_format.MessageToJson(proto) - ) - # Another SDK's proto JSON decodes to the same ref. - from_proto = temporalio.api.common.v1.Payload( - metadata={"encoding": b"json/plain"}, - data=google.protobuf.json_format.MessageToJson(proto).encode(), - ) - assert converter.from_payload(from_proto, streams.StreamRef) == ref - - -def test_snake_case_refs_still_decode(): - converter = temporalio.converter.DataConverter.default.payload_converter - payload = converter.to_payload( - {"id": "old", "workflow_id": "wf", "first_execution_run_id": "run"} - ) - assert converter.from_payload(payload, streams.StreamRef) == streams.StreamRef( - "old", workflow_id="wf", first_execution_run_id="run" - ) +def test_workflow_stream_has_no_reads(): + assert not hasattr(streams.WorkflowStream, "subscribe") + assert not hasattr(streams.WorkflowStream, "get_offset") def test_unsupported_capability_maps_to_its_error(): diff --git a/tests/contrib/streams/test_pydantic.py b/tests/contrib/streams/test_pydantic.py index 62f4cc36e..29c11e422 100644 --- a/tests/contrib/streams/test_pydantic.py +++ b/tests/contrib/streams/test_pydantic.py @@ -1,76 +1,73 @@ -import json import uuid +from datetime import timedelta import pydantic -from temporalio import workflow +from temporalio import activity, workflow from temporalio.client import Client from temporalio.contrib import streams from temporalio.contrib.pydantic import pydantic_data_converter +from temporalio.contrib.streams import workflow as stream_workflow from temporalio.contrib.streams.testing import MemoryStreamPlugin from tests.helpers import new_worker -class Turn(pydantic.BaseModel): - text: str - stream: streams.StreamRef - - -OWNED = streams.StreamRef("wf:run", workflow_id="wf", first_execution_run_id="run") +@streams.defn(name="events") +class Events: + tokens = streams.topic(str) -def test_model_dumps_the_proto_json(): - dumped = json.loads(Turn(text="hi", stream=OWNED).model_dump_json()) - assert dumped["stream"] == { - "id": "wf:run", - "workflowId": "wf", - "firstExecutionRunId": "run", - } - plain = json.loads(Turn(text="hi", stream=streams.StreamRef("s")).model_dump_json()) - assert plain["stream"] == {"id": "s"} +class CallModelInput(pydantic.BaseModel): + prompt: str + stream: streams.ActivityStream[Events] -def test_model_loads_proto_and_snake_case_names(): - for stream in ( - {"id": "wf:run", "workflowId": "wf", "firstExecutionRunId": "run"}, - {"id": "wf:run", "workflow_id": "wf", "first_execution_run_id": "run"}, - ): - loaded = Turn.model_validate_json(json.dumps({"text": "hi", "stream": stream})) - assert loaded.stream == OWNED +class TurnResult(pydantic.BaseModel): + text: str + stream: streams.ClientStream[Events] -def test_payload_converter_writes_the_proto_json(): - converter = pydantic_data_converter.payload_converter - payload = converter.to_payload(Turn(text="hi", stream=OWNED)) - assert json.loads(payload.data)["stream"]["workflowId"] == "wf" - assert converter.from_payload(payload, Turn) == Turn(text="hi", stream=OWNED) +@activity.defn +async def call_model(input: CallModelInput) -> int: + assert input.stream.definition is Events + return await input.stream.topic(Events.tokens).publish(*input.prompt.split()) @workflow.defn class TurnWorkflow: @workflow.run - async def run(self, turn: Turn) -> Turn: - scoped = streams.create_workflow_stream("turn") - scoped.topic("t", type=str).publish(turn.text) - return Turn(text=turn.text.upper(), stream=scoped.ref) + async def run(self, prompt: str) -> TurnResult: + stream = await stream_workflow.create_stream(Events) + await workflow.execute_activity( + call_model, + CallModelInput( + prompt=prompt, stream=streams.ActivityStream.from_handle(stream) + ), + start_to_close_timeout=timedelta(seconds=30), + ) + return TurnResult( + text=prompt.upper(), stream=streams.ClientStream.from_handle(stream) + ) -async def test_workflow_round_trip_on_the_memory_store(client: Client): +async def test_handles_in_models_round_trip_on_the_memory_store(client: Client): config = client.config() config["plugins"] = [MemoryStreamPlugin()] config["data_converter"] = pydantic_data_converter client = Client(**config) - async with new_worker(client, TurnWorkflow) as worker: + async with new_worker(client, TurnWorkflow, activities=[call_model]) as worker: handle = await client.start_workflow( TurnWorkflow.run, - Turn(text="hi", stream=streams.StreamRef("input")), + "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) result = await handle.result() - assert result.text == "HI" - assert result.stream.workflow_id == handle.id - assert result.stream.first_execution_run_id == handle.first_execution_run_id - stream = streams.for_ref(client, result.stream, end_with_run=True) - items = [item async for item in stream.subscribe(type=str)] - assert [i.data for i in items if isinstance(i, streams.Record)] == ["hi"] + 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 = streams.StreamClient(client).get_handle(Events, id=result.stream.id) + items = [item async for item in stream.subscribe()] + assert [i.data for i in items if isinstance(i, streams.Record)] == ["hi", "there"] + assert isinstance(items[-1], streams.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 a6973327f..338fe7024 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -1,5 +1,4 @@ import asyncio -import concurrent.futures import dataclasses import gc import uuid @@ -13,14 +12,12 @@ from temporalio import activity, workflow from temporalio.api.common.v1 import Payload -from temporalio.client import Client, WorkflowFailureError +from temporalio.client import Client, WorkflowFailureError, WorkflowHandle from temporalio.common import RawValue, RetryPolicy from temporalio.contrib import streams +from temporalio.contrib.streams import workflow as stream_workflow from temporalio.contrib.streams.proto import v1 as ext_pb -from temporalio.contrib.streams.redis import ( - StreamPlugin, - _WorkflowInbound, -) +from temporalio.contrib.streams.redis import StreamPlugin, _WorkflowInbound from temporalio.contrib.streams.testing import MemoryStreamPlugin from temporalio.converter import DataConverter, PayloadCodec from temporalio.exceptions import ApplicationError @@ -43,44 +40,65 @@ class Turn: text: str = "" +@streams.defn(name="events") +class Events: + turns = streams.topic(Turn) + tokens = streams.topic(str) + + +@dataclass +class CallModelInput: + prompt: str + stream: streams.ActivityStream[Events] + + +def workflow_stream( + client: Client, handle: WorkflowHandle[Any, Any], name: str | None = None +) -> streams.ClientStream[Events]: + return streams.StreamClient(client).workflow_stream(handle, Events, name=name) + + +async def read_all(stream: streams.ClientStream[Any]) -> list[Any]: + return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] + + +def records(items: list[Any]) -> list[streams.Record[Any]]: + return [item for item in items if isinstance(item, streams.Record)] + + @activity.defn -async def call_model(prompt: str) -> str: - tokens = streams.activity_stream().topic("tokens", type=str) - for word in prompt.split(): - tokens.publish(word) - return prompt.upper() +async def call_model(input: CallModelInput) -> str: + tokens = input.stream.topic(Events.tokens) + for word in input.prompt.split(): + await tokens.publish(word) + return input.prompt.upper() @workflow.defn class Agent: @workflow.run async def run(self, prompt: str) -> str: - turns = streams.workflow_stream().topic("turns", type=Turn) - turns.publish(Turn("model")) + stream = await stream_workflow.create_stream(Events) + turns = stream.topic(Events.turns) + await turns.publish(Turn("model")) reply = await workflow.execute_activity( - call_model, prompt, start_to_close_timeout=timedelta(seconds=30) + call_model, + CallModelInput(prompt, streams.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), ) - turns.publish(Turn("done", reply)) + await turns.publish(Turn("done", reply)) return reply -async def read_all(stream: streams.OutsideStream) -> list[streams.Record | object]: - return [item async for item in stream.subscribe()] - - -async def test_workflow_and_activity_write_the_default_stream(stream_client: Client): +async def test_workflow_and_activity_write_one_stream(stream_client: Client): async with new_worker(stream_client, Agent, activities=[call_model]) as worker: handle = await stream_client.start_workflow( Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) - stream = streams.for_handle(handle, wait_for_create=timedelta(seconds=20)) - turns = stream.topic("turns", type=Turn) - stream.topic("tokens", type=str) + stream = workflow_stream(stream_client, handle) items = await read_all(stream) assert await handle.result() == "HI THERE" - data = [ - (item.topic, item.data) for item in items if isinstance(item, streams.Record) - ] + data = [(item.topic, item.data) for item in records(items)] assert data == [ ("turns", Turn("model")), ("tokens", "hi"), @@ -89,7 +107,7 @@ async def test_workflow_and_activity_write_the_default_stream(stream_client: Cli ] closed = items[-1] assert isinstance(closed, streams.Closed) and closed.result == "HI THERE" - only_turns = [item async for item in turns.subscribe()] + only_turns = [item async for item in stream.subscribe(Events.turns)] assert [type(item) for item in only_turns] == [ streams.Record, streams.Record, @@ -97,105 +115,17 @@ async def test_workflow_and_activity_write_the_default_stream(stream_client: Cli ] -@workflow.defn -class Chat: - @workflow.run - async def run(self, turns: int) -> None: - index = streams.workflow_stream().topic("turns", type=int) - for n in range(turns): - with streams.create_workflow_stream(f"turn-{n}") as turn: - index.publish(n) - turn.topic("tokens", type=str).publish(f"token-{n}") - await workflow.sleep(timedelta(milliseconds=10)) - - -async def test_scoped_streams_close_at_the_end_of_their_block(stream_client: Client): - async with new_worker(stream_client, Chat) as worker: - handle = await stream_client.start_workflow( - Chat.run, 2, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue - ) - await handle.result() - turn = streams.for_handle(handle, name="turn-1") - items = [item async for item in turn.topic("tokens", type=str).subscribe()] - assert isinstance(items[0], streams.Record) and items[0].data == "token-1" - assert isinstance(items[-1], streams.Closed) - - -@workflow.defn -class WriteToRef: - @workflow.run - async def run(self, ref: streams.StreamRef) -> None: - streams.workflow_stream(ref).topic("t", type=str).publish("from workflow") - - -async def test_client_stream_takes_workflow_writes(stream_client: Client): - stream = await streams.create_stream(stream_client, f"chat-{uuid.uuid4()}") - stream.topic("t", type=str).publish("from client") - await stream.flush() - async with new_worker(stream_client, WriteToRef) as worker: - await stream_client.execute_workflow( - WriteToRef.run, - stream.ref, - id=f"wf-{uuid.uuid4()}", - task_queue=worker.task_queue, - ) - await stream.close() - items = [item async for item in stream.subscribe()] - data = [item.data for item in items if isinstance(item, streams.Record)] - assert data == ["from client", "from workflow"] - assert isinstance(items[-1], streams.Closed) - - attempts: dict[str, int] = {} -@activity.defn -async def flaky(key: str) -> None: - attempts[key] = attempts.get(key, 0) + 1 - tokens = streams.activity_stream().topic("tokens", type=str) - tokens.publish(f"attempt-{attempts[key]}") - await streams.activity_stream().flush() - if attempts[key] == 1: - raise RuntimeError("first attempt fails") - - -@workflow.defn -class Retrying: - @workflow.run - async def run(self, key: str) -> None: - streams.workflow_stream().topic("turns", type=str).publish("start") - await workflow.execute_activity( - flaky, - key, - start_to_close_timeout=timedelta(seconds=30), - retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), - ) - - -async def test_activity_retry_withdraws_the_failed_attempt(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_all(streams.for_handle(handle)) - records = [item for item in items if isinstance(item, streams.Record)] - superseded = [item for item in items if isinstance(item, streams.Superseded)] - assert [record.data for record in records] == ["start", "attempt-1", "attempt-2"] - assert len(superseded) == 1 - withdrawn = [record.data for record in records if superseded[0].withdraws(record)] - assert withdrawn == ["attempt-1"] - - @workflow.defn class PublishThenFailTask: @workflow.run async def run(self, key: str) -> None: - turns = streams.workflow_stream().topic("turns", type=str) - turns.publish("before") + turns = (await stream_workflow.create_stream(Events)).topic(Events.tokens) + await turns.publish("before") await workflow.sleep(timedelta(milliseconds=10)) - turns.publish("after") + await turns.publish("after") if attempts.setdefault(key, 0) == 0: attempts[key] = 1 raise RuntimeError("fail this Workflow Task once") @@ -216,34 +146,13 @@ async def test_failed_workflow_task_is_withdrawn(stream_client: Client): task_queue=worker.task_queue, ) await handle.result() - items = await read_all(streams.for_handle(handle)) - records = [item for item in items if isinstance(item, streams.Record)] + items = await read_all(workflow_stream(stream_client, handle)) # The failed task wrote nothing, so the reader gets each Record once. - assert [record.data for record in records] == ["before", "after"] + assert [record.data for record in records(items)] == ["before", "after"] assert not any(isinstance(item, streams.Superseded) for item in items) assert isinstance(items[-1], streams.Closed) -@workflow.defn -class Reading: - @workflow.run - async def run(self) -> str: - try: - streams.workflow_stream().topic("t", type=str).subscribe() - except TypeError as err: - return str(err) - return "read" - - -async def test_workflow_code_can_not_read(stream_client: Client): - async with new_worker(stream_client, Reading) as worker: - result = await stream_client.execute_workflow( - Reading.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue - ) - assert "can't read" in result - await asyncio.sleep(0) - - @activity.defn async def noop() -> None: pass @@ -273,48 +182,13 @@ async def test_plugin_leaves_the_workflow_random_alone(stream_client: Client): ) -retries: dict[str, int] = {} - - -@activity.defn -async def fails_once(key: str) -> None: - retries[key] = retries.get(key, 0) + 1 - if retries[key] == 1: - raise RuntimeError("first attempt fails") - - -@workflow.defn -class NoStreams: - @workflow.run - async def run(self, key: str) -> None: - await workflow.execute_activity( - fails_once, - key, - start_to_close_timeout=timedelta(seconds=30), - retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), - ) - - -async def test_workflow_without_streams_creates_none(stream_client: Client): - async with new_worker(stream_client, NoStreams, activities=[fails_once]) as w: - handle = await stream_client.start_workflow( - NoStreams.run, - str(uuid.uuid4()), - id=f"wf-{uuid.uuid4()}", - task_queue=w.task_queue, - ) - await handle.result() - # The retry withdrew nothing and the run's close found no Stream, so none exists. - with pytest.raises(streams.StreamNotFoundError): - await streams.for_handle(handle).get_offset() - - @workflow.defn class WritesThenContinues: @workflow.run async def run(self, write: bool) -> str: if write: - streams.workflow_stream().topic("t", type=str).publish("first run") + stream = await stream_workflow.create_stream(Events) + await stream.topic(Events.tokens).publish("first run") workflow.continue_as_new(False) return "second run" @@ -328,10 +202,8 @@ async def test_last_run_closes_a_stream_it_did_not_write(stream_client: Client): task_queue=worker.task_queue, ) assert await handle.result() == "second run" - stream = streams.for_handle(handle) - stream.topic("t", type=str) - items = await read_all(stream) - assert [i.data for i in items if isinstance(i, streams.Record)] == ["first run"] + items = await read_all(workflow_stream(stream_client, handle)) + assert [i.data for i in records(items)] == ["first run"] closed = items[-1] assert isinstance(closed, streams.Closed) and closed.result == "second run" @@ -340,8 +212,10 @@ async def test_last_run_closes_a_stream_it_did_not_write(stream_client: Client): class CancelledBlock: @workflow.run async def run(self) -> None: - with streams.create_workflow_stream("turn") as turn: - turn.topic("t", type=str).publish("started") + async with await stream_workflow.create_stream( + Events, name="turn", close_with_run=False + ) as turn: + await turn.topic(Events.tokens).publish("started") await workflow.wait_condition(lambda: False) @@ -350,12 +224,10 @@ async def test_cancelled_block_closes_with_a_failure(stream_client: Client): handle = await stream_client.start_workflow( CancelledBlock.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) - turn = streams.for_handle( - handle, name="turn", wait_for_create=timedelta(seconds=20) - ) - turn.topic("t", type=str) - first = await anext(aiter(turn.subscribe())) - assert isinstance(first, streams.Record) + turn = workflow_stream(stream_client, handle, "turn") + reader = turn.subscribe(create_wait=timedelta(seconds=20)) + assert isinstance(await anext(reader), streams.Record) + await reader.aclose() await handle.cancel() items = await read_all(turn) closed = items[-1] @@ -363,46 +235,21 @@ async def test_cancelled_block_closes_with_a_failure(stream_client: Client): assert closed.failure is not None -@activity.defn -def sync_tokens(count: int) -> None: - tokens = streams.activity_stream().topic("tokens", type=int) - for n in range(count): - tokens.publish(n) - - -@workflow.defn -class SyncWriter: - @workflow.run - async def run(self) -> None: - await workflow.execute_activity( - sync_tokens, 3, start_to_close_timeout=timedelta(seconds=30) - ) - - -async def test_sync_activity_publishes(stream_client: Client): - async with new_worker( - stream_client, - SyncWriter, - activities=[sync_tokens], - activity_executor=concurrent.futures.ThreadPoolExecutor(2), - ) as worker: - handle = await stream_client.start_workflow( - SyncWriter.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue - ) - await handle.result() - stream = streams.for_handle(handle) - stream.topic("tokens", type=int) - items = await read_all(stream) - assert [i.data for i in items if isinstance(i, streams.Record)] == [0, 1, 2] +@streams.defn(name="plain") +class Plain: + t = streams.topic(str) @workflow.defn class PlainWriter: @workflow.run async def run(self) -> str: - streams.workflow_stream().topic("t", type=str).publish("plain") - with streams.create_workflow_stream("scoped") as scoped: - scoped.topic("t", type=str).publish("scoped") + stream = await stream_workflow.create_stream(Plain) + await stream.topic(Plain.t).publish("plain") + async with await stream_workflow.create_stream( + Plain, name="scoped", close_with_run=False + ) as scoped: + await scoped.topic(Plain.t).publish("scoped") return "done" @@ -420,42 +267,42 @@ async def test_temporal_only_passthrough_keeps_records(stream_client: Client): PlainWriter.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) await handle.result() - items = await read_all(streams.for_handle(handle)) - scoped = await read_all(streams.for_handle(handle, name="scoped")) - assert [i.data for i in items if isinstance(i, streams.Record)] == ["plain"] + plain = streams.StreamClient(stream_client) + items = await read_all(plain.workflow_stream(handle, Plain)) + scoped = await read_all(plain.workflow_stream(handle, Plain, name="scoped")) + assert [i.data for i in records(items)] == ["plain"] assert isinstance(items[-1], streams.Closed) and items[-1].result == "done" - assert [i.data for i in scoped if isinstance(i, streams.Record)] == ["scoped"] + assert [i.data for i in records(scoped)] == ["scoped"] assert isinstance(scoped[-1], streams.Closed) @activity.defn -async def local_tokens(count: int) -> None: - tokens = streams.activity_stream().topic("tokens", type=int) - for n in range(count): - tokens.publish(n) +async def local_tokens(input: CallModelInput) -> None: + await input.stream.topic(Events.tokens).publish(*input.prompt.split()) @workflow.defn class LocalWriter: @workflow.run async def run(self) -> None: + stream = await stream_workflow.create_stream(Events) await workflow.execute_local_activity( - local_tokens, 2, start_to_close_timeout=timedelta(seconds=30) + local_tokens, + CallModelInput("0 1", streams.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), ) -async def test_local_activity_writes_the_default_stream(stream_client: Client): +async def test_local_activity_writes_with_the_task_epoch(stream_client: Client): async with new_worker(stream_client, LocalWriter, activities=[local_tokens]) as w: handle = await stream_client.start_workflow( LocalWriter.run, id=f"wf-{uuid.uuid4()}", task_queue=w.task_queue ) await handle.result() - stream = streams.for_handle(handle) - stream.topic("tokens", type=int) - items = await read_all(stream) - records = [item for item in items if isinstance(item, streams.Record)] - assert [record.data for record in records] == [0, 1] - producer = records[0].producer + items = await read_all(workflow_stream(stream_client, handle)) + written = records(items) + assert [record.data for record in written] == ["0", "1"] + producer = written[0].producer assert producer is not None and producer.id.startswith("la:") assert ( producer.epoch is not None and (producer.epoch.task_started_event_id or 0) > 0 @@ -466,9 +313,9 @@ async def test_local_activity_writes_the_default_stream(stream_client: Client): class FailingRun: @workflow.run async def run(self) -> str: - turns = streams.workflow_stream().topic("turns", type=str) + turns = (await stream_workflow.create_stream(Events)).topic(Events.tokens) info = workflow.info() - turns.publish(f"run-{info.attempt}") + await turns.publish(f"run-{info.attempt}") if info.attempt == 1: raise ApplicationError("first run fails") return "done" @@ -485,14 +332,12 @@ async def test_retry_run_withdraws_the_failed_run(stream_client: Client): ), ) assert await handle.result() == "done" - stream = streams.for_handle(handle) - stream.topic("turns", type=str) - items = await read_all(stream) - records = [item for item in items if isinstance(item, streams.Record)] + items = await read_all(workflow_stream(stream_client, handle)) + written = records(items) superseded = [item for item in items if isinstance(item, streams.Superseded)] - assert [record.data for record in records] == ["run-1", "run-2"] + assert [record.data for record in written] == ["run-1", "run-2"] assert len(superseded) == 1 - assert [r.data for r in records if superseded[0].withdraws(r)] == ["run-1"] + assert [r.data for r in written if superseded[0].withdraws(r)] == ["run-1"] # The failed run had retries left, so only the last run closes the Stream. closed = items[-1] assert isinstance(closed, streams.Closed) and closed.result == "done" @@ -530,18 +375,22 @@ async def test_records_go_through_the_payload_codec(stream_client: Client): Agent.run, "a b", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) await handle.result() - stream = streams.for_handle(handle) - stream.topic("turns", type=Turn) - stream.topic("tokens", type=str) + stream = workflow_stream(codec_client, handle) decoded = await read_all(stream) + # A Stream with no definition reads the raw Payloads. raw = [ item - async for item in streams.for_handle(handle) + async for item in streams.StreamClient(codec_client) + .get_handle(Plain, id=stream.id) .topic("tokens", type=RawValue) .subscribe(wait_new_records=False) ] - data = [item.data for item in decoded if isinstance(item, streams.Record)] - assert data == [Turn("model"), "a", "b", Turn("done", "A B")] + assert [item.data for item in records(decoded)] == [ + Turn("model"), + "a", + "b", + Turn("done", "A B"), + ] first = raw[0] assert isinstance(first, streams.Record) assert first.data.payload.metadata["encoding"] == b"json/plain" @@ -554,12 +403,14 @@ def __init__(self) -> None: @workflow.run async def run(self) -> str: - turns = streams.workflow_stream().topic("turns", type=str) - turns.publish("before restart") - with streams.create_workflow_stream("scoped") as scoped: - scoped.topic("t", type=str).publish("scoped") + turns = (await stream_workflow.create_stream(Events)).topic(Events.tokens) + await turns.publish("before restart") + async with await stream_workflow.create_stream( + Events, name="scoped", close_with_run=False + ) as scoped: + await scoped.topic(Events.tokens).publish("scoped") await workflow.wait_condition(lambda: self.go) - turns.publish("after restart") + await turns.publish("after restart") return "done" @workflow.signal @@ -573,19 +424,20 @@ async def test_new_worker_replays_and_writes_each_record_once(stream_client: Cli handle = await stream_client.start_workflow( WaitsForSignal.run, id=f"wf-{uuid.uuid4()}", task_queue=task_queue ) - first = streams.for_handle(handle, wait_for_create=timedelta(seconds=20)) - assert isinstance(await anext(aiter(first.subscribe())), streams.Record) + first = workflow_stream(stream_client, handle) + reader = first.subscribe(create_wait=timedelta(seconds=20)) + assert isinstance(await anext(reader), streams.Record) + await reader.aclose() # The new Worker has no cache, so it replays the run and sends the creates and the close # of the scoped Stream again. async with new_worker(stream_client, WaitsForSignal, task_queue=task_queue): await handle.signal(WaitsForSignal.proceed) assert await handle.result() == "done" - items = await read_all(streams.for_handle(handle)) - scoped = await read_all(streams.for_handle(handle, name="scoped")) - data = [i.data for i in items if isinstance(i, streams.Record)] - assert data == ["before restart", "after restart"] + items = await read_all(first) + scoped = await read_all(workflow_stream(stream_client, handle, "scoped")) + assert [i.data for i in records(items)] == ["before restart", "after restart"] assert isinstance(items[-1], streams.Closed) - assert [i.data for i in scoped if isinstance(i, streams.Record)] == ["scoped"] + assert [i.data for i in records(scoped)] == ["scoped"] assert [type(i) for i in scoped].count(streams.Closed) == 1 @@ -613,7 +465,8 @@ async def test_plugin_replays_a_history_made_without_it( class QuickWriter: @workflow.run async def run(self) -> str: - streams.workflow_stream().topic("t", type=str).publish("quick") + stream = await stream_workflow.create_stream(Events) + await stream.topic(Events.tokens).publish("quick") return "done" @@ -626,31 +479,37 @@ async def test_worker_shutdown_drains_closes(stream_client: Client): # Calls after the Worker exits still work, and the drain waited for the run's close. items = [ item - async for item in streams.for_handle(handle).subscribe(wait_new_records=False) + async for item in workflow_stream(stream_client, handle).subscribe( + wait_new_records=False + ) ] assert isinstance(items[-1], streams.Closed) and items[-1].result == "done" -async def test_plugin_shutdown_flushes_client_writers(stream_client: Client): +async def test_plugin_shutdown_after_client_writes(stream_client: Client): [plugin] = [ p for p in stream_client.config()["plugins"] if isinstance(p, StreamPlugin) ] - stream = await streams.create_stream(stream_client, f"shutdown-{uuid.uuid4()}") - stream.topic("t", type=str).publish("unflushed") + stream = await streams.StreamClient(stream_client).create( + Events, id=f"shutdown-{uuid.uuid4()}" + ) + async with streams.BufferedPublisher(stream.topic(Events.tokens)) as tokens: + tokens.publish("buffered") await plugin.shutdown() written = [ item.data async for item in stream.subscribe(wait_new_records=False) if isinstance(item, streams.Record) ] - assert written == ["unflushed"] + assert written == ["buffered"] @workflow.defn class WaitsForever: @workflow.run async def run(self) -> None: - streams.workflow_stream().topic("t", type=str).publish("waiting") + stream = await stream_workflow.create_stream(Events) + await stream.topic(Events.tokens).publish("waiting") await workflow.wait_condition(lambda: False) @@ -694,88 +553,17 @@ async def test_eviction_prints_nothing( handle = await stream_client.start_workflow( WaitsForever.run, id=f"wf-{uuid.uuid4()}", task_queue=w.task_queue ) - stream = streams.for_handle(handle, wait_for_create=timedelta(seconds=20)) - assert isinstance(await anext(aiter(stream.subscribe())), streams.Record) + reader = workflow_stream(stream_client, handle).subscribe( + create_wait=timedelta(seconds=20) + ) + assert isinstance(await anext(reader), streams.Record) + await reader.aclose() await handle.terminate() gc.collect() captured = capfd.readouterr() assert "Exception ignored" not in captured.err + captured.out -@workflow.defn -class CreatesEmpty: - @workflow.run - async def run(self) -> None: - streams.create_workflow_stream("empty") - await workflow.sleep(timedelta(milliseconds=10)) - - -async def test_get_offset_of_a_created_empty_stream(stream_client: Client): - async with new_worker(stream_client, CreatesEmpty) as worker: - handle = await stream_client.start_workflow( - CreatesEmpty.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue - ) - await handle.result() - # A created Stream with no Record gives the start offset, so a reader can begin there. - assert await streams.for_handle(handle, name="empty").get_offset() == 0 - # A Stream that nobody created is still an error. - with pytest.raises(streams.StreamNotFoundError): - await streams.for_handle(handle, name="missing").get_offset() - - -async def test_read_by_topic_prefix_on_redis(stream_client: Client): - stream = await streams.create_stream(stream_client, f"prefix-{uuid.uuid4()}") - for topic in ("agent/a", "agent/b", "other", "exact"): - stream.topic(topic, type=str).publish(topic) - await stream.flush() - read = [ - item.topic - async for item in stream.subscribe( - topics=[stream.topic("exact", type=str)], - topic_prefixes=["agent/"], - wait_new_records=False, - ) - if isinstance(item, streams.Record) - ] - assert read == ["agent/a", "agent/b", "exact"] - - -@workflow.defn -class LeavesStreamOpen: - @workflow.run - async def run(self) -> streams.StreamRef: - # Not in a with block, so nothing closes it when the run ends. - open_stream = streams.create_workflow_stream("open") - open_stream.topic("t", type=str).publish("only") - return open_stream.ref - - -async def test_for_ref_ends_with_the_owning_run(stream_client: Client): - async with new_worker(stream_client, LeavesStreamOpen) as worker: - handle = await stream_client.start_workflow( - LeavesStreamOpen.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue - ) - ref = await handle.result() - assert ref.workflow_id == handle.id - assert ref.first_execution_run_id == handle.first_execution_run_id - stream = streams.for_ref(stream_client, ref, end_with_run=True) - items = [item async for item in stream.subscribe(type=str)] - assert [i.data for i in items if isinstance(i, streams.Record)] == ["only"] - # A ref to a Stream that the finished run never created ends at once, not after the - # whole wait. - missing = streams.StreamRef( - ref.id + "-missing", - workflow_id=ref.workflow_id, - first_execution_run_id=ref.first_execution_run_id, - ) - waiting = streams.for_ref( - stream_client, missing, end_with_run=True, wait_for_create=timedelta(minutes=5) - ) - start = asyncio.get_running_loop().time() - assert [item async for item in waiting.subscribe()] == [] - assert asyncio.get_running_loop().time() - start < 30 - - async def test_workflow_calls_carry_the_owner(stream_client: Client): [plugin] = [ p for p in stream_client.config()["plugins"] if isinstance(p, StreamPlugin) @@ -813,7 +601,8 @@ class AlwaysFails: @workflow.run async def run(self, non_retryable: bool) -> None: attempt = workflow.info().attempt - streams.workflow_stream().topic("t", type=str).publish(f"run-{attempt}") + stream = await stream_workflow.create_stream(Events) + await stream.topic(Events.tokens).publish(f"run-{attempt}") raise ApplicationError(f"attempt {attempt} fails", non_retryable=non_retryable) @@ -833,11 +622,8 @@ async def test_final_failure_under_a_retry_policy( ) with pytest.raises(WorkflowFailureError): await handle.result() - reader = streams.for_handle( - handle, end_with_run=True, wait_for_create=timedelta(seconds=20) - ) - items = [item async for item in reader.subscribe(type=str)] - records = [item for item in items if isinstance(item, streams.Record)] + items = await read_all(workflow_stream(stream_client, handle)) + written = records(items) superseded = [item for item in items if isinstance(item, streams.Superseded)] closes = [item for item in items if isinstance(item, streams.Closed)] # Only the run that ends the chain closes the Stream, with its failure, and it is last. @@ -845,16 +631,16 @@ async def test_final_failure_under_a_retry_policy( assert closes[0].failure is not None if non_retryable: # A non-retryable failure ends the chain on attempt 1, so it closes at once. - assert [r.data for r in records] == ["run-1"] + assert [r.data for r in written] == ["run-1"] assert superseded == [] return # Attempts 1 and 2 left the Stream open, so attempts 2 and 3 could write. Each retry run # withdraws the run before it. - assert [r.data for r in records] == ["run-1", "run-2", "run-3"] + assert [r.data for r in written] == ["run-1", "run-2", "run-3"] assert len(superseded) == 2 - assert [r.data for r in records if superseded[0].withdraws(r)] == ["run-1"] - assert [r.data for r in records if superseded[1].withdraws(r)] == ["run-2"] - assert not any(s.withdraws(records[-1]) for s in superseded) + assert [r.data for r in written if superseded[0].withdraws(r)] == ["run-1"] + assert [r.data for r in written if superseded[1].withdraws(r)] == ["run-2"] + assert not any(s.withdraws(written[-1]) for s in superseded) async def test_missing_capability_fails_the_workflow_task(client: Client): @@ -869,3 +655,30 @@ async def test_missing_capability_fails_the_workflow_task(client: Client): ) await assert_task_fail_eventually(handle, message_contains="replica_wait") await handle.terminate() + + +async def test_reader_stops_on_its_own_timer_and_resumes(stream_client: Client): + stream = await streams.StreamClient(stream_client).create( + Events, id=f"timer-{uuid.uuid4()}" + ) + first = await stream.topic(Events.tokens).publish("one") + seen: list[streams.Record[Any]] = [] + + async def read_until_quiet() -> None: + async for item in stream.subscribe(): + assert isinstance(item, streams.Record) + seen.append(item) + + # A read has no idle end, so a reader that wants one sets its own timer. + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(read_until_quiet(), timeout=0.5) + assert [item.offset for item in seen] == [first] + await stream.topic(Events.tokens).publish("two") + resumed = [ + item.data + async for item in stream.subscribe( + from_offset=seen[-1].offset + 1, wait_new_records=False + ) + if isinstance(item, streams.Record) + ] + assert resumed == ["two"] diff --git a/tests/contrib/streams/test_vectors.py b/tests/contrib/streams/test_vectors.py index fbba96cc7..9801ab5ac 100644 --- a/tests/contrib/streams/test_vectors.py +++ b/tests/contrib/streams/test_vectors.py @@ -34,12 +34,12 @@ def test_vendored_protos_match_the_crate(relative: str): def test_stream_ids_match_the_crate(): data: Any = vectors("ids.json") for case in data["streams"]: + # Every Stream has a name, so the crate's cases with no name don't apply. if case["name"] is None: - built = _ids.default_stream_id(case["workflow_id"], case["first_run_id"]) - else: - built = _ids.scoped_stream_id( - case["workflow_id"], case["first_run_id"], case["name"] - ) + continue + built = _ids.scoped_stream_id( + case["workflow_id"], case["first_run_id"], case["name"] + ) assert built == case["stream_id"] diff --git a/tests/contrib/streams/test_workflow_module.py b/tests/contrib/streams/test_workflow_module.py index 1ea1def5a..62aadd51d 100644 --- a/tests/contrib/streams/test_workflow_module.py +++ b/tests/contrib/streams/test_workflow_module.py @@ -33,13 +33,11 @@ class Events: tokens = streams.topic(str) -async def read(handle: WorkflowHandle[Any, Any], name: str = "events") -> list[object]: - stream = streams.for_handle( - handle, name=name, wait_for_create=timedelta(seconds=20) - ) - stream.topic("turns", type=Turn) - stream.topic("tokens", type=str) - return [item async for item in stream.subscribe()] +async def read( + client: Client, handle: WorkflowHandle[Any, Any], name: str = "events" +) -> list[object]: + stream = streams.StreamClient(client).workflow_stream(handle, Events, name=name) + return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] def data(items: Sequence[object]) -> list[object]: @@ -63,7 +61,7 @@ async def test_create_publish_and_close_with_run(stream_client: Client): Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) assert await handle.result() == "HI THERE" - items = await read(handle) + items = await read(stream_client, handle) assert data(items) == [Turn("user", "hi there"), Turn("model"), "hi", "there"] closed = items[-1] assert isinstance(closed, streams.Closed) and closed.result == "HI THERE" @@ -92,8 +90,8 @@ async def test_async_with_closes_the_stream(stream_client: Client): Turns.run, 1, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) await handle.result() - first = await read(handle, "turn-0") - second = await read(handle, "turn-1") + first = await read(stream_client, handle, "turn-0") + second = await read(stream_client, handle, "turn-1") assert data(first) == ["token-0"] assert isinstance(first[-1], streams.Closed) and first[-1].failure is None assert data(second) == ["token-1"] @@ -121,7 +119,7 @@ async def test_continue_as_new_passes_the_stream(stream_client: Client): Continuing.run, None, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) assert await handle.result() == "done" - items = await read(handle) + items = await read(stream_client, handle) assert data(items) == ["run 1", "run 2"] closed = [item for item in items if isinstance(item, streams.Closed)] assert len(closed) == 1 and closed[0].result == "done" @@ -175,7 +173,7 @@ async def test_replay_reproducible_errors_raise_in_workflow_code( "stream_closed", ] assert await handle.query(Mistakes.publish_in_query) == "ReadOnlyContextError" - items = await read(handle) + items = await read(stream_client, handle) assert isinstance(items[-1], streams.Closed) and items[-1].result == "early" @@ -195,7 +193,7 @@ async def test_close_with_run_takes_the_failure(stream_client: Client): ) with pytest.raises(WorkflowFailureError): await handle.result() - items = await read(handle) + items = await read(stream_client, handle) assert data(items) == ["before"] assert isinstance(items[-1], streams.Closed) assert str(items[-1].failure) == "run failed" @@ -218,7 +216,7 @@ async def test_auto_close_closes_a_quiet_stream(stream_client: Client): ) await handle.result() await asyncio.sleep(1.5) - stream = streams.for_handle(handle, name="events") + stream = streams.StreamClient(stream_client).workflow_stream(handle, Events) items = [item async for item in stream.subscribe(wait_new_records=False)] assert data(items) == ["only"] assert isinstance(items[-1], streams.Closed) @@ -274,10 +272,8 @@ async def test_cron_runs_share_an_open_stream(stream_client: Client): task_queue=worker.task_queue, cron_schedule="@every 1s", ) - stream = streams.for_handle( - handle, name="events", wait_for_create=timedelta(seconds=20) - ) - reader = stream.subscribe() + stream = streams.StreamClient(stream_client).workflow_stream(handle, Events) + reader = stream.subscribe(create_wait=timedelta(seconds=20)) first = await reader.__anext__() second = await reader.__anext__() await reader.aclose() @@ -336,7 +332,7 @@ async def test_sync_handlers_publish_and_close_with_no_await(stream_client: Clie await handle.signal(SyncHandlers.token, "a") assert await handle.execute_update(SyncHandlers.finish) == "stream_closed" await handle.result() - items = await read(handle) + items = await read(stream_client, handle) assert data(items) == ["a", "p"] assert isinstance(items[-1], streams.Closed) and items[-1].result == "finished" @@ -363,6 +359,6 @@ async def test_init_creates_and_publishes_with_no_await(stream_client: Client): task_queue=worker.task_queue, ) assert await handle.result() == "done" - items = await read(handle) + items = await read(stream_client, handle) assert data(items) == ["hi", "HI"] assert isinstance(items[-1], streams.Closed) and items[-1].result == "done"