chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_xgboost_trainer.py

@@ -0,0 +1,138 @@
+import pandas as pd
+import pytest
+import xgboost
+from sklearn.datasets import load_breast_cancer
+from sklearn.model_selection import train_test_split
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.constants import TRAIN_DATASET_KEY
+from ray.train.v2._internal.constants import is_v2_enabled
+from ray.train.xgboost import RayTrainReportCallback, XGBoostTrainer
+
+assert is_v2_enabled()
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+scale_config = ScalingConfig(num_workers=2)
+
+data_raw = load_breast_cancer()
+dataset_df = pd.DataFrame(data_raw["data"], columns=data_raw["feature_names"])
+dataset_df["target"] = data_raw["target"]
+train_df, test_df = train_test_split(dataset_df, test_size=0.3)
+
+params = {
+    "tree_method": "approx",
+    "objective": "binary:logistic",
+    "eval_metric": ["logloss", "error"],
+}
+
+
+def test_fit(ray_start_4_cpus):
+    @ray.remote
+    class ValidationCollector:
+        def __init__(self):
+            self.validation_scores = {}
+
+        def report(self, rank, logloss, error):
+            self.validation_scores[rank] = {
+                "logloss": logloss,
+                "error": error,
+            }
+
+        def get_validation_scores(self):
+            return self.validation_scores
+
+    # Ensure all workers have the same model in data parallel training
+    # by comparing their validation scores.
+    # Comparing xgboost models directly seems less reliable.
+    collector = ValidationCollector.remote()
+
+    def xgboost_train_fn_per_worker(
+        label_column: str,
+        dataset_keys: set,
+    ):
+        checkpoint = ray.train.get_checkpoint()
+        starting_model = None
+        remaining_iters = 10
+        if checkpoint:
+            starting_model = RayTrainReportCallback.get_model(checkpoint)
+            starting_iter = starting_model.num_boosted_rounds()
+            remaining_iters = remaining_iters - starting_iter
+
+        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(
+            params,
+            dtrain=dtrain,
+            evals=evals,
+            evals_result=evals_result,
+            num_boost_round=remaining_iters,
+            xgb_model=starting_model,
+        )
+        collector.report.remote(
+            ray.train.get_context().get_world_rank(),
+            evals_result["valid"]["logloss"],
+            evals_result["valid"]["error"],
+        )
+
+    train_dataset = ray.data.from_pandas(train_df)
+    valid_dataset = ray.data.from_pandas(test_df)
+    trainer = XGBoostTrainer(
+        train_loop_per_worker=lambda: xgboost_train_fn_per_worker(
+            label_column="target",
+            dataset_keys={TRAIN_DATASET_KEY, "valid"},
+        ),
+        scaling_config=scale_config,
+        # Sharding the validation dataset across workers is ok because xgboost allreduces metrics.
+        # See https://github.com/dmlc/xgboost/blob/d0135d0f43ff91e738edcbcea54e44b50d336adf/python-package/xgboost/callback.py#L131.
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+    )
+    result = trainer.fit()
+    validation_scores = ray.get(collector.get_validation_scores.remote())
+    assert validation_scores[0]["logloss"] == pytest.approx(
+        validation_scores[1]["logloss"], abs=1e-6
+    )
+    assert validation_scores[0]["error"] == pytest.approx(
+        validation_scores[1]["error"], abs=1e-6
+    )
+    with pytest.raises(DeprecationWarning):
+        XGBoostTrainer.get_model(result.checkpoint)
+
+
+# TODO: Unit test RayTrainReportCallback
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))