Repository: mlflow/mlflow

Issue:
[BUG] 1.24.0 regression: Deadlock when downloading directory (async download)
# Regression: Downloading directories using SFTP may cause deadlocks on 1.24.0

### Willingness to contribute
The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

- [ ] Yes. I can contribute a fix for this bug independently.
- [x] Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
- [ ] No. I cannot contribute a bug fix at this time.

### System information
- **OS Platform and Distribution**: `Ubuntu 20.04.3 LTS`
- **MLflow installed from**: binary
- **MLflow version**: 1.24 (works on version 1.23)
- **Python version**: 3.8.10

### Describe the problem
Downloading directory artifacts causes a deadlock, and sometimes a paramiko issue (Garbage package received) when using sftp artifact storage.

### Code to reproduce issue
Working:
```python
from mlflow.tracking import MlflowClient
run_id = "run-id"
server = "server-uri"
# This works; fetches the contents of model sequentually.
MlflowClient(server).download_artifacts(run_id, "model/MLmodel", "out")
MlflowClient(server).download_artifacts(run_id, "model/conda.yaml", "out")
MlflowClient(server).download_artifacts(run_id, "model/model.pkl", "out")
MlflowClient(server).download_artifacts(run_id, "model/requirements.txt", "out")
```
Not working:
```python
# Hangs or throws exception
MlflowClient(server).download_artifacts(run_id, "model", "out2")
```
Both work on 1.23.1

### Other info / logs
Paramiko exception:
```
---------------------------------------------------------------------------
MlflowException                           Traceback (most recent call last)
<ipython-input-4-3293f4eca114> in <module>
----> 1 MlflowClient("<server url>").download_artifacts("run-id", "model", "out2")

~/.local/lib/python3.8/site-packages/mlflow/tracking/client.py in download_artifacts(self, run_id, path, dst_path)
   1411             Artifacts: ['features.txt']
   1412         """
-> 1413         return self._tracking_client.download_artifacts(run_id, path, dst_path)
   1414 
   1415     def set_terminated(

~/.local/lib/python3.8/site-packages/mlflow/tracking/_tracking_service/client.py in download_artifacts(self, run_id, path, dst_path)
    389         :return: Local path of desired artifact.
    390         """
--> 391         return self._get_artifact_repo(run_id).download_artifacts(path, dst_path)
    392 
    393     def set_terminated(self, run_id, status=None, end_time=None):

~/.local/lib/python3.8/site-packages/mlflow/store/artifact/artifact_repo.py in download_artifacts(self, artifact_path, dst_path)
    263 
    264         if len(failed_downloads) > 0:
--> 265             raise MlflowException(
    266                 message=(
    267                     "The following failures occurred while downloading one or more"

MlflowException: The following failures occurred while downloading one or more artifacts from sftp://server:/mlflow_ftp/upload/exp-id/run-id/artifacts: {'model/MLmodel': "SFTPError('Garbage packet received')", 'model/conda.yaml': "SFTPError('Garbage packet received')", 'model/model.pkl': "SFTPError('Garbage packet received')", 'model/requirements.txt': "SFTPError('Garbage packet received')"}
```

Exception when interrupting deadlock:
```
---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
<ipython-input-4-3293f4eca114> in <module>
----> 1 MlflowClient(server).download_artifacts(run_id, "model", "out2")

~/.local/lib/python3.8/site-packages/mlflow/tracking/client.py in download_artifacts(self, run_id, path, dst_path)
   1411             Artifacts: ['features.txt']
   1412         """
-> 1413         return self._tracking_client.download_artifacts(run_id, path, dst_path)
   1414 
   1415     def set_terminated(

~/.local/lib/python3.8/site-packages/mlflow/tracking/_tracking_service/client.py in download_artifacts(self, run_id, path, dst_path)
    389         :return: Local path of desired artifact.
    390         """
--> 391         return self._get_artifact_repo(run_id).download_artifacts(path, dst_path)
    392 
    393     def set_terminated(self, run_id, status=None, end_time=None):

~/.local/lib/python3.8/site-packages/mlflow/store/artifact/artifact_repo.py in download_artifacts(self, artifact_path, dst_path)
    258         for inflight_download in inflight_downloads:
    259             try:
--> 260                 inflight_download.download_future.result()
    261             except Exception as e:
    262                 failed_downloads[inflight_download.src_artifact_path] = repr(e)

/usr/lib/python3.8/concurrent/futures/_base.py in result(self, timeout)
    437                     return self.__get_result()
    438 
--> 439                 self._condition.wait(timeout)
    440 
    441                 if self._state in [CANCELLED, CANCELLED_AND_NOTIFIED]:

/usr/lib/python3.8/threading.py in wait(self, timeout)
    300         try:    # restore state no matter what (e.g., KeyboardInterrupt)
    301             if timeout is None:
--> 302                 waiter.acquire()
    303                 gotit = True
    304             else:

KeyboardInterrupt:
```

### What component(s), interfaces, languages, and integrations does this bug affect?
Components 
- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `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
- [x] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface 
- [ ] `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

Language 
- [ ] `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:
@rsundqvist https://github.com/mlflow/mlflow/commit/70fc2dd41162156d11f6070b2e75a478db7ad627 seems suspicious. We're investigating.
@rsundqvist Have you seen https://github.com/paramiko/paramiko/issues/395?
> @rsundqvist Have you seen [paramiko/paramiko#395](https://github.com/paramiko/paramiko/issues/395)?

Yes, saw it when googling. There may be two separate issues at play here, but I did not see any deadlocks running the old version whereas it happened 100% of the time running `1.24` except when the paramiko issue showed up.

We've had a lot of issues with paramiko in general but I don't recall seeing this particular error before. I couldn't reliably reproduce the issue running `paramiko==2.10.3`.
@rsundqvist Thanks for the reply. Can you create a reproducible example (e.g. by using docker compose)?
@rsundqvist We created a [branch](https://github.com/harupy/mlflow/commits/issue-5656
) that contains a fix for this issue. Can you install mlflow from this branch using the command below and check if it can fix the deadlock issue?

```
pip install git+https://github.com/harupy/mlflow.git@issue-5656
```


Could not use that command:
```console
ricsun@the-machine:~/git/repo-that-uses-mlflow$ pip install git+https://github.com/harupy/mlflow/commits/issue-5656
Defaulting to user installation because normal site-packages is not writeable
Looking in indexes: <company private index>, https://pypi.python.org/simple
Collecting git+https://github.com/harupy/mlflow/commits/issue-5656
  Cloning https://github.com/harupy/mlflow/commits/issue-5656 to /tmp/pip-req-build-hgv857xx
  Running command git clone -q https://github.com/harupy/mlflow/commits/issue-5656 /tmp/pip-req-build-hgv857xx
  fatal: repository 'https://github.com/harupy/mlflow/commits/issue-5656/' not found
```
Using `pip install git+https://github.com/harupy/mlflow.git@issue-5656` did the trick though, giving `1.24.1.dev0` which is what I would expect looking at your branch.

@harupy Doesn't seem to work unfortunately. Deadlocks in the same place.
> @rsundqvist Thanks for the reply. Can you create a reproducible example (e.g. by using docker compose)?

I'll see what I can do.
@rsundqvist Sorry for the wrong command, I fixed it.

> I'll see what I can do.

Thank you!
> @rsundqvist Sorry for the wrong command, I fixed it.
> 
> > I'll see what I can do.
> 
> Thank you!

Looked into this. The issue seems to be that the paramiko SFTP client is not thread safe which causes a deadlock. The `SFTPArtifactRepository` class probably needs a connection pool to support multithreading.

I've tested this a bit and it seems you can "corrupt" a `pysftp.Connection` by continously triggering multi-threaded downloads. Once in a bad state the error we discussed above is thrown even if trying to fetch "correctly" from a single thread.
@rsundqvist Thank you for diagnosing this! Would you be able to introduce a connection pool to `SFTPArtifactRepository` to address this issue?
@dbczumar I have an ad-hoc solution that works for the previously failing test case that I included in the initial report, however it will fail in cases where the `client` argument is given when creating the `SFTPArtifactRepository` instance. 

Any idea when (if?) this will happen? I didn't find any conditions under which `client` was anything other than `None` during my testing.
> @dbczumar I have an ad-hoc solution that works for the previously failing test case that I included in the initial report, however it will fail in cases where the `client` argument is given when creating the `SFTPArtifactRepository` instance.



I have the exact same issue working with mlflow when trying to load models from the previous runs. Can you share the ad-hoc method which is working for you?
@KrSiddhartha They're on master in [rsundqvist/mlflow](https://github.com/rsundqvist/mlflow). See the commits linked above in this thread :)

Can be installed with pip using `pip install git+https://github.com/rsundqvist/mlflow.git`
> @KrSiddhartha They're on master in [rsundqvist/mlflow](https://github.com/rsundqvist/mlflow). See the commits linked above in this thread :)
> 
> Can be installed with pip using `pip install git+https://github.com/rsundqvist/mlflow.git`

Hi, I had the same issue of deadlocks with sftp downloads. Your ad-hoc fix did solve the download issues for me and it seems like all the files are completely present but it still gets deadlocked.

From what I found the deadlock seems to occur in artifact_repo.py where the results of all download tasks are collected. There is always the same amount of download tasks that seem to work fine but at some point the download_future result does deadlock.
As an example you could say I have 100 files that are downloaded and always at the 30th the call of the result results in a deadlock. Interestingly increasing the workercount for the sftp pool increases this number but the effect still occurs.

In artifact_repo.py its line 258.
When I looked into the code from pythons concurrency Future it seems to be the wait. I dont know what this condition really is and why it waits for it forever.

``` python
    def result(self, timeout=None):
        with self._condition:
            if self._state in [CANCELLED, CANCELLED_AND_NOTIFIED]:
                raise CancelledError()
            elif self._state == FINISHED:
                return self.__get_result()

            self._condition.wait(timeout) # <-- This one

            if self._state in [CANCELLED, CANCELLED_AND_NOTIFIED]:
                raise CancelledError()
            elif self._state == FINISHED:
                return self.__get_result()
            else:
                raise TimeoutError()
```

Resolving patch:
diff --git a/mlflow/store/artifact/artifact_repo.py b/mlflow/store/artifact/artifact_repo.py
--- a/mlflow/store/artifact/artifact_repo.py
+++ b/mlflow/store/artifact/artifact_repo.py
@@ -34,9 +34,7 @@ def __init__(self, artifact_uri):
         # Limit the number of threads used for artifact uploads/downloads. Use at most
         # constants._NUM_MAX_THREADS threads or 2 * the number of CPU cores available on the
         # system (whichever is smaller)
-        num_cpus = os.cpu_count() or _NUM_DEFAULT_CPUS
-        num_artifact_workers = min(num_cpus * _NUM_MAX_THREADS_PER_CPU, _NUM_MAX_THREADS)
-        self.thread_pool = ThreadPoolExecutor(max_workers=num_artifact_workers)
+        self.thread_pool = ThreadPoolExecutor(max_workers=self.max_workers)
 
     @abstractmethod
     def log_artifact(self, local_file, artifact_path=None):
@@ -294,6 +292,12 @@ def delete_artifacts(self, artifact_path=None):
         """
         pass
 
+    @property
+    def max_workers(self) -> int:
+        """Compute the number of workers to use for multi-threading."""
+        num_cpus = os.cpu_count() or _NUM_DEFAULT_CPUS
+        return min(num_cpus * _NUM_MAX_THREADS_PER_CPU, _NUM_MAX_THREADS)
+
 
 def verify_artifact_path(artifact_path):
     if artifact_path and path_not_unique(artifact_path):
diff --git a/mlflow/store/artifact/sftp_artifact_repo.py b/mlflow/store/artifact/sftp_artifact_repo.py
--- a/mlflow/store/artifact/sftp_artifact_repo.py
+++ b/mlflow/store/artifact/sftp_artifact_repo.py
@@ -1,11 +1,13 @@
 import os
 import sys
+import threading
 
 import posixpath
 import urllib.parse
 
 from mlflow.entities import FileInfo
 from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from contextlib import contextmanager
 
 
 # Based on: https://stackoverflow.com/a/58466685
@@ -20,6 +22,20 @@ def _put_r_for_windows(sftp, local_dir, remote_dir, preserve_mtime=False):
             sftp.put(local_path, remote_path, preserve_mtime=preserve_mtime)
 
 
+class _SftpPool:
+    def __init__(self, connections):
+        self._connections = connections
+        self._cond = threading.Condition()
+
+    @contextmanager
+    def get_sfp_connection(self):
+        with self._cond:
+            self._cond.wait_for(lambda: bool(self._connections))
+            connection = self._connections.pop(-1)
+        yield connection
+        self._connections.append(connection)
+
+
 class SFTPArtifactRepository(ArtifactRepository):
     """Stores artifacts as files in a remote directory, via sftp."""
 
@@ -63,53 +79,63 @@ def __init__(self, artifact_uri):
         if "identityfile" in user_config:
             self.config["private_key"] = user_config["identityfile"][0]
 
-        self.sftp = pysftp.Connection(**self.config)
+        connections = [pysftp.Connection(**self.config) for _ in range(self.max_workers)]
+        self.pool = _SftpPool(connections)
 
         super().__init__(artifact_uri)
 
     def log_artifact(self, local_file, artifact_path=None):
         artifact_dir = posixpath.join(self.path, artifact_path) if artifact_path else self.path
-        self.sftp.makedirs(artifact_dir)
-        self.sftp.put(local_file, posixpath.join(artifact_dir, os.path.basename(local_file)))
+        with self.pool.get_sfp_connection() as sftp:
+            sftp.makedirs(artifact_dir)
+            sftp.put(local_file, posixpath.join(artifact_dir, os.path.basename(local_file)))
 
     def log_artifacts(self, local_dir, artifact_path=None):
         artifact_dir = posixpath.join(self.path, artifact_path) if artifact_path else self.path
-        self.sftp.makedirs(artifact_dir)
-        if sys.platform == "win32":
-            _put_r_for_windows(self.sftp, local_dir, artifact_dir)
-        else:
-            self.sftp.put_r(local_dir, artifact_dir)
+        with self.pool.get_sfp_connection() as sftp:
+            sftp.makedirs(artifact_dir)
+            if sys.platform == "win32":
+                _put_r_for_windows(sftp, local_dir, artifact_dir)
+            else:
+                sftp.put_r(local_dir, artifact_dir)
 
     def _is_directory(self, artifact_path):
         artifact_dir = self.path
         path = posixpath.join(artifact_dir, artifact_path) if artifact_path else artifact_dir
-        return self.sftp.isdir(path)
+        with self.pool.get_sfp_connection() as sftp:
+            return sftp.isdir(path)
 
     def list_artifacts(self, path=None):
         artifact_dir = self.path
         list_dir = posixpath.join(artifact_dir, path) if path else artifact_dir
-        if not self.sftp.isdir(list_dir):
-            return []
-        artifact_files = self.sftp.listdir(list_dir)
-        infos = []
-        for file_name in artifact_files:
-            file_path = file_name if path is None else posixpath.join(path, file_name)
-            full_file_path = posixpath.join(list_dir, file_name)
-            if self.sftp.isdir(full_file_path):
-                infos.append(FileInfo(file_path, True, None))
-            else:
-                infos.append(FileInfo(file_path, False, self.sftp.stat(full_file_path).st_size))
-        return infos
+        with self.pool.get_sfp_connection() as sftp:
+            if not sftp.isdir(list_dir):
+                return []
+            artifact_files = sftp.listdir(list_dir)
+            infos = []
+            for file_name in artifact_files:
+                file_path = file_name if path is None else posixpath.join(path, file_name)
+                full_file_path = posixpath.join(list_dir, file_name)
+                if sftp.isdir(full_file_path):
+                    infos.append(FileInfo(file_path, True, None))
+                else:
+                    infos.append(FileInfo(file_path, False, sftp.stat(full_file_path).st_size))
+            return infos
 
     def _download_file(self, remote_file_path, local_path):
         remote_full_path = posixpath.join(self.path, remote_file_path)
-        self.sftp.get(remote_full_path, local_path)
+        with self.pool.get_sfp_connection() as sftp:
+            sftp.get(remote_full_path, local_path)
 
     def delete_artifacts(self, artifact_path=None):
-        if self.sftp.isdir(artifact_path):
-            with self.sftp.cd(artifact_path):
-                for element in self.sftp.listdir():
-                    self.delete_artifacts(element)
-            self.sftp.rmdir(artifact_path)
-        elif self.sftp.isfile(artifact_path):
-            self.sftp.remove(artifact_path)
+        with self.pool.get_sfp_connection() as sftp:
+            self._delete_inner(artifact_path, sftp)
+
+    def _delete_inner(self, artifact_path, sftp):
+        if sftp.isdir(artifact_path):
+            with sftp.cd(artifact_path):
+                for element in sftp.listdir():
+                    self._delete_inner(element, sftp)
+            sftp.rmdir(artifact_path)
+        elif sftp.isfile(artifact_path):
+            sftp.remove(artifact_path)


Test patch:
diff --git a/tests/store/artifact/test_sftp_artifact_repo.py b/tests/store/artifact/test_sftp_artifact_repo.py
--- a/tests/store/artifact/test_sftp_artifact_repo.py
+++ b/tests/store/artifact/test_sftp_artifact_repo.py
@@ -4,6 +4,7 @@
 from mlflow.store.artifact.sftp_artifact_repo import SFTPArtifactRepository
 from mlflow.utils.file_utils import TempDir
 import os
+import mlflow
 import posixpath
 
 
@@ -196,3 +197,22 @@ def test_delete_selective_artifacts(artifact_path):
         assert not posixpath.exists(posixpath.join(remote_dir, file1))
         assert posixpath.isfile(posixpath.join(remote_dir, directory, file2))
         assert posixpath.isdir(remote_dir)
+
+
+def test_log_and_download_sklearn_model(tmp_path):
+    from sklearn.linear_model import LogisticRegression
+    from sklearn.datasets import load_iris
+    from numpy.testing import assert_allclose
+
+    X, y = load_iris(return_X_y=True)
+    original = LogisticRegression().fit(X, y)
+
+    experiment_id = mlflow.create_experiment(
+        name="sklearn-model-experiment",
+        artifact_location=f"sftp://{tmp_path}",
+    )
+    with mlflow.start_run(experiment_id=experiment_id):
+        model_uri = mlflow.sklearn.log_model(original, "model").model_uri
+        downloaded = mlflow.sklearn.load_model(model_uri)
+
+    assert_allclose(original.predict(X), downloaded.predict(X))