hatchet-sdk 1.12.3__py3-none-any.whl → 1.14.0__py3-none-any.whl

hatchet_sdk/v0/v2/callable.py

@@ -11,9 +11,7 @@ from typing import (
     Union,
 )
 
-from hatchet_sdk.v0.clients.admin import ChildTriggerWorkflowOptions
-from hatchet_sdk.v0.context.context import Context
-from hatchet_sdk.v0.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
+from hatchet_sdk.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
     CreateStepRateLimit,
     CreateWorkflowJobOpts,
     CreateWorkflowStepOpts,
@@ -23,8 +21,10 @@ from hatchet_sdk.v0.contracts.workflows_pb2 import (  # type: ignore[attr-define
     WorkflowConcurrencyOpts,
     WorkflowKind,
 )
+from hatchet_sdk.logger import logger
+from hatchet_sdk.v0.clients.admin import ChildTriggerWorkflowOptions
+from hatchet_sdk.v0.context.context import Context
 from hatchet_sdk.v0.labels import DesiredWorkerLabel
-from hatchet_sdk.v0.logger import logger
 from hatchet_sdk.v0.rate_limit import RateLimit
 from hatchet_sdk.v0.v2.concurrency import ConcurrencyFunction
 from hatchet_sdk.v0.workflow_run import RunRef

hatchet_sdk/v0/v2/concurrency.py

@@ -1,9 +1,9 @@
 from typing import Any, Callable
 
-from hatchet_sdk.v0.context.context import Context
-from hatchet_sdk.v0.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
+from hatchet_sdk.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
     ConcurrencyLimitStrategy,
 )
+from hatchet_sdk.v0.context.context import Context
 
 
 class ConcurrencyFunction:

hatchet_sdk/v0/v2/hatchet.py

@@ -1,11 +1,11 @@
 from typing import Any, Callable, TypeVar, Union
 
-from hatchet_sdk.v0 import Worker
-from hatchet_sdk.v0.context.context import Context
-from hatchet_sdk.v0.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
+from hatchet_sdk.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
     ConcurrencyLimitStrategy,
     StickyStrategy,
 )
+from hatchet_sdk.v0 import Worker
+from hatchet_sdk.v0.context.context import Context
 from hatchet_sdk.v0.hatchet import Hatchet as HatchetV1
 from hatchet_sdk.v0.hatchet import workflow
 from hatchet_sdk.v0.labels import DesiredWorkerLabel

hatchet_sdk/v0/worker/action_listener_process.py

@@ -8,6 +8,12 @@ from typing import Any, List, Mapping, Optional
 
 import grpc
 
