Repository: mlflow/mlflow
Pull request: Don't copy local artifacts in download_artifacts
## What changes are proposed in this pull request?
 
Trying to resolve https://github.com/mlflow/mlflow/issues/1283.

In this PR, we will no longer copy artifacts to a tmp directory when `LocalArtifactRepository.download_artifact` is invoked with `dst_path=None`.

In order to actually resolve the linked issue, we also had to change some of the behavior of the predict/serve functions for the rfunc/pyfunc flavors. In the old behavior, we would clean up the copied or downloaded artifacts upon shutdown of the python process. In the new behavior, we will copy it directly to the `/tmp` directory and it will be the OS/users responsibility to clean it up.

## How is this patch tested?
 
- New unit test verifying non-copying behavior on LocalArtifactRepository
- https://github.com/mlflow/mlflow/blob/master/tests/models/test_cli.py#L73 should catch any regressions on serving models logged to local artifact repository.
  
## Release Notes
 
### Is this a user-facing change? 

- [ ] No. You can skip the rest of this section.
- [x] Yes. Give a description of this change to be included in the release notes for MLflow users.
 
LocalArtifactRepository will no longer copy artifacts on `download_artifacts` if `dst_path=None`. This is a small breaking change since it means files returned by `download_artifact(dst_path=None)` are no longer mutable.

### What component(s) does this PR affect?
 
- [ ] UI
- [ ] CLI 
- [ ] API 
- [ ] REST-API 
- [ ] Examples 
- [ ] Docs
- [ ] Tracking
- [ ] Projects 
- [x] Artifacts 
- [x] Models 
- [ ] Scoring 
- [ ] Serving
- [ ] R
- [ ] Java
- [ ] Python

### How should the PR be classified in the release notes? Choose one:
 
- [x] `rn/breaking-change` - The PR will be mentioned in the "Breaking Changes" section
- [ ] `rn/none` - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
- [ ] `rn/feature` - A new user-facing feature worth mentioning in the release notes
- [x] `rn/bug-fix` - A user-facing bug fix worth mentioning in the release notes
- [ ] `rn/documentation` - A user-facing documentation change worth mentioning in the release notes


Problem statement:
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.

Hints from issue 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).

Commit message: fix lint

Patch to review:
diff --git a/mlflow/pyfunc/backend.py b/mlflow/pyfunc/backend.py
index c81ef7c32c790..e933183943d98 100644
--- a/mlflow/pyfunc/backend.py
+++ b/mlflow/pyfunc/backend.py
@@ -27,46 +27,44 @@ def predict(self, model_uri, input_path, output_path, content_type, json_format)
         Generate predictions using generic python model saved with MLflow.
         Return the prediction results as a JSON.
         """
-        with TempDir() as tmp:
-            local_path = _download_artifact_from_uri(model_uri, output_path=tmp.path())
-            if not self._no_conda and ENV in self._config:
-                conda_env_path = os.path.join(local_path, self._config[ENV])
-                # NOTE: We're calling main in the pyfunc scoring server belonging to the current
-                # conda environment. The model environment may contain mlflow with different version
-                # than the one in the current active environment. This is the intended behavior.
-                # We need to make sure the scoring server is consistent with the outside mlflow
-                # while the model that is being loaded may depend on a different version of mlflow.
-                # The hope is that the scoring server is self contained enough and does not have
-                # external mlflow dependencies that would be incompatible between mlflow versions.
-                if input_path is None:
-                    input_path = "__stdin__"
-                if output_path is None:
-                    output_path = "__stdout__"
-                command = "python {0} predict {1} {2} {3} {4} {5}".format(scoring_server.__file__,
-                                                                          shlex_quote(local_path),
-                                                                          shlex_quote(input_path),
-                                                                          shlex_quote(output_path),
-                                                                          content_type,
-                                                                          json_format)
-                return scoring_server._execute_in_conda_env(conda_env_path, command)
-            else:
-                scoring_server._predict(local_path, input_path, output_path, content_type,
-                                        json_format)
+        local_path = _download_artifact_from_uri(model_uri)
+        if not self._no_conda and ENV in self._config:
+            conda_env_path = os.path.join(local_path, self._config[ENV])
+            # NOTE: We're calling main in the pyfunc scoring server belonging to the current
+            # conda environment. The model environment may contain mlflow with different version
+            # than the one in the current active environment. This is the intended behavior.
+            # We need to make sure the scoring server is consistent with the outside mlflow
+            # while the model that is being loaded may depend on a different version of mlflow.
+            # The hope is that the scoring server is self contained enough and does not have
+            # external mlflow dependencies that would be incompatible between mlflow versions.
+            if input_path is None:
+                input_path = "__stdin__"
+            if output_path is None:
+                output_path = "__stdout__"
+            command = "python {0} predict {1} {2} {3} {4} {5}".format(scoring_server.__file__,
+                                                                      shlex_quote(local_path),
+                                                                      shlex_quote(input_path),
+                                                                      shlex_quote(output_path),
+                                                                      content_type,
+                                                                      json_format)
+            return scoring_server._execute_in_conda_env(conda_env_path, command)
+        else:
+            scoring_server._predict(local_path, input_path, output_path, content_type,
+                                    json_format)
 
     def serve(self, model_uri, port, host):
         """
         Serve pyfunc model locally.
         """
-        with TempDir() as tmp:
-            local_path = shlex_quote(_download_artifact_from_uri(model_uri, output_path=tmp.path()))
-            if not self._no_conda and ENV in self._config:
-                conda_env_path = os.path.join(local_path, self._config[ENV])
-                command = "python {0} serve {1} {2} {3}".format(scoring_server.__file__,
-                                                                shlex_quote(local_path),
-                                                                port, host)
-                return scoring_server._execute_in_conda_env(conda_env_path, command)
-            else:
-                scoring_server._serve(local_path, port, host)
+        local_path = shlex_quote(_download_artifact_from_uri(model_uri))
+        if not self._no_conda and ENV in self._config:
+            conda_env_path = os.path.join(local_path, self._config[ENV])
+            command = "python {0} serve {1} {2} {3}".format(scoring_server.__file__,
+                                                            shlex_quote(local_path),
+                                                            port, host)
+            return scoring_server._execute_in_conda_env(conda_env_path, command)
+        else:
+            scoring_server._serve(local_path, port, host)
 
     def can_score_model(self):
         if self._no_conda:
diff --git a/mlflow/rfunc/backend.py b/mlflow/rfunc/backend.py
index 25cb556275df2..86fdfd2e0ec97 100644
--- a/mlflow/rfunc/backend.py
+++ b/mlflow/rfunc/backend.py
@@ -23,26 +23,24 @@ def predict(self, model_uri, input_path, output_path, content_type, json_format)
         Generate predictions using R model saved with MLflow.
         Return the prediction results as a JSON.
         """
-        with TempDir() as tmp:
-            model_path = _download_artifact_from_uri(model_uri, output_path=tmp.path())
-            str_cmd = "mlflow:::mlflow_rfunc_predict(model_path = '{0}', input_path = {1}, " \
-                      "output_path = {2}, content_type = {3}, json_format = {4})"
-            command = str_cmd.format(shlex_quote(model_path),
-                                     _str_optional(input_path),
-                                     _str_optional(output_path),
-                                     _str_optional(content_type),
-                                     _str_optional(json_format))
-            _execute(command)
+        model_path = _download_artifact_from_uri(model_uri)
+        str_cmd = "mlflow:::mlflow_rfunc_predict(model_path = '{0}', input_path = {1}, " \
+                  "output_path = {2}, content_type = {3}, json_format = {4})"
+        command = str_cmd.format(shlex_quote(model_path),
+                                 _str_optional(input_path),
+                                 _str_optional(output_path),
+                                 _str_optional(content_type),
+                                 _str_optional(json_format))
+        _execute(command)
 
     def serve(self, model_uri, port, host):
         """
         Generate R model locally.
         """
-        with TempDir() as tmp:
-            model_path = _download_artifact_from_uri(model_uri, output_path=tmp.path())
-            command = "mlflow::mlflow_rfunc_serve('{0}', port = {1}, host = '{2}')".format(
-                shlex_quote(model_path), port, host)
-            _execute(command)
+        model_path = _download_artifact_from_uri(model_uri)
+        command = "mlflow::mlflow_rfunc_serve('{0}', port = {1}, host = '{2}')".format(
+            shlex_quote(model_path), port, host)
+        _execute(command)
 
     def can_score_model(self):
         process = subprocess.Popen(["Rscript", "--version"], close_fds=True,
diff --git a/mlflow/store/local_artifact_repo.py b/mlflow/store/local_artifact_repo.py
index 86b9bd15092aa..ddd564275bff1 100644
--- a/mlflow/store/local_artifact_repo.py
+++ b/mlflow/store/local_artifact_repo.py
@@ -39,6 +39,27 @@ def log_artifacts(self, local_dir, artifact_path=None):
             mkdir(artifact_dir)
         dir_util.copy_tree(src=local_dir, dst=artifact_dir)
 
+    def download_artifacts(self, artifact_path, dst_path=None):
+        """
+        Download an artifact file or directory to a local directory if applicable, and return a
+        local path for it.
+        The caller is responsible for managing the lifecycle of the downloaded artifacts. If
+        dst_path is unspecified, the downloaded artifact may not be mutated in general.
+
+        :param artifact_path: Relative source path to the desired artifacts.
+        :param dst_path: Absolute path of the local filesystem destination directory to which to
+                         download the specified artifacts. This directory must already exist. If
+                         unspecified, the absolute path of the local artifact will be returned.
+
+        :return: Absolute path of the local filesystem location containing the desired artifacts.
+        """
+        if dst_path:
+            return super(LocalArtifactRepository, self).download_artifacts(artifact_path, dst_path)
+        # NOTE: The remote_file_path is expected to be in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        local_artifact_path = os.path.join(self.artifact_dir, os.path.normpath(artifact_path))
+        return os.path.abspath(local_artifact_path)
+
     def list_artifacts(self, path=None):
         # NOTE: The path is expected to be in posix format.
         # Posix paths work fine on windows but just in case we normalize it here.
diff --git a/tests/store/test_local_artifact_repo.py b/tests/store/test_local_artifact_repo.py
index c863cbb16d0f2..84134bc48efc1 100644
--- a/tests/store/test_local_artifact_repo.py
+++ b/tests/store/test_local_artifact_repo.py
@@ -61,6 +61,24 @@ def test_download_artifacts(local_artifact_repo):
         assert open(dst_path).read() == artifact_text
 
 
+def test_download_artifacts_does_not_copy(local_artifact_repo):
+    """
+    The LocalArtifactRepository.download_artifact function should not copy the artifact if
+    the ``dst_path`` argument is None.
+    """
+    artifact_rel_path = "test.txt"
+    artifact_text = "hello world!"
+    with TempDir(chdr=True) as local_dir:
+        artifact_src_path = local_dir.path(artifact_rel_path)
+        with open(artifact_src_path, "w") as f:
+            f.write(artifact_text)
+        local_artifact_repo.log_artifact(artifact_src_path)
+        dst_path = local_artifact_repo.download_artifacts(artifact_path=artifact_rel_path)
+        assert open(dst_path).read() == artifact_text
+        assert dst_path.startswith(local_artifact_repo.artifact_dir), \
+            'downloaded artifact is not in local_artifact_repo.artifact_dir root'
+
+
 def test_download_artifacts_returns_absolute_paths(local_artifact_repo):
     artifact_rel_path = "test.txt"
     artifact_text = "hello world!"


Code review comments:
diff_hunk: @@ -39,6 +39,27 @@ def log_artifacts(self, local_dir, artifact_path=None):
             mkdir(artifact_dir)
         dir_util.copy_tree(src=local_dir, dst=artifact_dir)
 
+    def download_artifacts(self, artifact_path, dst_path=None):
+        """
+        Download an artifact file or directory to a local directory if applicable, and return a
+        local path for it.
+        The caller is responsible for managing the lifecycle of the downloaded artifacts. If
+        dst_path is unspecified, the downloaded artifact may not be mutated in general.
+
+        :param artifact_path: Relative source path to the desired artifacts.
+        :param dst_path: Absolute path of the local filesystem destination directory to which to
+                         download the specified artifacts. This directory must already exist. If
+                         unspecified, the absolute path of the local artifact will be returned.
+
+        :return: Absolute path of the local filesystem location containing the desired artifacts.
+        """
+        if dst_path:
+            return super(LocalArtifactRepository, self).download_artifacts(artifact_path, dst_path)
+        # NOTE: The remote_file_path is expected to be in posix format.
path: mlflow/store/local_artifact_repo.py
text: @user1:
Nit: I think this should say `artifact_path` instead of `remote_file_path`.

diff_hunk: @@ -39,6 +39,27 @@ def log_artifacts(self, local_dir, artifact_path=None):
             mkdir(artifact_dir)
         dir_util.copy_tree(src=local_dir, dst=artifact_dir)
 
+    def download_artifacts(self, artifact_path, dst_path=None):
+        """
+        Download an artifact file or directory to a local directory if applicable, and return a
+        local path for it.
+        The caller is responsible for managing the lifecycle of the downloaded artifacts. If
+        dst_path is unspecified, the downloaded artifact may not be mutated in general.
path: mlflow/store/local_artifact_repo.py
text: @user1:
To me, `mutated` implies that the artifact file / artifact directory contents will be somehow altered by `download_artifacts()`, which is not the case. Perhaps something like the following:

```
Artifacts tracked by ``LocalArtifactRepository`` already exist on the local filesystem. If ``dst_path`` is ``None``, the absolute filesystem path of the specified artifact is returned. If ``dst_path`` is not ``None``, the local artifact is copied to ``dst_path``.
```