hatchet-sdk 1.20.1__tar.gz → 1.21.0__tar.gz

{hatchet_sdk-1.20.1 → hatchet_sdk-1.21.0}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.1
 Name: hatchet-sdk
-Version: 1.20.1
-Summary: 
+Version: 1.21.0
+Summary: This is the official Python SDK for Hatchet, a distributed, fault-tolerant task queue. The SDK allows you to easily integrate Hatchet's task scheduling and workflow orchestration capabilities into your Python applications.
 License: MIT
 Author: Alexander Belanger
 Author-email: alexander@hatchet.run
@@ -14,10 +14,8 @@ Classifier: Programming Language :: Python :: 3.12
 Provides-Extra: otel
 Provides-Extra: v0-sdk
 Requires-Dist: aiohttp (>=3.10.5,<4.0.0)
-Requires-Dist: grpcio (>=1.64.1,!=1.68.*) ; python_version < "3.13"
-Requires-Dist: grpcio (>=1.69.0) ; python_version >= "3.13"
-Requires-Dist: grpcio-tools (>=1.64.1,!=1.68.*) ; python_version < "3.13"
-Requires-Dist: grpcio-tools (>=1.69.0) ; python_version >= "3.13"
+Requires-Dist: grpcio (>=1.76.0,<2.0.0)
+Requires-Dist: grpcio-tools (>=1.76.0,<2.0.0)
 Requires-Dist: opentelemetry-api (>=1.28.0,<2.0.0) ; extra == "otel"
 Requires-Dist: opentelemetry-distro (>=0.49b0) ; extra == "otel"
 Requires-Dist: opentelemetry-exporter-otlp (>=1.28.0,<2.0.0) ; extra == "otel"
@@ -25,7 +23,7 @@ Requires-Dist: opentelemetry-exporter-otlp-proto-http (>=1.28.0,<2.0.0) ; extra
 Requires-Dist: opentelemetry-instrumentation (>=0.49b0) ; extra == "otel"
 Requires-Dist: opentelemetry-sdk (>=1.28.0,<2.0.0) ; extra == "otel"
 Requires-Dist: prometheus-client (>=0.21.1)
-Requires-Dist: protobuf (>=5.29.5,<6.0.0)
+Requires-Dist: protobuf (>=6.30.0,<7.0.0)
 Requires-Dist: pydantic (>=2.6.3,<3.0.0)
 Requires-Dist: pydantic-settings (>=2.7.1,<3.0.0)
 Requires-Dist: python-dateutil (>=2.9.0.post0,<3.0.0)

{hatchet_sdk-1.20.1 → hatchet_sdk-1.21.0}/hatchet_sdk/context/context.py

@@ -4,6 +4,8 @@ from datetime import timedelta
 from typing import TYPE_CHECKING, Any, cast
 from warnings import warn
 
+from pydantic import TypeAdapter
+
 from hatchet_sdk.clients.admin import AdminClient
 from hatchet_sdk.clients.dispatcher.dispatcher import (  # type: ignore[attr-defined]
     Action,
@@ -25,7 +27,13 @@ from hatchet_sdk.exceptions import TaskRunError
 from hatchet_sdk.features.runs import RunsClient
 from hatchet_sdk.logger import logger
 from hatchet_sdk.utils.timedelta_to_expression import Duration, timedelta_to_expr
-from hatchet_sdk.utils.typing import JSONSerializableMapping, LogLevel
+from hatchet_sdk.utils.typing import (
+    JSONSerializableMapping,
+    LogLevel,
+    classify_output_validator,
+    is_basemodel_validator,
+    is_dataclass_validator,
+)
 from hatchet_sdk.worker.runner.utils.capture_logs import AsyncLogSender, LogRecord
 
 if TYPE_CHECKING:
@@ -106,7 +114,21 @@ class Context:
             raise ValueError(f"Step output for '{task.name}' not found") from e
 
         if parent_step_data and (v := task.validators.step_output):
-            return cast(R, v.model_validate(parent_step_data))
+            validator = classify_output_validator(v)
+
+            if is_dataclass_validator(validator):
+                return cast(
+                    R,
+                    TypeAdapter(validator.validator_type).validate_python(
+                        parent_step_data
+                    ),
+                )
+
+            if is_basemodel_validator(validator):
+                return cast(
+                    R,
+                    validator.validator_type.model_validate(parent_step_data),
+                )
 
         return parent_step_data
 

{hatchet_sdk-1.20.1 → hatchet_sdk-1.21.0}/hatchet_sdk/exceptions.py

@@ -167,3 +167,7 @@ class LoopAlreadyRunningError(Exception):
 
 class IllegalTaskOutputError(Exception):
     pass
+
+
+class LifespanSetupError(Exception):
+    pass

{hatchet_sdk-1.20.1 → hatchet_sdk-1.21.0}/hatchet_sdk/features/runs.py

@@ -405,6 +405,7 @@ class RunsClient(BaseRestClient):
         worker_id: str | None = None,
         parent_task_external_id: str | None = None,
         triggering_event_external_id: str | None = None,
+        include_payloads: bool = True,
     ) -> list[V1TaskSummary]:
         """
         List task runs according to a set of filters, paginating through days
@@ -420,6 +421,7 @@ class RunsClient(BaseRestClient):
         :param worker_id: The worker ID to filter task runs by.
         :param parent_task_external_id: The parent task external ID to filter task runs by.
         :param triggering_event_external_id: The event id that triggered the task run.
+        :param include_payloads: Whether to include payloads in the response.
 
         :return: A list of task runs matching the specified filters.
         """
@@ -446,6 +448,7 @@ class RunsClient(BaseRestClient):
                     worker_id=worker_id,
                     parent_task_external_id=parent_task_external_id,
                     triggering_event_external_id=triggering_event_external_id,
+                    include_payloads=include_payloads,
                 )
                 for s, u in date_ranges
             ]
@@ -475,6 +478,7 @@ class RunsClient(BaseRestClient):
         worker_id: str | None = None,
         parent_task_external_id: str | None = None,
         triggering_event_external_id: str | None = None,
+        include_payloads: bool = True,
     ) -> list[V1TaskSummary]:
         """
         List task runs according to a set of filters, paginating through days
@@ -490,6 +494,7 @@ class RunsClient(BaseRestClient):
         :param worker_id: The worker ID to filter task runs by.
         :param parent_task_external_id: The parent task external ID to filter task runs by.
         :param triggering_event_external_id: The event id that triggered the task run.
+        :param include_payloads: Whether to include payloads in the response.
 
         :return: A list of task runs matching the specified filters.
         """
@@ -517,6 +522,7 @@ class RunsClient(BaseRestClient):
                     worker_id=worker_id,
                     parent_task_external_id=parent_task_external_id,
                     triggering_event_external_id=triggering_event_external_id,
+                    include_payloads=include_payloads,
                 )
                 for s, u in date_ranges
             ]
@@ -550,6 +556,7 @@ class RunsClient(BaseRestClient):
         worker_id: str | None = None,
         parent_task_external_id: str | None = None,
         triggering_event_external_id: str | None = None,
+        include_payloads: bool = True,
     ) -> V1TaskSummaryList:
         """
         List task runs according to a set of filters.
@@ -565,6 +572,7 @@ class RunsClient(BaseRestClient):
         :param worker_id: The worker ID to filter task runs by.
         :param parent_task_external_id: The parent task external ID to filter task runs by.
         :param triggering_event_external_id: The event id that triggered the task run.
+        :param include_payloads: Whether to include payloads in the response.
 
         :return: A list of task runs matching the specified filters.
         """
@@ -581,6 +589,7 @@ class RunsClient(BaseRestClient):
             worker_id=worker_id,
             parent_task_external_id=parent_task_external_id,
             triggering_event_external_id=triggering_event_external_id,
+            include_payloads=include_payloads,
         )
 
     @retry
