def check_flow_is_pending_or_running(self, state: State) -> State: """ Checks if the flow is in either a Pending state or Running state. Either are valid starting points (because we allow simultaneous runs of the same flow run). Args: - state (State): the current state of this flow Returns: - State: the state of the flow after running the check Raises: - ENDRUN: if the flow is not pending or running """ # the flow run is already finished if state.is_finished() is True: self.logger.info("Flow run has already finished.") raise ENDRUN(state) # the flow run must be either pending or running (possibly redundant with above) elif not (state.is_pending() or state.is_running()): self.logger.info("Flow is not ready to run.") raise ENDRUN(state) return state
def set_task_to_running(self, state: State, inputs: Dict[str, Result]) -> State: """ Sets the task to running Args: - state (State): the current state of this task - inputs (Dict[str, Result]): a dictionary of inputs whose keys correspond to the task's `run()` arguments. Returns: - State: the state of the task after running the check Raises: - ENDRUN: if the task is not ready to run """ if not state.is_pending(): self.logger.debug( "Task '{name}': can't set state to Running because it " "isn't Pending; ending run.".format(name=prefect.context.get( "task_full_name", self.task.name))) raise ENDRUN(state) new_state = Running(message="Starting task run.") return new_state
def get_task_inputs( self, state: State, upstream_states: Dict[Edge, State] ) -> Dict[str, Result]: """ Given the task's current state and upstream states, generates the inputs for this task. Upstream state result values are used. If the current state has `cached_inputs`, they will override any upstream values. Args: - state (State): the task's current state. - upstream_states (Dict[Edge, State]): the upstream state_handlers Returns: - Dict[str, Result]: the task inputs """ task_inputs = {} # type: Dict[str, Result] results = {} # type: Dict[str, Result] for edge, upstream_state in upstream_states.items(): # construct task inputs if edge.key is not None: task_inputs[edge.key] = upstream_state._result # type: ignore if state.is_pending() and state.cached_inputs: task_inputs.update( { k: r for k, r in state.cached_inputs.items() if task_inputs.get(k, NoResult) == NoResult } ) return task_inputs
def check_task_is_ready(self, state: State) -> State: """ Checks to make sure the task is ready to run (Pending or Mapped). Args: - state (State): the current state of this task Returns: - State: the state of the task after running the check Raises: - ENDRUN: if the task is not ready to run """ # the task is ready if state.is_pending(): return state # the task is mapped, in which case we still proceed so that the children tasks # are generated (note that if the children tasks) elif state.is_mapped(): self.logger.debug( "Task '{name}': task is mapped, but run will proceed so children are generated.".format( name=prefect.context.get("task_full_name", self.task.name) ) ) return state # this task is already running elif state.is_running(): self.logger.debug( "Task '{name}': task is already running.".format( name=prefect.context.get("task_full_name", self.task.name) ) ) raise ENDRUN(state) elif state.is_cached(): return state # this task is already finished elif state.is_finished(): self.logger.debug( "Task '{name}': task is already finished.".format( name=prefect.context.get("task_full_name", self.task.name) ) ) raise ENDRUN(state) # this task is not pending else: self.logger.debug( "Task '{name}' is not ready to run or state was unrecognized ({state}).".format( name=prefect.context.get("task_full_name", self.task.name), state=state, ) ) raise ENDRUN(state)
def checkpoint_handler(task_runner: DSTaskRunner, old_state: State, new_state: State) -> State: """ A handler designed to implement result caching by filename. If the result handler's ``read`` method can be successfully run, this handler loads the result of that method as the task result and sets the task state to ``Success``. Similarly, on successful completion of the task, if the task was actually run and not loaded from cache, this handler will apply the result handler's ``write`` method to the task. Parameters ---------- task_runner : instance of DSTaskRunner The task runner associated with the flow the handler is used in. old_state : instance of prefect.engine.state.State The current state of the task. new_state : instance of prefect.engine.state.State The expected new state of the task. Returns ------- new_state : instance of prefect.engine.state.State The actual new state of the task. """ if "PREFECT__FLOWS__CHECKPOINTING" in os.environ and os.environ["PREFECT__FLOWS__CHECKPOINTING"] == "true": raise AttributeError("Cannot use standard prefect checkpointing with this handler") if task_runner.result_handler is not None and old_state.is_pending() and new_state.is_running(): if not hasattr(task_runner, "upstream_states"): raise TypeError( "upstream_states not found in task runner. Make sure to use " "prefect_ds.task_runner.DSTaskRunner." ) input_mapping = _create_input_mapping(task_runner.upstream_states) try: data = task_runner.task.result_handler.read(input_mapping=input_mapping) except FileNotFoundError: return new_state except TypeError: # unexpected argument input_mapping raise TypeError( "Result handler could not accept input_mapping argument. " "Please ensure that you are using a handler from prefect_ds." ) result = Result(value=data, result_handler=task_runner.task.result_handler) state = Success(result=result, message="Task loaded from disk.") return state if task_runner.result_handler is not None and old_state.is_running() and new_state.is_successful(): input_mapping = _create_input_mapping(task_runner.upstream_states) task_runner.task.result_handler.write(new_state.result, input_mapping=input_mapping) return new_state
def set_flow_to_running(self, state: State) -> State: """ Puts Pending flows in a Running state; leaves Running flows Running. Args: - state (State): the current state of this flow Returns: - State: the state of the flow after running the check Raises: - ENDRUN: if the flow is not pending or running """ if state.is_pending(): return Running(message="Running flow.") elif state.is_running(): return state else: raise ENDRUN(state)
def set_task_to_running(self, state: State) -> State: """ Sets the task to running Args: - state (State): the current state of this task Returns: - State: the state of the task after running the check Raises: - ENDRUN: if the task is not ready to run """ if not state.is_pending(): self.logger.debug( "Task '{name}': can't set state to Running because it " "isn't Pending; ending run.".format(name=prefect.context.get( "task_full_name", self.task.name))) raise ENDRUN(state) return Running(message="Starting task run.")
def get_task_inputs( self, state: State, upstream_states: Dict[Edge, State] ) -> Dict[str, Result]: """ Given the task's current state and upstream states, generates the inputs for this task. Upstream state result values are used. If the current state has `cached_inputs`, they will override any upstream values which are `NoResult`. Args: - state (State): the task's current state. - upstream_states (Dict[Edge, State]): the upstream state_handlers Returns: - Dict[str, Result]: the task inputs """ task_inputs = {} # type: Dict[str, Result] handlers = {} # type: Dict[str, ResultHandler] for edge, upstream_state in upstream_states.items(): # construct task inputs if edge.key is not None: handlers[edge.key] = handler = getattr( edge.upstream_task, "result_handler", None ) task_inputs[ # type: ignore edge.key ] = upstream_state._result.to_result( # type: ignore handler ) # type: ignore if state.is_pending() and state.cached_inputs is not None: # type: ignore task_inputs.update( { k: r.to_result(handlers.get(k)) for k, r in state.cached_inputs.items() # type: ignore if task_inputs.get(k, NoResult) == NoResult } ) return task_inputs