Missing edge in flow visualization for conditional branch

Description

The flow visualization for merged branches appears to be incorrect for noop branches.

Reproduction

The visualization in this case implies that the “False” branch doesn’t get merged with the multiplication task at the end, but it does:

@task
def a(x):
    return x + 1

with Flow('f') as flow:
    x = 1
    t = a(x)
    ifelse(x > 1, t, x)
    y = merge(t, x)
    z = y * 2

state = flow.run()
# `x` is returned as-is in the false branch and multiplied by 2 as expected
state.result[z].result
# > 2
# However, the visualization does not show that `x` would be multiplied
flow.visualize()

I would have expected the above to instead look like this, where the merge edge is incorporated if some mutation is applied to x:

@task
def a(x):
    return x + 1

@task
def b(x):
    return x + 2

with Flow('f') as flow:
    x = 1
    t1 = a(x)
    # Now the false branch includes some change to `x` rather than `x` itself
    t2 = b(x)
    ifelse(x > 1, t1, t2)
    y = merge(t1, t2)
    z = y * 2
flow.visualize()

Environment

> conda list | grep prefect
prefect                   0.11.2                     py_0    conda-forge

[cicdw]

Ah awesome - nice to “see you” again 😄

I can definitely explain the nitty gritty details here, but just want to caveat them by re-emphasizing that we’re playing with an edge case (how Prefect handles constants) that rarely arises in practice.

That being said, let’s dive in:

• constants as standard inputs: anytime a task references a non-Prefect constant as an input, Prefect stores the value on the Flow in the constants dictionary (code here). At runtime these constants are unpacked and wrapped in a Success state by the Flow Runner (code here).
• constants as mapped inputs: There is one caveat: if the task is mapping over a constant (e.g., my_task.map([1, 2, 3])), Prefect stores the constant as a proper Constant Task within the Flow (so you’ll actually see this dependency in the DAG view).
• constants as downstream dependencies: there are some circumstances where a constant is being referenced as a downstream dependency of a task. This is very rare, but easily happens with the conditional API (e.g., when using ifelse, both conditional tasks are downstream of the condition). In this case, Prefect creates a proper Constant Task within the Flow that is visible in the DAG.

In your code above, you have a mix of both situations. Using flow_1 as an example:

def flow_1():
    x = 1
    
    # 1 as input, stored in `flow.constants` dictionary
    t = a(x)

    # c is the hardcoded value `False`
    c = x > 1

    # `False` is added to `flow.constants`, 
    # and 1 is added downstream of the compare value task
    ifelse(c, t, x)

    # here again, the number 1 is being used a task input, 
    # and once again added to `flow.constants`
    merge(t, x)


with Flow("flow 1") as flow:
    flow_1()

print(flow.constants) 
# defaultdict(dict,
#            {<Task: a>: {'x': 1},
#             <Task: as_bool>: {'x': False},
#             <Task: Merge>: {'task_2': 1}})

[eric-czech]

Not from my perspective, thanks for the help.