[tune] add xgboost callbacks to integration module (#10502)

doc/source/conf.py

@@ -65,6 +65,7 @@ MOCK_MODULES = [
     "torch.utils.data",
     "torch.utils.data.distributed",
     "wandb",
+    "xgboost",
     "zoopt",
 ]
 import scipy.stats

doc/source/tune/_tutorials/tune-xgboost.rst

@@ -382,12 +382,19 @@ There are more parameters, which are explained in the
 :ref:`documentation <tune-scheduler-hyperband>`.
 
 Lastly, we have to report the loss metric to Tune. We do this with a ``Callback`` that
-XGBoost accepts and calls after each training iteration. We also tell XGBoost which
-loss metrics to calculate in the ``eval_metric`` parameter. These are the metrics
-available in ``env.evaluation_result_list`` below.
+XGBoost accepts and calls after each evaluation round. Ray Tune comes
+with :ref:`two XGBoost callbacks <tune-integration-xgboost>`
+we can use for this. The ``TuneReportCallback`` just reports the evaluation
+metrics back to Tune. The ``TuneReportCheckpointCallback`` would also save
+checkpoints after each evaluation round. We will just use the former in this
+example.
+
+We also tell XGBoost which loss metrics to calculate in the ``eval_metric``
+parameter in the config. These parameters are then reported to Tune
+via the callback.
 
 .. code-block:: python
-   :emphasize-lines: 11,12,13,26,42,44,45,46,47,48,49
+   :emphasize-lines: 9,26,42,44-49
 
     import numpy as np
     import sklearn.datasets
@@ -397,12 +404,7 @@ available in ``env.evaluation_result_list`` below.
     import xgboost as xgb
 
     from ray import tune
-
-
-    def XGBCallback(env):
-        # After every training iteration, report loss to Tune
-        tune.report(**dict(env.evaluation_result_list))
-
+    from ray.tune.integration.xgboost import TuneReportCallback
 
     def train_breast_cancer(config):
         # Load dataset
@@ -414,7 +416,12 @@ available in ``env.evaluation_result_list`` below.
         train_set = xgb.DMatrix(train_x, label=train_y)
         test_set = xgb.DMatrix(test_x, label=test_y)
         # Train the classifier
-        bst = xgb.train(config, train_set, evals=[(test_set, "eval")], verbose_eval=False, callbacks=[XGBCallback])
+        bst = xgb.train(
+            config,
+            train_set,
+            evals=[(test_set, "eval")],
+            verbose_eval=False,
+            callbacks=[TuneReportCallback()])
         # Predict labels for the test set
         preds = bst.predict(test_set)
         pred_labels = np.rint(preds)

doc/source/tune/api_docs/integration.rst

@@ -43,4 +43,14 @@ Weights and Biases (tune.integration.wandb)
 
 .. autoclass:: ray.tune.integration.wandb.WandbLogger
 
-.. autofunction:: ray.tune.integration.wandb.wandb_mixin
+.. autofunction:: ray.tune.integration.wandb.wandb_mixin
+
+
+.. _tune-integration-xgboost:
+
+XGBoost (tune.integration.xgboost)
+----------------------------------
+
+.. autoclass:: ray.tune.integration.xgboost.TuneReportCallback
+
+.. autoclass:: ray.tune.integration.xgboost.TuneReportCheckpointCallback

python/ray/tune/examples/xgboost_example.py

@@ -6,11 +6,7 @@ from sklearn.model_selection import train_test_split
 import xgboost as xgb
 
 from ray import tune
-
-
-def XGBCallback(env):
-    # After every training iteration, report loss to Tune
-    tune.report(**dict(env.evaluation_result_list))
+from ray.tune.integration.xgboost import TuneReportCallback
 
 
 def train_breast_cancer(config):
@@ -28,7 +24,7 @@ def train_breast_cancer(config):
         train_set,
         evals=[(test_set, "eval")],
         verbose_eval=False,
-        callbacks=[XGBCallback])
+        callbacks=[TuneReportCallback()])
     # Predict labels for the test set
     preds = bst.predict(test_set)
     pred_labels = np.rint(preds)

python/ray/tune/integration/xgboost.py

