[train] Updates to support xgboost==2.1.0

[justinvyu]

Why are these changes needed?

xgboost 2.1.0 was recently released, and it changed some of the distributed setup APIs.

In particular:

• CollectiveCommunicator was changed to not read from environment variables. Instead, an argument dict is required to be passed in: dmlc/xgboost@a5a5810#diff-a74bc610352aa00eda4ae89c1f3a51c33b934f50f055f366def49654efb42992
• The RabitTracker API was updated with a few renamed methods and changed behavior:
  • worker_envs() -> worker_args()
  • RabitTracker.wait_for must be run as a separate thread in order for worker_args to return properly: dmlc/xgboost@a5a5810#diff-94301e6ca68aefc564a0d617db4ab2de3425b2cca9d66fe95ad8c7ce97399c14R182-R186

This PR branches the setup logic between pre 2.1.0 and post 2.1.0. We should eventually drop pre-2.1.0 support.

Testing

This PR also updates the tested xgboost version to 2.1.0. Pre-2.1.0 has been tested manually.

Related issue number

Closes #46476

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[hongpeng-guo]

Thanks for the implementation! Two questions from me:

1. The RabitTracker class functions worker_args and worker_env resturn the same type of things Dict[str, Union[int, str]]. The only difference is the key of worker_env are uppercase letters but worker_args are lowercase letters. Our adaption to this change is to move from env settup to training context settup, is that correct?
2. The RabitTracker class doesn't maintain a thread itself, instead we need to create a main thread kind of thing using its wait_for method to wait for tracker.start by ourselves. My question is: Do we also need to distinguish xgboost version before/post 210 in the on_shutdown method. My first intuition is to using tracker.thread before 210, and using wait_for after 210. Currently, It seems we always use the wai_for method now.

[justinvyu]

The RabitTracker class functions worker_args and worker_env resturn the same type of things Dict[str, Union[int, str]]. The only difference is the key of worker_env are uppercase letters but worker_args are lowercase letters. Our adaption to this change is to move from env settup to training context settup, is that correct?

Yes. The API changed from accepting environment variables to only allowing you to pass the arguments directly as kwargs with those lower-case names.

[hongpeng-guo]

nit: it seems XGBoostingConfig has a hard coded backend_config field being "rabit", why do we still need an assertion here?

[justinvyu]

yeah, I can probably remove this field for now, since we don't support the "federated" option.

[hongpeng-guo]

Left some comments, most of them are nits and should be non-blocking.
It should be good to go if all unit tests look good.

[woshiyyya]

Are we able to save the xgboost_args into XGBoost config so we can avoid modifying the global variable?

[justinvyu]

Hmm, interesting. I actually don't understand why we need both BackendConfig and Backend classes. Any context here @matthewdeng ?

[matthewdeng]

The BackendConfig is the public API that the user could interact with. There is probably a better/cleaner way to organize the two.

[woshiyyya]

Yeah currently the dependency between BackendConfig and Backend are unidirectional. It's kind of hard to pass information from Backed -> BackendConfig.

[matthewdeng]

Should train_func_context be part of the Backend instead?

[matthewdeng]

Or at the very least the default one.

[matthewdeng]

Oh hm maybe that won't work because we construct the train loop before the backend...

[woshiyyya]

lgtm!

[hongpeng-guo]

LGTM!