hatchet-sdk 0.47.1__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

hatchet_sdk/worker/worker.py

@@ -2,33 +2,38 @@ import asyncio
 import multiprocessing
 import multiprocessing.context
 import os
+import re
 import signal
 import sys
-from concurrent.futures import Future
 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, Callable, TypeVar, get_type_hints
+from typing import Any, TypeVar, get_type_hints
 
 from aiohttp import web
 from aiohttp.web_request import Request
 from aiohttp.web_response import Response
-from prometheus_client import CONTENT_TYPE_LATEST, Gauge, generate_latest
+from prometheus_client import Gauge, generate_latest
+from pydantic import BaseModel
 
-from hatchet_sdk import Context
-from hatchet_sdk.client import Client, new_client_raw
-from hatchet_sdk.contracts.workflows_pb2 import CreateWorkflowVersionOpts
-from hatchet_sdk.loader import ClientConfig
+from hatchet_sdk.client import Client
+from hatchet_sdk.clients.dispatcher.action_listener import Action
+from hatchet_sdk.config import ClientConfig
+from hatchet_sdk.contracts.v1.workflows_pb2 import CreateWorkflowVersionRequest
 from hatchet_sdk.logger import logger
-from hatchet_sdk.utils.types import WorkflowValidator
-from hatchet_sdk.utils.typing import is_basemodel_subclass
-from hatchet_sdk.v2.callable import HatchetCallable
-from hatchet_sdk.v2.concurrency import ConcurrencyFunction
-from hatchet_sdk.worker.action_listener_process import worker_action_listener_process
-from hatchet_sdk.worker.runner.run_loop_manager import WorkerActionRunLoopManager
-from hatchet_sdk.workflow import WorkflowInterface
+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,
+)
+from hatchet_sdk.worker.runner.run_loop_manager import (
+    STOP_LOOP_TYPE,
+    WorkerActionRunLoopManager,
+)
 
 T = TypeVar("T")
 
@@ -45,62 +50,76 @@ class WorkerStartOptions:
     loop: asyncio.AbstractEventLoop | None = field(default=None)
 
 
-TWorkflow = TypeVar("TWorkflow", bound=object)
+class HealthCheckResponse(BaseModel):
+    status: str
+    name: str
+    slots: int
+    actions: list[str]
+    labels: dict[str, str | int]
+    python_version: str
 
 
 class Worker:
     def __init__(
         self,
         name: str,
-        config: ClientConfig = ClientConfig(),
-        max_runs: int | None = None,
+        config: ClientConfig,
+        slots: int | None = None,
         labels: dict[str, str | int] = {},
         debug: bool = False,
         owned_loop: bool = True,
         handle_kill: bool = True,
+        workflows: list[BaseWorkflow[Any]] = [],
     ) -> None:
-        self.name = name
         self.config = config
-        self.max_runs = max_runs
+        self.name = self.config.namespace + name
+        self.slots = slots
         self.debug = debug
         self.labels = labels
         self.handle_kill = handle_kill
         self.owned_loop = owned_loop
 
-        self.client: Client
+        self.action_registry: dict[str, Task[Any, Any]] = {}
+        self.durable_action_registry: dict[str, Task[Any, Any]] = {}
 
-        self.action_registry: dict[str, Callable[[Context], Any]] = {}
         self.validator_registry: dict[str, WorkflowValidator] = {}
 
         self.killing: bool = False
         self._status: WorkerStatus
 
-        self.action_listener_process: BaseProcess
-        self.action_listener_health_check: asyncio.Task[Any]
-        self.action_runner: WorkerActionRunLoopManager
+        self.action_listener_process: BaseProcess | None = None
+        self.durable_action_listener_process: BaseProcess | None = None
+
+        self.action_listener_health_check: asyncio.Task[None]
+
+        self.action_runner: WorkerActionRunLoopManager | None = None
+        self.durable_action_runner: WorkerActionRunLoopManager | None = None
 
         self.ctx = multiprocessing.get_context("spawn")
 
