From 69d8a57ee9d255075beb6bfc0183c65b6f5ebd27 Mon Sep 17 00:00:00 2001
From: Leo Borcherding <borchborchmail@gmail.com>
Date: Tue, 23 Jun 2026 08:22:20 -0500
Subject: [PATCH] Studio: lazy-import matplotlib so the server starts when the
 wheel is blocked (#6596)

* Studio: lazy-import matplotlib so the server starts when the wheel is blocked

matplotlib.pyplot was imported at the top of core/training/training.py, on the
server boot path. When matplotlib's native extension fails to load (e.g. an
unsigned wheel blocked by Windows Smart App Control), that import crashed the
whole Studio server at startup instead of just disabling loss plots.

Move it into a lazy _load_pyplot() helper called from _create_loss_plot, using
the headless Agg backend, and return None when matplotlib is unavailable so
plotting degrades gracefully. The plot return was already Optional, so callers
need no changes. Keep the type-only import under TYPE_CHECKING and quote the
annotations.

Fixes #6588

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: pin matplotlib==3.11.0

Pin matplotlib to the current latest so a new unsigned release does not
reintroduce the Smart App Control block on Windows. Belt-and-suspenders on
top of the lazy import. Pinned in both studio.txt and extras.txt.

* Pin matplotlib to 3.10.9 so Studio still installs on Python 3.10

matplotlib 3.11.0 requires Python >=3.11, so the pin had no installable wheel on
Python 3.10 (still supported) and pip install failed there. 3.10.9 is the latest
3.10.x (requires-python >=3.10) and covers Python 3.10 through 3.13. Also tighten
the lazy-import docstrings.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
---
 studio/backend/core/training/training.py | 41 +++++++++++++++++++++---
 studio/backend/requirements/extras.txt   |  2 +-
 studio/backend/requirements/studio.txt   |  2 +-
 3 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index 34bb4331a7..290b4090ec 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -24,9 +24,10 @@ from datetime import datetime, timezone
 from loggers import get_logger
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional, Tuple, Any
+from typing import Optional, Tuple, Any, TYPE_CHECKING
 
-import matplotlib.pyplot as plt
+if TYPE_CHECKING:
+    import matplotlib.pyplot as plt
 from utils.hardware import prepare_gpu_selection
 from utils.native_path_leases import (
     native_path_secret_removed_for_child_start,
@@ -36,6 +37,30 @@ from utils.paths import outputs_root
 
 logger = get_logger(__name__)
 
+_pyplot = None
+_pyplot_failed = False
+
+
+def _load_pyplot():
+    """Lazily import matplotlib.pyplot (headless Agg); return it, or None if
+    matplotlib is unavailable. Deferred so a blocked native wheel (e.g. Windows
+    Smart App Control) never breaks server startup, only loss plotting.
+    """
+    global _pyplot, _pyplot_failed
+    if _pyplot is not None or _pyplot_failed:
+        return _pyplot
+    try:
+        import matplotlib
+
+        matplotlib.use("Agg")  # headless backend
+        import matplotlib.pyplot as plt
+
+        _pyplot = plt
+    except Exception as e:
+        _pyplot_failed = True
+        logger.warning("matplotlib unavailable; loss plots disabled", error = str(e))
+    return _pyplot
+
 
 def _coerce_seed(value, default = 3407) -> int:
     """Normalize None / non-int to `default` (transformers.set_seed(None) raises)."""
@@ -655,7 +680,7 @@ class TrainingBackend:
         plot = self._create_loss_plot(progress, theme)
         return (plot, progress)
 
-    def refresh_plot_for_theme(self, theme: str) -> Optional[plt.Figure]:
+    def refresh_plot_for_theme(self, theme: str) -> "Optional[plt.Figure]":
         """Refresh plot with new theme."""
         if theme and isinstance(theme, str) and theme in ["light", "dark"]:
             self.current_theme = theme
@@ -1090,8 +1115,14 @@ class TrainingBackend:
         self,
         progress: TrainingProgress,
         theme: str = "light",
-    ) -> plt.Figure:
-        """Create training loss plot with theme-aware styling."""
+    ) -> "Optional[plt.Figure]":
+        """Create training loss plot with theme-aware styling.
+
+        matplotlib is loaded lazily; returns None if it is unavailable.
+        """
+        plt = _load_pyplot()
+        if plt is None:
+            return None
         plt.close("all")
 
         LIGHT_STYLE = {
diff --git a/studio/backend/requirements/extras.txt b/studio/backend/requirements/extras.txt
index 40737b0876..b56c2dfa40 100644
--- a/studio/backend/requirements/extras.txt
+++ b/studio/backend/requirements/extras.txt
@@ -39,7 +39,7 @@ ftfy
 importlib-resources
 librosa
 markdown2
-matplotlib
+matplotlib==3.10.9
 pystoi
 soundfile
 tensorboard
diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt
index 96fef60471..9efa10faef 100644
--- a/studio/backend/requirements/studio.txt
+++ b/studio/backend/requirements/studio.txt
@@ -4,7 +4,7 @@ fastapi
 uvicorn
 pydantic
 packaging
-matplotlib
+matplotlib==3.10.9
 pandas
 nest_asyncio
 datasets==4.3.0