Refactor restart() and restart_workers()

[crusaderky]

Refactor Client.restart() and Client.restart_workers() to use the same machinery for worker restart.

This is propaedeutic to #8537, where the scheduler will need to call restart_workers() internally.

• Blocks Restart workers when worker-ttl expires #8538

Changed behaviour

• restart() will no longer return self. It's weird and it was never advertised in the documentation anyway.
• in restart_workers(), the scheduler will now unilaterally remove workers where the nanny failed to terminate the worker within timeout seconds. This aligns its behaviour to restart().
• default timeout for restart_workers() has changed from infinity to distributed.comm.timeouts.connect * 4 (2 minutes by default). This aligns its behaviour to restart(). The method no longer accepts timeout=None.
• restart_workers() no longer forwards the exception to the client if the worker fails to restart. While this is, strictly speaking, a loss of functionality, IMHO I can think of very few of use cases when this can happen, none of which sound particularly interesting to me. Note that this is specifically a use case where the nanny is healthy and in the middle of an RPC call with the scheduler, while the worker is so busted that it fails to restart within the timeout. The only semi-realistic use case I can think of is a WorkerPlugin that fails to acquire an external resource (e.g. connect to a database) on startup.

[crusaderky]

This covers a race condition where Worker.close(nanny=True) ran successfully up until and including the message to the nanny, and then hanged or crashed - for example because of a bad plugin or extension teardown.

[github-actions]

Unit Test Results

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

    29 files  ± 0      29 suites  ±0   11h 28m 31s ⏱️ + 16m 19s
 4 056 tests + 1   3 934 ✅  -  1    109 💤 ±0  13 ❌ +2 
54 902 runs  +32  52 408 ✅ +36  2 409 💤  - 6  85 ❌ +2 

For more details on these failures, see this check.

Results for commit 3d3140f. ± Comparison against base commit 8927bfd.

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

distributed.tests.test_client ‑ test_restart_workers_exception[False]
distributed.tests.test_client ‑ test_restart_workers_exception[True]
distributed.tests.test_client ‑ test_restart_workers_timeout[False]
distributed.tests.test_client ‑ test_restart_workers_timeout[True]

distributed.tests.test_client ‑ test_restart_workers_exception
distributed.tests.test_client ‑ test_restart_workers_kill_timeout[False]
distributed.tests.test_client ‑ test_restart_workers_kill_timeout[True]
distributed.tests.test_client ‑ test_restart_workers_restart_timeout[False]
distributed.tests.test_client ‑ test_restart_workers_restart_timeout[True]

♻️ This comment has been updated with latest results.

[crusaderky]

There are so many places in the client code that just blindly assume it's not None. The json schema does not allow for None either.

[crusaderky]

All test failures are unrelated

[hendrikmakait]

restart_workers() no longer forwards the exception to the client if the worker fails to restart. While this is, strictly speaking, a loss of functionality, IMHO I can think of very few of use cases when this can happen, none of which sound particularly interesting to me. Note that this is specifically a use case where the nanny is healthy and in the middle of an RPC call with the scheduler, while the worker is so busted that it fails to restart within the timeout. The only semi-realistic use case I can think of is a WorkerPlugin that fails to acquire an external resource (e.g. connect to a database) on startup.

I suppose another realistic scenario would be a task running on the worker that repeatedly blocks the GIL for an extended period of time? Granted, there's probably not much added value to a simple timed out message in that case.

[hendrikmakait]

Thanks, @crusaderky! I love to see their behavior finally aligned. A few nits around code style/typing; feel free to ignore.

[hendrikmakait]

nit:

Suggested change

	):
	) -> dict[str, Literal["OK" , "removed", "timed out"]]:

[crusaderky]

No. We return Awaitable[dict] | dict depending if self.sync is true or false, and functions should never return Union (python/typing#566)

[hendrikmakait]

Right, thanks for the clarification.

[hendrikmakait]

nit: While we align functionality, would it make sense to return {worker address: "OK", "no nanny", or "timed out" or error message} here as well?

[crusaderky]

it makes sense but I'd rather treat it as material for a follow-up PR