dask_test.py

import imp
from os.path import expanduser
import dask
import numpy as np

from cesium import featurize, time_series


imp.reload(dask.async)
execute_task_old = dask.async.execute_task

def execute_and_log_task(key, task, data, queue, get_id, raise_on_exception=False):
#    print("key: ", key)
#    print("task: ", task)
#    print("data: ", data)
#    print("code: {} = {}({})".format(key, task[0].__name__,
#                                     ', '.join([str(t) for t in task[1:]])))
    return execute_task_old(key, task, data, queue, get_id, raise_on_exception=False)


if __name__ == '__main__':
    dask.async.execute_task = execute_and_log_task

#    def double(x):
#        return 2 * x
#    dsk = {'x': [1, 2, 3], 'two_x': (map, double, 'x'), 'total': (sum, 'two_x')}
#    value = dask.async.get_sync(dsk, 'total')

    listcomp = lambda f, x: [f(x_i) for x_i in x]
    dsk = {'uris': [expanduser('~/.local/cesium/ts_data/217801.nc'),
                    expanduser('~/.local/cesium/ts_data/224635.nc')],
           'all_time_series': (listcomp, time_series.from_netcdf, 'uris'),
           'all_features': (listcomp, lambda ts: featurize.featurize_single_ts(ts,
                                                     features_to_use=['maximum', 'minimum']),
                            'all_time_series'),
           'computed_fset': (featurize.assemble_featureset, 'all_features', 'all_time_series')
          }
    value = dask.async.get_sync(dsk, 'computed_fset')
    print(value)