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)