Add deterministic workflow task groups

README.md

@@ -879,15 +879,12 @@ as normal. Some asyncio features are disabled such as:
 * Calls that use anything external such as networking, subprocesses, disk IO, etc
 
 Also, there are some `asyncio` utilities that internally use `set()` which can make them non-deterministic from one
-worker to the next. Therefore the following `asyncio` functions have `workflow`-module alternatives that are
+worker to the next. Therefore the following `asyncio` utilities have `workflow`-module alternatives that are
 deterministic:
 
 * `asyncio.as_completed()` - use `workflow.as_completed()`
 * `asyncio.wait()` - use `workflow.wait()`
-
-`asyncio.TaskGroup` tracks tasks in a `set()`, so cancellation order can differ between execution and replay.
-This can cause non-determinism errors, even if the workflow catches the task group's exception and completes
-successfully. The sandbox warns when a task group is created.
+* `asyncio.TaskGroup` - use `workflow.TaskGroup`
 
 #### Asyncio Cancellation
 

changelog/added/wobbly-task-teapot.md

@@ -1 +1 @@
-Warn when `asyncio.TaskGroup` is created in a sandboxed workflow because its non-deterministic task cancellation order can cause replay failures.
+Add `workflow.TaskGroup`, a deterministic alternative to `asyncio.TaskGroup` that requests cancellation of unfinished tasks in the order they were added. Warn when `asyncio.TaskGroup` is created in a sandboxed workflow because its non-deterministic task cancellation order can cause replay failures.

temporalio/worker/workflow_sandbox/_restrictions.py

@@ -580,7 +580,7 @@ def _public_callables(parent: Any, *, exclude: set[str] = set()) -> set[str]:
                             match_self=True,
                             leaf_warning=UserWarning,
                             leaf_message=(
-                                "asyncio.TaskGroup() has non-deterministic task cancellation order"
+                                "asyncio.TaskGroup() has non-deterministic task cancellation order, use workflow.TaskGroup instead"
                             ),
                         )
                     },

temporalio/workflow/__init__.py

