Repository: dask/dask

Issue:
Generically rebuild a collection with different keys
I'm coding a generic library function that takes in input an arbitrary dask collection and produces in output a variant of that same collection, which computes to the same values but with a different dask graph *and cohexists on the original collection*.

An example is https://github.com/dask/dask/pull/6804.

I can't find how to implement it in a way that supports third-party collections - not without making assumptions beyond those stipulated by https://docs.dask.org/en/latest/spec.html and https://docs.dask.org/en/latest/custom-collections.html.

I am therefore proposing to tighten the definitions in both documents to align them to the de-facto implementation of all builtin dask collections. This is also a good thing because I expect the dask codebase to already blindly assume the below in several places simply because it's true for Array, Bag, Dataframe, Series and Delayed:

#### 1
In https://docs.dask.org/en/latest/spec.html#definitions, replace

> A key is any hashable value that is not a task:

with

> A key is either a string or a tuple of 1+ elements, whose first element is a string and the remainder are arbitrary hashables:

#### 2
 In https://docs.dask.org/en/latest/custom-collections.html#__dask_keys__, add the requirement that, if more than one key is returned, all keys must be tuples with an identical first element (commonly referred to as the *name* of the collection)

My problem here is that I don't know anything about how custom collections look like in other projects. However, I also strongly suspect that they already adhere to the de facto standard above, simply because otherwise the dask codebase would fall over in multiple points.

#### 3
In https://docs.dask.org/en/latest/custom-collections.html#__dask_postpersist__, add that the returned callable must accept an optional keyword parameter, ``name``, which will be used to replace the collection's name.


## Example
Here we're introducing a 1 second delay in the compute() call of a dask bag:

```python
>>> import time
>>> import dask.bag
>>> from dask import delayed
>>> from dask.base import tokenize
>>> from dask.highlevelgraph import HighLevelGraph

>>> a = dask.bag.from_sequence([1,2,3,4], npartitions=2)
>>> dict(a.__dask_graph__())
{('from_sequence-00e91bb38e7830800935abba6e2614f8', 0): [1, 2],
 ('from_sequence-00e91bb38e7830800935abba6e2614f8', 1): [3, 4]}
>>> s = delayed(time.sleep)(1)
>>> name = "with_sleep-" + tokenize(a, s)
>>> def _with_sleep(node, deps):
...     return node

# Simplified - actual implementation is a bit more involved to make it
# work for arbitrary collections, e.g. 2+ dimensional arrays
>>> layer = {(name, i): (_with_sleep, k, s.key) for i, k in enumerate(a.__dask_keys__())}

>>> dsk = HighLevelGraph.from_collections(name, layer, dependencies=(a, s))
>>> rebuild, args = a.__dask_postpersist__()
>>> b = rebuild(dsk, *args, name=name)
>>> dict(b.__dask_graph__())
{('from_sequence-00e91bb38e7830800935abba6e2614f8', 0): [1, 2],
 ('from_sequence-00e91bb38e7830800935abba6e2614f8', 1): [3, 4],
 'sleep-1485da81-3fc0-4a0d-b099-7dc21ccc7d92': (<function time.sleep>, 1),
 ('with_sleep-0e771ca8ce0b7fd08e0b0dc7c4dde9d7', 0): (
     <function _with_sleep>, 
     ('from_sequence-00e91bb38e7830800935abba6e2614f8', 0), 
     'sleep-1485da81-3fc0-4a0d-b099-7dc21ccc7d92'
 ),
 ('with_sleep-0e771ca8ce0b7fd08e0b0dc7c4dde9d7', 1): (
     <function _with_sleep>,
     ('from_sequence-00e91bb38e7830800935abba6e2614f8', 1),
     'sleep-1485da81-3fc0-4a0d-b099-7dc21ccc7d92'
 ),
}
```
---
Additionally, it doesn't seem that https://docs.dask.org/en/latest/custom-collections.html has been updated since the introduction of HighLevelGraph.

Specifically:
- it states that the output of ``__dask_graph__`` is a MutableMapping; however, HighLevelGraph is a Mapping.
- it makes no mention of the ``__dask_layers__`` method. I think it would be sensible to stipulate that, if ``__dask_layers__`` is defined, then ``__dask_graph__`` must return a HighLevelGraph?

Discussion:
jsignell: As a meta point, I would prefer if we tried to avoid names with violent implications. 
crusaderky: > As a meta point, I would prefer if we tried to avoid names with violent implications.

...again, happy to accept suggestions...

- graph_manipulation
- graph_edit
- graph_rebuild
- regraph
- other?

What I wanted to convey that we're altering an already existing graph but, unlike optimization, the output is not functionally equivalent to the input.
jsignell: I like `graph_manipulation`. For `choke` I might suggest `restrict` or `throttle`.


jrbourbeau: Also +1 for `graph_manipulation.py`. Another alternative for `choke` could be `checkpoint` -- though perhaps @crusaderky has a sense for which name aligns best with the functionality of that particular graph manipulation 
crusaderky: @jrbourbeau you may start code review (but see TODOs above)
jrbourbeau: Also cc @willirath who might find these graph manipulations methods useful 
jrbourbeau: Thanks for the updates @crusaderky! I look forward to reviewing them later today

Patch:
diff --git a/dask/base.py b/dask/base.py
index d520a24da..34f8fc286 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -18,7 +18,7 @@ from .compatibility import is_dataclass, dataclass_fields
 from .context import thread_state
 from .core import flatten, quote, get as simple_get, literal
 from .hashing import hash_buffer_hex
-from .utils import Dispatch, ensure_dict, apply
+from .utils import Dispatch, ensure_dict, apply, key_split
 from . import config, local, threaded
 
 
@@ -34,6 +34,7 @@ __all__ = (
     "normalize_token",
     "get_collection_name",
     "replace_name_in_key",
+    "clone_key",
 )
 
 
@@ -1239,3 +1240,26 @@ def replace_name_in_key(key, name: str):
     if isinstance(key, str):
         return name
     raise TypeError(f"Expected str or tuple[str, Hashable, ...]; got {key}")
+
+
+def clone_key(key, seed):
+    """Clone a key from a Dask collection, producing a new key with the same prefix and
+    indices and a token which a deterministic function of the previous token and seed.
+
+    Examples
+    --------
+    >>> clone_key("inc-cbb1eca3bafafbb3e8b2419c4eebb387", 123)
+    'inc-1d291de52f5045f8a969743daea271fd'
+    >>> clone_key(("sum-cbb1eca3bafafbb3e8b2419c4eebb387", 4, 3), 123)
+    ('sum-f0962cc58ef4415689a86cc1d4cc1723', 4, 3)
+    """
+    if isinstance(key, tuple) and key and isinstance(key[0], str):
+        return (clone_key(key[0], seed),) + key[1:]
+    if isinstance(key, str):
+        prefix = key_split(key)
+        token = key[len(prefix) + 1 :]
+        if token:
+            return prefix + "-" + tokenize(token, seed)
+        else:
+            return tokenize(key, seed)
+    raise TypeError(f"Expected str or tuple[str, Hashable, ...]; got {key}")
diff --git a/dask/graph_manipulation.py b/dask/graph_manipulation.py
new file mode 100644
index 000000000..f4e1fbf92
--- /dev/null
+++ b/dask/graph_manipulation.py
@@ -0,0 +1,426 @@
+"""Tools to modify already existing dask graphs. Unlike in :mod:`dask.optimization`, the
+output collections produced by this module are typically not functionally equivalent to
+their inputs.
+"""
+import uuid
+from typing import Callable, Hashable, Optional, Set, Tuple, TypeVar
+
+from .array import Array
+from .bag import Bag
+from .base import (
+    clone_key,
+    get_collection_name,
+    replace_name_in_key,
+    tokenize,
+    unpack_collections,
+)
+from .blockwise import blockwise
+from .core import flatten
+from .dataframe import DataFrame, Series
+from .delayed import Delayed, delayed
+from .highlevelgraph import BasicLayer, HighLevelGraph, Layer
+
+__all__ = ("bind", "checkpoint", "clone", "wait_on")
+
+T = TypeVar("T")
+
+
+def checkpoint(*collections) -> Delayed:
+    """Build a :doc:`delayed` which waits until all chunks of the input collection(s)
+    have been computed before returning None.
+
+    Parameters
+    ----------
+    collections
+        Zero or more Dask collections or nested data structures containing zero or more
+        collections
+
+    Returns
+    -------
+    :doc:`delayed` yielding None
+    """
+    collections, _ = unpack_collections(*collections)
+    if len(collections) != 1:
+        return delayed(chunks.checkpoint)(*(checkpoint(c) for c in collections))
+    (collection,) = collections
+
+    tok = tokenize(collection)
+    name = "checkpoint-" + tok
+    map_name = "checkpoint_map-" + tok
+
+    keys_iter = flatten(collection.__dask_keys__())
+    try:
+        next(keys_iter)
+        next(keys_iter)
+    except StopIteration:
+        # 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
+        map_layer = _build_map_layer(chunks.checkpoint, map_name, collection)
+        map_dsk = HighLevelGraph.from_collections(
+            map_name, map_layer, dependencies=(collection,)
+        )
+        reduce_layer = {name: (chunks.checkpoint, list(map_layer.get_output_keys()))}
+        reduce_dsk = HighLevelGraph(
+            {name: reduce_layer}, dependencies={name: {map_name}}
+        )
+        dsk = HighLevelGraph.merge(map_dsk, reduce_dsk)
+
+    return Delayed(name, dsk)
+
+
+def _build_map_layer(
+    func: Callable, name: str, collection, dependencies: Tuple[Delayed, ...] = ()
+) -> Layer:
+    """Apply func to all keys of collection. Create a Blockwise layer whenever possible;
+    fall back to BasicLayer otherwise.
+
+    Parameters
+    ----------
+    func
+        Callable to be invoked on the graph node
+    name : str
+        Layer name
+    collection
+        Arbitrary dask collection
+    dependencies
+        Zero or more Delayed objects, which will be passed as arbitrary variadic args to
+        func after the collection's chunk
+    """
+    if isinstance(collection, (Array, DataFrame, Series, Bag)):
+        # Use a Blockwise layer
+        if isinstance(collection, Array):
+            numblocks = collection.numblocks
+        else:
+            numblocks = (collection.npartitions,)
+        indices = tuple(i for i, _ in enumerate(numblocks))
+        kwargs = {"_deps": dependencies} if dependencies else {}
+        prev_name = get_collection_name(collection)
+        return blockwise(
+            func,
+            name,
+            indices,
+            prev_name,
+            indices,
+            numblocks={prev_name: numblocks},
+            dependencies=dependencies,
+            **kwargs,
+        )
+    else:
+        # Delayed, bag.Item, dataframe.core.Scalar, or third-party collection;
+        # fall back to BasicLayer
+        dep_keys = tuple(d.key for d in dependencies)
+        return BasicLayer(
+            {
+                replace_name_in_key(k, name): (func, k) + dep_keys
+                for k in flatten(collection.__dask_keys__())
+            }
+        )
+
+
+def bind(
+    children: T,
+    parents,
+    *,
+    omit=None,
+    seed: Hashable = None,
+    assume_layers: bool = True,
+) -> T:
+    """
+    Make ``children`` collection(s), optionally omitting sub-collections, dependent on
+    ``parents`` collection(s). Two examples follow.
+
+    The first example creates an array ``b2`` whose computation first computes an array
+    ``a`` completely and then computes ``b`` completely, recomputing ``a`` in the
+    process:
+
+    >>> from dask import array as da
+    >>> a = da.ones(4, chunks=2)
+    >>> b = a + 1
+    >>> b2 = bind(b, a)
+    >>> len(b2.dask)
+    11
+    >>> b2.compute()
+    array([2., 2., 2., 2.])
+
+    The second example creates arrays ``b3`` and ``c3``, whose computation first
+    computes an array ``a`` and then computes the additions, this time not
+    recomputing ``a`` in the process:
+
+    >>> c = a + 2
+    >>> b3, c3 = bind((b, c), a, omit=a)
+    >>> len(b3.dask), len(c3.dask)
+    (9, 9)
+    >>> dask.compute(b3, c3)
+    (array([2., 2., 2., 2.]), array([3., 3., 3., 3.]))
+
+    Parameters
+    ----------
+    children
+        Dask collection or nested structure of Dask collections
+    parents
+        Dask collection or nested structure of Dask collections
+    omit
+        Dask collection or nested structure of Dask collections
+    seed
+        Hashable used to seed the key regeneration. Omit to default to a random number
+        that will produce different keys at every call.
+    assume_layers
+        True
+            Use a fast algorithm that works at layer level, which assumes that all
+            collections in ``children`` and ``omit``
+
+            #. use :class:`~dask.highlevelgraph.HighLevelGraph`,
+            #. define the ``__dask_layers__()`` method, and
+            #. never had their graphs squashed and rebuilt between the creation of the
+               ``omit`` collections and the ``children`` collections; in other words if
+               the keys of the ``omit`` collections can be found among the keys of the
+               ``children`` collections, then the same must also hold true for the
+               layers.
+        False
+            Use a slower algorithm that works at keys level, which makes none of the
+            above assumptions.
+
+    Returns
+    -------
+    Same as ``children``
+        Dask collection or structure of dask collection equivalent to ``children``,
+        which compute to the same values. All keys of ``children`` will be regenerated,
+        up to and excluding the keys of ``omit``. Nodes immediately above ``omit``, or
+        the leaf nodes if the collections in ``omit`` are not found, are prevented from
+        computing until all collections in ``parents`` have been fully computed.
+    """
+    if seed is None:
+        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
+
+    omit, _ = unpack_collections(omit)
+    if assume_layers:
+        # Set of all the top-level layers of the collections in omit
+        omit_layers = {layer for coll in omit for layer in coll.__dask_layers__()}
+        omit_keys = set()
+    else:
+        omit_layers = set()
+        # Set of *all* the keys, not just the top-level ones, of the collections in omit
+        omit_keys = {key for coll in omit for key in coll.__dask_graph__()}
+
+    unpacked_children, repack = unpack_collections(children)
+    return repack(
+        [
+            _bind_one(child, blocker, omit_layers, omit_keys, seed)
+            for child in unpacked_children
+        ]
+    )[0]
+
+
+def _bind_one(
+    child: T,
+    blocker: Optional[Delayed],
+    omit_layers: Set[str],
+    omit_keys: Set[Hashable],
+    seed: Hashable,
+) -> T:
+    try:
+        prev_coll_name = get_collection_name(child)
+    except KeyError:
+        # Collection with no keys; this is a legitimate use case but, at the moment of
+        # writing, can only happen with third-party collections
+        return child
+
+    dsk = child.__dask_graph__()
+    new_layers = {}
+    new_deps = {}
+    if not isinstance(dsk, HighLevelGraph):
+        dsk = HighLevelGraph.from_collections(prev_coll_name, dsk)
+
+    clone_keys = dsk.get_all_external_keys() - omit_keys
+    for layer_name in omit_layers:
+        try:
+            layer = dsk.layers[layer_name]
+        except KeyError:
+            continue
+        clone_keys -= layer.get_output_keys()
+    # Note: when assume_layers=True, clone_keys can contain keys of the omit collections
+    # that are not top-level. This is OK, as they will never be encountered inside the
+    # values of their dependent layers.
+
+    if blocker is not None:
+        blocker_key = blocker.key
+        blocker_dsk = blocker.__dask_graph__()
+        assert isinstance(blocker_dsk, HighLevelGraph)
+        new_layers.update(blocker_dsk.layers)
+        new_deps.update(blocker_dsk.dependencies)
+    else:
+        blocker_key = None
+
+    try:
+        layers_to_clone = set(child.__dask_layers__())
+    except AttributeError:
+        layers_to_clone = {prev_coll_name}
+    layers_to_copy_verbatim = set()
+
+    while layers_to_clone:
+        prev_layer_name = layers_to_clone.pop()
+        new_layer_name = clone_key(prev_layer_name, seed=seed)
+        if new_layer_name in new_layers:
+            continue
+
+        layer = dsk.layers[prev_layer_name]
+        layer_deps = dsk.dependencies[prev_layer_name]
+        layer_deps_to_clone = layer_deps - omit_layers
+        layer_deps_to_omit = layer_deps & omit_layers
+        layers_to_clone |= layer_deps_to_clone
+        layers_to_copy_verbatim |= layer_deps_to_omit
+
+        new_layers[new_layer_name], is_bound = layer.clone(
+            keys=clone_keys, seed=seed, bind_to=blocker_key
+        )
+        new_deps[new_layer_name] = {
+            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)
+
+    # 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,
+    # omit[i].dsk). Also, collections in omit may or may not be parents of this specific
+    # child, or of any children at all.
+    while layers_to_copy_verbatim:
+        layer_name = layers_to_copy_verbatim.pop()
+        if layer_name in new_layers:
+            continue
+        layer_deps = dsk.dependencies[layer_name]
+        layers_to_copy_verbatim |= layer_deps
+        new_deps[layer_name] = layer_deps
+        new_layers[layer_name] = dsk.layers[layer_name]
+
+    rebuild, args = child.__dask_postpersist__()
+    return rebuild(
+        HighLevelGraph(new_layers, new_deps),
+        *args,
+        name=clone_key(prev_coll_name, seed),
+    )
+
+
+def clone(*collections, omit=None, seed: Hashable = None, assume_layers: bool = True):
+    """Clone dask collections, returning equivalent collections that are generated from
+    independent calculations.
+
+    Example
+    -------
+    (tokens have been simplified for the sake of brevity)
+
+    >>> from dask import array as da
+    >>> x_i = da.asarray([1, 1, 1, 1], chunks=2)
+    >>> y_i = x_i + 1
+    >>> z_i = y_i + 2
+    >>> dict(z_i.dask)
+    {('array-1', 0): array([1, 1]),
+     ('array-1', 1): array([1, 1]),
+     ('add-2', 0): (<function operator.add>, ('array-1', 0), 1),
+     ('add-2', 1): (<function operator.add>, ('array-1', 1), 1),
+     ('add-3', 0): (<function operator.add>, ('add-2', 0), 1),
+     ('add-3', 1): (<function operator.add>, ('add-2', 1), 1)}
+    >>> w_i = clone(z_i, omit=x_i)
+    >>> w_i.compute()
+    array([4., 4., 4., 4.])
+    >>> dict(w_i.dask)
+    {('array-1', 0): array([1, 1]),
+     ('array-1', 1): array([1, 1]),
+     ('add-4', 0): (<function operator.add>, ('array-1', 0), 1),
+     ('add-4', 1): (<function operator.add>, ('array-1', 1), 1),
+     ('add-5', 0): (<function operator.add>, ('add-4', 0), 1),
+     ('add-5', 1): (<function operator.add>, ('add-4', 1), 1)}
+
+    Parameters
+    ----------
+    collections
+        Zero or more Dask collections or nested structures of Dask collections
+    omit
+        Dask collection or nested structure of Dask collections which will not be cloned
+    seed
+        See :func:`bind`
+    assume_layers
+        See :func:`bind`
+
+    Returns
+    -------
+    Same as ``collections``
+        Dask collections of the same type as the inputs, which compute to the same
+        value, or nested structures equivalent to the inputs, where the original
+        collections have been replaced.
+    """
+    out = bind(
+        collections, parents=None, omit=omit, seed=seed, assume_layers=assume_layers
+    )
+    return out[0] if len(collections) == 1 else out
+
+
+def wait_on(*collections):
+    """Ensure that all chunks of all input collections have been computed before
+    computing the dependents of any of the chunks.
+
+    The following example creates a dask array ``u`` that, when used in a computation,
+    will only proceed when all chunks of the array ``x`` have been computed, but
+    otherwise matches ``x``:
+
+    >>> from dask import array as da
+    >>> x = da.ones(10, chunks=5)
+    >>> u = wait_on(x)
+
+    The following example will create two arrays ``u`` and ``v`` that, when used in a
+    computation, will only proceed when all chunks of the arrays ``x`` and ``y`` have
+    been computed but otherwise match ``x`` and ``y``:
+
+    >>> x = da.ones(10, chunks=5)
+    >>> y = da.zeros(10, chunks=5)
+    >>> u, v = wait_on(x, y)
+
+    Parameters
+    ----------
+    collections
+        Zero or more Dask collections or nested structures of Dask collections
+
+    Returns
+    -------
+    Same as ``collections``
+        Dask collection of the same type as the input, which computes to the same value,
+        or a nested structure equivalent to the input where the original collections
+        have been replaced.
+    """
+    blocker = checkpoint(*collections)
+
+    def block_one(coll):
+        name = "wait_on-" + tokenize(coll, blocker)
+        layer = _build_map_layer(chunks.bind, name, coll, dependencies=(blocker,))
+        dsk = HighLevelGraph.from_collections(name, layer, dependencies=(coll, blocker))
+        rebuild, args = coll.__dask_postpersist__()
+        return rebuild(dsk, *args, name=name)
+
+    unpacked, repack = unpack_collections(*collections)
+    out = repack([block_one(coll) for coll in unpacked])
+    return out[0] if len(collections) == 1 else out
+
+
+class chunks:
+    """Callables to be inserted in the Dask graph"""
+
+    @staticmethod
+    def bind(node: T, *args, **kwargs) -> T:
+        """Dummy graph node of :func:`bind` and :func:`wait_on`.
+        Wait for both node and all variadic args to complete; then return node.
+        """
+        return node
+
+    @staticmethod
+    def checkpoint(*args, **kwargs) -> None:
+        """Dummy graph node of :func:`checkpoint`.
+        Wait for all variadic args to complete; then return None.
+        """
+        pass
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index 7c9ae9dca..dceffa202 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -18,7 +18,7 @@ import tlz as toolz
 
 from . import config
 from .utils import ensure_dict, ignoring, stringify
-from .base import is_dask_collection
+from .base import clone_key, is_dask_collection
 from .core import reverse_dict, keys_in_tasks
 from .utils_test import add, inc  # noqa: F401
 
@@ -211,6 +211,81 @@ class Layer(collections.abc.Mapping):
         """
         return keys_in_tasks(all_hlg_keys, [self[key]])
 
+    def clone(
+        self,
+        keys: set,
+        seed: Hashable,
+        bind_to: Hashable = None,
+    ) -> "tuple[Layer, bool]":
+        """Clone selected keys in the layer, as well as references to keys in other
+        layers
+
+        Parameters
+        ----------
+        keys
+            Keys to be replaced. This never includes keys not listed by
+            :meth:`get_output_keys`. It must also include any keys that are outside
+            of this layer that may be referenced by it.
+        seed
+            Common hashable used to alter the keys; see :func:`dask.base.clone_key`
+        bind_to
+            Optional key to bind the leaf nodes to. A leaf node here is one that does
+            not reference any replaced keys; in other words it's a node where the
+            replacement graph traversal stops; it may still have dependencies on
+            non-replaced nodes.
+            A bound node will not be computed until after ``bind_to`` has been computed.
+
+        Returns
+        -------
+        - New layer
+        - True if the ``bind_to`` key was injected anywhere; False otherwise
+
+        Notes
+        -----
+        This method should be overridden by subclasses to avoid materializing the layer.
+        """
+        from .graph_manipulation import chunks
+
+        is_leaf: bool
+
+        def clone_value(o):
+            """Variant of distributed.utils_comm.subs_multiple, which allows injecting
+            bind_to
+            """
+            nonlocal is_leaf
+
+            typ = type(o)
+            if typ is tuple and o and callable(o[0]):
+                return (o[0],) + tuple(clone_value(i) for i in o[1:])
+            elif typ is list:
+                return [clone_value(i) for i in o]
+            elif typ is dict:
+                return {k: clone_value(v) for k, v in o.items()}
+            else:
+                try:
+                    if o not in keys:
+                        return o
+                except TypeError:
+                    return o
+                is_leaf = False
+                return clone_key(o, seed)
+
+        dsk_new = {}
+        bound = False
+
+        for key, value in self.items():
+            if key in keys:
+                key = clone_key(key, seed)
+                is_leaf = True
+                value = clone_value(value)
+                if bind_to is not None and is_leaf:
+                    value = (chunks.bind, value, bind_to)
+                    bound = True
+
+            dsk_new[key] = value
+
+        return BasicLayer(dsk_new), bound
+
     def __dask_distributed_pack__(self, client) -> Optional[Any]:
         """Pack the layer for scheduler communication in Distributed
 
diff --git a/docs/source/graph_manipulation.rst b/docs/source/graph_manipulation.rst
new file mode 100644
index 000000000..3953382ed
--- /dev/null
+++ b/docs/source/graph_manipulation.rst
@@ -0,0 +1,72 @@
+.. _graph_manipulation:
+
+Advanced graph manipulation
+===========================
+There are some situations where computations with Dask collections will result in
+suboptimal memory usage (e.g. an entire Dask DataFrame is loaded into memory).
+This may happen when Dask’s scheduler doesn’t automatically delay the computation of
+nodes in a task graph to avoid occupying memory with their output for prolonged periods
+of time, or in scenarios where recalculating nodes is much cheaper than holding their
+output in memory.
+
+This page highlights a set of graph manipulation utilities which can be used to help
+avoid these scenarios. In particular, the utilities described below rewrite the
+underlying Dask graph for Dask collections, producing equivalent collections with
+different sets of keys.
+
+Consider the following example:
+
+.. code-block:: python
+
+   >>> import dask.array as da
+   >>> x = da.random.normal(size=500_000_000, chunks=100_000)
+   >>> x_mean = x.mean()
+   >>> y = (x - x_mean).max().compute()
+
+The above example computes the largest value of a distribution after removing its bias.
+This involves loading the chunks of ``x`` into memory in order to compute ``x_mean``.
+However, since the ``x`` array is needed later in the computation to compute ``y``, the
+entire ``x`` array is kept in memory. For large Dask Arrays this can be very
+problematic.
+
+To alleviate the need for the entire ``x`` array to be kept in memory, one could rewrite
+the last line as follows:
+
+.. code-block:: python
+
+   >>> from dask.graph_manipulation import bind
+   >>> xb = bind(x, x_mean)
+   >>> y = (xb - x_mean).max().compute()
+
+Here we use :func:`~dask.graph_manipulation.bind` to create a new Dask Array, ``xb``,
+which produces exactly the same output as ``x``, but whose underlying Dask graph has
+different keys than ``x``, and will only be computed after ``x_mean`` has been
+calculated.
+
+This results in the chunks of ``x`` being computed and immediately individually reduced
+by ``mean``; then recomputed and again immediately pipelined into the subtraction
+followed by reduction with ``max``. This results in a much smaller peak memory usage as
+the full ``x`` array is no longer loaded into memory. However, the tradeoff is that the
+compute time increases as ``x`` is computed twice.
+
+
+API
+---
+
+.. currentmodule:: dask.graph_manipulation
+
+.. autosummary::
+
+   checkpoint
+   wait_on
+   bind
+   clone
+
+
+Definitions
+~~~~~~~~~~~
+
+.. autofunction:: checkpoint
+.. autofunction:: wait_on
+.. autofunction:: bind
+.. autofunction:: clone
diff --git a/docs/source/graphs.rst b/docs/source/graphs.rst
index 988a5c258..85ac94e50 100644
--- a/docs/source/graphs.rst
+++ b/docs/source/graphs.rst
@@ -15,6 +15,7 @@ often a better choice. If you are a *core developer*, then you should start here
    spec.rst
    custom-graphs.rst
    optimize.rst
