(
dsk,
filename="mydask",
traverse=True,
optimize_graph=False,
maxval=None,
o=None,
engine: Literal["cytoscape", "ipycytoscape", "graphviz"] | None = None,
limit=None,
**kwargs,
)
| 791 | |
| 792 | |
| 793 | def visualize_dsk( |
| 794 | dsk, |
| 795 | filename="mydask", |
| 796 | traverse=True, |
| 797 | optimize_graph=False, |
| 798 | maxval=None, |
| 799 | o=None, |
| 800 | engine: Literal["cytoscape", "ipycytoscape", "graphviz"] | None = None, |
| 801 | limit=None, |
| 802 | **kwargs, |
| 803 | ): |
| 804 | color = kwargs.get("color") |
| 805 | from dask.order import diagnostics, order |
| 806 | |
| 807 | if color in { |
| 808 | "order", |
| 809 | "order-age", |
| 810 | "order-freed", |
| 811 | "order-memoryincreases", |
| 812 | "order-memorydecreases", |
| 813 | "order-memorypressure", |
| 814 | "age", |
| 815 | "freed", |
| 816 | "memoryincreases", |
| 817 | "memorydecreases", |
| 818 | "memorypressure", |
| 819 | "critical", |
| 820 | "cpath", |
| 821 | }: |
| 822 | import matplotlib.pyplot as plt |
| 823 | |
| 824 | if o is None: |
| 825 | o_stats = order(dsk, return_stats=True) |
| 826 | o = {k: v.priority for k, v in o_stats.items()} |
| 827 | elif isinstance(next(iter(o.values())), int): |
| 828 | o_stats = order(dsk, return_stats=True) |
| 829 | else: |
| 830 | o_stats = o |
| 831 | o = {k: v.priority for k, v in o.items()} |
| 832 | |
| 833 | try: |
| 834 | cmap = kwargs.pop("cmap") |
| 835 | except KeyError: |
| 836 | cmap = plt.cm.plasma |
| 837 | if isinstance(cmap, str): |
| 838 | import matplotlib.pyplot as plt |
| 839 | |
| 840 | cmap = getattr(plt.cm, cmap) |
| 841 | |
| 842 | def label(x): |
| 843 | return str(values[x]) |
| 844 | |
| 845 | data_values = None |
| 846 | if color != "order": |
| 847 | info = diagnostics(dsk, o)[0] |
| 848 | if color.endswith("age"): |
| 849 | values = {key: val.age for key, val in info.items()} |
| 850 | elif color.endswith("freed"): |
searching dependent graphs…