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

hatchet_sdk/worker/runner/runner.py

@@ -3,23 +3,24 @@ import contextvars
 import ctypes
 import functools
 import json
-import time
 import traceback
 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, cast
+from typing import Any, Callable, Dict, Literal, cast, overload
 
 from pydantic import BaseModel
 
-from hatchet_sdk.client import new_client_raw
-from hatchet_sdk.clients.admin import new_admin
-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.client import Client
+from hatchet_sdk.clients.admin import AdminClient
+from hatchet_sdk.clients.dispatcher.action_listener import Action, ActionType
+from hatchet_sdk.clients.dispatcher.dispatcher import DispatcherClient
+from hatchet_sdk.clients.durable_event_listener import DurableEventListener
+from hatchet_sdk.clients.run_event_listener import RunEventListenerClient
 from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
-from hatchet_sdk.context import Context  # type: ignore[attr-defined]
+from hatchet_sdk.config import ClientConfig
+from hatchet_sdk.context.context import Context, DurableContext
 from hatchet_sdk.context.worker_context import WorkerContext
 from hatchet_sdk.contracts.dispatcher_pb2 import (
     GROUP_KEY_EVENT_TYPE_COMPLETED,
@@ -28,14 +29,20 @@ from hatchet_sdk.contracts.dispatcher_pb2 import (
     STEP_EVENT_TYPE_COMPLETED,
     STEP_EVENT_TYPE_FAILED,
     STEP_EVENT_TYPE_STARTED,
-    ActionType,
 )
-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 DurableContext
+from hatchet_sdk.runnables.contextvars import (
+    ctx_step_run_id,
+    ctx_worker_id,
+    ctx_workflow_run_id,
+    spawn_index_lock,
+    workflow_spawn_indices,
+)
+from hatchet_sdk.runnables.task import Task
+from hatchet_sdk.runnables.types import R, TWorkflowInput
+from hatchet_sdk.utils.typing import WorkflowValidator
 from hatchet_sdk.worker.action_listener_process import ActionEvent
-from hatchet_sdk.worker.runner.utils.capture_logs import copy_context_vars, sr, wr
+from hatchet_sdk.worker.runner.utils.capture_logs import copy_context_vars
 
 
 class WorkerStatus(Enum):
@@ -49,28 +56,28 @@ class Runner:
     def __init__(
         self,
         name: str,
-        event_queue: "Queue[Any]",
-        max_runs: int | None = None,
+        event_queue: "Queue[ActionEvent]",
+        config: ClientConfig,
+        slots: int | None = None,
         handle_kill: bool = True,
-        action_registry: dict[str, Callable[..., Any]] = {},
+        action_registry: dict[str, Task[TWorkflowInput, R]] = {},
         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.
         self.config = config
-        self.client = new_client_raw(config)
+        self.client = Client(config)
         self.name = self.client.config.namespace + name
-        self.max_runs = max_runs
+        self.slots = slots
         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.action_registry = action_registry
         self.validator_registry = validator_registry
 
         self.event_queue = event_queue
 
         # The thread pool is used for synchronous functions which need to run concurrently
-        self.thread_pool = ThreadPoolExecutor(max_workers=max_runs)
+        self.thread_pool = ThreadPoolExecutor(max_workers=slots)
         self.threads: Dict[str, Thread] = {}  # Store run ids and threads
 
         self.killing = False
@@ -78,13 +85,14 @@ class Runner:
 
         # We need to initialize a new admin and dispatcher client *after* we've started the event loop,
         # otherwise the grpc.aio methods will use a different event loop and we'll get a bunch of errors.
-        self.dispatcher_client = new_dispatcher(self.config)
-        self.admin_client = new_admin(self.config)
-        self.workflow_run_event_listener = new_listener(self.config)
+        self.dispatcher_client = DispatcherClient(self.config)
+        self.admin_client = AdminClient(self.config)
+        self.workflow_run_event_listener = RunEventListenerClient(self.config)
         self.client.workflow_listener = PooledWorkflowRunListener(self.config)
+        self.durable_event_listener = DurableEventListener(self.config)
 
         self.worker_context = WorkerContext(
-            labels=labels, client=new_client_raw(config).dispatcher
+            labels=labels, client=Client(config=config).dispatcher
         )
 
     def create_workflow_run_url(self, action: Action) -> str:
@@ -130,7 +138,7 @@ class Runner:
                     ActionEvent(
                         action=action,
                         type=STEP_EVENT_TYPE_FAILED,
-                        payload=str(errorWithTraceback(f"{e}", e)),
+                        payload=str(pretty_format_exception(f"{e}", e)),
                     )
                 )
 
@@ -172,7 +180,7 @@ class Runner:
                     ActionEvent(
                         action=action,
                         type=GROUP_KEY_EVENT_TYPE_FAILED,
-                        payload=str(errorWithTraceback(f"{e}", e)),
+                        payload=str(pretty_format_exception(f"{e}", e)),
                     )
                 )
 
@@ -195,10 +203,9 @@ class Runner:
 
         return inner_callback
 
