Studio: self-heal unsloth namespace shadows; clearer failed-load messages (#6532)

* Studio: self-heal unsloth namespace-package shadows in all subprocess workers

A directory named `unsloth` (or `unsloth_zoo`) without an __init__.py on
PYTHONPATH/sys.path, a stray source checkout or a polluted PYTHONPATH, makes
`import unsloth` resolve to an empty namespace package, so a worker's
`from unsloth import FastLanguageModel` dies with a cryptic
"cannot import name ... (unknown location)".

The LLM training path already recovered from this via `_ensure_real_packages`
in trainer.py (PR #6269), but the inference, export, and embedding-training
subprocesses imported Unsloth directly with no guard. Extract that helper into
a shared, dependency-free core/import_guards.py and call it before the Unsloth
import in every subprocess: it drops the offending sys.path entries, imports
the real packages (unsloth before unsloth_zoo so the pre-zoo GPU fixes run),
then restores sys.path. trainer.py now imports the shared helper instead of its
local copy.

Covers both unsloth and unsloth_zoo and both namespace origin forms (None and
"namespace"). The existing PR #6269 test now exercises the shared helper.

* Studio: distinguish a failed model load from no model in the attach gates

A failed load never sets the checkpoint, so the image and audio attach gates
fell through to "Load a model before adding images/audio", which reads as if
the user simply forgot to pick a model rather than that the load errored. Add a
dedicated lastModelLoadError to the chat runtime store, set only when an actual
load attempt fails (not on refresh, list, status, or unload errors, which keep
using modelsError) and cleared when the next load starts. The image gate (all
three call sites) and the audio gate now use it to report a failed load and
point at the server logs, while still blocking in exactly the same cases.

* Tighten namespace-shadow guard and load-error comments

studio/backend/core/export/worker.py

@@ -506,6 +506,11 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None
         if backend_path not in sys.path:
             sys.path.insert(0, backend_path)
 
+        # Recover from any namespace-package shadow before importing Unsloth.
+        from core.import_guards import ensure_real_packages
+
+        ensure_real_packages("unsloth_zoo", "unsloth")
+
         from core.export.export import ExportBackend
 
         import transformers

studio/backend/core/import_guards.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Recover `unsloth`/`unsloth_zoo` from a namespace-package shadow. Stdlib-only."""
+
+from __future__ import annotations
+
+import os
+import sys
+
+
+def ensure_real_packages(*names: str) -> None:
+    """Drop sys.path entries where a bare `<name>/` dir (no __init__.py) shadows
+    the installed package as a namespace, import the real packages, restore
+    sys.path. No-op without a shadow. Pass dependency-first (e.g. "unsloth_zoo",
+    "unsloth"); imports run dependency-last."""
+    import importlib
+    import importlib.util
+
+    bad: set = set()
+    shadowed: list = []
+    for name in names:
+        try:
+            spec = importlib.util.find_spec(name)
+        except (ImportError, ValueError, AttributeError):
+            spec = None
+        # real package -> spec.origin is its __init__; namespace shadow -> None/"namespace"
+        if spec is None or spec.origin not in (None, "namespace"):
+            continue
+        dirs = {os.path.realpath(d) for d in (spec.submodule_search_locations or [])}
+        if not dirs:
+            continue
+        shadowed.append(name)
+        for entry in sys.path:
+            pkg = os.path.join(entry or os.getcwd(), name)
+            if os.path.realpath(pkg) in dirs and not os.path.isfile(
+                os.path.join(pkg, "__init__.py")
+            ):
+                bad.add(entry)
+    if not bad:
+        return
+    saved = list(sys.path)
+    sys.path[:] = [e for e in sys.path if e not in bad]
+    for name in shadowed:
+        for cached in [m for m in list(sys.modules) if m == name or m.startswith(name + ".")]:
+            del sys.modules[cached]
+    try:
+        importlib.invalidate_caches()
+        # import unsloth before unsloth_zoo: unsloth.__init__ runs GPU/bnb fixes zoo relies on
+        for name in reversed(names):
+            importlib.import_module(name)
+    finally:
+        sys.path[:] = saved

studio/backend/core/inference/worker.py

@@ -716,6 +716,11 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
 
         _ensure_backend_on_path()
 
+        # Recover from any namespace-package shadow before importing Unsloth.
+        from core.import_guards import ensure_real_packages
+
+        ensure_real_packages("unsloth_zoo", "unsloth")
+
         from core.inference.inference import InferenceBackend
 
         import transformers

studio/backend/core/training/trainer.py

@@ -46,58 +46,8 @@ if hasattr(torch._dynamo.config, "recompile_limit"):
     torch._dynamo.config.recompile_limit = 64
 
 
-def _ensure_real_packages(*names: str) -> None:
-    """Stop `import <name>` from binding to a namespace-package shadow.
-
-    A directory named like the package but missing __init__.py on sys.path (a
-    stray checkout, a partial clone, or a polluted PYTHONPATH) makes the path
-    finder return a namespace package, so `from unsloth import FastLanguageModel`
-    dies with "cannot import name ... (unknown location)". A normal
-    site-packages install always wins, so only source/editable installs are
-    exposed. Drop the offending entries, import the real packages, then restore
-    sys.path so other modules on those entries keep importing.
-    """
-    import importlib
-    import importlib.util
-
-    bad: set = set()
-    shadowed: list = []
-    for name in names:
-        try:
-            spec = importlib.util.find_spec(name)
-        except (ImportError, ValueError, AttributeError):
-            spec = None
-        # a real package exposes its __init__ via spec.origin; a namespace
-        # shadow has origin None/"namespace" and only search locations
-        if spec is None or spec.origin not in (None, "namespace"):
-            continue
-        dirs = {os.path.realpath(d) for d in (spec.submodule_search_locations or [])}
-        if not dirs:
-            continue
-        shadowed.append(name)
-        for entry in sys.path:
-            pkg = os.path.join(entry or os.getcwd(), name)
-            if os.path.realpath(pkg) in dirs and not os.path.isfile(
-                os.path.join(pkg, "__init__.py")
-            ):
-                bad.add(entry)
-    if not bad:
-        return
-    saved = list(sys.path)
-    sys.path[:] = [e for e in sys.path if e not in bad]
-    for name in shadowed:
-        for cached in [m for m in list(sys.modules) if m == name or m.startswith(name + ".")]:
-            del sys.modules[cached]
-    try:
-        importlib.invalidate_caches()
-        # Import unsloth before unsloth_zoo (names are dependency-first):
-        # unsloth.__init__ runs ROCm/Windows bnb fixes before it imports zoo,
-        # so importing zoo first here would skip them. Repeat import is a no-op.
-        for name in reversed(names):
-            importlib.import_module(name)
-    finally:
-        sys.path[:] = saved
-
+# Drop any unsloth/unsloth_zoo namespace-package shadow before importing them.
+from core.import_guards import ensure_real_packages as _ensure_real_packages
 
 _ensure_real_packages("unsloth_zoo", "unsloth")
 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported

studio/backend/core/training/worker.py

@@ -3156,6 +3156,10 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
     # ── 1. Import embedding-specific libraries ──
     _send_status(event_queue, "Importing embedding libraries...")
     try:
+        # Recover from a namespace-package shadow (embedding imports unsloth directly).
+        from core.import_guards import ensure_real_packages
+
+        ensure_real_packages("unsloth_zoo", "unsloth")
         from unsloth import FastSentenceTransformer, is_bfloat16_supported
         from sentence_transformers import (
             SentenceTransformerTrainer,

studio/backend/tests/test_namespace_shadow_guard_pr6269.py

@@ -1,15 +1,18 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Verification tests for PR #6269 (training-worker namespace-shadow guard).
+"""Verification tests for PR #6269 (namespace-shadow guard).
 
-`_ensure_real_packages` (core/training/trainer.py) drops namespace-package
+`ensure_real_packages` (core/import_guards.py) drops namespace-package
 shadow dirs (a `unsloth`/`unsloth_zoo` dir with no __init__.py on sys.path)
 before `from unsloth import ...`. Order matters: `unsloth.__init__` runs its
 ROCm/Windows bnb fixes before importing unsloth_zoo, so the guard must import
 unsloth first. Each test runs the real guard (ast-extracted from source, no
 GPU/torch) in a subprocess, with fake packages reachable only via a meta path
 finder to mimic an editable/PEP 660 install where the shadow wins.
+
+Originally defined in core/training/trainer.py; extracted to the shared
+core/import_guards.py so the inference, export and embedding workers reuse it.
 """
 
 import json
@@ -21,7 +24,7 @@ from pathlib import Path
 
 import pytest
 
-TRAINER_PY = Path(__file__).resolve().parents[1] / "core" / "training" / "trainer.py"
+GUARD_PY = Path(__file__).resolve().parents[1] / "core" / "import_guards.py"
 
 
 # ── fake package bodies ──────────────────────────────────────────────
@@ -62,17 +65,17 @@ _DRIVER = textwrap.dedent(
 
     cfg = json.load(open(sys.argv[1]))
 
-    # Extract the real _ensure_real_packages from trainer.py source without
-    # importing the heavy module or its `from unsloth import ...` line.
-    src = open(cfg["trainer_py"]).read()
+    # Extract the real ensure_real_packages from import_guards.py source
+    # without importing the heavy module or its `from unsloth import ...` line.
+    src = open(cfg["guard_py"]).read()
     tree = ast.parse(src)
     fn = next(n for n in tree.body
-              if isinstance(n, ast.FunctionDef) and n.name == "_ensure_real_packages")
+              if isinstance(n, ast.FunctionDef) and n.name == "ensure_real_packages")
     mod = ast.Module(body=[fn], type_ignores=[])
     ast.fix_missing_locations(mod)
     ns = {"os": os, "sys": sys}
-    exec(compile(mod, cfg["trainer_py"], "exec"), ns)
-    _ensure_real_packages = ns["_ensure_real_packages"]
+    exec(compile(mod, cfg["guard_py"], "exec"), ns)
+    _ensure_real_packages = ns["ensure_real_packages"]
 
     # Under -S site-packages is off, so a shadow root placed first wins the
     # path finder; the real packages come only from the meta finder below.
@@ -148,7 +151,7 @@ def _run(
     shadow_roots,
     real: bool,
     names = ("unsloth_zoo", "unsloth"),
-    trainer_py: Path = TRAINER_PY,
+    guard_py: Path = GUARD_PY,
     raise_on_invalidate: bool = False,
 ):
     order_file = tmp_path / "order.txt"
@@ -159,7 +162,7 @@ def _run(
         _make_real_pkg(real_root)
 
     cfg = {
-        "trainer_py": str(trainer_py),
+        "guard_py": str(guard_py),
         "shadow_roots": [str(r) for r in shadow_roots],
         "real_root": str(real_root) if real else None,
         "names": list(names),

studio/frontend/src/features/chat/api/chat-adapter.ts

@@ -1926,6 +1926,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
           externalModelLabel: externalSelection?.modelId ?? null,
           loadedIsMultimodal: runtime.loadedIsMultimodal,
           modelLoaded: !!params.checkpoint && !runtime.modelLoading,
+          loadError: runtime.lastModelLoadError,
         });
         if (imageGateReason) {
           toast.error(imageGateReason);

studio/frontend/src/features/chat/audio-attachment-adapter.ts

@@ -39,7 +39,10 @@ export class AudioAttachmentAdapter implements AttachmentAdapter {
     const modelLoaded = !!checkpoint && !state.modelLoading;
     let unavailableReason: string | null = null;
     if (!modelLoaded) {
-      unavailableReason = "Load a model before adding audio files.";
+      // Mirror the image gate: flag a failed load vs "no model picked".
+      unavailableReason = state.lastModelLoadError
+        ? "The last model failed to load. Check the server logs, then load a model before adding audio files."
+        : "Load a model before adding audio files.";
     } else if (!activeModel?.hasAudioInput) {
       const label = activeModel?.name || checkpoint || "Current model";
       unavailableReason = `${label} cannot accept audio. Load an audio-input model before attaching audio files.`;

studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts

@@ -246,6 +246,9 @@ export function useChatModelRuntime() {
   const setLoras = useChatRuntimeStore((state) => state.setLoras);
   const setParams = useChatRuntimeStore((state) => state.setParams);
   const setModelsError = useChatRuntimeStore((state) => state.setModelsError);
+  const setLastModelLoadError = useChatRuntimeStore(
+    (state) => state.setLastModelLoadError,
+  );
   const setCheckpoint = useChatRuntimeStore((state) => state.setCheckpoint);
   const clearCheckpoint = useChatRuntimeStore((state) => state.clearCheckpoint);
 
@@ -493,6 +496,7 @@ export function useChatModelRuntime() {
         .filter(Boolean)
         .join(" ");
       setModelsError(null);
+      setLastModelLoadError(null); // clear prior failed-load marker
       setLoadToastDismissedState(false);
       const loadInfo = {
         id: modelId,
@@ -1184,6 +1188,7 @@ export function useChatModelRuntime() {
         const message =
           error instanceof Error ? error.message : "Failed to load model";
         setModelsError(message);
+        setLastModelLoadError(message); // load-specific failure for the attach gates
         if (throwOnError) {
           throw error instanceof Error ? error : new Error(message);
         }
@@ -1199,6 +1204,7 @@ export function useChatModelRuntime() {
       resetLoadingUi,
       setLoadToastDismissedState,
       setModelsError,
+      setLastModelLoadError,
       setParams,
     ],
   );

studio/frontend/src/features/chat/runtime-provider.tsx

@@ -114,6 +114,7 @@ class VisionImageAdapter implements AttachmentAdapter {
       externalModelLabel,
       loadedIsMultimodal: state.loadedIsMultimodal,
       modelLoaded,
+      loadError: state.lastModelLoadError,
     });
     if (unavailableReason) {
       toast.error(unavailableReason);

studio/frontend/src/features/chat/shared-composer.tsx

@@ -485,6 +485,7 @@ export function SharedComposer({
   const modelLoaded = useChatRuntimeStore(
     (s) => !!s.params.checkpoint && !s.modelLoading,
   );
+  const lastModelLoadError = useChatRuntimeStore((s) => s.lastModelLoadError);
   const loadedIsMultimodal = useChatRuntimeStore((s) => s.loadedIsMultimodal);
   const supportsReasoning = useChatRuntimeStore((s) => s.supportsReasoning);
   const reasoningAlwaysOn = useChatRuntimeStore((s) => s.reasoningAlwaysOn);
@@ -566,6 +567,7 @@ export function SharedComposer({
     externalModelLabel: externalSelection?.modelId ?? null,
     loadedIsMultimodal,
     modelLoaded,
+    loadError: lastModelLoadError,
   });
   const isCompareMode = Boolean(model1?.id || model2?.id);
   // Attach-time gate. Compare mode defers to send: the catalog can lag a

studio/frontend/src/features/chat/stores/chat-runtime-store.ts

@@ -482,6 +482,9 @@ type ChatRuntimeStore = {
   autoTitle: boolean;
   hfToken: string;
   modelsError: string | null;
+  // Set only when a LOAD fails (not refresh/list/unload, which use modelsError);
+  // lets the attach gates flag a failed load vs "no model picked".
+  lastModelLoadError: string | null;
   activeGgufVariant: string | null;
   ggufContextLength: number | null;
   ggufMaxContextLength: number | null;
@@ -660,6 +663,7 @@ type ChatRuntimeStore = {
   setAutoTitle: (enabled: boolean) => void;
   setHfToken: (token: string) => void;
   setModelsError: (error: string | null) => void;
+  setLastModelLoadError: (error: string | null) => void;
   setCheckpoint: (modelId: string, ggufVariant?: string | null) => void;
   setActiveThreadId: (threadId: string | null) => void;
   setActiveProjectId: (projectId: string | null) => void;
@@ -964,6 +968,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   autoTitle: false,
   hfToken: loadString(HF_TOKEN_KEY, ""),
   modelsError: null,
+  lastModelLoadError: null,
   activeGgufVariant: null,
   ggufContextLength: null,
   ggufMaxContextLength: null,
@@ -1159,6 +1164,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
     notifyHfTokenChanged(hfToken);
   },
   setModelsError: (modelsError) => set({ modelsError }),
+  setLastModelLoadError: (lastModelLoadError) => set({ lastModelLoadError }),
   setCheckpoint: (modelId, ggufVariant) =>
     set((state) => {
       // Persist external selections so they survive a refresh. Local ids are

studio/frontend/src/features/chat/utils/image-input-support.ts

@@ -10,6 +10,7 @@ export function getImageInputUnavailableReason({
   externalModelLabel,
   loadedIsMultimodal,
   modelLoaded,
+  loadError,
 }: {
   activeModel?: ChatModelSummary;
   isExternalModel: boolean;
@@ -21,6 +22,8 @@ export function getImageInputUnavailableReason({
   externalModelLabel?: string | null;
   loadedIsMultimodal: boolean;
   modelLoaded: boolean;
+  // Runtime lastModelLoadError; lets the no-model branch flag a failed load.
+  loadError?: string | null;
 }): string | null {
   if (isExternalModel) {
     const explicitlyNonVision =
@@ -39,7 +42,13 @@ export function getImageInputUnavailableReason({
     }
     return null;
   }
-  if (!modelLoaded) return "Load a model before adding images.";
+  if (!modelLoaded) {
+    // Distinguish a failed load from "no model picked yet".
+    if (loadError) {
+      return "The last model failed to load. Check the server logs, then load a model before adding images.";
+    }
+    return "Load a model before adding images.";
+  }
   // loadedIsMultimodal is true for vision OR audio; that one flag can't tell
   // them apart, so only block when activeModel confirms audio-only (audio
   // capability set AND isVision === false). Otherwise trust the load