From cba73457df6ae2ba8ddcd004b68d2ee23efee290 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 21 Jun 2026 22:43:31 -0700 Subject: [PATCH] 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 | 5 ++ studio/backend/core/import_guards.py | 53 ++++++++++++++++++ studio/backend/core/inference/worker.py | 5 ++ studio/backend/core/training/trainer.py | 54 +------------------ studio/backend/core/training/worker.py | 4 ++ .../test_namespace_shadow_guard_pr6269.py | 25 +++++---- .../src/features/chat/api/chat-adapter.ts | 1 + .../features/chat/audio-attachment-adapter.ts | 5 +- .../chat/hooks/use-chat-model-runtime.ts | 6 +++ .../src/features/chat/runtime-provider.tsx | 1 + .../src/features/chat/shared-composer.tsx | 2 + .../chat/stores/chat-runtime-store.ts | 6 +++ .../chat/utils/image-input-support.ts | 11 +++- 13 files changed, 113 insertions(+), 65 deletions(-) create mode 100644 studio/backend/core/import_guards.py diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 7216221f44..f03dcfa41a 100644 --- a/studio/backend/core/export/worker.py +++ b/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 diff --git a/studio/backend/core/import_guards.py b/studio/backend/core/import_guards.py new file mode 100644 index 0000000000..5b85a96cd2 --- /dev/null +++ b/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 `/` 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 diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 44c0f5b3be..eb31ec1062 100644 --- a/studio/backend/core/inference/worker.py +++ b/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 diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 0c14061be6..bdf2a6f030 100644 --- a/studio/backend/core/training/trainer.py +++ b/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 ` 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 diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 6374589930..4fbdbf21ad 100644 --- a/studio/backend/core/training/worker.py +++ b/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, diff --git a/studio/backend/tests/test_namespace_shadow_guard_pr6269.py b/studio/backend/tests/test_namespace_shadow_guard_pr6269.py index 3ed6636bda..f77345293e 100644 --- a/studio/backend/tests/test_namespace_shadow_guard_pr6269.py +++ b/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), diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index a9097527be..a58cdd94d7 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/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); diff --git a/studio/frontend/src/features/chat/audio-attachment-adapter.ts b/studio/frontend/src/features/chat/audio-attachment-adapter.ts index 148f025e57..99a26ed178 100644 --- a/studio/frontend/src/features/chat/audio-attachment-adapter.ts +++ b/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.`; diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 134878bdac..bbf106400c 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/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, ], ); diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index d08d5ef6af..fbaa4400ed 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/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); diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 2adc3ae25a..f35ccc6149 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/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 diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index d8e578e790..2ec741b1bf 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/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((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((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 diff --git a/studio/frontend/src/features/chat/utils/image-input-support.ts b/studio/frontend/src/features/chat/utils/image-input-support.ts index d1ea31fc76..9ce6de39ba 100644 --- a/studio/frontend/src/features/chat/utils/image-input-support.ts +++ b/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