chore: import upstream snapshot with attribution

python/ray/tune/examples/xgboost_dynamic_resources_example.py

@@ -0,0 +1,187 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+import sklearn.datasets
+import sklearn.metrics
+import xgboost as xgb
+from sklearn.model_selection import train_test_split
+
+import ray
+from ray import tune
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+from ray.tune.experiment import Trial
+from ray.tune.integration.xgboost import TuneReportCheckpointCallback
+from ray.tune.schedulers import ASHAScheduler, ResourceChangingScheduler
+
+if TYPE_CHECKING:
+    from ray.tune.execution.tune_controller import TuneController
+
+CHECKPOINT_FILENAME = "booster-checkpoint.json"
+
+
+def get_best_model_checkpoint(best_result: "ray.tune.Result"):
+    best_bst = TuneReportCheckpointCallback.get_model(
+        best_result.checkpoint, filename=CHECKPOINT_FILENAME
+    )
+
+    accuracy = 1.0 - best_result.metrics["eval-logloss"]
+    print(f"Best model parameters: {best_result.config}")
+    print(f"Best model total accuracy: {accuracy:.4f}")
+    return best_bst
+
+
+# our train function needs to be able to checkpoint
+# to work with ResourceChangingScheduler
+def train_breast_cancer(config: dict):
+    # This is a simple training function to be passed into Tune
+    # Load dataset
+    data, labels = sklearn.datasets.load_breast_cancer(return_X_y=True)
+    # Split into train and test set
+    train_x, test_x, train_y, test_y = train_test_split(data, labels, test_size=0.25)
+    # Build input matrices for XGBoost
+    train_set = xgb.DMatrix(train_x, label=train_y)
+    test_set = xgb.DMatrix(test_x, label=test_y)
+
+    # Checkpointing needs to be set up in order for dynamic
+    # resource allocation to work as intended
+    xgb_model = None
+    checkpoint = tune.get_checkpoint()
+    if checkpoint:
+        xgb_model = TuneReportCheckpointCallback.get_model(
+            checkpoint, filename=CHECKPOINT_FILENAME
+        )
+
+    # Set `nthread` to the number of CPUs available to the trial,
+    # which is assigned by the scheduler.
+    config["nthread"] = int(tune.get_context().get_trial_resources().head_cpus)
+    print(f"nthreads: {config['nthread']} xgb_model: {xgb_model}")
+    # Train the classifier, using the Tune callback
+    xgb.train(
+        config,
+        train_set,
+        evals=[(test_set, "eval")],
+        verbose_eval=False,
+        xgb_model=xgb_model,
+        callbacks=[
+            TuneReportCheckpointCallback(
+                # checkpointing should happen every iteration
+                # with dynamic resource allocation
+                frequency=1,
+                filename=CHECKPOINT_FILENAME,
+            )
+        ],
+    )
+
+
+def tune_xgboost():
+    search_space = {
+        # You can mix constants with search space objects.
+        "objective": "binary:logistic",
+        "eval_metric": ["logloss", "error"],
+        "max_depth": 9,
+        "learning_rate": 1,
+        "min_child_weight": tune.grid_search([2, 3]),
+        "subsample": tune.grid_search([0.8, 0.9]),
+        "colsample_bynode": tune.grid_search([0.8, 0.9]),
+        "random_state": 1,
+        "num_parallel_tree": 2000,
+    }
+    # This will enable aggressive early stopping of bad trials.
+    base_scheduler = ASHAScheduler(
+        max_t=16, grace_period=1, reduction_factor=2  # 16 training iterations
+    )
+
+    def example_resources_allocation_function(
+        tune_controller: "TuneController",
+        trial: Trial,
+        result: Dict[str, Any],
+        scheduler: "ResourceChangingScheduler",
+    ) -> Optional[PlacementGroupFactory]:
+        """This is a basic example of a resource allocating function.
+
+        The function naively balances available CPUs over live trials.
+
+        This function returns a new ``PlacementGroupFactory`` with updated
+        resource requirements, or None. If the returned
+        ``PlacementGroupFactory`` is equal by value to the one the
+        trial has currently, the scheduler will skip the update process
+        internally (same with None).
+
+        See :class:`DistributeResources` for a more complex,
+        robust approach.
+
+        Args:
+            tune_controller: Trial runner for this Tune run.
+                Can be used to obtain information about other trials.
+            trial: The trial to allocate new resources to.
+            result: The latest results of trial.
+            scheduler: The scheduler calling the function.
+
+        Returns:
+            A new ``PlacementGroupFactory`` with the updated resource
+            requirements, or ``None`` to leave the trial's resources unchanged.
+        """
+
+        # Get base trial resources as defined in
+        # ``tune.with_resources``
+        base_trial_resource = scheduler._base_trial_resources
+
+        # Don't bother if this is just the first iteration
+        if result["training_iteration"] < 1:
+            return None
+
+        # default values if resources_per_trial is unspecified
+        if base_trial_resource is None:
+            base_trial_resource = PlacementGroupFactory([{"CPU": 1, "GPU": 0}])
+
+        # Assume that the number of CPUs cannot go below what was
+        # specified in ``Tuner.fit()``.
+        min_cpu = base_trial_resource.required_resources.get("CPU", 0)
+
+        # Get the number of CPUs available in total (not just free)
+        total_available_cpus = tune_controller._resource_updater.get_num_cpus()
+
+        # Divide the free CPUs among all live trials
+        cpu_to_use = max(
+            min_cpu, total_available_cpus // len(tune_controller.get_live_trials())
+        )
+
+        # Assign new CPUs to the trial in a PlacementGroupFactory
+        return PlacementGroupFactory([{"CPU": cpu_to_use, "GPU": 0}])
+
+    # You can either define your own resources_allocation_function, or
+    # use the default one - DistributeResources
+
+    # from ray.tune.schedulers.resource_changing_scheduler import \
+    #    DistributeResources
+
+    scheduler = ResourceChangingScheduler(
+        base_scheduler=base_scheduler,
+        resources_allocation_function=example_resources_allocation_function,
+        # resources_allocation_function=DistributeResources()  # default
+    )
+
+    tuner = tune.Tuner(
+        tune.with_resources(
+            train_breast_cancer, resources=PlacementGroupFactory([{"CPU": 1, "GPU": 0}])
+        ),
+        tune_config=tune.TuneConfig(
+            metric="eval-logloss",
+            mode="min",
+            num_samples=1,
+            scheduler=scheduler,
+        ),
+        param_space=search_space,
+    )
+    results = tuner.fit()
+
+    return results.get_best_result()
+
+
+if __name__ == "__main__":
+    ray.init(num_cpus=8)
+
+    best_result = tune_xgboost()
+    best_bst = get_best_model_checkpoint(best_result)
+
+    # You could now do further predictions with
+    # best_bst.predict(...)