[train] Updates to support xgboost==2.1.0

python/ray/train/xgboost/config.py

@@ -1,9 +1,13 @@
 import json
 import logging
 import os
+import threading
 from contextlib import contextmanager
 from dataclasses import dataclass
+from typing import Optional
 
+import xgboost
+from packaging.version import Version
 from xgboost import RabitTracker
 from xgboost.collective import CommunicatorContext
 
@@ -37,7 +41,7 @@ class XGBoostConfig(BackendConfig):
     def train_func_context(self):
         @contextmanager
         def collective_communication_context():
-            with CommunicatorContext():
+            with CommunicatorContext(**_get_xgboost_args()):
                 yield
 
         return collective_communication_context
@@ -52,13 +56,10 @@ def backend_cls(self):
 
 class _XGBoostRabitBackend(Backend):
     def __init__(self):
-        self._tracker = None
-
-    def on_training_start(
-        self, worker_group: WorkerGroup, backend_config: XGBoostConfig
-    ):
-        assert backend_config.xgboost_communicator == "rabit"
+        self._tracker: Optional[RabitTracker] = None
+        self._wait_thread: Optional[threading.Thread] = None
 
+    def _setup_xgboost_less_than_210(self, worker_group: WorkerGroup):
         # Set up the rabit tracker on the Train driver.
         num_workers = len(worker_group)
         rabit_args = {"DMLC_NUM_WORKER": num_workers}
@@ -67,12 +68,14 @@ def on_training_start(
         # NOTE: sortby="task" is needed to ensure that the xgboost worker ranks
         # align with Ray Train worker ranks.
         # The worker ranks will be sorted by `DMLC_TASK_ID`,
-        # which is defined in `on_training_start`.
+        # which is defined below.
         self._tracker = RabitTracker(
-            host_ip=train_driver_ip, n_workers=num_workers, sortby="task"
+            n_workers=num_workers, host_ip=train_driver_ip, sortby="task"
         )
-        rabit_args.update(self._tracker.worker_envs())
-        self._tracker.start(num_workers)
+        self._tracker.start(n_workers=num_workers)
+
+        worker_args = self._tracker.worker_envs()
+        rabit_args.update(worker_args)
 
         start_log = (
             "RabitTracker coordinator started with parameters:\n"
@@ -95,13 +98,80 @@ def set_xgboost_env_vars():
 
         worker_group.execute(set_xgboost_env_vars)
 
+    def _setup_xgboost(self, worker_group: WorkerGroup):
+        # Set up the rabit tracker on the Train driver.
+        num_workers = len(worker_group)
+        rabit_args = {"n_workers": num_workers}
+        train_driver_ip = ray.util.get_node_ip_address()
+
+        # NOTE: sortby="task" is needed to ensure that the xgboost worker ranks
+        # align with Ray Train worker ranks.
+        # The worker ranks will be sorted by `dmlc_task_id`,
+        # which is defined below.
+        self._tracker = RabitTracker(
+            n_workers=num_workers, host_ip=train_driver_ip, sortby="task"
+        )
+        self._tracker.start()
+
+        # The RabitTracker is started in a separate thread, and the
+        # `wait_for` method must be called for `worker_args` to return.
+        self._wait_thread = threading.Thread(target=self._tracker.wait_for, daemon=True)
+        self._wait_thread.start()
+
+        rabit_args.update(self._tracker.worker_args())
+
+        start_log = (
+            "RabitTracker coordinator started with parameters:\n"
+            f"{json.dumps(rabit_args, indent=2)}"
+        )
+        logger.debug(start_log)
+
+        def set_xgboost_communicator_args(args):
+            import ray.train
+
+            args["dmlc_task_id"] = (
+                f"[xgboost.ray-rank={ray.train.get_context().get_world_rank():08}]:"
+                f"{ray.get_runtime_context().get_actor_id()}"
+            )
+
+            _set_xgboost_args(args)
+
+        worker_group.execute(set_xgboost_communicator_args, rabit_args)
+
+    def on_training_start(
+        self, worker_group: WorkerGroup, backend_config: XGBoostConfig
+    ):
+        assert backend_config.xgboost_communicator == "rabit"
+
+        if Version(xgboost.__version__) < Version("2.1.0"):
+            self._setup_xgboost_less_than_210(worker_group)
+        else:
+            self._setup_xgboost(worker_group)
+
     def on_shutdown(self, worker_group: WorkerGroup, backend_config: XGBoostConfig):
         timeout = 5
-        self._tracker.thread.join(timeout=timeout)
 
-        if self._tracker.thread.is_alive():
-            logger.warning(
-                "During shutdown, the RabitTracker thread failed to join "
-                f"within {timeout} seconds. "
-                "The process will still be terminated as part of Ray actor cleanup."
-            )
+        if self._wait_thread is not None:
+            self._wait_thread.join(timeout=timeout)
+
+            if self._wait_thread.is_alive():
+                logger.warning(
+                    "During shutdown, the RabitTracker thread failed to join "
+                    f"within {timeout} seconds. "
+                    "The process will still be terminated as part of Ray actor cleanup."
+                )
+
+
+_xgboost_args: dict = {}
+_xgboost_args_lock = threading.Lock()
+
+
+def _set_xgboost_args(args):
+    with _xgboost_args_lock:
+        global _xgboost_args
+        _xgboost_args = args
+
+
+def _get_xgboost_args() -> dict:
+    with _xgboost_args_lock:
+        return _xgboost_args

python/requirements/ml/core-requirements.txt

@@ -4,7 +4,7 @@ mlflow==2.9.2
 wandb==0.17.0
 
 # ML training frameworks
-xgboost==1.7.6
+xgboost==2.1.0
 lightgbm==3.3.5
 
 # Huggingface