@@ -597,6 +606,7 @@ class RunsClient(BaseRestClient):
         worker_id: str | None = None,
         parent_task_external_id: str | None = None,
         triggering_event_external_id: str | None = None,
+        include_payloads: bool = True,
     ) -> V1TaskSummaryList:
         """
         List task runs according to a set of filters.
@@ -612,6 +622,7 @@ class RunsClient(BaseRestClient):
         :param worker_id: The worker ID to filter task runs by.
         :param parent_task_external_id: The parent task external ID to filter task runs by.
         :param triggering_event_external_id: The event id that triggered the task run.
+        :param include_payloads: Whether to include payloads in the response.
 
         :return: A list of task runs matching the specified filters.
         """
@@ -643,6 +654,7 @@ class RunsClient(BaseRestClient):
                 worker_id=worker_id,
                 parent_task_external_id=parent_task_external_id,
                 triggering_event_external_id=triggering_event_external_id,
+                include_payloads=include_payloads,
             )
 
     def create(

{hatchet_sdk-1.20.1 → hatchet_sdk-1.21.0}/hatchet_sdk/hatchet.py

@@ -4,6 +4,8 @@ from collections.abc import Callable
 from datetime import timedelta
 from typing import Any, Concatenate, ParamSpec, cast, overload
 
+from pydantic import BaseModel
+
 from hatchet_sdk import Context, DurableContext
 from hatchet_sdk.client import Client
 from hatchet_sdk.clients.dispatcher.dispatcher import DispatcherClient
@@ -36,7 +38,7 @@ from hatchet_sdk.runnables.types import (
 )
 from hatchet_sdk.runnables.workflow import BaseWorkflow, Standalone, Workflow
 from hatchet_sdk.utils.timedelta_to_expression import Duration
-from hatchet_sdk.utils.typing import CoroutineLike
+from hatchet_sdk.utils.typing import CoroutineLike, DataclassInstance
 from hatchet_sdk.worker.worker import LifespanFn, Worker
 
 P = ParamSpec("P")
@@ -302,8 +304,10 @@ class Hatchet:
                 on_crons=on_crons or [],
                 sticky=sticky,
                 concurrency=concurrency,
-                input_validator=input_validator
-                or cast(type[TWorkflowInput], EmptyModel),
+                input_validator=cast(
+                    type[BaseModel] | type[DataclassInstance],
+                    input_validator or EmptyModel,
+                ),
                 task_defaults=task_defaults,
                 default_priority=default_priority,
                 default_filters=default_filters or [],
@@ -449,8 +453,10 @@ class Hatchet:
                     on_crons=on_crons or [],
                     sticky=sticky,
                     default_priority=default_priority,
-                    input_validator=input_validator
-                    or cast(type[TWorkflowInput], EmptyModel),
+                    input_validator=cast(
+                        type[BaseModel] | type[DataclassInstance],
+                        input_validator or EmptyModel,
+                    ),
                     default_filters=default_filters or [],
                 ),
                 self,
@@ -633,8 +639,10 @@ class Hatchet:
                     on_events=on_events or [],
                     on_crons=on_crons or [],
                     sticky=sticky,
-                    input_validator=input_validator
-                    or cast(type[TWorkflowInput], EmptyModel),
+                    input_validator=cast(
+                        type[BaseModel] | type[DataclassInstance],
+                        input_validator or EmptyModel,
+                    ),
                     default_priority=default_priority,
                     default_filters=default_filters or [],
                 ),

{hatchet_sdk-1.20.1 → hatchet_sdk-1.21.0}/hatchet_sdk/runnables/task.py

@@ -1,5 +1,6 @@
 import asyncio
 from collections.abc import Callable
