hatchet-sdk 1.12.3__py3-none-any.whl → 1.14.0__py3-none-any.whl

hatchet_sdk/worker/runner/runner.py

@@ -1,14 +1,14 @@
 import asyncio
-import contextvars
 import ctypes
 import functools
 import json
-import traceback
+from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import suppress
 from enum import Enum
 from multiprocessing import Queue
 from threading import Thread, current_thread
-from typing import Any, Callable, Dict, Literal, cast, overload
+from typing import Any, Literal, cast, overload
 
 from pydantic import BaseModel
 
@@ -30,7 +30,11 @@ from hatchet_sdk.contracts.dispatcher_pb2 import (
     STEP_EVENT_TYPE_FAILED,
     STEP_EVENT_TYPE_STARTED,
 )
-from hatchet_sdk.exceptions import NonRetryableException
+from hatchet_sdk.exceptions import (
+    IllegalTaskOutputError,
+    NonRetryableException,
+    TaskRunError,
+)
 from hatchet_sdk.features.runs import RunsClient
 from hatchet_sdk.logger import logger
 from hatchet_sdk.runnables.action import Action, ActionKey, ActionType
@@ -40,12 +44,17 @@ from hatchet_sdk.runnables.contextvars import (
     ctx_worker_id,
     ctx_workflow_run_id,
     spawn_index_lock,
+    task_count,
     workflow_spawn_indices,
 )
 from hatchet_sdk.runnables.task import Task
 from hatchet_sdk.runnables.types import R, TWorkflowInput
 from hatchet_sdk.worker.action_listener_process import ActionEvent
-from hatchet_sdk.worker.runner.utils.capture_logs import copy_context_vars
+from hatchet_sdk.worker.runner.utils.capture_logs import (
+    AsyncLogSender,
+    ContextVarToCopy,
+    copy_context_vars,
+)
 
 
 class WorkerStatus(Enum):
@@ -61,10 +70,11 @@ class Runner:
         event_queue: "Queue[ActionEvent]",
         config: ClientConfig,
         slots: int,
-        handle_kill: bool = True,
-        action_registry: dict[str, Task[TWorkflowInput, R]] = {},
-        labels: dict[str, str | int] = {},
-        lifespan_context: Any | None = None,
+        handle_kill: bool,
+        action_registry: dict[str, Task[TWorkflowInput, R]],
+        labels: dict[str, str | int] | None,
+        lifespan_context: Any | None,
+        log_sender: AsyncLogSender,
     ):
         # We store the config so we can dynamically create clients for the dispatcher client.
         self.config = config
@@ -72,13 +82,14 @@ class Runner:
         self.slots = slots
         self.tasks: dict[ActionKey, asyncio.Task[Any]] = {}  # Store run ids and futures
         self.contexts: dict[ActionKey, Context] = {}  # Store run ids and contexts
-        self.action_registry = action_registry
+        self.action_registry = action_registry or {}
 
         self.event_queue = event_queue
 
         # The thread pool is used for synchronous functions which need to run concurrently
         self.thread_pool = ThreadPoolExecutor(max_workers=slots)
-        self.threads: Dict[ActionKey, Thread] = {}  # Store run ids and threads
+        self.threads: dict[ActionKey, Thread] = {}  # Store run ids and threads
+        self.running_tasks = set[asyncio.Task[Exception | None]]()
 
         self.killing = False
         self.handle_kill = handle_kill
@@ -101,10 +112,11 @@ class Runner:
         self.durable_event_listener = DurableEventListener(self.config)
 
         self.worker_context = WorkerContext(
-            labels=labels, client=Client(config=config).dispatcher
+            labels=labels or {}, client=Client(config=config).dispatcher
         )
 
         self.lifespan_context = lifespan_context
+        self.log_sender = log_sender
 
         if self.config.enable_thread_pool_monitoring:
             self.start_background_monitoring()
@@ -116,68 +128,90 @@ class Runner:
         if self.worker_context.id() is None:
             self.worker_context._worker_id = action.worker_id
 
+        t: asyncio.Task[Exception | None] | None = None
         match action.action_type:
             case ActionType.START_STEP_RUN:
                 log = f"run: start step: {action.action_id}/{action.step_run_id}"
                 logger.info(log)
-                asyncio.create_task(self.handle_start_step_run(action))
+                t = asyncio.create_task(self.handle_start_step_run(action))
             case ActionType.CANCEL_STEP_RUN:
                 log = f"cancel: step run:  {action.action_id}/{action.step_run_id}/{action.retry_count}"
                 logger.info(log)
-                asyncio.create_task(self.handle_cancel_action(action))
+                t = asyncio.create_task(self.handle_cancel_action(action))
             case ActionType.START_GET_GROUP_KEY:
                 log = f"run: get group key:  {action.action_id}/{action.get_group_key_run_id}"
                 logger.info(log)
-                asyncio.create_task(self.handle_start_group_key_run(action))
+                t = asyncio.create_task(self.handle_start_group_key_run(action))
             case _:
                 log = f"unknown action type: {action.action_type}"
                 logger.error(log)
 
+        if t is not None:
+            self.running_tasks.add(t)
+            t.add_done_callback(lambda task: self.running_tasks.discard(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.key)
 
-            errored = False
-            cancelled = task.cancelled()
-            output = None
+            if task.cancelled():
+                return
 
-            # Get the output from the future
             try:
-                if not cancelled:
-                    output = task.result()
+                output = task.result()
             except Exception as e:
-                errored = True
-
                 should_not_retry = isinstance(e, NonRetryableException)
 
+                exc = TaskRunError.from_exception(e)
+
                 # This except is coming from the application itself, so we want to send that to the Hatchet instance
                 self.event_queue.put(
                     ActionEvent(
                         action=action,
                         type=STEP_EVENT_TYPE_FAILED,
-                        payload=str(pretty_format_exception(f"{e}", e)),
+                        payload=exc.serialize(),
                         should_not_retry=should_not_retry,
                     )
                 )
 
-                logger.error(
-                    f"failed step run: {action.action_id}/{action.step_run_id}"
+                log_with_level = logger.info if should_not_retry else logger.error
+
+                log_with_level(
+                    f"failed step run: {action.action_id}/{action.step_run_id}\n{exc.serialize()}"
                 )
 
-            if not errored and not cancelled:
+                return
+
+            try:
+                output = self.serialize_output(output)
+
                 self.event_queue.put(
                     ActionEvent(
                         action=action,
                         type=STEP_EVENT_TYPE_COMPLETED,
-                        payload=self.serialize_output(output),
+                        payload=output,
+                        should_not_retry=False,
+                    )
+                )
+            except IllegalTaskOutputError as e:
+                exc = TaskRunError.from_exception(e)
+                self.event_queue.put(
+                    ActionEvent(
+                        action=action,
+                        type=STEP_EVENT_TYPE_FAILED,
+                        payload=exc.serialize(),
                         should_not_retry=False,
                     )
                 )
 
-                logger.info(
-                    f"finished step run: {action.action_id}/{action.step_run_id}"
+                logger.error(
+                    f"failed step run: {action.action_id}/{action.step_run_id}\n{exc.serialize()}"
                 )
 
+                return
+
+            logger.info(f"finished step run: {action.action_id}/{action.step_run_id}")
+
         return inner_callback
 
     def group_key_run_callback(
@@ -186,51 +220,65 @@ class Runner:
         def inner_callback(task: asyncio.Task[Any]) -> None:
             self.cleanup_run_id(action.key)
 
-            errored = False
-            cancelled = task.cancelled()
-            output = None
+            if task.cancelled():
+                return
 
-            # Get the output from the future
             try:
-                if not cancelled:
-                    output = task.result()
+                output = task.result()
             except Exception as e:
-                errored = True
+                exc = TaskRunError.from_exception(e)
+
                 self.event_queue.put(
                     ActionEvent(
                         action=action,
                         type=GROUP_KEY_EVENT_TYPE_FAILED,
-                        payload=str(pretty_format_exception(f"{e}", e)),
+                        payload=exc.serialize(),
                         should_not_retry=False,
                     )
                 )
 
                 logger.error(
-                    f"failed step run: {action.action_id}/{action.step_run_id}"
+                    f"failed step run: {action.action_id}/{action.step_run_id}\n{exc.serialize()}"
                 )
 
-            if not errored and not cancelled:
+                return
+
+            try:
+                output = self.serialize_output(output)
+
                 self.event_queue.put(
                     ActionEvent(
                         action=action,
                         type=GROUP_KEY_EVENT_TYPE_COMPLETED,
-                        payload=self.serialize_output(output),
+                        payload=output,
+                        should_not_retry=False,
+                    )
+                )
+            except IllegalTaskOutputError as e:
+                exc = TaskRunError.from_exception(e)
+                self.event_queue.put(
+                    ActionEvent(
+                        action=action,
+                        type=STEP_EVENT_TYPE_FAILED,
+                        payload=exc.serialize(),
                         should_not_retry=False,
                     )
                 )
 
-                logger.info(
-                    f"finished step run: {action.action_id}/{action.step_run_id}"
+                logger.error(
+                    f"failed step run: {action.action_id}/{action.step_run_id}\n{exc.serialize()}"
                 )
 
+                return
+
+            logger.info(f"finished step run: {action.action_id}/{action.step_run_id}")
+
         return inner_callback
 
     def thread_action_func(
         self, ctx: Context, task: Task[TWorkflowInput, R], action: Action
     ) -> R:
-        if action.step_run_id:
-            self.threads[action.key] = current_thread()
-        elif action.get_group_key_run_id:
+        if action.step_run_id or action.get_group_key_run_id:
             self.threads[action.key] = current_thread()
 
         return task.call(ctx)
@@ -250,28 +298,36 @@ class Runner:
         try:
             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
-                    # contextvars
-                    copy_context_vars,
-                    contextvars.copy_context().items(),
-                    self.thread_action_func,
-                    ctx,
-                    task,
-                    action,
-                )
-
-                loop = asyncio.get_event_loop()
-                return await loop.run_in_executor(self.thread_pool, pfunc)
-        except Exception as e:
-            logger.error(
-                pretty_format_exception(
-                    f"exception raised in action ({action.action_id}, retry={action.retry_count}):\n{e}",
-                    e,
-                )
+            pfunc = functools.partial(
+                # we must copy the context vars to the new thread, as only asyncio natively supports
+                # contextvars
+                copy_context_vars,
+                [
+                    ContextVarToCopy(
+                        name="ctx_step_run_id",
+                        value=action.step_run_id,
+                    ),
+                    ContextVarToCopy(
+                        name="ctx_workflow_run_id",
+                        value=action.workflow_run_id,
+                    ),
+                    ContextVarToCopy(
+                        name="ctx_worker_id",
+                        value=action.worker_id,
+                    ),
+                    ContextVarToCopy(
+                        name="ctx_action_key",
+                        value=action.key,
+                    ),
+                ],
+                self.thread_action_func,
+                ctx,
+                task,
+                action,
             )
-            raise e
+
+            loop = asyncio.get_event_loop()
+            return await loop.run_in_executor(self.thread_pool, pfunc)
         finally:
             self.cleanup_run_id(action.key)
 
@@ -295,7 +351,7 @@ class Runner:
             while True:
                 await self.log_thread_pool_status()
 
-                for key in self.threads.keys():
+                for key in self.threads:
                     if key not in self.tasks:
                         logger.debug(f"Potential zombie thread found for key {key}")
 
@@ -350,6 +406,7 @@ class Runner:
             worker=self.worker_context,
             runs_client=self.runs_client,
             lifespan_context=self.lifespan_context,
+            log_sender=self.log_sender,
         )
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
@@ -361,7 +418,8 @@ class Runner:
 
         if action_func:
             context = self.create_context(
-                action, True if action_func.is_durable else False
+                action,
+                True if action_func.is_durable else False,  # noqa: SIM210
             )
 
             self.contexts[action.key] = context
@@ -382,11 +440,12 @@ class Runner:
             task.add_done_callback(self.step_run_callback(action))
             self.tasks[action.key] = task
 
-            try:
+            task_count.increment()
+
+            ## FIXME: Handle cancelled exceptions and other special exceptions
+            ## that we don't want to suppress here
+            with suppress(Exception):
                 await task
-            except Exception:
-                # do nothing, this should be caught in the callback
-                pass
 
         ## Once the step run completes, we need to remove the workflow spawn index
         ## so we don't leak memory
@@ -444,7 +503,7 @@ class Runner:
             res = ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(ident), exc)
             if res == 0:
                 raise ValueError("Invalid thread ID")
-            elif res != 1:
+            if res != 1:
                 logger.error("PyThreadState_SetAsyncExc failed")
 
                 # Call with exception set to 0 is needed to cleanup properly.
@@ -487,8 +546,16 @@ class Runner:
             self.cleanup_run_id(key)
 
     def serialize_output(self, output: Any) -> str:
+        if not output:
+            return ""
+
         if isinstance(output, BaseModel):
-            return output.model_dump_json()
+            output = output.model_dump()
+
+        if not isinstance(output, dict):
+            raise IllegalTaskOutputError(
+                f"Tasks must return either a dictionary or a Pydantic BaseModel which can be serialized to a JSON object. Got object of type {type(output)} instead."
+            )
 
         if output is not None:
             try:
@@ -505,8 +572,3 @@ class Runner:
             logger.info(f"waiting for {running} tasks to finish...")
             await asyncio.sleep(1)
             running = len(self.tasks.keys())
-
-
-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,73 +1,114 @@
+import asyncio
 import functools
 import logging
-from concurrent.futures import ThreadPoolExecutor
-from contextvars import ContextVar
+from collections.abc import Awaitable, Callable
 from io import StringIO
-from typing import Any, Awaitable, Callable, ItemsView, ParamSpec, TypeVar
+from typing import Literal, ParamSpec, TypeVar
+
+from pydantic import BaseModel
 
 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
+from hatchet_sdk.runnables.contextvars import (
+    ctx_action_key,
+    ctx_step_run_id,
+    ctx_worker_id,
+    ctx_workflow_run_id,
+)
+from hatchet_sdk.utils.typing import STOP_LOOP, STOP_LOOP_TYPE
 
 T = TypeVar("T")
 P = ParamSpec("P")
 
 
+class ContextVarToCopy(BaseModel):
+    name: Literal[
+        "ctx_workflow_run_id", "ctx_step_run_id", "ctx_action_key", "ctx_worker_id"
+    ]
+    value: str | None
+
+
 def copy_context_vars(
-    ctx_vars: ItemsView[ContextVar[Any], Any],
+    ctx_vars: list[ContextVarToCopy],
     func: Callable[P, T],
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> T:
-    for var, value in ctx_vars:
-        var.set(value)
+    for var in ctx_vars:
+        if var.name == "ctx_workflow_run_id":
+            ctx_workflow_run_id.set(var.value)
+        elif var.name == "ctx_step_run_id":
+            ctx_step_run_id.set(var.value)
+        elif var.name == "ctx_action_key":
+            ctx_action_key.set(var.value)
+        elif var.name == "ctx_worker_id":
+            ctx_worker_id.set(var.value)
+        else:
+            raise ValueError(f"Unknown context variable name: {var.name}")
+
     return 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: 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 LogRecord(BaseModel):
+    message: str
+    step_run_id: str
 
 
-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)
+class AsyncLogSender:
+    def __init__(self, event_client: EventClient):
         self.event_client = event_client
+        self.q = asyncio.Queue[LogRecord | STOP_LOOP_TYPE](maxsize=1000)
+
+    async def consume(self) -> None:
+        while True:
+            record = await self.q.get()
+
+            if record == STOP_LOOP:
+                break
 
-    def _log(self, line: str, step_run_id: str | None) -> None:
+            try:
+                self.event_client.log(
+                    message=record.message, step_run_id=record.step_run_id
+                )
+            except Exception as e:
+                logger.error(f"Error logging: {e}")
+
+    def publish(self, record: LogRecord | STOP_LOOP_TYPE) -> None:
         try:
-            if not step_run_id:
-                return
+            self.q.put_nowait(record)
+        except asyncio.QueueFull:
+            logger.warning("Log queue is full, dropping log message")
+
 
-            self.event_client.log(message=line, step_run_id=step_run_id)
-        except Exception as e:
-            logger.error(f"Error logging: {e}")
+class CustomLogHandler(logging.StreamHandler):  # type: ignore[type-arg]
+    def __init__(self, log_sender: AsyncLogSender, stream: StringIO):
+        super().__init__(stream)
+
+        self.log_sender = log_sender
 
     def emit(self, record: logging.LogRecord) -> None:
         super().emit(record)
 
         log_entry = self.format(record)
+        step_run_id = ctx_step_run_id.get()
+
+        if not step_run_id:
+            return
 
-        ## 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
+        self.log_sender.publish(LogRecord(message=log_entry, step_run_id=step_run_id))
 
 
 def capture_logs(
-    logger: logging.Logger, event_client: "EventClient", func: Callable[P, Awaitable[T]]
+    logger: logging.Logger, log_sender: AsyncLogSender, func: Callable[P, Awaitable[T]]
 ) -> Callable[P, Awaitable[T]]:
     @functools.wraps(func)
     async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
         log_stream = StringIO()
-        custom_handler = CustomLogHandler(event_client, log_stream)
+        custom_handler = CustomLogHandler(log_sender, log_stream)
         custom_handler.setLevel(logging.INFO)
-        custom_handler.addFilter(InjectingFilter())
-        logger.addHandler(custom_handler)
+
+        if not any(h for h in logger.handlers if isinstance(h, CustomLogHandler)):
+            logger.addHandler(custom_handler)
 
         try:
             result = await func(*args, **kwargs)