hatchet-sdk 1.6.3__tar.gz → 1.6.4__tar.gz

{hatchet_sdk-1.6.3 → hatchet_sdk-1.6.4}/PKG-INFO

@@ -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 → hatchet_sdk-1.6.4}/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-1.6.3 → hatchet_sdk-1.6.4}/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-1.6.3 → hatchet_sdk-1.6.4}/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-1.6.3 → hatchet_sdk-1.6.4}/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-1.6.3 → hatchet_sdk-1.6.4}/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 → hatchet_sdk-1.6.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "hatchet-sdk"
-version = "1.6.3"
+version = "1.6.4"
 description = ""
 authors = ["Alexander Belanger <alexander@hatchet.run>"]
 readme = "README.md"