Repository: dask/dask
Pull request #7282 (2021-02-26 18:40:21)

Add split_every to graph_manipulation
Although the outputs of the map stage of ``dask.graph_manipulation.checkpoint`` are just a bunch of None's, it's been observed that the distributed scheduler incorrectly concentrates the pre-map data (potentially gigabytes) onto the worker that computes the final node of ``checkpoint``. This will eventually need to be fixed in distributed but won't be an easy fix.

This PR sidesteps the problem, removing the nexus node and replacing it with a recursive aggregation equivalent to the one already implemented in dask.array, dask.bag, and dask.dataframe.

Follow-up: #7283

crusaderky: @jrbourbeau ready for review

Patch:
diff --git a/dask/graph_manipulation.py b/dask/graph_manipulation.py
index 3da0601ae..7440556fc 100644
--- a/dask/graph_manipulation.py
+++ b/dask/graph_manipulation.py
@@ -3,7 +3,18 @@ output collections produced by this module are typically not functionally equiva
 their inputs.
 """
 import uuid
-from typing import Callable, Hashable, Optional, Set, Tuple, TypeVar
+from numbers import Number
+from typing import (
+    AbstractSet,
+    Callable,
+    Dict,
+    Hashable,
+    Optional,
+    Set,
+    Tuple,
+    TypeVar,
+    Union,
+)
 
 from .base import (
     clone_key,
@@ -21,9 +32,15 @@ from .highlevelgraph import MaterializedLayer, HighLevelGraph, Layer
 __all__ = ("bind", "checkpoint", "clone", "wait_on")
 
 T = TypeVar("T")
+try:
+    from typing import Literal  # Python >= 3.8
+
+    SplitEvery = Union[Number, Literal[False], None]
+except ImportError:
+    SplitEvery = Union[Number, bool, None]  # type: ignore
 
 
-def checkpoint(*collections) -> Delayed:
+def checkpoint(*collections, split_every: SplitEvery = None) -> Delayed:
     """Build a :doc:`delayed` which waits until all chunks of the input collection(s)
     have been computed before returning None.
 
@@ -32,16 +49,37 @@ def checkpoint(*collections) -> Delayed:
     collections
         Zero or more Dask collections or nested data structures containing zero or more
         collections
+    split_every: int >= 2 or False, optional
+        Determines the depth of the recursive aggregation. If greater than the number of
+        input keys, the aggregation will be performed in multiple steps; the depth of
+        the aggregation graph will be :math:`log_{split_every}(input keys)`. Setting to
+        a low value can reduce cache size and network transfers, at the cost of more CPU
+        and a larger dask graph.
+
+        Set to False to disable. Defaults to 8.
 
     Returns
     -------
     :doc:`delayed` yielding None
     """
+    if split_every is None:
+        # FIXME https://github.com/python/typeshed/issues/5074
+        split_every = 8  # type: ignore
+    elif split_every is not False:
+        split_every = int(split_every)  # type: ignore
+        if split_every < 2:  # type: ignore
+            raise ValueError("split_every must be False, None, or >= 2")
+
     collections, _ = unpack_collections(*collections)
-    if len(collections) != 1:
-        return delayed(chunks.checkpoint)(*(checkpoint(c) for c in collections))
-    (collection,) = collections
+    if len(collections) == 1:
+        return _checkpoint_one(collections[0], split_every)
+    else:
+        return delayed(chunks.checkpoint)(
+            *(_checkpoint_one(c, split_every) for c in collections)
+        )
+
 
+def _checkpoint_one(collection, split_every) -> Delayed:
     tok = tokenize(collection)
     name = "checkpoint-" + tok
 
@@ -53,32 +91,36 @@ def checkpoint(*collections) -> Delayed:
         # Collection has 0 or 1 keys; no need for a map step
         layer = {name: (chunks.checkpoint, collection.__dask_keys__())}
         dsk = HighLevelGraph.from_collections(name, layer, dependencies=(collection,))
-    else:
-        # Collection has 2+ keys; apply a two-step map->reduce algorithm so that we
-        # transfer over the network and store in RAM only a handful of None's instead of
-        # the full computed collection's contents
-        dsks = []
-        map_names = set()
-        map_keys = []
-
-        for prev_name in get_collection_names(collection):
-            map_name = "checkpoint_map-" + tokenize(prev_name, tok)
-            map_names.add(map_name)
-            map_layer = _build_map_layer(
-                chunks.checkpoint, prev_name, map_name, collection
-            )
-            map_keys += list(map_layer.get_output_keys())
-            dsks.append(
-                HighLevelGraph.from_collections(
-                    map_name, map_layer, dependencies=(collection,)
-                )
-            )
-
-        reduce_layer = {name: (chunks.checkpoint, map_keys)}
+        return Delayed(name, dsk)
+
+    # Collection has 2+ keys; apply a two-step map->reduce algorithm so that we
+    # transfer over the network and store in RAM only a handful of None's instead of
+    # the full computed collection's contents
+    dsks = []
+    map_names = set()
+    map_keys = []
+
+    for prev_name in get_collection_names(collection):
+        map_name = "checkpoint_map-" + tokenize(prev_name, tok)
+        map_names.add(map_name)
+        map_layer = _build_map_layer(chunks.checkpoint, prev_name, map_name, collection)
+        map_keys += list(map_layer.get_output_keys())
         dsks.append(
-            HighLevelGraph({name: reduce_layer}, dependencies={name: map_names})
+            HighLevelGraph.from_collections(
+                map_name, map_layer, dependencies=(collection,)
+            )
         )
-        dsk = HighLevelGraph.merge(*dsks)
+
+    # recursive aggregation
+    reduce_layer: dict = {}
+    while split_every and len(map_keys) > split_every:
+        k = (name, len(reduce_layer))
+        reduce_layer[k] = (chunks.checkpoint, map_keys[:split_every])
+        map_keys = map_keys[split_every:] + [k]
+    reduce_layer[name] = (chunks.checkpoint, map_keys)
+
+    dsks.append(HighLevelGraph({name: reduce_layer}, dependencies={name: map_names}))
+    dsk = HighLevelGraph.merge(*dsks)
 
     return Delayed(name, dsk)
 
@@ -177,6 +219,7 @@ def bind(
     omit=None,
     seed: Hashable = None,
     assume_layers: bool = True,
+    split_every: SplitEvery = None,
 ) -> T:
     """
     Make ``children`` collection(s), optionally omitting sub-collections, dependent on
