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