[SC-16155] Fix async chart uploads by pre-serializing figure images

[juanmleng]

Pull Request Description

What and why?

Tests that generate charts or figures fail to upload results to ValidMind with:

RuntimeError: Timeout context manager should be used inside a task

and the customer confirmed that after the Kaleido 1.2.0 fix the process no longer hangs indefinitely but still fails to complete chart generation and upload the results.

This change pre-serializes figure PNGs before TestResult.log() enters the async upload flow, avoiding the Plotly/Kaleido and asyncio conflict in the customer-reported error path. It also centralizes figure PNG serialization behind a cached path so uploads, HTML rendering, and base64 generation all reuse the same image bytes consistently.

How to test

I was not able to reproduce this locally. I tested with Python 3.9 and also tried reproducing the issue outside ValidMind by serializing figures directly through the Figure object, but could not trigger the failure, so this may depend on the customer’s environment.

What needs special review?

Check whether it makes sense to take a defensive approach here, based on the customer-provided evidence, which includes clear stack traces showing the exact error path this fix addresses.

The fix is in principle low risk because it does not change behavior on systems where serialization already works; it only guarantees that to_image() is not called from within an async context where it may conflict.

Dependencies, breaking changes, and deployment notes

Release notes

Checklist

• What and why
• Screenshots or videos (Frontend)
• How to test
• What needs special review
• Dependencies, breaking changes, and deployment notes
• Labels applied
• PR linked to Shortcut
• Unit tests added (Backend)
• Tested locally
• Documentation updated (if required)
• Environment variable additions/changes documented (if required)

[johnwalz97]

lgtm... was also not able to reproduce but it seems like this could be a fix

[johnwalz97]

just fyi the bbox_inches param isn't maintained in the new implementation

[johnwalz97]

we might want the bbox_inches (see below) but not sure

[juanmleng]

Great catch @johnwalz97! I am inclined to not introduced too many more changes on this PR and perhaps leave it as-is for now. Previously bbox_inches="tight"only affected uploads, but with the new shared PNG path it would now affect all Matplotlib rendering, not just upload serialization. If we decide we want that back, we can always reintroduce it later?

[github-actions]

PR Summary

This PR introduces caching of PNG bytes for figures to ensure async-safe serialization, particularly addressing conflicts with Plotly's Kaleido library and asyncio event loops. A new private attribute (_cached_png_bytes) and helper method (_get_png_bytes) have been added to cache or generate PNG representations for figures, which can be of matplotlib, Plotly, or pre-rendered PNG image types.

Key changes include:

• A new method _get_png_bytes that either returns cached PNG bytes or generates them (using appropriate methods for matplotlib or Plotly figures).
• The introduction of pre_serialize(), which pre-computes and caches the PNG bytes before entering an async context. This avoids runtime conflicts when asynchronous methods are invoked.
• Updated implementations of to_html(), _get_b64_url(), and serialize_files() to make use of the cached PNG bytes. This unifies the image conversion process and reduces redundant code.
• In the result logging module, a pre-serialization call has been inserted for each figure to ensure that the async context is safe for operations that involve Plotly figures.

These changes enhance performance by avoiding repeated image generation and improve robustness in async environments by mitigating library conflicts.

Test Suggestions

• Test that calling pre_serialize() caches the PNG bytes and that subsequent calls use the cache.
• Verify that to_html() correctly encodes the cached or generated PNG bytes, including the inclusion of Plotly JSON metadata when appropriate.
• Test _get_b64_url() to ensure it returns a valid data URL with correct base64 encoding for different figure types.
• Validate that serialize_files() returns the correct file tuples for matplotlib, Plotly, and PNG figures, and includes a JSON file for Plotly figures.
• Confirm that an UnsupportedFigureError is raised when an unsupported figure type is passed.