Repository: mlflow/mlflow

Issue:
[BUG] mlflow.pyfunc does not support dict conda_env
Thank you for submitting an issue. Please refer to our [issue policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)
for information on what types of issues we address. For help with debugging your code, please refer to [Stack Overflow](https://stackoverflow.com/questions/tagged/mlflow).

Please fill in this template and do not delete it unless you are sure your issue is outside its scope.

### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**: No
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux Ubuntu 18.04.3 LTS
- **MLflow installed from (source or binary)**: Binary
- **MLflow version (run ``mlflow --version``)**: 1.5.0 and 1.7.0
- **Python version**: 3.6.10
- **npm version, if running the dev UI**: N/A
- **Exact command to reproduce**: `python -c "import mlflow.pyfunc; mlflow.pyfunc.log_model('x', 'x', conda_env={})"`

### Describe the problem
[The documentation](https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#mlflow.pyfunc.log_model) describes `conda_env` of `mlflow.pyfunc.log_model()` as:
> Either a dictionary representation of a Conda environment or the path to a Conda environment yaml file.

But `mlflow.pyfunc.log_model()` denies `conda_env` as a `dict`.

### Code to reproduce issue
```python
>>> import mlflow.pyfunc
>>> mlflow.pyfunc.log_model('x', 'x', conda_env={})
Traceback (most recent call last):
  ...
TypeError: stat: path should be string, bytes, os.PathLike or integer, not dict
```

### Other info / logs
N/A


Discussion:
@sublee
Can you paste the full stack trace?
@harupy Sure, here you are:

```python
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/opt/conda/lib/python3.6/site-packages/mlflow/pyfunc/__init__.py", line 691, in log_model
    registered_model_name=registered_model_name)
  File "/opt/conda/lib/python3.6/site-packages/mlflow/models/__init__.py", line 101, in log
    flavor.save_model(path=local_path, mlflow_model=mlflow_model, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/mlflow/pyfunc/__init__.py", line 597, in save_model
    code_paths=code_path, conda_env=conda_env, mlflow_model=mlflow_model)
  File "/opt/conda/lib/python3.6/site-packages/mlflow/pyfunc/__init__.py", line 732, in _save_model_with_loader_module_and_data_path
    shutil.copy(src=conda_env, dst=os.path.join(path, "mlflow_env.yml"))
  File "/opt/conda/lib/python3.6/shutil.py", line 245, in copy
    copyfile(src, dst, follow_symlinks=follow_symlinks)
  File "/opt/conda/lib/python3.6/shutil.py", line 103, in copyfile
    if _samefile(src, dst):
  File "/opt/conda/lib/python3.6/shutil.py", line 88, in _samefile
    return os.path.samefile(src, dst)
  File "/opt/conda/lib/python3.6/genericpath.py", line 96, in samefile
    s1 = os.stat(f1)
TypeError: stat: path should be string, bytes, os.PathLike or integer, not dict
```
This seems to be a bug. When `conda_env` is a dict, the code below does not work.

https://github.com/mlflow/mlflow/blob/27b57e8f655860846e451bf5b024d0ba46818dbf/mlflow/pyfunc/__init__.py#L730-L734

Resolving patch:
diff --git a/mlflow/pyfunc/__init__.py b/mlflow/pyfunc/__init__.py
--- a/mlflow/pyfunc/__init__.py
+++ b/mlflow/pyfunc/__init__.py
@@ -194,6 +194,7 @@
 import os
 import pandas
 import shutil
+import yaml
 from copy import deepcopy
 
 import mlflow
@@ -717,7 +718,6 @@ def _save_model_with_loader_module_and_data_path(path, loader_module, data_path=
 
     code = None
     data = None
-    env = None
 
     if data_path is not None:
         model_file = _copy_file_or_tree(src=data_path, dst=path, dst_dir="data")
@@ -728,12 +728,17 @@ def _save_model_with_loader_module_and_data_path(path, loader_module, data_path=
             _copy_file_or_tree(src=code_path, dst=path, dst_dir="code")
         code = "code"
 
-    if conda_env is not None:
-        shutil.copy(src=conda_env, dst=os.path.join(path, "mlflow_env.yml"))
-        env = "mlflow_env.yml"
+    conda_env_subpath = "mlflow_env.yml"
+    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)
 
     mlflow.pyfunc.add_to_model(
-        mlflow_model, loader_module=loader_module, code=code, data=data, env=env)
+        mlflow_model, loader_module=loader_module, code=code, data=data, env=conda_env_subpath)
     mlflow_model.save(os.path.join(path, 'MLmodel'))
     return mlflow_model
 


Test patch:
diff --git a/tests/pyfunc/test_model_export_with_loader_module_and_data_path.py b/tests/pyfunc/test_model_export_with_loader_module_and_data_path.py
--- a/tests/pyfunc/test_model_export_with_loader_module_and_data_path.py
+++ b/tests/pyfunc/test_model_export_with_loader_module_and_data_path.py
@@ -2,6 +2,7 @@
 
 import os
 import pickle
+import yaml
 
 import numpy as np
 import pytest
@@ -17,6 +18,8 @@
 from mlflow.exceptions import MlflowException
 from mlflow.models import Model
 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
 
 
 def _load_pyfunc(path):
@@ -27,6 +30,23 @@ def _load_pyfunc(path):
             return pickle.load(f, encoding='latin1')  # pylint: disable=unexpected-keyword-arg
 
 
+@pytest.fixture
+def pyfunc_custom_env_file(tmpdir):
+    conda_env = os.path.join(str(tmpdir), "conda_env.yml")
+    _mlflow_conda_env(
+        conda_env,
+        additional_conda_deps=["scikit-learn", "pytest", "cloudpickle"],
+        additional_pip_deps=["-e " + os.path.dirname(mlflow.__path__[0])])
+    return conda_env
+
+
+@pytest.fixture
+def pyfunc_custom_env_dict():
+    return _mlflow_conda_env(
+        additional_conda_deps=["scikit-learn", "pytest", "cloudpickle"],
+        additional_pip_deps=["-e " + os.path.dirname(mlflow.__path__[0])])
+
+
 @pytest.fixture(scope="module")
 def iris_data():
     iris = sklearn.datasets.load_iris()
@@ -131,3 +151,91 @@ def test_log_model_with_unsupported_argument_combinations_throws_exception():
         mlflow.pyfunc.log_model(artifact_path="pyfunc_model",
                                 data_path="/path/to/data")
     assert "Either `loader_module` or `python_model` must be specified" in str(exc_info)
+
+
+@pytest.mark.large
+def test_log_model_persists_specified_conda_env_file_in_mlflow_model_directory(
+        sklearn_knn_model, tmpdir, pyfunc_custom_env_file):
+    sk_model_path = os.path.join(str(tmpdir), "knn.pkl")
+    with open(sk_model_path, "wb") as f:
+        pickle.dump(sklearn_knn_model, f)
+
+    pyfunc_artifact_path = "pyfunc_model"
+    with mlflow.start_run():
+        mlflow.pyfunc.log_model(artifact_path=pyfunc_artifact_path,
+                                data_path=sk_model_path,
+                                loader_module=os.path.basename(__file__)[:-3],
+                                code_path=[__file__],
+                                conda_env=pyfunc_custom_env_file)
+        run_id = mlflow.active_run().info.run_id
+
+    pyfunc_model_path = _download_artifact_from_uri("runs:/{run_id}/{artifact_path}".format(
+        run_id=run_id, artifact_path=pyfunc_artifact_path))
+
+    pyfunc_conf = _get_flavor_configuration(
+        model_path=pyfunc_model_path, flavor_name=mlflow.pyfunc.FLAVOR_NAME)
+    saved_conda_env_path = os.path.join(pyfunc_model_path, pyfunc_conf[mlflow.pyfunc.ENV])
+    assert os.path.exists(saved_conda_env_path)
+    assert saved_conda_env_path != pyfunc_custom_env_file
+
+    with open(pyfunc_custom_env_file, "r") as f:
+        pyfunc_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 == pyfunc_custom_env_parsed
+
+
+@pytest.mark.large
+def test_log_model_persists_specified_conda_env_dict_in_mlflow_model_directory(
+        sklearn_knn_model, tmpdir, pyfunc_custom_env_dict):
+    sk_model_path = os.path.join(str(tmpdir), "knn.pkl")
+    with open(sk_model_path, "wb") as f:
+        pickle.dump(sklearn_knn_model, f)
+
+    pyfunc_artifact_path = "pyfunc_model"
+    with mlflow.start_run():
+        mlflow.pyfunc.log_model(artifact_path=pyfunc_artifact_path,
+                                data_path=sk_model_path,
+                                loader_module=os.path.basename(__file__)[:-3],
+                                code_path=[__file__],
+                                conda_env=pyfunc_custom_env_dict)
+        run_id = mlflow.active_run().info.run_id
+
+    pyfunc_model_path = _download_artifact_from_uri("runs:/{run_id}/{artifact_path}".format(
+        run_id=run_id, artifact_path=pyfunc_artifact_path))
+
+    pyfunc_conf = _get_flavor_configuration(
+        model_path=pyfunc_model_path, flavor_name=mlflow.pyfunc.FLAVOR_NAME)
+    saved_conda_env_path = os.path.join(pyfunc_model_path, pyfunc_conf[mlflow.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 == pyfunc_custom_env_dict
+
+
+@pytest.mark.large
+def test_log_model_without_specified_conda_env_uses_default_env_with_expected_dependencies(
+        sklearn_knn_model, tmpdir):
+    sk_model_path = os.path.join(str(tmpdir), "knn.pkl")
+    with open(sk_model_path, "wb") as f:
+        pickle.dump(sklearn_knn_model, f)
+
+    pyfunc_artifact_path = "pyfunc_model"
+    with mlflow.start_run():
+        mlflow.pyfunc.log_model(artifact_path=pyfunc_artifact_path,
+                                data_path=sk_model_path,
+                                loader_module=os.path.basename(__file__)[:-3],
+                                code_path=[__file__])
+        run_id = mlflow.active_run().info.run_id
+
+    pyfunc_model_path = _download_artifact_from_uri("runs:/{run_id}/{artifact_path}".format(
+        run_id=run_id, artifact_path=pyfunc_artifact_path))
+
+    pyfunc_conf = _get_flavor_configuration(
+        model_path=pyfunc_model_path, flavor_name=mlflow.pyfunc.FLAVOR_NAME)
+    conda_env_path = os.path.join(pyfunc_model_path, pyfunc_conf[mlflow.pyfunc.ENV])
+    with open(conda_env_path, "r") as f:
+        conda_env = yaml.safe_load(f)
+
+    assert conda_env == mlflow.pyfunc.model.get_default_conda_env()