[dask] Fix running on k8s.

[trivialfis]

Close #5765

• Use socket from Python.
• Only use workers associated with input data instead of all workers from client. This avoids accessing client.scheduler_info()['workers'].
• Avoid calling client.gather inside task function.

XGBoost still cannot be run on GKE yet. See dask/dask#6800 for details.

[jameslamb]

Hey @trivialfis , I'm really excited to see this! I have easy access to Dask clusters on EKS, so if you want me to test anything there I'd be happy to.

[trivialfis]

@jameslamb Thanks for the offer! So far I managed to get some examples working on GKE. But automated pytest is not on the table since everything just times out while waiting for scheduling. Feel free to test it out, it might help uncovering some other unknown issues on xgboost and dask and be useful to others. Also the issue is not reproducible on local deployment, so they are cluster specific. Your tests are definitely gonna be helpful!

I'm not sure about the performance impact and memory usage impact of current workaround, should be trivial even if there's any but I need to try it out.

[trivialfis]

A refactoring is coming.