-    ## TODO: Stricter type hinting here
     def thread_action_func(
-        self, context: Context, action_func: Callable[..., Any], action: Action
-    ) -> Any:
+        self, ctx: Context, task: Task[TWorkflowInput, R], action: Action
+    ) -> R:
         if action.step_run_id is not None and action.step_run_id != "":
             self.threads[action.step_run_id] = current_thread()
         elif (
@@ -207,25 +214,23 @@ class Runner:
         ):
             self.threads[action.get_group_key_run_id] = current_thread()
 
-        return action_func(context)
+        return task.call(ctx)
 
-    ## TODO: Stricter type hinting here
     # We wrap all actions in an async func
     async def async_wrapped_action_func(
         self,
-        context: Context,
-        action_func: Callable[..., Any],
+        ctx: Context,
+        task: Task[TWorkflowInput, R],
         action: Action,
         run_id: str,
-    ) -> Any:
-        wr.set(context.workflow_run_id())
-        sr.set(context.step_run_id)
+    ) -> R:
+        ctx_step_run_id.set(action.step_run_id)
+        ctx_workflow_run_id.set(action.workflow_run_id)
+        ctx_worker_id.set(action.worker_id)
 
         try:
-            if (
-                hasattr(action_func, "is_coroutine") and action_func.is_coroutine
-            ) or asyncio.iscoroutinefunction(action_func):
-                return await action_func(context)
+            if task.is_async_function:
+                return await task.aio_call(ctx)
             else:
                 pfunc = functools.partial(
                     # we must copy the context vars to the new thread, as only asyncio natively supports
@@ -233,8 +238,8 @@ class Runner:
                     copy_context_vars,
                     contextvars.copy_context().items(),
                     self.thread_action_func,
-                    context,
-                    action_func,
+                    ctx,
+                    task,
                     action,
                 )
 
@@ -242,7 +247,7 @@ class Runner:
                 return await loop.run_in_executor(self.thread_pool, pfunc)
         except Exception as e:
             logger.error(
-                errorWithTraceback(
+                pretty_format_exception(
                     f"exception raised in action ({action.action_id}, retry={action.retry_count}):\n{e}",
                     e,
                 )
@@ -261,30 +266,29 @@ class Runner:
         if run_id in self.contexts:
             del self.contexts[run_id]
 
+    @overload
+    def create_context(
+        self, action: Action, is_durable: Literal[True] = True
+    ) -> DurableContext: ...
+
+    @overload
     def create_context(
-        self, action: Action, action_func: Callable[..., Any] | None
+        self, action: Action, is_durable: Literal[False] = False
+    ) -> Context: ...
+
+    def create_context(
+        self, action: Action, is_durable: bool = True
     ) -> Context | DurableContext:
-        if hasattr(action_func, "durable") and getattr(action_func, "durable"):
-            return DurableContext(
-                action,
-                self.dispatcher_client,
-                self.admin_client,
-                self.client.event,
-                self.client.rest,
-                self.client.workflow_listener,
-                self.workflow_run_event_listener,
-                self.worker_context,
-                self.client.config.namespace,
-                validator_registry=self.validator_registry,
-            )
+        constructor = DurableContext if is_durable else Context
 
-        return Context(
+        return constructor(
             action,
             self.dispatcher_client,
             self.admin_client,
             self.client.event,
             self.client.rest,
             self.client.workflow_listener,
+            self.durable_event_listener,
             self.workflow_run_event_listener,
             self.worker_context,
             self.client.config.namespace,
@@ -292,22 +296,20 @@ class Runner:
         )
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
-    async def handle_start_step_run(self, action: Action) -> None | Exception:
+    async def handle_start_step_run(self, action: Action) -> None:
         action_name = action.action_id
 
         # Find the corresponding action function from the registry
         action_func = self.action_registry.get(action_name)
 
-        context = self.create_context(action, action_func)
-
-        self.contexts[action.step_run_id] = context
-
         if action_func:
+            context = self.create_context(
+                action, True if action_func.is_durable else False
+            )
+
+            self.contexts[action.step_run_id] = context
             self.event_queue.put(
-                ActionEvent(
-                    action=action,
-                    type=STEP_EVENT_TYPE_STARTED,
-                )
+                ActionEvent(action=action, type=STEP_EVENT_TYPE_STARTED, payload="")
             )
 
             loop = asyncio.get_event_loop()
@@ -322,25 +324,20 @@ class Runner:
 
             try:
                 await task
-            except Exception as e:
-                return e
+            except Exception:
+                # do nothing, this should be caught in the callback
+                pass
 
-        return None
+        ## Once the step run completes, we need to remove the workflow spawn index
+        ## so we don't leak memory
+        if action.workflow_run_id in workflow_spawn_indices:
+            async with spawn_index_lock:
+                workflow_spawn_indices.pop(action.workflow_run_id)
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
     async def handle_start_group_key_run(self, action: Action) -> Exception | None:
         action_name = action.action_id
-        context = Context(
-            action,
-            self.dispatcher_client,
-            self.admin_client,
-            self.client.event,
-            self.client.rest,
-            self.client.workflow_listener,
-            self.workflow_run_event_listener,
-            self.worker_context,
-            self.client.config.namespace,
-        )
+        context = self.create_context(action)
 
         self.contexts[action.get_group_key_run_id] = context
 
@@ -351,8 +348,7 @@ class Runner:
             # send an event that the group key run has started
             self.event_queue.put(
                 ActionEvent(
-                    action=action,
-                    type=GROUP_KEY_EVENT_TYPE_STARTED,
+                    action=action, type=GROUP_KEY_EVENT_TYPE_STARTED, payload=""
                 )
             )
 
@@ -434,13 +430,12 @@ class Runner:
             self.cleanup_run_id(run_id)
 
     def serialize_output(self, output: Any) -> str:
-
         if isinstance(output, BaseModel):
             return output.model_dump_json()
 
         if output is not None:
             try:
-                return json.dumps(output)
+                return json.dumps(output, default=str)
             except Exception as e:
                 logger.error(f"Could not serialize output: {e}")
                 return str(output)
@@ -455,6 +450,6 @@ class Runner:
             running = len(self.tasks.keys())
 
 
-def errorWithTraceback(message: str, e: Exception) -> str:
+def pretty_format_exception(message: str, e: Exception) -> str:
     trace = "".join(traceback.format_exception(type(e), e, e.__traceback__))
     return f"{message}\n{trace}"

hatchet_sdk/worker/runner/utils/capture_logs.py

@@ -1,22 +1,24 @@
-import contextvars
 import functools
 import logging
 from concurrent.futures import ThreadPoolExecutor
+from contextvars import ContextVar
 from io import StringIO
-from typing import Any, Coroutine
+from typing import Any, Awaitable, Callable, ItemsView, ParamSpec, TypeVar
 
-from hatchet_sdk import logger
 from hatchet_sdk.clients.events import EventClient
+from hatchet_sdk.logger import logger
+from hatchet_sdk.runnables.contextvars import ctx_step_run_id, ctx_workflow_run_id
 
-wr: contextvars.ContextVar[str | None] = contextvars.ContextVar(
-    "workflow_run_id", default=None
-)
-sr: contextvars.ContextVar[str | None] = contextvars.ContextVar(
-    "step_run_id", default=None
-)
+T = TypeVar("T")
+P = ParamSpec("P")
 
 
-def copy_context_vars(ctx_vars, func, *args, **kwargs):
+def copy_context_vars(
+    ctx_vars: ItemsView[ContextVar[Any], Any],
+    func: Callable[P, T],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> T:
     for var, value in ctx_vars:
         var.set(value)
     return func(*args, **kwargs)
@@ -25,19 +27,20 @@ def copy_context_vars(ctx_vars, func, *args, **kwargs):
 class InjectingFilter(logging.Filter):
     # For some reason, only the InjectingFilter has access to the contextvars method sr.get(),
     # otherwise we would use emit within the CustomLogHandler
-    def filter(self, record):
-        record.workflow_run_id = wr.get()
-        record.step_run_id = sr.get()
+    def filter(self, record: logging.LogRecord) -> bool:
+        ## TODO: Change how we do this to not assign to the log record
+        record.workflow_run_id = ctx_workflow_run_id.get()
+        record.step_run_id = ctx_step_run_id.get()
         return True
 
 
-class CustomLogHandler(logging.StreamHandler):
-    def __init__(self, event_client: EventClient, stream=None):
+class CustomLogHandler(logging.StreamHandler):  # type: ignore[type-arg]
+    def __init__(self, event_client: EventClient, stream: StringIO | None = None):
         super().__init__(stream)
         self.logger_thread_pool = ThreadPoolExecutor(max_workers=1)
         self.event_client = event_client
 
-    def _log(self, line: str, step_run_id: str | None):
+    def _log(self, line: str, step_run_id: str | None) -> None:
         try:
             if not step_run_id:
                 return
@@ -46,20 +49,20 @@ class CustomLogHandler(logging.StreamHandler):
         except Exception as e:
             logger.error(f"Error logging: {e}")
 
-    def emit(self, record):
+    def emit(self, record: logging.LogRecord) -> None:
         super().emit(record)
 
         log_entry = self.format(record)
-        self.logger_thread_pool.submit(self._log, log_entry, record.step_run_id)
+
+        ## TODO: Change how we do this to not assign to the log record
+        self.logger_thread_pool.submit(self._log, log_entry, record.step_run_id)  # type: ignore
 
 
 def capture_logs(
-    logger: logging.Logger,
-    event_client: EventClient,
-    func: Coroutine[Any, Any, Any],
-):
+    logger: logging.Logger, event_client: "EventClient", func: Callable[P, Awaitable[T]]
+) -> Callable[P, Awaitable[T]]:
     @functools.wraps(func)
-    async def wrapper(*args, **kwargs):
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
         if not logger:
             raise Exception("No logger configured on client")