Add drain option to dask-worker

[AlecThomson]

Closes #3141

• Tests added / passed
• 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]

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 75e98c3. ± Comparison against base commit 5589049.

♻️ 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

[jacobtomlinson]

I like the idea of having a flag you can set which allows workers to complete current tasks without accepting new ones. It's probably most useful when tasks a very large, which is not best practice, but common enough that we should make QoL improvements for those users.

I'm not sure on the drain terminology. In SLURM it means the same thing as it means here, finish existing work without accepting new work. However, in Kubernetes terminology it means evict all current work and reschedule tasks on another worker. So for Dask users in the Kubernetes community this may be confusing.

Perhaps this option could be made mode explicit like --complete-tasks-on-shutdown. Although that's pretty long, so if you can think of something better than that's great!

[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!