hatchet-sdk 0.46.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

hatchet_sdk/worker/worker.py

@@ -4,33 +4,36 @@ import multiprocessing.context
 import os
 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 hatchet_sdk import Context
 from hatchet_sdk.client import Client, new_client_raw
+from hatchet_sdk.clients.dispatcher.action_listener import Action
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.contracts.workflows_pb2 import CreateWorkflowVersionOpts
-from hatchet_sdk.loader import ClientConfig
 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.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,
+)
+from hatchet_sdk.workflow import BaseWorkflow, Step, StepType, Task
 
 T = TypeVar("T")
+TBaseWorkflow = TypeVar("TBaseWorkflow", bound=BaseWorkflow)
 
 
 class WorkerStatus(Enum):
@@ -45,9 +48,6 @@ class WorkerStartOptions:
     loop: asyncio.AbstractEventLoop | None = field(default=None)
 
 
-TWorkflow = TypeVar("TWorkflow", bound=object)
-
-
 class Worker:
     def __init__(
         self,
@@ -69,20 +69,20 @@ class Worker:
 
         self.client: Client
 
-        self.action_registry: dict[str, Callable[[Context], Any]] = {}
+        self.action_registry: dict[str, Step[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_listener_health_check: asyncio.Task[None]
         self.action_runner: WorkerActionRunLoopManager
 
         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.loop: asyncio.AbstractEventLoop
 
@@ -95,9 +95,6 @@ class Worker:
             "hatchet_worker_status", "Current status of the Hatchet worker"
         )
 
-    def register_function(self, action: str, func: Callable[[Context], Any]) -> None:
-        self.action_registry[action] = func
-
     def register_workflow_from_opts(
         self, name: str, opts: CreateWorkflowVersionOpts
     ) -> None:
@@ -108,10 +105,7 @@ 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: TBaseWorkflow) -> None:
         namespace = self.client.config.namespace
 
         try:
@@ -123,28 +117,42 @@ class Worker:
             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)
-
-            if asyncio.iscoroutinefunction(action_func):
-                setattr(action_function, "is_coroutine", True)
-            else:
-                setattr(action_function, "is_coroutine", False)
-
-            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")
+        for step in workflow.steps:
+            action_name = workflow.create_action_name(namespace, step)
+            self.action_registry[action_name] = step
+            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_function(self, function: Task[Any, Any]) -> None:
+        from hatchet_sdk.workflow import BaseWorkflow
+
+        declaration = function.hatchet.declare_workflow(
+            **function.workflow_config.model_dump()
+        )
+
+        class Workflow(BaseWorkflow):
+            config = declaration.config
+
+            @property
+            def default_steps(self) -> list[Step[Any]]:
+                return [function.step]
+
+            @property
+            def on_failure_steps(self) -> list[Step[Any]]:
+                if not function.on_failure_step:
+                    return []
+
+                step = function.on_failure_step.step
+                step.type = StepType.ON_FAILURE
+
+                return [step]
+
+        self.register_workflow(Workflow())
+
     def status(self) -> WorkerStatus:
         return self._status
 
@@ -173,7 +181,7 @@ class Worker:
         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
+        port = self.config.healthcheck.port
 
         app = web.Application()
         app.add_routes(
@@ -195,13 +203,11 @@ class Worker:
 
         logger.info(f"healthcheck server running on port {port}")
 
-    def start(
-        self, options: WorkerStartOptions = WorkerStartOptions()
-    ) -> Future[asyncio.Task[Any] | None]:
+    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(options, _from_start=True), self.loop
         )
 
         # start the loop and wait until its closed
@@ -211,14 +217,12 @@ class Worker:
             if self.handle_kill:
                 sys.exit(0)
 
-        return f
-
     ## Start methods
-    async def async_start(
+    async def aio_start(
         self,
         options: WorkerStartOptions = WorkerStartOptions(),
         _from_start: bool = False,
-    ) -> Any | None:
+    ) -> None:
         main_pid = os.getpid()
         logger.info("------------------------------------------")
         logger.info("STARTING HATCHET...")
@@ -236,7 +240,7 @@ class Worker:
         if not _from_start:
             self.setup_loop(options.loop)
 
-        if self.config.worker_healthcheck_enabled:
+        if self.config.healthcheck.enabled:
             await self.start_health_server()
 
         self.action_listener_process = self._start_listener()
@@ -247,7 +251,7 @@ class Worker:
             self._check_listener_health()
         )
 
-        return await self.action_listener_health_check
+        await self.action_listener_health_check
 
     def _run_action_runner(self) -> WorkerActionRunLoopManager:
         # Retrieve the shared queue
@@ -362,7 +366,8 @@ class Worker:
 
         logger.debug(f"forcefully stopping worker: {self.name}")
 
-        self.close()
+        ## TODO: `self.close` needs to be awaited / used
+        self.close()  # type: ignore[unused-coroutine]
 
         if self.action_listener_process:
             self.action_listener_process.kill()  # Forcefully kill the process
@@ -371,22 +376,3 @@ class Worker:
         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,
-        )
-
-    opts = callable.to_workflow_opts()
-
-    worker.register_workflow_from_opts(opts.name, opts)