hatchet-sdk 0.41.0__py3-none-any.whl → 0.42.1__py3-none-any.whl

hatchet_sdk/worker/runner/runner.py

@@ -8,9 +8,10 @@ from concurrent.futures import ThreadPoolExecutor
 from enum import Enum
 from multiprocessing import Queue
 from threading import Thread, current_thread
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, Literal, Type, TypeVar, cast, overload
 
 from opentelemetry.trace import StatusCode
+from pydantic import BaseModel
 
 from hatchet_sdk.client import new_client_raw
 from hatchet_sdk.clients.admin import new_admin
@@ -18,9 +19,9 @@ from hatchet_sdk.clients.dispatcher.action_listener import Action
 from hatchet_sdk.clients.dispatcher.dispatcher import new_dispatcher
 from hatchet_sdk.clients.run_event_listener import new_listener
 from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
-from hatchet_sdk.context import Context
+from hatchet_sdk.context import Context  # type: ignore[attr-defined]
 from hatchet_sdk.context.worker_context import WorkerContext
-from hatchet_sdk.contracts.dispatcher_pb2 import (
+from hatchet_sdk.contracts.dispatcher_pb2 import (  # type: ignore[attr-defined]
     GROUP_KEY_EVENT_TYPE_COMPLETED,
     GROUP_KEY_EVENT_TYPE_FAILED,
     GROUP_KEY_EVENT_TYPE_STARTED,
@@ -32,6 +33,7 @@ from hatchet_sdk.contracts.dispatcher_pb2 import (
 from hatchet_sdk.loader import ClientConfig
 from hatchet_sdk.logger import logger
 from hatchet_sdk.utils.tracing import create_tracer, parse_carrier_from_metadata
+from hatchet_sdk.utils.types import WorkflowValidator
 from hatchet_sdk.v2.callable import DurableContext
 from hatchet_sdk.worker.action_listener_process import ActionEvent
 from hatchet_sdk.worker.runner.utils.capture_logs import copy_context_vars, sr, wr
@@ -48,11 +50,12 @@ class Runner:
     def __init__(
         self,
         name: str,
-        event_queue: Queue,
+        event_queue: "Queue[Any]",
         max_runs: int | None = None,
         handle_kill: bool = True,
         action_registry: dict[str, Callable[..., Any]] = {},
-        config: ClientConfig = {},
+        validator_registry: dict[str, WorkflowValidator] = {},
+        config: ClientConfig = ClientConfig(),
         labels: dict[str, str | int] = {},
     ):
         # We store the config so we can dynamically create clients for the dispatcher client.
@@ -60,9 +63,10 @@ class Runner:
         self.client = new_client_raw(config)
         self.name = self.client.config.namespace + name
         self.max_runs = max_runs
-        self.tasks: Dict[str, asyncio.Task] = {}  # Store run ids and futures
-        self.contexts: Dict[str, Context] = {}  # Store run ids and contexts
+        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: dict[str, Callable[..., Any]] = action_registry
+        self.validator_registry = validator_registry
 
         self.event_queue = event_queue
 
@@ -89,7 +93,7 @@ class Runner:
     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}"
 
-    def run(self, action: Action):
+    def run(self, action: Action) -> None:
         ctx = parse_carrier_from_metadata(action.additional_metadata)
 
         with self.otel_tracer.start_as_current_span(
@@ -122,8 +126,8 @@ class Runner:
                     span.add_event(log)
                     logger.error(log)
 
-    def step_run_callback(self, action: Action):
-        def inner_callback(task: asyncio.Task):
+    def step_run_callback(self, action: Action) -> Callable[[asyncio.Task[Any]], None]:
+        def inner_callback(task: asyncio.Task[Any]) -> None:
             self.cleanup_run_id(action.step_run_id)
 
             errored = False
@@ -164,8 +168,10 @@ class Runner:
 
         return inner_callback
 
-    def group_key_run_callback(self, action: Action):
-        def inner_callback(task: asyncio.Task):
+    def group_key_run_callback(
+        self, action: Action
+    ) -> Callable[[asyncio.Task[Any]], None]:
+        def inner_callback(task: asyncio.Task[Any]) -> None:
             self.cleanup_run_id(action.get_group_key_run_id)
 
             errored = False
@@ -204,7 +210,10 @@ class Runner:
 
         return inner_callback
 
-    def thread_action_func(self, context, action_func, action: Action):
+    ## TODO: Stricter type hinting here
+    def thread_action_func(
+        self, context: Context, action_func: Callable[..., Any], action: Action
+    ) -> Any:
         if action.step_run_id is not None and action.step_run_id != "":
             self.threads[action.step_run_id] = current_thread()
         elif (
@@ -215,10 +224,15 @@ class Runner:
 
         return action_func(context)
 
+    ## TODO: Stricter type hinting here
     # We wrap all actions in an async func
     async def async_wrapped_action_func(
-        self, context: Context, action_func, action: Action, run_id: str
-    ):
+        self,
+        context: Context,
+        action_func: Callable[..., Any],
+        action: Action,
+        run_id: str,
+    ) -> Any:
         wr.set(context.workflow_run_id())
         sr.set(context.step_run_id)
 
@@ -240,9 +254,7 @@ class Runner:
                 )
 
                 loop = asyncio.get_event_loop()
-                res = await loop.run_in_executor(self.thread_pool, pfunc)
-
-                return res
+                return await loop.run_in_executor(self.thread_pool, pfunc)
         except Exception as e:
             logger.error(
                 errorWithTraceback(
@@ -254,7 +266,7 @@ class Runner:
         finally:
             self.cleanup_run_id(run_id)
 
-    def cleanup_run_id(self, run_id: str):
+    def cleanup_run_id(self, run_id: str | None) -> None:
         if run_id in self.tasks:
             del self.tasks[run_id]
 
@@ -267,7 +279,7 @@ class Runner:
     def create_context(
         self, action: Action, action_func: Callable[..., Any] | None
     ) -> Context | DurableContext:
-        if hasattr(action_func, "durable") and action_func.durable:
+        if hasattr(action_func, "durable") and getattr(action_func, "durable"):
             return DurableContext(
                 action,
                 self.dispatcher_client,
@@ -278,6 +290,7 @@ class Runner:
                 self.workflow_run_event_listener,
                 self.worker_context,
                 self.client.config.namespace,
+                validator_registry=self.validator_registry,
             )
 
         return Context(
@@ -290,9 +303,10 @@ class Runner:
             self.workflow_run_event_listener,
             self.worker_context,
             self.client.config.namespace,
+            validator_registry=self.validator_registry,
         )
 
-    async def handle_start_step_run(self, action: Action):
+    async def handle_start_step_run(self, action: Action) -> None:
         with self.otel_tracer.start_as_current_span(
             f"hatchet.worker.handle_start_step_run.{action.step_id}",
         ) as span:
@@ -336,7 +350,7 @@ class Runner:
 
             span.add_event("Finished step run")
 
-    async def handle_start_group_key_run(self, action: Action):
+    async def handle_start_group_key_run(self, action: Action) -> None:
         with self.otel_tracer.start_as_current_span(
             f"hatchet.worker.handle_start_step_run.{action.step_id}"
         ) as span:
@@ -353,6 +367,7 @@ class Runner:
                 self.worker_context,
                 self.client.config.namespace,
             )
+
             self.contexts[action.get_group_key_run_id] = context
 
             # Find the corresponding action function from the registry
@@ -387,18 +402,18 @@ class Runner:
 
             span.add_event("Finished group key run")
 
-    def force_kill_thread(self, thread):
+    def force_kill_thread(self, thread: Thread) -> None:
         """Terminate a python threading.Thread."""
         try:
             if not thread.is_alive():
                 return
 
-            logger.info(f"Forcefully terminating thread {thread.ident}")
+            ident = cast(int, thread.ident)
+
+            logger.info(f"Forcefully terminating thread {ident}")
 
             exc = ctypes.py_object(SystemExit)
-            res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
-                ctypes.c_long(thread.ident), exc
-            )
+            res = ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(ident), exc)
             if res == 0:
                 raise ValueError("Invalid thread ID")
             elif res != 1:
@@ -408,7 +423,7 @@ class Runner:
                 ctypes.pythonapi.PyThreadState_SetAsyncExc(thread.ident, 0)
                 raise SystemError("PyThreadState_SetAsyncExc failed")
 
-            logger.info(f"Successfully terminated thread {thread.ident}")
+            logger.info(f"Successfully terminated thread {ident}")
 
             # Immediately add a new thread to the thread pool, because we've actually killed a worker
             # in the ThreadPoolExecutor
@@ -416,7 +431,7 @@ class Runner:
         except Exception as e:
             logger.exception(f"Failed to terminate thread: {e}")
 
-    async def handle_cancel_action(self, run_id: str):
+    async def handle_cancel_action(self, run_id: str) -> None:
         with self.otel_tracer.start_as_current_span(
             "hatchet.worker.handle_cancel_action"
         ) as span:
@@ -427,7 +442,9 @@ class Runner:
                 # call cancel to signal the context to stop
                 if run_id in self.contexts:
                     context = self.contexts.get(run_id)
-                    context.cancel()
+
+                    if context:
+                        context.cancel()
 
                 await asyncio.sleep(1)
 
@@ -449,16 +466,20 @@ class Runner:
                 span.add_event(f"Finished cancelling run id: {run_id}")
 
     def serialize_output(self, output: Any) -> str:
-        output_bytes = ""
+
+        if isinstance(output, BaseModel):
+            return output.model_dump_json()
+
         if output is not None:
             try:
-                output_bytes = json.dumps(output)
+                return json.dumps(output)
             except Exception as e:
                 logger.error(f"Could not serialize output: {e}")
-                output_bytes = str(output)
-        return output_bytes
+                return str(output)
+
+        return ""
 
-    async def wait_for_tasks(self):
+    async def wait_for_tasks(self) -> None:
         running = len(self.tasks.keys())
         while running > 0:
             logger.info(f"waiting for {running} tasks to finish...")
@@ -466,6 +487,6 @@ class Runner:
             running = len(self.tasks.keys())
 
 
-def errorWithTraceback(message: str, e: Exception):
+def errorWithTraceback(message: str, e: Exception) -> str:
     trace = "".join(traceback.format_exception(type(e), e, e.__traceback__))
     return f"{message}\n{trace}"

hatchet_sdk/worker/worker.py

@@ -1,22 +1,32 @@
 import asyncio
 import multiprocessing
+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 Process, Queue
-from typing import Any, Callable, Dict, Optional
+from multiprocessing import Queue
+from multiprocessing.process import BaseProcess
+from types import FrameType
+from typing import Any, Callable, TypeVar, get_type_hints
 
+from hatchet_sdk import Context
 from hatchet_sdk.client import Client, new_client_raw
-from hatchet_sdk.context import Context
-from hatchet_sdk.contracts.workflows_pb2 import CreateWorkflowVersionOpts
+from hatchet_sdk.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
+    CreateWorkflowVersionOpts,
+)
 from hatchet_sdk.loader import ClientConfig
 from hatchet_sdk.logger import logger
