[dask] make Dask training resilient to worker restarts during network setup

[jameslamb]

Summary

LightGBM distributed training requires that all workers participating in training know about each other (https://lightgbm.readthedocs.io/en/latest/Parallel-Learning-Guide.html#preparation). This information is given via the parameter machines.

LightGBM's Dask module generates a list of workers by checking where pieces of the input data are stored.

LightGBM/python-package/lightgbm/dask.py

Line 192 in 706f2af

who_has = client.who_has(parts)

It then uses the addresses in that list to set up the machines parameters. This step can take a few seconds to a few minutes (#3766 (comment)), and if one of the workers is restarted in between when that list is generated and when training starts, you can get a hard-to-understand KeyError like this:

KeyError: 'tcp://10.10.0.20:44284'

Improving this experience would mean one or both of these:

• refactoring python-package/lightgbm/dask.py to reduce the risk of this, maybe by adding error handling for the case where workers are restarted
• detecting this condition and adding a more informative error message

Motivation

Adding this feature would make the Dask interface a bit more stable, reducing the need for users to retry training or try to debug low-level Dask details.

Notes for Reviewers

I added "during network setup" to the title here, because I think that making it possible for LightGBM training to continue if a worker disappears during training is a much bigger task and not limited to Dask. @guolinke, I'd love your thoughts on this when you have time.

References

This issue was originally posted as dask/dask-lightgbm#24, and more background is available there. Moving it here as part of dask/community#104.

[jameslamb]

Closing this, as we use #2302 to track feature requests. Leave a comment below if you'd like to contribute this feature, and we'll be happy to re-open it!