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.

    29 files  ±0      29 suites  ±0   11h 17m 37s ⏱️ + 3m 35s
 4 083 tests ±0   3 963 ✅  - 2    112 💤 ±0   8 ❌ +2 
55 230 runs  ±0  52 788 ✅  - 1  2 432 💤  - 1  10 ❌ +2 

For more details on these failures, see this check.

Results for commit 226c8b3. ± Comparison against base commit 8564dc7.

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