Repository: mlflow/mlflow

Issue:
[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


Discussion:

Resolving patch:
diff --git a/mlflow/tracking/client.py b/mlflow/tracking/client.py
--- a/mlflow/tracking/client.py
+++ b/mlflow/tracking/client.py
@@ -1157,7 +1157,15 @@ 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)))
 


Test patch:
diff --git a/tests/tracking/test_log_figure.py b/tests/tracking/test_log_figure.py
--- 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,40 @@ 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]
+
+
+@pytest.mark.parametrize("extension", ["", ".py"])
+def test_log_figure_raises_error_for_unsupported_file_extension(extension):
+    from plotly import graph_objects as go
+
+    filename = f"figure{extension}"
+    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: '{extension}'"
+    ):
+        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")