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