Fall back to placeholder PNG when Plotly static export fails

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()

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);

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,14 +153,46 @@ 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:
             raise UnsupportedFigureError(
                 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.