hatchet-sdk 1.0.2__py3-none-any.whl → 1.1.0__py3-none-any.whl

hatchet_sdk/client.py

@@ -1,14 +1,9 @@
-import asyncio
-
-import grpc
-
 from hatchet_sdk.clients.admin import AdminClient
 from hatchet_sdk.clients.dispatcher.dispatcher import DispatcherClient
-from hatchet_sdk.clients.events import EventClient, new_event
+from hatchet_sdk.clients.events import EventClient
 from hatchet_sdk.clients.run_event_listener import RunEventListenerClient
 from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
 from hatchet_sdk.config import ClientConfig
-from hatchet_sdk.connection import new_conn
 from hatchet_sdk.features.cron import CronClient
 from hatchet_sdk.features.logs import LogsClient
 from hatchet_sdk.features.metrics import MetricsClient
@@ -29,21 +24,13 @@ class Client:
         workflow_listener: PooledWorkflowRunListener | None | None = None,
         debug: bool = False,
     ):
-        try:
-            loop = asyncio.get_running_loop()
-        except RuntimeError:
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-
-        conn: grpc.Channel = new_conn(config, False)
-
         self.config = config
         self.admin = admin_client or AdminClient(config)
         self.dispatcher = dispatcher_client or DispatcherClient(config)
-        self.event = event_client or new_event(conn, config)
+        self.event = event_client or EventClient(config)
         self.listener = RunEventListenerClient(config)
         self.workflow_listener = workflow_listener
-        self.logInterceptor = config.logger
+        self.log_interceptor = config.logger
         self.debug = debug
 
         self.cron = CronClient(self.config)

hatchet_sdk/clients/admin.py

@@ -8,8 +8,6 @@ from google.protobuf import timestamp_pb2
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
-from hatchet_sdk.clients.run_event_listener import RunEventListenerClient
-from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
 from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.connection import new_conn
 from hatchet_sdk.contracts import workflows_pb2 as v0_workflow_protos
@@ -36,6 +34,7 @@ class ScheduleTriggerWorkflowOptions(BaseModel):
     child_index: int | None = None
     child_key: str | None = None
     namespace: str | None = None
+    additional_metadata: JSONSerializableMapping = Field(default_factory=dict)
 
 
 class TriggerWorkflowOptions(ScheduleTriggerWorkflowOptions):
@@ -63,14 +62,11 @@ class AdminClient:
     def __init__(self, config: ClientConfig):
         conn = new_conn(config, False)
         self.config = config
-        self.client = AdminServiceStub(conn)  # type: ignore[no-untyped-call]
-        self.v0_client = WorkflowServiceStub(conn)  # type: ignore[no-untyped-call]
+        self.client = AdminServiceStub(conn)
+        self.v0_client = WorkflowServiceStub(conn)
         self.token = config.token
-        self.listener_client = RunEventListenerClient(config=config)
         self.namespace = config.namespace
 
-        self.pooled_workflow_listener: PooledWorkflowRunListener | None = None
-
     class TriggerWorkflowRequest(BaseModel):
         model_config = ConfigDict(extra="ignore")
 
@@ -153,7 +149,11 @@ class AdminClient:
             name=name,
             schedules=[self._parse_schedule(schedule) for schedule in schedules],
             input=json.dumps(input),
-            **options.model_dump(),
+            parent_id=options.parent_id,
+            parent_step_run_id=options.parent_step_run_id,
+            child_index=options.child_index,
+            child_key=options.child_key,
+            additional_metadata=json.dumps(options.additional_metadata),
         )
 
     @tenacity_retry
@@ -302,9 +302,6 @@ class AdminClient:
     ) -> WorkflowRunRef:
         request = self._create_workflow_run_request(workflow_name, input, options)
 
-        if not self.pooled_workflow_listener:
-            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
-
         try:
             resp = cast(
                 v0_workflow_protos.TriggerWorkflowResponse,
@@ -320,8 +317,7 @@ class AdminClient:
 
         return WorkflowRunRef(
             workflow_run_id=resp.workflow_run_id,
-            workflow_listener=self.pooled_workflow_listener,
-            workflow_run_event_listener=self.listener_client,
+            config=self.config,
         )
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
@@ -338,9 +334,6 @@ class AdminClient:
         async with spawn_index_lock:
             request = self._create_workflow_run_request(workflow_name, input, options)
 
-        if not self.pooled_workflow_listener:
-            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
-
         try:
             resp = cast(
                 v0_workflow_protos.TriggerWorkflowResponse,
@@ -357,8 +350,7 @@ class AdminClient:
 
         return WorkflowRunRef(
             workflow_run_id=resp.workflow_run_id,
-            workflow_listener=self.pooled_workflow_listener,
-            workflow_run_event_listener=self.listener_client,
+            config=self.config,
         )
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
@@ -367,9 +359,6 @@ class AdminClient:
         self,
         workflows: list[WorkflowRunTriggerConfig],
     ) -> list[WorkflowRunRef]:
-        if not self.pooled_workflow_listener:
-            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
-
         bulk_request = v0_workflow_protos.BulkTriggerWorkflowRequest(
             workflows=[
                 self._create_workflow_run_request(
@@ -390,8 +379,7 @@ class AdminClient:
         return [
             WorkflowRunRef(
                 workflow_run_id=workflow_run_id,
-                workflow_listener=self.pooled_workflow_listener,
-                workflow_run_event_listener=self.listener_client,
+                config=self.config,
             )
             for workflow_run_id in resp.workflow_run_ids
         ]
@@ -404,9 +392,6 @@ class AdminClient:
         ## IMPORTANT: The `pooled_workflow_listener` must be created 1) lazily, and not at `init` time, and 2) on the
         ## main thread. If 1) is not followed, you'll get an error about something being attached to the wrong event
         ## loop. If 2) is not followed, you'll get an error about the event loop not being set up.
-        if not self.pooled_workflow_listener:
-            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
-
         async with spawn_index_lock:
             bulk_request = v0_workflow_protos.BulkTriggerWorkflowRequest(
                 workflows=[
@@ -428,18 +413,13 @@ class AdminClient:
         return [
             WorkflowRunRef(
                 workflow_run_id=workflow_run_id,
-                workflow_listener=self.pooled_workflow_listener,
-                workflow_run_event_listener=self.listener_client,
+                config=self.config,
             )
             for workflow_run_id in resp.workflow_run_ids
         ]
 
     def get_workflow_run(self, workflow_run_id: str) -> WorkflowRunRef:
-        if not self.pooled_workflow_listener:
-            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
-
         return WorkflowRunRef(
             workflow_run_id=workflow_run_id,
-            workflow_listener=self.pooled_workflow_listener,
-            workflow_run_event_listener=self.listener_client,
+            config=self.config,
         )

hatchet_sdk/clients/dispatcher/action_listener.py

@@ -152,7 +152,7 @@ class ActionListener:
         self.config = config
         self.worker_id = worker_id
 
-        self.aio_client = DispatcherStub(new_conn(self.config, True))  # type: ignore[no-untyped-call]
+        self.aio_client = DispatcherStub(new_conn(self.config, True))
         self.token = self.config.token
 
         self.retries = 0
@@ -232,14 +232,8 @@ class ActionListener:
         if self.heartbeat_task is not None:
             return
 
-        try:
-            loop = asyncio.get_event_loop()
-        except RuntimeError as e:
-            if str(e).startswith("There is no current event loop in thread"):
-                loop = asyncio.new_event_loop()
-                asyncio.set_event_loop(loop)
-            else:
-                raise e
+        loop = asyncio.get_event_loop()
+
         self.heartbeat_task = loop.create_task(self.heartbeat())
 
     def __aiter__(self) -> AsyncGenerator[Action | None, None]:
@@ -298,7 +292,9 @@ class ActionListener:
                             )
                         )
                     except (ValueError, json.JSONDecodeError) as e:
-                        raise ValueError(f"Error decoding payload: {e}")
+                        logger.error(f"Error decoding payload: {e}")
+
+                        action_payload = ActionPayload()
 
                     action = Action(
                         tenant_id=assigned_action.tenantId,
@@ -384,7 +380,7 @@ class ActionListener:
                 f"action listener connection interrupted, retrying... ({self.retries}/{DEFAULT_ACTION_LISTENER_RETRY_COUNT})"
             )
 
-        self.aio_client = DispatcherStub(new_conn(self.config, True))  # type: ignore[no-untyped-call]
+        self.aio_client = DispatcherStub(new_conn(self.config, True))
 
         if self.listen_strategy == "v2":
             # we should await for the listener to be established before

hatchet_sdk/clients/dispatcher/dispatcher.py

@@ -34,20 +34,23 @@ DEFAULT_REGISTER_TIMEOUT = 30
 
 
 class DispatcherClient:
-    config: ClientConfig
-
     def __init__(self, config: ClientConfig):
         conn = new_conn(config, False)
-        self.client = DispatcherStub(conn)  # type: ignore[no-untyped-call]
+        self.client = DispatcherStub(conn)
 
-        aio_conn = new_conn(config, True)
-        self.aio_client = DispatcherStub(aio_conn)  # type: ignore[no-untyped-call]
         self.token = config.token
         self.config = config
 
+        ## IMPORTANT: This needs to be created lazily so we don't require
+        ## an event loop to instantiate the client.
+        self.aio_client: DispatcherStub | None = None
+
     async def get_action_listener(
         self, req: GetActionListenerRequest
     ) -> ActionListener:
+        if not self.aio_client:
+            aio_conn = new_conn(self.config, True)
+            self.aio_client = DispatcherStub(aio_conn)
 
         # Override labels with the preset labels
         preset_labels = self.config.worker_preset_labels
@@ -95,6 +98,10 @@ class DispatcherClient:
     async def _try_send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
     ) -> grpc.aio.UnaryUnaryCall[StepActionEvent, ActionEventResponse]:
+        if not self.aio_client:
+            aio_conn = new_conn(self.config, True)
+            self.aio_client = DispatcherStub(aio_conn)
+
         event_timestamp = Timestamp()
         event_timestamp.GetCurrentTime()
 
@@ -122,6 +129,10 @@ class DispatcherClient:
     async def send_group_key_action_event(
         self, action: Action, event_type: GroupKeyActionEventType, payload: str
     ) -> grpc.aio.UnaryUnaryCall[GroupKeyActionEvent, ActionEventResponse]:
+        if not self.aio_client:
+            aio_conn = new_conn(self.config, True)
+            self.aio_client = DispatcherStub(aio_conn)
+
         event_timestamp = Timestamp()
         event_timestamp.GetCurrentTime()
 
@@ -191,6 +202,10 @@ class DispatcherClient:
         worker_id: str | None,
         labels: dict[str, str | int],
     ) -> None:
+        if not self.aio_client:
+            aio_conn = new_conn(self.config, True)
+            self.aio_client = DispatcherStub(aio_conn)
+
         worker_labels = {}
 
         for key, value in labels.items():

hatchet_sdk/clients/durable_event_listener.py

@@ -84,14 +84,6 @@ class RegisterDurableEventRequest(BaseModel):
 
 class DurableEventListener:
     def __init__(self, config: ClientConfig):
-        try:
-            asyncio.get_running_loop()
-        except RuntimeError:
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-
-        conn = new_conn(config, True)
-        self.client = V1DispatcherStub(conn)  # type: ignore[no-untyped-call]
         self.token = config.token
         self.config = config
 
@@ -129,11 +121,14 @@ class DurableEventListener:
             self.interrupt.set()
 
     async def _init_producer(self) -> None:
+        conn = new_conn(self.config, True)
+        client = V1DispatcherStub(conn)
+
         try:
             if not self.listener:
                 while True:
                     try:
-                        self.listener = await self._retry_subscribe()
+                        self.listener = await self._retry_subscribe(client)
 
                         logger.debug("Workflow run listener connected.")
 
@@ -282,6 +277,7 @@ class DurableEventListener:
 
     async def _retry_subscribe(
         self,
+        client: V1DispatcherStub,
     ) -> grpc.aio.UnaryStreamCall[ListenForDurableEventRequest, DurableEvent]:
         retries = 0
 
@@ -298,8 +294,8 @@ class DurableEventListener:
                     grpc.aio.UnaryStreamCall[
                         ListenForDurableEventRequest, DurableEvent
                     ],
-                    self.client.ListenForDurableEvent(
-                        self._request(),
+                    client.ListenForDurableEvent(
+                        self._request(),  # type: ignore[arg-type]
                         metadata=get_metadata(self.token),
                     ),
                 )
@@ -315,7 +311,10 @@ class DurableEventListener:
     def register_durable_event(
         self, request: RegisterDurableEventRequest
     ) -> Literal[True]:
-        self.client.RegisterDurableEvent(
+        conn = new_conn(self.config, True)
+        client = V1DispatcherStub(conn)
+
+        client.RegisterDurableEvent(
             request.to_proto(),
             timeout=5,
             metadata=get_metadata(self.token),

hatchet_sdk/clients/events.py

@@ -3,15 +3,16 @@ import datetime
 import json
 from typing import List, cast
 
-import grpc
 from google.protobuf import timestamp_pb2
 from pydantic import BaseModel, Field
 
 from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
 from hatchet_sdk.config import ClientConfig
+from hatchet_sdk.connection import new_conn
 from hatchet_sdk.contracts.events_pb2 import (
     BulkPushEventRequest,
     Event,
+    Events,
     PushEventRequest,
     PutLogRequest,
     PutStreamEventRequest,
@@ -21,13 +22,6 @@ from hatchet_sdk.metadata import get_metadata
 from hatchet_sdk.utils.typing import JSONSerializableMapping
 
 
-def new_event(conn: grpc.Channel, config: ClientConfig) -> "EventClient":
-    return EventClient(
-        client=EventsServiceStub(conn),  # type: ignore[no-untyped-call]
-        config=config,
-    )
-
-
 def proto_timestamp_now() -> timestamp_pb2.Timestamp:
     t = datetime.datetime.now().timestamp()
     seconds = int(t)
@@ -52,8 +46,10 @@ class BulkPushEventWithMetadata(BaseModel):
 
 
 class EventClient:
-    def __init__(self, client: EventsServiceStub, config: ClientConfig):
-        self.client = client
+    def __init__(self, config: ClientConfig):
+        conn = new_conn(config, False)
+        self.client = EventsServiceStub(conn)
+
         self.token = config.token
         self.namespace = config.namespace
 
@@ -146,11 +142,11 @@ class EventClient:
             ]
         )
 
-        response = self.client.BulkPush(bulk_request, metadata=get_metadata(self.token))
-
-        return cast(
-            list[Event],
-            response.events,
+        return list(
+            cast(
+                Events,
+                self.client.BulkPush(bulk_request, metadata=get_metadata(self.token)),
+            ).events
         )
 
     def log(self, message: str, step_run_id: str) -> None:

hatchet_sdk/clients/rest/models/tenant_resource.py

@@ -29,8 +29,10 @@ class TenantResource(str, Enum):
     allowed enum values
     """
     WORKER = "WORKER"
+    WORKER_SLOT = "WORKER_SLOT"
     EVENT = "EVENT"
     WORKFLOW_RUN = "WORKFLOW_RUN"
+    TASK_RUN = "TASK_RUN"
     CRON = "CRON"
     SCHEDULE = "SCHEDULE"
 

hatchet_sdk/clients/rest/models/workflow_runs_metrics.py

@@ -22,17 +22,13 @@ from typing import Any, ClassVar, Dict, List, Optional, Set
 from pydantic import BaseModel, ConfigDict
 from typing_extensions import Self
 
-from hatchet_sdk.clients.rest.models.workflow_runs_metrics_counts import (
-    WorkflowRunsMetricsCounts,
-)
-
 
 class WorkflowRunsMetrics(BaseModel):
     """
     WorkflowRunsMetrics
     """  # noqa: E501
 
-    counts: Optional[WorkflowRunsMetricsCounts] = None
+    counts: Optional[Dict[str, Any]] = None
     __properties: ClassVar[List[str]] = ["counts"]
 
     model_config = ConfigDict(

hatchet_sdk/clients/run_event_listener.py

@@ -1,7 +1,8 @@
 import asyncio
-import json
 from enum import Enum
-from typing import Any, AsyncGenerator, Callable, Generator, cast
+from queue import Empty, Queue
+from threading import Thread
+from typing import Any, AsyncGenerator, Callable, Generator, Literal, TypeVar, cast
 
 import grpc
 from pydantic import BaseModel
@@ -56,6 +57,8 @@ workflow_run_event_type_mapping = {
     ResourceEventType.RESOURCE_EVENT_TYPE_TIMED_OUT: WorkflowRunEventType.WORKFLOW_RUN_EVENT_TYPE_TIMED_OUT,
 }
 
+T = TypeVar("T")
+
 
 class StepRunEvent(BaseModel):
     type: StepRunEventType
@@ -65,18 +68,20 @@ class StepRunEvent(BaseModel):
 class RunEventListener:
     def __init__(
         self,
-        client: DispatcherStub,
-        token: str,
+        config: ClientConfig,
         workflow_run_id: str | None = None,
         additional_meta_kv: tuple[str, str] | None = None,
     ):
-        self.client = client
+        self.config = config
         self.stop_signal = False
-        self.token = token
 
         self.workflow_run_id = workflow_run_id
         self.additional_meta_kv = additional_meta_kv
 
+        ## IMPORTANT: This needs to be created lazily so we don't require
+        ## an event loop to instantiate the client.
+        self.client: DispatcherStub | None = None
+
     def abort(self) -> None:
         self.stop_signal = True
 
@@ -86,27 +91,46 @@ class RunEventListener:
     async def __anext__(self) -> StepRunEvent:
         return await self._generator().__anext__()
 
-    def __iter__(self) -> Generator[StepRunEvent, None, None]:
-        try:
-            loop = asyncio.get_event_loop()
-        except RuntimeError as e:
-            if str(e).startswith("There is no current event loop in thread"):
-                loop = asyncio.new_event_loop()
-                asyncio.set_event_loop(loop)
-            else:
-                raise e
+    def async_to_sync_thread(
+        self, async_iter: AsyncGenerator[T, None]
+    ) -> Generator[T, None, None]:
+        q = Queue[T | Literal["DONE"]]()
+        done_sentinel: Literal["DONE"] = "DONE"
+
+        def runner() -> None:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+            async def consume() -> None:
+                try:
+                    async for item in async_iter:
+                        q.put(item)
+                finally:
+                    q.put(done_sentinel)
 
-        async_iter = self.__aiter__()
+            try:
+                loop.run_until_complete(consume())
+            finally:
+                loop.stop()
+                loop.close()
+
+        thread = Thread(target=runner)
+        thread.start()
 
         while True:
             try:
-                future = asyncio.ensure_future(async_iter.__anext__())
-                yield loop.run_until_complete(future)
-            except StopAsyncIteration:
-                break
-            except Exception as e:
-                print(f"Error in synchronous iterator: {e}")
-                break
+                item = q.get(timeout=1)
+                if item == "DONE":
+                    break
+                yield item
+            except Empty:
+                continue
+
+        thread.join()
+
+    def __iter__(self) -> Generator[StepRunEvent, None, None]:
+        for item in self.async_to_sync_thread(self.__aiter__()):
+            yield item
 
     async def _generator(self) -> AsyncGenerator[StepRunEvent, None]:
         while True:
@@ -128,18 +152,10 @@ class RunEventListener:
                             raise Exception(
                                 f"Unknown event type: {workflow_event.eventType}"
                             )
-                        payload = None
-
-                        try:
-                            if workflow_event.eventPayload:
-                                payload = json.loads(workflow_event.eventPayload)
-                        except Exception:
-                            payload = workflow_event.eventPayload
-                            pass
 
-                        assert isinstance(payload, str)
-
-                        yield StepRunEvent(type=eventType, payload=payload)
+                        yield StepRunEvent(
+                            type=eventType, payload=workflow_event.eventPayload
+                        )
                     elif workflow_event.resourceType == RESOURCE_TYPE_WORKFLOW_RUN:
                         if workflow_event.eventType in step_run_event_type_mapping:
                             workflowRunEventType = step_run_event_type_mapping[
@@ -150,17 +166,10 @@ class RunEventListener:
                                 f"Unknown event type: {workflow_event.eventType}"
                             )
 
-                        payload = None
-
-                        try:
-                            if workflow_event.eventPayload:
-                                payload = json.loads(workflow_event.eventPayload)
-                        except Exception:
-                            pass
-
-                        assert isinstance(payload, str)
-
-                        yield StepRunEvent(type=workflowRunEventType, payload=payload)
+                        yield StepRunEvent(
+                            type=workflowRunEventType,
+                            payload=workflow_event.eventPayload,
+                        )
 
                     if workflow_event.hangup:
                         listener = None
@@ -188,6 +197,10 @@ class RunEventListener:
     async def retry_subscribe(self) -> AsyncGenerator[WorkflowEvent, None]:
         retries = 0
 
+        if self.client is None:
+            aio_conn = new_conn(self.config, True)
+            self.client = DispatcherStub(aio_conn)
+
         while retries < DEFAULT_ACTION_LISTENER_RETRY_COUNT:
             try:
                 if retries > 0:
@@ -200,7 +213,7 @@ class RunEventListener:
                             SubscribeToWorkflowEventsRequest(
                                 workflowRunId=self.workflow_run_id,
                             ),
-                            metadata=get_metadata(self.token),
+                            metadata=get_metadata(self.config.token),
                         ),
                     )
                 elif self.additional_meta_kv is not None:
@@ -211,7 +224,7 @@ class RunEventListener:
                                 additionalMetaKey=self.additional_meta_kv[0],
                                 additionalMetaValue=self.additional_meta_kv[1],
                             ),
-                            metadata=get_metadata(self.token),
+                            metadata=get_metadata(self.config.token),
                         ),
                     )
                 else:
@@ -228,33 +241,16 @@ class RunEventListener:
 
 class RunEventListenerClient:
     def __init__(self, config: ClientConfig):
-        self.token = config.token
         self.config = config
-        self.client: DispatcherStub | None = None
 
     def stream_by_run_id(self, workflow_run_id: str) -> RunEventListener:
         return self.stream(workflow_run_id)
 
     def stream(self, workflow_run_id: str) -> RunEventListener:
-        if not isinstance(workflow_run_id, str) and hasattr(workflow_run_id, "__str__"):
-            workflow_run_id = str(workflow_run_id)
-
-        if not self.client:
-            aio_conn = new_conn(self.config, True)
-            self.client = DispatcherStub(aio_conn)  # type: ignore[no-untyped-call]
-
-        return RunEventListener(
-            client=self.client, token=self.token, workflow_run_id=workflow_run_id
-        )
+        return RunEventListener(config=self.config, workflow_run_id=workflow_run_id)
 
     def stream_by_additional_metadata(self, key: str, value: str) -> RunEventListener:
-        if not self.client:
-            aio_conn = new_conn(self.config, True)
-            self.client = DispatcherStub(aio_conn)  # type: ignore[no-untyped-call]
-
-        return RunEventListener(
-            client=self.client, token=self.token, additional_meta_kv=(key, value)
-        )
+        return RunEventListener(config=self.config, additional_meta_kv=(key, value))
 
     async def on(
         self, workflow_run_id: str, handler: Callable[[StepRunEvent], Any] | None = None