chore: import upstream snapshot with attribution

python/ray/train/xgboost/__init__.py

@@ -0,0 +1,19 @@
+from ray.train.v2._internal.constants import is_v2_enabled
+from ray.train.xgboost._xgboost_utils import RayTrainReportCallback
+from ray.train.xgboost.config import XGBoostConfig
+from ray.train.xgboost.xgboost_checkpoint import XGBoostCheckpoint
+from ray.train.xgboost.xgboost_trainer import XGBoostTrainer
+
+if is_v2_enabled():
+    from ray.train.v2.xgboost.config import XGBoostConfig  # noqa: F811
+    from ray.train.v2.xgboost.xgboost_trainer import XGBoostTrainer  # noqa: F811
+
+__all__ = [
+    "RayTrainReportCallback",
+    "XGBoostCheckpoint",
+    "XGBoostConfig",
+    "XGBoostTrainer",
+]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/xgboost/_xgboost_utils.py

@@ -0,0 +1,251 @@
+import tempfile
+from abc import abstractmethod
+from collections import OrderedDict
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Callable, Dict, List, Optional, Union
+
+from xgboost.core import Booster
+
+import ray.train
+from ray.train import Checkpoint
+from ray.tune.utils import flatten_dict
+from ray.util.annotations import PublicAPI
+
+try:
+    from xgboost.callback import TrainingCallback
+except ImportError:
+
+    class TrainingCallback:
+        pass
+
+
+class RayReportCallback(TrainingCallback):
+    CHECKPOINT_NAME = "model.ubj"
+
+    def __init__(
+        self,
+        metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
+        filename: str = CHECKPOINT_NAME,
+        frequency: int = 0,
+        checkpoint_at_end: bool = True,
+        results_postprocessing_fn: Optional[
+            Callable[[Dict[str, Union[float, List[float]]]], Dict[str, float]]
+        ] = None,
+    ):
+        if isinstance(metrics, str):
+            metrics = [metrics]
+        self._metrics = metrics
+        self._filename = filename
+        self._frequency = frequency
+        self._checkpoint_at_end = checkpoint_at_end
+        self._results_postprocessing_fn = results_postprocessing_fn
+
+        # Keeps track of the eval metrics from the last iteration,
+        # so that the latest metrics can be reported with the checkpoint
+        # at the end of training.
+        self._evals_log = None
+        # Keep track of the last checkpoint iteration to avoid double-checkpointing
+        # when using `checkpoint_at_end=True`.
+        self._last_checkpoint_iteration = None
+
+    @classmethod
+    def get_model(
+        cls,
+        checkpoint: Checkpoint,
+        filename: str = CHECKPOINT_NAME,
+    ) -> Booster:
+        """Retrieve the model stored in a checkpoint reported by this callback.
+
+        Args:
+            checkpoint: The checkpoint object returned by a training run.
+                The checkpoint should be saved by an instance of this callback.
+            filename: The filename to load the model from, which should match
+                the filename used when creating the callback.
+
+        Returns:
+            The model loaded from the checkpoint.
+        """
+        with checkpoint.as_directory() as checkpoint_path:
+            booster = Booster()
+            booster.load_model(Path(checkpoint_path, filename).as_posix())
+            return booster
+
+    def _get_report_dict(self, evals_log):
+        if isinstance(evals_log, OrderedDict):
+            # xgboost>=1.3
+            result_dict = flatten_dict(evals_log, delimiter="-")
+            for k in list(result_dict):
+                result_dict[k] = result_dict[k][-1]
+        else:
+            # xgboost<1.3
+            result_dict = dict(evals_log)
+        if not self._metrics:
+            report_dict = result_dict
+        else:
+            report_dict = {}
+            for key in self._metrics:
+                if isinstance(self._metrics, dict):
+                    metric = self._metrics[key]
+                else:
+                    metric = key
+                report_dict[key] = result_dict[metric]
+
+        if self._results_postprocessing_fn:
+            report_dict = self._results_postprocessing_fn(report_dict)
+
+        return report_dict
+
+    @abstractmethod
+    def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]:
+        """Get checkpoint from model.
+
+        This method needs to be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster):
+        """Save checkpoint and report metrics corresonding to this checkpoint.
+
+        This method needs to be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _report_metrics(self, report_dict: Dict):
+        """Report Metrics.
+
+        This method needs to be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    def after_iteration(self, model: Booster, epoch: int, evals_log: Dict):
+        self._evals_log = evals_log
+
+        checkpointing_disabled = self._frequency == 0
+        # Ex: if frequency=2, checkpoint at epoch 1, 3, 5, ... (counting from 0)
+        should_checkpoint = (
+            not checkpointing_disabled and (epoch + 1) % self._frequency == 0
+        )
+
+        report_dict = self._get_report_dict(evals_log)
+        if should_checkpoint:
+            self._last_checkpoint_iteration = epoch
+            self._save_and_report_checkpoint(report_dict, model)
+
+        else:
+            self._report_metrics(report_dict)
+
+    def after_training(self, model: Booster) -> Booster:
+        if not self._checkpoint_at_end:
+            return model
+
+        if (
+            self._last_checkpoint_iteration is not None
+            and model.num_boosted_rounds() - 1 == self._last_checkpoint_iteration
+        ):
+            # Avoids a duplicate checkpoint if the checkpoint frequency happens
+            # to align with the last iteration.
+            return model
+
+        report_dict = self._get_report_dict(self._evals_log) if self._evals_log else {}
+        self._save_and_report_checkpoint(report_dict, model)
+
+        return model
+
+
+@PublicAPI(stability="beta")
+class RayTrainReportCallback(RayReportCallback):
+    """XGBoost callback to save checkpoints and report metrics.
+
+    Args:
+        metrics: Metrics to report. If this is a list,
+            each item describes the metric key reported to XGBoost,
+            and it will be reported under the same name.
+            This can also be a dict of {<key-to-report>: <xgboost-metric-key>},
+            which can be used to rename xgboost default metrics.
+        filename: Customize the saved checkpoint file type by passing
+            a filename. Defaults to "model.ubj".
+        frequency: How often to save checkpoints, in terms of iterations.
+            Defaults to 0 (no checkpoints are saved during training).
+        checkpoint_at_end: Whether or not to save a checkpoint at the end of training.
+        results_postprocessing_fn: An optional Callable that takes in
+            the metrics dict that will be reported (after it has been flattened)
+            and returns a modified dict. For example, this can be used to
+            average results across CV fold when using ``xgboost.cv``.
+
+    Examples:
+
+        Reporting checkpoints and metrics to Ray Tune when running many
+        independent xgboost trials (without data parallelism within a trial).
+
+        .. testcode::
+            :skipif: True
+
+            import xgboost
+
+            from ray.tune import Tuner
+            from ray.train.xgboost import RayTrainReportCallback
+
+            def train_fn(config):
+                # Report log loss to Ray Tune after each validation epoch.
+                bst = xgboost.train(
+                    ...,
+                    callbacks=[
+                        RayTrainReportCallback(
+                            metrics={"loss": "eval-logloss"}, frequency=1
+                        )
+                    ],
+                )
+
+            tuner = Tuner(train_fn)
+            results = tuner.fit()
+
+        Loading a model from a checkpoint reported by this callback.
+
+        .. testcode::
+            :skipif: True
+
+            from ray.train.xgboost import RayTrainReportCallback
+
+            # Get a `Checkpoint` object that is saved by the callback during training.
+            result = trainer.fit()
+            booster = RayTrainReportCallback.get_model(result.checkpoint)
+
+    """
+
+    def __init__(
+        self,
+        metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
+        filename: str = RayReportCallback.CHECKPOINT_NAME,
+        frequency: int = 0,
+        checkpoint_at_end: bool = True,
+        results_postprocessing_fn: Optional[
+            Callable[[Dict[str, Union[float, List[float]]]], Dict[str, float]]
+        ] = None,
+    ):
+        super().__init__(
+            metrics=metrics,
+            filename=filename,
+            frequency=frequency,
+            checkpoint_at_end=checkpoint_at_end,
+            results_postprocessing_fn=results_postprocessing_fn,
+        )
+
+    @contextmanager
+    def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]:
+        # NOTE: The world rank returns None for Tune usage without Train.
+        if ray.train.get_context().get_world_rank() in (0, None):
+            with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+                model.save_model(Path(temp_checkpoint_dir, self._filename).as_posix())
+                yield Checkpoint(temp_checkpoint_dir)
+        else:
+            yield None
+
+    def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster):
+        with self._get_checkpoint(model=model) as checkpoint:
+            ray.train.report(report_dict, checkpoint=checkpoint)
+
+    def _report_metrics(self, report_dict: Dict):
+        ray.train.report(report_dict)

python/ray/train/xgboost/config.py

@@ -0,0 +1,210 @@
+import json
+import logging
+import os
+import threading
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import xgboost
+from packaging.version import Version
+from xgboost import RabitTracker
+from xgboost.collective import CommunicatorContext
+
+import ray
+from ray.train._internal.base_worker_group import BaseWorkerGroup
+from ray.train.backend import Backend, BackendConfig
+from ray.train.v2._internal.util import TrainingFramework
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class XGBoostConfig(BackendConfig):
+    """Configuration for xgboost collective communication setup.
+
+    Ray Train will set up the necessary coordinator processes and environment
+    variables for your workers to communicate with each other.
+    Additional configuration options can be passed into the
+    `xgboost.collective.CommunicatorContext` that wraps your own `xgboost.train` code.
+
+    See the `xgboost.collective` module for more information:
+    https://github.com/dmlc/xgboost/blob/master/python-package/xgboost/collective.py
+
+    Args:
+        xgboost_communicator: The backend to use for collective communication for
+            distributed xgboost training. For now, only "rabit" is supported.
+    """
+
+    xgboost_communicator: str = "rabit"
+
+    @property
+    def train_func_context(self):
+        @contextmanager
+        def collective_communication_context():
+            with CommunicatorContext(**_get_xgboost_args()):
+                yield
+
+        return collective_communication_context
+
+    @property
+    def framework(self):
+        return TrainingFramework.XGBOOST
+
+    @property
+    def backend_cls(self):
+        if self.xgboost_communicator == "rabit":
+            return (
+                _XGBoostRabitBackend
+                if Version(xgboost.__version__) >= Version("2.1.0")
+                else _XGBoostRabitBackend_pre_xgb210
+            )
+
+        raise NotImplementedError(f"Unsupported backend: {self.xgboost_communicator}")
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {"xgboost_communicator": self.xgboost_communicator}
+
+
+class _XGBoostRabitBackend(Backend):
+    def __init__(self):
+        self._tracker: Optional[RabitTracker] = None
+        self._wait_thread: Optional[threading.Thread] = None
+
+    def _setup_xgboost_distributed_backend(self, worker_group: BaseWorkerGroup):
+        # 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: BaseWorkerGroup, backend_config: XGBoostConfig
+    ):
+        assert backend_config.xgboost_communicator == "rabit"
+        self._setup_xgboost_distributed_backend(worker_group)
+
+    def on_shutdown(self, worker_group: BaseWorkerGroup, backend_config: XGBoostConfig):
+        timeout = 5
+
+        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."
+                )
+
+
+class _XGBoostRabitBackend_pre_xgb210(Backend):
+    def __init__(self):
+        self._tracker: Optional[RabitTracker] = None
+
+    def _setup_xgboost_distributed_backend(self, worker_group: BaseWorkerGroup):
+        # Set up the rabit tracker on the Train driver.
+        num_workers = len(worker_group)
+        rabit_args = {"DMLC_NUM_WORKER": 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(n_workers=num_workers)
+
+        worker_args = self._tracker.worker_envs()
+        rabit_args.update(worker_args)
+
+        start_log = (
+            "RabitTracker coordinator started with parameters:\n"
+            f"{json.dumps(rabit_args, indent=2)}"
+        )
+        logger.debug(start_log)
+
+        def set_xgboost_env_vars():
+            import ray.train
+
+            for k, v in rabit_args.items():
+                os.environ[k] = str(v)
+
+            # Ranks are assigned in increasing order of the worker's task id.
+            # This task id will be sorted by increasing world rank.
+            os.environ["DMLC_TASK_ID"] = (
+                f"[xgboost.ray-rank={ray.train.get_context().get_world_rank():08}]:"
+                f"{ray.get_runtime_context().get_actor_id()}"
+            )
+
+        worker_group.execute(set_xgboost_env_vars)
+
+    def on_training_start(
+        self, worker_group: BaseWorkerGroup, backend_config: XGBoostConfig
+    ):
+        assert backend_config.xgboost_communicator == "rabit"
+        self._setup_xgboost_distributed_backend(worker_group)
+
+    def on_shutdown(self, worker_group: BaseWorkerGroup, backend_config: XGBoostConfig):
+        if not self._tracker:
+            return
+
+        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."
+            )
+
+
+_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/ray/train/xgboost/v2.py

@@ -0,0 +1,127 @@
+import logging
+from typing import Any, Callable, Dict, Optional, Union
+
+import ray.train
+from ray.train import Checkpoint
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.trainer import GenDataset
+from ray.train.xgboost import XGBoostConfig
+
+logger = logging.getLogger(__name__)
+
+
+class XGBoostTrainer(DataParallelTrainer):
+    """A Trainer for distributed data-parallel XGBoost training.
+
+    Example:
+
+        .. testcode::
+            :skipif: True
+
+            import xgboost
+
+            import ray.data
+            import ray.train
+            from ray.train.xgboost import RayTrainReportCallback, XGBoostTrainer
+
+            def train_fn_per_worker(config: dict):
+                # (Optional) Add logic to resume training state from a checkpoint.
+                # ray.train.get_checkpoint()
+
+                # 1. Get the dataset shard for the worker and convert to a `xgboost.DMatrix`
+                train_ds_iter, eval_ds_iter = (
+                    ray.train.get_dataset_shard("train"),
+                    ray.train.get_dataset_shard("validation"),
+                )
+                train_ds, eval_ds = train_ds_iter.materialize(), eval_ds_iter.materialize()
+
+                train_df, eval_df = train_ds.to_pandas(), eval_ds.to_pandas()
+                train_X, train_y = train_df.drop("y", axis=1), train_df["y"]
+                eval_X, eval_y = eval_df.drop("y", axis=1), eval_df["y"]
+
+                dtrain = xgboost.DMatrix(train_X, label=train_y)
+                deval = xgboost.DMatrix(eval_X, label=eval_y)
+
+                params = {
+                    "tree_method": "approx",
+                    "objective": "reg:squarederror",
+                    "eta": 1e-4,
+                    "subsample": 0.5,
+                    "max_depth": 2,
+                }
+
+                # 2. Do distributed data-parallel training.
+                # Ray Train sets up the necessary coordinator processes and
+                # environment variables for your workers to communicate with each other.
+                bst = xgboost.train(
+                    params,
+                    dtrain=dtrain,
+                    evals=[(deval, "validation")],
+                    num_boost_round=10,
+                    callbacks=[RayTrainReportCallback()],
+                )
+
+            train_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
+            eval_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(16)])
+            trainer = XGBoostTrainer(
+                train_fn_per_worker,
+                datasets={"train": train_ds, "validation": eval_ds},
+                scaling_config=ray.train.ScalingConfig(num_workers=4),
+            )
+            result = trainer.fit()
+            booster = RayTrainReportCallback.get_model(result.checkpoint)
+
+    Args:
+        train_loop_per_worker: The training function to execute on each worker.
+            This function can either take in zero arguments or a single ``Dict``
+            argument which is set by defining ``train_loop_config``.
+            Within this function you can use any of the
+            :ref:`Ray Train Loop utilities <train-loop-api>`.
+        train_loop_config: A configuration ``Dict`` to pass in as an argument to
+            ``train_loop_per_worker``.
+            This is typically used for specifying hyperparameters.
+        xgboost_config: The configuration for setting up the distributed xgboost
+            backend. Defaults to using the "rabit" backend.
+            See :class:`~ray.train.xgboost.XGBoostConfig` for more info.
+        scaling_config: The configuration for how to scale data parallel training.
+            ``num_workers`` determines how many Python processes are used for training,
+            and ``use_gpu`` determines whether or not each process should use GPUs.
+            See :class:`~ray.train.ScalingConfig` for more info.
+        run_config: The configuration for the execution of the training run.
+            See :class:`~ray.train.RunConfig` for more info.
+        datasets: The Ray Datasets to use for training and validation.
+        dataset_config: The configuration for ingesting the input ``datasets``.
+            By default, all the Ray Datasets are split equally across workers.
+            See :class:`~ray.train.DataConfig` for more details.
+        metadata: Dict that should be made available via
+            `ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
+            for checkpoints saved from this Trainer. Must be JSON-serializable.
+        resume_from_checkpoint: A checkpoint to resume training from.
+            This checkpoint can be accessed from within ``train_loop_per_worker``
+            by calling ``ray.train.get_checkpoint()``.
+    """
+
+    def __init__(
+        self,
+        train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]],
+        *,
+        train_loop_config: Optional[Dict] = None,
+        xgboost_config: Optional[XGBoostConfig] = None,
+        scaling_config: Optional[ray.train.ScalingConfig] = None,
+        run_config: Optional[ray.train.RunConfig] = None,
+        datasets: Optional[Dict[str, GenDataset]] = None,
+        dataset_config: Optional[ray.train.DataConfig] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        resume_from_checkpoint: Optional[Checkpoint] = None,
+    ):
+        super(XGBoostTrainer, self).__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            train_loop_config=train_loop_config,
+            backend_config=xgboost_config or XGBoostConfig(),
+            scaling_config=scaling_config,
+            dataset_config=dataset_config,
+            run_config=run_config,
+            datasets=datasets,
+            resume_from_checkpoint=resume_from_checkpoint,
+            metadata=metadata,
+        )

python/ray/train/xgboost/xgboost_checkpoint.py

@@ -0,0 +1,75 @@
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+import xgboost
+
+from ray.train._internal.framework_checkpoint import FrameworkCheckpoint
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.preprocessor import Preprocessor
+
+
+@PublicAPI(stability="beta")
+class XGBoostCheckpoint(FrameworkCheckpoint):
+    """A :py:class:`~ray.train.Checkpoint` with XGBoost-specific functionality."""
+
+    MODEL_FILENAME = "model.json"
+
+    @classmethod
+    def from_model(
+        cls,
+        booster: xgboost.Booster,
+        *,
+        preprocessor: Optional["Preprocessor"] = None,
+        path: Optional[str] = None,
+    ) -> "XGBoostCheckpoint":
+        """Create a :py:class:`~ray.train.Checkpoint` that stores an XGBoost
+        model.
+
+        Args:
+            booster: The XGBoost model to store in the checkpoint.
+            preprocessor: A fitted preprocessor to be applied before inference.
+            path: The path to the directory where the checkpoint file will be saved.
+                This should start as an empty directory, since the *entire*
+                directory will be treated as the checkpoint when reported.
+                By default, a temporary directory will be created.
+
+        Returns:
+            An :py:class:`XGBoostCheckpoint` containing the specified ``Estimator``.
+
+        Examples:
+
+            ... testcode::
+
+                import numpy as np
+                import ray
+                from ray.train.xgboost import XGBoostCheckpoint
+                import xgboost
+
+                train_X = np.array([[1, 2], [3, 4]])
+                train_y = np.array([0, 1])
+
+                model = xgboost.XGBClassifier().fit(train_X, train_y)
+                checkpoint = XGBoostCheckpoint.from_model(model.get_booster())
+
+        """
+        checkpoint_path = Path(path or tempfile.mkdtemp())
+
+        if not checkpoint_path.is_dir():
+            raise ValueError(f"`path` must be a directory, but got: {checkpoint_path}")
+
+        booster.save_model(checkpoint_path.joinpath(cls.MODEL_FILENAME).as_posix())
+
+        checkpoint = cls.from_directory(checkpoint_path.as_posix())
+        if preprocessor:
+            checkpoint.set_preprocessor(preprocessor)
+        return checkpoint
+
+    def get_model(self) -> xgboost.Booster:
+        """Retrieve the XGBoost model stored in this checkpoint."""
+        with self.as_directory() as checkpoint_path:
+            booster = xgboost.Booster()
+            booster.load_model(Path(checkpoint_path, self.MODEL_FILENAME).as_posix())
+            return booster

python/ray/train/xgboost/xgboost_trainer.py

@@ -0,0 +1,312 @@
+import logging
+from functools import partial
+from typing import Any, Callable, Dict, Optional, Union
+
+import xgboost
+from packaging.version import Version
+
+import ray.train
+from ray.train import Checkpoint
+from ray.train.constants import TRAIN_DATASET_KEY
+from ray.train.trainer import GenDataset
+from ray.train.utils import _log_deprecation_warning
+from ray.train.xgboost import RayTrainReportCallback, XGBoostConfig
+from ray.train.xgboost.v2 import XGBoostTrainer as SimpleXGBoostTrainer
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(__name__)
+
+
+LEGACY_XGBOOST_TRAINER_DEPRECATION_MESSAGE = (
+    "Passing in `xgboost.train` kwargs such as `params`, `num_boost_round`, "
+    "`label_column`, etc. to `XGBoostTrainer` is deprecated "
+    "in favor of the new API which accepts a training function, "
+    "similar to the other DataParallelTrainer APIs (ex: TorchTrainer). "
+    "See this issue for more context: "
+    "https://github.com/ray-project/ray/issues/50042"
+)
+
+
+def _xgboost_train_fn_per_worker(
+    config: dict,
+    label_column: str,
+    num_boost_round: int,
+    dataset_keys: set,
+    xgboost_train_kwargs: dict,
+):
+    checkpoint = ray.train.get_checkpoint()
+    starting_model = None
+    remaining_iters = num_boost_round
+    if checkpoint:
+        starting_model = RayTrainReportCallback.get_model(checkpoint)
+        starting_iter = starting_model.num_boosted_rounds()
+        remaining_iters = num_boost_round - starting_iter
+        logger.info(
+            f"Model loaded from checkpoint will train for "
+            f"additional {remaining_iters} iterations (trees) in order "
+            "to achieve the target number of iterations "
+            f"({num_boost_round=})."
+        )
+
+    train_ds_iter = ray.train.get_dataset_shard(TRAIN_DATASET_KEY)
+    train_df = train_ds_iter.materialize().to_pandas()
+
+    eval_ds_iters = {
+        k: ray.train.get_dataset_shard(k)
+        for k in dataset_keys
+        if k != TRAIN_DATASET_KEY
+    }
+    eval_dfs = {k: d.materialize().to_pandas() for k, d in eval_ds_iters.items()}
+
+    train_X, train_y = train_df.drop(label_column, axis=1), train_df[label_column]
+    dtrain = xgboost.DMatrix(train_X, label=train_y)
+
+    # NOTE: Include the training dataset in the evaluation datasets.
+    # This allows `train-*` metrics to be calculated and reported.
+    evals = [(dtrain, TRAIN_DATASET_KEY)]
+
+    for eval_name, eval_df in eval_dfs.items():
+        eval_X, eval_y = eval_df.drop(label_column, axis=1), eval_df[label_column]
+        evals.append((xgboost.DMatrix(eval_X, label=eval_y), eval_name))
+
+    evals_result = {}
+    xgboost.train(
+        config,
+        dtrain=dtrain,
+        evals=evals,
+        evals_result=evals_result,
+        num_boost_round=remaining_iters,
+        xgb_model=starting_model,
+        **xgboost_train_kwargs,
+    )
+
+
+@PublicAPI(stability="beta")
+class XGBoostTrainer(SimpleXGBoostTrainer):
+    """A Trainer for distributed data-parallel XGBoost training.
+
+    Example:
+
+        .. testcode::
+            :skipif: True
+
+            import xgboost
+
+            import ray.data
+            import ray.train
+            from ray.train.xgboost import RayTrainReportCallback, XGBoostTrainer
+
+            def train_fn_per_worker(config: dict):
+                # (Optional) Add logic to resume training state from a checkpoint.
+                # ray.train.get_checkpoint()
+
+                # 1. Get the dataset shard for the worker and convert to a `xgboost.DMatrix`
+                train_ds_iter, eval_ds_iter = (
+                    ray.train.get_dataset_shard("train"),
+                    ray.train.get_dataset_shard("validation"),
+                )
+                train_ds, eval_ds = train_ds_iter.materialize(), eval_ds_iter.materialize()
+
+                train_df, eval_df = train_ds.to_pandas(), eval_ds.to_pandas()
+                train_X, train_y = train_df.drop("y", axis=1), train_df["y"]
+                eval_X, eval_y = eval_df.drop("y", axis=1), eval_df["y"]
+
+                dtrain = xgboost.DMatrix(train_X, label=train_y)
+                deval = xgboost.DMatrix(eval_X, label=eval_y)
+
+                params = {
+                    "tree_method": "approx",
+                    "objective": "reg:squarederror",
+                    "eta": 1e-4,
+                    "subsample": 0.5,
+                    "max_depth": 2,
+                }
+
+                # 2. Do distributed data-parallel training.
+                # Ray Train sets up the necessary coordinator processes and
+                # environment variables for your workers to communicate with each other.
+                bst = xgboost.train(
+                    params,
+                    dtrain=dtrain,
+                    evals=[(deval, "validation")],
+                    num_boost_round=10,
+                    callbacks=[RayTrainReportCallback()],
+                )
+
+            train_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
+            eval_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(16)])
+            trainer = XGBoostTrainer(
+                train_fn_per_worker,
+                datasets={"train": train_ds, "validation": eval_ds},
+                scaling_config=ray.train.ScalingConfig(num_workers=4),
+            )
+            result = trainer.fit()
+            booster = RayTrainReportCallback.get_model(result.checkpoint)
+
+    Args:
+        train_loop_per_worker: The training function to execute on each worker.
+            This function can either take in zero arguments or a single ``Dict``
+            argument which is set by defining ``train_loop_config``.
+            Within this function you can use any of the
+            :ref:`Ray Train Loop utilities <train-loop-api>`.
+        train_loop_config: A configuration ``Dict`` to pass in as an argument to
+            ``train_loop_per_worker``.
+            This is typically used for specifying hyperparameters.
+        xgboost_config: The configuration for setting up the distributed xgboost
+            backend. Defaults to using the "rabit" backend.
+            See :class:`~ray.train.xgboost.XGBoostConfig` for more info.
+        scaling_config: The configuration for how to scale data parallel training.
+            ``num_workers`` determines how many Python processes are used for training,
+            and ``use_gpu`` determines whether or not each process should use GPUs.
+            See :class:`~ray.train.ScalingConfig` for more info.
+        run_config: The configuration for the execution of the training run.
+            See :class:`~ray.train.RunConfig` for more info.
+        datasets: The Ray Datasets to use for training and validation.
+        dataset_config: The configuration for ingesting the input ``datasets``.
+            By default, all the Ray Datasets are split equally across workers.
+            See :class:`~ray.train.DataConfig` for more details.
+        resume_from_checkpoint: A checkpoint to resume training from.
+            This checkpoint can be accessed from within ``train_loop_per_worker``
+            by calling ``ray.train.get_checkpoint()``.
+        metadata: Dict that should be made available via
+            `ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
+            for checkpoints saved from this Trainer. Must be JSON-serializable.
+        label_column: [Deprecated] Name of the label column. A column with this name
+            must be present in the training dataset.
+        params: [Deprecated] XGBoost training parameters.
+            Refer to `XGBoost documentation <https://xgboost.readthedocs.io/>`_
+            for a list of possible parameters.
+        num_boost_round: [Deprecated] Target number of boosting iterations (trees in the model).
+            Note that unlike in ``xgboost.train``, this is the target number
+            of trees, meaning that if you set ``num_boost_round=10`` and pass a model
+            that has already been trained for 5 iterations, it will be trained for 5
+            iterations more, instead of 10 more.
+        **train_kwargs: [Deprecated] Additional kwargs passed to ``xgboost.train()`` function.
+    """
+
+    _handles_checkpoint_freq = True
+    _handles_checkpoint_at_end = True
+
+    def __init__(
+        self,
+        train_loop_per_worker: Optional[
+            Union[Callable[[], None], Callable[[Dict], None]]
+        ] = None,
+        *,
+        train_loop_config: Optional[Dict] = None,
+        xgboost_config: Optional[XGBoostConfig] = None,
+        scaling_config: Optional[ray.train.ScalingConfig] = None,
+        run_config: Optional[ray.train.RunConfig] = None,
+        datasets: Optional[Dict[str, GenDataset]] = None,
+        dataset_config: Optional[ray.train.DataConfig] = None,
+        resume_from_checkpoint: Optional[Checkpoint] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        # TODO(justinvyu): [Deprecated] Legacy XGBoostTrainer API
+        label_column: Optional[str] = None,
+        params: Optional[Dict[str, Any]] = None,
+        num_boost_round: Optional[int] = None,
+        **train_kwargs,
+    ):
+        if Version(xgboost.__version__) < Version("1.7.0"):
+            raise ImportError(
+                "`XGBoostTrainer` requires the `xgboost` version to be >= 1.7.0. "
+                'Upgrade with: `pip install -U "xgboost>=1.7"`'
+            )
+
+        # TODO(justinvyu): [Deprecated] Legacy XGBoostTrainer API
+        legacy_api = train_loop_per_worker is None
+        if legacy_api:
+            train_loop_per_worker = self._get_legacy_train_fn_per_worker(
+                xgboost_train_kwargs=train_kwargs,
+                run_config=run_config,
+                label_column=label_column,
+                num_boost_round=num_boost_round,
+                datasets=datasets,
+            )
+            train_loop_config = params or {}
+        elif train_kwargs:
+            _log_deprecation_warning(
+                "Passing `xgboost.train` kwargs to `XGBoostTrainer` is deprecated. "
+                "In your training function, you can call `xgboost.train(**kwargs)` "
+                "with arbitrary arguments. "
+                f"{LEGACY_XGBOOST_TRAINER_DEPRECATION_MESSAGE}"
+            )
+
+        super(XGBoostTrainer, self).__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            train_loop_config=train_loop_config,
+            xgboost_config=xgboost_config,
+            scaling_config=scaling_config,
+            run_config=run_config,
+            datasets=datasets,
+            dataset_config=dataset_config,
+            resume_from_checkpoint=resume_from_checkpoint,
+            metadata=metadata,
+        )
+
+    def _get_legacy_train_fn_per_worker(
+        self,
+        xgboost_train_kwargs: Dict,
+        run_config: Optional[ray.train.RunConfig],
+        datasets: Optional[Dict[str, GenDataset]],
+        label_column: Optional[str],
+        num_boost_round: Optional[int],
+    ) -> Callable[[Dict], None]:
+        """Get the training function for the legacy XGBoostTrainer API."""
+
+        datasets = datasets or {}
+        if not datasets.get(TRAIN_DATASET_KEY):
+            raise ValueError(
+                "`datasets` must be provided for the XGBoostTrainer API "
+                "if `train_loop_per_worker` is not provided. "
+                "This dict must contain the training dataset under the "
+                f"key: '{TRAIN_DATASET_KEY}'. "
+                f"Got keys: {list(datasets.keys())}"
+            )
+        if not label_column:
+            raise ValueError(
+                "`label_column` must be provided for the XGBoostTrainer API "
+                "if `train_loop_per_worker` is not provided. "
+                "This is the column name of the label in the dataset."
+            )
+
+        num_boost_round = num_boost_round or 10
+
+        _log_deprecation_warning(LEGACY_XGBOOST_TRAINER_DEPRECATION_MESSAGE)
+
+        # Initialize a default Ray Train metrics/checkpoint reporting callback if needed
+        callbacks = xgboost_train_kwargs.get("callbacks", [])
+        user_supplied_callback = any(
+            isinstance(callback, RayTrainReportCallback) for callback in callbacks
+        )
+        callback_kwargs = {}
+        if run_config:
+            checkpoint_frequency = run_config.checkpoint_config.checkpoint_frequency
+            checkpoint_at_end = run_config.checkpoint_config.checkpoint_at_end
+
+            callback_kwargs["frequency"] = checkpoint_frequency
+            # Default `checkpoint_at_end=True` unless the user explicitly sets it.
+            callback_kwargs["checkpoint_at_end"] = (
+                checkpoint_at_end if checkpoint_at_end is not None else True
+            )
+
+        if not user_supplied_callback:
+            callbacks.append(RayTrainReportCallback(**callback_kwargs))
+        xgboost_train_kwargs["callbacks"] = callbacks
+
+        train_fn_per_worker = partial(
+            _xgboost_train_fn_per_worker,
+            label_column=label_column,
+            num_boost_round=num_boost_round,
+            dataset_keys=set(datasets),
+            xgboost_train_kwargs=xgboost_train_kwargs,
+        )
+        return train_fn_per_worker
+
+    @classmethod
+    def get_model(
+        cls,
+        checkpoint: Checkpoint,
+    ) -> xgboost.Booster:
+        """Retrieve the XGBoost model stored in this checkpoint."""
+        return RayTrainReportCallback.get_model(checkpoint)