+from dataclasses import asdict, is_dataclass
 from inspect import Parameter, iscoroutinefunction, signature
 from typing import (
     TYPE_CHECKING,
@@ -37,7 +38,6 @@ from hatchet_sdk.contracts.v1.workflows_pb2 import (
 from hatchet_sdk.exceptions import InvalidDependencyError
 from hatchet_sdk.runnables.types import (
     ConcurrencyExpression,
-    EmptyModel,
     R,
     StepType,
     TWorkflowInput,
@@ -289,11 +289,14 @@ class Task(Generic[TWorkflowInput, R]):
 
         additional_metadata = additional_metadata or {}
         parent_outputs = parent_outputs or {}
+        serialized_input: dict[str, Any] = {}
 
-        if input is None:
-            input = cast(TWorkflowInput, EmptyModel())
+        if is_dataclass(input):
+            serialized_input = asdict(input)
+        elif isinstance(input, BaseModel):
+            serialized_input = input.model_dump()
 
-        action_payload = ActionPayload(input=input.model_dump(), parents=parent_outputs)
+        action_payload = ActionPayload(input=serialized_input, parents=parent_outputs)
 
         action = Action(
             tenant_id=self.workflow.client.config.tenant_id,

{hatchet_sdk-1.20.1 → hatchet_sdk-1.21.0}/hatchet_sdk/runnables/types.py

@@ -4,15 +4,19 @@ from collections.abc import Callable, Mapping
 from enum import Enum
 from typing import Any, ParamSpec, TypeGuard, TypeVar
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, SkipValidation
 
 from hatchet_sdk.context.context import Context, DurableContext
 from hatchet_sdk.contracts.v1.workflows_pb2 import Concurrency
 from hatchet_sdk.contracts.v1.workflows_pb2 import DefaultFilter as DefaultFilterProto
 from hatchet_sdk.utils.timedelta_to_expression import Duration
-from hatchet_sdk.utils.typing import AwaitableLike, JSONSerializableMapping
+from hatchet_sdk.utils.typing import (
+    AwaitableLike,
+    DataclassInstance,
+    JSONSerializableMapping,
+)
 
-ValidTaskReturnType = BaseModel | Mapping[str, Any] | None
+ValidTaskReturnType = BaseModel | Mapping[str, Any] | DataclassInstance | None
 
 R = TypeVar("R", bound=ValidTaskReturnType)
 P = ParamSpec("P")
@@ -56,7 +60,7 @@ class ConcurrencyExpression(BaseModel):
         )
 
 
-TWorkflowInput = TypeVar("TWorkflowInput", bound=BaseModel)
+TWorkflowInput = TypeVar("TWorkflowInput", bound=BaseModel | DataclassInstance)
 
 
 class TaskDefaults(BaseModel):
@@ -93,7 +97,9 @@ class WorkflowConfig(BaseModel):
     on_crons: list[str] = Field(default_factory=list)
     sticky: StickyStrategy | None = None
     concurrency: ConcurrencyExpression | list[ConcurrencyExpression] | None = None
-    input_validator: type[BaseModel] = EmptyModel
+    input_validator: SkipValidation[type[BaseModel] | type[DataclassInstance]] = (
+        EmptyModel
+    )
     default_priority: int | None = None
 
     task_defaults: TaskDefaults = TaskDefaults()

{hatchet_sdk-1.20.1 → hatchet_sdk-1.21.0}/hatchet_sdk/runnables/workflow.py

@@ -1,5 +1,6 @@
 import asyncio
 from collections.abc import Callable
+from dataclasses import asdict
 from datetime import datetime, timedelta
 from functools import cached_property
 from typing import (
@@ -16,7 +17,7 @@ from typing import (
 )
 
 from google.protobuf import timestamp_pb2
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel, ConfigDict, SkipValidation, TypeAdapter, model_validator
 
 from hatchet_sdk.clients.admin import (
     ScheduleTriggerWorkflowOptions,
@@ -52,8 +53,12 @@ from hatchet_sdk.utils.proto_enums import convert_python_enum_to_proto
 from hatchet_sdk.utils.timedelta_to_expression import Duration
 from hatchet_sdk.utils.typing import (
     CoroutineLike,
+    DataclassInstance,
     JSONSerializableMapping,
-    is_basemodel_subclass,
+    classify_output_validator,
+    is_basemodel_validator,
+    is_dataclass_validator,
+    is_no_validator,
 )
 from hatchet_sdk.workflow_run import WorkflowRunRef
 
@@ -130,7 +135,8 @@ def transform_desired_worker_label(d: DesiredWorkerLabel) -> DesiredWorkerLabels
 
 
 class TypedTriggerWorkflowRunConfig(BaseModel, Generic[TWorkflowInput]):
-    input: TWorkflowInput
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    input: SkipValidation[TWorkflowInput]
     options: TriggerWorkflowOptions
 
 
@@ -217,10 +223,26 @@ class BaseWorkflow(Generic[TWorkflowInput]):
         )
 
     def _get_workflow_input(self, ctx: Context) -> TWorkflowInput:
-        return cast(
-            TWorkflowInput,
-            self.config.input_validator.model_validate(ctx.workflow_input),
-        )
+        validator = classify_output_validator(self.config.input_validator)
+
+        if is_dataclass_validator(validator):
+            return cast(
+                TWorkflowInput,
+                TypeAdapter(validator.validator_type).validate_python(
+                    ctx.workflow_input
+                ),
+            )
+
+        if is_basemodel_validator(validator):
+            return cast(
+                TWorkflowInput,
+                validator.validator_type.model_validate(ctx.workflow_input),
+            )
+
+        ## impossible to reach here since the input validator has to be either a BaseModel or dataclass
+
+        self.client.config.logger.error("input validator is of an unknown type")
+        return cast(TWorkflowInput, EmptyModel())
 
     @property
     def input_validator(self) -> type[TWorkflowInput]:
@@ -271,11 +293,16 @@ class BaseWorkflow(Generic[TWorkflowInput]):
         if not input:
             return {}
 
-        if isinstance(input, BaseModel):
-            return input.model_dump(mode="json")
+        validator = classify_output_validator(self.config.input_validator)
+
+        if is_dataclass_validator(validator):
+            return asdict(cast(DataclassInstance, input))
+
+        if is_basemodel_validator(validator):
+            return cast(BaseModel, input).model_dump(mode="json")
 
         raise ValueError(
-            f"Input must be a BaseModel or `None`, got {type(input)} instead."
+            f"Input must be a BaseModel or dataclass, got {type(input)} instead."
         )
 
     @cached_property
@@ -1203,12 +1230,26 @@ class Standalone(BaseWorkflow[TWorkflowInput], Generic[TWorkflowInput, R]):
 
         return_type = get_type_hints(self._task.fn).get("return")
 
-        self._output_validator = (
-            return_type if is_basemodel_subclass(return_type) else None
-        )
+        self._output_validator = self.get_output_validator(return_type)
 
         self.config = self._workflow.config
 
+    def get_output_validator(
+        self, return_type: Any | None
+    ) -> type[BaseModel] | type[DataclassInstance] | None:
+        validator = classify_output_validator(return_type)
+
+        if is_basemodel_validator(validator):
+            return validator.validator_type
+
+        if is_dataclass_validator(validator):
+            return validator.validator_type
+
+        if is_no_validator(validator):
+            return None
+
+        raise TypeError(f"Unhandled validator type: {validator}")
+
     @overload
     def _extract_result(self, result: dict[str, Any]) -> R: ...
 
@@ -1223,10 +1264,21 @@ class Standalone(BaseWorkflow[TWorkflowInput], Generic[TWorkflowInput, R]):
 
         output = result.get(self._task.name)
 
-        if not self._output_validator:
+        validator = classify_output_validator(self._output_validator)
+
+        if is_basemodel_validator(validator):
+            return cast(R, validator.validator_type.model_validate(output))
+
+        if is_dataclass_validator(validator):
+            return cast(
+                R,
+                TypeAdapter(validator.validator_type).validate_python(output),
+            )
+
+        if is_no_validator(validator):
             return cast(R, output)
 
-        return cast(R, self._output_validator.model_validate(output))
+        raise TypeError(f"Unhandled validator type: {validator}")
 
     def run(
         self,

hatchet_sdk-1.21.0/hatchet_sdk/utils/typing.py

@@ -0,0 +1,112 @@
+import sys
+from collections.abc import Awaitable, Coroutine, Generator
+from dataclasses import Field as DataclassField
+from dataclasses import dataclass, is_dataclass
+from enum import Enum
+from typing import Any, ClassVar, Literal, Protocol, TypeAlias, TypeGuard, TypeVar
+
+from pydantic import BaseModel, SkipValidation
+
+
+class DataclassInstance(Protocol):
+    __dataclass_fields__: ClassVar[dict[str, DataclassField[Any]]]
+
+
+def is_basemodel_subclass(model: Any) -> TypeGuard[type[BaseModel]]:
+    try:
+        return issubclass(model, BaseModel)
+    except TypeError:
+        return False
+
+
+@dataclass
+class PydanticModelValidator:
+    validator_type: type[BaseModel]
+    kind: Literal["basemodel"] = "basemodel"
+
+
+@dataclass
+class DataclassValidator:
+    validator_type: type[DataclassInstance]
+    kind: Literal["dataclass"] = "dataclass"
+
+
+@dataclass
+class NoValidator:
+    kind: Literal["none"] = "none"
+
+
+OutputValidator = PydanticModelValidator | DataclassValidator | NoValidator
+
+
+def is_basemodel_validator(
+    validator: OutputValidator,
+) -> TypeGuard[PydanticModelValidator]:
+    return validator.kind == "basemodel"
+
+
+def is_dataclass_validator(validator: OutputValidator) -> TypeGuard[DataclassValidator]:
+    return validator.kind == "dataclass"
+
+
+def is_no_validator(validator: OutputValidator) -> TypeGuard[NoValidator]:
+    return validator.kind == "none"
+
+
+def classify_output_validator(return_type: Any | None) -> OutputValidator:
+    if is_basemodel_subclass(return_type):
+        return PydanticModelValidator(validator_type=return_type)
+
+    if is_dataclass(return_type) and isinstance(return_type, type):
+        return DataclassValidator(validator_type=return_type)
+
+    return NoValidator()
+
+
+class TaskIOValidator(BaseModel):
+    workflow_input: SkipValidation[type[BaseModel] | type[DataclassInstance] | None] = (
+        None
+    )
+    step_output: SkipValidation[type[BaseModel] | type[DataclassInstance] | None] = None
+
+
+JSONSerializableMapping = dict[str, Any]
+
+
+_T_co = TypeVar("_T_co", covariant=True)
+
+if sys.version_info >= (3, 12):
+    AwaitableLike: TypeAlias = Awaitable[_T_co]
+    CoroutineLike: TypeAlias = Coroutine[Any, Any, _T_co]
+else:
+    AwaitableLike: TypeAlias = Generator[Any, None, _T_co] | Awaitable[_T_co]
+    CoroutineLike: TypeAlias = Generator[Any, None, _T_co] | Coroutine[Any, Any, _T_co]
+
+STOP_LOOP_TYPE = Literal["STOP_LOOP"]
+STOP_LOOP: STOP_LOOP_TYPE = "STOP_LOOP"  # Sentinel object to stop the loop
+
+
+class LogLevel(str, Enum):
+    DEBUG = "DEBUG"
+    INFO = "INFO"
+    WARN = "WARN"
+    ERROR = "ERROR"
+
+    @classmethod
+    def from_levelname(cls, levelname: str) -> "LogLevel":
+        levelname = levelname.upper()
+
+        if levelname == "DEBUG":
+            return cls.DEBUG
+
+        if levelname == "INFO":
+            return cls.INFO
+
+        if levelname in ["WARNING", "WARN"]:
+            return cls.WARN
+
+        if levelname == "ERROR":
+            return cls.ERROR
+
+        # fall back to INFO
+        return cls.INFO

{hatchet_sdk-1.20.1 → hatchet_sdk-1.21.0}/hatchet_sdk/worker/runner/runner.py

@@ -4,6 +4,7 @@ import functools
 import json
 from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
+from dataclasses import asdict, is_dataclass
 from enum import Enum
 from multiprocessing import Queue
 from textwrap import dedent
@@ -48,6 +49,7 @@ from hatchet_sdk.runnables.contextvars import (
 from hatchet_sdk.runnables.task import Task
 from hatchet_sdk.runnables.types import R, TWorkflowInput
 from hatchet_sdk.utils.serde import remove_null_unicode_character
+from hatchet_sdk.utils.typing import DataclassInstance
 from hatchet_sdk.worker.action_listener_process import ActionEvent
 from hatchet_sdk.worker.runner.utils.capture_logs import (
     AsyncLogSender,
@@ -479,10 +481,12 @@ class Runner:
 
         if isinstance(output, BaseModel):
             output = output.model_dump(mode="json")
+        elif is_dataclass(output):
+            output = asdict(cast(DataclassInstance, output))
 
         if not isinstance(output, dict):
             raise IllegalTaskOutputError(
-                f"Tasks must return either a dictionary or a Pydantic BaseModel which can be serialized to a JSON object. Got object of type {type(output)} instead."
+                f"Tasks must return either a dictionary, a Pydantic BaseModel, or a dataclass which can be serialized to a JSON object. Got object of type {type(output)} instead."
             )
 
         if output is None:

{hatchet_sdk-1.20.1 → hatchet_sdk-1.21.0}/hatchet_sdk/worker/worker.py

@@ -24,7 +24,7 @@ from pydantic import BaseModel
 from hatchet_sdk.client import Client
 from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.contracts.v1.workflows_pb2 import CreateWorkflowVersionRequest
-from hatchet_sdk.exceptions import LoopAlreadyRunningError
+from hatchet_sdk.exceptions import LifespanSetupError, LoopAlreadyRunningError
 from hatchet_sdk.logger import logger
 from hatchet_sdk.runnables.action import Action
 from hatchet_sdk.runnables.contextvars import task_count
@@ -273,12 +273,18 @@ class Worker:
                 "no actions registered, register workflows or actions before starting worker"
             )
 
-        if self.config.healthcheck.enabled:
-            await self._start_health_server()
-
         lifespan_context = None
         if self.lifespan:
-            lifespan_context = await self._setup_lifespan()
+            try:
+                lifespan_context = await self._setup_lifespan()
+            except LifespanSetupError as e:
+                logger.exception("lifespan setup failed")
+                if self.loop:
+                    self.loop.stop()
+                raise e
+
+        if self.config.healthcheck.enabled:
+            await self._start_health_server()
 
         if self.has_any_non_durable:
             self.action_listener_process = self._start_action_listener(is_durable=False)
@@ -328,8 +334,8 @@ class Worker:
 
         self.lifespan_stack = AsyncExitStack()
 
-        lifespan_gen = self.lifespan()
         try:
+            lifespan_gen = self.lifespan()
             context = await anext(lifespan_gen)
             await self.lifespan_stack.enter_async_context(
                 _create_async_context_manager(lifespan_gen)
@@ -337,10 +343,16 @@ class Worker:
             return context
         except StopAsyncIteration:
             return None
+        except Exception as e:
+            raise LifespanSetupError("An error occurred during lifespan setup") from e
 
     async def _cleanup_lifespan(self) -> None:
-        if self.lifespan_stack is not None:
-            await self.lifespan_stack.aclose()
+        try:
+            if self.lifespan_stack is not None:
+                await self.lifespan_stack.aclose()
+        except Exception as e:
+            logger.exception("error during lifespan cleanup")
+            raise LifespanSetupError("An error occurred during lifespan cleanup") from e
 
     def _start_action_listener(
         self, is_durable: bool
@@ -472,7 +484,10 @@ class Worker:
         ):
             self.durable_action_listener_process.kill()
 
-        await self._cleanup_lifespan()
+        try:
+            await self._cleanup_lifespan()
+        except LifespanSetupError:
+            logger.exception("lifespan cleanup failed")
 
         await self._close()
         if self.loop and self.owned_loop: