def from_model(
cls,
booster: xgboost.Booster,
*,
path: os.PathLike,
preprocessor: Optional["Preprocessor"] = None,
) -> "XGBoostCheckpoint":
"""Create a :py:class:`~ray.air.checkpoint.Checkpoint` that stores an XGBoost
model.
Args:
booster: The XGBoost model to store in the checkpoint.
path: The directory where the checkpoint will be stored.
preprocessor: A fitted preprocessor to be applied before inference.
Returns:
An :py:class:`XGBoostCheckpoint` containing the specified ``Estimator``.
Examples:
>>> from ray.train.xgboost import XGBoostCheckpoint
>>> import xgboost
>>>
>>> booster = xgboost.Booster()
>>> checkpoint = XGBoostCheckpoint.from_model(booster, path=".") # doctest: +SKIP # noqa: E501
You can use a :py:class:`XGBoostCheckpoint` to create an
:py:class:`~ray.train.xgboost.XGBoostPredictor` and preform inference.
>>> from ray.train.xgboost import XGBoostPredictor
>>>
>>> predictor = XGBoostPredictor.from_checkpoint(checkpoint) # doctest: +SKIP # noqa: E501
"""
booster.save_model(os.path.join(path, MODEL_KEY))
if preprocessor:
save_preprocessor_to_dir(preprocessor, path)
checkpoint = cls.from_directory(path)
return checkpoint
def to_air_checkpoint(
path: str,
booster: xgboost.Booster,
preprocessor: Optional["Preprocessor"] = None,
) -> Checkpoint:
"""Convert a pretrained model to AIR checkpoint for serve or inference.
Args:
path: The directory path where model and preprocessor steps are stored to.
booster: A pretrained xgboost model.
preprocessor: A fitted preprocessor. The preprocessing logic will
be applied to serve/inference.
Returns:
A Ray Air checkpoint.
"""
booster.save_model(os.path.join(path, MODEL_KEY))
if preprocessor:
save_preprocessor_to_dir(preprocessor, path)
checkpoint = Checkpoint.from_directory(path)
return checkpoint
def dump(self, model: xgboost.Booster) -> FilesContextManager:
with tempfile.TemporaryDirectory(prefix='ebonite_xgboost_dump') as f:
path = os.path.join(f, self.model_path)
model.save_model(path)
yield Blobs({self.model_path: LocalFileBlob(path)})
def _save_model(self, model: xgboost.Booster, path: str):
model.save_model(path)
