Workflow Update and Signal handlers concurrency sample

README.md

@@ -52,6 +52,7 @@ Some examples require extra dependencies. See each sample's directory for specif
   * [hello_signal](hello/hello_signal.py) - Send signals to a workflow.
 <!-- Keep this list in alphabetical order -->
 * [activity_worker](activity_worker) - Use Python activities from a workflow in another language.
+* [safe_message_handlers](updates_and_signals/safe_message_handlers/) - Safely handling updates and signals.
 * [bedrock](bedrock) - Orchestrate a chatbot with Amazon Bedrock.
 * [cloud_export_to_parquet](cloud_export_to_parquet) - Set up schedule workflow to process exported files on an hourly basis
 * [context_propagation](context_propagation) - Context propagation through workflows/activities via interceptor.

tests/updates_and_signals/safe_message_handlers_test.py

@@ -0,0 +1,107 @@
+import uuid
+
+from temporalio import common, workflow
+from temporalio.client import Client, WorkflowUpdateFailedError
+from temporalio.exceptions import ApplicationError
+from temporalio.worker import Worker
+
+from updates_and_signals.safe_message_handlers.activities import (
+    allocate_nodes_to_job,
+    deallocate_nodes_for_job,
+    find_bad_nodes,
+)
+from updates_and_signals.safe_message_handlers.starter import do_cluster_lifecycle
+from updates_and_signals.safe_message_handlers.workflow import (
+    ClusterManagerAllocateNNodesToJobInput,
+    ClusterManagerInput,
+    ClusterManagerWorkflow,
+)
+
+
+async def test_safe_message_handlers(client: Client):
+    task_queue = f"tq-{uuid.uuid4()}"
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[ClusterManagerWorkflow],
+        activities=[allocate_nodes_to_job, deallocate_nodes_for_job, find_bad_nodes],
+    ):
+        cluster_manager_handle = await client.start_workflow(
+            ClusterManagerWorkflow.run,
+            ClusterManagerInput(),
+            id=f"ClusterManagerWorkflow-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+        await do_cluster_lifecycle(cluster_manager_handle, delay_seconds=1)
+        result = await cluster_manager_handle.result()
+        assert result.max_assigned_nodes == 12
+        assert result.num_currently_assigned_nodes == 0
+
+
+async def test_update_idempotency(client: Client):
+    task_queue = f"tq-{uuid.uuid4()}"
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[ClusterManagerWorkflow],
+        activities=[allocate_nodes_to_job, deallocate_nodes_for_job, find_bad_nodes],
+    ):
+        cluster_manager_handle = await client.start_workflow(
+            ClusterManagerWorkflow.run,
+            ClusterManagerInput(),
+            id=f"ClusterManagerWorkflow-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+
+        await cluster_manager_handle.signal(ClusterManagerWorkflow.start_cluster)
+
+        nodes_1 = await cluster_manager_handle.execute_update(
+            ClusterManagerWorkflow.allocate_n_nodes_to_job,
+            ClusterManagerAllocateNNodesToJobInput(num_nodes=5, job_name=f"jobby-job"),
+        )
+        # simulate that in calling it twice, the operation is idempotent
+        nodes_2 = await cluster_manager_handle.execute_update(
+            ClusterManagerWorkflow.allocate_n_nodes_to_job,
+            ClusterManagerAllocateNNodesToJobInput(num_nodes=5, job_name=f"jobby-job"),
+        )
+        # the second call should not allocate more nodes (it may return fewer if the health check finds bad nodes
+        # in between the two signals.)
+        assert nodes_1 >= nodes_2
+
+
+async def test_update_failure(client: Client):
+    task_queue = f"tq-{uuid.uuid4()}"
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[ClusterManagerWorkflow],
+        activities=[allocate_nodes_to_job, deallocate_nodes_for_job, find_bad_nodes],
+    ):
+        cluster_manager_handle = await client.start_workflow(
+            ClusterManagerWorkflow.run,
+            ClusterManagerInput(),
+            id=f"ClusterManagerWorkflow-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+
+        await cluster_manager_handle.signal(ClusterManagerWorkflow.start_cluster)
+
+        await cluster_manager_handle.execute_update(
+            ClusterManagerWorkflow.allocate_n_nodes_to_job,
+            ClusterManagerAllocateNNodesToJobInput(num_nodes=24, job_name=f"big-task"),
+        )
+        try:
+            # Try to allocate too many nodes
+            await cluster_manager_handle.execute_update(
+                ClusterManagerWorkflow.allocate_n_nodes_to_job,
+                ClusterManagerAllocateNNodesToJobInput(
+                    num_nodes=3, job_name=f"little-task"
+                ),
+            )
+        except WorkflowUpdateFailedError as e:
+            assert isinstance(e.cause, ApplicationError)
+            assert e.cause.message == "Cannot allocate 3 nodes; have only 1 available"
+        finally:
+            await cluster_manager_handle.signal(ClusterManagerWorkflow.shutdown_cluster)
+            result = await cluster_manager_handle.result()
+            assert result.num_currently_assigned_nodes + result.num_bad_nodes == 24

updates_and_signals/safe_message_handlers/README.md

@@ -0,0 +1,21 @@
+# Atomic message handlers
+
+This sample shows off important techniques for handling signals and updates, aka messages.  In particular, it illustrates how message handlers can interleave or not be completed before the workflow completes, and how you can manage that.
+
+* Here, using workflow.wait_condition, signal and update handlers will only operate when the workflow is within a certain state--between cluster_started and cluster_shutdown.
+* You can run start_workflow with an initializer signal that you want to run before anything else other than the workflow's constructor.  This pattern is known as "signal-with-start."
+* Message handlers can block and their actions can be interleaved with one another and with the main workflow.  This can easily cause bugs, so we use a lock to protect shared state from interleaved access.
+* Message handlers should also finish before the workflow run completes.  One option is to use a lock.
+* An "Entity" workflow, i.e. a long-lived workflow, periodically "continues as new".  It must do this to prevent its history from growing too large, and it passes its state to the next workflow.  You can check `workflow.info().is_continue_as_new_suggested()` to see when it's time.  Just make sure message handlers have finished before doing so.  
+* Message handlers can be made idempotent.  See update `ClusterManager.allocate_n_nodes_to_job`.
+
+To run, first see [README.md](../../README.md) for prerequisites.
+
+Then, run the following from this directory to run the sample:
+
+```bash
+poetry run python worker.py
+poetry run python starter.py
+```
+
+This will start a worker to run your workflow and activities, then start a ClusterManagerWorkflow and put it through its paces.

updates_and_signals/safe_message_handlers/activities.py

@@ -0,0 +1,45 @@
+import asyncio
+from dataclasses import dataclass
+from typing import List
+
+from temporalio import activity
+
+
+@dataclass
+class AllocateNodesToJobInput:
+    nodes: List[str]
+    job_name: str
+
+
+@activity.defn
+async def allocate_nodes_to_job(input: AllocateNodesToJobInput):
+    print(f"Assigning nodes {input.nodes} to job {input.job_name}")
+    await asyncio.sleep(0.1)
+
+
+@dataclass
+class DeallocateNodesForJobInput:
+    nodes: List[str]
+    job_name: str
+
+
+@activity.defn
+async def deallocate_nodes_for_job(input: DeallocateNodesForJobInput):
+    print(f"Deallocating nodes {input.nodes} from job {input.job_name}")
+    await asyncio.sleep(0.1)
+
+
+@dataclass
+class FindBadNodesInput:
+    nodes_to_check: List[str]
+
+
+@activity.defn
+async def find_bad_nodes(input: FindBadNodesInput) -> List[str]:
+    await asyncio.sleep(0.1)
+    bad_nodes = [n for n in input.nodes_to_check if int(n) % 5 == 0]
+    if bad_nodes:
+        print(f"Found bad nodes: {bad_nodes}")
+    else:
+        print("No new bad nodes found.")
+    return bad_nodes

updates_and_signals/safe_message_handlers/starter.py

@@ -0,0 +1,80 @@
+import argparse
+import asyncio
+import logging
+import uuid
+from typing import Optional
+
+from temporalio import client, common
+from temporalio.client import Client, WorkflowHandle
+
+from updates_and_signals.safe_message_handlers.workflow import (
+    ClusterManagerAllocateNNodesToJobInput,
+    ClusterManagerDeleteJobInput,
+    ClusterManagerInput,
+    ClusterManagerWorkflow,
+)
+
+
+async def do_cluster_lifecycle(wf: WorkflowHandle, delay_seconds: Optional[int] = None):
+
+    await wf.signal(ClusterManagerWorkflow.start_cluster)
+
+    allocation_updates = []
+    for i in range(6):
+        allocation_updates.append(
+            wf.execute_update(
+                ClusterManagerWorkflow.allocate_n_nodes_to_job,
+                ClusterManagerAllocateNNodesToJobInput(
+                    num_nodes=2, job_name=f"task-{i}"
+                ),
+            )
+        )
+    await asyncio.gather(*allocation_updates)
+
+    if delay_seconds:
+        await asyncio.sleep(delay_seconds)
+
+    deletion_updates = []
+    for i in range(6):
+        deletion_updates.append(
+            wf.execute_update(
+                ClusterManagerWorkflow.delete_job,
+                ClusterManagerDeleteJobInput(job_name=f"task-{i}"),
+            )
+        )
+    await asyncio.gather(*deletion_updates)
+
+    await wf.signal(ClusterManagerWorkflow.shutdown_cluster)
+
+
+async def main(should_test_continue_as_new: bool):
+    # Connect to Temporal
+    client = await Client.connect("localhost:7233")
+
+    cluster_manager_handle = await client.start_workflow(
+        ClusterManagerWorkflow.run,
+        ClusterManagerInput(test_continue_as_new=should_test_continue_as_new),
+        id=f"ClusterManagerWorkflow-{uuid.uuid4()}",
+        task_queue="atomic-message-handlers-task-queue",
+        id_reuse_policy=common.WorkflowIDReusePolicy.TERMINATE_IF_RUNNING,
+    )
+    delay_seconds = 10 if should_test_continue_as_new else 1
+    await do_cluster_lifecycle(cluster_manager_handle, delay_seconds=delay_seconds)
+    result = await cluster_manager_handle.result()
+    print(
+        f"Cluster shut down successfully.  It peaked at {result.max_assigned_nodes} assigned nodes ."
+        f" It had {result.num_currently_assigned_nodes} nodes assigned at the end."
+    )
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    parser = argparse.ArgumentParser(description="Atomic message handlers")
+    parser.add_argument(
+        "--test-continue-as-new",
+        help="Make the ClusterManagerWorkflow continue as new before shutting down",
+        action="store_true",
+        default=False,
+    )
+    args = parser.parse_args()
+    asyncio.run(main(args.test_continue_as_new))

updates_and_signals/safe_message_handlers/worker.py

@@ -0,0 +1,41 @@
+import asyncio
+import logging
+
+from temporalio import activity, common, workflow
+from temporalio.client import Client, WorkflowHandle
+from temporalio.worker import Worker
+
+from updates_and_signals.safe_message_handlers.workflow import (
+    ClusterManagerWorkflow,
+    allocate_nodes_to_job,
+    deallocate_nodes_for_job,
+    find_bad_nodes,
+)
+
+interrupt_event = asyncio.Event()
+
+
+async def main():
+    # Connect client
+    client = await Client.connect("localhost:7233")
+
+    async with Worker(
+        client,
+        task_queue="atomic-message-handlers-task-queue",
+        workflows=[ClusterManagerWorkflow],
+        activities=[allocate_nodes_to_job, deallocate_nodes_for_job, find_bad_nodes],
+    ):
+        # Wait until interrupted
+        logging.info("ClusterManagerWorkflow worker started, ctrl+c to exit")
+        await interrupt_event.wait()
+        logging.info("Shutting down")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    loop = asyncio.new_event_loop()
+    try:
+        loop.run_until_complete(main())
+    except KeyboardInterrupt:
+        interrupt_event.set()
+        loop.run_until_complete(loop.shutdown_asyncgens())