+from hatchet_sdk.utils.types import WorkflowValidator
 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 WorkflowMeta
+from hatchet_sdk.workflow import WorkflowInterface
+
+T = TypeVar("T")
 
 
 class WorkerStatus(Enum):
@@ -28,46 +38,60 @@ class WorkerStatus(Enum):
 
 @dataclass
 class WorkerStartOptions:
-    loop: asyncio.AbstractEventLoop = field(default=None)
+    loop: asyncio.AbstractEventLoop | None = field(default=None)
 
 
-@dataclass
 class Worker:
-    name: str
-    config: ClientConfig = field(default_factory=dict)
-    max_runs: Optional[int] = None
-    debug: bool = False
-    labels: dict[str, str | int] = field(default_factory=dict)
-    handle_kill: bool = True
-
-    client: Client = field(init=False)
-    tasks: Dict[str, asyncio.Task] = field(default_factory=dict)
-    contexts: Dict[str, Context] = field(default_factory=dict)
-    action_registry: Dict[str, Callable[..., Any]] = field(default_factory=dict)
-    killing: bool = field(init=False, default=False)
-    _status: WorkerStatus = field(init=False, default=WorkerStatus.INITIALIZED)
-
-    action_listener_process: Process = field(init=False, default=None)
-    action_listener_health_check: asyncio.Task = field(init=False, default=None)
-    action_runner: WorkerActionRunLoopManager = field(init=False, default=None)
-    ctx = multiprocessing.get_context("spawn")
-
-    action_queue: Queue = field(init=False, default_factory=ctx.Queue)
-    event_queue: Queue = field(init=False, default_factory=ctx.Queue)
-
-    loop: asyncio.AbstractEventLoop = field(init=False, default=None)
-    owned_loop: bool = True
-
-    def __post_init__(self):
+    def __init__(
+        self,
+        name: str,
+        config: ClientConfig = ClientConfig(),
+        max_runs: int | None = None,
+        labels: dict[str, str | int] = {},
+        debug: bool = False,
+        owned_loop: bool = True,
+        handle_kill: bool = True,
+    ) -> None:
+        self.name = name
+        self.config = config
+        self.max_runs = max_runs
+        self.debug = debug
+        self.labels = labels
+        self.handle_kill = handle_kill
+        self.owned_loop = owned_loop
+
+        self.client: Client
+
+        self.action_registry: dict[str, Callable[[Context], T]] = {}
+        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.ctx = multiprocessing.get_context("spawn")
+
+        self.action_queue: "Queue[Any]" = self.ctx.Queue()
+        self.event_queue: "Queue[Any]" = self.ctx.Queue()
+
+        self.loop: asyncio.AbstractEventLoop
+
         self.client = new_client_raw(self.config, self.debug)
         self.name = self.client.config.namespace + self.name