@@ -0,0 +1,144 @@
+from typing import Dict, List, Union
+from ray import tune
+
+import os
+
+
+class TuneCallback:
+    """Base class for Tune's XGBoost callbacks."""
+    pass
+
+    def __call__(self, env):
+        raise NotImplementedError
+
+
+class TuneReportCallback(TuneCallback):
+    """XGBoost to Ray Tune reporting callback
+
+    Reports metrics to Ray Tune.
+
+    Args:
+        metrics (str|list|dict): Metrics to report to Tune. If this is a list,
+            each item describes the metric key reported to XGBoost,
+            and it will reported under the same name to Tune. If this is a
+            dict, each key will be the name reported to Tune and the respective
+            value will be the metric key reported to XGBoost. If this is None,
+            all metrics will be reported to Tune under their default names as
+            obtained from XGBoost.
+
+    Example:
+
+    .. code-block:: python
+
+        import xgboost
+        from ray.tune.integration.xgboost import TuneReportCallback
+
+        config = {
+            # ...
+            "eval_metric": ["auc", "logloss"]
+        }
+
+        # Report only log loss to Tune after each validation epoch:
+        bst = xgb.train(
+            config,
+            train_set,
+            evals=[(test_set, "eval")],
+            verbose_eval=False,
+            callbacks=[TuneReportCallback({"loss": "eval-logloss"})])
+
+    """
+
+    def __init__(self,
+                 metrics: Union[None, str, List[str], Dict[str, str]] = None):
+        if isinstance(metrics, str):
+            metrics = [metrics]
+        self._metrics = metrics
+
+    def __call__(self, env):
+        result_dict = dict(env.evaluation_result_list)
+        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]
+        tune.report(**report_dict)
+
+
+class _TuneCheckpointCallback(TuneCallback):
+    """XGBoost checkpoint callback
+
+    Saves checkpoints after each validation step.
+
+    Checkpoint are currently not registered if no ``tune.report()`` call
+    is made afterwards. Consider using ``TuneReportCheckpointCallback``
+    instead.
+
+    Args:
+        filename (str): Filename of the checkpoint within the checkpoint
+            directory. Defaults to "checkpoint".
+
+    """
+
+    def __init__(self, filename: str = "checkpoint"):
+        self._filename = filename
+
+    def __call__(self, env):
+        with tune.checkpoint_dir(step=env.iteration) as checkpoint_dir:
+            env.model.save_model(os.path.join(checkpoint_dir, self._filename))
+
+
+class TuneReportCheckpointCallback(TuneCallback):
+    """XGBoost report and checkpoint callback
+
+    Saves checkpoints after each validation step. Also reports metrics to Tune,
+    which is needed for checkpoint registration.
+
+    Args:
+        metrics (str|list|dict): Metrics to report to Tune. If this is a list,
+            each item describes the metric key reported to XGBoost,
+            and it will reported under the same name to Tune. If this is a
+            dict, each key will be the name reported to Tune and the respective
+            value will be the metric key reported to XGBoost.
+        filename (str): Filename of the checkpoint within the checkpoint
+            directory. Defaults to "checkpoint". If this is None,
+            all metrics will be reported to Tune under their default names as
+            obtained from XGBoost.
+
+    Example:
+
+    .. code-block:: python
+
+        import xgboost
+        from ray.tune.integration.xgboost import TuneReportCheckpointCallback
+
+        config = {
+            # ...
+            "eval_metric": ["auc", "logloss"]
+        }
+
+        # Report only log loss to Tune after each validation epoch.
+        # Save model as `xgboost.mdl`.
+        bst = xgb.train(
+            config,
+            train_set,
+            evals=[(test_set, "eval")],
+            verbose_eval=False,
+            callbacks=[TuneReportCheckpointCallback(
+                {"loss": "eval-logloss"}, "xgboost.mdl)])
+
+    """
+
+    def __init__(self,
+                 metrics: Union[None, str, List[str], Dict[str, str]] = None,
+                 filename: str = "checkpoint"):
+        self._checkpoint = _TuneCheckpointCallback(filename)
+        self._report = TuneReportCallback(metrics)
+
+    def __call__(self, env):
+        self._checkpoint(env)
+        self._report(env)