Fix P2P worker cleanup

[hendrikmakait]

Worker cleanup currently works "by accident" because shuffles either complete or worker state gets cleaned up on failure. However, ShuffleWorkerExtension.close() never gets triggered. This PR fixes that. It also refactors the setup such that the scheduler plugin is responsible of installing the worker plugin.

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

[hendrikmakait]

cc @fjetter

[github-actions]

Unit Test Results

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

       20 files  ±0         20 suites  ±0   10h 49m 46s ⏱️ - 24m 29s
  3 716 tests ±0    3 607 ✔️ ±0     106 💤 ±0  3 ❌ ±0 
35 942 runs  ±0  34 189 ✔️ ±0  1 749 💤  - 1  4 ❌ +1 

For more details on these failures, see this check.

Results for commit 74599ee. ± Comparison against base commit 7b21399.

This pull request removes 7 and adds 7 tests. Note that renamed tests count towards both.

distributed.shuffle.tests.test_shuffle_extension
distributed.shuffle.tests.test_shuffle_extension ‑ test_installation_on_scheduler
distributed.shuffle.tests.test_shuffle_extension ‑ test_installation_on_worker
distributed.shuffle.tests.test_shuffle_extension ‑ test_split_by_partition
distributed.shuffle.tests.test_shuffle_extension ‑ test_split_by_worker
distributed.shuffle.tests.test_shuffle_extension ‑ test_split_by_worker_empty
distributed.shuffle.tests.test_shuffle_extension ‑ test_split_by_worker_many_workers

distributed.shuffle.tests.test_shuffle_plugins
distributed.shuffle.tests.test_shuffle_plugins ‑ test_installation_on_scheduler
distributed.shuffle.tests.test_shuffle_plugins ‑ test_installation_on_worker
distributed.shuffle.tests.test_shuffle_plugins ‑ test_split_by_partition
distributed.shuffle.tests.test_shuffle_plugins ‑ test_split_by_worker
distributed.shuffle.tests.test_shuffle_plugins ‑ test_split_by_worker_empty
distributed.shuffle.tests.test_shuffle_plugins ‑ test_split_by_worker_many_workers

This pull request removes 1 skipped test and adds 1 skipped test. Note that renamed tests count towards both.

distributed.shuffle.tests.test_shuffle_extension

distributed.shuffle.tests.test_shuffle_plugins

♻️ This comment has been updated with latest results.

[wence-]

I think this looks right, although minor linting changes required.

[wence-]

So this is the magic that ensures things are cleaned up, because the client will teardown the scheduler which tears down all the worker plugins?

[wence-]

And pretty much everything else is downstream renaming changes....

[hendrikmakait]

Pretty much, yes. IIUC, extensions don't get cleaned up at all. Worker plugins will get torn down on worker close, but for that the method has to be called teardown (as seen below). Similarly, we need to move some of the initialization to setup.

[wence-]

Along with this so that the worker plugin now conforms to the WorkerPlugin interface.

[hendrikmakait]

although minor linting changes required

I'm not sure if I missed any additional linting problems, but I think the ones on here also happen on main (sigh).

[wence-]

I'm not sure if I missed any additional linting problems, but I think the ones on here also happen on main (sigh).

#7985