XGBoost-Ray

0.1.19 · active · verified Sun Apr 12

XGBoost-Ray provides a Ray backend for distributed XGBoost, enabling training and prediction on Ray clusters with minimal code changes. It extends the core XGBoost API to leverage distributed data representations and integrates seamlessly with other Ray libraries like Ray Tune for hyperparameter optimization and Ray Train for scalable ML workloads. The library is actively maintained, with frequent updates to ensure compatibility with recent XGBoost and Ray versions.

Warnings

breaking XGBoost-Ray v0.1.19 updated its API to work with XGBoost 2.0. This may introduce breaking changes if you are upgrading `xgboost-ray` and relying on specific API behaviors from older XGBoost versions.
Fix: Ensure your `xgboost` version is compatible with `xgboost-ray` (typically XGBoost 2.0+ for recent `xgboost-ray` versions). Review XGBoost 2.0 release notes for core API changes.
breaking XGBoost-Ray v0.1.12 introduced compatibility for `xgboost>=1.7.0`. Using `xgboost-ray` with `xgboost` versions older than 1.7.0 may lead to unexpected behavior or errors.
Fix: Upgrade your `xgboost` package to version 1.7.0 or newer when using `xgboost-ray` versions 0.1.12 and above.
gotcha Always use `xgboost_ray.RayDMatrix` instead of `xgboost.DMatrix` when passing data to `xgboost_ray.train`. `RayDMatrix` is essential for distributed data handling and sharding across Ray actors.
Fix: Change your import and data matrix creation from `xgb.DMatrix(...)` to `RayDMatrix(...)`.
gotcha Explicitly configure `num_actors` and `cpus_per_actor` in `RayParams`. While `xgboost-ray` attempts to auto-configure, manual specification is crucial for optimal performance, especially in heterogeneous clusters or multi-GPU setups. Ensure enough CPUs are available for Ray Data operations if used.
Fix: Set `RayParams(num_actors=..., cpus_per_actor=...)` based on your cluster resources and workload. Consider reserving some CPU for Ray Data if performing heavy data operations.
gotcha When using `RayDMatrix` with data sources that cannot be naturally sharded (e.g., a single large Parquet file), you may encounter a `RuntimeError` about insufficient shards. In such cases, enable centralized loading.
Fix: Pass `distributed=False` to `RayDMatrix` (e.g., `RayDMatrix(data, labels, distributed=False)`) to have the head node shard the data into the Ray object store.
gotcha On macOS, Ray's performance can degrade if the object store size exceeds 2.0GB. This may lead to warnings or slower execution for large datasets.
Fix: Either specify `object_store_memory` in `ray.init()` to limit the size, or set the environment variable `RAY_ENABLE_MAC_LARGE_OBJECT_STORE=1` to ignore the warning (use with caution).

Install

pip install xgboost-ray Install stable version

Imports

RayDMatrix
```
from xgboost_ray import RayDMatrix
```
xgboost_ray.RayDMatrix is required for data handling in a Ray distributed environment, replacing the native xgboost.DMatrix.
RayParams
```
from xgboost_ray import RayParams
```
Used to configure Ray-specific distributed training parameters, such as the number of actors.
train
```
from xgboost_ray import train
```
The distributed training function provided by xgboost-ray, a drop-in replacement for xgboost.train.
predict
```
from xgboost_ray import predict
```
The distributed prediction function for models trained with xgboost-ray.

Quickstart

This quickstart demonstrates how to train a distributed XGBoost model using `xgboost-ray`. It initializes a local Ray cluster, loads a sample dataset, prepares it as a `RayDMatrix`, configures `RayParams` for distributed execution, and then calls `xgboost_ray.train`.

import ray
from xgboost_ray import RayDMatrix, RayParams, train
from sklearn.datasets import load_breast_cancer
import os

# Initialize Ray (if not already running)
if ray.is_initialized():
    ray.shutdown()
ray.init(log_to_stdout=False) # log_to_stdout=False to suppress verbose Ray output in quickstart

# Load data
train_x, train_y = load_breast_cancer(return_X_y=True)

# Create RayDMatrix for distributed data handling
train_set = RayDMatrix(train_x, train_y)

# Configure Ray-specific training parameters
ray_params = RayParams(
    num_actors=2,
    cpus_per_actor=1
)

# Train the model using the xgboost-ray distributed train function
evals_result = {}
bst = train(
    {
        "objective": "binary:logistic",
        "eval_metric": ["logloss", "error"],
    },
    train_set,
    evals_result=evals_result,
    evals=[(train_set, "train")],
    verbose_eval=False,
    ray_params=ray_params
)

print(f"Final training error: {evals_result['train']['error'][-1]:.4f}")

# Shutdown Ray cluster
ray.shutdown()

view raw JSON →