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

hatchet_sdk/v0/worker/worker.py

@@ -0,0 +1,391 @@
+import asyncio
+import multiprocessing
+import multiprocessing.context
+import os
+import signal
+import sys
+from concurrent.futures import Future
+from dataclasses import dataclass, field
+from enum import Enum
+from multiprocessing import Queue
+from multiprocessing.process import BaseProcess
+from types import FrameType
+from typing import Any, Callable, TypeVar, get_type_hints
+
+from aiohttp import web
+from aiohttp.web_request import Request
+from aiohttp.web_response import Response
+from prometheus_client import CONTENT_TYPE_LATEST, Gauge, generate_latest
+
+from hatchet_sdk.v0 import Context
+from hatchet_sdk.v0.client import Client, new_client_raw
+from hatchet_sdk.v0.contracts.workflows_pb2 import CreateWorkflowVersionOpts
+from hatchet_sdk.v0.loader import ClientConfig
+from hatchet_sdk.v0.logger import logger
+from hatchet_sdk.v0.utils.types import WorkflowValidator
+from hatchet_sdk.v0.utils.typing import is_basemodel_subclass
+from hatchet_sdk.v0.v2.callable import HatchetCallable
+from hatchet_sdk.v0.v2.concurrency import ConcurrencyFunction
+from hatchet_sdk.v0.worker.action_listener_process import worker_action_listener_process
+from hatchet_sdk.v0.worker.runner.run_loop_manager import WorkerActionRunLoopManager
+from hatchet_sdk.v0.workflow import WorkflowInterface
+
+T = TypeVar("T")
+
+
+class WorkerStatus(Enum):
+    INITIALIZED = 1
+    STARTING = 2
+    HEALTHY = 3
+    UNHEALTHY = 4
+
+
+@dataclass
+class WorkerStartOptions:
+    loop: asyncio.AbstractEventLoop | None = field(default=None)
+
+
+TWorkflow = TypeVar("TWorkflow", bound=object)
+
+
+class Worker:
+    def __init__(
+        self,
+        name: str,
+        config: ClientConfig = ClientConfig(),
+        max_runs: int | None = None,
+        labels: dict[str, str | int] = {},
+        debug: bool = False,
+        owned_loop: bool = True,
+        handle_kill: bool = True,
+    ) -> None:
+        self.name = name
+        self.config = config
+        self.max_runs = max_runs
+        self.debug = debug
+        self.labels = labels
+        self.handle_kill = handle_kill
+        self.owned_loop = owned_loop
+
+        self.client: Client
+
+        self.action_registry: dict[str, Callable[[Context], Any]] = {}
+        self.validator_registry: dict[str, WorkflowValidator] = {}
+
+        self.killing: bool = False
+        self._status: WorkerStatus
+
+        self.action_listener_process: BaseProcess
+        self.action_listener_health_check: asyncio.Task[Any]
+        self.action_runner: WorkerActionRunLoopManager
+
+        self.ctx = multiprocessing.get_context("spawn")
+
+        self.action_queue: "Queue[Any]" = self.ctx.Queue()
+        self.event_queue: "Queue[Any]" = self.ctx.Queue()
+
+        self.loop: asyncio.AbstractEventLoop
+
+        self.client = new_client_raw(self.config, self.debug)
+        self.name = self.client.config.namespace + self.name
+
+        self._setup_signal_handlers()
+
+        self.worker_status_gauge = Gauge(
+            "hatchet_worker_status", "Current status of the Hatchet worker"
+        )
+
+    def register_function(self, action: str, func: Callable[[Context], Any]) -> None:
+        self.action_registry[action] = func
+
+    def register_workflow_from_opts(
+        self, name: str, opts: CreateWorkflowVersionOpts
+    ) -> None:
+        try:
+            self.client.admin.put_workflow(opts.name, opts)
+        except Exception as e:
+            logger.error(f"failed to register workflow: {opts.name}")
+            logger.error(e)
+            sys.exit(1)
+
+    def register_workflow(self, workflow: TWorkflow) -> None:
+        ## Hack for typing
+        assert isinstance(workflow, WorkflowInterface)
+
+        namespace = self.client.config.namespace
+
+        try:
+            self.client.admin.put_workflow(
+                workflow.get_name(namespace), workflow.get_create_opts(namespace)
+            )
+        except Exception as e:
+            logger.error(f"failed to register workflow: {workflow.get_name(namespace)}")
+            logger.error(e)
+            sys.exit(1)
+
+        def create_action_function(
+            action_func: Callable[..., T]
+        ) -> Callable[[Context], T]:
+            def action_function(context: Context) -> T:
+                return action_func(workflow, context)
+
+            if asyncio.iscoroutinefunction(action_func):
+                setattr(action_function, "is_coroutine", True)
+            else:
+                setattr(action_function, "is_coroutine", False)
+
+            return action_function
+
+        for action_name, action_func in workflow.get_actions(namespace):
+            self.action_registry[action_name] = create_action_function(action_func)
+            return_type = get_type_hints(action_func).get("return")
+
+            self.validator_registry[action_name] = WorkflowValidator(
+                workflow_input=workflow.input_validator,
+                step_output=return_type if is_basemodel_subclass(return_type) else None,
+            )
+
+    def status(self) -> WorkerStatus:
+        return self._status
+
+    def setup_loop(self, loop: asyncio.AbstractEventLoop | None = None) -> bool:
+        try:
+            loop = loop or asyncio.get_running_loop()
+            self.loop = loop
+            created_loop = False
+            logger.debug("using existing event loop")
+            return created_loop
+        except RuntimeError:
+            self.loop = asyncio.new_event_loop()
+            logger.debug("creating new event loop")
+            asyncio.set_event_loop(self.loop)
+            created_loop = True
+            return created_loop
+
+    async def health_check_handler(self, request: Request) -> Response:
+        status = self.status()
+
+        return web.json_response({"status": status.name})
+
+    async def metrics_handler(self, request: Request) -> Response:
+        self.worker_status_gauge.set(1 if self.status() == WorkerStatus.HEALTHY else 0)
+
+        return web.Response(body=generate_latest(), content_type="text/plain")
+
+    async def start_health_server(self) -> None:
+        port = self.config.worker_healthcheck_port or 8001
+
+        app = web.Application()
+        app.add_routes(
+            [
+                web.get("/health", self.health_check_handler),
+                web.get("/metrics", self.metrics_handler),
+            ]
+        )
+
+        runner = web.AppRunner(app)
+
+        try:
+            await runner.setup()
+            await web.TCPSite(runner, "0.0.0.0", port).start()
+        except Exception as e:
+            logger.error("failed to start healthcheck server")
+            logger.error(str(e))
+            return
+
+        logger.info(f"healthcheck server running on port {port}")
+
+    def start(
+        self, options: WorkerStartOptions = WorkerStartOptions()
+    ) -> Future[asyncio.Task[None]]:
+        self.owned_loop = self.setup_loop(options.loop)
+
+        f = asyncio.run_coroutine_threadsafe(
+            self._async_start(options, _from_start=True), self.loop
+        )
+
+        # start the loop and wait until its closed
+        if self.owned_loop:
+            self.loop.run_forever()
+
+            if self.handle_kill:
+                sys.exit(0)
+
+        return f
+
+    ## Start methods
+    async def _async_start(
+        self,
+        options: WorkerStartOptions = WorkerStartOptions(),
+        _from_start: bool = False,
+    ) -> asyncio.Task[None]:
+        main_pid = os.getpid()
+        logger.info("------------------------------------------")
+        logger.info("STARTING HATCHET...")
+        logger.debug(f"worker runtime starting on PID: {main_pid}")
+
+        self._status = WorkerStatus.STARTING
+
+        if len(self.action_registry.keys()) == 0:
+            raise ValueError(
+                "no actions registered, register workflows or actions before starting worker"
+            )
+
+        # non blocking setup
+        if not _from_start:
+            self.setup_loop(options.loop)
+
+        if self.config.worker_healthcheck_enabled:
+            await self.start_health_server()
+
+        self.action_listener_process = self._start_listener()
+
+        self.action_runner = self._run_action_runner()
+
+        self.action_listener_health_check = self.loop.create_task(
+            self._check_listener_health()
+        )
+
+        return await self.action_listener_health_check
+
+    def _run_action_runner(self) -> WorkerActionRunLoopManager:
+        # Retrieve the shared queue
+        return WorkerActionRunLoopManager(
+            self.name,
+            self.action_registry,
+            self.validator_registry,
+            self.max_runs,
+            self.config,
+            self.action_queue,
+            self.event_queue,
+            self.loop,
+            self.handle_kill,
+            self.client.debug,
+            self.labels,
+        )
+
+    def _start_listener(self) -> multiprocessing.context.SpawnProcess:
+        action_list = [str(key) for key in self.action_registry.keys()]
+
+        try:
+            process = self.ctx.Process(
+                target=worker_action_listener_process,
+                args=(
+                    self.name,
+                    action_list,
+                    self.max_runs,
+                    self.config,
+                    self.action_queue,
+                    self.event_queue,
+                    self.handle_kill,
+                    self.client.debug,
+                    self.labels,
+                ),
+            )
+            process.start()
+            logger.debug(f"action listener starting on PID: {process.pid}")
+
+            return process
+        except Exception as e:
+            logger.error(f"failed to start action listener: {e}")
+            sys.exit(1)
+
+    async def _check_listener_health(self) -> None:
+        logger.debug("starting action listener health check...")
+        try:
+            while not self.killing:
+                if (
+                    self.action_listener_process is None
+                    or not self.action_listener_process.is_alive()
+                ):
+                    logger.debug("child action listener process killed...")
+                    self._status = WorkerStatus.UNHEALTHY
+                    if not self.killing:
+                        self.loop.create_task(self.exit_gracefully())
+                    break
+                else:
+                    self._status = WorkerStatus.HEALTHY
+                await asyncio.sleep(1)
+        except Exception as e:
+            logger.error(f"error checking listener health: {e}")
+
+    ## Cleanup methods
+    def _setup_signal_handlers(self) -> None:
+        signal.signal(signal.SIGTERM, self._handle_exit_signal)
+        signal.signal(signal.SIGINT, self._handle_exit_signal)
+        signal.signal(signal.SIGQUIT, self._handle_force_quit_signal)
+
+    def _handle_exit_signal(self, signum: int, frame: FrameType | None) -> None:
+        sig_name = "SIGTERM" if signum == signal.SIGTERM else "SIGINT"
+        logger.info(f"received signal {sig_name}...")
+        self.loop.create_task(self.exit_gracefully())
+
+    def _handle_force_quit_signal(self, signum: int, frame: FrameType | None) -> None:
+        logger.info("received SIGQUIT...")
+        self.exit_forcefully()
+
+    async def close(self) -> None:
+        logger.info(f"closing worker '{self.name}'...")
+        self.killing = True
+        # self.action_queue.close()
+        # self.event_queue.close()
+
+        if self.action_runner is not None:
+            self.action_runner.cleanup()
+
+        await self.action_listener_health_check
+
+    async def exit_gracefully(self) -> None:
+        logger.debug(f"gracefully stopping worker: {self.name}")
+
+        if self.killing:
+            return self.exit_forcefully()
+
+        self.killing = True
+
+        await self.action_runner.wait_for_tasks()
+
+        await self.action_runner.exit_gracefully()
+
+        if self.action_listener_process and self.action_listener_process.is_alive():
+            self.action_listener_process.kill()
+
+        await self.close()
+        if self.loop and self.owned_loop:
+            self.loop.stop()
+
+        logger.info("👋")
+
+    def exit_forcefully(self) -> None:
+        self.killing = True
+
+        logger.debug(f"forcefully stopping worker: {self.name}")
+
+        self.close()
+
+        if self.action_listener_process:
+            self.action_listener_process.kill()  # Forcefully kill the process
+
+        logger.info("👋")
+        sys.exit(
+            1
+        )  # Exit immediately TODO - should we exit with 1 here, there may be other workers to cleanup
+
+
+def register_on_worker(callable: HatchetCallable[T], worker: Worker) -> None:
+    worker.register_function(callable.get_action_name(), callable)
+
+    if callable.function_on_failure is not None:
+        worker.register_function(
+            callable.function_on_failure.get_action_name(), callable.function_on_failure
+        )
+
+    if callable.function_concurrency is not None:
+        worker.register_function(
+            callable.function_concurrency.get_action_name(),
+            callable.function_concurrency,
+        )
+
+    opts = callable.to_workflow_opts()
+
+    worker.register_workflow_from_opts(opts.name, opts)

hatchet_sdk/v0/workflow.py

@@ -0,0 +1,261 @@
+import functools
+from typing import (
+    Any,
+    Callable,
+    Protocol,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    get_type_hints,
+    runtime_checkable,
+)
+
+from pydantic import BaseModel
+
+from hatchet_sdk.v0 import ConcurrencyLimitStrategy
+from hatchet_sdk.v0.contracts.workflows_pb2 import (
+    CreateWorkflowJobOpts,
+    CreateWorkflowStepOpts,
+    CreateWorkflowVersionOpts,
+    StickyStrategy,
+    WorkflowConcurrencyOpts,
+    WorkflowKind,
+)
+from hatchet_sdk.v0.logger import logger
+from hatchet_sdk.v0.utils.typing import is_basemodel_subclass
+
+
+class WorkflowStepProtocol(Protocol):
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+    __name__: str
+
+    _step_name: str
+    _step_timeout: str | None
+    _step_parents: list[str]
+    _step_retries: int | None
+    _step_rate_limits: list[str] | None
+    _step_desired_worker_labels: dict[str, str]
+    _step_backoff_factor: float | None
+    _step_backoff_max_seconds: int | None
+
+    _concurrency_fn_name: str
+    _concurrency_max_runs: int | None
+    _concurrency_limit_strategy: str | None
+
+    _on_failure_step_name: str
+    _on_failure_step_timeout: str | None
+    _on_failure_step_retries: int
+    _on_failure_step_rate_limits: list[str] | None
+    _on_failure_step_backoff_factor: float | None
+    _on_failure_step_backoff_max_seconds: int | None
+
+
+StepsType = list[tuple[str, WorkflowStepProtocol]]
+
+T = TypeVar("T")
+TW = TypeVar("TW", bound="WorkflowInterface")
+
+
+class ConcurrencyExpression:
+    """
+    Defines concurrency limits for a workflow using a CEL expression.
+
+    Args:
+        expression (str): CEL expression to determine concurrency grouping. (i.e. "input.user_id")
+        max_runs (int): Maximum number of concurrent workflow runs.
+        limit_strategy (ConcurrencyLimitStrategy): Strategy for handling limit violations.
+
+    Example:
+        ConcurrencyExpression("input.user_id", 5, ConcurrencyLimitStrategy.CANCEL_IN_PROGRESS)
+    """
+
+    def __init__(
+        self, expression: str, max_runs: int, limit_strategy: ConcurrencyLimitStrategy
+    ):
+        self.expression = expression
+        self.max_runs = max_runs
+        self.limit_strategy = limit_strategy
+
+
+@runtime_checkable
+class WorkflowInterface(Protocol):
+    def get_name(self, namespace: str) -> str: ...
+
+    def get_actions(self, namespace: str) -> list[tuple[str, Callable[..., Any]]]: ...
+
+    def get_create_opts(self, namespace: str) -> Any: ...
+
+    on_events: list[str] | None
+    on_crons: list[str] | None
+    name: str
+    version: str
+    timeout: str
+    schedule_timeout: str
+    sticky: Union[StickyStrategy.Value, None]  # type: ignore[name-defined]
+    default_priority: int | None
+    concurrency_expression: ConcurrencyExpression | None
+    input_validator: Type[BaseModel] | None
+
+
+class WorkflowMeta(type):
+    def __new__(
+        cls: Type["WorkflowMeta"],
+        name: str,
+        bases: tuple[type, ...],
+        attrs: dict[str, Any],
+    ) -> "WorkflowMeta":
+        def _create_steps_actions_list(name: str) -> StepsType:
+            return [
+                (getattr(func, name), attrs.pop(func_name))
+                for func_name, func in list(attrs.items())
+                if hasattr(func, name)
+            ]
+
+        concurrencyActions = _create_steps_actions_list("_concurrency_fn_name")
+        steps = _create_steps_actions_list("_step_name")
+
+        onFailureSteps = _create_steps_actions_list("_on_failure_step_name")
+
+        # Define __init__ and get_step_order methods
+        original_init = attrs.get("__init__")  # Get the original __init__ if it exists
+
+        def __init__(self: TW, *args: Any, **kwargs: Any) -> None:
+            if original_init:
+                original_init(self, *args, **kwargs)  # Call original __init__
+
+        def get_service_name(namespace: str) -> str:
+            return f"{namespace}{name.lower()}"
+
+        @functools.cache
+        def get_actions(self: TW, namespace: str) -> StepsType:
+            serviceName = get_service_name(namespace)
+
+            func_actions = [
+                (serviceName + ":" + func_name, func) for func_name, func in steps
+            ]
+            concurrency_actions = [
+                (serviceName + ":" + func_name, func)
+                for func_name, func in concurrencyActions
+            ]
+            onFailure_actions = [
+                (serviceName + ":" + func_name, func)
+                for func_name, func in onFailureSteps
+            ]
+
+            return func_actions + concurrency_actions + onFailure_actions
+
+        # Add these methods and steps to class attributes
+        attrs["__init__"] = __init__
+        attrs["get_actions"] = get_actions
+
+        for step_name, step_func in steps:
+            attrs[step_name] = step_func
+
+        def get_name(self: TW, namespace: str) -> str:
+            return namespace + cast(str, attrs["name"])
+
+        attrs["get_name"] = get_name
+
+        cron_triggers = attrs["on_crons"]
+        version = attrs["version"]
+        schedule_timeout = attrs["schedule_timeout"]
+        sticky = attrs["sticky"]
+        default_priority = attrs["default_priority"]
+
+        @functools.cache
+        def get_create_opts(self: TW, namespace: str) -> CreateWorkflowVersionOpts:
+            serviceName = get_service_name(namespace)
+            name = self.get_name(namespace)
+            event_triggers = [namespace + event for event in attrs["on_events"]]
+            createStepOpts: list[CreateWorkflowStepOpts] = [
+                CreateWorkflowStepOpts(
+                    readable_id=step_name,
+                    action=serviceName + ":" + step_name,
+                    timeout=func._step_timeout or "60s",
+                    inputs="{}",
+                    parents=[x for x in func._step_parents],
+                    retries=func._step_retries,
+                    rate_limits=func._step_rate_limits,  # type: ignore[arg-type]
+                    worker_labels=func._step_desired_worker_labels,  # type: ignore[arg-type]
+                    backoff_factor=func._step_backoff_factor,
+                    backoff_max_seconds=func._step_backoff_max_seconds,
+                )
+                for step_name, func in steps
+            ]
+
+            concurrency: WorkflowConcurrencyOpts | None = None
+
+            if len(concurrencyActions) > 0:
+                action = concurrencyActions[0]
+
+                concurrency = WorkflowConcurrencyOpts(
+                    action=serviceName + ":" + action[0],
+                    max_runs=action[1]._concurrency_max_runs,
+                    limit_strategy=action[1]._concurrency_limit_strategy,
+                )
+
+            if self.concurrency_expression:
+                concurrency = WorkflowConcurrencyOpts(
+                    expression=self.concurrency_expression.expression,
+                    max_runs=self.concurrency_expression.max_runs,
+                    limit_strategy=self.concurrency_expression.limit_strategy,
+                )
+
+            if len(concurrencyActions) > 0 and self.concurrency_expression:
+                raise ValueError(
+                    "Error: Both concurrencyActions and concurrency_expression are defined. Please use only one concurrency configuration method."
+                )
+
+            on_failure_job: CreateWorkflowJobOpts | None = None
+
+            if len(onFailureSteps) > 0:
+                func_name, func = onFailureSteps[0]
+                on_failure_job = CreateWorkflowJobOpts(
+                    name=name + "-on-failure",
+                    steps=[
+                        CreateWorkflowStepOpts(
+                            readable_id=func_name,
+                            action=serviceName + ":" + func_name,
+                            timeout=func._on_failure_step_timeout or "60s",
+                            inputs="{}",
+                            parents=[],
+                            retries=func._on_failure_step_retries,
+                            rate_limits=func._on_failure_step_rate_limits,  # type: ignore[arg-type]
+                            backoff_factor=func._on_failure_step_backoff_factor,
+                            backoff_max_seconds=func._on_failure_step_backoff_max_seconds,
+                        )
+                    ],
+                )
+
+            validated_priority = (
+                max(1, min(3, default_priority)) if default_priority else None
+            )
+            if validated_priority != default_priority:
+                logger.warning(
+                    "Warning: Default Priority Must be between 1 and 3 -- inclusively. Adjusted to be within the range."
+                )
+
+            return CreateWorkflowVersionOpts(
+                name=name,
+                kind=WorkflowKind.DAG,
+                version=version,
+                event_triggers=event_triggers,
+                cron_triggers=cron_triggers,
+                schedule_timeout=schedule_timeout,
+                sticky=sticky,
+                jobs=[
+                    CreateWorkflowJobOpts(
+                        name=name,
+                        steps=createStepOpts,
+                    )
+                ],
+                on_failure_job=on_failure_job,
+                concurrency=concurrency,
+                default_priority=validated_priority,
+            )
+
+        attrs["get_create_opts"] = get_create_opts
+
+        return super(WorkflowMeta, cls).__new__(cls, name, bases, attrs)

