Tracker for Debuggability Improvement

[ejguan]

🚀 Key Features

This is a tracker issue for the improvement for debuggability.

• Print out a graph of DataPipe from the Dictionary returned from traverse function #299

graph TD;
DP1-->DP2;
DP2-->DP3;
DP2-->DP4;
DP3-->DP5;
DP4-->DP6;
DP5-->DP6;
DP6-->output;

Loading

can be printed out as the following

>>> print_graph(traverse(dp6))
DP1 -> DP2 -> DP3 -> DP5
         \             \
         DP4 --------> DP6 ->

• Find a way to differentiate DataPipe instance in the graph
  • Provide shorter name for each DataPipe instance (name is now modified within exception)
  • Print input arguments of the DataPipe within exceptions
    [DataPipe] Improving debug message when exceptions are raised within IterDataPipe pytorch#75618
  • Provide unique name using incremental number for the instance of the same DataPipe class
• Improve back tracing
  • During Iteration, remove this profiling part within back tracing https://github.com/pytorch/pytorch/blob/11231b0f935c092598c994a4bab686485aac1856/torch/utils/data/_typing.py#L340
  • Reduce back tracing to highlight which DataPipe is the culprit when __len__ is invoked.
• Ability to track buffer statistics for .fork .shuffle .demux and similar pipes #210

Nice to Haves

This section is tracking potential features that we may want

• Handling mixed usage of IterDataPipe, MapDataPipe, torcharrow DataFrame
  • Are users able to clearly differentiate these when they are using a mixture of these classes?
• Connect profiling result with graph
  • Different colors for nodes based on their performance (similar to TensorBoard)

Motivation, pitch

This would help our users and developers to easily understand what's going on with the pipeline.
Feel free to post more request for debuggability.

Alternatives

No response

Additional context

No response

cc: @NivekT @VitalyFedyunin

[NivekT]

I made a similar Google Doc here. Not sure if it is easier to brainstorm using GitHub issue or the doc to flesh out what we want to do for each point.

[ejguan]

We should keep design in google doc. Adding this issue for OSS users to gather their feedback

[pmeier]

Regarding graph visualization: I needed this today and hacked something together:

from __future__ import annotations

import dataclasses
from typing import Optional, Any

import matplotlib.pyplot as plt
import networkx as nx
from torch.utils.data.graph import traverse


@dataclasses.dataclass(repr=False)
class Node:
    obj: Any
    child: Optional[Node] = None

    def __repr__(self):
        return type(self.obj).__name__

    def __hash__(self):
        return hash(self.obj)


def scan(graph, child=None):
    for node, parents in graph.items():
        current = Node(node, child)
        yield current
        yield from scan(parents, child=current)


def visualize_graph(dp):
    G = nx.DiGraph()
    for node in set(scan(traverse(dp))):
        if node.child is not None:
            G.add_edge(node, node.child)
    nx.draw_networkx(G)
    plt.show()

Simple example:

from torchdata.datapipes.iter import FileLister, FileOpener

dp = FileLister()
dp = FileOpener(dp).filter(bool).map(list)

visualize_graph(dp)

Complex example:

from torchvision.prototype import datasets

dp = datasets.load("coco")

visualize_graph(dp)

---

I've used networkx as backend here, but we can use any graph visualization library. Let me know, if I should prettify the plots and send a PR.

[ejguan]

Thanks @pmeier , This looks great.
For the backend lib, I personally would prefer tensorboard to visualize the graph because PyTorch is using it to visualize the model graph. https://pytorch.org/tutorials/intermediate/tensorboard_tutorial.html#inspect-the-model-using-tensorboard

Code reference from tensorboard about how to transform it into GraphDef proto: https://github.com/tensorflow/tensorboard/blob/b2b50a328f9bf6755a72cc6f9fd25bdd792b2bbe/tensorboard/plugins/graph/keras_util.py#L193-L203

[pmeier]

One thing that would also simplify debugging, is for each datapipe to have a __repr__. We can follow the approach nn.Module is going:

class IterDataPipe:
    def extra_repr(self) -> str:
        return ""

    def __repr__(self) -> str:
        return f"{type(self).__name__.replace('IterDataPipe', '')}({self.extra_repr()})"


class MinimalIterDataPipe(IterDataPipe):
    pass


def my_map(x):
    return x


class MapperIterDataPipe(IterDataPipe):
    def __init__(self, fn):
        self.fn = fn

    def extra_repr(self):
        return self.fn.__name__


print(MinimalIterDataPipe())
print(MapperIterDataPipe(my_map))

Minimal()
Mapper(my_map)

That would also improve the graph visualization from #299, since each node could contain the __repr__ of the datapipe rather than just its name.

[ejguan]

One thing that would also simplify debugging, is for each datapipe to have a __repr__. We can follow the approach nn.Module is going:

I would propose to change it to iter.ABC from ABCIterDataPipe because the same name can be shared by both iter and map DataPipe.

[pmeier]

Not sure what you mean. Could you give an example? Why does the functional interface need special handling? For example, calling dp.map(...) internally simply does Mapper(dp, ...), right? If so, the returned datapipe will have a __repr__ if Mapper has one.

[ejguan]

Yeah. I mean MapperMapDataPipe vs MapperIterDataPipe. If the whole pipeline contains both IterDataPipe and MapDataPipe ( I totally understand it's confusing ), we need to differentiate them in visualized graph/name.

iter_dp = IterableWrapper(range(10))  # IterDataPipe
iter_dp = iter_dp.map(fn)  # iter.Mapper
map_dp = IterToMapConverter(iter_dp)  # MapDataPipe
map_dp = map_dp.map(fn)  # map.Mapper

[pmeier]

Maybe just color the nodes according to their type? Since it seems this is only relevant for the graph visualization, let's move the discussion to #299 if there is more to it.