Repository: sktime/sktime

Issue:
[ENH] Adding sktime to mlflow
**Is your feature request related to a problem? Please describe.**
I got the idea to add `sktime` as a package to `mlflow` like the following packages that are already included:
https://www.mlflow.org/docs/latest/python_api/index.html

**Describe the solution you'd like**
Fot this we would have to contribute to `mlflow` ofc, but I thought anyway to open the issue here first as we might find here someone to work on it more likely maybe. We could also open an issue at `mlfow`.

**Additional context**
Outreach. We might also think of other packages where we could integrate `sktime` similarly to `mlflow`.


Discussion:
I can work on this one.
nice! Welcome, @kirilral!
Might need a plan of attack and/or a design, since this is a larger project?
might be worth to look into this doc: https://www.mlflow.org/docs/latest/models.html#custom-flavors
@kirilral are you still working on this?

@aiwalter Would be happy to work on this FR as I will soon have some capacity to contribute 
@aiwalter thanks for the opportunity to contribute!

If we want to create a mlflow custom flavor I  guess we should implement this in a separate repo `mlflow-sktime` as suggested in the mlflow docs. Probably best place to host this repo would be under https://github.com/sktime - would you agree?

Will be happy to start working on it once we have a repo in place.

Later we could also add this as a community plugin to mlflow docs under [MLflow Plugins](https://mlflow.org/docs/latest/plugins.html#community-plugins)
@benjaminbluhm I would actually prefer to have it in `sktime` main repo bcz then it will be better maintained and run with the same CI/CD and release process. We can have it probably best in path `sktime/utils/mlflow.py` and we can have `mlflow` as a soft-dependency. Instruction how to do that is [here](https://www.sktime.org/en/stable/developer_guide/dependencies.html?#adding-a-soft-dependency)

Otherwise I can also support here with this soft-dependency mechanism!
@aiwalter agree your approach makes sense, guess then it would be as simple as `pip install sktime[mlflow]`

Will familiarize with contributing guidelines and set up development environment. 
The possibility for doing `pip install sktime[mlflow]` is actually a different construction side :) I have just created an [issue ](https://github.com/sktime/sktime/issues/3833) for this
@aiwalter where you suggest should documentation live for mlflow-sktime custom model flavor?

I will probably raise a draft PR some time next week
Great! Depends a bit on the size of the doc. A short doc can definitely go directly in the docstring and some code snippets in the docstring examples would also be nice. A longer doc could go in a separate notebook in the "examples" folder which we can then add to the sphinx somewhere and it will be rendered automatically 🙂
Okay sounds good, thanks for guidance. Will follow your suggestions, guess we can align more on docs once draft PR is raised :-)

Resolving patch:
diff --git a/sktime/utils/mlflow_sktime.py b/sktime/utils/mlflow_sktime.py
new file mode 100644
--- /dev/null
+++ b/sktime/utils/mlflow_sktime.py
@@ -0,0 +1,804 @@
+#!/usr/bin/env python3 -u
+# -*- coding: utf-8 -*-
+"""The ``mlflow_sktime`` module provides an MLflow API for ``sktime`` forecasters.
+
+This module exports ``sktime`` models in the following formats:
+
+sktime (native) format
+    This is the main flavor that can be loaded back into sktime, which relies on pickle
+    internally to serialize a model.
+mlflow.pyfunc
+    Produced for use by generic pyfunc-based deployment tools and batch inference.
+
+    The `pyfunc` flavor of the model supports sktime predict methods `predict`,
+    `predict_interval`, `predict_proba`, `predict_quantiles`, `predict_var`.
+
+    The interface for utilizing a sktime model loaded as a `pyfunc` type for
+    generating forecasts requires passing an exogenous regressor as Pandas
+    DataFrame to the `pyfunc.predict()` method (an empty DataFrame must be
+    passed if no exogenous regressor is used). The configuration of predict
+    methods and parameter values passed to the predict methods is defined by
+    a dictionary to be saved as an attribute of the fitted sktime model
+    instance. If no prediction configuration is defined `pyfunc.predict()`
+    will return output from sktime `predict` method. Note that for `pyfunc`
+    flavor the forecasting horizon `fh` must be passed to the fit method.
+
+    Predict methods and parameter values for `pyfunc` flavor can be defined
+    in two ways: `Dict[str, dict]` if parameter values are passed to
+    `pyfunc.predict()`, for example
+    `{"predict_method": {"predict": {}, "predict_interval": {"coverage": [0.1, 0.9]}}`.
+    `Dict[str, list]`, with default parameters in predict method, for example
+    `{"predict_method": ["predict", "predict_interval"}` (Note: when including
+    `predict_proba` method the former appraoch must be followed as `quantiles`
+    parameter has to be provided by the user). If no prediction config is defined
+    `pyfunc.predict()` will return output from sktime `predict()` method.
+"""
+
+__author__ = ["benjaminbluhm"]
+__all__ = [
+    "get_default_pip_requirements",
+    "get_default_conda_env",
+    "save_model",
+    "log_model",
+    "load_model",
+]
+
+import logging
+import os
+import pickle
+
+import pandas as pd
+import yaml
+
+import sktime
+from sktime import utils
+from sktime.utils.multiindex import flatten_multiindex
+from sktime.utils.validation._dependencies import _check_soft_dependencies
+
+if _check_soft_dependencies("mlflow", severity="warning"):
+    from mlflow import pyfunc
+
+FLAVOR_NAME = "mlflow_sktime"
+
+PYFUNC_PREDICT_CONF = "pyfunc_predict_conf"
+PYFUNC_PREDICT_CONF_KEY = "predict_method"
+SKTIME_PREDICT = "predict"
+SKTIME_PREDICT_INTERVAL = "predict_interval"
+SKTIME_PREDICT_PROBA = "predict_proba"
+SKTIME_PREDICT_QUANTILES = "predict_quantiles"
+SKTIME_PREDICT_VAR = "predict_var"
+SUPPORTED_SKTIME_PREDICT_METHODS = [
+    SKTIME_PREDICT,
+    SKTIME_PREDICT_INTERVAL,
+    SKTIME_PREDICT_PROBA,
+    SKTIME_PREDICT_QUANTILES,
+    SKTIME_PREDICT_VAR,
+]
+
+SERIALIZATION_FORMAT_PICKLE = "pickle"
+SERIALIZATION_FORMAT_CLOUDPICKLE = "cloudpickle"
+SUPPORTED_SERIALIZATION_FORMATS = [
+    SERIALIZATION_FORMAT_PICKLE,
+    SERIALIZATION_FORMAT_CLOUDPICKLE,
+]
+
+_logger = logging.getLogger(__name__)
+
+
+def get_default_pip_requirements(include_cloudpickle=False):
+    """Create list of default pip requirements for MLflow Models.
+
+    Returns
+    -------
+    list of default pip requirements for MLflow Models produced by this flavor.
+    Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment
+    that, at a minimum, contains these requirements.
+    """
+    _check_soft_dependencies("mlflow", severity="error")
+    from mlflow.utils.requirements_utils import _get_pinned_requirement
+
+    pip_deps = [_get_pinned_requirement("sktime")]
+    if include_cloudpickle:
+        pip_deps += [_get_pinned_requirement("cloudpickle")]
+
+    return pip_deps
+
+
+def get_default_conda_env(include_cloudpickle=False):
+    """Return default Conda environment for MLflow Models.
+
+    Returns
+    -------
+    The default Conda environment for MLflow Models produced by calls to
+    :func:`save_model()` and :func:`log_model()`
+    """
+    _check_soft_dependencies("mlflow", severity="error")
+    from mlflow.utils.environment import _mlflow_conda_env
+
+    return _mlflow_conda_env(
+        additional_pip_deps=get_default_pip_requirements(include_cloudpickle)
+    )
+
+
+def save_model(
+    sktime_model,
+    path,
+    conda_env=None,
+    code_paths=None,
+    mlflow_model=None,
+    signature=None,
+    input_example=None,
+    pip_requirements=None,
+    extra_pip_requirements=None,
+    serialization_format=SERIALIZATION_FORMAT_PICKLE,
+):  # TODO: can we specify a type for fitted instance of sktime model below?
+    """Save a sktime model to a path on the local file system.
+
+    Parameters
+    ----------
+    sktime_model :
+        Fitted sktime model object.
+    path : str
+        Local path where the model is to be saved.
+    conda_env : Union[dict, str], optional (default=None)
+        Either a dictionary representation of a Conda environment or the path to a
+        conda environment yaml file.
+    code_paths : array-like, optional (default=None)
+        A list of local filesystem paths to Python file dependencies (or directories
+        containing file dependencies). These files are *prepended* to the system path
+        when the model is loaded.
+    mlflow_model: mlflow.models.Model, optional (default=None)
+        mlflow.models.Model configuration to which to add the python_function flavor.
+    signature : mlflow.models.signature.ModelSignature, optional (default=None)
+        Model Signature mlflow.models.ModelSignature describes
+        model input and output :py:class:`Schema <mlflow.types.Schema>`. The model
+        signature can be :py:func:`inferred <mlflow.models.infer_signature>` from
+        datasets with valid model input (e.g. the training dataset with target column
+        omitted) and valid model output (e.g. model predictions generated on the
+        training dataset), for example:
+
+        .. code-block:: py
+
+          from mlflow.models.signature import infer_signature
+          train = df.drop_column("target_label")
+          predictions = ... # compute model predictions
+          signature = infer_signature(train, predictions)
+
+        .. Warning:: if performing probabilistic forecasts (``predict_interval``,
+          ``predict_quantiles``) with a sktime model, the signature
+          on the returned prediction object will not be correctly inferred due
+          to the Pandas MultiIndex column type when using the these methods.
+          ``infer_schema`` will function correctly if using the ``pyfunc`` flavor
+          of the model, though.
+    input_example : Union[pandas.core.frame.DataFrame, numpy.ndarray, dict, list, csr_matrix, csc_matrix], optional (default=None)
+        Input example provides one or several instances of valid model input.
+        The example can be used as a hint of what data to feed the model. The given
+        example will be converted to a ``Pandas DataFrame`` and then serialized to json
+        using the ``Pandas`` split-oriented format. Bytes are base64-encoded.
+    pip_requirements : Union[Iterable, str], optional (default=None)
+        Either an iterable of pip requirement strings
+        (e.g. ["sktime", "-r requirements.txt", "-c constraints.txt"]) or the string
+        path to a pip requirements file on the local filesystem
+        (e.g. "requirements.txt")
+    extra_pip_requirements : Union[Iterable, str], optional (default=None)
+        Either an iterable of pip requirement strings
+        (e.g. ["pandas", "-r requirements.txt", "-c constraints.txt"]) or the string
+        path to a pip requirements file on the local filesystem
+        (e.g. "requirements.txt")
+    serialization_format : str, optional (default="pickle")
+        The format in which to serialize the model. This should be one of the formats
+        "pickle" or "cloudpickle"
+
+    References
+    ----------
+    .. [1] https://www.mlflow.org/docs/latest/python_api/mlflow.models.html#mlflow.models.Model.save
+
+    Examples
+    --------
+    >>> from sktime.datasets import load_airline  # doctest: +SKIP
+    >>> from sktime.forecasting.arima import ARIMA  # doctest: +SKIP
+    >>> from sktime.utils import mlflow_sktime  # doctest: +SKIP
+    >>> y = load_airline()  # doctest: +SKIP
+    >>> forecaster = ARIMA(  # doctest: +SKIP
+    ...     order=(1, 1, 0),
+    ...     seasonal_order=(0, 1, 0, 12),
+    ...     suppress_warnings=True)
+    >>> forecaster.fit(y)  # doctest: +SKIP
+    ARIMA(...)
+    >>> model_path = "model"  # doctest: +SKIP
+    >>> mlflow_sktime.save_model(  # doctest: +SKIP
+    ...     sktime_model=forecaster,
+    ...     path=model_path)  # doctest: +SKIP
+    >>> loaded_model = mlflow_sktime.load_model(model_uri=model_path)  # doctest: +SKIP
+    >>> loaded_model.predict(fh=[1, 2, 3])  # doctest: +SKIP
+    """  # noqa: E501
+    _check_soft_dependencies("mlflow", severity="error")
+    from mlflow.exceptions import MlflowException
+    from mlflow.models import Model
+    from mlflow.models.model import MLMODEL_FILE_NAME
+    from mlflow.models.utils import _save_example
+    from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+    from mlflow.utils.environment import (
+        _CONDA_ENV_FILE_NAME,
+        _CONSTRAINTS_FILE_NAME,
+        _PYTHON_ENV_FILE_NAME,
+        _REQUIREMENTS_FILE_NAME,
+        _process_conda_env,
+        _process_pip_requirements,
+        _PythonEnv,
+        _validate_env_arguments,
+    )
+    from mlflow.utils.file_utils import write_to
+    from mlflow.utils.model_utils import (
+        _validate_and_copy_code_paths,
+        _validate_and_prepare_target_save_path,
+    )
+
+    _validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
+
+    if serialization_format not in SUPPORTED_SERIALIZATION_FORMATS:
+        raise MlflowException(
+            message=(
+                "Unrecognized serialization format: {serialization_format}. "
+                "Please specify one of the following supported formats: "
+                "{supported_formats}.".format(
+                    serialization_format=serialization_format,
+                    supported_formats=SUPPORTED_SERIALIZATION_FORMATS,
+                )
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    _validate_and_prepare_target_save_path(path)
+    code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
+
+    if mlflow_model is None:
+        mlflow_model = Model()
+    if signature is not None:
+        mlflow_model.signature = signature
+    if input_example is not None:
+        _save_example(mlflow_model, input_example, path)
+
+    model_data_subpath = "model.pkl"
+    model_data_path = os.path.join(path, model_data_subpath)
+    _save_model(
+        sktime_model, model_data_path, serialization_format=serialization_format
+    )
+
+    pyfunc.add_to_model(
+        mlflow_model,
+        loader_module="sktime.utils.mlflow_sktime",
+        model_path=model_data_subpath,
+        conda_env=_CONDA_ENV_FILE_NAME,
+        python_env=_PYTHON_ENV_FILE_NAME,
+        code=code_dir_subpath,
+    )
+
+    mlflow_model.add_flavor(
+        FLAVOR_NAME,
+        pickled_model=model_data_subpath,
+        sktime_version=sktime.__version__,
+        serialization_format=serialization_format,
+        code=code_dir_subpath,
+    )
+    mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
+
+    if conda_env is None:
+        if pip_requirements is None:
+            include_cloudpickle = (
+                serialization_format == SERIALIZATION_FORMAT_CLOUDPICKLE
+            )
+            default_reqs = get_default_pip_requirements(include_cloudpickle)
+            default_reqs = sorted(default_reqs)
+        else:
+            default_reqs = None
+        conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
+            default_reqs, pip_requirements, extra_pip_requirements
+        )
+    else:
+        conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)
+
+    with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
+        yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
+
+    if pip_constraints:
+        write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
+
+    write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
+
+    _PythonEnv.current().to_yaml(os.path.join(path, _PYTHON_ENV_FILE_NAME))
+
+
+def log_model(
+    sktime_model,
+    artifact_path,
+    conda_env=None,
+    code_paths=None,
+    registered_model_name=None,
+    signature=None,
+    input_example=None,
+    await_registration_for=None,
+    pip_requirements=None,
+    extra_pip_requirements=None,
+    serialization_format=SERIALIZATION_FORMAT_PICKLE,
+    **kwargs,
+):  # TODO: can we specify a type for fitted instance of sktime model below?
+    """
+    Log a sktime model as an MLflow artifact for the current run.
+
+    Parameters
+    ----------
+    sktime_model : fitted sktime model
+        Fitted sktime model object.
+    artifact_path : str
+        Run-relative artifact path to save the model to.
+    conda_env : Union[dict, str], optional (default=None)
+        Either a dictionary representation of a Conda environment or the path to a
+        conda environment yaml file.
+    code_paths : array-like, optional (default=None)
+        A list of local filesystem paths to Python file dependencies (or directories
+        containing file dependencies). These files are *prepended* to the system path
+        when the model is loaded.
+    registered_model_name : str, optional (default=None)
+        If given, create a model version under ``registered_model_name``, also creating
+        a registered model if one with the given name does not exist.
+    signature : mlflow.models.signature.ModelSignature, optional (default=None)
+        Model Signature mlflow.models.ModelSignature describes
+        model input and output :py:class:`Schema <mlflow.types.Schema>`. The model
+        signature can be :py:func:`inferred <mlflow.models.infer_signature>` from
+        datasets with valid model input (e.g. the training dataset with target column
+        omitted) and valid model output (e.g. model predictions generated on the
+        training dataset), for example:
+
+        .. code-block:: py
+
+          from mlflow.models.signature import infer_signature
+          train = df.drop_column("target_label")
+          predictions = ... # compute model predictions
+          signature = infer_signature(train, predictions)
+
+        .. Warning:: if performing probabilistic forecasts (``predict_interval``,
+          ``predict_quantiles``) with a sktime model, the signature
+          on the returned prediction object will not be correctly inferred due
+          to the Pandas MultiIndex column type when using the these methods.
+          ``infer_schema`` will function correctly if using the ``pyfunc`` flavor
+          of the model, though. The ``pyfunc`` flavor of the model supports sktime
+          predict methods ``predict``, ``predict_interval``, ``predict_quantiles``
+          and ``predict_var`` while ``predict_proba`` and ``predict_residuals`` are
+          currently not supported.
+    input_example : Union[pandas.core.frame.DataFrame, numpy.ndarray, dict, list, csr_matrix, csc_matrix], optional (default=None)
+        Input example provides one or several instances of valid model input.
+        The example can be used as a hint of what data to feed the model. The given
+        example will be converted to a ``Pandas DataFrame`` and then serialized to json
+        using the ``Pandas`` split-oriented format. Bytes are base64-encoded.
+    await_registration_for : int, optional (default=None)
+        Number of seconds to wait for the model version to finish being created and is
+        in ``READY`` status. By default, the function waits for five minutes. Specify 0
+        or None to skip waiting.
+    pip_requirements : Union[Iterable, str], optional (default=None)
+        Either an iterable of pip requirement strings
+        (e.g. ["sktime", "-r requirements.txt", "-c constraints.txt"]) or the string
+        path to a pip requirements file on the local filesystem
+        (e.g. "requirements.txt")
+    extra_pip_requirements : Union[Iterable, str], optional (default=None)
+        Either an iterable of pip requirement strings
+        (e.g. ["pandas", "-r requirements.txt", "-c constraints.txt"]) or the string
+        path to a pip requirements file on the local filesystem
+        (e.g. "requirements.txt")
+    serialization_format : str, optional (default="pickle")
+        The format in which to serialize the model. This should be one of the formats
+        "pickle" or "cloudpickle"
+    kwargs:
+        Additional arguments for :py:class:`mlflow.models.model.Model`
+
+    Returns
+    -------
+    A :py:class:`ModelInfo <mlflow.models.model.ModelInfo>` instance that contains the
+    metadata of the logged model.
+
+    See Also
+    --------
+    MLflow
+
+    References
+    ----------
+    .. [1] https://www.mlflow.org/docs/latest/python_api/mlflow.models.html#mlflow.models.Model.log
+
+    >>> import mlflow  # doctest: +SKIP
+    >>> from mlflow.utils.environment import _mlflow_conda_env  # doctest: +SKIP
+    >>> from sktime.datasets import load_airline  # doctest: +SKIP
+    >>> from sktime.forecasting.arima import ARIMA  # doctest: +SKIP
+    >>> from sktime.utils import mlflow_sktime  # doctest: +SKIP
+    >>> y = load_airline()  # doctest: +SKIP
+    >>> forecaster = ARIMA(  # doctest: +SKIP
+    ...     order=(1, 1, 0),
+    ...     seasonal_order=(0, 1, 0, 12),
+    ...     suppress_warnings=True)
+    >>> forecaster.fit(y)  # doctest: +SKIP
+    ARIMA(...)
+    >>> mlflow.start_run()  # doctest: +SKIP
+    >>> artifact_path = "model"  # doctest: +SKIP
+    >>> model_info = mlflow_sktime.log_model(
+    ...     sktime_model=forecaster,
+    ...     artifact_path=artifact_path)  # doctest: +SKIP
+    """  # noqa: E501
+    _check_soft_dependencies("mlflow", severity="error")
+    from mlflow.models import Model
+
+    if await_registration_for is None:
+        from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
+
+        await_registration_for = DEFAULT_AWAIT_MAX_SLEEP_SECONDS
+
+    return Model.log(
+        artifact_path=artifact_path,
+        flavor=utils.mlflow_sktime,
+        registered_model_name=registered_model_name,
+        sktime_model=sktime_model,
+        conda_env=conda_env,
+        code_paths=code_paths,
+        signature=signature,
+        input_example=input_example,
+        await_registration_for=await_registration_for,
+        pip_requirements=pip_requirements,
+        extra_pip_requirements=extra_pip_requirements,
+        serialization_format=serialization_format,
+        **kwargs,
+    )
+
+
+def load_model(model_uri, dst_path=None):
+    """
+    Load a sktime model from a local file or a run.
+
+    Parameters
+    ----------
+    model_uri : str
+        The location, in URI format, of the MLflow model. For example:
+
+                    - ``/Users/me/path/to/local/model``
+                    - ``relative/path/to/local/model``
+                    - ``s3://my_bucket/path/to/model``
+                    - ``runs:/<mlflow_run_id>/run-relative/path/to/model``
+                    - ``mlflow-artifacts:/path/to/model``
+
+        For more information about supported URI schemes, see
+        `Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html#
+        artifact-locations>`_.
+    dst_path : str, optional (default=None)
+        The local filesystem path to which to download the model artifact.This
+        directory must already exist. If unspecified, a local output path will
+        be created.
+
+    Returns
+    -------
+    A sktime model instance.
+
+    References
+    ----------
+    .. [1] https://www.mlflow.org/docs/latest/python_api/mlflow.models.html#mlflow.models.Model.load
+
+    Examples
+    --------
+    >>> from sktime.datasets import load_airline
+    >>> from sktime.forecasting.arima import ARIMA
+    >>> from sktime.utils import mlflow_sktime  # doctest: +SKIP
+    >>> y = load_airline()  # doctest: +SKIP
+    >>> forecaster = ARIMA(  # doctest: +SKIP
+    ...     order=(1, 1, 0),
+    ...     seasonal_order=(0, 1, 0, 12),
+    ...     suppress_warnings=True)
+    >>> forecaster.fit(y)  # doctest: +SKIP
+    ARIMA(...)
+    >>> model_path = "model"  # doctest: +SKIP
+    >>> mlflow_sktime.save_model(  # doctest: +SKIP
+    ...     sktime_model=forecaster,
+    ...     path=model_path)
+    >>> loaded_model = mlflow_sktime.load_model(model_uri=model_path)  # doctest: +SKIP
+    """  # noqa: E501
+    _check_soft_dependencies("mlflow", severity="error")
+    from mlflow.tracking.artifact_utils import _download_artifact_from_uri
+    from mlflow.utils.model_utils import (
+        _add_code_from_conf_to_system_path,
+        _get_flavor_configuration,
+    )
+
+    local_model_path = _download_artifact_from_uri(
+        artifact_uri=model_uri, output_path=dst_path
+    )
+    flavor_conf = _get_flavor_configuration(
+        model_path=local_model_path, flavor_name=FLAVOR_NAME
+    )
+    _add_code_from_conf_to_system_path(local_model_path, flavor_conf)
+    sktime_model_file_path = os.path.join(
+        local_model_path, flavor_conf["pickled_model"]
+    )
+    serialization_format = flavor_conf.get(
+        "serialization_format", SERIALIZATION_FORMAT_PICKLE
+    )
+    return _load_model(
+        path=sktime_model_file_path, serialization_format=serialization_format
+    )
+
+
+def _save_model(model, path, serialization_format):
+
+    _check_soft_dependencies("mlflow", severity="error")
+    from mlflow.exceptions import MlflowException
+    from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+
+    with open(path, "wb") as out:
+        if serialization_format == SERIALIZATION_FORMAT_PICKLE:
+            pickle.dump(model, out)
+        elif serialization_format == SERIALIZATION_FORMAT_CLOUDPICKLE:
+            _check_soft_dependencies("cloudpickle", severity="error")
+            import cloudpickle
+
+            cloudpickle.dump(model, out)
+        else:
+            raise MlflowException(
+                message="Unrecognized serialization format: "
+                "{serialization_format}".format(
+                    serialization_format=serialization_format
+                ),
+                error_code=INTERNAL_ERROR,
+            )
+
+
+def _load_model(path, serialization_format):
+
+    _check_soft_dependencies("mlflow", severity="error")
+    from mlflow.exceptions import MlflowException
+    from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+    if serialization_format not in SUPPORTED_SERIALIZATION_FORMATS:
+        raise MlflowException(
+            message=(
+                "Unrecognized serialization format: {serialization_format}. "
+                "Please specify one of the following supported formats: "
+                "{supported_formats}.".format(
+                    serialization_format=serialization_format,
+                    supported_formats=SUPPORTED_SERIALIZATION_FORMATS,
+                )
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    with open(path, "rb") as pickled_model:
+        if serialization_format == SERIALIZATION_FORMAT_PICKLE:
+            return pickle.load(pickled_model)
+        elif serialization_format == SERIALIZATION_FORMAT_CLOUDPICKLE:
+            _check_soft_dependencies("cloudpickle", severity="error")
+            import cloudpickle
+
+            return cloudpickle.load(pickled_model)
+
+
+def _load_pyfunc(path):
+    """Load PyFunc implementation. Called by ``pyfunc.load_model``.
+
+    Parameters
+    ----------
+    path : str
+        Local filesystem path to the MLflow Model with the sktime flavor.
+
+    See Also
+    --------
+    MLflow
+
+    References
+    ----------
+    .. [1] https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#mlflow.pyfunc.load_model
+    """  # noqa: E501
+    _check_soft_dependencies("mlflow", severity="error")
+    from mlflow.exceptions import MlflowException
+    from mlflow.utils.model_utils import _get_flavor_configuration
+
+    if os.path.isfile(path):
+        serialization_format = SERIALIZATION_FORMAT_PICKLE
+        _logger.warning(
+            "Loading procedure in older versions of MLflow using pickle.load()"
+        )
+    else:
+        try:
+            sktime_flavor_conf = _get_flavor_configuration(
+                model_path=path, flavor_name=FLAVOR_NAME
+            )
+            serialization_format = sktime_flavor_conf.get(
+                "serialization_format", SERIALIZATION_FORMAT_PICKLE
+            )
+        except MlflowException:
+            _logger.warning(
+                "Could not find sktime flavor configuration during model "
+                "loading process. Assuming 'pickle' serialization format."
+            )
+            serialization_format = SERIALIZATION_FORMAT_PICKLE
+
+        pyfunc_flavor_conf = _get_flavor_configuration(
+            model_path=path, flavor_name=pyfunc.FLAVOR_NAME
+        )
+        path = os.path.join(path, pyfunc_flavor_conf["model_path"])
+
+    return _SktimeModelWrapper(
+        _load_model(path, serialization_format=serialization_format)
+    )
+
+
+class _SktimeModelWrapper:
+    def __init__(self, sktime_model):
+        _check_soft_dependencies("mlflow", severity="error")
+        self.sktime_model = sktime_model
+
+    def predict(self, X):
+
+        from mlflow.exceptions import MlflowException
+        from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+        X = None if X.empty else X
+        raw_predictions = {}
+
+        if not hasattr(self.sktime_model, "pyfunc_predict_conf"):
+            raw_predictions[SKTIME_PREDICT] = self.sktime_model.predict(X=X)
+
+        else:
+
+            if not isinstance(self.sktime_model.pyfunc_predict_conf, dict):
+                raise MlflowException(
+                    f"Attribute {PYFUNC_PREDICT_CONF} must be of type dict.",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+
+            if PYFUNC_PREDICT_CONF_KEY not in self.sktime_model.pyfunc_predict_conf:
+                raise MlflowException(
+                    f"Attribute {PYFUNC_PREDICT_CONF} must contain "
+                    f"a dictionary key {PYFUNC_PREDICT_CONF_KEY}.",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+
+            if isinstance(
+                self.sktime_model.pyfunc_predict_conf[PYFUNC_PREDICT_CONF_KEY], list
+            ):
+                predict_methods = self.sktime_model.pyfunc_predict_conf[
+                    PYFUNC_PREDICT_CONF_KEY
+                ]
+                predict_params = False
+            elif isinstance(
+                self.sktime_model.pyfunc_predict_conf[PYFUNC_PREDICT_CONF_KEY], dict
+            ):
+                predict_methods = list(
+                    self.sktime_model.pyfunc_predict_conf[
+                        PYFUNC_PREDICT_CONF_KEY
+                    ].keys()
+                )
+                predict_params = True
+            else:
+                raise MlflowException(
+                    "Dictionary value must be of type dict or list.",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+
+            if not set(predict_methods).issubset(set(SUPPORTED_SKTIME_PREDICT_METHODS)):
+                raise MlflowException(
+                    f"The provided {PYFUNC_PREDICT_CONF_KEY} values must be "
+                    f"a subset of {SUPPORTED_SKTIME_PREDICT_METHODS}",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+
+            if SKTIME_PREDICT in predict_methods:
+                raw_predictions[SKTIME_PREDICT] = self.sktime_model.predict(X=X)
+
+            if SKTIME_PREDICT_INTERVAL in predict_methods:
+                if predict_params:
+                    coverage = (
+                        0.9
+                        if "coverage"
+                        not in self.sktime_model.pyfunc_predict_conf[
+                            PYFUNC_PREDICT_CONF_KEY
+                        ][SKTIME_PREDICT_INTERVAL]
+                        else self.sktime_model.pyfunc_predict_conf[
+                            PYFUNC_PREDICT_CONF_KEY
+                        ][SKTIME_PREDICT_INTERVAL]["coverage"]
+                    )
+                else:
+                    coverage = 0.9
+
+                raw_predictions[
+                    SKTIME_PREDICT_INTERVAL
+                ] = self.sktime_model.predict_interval(X=X, coverage=coverage)
+
+            if SKTIME_PREDICT_PROBA in predict_methods:
+                if not isinstance(
+                    self.sktime_model.pyfunc_predict_conf[PYFUNC_PREDICT_CONF_KEY], dict
+                ):
+                    raise MlflowException(
+                        f"Method {SKTIME_PREDICT_PROBA} requires passing a dictionary.",
+                        error_code=INVALID_PARAMETER_VALUE,
+                    )
+
+                if (
+                    "quantiles"
+                    not in self.sktime_model.pyfunc_predict_conf[
+                        PYFUNC_PREDICT_CONF_KEY
+                    ][SKTIME_PREDICT_PROBA]
+                ):
+                    raise MlflowException(
+                        f"Method {SKTIME_PREDICT_PROBA} requires passing "
+                        f"quantile values.",
+                        error_code=INVALID_PARAMETER_VALUE,
+                    )
+
+                quantiles = self.sktime_model.pyfunc_predict_conf[
+                    PYFUNC_PREDICT_CONF_KEY
+                ][SKTIME_PREDICT_PROBA]["quantiles"]
+                marginal = (
+                    True
+                    if "marginal"
+                    not in self.sktime_model.pyfunc_predict_conf[
+                        PYFUNC_PREDICT_CONF_KEY
+                    ][SKTIME_PREDICT_PROBA]
+                    else self.sktime_model.pyfunc_predict_conf[PYFUNC_PREDICT_CONF_KEY][
+                        SKTIME_PREDICT_PROBA
+                    ]["marginal"]
+                )
+
+                y_pred_dist = self.sktime_model.predict_proba(X=X, marginal=marginal)
+                y_pred_dist_quantiles = pd.DataFrame(y_pred_dist.quantile(quantiles))
+                y_pred_dist_quantiles.columns = [f"Quantiles_{q}" for q in quantiles]
+                y_pred_dist_quantiles.index = y_pred_dist.parameters["loc"].index
+
+                raw_predictions[SKTIME_PREDICT_PROBA] = y_pred_dist_quantiles
+
+            if SKTIME_PREDICT_QUANTILES in predict_methods:
+                if predict_params:
+                    alpha = (
+                        None
+                        if "alpha"
+                        not in self.sktime_model.pyfunc_predict_conf[
+                            PYFUNC_PREDICT_CONF_KEY
+                        ][SKTIME_PREDICT_QUANTILES]
+                        else self.sktime_model.pyfunc_predict_conf[
+                            PYFUNC_PREDICT_CONF_KEY
+                        ][SKTIME_PREDICT_QUANTILES]["alpha"]
+                    )
+                else:
+                    alpha = None
+                raw_predictions[
+                    SKTIME_PREDICT_QUANTILES
+                ] = self.sktime_model.predict_quantiles(X=X, alpha=alpha)
+
+            if SKTIME_PREDICT_VAR in predict_methods:
+                if predict_params:
+                    cov = (
+                        False
+                        if "cov"
+                        not in self.sktime_model.pyfunc_predict_conf[
+                            PYFUNC_PREDICT_CONF_KEY
+                        ][SKTIME_PREDICT_VAR]
+                        else self.sktime_model.pyfunc_predict_conf[
+                            PYFUNC_PREDICT_CONF_KEY
+                        ][SKTIME_PREDICT_VAR]["cov"]
+                    )
+                else:
+                    cov = False
+                raw_predictions[SKTIME_PREDICT_VAR] = self.sktime_model.predict_var(
+                    X=X, cov=cov
+                )
+
+        for k, v in raw_predictions.items():
+            if hasattr(v, "columns") and isinstance(v.columns, pd.MultiIndex):
+                raw_predictions[k].columns = flatten_multiindex(v)
+
+        if len(raw_predictions) > 1:
+            predictions = pd.concat(
+                list(raw_predictions.values()),
+                axis=1,
+                keys=list(raw_predictions.keys()),
+            )
+            predictions.columns = flatten_multiindex(predictions)
+        else:
+            predictions = raw_predictions[list(raw_predictions.keys())[0]]
+
+        return predictions
+
+
+# TODO: Add support for autologging


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -148,7 +148,7 @@ jobs:
 
       - name: Install sktime and dependencies
         run: |
-          python -m pip install .[all_extras,dev]
+          python -m pip install .[all_extras,dev] --no-cache-dir
 
       - name: Show dependencies
         run: python -m pip list
diff --git a/sktime/utils/tests/test_mlflow_sktime_model_export.py b/sktime/utils/tests/test_mlflow_sktime_model_export.py
new file mode 100644
--- /dev/null
+++ b/sktime/utils/tests/test_mlflow_sktime_model_export.py
@@ -0,0 +1,714 @@
+# -*- coding: utf-8 -*-
+"""Tests for mlflow-sktime custom model flavor."""
+
+__author__ = ["benjaminbluhm"]
+
+import os
+import sys
+from pathlib import Path
+from unittest import mock
+
+import boto3
+import moto
+import numpy as np
+import pandas as pd
+import pytest
+from botocore.config import Config
+
+from sktime.datasets import load_airline, load_longley
+from sktime.forecasting.arima import AutoARIMA
+from sktime.forecasting.model_selection import temporal_train_test_split
+from sktime.forecasting.naive import NaiveForecaster
+from sktime.utils.multiindex import flatten_multiindex
+from sktime.utils.validation._dependencies import _check_soft_dependencies
+
+if not sys.platform.startswith("linux"):
+    pytest.skip("Skipping MLflow tests for Windows and macOS", allow_module_level=True)
+
+
+@pytest.fixture
+def model_path(tmp_path):
+    """Create a temporary path to save/log model."""
+    return tmp_path.joinpath("model")
+
+
+@pytest.fixture(scope="module")
+def mock_s3_bucket():
+    """Create a mock S3 bucket using moto.
+
+    Returns
+    -------
+    string with name of mock S3 bucket
+    """
+    with moto.mock_s3():
+        bucket_name = "mock-bucket"
+        my_config = Config(region_name="us-east-1")
+        s3_client = boto3.client("s3", config=my_config)
+        s3_client.create_bucket(Bucket=bucket_name)
+        yield bucket_name
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+@pytest.fixture
+def sktime_custom_env(tmp_path):
+    """Create a conda environment and returns path to conda environment yml file."""
+    from mlflow.utils.environment import _mlflow_conda_env
+
+    conda_env = tmp_path.joinpath("conda_env.yml")
+    _mlflow_conda_env(conda_env, additional_pip_deps=["sktime"])
+    return conda_env
+
+
+@pytest.fixture(scope="module")
+def test_data_airline():
+    """Create sample data for univariate model without exogenous regressor."""
+    return load_airline()
+
+
+@pytest.fixture(scope="module")
+def test_data_longley():
+    """Create sample data for univariate model with exogenous regressor."""
+    y, X = load_longley()
+    y_train, y_test, X_train, X_test = temporal_train_test_split(y, X)
+    return y_train, y_test, X_train, X_test
+
+
+@pytest.fixture(scope="module")
+def auto_arima_model(test_data_airline):
+    """Create instance of fitted auto arima model."""
+    return AutoARIMA(sp=12, d=0, max_p=2, max_q=2, suppress_warnings=True).fit(
+        test_data_airline, fh=[1, 2, 3]
+    )
+
+
+@pytest.fixture(scope="module")
+def naive_forecaster_model_with_regressor(test_data_longley):
+    """Create instance of fitted naive forecaster model."""
+    y_train, _, X_train, _ = test_data_longley
+    model = NaiveForecaster()
+    return model.fit(y_train, X_train, fh=[1, 2, 3])
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+@pytest.mark.parametrize("serialization_format", ["pickle", "cloudpickle"])
+def test_auto_arima_model_save_and_load(
+    auto_arima_model, model_path, serialization_format
+):
+    """Test saving and loading of native sktime auto_arima_model."""
+    from sktime.utils import mlflow_sktime
+
+    mlflow_sktime.save_model(
+        sktime_model=auto_arima_model,
+        path=model_path,
+        serialization_format=serialization_format,
+    )
+    loaded_model = mlflow_sktime.load_model(
+        model_uri=model_path,
+    )
+
+    np.testing.assert_array_equal(auto_arima_model.predict(), loaded_model.predict())
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+@pytest.mark.parametrize("serialization_format", ["pickle", "cloudpickle"])
+def test_auto_arima_model_pyfunc_output(
+    auto_arima_model, model_path, serialization_format
+):
+    """Test auto arima prediction of loaded pyfunc model."""
+    from sktime.utils import mlflow_sktime
+
+    auto_arima_model.pyfunc_predict_conf = {
+        "predict_method": [
+            "predict",
+            "predict_interval",
+            "predict_quantiles",
+            "predict_var",
+        ]
+    }
+    mlflow_sktime.save_model(
+        sktime_model=auto_arima_model,
+        path=model_path,
+        serialization_format=serialization_format,
+    )
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path)
+
+    model_predict = auto_arima_model.predict()
+    model_predict_interval = auto_arima_model.predict_interval()
+    model_predict_interval.columns = flatten_multiindex(model_predict_interval)
+    model_predict_quantiles = auto_arima_model.predict_quantiles()
+    model_predict_quantiles.columns = flatten_multiindex(model_predict_quantiles)
+    model_predict_var = auto_arima_model.predict_var()
+    model_predictions = pd.concat(
+        [
+            model_predict,
+            model_predict_interval,
+            model_predict_quantiles,
+            model_predict_var,
+        ],
+        axis=1,
+        keys=[
+            "model_predict",
+            "model_predict_interval",
+            "model_predict_quantiles",
+            "model_predict_var",
+        ],
+    )
+
+    pyfunc_predict = loaded_pyfunc.predict(pd.DataFrame())
+    np.testing.assert_array_equal(model_predictions.values, pyfunc_predict.values)
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_auto_arima_model_pyfunc_with_params_output(auto_arima_model, model_path):
+    """Test auto arima prediction of loaded pyfunc model with parameters."""
+    from sktime.utils import mlflow_sktime
+
+    auto_arima_model.pyfunc_predict_conf = {
+        "predict_method": {
+            "predict": {},
+            "predict_interval": {"coverage": [0.1, 0.9]},
+            "predict_proba": {"quantiles": [0.1, 0.9], "marginal": True},
+            "predict_quantiles": {"alpha": [0.1, 0.9]},
+            "predict_var": {"cov": True},
+        }
+    }
+    mlflow_sktime.save_model(
+        sktime_model=auto_arima_model,
+        path=model_path,
+    )
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path)
+
+    model_predict = auto_arima_model.predict()
+    model_predict_interval = auto_arima_model.predict_interval(coverage=[0.1, 0.9])
+    model_predict_interval.columns = flatten_multiindex(model_predict_interval)
+    model_predict_proba_dist = auto_arima_model.predict_proba()
+    model_predict_proba = pd.DataFrame(model_predict_proba_dist.quantile([0.1, 0.9]))
+    model_predict_proba.index = model_predict_proba_dist.parameters["loc"].index
+    model_predict_quantiles = auto_arima_model.predict_quantiles(alpha=[0.1, 0.9])
+    model_predict_quantiles.columns = flatten_multiindex(model_predict_quantiles)
+    model_predict_var = auto_arima_model.predict_var(cov=True)
+    model_predictions = pd.concat(
+        [
+            model_predict,
+            model_predict_interval,
+            model_predict_proba,
+            model_predict_quantiles,
+            model_predict_var,
+        ],
+        axis=1,
+        keys=[
+            "model_predict",
+            "model_predict_interval",
+            "model_predict_proba",
+            "model_predict_quantiles",
+            "model_predict_var",
+        ],
+    )
+
+    pyfunc_predict = loaded_pyfunc.predict(pd.DataFrame())
+    np.testing.assert_array_equal(model_predictions.values, pyfunc_predict.values)
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_auto_arima_model_pyfunc_without_params_output(auto_arima_model, model_path):
+    """Test auto arima prediction of loaded pyfunc model without parameters."""
+    from sktime.utils import mlflow_sktime
+
+    auto_arima_model.pyfunc_predict_conf = {
+        "predict_method": {
+            "predict": {},
+            "predict_interval": {},
+            "predict_proba": {"quantiles": [0.1, 0.9]},
+            "predict_quantiles": {},
+            "predict_var": {},
+        }
+    }
+    mlflow_sktime.save_model(
+        sktime_model=auto_arima_model,
+        path=model_path,
+    )
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path)
+
+    model_predict = auto_arima_model.predict()
+    model_predict_interval = auto_arima_model.predict_interval()
+    model_predict_interval.columns = flatten_multiindex(model_predict_interval)
+    model_predict_proba_dist = auto_arima_model.predict_proba()
+    model_predict_proba = pd.DataFrame(model_predict_proba_dist.quantile([0.1, 0.9]))
+    model_predict_proba.index = model_predict_proba_dist.parameters["loc"].index
+    model_predict_quantiles = auto_arima_model.predict_quantiles()
+    model_predict_quantiles.columns = flatten_multiindex(model_predict_quantiles)
+    model_predict_var = auto_arima_model.predict_var()
+    model_predictions = pd.concat(
+        [
+            model_predict,
+            model_predict_interval,
+            model_predict_proba,
+            model_predict_quantiles,
+            model_predict_var,
+        ],
+        axis=1,
+        keys=[
+            "model_predict",
+            "model_predict_interval",
+            "model_predict_proba",
+            "model_predict_quantiles",
+            "model_predict_var",
+        ],
+    )
+
+    pyfunc_predict = loaded_pyfunc.predict(pd.DataFrame())
+    np.testing.assert_array_equal(model_predictions.values, pyfunc_predict.values)
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_auto_arima_model_pyfunc_without_conf_output(auto_arima_model, model_path):
+    """Test auto arima prediction of loaded pyfunc model without config."""
+    from sktime.utils import mlflow_sktime
+
+    delattr(auto_arima_model, "pyfunc_predict_conf") if hasattr(
+        auto_arima_model, "pyfunc_predict_conf"
+    ) else None
+    mlflow_sktime.save_model(
+        sktime_model=auto_arima_model,
+        path=model_path,
+    )
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path)
+
+    model_predict = auto_arima_model.predict()
+
+    pyfunc_predict = loaded_pyfunc.predict(pd.DataFrame())
+    np.testing.assert_array_equal(model_predict.values, pyfunc_predict.values)
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_naive_forecaster_model_with_regressor_pyfunc_output(
+    naive_forecaster_model_with_regressor, model_path, test_data_longley
+):
+    """Test naive forecaster prediction of loaded pyfunc model."""
+    from sktime.utils import mlflow_sktime
+
+    _, _, _, X_test = test_data_longley
+
+    naive_forecaster_model_with_regressor.pyfunc_predict_conf = {
+        "predict_method": [
+            "predict",
+            "predict_interval",
+            "predict_quantiles",
+            "predict_var",
+        ]
+    }
+    mlflow_sktime.save_model(
+        sktime_model=naive_forecaster_model_with_regressor, path=model_path
+    )
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path)
+
+    model_predict = naive_forecaster_model_with_regressor.predict(X=X_test)
+    model_predict_interval = naive_forecaster_model_with_regressor.predict_interval(
+        X=X_test
+    )
+    model_predict_interval.columns = flatten_multiindex(model_predict_interval)
+    model_predict_quantiles = naive_forecaster_model_with_regressor.predict_quantiles(
+        X=X_test
+    )
+    model_predict_quantiles.columns = flatten_multiindex(model_predict_quantiles)
+    model_predict_var = naive_forecaster_model_with_regressor.predict_var(X=X_test)
+    model_predictions = pd.concat(
+        [
+            model_predict,
+            model_predict_interval,
+            model_predict_quantiles,
+            model_predict_var,
+        ],
+        axis=1,
+        keys=[
+            "model_predict",
+            "model_predict_interval",
+            "model_predict_quantiles",
+            "model_predict_var",
+        ],
+    )
+
+    pyfunc_predict = loaded_pyfunc.predict(X_test)
+    np.testing.assert_array_equal(model_predictions.values, pyfunc_predict.values)
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+@pytest.mark.parametrize("use_signature", [True, False])
+@pytest.mark.parametrize("use_example", [True, False])
+def test_signature_and_examples_saved_correctly(
+    auto_arima_model, test_data_airline, model_path, use_signature, use_example
+):
+    """Test saving of mlflow signature and example for native sktime predict method."""
+    from mlflow.models import Model, infer_signature
+    from mlflow.models.utils import _read_example
+
+    from sktime.utils import mlflow_sktime
+
+    # Note: Signature inference fails on native model predict_interval/predict_quantiles
+    prediction = auto_arima_model.predict()
+    signature = (
+        infer_signature(test_data_airline, prediction) if use_signature else None
+    )
+    example = (
+        pd.DataFrame(test_data_airline[0:5].copy(deep=False)) if use_example else None
+    )
+    mlflow_sktime.save_model(
+        auto_arima_model, path=model_path, signature=signature, input_example=example
+    )
+    mlflow_model = Model.load(model_path)
+    assert signature == mlflow_model.signature
+    if example is None:
+        assert mlflow_model.saved_input_example_info is None
+    else:
+        r_example = _read_example(mlflow_model, model_path).copy(deep=False)
+        np.testing.assert_array_equal(r_example, example)
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+@pytest.mark.parametrize("use_signature", [True, False])
+def test_predict_var_signature_saved_correctly(
+    auto_arima_model, test_data_airline, model_path, use_signature
+):
+    """Test saving of mlflow signature for native sktime predict_var method."""
+    from mlflow.models import Model, infer_signature
+
+    from sktime.utils import mlflow_sktime
+
+    # Note: Signature inference fails on native model predict_interval/predict_quantiles
+    prediction = auto_arima_model.predict_var()
+    signature = (
+        infer_signature(test_data_airline, prediction) if use_signature else None
+    )
+    mlflow_sktime.save_model(auto_arima_model, path=model_path, signature=signature)
+    mlflow_model = Model.load(model_path)
+    assert signature == mlflow_model.signature
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+@pytest.mark.parametrize("use_signature", [True, False])
+@pytest.mark.parametrize("use_example", [True, False])
+def test_signature_and_example_for_pyfunc_predict(
+    auto_arima_model, model_path, test_data_airline, use_signature, use_example
+):
+    """Test saving of mlflow signature and example for pyfunc predict."""
+    from mlflow.models import Model, infer_signature
+    from mlflow.models.utils import _read_example
+
+    from sktime.utils import mlflow_sktime
+
+    model_path_primary = model_path.joinpath("primary")
+    model_path_secondary = model_path.joinpath("secondary")
+    auto_arima_model.pyfunc_predict_conf = {
+        "predict_method": {
+            "predict": None,
+            "predict_interval": {"coverage": [0.1, 0.9]},
+            "predict_quantiles": {"alpha": [0.1, 0.9]},
+            "predict_var": {"cov": False},
+        }
+    }
+    mlflow_sktime.save_model(sktime_model=auto_arima_model, path=model_path_primary)
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path_primary)
+
+    forecast = loaded_pyfunc.predict(pd.DataFrame())
+    signature = infer_signature(test_data_airline, forecast) if use_signature else None
+    example = (
+        pd.DataFrame(test_data_airline[0:5].copy(deep=False)) if use_example else None
+    )
+    mlflow_sktime.save_model(
+        auto_arima_model,
+        path=model_path_secondary,
+        signature=signature,
+        input_example=example,
+    )
+    mlflow_model = Model.load(model_path_secondary)
+    assert signature == mlflow_model.signature
+    if example is None:
+        assert mlflow_model.saved_input_example_info is None
+    else:
+        r_example = _read_example(mlflow_model, model_path_secondary).copy(deep=False)
+        np.testing.assert_array_equal(r_example, example)
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_load_from_remote_uri_succeeds(auto_arima_model, model_path, mock_s3_bucket):
+    """Test loading native sktime model from mock S3 bucket."""
+    from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
+
+    from sktime.utils import mlflow_sktime
+
+    mlflow_sktime.save_model(sktime_model=auto_arima_model, path=model_path)
+
+    artifact_root = f"s3://{mock_s3_bucket}"
+    artifact_path = "model"
+    artifact_repo = S3ArtifactRepository(artifact_root)
+    artifact_repo.log_artifacts(model_path, artifact_path=artifact_path)
+
+    model_uri = os.path.join(artifact_root, artifact_path)
+    reloaded_sktime_model = mlflow_sktime.load_model(model_uri=model_uri)
+
+    np.testing.assert_array_equal(
+        auto_arima_model.predict(),
+        reloaded_sktime_model.predict(),
+    )
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+@pytest.mark.parametrize("should_start_run", [True, False])
+@pytest.mark.parametrize("serialization_format", ["pickle", "cloudpickle"])
+def test_log_model(auto_arima_model, tmp_path, should_start_run, serialization_format):
+    """Test logging and reloading sktime model."""
+    import mlflow
+    from mlflow import pyfunc
+    from mlflow.models import Model
+    from mlflow.tracking.artifact_utils import _download_artifact_from_uri
+    from mlflow.utils.environment import _mlflow_conda_env
+
+    from sktime.utils import mlflow_sktime
+
+    try:
+        if should_start_run:
+            mlflow.start_run()
+        artifact_path = "sktime"
+        conda_env = tmp_path.joinpath("conda_env.yaml")
+        _mlflow_conda_env(conda_env, additional_pip_deps=["sktime"])
+        model_info = mlflow_sktime.log_model(
+            sktime_model=auto_arima_model,
+            artifact_path=artifact_path,
+            conda_env=str(conda_env),
+            serialization_format=serialization_format,
+        )
+        model_uri = f"runs:/{mlflow.active_run().info.run_id}/{artifact_path}"
+        assert model_info.model_uri == model_uri
+        reloaded_model = mlflow_sktime.load_model(
+            model_uri=model_uri,
+        )
+        np.testing.assert_array_equal(
+            auto_arima_model.predict(), reloaded_model.predict()
+        )
+        model_path = Path(_download_artifact_from_uri(artifact_uri=model_uri))
+        model_config = Model.load(str(model_path.joinpath("MLmodel")))
+        assert pyfunc.FLAVOR_NAME in model_config.flavors
+    finally:
+        mlflow.end_run()
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_log_model_calls_register_model(auto_arima_model, tmp_path):
+    """Test log model calls register model."""
+    import mlflow
+    from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
+    from mlflow.utils.environment import _mlflow_conda_env
+
+    from sktime.utils import mlflow_sktime
+
+    artifact_path = "sktime"
+    register_model_patch = mock.patch("mlflow.register_model")
+    with mlflow.start_run(), register_model_patch:
+        conda_env = tmp_path.joinpath("conda_env.yaml")
+        _mlflow_conda_env(conda_env, additional_pip_deps=["sktime"])
+        mlflow_sktime.log_model(
+            sktime_model=auto_arima_model,
+            artifact_path=artifact_path,
+            conda_env=str(conda_env),
+            registered_model_name="SktimeModel",
+        )
+        model_uri = f"runs:/{mlflow.active_run().info.run_id}/{artifact_path}"
+        mlflow.register_model.assert_called_once_with(
+            model_uri,
+            "SktimeModel",
+            await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
+        )
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_log_model_no_registered_model_name(auto_arima_model, tmp_path):
+    """Test log model calls register model without registered model name."""
+    import mlflow
+    from mlflow.utils.environment import _mlflow_conda_env
+
+    from sktime.utils import mlflow_sktime
+
+    artifact_path = "sktime"
+    register_model_patch = mock.patch("mlflow.register_model")
+    with mlflow.start_run(), register_model_patch:
+        conda_env = tmp_path.joinpath("conda_env.yaml")
+        _mlflow_conda_env(conda_env, additional_pip_deps=["sktime"])
+        mlflow_sktime.log_model(
+            sktime_model=auto_arima_model,
+            artifact_path=artifact_path,
+            conda_env=str(conda_env),
+        )
+        mlflow.register_model.assert_not_called()
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_pyfunc_raises_invalid_attribute_type(auto_arima_model, model_path):
+    """Test pyfunc raises exception with invalid attribute type."""
+    from mlflow.exceptions import MlflowException
+
+    from sktime.utils import mlflow_sktime
+
+    auto_arima_model.pyfunc_predict_conf = ["predict"]
+    mlflow_sktime.save_model(sktime_model=auto_arima_model, path=model_path)
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path)
+
+    with pytest.raises(
+        MlflowException,
+        match=f"Attribute {mlflow_sktime.PYFUNC_PREDICT_CONF} must be of type dict.",
+    ):
+        loaded_pyfunc.predict(pd.DataFrame())
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_pyfunc_raises_invalid_dict_key(auto_arima_model, model_path):
+    """Test pyfunc raises exception with invalid dict key."""
+    from mlflow.exceptions import MlflowException
+
+    from sktime.utils import mlflow_sktime
+
+    auto_arima_model.pyfunc_predict_conf = {"prediction_method": ["predict"]}
+    mlflow_sktime.save_model(sktime_model=auto_arima_model, path=model_path)
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path)
+
+    with pytest.raises(
+        MlflowException,
+        match=f"Attribute {mlflow_sktime.PYFUNC_PREDICT_CONF} must contain ",
+    ):
+        loaded_pyfunc.predict(pd.DataFrame())
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_pyfunc_raises_invalid_dict_value_type(auto_arima_model, model_path):
+    """Test pyfunc raises exception with invalid dict value type."""
+    from mlflow.exceptions import MlflowException
+
+    from sktime.utils import mlflow_sktime
+
+    auto_arima_model.pyfunc_predict_conf = {"predict_method": "predict"}
+    mlflow_sktime.save_model(sktime_model=auto_arima_model, path=model_path)
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path)
+
+    with pytest.raises(
+        MlflowException, match="Dictionary value must be of type dict or list"
+    ):
+        loaded_pyfunc.predict(pd.DataFrame())
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_pyfunc_raises_invalid_dict_value(auto_arima_model, model_path):
+    """Test pyfunc raises exception with invalid dict value."""
+    from mlflow.exceptions import MlflowException
+
+    from sktime.utils import mlflow_sktime
+
+    auto_arima_model.pyfunc_predict_conf = {"predict_method": ["forecast"]}
+    mlflow_sktime.save_model(sktime_model=auto_arima_model, path=model_path)
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path)
+
+    with pytest.raises(
+        MlflowException,
+        match=f"The provided {mlflow_sktime.PYFUNC_PREDICT_CONF_KEY} values must be ",
+    ):
+        loaded_pyfunc.predict(pd.DataFrame())
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_pyfunc_predict_proba_raises_invalid_attribute_type(
+    auto_arima_model, model_path
+):
+    """Test pyfunc predict_proba raises exception with invalid attribute type."""
+    from mlflow.exceptions import MlflowException
+
+    from sktime.utils import mlflow_sktime
+
+    auto_arima_model.pyfunc_predict_conf = {"predict_method": ["predict_proba"]}
+    mlflow_sktime.save_model(sktime_model=auto_arima_model, path=model_path)
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path)
+
+    with pytest.raises(
+        MlflowException,
+        match=f"Method {mlflow_sktime.SKTIME_PREDICT_PROBA} requires passing a "
+        f"dictionary.",
+    ):
+        loaded_pyfunc.predict(pd.DataFrame())
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("mlflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_pyfunc_predict_proba_raises_invalid_dict_value(auto_arima_model, model_path):
+    """Test pyfunc predict_proba raises exception with invalid dict value."""
+    from mlflow.exceptions import MlflowException
+
+    from sktime.utils import mlflow_sktime
+
+    auto_arima_model.pyfunc_predict_conf = {
+        "predict_method": {"predict_proba": {"marginal": True}}
+    }
+    mlflow_sktime.save_model(sktime_model=auto_arima_model, path=model_path)
+    loaded_pyfunc = mlflow_sktime.pyfunc.load_model(model_uri=model_path)
+
+    with pytest.raises(
+        MlflowException,
+        match=f"Method {mlflow_sktime.SKTIME_PREDICT_PROBA} requires passing "
+        f"quantile values.",
+    ):
+        loaded_pyfunc.predict(pd.DataFrame())