From aa6c37d880d6f3e94714487ca75424b84a35e8d4 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Wed, 4 Nov 2026 10:46:12 -0500 Subject: [PATCH 1/2] Fall back to a placeholder PNG when Plotly static export fails Kaleido 1.x needs a working Chrome, which isn't available on Databricks serverless. Instead of failing the run, upload the Plotly JSON with a placeholder image, stop retrying after the first failure, and add VALIDMIND_PLOTLY_STATIC_EXPORT=false to skip the attempt. --- tests/test_figure_static_export.py | 61 +++++++++++++++++++++++++ validmind/ai/test_descriptions.py | 6 ++- validmind/vm_models/figure.py | 72 +++++++++++++++++++++++++++++- 3 files changed, 137 insertions(+), 2 deletions(-) create mode 100644 tests/test_figure_static_export.py 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..2bec1d720 100644 --- a/validmind/vm_models/figure.py +++ b/validmind/vm_models/figure.py @@ -18,9 +18,50 @@ 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 + + +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 + if _placeholder_png is None: + 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 +148,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 +156,35 @@ 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: + _static_export_unavailable = True + 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. From 519de90aff9fdc34de17697381cf40614d752c38 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Wed, 4 Nov 2026 10:48:14 -0500 Subject: [PATCH 2/2] Guard Plotly static export state with a lock --- validmind/vm_models/figure.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/validmind/vm_models/figure.py b/validmind/vm_models/figure.py index 2bec1d720..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 @@ -28,6 +29,7 @@ # 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: @@ -45,7 +47,10 @@ 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 - if _placeholder_png is None: + with _export_lock: + if _placeholder_png is not None: + return _placeholder_png + from matplotlib.figure import Figure as MplFigure fig = MplFigure(figsize=(6, 3)) @@ -60,7 +65,7 @@ def _get_placeholder_png() -> bytes: buffer = BytesIO() fig.savefig(buffer, format="png") _placeholder_png = buffer.getvalue() - return _placeholder_png + return _placeholder_png def is_matplotlib_figure(figure) -> bool: @@ -175,14 +180,17 @@ def _get_plotly_png_bytes(self) -> bytes: try: return self.figure.to_image(format="png") except Exception as e: - _static_export_unavailable = True - 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, - ) + 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):