chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,72 @@
|
||||
from typing import Dict
|
||||
|
||||
import ray.tune
|
||||
from ray.train.tensorflow import TensorflowCheckpoint
|
||||
from ray.train.tensorflow.keras import RayReportCallback
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
_DEPRECATION_MESSAGE = (
|
||||
"The `ray.tune.integration.keras` module is deprecated in favor of "
|
||||
"`ray.train.tensorflow.keras.ReportCheckpointCallback`."
|
||||
)
|
||||
|
||||
|
||||
class TuneReportCallback:
|
||||
"""Deprecated.
|
||||
Use :class:`ray.train.tensorflow.keras.ReportCheckpointCallback` instead."""
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
raise DeprecationWarning(_DEPRECATION_MESSAGE)
|
||||
|
||||
|
||||
class _TuneCheckpointCallback:
|
||||
"""Deprecated.
|
||||
Use :class:`ray.train.tensorflow.keras.ReportCheckpointCallback` instead."""
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
raise DeprecationWarning(_DEPRECATION_MESSAGE)
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class TuneReportCheckpointCallback(RayReportCallback):
|
||||
"""Keras callback for Ray Tune reporting and checkpointing.
|
||||
|
||||
.. note::
|
||||
Metrics are always reported with checkpoints, even if the event isn't specified
|
||||
in ``report_metrics_on``.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
############# Using it in Ray Tune ###############
|
||||
from ray.tune.integrations.keras import TuneReportCheckpointCallback
|
||||
|
||||
def train_fn():
|
||||
model = build_model()
|
||||
model.fit(dataset_shard, callbacks=[TuneReportCheckpointCallback()])
|
||||
|
||||
tuner = tune.Tuner(train_fn)
|
||||
results = tuner.fit()
|
||||
|
||||
Args:
|
||||
metrics: Metrics to report. If this is a list, each item describes
|
||||
the metric key reported to Keras, and it's reported under the
|
||||
same name. If this is a dict, each key is the name reported
|
||||
and the respective value is the metric key reported to Keras.
|
||||
If this is None, all Keras logs are reported.
|
||||
report_metrics_on: When to report metrics. Must be one of
|
||||
the Keras event hooks (less the ``on_``), e.g.
|
||||
"train_start" or "predict_end". Defaults to "epoch_end".
|
||||
checkpoint_on: When to save checkpoints. Must be one of the Keras event hooks
|
||||
(less the ``on_``), e.g. "train_start" or "predict_end". Defaults to
|
||||
"epoch_end".
|
||||
|
||||
"""
|
||||
|
||||
def _save_and_report_checkpoint(
|
||||
self, metrics: Dict, checkpoint: TensorflowCheckpoint
|
||||
):
|
||||
ray.tune.report(metrics, checkpoint=checkpoint)
|
||||
|
||||
def _report_metrics(self, metrics: Dict):
|
||||
ray.tune.report(metrics)
|
||||
@@ -0,0 +1,84 @@
|
||||
import tempfile
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
|
||||
from lightgbm import Booster
|
||||
|
||||
import ray.tune
|
||||
from ray.train.lightgbm._lightgbm_utils import RayReportCallback
|
||||
from ray.tune import Checkpoint
|
||||
from ray.util.annotations import Deprecated, PublicAPI
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class TuneReportCheckpointCallback(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.tune.integration.lightgbm import TuneReportCheckpointCallback
|
||||
|
||||
config = {
|
||||
# ...
|
||||
"metric": ["binary_logloss", "binary_error"],
|
||||
}
|
||||
|
||||
# Report only log loss to Tune after each validation epoch.
|
||||
bst = lightgbm.train(
|
||||
...,
|
||||
callbacks=[
|
||||
TuneReportCheckpointCallback(
|
||||
metrics={"loss": "eval-binary_logloss"}, frequency=1
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
"""
|
||||
|
||||
@contextmanager
|
||||
def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]:
|
||||
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)
|
||||
|
||||
def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster):
|
||||
with self._get_checkpoint(model=model) as checkpoint:
|
||||
ray.tune.report(report_dict, checkpoint=checkpoint)
|
||||
|
||||
def _report_metrics(self, report_dict: Dict):
|
||||
ray.tune.report(report_dict)
|
||||
|
||||
|
||||
@Deprecated
|
||||
class TuneReportCallback:
|
||||
def __new__(cls: type, *args, **kwargs):
|
||||
# TODO(justinvyu): [code_removal] Remove in 2.11.
|
||||
raise DeprecationWarning(
|
||||
"`TuneReportCallback` is deprecated. "
|
||||
"Use `ray.tune.integration.lightgbm.TuneReportCheckpointCallback` instead."
|
||||
)
|
||||
@@ -0,0 +1,206 @@
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import warnings
|
||||
from contextlib import contextmanager
|
||||
from typing import Dict, List, Optional, Type, Union
|
||||
|
||||
import ray.tune
|
||||
from ray.tune import Checkpoint
|
||||
from ray.util import log_once
|
||||
from ray.util.annotations import Deprecated, PublicAPI
|
||||
|
||||
try:
|
||||
from lightning.pytorch import Callback, LightningModule, Trainer
|
||||
except ModuleNotFoundError:
|
||||
from pytorch_lightning import Callback, LightningModule, Trainer
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Get all Pytorch Lightning Callback hooks based on whatever PTL version is being used.
|
||||
_allowed_hooks = {
|
||||
name
|
||||
for name, fn in inspect.getmembers(Callback, predicate=inspect.isfunction)
|
||||
if name.startswith("on_")
|
||||
}
|
||||
|
||||
|
||||
def _override_ptl_hooks(callback_cls: Type["TuneCallback"]) -> Type["TuneCallback"]:
|
||||
"""Overrides all allowed PTL Callback hooks with our custom handle logic."""
|
||||
|
||||
def generate_overridden_hook(fn_name):
|
||||
def overridden_hook(
|
||||
self,
|
||||
trainer: Trainer,
|
||||
*args,
|
||||
pl_module: Optional[LightningModule] = None,
|
||||
**kwargs,
|
||||
):
|
||||
if fn_name in self._on:
|
||||
self._handle(trainer=trainer, pl_module=pl_module)
|
||||
|
||||
return overridden_hook
|
||||
|
||||
# Set the overridden hook to all the allowed hooks in TuneCallback.
|
||||
for fn_name in _allowed_hooks:
|
||||
setattr(callback_cls, fn_name, generate_overridden_hook(fn_name))
|
||||
|
||||
return callback_cls
|
||||
|
||||
|
||||
@_override_ptl_hooks
|
||||
class TuneCallback(Callback):
|
||||
"""Base class for Tune's PyTorch Lightning callbacks.
|
||||
|
||||
Args:
|
||||
on: When to trigger checkpoint creations. Must be one of
|
||||
the PyTorch Lightning event hooks (less the ``on_``), e.g.
|
||||
"train_batch_start", or "train_end". Defaults to "validation_end"
|
||||
"""
|
||||
|
||||
def __init__(self, on: Union[str, List[str]] = "validation_end"):
|
||||
if not isinstance(on, list):
|
||||
on = [on]
|
||||
|
||||
for hook in on:
|
||||
if f"on_{hook}" not in _allowed_hooks:
|
||||
raise ValueError(
|
||||
f"Invalid hook selected: {hook}. Must be one of "
|
||||
f"{_allowed_hooks}"
|
||||
)
|
||||
|
||||
# Add back the "on_" prefix for internal consistency.
|
||||
on = [f"on_{hook}" for hook in on]
|
||||
|
||||
self._on = on
|
||||
|
||||
def _handle(self, trainer: Trainer, pl_module: Optional[LightningModule]):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@PublicAPI
|
||||
class TuneReportCheckpointCallback(TuneCallback):
|
||||
"""PyTorch Lightning report and checkpoint callback
|
||||
|
||||
Saves checkpoints after each validation step. Also reports metrics to Tune,
|
||||
which is needed for checkpoint registration.
|
||||
|
||||
Args:
|
||||
metrics: Metrics to report to Tune. If this is a list,
|
||||
each item describes the metric key reported to PyTorch Lightning,
|
||||
and it will reported under the same name to Tune. If this is a
|
||||
dict, each key will be the name reported to Tune and the respective
|
||||
value will be the metric key reported to PyTorch Lightning.
|
||||
filename: Filename of the checkpoint within the checkpoint
|
||||
directory. Defaults to "checkpoint".
|
||||
save_checkpoints: If True (default), checkpoints will be saved and
|
||||
reported to Ray. If False, only metrics will be reported.
|
||||
on: When to trigger checkpoint creations and metric reports. Must be one of
|
||||
the PyTorch Lightning event hooks (less the ``on_``), e.g.
|
||||
"train_batch_start", or "train_end". Defaults to "validation_end".
|
||||
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import lightning.pytorch as pl
|
||||
from ray.tune.integration.pytorch_lightning import (
|
||||
TuneReportCheckpointCallback)
|
||||
|
||||
# Save checkpoint after each training batch and after each
|
||||
# validation epoch.
|
||||
trainer = pl.Trainer(callbacks=[TuneReportCheckpointCallback(
|
||||
metrics={"loss": "val_loss", "mean_accuracy": "val_acc"},
|
||||
filename="trainer.ckpt", on="validation_end")])
|
||||
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
|
||||
filename: str = "checkpoint",
|
||||
save_checkpoints: bool = True,
|
||||
on: Union[str, List[str]] = "validation_end",
|
||||
):
|
||||
super(TuneReportCheckpointCallback, self).__init__(on=on)
|
||||
if isinstance(metrics, str):
|
||||
metrics = [metrics]
|
||||
self._save_checkpoints = save_checkpoints
|
||||
self._filename = filename
|
||||
self._metrics = metrics
|
||||
|
||||
def _get_report_dict(self, trainer: Trainer, pl_module: LightningModule):
|
||||
# Don't report if just doing initial validation sanity checks.
|
||||
if trainer.sanity_checking:
|
||||
return
|
||||
if not self._metrics:
|
||||
report_dict = {k: v.item() for k, v in trainer.callback_metrics.items()}
|
||||
else:
|
||||
report_dict = {}
|
||||
for key in self._metrics:
|
||||
if isinstance(self._metrics, dict):
|
||||
metric = self._metrics[key]
|
||||
else:
|
||||
metric = key
|
||||
if metric in trainer.callback_metrics:
|
||||
report_dict[key] = trainer.callback_metrics[metric].item()
|
||||
else:
|
||||
logger.warning(
|
||||
f"Metric {metric} does not exist in "
|
||||
"`trainer.callback_metrics."
|
||||
)
|
||||
|
||||
return report_dict
|
||||
|
||||
@contextmanager
|
||||
def _get_checkpoint(self, trainer: Trainer) -> Optional[Checkpoint]:
|
||||
if not self._save_checkpoints:
|
||||
yield None
|
||||
return
|
||||
|
||||
with tempfile.TemporaryDirectory() as checkpoint_dir:
|
||||
trainer.save_checkpoint(os.path.join(checkpoint_dir, self._filename))
|
||||
checkpoint = Checkpoint.from_directory(checkpoint_dir)
|
||||
yield checkpoint
|
||||
|
||||
def _handle(self, trainer: Trainer, pl_module: LightningModule):
|
||||
if trainer.sanity_checking:
|
||||
return
|
||||
|
||||
report_dict = self._get_report_dict(trainer, pl_module)
|
||||
if not report_dict:
|
||||
return
|
||||
|
||||
with self._get_checkpoint(trainer) as checkpoint:
|
||||
ray.tune.report(report_dict, checkpoint=checkpoint)
|
||||
|
||||
|
||||
class _TuneCheckpointCallback(TuneCallback):
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise DeprecationWarning(
|
||||
"`ray.tune.integration.pytorch_lightning._TuneCheckpointCallback` "
|
||||
"is deprecated."
|
||||
)
|
||||
|
||||
|
||||
@Deprecated
|
||||
class TuneReportCallback(TuneReportCheckpointCallback):
|
||||
def __init__(
|
||||
self,
|
||||
metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
|
||||
on: Union[str, List[str]] = "validation_end",
|
||||
):
|
||||
if log_once("tune_ptl_report_deprecated"):
|
||||
warnings.warn(
|
||||
"`ray.tune.integration.pytorch_lightning.TuneReportCallback` "
|
||||
"is deprecated. Use "
|
||||
"`ray.tune.integration.pytorch_lightning.TuneReportCheckpointCallback`"
|
||||
" instead."
|
||||
)
|
||||
super(TuneReportCallback, self).__init__(
|
||||
metrics=metrics, save_checkpoints=False, on=on
|
||||
)
|
||||
@@ -0,0 +1,40 @@
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from ray.train import Checkpoint as RayTrainCheckpoint
|
||||
from ray.train._internal.session import get_session
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2.api.callback import UserCallback
|
||||
from ray.tune.trainable.trainable_fn_utils import _in_tune_session
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
CHECKPOINT_PATH_KEY = "checkpoint_path"
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class TuneReportCallback(UserCallback):
|
||||
"""Propagate metrics and checkpoint paths from Ray Train workers to Ray Tune."""
|
||||
|
||||
def __init__(self):
|
||||
if not _in_tune_session():
|
||||
raise RuntimeError("TuneReportCallback must be used in a Tune session.")
|
||||
self._training_actor_item_queue = (
|
||||
get_session()._get_or_create_inter_actor_queue()
|
||||
)
|
||||
|
||||
def after_report(
|
||||
self,
|
||||
run_context: TrainRunContext,
|
||||
metrics: List[Dict[str, Any]],
|
||||
checkpoint: Optional[RayTrainCheckpoint],
|
||||
):
|
||||
# TODO: This can be changed to aggregate the metrics from all workers.
|
||||
# For now, just achieve feature parity with the old Tune+Train integration.
|
||||
metrics = metrics[0].copy()
|
||||
|
||||
# If a checkpoint is provided, add the checkpoint path to the metrics.
|
||||
# Don't report the checkpoint again since it's already been uploaded
|
||||
# to storage.
|
||||
if checkpoint:
|
||||
metrics[CHECKPOINT_PATH_KEY] = checkpoint.path
|
||||
|
||||
self._training_actor_item_queue.put(metrics)
|
||||
@@ -0,0 +1,101 @@
|
||||
import tempfile
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Callable, Dict, List, Optional, Union
|
||||
|
||||
from xgboost.core import Booster
|
||||
|
||||
import ray.tune
|
||||
from ray.train.xgboost._xgboost_utils import RayReportCallback
|
||||
from ray.tune import Checkpoint
|
||||
from ray.util.annotations import Deprecated, PublicAPI
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class TuneReportCheckpointCallback(RayReportCallback):
|
||||
"""XGBoost callback to save checkpoints and report metrics for Ray Tune.
|
||||
|
||||
Args:
|
||||
metrics: Metrics to report. If this is a list,
|
||||
each item describes the metric key reported to XGBoost,
|
||||
and it will be reported under the same name.
|
||||
This can also be a dict of {<key-to-report>: <xgboost-metric-key>},
|
||||
which can be used to rename xgboost default metrics.
|
||||
filename: Customize the saved checkpoint file type by passing
|
||||
a filename. Defaults to "model.ubj".
|
||||
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. For example, this can be used to
|
||||
average results across CV fold when using ``xgboost.cv``.
|
||||
|
||||
Examples:
|
||||
|
||||
Reporting checkpoints and metrics to Ray Tune when running many
|
||||
independent xgboost trials (without data parallelism within a trial).
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import xgboost
|
||||
|
||||
from ray.tune import Tuner
|
||||
from ray.tune.integration.xgboost import TuneReportCheckpointCallback
|
||||
|
||||
def train_fn(config):
|
||||
# Report log loss to Ray Tune after each validation epoch.
|
||||
bst = xgboost.train(
|
||||
...,
|
||||
callbacks=[
|
||||
TuneReportCheckpointCallback(
|
||||
metrics={"loss": "eval-logloss"}, frequency=1
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
tuner = Tuner(train_fn)
|
||||
results = tuner.fit()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
|
||||
filename: str = RayReportCallback.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,
|
||||
):
|
||||
super().__init__(
|
||||
metrics=metrics,
|
||||
filename=filename,
|
||||
frequency=frequency,
|
||||
checkpoint_at_end=checkpoint_at_end,
|
||||
results_postprocessing_fn=results_postprocessing_fn,
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]:
|
||||
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
||||
model.save_model(Path(temp_checkpoint_dir, self._filename).as_posix())
|
||||
yield Checkpoint(temp_checkpoint_dir)
|
||||
|
||||
def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster):
|
||||
with self._get_checkpoint(model=model) as checkpoint:
|
||||
ray.tune.report(report_dict, checkpoint=checkpoint)
|
||||
|
||||
def _report_metrics(self, report_dict: Dict):
|
||||
ray.tune.report(report_dict)
|
||||
|
||||
|
||||
@Deprecated
|
||||
class TuneReportCallback:
|
||||
def __new__(cls: type, *args, **kwargs):
|
||||
# TODO(justinvyu): [code_removal] Remove in 2.11.
|
||||
raise DeprecationWarning(
|
||||
"`TuneReportCallback` is deprecated. "
|
||||
"Use `ray.tune.integration.xgboost.TuneReportCheckpointCallback` instead."
|
||||
)
|
||||
Reference in New Issue
Block a user