hatchet-sdk 1.6.2__py3-none-any.whl → 1.6.4__py3-none-any.whl

hatchet_sdk/clients/dispatcher/action_listener.py

@@ -9,7 +9,11 @@ import grpc
 import grpc.aio
 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.event_ts import (
+    ThreadSafeEvent,
+    UnexpectedEOF,
+    read_with_interrupt,
+)
 from hatchet_sdk.clients.events import proto_timestamp_now
 from hatchet_sdk.clients.listeners.run_event_listener import (
     DEFAULT_ACTION_LISTENER_RETRY_INTERVAL,
@@ -84,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
@@ -137,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:
@@ -275,15 +294,17 @@ class ActionListener:
 
                         break
 
-                    assigned_action, _, is_eof = t.result()
+                    result = t.result()
 
-                    if is_eof:
+                    if isinstance(result, UnexpectedEOF):
                         logger.debug("Handling EOF in Action Listener")
                         self.retries = self.retries + 1
                         break
 
                     self.retries = 0
 
+                    assigned_action = result.data
+
                     try:
                         action_payload = (
                             ActionPayload()

hatchet_sdk/clients/event_ts.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Callable, TypeVar, cast, overload
+from typing import Callable, Generic, TypeVar, cast, overload
 
 import grpc.aio
 from grpc._cython import cygrpc  # type: ignore[attr-defined]
@@ -29,12 +29,23 @@ TRequest = TypeVar("TRequest")
 TResponse = TypeVar("TResponse")
 
 
+class ReadWithInterruptResult(Generic[TResponse]):
+    def __init__(self, data: TResponse, key: str):
+        self.data = data
+        self.key = key
+
+
+class UnexpectedEOF:
+    def __init__(self) -> None:
+        pass
+
+
 @overload
 async def read_with_interrupt(
     listener: grpc.aio.UnaryStreamCall[TRequest, TResponse],
     interrupt: ThreadSafeEvent,
     key_generator: Callable[[TResponse], str],
-) -> tuple[TResponse, str, bool]: ...
+) -> ReadWithInterruptResult[TResponse] | UnexpectedEOF: ...
 
 
 @overload
@@ -42,23 +53,23 @@ async def read_with_interrupt(
     listener: grpc.aio.UnaryStreamCall[TRequest, TResponse],
     interrupt: ThreadSafeEvent,
     key_generator: None = None,
-) -> tuple[TResponse, None, bool]: ...
+) -> ReadWithInterruptResult[TResponse] | UnexpectedEOF: ...
 
 
 async def read_with_interrupt(
     listener: grpc.aio.UnaryStreamCall[TRequest, TResponse],
     interrupt: ThreadSafeEvent,
     key_generator: Callable[[TResponse], str] | None = None,
-) -> tuple[TResponse, str | None, bool]:
+) -> ReadWithInterruptResult[TResponse] | UnexpectedEOF:
     try:
         result = cast(TResponse, await listener.read())
 
         if result is cygrpc.EOF:
             logger.warning("Received EOF from engine")
-            return cast(TResponse, None), None, True
+            return UnexpectedEOF()
 
-        key = key_generator(result) if key_generator else None
+        key = key_generator(result) if key_generator else ""
 
-        return result, key, False
+        return ReadWithInterruptResult(data=result, key=key)
     finally:
         interrupt.set()

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/clients/listeners/pooled_listener.py

@@ -6,7 +6,11 @@ from typing import Generic, Literal, TypeVar
 import grpc
 import grpc.aio
 
-from hatchet_sdk.clients.event_ts import ThreadSafeEvent, read_with_interrupt
+from hatchet_sdk.clients.event_ts import (
+    ThreadSafeEvent,
+    UnexpectedEOF,
+    read_with_interrupt,
+)
 from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.logger import logger
 from hatchet_sdk.metadata import get_metadata
@@ -130,18 +134,18 @@ class PooledListener(Generic[R, T, L], ABC):
                                 await asyncio.sleep(DEFAULT_LISTENER_RETRY_INTERVAL)
                                 break
 
-                            event, key, is_eof = t.result()
+                            event = t.result()
 
-                            if is_eof:
+                            if isinstance(event, UnexpectedEOF):
                                 logger.debug(
                                     f"Handling EOF in Pooled Listener {self.__class__.__name__}"
                                 )
                                 break
 
-                            subscriptions = self.to_subscriptions.get(key, [])
+                            subscriptions = self.to_subscriptions.get(event.key, [])
 
                             for subscription_id in subscriptions:
-                                await self.events[subscription_id].put(event)
+                                await self.events[subscription_id].put(event.data)
 
                     except grpc.RpcError as e:
                         logger.debug(f"grpc error in listener: {e}")

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/features/cron.py

@@ -26,11 +26,6 @@ from hatchet_sdk.utils.typing import JSONSerializableMapping
 class CreateCronTriggerConfig(BaseModel):
     """
     Schema for creating a workflow run triggered by a cron.
-
-    Attributes:
-        expression (str): The cron expression defining the schedule.
-        input (dict): The input data for the cron workflow.
-        additional_metadata (dict[str, str]): Additional metadata associated with the cron trigger (e.g. {"key1": "value1", "key2": "value2"}).
     """
 
     expression: str
@@ -43,14 +38,11 @@ class CreateCronTriggerConfig(BaseModel):
         """
         Validates the cron expression to ensure it adheres to the expected format.
 
-        Args:
-            v (str): The cron expression to validate.
+        :param v: The cron expression to validate.
 
-        Raises:
-            ValueError: If the expression is invalid.
+        :raises ValueError: If the expression is invalid
 
-        Returns:
-            str: The validated cron expression.
+        :return: The validated cron expression.
         """
         if not v:
             raise ValueError("Cron expression is required")
@@ -72,6 +64,10 @@ class CreateCronTriggerConfig(BaseModel):
 
 
 class CronClient(BaseRestClient):
+    """
+    The cron client is a client for managing cron workflows within Hatchet.
+    """
+
     def _wra(self, client: ApiClient) -> WorkflowRunApi:
         return WorkflowRunApi(client)
 
@@ -88,17 +84,16 @@ class CronClient(BaseRestClient):
         priority: int | None = None,
     ) -> CronWorkflows:
         """
-        Asynchronously creates a new workflow cron trigger.
+        Create a new workflow cron trigger.
 
-        Args:
-            workflow_name (str): The name of the workflow to trigger.
-            cron_name (str): The name of the cron trigger.
-            expression (str): The cron expression defining the schedule.
-            input (dict): The input data for the cron workflow.
-            additional_metadata (dict[str, str]): Additional metadata associated with the cron trigger (e.g. {"key1": "value1", "key2": "value2"}).
+        :param workflow_name: The name of the workflow to trigger.
+        :param cron_name: The name of the cron trigger.
+        :param expression: The cron expression defining the schedule.
+        :param input: The input data for the cron workflow.
+        :param additional_metadata: Additional metadata associated with the cron trigger.
+        :param priority: The priority of the cron workflow trigger.
 
-        Returns:
-            CronWorkflows: The created cron workflow instance.
+        :return: The created cron workflow instance.
         """
         validated_input = CreateCronTriggerConfig(
             expression=expression, input=input, additional_metadata=additional_metadata
@@ -126,6 +121,18 @@ class CronClient(BaseRestClient):
         additional_metadata: JSONSerializableMapping,
         priority: int | None = None,
     ) -> CronWorkflows:
+        """
+        Create a new workflow cron trigger.
+
+        :param workflow_name: The name of the workflow to trigger.
+        :param cron_name: The name of the cron trigger.
+        :param expression: The cron expression defining the schedule.
+        :param input: The input data for the cron workflow.
+        :param additional_metadata: Additional metadata associated with the cron trigger.
+        :param priority: The priority of the cron workflow trigger.
+
+        :return: The created cron workflow instance.
+        """
         return await asyncio.to_thread(
             self.create,
             workflow_name,
@@ -138,10 +145,10 @@ class CronClient(BaseRestClient):
 
     def delete(self, cron_id: str) -> None:
         """
-        Asynchronously deletes a workflow cron trigger.
+        Delete a workflow cron trigger.
 
-        Args:
-            cron_id (str): The cron trigger ID or CronWorkflows instance to delete.
+        :param cron_id: The ID of the cron trigger to delete.
+        :return: None
         """
         with self.client() as client:
             self._wa(client).workflow_cron_delete(
@@ -149,6 +156,12 @@ class CronClient(BaseRestClient):
             )
 
     async def aio_delete(self, cron_id: str) -> None:
+        """
+        Delete a workflow cron trigger.
+
+        :param cron_id: The ID of the cron trigger to delete.
+        :return: None
+        """
         return await asyncio.to_thread(self.delete, cron_id)
 
     async def aio_list(
@@ -161,18 +174,16 @@ class CronClient(BaseRestClient):
         order_by_direction: WorkflowRunOrderByDirection | None = None,
     ) -> CronWorkflowsList:
         """
-        Synchronously retrieves a list of all workflow cron triggers matching the criteria.
+        Retrieve a list of all workflow cron triggers matching the criteria.
 
-        Args:
-            offset (int | None): The offset to start the list from.
-            limit (int | None): The maximum number of items to return.
-            workflow_id (str | None): The ID of the workflow to filter by.
-            additional_metadata (list[str] | None): Filter by additional metadata keys (e.g. ["key1:value1", "key2:value2"]).
-            order_by_field (CronWorkflowsOrderByField | None): The field to order the list by.
-            order_by_direction (WorkflowRunOrderByDirection | None): The direction to order the list by.
+        :param offset: The offset to start the list from.
+        :param limit: The maximum number of items to return.
+        :param workflow_id: The ID of the workflow to filter by.
+        :param additional_metadata: Filter by additional metadata keys.
+        :param order_by_field: The field to order the list by.
+        :param order_by_direction: The direction to order the list by.
 
-        Returns:
-            CronWorkflowsList: A list of cron workflows.
+        :return: A list of cron workflows.
         """
         return await asyncio.to_thread(
             self.list,
@@ -194,18 +205,16 @@ class CronClient(BaseRestClient):
         order_by_direction: WorkflowRunOrderByDirection | None = None,
     ) -> CronWorkflowsList:
         """
-        Asynchronously retrieves a list of all workflow cron triggers matching the criteria.
+        Retrieve a list of all workflow cron triggers matching the criteria.
 
-        Args:
-            offset (int | None): The offset to start the list from.
-            limit (int | None): The maximum number of items to return.
-            workflow_id (str | None): The ID of the workflow to filter by.
-            additional_metadata (list[str] | None): Filter by additional metadata keys (e.g. ["key1:value1", "key2:value2"]).
-            order_by_field (CronWorkflowsOrderByField | None): The field to order the list by.
-            order_by_direction (WorkflowRunOrderByDirection | None): The direction to order the list by.
+        :param offset: The offset to start the list from.
+        :param limit: The maximum number of items to return.
+        :param workflow_id: The ID of the workflow to filter by.
+        :param additional_metadata: Filter by additional metadata keys.
+        :param order_by_field: The field to order the list by.
+        :param order_by_direction: The direction to order the list by.
 
-        Returns:
-            CronWorkflowsList: A list of cron workflows.
+        :return: A list of cron workflows.
         """
         with self.client() as client:
             return self._wa(client).cron_workflow_list(
@@ -222,13 +231,10 @@ class CronClient(BaseRestClient):
 
     def get(self, cron_id: str) -> CronWorkflows:
         """
-        Asynchronously retrieves a specific workflow cron trigger by ID.
-
-        Args:
-            cron_id (str): The cron trigger ID or CronWorkflows instance to retrieve.
+        Retrieve a specific workflow cron trigger by ID.
 
-        Returns:
-            CronWorkflows: The requested cron workflow instance.
+        :param cron_id: The cron trigger ID or CronWorkflows instance to retrieve.
+        :return: The requested cron workflow instance.
         """
         with self.client() as client:
             return self._wa(client).workflow_cron_get(
@@ -237,12 +243,9 @@ class CronClient(BaseRestClient):
 
     async def aio_get(self, cron_id: str) -> CronWorkflows:
         """
-        Synchronously retrieves a specific workflow cron trigger by ID.
-
-        Args:
-            cron_id (str): The cron trigger ID or CronWorkflows instance to retrieve.
+        Retrieve a specific workflow cron trigger by ID.
 
-        Returns:
-            CronWorkflows: The requested cron workflow instance.
+        :param cron_id: The cron trigger ID or CronWorkflows instance to retrieve.
+        :return: The requested cron workflow instance.
         """
         return await asyncio.to_thread(self.get, cron_id)

hatchet_sdk/features/logs.py

@@ -7,12 +7,28 @@ from hatchet_sdk.clients.v1.api_client import BaseRestClient
 
 
 class LogsClient(BaseRestClient):
+    """
+    The logs client is a client for interacting with Hatchet's logs API.
+    """
+
     def _la(self, client: ApiClient) -> LogApi:
         return LogApi(client)
 
     def list(self, task_run_id: str) -> V1LogLineList:
+        """
+        List log lines for a given task run.
+
+        :param task_run_id: The ID of the task run to list logs for.
+        :return: A list of log lines for the specified task run.
+        """
         with self.client() as client:
             return self._la(client).v1_log_line_list(task=task_run_id)
 
     async def aio_list(self, task_run_id: str) -> V1LogLineList:
+        """
+        List log lines for a given task run.
+
+        :param task_run_id: The ID of the task run to list logs for.
+        :return: A list of log lines for the specified task run.
+        """
         return await asyncio.to_thread(self.list, task_run_id)

hatchet_sdk/features/metrics.py

@@ -17,6 +17,10 @@ from hatchet_sdk.utils.typing import JSONSerializableMapping
 
 
 class MetricsClient(BaseRestClient):
+    """
+    The metrics client is a client for reading metrics out of Hatchet's metrics API.
+    """
+
     def _wa(self, client: ApiClient) -> WorkflowApi:
         return WorkflowApi(client)
 
@@ -29,6 +33,15 @@ class MetricsClient(BaseRestClient):
         status: WorkflowRunStatus | None = None,
         group_key: str | None = None,
     ) -> WorkflowMetrics:
+        """
+        Retrieve workflow metrics for a given workflow ID.
+
+        :param workflow_id: The ID of the workflow to retrieve metrics for.
+        :param status: The status of the workflow run to filter by.
+        :param group_key: The key to group the metrics by.
+
+        :return: Workflow metrics for the specified workflow ID.
+        """
         with self.client() as client:
             return self._wa(client).workflow_get_metrics(
                 workflow=workflow_id, status=status, group_key=group_key
@@ -40,6 +53,15 @@ class MetricsClient(BaseRestClient):
         status: WorkflowRunStatus | None = None,
         group_key: str | None = None,
     ) -> WorkflowMetrics:
+        """
+        Retrieve workflow metrics for a given workflow ID.
+
+        :param workflow_id: The ID of the workflow to retrieve metrics for.
+        :param status: The status of the workflow run to filter by.
+        :param group_key: The key to group the metrics by.
+
+        :return: Workflow metrics for the specified workflow ID.
+        """
         return await asyncio.to_thread(
             self.get_workflow_metrics, workflow_id, status, group_key
         )
@@ -49,6 +71,14 @@ class MetricsClient(BaseRestClient):
         workflow_ids: list[str] | None = None,
         additional_metadata: JSONSerializableMapping | None = None,
     ) -> TenantQueueMetrics:
+        """
+        Retrieve queue metrics for a set of workflow ids and additional metadata.
+
+        :param workflow_ids: A list of workflow IDs to retrieve metrics for.
+        :param additional_metadata: Additional metadata to filter the metrics by.
+
+        :return: Workflow metrics for the specified workflow IDs.
+        """
         with self.client() as client:
             return self._wa(client).tenant_get_queue_metrics(
                 tenant=self.client_config.tenant_id,
@@ -63,15 +93,33 @@ class MetricsClient(BaseRestClient):
         workflow_ids: list[str] | None = None,
         additional_metadata: JSONSerializableMapping | None = None,
     ) -> TenantQueueMetrics:
+        """
+        Retrieve queue metrics for a set of workflow ids and additional metadata.
+
+        :param workflow_ids: A list of workflow IDs to retrieve metrics for.
+        :param additional_metadata: Additional metadata to filter the metrics by.
+
+        :return: Workflow metrics for the specified workflow IDs.
+        """
         return await asyncio.to_thread(
             self.get_queue_metrics, workflow_ids, additional_metadata
         )
 
     def get_task_metrics(self) -> TenantStepRunQueueMetrics:
+        """
+        Retrieve queue metrics
+
+        :return: Step run queue metrics for the tenant
+        """
         with self.client() as client:
             return self._ta(client).tenant_get_step_run_queue_metrics(
                 tenant=self.client_config.tenant_id
             )
 
     async def aio_get_task_metrics(self) -> TenantStepRunQueueMetrics:
+        """
+        Retrieve queue metrics
+
+        :return: Step run queue metrics for the tenant
+        """
         return await asyncio.to_thread(self.get_task_metrics)

hatchet_sdk/features/rate_limits.py

@@ -12,6 +12,10 @@ from hatchet_sdk.utils.proto_enums import convert_python_enum_to_proto
 
 
 class RateLimitsClient(BaseRestClient):
+    """
+    The rate limits client is a wrapper for Hatchet's gRPC API that makes it easier to work with rate limits in Hatchet.
+    """
+
     @tenacity_retry
     def put(
         self,
@@ -19,6 +23,16 @@ class RateLimitsClient(BaseRestClient):
         limit: int,
         duration: RateLimitDuration = RateLimitDuration.SECOND,
     ) -> None:
+        """
+        Put a rate limit for a given key.
+
+        :param key: The key to set the rate limit for.
+        :param limit: The rate limit to set.
+        :param duration: The duration of the rate limit.
+
+        :return: None
+        """
+
         duration_proto = convert_python_enum_to_proto(
             duration, workflow_protos.RateLimitDuration
         )
@@ -42,4 +56,14 @@ class RateLimitsClient(BaseRestClient):
         limit: int,
         duration: RateLimitDuration = RateLimitDuration.SECOND,
     ) -> None:
+        """
+        Put a rate limit for a given key.
+
+        :param key: The key to set the rate limit for.
+        :param limit: The rate limit to set.
+        :param duration: The duration of the rate limit.
+
+        :return: None
+        """
+
         await asyncio.to_thread(self.put, key, limit, duration)