Tracker for Debuggability Improvement

🚀 Key Features

This is a tracker issue for the improvement for debuggability.

• #299

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

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 https://github.com/pytorch/pytorch/pull/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.
• #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

[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.

[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.