-        self.action_queue: "Queue[Any]" = self.ctx.Queue()
-        self.event_queue: "Queue[Any]" = self.ctx.Queue()
+        self.action_queue: "Queue[Action | STOP_LOOP_TYPE]" = self.ctx.Queue()
+        self.event_queue: "Queue[ActionEvent]" = self.ctx.Queue()
+
+        self.durable_action_queue: "Queue[Action | STOP_LOOP_TYPE]" = self.ctx.Queue()
+        self.durable_event_queue: "Queue[ActionEvent]" = self.ctx.Queue()
 
         self.loop: asyncio.AbstractEventLoop
 
-        self.client = new_client_raw(self.config, self.debug)
-        self.name = self.client.config.namespace + self.name
+        self.client = Client(config=self.config, debug=self.debug)
 
         self._setup_signal_handlers()
 
         self.worker_status_gauge = Gauge(
-            "hatchet_worker_status", "Current status of the Hatchet worker"
+            "hatchet_worker_status_" + re.sub(r"\W+", "", name),
+            "Current status of the Hatchet worker",
         )
 
-    def register_function(self, action: str, func: Callable[[Context], Any]) -> None:
-        self.action_registry[action] = func
+        self.has_any_durable = False
+        self.has_any_non_durable = False
 
-    def register_workflow_from_opts(
-        self, name: str, opts: CreateWorkflowVersionOpts
-    ) -> None:
+        self.register_workflows(workflows)
+
+    def register_workflow_from_opts(self, opts: CreateWorkflowVersionRequest) -> None:
         try:
             self.client.admin.put_workflow(opts.name, opts)
         except Exception as e:
@@ -108,78 +127,87 @@ class Worker:
             logger.error(e)
             sys.exit(1)
 
-    def register_workflow(self, workflow: TWorkflow) -> None:
-        ## Hack for typing
-        assert isinstance(workflow, WorkflowInterface)
-
+    def register_workflow(self, workflow: BaseWorkflow[Any]) -> None:
         namespace = self.client.config.namespace
 
+        opts = workflow._get_create_opts(namespace)
+        name = workflow._get_name(namespace)
+
         try:
-            self.client.admin.put_workflow(
-                workflow.get_name(namespace), workflow.get_create_opts(namespace)
-            )
+            self.client.admin.put_workflow(name, opts)
         except Exception as e:
-            logger.error(f"failed to register workflow: {workflow.get_name(namespace)}")
+            logger.error(
+                f"failed to register workflow: {workflow._get_name(namespace)}"
+            )
             logger.error(e)
             sys.exit(1)
 
-        def create_action_function(
-            action_func: Callable[..., T]
-        ) -> Callable[[Context], T]:
-            def action_function(context: Context) -> T:
-                return action_func(workflow, context)
+        for step in workflow.tasks:
+            action_name = workflow._create_action_name(namespace, step)
 
-            if asyncio.iscoroutinefunction(action_func):
-                setattr(action_function, "is_coroutine", True)
+            if workflow.is_durable:
+                self.has_any_durable = True
+                self.durable_action_registry[action_name] = step
             else:
-                setattr(action_function, "is_coroutine", False)
+                self.has_any_non_durable = True
+                self.action_registry[action_name] = step
 
-            return action_function
-
-        for action_name, action_func in workflow.get_actions(namespace):
-            self.action_registry[action_name] = create_action_function(action_func)
-            return_type = get_type_hints(action_func).get("return")
+            return_type = get_type_hints(step.fn).get("return")
 
             self.validator_registry[action_name] = WorkflowValidator(
-                workflow_input=workflow.input_validator,
+                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)
+
+    @property
     def status(self) -> WorkerStatus:
         return self._status
 
-    def setup_loop(self, loop: asyncio.AbstractEventLoop | None = None) -> bool:
+    def _setup_loop(self, loop: asyncio.AbstractEventLoop | None = None) -> bool:
         try:
-            loop = loop or asyncio.get_running_loop()
-            self.loop = loop
-            created_loop = False
+            self.loop = loop or asyncio.get_running_loop()
             logger.debug("using existing event loop")
-            return created_loop
+
+            created_loop = False
         except RuntimeError:
             self.loop = asyncio.new_event_loop()
+
             logger.debug("creating new event loop")
-            asyncio.set_event_loop(self.loop)
             created_loop = True
-            return created_loop
 
-    async def health_check_handler(self, request: Request) -> Response:
-        status = self.status()
+        asyncio.set_event_loop(self.loop)
+
+        return created_loop
+
+    async def _health_check_handler(self, request: Request) -> Response:
+        response = HealthCheckResponse(
+            status=self.status.name,
+            name=self.name,
+            slots=self.slots or 0,
+            actions=list(self.action_registry.keys()),
+            labels=self.labels,
+            python_version=sys.version,
+        ).model_dump()
 
-        return web.json_response({"status": status.name})
+        return web.json_response(response)
 
-    async def metrics_handler(self, request: Request) -> Response:
-        self.worker_status_gauge.set(1 if self.status() == WorkerStatus.HEALTHY else 0)
+    async def _metrics_handler(self, request: Request) -> Response:
+        self.worker_status_gauge.set(1 if self.status == WorkerStatus.HEALTHY else 0)
 
         return web.Response(body=generate_latest(), content_type="text/plain")
 
-    async def start_health_server(self) -> None:
-        port = self.config.worker_healthcheck_port or 8001
+    async def _start_health_server(self) -> None:
+        port = self.config.healthcheck.port
 
         app = web.Application()
         app.add_routes(
             [
-                web.get("/health", self.health_check_handler),
-                web.get("/metrics", self.metrics_handler),
+                web.get("/health", self._health_check_handler),
+                web.get("/metrics", self._metrics_handler),
             ]
         )
 
@@ -195,14 +223,10 @@ class Worker:
 
         logger.info(f"healthcheck server running on port {port}")
 
-    def start(
-        self, options: WorkerStartOptions = WorkerStartOptions()
-    ) -> Future[asyncio.Task[Any] | None]:
-        self.owned_loop = self.setup_loop(options.loop)
+    def start(self, options: WorkerStartOptions = WorkerStartOptions()) -> None:
+        self.owned_loop = self._setup_loop(options.loop)
 
-        f = asyncio.run_coroutine_threadsafe(
-            self.async_start(options, _from_start=True), self.loop
-        )
+        asyncio.run_coroutine_threadsafe(self._aio_start(), self.loop)
 
         # start the loop and wait until its closed
         if self.owned_loop:
@@ -211,73 +235,75 @@ class Worker:
             if self.handle_kill:
                 sys.exit(0)
 
-        return f
-
-    ## Start methods
-    async def async_start(
-        self,
-        options: WorkerStartOptions = WorkerStartOptions(),
-        _from_start: bool = False,
-    ) -> Any | None:
+    async def _aio_start(self) -> None:
         main_pid = os.getpid()
+
         logger.info("------------------------------------------")
         logger.info("STARTING HATCHET...")
         logger.debug(f"worker runtime starting on PID: {main_pid}")
 
         self._status = WorkerStatus.STARTING
 
