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

hatchet_sdk/clients/dispatcher/action_listener.py

@@ -1,19 +1,24 @@
 import asyncio
 import json
 import time
-from dataclasses import dataclass, field
-from typing import Any, AsyncGenerator, List, Optional
+from dataclasses import field
+from enum import Enum
+from typing import Any, AsyncGenerator, cast
 
 import grpc
-from grpc._cython import cygrpc
+import grpc.aio
+from grpc._cython import cygrpc  # type: ignore[attr-defined]
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
-from hatchet_sdk.clients.event_ts import Event_ts, read_with_interrupt
+from hatchet_sdk.clients.event_ts import ThreadSafeEvent, read_with_interrupt
+from hatchet_sdk.clients.events import proto_timestamp_now
 from hatchet_sdk.clients.run_event_listener import (
     DEFAULT_ACTION_LISTENER_RETRY_INTERVAL,
 )
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.connection import new_conn
+from hatchet_sdk.contracts.dispatcher_pb2 import ActionType as ActionTypeProto
 from hatchet_sdk.contracts.dispatcher_pb2 import (
-    ActionType,
     AssignedAction,
     HeartbeatRequest,
     WorkerLabels,
@@ -22,42 +27,65 @@ from hatchet_sdk.contracts.dispatcher_pb2 import (
 )
 from hatchet_sdk.contracts.dispatcher_pb2_grpc import DispatcherStub
 from hatchet_sdk.logger import logger
+from hatchet_sdk.metadata import get_metadata
 from hatchet_sdk.utils.backoff import exp_backoff_sleep
-from hatchet_sdk.utils.serialization import flatten
-
-from ...loader import ClientConfig
-from ...metadata import get_metadata
-from ..events import proto_timestamp_now
+from hatchet_sdk.utils.proto_enums import convert_proto_enum_to_python
+from hatchet_sdk.utils.typing import JSONSerializableMapping
 
 DEFAULT_ACTION_TIMEOUT = 600  # seconds
-
-
-DEFAULT_ACTION_LISTENER_RETRY_INTERVAL = 5  # 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: Optional[int] = None
-    _labels: dict[str, str | int] = field(default_factory=dict)
+    services: list[str]
+    actions: list[str]
+    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):
+    @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")
+
+    input: JSONSerializableMapping = Field(default_factory=dict)
+    parents: dict[str, JSONSerializableMapping] = Field(default_factory=dict)
+    overrides: JSONSerializableMapping = Field(default_factory=dict)
+    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"
+    )
+    @classmethod
+    def validate_fields(cls, v: Any) -> Any:
+        return v or {}
 
-@dataclass
-class Action:
+
+class ActionType(str, Enum):
+    START_STEP_RUN = "START_STEP_RUN"
+    CANCEL_STEP_RUN = "CANCEL_STEP_RUN"
+    START_GET_GROUP_KEY = "START_GET_GROUP_KEY"
+
+
+class Action(BaseModel):
     worker_id: str
     tenant_id: str
     workflow_run_id: str
@@ -68,31 +96,25 @@ class Action:
     step_id: str
     step_run_id: str
     action_id: str
-    action_payload: str
     action_type: ActionType
     retry_count: int
-    additional_metadata: dict[str, str] | None = None
+    action_payload: ActionPayload
+    additional_metadata: JSONSerializableMapping = field(default_factory=dict)
 
     child_workflow_index: int | None = None
     child_workflow_key: str | None = None
     parent_workflow_run_id: str | None = None
 
-    def __post_init__(self):
-        if isinstance(self.additional_metadata, str) and self.additional_metadata != "":
-            try:
-                self.additional_metadata = json.loads(self.additional_metadata)
-            except json.JSONDecodeError:
-                # If JSON decoding fails, keep the original string
-                pass
-
-        # Ensure additional_metadata is always a dictionary
-        if not isinstance(self.additional_metadata, dict):
-            self.additional_metadata = {}
+    def _dump_payload_to_str(self) -> str:
+        try:
+            return json.dumps(self.action_payload.model_dump(), default=str)
+        except Exception:
+            return str(self.action_payload)
 
     @property
     def otel_attributes(self) -> dict[str, str | int]:
         try:
-            payload_str = json.dumps(self.action_payload, default=str)
+            payload_str = json.dumps(self.action_payload.model_dump(), default=str)
         except Exception:
             payload_str = str(self.action_payload)
 
@@ -115,39 +137,38 @@ class Action:
         return {k: v for k, v in attrs.items() if v}
 
 
-START_STEP_RUN = 0
-CANCEL_STEP_RUN = 1
-START_GET_GROUP_KEY = 2
+def parse_additional_metadata(additional_metadata: str) -> JSONSerializableMapping:
+    try:
+        return cast(
+            JSONSerializableMapping,
+            json.loads(additional_metadata),
+        )
+    except json.JSONDecodeError:
+        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] = 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):
-        self.client = DispatcherStub(new_conn(self.config))
-        self.aio_client = DispatcherStub(new_conn(self.config, True))
+        self.aio_client = DispatcherStub(new_conn(self.config, True))  # type: ignore[no-untyped-call]
         self.token = self.config.token
 
-    def is_healthy(self):
+        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
 
-    async def heartbeat(self):
+    async def heartbeat(self) -> None:
         # send a heartbeat every 4 seconds
         heartbeat_delay = 4
 
@@ -207,7 +228,7 @@ class ActionListener:
                     break
             await asyncio.sleep(heartbeat_delay)
 
-    async def start_heartbeater(self):
+    async def start_heartbeater(self) -> None:
         if self.heartbeat_task is not None:
             return
 
@@ -221,10 +242,10 @@ class ActionListener:
                 raise e
         self.heartbeat_task = loop.create_task(self.heartbeat())
 
-    def __aiter__(self):
+    def __aiter__(self) -> AsyncGenerator[Action | None, None]:
         return self._generator()
 
-    async def _generator(self) -> AsyncGenerator[Action, None]:
+    async def _generator(self) -> AsyncGenerator[Action | None, None]:
         listener = None
 
         while not self.stop_signal:
@@ -239,7 +260,11 @@ class ActionListener:
 
             try:
                 while not self.stop_signal:
-                    self.interrupt = Event_ts()
+                    self.interrupt = ThreadSafeEvent()
+
+                    if listener is None:
+                        continue
+
                     t = asyncio.create_task(
                         read_with_interrupt(listener, self.interrupt)
                     )
@@ -252,7 +277,10 @@ class ActionListener:
                         )
 
                         t.cancel()
-                        listener.cancel()
+
+                        if listener:
+                            listener.cancel()
+
                         break
 
                     assigned_action = t.result()
@@ -262,20 +290,17 @@ class ActionListener:
                         break
 
                     self.retries = 0
-                    assigned_action: AssignedAction
-
-                    # Process the received action
-                    action_type = self.map_action_type(assigned_action.actionType)
 
-                    if (
-                        assigned_action.actionPayload is None
-                        or assigned_action.actionPayload == ""
-                    ):
-                        action_payload = None
-                    else:
-                        action_payload = 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,
@@ -289,9 +314,15 @@ class ActionListener:
                         step_run_id=assigned_action.stepRunId,
                         action_id=assigned_action.actionId,
                         action_payload=action_payload,
-                        action_type=action_type,
+                        action_type=convert_proto_enum_to_python(
+                            assigned_action.actionType,
+                            ActionType,
+                            ActionTypeProto,
+                        ),
                         retry_count=assigned_action.retryCount,
-                        additional_metadata=assigned_action.additional_metadata,
+                        additional_metadata=parse_additional_metadata(
+                            assigned_action.additional_metadata
+                        ),
                         child_workflow_index=assigned_action.child_workflow_index,
                         child_workflow_key=assigned_action.child_workflow_key,
                         parent_workflow_run_id=assigned_action.parent_workflow_run_id,
@@ -325,26 +356,10 @@ class ActionListener:
 
                     self.retries = self.retries + 1
 
-    def parse_action_payload(self, payload: str):
-        try:
-            payload_data = json.loads(payload)
-        except json.JSONDecodeError as e:
-            raise ValueError(f"Error decoding payload: {e}")
-        return payload_data
-
-    def map_action_type(self, action_type):
-        if action_type == ActionType.START_STEP_RUN:
-            return START_STEP_RUN
-        elif action_type == ActionType.CANCEL_STEP_RUN:
-            return CANCEL_STEP_RUN
-        elif action_type == ActionType.START_GET_GROUP_KEY:
-            return START_GET_GROUP_KEY
-        else:
-            # logger.error(f"Unknown action type: {action_type}")
-            return None
-
-    async def get_listen_client(self):
-        current_time = int(time.time())
+    async def get_listen_client(
+        self,
+    ) -> grpc.aio.UnaryStreamCall[WorkerListenRequest, AssignedAction]:
+        current_time = time.time()
 
         if (
             current_time - self.last_connection_attempt
@@ -371,7 +386,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))
+        self.aio_client = DispatcherStub(new_conn(self.config, True))  # type: ignore[no-untyped-call]
 
         if self.listen_strategy == "v2":
             # we should await for the listener to be established before
@@ -392,11 +407,14 @@ class ActionListener:
 
         self.last_connection_attempt = current_time
 
-        return listener
+        return cast(
+            grpc.aio.UnaryStreamCall[WorkerListenRequest, AssignedAction], listener
+        )
 
-    def cleanup(self):
+    def cleanup(self) -> None:
         self.run_heartbeat = False
-        self.heartbeat_task.cancel()
+        if self.heartbeat_task is not None:
+            self.heartbeat_task.cancel()
 
         try:
             self.unregister()
@@ -406,9 +424,11 @@ class ActionListener:
         if self.interrupt:
             self.interrupt.set()
 
-    def unregister(self):
+    def unregister(self) -> WorkerUnsubscribeRequest:
         self.run_heartbeat = False
-        self.heartbeat_task.cancel()
+
+        if self.heartbeat_task is not None:
+            self.heartbeat_task.cancel()
 
         try:
             req = self.aio_client.Unsubscribe(
@@ -416,8 +436,10 @@ class ActionListener:
                 timeout=5,
                 metadata=get_metadata(self.token),
             )
+
             if self.interrupt is not None:
                 self.interrupt.set()
-            return req
+
+            return cast(WorkerUnsubscribeRequest, req)
         except grpc.RpcError as e:
             raise Exception(f"Failed to unsubscribe: {e}")

hatchet_sdk/clients/dispatcher/dispatcher.py

@@ -1,5 +1,6 @@
-from typing import Any, cast
+from typing import cast
 
+import grpc.aio
 from google.protobuf.timestamp_pb2 import Timestamp
 
 from hatchet_sdk.clients.dispatcher.action_listener import (
@@ -8,6 +9,7 @@ from hatchet_sdk.clients.dispatcher.action_listener import (
     GetActionListenerRequest,
 )
 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.dispatcher_pb2 import (
     STEP_EVENT_TYPE_COMPLETED,
@@ -26,22 +28,16 @@ from hatchet_sdk.contracts.dispatcher_pb2 import (
     WorkerRegisterResponse,
 )
 from hatchet_sdk.contracts.dispatcher_pb2_grpc import DispatcherStub
-
-from ...loader import ClientConfig
-from ...metadata import get_metadata
+from hatchet_sdk.metadata import get_metadata
 
 DEFAULT_REGISTER_TIMEOUT = 30
 
 
-def new_dispatcher(config: ClientConfig) -> "DispatcherClient":
-    return DispatcherClient(config=config)
-
-
 class DispatcherClient:
     config: ClientConfig
 
     def __init__(self, config: ClientConfig):
-        conn = new_conn(config)
+        conn = new_conn(config, False)
         self.client = DispatcherStub(conn)  # type: ignore[no-untyped-call]
 
         aio_conn = new_conn(config, True)
@@ -59,24 +55,26 @@ 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)
 
     async def send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
-    ) -> Any:
+    ) -> grpc.aio.UnaryUnaryCall[StepActionEvent, ActionEventResponse] | None:
         try:
             return await self._try_send_step_action_event(action, event_type, payload)
         except Exception as e:
@@ -91,14 +89,14 @@ class DispatcherClient:
                     "Failed to send finished event: " + str(e),
                 )
 
-            return
+            return None
 
     @tenacity_retry
     async def _try_send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
-    ) -> Any:
-        eventTimestamp = Timestamp()
-        eventTimestamp.GetCurrentTime()
+    ) -> grpc.aio.UnaryUnaryCall[StepActionEvent, ActionEventResponse]:
+        event_timestamp = Timestamp()
+        event_timestamp.GetCurrentTime()
 
         event = StepActionEvent(
             workerId=action.worker_id,
@@ -107,38 +105,42 @@ 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,
         )
 
-        ## TODO: What does this return?
-        return await self.aio_client.SendStepActionEvent(
-            event,
-            metadata=get_metadata(self.token),
+        return cast(
+            grpc.aio.UnaryUnaryCall[StepActionEvent, ActionEventResponse],
+            await self.aio_client.SendStepActionEvent(
+                event,
+                metadata=get_metadata(self.token),
+            ),
         )
 
     async def send_group_key_action_event(
         self, action: Action, event_type: GroupKeyActionEventType, payload: str
-    ) -> Any:
-        eventTimestamp = Timestamp()
-        eventTimestamp.GetCurrentTime()
+    ) -> grpc.aio.UnaryUnaryCall[GroupKeyActionEvent, ActionEventResponse]:
+        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,
         )
 
-        ## TODO: What does this return?
-        return await self.aio_client.SendGroupKeyActionEvent(
-            event,
-            metadata=get_metadata(self.token),
+        return cast(
+            grpc.aio.UnaryUnaryCall[GroupKeyActionEvent, ActionEventResponse],
+            await self.aio_client.SendGroupKeyActionEvent(
+                event,
+                metadata=get_metadata(self.token),
+            ),
         )
 
     def put_overrides_data(self, data: OverridesData) -> ActionEventResponse: