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

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

        filename = 'dasks/part_{:0>4d}.png'.format(self.n)
        self.n += 1

        dot_graph(dsk, filename=filename, data_attributes=data,
                  function_attributes=func)

    def _pretask(self, key, dsk, state):
        self._plot(dsk, state)

    def _finish(self, dsk, state, errored):
        if not errored:
            self._plot(dsk, state)