+from hatchet_sdk.contracts.dispatcher_pb2 import (
+    GROUP_KEY_EVENT_TYPE_STARTED,
+    STEP_EVENT_TYPE_STARTED,
+    ActionType,
+)
+from hatchet_sdk.logger import logger
 from hatchet_sdk.v0.client import Client, new_client_raw
 from hatchet_sdk.v0.clients.dispatcher.action_listener import Action
 from hatchet_sdk.v0.clients.dispatcher.dispatcher import (
@@ -16,13 +22,7 @@ from hatchet_sdk.v0.clients.dispatcher.dispatcher import (
     new_dispatcher,
 )
 from hatchet_sdk.v0.clients.rest.models.update_worker_request import UpdateWorkerRequest
-from hatchet_sdk.v0.contracts.dispatcher_pb2 import (
-    GROUP_KEY_EVENT_TYPE_STARTED,
-    STEP_EVENT_TYPE_STARTED,
-    ActionType,
-)
 from hatchet_sdk.v0.loader import ClientConfig
-from hatchet_sdk.v0.logger import logger
 from hatchet_sdk.v0.utils.backoff import exp_backoff_sleep
 
 ACTION_EVENT_RETRY_COUNT = 5

hatchet_sdk/v0/worker/runner/run_loop_manager.py

@@ -4,11 +4,11 @@ from dataclasses import dataclass, field
 from multiprocessing import Queue
 from typing import Callable, TypeVar
 
+from hatchet_sdk.logger import logger
 from hatchet_sdk.v0 import Context
 from hatchet_sdk.v0.client import Client, new_client_raw
 from hatchet_sdk.v0.clients.dispatcher.action_listener import Action
 from hatchet_sdk.v0.loader import ClientConfig
-from hatchet_sdk.v0.logger import logger
 from hatchet_sdk.v0.utils.types import WorkflowValidator
 from hatchet_sdk.v0.worker.runner.runner import Runner
 from hatchet_sdk.v0.worker.runner.utils.capture_logs import capture_logs

hatchet_sdk/v0/worker/runner/runner.py

@@ -13,15 +13,7 @@ from typing import Any, Callable, Dict, cast
 
 from pydantic import BaseModel
 
-from hatchet_sdk.v0.client import new_client_raw
-from hatchet_sdk.v0.clients.admin import new_admin
-from hatchet_sdk.v0.clients.dispatcher.action_listener import Action
-from hatchet_sdk.v0.clients.dispatcher.dispatcher import new_dispatcher
-from hatchet_sdk.v0.clients.run_event_listener import new_listener
-from hatchet_sdk.v0.clients.workflow_listener import PooledWorkflowRunListener
-from hatchet_sdk.v0.context import Context  # type: ignore[attr-defined]
-from hatchet_sdk.v0.context.worker_context import WorkerContext
-from hatchet_sdk.v0.contracts.dispatcher_pb2 import (
+from hatchet_sdk.contracts.dispatcher_pb2 import (
     GROUP_KEY_EVENT_TYPE_COMPLETED,
     GROUP_KEY_EVENT_TYPE_FAILED,
     GROUP_KEY_EVENT_TYPE_STARTED,
@@ -30,8 +22,16 @@ from hatchet_sdk.v0.contracts.dispatcher_pb2 import (
     STEP_EVENT_TYPE_STARTED,
     ActionType,
 )
+from hatchet_sdk.logger import logger
+from hatchet_sdk.v0.client import new_client_raw
+from hatchet_sdk.v0.clients.admin import new_admin
+from hatchet_sdk.v0.clients.dispatcher.action_listener import Action
+from hatchet_sdk.v0.clients.dispatcher.dispatcher import new_dispatcher
+from hatchet_sdk.v0.clients.run_event_listener import new_listener
+from hatchet_sdk.v0.clients.workflow_listener import PooledWorkflowRunListener
+from hatchet_sdk.v0.context import Context  # type: ignore[attr-defined]
+from hatchet_sdk.v0.context.worker_context import WorkerContext
 from hatchet_sdk.v0.loader import ClientConfig
-from hatchet_sdk.v0.logger import logger
 from hatchet_sdk.v0.utils.types import WorkflowValidator
 from hatchet_sdk.v0.v2.callable import DurableContext
 from hatchet_sdk.v0.worker.action_listener_process import ActionEvent

hatchet_sdk/v0/worker/runner/utils/capture_logs.py

@@ -5,7 +5,7 @@ from concurrent.futures import ThreadPoolExecutor
 from io import StringIO
 from typing import Any, Coroutine
 
-from hatchet_sdk.v0 import logger
+from hatchet_sdk import logger
 from hatchet_sdk.v0.clients.events import EventClient
 
 wr: contextvars.ContextVar[str | None] = contextvars.ContextVar(

hatchet_sdk/v0/worker/worker.py

@@ -17,11 +17,11 @@ from aiohttp.web_request import Request
 from aiohttp.web_response import Response
 from prometheus_client import CONTENT_TYPE_LATEST, Gauge, generate_latest
 
+from hatchet_sdk.contracts.workflows_pb2 import CreateWorkflowVersionOpts
+from hatchet_sdk.logger import logger
 from hatchet_sdk.v0 import Context
 from hatchet_sdk.v0.client import Client, new_client_raw
-from hatchet_sdk.v0.contracts.workflows_pb2 import CreateWorkflowVersionOpts
 from hatchet_sdk.v0.loader import ClientConfig
-from hatchet_sdk.v0.logger import logger
 from hatchet_sdk.v0.utils.types import WorkflowValidator
 from hatchet_sdk.v0.utils.typing import is_basemodel_subclass
 from hatchet_sdk.v0.v2.callable import HatchetCallable

hatchet_sdk/v0/workflow.py

@@ -13,8 +13,7 @@ from typing import (
 
 from pydantic import BaseModel
 
-from hatchet_sdk.v0 import ConcurrencyLimitStrategy
-from hatchet_sdk.v0.contracts.workflows_pb2 import (
+from hatchet_sdk.contracts.workflows_pb2 import (
     CreateWorkflowJobOpts,
     CreateWorkflowStepOpts,
     CreateWorkflowVersionOpts,
@@ -22,7 +21,8 @@ from hatchet_sdk.v0.contracts.workflows_pb2 import (
     WorkflowConcurrencyOpts,
     WorkflowKind,
 )
-from hatchet_sdk.v0.logger import logger
+from hatchet_sdk.logger import logger
+from hatchet_sdk.v0 import ConcurrencyLimitStrategy
 from hatchet_sdk.v0.utils.typing import is_basemodel_subclass
 
 

hatchet_sdk/waits.py

@@ -6,6 +6,7 @@ from uuid import uuid4
 
 from pydantic import BaseModel, Field
 
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.contracts.v1.shared.condition_pb2 import Action as ProtoAction
 from hatchet_sdk.contracts.v1.shared.condition_pb2 import (
     BaseMatchCondition,
@@ -53,7 +54,7 @@ class Condition(ABC):
 
     @abstractmethod
     def to_proto(
-        self,
+        self, config: ClientConfig
     ) -> UserEventMatchCondition | ParentOverrideMatchCondition | SleepMatchCondition:
         pass
 
@@ -71,7 +72,7 @@ class SleepCondition(Condition):
 
         self.duration = duration
 
-    def to_proto(self) -> SleepMatchCondition:
+    def to_proto(self, config: ClientConfig) -> SleepMatchCondition:
         return SleepMatchCondition(
             base=self.base.to_proto(),
             sleep_for=timedelta_to_expr(self.duration),
@@ -95,10 +96,10 @@ class UserEventCondition(Condition):
         self.event_key = event_key
         self.expression = expression
 
-    def to_proto(self) -> UserEventMatchCondition:
+    def to_proto(self, config: ClientConfig) -> UserEventMatchCondition:
         return UserEventMatchCondition(
             base=self.base.to_proto(),
-            user_event_key=self.event_key,
+            user_event_key=config.apply_namespace(self.event_key),
         )
 
 
@@ -124,7 +125,7 @@ class ParentCondition(Condition):
 
         self.parent = parent
 
-    def to_proto(self) -> ParentOverrideMatchCondition:
+    def to_proto(self, config: ClientConfig) -> ParentOverrideMatchCondition:
         return ParentOverrideMatchCondition(
             base=self.base.to_proto(),
             parent_readable_id=self.parent.name,

hatchet_sdk/worker/action_listener_process.py

@@ -4,9 +4,10 @@ import signal
 import time
 from dataclasses import dataclass
 from multiprocessing import Queue
-from typing import Any, Literal
+from typing import Any
 
 import grpc
+from grpc.aio import UnaryUnaryCall
 
 from hatchet_sdk.client import Client
 from hatchet_sdk.clients.dispatcher.action_listener import (
@@ -19,6 +20,9 @@ from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.contracts.dispatcher_pb2 import (
     GROUP_KEY_EVENT_TYPE_STARTED,
     STEP_EVENT_TYPE_STARTED,
+    ActionEventResponse,
+    GroupKeyActionEvent,
+    StepActionEvent,
 )
 from hatchet_sdk.logger import logger
 from hatchet_sdk.runnables.action import Action, ActionType
@@ -29,6 +33,7 @@ from hatchet_sdk.runnables.contextvars import (
     ctx_workflow_run_id,
 )
 from hatchet_sdk.utils.backoff import exp_backoff_sleep
+from hatchet_sdk.utils.typing import STOP_LOOP, STOP_LOOP_TYPE
 
 ACTION_EVENT_RETRY_COUNT = 5
 
@@ -41,9 +46,6 @@ class ActionEvent:
     should_not_retry: bool
 
 
-STOP_LOOP_TYPE = Literal["STOP_LOOP"]
-STOP_LOOP: STOP_LOOP_TYPE = "STOP_LOOP"  # Sentinel object to stop the loop
-
 BLOCKED_THREAD_WARNING = "THE TIME TO START THE TASK RUN IS TOO LONG, THE EVENT LOOP MAY BE BLOCKED. See https://docs.hatchet.run/blog/warning-event-loop-blocked for details and debugging help."
 
 
@@ -56,9 +58,9 @@ class WorkerActionListenerProcess:
         config: ClientConfig,
         action_queue: "Queue[Action]",
         event_queue: "Queue[ActionEvent | STOP_LOOP_TYPE]",
-        handle_kill: bool = True,
-        debug: bool = False,
-        labels: dict[str, str | int] = {},
+        handle_kill: bool,
+        debug: bool,
+        labels: dict[str, str | int],
     ) -> None:
         self.name = name
         self.actions = actions
@@ -75,6 +77,14 @@ class WorkerActionListenerProcess:
         self.action_loop_task: asyncio.Task[None] | None = None
         self.event_send_loop_task: asyncio.Task[None] | None = None
         self.running_step_runs: dict[str, float] = {}
+        self.step_action_events: set[
+            asyncio.Task[UnaryUnaryCall[StepActionEvent, ActionEventResponse] | None]
+        ] = set()
+        self.group_key_action_events: set[
+            asyncio.Task[
+                UnaryUnaryCall[GroupKeyActionEvent, ActionEventResponse] | None
+            ]
+        ] = set()
 
         if self.debug:
             logger.setLevel(logging.DEBUG)
@@ -144,20 +154,21 @@ class WorkerActionListenerProcess:
                 break
 
             logger.debug(f"tx: event: {event.action.action_id}/{event.type}")
-            asyncio.create_task(self.send_event(event))
+            t = asyncio.create_task(self.send_event(event))
+            self.step_action_events.add(t)
+            t.add_done_callback(lambda t: self.step_action_events.discard(t))
 
     async def start_blocked_main_loop(self) -> None:
         threshold = 1
         while not self.killing:
             count = 0
-            for _, start_time in self.running_step_runs.items():
+            for start_time in self.running_step_runs.values():
                 diff = self.now() - start_time
                 if diff > threshold:
                     count += 1
 
             if count > 0:
                 logger.warning(f"{BLOCKED_THREAD_WARNING}: Waiting Steps {count}")
-                print(asyncio.current_task())
             await asyncio.sleep(1)
 
     async def send_event(self, event: ActionEvent, retry_attempt: int = 1) -> None:
@@ -187,7 +198,7 @@ class WorkerActionListenerProcess:
                                 self.now()
                             )
 
-                    asyncio.create_task(
+                    send_started_event_task = asyncio.create_task(
                         self.dispatcher_client.send_step_action_event(
                             event.action,
                             event.type,
@@ -195,14 +206,23 @@ class WorkerActionListenerProcess:
                             event.should_not_retry,
                         )
                     )
+
+                    self.step_action_events.add(send_started_event_task)
+                    send_started_event_task.add_done_callback(
+                        lambda t: self.step_action_events.discard(t)
+                    )
                 case ActionType.CANCEL_STEP_RUN:
                     logger.debug("unimplemented event send")
                 case ActionType.START_GET_GROUP_KEY:
-                    asyncio.create_task(
+                    get_group_key_task = asyncio.create_task(
                         self.dispatcher_client.send_group_key_action_event(
                             event.action, event.type, event.payload
                         )
                     )
+                    self.group_key_action_events.add(get_group_key_task)
+                    get_group_key_task.add_done_callback(
+                        lambda t: self.group_key_action_events.discard(t)
+                    )
                 case _:
                     logger.error("unknown action type for event send")
         except Exception as e:
@@ -317,7 +337,7 @@ def worker_action_listener_process(*args: Any, **kwargs: Any) -> None:
         process = WorkerActionListenerProcess(*args, **kwargs)
         await process.start()
         # Keep the process running
-        while not process.killing:
+        while not process.killing:  # noqa: ASYNC110
             await asyncio.sleep(0.1)
 
     asyncio.run(run())

hatchet_sdk/worker/runner/run_loop_manager.py

@@ -1,19 +1,17 @@
 import asyncio
 import logging
 from multiprocessing import Queue
-from typing import Any, Literal, TypeVar
+from typing import Any, TypeVar
 
 from hatchet_sdk.client import Client
 from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.logger import logger
 from hatchet_sdk.runnables.action import Action
 from hatchet_sdk.runnables.task import Task
+from hatchet_sdk.utils.typing import STOP_LOOP, STOP_LOOP_TYPE
 from hatchet_sdk.worker.action_listener_process import ActionEvent
 from hatchet_sdk.worker.runner.runner import Runner
-from hatchet_sdk.worker.runner.utils.capture_logs import capture_logs
-
-STOP_LOOP_TYPE = Literal["STOP_LOOP"]
-STOP_LOOP: STOP_LOOP_TYPE = "STOP_LOOP"
+from hatchet_sdk.worker.runner.utils.capture_logs import AsyncLogSender, capture_logs
 
 T = TypeVar("T")
 
@@ -28,10 +26,10 @@ class WorkerActionRunLoopManager:
         action_queue: "Queue[Action | STOP_LOOP_TYPE]",
         event_queue: "Queue[ActionEvent]",
         loop: asyncio.AbstractEventLoop,
-        handle_kill: bool = True,
-        debug: bool = False,
-        labels: dict[str, str | int] = {},
-        lifespan_context: Any | None = None,
+        handle_kill: bool,
+        debug: bool,
+        labels: dict[str, str | int] | None,
+        lifespan_context: Any | None,
     ) -> None:
         self.name = name
         self.action_registry = action_registry
@@ -52,15 +50,19 @@ class WorkerActionRunLoopManager:
         self.runner: Runner | None = None
 
         self.client = Client(config=self.config, debug=self.debug)
+        self.start_loop_manager_task: asyncio.Task[None] | None = None
+        self.log_sender = AsyncLogSender(self.client.event)
+        self.log_task = self.loop.create_task(self.log_sender.consume())
+
         self.start()
 
     def start(self) -> None:
-        k = self.loop.create_task(self.aio_start())  # noqa: F841
+        self.start_loop_manager_task = self.loop.create_task(self.aio_start())
 
     async def aio_start(self, retry_count: int = 1) -> None:
         await capture_logs(
             self.client.log_interceptor,
-            self.client.event,
+            self.log_sender,
             self._async_start,
         )()
 
@@ -75,6 +77,7 @@ class WorkerActionRunLoopManager:
         self.killing = True
 
         self.action_queue.put(STOP_LOOP)
+        self.log_sender.publish(STOP_LOOP)
 
     async def wait_for_tasks(self) -> None:
         if self.runner:
@@ -89,6 +92,7 @@ class WorkerActionRunLoopManager:
             self.action_registry,
             self.labels,
             self.lifespan_context,
+            self.log_sender,
         )
 
         logger.debug(f"'{self.name}' waiting for {list(self.action_registry.keys())}")