Reset state of ShuffleSchedulerExtension on restart

[hendrikmakait]

Fixes issue where forgotten shuffles could not be re-run after a cluster restart

• Tests added / passed
• Passes pre-commit run --all-files

[github-actions]

Unit Test Results

See test report for an extended history of previous test failures. This is useful for diagnosing flaky tests.

       22 files  ±  0         22 suites  ±0   10h 20m 46s ⏱️ + 50m 45s
  3 288 tests +  1    3 202 ✔️ +1       85 💤 ±0  1 ❌ ±0 
36 099 runs  +11  34 536 ✔️ +7  1 562 💤 +4  1 ❌ ±0 

For more details on these failures, see this check.

Results for commit cd8ddff. ± Comparison against base commit b5a2078.

♻️ This comment has been updated with latest results.

[fjetter]

Hold on, this is concerning. We cannot rerun the same shuffle even after it finished successfully?

[hendrikmakait]

Yes, due to the get-or-create style of shuffle_get, cancelled tasks could otherwise re-create a forgotten shuffle which would leave stale state indefinitely on the cluster. We should be able to solve this with something like #7372.

[hendrikmakait]

See also

distributed/distributed/shuffle/tests/test_shuffle.py

Lines 892 to 918 in 401b51d

	@pytest.mark.xfail(reason="Tombstone prohibits multiple calls to head")
	@gen_cluster(client=True, nthreads=[("127.0.0.1", 4)] * 2)
	async def test_repeat(c, s, a, b):
	df = dask.datasets.timeseries(
	start="2000-01-01",
	end="2000-01-10",
	dtypes={"x": float, "y": float},
	freq="100 s",
	)
	out = dd.shuffle.shuffle(df, "x", shuffle="p2p")
	await c.compute(out.head(compute=False))
	await clean_worker(a, timeout=2)
	await clean_worker(b, timeout=2)
	await clean_scheduler(s, timeout=2)
	await c.compute(out.tail(compute=False))
	await clean_worker(a, timeout=2)
	await clean_worker(b, timeout=2)
	await clean_scheduler(s, timeout=2)
	await c.compute(out.head(compute=False))
	await clean_worker(a, timeout=2)
	await clean_worker(b, timeout=2)
	await clean_scheduler(s, timeout=2)

[fjetter]

xref #7452