PyPI 官网下载 | arkitekt-0.1.90.tar.gz

from arkitekt.monitor.monitor import get_current_monitor from herre.console.context import get_current_console from uuid import uuid4 from arkitekt.packers.utils import expand_outputs, shrink_inputs from asyncio.futures import Future from arkitekt.messages.postman.log import LogLevel from contextvars import Context from arkitekt.messages.postman.reserve.reserve_transition import ReserveState from arkitekt.contracts.exceptions import AssignmentException from arkitekt.messages import * from arkitekt.schema.enums import NodeType from arkitekt.messages.postman.reserve.params import ReserveParams from arkitekt.postman import Postman from rich.table import Table from rich.panel import Panel from arkitekt.monitor import Monitor, current_monitor from arkitekt.registry import get_current_postman import asyncio import logging from herre.herre import Herre, get_current_herre from koil.koil import Koil, get_current_koil from koil.loop import koil, koil_gen logger = logging.getLogger(__name__) class UnknownMessageError(Exception): pass class ReservationError(Exception): pass class CouldNotReserveError(ReservationError): pass class IncorrectStateForAssignation(ReservationError): pass class AssignationError(Exception): pass class Omitted(AssignationError): pass def build_reserve_message( reference, node_id: str = None, template_id: str = None, provision: str = None, params_dict: dict = {}, with_log=False, context=None, ): assert reference is not None, "Must have a reference" assert ( node_id is not None or template_id is not None ), "Please provide either a node_id or template_id" data = { "node": node_id, "template": template_id, "provision": provision, "params": params_dict, } meta = { "reference": reference, "extensions": { "with_progress": with_log, }, } if context: meta = {**meta, "context": context} return BouncedReserveMessage(data=data, meta=meta) else: return ReserveMessage(data=data, meta=meta) def build_unreserve_messsage(reference, reservation, with_log=False, context=None): assert reference is not None, "Must have a reference" data = {"reservation": reservation} meta = {"reference": reference, "extensions": {"with_progress": with_log}} if context: meta = {**meta, "context": context} return BouncedUnreserveMessage(data=data, meta=meta) else: return UnreserveMessage(data=data, meta=meta) def build_assign_message( reference, reservation, args, kwargs, with_log=False, context=None, persist=False ): assert reference is not None, "Must have a reference" data = { "reservation": reservation, "args": args, "kwargs": kwargs, } meta = { "reference": reference, "extensions": {"with_progress": with_log, "persist": persist}, } if context: meta = {**meta, "context": context} return BouncedAssignMessage(data=data, meta=meta) else: return AssignMessage(data=data, meta=meta) def build_unassign_messsage( reference, assignation, provision, with_log=False, context=None, persist=False ): assert reference is not None, "Must have a reference" data = {"assignation": assignation, "provision": provision} meta = { "reference": reference, "extensions": {"with_progress": with_log, "persist": persist}, } if context: meta = {**meta, "context": context} return BouncedUnassignMessage(data=data, meta=meta) else: return UnassignMessage(data=data, meta=meta) class Reservation: def __init__( self, node, reference: str = None, provision: str = None, monitor: Monitor = None, ignore_node_exceptions=False, transition_hook=None, with_log=False, omit_on=[], enter_on=[ReserveState.ACTIVE], exit_on=[ReserveState.ERROR, ReserveState.CANCELLED, ReserveState.CRITICAL], context: Context = None, koil: Koil = None, herre: Herre = None, postman: Postman = None, **params, ) -> None: self.monitor: Monitor = monitor or get_current_monitor() self.panel = ( self.monitor.create_reservation_panel(self) if self.monitor else None ) self.console = get_current_console() self.herre = herre or get_current_herre() self.postman = postman or get_current_postman(force_creation=True) # Reservation Params self.reference = reference or str(uuid.uuid4()) self.provision = provision self.node = node self.params = ReserveParams(**params) self.with_log = with_log or (self.monitor.log if self.monitor else None) self.context = context # with_bounced allows us forward bounced checks if self.context: assert ( "can_forward_bounce" in self.herre.state.scopes ), "In order to use with_bounced forwarding you need to have the can_forward_bounced scope" # Exception Mangement self.ignore_node_exceptions = ignore_node_exceptions self.critical_error = None # State management self.transition_hook = transition_hook assert self.transition_hook is None or asyncio.iscoroutinefunction( self.transition_hook ), "Transition Hook must be either a coroutine or set to None" self.exit_states = exit_on self.enter_states = enter_on self.omit_states = omit_on self.current_state = ReserveState.STARTING def log(self, message: str, level: LogLevel = LogLevel.DEBUG): """Logs a Message The Logged Message will be display on the Monitor if running inside a Monitor and send to the logging output. Args: message (str): The Message level (LogLevel, optional): The LogLevel. Defaults to LogLevel.DEBUG. """ if self.panel: self.panel.log(message, level=level) logger.info(f"{level}: {message}") async def transition_state(self, message: ReserveTransitionMessage): # Once we acquire a reserved resource our contract (the inner part of the context can start) if self.transition_hook: await self.transition_hook(self, message.data.state) if message.data.state in self.exit_states: if self.enter_future.done(): self.log( f"We have transitioned to a critical State {message.data.message}. Terminating on Next Call" ) else: self.log("Cancelling Reservation") print(message) raise Exception(message.data.message) if not self.is_closing: self.log( f"Received Exitstate: {message.data.state}. Closing reservation at next assignment", level=LogLevel.CRITICAL, ) if message.data.state in self.enter_states: if self.enter_future.done(): logger.info("We are already entered.") else: self.enter_future.set_result(message.meta.reference) self.old_state = self.current_state self.current_state = message.data.state self.log( f"[red] {self.old_state} > {self.current_state}: {message.data.message}", LogLevel.INFO, ) async def assign_async( self, *args, bypass_shrink=False, bypass_expand=False, persist=True, with_log=True, context=None, raise_node_exceptions=True, **kwargs, ): logger.info(f"Assigning {args} {kwargs} ") assert ( self.node.type == NodeType.FUNCTION ), "You cannot assign to a Generator Node, use