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

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

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

@@ -1,10 +1,11 @@
 [tool.poetry]
 name = "hatchet-sdk"
-version = "1.3.2"
+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"