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