def check_task_ready_to_map(self, state: State, upstream_states: Dict[Edge, State]) -> State: """ Checks if the parent task is ready to proceed with mapping. Args: - state (State): the current state of this task - upstream_states (Dict[Edge, Union[State, List[State]]]): the upstream states Raises: - ENDRUN: either way, we dont continue past this point """ if state.is_mapped(): # this indicates we are executing a re-run of a mapped pipeline; # in this case, we populate both `map_states` and `cached_inputs` # to ensure the flow runner can properly regenerate the child tasks, # regardless of whether we mapped over an exchanged piece of data # or a non-data-exchanging upstream dependency if len(state.map_states ) == 0 and state.n_map_states > 0: # type: ignore state.map_states = [None] * state.n_map_states # type: ignore state.cached_inputs = { edge.key: state._result # type: ignore for edge, state in upstream_states.items() if edge.key } raise ENDRUN(state) # we can't map if there are no success states with iterables upstream if upstream_states and not any([ edge.mapped and state.is_successful() for edge, state in upstream_states.items() ]): new_state = Failed( "No upstream states can be mapped over.") # type: State raise ENDRUN(new_state) elif not all([ hasattr(state.result, "__getitem__") for edge, state in upstream_states.items() if state.is_successful() and not state.is_mapped() and edge.mapped ]): new_state = Failed( "At least one upstream state has an unmappable result.") raise ENDRUN(new_state) else: # compute and set n_map_states n_map_states = min( [ len(s.result) for e, s in upstream_states.items() if e.mapped and s.is_successful() and not s.is_mapped() ] + [ s.n_map_states # type: ignore for e, s in upstream_states.items() if e.mapped and s.is_mapped() ], default=0, ) new_state = Mapped("Ready to proceed with mapping.", n_map_states=n_map_states) raise ENDRUN(new_state)
def check_task_ready_to_map(self, state: State, upstream_states: Dict[Edge, State]) -> State: """ Checks if the parent task is ready to proceed with mapping. Args: - state (State): the current state of this task - upstream_states (Dict[Edge, Union[State, List[State]]]): the upstream states Raises: - ENDRUN: either way, we dont continue past this point """ if state.is_mapped(): if len(state.map_states ) == 0 and state.n_map_states > 0: # type: ignore state.map_states = [None] * state.n_map_states # type: ignore raise ENDRUN(state) # we can't map if there are no success states with iterables upstream if upstream_states and not any([ edge.mapped and state.is_successful() for edge, state in upstream_states.items() ]): new_state = Failed( "No upstream states can be mapped over.") # type: State raise ENDRUN(new_state) elif not all([ hasattr(state.result, "__getitem__") for edge, state in upstream_states.items() if state.is_successful() and not state.is_mapped() and edge.mapped ]): new_state = Failed( "At least one upstream state has an unmappable result.") raise ENDRUN(new_state) else: # compute and set n_map_states n_map_states = min( [ len(s.result) for e, s in upstream_states.items() if e.mapped and s.is_successful() and not s.is_mapped() ] + [ s.n_map_states # type: ignore for e, s in upstream_states.items() if e.mapped and s.is_mapped() ], default=0, ) new_state = Mapped("Ready to proceed with mapping.", n_map_states=n_map_states) raise ENDRUN(new_state)
def wait_for_mapped_task( self, state: State, executor: "prefect.engine.executors.Executor" ) -> State: """ Blocks until a mapped state's children have finished running. Args: - state (State): the current `Mapped` state - executor (Executor): the run's executor Returns: - State: the new state """ if state.is_mapped(): assert isinstance(state, Mapped) # mypy assert state.map_states = executor.wait(state.map_states) return state