@@ -0,0 +1,2951 @@
+%22%22%22Example of using callbacks in Dask%22%22%22%0Aimport numpy as np%0Aimport xgboost as xgb%0Afrom xgboost.dask import DaskDMatrix%0Afrom dask.distributed import Client%0Afrom dask.distributed import LocalCluster%0Afrom dask_ml.datasets import make_regression%0Afrom dask_ml.model_selection import train_test_split%0A%0A%0Adef probability_for_going_backward(epoch):%0A    return 0.999 / (1.0 + 0.05 * np.log(1.0 + epoch))%0A%0A%0A# All callback functions must inherit from TrainingCallback%0Aclass CustomEarlyStopping(xgb.callback.TrainingCallback):%0A    %22%22%22A custom early stopping class where early stopping is determined stochastically.%0A    In the beginning, allow the metric to become worse with a probability of 0.999.%0A    As boosting progresses, the probability should be adjusted downward%22%22%22%0A%0A    def __init__(self, *, validation_set, target_metric, maximize, seed):%0A        self.validation_set = validation_set%0A        self.target_metric = target_metric%0A        self.maximize = maximize%0A        self.seed = seed%0A        self.rng = np.random.default_rng(seed=seed)%0A        if maximize:%0A            self.better = lambda x, y: x %3E y%0A        else:%0A            self.better = lambda x, y: x %3C y%0A%0A    def after_iteration(self, model, epoch, evals_log):%0A        metric_history = evals_log%5Bself.validation_set%5D%5Bself.target_metric%5D%0A        if len(metric_history) %3C 2 or self.better(%0A            metric_history%5B-1%5D, metric_history%5B-2%5D%0A        ):%0A            return False  # continue training%0A        p = probability_for_going_backward(epoch)%0A        go_backward = self.rng.choice(2, size=(1,), replace=True, p=%5B1 - p, p%5D).astype(%0A            np.bool%0A        )%5B0%5D%0A        print(%0A            %22The validation metric went into the wrong direction. %22%0A            + f%22Stopping training with probability %7B1 - p%7D...%22%0A        )%0A        if go_backward:%0A            return False  # continue training%0A        else:%0A            return True  # stop training%0A%0A%0Adef main(client):%0A    m = 100000%0A    n = 100%0A    X, y = make_regression(n_samples=m, n_features=n, chunks=200, random_state=0)%0A    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)%0A%0A    dtrain = DaskDMatrix(client, X_train, y_train)%0A    dtest = DaskDMatrix(client, X_test, y_test)%0A%0A    output = xgb.dask.train(%0A        client,%0A        %7B%0A            %22verbosity%22: 1,%0A            %22tree_method%22: %22hist%22,%0A            %22objective%22: %22reg:squarederror%22,%0A            %22eval_metric%22: %22rmse%22,%0A            %22max_depth%22: 6,%0A            %22learning_rate%22: 1.0,%0A        %7D,%0A        dtrain,%0A        num_boost_round=1000,%0A        evals=%5B(dtrain, %22train%22), (dtest, %22test%22)%5D,%0A        callbacks=%5B%0A            CustomEarlyStopping(%0A                validation_set=%22test%22, target_metric=%22rmse%22, maximize=False, seed=0%0A            )%0A        %5D,%0A    )%0A%0A%0Aif __name__ == %22__main__%22:%0A    # or use other clusters for scaling%0A    with LocalCluster(n_workers=4, threads_per_worker=1) as cluster:%0A        with Client(cluster) as client:%0A            main(client)%0A

Add example of using custom callback in Dask (#6995)