inference/enterprise/workflows/complier/graph_parser.py

from collections import defaultdict
from typing import Any, Dict, List, Optional, Set, Tuple

import networkx as nx
from networkx import DiGraph

from inference.enterprise.workflows.complier.utils import (
    construct_input_selector,
    construct_output_name,
    construct_step_selector,
    get_nodes_of_specific_kind,
    get_step_input_selectors,
    get_step_selector_from_its_output,
    is_condition_step,
    is_step_output_selector,
)
from inference.enterprise.workflows.constants import (
    INPUT_NODE_KIND,
    OUTPUT_NODE_KIND,
    STEP_NODE_KIND,
)
from inference.enterprise.workflows.entities.outputs import JsonField
from inference.enterprise.workflows.entities.steps import StepInterface
from inference.enterprise.workflows.entities.validators import is_selector
from inference.enterprise.workflows.entities.workflows_specification import (
    InputType,
    StepType,
    WorkflowSpecificationV1,
)
from inference.enterprise.workflows.errors import (
    AmbiguousPathDetected,
    NodesNotReachingOutputError,
    NotAcyclicGraphError,
    SelectorToUndefinedNodeError,
)


def prepare_execution_graph(workflow_specification: WorkflowSpecificationV1) -> DiGraph:
    execution_graph = construct_graph(workflow_specification=workflow_specification)
    if not nx.is_directed_acyclic_graph(execution_graph):
        raise NotAcyclicGraphError(f"Detected cycle in execution graph.")
    verify_each_node_reach_at_least_one_output(execution_graph=execution_graph)
    verify_each_node_step_has_parent_in_the_same_branch(execution_graph=execution_graph)
    verify_that_steps_are_connected_with_compatible_inputs(
        execution_graph=execution_graph
    )
    return execution_graph


def construct_graph(workflow_specification: WorkflowSpecificationV1) -> DiGraph:
    execution_graph = nx.DiGraph()
    execution_graph = add_input_nodes_for_graph(
        inputs=workflow_specification.inputs, execution_graph=execution_graph
    )
    execution_graph = add_steps_nodes_for_graph(
        steps=workflow_specification.steps, execution_graph=execution_graph
    )
    execution_graph = add_output_nodes_for_graph(
        outputs=workflow_specification.outputs, execution_graph=execution_graph
    )
    execution_graph = add_steps_edges(
        workflow_specification=workflow_specification, execution_graph=execution_graph
    )
    return add_edges_for_outputs(
        workflow_specification=workflow_specification, execution_graph=execution_graph
    )


def add_input_nodes_for_graph(
    inputs: List[InputType],
    execution_graph: DiGraph,
) -> DiGraph:
    for input_spec in inputs:
        input_selector = construct_input_selector(input_name=input_spec.name)
        execution_graph.add_node(
            input_selector,
            kind=INPUT_NODE_KIND,
            definition=input_spec,
        )
    return execution_graph


def add_steps_nodes_for_graph(
    steps: List[StepType],
    execution_graph: DiGraph,
) -> DiGraph:
    for step in steps:
        step_selector = construct_step_selector(step_name=step.name)
        execution_graph.add_node(
            step_selector,
            kind=STEP_NODE_KIND,
            definition=step,
        )
    return execution_graph


def add_output_nodes_for_graph(
    outputs: List[JsonField],
    execution_graph: DiGraph,
) -> DiGraph:
    for output_spec in outputs:
        execution_graph.add_node(
            construct_output_name(name=output_spec.name),
            kind=OUTPUT_NODE_KIND,
            definition=output_spec,
        )
    return execution_graph


def add_steps_edges(
    workflow_specification: WorkflowSpecificationV1,
    execution_graph: DiGraph,
) -> DiGraph:
    for step in workflow_specification.steps:
        input_selectors = get_step_input_selectors(step=step)
        step_selector = construct_step_selector(step_name=step.name)
        execution_graph = add_edges_for_step_inputs(
            execution_graph=execution_graph,
            input_selectors=input_selectors,
            step_selector=step_selector,
        )
        if step.type == "Condition":
            verify_edge_is_created_between_existing_nodes(
                execution_graph=execution_graph,
                start=step_selector,
                end=step.step_if_true,
            )
            execution_graph.add_edge(step_selector, step.step_if_true)
            verify_edge_is_created_between_existing_nodes(
                execution_graph=execution_graph,
                start=step_selector,
                end=step.step_if_false,
            )
            execution_graph.add_edge(step_selector, step.step_if_false)
    return execution_graph


