Use workflow init from samples

bedrock/entity/workflows.py

@@ -18,20 +18,15 @@ class BedrockParams:
 
 @workflow.defn
 class EntityBedrockWorkflow:
-    def __init__(self) -> None:
+
+    @workflow.init
+    def __init__(self, params: BedrockParams) -> None:
         # List to store prompt history
         self.conversation_history: List[Tuple[str, str]] = []
         self.prompt_queue: Deque[str] = deque()
         self.conversation_summary: Optional[str] = None
         self.continue_as_new_per_turns: int = 6
         self.chat_ended: bool = False
-
-    @workflow.run
-    async def run(
-        self,
-        params: BedrockParams,
-    ) -> str:
-
         if params and params.conversation_summary:
             self.conversation_history.append(
                 ("conversation_summary", params.conversation_summary)
@@ -42,6 +37,12 @@ async def run(
         if params and params.prompt_queue:
             self.prompt_queue.extend(params.prompt_queue)
 
+
+    @workflow.run
+    async def run(
+        self,
+        params: BedrockParams,
+    ) -> str:
         while True:
             workflow.logger.info("Waiting for prompts...")
 

message_passing/introduction/workflows.py

@@ -33,7 +33,7 @@ class GreetingWorkflow:
 
     # 👉 This Workflow does not use any async handlers and so cannot use any
     # Activities. It only supports two languages, whose greetings are hardcoded
-    # in the Workflow definition. See GreetingWorkflowWithAsyncHandler below for
+    # in the Workflow definition. See ClusterManagerWorkflow for
     # a Workflow that uses an async Update handler to call an Activity.
 
     def __init__(self) -> None:

message_passing/safe_message_handlers/workflow.py

@@ -65,12 +65,20 @@ class ClusterManagerAssignNodesToJobResult:
 # These updates must run atomically.
 @workflow.defn
 class ClusterManagerWorkflow:
-    def __init__(self) -> None:
-        self.state = ClusterManagerState()
+    @workflow.init
+    def __init__(self, input: ClusterManagerInput) -> None:
         # Protects workflow state from interleaved access
         self.nodes_lock = asyncio.Lock()
+
+        # If we're continuing as new, we want to start from where we left off.
+        self.state = input.state or ClusterManagerState()
+
+        # For testing, we want to be able to more frequently continue as new.
         self.max_history_length: Optional[int] = None
         self.sleep_interval_seconds: int = 600
+        if input.test_continue_as_new:
+            self.max_history_length = 120
+            self.sleep_interval_seconds = 1
 
     @workflow.signal
     async def start_cluster(self) -> None:
@@ -202,15 +210,6 @@ async def perform_health_checks(self) -> None:
                     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:
         if workflow.info().is_continue_as_new_suggested():
             return True
@@ -224,7 +223,6 @@ def should_continue_as_new(self) -> bool:
 
     @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: