Repository: mlflow/mlflow
Pull request: Bugfix/log figure correct file extension
<!-- 🚨 We recommend pull requests be filed from a non-master branch on a repository fork (e.g. <username>:fix-xxx). 🚨 -->

## [BUG] log_figure in plotly format without ".html" extension in artifact_file #6039

<!--
Please reference any related feature requests, issues, or PRs here. For example, `#123`. To automatically close the referenced items when this PR is merged, please use a closing keyword (close, fix, or resolve). For example, `Close #123`. See https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue for more information.
-->

Closes #6039 

## What changes are proposed in this pull request?

If the figure is in plotly format, the artifact_file string must end with one of the specified extensions. If HTML the write_html method is used. If png, jpeg, webp, SVG or pdf, the write_image is used.

## How is this patch tested?

<!--
If you're unsure about what to test, where to add tests, or how to run tests, please feel free to ask. We'd be happy to help.
-->

- [X] I have written tests (not required for typo or doc fix) and confirmed the proposed feature/bug-fix/change works.
tests/tracking/test_log_figure.py
test_log_figure_plotly_png,
test_log_figure_raises_error_for_unsupported_file_extension

<!--
Please describe how you confirmed the proposed feature/bug-fix/change works here. For example, if you fixed an MLflow client API, you could attach the code that didn't work prior to the fix but works now, or if you added a new feature on MLflow UI, you could attach a video that demonstrates the feature.
-->

## Does this PR change the documentation?

- [X] No. You can skip the rest of this section.
- [ ] Yes. Make sure the changed pages / sections render correctly by following the steps below.

1. Check the status of the `ci/circleci: build_doc` check. If it's successful, proceed to the
   next step, otherwise fix it.
2. Click `Details` on the right to open the job page of CircleCI.
3. Click the `Artifacts` tab.
4. Click `docs/build/html/index.html`.
5. Find the changed pages / sections and make sure they render correctly.

## Release Notes

### Is this a user-facing change?

- [X] 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/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

<!--
Insert an empty named anchor here to allow jumping to this section with a fragment URL
(e.g. https://github.com/mlflow/mlflow/pull/123#user-content-release-note-category).
Note that GitHub prefixes anchor names in markdown with "user-content-".
-->
<a name="release-note-category"></a>
### 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
- [X] `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


Problem statement:
[BUG] log_figure in plotly format without ".html" extension in artifact_file
### Willingness to contribute

Yes. I can contribute a fix for this bug independently.

### MLflow version

1.25.1

### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: MacOS
- **Python version**: 3.10

### Describe the problem

When logging a figure in plotly format, the file extension in the artifact_file needs to be ".html".
Otherwise, the file will be saved in HTML format anyway, but the file extension will be None or ".png" or other that the user specified in the path string. The file can't be opened on the MLFlow UI later.
The ".html" extension could be added to the end of the path if not there, or the method "write_image" could be used when the path finishes with PNG, JPEG, SVG or PDF.
Cheers,

### Tracking information

_No response_

### Code to reproduce issue

    import mlflow
    import plotly.express as px

    with mlflow.start_run():

        df = px.data.gapminder().query("country=='Canada'")
        fig = px.line(df, x="year", y="lifeExp", title='Life expectancy in Canada')
        mlflow.log_figure(fig, 'fig_test')

### Other info / logs

_No response_

### What component(s) does this bug affect?

- [ ] `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

### What interface(s) does this bug affect?

- [X] `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

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

Hints from issue discussion:
@drsantos89 Can we throw an error if the file name doesn't end with `.html`?
@harupy 
Yes, that will work and maybe is more clear than just adding `.html`.
If we also want to support the `write_image` method (for png, jpeg, SVG, pdf).  We can have an if based on the artifact_file extensions and throw an error is not supported.
@drsantos89 Would you be interested in creating a PR?
Yes, I will do it

Commit message: Apply suggestions from code review

Signed-off-by: Diogo Santos <drsantos989@gmail.com>

Patch to review:
diff --git a/mlflow/tracking/client.py b/mlflow/tracking/client.py
index f6f11b8802334..a199ee275f989 100644
--- a/mlflow/tracking/client.py
+++ b/mlflow/tracking/client.py
@@ -1157,7 +1157,13 @@ def _is_plotly_figure(fig):
             if "matplotlib" in sys.modules and _is_matplotlib_figure(figure):
                 figure.savefig(tmp_path)
             elif "plotly" in sys.modules and _is_plotly_figure(figure):
-                figure.write_html(tmp_path, include_plotlyjs="cdn", auto_open=False)
+                file_extension = os.path.splitext(artifact_file)[1]
+                if file_extension == '.html':
+                    figure.write_html(tmp_path, include_plotlyjs="cdn", auto_open=False)
+                elif file_extension in ['.png', '.jpeg', '.webp', '.svg', '.pdf']:
+                    figure.write_image(tmp_path)
+                else:
+                    raise TypeError(f"Unsupported file extension for plotly figure: '{file_extension}'")
             else:
                 raise TypeError("Unsupported figure object type: '{}'".format(type(figure)))
 
diff --git a/tests/tracking/test_log_figure.py b/tests/tracking/test_log_figure.py
index 270e2445c3ca1..a3ce7cc6c7bc3 100644
--- a/tests/tracking/test_log_figure.py
+++ b/tests/tracking/test_log_figure.py
@@ -28,7 +28,7 @@ def test_log_figure_matplotlib(subdir):
 
 
 @pytest.mark.parametrize("subdir", [None, ".", "dir", "dir1/dir2", "dir/.."])
-def test_log_figure_plotly(subdir):
+def test_log_figure_plotly_html(subdir):
     from plotly import graph_objects as go
 
     filename = "figure.html"
@@ -45,6 +45,37 @@ def test_log_figure_plotly(subdir):
         assert os.listdir(run_artifact_dir) == [filename]
 
 
+@pytest.mark.parametrize("extension", ['png', 'jpeg', 'webp', 'svg', 'pdf'])
+def test_log_figure_plotly_image(extension):
+    from plotly import graph_objects as go
+    
+    subdir = "."
+    filename = f"figure.{extension}"
+    artifact_file = posixpath.join(subdir, filename)
+
+    fig = go.Figure(go.Scatter(x=[0, 1], y=[2, 3]))
+
+    with mlflow.start_run():
+        mlflow.log_figure(fig, artifact_file)
+
+        artifact_path = None if subdir is None else posixpath.normpath(subdir)
+        artifact_uri = mlflow.get_artifact_uri(artifact_path)
+        run_artifact_dir = local_file_uri_to_path(artifact_uri)
+        assert os.listdir(run_artifact_dir) == [filename]     
+        
+
+def test_log_figure_raises_error_for_unsupported_file_extension():
+    from plotly import graph_objects as go
+
+    filename = "figure"
+    artifact_file = posixpath.join(".", filename)
+
+    fig = go.Figure(go.Scatter(x=[0, 1], y=[2, 3]))
+
+    with mlflow.start_run(), pytest.raises(TypeError, match=f"Unsupported file extension for plotly figure: ''"):
+        mlflow.log_figure(fig, artifact_file)
+
+
 def test_log_figure_raises_error_for_unsupported_figure_object_type():
     with mlflow.start_run(), pytest.raises(TypeError, match="Unsupported figure object type"):
         mlflow.log_figure("not_figure", "figure.png")


Code review comments:
diff_hunk: @@ -45,6 +45,37 @@ def test_log_figure_plotly(subdir):
         assert os.listdir(run_artifact_dir) == [filename]
 
 
+@pytest.mark.parametrize("extension", ['png', 'jpeg', 'webp', 'svg', 'pdf'])
+def test_log_figure_plotly_image(extension):
+    from plotly import graph_objects as go
+    
+    subdir = "."
+    filename = f"figure.{extension}"
+    artifact_file = posixpath.join(subdir, filename)
+
+    fig = go.Figure(go.Scatter(x=[0, 1], y=[2, 3]))
+
+    with mlflow.start_run():
+        mlflow.log_figure(fig, artifact_file)
+
+        artifact_path = None if subdir is None else posixpath.normpath(subdir)
+        artifact_uri = mlflow.get_artifact_uri(artifact_path)
+        run_artifact_dir = local_file_uri_to_path(artifact_uri)
+        assert os.listdir(run_artifact_dir) == [filename]     
+        
+
+def test_log_figure_raises_error_for_unsupported_file_extension():
path: tests/tracking/test_log_figure.py
text: @user2:
```suggestion
@user1.mark.parametrize("extension", ['', '.py'])
def test_log_figure_raises_error_for_unsupported_file_extension(extension):
```

Can we also parametrize this test?

diff_hunk: @@ -45,6 +45,37 @@ def test_log_figure_plotly(subdir):
         assert os.listdir(run_artifact_dir) == [filename]
 
 
+@pytest.mark.parametrize("extension", ['png', 'jpeg', 'webp', 'svg', 'pdf'])
+def test_log_figure_plotly_image(extension):
+    from plotly import graph_objects as go
+    
+    subdir = "."
+    filename = f"figure.{extension}"
+    artifact_file = posixpath.join(subdir, filename)
+
+    fig = go.Figure(go.Scatter(x=[0, 1], y=[2, 3]))
+
+    with mlflow.start_run():
+        mlflow.log_figure(fig, artifact_file)
+
+        artifact_path = None if subdir is None else posixpath.normpath(subdir)
+        artifact_uri = mlflow.get_artifact_uri(artifact_path)
+        run_artifact_dir = local_file_uri_to_path(artifact_uri)
+        assert os.listdir(run_artifact_dir) == [filename]     
+        
+
+def test_log_figure_raises_error_for_unsupported_file_extension():
+    from plotly import graph_objects as go
+
+    filename = "figure"
+    artifact_file = posixpath.join(".", filename)
+
+    fig = go.Figure(go.Scatter(x=[0, 1], y=[2, 3]))
+
+    with mlflow.start_run(), pytest.raises(TypeError, match=f"Unsupported file extension for plotly figure: ''"):
path: tests/tracking/test_log_figure.py
text: @user1:
```suggestion
    with mlflow.start_run(), pytest.raises(TypeError, match="Unsupported file extension for plotly figure: ''"):
```

diff_hunk: @@ -45,6 +45,37 @@ def test_log_figure_plotly(subdir):
         assert os.listdir(run_artifact_dir) == [filename]
 
 
+@pytest.mark.parametrize("extension", ['png', 'jpeg', 'webp', 'svg', 'pdf'])
+def test_log_figure_plotly_image(extension):
+    from plotly import graph_objects as go
+    
+    subdir = "."
+    filename = f"figure.{extension}"
+    artifact_file = posixpath.join(subdir, filename)
+
+    fig = go.Figure(go.Scatter(x=[0, 1], y=[2, 3]))
+
+    with mlflow.start_run():
+        mlflow.log_figure(fig, artifact_file)
+
+        artifact_path = None if subdir is None else posixpath.normpath(subdir)
+        artifact_uri = mlflow.get_artifact_uri(artifact_path)
+        run_artifact_dir = local_file_uri_to_path(artifact_uri)
+        assert os.listdir(run_artifact_dir) == [filename]     
+        
+
+def test_log_figure_raises_error_for_unsupported_file_extension():
+    from plotly import graph_objects as go
+
+    filename = "figure"
path: tests/tracking/test_log_figure.py
text: @user1:
```suggestion
    filename = f"figure{extension}"
```