From 35b476d083438b3933c77000843fe1ef32512ae4 Mon Sep 17 00:00:00 2001 From: Drew Hoskins <166441821+drewhoskins-temporal@users.noreply.github.com> Date: Tue, 23 Jul 2024 18:17:43 -0700 Subject: [PATCH] Workflow Update and Signal handlers concurrency sample (#123) * Atomic message handlers sample * Remove resize jobs to reduce code size * Misc polish * Add test * Format code * Continue as new * Formatting * Feedback, readme, restructure files and directories * Format * More feedback. Add test-continue-as-new flag. * Feedback; throw ApplicationFailures from update handlers * Formatting * __init__.py * Fix lint issues * Dan Feedback * More typehints * s/atomic/safe/ * Fix and demo idempotency * Compatibility with 3.8 * More feedback * Re-add tests * Fix flaky test * Improve update and tests * list -> set to fix a test * Return a struct rather than a raw value from the list for better hygiene * Remove test dependency on race conditions between health check and adding nodes. * Ruff linting * Use consistent verbs, improve health check * poe format * Minor sample improvements * Skip update tests under Java test server --------- Co-authored-by: Drew Hoskins Co-authored-by: Dan Davison --- README.md | 1 + polling/frequent/README.md | 9 +- polling/infrequent/README.md | 10 +- polling/periodic_sequence/README.md | 10 +- .../safe_message_handlers/workflow_test.py | 155 +++++++++++ updates_and_signals/__init__.py | 0 .../safe_message_handlers/README.md | 22 ++ .../safe_message_handlers/__init__.py | 0 .../safe_message_handlers/activities.py | 45 +++ .../safe_message_handlers/starter.py | 84 ++++++ .../safe_message_handlers/worker.py | 40 +++ .../safe_message_handlers/workflow.py | 256 ++++++++++++++++++ 12 files changed, 620 insertions(+), 12 deletions(-) create mode 100644 tests/updates_and_signals/safe_message_handlers/workflow_test.py create mode 100644 updates_and_signals/__init__.py create mode 100644 updates_and_signals/safe_message_handlers/README.md create mode 100644 updates_and_signals/safe_message_handlers/__init__.py create mode 100644 updates_and_signals/safe_message_handlers/activities.py create mode 100644 updates_and_signals/safe_message_handlers/starter.py create mode 100644 updates_and_signals/safe_message_handlers/worker.py create mode 100644 updates_and_signals/safe_message_handlers/workflow.py diff --git a/README.md b/README.md index 9620132d..cdf88c6f 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,7 @@ Some examples require extra dependencies. See each sample's directory for specif * [polling](polling) - Recommended implementation of an activity that needs to periodically poll an external resource waiting its successful completion. * [prometheus](prometheus) - Configure Prometheus metrics on clients/workers. * [pydantic_converter](pydantic_converter) - Data converter for using Pydantic models. +* [safe_message_handlers](updates_and_signals/safe_message_handlers/) - Safely handling updates and signals. * [schedules](schedules) - Demonstrates a Workflow Execution that occurs according to a schedule. * [sentry](sentry) - Report errors to Sentry. * [worker_specific_task_queues](worker_specific_task_queues) - Use unique task queues to ensure activities run on specific workers. diff --git a/polling/frequent/README.md b/polling/frequent/README.md index 33b6bc0a..a54ea267 100644 --- a/polling/frequent/README.md +++ b/polling/frequent/README.md @@ -8,10 +8,11 @@ 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 run_worker.py -poetry run python run_frequent.py -``` + poetry run python run_worker.py + +Then, in another terminal, run the following to execute the workflow: + + poetry run python run_frequent.py The Workflow will continue to poll the service and heartbeat on every iteration until it succeeds. diff --git a/polling/infrequent/README.md b/polling/infrequent/README.md index 7c0a3225..c1b06d7f 100644 --- a/polling/infrequent/README.md +++ b/polling/infrequent/README.md @@ -13,10 +13,12 @@ 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 run_worker.py -poetry run python run_infrequent.py -``` + poetry run python run_worker.py + +Then, in another terminal, run the following to execute the workflow: + + poetry run python run_infrequent.py + Since the test service simulates being _down_ for four polling attempts and then returns _OK_ on the fifth poll attempt, the Workflow will perform four Activity retries with a 60-second poll interval, and then return the service result on the successful fifth attempt. diff --git a/polling/periodic_sequence/README.md b/polling/periodic_sequence/README.md index 65fe0669..f8416588 100644 --- a/polling/periodic_sequence/README.md +++ b/polling/periodic_sequence/README.md @@ -8,10 +8,12 @@ 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 run_worker.py -poetry run python run_periodic.py -``` + poetry run python run_worker.py + +Then, in another terminal, run the following to execute the workflow: + + poetry run python run_periodic.py + This will start a Workflow and Child Workflow to periodically poll an Activity. The Parent Workflow is not aware about the Child Workflow calling Continue-As-New, and it gets notified when it completes (or fails). \ No newline at end of file diff --git a/tests/updates_and_signals/safe_message_handlers/workflow_test.py b/tests/updates_and_signals/safe_message_handlers/workflow_test.py new file mode 100644 index 00000000..852419ef --- /dev/null +++ b/tests/updates_and_signals/safe_message_handlers/workflow_test.py @@ -0,0 +1,155 @@ +import asyncio +import uuid + +import pytest +from temporalio.client import Client, WorkflowUpdateFailedError +from temporalio.exceptions import ApplicationError +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Worker + +from updates_and_signals.safe_message_handlers.activities import ( + assign_nodes_to_job, + find_bad_nodes, + unassign_nodes_for_job, +) +from updates_and_signals.safe_message_handlers.workflow import ( + ClusterManagerAssignNodesToJobInput, + ClusterManagerDeleteJobInput, + ClusterManagerInput, + ClusterManagerWorkflow, +) + + +async def test_safe_message_handlers(client: Client, env: WorkflowEnvironment): + if env.supports_time_skipping: + pytest.skip( + "Java test server: https://github.com/temporalio/sdk-java/issues/1903" + ) + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ClusterManagerWorkflow], + activities=[assign_nodes_to_job, unassign_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) + + allocation_updates = [] + for i in range(6): + allocation_updates.append( + cluster_manager_handle.execute_update( + ClusterManagerWorkflow.assign_nodes_to_job, + ClusterManagerAssignNodesToJobInput( + total_num_nodes=2, job_name=f"task-{i}" + ), + ) + ) + results = await asyncio.gather(*allocation_updates) + for result in results: + assert len(result.nodes_assigned) == 2 + + await asyncio.sleep(1) + + deletion_updates = [] + for i in range(6): + deletion_updates.append( + cluster_manager_handle.execute_update( + ClusterManagerWorkflow.delete_job, + ClusterManagerDeleteJobInput(job_name=f"task-{i}"), + ) + ) + await asyncio.gather(*deletion_updates) + + await cluster_manager_handle.signal(ClusterManagerWorkflow.shutdown_cluster) + + result = await cluster_manager_handle.result() + assert result.num_currently_assigned_nodes == 0 + + +async def test_update_idempotency(client: Client, env: WorkflowEnvironment): + if env.supports_time_skipping: + pytest.skip( + "Java test server: https://github.com/temporalio/sdk-java/issues/1903" + ) + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ClusterManagerWorkflow], + activities=[assign_nodes_to_job, unassign_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) + + result_1 = await cluster_manager_handle.execute_update( + ClusterManagerWorkflow.assign_nodes_to_job, + ClusterManagerAssignNodesToJobInput( + total_num_nodes=5, job_name="jobby-job" + ), + ) + # simulate that in calling it twice, the operation is idempotent + result_2 = await cluster_manager_handle.execute_update( + ClusterManagerWorkflow.assign_nodes_to_job, + ClusterManagerAssignNodesToJobInput( + total_num_nodes=5, job_name="jobby-job" + ), + ) + # the second call should not assign more nodes (it may return fewer if the health check finds bad nodes + # in between the two signals.) + assert result_1.nodes_assigned >= result_2.nodes_assigned + + +async def test_update_failure(client: Client, env: WorkflowEnvironment): + if env.supports_time_skipping: + pytest.skip( + "Java test server: https://github.com/temporalio/sdk-java/issues/1903" + ) + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ClusterManagerWorkflow], + activities=[assign_nodes_to_job, unassign_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.assign_nodes_to_job, + ClusterManagerAssignNodesToJobInput( + total_num_nodes=24, job_name="big-task" + ), + ) + try: + # Try to assign too many nodes + await cluster_manager_handle.execute_update( + ClusterManagerWorkflow.assign_nodes_to_job, + ClusterManagerAssignNodesToJobInput( + total_num_nodes=3, job_name="little-task" + ), + ) + except WorkflowUpdateFailedError as e: + assert isinstance(e.cause, ApplicationError) + assert e.cause.message == "Cannot assign 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 diff --git a/updates_and_signals/__init__.py b/updates_and_signals/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/updates_and_signals/safe_message_handlers/README.md b/updates_and_signals/safe_message_handlers/README.md new file mode 100644 index 00000000..44bb9457 --- /dev/null +++ b/updates_and_signals/safe_message_handlers/README.md @@ -0,0 +1,22 @@ +# 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.assign_nodes_to_job`. + +To run, first see [README.md](../../README.md) for prerequisites. + +Then, run the following from this directory to run the worker: +\ + poetry run python worker.py + +Then, in another terminal, run the following to execute the workflow: + + 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. diff --git a/updates_and_signals/safe_message_handlers/__init__.py b/updates_and_signals/safe_message_handlers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/updates_and_signals/safe_message_handlers/activities.py b/updates_and_signals/safe_message_handlers/activities.py new file mode 100644 index 00000000..3a1c9cd2 --- /dev/null +++ b/updates_and_signals/safe_message_handlers/activities.py @@ -0,0 +1,45 @@ +import asyncio +from dataclasses import dataclass +from typing import List, Set + +from temporalio import activity + + +@dataclass +class AssignNodesToJobInput: + nodes: List[str] + job_name: str + + +@activity.defn +async def assign_nodes_to_job(input: AssignNodesToJobInput) -> None: + print(f"Assigning nodes {input.nodes} to job {input.job_name}") + await asyncio.sleep(0.1) + + +@dataclass +class UnassignNodesForJobInput: + nodes: List[str] + job_name: str + + +@activity.defn +async def unassign_nodes_for_job(input: UnassignNodesForJobInput) -> None: + print(f"Deallocating nodes {input.nodes} from job {input.job_name}") + await asyncio.sleep(0.1) + + +@dataclass +class FindBadNodesInput: + nodes_to_check: Set[str] + + +@activity.defn +async def find_bad_nodes(input: FindBadNodesInput) -> Set[str]: + await asyncio.sleep(0.1) + bad_nodes = set([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 diff --git a/updates_and_signals/safe_message_handlers/starter.py b/updates_and_signals/safe_message_handlers/starter.py new file mode 100644 index 00000000..adb56a66 --- /dev/null +++ b/updates_and_signals/safe_message_handlers/starter.py @@ -0,0 +1,84 @@ +import argparse +import asyncio +import logging +import uuid +from typing import Optional + +from temporalio import common +from temporalio.client import Client, WorkflowHandle + +from updates_and_signals.safe_message_handlers.workflow import ( + ClusterManagerAssignNodesToJobInput, + ClusterManagerDeleteJobInput, + ClusterManagerInput, + ClusterManagerWorkflow, +) + + +async def do_cluster_lifecycle(wf: WorkflowHandle, delay_seconds: Optional[int] = None): + + await wf.signal(ClusterManagerWorkflow.start_cluster) + + print("Assigning jobs to nodes...") + allocation_updates = [] + for i in range(6): + allocation_updates.append( + wf.execute_update( + ClusterManagerWorkflow.assign_nodes_to_job, + ClusterManagerAssignNodesToJobInput( + total_num_nodes=2, job_name=f"task-{i}" + ), + ) + ) + await asyncio.gather(*allocation_updates) + + print(f"Sleeping for {delay_seconds} second(s)") + if delay_seconds: + await asyncio.sleep(delay_seconds) + + print("Deleting jobs...") + 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") + + print("Starting cluster") + 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="safe-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." + 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)) diff --git a/updates_and_signals/safe_message_handlers/worker.py b/updates_and_signals/safe_message_handlers/worker.py new file mode 100644 index 00000000..5e28eca3 --- /dev/null +++ b/updates_and_signals/safe_message_handlers/worker.py @@ -0,0 +1,40 @@ +import asyncio +import logging + +from temporalio.client import Client +from temporalio.worker import Worker + +from updates_and_signals.safe_message_handlers.workflow import ( + ClusterManagerWorkflow, + assign_nodes_to_job, + find_bad_nodes, + unassign_nodes_for_job, +) + +interrupt_event = asyncio.Event() + + +async def main(): + # Connect client + client = await Client.connect("localhost:7233") + + async with Worker( + client, + task_queue="safe-message-handlers-task-queue", + workflows=[ClusterManagerWorkflow], + activities=[assign_nodes_to_job, unassign_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()) diff --git a/updates_and_signals/safe_message_handlers/workflow.py b/updates_and_signals/safe_message_handlers/workflow.py new file mode 100644 index 00000000..b8230578 --- /dev/null +++ b/updates_and_signals/safe_message_handlers/workflow.py @@ -0,0 +1,256 @@ +import asyncio +import dataclasses +from dataclasses import dataclass +from datetime import timedelta +from typing import Dict, List, Optional, Set + +from temporalio import workflow +from temporalio.common import RetryPolicy +from temporalio.exceptions import ApplicationError + +from updates_and_signals.safe_message_handlers.activities import ( + AssignNodesToJobInput, + FindBadNodesInput, + UnassignNodesForJobInput, + assign_nodes_to_job, + find_bad_nodes, + unassign_nodes_for_job, +) + + +# In workflows that continue-as-new, it's convenient to store all your state in one serializable structure +# to make it easier to pass between runs +@dataclass +class ClusterManagerState: + cluster_started: bool = False + cluster_shutdown: bool = False + nodes: Dict[str, Optional[str]] = dataclasses.field(default_factory=dict) + jobs_assigned: Set[str] = dataclasses.field(default_factory=set) + + +@dataclass +class ClusterManagerInput: + state: Optional[ClusterManagerState] = None + test_continue_as_new: bool = False + + +@dataclass +class ClusterManagerResult: + num_currently_assigned_nodes: int + num_bad_nodes: int + + +# Be in the habit of storing message inputs and outputs in serializable structures. +# This makes it easier to add more over time in a backward-compatible way. +@dataclass +class ClusterManagerAssignNodesToJobInput: + # If larger or smaller than previous amounts, will resize the job. + total_num_nodes: int + job_name: str + + +@dataclass +class ClusterManagerDeleteJobInput: + job_name: str + + +@dataclass +class ClusterManagerAssignNodesToJobResult: + nodes_assigned: Set[str] + + +# ClusterManagerWorkflow keeps track of the assignments of a cluster of nodes. +# Via signals, the cluster can be started and shutdown. +# Via updates, clients can also assign jobs to nodes and delete jobs. +# These updates must run atomically. +@workflow.defn +class ClusterManagerWorkflow: + def __init__(self) -> None: + self.state = ClusterManagerState() + # Protects workflow state from interleaved access + self.nodes_lock = asyncio.Lock() + self.max_history_length: Optional[int] = None + self.sleep_interval_seconds: int = 600 + + @workflow.signal + async def start_cluster(self) -> None: + self.state.cluster_started = True + self.state.nodes = {str(k): None for k in range(25)} + workflow.logger.info("Cluster started") + + @workflow.signal + async def shutdown_cluster(self) -> None: + await workflow.wait_condition(lambda: self.state.cluster_started) + self.state.cluster_shutdown = True + workflow.logger.info("Cluster shut down") + + # This is an update as opposed to a signal because the client may want to wait for nodes to be allocated + # before sending work to those nodes. + # Returns the list of node names that were allocated to the job. + @workflow.update + async def assign_nodes_to_job( + self, input: ClusterManagerAssignNodesToJobInput + ) -> ClusterManagerAssignNodesToJobResult: + await workflow.wait_condition(lambda: self.state.cluster_started) + if self.state.cluster_shutdown: + # If you want the client to receive a failure, either add an update validator and throw the + # exception from there, or raise an ApplicationError. Other exceptions in the main handler + # will cause the workflow to keep retrying and get it stuck. + raise ApplicationError( + "Cannot assign nodes to a job: Cluster is already shut down" + ) + + async with self.nodes_lock: + # Idempotency guard. + if input.job_name in self.state.jobs_assigned: + return ClusterManagerAssignNodesToJobResult( + self.get_assigned_nodes(job_name=input.job_name) + ) + unassigned_nodes = self.get_unassigned_nodes() + if len(unassigned_nodes) < input.total_num_nodes: + # If you want the client to receive a failure, either add an update validator and throw the + # exception from there, or raise an ApplicationError. Other exceptions in the main handler + # will cause the workflow to keep retrying and get it stuck. + raise ApplicationError( + f"Cannot assign {input.total_num_nodes} nodes; have only {len(unassigned_nodes)} available" + ) + nodes_to_assign = unassigned_nodes[: input.total_num_nodes] + # This await would be dangerous without nodes_lock because it yields control and allows interleaving + # with delete_job and perform_health_checks, which both touch self.state.nodes. + await self._assign_nodes_to_job(nodes_to_assign, input.job_name) + return ClusterManagerAssignNodesToJobResult( + nodes_assigned=self.get_assigned_nodes(job_name=input.job_name) + ) + + async def _assign_nodes_to_job( + self, assigned_nodes: List[str], job_name: str + ) -> None: + await workflow.execute_activity( + assign_nodes_to_job, + AssignNodesToJobInput(nodes=assigned_nodes, job_name=job_name), + start_to_close_timeout=timedelta(seconds=10), + ) + for node in assigned_nodes: + self.state.nodes[node] = job_name + self.state.jobs_assigned.add(job_name) + + # Even though it returns nothing, this is an update because the client may want to track it, for example + # to wait for nodes to be unassignd before reassigning them. + @workflow.update + async def delete_job(self, input: ClusterManagerDeleteJobInput) -> None: + await workflow.wait_condition(lambda: self.state.cluster_started) + if self.state.cluster_shutdown: + # If you want the client to receive a failure, either add an update validator and throw the + # exception from there, or raise an ApplicationError. Other exceptions in the main handler + # will cause the workflow to keep retrying and get it stuck. + raise ApplicationError("Cannot delete a job: Cluster is already shut down") + + async with self.nodes_lock: + nodes_to_unassign = [ + k for k, v in self.state.nodes.items() if v == input.job_name + ] + # This await would be dangerous without nodes_lock because it yields control and allows interleaving + # with assign_nodes_to_job and perform_health_checks, which all touch self.state.nodes. + await self._unassign_nodes_for_job(nodes_to_unassign, input.job_name) + + async def _unassign_nodes_for_job( + self, nodes_to_unassign: List[str], job_name: str + ): + await workflow.execute_activity( + unassign_nodes_for_job, + UnassignNodesForJobInput(nodes=nodes_to_unassign, job_name=job_name), + start_to_close_timeout=timedelta(seconds=10), + ) + for node in nodes_to_unassign: + self.state.nodes[node] = None + + def get_unassigned_nodes(self) -> List[str]: + return [k for k, v in self.state.nodes.items() if v is None] + + def get_bad_nodes(self) -> Set[str]: + return set([k for k, v in self.state.nodes.items() if v == "BAD!"]) + + def get_assigned_nodes(self, *, job_name: Optional[str] = None) -> Set[str]: + if job_name: + return set([k for k, v in self.state.nodes.items() if v == job_name]) + else: + return set( + [ + k + for k, v in self.state.nodes.items() + if v is not None and v != "BAD!" + ] + ) + + async def perform_health_checks(self) -> None: + async with self.nodes_lock: + assigned_nodes = self.get_assigned_nodes() + try: + # This await would be dangerous without nodes_lock because it yields control and allows interleaving + # with assign_nodes_to_job and delete_job, which both touch self.state.nodes. + bad_nodes = await workflow.execute_activity( + find_bad_nodes, + FindBadNodesInput(nodes_to_check=assigned_nodes), + start_to_close_timeout=timedelta(seconds=10), + # This health check is optional, and our lock would block the whole workflow if we let it retry forever. + retry_policy=RetryPolicy(maximum_attempts=1), + ) + for node in bad_nodes: + self.state.nodes[node] = "BAD!" + except Exception as e: + workflow.logger.warn( + f"Health check failed with error {type(e).__name__}:{e}" + ) + + # The cluster manager is a long-running "entity" workflow so we need to periodically checkpoint its state and + # continue-as-new. + def init(self, input: ClusterManagerInput) -> None: + if input.state: + self.state = input.state + if input.test_continue_as_new: + self.max_history_length = 120 + self.sleep_interval_seconds = 1 + + def should_continue_as_new(self) -> bool: + # We don't want to continue-as-new if we're in the middle of an update + if self.nodes_lock.locked(): + return False + if workflow.info().is_continue_as_new_suggested(): + return True + # This is just for ease-of-testing. In production, we trust temporal to tell us when to continue as new. + if ( + self.max_history_length + and workflow.info().get_current_history_length() > self.max_history_length + ): + return True + return False + + @workflow.run + async def run(self, input: ClusterManagerInput) -> ClusterManagerResult: + self.init(input) + await workflow.wait_condition(lambda: self.state.cluster_started) + # Perform health checks at intervals. + while True: + await self.perform_health_checks() + try: + await workflow.wait_condition( + lambda: self.state.cluster_shutdown + or self.should_continue_as_new(), + timeout=timedelta(seconds=self.sleep_interval_seconds), + ) + except asyncio.TimeoutError: + pass + if self.state.cluster_shutdown: + break + if self.should_continue_as_new(): + workflow.logger.info("Continuing as new") + workflow.continue_as_new( + ClusterManagerInput( + state=self.state, + test_continue_as_new=input.test_continue_as_new, + ) + ) + return ClusterManagerResult( + len(self.get_assigned_nodes()), + len(self.get_bad_nodes()), + )