| 875 | |
| 876 | |
| 877 | def visualize_dsk( |
| 878 | dsk, |
| 879 | filename="mydask", |
| 880 | traverse=True, |
| 881 | optimize_graph=False, |
| 882 | maxval=None, |
| 883 | o=None, |
| 884 | engine: Literal["cytoscape", "ipycytoscape", "graphviz"] | None = None, |
| 885 | limit=None, |
| 886 | **kwargs, |
| 887 | ): |
| 888 | color = kwargs.get("color") |
| 889 | from dask.order import diagnostics, order |
| 890 | |
| 891 | if color in { |
| 892 | "order", |
| 893 | "order-age", |
| 894 | "order-freed", |
| 895 | "order-memoryincreases", |
| 896 | "order-memorydecreases", |
| 897 | "order-memorypressure", |
| 898 | "age", |
| 899 | "freed", |
| 900 | "memoryincreases", |
| 901 | "memorydecreases", |
| 902 | "memorypressure", |
| 903 | "critical", |
| 904 | "cpath", |
| 905 | }: |
| 906 | import matplotlib.pyplot as plt |
| 907 | |
| 908 | if o is None: |
| 909 | o_stats = order(dsk, return_stats=True) |
| 910 | o = {k: v.priority for k, v in o_stats.items()} |
| 911 | elif isinstance(next(iter(o.values())), int): |
| 912 | o_stats = order(dsk, return_stats=True) |
| 913 | else: |
| 914 | o_stats = o |
| 915 | o = {k: v.priority for k, v in o.items()} |
| 916 | |
| 917 | try: |
| 918 | cmap = kwargs.pop("cmap") |
| 919 | except KeyError: |
| 920 | cmap = plt.cm.plasma |
| 921 | if isinstance(cmap, str): |
| 922 | import matplotlib.pyplot as plt |
| 923 | |
| 924 | cmap = getattr(plt.cm, cmap) |
| 925 | |
| 926 | def label(x): |
| 927 | return str(values[x]) |
| 928 | |
| 929 | data_values = None |
| 930 | if color != "order": |
| 931 | info = diagnostics(dsk, o)[0] |
| 932 | if color.endswith("age"): |
| 933 | values = {key: val.age for key, val in info.items()} |
| 934 | elif color.endswith("freed"): |