Repository: mlflow/mlflow

Issue:
[BUG] log_artifact timeout when using Azure Blob Storage
### 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
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**: Yes
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Windows 10
- **MLflow installed from (source or binary)**: Binary (pip)
- **MLflow version (run ``mlflow --version``)**: 1.12.0
- **Python version**: 3.6
- **npm version, if running the dev UI**:
- **Exact command to reproduce**:

```python
mlflow.log_artifact("/path/to/file", artifact_path="<<Azure Blob Storage URI>>")
```

### Describe the problem
Similar to issue  #3478 (except this is referring to the Python version), the Azure Blob Storage Python SDK causes a timeout if the upload takes too long. This can be solved by setting the `timeout` keyword when uploading a file, but MLFlow does not provide a way to set this value. I have implemented [solution](https://github.com/ismailuddin/mlflow/commit/e0e9148ea8a7e19ecb060de750bd3c5e02d71b34) in a fork of my own, which uses an environment value that can be set. If this approach is suitable, I can make a PR for it.

### Code to reproduce issue
```python
mlflow.log_artifact("/path/to/file", artifact_path="<<Azure Blob Storage URI>>")
```

### Other info / logs
Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.


### What component(s), interfaces, languages, and integrations does this bug affect?
Components 
- [x] `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`: Local serving, model deployment tools, spark UDFs
- [ ] `area/server-infra`: MLflow server, JavaScript dev server
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface 
- [ ] `area/uiux`: Front-end, user experience, JavaScript, plotting
- [ ] `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
- [x] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations



Discussion:
+1 on this. Using google-cloud-store, but same issue with the `log_artifact` step timing out, same error output as #3478.  I'm running [this](https://github.com/mlflow/mlflow/tree/master/examples/pytorch/BertNewsClassification) example code.  
In case it helps in reproduction, I'm running the example in Europe and the google-cloud-store is located in South Carolina, so it's taking a while to upload.
@savagej @ismailuddin For GCS I found that setting the upload/download chunk_size would also solve this issue instead of setting the timeouts. Would that also work for azure blob storage?
@morkru  how do you change the GCS chunk size?
@jiunjiunma-outdoorsy Using the chunk_size parameter for the gcs_bucket in the mlflow/store/artifact/gcs_artifact_repo.py  file (example solution: https://github.com/morkru/mlflow/commit/ce47e0856f67270a70b0f1c9190abf04a637eb55)

Resolving patch:
diff --git a/mlflow/environment_variables.py b/mlflow/environment_variables.py
--- a/mlflow/environment_variables.py
+++ b/mlflow/environment_variables.py
@@ -100,6 +100,7 @@ def get(self):
 #: ``google-cloud-storage`` package.
 MLFLOW_GCS_UPLOAD_CHUNK_SIZE = _EnvironmentVariable("MLFLOW_GCS_UPLOAD_CHUNK_SIZE", int, None)
 
+#: (Deprecated, please use ``MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT``)
 #: Specifies the default timeout to use when downloading/uploading a file from/to GCS
 #: (default: ``None``). If None, ``google.cloud.storage.constants._DEFAULT_TIMEOUT`` is used.
 MLFLOW_GCS_DEFAULT_TIMEOUT = _EnvironmentVariable("MLFLOW_GCS_DEFAULT_TIMEOUT", int, None)
@@ -186,3 +187,10 @@ def get(self):
 MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT = _EnvironmentVariable(
     "MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT", int, 60
 )
+
+#: (Experimental, may be changed or removed)
+#: Specifies the timeout to use when uploading or downloading a file
+#: (default: ``None``). If None, individual artifact stores will choose defaults.
+MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT = _EnvironmentVariable(
+    "MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT", int, None
+)
diff --git a/mlflow/store/artifact/azure_blob_artifact_repo.py b/mlflow/store/artifact/azure_blob_artifact_repo.py
--- a/mlflow/store/artifact/azure_blob_artifact_repo.py
+++ b/mlflow/store/artifact/azure_blob_artifact_repo.py
@@ -7,6 +7,8 @@
 from mlflow.exceptions import MlflowException
 from mlflow.store.artifact.artifact_repo import ArtifactRepository
 
+from mlflow.environment_variables import MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT
+
 
 class AzureBlobArtifactRepository(ArtifactRepository):
     """
@@ -23,6 +25,9 @@ class AzureBlobArtifactRepository(ArtifactRepository):
     def __init__(self, artifact_uri, client=None):
         super().__init__(artifact_uri)
 
+        _DEFAULT_TIMEOUT = 600  # 10 minutes
+        self.write_timeout = MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT.get() or _DEFAULT_TIMEOUT
+
         # Allow override for testing
         if client:
             self.client = client
@@ -86,7 +91,9 @@ def log_artifact(self, local_file, artifact_path=None):
             dest_path = posixpath.join(dest_path, artifact_path)
         dest_path = posixpath.join(dest_path, os.path.basename(local_file))
         with open(local_file, "rb") as file:
-            container_client.upload_blob(dest_path, file, overwrite=True)
+            container_client.upload_blob(
+                dest_path, file, overwrite=True, timeout=self.write_timeout
+            )
 
     def log_artifacts(self, local_dir, artifact_path=None):
         (container, _, dest_path, _) = self.parse_wasbs_uri(self.artifact_uri)
@@ -103,7 +110,9 @@ def log_artifacts(self, local_dir, artifact_path=None):
                 remote_file_path = posixpath.join(upload_path, f)
                 local_file_path = os.path.join(root, f)
                 with open(local_file_path, "rb") as file:
-                    container_client.upload_blob(remote_file_path, file, overwrite=True)
+                    container_client.upload_blob(
+                        remote_file_path, file, overwrite=True, timeout=self.write_timeout
+                    )
 
     def list_artifacts(self, path=None):
         # Newer versions of `azure-storage-blob` (>= 12.4.0) provide a public
diff --git a/mlflow/store/artifact/gcs_artifact_repo.py b/mlflow/store/artifact/gcs_artifact_repo.py
--- a/mlflow/store/artifact/gcs_artifact_repo.py
+++ b/mlflow/store/artifact/gcs_artifact_repo.py
@@ -7,6 +7,7 @@
 from mlflow.store.artifact.artifact_repo import ArtifactRepository
 from mlflow.utils.file_utils import relative_path_to_artifact_path
 from mlflow.environment_variables import (
+    MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT,
     MLFLOW_GCS_DEFAULT_TIMEOUT,
     MLFLOW_GCS_DOWNLOAD_CHUNK_SIZE,
     MLFLOW_GCS_UPLOAD_CHUNK_SIZE,
@@ -33,7 +34,11 @@ def __init__(self, artifact_uri, client=None):
 
         self._GCS_DOWNLOAD_CHUNK_SIZE = MLFLOW_GCS_DOWNLOAD_CHUNK_SIZE.get()
         self._GCS_UPLOAD_CHUNK_SIZE = MLFLOW_GCS_UPLOAD_CHUNK_SIZE.get()
-        self._GCS_DEFAULT_TIMEOUT = MLFLOW_GCS_DEFAULT_TIMEOUT.get() or _DEFAULT_TIMEOUT
+        self._GCS_DEFAULT_TIMEOUT = (
+            MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT.get()
+            or MLFLOW_GCS_DEFAULT_TIMEOUT.get()
+            or _DEFAULT_TIMEOUT
+        )
 
         # If the user-supplied timeout environment variable value is -1,
         # use `None` for `self._GCS_DEFAULT_TIMEOUT`


Test patch: