diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 60d35a56fc..344f516ba0 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -4,6 +4,7 @@ import asyncio import collections +import contextlib import contextvars import inspect import json @@ -13,6 +14,7 @@ import threading import time import traceback +import types import warnings from abc import ABC, abstractmethod from collections import deque @@ -47,6 +49,7 @@ import temporalio.activity import temporalio.api.common.v1 import temporalio.api.enums.v1 +import temporalio.api.nexus.v1 import temporalio.api.sdk.v1 import temporalio.api.stream.v1 import temporalio.bridge.proto.activity_result @@ -1184,6 +1187,34 @@ def _apply_resolve_nexus_operation( else: raise RuntimeError("Nexus operation did not have a result") + def _apply_nexus_operation_progress( + self, progress: temporalio.api.nexus.v1.NexusOperationProgress + ) -> None: + if not progress.HasField("operation_key"): + return + key = progress.operation_key + # Only a started operation that has not resolved has a pending handle + # with a token. A resolution in the same task comes before the + # attachments, so its progress finds no handle and is dropped. A token + # is unique only within one handler, so two operations with one key are + # one handler operation, and both get the progress. + for handle in self._pending_nexus_operations.values(): + if ( + handle._input.endpoint == key.endpoint + and handle._input.service == key.service + and handle._input.operation_name == key.operation + and handle.operation_token == key.operation_token + ): + handle._resolve_progress( + temporalio.workflow.NexusOperationProgress( + position=progress.position, + counter=progress.counter, + # Read-only, so Workflow code can't change what the next + # reader of latest_progress sees. + metadata=types.MappingProxyType(dict(progress.metadata)), + ) + ) + def _apply_resolve_request_cancel_external_workflow( self, job: temporalio.bridge.proto.workflow_activation.ResolveRequestCancelExternalWorkflow, @@ -1286,11 +1317,18 @@ def _apply_workflow_task_attachments( # Other attachments belong to features this Worker does not run, so # they are left alone. woke = False + progress: list[temporalio.api.nexus.v1.NexusOperationProgress] = [] for attachment in job.attachments: wake = temporalio.api.stream.v1.StreamWake() if attachment.Unpack(wake): self._stream_wakes.append(wake) woke = True + continue + one = temporalio.api.nexus.v1.NexusOperationProgress() + if attachment.Unpack(one): + progress.append(one) + for one in fold_nexus_operation_progress(progress): + self._apply_nexus_operation_progress(one) if woke: # A wake resolves no future, and wait conditions are checked only # after a ready callback ran, so give the loop one. @@ -3894,6 +3932,13 @@ def __init__( self._payload_converter = payload_converter self._failure_converter = failure_converter self._event_group_markers = _capture_event_group_markers(input.event_groups) + self._latest_progress: temporalio.workflow.NexusOperationProgress | None = None + self._progress_waiters: list[ + tuple[ + int, asyncio.Future[temporalio.workflow.NexusOperationProgress | None] + ] + ] = [] + self._result_fut.add_done_callback(lambda _: self._end_progress_waits()) @property def operation_token(self) -> str | None: @@ -3902,6 +3947,56 @@ def operation_token(self) -> str | None: except BaseException: return None + @property + def latest_progress(self) -> temporalio.workflow.NexusOperationProgress | None: + return self._latest_progress + + async def progress( + self, *, after_counter: int = 0 + ) -> temporalio.workflow.NexusOperationProgress | None: + latest = self._latest_progress + if latest is not None and latest.counter > after_counter: + return latest + if self._result_fut.done(): + return None + waiter: asyncio.Future[temporalio.workflow.NexusOperationProgress | None] = ( + self._instance.create_future() + ) + entry = (after_counter, waiter) + self._progress_waiters.append(entry) + try: + return await waiter + finally: + # A wait given up on, such as one under a timeout in a loop, would + # otherwise stay in the list until the next progress arrives. + with contextlib.suppress(ValueError): + self._progress_waiters.remove(entry) + + def _resolve_progress( + self, progress: temporalio.workflow.NexusOperationProgress + ) -> None: + if ( + self._latest_progress is not None + and progress.counter <= self._latest_progress.counter + ): + return + self._latest_progress = progress + still_waiting = [] + for after_counter, waiter in self._progress_waiters: + if waiter.done(): + continue + if progress.counter > after_counter: + waiter.set_result(progress) + else: + still_waiting.append((after_counter, waiter)) + self._progress_waiters = still_waiting + + def _end_progress_waits(self) -> None: + for _, waiter in self._progress_waiters: + if not waiter.done(): + waiter.set_result(None) + self._progress_waiters = [] + def __await__(self) -> Generator[Any, Any, OutputT]: return self._task.__await__() @@ -4032,6 +4127,31 @@ def _apply_command(self) -> None: command.event_group_markers.extend(self._event_group_markers) +def fold_nexus_operation_progress( + progress: list[temporalio.api.nexus.v1.NexusOperationProgress], +) -> list[temporalio.api.nexus.v1.NexusOperationProgress]: + """Fold the progress of one Workflow Task to one per operation key. + + The key is the endpoint, service, operation and token. The progress with + the highest counter for a key stays, and the first of equal counters + stays. Keys keep the order of their first progress. Progress with no key + is dropped. This is the fold rule that every SDK shares, and the + sdk-streams vectors ``testdata/nexus_progress_folds.json`` check it. + """ + folded: dict[ + tuple[str, str, str, str], temporalio.api.nexus.v1.NexusOperationProgress + ] = {} + for one in progress: + if not one.HasField("operation_key"): + continue + key = one.operation_key + id = (key.endpoint, key.service, key.operation, key.operation_token) + held = folded.get(id) + if held is None or one.counter > held.counter: + folded[id] = one + return list(folded.values()) + + def _encode_search_attributes( attributes: ( temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes diff --git a/temporalio/workflow/__init__.py b/temporalio/workflow/__init__.py index c631fd88e8..1221eb114f 100644 --- a/temporalio/workflow/__init__.py +++ b/temporalio/workflow/__init__.py @@ -137,6 +137,7 @@ NexusClient, NexusOperationCancellationType, NexusOperationHandle, + NexusOperationProgress, _NexusClient, create_nexus_client, ) @@ -253,6 +254,7 @@ "NexusClient", "NexusOperationCancellationType", "NexusOperationHandle", + "NexusOperationProgress", "create_nexus_client", "LoggerAdapter", "SandboxImportNotificationPolicy", diff --git a/temporalio/workflow/_nexus.py b/temporalio/workflow/_nexus.py index a23432545a..c8170210b6 100644 --- a/temporalio/workflow/_nexus.py +++ b/temporalio/workflow/_nexus.py @@ -2,6 +2,7 @@ from abc import ABC, abstractmethod from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence +from dataclasses import dataclass from datetime import timedelta from enum import IntEnum from typing import Any, Generic, overload @@ -21,10 +22,36 @@ "NexusClient", "NexusOperationCancellationType", "NexusOperationHandle", + "NexusOperationProgress", "create_nexus_client", ] +@dataclass(frozen=True) +class NexusOperationProgress: + """Progress a started Nexus operation reported to its caller. + + Progress says that the operation's output moved, not what the output is. + The Workflow reads the output itself, starting from :py:attr:`position`. + + .. warning:: + This API is experimental. + """ + + position: str + """Where the operation's output stands, in the handler's terms, such as a + stream cursor.""" + + counter: int + """Orders the progress of one operation. It only grows.""" + + metadata: Mapping[str, str] + """Small details from the handler.""" + + # The metadata mapping isn't hashable, so a hash would always raise. + __hash__ = None # type: ignore[assignment] + + class NexusOperationHandle(Generic[OutputT]): """Handle for interacting with a Nexus operation.""" @@ -43,6 +70,37 @@ def operation_token(self) -> str | None: """The operation token for this handle.""" raise NotImplementedError + @property + def latest_progress(self) -> NexusOperationProgress | None: + """The latest progress the operation reported, or ``None`` if none yet. + + .. warning:: + This API is experimental. + """ + raise NotImplementedError + + async def progress( + self, + *, + after_counter: int = 0, # type: ignore[reportUnusedParameter] + ) -> NexusOperationProgress | None: + """Wait for progress newer than ``after_counter``. + + Returns the latest progress as soon as its counter is higher than + ``after_counter``, which may be at once. Returns ``None`` if the + operation resolves first, since no newer progress can follow. Progress + comes from History, so a replay returns the same values at the same + points. + + Progress is a hint that the operation's output moved to + :py:attr:`NexusOperationProgress.position`. A burst of progress folds + to its highest counter, so intermediate counters may be skipped. + + .. warning:: + This API is experimental. + """ + raise NotImplementedError + class NexusOperationCancellationType(IntEnum): """Defines behavior of a Nexus operation when the caller workflow initiates cancellation. diff --git a/tests/contrib/streams/_support.py b/tests/contrib/streams/_support.py index f026803b99..3f0524d85a 100644 --- a/tests/contrib/streams/_support.py +++ b/tests/contrib/streams/_support.py @@ -8,7 +8,9 @@ import subprocess import uuid from pathlib import Path -from typing import Any +from typing import Any, TypeVar + +from google.protobuf.message import Message from temporalio.api.workflowservice.v1 import ( DescribeStreamNotifierRequest, @@ -19,6 +21,9 @@ from temporalio.contrib.streams._plugin import call, store_for_client from temporalio.contrib.streams.memory import MemoryStreams from temporalio.converter import DataConverter +from temporalio.service import RPCError, RPCStatusCode + +ResponseT = TypeVar("ResponseT", bound=Message) async def connect_with( @@ -57,6 +62,40 @@ async def describe_stream_notifier( return await call(service.describe_stream_notifier(request)) +def _serialize(message: Message) -> bytes: + return message.SerializeToString() + + +async def call_workflow_service( + client: Client, rpc: str, request: Message, response_type: type[ResponseT] +) -> ResponseT: + """Call a workflow service RPC that the client has no stub for, by its method path. + + The stream runtime calls the notifier RPCs this way. A test that plays a + producer or a handler calls them over a plain gRPC channel to the client's + server, and gets a failure as the ``RPCError`` that a stub raises. + """ + # Imported here, since the Workflow sandbox imports this module again. + import grpc + + async with grpc.aio.insecure_channel( + client.service_client.config.target_host + ) as channel: + method = channel.unary_unary( + f"/temporal.api.workflowservice.v1.WorkflowService/{rpc}", + request_serializer=_serialize, + response_deserializer=response_type.FromString, + ) + try: + return await method( + request, metadata=(("temporal-namespace", client.namespace),) + ) + except grpc.aio.AioRpcError as err: + raise RPCError( + err.details() or "", RPCStatusCode(err.code().value[0]), b"" + ) from None + + @functools.cache def streams_crate_dir() -> Path: """The checkout of the temporalio-streams crate the bridge builds against. diff --git a/tests/nexus/test_workflow_operation_progress_folds.py b/tests/nexus/test_workflow_operation_progress_folds.py new file mode 100644 index 0000000000..a37c25e244 --- /dev/null +++ b/tests/nexus/test_workflow_operation_progress_folds.py @@ -0,0 +1,130 @@ +"""Nexus operation progress as the server attaches it to a Workflow Task. + +Each progress is an ``Any`` of ``NexusOperationProgress`` in a +``WorkflowTaskAttachments`` job, keyed by endpoint, service, operation and +token. The instance folds the progress of a task to one per key and gives it +to each started, unresolved operation with that key. The fold and the +deliveries must give the results of the sdk-streams vectors. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from typing import Any, cast + +import pytest +from google.protobuf.any_pb2 import Any as AnyProto + +from temporalio import workflow +from temporalio.api.nexus.v1 import NexusOperationProgress +from temporalio.bridge.proto.workflow_activation import WorkflowTaskAttachments +from temporalio.worker._workflow_instance import ( + _NexusOperationHandle, + _WorkflowInstanceImpl, + fold_nexus_operation_progress, +) +from tests.contrib.streams._support import streams_crate_dir + + +def _vectors() -> dict[str, Any]: + with open(streams_crate_dir() / "testdata" / "nexus_progress_folds.json") as f: + return json.load(f) + + +def _progress(fields: dict[str, Any]) -> NexusOperationProgress: + fields = dict(fields) + key = fields.pop("key", None) + progress = NexusOperationProgress(**fields) + if key is not None: + progress.operation_key.CopyFrom(NexusOperationProgress.OperationKey(**key)) + return progress + + +@pytest.mark.parametrize("case", _vectors()["folds"], ids=lambda case: case["name"]) +def test_the_fold_gives_the_shared_vectors(case: dict[str, Any]): + folded = fold_nexus_operation_progress([_progress(p) for p in case["progress"]]) + assert folded == [_progress(p) for p in case["folded"]] + + +class _Handle: + """An operation handle with the real progress step, which records what it took.""" + + def __init__( + self, + seq: int, + endpoint: str, + service: str, + operation: str, + delivered: list[dict[str, Any]], + ): + self.seq = seq + self._delivered = delivered + self._input = SimpleNamespace( + endpoint=endpoint, service=service, operation_name=operation + ) + self.operation_token: str | None = None + self._latest_progress: workflow.NexusOperationProgress | None = None + self._progress_waiters: list[Any] = [] + + def _resolve_progress(self, progress: workflow.NexusOperationProgress) -> None: + before = self._latest_progress + _NexusOperationHandle._resolve_progress(self, progress) # type: ignore[arg-type] + if self._latest_progress is not before: + self._delivered.append( + { + "seq": self.seq, + "counter": progress.counter, + "position": progress.position, + } + ) + + +class _Instance: + """The instance's state that the attachments job reads, with its real steps.""" + + _apply_nexus_operation_progress = ( + _WorkflowInstanceImpl._apply_nexus_operation_progress + ) + + def __init__(self) -> None: + self._pending_nexus_operations: dict[int, _Handle] = {} + self._stream_wakes: list[Any] = [] + + def call_soon(self, callback: Any) -> None: + del callback + + +def _packed(progress: NexusOperationProgress) -> AnyProto: + attachment = AnyProto() + attachment.Pack(progress) + return attachment + + +@pytest.mark.parametrize( + "case", _vectors()["deliveries"], ids=lambda case: case["name"] +) +def test_the_deliveries_give_the_shared_vectors(case: dict[str, Any]): + instance = _Instance() + delivered: list[dict[str, Any]] = [] + for step in case["steps"]: + if "schedule" in step: + s = step["schedule"] + handle = _Handle( + s["seq"], s["endpoint"], s["service"], s["operation"], delivered + ) + instance._pending_nexus_operations[handle.seq] = handle + elif "start" in step: + started = instance._pending_nexus_operations[step["start"]["seq"]] + started.operation_token = step["start"]["operation_token"] + elif "resolve" in step: + instance._pending_nexus_operations.pop(step["resolve"]["seq"]) + else: + delivered.clear() + _WorkflowInstanceImpl._apply_workflow_task_attachments( + cast(Any, instance), + WorkflowTaskAttachments( + attachments=[_packed(_progress(p)) for p in step["attachments"]] + ), + ) + assert delivered == step["delivered"] diff --git a/tests/nexus/test_workflow_operation_progress_job.py b/tests/nexus/test_workflow_operation_progress_job.py new file mode 100644 index 0000000000..2317d6c9fd --- /dev/null +++ b/tests/nexus/test_workflow_operation_progress_job.py @@ -0,0 +1,168 @@ +"""A Worker that gets Nexus operation progress it doesn't act on. + +The server attaches an operation's progress to the caller's next Workflow +Task, and Core hands the attachments to the Workflow in one job. A Worker that +doesn't act on progress must drop it, not fail the task, so a server that sends +progress can't break it. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone + +import nexusrpc +from google.protobuf.any_pb2 import Any +from google.protobuf.duration_pb2 import Duration +from google.protobuf.timestamp_pb2 import Timestamp + +from temporalio import workflow +from temporalio.api.common.v1 import Payloads, WorkflowType +from temporalio.api.enums.v1 import EventType +from temporalio.api.history.v1 import ( + HistoryEvent, + NexusOperationCompletedEventAttributes, + NexusOperationScheduledEventAttributes, + NexusOperationStartedEventAttributes, + WorkflowExecutionCompletedEventAttributes, + WorkflowExecutionStartedEventAttributes, + WorkflowTaskCompletedEventAttributes, + WorkflowTaskScheduledEventAttributes, + WorkflowTaskStartedEventAttributes, +) +from temporalio.api.nexus.v1 import NexusOperationProgress +from temporalio.api.taskqueue.v1 import TaskQueue +from temporalio.client import WorkflowHistory +from temporalio.converter import DataConverter +from temporalio.worker import Replayer, UnsandboxedWorkflowRunner + +ENDPOINT = "progress-endpoint" +TASK_QUEUE = "progress-task-queue" +CONVERTER = DataConverter.default.payload_converter + + +@nexusrpc.service +class ProduceService: + produce: nexusrpc.Operation[str, str] + + +@workflow.defn(name="ResultOnlyCaller") +class ResultOnlyCaller: + @workflow.run + async def run(self) -> str: + client = workflow.create_nexus_client(service=ProduceService, endpoint=ENDPOINT) + handle = await client.start_operation(ProduceService.produce, "input") + return await handle + + +def packed(progress: NexusOperationProgress) -> Any: + attachment = Any() + attachment.Pack(progress) + return attachment + + +def history_with_progress() -> WorkflowHistory: + events: list[HistoryEvent] = [] + clock = datetime(2026, 10, 10, tzinfo=timezone.utc) + + def add(event_type: EventType.ValueType, **attributes: object) -> int: + nonlocal clock + clock += timedelta(milliseconds=10) + event_time = Timestamp() + event_time.FromDatetime(clock) + events.append( + HistoryEvent( + event_id=len(events) + 1, + event_type=event_type, + event_time=event_time, + **attributes, # type: ignore[arg-type] + ) + ) + return len(events) + + def task(*progress: NexusOperationProgress) -> int: + scheduled = add( + EventType.EVENT_TYPE_WORKFLOW_TASK_SCHEDULED, + workflow_task_scheduled_event_attributes=WorkflowTaskScheduledEventAttributes( + task_queue=TaskQueue(name=TASK_QUEUE), + start_to_close_timeout=Duration(seconds=10), + attempt=1, + attachments=[packed(p) for p in progress], + ), + ) + started = add( + EventType.EVENT_TYPE_WORKFLOW_TASK_STARTED, + workflow_task_started_event_attributes=WorkflowTaskStartedEventAttributes( + scheduled_event_id=scheduled + ), + ) + return add( + EventType.EVENT_TYPE_WORKFLOW_TASK_COMPLETED, + workflow_task_completed_event_attributes=WorkflowTaskCompletedEventAttributes( + scheduled_event_id=scheduled, started_event_id=started + ), + ) + + add( + EventType.EVENT_TYPE_WORKFLOW_EXECUTION_STARTED, + workflow_execution_started_event_attributes=WorkflowExecutionStartedEventAttributes( + workflow_type=WorkflowType(name="ResultOnlyCaller"), + task_queue=TaskQueue(name=TASK_QUEUE), + original_execution_run_id="run-id", + first_execution_run_id="run-id", + attempt=1, + ), + ) + completed = task() + scheduled = add( + EventType.EVENT_TYPE_NEXUS_OPERATION_SCHEDULED, + nexus_operation_scheduled_event_attributes=NexusOperationScheduledEventAttributes( + endpoint=ENDPOINT, + service="ProduceService", + operation="produce", + input=CONVERTER.to_payload("input"), + workflow_task_completed_event_id=completed, + request_id="request-id", + ), + ) + add( + EventType.EVENT_TYPE_NEXUS_OPERATION_STARTED, + nexus_operation_started_event_attributes=NexusOperationStartedEventAttributes( + scheduled_event_id=scheduled, + operation_token="operation-token", + request_id="request-id", + ), + ) + task( + NexusOperationProgress( + operation_key=NexusOperationProgress.OperationKey( + endpoint=ENDPOINT, + service="ProduceService", + operation="produce", + operation_token="operation-token", + ), + position="cursor-1", + counter=1, + ) + ) + add( + EventType.EVENT_TYPE_NEXUS_OPERATION_COMPLETED, + nexus_operation_completed_event_attributes=NexusOperationCompletedEventAttributes( + scheduled_event_id=scheduled, result=CONVERTER.to_payload("done") + ), + ) + completed = task() + add( + EventType.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED, + workflow_execution_completed_event_attributes=WorkflowExecutionCompletedEventAttributes( + result=Payloads(payloads=[CONVERTER.to_payload("done")]), + workflow_task_completed_event_id=completed, + ), + ) + return WorkflowHistory(workflow_id="result-only-caller", events=events) + + +async def test_progress_the_workflow_does_not_act_on_is_dropped() -> None: + result = await Replayer( + workflows=[ResultOnlyCaller], workflow_runner=UnsandboxedWorkflowRunner() + ).replay_workflow(history_with_progress(), raise_on_replay_failure=False) + assert result.replay_failure is None, result.replay_failure