Repository: dask/distributed

Issue:
[Tracking] Annotations-related issues with P2P
The implementation of P2P shuffling/rechunking relies on annotations. Unfortunately, there have been a number of issues caused by annotations getting dropped before the execution of the P2P shuffle/rechunk.

* https://github.com/dask/dask/issues/10118
* https://github.com/dask/dask/issues/10074
* #7715
* `da.rechunk` only works if `"optimization.fuse.active": False` is configured due to dask/dask#7036
* https://github.com/dask/dask/pull/10123

**Additional thoughts**

To get out of the mess we're currently in, I'm wondering if making annotations available from within tasks as suggested in https://github.com/dask/distributed/issues/6932 would make it easier to test that they do not get lost.

Together with the work on High-Level Expressions (dask/dask#7933), we should rethink the current API for annotations. It seems to be too easy to drop annotations somewhere along the way.

Discussion:
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/distributed/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

       26 files  ±0         26 suites  ±0   13h 27m 53s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "duration of all tests") -21s
  3 552 tests ±0    3 412 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "passed tests")  -   30     106 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "skipped / disabled tests") +  1    34 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "failed tests") +  29 
44 952 runs  ±0  42 466 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "passed tests")  - 355  2 137 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "skipped / disabled tests") +11  349 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "failed tests") +344 

For more details on these failures, see [this check](https://github.com/dask/distributed/runs/13017380956).

Results for commit 03590780. ± Comparison against base commit 76bbfaf9.

<details>
  <summary>This pull request <b>skips</b> 1 test.</summary>

```
distributed.shuffle.tests.test_shuffle ‑ test_raise_on_lost_annotation
```
</details>

[test-results]:data:application/gzip;base64,H4sIABwlSGQC/03MTQ7CIBCG4as0rF0MMEDxMgYoJMT+GEpXxrsLNWCX3zOT901CnP1O7gOTt4HsR8z/NR3J5LitZeOIihcq11zvXAjW5mM/nKuG9GLP+CpGQXYKJs7nWxef0pYKQZF0rDWMqM9ynS2MDKXs9gszylWjHtZNrmG3LUvMZRDgQoMawU1iYt5aNGAdBBm8sAqCogy8lsaTzxeZwL9RFwEAAA==

mrocklin: Cc @rjzamora and @d-v-b so that they're aware of this possibility
hendrikmakait: Here are results from an A/B test on this, it looks like there is no significant performance impact: 
<img width="1674" alt="Screenshot 2023-04-27 at 16 54 24" src="https://user-images.githubusercontent.com/2699097/234902129-7c80ecdf-824d-4513-a63b-1e4b2ce0f711.png">

fjetter: I think this all makes sense. I wonder if we still want to keep the annotations path around for the "happy case" where HLGs and annotations just work
hendrikmakait: > How do we deal with pre-existing worker restrictions that do not match the output worker we have chosen for the given partition?

If the worker restrictions are applied to the shuffle, they apply to all tasks. Thus, we can use the worker restrictions of the barrier task to bootstrap our set of valid workers. 

As things stand right now, I cannot think of another scenario that would add worker restrictions to tasks. With task fusion, annotations are still getting lost (dask/dask#7036). If this changes, we need to revisit this.

Patch:
diff --git a/distributed/client.py b/distributed/client.py
index a0fb7b78..003136df 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -47,6 +47,7 @@ from dask.utils import (
 from dask.widgets import get_template
 
 from distributed.core import ErrorMessage
+from distributed.protocol.serialize import _is_dumpable
 from distributed.utils import wait_for
 
 try:
@@ -2016,12 +2017,12 @@ class Client(SyncMethodMixin):
             Whether or not the function is pure.  Set ``pure=False`` for
             impure functions like ``np.random.random``.
             See :ref:`pure functions` for more details.
-        batch_size : int, optional
+        batch_size : int, optional (default: just one batch whose size is the entire iterable)
             Submit tasks to the scheduler in batches of (at most)
             ``batch_size``.
-            Larger batch sizes can be useful for very large ``iterables``,
-            as the cluster can start processing tasks while later ones are
-            submitted asynchronously.
+            The tradeoff in batch size is that large batches avoid more per-batch overhead,
+            but batches that are too big can take a long time to submit and unreasonably delay
+            the cluster from starting its processing.
         **kwargs : dict
             Extra keyword arguments to send to the function.
             Large values will be included explicitly in the task graph.
@@ -4412,6 +4413,10 @@ class Client(SyncMethodMixin):
         >>> from time import time
         >>> client.log_event("current-time", time())
         """
+        if not _is_dumpable(msg):
+            raise TypeError(
+                f"Message must be msgpack serializable. Got {type(msg)=} instead."
+            )
         return self.sync(self.scheduler.log_event, topic=topic, msg=msg)
 
     def get_events(self, topic: str | None = None):
diff --git a/distributed/exceptions.py b/distributed/exceptions.py
new file mode 100644
index 00000000..bd18fe57
--- /dev/null
+++ b/distributed/exceptions.py
@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+
+class Reschedule(Exception):
+    """Reschedule this task
+
+    Raising this exception will stop the current execution of the task and ask
+    the scheduler to reschedule this task, possibly on a different machine.
+
+    This does not guarantee that the task will move onto a different machine.
+    The scheduler will proceed through its normal heuristics to determine the
+    optimal machine to accept this task.  The machine will likely change if the
+    load across the cluster has significantly changed since first scheduling
+    the task.
+    """
diff --git a/distributed/nanny.py b/distributed/nanny.py
index 5bc120c3..0f9aa0c5 100644
--- a/distributed/nanny.py
+++ b/distributed/nanny.py
@@ -43,6 +43,7 @@ from distributed.node import ServerNode
 from distributed.process import AsyncProcess
 from distributed.proctitle import enable_proctitle_on_children
 from distributed.protocol import pickle
+from distributed.protocol.serialize import _is_dumpable
 from distributed.security import Security
 from distributed.utils import (
     get_ip,
@@ -610,6 +611,24 @@ class Nanny(ServerNode):
         )
 
     def log_event(self, topic, msg):
+        """Log an event under a given topic
+
+        Parameters
+        ----------
+        topic : str, list[str]
+            Name of the topic under which to log an event. To log the same
+            event under multiple topics, pass a list of topic names.
+        msg
+            Event message to log. Note this must be msgpack serializable.
+
+        See also
+        --------
+        Client.log_event
+        """
+        if not _is_dumpable(msg):
+            raise TypeError(
+                f"Message must be msgpack serializable. Got {type(msg)=} instead."
+            )
         self._ongoing_background_tasks.call_soon(self._log_event, topic, msg)
 
 
diff --git a/distributed/protocol/serialize.py b/distributed/protocol/serialize.py
index 0ba323a6..b2ca44b4 100644
--- a/distributed/protocol/serialize.py
+++ b/distributed/protocol/serialize.py
@@ -849,6 +849,27 @@ def _is_msgpack_serializable(v):
     )
 
 
+def _is_dumpable(v):
+    typ = type(v)
+    return (
+        v is None
+        or typ is str
+        or typ is bool
+        or typ is bytes
+        or typ is int
+        or typ is float
+        or typ is Pickled
+        or typ is Serialize
+        or typ is Serialized
+        or typ is ToPickle
+        or isinstance(v, dict)
+        and all(map(_is_dumpable, v.values()))
+        and all(type(x) is str for x in v.keys())
+        or isinstance(v, (list, tuple))
+        and all(map(_is_dumpable, v))
+    )
+
+
 class ObjectDictSerializer:
     def __init__(self, serializer):
         self.serializer = serializer
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index e8f5918c..c81b30c7 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -7818,6 +7818,20 @@ class Scheduler(SchedulerState, ServerNode):
         return results
 
     def log_event(self, topic: str | Collection[str], msg: Any) -> None:
+        """Log an event under a given topic
+
+        Parameters
+        ----------
+        topic : str, list[str]
+            Name of the topic under which to log an event. To log the same
+            event under multiple topics, pass a list of topic names.
+        msg
+            Event message to log. Note this must be msgpack serializable.
+
+        See also
+        --------
+        Client.log_event
+        """
         event = (time(), msg)
         if not isinstance(topic, str):
             for t in topic:
diff --git a/distributed/shuffle/_merge.py b/distributed/shuffle/_merge.py
index a4eeb4a9..51cbdb47 100644
--- a/distributed/shuffle/_merge.py
+++ b/distributed/shuffle/_merge.py
@@ -136,6 +136,7 @@ def merge_transfer(
     id: ShuffleId,
     input_partition: int,
     npartitions: int,
+    parts_out: set[int],
 ):
     return shuffle_transfer(
         input=input,
@@ -143,6 +144,7 @@ def merge_transfer(
         input_partition=input_partition,
         npartitions=npartitions,
         column=_HASH_COLUMN_NAME,
+        parts_out=parts_out,
     )
 
 
@@ -340,6 +342,7 @@ class HashJoinP2PLayer(Layer):
                 token_left,
                 i,
                 self.npartitions,
+                self.parts_out,
             )
         for i in range(self.n_partitions_right):
             transfer_keys_right.append((name_right, i))
@@ -349,6 +352,7 @@ class HashJoinP2PLayer(Layer):
                 token_right,
                 i,
                 self.npartitions,
+                self.parts_out,
             )
 
         _barrier_key_left = barrier_key(ShuffleId(token_left))
diff --git a/distributed/shuffle/_rechunk.py b/distributed/shuffle/_rechunk.py
index 6b644bcf..e2a7d2a4 100644
--- a/distributed/shuffle/_rechunk.py
+++ b/distributed/shuffle/_rechunk.py
@@ -8,6 +8,7 @@ import dask
 from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph, MaterializedLayer
 
+from distributed.exceptions import Reschedule
 from distributed.shuffle._shuffle import (
     ShuffleId,
     ShuffleType,
@@ -57,6 +58,8 @@ def rechunk_unpack(
         return _get_worker_extension().get_output_partition(
             id, barrier_run_id, output_chunk
         )
+    except Reschedule as e:
+        raise e
     except Exception as e:
         raise RuntimeError(f"rechunk_unpack failed during shuffle {id}") from e
 
@@ -70,12 +73,6 @@ def rechunk_p2p(x: da.Array, chunks: ChunkedAxes) -> da.Array:
         # Special case for empty array, as the algorithm below does not behave correctly
         return da.empty(x.shape, chunks=chunks, dtype=x.dtype)
 
-    if dask.config.get("optimization.fuse.active") is not False:
-        raise RuntimeError(
-            "P2P rechunking requires the fuse optimization to be turned off. "
-            "Set the 'optimization.fuse.active' config to False to deactivate."
-        )
-
     dsk: dict = {}
     token = tokenize(x, chunks)
     _barrier_key = barrier_key(ShuffleId(token))
diff --git a/distributed/shuffle/_scheduler_extension.py b/distributed/shuffle/_scheduler_extension.py
index 33e79c58..e990e67d 100644
--- a/distributed/shuffle/_scheduler_extension.py
+++ b/distributed/shuffle/_scheduler_extension.py
@@ -6,7 +6,9 @@ import itertools
 import logging
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, ClassVar
+from functools import partial
+from itertools import product
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Iterable, Sequence
 
 from distributed.diagnostics.plugin import SchedulerPlugin
 from distributed.shuffle._rechunk import ChunkedAxes, NIndex
@@ -105,6 +107,7 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
                 "shuffle_barrier": self.barrier,
                 "shuffle_get": self.get,
                 "shuffle_get_or_create": self.get_or_create,
+                "shuffle_restrict_task": self.restrict_task,
             }
         )
         self.heartbeats = defaultdict(lambda: defaultdict(dict))
@@ -122,6 +125,14 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
             msg=msg, workers=list(shuffle.participating_workers)
         )
 
+    def restrict_task(self, id: ShuffleId, run_id: int, key: str, worker: str) -> dict:
+        shuffle = self.states[id]
+        if shuffle.run_id != run_id:
+            return {"status": "error", "message": "Stale shuffle"}
+        ts = self.scheduler.tasks[key]
+        self._set_restriction(ts, worker)
+        return {"status": "OK"}
+
     def heartbeat(self, ws: WorkerState, data: dict) -> None:
         for shuffle_id, d in data.items():
             if shuffle_id in self.shuffle_ids():
@@ -129,7 +140,7 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
 
     def get(self, id: ShuffleId, worker: str) -> dict[str, Any]:
         if exception := self.erred_shuffles.get(id):
-            return {"status": "ERROR", "message": str(exception)}
+            return {"status": "error", "message": str(exception)}
         state = self.states[id]
         state.participating_workers.add(worker)
         return state.to_msg()
@@ -144,6 +155,11 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
         try:
             return self.get(id, worker)
         except KeyError:
+            # FIXME: The current implementation relies on the barrier task to be
+            # known by its name. If the name has been mangled, we cannot guarantee
+            # that the shuffle works as intended and should fail instead.
+            self._raise_if_barrier_unknown(id)
+
             state: ShuffleState
             if type == ShuffleType.DATAFRAME:
                 state = self._create_dataframe_shuffle_state(id, spec)
@@ -155,33 +171,33 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
             state.participating_workers.add(worker)
             return state.to_msg()
 
+    def _raise_if_barrier_unknown(self, id: ShuffleId) -> None:
+        key = barrier_key(id)
+        try:
+            self.scheduler.tasks[key]
+        except KeyError:
+            raise RuntimeError(
+                f"Barrier task with key {key!r} does not exist. This may be caused by "
+                "task fusion during graph generation. Please let us know that you ran "
+                "into this by leaving a comment at distributed#7816."
+            )
+
     def _create_dataframe_shuffle_state(
         self, id: ShuffleId, spec: dict[str, Any]
     ) -> DataFrameShuffleState:
         schema = spec["schema"]
         column = spec["column"]
         npartitions = spec["npartitions"]
+        parts_out = spec["parts_out"]
         assert schema is not None
         assert column is not None
         assert npartitions is not None
+        assert parts_out is not None
 
-        workers = list(self.scheduler.workers)
-        output_workers = set()
-
-        name = barrier_key(id)
-        mapping = {}
+        pick_worker = partial(get_worker_for_range_sharding, npartitions)
 
-        for ts in self.scheduler.tasks[name].dependents:
-            part = get_partition_id(ts)
-            if ts.worker_restrictions:
-                output_worker = list(ts.worker_restrictions)[0]
-            else:
-                output_worker = get_worker_for_range_sharding(
-                    part, workers, npartitions
-                )
-            mapping[part] = output_worker
-            output_workers.add(output_worker)
-            self.scheduler.set_restrictions({ts.key: {output_worker}})
+        mapping = self._pin_output_workers(id, parts_out, pick_worker)
+        output_workers = set(mapping.values())
 
         return DataFrameShuffleState(
             id=id,
@@ -193,6 +209,52 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
             participating_workers=output_workers.copy(),
         )
 
+    def _pin_output_workers(
+        self,
+        id: ShuffleId,
+        output_partitions: Iterable[Any],
+        pick: Callable[[Any, Sequence[str]], str],
+    ) -> dict[Any, str]:
+        """Pin the outputs of a P2P shuffle to specific workers.
+
+        Parameters
+        ----------
+        id: ID of the shuffle to pin
+        output_partitions: Output partition IDs to pin
+        pick: Function that picks a worker given a partition ID and sequence of worker
+
+        .. note:
+            This function assumes that the barrier task and the output tasks share
+            the same worker restrictions.
+        """
+        mapping = {}
+        barrier = self.scheduler.tasks[barrier_key(id)]
+
+        if barrier.worker_restrictions:
+            workers = list(barrier.worker_restrictions)
+        else:
+            workers = list(self.scheduler.workers)
+
+        for partition in output_partitions:
+            worker = pick(partition, workers)
+            mapping[partition] = worker
+
+        for dt in barrier.dependents:
+            try:
+                partition = dt.annotations["shuffle"]
+            except KeyError:
+                continue
+
+            if dt.worker_restrictions:
+                worker = pick(partition, list(dt.worker_restrictions))
+                mapping[partition] = worker
+            else:
+                worker = mapping[partition]
+
+            self._set_restriction(dt, worker)
+
+        return mapping
+
     def _create_array_rechunk_state(
         self, id: ShuffleId, spec: dict[str, Any]
     ) -> ArrayRechunkState:
@@ -201,21 +263,9 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
         assert old is not None
         assert new is not None
 
-        workers = list(self.scheduler.workers)
-        output_workers = set()
-
-        name = barrier_key(id)
-        mapping = {}
-
-        for ts in self.scheduler.tasks[name].dependents:
-            part = get_partition_id(ts)
-            if ts.worker_restrictions:
-                output_worker = list(ts.worker_restrictions)[0]
-            else:
-                output_worker = get_worker_for_hash_sharding(part, workers)
-            mapping[part] = output_worker
-            output_workers.add(output_worker)
-            self.scheduler.set_restrictions({ts.key: {output_worker}})
+        parts_out = product(*(range(len(c)) for c in new))
+        mapping = self._pin_output_workers(id, parts_out, get_worker_for_hash_sharding)
+        output_workers = set(mapping.values())
 
         return ArrayRechunkState(
             id=id,
@@ -227,6 +277,22 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
             participating_workers=output_workers.copy(),
         )
 
+    def _set_restriction(self, ts: TaskState, worker: str) -> None:
+        if "shuffle_original_restrictions" in ts.annotations:
+            # This may occur if multiple barriers share the same output task,
+            # e.g. in a hash join.
+            return
+        ts.annotations["shuffle_original_restrictions"] = ts.worker_restrictions.copy()
+        self.scheduler.set_restrictions({ts.key: {worker}})
+
+    def _unset_restriction(self, ts: TaskState) -> None:
+        # shuffle_original_restrictions is only set if the task was first scheduled
+        # on the wrong worker
+        if "shuffle_original_restrictions" not in ts.annotations:
+            return
+        original_restrictions = ts.annotations.pop("shuffle_original_restrictions")
+        self.scheduler.set_restrictions({ts.key: original_restrictions})
+
     def remove_worker(self, scheduler: Scheduler, worker: str) -> None:
         from time import time
 
@@ -247,7 +313,7 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
                 for dt in barrier_task.dependents:
                     if worker not in dt.worker_restrictions:
                         continue
-                    dt.worker_restrictions.clear()
+                    self._unset_restriction(dt)
                     recs.update({dt.key: "waiting"})
                 # TODO: Do we need to handle other states?
 
@@ -293,34 +359,27 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
         with contextlib.suppress(KeyError):
             del self.heartbeats[id]
 
+        barrier_task = self.scheduler.tasks[barrier_key(id)]
+        for dt in barrier_task.dependents:
+            self._unset_restriction(dt)
+
     def restart(self, scheduler: Scheduler) -> None:
         self.states.clear()
         self.heartbeats.clear()
         self.erred_shuffles.clear()
 
 
-def get_partition_id(ts: TaskState) -> Any:
-    """Get the output partition ID of this task state."""
-    try:
-        return ts.annotations["shuffle"]
-    except KeyError:
-        raise RuntimeError(
-            f"{ts} has lost its ``shuffle`` annotation. This may be caused by "
-            "unintended optimization during graph generation. "
-            "Please report this problem on GitHub and link it to "
-            "the tracking issue at https://github.com/dask/distributed/issues/7716."
-        )
-
-
 def get_worker_for_range_sharding(
-    output_partition: int, workers: list[str], npartitions: int
+    npartitions: int, output_partition: int, workers: Sequence[str]
 ) -> str:
     """Get address of target worker for this output partition using range sharding"""
     i = len(workers) * output_partition // npartitions
     return workers[i]
 
 
-def get_worker_for_hash_sharding(output_partition: NIndex, workers: list[str]) -> str:
+def get_worker_for_hash_sharding(
+    output_partition: NIndex, workers: Sequence[str]
+) -> str:
     """Get address of target worker for this output partition using hash sharding"""
     i = hash(output_partition) % len(workers)
     return workers[i]
diff --git a/distributed/shuffle/_shuffle.py b/distributed/shuffle/_shuffle.py
index 2619f7f5..7332097c 100644
--- a/distributed/shuffle/_shuffle.py
+++ b/distributed/shuffle/_shuffle.py
@@ -4,11 +4,11 @@ import logging
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Iterable, Iterator, NewType
 
-import dask
 from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import Layer
 
+from distributed.exceptions import Reschedule
 from distributed.shuffle._arrow import check_dtype_support, check_minimal_arrow_version
 
 logger = logging.getLogger("distributed.shuffle")
@@ -56,6 +56,7 @@ def shuffle_transfer(
     input_partition: int,
     npartitions: int,
     column: str,
+    parts_out: set[int],
 ) -> int:
     try:
         return _get_worker_extension().add_partition(
@@ -65,6 +66,7 @@ def shuffle_transfer(
             input_partition=input_partition,
             npartitions=npartitions,
             column=column,
+            parts_out=parts_out,
         )
     except Exception as e:
         raise RuntimeError(f"shuffle_transfer failed during shuffle {id}") from e
@@ -77,6 +79,8 @@ def shuffle_unpack(
         return _get_worker_extension().get_output_partition(
             id, barrier_run_id, output_partition
         )
+    except Reschedule as e:
+        raise e
     except Exception as e:
         raise RuntimeError(f"shuffle_unpack failed during shuffle {id}") from e
 
@@ -95,12 +99,6 @@ def rearrange_by_column_p2p(
 ) -> DataFrame:
     from dask.dataframe import DataFrame
 
-    if dask.config.get("optimization.fuse.active"):
-        raise RuntimeError(
-            "P2P shuffling requires the fuse optimization to be turned off. "
-            "Set the 'optimization.fuse.active' config to False to deactivate."
-        )
-
     check_dtype_support(df._meta)
     npartitions = npartitions or df.npartitions
     token = tokenize(df, column, npartitions)
@@ -145,8 +143,6 @@ class P2PShuffleLayer(Layer):
         annotations: dict | None = None,
     ):
         check_minimal_arrow_version()
-        annotations = annotations or {}
-        annotations.update({"shuffle": lambda key: key[1]})
         self.name = name
         self.column = column
         self.npartitions = npartitions
@@ -156,6 +152,8 @@ class P2PShuffleLayer(Layer):
         else:
             self.parts_out = set(range(self.npartitions))
         self.npartitions_input = npartitions_input
+        annotations = annotations or {}
+        annotations.update({"shuffle": lambda key: key[1]})
         super().__init__(annotations=annotations)
 
     def __repr__(self) -> str:
@@ -246,6 +244,7 @@ class P2PShuffleLayer(Layer):
                 i,
                 self.npartitions,
                 self.column,
+                self.parts_out,
             )
 
         dsk[_barrier_key] = (shuffle_barrier, token, transfer_keys)
diff --git a/distributed/shuffle/_worker_extension.py b/distributed/shuffle/_worker_extension.py
index 751d69e1..edf70df9 100644
--- a/distributed/shuffle/_worker_extension.py
+++ b/distributed/shuffle/_worker_extension.py
@@ -15,9 +15,11 @@ from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload
 
 import toolz
 
+from dask.context import thread_state
 from dask.utils import parse_bytes
 
 from distributed.core import PooledRPCCall
+from distributed.exceptions import Reschedule
 from distributed.protocol import to_serialize
 from distributed.shuffle._arrow import (
     convert_partition,
@@ -205,6 +207,22 @@ class ShuffleRun(Generic[T_transfer_shard_id, T_partition_id, T_partition_type])
     async def receive(self, data: list[tuple[T_transfer_shard_id, bytes]]) -> None:
         await self._receive(data)
 
+    async def _ensure_output_worker(self, i: T_partition_id, key: str) -> None:
+        assigned_worker = self._get_assigned_worker(i)
+
+        if assigned_worker != self.local_address:
+            result = await self.scheduler.shuffle_restrict_task(
+                id=self.id, run_id=self.run_id, key=key, worker=assigned_worker
+            )
+            if result["status"] == "error":
+                raise RuntimeError(result["message"])
+            assert result["status"] == "OK"
+            raise Reschedule()
+
+    @abc.abstractmethod
+    def _get_assigned_worker(self, i: T_partition_id) -> str:
+        """Get the address of the worker assigned to the output partition"""
+
     @abc.abstractmethod
     async def _receive(self, data: list[tuple[T_transfer_shard_id, bytes]]) -> None:
         """Receive shards belonging to output partitions of this shuffle run"""
@@ -216,7 +234,9 @@ class ShuffleRun(Generic[T_transfer_shard_id, T_partition_id, T_partition_type])
         """Add an input partition to the shuffle run"""
 
     @abc.abstractmethod
-    async def get_output_partition(self, i: T_partition_id) -> T_partition_type:
+    async def get_output_partition(
+        self, i: T_partition_id, key: str
+    ) -> T_partition_type:
         """Get an output partition to the shuffle run"""
 
 
@@ -354,16 +374,11 @@ class ArrayRechunkRun(ShuffleRun[ArrayRechunkShardID, NIndex, "np.ndarray"]):
         await self._write_to_comm(out)
         return self.run_id
 
-    async def get_output_partition(self, i: NIndex) -> np.ndarray:
+    async def get_output_partition(self, i: NIndex, key: str) -> np.ndarray:
         self.raise_if_closed()
         assert self.transferred, "`get_output_partition` called before barrier task"
 
-        assert self.worker_for[i] == self.local_address, (
-            f"Output partition {i} belongs on {self.worker_for[i]}, "
-            f"not {self.local_address}. "
-        )
-        # ^ NOTE: this check isn't necessary, just a nice validation to prevent incorrect
-        # data in the case something has gone very wrong
+        await self._ensure_output_worker(i, key)
 
         await self.flush_receive()
 
@@ -375,6 +390,9 @@ class ArrayRechunkRun(ShuffleRun[ArrayRechunkShardID, NIndex, "np.ndarray"]):
 
         return await self.offload(_)
 
+    def _get_assigned_worker(self, i: NIndex) -> str:
+        return self.worker_for[i]
+
 
 class DataFrameShuffleRun(ShuffleRun[int, int, "pd.DataFrame"]):
     """State for a single active shuffle execution
@@ -507,16 +525,11 @@ class DataFrameShuffleRun(ShuffleRun[int, int, "pd.DataFrame"]):
         await self._write_to_comm(out)
         return self.run_id
 
-    async def get_output_partition(self, i: int) -> pd.DataFrame:
+    async def get_output_partition(self, i: int, key: str) -> pd.DataFrame:
         self.raise_if_closed()
         assert self.transferred, "`get_output_partition` called before barrier task"
 
-        assert self.worker_for[i] == self.local_address, (
-            f"Output partition {i} belongs on {self.worker_for[i]}, "
-            f"not {self.local_address}. "
-        )
-        # ^ NOTE: this check isn't necessary, just a nice validation to prevent incorrect
-        # data in the case something has gone very wrong
+        await self._ensure_output_worker(i, key)
 
         await self.flush_receive()
         try:
@@ -531,6 +544,9 @@ class DataFrameShuffleRun(ShuffleRun[int, int, "pd.DataFrame"]):
             out = self.schema.empty_table().to_pandas()
         return out
 
+    def _get_assigned_worker(self, i: int) -> str:
+        return self.worker_for[i]
+
 
 class ShuffleWorkerExtension:
     """Interface between a Worker and a Shuffle.
@@ -763,6 +779,7 @@ class ShuffleWorkerExtension:
                     .to_pybytes(),
                     "npartitions": kwargs["npartitions"],
                     "column": kwargs["column"],
+                    "parts_out": kwargs["parts_out"],
                 },
                 worker=self.worker.address,
             )
@@ -776,7 +793,7 @@ class ShuffleWorkerExtension:
             )
         else:  # pragma: no cover
             raise TypeError(type)
-        if result["status"] == "ERROR":
+        if result["status"] == "error":
             raise RuntimeError(result["message"])
         assert result["status"] == "OK"
 
@@ -900,7 +917,10 @@ class ShuffleWorkerExtension:
         Calling this for a ``shuffle_id`` which is unknown or incomplete is an error.
         """
         shuffle = self.get_shuffle_run(shuffle_id, run_id)
-        return sync(self.worker.loop, shuffle.get_output_partition, output_partition)
+        key = thread_state.key
+        return sync(
+            self.worker.loop, shuffle.get_output_partition, output_partition, key
+        )
 
 
 def split_by_worker(
diff --git a/distributed/worker.py b/distributed/worker.py
index 9c6de71f..e6642d27 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -81,11 +81,13 @@ from distributed.core import send_recv
 from distributed.diagnostics import nvml, rmm
 from distributed.diagnostics.plugin import _get_plugin_name
 from distributed.diskutils import WorkSpace
+from distributed.exceptions import Reschedule
 from distributed.http import get_handlers
 from distributed.metrics import context_meter, thread_time, time
 from distributed.node import ServerNode
 from distributed.proctitle import setproctitle
 from distributed.protocol import pickle, to_serialize
+from distributed.protocol.serialize import _is_dumpable
 from distributed.pubsub import PubSubWorkerExtension
 from distributed.security import Security
 from distributed.shuffle import ShuffleWorkerExtension
@@ -957,6 +959,24 @@ class Worker(BaseWorker, ServerNode):
         return self._deque_handler.deque
 
     def log_event(self, topic: str | Collection[str], msg: Any) -> None:
+        """Log an event under a given topic
+
+        Parameters
+        ----------
+        topic : str, list[str]
+            Name of the topic under which to log an event. To log the same
+            event under multiple topics, pass a list of topic names.
+        msg
+            Event message to log. Note this must be msgpack serializable.
+
+        See also
+        --------
+        Client.log_event
+        """
+        if not _is_dumpable(msg):
+            raise TypeError(
+                f"Message must be msgpack serializable. Got {type(msg)=} instead."
+            )
         full_msg = {
             "op": "log-event",
             "topic": topic,
@@ -2771,20 +2791,6 @@ def secede():
     )
 
 
-class Reschedule(Exception):
-    """Reschedule this task
-
-    Raising this exception will stop the current execution of the task and ask
-    the scheduler to reschedule this task, possibly on a different machine.
-
-    This does not guarantee that the task will move onto a different machine.
-    The scheduler will proceed through its normal heuristics to determine the
-    optimal machine to accept this task.  The machine will likely change if the
-    load across the cluster has significantly changed since first scheduling
-    the task.
-    """
-
-
 @overload
 async def get_data_from_worker(
     rpc: ConnectionPool,


Test patch:
diff --git a/distributed/shuffle/tests/test_merge.py b/distributed/shuffle/tests/test_merge.py
index ee7824ce..4ea17c6d 100644
--- a/distributed/shuffle/tests/test_merge.py
+++ b/distributed/shuffle/tests/test_merge.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import pytest
 
 from distributed.shuffle._merge import hash_join
+from distributed.shuffle.tests.utils import invoke_annotation_chaos
 from distributed.utils_test import gen_cluster
 
 dd = pytest.importorskip("dask.dataframe")
@@ -13,6 +14,11 @@ from dask.dataframe.utils import assert_eq
 from dask.utils_test import hlg_layer_topological
 
 
+@pytest.fixture(params=[0, 0.3, 1], ids=["none", "some", "all"])
+def lose_annotations(request):
+    return request.param
+
+
 def list_eq(aa, bb):
     if isinstance(aa, dd.DataFrame):
         a = aa.compute(scheduler="sync")
@@ -36,7 +42,8 @@ def list_eq(aa, bb):
 
 @pytest.mark.parametrize("how", ["inner", "left", "right", "outer"])
 @gen_cluster(client=True)
-async def test_basic_merge(c, s, a, b, how):
+async def test_basic_merge(c, s, a, b, how, lose_annotations):
+    await invoke_annotation_chaos(lose_annotations, c)
     A = pd.DataFrame({"x": [1, 2, 3, 4, 5, 6], "y": [1, 1, 2, 2, 3, 4]})
     a = dd.repartition(A, [0, 4, 5])
 
@@ -72,7 +79,8 @@ async def test_basic_merge(c, s, a, b, how):
 
 @pytest.mark.parametrize("how", ["inner", "outer", "left", "right"])
 @gen_cluster(client=True)
-async def test_merge(c, s, a, b, how):
+async def test_merge(c, s, a, b, how, lose_annotations):
+    await invoke_annotation_chaos(lose_annotations, c)
     A = pd.DataFrame({"x": [1, 2, 3, 4, 5, 6], "y": [1, 1, 2, 2, 3, 4]})
     a = dd.repartition(A, [0, 4, 5])
 
diff --git a/distributed/shuffle/tests/test_rechunk.py b/distributed/shuffle/tests/test_rechunk.py
index d1f397b2..d9394145 100644
--- a/distributed/shuffle/tests/test_rechunk.py
+++ b/distributed/shuffle/tests/test_rechunk.py
@@ -18,7 +18,7 @@ from dask.array.utils import assert_eq
 
 from distributed.shuffle._limiter import ResourceLimiter
 from distributed.shuffle._rechunk import ShardID, rechunk_slicing
-from distributed.shuffle._scheduler_extension import get_worker_for_range_sharding
+from distributed.shuffle._scheduler_extension import get_worker_for_hash_sharding
 from distributed.shuffle._shuffle import ShuffleId
 from distributed.shuffle._worker_extension import ArrayRechunkRun
 from distributed.shuffle.tests.utils import AbstractShuffleTestPool
@@ -91,9 +91,7 @@ async def test_lowlevel_rechunk(
 
     new_indices = list(product(*(range(len(dim)) for dim in new)))
     for i, idx in enumerate(new_indices):
-        worker_for_mapping[idx] = get_worker_for_range_sharding(
-            i, workers, len(new_indices)
-        )
+        worker_for_mapping[idx] = get_worker_for_hash_sharding(i, workers)
 
     assert len(set(worker_for_mapping.values())) == min(n_workers, len(new_indices))
 
@@ -138,7 +136,7 @@ async def test_lowlevel_rechunk(
             all_chunks = np.empty(tuple(len(dim) for dim in new), dtype="O")
             for ix, worker in worker_for_mapping.items():
                 s = local_shuffle_pool.shuffles[worker]
-                all_chunks[ix] = await s.get_output_partition(ix)
+                all_chunks[ix] = await s.get_output_partition(ix, f"key-{ix}")
 
         finally:
             await asyncio.gather(*[s.close() for s in shuffles])
@@ -153,38 +151,9 @@ async def test_lowlevel_rechunk(
         )
 
 
-def test_raise_on_fuse_optimization():
-    a = np.random.uniform(0, 1, 30)
-    x = da.from_array(a, chunks=((10,) * 3,))
-    new = ((6,) * 5,)
-    with pytest.raises(RuntimeError, match="fuse optimization"):
-        rechunk(x, chunks=new, method="p2p")
-
-
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
-async def test_raise_on_lost_annotation(c, s, a, b):
-    a = np.random.uniform(0, 1, 30)
-    x = da.from_array(a, chunks=((10,) * 3,))
-    new = ((6,) * 5,)
-    x2 = rechunk(x, chunks=new, method="p2p")
-
-    # Manually drop "shuffle" annotation
-    for name, layer in x2.dask.layers.items():
-        if name.startswith("rechunk-p2p"):
-            del layer.annotations["shuffle"]
-
-    with raises_with_cause(
-        RuntimeError,
-        "rechunk_transfer failed",
-        RuntimeError,
-        "lost its ``shuffle`` annotation",
-    ):
-        await c.compute(x2)
-
-
 @pytest.mark.parametrize("config_value", ["tasks", "p2p", None])
 @pytest.mark.parametrize("keyword", ["tasks", "p2p", None])
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_configuration(c, s, *ws, config_value, keyword):
     """Try rechunking a random 1d matrix
 
@@ -208,7 +177,7 @@ async def test_rechunk_configuration(c, s, *ws, config_value, keyword):
     assert np.all(await c.compute(x2) == a)
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_2d(c, s, *ws):
     """Try rechunking a random 2d matrix
 
@@ -224,7 +193,7 @@ async def test_rechunk_2d(c, s, *ws):
     assert np.all(await c.compute(x2) == a)
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_4d(c, s, *ws):
     """Try rechunking a random 4d matrix
 
@@ -238,10 +207,15 @@ async def test_rechunk_4d(c, s, *ws):
     new = ((10,),) * 4
     x2 = rechunk(x, chunks=new, method="p2p")
     assert x2.chunks == new
-    assert np.all(await c.compute(x2) == a)
+    # FIXME: distributed#7816
+    with raises_with_cause(
+        RuntimeError, "rechunk_transfer failed", RuntimeError, "Barrier task"
+    ):
+        await c.compute(x2)
+    # assert np.all(await c.compute(x2) == a)
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_expand(c, s, *ws):
     """
     See Also
@@ -254,7 +228,7 @@ async def test_rechunk_expand(c, s, *ws):
     assert np.all(await c.compute(y) == a)
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_expand2(c, s, *ws):
     """
     See Also
@@ -274,7 +248,7 @@ async def test_rechunk_expand2(c, s, *ws):
             assert np.all(y == orig)
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_method(c, s, *ws):
     """Test rechunking can be done as a method of dask array.
 
@@ -291,7 +265,7 @@ async def test_rechunk_method(c, s, *ws):
     assert np.all(await c.compute(x2) == a)
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_blockshape(c, s, *ws):
     """Test that blockshape can be used.
 
@@ -309,7 +283,7 @@ async def test_rechunk_blockshape(c, s, *ws):
     assert np.all(await c.compute(check1) == a)
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_dtype(c, s, *ws):
     """
     See Also
@@ -320,7 +294,7 @@ async def test_dtype(c, s, *ws):
     assert x.rechunk(chunks=(1,), method="p2p").dtype == x.dtype
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_with_dict(c, s, *ws):
     """
     See Also
@@ -344,7 +318,7 @@ async def test_rechunk_with_dict(c, s, *ws):
     assert y.chunks == ((4, 4, 4, 4, 4, 4), (24,))
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_with_empty_input(c, s, *ws):
     """
     See Also
@@ -357,7 +331,7 @@ async def test_rechunk_with_empty_input(c, s, *ws):
         x.rechunk(chunks=(), method="p2p")
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_with_null_dimensions(c, s, *ws):
     """
     See Also
@@ -375,7 +349,7 @@ async def test_rechunk_with_null_dimensions(c, s, *ws):
     )
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_with_integer(c, s, *ws):
     """
     See Also
@@ -388,7 +362,7 @@ async def test_rechunk_with_integer(c, s, *ws):
     assert (await c.compute(x) == await c.compute(y)).all()
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_0d(c, s, *ws):
     """
     See Also
@@ -405,7 +379,7 @@ async def test_rechunk_0d(c, s, *ws):
 @pytest.mark.parametrize(
     "arr", [da.array([]), da.array([[], []]), da.array([[[]], [[]]])]
 )
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_empty_array(c, s, *ws, arr):
     """
     See Also
@@ -416,7 +390,7 @@ async def test_rechunk_empty_array(c, s, *ws, arr):
     assert arr.size == 0
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_empty(c, s, *ws):
     """
     See Also
@@ -429,7 +403,7 @@ async def test_rechunk_empty(c, s, *ws):
     assert_eq(await c.compute(x), await c.compute(y))
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_zero_dim_array(c, s, *ws):
     """
     See Also
@@ -442,7 +416,7 @@ async def test_rechunk_zero_dim_array(c, s, *ws):
     assert_eq(await c.compute(x), await c.compute(y))
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_zero_dim_array_II(c, s, *ws):
     """
     See Also
@@ -455,7 +429,7 @@ async def test_rechunk_zero_dim_array_II(c, s, *ws):
     assert_eq(await c.compute(x), await c.compute(y))
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_same(c, s, *ws):
     """
     See Also
@@ -467,7 +441,7 @@ async def test_rechunk_same(c, s, *ws):
     assert x is y
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_with_zero_placeholders(c, s, *ws):
     """
     See Also
@@ -480,7 +454,7 @@ async def test_rechunk_with_zero_placeholders(c, s, *ws):
     assert x.chunks == y.chunks
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_minus_one(c, s, *ws):
     """
     See Also
@@ -493,7 +467,7 @@ async def test_rechunk_minus_one(c, s, *ws):
     assert_eq(await c.compute(x), await c.compute(y))
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_warning(c, s, *ws):
     """
     See Also
@@ -507,7 +481,7 @@ async def test_rechunk_warning(c, s, *ws):
     assert not w
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_unknown_from_pandas(c, s, *ws):
     """
     See Also
@@ -529,7 +503,7 @@ async def test_rechunk_unknown_from_pandas(c, s, *ws):
     assert_eq(await c.compute(result), await c.compute(expected))
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_unknown_from_array(c, s, *ws):
     """
     See Also
@@ -562,7 +536,7 @@ async def test_rechunk_unknown_from_array(c, s, *ws):
         (da.ones(shape=(10, 10), chunks=(10, 2)), (None, (5, 5))),
     ],
 )
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_unknown(c, s, *ws, x, chunks):
     """
     See Also
@@ -578,7 +552,7 @@ async def test_rechunk_unknown(c, s, *ws, x, chunks):
     assert_eq(await c.compute(result), await c.compute(expected))
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_unknown_explicit(c, s, *ws):
     """
     See Also
@@ -602,7 +576,7 @@ def assert_chunks_match(left, right):
             assert x == y
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_unknown_raises(c, s, *ws):
     """
     See Also
@@ -616,7 +590,7 @@ async def test_rechunk_unknown_raises(c, s, *ws):
         x.rechunk((None, (5, 5, 5)), method="p2p")
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_zero_dim(c, s, *ws):
     """
     See Also
@@ -629,7 +603,7 @@ async def test_rechunk_zero_dim(c, s, *ws):
     assert len(await c.compute(x)) == 0
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_empty_chunks(c, s, *ws):
     """
     See Also
@@ -642,7 +616,7 @@ async def test_rechunk_empty_chunks(c, s, *ws):
 
 
 @pytest.mark.skip(reason="FIXME: We should avoid P2P in this case")
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_avoid_needless_chunking(c, s, *ws):
     x = da.ones(16, chunks=2)
     y = x.rechunk(8, method="p2p")
@@ -660,7 +634,7 @@ async def test_rechunk_avoid_needless_chunking(c, s, *ws):
         (20, (1, 1, 1, 1, 6, 2, 1, 7), 5, (5, 5, 5, 5)),
     ],
 )
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_auto_1d(c, s, *ws, shape, chunks, bs, expected):
     """
     See Also
@@ -672,7 +646,7 @@ async def test_rechunk_auto_1d(c, s, *ws, shape, chunks, bs, expected):
     assert y.chunks == (expected,)
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_auto_2d(c, s, *ws):
     """
     See Also
@@ -700,7 +674,7 @@ async def test_rechunk_auto_2d(c, s, *ws):
     assert y.chunks[0] == (4, 4, 4, 4, 4)  # limited by largest
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_auto_3d(c, s, *ws):
     """
     See Also
@@ -717,7 +691,7 @@ async def test_rechunk_auto_3d(c, s, *ws):
 
 
 @pytest.mark.parametrize("n", [100, 1000])
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_auto_image_stack(c, s, *ws, n):
     """
     See Also
@@ -743,7 +717,7 @@ async def test_rechunk_auto_image_stack(c, s, *ws, n):
         assert z.chunks == ((1,) * n, (362, 362, 276), (362, 362, 276))
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_down(c, s, *ws):
     """
     See Also
@@ -767,7 +741,7 @@ async def test_rechunk_down(c, s, *ws):
         assert z.chunks == ((10,) * 10, (104,) * 9 + (64,), (1000,))
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_zero(c, s, *ws):
     """
     See Also
@@ -780,7 +754,7 @@ async def test_rechunk_zero(c, s, *ws):
         assert y.chunks == ((1,) * 10,)
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_bad_keys(c, s, *ws):
     """
     See Also
@@ -811,7 +785,7 @@ async def test_rechunk_bad_keys(c, s, *ws):
     assert "-100" in str(info.value)
 
 
-@gen_cluster(client=True, config={"optimization.fuse.active": False})
+@gen_cluster(client=True)
 async def test_rechunk_with_zero(c, s, *ws):
     """
     See Also
diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
index d0534a64..b4672a8f 100644
--- a/distributed/shuffle/tests/test_shuffle.py
+++ b/distributed/shuffle/tests/test_shuffle.py
@@ -32,7 +32,7 @@ from distributed.shuffle._scheduler_extension import (
     ShuffleSchedulerExtension,
     get_worker_for_range_sharding,
 )
-from distributed.shuffle._shuffle import P2PShuffleLayer, ShuffleId, barrier_key
+from distributed.shuffle._shuffle import ShuffleId, barrier_key
 from distributed.shuffle._worker_extension import (
     DataFrameShuffleRun,
     ShuffleRun,
@@ -42,15 +42,12 @@ from distributed.shuffle._worker_extension import (
     split_by_partition,
     split_by_worker,
 )
-from distributed.shuffle.tests.utils import AbstractShuffleTestPool
-from distributed.utils import Deadline
-from distributed.utils_test import (
-    cluster,
-    gen_cluster,
-    gen_test,
-    raises_with_cause,
-    wait_for_state,
+from distributed.shuffle.tests.utils import (
+    AbstractShuffleTestPool,
+    invoke_annotation_chaos,
 )
+from distributed.utils import Deadline
+from distributed.utils_test import cluster, gen_cluster, gen_test, wait_for_state
 from distributed.worker_state_machine import TaskState as WorkerTaskState
 
 try:
@@ -59,6 +56,11 @@ except ImportError:
     pa = None
 
 
+@pytest.fixture(params=[0, 0.3, 1], ids=["none", "some", "all"])
+def lose_annotations(request):
+    return request.param
+
+
 async def clean_worker(
     worker: Worker, interval: float = 0.01, timeout: int | None = None
 ) -> None:
@@ -103,7 +105,8 @@ async def test_minimal_version(c, s, a, b):
 
 
 @gen_cluster(client=True)
-async def test_basic_integration(c, s, a, b):
+async def test_basic_integration(c, s, a, b, lose_annotations):
+    await invoke_annotation_chaos(lose_annotations, c)
     df = dask.datasets.timeseries(
         start="2000-01-01",
         end="2000-01-10",
@@ -121,46 +124,6 @@ async def test_basic_integration(c, s, a, b):
     await clean_scheduler(s)
 
 
-def test_raise_on_fuse_optimization():
-    df = dask.datasets.timeseries(
-        start="2000-01-01",
-        end="2000-01-10",
-        dtypes={"x": float, "y": float},
-        freq="10 s",
-    )
-    with dask.config.set({"optimization.fuse.active": True}):
-        with pytest.raises(RuntimeError, match="fuse optimization"):
-            dd.shuffle.shuffle(df, "x", shuffle="p2p")
-
-
-@gen_cluster(client=True)
-async def test_raise_on_lost_annotation(c, s, a, b):
-    df = dask.datasets.timeseries(
-        start="2000-01-01",
-        end="2000-01-10",
-        dtypes={"x": float, "y": float},
-        freq="10 s",
-    )
-    df = dd.shuffle.shuffle(df, "x", shuffle="p2p")
-
-    # Manually drop "shuffle" annotation
-    for layer in df.dask.layers.values():
-        if isinstance(layer, P2PShuffleLayer):
-            del layer.annotations["shuffle"]
-
-    with raises_with_cause(
-        RuntimeError,
-        "shuffle_transfer failed",
-        RuntimeError,
-        "lost its ``shuffle`` annotation",
-    ):
-        await c.compute(df)
-
-    await clean_worker(a)
-    await clean_worker(b)
-    await clean_scheduler(s)
-
-
 def test_shuffle_before_categorize(loop_in_thread):
     """Regression test for https://github.com/dask/distributed/issues/7615"""
     with cluster() as (s, [a, b]), Client(s["address"], loop=loop_in_thread) as c:
@@ -176,7 +139,8 @@ def test_shuffle_before_categorize(loop_in_thread):
 
 
 @gen_cluster(client=True)
-async def test_concurrent(c, s, a, b):
+async def test_concurrent(c, s, a, b, lose_annotations):
+    await invoke_annotation_chaos(lose_annotations, c)
     df = dask.datasets.timeseries(
         start="2000-01-01",
         end="2000-01-10",
@@ -231,7 +195,7 @@ async def wait_until_worker_has_tasks(
             [
                 key
                 for key, ts in scheduler.tasks.items()
-                if prefix in key and ts.state == "memory" and ws in ts.who_has
+                if prefix in key and ts.state == "memory" and {ws} == ts.who_has
             ]
         )
         < count
@@ -948,13 +912,17 @@ async def test_crashed_worker_after_shuffle(c, s, a):
         block_event = Event()
         with dask.annotate(workers=[n.worker_address], allow_other_workers=True):
             out = block(out, in_event, block_event)
-        fut = c.compute(out)
+        out = c.compute(out)
 
+        await wait_until_worker_has_tasks("shuffle-p2p", n.worker_address, 1, s)
         await in_event.wait()
         await n.process.process.kill()
-        block_event.set()
-        with pytest.raises(RuntimeError):
-            await fut
+        await block_event.set()
+
+        out = await out
+        result = out.x.size
+        expected = await c.compute(df.x.size)
+        assert result == expected
 
         await c.close()
         await clean_worker(a)
@@ -973,12 +941,16 @@ async def test_crashed_worker_after_shuffle_persisted(c, s, a):
         )
         out = dd.shuffle.shuffle(df, "x", shuffle="p2p")
         out = out.persist()
+
+        await wait_until_worker_has_tasks("shuffle-p2p", n.worker_address, 1, s)
         await out
 
         await n.process.process.kill()
 
-        with pytest.raises(RuntimeError):
-            await c.compute(out.sum())
+        result, expected = c.compute([out.x.size, df.x.size])
+        result = await result
+        expected = await expected
+        assert result == expected
 
         await c.close()
         await clean_worker(a)
@@ -1233,7 +1205,7 @@ async def test_basic_lowlevel_shuffle(
 
     for part in range(npartitions):
         worker_for_mapping[part] = get_worker_for_range_sharding(
-            part, workers, npartitions
+            npartitions, part, workers
         )
     assert len(set(worker_for_mapping.values())) == min(n_workers, npartitions)
     schema = pa.Schema.from_pandas(dfs[0])
@@ -1278,7 +1250,7 @@ async def test_basic_lowlevel_shuffle(
             all_parts = []
             for part, worker in worker_for_mapping.items():
                 s = local_shuffle_pool.shuffles[worker]
-                all_parts.append(s.get_output_partition(part))
+                all_parts.append(s.get_output_partition(part, f"key-{part}"))
 
             all_parts = await asyncio.gather(*all_parts)
 
@@ -1308,7 +1280,7 @@ async def test_error_offload(tmp_path, loop_in_thread):
 
     for part in range(npartitions):
         worker_for_mapping[part] = w = get_worker_for_range_sharding(
-            part, workers, npartitions
+            npartitions, part, workers
         )
         partitions_for_worker[w].append(part)
     schema = pa.Schema.from_pandas(dfs[0])
@@ -1362,7 +1334,7 @@ async def test_error_send(tmp_path, loop_in_thread):
 
     for part in range(npartitions):
         worker_for_mapping[part] = w = get_worker_for_range_sharding(
-            part, workers, npartitions
+            npartitions, part, workers
         )
         partitions_for_worker[w].append(part)
     schema = pa.Schema.from_pandas(dfs[0])
@@ -1415,7 +1387,7 @@ async def test_error_receive(tmp_path, loop_in_thread):
 
     for part in range(npartitions):
         worker_for_mapping[part] = w = get_worker_for_range_sharding(
-            part, workers, npartitions
+            npartitions, part, workers
         )
         partitions_for_worker[w].append(part)
     schema = pa.Schema.from_pandas(dfs[0])
diff --git a/distributed/shuffle/tests/test_shuffle_extension.py b/distributed/shuffle/tests/test_shuffle_extension.py
index 57447bd1..59462150 100644
--- a/distributed/shuffle/tests/test_shuffle_extension.py
+++ b/distributed/shuffle/tests/test_shuffle_extension.py
@@ -54,7 +54,7 @@ def test_split_by_worker():
     npartitions = 3
     for part in range(npartitions):
         worker_for_mapping[part] = get_worker_for_range_sharding(
-            part, workers, npartitions
+            npartitions, part, workers
         )
     worker_for = pd.Series(worker_for_mapping, name="_workers").astype("category")
     out = split_by_worker(df, "_partition", worker_for)
@@ -92,14 +92,14 @@ def test_split_by_worker_many_workers():
     worker_for_mapping = {}
     for part in range(npartitions):
         worker_for_mapping[part] = get_worker_for_range_sharding(
-            part, workers, npartitions
+            npartitions, part, workers
         )
     worker_for = pd.Series(worker_for_mapping, name="_workers").astype("category")
     out = split_by_worker(df, "_partition", worker_for)
-    assert get_worker_for_range_sharding(5, workers, npartitions) in out
-    assert get_worker_for_range_sharding(0, workers, npartitions) in out
-    assert get_worker_for_range_sharding(7, workers, npartitions) in out
-    assert get_worker_for_range_sharding(1, workers, npartitions) in out
+    assert get_worker_for_range_sharding(npartitions, 5, workers) in out
+    assert get_worker_for_range_sharding(npartitions, 0, workers) in out
+    assert get_worker_for_range_sharding(npartitions, 7, workers) in out
+    assert get_worker_for_range_sharding(npartitions, 1, workers) in out
 
     assert sum(map(len, out.values())) == len(df)
 
diff --git a/distributed/shuffle/tests/utils.py b/distributed/shuffle/tests/utils.py
index 85c0fe6e..8248592c 100644
--- a/distributed/shuffle/tests/utils.py
+++ b/distributed/shuffle/tests/utils.py
@@ -1,9 +1,13 @@
 from __future__ import annotations
 
 import itertools
+import random
 from typing import Any
 
+from distributed.client import Client
 from distributed.core import PooledRPCCall
+from distributed.diagnostics.plugin import SchedulerPlugin
+from distributed.scheduler import Scheduler, TaskStateState
 from distributed.shuffle._shuffle import ShuffleId
 from distributed.shuffle._worker_extension import ShuffleRun
 
@@ -43,3 +47,54 @@ class AbstractShuffleTestPool:
         for addr, s in self.shuffles.items():
             out[addr] = await s.inputs_done()
         return out
+
+
+class ShuffleAnnotationChaosPlugin(SchedulerPlugin):
+    #: Rate at which the plugin randomly drops shuffle annotations
+    rate: float
+    scheduler: Scheduler | None
+    seen: set
+
+    def __init__(self, rate: float):
+        self.rate = rate
+        self.scheduler = None
+        self.seen = set()
+
+    async def start(self, scheduler: Scheduler) -> None:
+        self.scheduler = scheduler
+
+    def transition(
+        self,
+        key: str,
+        start: TaskStateState,
+        finish: TaskStateState,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        assert self.scheduler
+        if finish != "waiting":
+            return
+        if not key.startswith("shuffle-barrier-"):
+            return
+        if key in self.seen:
+            return
+
+        self.seen.add(key)
+
+        barrier = self.scheduler.tasks[key]
+
+        if self._flip():
+            barrier.annotations.pop("shuffle", None)
+        for dt in barrier.dependents:
+            if self._flip():
+                dt.annotations.pop("shuffle", None)
+
+    def _flip(self) -> bool:
+        return random.random() < self.rate
+
+
+async def invoke_annotation_chaos(rate: float, client: Client) -> None:
+    if not rate:
+        return
+    plugin = ShuffleAnnotationChaosPlugin(rate)
+    await client.register_scheduler_plugin(plugin)
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
index 276f5ff6..4d42b8ee 100644
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -76,7 +76,7 @@ from distributed.client import (
 from distributed.cluster_dump import load_cluster_dump
 from distributed.comm import CommClosedError
 from distributed.compatibility import LINUX, WINDOWS
-from distributed.core import Status
+from distributed.core import Status, error_message
 from distributed.diagnostics.plugin import WorkerPlugin
 from distributed.metrics import time
 from distributed.scheduler import CollectTaskMetaDataPlugin, KilledWorker, Scheduler
@@ -7550,6 +7550,32 @@ async def test_log_event_warn(c, s, a, b):
         await c.submit(no_category)
 
 
+@gen_cluster(client=True, nthreads=[])
+async def test_log_event_msgpack(c, s, a, b):
+    await c.log_event("test-topic", "foo")
+    with pytest.raises(TypeError, match="msgpack"):
+
+        class C:
+            pass
+
+        await c.log_event("test-topic", C())
+    await c.log_event("test-topic", "bar")
+    await c.log_event("test-topic", error_message(Exception()))
+
+    # assertion reversed for mock.ANY.__eq__(Serialized())
+    assert [
+        "foo",
+        "bar",
+        {
+            "status": "error",
+            "exception": mock.ANY,
+            "traceback": mock.ANY,
+            "exception_text": "Exception()",
+            "traceback_text": "",
+        },
+    ] == [msg[1] for msg in s.get_events("test-topic")]
+
+
 @gen_cluster(client=True)
 async def test_log_event_warn_dask_warns(c, s, a, b):
     from dask.distributed import warn
diff --git a/distributed/tests/test_nanny.py b/distributed/tests/test_nanny.py
index 76b28a30..26719a34 100644
--- a/distributed/tests/test_nanny.py
+++ b/distributed/tests/test_nanny.py
@@ -21,7 +21,7 @@ from dask.utils import tmpfile
 
 from distributed import Nanny, Scheduler, Worker, profile, rpc, wait, worker
 from distributed.compatibility import LINUX, WINDOWS
-from distributed.core import CommClosedError, Status
+from distributed.core import CommClosedError, Status, error_message
 from distributed.diagnostics import SchedulerPlugin
 from distributed.diagnostics.plugin import WorkerPlugin
 from distributed.metrics import time
@@ -751,3 +751,33 @@ async def test_worker_inherits_temp_config(c, s):
         async with Nanny(s.address):
             out = await c.submit(lambda: dask.config.get("test123"))
             assert out == 123
+
+
+@gen_cluster(client=True, nthreads=[])
+async def test_log_event(c, s):
+    async with Nanny(s.address) as n:
+        n.log_event("test-topic1", "foo")
+
+        class C:
+            pass
+
+        with pytest.raises(TypeError, match="msgpack"):
+            n.log_event("test-topic2", C())
+        n.log_event("test-topic3", "bar")
+        n.log_event("test-topic4", error_message(Exception()))
+
+        # Worker unaffected
+        assert await c.submit(lambda x: x + 1, 1) == 2
+
+    assert [msg[1] for msg in s.get_events("test-topic1")] == ["foo"]
+    assert [msg[1] for msg in s.get_events("test-topic3")] == ["bar"]
+    # assertion reversed for mock.ANY.__eq__(Serialized())
+    assert [
+        {
+            "status": "error",
+            "exception": mock.ANY,
+            "traceback": mock.ANY,
+            "exception_text": "Exception()",
+            "traceback_text": "",
+        },
+    ] == [msg[1] for msg in s.get_events("test-topic4")]
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
index d40d7fb9..18a85bfd 100644
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -867,6 +867,39 @@ async def test_dont_overlap_communications_to_same_worker(c, s, a, b):
     assert l1["stop"] < l2["start"]
 
 
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_log_event(c, s, a):
+    def log_event(msg):
+        w = get_worker()
+        w.log_event("test-topic", msg)
+
+    await c.submit(log_event, "foo")
+
+    class C:
+        pass
+
+    with pytest.raises(TypeError, match="msgpack"):
+        await c.submit(log_event, C())
+
+    # Worker still works
+    await c.submit(log_event, "bar")
+    await c.submit(log_event, error_message(Exception()))
+
+    # assertion reversed for mock.ANY.__eq__(Serialized())
+    assert [
+        "foo",
+        "bar",
+        {
+            "status": "error",
+            "exception": mock.ANY,
+            "traceback": mock.ANY,
+            "exception_text": "Exception()",
+            "traceback_text": "",
+            "worker": a.address,
+        },
+    ] == [msg[1] for msg in s.get_events("test-topic")]
+
+
 @gen_cluster(client=True)
 async def test_log_exception_on_failed_task(c, s, a, b):
     with captured_logger("distributed.worker") as logger: