import inspect import json import traceback from concurrent.futures import Future, ThreadPoolExecutor from typing import Any, Generic, Type, TypeVar, cast, overload from warnings import warn from pydantic import BaseModel, StrictStr from hatchet_sdk.clients.events import EventClient from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry from hatchet_sdk.clients.rest_client import RestApi from hatchet_sdk.clients.run_event_listener import RunEventListenerClient from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener from hatchet_sdk.context.worker_context import WorkerContext from hatchet_sdk.contracts.dispatcher_pb2 import OverridesData from hatchet_sdk.contracts.workflows_pb2 import ( BulkTriggerWorkflowRequest, TriggerWorkflowRequest, ) from hatchet_sdk.utils.types import WorkflowValidator from hatchet_sdk.utils.typing import is_basemodel_subclass from hatchet_sdk.workflow_run import WorkflowRunRef from ..clients.admin import ( AdminClient, ChildTriggerWorkflowOptions, ChildWorkflowRunDict, TriggerWorkflowOptions, WorkflowRunDict, ) from ..clients.dispatcher.dispatcher import ( # type: ignore[attr-defined] Action, DispatcherClient, ) from ..logger import logger DEFAULT_WORKFLOW_POLLING_INTERVAL = 5 # Seconds T = TypeVar("T", bound=BaseModel) def get_caller_file_path() -> str: caller_frame = inspect.stack()[2] return caller_frame.filename class BaseContext: action: Action spawn_index: int def _prepare_workflow_options( self, key: str | None = None, options: ChildTriggerWorkflowOptions | None = None, worker_id: str | None = None, ) -> TriggerWorkflowOptions: workflow_run_id = self.action.workflow_run_id step_run_id = self.action.step_run_id desired_worker_id = None if options is not None and "sticky" in options and options["sticky"] == True: desired_worker_id = worker_id meta = None if options is not None and "additional_metadata" in options: meta = options["additional_metadata"] ## TODO: Pydantic here to simplify this trigger_options: TriggerWorkflowOptions = { "parent_id": workflow_run_id, "parent_step_run_id": step_run_id, "child_key": key, "child_index": self.spawn_index, "additional_metadata": meta, "desired_worker_id": desired_worker_id, } self.spawn_index += 1 return trigger_options class ContextAioImpl(BaseContext): def __init__( self, action: Action, dispatcher_client: DispatcherClient, admin_client: AdminClient, event_client: EventClient, rest_client: RestApi, workflow_listener: PooledWorkflowRunListener, workflow_run_event_listener: RunEventListenerClient, worker: WorkerContext, namespace: str = "", ): self.action = action self.dispatcher_client = dispatcher_client self.admin_client = admin_client self.event_client = event_client self.rest_client = rest_client self.workflow_listener = workflow_listener self.workflow_run_event_listener = workflow_run_event_listener self.namespace = namespace self.spawn_index = -1 self.worker = worker @tenacity_retry async def spawn_workflow( self, workflow_name: str, input: dict[str, Any] = {}, key: str | None = None, options: ChildTriggerWorkflowOptions | None = None, ) -> WorkflowRunRef: worker_id = self.worker.id() # if ( # options is not None # and "sticky" in options # and options["sticky"] == True # and not self.worker.has_workflow(workflow_name) # ): # raise Exception( # f"cannot run with sticky: workflow {workflow_name} is not registered on the worker" # ) trigger_options = self._prepare_workflow_options(key, options, worker_id) return await self.admin_client.aio.run_workflow( workflow_name, input, trigger_options ) @tenacity_retry async def spawn_workflows( self, child_workflow_runs: list[ChildWorkflowRunDict] ) -> list[WorkflowRunRef]: if len(child_workflow_runs) == 0: raise Exception("no child workflows to spawn") worker_id = self.worker.id() bulk_trigger_workflow_runs: list[WorkflowRunDict] = [] for child_workflow_run in child_workflow_runs: workflow_name = child_workflow_run["workflow_name"] input = child_workflow_run["input"] key = child_workflow_run.get("key") options = child_workflow_run.get("options", {}) trigger_options = self._prepare_workflow_options(key, options, worker_id) bulk_trigger_workflow_runs.append( WorkflowRunDict( workflow_name=workflow_name, input=input, options=trigger_options ) ) return await self.admin_client.aio.run_workflows(bulk_trigger_workflow_runs) class Context(BaseContext): spawn_index = -1 worker: WorkerContext def __init__( self, action: Action, dispatcher_client: DispatcherClient, admin_client: AdminClient, event_client: EventClient, rest_client: RestApi, workflow_listener: PooledWorkflowRunListener, workflow_run_event_listener: RunEventListenerClient, worker: WorkerContext, namespace: str = "", validator_registry: dict[str, WorkflowValidator] = {}, ): self.worker = worker self.validator_registry = validator_registry self.aio = ContextAioImpl( action, dispatcher_client, admin_client, event_client, rest_client, workflow_listener, workflow_run_event_listener, worker, namespace, ) # Check the type of action.action_payload before attempting to load it as JSON if isinstance(action.action_payload, (str, bytes, bytearray)): try: self.data = cast(dict[str, Any], json.loads(action.action_payload)) except Exception as e: logger.error(f"Error parsing action payload: {e}") # Assign an empty dictionary if parsing fails self.data: dict[str, Any] = {} # type: ignore[no-redef] else: # Directly assign the payload to self.data if it's already a dict self.data = ( action.action_payload if isinstance(action.action_payload, dict) else {} ) self.action = action # FIXME: stepRunId is a legacy field, we should remove it self.stepRunId = action.step_run_id self.step_run_id = action.step_run_id self.exit_flag = False self.dispatcher_client = dispatcher_client self.admin_client = admin_client self.event_client = event_client self.rest_client = rest_client self.workflow_listener = workflow_listener self.workflow_run_event_listener = workflow_run_event_listener self.namespace = namespace # FIXME: this limits the number of concurrent log requests to 1, which means we can do about # 100 log lines per second but this depends on network. self.logger_thread_pool = ThreadPoolExecutor(max_workers=1) self.stream_event_thread_pool = ThreadPoolExecutor(max_workers=1) # store each key in the overrides field in a lookup table # overrides_data is a dictionary of key-value pairs self.overrides_data = self.data.get("overrides", {}) if action.get_group_key_run_id != "": self.input = self.data else: self.input = self.data.get("input", {}) def step_output(self, step: str) -> dict[str, Any] | BaseModel: workflow_validator = next( (v for k, v in self.validator_registry.items() if k.split(":")[-1] == step), None, ) try: parent_step_data = cast(dict[str, Any], self.data["parents"][step]) except KeyError: raise ValueError(f"Step output for '{step}' not found") if workflow_validator and (v := workflow_validator.step_output): return v.model_validate(parent_step_data) return parent_step_data def triggered_by_event(self) -> bool: return cast(str, self.data.get("triggered_by", "")) == "event" def workflow_input(self) -> dict[str, Any] | T: if (r := self.validator_registry.get(self.action.action_id)) and ( i := r.workflow_input ): return cast( T, i.model_validate(self.input), ) return self.input def workflow_run_id(self) -> str: return self.action.workflow_run_id def cancel(self) -> None: logger.debug("cancelling step...") self.exit_flag = True # done returns true if the context has been cancelled def done(self) -> bool: return self.exit_flag def playground(self, name: str, default: str | None = None) -> str | None: # if the key exists in the overrides_data field, return the value if name in self.overrides_data: warn( "Use of `overrides_data` is deprecated.", DeprecationWarning, stacklevel=1, ) return str(self.overrides_data[name]) caller_file = get_caller_file_path() self.dispatcher_client.put_overrides_data( OverridesData( stepRunId=self.stepRunId, path=name, value=json.dumps(default), callerFilename=caller_file, ) ) return default def _log(self, line: str) -> tuple[bool, Exception | None]: try: self.event_client.log(message=line, step_run_id=self.stepRunId) return True, None except Exception as e: # we don't want to raise an exception here, as it will kill the log thread return False, e def log(self, line: Any, raise_on_error: bool = False) -> None: if self.stepRunId == "": return if not isinstance(line, str): try: line = json.dumps(line) except Exception: line = str(line) future = self.logger_thread_pool.submit(self._log, line) def handle_result(future: Future[tuple[bool, Exception | None]]) -> None: success, exception = future.result() if not success and exception: if raise_on_error: raise exception else: thread_trace = "".join( traceback.format_exception( type(exception), exception, exception.__traceback__ ) ) call_site_trace = "".join(traceback.format_stack()) logger.error( f"Error in log thread: {exception}\n{thread_trace}\nCalled from:\n{call_site_trace}" ) future.add_done_callback(handle_result) def release_slot(self) -> None: return self.dispatcher_client.release_slot(self.stepRunId) def _put_stream(self, data: str | bytes) -> None: try: self.event_client.stream(data=data, step_run_id=self.stepRunId) except Exception as e: logger.error(f"Error putting stream event: {e}") def put_stream(self, data: str | bytes) -> None: if self.stepRunId == "": return self.stream_event_thread_pool.submit(self._put_stream, data) def refresh_timeout(self, increment_by: str) -> None: try: return self.dispatcher_client.refresh_timeout( step_run_id=self.stepRunId, increment_by=increment_by ) except Exception as e: logger.error(f"Error refreshing timeout: {e}") def retry_count(self) -> int: return self.action.retry_count def additional_metadata(self) -> dict[str, Any] | None: return self.action.additional_metadata def child_index(self) -> int | None: return self.action.child_workflow_index def child_key(self) -> str | None: return self.action.child_workflow_key def parent_workflow_run_id(self) -> str | None: return self.action.parent_workflow_run_id def step_run_errors(self) -> dict[str, str]: errors = cast(dict[str, str], self.data.get("step_run_errors", {})) if not errors: logger.error( "No step run errors found. `context.step_run_errors` is intended to be run in an on-failure step, and will only work on engine versions more recent than v0.53.10" ) return errors def fetch_run_failures(self) -> list[dict[str, StrictStr]]: data = self.rest_client.workflow_run_get(self.action.workflow_run_id) other_job_runs = [ run for run in (data.job_runs or []) if run.job_id != self.action.job_id ] # TODO: Parse Step Runs using a Pydantic Model rather than a hand crafted dictionary return [ { "step_id": step_run.step_id, "step_run_action_name": step_run.step.action, "error": step_run.error, } for job_run in other_job_runs if job_run.step_runs for step_run in job_run.step_runs if step_run.error and step_run.step ] @tenacity_retry def spawn_workflow( self, workflow_name: str, input: dict[str, Any] = {}, key: str | None = None, options: ChildTriggerWorkflowOptions | None = None, ) -> WorkflowRunRef: worker_id = self.worker.id() trigger_options = self._prepare_workflow_options(key, options, worker_id) return self.admin_client.run_workflow(workflow_name, input, trigger_options) @tenacity_retry def spawn_workflows( self, child_workflow_runs: list[ChildWorkflowRunDict] ) -> list[WorkflowRunRef]: if len(child_workflow_runs) == 0: raise Exception("no child workflows to spawn") worker_id = self.worker.id() bulk_trigger_workflow_runs: list[WorkflowRunDict] = [] for child_workflow_run in child_workflow_runs: workflow_name = child_workflow_run["workflow_name"] input = child_workflow_run["input"] key = child_workflow_run.get("key") options = child_workflow_run.get("options", {}) trigger_options = self._prepare_workflow_options(key, options, worker_id) bulk_trigger_workflow_runs.append( WorkflowRunDict( workflow_name=workflow_name, input=input, options=trigger_options ) ) return self.admin_client.run_workflows(bulk_trigger_workflow_runs)