Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/dask/base.py: 16%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import dataclasses
4import uuid
5import warnings
6from collections import OrderedDict
7from collections.abc import Hashable, Iterable, Iterator, Mapping
8from concurrent.futures import Executor
9from contextlib import contextmanager, suppress
10from contextvars import ContextVar
11from functools import partial
12from numbers import Integral, Number
13from operator import getitem
14from typing import TYPE_CHECKING, Any, Literal, TypeVar
16from tlz import merge
18from dask import config, local
19from dask._compatibility import EMSCRIPTEN
20from dask._task_spec import DataNode, Dict, List, Task, TaskRef
21from dask.core import flatten
22from dask.core import get as simple_get
23from dask.system import CPU_COUNT
24from dask.typing import Key, SchedulerGetCallable
25from dask.utils import is_namedtuple_instance, key_split, shorten_traceback
27if TYPE_CHECKING:
28 from dask._expr import Expr
30_DistributedClient = None
31_get_distributed_client = None
32_DISTRIBUTED_AVAILABLE = None
35def _distributed_available() -> bool:
36 # Lazy import in get_scheduler can be expensive
37 global _DistributedClient, _get_distributed_client, _DISTRIBUTED_AVAILABLE
38 if _DISTRIBUTED_AVAILABLE is not None:
39 return _DISTRIBUTED_AVAILABLE # type: ignore[unreachable]
40 try:
41 from distributed import Client as _DistributedClient
42 from distributed.worker import get_client as _get_distributed_client
44 _DISTRIBUTED_AVAILABLE = True
45 except ImportError:
46 _DISTRIBUTED_AVAILABLE = False
47 return _DISTRIBUTED_AVAILABLE
50__all__ = (
51 "DaskMethodsMixin",
52 "annotate",
53 "get_annotations",
54 "is_dask_collection",
55 "compute",
56 "persist",
57 "optimize",
58 "visualize",
59 "tokenize",
60 "normalize_token",
61 "get_collection_names",
62 "get_name_from_key",
63 "replace_name_in_key",
64 "clone_key",
65)
67# Backwards compat
68from dask.tokenize import TokenizationError, normalize_token, tokenize # noqa: F401
70_annotations: ContextVar[dict[str, Any] | None] = ContextVar(
71 "annotations", default=None
72)
75def get_annotations() -> dict[str, Any]:
76 """Get current annotations.
78 Returns
79 -------
80 Dict of all current annotations
82 See Also
83 --------
84 annotate
85 """
86 return _annotations.get() or {}
89@contextmanager
90def annotate(**annotations: Any) -> Iterator[None]:
91 """Context Manager for setting HighLevelGraph Layer annotations.
93 Annotations are metadata or soft constraints associated with
94 tasks that dask schedulers may choose to respect: They signal intent
95 without enforcing hard constraints. As such, they are
96 primarily designed for use with the distributed scheduler.
98 Almost any object can serve as an annotation, but small Python objects
99 are preferred, while large objects such as NumPy arrays are discouraged.
101 Callables supplied as an annotation should take a single *key* argument and
102 produce the appropriate annotation. Individual task keys in the annotated collection
103 are supplied to the callable.
105 Parameters
106 ----------
107 **annotations : key-value pairs
109 Examples
110 --------
112 All tasks within array A should have priority 100 and be retried 3 times
113 on failure.
115 >>> import dask
116 >>> import dask.array as da
117 >>> with dask.annotate(priority=100, retries=3):
118 ... A = da.ones((10000, 10000))
120 Prioritise tasks within Array A on flattened block ID.
122 >>> nblocks = (10, 10)
123 >>> with dask.annotate(priority=lambda k: k[1]*nblocks[1] + k[2]):
124 ... A = da.ones((1000, 1000), chunks=(100, 100))
126 Annotations may be nested.
128 >>> with dask.annotate(priority=1):
129 ... with dask.annotate(retries=3):
130 ... A = da.ones((1000, 1000))
131 ... B = A + 1
133 See Also
134 --------
135 get_annotations
136 """
138 # Sanity check annotations used in place of
139 # legacy distributed Client.{submit, persist, compute} keywords
140 if "workers" in annotations:
141 if isinstance(annotations["workers"], (list, set, tuple)):
142 annotations["workers"] = list(annotations["workers"])
143 elif isinstance(annotations["workers"], str):
144 annotations["workers"] = [annotations["workers"]]
145 elif callable(annotations["workers"]):
146 pass
147 else:
148 raise TypeError(
149 "'workers' annotation must be a sequence of str, a str or a callable, but got {}.".format(
150 annotations["workers"]
151 )
152 )
154 if (
155 "priority" in annotations
156 and not isinstance(annotations["priority"], Number)
157 and not callable(annotations["priority"])
158 ):
159 raise TypeError(
160 "'priority' annotation must be a Number or a callable, but got {}".format(
161 annotations["priority"]
162 )
163 )
165 if (
166 "retries" in annotations
167 and not isinstance(annotations["retries"], Number)
168 and not callable(annotations["retries"])
169 ):
170 raise TypeError(
171 "'retries' annotation must be a Number or a callable, but got {}".format(
172 annotations["retries"]
173 )
174 )
176 if (
177 "resources" in annotations
178 and not isinstance(annotations["resources"], dict)
179 and not callable(annotations["resources"])
180 ):
181 raise TypeError(
182 "'resources' annotation must be a dict, but got {}".format(
183 annotations["resources"]
184 )
185 )
187 if (
188 "allow_other_workers" in annotations
189 and not isinstance(annotations["allow_other_workers"], bool)
190 and not callable(annotations["allow_other_workers"])
191 ):
192 raise TypeError(
193 "'allow_other_workers' annotations must be a bool or a callable, but got {}".format(
194 annotations["allow_other_workers"]
195 )
196 )
197 ctx_annot = _annotations.get()
198 if ctx_annot is None:
199 ctx_annot = {}
200 token = _annotations.set(merge(ctx_annot, annotations))
201 try:
202 yield
203 finally:
204 _annotations.reset(token)
207def is_dask_collection(x) -> bool:
208 """Returns ``True`` if ``x`` is a dask collection.
210 Parameters
211 ----------
212 x : Any
213 Object to test.
215 Returns
216 -------
217 result : bool
218 ``True`` if `x` is a Dask collection.
220 Notes
221 -----
222 The DaskCollection typing.Protocol implementation defines a Dask
223 collection as a class that returns a Mapping from the
224 ``__dask_graph__`` method. This helper function existed before the
225 implementation of the protocol.
227 """
228 if (
229 isinstance(x, type)
230 or not hasattr(x, "__dask_graph__")
231 or not callable(x.__dask_graph__)
232 ):
233 return False
235 pkg_name = getattr(type(x), "__module__", "")
236 if pkg_name.split(".")[0] == "dask_cudf":
237 # Temporary hack to avoid graph materialization. Note that this won't work with
238 # dask_expr.array objects wrapped by xarray or pint. By the time dask_expr.array
239 # is published, we hope to be able to rewrite this method completely.
240 # Read: https://github.com/dask/dask/pull/10676
241 return True
243 from dask._expr import Expr
245 if isinstance(x, Expr):
246 return True
248 try:
249 expr = x.expr
250 except AttributeError:
251 pass
252 else:
253 if isinstance(expr, Expr):
254 return True
256 # Wrappers like xarray expose their dask expressions via __dask_exprs__.
257 # When those are dask expressions we can answer cheaply without
258 # materializing. Otherwise (e.g. xarray wrapping a legacy HighLevelGraph
259 # array, whose children aren't Exprs) fall back to __dask_graph__.
260 dask_exprs = getattr(x, "__dask_exprs__", None)
261 if dask_exprs is not None:
262 exprs = dask_exprs()
263 if exprs is not None and any(isinstance(e, Expr) for e in exprs):
264 return True
266 return x.__dask_graph__() is not None
269class DaskMethodsMixin:
270 """A mixin adding standard dask collection methods"""
272 __slots__ = ("__weakref__",)
274 def visualize(self, filename="mydask", format=None, optimize_graph=False, **kwargs):
275 """Render the computation of this object's task graph using graphviz.
277 Requires ``graphviz`` to be installed.
279 Parameters
280 ----------
281 filename : str or None, optional
282 The name of the file to write to disk. If the provided `filename`
283 doesn't include an extension, '.png' will be used by default.
284 If `filename` is None, no file will be written, and we communicate
285 with dot using only pipes.
286 format : {'png', 'pdf', 'dot', 'svg', 'jpeg', 'jpg'}, optional
287 Format in which to write output file. Default is 'png'.
288 optimize_graph : bool, optional
289 If True, the graph is optimized before rendering. Otherwise,
290 the graph is displayed as is. Default is False.
291 color: {None, 'order'}, optional
292 Options to color nodes. Provide ``cmap=`` keyword for additional
293 colormap
294 **kwargs
295 Additional keyword arguments to forward to ``to_graphviz``.
297 Examples
298 --------
299 >>> x.visualize(filename='dask.pdf') # doctest: +SKIP
300 >>> x.visualize(filename='dask.pdf', color='order') # doctest: +SKIP
302 Returns
303 -------
304 result : IPython.display.Image, IPython.display.SVG, or None
305 See dask.dot.dot_graph for more information.
307 See Also
308 --------
309 dask.visualize
310 dask.dot.dot_graph
312 Notes
313 -----
314 For more information on optimization see here:
316 https://docs.dask.org/en/latest/optimize.html
317 """
318 return visualize(
319 self,
320 filename=filename,
321 format=format,
322 optimize_graph=optimize_graph,
323 **kwargs,
324 )
326 def persist(self, **kwargs):
327 """Persist this dask collection into memory
329 This turns a lazy Dask collection into a Dask collection with the same
330 metadata, but now with the results fully computed or actively computing
331 in the background.
333 The action of function differs significantly depending on the active
334 task scheduler. If the task scheduler supports asynchronous computing,
335 such as is the case of the dask.distributed scheduler, then persist
336 will return *immediately* and the return value's task graph will
337 contain Dask Future objects. However if the task scheduler only
338 supports blocking computation then the call to persist will *block*
339 and the return value's task graph will contain concrete Python results.
341 This function is particularly useful when using distributed systems,
342 because the results will be kept in distributed memory, rather than
343 returned to the local process as with compute.
345 Parameters
346 ----------
347 scheduler : string, optional
348 Which scheduler to use like "threads", "synchronous" or "processes".
349 If not provided, the default is to check the global settings first,
350 and then fall back to the collection defaults.
351 optimize_graph : bool, optional
352 If True [default], the graph is optimized before computation.
353 Otherwise the graph is run as is. This can be useful for debugging.
354 **kwargs
355 Extra keywords to forward to the scheduler function.
357 Returns
358 -------
359 New dask collections backed by in-memory data
361 See Also
362 --------
363 dask.persist
364 """
365 (result,) = persist(self, traverse=False, **kwargs)
366 return result
368 def compute(self, **kwargs):
369 """Compute this dask collection
371 This turns a lazy Dask collection into its in-memory equivalent.
372 For example a Dask array turns into a NumPy array and a Dask dataframe
373 turns into a Pandas dataframe. The entire dataset must fit into memory
374 before calling this operation.
376 Parameters
377 ----------
378 scheduler : string, optional
379 Which scheduler to use like "threads", "synchronous" or "processes".
380 If not provided, the default is to check the global settings first,
381 and then fall back to the collection defaults.
382 optimize_graph : bool, optional
383 If True [default], the graph is optimized before computation.
384 Otherwise the graph is run as is. This can be useful for debugging.
385 kwargs
386 Extra keywords to forward to the scheduler function.
388 See Also
389 --------
390 dask.compute
391 """
392 (result,) = compute(self, traverse=False, **kwargs)
393 return result
395 def __await__(self):
396 try:
397 from distributed import futures_of, wait
398 except ImportError as e:
399 raise type(e)(
400 "Using async/await with dask requires the `distributed` package"
401 ) from e
403 async def f():
404 if futures_of(self):
405 await wait(self)
406 return self
408 return f().__await__()
411def compute_as_if_collection(cls, dsk, keys, scheduler=None, get=None, **kwargs):
412 """Compute a graph as if it were of type cls.
414 Allows for applying the same optimizations and default scheduler."""
415 schedule = get_scheduler(scheduler=scheduler, cls=cls, get=get)
416 dsk2 = optimization_function(cls)(dsk, keys, **kwargs)
417 return schedule(dsk2, keys, **kwargs)
420def dont_optimize(dsk, keys, **kwargs):
421 return dsk
424def optimization_function(x):
425 return getattr(x, "__dask_optimize__", dont_optimize)
428def collections_to_expr(
429 collections: Iterable,
430 optimize_graph: bool = True,
431) -> Expr:
432 """
433 Convert many collections into a single dask expression.
435 Typically, users should not be required to interact with this function.
437 Parameters
438 ----------
439 collections : Iterable
440 An iterable of dask collections to be combined.
441 optimize_graph : bool, optional
442 If this is True and collections are encountered which are backed by
443 legacy HighLevelGraph objects, the returned Expression will run a low
444 level task optimization during materialization.
445 """
446 is_iterable = False
447 if isinstance(collections, (tuple, list, set)):
448 is_iterable = True
449 else:
450 collections = [collections]
451 if not collections:
452 raise ValueError("No collections provided")
453 from dask._expr import CompositeExpr, Expr, HLGExpr, _ExprSequence
455 graphs = []
456 for coll in collections:
457 from dask.delayed import Delayed
459 if isinstance(coll, Delayed):
460 graphs.append(HLGExpr.from_collection(coll, optimize_graph=optimize_graph))
461 else:
462 expr = getattr(coll, "expr", None)
463 if isinstance(expr, Expr):
464 graphs.append(expr)
465 continue
467 dask_exprs = getattr(coll, "__dask_exprs__", None)
468 if dask_exprs is not None:
469 exprs = dask_exprs()
470 if exprs is not None:
471 exprs = tuple(exprs)
472 if exprs:
473 graphs.append(CompositeExpr(coll, *exprs))
474 continue
475 graphs.append(HLGExpr.from_collection(coll, optimize_graph=optimize_graph))
477 if len(graphs) > 1 or is_iterable:
478 return _ExprSequence(*graphs)
479 else:
480 return graphs[0]
483def _rebuild_composite_collection(
484 collection, expr: Expr, dsk: dict, *, cull_to_child_keys: bool
485):
486 """Rebuild a collection that is represented by several child expressions.
488 Collections using ``__dask_exprs__`` expose their children separately so the
489 expression optimizer can still see and rewrite each child graph. After
490 optimization or persistence, rebuild each child collection from the new graph
491 and ask the original collection to assemble an equivalent composite object.
493 ``persist`` passes only computed key results and can select the child keys
494 directly. ``optimize`` passes a full optimized graph, so each child graph must
495 be culled from that graph before rebuilding the child collection.
496 """
497 from dask._collections import new_collection
499 rebuilt_exprs = []
500 for child_expr in expr.exprs:
501 child = new_collection(child_expr)
502 rebuild, state = child.__dask_postpersist__()
503 child_keys = list(flatten(child_expr.__dask_keys__()))
504 if cull_to_child_keys:
505 child_dsk = {key: dsk[key] for key in child_keys}
506 else:
507 from dask.optimization import cull
509 child_dsk, _ = cull(dsk, child_keys)
510 rebuilt = rebuild(child_dsk, *state)
511 rebuilt_exprs.append(rebuilt.expr)
512 return collection.__dask_rebuild_from_exprs__(rebuilt_exprs)
515def unpack_collections(*args, traverse=True):
516 """Extract collections in preparation for compute/persist/etc...
518 Intended use is to find all collections in a set of (possibly nested)
519 python objects, do something to them (compute, etc...), then repackage them
520 in equivalent python objects.
522 Parameters
523 ----------
524 *args
525 Any number of objects. If it is a dask collection, it's extracted and
526 added to the list of collections returned. By default, python builtin
527 collections are also traversed to look for dask collections (for more
528 information see the ``traverse`` keyword).
529 traverse : bool, optional
530 If True (default), builtin python collections are traversed looking for
531 any dask collections they might contain.
533 Returns
534 -------
535 collections : list
536 A list of all dask collections contained in ``args``
537 repack : callable
538 A function to call on the transformed collections to repackage them as
539 they were in the original ``args``.
540 """
542 collections = []
543 repack_dsk = {}
545 collections_token = uuid.uuid4().hex
547 def _unpack(expr):
548 if is_dask_collection(expr):
549 tok = tokenize(expr)
550 if tok not in repack_dsk:
551 repack_dsk[tok] = Task(
552 tok, getitem, TaskRef(collections_token), len(collections)
553 )
554 collections.append(expr)
555 return TaskRef(tok)
557 tok = uuid.uuid4().hex
558 tsk: DataNode | Task # type: ignore[annotation-unchecked]
559 if not traverse:
560 tsk = DataNode(None, expr)
561 else:
562 # Treat iterators like lists
563 typ = list if isinstance(expr, Iterator) else type(expr)
564 if typ in (list, tuple, set):
565 tsk = Task(tok, typ, List(*[_unpack(i) for i in expr]))
566 elif typ in (dict, OrderedDict):
567 tsk = Task(
568 tok, typ, Dict({_unpack(k): _unpack(v) for k, v in expr.items()})
569 )
570 elif dataclasses.is_dataclass(expr) and not isinstance(expr, type):
571 tsk = Task(
572 tok,
573 typ,
574 *[_unpack(getattr(expr, f.name)) for f in dataclasses.fields(expr)],
575 )
576 elif is_namedtuple_instance(expr):
577 tsk = Task(tok, typ, *[_unpack(i) for i in expr])
578 else:
579 return expr
581 repack_dsk[tok] = tsk
582 return TaskRef(tok)
584 out = uuid.uuid4().hex
585 repack_dsk[out] = Task(out, tuple, List(*[_unpack(i) for i in args]))
587 def repack(results):
588 dsk = repack_dsk.copy()
589 dsk[collections_token] = DataNode(collections_token, results)
590 return simple_get(dsk, out)
592 # The original `collections` is kept alive by the closure
593 # This causes the collection to be only freed by the garbage collector
594 collections2 = list(collections)
595 collections.clear()
596 return collections2, repack
599def optimize(*args, traverse=True, **kwargs):
600 """Optimize several dask collections at once.
602 Returns equivalent dask collections that all share the same merged and
603 optimized underlying graph. This can be useful if converting multiple
604 collections to delayed objects, or to manually apply the optimizations at
605 strategic points.
607 Note that in most cases you shouldn't need to call this function directly.
609 Warning::
611 This function triggers a materialization of the collections and looses
612 any annotations attached to HLG layers.
614 Parameters
615 ----------
616 *args : objects
617 Any number of objects. If a dask object, its graph is optimized and
618 merged with all those of all other dask objects before returning an
619 equivalent dask collection. Non-dask arguments are passed through
620 unchanged.
621 traverse : bool, optional
622 By default dask traverses builtin python collections looking for dask
623 objects passed to ``optimize``. For large collections this can be
624 expensive. If none of the arguments contain any dask objects, set
625 ``traverse=False`` to avoid doing this traversal.
626 optimizations : list of callables, optional
627 Additional optimization passes to perform.
628 **kwargs
629 Extra keyword arguments to forward to the optimization passes.
631 Examples
632 --------
633 >>> import dask
634 >>> import dask.array as da
635 >>> a = da.arange(10, chunks=2).sum()
636 >>> b = da.arange(10, chunks=2).mean()
637 >>> a2, b2 = dask.optimize(a, b)
639 >>> a2.compute() == a.compute()
640 np.True_
641 >>> b2.compute() == b.compute()
642 np.True_
643 """
644 # TODO: This API is problematic. The approach to using postpersist forces us
645 # to materialize the graph. Most low level optimizations will materialize as
646 # well
647 collections, repack = unpack_collections(*args, traverse=traverse)
648 if not collections:
649 return args
651 from dask._expr import CompositeExpr, _ExprSequence
653 dsk = collections_to_expr(collections)
654 collection_exprs = list(dsk.operands) if isinstance(dsk, _ExprSequence) else [dsk]
655 if len(collection_exprs) != len(collections):
656 raise RuntimeError(
657 "Expression collection count does not match input collection count: "
658 f"got {len(collection_exprs)} expressions for {len(collections)} "
659 "collections"
660 )
662 if any(isinstance(expr, CompositeExpr) for expr in collection_exprs):
663 dsk = dsk.optimize()
664 collection_exprs = (
665 list(dsk.operands) if isinstance(dsk, _ExprSequence) else [dsk]
666 )
667 if len(collection_exprs) != len(collections):
668 collection_exprs = [None] * len(collections)
670 graph = dsk.__dask_graph__()
672 postpersists = []
673 for a, aexpr in zip(collections, collection_exprs, strict=True):
674 if isinstance(aexpr, CompositeExpr):
675 postpersists.append(
676 _rebuild_composite_collection(a, aexpr, graph, cull_to_child_keys=False)
677 )
678 else:
679 r, s = a.__dask_postpersist__()
680 postpersists.append(r(graph, *s))
682 return repack(postpersists)
685def compute(
686 *args,
687 traverse=True,
688 optimize_graph=True,
689 scheduler=None,
690 get=None,
691 **kwargs,
692):
693 """Compute several dask collections at once.
695 Parameters
696 ----------
697 args : object
698 Any number of objects. If it is a dask object, it's computed and the
699 result is returned. By default, python builtin collections are also
700 traversed to look for dask objects (for more information see the
701 ``traverse`` keyword). Non-dask arguments are passed through unchanged.
702 traverse : bool, optional
703 By default dask traverses builtin python collections looking for dask
704 objects passed to ``compute``. For large collections this can be
705 expensive. If none of the arguments contain any dask objects, set
706 ``traverse=False`` to avoid doing this traversal.
707 scheduler : string, optional
708 Which scheduler to use like "threads", "synchronous" or "processes".
709 If not provided, the default is to check the global settings first,
710 and then fall back to the collection defaults.
711 optimize_graph : bool, optional
712 If True [default], the optimizations for each collection are applied
713 before computation. Otherwise the graph is run as is. This can be
714 useful for debugging.
715 get : ``None``
716 Should be left to ``None`` The get= keyword has been removed.
717 kwargs
718 Extra keywords to forward to the scheduler function.
720 Examples
721 --------
722 >>> import dask
723 >>> import dask.array as da
724 >>> a = da.arange(10, chunks=2).sum()
725 >>> b = da.arange(10, chunks=2).mean()
726 >>> dask.compute(a, b)
727 (np.int64(45), np.float64(4.5))
729 By default, dask objects inside python collections will also be computed:
731 >>> dask.compute({'a': a, 'b': b, 'c': 1})
732 ({'a': np.int64(45), 'b': np.float64(4.5), 'c': 1},)
733 """
735 collections, repack = unpack_collections(*args, traverse=traverse)
736 if not collections:
737 return args
739 schedule = get_scheduler(
740 scheduler=scheduler,
741 collections=collections,
742 get=get,
743 )
744 from dask._expr import FinalizeCompute
746 expr = collections_to_expr(collections, optimize_graph)
747 expr = FinalizeCompute(expr)
749 with shorten_traceback():
750 # The high level optimize will have to be called client side (for now)
751 # The optimize can internally trigger already a computation
752 # (e.g. parquet is reading some statistics). To move this to the
753 # scheduler we'd need some sort of scheduler-client to trigger a
754 # computation from inside the scheduler and continue with optimization
755 # once the results are in. An alternative could be to introduce a
756 # pre-optimize step for the Expressions that handles steps like these as
757 # a dedicated computation
759 # Another caveat is that optimize will only lock in the expression names
760 # after optimization. Names are determined using tokenize and tokenize
761 # is not cross-interpreter (let alone cross-host) stable such that we
762 # have to lock this in before sending stuff (otherwise we'd need to
763 # change the graph submission to a handshake which introduces all sorts
764 # of concurrency control issues)
766 expr = expr.optimize()
767 keys = list(flatten(expr.__dask_keys__()))
769 results = schedule(expr, keys, **kwargs)
771 return repack(results)
774def visualize(
775 *args,
776 filename="mydask",
777 traverse=True,
778 optimize_graph=False,
779 maxval=None,
780 engine: Literal["cytoscape", "ipycytoscape", "graphviz"] | None = None,
781 **kwargs,
782):
783 """
784 Visualize several dask graphs simultaneously.
786 Requires ``graphviz`` to be installed. All options that are not the dask
787 graph(s) should be passed as keyword arguments.
789 Parameters
790 ----------
791 args : object
792 Any number of objects. If it is a dask collection (for example, a
793 dask DataFrame, Array, Bag, or Delayed), its associated graph
794 will be included in the output of visualize. By default, python builtin
795 collections are also traversed to look for dask objects (for more
796 information see the ``traverse`` keyword). Arguments lacking an
797 associated graph will be ignored.
798 filename : str or None, optional
799 The name of the file to write to disk. If the provided `filename`
800 doesn't include an extension, '.png' will be used by default.
801 If `filename` is None, no file will be written, and we communicate
802 with dot using only pipes.
803 format : {'png', 'pdf', 'dot', 'svg', 'jpeg', 'jpg'}, optional
804 Format in which to write output file. Default is 'png'.
805 traverse : bool, optional
806 By default, dask traverses builtin python collections looking for dask
807 objects passed to ``visualize``. For large collections this can be
808 expensive. If none of the arguments contain any dask objects, set
809 ``traverse=False`` to avoid doing this traversal.
810 optimize_graph : bool, optional
811 If True, the graph is optimized before rendering. Otherwise,
812 the graph is displayed as is. Default is False.
813 color : {None, 'order', 'ages', 'freed', 'memoryincreases', 'memorydecreases', 'memorypressure'}, optional
814 Options to color nodes. colormap:
816 - None, the default, no colors.
817 - 'order', colors the nodes' border based on the order they appear in the graph.
818 - 'ages', how long the data of a node is held.
819 - 'freed', the number of dependencies released after running a node.
820 - 'memoryincreases', how many more outputs are held after the lifetime of a node.
821 Large values may indicate nodes that should have run later.
822 - 'memorydecreases', how many fewer outputs are held after the lifetime of a node.
823 Large values may indicate nodes that should have run sooner.
824 - 'memorypressure', the number of data held when the node is run (circle), or
825 the data is released (rectangle).
826 maxval : {int, float}, optional
827 Maximum value for colormap to normalize form 0 to 1.0. Default is ``None``
828 will make it the max number of values
829 collapse_outputs : bool, optional
830 Whether to collapse output boxes, which often have empty labels.
831 Default is False.
832 verbose : bool, optional
833 Whether to label output and input boxes even if the data aren't chunked.
834 Beware: these labels can get very long. Default is False.
835 engine : {"graphviz", "ipycytoscape", "cytoscape"}, optional.
836 The visualization engine to use. If not provided, this checks the dask config
837 value "visualization.engine". If that is not set, it tries to import ``graphviz``
838 and ``ipycytoscape``, using the first one to succeed.
839 **kwargs
840 Additional keyword arguments to forward to the visualization engine.
842 Examples
843 --------
844 >>> x.visualize(filename='dask.pdf') # doctest: +SKIP
845 >>> x.visualize(filename='dask.pdf', color='order') # doctest: +SKIP
847 Returns
848 -------
849 result : IPython.display.Image, IPython.display.SVG, or None
850 See dask.dot.dot_graph for more information.
852 See Also
853 --------
854 dask.dot.dot_graph
856 Notes
857 -----
858 For more information on optimization see here:
860 https://docs.dask.org/en/latest/optimize.html
861 """
862 args, _ = unpack_collections(*args, traverse=traverse)
864 dsk = collections_to_expr(args, optimize_graph=optimize_graph).__dask_graph__()
866 return visualize_dsk(
867 dsk=dsk,
868 filename=filename,
869 traverse=traverse,
870 optimize_graph=optimize_graph,
871 maxval=maxval,
872 engine=engine,
873 **kwargs,
874 )
877def visualize_dsk(
878 dsk,
879 filename="mydask",
880 traverse=True,
881 optimize_graph=False,
882 maxval=None,
883 o=None,
884 engine: Literal["cytoscape", "ipycytoscape", "graphviz"] | None = None,
885 limit=None,
886 **kwargs,
887):
888 color = kwargs.get("color")
889 from dask.order import diagnostics, order
891 if color in {
892 "order",
893 "order-age",
894 "order-freed",
895 "order-memoryincreases",
896 "order-memorydecreases",
897 "order-memorypressure",
898 "age",
899 "freed",
900 "memoryincreases",
901 "memorydecreases",
902 "memorypressure",
903 "critical",
904 "cpath",
905 }:
906 import matplotlib.pyplot as plt
908 if o is None:
909 o_stats = order(dsk, return_stats=True)
910 o = {k: v.priority for k, v in o_stats.items()}
911 elif isinstance(next(iter(o.values())), int):
912 o_stats = order(dsk, return_stats=True)
913 else:
914 o_stats = o
915 o = {k: v.priority for k, v in o.items()}
917 try:
918 cmap = kwargs.pop("cmap")
919 except KeyError:
920 cmap = plt.cm.plasma
921 if isinstance(cmap, str):
922 import matplotlib.pyplot as plt
924 cmap = getattr(plt.cm, cmap)
926 def label(x):
927 return str(values[x])
929 data_values = None
930 if color != "order":
931 info = diagnostics(dsk, o)[0]
932 if color.endswith("age"):
933 values = {key: val.age for key, val in info.items()}
934 elif color.endswith("freed"):
935 values = {key: val.num_dependencies_freed for key, val in info.items()}
936 elif color.endswith("memorypressure"):
937 values = {key: val.num_data_when_run for key, val in info.items()}
938 data_values = {
939 key: val.num_data_when_released for key, val in info.items()
940 }
941 elif color.endswith("memoryincreases"):
942 values = {
943 key: max(0, val.num_data_when_released - val.num_data_when_run)
944 for key, val in info.items()
945 }
946 elif color.endswith("memorydecreases"):
947 values = {
948 key: max(0, val.num_data_when_run - val.num_data_when_released)
949 for key, val in info.items()
950 }
951 elif color.split("-")[-1] in {"critical", "cpath"}:
952 values = {key: val.critical_path for key, val in o_stats.items()}
953 else:
954 raise NotImplementedError(color)
956 if color.startswith("order-"):
958 def label(x):
959 return f"{o[x]}-{values[x]}"
961 else:
962 values = o
963 if maxval is None:
964 maxval = max(1, *values.values())
965 colors = {
966 k: _colorize(tuple(map(int, cmap(v / maxval, bytes=True))))
967 for k, v in values.items()
968 }
969 if data_values is None:
970 data_colors = colors
971 else:
972 data_colors = {
973 k: _colorize(tuple(map(int, cmap(v / maxval, bytes=True))))
974 for k, v in values.items()
975 }
977 kwargs["function_attributes"] = {
978 k: {"color": v, "label": label(k)} for k, v in colors.items()
979 }
980 kwargs["data_attributes"] = {k: {"color": v} for k, v in data_colors.items()}
981 elif color:
982 raise NotImplementedError(f"Unknown value color={color}")
984 # Determine which engine to dispatch to, first checking the kwarg, then config,
985 # then whichever of graphviz or ipycytoscape are installed, in that order.
986 engine = engine or config.get("visualization.engine", None)
988 if not engine:
989 try:
990 import graphviz # noqa: F401
992 engine = "graphviz"
993 except ImportError:
994 try:
995 import ipycytoscape # noqa: F401
997 engine = "cytoscape"
998 except ImportError:
999 pass
1000 if engine == "graphviz":
1001 from dask.dot import dot_graph
1003 return dot_graph(dsk, filename=filename, **kwargs)
1004 elif engine in ("cytoscape", "ipycytoscape"):
1005 from dask.dot import cytoscape_graph
1007 return cytoscape_graph(dsk, filename=filename, **kwargs)
1008 elif engine is None:
1009 raise RuntimeError(
1010 "No visualization engine detected, please install graphviz or ipycytoscape"
1011 )
1012 else:
1013 raise ValueError(f"Visualization engine {engine} not recognized")
1016def persist(*args, traverse=True, optimize_graph=True, scheduler=None, **kwargs):
1017 """Persist multiple Dask collections into memory
1019 This turns lazy Dask collections into Dask collections with the same
1020 metadata, but now with their results fully computed or actively computing
1021 in the background.
1023 For example a lazy dask.array built up from many lazy calls will now be a
1024 dask.array of the same shape, dtype, chunks, etc., but now with all of
1025 those previously lazy tasks either computed in memory as many small :class:`numpy.array`
1026 (in the single-machine case) or asynchronously running in the
1027 background on a cluster (in the distributed case).
1029 This function operates differently if a ``dask.distributed.Client`` exists
1030 and is connected to a distributed scheduler. In this case this function
1031 will return as soon as the task graph has been submitted to the cluster,
1032 but before the computations have completed. Computations will continue
1033 asynchronously in the background. When using this function with the single
1034 machine scheduler it blocks until the computations have finished.
1036 When using Dask on a single machine you should ensure that the dataset fits
1037 entirely within memory.
1039 Examples
1040 --------
1041 >>> df = dd.read_csv('/path/to/*.csv') # doctest: +SKIP
1042 >>> df = df[df.name == 'Alice'] # doctest: +SKIP
1043 >>> df['in-debt'] = df.balance < 0 # doctest: +SKIP
1044 >>> df = df.persist() # triggers computation # doctest: +SKIP
1046 >>> df.value().min() # future computations are now fast # doctest: +SKIP
1047 -10
1048 >>> df.value().max() # doctest: +SKIP
1049 100
1051 >>> from dask import persist # use persist function on multiple collections
1052 >>> a, b = persist(a, b) # doctest: +SKIP
1054 Parameters
1055 ----------
1056 *args: Dask collections
1057 scheduler : string, optional
1058 Which scheduler to use like "threads", "synchronous" or "processes".
1059 If not provided, the default is to check the global settings first,
1060 and then fall back to the collection defaults.
1061 traverse : bool, optional
1062 By default dask traverses builtin python collections looking for dask
1063 objects passed to ``persist``. For large collections this can be
1064 expensive. If none of the arguments contain any dask objects, set
1065 ``traverse=False`` to avoid doing this traversal.
1066 optimize_graph : bool, optional
1067 If True [default], the graph is optimized before computation.
1068 Otherwise the graph is run as is. This can be useful for debugging.
1069 **kwargs
1070 Extra keywords to forward to the scheduler function.
1072 Returns
1073 -------
1074 New dask collections backed by in-memory data
1075 """
1076 collections, repack = unpack_collections(*args, traverse=traverse)
1077 if not collections:
1078 return args
1080 schedule = get_scheduler(scheduler=scheduler, collections=collections)
1082 # Protocol: scheduler can provide its own persist method for async behavior.
1083 # For Client-like objects, get_scheduler returns client.get (a bound method),
1084 # so we check __self__ for the actual client instance.
1085 client = getattr(schedule, "__self__", schedule)
1086 if hasattr(client, "persist") and callable(client.persist):
1087 results = client.persist(collections, optimize_graph=optimize_graph, **kwargs)
1088 return repack(results)
1090 from dask._expr import CompositeExpr, _ExprSequence
1092 expr = collections_to_expr(collections, optimize_graph)
1093 collection_exprs = (
1094 list(expr.operands) if isinstance(expr, _ExprSequence) else [expr]
1095 )
1096 if len(collection_exprs) != len(collections):
1097 raise RuntimeError(
1098 "Expression collection count does not match input collection count: "
1099 f"got {len(collection_exprs)} expressions for {len(collections)} "
1100 "collections"
1101 )
1102 has_composite = any(isinstance(expr, CompositeExpr) for expr in collection_exprs)
1103 expr = expr.optimize()
1104 if has_composite:
1105 collection_exprs = (
1106 list(expr.operands) if isinstance(expr, _ExprSequence) else [expr]
1107 )
1108 if len(collection_exprs) != len(collections):
1109 collection_exprs = [None] * len(collections)
1110 else:
1111 collection_exprs = [None] * len(collections)
1113 keys, postpersists = [], []
1114 for a, aexpr, akeys in zip(
1115 collections, collection_exprs, expr.__dask_keys__(), strict=True
1116 ):
1117 a_keys = list(flatten(akeys))
1118 keys.extend(a_keys)
1119 if isinstance(aexpr, CompositeExpr):
1120 postpersists.append((a, aexpr, None))
1121 else:
1122 rebuild, state = a.__dask_postpersist__()
1123 postpersists.append((rebuild, a_keys, state))
1125 with shorten_traceback():
1126 results = schedule(expr, keys, **kwargs)
1128 d = dict(zip(keys, results))
1129 results2 = []
1130 for rebuild, a_keys, state in postpersists:
1131 if state is None:
1132 results2.append(
1133 _rebuild_composite_collection(
1134 rebuild, a_keys, d, cull_to_child_keys=True
1135 )
1136 )
1137 else:
1138 results2.append(rebuild({k: d[k] for k in a_keys}, *state))
1139 return repack(results2)
1142def _colorize(t):
1143 """Convert (r, g, b) triple to "#RRGGBB" string
1145 For use with ``visualize(color=...)``
1147 Examples
1148 --------
1149 >>> _colorize((255, 255, 255))
1150 '#FFFFFF'
1151 >>> _colorize((0, 32, 128))
1152 '#002080'
1153 """
1154 t = t[:3]
1155 i = sum(v << 8 * i for i, v in enumerate(reversed(t)))
1156 return f"#{i:>06X}"
1159named_schedulers: dict[str, SchedulerGetCallable] = {
1160 "sync": local.get_sync,
1161 "synchronous": local.get_sync,
1162 "single-threaded": local.get_sync,
1163}
1165if not EMSCRIPTEN:
1166 from dask import threaded
1168 named_schedulers.update(
1169 {
1170 "threads": threaded.get,
1171 "threading": threaded.get,
1172 }
1173 )
1175 from dask import multiprocessing as dask_multiprocessing
1177 named_schedulers.update(
1178 {
1179 "processes": dask_multiprocessing.get,
1180 "multiprocessing": dask_multiprocessing.get,
1181 }
1182 )
1185get_err_msg = """
1186The get= keyword has been removed.
1188Please use the scheduler= keyword instead with the name of
1189the desired scheduler like 'threads' or 'processes'
1191 x.compute(scheduler='single-threaded')
1192 x.compute(scheduler='threads')
1193 x.compute(scheduler='processes')
1195or with a function that takes the graph and keys
1197 x.compute(scheduler=my_scheduler_function)
1199or with a Dask client
1201 x.compute(scheduler=client)
1202""".strip()
1205def _ensure_not_async(client):
1206 if client.asynchronous:
1207 if fallback := config.get("admin.async-client-fallback", None):
1208 warnings.warn(
1209 "Distributed Client detected but Client instance is "
1210 f"asynchronous. Falling back to `{fallback}` scheduler. "
1211 "To use an asynchronous Client, please use "
1212 "``Client.compute`` and ``Client.gather`` "
1213 "instead of the top level ``dask.compute``",
1214 UserWarning,
1215 )
1216 return get_scheduler(scheduler=fallback)
1217 else:
1218 raise RuntimeError(
1219 "Attempting to use an asynchronous "
1220 "Client in a synchronous context of `dask.compute`"
1221 )
1222 return client.get
1225def get_scheduler(get=None, scheduler=None, collections=None, cls=None):
1226 """Get scheduler function
1228 There are various ways to specify the scheduler to use:
1230 1. Passing in scheduler= parameters
1231 2. Passing these into global configuration
1232 3. Using a dask.distributed default Client
1233 4. Using defaults of a dask collection
1235 This function centralizes the logic to determine the right scheduler to use
1236 from those many options
1237 """
1238 if get:
1239 raise TypeError(get_err_msg)
1241 if scheduler is not None:
1242 if callable(scheduler):
1243 return scheduler
1244 elif "Client" in type(scheduler).__name__ and hasattr(scheduler, "get"):
1245 return _ensure_not_async(scheduler)
1246 elif isinstance(scheduler, str):
1247 scheduler = scheduler.lower()
1249 client_available = False
1250 if _distributed_available():
1251 assert _DistributedClient is not None
1252 with suppress(ValueError):
1253 _DistributedClient.current(allow_global=True)
1254 client_available = True
1255 if scheduler in named_schedulers:
1256 return named_schedulers[scheduler]
1257 elif scheduler in ("dask.distributed", "distributed"):
1258 if not client_available:
1259 raise RuntimeError(
1260 f"Requested {scheduler} scheduler but no Client active."
1261 )
1262 assert _get_distributed_client is not None
1263 client = _get_distributed_client()
1264 return _ensure_not_async(client)
1265 else:
1266 raise ValueError(
1267 "Expected one of [distributed, {}]".format(
1268 ", ".join(sorted(named_schedulers))
1269 )
1270 )
1271 elif isinstance(scheduler, Executor):
1272 # Get `num_workers` from `Executor`'s `_max_workers` attribute.
1273 # If undefined, fallback to `config` or worst case CPU_COUNT.
1274 num_workers = getattr(scheduler, "_max_workers", None)
1275 if num_workers is None:
1276 num_workers = config.get("num_workers", CPU_COUNT)
1277 assert isinstance(num_workers, Integral) and num_workers > 0
1278 return partial(local.get_async, scheduler.submit, num_workers)
1279 else:
1280 raise ValueError(f"Unexpected scheduler: {scheduler!r}")
1281 # else: # try to connect to remote scheduler with this name
1282 # return get_client(scheduler).get
1284 if config.get("scheduler", None):
1285 return get_scheduler(scheduler=config.get("scheduler", None))
1287 if config.get("get", None):
1288 raise ValueError(get_err_msg)
1290 try:
1291 from distributed import get_client
1293 return _ensure_not_async(get_client())
1294 except (ImportError, ValueError):
1295 pass
1297 if cls is not None:
1298 return cls.__dask_scheduler__
1300 if collections:
1301 collections = [c for c in collections if c is not None]
1302 if collections:
1303 get = collections[0].__dask_scheduler__
1304 if not all(c.__dask_scheduler__ == get for c in collections):
1305 raise ValueError(
1306 "Compute called on multiple collections with "
1307 "differing default schedulers. Please specify a "
1308 "scheduler=` parameter explicitly in compute or "
1309 "globally with `dask.config.set`."
1310 )
1311 return get
1313 return None
1316def wait(x, timeout=None, return_when="ALL_COMPLETED"):
1317 """Wait until computation has finished
1319 This is a compatibility alias for ``dask.distributed.wait``.
1320 If it is applied onto Dask collections without Dask Futures or if Dask
1321 distributed is not installed then it is a no-op
1322 """
1323 try:
1324 from distributed import wait
1326 return wait(x, timeout=timeout, return_when=return_when)
1327 except (ImportError, ValueError):
1328 return x
1331def get_collection_names(collection) -> set[str]:
1332 """Infer the collection names from the dask keys, under the assumption that all keys
1333 are either tuples with matching first element, and that element is a string, or
1334 there is exactly one key and it is a string.
1336 Examples
1337 --------
1338 >>> a.__dask_keys__() # doctest: +SKIP
1339 ["foo", "bar"]
1340 >>> get_collection_names(a) # doctest: +SKIP
1341 {"foo", "bar"}
1342 >>> b.__dask_keys__() # doctest: +SKIP
1343 [[("foo-123", 0, 0), ("foo-123", 0, 1)], [("foo-123", 1, 0), ("foo-123", 1, 1)]]
1344 >>> get_collection_names(b) # doctest: +SKIP
1345 {"foo-123"}
1346 """
1347 if not is_dask_collection(collection):
1348 raise TypeError(f"Expected Dask collection; got {type(collection)}")
1349 return {get_name_from_key(k) for k in flatten(collection.__dask_keys__())}
1352def get_name_from_key(key: Key) -> str:
1353 """Given a dask collection's key, extract the collection name.
1355 Parameters
1356 ----------
1357 key: string or tuple
1358 Dask collection's key, which must be either a single string or a tuple whose
1359 first element is a string (commonly referred to as a collection's 'name'),
1361 Examples
1362 --------
1363 >>> get_name_from_key("foo")
1364 'foo'
1365 >>> get_name_from_key(("foo-123", 1, 2))
1366 'foo-123'
1367 """
1368 if isinstance(key, tuple) and key and isinstance(key[0], str):
1369 return key[0]
1370 if isinstance(key, str):
1371 return key
1372 raise TypeError(f"Expected str or a tuple starting with str; got {key!r}")
1375KeyOrStrT = TypeVar("KeyOrStrT", Key, str)
1378def replace_name_in_key(key: KeyOrStrT, rename: Mapping[str, str]) -> KeyOrStrT:
1379 """Given a dask collection's key, replace the collection name with a new one.
1381 Parameters
1382 ----------
1383 key: string or tuple
1384 Dask collection's key, which must be either a single string or a tuple whose
1385 first element is a string (commonly referred to as a collection's 'name'),
1386 rename:
1387 Mapping of zero or more names from : to. Extraneous names will be ignored.
1388 Names not found in this mapping won't be replaced.
1390 Examples
1391 --------
1392 >>> replace_name_in_key("foo", {})
1393 'foo'
1394 >>> replace_name_in_key("foo", {"foo": "bar"})
1395 'bar'
1396 >>> replace_name_in_key(("foo-123", 1, 2), {"foo-123": "bar-456"})
1397 ('bar-456', 1, 2)
1398 """
1399 if isinstance(key, tuple) and key and isinstance(key[0], str):
1400 return (rename.get(key[0], key[0]),) + key[1:]
1401 if isinstance(key, str):
1402 return rename.get(key, key)
1403 raise TypeError(f"Expected str or a tuple starting with str; got {key!r}")
1406def clone_key(key: KeyOrStrT, seed: Hashable) -> KeyOrStrT:
1407 """Clone a key from a Dask collection, producing a new key with the same prefix and
1408 indices and a token which is a deterministic function of the previous key and seed.
1410 Examples
1411 --------
1412 >>> clone_key("x", 123) # doctest: +SKIP
1413 'x-c4fb64ccca807af85082413d7ef01721'
1414 >>> clone_key("inc-cbb1eca3bafafbb3e8b2419c4eebb387", 123) # doctest: +SKIP
1415 'inc-bc629c23014a4472e18b575fdaf29ee7'
1416 >>> clone_key(("sum-cbb1eca3bafafbb3e8b2419c4eebb387", 4, 3), 123) # doctest: +SKIP
1417 ('sum-c053f3774e09bd0f7de6044dbc40e71d', 4, 3)
1418 """
1419 if isinstance(key, tuple) and key and isinstance(key[0], str):
1420 return (clone_key(key[0], seed),) + key[1:]
1421 if isinstance(key, str):
1422 prefix = key_split(key)
1423 token = tokenize(key, seed)
1424 return f"{prefix}-{token}"
1425 raise TypeError(f"Expected str or a tuple starting with str; got {key!r}")