trace.py

import os

from dask.callbacks import Callback
from dask.dot import dot_graph


class Track(Callback):
    def __init__(self, path='dasks', save_every=1, **kwargs):
        self.path = path
        self.save_every = save_every
        self.n = 0
        os.mkdir(path)
        self.kwargs = kwargs

    def _plot(self, dsk, state):
        data = {}
        func = {}
        for key in state['released']:
            data[key] = {'color': 'blue'}
        for key in state['cache']:
            data[key] = {'color': 'red'}
        for key in state['finished']:
            func[key] = {'color': 'blue'}
        for key in state['running']:
            func[key] = {'color': 'red'}

        filename = os.path.join(self.path, 'part_{:0>4d}'.format(self.n))

        dot_graph(dsk, filename=filename,
                  data_attributes=data,
                  function_attributes=func, **self.kwargs)

    def _pretask(self, key, dsk, state):
        if self.n % self.save_every == 0:
            self._plot(dsk, state)
        self.n += 1

    def _finish(self, dsk, state, errored):
        self._plot(dsk, state)
        self.n += 1

from trace import Track
Track(save_every=1, rankdir='LR', path='my-dir', format='png', node_attr={'penwidth': '6'}).register()

x.compute()