from __future__ import absolute_import, division, print_function from collections import OrderedDict from functools import partial from hashlib import md5 from operator import getitem import inspect import pickle import os import threading import uuid from toolz import merge, groupby, curry, identity from toolz.functoolz import Compose from .compatibility import (apply, long, unicode, Iterator, is_dataclass, dataclass_fields, Mapping, cPickle) from .context import thread_state from .core import flatten, quote, get as simple_get from .hashing import hash_buffer_hex from .utils import Dispatch, ensure_dict from . import config, local, threaded __all__ = ("DaskMethodsMixin", "is_dask_collection", "compute", "persist", "optimize", "visualize", "tokenize", "normalize_token") def is_dask_collection(x): """Returns ``True`` if ``x`` is a dask collection""" try: return x.__dask_graph__() is not None except (AttributeError, TypeError): return False class DaskMethodsMixin(object): """A mixin adding standard dask collection methods""" __slots__ = () def visualize(self, filename='mydask', format=None, optimize_graph=False, **kwargs): """Render the computation of this object's task graph using graphviz. Requires ``graphviz`` to be installed. Parameters ---------- filename : str or None, optional The name (without an extension) of the file to write to disk. If `filename` is None, no file will be written, and we communicate with dot using only pipes. format : {'png', 'pdf', 'dot', 'svg', 'jpeg', 'jpg'}, optional Format in which to write output file. Default is 'png'. optimize_graph : bool, optional If True, the graph is optimized before rendering. Otherwise, the graph is displayed as is. Default is False. color: {None, 'order'}, optional Options to color nodes. Provide ``cmap=`` keyword for additional colormap **kwargs Additional keyword arguments to forward to ``to_graphviz``. Examples -------- >>> x.visualize(filename='dask.pdf') # doctest: +SKIP >>> x.visualize(filename='dask.pdf', color='order') # doctest: +SKIP Returns ------- result : IPython.diplay.Image, IPython.display.SVG, or None See dask.dot.dot_graph for more information. See Also -------- dask.base.visualize dask.dot.dot_graph Notes ----- For more information on optimization see here: https://docs.dask.org/en/latest/optimize.html """ return visualize(self, filename=filename, format=format, optimize_graph=optimize_graph, **kwargs) def persist(self, **kwargs): """Persist this dask collection into memory This turns a lazy Dask collection into a Dask collection with the same metadata, but now with the results fully computed or actively computing in the background. The action of function differs significantly depending on the active task scheduler. If the task scheduler supports asynchronous computing, such as is the case of the dask.distributed scheduler, then persist will return *immediately* and the return value's task graph will contain Dask Future objects. However if the task scheduler only supports blocking computation then the call to persist will *block* and the return value's task graph will contain concrete Python results. This function is particularly useful when using distributed systems, because the results will be kept in distributed memory, rather than returned to the local process as with compute. Parameters ---------- scheduler : string, optional Which scheduler to use like "threads", "synchronous" or "processes". If not provided, the default is to check the global settings first, and then fall back to the collection defaults. optimize_graph : bool, optional If True [default], the graph is optimized before computation. Otherwise the graph is run as is. This can be useful for debugging. **kwargs Extra keywords to forward to the scheduler function. Returns ------- New dask collections backed by in-memory data See Also -------- dask.base.persist """ (result,) = persist(self, traverse=False, **kwargs) return result def compute(self, **kwargs): """Compute this dask collection This turns a lazy Dask collection into its in-memory equivalent. For example a Dask.array turns into a :func:`numpy.array` and a Dask.dataframe turns into a Pandas dataframe. The entire dataset must fit into memory before calling this operation. Parameters ---------- scheduler : string, optional Which scheduler to use like "threads", "synchronous" or "processes". If not provided, the default is to check the global settings first, and then fall back to the collection defaults. optimize_graph : bool, optional If True [default], the graph is optimized before computation. Otherwise the graph is run as is. This can be useful for debugging. kwargs Extra keywords to forward to the scheduler function. See Also -------- dask.base.compute """ (result,) = compute(self, traverse=False, **kwargs) return result def compute_as_if_collection(cls, dsk, keys, scheduler=None, get=None, **kwargs): """Compute a graph as if it were of type cls. Allows for applying the same optimizations and default scheduler.""" schedule = get_scheduler(scheduler=scheduler, cls=cls, get=get) dsk2 = optimization_function(cls)(ensure_dict(dsk), keys, **kwargs) return schedule(dsk2, keys, **kwargs) def dont_optimize(dsk, keys, **kwargs): return dsk def optimization_function(x): return getattr(x, '__dask_optimize__', dont_optimize) def collections_to_dsk(collections, optimize_graph=True, **kwargs): """ Convert many collections into a single dask graph, after optimization """ optimizations = (kwargs.pop('optimizations', None) or config.get('optimizations', [])) if optimize_graph: groups = groupby(optimization_function, collections) groups = {opt: _extract_graph_and_keys(val) for opt, val in groups.items()} for opt in optimizations: groups = {k: (opt(dsk, keys), keys) for k, (dsk, keys) in groups.items()} dsk = merge(*map(ensure_dict, [opt(dsk, keys, **kwargs) for opt, (dsk, keys) in groups.items()])) else: dsk, _ = _extract_graph_and_keys(collections) return dsk def _extract_graph_and_keys(vals): """Given a list of dask vals, return a single graph and a list of keys such that ``get(dsk, keys)`` is equivalent to ``[v.compute() for v in vals]``.""" from .highlevelgraph import HighLevelGraph graphs = [v.__dask_graph__() for v in vals] keys = [v.__dask_keys__() for v in vals] if any(isinstance(graph, HighLevelGraph) for graph in graphs): graph = HighLevelGraph.merge(*graphs) else: graph = merge(*graphs) return graph, keys def unpack_collections(*args, **kwargs): """Extract collections in preparation for compute/persist/etc... Intended use is to find all collections in a set of (possibly nested) python objects, do something to them (compute, etc...), then repackage them in equivalent python objects. Parameters ---------- *args Any number of objects. If it is a dask collection, it's extracted and added to the list of collections returned. By default, python builtin collections are also traversed to look for dask collections (for more information see the ``traverse`` keyword). traverse : bool, optional If True (default), builtin python collections are traversed looking for any dask collections they might contain. Returns ------- collections : list A list of all dask collections contained in ``args`` repack : callable A function to call on the transformed collections to repackage them as they were in the original ``args``. """ traverse = kwargs.pop('traverse', True) collections = [] repack_dsk = {} collections_token = uuid.uuid4().hex def _unpack(expr): if is_dask_collection(expr): tok = tokenize(expr) if tok not in repack_dsk: repack_dsk[tok] = (getitem, collections_token, len(collections)) collections.append(expr) return tok tok = uuid.uuid4().hex if not traverse: tsk = quote(expr) else: # Treat iterators like lists typ = list if isinstance(expr, Iterator) else type(expr) if typ in (list, tuple, set): tsk = (typ, [_unpack(i) for i in expr]) elif typ is dict: tsk = (dict, [[_unpack(k), _unpack(v)] for k, v in expr.items()]) elif is_dataclass(expr): tsk = (apply, typ, (), (dict, [[f.name, _unpack(getattr(expr, f.name))] for f in dataclass_fields(expr)])) else: return expr repack_dsk[tok] = tsk return tok out = uuid.uuid4().hex repack_dsk[out] = (tuple, [_unpack(i) for i in args]) def repack(results): dsk = repack_dsk.copy() dsk[collections_token] = quote(results) return simple_get(dsk, out) return collections, repack def optimize(*args, **kwargs): """Optimize several dask collections at once. Returns equivalent dask collections that all share the same merged and optimized underlying graph. This can be useful if converting multiple collections to delayed objects, or to manually apply the optimizations at strategic points. Note that in most cases you shouldn't need to call this method directly. Parameters ---------- *args : objects Any number of objects. If a dask object, its graph is optimized and merged with all those of all other dask objects before returning an equivalent dask collection. Non-dask arguments are passed through unchanged. traverse : bool, optional By default dask traverses builtin python collections looking for dask objects passed to ``optimize``. For large collections this can be expensive. If none of the arguments contain any dask objects, set ``traverse=False`` to avoid doing this traversal. optimizations : list of callables, optional Additional optimization passes to perform. **kwargs Extra keyword arguments to forward to the optimization passes. Examples -------- >>> import dask.array as da >>> a = da.arange(10, chunks=2).sum() >>> b = da.arange(10, chunks=2).mean() >>> a2, b2 = optimize(a, b) >>> a2.compute() == a.compute() True >>> b2.compute() == b.compute() True """ collections, repack = unpack_collections(*args, **kwargs) if not collections: return args dsk = collections_to_dsk(collections, **kwargs) postpersists = [a.__dask_postpersist__() if is_dask_collection(a) else (None, a) for a in args] keys, postpersists = [], [] for a in collections: keys.extend(flatten(a.__dask_keys__())) postpersists.append(a.__dask_postpersist__()) return repack([r(dsk, *s) for r, s in postpersists]) def compute(*args, **kwargs): """Compute several dask collections at once. Parameters ---------- args : object Any number of objects. If it is a dask object, it's computed and the result is returned. By default, python builtin collections are also traversed to look for dask objects (for more information see the ``traverse`` keyword). Non-dask arguments are passed through unchanged. traverse : bool, optional By default dask traverses builtin python collections looking for dask objects passed to ``compute``. For large collections this can be expensive. If none of the arguments contain any dask objects, set ``traverse=False`` to avoid doing this traversal. scheduler : string, optional Which scheduler to use like "threads", "synchronous" or "processes". If not provided, the default is to check the global settings first, and then fall back to the collection defaults. optimize_graph : bool, optional If True [default], the optimizations for each collection are applied before computation. Otherwise the graph is run as is. This can be useful for debugging. kwargs Extra keywords to forward to the scheduler function. Examples -------- >>> import dask.array as da >>> a = da.arange(10, chunks=2).sum() >>> b = da.arange(10, chunks=2).mean() >>> compute(a, b) (45, 4.5) By default, dask objects inside python collections will also be computed: >>> compute({'a': a, 'b': b, 'c': 1}) # doctest: +SKIP ({'a': 45, 'b': 4.5, 'c': 1},) """ traverse = kwargs.pop('traverse', True) optimize_graph = kwargs.pop('optimize_graph', True) collections, repack = unpack_collections(*args, traverse=traverse) if not collections: return args schedule = get_scheduler(scheduler=kwargs.pop('scheduler', None), collections=collections, get=kwargs.pop('get', None)) dsk = collections_to_dsk(collections, optimize_graph, **kwargs) keys = [x.__dask_keys__() for x in collections] postcomputes = [x.__dask_postcompute__() for x in collections] results = schedule(dsk, keys, **kwargs) return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)]) def visualize(*args, **kwargs): """ Visualize several dask graphs at once. Requires ``graphviz`` to be installed. All options that are not the dask graph(s) should be passed as keyword arguments. Parameters ---------- dsk : dict(s) or collection(s) The dask graph(s) to visualize. filename : str or None, optional The name (without an extension) of the file to write to disk. If `filename` is None, no file will be written, and we communicate with dot using only pipes. format : {'png', 'pdf', 'dot', 'svg', 'jpeg', 'jpg'}, optional Format in which to write output file. Default is 'png'. optimize_graph : bool, optional If True, the graph is optimized before rendering. Otherwise, the graph is displayed as is. Default is False. color: {None, 'order'}, optional Options to color nodes. Provide ``cmap=`` keyword for additional colormap **kwargs Additional keyword arguments to forward to ``to_graphviz``. Examples -------- >>> x.visualize(filename='dask.pdf') # doctest: +SKIP >>> x.visualize(filename='dask.pdf', color='order') # doctest: +SKIP Returns ------- result : IPython.diplay.Image, IPython.display.SVG, or None See dask.dot.dot_graph for more information. See Also -------- dask.dot.dot_graph Notes ----- For more information on optimization see here: https://docs.dask.org/en/latest/optimize.html """ from dask.dot import dot_graph filename = kwargs.pop('filename', 'mydask') optimize_graph = kwargs.pop('optimize_graph', False) dsks = [arg for arg in args if isinstance(arg, Mapping)] args = [arg for arg in args if is_dask_collection(arg)] dsk = dict(collections_to_dsk(args, optimize_graph=optimize_graph)) for d in dsks: dsk.update(d) color = kwargs.get('color') if color == 'order': from .order import order import matplotlib.pyplot as plt o = order(dsk) try: cmap = kwargs.pop('cmap') except KeyError: cmap = plt.cm.RdBu if isinstance(cmap, str): import matplotlib.pyplot as plt cmap = getattr(plt.cm, cmap) mx = max(o.values()) + 1 colors = {k: _colorize(cmap(v / mx, bytes=True)) for k, v in o.items()} kwargs['function_attributes'] = {k: {'color': v, 'label': str(o[k])} for k, v in colors.items()} kwargs['data_attributes'] = {k: {'color': v} for k, v in colors.items()} elif color: raise NotImplementedError("Unknown value color=%s" % color) return dot_graph(dsk, filename=filename, **kwargs) def persist(*args, **kwargs): """ Persist multiple Dask collections into memory This turns lazy Dask collections into Dask collections with the same metadata, but now with their results fully computed or actively computing in the background. For example a lazy dask.array built up from many lazy calls will now be a dask.array of the same shape, dtype, chunks, etc., but now with all of those previously lazy tasks either computed in memory as many small :class:`numpy.array` (in the single-machine case) or asynchronously running in the background on a cluster (in the distributed case). This function operates differently if a ``dask.distributed.Client`` exists and is connected to a distributed scheduler. In this case this function will return as soon as the task graph has been submitted to the cluster, but before the computations have completed. Computations will continue asynchronously in the background. When using this function with the single machine scheduler it blocks until the computations have finished. When using Dask on a single machine you should ensure that the dataset fits entirely within memory. Examples -------- >>> df = dd.read_csv('/path/to/*.csv') # doctest: +SKIP >>> df = df[df.name == 'Alice'] # doctest: +SKIP >>> df['in-debt'] = df.balance < 0 # doctest: +SKIP >>> df = df.persist() # triggers computation # doctest: +SKIP >>> df.value().min() # future computations are now fast # doctest: +SKIP -10 >>> df.value().max() # doctest: +SKIP 100 >>> from dask import persist # use persist function on multiple collections >>> a, b = persist(a, b) # doctest: +SKIP Parameters ---------- *args: Dask collections scheduler : string, optional Which scheduler to use like "threads", "synchronous" or "processes". If not provided, the default is to check the global settings first, and then fall back to the collection defaults. traverse : bool, optional By default dask traverses builtin python collections looking for dask objects passed to ``persist``. For large collections this can be expensive. If none of the arguments contain any dask objects, set ``traverse=False`` to avoid doing this traversal. optimize_graph : bool, optional If True [default], the graph is optimized before computation. Otherwise the graph is run as is. This can be useful for debugging. **kwargs Extra keywords to forward to the scheduler function. Returns ------- New dask collections backed by in-memory data """ traverse = kwargs.pop('traverse', True) optimize_graph = kwargs.pop('optimize_graph', True) collections, repack = unpack_collections(*args, traverse=traverse) if not collections: return args schedule = get_scheduler(scheduler=kwargs.pop('scheduler', None), collections=collections) if inspect.ismethod(schedule): try: from distributed.client import default_client except ImportError: pass else: try: client = default_client() except ValueError: pass else: if client.get == schedule: results = client.persist(collections, optimize_graph=optimize_graph, **kwargs) return repack(results) dsk = collections_to_dsk(collections, optimize_graph, **kwargs) keys, postpersists = [], [] for a in collections: a_keys = list(flatten(a.__dask_keys__())) rebuild, state = a.__dask_postpersist__() keys.extend(a_keys) postpersists.append((rebuild, a_keys, state)) results = schedule(dsk, keys, **kwargs) d = dict(zip(keys, results)) results2 = [r({k: d[k] for k in ks}, *s) for r, ks, s in postpersists] return repack(results2) ############ # Tokenize # ############ def tokenize(*args, **kwargs): """ Deterministic token >>> tokenize([1, 2, '3']) '7d6a880cd9ec03506eee6973ff551339' >>> tokenize('Hello') == tokenize('Hello') True """ if kwargs: args = args + (kwargs,) return md5(str(tuple(map(normalize_token, args))).encode()).hexdigest() normalize_token = Dispatch() normalize_token.register((int, long, float, str, unicode, bytes, type(None), type, slice, complex, type(Ellipsis)), identity) @normalize_token.register(dict) def normalize_dict(d): return normalize_token(sorted(d.items(), key=str)) @normalize_token.register(OrderedDict) def normalize_ordered_dict(d): return type(d).__name__, normalize_token(list(d.items())) @normalize_token.register(set) def normalize_set(s): return normalize_token(sorted(s, key=str)) @normalize_token.register((tuple, list)) def normalize_seq(seq): return type(seq).__name__, list(map(normalize_token, seq)) @normalize_token.register(object) def normalize_object(o): method = getattr(o, '__dask_tokenize__', None) if method is not None: return method() return normalize_function(o) if callable(o) else uuid.uuid4().hex function_cache = {} function_cache_lock = threading.Lock() def normalize_function(func): try: return function_cache[func] except KeyError: result = _normalize_function(func) if len(function_cache) >= 500: # clear half of cache if full with function_cache_lock: if len(function_cache) >= 500: for k in list(function_cache)[::2]: del function_cache[k] function_cache[func] = result return result except TypeError: # not hashable return _normalize_function(func) def _normalize_function(func): if isinstance(func, curry): func = func._partial if isinstance(func, Compose): first = getattr(func, 'first', None) funcs = reversed((first,) + func.funcs) if first else func.funcs return tuple(normalize_function(f) for f in funcs) elif isinstance(func, partial): args = tuple(normalize_token(i) for i in func.args) if func.keywords: kws = tuple((k, normalize_token(v)) for k, v in sorted(func.keywords.items())) else: kws = None return (normalize_function(func.func), args, kws) else: try: result = pickle.dumps(func, protocol=0) if b'__main__' not in result: # abort on dynamic functions return result except Exception: pass try: import cloudpickle return cloudpickle.dumps(func, protocol=0) except Exception: return str(func) @normalize_token.register_lazy("pandas") def register_pandas(): import pandas as pd @normalize_token.register(pd.Index) def normalize_index(ind): return [ind.name, normalize_token(ind.values)] @normalize_token.register(pd.Categorical) def normalize_categorical(cat): return [normalize_token(cat.codes), normalize_token(cat.categories), cat.ordered] @normalize_token.register(pd.Series) def normalize_series(s): return [s.name, s.dtype, normalize_token(s._data.blocks[0].values), normalize_token(s.index)] @normalize_token.register(pd.DataFrame) def normalize_dataframe(df): data = [block.values for block in df._data.blocks] data += [df.columns, df.index] return list(map(normalize_token, data)) @normalize_token.register_lazy("numpy") def register_numpy(): import numpy as np @normalize_token.register(np.ndarray) def normalize_array(x): if not x.shape: return (str(x), x.dtype) if hasattr(x, 'mode') and getattr(x, 'filename', None): if hasattr(x.base, 'ctypes'): offset = (x.ctypes.get_as_parameter().value - x.base.ctypes.get_as_parameter().value) else: offset = 0 # root memmap's have mmap object as base return (x.filename, os.path.getmtime(x.filename), x.dtype, x.shape, x.strides, offset) if x.dtype.hasobject: try: try: # string fast-path data = hash_buffer_hex('-'.join(x.flat).encode(encoding='utf-8', errors='surrogatepass')) except UnicodeDecodeError: # bytes fast-path data = hash_buffer_hex(b'-'.join(x.flat)) except (TypeError, UnicodeDecodeError): # object data w/o fast-path, use fast cPickle try: data = hash_buffer_hex(cPickle.dumps(x, cPickle.HIGHEST_PROTOCOL)) except Exception: # pickling not supported, use UUID4-based fallback data = uuid.uuid4().hex else: try: data = hash_buffer_hex(x.ravel(order='K').view('i1')) except (BufferError, AttributeError, ValueError): data = hash_buffer_hex(x.copy().ravel(order='K').view('i1')) return (data, x.dtype, x.shape, x.strides) @normalize_token.register(np.matrix) def normalize_matrix(x): return type(x).__name__, normalize_array(x.view(type=np.ndarray)) normalize_token.register(np.dtype, repr) normalize_token.register(np.generic, repr) @normalize_token.register(np.ufunc) def normalize_ufunc(x): try: name = x.__name__ if getattr(np, name) is x: return 'np.' + name except AttributeError: return normalize_function(x) @normalize_token.register_lazy("scipy") def register_scipy(): import scipy.sparse as sp def normalize_sparse_matrix(x, attrs): return type(x).__name__, normalize_seq((normalize_token(getattr(x, key)) for key in attrs)) for cls, attrs in [(sp.dia_matrix, ('data', 'offsets', 'shape')), (sp.bsr_matrix, ('data', 'indices', 'indptr', 'blocksize', 'shape')), (sp.coo_matrix, ('data', 'row', 'col', 'shape')), (sp.csr_matrix, ('data', 'indices', 'indptr', 'shape')), (sp.csc_matrix, ('data', 'indices', 'indptr', 'shape')), (sp.lil_matrix, ('data', 'rows', 'shape'))]: normalize_token.register(cls, partial(normalize_sparse_matrix, attrs=attrs)) @normalize_token.register(sp.dok_matrix) def normalize_dok_matrix(x): return type(x).__name__, normalize_token(sorted(x.items())) def _colorize(t): """ Convert (r, g, b) triple to "#RRGGBB" string For use with ``visualize(color=...)`` Examples -------- >>> _colorize((255, 255, 255)) '#FFFFFF' >>> _colorize((0, 32, 128)) '#002080' """ t = t[:3] i = sum(v * 256 ** (len(t) - i - 1) for i, v in enumerate(t)) h = hex(int(i))[2:].upper() h = '0' * (6 - len(h)) + h return "#" + h named_schedulers = { 'sync': local.get_sync, 'synchronous': local.get_sync, 'single-threaded': local.get_sync, 'threads': threaded.get, 'threading': threaded.get, } try: from dask import multiprocessing as dask_multiprocessing except ImportError: pass else: named_schedulers.update({ 'processes': dask_multiprocessing.get, 'multiprocessing': dask_multiprocessing.get, }) get_err_msg = """ The get= keyword has been removed. Please use the scheduler= keyword instead with the name of the desired scheduler like 'threads' or 'processes' x.compute(scheduler='single-threaded') x.compute(scheduler='threads') x.compute(scheduler='processes') or with a function that takes the graph and keys x.compute(scheduler=my_scheduler_function) or with a Dask client x.compute(scheduler=client) """.strip() def get_scheduler(get=None, scheduler=None, collections=None, cls=None): """ Get scheduler function There are various ways to specify the scheduler to use: 1. Passing in scheduler= parameters 2. Passing these into global confiuration 3. Using defaults of a dask collection This function centralizes the logic to determine the right scheduler to use from those many options """ if get: raise TypeError(get_err_msg) if scheduler is not None: if callable(scheduler): return scheduler elif "Client" in type(scheduler).__name__ and hasattr(scheduler, 'get'): return scheduler.get elif scheduler.lower() in named_schedulers: return named_schedulers[scheduler.lower()] elif scheduler.lower() in ('dask.distributed', 'distributed'): from distributed.worker import get_client return get_client().get elif scheduler.lower() in ['processes', 'multiprocessing']: raise ValueError("Please install cloudpickle to use the '%s' scheduler." % scheduler) else: raise ValueError("Expected one of [distributed, %s]" % ', '.join(sorted(named_schedulers))) # else: # try to connect to remote scheduler with this name # return get_client(scheduler).get if config.get('scheduler', None): return get_scheduler(scheduler=config.get('scheduler', None)) if config.get('get', None): raise ValueError(get_err_msg) if getattr(thread_state, 'key', False): from distributed.worker import get_worker return get_worker().client.get if cls is not None: return cls.__dask_scheduler__ if collections: collections = [c for c in collections if c is not None] if collections: get = collections[0].__dask_scheduler__ if not all(c.__dask_scheduler__ == get for c in collections): raise ValueError("Compute called on multiple collections with " "differing default schedulers. Please specify a " "scheduler=` parameter explicitly in compute or " "globally with `dask.config.set`.") return get return None