Dask Execution Tracer

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):
        self.path = path
        self.save_every = save_every
        self.n = 0
        os.mkdir(path)

    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, format='gif',
                  data_attributes=data,
                  function_attributes=func)

    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

Saves a trace of dask graphs to a folder. These can then be stitched together to create a gif of the execution. One example gifsicle command might be:

$ gifsicle --delay 10 --loop=forever --colors 256 --scale=0.4 -O3 --merge dasks/part_*.gif > output.gif

Example of output:

Is this script out of date?

I am getting an error with dask dot.py when I run it

with Track():
graph.compute()

Traceback (most recent call last):

File "...python3.6/site-packages/dask/dot.py", line 179, in _get_display_cls
raise ValueError("Unknown format '%s' passed to dot_graph" % format)
ValueError: Unknown format 'gif' passed to dot_graph

The solution was to change line 28 to 'png' and then convert those pngs to a gif animation

I replaced gif in the script with png to get it to work