Repository: mlflow/mlflow

Issue:
[FR] Autologging functionality for PyTorch/fastai
## Describe the proposal

It'd be nice to add an mlflow.pytorch.autolog() API for automatically logging metrics, params & models generated via PyTorch. In particular, I'd like this to be incorporated with the fastai library (built on PyTorch), which also generates artifacts such as learning rate finder plots, and autolog those as well. 

### Motivation
PyTorch is heavily used in research for Deep Learning and is becoming very prevalent in industry as well with the fastai library. The fastai community has over 450 contributors, and it incorporates practical deep learning best practices + cutting edge algorithms in the library. In the fast.ai course that is run by USF, they use Excel spreadsheets to keep track of their runs. If we could just autolog everything once they import mlflow, it will help bring the adoption of MLflow to the fastai community. 

### Proposed Changes
See above - we propose adding a new mlflow.pytorch.autolog API


Discussion:
@shalomma @LarsHill would either of you be interested in working on this feature? If so I'm happy to write up / help collaborate on an API proposal :)
@smurching sorry, but in the near future I won't have time working on it. I just thought it would be a useful feature.
@smurching are you working on it ? I will be happy to help :)
@pierrelandry I haven't worked on it, that'd be awesome! Before writing code & opening a PR, a good first step would be to write up a description of the proposed API & its behavior. Happy to review a proposal if you're able to write one & share it here.

I'm new to Pytorch, but I'd recommend focusing on adding autologging for the most common Pytorch APIs first (will spend some time tomorrow learning more about Pytorch). Also, I'd suggest `mlflow.pytorch.autolog()` as the API name, to match the existing Keras & TF autologging APIs, but open to alternatives. Thanks for being willing to contribute :)!

@pierrelandry hope you enjoyed the holidays! Thanks again for your interest in helping with this - are you still interested in collaborating?
Hi, I'd love to work on the fastai integration, is there any proposed API?
I thought of adding support for logging metrics during training and saving artifacts generated from
[SaveModelCallback](https://docs.fast.ai/callbacks.html#SaveModelCallback) or [Learner.save](https://docs.fast.ai/basic_train.html#Learner.save) method.
@antoniomdk sorry for the delayed response, that'd be awesome - one strawman API proposal would be to add an `mlflow.pytorch.autolog()` API, and an `mlflow.fastai.autolog()` if appropriate - AFAICT fastai & PyTorch are analogous to Keras and TensorFlow (fastai seems to be a higher-level API implemented using PyTorch), so seems reasonable to have separate APIs for the two libraries. 

That being said, it's probably worth considering whether implementing PyTorch autologging could indirectly solve fastai autologging (i.e. could we reuse the implementation of PyTorch autologging for fastai autologging?). My guess is that it might be possible, but the right way to implement fastai autologging would be instead be to use public `fastai` APIs, e.g. by adding callbacks to save & log the user's model as you described - I don't know too much about either PyTorch or fastai so curious to hear your opinion.


Yes, you're right, Fastai is a higher-level API on top of Pytorch as Keras is on top of Tensorflow. Adding model saving autologging for Pytorch would solve it also for Fastai, as the latter relies on torch.save
https://github.com/fastai/fastai/blob/master/fastai/basic_train.py#L246
So, monkey patching [torch.save](https://pytorch.org/docs/stable/torch.html?highlight=save#torch.save) would work for both libraries. It'd also work for model checkpoints along epochs. 

On the other hand, I think we should go for adding a callback to handle metric logging on Fastai.
There is also a proposal from the Fastai community to add a callback for this: https://forums.fast.ai/t/mlflow-tracking-support/38893

For Pytorch I'm not sure how we should proceed because the training cycle is done manually and the evaluation depends on the user implementing it.
Adding autologging to PyTorch is much harder inherently because PyTorch's training cycle can be very flexible and can take many forms, like @antoniomdk mentioned. I took a look at what MLFlow autologging does for other frameworks (https://www.mlflow.org/docs/latest/tracking.html#automatic-logging-from-tensorflow-and-keras-experimental), and I think we might be able to do the following:

- Metrics

| Logging Item | Possible Implementation |
| ------------- | ------------- |
| training loss  | Monkey-patch [Tensor.backward](https://pytorch.org/docs/stable/autograd.html#torch.Tensor.backward) method and log the tensor |
| validation loss  | (This is difficult to track because we can't know programmatically which variable is the validation loss. We might have to resort to manual logging.) |
| user-specified metrics | (Not sure how we should map this concept from Keras to PyTorch.) |

- Parameters

| Logging Item | Possible Implementation |
| ------------- | ------------- |
| fit() parameters  | (Not sure how we should map this concept from Keras to PyTorch.) |
| optimizer name  | Find instantiations of [torch.optim.Optimizer](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer) subclasses in the code, and log the optimizer name |
| learning rate | Log the `lr` value of the optimizer found above |
| epsilon | Log the `eps` value of the optimizer found above |

- Artifacts

| Logging Item | Possible Implementation |
| ------------- | ------------- |
| model summary on training start | Monkey-patch `Tensor.backward` or `Optimizer.step` method and log the model summary on the first method call |
| MLFlow model (PyTorch model) on training end | The end of training is hard to detect, but maybe we can monkey-patch `Optimizer.step` or `torch.nn.Module.eval` method to trigger MLFlow model saving after every epoch |

> AFAICT fastai & PyTorch are analogous to Keras and TensorFlow (fastai seems to be a higher-level API implemented using PyTorch), so seems reasonable to have separate APIs for the two libraries.

+1 I think it's easier to add autologging to fastai as a separate API because callbacks are a cleaner way to implement autologging.

I think it'd be hard to detect the last epoch using Optimizer.step or torch.nn.Module.eval.
For model persistence it's easier to rely on [torch.save](https://pytorch.org/docs/stable/torch.html?highlight=save#torch.save).

For optimizer tracking we can monkey-patch torch.optim.Optimizer class to keep track of all instances created. We can use those instances to log optimizer parameters, as @jacobhess118 suggested.

For tracking loss I'm not sure if using Tensor.backward would be a good idea (could it make Tensor.backward subcalls?).
@jacobhess118 @antoniomdk thanks for the detailed investigation! Here are my impressions:
* Makes sense to implement fastai autologging via a separate `mlflow.fastai.autolog` API, where calling `mlflow.fastai.autolog` patches [`Learner.fit`](https://docs.fast.ai/basic_train.html#Learner.fit) & modifies the `callbacks` argument, injecting a callback that logs model & training information (see [analogous implementation](https://github.com/mlflow/mlflow/blob/1a762d56564f5f708bd7f9cbc5a48530fb354ae3/mlflow/keras.py#L528) in Keras autologging)
* PyTorch autologging: PyTorch is a popular framework (2x the GitHub stars of fastai) so would be nice to have autologging support, but if it seems likely to be brittle/complex (as it currently seems from @jacobhess118's description), we can punt on it for now - it's possible it'll become easier down the line

Thus @antoniomdk if you're interested in adding `fastai` autologging for now, would be happy to review a PR that does so :)

 > For tracking loss I'm not sure if using Tensor.backward would be a good idea (could it make Tensor.backward subcalls?).

Tensor.backward doesn’t make subcalls of itself at the moment (and likely won’t be in the future) because calling Python functions is slow. However I do think that it smells more like a hack, and we probably shouldn’t do it this way unless absolutely necessary.
> Thus @antoniomdk if you're interested in adding fastai autologging for now, would be happy to review a PR that does so :)

@smurching I'll work on that as soon as I've reviewed the contributing guideline and the Keras implementation. Thanks!
Hi @antoniomdk, just wanted to check in (no rush of course!) - did you get a chance to get started with this and/or have any questions?
Hi @smurching, I started to work on this recently. I've forked the repo and create a new branch but my advances are slow. I'm quite busy at this moment at work, but I hope to make a PR by the end of this month.
Awesome - totally makes sense & sounds good! Let us know if anything comes up along the way :)
Sorry for the late response, but I just want you to know that I'm working on the PR. I've already got cover most of the functionality of autologging and model exporting for Fastai.

I'm figuring out until which extend should we support autologging for specific callbacks.
There is a callback for saving the best model of the training, Early Stopping, etc.
So, I'm not sure how to proceed with those callbacks. Should we support just Early Stopping?
Or should we add [SaveModelCallback](SaveModelCallback) too?

I've added a custom callback to save the model params before training and the model itself after training. We can use that callback to track the best model along epochs.

What do you think guys?
I've created a PR for this @smurching . It has support for logging Early Stopping and OneCycleScheduler params (along with the other .fit params).

Resolving patch:
diff --git a/docs/source/conf.py b/docs/source/conf.py
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -250,6 +250,10 @@
      'Databricks', 'manual'),
 ]
 
+# Mock torch & fastai imports as per suggestion in
+# https://github.com/sphinx-doc/sphinx/issues/6521#issuecomment-505765893
+autodoc_mock_imports = ["torch", "fastai"]
+
 # The name of an image file (relative to this directory) to place at the top of
 # the title page.
 #latex_logo = None
diff --git a/mlflow/fastai.py b/mlflow/fastai.py
new file mode 100644
--- /dev/null
+++ b/mlflow/fastai.py
@@ -0,0 +1,194 @@
+"""
+The ``mlflow.fastai`` module provides an API for logging and loading fast.ai models. This module
+exports fast.ai models with the following flavors:
+
+fastai (native) format
+    This is the main flavor that can be loaded back into fastai.
+:py:mod:`mlflow.pyfunc`
+    Produced for use by generic pyfunc-based deployment tools and batch inference.
+
+.. _fastai.Learner:
+    https://docs.fast.ai/basic_train.html#Learner
+.. _fastai.Learner.export:
+    https://docs.fast.ai/basic_train.html#Learner.export
+"""
+
+from __future__ import absolute_import
+
+import os
+import yaml
+import pandas as pd
+import numpy as np
+
+from mlflow import pyfunc
+from mlflow.models import Model
+import mlflow.tracking
+from mlflow.exceptions import MlflowException
+from mlflow.tracking.artifact_utils import _download_artifact_from_uri
+from mlflow.utils.environment import _mlflow_conda_env
+from mlflow.utils.model_utils import _get_flavor_configuration
+
+from fastai.tabular import TabularList
+from fastai.basic_data import DatasetType
+
+
+FLAVOR_NAME = "fastai"
+
+
+def get_default_conda_env(include_cloudpickle=False):
+    """
+    :return: The default Conda environment for MLflow Models produced by calls to
+             :func:`save_model()` and :func:`log_model()`.
+    """
+    import fastai
+    pip_deps = None
+    if include_cloudpickle:
+        import cloudpickle
+        pip_deps = ["cloudpickle=={}".format(cloudpickle.__version__)]
+    return _mlflow_conda_env(
+        additional_conda_deps=[
+            "fastai={}".format(fastai.__version__),
+        ],
+        additional_pip_deps=pip_deps,
+        additional_conda_channels=None
+    )
+
+
+def save_model(fastai_learner, path, conda_env=None, mlflow_model=Model(), **kwargs):
+    """
+    Save a fastai Learner to a path on the local file system.
+
+    :param fastai_learner: fastai Learner to be saved.
+    :param path: Local path where the model is to be saved.
+    :param conda_env: Either a dictionary representation of a Conda environment or the path to a
+                      Conda environment yaml file. If provided, this describes the environment
+                      this model should be run in. At minimum, it should specify the
+                      dependencies contained in :func:`get_default_conda_env()`. If
+                      ``None``, the default :func:`get_default_conda_env()` environment is
+                      added to the model. The following is an *example* dictionary
+                      representation of a Conda environment::
+
+                        {
+                            'name': 'mlflow-env',
+                            'channels': ['defaults'],
+                            'dependencies': [
+                                'python=3.7.0',
+                                'fastai=1.0.60',
+                            ]
+                        }
+    :param mlflow_model: MLflow model config this flavor is being added to.
+    :param kwargs: kwargs to pass to ``Learner.save`` method.
+    """
+    import fastai
+    from pathlib import Path
+
+    path = os.path.abspath(path)
+    if os.path.exists(path):
+        raise MlflowException("Path '{}' already exists".format(path))
+    model_data_subpath = "model.fastai"
+    model_data_path = os.path.join(path, model_data_subpath)
+    model_data_path = Path(model_data_path)
+    os.makedirs(path)
+
+    # Save an Learner
+    fastai_learner.export(model_data_path, **kwargs)
+
+    conda_env_subpath = "conda.yaml"
+
+    if conda_env is None:
+        conda_env = get_default_conda_env()
+    elif not isinstance(conda_env, dict):
+        with open(conda_env, "r") as f:
+            conda_env = yaml.safe_load(f)
+    with open(os.path.join(path, conda_env_subpath), "w") as f:
+        yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
+
+    pyfunc.add_to_model(mlflow_model, loader_module="mlflow.fastai",
+                        data=model_data_subpath, env=conda_env_subpath)
+    mlflow_model.add_flavor(FLAVOR_NAME, fastai_version=fastai.__version__, data=model_data_subpath)
+    mlflow_model.save(os.path.join(path, "MLmodel"))
+
+
+def log_model(fastai_learner, artifact_path, conda_env=None, registered_model_name=None, **kwargs):
+    """
+    Log a fastai model as an MLflow artifact for the current run.
+
+    :param fastai_learner: Fastai model (an instance of `fastai.Learner`_) to be saved.
+    :param artifact_path: Run-relative artifact path.
+    :param conda_env: Either a dictionary representation of a Conda environment or the path to a
+                      Conda environment yaml file. If provided, this describes the environment
+                      this model should be run in. At minimum, it should specify the dependencies
+                      contained in :func:`get_default_conda_env()`. If ``None``, the default
+                      :func:`get_default_conda_env()` environment is added to the model.
+                      The following is an *example* dictionary representation of a Conda
+                      environment::
+
+                        {
+                            'name': 'mlflow-env',
+                            'channels': ['defaults'],
+                            'dependencies': [
+                                'python=3.7.0',
+                                'fastai=1.0.60',
+                            ]
+                        }
+    :param registered_model_name: Note:: Experimental: This argument may change or be removed in a
+                                  future release without warning. If given, create a model
+                                  version under ``registered_model_name``, also creating a
+                                  registered model if one with the given name does not exist.
+    :param kwargs: kwargs to pass to `fastai.Learner.export`_ method.
+    """
+    Model.log(artifact_path=artifact_path, flavor=mlflow.fastai,
+              registered_model_name=registered_model_name,
+              fastai_learner=fastai_learner, conda_env=conda_env, **kwargs)
+
+
+def _load_model(path):
+    from fastai.basic_train import load_learner
+    abspath = os.path.abspath(path)
+    path, file = os.path.split(abspath)
+    return load_learner(path, file)
+
+
+class _FastaiModelWrapper:
+    def __init__(self, learner):
+        self.learner = learner
+
+    def predict(self, dataframe):
+        test_data = TabularList.from_df(dataframe, cont_names=self.learner.data.cont_names)
+        self.learner.data.add_test(test_data)
+        preds, target = self.learner.get_preds(DatasetType.Test)
+        preds = pd.Series(map(np.array, preds.numpy()), name='predictions')
+        target = pd.Series(target.numpy(), name='target')
+        return pd.concat([preds, target], axis='columns')
+
+
+def _load_pyfunc(path):
+    """
+    Load PyFunc implementation. Called by ``pyfunc.load_pyfunc``.
+
+    :param path: Local filesystem path to the MLflow Model with the ``fastai`` flavor.
+    """
+    return _FastaiModelWrapper(_load_model(path))
+
+
+def load_model(model_uri):
+    """
+    Load a fastai model from a local file or a run.
+
+    :param model_uri: 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``
+
+                      For more information about supported URI schemes, see
+                      `Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html#
+                      artifact-locations>`_.
+
+    :return: A fastai model (an instance of `fastai.Learner`_).
+    """
+    local_model_path = _download_artifact_from_uri(artifact_uri=model_uri)
+    flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME)
+    model_file_path = os.path.join(local_model_path, flavor_conf.get("data", "model.fastai"))
+    return _load_model(path=model_file_path)
diff --git a/mlflow/utils/autologging_utils.py b/mlflow/utils/autologging_utils.py
--- a/mlflow/utils/autologging_utils.py
+++ b/mlflow/utils/autologging_utils.py
@@ -57,7 +57,12 @@ def log_fn_args_as_params(fn, args, kwargs, unlogged=[]):  # pylint: disable=W01
     """
     # all_default_values has length n, corresponding to values of the
     # last n elements in all_param_names
-    all_param_names, _, _, all_default_values = inspect.getargspec(fn)  # pylint: disable=W1505
+    pos_params, _, _, pos_defaults, kw_params, kw_defaults, _ = inspect.getfullargspec(fn)
+
+    kw_params = list(kw_params) if kw_params else []
+    pos_defaults = list(pos_defaults) if pos_defaults else []
+    all_param_names = pos_params + kw_params
+    all_default_values = pos_defaults + [kw_defaults[param] for param in kw_params]
 
     # Checking if default values are present for logging. Known bug that getargspec will return an
     # empty argspec for certain functions, despite the functions having an argspec.
@@ -71,9 +76,10 @@ def log_fn_args_as_params(fn, args, kwargs, unlogged=[]):  # pylint: disable=W01
 
     # Logging the arguments passed by the user
     args_dict = dict((param_name, param_val) for param_name, param_val
-                     in zip(all_param_names[:len(args)], args)
+                     in zip(all_param_names, args)
                      if param_name not in unlogged)
-    if len(args_dict.keys()) > 0:
+
+    if args_dict:
         try_mlflow_log(mlflow.log_params, args_dict)
 
     # Logging the kwargs passed by the user


Test patch:
diff --git a/tests/fastai/test_fastai_model_export.py b/tests/fastai/test_fastai_model_export.py
new file mode 100644
--- /dev/null
+++ b/tests/fastai/test_fastai_model_export.py
@@ -0,0 +1,283 @@
+from __future__ import print_function
+
+import mock
+import os
+import pytest
+import yaml
+from collections import namedtuple
+
+import numpy as np
+import pandas as pd
+import sklearn.datasets as datasets
+from fastai.tabular import tabular_learner, TabularList
+from fastai.metrics import accuracy
+
+import mlflow.fastai
+import mlflow.utils
+from mlflow import pyfunc
+from mlflow.models import Model
+from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
+from mlflow.tracking.artifact_utils import _download_artifact_from_uri
+from mlflow.utils.environment import _mlflow_conda_env
+from mlflow.utils.file_utils import TempDir
+from mlflow.utils.model_utils import _get_flavor_configuration
+from fastai.tabular import DatasetType
+
+from tests.helper_functions import set_boto_credentials  # pylint: disable=unused-import
+from tests.helper_functions import mock_s3_bucket  # pylint: disable=unused-import
+from tests.helper_functions import score_model_in_sagemaker_docker_container
+
+ModelWithData = namedtuple("ModelWithData", ["model", "inference_dataframe"])
+
+
+@pytest.fixture(scope="session")
+def fastai_model():
+    iris = datasets.load_iris()
+    X = pd.DataFrame(iris.data[:, :2], columns=iris.feature_names[:2])
+    y = pd.Series(iris.target, name='label')
+    data = (TabularList.from_df(pd.concat([X, y], axis=1), cont_names=list(X.columns))
+            .split_by_rand_pct(valid_pct=0.1, seed=42)
+            .label_from_df(cols='label')
+            .databunch())
+    model = tabular_learner(data, metrics=accuracy, layers=[3])
+    model.fit(1)
+    return ModelWithData(model=model, inference_dataframe=X)
+
+
+@pytest.fixture
+def model_path(tmpdir):
+    return os.path.join(str(tmpdir), "model")
+
+
+@pytest.fixture
+def fastai_custom_env(tmpdir):
+    conda_env = os.path.join(str(tmpdir), "conda_env.yml")
+    _mlflow_conda_env(
+        conda_env,
+        additional_pip_deps=["fastai", "pytest"])
+    return conda_env
+
+
+def compare_wrapper_results(wrapper1_results, wrapper2_results):
+    samples = wrapper1_results['predictions'].shape[0]
+    predictions1 = np.concatenate(wrapper1_results['predictions'], axis=0).reshape((samples, -1))
+    predictions2 = np.concatenate(wrapper2_results['predictions'], axis=0).reshape((samples, -1))
+    np.testing.assert_array_almost_equal(wrapper1_results['target'], wrapper2_results['target'])
+    np.testing.assert_array_almost_equal(predictions1, predictions2)
+
+
+@pytest.mark.large
+def test_model_save_load(fastai_model, model_path):
+    model = fastai_model.model
+
+    mlflow.fastai.save_model(fastai_learner=model, path=model_path)
+    reloaded_model = mlflow.fastai.load_model(model_uri=model_path)
+    reloaded_pyfunc = pyfunc.load_model(model_uri=model_path)
+
+    # Verify reloaded model computes same predictions as original model
+    test_data = TabularList.from_df(fastai_model.inference_dataframe)
+    model.data.add_test(test_data)
+    reloaded_model.data.add_test(test_data)
+
+    real_preds, real_target = map(lambda output: output.numpy(), model.get_preds(DatasetType.Test))
+    reloaded_preds, reloaded_target = map(lambda output: output.numpy(),
+                                          reloaded_model.get_preds(DatasetType.Test))
+
+    np.testing.assert_array_almost_equal(real_preds, reloaded_preds)
+    np.testing.assert_array_almost_equal(real_target, reloaded_target)
+
+    model_wrapper = mlflow.fastai._FastaiModelWrapper(model)
+    reloaded_model_wrapper = mlflow.fastai._FastaiModelWrapper(reloaded_model)
+
+    model_result = model_wrapper.predict(fastai_model.inference_dataframe)
+    reloaded_result = reloaded_model_wrapper.predict(fastai_model.inference_dataframe)
+    pyfunc_result = reloaded_pyfunc.predict(fastai_model.inference_dataframe)
+
+    compare_wrapper_results(model_result, reloaded_result)
+    compare_wrapper_results(reloaded_result, pyfunc_result)
+
+
+@pytest.mark.large
+def test_model_load_from_remote_uri_succeeds(fastai_model, model_path, mock_s3_bucket):
+    model = fastai_model.model
+
+    mlflow.fastai.save_model(fastai_learner=fastai_model.model, path=model_path)
+    artifact_root = "s3://{bucket_name}".format(bucket_name=mock_s3_bucket)
+    artifact_path = "model"
+    artifact_repo = S3ArtifactRepository(artifact_root)
+    artifact_repo.log_artifacts(model_path, artifact_path=artifact_path)
+
+    model_uri = artifact_root + "/" + artifact_path
+    reloaded_model = mlflow.fastai.load_model(model_uri=model_uri)
+
+    model_wrapper = mlflow.fastai._FastaiModelWrapper(model)
+    reloaded_model_wrapper = mlflow.fastai._FastaiModelWrapper(reloaded_model)
+
+    compare_wrapper_results(
+            model_wrapper.predict(fastai_model.inference_dataframe),
+            reloaded_model_wrapper.predict(fastai_model.inference_dataframe))
+
+
+@pytest.mark.large
+def test_model_log(fastai_model, model_path):
+    old_uri = mlflow.get_tracking_uri()
+    model = fastai_model.model
+    with TempDir(chdr=True, remove_on_exit=True) as tmp:
+        for should_start_run in [False, True]:
+            try:
+                mlflow.set_tracking_uri("test")
+                if should_start_run:
+                    mlflow.start_run()
+
+                artifact_path = "model"
+                conda_env = os.path.join(tmp.path(), "conda_env.yaml")
+                _mlflow_conda_env(conda_env, additional_pip_deps=["fastai"])
+
+                mlflow.fastai.log_model(
+                    fastai_learner=model,
+                    artifact_path=artifact_path,
+                    conda_env=conda_env)
+
+                model_uri = "runs:/{run_id}/{artifact_path}".format(
+                    run_id=mlflow.active_run().info.run_id,
+                    artifact_path=artifact_path)
+
+                reloaded_model = mlflow.fastai.load_model(model_uri=model_uri)
+
+                model_wrapper = mlflow.fastai._FastaiModelWrapper(model)
+                reloaded_model_wrapper = mlflow.fastai._FastaiModelWrapper(reloaded_model)
+
+                compare_wrapper_results(
+                    model_wrapper.predict(fastai_model.inference_dataframe),
+                    reloaded_model_wrapper.predict(fastai_model.inference_dataframe))
+
+                model_path = _download_artifact_from_uri(artifact_uri=model_uri)
+                model_config = Model.load(os.path.join(model_path, "MLmodel"))
+                assert pyfunc.FLAVOR_NAME in model_config.flavors
+                assert pyfunc.ENV in model_config.flavors[pyfunc.FLAVOR_NAME]
+                env_path = model_config.flavors[pyfunc.FLAVOR_NAME][pyfunc.ENV]
+                assert os.path.exists(os.path.join(model_path, env_path))
+
+            finally:
+                mlflow.end_run()
+                mlflow.set_tracking_uri(old_uri)
+
+
+def test_log_model_calls_register_model(fastai_model):
+    artifact_path = "model"
+    register_model_patch = mock.patch("mlflow.register_model")
+    with mlflow.start_run(), register_model_patch, TempDir(chdr=True, remove_on_exit=True) as tmp:
+        conda_env = os.path.join(tmp.path(), "conda_env.yaml")
+        _mlflow_conda_env(conda_env, additional_pip_deps=["fastai"])
+        mlflow.fastai.log_model(fastai_learner=fastai_model.model,
+                                artifact_path=artifact_path,
+                                conda_env=conda_env,
+                                registered_model_name="AdsModel1")
+        model_uri = "runs:/{run_id}/{artifact_path}".format(run_id=mlflow.active_run().info.run_id,
+                                                            artifact_path=artifact_path)
+        mlflow.register_model.assert_called_once_with(model_uri, "AdsModel1")
+
+
+def test_log_model_no_registered_model_name(fastai_model):
+    artifact_path = "model"
+    register_model_patch = mock.patch("mlflow.register_model")
+    with mlflow.start_run(), register_model_patch, TempDir(chdr=True, remove_on_exit=True) as tmp:
+        conda_env = os.path.join(tmp.path(), "conda_env.yaml")
+        _mlflow_conda_env(conda_env, additional_pip_deps=["fastai"])
+        mlflow.fastai.log_model(fastai_learner=fastai_model.model,
+                                artifact_path=artifact_path,
+                                conda_env=conda_env)
+        mlflow.register_model.assert_not_called()
+
+
+@pytest.mark.large
+def test_model_save_persists_specified_conda_env_in_mlflow_model_directory(
+        fastai_model, model_path, fastai_custom_env):
+    mlflow.fastai.save_model(
+        fastai_learner=fastai_model.model, path=model_path, conda_env=fastai_custom_env)
+
+    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
+    saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
+    assert os.path.exists(saved_conda_env_path)
+    assert saved_conda_env_path != fastai_custom_env
+
+    with open(fastai_custom_env, "r") as f:
+        fastai_custom_env_parsed = yaml.safe_load(f)
+    with open(saved_conda_env_path, "r") as f:
+        saved_conda_env_parsed = yaml.safe_load(f)
+    assert saved_conda_env_parsed == fastai_custom_env_parsed
+
+
+@pytest.mark.large
+def test_model_save_accepts_conda_env_as_dict(fastai_model, model_path):
+    conda_env = dict(mlflow.fastai.get_default_conda_env())
+    conda_env["dependencies"].append("pytest")
+    mlflow.fastai.save_model(fastai_learner=fastai_model.model,
+                             path=model_path, conda_env=conda_env)
+
+    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
+    saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
+    assert os.path.exists(saved_conda_env_path)
+
+    with open(saved_conda_env_path, "r") as f:
+        saved_conda_env_parsed = yaml.safe_load(f)
+    assert saved_conda_env_parsed == conda_env
+
+
+@pytest.mark.large
+def test_model_log_persists_specified_conda_env_in_mlflow_model_directory(
+        fastai_model, fastai_custom_env):
+    artifact_path = "model"
+    with mlflow.start_run():
+        mlflow.fastai.log_model(fastai_learner=fastai_model.model,
+                                artifact_path=artifact_path,
+                                conda_env=fastai_custom_env)
+        model_uri = "runs:/{run_id}/{artifact_path}".format(
+            run_id=mlflow.active_run().info.run_id,
+            artifact_path=artifact_path)
+
+    model_path = _download_artifact_from_uri(artifact_uri=model_uri)
+    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
+    saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
+    assert os.path.exists(saved_conda_env_path)
+    assert saved_conda_env_path != fastai_custom_env
+
+    with open(fastai_custom_env, "r") as f:
+        fastai_custom_env_parsed = yaml.safe_load(f)
+    with open(saved_conda_env_path, "r") as f:
+        saved_conda_env_parsed = yaml.safe_load(f)
+    assert saved_conda_env_parsed == fastai_custom_env_parsed
+
+
+@pytest.mark.large
+def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
+        fastai_model, model_path):
+    mlflow.fastai.save_model(fastai_learner=fastai_model.model,
+                             path=model_path, conda_env=None)
+
+    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
+    conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
+    with open(conda_env_path, "r") as f:
+        conda_env = yaml.safe_load(f)
+
+    assert conda_env == mlflow.fastai.get_default_conda_env()
+
+
+@pytest.mark.large
+def test_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies(
+        fastai_model):
+    artifact_path = "model"
+    with mlflow.start_run():
+        mlflow.fastai.log_model(fastai_learner=fastai_model.model,
+                                artifact_path=artifact_path, conda_env=None)
+        model_uri = "runs:/{run_id}/{artifact_path}".format(
+            run_id=mlflow.active_run().info.run_id,
+            artifact_path=artifact_path)
+
+    model_path = _download_artifact_from_uri(artifact_uri=model_uri)
+    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
+    conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
+    with open(conda_env_path, "r") as f:
+        conda_env = yaml.safe_load(f)
+
+    assert conda_env == mlflow.fastai.get_default_conda_env()