-        if len(self.action_registry.keys()) == 0:
-            logger.error(
+        if (
+            len(self.action_registry.keys()) == 0
+            and len(self.durable_action_registry.keys()) == 0
+        ):
+            raise ValueError(
                 "no actions registered, register workflows or actions before starting worker"
             )
-            return None
 
-        # non blocking setup
-        if not _from_start:
-            self.setup_loop(options.loop)
+        if self.config.healthcheck.enabled:
+            await self._start_health_server()
 
-        if self.config.worker_healthcheck_enabled:
-            await self.start_health_server()
+        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_listener_process = self._start_listener()
-
-        self.action_runner = self._run_action_runner()
+        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.action_listener_health_check = self.loop.create_task(
             self._check_listener_health()
         )
 
-        return await self.action_listener_health_check
+        await self.action_listener_health_check
 
-    def _run_action_runner(self) -> WorkerActionRunLoopManager:
+    def _run_action_runner(self, is_durable: bool) -> WorkerActionRunLoopManager:
         # Retrieve the shared queue
         return WorkerActionRunLoopManager(
-            self.name,
-            self.action_registry,
+            self.name + ("_durable" if is_durable else ""),
+            self.durable_action_registry if is_durable else self.action_registry,
             self.validator_registry,
-            self.max_runs,
+            1_000 if is_durable else self.slots,
             self.config,
-            self.action_queue,
-            self.event_queue,
+            self.durable_action_queue if is_durable else self.action_queue,
+            self.durable_event_queue if is_durable else self.event_queue,
             self.loop,
             self.handle_kill,
             self.client.debug,
             self.labels,
         )
 
-    def _start_listener(self) -> multiprocessing.context.SpawnProcess:
-        action_list = [str(key) for key in self.action_registry.keys()]
-
+    def _start_action_listener(
+        self, is_durable: bool
+    ) -> multiprocessing.context.SpawnProcess:
         try:
             process = self.ctx.Process(
                 target=worker_action_listener_process,
                 args=(
-                    self.name,
-                    action_list,
-                    self.max_runs,
+                    self.name + ("_durable" if is_durable else ""),
+                    (
+                        list(self.durable_action_registry.keys())
+                        if is_durable
+                        else list(self.action_registry.keys())
+                    ),
+                    1_000 if is_durable else self.slots,
                     self.config,
-                    self.action_queue,
-                    self.event_queue,
+                    self.durable_action_queue if is_durable else self.action_queue,
+                    self.durable_event_queue if is_durable else self.event_queue,
                     self.handle_kill,
                     self.client.debug,
                     self.labels,
@@ -296,8 +322,13 @@ class Worker:
         try:
             while not self.killing:
                 if (
-                    self.action_listener_process is None
-                    or not self.action_listener_process.is_alive()
+                    not self.action_listener_process
+                    and not self.durable_action_listener_process
+                ) or (
+                    self.action_listener_process
+                    and self.durable_action_listener_process
+                    and not self.action_listener_process.is_alive()
+                    and not self.durable_action_listener_process.is_alive()
                 ):
                     logger.debug("child action listener process killed...")
                     self._status = WorkerStatus.UNHEALTHY
@@ -310,7 +341,6 @@ class Worker:
         except Exception as e:
             logger.error(f"error checking listener health: {e}")
 
-    ## Cleanup methods
     def _setup_signal_handlers(self) -> None:
         signal.signal(signal.SIGTERM, self._handle_exit_signal)
         signal.signal(signal.SIGINT, self._handle_exit_signal)
@@ -323,70 +353,63 @@ class Worker:
 
     def _handle_force_quit_signal(self, signum: int, frame: FrameType | None) -> None:
         logger.info("received SIGQUIT...")
-        self.exit_forcefully()
+        self.loop.create_task(self._exit_forcefully())
 
-    async def close(self) -> None:
+    async def _close(self) -> None:
         logger.info(f"closing worker '{self.name}'...")
         self.killing = True
-        # self.action_queue.close()
-        # self.event_queue.close()
 
         if self.action_runner is not None:
             self.action_runner.cleanup()
 
+        if self.durable_action_runner is not None:
+            self.durable_action_runner.cleanup()
+
         await self.action_listener_health_check
 
     async def exit_gracefully(self) -> None:
         logger.debug(f"gracefully stopping worker: {self.name}")
 
         if self.killing:
-            return self.exit_forcefully()
+            return await self._exit_forcefully()
 
         self.killing = True
 
-        await self.action_runner.wait_for_tasks()
+        if self.action_runner:
+            await self.action_runner.wait_for_tasks()
+            await self.action_runner.exit_gracefully()
 
-        await self.action_runner.exit_gracefully()
+        if self.durable_action_runner:
+            await self.durable_action_runner.wait_for_tasks()
+            await self.durable_action_runner.exit_gracefully()
 
         if self.action_listener_process and self.action_listener_process.is_alive():
             self.action_listener_process.kill()
 
-        await self.close()
+        if (
+            self.durable_action_listener_process
+            and self.durable_action_listener_process.is_alive()
+        ):
+            self.durable_action_listener_process.kill()
+
+        await self._close()
         if self.loop and self.owned_loop:
             self.loop.stop()
 
         logger.info("👋")
 
-    def exit_forcefully(self) -> None:
+    async def _exit_forcefully(self) -> None:
         self.killing = True
 
         logger.debug(f"forcefully stopping worker: {self.name}")
 
-        self.close()
+        await self._close()
 
         if self.action_listener_process:
-            self.action_listener_process.kill()  # Forcefully kill the process
-
-        logger.info("👋")
-        sys.exit(
-            1
-        )  # Exit immediately TODO - should we exit with 1 here, there may be other workers to cleanup
-
-
-def register_on_worker(callable: HatchetCallable[T], worker: Worker) -> None:
-    worker.register_function(callable.get_action_name(), callable)
-
-    if callable.function_on_failure is not None:
-        worker.register_function(
-            callable.function_on_failure.get_action_name(), callable.function_on_failure
-        )
-
-    if callable.function_concurrency is not None:
-        worker.register_function(
-            callable.function_concurrency.get_action_name(),
-            callable.function_concurrency,
-        )
+            self.action_listener_process.kill()
 
-    opts = callable.to_workflow_opts()
+        if self.durable_action_listener_process:
+            self.durable_action_listener_process.kill()
 
-    worker.register_workflow_from_opts(opts.name, opts)
+        logger.info("👋")
+        sys.exit(1)

hatchet_sdk/workflow_run.py

@@ -1,17 +1,15 @@
 import asyncio
-from typing import Any, Coroutine, Generic, Optional, TypedDict, TypeVar
+from typing import Any
 
 from hatchet_sdk.clients.run_event_listener import (
     RunEventListener,
     RunEventListenerClient,
 )
 from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
-from hatchet_sdk.utils.aio_utils import EventLoopThread, get_active_event_loop
+from hatchet_sdk.utils.aio_utils import get_active_event_loop
 
 
 class WorkflowRunRef:
-    workflow_run_id: str
-
     def __init__(
         self,
         workflow_run_id: str,
@@ -22,17 +20,18 @@ class WorkflowRunRef:
         self.workflow_listener = workflow_listener
         self.workflow_run_event_listener = workflow_run_event_listener
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.workflow_run_id
 
     def stream(self) -> RunEventListener:
         return self.workflow_run_event_listener.stream(self.workflow_run_id)
 
-    def result(self) -> Coroutine:
-        return self.workflow_listener.result(self.workflow_run_id)
+    async def aio_result(self) -> dict[str, Any]:
+        return await self.workflow_listener.result(self.workflow_run_id)
 
-    def sync_result(self) -> dict:
+    def result(self) -> dict[str, Any]:
         coro = self.workflow_listener.result(self.workflow_run_id)
+
         loop = get_active_event_loop()
 
         if loop is None:
@@ -44,16 +43,3 @@ class WorkflowRunRef:
                 asyncio.set_event_loop(None)
         else:
             return loop.run_until_complete(coro)
-
-
-T = TypeVar("T")
-
-
-class RunRef(WorkflowRunRef, Generic[T]):
-    async def result(self) -> T:
-        res = await self.workflow_listener.result(self.workflow_run_id)
-
-        if len(res) == 1:
-            return list(res.values())[0]
-
-        return res

{hatchet_sdk-0.47.1.dist-info → hatchet_sdk-1.0.0a1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hatchet-sdk
-Version: 0.47.1
+Version: 1.0.0a1
 Summary: 
 Author: Alexander Belanger
 Author-email: alexander@hatchet.run
@@ -28,8 +28,8 @@ Requires-Dist: opentelemetry-sdk (>=1.28.0,<2.0.0) ; extra == "otel"
 Requires-Dist: prometheus-client (>=0.21.1,<0.22.0)
 Requires-Dist: protobuf (>=5.29.1,<6.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)
-Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)
 Requires-Dist: pyyaml (>=6.0.1,<7.0.0)
 Requires-Dist: tenacity (>=8.4.1)
 Requires-Dist: urllib3 (>=1.26.20)