diff --git a/tests/test_figure_static_export.py b/tests/test_figure_static_export.py new file mode 100644 index 000000000..2bd68b80e --- /dev/null +++ b/tests/test_figure_static_export.py @@ -0,0 +1,61 @@ +import os +import unittest +from unittest import mock + +import plotly.graph_objects as go + +from validmind.vm_models import figure as figure_module +from validmind.vm_models.figure import Figure + + +def _make_figure(): + return Figure(key="k", figure=go.Figure(go.Scatter(x=[1], y=[1])), ref_id="r") + + +class TestPlotlyStaticExportFallback(unittest.TestCase): + def setUp(self): + figure_module._static_export_unavailable = False + self.addCleanup(setattr, figure_module, "_static_export_unavailable", False) + + def test_falls_back_to_placeholder_and_stops_retrying(self): + fig = _make_figure() + with mock.patch.object( + go.FigureWidget, + "to_image", + side_effect=RuntimeError("Kaleido requires Google Chrome"), + ) as to_image: + first = fig._get_png_bytes() + second = _make_figure()._get_png_bytes() + + self.assertTrue(first.startswith(b"\x89PNG")) + self.assertEqual(first, second) + self.assertEqual(to_image.call_count, 1) + self.assertFalse(fig.has_static_image) + + def test_still_uploads_plotly_json(self): + with mock.patch.object( + go.FigureWidget, "to_image", side_effect=RuntimeError("no chrome") + ): + files = _make_figure().serialize_files() + self.assertEqual(set(files), {"image", "json_image"}) + + def test_env_var_skips_kaleido(self): + with mock.patch.dict(os.environ, {"VALIDMIND_PLOTLY_STATIC_EXPORT": "false"}): + with mock.patch.object(go.FigureWidget, "to_image") as to_image: + fig = _make_figure() + png = fig._get_png_bytes() + self.assertFalse(fig.has_static_image) + to_image.assert_not_called() + self.assertTrue(png.startswith(b"\x89PNG")) + + def test_working_export_is_used(self): + with mock.patch.object( + go.FigureWidget, "to_image", return_value=b"real-png" + ): + fig = _make_figure() + self.assertEqual(fig._get_png_bytes(), b"real-png") + self.assertTrue(fig.has_static_image) + + +if __name__ == "__main__": + unittest.main() diff --git a/validmind/ai/test_descriptions.py b/validmind/ai/test_descriptions.py index e7e62a4e7..6eee52a31 100644 --- a/validmind/ai/test_descriptions.py +++ b/validmind/ai/test_descriptions.py @@ -168,7 +168,11 @@ def generate_description( "test_description": test_description, "title": title, "summary": _truncate_summary(summary, test_id), - "figures": [figure._get_b64_url() for figure in figures or []], + "figures": [ + figure._get_b64_url() + for figure in figures or [] + if figure.has_static_image + ], "additional_context": additional_context, "instructions": instructions, # params can hold anything the test accepts (DataFrames, arrays); diff --git a/validmind/vm_models/figure.py b/validmind/vm_models/figure.py index 97c41b032..4f2e740ac 100644 --- a/validmind/vm_models/figure.py +++ b/validmind/vm_models/figure.py @@ -9,6 +9,7 @@ import base64 import json import os +import threading from dataclasses import dataclass from io import BytesIO from typing import Optional, Union @@ -18,9 +19,54 @@ from ..client_config import client_config from ..errors import UnsupportedFigureError +from ..logging import get_logger from ..utils import get_full_typename from .html_renderer import StatefulHTMLRenderer +logger = get_logger(__name__) + +# Set once Plotly's static image export (kaleido + Chrome) fails, so later +# figures skip it instead of retrying (and hanging) on every figure. +_static_export_unavailable = False +_placeholder_png: Optional[bytes] = None +_export_lock = threading.Lock() + + +def _static_export_disabled() -> bool: + """Static PNG export of Plotly figures can be turned off with + VALIDMIND_PLOTLY_STATIC_EXPORT=false (e.g. on Databricks serverless, + where Chrome can't run).""" + return os.getenv("VALIDMIND_PLOTLY_STATIC_EXPORT", "true").lower() in ( + "false", + "0", + "no", + ) + + +def _get_placeholder_png() -> bytes: + """A PNG standing in for a Plotly figure that couldn't be rendered. The + figure's Plotly JSON is still uploaded, so the platform can render it.""" + global _placeholder_png + with _export_lock: + if _placeholder_png is not None: + return _placeholder_png + + from matplotlib.figure import Figure as MplFigure + + fig = MplFigure(figsize=(6, 3)) + fig.text( + 0.5, + 0.5, + "Static preview unavailable.\nThis is an interactive Plotly figure.", + ha="center", + va="center", + color="#555555", + ) + buffer = BytesIO() + fig.savefig(buffer, format="png") + _placeholder_png = buffer.getvalue() + return _placeholder_png + def is_matplotlib_figure(figure) -> bool: return isinstance(figure, matplotlib.figure.Figure) @@ -107,7 +153,7 @@ def _get_png_bytes(self) -> bytes: buffer.seek(0) return buffer.read() elif is_plotly_figure(self.figure): - return self.figure.to_image(format="png") + return self._get_plotly_png_bytes() elif is_png_image(self.figure): return self.figure else: @@ -115,6 +161,38 @@ def _get_png_bytes(self) -> bytes: f"Unrecognized figure type: {get_full_typename(self.figure)}" ) + @property + def has_static_image(self) -> bool: + """False for a Plotly figure that can't be rendered to a real PNG.""" + if not is_plotly_figure(self.figure): + return True + return not (_static_export_unavailable or _static_export_disabled()) + + def _get_plotly_png_bytes(self) -> bytes: + """Render a Plotly figure to PNG, falling back to a placeholder when + kaleido/Chrome isn't available so the run doesn't fail. The figure's + JSON is uploaded either way.""" + global _static_export_unavailable + + if _static_export_unavailable or _static_export_disabled(): + return _get_placeholder_png() + + try: + return self.figure.to_image(format="png") + except Exception as e: + with _export_lock: + first_failure = not _static_export_unavailable + _static_export_unavailable = True + if first_failure: + logger.warning( + "Could not render Plotly figures to PNG (kaleido needs a " + "working Chrome install): %s. Figures will be sent as Plotly " + "JSON with a placeholder image. Set " + "VALIDMIND_PLOTLY_STATIC_EXPORT=false to skip this attempt.", + e, + ) + return _get_placeholder_png() + def pre_serialize(self): """Pre-serialize the figure to PNG bytes for async-safe upload.