Repository: mlflow/mlflow

Issue:
Avoid extensive copying of model files & artifacts in pyfunc
### Issue
I was trying out mlflow to serve a ML model with large model files (totalling about 2 GB) on sagemaker. I used `mlflow.pyfunc` to load the model data etc. specified via the artifacts-dictionary. 

However, when I built the Docker container and ran it locally, gunicorn always timed out (due to the hard-coded 60s timeout). 
Observing the `/tmp` folder in the Docker container during the run as well as digging into the mlflow code, I noticed that temporary folders containing the **full** model artifacts (in my case amounting to about 2 GB) were copied **twice** per worker, while the model files were already all lying safe and sound in the `/opt/ml/model` directory in the Docker container.

### Reason
The reason for this behavior is that each worker calls upon initialization `pyfunc.load_pyfunc("/opt/ml/model/")` (see `wsgi.py` in `mlflow/sagemaker/container/scoring_server`).
In `pyfunc.load_pyfunc`
i) the model files are "downloaded" (copied to a temporary directory in `/tmp` via `local_model_path = _download_artifact_from_uri(artifact_uri=model_uri)`. Then, the `_load_pyfunc` in the pyfunc loader_module (`model.py`) is called. This again,
ii) copies all artifact files to a temporary directory with suffix "artifacts"

### My workaround
As a quick fix for myself, I
i) set `local_model_path = local_file_uri_to_path(model_uri)`in `mlflow/pyfunc/__init__.py`
ii) set `artifacts[saved_artifact_name] = os.path.join(model_path, saved_artifact_info[CONFIG_KEY_ARTIFACT_RELATIVE_PATH])` instead of `tmp_artifact_path` in `mlflow/pyfunc/model.py` and commented out the creation of `tmp_artifact_path` which involves `_copy_file_or_tree` to copy the artifact files.

### Request
Since I believe I'm not the only one working with larger model artifacts, I suggest improving the handling of the model artifacts. Besides, it might be reasonable to allow for a custom gunicorn timeout in case model initialization is more involved.


Discussion:
+1 custom gunicorn timeout and
+1 to not duplicating large model artifacts. Currently have a model with a 4Gb artifact... it used to be a 50Gb kdtree 😳. 
Thanks for reporting this. _download_artifact_from_uri should actually not do anything special for local paths, so there seems to be something wrong there (maybe it got changed at some point).
https://github.com/mlflow/mlflow/pull/1307 should resolve half of this issue. We still need to do a bit of work to make sure we don't copy the model's dependent artifacts https://github.com/mlflow/mlflow/blob/6f7b3248b8ac98086113593effb2e3550d7a524b/mlflow/pyfunc/model.py#L217-L228 as @dbczumar pointed out.

Resolving patch:
diff --git a/mlflow/pyfunc/model.py b/mlflow/pyfunc/model.py
--- a/mlflow/pyfunc/model.py
+++ b/mlflow/pyfunc/model.py
@@ -4,7 +4,6 @@
 """
 
 import os
-import tempfile
 import shutil
 import yaml
 from abc import ABCMeta, abstractmethod
@@ -212,20 +211,11 @@ def _load_pyfunc(model_path):
     with open(os.path.join(model_path, python_model_subpath), "rb") as f:
         python_model = cloudpickle.load(f)
 
-    # TODO: If the longevity of the temporary directory prior becomes problematic, consider using
-    # an alternative solution.
-    tmp_artifacts_dir_path = tempfile.mkdtemp(suffix="artifacts")
     artifacts = {}
     for saved_artifact_name, saved_artifact_info in\
             pyfunc_config.get(CONFIG_KEY_ARTIFACTS, {}).items():
-        tmp_artifact_path = os.path.join(
-                tmp_artifacts_dir_path,
-                _copy_file_or_tree(
-                    src=os.path.join(
-                        model_path, saved_artifact_info[CONFIG_KEY_ARTIFACT_RELATIVE_PATH]),
-                    dst=tmp_artifacts_dir_path,
-                    dst_dir=saved_artifact_name))
-        artifacts[saved_artifact_name] = tmp_artifact_path
+        artifacts[saved_artifact_name] = os.path.join(
+            model_path, saved_artifact_info[CONFIG_KEY_ARTIFACT_RELATIVE_PATH])
 
     context = PythonModelContext(artifacts=artifacts)
     python_model.load_context(context=context)


Test patch: