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

hatchet_sdk/worker/action_listener_process.py

@@ -2,27 +2,32 @@ import asyncio
 import logging
 import signal
 import time
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from multiprocessing import Queue
-from typing import Any, List, Mapping, Optional
+from typing import Any, Literal
 
 import grpc
 
-from hatchet_sdk.client import Client, new_client_raw
-from hatchet_sdk.clients.dispatcher.action_listener import Action
-from hatchet_sdk.clients.dispatcher.dispatcher import (
+from hatchet_sdk.client import Client
+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.clients.rest.models.update_worker_request import UpdateWorkerRequest
+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.runnables.contextvars import (
+    ctx_step_run_id,
+    ctx_worker_id,
+    ctx_workflow_run_id,
+)
 from hatchet_sdk.utils.backoff import exp_backoff_sleep
 
 ACTION_EVENT_RETRY_COUNT = 5
@@ -32,10 +37,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 = (
@@ -43,32 +49,39 @@ BLOCKED_THREAD_WARNING = (
 )
 
 
-@dataclass
 class WorkerActionListenerProcess:
-    name: str
-    actions: List[str]
-    max_runs: int
-    config: ClientConfig
-    action_queue: Queue
-    event_queue: Queue
-    handle_kill: bool = True
-    debug: bool = False
-    labels: dict = field(default_factory=dict)
-
-    listener: ActionListener = field(init=False, default=None)
-
-    killing: bool = field(init=False, default=False)
+    def __init__(
+        self,
+        name: str,
+        actions: list[str],
+        slots: int,
+        config: ClientConfig,
+        action_queue: "Queue[Action]",
+        event_queue: "Queue[ActionEvent | STOP_LOOP_TYPE]",
+        handle_kill: bool = True,
+        debug: bool = False,
+        labels: dict[str, str | int] = {},
+    ) -> None:
+        self.name = name
+        self.actions = actions
+        self.slots = slots
+        self.config = config
+        self.action_queue = action_queue
+        self.event_queue = event_queue
+        self.debug = debug
+        self.labels = labels
+        self.handle_kill = handle_kill
+
+        self.listener: ActionListener | None = None
+        self.killing = False
+        self.action_loop_task: asyncio.Task[None] | None = None
+        self.event_send_loop_task: asyncio.Task[None] | None = None
+        self.running_step_runs: dict[str, float] = {}
 
-    action_loop_task: asyncio.Task = field(init=False, default=None)
-    event_send_loop_task: asyncio.Task = field(init=False, default=None)
-
-    running_step_runs: Mapping[str, float] = field(init=False, default_factory=dict)
-
-    def __post_init__(self):
         if self.debug:
             logger.setLevel(logging.DEBUG)
 
-        self.client = new_client_raw(self.config, self.debug)
+        self.client = Client(config=self.config, debug=self.debug)
 
         loop = asyncio.get_event_loop()
         loop.add_signal_handler(
@@ -81,7 +94,16 @@ class WorkerActionListenerProcess:
             signal.SIGQUIT, lambda: asyncio.create_task(self.exit_gracefully())
         )
 
-    async def start(self, retry_attempt=0):
+    async def pause_task_assignment(self) -> None:
+        if self.listener is None:
+            raise ValueError("listener not started")
+
+        await self.client.rest.worker_api.worker_update(
+            worker=self.listener.worker_id,
+            update_worker_request=UpdateWorkerRequest(isPaused=True),
+        )
+
+    async def start(self, retry_attempt: int = 0) -> None:
         if retry_attempt > 5:
             logger.error("could not start action listener")
             return
@@ -89,15 +111,15 @@ 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(
                     worker_name=self.name,
                     services=["default"],
                     actions=self.actions,
-                    max_runs=self.max_runs,
-                    _labels=self.labels,
+                    slots=self.slots,
+                    raw_labels=self.labels,
                 )
             )
 
@@ -112,13 +134,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
@@ -126,11 +148,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
@@ -139,7 +161,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
@@ -189,15 +211,22 @@ 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:
+        if self.listener is None:
+            raise ValueError("listener not started")
+
         try:
             async for action in self.listener:
                 if action is None:
                     break
 
+                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)
+
                 # Process the action here
                 match action.action_type:
                     case ActionType.START_STEP_RUN:
@@ -205,6 +234,7 @@ class WorkerActionListenerProcess:
                             ActionEvent(
                                 action=action,
                                 type=STEP_EVENT_TYPE_STARTED,  # TODO ack type
+                                payload="",
                             )
                         )
                         logger.info(
@@ -224,6 +254,7 @@ class WorkerActionListenerProcess:
                             ActionEvent(
                                 action=action,
                                 type=GROUP_KEY_EVENT_TYPE_STARTED,  # TODO ack type
+                                payload="",
                             )
                         )
                         logger.info(
@@ -243,9 +274,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:
@@ -253,13 +284,7 @@ class WorkerActionListenerProcess:
 
         self.event_queue.put(STOP_LOOP)
 
-    async def pause_task_assignment(self) -> None:
-        await self.client.rest.aio.worker_api.worker_update(
-            worker=self.listener.worker_id,
-            update_worker_request=UpdateWorkerRequest(isPaused=True),
-        )
-
-    async def exit_gracefully(self, skip_unregister=False):
+    async def exit_gracefully(self) -> None:
         await self.pause_task_assignment()
 
         if self.killing:
@@ -274,13 +299,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

@@ -1,52 +1,64 @@
 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.client import Client
 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.runnables.task import Task
+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
 
-STOP_LOOP = "STOP_LOOP"
+STOP_LOOP_TYPE = Literal["STOP_LOOP"]
+STOP_LOOP: STOP_LOOP_TYPE = "STOP_LOOP"
 
 T = TypeVar("T")
 
 
-@dataclass
 class WorkerActionRunLoopManager:
-    name: str
-    action_registry: dict[str, Callable[[Context], T]]
-    validator_registry: dict[str, WorkflowValidator]
-    max_runs: int | None
-    config: ClientConfig
-    action_queue: Queue
-    event_queue: Queue
-    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)
-
-    killing: bool = field(init=False, default=False)
-    runner: Runner = field(init=False, default=None)
-
-    def __post_init__(self):
+    def __init__(
+        self,
+        name: str,
+        action_registry: dict[str, Task[Any, Any]],
+        validator_registry: dict[str, WorkflowValidator],
+        slots: int | None,
+        config: ClientConfig,
+        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] = {},
+    ) -> None:
+        self.name = name
+        self.action_registry = action_registry
+        self.validator_registry = validator_registry
+        self.slots = slots
+        self.config = config
+        self.action_queue = action_queue
+        self.event_queue = event_queue
+        self.loop = loop
+        self.handle_kill = handle_kill
+        self.debug = debug
+        self.labels = labels
+
         if self.debug:
             logger.setLevel(logging.DEBUG)
-        self.client = new_client_raw(self.config, self.debug)
+
+        self.killing = False
+        self.runner: Runner | None = None
+
+        self.client = Client(config=self.config, debug=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,
@@ -73,17 +85,17 @@ class WorkerActionRunLoopManager:
         self.runner = Runner(
             self.name,
             self.event_queue,
-            self.max_runs,
+            self.config,
+            self.slots,
             self.handle_kill,
             self.action_registry,
             self.validator_registry,
-            self.config,
             self.labels,
         )
 
         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 +103,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: