Function checkpoint

dask/graph_manipulation.py:32–70

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 split_every: int

(
    *collections,
    split_every: float | Literal[False] | None = None,
)

def checkpoint(
    *collections,
    split_every: float | Literal[False] | None = None,
) -> 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
    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_\\text{split_every}(\\text{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:
        split_every = 8
    elif split_every is not False:
        split_every = int(split_every)
        if split_every < 2:
            raise ValueError("split_every must be False, None, or >= 2")

    collections, _ = unpack_collections(*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: