chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,23 @@
|
||||
from ray.train.lightgbm._lightgbm_utils import (
|
||||
RayTrainReportCallback,
|
||||
normalize_pandas_for_lightgbm,
|
||||
)
|
||||
from ray.train.lightgbm.config import LightGBMConfig, get_network_params
|
||||
from ray.train.lightgbm.lightgbm_checkpoint import LightGBMCheckpoint
|
||||
from ray.train.lightgbm.lightgbm_trainer import LightGBMTrainer
|
||||
from ray.train.v2._internal.constants import is_v2_enabled
|
||||
|
||||
if is_v2_enabled():
|
||||
from ray.train.v2.lightgbm.lightgbm_trainer import LightGBMTrainer # noqa: F811
|
||||
|
||||
__all__ = [
|
||||
"RayTrainReportCallback",
|
||||
"LightGBMCheckpoint",
|
||||
"LightGBMTrainer",
|
||||
"LightGBMConfig",
|
||||
"get_network_params",
|
||||
"normalize_pandas_for_lightgbm",
|
||||
]
|
||||
|
||||
|
||||
# DO NOT ADD ANYTHING AFTER THIS LINE.
|
||||
@@ -0,0 +1,256 @@
|
||||
import tempfile
|
||||
from abc import abstractmethod
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
|
||||
|
||||
from lightgbm.basic import Booster
|
||||
from lightgbm.callback import CallbackEnv
|
||||
|
||||
import ray.train
|
||||
from ray.train import Checkpoint
|
||||
from ray.tune.utils import flatten_dict
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pandas as pd
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
def normalize_pandas_for_lightgbm(df: "pd.DataFrame") -> "pd.DataFrame":
|
||||
"""Map Arrow-backed pandas dtypes to NumPy-nullable equivalents.
|
||||
|
||||
LightGBM's pandas input validation rejects Arrow-backed dtypes like
|
||||
``int64[pyarrow]``. Since Ray Data 2.56, ``Dataset.to_pandas()`` preserves
|
||||
Arrow-backed dtypes when the source was Arrow, so callers passing the
|
||||
resulting frame to ``lightgbm.Dataset`` must normalize first.
|
||||
|
||||
This helper is a faster alternative to
|
||||
``df.convert_dtypes(dtype_backend="numpy_nullable")``:
|
||||
|
||||
- It maps dtypes mechanically rather than scanning every value.
|
||||
- It only touches ``pd.ArrowDtype`` columns. NumPy-backed columns (e.g.
|
||||
from ``ray.data.from_pandas`` shards) keep their original buffers.
|
||||
|
||||
Only numeric and boolean Arrow dtypes are remapped. Other Arrow dtypes
|
||||
(string, decimal, timestamp) are left as-is; LightGBM doesn't accept them
|
||||
as features anyway.
|
||||
|
||||
Args:
|
||||
df: The pandas DataFrame to normalize.
|
||||
|
||||
Returns:
|
||||
A DataFrame with Arrow-backed numeric/boolean columns replaced by
|
||||
NumPy-nullable equivalents. Other columns are returned unchanged.
|
||||
"""
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
|
||||
dtype_mapping = {}
|
||||
for column, dtype in df.dtypes.items():
|
||||
if not isinstance(dtype, pd.ArrowDtype):
|
||||
continue
|
||||
arrow_dtype = dtype.pyarrow_dtype
|
||||
if pa.types.is_signed_integer(arrow_dtype):
|
||||
dtype_mapping[column] = f"Int{arrow_dtype.bit_width}"
|
||||
elif pa.types.is_unsigned_integer(arrow_dtype):
|
||||
dtype_mapping[column] = f"UInt{arrow_dtype.bit_width}"
|
||||
elif pa.types.is_floating(arrow_dtype):
|
||||
dtype_mapping[column] = f"Float{arrow_dtype.bit_width}"
|
||||
elif pa.types.is_boolean(arrow_dtype):
|
||||
dtype_mapping[column] = "boolean"
|
||||
if dtype_mapping:
|
||||
df = df.astype(dtype_mapping, copy=False)
|
||||
return df
|
||||
|
||||
|
||||
class RayReportCallback:
|
||||
CHECKPOINT_NAME = "model.txt"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
|
||||
filename: str = CHECKPOINT_NAME,
|
||||
frequency: int = 0,
|
||||
checkpoint_at_end: bool = True,
|
||||
results_postprocessing_fn: Optional[
|
||||
Callable[[Dict[str, Union[float, List[float]]]], Dict[str, float]]
|
||||
] = None,
|
||||
):
|
||||
if isinstance(metrics, str):
|
||||
metrics = [metrics]
|
||||
self._metrics = metrics
|
||||
self._filename = filename
|
||||
self._frequency = frequency
|
||||
self._checkpoint_at_end = checkpoint_at_end
|
||||
self._results_postprocessing_fn = results_postprocessing_fn
|
||||
|
||||
@classmethod
|
||||
def get_model(
|
||||
cls, checkpoint: Checkpoint, filename: str = CHECKPOINT_NAME
|
||||
) -> Booster:
|
||||
"""Retrieve the model stored in a checkpoint reported by this callback.
|
||||
|
||||
Args:
|
||||
checkpoint: The checkpoint object returned by a training run.
|
||||
The checkpoint should be saved by an instance of this callback.
|
||||
filename: The filename to load the model from, which should match
|
||||
the filename used when creating the callback.
|
||||
Returns:
|
||||
The model loaded from the checkpoint.
|
||||
"""
|
||||
with checkpoint.as_directory() as checkpoint_path:
|
||||
return Booster(model_file=Path(checkpoint_path, filename).as_posix())
|
||||
|
||||
def _get_report_dict(self, evals_log: Dict[str, Dict[str, list]]) -> dict:
|
||||
result_dict = flatten_dict(evals_log, delimiter="-")
|
||||
if not self._metrics:
|
||||
report_dict = result_dict
|
||||
else:
|
||||
report_dict = {}
|
||||
for key in self._metrics:
|
||||
if isinstance(self._metrics, dict):
|
||||
metric = self._metrics[key]
|
||||
else:
|
||||
metric = key
|
||||
report_dict[key] = result_dict[metric]
|
||||
if self._results_postprocessing_fn:
|
||||
report_dict = self._results_postprocessing_fn(report_dict)
|
||||
return report_dict
|
||||
|
||||
def _get_eval_result(self, env: CallbackEnv) -> dict:
|
||||
eval_result = {}
|
||||
for entry in env.evaluation_result_list:
|
||||
data_name, eval_name, result = entry[0:3]
|
||||
if len(entry) > 4:
|
||||
stdv = entry[4]
|
||||
suffix = "-mean"
|
||||
else:
|
||||
stdv = None
|
||||
suffix = ""
|
||||
if data_name not in eval_result:
|
||||
eval_result[data_name] = {}
|
||||
eval_result[data_name][eval_name + suffix] = result
|
||||
if stdv is not None:
|
||||
eval_result[data_name][eval_name + "-stdv"] = stdv
|
||||
return eval_result
|
||||
|
||||
@abstractmethod
|
||||
def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]:
|
||||
"""Get checkpoint from model.
|
||||
|
||||
This method needs to be implemented by subclasses.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster):
|
||||
"""Save checkpoint and report metrics corresonding to this checkpoint.
|
||||
|
||||
This method needs to be implemented by subclasses.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def _report_metrics(self, report_dict: Dict):
|
||||
"""Report Metrics.
|
||||
|
||||
This method needs to be implemented by subclasses.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def __call__(self, env: CallbackEnv) -> None:
|
||||
eval_result = self._get_eval_result(env)
|
||||
report_dict = self._get_report_dict(eval_result)
|
||||
|
||||
# Ex: if frequency=2, checkpoint_at_end=True and num_boost_rounds=11,
|
||||
# you will checkpoint at iterations 1, 3, 5, ..., 9, and 10 (checkpoint_at_end)
|
||||
# (iterations count from 0)
|
||||
on_last_iter = env.iteration == env.end_iteration - 1
|
||||
should_checkpoint_at_end = on_last_iter and self._checkpoint_at_end
|
||||
should_checkpoint_with_frequency = (
|
||||
self._frequency != 0 and (env.iteration + 1) % self._frequency == 0
|
||||
)
|
||||
should_checkpoint = should_checkpoint_at_end or should_checkpoint_with_frequency
|
||||
|
||||
if should_checkpoint:
|
||||
self._save_and_report_checkpoint(report_dict, env.model)
|
||||
else:
|
||||
self._report_metrics(report_dict)
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class RayTrainReportCallback(RayReportCallback):
|
||||
"""Creates a callback that reports metrics and checkpoints model.
|
||||
|
||||
Args:
|
||||
metrics: Metrics to report. If this is a list,
|
||||
each item should be a metric key reported by LightGBM,
|
||||
and it will be reported to Ray Train/Tune under the same name.
|
||||
This can also be a dict of {<key-to-report>: <lightgbm-metric-key>},
|
||||
which can be used to rename LightGBM default metrics.
|
||||
filename: Customize the saved checkpoint file type by passing
|
||||
a filename. Defaults to "model.txt".
|
||||
frequency: How often to save checkpoints, in terms of iterations.
|
||||
Defaults to 0 (no checkpoints are saved during training).
|
||||
checkpoint_at_end: Whether or not to save a checkpoint at the end of training.
|
||||
results_postprocessing_fn: An optional Callable that takes in
|
||||
the metrics dict that will be reported (after it has been flattened)
|
||||
and returns a modified dict.
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
Reporting checkpoints and metrics to Ray Tune when running many
|
||||
independent LightGBM trials (without data parallelism within a trial).
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import lightgbm
|
||||
|
||||
from ray.train.lightgbm import RayTrainReportCallback
|
||||
|
||||
config = {
|
||||
# ...
|
||||
"metric": ["binary_logloss", "binary_error"],
|
||||
}
|
||||
|
||||
# Report only log loss to Tune after each validation epoch.
|
||||
bst = lightgbm.train(
|
||||
...,
|
||||
callbacks=[
|
||||
RayTrainReportCallback(
|
||||
metrics={"loss": "eval-binary_logloss"}, frequency=1
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
Loading a model from a checkpoint reported by this callback.
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
from ray.train.lightgbm import RayTrainReportCallback
|
||||
|
||||
# Get a `Checkpoint` object that is saved by the callback during training.
|
||||
result = trainer.fit()
|
||||
booster = RayTrainReportCallback.get_model(result.checkpoint)
|
||||
|
||||
"""
|
||||
|
||||
@contextmanager
|
||||
def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]:
|
||||
if ray.train.get_context().get_world_rank() in (0, None):
|
||||
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
||||
model.save_model(Path(temp_checkpoint_dir, self._filename).as_posix())
|
||||
yield Checkpoint.from_directory(temp_checkpoint_dir)
|
||||
else:
|
||||
yield None
|
||||
|
||||
def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster):
|
||||
with self._get_checkpoint(model=model) as checkpoint:
|
||||
ray.train.report(report_dict, checkpoint=checkpoint)
|
||||
|
||||
def _report_metrics(self, report_dict: Dict):
|
||||
ray.train.report(report_dict)
|
||||
@@ -0,0 +1,98 @@
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import ray
|
||||
from ray._common.network_utils import build_address
|
||||
from ray.train._internal.base_worker_group import BaseWorkerGroup
|
||||
from ray.train._internal.utils import get_address_and_port
|
||||
from ray.train.backend import Backend, BackendConfig
|
||||
from ray.train.v2._internal.util import TrainingFramework
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Global LightGBM distributed network configuration for each worker process.
|
||||
_lightgbm_network_params: Optional[Dict[str, Any]] = None
|
||||
_lightgbm_network_params_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_network_params() -> Dict[str, Any]:
|
||||
"""Returns the network parameters to enable LightGBM distributed training."""
|
||||
global _lightgbm_network_params
|
||||
|
||||
with _lightgbm_network_params_lock:
|
||||
if not _lightgbm_network_params:
|
||||
logger.warning(
|
||||
"`ray.train.lightgbm.get_network_params` was called outside "
|
||||
"the context of a `ray.train.lightgbm.LightGBMTrainer`. "
|
||||
"The current process has no knowledge of the distributed training "
|
||||
"worker group, so this method will return an empty dict. "
|
||||
"Please call this within the training loop of a "
|
||||
"`ray.train.lightgbm.LightGBMTrainer`. "
|
||||
"If you are in fact calling this within a `LightGBMTrainer`, "
|
||||
"this is unexpected: please file a bug report to the Ray Team."
|
||||
)
|
||||
return {}
|
||||
|
||||
return _lightgbm_network_params.copy()
|
||||
|
||||
|
||||
def _set_network_params(
|
||||
num_machines: int,
|
||||
local_listen_port: int,
|
||||
machines: str,
|
||||
):
|
||||
global _lightgbm_network_params
|
||||
|
||||
with _lightgbm_network_params_lock:
|
||||
assert (
|
||||
_lightgbm_network_params is None
|
||||
), "LightGBM network params are already initialized."
|
||||
_lightgbm_network_params = dict(
|
||||
num_machines=num_machines,
|
||||
local_listen_port=local_listen_port,
|
||||
machines=machines,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LightGBMConfig(BackendConfig):
|
||||
"""Configuration for LightGBM distributed data-parallel training setup.
|
||||
|
||||
See the LightGBM docs for more information on the "network parameters"
|
||||
that Ray Train sets up for you:
|
||||
https://lightgbm.readthedocs.io/en/latest/Parameters.html#network-parameters
|
||||
"""
|
||||
|
||||
@property
|
||||
def backend_cls(self):
|
||||
return _LightGBMBackend
|
||||
|
||||
@property
|
||||
def framework(self):
|
||||
return TrainingFramework.LIGHTGBM
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {}
|
||||
|
||||
|
||||
class _LightGBMBackend(Backend):
|
||||
def on_training_start(
|
||||
self, worker_group: BaseWorkerGroup, backend_config: LightGBMConfig
|
||||
):
|
||||
node_ips_and_ports = worker_group.execute(get_address_and_port)
|
||||
ports = [port for _, port in node_ips_and_ports]
|
||||
machines = ",".join(
|
||||
[build_address(node_ip, port) for node_ip, port in node_ips_and_ports]
|
||||
)
|
||||
num_machines = len(worker_group)
|
||||
ray.get(
|
||||
[
|
||||
worker_group.execute_single_async(
|
||||
rank, _set_network_params, num_machines, ports[rank], machines
|
||||
)
|
||||
for rank in range(len(worker_group))
|
||||
]
|
||||
)
|
||||
@@ -0,0 +1,72 @@
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import lightgbm
|
||||
|
||||
from ray.train._internal.framework_checkpoint import FrameworkCheckpoint
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.preprocessor import Preprocessor
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class LightGBMCheckpoint(FrameworkCheckpoint):
|
||||
"""A :py:class:`~ray.train.Checkpoint` with LightGBM-specific functionality."""
|
||||
|
||||
MODEL_FILENAME = "model.txt"
|
||||
|
||||
@classmethod
|
||||
def from_model(
|
||||
cls,
|
||||
booster: lightgbm.Booster,
|
||||
*,
|
||||
preprocessor: Optional["Preprocessor"] = None,
|
||||
path: Optional[str] = None,
|
||||
) -> "LightGBMCheckpoint":
|
||||
"""Create a :py:class:`~ray.train.Checkpoint` that stores a LightGBM model.
|
||||
|
||||
Args:
|
||||
booster: The LightGBM model to store in the checkpoint.
|
||||
preprocessor: A fitted preprocessor to be applied before inference.
|
||||
path: The path to the directory where the checkpoint file will be saved.
|
||||
This should start as an empty directory, since the *entire*
|
||||
directory will be treated as the checkpoint when reported.
|
||||
By default, a temporary directory will be created.
|
||||
|
||||
Returns:
|
||||
An :py:class:`LightGBMCheckpoint` containing the specified ``Estimator``.
|
||||
|
||||
Examples:
|
||||
.. testcode::
|
||||
|
||||
import lightgbm
|
||||
import numpy as np
|
||||
from ray.train.lightgbm import LightGBMCheckpoint
|
||||
|
||||
train_X = np.array([[1, 2], [3, 4]])
|
||||
train_y = np.array([0, 1])
|
||||
|
||||
model = lightgbm.LGBMClassifier().fit(train_X, train_y)
|
||||
checkpoint = LightGBMCheckpoint.from_model(model.booster_)
|
||||
"""
|
||||
checkpoint_path = Path(path or tempfile.mkdtemp())
|
||||
|
||||
if not checkpoint_path.is_dir():
|
||||
raise ValueError(f"`path` must be a directory, but got: {checkpoint_path}")
|
||||
|
||||
booster.save_model(checkpoint_path.joinpath(cls.MODEL_FILENAME).as_posix())
|
||||
|
||||
checkpoint = cls.from_directory(checkpoint_path.as_posix())
|
||||
if preprocessor:
|
||||
checkpoint.set_preprocessor(preprocessor)
|
||||
|
||||
return checkpoint
|
||||
|
||||
def get_model(self) -> lightgbm.Booster:
|
||||
"""Retrieve the LightGBM model stored in this checkpoint."""
|
||||
with self.as_directory() as checkpoint_path:
|
||||
return lightgbm.Booster(
|
||||
model_file=Path(checkpoint_path, self.MODEL_FILENAME).as_posix()
|
||||
)
|
||||
@@ -0,0 +1,326 @@
|
||||
import logging
|
||||
from functools import partial
|
||||
from typing import Any, Callable, Dict, Optional, Union
|
||||
|
||||
import lightgbm
|
||||
|
||||
import ray
|
||||
from ray.train import Checkpoint
|
||||
from ray.train.constants import TRAIN_DATASET_KEY
|
||||
from ray.train.lightgbm._lightgbm_utils import (
|
||||
RayTrainReportCallback,
|
||||
normalize_pandas_for_lightgbm,
|
||||
)
|
||||
from ray.train.lightgbm.config import LightGBMConfig
|
||||
from ray.train.lightgbm.v2 import LightGBMTrainer as SimpleLightGBMTrainer
|
||||
from ray.train.trainer import GenDataset
|
||||
from ray.train.utils import _log_deprecation_warning
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
LEGACY_LIGHTGBM_TRAINER_DEPRECATION_MESSAGE = (
|
||||
"Passing in `lightgbm.train` kwargs such as `params`, `num_boost_round`, "
|
||||
"`label_column`, etc. to `LightGBMTrainer` is deprecated "
|
||||
"in favor of the new API which accepts a `train_loop_per_worker` argument, "
|
||||
"similar to the other DataParallelTrainer APIs (ex: TorchTrainer). "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/50042"
|
||||
)
|
||||
|
||||
|
||||
def _lightgbm_train_fn_per_worker(
|
||||
config: dict,
|
||||
label_column: str,
|
||||
num_boost_round: int,
|
||||
dataset_keys: set,
|
||||
lightgbm_train_kwargs: dict,
|
||||
):
|
||||
checkpoint = ray.train.get_checkpoint()
|
||||
starting_model = None
|
||||
remaining_iters = num_boost_round
|
||||
if checkpoint:
|
||||
starting_model = RayTrainReportCallback.get_model(checkpoint)
|
||||
starting_iter = starting_model.current_iteration()
|
||||
remaining_iters = num_boost_round - starting_iter
|
||||
logger.info(
|
||||
f"Model loaded from checkpoint will train for "
|
||||
f"additional {remaining_iters} iterations (trees) in order "
|
||||
"to achieve the target number of iterations "
|
||||
f"({num_boost_round=})."
|
||||
)
|
||||
|
||||
train_ds_iter = ray.train.get_dataset_shard(TRAIN_DATASET_KEY)
|
||||
train_df = normalize_pandas_for_lightgbm(train_ds_iter.materialize().to_pandas())
|
||||
|
||||
eval_ds_iters = {
|
||||
k: ray.train.get_dataset_shard(k)
|
||||
for k in dataset_keys
|
||||
if k != TRAIN_DATASET_KEY
|
||||
}
|
||||
eval_dfs = {
|
||||
k: normalize_pandas_for_lightgbm(d.materialize().to_pandas())
|
||||
for k, d in eval_ds_iters.items()
|
||||
}
|
||||
|
||||
train_X, train_y = train_df.drop(label_column, axis=1), train_df[label_column]
|
||||
train_set = lightgbm.Dataset(train_X, label=train_y)
|
||||
|
||||
# NOTE: Include the training dataset in the evaluation datasets.
|
||||
# This allows `train-*` metrics to be calculated and reported.
|
||||
valid_sets = [train_set]
|
||||
valid_names = [TRAIN_DATASET_KEY]
|
||||
|
||||
for eval_name, eval_df in eval_dfs.items():
|
||||
eval_X, eval_y = eval_df.drop(label_column, axis=1), eval_df[label_column]
|
||||
valid_sets.append(lightgbm.Dataset(eval_X, label=eval_y))
|
||||
valid_names.append(eval_name)
|
||||
|
||||
# Add network params of the worker group to enable distributed training.
|
||||
config.update(ray.train.lightgbm.get_network_params())
|
||||
config.setdefault("tree_learner", "data_parallel")
|
||||
config.setdefault("pre_partition", True)
|
||||
|
||||
lightgbm.train(
|
||||
params=config,
|
||||
train_set=train_set,
|
||||
num_boost_round=remaining_iters,
|
||||
valid_sets=valid_sets,
|
||||
valid_names=valid_names,
|
||||
init_model=starting_model,
|
||||
**lightgbm_train_kwargs,
|
||||
)
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class LightGBMTrainer(SimpleLightGBMTrainer):
|
||||
"""A Trainer for distributed data-parallel LightGBM training.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import lightgbm
|
||||
|
||||
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 `lightgbm.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"]
|
||||
dtrain = lightgbm.Dataset(train_X, label=train_y)
|
||||
deval = lightgbm.Dataset(eval_X, label=eval_y)
|
||||
|
||||
params = {
|
||||
"objective": "regression",
|
||||
"metric": "l2",
|
||||
"learning_rate": 1e-4,
|
||||
"subsample": 0.5,
|
||||
"max_depth": 2,
|
||||
# Adding the line below is the only change needed
|
||||
# for your `lgb.train` call!
|
||||
**ray.train.lightgbm.get_network_params(),
|
||||
}
|
||||
|
||||
# 2. Do distributed data-parallel training.
|
||||
# Ray Train sets up the necessary coordinator processes and
|
||||
# environment variables for your workers to communicate with each other.
|
||||
bst = lightgbm.train(
|
||||
params,
|
||||
train_set=dtrain,
|
||||
valid_sets=[deval],
|
||||
valid_names=["validation"],
|
||||
num_boost_round=10,
|
||||
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(16)])
|
||||
trainer = LightGBMTrainer(
|
||||
train_fn_per_worker,
|
||||
datasets={"train": train_ds, "validation": eval_ds},
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=4),
|
||||
)
|
||||
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. Defaults to using the "rabit" 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 use for training and validation.
|
||||
dataset_config: The configuration for ingesting the input ``datasets``.
|
||||
By default, all the Ray Datasets are split equally across workers.
|
||||
See :class:`~ray.train.DataConfig` for more details.
|
||||
resume_from_checkpoint: A checkpoint to resume training from.
|
||||
This checkpoint can be accessed from within ``train_loop_per_worker``
|
||||
by calling ``ray.train.get_checkpoint()``.
|
||||
metadata: Dict that should be made available via
|
||||
`ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
|
||||
for checkpoints saved from this Trainer. Must be JSON-serializable.
|
||||
label_column: [Deprecated] Name of the label column. A column with this name
|
||||
must be present in the training dataset.
|
||||
params: [Deprecated] LightGBM training parameters.
|
||||
Refer to `LightGBM documentation <https://lightgbm.readthedocs.io/>`_
|
||||
for a list of possible parameters.
|
||||
num_boost_round: [Deprecated] Target number of boosting iterations (trees in the model).
|
||||
Note that unlike in ``lightgbm.train``, this is the target number
|
||||
of trees, meaning that if you set ``num_boost_round=10`` and pass a model
|
||||
that has already been trained for 5 iterations, it will be trained for 5
|
||||
iterations more, instead of 10 more.
|
||||
**train_kwargs: [Deprecated] Additional kwargs passed to ``lightgbm.train()`` function.
|
||||
"""
|
||||
|
||||
_handles_checkpoint_freq = True
|
||||
_handles_checkpoint_at_end = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_loop_per_worker: Optional[
|
||||
Union[Callable[[], None], Callable[[Dict], None]]
|
||||
] = None,
|
||||
*,
|
||||
train_loop_config: Optional[Dict] = None,
|
||||
lightgbm_config: Optional[LightGBMConfig] = None,
|
||||
scaling_config: Optional[ray.train.ScalingConfig] = None,
|
||||
run_config: Optional[ray.train.RunConfig] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
dataset_config: Optional[ray.train.DataConfig] = None,
|
||||
resume_from_checkpoint: Optional[Checkpoint] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
# TODO: [Deprecated] Legacy LightGBMTrainer API
|
||||
label_column: Optional[str] = None,
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
num_boost_round: Optional[int] = None,
|
||||
**train_kwargs,
|
||||
):
|
||||
# TODO: [Deprecated] Legacy LightGBMTrainer API
|
||||
legacy_api = train_loop_per_worker is None
|
||||
if legacy_api:
|
||||
train_loop_per_worker = self._get_legacy_train_fn_per_worker(
|
||||
lightgbm_train_kwargs=train_kwargs,
|
||||
run_config=run_config,
|
||||
label_column=label_column,
|
||||
num_boost_round=num_boost_round,
|
||||
datasets=datasets,
|
||||
)
|
||||
train_loop_config = params or {}
|
||||
elif train_kwargs:
|
||||
_log_deprecation_warning(
|
||||
"Passing `lightgbm.train` kwargs to `LightGBMTrainer` is deprecated. "
|
||||
f"Got kwargs: {train_kwargs.keys()}\n"
|
||||
"In your training function, you can call `lightgbm.train(**kwargs)` "
|
||||
"with arbitrary arguments. "
|
||||
f"{LEGACY_LIGHTGBM_TRAINER_DEPRECATION_MESSAGE}"
|
||||
)
|
||||
|
||||
super(LightGBMTrainer, self).__init__(
|
||||
train_loop_per_worker=train_loop_per_worker,
|
||||
train_loop_config=train_loop_config,
|
||||
lightgbm_config=lightgbm_config,
|
||||
scaling_config=scaling_config,
|
||||
run_config=run_config,
|
||||
datasets=datasets,
|
||||
dataset_config=dataset_config,
|
||||
resume_from_checkpoint=resume_from_checkpoint,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
def _get_legacy_train_fn_per_worker(
|
||||
self,
|
||||
lightgbm_train_kwargs: Dict,
|
||||
run_config: Optional[ray.train.RunConfig],
|
||||
datasets: Optional[Dict[str, GenDataset]],
|
||||
label_column: Optional[str],
|
||||
num_boost_round: Optional[int],
|
||||
) -> Callable[[Dict], None]:
|
||||
"""Get the training function for the legacy LightGBMTrainer API."""
|
||||
|
||||
datasets = datasets or {}
|
||||
if not datasets.get(TRAIN_DATASET_KEY):
|
||||
raise ValueError(
|
||||
"`datasets` must be provided for the LightGBMTrainer API "
|
||||
"if `train_loop_per_worker` is not provided. "
|
||||
"This dict must contain the training dataset under the "
|
||||
f"key: '{TRAIN_DATASET_KEY}'. "
|
||||
f"Got keys: {list(datasets.keys())}"
|
||||
)
|
||||
if not label_column:
|
||||
raise ValueError(
|
||||
"`label_column` must be provided for the LightGBMTrainer API "
|
||||
"if `train_loop_per_worker` is not provided. "
|
||||
"This is the column name of the label in the dataset."
|
||||
)
|
||||
|
||||
num_boost_round = num_boost_round or 10
|
||||
|
||||
_log_deprecation_warning(LEGACY_LIGHTGBM_TRAINER_DEPRECATION_MESSAGE)
|
||||
|
||||
# Initialize a default Ray Train metrics/checkpoint reporting callback if needed
|
||||
callbacks = lightgbm_train_kwargs.get("callbacks", [])
|
||||
user_supplied_callback = any(
|
||||
isinstance(callback, RayTrainReportCallback) for callback in callbacks
|
||||
)
|
||||
callback_kwargs = {}
|
||||
if run_config:
|
||||
checkpoint_frequency = run_config.checkpoint_config.checkpoint_frequency
|
||||
checkpoint_at_end = run_config.checkpoint_config.checkpoint_at_end
|
||||
|
||||
callback_kwargs["frequency"] = checkpoint_frequency
|
||||
# Default `checkpoint_at_end=True` unless the user explicitly sets it.
|
||||
callback_kwargs["checkpoint_at_end"] = (
|
||||
checkpoint_at_end if checkpoint_at_end is not None else True
|
||||
)
|
||||
|
||||
if not user_supplied_callback:
|
||||
callbacks.append(RayTrainReportCallback(**callback_kwargs))
|
||||
lightgbm_train_kwargs["callbacks"] = callbacks
|
||||
|
||||
train_fn_per_worker = partial(
|
||||
_lightgbm_train_fn_per_worker,
|
||||
label_column=label_column,
|
||||
num_boost_round=num_boost_round,
|
||||
dataset_keys=set(datasets),
|
||||
lightgbm_train_kwargs=lightgbm_train_kwargs,
|
||||
)
|
||||
return train_fn_per_worker
|
||||
|
||||
@classmethod
|
||||
def get_model(
|
||||
cls,
|
||||
checkpoint: Checkpoint,
|
||||
) -> lightgbm.Booster:
|
||||
"""Retrieve the LightGBM model stored in this checkpoint."""
|
||||
return RayTrainReportCallback.get_model(checkpoint)
|
||||
@@ -0,0 +1,131 @@
|
||||
import logging
|
||||
from typing import Any, Callable, Dict, Optional, Union
|
||||
|
||||
import ray.train
|
||||
from ray.train import Checkpoint
|
||||
from ray.train.data_parallel_trainer import DataParallelTrainer
|
||||
from ray.train.lightgbm.config import LightGBMConfig, get_network_params # noqa: F401
|
||||
from ray.train.trainer import GenDataset
|
||||
|
||||
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 line below is the only change needed
|
||||
# for your `lgb.train` call!
|
||||
**ray.train.lightgbm.get_network_params(),
|
||||
}
|
||||
lgb.train(
|
||||
params,
|
||||
train_set,
|
||||
valid_sets=[eval_set],
|
||||
valid_names=["eval"],
|
||||
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=4),
|
||||
)
|
||||
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 use for training and validation.
|
||||
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.
|
||||
metadata: Dict that should be made available via
|
||||
`ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
|
||||
for checkpoints saved from this Trainer. Must be JSON-serializable.
|
||||
resume_from_checkpoint: A checkpoint to resume training from.
|
||||
This checkpoint can be accessed from within ``train_loop_per_worker``
|
||||
by calling ``ray.train.get_checkpoint()``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]],
|
||||
*,
|
||||
train_loop_config: Optional[Dict] = None,
|
||||
lightgbm_config: Optional[LightGBMConfig] = None,
|
||||
scaling_config: Optional[ray.train.ScalingConfig] = None,
|
||||
run_config: Optional[ray.train.RunConfig] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
dataset_config: Optional[ray.train.DataConfig] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
resume_from_checkpoint: Optional[Checkpoint] = None,
|
||||
):
|
||||
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,
|
||||
)
|
||||
Reference in New Issue
Block a user