@@ -233,6 +276,8 @@ def bind(
         False
             Use a slower algorithm that works at keys level, which makes none of the
             above assumptions.
+    split_every
+        See :func:`checkpoint`
 
     Returns
     -------
@@ -247,7 +292,9 @@ def bind(
         seed = uuid.uuid4().bytes
 
     # parents=None is a special case invoked by the one-liner wrapper clone() below
-    blocker = checkpoint(parents) if parents is not None else None
+    blocker = (
+        checkpoint(parents, split_every=split_every) if parents is not None else None
+    )
 
     omit, _ = unpack_collections(omit)
     if assume_layers:
@@ -281,15 +328,22 @@ def _bind_one(
         # writing, can only happen with third-party collections
         return child
 
-    dsk = child.__dask_graph__()
-    new_layers = {}
-    new_deps = {}
-    if not isinstance(dsk, HighLevelGraph):
+    dsk = child.__dask_graph__()  # type: ignore
+    new_layers: Dict[str, Layer] = {}
+    new_deps: Dict[str, AbstractSet[str]] = {}
+
+    if isinstance(dsk, HighLevelGraph):
+        try:
+            layers_to_clone = set(child.__dask_layers__())  # type: ignore
+        except AttributeError:
+            layers_to_clone = prev_coll_names.copy()
+    else:
         if len(prev_coll_names) == 1:
             hlg_name = next(iter(prev_coll_names))
         else:
             hlg_name = tokenize(*prev_coll_names)
         dsk = HighLevelGraph.from_collections(hlg_name, dsk)
+        layers_to_clone = {hlg_name}
 
     clone_keys = dsk.get_all_external_keys() - omit_keys
     for layer_name in omit_layers:
@@ -311,10 +365,6 @@ def _bind_one(
     else:
         blocker_key = None
 
-    try:
-        layers_to_clone = set(child.__dask_layers__())
-    except AttributeError:
-        layers_to_clone = prev_coll_names.copy()
     layers_to_copy_verbatim = set()
 
     while layers_to_clone:
@@ -333,11 +383,12 @@ def _bind_one(
         new_layers[new_layer_name], is_bound = layer.clone(
             keys=clone_keys, seed=seed, bind_to=blocker_key
         )
-        new_deps[new_layer_name] = {
+        new_dep = {
             clone_key(dep, seed=seed) for dep in layer_deps_to_clone
         } | layer_deps_to_omit
         if is_bound:
-            new_deps[new_layer_name].add(blocker_key)
+            new_dep.add(blocker_key)
+        new_deps[new_layer_name] = new_dep
 
     # Add the layers of the collections from omit from child.dsk. Note that, when
     # assume_layers=False, it would be unsafe to simply do HighLevelGraph.merge(dsk,
@@ -352,7 +403,7 @@ def _bind_one(
         new_deps[layer_name] = layer_deps
         new_layers[layer_name] = dsk.layers[layer_name]
 
-    rebuild, args = child.__dask_postpersist__()
+    rebuild, args = child.__dask_postpersist__()  # type: ignore
     return rebuild(
         HighLevelGraph(new_layers, new_deps),
         *args,
@@ -364,8 +415,8 @@ def clone(*collections, omit=None, seed: Hashable = None, assume_layers: bool =
     """Clone dask collections, returning equivalent collections that are generated from
     independent calculations.
 
-    Example
-    -------
+    Examples
+    --------
     (tokens have been simplified for the sake of brevity)
 
     >>> from dask import array as da
@@ -414,7 +465,7 @@ def clone(*collections, omit=None, seed: Hashable = None, assume_layers: bool =
     return out[0] if len(collections) == 1 else out
 
 
-def wait_on(*collections):
+def wait_on(*collections, split_every: SplitEvery = None):
     """Ensure that all chunks of all input collections have been computed before
     computing the dependents of any of the chunks.
 
@@ -438,6 +489,8 @@ def wait_on(*collections):
     ----------
     collections
         Zero or more Dask collections or nested structures of Dask collections
+    split_every
+        See :func:`checkpoint`
 
     Returns
     -------
@@ -446,7 +499,7 @@ def wait_on(*collections):
         or a nested structure equivalent to the input where the original collections
         have been replaced.
     """
-    blocker = checkpoint(*collections)
+    blocker = checkpoint(*collections, split_every=split_every)
 
     def block_one(coll):
         tok = tokenize(coll, blocker)
diff --git a/docs/source/optimize.rst b/docs/source/optimize.rst
index 500dc7493..792f96ae0 100644
--- a/docs/source/optimize.rst
+++ b/docs/source/optimize.rst
@@ -79,7 +79,7 @@ output keys to a scheduler ``get`` function:
     >>> outputs = ['print1', 'print2']
     >>> dsk1, dependencies = cull(dsk, outputs)  # remove unnecessary tasks from the graph
 
-    >>> results = get(dsk2, outputs)
+    >>> results = get(dsk1, outputs)
     word list has 2 occurrences of apple, out of 7 words
     word list has 2 occurrences of orange, out of 7 words
 


Test patch:
diff --git a/dask/tests/test_graph_manipulation.py b/dask/tests/test_graph_manipulation.py
index 933782ac8..9bd4a536d 100644
--- a/dask/tests/test_graph_manipulation.py
+++ b/dask/tests/test_graph_manipulation.py
@@ -215,15 +215,15 @@ def test_bind(layers):
     dsk4b = {"e": (cnt.f, 3)}
 
     if layers:
-        dsk1 = HighLevelGraph.from_collections("a-1", dsk1)
+        dsk1 = HighLevelGraph({"a-1": dsk1}, {"a-1": set()})
         dsk2 = HighLevelGraph(
-            {"a-1": dsk1, "b-1": dsk2}, dependencies={"a-1": set(), "b-1": {"a-1"}}
+            {"a-1": dsk1, "b-1": dsk2}, {"a-1": set(), "b-1": {"a-1"}}
         )
         dsk3 = HighLevelGraph(
             {"a-1": dsk1, "b-1": dsk2, "c-1": dsk3},
-            dependencies={"a-1": set(), "b-1": {"a-1"}, "c-1": {"b-1"}},
+            {"a-1": set(), "b-1": {"a-1"}, "c-1": {"b-1"}},
         )
-        dsk4 = HighLevelGraph({"d-1": dsk4, "e": dsk4b}, {})
+        dsk4 = HighLevelGraph({"d-1": dsk4, "e": dsk4b}, {"d-1": set(), "e": set()})
     else:
         dsk2.update(dsk1)
         dsk3.update(dsk2)
@@ -247,6 +247,15 @@ def test_bind(layers):
     assert bound2.compute() == (3,)
     assert cnt.n == 6
 
+    bound3 = bind(t4, t3, seed=1, assume_layers=layers)
+    cloned_d_name = clone_key("d-1", seed=1)
+    cloned_e_name = clone_key("e", seed=1)
+    assert bound3.__dask_graph__()[cloned_d_name, h1][0] is chunks.bind
+    assert bound3.__dask_graph__()[cloned_d_name, h2][0] is chunks.bind
+    assert bound3.__dask_graph__()[cloned_e_name][0] is chunks.bind
+    assert bound3.compute() == (1, 2, 3)
+    assert cnt.n == 9
+
 
 @pytest.mark.skipif("not da or not db or not dd")
 @pytest.mark.parametrize("func", [bind, clone])
@@ -312,3 +321,47 @@ def test_bind_clone_collections(func):
     assert_no_common_keys(ddf3c, ddf3, omit=ddf2, layers=True)
     assert_no_common_keys(ddf4c, ddf4, omit=ddf2, layers=True)
     assert_no_common_keys(ddf5c, ddf5, omit=ddf2, layers=True)
+
+
+@pytest.mark.parametrize(
+    "split_every,nkeys",
+    [
+        (2, 299),
+        (3, 250),
+        (8, 215),
+        (None, 215),  # default is 8
+        (8.1, 215),
+        (1e9, 201),
+        (False, 201),
+    ],
+)
+def test_split_every(split_every, nkeys):
+    dsk = {("a", i): i for i in range(100)}
+    t1 = Tuple(dsk, list(dsk))
+    c = checkpoint(t1, split_every=split_every)
+    assert len(c.__dask_graph__()) == nkeys
+    assert c.compute(scheduler="sync") is None
+
+    t2 = wait_on(t1, split_every=split_every)
+    assert len(t2.__dask_graph__()) == nkeys + 100
+    assert t2.compute(scheduler="sync") == tuple(range(100))
+
+    dsk3 = {"b": 1, "c": 2}
+    t3 = Tuple(dsk3, list(dsk3))
+    t4 = bind(t3, t1, split_every=split_every, assume_layers=False)
+    assert len(t4.__dask_graph__()) == nkeys + 2
+    assert t4.compute(scheduler="sync") == (1, 2)
+
+
+def test_split_every_invalid():
+    t = Tuple({"a": 1, "b": 2}, ["a", "b"])
+    with pytest.raises(ValueError):
+        checkpoint(t, split_every=1)
+    with pytest.raises(ValueError):
+        checkpoint(t, split_every=1.9)
+    with pytest.raises(ValueError):
+        checkpoint(t, split_every=0)  # Not to be confused with False or None
+    with pytest.raises(ValueError):
+        checkpoint(t, split_every=-2)
+    with pytest.raises(TypeError):
+        checkpoint(t, split_every={0: 2})  # This is legal for dask.array but not here