hatchet-sdk 1.0.0__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

hatchet_sdk/clients/dispatcher/action_listener.py

@@ -1,14 +1,14 @@
 import asyncio
 import json
 import time
-from dataclasses import dataclass, field
+from dataclasses import field
 from enum import Enum
-from typing import Any, AsyncGenerator, Optional, cast
+from typing import Any, AsyncGenerator, cast
 
 import grpc
 import grpc.aio
 from grpc._cython import cygrpc  # type: ignore[attr-defined]
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 from hatchet_sdk.clients.event_ts import ThreadSafeEvent, read_with_interrupt
 from hatchet_sdk.clients.events import proto_timestamp_now
@@ -36,25 +36,29 @@ DEFAULT_ACTION_TIMEOUT = 600  # seconds
 DEFAULT_ACTION_LISTENER_RETRY_COUNT = 15
 
 
-@dataclass
-class GetActionListenerRequest:
+class GetActionListenerRequest(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     worker_name: str
     services: list[str]
     actions: list[str]
-    max_runs: int | None = None
-    _labels: dict[str, str | int] = field(default_factory=dict)
+    slots: int = 100
+    raw_labels: dict[str, str | int] = Field(default_factory=dict)
 
-    labels: dict[str, WorkerLabels] = field(init=False)
+    labels: dict[str, WorkerLabels] = Field(default_factory=dict)
 
-    def __post_init__(self) -> None:
+    @model_validator(mode="after")
+    def validate_labels(self) -> "GetActionListenerRequest":
         self.labels = {}
 
-        for key, value in self._labels.items():
+        for key, value in self.raw_labels.items():
             if isinstance(value, int):
                 self.labels[key] = WorkerLabels(intValue=value)
             else:
                 self.labels[key] = WorkerLabels(strValue=str(value))
 
+        return self
+
 
 class ActionPayload(BaseModel):
     model_config = ConfigDict(extra="allow")
@@ -65,6 +69,7 @@ class ActionPayload(BaseModel):
     user_data: JSONSerializableMapping = Field(default_factory=dict)
     step_run_errors: dict[str, str] = Field(default_factory=dict)
     triggered_by: str | None = None
+    triggers: JSONSerializableMapping = Field(default_factory=dict)
 
     @field_validator(
         "input", "parents", "overrides", "user_data", "step_run_errors", mode="before"
@@ -142,30 +147,24 @@ def parse_additional_metadata(additional_metadata: str) -> JSONSerializableMappi
         return {}
 
 
-@dataclass
 class ActionListener:
-    config: ClientConfig
-    worker_id: str
+    def __init__(self, config: ClientConfig, worker_id: str) -> None:
+        self.config = config
+        self.worker_id = worker_id
 
-    client: DispatcherStub = field(init=False)
-    aio_client: DispatcherStub = field(init=False)
-    token: str = field(init=False)
-    retries: int = field(default=0, init=False)
-    last_connection_attempt: float = field(default=0, init=False)
-    last_heartbeat_succeeded: bool = field(default=True, init=False)
-    time_last_hb_succeeded: float = field(default=9999999999999, init=False)
-    heartbeat_task: Optional[asyncio.Task[None]] = field(default=None, init=False)
-    run_heartbeat: bool = field(default=True, init=False)
-    listen_strategy: str = field(default="v2", init=False)
-    stop_signal: bool = field(default=False, init=False)
-
-    missed_heartbeats: int = field(default=0, init=False)
-
-    def __post_init__(self) -> None:
-        self.client = DispatcherStub(new_conn(self.config, False))  # type: ignore[no-untyped-call]
         self.aio_client = DispatcherStub(new_conn(self.config, True))  # type: ignore[no-untyped-call]
         self.token = self.config.token
 
+        self.retries = 0
+        self.last_heartbeat_succeeded = True
+        self.time_last_hb_succeeded = 9999999999999.0
+        self.last_connection_attempt = 0.0
+        self.heartbeat_task: asyncio.Task[None] | None = None
+        self.run_heartbeat = True
+        self.listen_strategy = "v2"
+        self.stop_signal = False
+        self.missed_heartbeats = 0
+
     def is_healthy(self) -> bool:
         return self.last_heartbeat_succeeded
 
@@ -292,11 +291,16 @@ class ActionListener:
 
                     self.retries = 0
 
-                    action_payload = (
-                        {}
-                        if not assigned_action.actionPayload
-                        else self.parse_action_payload(assigned_action.actionPayload)
-                    )
+                    try:
+                        action_payload = (
+                            ActionPayload()
+                            if not assigned_action.actionPayload
+                            else ActionPayload.model_validate_json(
+                                assigned_action.actionPayload
+                            )
+                        )
+                    except (ValueError, json.JSONDecodeError) as e:
+                        raise ValueError(f"Error decoding payload: {e}")
 
                     action = Action(
                         tenant_id=assigned_action.tenantId,
@@ -309,7 +313,7 @@ class ActionListener:
                         step_id=assigned_action.stepId,
                         step_run_id=assigned_action.stepRunId,
                         action_id=assigned_action.actionId,
-                        action_payload=ActionPayload.model_validate(action_payload),
+                        action_payload=action_payload,
                         action_type=convert_proto_enum_to_python(
                             assigned_action.actionType,
                             ActionType,
@@ -352,16 +356,10 @@ class ActionListener:
 
                     self.retries = self.retries + 1
 
-    def parse_action_payload(self, payload: str) -> JSONSerializableMapping:
-        try:
-            return cast(JSONSerializableMapping, json.loads(payload))
-        except json.JSONDecodeError as e:
-            raise ValueError(f"Error decoding payload: {e}")
-
     async def get_listen_client(
         self,
     ) -> grpc.aio.UnaryStreamCall[WorkerListenRequest, AssignedAction]:
-        current_time = int(time.time())
+        current_time = time.time()
 
         if (
             current_time - self.last_connection_attempt
@@ -438,8 +436,10 @@ class ActionListener:
                 timeout=5,
                 metadata=get_metadata(self.token),
             )
+
             if self.interrupt is not None:
                 self.interrupt.set()
+
             return cast(WorkerUnsubscribeRequest, req)
         except grpc.RpcError as e:
             raise Exception(f"Failed to unsubscribe: {e}")

hatchet_sdk/clients/dispatcher/dispatcher.py

@@ -55,17 +55,19 @@ class DispatcherClient:
         for key, value in preset_labels.items():
             req.labels[key] = WorkerLabels(strValue=str(value))
 
-        # Register the worker
-        response: WorkerRegisterResponse = await self.aio_client.Register(
-            WorkerRegisterRequest(
-                workerName=req.worker_name,
-                actions=req.actions,
-                services=req.services,
-                maxRuns=req.max_runs,
-                labels=req.labels,
+        response = cast(
+            WorkerRegisterResponse,
+            await self.aio_client.Register(
+                WorkerRegisterRequest(
+                    workerName=req.worker_name,
+                    actions=req.actions,
+                    services=req.services,
+                    maxRuns=req.slots,
+                    labels=req.labels,
+                ),
+                timeout=DEFAULT_REGISTER_TIMEOUT,
+                metadata=get_metadata(self.token),
             ),
-            timeout=DEFAULT_REGISTER_TIMEOUT,
-            metadata=get_metadata(self.token),
         )
 
         return ActionListener(self.config, response.workerId)
@@ -93,8 +95,8 @@ class DispatcherClient:
     async def _try_send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
     ) -> grpc.aio.UnaryUnaryCall[StepActionEvent, ActionEventResponse]:
-        eventTimestamp = Timestamp()
-        eventTimestamp.GetCurrentTime()
+        event_timestamp = Timestamp()
+        event_timestamp.GetCurrentTime()
 
         event = StepActionEvent(
             workerId=action.worker_id,
@@ -103,7 +105,7 @@ class DispatcherClient:
             stepId=action.step_id,
             stepRunId=action.step_run_id,
             actionId=action.action_id,
-            eventTimestamp=eventTimestamp,
+            eventTimestamp=event_timestamp,
             eventType=event_type,
             eventPayload=payload,
             retryCount=action.retry_count,
@@ -120,15 +122,15 @@ class DispatcherClient:
     async def send_group_key_action_event(
         self, action: Action, event_type: GroupKeyActionEventType, payload: str
     ) -> grpc.aio.UnaryUnaryCall[GroupKeyActionEvent, ActionEventResponse]:
-        eventTimestamp = Timestamp()
-        eventTimestamp.GetCurrentTime()
+        event_timestamp = Timestamp()
+        event_timestamp.GetCurrentTime()
 
         event = GroupKeyActionEvent(
             workerId=action.worker_id,
             workflowRunId=action.workflow_run_id,
             getGroupKeyRunId=action.get_group_key_run_id,
             actionId=action.action_id,
-            eventTimestamp=eventTimestamp,
+            eventTimestamp=event_timestamp,
             eventType=event_type,
             eventPayload=payload,
         )

hatchet_sdk/clients/durable_event_listener.py

@@ -0,0 +1,327 @@
+import asyncio
+import json
+from collections.abc import AsyncIterator
+from typing import Any, Literal, cast
+
+import grpc
+import grpc.aio
+from grpc._cython import cygrpc  # type: ignore[attr-defined]
+from pydantic import BaseModel, ConfigDict
+
+from hatchet_sdk.clients.event_ts import ThreadSafeEvent, read_with_interrupt
+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.v1.dispatcher_pb2 import (
+    DurableEvent,
+    ListenForDurableEventRequest,
+)
+from hatchet_sdk.contracts.v1.dispatcher_pb2 import (
+    RegisterDurableEventRequest as RegisterDurableEventRequestProto,
+)
+from hatchet_sdk.contracts.v1.dispatcher_pb2_grpc import V1DispatcherStub
+from hatchet_sdk.contracts.v1.shared.condition_pb2 import DurableEventListenerConditions
+from hatchet_sdk.logger import logger
+from hatchet_sdk.metadata import get_metadata
+from hatchet_sdk.waits import SleepCondition, UserEventCondition
+
+DEFAULT_DURABLE_EVENT_LISTENER_RETRY_INTERVAL = 3  # seconds
+DEFAULT_DURABLE_EVENT_LISTENER_RETRY_COUNT = 5
+DEFAULT_DURABLE_EVENT_LISTENER_INTERRUPT_INTERVAL = 1800  # 30 minutes
+
+
+class _Subscription:
+    def __init__(self, id: int, task_id: str, signal_key: str):
+        self.id = id
+        self.task_id = task_id
+        self.signal_key = signal_key
+        self.queue: asyncio.Queue[DurableEvent | None] = asyncio.Queue()
+
+    async def __aiter__(self) -> "_Subscription":
+        return self
+
+    async def __anext__(self) -> DurableEvent | None:
+        return await self.queue.get()
+
+    async def get(self) -> DurableEvent:
+        event = await self.queue.get()
+
+        if event is None:
+            raise StopAsyncIteration
+
+        return event
+
+    async def put(self, item: DurableEvent) -> None:
+        await self.queue.put(item)
+
+    async def close(self) -> None:
+        await self.queue.put(None)
+
+
+class RegisterDurableEventRequest(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    task_id: str
+    signal_key: str
+    conditions: list[SleepCondition | UserEventCondition]
+
+    def to_proto(self) -> RegisterDurableEventRequestProto:
+        return RegisterDurableEventRequestProto(
+            task_id=self.task_id,
+            signal_key=self.signal_key,
+            conditions=DurableEventListenerConditions(
+                sleep_conditions=[
+                    c.to_pb() for c in self.conditions if isinstance(c, SleepCondition)
+                ],
+                user_event_conditions=[
+                    c.to_pb()
+                    for c in self.conditions
+                    if isinstance(c, UserEventCondition)
+                ],
+            ),
+        )
+
+
+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
+
+        # list of all active subscriptions, mapping from a subscription id to a task id and signal key
+        self.subscriptions_to_task_id_signal_key: dict[int, tuple[str, str]] = {}
+
+        # task id-signal key tuples mapped to an array of subscription ids
+        self.task_id_signal_key_to_subscriptions: dict[tuple[str, str], list[int]] = {}
+
+        self.subscription_counter: int = 0
+        self.subscription_counter_lock: asyncio.Lock = asyncio.Lock()
+
+        self.requests: asyncio.Queue[ListenForDurableEventRequest | int] = (
+            asyncio.Queue()
+        )
+
+        self.listener: (
+            grpc.aio.UnaryStreamCall[ListenForDurableEventRequest, DurableEvent] | None
+        ) = None
+        self.listener_task: asyncio.Task[None] | None = None
+
+        self.curr_requester: int = 0
+
+        self.events: dict[int, _Subscription] = {}
+
+        self.interrupter: asyncio.Task[None] | None = None
+
+    async def _interrupter(self) -> None:
+        """
+        _interrupter runs in a separate thread and interrupts the listener according to a configurable duration.
+        """
+        await asyncio.sleep(DEFAULT_DURABLE_EVENT_LISTENER_INTERRUPT_INTERVAL)
+
+        if self.interrupt is not None:
+            self.interrupt.set()
+
+    async def _init_producer(self) -> None:
+        try:
+            if not self.listener:
+                while True:
+                    try:
+                        self.listener = await self._retry_subscribe()
+
+                        logger.debug("Workflow run listener connected.")
+
+                        # spawn an interrupter task
+                        if self.interrupter is not None and not self.interrupter.done():
+                            self.interrupter.cancel()
+
+                        self.interrupter = asyncio.create_task(self._interrupter())
+
+                        while True:
+                            self.interrupt = ThreadSafeEvent()
+                            if self.listener is None:
+                                continue
+
+                            t = asyncio.create_task(
+                                read_with_interrupt(self.listener, self.interrupt)
+                            )
+                            await self.interrupt.wait()
+
+                            if not t.done():
+                                logger.warning(
+                                    "Interrupted read_with_interrupt task of durable event listener"
+                                )
+
+                                t.cancel()
+                                if self.listener:
+                                    self.listener.cancel()
+                                await asyncio.sleep(
+                                    DEFAULT_DURABLE_EVENT_LISTENER_RETRY_INTERVAL
+                                )
+                                break
+
+                            event = t.result()
+
+                            if event is cygrpc.EOF:
+                                break
+
+                            # get a list of subscriptions for this task-signal pair
+                            subscriptions = (
+                                self.task_id_signal_key_to_subscriptions.get(
+                                    (event.task_id, event.signal_key), []
+                                )
+                            )
+
+                            for subscription_id in subscriptions:
+                                await self.events[subscription_id].put(event)
+
+                    except grpc.RpcError as e:
+                        logger.debug(f"grpc error in durable event listener: {e}")
+                        await asyncio.sleep(
+                            DEFAULT_DURABLE_EVENT_LISTENER_RETRY_INTERVAL
+                        )
+                        continue
+
+        except Exception as e:
+            logger.error(f"Error in durable event listener: {e}")
+
+            self.listener = None
+
+            # close all subscriptions
+            for subscription_id in self.events:
+                await self.events[subscription_id].close()
+
+            raise e
+
+    async def _request(self) -> AsyncIterator[ListenForDurableEventRequest]:
+        self.curr_requester = self.curr_requester + 1
+
+        # replay all existing subscriptions
+        for task_id, signal_key in set(
+            self.subscriptions_to_task_id_signal_key.values()
+        ):
+            yield ListenForDurableEventRequest(
+                task_id=task_id,
+                signal_key=signal_key,
+            )
+
+        while True:
+            request = await self.requests.get()
+
+            # if the request is an int which matches the current requester, then we should stop
+            if request == self.curr_requester:
+                break
+
+            # if we've gotten an int that doesn't match the current requester, then we should ignore it
+            if isinstance(request, int):
+                continue
+
+            yield request
+            self.requests.task_done()
+
+    def cleanup_subscription(self, subscription_id: int) -> None:
+        task_id_signal_key = self.subscriptions_to_task_id_signal_key[subscription_id]
+
+        if task_id_signal_key in self.task_id_signal_key_to_subscriptions:
+            self.task_id_signal_key_to_subscriptions[task_id_signal_key].remove(
+                subscription_id
+            )
+
+        del self.subscriptions_to_task_id_signal_key[subscription_id]
+        del self.events[subscription_id]
+
+    async def subscribe(self, task_id: str, signal_key: str) -> DurableEvent:
+        try:
+            # create a new subscription id, place a mutex on the counter
+            async with self.subscription_counter_lock:
+                self.subscription_counter += 1
+                subscription_id = self.subscription_counter
+
+            self.subscriptions_to_task_id_signal_key[subscription_id] = (
+                task_id,
+                signal_key,
+            )
+
+            if (task_id, signal_key) not in self.task_id_signal_key_to_subscriptions:
+                self.task_id_signal_key_to_subscriptions[(task_id, signal_key)] = [
+                    subscription_id
+                ]
+            else:
+                self.task_id_signal_key_to_subscriptions[(task_id, signal_key)].append(
+                    subscription_id
+                )
+
+            self.events[subscription_id] = _Subscription(
+                subscription_id, task_id, signal_key
+            )
+
+            await self.requests.put(
+                ListenForDurableEventRequest(
+                    task_id=task_id,
+                    signal_key=signal_key,
+                )
+            )
+
+            if not self.listener_task or self.listener_task.done():
+                self.listener_task = asyncio.create_task(self._init_producer())
+
+            return await self.events[subscription_id].get()
+        except asyncio.CancelledError:
+            raise
+        finally:
+            self.cleanup_subscription(subscription_id)
+
+    async def _retry_subscribe(
+        self,
+    ) -> grpc.aio.UnaryStreamCall[ListenForDurableEventRequest, DurableEvent]:
+        retries = 0
+
+        while retries < DEFAULT_DURABLE_EVENT_LISTENER_RETRY_COUNT:
+            try:
+                if retries > 0:
+                    await asyncio.sleep(DEFAULT_DURABLE_EVENT_LISTENER_RETRY_INTERVAL)
+
+                # signal previous async iterator to stop
+                if self.curr_requester != 0:
+                    self.requests.put_nowait(self.curr_requester)
+
+                return cast(
+                    grpc.aio.UnaryStreamCall[
+                        ListenForDurableEventRequest, DurableEvent
+                    ],
+                    self.client.ListenForDurableEvent(
+                        self._request(),
+                        metadata=get_metadata(self.token),
+                    ),
+                )
+            except grpc.RpcError as e:
+                if e.code() == grpc.StatusCode.UNAVAILABLE:
+                    retries = retries + 1
+                else:
+                    raise ValueError(f"gRPC error: {e}")
+
+        raise ValueError("Failed to connect to durable event listener")
+
+    @tenacity_retry
+    def register_durable_event(
+        self, request: RegisterDurableEventRequest
+    ) -> Literal[True]:
+        self.client.RegisterDurableEvent(
+            request.to_proto(),
+            timeout=5,
+            metadata=get_metadata(self.token),
+        )
+
+        return True
+
+    @tenacity_retry
+    async def result(self, task_id: str, signal_key: str) -> dict[str, Any]:
+        event = await self.subscribe(task_id, signal_key)
+
+        return cast(dict[str, Any], json.loads(event.data.decode("utf-8")))

hatchet_sdk/clients/rest/__init__.py

@@ -230,8 +230,10 @@ from hatchet_sdk.clients.rest.models.user_tenant_memberships_list import (
 from hatchet_sdk.clients.rest.models.user_tenant_public import UserTenantPublic
 from hatchet_sdk.clients.rest.models.v1_cancel_task_request import V1CancelTaskRequest
 from hatchet_sdk.clients.rest.models.v1_dag_children import V1DagChildren
+from hatchet_sdk.clients.rest.models.v1_log_line import V1LogLine
+from hatchet_sdk.clients.rest.models.v1_log_line_level import V1LogLineLevel
+from hatchet_sdk.clients.rest.models.v1_log_line_list import V1LogLineList
 from hatchet_sdk.clients.rest.models.v1_replay_task_request import V1ReplayTaskRequest
-from hatchet_sdk.clients.rest.models.v1_task import V1Task
 from hatchet_sdk.clients.rest.models.v1_task_event import V1TaskEvent
 from hatchet_sdk.clients.rest.models.v1_task_event_list import V1TaskEventList
 from hatchet_sdk.clients.rest.models.v1_task_event_type import V1TaskEventType
@@ -243,8 +245,17 @@ from hatchet_sdk.clients.rest.models.v1_task_run_status import V1TaskRunStatus
 from hatchet_sdk.clients.rest.models.v1_task_status import V1TaskStatus
 from hatchet_sdk.clients.rest.models.v1_task_summary import V1TaskSummary
 from hatchet_sdk.clients.rest.models.v1_task_summary_list import V1TaskSummaryList
+from hatchet_sdk.clients.rest.models.v1_trigger_workflow_run_request import (
+    V1TriggerWorkflowRunRequest,
+)
 from hatchet_sdk.clients.rest.models.v1_workflow_run import V1WorkflowRun
 from hatchet_sdk.clients.rest.models.v1_workflow_run_details import V1WorkflowRunDetails
+from hatchet_sdk.clients.rest.models.v1_workflow_run_display_name import (
+    V1WorkflowRunDisplayName,
+)
+from hatchet_sdk.clients.rest.models.v1_workflow_run_display_name_list import (
+    V1WorkflowRunDisplayNameList,
+)
 from hatchet_sdk.clients.rest.models.v1_workflow_type import V1WorkflowType
 from hatchet_sdk.clients.rest.models.webhook_worker import WebhookWorker
 from hatchet_sdk.clients.rest.models.webhook_worker_create_request import (