+   graph_manipulation.rst
    custom-collections.rst
    high-level-graphs.rst
 


Test patch:
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index dacedca80..22353265a 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -27,20 +27,15 @@ from dask.base import (
     collections_to_dsk,
     get_collection_name,
     replace_name_in_key,
+    clone_key,
 )
 from dask.core import literal
 from dask.delayed import Delayed
-from dask.utils import tmpdir, tmpfile, ignoring
-from dask.utils_test import inc, dec
+from dask.utils import tmpdir, tmpfile
+from dask.utils_test import dec, inc, import_or_none
 from dask.diagnostics import Profiler
 
 
-def import_or_none(path):
-    with ignoring(BaseException):
-        return pytest.importorskip(path)
-    return None
-
-
 tz = pytest.importorskip("tlz")
 da = import_or_none("dask.array")
 db = import_or_none("dask.bag")
@@ -615,6 +610,9 @@ class Tuple(DaskMethodsMixin):
     def __dask_keys__(self):
         return self._keys
 
+    def __dask_layers__(self):
+        return (get_collection_name(self),)
+
     def __dask_tokenize__(self):
         return self._keys
 
@@ -1354,3 +1352,17 @@ def test_optimizations_ctd():
         dsk2 = collections_to_dsk([x])
 
     assert dsk1 == dsk2
