trainer/examples/xgboost/distributed-training/xgboost-distributed.ipynb

Distributed XGBoost Training with Kubeflow Trainer

This notebook demonstrates how to run distributed XGBoost training on Kubernetes using the Kubeflow Trainer SDK.

XGBoost uses the Rabit/Collective protocol for distributed coordination. The Kubeflow XGBoost runtime automatically injects the required DMLC_* environment variables into each worker pod.

Install the Kubeflow SDK

You need to install the Kubeflow SDK to interact with Kubeflow Trainer APIs:

In [ ]:

 #!pip install -U kubeflow

Define Training Function

This function will be serialized and executed on each XGBoost worker node. It uses:

• XGBoost Collective API (xgb.collective) for distributed coordination
• RabitTracker on rank-0 for worker coordination
• sklearn for generating synthetic classification data

The DMLC_* environment variables are automatically injected by the Kubeflow XGBoost runtime plugin.

In [ ]:

def xgboost_train_classification():
    """
    Distributed XGBoost training function using the modern Collective API.

    DMLC_* env vars are injected by the Kubeflow Trainer XGBoost plugin:
      - DMLC_TRACKER_URI:  DNS name of the rank-0 pod running the tracker
      - DMLC_TRACKER_PORT: Port for tracker communication
      - DMLC_TASK_ID:      Worker rank (0, 1, 2, ...)
      - DMLC_NUM_WORKER:   Total number of workers
    """
    import os
    import xgboost as xgb
    from xgboost import collective as coll
    from xgboost.tracker import RabitTracker
    from sklearn.datasets import make_classification
    from sklearn.model_selection import train_test_split
    from sklearn.metrics import accuracy_score

    # Read injected environment variables.
    rank = int(os.environ["DMLC_TASK_ID"])
    world_size = int(os.environ["DMLC_NUM_WORKER"])
    tracker_uri = os.environ["DMLC_TRACKER_URI"]
    tracker_port = int(os.environ["DMLC_TRACKER_PORT"])

    print(f"Worker {rank}/{world_size} starting...")
    print(f"Tracker: {tracker_uri}:{tracker_port}")

    # Rank 0 starts the Rabit tracker (required for coordination).
    tracker = None
    if rank == 0:
        tracker = RabitTracker(host_ip="0.0.0.0", n_workers=world_size, port=tracker_port)
        tracker.start()
        print(f"Tracker started on {tracker_uri}:{tracker_port}")

    # All workers connect to the tracker via the Collective communicator.
    with coll.CommunicatorContext(
        dmlc_tracker_uri=tracker_uri,
        dmlc_tracker_port=tracker_port,
        dmlc_task_id=str(rank),
    ):
        print(f"Worker {coll.get_rank()}/{coll.get_world_size()} connected")

        # Generate synthetic classification data.
        # Each worker generates a different shard using rank-based random seed.
        X, y = make_classification(
            n_samples=10000, n_features=20, n_informative=10,
            n_classes=2, random_state=42 + rank,
        )
        X_train, X_valid, y_train, y_valid = train_test_split(
            X, y, test_size=0.2, random_state=42,
        )

        # NOTE: QuantileDMatrix construction MUST be inside the communicator context
        # because it involves cross-worker synchronization for quantization.
        dtrain = xgb.QuantileDMatrix(X_train, label=y_train)
        dvalid = xgb.QuantileDMatrix(X_valid, label=y_valid, ref=dtrain)

        # Training parameters (for GPU training, add device="cuda").
        params = {
            "objective": "binary:logistic",
            "tree_method": "hist",
            "max_depth": 6,
            "eta": 0.1,
            "eval_metric": "logloss",
        }

        # Distributed training - collective operations synchronize gradients.
        model = xgb.train(
            params, dtrain,
            num_boost_round=50,
            evals=[(dvalid, "validation")],
            verbose_eval=10,
        )

        # Evaluate on validation set.
        preds = model.predict(dvalid)
        predictions = [1 if p > 0.5 else 0 for p in preds]
        accuracy = accuracy_score(y_valid, predictions)
        print(f"Worker {rank} - Validation Accuracy: {accuracy:.4f}")

        # Save model from rank 0 only.
        if coll.get_rank() == 0:
            model.save_model("/workspace/xgboost_model.json")
            print("Model saved to /workspace/xgboost_model.json")

    # Wait for tracker to finish (rank 0 only).
    if tracker is not None:
        tracker.wait_for()

    print(f"Worker {rank}: Training complete!")

Scale XGBoost with Kubeflow TrainJob

You can use TrainerClient() from the Kubeflow SDK to communicate with Kubeflow Trainer APIs and scale your training function across multiple XGBoost worker nodes.

TrainerClient() verifies that you have required access to the Kubernetes cluster.

Kubeflow Trainer creates a TrainJob resource and automatically sets the DMLC_* environment variables to set up XGBoost distributed training.

In [ ]:

from kubeflow.trainer import CustomTrainer, TrainerClient

client = TrainerClient()

List the Training Runtimes

You can get the list of available Training Runtimes to start your TrainJob.

Additionally, it might show available accelerator type and number of available resources.

In [ ]:

for runtime in client.list_runtimes():
    print(runtime)

Run the Distributed TrainJob

Kubeflow TrainJob will train the above model on 3 XGBoost worker nodes.

Each worker generates its own shard of synthetic data and participates in distributed gradient synchronization via the Rabit protocol.

In [ ]:

#parameters
num_cpu=3
num_gpu=0
num_nodes=3

In [ ]:

resources_per_node = {
    "cpu": num_cpu,
}
if num_gpu > 0:
    resources_per_node["gpu"] = num_gpu

job_name = client.train(
    trainer=CustomTrainer(
        func=xgboost_train_classification,
        # Set how many XGBoost workers you want for distributed training.
        num_nodes=num_nodes,
        resources_per_node=resources_per_node,
    ),
    runtime="xgboost-distributed",
)

print(f"Training job {job_name} submitted with {num_cpu} CPU and {num_gpu} GPU")

Check the TrainJob Steps

You can check the components of TrainJob that's created.

Since the TrainJob performs distributed training across 3 nodes, it generates 3 steps: trainer-node-0 .. trainer-node-2.

You can get the individual status for each of these steps.

In [ ]:

# Wait for the running status.
client.wait_for_job_status(name=job_name, status={"Running"})

In [ ]:

for c in client.get_job(name=job_name).steps:
    print(f"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\n")

Watch the TrainJob Logs

We can use the get_job_logs() API to get the TrainJob logs.

Since we run training on 3 nodes, every XGBoost worker generates its own data shard and participates in distributed training.

In [ ]:

for i in range(num_nodes):
    print(f"\n** Distributed XGBoost env on node-{i} **")
    print(f"=========================================")
    print("\n".join(TrainerClient().get_job_logs(name=job_name, follow=True, step=f"node-{i}")))

Verify TrainJob Completion

Check the TrainJob status to ensure it completed successfully.

In [ ]:

client.wait_for_job_status(name=job_name, timeout=20)

Delete the TrainJob

When TrainJob is finished, you can delete the resource.

In [ ]:

#client.delete_job(job_name)