Fix flaky test_download_artifacts in tests/store/artifact/test_http_artifact_repo.py

[harupy]

Signed-off-by: harupy hkawamura0130@gmail.com

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Fix flaky test_download_artifacts in tests/store/artifact/test_http_artifact_repo.py.

https://github.com/mlflow/mlflow/runs/8187814349?check_suite_focus=true

    def list_artifacts(self, path=None):
        endpoint = "/mlflow-artifacts/artifacts"
        url, tail = self.artifact_uri.split(endpoint, maxsplit=1)
        root = tail.lstrip("/")
        params = {"path": posixpath.join(root, path) if path else root}
        host_creds = _get_default_host_creds(url)
        resp = http_request(host_creds, endpoint, "GET", params=params)
        augmented_raise_for_status(resp)
        file_infos = []
>       for f in resp.json().get("files", []):
E       AttributeError: 'str' object has no attribute 'get'

How is this patch tested?

• I have written tests (not required for typo or doc fix) and confirmed the proposed feature/bug-fix/change works.

Updated test

Does this PR change the documentation?

• No. You can skip the rest of this section.
• Yes. Make sure the changed pages / sections render correctly by following the steps below.

1. Click the Details link on the Preview docs check.
2. Find the changed pages / sections and make sure they render correctly.

Release Notes

Is this a user-facing change?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release notes for MLflow users.

(Details in 1-2 sentences. You can just refer to another PR with a description if this PR is part of a larger change.)

What component(s), interfaces, languages, and integrations does this PR 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/pipelines: Pipelines, Pipeline APIs, Pipeline configs, Pipeline Templates
• area/projects: MLproject format, project running backends
• area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• area/server-infra: MLflow Tracking server backend
• 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

How should the PR be classified in the release notes? Choose one:

• 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
• 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

[harupy]

The traceback below indicates that a list-artifacts request (to list artifacts in the dir directory) was responded with MockStreamResponse("data_a", 200).

    def list_artifacts(self, path=None):
        endpoint = "/mlflow-artifacts/artifacts"
        url, tail = self.artifact_uri.split(endpoint, maxsplit=1)
        root = tail.lstrip("/")
        params = {"path": posixpath.join(root, path) if path else root}
        host_creds = _get_default_host_creds(url)
        resp = http_request(host_creds, endpoint, "GET", params=params)
        augmented_raise_for_status(resp)
        file_infos = []
>       for f in resp.json().get("files", []):
E       AttributeError: 'str' object has no attribute 'get'

endpoint   = '/mlflow-artifacts/artifacts'
file_infos = []
host_creds = <mlflow.utils.rest_utils.MlflowHostCreds object at 0x0000022673C56B88>
params     = {'path': 'dir'}
path       = 'dir'
resp       = <tests.store.artifact.test_mlflow_artifact_repo.MockStreamResponse object at 0x0000022671828D88>
root       = ''
self       = <mlflow.store.artifact.mlflow_artifacts_repo.MlflowArtifactsRepository object at 0x00000226718288C8>
tail       = ''
url        = 'http://localhost:5000/api/2.0'

[harupy]

This error can be reproduced by modifying async_download_artifact like this:

diff --git a/mlflow/store/artifact/artifact_repo.py b/mlflow/store/artifact/artifact_repo.py
index 63a0195c3..a0c52e242 100644
--- a/mlflow/store/artifact/artifact_repo.py
+++ b/mlflow/store/artifact/artifact_repo.py
@@ -152,8 +152,14 @@ class ArtifactRepository:
             local_destination_file_path = self._create_download_destination(
                 src_artifact_path=src_artifact_path, dst_local_dir_path=dst_local_dir_path
             )
+            import time
+
+            def delayed_download_file(remote_file_path, local_path):
+                time.sleep(1)
+                return self._download_file(remote_file_path, local_path)
+
             download_future = self.thread_pool.submit(
-                self._download_file,
+                delayed_download_file,
                 remote_file_path=src_artifact_path,
                 local_path=local_destination_file_path,
             )

[harupy]

This test assumes the following request order:

1. list_artifacts
2. list_artifacts
3. download_artifact <-
4. list_artifacts.   <- 
5. download_artifact

but request 3 and 4 could be swapped because a download-artifact operation is performed asynchronously.

[harupy]

In the fixed test, the request order doesn't matter.

[harupy]

I've observed this error a few times before on Windows, but not sure why it only occurs on Windows.

[WeichenXu123]

LGTM