hatchet-sdk 1.3.2__py3-none-any.whl → 1.4.1__py3-none-any.whl

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/context/context.py

@@ -37,6 +37,7 @@ class Context:
         durable_event_listener: DurableEventListener | None,
         worker: WorkerContext,
         runs_client: RunsClient,
+        lifespan_context: Any | None,
     ):
         self.worker = worker
 
@@ -59,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)
 
@@ -116,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/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/waits.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from datetime import datetime
 from enum import Enum
 from typing import TYPE_CHECKING
 from uuid import uuid4
@@ -58,10 +59,13 @@ class Condition(ABC):
 
 
 class SleepCondition(Condition):
-    def __init__(self, duration: Duration) -> None:
+    def __init__(
+        self, duration: Duration, readable_data_key: str | None = None
+    ) -> None:
         super().__init__(
             BaseCondition(
-                readable_data_key=f"sleep:{timedelta_to_expr(duration)}",
+                readable_data_key=readable_data_key
+                or f"sleep:{timedelta_to_expr(duration)}",
             )
         )
 
@@ -75,10 +79,15 @@ class SleepCondition(Condition):
 
 
 class UserEventCondition(Condition):
-    def __init__(self, event_key: str, expression: str | None = None) -> None:
+    def __init__(
+        self,
+        event_key: str,
+        expression: str | None = None,
+        readable_data_key: str | None = None,
+    ) -> None:
         super().__init__(
             BaseCondition(
-                readable_data_key=event_key,
+                readable_data_key=readable_data_key or event_key,
                 expression=expression,
             )
         )
@@ -95,10 +104,22 @@ class UserEventCondition(Condition):
 
 class ParentCondition(Condition):
     def __init__(
-        self, parent: "Task[TWorkflowInput, R]", expression: str | None = None
+        self,
+        parent: "Task[TWorkflowInput, R]",
+        expression: str | None = None,
+        readable_data_key: str | None = None,
     ) -> None:
         super().__init__(
-            BaseCondition(readable_data_key=parent.name, expression=expression)
+            BaseCondition(
+                readable_data_key=readable_data_key
+                or (
+                    parent.name
+                    + (f":{expression}" if expression else "")
+                    + ":"
+                    + datetime.now().isoformat()
+                ),
+                expression=expression,
+            )
         )
 
         self.parent = parent

hatchet_sdk/worker/runner/run_loop_manager.py

@@ -23,7 +23,7 @@ class WorkerActionRunLoopManager:
         self,
         name: str,
         action_registry: dict[str, Task[Any, Any]],
-        slots: int | None,
+        slots: int,
         config: ClientConfig,
         action_queue: "Queue[Action | STOP_LOOP_TYPE]",
         event_queue: "Queue[ActionEvent]",
@@ -31,6 +31,7 @@ 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
@@ -42,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)
@@ -86,6 +88,7 @@ class WorkerActionRunLoopManager:
             self.handle_kill,
             self.action_registry,
             self.labels,
+            self.lifespan_context,
         )
 
         logger.debug(f"'{self.name}' waiting for {list(self.action_registry.keys())}")

hatchet_sdk/worker/runner/runner.py

@@ -59,10 +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]] = {},
         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
@@ -102,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}"
 
@@ -303,6 +306,7 @@ class Runner:
             durable_event_listener=self.durable_event_listener,
             worker=self.worker_context,
             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/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
+from typing import Any, AsyncGenerator, Callable, TypeVar
 from warnings import warn
 
 from aiohttp import web
@@ -63,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
@@ -119,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:
@@ -130,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
 
@@ -213,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",
@@ -252,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(
@@ -269,13 +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,
-                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,
@@ -283,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:
@@ -300,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,
@@ -394,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.2.dist-info → hatchet_sdk-1.4.1.dist-info}/METADATA

@@ -1,10 +1,12 @@
 Metadata-Version: 2.1
 Name: hatchet-sdk
-Version: 1.3.2
+Version: 1.4.1
 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.2.dist-info → hatchet_sdk-1.4.1.dist-info}/RECORD

@@ -1,7 +1,7 @@
 hatchet_sdk/__init__.py,sha256=LUj6VyGVSHCYTQTaoyiVhjyJLOfv6gMCmb-s4hRyISM,10031
 hatchet_sdk/client.py,sha256=tbOeMuaJmgpyYSQg8QUz_J4AdqRNvV9E0aEZpgsiZTE,2207
 hatchet_sdk/clients/admin.py,sha256=qY9-nB8o-jCpQQF65nooemH8HOioXjntp9ethpGke9o,17266
-hatchet_sdk/clients/dispatcher/action_listener.py,sha256=sAboL2Dr59MAPBR3KcJ7sSBfOwnTR3rdsRTJ0bDRuuc,16279
+hatchet_sdk/clients/dispatcher/action_listener.py,sha256=yJhHqnp_RAhUfOJUpBCLIpqMyXaLX-CSUsm0TgPxGEs,16273
 hatchet_sdk/clients/dispatcher/dispatcher.py,sha256=IL-hDXG8Lzas9FieVuNr47E_3Gvpc-aL4Xu_l385Vp8,8140
 hatchet_sdk/clients/event_ts.py,sha256=OjYc_y1K3fNNZ7QLRJYBDnfgsyUVERkxofYldO4QVVs,1672
 hatchet_sdk/clients/events.py,sha256=sKqzGwkrW_AKWW0Q_rxfwo9jXV5EUpLx7fTQvCNHaVo,5443
@@ -221,7 +221,7 @@ hatchet_sdk/clients/v1/api_client.py,sha256=mJQUZ3cOxlFJiwWKK5F8jBxcpNZ7A2292Huc
 hatchet_sdk/config.py,sha256=jJA76BOvVdfOQHy6TKclAvr2qyblcM-Pz5J-hVAdpQ4,3588
 hatchet_sdk/connection.py,sha256=B5gT5NL9BBB5-l9U_cN6pMlraQk880rEYMnqaK_dgL0,2590
 hatchet_sdk/context/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-hatchet_sdk/context/context.py,sha256=povQ0iO-QOi_lzntSWUiZPGbtpNTq-1m80b3zjIX74A,9163
+hatchet_sdk/context/context.py,sha256=DWaHE3AuXFvRKPxYaKW5-CUY8R9lz-lF4t3sgucwZPs,9336
 hatchet_sdk/context/worker_context.py,sha256=OVcEWvdT_Kpd0nlg61VAPUgIPSFzSLs0aSrXWj-1GX4,974
 hatchet_sdk/contracts/dispatcher_pb2.py,sha256=SN4CIKeQwYkrbfRGhdhZo2uBn4nGzjUWIC1vvXdDeOQ,14503
 hatchet_sdk/contracts/dispatcher_pb2.pyi,sha256=ZSGio5eYxkw-QuQx2C5ASTNcKzeMQn5JTnWaRiThafM,18455
@@ -250,7 +250,7 @@ hatchet_sdk/features/runs.py,sha256=2069-4YUaouYloFUxa3-WmtkWhCjEIGMeGrN_nlnYCA,
 hatchet_sdk/features/scheduled.py,sha256=1_GMP7jSC9UCFFWLVES5HtTGQInQd9cJeqJDhq0Hnmg,8813
 hatchet_sdk/features/workers.py,sha256=fG8zoRXMm-lNNdNfeovKKpmJ0RzjZd6CIq_dVsu-ih0,1515
 hatchet_sdk/features/workflows.py,sha256=x1Z_HZM_ZQZbWZhOfzJaEiPfRNSn6YX3IaTmFZ5D7MY,2069
-hatchet_sdk/hatchet.py,sha256=7wDYLGBlPdQ0IvnJUdAe8Ms2B9IZWpedRED4AflZCW0,22982
+hatchet_sdk/hatchet.py,sha256=oHE27SRgYddR1KiJyQwGAa0eV_uEH73w4_jbXxFIUqU,23146
 hatchet_sdk/labels.py,sha256=nATgxWE3lFxRTnfISEpoIRLGbMfAZsHF4lZTuG4Mfic,182
 hatchet_sdk/logger.py,sha256=5uOr52T4mImSQm1QvWT8HvZFK5WfPNh3Y1cBQZRFgUQ,333
 hatchet_sdk/metadata.py,sha256=XkRbhnghJJGCdVvF-uzyGBcNaTqpeQ3uiQvNNP1wyBc,107
@@ -498,15 +498,15 @@ hatchet_sdk/v0/worker/runner/utils/error_with_traceback.py,sha256=Iih_s8JNqrinXE
 hatchet_sdk/v0/worker/worker.py,sha256=0yU0z-0si7NzG0U9et9J0tiwfVBSHl4QSiOW-WNmTQM,13027
 hatchet_sdk/v0/workflow.py,sha256=d4o425efk7J3JgLIge34MW_A3pzwnwSRtwEOgIqM2pc,9387
 hatchet_sdk/v0/workflow_run.py,sha256=jsEZprXshrSV7i_TtL5uoCL03D18zQ3NeJCq7mp97Dg,1752
-hatchet_sdk/waits.py,sha256=dTrelK8Lk0W5anq3TlpJwAK5z8TJ0FW7XgoSg8bf8fA,3351
+hatchet_sdk/waits.py,sha256=L2xZUcmrQX-pTVXWv1W8suMoYU_eA0uowpollauQmOM,3893
 hatchet_sdk/worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hatchet_sdk/worker/action_listener_process.py,sha256=KxS7-wBpfKnsq0LNSvk-MG442Lh60iQMy3VpD1FW3mU,11703
-hatchet_sdk/worker/runner/run_loop_manager.py,sha256=QrhhAMnQmyThDMgohZwFw-0AVPHvhT4seK2FVCzvC7c,3690
-hatchet_sdk/worker/runner/runner.py,sha256=t8U4QK2GYNjgxGHFmD9RW12e3vungu_bvBSBBdIb_cs,17219
+hatchet_sdk/worker/runner/run_loop_manager.py,sha256=RNWKDCjR57nJ0LCoLUMi0_3pnmpqyo80mz_RaxHYGIc,3812
+hatchet_sdk/worker/runner/runner.py,sha256=J-gTpoYF9WphB1zZhpPJryCv10VjU-MiCdJgMYH-nds,17352
 hatchet_sdk/worker/runner/utils/capture_logs.py,sha256=nHRPSiDBqzhObM7i2X7t03OupVFnE7kQBdR2Ckgg-2w,2709
-hatchet_sdk/worker/worker.py,sha256=5DfmNNjGuYtn-n9Qce1B2XCLKjnIXzNwF8VLW-wtI5Y,14161
+hatchet_sdk/worker/worker.py,sha256=m11u3QPyAqXJpF6Y3lTUEVqTKdIigGKTXTymNPqq6hw,16149
 hatchet_sdk/workflow_run.py,sha256=ZwH0HLFGFVXz6jbiqSv4w0Om2XuR52Tzzw6LH4y65jQ,2765
-hatchet_sdk-1.3.2.dist-info/METADATA,sha256=tWsWj0cUquv41Ih8O3mIM3Pwk6jl3mp0rIBWExejSI0,3571
-hatchet_sdk-1.3.2.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-hatchet_sdk-1.3.2.dist-info/entry_points.txt,sha256=5mTp_AsCWK5raiVxP_MU9eBCgkRGl4OsN6chpHcvm7o,1235
-hatchet_sdk-1.3.2.dist-info/RECORD,,
+hatchet_sdk-1.4.1.dist-info/METADATA,sha256=vACzxEl9nCc_VuqANxq7gsEgrDHrEdOT_2Xsa6YIha8,3635
+hatchet_sdk-1.4.1.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+hatchet_sdk-1.4.1.dist-info/entry_points.txt,sha256=uYc-ivF7fUATjrWnwHP32ojgMCmogIs6yFOU8lEA4mk,1276
+hatchet_sdk-1.4.1.dist-info/RECORD,,

{hatchet_sdk-1.3.2.dist-info → hatchet_sdk-1.4.1.dist-info}/entry_points.txt

@@ -14,6 +14,7 @@ events=examples.events.worker:main
 existing_loop=examples.worker_existing_loop.worker:main
 fanout=examples.fanout.worker:main
 fanout_sync=examples.fanout_sync.worker:main
+lifespans=examples.lifespans.worker:main
 logger=examples.logger.worker:main
 manual_trigger=examples.manual_trigger.worker:main
 on_failure=examples.on_failure.worker:main