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

hatchet_sdk/v0/clients/run_event_listener.py

@@ -1,260 +0,0 @@
-import asyncio
-import json
-from typing import AsyncGenerator
-
-import grpc
-
-from hatchet_sdk.contracts.dispatcher_pb2 import (
-    RESOURCE_TYPE_STEP_RUN,
-    RESOURCE_TYPE_WORKFLOW_RUN,
-    ResourceEventType,
-    SubscribeToWorkflowEventsRequest,
-    WorkflowEvent,
-)
-from hatchet_sdk.contracts.dispatcher_pb2_grpc import DispatcherStub
-from hatchet_sdk.v0.connection import new_conn
-
-from ..loader import ClientConfig
-from ..metadata import get_metadata
-
-DEFAULT_ACTION_LISTENER_RETRY_INTERVAL = 5  # seconds
-DEFAULT_ACTION_LISTENER_RETRY_COUNT = 5
-
-
-class StepRunEventType:
-    STEP_RUN_EVENT_TYPE_STARTED = "STEP_RUN_EVENT_TYPE_STARTED"
-    STEP_RUN_EVENT_TYPE_COMPLETED = "STEP_RUN_EVENT_TYPE_COMPLETED"
-    STEP_RUN_EVENT_TYPE_FAILED = "STEP_RUN_EVENT_TYPE_FAILED"
-    STEP_RUN_EVENT_TYPE_CANCELLED = "STEP_RUN_EVENT_TYPE_CANCELLED"
-    STEP_RUN_EVENT_TYPE_TIMED_OUT = "STEP_RUN_EVENT_TYPE_TIMED_OUT"
-    STEP_RUN_EVENT_TYPE_STREAM = "STEP_RUN_EVENT_TYPE_STREAM"
-
-
-class WorkflowRunEventType:
-    WORKFLOW_RUN_EVENT_TYPE_STARTED = "WORKFLOW_RUN_EVENT_TYPE_STARTED"
-    WORKFLOW_RUN_EVENT_TYPE_COMPLETED = "WORKFLOW_RUN_EVENT_TYPE_COMPLETED"
-    WORKFLOW_RUN_EVENT_TYPE_FAILED = "WORKFLOW_RUN_EVENT_TYPE_FAILED"
-    WORKFLOW_RUN_EVENT_TYPE_CANCELLED = "WORKFLOW_RUN_EVENT_TYPE_CANCELLED"
-    WORKFLOW_RUN_EVENT_TYPE_TIMED_OUT = "WORKFLOW_RUN_EVENT_TYPE_TIMED_OUT"
-
-
-step_run_event_type_mapping = {
-    ResourceEventType.RESOURCE_EVENT_TYPE_STARTED: StepRunEventType.STEP_RUN_EVENT_TYPE_STARTED,
-    ResourceEventType.RESOURCE_EVENT_TYPE_COMPLETED: StepRunEventType.STEP_RUN_EVENT_TYPE_COMPLETED,
-    ResourceEventType.RESOURCE_EVENT_TYPE_FAILED: StepRunEventType.STEP_RUN_EVENT_TYPE_FAILED,
-    ResourceEventType.RESOURCE_EVENT_TYPE_CANCELLED: StepRunEventType.STEP_RUN_EVENT_TYPE_CANCELLED,
-    ResourceEventType.RESOURCE_EVENT_TYPE_TIMED_OUT: StepRunEventType.STEP_RUN_EVENT_TYPE_TIMED_OUT,
-    ResourceEventType.RESOURCE_EVENT_TYPE_STREAM: StepRunEventType.STEP_RUN_EVENT_TYPE_STREAM,
-}
-
-workflow_run_event_type_mapping = {
-    ResourceEventType.RESOURCE_EVENT_TYPE_STARTED: WorkflowRunEventType.WORKFLOW_RUN_EVENT_TYPE_STARTED,
-    ResourceEventType.RESOURCE_EVENT_TYPE_COMPLETED: WorkflowRunEventType.WORKFLOW_RUN_EVENT_TYPE_COMPLETED,
-    ResourceEventType.RESOURCE_EVENT_TYPE_FAILED: WorkflowRunEventType.WORKFLOW_RUN_EVENT_TYPE_FAILED,
-    ResourceEventType.RESOURCE_EVENT_TYPE_CANCELLED: WorkflowRunEventType.WORKFLOW_RUN_EVENT_TYPE_CANCELLED,
-    ResourceEventType.RESOURCE_EVENT_TYPE_TIMED_OUT: WorkflowRunEventType.WORKFLOW_RUN_EVENT_TYPE_TIMED_OUT,
-}
-
-
-class StepRunEvent:
-    def __init__(self, type: StepRunEventType, payload: str):
-        self.type = type
-        self.payload = payload
-
-
-def new_listener(config: ClientConfig):
-    return RunEventListenerClient(config=config)
-
-
-class RunEventListener:
-
-    workflow_run_id: str = None
-    additional_meta_kv: tuple[str, str] = None
-
-    def __init__(self, client: DispatcherStub, token: str):
-        self.client = client
-        self.stop_signal = False
-        self.token = token
-
-    @classmethod
-    def for_run_id(cls, workflow_run_id: str, client: DispatcherStub, token: str):
-        listener = RunEventListener(client, token)
-        listener.workflow_run_id = workflow_run_id
-        return listener
-
-    @classmethod
-    def for_additional_meta(
-        cls, key: str, value: str, client: DispatcherStub, token: str
-    ):
-        listener = RunEventListener(client, token)
-        listener.additional_meta_kv = (key, value)
-        return listener
-
-    def abort(self):
-        self.stop_signal = True
-
-    def __aiter__(self):
-        return self._generator()
-
-    async def __anext__(self):
-        return await self._generator().__anext__()
-
-    def __iter__(self):
-        try:
-            loop = asyncio.get_event_loop()
-        except RuntimeError as e:
-            if str(e).startswith("There is no current event loop in thread"):
-                loop = asyncio.new_event_loop()
-                asyncio.set_event_loop(loop)
-            else:
-                raise e
-
-        async_iter = self.__aiter__()
-
-        while True:
-            try:
-                future = asyncio.ensure_future(async_iter.__anext__())
-                yield loop.run_until_complete(future)
-            except StopAsyncIteration:
-                break
-            except Exception as e:
-                print(f"Error in synchronous iterator: {e}")
-                break
-
-    async def _generator(self) -> AsyncGenerator[StepRunEvent, None]:
-        while True:
-            if self.stop_signal:
-                listener = None
-                break
-
-            listener = await self.retry_subscribe()
-
-            try:
-                async for workflow_event in listener:
-                    eventType = None
-                    if workflow_event.resourceType == RESOURCE_TYPE_STEP_RUN:
-                        if workflow_event.eventType in step_run_event_type_mapping:
-                            eventType = step_run_event_type_mapping[
-                                workflow_event.eventType
-                            ]
-                        else:
-                            raise Exception(
-                                f"Unknown event type: {workflow_event.eventType}"
-                            )
-                        payload = None
-
-                        try:
-                            if workflow_event.eventPayload:
-                                payload = json.loads(workflow_event.eventPayload)
-                        except Exception as e:
-                            payload = workflow_event.eventPayload
-                            pass
-
-                        yield StepRunEvent(type=eventType, payload=payload)
-                    elif workflow_event.resourceType == RESOURCE_TYPE_WORKFLOW_RUN:
-                        if workflow_event.eventType in workflow_run_event_type_mapping:
-                            eventType = workflow_run_event_type_mapping[
-                                workflow_event.eventType
-                            ]
-                        else:
-                            raise Exception(
-                                f"Unknown event type: {workflow_event.eventType}"
-                            )
-
-                        payload = None
-
-                        try:
-                            if workflow_event.eventPayload:
-                                payload = json.loads(workflow_event.eventPayload)
-                        except Exception as e:
-                            pass
-
-                        yield StepRunEvent(type=eventType, payload=payload)
-
-                    if workflow_event.hangup:
-                        listener = None
-                        break
-
-                break
-            except grpc.RpcError as e:
-                # Handle different types of errors
-                if e.code() == grpc.StatusCode.CANCELLED:
-                    # Context cancelled, unsubscribe and close
-                    break
-                elif e.code() == grpc.StatusCode.UNAVAILABLE:
-                    # Retry logic
-                    # logger.info("Could not connect to Hatchet, retrying...")
-                    listener = await self.retry_subscribe()
-                elif e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
-                    # logger.info("Deadline exceeded, retrying subscription")
-                    continue
-                else:
-                    # Unknown error, report and break
-                    # logger.error(f"Failed to receive message: {e}")
-                    break
-                # Raise StopAsyncIteration to properly end the generator
-
-    async def retry_subscribe(self):
-        retries = 0
-
-        while retries < DEFAULT_ACTION_LISTENER_RETRY_COUNT:
-            try:
-                if retries > 0:
-                    await asyncio.sleep(DEFAULT_ACTION_LISTENER_RETRY_INTERVAL)
-
-                if self.workflow_run_id is not None:
-                    return self.client.SubscribeToWorkflowEvents(
-                        SubscribeToWorkflowEventsRequest(
-                            workflowRunId=self.workflow_run_id,
-                        ),
-                        metadata=get_metadata(self.token),
-                    )
-                elif self.additional_meta_kv is not None:
-                    return self.client.SubscribeToWorkflowEvents(
-                        SubscribeToWorkflowEventsRequest(
-                            additionalMetaKey=self.additional_meta_kv[0],
-                            additionalMetaValue=self.additional_meta_kv[1],
-                        ),
-                        metadata=get_metadata(self.token),
-                    )
-                else:
-                    raise Exception("no listener method provided")
-
-            except grpc.RpcError as e:
-                if e.code() == grpc.StatusCode.UNAVAILABLE:
-                    retries = retries + 1
-                else:
-                    raise ValueError(f"gRPC error: {e}")
-
-
-class RunEventListenerClient:
-    def __init__(self, config: ClientConfig):
-        self.token = config.token
-        self.config = config
-        self.client: DispatcherStub = None
-
-    def stream_by_run_id(self, workflow_run_id: str):
-        return self.stream(workflow_run_id)
-
-    def stream(self, workflow_run_id: str):
-        if not isinstance(workflow_run_id, str) and hasattr(workflow_run_id, "__str__"):
-            workflow_run_id = str(workflow_run_id)
-
-        if not self.client:
-            aio_conn = new_conn(self.config, True)
-            self.client = DispatcherStub(aio_conn)
-
-        return RunEventListener.for_run_id(workflow_run_id, self.client, self.token)
-
-    def stream_by_additional_metadata(self, key: str, value: str):
-        if not self.client:
-            aio_conn = new_conn(self.config, True)
-            self.client = DispatcherStub(aio_conn)
-
-        return RunEventListener.for_additional_meta(key, value, self.client, self.token)
-
-    async def on(self, workflow_run_id: str, handler: callable = None):
-        async for event in self.stream(workflow_run_id):
-            # call the handler if provided
-            if handler:
-                handler(event)

