Repository: mlflow/mlflow

Issue:
[FR] [Roadmap] Create a detailed example of creating a custom model flavor
## MLflow Roadmap Item

This is an MLflow Roadmap item that has been prioritized by the MLflow maintainers. We’ve identified this feature as a highly requested addition to the MLflow package based on community feedback. 
We're seeking a community contribution for the implementation of this feature and will enthusiastically support the development and review of a submitted PR for this.

## Contribution Note

As with other roadmap items, there may be a desire for multiple contributors to work on an issue. While we don’t discourage collaboration, we strongly encourage that a primary contributor is assigned to roadmap issues to simplify the merging process. 
The items on the roadmap are of a high priority. Due to the wide-spread demand of roadmap features, we encourage potential contributors to only agree to take on the work of creating a PR, making changes, and ensuring that test coverage is adequately created for the feature if they are willing and able to see the implementation through to a merged state. 

## Feature scope

This roadmap feature’s complexity is classified as:

- [X] `good-first-issue`: This feature is limited in complexity and effort required to implement.
- [ ] `simple`: This feature does not require a large amount of effort to implement and / or is clear enough to not need a design discussion with maintainers.
- [ ] `involved`: This feature will require a substantial amount of development effort but does not require an agreed-upon design from the maintainers. The feedback given during the PR phase may be involved and necessitate multiple iterations before approval. (Please bear with us as we collaborate with you to make a great contribution)
- [ ] `design-recommended`: This is a substantial feature that should have a design document approved prior to working on an implementation (to save your time, not ours). After agreeing to work on this feature, a maintainer will be assigned to support you throughout the development process.

## Proposal Summary

The current example for custom model flavors is inadequate as a guide. This is a common question that gets raised amongst users. 
This FR is a request to create a much more in-depth example of a custom model flavor that shows its construction in-line within the [docs](https://mlflow.org/docs/latest/models.html#custom-flavors), provides explanations for what is required, and shows usage of it with a screenshot of the model within the UI (namely the serialized model artifact within the run page). Custom flavors should be introduced as separate GitHub repositories with documentation provided in https://mlflow.org/docs/latest/plugins.html#community-plugins.

## Motivation
> #### What is the use case for this feature? 
  Provide an example and better explanation for a common ask that users struggle with.

> #### Why is this use case valuable to support for MLflow users in general? 
  Many users ask how to do this to incorporate a non-officially supported model flavor or request inclusion of esoteric libraries that will not be considered for official inclusion into MLflow due to low usage.

### What component(s), interfaces, languages, and integrations does this feature affect?
Components 
- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [X] `area/docs`: MLflow documentation pages
- [X] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interfaces
- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Languages 
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations



Discussion:
For any questions, concerns, or clarification on implementing this issue, please ping @WeichenXu123 
@BenWilson2 I would love to do this task, can you assign the issue to me?
Hi @ikrizanic 
Few days ago [Lakshika Parihar](https://app.slack.com/team/U01TDNYQU21) ping me and also want to do this task. Let's wait for his response first.
Hey @WeichenXu123 , Yeah i would love to do this task , you can assign the issue to me
@lakshikaparihar are you still working on this issue??
@AMMAR-62 Yeah, it's assigned to me today only. so I will start working on it.
Hi @lakshikaparihar, are you still working on this item? If not, @AMMAR-62 , are you interested in contributing it?
@dbczumar If @lakshikaparihar or @AMMAR-62 won't be able to do it at the moment, I would be happy to contribute 
@ikrizanic Sounds good! I will assign you to the issue :). @AMMAR-62 @lakshikaparihar , feel free to collaborate with @ikrizanic.
@dbczumar would be very happy to work on this issue if @ikrizanic cannot devote any capacity at the moment. 

I could start dedicating significant capacity beginning of December, but would definitely need some guidance/alignment before starting to contribute 
@benjaminbluhm @dbczumar My apologies, I forgot about this issue and currently don't have any capacity to work on this, so I'll be happy to hand over it to you. 
@BenWilson2 @dbczumar I would have a few questions to better understand your idea for the custom flavor example to be added to the docs. 

Do you suggest we create the custom flavor example using a simple fake model (along the lines of ["add n" model](https://mlflow.org/docs/latest/models.html#example-creating-a-custom-add-n-model) in the docs) or would you prefer we create this for a well established ML library?

My thought was to first create the custom flavor example for this FR in a separate github repo and once this is implemented I could start working on the PR for this FR- do you think this strategy makes sense or would you suggest a different route?

And would it make sense to align on a template for custom flavor repo that can be followed by others (scikit learn has [project template](https://github.com/scikit-learn-contrib/project-template) for scikit-learn compatible extensions so I thought having some kind of template could also help users to create their own mlflow custom flavor repo)? I can look at the custom flavor PRs in mlflow repo but at this moment I am not quite sure what would be a suitable folder and module structure for a custom flavor in a dedicated repo, would be great to get some guidance from you on this aspect.

I might have some more questions but rather wait for your initial feedback.

Thanks a lot!
Hi @benjaminbluhm I think the best place for it to live for a step-by-step working example would be within the `examples` section of Mlflow, here: https://github.com/mlflow/mlflow/tree/master/examples . 
The main crux of this FR isn't just a code example though. It's in building a step-by-step walkthrough of the components, nuances, and 'gotchas' surrounding custom model creation and utilization to support use cases that aren't directly supported natively within named flavors. 
This could be expanding the functionality of a named flavor (the simpler, the better; no need to try to demonstrate graph embeddings on a TF core model) or showing support for an unsupported but popular modeling framework (if going down this route, I'd check in here first to make sure that it's something that is installable and testable within the examples test suite). 
Hi @dbczumar @benjaminbluhm @BenWilson2 , sorry forgot about this FR I do
have One Usecase in mind i.e "How to do prediction using .predict() for any
type of model/framework", lemme know if you want to work on it, would be
happy to collaborate.

On Thu, 24 Nov 2022, 01:35 Ben Wilson, ***@***.***> wrote:

> Hi @benjaminbluhm <https://github.com/benjaminbluhm> I think the best
> place for it to live for a step-by-step working example would be within the
> examples section of Mlflow, here:
> https://github.com/mlflow/mlflow/tree/master/examples .
> The main crux of this FR isn't just a code example though. It's in
> building a step-by-step walkthrough of the components, nuances, and
> 'gotchas' surrounding custom model creation and utilization to support use
> cases that aren't directly supported natively within named flavors.
> This could be expanding the functionality of a named flavor (the simpler,
> the better; no need to try to demonstrate graph embeddings on a TF core
> model) or showing support for an unsupported but popular modeling framework
> (if going down this route, I'd check in here first to make sure that it's
> something that is installable and testable within the examples test suite).
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/mlflow/mlflow/issues/6102#issuecomment-1325598964>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AHT754ZY7JAHV3SLNU25SDTWJZ2JVANCNFSM5ZDMALZA>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>

Hi @BenWilson2 thanks for clarification and explaining the scope of having in-depth step-by-step walkthrough. After your feedback yesterday I connected with @aiwalter from [sktime](https://github.com/sktime/sktime) team, plan is to create custom model flavor for sktime (see [feature request](https://github.com/sktime/sktime/issues/1215) in their repo). I think sktime could be a nice candidate for this FR given it is a very popular time series ML library. In case you like the idea and you can confirm that it is installable and testable within the examples test suite my suggestion would be to first work on custom flavor implementation in sktime and afterwards I could use my insights to work on this FR to provide a detailed step-by-step walkthrough for sktime.

In case you like more the idea from @lakshikaparihar I am also happy to hand over the issue and see if/how I can contribute 



Hi @BenWilson2 @benjaminbluhm, sktime sounds like a great candidate framework. @benjaminbluhm is this still something you'd be interested in contributing?
Hi @dbczumar for sktime mlflow implementation we have pretty much followed the standard from mlflow built-in model flavors for `save_model()`, `log_model()` and `load_model()` functions (very little modification). The `_SktimeModelWrapper` class contains the custom inference logic - this is where I could presumably elaborate the most in the step-by-step walk-through.

If this sounds good to you I would be happy to contribute. To my understanding the PR will affect the following components:

- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code

I guess first step of the PR could be to add the model flavor code in a `flavor.py` module under `mlflow/examples/sktime` folder and also include a `train.py` module similar to existing examples.

Writing the walk-through as part of the docs I would see a second step. Please let me know if you think this plan makes sense or you have a different proposal.
Hi @benjaminbluhm, apologies for the delay. That plan sounds great!  
Hi @dbczumar, sounds good! Will raise a draft PR in the coming weeks when I have a first version of the custom flavor code ready to get some preliminary feedback

Resolving patch:
diff --git a/examples/sktime/flavor.py b/examples/sktime/flavor.py
new file mode 100644
--- /dev/null
+++ b/examples/sktime/flavor.py
@@ -0,0 +1,540 @@
+"""The ``flavor`` module provides an example for a custom model flavor for ``sktime`` library.
+
+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.
+
+    Note that pickle serialization requires using the same python environment (version) in
+    whatever environment you're going to use this model for inference to ensure that the model
+    will load with appropriate version of pickle.
+mlflow.pyfunc
+    Produced for use by generic pyfunc-based deployment tools and batch inference.
+
+    The interface for utilizing a ``sktime`` model loaded as a ``pyfunc`` type for generating
+    forecast predictions uses a *single-row* ``Pandas DataFrame`` configuration argument. The
+    following columns in this configuration ``Pandas DataFrame`` are supported:
+
+    * ``predict_method`` (required) - specifies the ``sktime`` predict method. The
+        supported predict methods in this example flavor are ``predict``, ``predict_interval``,
+        ``predict_quantiles``, ``predict_var``. Additional methods (e.g. ``predict_proba``) could
+        be added in a similar fashion.
+    * ``fh`` (optional) - specifies the number of future periods to generate starting from
+        the last datetime value of the training dataset, utilizing the frequency of the input
+        training series when the model was trained. (for example, if the training data series
+        elements represent one value per hour, in order to forecast 3 hours of future data, set
+        the column ``fh`` to ``[1,2,3]``. If the paramter is not provided it must be passed
+        during fit(). (Default: ``None``)
+    * ``X`` (optional) - exogenous regressor values as a 2D numpy ndarray or list of values for future
+        time period events. For more information, read the underlying library explanation
+        https://www.sktime.net/en/latest/examples/AA_datatypes_and_datasets.html#Section-1:-in-memory-data-containers.
+        (Default: ``None``)
+    * ``coverage`` (optional) - the nominal coverage value for calculating prediction interval
+        forecasts. Can only be provided in combination with predict method ``predict_interval``.
+        (Default: ``0.9``)
+    * ``alpha`` (optional) - the probability value for calculating prediction quantile forecasts.
+        Can only be provided in combination with predict method ``predict_quantiles``.
+        (Default: ``None``)
+    * ``cov`` (optional) - if True, computes covariance matrix forecast.
+        Can only be provided in combination with predict method ``predict_var``.
+        (Default: ``False``)
+
+An example configuration for the ``pyfunc`` predict of a ``sktime`` model is shown below, using an
+interval forecast with nominal coverage value ``[0.9,0.95]``, a future forecast horizon of 3 periods,
+and no exogenous regressor elements:
+
+====== ================= ============ ========
+Index  predict_method    coverage     fh
+====== ================= ============ ========
+0      predict_interval  [0.9,0.95]   [1,2,3]
+====== ================= ============ ========
+"""
+import logging
+import os
+import pickle
+
+import flavor
+import mlflow
+import numpy as np
+import pandas as pd
+import sktime
+import yaml
+from mlflow import pyfunc
+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 INTERNAL_ERROR, INVALID_PARAMETER_VALUE
+from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
+from mlflow.tracking.artifact_utils import _download_artifact_from_uri
+from mlflow.utils.environment import (
+    _CONDA_ENV_FILE_NAME,
+    _CONSTRAINTS_FILE_NAME,
+    _PYTHON_ENV_FILE_NAME,
+    _REQUIREMENTS_FILE_NAME,
+    _mlflow_conda_env,
+    _process_conda_env,
+    _process_pip_requirements,
+    _PythonEnv,
+    _validate_env_arguments,
+)
+from mlflow.utils.file_utils import write_to
+from mlflow.utils.model_utils import (
+    _add_code_from_conf_to_system_path,
+    _get_flavor_configuration,
+    _validate_and_copy_code_paths,
+    _validate_and_prepare_target_save_path,
+)
+from mlflow.utils.requirements_utils import _get_pinned_requirement
+from sktime.utils.multiindex import flatten_multiindex
+
+FLAVOR_NAME = "sktime"
+
+SKTIME_PREDICT = "predict"
+SKTIME_PREDICT_INTERVAL = "predict_interval"
+SKTIME_PREDICT_QUANTILES = "predict_quantiles"
+SKTIME_PREDICT_VAR = "predict_var"
+SUPPORTED_SKTIME_PREDICT_METHODS = [
+    SKTIME_PREDICT,
+    SKTIME_PREDICT_INTERVAL,
+    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.
+    """
+    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()`
+    """
+    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,
+):
+    """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"
+    """
+    _validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
+
+    if serialization_format not in SUPPORTED_SERIALIZATION_FORMATS:
+        raise MlflowException(
+            message=(
+                f"Unrecognized serialization format: {serialization_format}. "
+                "Please specify one of the following 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="flavor",
+        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)
+            inferred_reqs = mlflow.models.infer_pip_requirements(
+                path, FLAVOR_NAME, fallback=default_reqs
+            )
+            default_reqs = sorted(set(inferred_reqs).union(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=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
+    pip_requirements=None,
+    extra_pip_requirements=None,
+    serialization_format=SERIALIZATION_FORMAT_PICKLE,
+    **kwargs,
+):
+    """
+    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.
+    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.
+    """
+    return Model.log(
+        artifact_path=artifact_path,
+        flavor=flavor,
+        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.
+    """
+    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):
+    with open(path, "wb") as out:
+        if serialization_format == SERIALIZATION_FORMAT_PICKLE:
+            pickle.dump(model, out)
+        else:
+            import cloudpickle
+
+            cloudpickle.dump(model, out)
+
+
+def _load_model(path, serialization_format):
+    with open(path, "rb") as pickled_model:
+        if serialization_format == SERIALIZATION_FORMAT_PICKLE:
+            return pickle.load(pickled_model)
+        elif serialization_format == SERIALIZATION_FORMAT_CLOUDPICKLE:
+            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.
+
+    """
+    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):
+        self.sktime_model = sktime_model
+
+    def predict(self, dataframe) -> pd.DataFrame:
+        df_schema = dataframe.columns.values.tolist()
+
+        if len(dataframe) > 1:
+            raise MlflowException(
+                f"The provided prediction pd.DataFrame contains {len(dataframe)} rows. "
+                "Only 1 row should be supplied.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        # Convert the configuration dataframe into a dictionary to simplify the
+        # extraction of parameters passed to the sktime predcition methods.
+        attrs = dataframe.to_dict(orient="index").get(0)
+        predict_method = attrs.get("predict_method")
+
+        if not predict_method:
+            raise MlflowException(
+                f"The provided prediction configuration pd.DataFrame columns ({df_schema}) do not "
+                "contain the required column `predict_method` for specifying the prediction method.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        if predict_method not in SUPPORTED_SKTIME_PREDICT_METHODS:
+            raise MlflowException(
+                "Invalid `predict_method` value."
+                f"The supported prediction methods are {SUPPORTED_SKTIME_PREDICT_METHODS}",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        # For inference parameters 'fh', 'X', 'coverage', 'alpha', and 'cov'
+        # the respective sktime default value is used if the value was not
+        # provided in the configuration dataframe.
+        fh = attrs.get("fh", None)
+
+        # Any model that is trained with exogenous regressor elements will need
+        # to provide `X` entries as a numpy ndarray to the predict method.
+        X = attrs.get("X", None)
+
+        # When the model is served via REST API the exogenous regressor must be
+        # provided as a list to the configuration DataFrame to be JSON serializable.
+        # Below we convert the list back to ndarray type as required by sktime
+        # predict methods.
+        if isinstance(X, list):
+            X = np.array(X)
+
+        # For illustration purposes only a subset of the available sktime prediction
+        # methods is exposed. Additional methods (e.g. predict_proba) could be added
+        # in a similar fashion.
+        if predict_method == SKTIME_PREDICT:
+            predictions = self.sktime_model.predict(fh=fh, X=X)
+
+        if predict_method == SKTIME_PREDICT_INTERVAL:
+            coverage = attrs.get("coverage", 0.9)
+            predictions = self.sktime_model.predict_interval(fh=fh, X=X, coverage=coverage)
+
+        if predict_method == SKTIME_PREDICT_QUANTILES:
+            alpha = attrs.get("alpha", None)
+            predictions = self.sktime_model.predict_quantiles(fh=fh, X=X, alpha=alpha)
+
+        if predict_method == SKTIME_PREDICT_VAR:
+            cov = attrs.get("cov", False)
+            predictions = self.sktime_model.predict_var(fh=fh, X=X, cov=cov)
+
+        # Methods predict_interval() and predict_quantiles() return a pandas
+        # MultiIndex column structure. As MLflow signature inference does not
+        # support MultiIndex column structure the columns must be flattened.
+        if predict_method in [SKTIME_PREDICT_INTERVAL, SKTIME_PREDICT_QUANTILES]:
+            predictions.columns = flatten_multiindex(predictions)
+
+        return predictions
diff --git a/examples/sktime/score_model.py b/examples/sktime/score_model.py
new file mode 100644
--- /dev/null
+++ b/examples/sktime/score_model.py
@@ -0,0 +1,37 @@
+import pandas as pd
+import requests
+
+from sktime.datasets import load_longley
+from sktime.forecasting.model_selection import temporal_train_test_split
+
+y, X = load_longley()
+y_train, y_test, X_train, X_test = temporal_train_test_split(y, X)
+
+# Define local host and endpoint url
+host = "127.0.0.1"
+url = f"http://{host}:5000/invocations"
+
+# Model scoring via REST API requires transforming the configuration DataFrame
+# into JSON format. As numpy ndarray type is not JSON serializable we need to
+# convert the exogenous regressor into a list. The wrapper instance will convert
+# the list back to ndarray type as required by sktime predict methods. For more
+# details read the MLflow deployment API reference.
+# (https://mlflow.org/docs/latest/models.html#deploy-mlflow-models)
+X_test_list = X_test.to_numpy().tolist()
+predict_conf = pd.DataFrame(
+    [
+        {
+            "fh": [1, 2, 3, 4],
+            "predict_method": "predict_interval",
+            "coverage": [0.9, 0.95],
+            "X": X_test_list,
+        }
+    ]
+)
+
+# Create dictionary with pandas DataFrame in the split orientation
+json_data = {"dataframe_split": predict_conf.to_dict(orient="split")}
+
+# Score model
+response = requests.post(url, json=json_data)
+print(f"\nPyfunc 'predict_interval':\n${response.json()}")
diff --git a/examples/sktime/train.py b/examples/sktime/train.py
new file mode 100644
--- /dev/null
+++ b/examples/sktime/train.py
@@ -0,0 +1,86 @@
+import json
+
+import flavor
+import pandas as pd
+from sktime.datasets import load_longley
+from sktime.forecasting.model_selection import temporal_train_test_split
+from sktime.forecasting.naive import NaiveForecaster
+from sktime.performance_metrics.forecasting import (
+    mean_absolute_error,
+    mean_absolute_percentage_error,
+)
+
+import mlflow
+
+ARTIFACT_PATH = "model"
+
+with mlflow.start_run() as run:
+    y, X = load_longley()
+    y_train, y_test, X_train, X_test = temporal_train_test_split(y, X)
+
+    forecaster = NaiveForecaster()
+    forecaster.fit(
+        y_train,
+        X=X_train,
+    )
+
+    # Extract parameters
+    parameters = forecaster.get_params()
+
+    # Evaluate model
+    y_pred = forecaster.predict(fh=[1, 2, 3, 4], X=X_test)
+    metrics = {
+        "mae": mean_absolute_error(y_test, y_pred),
+        "mape": mean_absolute_percentage_error(y_test, y_pred),
+    }
+
+    print(f"Parameters: \n{json.dumps(parameters, indent=2)}")
+    print(f"Metrics: \n{json.dumps(metrics, indent=2)}")
+
+    # Log parameters and metrics
+    mlflow.log_params(parameters)
+    mlflow.log_metrics(metrics)
+
+    # Log model using custom model flavor with pickle serialization (default).
+    # Note that pickle serialization requires using the same python environment
+    # (version) in whatever environment you're going to use this model for
+    # inference to ensure that the model will load with appropriate version of
+    # pickle.
+    flavor.log_model(
+        sktime_model=forecaster,
+        artifact_path=ARTIFACT_PATH,
+        serialization_format="pickle",
+    )
+    model_uri = mlflow.get_artifact_uri(ARTIFACT_PATH)
+
+# Load model in native sktime flavor and pyfunc flavor
+loaded_model = flavor.load_model(model_uri=model_uri)
+loaded_pyfunc = flavor.pyfunc.load_model(model_uri=model_uri)
+
+# Convert test data to 2D numpy array so it can be passed to pyfunc predict using
+# a single-row Pandas DataFrame configuration argument
+X_test_array = X_test.to_numpy()
+
+# Create configuration DataFrame for interval forecast with nominal coverage
+# value [0.9,0.95], future forecast horizon of 4 periods, and exogenous regressor.
+# Read more in the flavor.py module docstrings about the possible configurations.
+predict_conf = pd.DataFrame(
+    [
+        {
+            "fh": [1, 2, 3, 4],
+            "predict_method": "predict_interval",
+            "coverage": [0.9, 0.95],
+            "X": X_test_array,
+        }
+    ]
+)
+
+# Generate interval forecasts with native sktime flavor and pyfunc flavor
+print(
+    f"\nNative sktime 'predict_interval':\n${loaded_model.predict_interval(fh=[1, 2, 3], X=X_test, coverage=[0.9, 0.95])}"
+)
+print(f"\nPyfunc 'predict_interval':\n${loaded_pyfunc.predict(predict_conf)}")
+
+# Print the run id wich is used for serving the model to a local REST API endpoint
+# in the score_model.py module
+print(f"\nMLflow run id:\n{run.info.run_id}")


Test patch:
diff --git a/examples/sktime/test_sktime_model_export.py b/examples/sktime/test_sktime_model_export.py
new file mode 100644
--- /dev/null
+++ b/examples/sktime/test_sktime_model_export.py
@@ -0,0 +1,416 @@
+"""Tests for sktime custom model flavor."""
+
+import flavor
+import os
+from pathlib import Path
+from unittest import mock
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from sktime.datatypes import convert
+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
+import boto3
+import moto
+from botocore.config import Config
+from mlflow.utils.environment import _mlflow_conda_env
+from mlflow.models import Model, infer_signature
+from mlflow.models.utils import _read_example
+from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
+import mlflow
+from mlflow import pyfunc
+from mlflow.tracking.artifact_utils import _download_artifact_from_uri
+from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
+from mlflow.exceptions import MlflowException
+
+FH = [1, 2, 3]
+COVERAGE = [0.1, 0.5, 0.9]
+ALPHA = [0.1, 0.5, 0.9]
+COV = False
+
+
+@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.fixture
+def sktime_custom_env(tmp_path):
+    """Create a conda environment and returns path to conda environment yml file."""
+    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 data_airline():
+    """Create sample data for univariate model without exogenous regressor."""
+    return load_airline()
+
+
+@pytest.fixture(scope="module")
+def 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(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(data_airline)
+
+
+@pytest.fixture(scope="module")
+def naive_forecaster_model_with_regressor(data_longley):
+    """Create instance of fitted naive forecaster model."""
+    y_train, _, X_train, _ = data_longley
+    model = NaiveForecaster()
+    return model.fit(y_train, X_train)
+
+
+@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."""
+    flavor.save_model(
+        sktime_model=auto_arima_model,
+        path=model_path,
+        serialization_format=serialization_format,
+    )
+    loaded_model = flavor.load_model(
+        model_uri=model_path,
+    )
+
+    np.testing.assert_array_equal(auto_arima_model.predict(fh=FH), loaded_model.predict(fh=FH))
+
+
+@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."""
+    flavor.save_model(
+        sktime_model=auto_arima_model,
+        path=model_path,
+        serialization_format=serialization_format,
+    )
+    loaded_pyfunc = flavor.pyfunc.load_model(model_uri=model_path)
+
+    model_predict = auto_arima_model.predict(fh=FH)
+    predict_conf = pd.DataFrame([{"fh": FH, "predict_method": "predict"}])
+    pyfunc_predict = loaded_pyfunc.predict(predict_conf)
+    np.testing.assert_array_equal(model_predict, pyfunc_predict)
+
+    model_predict_interval = auto_arima_model.predict_interval(fh=FH, coverage=COVERAGE)
+    predict_interval_conf = pd.DataFrame(
+        [
+            {
+                "fh": FH,
+                "predict_method": "predict_interval",
+                "coverage": COVERAGE,
+            }
+        ]
+    )
+    pyfunc_predict_interval = loaded_pyfunc.predict(predict_interval_conf)
+    np.testing.assert_array_equal(model_predict_interval.values, pyfunc_predict_interval.values)
+
+    model_predict_quantiles = auto_arima_model.predict_quantiles(fh=FH, alpha=ALPHA)
+    predict_quantiles_conf = pd.DataFrame(
+        [
+            {
+                "fh": FH,
+                "predict_method": "predict_quantiles",
+                "alpha": ALPHA,
+            }
+        ]
+    )
+    pyfunc_predict_quantiles = loaded_pyfunc.predict(predict_quantiles_conf)
+    np.testing.assert_array_equal(model_predict_quantiles.values, pyfunc_predict_quantiles.values)
+
+    model_predict_var = auto_arima_model.predict_var(fh=FH, cov=COV)
+    predict_var_conf = pd.DataFrame([{"fh": FH, "predict_method": "predict_var", "cov": COV}])
+    pyfunc_predict_var = loaded_pyfunc.predict(predict_var_conf)
+    np.testing.assert_array_equal(model_predict_var.values, pyfunc_predict_var.values)
+
+
+def test_naive_forecaster_model_with_regressor_pyfunc_output(
+    naive_forecaster_model_with_regressor, model_path, data_longley
+):
+    """Test naive forecaster prediction of loaded pyfunc model."""
+    _, _, _, X_test = data_longley
+
+    flavor.save_model(sktime_model=naive_forecaster_model_with_regressor, path=model_path)
+    loaded_pyfunc = flavor.pyfunc.load_model(model_uri=model_path)
+
+    X_test_array = convert(X_test, "pd.DataFrame", "np.ndarray")
+
+    model_predict = naive_forecaster_model_with_regressor.predict(fh=FH, X=X_test)
+    predict_conf = pd.DataFrame([{"fh": FH, "predict_method": "predict", "X": X_test_array}])
+    pyfunc_predict = loaded_pyfunc.predict(predict_conf)
+    np.testing.assert_array_equal(model_predict, pyfunc_predict)
+
+    model_predict_interval = naive_forecaster_model_with_regressor.predict_interval(
+        fh=FH, coverage=COVERAGE, X=X_test
+    )
+    predict_interval_conf = pd.DataFrame(
+        [
+            {
+                "fh": FH,
+                "predict_method": "predict_interval",
+                "coverage": COVERAGE,
+                "X": X_test_array,
+            }
+        ]
+    )
+    pyfunc_predict_interval = loaded_pyfunc.predict(predict_interval_conf)
+    np.testing.assert_array_equal(model_predict_interval.values, pyfunc_predict_interval.values)
+
+    model_predict_quantiles = naive_forecaster_model_with_regressor.predict_quantiles(
+        fh=FH, alpha=ALPHA, X=X_test
+    )
+    predict_quantiles_conf = pd.DataFrame(
+        [
+            {
+                "fh": FH,
+                "predict_method": "predict_quantiles",
+                "alpha": ALPHA,
+                "X": X_test_array,
+            }
+        ]
+    )
+    pyfunc_predict_quantiles = loaded_pyfunc.predict(predict_quantiles_conf)
+    np.testing.assert_array_equal(model_predict_quantiles.values, pyfunc_predict_quantiles.values)
+
+    model_predict_var = naive_forecaster_model_with_regressor.predict_var(fh=FH, cov=COV, X=X_test)
+    predict_var_conf = pd.DataFrame(
+        [
+            {
+                "fh": FH,
+                "predict_method": "predict_var",
+                "cov": COV,
+                "X": X_test_array,
+            }
+        ]
+    )
+    pyfunc_predict_var = loaded_pyfunc.predict(predict_var_conf)
+    np.testing.assert_array_equal(model_predict_var.values, pyfunc_predict_var.values)
+
+
+@pytest.mark.parametrize("use_signature", [True, False])
+@pytest.mark.parametrize("use_example", [True, False])
+def test_signature_and_examples_saved_correctly(
+    auto_arima_model, data_airline, model_path, use_signature, use_example
+):
+    """Test saving of mlflow signature and example for native sktime predict method."""
+    # Note: Signature inference fails on native model predict_interval/predict_quantiles
+    prediction = auto_arima_model.predict(fh=FH)
+    signature = infer_signature(data_airline, prediction) if use_signature else None
+    example = pd.DataFrame(data_airline[0:5].copy(deep=False)) if use_example else None
+    flavor.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.parametrize("use_signature", [True, False])
+def test_predict_var_signature_saved_correctly(
+    auto_arima_model, data_airline, model_path, use_signature
+):
+    """Test saving of mlflow signature for native sktime predict_var method."""
+    prediction = auto_arima_model.predict_var(fh=FH)
+    signature = infer_signature(data_airline, prediction) if use_signature else None
+    flavor.save_model(auto_arima_model, path=model_path, signature=signature)
+    mlflow_model = Model.load(model_path)
+    assert signature == mlflow_model.signature
+
+
+@pytest.mark.parametrize("use_signature", [True, False])
+@pytest.mark.parametrize("use_example", [True, False])
+def test_signature_and_example_for_pyfunc_predict_inteval(
+    auto_arima_model, model_path, data_airline, use_signature, use_example
+):
+    """Test saving of mlflow signature and example for pyfunc predict."""
+    model_path_primary = model_path.joinpath("primary")
+    model_path_secondary = model_path.joinpath("secondary")
+    flavor.save_model(sktime_model=auto_arima_model, path=model_path_primary)
+    loaded_pyfunc = flavor.pyfunc.load_model(model_uri=model_path_primary)
+    predict_conf = pd.DataFrame(
+        [
+            {
+                "fh": FH,
+                "predict_method": "predict_interval",
+                "coverage": COVERAGE,
+            }
+        ]
+    )
+    forecast = loaded_pyfunc.predict(predict_conf)
+    signature = infer_signature(data_airline, forecast) if use_signature else None
+    example = pd.DataFrame(data_airline[0:5].copy(deep=False)) if use_example else None
+    flavor.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.parametrize("use_signature", [True, False])
+def test_signature_for_pyfunc_predict_quantiles(
+    auto_arima_model, model_path, data_airline, use_signature
+):
+    """Test saving of mlflow signature for pyfunc sktime predict_quantiles method."""
+    model_path_primary = model_path.joinpath("primary")
+    model_path_secondary = model_path.joinpath("secondary")
+    flavor.save_model(sktime_model=auto_arima_model, path=model_path_primary)
+    loaded_pyfunc = flavor.pyfunc.load_model(model_uri=model_path_primary)
+    predict_conf = pd.DataFrame(
+        [
+            {
+                "fh": FH,
+                "predict_method": "predict_quantiles",
+                "alpha": ALPHA,
+            }
+        ]
+    )
+    forecast = loaded_pyfunc.predict(predict_conf)
+    signature = infer_signature(data_airline, forecast) if use_signature else None
+    flavor.save_model(auto_arima_model, path=model_path_secondary, signature=signature)
+    mlflow_model = Model.load(model_path_secondary)
+    assert signature == mlflow_model.signature
+
+
+def test_load_from_remote_uri_succeeds(auto_arima_model, model_path, mock_s3_bucket):
+    """Test loading native sktime model from mock S3 bucket."""
+    flavor.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 = flavor.load_model(model_uri=model_uri)
+
+    np.testing.assert_array_equal(
+        auto_arima_model.predict(fh=FH),
+        reloaded_sktime_model.predict(fh=FH),
+    )
+
+
+@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."""
+    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 = flavor.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 = flavor.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()
+
+
+def test_log_model_calls_register_model(auto_arima_model, tmp_path):
+    """Test log model calls register model."""
+    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"])
+        flavor.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,
+        )
+
+
+def test_log_model_no_registered_model_name(auto_arima_model, tmp_path):
+    """Test log model calls register model without registered model name."""
+    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"])
+        flavor.log_model(
+            sktime_model=auto_arima_model,
+            artifact_path=artifact_path,
+            conda_env=str(conda_env),
+        )
+        mlflow.register_model.assert_not_called()
+
+
+def test_sktime_pyfunc_raises_invalid_df_input(auto_arima_model, model_path):
+    """Test pyfunc call raises error with invalid dataframe configuration."""
+    flavor.save_model(sktime_model=auto_arima_model, path=model_path)
+    loaded_pyfunc = flavor.pyfunc.load_model(model_uri=model_path)
+
+    with pytest.raises(MlflowException, match="The provided prediction pd.DataFrame "):
+        loaded_pyfunc.predict(pd.DataFrame([{"predict_method": "predict"}, {"fh": FH}]))
+
+    with pytest.raises(MlflowException, match="The provided prediction configuration "):
+        loaded_pyfunc.predict(pd.DataFrame([{"invalid": True}]))
+
+    with pytest.raises(MlflowException, match="Invalid `predict_method` value"):
+        loaded_pyfunc.predict(pd.DataFrame([{"predict_method": "predict_proba"}]))
+
+
+def test_sktime_save_model_raises_invalid_serialization_format(auto_arima_model, model_path):
+    """Test save_model call raises error with invalid serialization format."""
+    with pytest.raises(MlflowException, match="Unrecognized serialization format: "):
+        flavor.save_model(
+            sktime_model=auto_arima_model, path=model_path, serialization_format="json"
+        )
diff --git a/tests/examples/test_examples.py b/tests/examples/test_examples.py
--- a/tests/examples/test_examples.py
+++ b/tests/examples/test_examples.py
@@ -162,6 +162,7 @@ def test_mlflow_run_example(directory, params, tmp_path):
         ("diviner", ["python", "train.py"]),
         ("spark_udf", ["python", "spark_udf_datetime.py"]),
         ("pyfunc", ["python", "train.py"]),
+        ("sktime", ["python", "train.py"]),
     ],
 )
 def test_command_example(directory, command):