Repository: mlflow/mlflow

Issue:
[FR] Allow Poetry MathJax Specification
### Willingness to contribute

Yes. I can contribute this feature independently.

### Proposal Summary

While logging a Plotly figure with MlFlow using the `mlflow.log_figure(...)` method works well in principle, it is not possible to specify further arguments available by the `.write_html(...)` method (see [documentation](https://plotly.com/python/interactive-html-export/#full-parameter-documentation) in Plotly.
This however is necessary if the figure includes Latex/MathJax formulas as those are not being rendered correctly in the resulting html artifact (see [this Stackoverflow question](https://stackoverflow.com/questions/71735053/save-plotly-figure-interactively-html-whilst-preserving-latex-font)).
Therefore I propose to add either `**kwargs` or the more specific `include_mathjax` argument to `mlflow.log_figure(...)` that then can be passed to the underlying `.write_html(...)` method.

### Motivation

> #### What is the use case for this feature?
Logging interactive (html) Plotly figures that include Latex/MathJax formulas.

> #### Why is this use case valuable to support for MLflow users in general?
Otherwise, formulas in interactive figures are not being rendered correctly.

> #### Why is this use case valuable to support for your project(s) or organization?
See above.

> #### Why is it currently difficult to achieve this use case?
To my best knowledge, the only possible workaround is:
- writing a json based plotly figure
- loading this figure in a separate python script
- injecting the above mentioned `include_mathjax` argument
- saving the figure using the Plotly API directly
This is way more difficult compared to directly specifying the required argument in `mlflow.log_figure(...)

### Details

_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/gateway`: AI Gateway service, Gateway client APIs, third-party Gateway integrations
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe 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
- [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:
@stroblme Thanks for the proposal! It seems reasonable to me :) Adding `**kwargs` sounds more appealing because `log_figure` also supports matplotlib.

Resolving patch:
diff --git a/mlflow/tracking/client.py b/mlflow/tracking/client.py
--- a/mlflow/tracking/client.py
+++ b/mlflow/tracking/client.py
@@ -1234,6 +1234,8 @@ def log_figure(
         run_id: str,
         figure: Union["matplotlib.figure.Figure", "plotly.graph_objects.Figure"],
         artifact_file: str,
+        *,
+        save_kwargs: Optional[Dict[str, Any]] = None,
     ) -> None:
         """
         Log a figure as an artifact. The following figure objects are supported:
@@ -1251,6 +1253,7 @@ def log_figure(
         :param figure: Figure to log.
         :param artifact_file: The run-relative artifact file path in posixpath format to which
                               the figure is saved (e.g. "dir/file.png").
+        :param save_kwargs: Additional keyword arguments passed to the method that saves the figure.
 
         .. code-block:: python
             :caption: Matplotlib Example
@@ -1286,18 +1289,21 @@ def _is_plotly_figure(fig):
 
             return isinstance(fig, plotly.graph_objects.Figure)
 
+        save_kwargs = save_kwargs or {}
         with self._log_artifact_helper(run_id, artifact_file) as tmp_path:
             # `is_matplotlib_figure` is executed only when `matplotlib` is found in `sys.modules`.
             # This allows logging a `plotly` figure in an environment where `matplotlib` is not
             # installed.
             if "matplotlib" in sys.modules and _is_matplotlib_figure(figure):
-                figure.savefig(tmp_path)
+                figure.savefig(tmp_path, **save_kwargs)
             elif "plotly" in sys.modules and _is_plotly_figure(figure):
                 file_extension = os.path.splitext(artifact_file)[1]
                 if file_extension == ".html":
-                    figure.write_html(tmp_path, include_plotlyjs="cdn", auto_open=False)
+                    save_kwargs.setdefault("include_plotlyjs", "cdn")
+                    save_kwargs.setdefault("auto_open", False)
+                    figure.write_html(tmp_path, **save_kwargs)
                 elif file_extension in [".png", ".jpeg", ".webp", ".svg", ".pdf"]:
-                    figure.write_image(tmp_path)
+                    figure.write_image(tmp_path, **save_kwargs)
                 else:
                     raise TypeError(
                         f"Unsupported file extension for plotly figure: '{file_extension}'"
diff --git a/mlflow/tracking/fluent.py b/mlflow/tracking/fluent.py
--- a/mlflow/tracking/fluent.py
+++ b/mlflow/tracking/fluent.py
@@ -973,7 +973,10 @@ def log_dict(dictionary: Dict[str, Any], artifact_file: str) -> None:
 
 
 def log_figure(
-    figure: Union["matplotlib.figure.Figure", "plotly.graph_objects.Figure"], artifact_file: str
+    figure: Union["matplotlib.figure.Figure", "plotly.graph_objects.Figure"],
+    artifact_file: str,
+    *,
+    save_kwargs: Optional[Dict[str, Any]] = None,
 ) -> None:
     """
     Log a figure as an artifact. The following figure objects are supported:
@@ -990,6 +993,7 @@ def log_figure(
     :param figure: Figure to log.
     :param artifact_file: The run-relative artifact file path in posixpath format to which
                           the figure is saved (e.g. "dir/file.png").
+    :param save_kwargs: Additional keyword arguments passed to the method that saves the figure.
 
     .. test-code-block:: python
         :caption: Matplotlib Example
@@ -1015,7 +1019,7 @@ def log_figure(
             mlflow.log_figure(fig, "figure.html")
     """
     run_id = _get_or_start_run().info.run_id
-    MlflowClient().log_figure(run_id, figure, artifact_file)
+    MlflowClient().log_figure(run_id, figure, artifact_file, save_kwargs=save_kwargs)
 
 
 def log_image(image: Union["numpy.ndarray", "PIL.Image.Image"], artifact_file: str) -> None:


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
@@ -1,5 +1,6 @@
 import os
 import posixpath
+import uuid
 
 import pytest
 
@@ -66,6 +67,20 @@ def test_log_figure_plotly_image(extension):
         assert os.listdir(run_artifact_dir) == [filename]
 
 
+def test_log_figure_save_kwargs():
+    from plotly import graph_objects as go
+
+    fig = go.Figure(go.Scatter(x=[0, 1], y=[2, 3]))
+    with mlflow.start_run():
+        name = "figure.html"
+        div_id = uuid.uuid4().hex
+        mlflow.log_figure(fig, name, save_kwargs={"div_id": div_id})
+        artifact_uri = mlflow.get_artifact_uri(name)
+        local_path = local_file_uri_to_path(artifact_uri)
+        with open(local_path) as f:
+            assert div_id in f.read()
+
+
 @pytest.mark.parametrize("extension", ["", ".py"])
 def test_log_figure_raises_error_for_unsupported_file_extension(extension):
     from plotly import graph_objects as go