@@ -52,6 +52,7 @@
 )
 from ._asyncio import (
     _FT,
+    TaskGroup,
     _release_waiter,
     _wait,
     as_completed,
@@ -195,6 +196,7 @@
     "start_local_activity",
     "start_local_activity_class",
     "start_local_activity_method",
+    "TaskGroup",
     "as_completed",
     "wait",
     "Info",

temporalio/workflow/_asyncio.py

@@ -1,17 +1,57 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Awaitable, Iterable, Iterator
+from collections.abc import Awaitable, Iterable, Iterator, MutableSet
 from typing import TYPE_CHECKING, Any, TypeVar, overload
 
 from ..types import AnyType
 
 __all__ = [
+    "TaskGroup",
     "as_completed",
     "wait",
 ]
 
 
+class _OrderedTasks(MutableSet[asyncio.Task[Any]]):
+    def __init__(self) -> None:
+        self._items: dict[asyncio.Task[Any], None] = {}
+
+    def __contains__(self, task: object) -> bool:
+        return task in self._items
+
+    def __iter__(self) -> Iterator[asyncio.Task[Any]]:
+        return iter(self._items)
+
+    def __len__(self) -> int:
+        return len(self._items)
+
+    def add(self, value: asyncio.Task[Any]) -> None:
+        self._items[value] = None
+
+    def discard(self, value: asyncio.Task[Any]) -> None:
+        self._items.pop(value, None)
+
+
+class TaskGroup(asyncio.TaskGroup):
+    """A deterministic version of :py:class:`asyncio.TaskGroup` for workflows.
+
+    Cancellation is requested for unfinished tasks in the order they were added
+    to the group. This does not guarantee the order in which their cleanup finishes.
+    All other behavior follows the running Python version's ``asyncio.TaskGroup``,
+    including waiting for tasks to finish and grouping their exceptions.
+
+    As with other workflow code, tasks must not perform non-deterministic operations.
+    """
+
+    def __init__(self) -> None:
+        """Create a task group."""
+        super().__init__()
+        # Keep asyncio's version-specific cancellation and exception handling,
+        # but remove the set iteration that makes cancellation non-deterministic.
+        self._tasks = _OrderedTasks()  # type: ignore[assignment]
+
+
 def as_completed(
     fs: Iterable[Awaitable[AnyType]], *, timeout: float | None = None
 ) -> Iterator[Awaitable[AnyType]]:

tests/test_workflow_exports.py

@@ -47,6 +47,7 @@
     "RootInfo",
     "SandboxImportNotificationPolicy",
     "SelfType",
+    "TaskGroup",
     "UnfinishedSignalHandlersWarning",
     "UnfinishedUpdateHandlersWarning",
     "UpdateInfo",

tests/test_workflow_task_group.py

@@ -0,0 +1,249 @@
+import asyncio
+import sys
+import uuid
+from datetime import timedelta
+
+import pytest
+
+import temporalio.exceptions
+from temporalio import activity, workflow
+from temporalio.api.enums.v1 import EventType
+from temporalio.client import Client
+from temporalio.worker import Replayer, Worker
+from temporalio.workflow._asyncio import _OrderedTasks
+
+
+@pytest.mark.parametrize("failure_in_body", [False, True])
+async def test_task_group_cancellation_order(failure_in_body: bool):
+    cancelled = []
+
+    async def wait(index: int) -> None:
+        try:
+            await asyncio.Future()
+        finally:
+            cancelled.append(index)
+            await asyncio.sleep(0)
+
+    async def fail() -> None:
+        raise ValueError("planned failure")
+
+    tasks = []
+    with pytest.raises(ExceptionGroup) as exc:
+        async with workflow.TaskGroup() as group:
+            tasks = [group.create_task(wait(index)) for index in range(3)]
+            await asyncio.sleep(0)
+            if failure_in_body:
+                raise ValueError("planned failure")
+            group.create_task(fail())
+            await asyncio.Future()
+    assert [str(error) for error in exc.value.exceptions] == ["planned failure"]
+    assert all(task.done() for task in tasks)
+    assert cancelled == [0, 1, 2]
+    # An internal cancellation must not leak out after the exception is caught.
+    await asyncio.sleep(0)
+
+
+async def test_task_group_ordered_registry():
+    tasks = [asyncio.create_task(asyncio.sleep(0)) for _ in range(3)]
+    await asyncio.gather(*tasks)
+    registry = _OrderedTasks()
+    for task in tasks:
+        registry.add(task)
+    registry.add(tasks[1])
+    assert list(registry) == tasks
+    assert len(registry) == 3
+
+    registry.discard(tasks[1])
+    registry.discard(tasks[1])
+    assert list(registry) == [tasks[0], tasks[2]]
+    assert tasks[1] not in registry
+
+    registry.add(tasks[1])
+    assert list(registry) == [tasks[0], tasks[2], tasks[1]]
+    assert tasks[1] in registry
+    for task in tasks:
+        registry.discard(task)
+    assert not registry
+
+
+async def test_task_group_external_cancellation():
+    started = asyncio.Event()
+    cleaned_up = asyncio.Event()
+
+    async def wait() -> None:
+        started.set()
+        try:
+            await asyncio.Future()
+        finally:
+            await asyncio.sleep(0)
+            cleaned_up.set()
+
+    async def run() -> None:
+        async with workflow.TaskGroup() as group:
+            group.create_task(wait())
+
+    parent = asyncio.create_task(run())
+    await started.wait()
+    parent.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await parent
+    assert cleaned_up.is_set()
+    assert parent.cancelling() == 1
+
+
+async def test_task_group_nested_failure():
+    cancelled = []
+
+    async def wait(name: str) -> None:
+        try:
+            await asyncio.Future()
+        finally:
+            cancelled.append(name)
+
+    async def fail() -> None:
+        await asyncio.sleep(0)
+        raise ValueError("nested failure")
+
+    async def inner() -> None:
+        async with workflow.TaskGroup() as group:
+            group.create_task(wait("inner"))
+            group.create_task(fail())
+
+    with pytest.raises(ExceptionGroup) as exc:
+        async with workflow.TaskGroup() as group:
+            group.create_task(wait("outer"))
+            group.create_task(inner())
+    assert cancelled == ["inner", "outer"]
+    nested = exc.value.exceptions[0]
+    assert isinstance(nested, ExceptionGroup)
+    assert isinstance(nested.exceptions[0], ValueError)
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 13),
+    reason="asyncio.TaskGroup preserves simultaneous external cancellation on Python 3.13+",
+)
+async def test_task_group_failure_preserves_external_cancellation():
+    cleanup_started = asyncio.Event()
+    release_cleanup = asyncio.Event()
+
+    async def wait() -> None:
+        try:
+            await asyncio.Future()
+        finally:
+            cleanup_started.set()
+            await release_cleanup.wait()
+
+    async def fail() -> None:
+        raise ValueError("failure")
+
+    async def run() -> None:
+        with pytest.raises(ExceptionGroup) as exc:
+            async with workflow.TaskGroup() as group:
+                group.create_task(wait())
+                group.create_task(fail())
+        assert [str(error) for error in exc.value.exceptions] == ["failure"]
+        with pytest.raises(asyncio.CancelledError):
+            await asyncio.sleep(0)
+
+    parent = asyncio.create_task(run())
+    await cleanup_started.wait()
+    # Cancel while the group is already waiting for a failed child's siblings.
+    parent.cancel()
+    release_cleanup.set()
+    await parent
+    assert parent.cancelling() == 1
+
+
+@workflow.defn
+class TaskGroupWaitingChild:
+    @workflow.run
+    async def run(self) -> None:
+        await asyncio.Future()
+
+
+@activity.defn
+async def task_group_waiting_activity() -> None:
+    await asyncio.Future()
+
+
+@workflow.defn
+class TaskGroupMixedOperationsWorkflow:
+    @workflow.run
+    async def run(self) -> None:
+        children_started = 0
+
+        async def child(index: int) -> None:
+            nonlocal children_started
+            handle = await workflow.start_child_workflow(
+                TaskGroupWaitingChild.run,
+                id=f"{workflow.info().workflow_id}-child-{index}",
+            )
+            children_started += 1
+            await handle
+
+        async def fail() -> None:
+            # Wait for child starts to be recorded before testing cancellation.
+            await workflow.wait_condition(lambda: children_started == 2)
+            raise ValueError("planned failure")
+
+        try:
+            async with workflow.TaskGroup() as group:
+                for index in range(2):
+                    group.create_task(workflow.sleep(1000))
+                    group.create_task(
+                        workflow.execute_activity(
+                            task_group_waiting_activity,
+                            start_to_close_timeout=timedelta(minutes=5),
+                        )
+                    )
+                    group.create_task(child(index))
+                group.create_task(fail())
+        except ExceptionGroup as err:
+            # Temporal wraps activity and child cancellations in failure exceptions.
+            failures = [
+                error
+                for error in err.exceptions
+                if not temporalio.exceptions.is_cancelled_exception(error)
+            ]
+            assert len(failures) == 1
+            assert isinstance(failures[0], ValueError)
+            assert str(failures[0]) == "planned failure"
+        await workflow.sleep(0.01)
+
+
+@pytest.mark.parametrize("max_cached_workflows", [0, 1000])
+async def test_task_group_mixed_operations_replay(
+    client: Client, max_cached_workflows: int
+):
+    async with Worker(
+        client,
+        task_queue=str(uuid.uuid4()),
+        workflows=[TaskGroupMixedOperationsWorkflow, TaskGroupWaitingChild],
+        activities=[task_group_waiting_activity],
+        max_cached_workflows=max_cached_workflows,
+    ) as worker:
+        handle = await client.start_workflow(
+            TaskGroupMixedOperationsWorkflow.run,
+            id=str(uuid.uuid4()),
+            task_queue=worker.task_queue,
+        )
+        await handle.result()
+        history = await handle.fetch_history()
+
+    cancel_types = {
+        EventType.EVENT_TYPE_TIMER_CANCELED,
+        EventType.EVENT_TYPE_ACTIVITY_TASK_CANCEL_REQUESTED,
+        EventType.EVENT_TYPE_REQUEST_CANCEL_EXTERNAL_WORKFLOW_EXECUTION_INITIATED,
+    }
+    assert [
+        event.event_type for event in history.events if event.event_type in cancel_types
+    ] == [
+        EventType.EVENT_TYPE_TIMER_CANCELED,
+        EventType.EVENT_TYPE_ACTIVITY_TASK_CANCEL_REQUESTED,
+        EventType.EVENT_TYPE_REQUEST_CANCEL_EXTERNAL_WORKFLOW_EXECUTION_INITIATED,
+    ] * 2
+    for _ in range(3):
+        await Replayer(
+            workflows=[TaskGroupMixedOperationsWorkflow, TaskGroupWaitingChild]
+        ).replay_workflow(history)