-        if self.owned_loop:
-            self._setup_signal_handlers()
 
-    def register_function(self, action: str, func: HatchetCallable):
+        self._setup_signal_handlers()
+
+    def register_function(
+        self, action: str, func: HatchetCallable[Any] | ConcurrencyFunction
+    ) -> None:
         self.action_registry[action] = func
 
-    def register_workflow_from_opts(self, name: str, opts: CreateWorkflowVersionOpts):
+    def register_workflow_from_opts(
+        self, name: str, opts: CreateWorkflowVersionOpts
+    ) -> None:
         try:
             self.client.admin.put_workflow(opts.name, opts)
         except Exception as e:
@@ -75,7 +99,7 @@ class Worker:
             logger.error(e)
             sys.exit(1)
 
-    def register_workflow(self, workflow: WorkflowMeta):
+    def register_workflow(self, workflow: WorkflowInterface) -> None:
         namespace = self.client.config.namespace
 
         try:
@@ -87,24 +111,30 @@ class Worker:
             logger.error(e)
             sys.exit(1)
 
-        def create_action_function(action_func):
-            def action_function(context):
+        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):
-                action_function.is_coroutine = True
+                setattr(action_function, "is_coroutine", True)
             else:
-                action_function.is_coroutine = False
+                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")
+            self.validator_registry[action_name] = WorkflowValidator(
+                workflow_input=workflow.input_validator, step_output=return_type
+            )
 
     def status(self) -> WorkerStatus:
         return self._status
 
-    def setup_loop(self, loop: asyncio.AbstractEventLoop = None):
+    def setup_loop(self, loop: asyncio.AbstractEventLoop | None = None) -> bool:
         try:
             loop = loop or asyncio.get_running_loop()
             self.loop = loop
@@ -118,17 +148,22 @@ class Worker:
             created_loop = True
             return created_loop
 
-    def start(self, options: WorkerStartOptions = WorkerStartOptions()):
+    def start(
+        self, options: WorkerStartOptions = WorkerStartOptions()
+    ) -> Future[asyncio.Task[Any] | None]:
         self.owned_loop = self.setup_loop(options.loop)
+
         f = asyncio.run_coroutine_threadsafe(
             self.async_start(options, _from_start=True), self.loop
         )
+
         # start the loop and wait until its closed
         if self.owned_loop:
             self.loop.run_forever()
 
             if self.handle_kill:
                 sys.exit(0)
+
         return f
 
     ## Start methods
@@ -136,7 +171,7 @@ class Worker:
         self,
         options: WorkerStartOptions = WorkerStartOptions(),
         _from_start: bool = False,
-    ):
+    ) -> Any | None:
         main_pid = os.getpid()
         logger.info("------------------------------------------")
         logger.info("STARTING HATCHET...")
@@ -148,25 +183,28 @@ class Worker:
             logger.error(
                 "no actions registered, register workflows or actions before starting worker"
             )
