Polyaxon & XGBoost Experiment Tracking

How to use Polyaxon and XGBoost Experiment Tracking together

+

Track XGBoost experiments with Polyaxon's PolyaxonCallback: record metrics at each boosting round, save the trained model, and compare runs alongside their training parameters. You can use the callback in a Polyaxon job or in a training script running elsewhere.

Tracking API

The callback reports XGBoost evaluation metrics to the active Polyaxon run. Use the tracking API alongside it for your own parameters, evaluation results, dataset references, and artifacts.

For example, compare validation RMSE across different learning rates and tree depths, then inspect the corresponding model artifact before choosing a run to deploy.

Setup

Install the Polyaxon library and the dependencies used by the example:

pip install polyaxon xgboost scikit-learn

Use a Polyaxon SDK version that includes PolyaxonCallback. If the import is unavailable in an older environment, update to a version compatible with your deployment. The example uses XGBoost's class-based training callback API.

Initialize your script with Polyaxon

Inside a Polyaxon run, tracking.init() uses the injected run context. Outside the cluster, first configure authentication and select your project using the Python client setup guide.

For an API-free local example, replace tracking.init() in the script below with:

from polyaxon import tracking

tracking.init(project="local/xgboost", is_new=True, is_offline=True)

Offline tracking writes locally; it does not make the run appear in the online dashboard automatically. See offline mode for persistence and resuming a local run. The environment-variable equivalent is POLYAXON_IS_OFFLINE=true, not POLYAXON_OFFLINE.

XGBoost callback

Pass a new PolyaxonCallback instance in the callbacks argument to xgboost.train. Supply named evaluation datasets through evals so the callback can record their metrics. With the example below, the run receives train-mae, train-rmse, validation-mae, and validation-rmse at each boosting round.

The older function-style polyaxon_callback exists for legacy XGBoost environments. Use the class callback for the current training API, and create a fresh callback for each training invocation rather than sharing its state across runs.

Customizing the callback

• log_model=True saves and records the trained model after training.
• log_importance=False skips feature-importance plotting. The example uses this setting to keep visualization dependencies out of the setup.
• run=... associates the callback with an explicitly created Run instead of the active tracking run.

If you enable log_importance=True, install the plotting dependencies used by the tracking library, including Matplotlib and Plotly. max_num_features limits the number of features shown in the importance chart.

Manual logging

Use tracking.log_inputs() for the training configuration and tracking.log_metrics() for additional numeric results.

For learning curves, include a step value. For a final evaluation result, use a distinct name so it is not confused with the validation metrics recorded during training. See tracking metadata and events for other logging methods.

Example

This script generates a regression dataset locally, holds out a validation split, and logs the training configuration, learning curves, and model. It does not require a dataset download. Save it as train_xgboost.py and configure a run context as described above before running it.

import xgboost as xgb
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split

from polyaxon import tracking
from polyaxon.tracking.contrib.xgboost import PolyaxonCallback


def main():
    tracking.init()

    features, targets = make_regression(
        n_samples=1000,
        n_features=12,
        n_informative=8,
        noise=5.0,
        random_state=42,
    )
    x_train, x_validation, y_train, y_validation = train_test_split(
        features, targets, test_size=0.2, random_state=42
    )
    dtrain = xgb.DMatrix(x_train, label=y_train)
    dvalidation = xgb.DMatrix(x_validation, label=y_validation)

    params = {
        "objective": "reg:squarederror",
        "eval_metric": ["mae", "rmse"],
        "max_depth": 5,
        "eta": 0.1,
        "subsample": 0.8,
        "seed": 42,
    }
    num_boost_round = 50
    tracking.log_inputs(
        **params,
        num_boost_round=num_boost_round,
        dataset="synthetic-regression",
        dataset_seed=42,
        validation_fraction=0.2,
    )

    xgb.train(
        params=params,
        dtrain=dtrain,
        num_boost_round=num_boost_round,
        evals=[(dtrain, "train"), (dvalidation, "validation")],
        callbacks=[PolyaxonCallback(log_model=True, log_importance=False)],
    )
    tracking.end()


if __name__ == "__main__":
    main()

After an online run completes, open its metrics and artifacts in Polyaxon. Repeat with a different eta or max_depth in a new run to compare validation curves. Synthetic data demonstrates the tracking workflow; it is not a benchmark for selecting a production model.

To move from individual experiments to managed workloads, continue with the job runtime reference and experiment tracking overview.