Repository: dask/dask
Pull request #11347 (2024-08-26 11:59:56)

Full support for task spec in dask.order
This builds on https://github.com/dask/distributed/pull/8797

The relevant commit is https://github.com/dask/dask/commit/b2ee43efaf6f0f847cf4e9ec27fe222c7b10ff5e

It allows dask.order to fully benefit from the task_spec implementation and removes the need to recompute the dependencies of a given graph. This dependency computation is one of the driving runtime cost (particularly since it's been done multiple times).

It also removes a quite bad code path where we previously triggered a deepcopy which was very expensive (for very large graphs)


Sibling PR that enables this in distributed https://github.com/dask/distributed/pull/8842

This PR _should_ work without the distributed change but not the other way round


Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 695acdbd8..77c388735 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -82,8 +82,6 @@ from contextlib import contextmanager
 from functools import partial
 from typing import Any, TypeVar, cast, overload
 
-from dask.base import tokenize
-from dask.core import reverse_dict
 from dask.sizeof import sizeof
 from dask.typing import Key as KeyType
 from dask.utils import is_namedtuple_instance
@@ -526,6 +524,8 @@ class DataNode(GraphNode):
         return f"DataNode({self.key}, type={self.typ}, {self.value})"
 
     def __dask_tokenize__(self):
+        from dask.base import tokenize
+
         return (type(self).__name__, tokenize(self.value))
 
     def __reduce__(self) -> str | tuple[Any, ...]:
@@ -608,7 +608,7 @@ class Task(GraphNode):
         self.kwargs = parse_input(kwargs)
         dependencies: set = set()
         dependencies.update(_get_dependencies(self.args))
-        dependencies.update(_get_dependencies(self.kwargs.values()))
+        dependencies.update(_get_dependencies(tuple(self.kwargs.values())))
         if dependencies:
             self.dependencies = dependencies
         else:
@@ -809,50 +809,36 @@ class Task(GraphNode):
         return self._is_coro
 
 
-class DependenciesMapping(Mapping):
+class DependenciesMapping(MutableMapping):
     def __init__(self, dsk):
         self.dsk = dsk
-        self.blocklist = None
-        self.removed_keys = set()
+        self._removed = set()
 
     def __getitem__(self, key):
-        if key in self.removed_keys:
-            raise KeyError(key)
         v = self.dsk[key]
         if not isinstance(v, GraphNode):
             from dask.core import get_dependencies
 
-            return get_dependencies(self.dsk, task=self.dsk[key])
-        if self.blocklist and self.blocklist[key]:
-            return self.dsk[key].dependencies - self.blocklist[key]
-        return self.dsk[key].dependencies
+            deps = get_dependencies(self.dsk, task=self.dsk[key])
+        else:
+            deps = self.dsk[key].dependencies
+        if self._removed:
+            deps -= self._removed
+        return deps
 
     def __iter__(self):
         return iter(self.dsk)
 
-    def copy(self):
-        return DependenciesMapping(self.dsk)
+    def __delitem__(self, key: Any) -> None:
+        self._removed.add(key)
 
-    def __delitem__(self, key):
-        self.removed_keys.add(key)
-
-    def remove_dependency(self, key, dep):
-        if self.blocklist is None:
-            self.blocklist = defaultdict(set)
-        self.blocklist[key].add(dep)
+    def __setitem__(self, key: Any, value: Any) -> None:
+        raise NotImplementedError
 
     def __len__(self) -> int:
         return len(self.dsk)
 
 
-def get_deps(dsk):
-    # FIXME: I think we don't need this function
-    assert all(isinstance(v, GraphNode) for v in dsk.values())
-    dependencies = DependenciesMapping(dsk)
-    dependents = reverse_dict(dependencies)
-    return dependencies, dependents
-
-
 class _DevNullMapping(MutableMapping):
     def __getitem__(self, key):
         raise KeyError(key)
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index b1536240c..8299f4224 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -12,7 +12,7 @@ import numpy as np
 from dask import config
 from dask.array.chunk import getitem
 from dask.array.core import Array, unknown_chunk_message
-from dask.array.dispatch import concatenate_lookup
+from dask.array.dispatch import concatenate_lookup, take_lookup
 from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
 
@@ -294,8 +294,11 @@ def _getitem(obj, index):
 
 
 def concatenate_arrays(arrs, sorter, axis):
-    concatenate = concatenate_lookup.dispatch(type(arrs[0]))
-    return np.take(concatenate(arrs, axis=axis), np.argsort(sorter[1]), axis=axis)
+    return take_lookup(
+        concatenate_lookup.dispatch(type(arrs[0]))(arrs, axis=axis),
+        np.argsort(sorter[1]),
+        axis=axis,
+    )
 
 
 def convert_key(key, chunk, axis):
diff --git a/dask/array/backends.py b/dask/array/backends.py
index e033d90d8..0c114beef 100644
--- a/dask/array/backends.py
+++ b/dask/array/backends.py
@@ -14,6 +14,7 @@ from dask.array.dispatch import (
     nannumel_lookup,
     numel_lookup,
     percentile_lookup,
+    take_lookup,
     tensordot_lookup,
     to_cupy_dispatch,
     to_numpy_dispatch,
@@ -24,6 +25,7 @@ from dask.array.percentile import _percentile
 from dask.backends import CreationDispatch, DaskBackendEntrypoint
 
 concatenate_lookup.register((object, np.ndarray), np.concatenate)
+take_lookup.register((object, np.ndarray, np.ma.masked_array), np.take)
 tensordot_lookup.register((object, np.ndarray), np.tensordot)
 einsum_lookup.register((object, np.ndarray), np.einsum)
 empty_lookup.register((object, np.ndarray), np.empty)
@@ -123,6 +125,7 @@ def _tensordot(a, b, axes=2):
 
 @tensordot_lookup.register_lazy("cupy")
 @concatenate_lookup.register_lazy("cupy")
+@take_lookup.register_lazy("cupy")
 @nannumel_lookup.register_lazy("cupy")
 @numel_lookup.register_lazy("cupy")
 @to_numpy_dispatch.register_lazy("cupy")
@@ -130,6 +133,7 @@ def register_cupy():
     import cupy
 
     concatenate_lookup.register(cupy.ndarray, cupy.concatenate)
+    take_lookup.register(cupy.ndarray, cupy.take)
     tensordot_lookup.register(cupy.ndarray, cupy.tensordot)
     percentile_lookup.register(cupy.ndarray, percentile)
     numel_lookup.register(cupy.ndarray, _numel_arraylike)
@@ -151,6 +155,13 @@ def register_cupy():
         return cupy.einsum(*args, **kwargs)
 
 
+def sparse_take(array, idx, axis=0):
+    if axis not in {0, 1}:  # pragma: no cover
+        raise ValueError("Sparse matrices can only be concatenated along axis 0 or 1")
+    indexer = (slice(None), idx) if axis else (idx, slice(None))
+    return array[indexer]
+
+
 @tensordot_lookup.register_lazy("cupyx")
 @concatenate_lookup.register_lazy("cupyx")
 def register_cupyx():
@@ -177,12 +188,14 @@ def register_cupyx():
 
     concatenate_lookup.register(spmatrix, _concat_cupy_sparse)
     tensordot_lookup.register(spmatrix, _tensordot_scipy_sparse)
+    take_lookup.register(spmatrix, sparse_take)
 
 
 @tensordot_lookup.register_lazy("sparse")
 @concatenate_lookup.register_lazy("sparse")
 @nannumel_lookup.register_lazy("sparse")
 @numel_lookup.register_lazy("sparse")
+@take_lookup.register_lazy("sparse")
 def register_sparse():
     import sparse
 
@@ -193,10 +206,12 @@ def register_sparse():
     # https://github.com/dask/dask/issues/7169
     numel_lookup.register(sparse.COO, _numel_ndarray)
     nannumel_lookup.register(sparse.COO, _nannumel_sparse)
+    take_lookup.register(sparse.COO, np.take)
 
 
 @tensordot_lookup.register_lazy("scipy")
 @concatenate_lookup.register_lazy("scipy")
+@take_lookup.register_lazy("scipy")
 def register_scipy_sparse():
     import scipy.sparse
 
@@ -214,6 +229,7 @@ def register_scipy_sparse():
 
     concatenate_lookup.register(scipy.sparse.spmatrix, _concatenate)
     tensordot_lookup.register(scipy.sparse.spmatrix, _tensordot_scipy_sparse)
+    take_lookup.register(scipy.sparse.spmatrix, sparse_take)
 
 
 def _tensordot_scipy_sparse(a, b, axes):
diff --git a/dask/array/core.py b/dask/array/core.py
index 7d4f023fa..faedbb6ef 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -4795,7 +4795,7 @@ def broadcast_shapes(*shapes):
         if np.isnan(sizes).any():
             dim = np.nan
         else:
-            dim = 0 if 0 in sizes else np.max(sizes)
+            dim = 0 if 0 in sizes else np.max(sizes).item()
         if any(i not in [-1, 0, 1, dim] and not np.isnan(i) for i in sizes):
             raise ValueError(
                 "operands could not be broadcast together with "
diff --git a/dask/array/dispatch.py b/dask/array/dispatch.py
index 44328ad67..7fa99b673 100644
--- a/dask/array/dispatch.py
+++ b/dask/array/dispatch.py
@@ -9,6 +9,7 @@ from __future__ import annotations
 from dask.utils import Dispatch
 
 concatenate_lookup = Dispatch("concatenate")
+take_lookup = Dispatch("take")
 tensordot_lookup = Dispatch("tensordot")
 einsum_lookup = Dispatch("einsum")
 empty_lookup = Dispatch("empty")
diff --git a/dask/core.py b/dask/core.py
index b2075ba7e..eb126cea4 100644
--- a/dask/core.py
+++ b/dask/core.py
@@ -4,6 +4,7 @@ from collections import defaultdict
 from collections.abc import Collection, Iterable, Mapping
 from typing import Any, Literal, TypeVar, cast, overload
 
+from dask._task_spec import DependenciesMapping
 from dask.typing import Graph, Key, NoDefault, no_default
 
 
@@ -418,6 +419,7 @@ def subs(task, key, val):
 
 
 def _toposort(dsk, keys=None, returncycle=False, dependencies=None):
+
     # Stack-based depth-first search traversal.  This is based on Tarjan's
     # method for topological sorting (see wikipedia for pseudocode)
     if keys is None:
@@ -439,7 +441,8 @@ def _toposort(dsk, keys=None, returncycle=False, dependencies=None):
     seen = set()
 
     if dependencies is None:
-        dependencies = {k: get_dependencies(dsk, k) for k in dsk}
+
+        dependencies = DependenciesMapping(dsk)
 
     for key in keys:
         if key in completed:
diff --git a/dask/order.py b/dask/order.py
index a261f4121..d80220338 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -46,12 +46,12 @@ graph types laid out very carefully to show the kinds of situations that often
 arise, and the order we would like to be determined.
 
 """
-import copy
 from collections import defaultdict, deque, namedtuple
 from collections.abc import Callable, Iterable, Mapping, MutableMapping
 from typing import Any, Literal, NamedTuple, overload
 
-from dask.core import get_dependencies, get_deps, getcycle, istask, reverse_dict
+from dask._task_spec import DataNode, DependenciesMapping
+from dask.core import get_deps, getcycle, istask, reverse_dict
 from dask.typing import Key
 
 
@@ -105,20 +105,21 @@ def order(
         return {}
 
     dsk = dict(dsk)
-    expected_len = len(dsk)
-
-    if dependencies is None:
-        dependencies_are_copy = True
-        dependencies = {k: get_dependencies(dsk, k) for k in dsk}
-    else:
-        # Below we're removing items from the sets in this dict
-        # We need a deepcopy for that but we only want to do this if necessary
-        # since this only happens for special cases.
-        dependencies_are_copy = False
-        dependencies = dict(dependencies)
 
+    dependencies = DependenciesMapping(dsk)
     dependents = reverse_dict(dependencies)
 
+    external_keys = set()
+    if len(dependents) != len(dependencies):
+        # There are external keys. Let's add a DataNode to ensure the algo below
+        # is encountering a complete graph. Those artificial data nodes will be
+        # ignored when assigning results
+        for k in dependents:
+            if k not in dependencies:
+                external_keys.add(k)
+                dsk[k] = DataNode(k, object())
+
+    expected_len = len(dsk)
     leaf_nodes = {k for k, v in dependents.items() if not v}
     root_nodes = {k for k, v in dependencies.items() if not v}
 
@@ -136,6 +137,7 @@ def order(
     all_tasks = False
     n_removed_leaves = 0
     requires_data_task = defaultdict(set)
+
     while not all_tasks:
         all_tasks = True
         for leaf in list(leaf_nodes):
@@ -156,29 +158,27 @@ def order(
                     result[leaf] = prio
                 n_removed_leaves += 1
                 leaf_nodes.remove(leaf)
-                del dsk[leaf]
-                del dependents[leaf]
                 for dep in dependencies[leaf]:
                     dependents[dep].remove(leaf)
                     if not dependents[dep]:
                         leaf_nodes.add(dep)
+                del dsk[leaf]
+                del dependencies[leaf]
+                del dependents[leaf]
 
         for root in list(root_nodes):
             if root in leaf_nodes:
                 continue
-            if not istask(dsk[root]) and len(dependents[root]) > 1:
-                if not dependencies_are_copy:
-                    dependencies_are_copy = True
-                    dependencies = copy.deepcopy(dependencies)
+            deps_root = dependents[root]
+            if not istask(dsk[root]) and len(deps_root) > 1:
+                del dsk[root]
+                del dependencies[root]
                 root_nodes.remove(root)
-                for dep in dependents[root]:
+                del dependents[root]
+                for dep in deps_root:
                     requires_data_task[dep].add(root)
-                    dependencies[dep].remove(root)
                     if not dependencies[dep]:
                         root_nodes.add(dep)
-                del dsk[root]
-                del dependencies[root]
-                del dependents[root]
 
     num_needed, total_dependencies = ndependencies(dependencies, dependents)
     if len(total_dependencies) != len(dsk):
@@ -242,16 +242,14 @@ def order(
 
             while requires_data_task[item]:
                 add_to_result(requires_data_task[item].pop())
-
             if return_stats:
                 result[item] = Order(i, crit_path_counter - _crit_path_counter_offset)
             else:
                 result[item] = i
-
+            if item not in external_keys:
+                i += 1
             if item in root_nodes:
                 processed_roots.add(item)
-
-            i += 1
             # Note: This is a `set` and therefore this introduces a certain
             # randomness. However, this randomness should not have any impact on
             # the final result since the `process_runnable` should produce
@@ -604,6 +602,8 @@ def order(
         process_runnables()
 
     assert len(result) == expected_len
+    for k in external_keys:
+        del result[k]
     return result  # type: ignore
 
 
diff --git a/dask/tokenize.py b/dask/tokenize.py
index 0c733db5a..72ba96b0b 100644
--- a/dask/tokenize.py
+++ b/dask/tokenize.py
@@ -13,6 +13,7 @@ import types
 import uuid
 from collections import OrderedDict
 from collections.abc import Iterable
+from contextvars import ContextVar
 from functools import partial
 
 import cloudpickle
@@ -40,7 +41,7 @@ def _tokenize(*args: object, **kwargs: object) -> str:
 
 tokenize_lock = threading.RLock()
 _SEEN: dict[int, tuple[int, object]] = {}
-_ENSURE_DETERMINISTIC = None
+_ENSURE_DETERMINISTIC: ContextVar[bool | None] = ContextVar("_ENSURE_DETERMINISTIC")
 
 
 def tokenize(
@@ -66,20 +67,25 @@ def tokenize(
     global _SEEN, _ENSURE_DETERMINISTIC
     with tokenize_lock:
         seen_before, _SEEN = _SEEN, {}
-        _ENSURE_DETERMINISTIC = ensure_deterministic
+        token = None
+        try:
+            _ENSURE_DETERMINISTIC.get()
+        except LookupError:
+            token = _ENSURE_DETERMINISTIC.set(ensure_deterministic)
         try:
             return _tokenize(*args, **kwargs)
         finally:
+            if token:
+                _ENSURE_DETERMINISTIC.reset(token)
             _SEEN = seen_before
-            _ENSURE_DETERMINISTIC = None
 
 
 def _maybe_raise_nondeterministic(msg: str) -> None:
-    if (
-        _ENSURE_DETERMINISTIC
-        or _ENSURE_DETERMINISTIC is None
-        and config.get("tokenize.ensure-deterministic")
-    ):
+    try:
+        val = _ENSURE_DETERMINISTIC.get()
+    except LookupError:
+        val = None
+    if val or val is None and config.get("tokenize.ensure-deterministic"):
         raise TokenizationError(msg)
 
 
diff --git a/dask/utils.py b/dask/utils.py
index e911047aa..fc9322947 100644
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -15,7 +15,7 @@ from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Set
 from contextlib import contextmanager, nullcontext, suppress
 from datetime import datetime, timedelta
 from errno import ENOENT
-from functools import lru_cache, wraps
+from functools import wraps
 from importlib import import_module
 from numbers import Integral, Number
 from operator import add
@@ -26,7 +26,6 @@ from weakref import WeakValueDictionary
 import tlz as toolz
 
 from dask import config
-from dask.core import get_deps
 from dask.typing import no_default
 
 K = TypeVar("K")
@@ -1161,20 +1160,6 @@ def insert(tup, loc, val):
     return tuple(L)
 
 
-def dependency_depth(dsk):
-    deps, _ = get_deps(dsk)
-
-    @lru_cache(maxsize=None)
-    def max_depth_by_deps(key):
-        if not deps[key]:
-            return 1
-
-        d = 1 + max(max_depth_by_deps(dep_key) for dep_key in deps[key])
-        return d
-
-    return max(max_depth_by_deps(dep_key) for dep_key in deps.keys())
-
-
 def memory_repr(num):
     for x in ["bytes", "KB", "MB", "GB", "TB"]:
         if num < 1024.0:
diff --git a/docs/release-procedure.md b/docs/release-procedure.md
index ec2531151..bece1acce 100644
--- a/docs/release-procedure.md
+++ b/docs/release-procedure.md
@@ -31,7 +31,7 @@ Releasing dask and distributed:
 
 *   Commit
 
-        git commit -a -m "bump version to YYYY.M.X"
+        git commit -a -m "Version YYYY.M.X"
 
 *   Tag commit
 
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 4011c6874..c73856fcc 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -1,6 +1,39 @@
 Changelog
 =========
 
+.. note::
+
+    This is not exhaustive. For an exhaustive list of changes, see the git log.
+
+.. _v2024.10.0:
+
+2024.10.0
+---------
+
+Notable Changes
+^^^^^^^^^^^^^^^
+
+- Zarr-Python 3 compatibility  (:pr:`11388`)
+- Avoid exponentially increasing taskgraph in overlap (:pr:`11423`)
+- Ensure numba tokenization does not use slow pickle path (:pr:`11419`)
+
+.. dropdown:: Additional changes
+
+  - Ensure `broadcast_shapes()` returns integers, not NumPy scalars. (:pr:`11434`) `Martin Yeo`_
+  - (fix): sparse indexing (:pr:`11430`) `Ilan Gold`_
+  - Ensure that recursively calling tokenize respects ensure_deterministic (:pr:`11431`) `Florian Jetter`_
+  - Make P2P more configurable (:pr-distributed:`8469`) `Hendrik Makait`_
+  - Fit Dashboard worker table to page width (:pr-distributed:`8897`) `Jacob Tomlinson`_
+  - Raise helpful error when using the wrong plugin base classes (:pr-distributed:`8893`) `Jacob Tomlinson`_
+  - Fix url escaping on exceptions dashboard for non-string keys (:pr-distributed:`8891`) `Patrick Hoefler`_
+  - Add meaningful error for out of disk exception during write (:pr-distributed:`8886`) `Hendrik Makait`_
+  - Fix binary operations with scalar on the left (:pr-expr:`1150`) `Patrick Hoefler`_
+  - Raise exception when calculating divisons (:pr-expr:`1149`) `Patrick Hoefler`_
+  - Fix merge_asof for single partition (:pr-expr:`1145`) `Patrick Hoefler`_
+  - Improve handling of optional dependencies in `analyze` and `explain` (:pr-expr:`1146`) `Hendrik Makait`_
+  - Fix alignment issue with groupby index accessors (:pr-expr:`1142`) `Patrick Hoefler`_
+  - Fix displaying timestamp scalar (:pr-expr:`1141`) `Patrick Hoefler`_
+
 .. _v2024.9.1:
 
 2024.9.1
@@ -14,7 +47,7 @@ Improved adaptive scaling resilience
 Adaptive scaling clusters now recover from spurious errors during scaling.
 
 See :pr-distributed:`8871` by `Hendrik Makait`_ for more details.
-  
+
 .. dropdown:: Additional changes
 
   - Improve error message for incorrect columns order in meta information (:pr:`11393`) `Dmitry Balabka`_
@@ -95,7 +128,7 @@ To disable automatic selection, users can select their preferred method via the
 
     import dask.config
     # Choose either "tasks" or "p2p"
-    dask.config.set({"array.rechunk.method": "tasks"}) 
+    dask.config.set({"array.rechunk.method": "tasks"})
 
 or when rechunking
 
@@ -103,10 +136,10 @@ or when rechunking
 
     import dask.array as da
     arr = da.random.random(size=(1000, 1000, 365), chunks=(-1, -1, "auto"))
-    # Choose either "tasks" or "p2p" 
+    # Choose either "tasks" or "p2p"
     arr = arr.rechunk(("auto", "auto", -1), method="tasks")
 
-See :pr:`11337` by `Hendrik Makait`_ for more details. 
+See :pr:`11337` by `Hendrik Makait`_ for more details.
 
 New shuffle API for Dask Arrays
 """""""""""""""""""""""""""""""
@@ -436,7 +469,7 @@ More resilient distributed lock
 :py:class:`distributed.Lock` is now resilient to worker failures.
 Previously deadlocks were possible in cases where a lock-holding worker
 was lost and/or failed to release the lock due to an error.
-  
+
 See :pr-distributed:`8770` by `Florian Jetter`_ for more details.
 
 .. dropdown:: Additional changes
@@ -486,7 +519,7 @@ is now the required minimum version to run Dask DataFrame.
 
 The mimimum version of ``partd`` was also raised to 1.4.0. Versions before 1.4
 are not compatible with ``pandas`` 2.
-  
+
 See :pr:`11199` by `Patrick Hoefler`_ for more details.
 
 Publish-subscribe APIs deprecated
@@ -8845,3 +8878,5 @@ Other
 .. _`David Stansby`: https://github.com/dstansby
 .. _`Mario Linker`: https://github.com/maldag
 .. _`Dmitry Balabka`: https://github.com/dbalabka
+.. _`Martin Yeo`: https://github.com/trexfeathers
+.. _`Ilan Gold`: https://github.com/ilan-gold
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 8e3874e68..1badbb40c 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -317,9 +317,10 @@ epub_copyright = "2014-2018, Anaconda, Inc. and contributors"
 # epub_tocdup = True
 
 extlinks = {
-    "issue": ("https://github.com/dask/dask/issues/%s", "GH#"),
-    "pr": ("https://github.com/dask/dask/pull/%s", "GH#"),
-    "pr-distributed": ("https://github.com/dask/distributed/pull/%s", "GH#"),
+    "issue": ("https://github.com/dask/dask/issues/%s", "dask#"),
+    "pr": ("https://github.com/dask/dask/pull/%s", "dask#"),
+    "pr-distributed": ("https://github.com/dask/distributed/pull/%s", "distributed#"),
+    "pr-expr": ("https://github.com/dask/dask-expr/pull/%s", "dask-expr#"),
 }
 
 #  --Options for sphinx extensions -----------------------------------------------
diff --git a/docs/source/how-to/selecting-the-collection-backend.rst b/docs/source/how-to/selecting-the-collection-backend.rst
index be9eed924..4674a389d 100644
--- a/docs/source/how-to/selecting-the-collection-backend.rst
+++ b/docs/source/how-to/selecting-the-collection-backend.rst
@@ -128,6 +128,7 @@ Dask-Array compute-based dispatch functions (as defined in ``dask.array.dispatch
    - numel_lookup
    - percentile_lookup
    - tensordot_lookup
+   - take_lookup
 
 Dask-Dataframe compute-based dispatch functions (as defined in ``dask.dataframe.dispatch``, and defined for Pandas in ``dask.dataframe.backends``):
 
diff --git a/pyproject.toml b/pyproject.toml
index d9f77a2d7..73d761cc8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -54,7 +54,7 @@ dataframe = [
     "pandas >= 2.0",
     "dask-expr >= 1.1, <1.2",  # dask-expr pins the dask version
 ]
-distributed = ["distributed == 2024.9.1"]
+distributed = ["distributed == 2024.10.0"]
 diagnostics = [
     "bokeh >= 3.1.0",
     "jinja2 >= 2.10.3",


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 8d56ead01..121f8a698 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -10,6 +10,7 @@ import pytest
 
 np = pytest.importorskip("numpy")
 
+import itertools
 import math
 import operator
 import os
@@ -871,6 +872,10 @@ def test_broadcast_shapes():
     assert (3, 4) == broadcast_shapes((3, 1), (1, 4), (4,))
     assert (5, 6, 7, 3, 4) == broadcast_shapes((3, 1), (), (5, 6, 7, 1, 4))
 
+    assert all(
+        isinstance(i, int) for i in itertools.chain(*broadcast_shapes([(1, 2), (3, 1)]))
+    )
+
     pytest.raises(ValueError, lambda: broadcast_shapes((3,), (3, 4)))
     pytest.raises(ValueError, lambda: broadcast_shapes((2, 3), (2, 3, 1)))
     pytest.raises(ValueError, lambda: broadcast_shapes((2, 3), (1, np.nan)))
@@ -4905,6 +4910,30 @@ def test_dask_array_holds_scipy_sparse_containers():
     assert (zz == xx.T).all()
 
 
+@pytest.mark.parametrize(
+    "index",
+    [
+        [5, 8],
+        0,
+        slice(5, 8),
+        np.array([5, 8]),
+        np.array([True, False] * 500),
+        [True, False] * 500,
+    ],
+)
+@pytest.mark.parametrize("sparse_module_path", ["scipy.sparse", "cupyx.scipy.sparse"])
+def test_scipy_sparse_indexing(index, sparse_module_path):
+    sp = pytest.importorskip(sparse_module_path)
+    x = da.random.default_rng().random((1000, 10), chunks=(100, 10))
+    x[x < 0.9] = 0
+    y = x.map_blocks(sp.csr_matrix)
+
+    assert not (
+        x[index, :].compute(scheduler="single-threaded")
+        != y[index, :].compute(scheduler="single-threaded")
+    ).sum()
+
+
 @pytest.mark.parametrize("axis", [0, 1])
 def test_scipy_sparse_concatenate(axis):
     pytest.importorskip("scipy.sparse")
diff --git a/dask/dataframe/io/tests/test_hdf.py b/dask/dataframe/io/tests/test_hdf.py
index 929b5d82b..07cbacc14 100644
--- a/dask/dataframe/io/tests/test_hdf.py
+++ b/dask/dataframe/io/tests/test_hdf.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import os
 import pathlib
+from functools import lru_cache
 from time import sleep
 
 import numpy as np
@@ -12,11 +13,12 @@ from packaging.version import Version
 import dask
 import dask.dataframe as dd
 from dask._compatibility import PY_VERSION
+from dask.core import get_deps
 from dask.dataframe._compat import tm
 from dask.dataframe.optimize import optimize_dataframe_getitem
 from dask.dataframe.utils import assert_eq
 from dask.layers import DataFrameIOLayer
-from dask.utils import dependency_depth, tmpdir, tmpfile
+from dask.utils import tmpdir, tmpfile
 
 # there's no support in upstream for writing HDF with extension dtypes yet.
 # see https://github.com/pandas-dev/pandas/issues/31199
@@ -333,6 +335,20 @@ def test_to_hdf_modes_multiple_files():
         assert_eq(dd.concat([df, df]), out)
 
 
+def dependency_depth(dsk):
+    deps, _ = get_deps(dsk)
+
+    @lru_cache(maxsize=None)
+    def max_depth_by_deps(key):
+        if not deps[key]:
+            return 1
+
+        d = 1 + max(max_depth_by_deps(dep_key) for dep_key in deps[key])
+        return d
+
+    return max(max_depth_by_deps(dep_key) for dep_key in deps.keys())
+
+
 def test_to_hdf_link_optimizations():
     """testing dask link levels is correct by calculating the depth of the dask graph"""
     pytest.importorskip("tables")
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 694749818..6f04138fb 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -124,7 +124,8 @@ def test_avoid_broker_nodes(abcde):
     assert o[(a, 1)] < o[(b, 0)] or (o[(b, 1)] < o[(a, 0)] and o[(b, 2)] < o[(a, 0)])
 
 
-def test_base_of_reduce_preferred(abcde):
+@pytest.mark.parametrize("data_root", [True, False])
+def test_base_of_reduce_preferred(abcde, data_root):
     r"""
                a3
               /|
@@ -144,7 +145,10 @@ def test_base_of_reduce_preferred(abcde):
     dsk = {(a, i): (f, (a, i - 1), (b, i)) for i in [1, 2, 3]}
     dsk[(a, 0)] = (f, (b, 0))
     dsk.update({(b, i): (f, c, 1) for i in [0, 1, 2, 3]})
-    dsk[c] = 1
+    if data_root:
+        dsk[c] = 1
+    else:
+        dsk[c] = (f, 1)
 
     o = order(dsk)
     assert_topological_sort(dsk, o)
@@ -243,7 +247,13 @@ def test_deep_bases_win_over_dependents(abcde):
        e    d
     """
     a, b, c, d, e = abcde
-    dsk = {a: (f, b, c, d), b: (f, d, e), c: (f, d), d: 1, e: 2}
+    dsk = {
+        a: (f, b, c, d),
+        b: (f, d, e),
+        c: (f, d),
+        d: (f, 1),
+        e: (f, 2),
+    }
 
     o = order(dsk)
     assert_topological_sort(dsk, o)
@@ -2651,3 +2661,14 @@ def test_stackstac():
         # Ensure that we're not processing the entire graph using
         # process_runnables (fractional values) but are using the critical path
         assert o[k].critical_path in {1, 2}
+
+
+def test_handle_out_of_graph_dependencies():
+    from dask._task_spec import Task
+
+    ta = Task("a", f, ())
+    tb = Task("b", f, (ta.ref(),))
+    tc = Task("c", f, (tb.ref(),))
+    dsk = {t.key: t for t in [tb, tc]}
+    o = order(dsk)
+    assert len(o) == 2
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index efa2c5265..da361701e 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -13,6 +13,7 @@ from dask._task_spec import (
     DependenciesMapping,
     Task,
     TaskRef,
+    _get_dependencies,
     convert_legacy_graph,
     convert_legacy_task,
     execute_graph,
@@ -341,6 +342,38 @@ def test_runnable_as_kwarg():
     assert t() == "ab3"
 
 
+def test_dependency_as_kwarg():
+    def func_kwarg(a, b, c=""):
+        return a + b + str(c)
+
+    t1 = Task("key-1", sum, [1, 2])
+    t2 = Task(
+        "key-2",
+        func_kwarg,
+        "a",
+        "b",
+        c=t1.ref(),
+    )
+    with pytest.raises(RuntimeError, match="missing"):
+        t2()
+    # It isn't sufficient to raise. We also rely on the attribute being set
+    # properly since distribute will use this to infer actual dependencies The
+    # exception may be raised recursively
+    assert t2.dependencies
+    assert t2({"key-1": t1()}) == "ab3"
+
+
+def test_array_as_argument():
+    np = pytest.importorskip("numpy")
+    t = Task("key-1", func, np.array([1, 2]), "b")
+    assert t() == "[1 2]-b"
+
+    # This will **not** work since we do not want to recurse into an array!
+    t2 = Task("key-2", func, np.array([1, t.ref()]), "b")
+    assert t2({"key-1": "foo"}) != "[1 foo]-b"
+    assert not _get_dependencies(np.array([1, t.ref()]))
+
+
 def test_subgraph_callable():
     def add(a, b):
         return a + b
@@ -724,3 +757,8 @@ def test_execute_tasks_in_graph():
     res = execute_graph(dsk)
     assert len(res) == 1
     assert res["key-4"] == "foo-bar-a-b=c"
+
+
+def test_deterministic_tokenization_respected():
+    with pytest.raises(RuntimeError, match="deterministic"):
+        tokenize(Task("key", func, object()), ensure_deterministic=True)
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index a3f447027..d361b6bc8 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -38,10 +38,12 @@ def check_clean_state():
     from dask.tokenize import _ENSURE_DETERMINISTIC, _SEEN
 
     assert not _SEEN
-    assert _ENSURE_DETERMINISTIC is None
+    with pytest.raises(LookupError):
+        _ENSURE_DETERMINISTIC.get()
     yield
     assert not _SEEN
-    assert _ENSURE_DETERMINISTIC is None
+    with pytest.raises(LookupError):
+        _ENSURE_DETERMINISTIC.get()
 
 
 def check_tokenize(*args, **kwargs):
@@ -1423,3 +1425,12 @@ def test_tokenize_pandas_arrow_strings():
     tokens = normalize_token(ser)
     # Maybe a little brittle but will do for now
     assert any(str(tok) == "string" for tok in flatten(tokens))
+
+
+def test_tokenize_recursive_respects_ensure_deterministic():
+    class Foo:
+        def __dask_tokenize__(self):
+            return tokenize(object())
+
+    with pytest.raises(RuntimeError):
+        tokenize(Foo(), ensure_deterministic=True)