Reset state of ShuffleSchedulerExtension on restart (#7446)

distributed/shuffle/_scheduler_extension.py

@@ -227,6 +227,12 @@ def _clean_on_scheduler(self, id: ShuffleId) -> None:
         with contextlib.suppress(KeyError):
             del self.heartbeats[id]
 
+    def restart(self, scheduler: Scheduler) -> None:
+        self.states.clear()
+        self.heartbeats.clear()
+        self.tombstones.clear()
+        self.erred_shuffles.clear()
+
 
 def get_worker_for(output_partition: int, workers: list[str], npartitions: int) -> str:
     "Get the address of the worker which should hold this output partition number"

distributed/shuffle/tests/test_shuffle.py

@@ -1010,6 +1010,41 @@ async def test_closed_worker_between_repeats(c, s, w1, w2, w3):
     await clean_scheduler(s)
 
 
+@pytest.mark.slow
+@gen_cluster(client=True, nthreads=[("", 1)], Worker=Nanny)
+async def test_restart_cluster_between_repeats(c, s, a):
+    scheduler_extension = s.extensions["shuffle"]
+    df = dask.datasets.timeseries(
+        start="2000-01-01",
+        end="2000-01-10",
+        dtypes={"x": float, "y": float},
+        freq="100 s",
+        seed=42,
+    )
+
+    await c.compute(dd.shuffle.shuffle(df, "y", shuffle="p2p"))
+    await clean_scheduler(s)
+
+    assert scheduler_extension.tombstones
+
+    # Cannot rerun forgotten shuffle due to tombstone
+    with pytest.raises(RuntimeError, match="shuffle_transfer"):
+        await c.compute(dd.shuffle.shuffle(df, "y", shuffle="p2p"))
+
+    out = dd.shuffle.shuffle(df, "x", shuffle="p2p")
+    await c.compute(out)
+    assert scheduler_extension.shuffle_ids()
+
+    await c.restart()
+    await clean_scheduler(s)
+    assert not scheduler_extension.tombstones
+
+    await c.compute(out)
+    await c.compute(dd.shuffle.shuffle(df, "y", shuffle="p2p"))
+    del out
+    await clean_scheduler(s)
+
+
 @gen_cluster(client=True)
 async def test_new_worker(c, s, a, b):
     df = dask.datasets.timeseries(