poe format

nagl-temporal · nagl-temporal · commit e6772d4e5c2d · 2025-04-01T16:56:56.000-07:00
diff --git a/resource_locking/lock_manager_workflow.py b/resource_locking/lock_manager_workflow.py
@@ -3,37 +3,44 @@
 
 from temporalio import workflow
 
+
 @dataclass
 class AssignedResource:
     resource: str
 
+
 @dataclass
 class AcquireRequest:
     workflow_id: str
     run_id: str
 
+
 @dataclass
 class ReleaseRequest:
     resource: str
     workflow_id: str
     run_id: str
 
+
 @dataclass
 class HandoffRequest:
     resource: str
     workflow_id: str
     old_run_id: str
     new_run_id: str
 
+
 LOCK_MANAGER_WORKFLOW_ID = "lock_manager"
 
+
 @dataclass
 class LockManagerWorkflowInput:
     # Key is resource, value is users of the resource. The first item in each list is the current holder of the lease
     # on that resource. A similar data structure could allow for multiple holders (perhaps the first n items are the
     # current holders).
     resource_queues: dict[str, list[AcquireRequest]]
 
+
 @workflow.defn
 class LockManagerWorkflow:
     def __init__(self):
@@ -47,73 +54,117 @@ async def acquire_resource(self, request: AcquireRequest):
         for resource in self.resource_queues:
             # Naively give out the first free resource, if we have one
             if len(self.resource_queues[resource]) == 0:
-                workflow.logger.info(f"workflow_id={request.workflow_id} run_id={request.run_id} acquired resource {resource}")
+                workflow.logger.info(
+                    f"workflow_id={request.workflow_id} run_id={request.run_id} acquired resource {resource}"
+                )
                 self.resource_queues[resource].append(request)
-                requester = workflow.get_external_workflow_handle(request.workflow_id, run_id=request.run_id)
+                requester = workflow.get_external_workflow_handle(
+                    request.workflow_id, run_id=request.run_id
+                )
                 await requester.signal("assign_resource", AssignedResource(resource))
                 return
 
         # Otherwise put this resource in a random queue.
         resource = workflow.random().choice(list(self.resource_queues.keys()))
-        workflow.logger.info(f"workflow_id={request.workflow_id} run_id={request.run_id} is waiting for resource {resource}")
+        workflow.logger.info(
+            f"workflow_id={request.workflow_id} run_id={request.run_id} is waiting for resource {resource}"
+        )
         self.resource_queues[resource].append(request)
 
     @workflow.signal
     async def release_resource(self, request: ReleaseRequest):
         queue = self.resource_queues[request.resource]
         if queue is None:
-            workflow.logger.warning(f"Ignoring request from {request.workflow_id} to release non-existent resource: {request.resource}")
+            workflow.logger.warning(
+                f"Ignoring request from {request.workflow_id} to release non-existent resource: {request.resource}"
+            )
             return
 
         if len(queue) == 0:
-            workflow.logger.warning(f"Ignoring request from {request.workflow_id} to release resource that is not held: {request.resource}")
+            workflow.logger.warning(
+                f"Ignoring request from {request.workflow_id} to release resource that is not held: {request.resource}"
+            )
             return
 
         holder = queue[0]
-        if not (holder.workflow_id == request.workflow_id and holder.run_id == request.run_id):
-            workflow.logger.warning(f"Ignoring request from non-holder to release resource {request.resource}")
-            workflow.logger.warning(f"resource is currently held by wf_id={holder.workflow_id} run_id={holder.run_id}")
-            workflow.logger.warning(f"request was from wf_id={request.workflow_id} run_id={request.run_id}")
+        if not (
+            holder.workflow_id == request.workflow_id
+            and holder.run_id == request.run_id
+        ):
+            workflow.logger.warning(
+                f"Ignoring request from non-holder to release resource {request.resource}"
+            )
+            workflow.logger.warning(
+                f"resource is currently held by wf_id={holder.workflow_id} run_id={holder.run_id}"
+            )
+            workflow.logger.warning(
+                f"request was from wf_id={request.workflow_id} run_id={request.run_id}"
+            )
             return
 
         # Remove the current holder from the head of the queue
-        workflow.logger.info(f"workflow_id={request.workflow_id} run_id={request.run_id} released resource {request.resource}")
+        workflow.logger.info(
+            f"workflow_id={request.workflow_id} run_id={request.run_id} released resource {request.resource}"
+        )
         queue = queue[1:]
         self.resource_queues[request.resource] = queue
 
         # If there are queued requests, assign the resource to the next one
         if len(queue) > 0:
             next_holder = queue[0]
-            workflow.logger.info(f"workflow_id={next_holder.workflow_id} run_id={next_holder.run_id} acquired resource {request.resource} after waiting")
-            requester = workflow.get_external_workflow_handle(next_holder.workflow_id, run_id=next_holder.run_id)
-            await requester.signal("assign_resource", AssignedResource(request.resource))
+            workflow.logger.info(
+                f"workflow_id={next_holder.workflow_id} run_id={next_holder.run_id} acquired resource {request.resource} after waiting"
+            )
+            requester = workflow.get_external_workflow_handle(
+                next_holder.workflow_id, run_id=next_holder.run_id
+            )
+            await requester.signal(
+                "assign_resource", AssignedResource(request.resource)
+            )
 
     @workflow.signal
     async def handoff_resource(self, request: HandoffRequest):
         queue = self.resource_queues[request.resource]
         if queue is None:
-            workflow.logger.warning(f"Ignoring request from {request.workflow_id} to hand off non-existent resource: {request.resource}")
+            workflow.logger.warning(
+                f"Ignoring request from {request.workflow_id} to hand off non-existent resource: {request.resource}"
+            )
             return
 
         if len(queue) == 0:
-            workflow.logger.warning(f"Ignoring request from {request.workflow_id} to hand off resource that is not held: {request.resource}")
+            workflow.logger.warning(
+                f"Ignoring request from {request.workflow_id} to hand off resource that is not held: {request.resource}"
+            )
             return
 
         holder = queue[0]
-        if not (holder.workflow_id == request.workflow_id and holder.run_id == request.old_run_id):
-            workflow.logger.warning(f"Ignoring request from non-holder to hand off resource {request.resource}")
-            workflow.logger.warning(f"resource is currently held by wf_id={holder.workflow_id} run_id={holder.run_id}")
-            workflow.logger.warning(f"request was from wf_id={request.workflow_id} run_id={request.old_run_id}")
+        if not (
+            holder.workflow_id == request.workflow_id
+            and holder.run_id == request.old_run_id
+        ):
+            workflow.logger.warning(
+                f"Ignoring request from non-holder to hand off resource {request.resource}"
+            )
+            workflow.logger.warning(
+                f"resource is currently held by wf_id={holder.workflow_id} run_id={holder.run_id}"
+            )
+            workflow.logger.warning(
+                f"request was from wf_id={request.workflow_id} run_id={request.old_run_id}"
+            )
             return
 
-        workflow.logger.info(f"workflow_id={request.workflow_id} handed off resource {request.resource} from run_id={request.old_run_id} to run_id={request.new_run_id}")
+        workflow.logger.info(
+            f"workflow_id={request.workflow_id} handed off resource {request.resource} from run_id={request.old_run_id} to run_id={request.new_run_id}"
+        )
         queue[0] = AcquireRequest(request.workflow_id, request.new_run_id)
-        requester = workflow.get_external_workflow_handle(request.workflow_id, run_id=request.new_run_id)
+        requester = workflow.get_external_workflow_handle(
+            request.workflow_id, run_id=request.new_run_id
+        )
         await requester.signal("assign_resource", AssignedResource(request.resource))
 
     @workflow.query
     def get_current_holders(self) -> dict[str, AcquireRequest]:
-        return { k: v[0] if v else None for k, v in self.resource_queues.items() }
+        return {k: v[0] if v else None for k, v in self.resource_queues.items()}
 
     @workflow.run
     async def run(self, input: LockManagerWorkflowInput) -> None:
@@ -125,4 +176,4 @@ async def run(self, input: LockManagerWorkflowInput) -> None:
             timeout=timedelta(hours=12),
         )
 
-        workflow.continue_as_new(LockManagerWorkflowInput(self.resource_queues))
+        workflow.continue_as_new(LockManagerWorkflowInput(self.resource_queues))
diff --git a/resource_locking/resource_locking_workflow.py b/resource_locking/resource_locking_workflow.py
@@ -5,20 +5,32 @@
 
 from temporalio import activity, workflow
 
-from resource_locking.lock_manager_workflow import AssignedResource, LOCK_MANAGER_WORKFLOW_ID, \
-    ReleaseRequest, AcquireRequest, HandoffRequest
+from resource_locking.lock_manager_workflow import (
+    LOCK_MANAGER_WORKFLOW_ID,
+    AcquireRequest,
+    AssignedResource,
+    HandoffRequest,
+    ReleaseRequest,
+)
+
 
 @dataclass
 class UseResourceActivityInput:
     resource: str
     iteration: str
 
+
 @activity.defn
 async def use_resource(input: UseResourceActivityInput) -> None:
     info = activity.info()
-    activity.logger.info(f"{info.workflow_id} starts using {input.resource} the {input.iteration} time")
+    activity.logger.info(
+        f"{info.workflow_id} starts using {input.resource} the {input.iteration} time"
+    )
     await asyncio.sleep(3)
-    activity.logger.info(f"{info.workflow_id} done using {input.resource} the {input.iteration} time")
+    activity.logger.info(
+        f"{info.workflow_id} done using {input.resource} the {input.iteration} time"
+    )
+
 
 @dataclass
 class ResourceLockingWorkflowInput:
@@ -32,17 +44,17 @@ class ResourceLockingWorkflowInput:
     # Used to transfer resource ownership between iterations during continue_as_new
     already_owned_resource: Optional[str]
 
+
 class FailWorkflowException(Exception):
     pass
 
+
 # Wait this long for a resource before giving up
 MAX_RESOURCE_WAIT_TIME = timedelta(minutes=5)
 
-@workflow.defn(
-  failure_exception_types=[FailWorkflowException]
-)
-class ResourceLockingWorkflow:
 
+@workflow.defn(failure_exception_types=[FailWorkflowException])
+class ResourceLockingWorkflow:
     def __init__(self):
         self.assigned_resource: Optional[str] = None
 
@@ -55,24 +67,42 @@ async def run(self, input: ResourceLockingWorkflowInput):
         workflow.info()
         if has_timeout(workflow.info().run_timeout):
             # See "locking" comment below for rationale
-            raise FailWorkflowException(f"ResourceLockingWorkflow cannot have a run_timeout (found {workflow.info().run_timeout})")
+            raise FailWorkflowException(
+                f"ResourceLockingWorkflow cannot have a run_timeout (found {workflow.info().run_timeout})"
+            )
         if has_timeout(workflow.info().execution_timeout):
-            raise FailWorkflowException(f"ResourceLockingWorkflow cannot have an execution_timeout (found {workflow.info().execution_timeout})")
+            raise FailWorkflowException(
+                f"ResourceLockingWorkflow cannot have an execution_timeout (found {workflow.info().execution_timeout})"
+            )
 
         sem_handle = workflow.get_external_workflow_handle(LOCK_MANAGER_WORKFLOW_ID)
 
         info = workflow.info()
         if input.already_owned_resource is None:
-            await sem_handle.signal("acquire_resource", AcquireRequest(info.workflow_id, info.run_id))
+            await sem_handle.signal(
+                "acquire_resource", AcquireRequest(info.workflow_id, info.run_id)
+            )
         else:
             # If we continued as new, we already have a resource. We need to transfer ownership from our predecessor to
             # ourselves.
-            await sem_handle.signal("handoff_resource", HandoffRequest(input.already_owned_resource, info.workflow_id, info.continued_run_id, info.run_id))
+            await sem_handle.signal(
+                "handoff_resource",
+                HandoffRequest(
+                    input.already_owned_resource,
+                    info.workflow_id,
+                    info.continued_run_id,
+                    info.run_id,
+                ),
+            )
 
         # Both branches above should cause us to receive an "assign_resource" signal.
-        await workflow.wait_condition(lambda: self.assigned_resource is not None, timeout=MAX_RESOURCE_WAIT_TIME)
+        await workflow.wait_condition(
+            lambda: self.assigned_resource is not None, timeout=MAX_RESOURCE_WAIT_TIME
+        )
         if self.assigned_resource is None:
-            raise FailWorkflowException(f"No resource was assigned after {MAX_RESOURCE_WAIT_TIME}")
+            raise FailWorkflowException(
+                f"No resource was assigned after {MAX_RESOURCE_WAIT_TIME}"
+            )
 
         # From this point forward, we own the resource. Note that this is a lock, not a lease! Our finally block will
         # release the resource if an activity fails. This is why we asserted the lack of workflow-level timeouts
@@ -86,7 +116,9 @@ async def run(self, input: ResourceLockingWorkflowInput):
                 )
 
                 if iteration == input.iteration_to_fail_after:
-                    workflow.logger.info(f"Failing after iteration {input.iteration_to_fail_after}")
+                    workflow.logger.info(
+                        f"Failing after iteration {input.iteration_to_fail_after}"
+                    )
                     raise FailWorkflowException()
 
             if input.should_continue_as_new:
@@ -101,9 +133,15 @@ async def run(self, input: ResourceLockingWorkflowInput):
             # execution, but this code in this finally block will still run. It wouldn't successfully send the signal...
             # the if statement just avoids some warnings in the log.
             if not input.should_continue_as_new:
-                await sem_handle.signal("release_resource", ReleaseRequest(self.assigned_resource, info.workflow_id, info.run_id))
+                await sem_handle.signal(
+                    "release_resource",
+                    ReleaseRequest(
+                        self.assigned_resource, info.workflow_id, info.run_id
+                    ),
+                )
+
 
 def has_timeout(timeout: Optional[timedelta]) -> bool:
     # After continue_as_new, timeouts are 0, even if they were None before continue_as_new (and were not set in the
     # continue_as_new call).
-    return timeout is not None and timeout > timedelta(0)
+    return timeout is not None and timeout > timedelta(0)
diff --git a/resource_locking/starter.py b/resource_locking/starter.py
@@ -1,10 +1,17 @@
 import asyncio
 from typing import Any
 
-from temporalio.client import Client, WorkflowHandle, WorkflowFailureError
+from temporalio.client import Client, WorkflowFailureError, WorkflowHandle
 
-from resource_locking.resource_locking_workflow import ResourceLockingWorkflow, ResourceLockingWorkflowInput
-from resource_locking.lock_manager_workflow import LockManagerWorkflow, LockManagerWorkflowInput, LOCK_MANAGER_WORKFLOW_ID
+from resource_locking.lock_manager_workflow import (
+    LOCK_MANAGER_WORKFLOW_ID,
+    LockManagerWorkflow,
+    LockManagerWorkflowInput,
+)
+from resource_locking.resource_locking_workflow import (
+    ResourceLockingWorkflow,
+    ResourceLockingWorkflowInput,
+)
 
 
 async def main():
@@ -14,18 +21,24 @@ async def main():
     # Start the LockManagerWorkflow
     lock_manager_handle = await client.start_workflow(
         workflow=LockManagerWorkflow.run,
-        arg=LockManagerWorkflowInput({
-            "resource_a": [],
-            "resource_b": [],
-        }),
+        arg=LockManagerWorkflowInput(
+            {
+                "resource_a": [],
+                "resource_b": [],
+            }
+        ),
         id=LOCK_MANAGER_WORKFLOW_ID,
         task_queue="default",
     )
 
     # Start the ResourceLockingWorkflows
     resource_locking_handles: list[WorkflowHandle[Any, Any]] = []
     for i in range(0, 4):
-        input = ResourceLockingWorkflowInput(iteration_to_fail_after=None, should_continue_as_new=False, already_owned_resource=None)
+        input = ResourceLockingWorkflowInput(
+            iteration_to_fail_after=None,
+            should_continue_as_new=False,
+            already_owned_resource=None,
+        )
         if i == 0:
             input.should_continue_as_new = True
         if i == 1:
@@ -48,5 +61,6 @@ async def main():
     # Clean up after ourselves. In the real world, the lock manager workflow would run forever.
     await lock_manager_handle.terminate()
 
+
 if __name__ == "__main__":
     asyncio.run(main())
diff --git a/resource_locking/worker.py b/resource_locking/worker.py
