hatchet-sdk 1.3.1__tar.gz → 1.4.0__tar.gz

{hatchet_sdk-1.3.1 → hatchet_sdk-1.4.0}/PKG-INFO

@@ -1,10 +1,12 @@
 Metadata-Version: 2.1
 Name: hatchet-sdk
-Version: 1.3.1
+Version: 1.4.0
 Summary: 
+License: MIT
 Author: Alexander Belanger
 Author-email: alexander@hatchet.run
 Requires-Python: >=3.10,<4.0
+Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11

{hatchet_sdk-1.3.1 → hatchet_sdk-1.4.0}/hatchet_sdk/clients/dispatcher/action_listener.py

@@ -42,7 +42,7 @@ class GetActionListenerRequest(BaseModel):
     worker_name: str
     services: list[str]
     actions: list[str]
-    slots: int = 100
+    slots: int
     raw_labels: dict[str, str | int] = Field(default_factory=dict)
 
     labels: dict[str, WorkerLabels] = Field(default_factory=dict)

{hatchet_sdk-1.3.1 → hatchet_sdk-1.4.0}/hatchet_sdk/context/context.py

@@ -19,7 +19,7 @@ from hatchet_sdk.context.worker_context import WorkerContext
 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, WorkflowValidator
+from hatchet_sdk.utils.typing import JSONSerializableMapping
 from hatchet_sdk.waits import SleepCondition, UserEventCondition
 
 if TYPE_CHECKING:
@@ -37,10 +37,9 @@ class Context:
         durable_event_listener: DurableEventListener | None,
         worker: WorkerContext,
         runs_client: RunsClient,
-        validator_registry: dict[str, WorkflowValidator] = {},
+        lifespan_context: Any | None,
     ):
         self.worker = worker
-        self.validator_registry = validator_registry
 
         self.data = action.action_payload
 
@@ -61,6 +60,8 @@ class Context:
 
         self.input = self.data.input
 
+        self._lifespan_context = lifespan_context
+
     def was_skipped(self, task: "Task[TWorkflowInput, R]") -> bool:
         return self.data.parents.get(task.name, {}).get("skipped", False)
 
@@ -74,27 +75,12 @@ class Context:
         if self.was_skipped(task):
             raise ValueError("{task.name} was skipped")
 
-        action_prefix = self.action.action_id.split(":")[0]
-
-        workflow_validator = next(
-            (
-                v
-                for k, v in self.validator_registry.items()
-                if k == f"{action_prefix}:{task.name}"
-            ),
-            None,
-        )
-
         try:
             parent_step_data = cast(R, self.data.parents[task.name])
         except KeyError:
             raise ValueError(f"Step output for '{task.name}' not found")
 
-        if (
-            parent_step_data
-            and workflow_validator
-            and (v := workflow_validator.step_output)
-        ):
+        if parent_step_data and (v := task.validators.step_output):
             return cast(R, v.model_validate(parent_step_data))
 
         return parent_step_data
@@ -133,6 +119,10 @@ class Context:
     def workflow_input(self) -> JSONSerializableMapping:
         return self.input
 
+    @property
+    def lifespan(self) -> Any:
+        return self._lifespan_context
+
     @property
     def workflow_run_id(self) -> str:
         return self.action.workflow_run_id

{hatchet_sdk-1.3.1 → hatchet_sdk-1.4.0}/hatchet_sdk/hatchet.py

@@ -33,7 +33,7 @@ from hatchet_sdk.runnables.types import (
 )
 from hatchet_sdk.runnables.workflow import BaseWorkflow, Workflow
 from hatchet_sdk.utils.timedelta_to_expression import Duration
-from hatchet_sdk.worker.worker import Worker
+from hatchet_sdk.worker.worker import LifespanFn, Worker
 
 
 class Hatchet:
@@ -134,8 +134,10 @@ class Hatchet:
         self,
         name: str,
         slots: int = 100,
