Support Flow Visualization about metaflow

[jaychia]

I wrote a small custom step decorator for myself to tag each task with its parent tasks. I had to do this because I couldn't find any mechanism built into the Metaflow API to track task lineage (do correct me if I'm wrong but it seems easy to check which Step a task came from with .parent, but it doesn't seem possible to check which task spawned the current task?)

def _track_dag(step_func: Callable) -> Callable:
    """Decorator on steps to update self.parent_tasks and self.curr_task so as to enable
    visualization of the DAG during analysis
    """

    @wraps(step_func)
    def wrapper(*args: Any, **kwargs: Any) -> None:
        self_ = args[0]
        curr_task: Optional[str] = self_.curr_task if hasattr(self_, "curr_task") else None
        if "inputs" in signature(step_func).parameters:
            inputs = args[1] if len(args) > 1 else kwargs["inputs"]
            self_.parent_tasks = [i.curr_task for i in inputs]
        else:
            self_.parent_tasks = [curr_task] if curr_task else []
        self_.curr_task = current.pathspec
        return step_func(*args, **kwargs)
    return wrapper

So my flows look like:

class MyFlow(FlowSpec):
    @step
    @_track_dag
    def step(self):
        pass
    ...

And when I'm looking at my run in a notebook, I have another helper that uses PyGraphViz to use the parent pointers and display the DAG, along with some custom coloring to show status of the tasks:

from IPython.display import Image
import pygraphviz as pgv

class MetaflowUtils:
    @staticmethod
    def show_graph(run):
        G = pgv.AGraph(directed=True)

        for s in list(run):
            for t in list(s):
                color = None
                if not t.finished:
                    color = 'black'
                elif t.successful:
                    color = 'green'
                else:
                    color = 'red'
                G.add_node(parse_pathspec(t.data.curr_task), color=color)

        for s in list(run):
            for t in list(s):
                for p in t.data.parent_tasks:
                    G.add_edge(parse_pathspec(p), parse_pathspec(t.data.curr_task))

        return Image(G.draw(format='png', prog='dot'))
        
MetaflowUtils.show_graph(run)

Is this perhaps close to what you're looking for @SteNicholas ?

[dvukolov]

There's some related built-in functionality in Metaflow, which requires graphviz:

$ python playlist.py output-dot | dot -Tpng -o playlist.png

The format can be changed, e.g. to SVG, if you're having issues with PNG dependencies.

[savingoyal]

@SteNicholas Can you elaborate on your ask? You should be able to stand-up custom dashboards using the metaflow.client in a Jupyter notebook easily. A very basic example is in the docs. The client gives you access to the flow state and all the data artifacts. In my experience, dashboards and visualizations depend on a particular business context and visualizing via a Notebook is a very good first step.

[SteNicholas]

@jaychia Yes, I will close this issue.