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
