Fix regression in downloading single files from models artifact store #10362

BenWilson2 · 2023-11-11T03:56:06Z

🛠 DevTools 🛠

Install mlflow from this PR

pip install git+https://github.com/mlflow/mlflow.git@refs/pull/10362/merge

Checkout with GitHub CLI

gh pr checkout 10362

Related Issues/PRs

Resolve #10230

What changes are proposed in this pull request?

Fixed the regression introduced in #9402 wherein no handling logic is provided for non-directory source file resolution, which would attempt to create a registration metadata file using a destination file as a directory path (which raises an Exception).

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

Validated the basic repro is broken in 2.8.0 and is fixed in this PR (test added to validate as well):

import mlflow
from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository

mlflow.set_tracking_uri("sqlite:///:memory:")

with mlflow.start_run():
    mlflow.sklearn.log_model("I am a model", "model", registered_model_name="foo")

model_uri = f"models:/foo/1"

artifact_repository = get_artifact_repository(model_uri) 
artifact_repository.download_artifacts("requirements.txt")

Does this PR require documentation update?

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.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

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/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/breaking-change - The PR will be mentioned in the "Breaking Changes" 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

github-actions · 2023-11-11T03:56:28Z

Documentation preview for 0e72186 will be available here when this CircleCI job completes successfully.

More info

Ignore this comment if this PR does not change the documentation.
It takes a few minutes for the preview to be available.
The preview is updated when a new commit is pushed to this PR.
This comment was created by https://github.com/mlflow/mlflow/actions/runs/6859271985.

harupy · 2023-11-11T19:29:09Z

tests/store/artifact/test_models_artifact_repo.py

    model_name = "MyModel"
    model_version = "12"

    with mock.patch.object(
        MlflowClient, "get_model_version_download_uri", return_value=artifact_location
-    ), mock.patch("mlflow.store.artifact.models_artifact_repo.write_yaml") as write_yaml_mock:
+    ), mock.patch("os.path.isdir", return_value=True), mock.patch(


Do we really need to patch os.path.isdir? Patching the built-in functions can lead to bugs. I think we can download a file instead of a directory.

refactored tests to use tmp_path location

harupy · 2023-11-11T19:30:32Z

mlflow/store/artifact/models_artifact_repo.py

+            _logger.warning(
+                "Registered Model Metadata file not able to be written. Destination path "
+                f"'{model_path}' is not a directory."
+            )


This warning might scare users or make them wonder what Registered Model Metadata file unnecessarily.

Agreed. Removed this. It's really only valid for deploying a root model directory and only if deploying to MLserver anyway. I'll add a note in the code about why the file might not be created.

harupy · 2023-11-11T19:32:06Z

tests/store/artifact/test_models_artifact_repo.py

+    ) as write_yaml_mock, mock.patch(
+        "os.path.isdir", return_value=False
+    ), mock.patch(


Same here: https://github.com/mlflow/mlflow/pull/10362/files#r1390275186

harupy · 2023-11-11T19:38:26Z

tests/store/artifact/test_models_artifact_repo.py

+    # Calling the download_artifacts method on local FileStore will create an ./mlruns directory
+    # which is a test side effect. Clean this up.
+    mlruns_dir = "./mlruns"
+    if os.path.exists(mlruns_dir):
+        shutil.rmtree(mlruns_dir)


Instead, can we set the artifact location to a temporary directory? If we hit an error before this block, ./mlruns won't be cleaned up.

yep! refactored

harupy · 2023-11-13T03:07:17Z

mlflow/store/artifact/models_artifact_repo.py

-            overwrite=True,
-            ensure_yaml_extension=False,
-        )
+        if os.path.isdir(model_path):


Is this file created when we download a subdirectory in the model directory?

- model - subdir <- this directory - MLmodel - ...

Could we update 'download_artifacts' function to only call this method when 'model_path' is a directory instead? Then we also don't need below warning as it doesn't make sense for non-directory cases.

@harupy good call. Adjusted the logic to only create this registered metadata file iff the model_path resolution after download contains the MLmodel file. Add a test to verify this behavior.

@serena-ruan great idea - updated!

BenWilson2 · 2023-11-13T17:36:35Z

mlflow/store/artifact/models_artifact_repo.py

@@ -169,8 +173,14 @@ def download_artifacts(self, artifact_path, dst_path=None):
        :return: Absolute path of the local filesystem location containing the desired artifacts.
        """

+        from mlflow.models.model import MLMODEL_FILE_NAME


Cannot do a non-local import due to circular reference. We really should migrate constants like these to a separate 'neutral zero-dependency' location, but that is out of scope for this PR.

harupy · 2023-11-13T22:54:27Z

tests/store/artifact/test_models_artifact_repo.py

+
+        models_repo.download_artifacts(artifact_location, str(tmp_path))
+
+        add_meta_mock.assert_called_once()


Instead, can we check that the metadata file exists?

harupy · 2023-11-13T22:55:33Z

tests/store/artifact/test_models_artifact_repo.py

+    dummy_dir = tmp_path / artifact_path
+    dummy_dir.mkdir()
+    dummy_file = dummy_dir / "dummy_file.txt"
+    dummy_file.write_text("dummy content")


Suggested change

dummy_file.write_text("dummy content")

dummy_file.touch()

We can use touch if the content is not important.

harupy

LGTM

harupy · 2023-11-14T02:21:04Z

tests/store/artifact/test_models_artifact_repo.py

+    model_dir = temp_remote_storage / "model_dir"
+    model_dir.mkdir(parents=True)
+    mlmodel_path = model_dir / "MLmodel"
+    mlmodel_path.write_text("dummy content")


Suggested change

mlmodel_path.write_text("dummy content")

mlmodel_path.touch()

we can use more touch.

harupy · 2023-11-14T02:21:15Z

tests/store/artifact/test_models_artifact_repo.py

-    artifact_dst_path = f"{dst_path}/{artifact_path}"
+
+    dummy_file = tmp_path / artifact_path
+    dummy_file.write_text("dummy content")


harupy · 2023-11-14T02:21:26Z

tests/store/artifact/test_models_artifact_repo.py

    """
    artifact_location = "s3://blah_bucket/"
+    dummy_file = tmp_path / "dummy_file.txt"
+    dummy_file.write_text("dummy content")


Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

…#10362) Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

…mlflow#10362) Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com> Signed-off-by: swathi <konakanchi.swathi@gmail.com>

BenWilson2 requested a review from smurching November 11, 2023 03:56

github-actions bot added area/artifacts Artifact stores and artifact logging rn/bug-fix Mention under Bug Fixes in Changelogs. labels Nov 11, 2023

BenWilson2 requested review from serena-ruan, harupy, B-Step62 and daniellok-db November 11, 2023 03:57

harupy reviewed Nov 11, 2023

View reviewed changes

harupy reviewed Nov 13, 2023

View reviewed changes

BenWilson2 commented Nov 13, 2023

View reviewed changes

BenWilson2 requested a review from harupy November 13, 2023 17:36

harupy reviewed Nov 13, 2023

View reviewed changes

harupy approved these changes Nov 14, 2023

View reviewed changes

harupy reviewed Nov 14, 2023

View reviewed changes

BenWilson2 added 5 commits November 13, 2023 23:31

fix non directory download

7f13bd0

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

adjustments

a823a39

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

disable local artifact tmp path access on Windows CI

ec949de

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

validate file creation explicitly

fc856b5

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

don't actually write to files

0e72186

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

BenWilson2 force-pushed the fix-metadata-write branch from fa68683 to 0e72186 Compare November 14, 2023 04:36

BenWilson2 merged commit d668245 into mlflow:master Nov 14, 2023
45 checks passed

BenWilson2 deleted the fix-metadata-write branch November 14, 2023 16:00

BenWilson2 added a commit that referenced this pull request Nov 14, 2023

Fix regression in downloading single files from models artifact store (…

a648634

…#10362) Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix regression in downloading single files from models artifact store #10362

Fix regression in downloading single files from models artifact store #10362

BenWilson2 commented Nov 11, 2023 •

edited by harupy

Loading

github-actions bot commented Nov 11, 2023 •

edited

Loading

harupy Nov 11, 2023 •

edited

Loading

BenWilson2 Nov 13, 2023

harupy Nov 11, 2023 •

edited

Loading

BenWilson2 Nov 13, 2023

harupy Nov 11, 2023

harupy Nov 11, 2023

BenWilson2 Nov 13, 2023

harupy Nov 13, 2023

serena-ruan Nov 13, 2023 •

edited

Loading

BenWilson2 Nov 13, 2023

BenWilson2 Nov 13, 2023

harupy Nov 13, 2023

harupy Nov 13, 2023 •

edited

Loading

harupy left a comment

harupy Nov 14, 2023

harupy Nov 14, 2023

harupy Nov 14, 2023


		models_repo.download_artifacts(artifact_location, str(tmp_path))

		add_meta_mock.assert_called_once()

	mlmodel_path.write_text("dummy content")
	mlmodel_path.touch()

Fix regression in downloading single files from models artifact store #10362

Fix regression in downloading single files from models artifact store #10362

Conversation

BenWilson2 commented Nov 11, 2023 • edited by harupy Loading

Install mlflow from this PR

Checkout with GitHub CLI

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

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

github-actions bot commented Nov 11, 2023 • edited Loading

harupy Nov 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

harupy Nov 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

serena-ruan Nov 13, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

harupy Nov 13, 2023 • edited Loading

Choose a reason for hiding this comment

harupy left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

BenWilson2 commented Nov 11, 2023 •

edited by harupy

Loading

github-actions bot commented Nov 11, 2023 •

edited

Loading

harupy Nov 11, 2023 •

edited

Loading

harupy Nov 11, 2023 •

edited

Loading

serena-ruan Nov 13, 2023 •

edited

Loading

harupy Nov 13, 2023 •

edited

Loading