Repository: dask/dask
Pull request #11736 (2025-02-11 16:33:37)

Wrap HLGs in an Expr to avoid Client side materialization
This is a prototype to wrap HLGs in an expression. This helps with streamlining our `compute` calls and the serialization path to the scheduler.

I only tested this on a single test so far using a simple threaded scheduler. There are likely still tons of issues and the new expression classes have to be cleaned up, etc.

The top-level `compute` is implementing the new path and includes some comments on how this will look like when migrated to the client

sibling https://github.com/dask/distributed/pull/9008

Related PRs
- https://github.com/dask/dask/pull/11790
- https://github.com/dask/dask/pull/11791
- https://github.com/dask/dask/pull/11793
- https://github.com/dask/dask/pull/11835
- https://github.com/dask/dask/pull/11836
- https://github.com/dask/dask/pull/11837
- https://github.com/dask/dask/pull/11840


Closes https://github.com/dask/distributed/issues/7964
Closes https://github.com/dask/dask-expr/issues/14


fjetter: I believe all the test cases are now working fine except of some distribtued integration tests. CI seems to not install the right version of the distributed branch since they pass locally. 
There are many smaller changes that make sense unrelated to the HLG pickle change and I will start breaking those out into separate PRs to put this one in a reviewable state.
fjetter: The remaining test failures are a distributed issue. The simple APIs like Client.submit, map, etc. are still submitting low level graphs which the scheduler internally cannot handle on the other PR. That's easy to fix and something I intentionally postponed until the end.

This is now ready for review. I will be out for a couple of weeks and will only be able to address review comments then. 
fjetter: Somehow I broke dask.order... fun!
fjetter: While CI is green, there are a bunch of things that are only tested in distributed so this isn't done, yet
fjetter: On distributed side, there is one (maybe two) issues left that relate to how futures are serialized and registered. I doubt this will have a major impact on the code here. 
I'll give it another pass to see if there are smaller changes that are worth breaking out but review can start already.

fjetter: I summarized the high level changes in the changelog already. I recommend to start there to get an idea of the changes
fjetter: Also important in case somebody wants to test: To get a meaningful kick out of this when working with arrays we will first have to rework da.store since that is still explicitly optimizing and materializing (the only place in the code base doing this differently). While previous attempts to migrate this to blockwise failed, wrapping it up in these classes is simpler and possible. I already have an early prototype and will look into that once the final test cases are wrapped up.

(i.e. coiled benchmarks are using to_zarr / store in almost all cases so to measure the difference here the benchmarks have to be rewritten a bit. I could already confirm there that it is indeed not materializing locally which saves transmission time and removes any relevant constraints on memory usage on client side we had on the very large examples)

Patch:
diff --git a/dask/_expr.py b/dask/_expr.py
index 640f4b679..f79da9424 100644
--- a/dask/_expr.py
+++ b/dask/_expr.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import functools
 import os
+import warnings
 import weakref
 from collections import defaultdict
 from collections.abc import Generator
@@ -11,9 +12,10 @@ import toolz
 
 import dask
 from dask._task_spec import Task
+from dask.core import reverse_dict
 from dask.tokenize import _tokenize_deterministic
 from dask.typing import Key
-from dask.utils import funcname, import_required
+from dask.utils import ensure_dict, funcname, import_required
 
 if TYPE_CHECKING:
     # TODO import from typing (requires Python >=3.10)
@@ -73,8 +75,13 @@ class Expr:
     def _tune_up(self, parent):
         return None
 
+    def finalize_compute(self):
+        return self
+
     def _operands_for_repr(self):
-        raise NotImplementedError("Subclasses should implement this method")
+        return [
+            f"{param}={repr(op)}" for param, op in zip(self._parameters, self.operands)
+        ]
 
     def __str__(self):
         s = ", ".join(self._operands_for_repr())
@@ -99,7 +106,7 @@ class Expr:
         return header
 
     def _tree_repr_lines(self, indent=0, recursive=True):
-        raise NotImplementedError("Subclasses should implement this method")
+        return " " * indent + repr(self)
 
     def tree_repr(self):
         return os.linesep.join(self._tree_repr_lines())
@@ -140,7 +147,7 @@ class Expr:
         if dask.config.get("dask-expr-no-serialize", False):
             raise RuntimeError(f"Serializing a {type(self)} object")
         return Expr._reconstruct, tuple(
-            [type(self)] + self.operands + [self.deterministic_token]
+            [type(self), *self.operands, self.deterministic_token]
         )
 
     def _depth(self, cache=None):
@@ -498,6 +505,9 @@ class Expr:
     def _meta(self):
         raise NotImplementedError()
 
+    def __dask_annotations__(self):
+        return {}
+
     def __dask_graph__(self):
         """Traverse expression tree, collect layers"""
         stack = [self]
@@ -862,3 +872,314 @@ def optimize_until(expr: Expr, stage: OptimizerStage) -> Expr:
         return expr
 
     raise ValueError(f"Stage {stage!r} not supported.")
+
+
+class LLGExpr(Expr):
+    """Low Level Graph Expression"""
+
+    _parameters = ["dsk"]
+
+    def __dask_keys__(self):
+        return list(self.operand("dsk"))
+
+    def __dask_tokenize__(self):
+        return str(id(self))
+
+    def _layer(self) -> dict:
+        return ensure_dict(self.operand("dsk"))
+
+
+class HLGExpr(Expr):
+    _parameters = [
+        "dsk",
+        "low_level_optimizer",
+        "output_keys",
+        "postcompute",
+        "_cached_optimized",
+    ]
+    _defaults = {
+        "low_level_optimizer": None,
+        "output_keys": None,
+        "postcompute": None,
+        "_cached_optimized": None,
+    }
+
+    @staticmethod
+    def from_collection(collection, optimize_graph=True):
+        from dask.highlevelgraph import HighLevelGraph
+
+        if hasattr(collection, "dask"):
+            dsk = collection.dask.copy()
+        else:
+            dsk = collection.__dask_graph__()
+
+        # Delayed objects still ship with low level graphs as `dask` when going
+        # through optimize / persist
+        if not isinstance(dsk, HighLevelGraph):
+
+            dsk = HighLevelGraph.from_collections(
+                str(id(collection)), dsk, dependencies=()
+            )
+        if optimize_graph and not hasattr(collection, "__dask_optimize__"):
+            warnings.warn(
+                f"Collection {type(collection)} does not define a "
+                "`__dask_optimize__` method. In the future this will raise. "
+                "If no optimization is desired, please set this to `None`.",
+                PendingDeprecationWarning,
+            )
+            low_level_optimizer = None
+        else:
+            low_level_optimizer = (
+                collection.__dask_optimize__ if optimize_graph else None
+            )
+        return HLGExpr(
+            dsk=dsk,
+            low_level_optimizer=low_level_optimizer,
+            output_keys=collection.__dask_keys__(),
+            postcompute=collection.__dask_postcompute__,
+        )
+
+    def finalize_compute(self):
+        return HLGFinalizeCompute(self)
+
+    def __dask_annotations__(self) -> dict[str, dict[Key, object]]:
+        # optimization has to be called (and cached) since blockwise fusion can
+        # alter annotations
+        # see `dask.blockwise.(_fuse_annotations|_can_fuse_annotations)`
+        dsk = self._optimized_dsk()
+        if isinstance(dsk, dict):
+            dsk = self.dsk
+        annotations_by_type: defaultdict[str, dict[Key, object]] = defaultdict(dict)
+        for layer in dsk.layers.values():
+            if layer.annotations:
+                annot = layer.annotations
+                for annot_type, value in annot.items():
+                    annotations_by_type[annot_type].update(
+                        {k: (value(k) if callable(value) else value) for k in layer}
+                    )
+        return dict(annotations_by_type)
+
+    def __dask_keys__(self):
+        if keys := self.operand("output_keys"):
+            return keys
+        dsk = self.operand("dsk")
+        # Note: This will materialize
+        dependencies = dsk.get_all_dependencies()
+        dependents = reverse_dict(dependencies)
+        keys = [d for d in dependents if not dependents[d] and d in dsk]
+        self.output_keys = keys
+        return keys
+
+    def __dask_tokenize__(self):
+        # There is currently not way to hash a HighLevelGraph fast and reliably.
+        # It is important for dask-expr for this to not be duplicated so we'll
+        # just use the ID.
+        return str(id(self))
+
+    def _optimized_dsk(self):
+        if self._cached_optimized:
+            return self._cached_optimized
+        keys = self.output_keys
+        optimizer = self.low_level_optimizer
+        if keys is None and optimizer is not None:
+            keys = self.__dask_keys__()
+        dsk = self.dsk
+        if (optimizer := self.low_level_optimizer) is not None:
+            dsk = optimizer(dsk, keys)
+        self._cached_optimized = dsk
+        return dsk
+
+    def _layer(self) -> dict:
+        dsk = self._optimized_dsk()
+        return ensure_dict(dsk)
+
+
+class _HLGExprSequence(Expr):
+
+    def __getitem__(self, other):
+        return self.operands[other]
+
+    def _operands_for_repr(self):
+        return [
+            f"name={self.operand('name')!r}",
+            f"dsk={self.operand('dsk')!r}",
+        ]
+
+    def _tree_repr_lines(self, indent=0, recursive=True):
+        return self._operands_for_repr()
+
+    def finalize_compute(self):
+        return HLGFinalizeCompute(self)
+
+    def __dask_graph__(self):
+        # This class has to override this and not just _layer to ensure the HLGs
+        # are not optimized individually
+        from dask.highlevelgraph import HighLevelGraph
+
+        groups = toolz.groupby(
+            lambda x: x.low_level_optimizer if isinstance(x, HLGExpr) else None,
+            self.operands,
+        )
+        outer_graphs = []
+        for optimizer, group in groups.items():
+            graphs = []
+            for hlg in group:
+                if isinstance(hlg, HLGExpr):
+                    graphs.append(hlg.dsk)
+                else:
+                    # FinalizeCompute
+                    graphs.append(hlg._layer())
+
+            dsk = HighLevelGraph.merge(*graphs)
+            keys = [v.__dask_keys__() for v in group]
+            if optimizer is not None:
+                dsk = optimizer(dsk, keys)
+            outer_graphs.append(dsk)
+
+        dsk = HighLevelGraph.merge(*outer_graphs)
+        return ensure_dict(dsk)
+
+    _layer = __dask_graph__
+
+    def __dask_annotations__(self):
+        annotations_by_type = {}
+        for hlg in self.operands:
+            for k, v in hlg.__dask_annotations__().items():
+                annotations_by_type.setdefault(k, {}).update(v)
+        return annotations_by_type
+
+    def __dask_keys__(self) -> list:
+        all_keys = []
+        for op in self.operands:
+            all_keys.append(op.__dask_keys__())
+        return all_keys
+
+
+class _ExprSequence(Expr):
+    """A sequence of expressions
+
+    This is used to be able to optimize multiple collections combined, e.g. when
+    being computed simultaneously with ``dask.compute((Expr1, Expr2))``.
+    """
+
+    def __getitem__(self, other):
+        return self.operands[other]
+
+    def _layer(self) -> dict:
+        return toolz.merge(op._layer() for op in self.operands)
+
+    def __dask_keys__(self) -> list:
+        all_keys = []
+        for op in self.operands:
+            all_keys.append(op.__dask_keys__())
+        return all_keys
+
+    def finalize_compute(self):
+        return _ExprSequence(
+            *(op.finalize_compute() for op in self.operands),
+        )
+
+    def __dask_annotations__(self):
+        annotations_by_type = {}
+        for op in self.operands:
+            for k, v in op.__dask_annotations__().items():
+                annotations_by_type.setdefault(k, {}).update(v)
+        return annotations_by_type
+
+    def __len__(self):
+        return len(self.operands)
+
+    def __iter__(self):
+        return iter(self.operands)
+
+    def _simplify_down(self):
+        from dask.highlevelgraph import HighLevelGraph
+
+        issue_warning = False
+        hlgs = []
+        for op in self.operands:
+            if isinstance(op, (HLGExpr, HLGFinalizeCompute)):
+                hlgs.append(op)
+            elif isinstance(op, dict):
+                hlgs.append(
+                    HLGExpr(
+                        dsk=HighLevelGraph.from_collections(
+                            str(id(op)), op, dependencies=()
+                        )
+                    )
+                )
+            elif hlgs:
+                issue_warning = True
+                opt = op.optimize()
+                hlgs.append(
+                    HLGExpr(
+                        dsk=HighLevelGraph.from_collections(
+                            opt._name, opt.__dask_graph__(), dependencies=()
+                        )
+                    )
+                )
+        if issue_warning:
+            warnings.warn(
+                "Computing mixed collections that are backed by "
+                "HighlevelGraphs/dicts and Expressions. "
+                "This forces Expressions to be materialized. "
+                "It is recommended to use only one type and separate the dask."
+                "compute calls if necessary.",
+                UserWarning,
+            )
+        if not hlgs:
+            return None
+        return _HLGExprSequence(*hlgs)
+
+
+class FinalizeCompute(Expr):
+    _parameters = ["expr"]
+
+    def _simplify_down(self):
+        return self.expr.finalize_compute()
+
+
+def _convert_dask_keys(keys):
+    from dask._task_spec import List, TaskRef
+
+    assert isinstance(keys, list)
+    new_keys = []
+    for key in keys:
+        if isinstance(key, list):
+            new_keys.append(_convert_dask_keys(key))
+        else:
+            new_keys.append(TaskRef(key))
+    return List(*new_keys)
+
+
+class HLGFinalizeCompute(Expr):
+    _parameters = ["dsk"]
+
+    def __dask_annotations__(self):
+        return self.dsk.__dask_annotations__()
+
+    def _simplify_down(self):
+        from dask.delayed import Delayed
+
+        # Skip finalization for Delayed
+        if self.dsk.postcompute() == Delayed.__dask_postcompute__(self.dsk):
+            return self.dsk
+        return self
+
+    @property
+    def _name(self):
+        return f"finalize-{self.deterministic_token}"
+
+    def _layer(self) -> dict:
+        expr = self.operand("dsk")
+        dsk = expr._layer().copy()
+
+        func, extra_args = expr.postcompute()
+        keys = expr.__dask_keys__()
+
+        t = Task(self._name, func, _convert_dask_keys(keys), *extra_args)
+        dsk[t.key] = t
+        return dsk
+
+    def __dask_keys__(self):
+        return [self._name]
diff --git a/dask/array/_array_expr/_collection.py b/dask/array/_array_expr/_collection.py
index 99e1e476f..6c51cc6ff 100644
--- a/dask/array/_array_expr/_collection.py
+++ b/dask/array/_array_expr/_collection.py
@@ -51,6 +51,7 @@ class Array(DaskMethodsMixin):
         return from_graph, (
             state._meta,
             state.chunks,
+            # FIXME: This is using keys of the unoptimized graph
             list(flatten(state.__dask_keys__())),
             key_split(state._name),
         )
diff --git a/dask/array/_array_expr/_expr.py b/dask/array/_array_expr/_expr.py
index 52810f06a..b727e763b 100644
--- a/dask/array/_array_expr/_expr.py
+++ b/dask/array/_array_expr/_expr.py
@@ -10,7 +10,7 @@ import numpy as np
 import toolz
 from tlz import accumulate
 
-from dask._expr import Expr
+from dask._expr import Expr, FinalizeCompute
 from dask._task_spec import List, Task, TaskRef
 from dask.array.chunk import getitem
 from dask.array.core import T_IntOrNaN, common_blockdim, unknown_chunk_message
@@ -131,6 +131,9 @@ class ArrayExpr(Expr):
         result.chunks
         return result
 
+    def finalize_compute(self):
+        return FinalizeComputeArray(self)
+
 
 def unify_chunks_expr(*args):
     # TODO(expr): This should probably be a dedicated expression
@@ -267,3 +270,22 @@ class Concatenate(ArrayExpr):
         ]
 
         return dict(zip(keys, values))
+
+
+class FinalizeComputeArray(FinalizeCompute, ArrayExpr):
+    _parameters = ["arr"]
+
+    def chunks(self):
+        return (self.arr.shape,)
+
+    def _simplify_down(self):
+        if self.arr.numblocks in ((), (1,)):
+            return self.arr
+        else:
+            from dask.array._array_expr._rechunk import Rechunk
+
+            return Rechunk(
+                self.arr,
+                tuple(-1 for _ in range(self.arr.ndim)),
+                method="tasks",
+            )
diff --git a/dask/base.py b/dask/base.py
index ba7afe4d3..020523c4c 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -5,16 +5,16 @@ import inspect
 import uuid
 import warnings
 from collections import OrderedDict
-from collections.abc import Hashable, Iterator, Mapping
+from collections.abc import Hashable, Iterable, Iterator, Mapping
 from concurrent.futures import Executor
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
 from functools import partial
 from numbers import Integral, Number
 from operator import getitem
-from typing import Any, Literal, TypeVar
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
-from tlz import groupby, merge
+from tlz import merge
 
 from dask import config, local
 from dask._compatibility import EMSCRIPTEN
@@ -23,7 +23,10 @@ from dask.core import flatten
 from dask.core import get as simple_get
 from dask.system import CPU_COUNT
 from dask.typing import Key, SchedulerGetCallable
-from dask.utils import ensure_dict, is_namedtuple_instance, key_split, shorten_traceback
+from dask.utils import is_namedtuple_instance, key_split, shorten_traceback
+
+if TYPE_CHECKING:
+    from dask._expr import Expr
 
 _DistributedClient = None
 _get_distributed_client = None
@@ -403,54 +406,46 @@ def optimization_function(x):
     return getattr(x, "__dask_optimize__", dont_optimize)
 
 
-def collections_to_dsk(collections, optimize_graph=True, optimizations=(), **kwargs):
-    """
-    Convert many collections into a single dask graph, after optimization
+def collections_to_expr(
+    collections: Iterable,
+    optimize_graph: bool = True,
+) -> Expr:
     """
-    from dask.highlevelgraph import HighLevelGraph
-
-    optimizations = tuple(optimizations) + tuple(config.get("optimizations", ()))
-
-    if optimize_graph:
-        groups = groupby(optimization_function, collections)
+    Convert many collections into a single dask expression.
 
-        graphs = []
-        for opt, val in groups.items():
-            dsk, keys = _extract_graph_and_keys(val)
-            dsk = opt(dsk, keys, **kwargs)
+    Typically, users should not be required to interact with this function.
 
-            for opt_inner in optimizations:
-                dsk = opt_inner(dsk, keys, **kwargs)
-
-            graphs.append(dsk)
-
-        # Merge all graphs
-        if any(isinstance(graph, HighLevelGraph) for graph in graphs):
-            dsk = HighLevelGraph.merge(*graphs)
-        else:
-            dsk = merge(*map(ensure_dict, graphs))
+    Parameters
+    ----------
+    collections : Iterable
+        An iterable of dask collections to be combined.
+    optimize_graph : bool, optional
+        If this is True and collections are encountered which are backed by
+        legacy HighLevelGraph objects, the returned Expression will run a low
+        level task optimization during materialization.
+    """
+    is_iterable = False
+    if isinstance(collections, (tuple, list, set)):
+        is_iterable = True
     else:
-        dsk, _ = _extract_graph_and_keys(collections)
-
-    return dsk
-
+        collections = [collections]
+    if not collections:
+        raise ValueError("No collections provided")
+    from dask._expr import HLGExpr, _ExprSequence
 
-def _extract_graph_and_keys(vals):
-    """Given a list of dask vals, return a single graph and a list of keys such
-    that ``get(dsk, keys)`` is equivalent to ``[v.compute() for v in vals]``."""
-    from dask.highlevelgraph import HighLevelGraph
+    graphs = []
+    for coll in collections:
+        from dask.delayed import Delayed
 
-    graphs, keys = [], []
-    for v in vals:
-        graphs.append(v.__dask_graph__())
-        keys.append(v.__dask_keys__())
+        if isinstance(coll, Delayed) or not hasattr(coll, "expr"):
+            graphs.append(HLGExpr.from_collection(coll, optimize_graph=optimize_graph))
+        else:
+            graphs.append(coll.expr)
 
-    if any(isinstance(graph, HighLevelGraph) for graph in graphs):
-        graph = HighLevelGraph.merge(*graphs)
+    if len(graphs) > 1 or is_iterable:
+        return _ExprSequence(*graphs)
     else:
-        graph = merge(*map(ensure_dict, graphs))
-
-    return graph, keys
+        return graphs[0]
 
 
 def unpack_collections(*args, traverse=True):
@@ -545,7 +540,12 @@ def optimize(*args, traverse=True, **kwargs):
     collections to delayed objects, or to manually apply the optimizations at
     strategic points.
 
-    Note that in most cases you shouldn't need to call this method directly.
+    Note that in most cases you shouldn't need to call this function directly.
+
+    Warning::
+
+        This function triggers a materialization of the collections and looses
+        any annotations attached to HLG layers.
 
     Parameters
     ----------
@@ -577,22 +577,30 @@ def optimize(*args, traverse=True, **kwargs):
     >>> b2.compute() == b.compute()
     np.True_
     """
+    # TODO: This API is problematic. The approach to using postpersist forces us
+    # to materialize the graph. Most low level optimizations will materialize as
+    # well
     collections, repack = unpack_collections(*args, traverse=traverse)
     if not collections:
         return args
 
-    dsk = collections_to_dsk(collections, **kwargs)
+    dsk = collections_to_expr(collections)
 
     postpersists = []
     for a in collections:
         r, s = a.__dask_postpersist__()
-        postpersists.append(r(dsk, *s))
+        postpersists.append(r(dsk.__dask_graph__(), *s))
 
     return repack(postpersists)
 
 
 def compute(
-    *args, traverse=True, optimize_graph=True, scheduler=None, get=None, **kwargs
+    *args,
+    traverse=True,
+    optimize_graph=True,
+    scheduler=None,
+    get=None,
+    **kwargs,
 ):
     """Compute several dask collections at once.
 
@@ -645,17 +653,34 @@ def compute(
         collections=collections,
         get=get,
     )
+    from dask._expr import FinalizeCompute
 
-    dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
-    keys, postcomputes = [], []
-    for x in collections:
-        keys.append(x.__dask_keys__())
-        postcomputes.append(x.__dask_postcompute__())
+    expr = collections_to_expr(collections, optimize_graph)
+    expr = FinalizeCompute(expr)
 
     with shorten_traceback():
-        results = schedule(dsk, keys, **kwargs)
+        # The high level optimize will have to be called client side (for now)
+        # The optimize can internally trigger already a computation
+        # (e.g. parquet is reading some statistics). To move this to the
+        # scheduler we'd need some sort of scheduler-client to trigger a
+        # computation from inside the scheduler and continue with optimization
+        # once the results are in. An alternative could be to introduce a
+        # pre-optimize step for the Expressions that handles steps like these as
+        # a dedicated computation
+
+        # Another caveat is that optimize will only lock in the expression names
+        # after optimization. Names are determined using tokenize and tokenize
+        # is not cross-interpreter (let alone cross-host) stable such that we
+        # have to lock this in before sending stuff (otherwise we'd need to
+        # change the graph submission to a handshake which introduces all sorts
+        # of concurrency control issues)
 
-    return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
+        expr = expr.optimize()
+        keys = list(flatten(expr.__dask_keys__()))
+
+        results = schedule(expr, keys, **kwargs)
+
+    return repack(results)
 
 
 def visualize(
@@ -748,7 +773,7 @@ def visualize(
     """
     args, _ = unpack_collections(*args, traverse=traverse)
 
-    dsk = dict(collections_to_dsk(args, optimize_graph=optimize_graph))
+    dsk = collections_to_expr(args, optimize_graph=optimize_graph).__dask_graph__()
 
     return visualize_dsk(
         dsk=dsk,
@@ -983,7 +1008,8 @@ def persist(*args, traverse=True, optimize_graph=True, scheduler=None, **kwargs)
                     )
                     return repack(results)
 
-    dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
+    expr = collections_to_expr(collections, optimize_graph)
+    expr = expr.optimize()
     keys, postpersists = [], []
     for a in collections:
         a_keys = list(flatten(a.__dask_keys__()))
@@ -992,7 +1018,7 @@ def persist(*args, traverse=True, optimize_graph=True, scheduler=None, **kwargs)
         postpersists.append((rebuild, a_keys, state))
 
     with shorten_traceback():
-        results = schedule(dsk, keys, **kwargs)
+        results = schedule(expr, keys, **kwargs)
 
     d = dict(zip(keys, results))
     results2 = [r({k: d[k] for k in ks}, *s) for r, ks, s in postpersists]
@@ -1115,11 +1141,6 @@ def get_scheduler(get=None, scheduler=None, collections=None, cls=None):
                     _DistributedClient.current(allow_global=True)
                     client_available = True
             if scheduler in named_schedulers:
-                if client_available:
-                    warnings.warn(
-                        "Running on a single-machine scheduler when a distributed client "
-                        "is active might lead to unexpected results."
-                    )
                 return named_schedulers[scheduler]
             elif scheduler in ("dask.distributed", "distributed"):
                 if not client_available:
diff --git a/dask/dataframe/dask_expr/_collection.py b/dask/dataframe/dask_expr/_collection.py
index dabbe4602..92b82b64c 100644
--- a/dask/dataframe/dask_expr/_collection.py
+++ b/dask/dataframe/dask_expr/_collection.py
@@ -457,39 +457,6 @@ Expr={expr}"""
         out = self.optimize(fuse=fuse)
         return DaskMethodsMixin.persist(out, **kwargs)
 
-    def compute(self, fuse=True, concatenate=True, **kwargs):
-        """Compute this DataFrame.
-
-        This turns a lazy Dask DataFrame into an in-memory pandas DataFrame.
-        The entire dataset must fit into memory before calling this operation.
-
-        The optimizer runs over the DataFrame before triggering the computation.
-        The optimizer injects a repartition operation that reduces the partition
-        count to 1 to enable better optimization strategies.
-
-        Parameters
-        ----------
-        fuse : bool, default True
-            Whether to fuse the expression tree before computing. Fusing significantly
-            reduces the number of tasks and improves performance. It shouldn't be
-            disabled unless absolutely necessary.
-        concatenate : bool, default True
-            Whether to concatenate all partitions into a single one before computing.
-            Concatenating enables more powerful optimizations but it also incurs additional
-            data transfer cost. Generally, it should be enabled.
-        kwargs
-            Extra keywords to forward to the base compute function.
-
-        See Also
-        --------
-        dask.compute
-        """
-        out = self
-        if not isinstance(out, Scalar) and concatenate:
-            out = out.repartition(npartitions=1)
-        out = out.optimize(fuse=fuse)
-        return DaskMethodsMixin.compute(out, **kwargs)
-
     def analyze(self, filename: str | None = None, format: str | None = None) -> None:
         """Outputs statistics about every node in the expression.
 
@@ -610,13 +577,23 @@ Expr={expr}"""
             return state.__dask_postcompute__()
         return _concat, ()
 
+    @staticmethod
+    def _postpersist(futures, meta, divisions, name):
+        return from_graph(
+            futures,
+            meta,
+            divisions,
+            sorted(futures),
+            name,
+        )
+
     def __dask_postpersist__(self):
-        state = new_collection(self.expr.lower_completely())
-        return from_graph, (
-            state._meta,
-            state.divisions,
-            state.__dask_keys__(),
-            key_split(state._name),
+        return FrameBase._postpersist, (
+            self._meta,
+            self.divisions,
+            # Note: This prefix is wrong since optimization may actually yield a
+            # different one. That's should only be an issue for visualization.
+            key_split(self._name),
         )
 
     def __getattr__(self, key):
@@ -854,6 +831,7 @@ Expr={expr}"""
         npartitions: int | None = None,
         shuffle_method: str | None = None,
         on_index: bool = False,
+        force: bool = False,
         **options,
     ):
         """Rearrange DataFrame into new partitions
@@ -876,6 +854,9 @@ Expr={expr}"""
         on_index : bool, default False
             Whether to shuffle on the index. Mutually exclusive with 'on'.
             Set this to ``True`` if 'on' is not provided.
+        force : bool, default False
+            This forces the optimizer to keep the shuffle even if the final
+            expression could be further simplified.
         **options : optional
             Algorithm-specific options.
 
@@ -937,7 +918,7 @@ Expr={expr}"""
                     f"p2p requires all column names to be str, found: {unsupported}",
                 )
         # Returned shuffled result
-        return new_collection(
+        res = new_collection(
             RearrangeByColumn(
                 self,
                 on,
@@ -948,6 +929,12 @@ Expr={expr}"""
                 index_shuffle=on_index,
             )
         )
+        if force:
+            # TODO: This forces the optimizer to not remove the shuffle It would
+            # be nice to teach the optimizer directly (e.g. to avoid key
+            # renames, etc.)
+            return res.map_partitions(lambda x: x, meta=res._meta)
+        return res
 
     @derived_from(pd.DataFrame)
     def resample(self, rule, closed=None, label=None):
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index dd67b09d9..408bab56b 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -15,6 +15,7 @@ from pandas.errors import PerformanceWarning
 from tlz import merge_sorted, partition, unique
 
 from dask import _expr as core
+from dask._expr import FinalizeCompute
 from dask._task_spec import Alias, DataNode, Task, TaskRef, execute_graph
 from dask.array import Array
 from dask.core import flatten
@@ -506,6 +507,18 @@ class Expr(core.Expr):
     def fuse(self):
         return optimize_blockwise_fusion(self)
 
+    def finalize_compute(self):
+        return FinalizeComputeDF(self)
+
+
+class FinalizeComputeDF(FinalizeCompute, Expr):
+    _parameters = ["frame"]
+
+    def _simplify_down(self):
+        from dask.dataframe.dask_expr._repartition import Repartition
+
+        return Repartition(self.frame, 1)
+
 
 class Literal(Expr):
     """Represent a literal (known) value as an `Expr`"""
diff --git a/dask/dataframe/dask_expr/_shuffle.py b/dask/dataframe/dask_expr/_shuffle.py
index 271f52aa6..c85726df8 100644
--- a/dask/dataframe/dask_expr/_shuffle.py
+++ b/dask/dataframe/dask_expr/_shuffle.py
@@ -1181,19 +1181,21 @@ class SetPartition(SetIndex):
             drop, set_name = True, "_index"
         else:
             drop, set_name = self.drop, self.other._meta.name
-        kwargs = {
-            "other": self.other._name,
-            "partitions": self._npartitions_input,
-            "ascending": self.ascending,
-            "upsample": self.upsample,
-        }
+        lru_key = (
+            self.other._name,
+            self._npartitions_input,
+            self.ascending,
+            128e6,
+            self.upsample,
+        )
+        computed_divisions = divisions_lru.get(lru_key)
         index_set = _SetIndexPost(
             shuffled,
             self.other._meta.name,
             drop,
             set_name,
             self.frame._meta.columns.dtype,
-            kwargs,
+            computed_divisions,
             self.user_divisions,
         )
         return SortIndexBlockwise(index_set)
@@ -1217,7 +1219,7 @@ class _SetIndexPost(Blockwise):
         "drop",
         "set_name",
         "column_dtype",
-        "key_kwargs",
+        "computed_divisions",
         "user_divisions",
     ]
     _is_length_preserving = True
@@ -1253,16 +1255,8 @@ class _SetIndexPost(Blockwise):
     def _divisions(self):
         if self.operand("user_divisions") is not None:
             return self._get_culled_divisions(self.operand("user_divisions"))
-        kwargs = self.key_kwargs
-        key = (
-            kwargs["other"],
-            kwargs["partitions"],
-            kwargs["ascending"],
-            128e6,
-            kwargs["upsample"],
-        )
-        assert key in divisions_lru
-        return self._get_culled_divisions(divisions_lru[key][0])
+        assert self.computed_divisions is not None
+        return self._get_culled_divisions(self.computed_divisions[0])
 
 
 class SortIndexBlockwise(Blockwise):
diff --git a/dask/dataframe/dask_expr/io/io.py b/dask/dataframe/dask_expr/io/io.py
index 0578dd9a9..4345e8ddd 100644
--- a/dask/dataframe/dask_expr/io/io.py
+++ b/dask/dataframe/dask_expr/io/io.py
@@ -8,7 +8,7 @@ import numpy as np
 import pyarrow as pa
 
 from dask import delayed
-from dask._task_spec import DataNode, List, Task
+from dask._task_spec import Alias, DataNode, List, Task
 from dask.dataframe import methods
 from dask.dataframe._pyarrow import to_pyarrow_string
 from dask.dataframe.core import apply_and_enforce, is_dataframe_like
@@ -26,7 +26,6 @@ from dask.dataframe.dask_expr._reductions import Len
 from dask.dataframe.dask_expr._util import _BackendData, _convert_to_list
 from dask.dataframe.dispatch import make_meta
 from dask.dataframe.io.io import _meta_from_array, sorted_division_locations
-from dask.tokenize import _tokenize_deterministic
 from dask.typing import Key
 from dask.utils import funcname, is_series_like
 
@@ -54,16 +53,14 @@ class FromGraph(IO):
 
     @functools.cached_property
     def _name(self):
-        return (
-            self.operand("name_prefix") + "-" + _tokenize_deterministic(*self.operands)
-        )
+        return self.operand("name_prefix") + "-" + self.deterministic_token
 
     def _layer(self):
         dsk = dict(self.operand("layer"))
         # The name may not actually match the layers name therefore rewrite this
         # using an alias
-        for part, k in enumerate(self.operand("keys")):
-            dsk[(self._name, part)] = k
+        for new, old in zip(self.__dask_keys__(), self.operand("keys")):
+            dsk[new] = Alias(new, old)
         return dsk
 
 
@@ -628,7 +625,6 @@ class FromArray(PartitionsFiltered, BlockwiseIO):
         "columns": None,
         "_partitions": None,
     }
-    _pd_length_stats = None
     _absorb_projections = True
 
     @functools.cached_property
diff --git a/dask/dataframe/dask_expr/io/parquet.py b/dask/dataframe/dask_expr/io/parquet.py
index 2c549560d..8e12e9abd 100644
--- a/dask/dataframe/dask_expr/io/parquet.py
+++ b/dask/dataframe/dask_expr/io/parquet.py
@@ -72,7 +72,6 @@ def _tokenize_fileinfo(fileinfo):
         fileinfo.path,
         fileinfo.size,
         fileinfo.mtime_ns,
-        fileinfo.size,
     )
 
 
diff --git a/dask/dataframe/utils.py b/dask/dataframe/utils.py
index 7260ae0d6..b7b06aae2 100644
--- a/dask/dataframe/utils.py
+++ b/dask/dataframe/utils.py
@@ -607,15 +607,16 @@ def assert_divisions(ddf, scheduler=None):
 
 
 def assert_sane_keynames(ddf):
-    if not hasattr(ddf, "dask"):
-        return
-    for k in ddf.dask.keys():
-        while isinstance(k, tuple):
-            k = k[0]
-        assert isinstance(k, (str, bytes))
-        assert len(k) < 100
-        assert " " not in k
-        assert k.split("-")[0].isidentifier(), k
+    if hasattr(ddf, "__dask_graph__"):
+        dsk = ddf.optimize()
+        dsk = dsk.__dask_graph__()
+        for k in dsk:
+            while isinstance(k, tuple):
+                k = k[0]
+            assert isinstance(k, (str, bytes))
+            assert len(k) < 100
+            assert " " not in k
+            assert k.split("-")[0].isidentifier(), k
 
 
 def assert_dask_dtypes(ddf, res, numeric_equal=True):
diff --git a/dask/delayed.py b/dask/delayed.py
index 1718ca88a..24f8894c2 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -11,9 +11,10 @@ from functools import partial
 from tlz import concat, curry, merge, unique
 
 from dask import config
+from dask._expr import FinalizeCompute
 from dask.base import (
     DaskMethodsMixin,
-    dont_optimize,
+    collections_to_expr,
     is_dask_collection,
     named_schedulers,
     replace_name_in_key,
@@ -206,13 +207,11 @@ def to_task_dask(expr):
         return expr.key, expr.dask
 
     if is_dask_collection(expr):
-        name = "finalize-" + tokenize(expr, pure=True)
-        keys = expr.__dask_keys__()
-        opt = getattr(expr, "__dask_optimize__", dont_optimize)
-        finalize, args = expr.__dask_postcompute__()
-        dsk = {name: (finalize, keys) + args}
-        dsk.update(opt(expr.__dask_graph__(), keys))
-        return name, dsk
+        expr = collections_to_expr(expr)
+        expr = FinalizeCompute(expr)
+        expr = expr.optimize()
+        (name,) = expr.__dask_keys__()
+        return name, expr.__dask_graph__()
 
     if type(expr) is type(iter(list())):
         expr = list(expr)
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index b40fbaf58..4a8603537 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -318,6 +318,8 @@ class MaterializedLayer(Layer):
         super().__init__(
             annotations=annotations, collection_annotations=collection_annotations
         )
+        if not isinstance(mapping, Mapping):
+            raise TypeError(f"mapping must be a Mapping. Instead got {type(mapping)}")
         self.mapping = mapping
 
     def __contains__(self, k):
@@ -996,4 +998,6 @@ def _get_some_layer_name(collection) -> str:
 
 @normalize_token.register(HighLevelGraph)
 def register_highlevelgraph(hlg):
+    # Note: Layer keys are not necessarily identifying HLGs uniquely
+    # see https://github.com/dask/dask/issues/9888
     return normalize_token(list(hlg.layers.keys()))
diff --git a/dask/local.py b/dask/local.py
index 8a8ecaa42..d701615da 100644
--- a/dask/local.py
+++ b/dask/local.py
@@ -426,7 +426,9 @@ def get_async(
         result_flat = {result}
     results = set(result_flat)
 
-    dsk = dict(convert_legacy_graph(dsk))
+    if not isinstance(dsk, Mapping):
+        dsk = dsk.__dask_graph__()
+    dsk = convert_legacy_graph(dsk)
     with local_callbacks(callbacks) as callbacks:
         _, _, pretask_cbs, posttask_cbs, _ = unpack_callbacks(callbacks)
         started_cbs = []
diff --git a/dask/multiprocessing.py b/dask/multiprocessing.py
index 19b024bfa..e40af4e0a 100644
--- a/dask/multiprocessing.py
+++ b/dask/multiprocessing.py
@@ -213,7 +213,9 @@ def get(
             pool = MultiprocessingPoolExecutor(pool)
         cleanup = False
 
-    # Optimize Dask
+    if hasattr(dsk, "__dask_graph__"):
+        dsk = dsk.__dask_graph__()
+
     dsk = ensure_dict(dsk)
     dsk2, dependencies = cull(dsk, keys)
     if optimize_graph:
diff --git a/dask/tokenize.py b/dask/tokenize.py
index 9c7af3cc3..e1b434c41 100644
--- a/dask/tokenize.py
+++ b/dask/tokenize.py
@@ -210,9 +210,8 @@ def normalize_object(o):
         return _normalize_pickle(o)
     except Exception:
         _maybe_raise_nondeterministic(
-            f"Object {o!r} cannot be deterministically hashed. See "
-            "https://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing "
-            "for more information."
+            f"Object {o!r} cannot be deterministically hashed. This likely "
+            "indicates that the object cannot be serialized deterministically."
         )
         return uuid.uuid4().hex
 
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 624bdbbeb..58534b3ac 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -5,6 +5,35 @@ Changelog
 
     This is not exhaustive. For an exhaustive list of changes, see the git log.
 
+Next Release
+------------
+
+Highlights
+^^^^^^^^^^
+- When computing multiple dask Expr backed collections like DataFrames, they are
+  now optimized together instead of individually.
+- Graph materialization and low level optimization is now being performed on the
+  scheduler of a distributed cluster (if available)
+- New kwarg `force` for ``DataFrame.shuffle`` which signals the optimizer to not
+  drop the shuffle during optimization.
+
+Breaking changes
+^^^^^^^^^^^^^^^^
+- Support for custom low level optimizers removed
+- Top level dask.optimize will now always trigger graph materialization.
+  Previously this was not always the case. This also causes any low level HLG
+  annotations to be dropped.
+- DataFrame and Array compute results are now always concatenated on the
+  cluster. Previously, the behavior was dependent on the API used to call
+  compute (dask.compute, DaskCollection.compute or Client.compute).
+- ``dask.base.collections_to_dsk`` has been renamed to `collections_to_expr` and
+  no longer returns a ``HighLevelGraph`` or ``dict`` object but instead
+  guarantees an ``dask._expr.Expr`` object. Further, it no longer performs low
+  level optimization immediately but instead delays until the ``Expr`` instance
+  is materialized, i.e. the returned object is no longer a mapping such that
+  converting it to `dict` or iterating over it is not possible any more.
+
+
 .. _v2025.3.0:
 
 2025.3.0
@@ -193,7 +222,7 @@ See :pr-distributed:`8991` by `Hendrik Makait`_ for more details.
 .. _v2025.1.0:
 
 2025.1.0
----------
+--------
 
 Highlights
 ^^^^^^^^^^
@@ -244,7 +273,7 @@ Improved scheduler responsiveness for large task graphs
 This release reduces the number of Python object references
 related to tracking tasks by the Dask scheduler. This increases
 scheduler responsiveness by reducing the time needed to run
-garbage collection on the scheduler. 
+garbage collection on the scheduler.
 
 See :issue:`8958`, :pr:`11608`, :pr:`11600`, :pr:`11598`,
 :pr:`11597`, and :pr-distributed:`8963` from `Hendrik Makait`_ for more details.


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index c894ca695..27bcbe56c 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -55,7 +55,7 @@ from dask.array.core import (
 from dask.array.numpy_compat import NUMPY_GE_200, NUMPY_GE_210
 from dask.array.reshape import _not_implemented_message
 from dask.array.utils import assert_eq, same_keys
-from dask.base import collections_to_dsk, compute_as_if_collection, tokenize
+from dask.base import collections_to_expr, compute_as_if_collection, tokenize
 from dask.blockwise import (
     _make_blockwise_graph,
     broadcast_dimensions,
@@ -5834,7 +5834,7 @@ def test_from_array_xarray_dataarray():
     xr = pytest.importorskip("xarray")
     arr = xr.DataArray(da.random.random((1000, 1000), chunks=(50, 50)))
     dask_array = da.from_array(arr)
-    dsk = collections_to_dsk([dask_array])
+    dsk = collections_to_expr([dask_array]).__dask_graph__()
     assert len(dsk) == 400
     assert all(k[0].startswith("random_sample") for k in dsk)
     assert_eq(dask_array, arr.data)
diff --git a/dask/array/tests/test_creation.py b/dask/array/tests/test_creation.py
index d0fca5540..efbfb72e0 100644
--- a/dask/array/tests/test_creation.py
+++ b/dask/array/tests/test_creation.py
@@ -5,7 +5,7 @@ import pickle
 import pytest
 
 from dask._task_spec import Alias
-from dask.base import collections_to_dsk
+from dask.base import collections_to_expr
 
 pytest.importorskip("numpy")
 
@@ -567,26 +567,26 @@ def test_diag_extraction(k):
 def test_creation_data_producers():
     x = np.arange(64).reshape((8, 8))
     d = da.from_array(x, chunks=(4, 4))
-    dsk = collections_to_dsk([d])
+    dsk = collections_to_expr([d]).__dask_graph__()
     assert all(v.data_producer for v in dsk.values())
 
     # blockwise fusion
     x = d.astype("float64")
-    dsk = collections_to_dsk([x])
+    dsk = collections_to_expr([x]).__dask_graph__()
     assert sum(v.data_producer for v in dsk.values()) == 4
     assert sum(isinstance(v, Alias) for v in dsk.values()) == 4
     assert len(dsk) == 8
 
     # linear fusion
     x = d[slice(0, 6), None].astype("float64")
-    dsk = collections_to_dsk([x])
+    dsk = collections_to_expr([x]).__dask_graph__()
     assert sum(v.data_producer for v in dsk.values()) == 4
     assert sum(isinstance(v, Alias) for v in dsk.values()) == 4
     assert len(dsk) == 8
 
     # no fusion
     x = d[[1, 3, 5, 7, 6, 4, 2, 0]].astype("float64")
-    dsk = collections_to_dsk([x])
+    dsk = collections_to_expr([x]).__dask_graph__()
     assert sum(v.data_producer and "array-" in k[0] for k, v in dsk.items()) == 4
     assert sum(v.data_producer for v in dsk.values()) == 8  # getitem data nodes
     assert len(dsk) == 24
@@ -1057,7 +1057,7 @@ def test_from_array_getitem_fused():
     arr = np.arange(100).reshape(10, 10)
     darr = da.from_array(arr, chunks=(5, 5))
     result = darr[slice(1, 5), :][slice(1, 3), :]
-    dsk = collections_to_dsk([result])
+    dsk = collections_to_expr([result]).__dask_graph__()
     # Ensure that slices are merged properly
     key = [k for k in dsk if "array-getitem" in k[0]][0]
     key_2 = [
diff --git a/dask/array/tests/test_map_blocks.py b/dask/array/tests/test_map_blocks.py
index e27c74daf..602000192 100644
--- a/dask/array/tests/test_map_blocks.py
+++ b/dask/array/tests/test_map_blocks.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import dask.array as da
-from dask.base import collections_to_dsk
+from dask.base import collections_to_expr
 
 
 def test_map_blocks_block_id_fusion():
@@ -11,5 +11,5 @@ def test_map_blocks_block_id_fusion():
         return x
 
     result = arr.map_blocks(dummy).astype("f8")
-    dsk = collections_to_dsk([result])
-    assert len(dsk) == 20
+    dsk = collections_to_expr([result])
+    assert len(dsk.__dask_graph__()) == 20
diff --git a/dask/array/tests/test_overlap.py b/dask/array/tests/test_overlap.py
index c3d6019c3..f588ca832 100644
--- a/dask/array/tests/test_overlap.py
+++ b/dask/array/tests/test_overlap.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 import pytest
 
 from dask.array import from_array
-from dask.base import collections_to_dsk
+from dask.base import collections_to_expr
 
 pytest.importorskip("numpy")
 
@@ -125,7 +125,7 @@ def test_overlap_internal_asymmetric_small():
 def test_trim_internal():
     d = da.ones((40, 60), chunks=(10, 10))
     e = trim_internal(d, axes={0: 1, 1: 2}, boundary="reflect")
-    assert len(collections_to_dsk([e])) == 24
+    assert len(collections_to_expr([e]).__dask_graph__()) == 24
     assert e.chunks == ((8, 8, 8, 8), (6, 6, 6, 6, 6, 6))
 
 
@@ -960,5 +960,5 @@ def test_overlap_not_blowing_up_graph():
         dims=["x", "y", "z", "year"],
     )
     result = arr.rolling(dim={"year": 11}, center=True).construct("window_dim")
-    dc = collections_to_dsk([result.data])
+    dc = collections_to_expr([result.data]).__dask_graph__()
     assert len(dc) <= 1651  # previously 3000
diff --git a/dask/bag/tests/test_random.py b/dask/bag/tests/test_random.py
index 6a30ba255..158b1b6fd 100644
--- a/dask/bag/tests/test_random.py
+++ b/dask/bag/tests/test_random.py
@@ -30,7 +30,7 @@ def test_choices_empty_partition():
     sut = db.from_sequence(seq, partition_size=9)
     sut = sut.repartition(3)
     li = list(random.choices(sut, k=2).compute())
-    assert sut.map_partitions(len).compute() == (9, 0, 1)
+    assert sut.map_partitions(len).compute() == [9, 0, 1]
     assert len(li) == 2
     assert all(i in seq for i in li)
 
@@ -39,7 +39,7 @@ def test_choices_k_bigger_than_smallest_partition_size():
     seq = range(10)
     sut = db.from_sequence(seq, partition_size=9)
     li = list(random.choices(sut, k=2).compute())
-    assert sut.map_partitions(len).compute() == (9, 1)
+    assert sut.map_partitions(len).compute() == [9, 1]
     assert len(li) == 2
     assert all(i in seq for i in li)
 
@@ -48,7 +48,7 @@ def test_choices_k_equal_bag_size_with_unbalanced_partitions():
     seq = range(10)
     sut = db.from_sequence(seq, partition_size=9)
     li = list(random.choices(sut, k=10).compute())
-    assert sut.map_partitions(len).compute() == (9, 1)
+    assert sut.map_partitions(len).compute() == [9, 1]
     assert len(li) == 10
     assert all(i in seq for i in li)
 
@@ -58,7 +58,7 @@ def test_choices_with_more_bag_partitons():
     seq = range(100)
     sut = db.from_sequence(seq, npartitions=10)
     li = list(random.choices(sut, k=10, split_every=8).compute())
-    assert sut.map_partitions(len).compute() == (10, 10, 10, 10, 10, 10, 10, 10, 10, 10)
+    assert sut.map_partitions(len).compute() == [10, 10, 10, 10, 10, 10, 10, 10, 10, 10]
     assert len(li) == 10
     assert all(i in seq for i in li)
 
@@ -68,7 +68,7 @@ def test_sample_with_more_bag_partitons():
     seq = range(100)
     sut = db.from_sequence(seq, npartitions=10)
     li = list(random.sample(sut, k=10, split_every=8).compute())
-    assert sut.map_partitions(len).compute() == (10, 10, 10, 10, 10, 10, 10, 10, 10, 10)
+    assert sut.map_partitions(len).compute() == [10, 10, 10, 10, 10, 10, 10, 10, 10, 10]
     assert len(li) == 10
     assert all(i in seq for i in li)
     assert len(set(li)) == len(li)
@@ -78,7 +78,7 @@ def test_sample_size_exactly_k():
     seq = range(20)
     sut = db.from_sequence(seq, npartitions=3)
     li = list(random.sample(sut, k=2).compute())
-    assert sut.map_partitions(len).compute() == (7, 7, 6)
+    assert sut.map_partitions(len).compute() == [7, 7, 6]
     assert len(li) == 2
     assert all(i in seq for i in li)
     assert len(set(li)) == len(li)
@@ -97,7 +97,7 @@ def test_sample_empty_partition():
     sut = db.from_sequence(seq, partition_size=9)
     sut = sut.repartition(3)
     li = list(random.sample(sut, k=2).compute())
-    assert sut.map_partitions(len).compute() == (9, 0, 1)
+    assert sut.map_partitions(len).compute() == [9, 0, 1]
     assert len(li) == 2
     assert all(i in seq for i in li)
     assert len(set(li)) == len(li)
@@ -107,7 +107,7 @@ def test_sample_size_k_bigger_than_smallest_partition_size():
     seq = range(10)
     sut = db.from_sequence(seq, partition_size=9)
     li = list(random.sample(sut, k=2).compute())
-    assert sut.map_partitions(len).compute() == (9, 1)
+    assert sut.map_partitions(len).compute() == [9, 1]
     assert len(li) == 2
     assert all(i in seq for i in li)
     assert len(set(li)) == len(li)
@@ -117,7 +117,7 @@ def test_sample_k_equal_bag_size_with_unbalanced_partitions():
     seq = range(10)
     sut = db.from_sequence(seq, partition_size=9)
     li = list(random.sample(sut, k=10).compute())
-    assert sut.map_partitions(len).compute() == (9, 1)
+    assert sut.map_partitions(len).compute() == [9, 1]
     assert len(li) == 10
     assert all(i in seq for i in li)
     assert len(set(li)) == len(li)
diff --git a/dask/dataframe/dask_expr/io/tests/test_distributed.py b/dask/dataframe/dask_expr/io/tests/test_distributed.py
index dc85ab232..e5cd30da2 100644
--- a/dask/dataframe/dask_expr/io/tests/test_distributed.py
+++ b/dask/dataframe/dask_expr/io/tests/test_distributed.py
@@ -4,16 +4,15 @@ import os
 
 import pytest
 
-from dask.dataframe._compat import PYARROW_GE_1500
-from dask.dataframe.dask_expr import read_parquet
-from dask.dataframe.dask_expr.tests._util import _backend_library, assert_eq
-
 distributed = pytest.importorskip("distributed")
 
-from distributed.utils_test import client as c  # noqa F401
+from distributed.utils_test import *  # noqa F401, F403
 from distributed.utils_test import gen_cluster
 
 import dask.dataframe as dd
+from dask.dataframe._compat import PYARROW_GE_1500
+from dask.dataframe.dask_expr import read_parquet
+from dask.dataframe.dask_expr.tests._util import _backend_library, assert_eq
 
 pd = _backend_library()
 
@@ -31,7 +30,7 @@ def _make_file(dir, df=None, filename="myfile.parquet", **kwargs):
     return fn
 
 
-@gen_cluster(client=True)
+@gen_cluster(client=True, clean_kwargs={"threads": False})
 async def test_io_fusion_merge(c, s, a, b, tmpdir):
     pdf = pd.DataFrame({c: range(100) for c in "abcdefghij"})
     p = dd.from_pandas(pdf, 2).to_parquet(tmpdir, compute=False)
@@ -51,11 +50,10 @@ async def test_io_fusion_merge(c, s, a, b, tmpdir):
 
 @pytest.mark.skipif(not PYARROW_GE_1500, reason="requires 15.0.0")
 @pytest.mark.filterwarnings("error")
-@gen_cluster(client=True)
-async def test_parquet_distriuted(c, s, a, b, tmpdir, filesystem):
+def test_parquet_distributed(c, tmpdir, filesystem):
     pdf = pd.DataFrame({"x": [1, 4, 3, 2, 0, 5]})
     df = read_parquet(_make_file(tmpdir, df=pdf), filesystem=filesystem)
-    assert_eq(await c.gather(c.compute(df.optimize())), pdf)
+    assert_eq(c.gather(c.compute(df)), pdf)
 
 
 @pytest.mark.skipif(not PYARROW_GE_1500, reason="requires 15.0.0")
diff --git a/dask/dataframe/dask_expr/tests/test_distributed.py b/dask/dataframe/dask_expr/tests/test_distributed.py
index 4033e2012..a1333c1ea 100644
--- a/dask/dataframe/dask_expr/tests/test_distributed.py
+++ b/dask/dataframe/dask_expr/tests/test_distributed.py
@@ -11,11 +11,9 @@ from dask.dataframe.dask_expr.tests._util import _backend_library
 
 distributed = pytest.importorskip("distributed")
 
-from distributed import Client, LocalCluster, SchedulerPlugin
+from distributed import Client, LocalCluster
 from distributed.shuffle._core import id_from_key
-from distributed.utils_test import cleanup  # noqa F401
-from distributed.utils_test import client as c  # noqa F401
-from distributed.utils_test import gen_cluster, loop, loop_in_thread  # noqa F401
+from distributed.utils_test import gen_cluster  # noqa F401
 
 import dask.dataframe as dd
 
@@ -458,39 +456,3 @@ def test_respect_context_shuffle(df, pdf, func):
     with dask.config.set({"dataframe.shuffle.method": "tasks"}):
         result = q.optimize(fuse=False)
     assert len([x for x in result.walk() if isinstance(x, P2PShuffle)]) > 0
-
-
-@pytest.mark.parametrize("concatenate", [True, False])
-def test_compute_concatenates(loop, concatenate):
-    pdf = pd.DataFrame({"a": np.random.randint(1, 100, (100,)), "b": 1})
-    df = from_pandas(pdf, npartitions=10)
-
-    class Plugin(SchedulerPlugin):
-        def start(self, *args, **kwargs):
-            self.repartition_in_tasks = False
-
-        def update_graph(
-            self,
-            scheduler,
-            *,
-            client,
-            keys,
-            tasks,
-            annotations,
-            priority,
-            dependencies,
-            **kwargs,
-        ):
-            for key in dependencies:
-                if not isinstance(key, tuple):
-                    continue
-                group = key[0]
-                if not isinstance(group, str):
-                    continue
-                self.repartition_in_tasks |= group.startswith("repartitiontofewer")
-
-    with Client(loop=loop) as c:
-        c.register_plugin(Plugin(), name="tracker")
-        df.compute(fuse=False, concatenate=concatenate)
-        plugin = c.cluster.scheduler.plugins["tracker"]
-        assert plugin.repartition_in_tasks is concatenate
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index 55e448357..44d6335e8 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -20,7 +20,7 @@ import dask.dataframe as dd
 import dask.dataframe.groupby
 from dask import delayed
 from dask._compatibility import WINDOWS
-from dask.base import collections_to_dsk
+from dask.base import collections_to_expr
 from dask.dataframe import _compat, methods
 from dask.dataframe._compat import PANDAS_GE_210, PANDAS_GE_220, PANDAS_GE_300, tm
 from dask.dataframe._pyarrow import to_pyarrow_string
@@ -5318,7 +5318,9 @@ def test_from_xarray():
     }
 
     result = a.reindex(**sel_coords).to_dask_dataframe()
-    assert len(collections_to_dsk([result])) < 30  # previously 3000
+    assert (
+        len(collections_to_expr([result]).optimize().__dask_graph__()) < 30
+    )  # previously 3000
 
 
 def test_from_xarray_string_conversion():
diff --git a/dask/dataframe/tests/test_shuffle.py b/dask/dataframe/tests/test_shuffle.py
index 497e0520b..b5cb297e9 100644
--- a/dask/dataframe/tests/test_shuffle.py
+++ b/dask/dataframe/tests/test_shuffle.py
@@ -34,8 +34,8 @@ meta = make_meta(
     {"a": "i8", "b": "i8"}, index=pd.Index([], "i8"), parent_meta=pd.DataFrame()
 )
 d = dd.repartition(pd.concat(dsk.values()), [0, 4, 9, 9])
-shuffle_func = lambda df, *args, **kwargs: df.shuffle(*args, **kwargs)
-shuffle = lambda df, *args, **kwargs: df.shuffle(*args, **kwargs)
+shuffle_func = lambda df, *args, **kwargs: df.shuffle(*args, **kwargs, force=True)
+shuffle = lambda df, *args, **kwargs: df.shuffle(*args, **kwargs, force=True)
 full = d.compute()
 
 
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index 0a07ce438..76df1b678 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -19,7 +19,7 @@ import dask.bag as db
 from dask.base import (
     DaskMethodsMixin,
     clone_key,
-    collections_to_dsk,
+    collections_to_expr,
     compute,
     compute_as_if_collection,
     get_collection_names,
@@ -234,6 +234,7 @@ def test_replace_name_in_keys():
 class Tuple(DaskMethodsMixin):
     __slots__ = ("_dask", "_keys")
     __dask_scheduler__ = staticmethod(dask.threaded.get)
+    __dask_optimize__ = None
 
     def __init__(self, dsk, keys):
         self._dask = dsk
@@ -463,7 +464,8 @@ def test_compute_array_dataframe():
     darr = da.from_array(arr, chunks=(5, 5)) + 1
     df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [5, 5, 3, 3]})
     ddf = dd.from_pandas(df, npartitions=2).a + 2
-    arr_out, df_out = compute(darr, ddf)
+    with pytest.warns(UserWarning, match="mixed.*materialize"):
+        arr_out, df_out = compute(darr, ddf)
     assert np.allclose(arr_out, arr + 1)
     dd.utils.assert_eq(df_out, df.a + 2)
 
@@ -606,16 +608,6 @@ def inc_to_dec(dsk, keys):
     return dsk
 
 
-def test_optimizations_keyword():
-    x = dask.delayed(inc)(1)
-    assert x.compute() == 2
-
-    with dask.config.set(optimizations=[inc_to_dec]):
-        assert x.compute() == 0
-
-    assert x.compute() == 2
-
-
 def test_optimize():
     x = dask.delayed(inc)(1)
     y = dask.delayed(inc)(x)
@@ -632,18 +624,6 @@ def test_optimize():
     # Computationally equivalent
     assert dask.compute(x2, y2, z2) == dask.compute(x, y, z)
 
-    # Applying optimizations before compute and during compute gives
-    # same results. Shows optimizations are occurring.
-    sols = dask.compute(x, y, z, optimizations=[inc_to_dec])
-    x3, y3, z3 = optimize(x, y, z, optimizations=[inc_to_dec])
-    assert dask.compute(x3, y3, z3) == sols
-
-    # Optimize respects global optimizations as well
-    with dask.config.set(optimizations=[inc_to_dec]):
-        x4, y4, z4 = optimize(x, y, z)
-    for a, b in zip([x3, y3, z3], [x4, y4, z4]):
-        assert dict(a.dask) == dict(b.dask)
-
 
 def test_optimize_nested():
     a = dask.delayed(inc)(1)
@@ -866,7 +846,8 @@ def test_optimize_None():
     y = x[:9][1:8][::2] + 1  # normally these slices would be fused
 
     def my_get(dsk, keys):
-        assert dsk == dict(y.dask)  # but they aren't
+        # but they aren't. +1 for the finalize task
+        assert len(dsk.__dask_graph__()) == len(y.dask) + 1
         return dask.get(dsk, keys)
 
     with dask.config.set(array_optimize=None, scheduler=my_get):
@@ -875,7 +856,7 @@ def test_optimize_None():
 
 def test_scheduler_keyword():
     def schedule(dsk, keys, **kwargs):
-        return [[123]]
+        return [123]
 
     named_schedulers["foo"] = schedule
 
@@ -954,9 +935,9 @@ def test_optimizations_ctd():
     pytest.importorskip("numpy")
     da = pytest.importorskip("dask.array")
     x = da.arange(2, chunks=1)[:1]
-    dsk1 = collections_to_dsk([x])
+    dsk1 = collections_to_expr([x])
     with dask.config.set({"optimizations": [lambda dsk, keys: dsk]}):
-        dsk2 = collections_to_dsk([x])
+        dsk2 = collections_to_expr([x])
     assert dsk1 == dsk2
 
 
diff --git a/dask/tests/test_delayed.py b/dask/tests/test_delayed.py
index 4ce1dbc8e..41a32817c 100644
--- a/dask/tests/test_delayed.py
+++ b/dask/tests/test_delayed.py
@@ -16,7 +16,7 @@ from tlz import merge
 import dask
 import dask.bag as db
 from dask import compute
-from dask.base import collections_to_dsk
+from dask.base import collections_to_expr
 from dask.delayed import Delayed, delayed, to_task_dask
 from dask.highlevelgraph import HighLevelGraph
 from dask.utils_test import inc
@@ -24,6 +24,7 @@ from dask.utils_test import inc
 
 class Tuple:
     __dask_scheduler__ = staticmethod(dask.threaded.get)
+    __dask_optimize__ = None
 
     def __init__(self, dsk, keys):
         self._dask = dsk
@@ -80,7 +81,8 @@ def test_to_task_dask():
     task, dask = to_task_dask(x)
     assert task in dask
     f = dask.pop(task)
-    assert f == (tuple, ["a", "b", "c"])
+    assert f.func == tuple
+    assert f.dependencies == {"a", "b", "c"}
     assert dask == x._dask
 
 
@@ -328,16 +330,6 @@ def test_common_subexpressions():
     assert len(res.dask) == 3
 
 
-def test_delayed_optimize():
-    x = Delayed("b", {"a": 1, "b": (inc, "a"), "c": (inc, "b")})
-    (x2,) = dask.optimize(x)
-    # Delayed's __dask_optimize__ culls out 'c'
-    assert sorted(x2.dask.keys()) == ["a", "b"]
-    assert x2._layer != x2._key
-    # Optimize generates its own layer name, which doesn't match the key.
-    # `Delayed._rebuild` handles this.
-
-
 def test_lists():
     a = delayed(1)
     b = delayed(2)
@@ -837,14 +829,10 @@ def test_annotations_survive_optimization():
     assert len(d.dask.layers) == 1
     assert len(d.dask.layers["b"]) == 3
     assert d.dask.layers["b"].annotations == {"foo": "bar"}
-
-    # Ensure optimizing a Delayed object returns a HighLevelGraph
-    # and doesn't loose annotations
-    (d_opt,) = dask.optimize(d)
-    assert type(d_opt.dask) is HighLevelGraph
-    assert len(d_opt.dask.layers) == 1
-    assert len(d_opt.dask.layers["b"]) == 2  # c is culled
-    assert d_opt.dask.layers["b"].annotations == {"foo": "bar"}
+    optimized = collections_to_expr([d]).optimize()
+    assert optimized.__dask_annotations__() == {
+        "foo": {k: "bar" for k in optimized.__dask_graph__()}
+    }
 
 
 def test_delayed_function_attributes_forwarded():
@@ -872,12 +860,12 @@ def test_delayed_fusion():
         return i + 3
 
     obj = test3(test2(test(10)))
-    dsk = dict(collections_to_dsk([obj]))
+    dsk = collections_to_expr([obj]).__dask_graph__()
     assert len(dsk) == 3
 
     obj2 = test3(test2(test(10)))
     with dask.config.set({"optimization.fuse.delayed": True}):
-        dsk2 = dict(collections_to_dsk([obj]))
+        dsk2 = collections_to_expr([obj]).__dask_graph__()
         result = dask.compute(obj2)
     assert len(dsk2) == 2
     assert dask.compute(obj) == result
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index d92ba547a..63635577e 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -892,23 +892,10 @@ def test_get_scheduler_without_distributed_raises():
         get_scheduler(scheduler="distributed")
 
 
-def test_get_scheduler_with_distributed_active(c):
-    assert get_scheduler() == c.get
-    warning_message = (
-        "Running on a single-machine scheduler when a distributed client "
-        "is active might lead to unexpected results."
-    )
-    with pytest.warns(UserWarning, match=warning_message) as user_warnings_a:
-        get_scheduler(scheduler="threads")
-        get_scheduler(scheduler="sync")
-    assert len(user_warnings_a) == 2
-
-
 def test_get_scheduler_with_distributed_active_reset_config(c):
     assert get_scheduler() == c.get
     with dask.config.set(scheduler="threads"):
-        with pytest.warns(UserWarning):
-            assert get_scheduler() != c.get
+        assert get_scheduler() != c.get
         with dask.config.set(scheduler=None):
             assert get_scheduler() == c.get
 
@@ -981,7 +968,7 @@ def test_get_scheduler_default_client_config_interleaving(s):
     # This test is using context managers intentionally. We should not refactor
     # this to use it in more places to make the client closing cleaner.
     s_address = s["address"]
-    with pytest.warns(UserWarning), dask.config.set(scheduler="sync"):
+    with dask.config.set(scheduler="sync"):
         assert dask.base.get_scheduler() == dask.local.get_sync
         with dask.config.set(scheduler="threads"):
             assert dask.base.get_scheduler() == dask.threaded.get
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index c1711d537..4074c5b78 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -9,7 +9,7 @@ from packaging.version import Version
 import dask
 from dask import delayed
 from dask._compatibility import WINDOWS
-from dask.base import collections_to_dsk, key_split, visualize_dsk
+from dask.base import collections_to_expr, key_split, visualize_dsk
 from dask.core import get_deps
 from dask.order import _connecting_to_roots, diagnostics, ndependencies, order
 from dask.utils_test import add, inc
@@ -36,7 +36,9 @@ def visualize(dsk, suffix="", **kwargs):
     measure used."""
     funcname = inspect.stack()[1][3] + suffix
     if hasattr(dsk, "__dask_graph__"):
-        dsk = collections_to_dsk([dsk], optimize_graph=True)
+        dsk = (
+            collections_to_expr([dsk], optimize_graph=True).optimize().__dask_graph__()
+        )
 
     node_attrs = {"penwidth": "6"}
     visualize_dsk(dsk, filename=f"{funcname}.pdf", node_attr=node_attrs, **kwargs)
@@ -1413,9 +1415,13 @@ def test_array_vs_dataframe(optimize):
     quad = ds**2
     quad["uv"] = ds.anom_u * ds.anom_v
     mean = quad.mean("time")
-    diag_array = diagnostics(collections_to_dsk([mean], optimize_graph=optimize))
+    diag_array = diagnostics(
+        collections_to_expr([mean], optimize_graph=optimize).optimize().__dask_graph__()
+    )
     diag_df = diagnostics(
-        collections_to_dsk([mean.to_dask_dataframe()], optimize_graph=optimize)
+        collections_to_expr([mean.to_dask_dataframe()], optimize_graph=optimize)
+        .optimize()
+        .__dask_graph__()
     )
     assert max(diag_df[1]) == 15
     assert max(diag_array[1]) == 38
@@ -1440,8 +1446,8 @@ def test_anom_mean():
     clim = arr.groupby("day").mean(dim="time")
     anom = arr.groupby("day") - clim
     anom_mean = anom.mean(dim="time")
-    graph = collections_to_dsk([anom_mean])
-    dependencies, dependents = get_deps(graph)
+    graph = collections_to_expr([anom_mean]).optimize().__dask_graph__()
+    _, dependents = get_deps(graph)
     diags, pressure = diagnostics(graph)
     # Encoding the "best" ordering for this graph is tricky. When inspecting the
     # visualization, one sees that there are small, connected tree-like steps at
@@ -2019,9 +2025,17 @@ def test_reduce_with_many_common_dependents(optimize, keep_self, ndeps, n_reduce
 
     if keep_self:
         # Keeping self adds a layer that cannot be fused
-        dsk = collections_to_dsk([x, graph], optimize_graph=optimize)
+        dsk = (
+            collections_to_expr([x, graph], optimize_graph=optimize)
+            .optimize()
+            .__dask_graph__()
+        )
     else:
-        dsk = collections_to_dsk([graph], optimize_graph=optimize)
+        dsk = (
+            collections_to_expr([graph], optimize_graph=optimize)
+            .optimize()
+            .__dask_graph__()
+        )
     dependencies, dependents = get_deps(dsk)
     # Verify assumptions
     before = len(dsk)