Sub-orchestrations, when_all, and when_any support (#5)

cgillum · cgillum · commit 62ec41205068 · 2023-04-14T00:29:37.000Z
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -1,9 +1,18 @@
 {
-    "editor.formatOnSave": true,
-    "python.analysis.typeCheckingMode": "basic",
-    "python.formatting.autopep8Args": [
-        "--max-line-length=120"
+    "[python]": {
+        "editor.defaultFormatter": "ms-python.autopep8",
+        "editor.formatOnSave": true,
+        "editor.codeActionsOnSave": {
+            "source.organizeImports": true,
+        },
+        "editor.rulers": [
+            119
+        ],
+    },
+    "autopep8.args": [
+        "--max-line-length=119"
     ],
+    "python.analysis.typeCheckingMode": "basic",
     "python.testing.pytestArgs": [
         "-v",
         "--cov=durabletask/",
diff --git a/durabletask/protos/helpers.py b/durabletask/protos/helpers.py
@@ -69,6 +69,37 @@ def new_task_failed_event(event_id: int, ex: Exception) -> HistoryEvent:
     )
 
 
+def new_sub_orchestration_created_event(event_id: int, name: str, instance_id: str, encoded_input: str | None = None) -> HistoryEvent:
+    return HistoryEvent(
+        eventId=event_id,
+        timestamp=timestamp_pb2.Timestamp(),
+        subOrchestrationInstanceCreated=SubOrchestrationInstanceCreatedEvent(
+            name=name,
+            input=get_string_value(encoded_input),
+            instanceId=instance_id)
+    )
+
+
+def new_sub_orchestration_completed_event(event_id: int, encoded_output: str | None = None) -> HistoryEvent:
+    return HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        subOrchestrationInstanceCompleted=SubOrchestrationInstanceCompletedEvent(
+            result=get_string_value(encoded_output),
+            taskScheduledId=event_id)
+    )
+
+
+def new_sub_orchestration_failed_event(event_id: int, ex: Exception) -> HistoryEvent:
+    return HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        subOrchestrationInstanceFailed=SubOrchestrationInstanceFailedEvent(
+            failureDetails=new_failure_details(ex),
+            taskScheduledId=event_id)
+    )
+
+
 def new_failure_details(ex: Exception) -> TaskFailureDetails:
     return TaskFailureDetails(
         errorType=type(ex).__name__,
@@ -120,5 +151,14 @@ def new_timestamp(dt: datetime) -> timestamp_pb2.Timestamp:
     return ts
 
 
+def new_create_sub_orchestration_action(id: int, name: str, instance_id: str | None, input: Any) -> OrchestratorAction:
+    encoded_input = json.dumps(input) if input is not None else None
+    return OrchestratorAction(id=id, createSubOrchestration=CreateSubOrchestrationAction(
+        name=name,
+        instanceId=instance_id,
+        input=get_string_value(encoded_input)
+    ))
+
+
 def is_empty(v: wrappers_pb2.StringValue):
     return v is None or v.value == ''
diff --git a/durabletask/task/execution.py b/durabletask/task/execution.py
@@ -1,16 +1,14 @@
-import inspect
-from logging import Logger
-import simplejson as json
-
 from datetime import datetime
+from logging import Logger
 from types import GeneratorType
 from typing import Any, Dict, Generator, Iterable, List, TypeVar
 
+import simplejson as json
+
 import durabletask.protos.helpers as ph
 import durabletask.protos.orchestrator_service_pb2 as pb
-from durabletask.task.activities import Activity, ActivityContext
 import durabletask.task.task as task
-
+from durabletask.task.activities import Activity, ActivityContext
 from durabletask.task.orchestration import OrchestrationContext, Orchestrator
 from durabletask.task.registry import Registry, get_name
 from durabletask.task.task import Task
@@ -67,11 +65,11 @@ def resume(self):
         # case is if the user yielded on a WhenAll task and there are still
         # outstanding child tasks that need to be completed.
         if self._previous_task is not None:
-            if self._previous_task.is_failed():
+            if self._previous_task.is_failed:
                 # Raise the failure as an exception to the generator. The orchestrator can then either
                 # handle the exception or allow it to fail the orchestration.
                 self._generator.throw(self._previous_task.get_exception())
-            elif self._previous_task.is_complete():
+            elif self._previous_task.is_complete:
                 # Resume the generator. This will either return a Task or raise StopIteration if it's done.
                 next_task = self._generator.send(self._previous_task.get_result())
                 # TODO: Validate the return value
@@ -138,6 +136,21 @@ def call_activity(self, activity: Activity[TInput, TOutput], *,
         self._pending_tasks[id] = activity_task
         return activity_task
 
+    def call_sub_orchestrator(self, orchestrator: Orchestrator[TInput, TOutput], *,
+                              input: TInput | None = None,
+                              instance_id: str | None = None) -> task.Task[TOutput]:
+        id = self.next_sequence_number()
+        name = get_name(orchestrator)
+        if instance_id is None:
+            # Create a deteministic instance ID based on the parent instance ID
+            instance_id = f"{self.instance_id}:{id:04x}"
+        action = ph.new_create_sub_orchestration_action(id, name, instance_id, input)
+        self._pending_actions[id] = action
+
+        sub_orch_task = task.CompletableTask[TOutput]()
+        self._pending_tasks[id] = sub_orch_task
+        return sub_orch_task
+
 
 class OrchestrationExecutor:
     _generator: Orchestrator | None
@@ -252,6 +265,45 @@ def process_event(self, ctx: RuntimeOrchestrationContext, event: pb.HistoryEvent
                     return
                 activity_task.fail(event.taskFailed.failureDetails)
                 ctx.resume()
+            elif event.HasField("subOrchestrationInstanceCreated"):
+                # This history event confirms that the sub-orchestration execution was successfully scheduled.
+                # Remove the subOrchestrationInstanceCreated event from the pending action list so we don't schedule it again.
+                task_id = event.eventId
+                action = ctx._pending_actions.pop(task_id, None)
+                if not action:
+                    raise _get_non_determinism_error(task_id, get_name(ctx.call_sub_orchestrator))
+                elif not action.HasField("createSubOrchestration"):
+                    expected_method_name = get_name(ctx.call_sub_orchestrator)
+                    raise _get_wrong_action_type_error(task_id, expected_method_name, action)
+                elif action.createSubOrchestration.name != event.subOrchestrationInstanceCreated.name:
+                    raise _get_wrong_action_name_error(
+                        task_id,
+                        method_name=get_name(ctx.call_sub_orchestrator),
+                        expected_task_name=event.subOrchestrationInstanceCreated.name,
+                        actual_task_name=action.createSubOrchestration.name)
+            elif event.HasField("subOrchestrationInstanceCompleted"):
+                task_id = event.subOrchestrationInstanceCompleted.taskScheduledId
+                sub_orch_task = ctx._pending_tasks.pop(task_id, None)
+                if not sub_orch_task:
+                    # TODO: Should this be an error? When would it ever happen?
+                    self._logger.warning(
+                        f"Ignoring unexpected subOrchestrationInstanceCompleted event for '{ctx.instance_id}' with ID = {task_id}.")
+                    return
+                result = None
+                if not ph.is_empty(event.subOrchestrationInstanceCompleted.result):
+                    result = json.loads(event.subOrchestrationInstanceCompleted.result.value)
+                sub_orch_task.complete(result)
+                ctx.resume()
+            elif event.HasField("subOrchestrationInstanceFailed"):
+                task_id = event.subOrchestrationInstanceFailed.taskScheduledId
+                sub_orch_task = ctx._pending_tasks.pop(task_id, None)
+                if not sub_orch_task:
+                    # TODO: Should this be an error? When would it ever happen?
+                    self._logger.warning(
+                        f"Ignoring unexpected subOrchestrationInstanceFailed event for '{ctx.instance_id}' with ID = {task_id}.")
+                    return
+                sub_orch_task.fail(event.subOrchestrationInstanceFailed.failureDetails)
+                ctx.resume()
             else:
                 eventType = event.WhichOneof("eventType")
                 raise OrchestrationStateError(f"Don't know how to handle event of type '{eventType}'")
diff --git a/durabletask/task/orchestration.py b/durabletask/task/orchestration.py
@@ -1,3 +1,6 @@
+# See https://peps.python.org/pep-0563/
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import Any, Callable, Generator, TypeVar
@@ -99,6 +102,29 @@ def call_activity(self, activity: Activity[TInput, TOutput], *,
         """
         pass
 
+    @abstractmethod
+    def call_sub_orchestrator(self,  orchestrator: Orchestrator[TInput, TOutput], *,
+                              input: TInput | None = None,
+                              instance_id: str | None = None) -> task.Task[TOutput]:
+        """Schedule sub-orchestrator function for execution.
+
+        Parameters
+        ----------
+        orchestrator: Orchestrator[TInput, TOutput]
+            A reference to the orchestrator function to call.
+        input: TInput | None
+            The optional JSON-serializable input to pass to the orchestrator function.
+        instance_id: str | None
+            A unique ID to use for the sub-orchestration instance. If not specified, a
+            random UUID will be used.
+
+        Returns
+        -------
+        Task
+            A Durable Task that completes when the called sub-orchestrator completes or fails.
+        """
+        pass
+
 
 # Orchestrators are generators that yield tasks and receive/return any type
 Orchestrator = Callable[[OrchestrationContext, TInput], Generator[task.Task, Any, Any] | TOutput]
diff --git a/durabletask/task/task.py b/durabletask/task/task.py
@@ -1,8 +1,11 @@
+# See https://peps.python.org/pep-0563/
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import Type, TypeVar, Generic
+from typing import Generic, List, TypeVar
 
-import durabletask.protos.orchestrator_service_pb2 as pb
 import durabletask.protos.helpers as pbh
+import durabletask.protos.orchestrator_service_pb2 as pb
 
 T = TypeVar('T')
 
@@ -31,57 +34,132 @@ def stack_trace(self) -> str | None:
 
 class Task(ABC, Generic[T]):
     """Abstract base class for asynchronous tasks in a durable orchestration."""
+    _result: T
+    _exception: TaskFailedError | None
+    _parent: CompositeTask[T] | None
 
     def __init__(self) -> None:
         super().__init__()
-        pass
+        self._is_complete = False
+        self._exception = None
+        self._parent = None
 
-    @abstractmethod
+    @property
     def is_complete(self) -> bool:
-        pass
+        """Returns True if the task has completed, False otherwise."""
+        return self._is_complete
 
-    @abstractmethod
+    @property
     def is_failed(self) -> bool:
-        pass
+        """Returns True if the task has failed, False otherwise."""
+        return self._exception is not None
 
-    @abstractmethod
     def get_result(self) -> T:
-        pass
+        """Returns the result of the task."""
+        if not self._is_complete:
+            raise ValueError('The task has not completed.')
+        elif self._exception is not None:
+            raise self._exception
+        return self._result
 
-    @abstractmethod
     def get_exception(self) -> TaskFailedError:
+        """Returns the exception that caused the task to fail."""
+        if self._exception is None:
+            raise ValueError('The task has not failed.')
+        return self._exception
+
+
+class CompositeTask(Task[T]):
+    """A task that is composed of other tasks."""
+    _tasks: List[Task]
+
+    def __init__(self, tasks: List[Task]):
+        super().__init__()
+        self._tasks = tasks
+        self._completed_tasks = 0
+        self._failed_tasks = 0
+        for task in tasks:
+            task._parent = self
+
+    def get_tasks(self) -> List[Task]:
+        return self._tasks
+
+    @abstractmethod
+    def on_child_completed(self, task: Task[T]):
         pass
 
 
 class CompletableTask(Task[T]):
-    _result: T | None
-    _exception: TaskFailedError | None
 
     def __init__(self):
         super().__init__()
-        self._is_complete = False
-        self._result = None
-        self._exception = None
 
     def complete(self, result: T):
+        if self._is_complete:
+            raise ValueError('The task has already completed.')
         self._result = result
         self._is_complete = True
+        if self._parent is not None:
+            self._parent.on_child_completed(self)
 
     def fail(self, details: pb.TaskFailureDetails):
+        if self._is_complete:
+            raise ValueError('The task has already completed.')
         self._exception = TaskFailedError(
             details.errorMessage,
             details.errorType,
             details.stackTrace.value if not pbh.is_empty(details.stackTrace) else None)
         self._is_complete = True
+        if self._parent is not None:
+            self._parent.on_child_completed(self)
 
-    def is_complete(self) -> bool:
-        return self._is_complete
 
-    def is_failed(self) -> bool:
-        return self._exception is not None
+class WhenAllTask(CompositeTask[List[T]]):
+    """A task that completes when all of its child tasks complete."""
 
-    def get_result(self) -> T | None:
-        return self._result
+    def __init__(self, tasks: List[Task[T]]):
+        super().__init__(tasks)
+        self._completed_tasks = 0
+        self._failed_tasks = 0
 
-    def get_exception(self) -> TaskFailedError | None:
-        return self._exception
+    @property
+    def pending_tasks(self) -> int:
+        """Returns the number of tasks that have not yet completed."""
+        return len(self._tasks) - self._completed_tasks
+
+    def on_child_completed(self, task: Task[T]):
+        if self.is_complete:
+            raise ValueError('The task has already completed.')
+        self._completed_tasks += 1
+        if task.is_failed and self._exception is None:
+            self._exception = task.get_exception()
+            self._is_complete = True
+        if self._completed_tasks == len(self._tasks):
+            # The order of the result MUST match the order of the tasks provided to the constructor.
+            self._result = [task.get_result() for task in self._tasks]
+            self._is_complete = True
+
+    def get_completed_tasks(self) -> int:
+        return self._completed_tasks
+
+
+class WhenAnyTask(CompositeTask[Task]):
+    """A task that completes when any of its child tasks complete."""
+
+    def __init__(self, tasks: List[Task]):
+        super().__init__(tasks)
+
+    def on_child_completed(self, task: Task):
+        if not self.is_complete:
+            self._is_complete = True
+            self._result = task
+
+
+def when_all(tasks: List[Task[T]]) -> WhenAllTask[T]:
+    """Returns a task that completes when all of the provided tasks complete or when one of the tasks fail."""
+    return WhenAllTask(tasks)
+
+
+def when_any(tasks: List[Task]) -> WhenAnyTask:
+    """Returns a task that completes when any of the provided tasks complete or fail."""
+    return WhenAnyTask(tasks)
diff --git a/tests/test_orchestration_e2e.py b/tests/test_orchestration_e2e.py
diff --git a/tests/test_orchestration_executor.py b/tests/test_orchestration_executor.py
