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

hatchet_sdk/utils/aio_utils.py

@@ -1,83 +1,4 @@
 import asyncio
-import inspect
-from concurrent.futures import Executor
-from functools import partial, wraps
-from typing import Any
-
-
-## TODO: Stricter typing here
-def sync_to_async(func: Any) -> Any:
-    """
-    A decorator to run a synchronous function or coroutine in an asynchronous context with added
-    asyncio loop safety.
-
-    This decorator allows you to safely call synchronous functions or coroutines from an
-    asynchronous function by running them in an executor.
-
-    Args:
-        func (callable): The synchronous function or coroutine to be run asynchronously.
-
-    Returns:
-        callable: An asynchronous wrapper function that runs the given function in an executor.
-
-    Example:
-        @sync_to_async
-        def sync_function(x, y):
-            return x + y
-
-        @sync_to_async
-        async def async_function(x, y):
-            return x + y
-
-
-        def undecorated_function(x, y):
-            return x + y
-
-        async def main():
-            result1 = await sync_function(1, 2)
-            result2 = await async_function(3, 4)
-            result3 = await sync_to_async(undecorated_function)(5, 6)
-            print(result1, result2, result3)
-
-        asyncio.run(main())
-    """
-
-    ## TODO: Stricter typing here
-    @wraps(func)
-    async def run(
-        *args: Any,
-        loop: asyncio.AbstractEventLoop | None = None,
-        executor: Executor | None = None,
-        **kwargs: Any
-    ) -> Any:
-        """
-        The asynchronous wrapper function that runs the given function in an executor.
-
-        Args:
-            *args: Positional arguments to pass to the function.
-            loop (asyncio.AbstractEventLoop, optional): The event loop to use. If None, the current running loop is used.
-            executor (concurrent.futures.Executor, optional): The executor to use. If None, the default executor is used.
-            **kwargs: Keyword arguments to pass to the function.
-
-        Returns:
-            The result of the function call.
-        """
-        if loop is None:
-            loop = asyncio.get_running_loop()
-
-        if inspect.iscoroutinefunction(func):
-            # Wrap the coroutine to run it in an executor
-            async def wrapper() -> Any:
-                return await func(*args, **kwargs)
-
-            pfunc = partial(asyncio.run, wrapper())
-            return await loop.run_in_executor(executor, pfunc)
-        else:
-            # Run the synchronous function in an executor
-            pfunc = partial(func, *args, **kwargs)
-            return await loop.run_in_executor(executor, pfunc)
-
-    return run
 
 
 def get_active_event_loop() -> asyncio.AbstractEventLoop | None:

hatchet_sdk/utils/proto_enums.py

@@ -45,10 +45,3 @@ def convert_proto_enum_to_python(
         return None
 
     return python_enum_class[proto_enum.Name(value)]
-
-
-def maybe_int_to_str(value: int | None) -> str | None:
-    if value is None:
-        return None
-
-    return str(value)

hatchet_sdk/utils/timedelta_to_expression.py

@@ -0,0 +1,23 @@
+from datetime import timedelta
+
+DAY = 86400
+HOUR = 3600
+MINUTE = 60
+
+Duration = timedelta | str
+
+
+def timedelta_to_expr(td: Duration) -> str:
+    if isinstance(td, str):
+        return td
+
+    seconds = td.seconds
+
+    if seconds % DAY == 0:
+        return f"{seconds // DAY}d"
+    elif seconds % HOUR == 0:
+        return f"{seconds // HOUR}h"
+    elif seconds % MINUTE == 0:
+        return f"{seconds // MINUTE}m"
+    else:
+        return f"{seconds}s"

hatchet_sdk/utils/typing.py

@@ -1,11 +1,11 @@
-from typing import Any, Mapping, Type, TypeVar
+from typing import Any, Mapping, Type, TypeGuard, TypeVar
 
 from pydantic import BaseModel
 
 T = TypeVar("T", bound=BaseModel)
 
 
-def is_basemodel_subclass(model: Any) -> bool:
+def is_basemodel_subclass(model: Any) -> TypeGuard[Type[BaseModel]]:
     try:
         return issubclass(model, BaseModel)
     except TypeError:

hatchet_sdk/v0/clients/rest_client.py

@@ -9,6 +9,7 @@ from pydantic import StrictInt
 from hatchet_sdk.v0.clients.rest.api.event_api import EventApi
 from hatchet_sdk.v0.clients.rest.api.log_api import LogApi
 from hatchet_sdk.v0.clients.rest.api.step_run_api import StepRunApi
+from hatchet_sdk.v0.clients.rest.api.worker_api import WorkerApi
 from hatchet_sdk.v0.clients.rest.api.workflow_api import WorkflowApi
 from hatchet_sdk.v0.clients.rest.api.workflow_run_api import WorkflowRunApi
 from hatchet_sdk.v0.clients.rest.api.workflow_runs_api import WorkflowRunsApi
@@ -85,6 +86,7 @@ class AsyncRestApi:
         self._step_run_api = None
         self._event_api = None
         self._log_api = None
+        self._worker_api = None
 
     @property
     def api_client(self):
@@ -104,6 +106,13 @@ class AsyncRestApi:
             self._workflow_run_api = WorkflowRunApi(self.api_client)
         return self._workflow_run_api
 
+    @property
+    def worker_api(self):
+        if self._worker_api is None:
+            self._worker_api = WorkerApi(self.api_client)
+
+        return self._worker_api
+
     @property
     def step_run_api(self):
         if self._step_run_api is None:

hatchet_sdk/v0/worker/action_listener_process.py

@@ -8,12 +8,14 @@ from typing import Any, List, Mapping, Optional
 
 import grpc
 
+from hatchet_sdk.v0.client import Client, new_client_raw
 from hatchet_sdk.v0.clients.dispatcher.action_listener import Action
 from hatchet_sdk.v0.clients.dispatcher.dispatcher import (
     ActionListener,
     GetActionListenerRequest,
     new_dispatcher,
 )
+from hatchet_sdk.v0.clients.rest.models.update_worker_request import UpdateWorkerRequest
 from hatchet_sdk.v0.contracts.dispatcher_pb2 import (
     GROUP_KEY_EVENT_TYPE_STARTED,
     STEP_EVENT_TYPE_STARTED,
@@ -70,9 +72,15 @@ class WorkerActionListenerProcess:
         if self.debug:
             logger.setLevel(logging.DEBUG)
 
+        self.client = new_client_raw(self.config, self.debug)
+
         loop = asyncio.get_event_loop()
-        loop.add_signal_handler(signal.SIGINT, noop_handler)
-        loop.add_signal_handler(signal.SIGTERM, noop_handler)
+        loop.add_signal_handler(
+            signal.SIGINT, lambda: asyncio.create_task(self.pause_task_assignment())
+        )
+        loop.add_signal_handler(
+            signal.SIGTERM, lambda: asyncio.create_task(self.pause_task_assignment())
+        )
         loop.add_signal_handler(
             signal.SIGQUIT, lambda: asyncio.create_task(self.exit_gracefully())
         )
@@ -249,7 +257,15 @@ 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):
+        await self.pause_task_assignment()
+
         if self.killing:
             return
 

hatchet_sdk/waits.py

@@ -0,0 +1,120 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+from hatchet_sdk.contracts.v1.shared.condition_pb2 import Action as ProtoAction
+from hatchet_sdk.contracts.v1.shared.condition_pb2 import (
+    BaseMatchCondition,
+    ParentOverrideMatchCondition,
+    SleepMatchCondition,
+    UserEventMatchCondition,
+)
+from hatchet_sdk.utils.proto_enums import convert_python_enum_to_proto
+from hatchet_sdk.utils.timedelta_to_expression import Duration, timedelta_to_expr
+
+if TYPE_CHECKING:
+    from hatchet_sdk.runnables.task import Task
+    from hatchet_sdk.runnables.types import R, TWorkflowInput
+
+
+def generate_or_group_id() -> str:
+    return str(uuid4())
+
+
+class Action(Enum):
+    CREATE = 0
+    QUEUE = 1
+    CANCEL = 2
+    SKIP = 3
+
+
+class BaseCondition(BaseModel):
+    readable_data_key: str
+    action: Action | None = None
+    or_group_id: str = Field(default_factory=generate_or_group_id)
+    expression: str | None = None
+
+    def to_pb(self) -> BaseMatchCondition:
+        return BaseMatchCondition(
+            readable_data_key=self.readable_data_key,
+            action=convert_python_enum_to_proto(self.action, ProtoAction),  # type: ignore[arg-type]
+            or_group_id=self.or_group_id,
+            expression=self.expression,
+        )
+
+
+class Condition(ABC):
+    def __init__(self, base: BaseCondition):
+        self.base = base
+
+    @abstractmethod
+    def to_pb(
+        self,
+    ) -> UserEventMatchCondition | ParentOverrideMatchCondition | SleepMatchCondition:
+        pass
+
+
+class SleepCondition(Condition):
+    def __init__(self, duration: Duration) -> None:
+        super().__init__(
+            BaseCondition(
+                readable_data_key=f"sleep:{timedelta_to_expr(duration)}",
+            )
+        )
+
+        self.duration = duration
+
+    def to_pb(self) -> SleepMatchCondition:
+        return SleepMatchCondition(
+            base=self.base.to_pb(),
+            sleep_for=timedelta_to_expr(self.duration),
+        )
+
+
+class UserEventCondition(Condition):
+    def __init__(self, event_key: str, expression: str | None = None) -> None:
+        super().__init__(
+            BaseCondition(
+                readable_data_key=event_key,
+                expression=expression,
+            )
+        )
+
+        self.event_key = event_key
+        self.expression = expression
+
+    def to_pb(self) -> UserEventMatchCondition:
+        return UserEventMatchCondition(
+            base=self.base.to_pb(),
+            user_event_key=self.event_key,
+        )
+
+
+class ParentCondition(Condition):
+    def __init__(
+        self, parent: "Task[TWorkflowInput, R]", expression: str | None = None
+    ) -> None:
+        super().__init__(
+            BaseCondition(readable_data_key=parent.name, expression=expression)
+        )
+
+        self.parent = parent
+
+    def to_pb(self) -> ParentOverrideMatchCondition:
+        return ParentOverrideMatchCondition(
+            base=self.base.to_pb(),
+            parent_readable_id=self.parent.name,
+        )
+
+
+class OrGroup:
+    def __init__(self, conditions: list[Condition]) -> None:
+        self.or_group_id = generate_or_group_id()
+        self.conditions = conditions
+
+
+def or_(*conditions: Condition) -> OrGroup:
+    return OrGroup(conditions=list(conditions))

hatchet_sdk/worker/action_listener_process.py

@@ -2,12 +2,13 @@ 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, Literal
+from typing import Any, Literal
 
 import grpc
 
+from hatchet_sdk.client import Client
 from hatchet_sdk.clients.dispatcher.action_listener import (
     Action,
     ActionListener,
@@ -15,12 +16,18 @@ from hatchet_sdk.clients.dispatcher.action_listener import (
     GetActionListenerRequest,
 )
 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,
 )
 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
@@ -42,42 +49,60 @@ BLOCKED_THREAD_WARNING = (
 )
 
 
-def noop_handler() -> None:
-    pass
-
-
-@dataclass
 class WorkerActionListenerProcess:
-    name: str
-    actions: List[str]
-    max_runs: 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] = field(default_factory=dict)
-
-    listener: ActionListener = field(init=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] = {}
 
-    killing: bool = field(init=False, default=False)
-
-    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: dict[str, float] = field(init=False, default_factory=dict)
-
-    def __post_init__(self) -> None:
         if self.debug:
             logger.setLevel(logging.DEBUG)
 
+        self.client = Client(config=self.config, debug=self.debug)
+
         loop = asyncio.get_event_loop()
-        loop.add_signal_handler(signal.SIGINT, noop_handler)
-        loop.add_signal_handler(signal.SIGTERM, noop_handler)
+        loop.add_signal_handler(
+            signal.SIGINT, lambda: asyncio.create_task(self.pause_task_assignment())
+        )
+        loop.add_signal_handler(
+            signal.SIGTERM, lambda: asyncio.create_task(self.pause_task_assignment())
+        )
         loop.add_signal_handler(
             signal.SIGQUIT, lambda: asyncio.create_task(self.exit_gracefully())
         )
 
+    async def pause_task_assignment(self) -> None:
+        if self.listener is None:
+            raise ValueError("listener not started")
+
+        await self.client.workers.aio_update(
+            worker_id=self.listener.worker_id,
+            opts=UpdateWorkerRequest(isPaused=True),
+        )
+
     async def start(self, retry_attempt: int = 0) -> None:
         if retry_attempt > 5:
             logger.error("could not start action listener")
@@ -93,8 +118,8 @@ class WorkerActionListenerProcess:
                     worker_name=self.name,
                     services=["default"],
                     actions=self.actions,
-                    max_runs=self.max_runs,
-                    _labels=self.labels,
+                    slots=self.slots,
+                    raw_labels=self.labels,
                 )
             )
 
@@ -190,11 +215,18 @@ class WorkerActionListenerProcess:
         return time.time()
 
     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:
@@ -253,6 +285,8 @@ class WorkerActionListenerProcess:
         self.event_queue.put(STOP_LOOP)
 
     async def exit_gracefully(self) -> None:
+        await self.pause_task_assignment()
+
         if self.killing:
             return
 

hatchet_sdk/worker/runner/run_loop_manager.py

@@ -1,18 +1,17 @@
 import asyncio
 import logging
-from dataclasses import dataclass, field
 from multiprocessing import Queue
 from typing import Any, Literal, TypeVar
 
-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.config import ClientConfig
 from hatchet_sdk.logger import logger
+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
-from hatchet_sdk.workflow import Step
 
 STOP_LOOP_TYPE = Literal["STOP_LOOP"]
 STOP_LOOP: STOP_LOOP_TYPE = "STOP_LOOP"
@@ -20,29 +19,40 @@ STOP_LOOP: STOP_LOOP_TYPE = "STOP_LOOP"
 T = TypeVar("T")
 
 
-@dataclass
 class WorkerActionRunLoopManager:
-    name: str
-    action_registry: dict[str, Step[Any]]
-    validator_registry: dict[str, WorkflowValidator]
-    max_runs: 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] = field(default_factory=dict)
-
-    client: Client = field(init=False)
-
-    killing: bool = field(init=False, default=False)
-    runner: Runner | None = field(init=False, default=None)
-
-    def __post_init__(self) -> None:
+    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: int = 1) -> None:
@@ -73,13 +83,12 @@ class WorkerActionRunLoopManager:
 
     async def _start_action_loop(self) -> None:
         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,
         )