+        durable_slots: int = 1_000,
         labels: dict[str, str | int] = {},
         workflows: list[BaseWorkflow[Any]] = [],
+        lifespan: LifespanFn | None = None,
     ) -> Worker:
         """
         Create a Hatchet worker on which to run workflows.
@@ -165,11 +167,13 @@ class Hatchet:
         return Worker(
             name=name,
             slots=slots,
+            durable_slots=durable_slots,
             labels=labels,
             config=self._client.config,
             debug=self._client.debug,
             owned_loop=loop is None,
             workflows=workflows,
+            lifespan=lifespan,
         )
 
     @overload

{hatchet_sdk-1.3.1 → hatchet_sdk-1.4.0}/hatchet_sdk/runnables/task.py

@@ -7,6 +7,7 @@ from typing import (
     TypeVar,
     Union,
     cast,
+    get_type_hints,
 )
 
 from hatchet_sdk.context.context import Context, DurableContext
@@ -28,6 +29,7 @@ from hatchet_sdk.runnables.types import (
     is_sync_fn,
 )
 from hatchet_sdk.utils.timedelta_to_expression import Duration, timedelta_to_expr
+from hatchet_sdk.utils.typing import TaskIOValidator, is_basemodel_subclass
 from hatchet_sdk.waits import (
     Action,
     Condition,
@@ -106,6 +108,13 @@ class Task(Generic[TWorkflowInput, R]):
         self.skip_if = self._flatten_conditions(skip_if)
         self.cancel_if = self._flatten_conditions(cancel_if)
 
+        return_type = get_type_hints(_fn).get("return")
+
+        self.validators: TaskIOValidator = TaskIOValidator(
+            workflow_input=workflow.config.input_validator,
+            step_output=return_type if is_basemodel_subclass(return_type) else None,
+        )
+
     def _flatten_conditions(
         self, conditions: list[Condition | OrGroup]
     ) -> list[Condition]:

{hatchet_sdk-1.3.1 → hatchet_sdk-1.4.0}/hatchet_sdk/utils/typing.py

@@ -1,9 +1,7 @@
-from typing import Any, Mapping, Type, TypeGuard, TypeVar
+from typing import Any, Mapping, Type, TypeGuard
 
 from pydantic import BaseModel
 
-T = TypeVar("T", bound=BaseModel)
-
 
 def is_basemodel_subclass(model: Any) -> TypeGuard[Type[BaseModel]]:
     try:
@@ -12,7 +10,7 @@ def is_basemodel_subclass(model: Any) -> TypeGuard[Type[BaseModel]]:
         return False
 
 
-class WorkflowValidator(BaseModel):
+class TaskIOValidator(BaseModel):
     workflow_input: Type[BaseModel] | None = None
     step_output: Type[BaseModel] | None = None
 

{hatchet_sdk-1.3.1 → hatchet_sdk-1.4.0}/hatchet_sdk/worker/runner/run_loop_manager.py

@@ -8,7 +8,6 @@ from hatchet_sdk.clients.dispatcher.action_listener import Action
 from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.logger import logger
 from hatchet_sdk.runnables.task import Task
-from hatchet_sdk.utils.typing import WorkflowValidator
 from hatchet_sdk.worker.action_listener_process import ActionEvent
 from hatchet_sdk.worker.runner.runner import Runner
 from hatchet_sdk.worker.runner.utils.capture_logs import capture_logs
@@ -24,8 +23,7 @@ class WorkerActionRunLoopManager:
         self,
         name: str,
         action_registry: dict[str, Task[Any, Any]],
-        validator_registry: dict[str, WorkflowValidator],
-        slots: int | None,
+        slots: int,
         config: ClientConfig,
         action_queue: "Queue[Action | STOP_LOOP_TYPE]",
         event_queue: "Queue[ActionEvent]",
@@ -33,10 +31,10 @@ class WorkerActionRunLoopManager:
         handle_kill: bool = True,
         debug: bool = False,
         labels: dict[str, str | int] = {},
+        lifespan_context: Any | None = None,
     ) -> None:
         self.name = name
         self.action_registry = action_registry
-        self.validator_registry = validator_registry
         self.slots = slots
         self.config = config
         self.action_queue = action_queue
@@ -45,6 +43,7 @@ class WorkerActionRunLoopManager:
         self.handle_kill = handle_kill
         self.debug = debug
         self.labels = labels
+        self.lifespan_context = lifespan_context
 
         if self.debug:
             logger.setLevel(logging.DEBUG)
@@ -88,8 +87,8 @@ class WorkerActionRunLoopManager:
             self.slots,
             self.handle_kill,
             self.action_registry,
-            self.validator_registry,
             self.labels,
+            self.lifespan_context,
         )
 
         logger.debug(f"'{self.name}' waiting for {list(self.action_registry.keys())}")

{hatchet_sdk-1.3.1 → hatchet_sdk-1.4.0}/hatchet_sdk/worker/runner/runner.py

@@ -43,7 +43,6 @@ 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.typing import WorkflowValidator
 from hatchet_sdk.worker.action_listener_process import ActionEvent
 from hatchet_sdk.worker.runner.utils.capture_logs import copy_context_vars
 
@@ -60,11 +59,11 @@ class Runner:
         self,
         event_queue: "Queue[ActionEvent]",
         config: ClientConfig,
-        slots: int | None = None,
+        slots: int,
         handle_kill: bool = True,
         action_registry: dict[str, Task[TWorkflowInput, R]] = {},
-        validator_registry: dict[str, WorkflowValidator] = {},
         labels: dict[str, str | int] = {},
+        lifespan_context: Any | None = None,
     ):
         # We store the config so we can dynamically create clients for the dispatcher client.
         self.config = config
@@ -73,7 +72,6 @@ class Runner:
         self.tasks: dict[str, asyncio.Task[Any]] = {}  # Store run ids and futures
         self.contexts: dict[str, Context] = {}  # Store run ids and contexts
         self.action_registry = action_registry
-        self.validator_registry = validator_registry
 
         self.event_queue = event_queue
 
@@ -105,6 +103,8 @@ class Runner:
             labels=labels, client=Client(config=config).dispatcher
         )
 
+        self.lifespan_context = lifespan_context
+
     def create_workflow_run_url(self, action: Action) -> str:
         return f"{self.config.server_url}/workflow-runs/{action.workflow_run_id}?tenant={action.tenant_id}"
 
@@ -305,8 +305,8 @@ class Runner:
             event_client=self.event_client,
             durable_event_listener=self.durable_event_listener,
             worker=self.worker_context,
-            validator_registry=self.validator_registry,
             runs_client=self.runs_client,
+            lifespan_context=self.lifespan_context,
         )
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor

{hatchet_sdk-1.3.1 → hatchet_sdk-1.4.0}/hatchet_sdk/worker/worker.py

@@ -5,12 +5,13 @@ import os
 import re
 import signal
 import sys
+from contextlib import AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass, field
 from enum import Enum
 from multiprocessing import Queue
 from multiprocessing.process import BaseProcess
 from types import FrameType
-from typing import Any, TypeVar, get_type_hints
+from typing import Any, AsyncGenerator, Callable, TypeVar
 from warnings import warn
 
 from aiohttp import web
@@ -26,7 +27,6 @@ from hatchet_sdk.contracts.v1.workflows_pb2 import CreateWorkflowVersionRequest
 from hatchet_sdk.logger import logger
 from hatchet_sdk.runnables.task import Task
 from hatchet_sdk.runnables.workflow import BaseWorkflow
-from hatchet_sdk.utils.typing import WorkflowValidator, is_basemodel_subclass
 from hatchet_sdk.worker.action_listener_process import (
     ActionEvent,
     worker_action_listener_process,
@@ -64,21 +64,41 @@ class HealthCheckResponse(BaseModel):
     python_version: str
 
 
+LifespanGenerator = AsyncGenerator[Any, Any]
+LifespanFn = Callable[[], LifespanGenerator]
+
+
+@asynccontextmanager
+async def _create_async_context_manager(
+    gen: LifespanGenerator,
+) -> AsyncGenerator[None, None]:
+    try:
+        yield
+    finally:
+        try:
+            await anext(gen)
+        except StopAsyncIteration:
+            pass
+
+
 class Worker:
     def __init__(
         self,
         name: str,
         config: ClientConfig,
-        slots: int | None = None,
+        slots: int,
+        durable_slots: int,
         labels: dict[str, str | int] = {},
         debug: bool = False,
         owned_loop: bool = True,
         handle_kill: bool = True,
         workflows: list[BaseWorkflow[Any]] = [],
+        lifespan: LifespanFn | None = None,
     ) -> None:
         self.config = config
         self.name = self.config.namespace + name
         self.slots = slots
+        self.durable_slots = durable_slots
         self.debug = debug
         self.labels = labels
         self.handle_kill = handle_kill
@@ -87,8 +107,6 @@ class Worker:
         self.action_registry: dict[str, Task[Any, Any]] = {}
         self.durable_action_registry: dict[str, Task[Any, Any]] = {}
 
-        self.validator_registry: dict[str, WorkflowValidator] = {}
-
         self.killing: bool = False
         self._status: WorkerStatus
 
@@ -122,6 +140,9 @@ class Worker:
         self.has_any_durable = False
         self.has_any_non_durable = False
 
+        self.lifespan = lifespan
+        self.lifespan_stack: AsyncExitStack | None = None
+
         self.register_workflows(workflows)
 
     def register_workflow_from_opts(self, opts: CreateWorkflowVersionRequest) -> None:
@@ -133,6 +154,11 @@ class Worker:
             sys.exit(1)
 
     def register_workflow(self, workflow: BaseWorkflow[Any]) -> None:
+        if not workflow.tasks:
+            raise ValueError(
+                "workflow must have at least one task registered before registering"
+            )
+
         opts = workflow.to_proto()
         name = workflow.name
 
@@ -153,13 +179,6 @@ class Worker:
                 self.has_any_non_durable = True
                 self.action_registry[action_name] = step
 
-            return_type = get_type_hints(step.fn).get("return")
-
-            self.validator_registry[action_name] = WorkflowValidator(
-                workflow_input=workflow.config.input_validator,
-                step_output=return_type if is_basemodel_subclass(return_type) else None,
-            )
-
     def register_workflows(self, workflows: list[BaseWorkflow[Any]]) -> None:
         for workflow in workflows:
             self.register_workflow(workflow)
@@ -223,6 +242,11 @@ class Worker:
         logger.info(f"healthcheck server running on port {port}")
 
     def start(self, options: WorkerStartOptions = WorkerStartOptions()) -> None:
+        if not (self.action_registry or self.durable_action_registry):
+            raise ValueError(
+                "no actions registered, register workflows before starting worker"
+            )
+
         if options.loop is not None:
             warn(
                 "Passing a custom event loop is deprecated and will be removed in the future. This option no longer has any effect",
@@ -262,15 +286,23 @@ class Worker:
         if self.config.healthcheck.enabled:
             await self._start_health_server()
 
+        lifespan_context = None
+        if self.lifespan:
+            lifespan_context = await self._setup_lifespan()
+
         if self.has_any_non_durable:
             self.action_listener_process = self._start_action_listener(is_durable=False)
-            self.action_runner = self._run_action_runner(is_durable=False)
+            self.action_runner = self._run_action_runner(
+                is_durable=False, lifespan_context=lifespan_context
+            )
 
         if self.has_any_durable:
             self.durable_action_listener_process = self._start_action_listener(
                 is_durable=True
             )
-            self.durable_action_runner = self._run_action_runner(is_durable=True)
+            self.durable_action_runner = self._run_action_runner(
+                is_durable=True, lifespan_context=lifespan_context
+            )
 
         if self.loop:
             self.action_listener_health_check = self.loop.create_task(
@@ -279,14 +311,15 @@ class Worker:
 
             await self.action_listener_health_check
 
-    def _run_action_runner(self, is_durable: bool) -> WorkerActionRunLoopManager:
+    def _run_action_runner(
+        self, is_durable: bool, lifespan_context: Any | None
+    ) -> WorkerActionRunLoopManager:
         # Retrieve the shared queue
         if self.loop:
             return WorkerActionRunLoopManager(
                 self.name + ("_durable" if is_durable else ""),
                 self.durable_action_registry if is_durable else self.action_registry,
-                self.validator_registry,
-                1_000 if is_durable else self.slots,
+                self.durable_slots if is_durable else self.slots,
                 self.config,
                 self.durable_action_queue if is_durable else self.action_queue,
                 self.durable_event_queue if is_durable else self.event_queue,
@@ -294,10 +327,31 @@ class Worker:
                 self.handle_kill,
                 self.client.debug,
                 self.labels,
+                lifespan_context,
             )
 
         raise RuntimeError("event loop not set, cannot start action runner")
 
+    async def _setup_lifespan(self) -> Any:
+        if self.lifespan is None:
+            return None
+
+        self.lifespan_stack = AsyncExitStack()
+
+        lifespan_gen = self.lifespan()
+        try:
+            context = await anext(lifespan_gen)
+            await self.lifespan_stack.enter_async_context(
+                _create_async_context_manager(lifespan_gen)
+            )
+            return context
+        except StopAsyncIteration:
+            return None
+
+    async def _cleanup_lifespan(self) -> None:
+        if self.lifespan_stack is not None:
+            await self.lifespan_stack.aclose()
+
     def _start_action_listener(
         self, is_durable: bool
     ) -> multiprocessing.context.SpawnProcess:
@@ -311,7 +365,7 @@ class Worker:
                         if is_durable
                         else list(self.action_registry.keys())
                     ),
-                    1_000 if is_durable else self.slots,
+                    self.durable_slots if is_durable else self.slots,
                     self.config,
                     self.durable_action_queue if is_durable else self.action_queue,
                     self.durable_event_queue if is_durable else self.event_queue,
@@ -405,6 +459,8 @@ class Worker:
         ):
             self.durable_action_listener_process.kill()
 
+        await self._cleanup_lifespan()
+
         await self._close()
         if self.loop and self.owned_loop:
             self.loop.stop()

{hatchet_sdk-1.3.1 → hatchet_sdk-1.4.0}/pyproject.toml

@@ -1,10 +1,11 @@
 [tool.poetry]
 name = "hatchet-sdk"
-version = "1.3.1"
+version = "1.4.0"
 description = ""
 authors = ["Alexander Belanger <alexander@hatchet.run>"]
 readme = "README.md"
 include = ["hatchet_sdk/py.typed"]
+license = "MIT"
 
 [tool.poetry.dependencies]
 python = "^3.10"
@@ -50,6 +51,7 @@ pytest = "^8.3.5"
 pytest-asyncio = "^0.25.3"
 pytest-env = "^1.1.5"
 pytest-retry = "^1.7.0"
+psycopg = { extras = ["pool"], version = "^3.2.6" }
 
 [tool.poetry.extras]
 otel = [
@@ -120,6 +122,7 @@ exclude = [
     "hatchet_sdk/clients/rest/rest.py",
     "hatchet_sdk/v0/*",
 ]
+target-version = "py310"
 
 [tool.poetry.scripts]
 
@@ -151,3 +154,4 @@ otel = "examples.opentelemetry_instrumentation.worker:main"
 waits = "examples.waits.worker:main"
 durable = "examples.durable.worker:main"
 streaming = "examples.streaming.worker:main"
+lifespans = "examples.lifespans.worker:main"