hatchet-sdk 1.18.1__py3-none-any.whl → 1.20.0__py3-none-any.whl

hatchet_sdk/v0/worker/worker.py

@@ -1,391 +0,0 @@
-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.contracts.workflows_pb2 import CreateWorkflowVersionOpts
-from hatchet_sdk.logger import logger
-from hatchet_sdk.v0 import Context
-from hatchet_sdk.v0.client import Client, new_client_raw
-from hatchet_sdk.v0.loader import ClientConfig
-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

@@ -1,261 +0,0 @@
-import functools
-from typing import (
-    Any,
-    Callable,
-    Protocol,
-    Type,
-    TypeVar,
-    Union,
-    cast,
-    get_type_hints,
-    runtime_checkable,
-)
-
-from pydantic import BaseModel
-
-from hatchet_sdk.contracts.workflows_pb2 import (
-    CreateWorkflowJobOpts,
-    CreateWorkflowStepOpts,
-    CreateWorkflowVersionOpts,
-    StickyStrategy,
-    WorkflowConcurrencyOpts,
-    WorkflowKind,
-)
-from hatchet_sdk.logger import logger
-from hatchet_sdk.v0 import ConcurrencyLimitStrategy
-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

@@ -1,59 +0,0 @@
-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