+
+
+def test_clone_key():
+    h = object()  # arbitrary hashable
+    assert clone_key("inc-1-2-3", 123) == "inc-27b6e15b795fcaff169e0e0df14af97a"
+    assert clone_key("x", 123) == "dc2b8d1c184c72c19faa81c797f8c6b0"
+    assert clone_key("x", 456) == "b76f061b547b00d18b9c7a18ccc47e2d"
+    assert clone_key(("sum-1-2-3", h, 1), 123) == (
+        "sum-27b6e15b795fcaff169e0e0df14af97a",
+        h,
+        1,
+    )
+    with pytest.raises(TypeError):
+        clone_key(1, 2)
diff --git a/dask/tests/test_graph_manipulation.py b/dask/tests/test_graph_manipulation.py
new file mode 100644
index 000000000..69f6b44d4
--- /dev/null
+++ b/dask/tests/test_graph_manipulation.py
@@ -0,0 +1,299 @@
+import random
+import time
+from operator import add
+
+import pytest
+
+import dask
+from dask import delayed
+from dask.base import clone_key
+from dask.graph_manipulation import bind, checkpoint, chunks, clone, wait_on
+from dask.highlevelgraph import HighLevelGraph
+from dask.tests.test_base import Tuple
+from dask.utils_test import import_or_none
+
+da = import_or_none("dask.array")
+db = import_or_none("dask.bag")
+dd = import_or_none("dask.dataframe")
+pd = import_or_none("pandas")
+
+
+class NodeCounter:
+    def __init__(self):
+        self.n = 0
+
+    def f(self, x):
+        time.sleep(random.random() / 100)
+        self.n += 1
+        return x
+
+
+def assert_no_common_keys(a, b, omit=None, *, layers: bool) -> None:
+    dsk1 = a.__dask_graph__()
+    dsk2 = b.__dask_graph__()
+
+    if omit is not None:
+        dsko = omit.__dask_graph__()
+        assert not (dsk1.keys() - dsko.keys()) & dsk2.keys()
+        assert not dsko.keys() - dsk1.keys()
+        assert not dsko.keys() - dsk2.keys()
+        if layers:
+            assert not (dsk1.layers.keys() - dsko.layers.keys()) & dsk2.layers.keys()
+            assert (
+                not (dsk1.dependencies.keys() - dsko.dependencies.keys())
+                & dsk2.dependencies.keys()
+            )
+            assert not dsko.layers.keys() - dsk1.layers.keys()
+            assert not dsko.layers.keys() - dsk2.layers.keys()
+            assert not dsko.dependencies.keys() - dsk1.dependencies.keys()
+            assert not dsko.dependencies.keys() - dsk2.dependencies.keys()
+
+    else:
+        assert not dsk1.keys() & dsk2.keys()
+        if layers:
+            assert not dsk1.layers.keys() & dsk2.layers.keys()
+            assert not dsk1.dependencies.keys() & dsk2.dependencies.keys()
+
+
+# Generic hashables
+h1 = object()
+h2 = object()
+
+
+def collections_with_node_counters():
+    cnt = NodeCounter()
+    df = pd.DataFrame({"x": list(range(10))})
+
+    # Test two samples of all collections where applicable, one with multiple chunks
+    # and one with a single chunk
+    colls = [
+        # dask.delayed
+        delayed(cnt.f)("Hello 1"),  # 1 chunk
+        # dask.array
+        da.ones((10, 10), chunks=5).map_blocks(cnt.f),  # 4 chunks
+        da.ones((1,), chunks=-1).map_blocks(cnt.f),  # 1 chunk
+        # dask.bag
+        db.from_sequence([1, 2], npartitions=2).map(cnt.f),  # 2 chunks
+        db.from_sequence([1], npartitions=1).map(cnt.f),  # 1 chunk
+        db.Item.from_delayed(delayed(cnt.f)("Hello 2")),  # 1 chunk
+        # dask.dataframe
+        dd.from_pandas(df, npartitions=2).map_partitions(cnt.f),  # 2 chunks
+        dd.from_pandas(df, npartitions=1).map_partitions(cnt.f),  # 1 chunk
+        dd.from_pandas(df["x"], npartitions=2).map_partitions(cnt.f),  # 2 chunks
+        dd.from_pandas(df["x"], npartitions=1).map_partitions(cnt.f),  # 1 chunk
+    ]
+    cnt.n = 0
+    return colls, cnt
+
+
+def test_checkpoint():
+    cnt = NodeCounter()
+    dsk1 = {("a", h1): (cnt.f, 1), ("a", h2): (cnt.f, 2)}
+    dsk2 = {"b": (cnt.f, 2)}
+    cp = checkpoint(Tuple(dsk1, list(dsk1)), {"x": [Tuple(dsk2, list(dsk2))]})
+    assert cp.compute(scheduler="sync") is None
+    assert cnt.n == 3
+
+
+@pytest.mark.skipif("not da or not db or not dd")
+def test_checkpoint_collections():
+    colls, cnt = collections_with_node_counters()
+    cp = checkpoint(*colls)
+    cp.compute(scheduler="sync")
+    assert cnt.n == 16
+
+
+def test_wait_on_one():
+    cnt = NodeCounter()
+    dsk = {("a", h1): (cnt.f, 1), ("a", h2): (cnt.f, 2)}
+    t = wait_on(Tuple(dsk, list(dsk)))
+    assert t.compute(scheduler="sync") == (1, 2)
+    assert cnt.n == 2
+
+
+def test_wait_on_many():
+    cnt = NodeCounter()
+    dsk1 = {("a", h1): (cnt.f, 1), ("a", h2): (cnt.f, 2)}
+    dsk2 = {"b": (cnt.f, 3)}
+    out = wait_on(Tuple(dsk1, list(dsk1)), {"x": [Tuple(dsk2, list(dsk2))]})
+    assert dask.compute(*out, scheduler="sync") == ((1, 2), {"x": [(3,)]})
+    assert cnt.n == 3
+
+
+@pytest.mark.skipif("not da or not db or not dd")
+def test_wait_on_collections():
+    colls, cnt = collections_with_node_counters()
+
+    # Create a delayed that depends on a single one among all collections
+    @delayed
+    def f(x):
+        pass
+
+    colls2 = wait_on(*colls)
+    f(colls2[0]).compute()
+    assert cnt.n == 16
+
+    # dask.delayed
+    assert colls2[0].compute() == colls[0].compute()
+    # dask.array
+    da.utils.assert_eq(colls2[1], colls[1])
+    da.utils.assert_eq(colls2[2], colls[2])
+    # dask.bag
+    db.utils.assert_eq(colls2[3], colls[3])
+    db.utils.assert_eq(colls2[4], colls[4])
+    db.utils.assert_eq(colls2[5], colls[5])
+    # dask.dataframe
+    dd.utils.assert_eq(colls2[6], colls[6])
+    dd.utils.assert_eq(colls2[7], colls[7])
+    dd.utils.assert_eq(colls2[8], colls[8])
+    dd.utils.assert_eq(colls2[9], colls[9])
+
+
+@pytest.mark.parametrize("layers", [False, True])
+def test_clone(layers):
+    dsk1 = {("a", h1): 1, ("a", h2): 2}
+    dsk2 = {"b": (add, ("a", h1), ("a", h2))}
+    dsk3 = {"c": 1}
+    if layers:
+        dsk1 = HighLevelGraph.from_collections("a", dsk1)
+        dsk2 = HighLevelGraph(
+            {"a": dsk1, "b": dsk2}, dependencies={"a": set(), "b": {"a"}}
+        )
+        dsk3 = HighLevelGraph.from_collections("c", dsk3)
+    else:
+        dsk2.update(dsk1)
+
+    t1 = Tuple(dsk1, [("a", h1), ("a", h2)])
+    t2 = Tuple(dsk2, ["b"])
+    t3 = Tuple(dsk3, ["c"])
+
+    c1 = clone(t2, seed=1, assume_layers=layers)
+    c2 = clone(t2, seed=1, assume_layers=layers)
+    c3 = clone(t2, seed=2, assume_layers=layers)
+    c4 = clone(c1, seed=1, assume_layers=layers)  # Clone of a clone has different keys
+    c5 = clone(t2, assume_layers=layers)  # Random seed
+    c6 = clone(t2, assume_layers=layers)  # Random seed
+    c7 = clone(t2, omit=t1, seed=1, assume_layers=layers)
+
+    assert c1.__dask_graph__() == c2.__dask_graph__()
+    assert_no_common_keys(c1, t2, layers=layers)
+    assert_no_common_keys(c1, c3, layers=layers)
+    assert_no_common_keys(c1, c4, layers=layers)
+    assert_no_common_keys(c1, c5, layers=layers)
+    assert_no_common_keys(c5, c6, layers=layers)
+    assert_no_common_keys(c7, t2, omit=t1, layers=layers)
+    assert dask.compute(t2, c1, c2, c3, c4, c5, c6, c7) == ((3,),) * 8
+
+    # Clone nested; some of the collections in omit are unrelated
+    out = clone({"x": [t2]}, omit={"y": [t1, t3]}, assume_layers=layers)
+    assert dask.compute(out) == ({"x": [(3,)]},)
+    c8 = out["x"][0]
+    assert_no_common_keys(c8, t2, omit=t1, layers=layers)
+    assert_no_common_keys(c8, t3, layers=layers)
+
+
+@pytest.mark.parametrize("layers", [False, True])
+def test_bind(layers):
+    dsk1 = {("a-1", h1): 1, ("a-1", h2): 2}
+    dsk2 = {"b-1": (add, ("a-1", h1), ("a-1", h2))}
+    dsk3 = {"c-1": "b-1"}
+    cnt = NodeCounter()
+    dsk4 = {("d-1", h1): (cnt.f, 1), ("d-1", h2): (cnt.f, 2)}
+
+    if layers:
+        dsk1 = HighLevelGraph.from_collections("a-1", dsk1)
+        dsk2 = HighLevelGraph(
+            {"a-1": dsk1, "b-1": dsk2}, dependencies={"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"}},
+        )
+        dsk4 = HighLevelGraph.from_collections("d-1", dsk4)
+    else:
+        dsk2.update(dsk1)
+        dsk3.update(dsk2)
+
+    # t1 = Tuple(dsk1, [("a", h1), ("a", h2)])
+    t2 = Tuple(dsk2, ["b-1"])
+    t3 = Tuple(dsk3, ["c-1"])
+    t4 = Tuple(dsk4, [("d-1", h1), ("d-1", h2)])
+
+    bound1 = bind(t3, t4, seed=1, assume_layers=layers)
+    cloned_a_name = clone_key("a-1", seed=1)
+    assert bound1.__dask_graph__()[cloned_a_name, h1][0] is chunks.bind
+    assert bound1.__dask_graph__()[cloned_a_name, h2][0] is chunks.bind
+    assert bound1.compute() == (3,)
+    assert cnt.n == 2
+
+    bound2 = bind(t3, t4, omit=t2, seed=1, assume_layers=layers)
+    cloned_c_name = clone_key("c-1", seed=1)
+    assert bound2.__dask_graph__()[cloned_c_name][0] is chunks.bind
+    assert bound2.compute() == (3,)
+    assert cnt.n == 4
+
+
+@pytest.mark.skipif("not da or not db or not dd")
+@pytest.mark.parametrize("func", [bind, clone])
+def test_bind_clone_collections(func):
+    @delayed
+    def double(x):
+        return x * 2
+
+    # dask.delayed
+    d1 = double(2)
+    d2 = double(d1)
+    # dask.array
+    a1 = da.ones((10, 10), chunks=5)
+    a2 = a1 + 1
+    a3 = a2.T
+    # dask.bag
+    b1 = db.from_sequence([1, 2], npartitions=2)
+    # b1's tasks are not callable, so we need an extra step to properly test bind
+    b2 = b1.map(lambda x: x * 2)
+    b3 = b2.map(lambda x: x + 1)
+    b4 = b3.min()
+    # dask.dataframe
+    df = pd.DataFrame({"x": list(range(10))})
+    ddf1 = dd.from_pandas(df, npartitions=2)
+    # ddf1's tasks are not callable, so we need an extra step to properly test bind
+    ddf2 = ddf1.map_partitions(lambda x: x * 2)
+    ddf3 = ddf2.map_partitions(lambda x: x + 1)
+    ddf4 = ddf3["x"]  # dd.Series
+    ddf5 = ddf4.min()  # dd.Scalar
+
+    cnt = NodeCounter()
+    if func is bind:
+        parent = da.ones((10, 10), chunks=5).map_blocks(cnt.f)
+        cnt.n = 0
+        d2c, a3c, b3c, b4c, ddf3c, ddf4c, ddf5c = bind(
+            children=(d2, a3, b3, b4, ddf3, ddf4, ddf5),
+            parents=parent,
+            omit=(d1, a1, b2, ddf2),
+        )
+    else:
+        d2c, a3c, b3c, b4c, ddf3c, ddf4c, ddf5c = clone(
+            d2, a3, b3, b4, ddf3, ddf4, ddf5, omit=(d1, a1, b2, ddf2)
+        )
+
+    assert d2.compute() == d2c.compute()
+    assert cnt.n == 4 or func is clone
+    da.utils.assert_eq(a3c, a3)
+    assert cnt.n == 8 or func is clone
+    db.utils.assert_eq(b3c, b3)
+    assert cnt.n == 12 or func is clone
+    db.utils.assert_eq(b4c, b4)
+    assert cnt.n == 16 or func is clone
+    dd.utils.assert_eq(ddf3c, ddf3)
+    assert cnt.n == 24 or func is clone  # dd.utils.assert_eq calls compute() twice
+    dd.utils.assert_eq(ddf4c, ddf4)
+    assert cnt.n == 32 or func is clone  # dd.utils.assert_eq calls compute() twice
+    dd.utils.assert_eq(ddf5c, ddf5)
+    assert cnt.n == 36 or func is clone
+
+    assert_no_common_keys(d2c, d2, omit=d1, layers=True)
+    assert_no_common_keys(a3c, a3, omit=a1, layers=True)
+    assert_no_common_keys(b3c, b3, omit=b2, layers=True)
+    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)
diff --git a/dask/utils_test.py b/dask/utils_test.py
index 28da57020..d7189806b 100644
--- a/dask/utils_test.py
+++ b/dask/utils_test.py
@@ -1,3 +1,6 @@
+import importlib
+
+
 def inc(x):
     return x + 1
 
@@ -112,3 +115,11 @@ class GetFunctionTestMixin:
         dependencies = {"a": (), "b": {"a"}}
         graph = HighLevelGraph(layers, dependencies)
         assert self.get(graph, "z") == 4
+
+
+def import_or_none(name):
+    """Import a module and return it; in case of failure; return None"""
+    try:
+        return importlib.import_module(name)
+    except (ImportError, AttributeError):
+        return None