ray-project--ray/python/ray/train/v2/lightgbm/lightgbm_trainer.py

import logging
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union

import ray.train
from ray.train import Checkpoint
from ray.train.trainer import GenDataset
from ray.train.v2.api.config import RunConfig, ScalingConfig
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
from ray.train.v2.api.validation_config import ValidationConfig
from ray.util.annotations import Deprecated

if TYPE_CHECKING:
    from ray.train.lightgbm import LightGBMConfig

logger = logging.getLogger(__name__)


class LightGBMTrainer(DataParallelTrainer):
    """A Trainer for distributed data-parallel LightGBM training.

    Example:

        .. testcode::
            :skipif: True

            import lightgbm as lgb

            import ray.data
            import ray.train
            from ray.train.lightgbm import (
                LightGBMTrainer,
                RayTrainReportCallback,
                normalize_pandas_for_lightgbm,
            )


            def train_fn_per_worker(config: dict):
                # (Optional) Add logic to resume training state from a checkpoint.
                # ray.train.get_checkpoint()

                # 1. Get the dataset shard for the worker and convert to a `lgb.Dataset`
                train_ds_iter, eval_ds_iter = (
                    ray.train.get_dataset_shard("train"),
                    ray.train.get_dataset_shard("validation"),
                )
                train_ds, eval_ds = train_ds_iter.materialize(), eval_ds_iter.materialize()
                train_df = normalize_pandas_for_lightgbm(train_ds.to_pandas())
                eval_df = normalize_pandas_for_lightgbm(eval_ds.to_pandas())
                train_X, train_y = train_df.drop("y", axis=1), train_df["y"]
                eval_X, eval_y = eval_df.drop("y", axis=1), eval_df["y"]

                train_set = lgb.Dataset(train_X, label=train_y)
                eval_set = lgb.Dataset(eval_X, label=eval_y)

                # 2. Run distributed data-parallel training.
                # `get_network_params` sets up the necessary configurations for LightGBM
                # to set up the data parallel training worker group on your Ray cluster.
                params = {
                    "objective": "regression",
                    # Adding the lines below are the only changes needed
                    # for your `lgb.train` call!
                    "tree_learner": "data_parallel",
                    "pre_partition": True,
                    **ray.train.lightgbm.get_network_params(),
                }
                lgb.train(
                    params,
                    train_set,
                    valid_sets=[eval_set],
                    valid_names=["eval"],
                    num_boost_round=1,
                    # To access the checkpoint from trainer, you need this callback.
                    callbacks=[RayTrainReportCallback()],
                )

            train_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
            eval_ds = ray.data.from_items(
                [{"x": x, "y": x + 1} for x in range(32, 32 + 16)]
            )
            trainer = LightGBMTrainer(
                train_fn_per_worker,
                datasets={"train": train_ds, "validation": eval_ds},
                scaling_config=ray.train.ScalingConfig(num_workers=2),
            )
            result = trainer.fit()
            booster = RayTrainReportCallback.get_model(result.checkpoint)

    Args:
        train_loop_per_worker: The training function to execute on each worker.
            This function can either take in zero arguments or a single ``Dict``
            argument which is set by defining ``train_loop_config``.
            Within this function you can use any of the
            :ref:`Ray Train Loop utilities <train-loop-api>`.
        train_loop_config: A configuration ``Dict`` to pass in as an argument to
            ``train_loop_per_worker``.
            This is typically used for specifying hyperparameters.
        lightgbm_config: The configuration for setting up the distributed lightgbm
            backend. See :class:`~ray.train.lightgbm.LightGBMConfig` for more info.
        scaling_config: The configuration for how to scale data parallel training.
            ``num_workers`` determines how many Python processes are used for training,
            and ``use_gpu`` determines whether or not each process should use GPUs.
            See :class:`~ray.train.ScalingConfig` for more info.
        run_config: The configuration for the execution of the training run.
            See :class:`~ray.train.RunConfig` for more info.
        datasets: The Ray Datasets to ingest for training.
            Datasets are keyed by name (``{name: dataset}``).
            Each dataset can be accessed from within the ``train_loop_per_worker``
            by calling ``ray.train.get_dataset_shard(name)``.
            Sharding and additional configuration can be done by
            passing in a ``dataset_config``.
        dataset_config: The configuration for ingesting the input ``datasets``.
            By default, all the Ray Dataset are split equally across workers.
            See :class:`~ray.train.DataConfig` for more details.
        validation_config: [Alpha] Configuration for checkpoint validation.
            If provided and ``ray.train.report`` is called with the ``validation``
            argument, Ray Train will validate the reported checkpoint using
            the validation function specified in this config.
        metadata: [Deprecated]
        resume_from_checkpoint: [Deprecated]
        label_column: [Deprecated] Legacy LightGBMTrainer API.
        params: [Deprecated] Legacy LightGBMTrainer API.
        num_boost_round: [Deprecated] Legacy LightGBMTrainer API.
    """

    def __init__(
        self,
        train_loop_per_worker: Union[Callable[[], Any], Callable[[Dict], Any]],
        *,
        train_loop_config: Optional[Dict] = None,
        lightgbm_config: Optional["LightGBMConfig"] = None,
        scaling_config: Optional[ScalingConfig] = None,
        run_config: Optional[RunConfig] = None,
        datasets: Optional[Dict[str, GenDataset]] = None,
        dataset_config: Optional[ray.train.DataConfig] = None,
        validation_config: Optional[ValidationConfig] = None,
        # TODO: [Deprecated]
        metadata: Optional[Dict[str, Any]] = None,
        resume_from_checkpoint: Optional[Checkpoint] = None,
        # TODO: [Deprecated] Legacy LightGBMTrainer API
        label_column: Optional[str] = None,
        params: Optional[Dict[str, Any]] = None,
        num_boost_round: Optional[int] = None,
    ):
        if (
            label_column is not None
            or params is not None
            or num_boost_round is not None
        ):
            raise DeprecationWarning(
                "The legacy LightGBMTrainer API is deprecated. "
                "Please switch to passing in a custom `train_loop_per_worker` "
                "function instead. "
                "See this issue for more context: "
                "https://github.com/ray-project/ray/issues/50042"
            )
        from ray.train.lightgbm import LightGBMConfig

        super(LightGBMTrainer, self).__init__(
            train_loop_per_worker=train_loop_per_worker,
            train_loop_config=train_loop_config,
            backend_config=lightgbm_config or LightGBMConfig(),
            scaling_config=scaling_config,
            dataset_config=dataset_config,
            run_config=run_config,
            datasets=datasets,
            resume_from_checkpoint=resume_from_checkpoint,
            metadata=metadata,
            validation_config=validation_config,
        )

    @classmethod
    @Deprecated
    def get_model(
        cls,
        checkpoint: Checkpoint,
    ):
        """Retrieve the LightGBM model stored in this checkpoint.

        This API is deprecated. Use `RayTrainReportCallback.get_model` instead.
        """
        raise DeprecationWarning(
            "`LightGBMTrainer.get_model` is deprecated. "
            "Use `RayTrainReportCallback.get_model` instead."
        )