-            return
+            return None
 
         # non blocking setup
         if not _from_start:
             self.setup_loop(options.loop)
 
         self.action_listener_process = self._start_listener()
+
         self.action_runner = self._run_action_runner()
+
         self.action_listener_health_check = self.loop.create_task(
             self._check_listener_health()
         )
 
         return await self.action_listener_health_check
 
-    def _run_action_runner(self):
+    def _run_action_runner(self) -> WorkerActionRunLoopManager:
         # Retrieve the shared queue
-        runner = WorkerActionRunLoopManager(
+        return WorkerActionRunLoopManager(
             self.name,
             self.action_registry,
+            self.validator_registry,
             self.max_runs,
             self.config,
             self.action_queue,
@@ -177,10 +215,9 @@ class Worker:
             self.labels,
         )
 
-        return runner
-
-    def _start_listener(self):
+    def _start_listener(self) -> multiprocessing.context.SpawnProcess:
         action_list = [str(key) for key in self.action_registry.keys()]
+
         try:
             process = self.ctx.Process(
                 target=worker_action_listener_process,
@@ -204,7 +241,7 @@ class Worker:
             logger.error(f"failed to start action listener: {e}")
             sys.exit(1)
 
-    async def _check_listener_health(self):
+    async def _check_listener_health(self) -> None:
         logger.debug("starting action listener health check...")
         try:
             while not self.killing:
@@ -224,21 +261,21 @@ class Worker:
             logger.error(f"error checking listener health: {e}")
 
     ## Cleanup methods
-    def _setup_signal_handlers(self):
+    def _setup_signal_handlers(self) -> None:
         signal.signal(signal.SIGTERM, self._handle_exit_signal)
         signal.signal(signal.SIGINT, self._handle_exit_signal)
         signal.signal(signal.SIGQUIT, self._handle_force_quit_signal)
 
-    def _handle_exit_signal(self, signum, frame):
+    def _handle_exit_signal(self, signum: int, frame: FrameType | None) -> None:
         sig_name = "SIGTERM" if signum == signal.SIGTERM else "SIGINT"
         logger.info(f"received signal {sig_name}...")
         self.loop.create_task(self.exit_gracefully())
 
-    def _handle_force_quit_signal(self, signum, frame):
+    def _handle_force_quit_signal(self, signum: int, frame: FrameType | None) -> None:
         logger.info("received SIGQUIT...")
         self.exit_forcefully()
 
-    async def close(self):
+    async def close(self) -> None:
         logger.info(f"closing worker '{self.name}'...")
         self.killing = True
         # self.action_queue.close()
@@ -249,7 +286,7 @@ class Worker:
 
         await self.action_listener_health_check
 
-    async def exit_gracefully(self):
+    async def exit_gracefully(self) -> None:
         logger.debug(f"gracefully stopping worker: {self.name}")
 
         if self.killing:
@@ -270,7 +307,7 @@ class Worker:
 
         logger.info("👋")
 
-    def exit_forcefully(self):
+    def exit_forcefully(self) -> None:
         self.killing = True
 
         logger.debug(f"forcefully stopping worker: {self.name}")
@@ -286,7 +323,7 @@ class Worker:
         )  # Exit immediately TODO - should we exit with 1 here, there may be other workers to cleanup
 
 
-def register_on_worker(callable: HatchetCallable, worker: Worker):
+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: