save_model¶

def save_model(
    forecaster: Any,
    path: str | os.PathLike[str],
    *,
    extra_trusted_types: Iterable[str] | None = None,
    extra_trusted_prefixes: Iterable[str] | None = None,
    extra_version_rules: Mapping[str, str] | None = None,
    signature: ModelSignature | None = None,
    conda_env: Any = None,
    mlflow_model: Model | None = None,
    pip_requirements: Any = None,
    extra_pip_requirements: Any = None,
    metadata: dict[str, Any] | None = None,
) -> None:
    """Save a fitted yohou forecaster as an MLflow model directory.

    The forecaster is written with skops, never pickle. Before anything is kept, the
    written file is checked against the trust policy and loaded back, and its
    predictions are compared with the original's. Any failure removes the directory.

    Parameters
    ----------
    forecaster : yohou forecaster
        A fitted yohou forecaster, including the state it accumulated through
        ``observe``.
    path : str or path-like
        Directory to create. It must not exist, or must be empty.
    extra_trusted_types : iterable of str or None, default=None
        Exact type names to trust in addition to the built-in policy, for example a
        third-party regressor inside the forecaster. They are not stored: pass the
        same names every time the model is loaded.
    extra_trusted_prefixes : iterable of str or None, default=None
        Module prefixes to trust, each ending in ``"."``, for example the caller's own
        package. Not stored: pass the same prefixes every time the model is loaded.
    extra_version_rules : mapping of str to str or None, default=None
        Distribution name to ``"exact"`` or ``"major.minor"``, for packages outside the
        built-in rules. A package's version is recorded only when the forecaster holds
        a type from it, and compared at load only under the loader's own rules.
    signature : mlflow.models.ModelSignature or None, default=None
        Optional signature. It may carry outputs and extra params, but no inputs.
        The flavour's four params are always added.
    conda_env : dict, str or None, default=None
        Conda environment, as for MLflow's built-in flavours.
    mlflow_model : mlflow.models.Model or None, default=None
        Model configuration to add the flavour to. Used by ``log_model``.
    pip_requirements : iterable of str, str or None, default=None
        Replaces the default pip requirements.
    extra_pip_requirements : iterable of str, str or None, default=None
        Added to the default pip requirements.
    metadata : dict or None, default=None
        Custom metadata stored in the ``MLmodel`` file.

    Raises
    ------
    TypeError
        If ``forecaster`` is not a yohou forecaster.
    sklearn.exceptions.NotFittedError
        If ``forecaster`` is not fitted. Nothing is written.
    ValueError
        If ``signature`` has an input schema or redefines a flavour param, the
        forecaster supports none of the ``point``, ``interval`` and ``class_proba``
        prediction types, a trusted prefix is not a module path ending in ``"."``, or a version rule
        is unknown or names a built-in package.
    UntrustedTypesError
        If the forecaster holds types outside the trust policy and the caller's
        additions.
    SaveVerificationError
        If the written model cannot be loaded back, fails to predict once reloaded, or
        loads back with different predictions.

    See Also
    --------
    log_model : Save the forecaster into an MLflow run and optionally register it.
    load_model : Load a saved forecaster.
    """
    from yohou.base import BaseForecaster

    if not isinstance(forecaster, BaseForecaster):
        msg = (
            f"save_model expects a yohou forecaster, got {type(forecaster).__module__}.{type(forecaster).__qualname__}."
        )
        raise TypeError(msg)
    check_is_fitted(forecaster)
    validate_prefixes(extra_trusted_prefixes)
    validate_version_rules(extra_version_rules)
    _validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
    supported = supported_prediction_types(forecaster)
    default_type = default_prediction_type(supported)
    signature = merge_signature(signature, default_type)

    target = os.path.abspath(os.fspath(path))
    _validate_and_prepare_target_save_path(target)
    try:
        _write_model(
            forecaster,
            target,
            extra_trusted_types=extra_trusted_types,
            extra_trusted_prefixes=extra_trusted_prefixes,
            extra_version_rules=extra_version_rules,
            signature=signature,
            supported=supported,
            default_type=default_type,
            conda_env=conda_env,
            mlflow_model=mlflow_model,
            pip_requirements=pip_requirements,
            extra_pip_requirements=extra_pip_requirements,
            metadata=metadata,
        )
    except BaseException:
        shutil.rmtree(target, ignore_errors=True)
        raise