hatchet_sdk/v0/workflow_run.py

@@ -0,0 +1,59 @@
+import asyncio
+from typing import Any, Coroutine, Generic, Optional, TypedDict, TypeVar
+
+from hatchet_sdk.v0.clients.run_event_listener import (
+    RunEventListener,
+    RunEventListenerClient,
+)
+from hatchet_sdk.v0.clients.workflow_listener import PooledWorkflowRunListener
+from hatchet_sdk.v0.utils.aio_utils import EventLoopThread, get_active_event_loop
+
+
+class WorkflowRunRef:
+    workflow_run_id: str
+
+    def __init__(
+        self,
+        workflow_run_id: str,
+        workflow_listener: PooledWorkflowRunListener,
+        workflow_run_event_listener: RunEventListenerClient,
+    ):
+        self.workflow_run_id = workflow_run_id
+        self.workflow_listener = workflow_listener
+        self.workflow_run_event_listener = workflow_run_event_listener
+
+    def __str__(self):
+        return self.workflow_run_id
+
+    def stream(self) -> RunEventListener:
+        return self.workflow_run_event_listener.stream(self.workflow_run_id)
+
+    def result(self) -> Coroutine:
+        return self.workflow_listener.result(self.workflow_run_id)
+
+    def sync_result(self) -> dict:
+        coro = self.workflow_listener.result(self.workflow_run_id)
+        loop = get_active_event_loop()
+
+        if loop is None:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            try:
+                return loop.run_until_complete(coro)
+            finally:
+                asyncio.set_event_loop(None)
+        else:
+            return loop.run_until_complete(coro)
+
+
+T = TypeVar("T")
+
+
+class RunRef(WorkflowRunRef, Generic[T]):
+    async def result(self) -> T:
+        res = await self.workflow_listener.result(self.workflow_run_id)
+
+        if len(res) == 1:
+            return list(res.values())[0]
+
+        return res