hatchet-sdk 1.6.3__py3-none-any.whl → 1.6.4__py3-none-any.whl

hatchet_sdk/clients/dispatcher/action_listener.py

@@ -88,6 +88,9 @@ class ActionType(str, Enum):
     START_GET_GROUP_KEY = "START_GET_GROUP_KEY"
 
 
+ActionKey = str
+
+
 class Action(BaseModel):
     worker_id: str
     tenant_id: str
@@ -141,6 +144,18 @@ class Action(BaseModel):
 
         return {k: v for k, v in attrs.items() if v}
 
+    @property
+    def key(self) -> ActionKey:
+        """
+        This key is used to uniquely identify a single step run by its id + retry count.
+        It's used when storing references to a task, a context, etc. in a dictionary so that
+        we can look up those items in the dictionary by a unique key.
+        """
+        if self.action_type == ActionType.START_GET_GROUP_KEY:
+            return f"{self.get_group_key_run_id}/{self.retry_count}"
+        else:
+            return f"{self.step_run_id}/{self.retry_count}"
+
 
 def parse_additional_metadata(additional_metadata: str) -> JSONSerializableMapping:
     try:

hatchet_sdk/clients/events.py

@@ -149,6 +149,7 @@ class EventClient:
             ).events
         )
 
+    @tenacity_retry
     def log(self, message: str, step_run_id: str) -> None:
         request = PutLogRequest(
             stepRunId=step_run_id,
@@ -158,6 +159,7 @@ class EventClient:
 
         self.client.PutLog(request, metadata=get_metadata(self.token))
 
+    @tenacity_retry
     def stream(self, data: str | bytes, step_run_id: str) -> None:
         if isinstance(data, str):
             data_bytes = data.encode("utf-8")

hatchet_sdk/context/context.py

@@ -127,15 +127,18 @@ class Context:
     def workflow_run_id(self) -> str:
         return self.action.workflow_run_id
 
+    def _set_cancellation_flag(self) -> None:
+        self.exit_flag = True
+
     def cancel(self) -> None:
         logger.debug("cancelling step...")
         self.runs_client.cancel(self.step_run_id)
-        self.exit_flag = True
+        self._set_cancellation_flag()
 
     async def aio_cancel(self) -> None:
         logger.debug("cancelling step...")
         await self.runs_client.aio_cancel(self.step_run_id)
-        self.exit_flag = True
+        self._set_cancellation_flag()
 
     # done returns true if the context has been cancelled
     def done(self) -> bool:

hatchet_sdk/opentelemetry/instrumentor.py

@@ -267,17 +267,17 @@ class HatchetInstrumentor(BaseInstrumentor):  # type: ignore[misc]
     ## IMPORTANT: Keep these types in sync with the wrapped method's signature
     async def _wrap_handle_cancel_action(
         self,
-        wrapped: Callable[[str], Coroutine[None, None, Exception | None]],
+        wrapped: Callable[[Action], Coroutine[None, None, Exception | None]],
         instance: Runner,
-        args: tuple[str],
+        args: tuple[Action],
         kwargs: Any,
     ) -> Exception | None:
-        step_run_id = args[0]
+        action = args[0]
 
         with self._tracer.start_as_current_span(
             "hatchet.cancel_step_run",
             attributes={
-                "hatchet.step_run_id": step_run_id,
+                "hatchet.step_run_id": action.step_run_id,
             },
         ):
             return await wrapped(*args, **kwargs)

hatchet_sdk/worker/runner/runner.py

@@ -14,7 +14,7 @@ from pydantic import BaseModel
 
 from hatchet_sdk.client import Client
 from hatchet_sdk.clients.admin import AdminClient
-from hatchet_sdk.clients.dispatcher.action_listener import Action, ActionType
+from hatchet_sdk.clients.dispatcher.action_listener import Action, ActionKey, ActionType
 from hatchet_sdk.clients.dispatcher.dispatcher import DispatcherClient
 from hatchet_sdk.clients.events import EventClient
 from hatchet_sdk.clients.listeners.durable_event_listener import DurableEventListener
@@ -69,15 +69,15 @@ class Runner:
         self.config = config
 
         self.slots = slots
-        self.tasks: dict[str, asyncio.Task[Any]] = {}  # Store run ids and futures
-        self.contexts: dict[str, Context] = {}  # Store run ids and contexts
+        self.tasks: dict[ActionKey, asyncio.Task[Any]] = {}  # Store run ids and futures
+        self.contexts: dict[ActionKey, Context] = {}  # Store run ids and contexts
         self.action_registry = action_registry
 
         self.event_queue = event_queue
 
         # The thread pool is used for synchronous functions which need to run concurrently
         self.thread_pool = ThreadPoolExecutor(max_workers=slots)
-        self.threads: Dict[str, Thread] = {}  # Store run ids and threads
+        self.threads: Dict[ActionKey, Thread] = {}  # Store run ids and threads
 
         self.killing = False
         self.handle_kill = handle_kill
@@ -118,9 +118,9 @@ class Runner:
                 logger.info(log)
                 asyncio.create_task(self.handle_start_step_run(action))
             case ActionType.CANCEL_STEP_RUN:
-                log = f"cancel: step run:  {action.action_id}/{action.step_run_id}"
+                log = f"cancel: step run:  {action.action_id}/{action.step_run_id}/{action.retry_count}"
                 logger.info(log)
-                asyncio.create_task(self.handle_cancel_action(action.step_run_id))
+                asyncio.create_task(self.handle_cancel_action(action))
             case ActionType.START_GET_GROUP_KEY:
                 log = f"run: get group key:  {action.action_id}/{action.get_group_key_run_id}"
                 logger.info(log)
@@ -129,11 +129,9 @@ class Runner:
                 log = f"unknown action type: {action.action_type}"
                 logger.error(log)
 
-    def step_run_callback(
-        self, action: Action, action_task: "Task[TWorkflowInput, R]"
-    ) -> Callable[[asyncio.Task[Any]], None]:
+    def step_run_callback(self, action: Action) -> Callable[[asyncio.Task[Any]], None]:
         def inner_callback(task: asyncio.Task[Any]) -> None:
-            self.cleanup_run_id(action.step_run_id)
+            self.cleanup_run_id(action.key)
 
             errored = False
             cancelled = task.cancelled()
@@ -182,7 +180,7 @@ class Runner:
         self, action: Action
     ) -> Callable[[asyncio.Task[Any]], None]:
         def inner_callback(task: asyncio.Task[Any]) -> None:
-            self.cleanup_run_id(action.get_group_key_run_id)
+            self.cleanup_run_id(action.key)
 
             errored = False
             cancelled = task.cancelled()
@@ -227,9 +225,9 @@ class Runner:
         self, ctx: Context, task: Task[TWorkflowInput, R], action: Action
     ) -> R:
         if action.step_run_id:
-            self.threads[action.step_run_id] = current_thread()
+            self.threads[action.key] = current_thread()
         elif action.get_group_key_run_id:
-            self.threads[action.get_group_key_run_id] = current_thread()
+            self.threads[action.key] = current_thread()
 
         return task.call(ctx)
 
@@ -239,7 +237,6 @@ class Runner:
         ctx: Context,
         task: Task[TWorkflowInput, R],
         action: Action,
-        run_id: str,
     ) -> R:
         ctx_step_run_id.set(action.step_run_id)
         ctx_workflow_run_id.set(action.workflow_run_id)
@@ -271,17 +268,17 @@ class Runner:
             )
             raise e
         finally:
-            self.cleanup_run_id(run_id)
+            self.cleanup_run_id(action.key)
 
-    def cleanup_run_id(self, run_id: str | None) -> None:
-        if run_id in self.tasks:
-            del self.tasks[run_id]
+    def cleanup_run_id(self, key: ActionKey) -> None:
+        if key in self.tasks:
+            del self.tasks[key]
 
-        if run_id in self.threads:
-            del self.threads[run_id]
+        if key in self.threads:
+            del self.threads[key]
 
-        if run_id in self.contexts:
-            del self.contexts[run_id]
+        if key in self.contexts:
+            del self.contexts[key]
 
     @overload
     def create_context(
@@ -321,7 +318,7 @@ class Runner:
                 action, True if action_func.is_durable else False
             )
 
-            self.contexts[action.step_run_id] = context
+            self.contexts[action.key] = context
             self.event_queue.put(
                 ActionEvent(
                     action=action,
@@ -333,13 +330,11 @@ class Runner:
 
             loop = asyncio.get_event_loop()
             task = loop.create_task(
-                self.async_wrapped_action_func(
-                    context, action_func, action, action.step_run_id
-                )
+                self.async_wrapped_action_func(context, action_func, action)
             )
 
-            task.add_done_callback(self.step_run_callback(action, action_func))
-            self.tasks[action.step_run_id] = task
+            task.add_done_callback(self.step_run_callback(action))
+            self.tasks[action.key] = task
 
             try:
                 await task
@@ -358,7 +353,7 @@ class Runner:
         action_name = action.action_id
         context = self.create_context(action)
 
-        self.contexts[action.get_group_key_run_id] = context
+        self.contexts[action.key] = context
 
         # Find the corresponding action function from the registry
         action_func = self.action_registry.get(action_name)
@@ -376,13 +371,11 @@ class Runner:
 
             loop = asyncio.get_event_loop()
             task = loop.create_task(
-                self.async_wrapped_action_func(
-                    context, action_func, action, action.get_group_key_run_id
-                )
+                self.async_wrapped_action_func(context, action_func, action)
             )
 
             task.add_done_callback(self.group_key_run_callback(action))
-            self.tasks[action.get_group_key_run_id] = task
+            self.tasks[action.key] = task
 
             try:
                 await task
@@ -421,35 +414,36 @@ class Runner:
             logger.exception(f"Failed to terminate thread: {e}")
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
-    async def handle_cancel_action(self, run_id: str) -> None:
+    async def handle_cancel_action(self, action: Action) -> None:
+        key = action.key
         try:
             # call cancel to signal the context to stop
-            if run_id in self.contexts:
-                context = self.contexts.get(run_id)
+            if key in self.contexts:
+                context = self.contexts.get(key)
 
                 if context:
-                    await context.aio_cancel()
+                    context._set_cancellation_flag()
 
             await asyncio.sleep(1)
 
-            if run_id in self.tasks:
-                future = self.tasks.get(run_id)
+            if key in self.tasks:
+                future = self.tasks.get(key)
 
                 if future:
                     future.cancel()
 
             # check if thread is still running, if so, print a warning
-            if run_id in self.threads:
-                thread = self.threads.get(run_id)
+            if key in self.threads:
+                thread = self.threads.get(key)
                 if thread and self.config.enable_force_kill_sync_threads:
                     self.force_kill_thread(thread)
                     await asyncio.sleep(1)
 
                 logger.warning(
-                    f"Thread {self.threads[run_id].ident} with run id {run_id} is still running after cancellation. This could cause the thread pool to get blocked and prevent new tasks from running."
+                    f"Thread {self.threads[key].ident} with key {key} is still running after cancellation. This could cause the thread pool to get blocked and prevent new tasks from running."
                 )
         finally:
-            self.cleanup_run_id(run_id)
+            self.cleanup_run_id(key)
 
     def serialize_output(self, output: Any) -> str:
         if isinstance(output, BaseModel):

{hatchet_sdk-1.6.3.dist-info → hatchet_sdk-1.6.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hatchet-sdk
-Version: 1.6.3
+Version: 1.6.4
 Summary: 
 License: MIT
 Author: Alexander Belanger

{hatchet_sdk-1.6.3.dist-info → hatchet_sdk-1.6.4.dist-info}/RECORD

@@ -1,10 +1,10 @@
 hatchet_sdk/__init__.py,sha256=LUj6VyGVSHCYTQTaoyiVhjyJLOfv6gMCmb-s4hRyISM,10031
 hatchet_sdk/client.py,sha256=tbOeMuaJmgpyYSQg8QUz_J4AdqRNvV9E0aEZpgsiZTE,2207
 hatchet_sdk/clients/admin.py,sha256=zZKqjqjiq6_D2S5INuJBXdkyh5Owk65FMz6V6PC5Ta4,17692
-hatchet_sdk/clients/dispatcher/action_listener.py,sha256=QPBWNnnOGODaKaHtDd99imf1PdqfxzvNwnSJ6M3C7-M,16408
+hatchet_sdk/clients/dispatcher/action_listener.py,sha256=IwI2DKXTaPjqekZW8i6LXw4GOTtI2YdqR8PLUIhxooA,16955
 hatchet_sdk/clients/dispatcher/dispatcher.py,sha256=IL-hDXG8Lzas9FieVuNr47E_3Gvpc-aL4Xu_l385Vp8,8140
 hatchet_sdk/clients/event_ts.py,sha256=MudFszIb9IcPKQYvBTzcatPkcWEy3nxbAtEQ0_NYxMg,2094
-hatchet_sdk/clients/events.py,sha256=sKqzGwkrW_AKWW0Q_rxfwo9jXV5EUpLx7fTQvCNHaVo,5443
+hatchet_sdk/clients/events.py,sha256=4l_xiBZN7R9sxK7PzPHRj6C7Mrg1xCP14AzkiQyzKZY,5483
 hatchet_sdk/clients/listeners/durable_event_listener.py,sha256=jpqnbZsuouWk3XaOIYL9apaGtVk65eKKq66eBP9klBs,4085
 hatchet_sdk/clients/listeners/pooled_listener.py,sha256=1rodfIeqmHRF-u-PB6cBJbOU8NrvToLTyGigJMydpGo,8496
 hatchet_sdk/clients/listeners/run_event_listener.py,sha256=rIjBLRF7d7FBoEq7RKbmbOA84lX_hHSU26trwnthqV8,10230
@@ -221,7 +221,7 @@ hatchet_sdk/clients/v1/api_client.py,sha256=mJQUZ3cOxlFJiwWKK5F8jBxcpNZ7A2292Huc
 hatchet_sdk/config.py,sha256=jJA76BOvVdfOQHy6TKclAvr2qyblcM-Pz5J-hVAdpQ4,3588
 hatchet_sdk/connection.py,sha256=B5gT5NL9BBB5-l9U_cN6pMlraQk880rEYMnqaK_dgL0,2590
 hatchet_sdk/context/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-hatchet_sdk/context/context.py,sha256=EH7gCJtizMae06bu-pHDFaxYE2n8VkEQTPgFTveKGfA,9425
+hatchet_sdk/context/context.py,sha256=djzppHK3VINAtnrmPwLJhpIqzSMA5qKAVd6qB1r6gfM,9518
 hatchet_sdk/context/worker_context.py,sha256=OVcEWvdT_Kpd0nlg61VAPUgIPSFzSLs0aSrXWj-1GX4,974
 hatchet_sdk/contracts/dispatcher_pb2.py,sha256=wO-x-NiHjMuC55RArpcvQX7QAjpSrjq1n70mFVwOD0Q,14543
 hatchet_sdk/contracts/dispatcher_pb2.pyi,sha256=iOcGfGtoyvmT58yCrQDvst_o0VPbq1-9rir-0jH_140,18559
@@ -254,7 +254,7 @@ hatchet_sdk/hatchet.py,sha256=kIRcvHiLgyOc1OY1uytNatUZ8qaYA4beqoKEiqFcRpE,21873
 hatchet_sdk/labels.py,sha256=nATgxWE3lFxRTnfISEpoIRLGbMfAZsHF4lZTuG4Mfic,182
 hatchet_sdk/logger.py,sha256=5uOr52T4mImSQm1QvWT8HvZFK5WfPNh3Y1cBQZRFgUQ,333
 hatchet_sdk/metadata.py,sha256=XkRbhnghJJGCdVvF-uzyGBcNaTqpeQ3uiQvNNP1wyBc,107
-hatchet_sdk/opentelemetry/instrumentor.py,sha256=uRAQ9cRMIMyGFRW9Hy1G4CgW2e-Z0R64XJ8CgP4jrxk,13807
+hatchet_sdk/opentelemetry/instrumentor.py,sha256=GbsMZ1c9s0VRE7wwq-Iz5U9sT9fl7ZdXJ_ddqOqO5B0,13815
 hatchet_sdk/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hatchet_sdk/rate_limit.py,sha256=TwbCuggiZaWpYuo4mjVLlE-z1OfQ2mRBiVvCSaG3lv4,3919
 hatchet_sdk/runnables/contextvars.py,sha256=6MDocAMmlyiRW37oQ1jyx10tAlJs-xgDjR3xPoPz05g,426
@@ -502,11 +502,11 @@ hatchet_sdk/waits.py,sha256=L2xZUcmrQX-pTVXWv1W8suMoYU_eA0uowpollauQmOM,3893
 hatchet_sdk/worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hatchet_sdk/worker/action_listener_process.py,sha256=KxS7-wBpfKnsq0LNSvk-MG442Lh60iQMy3VpD1FW3mU,11703
 hatchet_sdk/worker/runner/run_loop_manager.py,sha256=RNWKDCjR57nJ0LCoLUMi0_3pnmpqyo80mz_RaxHYGIc,3812
-hatchet_sdk/worker/runner/runner.py,sha256=J-gTpoYF9WphB1zZhpPJryCv10VjU-MiCdJgMYH-nds,17352
+hatchet_sdk/worker/runner/runner.py,sha256=z8ri-viK_avAfF6zgbVNBc-rztFDbxSwng3RHsof92w,17063
 hatchet_sdk/worker/runner/utils/capture_logs.py,sha256=nHRPSiDBqzhObM7i2X7t03OupVFnE7kQBdR2Ckgg-2w,2709
 hatchet_sdk/worker/worker.py,sha256=AdkYo167jNsQ3CAaxuBcqA8YjmnKEu33Lp9TKK15amg,16168
 hatchet_sdk/workflow_run.py,sha256=ZwH0HLFGFVXz6jbiqSv4w0Om2XuR52Tzzw6LH4y65jQ,2765
-hatchet_sdk-1.6.3.dist-info/METADATA,sha256=T2dHw4jxT41_A6K9MiDuNS4BKFnyRpGQJt8SUPc5E48,3635
-hatchet_sdk-1.6.3.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-hatchet_sdk-1.6.3.dist-info/entry_points.txt,sha256=Un_76pcLse-ZGBlwebhQpnTPyQrripeHW8J7qmEpGOk,1400
-hatchet_sdk-1.6.3.dist-info/RECORD,,
+hatchet_sdk-1.6.4.dist-info/METADATA,sha256=NumPfRGN0_oh4B0nBlGiBqe0l4K2BbOz1MEgq3W4ucs,3635
+hatchet_sdk-1.6.4.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+hatchet_sdk-1.6.4.dist-info/entry_points.txt,sha256=Un_76pcLse-ZGBlwebhQpnTPyQrripeHW8J7qmEpGOk,1400
+hatchet_sdk-1.6.4.dist-info/RECORD,,