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

hatchet_sdk/worker/__init__.py

@@ -1 +0,0 @@
-from .worker import Worker, WorkerStartOptions, WorkerStatus

hatchet_sdk/worker/action_listener_process.py

@@ -4,22 +4,22 @@ import signal
 import time
 from dataclasses import dataclass, field
 from multiprocessing import Queue
-from typing import Any, List, Mapping, Optional
+from typing import Any, List, Literal
 
 import grpc
 
-from hatchet_sdk.clients.dispatcher.action_listener import Action
-from hatchet_sdk.clients.dispatcher.dispatcher import (
+from hatchet_sdk.clients.dispatcher.action_listener import (
+    Action,
     ActionListener,
+    ActionType,
     GetActionListenerRequest,
-    new_dispatcher,
 )
+from hatchet_sdk.clients.dispatcher.dispatcher import DispatcherClient
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.contracts.dispatcher_pb2 import (
     GROUP_KEY_EVENT_TYPE_STARTED,
     STEP_EVENT_TYPE_STARTED,
-    ActionType,
 )
-from hatchet_sdk.loader import ClientConfig
 from hatchet_sdk.logger import logger
 from hatchet_sdk.utils.backoff import exp_backoff_sleep
 
@@ -30,10 +30,11 @@ ACTION_EVENT_RETRY_COUNT = 5
 class ActionEvent:
     action: Action
     type: Any  # TODO type
-    payload: Optional[str] = None
+    payload: str
 
 
-STOP_LOOP = "STOP_LOOP"  # Sentinel object to stop the loop
+STOP_LOOP_TYPE = Literal["STOP_LOOP"]
+STOP_LOOP: STOP_LOOP_TYPE = "STOP_LOOP"  # Sentinel object to stop the loop
 
 # TODO link to a block post
 BLOCKED_THREAD_WARNING = (
@@ -41,7 +42,7 @@ BLOCKED_THREAD_WARNING = (
 )
 
 
-def noop_handler():
+def noop_handler() -> None:
     pass
 
 
@@ -51,22 +52,22 @@ class WorkerActionListenerProcess:
     actions: List[str]
     max_runs: int
     config: ClientConfig
-    action_queue: Queue
-    event_queue: Queue
+    action_queue: "Queue[Action]"
+    event_queue: "Queue[ActionEvent | STOP_LOOP_TYPE]"
     handle_kill: bool = True
     debug: bool = False
-    labels: dict = field(default_factory=dict)
+    labels: dict[str, str | int] = field(default_factory=dict)
 
-    listener: ActionListener = field(init=False, default=None)
+    listener: ActionListener = field(init=False)
 
     killing: bool = field(init=False, default=False)
 
-    action_loop_task: asyncio.Task = field(init=False, default=None)
-    event_send_loop_task: asyncio.Task = field(init=False, default=None)
+    action_loop_task: asyncio.Task[None] | None = field(init=False, default=None)
+    event_send_loop_task: asyncio.Task[None] | None = field(init=False, default=None)
 
-    running_step_runs: Mapping[str, float] = field(init=False, default_factory=dict)
+    running_step_runs: dict[str, float] = field(init=False, default_factory=dict)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.debug:
             logger.setLevel(logging.DEBUG)
 
@@ -77,7 +78,7 @@ class WorkerActionListenerProcess:
             signal.SIGQUIT, lambda: asyncio.create_task(self.exit_gracefully())
         )
 
-    async def start(self, retry_attempt=0):
+    async def start(self, retry_attempt: int = 0) -> None:
         if retry_attempt > 5:
             logger.error("could not start action listener")
             return
@@ -85,7 +86,7 @@ class WorkerActionListenerProcess:
         logger.debug(f"starting action listener: {self.name}")
 
         try:
-            self.dispatcher_client = new_dispatcher(self.config)
+            self.dispatcher_client = DispatcherClient(self.config)
 
             self.listener = await self.dispatcher_client.get_action_listener(
                 GetActionListenerRequest(
@@ -108,13 +109,13 @@ class WorkerActionListenerProcess:
         self.blocked_main_loop = asyncio.create_task(self.start_blocked_main_loop())
 
     # TODO move event methods to separate class
-    async def _get_event(self):
+    async def _get_event(self) -> ActionEvent | STOP_LOOP_TYPE:
         loop = asyncio.get_running_loop()
         return await loop.run_in_executor(None, self.event_queue.get)
 
-    async def start_event_send_loop(self):
+    async def start_event_send_loop(self) -> None:
         while True:
-            event: ActionEvent = await self._get_event()
+            event = await self._get_event()
             if event == STOP_LOOP:
                 logger.debug("stopping event send loop...")
                 break
@@ -122,11 +123,11 @@ class WorkerActionListenerProcess:
             logger.debug(f"tx: event: {event.action.action_id}/{event.type}")
             asyncio.create_task(self.send_event(event))
 
-    async def start_blocked_main_loop(self):
+    async def start_blocked_main_loop(self) -> None:
         threshold = 1
         while not self.killing:
             count = 0
-            for step_run_id, start_time in self.running_step_runs.items():
+            for _, start_time in self.running_step_runs.items():
                 diff = self.now() - start_time
                 if diff > threshold:
                     count += 1
@@ -135,7 +136,7 @@ class WorkerActionListenerProcess:
                 logger.warning(f"{BLOCKED_THREAD_WARNING}: Waiting Steps {count}")
             await asyncio.sleep(1)
 
-    async def send_event(self, event: ActionEvent, retry_attempt: int = 1):
+    async def send_event(self, event: ActionEvent, retry_attempt: int = 1) -> None:
         try:
             match event.action.action_type:
                 # FIXME: all events sent from an execution of a function are of type ActionType.START_STEP_RUN since
@@ -185,10 +186,10 @@ class WorkerActionListenerProcess:
                 await exp_backoff_sleep(retry_attempt, 1)
                 await self.send_event(event, retry_attempt + 1)
 
-    def now(self):
+    def now(self) -> float:
         return time.time()
 
-    async def start_action_loop(self):
+    async def start_action_loop(self) -> None:
         try:
             async for action in self.listener:
                 if action is None:
@@ -201,6 +202,7 @@ class WorkerActionListenerProcess:
                             ActionEvent(
                                 action=action,
                                 type=STEP_EVENT_TYPE_STARTED,  # TODO ack type
+                                payload="",
                             )
                         )
                         logger.info(
@@ -220,6 +222,7 @@ class WorkerActionListenerProcess:
                             ActionEvent(
                                 action=action,
                                 type=GROUP_KEY_EVENT_TYPE_STARTED,  # TODO ack type
+                                payload="",
                             )
                         )
                         logger.info(
@@ -239,9 +242,9 @@ class WorkerActionListenerProcess:
         finally:
             logger.info("action loop closed")
             if not self.killing:
-                await self.exit_gracefully(skip_unregister=True)
+                await self.exit_gracefully()
 
-    async def cleanup(self):
+    async def cleanup(self) -> None:
         self.killing = True
 
         if self.listener is not None:
@@ -249,7 +252,7 @@ class WorkerActionListenerProcess:
 
         self.event_queue.put(STOP_LOOP)
 
-    async def exit_gracefully(self, skip_unregister=False):
+    async def exit_gracefully(self) -> None:
         if self.killing:
             return
 
@@ -262,13 +265,13 @@ class WorkerActionListenerProcess:
 
         logger.info("action listener closed")
 
-    def exit_forcefully(self):
+    def exit_forcefully(self) -> None:
         asyncio.run(self.cleanup())
         logger.debug("forcefully closing listener...")
 
 
-def worker_action_listener_process(*args, **kwargs):
-    async def run():
+def worker_action_listener_process(*args: Any, **kwargs: Any) -> None:
+    async def run() -> None:
         process = WorkerActionListenerProcess(*args, **kwargs)
         await process.start()
         # Keep the process running

hatchet_sdk/worker/runner/run_loop_manager.py

@@ -2,18 +2,20 @@ import asyncio
 import logging
 from dataclasses import dataclass, field
 from multiprocessing import Queue
-from typing import Callable, TypeVar
+from typing import Any, Literal, TypeVar
 
-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.loader import ClientConfig
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.logger import logger
-from hatchet_sdk.utils.types import WorkflowValidator
+from hatchet_sdk.utils.typing import WorkflowValidator
+from hatchet_sdk.worker.action_listener_process import ActionEvent
 from hatchet_sdk.worker.runner.runner import Runner
 from hatchet_sdk.worker.runner.utils.capture_logs import capture_logs
+from hatchet_sdk.workflow import Step
 
-STOP_LOOP = "STOP_LOOP"
+STOP_LOOP_TYPE = Literal["STOP_LOOP"]
+STOP_LOOP: STOP_LOOP_TYPE = "STOP_LOOP"
 
 T = TypeVar("T")
 
@@ -21,32 +23,32 @@ T = TypeVar("T")
 @dataclass
 class WorkerActionRunLoopManager:
     name: str
-    action_registry: dict[str, Callable[[Context], T]]
+    action_registry: dict[str, Step[Any]]
     validator_registry: dict[str, WorkflowValidator]
     max_runs: int | None
     config: ClientConfig
-    action_queue: Queue
-    event_queue: Queue
+    action_queue: "Queue[Action | STOP_LOOP_TYPE]"
+    event_queue: "Queue[ActionEvent]"
     loop: asyncio.AbstractEventLoop
     handle_kill: bool = True
     debug: bool = False
     labels: dict[str, str | int] = field(default_factory=dict)
 
-    client: Client = field(init=False, default=None)
+    client: Client = field(init=False)
 
     killing: bool = field(init=False, default=False)
-    runner: Runner = field(init=False, default=None)
+    runner: Runner | None = field(init=False, default=None)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.debug:
             logger.setLevel(logging.DEBUG)
         self.client = new_client_raw(self.config, self.debug)
         self.start()
 
-    def start(self, retry_count=1):
-        k = self.loop.create_task(self.async_start(retry_count))
+    def start(self, retry_count: int = 1) -> None:
+        k = self.loop.create_task(self.aio_start(retry_count))  # noqa: F841
 
-    async def async_start(self, retry_count=1):
+    async def aio_start(self, retry_count: int = 1) -> None:
         await capture_logs(
             self.client.logInterceptor,
             self.client.event,
@@ -83,7 +85,7 @@ class WorkerActionRunLoopManager:
 
         logger.debug(f"'{self.name}' waiting for {list(self.action_registry.keys())}")
         while not self.killing:
-            action: Action = await self._get_action()
+            action = await self._get_action()
             if action == STOP_LOOP:
                 logger.debug("stopping action runner loop...")
                 break
@@ -91,7 +93,7 @@ class WorkerActionRunLoopManager:
             self.runner.run(action)
         logger.debug("action runner loop stopped")
 
-    async def _get_action(self):
+    async def _get_action(self) -> Action | STOP_LOOP_TYPE:
         return await self.loop.run_in_executor(None, self.action_queue.get)
 
     async def exit_gracefully(self) -> None:

hatchet_sdk/worker/runner/runner.py

@@ -3,23 +3,23 @@ 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, TypeVar, cast
 
 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.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.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
 from hatchet_sdk.context.worker_context import WorkerContext
 from hatchet_sdk.contracts.dispatcher_pb2 import (
     GROUP_KEY_EVENT_TYPE_COMPLETED,
@@ -28,14 +28,14 @@ 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.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.workflow import Step
+
+T = TypeVar("T")
 
 
 class WorkerStatus(Enum):
@@ -52,7 +52,7 @@ class Runner:
         event_queue: "Queue[Any]",
         max_runs: int | None = None,
         handle_kill: bool = True,
-        action_registry: dict[str, Callable[..., Any]] = {},
+        action_registry: dict[str, Step[T]] = {},
         validator_registry: dict[str, WorkflowValidator] = {},
         config: ClientConfig = ClientConfig(),
         labels: dict[str, str | int] = {},
@@ -64,7 +64,7 @@ class Runner:
         self.max_runs = max_runs
         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: dict[str, Step[T]] = action_registry
         self.validator_registry = validator_registry
 
         self.event_queue = event_queue
@@ -78,9 +78,9 @@ 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.worker_context = WorkerContext(
@@ -195,10 +195,7 @@ class Runner:
 
         return inner_callback
 
-    ## TODO: Stricter type hinting here
-    def thread_action_func(
-        self, context: Context, action_func: Callable[..., Any], action: Action
-    ) -> Any:
+    def thread_action_func(self, context: Context, step: Step[T], action: Action) -> T:
         if action.step_run_id is not None and action.step_run_id != "":
             self.threads[action.step_run_id] = current_thread()
         elif (
@@ -207,25 +204,22 @@ class Runner:
         ):
             self.threads[action.get_group_key_run_id] = current_thread()
 
-        return action_func(context)
+        return step.call(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: Callable[..., Any],
+        step: Step[T],
         action: Action,
         run_id: str,
-    ) -> Any:
-        wr.set(context.workflow_run_id())
+    ) -> T:
+        wr.set(context.workflow_run_id)
         sr.set(context.step_run_id)
 
         try:
-            if (
-                hasattr(action_func, "is_coroutine") and action_func.is_coroutine
-            ) or asyncio.iscoroutinefunction(action_func):
-                return await action_func(context)
+            if step.is_async_function:
+                return await step.aio_call(context)
             else:
                 pfunc = functools.partial(
                     # we must copy the context vars to the new thread, as only asyncio natively supports
@@ -234,7 +228,7 @@ class Runner:
                     contextvars.copy_context().items(),
                     self.thread_action_func,
                     context,
-                    action_func,
+                    step,
                     action,
                 )
 
@@ -261,23 +255,7 @@ class Runner:
         if run_id in self.contexts:
             del self.contexts[run_id]
 
-    def create_context(
-        self, action: Action, action_func: Callable[..., Any] | None
-    ) -> 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,
-            )
-
+    def create_context(self, action: Action) -> Context:
         return Context(
             action,
             self.dispatcher_client,
@@ -292,22 +270,19 @@ 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)
+        context = self.create_context(action)
 
         self.contexts[action.step_run_id] = context
 
         if action_func:
             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,10 +297,9 @@ class Runner:
 
             try:
                 await task
-            except Exception as e:
-                return e
-
-        return None
+            except Exception:
+                # do nothing, this should be caught in the callback
+                pass
 
     ## 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:
@@ -351,8 +325,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,7 +407,6 @@ class Runner:
             self.cleanup_run_id(run_id)
 
     def serialize_output(self, output: Any) -> str:
-
         if isinstance(output, BaseModel):
             return output.model_dump_json()
 

hatchet_sdk/worker/runner/utils/capture_logs.py

@@ -2,11 +2,12 @@ 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
 
 wr: contextvars.ContextVar[str | None] = contextvars.ContextVar(
     "workflow_run_id", default=None
@@ -16,7 +17,16 @@ sr: contextvars.ContextVar[str | None] = contextvars.ContextVar(
 )
 
 
-def copy_context_vars(ctx_vars, func, *args, **kwargs):
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+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 +35,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):
+    def filter(self, record: logging.LogRecord) -> bool:
+        ## TODO: Change how we do this to not assign to the log record
         record.workflow_run_id = wr.get()
         record.step_run_id = sr.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 +57,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")
 

hatchet_sdk/worker/runner/utils/error_with_traceback.py

@@ -1,6 +1,6 @@
 import traceback
 
 
-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}"