Add `drain` option to dask-worker

[AlecThomson]

Closes #3141

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

Need some help on adding appropriate tests for this option

[GPUtester]

Can one of the admins verify this patch?

Admins can comment ok to test to allow this one PR to run or add to allowlist to allow all future PRs from the same author to run.

[github-actions[bot]]

Unit Test Results

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

    27 files  ±0      27 suites  ±0   11h 29m 34s ⏱️ - 10m 9s  4 116 tests ±0   4 003 ✅ ±0    111 💤 ±0  2 ❌ ±0  51 615 runs  ±0  49 317 ✅ +1  2 296 💤  - 1  2 ❌ ±0 

For more details on these failures, see this check.

Results for commit 75e98c38. ± Comparison against base commit 55890493.

:recycle: This comment has been updated with latest results.

[AlecThomson]

Found that drain was not blocking additional tasks from being scheduled. Fixed with additional call to self.scheduler.retire_workers

[fsteinmetz]

I have successfully tested this patch on a sample case where I run 2-min jobs with a 1-min lifetime : tasks complete as expected when activating --drain. This feature is important in our group since we have to deal with sometimes relatively long tasks, where killing workers could lead to a significantly reduced efficiency. I hope it can be merged soon!