def add_edges_for_step_inputs(
    execution_graph: DiGraph,
    input_selectors: Set[str],
    step_selector: str,
) -> DiGraph:
    for input_selector in input_selectors:
        if is_step_output_selector(selector_or_value=input_selector):
            input_selector = get_step_selector_from_its_output(
                step_output_selector=input_selector
            )
        verify_edge_is_created_between_existing_nodes(
            execution_graph=execution_graph,
            start=input_selector,
            end=step_selector,
        )
        execution_graph.add_edge(input_selector, step_selector)
    return execution_graph


def add_edges_for_outputs(
    workflow_specification: WorkflowSpecificationV1,
    execution_graph: DiGraph,
) -> DiGraph:
    for output in workflow_specification.outputs:
        output_selector = output.selector
        if is_step_output_selector(selector_or_value=output_selector):
            output_selector = get_step_selector_from_its_output(
                step_output_selector=output_selector
            )
        output_name = construct_output_name(name=output.name)
        verify_edge_is_created_between_existing_nodes(
            execution_graph=execution_graph,
            start=output_selector,
            end=output_name,
        )
        execution_graph.add_edge(output_selector, output_name)
    return execution_graph


def verify_edge_is_created_between_existing_nodes(
    execution_graph: DiGraph,
    start: str,
    end: str,
) -> None:
    if not execution_graph.has_node(start):
        raise SelectorToUndefinedNodeError(
            f"Graph definition contains selector {start} that points to not defined element."
        )
    if not execution_graph.has_node(end):
        raise SelectorToUndefinedNodeError(
            f"Graph definition contains selector {end} that points to not defined element."
        )


def verify_each_node_reach_at_least_one_output(
    execution_graph: DiGraph,
) -> None:
    all_nodes = set(execution_graph.nodes())
    output_nodes = get_nodes_of_specific_kind(
        execution_graph=execution_graph, kind=OUTPUT_NODE_KIND
    )
    nodes_without_outputs = get_nodes_that_do_not_produce_outputs(
        execution_graph=execution_graph
    )
    nodes_that_must_be_reached = output_nodes.union(nodes_without_outputs)
    nodes_reaching_output = (
        get_nodes_that_are_reachable_from_pointed_ones_in_reversed_graph(
            execution_graph=execution_graph,
            pointed_nodes=nodes_that_must_be_reached,
        )
    )
    nodes_not_reaching_output = all_nodes.difference(nodes_reaching_output)
    if len(nodes_not_reaching_output) > 0:
        raise NodesNotReachingOutputError(
            f"Detected {len(nodes_not_reaching_output)} nodes not reaching any of output node:"
            f"{nodes_not_reaching_output}."
        )


def get_nodes_that_do_not_produce_outputs(execution_graph: DiGraph) -> Set[str]:
    # assumption is that nodes without outputs will produce some side effect and shall be
    # treated as output nodes while checking if there is no dangling steps in graph
    step_nodes = get_nodes_of_specific_kind(
        execution_graph=execution_graph, kind=STEP_NODE_KIND
    )
    return {
        step_node
        for step_node in step_nodes
        if len(execution_graph.nodes[step_node]["definition"].get_output_names()) == 0
    }


def get_nodes_that_are_reachable_from_pointed_ones_in_reversed_graph(
    execution_graph: DiGraph,
    pointed_nodes: Set[str],
) -> Set[str]:
    result = set()
    reversed_graph = execution_graph.reverse(copy=True)
    for pointed_node in pointed_nodes:
        nodes_reaching_pointed_one = list(
            nx.dfs_postorder_nodes(reversed_graph, pointed_node)
        )
        result.update(nodes_reaching_pointed_one)
    return result


def verify_each_node_step_has_parent_in_the_same_branch(
    execution_graph: DiGraph,
) -> None:
    """
    Conditional branching creates a bit of mess, in terms of determining which
    steps to execute.
    Let's imagine graph:
              / -> B -> C -> D
    A -> IF <             \
              \ -> E -> F -> G -> H
    where node G requires node C even though IF branched the execution. In other
    words - the problem emerges if a node of kind STEP has a parent (node from which
    it can be achieved) of kind STEP and this parent is in a different branch (point out that
    we allow for a single step to have multiple steps as input, but they must be at the same
    execution path - for instance if D requires an output from C and B - this is allowed).
    Additionally, we must prevent situation when outcomes of branches started by two or more
    condition steps merge with each other, as condition eval may result in contradictory
    execution (2).


    We need to detect that situation upfront, such that we can raise error of ambiguous execution path
    rather than run time-consuming computations that will end up in error.

    To detect problem, first we detect steps with more than one parent step.
    From those steps we trace what sequence of steps would lead to execution of problematic one.
    For each problematic node we take its parent nodes. Then, we analyse paths from
    those parent nodes in reversed topological order (from those nodes towards entry nodes
    of execution graph). While our analysis, on each path we denote `Condition` steps and
    result of condition evaluation that must have been observed in runtime, to reach
    the problematic node while graph execution in normal direction. If we detect that
    for any `Condition` step we would need to output both True and False (more than one registered
    next step of `Condition` step) - we raise error.
    To detect problem (2) - we only let number of different condition steps considered be the number of
    max condition steps in a single path from origin to parent of problematic step.

    Beware that the latter part of algorithm has quite bad time complexity in general case.
    Worst part of algorithm runs at O(V^4) - at least taking coarse, worst-case estimations.
    In fact, there is not so bad:
    * The number of step nodes with multiple parents that we loop over in main loop, reduces the number of
    steps we iterate through in inner loops, as we are dealing with DAG (with quite limited amount of edges)
    and for each multi-parent node takes at least two other nodes (to construct a suspicious group) -
    so expected number of iterations in main loop is low - let's say 1-3 for a real graph.
    * for any reasonable execution graph, the complexity should be acceptable.
    """
    steps_with_more_than_one_parent = detect_steps_with_more_than_one_parent_step(
        execution_graph=execution_graph
    )  # O(V+E)
    if len(steps_with_more_than_one_parent) == 0:
        return None
    reversed_steps_graph = construct_reversed_steps_graph(
        execution_graph=execution_graph
    )  # O(V+E)
    reversed_topological_order = list(
        nx.topological_sort(reversed_steps_graph)
    )  # O(V+E)
    for step in steps_with_more_than_one_parent:  # O(V)
        verify_multi_parent_step_execution_paths(
            reversed_steps_graph=reversed_steps_graph,
            reversed_topological_order=reversed_topological_order,
            step=step,
        )


def detect_steps_with_more_than_one_parent_step(execution_graph: DiGraph) -> Set[str]:
    steps_nodes = get_nodes_of_specific_kind(
        execution_graph=execution_graph,
        kind=STEP_NODE_KIND,
    )
    edges_of_steps_nodes = [edge for edge in execution_graph.edges()]
    steps_parents = defaultdict(set)
    for edge in edges_of_steps_nodes:
        parent, child = edge
        if parent not in steps_nodes or child not in steps_nodes:
            continue
        steps_parents[child].add(parent)
    return {key for key, value in steps_parents.items() if len(value) > 1}


def construct_reversed_steps_graph(execution_graph: DiGraph) -> DiGraph:
    reversed_steps_graph = execution_graph.reverse()
    for node, node_data in list(reversed_steps_graph.nodes(data=True)):
        if node_data.get("kind") != STEP_NODE_KIND:
            reversed_steps_graph.remove_node(node)
    return reversed_steps_graph


def verify_multi_parent_step_execution_paths(
    reversed_steps_graph: nx.DiGraph,
    reversed_topological_order: List[str],
    step: str,
) -> None:
    condition_steps_successors = defaultdict(set)
    max_conditions_steps = 0
    for normal_flow_predecessor in reversed_steps_graph.successors(step):  # O(V)
        reversed_flow_path = (
            construct_reversed_path_to_multi_parent_step_parent(  # O(E) -> O(V^2)
                reversed_steps_graph=reversed_steps_graph,
                reversed_topological_order=reversed_topological_order,
                parent_step=normal_flow_predecessor,
                step=step,
            )
        )
        (
            condition_steps_successors,
            condition_steps,
        ) = denote_condition_steps_successors_in_normal_flow(  # O(V)
            reversed_steps_graph=reversed_steps_graph,
            reversed_flow_path=reversed_flow_path,
            condition_steps_successors=condition_steps_successors,
        )
        max_conditions_steps = max(condition_steps, max_conditions_steps)
    if len(condition_steps_successors) > max_conditions_steps:
        raise AmbiguousPathDetected(
            f"In execution graph, detected collision of branches that originate in different condition steps."
        )
    for condition_step, potential_next_steps in condition_steps_successors.items():
        if len(potential_next_steps) > 1:
            raise AmbiguousPathDetected(
                f"In execution graph, condition step: {condition_step} creates ambiguous execution paths."
            )


def construct_reversed_path_to_multi_parent_step_parent(
    reversed_steps_graph: nx.DiGraph,
    reversed_topological_order: List[str],
    parent_step: str,
    step: str,
) -> List[str]:
    normal_flow_path_nodes = nx.descendants(reversed_steps_graph, parent_step)
    normal_flow_path_nodes.add(parent_step)
    normal_flow_path_nodes.add(step)
    return [n for n in reversed_topological_order if n in normal_flow_path_nodes]


def denote_condition_steps_successors_in_normal_flow(
    reversed_steps_graph: nx.DiGraph,
    reversed_flow_path: List[str],
    condition_steps_successors: Dict[str, Set[str]],
) -> Tuple[Dict[str, Set[str]], int]:
    conditions_steps = 0
    if len(reversed_flow_path) == 0:
        return condition_steps_successors, conditions_steps
    previous_node = reversed_flow_path[0]
    for node in reversed_flow_path[1:]:
        if is_condition_step(execution_graph=reversed_steps_graph, node=node):
            condition_steps_successors[node].add(previous_node)
            conditions_steps += 1
        previous_node = node
    return condition_steps_successors, conditions_steps


def verify_that_steps_are_connected_with_compatible_inputs(
    execution_graph: nx.DiGraph,
) -> None:
    steps_nodes = get_nodes_of_specific_kind(
        execution_graph=execution_graph,
        kind=STEP_NODE_KIND,
    )
    for step in steps_nodes:
        verify_step_inputs_selectors(step=step, execution_graph=execution_graph)


def verify_step_inputs_selectors(step: str, execution_graph: nx.DiGraph) -> None:
    step_definition = execution_graph.nodes[step]["definition"]
    all_inputs = step_definition.get_input_names()
    for input_step in all_inputs:
        input_selector_or_value = getattr(step_definition, input_step)
        if issubclass(type(input_selector_or_value), list):
            for idx, single_selector_or_value in enumerate(input_selector_or_value):
                validate_step_definition_input(
                    step_definition=step_definition,
                    input_name=input_step,
                    execution_graph=execution_graph,
                    input_selector_or_value=single_selector_or_value,
                    index=idx,
                )
        else:
            validate_step_definition_input(
                step_definition=step_definition,
                input_name=input_step,
                execution_graph=execution_graph,
                input_selector_or_value=input_selector_or_value,
            )


def validate_step_definition_input(
    step_definition: StepInterface,
    input_name: str,
    execution_graph: nx.DiGraph,
    input_selector_or_value: Any,
    index: Optional[int] = None,
) -> None:
    if not is_selector(selector_or_value=input_selector_or_value):
        return None
    if is_step_output_selector(selector_or_value=input_selector_or_value):
        input_selector_or_value = get_step_selector_from_its_output(
            step_output_selector=input_selector_or_value
        )
    input_node_definition = execution_graph.nodes[input_selector_or_value]["definition"]
    step_definition.validate_field_selector(
        field_name=input_name,
        input_step=input_node_definition,
        index=index,
    )