hatchet_sdk/v0/clients/workflow_listener.py

@@ -1,277 +0,0 @@
-import asyncio
-import json
-from collections.abc import AsyncIterator
-from typing import AsyncGenerator
-
-import grpc
-from grpc._cython import cygrpc
-
-from hatchet_sdk.contracts.dispatcher_pb2 import (
-    SubscribeToWorkflowRunsRequest,
-    WorkflowRunEvent,
-)
-from hatchet_sdk.contracts.dispatcher_pb2_grpc import DispatcherStub
-from hatchet_sdk.v0.clients.event_ts import ThreadSafeEvent, read_with_interrupt
-from hatchet_sdk.v0.connection import new_conn
-
-from ...logger import logger
-from ..loader import ClientConfig
-from ..metadata import get_metadata
-
-DEFAULT_WORKFLOW_LISTENER_RETRY_INTERVAL = 3  # seconds
-DEFAULT_WORKFLOW_LISTENER_RETRY_COUNT = 5
-DEFAULT_WORKFLOW_LISTENER_INTERRUPT_INTERVAL = 1800  # 30 minutes
-
-DEDUPE_MESSAGE = "DUPLICATE_WORKFLOW_RUN"
-
-
-class _Subscription:
-    def __init__(self, id: int, workflow_run_id: str):
-        self.id = id
-        self.workflow_run_id = workflow_run_id
-        self.queue: asyncio.Queue[WorkflowRunEvent | None] = asyncio.Queue()
-
-    async def __aiter__(self):
-        return self
-
-    async def __anext__(self) -> WorkflowRunEvent:
-        return await self.queue.get()
-
-    async def get(self) -> WorkflowRunEvent:
-        event = await self.queue.get()
-
-        if event is None:
-            raise StopAsyncIteration
-
-        return event
-
-    async def put(self, item: WorkflowRunEvent):
-        await self.queue.put(item)
-
-    async def close(self):
-        await self.queue.put(None)
-
-
-class PooledWorkflowRunListener:
-    # list of all active subscriptions, mapping from a subscription id to a workflow run id
-    subscriptionsToWorkflows: dict[int, str] = {}
-
-    # list of workflow run ids mapped to an array of subscription ids
-    workflowsToSubscriptions: dict[str, list[int]] = {}
-
-    subscription_counter: int = 0
-    subscription_counter_lock: asyncio.Lock = asyncio.Lock()
-
-    requests: asyncio.Queue[SubscribeToWorkflowRunsRequest] = asyncio.Queue()
-
-    listener: AsyncGenerator[WorkflowRunEvent, None] = None
-    listener_task: asyncio.Task = None
-
-    curr_requester: int = 0
-
-    # events have keys of the format workflow_run_id + subscription_id
-    events: dict[int, _Subscription] = {}
-
-    interrupter: asyncio.Task = None
-
-    def __init__(self, config: ClientConfig):
-        try:
-            asyncio.get_running_loop()
-        except RuntimeError:
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-
-        conn = new_conn(config, True)
-        self.client = DispatcherStub(conn)
-        self.token = config.token
-        self.config = config
-
-    async def _interrupter(self):
-        """
-        _interrupter runs in a separate thread and interrupts the listener according to a configurable duration.
-        """
-        await asyncio.sleep(DEFAULT_WORKFLOW_LISTENER_INTERRUPT_INTERVAL)
-
-        if self.interrupt is not None:
-            self.interrupt.set()
-
-    async def _init_producer(self):
-        try:
-            if not self.listener:
-                while True:
-                    try:
-                        self.listener = await self._retry_subscribe()
-
-                        logger.debug("Workflow run listener connected.")
-
-                        # spawn an interrupter task
-                        if self.interrupter is not None and not self.interrupter.done():
-                            self.interrupter.cancel()
-
-                        self.interrupter = asyncio.create_task(self._interrupter())
-
-                        while True:
-                            self.interrupt = ThreadSafeEvent()
-                            t = asyncio.create_task(
-                                read_with_interrupt(self.listener, self.interrupt)
-                            )
-                            await self.interrupt.wait()
-
-                            if not t.done():
-                                # print a warning
-                                logger.warning(
-                                    "Interrupted read_with_interrupt task of workflow run listener"
-                                )
-
-                                t.cancel()
-                                self.listener.cancel()
-                                await asyncio.sleep(
-                                    DEFAULT_WORKFLOW_LISTENER_RETRY_INTERVAL
-                                )
-                                break
-
-                            workflow_event: WorkflowRunEvent = t.result()
-
-                            if workflow_event is cygrpc.EOF:
-                                break
-
-                            # get a list of subscriptions for this workflow
-                            subscriptions = self.workflowsToSubscriptions.get(
-                                workflow_event.workflowRunId, []
-                            )
-
-                            for subscription_id in subscriptions:
-                                await self.events[subscription_id].put(workflow_event)
-
-                    except grpc.RpcError as e:
-                        logger.debug(f"grpc error in workflow run listener: {e}")
-                        await asyncio.sleep(DEFAULT_WORKFLOW_LISTENER_RETRY_INTERVAL)
-                        continue
-
-        except Exception as e:
-            logger.error(f"Error in workflow run listener: {e}")
-
-            self.listener = None
-
-            # close all subscriptions
-            for subscription_id in self.events:
-                await self.events[subscription_id].close()
-
-            raise e
-
-    async def _request(self) -> AsyncIterator[SubscribeToWorkflowRunsRequest]:
-        self.curr_requester = self.curr_requester + 1
-
-        # replay all existing subscriptions
-        workflow_run_set = set(self.subscriptionsToWorkflows.values())
-
-        for workflow_run_id in workflow_run_set:
-            yield SubscribeToWorkflowRunsRequest(
-                workflowRunId=workflow_run_id,
-            )
-
-        while True:
-            request = await self.requests.get()
-
-            # if the request is an int which matches the current requester, then we should stop
-            if request == self.curr_requester:
-                break
-
-            # if we've gotten an int that doesn't match the current requester, then we should ignore it
-            if isinstance(request, int):
-                continue
-
-            yield request
-            self.requests.task_done()
-
-    def cleanup_subscription(self, subscription_id: int):
-        workflow_run_id = self.subscriptionsToWorkflows[subscription_id]
-
-        if workflow_run_id in self.workflowsToSubscriptions:
-            self.workflowsToSubscriptions[workflow_run_id].remove(subscription_id)
-
-        del self.subscriptionsToWorkflows[subscription_id]
-        del self.events[subscription_id]
-
-    async def subscribe(self, workflow_run_id: str):
-        init_producer: asyncio.Task = None
-        try:
-            # create a new subscription id, place a mutex on the counter
-            await self.subscription_counter_lock.acquire()
-            self.subscription_counter += 1
-            subscription_id = self.subscription_counter
-            self.subscription_counter_lock.release()
-
-            self.subscriptionsToWorkflows[subscription_id] = workflow_run_id
-
-            if workflow_run_id not in self.workflowsToSubscriptions:
-                self.workflowsToSubscriptions[workflow_run_id] = [subscription_id]
-            else:
-                self.workflowsToSubscriptions[workflow_run_id].append(subscription_id)
-
-            self.events[subscription_id] = _Subscription(
-                subscription_id, workflow_run_id
-            )
-
-            await self.requests.put(
-                SubscribeToWorkflowRunsRequest(
-                    workflowRunId=workflow_run_id,
-                )
-            )
-
-            if not self.listener_task or self.listener_task.done():
-                self.listener_task = asyncio.create_task(self._init_producer())
-
-            event = await self.events[subscription_id].get()
-
-            return event
-        except asyncio.CancelledError:
-            raise
-        finally:
-            self.cleanup_subscription(subscription_id)
-
-    async def result(self, workflow_run_id: str):
-        from hatchet_sdk.v0.clients.admin import DedupeViolationErr
-
-        event = await self.subscribe(workflow_run_id)
-
-        errors = []
-
-        if event.results:
-            errors = [result.error for result in event.results if result.error]
-
-        if errors:
-            if DEDUPE_MESSAGE in errors[0]:
-                raise DedupeViolationErr(errors[0])
-            else:
-                raise Exception(f"Workflow Errors: {errors}")
-
-        results = {
-            result.stepReadableId: json.loads(result.output)
-            for result in event.results
-            if result.output
-        }
-
-        return results
-
-    async def _retry_subscribe(self):
-        retries = 0
-
-        while retries < DEFAULT_WORKFLOW_LISTENER_RETRY_COUNT:
-            try:
-                if retries > 0:
-                    await asyncio.sleep(DEFAULT_WORKFLOW_LISTENER_RETRY_INTERVAL)
-
-                # signal previous async iterator to stop
-                if self.curr_requester != 0:
-                    self.requests.put_nowait(self.curr_requester)
-
-                return self.client.SubscribeToWorkflowRuns(
-                    self._request(),
-                    metadata=get_metadata(self.token),
-                )
-            except grpc.RpcError as e:
-                if e.code() == grpc.StatusCode.UNAVAILABLE:
-                    retries = retries + 1
-                else:
-                    raise ValueError(f"gRPC error: {e}")

hatchet_sdk/v0/connection.py

@@ -1,63 +0,0 @@
-import os
-from typing import TYPE_CHECKING
-
-import grpc
-
-if TYPE_CHECKING:
-    from hatchet_sdk.v0.loader import ClientConfig
-
-
-def new_conn(config: "ClientConfig", aio=False):
-    credentials: grpc.ChannelCredentials | None = None
-
-    # load channel credentials
-    if config.tls_config.tls_strategy == "tls":
-        root: bytes | None = None
-
-        if config.tls_config.ca_file:
-            root = open(config.tls_config.ca_file, "rb").read()
-
-        credentials = grpc.ssl_channel_credentials(root_certificates=root)
-    elif config.tls_config.tls_strategy == "mtls":
-        root = open(config.tls_config.ca_file, "rb").read()
-        private_key = open(config.tls_config.key_file, "rb").read()
-        certificate_chain = open(config.tls_config.cert_file, "rb").read()
-
-        credentials = grpc.ssl_channel_credentials(
-            root_certificates=root,
-            private_key=private_key,
-            certificate_chain=certificate_chain,
-        )
-
-    start = grpc if not aio else grpc.aio
-
-    channel_options = [
-        ("grpc.max_send_message_length", config.grpc_max_send_message_length),
-        ("grpc.max_receive_message_length", config.grpc_max_recv_message_length),
-        ("grpc.keepalive_time_ms", 10 * 1000),
-        ("grpc.keepalive_timeout_ms", 60 * 1000),
-        ("grpc.client_idle_timeout_ms", 60 * 1000),
-        ("grpc.http2.max_pings_without_data", 0),
-        ("grpc.keepalive_permit_without_calls", 1),
-    ]
-
-    # Set environment variable to disable fork support. Reference: https://github.com/grpc/grpc/issues/28557
-    # When steps execute via os.fork, we see `TSI_DATA_CORRUPTED` errors.
-    os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "False"
-
-    if config.tls_config.tls_strategy == "none":
-        conn = start.insecure_channel(
-            target=config.host_port,
-            options=channel_options,
-        )
-    else:
-        channel_options.append(
-            ("grpc.ssl_target_name_override", config.tls_config.server_name)
-        )
-
-        conn = start.secure_channel(
-            target=config.host_port,
-            credentials=credentials,
-            options=channel_options,
-        )
-    return conn

hatchet_sdk/v0/context/__init__.py

@@ -1 +0,0 @@
-from .context import Context