Skip to content

Instantly share code, notes, and snippets.

@mrocklin
Forked from jcrist/trace.py
Created July 11, 2016 21:03
Show Gist options
  • Save mrocklin/056f7a6c6ad5fa8ecc85de3e85a63550 to your computer and use it in GitHub Desktop.
Save mrocklin/056f7a6c6ad5fa8ecc85de3e85a63550 to your computer and use it in GitHub Desktop.
Dask Execution Tracer
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment