def get_flow_run_state(
self,
state: State,
task_states: Dict[Task, State],
task_contexts: Dict[Task, Dict[str, Any]],
return_tasks: Set[Task],
task_runner_state_handlers: Iterable[Callable],
executor: "prefect.executors.base.Executor",
) -> State:
"""
Runs the flow.
Args:
- state (State): starting state for the Flow. Defaults to
`Pending`
- task_states (dict): dictionary of task states to begin
computation with, with keys being Tasks and values their corresponding state
- task_contexts (Dict[Task, Dict[str, Any]]): contexts that will be provided to
each task
- return_tasks ([Task], optional): list of Tasks to include in the
final returned Flow state. Defaults to `None`
- task_runner_state_handlers (Iterable[Callable]): A list of state change handlers
that will be provided to the task_runner, and called whenever a task changes
state.
- executor (Executor): executor to use when performing computation; defaults to the
executor provided in your prefect configuration
Returns:
- State: `State` representing the final post-run state of the `Flow`.
"""
# this dictionary is used for tracking the states of "children" mapped tasks;
# when running on Dask, we want to avoid serializing futures, so instead
# of storing child task states in the `map_states` attribute we instead store
# in this dictionary and only after they are resolved do we attach them to the Mapped state
mapped_children = dict() # type: Dict[Task, list]
if not state.is_running():
self.logger.info("Flow is not in a Running state.")
raise ENDRUN(state)
if return_tasks is None:
return_tasks = set()
if set(return_tasks).difference(self.flow.tasks):
raise ValueError("Some tasks in return_tasks were not found in the flow.")
def extra_context(task: Task, task_index: int = None) -> dict:
return {
"task_name": task.name,
"task_tags": task.tags,
"task_index": task_index,
}
# -- process each task in order
with self.check_for_cancellation(), executor.start():
for task in self.flow.sorted_tasks():
task_state = task_states.get(task)
# if a task is a constant task, we already know its return value
# no need to use up resources by running it through a task runner
if task_state is None and isinstance(
task, prefect.tasks.core.constants.Constant
):
task_states[task] = task_state = Success(result=task.value)
# Always restart completed resource setup/cleanup tasks and
# secret tasks unless they were explicitly cached.
# TODO: we only need to rerun these tasks if any pending
# downstream tasks depend on them.
if (
isinstance(
task,
(
prefect.tasks.core.resource_manager.ResourceSetupTask,
prefect.tasks.core.resource_manager.ResourceCleanupTask,
prefect.tasks.secrets.SecretBase,
),
)
and task_state is not None
and task_state.is_finished()
and not task_state.is_cached()
):
task_states[task] = task_state = Pending()
# if the state is finished, don't run the task, just use the provided state if
# the state is cached / mapped, we still want to run the task runner pipeline
# steps to either ensure the cache is still valid / or to recreate the mapped
# pipeline for possible retries
if (
isinstance(task_state, State)
and task_state.is_finished()
and not task_state.is_cached()
and not task_state.is_mapped()
):
continue
upstream_states = {} # type: Dict[Edge, State]
# this dictionary is used exclusively for "reduce" tasks in particular we store
# the states / futures corresponding to the upstream children, and if running
# on Dask, let Dask resolve them at the appropriate time.
# Note: this is an optimization that allows Dask to resolve the mapped
# dependencies by "elevating" them to a function argument.
upstream_mapped_states = {} # type: Dict[Edge, list]
# -- process each edge to the task
for edge in self.flow.edges_to(task):
# load the upstream task states (supplying Pending as a default)
upstream_states[edge] = task_states.get(
edge.upstream_task, Pending(message="Task state not available.")
)
# if the edge is flattened and not the result of a map, then we
# preprocess the upstream states. If it IS the result of a
# map, it will be handled in `prepare_upstream_states_for_mapping`
if edge.flattened:
if not isinstance(upstream_states[edge], Mapped):
upstream_states[edge] = executor.submit(
executors.flatten_upstream_state, upstream_states[edge]
)
# this checks whether the task is a "reduce" task for a mapped pipeline
# and if so, collects the appropriate upstream children
if not edge.mapped and isinstance(upstream_states[edge], Mapped):
children = mapped_children.get(edge.upstream_task, [])
# if the edge is flattened, then we need to wait for the mapped children
# to complete and then flatten them
if edge.flattened:
children = executors.flatten_mapped_children(
mapped_children=children, executor=executor
)
upstream_mapped_states[edge] = children
# augment edges with upstream constants
for key, val in self.flow.constants[task].items():
edge = Edge(
upstream_task=prefect.tasks.core.constants.Constant(val),
downstream_task=task,
key=key,
)
upstream_states[edge] = Success(
"Auto-generated constant value",
result=ConstantResult(value=val),
)
# handle mapped tasks
if any(edge.mapped for edge in upstream_states.keys()):
# wait on upstream states to determine the width of the pipeline
# this is the key to depth-first execution
upstream_states = executor.wait(
{e: state for e, state in upstream_states.items()}
)
# we submit the task to the task runner to determine if
# we can proceed with mapping - if the new task state is not a Mapped
# state then we don't proceed
task_states[task] = executor.wait(
executor.submit(
run_task,
task=task,
state=task_state, # original state
upstream_states=upstream_states,
context=dict(
prefect.context, **task_contexts.get(task, {})
),
flow_result=self.flow.result,
task_runner_cls=self.task_runner_cls,
task_runner_state_handlers=task_runner_state_handlers,
upstream_mapped_states=upstream_mapped_states,
is_mapped_parent=True,
extra_context=extra_context(task),
)
)
# either way, we should now have enough resolved states to restructure
# the upstream states into a list of upstream state dictionaries to iterate over
list_of_upstream_states = (
executors.prepare_upstream_states_for_mapping(
task_states[task],
upstream_states,
mapped_children,
executor=executor,
)
)
submitted_states = []
for idx, states in enumerate(list_of_upstream_states):
# if we are on a future rerun of a partially complete flow run,
# there might be mapped children in a retrying state; this check
# looks into the current task state's map_states for such info
if (
isinstance(task_state, Mapped)
and len(task_state.map_states) >= idx + 1
):
current_state = task_state.map_states[
idx
] # type: Optional[State]
elif isinstance(task_state, Mapped):
current_state = None
else:
current_state = task_state
# this is where each child is submitted for actual work
submitted_states.append(
executor.submit(
run_task,
task=task,
state=current_state,
upstream_states=states,
context=dict(
prefect.context,
**task_contexts.get(task, {}),
map_index=idx,
),
flow_result=self.flow.result,
task_runner_cls=self.task_runner_cls,
task_runner_state_handlers=task_runner_state_handlers,
upstream_mapped_states=upstream_mapped_states,
extra_context=extra_context(task, task_index=idx),
)
)
if isinstance(task_states.get(task), Mapped):
mapped_children[task] = submitted_states # type: ignore
else:
task_states[task] = executor.submit(
run_task,
task=task,
state=task_state,
upstream_states=upstream_states,
context=dict(prefect.context, **task_contexts.get(task, {})),
flow_result=self.flow.result,
task_runner_cls=self.task_runner_cls,
task_runner_state_handlers=task_runner_state_handlers,
upstream_mapped_states=upstream_mapped_states,
extra_context=extra_context(task),
)
# ---------------------------------------------
# Collect results
# ---------------------------------------------
# terminal tasks determine if the flow is finished
terminal_tasks = self.flow.terminal_tasks()
# reference tasks determine flow state
reference_tasks = self.flow.reference_tasks()
# wait until all terminal tasks are finished
final_tasks = terminal_tasks.union(reference_tasks).union(return_tasks)
final_states = executor.wait(
{
t: task_states.get(t, Pending("Task not evaluated by FlowRunner."))
for t in final_tasks
}
)
# also wait for any children of Mapped tasks to finish, and add them
# to the dictionary to determine flow state
all_final_states = final_states.copy()
for t, s in list(final_states.items()):
if s.is_mapped():
# ensure we wait for any mapped children to complete
if t in mapped_children:
s.map_states = executor.wait(mapped_children[t])
s.result = [ms.result for ms in s.map_states]
all_final_states[t] = s.map_states
assert isinstance(final_states, dict)
key_states = set(flatten_seq([all_final_states[t] for t in reference_tasks]))
terminal_states = set(
flatten_seq([all_final_states[t] for t in terminal_tasks])
)
return_states = {t: final_states[t] for t in return_tasks}
state = self.determine_final_state(
state=state,
key_states=key_states,
return_states=return_states,
terminal_states=terminal_states,
)
return state
if edge.flattened:
if not isinstance(upstream_states[edge], Mapped):
upstream_states[edge] = executor.submit(
executors.flatten_upstream_state, upstream_states[edge]
)
# this checks whether the task is a "reduce" task for a mapped pipeline
# and if so, collects the appropriate upstream children
if not edge.mapped and isinstance(upstream_states[edge], Mapped):
children = mapped_children.get(edge.upstream_task, [])
# if the edge is flattened, then we need to wait for the mapped children
# to complete and then flatten them
if edge.flattened:
children = executors.flatten_mapped_children(
mapped_children=children, executor=executor
)
upstream_mapped_states[edge] = children
# augment edges with upstream constants
for key, val in self.flow.constants[task].items():
edge = Edge(
upstream_task=prefect.tasks.core.constants.Constant(val),
downstream_task=task,
key=key,
)
upstream_states[edge] = Success(
"Auto-generated constant value",
result=ConstantResult(value=val),
)