WIP: nexus worker

Author: dandavison

temporalio/nexus/__init__.py

@@ -0,0 +1,12 @@
+import logging
+from collections.abc import Mapping
+from typing import Any, Optional
+
+
+class LoggerAdapter(logging.LoggerAdapter):
+    def __init__(self, logger: logging.Logger, extra: Optional[Mapping[str, Any]]):
+        super().__init__(logger, extra or {})
+
+
+logger = LoggerAdapter(logging.getLogger(__name__), None)
+"""Logger that has additional details regarding the current Nexus operation."""

temporalio/worker/_activity.py

@@ -201,6 +201,7 @@ async def drain_poll_queue(self) -> None:
 
     # Only call this after run()/drain_poll_queue() have returned. This will not
     # raise an exception.
+    # TODO(dan): check accuracy of this comment; I would say it *does* raise an exception.
     async def wait_all_completed(self) -> None:
         running_tasks = [v.task for v in self._running_activities.values() if v.task]
         if running_tasks:

temporalio/worker/_nexus.py

@@ -8,6 +8,7 @@
 import pprint
 from typing import (
     Any,
+    Awaitable,
     Callable,
     Sequence,
     Union,
@@ -67,6 +68,7 @@ def __init__(
         self._interceptors = interceptors
         # TODO(dan): metric_meter
         self._metric_meter = metric_meter
+        self._running_operations: dict[bytes, asyncio.Task] = {}
 
     def _validate_nexus_services(
         self, nexus_services: Sequence[Any]
@@ -163,6 +165,11 @@ async def drain_poll_queue(self) -> None:
             except temporalio.bridge.worker.PollShutdownError:
                 return
 
+    async def wait_all_completed(self) -> None:
+        await asyncio.gather(
+            *self._running_operations.values(), return_exceptions=False
+        )
+
     # TODO(dan): is it correct to import from temporalio.api.nexus?
     # Why are these things not exposed in temporalio.bridge?
     async def _handle_start_operation(
@@ -194,53 +201,77 @@ async def _handle_start_operation(
             print(
                 f"🌈@@ worker received task with link: {google.protobuf.json_format.MessageToJson(l)}"
             )
+
+        # TODO(dan): shouldn't this be set in the _run_nexus_operation context? (that doesn't work currently)
         temporalio.nexus.handler._current_context.set(
             temporalio.nexus.handler._Context(
                 client=self._client,
                 task_queue=self._task_queue,
             )
         )
+        self._running_operations[task_token] = asyncio.create_task(
+            self._run_nexus_operation(task_token, operation.start, input, options)
+        )
 
-        # message NexusTaskCompletion {
-        #     bytes task_token = 1;
-        #     oneof status {
-        #         temporal.api.nexus.v1.Response completed = 2;
-        #         temporal.api.nexus.v1.HandlerError error = 3;
-        #         bool ack_cancel = 4;
-        #     }
-        # }
+    # TODO(dan): start type
+    async def _run_nexus_operation(
+        self,
+        task_token: bytes,
+        start: Callable[..., Awaitable[Any]],
+        input: Any,
+        options: nexusrpc.handler.StartOperationOptions,
+    ) -> None:
+        try:
+            result = await start(input, options)
+        except BaseException:
+            # TODO(dan): mirror appropriate aspects of _run_activity error handling
+            raise NotImplementedError(
+                "TODO: Nexus operation error handling not implemented"
+            )
 
-        result = await operation.start(input, options)
+        try:
+            # Send task completion
+            if isinstance(result, nexusrpc.handler.StartOperationAsyncResult):
+                print(f"🟢 Nexus operation started with async response {result}")
+                op_resp = temporalio.api.nexus.v1.StartOperationResponse(
+                    async_success=temporalio.api.nexus.v1.StartOperationResponse.Async(
+                        operation_token=result.token,
+                        links=[
+                            temporalio.api.nexus.v1.Link(url=l.url, type=l.type)
+                            for l in result.links
+                        ],
+                    )
+                )
+            else:
+                # TODO(dan): are we going to use StartOperationSyncResult from nexusrpc?
+                # (contains links and headers in addition to result) IIRC Go does something
+                # like that.
+                [payload] = await self._data_converter.encode([result])
+                op_resp = temporalio.api.nexus.v1.StartOperationResponse(
+                    sync_success=temporalio.api.nexus.v1.StartOperationResponse.Sync(
+                        payload=payload
+                    )
+                )
+            # message NexusTaskCompletion {
+            #     bytes task_token = 1;
+            #     oneof status {
+            #         temporal.api.nexus.v1.Response completed = 2;
+            #         temporal.api.nexus.v1.HandlerError error = 3;
+            #         bool ack_cancel = 4;
+            #     }
+            # }
 
-        # Send task completion
-        if isinstance(result, nexusrpc.handler.StartOperationAsyncResult):
-            print(
-                f"🟢 Nexus operation {request.operation} started with async response {result}"
-            )
-            op_resp = temporalio.api.nexus.v1.StartOperationResponse(
-                async_success=temporalio.api.nexus.v1.StartOperationResponse.Async(
-                    operation_token=result.token,
-                    links=[
-                        temporalio.api.nexus.v1.Link(url=l.url, type=l.type)
-                        for l in result.links
-                    ],
+            await self._bridge_worker().complete_nexus_task(
+                temporalio.bridge.proto.nexus.NexusTaskCompletion(
+                    task_token=task_token,
+                    completed=temporalio.api.nexus.v1.Response(start_operation=op_resp),
                 )
             )
-        else:
-            # TODO(dan): are we going to use StartOperationSyncResult from nexusrpc?
-            # (contains links and headers in addition to result) IIRC Go does something
-            # like that.
-            [payload] = await self._data_converter.encode([result])
-            op_resp = temporalio.api.nexus.v1.StartOperationResponse(
-                sync_success=temporalio.api.nexus.v1.StartOperationResponse.Sync(
-                    payload=payload
-                )
+            del self._running_operations[task_token]
+        except Exception:
+            temporalio.nexus.logger.exception(
+                "Failed to send Nexus operation completion"
             )
-        completion = temporalio.bridge.proto.nexus.NexusTaskCompletion(
-            task_token=task_token,
-            completed=temporalio.api.nexus.v1.Response(start_operation=op_resp),
-        )
-        await self._bridge_worker().complete_nexus_task(completion)
 
     async def _handle_cancel_operation(
         self, request: temporalio.api.nexus.v1.CancelOperationRequest, task_token: bytes

temporalio/worker/_worker.py

@@ -146,6 +146,8 @@ def __init__(
             maximum=5
         ),
     ) -> None:
+        # TODO(dan): consider not allowing max_workers < max_concurrent_nexus_operations?
+        # TODO(dan): Nexus tuner support?
         """Create a worker to process workflows and/or activities.
 
         Args:
@@ -171,6 +173,14 @@ def __init__(
                 executor should at least be ``max_concurrent_activities`` or a
                 warning is issued. Note, a broken-executor failure from this
                 executor will cause the worker to fail and shutdown.
+            nexus_operation_executor: Concurrent executor to use for non-async
+                Nexus operations. This is required if any operation start methods
+                are non-async. :py:class:`concurrent.futures.ThreadPoolExecutor`
+                is recommended. If this is a
+                :py:class:`concurrent.futures.ProcessPoolExecutor`, all non-async
+                start methods must be picklable. ``max_workers`` on the executor
+                should at least be ``max_concurrent_nexus_operations`` or a warning
+                is issued.
             workflow_task_executor: Thread pool executor for workflow tasks. If
                 this is not present, a new
                 :py:class:`concurrent.futures.ThreadPoolExecutor` will be
@@ -199,6 +209,8 @@ def __init__(
                 will ever be given to the activity worker concurrently. Mutually exclusive with ``tuner``.
             max_concurrent_local_activities: Maximum number of local activity
                 tasks that will ever be given to the activity worker concurrently. Mutually exclusive with ``tuner``.
+            max_concurrent_nexus_operations: Maximum number of Nexus operations that
+                will ever be given to the Nexus worker concurrently. Mutually exclusive with ``tuner``.
             max_concurrent_workflow_tasks: Maximum allowed number of
                 tasks that will ever be given to the workflow worker at one time. Mutually exclusive with ``tuner``.
             tuner:  Provide a custom :py:class:`WorkerTuner`. Mutually exclusive with the
@@ -688,12 +700,15 @@ async def raise_on_shutdown():
         for task in tasks.values():
             task.cancel()
 
-        # If there's an activity worker, we have to let all activity completions
-        # finish. We cannot guarantee that because poll shutdown completed
-        # (which means activities completed) that they got flushed to the
-        # server.
+        # Let all activity / nexus operations completions finish. We cannot guarantee that
+        # because poll shutdown completed (which means activities/operations completed)
+        # that they got flushed to the server.
         if self._activity_worker:
             await self._activity_worker.wait_all_completed()
+        if self._nexus_worker:
+            await self._nexus_worker.wait_all_completed()
+
+        # TODO(dan): check that we do all appropriate things for nexus worker that we do for activity worker
 
         # Do final shutdown
         try:

tests/worker/test_nexus.py

@@ -31,6 +31,11 @@
 from temporalio.exceptions import CancelledError, NexusOperationError
 from temporalio.worker import UnsandboxedWorkflowRunner, Worker
 
+# TODO(dan): test availability of Temporal client etc in async context set by worker
+# TODO(dan): test worker shutdown, wait_all_completed, drain etc
+# TODO(dan): test worker op handling failure
+# TODO(dan): test contextual logger
+
 # -----------------------------------------------------------------------------
 # Test definition
 #