def load_results( self, state: State, upstream_states: Dict[Edge, State] ) -> Tuple[State, Dict[Edge, State]]: """ Given the task's current state and upstream states, populates all relevant result objects for this task run. Args: - state (State): the task's current state. - upstream_states (Dict[Edge, State]): the upstream state_handlers Returns: - Tuple[State, dict]: a tuple of (state, upstream_states) """ upstream_results = {} try: if state.is_mapped(): # ensures mapped children are only loaded once state = state.load_result(self.result) for edge, upstream_state in upstream_states.items(): upstream_states[edge] = upstream_state.load_result( edge.upstream_task.result or self.flow_result ) if edge.key is not None: upstream_results[edge.key] = ( edge.upstream_task.result or self.flow_result ) state.load_cached_results(upstream_results) return state, upstream_states except Exception as exc: new_state = Failed( message=f"Failed to retrieve task results: {exc}", result=exc ) final_state = self.handle_state_change(old_state=state, new_state=new_state) raise ENDRUN(final_state) from exc
def check_task_is_cached(self, state: State, inputs: Dict[str, Result]) -> State: """ Checks if task is cached in the DB and whether any of the caches are still valid. 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 state.is_cached() is True: assert isinstance(state, Cached) # mypy assert sanitized_inputs = {key: res.value for key, res in inputs.items()} if self.task.cache_validator(state, sanitized_inputs, prefect.context.get("parameters")): state = state.load_result(self.result) return state if self.task.cache_for is not None: oldest_valid_cache = datetime.datetime.utcnow( ) - self.task.cache_for cached_states = self.client.get_latest_cached_states( task_id=prefect.context.get("task_id", ""), cache_key=self.task.cache_key, created_after=oldest_valid_cache, ) if cached_states: self.logger.debug( "Task '{name}': {num} candidate cached states were found". format( name=prefect.context.get("task_full_name", self.task.name), num=len(cached_states), )) for candidate_state in cached_states: assert isinstance(candidate_state, Cached) # mypy assert candidate_state.load_cached_results(inputs) sanitized_inputs = { key: res.value for key, res in inputs.items() } if self.task.cache_validator( candidate_state, sanitized_inputs, prefect.context.get("parameters"), ): return candidate_state.load_result(self.result) self.logger.debug( "Task '{name}': can't use cache because no candidate Cached states " "were valid".format(name=prefect.context.get( "task_full_name", self.task.name))) else: self.logger.debug( "Task '{name}': can't use cache because no Cached states were found" .format(name=prefect.context.get("task_full_name", self.task.name))) return state